Trigger BQ CDC IT. Fix impl

CHANGES.md

@@ -133,6 +133,9 @@
 * (Python) DataFrame API: `agg` with a list that includes `'count'` or `'size'` now labels those rows `count`/`size` instead of `sum` ([#40504](https://github.com/apache/beam/issues/40504)).
 * (Python) `ConcatSource` no longer raises `ZeroDivisionError` when every source bundle reports a weight of 0, e.g. bundles smaller than 1 MB ([#21316](https://github.com/apache/beam/issues/21316)).
 * (Python) DataFrame API: `Series.value_counts(bins=...)` no longer always fails with a `TypeError` ([#40503](https://github.com/apache/beam/issues/40503)).
+* (Python) `ReadBigQueryChangeHistory` no longer fails at pipeline construction with `TypeCheckError: The transform 'ExecuteQueries' requires PCollections of type ... _QueryRange` when pipeline type checking is enabled (the default); a regression in 2.76.0 ([#40618](https://github.com/apache/beam/issues/40618)).
+* (Python) `ReadBigQueryChangeHistory` with `change_function='CHANGES'` no longer fails reading its temp table with a Storage Read API `500 InternalServerError`: the `_CHANGE_IS_FOR_UPDATE` pseudo-column is now excluded from the query instead of being written to the temp table ([#40618](https://github.com/apache/beam/issues/40618)).
+* (Python) `ReadBigQueryChangeHistory` now backtick-quotes the column names it puts in the generated CHANGES/APPENDS query (`columns`, `change_type_column`, `change_timestamp_column`), so reserved words and names with spaces work, and `columns=[]` selects only the pseudo-columns instead of generating invalid SQL ([#40618](https://github.com/apache/beam/issues/40618)).
 * (Python) Setting `max_batch_weight` on a RunInference `ModelHandler` no longer fails with a `TypeError`: `BatchElements` now supports `max_batch_weight`, which limits the total `element_size_fn` of a batch in addition to its number of elements ([#40644](https://github.com/apache/beam/issues/40644)).
 * (Python) `BatchElements` now honors `element_size_fn` when `max_batch_duration_secs` is set, and no longer emits an empty batch before an element that exceeds the batch size on its own ([#40644](https://github.com/apache/beam/issues/40644)).
 * (Python) `KeyedModelHandler` and `MaybeKeyedModelHandler` now apply the `element_size_fn` and `batch_length_fn` of the wrapped `ModelHandler` to examples, instead of to `(key, example)` pairs ([#40644](https://github.com/apache/beam/issues/40644)).

sdks/python/apache_beam/io/gcp/bigquery_change_history.py

@@ -299,6 +299,12 @@ def create_watermark_estimator(
     return _PollWatermarkEstimator(estimator_state)
 
 
+def _quote_identifier(name: str) -> str:
+  if '`' in name:
+    raise ValueError(f'Invalid column name (contains backtick): {name!r}')
+  return f'`{name}`'
+
+
 def build_changes_query(
     table: str,
     start: Timestamp,
@@ -334,12 +340,18 @@ def build_changes_query(
   # destination tables with their original names. Rename them so they can
   # be persisted to the temp table for Storage Read API reading.
   pseudo = (
-      f"_CHANGE_TYPE AS {change_type_column}, "
-      f"_CHANGE_TIMESTAMP AS {change_timestamp_column}")
+      f"_CHANGE_TYPE AS {_quote_identifier(change_type_column)}, "
+      f"_CHANGE_TIMESTAMP AS {_quote_identifier(change_timestamp_column)}")
   if columns is None:
-    select = f"SELECT * EXCEPT(_CHANGE_TYPE, _CHANGE_TIMESTAMP), {pseudo}"
+    hidden = "_CHANGE_TYPE, _CHANGE_TIMESTAMP"
+    if change_function == 'CHANGES':
+      hidden += ", _CHANGE_IS_FOR_UPDATE"
+    select = f"SELECT * EXCEPT({hidden}), {pseudo}"
+  elif columns:
+    quoted = ', '.join(_quote_identifier(c) for c in columns)
+    select = f"SELECT {quoted}, {pseudo}"
   else:
-    select = f"SELECT {', '.join(columns)}, {pseudo}"
+    select = f"SELECT {pseudo}"
   from_clause = (
       f"FROM {change_function}"
       f"(TABLE `{table}`, "
@@ -499,8 +511,7 @@ def _emit_query_ranges(
       self,
       start_ts: Timestamp,
       end_ts: Timestamp,
-      watermark_estimator: _PollWatermarkEstimator
-  ) -> Iterable[TimestampedValue[_QueryRange]]:
+      watermark_estimator: _PollWatermarkEstimator):
     """Compute and yield _QueryRange elements, advancing estimator state."""
     ranges = compute_ranges(start_ts, end_ts, self._change_function)
     _LOGGER.info(
@@ -527,9 +538,7 @@ def process(
       _: _PollConfig,
       restriction_tracker=beam.DoFn.RestrictionParam(),
       watermark_estimator=beam.DoFn.WatermarkEstimatorParam(
-          _PollWatermarkEstimatorProvider())
-  ) -> Iterable[TimestampedValue[_QueryRange]]:
-
+          _PollWatermarkEstimatorProvider())):
     now = time.time()
     start_ts = watermark_estimator.poll_cursor()
 
@@ -1041,10 +1050,7 @@ class _DecompressArrowBatchesFn(beam.DoFn):
   def __init__(self, change_timestamp_column: str = 'change_timestamp') -> None:
     self._change_timestamp_column = change_timestamp_column
 
-  def process(
-      self,
-      element: tuple[bytes,
-                     bytes]) -> Iterable[TimestampedValue[dict[str, Any]]]:
+  def process(self, element: tuple[bytes, bytes]):
     schema_bytes, batch_bytes = element
     schema = pyarrow.ipc.read_schema(pyarrow.py_buffer(schema_bytes))
     batch = pyarrow.ipc.read_record_batch(
@@ -1307,7 +1313,8 @@ def expand(self, pbegin: beam.pvalue.PBegin) -> beam.PCollection:
                 stop_time=stop_time,
                 poll_interval=poll_interval,
                 location=self._location,
-                quota_project_id=quota_project_id)))
+                quota_project_id=quota_project_id)).with_output_types(
+                    _QueryRange))
 
     # CommitQueryResults: Reshuffle commits _QueryResult (temp table ref)
     # so that if the Read SDF retries, it re-reads the existing temp table
@@ -1364,7 +1371,8 @@ def expand(self, pbegin: beam.pvalue.PBegin) -> beam.PCollection:
           | 'ReshuffleForFanout' >> beam.Reshuffle()
           | 'DecompressBatches' >> beam.ParDo(
               _DecompressArrowBatchesFn(
-                  change_timestamp_column=(self._change_timestamp_column))))
+                  change_timestamp_column=self._change_timestamp_column)).
+          with_output_types(dict[str, Any]))
       return rows
     else:
       return read_outputs['rows']

sdks/python/apache_beam/io/gcp/bigquery_change_history_it_test.py

@@ -24,6 +24,8 @@
 import unittest
 import uuid
 
+import pytest
+
 import apache_beam as beam
 from apache_beam.io.gcp import bigquery_tools
 from apache_beam.io.gcp.bigquery_change_history import ReadBigQueryChangeHistory
@@ -168,6 +170,7 @@ def _run_dml(cls, sql):
 
 class CleanupTempTablesFnTest(BigQueryChangeHistoryIntegrationBase):
   """Integration tests for _CleanupTempTablesFn against real BigQuery."""
+  @pytest.mark.it_postcommit
   def test_single_complete_signal_deletes_table(self):
     """A single signal with streams_read == total deletes the temp table."""
     table_id = f'cleanup_test_{secrets.token_hex(4)}'
@@ -189,6 +192,7 @@ def test_single_complete_signal_deletes_table(self):
     with self.assertRaises(Exception):
       self.bq_wrapper.get_table(self.project, self.temp_dataset, table_id)
 
+  @pytest.mark.it_postcommit
   def test_partial_signals_then_complete(self):
     """Partial signals don't delete; final signal triggers cleanup."""
     table_id = f'cleanup_partial_{secrets.token_hex(4)}'
@@ -228,6 +232,7 @@ class ReadStorageStreamsSDFTest(BigQueryChangeHistoryIntegrationBase):
       ('change_type', 'STRING'),
   ]
 
+  @pytest.mark.it_postcommit
   def test_reads_rows_from_temp_table(self):
     """SDF reads rows from a real temp table via Storage Read API."""
     table_id = f'sdf_test_{secrets.token_hex(4)}'
@@ -277,6 +282,7 @@ def test_reads_rows_from_temp_table(self):
           | 'CountRows' >> beam.combiners.Count.Globally())
       assert_that(row_count, equal_to([3]), label='CheckRowCount')
 
+  @pytest.mark.it_postcommit
   def test_cleanup_signal_emitted(self):
     """SDF emits cleanup signal with correct counts."""
     table_id = f'sdf_cleanup_{secrets.token_hex(4)}'
@@ -304,15 +310,27 @@ def test_cleanup_signal_emitted(self):
           | beam.ParDo(_ReadStorageStreamsSDF()).with_outputs(
               'cleanup', main='rows'))
 
-      # Verify cleanup signal
-      cleanup_table_keys = (
-          outputs['cleanup']
-          | 'ExtractKey' >> beam.Map(lambda x: x[0]))
+      # Each stream split emits its own (table_key, (streams_read, total))
+      # signal; _CleanupTempTablesFn sums streams_read across them.
+      expected_key = bigquery_tools.get_hashable_destination(table_ref)
+
+      def check_cleanup_signals(signals):
+        assert signals, 'Expected at least one cleanup signal'
+        keys = {key for key, _ in signals}
+        assert keys == {expected_key}, f'Unexpected table keys: {keys}'
+        totals = {total for _, (_, total) in signals}
+        assert len(totals) == 1, f'Inconsistent total_streams: {totals}'
+        total_streams = totals.pop()
+        streams_read = sum(count for _, (count, _) in signals)
+        assert streams_read == total_streams, (
+            f'streams_read={streams_read} != total_streams={total_streams}')
+
       assert_that(
-          cleanup_table_keys,
-          equal_to([bigquery_tools.get_hashable_destination(table_ref)]),
-          label='CheckCleanupKey')
+          outputs['cleanup'],
+          check_cleanup_signals,
+          label='CheckCleanupSignals')
 
+  @pytest.mark.it_postcommit
   def test_empty_table(self):
     """Empty table produces 0 rows and cleanup signal."""
     table_id = f'sdf_empty_{secrets.token_hex(4)}'
@@ -340,6 +358,7 @@ def test_empty_table(self):
 
 class PollChangeHistoryFnTest(BigQueryChangeHistoryIntegrationBase):
   """Integration test for _PollChangeHistoryFn in isolation."""
+  @pytest.mark.it_postcommit
   def test_poll_emits_query_ranges(self):
     """Poll SDF emits _QueryRange elements with valid time ranges."""
     table_str = f'{self.project}:{self.dataset}.nonexistent'
@@ -389,6 +408,7 @@ def setUpClass(cls):
     cls.insert_time = time.time()
     time.sleep(10)
 
+  @pytest.mark.it_postcommit
   def test_execute_query_produces_query_result(self):
     """ExecuteQueryFn creates a temp table from a _QueryRange."""
     table_str = f'{self.project}:{self.dataset}.{self.test_table_id}'
@@ -439,6 +459,7 @@ def setUpClass(cls):
     _LOGGER.info('Waiting for change history propagation...')
     time.sleep(15)
 
+  @pytest.mark.it_postcommit
   def test_changes_captures_insert_update_delete(self):
     """ReadBigQueryChangeHistory with CHANGES sees all mutation types."""
     table_str = f'{self.project}:{self.dataset}.{self.test_table_id}'
@@ -531,6 +552,7 @@ def setUpClass(cls):
     _LOGGER.info('Waiting for streaming buffer to flush...')
     time.sleep(15)
 
+  @pytest.mark.it_postcommit
   def test_public_api_reads_inserted_rows(self):
     """ReadBigQueryChangeHistory PTransform with polling SDF."""
     table_str = f'{self.project}:{self.dataset}.{self.test_table_id}'
@@ -564,6 +586,7 @@ def check_rows(actual):
 
       assert_that(rows, check_rows)
 
+  @pytest.mark.it_postcommit
   def test_public_api_reads_inserted_row_inline_decompress(self):
     """ReadBigQueryChangeHistory with inline decompression (no reshuffle)."""
     table_str = f'{self.project}:{self.dataset}.{self.test_table_id}'
@@ -598,6 +621,7 @@ def check_rows(actual):
 
       assert_that(rows, check_rows)
 
+  @pytest.mark.it_postcommit
   def test_public_api_reads_inserted_row_with_split_streams(self):
     """ReadBigQueryChangeHistory with max_split_rounds=1."""
     table_str = f'{self.project}:{self.dataset}.{self.test_table_id}'

sdks/python/apache_beam/io/gcp/bigquery_change_history_test.py

@@ -28,6 +28,8 @@
 import unittest
 from unittest import mock
 
+import apache_beam as beam
+from apache_beam import typehints
 from apache_beam.io.gcp import bigquery_change_history as bqch
 from apache_beam.io.gcp import bigquery_tools
 from apache_beam.io.gcp.bigquery_change_history import ReadBigQueryChangeHistory
@@ -73,14 +75,55 @@ def test_changes_query_format(self):
     self.assertIn('CHANGES', sql)
     self.assertIn('TABLE `proj.ds.tbl`', sql)
 
+  def test_changes_hides_is_for_update(self):
+    start = _ts(2025, 1, 1)
+    end = _ts(2025, 1, 2)
+    sql = build_changes_query('proj.ds.tbl', start, end, 'CHANGES')
+    self.assertIn(
+        'EXCEPT(_CHANGE_TYPE, _CHANGE_TIMESTAMP, _CHANGE_IS_FOR_UPDATE)', sql)
+
+  def test_appends_has_no_is_for_update(self):
+    """APPENDS() has no _CHANGE_IS_FOR_UPDATE; EXCEPT on it would fail."""
+    start = _ts(2025, 1, 1)
+    end = _ts(2025, 1, 2)
+    sql = build_changes_query('proj.ds.tbl', start, end, 'APPENDS')
+    self.assertIn('EXCEPT(_CHANGE_TYPE, _CHANGE_TIMESTAMP)', sql)
+    self.assertNotIn('_CHANGE_IS_FOR_UPDATE', sql)
+
   def test_columns_select(self):
     start = _ts(2025, 1, 1)
     end = _ts(2025, 1, 2)
     sql = build_changes_query(
         'proj.ds.tbl', start, end, 'APPENDS', columns=['col_a', 'col_b'])
-    self.assertIn('SELECT col_a, col_b, _CHANGE_TYPE AS', sql)
+    self.assertIn('SELECT `col_a`, `col_b`, _CHANGE_TYPE AS', sql)
     self.assertNotIn('EXCEPT', sql)
 
+  def test_columns_empty_selects_only_pseudo_columns(self):
+    start = _ts(2025, 1, 1)
+    end = _ts(2025, 1, 2)
+    sql = build_changes_query('proj.ds.tbl', start, end, 'APPENDS', columns=[])
+    self.assertIn('SELECT _CHANGE_TYPE AS', sql)
+
+  def test_pseudo_column_names_quoted(self):
+    start = _ts(2025, 1, 1)
+    end = _ts(2025, 1, 2)
+    sql = build_changes_query(
+        'proj.ds.tbl',
+        start,
+        end,
+        'APPENDS',
+        change_type_column='op type',
+        change_timestamp_column='from')
+    self.assertIn('_CHANGE_TYPE AS `op type`', sql)
+    self.assertIn('_CHANGE_TIMESTAMP AS `from`', sql)
+
+  def test_column_with_backtick_rejected(self):
+    start = _ts(2025, 1, 1)
+    end = _ts(2025, 1, 2)
+    with self.assertRaises(ValueError):
+      build_changes_query(
+          'proj.ds.tbl', start, end, 'APPENDS', columns=['a` FROM x --'])
+
   def test_columns_none_selects_all(self):
     start = _ts(2025, 1, 1)
     end = _ts(2025, 1, 2)
@@ -112,7 +155,7 @@ def test_columns_and_row_filter(self):
         'CHANGES',
         columns=['id', 'name'],
         row_filter='id > 100')
-    self.assertIn('SELECT id, name, _CHANGE_TYPE AS', sql)
+    self.assertIn('SELECT `id`, `name`, _CHANGE_TYPE AS', sql)
     self.assertNotIn('EXCEPT', sql)
     self.assertIn('WHERE id > 100', sql)
 
@@ -280,5 +323,28 @@ def test_expand_resolves_quota_project_from_options(self):
           m.call_args.kwargs['quota_project_id'], self._QUOTA, name)
 
 
+@unittest.skipIf(HttpError is None, 'GCP dependencies are not installed')
+class PipelineConstructionTest(unittest.TestCase):
+  """Applies the transform with pipeline type checking on (the default)."""
+  def _apply(self, reshuffle_decompress):
+    p = beam.Pipeline()
+    return p | ReadBigQueryChangeHistory(
+        table='p:d.t',
+        project='p',
+        temp_dataset='tmp',
+        location='US',
+        start_time=1700000000,
+        stop_time=1700000060,
+        reshuffle_decompress=reshuffle_decompress)
+
+  def test_reshuffle_decompress_output_type(self):
+    rows = self._apply(reshuffle_decompress=True)
+    self.assertEqual(rows.element_type, typehints.Dict[str, typehints.Any])
+
+  def test_inline_decompress_output_type(self):
+    rows = self._apply(reshuffle_decompress=False)
+    self.assertEqual(rows.element_type, typehints.Dict[str, typehints.Any])
+
+
 if __name__ == '__main__':
   unittest.main()