diff --git a/.github/trigger_files/beam_PostCommit_Python_Xlang_Gcp_Dataflow.json b/.github/trigger_files/beam_PostCommit_Python_Xlang_Gcp_Dataflow.json index 83346d34aee0..48379c82f94c 100644 --- a/.github/trigger_files/beam_PostCommit_Python_Xlang_Gcp_Dataflow.json +++ b/.github/trigger_files/beam_PostCommit_Python_Xlang_Gcp_Dataflow.json @@ -1,4 +1,4 @@ { "comment": "Modify this file in a trivial way to cause this test suite to run", - "modification": 16 + "modification": 21 } diff --git a/sdks/python/apache_beam/io/external/xlang_bigqueryio_it_test.py b/sdks/python/apache_beam/io/external/xlang_bigqueryio_it_test.py index 49725d54e990..da8daf75dec3 100644 --- a/sdks/python/apache_beam/io/external/xlang_bigqueryio_it_test.py +++ b/sdks/python/apache_beam/io/external/xlang_bigqueryio_it_test.py @@ -31,6 +31,7 @@ from hamcrest.core.core.allof import all_of import apache_beam as beam +from apache_beam.io.gcp import bigquery from apache_beam.io.gcp.bigquery import StorageWriteToBigQuery from apache_beam.io.gcp.bigquery_tools import BigQueryWrapper from apache_beam.io.gcp.tests.bigquery_matcher import BigqueryFullResultMatcher @@ -483,6 +484,74 @@ def test_write_to_dynamic_destinations(self): use_at_least_once=False)) hamcrest_assert(p, all_of(*bq_matchers)) + def test_write_to_dynamic_destinations_with_dynamic_schema(self): + base_table_spec = '{}.dynamic_dest_dyn_schema_'.format(self.dataset_id) + spec_with_project = '{}:{}'.format(self.project, base_table_spec) + table_a = base_table_spec + 'users' + table_b = base_table_spec + 'scores' + + schema_a = "id:INTEGER,name:STRING" + schema_b = "id:INTEGER,score:INTEGER,active:BOOLEAN" + + elements_a = [ + { + 'id': 1, 'name': 'alice' + }, + { + 'id': 2, 'name': 'bob' + }, + ] + elements_b = [ + { + 'id': 101, 'score': 95, 'active': True + }, + { + 'id': 102, 'score': 80, 'active': False + }, + ] + elements = elements_a + elements_b + + schema_map = { + spec_with_project + 'users': schema_a, + spec_with_project + 'scores': schema_b, + } + + bq_matchers = [ + BigqueryFullResultMatcher( + project=self.project, + query="SELECT id, name FROM %s" % table_a, + data=self.parse_expected_data(elements_a)), + BigqueryFullResultMatcher( + project=self.project, + query="SELECT id, score, active FROM %s" % table_b, + data=self.parse_expected_data(elements_b)), + ] + + def get_destination(record): + if 'name' in record: + return spec_with_project + 'users' + return spec_with_project + 'scores' + + def get_schema_raw(dest, side_map): + return side_map[dest] + + get_schema = bigquery.dynamic_schema( + get_schema_raw, + union_schema="id:INTEGER,name:STRING,score:INTEGER,active:BOOLEAN") + + with beam.Pipeline(argv=self.args) as p: + schema_pc = p | "CreateSchema" >> beam.Create([schema_map]) + _ = ( + p + | "CreateElements" >> beam.Create(elements) + | beam.io.WriteToBigQuery( + table=get_destination, + method=beam.io.WriteToBigQuery.Method.STORAGE_WRITE_API, + schema=get_schema, + schema_side_inputs=(beam.pvalue.AsSingleton(schema_pc), ), + use_at_least_once=False)) + hamcrest_assert(p, all_of(*bq_matchers)) + def test_write_to_dynamic_destinations_with_beam_rows(self): base_table_spec = '{}.dynamic_dest_'.format(self.dataset_id) spec_with_project = '{}:{}'.format(self.project, base_table_spec) diff --git a/sdks/python/apache_beam/io/gcp/bigquery.py b/sdks/python/apache_beam/io/gcp/bigquery.py index a2d17f12569e..d4013095a717 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery.py +++ b/sdks/python/apache_beam/io/gcp/bigquery.py @@ -197,6 +197,43 @@ def compute_table_name(row): a tuple of PCollectionViews to be passed to the schema callable (much like the `table_side_inputs` parameter). +Dynamic Schemas with Storage Write API +-------------------------------------- +When writing to dynamic destinations with `method=STORAGE_WRITE_API`, a union schema +containing all fields across destination tables is required at the PCollection level +for cross-language type inference and runtime row serialization. + +The recommended best-practice is to use the `dynamic_schema` helper: + +* **Using a dictionary map**: If schemas are known at pipeline construction time, pass + a dictionary mapping destinations to schemas. The helper automatically infers and merges + all fields into the required union schema:: + + schema_map = { + 'my_project:dataset.users': 'id:INTEGER,name:STRING', + 'my_project:dataset.scores': 'id:INTEGER,score:INTEGER,active:BOOLEAN' + } + + elements | WriteToBigQuery( + table=get_destination, + method=WriteToBigQuery.Method.STORAGE_WRITE_API, + schema=dynamic_schema(schema_map)) + +* **Using a callable function or side inputs**: If schemas are determined dynamically + via a callable function, wrap the callable with `dynamic_schema` and explicitly pass + `union_schema`:: + + def get_schema(destination, schema_dict): + return schema_dict[destination] + + elements | WriteToBigQuery( + table=get_destination, + method=WriteToBigQuery.Method.STORAGE_WRITE_API, + schema=dynamic_schema( + get_schema, + union_schema='id:INTEGER,name:STRING,score:INTEGER,active:BOOLEAN'), + schema_side_inputs=(schema_dict_side_input,)) + Additional Parameters for BigQuery Tables ----------------------------------------- @@ -356,6 +393,7 @@ def chain_after(result): # pytype: skip-file import collections +import copy import io import itertools import json @@ -1956,6 +1994,120 @@ def _restore_table_ref(sharded_table_ref_elems_kv): SCHEMA_AUTODETECT = 'SCHEMA_AUTODETECT' +def dynamic_schema(schema_fn_or_map, union_schema=None): + """Helper to construct a dynamic schema callable with a union schema hint. + + When using the BigQuery Storage Write API (`method=STORAGE_WRITE_API`) with + dynamic destinations, the cross-language transform requires a PCollection-level + union schema containing all fields across all target tables for protobuf + serialization and type inference. + + This helper provides the recommended best practice for constructing dynamic + schemas: + + 1. **Dictionary Map**: If destination table schemas are provided as a dictionary + mapping table names/specs to schemas (str, dict, or TableSchema), this helper + automatically merges all fields into a single union schema. + 2. **Callable**: If a callable function is used, this helper attaches the provided + `union_schema` to the callable as a schema hint (`_union_schema`). + + Example using a dictionary map (union schema is auto-inferred):: + + schema_map = { + 'project:dataset.users': 'id:INTEGER,name:STRING', + 'project:dataset.scores': 'id:INTEGER,score:INTEGER,active:BOOLEAN' + } + + def get_destination(record): + if 'name' in record: + return 'project:dataset.users' + return 'project:dataset.scores' + + schema_callable = dynamic_schema(schema_map) + + elements | WriteToBigQuery( + table=get_destination, + method=WriteToBigQuery.Method.STORAGE_WRITE_API, + schema=schema_callable) + + Example using a callable with explicit union schema:: + + def get_schema(destination, schema_side_input): + return schema_side_input[destination] + + schema_callable = dynamic_schema( + get_schema, + union_schema='id:INTEGER,name:STRING,score:INTEGER,active:BOOLEAN') + + elements | WriteToBigQuery( + table=get_destination, + method=WriteToBigQuery.Method.STORAGE_WRITE_API, + schema=schema_callable, + schema_side_inputs=(schema_side_input,)) + + Args: + schema_fn_or_map: A callable `(destination, *side_inputs) -> schema` + or a dictionary mapping destination strings to schemas (str, dict, or + TableSchema). + union_schema: (Optional) The union schema containing all fields across + target tables. Can be a string, dict, or TableSchema object. Required if + `schema_fn_or_map` is a callable. + + Returns: + A callable with the attached `_union_schema` attribute for Storage Write API. + """ + if isinstance(schema_fn_or_map, dict): + if union_schema is None: + bq_schemas = [ + copy.deepcopy(bigquery_tools.get_bq_tableschema(s)) + for s in schema_fn_or_map.values() + ] + + def _merge_fields(field_a, field_b): + if field_a.type != field_b.type: + raise ValueError( + f"Conflicting types for field '{field_a.name}': " + f"{field_a.type} vs {field_b.type}") + if field_a.type in ('RECORD', 'STRUCT'): + merged_subfields = {} + for f in (field_a.fields or []): + merged_subfields[f.name] = f + for f in (field_b.fields or []): + if f.name in merged_subfields: + merged_subfields[f.name] = _merge_fields( + merged_subfields[f.name], f) + else: + merged_subfields[f.name] = f + field_a.fields = list(merged_subfields.values()) + return field_a + + merged_fields = {} + for schema in bq_schemas: + for field in schema.fields: + name = field.name + if name in merged_fields: + merged_fields[name] = _merge_fields(merged_fields[name], field) + else: + merged_fields[name] = field + union_schema = bigquery.TableSchema(fields=list(merged_fields.values())) + + def lookup_schema(destination, *args): + return schema_fn_or_map[destination] + + schema_callable = lookup_schema + elif callable(schema_fn_or_map): + if union_schema is None: + raise ValueError( + "union_schema must be explicitly provided when schema_fn_or_map " + "is a callable.") + schema_callable = schema_fn_or_map + else: + raise TypeError("schema_fn_or_map must be a callable or a dictionary.") + + schema_callable._union_schema = union_schema + return schema_callable + + class WriteToBigQuery(PTransform): """Write data to BigQuery. @@ -2388,6 +2540,7 @@ def find_in_nested_dict(schema): table=self.table_reference, schema=self.schema, table_side_inputs=self.table_side_inputs, + schema_side_inputs=self.schema_side_inputs, create_disposition=self.create_disposition, write_disposition=self.write_disposition, additional_bq_parameters=self.additional_bq_parameters, @@ -2616,7 +2769,7 @@ def __getitem__(self, key): class StorageWriteToBigQuery(PTransform): """Writes data to BigQuery using Storage API. - Supports dynamic destinations. Dynamic schemas are not supported yet. + Supports dynamic destinations and dynamic schemas. Experimental; no backwards compatibility guarantees. """ @@ -2638,6 +2791,7 @@ def __init__( table, table_side_inputs=None, schema=None, + schema_side_inputs=None, create_disposition=BigQueryDisposition.CREATE_IF_NEEDED, write_disposition=BigQueryDisposition.WRITE_APPEND, additional_bq_parameters=None, @@ -2651,8 +2805,9 @@ def __init__( expansion_service=None, type_overrides=None): self._table = table - self._table_side_inputs = table_side_inputs + self._table_side_inputs = table_side_inputs or () self._schema = schema + self._schema_side_inputs = schema_side_inputs or () self._create_disposition = create_disposition self._write_disposition = write_disposition self.additional_bq_parameters = additional_bq_parameters @@ -2677,9 +2832,8 @@ def expand(self, input): "A schema is required in order to prepare rows " "for writing with STORAGE_WRITE_API.") from exn elif callable(self._schema): - raise NotImplementedError( - "Writing with dynamic schemas is not " - "supported for this write method.") + schema = self._schema + is_rows = False elif isinstance(self._schema, vp.ValueProvider): schema = self._schema.get() is_rows = False @@ -2691,6 +2845,10 @@ def expand(self, input): # if writing to one destination, just convert to Beam rows and send over if not callable(table): + if callable(schema): + raise ValueError( + "Writing with a dynamic schema is only supported when writing to " + "dynamic destinations.") if is_rows: input_beam_rows = input else: @@ -2729,7 +2887,11 @@ def expand(self, input): input_beam_rows = ( input_rows | "Convert dict to Beam Row" >> self.ConvertToBeamRows( - schema, True, self._type_overrides).with_output_types()) + schema, + True, + self._type_overrides, + schema_side_inputs=self._schema_side_inputs).with_output_types( + )) # communicate to Java that this write should use dynamic destinations table = StorageWriteToBigQuery.DYNAMIC_DESTINATIONS @@ -2797,24 +2959,84 @@ def __exit__(self, *args): pass class ConvertToBeamRows(PTransform): - def __init__(self, schema, dynamic_destinations, type_overrides=None): + def __init__( + self, + schema, + dynamic_destinations, + type_overrides=None, + schema_side_inputs=None): self.schema = schema self.dynamic_destinations = dynamic_destinations self.type_overrides = type_overrides + self.schema_side_inputs = schema_side_inputs or () + + def _get_record_type_hint(self): + if callable(self.schema): + schema_hint = ( + getattr(self.schema, '_union_schema', None) or + getattr(self.schema, '_table_schema', None) or + getattr(self.schema, '_beam_schema', None) or + getattr(self.schema, '_schema_hint', None) or + getattr(self.schema, '_output_types', None) or + getattr(self.schema, 'table_schema', None) or + getattr(self.schema, 'schema', None)) + if schema_hint is not None: + if isinstance( + schema_hint, + (bigquery.TableSchema, bigquery.TableFieldSchema, str, dict)): + row_type_hints = bigquery_tools.get_beam_typehints_from_tableschema( + schema_hint, self.type_overrides) + return RowTypeConstraint.from_fields(row_type_hints) + elif isinstance(schema_hint, RowTypeConstraint): + return schema_hint + return RowTypeConstraint.from_fields([]) + else: + row_type_hints = bigquery_tools.get_beam_typehints_from_tableschema( + self.schema, self.type_overrides) + return RowTypeConstraint.from_fields(row_type_hints) def expand(self, input_dicts): if self.dynamic_destinations: - return ( - input_dicts - | "Convert dict to Beam Row" >> beam.Map( - lambda row, schema=DoFn.SetupContextParam( - StorageWriteToBigQuery.ConvertToBeamRowsSetupSchema, args= - [self.schema]): beam.Row( - **{ - StorageWriteToBigQuery.DESTINATION: row[0], - StorageWriteToBigQuery.RECORD: bigquery_tools. - beam_row_from_dict(row[1], schema) - }))) + if callable(self.schema): + record_hint = self._get_record_type_hint() + union_field_names = [ + name for name, _ in getattr(record_hint, '_fields', ()) + ] + + def convert_dynamic_row(row, *schema_side_inputs): + dest, dict_row = row[0], row[1] + record_schema = self.schema(dest, *schema_side_inputs) + record_row = bigquery_tools.beam_row_from_dict( + dict_row, record_schema) + if union_field_names: + record_dict = record_row._asdict() + record_row = beam.Row( + **{ + name: record_dict.get(name, None) + for name in union_field_names + }) + return beam.Row( + **{ + StorageWriteToBigQuery.DESTINATION: dest, + StorageWriteToBigQuery.RECORD: record_row + }) + + return ( + input_dicts + | "Convert dict to Beam Row" >> beam.Map( + convert_dynamic_row, *self.schema_side_inputs)) + else: + return ( + input_dicts + | "Convert dict to Beam Row" >> beam.Map( + lambda row, schema=DoFn.SetupContextParam( + StorageWriteToBigQuery.ConvertToBeamRowsSetupSchema, args= + [self.schema]): beam.Row( + **{ + StorageWriteToBigQuery.DESTINATION: row[0], + StorageWriteToBigQuery.RECORD: bigquery_tools. + beam_row_from_dict(row[1], schema) + }))) else: return ( input_dicts @@ -2825,17 +3047,14 @@ def expand(self, input_dicts): ]): bigquery_tools.beam_row_from_dict(row, schema))) def with_output_types(self): - row_type_hints = bigquery_tools.get_beam_typehints_from_tableschema( - self.schema, self.type_overrides) + record_hint = self._get_record_type_hint() if self.dynamic_destinations: type_hint = RowTypeConstraint.from_fields([ (StorageWriteToBigQuery.DESTINATION, str), - ( - StorageWriteToBigQuery.RECORD, - RowTypeConstraint.from_fields(row_type_hints)) + (StorageWriteToBigQuery.RECORD, record_hint) ]) else: - type_hint = RowTypeConstraint.from_fields(row_type_hints) + type_hint = record_hint return super().with_output_types(type_hint) diff --git a/sdks/python/apache_beam/io/gcp/bigquery_storage_write_test.py b/sdks/python/apache_beam/io/gcp/bigquery_storage_write_test.py new file mode 100644 index 000000000000..dcce144a0663 --- /dev/null +++ b/sdks/python/apache_beam/io/gcp/bigquery_storage_write_test.py @@ -0,0 +1,343 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +"""Unit tests for BigQuery Storage Write API dynamic schemas.""" + +import unittest +from unittest import mock + +import apache_beam as beam +from apache_beam.io.gcp import bigquery +from apache_beam.io.gcp import bigquery_tools +from apache_beam.testing.test_pipeline import TestPipeline +from apache_beam.testing.util import assert_that +from apache_beam.testing.util import equal_to +from apache_beam.typehints.row_type import RowTypeConstraint + +try: + from google.api_core.exceptions import GoogleAPICallError +except ImportError: + GoogleAPICallError = None + + +@unittest.skipIf( + GoogleAPICallError is None, 'GCP dependencies are not installed') +@mock.patch('apache_beam.io.gcp.bigquery.BeamJarExpansionService') +class BigQueryStorageWriteDynamicSchemaTest(unittest.TestCase): + """Test dynamic schema support in BigQuery Storage Write API.""" + def test_storage_write_init_with_schema_side_inputs( + self, mock_expansion_service): + """Test that StorageWriteToBigQuery accepts schema_side_inputs.""" + transform = bigquery.StorageWriteToBigQuery( + table='test-project:test_dataset.test_table', + schema=lambda dest: None, + schema_side_inputs=('side_input_1', )) + self.assertEqual(transform._schema_side_inputs, ('side_input_1', )) + self.assertEqual(transform._table_side_inputs, ()) + + def test_convert_to_beam_rows_dynamic_destinations_dynamic_schema( + self, mock_expansion_service): + """Test ConvertToBeamRows with dynamic destinations and dynamic schema.""" + schema1 = { + 'fields': [ + { + 'name': 'id', 'type': 'INTEGER' + }, + { + 'name': 'name', 'type': 'STRING' + }, + ] + } + schema2 = { + 'fields': [ + { + 'name': 'id', 'type': 'INTEGER' + }, + { + 'name': 'score', 'type': 'FLOAT' + }, + ] + } + schema_map = {'table1': schema1, 'table2': schema2} + + converter = bigquery.StorageWriteToBigQuery.ConvertToBeamRows( + schema=lambda dest: schema_map[dest], dynamic_destinations=True) + + with TestPipeline() as p: + input_data = [ + ('table1', { + 'id': 1, 'name': 'foo' + }), + ('table2', { + 'id': 2, 'score': 3.14 + }), + ] + res = p | "CreateInput" >> beam.Create(input_data) | converter + + expected_rows = [ + beam.Row(destination='table1', record=beam.Row(id=1, name='foo')), + beam.Row(destination='table2', record=beam.Row(id=2, score=3.14)), + ] + assert_that(res, equal_to(expected_rows)) + + def test_convert_to_beam_rows_dynamic_destinations_with_side_inputs( + self, mock_expansion_service): + """Test ConvertToBeamRows with dynamic schema and side inputs.""" + schema1 = { + 'fields': [ + { + 'name': 'id', 'type': 'INTEGER' + }, + { + 'name': 'name', 'type': 'STRING' + }, + ] + } + schema2 = { + 'fields': [ + { + 'name': 'id', 'type': 'INTEGER' + }, + { + 'name': 'score', 'type': 'FLOAT' + }, + ] + } + + with TestPipeline() as p: + side_pcoll = ( + p + | "CreateSide" >> beam.Create([{ + 'table1': schema1, 'table2': schema2 + }])) + + converter = bigquery.StorageWriteToBigQuery.ConvertToBeamRows( + schema=lambda dest, side_map: side_map[dest], + dynamic_destinations=True, + schema_side_inputs=(beam.pvalue.AsSingleton(side_pcoll), )) + + input_data = [ + ('table1', { + 'id': 1, 'name': 'foo' + }), + ('table2', { + 'id': 2, 'score': 3.14 + }), + ] + res = p | "CreateInput" >> beam.Create(input_data) | converter + + expected_rows = [ + beam.Row(destination='table1', record=beam.Row(id=1, name='foo')), + beam.Row(destination='table2', record=beam.Row(id=2, score=3.14)), + ] + assert_that(res, equal_to(expected_rows)) + + def test_storage_write_static_destination_dynamic_schema_raises_error( + self, mock_expansion_service): + """Test that static destination with dynamic schema raises ValueError.""" + transform = bigquery.StorageWriteToBigQuery( + table='test-project:test_dataset.test_table', schema=lambda dest: None) + with self.assertRaisesRegex( + ValueError, + "Writing with a dynamic schema is only supported when writing to " + "dynamic destinations."): + with TestPipeline() as p: + _ = p | "CreateInput" >> beam.Create([{'id': 1}]) | transform + + def test_convert_to_beam_rows_with_output_types_dynamic_schema( + self, mock_expansion_service): + """Test with_output_types when schema is callable.""" + converter_dyn = bigquery.StorageWriteToBigQuery.ConvertToBeamRows( + schema=lambda dest: None, dynamic_destinations=True) + type_hint_dyn = converter_dyn.with_output_types().get_type_hints( + ).simple_output_type('') + self.assertIsInstance(type_hint_dyn, RowTypeConstraint) + self.assertEqual( + type_hint_dyn._fields, + ( + (bigquery.StorageWriteToBigQuery.DESTINATION, str), + ( + bigquery.StorageWriteToBigQuery.RECORD, + RowTypeConstraint.from_fields([])), + )) + + def test_convert_to_beam_rows_with_output_types_dynamic_schema_hint( + self, mock_expansion_service): + """Test with_output_types when schema is callable with _union_schema.""" + def dyn_schema(dest): + return None + + dyn_schema._union_schema = 'id:INTEGER,name:STRING' + converter_dyn = bigquery.StorageWriteToBigQuery.ConvertToBeamRows( + schema=dyn_schema, dynamic_destinations=True) + type_hint_dyn = converter_dyn.with_output_types().get_type_hints( + ).simple_output_type('') + self.assertIsInstance(type_hint_dyn, RowTypeConstraint) + self.assertEqual( + type_hint_dyn._fields[0], + (bigquery.StorageWriteToBigQuery.DESTINATION, str)) + self.assertEqual( + type_hint_dyn._fields[1][0], bigquery.StorageWriteToBigQuery.RECORD) + expected_record_hint = RowTypeConstraint.from_fields( + bigquery_tools.get_beam_typehints_from_tableschema( + 'id:INTEGER,name:STRING')) + self.assertEqual( + type_hint_dyn._fields[1][1]._fields, expected_record_hint._fields) + + def test_convert_to_beam_rows_union_schema_fills_missing_attributes( + self, mock_expansion_service): + """Test ConvertToBeamRows fills None for fields in union schema not in row.""" + def dyn_schema(dest): + if 'users' in dest: + return 'id:INTEGER,name:STRING' + return 'id:INTEGER,score:INTEGER' + + dyn_schema._union_schema = 'id:INTEGER,name:STRING,score:INTEGER' + converter = bigquery.StorageWriteToBigQuery.ConvertToBeamRows( + schema=dyn_schema, dynamic_destinations=True) + + with TestPipeline() as p: + rows = ( + p + | beam.Create([ + ('dest_users', { + 'id': 1, 'name': 'alice' + }), + ('dest_scores', { + 'id': 2, 'score': 95 + }), + ]) + | converter) + + def check_rows(actual): + actual_list = list(actual) + assert len(actual_list) == 2 + r1, r2 = actual_list[0], actual_list[1] + if r1.destination == 'dest_scores': + r1, r2 = r2, r1 + assert r1.destination == 'dest_users' + assert r1.record.id == 1 + assert r1.record.name == 'alice' + assert r1.record.score is None + assert r2.destination == 'dest_scores' + assert r2.record.id == 2 + assert r2.record.name is None + assert r2.record.score == 95 + + assert_that(rows, check_rows) + + def test_storage_write_to_bigquery_expand_dynamic_schema( + self, mock_expansion_service): + """Test StorageWriteToBigQuery expand does not fail for callable schema.""" + schema1 = { + 'fields': [ + { + 'name': 'id', 'type': 'INTEGER' + }, + ] + } + + class _DummyExternalTransform(beam.PTransform): + def expand(self, pcoll): + return { + bigquery.StorageWriteToBigQuery.FAILED_ROWS_WITH_ERRORS: ( + pcoll.pipeline | "CreateErrors" >> beam.Create([])) + } + + with mock.patch.object(bigquery, + 'SchemaAwareExternalTransform', + autospec=True) as mock_ext: + mock_ext.return_value = _DummyExternalTransform() + transform = bigquery.StorageWriteToBigQuery( + table=lambda record: 'table1', schema=lambda dest: schema1) + + with TestPipeline() as p: + _ = p | "CreateInput" >> beam.Create([{'id': 1}]) | transform + + mock_ext.assert_called_once() + _, kwargs = mock_ext.call_args + self.assertEqual( + kwargs['table'], bigquery.StorageWriteToBigQuery.DYNAMIC_DESTINATIONS) + + def test_write_to_bigquery_storage_api_passes_schema_side_inputs( + self, mock_expansion_service): + """Test WriteToBigQuery passes schema_side_inputs to StorageWriteToBigQuery.""" + with mock.patch.object(bigquery, 'StorageWriteToBigQuery', + autospec=True) as mock_storage_write: + mock_storage_write.return_value = beam.Map(lambda x: x) + with TestPipeline() as p: + side_pc = p | "CreateSide" >> beam.Create([1]) + write_transform = bigquery.WriteToBigQuery( + table='proj:ds.table', + method=bigquery.WriteToBigQuery.Method.STORAGE_WRITE_API, + schema=lambda dest: None, + schema_side_inputs=(beam.pvalue.AsSingleton(side_pc), )) + _ = p | "CreateInput" >> beam.Create([{'id': 1}]) | write_transform + + mock_storage_write.assert_called_once() + _, kwargs = mock_storage_write.call_args + self.assertEqual(len(kwargs['schema_side_inputs']), 1) + + def test_dynamic_schema_helper_with_dictionary(self, mock_expansion_service): + """Test dynamic_schema auto-merges fields from a dictionary map.""" + schema_map = { + 'table_a': 'id:INTEGER,name:STRING', + 'table_b': 'id:INTEGER,score:INTEGER,active:BOOLEAN' + } + schema_callable = bigquery.dynamic_schema(schema_map) + + # 1. Verify it returns the correct schema per destination + self.assertEqual(schema_callable('table_a'), 'id:INTEGER,name:STRING') + self.assertEqual( + schema_callable('table_b'), 'id:INTEGER,score:INTEGER,active:BOOLEAN') + + # 2. Verify it auto-merged all unique fields into _union_schema + expected_union = bigquery_tools.get_bq_tableschema( + 'id:INTEGER,name:STRING,score:INTEGER,active:BOOLEAN') + self.assertEqual(schema_callable._union_schema, expected_union) + + def test_dynamic_schema_helper_with_callable_and_explicit_union( + self, mock_expansion_service): + """Test dynamic_schema attaches union_schema explicitly to a callable.""" + def get_schema(dest): + return 'id:INTEGER,name:STRING' + + union_schema = 'id:INTEGER,name:STRING,score:INTEGER' + schema_callable = bigquery.dynamic_schema( + get_schema, union_schema=union_schema) + + self.assertEqual(schema_callable('table_a'), 'id:INTEGER,name:STRING') + self.assertEqual(schema_callable._union_schema, union_schema) + + def test_dynamic_schema_helper_missing_union_on_callable_raises_error( + self, mock_expansion_service): + """Test dynamic_schema raises ValueError if union_schema is missing on callable.""" + def get_schema(dest): + return 'id:INTEGER' + + with self.assertRaises(ValueError): + bigquery.dynamic_schema(get_schema) + + def test_dynamic_schema_helper_with_invalid_type_raises_error( + self, mock_expansion_service): + """Test dynamic_schema raises TypeError if passed an invalid type.""" + with self.assertRaises(TypeError): + bigquery.dynamic_schema('id:INTEGER') + + +if __name__ == '__main__': + unittest.main()