This is an automated email from the ASF dual-hosted git repository.
claudevdm pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new eb8f207402d Trigger BQ CDC IT. Fix impl (#40618)
eb8f207402d is described below
commit eb8f207402d446ca1632438a1a179164af3667bc
Author: claudevdm <[email protected]>
AuthorDate: Fri Oct 23 17:24:53 2026 -0400
Trigger BQ CDC IT. Fix impl (#40618)
* Trigger BQ CDC IT. Fix impl
* comments
---
CHANGES.md | 3 +
.../apache_beam/io/gcp/bigquery_change_history.py | 38 +++++++-----
.../io/gcp/bigquery_change_history_it_test.py | 38 +++++++++---
.../io/gcp/bigquery_change_history_test.py | 70 +++++++++++++++++++++-
4 files changed, 125 insertions(+), 24 deletions(-)
diff --git a/CHANGES.md b/CHANGES.md
index c9350da11d6..8fa59e7d5ea 100644
--- a/CHANGES.md
+++ b/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)).
diff --git a/sdks/python/apache_beam/io/gcp/bigquery_change_history.py
b/sdks/python/apache_beam/io/gcp/bigquery_change_history.py
index 9f836190d9c..f6e86b07e71 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery_change_history.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery_change_history.py
@@ -299,6 +299,12 @@ class
_PollWatermarkEstimatorProvider(WatermarkEstimatorProvider):
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 @@ class _PollChangeHistoryFn(beam.DoFn,
beam.transforms.core.RestrictionProvider):
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 @@ class _PollChangeHistoryFn(beam.DoFn,
beam.transforms.core.RestrictionProvider):
_: _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 @@ class ReadBigQueryChangeHistory(beam.PTransform):
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 @@ class ReadBigQueryChangeHistory(beam.PTransform):
| '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']
diff --git a/sdks/python/apache_beam/io/gcp/bigquery_change_history_it_test.py
b/sdks/python/apache_beam/io/gcp/bigquery_change_history_it_test.py
index 8cc7c7bb9c3..5bf2486f773 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery_change_history_it_test.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery_change_history_it_test.py
@@ -24,6 +24,8 @@ import time
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 @@ class
BigQueryChangeHistoryIntegrationBase(unittest.TestCase):
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 @@ class
CleanupTempTablesFnTest(BigQueryChangeHistoryIntegrationBase):
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 @@ class
ReadStorageStreamsSDFTest(BigQueryChangeHistoryIntegrationBase):
| '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 @@ class
ReadStorageStreamsSDFTest(BigQueryChangeHistoryIntegrationBase):
| 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 @@ class
ReadStorageStreamsSDFTest(BigQueryChangeHistoryIntegrationBase):
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 @@ class
ExecuteQueryFnTest(BigQueryChangeHistoryIntegrationBase):
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 @@ class
ChangesEndToEndTest(BigQueryChangeHistoryIntegrationBase):
_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 @@ class EndToEndTest(BigQueryChangeHistoryIntegrationBase):
_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 @@ class EndToEndTest(BigQueryChangeHistoryIntegrationBase):
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 @@ class EndToEndTest(BigQueryChangeHistoryIntegrationBase):
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}'
diff --git a/sdks/python/apache_beam/io/gcp/bigquery_change_history_test.py
b/sdks/python/apache_beam/io/gcp/bigquery_change_history_test.py
index e1b1eb69ae3..20d614c9362 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery_change_history_test.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery_change_history_test.py
@@ -28,6 +28,8 @@ import datetime
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 @@ class BuildChangesQueryTest(unittest.TestCase):
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 @@ class BuildChangesQueryTest(unittest.TestCase):
'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 @@ class QuotaProjectTest(unittest.TestCase):
m.call_args.kwargs['quota_project_id'], self._QUOTA, name)
[email protected](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()
