diff --git a/sdks/python/apache_beam/yaml/examples/testing/examples_test.py b/sdks/python/apache_beam/yaml/examples/testing/examples_test.py index 97acf427f548..3ea7e0f07d80 100644 --- a/sdks/python/apache_beam/yaml/examples/testing/examples_test.py +++ b/sdks/python/apache_beam/yaml/examples/testing/examples_test.py @@ -669,6 +669,7 @@ def _kafka_test_preprocessor( 'test_gcs_text_to_bigquery_yaml', 'test_sqlserver_to_bigquery_yaml', 'test_postgres_to_bigquery_yaml', + 'test_cloudsql_to_bigquery_yaml', 'test_kafka_to_iceberg_yaml', 'test_pubsub_to_iceberg_yaml', 'test_oracle_to_bigquery_yaml', @@ -931,6 +932,7 @@ def __sqlserver_io_read_test_preprocessor( @YamlExamplesTestSuite.register_test_preprocessor([ 'test_postgres_to_bigquery_yaml', + 'test_cloudsql_to_bigquery_yaml', ]) def __postgres_io_read_test_preprocessor( test_spec: dict, expected: list[str], env: TestEnvironment): diff --git a/sdks/python/apache_beam/yaml/examples/transforms/blueprint/cloudsql_to_bigquery.yaml b/sdks/python/apache_beam/yaml/examples/transforms/blueprint/cloudsql_to_bigquery.yaml new file mode 100644 index 000000000000..ee8885f28153 --- /dev/null +++ b/sdks/python/apache_beam/yaml/examples/transforms/blueprint/cloudsql_to_bigquery.yaml @@ -0,0 +1,52 @@ +# coding=utf-8 +# +# 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. + +# This is an example of a Beam YAML pipeline that reads from a Google Cloud SQL +# PostgreSQL database using the Cloud SQL JDBC Socket Factory and writes to BigQuery. + +pipeline: + type: composite + transforms: + # Step 1: Reading data from Cloud SQL Postgres + - type: ReadFromPostgres + name: ReadFromCloudSqlPostgres + config: + url: "jdbc:postgresql:///shipment?cloudSqlInstance=my-project:us-central1:my-instance&socketFactory=com.google.cloud.sql.postgres.SocketFactory" + query: "SELECT * FROM shipments" + driver_class_name: "org.postgresql.Driver" + username: "my-username" + password: "my-password" + # Step 2: Write records out to BigQuery + - type: WriteToBigQuery + name: WriteShipments + input: ReadFromCloudSqlPostgres + config: + table: "apache-beam-testing.yaml_test.shipments" + create_disposition: "CREATE_NEVER" + write_disposition: "WRITE_APPEND" + num_streams: 1 + +options: + temp_location: "gs://apache-beam-testing/temp" + +# Expected: +# Row(shipment_id='S1', customer_id='C1', shipment_date='2023-05-01', shipment_cost=150.0, customer_name='Alice', customer_email='alice@example.com') +# Row(shipment_id='S2', customer_id='C2', shipment_date='2023-06-12', shipment_cost=300.0, customer_name='Bob', customer_email='bob@example.com') +# Row(shipment_id='S3', customer_id='C1', shipment_date='2023-05-10', shipment_cost=20.0, customer_name='Alice', customer_email='alice@example.com') +# Row(shipment_id='S4', customer_id='C4', shipment_date='2024-07-01', shipment_cost=150.0, customer_name='Derek', customer_email='derek@example.com') +# Row(shipment_id='S5', customer_id='C5', shipment_date='2023-05-09', shipment_cost=300.0, customer_name='Erin', customer_email='erin@example.com') +# Row(shipment_id='S6', customer_id='C4', shipment_date='2024-07-02', shipment_cost=150.0, customer_name='Derek', customer_email='derek@example.com') diff --git a/sdks/python/apache_beam/yaml/examples/transforms/blueprint/postgres_to_bigquery.yaml b/sdks/python/apache_beam/yaml/examples/transforms/blueprint/postgres_to_bigquery.yaml index b532636f46ee..4ca4cd2461ce 100644 --- a/sdks/python/apache_beam/yaml/examples/transforms/blueprint/postgres_to_bigquery.yaml +++ b/sdks/python/apache_beam/yaml/examples/transforms/blueprint/postgres_to_bigquery.yaml @@ -15,9 +15,8 @@ # See the License for the specific language governing permissions and # limitations under the License. -# This is an example of a Beam YAML pipeline that reads from spanner database -# and writes to GCS avro files. This matches the Dataflow Template located -# here - https://cloud.google.com/dataflow/docs/guides/templates/provided/cloud-spanner-to-avro +# This is an example of a Beam YAML pipeline that reads from a PostgreSQL database +# and writes to BigQuery. pipeline: type: composite diff --git a/sdks/python/apache_beam/yaml/extended_tests/e2e/cloudsql_to_bigquery.yaml b/sdks/python/apache_beam/yaml/extended_tests/e2e/cloudsql_to_bigquery.yaml new file mode 100644 index 000000000000..a409359626d2 --- /dev/null +++ b/sdks/python/apache_beam/yaml/extended_tests/e2e/cloudsql_to_bigquery.yaml @@ -0,0 +1,82 @@ +# +# 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. +# + +fixtures: + - name: CLOUDSQL + type: "apache_beam.yaml.integration_tests.cloudsql_postgres_fixture" + - name: BQ_TABLE + type: "apache_beam.yaml.integration_tests.temp_bigquery_table" + config: + project: "apache-beam-testing" + - name: TEMP_DIR + type: "apache_beam.yaml.integration_tests.gcs_temp_dir" + config: + bucket: "gs://temp-storage-for-end-to-end-tests/temp-it" + +pipelines: + # Pipeline 1: Write records to Cloud SQL Postgres + - pipeline: + type: chain + transforms: + - type: Create + config: + elements: + - {id: 1, name: "Alice", score: 95.5} + - {id: 2, name: "Bob", score: 88.0} + - {id: 3, name: "Charlie", score: 92.3} + - type: WriteToPostgres + config: + url: "{CLOUDSQL[JDBC_URL]}" + driver_class_name: "{CLOUDSQL[DRIVER_CLASS_NAME]}" + username: "{CLOUDSQL[USERNAME]}" + password: "{CLOUDSQL[PASSWORD]}" + table: "{CLOUDSQL[TABLE]}" + + # Pipeline 2: Read from Cloud SQL Postgres and write to BigQuery + - pipeline: + type: chain + transforms: + - type: ReadFromPostgres + config: + url: "{CLOUDSQL[JDBC_URL]}" + driver_class_name: "{CLOUDSQL[DRIVER_CLASS_NAME]}" + username: "{CLOUDSQL[USERNAME]}" + password: "{CLOUDSQL[PASSWORD]}" + query: "SELECT id, name, score FROM {CLOUDSQL[TABLE]}" + - type: WriteToBigQuery + config: + table: "{BQ_TABLE}" + options: + project: "apache-beam-testing" + temp_location: "{TEMP_DIR}" + + # Pipeline 3: Read from BigQuery and verify records + - pipeline: + type: chain + transforms: + - type: ReadFromBigQuery + config: + table: "{BQ_TABLE}" + - type: AssertEqual + config: + elements: + - {id: 1, name: "Alice", score: 95.5} + - {id: 2, name: "Bob", score: 88.0} + - {id: 3, name: "Charlie", score: 92.3} + options: + project: "apache-beam-testing" + temp_location: "{TEMP_DIR}" diff --git a/sdks/python/apache_beam/yaml/integration_tests.py b/sdks/python/apache_beam/yaml/integration_tests.py index 32794f4588d2..0d112085901f 100644 --- a/sdks/python/apache_beam/yaml/integration_tests.py +++ b/sdks/python/apache_beam/yaml/integration_tests.py @@ -656,6 +656,41 @@ def temp_postgres_database_with_secret_manager( _LOGGER.warning("Could not delete GCP secret %s: %s", secret_path, err) +@contextlib.contextmanager +def cloudsql_postgres_fixture(): + """Context manager to provide a PostgreSQL testcontainer database for testing.""" + default_port = 5432 + with PostgresContainer(port=default_port) as postgres_container: + try: + engine = sqlalchemy.create_engine(postgres_container.get_connection_url()) + with engine.begin() as connection: + connection.execute( + sqlalchemy.text( + "CREATE TABLE tmp_table (id INTEGER, name VARCHAR(255), score FLOAT);" + )) + + jdbc_url = ( + f"jdbc:postgresql://{postgres_container.get_container_host_ip()}:" + f"{postgres_container.get_exposed_port(default_port)}/" + f"{postgres_container.dbname}?" + f"user={postgres_container.username}&" + f"password={postgres_container.password}") + + yield { + 'JDBC_URL': jdbc_url, + 'DRIVER_CLASS_NAME': 'org.postgresql.Driver', + 'USERNAME': postgres_container.username, + 'PASSWORD': postgres_container.password, + 'DATABASE': postgres_container.dbname, + 'TABLE': 'tmp_table', + } + except (psycopg2.Error, Exception) as err: + logging.error( + "Error interacting with temporary Postgres DB in cloudsql_postgres_fixture: %s", + err) + raise err + + @contextlib.contextmanager def temp_sqlserver_database(): """Context manager to provide a temporary SQL Server database for testing.