sqlspec 0.16.0__cp311-cp311-macosx_13_0_x86_64.whl

sqlspec/driver/_sync.py

@@ -0,0 +1,473 @@
+"""Synchronous driver protocol implementation.
+
+This module provides the sync driver infrastructure for database adapters,
+including connection management, transaction support, and result processing.
+"""
+
+from abc import abstractmethod
+from typing import TYPE_CHECKING, Any, Optional, Union, cast, overload
+
+from sqlspec.core import SQL
+from sqlspec.driver._common import CommonDriverAttributesMixin, ExecutionResult
+from sqlspec.driver.mixins import SQLTranslatorMixin, ToSchemaMixin
+from sqlspec.exceptions import NotFoundError
+from sqlspec.utils.logging import get_logger
+from sqlspec.utils.type_guards import is_dict_row, is_indexable_row
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+    from contextlib import AbstractContextManager
+
+    from sqlspec.builder import QueryBuilder
+    from sqlspec.core import SQLResult, Statement, StatementConfig, StatementFilter
+    from sqlspec.typing import ModelDTOT, ModelT, RowT, StatementParameters
+
+logger = get_logger("sqlspec")
+
+__all__ = ("SyncDriverAdapterBase",)
+
+
+EMPTY_FILTERS: "list[StatementFilter]" = []
+
+
+class SyncDriverAdapterBase(CommonDriverAttributesMixin, SQLTranslatorMixin, ToSchemaMixin):
+    """Base class for synchronous database drivers.
+
+    Provides the foundation for sync database adapters, including connection management,
+    transaction support, and SQL execution methods. All database operations are performed
+    synchronously and support context manager patterns for proper resource cleanup.
+    """
+
+    __slots__ = ()
+
+    def dispatch_statement_execution(self, statement: "SQL", connection: "Any") -> "SQLResult":
+        """Central execution dispatcher using the Template Method Pattern.
+
+        Orchestrates the common execution flow, delegating database-specific steps
+        to abstract methods that concrete adapters must implement.
+        All database operations are wrapped in exception handling.
+
+        Args:
+            statement: The SQL statement to execute
+            connection: The database connection to use
+
+        Returns:
+            The result of the SQL execution
+        """
+        with self.handle_database_exceptions(), self.with_cursor(connection) as cursor:
+            special_result = self._try_special_handling(cursor, statement)
+            if special_result is not None:
+                return special_result
+
+            if statement.is_script:
+                execution_result = self._execute_script(cursor, statement)
+            elif statement.is_many:
+                execution_result = self._execute_many(cursor, statement)
+            else:
+                execution_result = self._execute_statement(cursor, statement)
+
+            return self.build_statement_result(statement, execution_result)
+
+    @abstractmethod
+    def with_cursor(self, connection: Any) -> Any:
+        """Create and return a context manager for cursor acquisition and cleanup.
+
+        Returns a context manager that yields a cursor for database operations.
+        Concrete implementations handle database-specific cursor creation and cleanup.
+        """
+
+    @abstractmethod
+    def handle_database_exceptions(self) -> "AbstractContextManager[None]":
+        """Handle database-specific exceptions and wrap them appropriately.
+
+        Returns:
+            ContextManager that can be used in with statements
+        """
+
+    @abstractmethod
+    def begin(self) -> None:
+        """Begin a database transaction on the current connection."""
+
+    @abstractmethod
+    def rollback(self) -> None:
+        """Rollback the current transaction on the current connection."""
+
+    @abstractmethod
+    def commit(self) -> None:
+        """Commit the current transaction on the current connection."""
+
+    @abstractmethod
+    def _try_special_handling(self, cursor: Any, statement: "SQL") -> "Optional[SQLResult]":
+        """Hook for database-specific special operations (e.g., PostgreSQL COPY, bulk operations).
+
+        This method is called first in dispatch_statement_execution() to allow drivers to handle
+        special operations that don't follow the standard SQL execution pattern.
+
+        Args:
+            cursor: Database cursor/connection object
+            statement: SQL statement to analyze
+
+        Returns:
+            SQLResult if the special operation was handled and completed,
+            None if standard execution should proceed
+        """
+
+    def _execute_script(self, cursor: Any, statement: "SQL") -> ExecutionResult:
+        """Execute a SQL script containing multiple statements.
+
+        Default implementation splits the script and executes statements individually.
+        Drivers can override for database-specific script execution methods.
+
+        Args:
+            cursor: Database cursor/connection object
+            statement: SQL statement object with all necessary data and configuration
+
+        Returns:
+            ExecutionResult with script execution data including statement counts
+        """
+        sql, prepared_parameters = self._get_compiled_sql(statement, self.statement_config)
+        statements = self.split_script_statements(sql, self.statement_config, strip_trailing_semicolon=True)
+
+        for stmt in statements:
+            single_stmt = statement.copy(statement=stmt, parameters=prepared_parameters)
+            self._execute_statement(cursor, single_stmt)
+
+        return self.create_execution_result(
+            cursor, statement_count=len(statements), successful_statements=len(statements), is_script_result=True
+        )
+
+    @abstractmethod
+    def _execute_many(self, cursor: Any, statement: "SQL") -> ExecutionResult:
+        """Execute SQL with multiple parameter sets (executemany).
+
+        Must be implemented by each driver for database-specific executemany logic.
+
+        Args:
+            cursor: Database cursor/connection object
+            statement: SQL statement object with all necessary data and configuration
+
+        Returns:
+            ExecutionResult with execution data for the many operation
+        """
+
+    @abstractmethod
+    def _execute_statement(self, cursor: Any, statement: "SQL") -> ExecutionResult:
+        """Execute a single SQL statement.
+
+        Must be implemented by each driver for database-specific execution logic.
+
+        Args:
+            cursor: Database cursor/connection object
+            statement: SQL statement object with all necessary data and configuration
+
+        Returns:
+            ExecutionResult with execution data
+        """
+
+    def execute(
+        self,
+        statement: "Union[SQL, Statement, QueryBuilder]",
+        /,
+        *parameters: "Union[StatementParameters, StatementFilter]",
+        statement_config: "Optional[StatementConfig]" = None,
+        **kwargs: Any,
+    ) -> "SQLResult":
+        """Execute a statement with parameter handling."""
+        sql_statement = self.prepare_statement(
+            statement, parameters, statement_config=statement_config or self.statement_config, kwargs=kwargs
+        )
+        return self.dispatch_statement_execution(statement=sql_statement, connection=self.connection)
+
+    def execute_many(
+        self,
+        statement: "Union[SQL, Statement, QueryBuilder]",
+        /,
+        parameters: "Sequence[StatementParameters]",
+        *filters: "Union[StatementParameters, StatementFilter]",
+        statement_config: "Optional[StatementConfig]" = None,
+        **kwargs: Any,
+    ) -> "SQLResult":
+        """Execute statement multiple times with different parameters.
+
+        Parameters passed will be used as the batch execution sequence.
+        """
+        config = statement_config or self.statement_config
+
+        if isinstance(statement, SQL):
+            sql_statement = SQL(statement._raw_sql, parameters, statement_config=config, is_many=True, **kwargs)
+        else:
+            base_statement = self.prepare_statement(statement, filters, statement_config=config, kwargs=kwargs)
+            sql_statement = SQL(base_statement._raw_sql, parameters, statement_config=config, is_many=True, **kwargs)
+
+        return self.dispatch_statement_execution(statement=sql_statement, connection=self.connection)
+
+    def execute_script(
+        self,
+        statement: "Union[str, SQL]",
+        /,
+        *parameters: "Union[StatementParameters, StatementFilter]",
+        statement_config: "Optional[StatementConfig]" = None,
+        **kwargs: Any,
+    ) -> "SQLResult":
+        """Execute a multi-statement script.
+
+        By default, validates each statement and logs warnings for dangerous
+        operations. Use suppress_warnings=True for migrations and admin scripts.
+        """
+        script_config = statement_config or self.statement_config
+        sql_statement = self.prepare_statement(statement, parameters, statement_config=script_config, kwargs=kwargs)
+
+        return self.dispatch_statement_execution(statement=sql_statement.as_script(), connection=self.connection)
+
+    @overload
+    def select_one(
+        self,
+        statement: "Union[Statement, QueryBuilder]",
+        /,
+        *parameters: "Union[StatementParameters, StatementFilter]",
+        schema_type: "type[ModelDTOT]",
+        statement_config: "Optional[StatementConfig]" = None,
+        **kwargs: Any,
+    ) -> "ModelDTOT": ...
+
+    @overload
+    def select_one(
+        self,
+        statement: "Union[Statement, QueryBuilder]",
+        /,
+        *parameters: "Union[StatementParameters, StatementFilter]",
+        schema_type: None = None,
+        statement_config: "Optional[StatementConfig]" = None,
+        **kwargs: Any,
+    ) -> "Union[ModelT, RowT, dict[str, Any]]": ...  # pyright: ignore[reportInvalidTypeVarUse]
+
+    def select_one(
+        self,
+        statement: "Union[Statement, QueryBuilder]",
+        /,
+        *parameters: "Union[StatementParameters, StatementFilter]",
+        schema_type: "Optional[type[ModelDTOT]]" = None,
+        statement_config: "Optional[StatementConfig]" = None,
+        **kwargs: Any,
+    ) -> "Union[ModelT, RowT, ModelDTOT]":  # pyright: ignore[reportInvalidTypeVarUse]
+        """Execute a select statement and return exactly one row.
+
+        Raises an exception if no rows or more than one row is returned.
+        """
+        result = self.execute(statement, *parameters, statement_config=statement_config, **kwargs)
+        data = result.get_data()
+        if not data:
+            msg = "No rows found"
+            raise NotFoundError(msg)
+        if len(data) > 1:
+            msg = f"Expected exactly one row, found {len(data)}"
+            raise ValueError(msg)
+        return cast(
+            "Union[ModelT, RowT, ModelDTOT]",
+            self.to_schema(data[0], schema_type=schema_type) if schema_type else data[0],
+        )
+
+    @overload
+    def select_one_or_none(
+        self,
+        statement: "Union[Statement, QueryBuilder]",
+        /,
+        *parameters: "Union[StatementParameters, StatementFilter]",
+        schema_type: "type[ModelDTOT]",
+        statement_config: "Optional[StatementConfig]" = None,
+        **kwargs: Any,
+    ) -> "Optional[ModelDTOT]": ...
+
+    @overload
+    def select_one_or_none(
+        self,
+        statement: "Union[Statement, QueryBuilder]",
+        /,
+        *parameters: "Union[StatementParameters, StatementFilter]",
+        schema_type: None = None,
+        statement_config: "Optional[StatementConfig]" = None,
+        **kwargs: Any,
+    ) -> "Optional[ModelT]": ...  # pyright: ignore[reportInvalidTypeVarUse]
+
+    def select_one_or_none(
+        self,
+        statement: "Union[Statement, QueryBuilder]",
+        /,
+        *parameters: "Union[StatementParameters, StatementFilter]",
+        schema_type: "Optional[type[ModelDTOT]]" = None,
+        statement_config: "Optional[StatementConfig]" = None,
+        **kwargs: Any,
+    ) -> "Optional[Union[ModelT, ModelDTOT]]":  # pyright: ignore[reportInvalidTypeVarUse]
+        """Execute a select statement and return at most one row.
+
+        Returns None if no rows are found.
+        Raises an exception if more than one row is returned.
+        """
+        result = self.execute(statement, *parameters, statement_config=statement_config, **kwargs)
+        data = result.get_data()
+        if not data:
+            return None
+        if len(data) > 1:
+            msg = f"Expected at most one row, found {len(data)}"
+            raise ValueError(msg)
+        return cast("Optional[Union[ModelT, ModelDTOT]]", self.to_schema(data[0], schema_type=schema_type))
+
+    @overload
+    def select(
+        self,
+        statement: "Union[Statement, QueryBuilder]",
+        /,
+        *parameters: "Union[StatementParameters, StatementFilter]",
+        schema_type: "type[ModelDTOT]",
+        statement_config: "Optional[StatementConfig]" = None,
+        **kwargs: Any,
+    ) -> "list[ModelDTOT]": ...
+
+    @overload
+    def select(
+        self,
+        statement: "Union[Statement, QueryBuilder]",
+        /,
+        *parameters: "Union[StatementParameters, StatementFilter]",
+        schema_type: None = None,
+        statement_config: "Optional[StatementConfig]" = None,
+        **kwargs: Any,
+    ) -> "list[ModelT]": ...  # pyright: ignore[reportInvalidTypeVarUse]
+
+    def select(
+        self,
+        statement: "Union[Statement, QueryBuilder]",
+        /,
+        *parameters: "Union[StatementParameters, StatementFilter]",
+        schema_type: "Optional[type[ModelDTOT]]" = None,
+        statement_config: "Optional[StatementConfig]" = None,
+        **kwargs: Any,
+    ) -> "Union[list[ModelT], list[ModelDTOT]]":  # pyright: ignore[reportInvalidTypeVarUse]
+        """Execute a select statement and return all rows."""
+        result = self.execute(statement, *parameters, statement_config=statement_config, **kwargs)
+        return cast(
+            "Union[list[ModelT], list[ModelDTOT]]",
+            self.to_schema(cast("list[ModelT]", result.get_data()), schema_type=schema_type),
+        )
+
+    def select_value(
+        self,
+        statement: "Union[Statement, QueryBuilder]",
+        /,
+        *parameters: "Union[StatementParameters, StatementFilter]",
+        statement_config: "Optional[StatementConfig]" = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Execute a select statement and return a single scalar value.
+
+        Expects exactly one row with one column.
+        Raises an exception if no rows or more than one row/column is returned.
+        """
+        result = self.execute(statement, *parameters, statement_config=statement_config, **kwargs)
+        try:
+            row = result.one()
+        except ValueError as e:
+            msg = "No rows found"
+            raise NotFoundError(msg) from e
+        if not row:
+            msg = "No rows found"
+            raise NotFoundError(msg)
+        if is_dict_row(row):
+            if not row:
+                msg = "Row has no columns"
+                raise ValueError(msg)
+            return next(iter(row.values()))
+        if is_indexable_row(row):
+            if not row:
+                msg = "Row has no columns"
+                raise ValueError(msg)
+            return row[0]
+        msg = f"Unexpected row type: {type(row)}"
+        raise ValueError(msg)
+
+    def select_value_or_none(
+        self,
+        statement: "Union[Statement, QueryBuilder]",
+        /,
+        *parameters: "Union[StatementParameters, StatementFilter]",
+        statement_config: "Optional[StatementConfig]" = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Execute a select statement and return a single scalar value or None.
+
+        Returns None if no rows are found.
+        Expects at most one row with one column.
+        Raises an exception if more than one row is returned.
+        """
+        result = self.execute(statement, *parameters, statement_config=statement_config, **kwargs)
+        data = result.get_data()
+        if not data:
+            return None
+        if len(data) > 1:
+            msg = f"Expected at most one row, found {len(data)}"
+            raise ValueError(msg)
+        row = data[0]
+        if isinstance(row, dict):
+            if not row:
+                return None
+            return next(iter(row.values()))
+        if isinstance(row, (tuple, list)):
+            return row[0]
+        msg = f"Cannot extract value from row type {type(row).__name__}"
+        raise TypeError(msg)
+
+    @overload
+    def select_with_total(
+        self,
+        statement: "Union[Statement, QueryBuilder]",
+        /,
+        *parameters: "Union[StatementParameters, StatementFilter]",
+        schema_type: "type[ModelDTOT]",
+        statement_config: "Optional[StatementConfig]" = None,
+        **kwargs: Any,
+    ) -> "tuple[list[ModelDTOT], int]": ...
+
+    @overload
+    def select_with_total(
+        self,
+        statement: "Union[Statement, QueryBuilder]",
+        /,
+        *parameters: "Union[StatementParameters, StatementFilter]",
+        schema_type: None = None,
+        statement_config: "Optional[StatementConfig]" = None,
+        **kwargs: Any,
+    ) -> "tuple[list[dict[str, Any]], int]": ...
+
+    def select_with_total(
+        self,
+        statement: "Union[Statement, QueryBuilder]",
+        /,
+        *parameters: "Union[StatementParameters, StatementFilter]",
+        schema_type: "Optional[type[ModelDTOT]]" = None,
+        statement_config: "Optional[StatementConfig]" = None,
+        **kwargs: Any,
+    ) -> "tuple[Union[list[dict[str, Any]], list[ModelDTOT]], int]":
+        """Execute a select statement and return both the data and total count.
+
+        This method is designed for pagination scenarios where you need both
+        the current page of data and the total number of rows that match the query.
+
+        Args:
+            statement: The SQL statement, QueryBuilder, or raw SQL string
+            *parameters: Parameters for the SQL statement
+            schema_type: Optional schema type for data transformation
+            statement_config: Optional SQL configuration
+            **kwargs: Additional keyword arguments
+
+        Returns:
+            A tuple containing:
+            - List of data rows (transformed by schema_type if provided)
+            - Total count of rows matching the query (ignoring LIMIT/OFFSET)
+        """
+        sql_statement = self.prepare_statement(
+            statement, parameters, statement_config=statement_config or self.statement_config, kwargs=kwargs
+        )
+        count_result = self.dispatch_statement_execution(self._create_count_query(sql_statement), self.connection)
+        select_result = self.execute(sql_statement)
+
+        return (self.to_schema(select_result.get_data(), schema_type=schema_type), count_result.scalar())

sqlspec/driver/mixins/__init__.py

@@ -0,0 +1,6 @@
+"""Driver mixins for instrumentation, storage, and utilities."""
+
+from sqlspec.driver.mixins._result_tools import ToSchemaMixin
+from sqlspec.driver.mixins._sql_translator import SQLTranslatorMixin
+
+__all__ = ("SQLTranslatorMixin", "ToSchemaMixin")

sqlspec/driver/mixins/_result_tools.py

@@ -0,0 +1,164 @@
+# pyright: reportCallIssue=false, reportAttributeAccessIssue=false, reportArgumentType=false
+import datetime
+import logging
+from collections.abc import Sequence
+from enum import Enum
+from functools import partial
+from pathlib import Path, PurePath
+from typing import Any, Callable, Optional, overload
+from uuid import UUID
+
+from mypy_extensions import trait
+
+from sqlspec.exceptions import SQLSpecError, wrap_exceptions
+from sqlspec.typing import (
+    CATTRS_INSTALLED,
+    ModelDTOT,
+    ModelT,
+    attrs_asdict,
+    cattrs_structure,
+    cattrs_unstructure,
+    convert,
+    get_type_adapter,
+)
+from sqlspec.utils.type_guards import is_attrs_schema, is_dataclass, is_msgspec_struct, is_pydantic_model
+
+__all__ = ("_DEFAULT_TYPE_DECODERS", "_default_msgspec_deserializer")
+
+
+logger = logging.getLogger(__name__)
+_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.
+
+    Converts values to appropriate types for msgspec deserialization, including
+    UUID, datetime, date, time, Enum, Path, and PurePath types.
+    """
+    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
+
+
+@trait
+class ToSchemaMixin:
+    __slots__ = ()
+
+    # Schema conversion overloads - handle common cases first
+    @overload
+    @staticmethod
+    def to_schema(data: "list[dict[str, Any]]") -> "list[dict[str, Any]]": ...
+    @overload
+    @staticmethod
+    def to_schema(data: "list[dict[str, Any]]", *, schema_type: "type[ModelDTOT]") -> "list[ModelDTOT]": ...
+    @overload
+    @staticmethod
+    def to_schema(data: "list[dict[str, Any]]", *, schema_type: None = None) -> "list[dict[str, Any]]": ...
+    @overload
+    @staticmethod
+    def to_schema(data: "dict[str, Any]") -> "dict[str, Any]": ...
+    @overload
+    @staticmethod
+    def to_schema(data: "dict[str, Any]", *, schema_type: "type[ModelDTOT]") -> "ModelDTOT": ...
+    @overload
+    @staticmethod
+    def to_schema(data: "dict[str, Any]", *, schema_type: None = None) -> "dict[str, Any]": ...
+    @overload
+    @staticmethod
+    def to_schema(data: "list[ModelT]") -> "list[ModelT]": ...
+    @overload
+    @staticmethod
+    def to_schema(data: "list[ModelT]", *, schema_type: "type[ModelDTOT]") -> "list[ModelDTOT]": ...
+    @overload
+    @staticmethod
+    def to_schema(data: "list[ModelT]", *, schema_type: None = None) -> "list[ModelT]": ...
+    @overload
+    @staticmethod
+    def to_schema(data: "ModelT") -> "ModelT": ...
+    @overload
+    @staticmethod
+    def to_schema(data: Any, *, schema_type: None = None) -> Any: ...
+
+    @staticmethod
+    def to_schema(data: Any, *, schema_type: "Optional[type[ModelDTOT]]" = None) -> Any:
+        """Convert data to a specified schema type.
+
+        Supports conversion to dataclasses, msgspec structs, Pydantic models, and attrs classes.
+        Handles both single objects and sequences.
+
+        Raises:
+            SQLSpecError if `schema_type` is not a valid type.
+
+        Returns:
+            Converted data in the specified schema type.
+
+        """
+        if schema_type is None:
+            return data
+        if is_dataclass(schema_type):
+            if isinstance(data, list):
+                return [schema_type(**dict(item) if hasattr(item, "keys") else item) for item in data]  # type: ignore[operator]
+            if hasattr(data, "keys"):
+                return schema_type(**dict(data))  # type: ignore[operator]
+            if isinstance(data, dict):
+                return schema_type(**data)  # type: ignore[operator]
+            # Fallback for other types
+            return data
+        if is_msgspec_struct(schema_type):
+            if not isinstance(data, Sequence):
+                return convert(
+                    obj=data,
+                    type=schema_type,
+                    from_attributes=True,
+                    dec_hook=partial(_default_msgspec_deserializer, type_decoders=_DEFAULT_TYPE_DECODERS),
+                )
+            return 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 is_pydantic_model(schema_type):
+            if not isinstance(data, Sequence):
+                return get_type_adapter(schema_type).validate_python(data, from_attributes=True)  # pyright: ignore
+            return get_type_adapter(list[schema_type]).validate_python(data, from_attributes=True)  # type: ignore[valid-type]  # pyright: ignore
+        if is_attrs_schema(schema_type):
+            if CATTRS_INSTALLED:
+                if isinstance(data, Sequence):
+                    return cattrs_structure(data, list[schema_type])  # type: ignore[valid-type]  # pyright: ignore
+                # If data is already structured (attrs instance), unstructure it first
+                if hasattr(data, "__attrs_attrs__"):
+                    data = cattrs_unstructure(data)
+                return cattrs_structure(data, schema_type)  # pyright: ignore
+            if isinstance(data, list):
+                return [schema_type(**dict(item) if hasattr(item, "keys") else attrs_asdict(item)) for item in data]
+            if hasattr(data, "keys"):
+                return schema_type(**dict(data))
+            if isinstance(data, dict):
+                return schema_type(**data)
+            # Fallback for other types
+            return data
+        msg = "`schema_type` should be a valid Dataclass, Pydantic model, Msgspec struct, or Attrs class"
+        raise SQLSpecError(msg)

sqlspec/driver/mixins/_sql_translator.py

@@ -0,0 +1,36 @@
+from mypy_extensions import trait
+from sqlglot import exp, parse_one
+from sqlglot.dialects.dialect import DialectType
+
+from sqlspec.core.statement import SQL, Statement
+from sqlspec.exceptions import SQLConversionError
+
+__all__ = ("SQLTranslatorMixin",)
+
+
+@trait
+class SQLTranslatorMixin:
+    """Mixin for drivers supporting SQL translation."""
+
+    __slots__ = ()
+
+    def convert_to_dialect(self, statement: "Statement", to_dialect: DialectType = None, pretty: bool = True) -> str:
+        if statement is not None and isinstance(statement, SQL):
+            if statement.expression is None:
+                msg = "Statement could not be parsed"
+                raise SQLConversionError(msg)
+            parsed_expression = statement.expression
+        elif isinstance(statement, exp.Expression):
+            parsed_expression = statement
+        else:
+            try:
+                parsed_expression = parse_one(statement, dialect=self.dialect)  # type: ignore[attr-defined]
+            except Exception as e:
+                error_msg = f"Failed to parse SQL statement: {e!s}"
+                raise SQLConversionError(error_msg) from e
+        target_dialect = to_dialect or self.dialect  # type: ignore[attr-defined]
+        try:
+            return parsed_expression.sql(dialect=target_dialect, pretty=pretty)
+        except Exception as e:
+            error_msg = f"Failed to convert SQL expression to {target_dialect}: {e!s}"
+            raise SQLConversionError(error_msg) from e