sdks/python: enrich data with CloudSQL

Author: mohamedawnallah

sdks/python/apache_beam/transforms/enrichment_handlers/cloudsql.py

@@ -0,0 +1,166 @@
+#
+# 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.
+#
+import logging
+from collections.abc import Callable
+from enum import Enum
+from typing import Any
+from typing import Optional
+
+from google.cloud.sql.connector import Connector
+
+import apache_beam as beam
+from apache_beam.transforms.enrichment import EnrichmentSourceHandler
+from apache_beam.transforms.enrichment_handlers.utils import ExceptionLevel
+
+__all__ = [
+    'CloudSQLEnrichmentHandler',
+]
+
+RowKeyFn = Callable[[beam.Row], str]
+
+_LOGGER = logging.getLogger(__name__)
+
+class DatabaseTypeAdapter(Enum):
+    POSTGRESQL = "pg8000"
+    MYSQL = "pymysql"
+    SQLSERVER = "pytds"
+
+    def __str__(self):
+        return self.value
+
+class CloudSQLEnrichmentHandler(EnrichmentSourceHandler[beam.Row, beam.Row]):
+  """A handler for :class:`apache_beam.transforms.enrichment.Enrichment`
+  transform to interact with Google Cloud SQL databases.
+
+  Args:
+    project_id (str): GCP project-id of the Cloud SQL instance.
+    region_id (str): GCP region-id of the Cloud SQL instance.
+    instance_id (str): GCP instance-id of the Cloud SQL instance.
+    database_type_adapter (DatabaseTypeAdapter): The type of database adapter to use.
+      Supported adapters are: POSTGRESQL (pg8000), MYSQL (pymysql), and SQLSERVER (pytds).
+    database_name (str): The name of the database to connect to.
+    database_user (str): The username for connecting to the database.
+    database_password (str): The password for connecting to the database.
+    table_id (str): The name of the table to query.
+    row_key (str): Field name from the input `beam.Row` object to use as 
+      identifier for database querying.
+    row_key_fn: A lambda function that returns a string key from the
+      input row. Used to build/extract the identifier for the database query.
+    exception_level: A `enum.Enum` value from
+      ``apache_beam.transforms.enrichment_handlers.utils.ExceptionLevel``
+      to set the level when no matching record is found from the database query.
+      Defaults to ``ExceptionLevel.WARN``.
+  """
+  def __init__(
+      self,
+      region_id: str,
+      project_id: str,
+      instance_id: str,
+      database_type_adapter: DatabaseTypeAdapter,
+      database_name: str,
+      database_user: str,
+      database_password: str,
+      table_id: str,
+      row_key: str = "",
+      *,
+      row_key_fn: Optional[RowKeyFn] = None,
+      exception_level: ExceptionLevel = ExceptionLevel.WARN,
+  ):
+    self._project_id = project_id
+    self._region_id = region_id
+    self._instance_id = instance_id
+    self._database_type_adapter = database_type_adapter
+    self._database_name = database_name
+    self._database_user = database_user
+    self._database_password = database_password
+    self._table_id = table_id
+    self._row_key = row_key
+    self._row_key_fn = row_key_fn
+    self._exception_level = exception_level
+    if ((not self._row_key_fn and not self._row_key) or
+        bool(self._row_key_fn and self._row_key)):
+      raise ValueError(
+          "Please specify exactly one of `row_key` or a lambda "
+          "function with `row_key_fn` to extract the row key "
+          "from the input row.")
+
+  def __enter__(self):
+    """Connect to the the Cloud SQL instance."""
+    connector = Connector()
+    self.client = connector.connect(
+      f"{self._project_id}:{self._region_id}:{self._instance_id}",
+      driver=self._database_type_adapter.value,
+      user=self._database_user,
+      password=self._database_password,
+      db=self._database_name
+    )
+    self.cursor = self.client.cursor()
+
+  def __call__(self, request: beam.Row, *args, **kwargs):
+    """
+    Executes a query to the Cloud SQL instance and returns
+    a `Tuple` of request and response.
+
+    Args:
+    request: the input `beam.Row` to enrich.
+    """
+    response_dict: dict[str, Any] = {}
+    row_key_str: str = ""
+
+    try:
+      if self._row_key_fn:
+        row_key = self._row_key_fn(request)
+      else:
+        request_dict = request._asdict()
+        row_key_str = str(request_dict[self._row_key])
+        row_key = row_key_str
+      
+      query = f"SELECT * FROM {self._table_id} WHERE {self._row_key} = %s"
+      self.cursor.execute(query, (row_key,))
+      result = self.cursor.fetchone()
+      
+      if result:
+        columns = [col[0] for col in self.cursor.description]
+        for i, value in enumerate(result):
+          response_dict[columns[i]] = value
+      elif self._exception_level == ExceptionLevel.WARN:
+            _LOGGER.warning(
+                'No matching record found for row_key: %s in table: %s',
+                row_key_str, self._table_id)
+      elif self._exception_level == ExceptionLevel.RAISE:
+            raise ValueError(
+                'No matching record found for row_key: %s in table: %s' % 
+                (row_key_str, self._table_id))
+    except KeyError:
+        raise KeyError('row_key %s not found in input PCollection.' % row_key_str)
+    except Exception as e:
+        raise e
+      
+    return request, beam.Row(**response_dict)
+
+  def __exit__(self, exc_type, exc_val, exc_tb):
+    """Clean the instantiated the Cloud SQL client."""
+    self.cursor.close()
+    self.client.close()
+    self.cursor, self.client = None, None
+
+def get_cache_key(self, request: beam.Row) -> str:
+    """Returns a string formatted with row key since it is unique to
+    a request made to the Cloud SQL instance."""
+    if self._row_key_fn:
+        return f"row_key: {str(self._row_key_fn(request))}"
+    return f"{self._row_key}: {request._asdict()[self._row_key]}"