sqlspec 0.25.0__py3-none-any.whl → 0.27.0__py3-none-any.whl

sqlspec/driver/_common.py

@@ -1,6 +1,8 @@
 """Common driver attributes and utilities."""
 
-from typing import TYPE_CHECKING, Any, Final, NamedTuple, Optional, Union, cast
+import re
+from contextlib import suppress
+from typing import TYPE_CHECKING, Any, Final, NamedTuple, NoReturn, Optional, TypeVar, cast
 
 from mypy_extensions import trait
 from sqlglot import exp
@@ -9,7 +11,7 @@ from sqlspec.builder import QueryBuilder
 from sqlspec.core import SQL, ParameterStyle, SQLResult, Statement, StatementConfig, TypedParameter
 from sqlspec.core.cache import CachedStatement, get_cache, get_cache_config
 from sqlspec.core.splitter import split_sql_script
-from sqlspec.exceptions import ImproperConfigurationError
+from sqlspec.exceptions import ImproperConfigurationError, NotFoundError
 from sqlspec.utils.logging import get_logger
 
 if TYPE_CHECKING:
@@ -25,19 +27,229 @@ __all__ = (
     "EXEC_ROWCOUNT_OVERRIDE",
     "EXEC_SPECIAL_DATA",
     "CommonDriverAttributesMixin",
+    "DataDictionaryMixin",
     "ExecutionResult",
     "ScriptExecutionResult",
+    "VersionInfo",
+    "handle_single_row_error",
+    "make_cache_key_hashable",
 )
 
 
 logger = get_logger("driver")
 
+DriverT = TypeVar("DriverT")
+VERSION_GROUPS_MIN_FOR_MINOR = 1
+VERSION_GROUPS_MIN_FOR_PATCH = 2
+
+
+def make_cache_key_hashable(obj: Any) -> Any:
+    """Recursively convert unhashable types to hashable ones for cache keys.
+
+    For array-like objects (NumPy arrays, Python arrays, etc.), we use structural
+    info (dtype + shape or typecode + length) rather than content for cache keys.
+    This ensures high cache hit rates for parameterized queries with different
+    vector values while avoiding expensive content hashing.
+
+    Args:
+        obj: Object to make hashable.
+
+    Returns:
+        A hashable representation of the object. Collections become tuples,
+        arrays become structural tuples like ("ndarray", dtype, shape).
+
+    Examples:
+        >>> make_cache_key_hashable([1, 2, 3])
+        (1, 2, 3)
+        >>> make_cache_key_hashable({"a": 1, "b": 2})
+        (('a', 1), ('b', 2))
+    """
+    if isinstance(obj, (list, tuple)):
+        return tuple(make_cache_key_hashable(item) for item in obj)
+    if isinstance(obj, dict):
+        return tuple(sorted((k, make_cache_key_hashable(v)) for k, v in obj.items()))
+    if isinstance(obj, set):
+        return frozenset(make_cache_key_hashable(item) for item in obj)
+
+    typecode = getattr(obj, "typecode", None)
+    if typecode is not None:
+        try:
+            length = len(obj)
+        except (AttributeError, TypeError):
+            return ("array", typecode)
+        else:
+            return ("array", typecode, length)
+
+    if hasattr(obj, "__array__"):
+        try:
+            dtype_str = getattr(obj.dtype, "str", str(type(obj)))
+            shape = tuple(int(s) for s in obj.shape)
+        except (AttributeError, TypeError):
+            try:
+                length = len(obj)
+            except (AttributeError, TypeError):
+                return ("array_like", type(obj).__name__)
+            else:
+                return ("array_like", type(obj).__name__, length)
+        else:
+            return ("ndarray", dtype_str, shape)
+    return obj
+
+
+def handle_single_row_error(error: ValueError) -> "NoReturn":
+    """Normalize single-row selection errors to SQLSpec exceptions."""
+
+    message = str(error)
+    if message.startswith("No result found"):
+        msg = "No rows found"
+        raise NotFoundError(msg) from error
+    raise error
+
+
+class VersionInfo:
+    """Database version information."""
+
+    def __init__(self, major: int, minor: int = 0, patch: int = 0) -> None:
+        """Initialize version info.
+
+        Args:
+            major: Major version number
+            minor: Minor version number
+            patch: Patch version number
+        """
+        self.major = major
+        self.minor = minor
+        self.patch = patch
+
+    @property
+    def version_tuple(self) -> "tuple[int, int, int]":
+        """Get version as tuple for comparison."""
+        return (self.major, self.minor, self.patch)
+
+    def __str__(self) -> str:
+        """String representation of version info."""
+        return f"{self.major}.{self.minor}.{self.patch}"
+
+    def __repr__(self) -> str:
+        """Detailed string representation."""
+        return f"VersionInfo({self.major}, {self.minor}, {self.patch})"
+
+    def __eq__(self, other: object) -> bool:
+        """Check version equality."""
+        if not isinstance(other, VersionInfo):
+            return NotImplemented
+        return self.version_tuple == other.version_tuple
+
+    def __lt__(self, other: "VersionInfo") -> bool:
+        """Check if this version is less than another."""
+        return self.version_tuple < other.version_tuple
+
+    def __le__(self, other: "VersionInfo") -> bool:
+        """Check if this version is less than or equal to another."""
+        return self.version_tuple <= other.version_tuple
+
+    def __gt__(self, other: "VersionInfo") -> bool:
+        """Check if this version is greater than another."""
+        return self.version_tuple > other.version_tuple
+
+    def __ge__(self, other: "VersionInfo") -> bool:
+        """Check if this version is greater than or equal to another."""
+        return self.version_tuple >= other.version_tuple
+
+    def __hash__(self) -> int:
+        """Make VersionInfo hashable based on version tuple."""
+        return hash(self.version_tuple)
+
+
+@trait
+class DataDictionaryMixin:
+    """Mixin providing common data dictionary functionality."""
+
+    def parse_version_string(self, version_str: str) -> "VersionInfo | None":
+        """Parse version string into VersionInfo.
+
+        Args:
+            version_str: Raw version string from database
+
+        Returns:
+            VersionInfo instance or None if parsing fails
+        """
+        # Try common version patterns
+        patterns = [
+            r"(\d+)\.(\d+)\.(\d+)",  # x.y.z
+            r"(\d+)\.(\d+)",  # x.y
+            r"(\d+)",  # x
+        ]
+
+        for pattern in patterns:
+            match = re.search(pattern, version_str)
+            if match:
+                groups = match.groups()
+
+                major = int(groups[0])
+                minor = int(groups[1]) if len(groups) > VERSION_GROUPS_MIN_FOR_MINOR else 0
+                patch = int(groups[2]) if len(groups) > VERSION_GROUPS_MIN_FOR_PATCH else 0
+                return VersionInfo(major, minor, patch)
+
+        return None
+
+    def detect_version_with_queries(self, driver: Any, queries: "list[str]") -> "VersionInfo | None":
+        """Try multiple version queries to detect database version.
+
+        Args:
+            driver: Database driver instance
+            queries: List of SQL queries to try
+
+        Returns:
+            Version information or None if detection fails
+        """
+        for query in queries:
+            with suppress(Exception):
+                result = driver.execute(query)
+                if result.data:
+                    version_str = str(result.data[0])
+                    if isinstance(result.data[0], dict):
+                        version_str = str(next(iter(result.data[0].values())))
+                    elif isinstance(result.data[0], (list, tuple)):
+                        version_str = str(result.data[0][0])
+
+                    parsed_version = self.parse_version_string(version_str)
+                    if parsed_version:
+                        logger.debug("Detected database version: %s", parsed_version)
+                        return parsed_version
+
+        logger.warning("Could not detect database version")
+        return None
+
+    def get_default_type_mapping(self) -> "dict[str, str]":
+        """Get default type mappings for common categories.
+
+        Returns:
+            Dictionary mapping type categories to generic SQL types
+        """
+        return {
+            "json": "TEXT",
+            "uuid": "VARCHAR(36)",
+            "boolean": "INTEGER",
+            "timestamp": "TIMESTAMP",
+            "text": "TEXT",
+            "blob": "BLOB",
+        }
+
+    def get_default_features(self) -> "list[str]":
+        """Get default feature flags supported by most databases.
+
+        Returns:
+            List of commonly supported feature names
+        """
+        return ["supports_transactions", "supports_prepared_statements"]
+
 
 class ScriptExecutionResult(NamedTuple):
     """Result from script execution with statement count information."""
 
     cursor_result: Any
-    rowcount_override: Optional[int]
+    rowcount_override: int | None
     special_data: Any
     statement_count: int
     successful_statements: int
@@ -47,23 +259,23 @@ class ExecutionResult(NamedTuple):
     """Execution result containing all data needed for SQLResult building."""
 
     cursor_result: Any
-    rowcount_override: Optional[int]
+    rowcount_override: int | None
     special_data: Any
     selected_data: Optional["list[dict[str, Any]]"]
     column_names: Optional["list[str]"]
-    data_row_count: Optional[int]
-    statement_count: Optional[int]
-    successful_statements: Optional[int]
+    data_row_count: int | None
+    statement_count: int | None
+    successful_statements: int | None
     is_script_result: bool
     is_select_result: bool
     is_many_result: bool
-    last_inserted_id: Optional[Union[int, str]] = None
+    last_inserted_id: int | str | None = None
 
 
 EXEC_CURSOR_RESULT: Final[int] = 0
 EXEC_ROWCOUNT_OVERRIDE: Final[int] = 1
 EXEC_SPECIAL_DATA: Final[int] = 2
-DEFAULT_EXECUTION_RESULT: Final[tuple[Any, Optional[int], Any]] = (None, None, None)
+DEFAULT_EXECUTION_RESULT: Final[tuple[Any, int | None, Any]] = (None, None, None)
 
 
 @trait
@@ -76,7 +288,7 @@ class CommonDriverAttributesMixin:
     driver_features: "dict[str, Any]"
 
     def __init__(
-        self, connection: "Any", statement_config: "StatementConfig", driver_features: "Optional[dict[str, Any]]" = None
+        self, connection: "Any", statement_config: "StatementConfig", driver_features: "dict[str, Any] | None" = None
     ) -> None:
         """Initialize driver adapter with connection and configuration.
 
@@ -93,17 +305,17 @@ class CommonDriverAttributesMixin:
         self,
         cursor_result: Any,
         *,
-        rowcount_override: Optional[int] = None,
+        rowcount_override: int | None = None,
         special_data: Any = None,
         selected_data: Optional["list[dict[str, Any]]"] = None,
         column_names: Optional["list[str]"] = None,
-        data_row_count: Optional[int] = None,
-        statement_count: Optional[int] = None,
-        successful_statements: Optional[int] = None,
+        data_row_count: int | None = None,
+        statement_count: int | None = None,
+        successful_statements: int | None = None,
         is_script_result: bool = False,
         is_select_result: bool = False,
         is_many_result: bool = False,
-        last_inserted_id: Optional[Union[int, str]] = None,
+        last_inserted_id: int | str | None = None,
     ) -> ExecutionResult:
         """Create ExecutionResult with all necessary data for any operation type.
 
@@ -181,11 +393,11 @@ class CommonDriverAttributesMixin:
 
     def prepare_statement(
         self,
-        statement: "Union[Statement, QueryBuilder]",
-        parameters: "tuple[Union[StatementParameters, StatementFilter], ...]" = (),
+        statement: "Statement | QueryBuilder",
+        parameters: "tuple[StatementParameters | StatementFilter, ...]" = (),
         *,
         statement_config: "StatementConfig",
-        kwargs: "Optional[dict[str, Any]]" = None,
+        kwargs: "dict[str, Any] | None" = None,
     ) -> "SQL":
         """Build SQL statement from various input types.
 
@@ -270,7 +482,11 @@ class CommonDriverAttributesMixin:
         ]
 
     def prepare_driver_parameters(
-        self, parameters: Any, statement_config: "StatementConfig", is_many: bool = False
+        self,
+        parameters: Any,
+        statement_config: "StatementConfig",
+        is_many: bool = False,
+        prepared_statement: Any | None = None,  # pyright: ignore[reportUnusedParameter]
     ) -> Any:
         """Prepare parameters for database driver consumption.
 
@@ -281,6 +497,7 @@ class CommonDriverAttributesMixin:
             parameters: Parameters in any format (dict, list, tuple, scalar, TypedParameter)
             statement_config: Statement configuration for parameter style detection
             is_many: If True, handle as executemany parameter sequence
+            prepared_statement: Optional prepared statement containing metadata for parameter processing
 
         Returns:
             Parameters with TypedParameter objects unwrapped to primitive values
@@ -297,6 +514,23 @@ class CommonDriverAttributesMixin:
             return [self._format_parameter_set_for_many(parameters, statement_config)]
         return self._format_parameter_set(parameters, statement_config)
 
+    def _apply_coercion(self, value: Any, statement_config: "StatementConfig") -> Any:
+        """Apply type coercion to a single value.
+
+        Args:
+            value: Value to coerce (may be TypedParameter or raw value)
+            statement_config: Statement configuration for type coercion map
+
+        Returns:
+            Coerced value with TypedParameter unwrapped
+        """
+        unwrapped_value = value.value if isinstance(value, TypedParameter) else value
+        if statement_config.parameter_config.type_coercion_map:
+            for type_check, converter in statement_config.parameter_config.type_coercion_map.items():
+                if isinstance(unwrapped_value, type_check):
+                    return converter(unwrapped_value)
+        return unwrapped_value
+
     def _format_parameter_set_for_many(self, parameters: Any, statement_config: "StatementConfig") -> Any:
         """Prepare a single parameter set for execute_many operations.
 
@@ -313,27 +547,14 @@ class CommonDriverAttributesMixin:
         if not parameters:
             return []
 
-        def apply_type_coercion(value: Any) -> Any:
-            """Apply type coercion to a single value."""
-            unwrapped_value = value.value if isinstance(value, TypedParameter) else value
-
-            if statement_config.parameter_config.type_coercion_map:
-                for type_check, converter in statement_config.parameter_config.type_coercion_map.items():
-                    if type_check in {list, tuple} and isinstance(unwrapped_value, (list, tuple)):
-                        continue
-                    if isinstance(unwrapped_value, type_check):
-                        return converter(unwrapped_value)
-
-            return unwrapped_value
+        if not isinstance(parameters, (dict, list, tuple)):
+            return self._apply_coercion(parameters, statement_config)
 
         if isinstance(parameters, dict):
-            return {k: apply_type_coercion(v) for k, v in parameters.items()}
-
-        if isinstance(parameters, (list, tuple)):
-            coerced_params = [apply_type_coercion(p) for p in parameters]
-            return tuple(coerced_params) if isinstance(parameters, tuple) else coerced_params
+            return {k: self._apply_coercion(v, statement_config) for k, v in parameters.items()}
 
-        return apply_type_coercion(parameters)
+        coerced_params = [self._apply_coercion(p, statement_config) for p in parameters]
+        return tuple(coerced_params) if isinstance(parameters, tuple) else coerced_params
 
     def _format_parameter_set(self, parameters: Any, statement_config: "StatementConfig") -> Any:
         """Prepare a single parameter set for database driver consumption.
@@ -348,50 +569,34 @@ class CommonDriverAttributesMixin:
         if not parameters:
             return []
 
-        def apply_type_coercion(value: Any) -> Any:
-            """Apply type coercion to a single value."""
-            unwrapped_value = value.value if isinstance(value, TypedParameter) else value
-
-            if statement_config.parameter_config.type_coercion_map:
-                for type_check, converter in statement_config.parameter_config.type_coercion_map.items():
-                    if isinstance(unwrapped_value, type_check):
-                        return converter(unwrapped_value)
-
-            return unwrapped_value
+        if not isinstance(parameters, (dict, list, tuple)):
+            return [self._apply_coercion(parameters, statement_config)]
 
         if isinstance(parameters, dict):
-            if not parameters:
-                return []
             if statement_config.parameter_config.supported_execution_parameter_styles and (
                 ParameterStyle.NAMED_PYFORMAT in statement_config.parameter_config.supported_execution_parameter_styles
                 or ParameterStyle.NAMED_COLON in statement_config.parameter_config.supported_execution_parameter_styles
             ):
-                return {k: apply_type_coercion(v) for k, v in parameters.items()}
+                return {k: self._apply_coercion(v, statement_config) for k, v in parameters.items()}
             if statement_config.parameter_config.default_parameter_style in {
                 ParameterStyle.NUMERIC,
                 ParameterStyle.QMARK,
                 ParameterStyle.POSITIONAL_PYFORMAT,
             }:
-                ordered_parameters = []
                 sorted_items = sorted(
                     parameters.items(),
                     key=lambda item: int(item[0])
                     if item[0].isdigit()
                     else (int(item[0][6:]) if item[0].startswith("param_") and item[0][6:].isdigit() else float("inf")),
                 )
-                for _, value in sorted_items:
-                    ordered_parameters.append(apply_type_coercion(value))
-                return ordered_parameters
+                return [self._apply_coercion(value, statement_config) for _, value in sorted_items]
 
-            return {k: apply_type_coercion(v) for k, v in parameters.items()}
+            return {k: self._apply_coercion(v, statement_config) for k, v in parameters.items()}
 
-        if isinstance(parameters, (list, tuple)):
-            coerced_params = [apply_type_coercion(p) for p in parameters]
-            if statement_config.parameter_config.preserve_parameter_format and isinstance(parameters, tuple):
-                return tuple(coerced_params)
-            return coerced_params
-
-        return [apply_type_coercion(parameters)]
+        coerced_params = [self._apply_coercion(p, statement_config) for p in parameters]
+        if statement_config.parameter_config.preserve_parameter_format and isinstance(parameters, tuple):
+            return tuple(coerced_params)
+        return coerced_params
 
     def _get_compiled_sql(
         self, statement: "SQL", statement_config: "StatementConfig", flatten_single_parameters: bool = False
@@ -422,7 +627,7 @@ class CommonDriverAttributesMixin:
         compiled_sql, execution_parameters = prepared_statement.compile()
 
         prepared_parameters = self.prepare_driver_parameters(
-            execution_parameters, statement_config, is_many=statement.is_many
+            execution_parameters, statement_config, is_many=statement.is_many, prepared_statement=statement
         )
 
         if statement_config.parameter_config.output_transformer:
@@ -438,7 +643,7 @@ class CommonDriverAttributesMixin:
                 if isinstance(prepared_parameters, list)
                 else (
                     prepared_parameters
-                    if prepared_parameters is None
+                    if prepared_parameters is None or isinstance(prepared_parameters, dict)
                     else (
                         tuple(prepared_parameters)
                         if not isinstance(prepared_parameters, tuple)
@@ -473,26 +678,26 @@ class CommonDriverAttributesMixin:
         )
 
         params = statement.parameters
-        params_key: Any
-
-        def make_hashable(obj: Any) -> Any:
-            """Recursively convert unhashable types to hashable ones."""
-            if isinstance(obj, (list, tuple)):
-                return tuple(make_hashable(item) for item in obj)
-            if isinstance(obj, dict):
-                return tuple(sorted((k, make_hashable(v)) for k, v in obj.items()))
-            if isinstance(obj, set):
-                return frozenset(make_hashable(item) for item in obj)
-            return obj
+
+        if params is None or (isinstance(params, (list, tuple, dict)) and not params):
+            return f"compiled:{hash(statement.sql)}:{context_hash}"
+
+        if isinstance(params, tuple) and all(isinstance(p, (int, str, bytes, bool, type(None))) for p in params):
+            try:
+                return (
+                    f"compiled:{hash((statement.sql, params, statement.is_many, statement.is_script))}:{context_hash}"
+                )
+            except TypeError:
+                pass
 
         try:
             if isinstance(params, dict):
-                params_key = make_hashable(params)
+                params_key = make_cache_key_hashable(params)
             elif isinstance(params, (list, tuple)) and params:
                 if isinstance(params[0], dict):
-                    params_key = tuple(make_hashable(d) for d in params)
+                    params_key = tuple(make_cache_key_hashable(d) for d in params)
                 else:
-                    params_key = make_hashable(params)
+                    params_key = make_cache_key_hashable(params)
             elif isinstance(params, (list, tuple)):
                 params_key = ()
             else:
@@ -503,7 +708,7 @@ class CommonDriverAttributesMixin:
         base_hash = hash((statement.sql, params_key, statement.is_many, statement.is_script))
         return f"compiled:{base_hash}:{context_hash}"
 
-    def _get_dominant_parameter_style(self, parameters: "list[Any]") -> "Optional[ParameterStyle]":
+    def _get_dominant_parameter_style(self, parameters: "list[Any]") -> "ParameterStyle | None":
         """Determine the dominant parameter style from parameter info list.
 
         Args:
@@ -536,7 +741,7 @@ class CommonDriverAttributesMixin:
     def find_filter(
         filter_type: "type[FilterTypeT]",
         filters: "Sequence[StatementFilter | StatementParameters] | Sequence[StatementFilter]",
-    ) -> "Optional[FilterTypeT]":
+    ) -> "FilterTypeT | None":
         """Get the filter specified by filter type from the filters.
 
         Args: