feat: make it easier to disable best-effort deduplication with streaming inserts (#734)

* feat: make it easier to disable row insert IDs

* Also accept any iterables for row_ids

Author: plamut

google/cloud/bigquery/__init__.py

@@ -37,6 +37,7 @@
 from google.cloud.bigquery.dataset import Dataset
 from google.cloud.bigquery.dataset import DatasetReference
 from google.cloud.bigquery import enums
+from google.cloud.bigquery.enums import AutoRowIDs
 from google.cloud.bigquery.enums import KeyResultStatementKind
 from google.cloud.bigquery.enums import SqlTypeNames
 from google.cloud.bigquery.enums import StandardSqlDataTypes
@@ -144,6 +145,7 @@
     "DEFAULT_RETRY",
     # Enum Constants
     "enums",
+    "AutoRowIDs",
     "Compression",
     "CreateDisposition",
     "DestinationFormat",

google/cloud/bigquery/client.py

@@ -68,6 +68,7 @@
 from google.cloud.bigquery.dataset import Dataset
 from google.cloud.bigquery.dataset import DatasetListItem
 from google.cloud.bigquery.dataset import DatasetReference
+from google.cloud.bigquery.enums import AutoRowIDs
 from google.cloud.bigquery.exceptions import LegacyBigQueryStorageError
 from google.cloud.bigquery.opentelemetry_tracing import create_span
 from google.cloud.bigquery import job
@@ -3349,7 +3350,7 @@ def insert_rows_json(
         self,
         table: Union[Table, TableReference, str],
         json_rows: Sequence[Dict],
-        row_ids: Sequence[str] = None,
+        row_ids: Union[Iterable[str], AutoRowIDs, None] = AutoRowIDs.GENERATE_UUID,
         skip_invalid_rows: bool = None,
         ignore_unknown_values: bool = None,
         template_suffix: str = None,
@@ -3371,11 +3372,20 @@ def insert_rows_json(
             json_rows (Sequence[Dict]):
                 Row data to be inserted. Keys must match the table schema fields
                 and values must be JSON-compatible representations.
-            row_ids (Optional[Sequence[Optional[str]]]):
+            row_ids (Union[Iterable[str], AutoRowIDs, None]):
                 Unique IDs, one per row being inserted. An ID can also be
                 ``None``, indicating that an explicit insert ID should **not**
                 be used for that row. If the argument is omitted altogether,
                 unique IDs are created automatically.
+
+                .. versionchanged:: 2.21.0
+                    Can also be an iterable, not just a sequence, or an
+                    :class:`AutoRowIDs` enum member.
+
+                .. deprecated:: 2.21.0
+                    Passing ``None`` to explicitly request autogenerating insert IDs is
+                    deprecated, use :attr:`AutoRowIDs.GENERATE_UUID` instead.
+
             skip_invalid_rows (Optional[bool]):
                 Insert all valid rows of a request, even if invalid rows exist.
                 The default value is ``False``, which causes the entire request
@@ -3415,12 +3425,37 @@ def insert_rows_json(
         rows_info = []
         data = {"rows": rows_info}
 
-        for index, row in enumerate(json_rows):
+        if row_ids is None:
+            warnings.warn(
+                "Passing None for row_ids is deprecated. To explicitly request "
+                "autogenerated insert IDs, use AutoRowIDs.GENERATE_UUID instead",
+                category=DeprecationWarning,
+            )
+            row_ids = AutoRowIDs.GENERATE_UUID
+
+        if not isinstance(row_ids, AutoRowIDs):
+            try:
+                row_ids_iter = iter(row_ids)
+            except TypeError:
+                msg = "row_ids is neither an iterable nor an AutoRowIDs enum member"
+                raise TypeError(msg)
+
+        for i, row in enumerate(json_rows):
             info = {"json": row}
-            if row_ids is not None:
-                info["insertId"] = row_ids[index]
-            else:
+
+            if row_ids is AutoRowIDs.GENERATE_UUID:
                 info["insertId"] = str(uuid.uuid4())
+            elif row_ids is AutoRowIDs.DISABLED:
+                info["insertId"] = None
+            else:
+                try:
+                    insert_id = next(row_ids_iter)
+                except StopIteration:
+                    msg = f"row_ids did not generate enough IDs, error at index {i}"
+                    raise ValueError(msg)
+                else:
+                    info["insertId"] = insert_id
+
             rows_info.append(info)
 
         if skip_invalid_rows is not None:

google/cloud/bigquery/enums.py

@@ -21,6 +21,13 @@
 from google.cloud.bigquery.query import ScalarQueryParameterType
 
 
+class AutoRowIDs(enum.Enum):
+    """How to handle automatic insert IDs when inserting rows as a stream."""
+
+    DISABLED = enum.auto()
+    GENERATE_UUID = enum.auto()
+
+
 class Compression(object):
     """The compression type to use for exported files. The default value is
     :attr:`NONE`.

tests/unit/test_client.py

@@ -5434,7 +5434,7 @@ def test_insert_rows_from_dataframe_w_explicit_none_insert_ids(self):
             method="POST", path=API_PATH, data=EXPECTED_SENT_DATA, timeout=None
         )
 
-    def test_insert_rows_json(self):
+    def test_insert_rows_json_default_behavior(self):
         from google.cloud.bigquery.dataset import DatasetReference
         from google.cloud.bigquery.schema import SchemaField
         from google.cloud.bigquery.table import Table
@@ -5481,29 +5481,127 @@ def test_insert_rows_json(self):
             method="POST", path="/%s" % PATH, data=SENT, timeout=7.5,
         )
 
-    def test_insert_rows_json_with_string_id(self):
-        rows = [{"col1": "val1"}]
+    def test_insert_rows_json_w_explicitly_requested_autogenerated_insert_ids(self):
+        from google.cloud.bigquery import AutoRowIDs
+
+        rows = [{"col1": "val1"}, {"col2": "val2"}]
         creds = _make_credentials()
         http = object()
         client = self._make_one(
             project="default-project", credentials=creds, _http=http
         )
         conn = client._connection = make_connection({})
 
-        with mock.patch("uuid.uuid4", side_effect=map(str, range(len(rows)))):
-            errors = client.insert_rows_json("proj.dset.tbl", rows)
+        uuid_patcher = mock.patch("uuid.uuid4", side_effect=map(str, range(len(rows))))
+        with uuid_patcher:
+            errors = client.insert_rows_json(
+                "proj.dset.tbl", rows, row_ids=AutoRowIDs.GENERATE_UUID
+            )
 
         self.assertEqual(len(errors), 0)
-        expected = {
-            "rows": [{"json": row, "insertId": str(i)} for i, row in enumerate(rows)]
+
+        # Check row data sent to the backend.
+        expected_row_data = {
+            "rows": [
+                {"json": {"col1": "val1"}, "insertId": "0"},
+                {"json": {"col2": "val2"}, "insertId": "1"},
+            ]
         }
         conn.api_request.assert_called_once_with(
             method="POST",
             path="/projects/proj/datasets/dset/tables/tbl/insertAll",
-            data=expected,
+            data=expected_row_data,
+            timeout=None,
+        )
+
+    def test_insert_rows_json_w_explicitly_disabled_insert_ids(self):
+        from google.cloud.bigquery import AutoRowIDs
+
+        rows = [{"col1": "val1"}, {"col2": "val2"}]
+        creds = _make_credentials()
+        http = object()
+        client = self._make_one(
+            project="default-project", credentials=creds, _http=http
+        )
+        conn = client._connection = make_connection({})
+
+        errors = client.insert_rows_json(
+            "proj.dset.tbl", rows, row_ids=AutoRowIDs.DISABLED,
+        )
+
+        self.assertEqual(len(errors), 0)
+
+        expected_row_data = {
+            "rows": [
+                {"json": {"col1": "val1"}, "insertId": None},
+                {"json": {"col2": "val2"}, "insertId": None},
+            ]
+        }
+        conn.api_request.assert_called_once_with(
+            method="POST",
+            path="/projects/proj/datasets/dset/tables/tbl/insertAll",
+            data=expected_row_data,
+            timeout=None,
+        )
+
+    def test_insert_rows_json_with_iterator_row_ids(self):
+        rows = [{"col1": "val1"}, {"col2": "val2"}, {"col3": "val3"}]
+        creds = _make_credentials()
+        http = object()
+        client = self._make_one(
+            project="default-project", credentials=creds, _http=http
+        )
+        conn = client._connection = make_connection({})
+
+        row_ids_iter = map(str, itertools.count(42))
+        errors = client.insert_rows_json("proj.dset.tbl", rows, row_ids=row_ids_iter)
+
+        self.assertEqual(len(errors), 0)
+        expected_row_data = {
+            "rows": [
+                {"json": {"col1": "val1"}, "insertId": "42"},
+                {"json": {"col2": "val2"}, "insertId": "43"},
+                {"json": {"col3": "val3"}, "insertId": "44"},
+            ]
+        }
+        conn.api_request.assert_called_once_with(
+            method="POST",
+            path="/projects/proj/datasets/dset/tables/tbl/insertAll",
+            data=expected_row_data,
             timeout=None,
         )
 
+    def test_insert_rows_json_with_non_iterable_row_ids(self):
+        rows = [{"col1": "val1"}]
+        creds = _make_credentials()
+        http = object()
+        client = self._make_one(
+            project="default-project", credentials=creds, _http=http
+        )
+        client._connection = make_connection({})
+
+        with self.assertRaises(TypeError) as exc:
+            client.insert_rows_json("proj.dset.tbl", rows, row_ids=object())
+
+        err_msg = str(exc.exception)
+        self.assertIn("row_ids", err_msg)
+        self.assertIn("iterable", err_msg)
+
+    def test_insert_rows_json_with_too_few_row_ids(self):
+        rows = [{"col1": "val1"}, {"col2": "val2"}, {"col3": "val3"}]
+        creds = _make_credentials()
+        http = object()
+        client = self._make_one(
+            project="default-project", credentials=creds, _http=http
+        )
+        client._connection = make_connection({})
+
+        insert_ids = ["10", "20"]
+
+        error_msg_pattern = "row_ids did not generate enough IDs.*index 2"
+        with self.assertRaisesRegex(ValueError, error_msg_pattern):
+            client.insert_rows_json("proj.dset.tbl", rows, row_ids=insert_ids)
+
     def test_insert_rows_json_w_explicit_none_insert_ids(self):
         rows = [{"col1": "val1"}, {"col2": "val2"}]
         creds = _make_credentials()
@@ -5526,6 +5624,45 @@ def test_insert_rows_json_w_explicit_none_insert_ids(self):
             timeout=None,
         )
 
+    def test_insert_rows_json_w_none_insert_ids_sequence(self):
+        rows = [{"col1": "val1"}, {"col2": "val2"}]
+        creds = _make_credentials()
+        http = object()
+        client = self._make_one(
+            project="default-project", credentials=creds, _http=http
+        )
+        conn = client._connection = make_connection({})
+
+        uuid_patcher = mock.patch("uuid.uuid4", side_effect=map(str, range(len(rows))))
+        with warnings.catch_warnings(record=True) as warned, uuid_patcher:
+            errors = client.insert_rows_json("proj.dset.tbl", rows, row_ids=None)
+
+        self.assertEqual(len(errors), 0)
+
+        # Passing row_ids=None should have resulted in a deprecation warning.
+        matches = [
+            warning
+            for warning in warned
+            if issubclass(warning.category, DeprecationWarning)
+            and "row_ids" in str(warning)
+            and "AutoRowIDs.GENERATE_UUID" in str(warning)
+        ]
+        assert matches, "The expected deprecation warning was not raised."
+
+        # Check row data sent to the backend.
+        expected_row_data = {
+            "rows": [
+                {"json": {"col1": "val1"}, "insertId": "0"},
+                {"json": {"col2": "val2"}, "insertId": "1"},
+            ]
+        }
+        conn.api_request.assert_called_once_with(
+            method="POST",
+            path="/projects/proj/datasets/dset/tables/tbl/insertAll",
+            data=expected_row_data,
+            timeout=None,
+        )
+
     def test_insert_rows_w_wrong_arg(self):
         from google.cloud.bigquery.dataset import DatasetReference
         from google.cloud.bigquery.schema import SchemaField