api-foundry-query-engine 0.8.81__tar.gz → 0.8.83__tar.gz

{api_foundry_query_engine-0.8.81 → api_foundry_query_engine-0.8.83}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: api-foundry-query-engine
-Version: 0.8.81
+Version: 0.8.83
 Summary: The AWS lambda service handler use by the `api_foundry` project is a powerful tool designed to automate the deployment of REST APIs on AWS using Lambda services to access and interact with relational databases (RDBMS). This project leverages the OpenAPI specification to define and manage the APIs
 Project-URL: Documentation, https://github.com/DanRepik/api-foundry
 Project-URL: Source, https://github.com/DanRepik/api-foundry
@@ -14,8 +14,10 @@ Requires-Dist: pyhumps
 Requires-Dist: pyyaml
 Provides-Extra: dev
 Requires-Dist: black>=23.0.0; extra == 'dev'
+Requires-Dist: docker; extra == 'dev'
 Requires-Dist: fixture-foundry>=0.1.0; extra == 'dev'
 Requires-Dist: isort>=5.12.0; extra == 'dev'
+Requires-Dist: psycopg2-binary; extra == 'dev'
 Requires-Dist: pulumi>=3.224.0; extra == 'dev'
 Requires-Dist: pytest-cov>=4.0.0; extra == 'dev'
 Requires-Dist: pytest>=7.0; extra == 'dev'

api_foundry_query_engine-0.8.83/api_foundry_query_engine/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.8.83"

api_foundry_query_engine-0.8.83/api_foundry_query_engine/connectors/connection.py

@@ -0,0 +1,106 @@
+import os
+from collections import defaultdict, deque
+from typing import List, Optional, Tuple
+
+from api_foundry_query_engine.utils.logger import logger
+
+# Initialize the logger
+log = logger(__name__)
+
+db_config_map = dict()
+
+# The driver a bare "engine" value (no ":driver" suffix) resolves to for
+# each dialect, when nothing else overrides it. This only names a driver --
+# it doesn't guarantee one is actually registered in ConnectionFactory's
+# CONNECTOR_REGISTRY yet (oracle/mysql have neither a real connector nor
+# tests exercising one, only dialect-specific SQL generation in
+# SQLQueryHandler.placeholder/concurrency_generator). Using ConnectionFactory
+# with such a dialect fails there instead, with a clear "unsupported
+# engine" error rather than a "no default driver" one.
+DEFAULT_DRIVERS = {
+    "postgres": "psycopg2",
+    "oracle": "cx_oracle",
+    "mysql": "mysqlclient",
+}
+
+
+def parse_engine(engine: str) -> Tuple[str, str]:
+    """Split an `engine` config value into (dialect, driver).
+
+    `engine` was originally meant to select a SQL dialect ("postgres",
+    "oracle", "mysql") for dialect-specific SQL generation (see
+    SQLQueryHandler.placeholder/concurrency_generator). Once a second
+    connector (the RDS Data API) existed for the same dialect, driver
+    selection got bolted onto the same string as a distinct value
+    ("postgres-data-api") instead of an orthogonal axis -- indistinguishable
+    from a dialect at a glance, and unable to express e.g. an Oracle Data
+    API-style driver without inventing another one-off name.
+
+    The explicit form is "{dialect}:{driver}", e.g. "postgres:data-api".
+    A bare dialect ("postgres") resolves its driver from the
+    "{DIALECT}_DEFAULT_DRIVER" environment variable if set, else
+    DEFAULT_DRIVERS -- so existing config/tests that only ever named a
+    dialect keep working, and a deployment can switch its default driver
+    (e.g. postgres -> data-api) without touching every config value that
+    names the dialect alone.
+    """
+    if ":" in engine:
+        dialect, driver = engine.split(":", 1)
+        return dialect, driver
+
+    dialect = engine
+    driver = os.environ.get(f"{dialect.upper()}_DEFAULT_DRIVER") or DEFAULT_DRIVERS.get(dialect)
+    if not driver:
+        raise ValueError(
+            f"No default driver configured for dialect {dialect!r} -- set "
+            f"{dialect.upper()}_DEFAULT_DRIVER or add it to DEFAULT_DRIVERS, "
+            f"or specify one explicitly as \"{dialect}:<driver>\"."
+        )
+    return dialect, driver
+
+
+def map_columns_to_selection_keys(column_names: List[str], selection_results: dict) -> List[Optional[str]]:
+    """Match a result set's bare column names, in cursor order, to
+    selection_results' keys (which may be table-alias-qualified, e.g.
+    "m.media_type_id" -- cursor/column metadata never returns qualified
+    names).
+
+    Returns one entry per column: selection_results' own key for that
+    column, or None if it isn't selected. A join can select the same bare
+    column from multiple tables (e.g. both "i.customer_id" and
+    "c.customer_id"), so this consumes one queued key per bare name per
+    occurrence -- in selection_results' insertion order, which matches the
+    SELECT list's column order -- rather than a plain name lookup, which
+    would collapse duplicates and silently drop a column's value.
+    """
+    queues: dict[str, deque] = defaultdict(deque)
+    for key in selection_results:
+        queues[key.rsplit(".", 1)[-1]].append(key)
+
+    return [queues[name].popleft() if queues.get(name) else None for name in column_names]
+
+
+class Cursor:
+    def execute(self, sql: str, params: dict, selection_results: dict) -> list[dict]:
+        raise NotImplementedError
+
+    def close(self):
+        raise NotImplementedError
+
+
+class Connection:
+    def __init__(self, db_config: dict) -> None:
+        super().__init__()
+        self.db_config = db_config
+
+    def engine(self) -> str:
+        return self.db_config["engine"]
+
+    def cursor(self) -> Cursor:
+        raise NotImplementedError
+
+    def commit(self):
+        raise NotImplementedError
+
+    def close(self):
+        raise NotImplementedError

api_foundry_query_engine-0.8.83/api_foundry_query_engine/connectors/connection_factory.py

@@ -0,0 +1,169 @@
+from typing import Callable, Mapping, Tuple, Type
+import boto3
+import json
+
+from api_foundry_query_engine.connectors.connection import Connection, parse_engine
+from api_foundry_query_engine.utils.app_exception import ApplicationException
+from api_foundry_query_engine.utils.logger import logger
+
+log = logger(__name__)
+
+
+def _load_postgres_connection() -> Type[Connection]:
+    from api_foundry_query_engine.connectors.postgres_connection import PostgresConnection
+
+    return PostgresConnection
+
+
+def _load_data_api_connection() -> Type[Connection]:
+    from api_foundry_query_engine.connectors.data_api_connection import DataApiConnection
+
+    return DataApiConnection
+
+
+# (dialect, driver) -> a loader returning the Connection subclass to
+# instantiate. A function rather than the class itself so each connector's
+# module (and its own dependencies, e.g. psycopg2) is only imported once a
+# config value actually selects it. Add new dialect/driver combinations
+# here rather than growing an if/elif chain.
+CONNECTOR_REGISTRY: dict[Tuple[str, str], Callable[[], Type[Connection]]] = {
+    ("postgres", "psycopg2"): _load_postgres_connection,
+    ("postgres", "data-api"): _load_data_api_connection,
+}
+
+
+class ConnectionFactory:
+    db_config_map: dict[str, dict]
+    config: Mapping[str, str]
+
+    def __init__(self, config: Mapping[str, str] = {}):
+        self.db_config_map = dict()
+        self.config = config
+
+    def get_connection(self, database: str) -> Connection:
+        """
+        Factory function to create a database connector based on the
+        specified engine and schema.
+
+        Args:
+        - engine (str): The database engine type
+                ('postgres', 'oracle', or 'mysql').
+        - schema (str): The schema for the database.
+
+        Returns:
+        - Connector: An instance of the appropriate Connector subclass.
+        """
+
+        # Get the secret name based on the engine and database from the secrets map
+        log.info("database: %s", database)
+        db_config = self.db_config_map.get(database)
+        if not db_config:
+            # A <DATABASE>_DATA_API_CLUSTER_ARN config value selects the
+            # RDS Data API connector and, like the DSN branch below, never
+            # fetches the secret itself -- the Data API resolves the
+            # credentials server-side from secret_arn, so no password is
+            # ever in this Lambda's environment or memory. The Data API
+            # does that lookup with the *caller's* IAM permissions, though,
+            # so the role still needs secretsmanager:GetSecretValue on
+            # secret_arn alongside the rds-data:* actions.
+            data_api_cluster_arn = self.config.get(f"{database.upper()}_DATA_API_CLUSTER_ARN")
+            if data_api_cluster_arn:
+                db_config = {
+                    "engine": "postgres:data-api",
+                    "resource_arn": data_api_cluster_arn,
+                    "secret_arn": self.config.get(f"{database.upper()}_DATA_API_SECRET_ARN"),
+                    "database": self.config.get(f"{database.upper()}_DATA_API_DATABASE"),
+                }
+                schema = self.config.get(f"{database.upper()}_DATA_API_SCHEMA")
+                if schema:
+                    db_config["schema"] = schema
+                endpoint_url = self.config.get("AWS_ENDPOINT_URL")  # LocalStack endpoint
+                if endpoint_url:
+                    db_config["endpoint_url"] = endpoint_url
+            else:
+                # A <DATABASE>_DSN config value (e.g. CONTRACT_DB_DSN) skips
+                # Secrets Manager entirely. This matters for a Lambda that's
+                # VPC-attached for private database access but has no route
+                # to the Secrets Manager API (no NAT gateway, no interface
+                # endpoint) -- the SECRETS/get_secret_value path would just
+                # hang. PostgresConnection already accepts a raw "dsn" key.
+                dsn = self.config.get(f"{database.upper()}_DSN")
+                if dsn:
+                    db_config = {"engine": self.config.get(f"{database.upper()}_ENGINE", "postgres"), "dsn": dsn}
+                else:
+                    # Use config dict for secrets
+                    secrets_map = self.config.get("SECRETS", {})
+                    if isinstance(secrets_map, str):
+                        secrets_map = json.loads(secrets_map)
+                    secret_name = secrets_map.get(database)
+                    log.debug("secret_name: %s", secret_name)
+
+                    if secret_name:
+                        db_config = self.__get_secret(secret_name)
+                    else:
+                        raise ValueError(f"Secret not found for database: {database}")
+
+        engine = db_config.get("engine")
+        if not engine:
+            raise ApplicationException(500, "Database 'engine' is not defined in the secret.")
+
+        dialect, driver = parse_engine(engine)
+        loader = CONNECTOR_REGISTRY.get((dialect, driver))
+        if loader is None:
+            raise ValueError(
+                f"Unsupported database engine: {engine!r} (dialect={dialect!r}, driver={driver!r}). "
+                f"Registered: {sorted(CONNECTOR_REGISTRY.keys())}"
+            )
+        return loader()(db_config)
+
+    def __get_secret(self, db_secret_name: str):
+        """
+        Get the secret from AWS Secrets Manager.
+
+        Parameters:
+        - db_secret_name (str): The name of the AWS Secrets Manager secret.
+
+        Returns:
+        - dict: The database configuration obtained from the secret.
+        """
+        if self.config.get(db_secret_name):
+            return self.config.get(db_secret_name)
+
+        endpoint_url = self.config.get("AWS_ENDPOINT_URL")  # LocalStack endpoint
+        sts_client = boto3.client("sts", endpoint_url=endpoint_url)
+
+        secret_account_id = self.config.get("SECRET_ACCOUNT_ID", None)
+        log.debug("secret_account_id: %s", secret_account_id)
+
+        if secret_account_id:
+            # If a secret account ID is provided, assume a role in that account
+            secret_role = self.config.get("ROLE_NAME", None)
+            assume_role_response = sts_client.assume_role(
+                RoleArn=f"arn:aws:iam::{secret_account_id}:role/{secret_role}",
+                RoleSessionName="AssumeRoleSession",
+            )
+
+            credentials = assume_role_response["Credentials"]
+
+            secretsmanager = boto3.client(
+                "secretsmanager",
+                aws_access_key_id=credentials["AccessKeyId"],
+                aws_secret_access_key=credentials["SecretAccessKey"],
+                aws_session_token=credentials["SessionToken"],
+                endpoint_url=endpoint_url,
+            )
+        else:
+            # If no secret account ID is provided, use the default account
+            log.info("endpoint_url: %s", endpoint_url)
+            secretsmanager = boto3.client(
+                "secretsmanager",
+                endpoint_url=endpoint_url,
+            )
+
+        # Get the secret value from AWS Secrets Manager
+        log.info("db_secret_name: %s", db_secret_name)
+        db_secret = secretsmanager.get_secret_value(SecretId=db_secret_name)
+        log.debug("loading secret name: %s", db_secret)
+
+        # Return the parsed JSON secret string
+        return json.loads(db_secret.get("SecretString"))

