sqlspec 0.26.0__py3-none-any.whl → 0.28.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,8 +11,9 @@ 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
+from sqlspec.utils.type_guards import is_statement_filter
 
 if TYPE_CHECKING:
     from collections.abc import Sequence
@@ -31,6 +32,8 @@ __all__ = (
     "ExecutionResult",
     "ScriptExecutionResult",
     "VersionInfo",
+    "handle_single_row_error",
+    "make_cache_key_hashable",
 )
 
 
@@ -41,6 +44,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 +166,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 +194,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 +250,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 +260,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 +289,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 +306,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 +394,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.
 
@@ -349,48 +415,82 @@ class CommonDriverAttributesMixin:
         """
         kwargs = kwargs or {}
 
+        filters: list[StatementFilter] = []
+        data_parameters: list[StatementParameters] = []
+
+        for param in parameters:
+            if is_statement_filter(param):
+                filters.append(param)
+            else:
+                data_parameters.append(param)
+
         if isinstance(statement, QueryBuilder):
             sql_statement = statement.to_statement(statement_config)
-            if parameters or kwargs:
+            if data_parameters or kwargs:
                 merged_parameters = (
-                    (*sql_statement.positional_parameters, *parameters)
-                    if parameters
+                    (*sql_statement.positional_parameters, *tuple(data_parameters))
+                    if data_parameters
                     else sql_statement.positional_parameters
                 )
-                return SQL(sql_statement.sql, *merged_parameters, statement_config=statement_config, **kwargs)
+                sql_statement = SQL(sql_statement.sql, *merged_parameters, statement_config=statement_config, **kwargs)
+
+            for filter_obj in filters:
+                sql_statement = filter_obj.append_to_statement(sql_statement)
+
             return sql_statement
+
         if isinstance(statement, SQL):
-            if parameters or kwargs:
+            sql_statement = statement
+
+            if data_parameters or kwargs:
                 merged_parameters = (
-                    (*statement.positional_parameters, *parameters) if parameters else statement.positional_parameters
+                    (*sql_statement.positional_parameters, *tuple(data_parameters))
+                    if data_parameters
+                    else sql_statement.positional_parameters
                 )
-                return SQL(statement.sql, *merged_parameters, statement_config=statement_config, **kwargs)
-            needs_rebuild = False
-
-            if statement_config.dialect and (
-                not statement.statement_config.dialect or statement.statement_config.dialect != statement_config.dialect
-            ):
-                needs_rebuild = True
+                sql_statement = SQL(sql_statement.sql, *merged_parameters, statement_config=statement_config, **kwargs)
+            else:
+                needs_rebuild = False
+
+                if statement_config.dialect and (
+                    not sql_statement.statement_config.dialect
+                    or sql_statement.statement_config.dialect != statement_config.dialect
+                ):
+                    needs_rebuild = True
+
+                if (
+                    sql_statement.statement_config.parameter_config.default_execution_parameter_style
+                    != statement_config.parameter_config.default_execution_parameter_style
+                ):
+                    needs_rebuild = True
+
+                if needs_rebuild:
+                    sql_text = sql_statement.raw_sql or sql_statement.sql
+
+                    if sql_statement.is_many and sql_statement.parameters:
+                        sql_statement = SQL(
+                            sql_text, sql_statement.parameters, statement_config=statement_config, is_many=True
+                        )
+                    elif sql_statement.named_parameters:
+                        sql_statement = SQL(
+                            sql_text, statement_config=statement_config, **sql_statement.named_parameters
+                        )
+                    else:
+                        sql_statement = SQL(
+                            sql_text, *sql_statement.positional_parameters, statement_config=statement_config
+                        )
+
+            for filter_obj in filters:
+                sql_statement = filter_obj.append_to_statement(sql_statement)
 
-            if (
-                statement.statement_config.parameter_config.default_execution_parameter_style
-                != statement_config.parameter_config.default_execution_parameter_style
-            ):
-                needs_rebuild = True
+            return sql_statement
 
-            if needs_rebuild:
-                sql_text = statement.raw_sql or statement.sql
+        sql_statement = SQL(statement, *tuple(data_parameters), statement_config=statement_config, **kwargs)
 
-                if statement.is_many and statement.parameters:
-                    new_sql = SQL(sql_text, statement.parameters, statement_config=statement_config, is_many=True)
-                elif statement.named_parameters:
-                    new_sql = SQL(sql_text, statement_config=statement_config, **statement.named_parameters)
-                else:
-                    new_sql = SQL(sql_text, *statement.positional_parameters, statement_config=statement_config)
+        for filter_obj in filters:
+            sql_statement = filter_obj.append_to_statement(sql_statement)
 
-                return new_sql
-            return statement
-        return SQL(statement, *parameters, statement_config=statement_config, **kwargs)
+        return sql_statement
 
     def split_script_statements(
         self, script: str, statement_config: "StatementConfig", strip_trailing_semicolon: bool = False
@@ -421,7 +521,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 +549,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 +582,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 +604,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
@@ -625,26 +713,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 +743,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 +776,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: