sqlspec 0.13.1__py3-none-any.whl → 0.16.2__py3-none-any.whl

sqlspec/driver/mixins/_pipeline.py

@@ -1,512 +0,0 @@
-"""Pipeline execution mixin for batch database operations.
-
-This module provides mixins that enable pipelined execution of SQL statements,
-allowing multiple operations to be sent to the database in a single network
-round-trip for improved performance.
-
-The implementation leverages native driver support where available (psycopg, asyncpg, oracledb)
-and provides high-quality simulated behavior for others.
-"""
-
-from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, Optional, Union, cast
-
-from sqlspec.exceptions import PipelineExecutionError
-from sqlspec.statement.filters import StatementFilter
-from sqlspec.statement.result import SQLResult
-from sqlspec.statement.sql import SQL
-from sqlspec.utils.logging import get_logger
-from sqlspec.utils.type_guards import (
-    is_async_pipeline_capable_driver,
-    is_async_transaction_state_capable,
-    is_sync_pipeline_capable_driver,
-    is_sync_transaction_state_capable,
-)
-
-if TYPE_CHECKING:
-    from typing import Literal
-
-    from sqlspec.config import DriverT
-    from sqlspec.driver import AsyncDriverAdapterProtocol, SyncDriverAdapterProtocol
-    from sqlspec.typing import StatementParameters
-
-__all__ = (
-    "AsyncPipeline",
-    "AsyncPipelinedExecutionMixin",
-    "Pipeline",
-    "PipelineOperation",
-    "SyncPipelinedExecutionMixin",
-)
-
-logger = get_logger(__name__)
-
-
-@dataclass
-class PipelineOperation:
-    """Container for a queued pipeline operation."""
-
-    sql: SQL
-    operation_type: "Literal['execute', 'execute_many', 'execute_script', 'select']"
-    filters: "Optional[list[StatementFilter]]" = None
-    original_params: "Optional[Any]" = None
-
-
-class SyncPipelinedExecutionMixin:
-    """Mixin providing pipeline execution for sync drivers."""
-
-    __slots__ = ()
-
-    def pipeline(
-        self,
-        *,
-        isolation_level: "Optional[str]" = None,
-        continue_on_error: bool = False,
-        max_operations: int = 1000,
-        **options: Any,
-    ) -> "Pipeline":
-        """Create a new pipeline for batch operations.
-
-        Args:
-            isolation_level: Transaction isolation level
-            continue_on_error: Continue processing after errors
-            max_operations: Maximum operations before auto-flush
-            **options: Driver-specific pipeline options
-
-        Returns:
-            A new Pipeline instance for queuing operations
-        """
-        return Pipeline(
-            driver=cast("SyncDriverAdapterProtocol[Any, Any]", self),
-            isolation_level=isolation_level,
-            continue_on_error=continue_on_error,
-            max_operations=max_operations,
-            options=options,
-        )
-
-
-class AsyncPipelinedExecutionMixin:
-    """Async version of pipeline execution mixin."""
-
-    __slots__ = ()
-
-    def pipeline(
-        self,
-        *,
-        isolation_level: "Optional[str]" = None,
-        continue_on_error: bool = False,
-        max_operations: int = 1000,
-        **options: Any,
-    ) -> "AsyncPipeline":
-        """Create a new async pipeline for batch operations."""
-        return AsyncPipeline(
-            driver=cast("AsyncDriverAdapterProtocol[Any, Any]", self),
-            isolation_level=isolation_level,
-            continue_on_error=continue_on_error,
-            max_operations=max_operations,
-            options=options,
-        )
-
-
-class Pipeline:
-    """Synchronous pipeline with enhanced parameter handling."""
-
-    def __init__(
-        self,
-        driver: "DriverT",  # pyright: ignore
-        isolation_level: "Optional[str]" = None,
-        continue_on_error: bool = False,
-        max_operations: int = 1000,
-        options: "Optional[dict[str, Any]]" = None,
-    ) -> None:
-        self.driver = driver
-        self.isolation_level = isolation_level
-        self.continue_on_error = continue_on_error
-        self.max_operations = max_operations
-        self.options = options or {}
-        self._operations: list[PipelineOperation] = []
-        self._results: Optional[list[SQLResult[Any]]] = None
-        self._simulation_logged = False
-
-    def add_execute(
-        self, statement: "Union[str, SQL]", /, *parameters: "Union[StatementParameters, StatementFilter]", **kwargs: Any
-    ) -> "Pipeline":
-        """Add an execute operation to the pipeline.
-
-        Args:
-            statement: SQL statement to execute
-            *parameters: Mixed positional args containing parameters and filters
-            **kwargs: Named parameters
-
-        Returns:
-            Self for fluent API
-        """
-        self._operations.append(
-            PipelineOperation(
-                sql=SQL(statement, parameters=parameters or None, config=self.driver.config, **kwargs),
-                operation_type="execute",
-            )
-        )
-
-        if len(self._operations) >= self.max_operations:
-            logger.warning("Pipeline auto-flushing at %s operations", len(self._operations))
-            self.process()
-
-        return self
-
-    def add_select(
-        self, statement: "Union[str, SQL]", /, *parameters: "Union[StatementParameters, StatementFilter]", **kwargs: Any
-    ) -> "Pipeline":
-        """Add a select operation to the pipeline."""
-        self._operations.append(
-            PipelineOperation(
-                sql=SQL(statement, parameters=parameters or None, config=self.driver.config, **kwargs),
-                operation_type="select",
-            )
-        )
-        return self
-
-    def add_execute_many(
-        self, statement: "Union[str, SQL]", /, *parameters: "Union[StatementParameters, StatementFilter]", **kwargs: Any
-    ) -> "Pipeline":
-        """Add batch execution preserving parameter types.
-
-        Args:
-            statement: SQL statement to execute multiple times
-            *parameters: First arg should be batch data (list of param sets),
-                        followed by optional StatementFilter instances
-            **kwargs: Not typically used for execute_many
-        """
-        # First parameter should be the batch data
-        if not parameters or not isinstance(parameters[0], (list, tuple)):
-            msg = "execute_many requires a sequence of parameter sets as first parameter"
-            raise ValueError(msg)
-
-        batch_params = parameters[0]
-        if isinstance(batch_params, tuple):
-            batch_params = list(batch_params)
-        sql_obj = SQL(
-            statement, parameters=parameters[1:] if len(parameters) > 1 else None, config=self.driver.config, **kwargs
-        ).as_many(batch_params)
-
-        self._operations.append(PipelineOperation(sql=sql_obj, operation_type="execute_many"))
-        return self
-
-    def add_execute_script(self, script: "Union[str, SQL]", *filters: StatementFilter, **kwargs: Any) -> "Pipeline":
-        """Add a multi-statement script to the pipeline."""
-        if isinstance(script, SQL):
-            sql_obj = script.as_script()
-        else:
-            sql_obj = SQL(script, parameters=filters or None, config=self.driver.config, **kwargs).as_script()
-
-        self._operations.append(PipelineOperation(sql=sql_obj, operation_type="execute_script"))
-        return self
-
-    def process(self, filters: "Optional[list[StatementFilter]]" = None) -> "list[SQLResult]":
-        """Execute all queued operations.
-
-        Args:
-            filters: Global filters to apply to all operations
-
-        Returns:
-            List of results from all operations
-        """
-        if not self._operations:
-            return []
-
-        if filters:
-            self._apply_global_filters(filters)
-
-        if is_sync_pipeline_capable_driver(self.driver):
-            results = self.driver._execute_pipeline_native(self._operations, **self.options)
-        else:
-            results = self._execute_pipeline_simulated()
-
-        self._results = results
-        self._operations.clear()
-        return cast("list[SQLResult]", results)
-
-    def _execute_pipeline_simulated(self) -> "list[SQLResult]":
-        """Enhanced simulation with transaction support and error handling."""
-        results: list[SQLResult[Any]] = []
-        connection = None
-        auto_transaction = False
-
-        if not self._simulation_logged:
-            logger.info(
-                "%s using simulated pipeline. Native support: %s",
-                self.driver.__class__.__name__,
-                self._has_native_support(),
-            )
-            self._simulation_logged = True
-
-        try:
-            connection = self.driver._connection()
-
-            if is_sync_transaction_state_capable(connection) and not connection.in_transaction():
-                connection.begin()
-                auto_transaction = True
-
-            for i, op in enumerate(self._operations):
-                self._execute_single_operation(i, op, results, connection, auto_transaction)
-
-            # Commit if we started the transaction
-            if auto_transaction and is_sync_transaction_state_capable(connection):
-                connection.commit()
-
-        except Exception as e:
-            if connection and auto_transaction and is_sync_transaction_state_capable(connection):
-                connection.rollback()
-            if not isinstance(e, PipelineExecutionError):
-                msg = f"Pipeline execution failed: {e}"
-                raise PipelineExecutionError(msg) from e
-            raise
-
-        return results
-
-    def _execute_single_operation(
-        self, i: int, op: PipelineOperation, results: "list[SQLResult[Any]]", connection: Any, auto_transaction: bool
-    ) -> None:
-        """Execute a single pipeline operation with error handling."""
-        try:
-            # Execute based on operation type
-            result: SQLResult[Any]
-            if op.operation_type == "execute_script":
-                result = cast("SQLResult[Any]", self.driver.execute_script(op.sql, _connection=connection))
-            elif op.operation_type == "execute_many":
-                result = cast("SQLResult[Any]", self.driver.execute_many(op.sql, _connection=connection))
-            else:
-                result = cast("SQLResult[Any]", self.driver.execute(op.sql, _connection=connection))
-
-            result.operation_index = i
-            result.pipeline_sql = op.sql
-            results.append(result)
-
-        except Exception as e:
-            if self.continue_on_error:
-                error_result = SQLResult(
-                    statement=op.sql, data=[], error=e, operation_index=i, parameters=op.sql.parameters
-                )
-                results.append(error_result)
-            else:
-                if auto_transaction and is_sync_transaction_state_capable(connection):
-                    connection.rollback()
-                msg = f"Pipeline failed at operation {i}: {e}"
-                raise PipelineExecutionError(
-                    msg, operation_index=i, partial_results=results, failed_operation=op
-                ) from e
-
-    def _apply_global_filters(self, filters: "list[StatementFilter]") -> None:
-        """Apply global filters to all operations."""
-        for operation in self._operations:
-            if operation.filters is None:
-                operation.filters = []
-            operation.filters.extend(filters)
-
-    def _apply_operation_filters(self, sql: SQL, filters: "list[StatementFilter]") -> SQL:
-        """Apply filters to a SQL object."""
-        result = sql
-        for filter_obj in filters:
-            result = filter_obj.append_to_statement(result)
-        return result
-
-    def _has_native_support(self) -> bool:
-        """Check if driver has native pipeline support."""
-        return is_sync_pipeline_capable_driver(self.driver)
-
-    def _process_parameters(self, params: tuple[Any, ...]) -> tuple["list[StatementFilter]", "Optional[Any]"]:
-        """Extract filters and parameters from mixed args.
-
-        Returns:
-            Tuple of (filters, parameters)
-        """
-        filters: list[StatementFilter] = []
-        parameters: list[Any] = []
-
-        for param in params:
-            if isinstance(param, StatementFilter):
-                filters.append(param)
-            else:
-                parameters.append(param)
-
-        if not parameters:
-            return filters, None
-        if len(parameters) == 1:
-            return filters, parameters[0]
-        return filters, parameters
-
-    @property
-    def operations(self) -> "list[PipelineOperation]":
-        """Get the current list of queued operations."""
-        return self._operations.copy()
-
-
-class AsyncPipeline:
-    """Asynchronous pipeline with identical structure to Pipeline."""
-
-    def __init__(
-        self,
-        driver: "AsyncDriverAdapterProtocol[Any, Any]",
-        isolation_level: "Optional[str]" = None,
-        continue_on_error: bool = False,
-        max_operations: int = 1000,
-        options: "Optional[dict[str, Any]]" = None,
-    ) -> None:
-        self.driver = driver
-        self.isolation_level = isolation_level
-        self.continue_on_error = continue_on_error
-        self.max_operations = max_operations
-        self.options = options or {}
-        self._operations: list[PipelineOperation] = []
-        self._results: Optional[list[SQLResult[Any]]] = None
-        self._simulation_logged = False
-
-    async def add_execute(
-        self, statement: "Union[str, SQL]", /, *parameters: "Union[StatementParameters, StatementFilter]", **kwargs: Any
-    ) -> "AsyncPipeline":
-        """Add an execute operation to the async pipeline."""
-        self._operations.append(
-            PipelineOperation(
-                sql=SQL(statement, parameters=parameters or None, config=self.driver.config, **kwargs),
-                operation_type="execute",
-            )
-        )
-
-        if len(self._operations) >= self.max_operations:
-            logger.warning("Async pipeline auto-flushing at %s operations", len(self._operations))
-            await self.process()
-
-        return self
-
-    async def add_select(
-        self, statement: "Union[str, SQL]", /, *parameters: "Union[StatementParameters, StatementFilter]", **kwargs: Any
-    ) -> "AsyncPipeline":
-        """Add a select operation to the async pipeline."""
-        self._operations.append(
-            PipelineOperation(
-                sql=SQL(statement, parameters=parameters or None, config=self.driver.config, **kwargs),
-                operation_type="select",
-            )
-        )
-        return self
-
-    async def add_execute_many(
-        self, statement: "Union[str, SQL]", /, *parameters: "Union[StatementParameters, StatementFilter]", **kwargs: Any
-    ) -> "AsyncPipeline":
-        """Add batch execution to the async pipeline."""
-        # First parameter should be the batch data
-        if not parameters or not isinstance(parameters[0], (list, tuple)):
-            msg = "execute_many requires a sequence of parameter sets as first parameter"
-            raise ValueError(msg)
-
-        batch_params = parameters[0]
-        if isinstance(batch_params, tuple):
-            batch_params = list(batch_params)
-        sql_obj = SQL(
-            statement, parameters=parameters[1:] if len(parameters) > 1 else None, config=self.driver.config, **kwargs
-        ).as_many(batch_params)
-
-        self._operations.append(PipelineOperation(sql=sql_obj, operation_type="execute_many"))
-        return self
-
-    async def add_execute_script(
-        self, script: "Union[str, SQL]", *filters: StatementFilter, **kwargs: Any
-    ) -> "AsyncPipeline":
-        """Add a script to the async pipeline."""
-        if isinstance(script, SQL):
-            sql_obj = script.as_script()
-        else:
-            sql_obj = SQL(script, parameters=filters or None, config=self.driver.config, **kwargs).as_script()
-
-        self._operations.append(PipelineOperation(sql=sql_obj, operation_type="execute_script"))
-        return self
-
-    async def process(self, filters: "Optional[list[StatementFilter]]" = None) -> "list[SQLResult]":
-        """Execute all queued operations asynchronously."""
-        if not self._operations:
-            return []
-
-        if is_async_pipeline_capable_driver(self.driver):
-            results = await cast("Any", self.driver)._execute_pipeline_native(self._operations, **self.options)
-        else:
-            results = await self._execute_pipeline_simulated()
-
-        self._results = results
-        self._operations.clear()
-        return cast("list[SQLResult]", results)
-
-    async def _execute_pipeline_simulated(self) -> "list[SQLResult]":
-        """Async version of simulated pipeline execution."""
-        results: list[SQLResult[Any]] = []
-        connection = None
-        auto_transaction = False
-
-        if not self._simulation_logged:
-            logger.info(
-                "%s using simulated async pipeline. Native support: %s",
-                self.driver.__class__.__name__,
-                self._has_native_support(),
-            )
-            self._simulation_logged = True
-
-        try:
-            connection = self.driver._connection()
-
-            if is_async_transaction_state_capable(connection) and not connection.in_transaction():
-                await connection.begin()
-                auto_transaction = True
-
-            for i, op in enumerate(self._operations):
-                await self._execute_single_operation_async(i, op, results, connection, auto_transaction)
-
-            if auto_transaction and is_async_transaction_state_capable(connection):
-                await connection.commit()
-
-        except Exception as e:
-            if connection and auto_transaction and is_async_transaction_state_capable(connection):
-                await connection.rollback()
-            if not isinstance(e, PipelineExecutionError):
-                msg = f"Async pipeline execution failed: {e}"
-                raise PipelineExecutionError(msg) from e
-            raise
-
-        return results
-
-    async def _execute_single_operation_async(
-        self, i: int, op: PipelineOperation, results: "list[SQLResult[Any]]", connection: Any, auto_transaction: bool
-    ) -> None:
-        """Execute a single async pipeline operation with error handling."""
-        try:
-            result: SQLResult[Any]
-            if op.operation_type == "execute_script":
-                result = await self.driver.execute_script(op.sql, _connection=connection)
-            elif op.operation_type == "execute_many":
-                result = await self.driver.execute_many(op.sql, _connection=connection)
-            else:
-                result = await self.driver.execute(op.sql, _connection=connection)
-
-            result.operation_index = i
-            result.pipeline_sql = op.sql
-            results.append(result)
-
-        except Exception as e:
-            if self.continue_on_error:
-                error_result = SQLResult(
-                    statement=op.sql, data=[], error=e, operation_index=i, parameters=op.sql.parameters
-                )
-                results.append(error_result)
-            else:
-                if auto_transaction and is_async_transaction_state_capable(connection):
-                    await connection.rollback()
-                msg = f"Async pipeline failed at operation {i}: {e}"
-                raise PipelineExecutionError(
-                    msg, operation_index=i, partial_results=results, failed_operation=op
-                ) from e
-
-    def _has_native_support(self) -> bool:
-        """Check if driver has native pipeline support."""
-        return is_async_pipeline_capable_driver(self.driver)
-
-    @property
-    def operations(self) -> "list[PipelineOperation]":
-        """Get the current list of queued operations."""
-        return self._operations.copy()

sqlspec/driver/mixins/_result_utils.py

@@ -1,140 +0,0 @@
-"""Result conversion utilities for unified storage architecture.
-
-This module contains the result conversion functionality integrated with the unified
-storage architecture.
-"""
-
-import datetime
-from collections.abc import Sequence
-from enum import Enum
-from functools import partial
-from pathlib import Path, PurePath
-from typing import Any, Callable, Optional, Union, cast, overload
-from uuid import UUID
-
-from sqlspec.exceptions import SQLSpecError, wrap_exceptions
-from sqlspec.statement.result import OperationType
-from sqlspec.typing import ModelDTOT, ModelT, convert, get_type_adapter
-from sqlspec.utils.type_guards import is_dataclass, is_msgspec_struct, is_pydantic_model
-
-__all__ = ("_DEFAULT_TYPE_DECODERS", "ToSchemaMixin", "_default_msgspec_deserializer")
-
-
-_DEFAULT_TYPE_DECODERS: list[tuple[Callable[[Any], bool], Callable[[Any, Any], Any]]] = [
-    (lambda x: x is UUID, lambda t, v: t(v.hex)),
-    (lambda x: x is datetime.datetime, lambda t, v: t(v.isoformat())),
-    (lambda x: x is datetime.date, lambda t, v: t(v.isoformat())),
-    (lambda x: x is datetime.time, lambda t, v: t(v.isoformat())),
-    (lambda x: x is Enum, lambda t, v: t(v.value)),
-]
-
-
-def _default_msgspec_deserializer(
-    target_type: Any, value: Any, type_decoders: "Optional[Sequence[tuple[Any, Any]]]" = None
-) -> Any:
-    """Default msgspec deserializer with type conversion support."""
-    if type_decoders:
-        for predicate, decoder in type_decoders:
-            if predicate(target_type):
-                return decoder(target_type, value)
-    if target_type is UUID and isinstance(value, UUID):
-        return value.hex
-    if target_type in {datetime.datetime, datetime.date, datetime.time}:
-        with wrap_exceptions(suppress=AttributeError):
-            return value.isoformat()
-    if isinstance(target_type, type) and issubclass(target_type, Enum) and isinstance(value, Enum):
-        return value.value
-    if isinstance(value, target_type):
-        return value
-    if issubclass(target_type, (Path, PurePath, UUID)):
-        return target_type(value)
-    return value
-
-
-class ToSchemaMixin:
-    __slots__ = ()
-
-    @staticmethod
-    def _determine_operation_type(statement: "Any") -> OperationType:
-        """Determine operation type from SQL statement expression.
-
-        Args:
-            statement: SQL statement object with expression attribute
-
-        Returns:
-            OperationType literal value
-        """
-        if not hasattr(statement, "expression") or not statement.expression:
-            return "EXECUTE"
-
-        expr_type = type(statement.expression).__name__.upper()
-        if "INSERT" in expr_type:
-            return "INSERT"
-        if "UPDATE" in expr_type:
-            return "UPDATE"
-        if "DELETE" in expr_type:
-            return "DELETE"
-        if "SELECT" in expr_type:
-            return "SELECT"
-        return "EXECUTE"
-
-    @overload
-    @staticmethod
-    def to_schema(data: "ModelT", *, schema_type: None = None) -> "ModelT": ...
-    @overload
-    @staticmethod
-    def to_schema(data: "dict[str, Any]", *, schema_type: "type[ModelDTOT]") -> "ModelDTOT": ...
-    @overload
-    @staticmethod
-    def to_schema(data: "Sequence[ModelT]", *, schema_type: None = None) -> "Sequence[ModelT]": ...
-    @overload
-    @staticmethod
-    def to_schema(data: "Sequence[dict[str, Any]]", *, schema_type: "type[ModelDTOT]") -> "Sequence[ModelDTOT]": ...
-
-    @staticmethod
-    def to_schema(
-        data: "Union[ModelT, dict[str, Any], Sequence[ModelT], Sequence[dict[str, Any]]]",
-        *,
-        schema_type: "Optional[type[ModelDTOT]]" = None,
-    ) -> "Union[ModelT, ModelDTOT, Sequence[ModelT], Sequence[ModelDTOT]]":
-        """Convert data to a specified schema type."""
-        if schema_type is None:
-            if not isinstance(data, Sequence):
-                return cast("ModelT", data)
-            return cast("Sequence[ModelT]", data)
-        if is_dataclass(schema_type):
-            if not isinstance(data, Sequence):
-                return cast("ModelDTOT", schema_type(**data))  # type: ignore[operator]
-            return cast("Sequence[ModelDTOT]", [schema_type(**item) for item in data])  # type: ignore[operator]
-        if is_msgspec_struct(schema_type):
-            if not isinstance(data, Sequence):
-                return cast(
-                    "ModelDTOT",
-                    convert(
-                        obj=data,
-                        type=schema_type,
-                        from_attributes=True,
-                        dec_hook=partial(_default_msgspec_deserializer, type_decoders=_DEFAULT_TYPE_DECODERS),
-                    ),
-                )
-            return cast(
-                "Sequence[ModelDTOT]",
-                convert(
-                    obj=data,
-                    type=list[schema_type],  # type: ignore[valid-type]  # pyright: ignore
-                    from_attributes=True,
-                    dec_hook=partial(_default_msgspec_deserializer, type_decoders=_DEFAULT_TYPE_DECODERS),
-                ),
-            )
-        if schema_type is not None and is_pydantic_model(schema_type):
-            if not isinstance(data, Sequence):
-                return cast(
-                    "ModelDTOT",
-                    get_type_adapter(schema_type).validate_python(data, from_attributes=True),  # pyright: ignore
-                )
-            return cast(
-                "Sequence[ModelDTOT]",
-                get_type_adapter(list[schema_type]).validate_python(data, from_attributes=True),  # type: ignore[valid-type]  # pyright: ignore
-            )
-        msg = "`schema_type` should be a valid Dataclass, Pydantic model or Msgspec struct"
-        raise SQLSpecError(msg)