From cfa21a7c32757622a53155491c635dc49fa18488 Mon Sep 17 00:00:00 2001 From: liferoad Date: Sat, 20 Sep 2025 20:16:12 -0400 Subject: [PATCH 01/14] feat(bigquery): add BigLake configuration support for Storage Write API add BigLake configuration parameter to StorageWriteToBigQuery transform in both Java and Python SDKs include validation tests for BigLake configuration parameters --- ...torageWriteApiSchemaTransformProvider.java | 3 + .../providers/BigQueryWriteConfiguration.java | 10 ++ sdks/python/apache_beam/io/gcp/bigquery.py | 3 + .../io/gcp/bigquery_biglake_test.py | 117 ++++++++++++++++++ 4 files changed, 133 insertions(+) create mode 100644 sdks/python/apache_beam/io/gcp/bigquery_biglake_test.py diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/providers/BigQueryStorageWriteApiSchemaTransformProvider.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/providers/BigQueryStorageWriteApiSchemaTransformProvider.java index bb8f72003429..f9f86cc80186 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/providers/BigQueryStorageWriteApiSchemaTransformProvider.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/providers/BigQueryStorageWriteApiSchemaTransformProvider.java @@ -315,6 +315,9 @@ BigQueryIO.Write createStorageWriteApiTransform(Schema schema) { if (!Strings.isNullOrEmpty(configuration.getKmsKey())) { write = write.withKmsKey(configuration.getKmsKey()); } + if (configuration.getBigLakeConfiguration() != null) { + write = write.withBigLakeConfiguration(configuration.getBigLakeConfiguration()); + } if (this.testBigQueryServices != null) { write = write.withTestServices(testBigQueryServices); } diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/providers/BigQueryWriteConfiguration.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/providers/BigQueryWriteConfiguration.java index 5df6e1f6afcd..ec3fb7663546 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/providers/BigQueryWriteConfiguration.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/providers/BigQueryWriteConfiguration.java @@ -197,6 +197,14 @@ public static Builder builder() { @SchemaFieldDescription("A list of columns to cluster the BigQuery table by.") public abstract @Nullable List getClusteringFields(); + @SchemaFieldDescription( + "Configuration for creating BigLake tables. The following options are available: " + + "connectionId (REQUIRED): the name of your cloud resource connection, " + + "storageUri (REQUIRED): the path to your GCS folder where data will be written to, " + + "fileFormat (OPTIONAL): defaults to 'parquet', " + + "tableFormat (OPTIONAL): defaults to 'iceberg'.") + public abstract @Nullable java.util.Map getBigLakeConfiguration(); + /** Builder for {@link BigQueryWriteConfiguration}. */ @AutoValue.Builder public abstract static class Builder { @@ -231,6 +239,8 @@ public abstract static class Builder { public abstract Builder setClusteringFields(List clusteringFields); + public abstract Builder setBigLakeConfiguration(java.util.Map bigLakeConfiguration); + /** Builds a {@link BigQueryWriteConfiguration} instance. */ public abstract BigQueryWriteConfiguration build(); } diff --git a/sdks/python/apache_beam/io/gcp/bigquery.py b/sdks/python/apache_beam/io/gcp/bigquery.py index aa0ebc12ef18..828e707fc5cd 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery.py +++ b/sdks/python/apache_beam/io/gcp/bigquery.py @@ -2626,6 +2626,7 @@ def __init__( num_storage_api_streams=0, use_cdc_writes: bool = False, primary_key: List[str] = None, + big_lake_configuration=None, expansion_service=None): self._table = table self._table_side_inputs = table_side_inputs @@ -2639,6 +2640,7 @@ def __init__( self._num_storage_api_streams = num_storage_api_streams self._use_cdc_writes = use_cdc_writes self._primary_key = primary_key + self._big_lake_configuration = big_lake_configuration self._expansion_service = expansion_service or BeamJarExpansionService( 'sdks:java:io:google-cloud-platform:expansion-service:build') @@ -2733,6 +2735,7 @@ def expand(self, input): use_cdc_writes=self._use_cdc_writes, primary_key=self._primary_key, clustering_fields=clustering_fields, + big_lake_configuration=self._big_lake_configuration, error_handling={ 'output': StorageWriteToBigQuery.FAILED_ROWS_WITH_ERRORS })) diff --git a/sdks/python/apache_beam/io/gcp/bigquery_biglake_test.py b/sdks/python/apache_beam/io/gcp/bigquery_biglake_test.py new file mode 100644 index 000000000000..80792fdbeabb --- /dev/null +++ b/sdks/python/apache_beam/io/gcp/bigquery_biglake_test.py @@ -0,0 +1,117 @@ +# +# 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 BigLake configuration.""" + +import unittest +import mock + +import apache_beam as beam +from apache_beam.io.gcp import bigquery +from apache_beam.testing.test_pipeline import TestPipeline + + +class BigQueryBigLakeTest(unittest.TestCase): + """Test BigLake configuration support in BigQuery Storage Write API.""" + + def test_storage_write_to_bigquery_with_biglake_config(self): + """Test that StorageWriteToBigQuery accepts bigLakeConfiguration parameter.""" + big_lake_config = { + 'connectionId': 'projects/test-project/locations/us/connections/test-connection', + 'storageUri': 'gs://test-bucket/test-path', + 'fileFormat': 'parquet', + 'tableFormat': 'iceberg' + } + + # Test that the constructor accepts the bigLakeConfiguration parameter + transform = bigquery.StorageWriteToBigQuery( + table='test-project:test_dataset.test_table', + big_lake_configuration=big_lake_config + ) + + # Verify the configuration is stored + self.assertEqual(transform._big_lake_configuration, big_lake_config) + + def test_storage_write_to_bigquery_without_biglake_config(self): + """Test that StorageWriteToBigQuery works without bigLakeConfiguration parameter.""" + transform = bigquery.StorageWriteToBigQuery( + table='test-project:test_dataset.test_table' + ) + + # Verify the configuration is None by default + self.assertIsNone(transform._big_lake_configuration) + + @mock.patch('apache_beam.io.gcp.bigquery.SchemaAwareExternalTransform') + def test_biglake_config_passed_to_external_transform(self, mock_external_transform): + """Test that bigLakeConfiguration is passed to the external transform.""" + big_lake_config = { + 'connectionId': 'projects/test-project/locations/us/connections/test-connection', + 'storageUri': 'gs://test-bucket/test-path' + } + + with TestPipeline() as p: + input_data = p | 'Create' >> beam.Create([ + beam.Row(name='Alice', age=30), + beam.Row(name='Bob', age=25) + ]) + + transform = bigquery.StorageWriteToBigQuery( + table='test-project:test_dataset.test_table', + big_lake_configuration=big_lake_config + ) + + # Apply the transform + _ = input_data | transform + + # Verify that SchemaAwareExternalTransform was called with bigLakeConfiguration + mock_external_transform.assert_called_once() + call_kwargs = mock_external_transform.call_args[1] + self.assertEqual(call_kwargs['big_lake_configuration'], big_lake_config) + + def test_biglake_config_validation(self): + """Test validation of bigLakeConfiguration parameters.""" + # Test with minimal required configuration + minimal_config = { + 'connectionId': 'projects/test-project/locations/us/connections/test-connection', + 'storageUri': 'gs://test-bucket/test-path' + } + + transform = bigquery.StorageWriteToBigQuery( + table='test-project:test_dataset.test_table', + big_lake_configuration=minimal_config + ) + + self.assertEqual(transform._big_lake_configuration, minimal_config) + + # Test with full configuration + full_config = { + 'connectionId': 'projects/test-project/locations/us/connections/test-connection', + 'storageUri': 'gs://test-bucket/test-path', + 'fileFormat': 'parquet', + 'tableFormat': 'iceberg' + } + + transform = bigquery.StorageWriteToBigQuery( + table='test-project:test_dataset.test_table', + big_lake_configuration=full_config + ) + + self.assertEqual(transform._big_lake_configuration, full_config) + + +if __name__ == '__main__': + unittest.main() \ No newline at end of file From 9bc296a826ed9fb70d9be33f76074b02d224960d Mon Sep 17 00:00:00 2001 From: liferoad Date: Sun, 21 Sep 2025 08:29:17 -0400 Subject: [PATCH 02/14] fixed errors --- .../io/gcp/bigquery/providers/BigQueryWriteConfiguration.java | 2 +- sdks/python/apache_beam/io/gcp/bigquery_biglake_test.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/providers/BigQueryWriteConfiguration.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/providers/BigQueryWriteConfiguration.java index ec3fb7663546..3c803f3dd42e 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/providers/BigQueryWriteConfiguration.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/providers/BigQueryWriteConfiguration.java @@ -203,7 +203,7 @@ public static Builder builder() { + "storageUri (REQUIRED): the path to your GCS folder where data will be written to, " + "fileFormat (OPTIONAL): defaults to 'parquet', " + "tableFormat (OPTIONAL): defaults to 'iceberg'.") - public abstract @Nullable java.util.Map getBigLakeConfiguration(); + public abstract java.util.@Nullable Map getBigLakeConfiguration(); /** Builder for {@link BigQueryWriteConfiguration}. */ @AutoValue.Builder diff --git a/sdks/python/apache_beam/io/gcp/bigquery_biglake_test.py b/sdks/python/apache_beam/io/gcp/bigquery_biglake_test.py index 80792fdbeabb..c1255be3fad9 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_biglake_test.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_biglake_test.py @@ -36,7 +36,7 @@ def test_storage_write_to_bigquery_with_biglake_config(self): 'fileFormat': 'parquet', 'tableFormat': 'iceberg' } - + # Test that the constructor accepts the bigLakeConfiguration parameter transform = bigquery.StorageWriteToBigQuery( table='test-project:test_dataset.test_table', From 1a7e5fa9535bb62ccb04fccb6cda9d14d689b1e2 Mon Sep 17 00:00:00 2001 From: liferoad Date: Sun, 21 Sep 2025 11:16:39 -0400 Subject: [PATCH 03/14] mock --- .../io/gcp/bigquery_biglake_test.py | 64 +++++++++++-------- 1 file changed, 36 insertions(+), 28 deletions(-) diff --git a/sdks/python/apache_beam/io/gcp/bigquery_biglake_test.py b/sdks/python/apache_beam/io/gcp/bigquery_biglake_test.py index c1255be3fad9..4d882ac9410e 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_biglake_test.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_biglake_test.py @@ -31,7 +31,8 @@ class BigQueryBigLakeTest(unittest.TestCase): def test_storage_write_to_bigquery_with_biglake_config(self): """Test that StorageWriteToBigQuery accepts bigLakeConfiguration parameter.""" big_lake_config = { - 'connectionId': 'projects/test-project/locations/us/connections/test-connection', + 'connectionId': ( + 'projects/test-project/locations/us/connections/test-connection'), 'storageUri': 'gs://test-bucket/test-path', 'fileFormat': 'parquet', 'tableFormat': 'iceberg' @@ -40,44 +41,51 @@ def test_storage_write_to_bigquery_with_biglake_config(self): # Test that the constructor accepts the bigLakeConfiguration parameter transform = bigquery.StorageWriteToBigQuery( table='test-project:test_dataset.test_table', - big_lake_configuration=big_lake_config - ) - + big_lake_configuration=big_lake_config) + # Verify the configuration is stored self.assertEqual(transform._big_lake_configuration, big_lake_config) def test_storage_write_to_bigquery_without_biglake_config(self): - """Test that StorageWriteToBigQuery works without bigLakeConfiguration parameter.""" + """Test that StorageWriteToBigQuery works without bigLakeConfiguration.""" transform = bigquery.StorageWriteToBigQuery( - table='test-project:test_dataset.test_table' - ) - + table='test-project:test_dataset.test_table') + # Verify the configuration is None by default self.assertIsNone(transform._big_lake_configuration) @mock.patch('apache_beam.io.gcp.bigquery.SchemaAwareExternalTransform') - def test_biglake_config_passed_to_external_transform(self, mock_external_transform): + def test_biglake_config_passed_to_external_transform(self, + mock_external_transform): """Test that bigLakeConfiguration is passed to the external transform.""" big_lake_config = { - 'connectionId': 'projects/test-project/locations/us/connections/test-connection', + 'connectionId': ( + 'projects/test-project/locations/us/connections/test-connection'), 'storageUri': 'gs://test-bucket/test-path' } - + + # Mock the external transform to return a dummy PCollection + mock_transform_instance = mock.MagicMock() + mock_external_transform.return_value = mock_transform_instance + mock_transform_instance.__ror__ = mock.MagicMock() + mock_transform_instance.__ror__.return_value = { + bigquery.StorageWriteToBigQuery.FAILED_ROWS_WITH_ERRORS: [] + } + with TestPipeline() as p: input_data = p | 'Create' >> beam.Create([ beam.Row(name='Alice', age=30), beam.Row(name='Bob', age=25) ]) - + transform = bigquery.StorageWriteToBigQuery( table='test-project:test_dataset.test_table', - big_lake_configuration=big_lake_config - ) - + big_lake_configuration=big_lake_config) + # Apply the transform _ = input_data | transform - - # Verify that SchemaAwareExternalTransform was called with bigLakeConfiguration + + # Verify that SchemaAwareExternalTransform was called with config mock_external_transform.assert_called_once() call_kwargs = mock_external_transform.call_args[1] self.assertEqual(call_kwargs['big_lake_configuration'], big_lake_config) @@ -86,30 +94,30 @@ def test_biglake_config_validation(self): """Test validation of bigLakeConfiguration parameters.""" # Test with minimal required configuration minimal_config = { - 'connectionId': 'projects/test-project/locations/us/connections/test-connection', + 'connectionId': ( + 'projects/test-project/locations/us/connections/test-connection'), 'storageUri': 'gs://test-bucket/test-path' } - + transform = bigquery.StorageWriteToBigQuery( table='test-project:test_dataset.test_table', - big_lake_configuration=minimal_config - ) - + big_lake_configuration=minimal_config) + self.assertEqual(transform._big_lake_configuration, minimal_config) - + # Test with full configuration full_config = { - 'connectionId': 'projects/test-project/locations/us/connections/test-connection', + 'connectionId': ( + 'projects/test-project/locations/us/connections/test-connection'), 'storageUri': 'gs://test-bucket/test-path', 'fileFormat': 'parquet', 'tableFormat': 'iceberg' } - + transform = bigquery.StorageWriteToBigQuery( table='test-project:test_dataset.test_table', - big_lake_configuration=full_config - ) - + big_lake_configuration=full_config) + self.assertEqual(transform._big_lake_configuration, full_config) From 81ea94823c9dee597c8a77a681e39b7e487fafad Mon Sep 17 00:00:00 2001 From: liferoad Date: Mon, 22 Sep 2025 10:17:07 -0400 Subject: [PATCH 04/14] fixed tests --- ..._PostCommit_Python_Xlang_Gcp_Dataflow.json | 2 +- ...m_PostCommit_Python_Xlang_IO_Dataflow.json | 2 +- .../io/gcp/bigquery_biglake_test.py | 53 +++++++------------ 3 files changed, 20 insertions(+), 37 deletions(-) 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 95fef3e26ca2..99a8fc8ff6d5 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": 13 + "modification": 14 } diff --git a/.github/trigger_files/beam_PostCommit_Python_Xlang_IO_Dataflow.json b/.github/trigger_files/beam_PostCommit_Python_Xlang_IO_Dataflow.json index e0266d62f2e0..f1ba03a243ee 100644 --- a/.github/trigger_files/beam_PostCommit_Python_Xlang_IO_Dataflow.json +++ b/.github/trigger_files/beam_PostCommit_Python_Xlang_IO_Dataflow.json @@ -1,4 +1,4 @@ { "comment": "Modify this file in a trivial way to cause this test suite to run", - "modification": 4 + "modification": 5 } diff --git a/sdks/python/apache_beam/io/gcp/bigquery_biglake_test.py b/sdks/python/apache_beam/io/gcp/bigquery_biglake_test.py index 4d882ac9410e..58e0f650272d 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_biglake_test.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_biglake_test.py @@ -18,18 +18,14 @@ """Unit tests for BigQuery BigLake configuration.""" import unittest -import mock -import apache_beam as beam from apache_beam.io.gcp import bigquery -from apache_beam.testing.test_pipeline import TestPipeline class BigQueryBigLakeTest(unittest.TestCase): """Test BigLake configuration support in BigQuery Storage Write API.""" - def test_storage_write_to_bigquery_with_biglake_config(self): - """Test that StorageWriteToBigQuery accepts bigLakeConfiguration parameter.""" + """Test that StorageWriteToBigQuery accepts bigLakeConfiguration.""" big_lake_config = { 'connectionId': ( 'projects/test-project/locations/us/connections/test-connection'), @@ -54,41 +50,28 @@ def test_storage_write_to_bigquery_without_biglake_config(self): # Verify the configuration is None by default self.assertIsNone(transform._big_lake_configuration) - @mock.patch('apache_beam.io.gcp.bigquery.SchemaAwareExternalTransform') - def test_biglake_config_passed_to_external_transform(self, - mock_external_transform): - """Test that bigLakeConfiguration is passed to the external transform.""" + def test_biglake_config_passed_to_external_transform(self): + """Test that StorageWriteToBigQuery accepts bigLakeConfiguration.""" big_lake_config = { - 'connectionId': ( - 'projects/test-project/locations/us/connections/test-connection'), - 'storageUri': 'gs://test-bucket/test-path' + 'connection_id': 'projects/my-project/locations/us/connections/my-conn', + 'table_format': 'ICEBERG' } - # Mock the external transform to return a dummy PCollection - mock_transform_instance = mock.MagicMock() - mock_external_transform.return_value = mock_transform_instance - mock_transform_instance.__ror__ = mock.MagicMock() - mock_transform_instance.__ror__.return_value = { - bigquery.StorageWriteToBigQuery.FAILED_ROWS_WITH_ERRORS: [] - } - - with TestPipeline() as p: - input_data = p | 'Create' >> beam.Create([ - beam.Row(name='Alice', age=30), - beam.Row(name='Bob', age=25) - ]) + # Create the transform + transform = bigquery.StorageWriteToBigQuery( + table='my-project:my_dataset.my_table', + big_lake_configuration=big_lake_config) - transform = bigquery.StorageWriteToBigQuery( - table='test-project:test_dataset.test_table', - big_lake_configuration=big_lake_config) + # Verify the big_lake_configuration is stored correctly + self.assertEqual(transform._big_lake_configuration, big_lake_config) - # Apply the transform - _ = input_data | transform + # Verify that the transform has the expected identifier + self.assertEqual( + transform.IDENTIFIER, + "beam:schematransform:org.apache.beam:bigquery_storage_write:v2") - # Verify that SchemaAwareExternalTransform was called with config - mock_external_transform.assert_called_once() - call_kwargs = mock_external_transform.call_args[1] - self.assertEqual(call_kwargs['big_lake_configuration'], big_lake_config) + # Verify that the expansion service is set up correctly + self.assertIsNotNone(transform._expansion_service) def test_biglake_config_validation(self): """Test validation of bigLakeConfiguration parameters.""" @@ -122,4 +105,4 @@ def test_biglake_config_validation(self): if __name__ == '__main__': - unittest.main() \ No newline at end of file + unittest.main() From 804b9d3f443af17140e46b876a659ae95352880e Mon Sep 17 00:00:00 2001 From: liferoad Date: Mon, 22 Sep 2025 14:37:06 -0400 Subject: [PATCH 05/14] fixed tests --- .../apache/beam/gradle/BeamModulePlugin.groovy | 16 +++++++++++++--- .../apache_beam/io/gcp/bigquery_biglake_test.py | 13 ++++++++++--- 2 files changed, 23 insertions(+), 6 deletions(-) diff --git a/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy b/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy index 103405a57931..06cf43e36791 100644 --- a/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy +++ b/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy @@ -3199,6 +3199,16 @@ class BeamModulePlugin implements Plugin { testJavaHome = project.findProperty("java${testJavaVersion}Home") } + // Detect macOS and append '-macos' to tox environment to avoid pip check issues + def actualToxEnv = tox_env + def osName = System.getProperty("os.name").toLowerCase() + if (osName.contains("mac")) { + // Only append -macos for standard python environments (py39, py310, etc.) + if (tox_env.matches("py\\d+")) { + actualToxEnv = "${tox_env}-macos" + } + } + if (project.hasProperty('useWheelDistribution')) { def pythonVersionNumber = project.ext.pythonVersion.replace('.', '') dependsOn ":sdks:python:bdistPy${pythonVersionNumber}linux" @@ -3214,7 +3224,7 @@ class BeamModulePlugin implements Plugin { environment "JAVA_HOME", testJavaHome } executable 'sh' - args '-c', ". ${project.ext.envdir}/bin/activate && cd ${copiedPyRoot} && scripts/run_tox.sh $tox_env ${packageFilename} '$posargs' " + args '-c', ". ${project.ext.envdir}/bin/activate && cd ${copiedPyRoot} && scripts/run_tox.sh $actualToxEnv ${packageFilename} '$posargs' " } } } else { @@ -3227,12 +3237,12 @@ class BeamModulePlugin implements Plugin { environment "JAVA_HOME", testJavaHome } executable 'sh' - args '-c', ". ${project.ext.envdir}/bin/activate && cd ${copiedPyRoot} && scripts/run_tox.sh $tox_env '$posargs'" + args '-c', ". ${project.ext.envdir}/bin/activate && cd ${copiedPyRoot} && scripts/run_tox.sh $actualToxEnv '$posargs'" } } } inputs.files project.pythonSdkDeps - outputs.files project.fileTree(dir: "${pythonRootDir}/target/.tox/${tox_env}/log/") + outputs.files project.fileTree(dir: "${pythonRootDir}/target/.tox/${actualToxEnv}/log/") } } // Run single or a set of integration tests with provided test options and pipeline options. diff --git a/sdks/python/apache_beam/io/gcp/bigquery_biglake_test.py b/sdks/python/apache_beam/io/gcp/bigquery_biglake_test.py index 58e0f650272d..0f42591105cd 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_biglake_test.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_biglake_test.py @@ -18,6 +18,7 @@ """Unit tests for BigQuery BigLake configuration.""" import unittest +from unittest import mock from apache_beam.io.gcp import bigquery @@ -50,13 +51,18 @@ def test_storage_write_to_bigquery_without_biglake_config(self): # Verify the configuration is None by default self.assertIsNone(transform._big_lake_configuration) - def test_biglake_config_passed_to_external_transform(self): + @mock.patch('apache_beam.io.gcp.bigquery.BeamJarExpansionService') + def test_biglake_config_passed_to_external_transform( + self, mock_expansion_service): """Test that StorageWriteToBigQuery accepts bigLakeConfiguration.""" big_lake_config = { 'connection_id': 'projects/my-project/locations/us/connections/my-conn', 'table_format': 'ICEBERG' } + # Mock the expansion service to avoid JAR dependency + mock_expansion_service.return_value = mock.MagicMock() + # Create the transform transform = bigquery.StorageWriteToBigQuery( table='my-project:my_dataset.my_table', @@ -70,8 +76,9 @@ def test_biglake_config_passed_to_external_transform(self): transform.IDENTIFIER, "beam:schematransform:org.apache.beam:bigquery_storage_write:v2") - # Verify that the expansion service is set up correctly - self.assertIsNotNone(transform._expansion_service) + # Verify that the expansion service was created (mocked) + mock_expansion_service.assert_called_once_with( + 'sdks:java:io:google-cloud-platform:expansion-service:build') def test_biglake_config_validation(self): """Test validation of bigLakeConfiguration parameters.""" From 96ff0c0863b42d4520707be101237632e054ddf3 Mon Sep 17 00:00:00 2001 From: liferoad Date: Mon, 22 Sep 2025 15:33:11 -0400 Subject: [PATCH 06/14] mock tests --- .../bigquery/providers/BigQueryWriteConfiguration.java | 3 ++- .../python/apache_beam/io/gcp/bigquery_biglake_test.py | 10 ++++++---- 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/providers/BigQueryWriteConfiguration.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/providers/BigQueryWriteConfiguration.java index 3c803f3dd42e..9e6a26399fa9 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/providers/BigQueryWriteConfiguration.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/providers/BigQueryWriteConfiguration.java @@ -239,7 +239,8 @@ public abstract static class Builder { public abstract Builder setClusteringFields(List clusteringFields); - public abstract Builder setBigLakeConfiguration(java.util.Map bigLakeConfiguration); + public abstract Builder setBigLakeConfiguration( + java.util.Map bigLakeConfiguration); /** Builds a {@link BigQueryWriteConfiguration} instance. */ public abstract BigQueryWriteConfiguration build(); diff --git a/sdks/python/apache_beam/io/gcp/bigquery_biglake_test.py b/sdks/python/apache_beam/io/gcp/bigquery_biglake_test.py index 0f42591105cd..773523fcedd9 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_biglake_test.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_biglake_test.py @@ -23,9 +23,11 @@ from apache_beam.io.gcp import bigquery +@mock.patch('apache_beam.io.gcp.bigquery.BeamJarExpansionService') class BigQueryBigLakeTest(unittest.TestCase): """Test BigLake configuration support in BigQuery Storage Write API.""" - def test_storage_write_to_bigquery_with_biglake_config(self): + def test_storage_write_to_bigquery_with_biglake_config( + self, mock_expansion_service): """Test that StorageWriteToBigQuery accepts bigLakeConfiguration.""" big_lake_config = { 'connectionId': ( @@ -43,7 +45,8 @@ def test_storage_write_to_bigquery_with_biglake_config(self): # Verify the configuration is stored self.assertEqual(transform._big_lake_configuration, big_lake_config) - def test_storage_write_to_bigquery_without_biglake_config(self): + def test_storage_write_to_bigquery_without_biglake_config( + self, mock_expansion_service): """Test that StorageWriteToBigQuery works without bigLakeConfiguration.""" transform = bigquery.StorageWriteToBigQuery( table='test-project:test_dataset.test_table') @@ -51,7 +54,6 @@ def test_storage_write_to_bigquery_without_biglake_config(self): # Verify the configuration is None by default self.assertIsNone(transform._big_lake_configuration) - @mock.patch('apache_beam.io.gcp.bigquery.BeamJarExpansionService') def test_biglake_config_passed_to_external_transform( self, mock_expansion_service): """Test that StorageWriteToBigQuery accepts bigLakeConfiguration.""" @@ -80,7 +82,7 @@ def test_biglake_config_passed_to_external_transform( mock_expansion_service.assert_called_once_with( 'sdks:java:io:google-cloud-platform:expansion-service:build') - def test_biglake_config_validation(self): + def test_biglake_config_validation(self, mock_expansion_service): """Test validation of bigLakeConfiguration parameters.""" # Test with minimal required configuration minimal_config = { From 1596361b44c7d474cd3291edadb23115ad34943e Mon Sep 17 00:00:00 2001 From: liferoad Date: Tue, 23 Sep 2025 22:04:19 -0400 Subject: [PATCH 07/14] added one IT --- .../providers/BigQueryWriteConfiguration.java | 10 +++--- .../io/external/xlang_bigqueryio_it_test.py | 33 +++++++++++++++++++ 2 files changed, 38 insertions(+), 5 deletions(-) diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/providers/BigQueryWriteConfiguration.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/providers/BigQueryWriteConfiguration.java index 9e6a26399fa9..55d7f7c8d72a 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/providers/BigQueryWriteConfiguration.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/providers/BigQueryWriteConfiguration.java @@ -198,11 +198,11 @@ public static Builder builder() { public abstract @Nullable List getClusteringFields(); @SchemaFieldDescription( - "Configuration for creating BigLake tables. The following options are available: " - + "connectionId (REQUIRED): the name of your cloud resource connection, " - + "storageUri (REQUIRED): the path to your GCS folder where data will be written to, " - + "fileFormat (OPTIONAL): defaults to 'parquet', " - + "tableFormat (OPTIONAL): defaults to 'iceberg'.") + "Configuration for creating BigLake tables. The following options are available:" + + "\n - connectionId (REQUIRED): the name of your cloud resource connection," + + "\n - storageUri (REQUIRED): the path to your GCS folder where data will be written to," + + "\n - fileFormat (OPTIONAL): defaults to 'parquet'," + + "\n - tableFormat (OPTIONAL): defaults to 'iceberg'.") public abstract java.util.@Nullable Map getBigLakeConfiguration(); /** Builder for {@link BigQueryWriteConfiguration}. */ 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 38d9174cef2b..0e554cec5e2a 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 @@ -511,6 +511,39 @@ def test_streaming_with_at_least_once(self): table = 'streaming_with_at_least_once' self.run_streaming(table_name=table, use_at_least_once=True) + def test_write_with_big_lake_configuration(self): + """Test BigQuery Storage Write API with BigLake configuration.""" + table = 'write_with_big_lake_config' + table_id = '{}:{}.{}'.format(self.project, self.dataset_id, table) + + # BigLake configuration with required parameters + big_lake_config = { + 'connectionId': 'projects/{}/locations/us/connections/test-connection'. + format(self.project), + 'storageUri': 'gs://test-bucket-{}/biglake-data'.format(self.project), + 'fileFormat': 'parquet', + 'tableFormat': 'iceberg' + } + + bq_matcher = BigqueryFullResultMatcher( + project=self.project, + query="SELECT * FROM {}.{}".format(self.dataset_id, table), + data=self.parse_expected_data(self.ELEMENTS)) + + with beam.Pipeline(argv=self.args) as p: + _ = ( + p + | "Create test data" >> beam.Create(self.ELEMENTS) + | beam.io.WriteToBigQuery( + table=table_id, + method=beam.io.WriteToBigQuery.Method.STORAGE_WRITE_API, + schema=self.ALL_TYPES_SCHEMA, + create_disposition='CREATE_IF_NEEDED', + write_disposition='WRITE_TRUNCATE', + big_lake_configuration=big_lake_config)) + + hamcrest_assert(p, bq_matcher) + if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) From 2a9ecdfec1a8f32911a74362b64b8802bed5bb9b Mon Sep 17 00:00:00 2001 From: liferoad Date: Wed, 24 Sep 2025 20:52:46 -0400 Subject: [PATCH 08/14] added the gcs checks --- .../io/external/xlang_bigqueryio_it_test.py | 52 +++++++++++++++++++ 1 file changed, 52 insertions(+) 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 0e554cec5e2a..d05ae6d00c7b 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 @@ -33,6 +33,7 @@ import apache_beam as beam from apache_beam.io.gcp.bigquery import StorageWriteToBigQuery from apache_beam.io.gcp.bigquery_tools import BigQueryWrapper +from apache_beam.io.gcp.gcsio import GcsIO from apache_beam.io.gcp.tests.bigquery_matcher import BigqueryFullResultMatcher from apache_beam.io.gcp.tests.bigquery_matcher import BigqueryFullResultStreamingMatcher from apache_beam.options.pipeline_options import PipelineOptions @@ -145,6 +146,54 @@ def parse_expected_data(self, expected_elements): return data + def assert_iceberg_tables_created( + self, table_prefix, storage_uri, expected_count=1): + """Verify that Iceberg table directories are created in + the warehouse location. + + Args: + table_prefix: The table name prefix to look for + storage_uri: The GCS storage URI (e.g., 'gs://bucket/path') + expected_count: Expected number of table directories + """ + gcs_io = GcsIO() + + # Parse the storage URI to get bucket and prefix + if not storage_uri.startswith('gs://'): + raise ValueError(f'Storage URI must start with gs://, got: {storage_uri}') + + # Remove 'gs://' prefix and split bucket from path + path_parts = storage_uri[5:].split('/', 1) + bucket_name = path_parts[0] + base_prefix = path_parts[1] if len(path_parts) > 1 else '' + + # Construct the full prefix to search for table directories + # Following the pattern: + # {base_prefix}/{class_name}/{project}/{dataset}/{table_prefix} + search_prefix = ( + f"{base_prefix}/{self.__class__.__name__}/" + f"{self.project}/{self.dataset_id}/{table_prefix}") + + # List objects in the bucket with the constructed prefix + try: + objects = gcs_io.list_prefix(f"gs://{bucket_name}/{search_prefix}") + object_count = len(list(objects)) + + if object_count < expected_count: + raise AssertionError( + f"Expected at least {expected_count} objects in warehouse " + f"location gs://{bucket_name}/{search_prefix}, but found " + f"{object_count}") + + _LOGGER.info( + f"Successfully verified {object_count} objects created in " + f"warehouse location gs://{bucket_name}/{search_prefix}") + + except Exception as e: + raise AssertionError( + f"Failed to verify table creation in warehouse location " + f"gs://{bucket_name}/{search_prefix}: {str(e)}") + def run_storage_write_test( self, table_name, items, schema, use_at_least_once=False): table_id = '{}:{}.{}'.format(self.project, self.dataset_id, table_name) @@ -544,6 +593,9 @@ def test_write_with_big_lake_configuration(self): hamcrest_assert(p, bq_matcher) + # Verify that the table directory was created in the warehouse location + self.assert_iceberg_tables_created(table, big_lake_config['storageUri']) + if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) From fe5a3447ef3359787c5418ef4404591e2b8488d5 Mon Sep 17 00:00:00 2001 From: liferoad Date: Sat, 27 Sep 2025 09:48:46 -0400 Subject: [PATCH 09/14] fixed tests --- .../io/external/xlang_bigqueryio_it_test.py | 23 ++++++++++++++----- .../io/gcp/bigquery_biglake_test.py | 6 ++--- 2 files changed, 19 insertions(+), 10 deletions(-) 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 d05ae6d00c7b..3879ecb53701 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 @@ -33,7 +33,6 @@ import apache_beam as beam from apache_beam.io.gcp.bigquery import StorageWriteToBigQuery from apache_beam.io.gcp.bigquery_tools import BigQueryWrapper -from apache_beam.io.gcp.gcsio import GcsIO from apache_beam.io.gcp.tests.bigquery_matcher import BigqueryFullResultMatcher from apache_beam.io.gcp.tests.bigquery_matcher import BigqueryFullResultStreamingMatcher from apache_beam.options.pipeline_options import PipelineOptions @@ -48,6 +47,11 @@ from apitools.base.py.exceptions import HttpError except ImportError: HttpError = None + +try: + from apache_beam.io.gcp.gcsio import GcsIO +except ImportError: + GcsIO = None # pylint: enable=wrong-import-order, wrong-import-position _LOGGER = logging.getLogger(__name__) @@ -148,7 +152,7 @@ def parse_expected_data(self, expected_elements): def assert_iceberg_tables_created( self, table_prefix, storage_uri, expected_count=1): - """Verify that Iceberg table directories are created in + """Verify that Iceberg table directories are created in the warehouse location. Args: @@ -156,6 +160,11 @@ def assert_iceberg_tables_created( storage_uri: The GCS storage URI (e.g., 'gs://bucket/path') expected_count: Expected number of table directories """ + if GcsIO is None: + _LOGGER.warning( + "GcsIO not available, skipping warehouse location verification") + return + gcs_io = GcsIO() # Parse the storage URI to get bucket and prefix @@ -168,8 +177,7 @@ def assert_iceberg_tables_created( base_prefix = path_parts[1] if len(path_parts) > 1 else '' # Construct the full prefix to search for table directories - # Following the pattern: - # {base_prefix}/{class_name}/{project}/{dataset}/{table_prefix} + # Following the pattern: {base_prefix}/{class_name}/{project}/{dataset}/{table_prefix} search_prefix = ( f"{base_prefix}/{self.__class__.__name__}/" f"{self.project}/{self.dataset_id}/{table_prefix}") @@ -186,8 +194,11 @@ def assert_iceberg_tables_created( f"{object_count}") _LOGGER.info( - f"Successfully verified {object_count} objects created in " - f"warehouse location gs://{bucket_name}/{search_prefix}") + "Successfully verified %s objects created in " + "warehouse location gs://%s/%s", + object_count, + bucket_name, + search_prefix) except Exception as e: raise AssertionError( diff --git a/sdks/python/apache_beam/io/gcp/bigquery_biglake_test.py b/sdks/python/apache_beam/io/gcp/bigquery_biglake_test.py index 773523fcedd9..4567ddf20e64 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_biglake_test.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_biglake_test.py @@ -26,8 +26,7 @@ @mock.patch('apache_beam.io.gcp.bigquery.BeamJarExpansionService') class BigQueryBigLakeTest(unittest.TestCase): """Test BigLake configuration support in BigQuery Storage Write API.""" - def test_storage_write_to_bigquery_with_biglake_config( - self, mock_expansion_service): + def test_storage_write_to_bigquery_with_biglake_config(self, mock_expansion_service): """Test that StorageWriteToBigQuery accepts bigLakeConfiguration.""" big_lake_config = { 'connectionId': ( @@ -45,8 +44,7 @@ def test_storage_write_to_bigquery_with_biglake_config( # Verify the configuration is stored self.assertEqual(transform._big_lake_configuration, big_lake_config) - def test_storage_write_to_bigquery_without_biglake_config( - self, mock_expansion_service): + def test_storage_write_to_bigquery_without_biglake_config(self, mock_expansion_service): """Test that StorageWriteToBigQuery works without bigLakeConfiguration.""" transform = bigquery.StorageWriteToBigQuery( table='test-project:test_dataset.test_table') From 909898d95420805d4a73465f2584c8491e4d8ddb Mon Sep 17 00:00:00 2001 From: liferoad Date: Sat, 27 Sep 2025 16:57:00 -0400 Subject: [PATCH 10/14] lint --- .../io/external/xlang_bigqueryio_it_test.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) 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 3879ecb53701..c8ab4185c10c 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 @@ -44,14 +44,14 @@ # pylint: disable=wrong-import-order, wrong-import-position try: - from apitools.base.py.exceptions import HttpError + from apache_beam.io.gcp.gcsio import GcsIO except ImportError: - HttpError = None + GcsIO = None try: - from apache_beam.io.gcp.gcsio import GcsIO + from apitools.base.py.exceptions import HttpError except ImportError: - GcsIO = None + HttpError = None # pylint: enable=wrong-import-order, wrong-import-position _LOGGER = logging.getLogger(__name__) @@ -177,7 +177,8 @@ def assert_iceberg_tables_created( base_prefix = path_parts[1] if len(path_parts) > 1 else '' # Construct the full prefix to search for table directories - # Following the pattern: {base_prefix}/{class_name}/{project}/{dataset}/{table_prefix} + # Following the pattern: + # {base_prefix}/{class_name}/{project}/{dataset}/{table_prefix} search_prefix = ( f"{base_prefix}/{self.__class__.__name__}/" f"{self.project}/{self.dataset_id}/{table_prefix}") From 9bc57ae8b24d54888f804a576f4cb9e84aad2d06 Mon Sep 17 00:00:00 2001 From: liferoad Date: Sat, 27 Sep 2025 20:29:33 -0400 Subject: [PATCH 11/14] lint --- .../apache_beam/io/external/xlang_bigqueryio_it_test.py | 2 +- sdks/python/apache_beam/io/gcp/bigquery_biglake_test.py | 6 ++++-- 2 files changed, 5 insertions(+), 3 deletions(-) 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 c8ab4185c10c..f6af917afaf7 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 @@ -177,7 +177,7 @@ def assert_iceberg_tables_created( base_prefix = path_parts[1] if len(path_parts) > 1 else '' # Construct the full prefix to search for table directories - # Following the pattern: + # Following the pattern: # {base_prefix}/{class_name}/{project}/{dataset}/{table_prefix} search_prefix = ( f"{base_prefix}/{self.__class__.__name__}/" diff --git a/sdks/python/apache_beam/io/gcp/bigquery_biglake_test.py b/sdks/python/apache_beam/io/gcp/bigquery_biglake_test.py index 4567ddf20e64..773523fcedd9 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_biglake_test.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_biglake_test.py @@ -26,7 +26,8 @@ @mock.patch('apache_beam.io.gcp.bigquery.BeamJarExpansionService') class BigQueryBigLakeTest(unittest.TestCase): """Test BigLake configuration support in BigQuery Storage Write API.""" - def test_storage_write_to_bigquery_with_biglake_config(self, mock_expansion_service): + def test_storage_write_to_bigquery_with_biglake_config( + self, mock_expansion_service): """Test that StorageWriteToBigQuery accepts bigLakeConfiguration.""" big_lake_config = { 'connectionId': ( @@ -44,7 +45,8 @@ def test_storage_write_to_bigquery_with_biglake_config(self, mock_expansion_serv # Verify the configuration is stored self.assertEqual(transform._big_lake_configuration, big_lake_config) - def test_storage_write_to_bigquery_without_biglake_config(self, mock_expansion_service): + def test_storage_write_to_bigquery_without_biglake_config( + self, mock_expansion_service): """Test that StorageWriteToBigQuery works without bigLakeConfiguration.""" transform = bigquery.StorageWriteToBigQuery( table='test-project:test_dataset.test_table') From 0a91bbf56902155036209e74a7f1192cde6c48bd Mon Sep 17 00:00:00 2001 From: liferoad Date: Sat, 27 Sep 2025 23:03:59 -0400 Subject: [PATCH 12/14] fixed the tests --- .../apache_beam/io/external/xlang_bigqueryio_it_test.py | 7 +++---- sdks/python/apache_beam/io/gcp/bigquery.py | 5 ++++- 2 files changed, 7 insertions(+), 5 deletions(-) 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 f6af917afaf7..7bb16ab4960d 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 @@ -577,11 +577,10 @@ def test_write_with_big_lake_configuration(self): table = 'write_with_big_lake_config' table_id = '{}:{}.{}'.format(self.project, self.dataset_id, table) - # BigLake configuration with required parameters + # BigLake configuration with required parameters (matching Java test) big_lake_config = { - 'connectionId': 'projects/{}/locations/us/connections/test-connection'. - format(self.project), - 'storageUri': 'gs://test-bucket-{}/biglake-data'.format(self.project), + 'connectionId': 'projects/apache-beam-testing/locations/us/connections/apache-beam-testing-storageapi-biglake-nodelete', # pylint: disable=line-too-long + 'storageUri': 'gs://apache-beam-testing-bq-biglake/BigQueryXlangStorageWriteIT', 'fileFormat': 'parquet', 'tableFormat': 'iceberg' } diff --git a/sdks/python/apache_beam/io/gcp/bigquery.py b/sdks/python/apache_beam/io/gcp/bigquery.py index 828e707fc5cd..0905ba764deb 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery.py +++ b/sdks/python/apache_beam/io/gcp/bigquery.py @@ -1995,7 +1995,8 @@ def __init__( num_streaming_keys=DEFAULT_SHARDS_PER_DESTINATION, use_cdc_writes: bool = False, primary_key: List[str] = None, - expansion_service=None): + expansion_service=None, + big_lake_configuration=None): """Initialize a WriteToBigQuery transform. Args: @@ -2216,6 +2217,7 @@ def __init__( self._num_streaming_keys = num_streaming_keys self._use_cdc_writes = use_cdc_writes self._primary_key = primary_key + self._big_lake_configuration = big_lake_configuration # Dict/schema methods were moved to bigquery_tools, but keep references # here for backward compatibility. @@ -2378,6 +2380,7 @@ def find_in_nested_dict(schema): num_storage_api_streams=self._num_storage_api_streams, use_cdc_writes=self._use_cdc_writes, primary_key=self._primary_key, + big_lake_configuration=self._big_lake_configuration, expansion_service=self.expansion_service) else: raise ValueError(f"Unsupported method {method_to_use}") From 54de061d5c390f0866a930e9381bc330608b2221 Mon Sep 17 00:00:00 2001 From: liferoad Date: Sat, 27 Sep 2025 23:33:19 -0400 Subject: [PATCH 13/14] fixed the test --- .../apache_beam/io/external/xlang_bigqueryio_it_test.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 7bb16ab4960d..d1c7d9dc8dc0 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 @@ -178,9 +178,9 @@ def assert_iceberg_tables_created( # Construct the full prefix to search for table directories # Following the pattern: - # {base_prefix}/{class_name}/{project}/{dataset}/{table_prefix} + # {base_prefix}/{project}/{dataset}/{table_prefix} search_prefix = ( - f"{base_prefix}/{self.__class__.__name__}/" + f"{base_prefix}/" f"{self.project}/{self.dataset_id}/{table_prefix}") # List objects in the bucket with the constructed prefix From 1b035e5feebffd8a8a7c9ba7a49701e07fd9eeed Mon Sep 17 00:00:00 2001 From: liferoad Date: Sun, 28 Sep 2025 08:09:46 -0400 Subject: [PATCH 14/14] lint --- sdks/python/apache_beam/io/external/xlang_bigqueryio_it_test.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 d1c7d9dc8dc0..51ae97b99175 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 @@ -580,7 +580,7 @@ def test_write_with_big_lake_configuration(self): # BigLake configuration with required parameters (matching Java test) big_lake_config = { 'connectionId': 'projects/apache-beam-testing/locations/us/connections/apache-beam-testing-storageapi-biglake-nodelete', # pylint: disable=line-too-long - 'storageUri': 'gs://apache-beam-testing-bq-biglake/BigQueryXlangStorageWriteIT', + 'storageUri': 'gs://apache-beam-testing-bq-biglake/BigQueryXlangStorageWriteIT', # pylint: disable=line-too-long 'fileFormat': 'parquet', 'tableFormat': 'iceberg' }