api_foundry_query_engine-0.8.83/api_foundry_query_engine/connectors/data_api_connection.py

@@ -0,0 +1,329 @@
+import json
+import re
+import time
+from datetime import date, datetime, time as time_type, timezone
+from decimal import Decimal
+from typing import Any, Optional
+
+import boto3
+from botocore.config import Config as BotoConfig
+
+from api_foundry_query_engine.connectors.connection import (
+    Connection,
+    Cursor,
+    map_columns_to_selection_keys,
+)
+from api_foundry_query_engine.utils.app_exception import ApplicationException
+from api_foundry_query_engine.utils.logger import logger
+
+log = logger(__name__)
+
+# The Data API's own statement timeout is 45s (AWS docs); a boto3
+# read_timeout below that would surface as a generic SDK timeout instead
+# of the service's own, more useful error, so stay comfortably above it.
+READ_TIMEOUT_SECONDS = 65
+
+# A paused Aurora Serverless v2 cluster resumes in well under this budget
+# (AWS docs put it at ~15s), but the caller (API Gateway) has a hard 29s
+# cap end to end, so don't retry for longer than that leaves room for.
+RESUME_RETRY_BUDGET_SECONDS = 20
+RESUME_RETRY_INTERVAL_SECONDS = 2
+
+_WRITE_KEYWORDS = ("INSERT", "UPDATE", "DELETE")
+
+
+def _is_write_statement(sql: str) -> bool:
+    stripped = sql.strip()
+    if not stripped:
+        return False
+    return stripped.split(None, 1)[0].upper() in _WRITE_KEYWORDS
+
+
+def _encode_parameter(value: Any) -> dict:
+    """Encode a Python value into an RDS Data API parameter `value` field.
+
+    Values arriving here have already been through
+    SchemaObjectProperty.convert_to_db_value (or json.dumps, for embedded
+    objects/arrays), so this only needs to get each Python type onto the
+    wire faithfully -- columns a plain string can't be implicitly coerced
+    into (uuid, jsonb, timestamptz, ...) are handled by an explicit SQL
+    cast on the placeholder instead (see SQLQueryHandler.placeholder).
+    """
+    if value is None:
+        return {"isNull": True}
+    if isinstance(value, bool):
+        return {"booleanValue": value}
+    if isinstance(value, int):
+        return {"longValue": value}
+    if isinstance(value, float):
+        return {"doubleValue": value}
+    if isinstance(value, Decimal):
+        return {"stringValue": str(value)}
+    if isinstance(value, datetime):
+        return {"stringValue": value.strftime("%Y-%m-%d %H:%M:%S.%f")}
+    if isinstance(value, date):
+        return {"stringValue": value.isoformat()}
+    if isinstance(value, time_type):
+        return {"stringValue": value.strftime("%H:%M:%S.%f")}
+    if isinstance(value, (dict, list)):
+        # Not expected on this path today (object/array properties are
+        # already JSON-text by the time they reach here), but dumping to
+        # JSON is the only non-lossy fallback if one ever does.
+        return {"stringValue": json.dumps(value)}
+    return {"stringValue": str(value)}
+
+
+def _build_parameters(params: dict) -> list:
+    return [{"name": name, "value": _encode_parameter(value)} for name, value in (params or {}).items()]
+
+
+def _parse_data_api_timestamp(value: str) -> datetime:
+    for fmt in ("%Y-%m-%d %H:%M:%S.%f", "%Y-%m-%d %H:%M:%S"):
+        try:
+            return datetime.strptime(value, fmt)
+        except ValueError:
+            continue
+    raise ApplicationException(500, f"Unrecognized Data API timestamp value: {value!r}")
+
+
+def _parse_data_api_time(value: str) -> time_type:
+    for fmt in ("%H:%M:%S.%f", "%H:%M:%S"):
+        try:
+            return datetime.strptime(value, fmt).time()
+        except ValueError:
+            continue
+    raise ApplicationException(500, f"Unrecognized Data API time value: {value!r}")
+
+
+def _decode_array(array_value: dict) -> list:
+    for key in ("stringValues", "longValues", "doubleValues", "booleanValues"):
+        values = array_value.get(key)
+        if values is not None:
+            return list(values)
+    nested = array_value.get("arrayValues")
+    if nested is not None:
+        return [_decode_array(v) for v in nested]
+    return []
+
+
+def _rehydrate_string(value: str, type_name: Optional[str]) -> Any:
+    """Turn a Data API stringValue back into the Python type psycopg2
+    would have produced for the same column, so marshal_record and
+    SchemaObjectProperty.convert_to_api_value behave identically
+    regardless of which connector ran the query. Without this, e.g. a
+    jsonb column would come back as an escaped string instead of a dict,
+    and a timestamptz would lose its offset.
+    """
+    if type_name in ("json", "jsonb"):
+        return json.loads(value)
+    if type_name == "timestamptz":
+        return _parse_data_api_timestamp(value).replace(tzinfo=timezone.utc)
+    if type_name == "timestamp":
+        return _parse_data_api_timestamp(value)
+    if type_name == "date":
+        return date.fromisoformat(value)
+    if type_name in ("time", "timetz"):
+        return _parse_data_api_time(value)
+    return value
+
+
+def _decode_field(field: dict, type_name: Optional[str]) -> Any:
+    if field.get("isNull"):
+        return None
+    if "booleanValue" in field:
+        return field["booleanValue"]
+    if "longValue" in field:
+        return field["longValue"]
+    if "doubleValue" in field:
+        return field["doubleValue"]
+    if "blobValue" in field:
+        return field["blobValue"]
+    if "arrayValue" in field:
+        return _decode_array(field["arrayValue"])
+    value = field.get("stringValue")
+    if value is None:
+        return None
+    return _rehydrate_string(value, type_name)
+
+
+def _sqlstate_class(message: str) -> str:
+    match = re.search(r"SQLState:\s*([0-9A-Za-z]+)", message or "")
+    return match.group(1)[:2] if match else ""
+
+
+def _map_database_error(message: str) -> ApplicationException:
+    # Mirrors PostgresCursor's IntegrityError/ProgrammingError/Error
+    # mapping, driven off the SQLState class Postgres appends to the
+    # DatabaseErrorException message instead of a driver-specific type.
+    sqlstate_class = _sqlstate_class(message)
+    if sqlstate_class == "23":
+        return ApplicationException(409, message)
+    if sqlstate_class == "42":
+        return ApplicationException(400, message)
+    return ApplicationException(500, message)
+
+
+class DataApiCursor(Cursor):
+    def __init__(self, connection: "DataApiConnection"):
+        self.__connection = connection
+
+    def execute(self, sql: str, params: dict, selection_results: dict) -> list[dict]:
+        log.info("sql: %s", sql)
+
+        transaction_id = self.__connection.begin_if_needed(sql)
+        response = self.__connection.execute_statement(sql, _build_parameters(params), transaction_id)
+
+        records = response.get("records")
+        if not records:
+            return []
+
+        column_metadata = response.get("columnMetadata") or []
+        columns = [(meta.get("name"), meta.get("typeName")) for meta in column_metadata]
+        # column_metadata's "name" is always bare, even when
+        # selection_results' keys are table-alias-qualified (e.g.
+        # "m.media_type_id") -- see postgres_connection.py's PostgresCursor
+        # for the same mapping against psycopg2's cursor.description.
+        position_keys = map_columns_to_selection_keys([name for name, _ in columns], selection_results)
+
+        result = []
+        for record in records:
+            row = {}
+            for key, (_, type_name), field in zip(position_keys, columns, record):
+                if key is not None:
+                    row[key] = _decode_field(field, type_name)
+            result.append(row)
+        return result
+
+    def close(self):
+        # Each execute() is a single stateless HTTPS call; there is no
+        # per-cursor resource to release. The connection, not the cursor,
+        # owns the transaction lifecycle -- see DataApiConnection.close().
+        pass
+
+
+class DataApiConnection(Connection):
+    """RDS Data API connection.
+
+    Lets the query engine reach a private, PubliclyAccessible=False Aurora
+    cluster from a Lambda that is *not* VPC-attached: every call is a
+    plain HTTPS request to the `rds-data` API, authenticated by IAM, with
+    the database credentials looked up server-side from the secret named
+    by `secret_arn`. Unlike the VPC-attached DSN-in-env-var workaround
+    this replaces, the password never reaches the Lambda's environment or
+    memory. The Data API reads the secret with the caller's own IAM
+    permissions, so the role needs `secretsmanager:GetSecretValue` on
+    `secret_arn` as well as the `rds-data` actions (ExecuteStatement,
+    BeginTransaction, CommitTransaction, RollbackTransaction) on the
+    cluster.
+
+    db_config keys: `resource_arn` (the cluster ARN), `secret_arn`,
+    `database`, and optionally `schema` and `endpoint_url` (LocalStack).
+    """
+
+    def __init__(self, db_config: dict) -> None:
+        super().__init__(db_config)
+        self.resource_arn = db_config["resource_arn"]
+        self.secret_arn = db_config["secret_arn"]
+        self.database = db_config["database"]
+        self.schema = db_config.get("schema")
+        self.transaction_id: Optional[str] = None
+        self.client = boto3.client(
+            "rds-data",
+            endpoint_url=db_config.get("endpoint_url"),
+            config=BotoConfig(read_timeout=READ_TIMEOUT_SECONDS, retries={"max_attempts": 1}),
+        )
+
+    def cursor(self) -> Cursor:
+        return DataApiCursor(self)
+
+    def begin_if_needed(self, sql: str) -> Optional[str]:
+        """Begin a transaction on the first write seen on this connection,
+        and reuse it for every statement (read or write) after that, so a
+        read that depends on an earlier write in the same request or
+        batch sees it. A pure read sequence never begins one, avoiding
+        the extra begin/commit round trips for the common read-only path.
+        """
+        if self.transaction_id is not None:
+            return self.transaction_id
+        if not _is_write_statement(sql):
+            return None
+
+        kwargs = {"resourceArn": self.resource_arn, "secretArn": self.secret_arn, "database": self.database}
+        if self.schema:
+            kwargs["schema"] = self.schema
+        response = self.client.begin_transaction(**kwargs)
+        self.transaction_id = response["transactionId"]
+        return self.transaction_id
+
+    def execute_statement(self, sql: str, parameters: list, transaction_id: Optional[str]) -> dict:
+        kwargs = {
+            "resourceArn": self.resource_arn,
+            "secretArn": self.secret_arn,
+            "database": self.database,
+            "sql": sql,
+            "parameters": parameters,
+            "includeResultMetadata": True,
+        }
+        if self.schema:
+            kwargs["schema"] = self.schema
+        if transaction_id:
+            kwargs["transactionId"] = transaction_id
+
+        deadline = time.monotonic() + RESUME_RETRY_BUDGET_SECONDS
+        while True:
+            try:
+                return self.client.execute_statement(**kwargs)
+            except self.client.exceptions.DatabaseResumingException:
+                if time.monotonic() >= deadline:
+                    raise ApplicationException(
+                        503,
+                        "Database is resuming from auto-pause; retry the request shortly.",
+                    )
+                log.info("database resuming, retrying in %ss", RESUME_RETRY_INTERVAL_SECONDS)
+                time.sleep(RESUME_RETRY_INTERVAL_SECONDS)
+            except self.client.exceptions.UnsupportedResultException as err:
+                raise ApplicationException(
+                    400,
+                    "Result exceeds the RDS Data API's response size limit; " "narrow the query or add pagination.",
+                ) from err
+            except self.client.exceptions.DatabaseErrorException as err:
+                raise _map_database_error(_error_message(err)) from err
+            except self.client.exceptions.BadRequestException as err:
+                raise ApplicationException(400, _error_message(err)) from err
+            except self.client.exceptions.TransactionNotFoundException as err:
+                # The transaction's idle timeout elapsed server-side; the
+                # connection can't keep using it.
+                self.transaction_id = None
+                raise ApplicationException(500, "Database transaction expired before the request completed.") from err
+            except self.client.exceptions.StatementTimeoutException as err:
+                raise ApplicationException(504, _error_message(err)) from err
+
+    def commit(self):
+        if self.transaction_id is None:
+            return
+        self.client.commit_transaction(
+            resourceArn=self.resource_arn, secretArn=self.secret_arn, transactionId=self.transaction_id
+        )
+        self.transaction_id = None
+
+    def rollback(self):
+        if self.transaction_id is None:
+            return
+        self.client.rollback_transaction(
+            resourceArn=self.resource_arn, secretArn=self.secret_arn, transactionId=self.transaction_id
+        )
+        self.transaction_id = None
+
+    def close(self):
+        # Mirrors psycopg2: closing a connection with an open, uncommitted
+        # transaction implicitly rolls it back, rather than leaving a Data
+        # API transaction dangling until its own idle timeout.
+        if self.transaction_id is not None:
+            try:
+                self.rollback()
+            except Exception:
+                log.error("failed to roll back open Data API transaction on close")
+
+
+def _error_message(err) -> str:
+    return err.response.get("Error", {}).get("Message", str(err))

