[2/2] sdks/python: enrich data with CloudSQL [PostgreSQL, MySQL, SQLServer]

.github/workflows/beam_PreCommit_Python_Coverage.yml

@@ -54,7 +54,6 @@ env:
   GRADLE_ENTERPRISE_CACHE_USERNAME: ${{ secrets.GE_CACHE_USERNAME }}
   GRADLE_ENTERPRISE_CACHE_PASSWORD: ${{ secrets.GE_CACHE_PASSWORD }}
   HF_INFERENCE_TOKEN: ${{ secrets.HF_INFERENCE_TOKEN }}
-  ALLOYDB_PASSWORD: ${{ secrets.ALLOYDB_PASSWORD }}
 
 
 jobs:
@@ -113,7 +112,7 @@ jobs:
           TESTCONTAINERS_HOST_OVERRIDE: ${{ contains(matrix.os, 'self-hosted') && env.DIND_IP || '' }}
           TESTCONTAINERS_DOCKER_SOCKET_OVERRIDE: "/var/run/docker.sock"
           TESTCONTAINERS_RYUK_DISABLED: "false"
-          TESTCONTAINERS_RYUK_CONTAINER_PRIVILEGED: "true"          
+          TESTCONTAINERS_RYUK_CONTAINER_PRIVILEGED: "true"
           PYTEST_ADDOPTS: "-v --tb=short --maxfail=3 --durations=20 --reruns=2 --reruns-delay=5"
           TC_TIMEOUT: "120"
           TC_MAX_TRIES: "120"

CHANGES.md

