api-foundry-query-engine 0.8.39__py3-none-any.whl

api_foundry_query_engine/connectors/postgres_connection.py

@@ -0,0 +1,142 @@
+from api_foundry_query_engine.connectors.connection import Connection, Cursor
+from api_foundry_query_engine.utils.logger import logger
+
+# Initialize the logger
+log = logger(__name__)
+
+
+class PostgresCursor(Cursor):
+    def __init__(self, cursor):
+        self.__cursor = cursor
+
+    def execute(self, sql: str, params: dict, selection_results: dict) -> list[dict]:
+        """
+        Execute SQL statements on the PostgreSQL database.
+
+        Parameters:
+        - cursor: The database cursor.
+        - sql (str): The SQL statement to execute.
+        - params (dict): Parameters to be used in the SQL statement.
+        - selection_results (dict): Mapping of result columns.
+
+        Returns:
+        - list[dict]: List of result records as dictionaries.
+
+        Raises:
+        - AppException: Custom exception for handling database-related errors.
+        """
+        from psycopg2 import Error, IntegrityError, ProgrammingError
+
+        log.info("sql: %s", sql)
+
+        try:
+            # Execute the SQL statement with parameters
+            self.__cursor.execute(sql, params)
+            result = []
+            for record in self.__cursor:
+                # Convert record tuple to dictionary using selection_results
+                result.append(
+                    {col: value for col, value in zip(selection_results, record)}
+                )
+
+            return result
+        except IntegrityError as err:
+            # Handle integrity constraint violation (e.g., duplicate key)
+            from api_foundry_query_engine.utils.app_exception import (
+                ApplicationException,
+            )
+
+            raise ApplicationException(409, err.pgerror)
+        except ProgrammingError as err:
+            # Handle programming errors (e.g., syntax error in SQL)
+            from api_foundry_query_engine.utils.app_exception import (
+                ApplicationException,
+            )
+
+            raise ApplicationException(400, err.pgerror)
+        except Error as err:
+            # Handle other database errors
+            from api_foundry_query_engine.utils.app_exception import (
+                ApplicationException,
+            )
+
+            raise ApplicationException(500, err.pgerror)
+
+    def close(self):
+        self.__cursor.close()
+
+
+class PostgresConnection(Connection):
+    """
+    PostgreSQL database connection wrapper.
+
+    Supports two configuration formats:
+    1. DSN-based (preferred for testing with fixture_foundry):
+       {"dsn": "postgresql://user:pass@host:port/dbname"}
+
+    2. Individual parameters (for production AWS Secrets Manager):
+       {"host": "...", "port": 5432, "database": "...", "username": "...", "password": "..."}
+
+    The get_connection() method prioritizes DSN if present, otherwise builds connection
+    from individual parameters.
+    """
+
+    def __init__(self, db_config: dict) -> None:
+        super().__init__(db_config)
+        self.__connection = self.get_connection()
+
+    def cursor(self) -> Cursor:
+        return PostgresCursor(self.__connection.cursor())
+
+    def close(self):
+        self.__connection.close()
+
+    def commit(self):
+        self.__connection.commit()
+
+    def rollback(self):
+        self.__connection.rollback()
+
+    def get_connection(self):
+        """
+        Get a connection to the PostgreSQL database.
+
+        Parameters:
+        - schema (str, optional): The database schema to set for the
+            connection.
+
+        Returns:
+        - connection: A connection to the PostgreSQL database.
+        """
+        from psycopg2 import connect
+
+        # If DSN is provided, use it directly (simplifies fixture_foundry integration)
+        if "dsn" in self.db_config:
+            log.info("Connecting using DSN: %s", self.db_config["dsn"])
+            return connect(self.db_config["dsn"])
+
+        # Otherwise, build connection from individual parameters
+        dbname = self.db_config["database"]
+        user = self.db_config["username"]
+        password = self.db_config["password"]
+        host = self.db_config.get("host", "localhost")
+        port = self.db_config.get("port", 5432)
+        additional_config = self.db_config.get("configuration", {})
+
+        # Merge additional configuration parameters with the main connection parameters
+        connection_params = {
+            "dbname": dbname,
+            "user": user,
+            "password": password,
+            "host": host,
+            "port": port,
+        }
+
+        connection_params.update(additional_config)
+
+        log.info(
+            f"connection_params: dbname: {dbname}, user: {user}, host: {host}, port: {port}"
+        )
+
+        # Create a connection to the PostgreSQL database
+        return connect(**connection_params)

