fix(bigquery): limit ssl error workaround to insert_rows

Author: daniel-sanche

packages/google-cloud-bigquery/google/cloud/bigquery/__init__.py

@@ -96,6 +96,7 @@
 from google.cloud.bigquery.query import StructQueryParameter
 from google.cloud.bigquery.query import StructQueryParameterType
 from google.cloud.bigquery.query import UDFResource
+from google.cloud.bigquery.retry import DEFAULT_INSERT_ROWS_RETRY
 from google.cloud.bigquery.retry import DEFAULT_RETRY
 from google.cloud.bigquery.routine import DeterminismLevel
 from google.cloud.bigquery.routine import Routine
@@ -204,6 +205,7 @@
     "ParquetOptions",
     "ScriptOptions",
     "TransactionInfo",
+    "DEFAULT_INSERT_ROWS_RETRY",
     "DEFAULT_RETRY",
     # Standard SQL types
     "StandardSqlDataType",

packages/google-cloud-bigquery/google/cloud/bigquery/client.py

@@ -111,6 +111,7 @@
 from google.cloud.bigquery.query import _QueryResults
 from google.cloud.bigquery.retry import (
     DEFAULT_GET_JOB_TIMEOUT,
+    DEFAULT_INSERT_ROWS_RETRY,
     DEFAULT_JOB_RETRY,
     DEFAULT_RETRY,
     DEFAULT_TIMEOUT,
@@ -3956,7 +3957,7 @@ def insert_rows_json(
         skip_invalid_rows: Optional[bool] = None,
         ignore_unknown_values: Optional[bool] = None,
         template_suffix: Optional[str] = None,
-        retry: retries.Retry = DEFAULT_RETRY,
+        retry: retries.Retry = DEFAULT_INSERT_ROWS_RETRY,
         timeout: TimeoutType = DEFAULT_TIMEOUT,
     ) -> Sequence[dict]:
         """Insert rows into a table without applying local type conversions.
@@ -4080,6 +4081,9 @@ def insert_rows_json(
         path = "%s/insertAll" % table.path
         # We can always retry, because every row has an insert ID.
         span_attributes = {"path": path}
+        if retry is DEFAULT_RETRY:
+            retry = DEFAULT_INSERT_ROWS_RETRY
+
         try:
             response = self._call_api(
                 retry,
@@ -4090,12 +4094,15 @@ def insert_rows_json(
                 data=data,
                 timeout=timeout,
             )
-        except requests.exceptions.SSLError as exc:
-            msg = (
-                "An SSL/Connection error occurred while streaming rows. This "
-                "could be due to an invalid request (e.g., invalid table schema)."
-            )
-            raise requests.exceptions.SSLError(msg) from exc
+        except (requests.exceptions.SSLError, core_exceptions.RetryError) as exc:
+            cause = exc.cause if isinstance(exc, core_exceptions.RetryError) else exc
+            if isinstance(cause, requests.exceptions.SSLError):
+                msg = (
+                    "An SSL/Connection error occurred while streaming rows. This "
+                    "could be due to an invalid request (e.g., invalid table schema)."
+                )
+                raise requests.exceptions.SSLError(msg) from exc
+            raise
         errors = []
 
         for error in response.get("insertErrors", ()):

packages/google-cloud-bigquery/google/cloud/bigquery/retry.py

@@ -40,11 +40,6 @@
 
 _DEFAULT_RETRY_DEADLINE = 10.0 * 60.0  # 10 minutes
 
-# Exceptions that are subclasses of types in _UNSTRUCTURED_RETRYABLE_TYPES
-# but should not be retried because they typically indicate persistent
-# configuration or security issues.
-_UNSTRUCTURED_NON_RETRYABLE_TYPES = (requests.exceptions.SSLError,)
-
 # Ambiguous errors (e.g. internalError, backendError, rateLimitExceeded) retry
 # until the full `_DEFAULT_RETRY_DEADLINE`. This is because the
 # `jobs.getQueryResults` REST API translates a job failure into an HTTP error.
@@ -69,13 +64,9 @@
 def _should_retry(exc):
     """Predicate for determining when to retry.
 
-    We retry if the 'reason' is in _RETRYABLE_REASONS or if the exception
-    is an instance of one of the _UNSTRUCTURED_RETRYABLE_TYPES, unless it
-    is explicitly excluded by being in _UNSTRUCTURED_NON_RETRYABLE_TYPES.
+    We retry if and only if the 'reason' is in _RETRYABLE_REASONS or is
+    in _UNSTRUCTURED_RETRYABLE_TYPES.
     """
-    if isinstance(exc, _UNSTRUCTURED_NON_RETRYABLE_TYPES):
-        return False
-
     try:
         reason = exc.errors[0]["reason"]
     except (AttributeError, IndexError, TypeError, KeyError):
@@ -98,6 +89,23 @@ def _should_retry(exc):
 """
 
 
+def _should_retry_insert_rows(exc):
+    """Predicate for determining when to retry streaming inserts (insertAll).
+
+    Unlike standard API calls, tabledata.insertAll failures due to schema
+    mismatches often manifest as an SSLError because the server abruptly terminates
+    the connection. These errors will not resolve on retry and should fail
+    immediately with descriptive guidance.
+    """
+    if isinstance(exc, requests.exceptions.SSLError):
+        return False
+    return _should_retry(exc)
+
+
+DEFAULT_INSERT_ROWS_RETRY = DEFAULT_RETRY.with_predicate(_should_retry_insert_rows)
+"""The default retry object for streaming inserts (insert_rows_json / insert_rows)."""
+
+
 def _should_retry_get_job_conflict(exc):
     """Predicate for determining when to retry a jobs.get call after a conflict error.
 

packages/google-cloud-bigquery/tests/system/test_ssl_retry.py

@@ -44,7 +44,7 @@ def mock_api_request(*args, **kwargs):
             bigquery_client._connection, "api_request", side_effect=mock_api_request
         ):
             # Use a reasonably short deadline for the test, although it should fail on the first attempt anyway.
-            retry = bigquery.DEFAULT_RETRY.with_deadline(5.0)
+            retry = bigquery.DEFAULT_INSERT_ROWS_RETRY.with_deadline(5.0)
 
             start_time = time.time()
             with pytest.raises(requests.exceptions.SSLError) as excinfo:

packages/google-cloud-bigquery/tests/unit/test_client.py

@@ -6898,6 +6898,38 @@ def test_insert_rows_json_w_ssl_error(self):
         self.assertIn("invalid table schema", str(context.exception))
         self.assertIn("SSL/Connection error occurred", str(context.exception))
 
+    def test_insert_rows_json_w_ssl_error_explicit_default_retry(self):
+        import requests.exceptions
+        from google.cloud.bigquery.dataset import DatasetReference
+        from google.cloud.bigquery.schema import SchemaField
+        from google.cloud.bigquery.table import Table
+        from google.cloud.bigquery.retry import DEFAULT_RETRY
+
+        PROJECT = "PROJECT"
+        DS_ID = "DS_ID"
+        TABLE_ID = "TABLE_ID"
+        ROWS = [{"full_name": "Bhettye Rhubble", "age": "27", "joined": None}]
+
+        creds = _make_credentials()
+        client = self._make_one(project=PROJECT, credentials=creds, _http=object())
+        conn = client._connection = make_connection({})
+
+        conn.api_request.side_effect = requests.exceptions.SSLError("EOF occurred")
+
+        table_ref = DatasetReference(PROJECT, DS_ID).table(TABLE_ID)
+        schema = [
+            SchemaField("full_name", "STRING", mode="REQUIRED"),
+            SchemaField("age", "INTEGER", mode="REQUIRED"),
+            SchemaField("joined", "TIMESTAMP", mode="NULLABLE"),
+        ]
+        table = Table(table_ref, schema=schema)
+
+        with self.assertRaises(requests.exceptions.SSLError) as context:
+            client.insert_rows_json(table, ROWS, retry=DEFAULT_RETRY)
+
+        self.assertIn("invalid table schema", str(context.exception))
+        self.assertIn("SSL/Connection error occurred", str(context.exception))
+
     def test_list_partitions(self):
         from google.cloud.bigquery.table import Table
 

packages/google-cloud-bigquery/tests/unit/test_retry.py

@@ -53,7 +53,7 @@ def test_w_unstructured_requests_connectionerror(self):
 
     def test_w_unstructured_requests_sslerror(self):
         exc = requests.exceptions.SSLError()
-        self.assertFalse(self._call_fut(exc))
+        self.assertTrue(self._call_fut(exc))
 
     def test_w_unstructured_requests_chunked_encoding_error(self):
         exc = requests.exceptions.ChunkedEncodingError()
@@ -160,3 +160,32 @@ def test_DEFAULT_JOB_RETRY_job_rate_limit_exceeded_retry_predicate():
     assert DEFAULT_JOB_RETRY._predicate(
         ClientError("fail", errors=[dict(reason="backendError")])
     )
+
+
+class Test_should_retry_insert_rows(unittest.TestCase):
+    def _call_fut(self, exc):
+        from google.cloud.bigquery.retry import _should_retry_insert_rows
+
+        return _should_retry_insert_rows(exc)
+
+    def test_w_unstructured_requests_sslerror(self):
+        exc = requests.exceptions.SSLError()
+        self.assertFalse(self._call_fut(exc))
+
+    def test_w_unstructured_requests_connectionerror(self):
+        exc = requests.exceptions.ConnectionError()
+        self.assertTrue(self._call_fut(exc))
+
+    def test_w_backendError(self):
+        exc = mock.Mock(errors=[{"reason": "backendError"}], spec=["errors"])
+        self.assertTrue(self._call_fut(exc))
+
+
+def test_DEFAULT_INSERT_ROWS_RETRY_predicate():
+    from google.cloud.bigquery.retry import DEFAULT_INSERT_ROWS_RETRY
+
+    exc_ssl = requests.exceptions.SSLError()
+    assert not DEFAULT_INSERT_ROWS_RETRY._predicate(exc_ssl)
+
+    exc_conn = requests.exceptions.ConnectionError()
+    assert DEFAULT_INSERT_ROWS_RETRY._predicate(exc_conn)