{api_foundry_query_engine-0.8.81 → api_foundry_query_engine-0.8.83}/api_foundry_query_engine/connectors/postgres_connection.py

@@ -1,4 +1,8 @@
-from api_foundry_query_engine.connectors.connection import Connection, Cursor
+from api_foundry_query_engine.connectors.connection import (
+    Connection,
+    Cursor,
+    map_columns_to_selection_keys,
+)
 from api_foundry_query_engine.utils.logger import logger
 
 # Initialize the logger
@@ -32,16 +36,16 @@ class PostgresCursor(Cursor):
         try:
             # Execute the SQL statement with parameters
             self.__cursor.execute(sql, params)
-            selected_columns = set(selection_results)
+            # cursor.description always returns bare column names, even
+            # when selection_results' keys are table-alias-qualified (e.g.
+            # "m.media_type_id"). Match on the bare name but keep
+            # selection_results' own key (qualified or not) in the output,
+            # since downstream code looks records up by that exact key.
             column_names = [desc[0] for desc in (self.__cursor.description or ())]
+            position_keys = map_columns_to_selection_keys(column_names, selection_results)
             result = []
             for record in self.__cursor:
-                # Use cursor metadata so SQL aliases define the output mapping.
-                row = {
-                    col: value
-                    for col, value in zip(column_names, record)
-                    if col in selected_columns
-                }
+                row = {key: value for key, value in zip(position_keys, record) if key is not None}
                 result.append(row)
 
             return result