api_foundry_query_engine/dao/batch_operation_handler.py

@@ -0,0 +1,295 @@
+"""
+Batch Operation Handler
+
+Orchestrates execution of multiple operations with dependency resolution,
+reference substitution, and transaction management.
+"""
+
+from typing import Dict, Any, List, Optional
+from api_foundry_query_engine.operation import Operation
+from api_foundry_query_engine.dao.operation_dao import OperationDAO
+from api_foundry_query_engine.utils.dependency_resolver import DependencyResolver
+from api_foundry_query_engine.utils.reference_resolver import ReferenceResolver
+from api_foundry_query_engine.utils.app_exception import ApplicationException
+from api_foundry_query_engine.utils.logger import logger
+
+log = logger(__name__)
+
+
+class BatchOperationHandler:
+    """Handles execution of batch operations with dependencies."""
+
+    def __init__(self, batch_request: Dict[str, Any], connection, engine: str):
+        """
+        Initialize the batch operation handler.
+
+        Args:
+            batch_request: Batch request with 'operations' and 'options'
+            connection: Database connection for executing operations
+            engine: Database engine type (postgres, mysql, etc.)
+        """
+        self.operations = batch_request.get("operations", [])
+        self.options = batch_request.get("options", {})
+        self.connection = connection
+        self.engine = engine
+        self.results: Dict[str, Dict[str, Any]] = {}
+        self.failed_operations: List[str] = []
+
+        # Validate batch request
+        self._validate_batch_request()
+
+        # Resolve execution order
+        self.resolver = DependencyResolver(self.operations)
+        self.execution_order = self.resolver.get_execution_order()
+
+    def _validate_batch_request(self):
+        """Validate the batch request structure."""
+        if not self.operations:
+            raise ApplicationException(
+                400, "Batch request must contain at least one operation"
+            )
+
+        if len(self.operations) > 100:
+            raise ApplicationException(
+                400,
+                f"Batch size exceeds maximum (100). "
+                f"Requested: {len(self.operations)}",
+            )
+
+        # Auto-generate IDs for operations that don't have them
+        seen_ids = set()
+        for i, op in enumerate(self.operations):
+            # Auto-generate ID if not provided
+            if "id" not in op or not op["id"]:
+                op["id"] = f"op_{i}"
+
+            # Check for duplicate IDs
+            if op["id"] in seen_ids:
+                raise ApplicationException(
+                    400,
+                    f"Duplicate operation ID '{op['id']}' found",
+                )
+            seen_ids.add(op["id"])
+
+            # Validate required fields
+            if "entity" not in op:
+                raise ApplicationException(
+                    400,
+                    f"Operation '{op['id']}' missing required field 'entity'",
+                )
+            if "action" not in op:
+                raise ApplicationException(
+                    400,
+                    f"Operation '{op['id']}' missing required field 'action'",
+                )
+            if op["action"] not in ["create", "read", "update", "delete"]:
+                raise ApplicationException(
+                    400,
+                    f"Operation '{op['id']}' has invalid action " f"'{op['action']}'",
+                )
+
+    def execute(self) -> Dict[str, Any]:
+        """
+        Execute all operations in dependency order.
+
+        Returns:
+            Dictionary with 'success', 'results', and optionally 'errors'
+
+        Raises:
+            ApplicationException: On validation or execution errors
+        """
+        atomic = self.options.get("atomic", True)
+        continue_on_error = self.options.get("continueOnError", False)
+
+        log.info(
+            "Executing batch with %d operations (atomic=%s, continue=%s)",
+            len(self.operations),
+            atomic,
+            continue_on_error,
+        )
+
+        try:
+            for op_id in self.execution_order:
+                # Check if we should skip this operation
+                if self._should_skip_operation(op_id):
+                    self.results[op_id] = {
+                        "status": "skipped",
+                        "reason": "Dependency failed",
+                    }
+                    continue
+
+                # Execute the operation
+                try:
+                    result = self._execute_operation(op_id)
+                    # Unwrap single-item lists for easier reference access
+                    # e.g., $ref:op_0.invoice_id instead of $ref:op_0.0.invoice_id
+                    if isinstance(result, list) and len(result) == 1:
+                        result = result[0]
+                    self.results[op_id] = {"status": "completed", "data": result}
+                    log.info("Operation '%s' completed successfully", op_id)
+
+                except ApplicationException as e:
+                    log.error("Operation '%s' failed: %s", op_id, e.message)
+                    self.results[op_id] = {
+                        "status": "failed",
+                        "error": e.message,
+                        "statusCode": e.status_code,
+                    }
+                    self.failed_operations.append(op_id)
+
+                    # Rollback the failed transaction
+                    # In non-atomic mode, this allows next operation to proceed
+                    if not atomic:
+                        self.connection.rollback()
+                        log.info("Failed operation rolled back (non-atomic mode)")
+
+                    # Stop on error if not continuing
+                    if not continue_on_error:
+                        if atomic:
+                            self.connection.rollback()
+                            log.info("Transaction rolled back")
+                        raise ApplicationException(
+                            400,
+                            f"Batch failed at operation '{op_id}': {e.message}",
+                        ) from e
+
+                # Commit each operation if not atomic and it succeeded
+                if not atomic and op_id not in self.failed_operations:
+                    self.connection.commit()
+                    log.info("Operation '%s' committed (non-atomic mode)", op_id)
+
+            # Commit transaction if atomic and all succeeded
+            if atomic and not self.failed_operations:
+                self.connection.commit()
+                log.info("Batch transaction committed successfully")
+            elif atomic and self.failed_operations and continue_on_error:
+                self.connection.rollback()
+                log.info("Batch had failures, transaction rolled back")
+
+            # Build response
+            success = len(self.failed_operations) == 0
+            response = {"success": success, "results": self.results}
+
+            if self.failed_operations:
+                response["failedOperations"] = self.failed_operations
+
+            return response
+
+        except Exception as e:
+            # Rollback on any unexpected error
+            if atomic:
+                self.connection.rollback()
+                log.error("Batch failed, transaction rolled back")
+            raise e
+
+    def _should_skip_operation(self, op_id: str) -> bool:
+        """
+        Determine if operation should be skipped due to failed dependencies.
+
+        Args:
+            op_id: Operation ID to check
+
+        Returns:
+            True if operation should be skipped
+        """
+        op = next(o for o in self.operations if o["id"] == op_id)
+        depends_on = op.get("depends_on", [])
+
+        for dep_id in depends_on:
+            if dep_id in self.failed_operations:
+                log.info(
+                    "Skipping operation '%s' due to failed dependency '%s'",
+                    op_id,
+                    dep_id,
+                )
+                return True
+            if dep_id not in self.results:
+                log.warning(
+                    "Operation '%s' depends on '%s' which hasn't executed yet",
+                    op_id,
+                    dep_id,
+                )
+                return True
+            if self.results[dep_id].get("status") != "completed":
+                log.info(
+                    "Skipping operation '%s' due to incomplete dependency '%s'",
+                    op_id,
+                    dep_id,
+                )
+                return True
+
+        return False
+
+    def _execute_operation(self, op_id: str) -> Any:
+        """
+        Execute a single operation.
+
+        Args:
+            op_id: Operation ID to execute
+
+        Returns:
+            Operation result data
+
+        Raises:
+            ApplicationException: On execution error
+        """
+        # Get operation definition
+        op_def = next(o for o in self.operations if o["id"] == op_id)
+
+        # Resolve references in parameters
+        ref_resolver = ReferenceResolver(self.results)
+        query_params = ref_resolver.resolve_parameters(
+            op_def.get("query_params", {}), op_id
+        )
+        store_params = ref_resolver.resolve_parameters(
+            op_def.get("store_params", {}), op_id
+        )
+        metadata_params = ref_resolver.resolve_parameters(
+            op_def.get("metadata_params", {}), op_id
+        )
+
+        # Get claims from operation or use empty dict
+        claims = op_def.get("claims", {})
+
+        # Create Operation object
+        operation = Operation(
+            entity=op_def["entity"],
+            action=op_def["action"],
+            query_params=query_params,
+            store_params=store_params,
+            metadata_params=metadata_params,
+            claims=claims,
+        )
+
+        log.debug(
+            "Executing operation '%s': %s %s",
+            op_id,
+            operation.action,
+            operation.entity,
+        )
+
+        # Execute through OperationDAO
+        dao = OperationDAO(operation, self.engine)
+        result = dao.execute(self.connection)
+
+        return result
+
+    def get_operation_summary(self) -> Dict[str, int]:
+        """
+        Get summary statistics of batch execution.
+
+        Returns:
+            Dictionary with counts of completed, failed, skipped operations
+        """
+        summary = {"total": len(self.operations), "completed": 0, "failed": 0}
+
+        for op_result in self.results.values():
+            status = op_result.get("status")
+            if status == "completed":
+                summary["completed"] += 1
+            elif status == "failed":
+                summary["failed"] += 1
+            elif status == "skipped":
+                summary["skipped"] = summary.get("skipped", 0) + 1
+
+        return summary