@@ -74,6 +74,7 @@
 ## New Features / Improvements
 
 * X feature added (Java/Python) ([#X](https://github.com/apache/beam/issues/X)).
+* Python examples added for CloudSQL enrichment handler on [Beam website](https://beam.apache.org/documentation/transforms/python/elementwise/enrichment-cloudsql/) (Python) ([#35473](https://github.com/apache/beam/issues/36095)).
 
 ## Breaking Changes
 

sdks/python/apache_beam/examples/snippets/transforms/elementwise/enrichment.py

@@ -116,3 +116,214 @@ def enrichment_with_vertex_ai_legacy():
         | "Enrich W/ Vertex AI" >> Enrichment(vertex_ai_handler)
         | "Print" >> beam.Map(print))
   # [END enrichment_with_vertex_ai_legacy]
+
+
+def enrichment_with_google_cloudsql_pg():
+  # [START enrichment_with_google_cloudsql_pg]
+  import apache_beam as beam
+  from apache_beam.transforms.enrichment import Enrichment
+  from apache_beam.transforms.enrichment_handlers.cloudsql import (
+      CloudSQLEnrichmentHandler,
+      DatabaseTypeAdapter,
+      TableFieldsQueryConfig,
+      CloudSQLConnectionConfig)
+  import os
+
+  database_adapter = DatabaseTypeAdapter.POSTGRESQL
+  database_uri = os.environ.get("GOOGLE_CLOUD_SQL_DB_URI")
+  database_user = os.environ.get("GOOGLE_CLOUD_SQL_DB_USER")
+  database_password = os.environ.get("GOOGLE_CLOUD_SQL_DB_PASSWORD")
+  database_id = os.environ.get("GOOGLE_CLOUD_SQL_DB_ID")
+  table_id = os.environ.get("GOOGLE_CLOUD_SQL_DB_TABLE_ID")
+  where_clause_template = "product_id = :pid"
+  where_clause_fields = ["product_id"]
+
+  data = [
+      beam.Row(product_id=1, name='A'),
+      beam.Row(product_id=2, name='B'),
+      beam.Row(product_id=3, name='C'),
+  ]
+
+  connection_config = CloudSQLConnectionConfig(
+      db_adapter=database_adapter,
+      instance_connection_uri=database_uri,
+      user=database_user,
+      password=database_password,
+      db_id=database_id)
+
+  query_config = TableFieldsQueryConfig(
+      table_id=table_id,
+      where_clause_template=where_clause_template,
+      where_clause_fields=where_clause_fields)
+
+  cloudsql_handler = CloudSQLEnrichmentHandler(
+      connection_config=connection_config,
+      table_id=table_id,
+      query_config=query_config)
+  with beam.Pipeline() as p:
+    _ = (
+        p
+        | "Create" >> beam.Create(data)
+        |
+        "Enrich W/ Google CloudSQL PostgreSQL" >> Enrichment(cloudsql_handler)
+        | "Print" >> beam.Map(print))
+  # [END enrichment_with_google_cloudsql_pg]
+
+
+def enrichment_with_external_pg():
+  # [START enrichment_with_external_pg]
+  import apache_beam as beam
+  from apache_beam.transforms.enrichment import Enrichment
+  from apache_beam.transforms.enrichment_handlers.cloudsql import (
+      CloudSQLEnrichmentHandler,
+      DatabaseTypeAdapter,
+      TableFieldsQueryConfig,
+      ExternalSQLDBConnectionConfig)
+  import os
+
+  database_adapter = DatabaseTypeAdapter.POSTGRESQL
+  database_host = os.environ.get("EXTERNAL_SQL_DB_HOST")
+  database_port = int(os.environ.get("EXTERNAL_SQL_DB_PORT"))
+  database_user = os.environ.get("EXTERNAL_SQL_DB_USER")
+  database_password = os.environ.get("EXTERNAL_SQL_DB_PASSWORD")
+  database_id = os.environ.get("EXTERNAL_SQL_DB_ID")
+  table_id = os.environ.get("EXTERNAL_SQL_DB_TABLE_ID")
+  where_clause_template = "product_id = :pid"
+  where_clause_fields = ["product_id"]
+
+  data = [
+      beam.Row(product_id=1, name='A'),
+      beam.Row(product_id=2, name='B'),
+      beam.Row(product_id=3, name='C'),
+  ]
+
+  connection_config = ExternalSQLDBConnectionConfig(
+      db_adapter=database_adapter,
+      host=database_host,
+      port=database_port,
+      user=database_user,
+      password=database_password,
+      db_id=database_id)
+
+  query_config = TableFieldsQueryConfig(
+      table_id=table_id,
+      where_clause_template=where_clause_template,
+      where_clause_fields=where_clause_fields)
+
+  cloudsql_handler = CloudSQLEnrichmentHandler(
+      connection_config=connection_config,
+      table_id=table_id,
+      query_config=query_config)
+  with beam.Pipeline() as p:
+    _ = (
+        p
+        | "Create" >> beam.Create(data)
+        | "Enrich W/ Unmanaged PostgreSQL" >> Enrichment(cloudsql_handler)
+        | "Print" >> beam.Map(print))
+  # [END enrichment_with_external_pg]
+
+
+def enrichment_with_external_mysql():
+  # [START enrichment_with_external_mysql]
+  import apache_beam as beam
+  from apache_beam.transforms.enrichment import Enrichment
+  from apache_beam.transforms.enrichment_handlers.cloudsql import (
+      CloudSQLEnrichmentHandler,
+      DatabaseTypeAdapter,
+      TableFieldsQueryConfig,
+      ExternalSQLDBConnectionConfig)
+  import os
+
+  database_adapter = DatabaseTypeAdapter.MYSQL
+  database_host = os.environ.get("EXTERNAL_SQL_DB_HOST")
+  database_port = int(os.environ.get("EXTERNAL_SQL_DB_PORT"))
+  database_user = os.environ.get("EXTERNAL_SQL_DB_USER")
+  database_password = os.environ.get("EXTERNAL_SQL_DB_PASSWORD")
+  database_id = os.environ.get("EXTERNAL_SQL_DB_ID")
+  table_id = os.environ.get("EXTERNAL_SQL_DB_TABLE_ID")
+  where_clause_template = "product_id = :pid"
+  where_clause_fields = ["product_id"]
+
+  data = [
+      beam.Row(product_id=1, name='A'),
+      beam.Row(product_id=2, name='B'),
+      beam.Row(product_id=3, name='C'),
+  ]
+
+  connection_config = ExternalSQLDBConnectionConfig(
+      db_adapter=database_adapter,
+      host=database_host,
+      port=database_port,
+      user=database_user,
+      password=database_password,
+      db_id=database_id)
+
+  query_config = TableFieldsQueryConfig(
+      table_id=table_id,
+      where_clause_template=where_clause_template,
+      where_clause_fields=where_clause_fields)
+
+  cloudsql_handler = CloudSQLEnrichmentHandler(
+      connection_config=connection_config,
+      table_id=table_id,
+      query_config=query_config)
+  with beam.Pipeline() as p:
+    _ = (
+        p
+        | "Create" >> beam.Create(data)
+        | "Enrich W/ Unmanaged MySQL" >> Enrichment(cloudsql_handler)
+        | "Print" >> beam.Map(print))
+  # [END enrichment_with_external_mysql]
+
+
+def enrichment_with_external_sqlserver():
+  # [START enrichment_with_external_sqlserver]
+  import apache_beam as beam
+  from apache_beam.transforms.enrichment import Enrichment
+  from apache_beam.transforms.enrichment_handlers.cloudsql import (
+      CloudSQLEnrichmentHandler,
+      DatabaseTypeAdapter,
+      TableFieldsQueryConfig,
+      ExternalSQLDBConnectionConfig)
+  import os
+
+  database_adapter = DatabaseTypeAdapter.SQLSERVER
+  database_host = os.environ.get("EXTERNAL_SQL_DB_HOST")
+  database_port = int(os.environ.get("EXTERNAL_SQL_DB_PORT"))
+  database_user = os.environ.get("EXTERNAL_SQL_DB_USER")
+  database_password = os.environ.get("EXTERNAL_SQL_DB_PASSWORD")
+  database_id = os.environ.get("EXTERNAL_SQL_DB_ID")
+  table_id = os.environ.get("EXTERNAL_SQL_DB_TABLE_ID")
+  where_clause_template = "product_id = :pid"
+  where_clause_fields = ["product_id"]
+
+  data = [
+      beam.Row(product_id=1, name='A'),
+      beam.Row(product_id=2, name='B'),
+      beam.Row(product_id=3, name='C'),
+  ]
+
+  connection_config = ExternalSQLDBConnectionConfig(
+      db_adapter=database_adapter,
+      host=database_host,
+      port=database_port,
+      user=database_user,
+      password=database_password,
+      db_id=database_id)
+
+  query_config = TableFieldsQueryConfig(
+      table_id=table_id,
+      where_clause_template=where_clause_template,
+      where_clause_fields=where_clause_fields)
+
+  cloudsql_handler = CloudSQLEnrichmentHandler(
+      connection_config=connection_config,
+      table_id=table_id,
+      query_config=query_config)
+  with beam.Pipeline() as p:
+    _ = (
+        p
+        | "Create" >> beam.Create(data)
+        | "Enrich W/ Unmanaged SQL Server" >> Enrichment(cloudsql_handler)
+        | "Print" >> beam.Map(print))
+  # [END enrichment_with_external_sqlserver]