{api_foundry_query_engine-0.8.81 → api_foundry_query_engine-0.8.83}/api_foundry_query_engine/dao/sql_query_handler.py

@@ -4,6 +4,7 @@ import uuid
 from typing import Optional, List, Dict
 from datetime import datetime, date
 
+from api_foundry_query_engine.connectors.connection import parse_engine
 from api_foundry_query_engine.utils.app_exception import ApplicationException
 from api_foundry_query_engine.operation import Operation
 from api_foundry_query_engine.utils.api_model import SchemaObject, SchemaObjectProperty
@@ -108,14 +109,35 @@ RELATIONAL_TYPES = {
     "not-like": "not-like",
 }
 
+# The RDS Data API sends every parameter as a plain typed value (string,
+# long, double, boolean, null) with no column context, so Postgres can't
+# implicitly coerce it the way it does a psycopg2 %s literal -- comparing
+# an untyped string against a uuid/jsonb/timestamptz column fails with
+# "operator does not exist". Casting the placeholder itself (`:x::uuid`)
+# fixes this under the Data API and is a no-op under psycopg2, which
+# never sees this branch. Keyed on SchemaObjectProperty.column_type.
+DATA_API_COLUMN_CASTS = {
+    "uuid": "uuid",
+    "date": "date",
+    "time": "time",
+    "timetz": "timetz",
+    "date-time": "timestamptz",
+    "datetime": "timestamptz",
+    "timestamp": "timestamptz",
+    "timestamptz": "timestamptz",
+}
+
 
 class SQLQueryHandler:
     operation: Operation
     engine: str