api_foundry_query_engine/dao/dao.py

@@ -0,0 +1,23 @@
+import abc
+from typing import Union
+
+from api_foundry_query_engine.connectors.connection import Connection
+from api_foundry_query_engine.operation import Operation
+
+
+class DAO(metaclass=abc.ABCMeta):
+    @classmethod
+    def __subclasshook__(cls, __subclass: type) -> bool:
+        return hasattr(__subclass, "execute") and callable(__subclass.execute)
+
+    def execute(
+        self, connector: Connection, operation: Operation
+    ) -> Union[list[dict], dict]:
+        raise NotImplementedError
+
+
+class DAOAdapter(DAO):
+    def execute(
+        self, connector: Connection, operation: Operation
+    ) -> Union[list[dict], dict]:
+        return super().execute(connector, operation)

api_foundry_query_engine/dao/operation_dao.py

@@ -0,0 +1,186 @@
+from typing import Union
+
+from api_foundry_query_engine.dao.sql_custom_query_handler import SQLCustomQueryHandler
+from api_foundry_query_engine.dao.sql_delete_query_handler import (
+    SQLDeleteSchemaQueryHandler,
+)
+from api_foundry_query_engine.dao.sql_insert_query_handler import (
+    SQLInsertSchemaQueryHandler,
+)
+from api_foundry_query_engine.dao.sql_select_query_handler import (
+    SQLSelectSchemaQueryHandler,
+)
+from api_foundry_query_engine.dao.sql_subselect_query_handler import (
+    SQLSubselectSchemaQueryHandler,
+)
+from api_foundry_query_engine.dao.sql_update_query_handler import (
+    SQLUpdateSchemaQueryHandler,
+)
+from api_foundry_query_engine.dao.sql_restore_query_handler import (
+    SQLRestoreSchemaQueryHandler,
+)
+from api_foundry_query_engine.utils.app_exception import ApplicationException
+from api_foundry_query_engine.dao.dao import DAO
+from api_foundry_query_engine.connectors.connection import Cursor
+from api_foundry_query_engine.operation import Operation
+from api_foundry_query_engine.utils.api_model import (
+    get_path_operation,
+    get_schema_object,
+)
+from api_foundry_query_engine.dao.sql_query_handler import SQLQueryHandler
+
+
+class OperationDAO(DAO):
+    """
+    A class to handle database operations based on the provided
+    Operation object.
+
+    Attributes:
+        operation (Operation): The operation to perform.
+    """
+
+    def __init__(self, operation: Operation, engine: str) -> None:
+        """
+        Initialize the OperationDAO with the provided Operation object.
+
+        Args:
+            operation (Operation): The operation to perform.
+        """
+        super().__init__()
+        self.operation = operation
+        self.engine = engine
+
+    @property
+    def query_handler(self) -> SQLQueryHandler:
+        if not hasattr(self, "_query_handler"):
+            path_operation = get_path_operation(
+                self.operation.entity, self.operation.action
+            )
+            if path_operation:
+                self._query_handler = SQLCustomQueryHandler(
+                    self.operation, path_operation, self.engine
+                )
+                return self._query_handler
+
+            schema_object = get_schema_object(self.operation.entity)
+            if not schema_object:
+                raise ApplicationException(
+                    500, f"Unknown operation: {self.operation.entity}"
+                )
+            if self.operation.action == "read":
+                self._query_handler = SQLSelectSchemaQueryHandler(
+                    self.operation, schema_object, self.engine
+                )
+            elif self.operation.action == "create":
+                self._query_handler = SQLInsertSchemaQueryHandler(
+                    self.operation, schema_object, self.engine
+                )
+            elif self.operation.action == "update":
+                self._query_handler = SQLUpdateSchemaQueryHandler(
+                    self.operation, schema_object, self.engine
+                )
+            elif self.operation.action == "delete":
+                self._query_handler = SQLDeleteSchemaQueryHandler(
+                    self.operation, schema_object, self.engine
+                )
+            elif self.operation.action == "restore":
+                self._query_handler = SQLRestoreSchemaQueryHandler(
+                    self.operation, schema_object, self.engine
+                )
+            else:
+                raise ApplicationException(
+                    400, f"Invalid operation action: {self.operation.action}"
+                )
+        return self._query_handler
+
+    def execute(self, connector, operation=None) -> Union[list[dict], dict]:
+        """
+        Execute the database operation based on the provided connector.
+
+        Args:
+            connector (Connection): The database connection.
+            operation (Operation, optional): The operation to perform.
+
+        Returns:
+            list[dict]: A list of dictionaries containing the results
+            of the operation.
+        """
+
+        # Use self.operation if operation is not provided
+        op = operation if operation is not None else self.operation
+
+        # Check if this is a batch operation
+        if op.entity == "batch" and op.action == "create":
+            from api_foundry_query_engine.dao.batch_operation_handler import (
+                BatchOperationHandler,
+            )
+
+            # Extract batch request from store_params
+            batch_request = op.store_params
+
+            # Execute batch
+            handler = BatchOperationHandler(batch_request, connector, self.engine)
+            return handler.execute()
+
+        # Standard operation handling
+        # Assume connector has a 'cursor()' method to get a Cursor
+        cursor = connector.cursor()
+
+        result = self.__fetch_record_set(self.query_handler, cursor)
+
+        if op.action == "read":
+            if op.metadata_params.get("count", False):
+                return result[0]
+            self.__fetch_many(result, cursor)
+        elif op.action in ["update", "delete", "restore"] and len(result) == 0:
+            raise ApplicationException(400, "No records were modified")
+
+        return result
+
+    def __fetch_many(self, parent_set: list[dict], cursor: Cursor):
+        if "properties" not in self.operation.metadata_params:
+            return
+
+        schema_object = get_schema_object(self.operation.entity)
+        for name, relation in schema_object.relations.items():
+            if relation.type == "object":
+                continue
+
+            child_set = self.__fetch_record_set(
+                SQLSubselectSchemaQueryHandler(
+                    self.operation, relation, self.query_handler  # type: ignore
+                ),
+                cursor,
+            )
+            if len(child_set) == 0:
+                continue
+
+            for parent in parent_set:
+                parent[name] = []
+
+            parents = {}
+            for parent in parent_set:
+                parents[parent[relation.parent_property]] = parent
+
+            for child in child_set:
+                parent_id = child[relation.child_property]
+                parent = parents.get(parent_id)
+                if parent:
+                    parent[name].append(child)
+
+    def __fetch_record_set(
+        self, query_handler: SQLQueryHandler, cursor: Cursor
+    ) -> list[dict]:
+        sql = query_handler.sql
+        if not sql:
+            return []
+
+        record_set = cursor.execute(
+            sql, query_handler.placeholders, query_handler.selection_results
+        )
+        result = []
+        for record in record_set:
+            object = query_handler.marshal_record(record)
+            result.append(object)
+
+        return result