diff --git a/CHANGES.md b/CHANGES.md index c9350da11d69..8fa59e7d5ea0 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 9f836190d9ca..f6e86b07e71f 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 @@ 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'] 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 8cc7c7bb9c34..5bf2486f7734 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 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}' 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 e1b1eb69ae38..20d614c93622 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 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()