sdks/python: enrich data with CloudSQL

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

@@ -51,6 +51,51 @@ def enrichment_with_bigtable():
   # [END enrichment_with_bigtable]
 
 
+def enrichment_with_cloudsql():
+  # [START enrichment_with_cloudsql]
+  import apache_beam as beam
+  from apache_beam.transforms.enrichment import Enrichment
+  from apache_beam.transforms.enrichment_handlers.cloudsql import (
+      CloudSQLEnrichmentHandler, DatabaseTypeAdapter, TableFieldsQueryConfig)
+  import os
+
+  database_type_adapter = DatabaseTypeAdapter[os.environ.get("SQL_DB_TYPE")]
+  database_address = os.environ.get("SQL_DB_ADDRESS")
+  database_user = os.environ.get("SQL_DB_USER")
+  database_password = os.environ.get("SQL_DB_PASSWORD")
+  database_id = os.environ.get("SQL_DB_ID")
+  table_id = "products"
+  where_clause_template = "product_id = {}"
+  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'),
+  ]
+
+  query_config = TableFieldsQueryConfig(
+      table_id=table_id,
+      where_clause_template=where_clause_template,
+      where_clause_fields=where_clause_fields)
+
+  cloudsql_handler = CloudSQLEnrichmentHandler(
+      database_type_adapter=database_type_adapter,
+      database_address=database_address,
+      database_user=database_user,
+      database_password=database_password,
+      database_id=database_id,
+      table_id=table_id,
+      query_config=query_config)
+  with beam.Pipeline() as p:
+    _ = (
+        p
+        | "Create" >> beam.Create(data)
+        | "Enrich W/ CloudSQL" >> Enrichment(cloudsql_handler)
+        | "Print" >> beam.Map(print))
+  # [END enrichment_with_cloudsql]
+
+
 def enrichment_with_vertex_ai():
   # [START enrichment_with_vertex_ai]
   import apache_beam as beam

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

@@ -18,16 +18,24 @@
 # pytype: skip-file
 # pylint: disable=line-too-long
 
+import os
 import unittest
 from io import StringIO
 
 import mock
+import pytest
 
 # pylint: disable=unused-import
 try:
-  from apache_beam.examples.snippets.transforms.elementwise.enrichment import enrichment_with_bigtable, \
-  enrichment_with_vertex_ai_legacy
-  from apache_beam.examples.snippets.transforms.elementwise.enrichment import enrichment_with_vertex_ai
+  from sqlalchemy import Column, Integer, String, Engine
+  from apache_beam.examples.snippets.transforms.elementwise.enrichment import (
+      enrichment_with_bigtable, enrichment_with_vertex_ai_legacy)
+  from apache_beam.examples.snippets.transforms.elementwise.enrichment import (
+      enrichment_with_vertex_ai, enrichment_with_cloudsql)
+  from apache_beam.transforms.enrichment_handlers.cloudsql import (
+      DatabaseTypeAdapter)
+  from apache_beam.transforms.enrichment_handlers.cloudsql_it_test import (
+      CloudSQLEnrichmentTestHelper, SQLDBContainerInfo)
   from apache_beam.io.requestresponse import RequestResponseIO
 except ImportError:
   raise unittest.SkipTest('RequestResponseIO dependencies are not installed')
@@ -42,6 +50,15 @@ def validate_enrichment_with_bigtable():
   return expected
 
 
+def validate_enrichment_with_cloudsql():
+  expected = '''[START enrichment_with_cloudsql]
+Row(product_id=1, name='A', quantity=2, region_id=3)
+Row(product_id=2, name='B', quantity=3, region_id=1)
+Row(product_id=3, name='C', quantity=10, region_id=4)
+  [END enrichment_with_cloudsql]'''.splitlines()[1:-1]
+  return expected
+
+
 def validate_enrichment_with_vertex_ai():
   expected = '''[START enrichment_with_vertex_ai]
 Row(user_id='2963', product_id=14235, sale_price=15.0, age=12.0, state='1', gender='1', country='1')
@@ -61,13 +78,28 @@ def validate_enrichment_with_vertex_ai_legacy():
 
 
 @mock.patch('sys.stdout', new_callable=StringIO)
+@pytest.mark.uses_testcontainer
 class EnrichmentTest(unittest.TestCase):
   def test_enrichment_with_bigtable(self, mock_stdout):
     enrichment_with_bigtable()
     output = mock_stdout.getvalue().splitlines()
     expected = validate_enrichment_with_bigtable()
     self.assertEqual(output, expected)
 
+  def test_enrichment_with_cloudsql(self, mock_stdout):
+    db, engine = None, None
+    try:
+      db, engine = self.pre_cloudsql_enrichment_test()
+      enrichment_with_cloudsql()
+      output = mock_stdout.getvalue().splitlines()
+      expected = validate_enrichment_with_cloudsql()
+      self.assertEqual(output, expected)
+    except Exception as e:
+      self.fail(f"Test failed with unexpected error: {e}")
+    finally:
+      if db and engine:
+        self.post_cloudsql_enrichment_test(db, engine)
+
   def test_enrichment_with_vertex_ai(self, mock_stdout):
     enrichment_with_vertex_ai()
     output = mock_stdout.getvalue().splitlines()
@@ -83,6 +115,51 @@ def test_enrichment_with_vertex_ai_legacy(self, mock_stdout):
     self.maxDiff = None
     self.assertEqual(output, expected)
 
+  def pre_cloudsql_enrichment_test(self):
+    table_id = "products"
+    columns = [
+        Column("product_id", Integer, primary_key=True),
+        Column("name", String, nullable=False),
+        Column("quantity", Integer, nullable=False),
+        Column("region_id", Integer, nullable=False),
+    ]
+    table_data = [
+        {
+            "product_id": 1, "name": "A", 'quantity': 2, 'region_id': 3
+        },
+        {
+            "product_id": 2, "name": "B", 'quantity': 3, 'region_id': 1
+        },
+        {
+            "product_id": 3, "name": "C", 'quantity': 10, 'region_id': 4
+        },
+    ]
+    db_adapter = DatabaseTypeAdapter.POSTGRESQL
+    db = CloudSQLEnrichmentTestHelper.start_sql_db_container(db_adapter)
+    os.environ['SQL_DB_TYPE'] = db.adapter.name
+    os.environ['SQL_DB_ADDRESS'] = db.address
+    os.environ['SQL_DB_USER'] = db.user
+    os.environ['SQL_DB_PASSWORD'] = db.password
+    os.environ['SQL_DB_ID'] = db.id
+    os.environ['SQL_DB_URL'] = db.url
+    engine = CloudSQLEnrichmentTestHelper.create_table(
+        db_url=os.environ.get("SQL_DB_URL"),
+        table_id=table_id,
+        columns=columns,
+        table_data=table_data)
+    return db, engine
+
+  def post_cloudsql_enrichment_test(
+      self, db: SQLDBContainerInfo, engine: Engine):
+    engine.dispose(close=True)
+    CloudSQLEnrichmentTestHelper.stop_sql_db_container(db)
+    os.environ.pop('SQL_DB_TYPE', None)
+    os.environ.pop('SQL_DB_ADDRESS', None)
+    os.environ.pop('SQL_DB_USER', None)
+    os.environ.pop('SQL_DB_PASSWORD', None)
+    os.environ.pop('SQL_DB_ID', None)
+    os.environ.pop('SQL_DB_URL', None)
+
 
 if __name__ == '__main__':
   unittest.main()