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/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy b/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy index f192d5301722..6e3bdef6f136 100644 --- a/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy +++ b/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy @@ -3203,6 +3203,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" @@ -3218,7 +3228,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 { @@ -3231,12 +3241,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/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..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 @@ -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:" + + "\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}. */ @AutoValue.Builder public abstract static class Builder { @@ -231,6 +239,9 @@ 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/external/xlang_bigqueryio_it_test.py b/sdks/python/apache_beam/io/external/xlang_bigqueryio_it_test.py index 38d9174cef2b..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 @@ -43,6 +43,11 @@ # Protect against environments where bigquery library is not available. # pylint: disable=wrong-import-order, wrong-import-position +try: + from apache_beam.io.gcp.gcsio import GcsIO +except ImportError: + GcsIO = None + try: from apitools.base.py.exceptions import HttpError except ImportError: @@ -145,6 +150,62 @@ 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 + """ + 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 + 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}/{project}/{dataset}/{table_prefix} + search_prefix = ( + f"{base_prefix}/" + 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( + "Successfully verified %s objects created in " + "warehouse location gs://%s/%s", + object_count, + 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) @@ -511,6 +572,41 @@ 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 (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', # pylint: disable=line-too-long + '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) + + # 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) diff --git a/sdks/python/apache_beam/io/gcp/bigquery.py b/sdks/python/apache_beam/io/gcp/bigquery.py index aa0ebc12ef18..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}") @@ -2626,6 +2629,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 +2643,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 +2738,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..773523fcedd9 --- /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 +from unittest import mock + +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, mock_expansion_service): + """Test that StorageWriteToBigQuery accepts bigLakeConfiguration.""" + 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, mock_expansion_service): + """Test that StorageWriteToBigQuery works without bigLakeConfiguration.""" + transform = bigquery.StorageWriteToBigQuery( + table='test-project:test_dataset.test_table') + + # Verify the configuration is None by default + self.assertIsNone(transform._big_lake_configuration) + + 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', + big_lake_configuration=big_lake_config) + + # Verify the big_lake_configuration is stored correctly + self.assertEqual(transform._big_lake_configuration, big_lake_config) + + # Verify that the transform has the expected identifier + self.assertEqual( + transform.IDENTIFIER, + "beam:schematransform:org.apache.beam:bigquery_storage_write:v2") + + # 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, mock_expansion_service): + """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()