Add examples for cloudsql

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):

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')

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

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}"

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.