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

sqlspec/driver/_common.py

@@ -2,7 +2,7 @@
 
 import re
 from contextlib import suppress
-from typing import TYPE_CHECKING, Any, Final, NamedTuple, Optional, TypeVar, Union, cast
+from typing import TYPE_CHECKING, Any, Final, NamedTuple, NoReturn, Optional, TypeVar, cast
 
 from mypy_extensions import trait
 from sqlglot import exp
@@ -11,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:
@@ -31,6 +31,8 @@ __all__ = (
     "ExecutionResult",
     "ScriptExecutionResult",
     "VersionInfo",
+    "handle_single_row_error",
+    "make_cache_key_hashable",
 )
 
 
@@ -41,6 +43,69 @@ 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."""
 
@@ -100,7 +165,7 @@ class VersionInfo:
 class DataDictionaryMixin:
     """Mixin providing common data dictionary functionality."""
 
-    def parse_version_string(self, version_str: str) -> "Optional[VersionInfo]":
+    def parse_version_string(self, version_str: str) -> "VersionInfo | None":
         """Parse version string into VersionInfo.
 
         Args:
@@ -128,7 +193,7 @@ class DataDictionaryMixin:
 
         return None
 
-    def detect_version_with_queries(self, driver: Any, queries: "list[str]") -> "Optional[VersionInfo]":
+    def detect_version_with_queries(self, driver: Any, queries: "list[str]") -> "VersionInfo | None":
         """Try multiple version queries to detect database version.
 
         Args:
@@ -184,7 +249,7 @@ 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
@@ -194,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
@@ -223,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.
 
@@ -240,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.
 
@@ -328,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.
 
@@ -421,7 +486,7 @@ class CommonDriverAttributesMixin:
         parameters: Any,
         statement_config: "StatementConfig",
         is_many: bool = False,
-        prepared_statement: Optional[Any] = None,  # pyright: ignore[reportUnusedParameter]
+        prepared_statement: Any | None = None,  # pyright: ignore[reportUnusedParameter]
     ) -> Any:
         """Prepare parameters for database driver consumption.
 
@@ -449,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.
 
@@ -465,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.
@@ -500,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 {k: apply_type_coercion(v) for k, v in parameters.items()}
+                return [self._apply_coercion(value, statement_config) for _, value in sorted_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 {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]
+        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
@@ -625,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:
@@ -655,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:
@@ -688,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: