cloudsql.py

#
# 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 re
from abc import ABC
from abc import abstractmethod
from collections.abc import Callable
from collections.abc import Mapping
from dataclasses import dataclass
from dataclasses import field
from enum import Enum
from typing import Any
from typing import Dict
from typing import List
from typing import Optional
from typing import Union

import pg8000
import pymysql
import pytds
from google.cloud.sql.connector import Connector as CloudSQLConnector
from google.cloud.sql.connector.enums import RefreshStrategy
from sqlalchemy import create_engine
from sqlalchemy import text
from sqlalchemy.engine import Connection as DBAPIConnection

import apache_beam as beam
from apache_beam.transforms.enrichment import EnrichmentSourceHandler

QueryFn = Callable[[beam.Row], str]
ConditionValueFn = Callable[[beam.Row], list[Any]]


@dataclass
class CustomQueryConfig:
  """Configuration for using a custom query function."""
  query_fn: QueryFn

  def __post_init__(self):
    if not self.query_fn:
      raise ValueError("CustomQueryConfig must provide a valid query_fn")


@dataclass
class TableFieldsQueryConfig:
  """Configuration for using table name, where clause, and field names."""
  table_id: str
  where_clause_template: str
  where_clause_fields: List[str]

  def __post_init__(self):
    if not self.table_id or not self.where_clause_template:
      raise ValueError(
          "TableFieldsQueryConfig must provide table_id and " +
          "where_clause_template")

    if not self.where_clause_fields:
      raise ValueError(
          "TableFieldsQueryConfig must provide non-empty " +
          "where_clause_fields")


@dataclass
class TableFunctionQueryConfig:
  """Configuration for using table name, where clause, and a value function."""
  table_id: str
  where_clause_template: str
  where_clause_value_fn: ConditionValueFn

  def __post_init__(self):
    if not self.table_id or not self.where_clause_template:
      raise ValueError(
          "TableFunctionQueryConfig must provide table_id and " +
          "where_clause_template")

    if not self.where_clause_value_fn:
      raise ValueError(
          "TableFunctionQueryConfig must provide " + "where_clause_value_fn")


class DatabaseTypeAdapter(Enum):
  POSTGRESQL = "pg8000"
  MYSQL = "pymysql"
  SQLSERVER = "pytds"

  def to_sqlalchemy_dialect(self):
    """Map the adapter type to its corresponding SQLAlchemy dialect.

    Returns:
        str: SQLAlchemy dialect string.
    """
    if self == DatabaseTypeAdapter.POSTGRESQL:
      return f"postgresql+{self.value}"
    elif self == DatabaseTypeAdapter.MYSQL:
      return f"mysql+{self.value}"
    elif self == DatabaseTypeAdapter.SQLSERVER:
      return f"mssql+{self.value}"
    else:
      raise ValueError(f"Unsupported database adapter type: {self.name}")


class ConnectionConfig(ABC):
  @abstractmethod
  def get_connector_handler(self) -> Callable[[], DBAPIConnection]:
    pass

  @abstractmethod
  def get_db_url(self) -> str:
    pass


@dataclass
class CloudSQLConnectionConfig(ConnectionConfig):
  """Connects to Google Cloud SQL using Cloud SQL Python Connector.

    Args:
        db_adapter: The database adapter type (PostgreSQL, MySQL, SQL Server).
        instance_connection_uri: URI for connecting to the Cloud SQL instance.
        user: Username for authentication.
        password: Password for authentication. Defaults to None.
        db_id: Database identifier/name.
        refresh_strategy: Strategy for refreshing connection (default: LAZY).
        connector_kwargs: Additional keyword arguments for the
          Cloud SQL Python Connector. Enables forward compatibility.
        connect_kwargs: Additional keyword arguments for the client connect
          method. Enables forward compatibility.
    """
  db_adapter: DatabaseTypeAdapter
  instance_connection_uri: str
  user: str = field(default_factory=str)
  password: str = field(default_factory=str)
  db_id: str = field(default_factory=str)
  refresh_strategy: RefreshStrategy = RefreshStrategy.LAZY
  connector_kwargs: Dict[str, Any] = field(default_factory=dict)
  connect_kwargs: Dict[str, Any] = field(default_factory=dict)

  def __post_init__(self):
    if not self.instance_connection_uri:
      raise ValueError("Instance connection URI cannot be empty")

  def get_connector_handler(self) -> Callable[[], DBAPIConnection]:
    """Returns a function that creates a new database connection.

      The returned connector function creates database connections that should
      be properly closed by the caller when no longer needed.
      """
    cloudsql_client = CloudSQLConnector(
        refresh_strategy=self.refresh_strategy, **self.connector_kwargs)

    cloudsql_connector = lambda: cloudsql_client.connect(
        instance_connection_string=self.instance_connection_uri, driver=self.
        db_adapter.value, user=self.user, password=self.password, db=self.db_id,
        **self.connect_kwargs)

    return cloudsql_connector

  def get_db_url(self) -> str:
    return self.db_adapter.to_sqlalchemy_dialect() + "://"


@dataclass
class ExternalSQLDBConnectionConfig(ConnectionConfig):
  """Connects to External SQL DBs (PostgreSQL, MySQL, SQL Server) over TCP.

    Args:
        db_adapter: The database adapter type (PostgreSQL, MySQL, SQL Server).
        host: Hostname or IP address of the database server.
        port: Port number for the database connection.
        user: Username for authentication.
        password: Password for authentication.
        db_id: Database identifier/name.
        connect_kwargs: Additional keyword arguments for the client connect
          method. Enables forward compatibility.
    """
  db_adapter: DatabaseTypeAdapter
  host: str
  port: int
  user: str = field(default_factory=str)
  password: str = field(default_factory=str)
  db_id: str = field(default_factory=str)
  connect_kwargs: Dict[str, Any] = field(default_factory=dict)

  def __post_init__(self):
    if not self.host:
      raise ValueError("Database host cannot be empty")

  def get_connector_handler(self) -> Callable[[], DBAPIConnection]:
    """Returns a function that creates a new database connection.

      The returned connector function creates database connections that should
      be properly closed by the caller when no longer needed.
      """
    if self.db_adapter == DatabaseTypeAdapter.POSTGRESQL:
      return lambda: pg8000.connect(
          host=self.host, port=self.port, database=self.db_id, user=self.user,
          password=self.password, **self.connect_kwargs)
    elif self.db_adapter == DatabaseTypeAdapter.MYSQL:
      return lambda: pymysql.connect(
          host=self.host, port=self.port, database=self.db_id, user=self.user,
          password=self.password, **self.connect_kwargs)
    elif self.db_adapter == DatabaseTypeAdapter.SQLSERVER:
      return lambda: pytds.connect(
          dsn=self.host, port=self.port, database=self.db_id, user=self.user,
          password=self.password, **self.connect_kwargs)
    else:
      raise ValueError(f"Unsupported database adapter: {self.db_adapter}")

  def get_db_url(self) -> str:
    return self.db_adapter.to_sqlalchemy_dialect() + "://"


QueryConfig = Union[CustomQueryConfig,
                    TableFieldsQueryConfig,
                    TableFunctionQueryConfig]


class CloudSQLEnrichmentHandler(EnrichmentSourceHandler[beam.Row, beam.Row]):
  """Enrichment handler for Cloud SQL databases.

  This handler is designed to work with the
  :class:`apache_beam.transforms.enrichment.Enrichment` transform.

  To use this handler, you need to provide one of the following query configs:
    * CustomQueryConfig - For providing a custom query function
    * TableFieldsQueryConfig - For specifying table, where clause, and fields
    * TableFunctionQueryConfig - For specifying table, where clause, and val fn

  By default, the handler retrieves all columns from the specified table.
  To limit the columns, use the `column_names` parameter to specify
  the desired column names.

  This handler queries the Cloud SQL database per element by default.
  To enable batching, set the `min_batch_size`, `max_batch_size`, and
  `max_batch_duration_secs` parameters. These values control batching behavior
  in the
  :class:`apache_beam.transforms.utils.BatchElements` transform.

  NOTE: Batching is not supported when using the CustomQueryConfig.
  """
  def __init__(
      self,
      connection_config: ConnectionConfig,
      *,
      query_config: QueryConfig,
      column_names: Optional[list[str]] = None,
      min_batch_size: int = 1,
      max_batch_size: int = 10000,
      max_batch_duration_secs: Optional[float] = None,
      **kwargs,
  ):
    """
    Example usage::

      connection_config = CloudSQLConnectionConfig(
        db_adapter=DatabaseTypeAdapter.POSTGRESQL,
        instance_connection_uri="apache-beam-testing:us-central1:itests",
        user='postgres',
        password= os.getenv("CLOUDSQL_PG_PASSWORD"))
      query_config=TableFieldsQueryConfig('my_table',"id = :param0",['id']),
      cloudsql_handler = CloudSQLEnrichmentHandler(
        connection_config=connection_config,
        query_config=query_config,
        min_batch_size=2,
        max_batch_size=100)

    Args:
      connection_config (ConnectionConfig): Configuration for connecting to
        the SQL database. Must be an instance of a subclass of
        `ConnectionConfig`, such as `CloudSQLConnectionConfig` or
        `ExternalSQLDBConnectionConfig`. This determines how the handler
        connects to the target SQL database.
      query_config: Configuration for database queries. Must be one of:
      * CustomQueryConfig: For providing a custom query function
      * TableFieldsQueryConfig: specifies table, where clause, and field names
      * TableFunctionQueryConfig: specifies table, where clause, and val func
      column_names (Optional[list[str]]): List of column names to select from
        the Cloud SQL table. If not provided, all columns (`*`) are selected.
      min_batch_size (int): Minimum number of rows to batch together when
        querying the database. Defaults to 1 if `query_fn` is not used.
      max_batch_size (int): Maximum number of rows to batch together. Defaults
        to 10,000 if `query_fn` is not used.
      max_batch_duration_secs (float): Maximum amount of time in seconds to
        buffer a batch before emitting it. If not provided, batching duration
        is determined by `BatchElements` defaults.
      **kwargs: Additional keyword arguments for database connection or query
        handling.

    Note:
      * `min_batch_size`, `max_batch_size`, and `max_batch_duration_secs`
        are not used with `query_fn`.
      * Either `where_clause_fields` or `where_clause_value_fn` must be provided
        for query construction if `query_fn` is not provided.
      * Ensure that the database user has the necessary permissions to query the
        specified table.
    """
    self._connection_config = connection_config
    self._query_config = query_config
    self._column_names = ",".join(column_names) if column_names else "*"
    self.kwargs = kwargs
    self._batching_kwargs = {}
    table_query_configs = (TableFieldsQueryConfig, TableFunctionQueryConfig)
    if isinstance(query_config, table_query_configs):
      self.query_template = (
          f"SELECT {self._column_names} "
          f"FROM {query_config.table_id} "
          f"WHERE {query_config.where_clause_template}")
      self._batching_kwargs['min_batch_size'] = min_batch_size
      self._batching_kwargs['max_batch_size'] = max_batch_size
      if max_batch_duration_secs is not None:
        self._batching_kwargs[
            'max_batch_duration_secs'] = max_batch_duration_secs

  def __enter__(self):
    connector = self._connection_config.get_connector_handler()
    self._engine = create_engine(
        url=self._connection_config.get_db_url(), creator=connector)

  def __call__(
      self, request: Union[beam.Row, list[beam.Row]], *_args, **_kwargs):
    """Handle requests by delegating to single or batch processing."""
    if isinstance(request, list):
      return self._process_batch_request(request)
    else:
      return self._process_single_request(request)

  def _process_single_request(self, request: beam.Row):
    """Process a single request and return with its response."""
    response: Union[List[Dict[str, Any]], Dict[str, Any]]
    if isinstance(self._query_config, CustomQueryConfig):
      query = self._query_config.query_fn(request)
      response = self._execute_query(query, is_batch=False)
    else:
      values = self._extract_values_from_request(request)
      param_dict = self._build_single_param_dict(values)
      response = self._execute_query(
          self.query_template, params=param_dict, is_batch=False)
    return request, beam.Row(**response)  # type: ignore[arg-type]

  def _process_batch_request(self, requests: list[beam.Row]):
    """Process batch requests and match responses to original requests."""
    values, responses = [], []
    requests_map: dict[Any, Any] = {}
    batch_size = len(requests)

    # Build the appropriate query (single or batched).
    raw_query = self._build_batch_query(requests, batch_size)

    # Extract where_clause_fields values and map the generated request key to
    # the original request object..
    for req in requests:
      current_values = self._extract_values_from_request(req)
      values.extend(current_values)
      requests_map[self.create_row_key(req)] = req

    # Build named parameters dictionary for parameterized query.
    param_dict = self._build_parameters_dict(requests, batch_size)

    # Execute the parameterized query with validated parameters.
    result: Union[List[Dict[str, Any]], Dict[str, Any]] = self._execute_query(
        raw_query, params=param_dict, is_batch=True)
    for response in result:
      response_row = beam.Row(**response)  # type: ignore[arg-type]
      response_key = self.create_row_key(response_row)
      if response_key in requests_map:
        responses.append((requests_map[response_key], response_row))
    return responses

  def _execute_query(
      self,
      query: str,
      params: Optional[dict] = None,
      is_batch: bool = False) -> Union[List[Dict[str, Any]], Dict[str, Any]]:
    connection = None
    try:
      connection = self._engine.connect()
      transaction = connection.begin()
      try:
        if params:
          result = connection.execute(text(query), params)
        else:
          result = connection.execute(text(query))
        # Materialize results while transaction is active.
        data: Union[List[Dict[str, Any]], Dict[str, Any]]
        if is_batch:
          data = [row._asdict() for row in result]
        else:
          result_row = result.first()
          data = result_row._asdict() if result_row else {}
        # Explicitly commit the transaction.
        transaction.commit()
        return data
      except Exception as e:
        transaction.rollback()
        raise RuntimeError(f"Database operation failed: {e}") from e
    except Exception as e:
      raise Exception(
          f'Could not execute the query. Please check if the query is properly '
          f'formatted and the table exists. {e}') from e
    finally:
      if connection:
        connection.close()

  def _build_batch_query(
      self, requests: list[beam.Row], batch_size: int) -> str:
    """Build batched query with unique parameter names for multiple requests.

    This method extracts parameter placeholders from the where_clause_template
    using regex and creates unique parameter names for each batch item. The
    parameter names in the template can be any valid identifiers (e.g., :id,
    :param_0, :user_name) and don't need to match field names exactly.

    For batch queries, placeholders are replaced with unique names like
    :batch_0_id, :batch_1_param_0, etc., based on the actual parameter names
    found in the template.

    Args:
      requests: List of beam.Row requests to process
      batch_size: Number of requests in the batch

    Returns:
      SQL query string with batched WHERE clauses using unique parameter names
    """
    # Single request - return original query.
    if batch_size <= 1:
      return self.query_template

    # Only batch table-based query configs.
    table_query_configs = (TableFieldsQueryConfig, TableFunctionQueryConfig)
    if not isinstance(self._query_config, table_query_configs):
      return self.query_template

    # Build batched WHERE clauses.
    where_clauses = [self._create_batch_clause(i) for i in range(batch_size)]

    # Combine clauses and update query.
    where_clause_batched = ' OR '.join(where_clauses)
    # We know this is a table-based config from the check above.
    assert isinstance(self._query_config, table_query_configs)
    return self.query_template.replace(
        self._query_config.where_clause_template, where_clause_batched)

  def _create_batch_clause(self, batch_index: int) -> str:
    """Create a WHERE clause for a single batch item with unique parameter
    names."""
    # This method is only called for table-based query configs
    table_query_configs = (TableFieldsQueryConfig, TableFunctionQueryConfig)
    assert isinstance(self._query_config, table_query_configs)
    clause = self._query_config.where_clause_template

    # Extract parameter names from the template using regex.
    param_names = self._extract_parameter_names(
        self._query_config.where_clause_template)
    for param_name in param_names:
      old_param = f':{param_name}'
      new_param = f':batch_{batch_index}_{param_name}'
      clause = clause.replace(old_param, new_param)

    return f'({clause})'

  def _build_parameters_dict(
      self, requests: list[beam.Row], batch_size: int) -> dict:
    """Build named parameters dictionary for parameterized queries.

    Args:
      requests: List of beam.Row requests to process
      batch_size: Number of requests in the batch

    Returns:
      Dictionary mapping parameter names to validated values
    """
    param_dict = {}
    for i, req in enumerate(requests):
      current_values = self._extract_values_from_request(req)

      # For batched queries, use unique parameter names per batch item.
      if batch_size > 1:
        # Batching is only used with table-based query configs.
        table_query_configs = (TableFieldsQueryConfig, TableFunctionQueryConfig)
        assert isinstance(self._query_config, table_query_configs)
        batch_param_dict = self._build_single_param_dict(current_values)
        # Prefix batch parameters to make them globally unique.
        for param_name, val in batch_param_dict.items():
          param_dict[f'batch_{i}_{param_name}'] = val
      else:
        single_param_dict = self._build_single_param_dict(current_values)
        param_dict.update(single_param_dict)

    return param_dict

  def _build_single_param_dict(self, values: list[Any]) -> dict[str, Any]:
    """Build parameter dictionary for single request processing.

    Args:
      values: List of parameter values

    Returns:
      Dictionary mapping parameter names to values
    """
    table_query_configs = (TableFieldsQueryConfig, TableFunctionQueryConfig)
    if not isinstance(self._query_config, table_query_configs):
      raise ValueError(
          f"Parameter binding not supported for "
          f"{type(self._query_config).__name__}")

    _, param_dict = self._get_unique_template_and_params(
        self._query_config.where_clause_template, values)
    return param_dict

  def _get_unique_template_and_params(
      self, template: str, values: list[Any]) -> tuple[str, dict[str, Any]]:
    """Generate unique binding parameter names for duplicate templates.

    Args:
      template: SQL template with potentially duplicate binding parameter names
      values: List of parameter values

    Returns:
      Tuple of (updated_template, param_dict) with unique binding names.
    """
    param_names = self._extract_parameter_names(template)
    unique_param_names = [
        f"{param_name}_{i}" if param_names.count(param_name) > 1 else param_name
        for i, param_name in enumerate(param_names)
    ]

    # Update template by replacing each parameter occurrence in order.
    updated_template = template
    param_positions = []

    # Find all parameter positions.
    for match in re.finditer(r':(\w+)', template):
      param_positions.append((match.start(), match.end(), match.group(1)))

    # Replace parameters from right to left to avoid position shifts.
    for i in reversed(range(len(param_positions))):
      start, end, _ = param_positions[i]
      unique_name = unique_param_names[i]
      updated_template = (
          updated_template[:start] + f':{unique_name}' + updated_template[end:])

    # Build parameter dictionary.
    param_dict = {
        unique_name: val
        for unique_name, val in zip(unique_param_names, values)
    }

    return updated_template, param_dict

  def _extract_values_from_request(self, request: beam.Row) -> list[Any]:
    """Extract parameter values from a request based on query configuration.

    Args:
      request: The beam.Row request to extract values from

    Returns:
      List of parameter values

    Raises:
      KeyError: If required fields are missing from the request
    """
    try:
      if isinstance(self._query_config, TableFunctionQueryConfig):
        return [
            val for val in self._query_config.where_clause_value_fn(request)
        ]
      elif isinstance(self._query_config, TableFieldsQueryConfig):
        request_dict = request._asdict()
        return [
            request_dict[field]
            for field in self._query_config.where_clause_fields
        ]
      else:
        raise ValueError("Unsupported query configuration type")
    except KeyError as e:
      raise KeyError(
          "Make sure the values passed in `where_clause_fields` are "
          "the keys in the input `beam.Row`." + str(e))

  def _extract_parameter_names(self, template: str) -> list[str]:
    """Extract parameter names from a SQL template string.

    Args:
      template: SQL template string with named parameters (e.g., "id = :id")

    Returns:
      List of parameter names found in the template (e.g., ["id"])
    """
    return re.findall(r':(\w+)', template)

  def create_row_key(self, row: beam.Row):
    if isinstance(self._query_config, TableFunctionQueryConfig):
      return tuple(self._query_config.where_clause_value_fn(row))
    if isinstance(self._query_config, TableFieldsQueryConfig):
      row_dict = row._asdict()
      return (
          tuple(
              row_dict[where_clause_field]
              for where_clause_field in self._query_config.where_clause_fields))
    raise ValueError(
        "Either where_clause_fields or where_clause_value_fn must be specified")

  def get_cache_key(self, request: Union[beam.Row, list[beam.Row]]):
    if isinstance(self._query_config, CustomQueryConfig):
      raise NotImplementedError(
          "Caching is not supported for CustomQueryConfig. "
          "Consider using TableFieldsQueryConfig or " +
          "TableFunctionQueryConfig instead.")

    if isinstance(request, list):
      cache_keys = []
      for req in request:
        req_dict = req._asdict()
        try:
          if isinstance(self._query_config, TableFunctionQueryConfig):
            current_values = self._query_config.where_clause_value_fn(req)
          elif isinstance(self._query_config, TableFieldsQueryConfig):
            current_values = [
                req_dict[field]
                for field in self._query_config.where_clause_fields
            ]
          key = ';'.join(map(repr, current_values))
          cache_keys.append(key)
        except KeyError as e:
          raise KeyError(
              "Make sure the values passed in `where_clause_fields` are the "
              "keys in the input `beam.Row`." + str(e))
      return cache_keys
    else:
      req_dict = request._asdict()
      try:
        if isinstance(self._query_config, TableFunctionQueryConfig):
          current_values = self._query_config.where_clause_value_fn(request)
        else:  # TableFieldsQueryConfig.
          current_values = [
              req_dict[field]
              for field in self._query_config.where_clause_fields
          ]
        key = ";".join(["%s"] * len(current_values))
        cache_key = key % tuple(current_values)
      except KeyError as e:
        raise KeyError(
            "Make sure the values passed in `where_clause_fields` are the "
            "keys in the input `beam.Row`." + str(e))
      return cache_key

  def __exit__(self, _exc_type, _exc_val, _exc_tb):
    self._engine.dispose(close=True)
    self._engine = None

  def batch_elements_kwargs(self) -> Mapping[str, Any]:
    """Returns a kwargs suitable for `beam.BatchElements`."""
    return self._batching_kwargs