+    dialect: str
+    driver: str
 
     def __init__(self, operation: Operation, engine: str):  # , schema_object: SchemaObject):
         self.operation = operation
         self.engine = engine
+        self.dialect, self.driver = parse_engine(engine)
         self.__select_list_columns = None
 
     @property
@@ -155,7 +177,7 @@ class SQLQueryHandler:
         if len(param) == 0:
             param = property.api_name if property.api_name is not None else ""
 
-        if self.engine == "oracle":
+        if self.dialect == "oracle":
             if property.column_type == "date":
                 return f"TO_DATE(:{param}, 'YYYY-MM-DD')"
             elif property.column_type == "datetime":
@@ -163,6 +185,16 @@ class SQLQueryHandler:
             elif property.column_type == "time":
                 return f"TO_TIME(:{param}, 'HH24:MI:SS.FF')"
             return f":{param}"
+        if self.driver == "data-api":
+            # JSON-serialized properties (embedded objects/arrays, see
+            # SQLInsertSchemaQueryHandler.insert_values) always need a
+            # jsonb cast regardless of the declared column_type, since
+            # they arrive as a JSON-text string rather than going through
+            # convert_to_db_value.
+            if property.api_type in ("object", "array"):
+                return f":{param}::jsonb"
+            cast = DATA_API_COLUMN_CASTS.get(property.column_type)
+            return f":{param}::{cast}" if cast else f":{param}"
         return f"%({param})s"
 
     def check_permissions(
@@ -766,9 +798,9 @@ class SQLSchemaQueryHandler(SQLQueryHandler):
         elif property.api_type == "integer":
             return f"{property.column_name} + 1"
         elif property.api_type in ["string", "uuid"]:
-            if self.engine == "oracle":
+            if self.dialect == "oracle":
                 return "SYS_GUID()"
-            if self.engine == "mysql":
+            if self.dialect == "mysql":
                 return "UUID()"
             return "gen_random_uuid()"
         raise ApplicationException(

{api_foundry_query_engine-0.8.81 → api_foundry_query_engine-0.8.83}/api_foundry_query_engine/dao/sql_select_query_handler.py

@@ -281,21 +281,10 @@ class SQLSelectSchemaQueryHandler(SQLSchemaQueryHandler):
                     schema_object.permissions,
                     schema_object.properties,
                 )
-                # Only genuine joined relations need their columns
-                # qualified (e.g. "customer.name", to disambiguate from
-                # the parent's own "name") -- the primary entity's own
-                # properties come back from the cursor under their bare
-                # column names, and marshal_record/the post-fetch column
-                # filter in PostgresCursor.execute both key off that.
-                # Prefixing them too (self.prefix_map[relation], as
-                # before) produced selection_results keys like "q.slug"
-                # that never matched any real result column, so every
-                # row came back empty and marshal_record KeyError'd
-                # trying to unwrap it.
                 filtered_keys = self.filter_and_prefix_keys(
                     reg_exs,
                     allowed_properties,
-                    self.prefix_map[relation] if relation_property else None,
+                    self.prefix_map[relation],
                 )
 
                 # Extend the result map with the filtered keys

{api_foundry_query_engine-0.8.81 → api_foundry_query_engine-0.8.83}/api_foundry_query_engine/lambda_handler.py

@@ -13,21 +13,54 @@ logging.basicConfig(level=os.getenv("LOG_LEVEL", "INFO"))
 log = logging.getLogger(__name__)
 
 
+# Lambda rejects synchronous responses over 6 MB (6,291,556 bytes including
+# its own envelope); API Gateway then surfaces that as an opaque
+# "502 Internal server error". The margin leaves room for the statusCode/
+# headers wrapper so a body under the limit can't still be refused.
+DEFAULT_MAX_RESPONSE_BYTES = 6 * 1024 * 1024 - 4096
+
+
 class QueryEngine:
     def __init__(self, config: Mapping[str, str]):
         self.adapter = GatewayAdapter(config)
+        self.max_response_bytes = int(
+            config.get("MAX_RESPONSE_BYTES") or DEFAULT_MAX_RESPONSE_BYTES
+        )
 
     def handler(self, event) -> dict[str, Any]:
         log.debug("event: %s", event)
         try:
             response = self.adapter.process_event(event)
+            body = json.dumps(response)
+
+            # json.dumps escapes to ASCII by default, so len(body) is bytes.
+            if len(body) > self.max_response_bytes:
+                log.error(
+                    "response of %d bytes exceeds limit of %d bytes",
+                    len(body),
+                    self.max_response_bytes,
+                )
+                return {
+                    "isBase64Encoded": False,
+                    "statusCode": 502,
+                    "headers": {"Content-Type": "application/json"},
+                    "body": json.dumps(
+                        {
+                            "message": (
+                                f"response of {len(body)} bytes exceeds the "
+                                f"maximum of {self.max_response_bytes} bytes; "
+                                "narrow the request with filters"
+                            )
+                        }
+                    ),
+                }
 
             # Ensure the response conforms to API Gateway requirements
             return {
                 "isBase64Encoded": False,
                 "statusCode": 200,
                 "headers": {"Content-Type": "application/json"},
-                "body": json.dumps(response),
+                "body": body,
             }
         except ApplicationException as e:
             log.error("exception: %s", e, exc_info=True)

{api_foundry_query_engine-0.8.81 → api_foundry_query_engine-0.8.83}/pyproject.toml

@@ -1,5 +1,10 @@
 [build-system]
-requires = ["hatchling>=1.0"]
+# hatchling>=1.32.0 emits Metadata-Version: 2.5, which PyPI's Warehouse
+# rejects outright with a 400 on every upload regardless of package
+# version -- see https://github.com/pypi/warehouse/issues/19083 (open,
+# unresolved). Pinned below that until Warehouse adds 2.5 support (same
+# fix applied to cloud_foundry, see DanRepik/cloud_foundry#10).
+requires = ["hatchling>=1.0,<1.32.0"]
 build-backend = "hatchling.build"
 
 [project]
@@ -28,6 +33,14 @@ dev = [
     "fixture_foundry>=0.1.0",
     "simple-oauth-server",
     "pulumi>=3.224.0",
+    # tests/conftest.py imports this directly; was missing here entirely
+    # until CI first ran the test suite and caught it (same gap found in
+    # api_foundry, which uses psycopg2-binary the same way).
+    "psycopg2-binary",
+    # fixture_foundry imports this at module load time (fixtures.py);
+    # missing here caused an ImportError from conftest.py's own
+    # `from fixture_foundry import deploy` before any test could run.
+    "docker",
 ]
 
 [project.urls]

api_foundry_query_engine-0.8.81/api_foundry_query_engine/__init__.py

@@ -1 +0,0 @@
-__version__ = "0.8.81"

api_foundry_query_engine-0.8.81/api_foundry_query_engine/connectors/connection.py

@@ -1,32 +0,0 @@
-from api_foundry_query_engine.utils.logger import logger
-
-# Initialize the logger
-log = logger(__name__)
-
-db_config_map = dict()
-
-
-class Cursor:
-    def execute(self, sql: str, params: dict, selection_results: dict) -> list[dict]:
-        raise NotImplementedError
-
-    def close(self):
-        raise NotImplementedError
-
-
-class Connection:
-    def __init__(self, db_config: dict) -> None:
-        super().__init__()
-        self.db_config = db_config
-
-    def engine(self) -> str:
-        return self.db_config["engine"]
-
-    def cursor(self) -> Cursor:
-        raise NotImplementedError
-
-    def commit(self):
-        raise NotImplementedError
-
-    def close(self):
-        raise NotImplementedError

api_foundry_query_engine-0.8.81/api_foundry_query_engine/connectors/connection_factory.py

@@ -1,123 +0,0 @@
-from typing import Mapping
-import boto3
-import json
-
-from api_foundry_query_engine.connectors.connection import Connection
-from api_foundry_query_engine.utils.app_exception import ApplicationException
-from api_foundry_query_engine.utils.logger import logger
-
-log = logger(__name__)
-
-
-class ConnectionFactory:
-    db_config_map: dict[str, dict]
-    config: Mapping[str, str]
-
-    def __init__(self, config: Mapping[str, str] = {}):
-        self.db_config_map = dict()
-        self.config = config
-
-    def get_connection(self, database: str) -> Connection:
-        """
-        Factory function to create a database connector based on the
-        specified engine and schema.
-
-        Args:
-        - engine (str): The database engine type
-                ('postgres', 'oracle', or 'mysql').
-        - schema (str): The schema for the database.
-
-        Returns:
-        - Connector: An instance of the appropriate Connector subclass.
-        """
-
-        # Get the secret name based on the engine and database from the secrets map
-        log.info("database: %s", database)
-        db_config = self.db_config_map.get(database)
-        if not db_config:
-            # A <DATABASE>_DSN config value (e.g. CONTRACT_DB_DSN) skips
-            # Secrets Manager entirely. This matters for a Lambda that's
-            # VPC-attached for private database access but has no route
-            # to the Secrets Manager API (no NAT gateway, no interface
-            # endpoint) -- the SECRETS/get_secret_value path would just
-            # hang. PostgresConnection already accepts a raw "dsn" key.
-            dsn = self.config.get(f"{database.upper()}_DSN")
-            if dsn:
-                db_config = {"engine": self.config.get(f"{database.upper()}_ENGINE", "postgres"), "dsn": dsn}
-            else:
-                # Use config dict for secrets
-                secrets_map = self.config.get("SECRETS", {})
-                if isinstance(secrets_map, str):
-                    secrets_map = json.loads(secrets_map)
-                secret_name = secrets_map.get(database)
-                log.debug("secret_name: %s", secret_name)
-
-                if secret_name:
-                    db_config = self.__get_secret(secret_name)
-                else:
-                    raise ValueError(f"Secret not found for database: {database}")
-
-        engine = db_config.get("engine")
-        if not engine:
-            raise ApplicationException(500, "Database 'engine' is not defined in the secret.")
-
-        if engine == "postgres":
-            from .postgres_connection import PostgresConnection
-
-            return PostgresConnection(db_config)
-
-        # Add support for other engines here if needed in the future
-
-        raise ValueError(f"Unsupported database engine: {engine}")
-
-    def __get_secret(self, db_secret_name: str):
-        """
-        Get the secret from AWS Secrets Manager.
-
-        Parameters:
-        - db_secret_name (str): The name of the AWS Secrets Manager secret.
-
-        Returns:
-        - dict: The database configuration obtained from the secret.
-        """
-        if self.config.get(db_secret_name):
-            return self.config.get(db_secret_name)
-
-        endpoint_url = self.config.get("AWS_ENDPOINT_URL")  # LocalStack endpoint
-        sts_client = boto3.client("sts", endpoint_url=endpoint_url)
-
-        secret_account_id = self.config.get("SECRET_ACCOUNT_ID", None)
-        log.debug("secret_account_id: %s", secret_account_id)
-
-        if secret_account_id:
-            # If a secret account ID is provided, assume a role in that account
-            secret_role = self.config.get("ROLE_NAME", None)
-            assume_role_response = sts_client.assume_role(
-                RoleArn=f"arn:aws:iam::{secret_account_id}:role/{secret_role}",
-                RoleSessionName="AssumeRoleSession",
-            )
-
-            credentials = assume_role_response["Credentials"]
-
-            secretsmanager = boto3.client(
-                "secretsmanager",
-                aws_access_key_id=credentials["AccessKeyId"],
-                aws_secret_access_key=credentials["SecretAccessKey"],
-                aws_session_token=credentials["SessionToken"],
-                endpoint_url=endpoint_url,
-            )
-        else:
-            # If no secret account ID is provided, use the default account
-            log.info("endpoint_url: %s", endpoint_url)
-            secretsmanager = boto3.client(
-                "secretsmanager",
-                endpoint_url=endpoint_url,
-            )
-
-        # Get the secret value from AWS Secrets Manager
-        log.info("db_secret_name: %s", db_secret_name)
-        db_secret = secretsmanager.get_secret_value(SecretId=db_secret_name)
-        log.debug("loading secret name: %s", db_secret)
-
-        # Return the parsed JSON secret string
-        return json.loads(db_secret.get("SecretString"))