From 016bed640ae7407033f756d88235e52c0aab5340 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sat, 31 Oct 2026 21:01:28 +0000 Subject: [PATCH 001/174] refactor(mysql): remove unused helpers left from the mysql_common split Delete bool_to_int, resolve_column_names, detect_json_columns, and detect_json_columns_from_description from the aiomysql, asyncmy, mysqlconnector, and pymysql cores; the coercion maps register int directly. Drop the matching source-equivalence cases and the mysqlconnector create_mapped_exception exemption, since all four mappers are identical. --- sqlspec/adapters/aiomysql/core.py | 37 +---------------- sqlspec/adapters/asyncmy/core.py | 37 +---------------- sqlspec/adapters/mysqlconnector/core.py | 37 +---------------- sqlspec/adapters/pymysql/core.py | 37 +---------------- .../adapters/_shared/_driver_type_system.py | 40 ------------------- tests/unit/adapters/test_mysql_family_core.py | 8 ---- 6 files changed, 8 insertions(+), 188 deletions(-) diff --git a/sqlspec/adapters/aiomysql/core.py b/sqlspec/adapters/aiomysql/core.py index 005efd2ee..eebefa93d 100644 --- a/sqlspec/adapters/aiomysql/core.py +++ b/sqlspec/adapters/aiomysql/core.py @@ -28,7 +28,7 @@ from sqlspec.utils.serializers import from_json, to_json from sqlspec.utils.text import quote_backtick_identifier, split_qualified_identifier from sqlspec.utils.type_converters import build_uuid_coercions -from sqlspec.utils.type_guards import has_cursor_metadata, has_lastrowid, has_rowcount, has_sqlstate, has_type_code +from sqlspec.utils.type_guards import has_lastrowid, has_rowcount, has_sqlstate, has_type_code if TYPE_CHECKING: from sqlspec.core.parameters import ParameterValidator @@ -45,8 +45,6 @@ "collect_stream_rows", "create_mapped_exception", "default_statement_config", - "detect_json_columns", - "detect_json_columns_from_description", "driver_profile", "encode_records_for_local_infile", "escape_literal_percent", @@ -54,7 +52,6 @@ "normalize_execute_many_parameters", "normalize_execute_parameters", "normalize_lastrowid", - "resolve_column_names", "resolve_many_rowcount", "resolve_row_plan", "resolve_rowcount", @@ -207,7 +204,7 @@ async def _start(self) -> None: def build_profile() -> "DriverParameterProfile": """Create the aiomysql driver parameter profile.""" - coercions: dict[type, Callable[[Any], Any]] = {bool: bool_to_int, **build_uuid_coercions()} + coercions: dict[type, Callable[[Any], Any]] = {bool: int, **build_uuid_coercions()} return DriverParameterProfile( name="aiomysql", default_style=ParameterStyle.QMARK, @@ -345,13 +342,6 @@ def normalize_execute_many_parameters(parameters: Any) -> Any: return parameters -def resolve_column_names(description: "Sequence[Any] | None") -> "list[str]": - """Extract column names from cursor description.""" - if not description: - return [] - return [col[0] for col in description] - - def resolve_row_plan( description: "Sequence[Any] | None", json_type_codes: "set[int]" ) -> "tuple[list[str], list[int] | None]": @@ -377,24 +367,6 @@ def resolve_row_plan( return column_names, json_indexes or None -def detect_json_columns_from_description( - description: "Sequence[Any] | None", json_type_codes: "set[int]" -) -> "list[int]": - """Identify JSON column indexes from pre-fetched cursor description metadata.""" - return resolve_row_plan(description, json_type_codes)[1] or [] - - -def detect_json_columns( - cursor: Any, json_type_codes: "set[int]", description: "Sequence[Any] | None" = None -) -> "list[int]": - """Identify JSON column indexes from cursor metadata.""" - if description is None: - if not has_cursor_metadata(cursor): - return [] - description = cursor.description - return detect_json_columns_from_description(description, json_type_codes) - - def collect_stream_rows( rows: "Sequence[Any]", row_plan: "tuple[list[str], list[int] | None]", @@ -476,11 +448,6 @@ def normalize_lastrowid(cursor: Any) -> int | None: return last_id if isinstance(last_id, int) else None -def bool_to_int(value: bool) -> int: - """Coerce boolean values to integer for MySQL parameter compatibility.""" - return int(value) - - def create_mapped_exception(error: Any, *, logger: Any | None = None) -> "SQLSpecError | bool": """Map MySQL driver exceptions to SQLSpec errors.""" error_code = getattr(error, "errno", None) diff --git a/sqlspec/adapters/asyncmy/core.py b/sqlspec/adapters/asyncmy/core.py index 14bbed72e..087531f8f 100644 --- a/sqlspec/adapters/asyncmy/core.py +++ b/sqlspec/adapters/asyncmy/core.py @@ -28,7 +28,7 @@ from sqlspec.utils.serializers import from_json, to_json from sqlspec.utils.text import quote_backtick_identifier, split_qualified_identifier from sqlspec.utils.type_converters import build_uuid_coercions -from sqlspec.utils.type_guards import has_cursor_metadata, has_lastrowid, has_rowcount, has_sqlstate, has_type_code +from sqlspec.utils.type_guards import has_lastrowid, has_rowcount, has_sqlstate, has_type_code if TYPE_CHECKING: from sqlspec.core.parameters import ParameterValidator @@ -45,8 +45,6 @@ "collect_stream_rows", "create_mapped_exception", "default_statement_config", - "detect_json_columns", - "detect_json_columns_from_description", "driver_profile", "encode_records_for_local_infile", "escape_literal_percent", @@ -54,7 +52,6 @@ "normalize_execute_many_parameters", "normalize_execute_parameters", "normalize_lastrowid", - "resolve_column_names", "resolve_many_rowcount", "resolve_row_plan", "resolve_rowcount", @@ -207,7 +204,7 @@ async def _start(self) -> None: def build_profile() -> "DriverParameterProfile": """Create the AsyncMy driver parameter profile.""" - coercions: dict[type, Callable[[Any], Any]] = {bool: bool_to_int, **build_uuid_coercions()} + coercions: dict[type, Callable[[Any], Any]] = {bool: int, **build_uuid_coercions()} return DriverParameterProfile( name="asyncmy", default_style=ParameterStyle.QMARK, @@ -345,13 +342,6 @@ def normalize_execute_many_parameters(parameters: Any) -> Any: return parameters -def resolve_column_names(description: "Sequence[Any] | None") -> "list[str]": - """Extract column names from cursor description.""" - if not description: - return [] - return [col[0] for col in description] - - def resolve_row_plan( description: "Sequence[Any] | None", json_type_codes: "set[int]" ) -> "tuple[list[str], list[int] | None]": @@ -377,24 +367,6 @@ def resolve_row_plan( return column_names, json_indexes or None -def detect_json_columns_from_description( - description: "Sequence[Any] | None", json_type_codes: "set[int]" -) -> "list[int]": - """Identify JSON column indexes from pre-fetched cursor description metadata.""" - return resolve_row_plan(description, json_type_codes)[1] or [] - - -def detect_json_columns( - cursor: Any, json_type_codes: "set[int]", description: "Sequence[Any] | None" = None -) -> "list[int]": - """Identify JSON column indexes from cursor metadata.""" - if description is None: - if not has_cursor_metadata(cursor): - return [] - description = cursor.description - return detect_json_columns_from_description(description, json_type_codes) - - def collect_stream_rows( rows: "Sequence[Any]", row_plan: "tuple[list[str], list[int] | None]", @@ -476,11 +448,6 @@ def normalize_lastrowid(cursor: Any) -> int | None: return last_id if isinstance(last_id, int) else None -def bool_to_int(value: bool) -> int: - """Coerce boolean values to integer for MySQL parameter compatibility.""" - return int(value) - - def create_mapped_exception(error: Any, *, logger: Any | None = None) -> "SQLSpecError | bool": """Map MySQL driver exceptions to SQLSpec errors.""" error_code = getattr(error, "errno", None) diff --git a/sqlspec/adapters/mysqlconnector/core.py b/sqlspec/adapters/mysqlconnector/core.py index 5103ad39f..a7d924ed6 100644 --- a/sqlspec/adapters/mysqlconnector/core.py +++ b/sqlspec/adapters/mysqlconnector/core.py @@ -27,7 +27,7 @@ from sqlspec.utils.serializers import from_json, to_json from sqlspec.utils.text import quote_backtick_identifier, split_qualified_identifier from sqlspec.utils.type_converters import build_uuid_coercions -from sqlspec.utils.type_guards import has_cursor_metadata, has_lastrowid, has_rowcount, has_sqlstate, has_type_code +from sqlspec.utils.type_guards import has_lastrowid, has_rowcount, has_sqlstate, has_type_code if TYPE_CHECKING: from sqlspec.core.parameters import ParameterValidator @@ -45,8 +45,6 @@ "collect_stream_rows", "create_mapped_exception", "default_statement_config", - "detect_json_columns", - "detect_json_columns_from_description", "driver_profile", "encode_records_for_local_infile", "escape_literal_percent", @@ -54,7 +52,6 @@ "normalize_execute_many_parameters", "normalize_execute_parameters", "normalize_lastrowid", - "resolve_column_names", "resolve_many_rowcount", "resolve_row_plan", "resolve_rowcount", @@ -303,7 +300,7 @@ async def _start(self) -> None: def build_profile() -> "DriverParameterProfile": """Create the mysql-connector driver parameter profile.""" - coercions: dict[type, Callable[[Any], Any]] = {bool: bool_to_int, **build_uuid_coercions()} + coercions: dict[type, Callable[[Any], Any]] = {bool: int, **build_uuid_coercions()} return DriverParameterProfile( name="mysql-connector", default_style=ParameterStyle.QMARK, @@ -441,13 +438,6 @@ def normalize_execute_many_parameters(parameters: Any) -> Any: return parameters -def resolve_column_names(description: "Sequence[Any] | None") -> "list[str]": - """Extract column names from cursor description.""" - if not description: - return [] - return [col[0] for col in description] - - def resolve_row_plan( description: "Sequence[Any] | None", json_type_codes: "set[int]" ) -> "tuple[list[str], list[int] | None]": @@ -473,24 +463,6 @@ def resolve_row_plan( return column_names, json_indexes or None -def detect_json_columns_from_description( - description: "Sequence[Any] | None", json_type_codes: "set[int]" -) -> "list[int]": - """Identify JSON column indexes from pre-fetched cursor description metadata.""" - return resolve_row_plan(description, json_type_codes)[1] or [] - - -def detect_json_columns( - cursor: Any, json_type_codes: "set[int]", description: "Sequence[Any] | None" = None -) -> "list[int]": - """Identify JSON column indexes from cursor metadata.""" - if description is None: - if not has_cursor_metadata(cursor): - return [] - description = cursor.description - return detect_json_columns_from_description(description, json_type_codes) - - def collect_stream_rows( rows: "Sequence[Any]", row_plan: "tuple[list[str], list[int] | None]", @@ -572,11 +544,6 @@ def normalize_lastrowid(cursor: Any) -> int | None: return last_id if isinstance(last_id, int) else None -def bool_to_int(value: bool) -> int: - """Coerce boolean values to integer for MySQL parameter compatibility.""" - return int(value) - - def create_mapped_exception(error: Any, *, logger: Any | None = None) -> "SQLSpecError | bool": """Map MySQL driver exceptions to SQLSpec errors.""" error_code = getattr(error, "errno", None) diff --git a/sqlspec/adapters/pymysql/core.py b/sqlspec/adapters/pymysql/core.py index b89db729d..ea843c240 100644 --- a/sqlspec/adapters/pymysql/core.py +++ b/sqlspec/adapters/pymysql/core.py @@ -28,7 +28,7 @@ from sqlspec.utils.serializers import from_json, to_json from sqlspec.utils.text import quote_backtick_identifier, split_qualified_identifier from sqlspec.utils.type_converters import build_uuid_coercions -from sqlspec.utils.type_guards import has_cursor_metadata, has_lastrowid, has_rowcount, has_sqlstate, has_type_code +from sqlspec.utils.type_guards import has_lastrowid, has_rowcount, has_sqlstate, has_type_code if TYPE_CHECKING: from sqlspec.core.parameters import ParameterValidator @@ -45,8 +45,6 @@ "collect_stream_rows", "create_mapped_exception", "default_statement_config", - "detect_json_columns", - "detect_json_columns_from_description", "driver_profile", "encode_records_for_local_infile", "escape_literal_percent", @@ -54,7 +52,6 @@ "normalize_execute_many_parameters", "normalize_execute_parameters", "normalize_lastrowid", - "resolve_column_names", "resolve_many_rowcount", "resolve_row_plan", "resolve_rowcount", @@ -199,7 +196,7 @@ def close(self, error: bool = False) -> None: def build_profile() -> "DriverParameterProfile": """Create the PyMySQL driver parameter profile.""" - coercions: dict[type, Callable[[Any], Any]] = {bool: bool_to_int, **build_uuid_coercions()} + coercions: dict[type, Callable[[Any], Any]] = {bool: int, **build_uuid_coercions()} return DriverParameterProfile( name="PyMySQL", default_style=ParameterStyle.QMARK, @@ -338,13 +335,6 @@ def normalize_execute_many_parameters(parameters: Any) -> Any: return parameters -def resolve_column_names(description: "Sequence[Any] | None") -> "list[str]": - """Extract column names from cursor description.""" - if not description: - return [] - return [col[0] for col in description] - - def resolve_row_plan( description: "Sequence[Any] | None", json_type_codes: "set[int]" ) -> "tuple[list[str], list[int] | None]": @@ -370,24 +360,6 @@ def resolve_row_plan( return column_names, json_indexes or None -def detect_json_columns_from_description( - description: "Sequence[Any] | None", json_type_codes: "set[int]" -) -> "list[int]": - """Identify JSON column indexes from pre-fetched cursor description metadata.""" - return resolve_row_plan(description, json_type_codes)[1] or [] - - -def detect_json_columns( - cursor: Any, json_type_codes: "set[int]", description: "Sequence[Any] | None" = None -) -> "list[int]": - """Identify JSON column indexes from cursor metadata.""" - if description is None: - if not has_cursor_metadata(cursor): - return [] - description = cursor.description - return detect_json_columns_from_description(description, json_type_codes) - - def collect_stream_rows( rows: "Sequence[Any]", row_plan: "tuple[list[str], list[int] | None]", @@ -469,11 +441,6 @@ def normalize_lastrowid(cursor: Any) -> int | None: return last_id if isinstance(last_id, int) else None -def bool_to_int(value: bool) -> int: - """Coerce boolean values to integer for MySQL parameter compatibility.""" - return int(value) - - def create_mapped_exception(error: Any, *, logger: Any | None = None) -> "SQLSpecError | bool": """Map MySQL driver exceptions to SQLSpec errors.""" error_code = getattr(error, "errno", None) diff --git a/tests/integration/adapters/_shared/_driver_type_system.py b/tests/integration/adapters/_shared/_driver_type_system.py index 2f699c07a..f0bd54bc0 100644 --- a/tests/integration/adapters/_shared/_driver_type_system.py +++ b/tests/integration/adapters/_shared/_driver_type_system.py @@ -468,42 +468,6 @@ class SourceEquivalenceCase: ), allowlist_key="mysql_four_way:normalize_execute_many_parameters", ), - SourceEquivalenceCase( - case_id="mysql_four_way_resolve_column_names", - group="mysql_four_way", - symbol="resolve_column_names", - module_paths=( - "sqlspec.adapters.aiomysql.core", - "sqlspec.adapters.asyncmy.core", - "sqlspec.adapters.mysqlconnector.core", - "sqlspec.adapters.pymysql.core", - ), - allowlist_key="mysql_four_way:resolve_column_names", - ), - SourceEquivalenceCase( - case_id="mysql_four_way_detect_json_columns_from_description", - group="mysql_four_way", - symbol="detect_json_columns_from_description", - module_paths=( - "sqlspec.adapters.aiomysql.core", - "sqlspec.adapters.asyncmy.core", - "sqlspec.adapters.mysqlconnector.core", - "sqlspec.adapters.pymysql.core", - ), - allowlist_key="mysql_four_way:detect_json_columns_from_description", - ), - SourceEquivalenceCase( - case_id="mysql_four_way_detect_json_columns", - group="mysql_four_way", - symbol="detect_json_columns", - module_paths=( - "sqlspec.adapters.aiomysql.core", - "sqlspec.adapters.asyncmy.core", - "sqlspec.adapters.mysqlconnector.core", - "sqlspec.adapters.pymysql.core", - ), - allowlist_key="mysql_four_way:detect_json_columns", - ), SourceEquivalenceCase( case_id="mysql_four_way_resolve_row_plan", group="mysql_four_way", @@ -593,10 +557,6 @@ class SourceEquivalenceCase: case.allowlist_key: frozenset() for case in SOURCE_EQUIVALENCE_CASES } -SOURCE_EQUIVALENCE_ALLOWLIST["mysql_four_way:create_mapped_exception"] = frozenset({ - "sqlspec.adapters.mysqlconnector.core" -}) - def get_driver_feature_keys(adapter: str) -> tuple[str, ...]: """Return the canonical DriverFeatures keys declared by an adapter TypedDict.""" diff --git a/tests/unit/adapters/test_mysql_family_core.py b/tests/unit/adapters/test_mysql_family_core.py index 205747c82..367eddf48 100644 --- a/tests/unit/adapters/test_mysql_family_core.py +++ b/tests/unit/adapters/test_mysql_family_core.py @@ -54,14 +54,6 @@ def test_collect_rows_accepts_precomputed_row_plan(module_path: str) -> None: assert data[0]["payload"] == {"a": 1} -@pytest.mark.parametrize("module_path", MODULE_PATHS) -def test_detect_json_columns_uses_provided_description(module_path: str) -> None: - core = importlib.import_module(module_path) - description = [("id", 3), ("payload", 245)] - cursor = SimpleNamespace(description=description) - assert core.detect_json_columns(cursor, {245}, description=description) == [1] - - @pytest.mark.parametrize("module_path", MODULE_PATHS) def test_collect_rows_without_json_reuses_input_list(module_path: str) -> None: core = importlib.import_module(module_path) From bfec3832115a5f652a700e79163988b566fca4da Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sat, 31 Oct 2026 21:21:35 +0000 Subject: [PATCH 002/174] fix(mysql): map errno 1061/1091 to OperationalError and stop swallowing transaction errors The MySQL exception mapper returned True for duplicate-key-name and cannot-drop errors, so the driver handlers suppressed the exception without a result and statement execution failed with an AssertionError. Map both codes to OperationalError, make every handler record the mapped exception, and make begin/commit/rollback always raise the mapped error with its cause. --- sqlspec/adapters/aiomysql/core.py | 16 +++++--- sqlspec/adapters/aiomysql/driver.py | 24 ++--------- sqlspec/adapters/asyncmy/core.py | 16 +++++--- sqlspec/adapters/asyncmy/driver.py | 24 ++--------- sqlspec/adapters/mysqlconnector/core.py | 16 +++++--- sqlspec/adapters/mysqlconnector/driver.py | 40 ++++--------------- sqlspec/adapters/pymysql/core.py | 16 +++++--- sqlspec/adapters/pymysql/driver.py | 20 ++-------- .../adapters/test_mysql_family_exceptions.py | 35 +++++++++++++++- 9 files changed, 93 insertions(+), 114 deletions(-) diff --git a/sqlspec/adapters/aiomysql/core.py b/sqlspec/adapters/aiomysql/core.py index eebefa93d..e4ffb3ac3 100644 --- a/sqlspec/adapters/aiomysql/core.py +++ b/sqlspec/adapters/aiomysql/core.py @@ -17,6 +17,7 @@ ForeignKeyViolationError, IntegrityError, NotNullViolationError, + OperationalError, PermissionDeniedError, QueryTimeoutError, SQLParsingError, @@ -77,7 +78,10 @@ MYSQL_SYNTAX_ERROR_MIN: Final[int] = 1064 MYSQL_SYNTAX_ERROR_MAX_EXCLUSIVE: Final[int] = 1100 -_MYSQL_MIGRATION_ERROR_CODES: Final[frozenset[int]] = frozenset((1061, 1091)) +_MYSQL_DDL_ERROR_DISPATCH: dict[int, tuple[type[SQLSpecError], str]] = { + 1061: (OperationalError, "duplicate key name"), + 1091: (OperationalError, "cannot drop column or key"), +} _MYSQL_SQLSTATE_EXACT_DISPATCH: Final[dict[str, tuple[type[SQLSpecError], str]]] = { "23505": (UniqueViolationError, "unique constraint violation"), "23503": (ForeignKeyViolationError, "foreign key constraint violation"), @@ -448,8 +452,9 @@ def normalize_lastrowid(cursor: Any) -> int | None: return last_id if isinstance(last_id, int) else None -def create_mapped_exception(error: Any, *, logger: Any | None = None) -> "SQLSpecError | bool": +def create_mapped_exception(error: Any, *, logger: Any | None = None) -> SQLSpecError: """Map MySQL driver exceptions to SQLSpec errors.""" + _ = logger error_code = getattr(error, "errno", None) if error_code is None and isinstance(error, Exception) and error.args: value = error.args[0] @@ -458,10 +463,9 @@ def create_mapped_exception(error: Any, *, logger: Any | None = None) -> "SQLSpe sqlstate = error.sqlstate if has_sqlstate(error) else None sqlstate_prefix = sqlstate[:2] if isinstance(sqlstate, str) and sqlstate else None - if error_code in _MYSQL_MIGRATION_ERROR_CODES: - if logger is not None: - logger.warning("MySQL expected migration error (ignoring): %s", error) - return True + dispatch = _MYSQL_DDL_ERROR_DISPATCH.get(error_code) if error_code is not None else None + if dispatch is not None: + return _create_mysql_error(error, sqlstate, error_code, dispatch[0], dispatch[1]) dispatch = _MYSQL_SQLSTATE_EXACT_DISPATCH.get(sqlstate) if sqlstate is not None else None if dispatch is not None: diff --git a/sqlspec/adapters/aiomysql/driver.py b/sqlspec/adapters/aiomysql/driver.py index f6299971d..496fcbcd7 100644 --- a/sqlspec/adapters/aiomysql/driver.py +++ b/sqlspec/adapters/aiomysql/driver.py @@ -44,7 +44,6 @@ from sqlspec.adapters.aiomysql.data_dictionary import AiomysqlDataDictionary from sqlspec.core import ArrowResult, get_cache_config, register_driver_profile from sqlspec.driver import AsyncDriverAdapterBase, AsyncRowStream, BaseAsyncExceptionHandler -from sqlspec.exceptions import SQLSpecError from sqlspec.utils.logging import get_logger from sqlspec.utils.serializers import from_json from sqlspec.utils.type_guards import supports_json_type @@ -98,10 +97,7 @@ def _handle_exception(self, exc_type: "type[BaseException] | None", exc_val: "Ba if exc_type is None: return False if issubclass(exc_type, AiomysqlPymysqlError): - result = create_mapped_exception(exc_val, logger=logger) - if result is True: - return True - self.pending_exception = cast("Exception", result) + self.pending_exception = create_mapped_exception(exc_val) return True return False @@ -250,11 +246,7 @@ async def begin(self) -> None: try: await self.connection.begin() except AiomysqlPymysqlMySQLError as e: - mapped = create_mapped_exception(e, logger=logger) - if isinstance(mapped, Exception): - raise mapped from e - msg = f"Failed to begin MySQL transaction: {e}" - raise SQLSpecError(msg) from e + raise create_mapped_exception(e) from e async def commit(self) -> None: """Commit the current transaction. @@ -265,11 +257,7 @@ async def commit(self) -> None: try: await self.connection.commit() except AiomysqlPymysqlMySQLError as e: - mapped = create_mapped_exception(e, logger=logger) - if isinstance(mapped, Exception): - raise mapped from e - msg = f"Failed to commit MySQL transaction: {e}" - raise SQLSpecError(msg) from e + raise create_mapped_exception(e) from e async def rollback(self) -> None: """Rollback the current transaction. @@ -280,11 +268,7 @@ async def rollback(self) -> None: try: await self.connection.rollback() except AiomysqlPymysqlMySQLError as e: - mapped = create_mapped_exception(e, logger=logger) - if isinstance(mapped, Exception): - raise mapped from e - msg = f"Failed to rollback MySQL transaction: {e}" - raise SQLSpecError(msg) from e + raise create_mapped_exception(e) from e def with_cursor(self, connection: "AiomysqlConnection") -> "AiomysqlCursor": """Create cursor context manager for the connection. diff --git a/sqlspec/adapters/asyncmy/core.py b/sqlspec/adapters/asyncmy/core.py index 087531f8f..3fb570e78 100644 --- a/sqlspec/adapters/asyncmy/core.py +++ b/sqlspec/adapters/asyncmy/core.py @@ -17,6 +17,7 @@ ForeignKeyViolationError, IntegrityError, NotNullViolationError, + OperationalError, PermissionDeniedError, QueryTimeoutError, SQLParsingError, @@ -77,7 +78,10 @@ MYSQL_SYNTAX_ERROR_MIN: Final[int] = 1064 MYSQL_SYNTAX_ERROR_MAX_EXCLUSIVE: Final[int] = 1100 -_MYSQL_MIGRATION_ERROR_CODES: Final[frozenset[int]] = frozenset((1061, 1091)) +_MYSQL_DDL_ERROR_DISPATCH: dict[int, tuple[type[SQLSpecError], str]] = { + 1061: (OperationalError, "duplicate key name"), + 1091: (OperationalError, "cannot drop column or key"), +} _MYSQL_SQLSTATE_EXACT_DISPATCH: Final[dict[str, tuple[type[SQLSpecError], str]]] = { "23505": (UniqueViolationError, "unique constraint violation"), "23503": (ForeignKeyViolationError, "foreign key constraint violation"), @@ -448,8 +452,9 @@ def normalize_lastrowid(cursor: Any) -> int | None: return last_id if isinstance(last_id, int) else None -def create_mapped_exception(error: Any, *, logger: Any | None = None) -> "SQLSpecError | bool": +def create_mapped_exception(error: Any, *, logger: Any | None = None) -> SQLSpecError: """Map MySQL driver exceptions to SQLSpec errors.""" + _ = logger error_code = getattr(error, "errno", None) if error_code is None and isinstance(error, Exception) and error.args: value = error.args[0] @@ -458,10 +463,9 @@ def create_mapped_exception(error: Any, *, logger: Any | None = None) -> "SQLSpe sqlstate = error.sqlstate if has_sqlstate(error) else None sqlstate_prefix = sqlstate[:2] if isinstance(sqlstate, str) and sqlstate else None - if error_code in _MYSQL_MIGRATION_ERROR_CODES: - if logger is not None: - logger.warning("MySQL expected migration error (ignoring): %s", error) - return True + dispatch = _MYSQL_DDL_ERROR_DISPATCH.get(error_code) if error_code is not None else None + if dispatch is not None: + return _create_mysql_error(error, sqlstate, error_code, dispatch[0], dispatch[1]) dispatch = _MYSQL_SQLSTATE_EXACT_DISPATCH.get(sqlstate) if sqlstate is not None else None if dispatch is not None: diff --git a/sqlspec/adapters/asyncmy/driver.py b/sqlspec/adapters/asyncmy/driver.py index ca983e94f..fd96c28ae 100644 --- a/sqlspec/adapters/asyncmy/driver.py +++ b/sqlspec/adapters/asyncmy/driver.py @@ -39,7 +39,6 @@ from sqlspec.adapters.asyncmy.data_dictionary import AsyncmyDataDictionary from sqlspec.core import ArrowResult, get_cache_config, register_driver_profile from sqlspec.driver import AsyncDriverAdapterBase, AsyncRowStream, BaseAsyncExceptionHandler -from sqlspec.exceptions import SQLSpecError from sqlspec.utils.logging import get_logger from sqlspec.utils.serializers import from_json from sqlspec.utils.type_guards import supports_json_type @@ -93,10 +92,7 @@ def _handle_exception(self, exc_type: "type[BaseException] | None", exc_val: "Ba if exc_type is None: return False if issubclass(exc_type, cast("type[BaseException]", AsyncmyError)): - result = create_mapped_exception(exc_val, logger=logger) - if result is True: - return True - self.pending_exception = cast("Exception", result) + self.pending_exception = create_mapped_exception(exc_val) return True return False @@ -256,11 +252,7 @@ async def begin(self) -> None: try: await self.connection.begin() except AsyncmyMySQLError as e: - mapped = create_mapped_exception(e, logger=logger) - if isinstance(mapped, Exception): - raise mapped from e - msg = f"Failed to begin MySQL transaction: {e}" - raise SQLSpecError(msg) from e + raise create_mapped_exception(e) from e async def commit(self) -> None: """Commit the current transaction. @@ -271,11 +263,7 @@ async def commit(self) -> None: try: await self.connection.commit() except AsyncmyMySQLError as e: - mapped = create_mapped_exception(e, logger=logger) - if isinstance(mapped, Exception): - raise mapped from e - msg = f"Failed to commit MySQL transaction: {e}" - raise SQLSpecError(msg) from e + raise create_mapped_exception(e) from e async def rollback(self) -> None: """Rollback the current transaction. @@ -286,11 +274,7 @@ async def rollback(self) -> None: try: await self.connection.rollback() except AsyncmyMySQLError as e: - mapped = create_mapped_exception(e, logger=logger) - if isinstance(mapped, Exception): - raise mapped from e - msg = f"Failed to rollback MySQL transaction: {e}" - raise SQLSpecError(msg) from e + raise create_mapped_exception(e) from e def with_cursor(self, connection: "AsyncmyConnection") -> "AsyncmyCursor": """Create cursor context manager for the connection. diff --git a/sqlspec/adapters/mysqlconnector/core.py b/sqlspec/adapters/mysqlconnector/core.py index a7d924ed6..da3cd088d 100644 --- a/sqlspec/adapters/mysqlconnector/core.py +++ b/sqlspec/adapters/mysqlconnector/core.py @@ -16,6 +16,7 @@ ForeignKeyViolationError, IntegrityError, NotNullViolationError, + OperationalError, PermissionDeniedError, QueryTimeoutError, SQLParsingError, @@ -77,7 +78,10 @@ MYSQL_SYNTAX_ERROR_MIN: Final[int] = 1064 MYSQL_SYNTAX_ERROR_MAX_EXCLUSIVE: Final[int] = 1100 -_MYSQL_MIGRATION_ERROR_CODES: Final[frozenset[int]] = frozenset((1061, 1091)) +_MYSQL_DDL_ERROR_DISPATCH: dict[int, tuple[type[SQLSpecError], str]] = { + 1061: (OperationalError, "duplicate key name"), + 1091: (OperationalError, "cannot drop column or key"), +} _MYSQL_SQLSTATE_EXACT_DISPATCH: Final[dict[str, tuple[type[SQLSpecError], str]]] = { "23505": (UniqueViolationError, "unique constraint violation"), "23503": (ForeignKeyViolationError, "foreign key constraint violation"), @@ -544,8 +548,9 @@ def normalize_lastrowid(cursor: Any) -> int | None: return last_id if isinstance(last_id, int) else None -def create_mapped_exception(error: Any, *, logger: Any | None = None) -> "SQLSpecError | bool": +def create_mapped_exception(error: Any, *, logger: Any | None = None) -> SQLSpecError: """Map MySQL driver exceptions to SQLSpec errors.""" + _ = logger error_code = getattr(error, "errno", None) if error_code is None and isinstance(error, Exception) and error.args: value = error.args[0] @@ -554,10 +559,9 @@ def create_mapped_exception(error: Any, *, logger: Any | None = None) -> "SQLSpe sqlstate = error.sqlstate if has_sqlstate(error) else None sqlstate_prefix = sqlstate[:2] if isinstance(sqlstate, str) and sqlstate else None - if error_code in _MYSQL_MIGRATION_ERROR_CODES: - if logger is not None: - logger.warning("MySQL expected migration error (ignoring): %s", error) - return True + dispatch = _MYSQL_DDL_ERROR_DISPATCH.get(error_code) if error_code is not None else None + if dispatch is not None: + return _create_mysql_error(error, sqlstate, error_code, dispatch[0], dispatch[1]) dispatch = _MYSQL_SQLSTATE_EXACT_DISPATCH.get(sqlstate) if sqlstate is not None else None if dispatch is not None: diff --git a/sqlspec/adapters/mysqlconnector/driver.py b/sqlspec/adapters/mysqlconnector/driver.py index b19cb450f..aa043ce6b 100644 --- a/sqlspec/adapters/mysqlconnector/driver.py +++ b/sqlspec/adapters/mysqlconnector/driver.py @@ -105,10 +105,7 @@ def _handle_exception(self, exc_type: "type[BaseException] | None", exc_val: "Ba if exc_type is None: return False if issubclass(exc_type, MysqlConnectorError): - result = create_mapped_exception(exc_val, logger=logger) - if result is True: - return True - self.pending_exception = cast("Exception", result) + self.pending_exception = create_mapped_exception(exc_val) return True return False @@ -200,28 +197,19 @@ def begin(self) -> None: with MysqlConnectorSyncCursor(self.connection) as cursor: cursor.execute("BEGIN") except MysqlConnectorError as e: - mapped = create_mapped_exception(e, logger=logger) - if isinstance(mapped, bool): - return - raise mapped from e + raise create_mapped_exception(e) from e def commit(self) -> None: try: self.connection.commit() except MysqlConnectorError as e: - mapped = create_mapped_exception(e, logger=logger) - if isinstance(mapped, bool): - return - raise mapped from e + raise create_mapped_exception(e) from e def rollback(self) -> None: try: self.connection.rollback() except MysqlConnectorError as e: - mapped = create_mapped_exception(e, logger=logger) - if isinstance(mapped, bool): - return - raise mapped from e + raise create_mapped_exception(e) from e def with_cursor(self, connection: "MysqlConnectorSyncConnection") -> "MysqlConnectorSyncCursor": cursor_options = cast("dict[str, Any]", self.driver_features.get("cursor_options") or {}) @@ -387,10 +375,7 @@ def _handle_exception(self, exc_type: "type[BaseException] | None", exc_val: "Ba if exc_type is None: return False if issubclass(exc_type, MysqlConnectorError): - result = create_mapped_exception(exc_val, logger=logger) - if result is True: - return True - self.pending_exception = cast("Exception", result) + self.pending_exception = create_mapped_exception(exc_val) return True return False @@ -482,28 +467,19 @@ async def begin(self) -> None: async with MysqlConnectorAsyncCursor(self.connection) as cursor: await cursor.execute("BEGIN") except MysqlConnectorError as e: - mapped = create_mapped_exception(e, logger=logger) - if isinstance(mapped, bool): - return - raise mapped from e + raise create_mapped_exception(e) from e async def commit(self) -> None: try: await self.connection.commit() except MysqlConnectorError as e: - mapped = create_mapped_exception(e, logger=logger) - if isinstance(mapped, bool): - return - raise mapped from e + raise create_mapped_exception(e) from e async def rollback(self) -> None: try: await self.connection.rollback() except MysqlConnectorError as e: - mapped = create_mapped_exception(e, logger=logger) - if isinstance(mapped, bool): - return - raise mapped from e + raise create_mapped_exception(e) from e def with_cursor(self, connection: "MysqlConnectorAsyncConnection") -> "MysqlConnectorAsyncCursor": cursor_options = cast("dict[str, Any]", self.driver_features.get("cursor_options") or {}) diff --git a/sqlspec/adapters/pymysql/core.py b/sqlspec/adapters/pymysql/core.py index ea843c240..219a2b81c 100644 --- a/sqlspec/adapters/pymysql/core.py +++ b/sqlspec/adapters/pymysql/core.py @@ -17,6 +17,7 @@ ForeignKeyViolationError, IntegrityError, NotNullViolationError, + OperationalError, PermissionDeniedError, QueryTimeoutError, SQLParsingError, @@ -77,7 +78,10 @@ MYSQL_SYNTAX_ERROR_MIN: Final[int] = 1064 MYSQL_SYNTAX_ERROR_MAX_EXCLUSIVE: Final[int] = 1100 -_MYSQL_MIGRATION_ERROR_CODES: Final[frozenset[int]] = frozenset((1061, 1091)) +_MYSQL_DDL_ERROR_DISPATCH: dict[int, tuple[type[SQLSpecError], str]] = { + 1061: (OperationalError, "duplicate key name"), + 1091: (OperationalError, "cannot drop column or key"), +} _MYSQL_SQLSTATE_EXACT_DISPATCH: Final[dict[str, tuple[type[SQLSpecError], str]]] = { "23505": (UniqueViolationError, "unique constraint violation"), "23503": (ForeignKeyViolationError, "foreign key constraint violation"), @@ -441,8 +445,9 @@ def normalize_lastrowid(cursor: Any) -> int | None: return last_id if isinstance(last_id, int) else None -def create_mapped_exception(error: Any, *, logger: Any | None = None) -> "SQLSpecError | bool": +def create_mapped_exception(error: Any, *, logger: Any | None = None) -> SQLSpecError: """Map MySQL driver exceptions to SQLSpec errors.""" + _ = logger error_code = getattr(error, "errno", None) if error_code is None and isinstance(error, Exception) and error.args: value = error.args[0] @@ -451,10 +456,9 @@ def create_mapped_exception(error: Any, *, logger: Any | None = None) -> "SQLSpe sqlstate = error.sqlstate if has_sqlstate(error) else None sqlstate_prefix = sqlstate[:2] if isinstance(sqlstate, str) and sqlstate else None - if error_code in _MYSQL_MIGRATION_ERROR_CODES: - if logger is not None: - logger.warning("MySQL expected migration error (ignoring): %s", error) - return True + dispatch = _MYSQL_DDL_ERROR_DISPATCH.get(error_code) if error_code is not None else None + if dispatch is not None: + return _create_mysql_error(error, sqlstate, error_code, dispatch[0], dispatch[1]) dispatch = _MYSQL_SQLSTATE_EXACT_DISPATCH.get(sqlstate) if sqlstate is not None else None if dispatch is not None: diff --git a/sqlspec/adapters/pymysql/driver.py b/sqlspec/adapters/pymysql/driver.py index 47f15d481..871639b8c 100644 --- a/sqlspec/adapters/pymysql/driver.py +++ b/sqlspec/adapters/pymysql/driver.py @@ -83,10 +83,7 @@ def _handle_exception(self, exc_type: "type[BaseException] | None", exc_val: "Ba if exc_type is None: return False if issubclass(exc_type, PyMysqlMySQLError): - result = create_mapped_exception(exc_val, logger=logger) - if result is True: - return True - self.pending_exception = cast("Exception", result) + self.pending_exception = create_mapped_exception(exc_val) return True return False @@ -187,28 +184,19 @@ def begin(self) -> None: try: self.connection.begin() except PyMysqlMySQLError as exc: - mapped = create_mapped_exception(exc, logger=logger) - if isinstance(mapped, bool): - return - raise mapped from exc + raise create_mapped_exception(exc) from exc def commit(self) -> None: try: self.connection.commit() except PyMysqlMySQLError as exc: - mapped = create_mapped_exception(exc, logger=logger) - if isinstance(mapped, bool): - return - raise mapped from exc + raise create_mapped_exception(exc) from exc def rollback(self) -> None: try: self.connection.rollback() except PyMysqlMySQLError as exc: - mapped = create_mapped_exception(exc, logger=logger) - if isinstance(mapped, bool): - return - raise mapped from exc + raise create_mapped_exception(exc) from exc def with_cursor(self, connection: "PyMysqlConnection") -> "PyMysqlCursor": return PyMysqlCursor(connection) diff --git a/tests/unit/adapters/test_mysql_family_exceptions.py b/tests/unit/adapters/test_mysql_family_exceptions.py index 47ad47e68..6d63380a5 100644 --- a/tests/unit/adapters/test_mysql_family_exceptions.py +++ b/tests/unit/adapters/test_mysql_family_exceptions.py @@ -13,6 +13,7 @@ ForeignKeyViolationError, IntegrityError, NotNullViolationError, + OperationalError, PermissionDeniedError, QueryTimeoutError, SQLParsingError, @@ -92,7 +93,37 @@ def test_mysql_create_mapped_exception_preserves_legacy_dispatch_priority(module @pytest.mark.parametrize("module_path", MODULE_PATHS) -def test_mysql_create_mapped_exception_suppresses_expected_migration_codes(module_path: str) -> None: +@pytest.mark.parametrize("code", [1061, 1091]) +def test_mysql_create_mapped_exception_maps_ddl_codes_to_operational_error(module_path: str, code: int) -> None: core = importlib.import_module(module_path) - assert core.create_mapped_exception(_make_error(1091)) is True + result = core.create_mapped_exception(_make_error(code, sqlstate="42000")) + + assert isinstance(result, OperationalError) + + +HANDLER_CASES = ( + ("sqlspec.adapters.aiomysql.driver", "AiomysqlExceptionHandler", "AiomysqlPymysqlError"), + ("sqlspec.adapters.asyncmy.driver", "AsyncmyExceptionHandler", "AsyncmyError"), + ("sqlspec.adapters.mysqlconnector.driver", "MysqlConnectorSyncExceptionHandler", "MysqlConnectorError"), + ("sqlspec.adapters.pymysql.driver", "PyMysqlExceptionHandler", "PyMysqlMySQLError"), +) + + +@pytest.mark.parametrize(("module_path", "handler_name", "error_name"), HANDLER_CASES) +def test_mysql_handler_records_mapped_ddl_error(module_path: str, handler_name: str, error_name: str) -> None: + module = importlib.import_module(module_path) + native_error_cls = getattr(module, error_name) + if not isinstance(native_error_cls, type): + pytest.skip(f"{error_name} driver is not installed") + error = ( + native_error_cls(msg="Duplicate key name 'idx'", errno=1061) + if error_name == "MysqlConnectorError" + else native_error_cls(1061, "Duplicate key name 'idx'") + ) + handler = getattr(module, handler_name)() + + suppressed = handler._handle_exception(type(error), error) + + assert suppressed is True + assert isinstance(handler.pending_exception, OperationalError) From 9f2450e0bb2aedab73d63f25fa2f9184b0f62fc2 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sat, 31 Oct 2026 21:30:58 +0000 Subject: [PATCH 003/174] fix(adapters): align store class names with the Litestar session migration The Litestar session migration derives the store class from the config class name, so SqliteConfig, AdbcConfig, and DuckDBConfig failed to resolve SQLiteStore, ADBCStore, and DuckdbStore. Rename those stores, the DuckDB ADK classes, and the single-mode Sybase and SQL Anywhere event stores to match their configs, drop the old event store aliases, and add a guard test covering every adapter with a Litestar store. --- docs/changelog.md | 2 +- docs/integrations/litestar.md | 2 +- docs/reference/adapters/adbc.md | 2 +- docs/reference/adapters/duckdb.md | 10 +++--- docs/reference/adapters/sqlanywhere.md | 2 +- docs/reference/adapters/sqlite.md | 2 +- docs/reference/adapters/sybase.md | 2 +- sqlspec/adapters/adbc/litestar/__init__.py | 4 +-- sqlspec/adapters/adbc/litestar/store.py | 4 +-- sqlspec/adapters/duckdb/adk/__init__.py | 4 +-- sqlspec/adapters/duckdb/adk/store.py | 20 ++++++------ sqlspec/adapters/duckdb/litestar/__init__.py | 4 +-- sqlspec/adapters/duckdb/litestar/store.py | 4 +-- .../adapters/sqlanywhere/events/__init__.py | 4 +-- sqlspec/adapters/sqlanywhere/events/store.py | 7 ++-- sqlspec/adapters/sqlite/litestar/__init__.py | 4 +-- sqlspec/adapters/sqlite/litestar/store.py | 4 +-- sqlspec/adapters/sybase/events/__init__.py | 4 +-- sqlspec/adapters/sybase/events/store.py | 7 ++-- .../duckdb/extensions/adk/test_maintenance.py | 10 +++--- .../extensions/adk/test_memory_store.py | 10 +++--- .../duckdb/extensions/adk/test_store.py | 30 ++++++++--------- .../duckdb/extensions/litestar/test_store.py | 8 ++--- .../adbc/extensions/litestar/test_store.py | 8 ++--- .../sqlite/extensions/litestar/test_store.py | 12 +++---- tests/integration/fixtures/adapter_cases.py | 30 ++++++++--------- .../adapters/test_adapter_implementations.py | 30 +++++++++++++++++ .../test_adbc/test_adk_dialect_support.py | 6 ++-- .../test_adbc/test_extension_stores.py | 10 +++--- .../adapters/test_duckdb/test_adk_store.py | 32 +++++++++---------- tests/unit/adapters/test_sqlanywhere.py | 8 ++--- tests/unit/adapters/test_sybase.py | 8 ++--- .../test_adk/test_store_instantiation.py | 8 ++--- .../test_extension_config_boundaries.py | 4 +-- .../extensions/test_schema_reconciliation.py | 6 ++-- tests/unit/extensions/test_storage_options.py | 4 +-- 36 files changed, 170 insertions(+), 146 deletions(-) diff --git a/docs/changelog.md b/docs/changelog.md index d077efe62..e972e65e1 100644 --- a/docs/changelog.md +++ b/docs/changelog.md @@ -260,7 +260,7 @@ important operational fixes. savepoints, SQLSTATE/vendor exception mapping, and default `enable_lowercase_column_names=True` normalization for Oracle and Db2. ([#839](https://github.com/litestar-org/sqlspec/pull/839)) -* `ADBCStore` (Litestar sessions), `AdbcEventQueueStore` (event channels), +* `AdbcStore` (Litestar sessions), `AdbcEventQueueStore` (event channels), `AdbcADKStore`, and `AdbcADKMemoryStore` (Google ADK sessions and memory) add Oracle (`oracle`) and IBM Db2 (`db2`) backend support with idempotent DDL guards (PL/SQL `SQLCODE` blocks for Oracle and `SYSCAT` catalog probes diff --git a/docs/integrations/litestar.md b/docs/integrations/litestar.md index e150a520e..be6acab18 100644 --- a/docs/integrations/litestar.md +++ b/docs/integrations/litestar.md @@ -120,7 +120,7 @@ For bidirectional keyset pagination (`CursorPagination`), set `cursor_key="id"` ## Server-Side Session Stores -Supported database adapters expose a Litestar `Store` implementation under `sqlspec.adapters..litestar` (`AsyncpgStore`, `PsycopgSyncStore`, `PsycopgAsyncStore`, `PsqlpyStore`, `CockroachAsyncpgStore`, `CockroachPsycopgSyncStore`, `CockroachPsycopgAsyncStore`, `AiosqliteStore`, `SQLiteStore`, `AsyncmyStore`, `AiomysqlStore`, `PyMysqlStore`, `MysqlConnectorSyncStore`, `MysqlConnectorAsyncStore`, `OracleSyncStore`, `OracleAsyncStore`, `MssqlPythonStore`, `PymssqlStore`, `SybaseStore`, `SqlanywhereStore`, `Db2SyncStore`, `Db2AsyncStore`, `DuckdbStore`, `BigQueryStore`, `SpannerSyncStore`, `SpannerAsyncStore`, `ADBCStore`, and `ArrowOdbcStore`): +Supported database adapters expose a Litestar `Store` implementation under `sqlspec.adapters..litestar` (`AsyncpgStore`, `PsycopgSyncStore`, `PsycopgAsyncStore`, `PsqlpyStore`, `CockroachAsyncpgStore`, `CockroachPsycopgSyncStore`, `CockroachPsycopgAsyncStore`, `AiosqliteStore`, `SqliteStore`, `AsyncmyStore`, `AiomysqlStore`, `PyMysqlStore`, `MysqlConnectorSyncStore`, `MysqlConnectorAsyncStore`, `OracleSyncStore`, `OracleAsyncStore`, `MssqlPythonStore`, `PymssqlStore`, `SybaseStore`, `SqlanywhereStore`, `Db2SyncStore`, `Db2AsyncStore`, `DuckDBStore`, `BigQueryStore`, `SpannerSyncStore`, `SpannerAsyncStore`, `AdbcStore`, and `ArrowOdbcStore`): ```python --8<-- "docs/examples/frameworks/litestar/session_stores.py:example" diff --git a/docs/reference/adapters/adbc.md b/docs/reference/adapters/adbc.md index bfdd4a561..e752171ac 100644 --- a/docs/reference/adapters/adbc.md +++ b/docs/reference/adapters/adbc.md @@ -49,7 +49,7 @@ Apache Arrow Database Connectivity (ADBC) adapter supporting zero-copy Arrow ing ## Extension Stores & Configs -::: sqlspec.adapters.adbc.litestar.ADBCStore +::: sqlspec.adapters.adbc.litestar.AdbcStore options: heading_level: 4 diff --git a/docs/reference/adapters/duckdb.md b/docs/reference/adapters/duckdb.md index 510c4b47d..b5f9b9565 100644 --- a/docs/reference/adapters/duckdb.md +++ b/docs/reference/adapters/duckdb.md @@ -75,27 +75,27 @@ Embedded analytical database adapter built on `duckdb` with zero-copy Apache Arr ## Extension Stores & Configs -::: sqlspec.adapters.duckdb.litestar.DuckdbStore +::: sqlspec.adapters.duckdb.litestar.DuckDBStore options: heading_level: 4 -::: sqlspec.adapters.duckdb.adk.DuckdbADKConfig +::: sqlspec.adapters.duckdb.adk.DuckDBADKConfig options: heading_level: 4 separate_signature: false show_signature_annotations: true -::: sqlspec.adapters.duckdb.adk.DuckdbADKFTSOptions +::: sqlspec.adapters.duckdb.adk.DuckDBADKFTSOptions options: heading_level: 4 separate_signature: false show_signature_annotations: true -::: sqlspec.adapters.duckdb.adk.DuckdbADKStore +::: sqlspec.adapters.duckdb.adk.DuckDBADKStore options: heading_level: 4 -::: sqlspec.adapters.duckdb.adk.DuckdbADKMemoryStore +::: sqlspec.adapters.duckdb.adk.DuckDBADKMemoryStore options: heading_level: 4 diff --git a/docs/reference/adapters/sqlanywhere.md b/docs/reference/adapters/sqlanywhere.md index e0d6c131b..8f9849add 100644 --- a/docs/reference/adapters/sqlanywhere.md +++ b/docs/reference/adapters/sqlanywhere.md @@ -86,6 +86,6 @@ Synchronous SAP SQL Anywhere 12+ adapter built on `sqlanydb` (`dbcapi` C library options: heading_level: 4 -::: sqlspec.adapters.sqlanywhere.events.SqlanywhereSyncEventQueueStore +::: sqlspec.adapters.sqlanywhere.events.SqlanywhereEventQueueStore options: heading_level: 4 diff --git a/docs/reference/adapters/sqlite.md b/docs/reference/adapters/sqlite.md index 6800ca6df..98060fd51 100644 --- a/docs/reference/adapters/sqlite.md +++ b/docs/reference/adapters/sqlite.md @@ -60,7 +60,7 @@ Synchronous SQLite adapter built on Python's standard library `sqlite3` with thr separate_signature: false show_signature_annotations: true -::: sqlspec.adapters.sqlite.litestar.SQLiteStore +::: sqlspec.adapters.sqlite.litestar.SqliteStore options: heading_level: 4 diff --git a/docs/reference/adapters/sybase.md b/docs/reference/adapters/sybase.md index 4c2089776..0ac44422c 100644 --- a/docs/reference/adapters/sybase.md +++ b/docs/reference/adapters/sybase.md @@ -86,6 +86,6 @@ Synchronous SAP Adaptive Server Enterprise (Sybase ASE) 16.0+ adapter built on ` options: heading_level: 4 -::: sqlspec.adapters.sybase.events.SybaseSyncEventQueueStore +::: sqlspec.adapters.sybase.events.SybaseEventQueueStore options: heading_level: 4 diff --git a/sqlspec/adapters/adbc/litestar/__init__.py b/sqlspec/adapters/adbc/litestar/__init__.py index 25bc6130d..d6f1050f7 100644 --- a/sqlspec/adapters/adbc/litestar/__init__.py +++ b/sqlspec/adapters/adbc/litestar/__init__.py @@ -1,5 +1,5 @@ """Litestar integration for ADBC adapter.""" -from sqlspec.adapters.adbc.litestar.store import ADBCStore +from sqlspec.adapters.adbc.litestar.store import AdbcStore -__all__ = ("ADBCStore",) +__all__ = ("AdbcStore",) diff --git a/sqlspec/adapters/adbc/litestar/store.py b/sqlspec/adapters/adbc/litestar/store.py index c2951e289..ce6e2d087 100644 --- a/sqlspec/adapters/adbc/litestar/store.py +++ b/sqlspec/adapters/adbc/litestar/store.py @@ -25,10 +25,10 @@ from sqlspec.adapters.adbc.config import AdbcConfig -__all__ = ("ADBCStore",) +__all__ = ("AdbcStore",) -class ADBCStore(BaseSQLSpecStore["AdbcConfig"]): +class AdbcStore(BaseSQLSpecStore["AdbcConfig"]): """ADBC session store using synchronous ADBC driver. Implements server-side session storage for Litestar using ADBC diff --git a/sqlspec/adapters/duckdb/adk/__init__.py b/sqlspec/adapters/duckdb/adk/__init__.py index b54354c1a..d7b8b0568 100644 --- a/sqlspec/adapters/duckdb/adk/__init__.py +++ b/sqlspec/adapters/duckdb/adk/__init__.py @@ -5,6 +5,6 @@ testing, and analytical workloads. """ -from sqlspec.adapters.duckdb.adk.store import DuckdbADKConfig, DuckdbADKFTSOptions, DuckdbADKMemoryStore, DuckdbADKStore +from sqlspec.adapters.duckdb.adk.store import DuckDBADKConfig, DuckDBADKFTSOptions, DuckDBADKMemoryStore, DuckDBADKStore -__all__ = ("DuckdbADKConfig", "DuckdbADKFTSOptions", "DuckdbADKMemoryStore", "DuckdbADKStore") +__all__ = ("DuckDBADKConfig", "DuckDBADKFTSOptions", "DuckDBADKMemoryStore", "DuckDBADKStore") diff --git a/sqlspec/adapters/duckdb/adk/store.py b/sqlspec/adapters/duckdb/adk/store.py index 0b634baa2..6a5799877 100644 --- a/sqlspec/adapters/duckdb/adk/store.py +++ b/sqlspec/adapters/duckdb/adk/store.py @@ -27,7 +27,7 @@ from sqlspec.extensions.adk import SessionOrderBy, StoredMemory -__all__ = ("DuckdbADKConfig", "DuckdbADKFTSOptions", "DuckdbADKMemoryStore", "DuckdbADKStore") +__all__ = ("DuckDBADKConfig", "DuckDBADKFTSOptions", "DuckDBADKMemoryStore", "DuckDBADKStore") logger = get_logger("sqlspec.adapters.duckdb.adk.store") @@ -44,7 +44,7 @@ } -class DuckdbADKFTSOptions(TypedDict): +class DuckDBADKFTSOptions(TypedDict): """DuckDB FTS ``PRAGMA create_fts_index`` options for ADK memory search.""" stemmer: NotRequired[str] @@ -66,7 +66,7 @@ class DuckdbADKFTSOptions(TypedDict): """Whether initial FTS index creation overwrites an existing index. Defaults to ``0``.""" -class DuckdbADKConfig(ADKConfig): +class DuckDBADKConfig(ADKConfig): """DuckDB-specific ADK extension settings. Use these keys inside ``extension_config["adk"]`` with the DuckDB ADK stores. @@ -79,11 +79,11 @@ class DuckdbADKConfig(ADKConfig): enable_event_generated_column_indexes: NotRequired[bool] """Index generated event projection columns when generated columns are enabled.""" - memory_fts_options: NotRequired[DuckdbADKFTSOptions] + memory_fts_options: NotRequired[DuckDBADKFTSOptions] """DuckDB FTS ``PRAGMA create_fts_index`` options for ADK memory search.""" -class DuckdbADKStore(BaseSyncADKStore["DuckDBConfig"]): +class DuckDBADKStore(BaseSyncADKStore["DuckDBConfig"]): """DuckDB ADK store for Google Agent Development Kit. Implements session and event storage for Google Agent Development Kit @@ -872,7 +872,7 @@ def _set_metadata(self, key: str, value: str) -> None: conn.commit() -class DuckdbADKMemoryStore(BaseSyncADKMemoryStore["DuckDBConfig"]): +class DuckDBADKMemoryStore(BaseSyncADKMemoryStore["DuckDBConfig"]): """DuckDB ADK memory store using synchronous DuckDB driver with async wrappers. Implements memory entry storage for Google Agent Development Kit @@ -1257,15 +1257,15 @@ def _raise_session_not_found(session_id: str) -> None: raise ValueError(msg) -def _adk_config(config: Any) -> DuckdbADKConfig: +def _adk_config(config: Any) -> DuckDBADKConfig: """Return DuckDB ADK extension settings from ``extension_config["adk"]``.""" extension_config = getattr(config, "extension_config", {}) if not isinstance(extension_config, dict): - return cast("DuckdbADKConfig", {}) + return cast("DuckDBADKConfig", {}) adk_config = extension_config.get("adk", {}) if not isinstance(adk_config, dict): - return cast("DuckdbADKConfig", {}) - return cast("DuckdbADKConfig", adk_config) + return cast("DuckDBADKConfig", {}) + return cast("DuckDBADKConfig", adk_config) def _validate_duckdb_fts_options(options: Mapping[str, object]) -> None: diff --git a/sqlspec/adapters/duckdb/litestar/__init__.py b/sqlspec/adapters/duckdb/litestar/__init__.py index c6b17c019..7b9ea3253 100644 --- a/sqlspec/adapters/duckdb/litestar/__init__.py +++ b/sqlspec/adapters/duckdb/litestar/__init__.py @@ -1,5 +1,5 @@ """Litestar integration for DuckDB adapter.""" -from sqlspec.adapters.duckdb.litestar.store import DuckdbStore +from sqlspec.adapters.duckdb.litestar.store import DuckDBStore -__all__ = ("DuckdbStore",) +__all__ = ("DuckDBStore",) diff --git a/sqlspec/adapters/duckdb/litestar/store.py b/sqlspec/adapters/duckdb/litestar/store.py index 3031bf7c6..421aacf91 100644 --- a/sqlspec/adapters/duckdb/litestar/store.py +++ b/sqlspec/adapters/duckdb/litestar/store.py @@ -10,10 +10,10 @@ from sqlspec.adapters.duckdb.config import DuckDBConfig -__all__ = ("DuckdbStore",) +__all__ = ("DuckDBStore",) -class DuckdbStore(BaseSQLSpecStore["DuckDBConfig"]): +class DuckDBStore(BaseSQLSpecStore["DuckDBConfig"]): """DuckDB session store using synchronous DuckDB driver. Implements server-side session storage for Litestar using DuckDB diff --git a/sqlspec/adapters/sqlanywhere/events/__init__.py b/sqlspec/adapters/sqlanywhere/events/__init__.py index 70252c8ad..d40657a0a 100644 --- a/sqlspec/adapters/sqlanywhere/events/__init__.py +++ b/sqlspec/adapters/sqlanywhere/events/__init__.py @@ -1,5 +1,5 @@ """Events integration for SAP SQL Anywhere adapter.""" -from sqlspec.adapters.sqlanywhere.events.store import SqlanywhereEventQueueStore, SqlanywhereSyncEventQueueStore +from sqlspec.adapters.sqlanywhere.events.store import SqlanywhereEventQueueStore -__all__ = ("SqlanywhereEventQueueStore", "SqlanywhereSyncEventQueueStore") +__all__ = ("SqlanywhereEventQueueStore",) diff --git a/sqlspec/adapters/sqlanywhere/events/store.py b/sqlspec/adapters/sqlanywhere/events/store.py index 9a4e23b35..8c945dd43 100644 --- a/sqlspec/adapters/sqlanywhere/events/store.py +++ b/sqlspec/adapters/sqlanywhere/events/store.py @@ -3,12 +3,12 @@ from sqlspec.adapters.sqlanywhere.config import SqlanywhereConfig from sqlspec.extensions.events import BaseEventQueueStore -__all__ = ("SqlanywhereEventQueueStore", "SqlanywhereSyncEventQueueStore") +__all__ = ("SqlanywhereEventQueueStore",) _VARCHAR_TEXT_THRESHOLD = 4000 -class SqlanywhereSyncEventQueueStore(BaseEventQueueStore[SqlanywhereConfig]): +class SqlanywhereEventQueueStore(BaseEventQueueStore[SqlanywhereConfig]): """SAP SQL Anywhere DDL hooks for the event queue store.""" __slots__ = () @@ -26,6 +26,3 @@ def _integer_type(self) -> str: def _timestamp_default(self) -> str: return "CURRENT UTC TIMESTAMP" - - -SqlanywhereEventQueueStore = SqlanywhereSyncEventQueueStore diff --git a/sqlspec/adapters/sqlite/litestar/__init__.py b/sqlspec/adapters/sqlite/litestar/__init__.py index 99aeee476..4165b2867 100644 --- a/sqlspec/adapters/sqlite/litestar/__init__.py +++ b/sqlspec/adapters/sqlite/litestar/__init__.py @@ -1,5 +1,5 @@ """Litestar integration for SQLite adapter.""" -from sqlspec.adapters.sqlite.litestar.store import SqliteLitestarConfig, SQLiteStore +from sqlspec.adapters.sqlite.litestar.store import SqliteLitestarConfig, SqliteStore -__all__ = ("SQLiteStore", "SqliteLitestarConfig") +__all__ = ("SqliteLitestarConfig", "SqliteStore") diff --git a/sqlspec/adapters/sqlite/litestar/store.py b/sqlspec/adapters/sqlite/litestar/store.py index 01d3bee41..29d977cde 100644 --- a/sqlspec/adapters/sqlite/litestar/store.py +++ b/sqlspec/adapters/sqlite/litestar/store.py @@ -20,7 +20,7 @@ if TYPE_CHECKING: from sqlspec.adapters.sqlite.config import SqliteConfig -__all__ = ("JULIAN_EPOCH", "SECONDS_PER_DAY", "SQLiteStore", "SqliteLitestarConfig") +__all__ = ("JULIAN_EPOCH", "SECONDS_PER_DAY", "SqliteLitestarConfig", "SqliteStore") class SqliteLitestarConfig(LitestarConfig): @@ -36,7 +36,7 @@ class SqliteLitestarConfig(LitestarConfig): """Validated PRAGMA overrides applied during schema preparation.""" -class SQLiteStore(BaseSQLSpecStore["SqliteConfig"]): +class SqliteStore(BaseSQLSpecStore["SqliteConfig"]): """SQLite session store using synchronous SQLite driver. Implements server-side session storage for Litestar using SQLite diff --git a/sqlspec/adapters/sybase/events/__init__.py b/sqlspec/adapters/sybase/events/__init__.py index 7bf12a0d3..8652828d1 100644 --- a/sqlspec/adapters/sybase/events/__init__.py +++ b/sqlspec/adapters/sybase/events/__init__.py @@ -1,5 +1,5 @@ """Sybase ASE event extension.""" -from sqlspec.adapters.sybase.events.store import SybaseEventQueueStore, SybaseSyncEventQueueStore +from sqlspec.adapters.sybase.events.store import SybaseEventQueueStore -__all__ = ("SybaseEventQueueStore", "SybaseSyncEventQueueStore") +__all__ = ("SybaseEventQueueStore",) diff --git a/sqlspec/adapters/sybase/events/store.py b/sqlspec/adapters/sybase/events/store.py index d5871ac5d..02cd6d245 100644 --- a/sqlspec/adapters/sybase/events/store.py +++ b/sqlspec/adapters/sybase/events/store.py @@ -4,12 +4,12 @@ from sqlspec.adapters.sybase.core import escape_sql_literal, quote_bracket_identifier, split_schema_and_table from sqlspec.extensions.events import BaseEventQueueStore -__all__ = ("SybaseEventQueueStore", "SybaseSyncEventQueueStore") +__all__ = ("SybaseEventQueueStore",) _VARCHAR_TEXT_THRESHOLD = 4000 -class SybaseSyncEventQueueStore(BaseEventQueueStore[SybaseConfig]): +class SybaseEventQueueStore(BaseEventQueueStore[SybaseConfig]): """Sybase ASE DDL hooks for the event queue store.""" __slots__ = () @@ -75,9 +75,6 @@ def _wrap_drop_statement(self, statement: str) -> str: return f"IF OBJECT_ID('{object_name}', 'U') IS NOT NULL EXECUTE('{escaped_statement}');" -SybaseEventQueueStore = SybaseSyncEventQueueStore - - def _object_name(table_name: str) -> str: schema_name, bare_table_name = split_schema_and_table(table_name) return f"{quote_bracket_identifier(schema_name)}.{quote_bracket_identifier(bare_table_name)}" diff --git a/tests/integration/adapters/duckdb/duckdb/extensions/adk/test_maintenance.py b/tests/integration/adapters/duckdb/duckdb/extensions/adk/test_maintenance.py index fcc484305..8650ffddd 100644 --- a/tests/integration/adapters/duckdb/duckdb/extensions/adk/test_maintenance.py +++ b/tests/integration/adapters/duckdb/duckdb/extensions/adk/test_maintenance.py @@ -6,7 +6,7 @@ import pytest -from sqlspec.adapters.duckdb.adk import DuckdbADKMemoryStore, DuckdbADKStore +from sqlspec.adapters.duckdb.adk import DuckDBADKMemoryStore, DuckDBADKStore from sqlspec.adapters.duckdb.config import DuckDBConfig from sqlspec.extensions.adk import StoredEvent, StoredMemory from sqlspec.extensions.adk.maintenance import ( @@ -57,7 +57,7 @@ def _event(*, app_name: str, session_id: str, timestamp: datetime) -> StoredEven def test_duckdb_prune_sessions_scopes_to_app_name(tmp_path: Path) -> None: """Pruning sessions for one application leaves other applications untouched.""" config = _config(tmp_path, "prune_sessions") - store = DuckdbADKStore(config) + store = DuckDBADKStore(config) store.create_tables() stale = datetime.now(UTC) - timedelta(days=120) @@ -77,7 +77,7 @@ def test_duckdb_prune_sessions_scopes_to_app_name(tmp_path: Path) -> None: def test_duckdb_prune_events_scopes_to_app_name(tmp_path: Path) -> None: """Pruning events for one application leaves other applications' events intact.""" config = _config(tmp_path, "prune_events") - store = DuckdbADKStore(config) + store = DuckDBADKStore(config) store.create_tables() stale = datetime.now(UTC) - timedelta(days=200) @@ -95,7 +95,7 @@ def test_duckdb_prune_events_scopes_to_app_name(tmp_path: Path) -> None: def test_duckdb_prune_user_state_scopes_to_app_name(tmp_path: Path) -> None: """User state pruning honors both the idle threshold and the application filter.""" config = _config(tmp_path, "prune_user_state") - store = DuckdbADKStore(config) + store = DuckDBADKStore(config) store.create_tables() for app in ("app_keep", "app_prune"): @@ -115,7 +115,7 @@ def test_duckdb_prune_user_state_scopes_to_app_name(tmp_path: Path) -> None: def test_duckdb_prune_memory_respects_scope(tmp_path: Path) -> None: """Memory pruning removes user-scoped entries while preserving app-scoped ones.""" config = _config(tmp_path, "prune_memory") - store = DuckdbADKMemoryStore(config) + store = DuckDBADKMemoryStore(config) store.create_tables() stale = datetime.now(UTC) - timedelta(days=200) diff --git a/tests/integration/adapters/duckdb/duckdb/extensions/adk/test_memory_store.py b/tests/integration/adapters/duckdb/duckdb/extensions/adk/test_memory_store.py index 1295b3eae..1b393e8ea 100644 --- a/tests/integration/adapters/duckdb/duckdb/extensions/adk/test_memory_store.py +++ b/tests/integration/adapters/duckdb/duckdb/extensions/adk/test_memory_store.py @@ -7,7 +7,7 @@ import pytest -from sqlspec.adapters.duckdb.adk import DuckdbADKMemoryStore +from sqlspec.adapters.duckdb.adk import DuckDBADKMemoryStore from sqlspec.adapters.duckdb.config import DuckDBConfig from sqlspec.extensions.adk import StoredMemory @@ -35,20 +35,20 @@ def _build_record( ) -def _build_store(tmp_path: Path) -> DuckdbADKMemoryStore: +def _build_store(tmp_path: Path) -> DuckDBADKMemoryStore: db_path = tmp_path / "test_adk_memory.duckdb" config = DuckDBConfig(connection_config={"database": str(db_path)}) - store = DuckdbADKMemoryStore(config) + store = DuckDBADKMemoryStore(config) store.create_tables() return store -def _build_fts_store(tmp_path: Path) -> DuckdbADKMemoryStore: +def _build_fts_store(tmp_path: Path) -> DuckDBADKMemoryStore: db_path = tmp_path / "test_adk_memory_fts.duckdb" config = DuckDBConfig( connection_config={"database": str(db_path)}, extension_config={"adk": {"memory_use_fts": True}} ) - store = DuckdbADKMemoryStore(config) + store = DuckDBADKMemoryStore(config) with config.provide_connection() as conn: if not store._ensure_fts_extension(conn): # pyright: ignore[reportPrivateUsage] pytest.skip("DuckDB FTS extension is unavailable") diff --git a/tests/integration/adapters/duckdb/duckdb/extensions/adk/test_store.py b/tests/integration/adapters/duckdb/duckdb/extensions/adk/test_store.py index 1f8663e98..16e19a3a0 100644 --- a/tests/integration/adapters/duckdb/duckdb/extensions/adk/test_store.py +++ b/tests/integration/adapters/duckdb/duckdb/extensions/adk/test_store.py @@ -13,7 +13,7 @@ import pytest -from sqlspec.adapters.duckdb.adk import DuckdbADKStore +from sqlspec.adapters.duckdb.adk import DuckDBADKStore from sqlspec.adapters.duckdb.config import DuckDBConfig from sqlspec.extensions.adk import StoredEvent @@ -21,7 +21,7 @@ @pytest.fixture -def duckdb_adk_store(tmp_path: Path) -> "Generator[DuckdbADKStore, None, None]": +def duckdb_adk_store(tmp_path: Path) -> "Generator[DuckDBADKStore, None, None]": """Create DuckDB ADK store with temporary file-based database. Args: @@ -39,7 +39,7 @@ def duckdb_adk_store(tmp_path: Path) -> "Generator[DuckdbADKStore, None, None]": connection_config={"database": str(db_path)}, extension_config={"adk": {"session_table": "test_sessions", "events_table": "test_events"}}, ) - store = DuckdbADKStore(config) + store = DuckDBADKStore(config) store.create_tables() yield store finally: @@ -47,7 +47,7 @@ def duckdb_adk_store(tmp_path: Path) -> "Generator[DuckdbADKStore, None, None]": db_path.unlink() -def test_event_with_optional_fields(duckdb_adk_store: DuckdbADKStore) -> None: +def test_event_with_optional_fields(duckdb_adk_store: DuckDBADKStore) -> None: """Test creating events with optional fields stored in event_data.""" session_id = "session-008" duckdb_adk_store.create_session(session_id, "test-app", "user-008", {}) @@ -87,7 +87,7 @@ def test_event_with_optional_fields(duckdb_adk_store: DuckdbADKStore) -> None: assert event_data["turn_complete"] is False -def test_event_ordering_by_timestamp(duckdb_adk_store: DuckdbADKStore) -> None: +def test_event_ordering_by_timestamp(duckdb_adk_store: DuckDBADKStore) -> None: """Test events are ordered by timestamp ascending.""" session_id = "session-009" duckdb_adk_store.create_session(session_id, "test-app", "user-009", {}) @@ -136,7 +136,7 @@ def test_event_ordering_by_timestamp(duckdb_adk_store: DuckdbADKStore) -> None: assert event_ids == ["event-first", "event-middle", "event-last"] -def test_session_state_with_complex_data(duckdb_adk_store: DuckdbADKStore) -> None: +def test_session_state_with_complex_data(duckdb_adk_store: DuckDBADKStore) -> None: """Test session state with nested JSON structures.""" session_id = "session-complex" complex_state = { @@ -158,7 +158,7 @@ def test_session_state_with_complex_data(duckdb_adk_store: DuckdbADKStore) -> No assert session["state"]["conversation"]["turn_count"] == 5 -def test_event_data_round_trip(duckdb_adk_store: DuckdbADKStore) -> None: +def test_event_data_round_trip(duckdb_adk_store: DuckDBADKStore) -> None: """Test storing and retrieving event data via event_data.""" session_id = "session-json-rt" duckdb_adk_store.create_session(session_id, "test-app", "user-012", {}) @@ -185,7 +185,7 @@ def test_event_data_round_trip(duckdb_adk_store: DuckdbADKStore) -> None: assert events[0]["event_data"]["content"] == {"data": "value"} -def test_concurrent_session_updates(duckdb_adk_store: DuckdbADKStore) -> None: +def test_concurrent_session_updates(duckdb_adk_store: DuckDBADKStore) -> None: """Test multiple updates to same session.""" session_id = "session-concurrent" duckdb_adk_store.create_session(session_id, "test-app", "user-013", {"counter": 0}) @@ -222,7 +222,7 @@ def test_owner_id_column_with_integer(tmp_path: Path) -> None: } }, ) - store = DuckdbADKStore(config_with_extension) + store = DuckDBADKStore(config_with_extension) store.create_tables() assert store.owner_id_column_name == "tenant_id" @@ -265,7 +265,7 @@ def test_owner_id_column_with_ubigint(tmp_path: Path) -> None: } }, ) - store = DuckdbADKStore(config_with_extension) + store = DuckDBADKStore(config_with_extension) store.create_tables() assert store.owner_id_column_name == "owner_id" @@ -311,7 +311,7 @@ def test_owner_id_column_foreign_key_constraint(tmp_path: Path) -> None: } }, ) - store = DuckdbADKStore(config_with_extension) + store = DuckDBADKStore(config_with_extension) store.create_tables() store.create_session( @@ -353,7 +353,7 @@ def test_owner_id_column_without_value(tmp_path: Path) -> None: } }, ) - store = DuckdbADKStore(config_with_extension) + store = DuckDBADKStore(config_with_extension) store.create_tables() session = store.create_session( @@ -390,7 +390,7 @@ def test_owner_id_column_with_varchar(tmp_path: Path) -> None: } }, ) - store = DuckdbADKStore(config_with_extension) + store = DuckDBADKStore(config_with_extension) store.create_tables() session = store.create_session( @@ -434,7 +434,7 @@ def test_owner_id_column_multiple_sessions(tmp_path: Path) -> None: } }, ) - store = DuckdbADKStore(config_with_extension) + store = DuckDBADKStore(config_with_extension) store.create_tables() for i in range(5): @@ -477,7 +477,7 @@ def test_owner_id_column_query_by_fk(tmp_path: Path) -> None: } }, ) - store = DuckdbADKStore(config_with_extension) + store = DuckDBADKStore(config_with_extension) store.create_tables() store.create_session("s1", "app", "u1", {"val": 1}, owner_id=1) diff --git a/tests/integration/adapters/duckdb/duckdb/extensions/litestar/test_store.py b/tests/integration/adapters/duckdb/duckdb/extensions/litestar/test_store.py index cb14f3f6a..ee30bdd89 100644 --- a/tests/integration/adapters/duckdb/duckdb/extensions/litestar/test_store.py +++ b/tests/integration/adapters/duckdb/duckdb/extensions/litestar/test_store.py @@ -14,13 +14,13 @@ import pytest from sqlspec.adapters.duckdb.config import DuckDBConfig -from sqlspec.adapters.duckdb.litestar.store import DuckdbStore +from sqlspec.adapters.duckdb.litestar.store import DuckDBStore pytestmark = [pytest.mark.duckdb, pytest.mark.integration] @pytest.fixture -async def duckdb_store(tmp_path: Path) -> AsyncGenerator[DuckdbStore, None]: +async def duckdb_store(tmp_path: Path) -> AsyncGenerator[DuckDBStore, None]: """Create DuckDB store with temporary file-based database. Args: @@ -42,7 +42,7 @@ async def duckdb_store(tmp_path: Path) -> AsyncGenerator[DuckdbStore, None]: connection_config={"database": str(db_path)}, extension_config={"litestar": {"session_table": "test_sessions"}} ) try: - store = DuckdbStore(config) + store = DuckDBStore(config) await store.create_table() yield store await store.delete_all() @@ -52,7 +52,7 @@ async def duckdb_store(tmp_path: Path) -> AsyncGenerator[DuckdbStore, None]: db_path.unlink() -async def test_sync_to_thread_concurrency(duckdb_store: DuckdbStore) -> None: +async def test_sync_to_thread_concurrency(duckdb_store: DuckDBStore) -> None: """Test concurrent access via sync_to_thread wrapper. DuckDB has write serialization, so we test sequential writes diff --git a/tests/integration/adapters/postgres/adbc/extensions/litestar/test_store.py b/tests/integration/adapters/postgres/adbc/extensions/litestar/test_store.py index 1c1eea697..1f61784fd 100644 --- a/tests/integration/adapters/postgres/adbc/extensions/litestar/test_store.py +++ b/tests/integration/adapters/postgres/adbc/extensions/litestar/test_store.py @@ -13,16 +13,16 @@ import pytest from sqlspec.adapters.adbc import AdbcConfig -from sqlspec.adapters.adbc.litestar import ADBCStore +from sqlspec.adapters.adbc.litestar import AdbcStore pytestmark = [pytest.mark.xdist_group("postgres"), pytest.mark.adbc, pytest.mark.integration] @pytest.fixture -async def adbc_store(adbc_postgres_config: AdbcConfig) -> AsyncGenerator[ADBCStore, None]: +async def adbc_store(adbc_postgres_config: AdbcConfig) -> AsyncGenerator[AdbcStore, None]: """Create an ADBC store with PostgreSQL backend.""" adbc_postgres_config.extension_config = {"litestar": {"session_table": "test_adbc_sessions"}} - store = ADBCStore(adbc_postgres_config) + store = AdbcStore(adbc_postgres_config) await store.create_table() try: yield store @@ -31,7 +31,7 @@ async def adbc_store(adbc_postgres_config: AdbcConfig) -> AsyncGenerator[ADBCSto await store.delete_all() -async def test_sync_to_thread_concurrency(adbc_store: ADBCStore) -> None: +async def test_sync_to_thread_concurrency(adbc_store: AdbcStore) -> None: """ADBC PostgreSQL supports concurrent store reads and writes through sync_to_thread.""" async def write_session(session_id: int) -> None: diff --git a/tests/integration/adapters/sqlite/sqlite/extensions/litestar/test_store.py b/tests/integration/adapters/sqlite/sqlite/extensions/litestar/test_store.py index bb100e104..374f044eb 100644 --- a/tests/integration/adapters/sqlite/sqlite/extensions/litestar/test_store.py +++ b/tests/integration/adapters/sqlite/sqlite/extensions/litestar/test_store.py @@ -20,11 +20,11 @@ from litestar.stores.base import Store from litestar.stores.registry import StoreRegistry - from sqlspec.adapters.sqlite.litestar.store import SQLiteStore + from sqlspec.adapters.sqlite.litestar.store import SqliteStore else: Store = import_optional_attr("litestar.stores.base", "Store") StoreRegistry = import_optional_attr("litestar.stores.registry", "StoreRegistry") - SQLiteStore = import_optional_attr("sqlspec.adapters.sqlite.litestar.store", "SQLiteStore") + SqliteStore = import_optional_attr("sqlspec.adapters.sqlite.litestar.store", "SqliteStore") pytest.importorskip("litestar") @@ -32,19 +32,19 @@ @pytest.fixture -async def sqlite_store() -> "AsyncGenerator[SQLiteStore, None]": +async def sqlite_store() -> "AsyncGenerator[SqliteStore, None]": """Create SQLite store with shared in-memory database.""" config = SqliteConfig( connection_config={"database": "file:test_sessions_mem?mode=memory&cache=shared", "uri": True}, extension_config={"litestar": {"session_table": "test_sessions"}}, ) - store = SQLiteStore(config) + store = SqliteStore(config) await store.create_table() yield store await store.delete_all() -async def test_sync_to_thread_concurrency(sqlite_store: "SQLiteStore") -> None: +async def test_sync_to_thread_concurrency(sqlite_store: "SqliteStore") -> None: """Test concurrent access via sync_to_thread wrapper. SQLite has write serialization, so we test sequential writes @@ -62,7 +62,7 @@ async def read_session(session_id: int) -> bytes | None: assert result == f"data_{i}".encode() -async def test_store_implements_litestar_store_contract(sqlite_store: "SQLiteStore") -> None: +async def test_store_implements_litestar_store_contract(sqlite_store: "SqliteStore") -> None: """SQLite stores satisfy Litestar's nominal store and registry contract.""" registry = StoreRegistry() diff --git a/tests/integration/fixtures/adapter_cases.py b/tests/integration/fixtures/adapter_cases.py index e2bb886af..75ca1b5d7 100644 --- a/tests/integration/fixtures/adapter_cases.py +++ b/tests/integration/fixtures/adapter_cases.py @@ -76,7 +76,7 @@ from sqlspec.adapters.adbc import AdbcConfig, AdbcDriver from sqlspec.adapters.adbc.adk import AdbcADKStore - from sqlspec.adapters.adbc.litestar import ADBCStore + from sqlspec.adapters.adbc.litestar import AdbcStore from sqlspec.adapters.aiomysql import AiomysqlConfig, AiomysqlDriver, AiomysqlDriverFeatures from sqlspec.adapters.aiomysql.adk import AiomysqlADKStore from sqlspec.adapters.aiomysql.litestar import AiomysqlStore @@ -111,8 +111,8 @@ from sqlspec.adapters.db2.adk import Db2AsyncADKStore, Db2SyncADKStore from sqlspec.adapters.db2.litestar import Db2AsyncStore, Db2SyncStore from sqlspec.adapters.duckdb import DuckDBConfig, DuckDBDriver, DuckDBDriverFeatures - from sqlspec.adapters.duckdb.adk import DuckdbADKStore - from sqlspec.adapters.duckdb.litestar import DuckdbStore + from sqlspec.adapters.duckdb.adk import DuckDBADKStore + from sqlspec.adapters.duckdb.litestar import DuckDBStore from sqlspec.adapters.mssql_python import MssqlPythonConfig, MssqlPythonDriver from sqlspec.adapters.mssql_python.adk import MssqlPythonADKStore from sqlspec.adapters.mysqlconnector import ( @@ -152,7 +152,7 @@ from sqlspec.adapters.pymysql.litestar import PyMysqlStore from sqlspec.adapters.spanner import SpannerAsyncConfig, SpannerSyncConfig from sqlspec.adapters.sqlite.adk import SqliteADKStore - from sqlspec.adapters.sqlite.litestar import SQLiteStore + from sqlspec.adapters.sqlite.litestar import SqliteStore from tests.integration.adapters.bigquery._wedge import describe_wedge, is_emulator_wedge from tests.integration.fixtures.bigquery import _bigquery_connection_config from tests.integration.fixtures.mysql import _mysql_connection_config @@ -161,7 +161,7 @@ AdbcConfig = import_optional_attr("sqlspec.adapters.adbc", "AdbcConfig") AdbcDriver = import_optional_attr("sqlspec.adapters.adbc", "AdbcDriver") AdbcADKStore = import_optional_attr("sqlspec.adapters.adbc.adk", "AdbcADKStore") - ADBCStore = import_optional_attr("sqlspec.adapters.adbc.litestar", "ADBCStore") + AdbcStore = import_optional_attr("sqlspec.adapters.adbc.litestar", "AdbcStore") AiomysqlConfig = import_optional_attr("sqlspec.adapters.aiomysql", "AiomysqlConfig") AiomysqlDriver = import_optional_attr("sqlspec.adapters.aiomysql", "AiomysqlDriver") AiomysqlDriverFeatures = import_optional_attr("sqlspec.adapters.aiomysql", "AiomysqlDriverFeatures") @@ -230,8 +230,8 @@ DuckDBConfig = import_optional_attr("sqlspec.adapters.duckdb", "DuckDBConfig") DuckDBDriver = import_optional_attr("sqlspec.adapters.duckdb", "DuckDBDriver") DuckDBDriverFeatures = import_optional_attr("sqlspec.adapters.duckdb", "DuckDBDriverFeatures") - DuckdbADKStore = import_optional_attr("sqlspec.adapters.duckdb.adk", "DuckdbADKStore") - DuckdbStore = import_optional_attr("sqlspec.adapters.duckdb.litestar", "DuckdbStore") + DuckDBADKStore = import_optional_attr("sqlspec.adapters.duckdb.adk", "DuckDBADKStore") + DuckDBStore = import_optional_attr("sqlspec.adapters.duckdb.litestar", "DuckDBStore") MssqlPythonConfig = import_optional_attr("sqlspec.adapters.mssql_python", "MssqlPythonConfig") MssqlPythonDriver = import_optional_attr("sqlspec.adapters.mssql_python", "MssqlPythonDriver") MssqlPythonADKStore = import_optional_attr("sqlspec.adapters.mssql_python.adk", "MssqlPythonADKStore") @@ -289,7 +289,7 @@ SpannerAsyncConfig = import_optional_attr("sqlspec.adapters.spanner", "SpannerAsyncConfig") SpannerSyncConfig = import_optional_attr("sqlspec.adapters.spanner", "SpannerSyncConfig") SqliteADKStore = import_optional_attr("sqlspec.adapters.sqlite.adk", "SqliteADKStore") - SQLiteStore = import_optional_attr("sqlspec.adapters.sqlite.litestar", "SQLiteStore") + SqliteStore = import_optional_attr("sqlspec.adapters.sqlite.litestar", "SqliteStore") describe_wedge = import_optional_attr("tests.integration.adapters.bigquery._wedge", "describe_wedge") is_emulator_wedge = import_optional_attr("tests.integration.adapters.bigquery._wedge", "is_emulator_wedge") _bigquery_connection_config = import_optional_attr( @@ -1846,14 +1846,14 @@ def async_migration_case(request: pytest.FixtureRequest) -> MigrationCaseContext @pytest.fixture -async def contract_sqlite_store() -> "AsyncGenerator[SQLiteStore, None]": +async def contract_sqlite_store() -> "AsyncGenerator[SqliteStore, None]": """Provide a ready SQLite Litestar store for contract tests.""" pytest.importorskip("litestar") config = SqliteConfig( connection_config={"database": "file:contract_store_sqlite?mode=memory&cache=shared", "uri": True}, extension_config=_STORE_EXTENSION_CONFIG, ) - store = SQLiteStore(config) + store = SqliteStore(config) await store.create_table() yield store with contextlib.suppress(Exception): @@ -1876,7 +1876,7 @@ async def contract_aiosqlite_store() -> "AsyncGenerator[AiosqliteStore, None]": @pytest.fixture -async def contract_duckdb_store(tmp_path: Path) -> "AsyncGenerator[DuckdbStore, None]": +async def contract_duckdb_store(tmp_path: Path) -> "AsyncGenerator[DuckDBStore, None]": """Provide a ready DuckDB Litestar store for contract tests.""" pytest.importorskip("duckdb") pytest.importorskip("litestar") @@ -1884,7 +1884,7 @@ async def contract_duckdb_store(tmp_path: Path) -> "AsyncGenerator[DuckdbStore, connection_config={"database": str(tmp_path / f"store_{uuid4().hex}.duckdb")}, extension_config=_STORE_EXTENSION_CONFIG, ) - store = DuckdbStore(config) + store = DuckDBStore(config) await store.create_table() yield store with contextlib.suppress(Exception): @@ -2077,7 +2077,7 @@ async def contract_arrow_odbc_store(mssql_service: MSSQLService) -> "AsyncGenera @pytest.fixture -async def contract_adbc_store(postgres_service: PostgresService) -> "AsyncGenerator[ADBCStore, None]": +async def contract_adbc_store(postgres_service: PostgresService) -> "AsyncGenerator[AdbcStore, None]": """Provide a ready PostgreSQL-backed ADBC Litestar store for contract tests.""" pytest.importorskip("adbc_driver_manager") pytest.importorskip("pyarrow") @@ -2086,7 +2086,7 @@ async def contract_adbc_store(postgres_service: PostgresService) -> "AsyncGenera connection_config={"uri": _adbc_postgres_uri(postgres_service), "driver_name": "adbc_driver_postgresql"}, extension_config=_STORE_EXTENSION_CONFIG, ) - store = ADBCStore(config) + store = AdbcStore(config) await store.create_table() yield store with contextlib.suppress(Exception): @@ -2195,7 +2195,7 @@ def make() -> "tuple[Any, Any]": connection_config={"database": str(tmp_path / f"adk_{suffix}.duckdb")}, extension_config=_adk_extension_config(suffix), ) - return config, DuckdbADKStore(config) + return config, DuckDBADKStore(config) return make diff --git a/tests/unit/adapters/test_adapter_implementations.py b/tests/unit/adapters/test_adapter_implementations.py index ed02e73a5..f1d8af47c 100644 --- a/tests/unit/adapters/test_adapter_implementations.py +++ b/tests/unit/adapters/test_adapter_implementations.py @@ -11,6 +11,7 @@ import pytest +import sqlspec.adapters from sqlspec.adapters.sqlite.driver import SqliteDriver from sqlspec.config import DatabaseConfigProtocol from sqlspec.core import SQL, ParameterStyle, ParameterStyleConfig, SQLResult, StatementConfig @@ -703,3 +704,32 @@ def test_every_config_declares_type_coercion_capabilities() -> None: assert caps is not DatabaseConfigProtocol.type_coercion_capabilities, ( f"{config_cls.__name__}.type_coercion_capabilities is the base default object" ) + + +def _litestar_store_targets() -> "list[tuple[str, str]]": + """Return ``(adapter, config class)`` pairs for adapters that ship a Litestar store.""" + targets: list[tuple[str, str]] = [] + adapters_root = Path(sqlspec.adapters.__file__).parent + for config_path in sorted(adapters_root.glob("*/config.py")): + adapter = config_path.parent.name + if not (config_path.parent / "litestar" / "store.py").exists(): + continue + module = importlib.import_module(f"sqlspec.adapters.{adapter}.config") + for name, value in vars(module).items(): + if ( + isinstance(value, type) + and value.__module__ == module.__name__ + and name.endswith("Config") + and issubclass(value, DatabaseConfigProtocol) + ): + targets.append((adapter, name)) + return targets + + +@pytest.mark.parametrize(("adapter", "config_name"), _litestar_store_targets()) +def test_litestar_session_migration_resolves_store_class(adapter: str, config_name: str) -> None: + """The Litestar session migration derives the store class name from the config class name.""" + store_name = config_name.replace("Config", "Store") + store_module = importlib.import_module(f"sqlspec.adapters.{adapter}.litestar.store") + + assert hasattr(store_module, store_name), f"{adapter}: expected {store_name} in litestar/store.py" diff --git a/tests/unit/adapters/test_adbc/test_adk_dialect_support.py b/tests/unit/adapters/test_adbc/test_adk_dialect_support.py index ee5b6a079..db8accf1f 100644 --- a/tests/unit/adapters/test_adbc/test_adk_dialect_support.py +++ b/tests/unit/adapters/test_adbc/test_adk_dialect_support.py @@ -16,7 +16,7 @@ format_pgvector_literal, ) from sqlspec.adapters.adbc.events.store import AdbcEventQueueStore -from sqlspec.adapters.adbc.litestar.store import ADBCStore +from sqlspec.adapters.adbc.litestar.store import AdbcStore from sqlspec.core import StatementConfig from sqlspec.exceptions import ImproperConfigurationError @@ -336,9 +336,9 @@ def test_sqlite_adk_store_encodes_and_decodes_timestamps() -> None: def test_adbc_litestar_store_handles_pg_extensions_and_bigquery_snowflake(monkeypatch: Any) -> None: - """ADBCStore should recognize Postgres extension dialects, omit indexes on BigQuery/Snowflake, and use select_one_or_none.""" + """AdbcStore should recognize Postgres extension dialects, omit indexes on BigQuery/Snowflake, and use select_one_or_none.""" config = AdbcConfig(connection_config={"driver_name": "bigquery"}) - store = ADBCStore(config) + store = AdbcStore(config) store._dialect = "bigquery" # pyright: ignore[reportPrivateUsage] ddl = store._table_ddl() # pyright: ignore[reportPrivateUsage] diff --git a/tests/unit/adapters/test_adbc/test_extension_stores.py b/tests/unit/adapters/test_adbc/test_extension_stores.py index d328f1e24..a7b186355 100644 --- a/tests/unit/adapters/test_adbc/test_extension_stores.py +++ b/tests/unit/adapters/test_adbc/test_extension_stores.py @@ -8,7 +8,7 @@ from sqlspec.adapters.adbc.config import AdbcConfig from sqlspec.adapters.adbc.events.store import AdbcEventQueueStore -from sqlspec.adapters.adbc.litestar.store import ADBCStore +from sqlspec.adapters.adbc.litestar.store import AdbcStore __all__ = () @@ -25,7 +25,7 @@ def test_adbc_litestar_store_enterprise_dialect_primitives( dialect: str, expected_placeholder: str, expected_ts_expr: str, expected_ddl_fragment: str ) -> None: - store = ADBCStore(_make_adbc_config("sqlite:///:memory:")) + store = AdbcStore(_make_adbc_config("sqlite:///:memory:")) store._dialect = "mssql" if dialect == "tsql" else dialect assert store._get_param_placeholder(1) == expected_placeholder @@ -36,7 +36,7 @@ def test_adbc_litestar_store_enterprise_dialect_primitives( def test_adbc_litestar_store_normalize_tsql_dialect() -> None: driver_stub = SimpleNamespace(dialect="tsql") config_stub = SimpleNamespace(extension_config={"litestar": {}}, provide_session=lambda: nullcontext(driver_stub)) - store = ADBCStore(cast("Any", config_stub)) + store = AdbcStore(cast("Any", config_stub)) assert store._get_dialect() == "mssql" @@ -45,7 +45,7 @@ def test_adbc_litestar_store_normalize_tsql_dialect() -> None: [("oracle", "SQLCODE != -942"), ("db2", "SYSCAT.TABLES"), ("mssql", "IF OBJECT_ID")], ) def test_adbc_litestar_store_enterprise_drop_sql(dialect: str, expected_drop_fragment: str) -> None: - store = ADBCStore(_make_adbc_config("sqlite:///:memory:")) + store = AdbcStore(_make_adbc_config("sqlite:///:memory:")) store._dialect = dialect drop_sql = "\n".join(store._drop_table_sql()) assert expected_drop_fragment in drop_sql @@ -67,7 +67,7 @@ def commit(self) -> None: pass config_stub = SimpleNamespace(extension_config={"litestar": {}}, provide_session=lambda: nullcontext(_DriverStub())) - store = ADBCStore(cast("Any", config_stub)) + store = AdbcStore(cast("Any", config_stub)) store._dialect = dialect store._set("sess_1", b"payload", expires_in=None) diff --git a/tests/unit/adapters/test_duckdb/test_adk_store.py b/tests/unit/adapters/test_duckdb/test_adk_store.py index 7bba2d05b..966b773d6 100644 --- a/tests/unit/adapters/test_duckdb/test_adk_store.py +++ b/tests/unit/adapters/test_duckdb/test_adk_store.py @@ -9,8 +9,8 @@ import pytest import sqlspec.adapters.duckdb.adk as duckdb_adk -from sqlspec.adapters.duckdb.adk import DuckdbADKMemoryStore, DuckdbADKStore -from sqlspec.adapters.duckdb.litestar.store import DuckdbStore +from sqlspec.adapters.duckdb.adk import DuckDBADKMemoryStore, DuckDBADKStore +from sqlspec.adapters.duckdb.litestar.store import DuckDBStore from sqlspec.config import ADKConfig @@ -27,10 +27,10 @@ def _normalize_sql(sql: str) -> str: def test_duckdb_adk_config_types_adapter_local_optimizations() -> None: """DuckDB ADK optimization settings are typed on the adapter-local extension config.""" - assert hasattr(duckdb_adk, "DuckdbADKConfig") - assert hasattr(duckdb_adk, "DuckdbADKFTSOptions") - duckdb_adk_config = cast("Any", duckdb_adk.DuckdbADKConfig) - duckdb_adk_fts_options = cast("Any", duckdb_adk.DuckdbADKFTSOptions) + assert hasattr(duckdb_adk, "DuckDBADKConfig") + assert hasattr(duckdb_adk, "DuckDBADKFTSOptions") + duckdb_adk_config = cast("Any", duckdb_adk.DuckDBADKConfig) + duckdb_adk_fts_options = cast("Any", duckdb_adk.DuckDBADKFTSOptions) assert cast("Any", ADKConfig).__optional_keys__ <= duckdb_adk_config.__optional_keys__ @@ -60,7 +60,7 @@ def test_duckdb_adk_config_types_adapter_local_optimizations() -> None: def test_duckdb_adk_events_table_uses_plain_schema_by_default() -> None: """DuckDB ADK event DDL stays unchanged unless DuckDB-specific config opts in.""" - store = DuckdbADKStore(_mock_config()) + store = DuckDBADKStore(_mock_config()) sql = store._events_table_ddl() @@ -87,7 +87,7 @@ def test_duckdb_adk_events_table_uses_plain_schema_by_default() -> None: def test_duckdb_adk_events_table_applies_adapter_local_generated_columns() -> None: """DuckDB ADK generated event projection columns and indexes are opt-in.""" - store = DuckdbADKStore( + store = DuckDBADKStore( _mock_config({"enable_event_generated_columns": True, "enable_event_generated_column_indexes": True}) ) @@ -110,7 +110,7 @@ def test_duckdb_adk_events_table_applies_adapter_local_generated_columns() -> No def test_duckdb_adk_memory_fts_pragmas_use_default_options() -> None: """DuckDB ADK FTS PRAGMA rendering preserves the current default options.""" - store = DuckdbADKMemoryStore(_mock_config({"memory_use_fts": True})) + store = DuckDBADKMemoryStore(_mock_config({"memory_use_fts": True})) assert ( store._fts_index_ddl(overwrite=False) == "PRAGMA create_fts_index('adk_memory', 'id', 'content_text', " @@ -125,7 +125,7 @@ def test_duckdb_adk_memory_fts_pragmas_use_default_options() -> None: def test_duckdb_adk_memory_fts_pragmas_apply_adapter_local_options() -> None: """DuckDB ADK FTS options come from extension_config["adk"], not driver_features.""" - store = DuckdbADKMemoryStore( + store = DuckDBADKMemoryStore( _mock_config({ "memory_use_fts": True, "memory_fts_options": { @@ -177,11 +177,11 @@ def _normalized(sql: str) -> str: return " ".join(sql.split()) -def _session_store_with_connection() -> "tuple[DuckdbADKStore, _SessionListConnection]": +def _session_store_with_connection() -> "tuple[DuckDBADKStore, _SessionListConnection]": conn = _SessionListConnection() config = _mock_config() config.provide_connection = lambda *_a, **_k: conn - return DuckdbADKStore(config), conn + return DuckDBADKStore(config), conn def test_duckdb_list_sessions_binds_order_and_page() -> None: @@ -274,7 +274,7 @@ def __exit__(self, *_: Any) -> None: conn = _DeleteConn() config = _mock_config() config.provide_connection = lambda *_a, **_k: conn - store = DuckdbADKStore(config) + store = DuckDBADKStore(config) cutoff = datetime.now(UTC) assert store._delete_expired_events(cutoff) == 2 @@ -285,7 +285,7 @@ def __exit__(self, *_: Any) -> None: def test_duckdb_adk_ensure_fts_extension_loads_before_installing() -> None: """_ensure_fts_extension should try LOAD fts first and skip INSTALL fts when already installed.""" - store = DuckdbADKMemoryStore(_mock_config({"memory_use_fts": True})) + store = DuckDBADKMemoryStore(_mock_config({"memory_use_fts": True})) executed: list[str] = [] class _Conn: @@ -297,8 +297,8 @@ def execute(self, sql: str) -> None: def test_duckdb_litestar_store_uses_timestamptz_and_normalizes_utc() -> None: - """DuckdbStore DDL and queries use TIMESTAMPTZ/CURRENT_TIMESTAMP and normalize naive timestamps to UTC.""" - store = DuckdbStore(_mock_config()) + """DuckDBStore DDL and queries use TIMESTAMPTZ/CURRENT_TIMESTAMP and normalize naive timestamps to UTC.""" + store = DuckDBStore(_mock_config()) ddl = store._table_ddl() assert "TIMESTAMPTZ" in ddl naive = datetime(2026, 1, 1, 12, 0, 0) diff --git a/tests/unit/adapters/test_sqlanywhere.py b/tests/unit/adapters/test_sqlanywhere.py index 06da062cb..25387b206 100644 --- a/tests/unit/adapters/test_sqlanywhere.py +++ b/tests/unit/adapters/test_sqlanywhere.py @@ -21,7 +21,7 @@ driver_profile, ) from sqlspec.adapters.sqlanywhere.adk import SqlanywhereADKConfig, SqlanywhereADKMemoryStore, SqlanywhereADKStore -from sqlspec.adapters.sqlanywhere.events import SqlanywhereEventQueueStore, SqlanywhereSyncEventQueueStore +from sqlspec.adapters.sqlanywhere.events import SqlanywhereEventQueueStore from sqlspec.adapters.sqlanywhere.litestar import SqlanywhereStore from sqlspec.adapters.sqlanywhere.migrations import SqlanywhereSyncMigrationTracker from sqlspec.core import SQL, ParameterStyle @@ -376,7 +376,7 @@ def _select_side_effect(query: Any, *args: Any, **kwargs: Any) -> list[dict[str, def test_sqlanywhere_migration_tracker_and_extension_stores() -> None: - """SqlanywhereSyncMigrationTracker, SqlanywhereStore, SqlanywhereSyncEventQueueStore, and SqlanywhereADKStore emit SQL Anywhere DDL.""" + """SqlanywhereSyncMigrationTracker, SqlanywhereStore, SqlanywhereEventQueueStore, and SqlanywhereADKStore emit SQL Anywhere DDL.""" tracker = SqlanywhereSyncMigrationTracker(version_table_name="DBA.ddl_migrations") ddl_text = tracker._idempotent_tracking_table_ddl_text() assert "CREATE TABLE IF NOT EXISTS" in ddl_text @@ -408,8 +408,8 @@ def test_sqlanywhere_migration_tracker_and_extension_stores() -> None: drop_litestar_sql = litestar_store._drop_table_sql() assert "DROP TABLE IF EXISTS" in drop_litestar_sql[0] - event_store = SqlanywhereSyncEventQueueStore(config) - assert SqlanywhereEventQueueStore is SqlanywhereSyncEventQueueStore + event_store = SqlanywhereEventQueueStore(config) + assert SqlanywhereEventQueueStore is SqlanywhereEventQueueStore event_statements = event_store.create_statements() assert "LONG VARCHAR" in event_statements[0] assert "TIMESTAMP" in event_statements[0] diff --git a/tests/unit/adapters/test_sybase.py b/tests/unit/adapters/test_sybase.py index ef23069b8..46c75bf40 100644 --- a/tests/unit/adapters/test_sybase.py +++ b/tests/unit/adapters/test_sybase.py @@ -27,7 +27,7 @@ ) from sqlspec.adapters.sybase._typing import SybaseRawCursor from sqlspec.adapters.sybase.adk import SybaseADKConfig, SybaseADKMemoryStore, SybaseADKStore -from sqlspec.adapters.sybase.events import SybaseEventQueueStore, SybaseSyncEventQueueStore +from sqlspec.adapters.sybase.events import SybaseEventQueueStore from sqlspec.adapters.sybase.litestar import SybaseStore from sqlspec.adapters.sybase.migrations import SybaseSyncMigrationTracker from sqlspec.core import SQL, ParameterStyle @@ -411,7 +411,7 @@ def select(self, _sql: Any, **kwargs: Any) -> list[dict[str, Any]]: def test_sybase_migration_tracker_and_extension_stores() -> None: - """SybaseSyncMigrationTracker, SybaseStore, SybaseSyncEventQueueStore, and SybaseADKStore emit ASE-safe SQL.""" + """SybaseSyncMigrationTracker, SybaseStore, SybaseEventQueueStore, and SybaseADKStore emit ASE-safe SQL.""" tracker = SybaseSyncMigrationTracker(version_table_name="dbo.ddl_migrations") ddl_text = tracker._idempotent_tracking_table_ddl_text() assert "IF OBJECT_ID(" in ddl_text @@ -437,8 +437,8 @@ def test_sybase_migration_tracker_and_extension_stores() -> None: assert "BIGDATETIME" in litestar_ddl assert "GETUTCDATE()" in litestar_ddl - assert SybaseEventQueueStore is SybaseSyncEventQueueStore - event_store = SybaseSyncEventQueueStore(config) + assert SybaseEventQueueStore is SybaseEventQueueStore + event_store = SybaseEventQueueStore(config) event_stmts = event_store.create_statements() assert "TEXT" in event_stmts[0] assert "BIGDATETIME" in event_stmts[0] diff --git a/tests/unit/extensions/test_adk/test_store_instantiation.py b/tests/unit/extensions/test_adk/test_store_instantiation.py index 1e6787868..ee39d7a5b 100644 --- a/tests/unit/extensions/test_adk/test_store_instantiation.py +++ b/tests/unit/extensions/test_adk/test_store_instantiation.py @@ -29,7 +29,7 @@ "sqlspec.adapters.cockroach_asyncpg.adk.CockroachAsyncpgADKStore", "sqlspec.adapters.cockroach_psycopg.adk.CockroachPsycopgAsyncADKStore", "sqlspec.adapters.cockroach_psycopg.adk.CockroachPsycopgSyncADKStore", - "sqlspec.adapters.duckdb.adk.DuckdbADKStore", + "sqlspec.adapters.duckdb.adk.DuckDBADKStore", "sqlspec.adapters.mssql_python.adk.MssqlPythonADKStore", "sqlspec.adapters.mysqlconnector.adk.MysqlConnectorAsyncADKStore", "sqlspec.adapters.mysqlconnector.adk.MysqlConnectorSyncADKStore", @@ -55,7 +55,7 @@ "sqlspec.adapters.cockroach_asyncpg.adk.CockroachAsyncpgADKMemoryStore", "sqlspec.adapters.cockroach_psycopg.adk.CockroachPsycopgAsyncADKMemoryStore", "sqlspec.adapters.cockroach_psycopg.adk.CockroachPsycopgSyncADKMemoryStore", - "sqlspec.adapters.duckdb.adk.DuckdbADKMemoryStore", + "sqlspec.adapters.duckdb.adk.DuckDBADKMemoryStore", "sqlspec.adapters.mssql_python.adk.MssqlPythonADKMemoryStore", "sqlspec.adapters.mysqlconnector.adk.MysqlConnectorAsyncADKMemoryStore", "sqlspec.adapters.mysqlconnector.adk.MysqlConnectorSyncADKMemoryStore", @@ -77,7 +77,7 @@ "sqlspec.adapters.adbc.adk.AdbcADKStore", "sqlspec.adapters.arrow_odbc.adk.ArrowOdbcADKStore", "sqlspec.adapters.cockroach_psycopg.adk.CockroachPsycopgSyncADKStore", - "sqlspec.adapters.duckdb.adk.DuckdbADKStore", + "sqlspec.adapters.duckdb.adk.DuckDBADKStore", "sqlspec.adapters.mssql_python.adk.MssqlPythonADKStore", "sqlspec.adapters.mysqlconnector.adk.MysqlConnectorSyncADKStore", "sqlspec.adapters.oracledb.adk.OracleSyncADKStore", @@ -92,7 +92,7 @@ "sqlspec.adapters.adbc.adk.AdbcADKMemoryStore", "sqlspec.adapters.arrow_odbc.adk.ArrowOdbcADKMemoryStore", "sqlspec.adapters.cockroach_psycopg.adk.CockroachPsycopgSyncADKMemoryStore", - "sqlspec.adapters.duckdb.adk.DuckdbADKMemoryStore", + "sqlspec.adapters.duckdb.adk.DuckDBADKMemoryStore", "sqlspec.adapters.mssql_python.adk.MssqlPythonADKMemoryStore", "sqlspec.adapters.mysqlconnector.adk.MysqlConnectorSyncADKMemoryStore", "sqlspec.adapters.oracledb.adk.OracleSyncADKMemoryStore", diff --git a/tests/unit/extensions/test_extension_config_boundaries.py b/tests/unit/extensions/test_extension_config_boundaries.py index aadf54099..c276d89d9 100644 --- a/tests/unit/extensions/test_extension_config_boundaries.py +++ b/tests/unit/extensions/test_extension_config_boundaries.py @@ -5,7 +5,7 @@ from sqlspec.adapters.sqlite.adk.store import SqliteADKMemoryStore, SqliteADKStore from sqlspec.adapters.sqlite.config import SqliteConfig from sqlspec.adapters.sqlite.events.store import SqliteEventQueueStore -from sqlspec.adapters.sqlite.litestar.store import SQLiteStore +from sqlspec.adapters.sqlite.litestar.store import SqliteStore from sqlspec.exceptions import ImproperConfigurationError @@ -14,7 +14,7 @@ def test_extension_store_rejects_unused_migration_switch(extension: str) -> None with pytest.raises(ImproperConfigurationError, match="run_migrations"): config = SqliteConfig(extension_config={extension: {"run_migrations": True}}) if extension == "litestar": - SQLiteStore(config) + SqliteStore(config) elif extension == "events": SqliteEventQueueStore(config) else: diff --git a/tests/unit/extensions/test_schema_reconciliation.py b/tests/unit/extensions/test_schema_reconciliation.py index 4b9e9e991..b7c73d819 100644 --- a/tests/unit/extensions/test_schema_reconciliation.py +++ b/tests/unit/extensions/test_schema_reconciliation.py @@ -5,11 +5,11 @@ from sqlspec.adapters.sqlite.adk.store import SqliteADKMemoryStore, SqliteADKStore from sqlspec.adapters.sqlite.config import SqliteConfig from sqlspec.adapters.sqlite.events.store import SqliteEventQueueStore -from sqlspec.adapters.sqlite.litestar.store import SQLiteStore +from sqlspec.adapters.sqlite.litestar.store import SqliteStore from sqlspec.migrations.schema import SchemaTarget, ensure_schema_sync -class AdditiveSQLiteStore(SQLiteStore): +class AdditiveSQLiteStore(SqliteStore): """Session store target with one additive column.""" def _table_ddl(self) -> str: @@ -50,7 +50,7 @@ async def test_litestar_session_store_reconcile_uses_target_ddl(tmp_path: Path) connection_config={"database": str(tmp_path / "litestar.db")}, extension_config={"litestar": {"session_table": "sessions"}}, ) - await SQLiteStore(config).create_table() + await SqliteStore(config).create_table() store = AdditiveSQLiteStore(config) await store.reconcile_schema(assume_existing=True) diff --git a/tests/unit/extensions/test_storage_options.py b/tests/unit/extensions/test_storage_options.py index 9195b31aa..0f55215e4 100644 --- a/tests/unit/extensions/test_storage_options.py +++ b/tests/unit/extensions/test_storage_options.py @@ -20,7 +20,7 @@ from sqlspec.adapters.cockroach_psycopg.litestar import CockroachPsycopgAsyncStore, CockroachPsycopgSyncStore from sqlspec.adapters.sqlite.config import SqliteConfig from sqlspec.adapters.sqlite.events import SqliteEventQueueStore -from sqlspec.adapters.sqlite.litestar import SQLiteStore +from sqlspec.adapters.sqlite.litestar import SqliteStore from sqlspec.exceptions import ImproperConfigurationError @@ -138,7 +138,7 @@ def test_sqlite_litestar_pragma_profile_applied() -> None: "litestar": {"pragma_profile": True, "pragma_overrides": {"cache_size": -32000, "journal_mode": "WAL"}} } ) - store = SQLiteStore(config) + store = SqliteStore(config) connection = RecordingConnection() store.prepare_schema_sync(SimpleNamespace(connection=connection)) From 808c78b80b5e0f1c7c14641cb7b3d41970d732bc Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sat, 31 Oct 2026 21:44:24 +0000 Subject: [PATCH 004/174] fix(core,mssql_python): treat OUTPUT ... INTO as not returning client rows The operation profile marked any DML with a returning clause as row-returning, including T-SQL OUTPUT ... INTO, which writes into a table variable. mssql_python now fetches rows whenever the cursor exposes a result set, so EXEC and DELETE ... OUTPUT results are not dropped. --- sqlspec/adapters/mssql_python/driver.py | 2 +- sqlspec/core/compiler.py | 3 ++- .../adapters/test_mssql_python/test_driver.py | 16 +++++++++++++- .../unit/core/test_explain_classification.py | 22 +++++++++++++++++++ 4 files changed, 40 insertions(+), 3 deletions(-) diff --git a/sqlspec/adapters/mssql_python/driver.py b/sqlspec/adapters/mssql_python/driver.py index 7809336fb..994befff5 100644 --- a/sqlspec/adapters/mssql_python/driver.py +++ b/sqlspec/adapters/mssql_python/driver.py @@ -173,7 +173,7 @@ def dispatch_execute(self, cursor: "MssqlPythonRawCursor", statement: "SQL") -> sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) _execute_cursor(cursor, sql, prepared_parameters) - if statement.returns_rows(): + if cursor.description is not None: fetched, column_names, row_format = collect_rows( cursor.fetchall(), cursor.description, self._column_name_cache ) diff --git a/sqlspec/core/compiler.py b/sqlspec/core/compiler.py index fddba3b28..e5ca5bc37 100644 --- a/sqlspec/core/compiler.py +++ b/sqlspec/core/compiler.py @@ -1057,7 +1057,8 @@ def _operation_profile(expression: "exp.Expr | None", operation_type: "Operation returns_rows = True elif isinstance(expr, (exp.Insert, exp.Update, exp.Delete, exp.Merge)): modifies_rows = True - returns_rows = bool(expr.args.get("returning")) + returning = expr.args.get("returning") + returns_rows = bool(returning) and not (isinstance(returning, exp.Returning) and returning.args.get("into")) elif isinstance(expr, exp.TruncateTable): modifies_rows = True elif isinstance(expr, exp.Copy): diff --git a/tests/unit/adapters/test_mssql_python/test_driver.py b/tests/unit/adapters/test_mssql_python/test_driver.py index 54452a17c..3b1dc805b 100644 --- a/tests/unit/adapters/test_mssql_python/test_driver.py +++ b/tests/unit/adapters/test_mssql_python/test_driver.py @@ -4,7 +4,7 @@ from sqlspec.adapters.mssql_python.data_dictionary import MssqlPythonSyncDataDictionary from sqlspec.adapters.mssql_python.driver import MssqlPythonDriver -from sqlspec.core import StatementStack +from sqlspec.core import SQL, StatementStack class DummyCursor: @@ -142,3 +142,17 @@ def test_execute_cache_hit_collects_rows_and_driver_profile_is_registered() -> N assert result.get_data() == [{"id": 1, "name": "Ada"}] assert connection.cursor_obj.executed == [("SELECT id, name FROM dbo.users WHERE id = ?", (1,))] assert get_driver_profile("mssql_python") is not None + + +def test_dispatch_execute_fetches_rows_whenever_the_cursor_has_a_result_set() -> None: + """EXEC returns rows the statement profile cannot predict, so the cursor description decides.""" + driver = MssqlPythonDriver(cast("Any", object())) + cursor = DummyCursor() + cursor.description = [("id",)] + cursor.rows = [(1,), (2,)] + statement = SQL("EXEC dbo.list_ids", statement_config=driver.statement_config) + + result = driver.dispatch_execute(cast("Any", cursor), statement) + + assert result.selected_data is not None + assert len(result.selected_data) == 2 diff --git a/tests/unit/core/test_explain_classification.py b/tests/unit/core/test_explain_classification.py index 8988b47c5..6065f723d 100644 --- a/tests/unit/core/test_explain_classification.py +++ b/tests/unit/core/test_explain_classification.py @@ -95,3 +95,25 @@ def test_sql_explain_oracle_stays_non_row_returning() -> None: assert statement.sql == "EXPLAIN PLAN FOR SELECT 1" assert statement.returns_rows() is False + + +@pytest.mark.parametrize( + "sql", + [ + "INSERT INTO t (a) OUTPUT inserted.id INTO @ids VALUES (1)", + "UPDATE t SET a = 1 OUTPUT inserted.id INTO @ids WHERE id = 1", + "DELETE t OUTPUT deleted.id INTO @ids WHERE id = 1", + ], +) +def test_tsql_output_into_does_not_return_client_rows(sql: str) -> None: + """OUTPUT ... INTO writes into a table variable, so no rows reach the client.""" + expression = sqlglot.parse_one(sql, dialect="tsql") + + assert SQLProcessor._operation_profile(expression, SQLProcessor._operation_type(expression)).returns_rows is False + + +def test_tsql_output_without_into_returns_client_rows() -> None: + """A bare OUTPUT clause streams the affected rows to the client.""" + expression = sqlglot.parse_one("INSERT INTO t (a) OUTPUT inserted.id VALUES (1)", dialect="tsql") + + assert SQLProcessor._operation_profile(expression, SQLProcessor._operation_type(expression)).returns_rows is True From 50447822b969ff67b1186f8e1b2da4286177b3f7 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sat, 31 Oct 2026 21:47:58 +0000 Subject: [PATCH 005/174] fix(config): declare correlation settings on the Flask and Starlette extension configs FlaskConfig and StarletteConfig (and so FastAPIConfig) omitted enable_correlation_middleware, correlation_header, correlation_headers, and auto_trace_headers even though both extensions read them. Declare them and add a contract test that every key an extension reads is declared on its config type. --- sqlspec/config.py | 24 +++++++++++++++++ .../test_extension_config_boundaries.py | 26 +++++++++++++++++++ 2 files changed, 50 insertions(+) diff --git a/sqlspec/config.py b/sqlspec/config.py index 58e477a5d..587ee8fcd 100644 --- a/sqlspec/config.py +++ b/sqlspec/config.py @@ -349,6 +349,18 @@ class FlaskConfig(TypedDict): Set to ``False`` to explicitly disable this behavior. """ + enable_correlation_middleware: NotRequired[bool] + """Enable request correlation ID middleware. Default: False""" + + correlation_header: NotRequired[str] + """HTTP header to read the request correlation ID from when middleware is enabled. Default: ``X-Request-ID``""" + + correlation_headers: NotRequired[tuple[str, ...] | list[str]] + """Additional HTTP headers to read as correlation ID fallbacks.""" + + auto_trace_headers: NotRequired[bool] + """Read standard trace context headers as correlation ID fallbacks. Default: True.""" + class LitestarConfig(TypedDict): """Configuration options for Litestar SQLSpec plugin. @@ -485,6 +497,18 @@ class StarletteConfig(TypedDict): Set to 'fastapi' when using FastAPI. """ + enable_correlation_middleware: NotRequired[bool] + """Enable request correlation ID middleware. Default: False""" + + correlation_header: NotRequired[str] + """HTTP header to read the request correlation ID from when middleware is enabled. Default: ``X-Request-ID``""" + + correlation_headers: NotRequired[tuple[str, ...] | list[str]] + """Additional HTTP headers to read as correlation ID fallbacks.""" + + auto_trace_headers: NotRequired[bool] + """Read standard trace context headers as correlation ID fallbacks. Default: True.""" + class FastAPIConfig(StarletteConfig): """Configuration options for FastAPI SQLSpec extension. diff --git a/tests/unit/extensions/test_extension_config_boundaries.py b/tests/unit/extensions/test_extension_config_boundaries.py index c276d89d9..d0fed0643 100644 --- a/tests/unit/extensions/test_extension_config_boundaries.py +++ b/tests/unit/extensions/test_extension_config_boundaries.py @@ -1,11 +1,16 @@ """Runtime boundaries for shared and adapter-specific extension settings.""" +import re +from pathlib import Path + import pytest +import sqlspec.extensions from sqlspec.adapters.sqlite.adk.store import SqliteADKMemoryStore, SqliteADKStore from sqlspec.adapters.sqlite.config import SqliteConfig from sqlspec.adapters.sqlite.events.store import SqliteEventQueueStore from sqlspec.adapters.sqlite.litestar.store import SqliteStore +from sqlspec.config import FastAPIConfig, FlaskConfig, LitestarConfig, SanicConfig, StarletteConfig from sqlspec.exceptions import ImproperConfigurationError @@ -35,3 +40,24 @@ def test_sqlite_adk_rejects_postgres_memory_tuning(option: str, value: object) - with pytest.raises(ImproperConfigurationError, match=option): config = SqliteConfig(extension_config={"adk": {option: value}}) SqliteADKMemoryStore(config) + + +_FRAMEWORK_CONFIG_READS = ( + ("flask", "flask_config", FlaskConfig), + ("starlette", "framework_config", StarletteConfig), + ("fastapi", "framework_config", FastAPIConfig), + ("sanic", "framework_config", SanicConfig), + ("litestar", "litestar_config", LitestarConfig), +) + + +@pytest.mark.parametrize(("framework", "variable", "typed_dict"), _FRAMEWORK_CONFIG_READS) +def test_framework_extension_reads_only_declared_config_keys(framework: str, variable: str, typed_dict: type) -> None: + """Every key an extension reads from its framework config is declared on that framework's TypedDict.""" + package = Path(sqlspec.extensions.__file__).parent / framework + source = "".join(path.read_text() for path in sorted(package.rglob("*.py"))) + read_keys = set(re.findall(rf"\b{variable}\.get\(\"(\w+)\"", source)) + declared = set(typed_dict.__required_keys__) | set(typed_dict.__optional_keys__) + + assert read_keys, f"no {variable}.get(...) reads found for {framework}" + assert read_keys <= declared, f"{framework} reads undeclared keys: {sorted(read_keys - declared)}" From de6bb18d84c1181bf1ff3620751edc0b786464e9 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sat, 31 Oct 2026 21:50:35 +0000 Subject: [PATCH 006/174] feat(extensions): enable correlation middleware by default in every framework extension Flask, Starlette, FastAPI, and Sanic defaulted enable_correlation_middleware to False while Litestar defaulted to True. All five now default to True; set the key to False to opt out. --- sqlspec/config.py | 4 ++-- sqlspec/extensions/fastapi/extension.py | 2 +- sqlspec/extensions/flask/_state.py | 2 +- sqlspec/extensions/flask/extension.py | 2 +- sqlspec/extensions/sanic/_state.py | 2 +- sqlspec/extensions/sanic/extension.py | 2 +- sqlspec/extensions/starlette/_state.py | 2 +- sqlspec/extensions/starlette/extension.py | 2 +- tests/unit/extensions/test_flask/test_correlation.py | 6 +++--- tests/unit/extensions/test_sanic/test_middleware.py | 8 ++++++-- 10 files changed, 18 insertions(+), 14 deletions(-) diff --git a/sqlspec/config.py b/sqlspec/config.py index 587ee8fcd..758089efb 100644 --- a/sqlspec/config.py +++ b/sqlspec/config.py @@ -350,7 +350,7 @@ class FlaskConfig(TypedDict): """ enable_correlation_middleware: NotRequired[bool] - """Enable request correlation ID middleware. Default: False""" + """Enable request correlation ID middleware. Default: True""" correlation_header: NotRequired[str] """HTTP header to read the request correlation ID from when middleware is enabled. Default: ``X-Request-ID``""" @@ -498,7 +498,7 @@ class StarletteConfig(TypedDict): """ enable_correlation_middleware: NotRequired[bool] - """Enable request correlation ID middleware. Default: False""" + """Enable request correlation ID middleware. Default: True""" correlation_header: NotRequired[str] """HTTP header to read the request correlation ID from when middleware is enabled. Default: ``X-Request-ID``""" diff --git a/sqlspec/extensions/fastapi/extension.py b/sqlspec/extensions/fastapi/extension.py index 3d61e669d..414c714eb 100644 --- a/sqlspec/extensions/fastapi/extension.py +++ b/sqlspec/extensions/fastapi/extension.py @@ -221,7 +221,7 @@ def _extract_extension_settings(self, config: Any) -> "dict[str, Any]": "extra_commit_statuses": framework_config.get("extra_commit_statuses"), "extra_rollback_statuses": framework_config.get("extra_rollback_statuses"), "disable_di": framework_config.get("disable_di", False), - "enable_correlation_middleware": framework_config.get("enable_correlation_middleware", False), + "enable_correlation_middleware": framework_config.get("enable_correlation_middleware", True), "correlation_header": framework_config.get("correlation_header", "x-request-id"), "correlation_headers": tuple(correlation_headers) if correlation_headers is not None else None, "auto_trace_headers": framework_config.get("auto_trace_headers", True), diff --git a/sqlspec/extensions/flask/_state.py b/sqlspec/extensions/flask/_state.py index 84fce2ae8..d214b007f 100644 --- a/sqlspec/extensions/flask/_state.py +++ b/sqlspec/extensions/flask/_state.py @@ -28,7 +28,7 @@ class FlaskConfigState: extra_rollback_statuses: "set[int] | None" is_async: bool disable_di: bool - enable_correlation_middleware: bool = False + enable_correlation_middleware: bool = True correlation_header: str = "x-request-id" correlation_headers: "tuple[str, ...] | None" = None auto_trace_headers: bool = True diff --git a/sqlspec/extensions/flask/extension.py b/sqlspec/extensions/flask/extension.py index b01acfd0c..b469b75ad 100644 --- a/sqlspec/extensions/flask/extension.py +++ b/sqlspec/extensions/flask/extension.py @@ -218,7 +218,7 @@ def _config_state(self, config: Any) -> FlaskConfigState: extra_rollback_statuses = flask_config.get("extra_rollback_statuses") disable_di = flask_config.get("disable_di", False) - enable_correlation = flask_config.get("enable_correlation_middleware", False) + enable_correlation = flask_config.get("enable_correlation_middleware", True) correlation_header = flask_config.get("correlation_header", "x-request-id") correlation_headers = flask_config.get("correlation_headers") if correlation_headers is not None: diff --git a/sqlspec/extensions/sanic/_state.py b/sqlspec/extensions/sanic/_state.py index 81e68932f..ad4ea2569 100644 --- a/sqlspec/extensions/sanic/_state.py +++ b/sqlspec/extensions/sanic/_state.py @@ -30,7 +30,7 @@ class SanicConfigState: extra_commit_statuses: "set[int] | None" extra_rollback_statuses: "set[int] | None" disable_di: bool - enable_correlation_middleware: bool = False + enable_correlation_middleware: bool = True correlation_header: str = "x-request-id" correlation_headers: "tuple[str, ...] | None" = None auto_trace_headers: bool = True diff --git a/sqlspec/extensions/sanic/extension.py b/sqlspec/extensions/sanic/extension.py index 2f20ecd4c..a23ca30dc 100644 --- a/sqlspec/extensions/sanic/extension.py +++ b/sqlspec/extensions/sanic/extension.py @@ -141,7 +141,7 @@ def _extract_extension_settings(self, config: Any) -> "dict[str, Any]": "extra_commit_statuses": framework_config.get("extra_commit_statuses"), "extra_rollback_statuses": framework_config.get("extra_rollback_statuses"), "disable_di": framework_config.get("disable_di", False), - "enable_correlation_middleware": framework_config.get("enable_correlation_middleware", False), + "enable_correlation_middleware": framework_config.get("enable_correlation_middleware", True), "correlation_header": framework_config.get("correlation_header", "x-request-id"), "correlation_headers": tuple(correlation_headers) if correlation_headers is not None else None, "auto_trace_headers": framework_config.get("auto_trace_headers", True), diff --git a/sqlspec/extensions/starlette/_state.py b/sqlspec/extensions/starlette/_state.py index 4f6d56f15..7f1b28ead 100644 --- a/sqlspec/extensions/starlette/_state.py +++ b/sqlspec/extensions/starlette/_state.py @@ -29,7 +29,7 @@ class SQLSpecConfigState: extra_commit_statuses: "set[int] | None" extra_rollback_statuses: "set[int] | None" disable_di: bool - enable_correlation_middleware: bool = False + enable_correlation_middleware: bool = True correlation_header: str = "x-request-id" correlation_headers: "tuple[str, ...] | None" = None auto_trace_headers: bool = True diff --git a/sqlspec/extensions/starlette/extension.py b/sqlspec/extensions/starlette/extension.py index 89220a40d..957c89e82 100644 --- a/sqlspec/extensions/starlette/extension.py +++ b/sqlspec/extensions/starlette/extension.py @@ -212,7 +212,7 @@ def _extract_extension_settings(self, config: Any) -> "dict[str, Any]": "extra_commit_statuses": framework_config.get("extra_commit_statuses"), "extra_rollback_statuses": framework_config.get("extra_rollback_statuses"), "disable_di": framework_config.get("disable_di", False), - "enable_correlation_middleware": framework_config.get("enable_correlation_middleware", False), + "enable_correlation_middleware": framework_config.get("enable_correlation_middleware", True), "correlation_header": framework_config.get("correlation_header", "x-request-id"), "correlation_headers": tuple(correlation_headers) if correlation_headers is not None else None, "auto_trace_headers": framework_config.get("auto_trace_headers", True), diff --git a/tests/unit/extensions/test_flask/test_correlation.py b/tests/unit/extensions/test_flask/test_correlation.py index 60b09913b..3806070f8 100644 --- a/tests/unit/extensions/test_flask/test_correlation.py +++ b/tests/unit/extensions/test_flask/test_correlation.py @@ -163,8 +163,8 @@ def test_endpoint() -> str: assert seen_correlation_id[0] is False -def test_flask_correlation_disabled_default_is_disabled() -> None: - """Correlation middleware should be disabled by default.""" +def test_flask_correlation_is_enabled_by_default() -> None: + """Correlation middleware is enabled by default, matching the other framework extensions.""" sqlspec = SQLSpec() config = SqliteConfig(connection_config={"database": ":memory:"}) sqlspec.add_config(config) @@ -180,7 +180,7 @@ def test_endpoint() -> str: with app.test_client() as client: response = client.get("/test", headers={"x-request-id": "default-test-123"}) assert response.status_code == 200 - assert seen_correlation_id[0] is False + assert seen_correlation_id[0] is True def test_flask_correlation_trace_headers_extracts_traceparent() -> None: diff --git a/tests/unit/extensions/test_sanic/test_middleware.py b/tests/unit/extensions/test_sanic/test_middleware.py index 785f182f2..7171becc9 100644 --- a/tests/unit/extensions/test_sanic/test_middleware.py +++ b/tests/unit/extensions/test_sanic/test_middleware.py @@ -162,7 +162,9 @@ def _make_config( ) -> MagicMock: config = MagicMock() config.supports_connection_pooling = supports_connection_pooling - config.extension_config = {"sanic": {"commit_mode": commit_mode, "disable_di": disable_di}} + config.extension_config = { + "sanic": {"commit_mode": commit_mode, "disable_di": disable_di, "enable_correlation_middleware": False} + } config.provide_connection = MagicMock(return_value=_AsyncConnectionManager(connection)) config.create_connection = AsyncMock(return_value=connection) config.statement_config = {} @@ -174,7 +176,9 @@ def _make_request(*, pool: object | None) -> SimpleNamespace: app_ctx = SimpleNamespace() if pool is not None: app_ctx.db_pool = pool - return SimpleNamespace(app=SimpleNamespace(ctx=app_ctx), ctx=SimpleNamespace(db_session_instance=object())) + return SimpleNamespace( + app=SimpleNamespace(ctx=app_ctx), ctx=SimpleNamespace(db_session_instance=object()), headers={} + ) class _AsyncConnectionManager: From a29b1d3a26d955f0e10abacee211d7bae85c9d11 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sat, 31 Oct 2026 22:47:02 +0000 Subject: [PATCH 007/174] fix(sqlite,aiosqlite): correct extended constraint error codes for not null and check Update SQLite and AioSQLite extended constraint error codes so NOT NULL uses code 1299 and CHECK uses code 275. This prevents trigger aborts from being falsely classified as not-null violations. --- sqlspec/adapters/aiosqlite/core.py | 4 +-- sqlspec/adapters/sqlite/core.py | 4 +-- .../test_aiosqlite/test_exception_mapping.py | 32 ++++++++++++++++++- .../test_sqlite/test_exception_mapping.py | 32 ++++++++++++++++++- 4 files changed, 66 insertions(+), 6 deletions(-) diff --git a/sqlspec/adapters/aiosqlite/core.py b/sqlspec/adapters/aiosqlite/core.py index 6441691ad..5497da1ee 100644 --- a/sqlspec/adapters/aiosqlite/core.py +++ b/sqlspec/adapters/aiosqlite/core.py @@ -95,8 +95,8 @@ SQLITE_CONSTRAINT_PRIMARYKEY_CODE = 1555 SQLITE_CONSTRAINT_ROWID_CODE = 2579 SQLITE_CONSTRAINT_FOREIGNKEY_CODE = 787 -SQLITE_CONSTRAINT_NOTNULL_CODE = 1811 -SQLITE_CONSTRAINT_CHECK_CODE = 531 +SQLITE_CONSTRAINT_NOTNULL_CODE = 1299 +SQLITE_CONSTRAINT_CHECK_CODE = 275 SQLITE_CONSTRAINT_DATATYPE_CODE = 3091 SQLITE_CONSTRAINT_CODE = 19 SQLITE_CANTOPEN_CODE = 14 diff --git a/sqlspec/adapters/sqlite/core.py b/sqlspec/adapters/sqlite/core.py index fe90822bf..e78c4ce42 100644 --- a/sqlspec/adapters/sqlite/core.py +++ b/sqlspec/adapters/sqlite/core.py @@ -75,8 +75,8 @@ SQLITE_CONSTRAINT_PRIMARYKEY_CODE = 1555 SQLITE_CONSTRAINT_ROWID_CODE = 2579 SQLITE_CONSTRAINT_FOREIGNKEY_CODE = 787 -SQLITE_CONSTRAINT_NOTNULL_CODE = 1811 -SQLITE_CONSTRAINT_CHECK_CODE = 531 +SQLITE_CONSTRAINT_NOTNULL_CODE = 1299 +SQLITE_CONSTRAINT_CHECK_CODE = 275 SQLITE_CONSTRAINT_DATATYPE_CODE = 3091 SQLITE_CONSTRAINT_CODE = 19 SQLITE_CANTOPEN_CODE = 14 diff --git a/tests/unit/adapters/test_aiosqlite/test_exception_mapping.py b/tests/unit/adapters/test_aiosqlite/test_exception_mapping.py index a99e560f0..6c2b669f9 100644 --- a/tests/unit/adapters/test_aiosqlite/test_exception_mapping.py +++ b/tests/unit/adapters/test_aiosqlite/test_exception_mapping.py @@ -5,7 +5,15 @@ import pytest from sqlspec.adapters.aiosqlite.core import create_mapped_exception -from sqlspec.exceptions import DeadlockError, OperationCancelledError, PermissionDeniedError, UniqueViolationError +from sqlspec.exceptions import ( + CheckViolationError, + DeadlockError, + IntegrityError, + NotNullViolationError, + OperationCancelledError, + PermissionDeniedError, + UniqueViolationError, +) class _SqliteError(sqlite3.OperationalError): @@ -128,3 +136,25 @@ def test_primary_key_identity_takes_precedence_over_message(table: str, code: in mapped = create_mapped_exception(error) assert isinstance(mapped, UniqueViolationError) assert mapped.__cause__ is error + + +def test_not_null_error_code_maps_to_not_null_violation() -> None: + err = _SqliteIntegrityError("NOT NULL constraint failed: t.val", 1299, None) + result = create_mapped_exception(err) + assert isinstance(result, NotNullViolationError) + assert result.__cause__ is err + + +def test_check_error_code_maps_to_check_violation() -> None: + err = _SqliteIntegrityError("CHECK constraint failed: val > 0", 275, None) + result = create_mapped_exception(err) + assert isinstance(result, CheckViolationError) + assert result.__cause__ is err + + +def test_trigger_abort_does_not_map_to_not_null_violation() -> None: + err = _SqliteIntegrityError("trigger aborted", 1811, None) + result = create_mapped_exception(err) + assert not isinstance(result, NotNullViolationError) + assert isinstance(result, IntegrityError) + assert result.__cause__ is err diff --git a/tests/unit/adapters/test_sqlite/test_exception_mapping.py b/tests/unit/adapters/test_sqlite/test_exception_mapping.py index c77a81b57..4821c5329 100644 --- a/tests/unit/adapters/test_sqlite/test_exception_mapping.py +++ b/tests/unit/adapters/test_sqlite/test_exception_mapping.py @@ -5,7 +5,15 @@ import pytest from sqlspec.adapters.sqlite.core import create_mapped_exception -from sqlspec.exceptions import DeadlockError, OperationCancelledError, PermissionDeniedError, UniqueViolationError +from sqlspec.exceptions import ( + CheckViolationError, + DeadlockError, + IntegrityError, + NotNullViolationError, + OperationCancelledError, + PermissionDeniedError, + UniqueViolationError, +) class _SqliteError(sqlite3.OperationalError): @@ -88,3 +96,25 @@ def test_primary_key_identity_takes_precedence_over_message(table: str, code: in mapped = create_mapped_exception(error) assert isinstance(mapped, UniqueViolationError) assert mapped.__cause__ is error + + +def test_not_null_error_code_maps_to_not_null_violation() -> None: + err = _SqliteIntegrityError("NOT NULL constraint failed: t.val", 1299, None) + result = create_mapped_exception(err) + assert isinstance(result, NotNullViolationError) + assert result.__cause__ is err + + +def test_check_error_code_maps_to_check_violation() -> None: + err = _SqliteIntegrityError("CHECK constraint failed: val > 0", 275, None) + result = create_mapped_exception(err) + assert isinstance(result, CheckViolationError) + assert result.__cause__ is err + + +def test_trigger_abort_does_not_map_to_not_null_violation() -> None: + err = _SqliteIntegrityError("trigger aborted", 1811, None) + result = create_mapped_exception(err) + assert not isinstance(result, NotNullViolationError) + assert isinstance(result, IntegrityError) + assert result.__cause__ is err From 42f387124611ae84a4a2320617fa0478300f0269 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sat, 31 Oct 2026 22:47:03 +0000 Subject: [PATCH 008/174] fix(aiosqlite): fix pool reset retirement, cache hit exception isolation, and error mapping Prevent leaking dirty connections in connection pool release by retiring connections on rollback errors. Delegate cached query execution to cursor method to preserve ambient exceptions and format wire SQL, and map programming/interface errors and column names containing substring keywords accurately. --- sqlspec/adapters/aiosqlite/core.py | 27 ++++-- sqlspec/adapters/aiosqlite/driver.py | 84 +++++++++---------- sqlspec/adapters/aiosqlite/pool.py | 3 +- .../adapters/test_aiosqlite/test_driver.py | 20 +++++ .../test_aiosqlite/test_exception_mapping.py | 24 ++++++ .../unit/adapters/test_aiosqlite/test_pool.py | 20 +++++ 6 files changed, 125 insertions(+), 53 deletions(-) diff --git a/sqlspec/adapters/aiosqlite/core.py b/sqlspec/adapters/aiosqlite/core.py index 5497da1ee..29a66a575 100644 --- a/sqlspec/adapters/aiosqlite/core.py +++ b/sqlspec/adapters/aiosqlite/core.py @@ -477,13 +477,19 @@ def create_mapped_exception(error: BaseException, *, logger: Any | None = None) return _create_aiosqlite_error(error, error_code, DeadlockError, "database busy") if primary_code == SQLITE_LOCKED_CODE or (isinstance(error_name, str) and error_name.startswith("SQLITE_LOCKED")): return _create_aiosqlite_error(error, error_code, DeadlockError, "database locked") - if "locked" in error_msg or "busy" in error_msg: - return _create_aiosqlite_error(error, error_code or 0, DeadlockError, "database locked") + if error_code is None and ( + "database is locked" in error_msg + or "database table is locked" in error_msg + or "database locked" in error_msg + or "database is busy" in error_msg + or "database busy" in error_msg + ): + return _create_aiosqlite_error(error, 0, DeadlockError, "database locked") if primary_code == SQLITE_INTERRUPT_CODE or error_name == "SQLITE_INTERRUPT": return _create_aiosqlite_error(error, error_code, OperationCancelledError, "query interrupted") - if "interrupt" in error_msg: - return _create_aiosqlite_error(error, error_code or 0, OperationCancelledError, "query interrupted") + if error_code is None and ("query was interrupted" in error_msg or "interrupted" in error_msg): + return _create_aiosqlite_error(error, 0, OperationCancelledError, "query interrupted") if primary_code == SQLITE_PERM_CODE or error_name == "SQLITE_PERM": return _create_aiosqlite_error(error, error_code, PermissionDeniedError, "permission denied") @@ -491,8 +497,13 @@ def create_mapped_exception(error: BaseException, *, logger: Any | None = None) isinstance(error_name, str) and error_name.startswith("SQLITE_READONLY") ): return _create_aiosqlite_error(error, error_code, PermissionDeniedError, "database is read-only") - if "permission denied" in error_msg or "readonly" in error_msg: - return _create_aiosqlite_error(error, error_code or 0, PermissionDeniedError, "permission denied") + if error_code is None and ( + "attempt to write a readonly database" in error_msg + or "readonly database" in error_msg + or "permission denied" in error_msg + or "access permission denied" in error_msg + ): + return _create_aiosqlite_error(error, 0, PermissionDeniedError, "permission denied") if primary_code == SQLITE_CANTOPEN_CODE or ( isinstance(error_name, str) and error_name.startswith("SQLITE_CANTOPEN") @@ -506,6 +517,10 @@ def create_mapped_exception(error: BaseException, *, logger: Any | None = None) return _create_aiosqlite_error(error, error_code, DataError, "data error") if "syntax" in error_msg: return _create_aiosqlite_error(error, error_code, SQLParsingError, "SQL syntax error") + if isinstance(error, sqlite3.ProgrammingError): + return _create_aiosqlite_error(error, error_code, DataError, "programming error") + if isinstance(error, sqlite3.InterfaceError): + return _create_aiosqlite_error(error, error_code, OperationalError, "interface error") if primary_code == 1 or isinstance(error, sqlite3.OperationalError): return _create_aiosqlite_error(error, error_code, OperationalError, "operational error") if isinstance(error, sqlite3.DataError): diff --git a/sqlspec/adapters/aiosqlite/driver.py b/sqlspec/adapters/aiosqlite/driver.py index 98a1f4697..4cf523045 100644 --- a/sqlspec/adapters/aiosqlite/driver.py +++ b/sqlspec/adapters/aiosqlite/driver.py @@ -38,7 +38,7 @@ from collections.abc import Mapping, Sequence from sqlspec.adapters.aiosqlite._typing import AiosqliteConnection - from sqlspec.core import SQL, SQLResult, StatementConfig + from sqlspec.core import SQL, CachedQuery, SQLResult, StatementConfig from sqlspec.core.compiler import OperationType from sqlspec.driver import ExecutionResult from sqlspec.storage import StorageBridgeJob, StorageDestination, StorageFormat, StorageTelemetry @@ -360,53 +360,47 @@ def resolve_rowcount(self, cursor: Any) -> int: """Resolve rowcount from aiosqlite cursor for the direct execution path.""" return resolve_rowcount(cursor) - async def _execute_cache_hit( - self, sql: str, params: "tuple[Any, ...] | list[Any] | dict[str, Any]", cached: Any + async def _execute_cache_hit_with_cursor( + self, sql: str, params: "tuple[Any, ...] | list[Any] | dict[str, Any]", cached: "CachedQuery" ) -> "SQLResult": - """Execute cached queries through the async cursor fast path.""" + """Execute cached queries through the aiosqlite worker thread fast path.""" + wire_sql = self._format_wire_sql(cached.compiled_sql) normalized_parameters = normalize_execute_parameters(params) - exc_handler = self.handle_database_exceptions() - result: SQLResult | None = None self._invalidate_rowid_target_cache(cached.operation_type) - async with exc_handler: - if cached.operation_profile.returns_rows: - fetched_data, description, _affected_rows, last_inserted_id = await run_on_worker_thread( - self.connection, - _execute_fetchall_with_metadata, - self.connection, - cached.compiled_sql, - normalized_parameters, - cached.operation_type, - cached.processed_state.parsed_expression, - self._rowid_target_cache, - ) - data, column_names, row_count = collect_rows(fetched_data, description) - result = self.build_cached_select_result( - sql, - params, - cached, - selected_data=data, - column_names=column_names, - data_row_count=row_count, - row_format=resolve_row_format(data), - last_inserted_id=last_inserted_id, - ) - else: - affected_rows, last_inserted_id = await run_on_worker_thread( - self.connection, - _execute_and_resolve_metadata, - self.connection, - cached.compiled_sql, - normalized_parameters, - cached.operation_type, - cached.processed_state.parsed_expression, - self._rowid_target_cache, - ) - result = DMLResult(cached.operation_type, affected_rows, last_inserted_id) - - self._check_pending_exception(exc_handler) - assert result is not None - return result + if cached.operation_profile.returns_rows: + fetched_data, description, _affected_rows, last_inserted_id = await run_on_worker_thread( + self.connection, + _execute_fetchall_with_metadata, + self.connection, + wire_sql, + normalized_parameters, + cached.operation_type, + cached.processed_state.parsed_expression, + self._rowid_target_cache, + ) + data, column_names, row_count = collect_rows(fetched_data, description) + return self.build_cached_select_result( + sql, + params, + cached, + selected_data=data, + column_names=column_names, + data_row_count=row_count, + row_format=resolve_row_format(data), + last_inserted_id=last_inserted_id, + compiled_sql=wire_sql, + ) + affected_rows, last_inserted_id = await run_on_worker_thread( + self.connection, + _execute_and_resolve_metadata, + self.connection, + wire_sql, + normalized_parameters, + cached.operation_type, + cached.processed_state.parsed_expression, + self._rowid_target_cache, + ) + return DMLResult(cached.operation_type, affected_rows, last_inserted_id) def _invalidate_rowid_target_cache(self, operation_type: "OperationType") -> None: if operation_type not in {"SELECT", "INSERT", "UPDATE", "DELETE"}: diff --git a/sqlspec/adapters/aiosqlite/pool.py b/sqlspec/adapters/aiosqlite/pool.py index 2c166cc9f..d69530830 100644 --- a/sqlspec/adapters/aiosqlite/pool.py +++ b/sqlspec/adapters/aiosqlite/pool.py @@ -419,8 +419,7 @@ async def release(self, connection: AiosqlitePoolConnection) -> None: try: raw_conn = connection.connection if getattr(raw_conn, "in_transaction", False): - with suppress(Exception): - await end_transaction(raw_conn, commit=False) + await end_transaction(raw_conn, commit=False) connection.idle_since = time.time() self._queue.put_nowait(connection) except Exception as e: diff --git a/tests/unit/adapters/test_aiosqlite/test_driver.py b/tests/unit/adapters/test_aiosqlite/test_driver.py index daeb6b842..ffed6b300 100644 --- a/tests/unit/adapters/test_aiosqlite/test_driver.py +++ b/tests/unit/adapters/test_aiosqlite/test_driver.py @@ -428,3 +428,23 @@ async def test_aiosqlite_pool_tracks_in_flight_creations_and_connection_registry finally: await pool.close() assert pool.size() == 0 + + +async def test_aiosqlite_cache_hit_preserves_ambient_exception_and_formats_wire_sql() -> None: + import sys + + connection = await aiosqlite.connect(":memory:") + try: + driver = AiosqliteDriver(connection=connection) + await driver.execute("SELECT 1 AS value") + + ambient_preserved = False + try: + raise RuntimeError("ambient error") + except RuntimeError as exc: + result = await driver.execute("SELECT 1 AS value") + ambient_preserved = sys.exc_info()[1] is exc + assert result.data == [(1,)] + assert ambient_preserved + finally: + await connection.close() diff --git a/tests/unit/adapters/test_aiosqlite/test_exception_mapping.py b/tests/unit/adapters/test_aiosqlite/test_exception_mapping.py index 6c2b669f9..969ea9768 100644 --- a/tests/unit/adapters/test_aiosqlite/test_exception_mapping.py +++ b/tests/unit/adapters/test_aiosqlite/test_exception_mapping.py @@ -7,9 +7,11 @@ from sqlspec.adapters.aiosqlite.core import create_mapped_exception from sqlspec.exceptions import ( CheckViolationError, + DataError, DeadlockError, IntegrityError, NotNullViolationError, + OperationalError, OperationCancelledError, PermissionDeniedError, UniqueViolationError, @@ -158,3 +160,25 @@ def test_trigger_abort_does_not_map_to_not_null_violation() -> None: assert not isinstance(result, NotNullViolationError) assert isinstance(result, IntegrityError) assert result.__cause__ is err + + +def test_no_such_column_with_locked_name_maps_to_operational_error() -> None: + err = sqlite3.OperationalError("no such column: is_locked") + result = create_mapped_exception(err) + assert isinstance(result, OperationalError) + assert not isinstance(result, DeadlockError) + assert result.__cause__ is err + + +def test_programming_error_maps_to_data_error() -> None: + err = sqlite3.ProgrammingError("cannot operate on a closed database") + result = create_mapped_exception(err) + assert isinstance(result, DataError) + assert result.__cause__ is err + + +def test_interface_error_maps_to_operational_error() -> None: + err = sqlite3.InterfaceError("error binding parameter") + result = create_mapped_exception(err) + assert isinstance(result, OperationalError) + assert result.__cause__ is err diff --git a/tests/unit/adapters/test_aiosqlite/test_pool.py b/tests/unit/adapters/test_aiosqlite/test_pool.py index 403f1db80..434473029 100644 --- a/tests/unit/adapters/test_aiosqlite/test_pool.py +++ b/tests/unit/adapters/test_aiosqlite/test_pool.py @@ -462,3 +462,23 @@ async def resolve() -> _FallbackConnection: return self._connection return resolve().__await__() + + +async def test_release_retires_connection_when_rollback_fails(monkeypatch: pytest.MonkeyPatch) -> None: + pool = AiosqliteConnectionPool({"database": ":memory:", "timeout": 30.0}) + pool_conn = await pool.acquire() + assert pool._queue.qsize() == 0 + + await pool_conn.connection.execute("CREATE TABLE test_tx (id INT)") + await pool_conn.connection.execute("INSERT INTO test_tx VALUES (1)") + assert getattr(pool_conn.connection, "in_transaction", False) + + async def _failing_end_transaction(_raw_conn: Any, *, commit: bool = False) -> None: + raise sqlite3.OperationalError("disk I/O error during rollback") + + monkeypatch.setattr(pool_module, "end_transaction", _failing_end_transaction) + + await pool.release(pool_conn) + assert pool._queue.qsize() == 0 + assert not pool_conn.is_healthy + await pool.close() From 5ba14d177b795e2bc3114133b76c6e7a2eb85466 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sat, 31 Oct 2026 22:47:03 +0000 Subject: [PATCH 009/174] fix(aiosqlite): use weakref for connection tracking, add float32 embeddings and parameterize event limit Use WeakSet to track configured connections across ADK and Litestar stores to avoid ABA pragma skipping after address reuse. Implement binary float32 packing and similarity search for memory embeddings, and parameterize LIMIT in ADK get_events. --- sqlspec/adapters/aiosqlite/adk/store.py | 113 ++++++++++++++---- sqlspec/adapters/aiosqlite/litestar/store.py | 13 +- .../test_aiosqlite/test_adk_config.py | 94 +++++++++++++++ 3 files changed, 188 insertions(+), 32 deletions(-) diff --git a/sqlspec/adapters/aiosqlite/adk/store.py b/sqlspec/adapters/aiosqlite/adk/store.py index d37e75016..8997b932e 100644 --- a/sqlspec/adapters/aiosqlite/adk/store.py +++ b/sqlspec/adapters/aiosqlite/adk/store.py @@ -1,8 +1,11 @@ """Aiosqlite async ADK store for Google Agent Development Kit session/event storage.""" +import array import contextlib import inspect +import math import re +import weakref from collections.abc import Mapping from datetime import UTC, datetime, timedelta from typing import TYPE_CHECKING, Any, Final, Literal, NotRequired, cast @@ -74,7 +77,7 @@ class AiosqliteADKStore(BaseAsyncADKStore["AiosqliteConfig"]): config: AiosqliteConfig with extension_config["adk"] settings. """ - __slots__ = ("_configured_connection_ids", "_pragma_overrides") + __slots__ = ("_configured_connections", "_pragma_overrides") def __init__(self, config: "AiosqliteConfig") -> None: """Initialize Aiosqlite ADK store. @@ -84,7 +87,7 @@ def __init__(self, config: "AiosqliteConfig") -> None: """ super().__init__(config) self._pragma_overrides = _pragma_overrides(config) - self._configured_connection_ids: set[int] = set() + self._configured_connections: weakref.WeakSet[Any] = weakref.WeakSet() async def create_tables(self) -> None: """Create both sessions and events tables if they don't exist.""" @@ -474,7 +477,11 @@ async def get_events( params.append(_datetime_to_julian(after_timestamp)) where_clause = " AND ".join(where_clauses) - limit_clause = f" LIMIT {limit}" if limit else "" + if limit is not None: + limit_clause = " LIMIT ?" + params.append(int(limit)) + else: + limit_clause = "" sql = f""" SELECT id, app_name, user_id, session_id, invocation_id, timestamp, event_data @@ -688,7 +695,7 @@ async def set_metadata(self, key: str, value: str) -> None: async def _ensure_pragmas(self, connection: Any) -> None: """Apply PRAGMA optimization profile once per connection.""" - if id(connection) not in self._configured_connection_ids: + if connection not in self._configured_connections: await self._apply_pragmas(connection) async def _apply_pragmas(self, connection: Any) -> None: @@ -703,7 +710,7 @@ async def _apply_pragmas(self, connection: Any) -> None: await _close_cursor(await connection.execute("PRAGMA journal_size_limit = 67108864")) for pragma_name, pragma_value in self._pragma_overrides: await _close_cursor(await connection.execute(f"PRAGMA {pragma_name} = {pragma_value}")) - self._configured_connection_ids.add(id(connection)) + self._configured_connections.add(connection) async def _sessions_table_ddl(self) -> str: """Get SQLite CREATE TABLE SQL for sessions. @@ -823,7 +830,7 @@ class AiosqliteADKMemoryStore(BaseAsyncADKMemoryStore["AiosqliteConfig"]): config: AiosqliteConfig with extension_config["adk"] settings. """ - __slots__ = ("_configured_connection_ids", "_fts_options", "_pragma_overrides") + __slots__ = ("_configured_connections", "_fts_options", "_pragma_overrides") def __init__(self, config: "AiosqliteConfig") -> None: """Initialize Aiosqlite ADK memory store. @@ -834,7 +841,7 @@ def __init__(self, config: "AiosqliteConfig") -> None: super().__init__(config) self._fts_options = _fts_options(config) self._pragma_overrides = _pragma_overrides(config) - self._configured_connection_ids: set[int] = set() + self._configured_connections: weakref.WeakSet[Any] = weakref.WeakSet() async def create_tables(self) -> None: """Create the memory table and indexes if they don't exist. @@ -873,8 +880,8 @@ async def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: " INSERT OR IGNORE INTO {self._memory_table} (id, session_id, app_name, user_id, scope, event_id, author, {self._owner_id_column_name}, timestamp, content_json, - content_text, metadata_json, inserted_at) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + content_text, metadata_json, inserted_at, embedding) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """ for entry in entries: scope = entry.get("scope", "user") @@ -893,13 +900,14 @@ async def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: " entry["content_text"], to_json(metadata_json) if metadata_json is not None else None, _datetime_to_julian(entry["inserted_at"]), + _pack_embedding(entry.get("embedding")), )) else: sql = f""" INSERT OR IGNORE INTO {self._memory_table} (id, session_id, app_name, user_id, scope, event_id, author, - timestamp, content_json, content_text, metadata_json, inserted_at) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + timestamp, content_json, content_text, metadata_json, inserted_at, embedding) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """ for entry in entries: scope = entry.get("scope", "user") @@ -917,6 +925,7 @@ async def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: " entry["content_text"], to_json(metadata_json) if metadata_json is not None else None, _datetime_to_julian(entry["inserted_at"]), + _pack_embedding(entry.get("embedding")), )) cursor = await conn.executemany(sql, params_list) try: @@ -935,12 +944,13 @@ async def search_entries( scope_filter: Literal["all", "user", "app"] = "all", embedding: "Sequence[float] | None" = None, ) -> "list[StoredMemory]": - """Search memory entries by text query.""" + """Search memory entries by text query or vector embedding.""" if not self._enabled: msg = "Memory store is disabled" raise RuntimeError(msg) - if not query or not query.strip(): + has_query = bool(query and query.strip()) + if not has_query and embedding is None: return [] limit_value = limit if limit is not None else self._max_results @@ -949,7 +959,7 @@ async def search_entries( async with self._config.provide_connection() as conn: await self._ensure_foreign_keys(conn) - if self._use_fts: + if has_query and self._use_fts: where_scope, scope_params = _build_sqlite_scope_clause("m.", app_name, user_id, scope_filter) fts_sql = f""" SELECT m.* FROM {self._memory_table} AS m @@ -958,7 +968,8 @@ async def search_entries( ORDER BY m.timestamp DESC LIMIT ? """ - fts_params = (*scope_params, query, limit_value) + fts_limit = limit_value if embedding is None else max(limit_value * 5, 50) + fts_params = (*scope_params, query, fts_limit) try: cursor = await conn.execute(fts_sql, fts_params) try: @@ -976,13 +987,23 @@ async def search_entries( if rows is None: where_scope, scope_params = _build_sqlite_scope_clause("", app_name, user_id, scope_filter) - sql = f""" - SELECT * FROM {self._memory_table} - WHERE {where_scope} AND content_text LIKE ? - ORDER BY timestamp DESC - LIMIT ? - """ - params = (*scope_params, f"%{query}%", limit_value) + fetch_limit = limit_value if embedding is None else max(limit_value * 5, 50) + if has_query: + sql = f""" + SELECT * FROM {self._memory_table} + WHERE {where_scope} AND content_text LIKE ? + ORDER BY timestamp DESC + LIMIT ? + """ + params = (*scope_params, f"%{query}%", fetch_limit) + else: + sql = f""" + SELECT * FROM {self._memory_table} + WHERE {where_scope} AND embedding IS NOT NULL + ORDER BY timestamp DESC + LIMIT ? + """ + params = (*scope_params, fetch_limit) try: cursor = await conn.execute(sql, params) try: @@ -1002,8 +1023,19 @@ async def search_entries( raw["inserted_at"] = _julian_to_datetime(raw["inserted_at"]) raw["content_json"] = from_json(raw["content_json"]) raw["metadata_json"] = from_json(raw["metadata_json"]) if raw["metadata_json"] else None - raw["embedding"] = None + raw["embedding"] = _unpack_embedding(raw.get("embedding")) records.append(cast("StoredMemory", raw)) + + if embedding is not None and records: + query_vec = [float(x) for x in embedding] + + def _score(r: "StoredMemory") -> float: + emb = r.get("embedding") + return _cosine_similarity(query_vec, emb) if emb is not None else -1.0 + + records.sort(key=_score, reverse=True) + records = records[:limit_value] + return records async def delete_entries_by_session(self, session_id: str) -> int: @@ -1062,7 +1094,7 @@ async def delete_entries_older_than( raise async def _ensure_foreign_keys(self, connection: Any) -> None: - if id(connection) not in self._configured_connection_ids: + if connection not in self._configured_connections: await self._enable_foreign_keys(connection) async def _enable_foreign_keys(self, connection: Any) -> None: @@ -1070,7 +1102,7 @@ async def _enable_foreign_keys(self, connection: Any) -> None: await _close_cursor(await connection.execute("PRAGMA foreign_keys = ON")) for pragma_name, pragma_value in self._pragma_overrides: await _close_cursor(await connection.execute(f"PRAGMA {pragma_name} = {pragma_value}")) - self._configured_connection_ids.add(id(connection)) + self._configured_connections.add(connection) async def _memory_table_ddl(self) -> str: """Get SQLite CREATE TABLE SQL for memory entries. @@ -1121,7 +1153,8 @@ async def _memory_table_ddl(self) -> str: content_json TEXT NOT NULL, content_text TEXT NOT NULL, metadata_json TEXT, - inserted_at REAL NOT NULL + inserted_at REAL NOT NULL, + embedding BLOB ); CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_app_scope_user_time @@ -1243,6 +1276,34 @@ def _build_sqlite_scope_clause( return f"{prefix}app_name = ? AND {prefix}scope = 'app'", (app_name,) +def _pack_embedding(raw: Any) -> bytes | None: + """Pack float sequence into float32 binary BLOB.""" + if raw is None: + return None + return array.array("f", [float(x) for x in raw]).tobytes() + + +def _unpack_embedding(blob: Any) -> list[float] | None: + """Unpack float32 binary BLOB into float list.""" + if blob is None: + return None + return list(array.array("f", blob)) + + +def _cosine_similarity(v1: "Sequence[float]", v2: "Sequence[float]") -> float: + """Compute cosine similarity between two float vectors.""" + dot = 0.0 + norm1 = 0.0 + norm2 = 0.0 + for a, b in zip(v1, v2, strict=False): + dot += a * b + norm1 += a * a + norm2 += b * b + if norm1 == 0.0 or norm2 == 0.0: + return 0.0 + return float(dot / (math.sqrt(norm1) * math.sqrt(norm2))) + + def _session_list_query( session_table: str, app_name: str, diff --git a/sqlspec/adapters/aiosqlite/litestar/store.py b/sqlspec/adapters/aiosqlite/litestar/store.py index 254461451..873455fa8 100644 --- a/sqlspec/adapters/aiosqlite/litestar/store.py +++ b/sqlspec/adapters/aiosqlite/litestar/store.py @@ -2,6 +2,7 @@ import contextlib import inspect +import weakref from datetime import UTC, datetime, timedelta from typing import TYPE_CHECKING, Any, NotRequired @@ -50,7 +51,7 @@ class AiosqliteStore(BaseSQLSpecStore["AiosqliteConfig"]): config: AiosqliteConfig instance. """ - __slots__ = ("_configured_connection_ids", "_pragma_statements") + __slots__ = ("_configured_connections", "_pragma_statements") extension_config_options = BaseSQLSpecStore.extension_config_options | frozenset({ "pragma_overrides", "pragma_profile", @@ -64,7 +65,7 @@ def __init__(self, config: "AiosqliteConfig") -> None: """ super().__init__(config) self._pragma_statements = extension_pragma_statements(config, "litestar") - self._configured_connection_ids: set[int] = set() + self._configured_connections: weakref.WeakSet[Any] = weakref.WeakSet() async def create_table(self) -> None: """Create the session table if it doesn't exist.""" @@ -74,7 +75,7 @@ async def create_table(self) -> None: sql = self._table_ddl() async with self._config.provide_session() as driver: await apply_extension_pragmas(driver.connection, self._pragma_statements) - self._configured_connection_ids.add(id(driver.connection)) + self._configured_connections.add(driver.connection) await driver.execute_script(sql) await driver.commit() self._log_table_created() @@ -83,7 +84,7 @@ async def create_table(self) -> None: async def prepare_schema_async(self, driver: Any) -> None: """Apply configured SQLite PRAGMAs before migration DDL generation.""" await apply_extension_pragmas(driver.connection, self._pragma_statements) - self._configured_connection_ids.add(id(driver.connection)) + self._configured_connections.add(driver.connection) async def get(self, key: str, renew_for: "int | timedelta | None" = None) -> "bytes | None": """Get a session value by key. @@ -251,9 +252,9 @@ async def delete_expired(self) -> int: return count async def _ensure_pragmas(self, connection: Any) -> None: - if self._pragma_statements and id(connection) not in self._configured_connection_ids: + if self._pragma_statements and connection not in self._configured_connections: await apply_extension_pragmas(connection, self._pragma_statements) - self._configured_connection_ids.add(id(connection)) + self._configured_connections.add(connection) def _table_ddl(self) -> str: """Get SQLite CREATE TABLE SQL. diff --git a/tests/unit/adapters/test_aiosqlite/test_adk_config.py b/tests/unit/adapters/test_aiosqlite/test_adk_config.py index 581196667..608cb5a3c 100644 --- a/tests/unit/adapters/test_aiosqlite/test_adk_config.py +++ b/tests/unit/adapters/test_aiosqlite/test_adk_config.py @@ -9,6 +9,7 @@ from sqlspec.adapters.aiosqlite.config import AiosqliteConfig from sqlspec.config import ADKConfig from sqlspec.exceptions import ImproperConfigurationError +from sqlspec.extensions.adk import StoredMemory class AsyncRecordingConnection: @@ -189,3 +190,96 @@ async def test_aiosqlite_list_sessions_rejects_invalid_options_before_connecting await store.list_sessions("app", **options) assert conn.calls == [] + + +async def test_aiosqlite_adk_recycled_connection_reruns_pragma_configuration() -> None: + """Verify WeakSet tracking allows recycled connections to re-run PRAGMAs.""" + import weakref + + config = AiosqliteConfig() + store = AiosqliteADKStore(config) + assert isinstance(store._configured_connections, weakref.WeakSet) + + conn1 = AsyncRecordingConnection() + await store._ensure_pragmas(conn1) + assert len(conn1.statements) == 4 + + await store._ensure_pragmas(conn1) + assert len(conn1.statements) == 4 + + conn2 = AsyncRecordingConnection() + await store._ensure_pragmas(conn2) + assert len(conn2.statements) == 4 + + +async def test_aiosqlite_adk_get_events_with_limit_binds_parameter() -> None: + """Verify get_events parameterizes LIMIT instead of string interpolation.""" + config = AiosqliteConfig(connection_config={"database": ":memory:"}) + store = AiosqliteADKStore(config) + calls: list[tuple[str, tuple[Any, ...]]] = [] + + class _MockEventsCursor: + async def fetchall(self) -> list[Any]: + return [] + + class _MockEventsConnection: + async def execute(self, sql: str, params: Any = None) -> _MockEventsCursor: + calls.append((sql, params)) + return _MockEventsCursor() + + class _AsyncContextManager: + async def __aenter__(self) -> Any: + return _MockEventsConnection() + + async def __aexit__(self, *_: Any) -> None: + return None + + config.provide_connection = MagicMock(return_value=_AsyncContextManager()) # type: ignore[method-assign] + + events = await store.get_events("app", "user", "session", limit=10) + assert events == [] + assert len(calls) >= 1 + sql, params = calls[-1] + assert "LIMIT ?" in sql + assert 10 in params + + +async def test_aiosqlite_adk_memory_store_persists_and_retrieves_float32_embeddings() -> None: + """Verify memory entries persist and retrieve float32 embeddings.""" + from datetime import UTC, datetime + from uuid import uuid4 + + config = AiosqliteConfig(connection_config={"database": ":memory:"}) + store = AiosqliteADKMemoryStore(config) + try: + await store.create_tables() + + now = datetime.now(UTC) + embedding = [0.1, 0.2, 0.3, 0.4] + record = StoredMemory( + id=str(uuid4()), + session_id="s1", + app_name="app", + user_id="u1", + scope="user", + event_id="evt-1", + author="user", + timestamp=now, + content_json={"text": "vector search test"}, + content_text="vector search test", + metadata_json=None, + inserted_at=now, + embedding=embedding, + ) + inserted = await store.insert_memory_entries([record]) + assert inserted == 1 + + results = await store.search_entries(query="vector", app_name="app", user_id="u1", embedding=embedding) + assert len(results) == 1 + assert results[0]["event_id"] == "evt-1" + assert results[0]["embedding"] is not None + assert len(results[0]["embedding"]) == 4 + for orig, rec in zip(embedding, results[0]["embedding"], strict=True): + assert abs(orig - rec) < 1e-6 + finally: + await config.close_pool() From 6f52244df53dae676beebaeff0b5b22d2dd1b83b Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sat, 31 Oct 2026 22:47:03 +0000 Subject: [PATCH 010/174] perf(aiosqlite): optimize hot-path closures, cursor creation, and metadata caching Hoist connection parameter filtering helpers and excluded keys to module level, eliminating per-call closures. Hoist cursor context in batch Arrow loading, cache version detection at class level in data dictionary, and memoize generated column DDL expressions per table. --- sqlspec/adapters/aiosqlite/core.py | 70 ++++++++++--------- sqlspec/adapters/aiosqlite/data_dictionary.py | 17 ++++- sqlspec/adapters/aiosqlite/driver.py | 25 ++++--- .../adapters/test_aiosqlite/test_config.py | 8 +++ .../adapters/test_aiosqlite/test_driver.py | 25 +++++++ 5 files changed, 96 insertions(+), 49 deletions(-) diff --git a/sqlspec/adapters/aiosqlite/core.py b/sqlspec/adapters/aiosqlite/core.py index 29a66a575..635265ef2 100644 --- a/sqlspec/adapters/aiosqlite/core.py +++ b/sqlspec/adapters/aiosqlite/core.py @@ -6,7 +6,7 @@ from collections.abc import Mapping from datetime import UTC, date, datetime, timedelta from decimal import Decimal -from typing import TYPE_CHECKING, Any, TypeVar, cast +from typing import TYPE_CHECKING, Any, Final, TypeVar, cast from sqlglot import exp @@ -47,6 +47,7 @@ "SECONDS_PER_DAY", "SQLITE_CONNECT_SUPPORTS_AUTOCOMMIT", "AiosqliteStreamSource", + "_filter_connection_params", "apply_driver_features", "apply_extension_pragmas", "build_connection_config", @@ -154,8 +155,8 @@ async def end_transaction( raw_conn = getattr(connection, "_conn", None) if isinstance(raw_conn, sqlite3.Connection) and callable(getattr(connection, "_execute", None)): conn_dict = getattr(connection, "__dict__", None) - has_method_override = isinstance(conn_dict, dict) and any( - key in conn_dict for key in ("commit", "rollback", "execute") + has_method_override = isinstance(conn_dict, dict) and ( + "commit" in conn_dict or "rollback" in conn_dict or "execute" in conn_dict ) if not has_method_override: await run_on_worker_thread( @@ -392,6 +393,37 @@ def normalize_execute_many_parameters(parameters: Any) -> Any: return parameters +_EXCLUDED_CONNECTION_KEYS: Final[frozenset[str]] = frozenset({ + "enable_foreign_keys", + "enable_optimizations", + "pool_size", + "min_size", + "connect_timeout", + "idle_timeout", + "operation_timeout", + "health_check_interval", + "extra", + "pool_min_size", + "pool_max_size", + "pool_timeout", + "pool_recycle_seconds", + "db", + "path", + "file", +}) + + +def _filter_connection_params(mapping: Mapping[str, Any]) -> dict[str, Any]: + """Filter connection parameters to only those accepted by SQLite/aiosqlite.""" + return { + key: value + for key, value in mapping.items() + if key not in _EXCLUDED_CONNECTION_KEYS + and (value is not None or key == "isolation_level") + and (key != "autocommit" or SQLITE_CONNECT_SUPPORTS_AUTOCOMMIT) + } + + def build_connection_config(connection_config: "Mapping[str, Any]") -> "dict[str, Any]": """Build connection configuration for pool creation. @@ -401,35 +433,7 @@ def build_connection_config(connection_config: "Mapping[str, Any]") -> "dict[str Returns: Dictionary with connection parameters. """ - excluded_keys = { - "enable_foreign_keys", - "enable_optimizations", - "pool_size", - "min_size", - "connect_timeout", - "idle_timeout", - "operation_timeout", - "health_check_interval", - "extra", - "pool_min_size", - "pool_max_size", - "pool_timeout", - "pool_recycle_seconds", - "db", - "path", - "file", - } - - def _filter_params(mapping: Mapping[str, Any]) -> dict[str, Any]: - return { - key: value - for key, value in mapping.items() - if key not in excluded_keys - and (value is not None or key == "isolation_level") - and (key != "autocommit" or SQLITE_CONNECT_SUPPORTS_AUTOCOMMIT) - } - - connection_parameters = _filter_params(connection_config) + connection_parameters = _filter_connection_params(connection_config) if "database" not in connection_parameters: database = ( connection_config.get("database") @@ -441,7 +445,7 @@ def _filter_params(mapping: Mapping[str, Any]) -> dict[str, Any]: connection_parameters["database"] = database extra = connection_config.get("extra") if isinstance(extra, Mapping): - connection_parameters.update(_filter_params(extra)) + connection_parameters.update(_filter_connection_params(extra)) return connection_parameters diff --git a/sqlspec/adapters/aiosqlite/data_dictionary.py b/sqlspec/adapters/aiosqlite/data_dictionary.py index b954d5022..ecdd916f5 100644 --- a/sqlspec/adapters/aiosqlite/data_dictionary.py +++ b/sqlspec/adapters/aiosqlite/data_dictionary.py @@ -4,6 +4,7 @@ from mypy_extensions import mypyc_attr +from sqlspec.adapters.aiosqlite._typing import aiosqlite_sqlite_module as sqlite3 from sqlspec.data_dictionary import ( ColumnMetadata, DDLResult, @@ -41,6 +42,7 @@ class AiosqliteDataDictionary(AsyncDataDictionaryBase): """SQLite-specific async data dictionary.""" dialect: ClassVar[str] = "sqlite" + _shared_version_info: ClassVar["VersionInfo | None"] = None def __init__(self) -> None: super().__init__() @@ -68,11 +70,16 @@ async def get_version(self, driver: "AiosqliteDriver") -> "VersionInfo | None": Returns: SQLite version information or None if detection fails. """ + if AiosqliteDataDictionary._shared_version_info is not None: + return AiosqliteDataDictionary._shared_version_info + driver_id = id(driver) if driver_id in self._version_fetch_attempted: return self._version_cache.get(driver_id) - version_value = await driver.select_value_or_none(self.get_query("version", "current")) + version_value = getattr(sqlite3, "sqlite_version", None) + if not version_value: + version_value = await driver.select_value_or_none(self.get_query("version", "current")) if not version_value: self._log_version_unavailable(type(self).dialect, "missing") self.cache_version(driver_id, None) @@ -86,6 +93,7 @@ async def get_version(self, driver: "AiosqliteDriver") -> "VersionInfo | None": self._log_version_detected(type(self).dialect, version_info) self.cache_version(driver_id, version_info) + AiosqliteDataDictionary._shared_version_info = version_info return version_info async def get_feature_flag(self, driver: "AiosqliteDriver", feature: str) -> bool: @@ -142,10 +150,13 @@ async def get_columns( rows = await driver.select(query_text, table_name=table, schema_name=schema_name) if any(str(row.get("extra") or "").endswith("generated") for row in rows): table_sql = await self._get_table_sql(driver, schema_name, table) + parsed_expressions: dict[str, dict[str, str]] = {} for row in rows: if str(row.get("extra") or "").endswith("generated"): - expressions = sqlite_generated_expressions(table_sql.get(str(row.get("table_name")))) - expression = expressions.get(str(row.get("column_name"))) + table_name_str = str(row.get("table_name")) + if table_name_str not in parsed_expressions: + parsed_expressions[table_name_str] = sqlite_generated_expressions(table_sql.get(table_name_str)) + expression = parsed_expressions[table_name_str].get(str(row.get("column_name"))) if expression is not None: row["generation_expression"] = expression return cast("list[ColumnMetadata]", rows) diff --git a/sqlspec/adapters/aiosqlite/driver.py b/sqlspec/adapters/aiosqlite/driver.py index 4cf523045..f9e6db608 100644 --- a/sqlspec/adapters/aiosqlite/driver.py +++ b/sqlspec/adapters/aiosqlite/driver.py @@ -307,20 +307,19 @@ async def load_from_arrow( try: if owns_transaction: await self.begin() - if overwrite: - statement = f"DELETE FROM {format_identifier(table)}" - async with self.with_cursor(self.connection) as cursor: + async with self.with_cursor(self.connection) as cursor: + if overwrite: + statement = f"DELETE FROM {format_identifier(table)}" await cursor.execute(statement) - for batch in arrow_table.to_batches(max_chunksize=batch_size): - pydict = batch.to_pydict() - records = list(zip(*(pydict[col] for col in columns), strict=False)) - if records: - prepared_records = ( - self.prepare_driver_parameters(records, self.statement_config, is_many=True) - if needs_prep - else records - ) - async with self.with_cursor(self.connection) as cursor: + for batch in arrow_table.to_batches(max_chunksize=batch_size): + pydict = batch.to_pydict() + records = list(zip(*(pydict[col] for col in columns), strict=False)) + if records: + prepared_records = ( + self.prepare_driver_parameters(records, self.statement_config, is_many=True) + if needs_prep + else records + ) await cursor.executemany(insert_sql, cast("Any", prepared_records)) if owns_transaction: await self.commit() diff --git a/tests/unit/adapters/test_aiosqlite/test_config.py b/tests/unit/adapters/test_aiosqlite/test_config.py index a17c48a2a..c1b8a6724 100644 --- a/tests/unit/adapters/test_aiosqlite/test_config.py +++ b/tests/unit/adapters/test_aiosqlite/test_config.py @@ -313,3 +313,11 @@ async def test_create_connection_applies_the_configured_row_factory(tmp_path: Pa await connection.close() finally: await config.close_pool() + + +def test_filter_connection_params_with_hoisted_helper() -> None: + from sqlspec.adapters.aiosqlite.core import _filter_connection_params + + raw = {"database": ":memory:", "extra": {}, "pool_size": 10, "timeout": 30.0, "isolation_level": None} + filtered = _filter_connection_params(raw) + assert filtered == {"database": ":memory:", "timeout": 30.0, "isolation_level": None} diff --git a/tests/unit/adapters/test_aiosqlite/test_driver.py b/tests/unit/adapters/test_aiosqlite/test_driver.py index ffed6b300..7e9233068 100644 --- a/tests/unit/adapters/test_aiosqlite/test_driver.py +++ b/tests/unit/adapters/test_aiosqlite/test_driver.py @@ -448,3 +448,28 @@ async def test_aiosqlite_cache_hit_preserves_ambient_exception_and_formats_wire_ assert ambient_preserved finally: await connection.close() + + +async def test_aiosqlite_cursor_lazy_allocation() -> None: + conn = await aiosqlite.connect(":memory:") + try: + driver = AiosqliteDriver(connection=conn) + result = await driver.execute("SELECT 1 AS probe") + assert result.data == [(1,)] + finally: + await conn.close() + + +async def test_aiosqlite_data_dictionary_version_cached_across_drivers() -> None: + conn1 = await aiosqlite.connect(":memory:") + conn2 = await aiosqlite.connect(":memory:") + try: + driver1 = AiosqliteDriver(connection=conn1) + driver2 = AiosqliteDriver(connection=conn2) + v1 = await driver1.data_dictionary.get_version(driver1) + v2 = await driver2.data_dictionary.get_version(driver2) + assert v1 is not None + assert v1 is v2 + finally: + await conn1.close() + await conn2.close() From c601ffd7a5b7b1fbc207e695d73b4455a204ddbd Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sat, 31 Oct 2026 22:47:04 +0000 Subject: [PATCH 011/174] refactor(aiosqlite): clean up dead code, redundant wrappers, and error mappings Delegate script execution to base class with rowid cache invalidation, remove redundant pass-through load_from_storage, replace truthy or in pool creation with explicit is None checks, remove dead Python 3.11 version check, and add missing error mappings for auth, range, toobig, and notadb. --- sqlspec/adapters/aiosqlite/config.py | 21 +++++++----- sqlspec/adapters/aiosqlite/core.py | 14 ++++++-- sqlspec/adapters/aiosqlite/driver.py | 30 +---------------- sqlspec/adapters/aiosqlite/type_converter.py | 3 -- .../adapters/test_aiosqlite/test_config.py | 18 ++++++++++ .../adapters/test_aiosqlite/test_driver.py | 11 +++++++ .../test_aiosqlite/test_exception_mapping.py | 33 +++++++++++++++++++ 7 files changed, 87 insertions(+), 43 deletions(-) diff --git a/sqlspec/adapters/aiosqlite/config.py b/sqlspec/adapters/aiosqlite/config.py index b8c04305c..a727d086a 100644 --- a/sqlspec/adapters/aiosqlite/config.py +++ b/sqlspec/adapters/aiosqlite/config.py @@ -17,7 +17,6 @@ build_connection_config, default_statement_config, render_pragmas, - require_python_version, ) from sqlspec.adapters.aiosqlite.driver import AiosqliteDriver, AiosqliteExceptionHandler from sqlspec.adapters.aiosqlite.pool import ( @@ -101,7 +100,7 @@ class AiosqlitePoolParams(AiosqliteConnectionParams): enable_optimizations: NotRequired[bool] """Apply WAL journal mode, ``synchronous=NORMAL``, and busy-timeout PRAGMAs on connect.""" enable_foreign_keys: NotRequired[bool] - """Execute ``PRAGMA foreign_keys = ON`` on every pooled connection (defaults to ``True``).""" + """Execute ``PRAGMA foreign_keys = ON`` on every pooled connection (defaults to ``False``).""" extra: NotRequired["dict[str, Any]"] """Additional keyword arguments forwarded directly to ``aiosqlite.connect``.""" @@ -408,13 +407,21 @@ async def _create_pool(self) -> AiosqliteConnectionPool: Returns: AiosqliteConnectionPool: The connection pool instance. """ - pool_size = self.connection_config.get("pool_size") or 5 + pool_size = self.connection_config.get("pool_size") + if pool_size is None: + pool_size = 5 min_size = self.connection_config.get("min_size") if min_size is None: min_size = 0 - connect_timeout = self.connection_config.get("connect_timeout") or 30.0 - idle_timeout = self.connection_config.get("idle_timeout") or 24 * 60 * 60 - operation_timeout = self.connection_config.get("operation_timeout") or 10.0 + connect_timeout = self.connection_config.get("connect_timeout") + if connect_timeout is None: + connect_timeout = 30.0 + idle_timeout = self.connection_config.get("idle_timeout") + if idle_timeout is None: + idle_timeout = 24 * 60 * 60 + operation_timeout = self.connection_config.get("operation_timeout") + if operation_timeout is None: + operation_timeout = 10.0 health_check_interval = self.connection_config.get("health_check_interval") if health_check_interval is None: health_check_interval = 30.0 @@ -518,8 +525,6 @@ def _build_runtime_setup(features: "dict[str, Any]") -> "dict[str, Any] | None": ) window_functions = runtime_setup.get("custom_window_functions", ()) _validate_entries(window_functions, ("name", "narg", "window_class"), "custom_window_functions") - if window_functions: - require_python_version("Custom SQLite window functions", (3, 11)) interval = runtime_setup.get("progress_handler_interval") if interval is not None and (not isinstance(interval, int) or isinstance(interval, bool) or interval < 1): diff --git a/sqlspec/adapters/aiosqlite/core.py b/sqlspec/adapters/aiosqlite/core.py index 635265ef2..25bb989df 100644 --- a/sqlspec/adapters/aiosqlite/core.py +++ b/sqlspec/adapters/aiosqlite/core.py @@ -109,6 +109,10 @@ SQLITE_INTERRUPT_CODE = 9 SQLITE_PERM_CODE = 3 SQLITE_READONLY_CODE = 8 +SQLITE_TOOBIG_CODE = 18 +SQLITE_AUTH_CODE = 23 +SQLITE_RANGE_CODE = 25 +SQLITE_NOTADB_CODE = 26 SQLITE_CONNECT_SUPPORTS_AUTOCOMMIT = sys.version_info >= (3, 12) SQLITE_DATABASE_LIST_MIN_COLUMNS = 2 SQLITE_TABLE_LIST_MIN_COLUMNS = 5 @@ -463,7 +467,7 @@ def create_mapped_exception(error: BaseException, *, logger: Any | None = None) Returns: A SQLSpec exception that wraps the original error """ - del logger + _ = logger if has_sqlite_error(error): error_code = error.sqlite_errorcode error_name = error.sqlite_errorname @@ -509,8 +513,12 @@ def create_mapped_exception(error: BaseException, *, logger: Any | None = None) ): return _create_aiosqlite_error(error, 0, PermissionDeniedError, "permission denied") - if primary_code == SQLITE_CANTOPEN_CODE or ( - isinstance(error_name, str) and error_name.startswith("SQLITE_CANTOPEN") + if primary_code == SQLITE_AUTH_CODE or error_name == "SQLITE_AUTH": + return _create_aiosqlite_error(error, error_code, PermissionDeniedError, "authorization denied") + if primary_code in (SQLITE_RANGE_CODE, SQLITE_TOOBIG_CODE) or error_name in ("SQLITE_RANGE", "SQLITE_TOOBIG"): + return _create_aiosqlite_error(error, error_code, DataError, "data range or size exceeded") + if primary_code in (SQLITE_CANTOPEN_CODE, SQLITE_NOTADB_CODE) or ( + isinstance(error_name, str) and error_name.startswith(("SQLITE_CANTOPEN", "SQLITE_NOTADB")) ): return _create_aiosqlite_error(error, error_code, DatabaseConnectionError, "connection error") if primary_code in (SQLITE_IOERR_CODE, SQLITE_CORRUPT_CODE) or ( diff --git a/sqlspec/adapters/aiosqlite/driver.py b/sqlspec/adapters/aiosqlite/driver.py index f9e6db608..97240fc14 100644 --- a/sqlspec/adapters/aiosqlite/driver.py +++ b/sqlspec/adapters/aiosqlite/driver.py @@ -161,23 +161,11 @@ async def dispatch_execute_many(self, cursor: "AiosqliteRawCursor", statement: " async def dispatch_execute_script(self, cursor: "AiosqliteRawCursor", statement: "SQL") -> "ExecutionResult": """Execute SQL script.""" self._rowid_target_cache.clear() - sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) - statements = self.split_script_statements(sql, statement.statement_config, strip_trailing_semicolon=True) - - successful_count = 0 - last_cursor = cursor - try: - for stmt in statements: - await cursor.execute(stmt, normalize_execute_parameters(prepared_parameters)) - successful_count += 1 + return await super().dispatch_execute_script(cursor, statement) finally: self._rowid_target_cache.clear() - return self.create_execution_result( - last_cursor, statement_count=len(statements), successful_statements=successful_count, is_script_result=True - ) - async def begin(self, mode: "Literal['DEFERRED', 'IMMEDIATE', 'EXCLUSIVE'] | None" = None) -> None: """Begin a database transaction. @@ -335,22 +323,6 @@ async def load_from_arrow( self._attach_partition_telemetry(telemetry_payload, partitioner) return self._storage_job(telemetry_payload, telemetry) - async def load_from_storage( - self, - table: str, - source: "StorageDestination", - *, - file_format: "StorageFormat", - partitioner: "dict[str, object] | None" = None, - overwrite: bool = False, - **kwargs: Any, - ) -> "StorageBridgeJob": - """Load staged artifacts from storage into SQLite.""" - - return await super().load_from_storage( - table, source, file_format=file_format, partitioner=partitioner, overwrite=overwrite, **kwargs - ) - def collect_rows(self, cursor: Any, fetched: "list[Any]") -> "tuple[list[Any], list[str], int]": """Collect aiosqlite rows for the direct execution path.""" return collect_rows(fetched, cursor.description) diff --git a/sqlspec/adapters/aiosqlite/type_converter.py b/sqlspec/adapters/aiosqlite/type_converter.py index 23fb2d079..e83d86713 100644 --- a/sqlspec/adapters/aiosqlite/type_converter.py +++ b/sqlspec/adapters/aiosqlite/type_converter.py @@ -13,15 +13,12 @@ from typing import TYPE_CHECKING, Any from sqlspec.adapters.aiosqlite._typing import aiosqlite_sqlite_module as sqlite3 -from sqlspec.utils.logging import get_logger if TYPE_CHECKING: from collections.abc import Callable __all__ = ("json_adapter", "json_converter", "register_type_handlers", "unregister_type_handlers") -logger = get_logger(__name__) - DEFAULT_JSON_TYPE = "JSON" diff --git a/tests/unit/adapters/test_aiosqlite/test_config.py b/tests/unit/adapters/test_aiosqlite/test_config.py index c1b8a6724..09f2be0b2 100644 --- a/tests/unit/adapters/test_aiosqlite/test_config.py +++ b/tests/unit/adapters/test_aiosqlite/test_config.py @@ -321,3 +321,21 @@ def test_filter_connection_params_with_hoisted_helper() -> None: raw = {"database": ":memory:", "extra": {}, "pool_size": 10, "timeout": 30.0, "isolation_level": None} filtered = _filter_connection_params(raw) assert filtered == {"database": ":memory:", "timeout": 30.0, "isolation_level": None} + + +async def test_create_pool_preserves_zero_timeouts(tmp_path: Path) -> None: + config = AiosqliteConfig( + connection_config={ + "database": tmp_path / "zero.db", + "connect_timeout": 0.0, + "idle_timeout": 0.0, + "operation_timeout": 0.0, + } + ) + pool = await config._create_pool() + try: + assert pool._connect_timeout == 0.0 + assert pool._idle_timeout == 0.0 + assert pool._operation_timeout == 0.0 + finally: + await pool.close() diff --git a/tests/unit/adapters/test_aiosqlite/test_driver.py b/tests/unit/adapters/test_aiosqlite/test_driver.py index 7e9233068..c75014cfb 100644 --- a/tests/unit/adapters/test_aiosqlite/test_driver.py +++ b/tests/unit/adapters/test_aiosqlite/test_driver.py @@ -473,3 +473,14 @@ async def test_aiosqlite_data_dictionary_version_cached_across_drivers() -> None finally: await conn1.close() await conn2.close() + + +async def test_aiosqlite_dispatch_execute_script_emits_events() -> None: + conn = await aiosqlite.connect(":memory:") + try: + driver = AiosqliteDriver(connection=conn) + result = await driver.execute_script("CREATE TABLE scr_t (x INT); INSERT INTO scr_t VALUES (1);") + assert result.total_statements == 2 + assert result.successful_statements == 2 + finally: + await conn.close() diff --git a/tests/unit/adapters/test_aiosqlite/test_exception_mapping.py b/tests/unit/adapters/test_aiosqlite/test_exception_mapping.py index 969ea9768..20c08fb60 100644 --- a/tests/unit/adapters/test_aiosqlite/test_exception_mapping.py +++ b/tests/unit/adapters/test_aiosqlite/test_exception_mapping.py @@ -7,6 +7,7 @@ from sqlspec.adapters.aiosqlite.core import create_mapped_exception from sqlspec.exceptions import ( CheckViolationError, + DatabaseConnectionError, DataError, DeadlockError, IntegrityError, @@ -182,3 +183,35 @@ def test_interface_error_maps_to_operational_error() -> None: result = create_mapped_exception(err) assert isinstance(result, OperationalError) assert result.__cause__ is err + + +def test_auth_error_maps_to_permission_denied() -> None: + err = _SqliteError("not authorized", 23, "SQLITE_AUTH") + result = create_mapped_exception(err) + assert isinstance(result, PermissionDeniedError) + assert result.__cause__ is err + + +def test_range_and_toobig_errors_map_to_data_error() -> None: + err_range = _SqliteError("column index out of range", 25, "SQLITE_RANGE") + res_range = create_mapped_exception(err_range) + assert isinstance(res_range, DataError) + assert res_range.__cause__ is err_range + + err_toobig = _SqliteError("string or blob too big", 18, "SQLITE_TOOBIG") + res_toobig = create_mapped_exception(err_toobig) + assert isinstance(res_toobig, DataError) + assert res_toobig.__cause__ is err_toobig + + +def test_notadb_error_maps_to_database_connection_error() -> None: + err = _SqliteError("file is not a database", 26, "SQLITE_NOTADB") + result = create_mapped_exception(err) + assert isinstance(result, DatabaseConnectionError) + assert result.__cause__ is err + + +def test_create_mapped_exception_accepts_logger_parameter() -> None: + err = sqlite3.OperationalError("generic") + result = create_mapped_exception(err, logger=None) + assert isinstance(result, OperationalError) From b007c51ad10f7375733b2076fa05680ef21dfbb4 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sat, 31 Oct 2026 22:47:04 +0000 Subject: [PATCH 012/174] fix(sqlite): fix transaction rollback on exit, reentrancy depth, shared memory cache, and error mappings Roll back uncommitted transactions on exit in get_connection unless explicitly committed, track thread-local reentrancy depth, default memory database to cache=shared, delegate script execution to base driver with rowid cache invalidation, and map substring errors only when error_code is None. --- sqlspec/adapters/sqlite/config.py | 2 +- sqlspec/adapters/sqlite/core.py | 33 +++++++++++--- sqlspec/adapters/sqlite/driver.py | 26 +++-------- sqlspec/adapters/sqlite/pool.py | 20 ++++++--- .../unit/adapters/test_sqlite/test_config.py | 26 +++++++++++ .../unit/adapters/test_sqlite/test_driver.py | 11 +++++ .../test_sqlite/test_exception_mapping.py | 31 +++++++++++++ tests/unit/adapters/test_sqlite/test_pool.py | 44 +++++++++++++++++-- 8 files changed, 154 insertions(+), 39 deletions(-) diff --git a/sqlspec/adapters/sqlite/config.py b/sqlspec/adapters/sqlite/config.py index b17cb99fe..0260d8c81 100644 --- a/sqlspec/adapters/sqlite/config.py +++ b/sqlspec/adapters/sqlite/config.py @@ -251,7 +251,7 @@ def __init__( if database is not None: config_dict["database"] = database if "database" not in config_dict or config_dict["database"] == ":memory:": - config_dict["database"] = f"file:memory_{uuid4().hex}?mode=memory&cache=private" + config_dict["database"] = f"file:memory_{uuid4().hex}?mode=memory&cache=shared" config_dict["uri"] = True elif "database" in config_dict: database_path = str(config_dict["database"]) diff --git a/sqlspec/adapters/sqlite/core.py b/sqlspec/adapters/sqlite/core.py index e78c4ce42..0e7fad619 100644 --- a/sqlspec/adapters/sqlite/core.py +++ b/sqlspec/adapters/sqlite/core.py @@ -83,6 +83,7 @@ SQLITE_IOERR_CODE = 10 SQLITE_CORRUPT_CODE = 11 SQLITE_MISMATCH_CODE = 20 +SQLITE_MISUSE_CODE = 21 SQLITE_BUSY_CODE = 5 SQLITE_LOCKED_CODE = 6 SQLITE_INTERRUPT_CODE = 9 @@ -475,7 +476,7 @@ def create_mapped_exception(error: BaseException, *, logger: Any | None = None) Returns: A SQLSpec exception that wraps the original error """ - del logger + _ = logger if has_sqlite_error(error): error_code = error.sqlite_errorcode error_name = error.sqlite_errorname @@ -493,13 +494,19 @@ def create_mapped_exception(error: BaseException, *, logger: Any | None = None) return _create_sqlite_error(error, error_code, DeadlockError, "database busy") if primary_code == SQLITE_LOCKED_CODE or (isinstance(error_name, str) and error_name.startswith("SQLITE_LOCKED")): return _create_sqlite_error(error, error_code, DeadlockError, "database locked") - if "locked" in error_msg or "busy" in error_msg: - return _create_sqlite_error(error, error_code or 0, DeadlockError, "database locked") + if error_code is None and ( + "database is locked" in error_msg + or "database table is locked" in error_msg + or "database locked" in error_msg + or "database is busy" in error_msg + or "database busy" in error_msg + ): + return _create_sqlite_error(error, 0, DeadlockError, "database locked") if primary_code == SQLITE_INTERRUPT_CODE or error_name == "SQLITE_INTERRUPT": return _create_sqlite_error(error, error_code, OperationCancelledError, "query interrupted") - if "interrupt" in error_msg: - return _create_sqlite_error(error, error_code or 0, OperationCancelledError, "query interrupted") + if error_code is None and ("query was interrupted" in error_msg or "interrupted" in error_msg): + return _create_sqlite_error(error, 0, OperationCancelledError, "query interrupted") if primary_code == SQLITE_PERM_CODE or error_name == "SQLITE_PERM": return _create_sqlite_error(error, error_code, PermissionDeniedError, "permission denied") @@ -507,8 +514,13 @@ def create_mapped_exception(error: BaseException, *, logger: Any | None = None) isinstance(error_name, str) and error_name.startswith("SQLITE_READONLY") ): return _create_sqlite_error(error, error_code, PermissionDeniedError, "database is read-only") - if "permission denied" in error_msg or "readonly" in error_msg: - return _create_sqlite_error(error, error_code or 0, PermissionDeniedError, "permission denied") + if error_code is None and ( + "attempt to write a readonly database" in error_msg + or "readonly database" in error_msg + or "permission denied" in error_msg + or "access permission denied" in error_msg + ): + return _create_sqlite_error(error, 0, PermissionDeniedError, "permission denied") if primary_code == SQLITE_CANTOPEN_CODE or ( isinstance(error_name, str) and error_name.startswith("SQLITE_CANTOPEN") @@ -522,9 +534,16 @@ def create_mapped_exception(error: BaseException, *, logger: Any | None = None) if primary_code == SQLITE_MISMATCH_CODE or error_name == "SQLITE_MISMATCH": return _create_sqlite_error(error, error_code, DataError, "data error") + if primary_code == SQLITE_MISUSE_CODE or error_name == "SQLITE_MISUSE": + return _create_sqlite_error(error, error_code, DataError, "programming error") + if "syntax" in error_msg: return _create_sqlite_error(error, error_code, SQLParsingError, "SQL syntax error") + if isinstance(error, sqlite3.ProgrammingError): + return _create_sqlite_error(error, error_code, DataError, "programming error") + if isinstance(error, sqlite3.InterfaceError): + return _create_sqlite_error(error, error_code, OperationalError, "interface error") if primary_code == 1 or isinstance(error, sqlite3.OperationalError): return _create_sqlite_error(error, error_code, OperationalError, "operational error") if isinstance(error, sqlite3.DataError): diff --git a/sqlspec/adapters/sqlite/driver.py b/sqlspec/adapters/sqlite/driver.py index 1729a5138..af3397eba 100644 --- a/sqlspec/adapters/sqlite/driver.py +++ b/sqlspec/adapters/sqlite/driver.py @@ -175,28 +175,12 @@ def dispatch_execute_many(self, cursor: Any, statement: "SQL") -> "ExecutionResu return self.create_execution_result(cursor, rowcount_override=affected_rows, is_many_result=True) def dispatch_execute_script(self, cursor: Any, statement: "SQL") -> "ExecutionResult": - """Execute SQL script with statement splitting and parameter handling. - - Args: - cursor: SQLite cursor object - statement: SQL statement containing multiple statements - - Returns: - ExecutionResult with script execution details - """ + """Execute SQL script with statement splitting and parameter handling.""" self._rowid_target_cache.clear() - sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) - statements = self.split_script_statements(sql, statement.statement_config, strip_trailing_semicolon=True) - - successful_count = 0 - - for stmt in statements: - cursor.execute(stmt, normalize_execute_parameters(prepared_parameters)) - successful_count += 1 - - return self.create_execution_result( - cursor, statement_count=len(statements), successful_statements=successful_count, is_script_result=True - ) + try: + return super().dispatch_execute_script(cursor, statement) + finally: + self._rowid_target_cache.clear() def begin(self, mode: "Literal['DEFERRED', 'IMMEDIATE', 'EXCLUSIVE'] | None" = None) -> None: """Begin a database transaction. diff --git a/sqlspec/adapters/sqlite/pool.py b/sqlspec/adapters/sqlite/pool.py index a2de6cbf0..3de1a227a 100644 --- a/sqlspec/adapters/sqlite/pool.py +++ b/sqlspec/adapters/sqlite/pool.py @@ -153,19 +153,25 @@ def get_connection(self) -> "Generator[SqliteConnection, None, None]": SqliteConnection: A thread-local connection. """ connection = self._get_thread_connection() + depth = getattr(self._thread_local, "depth", 0) + self._thread_local.depth = depth + 1 try: yield connection except BaseException: - with contextlib.suppress(Exception): - end_transaction(connection, commit=False) + if self._thread_local.depth == 1: + with contextlib.suppress(Exception): + end_transaction(connection, commit=False) raise else: - try: - end_transaction(connection, commit=True) - except BaseException: - with contextlib.suppress(Exception): + if self._thread_local.depth == 1: + try: end_transaction(connection, commit=False) - raise + except BaseException: + with contextlib.suppress(Exception): + end_transaction(connection, commit=False) + raise + finally: + self._thread_local.depth -= 1 def close(self) -> None: """Close every connection this pool opened, on any thread.""" diff --git a/tests/unit/adapters/test_sqlite/test_config.py b/tests/unit/adapters/test_sqlite/test_config.py index 7d607481c..d2087278a 100644 --- a/tests/unit/adapters/test_sqlite/test_config.py +++ b/tests/unit/adapters/test_sqlite/test_config.py @@ -1,5 +1,6 @@ """SQLite configuration normalization tests.""" +import threading from os import PathLike from pathlib import Path from typing import Literal, get_args, get_origin, get_type_hints @@ -275,3 +276,28 @@ def test_create_connection_returns_a_caller_owned_connection(tmp_path: Path) -> pooled.execute("SELECT 1") finally: config.close_pool() + + +def test_default_memory_database_uses_shared_cache() -> None: + config = SqliteConfig() + assert "cache=shared" in config.connection_config["database"] + pool = config.provide_pool() + try: + with pool.get_connection() as conn: + conn.execute("CREATE TABLE shared_probe (id INT)") + conn.commit() + + results: list[int] = [] + + def worker() -> None: + with pool.get_connection() as conn: + row = conn.execute("SELECT count(*) FROM shared_probe").fetchone() + if row: + results.append(row[0]) + + t = threading.Thread(target=worker) + t.start() + t.join() + assert results == [0] + finally: + pool.close() diff --git a/tests/unit/adapters/test_sqlite/test_driver.py b/tests/unit/adapters/test_sqlite/test_driver.py index 7d0b46fc9..dae39f8bb 100644 --- a/tests/unit/adapters/test_sqlite/test_driver.py +++ b/tests/unit/adapters/test_sqlite/test_driver.py @@ -342,3 +342,14 @@ def test_sqlite_json_converter_and_unregister() -> None: register_type_handlers() unregister_type_handlers() assert (dict, sqlite3.PrepareProtocol) not in sqlite3.adapters + + +def test_sqlite_dispatch_execute_script_emits_events() -> None: + connection = sqlite3.connect(":memory:") + try: + driver = SqliteDriver(connection=connection) + result = driver.execute_script("CREATE TABLE scr_sync (x INT); INSERT INTO scr_sync VALUES (1);") + assert result.total_statements == 2 + assert result.successful_statements == 2 + finally: + connection.close() diff --git a/tests/unit/adapters/test_sqlite/test_exception_mapping.py b/tests/unit/adapters/test_sqlite/test_exception_mapping.py index 4821c5329..b2eb60a9e 100644 --- a/tests/unit/adapters/test_sqlite/test_exception_mapping.py +++ b/tests/unit/adapters/test_sqlite/test_exception_mapping.py @@ -7,9 +7,11 @@ from sqlspec.adapters.sqlite.core import create_mapped_exception from sqlspec.exceptions import ( CheckViolationError, + DataError, DeadlockError, IntegrityError, NotNullViolationError, + OperationalError, OperationCancelledError, PermissionDeniedError, UniqueViolationError, @@ -118,3 +120,32 @@ def test_trigger_abort_does_not_map_to_not_null_violation() -> None: assert not isinstance(result, NotNullViolationError) assert isinstance(result, IntegrityError) assert result.__cause__ is err + + +def test_no_such_column_with_locked_name_maps_to_operational_error() -> None: + err = sqlite3.OperationalError("no such column: is_locked") + result = create_mapped_exception(err) + assert isinstance(result, OperationalError) + assert not isinstance(result, DeadlockError) + assert result.__cause__ is err + + +def test_programming_error_maps_to_data_error() -> None: + err = sqlite3.ProgrammingError("cannot operate on a closed database") + result = create_mapped_exception(err) + assert isinstance(result, DataError) + assert result.__cause__ is err + + +def test_interface_error_maps_to_operational_error() -> None: + err = sqlite3.InterfaceError("error binding parameter") + result = create_mapped_exception(err) + assert isinstance(result, OperationalError) + assert result.__cause__ is err + + +def test_misuse_maps_to_data_error() -> None: + err = _SqliteError("bad call", 21, "SQLITE_MISUSE") + result = create_mapped_exception(err) + assert isinstance(result, DataError) + assert result.__cause__ is err diff --git a/tests/unit/adapters/test_sqlite/test_pool.py b/tests/unit/adapters/test_sqlite/test_pool.py index 8b55d89f8..d8a49b718 100644 --- a/tests/unit/adapters/test_sqlite/test_pool.py +++ b/tests/unit/adapters/test_sqlite/test_pool.py @@ -155,8 +155,8 @@ class _BoomError(RuntimeError): pool.close() -def test_get_connection_commits_open_transaction_on_clean_exit() -> None: - """A clean with-block exit commits the open transaction.""" +def test_get_connection_rolls_back_open_transaction_on_clean_exit() -> None: + """A clean with-block exit rolls back the uncommitted transaction.""" pool = SqliteConnectionPool(connection_parameters={"database": ":memory:"}, enable_optimizations=False) with pool.get_connection() as connection: connection.execute("CREATE TABLE widgets (id INTEGER PRIMARY KEY)") @@ -167,7 +167,7 @@ def test_get_connection_commits_open_transaction_on_clean_exit() -> None: with pool.get_connection() as connection: count = connection.execute("SELECT COUNT(*) FROM widgets").fetchone()[0] - assert count == 1 + assert count == 0 pool.close() @@ -389,3 +389,41 @@ def test_memory_pool_replaces_a_closed_connection() -> None: assert replacement.execute("SELECT 1").fetchone() == (1,) finally: pool.close() + + +def test_uncommitted_writes_rollback_on_exit() -> None: + pool = SqliteConnectionPool({"database": ":memory:"}) + try: + with pool.get_connection() as conn: + conn.execute("CREATE TABLE test_rb (x INT)") + conn.execute("INSERT INTO test_rb VALUES (1)") + conn.commit() + + with pool.get_connection() as conn: + conn.execute("INSERT INTO test_rb VALUES (2)") + + with pool.get_connection() as conn: + rows = conn.execute("SELECT * FROM test_rb").fetchall() + assert rows == [(1,)] + finally: + pool.close() + + +def test_nested_get_connection_reentrancy_depth() -> None: + pool = SqliteConnectionPool({"database": ":memory:"}) + try: + with pool.get_connection() as conn1: + conn1.execute("CREATE TABLE test_depth (x INT)") + conn1.commit() + conn1.execute("INSERT INTO test_depth VALUES (1)") + with pool.get_connection() as conn2: + assert conn1 is conn2 + conn2.execute("INSERT INTO test_depth VALUES (2)") + assert getattr(conn1, "in_transaction", False) + conn1.commit() + + with pool.get_connection() as conn: + rows = conn.execute("SELECT * FROM test_depth").fetchall() + assert rows == [(1,), (2,)] + finally: + pool.close() From c976923b87406b391d1e349e391988827f8839d7 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sat, 31 Oct 2026 22:47:05 +0000 Subject: [PATCH 013/174] fix(sqlite): serialize float32 embeddings in ADK store and parameterize events limit Implement binary float32 packing and vector similarity ranking for SQLite ADK memory store, parameterize LIMIT in get_events, and retain connection ID tracking for non-weakref-capable sqlite3.Connection instances. --- sqlspec/adapters/sqlite/adk/store.py | 106 +++++++++++++++--- .../adapters/test_sqlite/test_adk_config.py | 74 ++++++++++++ .../extensions/test_schema_reconciliation.py | 2 +- 3 files changed, 164 insertions(+), 18 deletions(-) diff --git a/sqlspec/adapters/sqlite/adk/store.py b/sqlspec/adapters/sqlite/adk/store.py index a8e563ea6..6a919cfbe 100644 --- a/sqlspec/adapters/sqlite/adk/store.py +++ b/sqlspec/adapters/sqlite/adk/store.py @@ -1,6 +1,8 @@ """SQLite sync ADK store for Google Agent Development Kit session/event storage.""" +import array import contextlib +import math import re from collections.abc import Mapping from datetime import UTC, datetime, timedelta @@ -478,7 +480,11 @@ def get_events( params.append(_datetime_to_julian(after_timestamp)) where_clause = " AND ".join(where_clauses) - limit_clause = f" LIMIT {limit}" if limit else "" + if limit is not None: + limit_clause = " LIMIT ?" + params.append(int(limit)) + else: + limit_clause = "" sql = f""" SELECT id, app_name, user_id, session_id, invocation_id, timestamp, event_data @@ -887,8 +893,8 @@ def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "object INSERT OR IGNORE INTO {self._memory_table} (id, session_id, app_name, user_id, scope, event_id, author, {self._owner_id_column_name}, timestamp, content_json, - content_text, metadata_json, inserted_at) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + content_text, metadata_json, inserted_at, embedding) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """ for entry in entries: timestamp_julian = _datetime_to_julian(entry["timestamp"]) @@ -911,13 +917,14 @@ def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "object entry["content_text"], metadata_json_str, inserted_at_julian, + _pack_embedding(entry.get("embedding")), )) else: sql = f""" INSERT OR IGNORE INTO {self._memory_table} (id, session_id, app_name, user_id, scope, event_id, author, - timestamp, content_json, content_text, metadata_json, inserted_at) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + timestamp, content_json, content_text, metadata_json, inserted_at, embedding) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """ for entry in entries: timestamp_julian = _datetime_to_julian(entry["timestamp"]) @@ -939,6 +946,7 @@ def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "object entry["content_text"], metadata_json_str, inserted_at_julian, + _pack_embedding(entry.get("embedding")), )) cursor = conn.executemany(sql, params_list) @@ -959,19 +967,22 @@ def search_entries( scope_filter: Literal["all", "user", "app"] = "all", embedding: "Sequence[float] | None" = None, ) -> "list[StoredMemory]": - """Search memory entries by text query.""" + """Search memory entries by text query or vector embedding.""" if not self._enabled: msg = "Memory store is disabled" raise RuntimeError(msg) - if not query or not query.strip(): + has_query = bool(query and query.strip()) + if not has_query and embedding is None: return [] effective_limit = limit if limit is not None else self._max_results + fetch_limit = effective_limit if embedding is None else max(effective_limit * 5, 50) - if self._use_fts: + records: list[StoredMemory] | None = None + if has_query and self._use_fts: try: - return self._search_entries_fts(query, app_name, user_id, effective_limit, scope_filter) + records = self._search_entries_fts(query, app_name, user_id, fetch_limit, scope_filter) except sqlite3.OperationalError as exc: error_text = str(exc) if SQLITE_TABLE_NOT_FOUND_ERROR in error_text and f"{self._memory_table}_fts" not in error_text: @@ -979,12 +990,29 @@ def search_entries( logger.warning("FTS search failed; falling back to simple search: %s", exc) except Exception as exc: logger.warning("FTS search failed; falling back to simple search: %s", exc) - try: - return self._search_entries_simple(query, app_name, user_id, effective_limit, scope_filter) - except sqlite3.OperationalError as exc: - if SQLITE_TABLE_NOT_FOUND_ERROR in str(exc): - return [] - raise + + if records is None: + try: + if has_query: + records = self._search_entries_simple(query, app_name, user_id, fetch_limit, scope_filter) + else: + records = self._search_entries_vector(app_name, user_id, fetch_limit, scope_filter) + except sqlite3.OperationalError as exc: + if SQLITE_TABLE_NOT_FOUND_ERROR in str(exc): + return [] + raise + + if embedding is not None and records: + query_vec = [float(x) for x in embedding] + + def _score(r: "StoredMemory") -> float: + emb = r.get("embedding") + return _cosine_similarity(query_vec, emb) if emb is not None else -1.0 + + records.sort(key=_score, reverse=True) + records = records[:effective_limit] + + return records def delete_entries_by_session(self, session_id: str) -> int: """Delete all memory entries for a specific session.""" @@ -1092,7 +1120,8 @@ def _memory_table_ddl(self) -> str: content_json TEXT NOT NULL, content_text TEXT NOT NULL, metadata_json TEXT, - inserted_at REAL NOT NULL + inserted_at REAL NOT NULL, + embedding BLOB ); CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_app_scope_user_time @@ -1169,6 +1198,21 @@ def _search_entries_simple( params = (*scope_params, pattern, limit) return self._fetch_records(sql, params) + def _search_entries_vector( + self, app_name: str, user_id: str, limit: int, scope_filter: Literal["all", "user", "app"] = "all" + ) -> "list[StoredMemory]": + where_scope, scope_params = _build_sqlite_scope_clause("", app_name, user_id, scope_filter) + sql = f""" + SELECT * + FROM {self._memory_table} + WHERE {where_scope} + AND embedding IS NOT NULL + ORDER BY timestamp DESC + LIMIT ? + """ + params = (*scope_params, limit) + return self._fetch_records(sql, params) + def _fetch_records(self, sql: str, params: "tuple[Any, ...]") -> "list[StoredMemory]": with self._config.provide_connection() as conn: self._ensure_foreign_keys(conn) @@ -1205,7 +1249,7 @@ def _fetch_records(self, sql: str, params: "tuple[Any, ...]") -> "list[StoredMem raw["inserted_at"] = _julian_to_datetime(raw["inserted_at"]) raw["content_json"] = from_json(raw["content_json"]) if raw["content_json"] else {} raw["metadata_json"] = from_json(raw["metadata_json"]) if raw["metadata_json"] else None - raw["embedding"] = None + raw["embedding"] = _unpack_embedding(raw.get("embedding")) records.append(cast("StoredMemory", raw)) return records @@ -1303,6 +1347,34 @@ def _build_sqlite_scope_clause( return f"{prefix}app_name = ? AND {prefix}scope = 'app'", (app_name,) +def _pack_embedding(raw: Any) -> bytes | None: + """Pack float sequence into float32 binary BLOB.""" + if raw is None: + return None + return array.array("f", [float(x) for x in raw]).tobytes() + + +def _unpack_embedding(blob: Any) -> list[float] | None: + """Unpack float32 binary BLOB into float list.""" + if blob is None: + return None + return list(array.array("f", blob)) + + +def _cosine_similarity(v1: "Sequence[float]", v2: "Sequence[float]") -> float: + """Compute cosine similarity between two float vectors.""" + dot = 0.0 + norm1 = 0.0 + norm2 = 0.0 + for a, b in zip(v1, v2, strict=False): + dot += a * b + norm1 += a * a + norm2 += b * b + if norm1 == 0.0 or norm2 == 0.0: + return 0.0 + return float(dot / (math.sqrt(norm1) * math.sqrt(norm2))) + + def _session_list_query( session_table: str, app_name: str, diff --git a/tests/unit/adapters/test_sqlite/test_adk_config.py b/tests/unit/adapters/test_sqlite/test_adk_config.py index 22c846449..06b47814f 100644 --- a/tests/unit/adapters/test_sqlite/test_adk_config.py +++ b/tests/unit/adapters/test_sqlite/test_adk_config.py @@ -9,6 +9,7 @@ from sqlspec.adapters.sqlite.config import SqliteConfig from sqlspec.config import ADKConfig from sqlspec.exceptions import ImproperConfigurationError +from sqlspec.extensions.adk import StoredMemory class RecordingConnection: @@ -189,3 +190,76 @@ def test_sqlite_list_sessions_rejects_invalid_options_before_connecting(options: store.list_sessions("app", **options) assert conn.calls == [] + + +def test_sqlite_adk_get_events_with_limit_binds_parameter() -> None: + """Verify get_events parameterizes LIMIT instead of string interpolation.""" + config = SqliteConfig(connection_config={"database": ":memory:"}) + store = SqliteADKStore(config) + calls: list[tuple[str, tuple[Any, ...]]] = [] + + class _MockEventsCursor: + def fetchall(self) -> list[Any]: + return [] + + class _MockEventsConnection: + def execute(self, sql: str, params: Any = None) -> _MockEventsCursor: + calls.append((sql, params)) + return _MockEventsCursor() + + class _SyncContextManager: + def __enter__(self) -> Any: + return _MockEventsConnection() + + def __exit__(self, *_: Any) -> None: + return None + + config.provide_connection = MagicMock(return_value=_SyncContextManager()) # type: ignore[method-assign] + + events = store.get_events("app", "user", "session", limit=10) + assert events == [] + assert len(calls) >= 1 + sql, params = calls[-1] + assert "LIMIT ?" in sql + assert 10 in params + + +def test_sqlite_adk_memory_store_persists_and_retrieves_float32_embeddings() -> None: + """Verify memory entries persist and retrieve float32 embeddings.""" + from datetime import UTC, datetime + from uuid import uuid4 + + config = SqliteConfig(connection_config={"database": ":memory:"}) + store = SqliteADKMemoryStore(config) + try: + store.create_tables() + + now = datetime.now(UTC) + embedding = [0.1, 0.2, 0.3, 0.4] + record = StoredMemory( + id=str(uuid4()), + session_id="s1", + app_name="app", + user_id="u1", + scope="user", + event_id="evt-1", + author="user", + timestamp=now, + content_json={"text": "vector search test"}, + content_text="vector search test", + metadata_json=None, + inserted_at=now, + embedding=embedding, + ) + inserted = store.insert_memory_entries([record]) + assert inserted == 1 + + results = store.search_entries(query="vector", app_name="app", user_id="u1", embedding=embedding) + assert len(results) == 1 + assert results[0]["event_id"] == "evt-1" + assert results[0]["embedding"] is not None + assert len(results[0]["embedding"]) == 4 + for orig, rec in zip(embedding, results[0]["embedding"], strict=True): + assert abs(orig - rec) < 1e-6 + finally: + config.close_pool() diff --git a/tests/unit/extensions/test_schema_reconciliation.py b/tests/unit/extensions/test_schema_reconciliation.py index b7c73d819..33c196c82 100644 --- a/tests/unit/extensions/test_schema_reconciliation.py +++ b/tests/unit/extensions/test_schema_reconciliation.py @@ -34,7 +34,7 @@ def _memory_table_ddl(self) -> str: return ( super() ._memory_table_ddl() - .replace("inserted_at REAL NOT NULL\n", "inserted_at REAL NOT NULL,\n schema_tag TEXT\n") + .replace("embedding BLOB\n", "embedding BLOB,\n schema_tag TEXT\n") ) From e156fd5856fc18f572e2c6a9606c817614bc5d8c Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sat, 31 Oct 2026 22:47:05 +0000 Subject: [PATCH 014/174] fix(sqlite): memoize generated columns, coalesce schema to main, and raise SQLFileNotFoundError Memoize generated column expression parsing per table in SQLite and AioSQLite data dictionaries. Coalesce schema_name to 'main' in SQLite index and column introspection queries, and raise SQLFileNotFoundError when data dictionary queries are missing. --- sqlspec/adapters/aiosqlite/data_dictionary.py | 3 +- sqlspec/adapters/sqlite/data_dictionary.py | 10 ++- .../dialects/sqlite/sql/columns.sql | 4 +- .../dialects/sqlite/sql/indexes.sql | 10 +-- .../test_sqlite_duckdb_metadata.py | 71 +++++++++++++++++++ 5 files changed, 87 insertions(+), 11 deletions(-) diff --git a/sqlspec/adapters/aiosqlite/data_dictionary.py b/sqlspec/adapters/aiosqlite/data_dictionary.py index ecdd916f5..d2f655ea5 100644 --- a/sqlspec/adapters/aiosqlite/data_dictionary.py +++ b/sqlspec/adapters/aiosqlite/data_dictionary.py @@ -27,6 +27,7 @@ sqlite_schema_ddl_results, ) from sqlspec.driver import AsyncDataDictionaryBase +from sqlspec.exceptions import SQLFileNotFoundError if TYPE_CHECKING: from sqlspec.adapters.aiosqlite.driver import AiosqliteDriver @@ -289,7 +290,7 @@ def _get_domain_query_text(self, domain: str, query_name: str) -> str: query_text = get_data_dictionary_loader().get_domain_query_text(type(self).dialect, domain, query_name) if query_text is None: msg = f"Missing SQLite data-dictionary query: {domain}/{query_name}" - raise RuntimeError(msg) + raise SQLFileNotFoundError(msg) return query_text diff --git a/sqlspec/adapters/sqlite/data_dictionary.py b/sqlspec/adapters/sqlite/data_dictionary.py index 89a43a65a..6267502d8 100644 --- a/sqlspec/adapters/sqlite/data_dictionary.py +++ b/sqlspec/adapters/sqlite/data_dictionary.py @@ -26,6 +26,7 @@ sqlite_schema_ddl_results, ) from sqlspec.driver import SyncDataDictionaryBase +from sqlspec.exceptions import SQLFileNotFoundError if TYPE_CHECKING: from sqlspec.adapters.sqlite.driver import SqliteDriver @@ -138,10 +139,13 @@ def get_columns( rows = driver.select(query_text, table_name=table, schema_name=schema_name) if any(str(row.get("extra") or "").endswith("generated") for row in rows): table_sql = self._get_table_sql(driver, schema_name, table) + parsed_expressions: dict[str, dict[str, str]] = {} for row in rows: if str(row.get("extra") or "").endswith("generated"): - expressions = sqlite_generated_expressions(table_sql.get(str(row.get("table_name")))) - expression = expressions.get(str(row.get("column_name"))) + table_name_str = str(row.get("table_name")) + if table_name_str not in parsed_expressions: + parsed_expressions[table_name_str] = sqlite_generated_expressions(table_sql.get(table_name_str)) + expression = parsed_expressions[table_name_str].get(str(row.get("column_name"))) if expression is not None: row["generation_expression"] = expression return cast("list[ColumnMetadata]", rows) @@ -270,7 +274,7 @@ def _get_domain_query_text(self, domain: str, query_name: str) -> str: query_text = get_data_dictionary_loader().get_domain_query_text(type(self).dialect, domain, query_name) if query_text is None: msg = f"Missing SQLite data-dictionary query: {domain}/{query_name}" - raise RuntimeError(msg) + raise SQLFileNotFoundError(msg) return query_text diff --git a/sqlspec/data_dictionary/dialects/sqlite/sql/columns.sql b/sqlspec/data_dictionary/dialects/sqlite/sql/columns.sql index dcae8f703..510bac0d8 100644 --- a/sqlspec/data_dictionary/dialects/sqlite/sql/columns.sql +++ b/sqlspec/data_dictionary/dialects/sqlite/sql/columns.sql @@ -33,7 +33,7 @@ ORDER BY tl.name, ti.cid; -- name: by_table -- dialect: sqlite SELECT - :schema_name AS schema_name, + COALESCE(:schema_name, 'main') AS schema_name, :table_name AS table_name, ti.name AS column_name, ti.type AS data_type, @@ -54,6 +54,6 @@ SELECT END AS extra, ti.hidden IN (2, 3) AS is_generated, CASE WHEN ti.hidden = 1 THEN 1 ELSE 0 END AS is_hidden -FROM pragma_table_xinfo(:table_name, :schema_name) AS ti +FROM pragma_table_xinfo(:table_name, COALESCE(:schema_name, 'main')) AS ti LEFT JOIN {schema_prefix}sqlite_schema AS sm ON sm.type = 'table' AND sm.name = :table_name ORDER BY ti.cid; diff --git a/sqlspec/data_dictionary/dialects/sqlite/sql/indexes.sql b/sqlspec/data_dictionary/dialects/sqlite/sql/indexes.sql index ccc3f9114..f777f4dd8 100644 --- a/sqlspec/data_dictionary/dialects/sqlite/sql/indexes.sql +++ b/sqlspec/data_dictionary/dialects/sqlite/sql/indexes.sql @@ -30,7 +30,7 @@ ORDER BY tl.name, il.seq, ix.seqno; -- name: by_table -- dialect: sqlite SELECT - :schema_name AS schema_name, + COALESCE(:schema_name, 'main') AS schema_name, :table_name AS table_name, il.name AS index_name, il."unique" AS is_unique, @@ -44,8 +44,8 @@ SELECT ix.coll AS collation, ix.key AS is_key_column, sm.sql AS index_sql -FROM pragma_index_list(:table_name, :schema_name) AS il -JOIN pragma_index_xinfo(il.name, :schema_name) AS ix +FROM pragma_index_list(:table_name, COALESCE(:schema_name, 'main')) AS il +JOIN pragma_index_xinfo(il.name, COALESCE(:schema_name, 'main')) AS ix LEFT JOIN {schema_prefix}sqlite_schema AS sm ON sm.type = 'index' AND sm.name = il.name @@ -55,7 +55,7 @@ ORDER BY il.seq, ix.seqno; -- name: columns_by_index -- dialect: sqlite SELECT - :schema_name AS schema_name, + COALESCE(:schema_name, 'main') AS schema_name, :index_name AS index_name, ix.seqno, ix.cid AS column_id, @@ -63,7 +63,7 @@ SELECT ix.desc AS is_descending, ix.coll AS collation, ix.key AS is_key_column -FROM pragma_index_xinfo(:index_name, :schema_name) AS ix +FROM pragma_index_xinfo(:index_name, COALESCE(:schema_name, 'main')) AS ix ORDER BY ix.seqno; -- name: index_columns_by_index diff --git a/tests/unit/data_dictionary/test_sqlite_duckdb_metadata.py b/tests/unit/data_dictionary/test_sqlite_duckdb_metadata.py index 933f10c19..9ff8c9f4f 100644 --- a/tests/unit/data_dictionary/test_sqlite_duckdb_metadata.py +++ b/tests/unit/data_dictionary/test_sqlite_duckdb_metadata.py @@ -12,6 +12,7 @@ from sqlspec.adapters.sqlite import SqliteConfig, SqliteDriver from sqlspec.adapters.sqlite.data_dictionary import SqliteDataDictionary from sqlspec.data_dictionary import DataDictionaryLoader, DependencyMetadata, MetadataSupport +from sqlspec.exceptions import SQLFileNotFoundError from tests.conftest import requires_interpreted pytestmark = requires_interpreted @@ -396,3 +397,73 @@ def test_sqlite_constraints_include_rowid_alias_primary_key() -> None: types = {item["constraint_type"]: item["columns"] for item in constraint_items} assert types["PRIMARY KEY"] == ["id"] assert types["UNIQUE"] == ["code"] + + +def test_sqlite_get_columns_and_indexes_with_none_schema() -> None: + """Introspection with schema=None should coalesce to main.""" + connection = sqlite3.connect(":memory:") + connection.row_factory = sqlite3.Row + driver = SqliteDriver(connection) + try: + driver.execute_script(""" + CREATE TABLE users ( + id INTEGER PRIMARY KEY, + name TEXT NOT NULL, + email TEXT NOT NULL + ); + CREATE INDEX idx_users_name ON users(name); + """) + + cols = driver.data_dictionary.get_columns(driver, table="users", schema=None) + assert len(cols) == 3 + assert all(c["schema_name"] == "main" for c in cols) + + indexes = driver.data_dictionary.get_indexes(driver, table="users", schema=None) + assert len(indexes) == 1 + assert indexes[0]["schema_name"] == "main" + assert indexes[0]["index_name"] == "idx_users_name" + finally: + connection.close() + + +def test_sqlite_data_dictionary_missing_query_raises_sql_file_not_found() -> None: + """Missing query in SQLite data dictionary raises SQLFileNotFoundError.""" + dd = SqliteDataDictionary() + with pytest.raises(SQLFileNotFoundError, match="Missing SQLite data-dictionary query"): + dd._get_domain_query_text("nonexistent_domain", "nonexistent_query") + + +def test_sqlite_generated_column_memoization(monkeypatch: pytest.MonkeyPatch) -> None: + """Verify generated column expressions are parsed once per table.""" + from sqlspec.adapters.sqlite import data_dictionary as dd_module + + connection = sqlite3.connect(":memory:") + connection.row_factory = sqlite3.Row + driver = SqliteDriver(connection) + try: + driver.execute_script(""" + CREATE TABLE gen_test ( + id INTEGER PRIMARY KEY, + a INT, + b INT, + c INT GENERATED ALWAYS AS (a + b) VIRTUAL, + d INT GENERATED ALWAYS AS (a * b) STORED + ); + """) + + call_count = 0 + original_fn = dd_module.sqlite_generated_expressions + + def counting_parser(sql: str | None) -> dict[str, str]: + nonlocal call_count + call_count += 1 + return original_fn(sql) + + monkeypatch.setattr(dd_module, "sqlite_generated_expressions", counting_parser) + + cols = driver.data_dictionary.get_columns(driver, table="gen_test") + assert call_count == 1 + gen_cols = {c["column_name"]: c["generation_expression"] for c in cols if c.get("is_generated")} + assert gen_cols == {"c": "a + b", "d": "a * b"} + finally: + connection.close() From 19db81351ecc3a2ed606f44a94f31e03af6870bd Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sat, 31 Oct 2026 22:47:06 +0000 Subject: [PATCH 015/174] perf(sqlite): optimize hot paths, parameter normalization, and savepoint handling Hoist connection parameter filtering helpers, epoch constants, and schema lists to module scope. Normalize cached query parameters, hoist cursor context in Arrow loading, and implement direct savepoint create/release without script splitting or rowid cache invalidation. --- sqlspec/adapters/sqlite/core.py | 65 ++++++++++--------- sqlspec/adapters/sqlite/driver.py | 59 ++++++++++++----- .../unit/adapters/test_sqlite/test_config.py | 8 +++ .../unit/adapters/test_sqlite/test_driver.py | 25 +++++++ .../test_sqlite/test_type_handlers.py | 15 +++++ 5 files changed, 125 insertions(+), 47 deletions(-) diff --git a/sqlspec/adapters/sqlite/core.py b/sqlspec/adapters/sqlite/core.py index 0e7fad619..aecdbb868 100644 --- a/sqlspec/adapters/sqlite/core.py +++ b/sqlspec/adapters/sqlite/core.py @@ -6,7 +6,7 @@ from collections.abc import Mapping from datetime import UTC, date, datetime, timedelta from decimal import Decimal -from typing import TYPE_CHECKING, Any, cast +from typing import TYPE_CHECKING, Any, Final, cast from sqlglot import exp @@ -45,6 +45,7 @@ "SECONDS_PER_DAY", "SQLITE_CONNECT_SUPPORTS_AUTOCOMMIT", "SqliteStreamSource", + "_filter_connection_params", "apply_driver_features", "apply_extension_pragmas", "build_connection_config", @@ -102,6 +103,8 @@ "PRAGMA mmap_size = 30000000", "PRAGMA journal_size_limit = 67108864", ) +_UNIX_EPOCH: Final[datetime] = datetime(1970, 1, 1, tzinfo=UTC) +_DEFAULT_SCHEMAS: Final[frozenset[str]] = frozenset({"temp", "main"}) def datetime_to_julian(dt: "datetime | None") -> "float | None": @@ -110,8 +113,7 @@ def datetime_to_julian(dt: "datetime | None") -> "float | None": return None if dt.tzinfo is None: dt = dt.replace(tzinfo=UTC) - epoch = datetime(1970, 1, 1, tzinfo=UTC) - delta_days = (dt - epoch).total_seconds() / SECONDS_PER_DAY + delta_days = (dt - _UNIX_EPOCH).total_seconds() / SECONDS_PER_DAY return JULIAN_EPOCH + delta_days @@ -121,7 +123,7 @@ def julian_to_datetime(julian: "float | None") -> "datetime | None": return None delta_days = julian - JULIAN_EPOCH timestamp = delta_days * SECONDS_PER_DAY - return datetime(1970, 1, 1, tzinfo=UTC) + timedelta(seconds=timestamp) + return _UNIX_EPOCH + timedelta(seconds=timestamp) def end_transaction( @@ -405,6 +407,32 @@ def normalize_execute_many_parameters(parameters: Any) -> Any: return parameters +_EXCLUDED_CONNECTION_KEYS: Final[frozenset[str]] = frozenset({ + "enable_foreign_keys", + "enable_optimizations", + "health_check_interval", + "pool_min_size", + "pool_max_size", + "pool_timeout", + "pool_recycle_seconds", + "extra", + "db", + "path", + "file", +}) + + +def _filter_connection_params(mapping: Mapping[str, Any]) -> dict[str, Any]: + """Filter connection parameters for SQLite.""" + return { + key: value + for key, value in mapping.items() + if key not in _EXCLUDED_CONNECTION_KEYS + and (value is not None or key == "isolation_level") + and (key != "autocommit" or SQLITE_CONNECT_SUPPORTS_AUTOCOMMIT) + } + + def build_connection_config(connection_config: "Mapping[str, Any]") -> "dict[str, Any]": """Build connection configuration for pool creation. @@ -414,30 +442,7 @@ def build_connection_config(connection_config: "Mapping[str, Any]") -> "dict[str Returns: Dictionary with connection parameters. """ - excluded_keys = { - "enable_foreign_keys", - "enable_optimizations", - "health_check_interval", - "pool_min_size", - "pool_max_size", - "pool_timeout", - "pool_recycle_seconds", - "extra", - "db", - "path", - "file", - } - - def _filter_params(mapping: Mapping[str, Any]) -> dict[str, Any]: - return { - key: value - for key, value in mapping.items() - if key not in excluded_keys - and (value is not None or key == "isolation_level") - and (key != "autocommit" or SQLITE_CONNECT_SUPPORTS_AUTOCOMMIT) - } - - connection_parameters = _filter_params(connection_config) + connection_parameters = _filter_connection_params(connection_config) if "database" not in connection_parameters: database = ( @@ -451,7 +456,7 @@ def _filter_params(mapping: Mapping[str, Any]) -> dict[str, Any]: extra = connection_config.get("extra") if isinstance(extra, Mapping): - connection_parameters.update(_filter_params(extra)) + connection_parameters.update(_filter_connection_params(extra)) return connection_parameters @@ -663,7 +668,7 @@ def _target_supports_rowid(connection: Any, target: "tuple[str | None, str]") -> return _target_supports_rowid_legacy(connection, target) schema_order = ["temp", "main"] - if not any(row[0].casefold() in {"temp", "main"} for row in candidates): + if not any(row[0].casefold() in _DEFAULT_SCHEMAS for row in candidates): try: database_cursor = connection.execute("PRAGMA database_list") try: diff --git a/sqlspec/adapters/sqlite/driver.py b/sqlspec/adapters/sqlite/driver.py index af3397eba..b5eacf16a 100644 --- a/sqlspec/adapters/sqlite/driver.py +++ b/sqlspec/adapters/sqlite/driver.py @@ -1,7 +1,7 @@ """SQLite driver implementation.""" import contextlib -from typing import TYPE_CHECKING, Any, Literal, cast +from typing import TYPE_CHECKING, Any, Final, Literal, cast from mypy_extensions import mypyc_attr @@ -25,6 +25,7 @@ from sqlspec.core.cache import get_cache_config from sqlspec.core.parameters._registry import register_driver_profile from sqlspec.core.result._base import ArrowResult, DMLResult +from sqlspec.driver._common import validate_savepoint_name from sqlspec.driver._exception_handler import BaseSyncExceptionHandler from sqlspec.driver._stream import SyncRowStream from sqlspec.driver._sync import SyncDriverAdapterBase @@ -44,6 +45,8 @@ __all__ = ("SqliteCursor", "SqliteDriver", "SqliteExceptionHandler", "SqliteSessionContext") +_ROWID_SAFE_OPERATIONS: Final[frozenset[str]] = frozenset({"SELECT", "INSERT", "UPDATE", "DELETE"}) + @mypyc_attr(allow_interpreted_subclasses=True) class SqliteExceptionHandler(BaseSyncExceptionHandler): @@ -242,6 +245,28 @@ def rollback(self) -> None: self._rowid_target_cache.clear() raise create_mapped_exception(e) from e + def create_savepoint(self, name: str) -> None: + """Create a savepoint within the current transaction without script splitting.""" + safe_name = validate_savepoint_name(name) + try: + cursor = self.connection.execute(f"SAVEPOINT {safe_name}") + if cursor is not None: + with contextlib.suppress(Exception): + cursor.close() + except sqlite3.Error as exc: + raise create_mapped_exception(exc) from exc + + def release_savepoint(self, name: str) -> None: + """Release a savepoint without clearing rowid target cache or splitting script.""" + safe_name = validate_savepoint_name(name) + try: + cursor = self.connection.execute(f"RELEASE SAVEPOINT {safe_name}") + if cursor is not None: + with contextlib.suppress(Exception): + cursor.close() + except sqlite3.Error as exc: + raise create_mapped_exception(exc) from exc + def with_cursor(self, connection: "SqliteConnection") -> "SqliteCursor": """Create context manager for SQLite cursor. @@ -335,7 +360,7 @@ def load_from_arrow( insert_sql = build_insert_statement(table, columns) needs_prep = self._arrow_rows_need_preparation(arrow_table) - owns_transaction = not self.connection.in_transaction + owns_transaction = not self._connection_in_transaction() try: if owns_transaction: transaction_mode = cast( @@ -343,20 +368,19 @@ def load_from_arrow( self.driver_features.get("default_transaction_mode") or "IMMEDIATE", ) self.begin(mode=transaction_mode) - if overwrite: - statement = f"DELETE FROM {format_identifier(table)}" - with self.with_cursor(self.connection) as cursor: + with self.with_cursor(self.connection) as cursor: + if overwrite: + statement = f"DELETE FROM {format_identifier(table)}" cursor.execute(statement) - for batch in arrow_table.to_batches(max_chunksize=batch_size): - pydict = batch.to_pydict() - records = list(zip(*[pydict[col] for col in columns], strict=False)) - if records: - prepared_records = ( - self.prepare_driver_parameters(records, self.statement_config, is_many=True) - if needs_prep - else records - ) - with self.with_cursor(self.connection) as cursor: + for batch in arrow_table.to_batches(max_chunksize=batch_size): + pydict = batch.to_pydict() + records = list(zip(*[pydict[col] for col in columns], strict=False)) + if records: + prepared_records = ( + self.prepare_driver_parameters(records, self.statement_config, is_many=True) + if needs_prep + else records + ) cursor.executemany(insert_sql, cast("Any", prepared_records)) if owns_transaction: self.commit() @@ -415,10 +439,11 @@ def _execute_cache_hit( """ cursor = None returns_rows = cached.operation_profile.returns_rows + normalized_parameters = normalize_execute_parameters(params) self._invalidate_rowid_target_cache(cached.operation_type) try: try: - cursor = self.connection.execute(cached.compiled_sql, params) + cursor = self.connection.execute(cached.compiled_sql, normalized_parameters) except sqlite3.Error as exc: raise create_mapped_exception(exc) from exc @@ -467,7 +492,7 @@ def _execute_cache_hit( cursor.close() def _invalidate_rowid_target_cache(self, operation_type: "OperationType") -> None: - if operation_type not in {"SELECT", "INSERT", "UPDATE", "DELETE"}: + if operation_type not in _ROWID_SAFE_OPERATIONS: self._rowid_target_cache.clear() def _connection_in_transaction(self) -> bool: diff --git a/tests/unit/adapters/test_sqlite/test_config.py b/tests/unit/adapters/test_sqlite/test_config.py index d2087278a..f6b3b8d2a 100644 --- a/tests/unit/adapters/test_sqlite/test_config.py +++ b/tests/unit/adapters/test_sqlite/test_config.py @@ -301,3 +301,11 @@ def worker() -> None: assert results == [0] finally: pool.close() + + +def test_filter_connection_params_with_hoisted_helper() -> None: + from sqlspec.adapters.sqlite.core import _filter_connection_params + + raw = {"database": ":memory:", "extra": {}, "pool_min_size": 10, "timeout": 30.0, "isolation_level": None} + filtered = _filter_connection_params(raw) + assert filtered == {"database": ":memory:", "timeout": 30.0, "isolation_level": None} diff --git a/tests/unit/adapters/test_sqlite/test_driver.py b/tests/unit/adapters/test_sqlite/test_driver.py index dae39f8bb..9de1917a3 100644 --- a/tests/unit/adapters/test_sqlite/test_driver.py +++ b/tests/unit/adapters/test_sqlite/test_driver.py @@ -353,3 +353,28 @@ def test_sqlite_dispatch_execute_script_emits_events() -> None: assert result.successful_statements == 2 finally: connection.close() + + +def test_sqlite_savepoint_does_not_clear_rowid_target_cache() -> None: + connection = sqlite3.connect(":memory:") + try: + driver = SqliteDriver(connection=connection) + driver._rowid_target_cache[(None, "test_target")] = True + driver.create_savepoint("sp1") + assert (None, "test_target") in driver._rowid_target_cache + driver.release_savepoint("sp1") + assert (None, "test_target") in driver._rowid_target_cache + finally: + connection.close() + + +def test_sqlite_cached_query_normalizes_parameters() -> None: + connection = sqlite3.connect(":memory:") + try: + driver = SqliteDriver(connection=connection) + driver.execute("CREATE TABLE norm_test (val INT)") + driver.execute("SELECT * FROM norm_test WHERE val = ?", [1]) + res = driver.execute("SELECT * FROM norm_test WHERE val = ?", [1]) + assert res.data == [] + finally: + connection.close() diff --git a/tests/unit/adapters/test_sqlite/test_type_handlers.py b/tests/unit/adapters/test_sqlite/test_type_handlers.py index 6e6fbe78b..7a962b829 100644 --- a/tests/unit/adapters/test_sqlite/test_type_handlers.py +++ b/tests/unit/adapters/test_sqlite/test_type_handlers.py @@ -53,3 +53,18 @@ def test_sqlite_config_does_not_register_type_handlers_by_default() -> None: config._create_pool() # pyright: ignore[reportPrivateUsage] mock_register.assert_not_called() + + +def test_julian_conversions_match_hoisted_epoch() -> None: + """Verify datetime_to_julian and julian_to_datetime use hoisted epoch.""" + from datetime import UTC, datetime + + from sqlspec.adapters.sqlite.core import _UNIX_EPOCH, datetime_to_julian, julian_to_datetime + + now = datetime(2026, 1, 1, 12, 0, 0, tzinfo=UTC) + julian = datetime_to_julian(now) + assert julian is not None + recovered = julian_to_datetime(julian) + assert recovered is not None + assert recovered == now + assert datetime_to_julian(_UNIX_EPOCH) == 2440587.5 From bf904d3560161087431803e36fa0a6744d6a3079 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sat, 31 Oct 2026 22:47:06 +0000 Subject: [PATCH 016/174] refactor(sqlite): clean up dead code, redundant wrappers, and connection lifecycle Track active connection checkouts, prune dead thread connections on size inspection, default transaction mode to IMMEDIATE with ImproperConfigurationError validation, delegate commit and rollback to end_transaction, remove redundant load_from_storage, and remove unused logger. --- sqlspec/adapters/sqlite/_typing.py | 15 ++++-- sqlspec/adapters/sqlite/config.py | 2 +- sqlspec/adapters/sqlite/driver.py | 52 ++++++------------- sqlspec/adapters/sqlite/pool.py | 40 +++++++++++++- sqlspec/adapters/sqlite/type_converter.py | 3 -- .../unit/adapters/test_sqlite/test_driver.py | 52 ++++++++++++++++++- tests/unit/adapters/test_sqlite/test_pool.py | 22 ++++++++ 7 files changed, 140 insertions(+), 46 deletions(-) diff --git a/sqlspec/adapters/sqlite/_typing.py b/sqlspec/adapters/sqlite/_typing.py index e2661eff7..83abb5df9 100644 --- a/sqlspec/adapters/sqlite/_typing.py +++ b/sqlspec/adapters/sqlite/_typing.py @@ -6,7 +6,7 @@ import contextlib import sqlite3 as sqlite_module -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, ClassVar _SqliteConnection = sqlite_module.Connection @@ -90,12 +90,15 @@ class SqliteSessionContext: "_acquire_connection", "_connection", "_driver", + "_driver_class", "_driver_features", "_prepare_driver", "_release_connection", "_statement_config", ) + _default_driver_class: ClassVar[Any] = None + def __init__( self, acquire_connection: "Callable[[], Any]", @@ -103,20 +106,26 @@ def __init__( statement_config: "StatementConfig", driver_features: "dict[str, Any]", prepare_driver: "Callable[[SqliteDriver], SqliteDriver]", + driver_class: "type[SqliteDriver] | None" = None, ) -> None: self._acquire_connection = acquire_connection self._release_connection = release_connection self._statement_config = statement_config self._driver_features = driver_features self._prepare_driver = prepare_driver + self._driver_class = driver_class or self._default_driver_class self._connection: Any = None self._driver: SqliteDriver | None = None def __enter__(self) -> "SqliteDriver": - from sqlspec.adapters.sqlite.driver import SqliteDriver + driver_cls = self._driver_class + if driver_cls is None: + from sqlspec.adapters.sqlite.driver import SqliteDriver + + driver_cls = SqliteDriver self._connection = self._acquire_connection() - self._driver = SqliteDriver( + self._driver = driver_cls( connection=self._connection, statement_config=self._statement_config, driver_features=self._driver_features ) return self._prepare_driver(self._driver) diff --git a/sqlspec/adapters/sqlite/config.py b/sqlspec/adapters/sqlite/config.py index 0260d8c81..b7a30de6a 100644 --- a/sqlspec/adapters/sqlite/config.py +++ b/sqlspec/adapters/sqlite/config.py @@ -76,7 +76,7 @@ class SqliteConnectionParams(TypedDict): enable_optimizations: NotRequired[bool] """Apply recommended SQLite ``PRAGMA`` settings (WAL journal mode, synchronous=NORMAL, busy_timeout) on connect.""" enable_foreign_keys: NotRequired[bool] - """Execute ``PRAGMA foreign_keys = ON`` when initializing each connection (defaults to ``True``).""" + """Execute ``PRAGMA foreign_keys = ON`` when initializing each connection (defaults to ``False``).""" extra: NotRequired[dict[str, Any]] """Additional keyword arguments forwarded to ``sqlite3.connect``.""" diff --git a/sqlspec/adapters/sqlite/driver.py b/sqlspec/adapters/sqlite/driver.py index b5eacf16a..6165c0076 100644 --- a/sqlspec/adapters/sqlite/driver.py +++ b/sqlspec/adapters/sqlite/driver.py @@ -15,6 +15,7 @@ create_mapped_exception, default_statement_config, driver_profile, + end_transaction, format_identifier, normalize_execute_many_parameters, normalize_execute_parameters, @@ -29,6 +30,7 @@ from sqlspec.driver._exception_handler import BaseSyncExceptionHandler from sqlspec.driver._stream import SyncRowStream from sqlspec.driver._sync import SyncDriverAdapterBase +from sqlspec.exceptions import ImproperConfigurationError from sqlspec.utils.type_guards import resolve_row_format if TYPE_CHECKING: @@ -195,16 +197,21 @@ def begin(self, mode: "Literal['DEFERRED', 'IMMEDIATE', 'EXCLUSIVE'] | None" = N Raises: SQLSpecError: If transaction cannot be started """ - transaction_mode = mode if mode is not None else self.driver_features.get("default_transaction_mode") + transaction_mode = ( + mode if mode is not None else self.driver_features.get("default_transaction_mode", "IMMEDIATE") + ) if transaction_mode is not None and transaction_mode not in {"DEFERRED", "IMMEDIATE", "EXCLUSIVE"}: msg = "Transaction mode must be DEFERRED, IMMEDIATE, or EXCLUSIVE" - raise ValueError(msg) + raise ImproperConfigurationError(msg) try: if not self.connection.in_transaction: stmt = f"BEGIN {transaction_mode}" if transaction_mode else "BEGIN" cursor = self.connection.execute(stmt) - with contextlib.suppress(Exception): - cursor.close() + with contextlib.suppress(AttributeError, TypeError): + self.connection.in_transaction = True + if cursor is not None: + with contextlib.suppress(Exception): + cursor.close() except sqlite3.Error as e: raise create_mapped_exception(e) from e @@ -215,13 +222,9 @@ def commit(self) -> None: SQLSpecError: If transaction cannot be committed """ try: - if self._in_autocommit_mode(): - if self.connection.in_transaction: - cursor = self.connection.execute("COMMIT") - with contextlib.suppress(Exception): - cursor.close() - return - self.connection.commit() + end_transaction(self.connection, commit=True, supports_autocommit=self._in_autocommit_mode()) + with contextlib.suppress(AttributeError, TypeError): + self.connection.in_transaction = False except sqlite3.Error as e: raise create_mapped_exception(e) from e @@ -232,14 +235,9 @@ def rollback(self) -> None: SQLSpecError: If transaction cannot be rolled back """ try: - if self._in_autocommit_mode(): - if self.connection.in_transaction: - cursor = self.connection.execute("ROLLBACK") - with contextlib.suppress(Exception): - cursor.close() - self._rowid_target_cache.clear() - return - self.connection.rollback() + end_transaction(self.connection, commit=False, supports_autocommit=self._in_autocommit_mode()) + with contextlib.suppress(AttributeError, TypeError): + self.connection.in_transaction = False self._rowid_target_cache.clear() except sqlite3.Error as e: self._rowid_target_cache.clear() @@ -396,22 +394,6 @@ def load_from_arrow( self._attach_partition_telemetry(telemetry_payload, partitioner) return self._storage_job(telemetry_payload, telemetry) - def load_from_storage( - self, - table: str, - source: "StorageDestination", - *, - file_format: "StorageFormat", - partitioner: "dict[str, object] | None" = None, - overwrite: bool = False, - **kwargs: Any, - ) -> "StorageBridgeJob": - """Load staged artifacts from storage into SQLite.""" - - return super().load_from_storage( - table, source, file_format=file_format, partitioner=partitioner, overwrite=overwrite, **kwargs - ) - def collect_rows(self, cursor: Any, fetched: "list[Any]") -> "tuple[list[Any], list[str], int]": """Collect SQLite rows for the direct execution path.""" return collect_rows(fetched, cursor.description) diff --git a/sqlspec/adapters/sqlite/pool.py b/sqlspec/adapters/sqlite/pool.py index 3de1a227a..9d8487cab 100644 --- a/sqlspec/adapters/sqlite/pool.py +++ b/sqlspec/adapters/sqlite/pool.py @@ -39,6 +39,7 @@ class SqliteConnectionPool: __slots__ = ( "_anchor_connection", + "_checked_out_count", "_connection_parameters", "_connection_registry", "_enable_foreign_keys", @@ -52,6 +53,7 @@ class SqliteConnectionPool: "_recycle_seconds", "_registry_lock", "_runtime_setup", + "_thread_connections", "_thread_local", ) @@ -86,6 +88,8 @@ def __init__( self._anchor_connection: SqliteConnection | None = None self._thread_local = threading.local() self._connection_registry: set[SqliteConnection] = set() + self._thread_connections: dict[threading.Thread, SqliteConnection] = {} + self._checked_out_count = 0 self._generation = 0 self._registry_lock = threading.Lock() self._enable_optimizations = enable_optimizations @@ -155,6 +159,9 @@ def get_connection(self) -> "Generator[SqliteConnection, None, None]": connection = self._get_thread_connection() depth = getattr(self._thread_local, "depth", 0) self._thread_local.depth = depth + 1 + if self._thread_local.depth == 1: + with self._registry_lock: + self._checked_out_count += 1 try: yield connection except BaseException: @@ -171,6 +178,9 @@ def get_connection(self) -> "Generator[SqliteConnection, None, None]": end_transaction(connection, commit=False) raise finally: + if self._thread_local.depth == 1: + with self._registry_lock: + self._checked_out_count -= 1 self._thread_local.depth -= 1 def close(self) -> None: @@ -179,6 +189,8 @@ def close(self) -> None: with self._registry_lock: orphaned = list(self._connection_registry) self._connection_registry.clear() + self._thread_connections.clear() + self._checked_out_count = 0 anchor = self._anchor_connection self._anchor_connection = None self._generation += 1 @@ -206,14 +218,27 @@ def release(self, connection: SqliteConnection) -> None: with contextlib.suppress(Exception): end_transaction(connection, commit=False) + def _prune_dead_threads(self) -> None: + """Prune connections owned by dead threads.""" + with self._registry_lock: + dead_threads = [t for t in self._thread_connections if not t.is_alive()] + for t in dead_threads: + conn = self._thread_connections.pop(t, None) + if conn is not None: + self._connection_registry.discard(conn) + with contextlib.suppress(Exception): + conn.close() + def size(self) -> int: """Get total number of active thread-local connections in pool.""" + self._prune_dead_threads() with self._registry_lock: return len(self._connection_registry) def checked_out(self) -> int: - """Get number of checked out connections (always 0).""" - return 0 + """Get number of checked out connections.""" + with self._registry_lock: + return self._checked_out_count def _create_connection(self) -> SqliteConnection: """Create a pool-owned connection and record it for shutdown.""" @@ -263,6 +288,8 @@ def _get_thread_connection(self) -> SqliteConnection: self._thread_local.connection = conn self._thread_local.created_at = now self._thread_local.last_used = now + with self._registry_lock: + self._thread_connections[threading.current_thread()] = conn return conn created_at = getattr(self._thread_local, "created_at", 0.0) @@ -282,6 +309,8 @@ def _get_thread_connection(self) -> SqliteConnection: self._thread_local.connection = conn self._thread_local.created_at = now self._thread_local.last_used = now + with self._registry_lock: + self._thread_connections[threading.current_thread()] = conn return conn last_used = getattr(self._thread_local, "last_used", 0.0) @@ -302,15 +331,22 @@ def _get_thread_connection(self) -> SqliteConnection: self._thread_local.connection = conn self._thread_local.created_at = now self._thread_local.last_used = now + with self._registry_lock: + self._thread_connections[threading.current_thread()] = conn return conn self._thread_local.last_used = now + with self._registry_lock: + self._thread_connections[threading.current_thread()] = conn return cast("SqliteConnection", conn) def _retire_connection(self, connection: SqliteConnection) -> None: """Close a pool-owned connection and drop it from the shutdown registry.""" with self._registry_lock: self._connection_registry.discard(connection) + dead_threads = [t for t, c in self._thread_connections.items() if c is connection] + for t in dead_threads: + self._thread_connections.pop(t, None) with contextlib.suppress(Exception): connection.close() diff --git a/sqlspec/adapters/sqlite/type_converter.py b/sqlspec/adapters/sqlite/type_converter.py index 74f1c7d91..d47069d4a 100644 --- a/sqlspec/adapters/sqlite/type_converter.py +++ b/sqlspec/adapters/sqlite/type_converter.py @@ -13,15 +13,12 @@ from typing import TYPE_CHECKING, Any from sqlspec.adapters.sqlite._typing import sqlite_module as sqlite3 -from sqlspec.utils.logging import get_logger if TYPE_CHECKING: from collections.abc import Callable __all__ = ("json_adapter", "json_converter", "register_type_handlers", "unregister_type_handlers") -logger = get_logger(__name__) - DEFAULT_JSON_TYPE = "JSON" diff --git a/tests/unit/adapters/test_sqlite/test_driver.py b/tests/unit/adapters/test_sqlite/test_driver.py index 9de1917a3..c5b36e877 100644 --- a/tests/unit/adapters/test_sqlite/test_driver.py +++ b/tests/unit/adapters/test_sqlite/test_driver.py @@ -5,7 +5,7 @@ import pytest from sqlspec.adapters.sqlite import SqliteConfig, SqliteDriver -from sqlspec.adapters.sqlite.core import SQLITE_CONNECT_SUPPORTS_AUTOCOMMIT +from sqlspec.adapters.sqlite.core import SQLITE_CONNECT_SUPPORTS_AUTOCOMMIT, default_statement_config from sqlspec.adapters.sqlite.pool import SqliteConnectionPool from sqlspec.adapters.sqlite.type_converter import json_converter, register_type_handlers, unregister_type_handlers from sqlspec.exceptions import ImproperConfigurationError @@ -271,7 +271,7 @@ def test_sqlite_begin_mode_and_arrow_batch_validation() -> None: driver.rollback() driver.begin(mode="DEFERRED") driver.rollback() - with pytest.raises(ValueError, match="Transaction mode"): + with pytest.raises((ValueError, ImproperConfigurationError), match="Transaction mode"): driver.begin(mode=cast("Any", "INVALID")) with pytest.raises(ValueError, match="batch_size"): driver.load_from_arrow("sample", cast("Any", None), batch_size=0) @@ -378,3 +378,51 @@ def test_sqlite_cached_query_normalizes_parameters() -> None: assert res.data == [] finally: connection.close() + + +def test_sqlite_begin_defaults_to_immediate_and_rejects_invalid_mode() -> None: + connection = sqlite3.connect(":memory:") + try: + driver = SqliteDriver(connection=connection) + with pytest.raises(ImproperConfigurationError, match="Transaction mode"): + driver.begin(mode=cast("Any", "INVALID")) + driver.begin() + assert connection.in_transaction + driver.rollback() + finally: + connection.close() + + +def test_sqlite_commit_and_rollback_delegate_to_end_transaction() -> None: + connection = sqlite3.connect(":memory:") + try: + driver = SqliteDriver(connection=connection) + driver.execute("CREATE TABLE cr_test (x INT)") + driver.begin() + driver.execute("INSERT INTO cr_test VALUES (1)") + driver.commit() + assert not connection.in_transaction + + driver.begin() + driver.execute("INSERT INTO cr_test VALUES (2)") + driver.rollback() + assert not connection.in_transaction + + res = driver.execute("SELECT count(*) FROM cr_test") + assert res.data == [(1,)] + finally: + connection.close() + + +def test_sqlite_session_context_initializes_with_driver_class() -> None: + from sqlspec.adapters.sqlite._typing import SqliteSessionContext + + ctx = SqliteSessionContext( + acquire_connection=lambda: None, + release_connection=lambda *_: None, + statement_config=default_statement_config, + driver_features={}, + prepare_driver=lambda d: d, + driver_class=SqliteDriver, + ) + assert ctx._driver_class is SqliteDriver diff --git a/tests/unit/adapters/test_sqlite/test_pool.py b/tests/unit/adapters/test_sqlite/test_pool.py index d8a49b718..1f09f9353 100644 --- a/tests/unit/adapters/test_sqlite/test_pool.py +++ b/tests/unit/adapters/test_sqlite/test_pool.py @@ -427,3 +427,25 @@ def test_nested_get_connection_reentrancy_depth() -> None: assert rows == [(1,), (2,)] finally: pool.close() + + +def test_sqlite_pool_checked_out_reflects_active_sessions_and_prunes_dead_threads() -> None: + pool = SqliteConnectionPool({"database": ":memory:"}) + try: + assert pool.checked_out() == 0 + with pool.get_connection(): + assert pool.checked_out() == 1 + assert pool.checked_out() == 0 + + def worker() -> None: + with pool.get_connection() as conn: + conn.execute("SELECT 1") + + t = threading.Thread(target=worker) + t.start() + t.join() + + assert not t.is_alive() + assert pool.size() == 1 + finally: + pool.close() From 8a9b64a87e34e23e5cf242f8154192d684aae33a Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sat, 31 Oct 2026 22:52:36 +0000 Subject: [PATCH 017/174] fix(aiomysql): escape LOAD DATA statement percent signs and restore cache hit fast path Ensure LOAD DATA statements escape percent characters in identifiers for PyMySQL pyformat interpolation. Check literal percent escaping in cached queries so clean queries delegate directly to the driver fast path without redundant statement allocations. --- sqlspec/adapters/aiomysql/driver.py | 18 +++++--- .../adapters/test_aiomysql/test_driver.py | 42 +++++++++++++++++++ .../test_aiomysql/test_load_from_arrow.py | 11 +++++ 3 files changed, 66 insertions(+), 5 deletions(-) diff --git a/sqlspec/adapters/aiomysql/driver.py b/sqlspec/adapters/aiomysql/driver.py index 496fcbcd7..36bd446ad 100644 --- a/sqlspec/adapters/aiomysql/driver.py +++ b/sqlspec/adapters/aiomysql/driver.py @@ -380,7 +380,7 @@ async def load_from_arrow( try: with os.fdopen(fd, "wb") as tmp: tmp.write(payload) - load_sql = build_load_data_statement(table, columns) + load_sql = build_load_data_statement(table, columns, escape_percent=True) exc_handler = self.handle_database_exceptions() await self._run_with_exception_handler( exc_handler, self._execute_cursor_statement, load_sql, (tmp_name,) @@ -436,10 +436,18 @@ def resolve_rowcount(self, cursor: Any) -> int: async def _execute_cache_hit( self, sql: str, params: "tuple[Any, ...] | list[Any] | dict[str, Any]", cached: "CachedQuery" ) -> "SQLResult": - statement = self._cached_statement( - sql, params, cached, params, params_are_simple=True, compiled_sql=cached.compiled_sql - ) - return await self._execute_cached_statement(statement) + if ( + "%" in cached.compiled_sql + and escape_literal_percent( + cached.compiled_sql, params or (None,), self.statement_config.parameter_validator + ) + != cached.compiled_sql + ): + statement = self._cached_statement( + sql, params, cached, params, params_are_simple=True, compiled_sql=cached.compiled_sql + ) + return await self._execute_cached_statement(statement) + return await super()._execute_cache_hit(sql, params, cached) async def _execute_cursor_statement(self, statement: str, parameters: Any = None) -> None: async with self.with_cursor(self.connection) as cursor: diff --git a/tests/unit/adapters/test_aiomysql/test_driver.py b/tests/unit/adapters/test_aiomysql/test_driver.py index 304d3f3ec..30402407b 100644 --- a/tests/unit/adapters/test_aiomysql/test_driver.py +++ b/tests/unit/adapters/test_aiomysql/test_driver.py @@ -97,3 +97,45 @@ def test_aiomysql_signature_namespace_includes_cursors() -> None: ns = config.get_signature_namespace() assert "AiomysqlDictCursor" in ns assert "AiomysqlRawCursor" in ns + + +async def test_aiomysql_execute_cache_hit_invokes_collect_rows_on_clean_query(monkeypatch: pytest.MonkeyPatch) -> None: + """Cached query without unescaped literal percent delegates to fast path and invokes collect_rows.""" + from unittest.mock import AsyncMock, MagicMock + + from sqlspec.core import OperationProfile, ParameterProfile, ProcessedState + from sqlspec.core.cache import CachedQuery + + cursor = MagicMock() + cursor.description = [("id", 3, None, None, None, None, None)] + cursor.fetchall = AsyncMock(return_value=[(1,)]) + cursor.execute = AsyncMock(return_value=None) + + conn = MagicMock() + conn.cursor = AsyncMock(return_value=cursor) + + driver = AiomysqlDriver(connection=cast("Any", conn)) + collect_called: list[bool] = [] + original_collect_rows = AiomysqlDriver.collect_rows + + def _spy_collect_rows(*args: Any, **kwargs: Any) -> Any: + collect_called.append(True) + return original_collect_rows(*args, **kwargs) + + monkeypatch.setattr(AiomysqlDriver, "collect_rows", _spy_collect_rows) + + cached = CachedQuery( + compiled_sql="SELECT 1", + parameter_profile=ParameterProfile(), + input_named_parameters=(), + applied_wrap_types=False, + parameter_casts={}, + operation_type="SELECT", + operation_profile=OperationProfile(returns_rows=True, modifies_rows=False), + param_count=0, + processed_state=ProcessedState(compiled_sql="SELECT 1", execution_parameters=[], operation_type="SELECT"), + ) + + result = await driver._execute_cache_hit("SELECT 1", (), cached) + assert bool(collect_called) is True + assert result.data == [(1,)] diff --git a/tests/unit/adapters/test_aiomysql/test_load_from_arrow.py b/tests/unit/adapters/test_aiomysql/test_load_from_arrow.py index 6c2937303..8fb2715c3 100644 --- a/tests/unit/adapters/test_aiomysql/test_load_from_arrow.py +++ b/tests/unit/adapters/test_aiomysql/test_load_from_arrow.py @@ -122,3 +122,14 @@ async def test_load_from_arrow_binds_the_infile_path_as_a_parameter() -> None: assert "LOAD DATA LOCAL INFILE %s INTO TABLE" in statement assert parameters[0].endswith(".tsv") assert parameters[0] not in statement + + +async def test_load_from_arrow_percent_in_column_escapes_pyformat() -> None: + """Percent signs in table or column identifiers must be escaped for pyformat.""" + conn = _FakeConnection() + driver = _make_driver(conn, enable_local_infile=True) + + job = await driver.load_from_arrow("orders", pa.table({"discount%rate": [1, 2]})) + + assert job.telemetry["rows_processed"] == 2 + assert "(`discount%%rate`)" in conn._cursor.execute_calls[0] From 921e8725facc21f9b5eb69800da7c44e56c00a4a Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sat, 31 Oct 2026 22:52:36 +0000 Subject: [PATCH 018/174] fix(aiomysql): batch ADK memory insertions and use cross-version MySQL upsert syntax Batch ADK memory entries with cursor.executemany rather than sequential execute calls in a loop. Drop redundant DROP INDEX statement from session table teardown and use VALUES(col) in Litestar session upsert statements for cross-version MySQL and MariaDB compatibility. --- sqlspec/adapters/aiomysql/adk/store.py | 88 ++++++++++--------- sqlspec/adapters/aiomysql/litestar/store.py | 13 ++- .../adapters/test_aiomysql/test_adk_store.py | 48 +++++++++- .../test_aiomysql/test_litestar_store.py | 35 ++++++++ 4 files changed, 134 insertions(+), 50 deletions(-) diff --git a/sqlspec/adapters/aiomysql/adk/store.py b/sqlspec/adapters/aiomysql/adk/store.py index 3a9016b6e..d0f4458d6 100644 --- a/sqlspec/adapters/aiomysql/adk/store.py +++ b/sqlspec/adapters/aiomysql/adk/store.py @@ -5,7 +5,7 @@ from typing import TYPE_CHECKING, Any, Final, Literal, NotRequired, cast from sqlspec.adapters.aiomysql._typing import AiomysqlCursor, AiomysqlProgrammingError, AiomysqlRawCursor -from sqlspec.adapters.aiomysql.core import resolve_rowcount +from sqlspec.adapters.aiomysql.core import resolve_many_rowcount, resolve_rowcount from sqlspec.config import ADKConfig from sqlspec.extensions.adk import BaseAsyncADKStore, StoredEvent, StoredSession, normalize_session_list_options from sqlspec.extensions.adk.memory.store import BaseAsyncADKMemoryStore @@ -566,46 +566,52 @@ async def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: " ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) """ - async with self._config.provide_connection() as conn: - async with AiomysqlCursor(conn, cursor_class=AiomysqlRawCursor) as cursor: - for entry in entries: - content_json = self._json_serializer(entry["content_json"]) - meta = entry.get("metadata_json") - metadata_json = self._json_serializer(meta) if meta is not None else None - params: tuple[Any, ...] - if self._owner_id_column_name: - params = ( - entry["id"], - entry["session_id"], - entry["app_name"], - entry["user_id"], - entry.get("scope", "user"), - entry["event_id"], - entry.get("author"), - owner_id, - entry["timestamp"], - content_json, - entry["content_text"], - metadata_json, - entry["inserted_at"], - ) - else: - params = ( - entry["id"], - entry["session_id"], - entry["app_name"], - entry["user_id"], - entry.get("scope", "user"), - entry["event_id"], - entry.get("author"), - entry["timestamp"], - content_json, - entry["content_text"], - metadata_json, - entry["inserted_at"], - ) - await cursor.execute(sql, params) - inserted_count += resolve_rowcount(cursor) + params_list: list[tuple[Any, ...]] = [] + for entry in entries: + content_json = self._json_serializer(entry["content_json"]) + meta = entry.get("metadata_json") + metadata_json = self._json_serializer(meta) if meta is not None else None + if self._owner_id_column_name: + params_list.append(( + entry["id"], + entry["session_id"], + entry["app_name"], + entry["user_id"], + entry.get("scope", "user"), + entry["event_id"], + entry.get("author"), + owner_id, + entry["timestamp"], + content_json, + entry["content_text"], + metadata_json, + entry["inserted_at"], + )) + else: + params_list.append(( + entry["id"], + entry["session_id"], + entry["app_name"], + entry["user_id"], + entry.get("scope", "user"), + entry["event_id"], + entry.get("author"), + entry["timestamp"], + content_json, + entry["content_text"], + metadata_json, + entry["inserted_at"], + )) + + if not params_list: + return 0 + + async with ( + self._config.provide_connection() as conn, + AiomysqlCursor(conn, cursor_class=AiomysqlRawCursor) as cursor, + ): + await cursor.executemany(sql, params_list) + inserted_count = resolve_many_rowcount(cursor, params_list, fallback_count=len(params_list)) await conn.commit() return inserted_count diff --git a/sqlspec/adapters/aiomysql/litestar/store.py b/sqlspec/adapters/aiomysql/litestar/store.py index 8f8b449f7..b2b154786 100644 --- a/sqlspec/adapters/aiomysql/litestar/store.py +++ b/sqlspec/adapters/aiomysql/litestar/store.py @@ -140,10 +140,10 @@ async def set(self, key: str, value: "str | bytes", expires_in: "int | timedelta sql = f""" INSERT INTO {self._table_name} (session_id, data, expires_at) - VALUES (%s, %s, %s) AS new + VALUES (%s, %s, %s) ON DUPLICATE KEY UPDATE - data = new.data, - expires_at = new.expires_at, + data = VALUES(data), + expires_at = VALUES(expires_at), updated_at = UTC_TIMESTAMP(6) """ @@ -301,12 +301,9 @@ def _drop_table_sql(self) -> "list[str]": """Get MySQL/MariaDB DROP TABLE SQL statements. Returns: - List of SQL statements to drop indexes and table. + List of SQL statements to drop table. """ - return [ - f"DROP INDEX idx_{self._table_name}_expires_at ON {self._table_name}", - f"DROP TABLE IF EXISTS {self._table_name}", - ] + return [f"DROP TABLE IF EXISTS {self._table_name}"] def _is_mysql_table_missing(exc: BaseException) -> bool: diff --git a/tests/unit/adapters/test_aiomysql/test_adk_store.py b/tests/unit/adapters/test_aiomysql/test_adk_store.py index ea75d6e65..fabfb7049 100644 --- a/tests/unit/adapters/test_aiomysql/test_adk_store.py +++ b/tests/unit/adapters/test_aiomysql/test_adk_store.py @@ -223,6 +223,7 @@ async def test_aiomysql_adk_memory_store_insert_and_search_json_handling() -> No cursor = MagicMock() cursor.rowcount = 1 cursor.execute = AsyncMock() + cursor.executemany = AsyncMock() cursor.description = [ ("id",), ("session_id",), @@ -284,7 +285,7 @@ async def test_aiomysql_adk_memory_store_insert_and_search_json_handling() -> No records = await store.search_entries("hello", "app", "user-1") assert inserted == 1 - insert_params = cursor.execute.call_args_list[0].args[1] + insert_params = cursor.executemany.call_args_list[0].args[1][0] assert insert_params[6] is None assert insert_params[11] is None assert len(records) == 1 @@ -292,6 +293,51 @@ async def test_aiomysql_adk_memory_store_insert_and_search_json_handling() -> No assert records[0]["metadata_json"] == {"source": "unit"} +async def test_aiomysql_insert_memory_entries_batches_with_executemany() -> None: + """insert_memory_entries batches multiple entries in a single cursor.executemany call.""" + from datetime import UTC, datetime + from unittest.mock import AsyncMock + + now = datetime.now(tz=UTC) + cursor = MagicMock() + cursor.rowcount = 2 + cursor.executemany = AsyncMock() + conn = MagicMock() + conn.cursor = AsyncMock(return_value=cursor) + conn.commit = AsyncMock() + conn.__aenter__ = AsyncMock(return_value=conn) + conn.__aexit__ = AsyncMock(return_value=None) + config = _mock_config() + config.provide_connection = lambda *_a, **_k: conn + store = AiomysqlADKMemoryStore(config) + + entries = [ + cast( + "Any", + { + "id": f"mem-{i}", + "session_id": "sess-1", + "app_name": "app", + "user_id": "user-1", + "event_id": f"evt-{i}", + "timestamp": now, + "content_json": {"idx": i}, + "content_text": f"text {i}", + "metadata_json": None, + "inserted_at": now, + }, + ) + for i in range(2) + ] + + inserted = await store.insert_memory_entries(entries) + + assert inserted == 2 + assert cursor.executemany.call_count == 1 + batch_params = cursor.executemany.call_args_list[0].args[1] + assert len(batch_params) == 2 + + async def test_aiomysql_stream_source_closes_cursor_when_execute_raises() -> None: """AiomysqlStreamSource.start() closes the cursor if execute() fails.""" from unittest.mock import AsyncMock diff --git a/tests/unit/adapters/test_aiomysql/test_litestar_store.py b/tests/unit/adapters/test_aiomysql/test_litestar_store.py index 38a61079c..8d7c74895 100644 --- a/tests/unit/adapters/test_aiomysql/test_litestar_store.py +++ b/tests/unit/adapters/test_aiomysql/test_litestar_store.py @@ -30,3 +30,38 @@ def test_aiomysql_litestar_store_honors_table_options() -> None: ddl = store._table_ddl() assert f"{DEFAULT_ENGINE_CLAUSE} COMMENT='litestar-session'" in ddl + + +def test_aiomysql_litestar_store_drop_table_sql_omits_drop_index() -> None: + """_drop_table_sql must emit only DROP TABLE IF EXISTS without preceding DROP INDEX.""" + store = AiomysqlStore(_mock_config()) + + statements = store._drop_table_sql() + + assert len(statements) == 1 + assert statements[0] == "DROP TABLE IF EXISTS litestar_session" + + +async def test_aiomysql_litestar_store_set_uses_mariadb_compatible_upsert() -> None: + """_set must use VALUES(col) in ON DUPLICATE KEY UPDATE for MySQL and MariaDB compatibility.""" + from unittest.mock import AsyncMock + + cursor = MagicMock() + cursor.execute = AsyncMock() + conn = MagicMock() + conn.cursor = AsyncMock(return_value=cursor) + conn.commit = AsyncMock() + conn.__aenter__ = AsyncMock(return_value=conn) + conn.__aexit__ = AsyncMock(return_value=None) + config = _mock_config() + config.provide_connection = lambda *_a, **_k: conn + + store = AiomysqlStore(config) + await store.set("session-1", b"payload", 300) + + assert cursor.execute.call_count == 1 + sql = cursor.execute.call_args_list[0].args[0] + assert "ON DUPLICATE KEY UPDATE" in sql + assert "data = VALUES(data)" in sql + assert "expires_at = VALUES(expires_at)" in sql + assert "AS new" not in sql From b6717b55ab026596645375a0b12c80361774452a Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sat, 31 Oct 2026 22:52:36 +0000 Subject: [PATCH 019/174] perf(aiomysql): add fast path to literal percent escaping and eliminate generator closures in compiled helpers Add O(1) fast path to escape_literal_percent that returns the SQL string unchanged when parameter and percent counts match. Eliminate generator expressions from localized core helpers to avoid heap closures under MyPyC, and offload LOCAL INFILE tempfile creation and payload writes to asyncio.to_thread. --- sqlspec/adapters/aiomysql/config.py | 7 ++++++ sqlspec/adapters/aiomysql/core.py | 24 ++++++++++++------- sqlspec/adapters/aiomysql/driver.py | 12 +++++++--- .../adapters/test_aiomysql/test_config.py | 7 ++++++ .../adapters/test_mysql_parameter_profiles.py | 9 +++++++ 5 files changed, 48 insertions(+), 11 deletions(-) diff --git a/sqlspec/adapters/aiomysql/config.py b/sqlspec/adapters/aiomysql/config.py index 701380930..31fa29daa 100644 --- a/sqlspec/adapters/aiomysql/config.py +++ b/sqlspec/adapters/aiomysql/config.py @@ -171,6 +171,13 @@ def build_connection_config( class _AiomysqlSessionFactory(AsyncPoolSessionFactory): + """Session factory for aiomysql connections. + + The ``_contexts`` dictionary maps connection IDs to pool acquire contexts + to track and release connections safely across concurrent session checkouts + from the same factory instance. + """ + __slots__ = ("_contexts",) def __init__(self, config: "AiomysqlConfig") -> None: diff --git a/sqlspec/adapters/aiomysql/core.py b/sqlspec/adapters/aiomysql/core.py index e4ffb3ac3..bb98910ac 100644 --- a/sqlspec/adapters/aiomysql/core.py +++ b/sqlspec/adapters/aiomysql/core.py @@ -262,14 +262,14 @@ def format_identifier(identifier: str) -> str: msg = "Table name must not be empty" raise SQLSpecError(msg) parts = split_qualified_identifier(cleaned, quote_chars="`", allow_bracket_quotes=False) - formatted = ".".join(quote_backtick_identifier(part) for part in parts) + formatted = ".".join([quote_backtick_identifier(part) for part in parts]) return formatted or quote_backtick_identifier(cleaned) def build_insert_statement(table: str, columns: "list[str]") -> str: """Construct an INSERT statement with positional placeholders.""" - column_clause = ", ".join(quote_backtick_identifier(column) for column in columns) - placeholders = ", ".join("%s" for _ in columns) + column_clause = ", ".join([quote_backtick_identifier(column) for column in columns]) + placeholders = ", ".join(["%s" for _ in columns]) return f"INSERT INTO {format_identifier(table)} ({column_clause}) VALUES ({placeholders})" @@ -302,7 +302,11 @@ def can_use_local_infile_records(arrow_table: Any, records: "Sequence[Sequence[A all_safe_types = False if all_safe_types: return True - return not any(isinstance(value, (bytes, bytearray, memoryview, timedelta)) for row in records for value in row) + for row in records: + for value in row: + if isinstance(value, (bytes, bytearray, memoryview, timedelta)): + return False + return True def encode_records_for_local_infile(records: "list[tuple[Any, ...]]") -> bytes: @@ -355,7 +359,7 @@ def resolve_row_plan( column_names: list[str] = [] if not json_type_codes: - column_names.extend(column[0] for column in description) + column_names.extend([column[0] for column in description]) return column_names, None json_indexes: list[int] = [] @@ -520,8 +524,12 @@ def escape_literal_percent(sql: str, parameters: Any, validator: "ParameterValid """Escape literal percent characters before driver-side interpolation.""" if not parameters or "%" not in sql: return sql + if isinstance(parameters, Sized) and sql.count("%") == len(parameters): + return sql keep = {info.position for info in validator.extract_parameters(sql)} - return "".join("%%" if char == "%" and index not in keep else char for index, char in enumerate(sql)) + if len(keep) == sql.count("%"): + return sql + return "".join(["%%" if char == "%" and index not in keep else char for index, char in enumerate(sql)]) def build_load_data_statement( @@ -531,10 +539,10 @@ def build_load_data_statement( table_sql = format_identifier(table) if escape_percent: table_sql = table_sql.replace("%", "%%") - column_list = ", ".join( + column_list = ", ".join([ quote_backtick_identifier(column).replace("%", "%%") if escape_percent else quote_backtick_identifier(column) for column in columns - ) + ]) return ( f"LOAD DATA LOCAL INFILE {placeholder} INTO TABLE {table_sql} " "CHARACTER SET utf8mb4 FIELDS TERMINATED BY '\\t' ESCAPED BY '\\\\' " diff --git a/sqlspec/adapters/aiomysql/driver.py b/sqlspec/adapters/aiomysql/driver.py index 36bd446ad..f57a75112 100644 --- a/sqlspec/adapters/aiomysql/driver.py +++ b/sqlspec/adapters/aiomysql/driver.py @@ -376,10 +376,8 @@ async def load_from_arrow( ) if use_infile: payload = encode_records_for_local_infile(records) - fd, tmp_name = tempfile.mkstemp(suffix=".tsv") + tmp_name = await asyncio.to_thread(_write_local_infile_tempfile, payload) try: - with os.fdopen(fd, "wb") as tmp: - tmp.write(payload) load_sql = build_load_data_statement(table, columns, escape_percent=True) exc_handler = self.handle_database_exceptions() await self._run_with_exception_handler( @@ -469,6 +467,14 @@ def _connection_in_transaction(self) -> bool: return bool(self.connection.get_transaction_status()) +def _write_local_infile_tempfile(payload: bytes) -> str: + """Write LOCAL INFILE payload to a temporary file.""" + fd, tmp_name = tempfile.mkstemp(suffix=".tsv") + with os.fdopen(fd, "wb") as tmp: + tmp.write(payload) + return tmp_name + + def _resolve_column_types(description: Any) -> "dict[str, str] | None": """Map MySQL cursor column FIELD_TYPE codes to neutral Arrow type tokens. diff --git a/tests/unit/adapters/test_aiomysql/test_config.py b/tests/unit/adapters/test_aiomysql/test_config.py index 71925992a..612e9ff64 100644 --- a/tests/unit/adapters/test_aiomysql/test_config.py +++ b/tests/unit/adapters/test_aiomysql/test_config.py @@ -320,3 +320,10 @@ def test_aiomysql_config_with_dsn_applies_to_connection_kwargs() -> None: assert kwargs["password"] == "pass1" assert kwargs["host"] == "dbhost" assert kwargs["db"] == "production" + + +def test_aiomysql_session_factory_slots() -> None: + """_AiomysqlSessionFactory declares __slots__ for connection context tracking.""" + from sqlspec.adapters.aiomysql.config import _AiomysqlSessionFactory + + assert _AiomysqlSessionFactory.__slots__ == ("_contexts",) diff --git a/tests/unit/adapters/test_mysql_parameter_profiles.py b/tests/unit/adapters/test_mysql_parameter_profiles.py index 100e3b38c..8b9a37b0a 100644 --- a/tests/unit/adapters/test_mysql_parameter_profiles.py +++ b/tests/unit/adapters/test_mysql_parameter_profiles.py @@ -81,3 +81,12 @@ async def test_cached_literal_placeholder_is_escaped(adapter: str, class_name: s result = await result assert result == "result" execute.assert_called_once() + + +def test_escape_literal_percent_fast_path_returns_same_string_reference() -> None: + from sqlspec.adapters.aiomysql.core import default_statement_config, escape_literal_percent + + sql = "SELECT id, name FROM users WHERE id = %s AND status = %s" + params = (1, "active") + escaped = escape_literal_percent(sql, params, default_statement_config.parameter_validator) + assert escaped is sql From 86c8f0f7a1fc27f724b96dc6542d7d78a1db5ce8 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sat, 31 Oct 2026 22:52:37 +0000 Subject: [PATCH 020/174] refactor(aiomysql): use statement config in execution methods, delegate script execution, and clean up connection handling Pass statement.statement_config to statement compilation and parameter validation across execution methods. Delegate dispatch_execute_script to the base driver implementation to emit statement observability events, close standalone connections on creation failures, catch AiomysqlPymysqlError consistently in transaction methods, and remove the redundant load_from_storage override. --- sqlspec/adapters/aiomysql/config.py | 6 +- sqlspec/adapters/aiomysql/driver.py | 59 ++++---------- .../adapters/test_aiomysql/test_config.py | 21 +++++ .../adapters/test_aiomysql/test_driver.py | 76 +++++++++++++++++++ 4 files changed, 115 insertions(+), 47 deletions(-) diff --git a/sqlspec/adapters/aiomysql/config.py b/sqlspec/adapters/aiomysql/config.py index 31fa29daa..e5655fbdb 100644 --- a/sqlspec/adapters/aiomysql/config.py +++ b/sqlspec/adapters/aiomysql/config.py @@ -335,7 +335,11 @@ async def create_connection(self) -> AiomysqlConnection: An aiomysql connection instance. """ connection = await aiomysql.connect(**self._connection_kwargs()) - await self._ensure_connection(connection) + try: + await self._ensure_connection(connection) + except BaseException: + connection.close() + raise return connection async def provide_pool(self, *args: Any, **kwargs: Any) -> "AiomysqlPool": diff --git a/sqlspec/adapters/aiomysql/driver.py b/sqlspec/adapters/aiomysql/driver.py index f57a75112..e6734526d 100644 --- a/sqlspec/adapters/aiomysql/driver.py +++ b/sqlspec/adapters/aiomysql/driver.py @@ -19,7 +19,6 @@ AiomysqlCursor, AiomysqlFieldType, AiomysqlPymysqlError, - AiomysqlPymysqlMySQLError, AiomysqlSessionContext, ) from sqlspec.adapters.aiomysql.core import ( @@ -152,8 +151,8 @@ async def dispatch_execute(self, cursor: Any, statement: "SQL") -> "ExecutionRes Returns: ExecutionResult: Statement execution results with data or row counts """ - sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) - sql = escape_literal_percent(sql, prepared_parameters, self.statement_config.parameter_validator) + sql, prepared_parameters = self._compiled_sql(statement, statement.statement_config) + sql = escape_literal_percent(sql, prepared_parameters, statement.statement_config.parameter_validator) await cursor.execute(sql, normalize_execute_parameters(prepared_parameters)) description = getattr(cursor, "description", None) or None @@ -189,15 +188,17 @@ async def dispatch_execute_many(self, cursor: Any, statement: "SQL") -> "Executi Returns: ExecutionResult: Batch execution results """ - sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) + sql, prepared_parameters = self._compiled_sql(statement, statement.statement_config) match = AIOMYSQL_INSERT_VALUES_PATTERN.match(sql) if match: sql = ( - escape_literal_percent(match.group(1), prepared_parameters, self.statement_config.parameter_validator) + escape_literal_percent( + match.group(1), prepared_parameters, statement.statement_config.parameter_validator + ) + sql[match.end(1) :] ) else: - sql = escape_literal_percent(sql, prepared_parameters, self.statement_config.parameter_validator) + sql = escape_literal_percent(sql, prepared_parameters, statement.statement_config.parameter_validator) prepared_parameters = normalize_execute_many_parameters(prepared_parameters) parameter_count = len(prepared_parameters) if isinstance(prepared_parameters, Sized) else None @@ -217,25 +218,7 @@ async def dispatch_execute_script(self, cursor: Any, statement: "SQL") -> "Execu Returns: ExecutionResult: Script execution results with statement count """ - sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) - statements = self.split_script_statements(sql, statement.statement_config, strip_trailing_semicolon=True) - - successful_count = 0 - last_cursor = cursor - normalized_params = normalize_execute_parameters(prepared_parameters) - - for stmt in statements: - stmt_sql = ( - escape_literal_percent(stmt, prepared_parameters, self.statement_config.parameter_validator) - if normalized_params is not None - else stmt - ) - await cursor.execute(stmt_sql, normalized_params) - successful_count += 1 - - return self.create_execution_result( - last_cursor, statement_count=len(statements), successful_statements=successful_count, is_script_result=True - ) + return await super().dispatch_execute_script(cursor, statement) async def begin(self) -> None: """Begin a database transaction. @@ -245,7 +228,7 @@ async def begin(self) -> None: """ try: await self.connection.begin() - except AiomysqlPymysqlMySQLError as e: + except AiomysqlPymysqlError as e: raise create_mapped_exception(e) from e async def commit(self) -> None: @@ -256,7 +239,7 @@ async def commit(self) -> None: """ try: await self.connection.commit() - except AiomysqlPymysqlMySQLError as e: + except AiomysqlPymysqlError as e: raise create_mapped_exception(e) from e async def rollback(self) -> None: @@ -267,7 +250,7 @@ async def rollback(self) -> None: """ try: await self.connection.rollback() - except AiomysqlPymysqlMySQLError as e: + except AiomysqlPymysqlError as e: raise create_mapped_exception(e) from e def with_cursor(self, connection: "AiomysqlConnection") -> "AiomysqlCursor": @@ -285,8 +268,8 @@ def dispatch_select_stream(self, statement: "SQL", chunk_size: int) -> "AsyncRow """Return a native aiomysql row stream backed by an unbuffered ``SSCursor``.""" if not statement.returns_rows(): return None - sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) - sql = escape_literal_percent(sql, prepared_parameters, self.statement_config.parameter_validator) + sql, prepared_parameters = self._compiled_sql(statement, statement.statement_config) + sql = escape_literal_percent(sql, prepared_parameters, statement.statement_config.parameter_validator) return AsyncRowStream( AiomysqlStreamSource(self, sql, prepared_parameters, chunk_size, AIOMYSQL_JSON_TYPE_CODES) ) @@ -404,22 +387,6 @@ async def load_from_arrow( self._attach_partition_telemetry(telemetry_payload, partitioner) return self._storage_job(telemetry_payload, telemetry) - async def load_from_storage( - self, - table: str, - source: "StorageDestination", - *, - file_format: "StorageFormat", - partitioner: "dict[str, object] | None" = None, - overwrite: bool = False, - **kwargs: Any, - ) -> "StorageBridgeJob": - """Load staged artifacts from storage into MySQL.""" - - return await super().load_from_storage( - table, source, file_format=file_format, partitioner=partitioner, overwrite=overwrite, **kwargs - ) - def collect_rows(self, cursor: Any, fetched: "list[Any]") -> "tuple[list[Any], list[str], int]": """Collect aiomysql rows for the direct execution path.""" description = cursor.description or None diff --git a/tests/unit/adapters/test_aiomysql/test_config.py b/tests/unit/adapters/test_aiomysql/test_config.py index 612e9ff64..f3c4710ae 100644 --- a/tests/unit/adapters/test_aiomysql/test_config.py +++ b/tests/unit/adapters/test_aiomysql/test_config.py @@ -327,3 +327,24 @@ def test_aiomysql_session_factory_slots() -> None: from sqlspec.adapters.aiomysql.config import _AiomysqlSessionFactory assert _AiomysqlSessionFactory.__slots__ == ("_contexts",) + + +async def test_aiomysql_create_connection_closes_socket_on_error(monkeypatch: pytest.MonkeyPatch) -> None: + """create_connection closes the standalone connection if _ensure_connection raises.""" + from unittest.mock import AsyncMock, MagicMock + + import aiomysql + + mock_conn = MagicMock() + mock_conn.close = MagicMock() + monkeypatch.setattr(aiomysql, "connect", AsyncMock(return_value=mock_conn)) + + config = AiomysqlConfig( + connection_config={"host": "localhost", "db": "test"}, + driver_features={"on_connection_create": AsyncMock(side_effect=RuntimeError("hook error"))}, + ) + + with pytest.raises(RuntimeError, match="hook error"): + await config.create_connection() + + assert mock_conn.close.called diff --git a/tests/unit/adapters/test_aiomysql/test_driver.py b/tests/unit/adapters/test_aiomysql/test_driver.py index 30402407b..caefe1d50 100644 --- a/tests/unit/adapters/test_aiomysql/test_driver.py +++ b/tests/unit/adapters/test_aiomysql/test_driver.py @@ -139,3 +139,79 @@ def _spy_collect_rows(*args: Any, **kwargs: Any) -> Any: result = await driver._execute_cache_hit("SELECT 1", (), cached) assert bool(collect_called) is True assert result.data == [(1,)] + + +async def test_aiomysql_dispatch_execute_uses_statement_statement_config() -> None: + """dispatch_execute compiles using statement.statement_config.""" + from unittest.mock import AsyncMock, MagicMock + + from sqlspec import SQL + from sqlspec.core import StatementConfig + + custom_config = StatementConfig(enable_caching=False) + statement = SQL("SELECT 1", statement_config=custom_config) + cursor = MagicMock() + cursor.description = [("1", 3, None, None, None, None, None)] + cursor.execute = AsyncMock() + cursor.fetchall = AsyncMock(return_value=[(1,)]) + cursor.rowcount = 1 + cursor.lastrowid = None + + driver = AiomysqlDriver(connection=MagicMock()) + result = await driver.dispatch_execute(cursor, statement) + assert cursor.execute.called + assert result.data_row_count == 1 + + +async def test_aiomysql_dispatch_select_stream_uses_statement_config() -> None: + """dispatch_select_stream compiles using statement.statement_config.""" + from unittest.mock import MagicMock + + from sqlspec import SQL + from sqlspec.core import StatementConfig + + custom_config = StatementConfig(enable_caching=False) + statement = SQL("SELECT 1", statement_config=custom_config) + driver = AiomysqlDriver(connection=MagicMock()) + stream = driver.dispatch_select_stream(statement, chunk_size=10) + assert stream is not None + + +async def test_aiomysql_dispatch_execute_script_delegates_to_super() -> None: + """dispatch_execute_script emits script statement events via super().""" + from unittest.mock import AsyncMock, MagicMock + + from sqlspec import SQL + + cursor = MagicMock() + cursor.execute = AsyncMock() + cursor.description = None + cursor.rowcount = 0 + cursor.lastrowid = None + + driver = AiomysqlDriver(connection=MagicMock()) + statement = SQL("SELECT 1; SELECT 2").as_script() + result = await driver.dispatch_execute_script(cursor, statement) + assert result.statement_count == 2 + assert result.successful_statements == 2 + + +async def test_aiomysql_transaction_methods_catch_pymysql_error() -> None: + """begin, commit, and rollback catch AiomysqlPymysqlError and raise mapped exception.""" + from unittest.mock import AsyncMock, MagicMock + + from sqlspec.adapters.aiomysql._typing import AiomysqlPymysqlError + from sqlspec.exceptions import SQLSpecError + + conn = MagicMock() + conn.begin = AsyncMock(side_effect=AiomysqlPymysqlError(1045, "Access denied")) + conn.commit = AsyncMock(side_effect=AiomysqlPymysqlError(1045, "Access denied")) + conn.rollback = AsyncMock(side_effect=AiomysqlPymysqlError(1045, "Access denied")) + + driver = AiomysqlDriver(connection=conn) + with pytest.raises(SQLSpecError): + await driver.begin() + with pytest.raises(SQLSpecError): + await driver.commit() + with pytest.raises(SQLSpecError): + await driver.rollback() From ba149bc20ea4e92fd825df0c3629b8ecb688c445 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sat, 31 Oct 2026 22:52:37 +0000 Subject: [PATCH 021/174] fix(asyncmy): restore cursor fast path for clean cached query execution Delegate _execute_cache_hit to super()._execute_cache_hit when cached SQL queries do not require literal percent escaping. This eliminates unnecessary statement and result allocations and restores direct cursor execution via collect_rows and resolve_rowcount. --- sqlspec/adapters/asyncmy/driver.py | 16 +++++-- .../unit/adapters/test_asyncmy/test_driver.py | 42 +++++++++++++++++++ 2 files changed, 54 insertions(+), 4 deletions(-) diff --git a/sqlspec/adapters/asyncmy/driver.py b/sqlspec/adapters/asyncmy/driver.py index fd96c28ae..d5596e4e1 100644 --- a/sqlspec/adapters/asyncmy/driver.py +++ b/sqlspec/adapters/asyncmy/driver.py @@ -428,10 +428,18 @@ def resolve_rowcount(self, cursor: Any) -> int: async def _execute_cache_hit( self, sql: str, params: "tuple[Any, ...] | list[Any] | dict[str, Any]", cached: "CachedQuery" ) -> "SQLResult": - statement = self._cached_statement( - sql, params, cached, params, params_are_simple=True, compiled_sql=cached.compiled_sql - ) - return await self._execute_cached_statement(statement) + if ( + "%" in cached.compiled_sql + and escape_literal_percent( + cached.compiled_sql, params or (None,), self.statement_config.parameter_validator + ) + != cached.compiled_sql + ): + statement = self._cached_statement( + sql, params, cached, params, params_are_simple=True, compiled_sql=cached.compiled_sql + ) + return await self._execute_cached_statement(statement) + return await super()._execute_cache_hit(sql, params, cached) async def _execute_cursor_statement(self, statement: str, parameters: Any = None) -> None: async with self.with_cursor(self.connection) as cursor: diff --git a/tests/unit/adapters/test_asyncmy/test_driver.py b/tests/unit/adapters/test_asyncmy/test_driver.py index 454265359..72b3e2a55 100644 --- a/tests/unit/adapters/test_asyncmy/test_driver.py +++ b/tests/unit/adapters/test_asyncmy/test_driver.py @@ -98,3 +98,45 @@ def acquire(self) -> _AcquireCtx: assert cast("object", acquired) is conn assert conn.rolled_back is True assert pool.released == [conn] + + +async def test_asyncmy_execute_cache_hit_invokes_collect_rows_on_clean_query(monkeypatch: pytest.MonkeyPatch) -> None: + """Cached query without unescaped literal percent delegates to fast path and invokes collect_rows.""" + from unittest.mock import AsyncMock, MagicMock + + from sqlspec.core import OperationProfile, ParameterProfile, ProcessedState + from sqlspec.core.cache import CachedQuery + + cursor = MagicMock() + cursor.description = [("id", 3, None, None, None, None, None)] + cursor.fetchall = AsyncMock(return_value=[(1,)]) + cursor.execute = AsyncMock(return_value=None) + + conn = MagicMock() + conn.cursor = MagicMock(return_value=cursor) + + driver = AsyncmyDriver(connection=cast("Any", conn)) + collect_called: list[bool] = [] + original_collect_rows = AsyncmyDriver.collect_rows + + def _spy_collect_rows(*args: Any, **kwargs: Any) -> Any: + collect_called.append(True) + return original_collect_rows(*args, **kwargs) + + monkeypatch.setattr(AsyncmyDriver, "collect_rows", _spy_collect_rows) + + cached = CachedQuery( + compiled_sql="SELECT 1", + parameter_profile=ParameterProfile(), + input_named_parameters=(), + applied_wrap_types=False, + parameter_casts={}, + operation_type="SELECT", + operation_profile=OperationProfile(returns_rows=True, modifies_rows=False), + param_count=0, + processed_state=ProcessedState(compiled_sql="SELECT 1", execution_parameters=[], operation_type="SELECT"), + ) + + result = await driver._execute_cache_hit("SELECT 1", (), cached) + assert bool(collect_called) is True + assert result.data == [(1,)] From f41dca27522c133fd2c4795ba781d5c98fd61125 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sat, 31 Oct 2026 22:52:37 +0000 Subject: [PATCH 022/174] fix(asyncmy): batch ADK memory insertions, escape LIKE wildcards, and use cross-version upsert syntax Batch ADK memory entries with cursor.executemany rather than sequential execute calls in a loop. Escape %, _, and \\ wildcards in memory search with an explicit ESCAPE clause, use LONGTEXT for content_text in the memory table DDL, remove redundant DROP INDEX from Litestar session teardown, and use VALUES(col) in Litestar session upsert statements for cross-version MySQL and MariaDB compatibility. --- sqlspec/adapters/asyncmy/adk/store.py | 95 ++++++++++--------- sqlspec/adapters/asyncmy/litestar/store.py | 13 +-- .../adapters/test_asyncmy/test_adk_store.py | 81 +++++++++++++++- .../test_asyncmy/test_litestar_store.py | 35 +++++++ 4 files changed, 171 insertions(+), 53 deletions(-) diff --git a/sqlspec/adapters/asyncmy/adk/store.py b/sqlspec/adapters/asyncmy/adk/store.py index 3b335dc4b..c2ed507cd 100644 --- a/sqlspec/adapters/asyncmy/adk/store.py +++ b/sqlspec/adapters/asyncmy/adk/store.py @@ -5,7 +5,7 @@ from typing import TYPE_CHECKING, Any, Final, Literal, NotRequired, cast from sqlspec.adapters.asyncmy._typing import AsyncmyCursor, AsyncmyProgrammingError, AsyncmyRawCursor -from sqlspec.adapters.asyncmy.core import resolve_rowcount +from sqlspec.adapters.asyncmy.core import resolve_many_rowcount, resolve_rowcount from sqlspec.config import ADKConfig from sqlspec.extensions.adk import BaseAsyncADKStore, StoredEvent, StoredSession, normalize_session_list_options from sqlspec.extensions.adk.memory.store import BaseAsyncADKMemoryStore @@ -565,46 +565,52 @@ async def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: " ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) """ - async with self._config.provide_connection() as conn: - async with AsyncmyCursor(conn, cursor_class=AsyncmyRawCursor) as cursor: - for entry in entries: - content_json = self._json_serializer(entry["content_json"]) - meta = entry.get("metadata_json") - metadata_json = self._json_serializer(meta) if meta is not None else None - params: tuple[Any, ...] - if self._owner_id_column_name: - params = ( - entry["id"], - entry["session_id"], - entry["app_name"], - entry["user_id"], - entry.get("scope", "user"), - entry["event_id"], - entry.get("author"), - owner_id, - entry["timestamp"], - content_json, - entry["content_text"], - metadata_json, - entry["inserted_at"], - ) - else: - params = ( - entry["id"], - entry["session_id"], - entry["app_name"], - entry["user_id"], - entry.get("scope", "user"), - entry["event_id"], - entry.get("author"), - entry["timestamp"], - content_json, - entry["content_text"], - metadata_json, - entry["inserted_at"], - ) - await cursor.execute(sql, params) - inserted_count += resolve_rowcount(cursor) + params_list: list[tuple[Any, ...]] = [] + for entry in entries: + content_json = self._json_serializer(entry["content_json"]) + meta = entry.get("metadata_json") + metadata_json = self._json_serializer(meta) if meta is not None else None + if self._owner_id_column_name: + params_list.append(( + entry["id"], + entry["session_id"], + entry["app_name"], + entry["user_id"], + entry.get("scope", "user"), + entry["event_id"], + entry.get("author"), + owner_id, + entry["timestamp"], + content_json, + entry["content_text"], + metadata_json, + entry["inserted_at"], + )) + else: + params_list.append(( + entry["id"], + entry["session_id"], + entry["app_name"], + entry["user_id"], + entry.get("scope", "user"), + entry["event_id"], + entry.get("author"), + entry["timestamp"], + content_json, + entry["content_text"], + metadata_json, + entry["inserted_at"], + )) + + if not params_list: + return 0 + + async with ( + self._config.provide_connection() as conn, + AsyncmyCursor(conn, cursor_class=AsyncmyRawCursor) as cursor, + ): + await cursor.executemany(sql, params_list) + inserted_count = resolve_many_rowcount(cursor, params_list, fallback_count=len(params_list)) await conn.commit() return inserted_count @@ -637,13 +643,14 @@ async def search_entries( """ params = (*scope_params, query, limit_value) else: + escaped_query = query.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_") sql = f""" SELECT * FROM {self._memory_table} - WHERE {where_scope} AND content_text LIKE %s + WHERE {where_scope} AND content_text LIKE %s ESCAPE '\\\\' ORDER BY timestamp DESC LIMIT %s """ - params = (*scope_params, f"%{query}%", limit_value) + params = (*scope_params, f"%{escaped_query}%", limit_value) try: async with ( @@ -750,7 +757,7 @@ async def _memory_table_ddl(self) -> str: author VARCHAR(256){owner_id_line}, timestamp TIMESTAMP(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6), content_json JSON NOT NULL, - content_text TEXT NOT NULL, + content_text LONGTEXT NOT NULL, metadata_json JSON, inserted_at TIMESTAMP(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6), INDEX idx_{self._memory_table}_app_scope_user_time (app_name, scope, user_id, timestamp), diff --git a/sqlspec/adapters/asyncmy/litestar/store.py b/sqlspec/adapters/asyncmy/litestar/store.py index 0c0441c54..865664ee7 100644 --- a/sqlspec/adapters/asyncmy/litestar/store.py +++ b/sqlspec/adapters/asyncmy/litestar/store.py @@ -141,10 +141,10 @@ async def set(self, key: str, value: "str | bytes", expires_in: "int | timedelta sql = f""" INSERT INTO {self._table_name} (session_id, data, expires_at) - VALUES (%s, %s, %s) AS new + VALUES (%s, %s, %s) ON DUPLICATE KEY UPDATE - data = new.data, - expires_at = new.expires_at, + data = VALUES(data), + expires_at = VALUES(expires_at), updated_at = UTC_TIMESTAMP(6) """ @@ -304,12 +304,9 @@ def _drop_table_sql(self) -> "list[str]": """Get MySQL/MariaDB DROP TABLE SQL statements. Returns: - List of SQL statements to drop indexes and table. + List of SQL statements to drop table. """ - return [ - f"DROP INDEX idx_{self._table_name}_expires_at ON {self._table_name}", - f"DROP TABLE IF EXISTS {self._table_name}", - ] + return [f"DROP TABLE IF EXISTS {self._table_name}"] def _is_mysql_table_missing(exc: BaseException) -> bool: diff --git a/tests/unit/adapters/test_asyncmy/test_adk_store.py b/tests/unit/adapters/test_asyncmy/test_adk_store.py index e3b2d6759..d60897090 100644 --- a/tests/unit/adapters/test_asyncmy/test_adk_store.py +++ b/tests/unit/adapters/test_asyncmy/test_adk_store.py @@ -226,6 +226,7 @@ async def test_asyncmy_adk_memory_store_insert_and_search_json_handling() -> Non cursor = MagicMock() cursor.rowcount = 1 cursor.execute = AsyncMock() + cursor.executemany = AsyncMock() cursor.description = [ ("id",), ("session_id",), @@ -288,7 +289,7 @@ async def test_asyncmy_adk_memory_store_insert_and_search_json_handling() -> Non records = await store.search_entries("hello", "app", "user-1") assert inserted == 1 - insert_params = cursor.execute.call_args_list[0].args[1] + insert_params = cursor.executemany.call_args_list[0].args[1][0] assert insert_params[6] is None assert insert_params[11] is None assert len(records) == 1 @@ -296,6 +297,84 @@ async def test_asyncmy_adk_memory_store_insert_and_search_json_handling() -> Non assert records[0]["metadata_json"] == {"source": "unit"} +async def test_asyncmy_insert_memory_entries_batches_with_executemany() -> None: + """insert_memory_entries batches multiple entries in a single cursor.executemany call.""" + from datetime import UTC, datetime + from unittest.mock import AsyncMock + + now = datetime.now(tz=UTC) + cursor = MagicMock() + cursor.rowcount = 2 + cursor.executemany = AsyncMock() + conn = MagicMock() + conn.cursor = MagicMock(return_value=cursor) + conn.commit = AsyncMock() + conn.__aenter__ = AsyncMock(return_value=conn) + conn.__aexit__ = AsyncMock(return_value=None) + config = _mock_config() + config.provide_connection = lambda *_a, **_k: conn + store = AsyncmyADKMemoryStore(config) + + entries = [ + cast( + "Any", + { + "id": f"mem-{i}", + "session_id": "sess-1", + "app_name": "app", + "user_id": "user-1", + "event_id": f"evt-{i}", + "timestamp": now, + "content_json": {"idx": i}, + "content_text": f"text {i}", + "metadata_json": None, + "inserted_at": now, + }, + ) + for i in range(2) + ] + + inserted = await store.insert_memory_entries(entries) + + assert inserted == 2 + assert cursor.executemany.call_count == 1 + batch_params = cursor.executemany.call_args_list[0].args[1] + assert len(batch_params) == 2 + + +async def test_asyncmy_search_entries_escapes_like_wildcards() -> None: + """search_entries escapes %, _, and \\ wildcards and appends ESCAPE clause.""" + from unittest.mock import AsyncMock + + cursor = MagicMock() + cursor.fetchall = AsyncMock(return_value=[]) + cursor.description = [] + cursor.execute = AsyncMock() + conn = MagicMock() + conn.cursor = MagicMock(return_value=cursor) + conn.__aenter__ = AsyncMock(return_value=conn) + conn.__aexit__ = AsyncMock(return_value=None) + config = _mock_config() + config.provide_connection = lambda *_a, **_k: conn + store = AsyncmyADKMemoryStore(config) + + await store.search_entries("100%_val\\id", "app", "user-1") + + assert cursor.execute.call_count == 1 + executed_sql = cursor.execute.call_args_list[0].args[0] + executed_params = cursor.execute.call_args_list[0].args[1] + assert "LIKE %s ESCAPE '\\\\'" in executed_sql + assert "%100\\%\\_val\\\\id%" in executed_params + + +async def test_asyncmy_memory_table_ddl_uses_longtext() -> None: + """_memory_table_ddl uses LONGTEXT for content_text.""" + config = _mock_config() + store = AsyncmyADKMemoryStore(config) + ddl = await store._memory_table_ddl() + assert "content_text LONGTEXT NOT NULL" in ddl + + async def test_asyncmy_stream_source_closes_cursor_when_execute_raises() -> None: """AsyncmyStreamSource.start() closes the cursor if execute() fails.""" from unittest.mock import AsyncMock diff --git a/tests/unit/adapters/test_asyncmy/test_litestar_store.py b/tests/unit/adapters/test_asyncmy/test_litestar_store.py index f47c370de..b44d9a91d 100644 --- a/tests/unit/adapters/test_asyncmy/test_litestar_store.py +++ b/tests/unit/adapters/test_asyncmy/test_litestar_store.py @@ -30,3 +30,38 @@ def test_asyncmy_litestar_store_honors_table_options() -> None: ddl = store._table_ddl() assert f"{DEFAULT_ENGINE_CLAUSE} COMMENT='litestar-session'" in ddl + + +def test_asyncmy_litestar_store_drop_table_sql_omits_drop_index() -> None: + """_drop_table_sql must emit only DROP TABLE IF EXISTS without preceding DROP INDEX.""" + store = AsyncmyStore(_mock_config()) + + statements = store._drop_table_sql() + + assert len(statements) == 1 + assert statements[0] == "DROP TABLE IF EXISTS litestar_session" + + +async def test_asyncmy_litestar_store_set_uses_mariadb_compatible_upsert() -> None: + """_set must use VALUES(col) in ON DUPLICATE KEY UPDATE for MySQL and MariaDB compatibility.""" + from unittest.mock import AsyncMock + + cursor = MagicMock() + cursor.execute = AsyncMock() + conn = MagicMock() + conn.cursor = MagicMock(return_value=cursor) + conn.commit = AsyncMock() + conn.__aenter__ = AsyncMock(return_value=conn) + conn.__aexit__ = AsyncMock(return_value=None) + config = _mock_config() + config.provide_connection = lambda *_a, **_k: conn + + store = AsyncmyStore(config) + await store.set("session-1", b"payload", 300) + + assert cursor.execute.call_count == 1 + sql = cursor.execute.call_args_list[0].args[0] + assert "ON DUPLICATE KEY UPDATE" in sql + assert "data = VALUES(data)" in sql + assert "expires_at = VALUES(expires_at)" in sql + assert "AS new" not in sql From da7e2ec9e8bce07e7bf2767cb606cba2582bb133 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sat, 31 Oct 2026 22:52:38 +0000 Subject: [PATCH 023/174] refactor(data_dictionary): combine MySQL DDL setup scalar queries and add object_name filter to privileges Combine VERSION(), @@sql_mode, and @@sql_quote_show_create scalar queries into a single SELECT statement and cache the resulting DDL context per driver session to minimize round-trips before SHOW CREATE. Add object_name filtering support to the privileges domain query in MySQL data dictionaries. --- .../dialects/mysql/dictionary.py | 96 +++++++++++++++---- .../dialects/mysql/sql/privileges.sql | 2 + .../data_dictionary/test_mysql_metadata.py | 35 +++++++ 3 files changed, 117 insertions(+), 16 deletions(-) diff --git a/sqlspec/data_dictionary/dialects/mysql/dictionary.py b/sqlspec/data_dictionary/dialects/mysql/dictionary.py index 207129d2c..ec2bbea69 100644 --- a/sqlspec/data_dictionary/dialects/mysql/dictionary.py +++ b/sqlspec/data_dictionary/dialects/mysql/dictionary.py @@ -56,6 +56,7 @@ class MySQLAsyncDataDictionary(AsyncDataDictionaryBase): def __init__(self) -> None: super().__init__() self._engine_version_cache: dict[int, MySQLEngineVersion | None] = {} + self._ddl_context_cache: dict[int, tuple[str | None, str | None, str | None]] = {} async def get_version(self, driver: "AsyncDriverAdapterBase") -> "VersionInfo | None": """Get MySQL database version information.""" @@ -138,7 +139,9 @@ async def get_privileges( """Get privilege metadata from INFORMATION_SCHEMA.""" dialect = await self._get_metadata_dialect(driver) schema_name = self._resolve_metadata_schema(schema) - return await self._select_domain_result(driver, dialect, "privileges", "by_schema", schema_name=schema_name) + return await self._select_domain_result( + driver, dialect, "privileges", "by_schema", schema_name=schema_name, object_name=object_name + ) async def get_system_metadata( self, driver: "AsyncDriverAdapterBase", request: "SystemMetadataRequest | str | None" = None, **kwargs: Any @@ -250,9 +253,35 @@ async def _get_native_object_ddl( _ = include_dependencies, redact if object_type.strip().upper() not in MYSQL_SHOW_CREATE_OBJECT_TYPES: return None - raw_version = await driver.select_value_or_none(self.get_query_text("version", "current")) - sql_mode = await driver.select_value_or_none("SELECT @@sql_mode") - sql_quote_show_create = await driver.select_value_or_none("SELECT @@sql_quote_show_create") + cache_key = _get_driver_cache_key(driver) + ddl_ctx = self._ddl_context_cache.get(cache_key) + if ddl_ctx is None: + try: + row_ctx = await driver.select_one( + "SELECT VERSION() AS raw_version, @@sql_mode AS sql_mode, @@sql_quote_show_create AS sql_quote_show_create", + schema_type=dict, + ) + ddl_ctx = ( + str(row_ctx.get("raw_version")) if row_ctx.get("raw_version") is not None else None, + str(row_ctx.get("sql_mode")) if row_ctx.get("sql_mode") is not None else None, + str(row_ctx.get("sql_quote_show_create")) + if row_ctx.get("sql_quote_show_create") is not None + else None, + ) + except Exception: + raw_version = await driver.select_value_or_none(self.get_query_text("version", "current")) + sql_mode = await driver.select_value_or_none("SELECT @@sql_mode") + sql_quote_show_create = await driver.select_value_or_none("SELECT @@sql_quote_show_create") + ddl_ctx = ( + str(raw_version) if raw_version is not None else None, + str(sql_mode) if sql_mode is not None else None, + str(sql_quote_show_create) if sql_quote_show_create is not None else None, + ) + if len(self._ddl_context_cache) >= _MAX_ENGINE_VERSION_CACHE_SIZE: + self._ddl_context_cache.pop(next(iter(self._ddl_context_cache))) + self._ddl_context_cache[cache_key] = ddl_ctx + + raw_version, sql_mode, sql_quote_show_create = ddl_ctx statement = build_mysql_show_create_statement(object_name, schema, object_type) try: row = await driver.select_one(statement, schema_type=dict) @@ -317,15 +346,18 @@ async def _select_domain_result( *, schema_name: "str | None", table_name: "str | None" = None, + object_name: "str | None" = None, ) -> "MetadataResult": query = get_data_dictionary_loader().get_domain_query(dialect, domain, query_name) if not query.is_supported or query.sql is None: return MetadataResult.unsupported(domain) capability = build_mysql_metadata_capability_profile(dialect, type(self).__name__).get(domain) - if table_name is None: - rows = await driver.select(query.sql, schema_name=schema_name) - else: - rows = await driver.select(query.sql, schema_name=schema_name, table_name=table_name) + params: dict[str, Any] = {"schema_name": schema_name} + if table_name is not None: + params["table_name"] = table_name + if object_name is not None: + params["object_name"] = object_name + rows = await driver.select(query.sql, **params) return MetadataResult(domain, capability=capability, items=tuple(rows)) @@ -338,6 +370,7 @@ class MySQLSyncDataDictionary(SyncDataDictionaryBase): def __init__(self) -> None: super().__init__() self._engine_version_cache: dict[int, MySQLEngineVersion | None] = {} + self._ddl_context_cache: dict[int, tuple[str | None, str | None, str | None]] = {} def get_version(self, driver: "SyncDriverAdapterBase") -> "VersionInfo | None": """Get MySQL database version information.""" @@ -420,7 +453,9 @@ def get_privileges( """Get privilege metadata from INFORMATION_SCHEMA.""" dialect = self._get_metadata_dialect(driver) schema_name = self._resolve_metadata_schema(schema) - return self._select_domain_result(driver, dialect, "privileges", "by_schema", schema_name=schema_name) + return self._select_domain_result( + driver, dialect, "privileges", "by_schema", schema_name=schema_name, object_name=object_name + ) def get_system_metadata( self, driver: "SyncDriverAdapterBase", request: "SystemMetadataRequest | str | None" = None, **kwargs: Any @@ -532,9 +567,35 @@ def _get_native_object_ddl( _ = include_dependencies, redact if object_type.strip().upper() not in MYSQL_SHOW_CREATE_OBJECT_TYPES: return None - raw_version = driver.select_value_or_none(self.get_query_text("version", "current")) - sql_mode = driver.select_value_or_none("SELECT @@sql_mode") - sql_quote_show_create = driver.select_value_or_none("SELECT @@sql_quote_show_create") + cache_key = _get_driver_cache_key(driver) + ddl_ctx = self._ddl_context_cache.get(cache_key) + if ddl_ctx is None: + try: + row_ctx = driver.select_one( + "SELECT VERSION() AS raw_version, @@sql_mode AS sql_mode, @@sql_quote_show_create AS sql_quote_show_create", + schema_type=dict, + ) + ddl_ctx = ( + str(row_ctx.get("raw_version")) if row_ctx.get("raw_version") is not None else None, + str(row_ctx.get("sql_mode")) if row_ctx.get("sql_mode") is not None else None, + str(row_ctx.get("sql_quote_show_create")) + if row_ctx.get("sql_quote_show_create") is not None + else None, + ) + except Exception: + raw_version = driver.select_value_or_none(self.get_query_text("version", "current")) + sql_mode = driver.select_value_or_none("SELECT @@sql_mode") + sql_quote_show_create = driver.select_value_or_none("SELECT @@sql_quote_show_create") + ddl_ctx = ( + str(raw_version) if raw_version is not None else None, + str(sql_mode) if sql_mode is not None else None, + str(sql_quote_show_create) if sql_quote_show_create is not None else None, + ) + if len(self._ddl_context_cache) >= _MAX_ENGINE_VERSION_CACHE_SIZE: + self._ddl_context_cache.pop(next(iter(self._ddl_context_cache))) + self._ddl_context_cache[cache_key] = ddl_ctx + + raw_version, sql_mode, sql_quote_show_create = ddl_ctx statement = build_mysql_show_create_statement(object_name, schema, object_type) try: row = driver.select_one(statement, schema_type=dict) @@ -599,15 +660,18 @@ def _select_domain_result( *, schema_name: "str | None", table_name: "str | None" = None, + object_name: "str | None" = None, ) -> "MetadataResult": query = get_data_dictionary_loader().get_domain_query(dialect, domain, query_name) if not query.is_supported or query.sql is None: return MetadataResult.unsupported(domain) capability = build_mysql_metadata_capability_profile(dialect, type(self).__name__).get(domain) - if table_name is None: - rows = driver.select(query.sql, schema_name=schema_name) - else: - rows = driver.select(query.sql, schema_name=schema_name, table_name=table_name) + params: dict[str, Any] = {"schema_name": schema_name} + if table_name is not None: + params["table_name"] = table_name + if object_name is not None: + params["object_name"] = object_name + rows = driver.select(query.sql, **params) return MetadataResult(domain, capability=capability, items=tuple(rows)) diff --git a/sqlspec/data_dictionary/dialects/mysql/sql/privileges.sql b/sqlspec/data_dictionary/dialects/mysql/sql/privileges.sql index c7900fa34..1002dab00 100644 --- a/sqlspec/data_dictionary/dialects/mysql/sql/privileges.sql +++ b/sqlspec/data_dictionary/dialects/mysql/sql/privileges.sql @@ -10,6 +10,7 @@ SELECT NULL AS column_name FROM information_schema.table_privileges WHERE table_schema = COALESCE(:schema_name, DATABASE()) + AND (:object_name IS NULL OR table_name = :object_name) UNION ALL SELECT grantee, @@ -21,4 +22,5 @@ SELECT column_name FROM information_schema.column_privileges WHERE table_schema = COALESCE(:schema_name, DATABASE()) + AND (:object_name IS NULL OR table_name = :object_name) ORDER BY schema_name, table_name, column_name, grantee, privilege_type; diff --git a/tests/unit/data_dictionary/test_mysql_metadata.py b/tests/unit/data_dictionary/test_mysql_metadata.py index 0ec27bf81..f1a01bffc 100644 --- a/tests/unit/data_dictionary/test_mysql_metadata.py +++ b/tests/unit/data_dictionary/test_mysql_metadata.py @@ -31,6 +31,13 @@ def __init__(self) -> None: def select_one(self, statement: str, *args: Any, **kwargs: Any) -> dict[str, str]: self.statements.append(str(statement)) + query = str(statement) + if "SELECT VERSION()" in query: + return { + "raw_version": "8.4.0", + "sql_mode": "STRICT_TRANS_TABLES,NO_ENGINE_SUBSTITUTION", + "sql_quote_show_create": "1", + } return {"Table": "orders", "Create Table": "CREATE TABLE `orders` (`id` int)"} def select_value_or_none(self, statement: object, *args: Any, **kwargs: Any) -> str | None: @@ -259,3 +266,31 @@ def select(self, statement: str, *args: Any, **kwargs: Any) -> list[dict[str, ob synthesized = dictionary.get_object_ddl(cast(Any, synth_driver), "orders", schema="shop", prefer_native=False) assert synthesized.status == MetadataSupport.UNSUPPORTED assert not any(statement.startswith("SHOW CREATE") for statement in synth_driver.statements) + + +def test_mysql_native_ddl_single_roundtrip_setup_and_cache() -> None: + """_get_native_object_ddl combines scalar queries into one round-trip and caches per session.""" + driver = _SyncDDLDriver() + dictionary = PyMysqlDataDictionary() + + result1 = dictionary.get_ddl(cast(Any, driver), "orders", schema="shop") + assert isinstance(result1, DDLResult) + setup_queries_count = sum(1 for s in driver.statements if "SELECT VERSION()" in s) + assert setup_queries_count == 1 + + result2 = dictionary.get_ddl(cast(Any, driver), "orders", schema="shop") + assert isinstance(result2, DDLResult) + setup_queries_count_after = sum(1 for s in driver.statements if "SELECT VERSION()" in s) + assert setup_queries_count_after == 1 + + +def test_mysql_get_privileges_filters_by_object_name() -> None: + """get_privileges passes object_name to domain query.""" + driver = _SyncDDLDriver() + dictionary = PyMysqlDataDictionary() + + dictionary.get_privileges(cast(Any, driver), object_name="orders", schema="shop") + privilege_queries = [ + s for s in driver.statements if "information_schema.table_privileges" in s or "table_schema" in s + ] + assert len(privilege_queries) > 0 From d5d1af5d2d6aea3eb5b7756541429db402975810 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sat, 31 Oct 2026 22:52:38 +0000 Subject: [PATCH 024/174] perf(asyncmy): offload local infile payload write to worker thread and eliminate redundant dict copy Offload LOCAL INFILE temporary file creation and payload writing to asyncio.to_thread in _load_from_arrow_via_local_infile, ensuring asynchronous cleanup with shutil.rmtree on both success and failure. Eliminate redundant dict copy of driver_features in AsyncmyConfig.__init__. --- sqlspec/adapters/asyncmy/config.py | 3 +-- sqlspec/adapters/asyncmy/driver.py | 24 +++++++++++++++---- .../unit/adapters/test_asyncmy/test_config.py | 7 ++++++ 3 files changed, 28 insertions(+), 6 deletions(-) diff --git a/sqlspec/adapters/asyncmy/config.py b/sqlspec/adapters/asyncmy/config.py index 69dfd4052..b02d2056d 100644 --- a/sqlspec/adapters/asyncmy/config.py +++ b/sqlspec/adapters/asyncmy/config.py @@ -320,8 +320,7 @@ def __init__( connection_config = build_connection_config(connection_config) statement_config = statement_config or default_statement_config - statement_config, driver_features = apply_driver_features(statement_config, driver_features) - features_dict = dict(driver_features) if driver_features else {} + statement_config, features_dict = apply_driver_features(statement_config, driver_features) self._user_connection_hook: Callable[[AsyncmyConnection], Awaitable[None]] | None = features_dict.pop( "on_connection_create", None ) diff --git a/sqlspec/adapters/asyncmy/driver.py b/sqlspec/adapters/asyncmy/driver.py index d5596e4e1..0ed595fb3 100644 --- a/sqlspec/adapters/asyncmy/driver.py +++ b/sqlspec/adapters/asyncmy/driver.py @@ -4,6 +4,8 @@ type coercion, error handling, and transaction management. """ +import asyncio +import shutil import tempfile from collections.abc import Sized from typing import TYPE_CHECKING, Any, Final, Literal, cast @@ -460,14 +462,14 @@ async def _execute_local_infile_cursor(self, statement: str, filename: str) -> N async def _load_from_arrow_via_local_infile( self, table: str, columns: "list[str]", records: "list[tuple[Any, ...]]" ) -> None: - with tempfile.TemporaryDirectory(prefix="sqlspec-asyncmy-") as directory: - with tempfile.NamedTemporaryFile(dir=directory, suffix=".tsv", delete=False) as payload: - payload.write(encode_records_for_local_infile(records)) - filename = payload.name + directory, filename = await asyncio.to_thread(_write_asyncmy_local_infile_payload, records) + try: statement = build_load_data_statement(table, columns) exc_handler = self.handle_database_exceptions() await self._run_with_exception_handler(exc_handler, self._execute_local_infile_cursor, statement, filename) self._check_pending_exception(exc_handler) + finally: + await asyncio.to_thread(shutil.rmtree, directory, ignore_errors=True) def _connection_in_transaction(self) -> bool: """Check if connection is in transaction. @@ -478,6 +480,20 @@ def _connection_in_transaction(self) -> bool: return bool(getattr(self.connection, "server_status", 0) & 1) +def _write_asyncmy_local_infile_payload(records: "list[tuple[Any, ...]]") -> tuple[str, str]: + """Write LOCAL INFILE payload to a temporary file in a dedicated directory.""" + directory = tempfile.mkdtemp(prefix="sqlspec-asyncmy-") + try: + with tempfile.NamedTemporaryFile(dir=directory, suffix=".tsv", delete=False) as payload: + payload.write(encode_records_for_local_infile(records)) + filename = payload.name + except BaseException: + shutil.rmtree(directory, ignore_errors=True) + raise + else: + return directory, filename + + def _resolve_column_types(description: Any) -> "dict[str, str] | None": """Map MySQL cursor column FIELD_TYPE codes to neutral Arrow type tokens. diff --git a/tests/unit/adapters/test_asyncmy/test_config.py b/tests/unit/adapters/test_asyncmy/test_config.py index 5d091f1d2..18472366d 100644 --- a/tests/unit/adapters/test_asyncmy/test_config.py +++ b/tests/unit/adapters/test_asyncmy/test_config.py @@ -329,3 +329,10 @@ def test_asyncmy_config_with_dsn_applies_to_connection_kwargs() -> None: assert conn_kwargs["password"] == "pass1" assert conn_kwargs["host"] == "dbhost" assert conn_kwargs["database"] == "production" + + +def test_asyncmy_config_driver_features_initialization() -> None: + """AsyncmyConfig initializes driver_features without redundant copying.""" + config = AsyncmyConfig(connection_config={"local_infile": True}, driver_features={"json_serializer": str}) + assert config.driver_features["json_serializer"] is str + assert config.driver_features["enable_local_infile_bulk_load"] is True From 832998b387a733113d75ef2183bca6ca36f62936 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sat, 31 Oct 2026 22:52:39 +0000 Subject: [PATCH 025/174] refactor(asyncmy): improve exception mapping, guard optional driver imports, and streamline execution hooks Add explicit MySQL errno mappings for 1054, 1146, 1040, 1290, and 1836 to handle asyncmy errors without SQLSTATE, remove Postgres 235xx codes, and close connection on SSCursor close error. Guard top-level asyncmy imports, relocate AsyncmySessionContext to config.py, set native Arrow/Parquet capabilities accurately to False, delegate script execution to base driver, and close standalone connections on creation failure. --- sqlspec/adapters/asyncmy/_typing.py | 133 +++++++----------- sqlspec/adapters/asyncmy/config.py | 87 ++++++++++-- sqlspec/adapters/asyncmy/core.py | 23 ++- sqlspec/adapters/asyncmy/driver.py | 63 ++------- .../unit/adapters/test_asyncmy/test_config.py | 52 +++++++ .../unit/adapters/test_asyncmy/test_driver.py | 92 ++++++++++++ 6 files changed, 299 insertions(+), 151 deletions(-) diff --git a/sqlspec/adapters/asyncmy/_typing.py b/sqlspec/adapters/asyncmy/_typing.py index 0ebf568b0..3807eafbc 100644 --- a/sqlspec/adapters/asyncmy/_typing.py +++ b/sqlspec/adapters/asyncmy/_typing.py @@ -8,30 +8,61 @@ import os from typing import TYPE_CHECKING, Any -import asyncmy as _asyncmy -from asyncmy import Connection -from asyncmy import errors as _asyncmy_errors -from asyncmy.connection import LoadLocalFile -from asyncmy.connection import MySQLResult as _AsyncmyResult -from asyncmy.constants import FIELD_TYPE as _ASYNCMY_FIELD_TYPE -from asyncmy.cursors import RE_INSERT_VALUES as ASYNCMY_INSERT_VALUES_PATTERN -from asyncmy.cursors import Cursor as _AsyncmyCursor -from asyncmy.cursors import DictCursor as _AsyncmyDictCursor -from asyncmy.cursors import SSCursor as AsyncmySSCursor -from asyncmy.errors import ProgrammingError as _AsyncmyProgrammingError -from asyncmy.pool import Pool as _AsyncmyPool -from asyncmy.protocol import LoadLocalPacketWrapper as _LoadLocalPacketWrapper +if TYPE_CHECKING: + import asyncmy as _asyncmy + from asyncmy import Connection + from asyncmy import errors as _asyncmy_errors + from asyncmy.connection import LoadLocalFile + from asyncmy.connection import MySQLResult as _AsyncmyResult + from asyncmy.constants import FIELD_TYPE as _ASYNCMY_FIELD_TYPE + from asyncmy.cursors import RE_INSERT_VALUES as ASYNCMY_INSERT_VALUES_PATTERN + from asyncmy.cursors import Cursor as _AsyncmyCursor + from asyncmy.cursors import DictCursor as _AsyncmyDictCursor + from asyncmy.cursors import SSCursor as AsyncmySSCursor + from asyncmy.errors import ProgrammingError as _AsyncmyProgrammingError + from asyncmy.pool import Pool as _AsyncmyPool + from asyncmy.protocol import LoadLocalPacketWrapper as _LoadLocalPacketWrapper +else: + try: + import asyncmy as _asyncmy + from asyncmy import Connection + from asyncmy import errors as _asyncmy_errors + from asyncmy.connection import LoadLocalFile + from asyncmy.connection import MySQLResult as _AsyncmyResult + from asyncmy.constants import FIELD_TYPE as _ASYNCMY_FIELD_TYPE + from asyncmy.cursors import RE_INSERT_VALUES as ASYNCMY_INSERT_VALUES_PATTERN + from asyncmy.cursors import Cursor as _AsyncmyCursor + from asyncmy.cursors import DictCursor as _AsyncmyDictCursor + from asyncmy.cursors import SSCursor as AsyncmySSCursor + from asyncmy.errors import ProgrammingError as _AsyncmyProgrammingError + from asyncmy.pool import Pool as _AsyncmyPool + from asyncmy.protocol import LoadLocalPacketWrapper as _LoadLocalPacketWrapper + except ImportError: + import re + + _asyncmy = None + Connection = None + _asyncmy_errors = None + LoadLocalFile = None + _AsyncmyResult = object + _ASYNCMY_FIELD_TYPE = None + ASYNCMY_INSERT_VALUES_PATTERN = re.compile( + r"(VALUES\s*\()((?:(?!\bVALUES\b).)*?)(\)(?:\s*(?:ON\s+DUPLICATE\s+KEY\s+UPDATE|RETURNING)\b.*)?\Z", + re.IGNORECASE | re.DOTALL, + ) + _AsyncmyCursor = None + _AsyncmyDictCursor = None + AsyncmySSCursor = None + _AsyncmyProgrammingError = Exception + _AsyncmyPool = None + _LoadLocalPacketWrapper = None from sqlspec.exceptions import SQLSpecError if TYPE_CHECKING: - from collections.abc import Awaitable, Callable, Iterator - from types import TracebackType + from collections.abc import Iterator from typing import Protocol, TypeAlias - from sqlspec.adapters.asyncmy.driver import AsyncmyDriver - from sqlspec.core import StatementConfig - class AsyncmyMySQLError(Exception): errno: int @@ -93,12 +124,9 @@ class AsyncmyFieldTypeProtocol(Protocol): "AsyncmyProgrammingError", "AsyncmyRawCursor", "AsyncmySSCursor", - "AsyncmySessionContext", "asyncmy_local_infile", ) -_resolved_driver_type: "type[AsyncmyDriver] | None" = None - class AsyncmyCursor: """Context manager for AsyncMy cursor operations. @@ -133,60 +161,6 @@ async def __aexit__(self, *_: object) -> None: await raw_cursor.__aexit__(None, None, None) -class AsyncmySessionContext: - """Async context manager for AsyncMy sessions. - - This class is intentionally excluded from mypyc compilation to avoid ABI - boundary issues. It receives callables from uncompiled config classes and - instantiates compiled Driver objects, acting as a bridge between compiled - and uncompiled code. - - Uses callable-based connection management to decouple from config implementation. - """ - - __slots__ = ( - "_acquire_connection", - "_connection", - "_driver", - "_driver_features", - "_prepare_driver", - "_release_connection", - "_statement_config", - ) - - def __init__( - self, - acquire_connection: "Callable[[], Awaitable[AsyncmyConnection]]", - release_connection: "Callable[..., Any]", - statement_config: "StatementConfig", - driver_features: "dict[str, Any]", - prepare_driver: "Callable[[AsyncmyDriver], AsyncmyDriver]", - ) -> None: - self._acquire_connection = acquire_connection - self._release_connection = release_connection - self._statement_config = statement_config - self._driver_features = driver_features - self._prepare_driver = prepare_driver - self._connection: AsyncmyConnection | None = None - self._driver: AsyncmyDriver | None = None - - async def __aenter__(self) -> "AsyncmyDriver": - driver_cls = _get_asyncmy_driver_type() - self._connection = await self._acquire_connection() - self._driver = driver_cls( - connection=self._connection, statement_config=self._statement_config, driver_features=self._driver_features - ) - return self._prepare_driver(self._driver) - - async def __aexit__( - self, exc_type: "type[BaseException] | None", exc_val: "BaseException | None", exc_tb: "TracebackType | None" - ) -> "bool | None": - if self._connection is not None: - await self._release_connection(self._connection, exc_type=exc_type, exc_val=exc_val, exc_tb=exc_tb) - self._connection = None - return None - - @contextlib.contextmanager def asyncmy_local_infile(connection: "AsyncmyConnection", filename: str) -> "Iterator[None]": """Scope the asyncmy 0.2.13/0.2.14 filename handoff fix to one native load. @@ -256,12 +230,3 @@ async def _read_load_local_packet(self, first_packet: Any) -> None: raise SQLSpecError(msg) raw_self: Any = self raw_self._read_ok_packet(packet) - - -def _get_asyncmy_driver_type() -> "type[AsyncmyDriver]": - global _resolved_driver_type - if _resolved_driver_type is None: - from sqlspec.adapters.asyncmy.driver import AsyncmyDriver as _Driver - - _resolved_driver_type = _Driver - return _resolved_driver_type diff --git a/sqlspec/adapters/asyncmy/config.py b/sqlspec/adapters/asyncmy/config.py index b02d2056d..934e13bb0 100644 --- a/sqlspec/adapters/asyncmy/config.py +++ b/sqlspec/adapters/asyncmy/config.py @@ -14,7 +14,6 @@ AsyncmyModule, AsyncmyPool, AsyncmyRawCursor, - AsyncmySessionContext, ) from sqlspec.adapters.asyncmy.core import apply_driver_features, default_statement_config from sqlspec.adapters.asyncmy.driver import AsyncmyDriver, AsyncmyExceptionHandler @@ -40,6 +39,7 @@ "AsyncmyDriverFeatures", "AsyncmyPoolParams", "AsyncmySSLParams", + "AsyncmySessionContext", "build_connection_config", ) @@ -222,9 +222,13 @@ async def acquire_connection(self) -> "AsyncmyConnection": async def release_connection(self, _conn: "AsyncmyConnection", **kwargs: Any) -> None: ctx = self._contexts.pop(id(_conn), None) if ctx is not None: - if _conn is not None and bool(getattr(_conn, "server_status", 0) & 1): - with contextlib.suppress(Exception): - await _conn.rollback() + if _conn is not None: + in_transaction = False + with contextlib.suppress(AttributeError): + in_transaction = bool(_conn.server_status & 1) + if in_transaction: + with contextlib.suppress(Exception): + await _conn.rollback() await ctx.__aexit__(kwargs.get("exc_type"), kwargs.get("exc_val"), kwargs.get("exc_tb")) @@ -259,9 +263,13 @@ async def __aexit__( ) -> bool | None: conn = self._connection self._connection = None - if conn is not None and bool(getattr(conn, "server_status", 0) & 1): - with contextlib.suppress(Exception): - await conn.rollback() + if conn is not None: + in_transaction = False + with contextlib.suppress(AttributeError): + in_transaction = bool(conn.server_status & 1) + if in_transaction: + with contextlib.suppress(Exception): + await conn.rollback() if self._ctx: ctx = self._ctx self._ctx = None @@ -269,6 +277,57 @@ async def __aexit__( return None +class AsyncmySessionContext: + """Async context manager for AsyncMy sessions. + + This class is intentionally excluded from mypyc compilation to avoid ABI + boundary issues. It receives callables from uncompiled config classes and + instantiates compiled Driver objects, acting as a bridge between compiled + and uncompiled code. + """ + + __slots__ = ( + "_acquire_connection", + "_connection", + "_driver", + "_driver_features", + "_prepare_driver", + "_release_connection", + "_statement_config", + ) + + def __init__( + self, + acquire_connection: "Callable[[], Awaitable[AsyncmyConnection]]", + release_connection: "Callable[..., Any]", + statement_config: "StatementConfig", + driver_features: "dict[str, Any]", + prepare_driver: "Callable[[AsyncmyDriver], AsyncmyDriver]", + ) -> None: + self._acquire_connection = acquire_connection + self._release_connection = release_connection + self._statement_config = statement_config + self._driver_features = driver_features + self._prepare_driver = prepare_driver + self._connection: AsyncmyConnection | None = None + self._driver: AsyncmyDriver | None = None + + async def __aenter__(self) -> "AsyncmyDriver": + self._connection = await self._acquire_connection() + self._driver = AsyncmyDriver( + connection=self._connection, statement_config=self._statement_config, driver_features=self._driver_features + ) + return self._prepare_driver(self._driver) + + async def __aexit__( + self, exc_type: "type[BaseException] | None", exc_val: "BaseException | None", exc_tb: "TracebackType | None" + ) -> "bool | None": + if self._connection is not None: + await self._release_connection(self._connection, exc_type=exc_type, exc_val=exc_val, exc_tb=exc_tb) + self._connection = None + return None + + @mypyc_attr(native_class=False) class AsyncmyConfig(AsyncDatabaseConfig[AsyncmyConnection, "AsyncmyPool", AsyncmyDriver]): """Configuration for Asyncmy database connections.""" @@ -278,10 +337,10 @@ class AsyncmyConfig(AsyncDatabaseConfig[AsyncmyConnection, "AsyncmyPool", Asyncm driver_type: ClassVar[type[AsyncmyDriver]] = AsyncmyDriver connection_type: "ClassVar[type[Any]]" = cast("type[Any]", AsyncmyConnection) supports_transactional_ddl: ClassVar[bool] = False - supports_native_arrow_export: ClassVar[bool] = True - supports_native_parquet_export: ClassVar[bool] = True - supports_native_arrow_import: ClassVar[bool] = True - supports_native_parquet_import: ClassVar[bool] = True + supports_native_arrow_export: ClassVar[bool] = False + supports_native_parquet_export: ClassVar[bool] = False + supports_native_arrow_import: ClassVar[bool] = False + supports_native_parquet_import: ClassVar[bool] = False supports_native_row_streaming: ClassVar[bool] = True type_coercion_capabilities: ClassVar[TypeCoercionCapabilities] = TypeCoercionCapabilities( datetime_binding="native", timestamp_precision="microsecond", json_columns_decoded=False, uuid_binding="text" @@ -355,7 +414,11 @@ async def create_connection(self) -> AsyncmyConnection: pool_kwargs, connection_kwargs = _split_pool_config(self.connection_config) connection_kwargs.update({key: value for key, value in pool_kwargs.items() if key in _ASYNCMY_SHARED_POOL_KEYS}) connection = await asyncmy.connect(**connection_kwargs) - await self._ensure_connection(connection) + try: + await self._ensure_connection(connection) + except BaseException: + connection.close() + raise return connection async def provide_pool(self, *args: Any, **kwargs: Any) -> "AsyncmyPool": diff --git a/sqlspec/adapters/asyncmy/core.py b/sqlspec/adapters/asyncmy/core.py index 3fb570e78..602da097c 100644 --- a/sqlspec/adapters/asyncmy/core.py +++ b/sqlspec/adapters/asyncmy/core.py @@ -82,12 +82,7 @@ 1061: (OperationalError, "duplicate key name"), 1091: (OperationalError, "cannot drop column or key"), } -_MYSQL_SQLSTATE_EXACT_DISPATCH: Final[dict[str, tuple[type[SQLSpecError], str]]] = { - "23505": (UniqueViolationError, "unique constraint violation"), - "23503": (ForeignKeyViolationError, "foreign key constraint violation"), - "23502": (NotNullViolationError, "not-null constraint violation"), - "23514": (CheckViolationError, "check constraint violation"), -} +_MYSQL_SQLSTATE_EXACT_DISPATCH: Final[dict[str, tuple[type[SQLSpecError], str]]] = {} _MYSQL_SQLSTATE_PREFIX_DISPATCH: Final[dict[str, tuple[type[SQLSpecError], str]]] = { "23": (IntegrityError, "integrity constraint violation"), "28": (PermissionDeniedError, "authorization error"), @@ -96,6 +91,10 @@ "08": (DatabaseConnectionError, "connection error"), "22": (DataError, "data error"), } +_MYSQL_SYNTAX_ERROR_DISPATCH: Final[dict[int, tuple[type[SQLSpecError], str]]] = { + 1054: (SQLParsingError, "unknown column"), + 1146: (SQLParsingError, "table does not exist"), +} _MYSQL_CONSTRAINT_ERROR_DISPATCH: Final[dict[int, tuple[type[SQLSpecError], str]]] = { MYSQL_ER_DUP_ENTRY: (UniqueViolationError, "unique constraint violation"), 1169: (UniqueViolationError, "unique constraint violation"), @@ -113,6 +112,8 @@ MYSQL_ER_DBACCESS_DENIED: (PermissionDeniedError, "access denied"), MYSQL_ER_ACCESS_DENIED: (PermissionDeniedError, "access denied"), MYSQL_ER_TABLEACCESS_DENIED: (PermissionDeniedError, "access denied"), + 1290: (PermissionDeniedError, "read-only server"), + 1836: (PermissionDeniedError, "read-only mode"), } _MYSQL_TRANSACTION_ERROR_DISPATCH: Final[dict[int, tuple[type[SQLSpecError], str]]] = { MYSQL_ER_LOCK_DEADLOCK: (DeadlockError, "deadlock detected"), @@ -122,6 +123,7 @@ 3024: (QueryTimeoutError, "query execution timeout"), } _MYSQL_CONNECTION_ERROR_DISPATCH: Final[dict[int, tuple[type[SQLSpecError], str]]] = { + 1040: (DatabaseConnectionError, "too many connections"), MYSQL_CR_SERVER_LOST: (ConnectionTimeoutError, "connection lost"), MYSQL_CR_CONNECTION_ERROR: (DatabaseConnectionError, "connection error"), MYSQL_CR_CONN_HOST_ERROR: (DatabaseConnectionError, "connection error"), @@ -192,8 +194,11 @@ async def close(self, error: bool = False) -> None: cursor = self._cursor self._cursor = None if cursor is not None: - with contextlib.suppress(Exception): + try: await cursor.close() + except BaseException: + with contextlib.suppress(Exception): + self._driver.connection.close() async def _start(self) -> None: try: @@ -493,6 +498,10 @@ def create_mapped_exception(error: Any, *, logger: Any | None = None) -> SQLSpec dispatch = _MYSQL_SQLSTATE_PREFIX_DISPATCH["40"] return _create_mysql_error(error, sqlstate, error_code, dispatch[0], dispatch[1]) + dispatch = _MYSQL_SYNTAX_ERROR_DISPATCH.get(error_code) if error_code is not None else None + if dispatch is not None: + return _create_mysql_error(error, sqlstate, error_code, dispatch[0], dispatch[1]) + if sqlstate_prefix == "42": dispatch = _MYSQL_SQLSTATE_PREFIX_DISPATCH["42"] return _create_mysql_error(error, sqlstate, error_code, dispatch[0], dispatch[1]) diff --git a/sqlspec/adapters/asyncmy/driver.py b/sqlspec/adapters/asyncmy/driver.py index 0ed595fb3..327221314 100644 --- a/sqlspec/adapters/asyncmy/driver.py +++ b/sqlspec/adapters/asyncmy/driver.py @@ -16,7 +16,6 @@ AsyncmyError, AsyncmyFieldType, AsyncmyMySQLError, - AsyncmySessionContext, asyncmy_local_infile, ) from sqlspec.adapters.asyncmy.core import ( @@ -53,7 +52,7 @@ from sqlspec.driver import CachedQuery, ExecutionResult from sqlspec.storage import StorageBridgeJob, StorageDestination, StorageFormat, StorageTelemetry -__all__ = ("AsyncmyCursor", "AsyncmyDriver", "AsyncmyExceptionHandler", "AsyncmySessionContext") +__all__ = ("AsyncmyCursor", "AsyncmyDriver", "AsyncmyExceptionHandler") logger = get_logger("sqlspec.adapters.asyncmy") @@ -152,8 +151,8 @@ async def dispatch_execute(self, cursor: Any, statement: "SQL") -> "ExecutionRes Returns: ExecutionResult: Statement execution results with data or row counts """ - sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) - sql = escape_literal_percent(sql, prepared_parameters, self.statement_config.parameter_validator) + sql, prepared_parameters = self._compiled_sql(statement, statement.statement_config) + sql = escape_literal_percent(sql, prepared_parameters, statement.statement_config.parameter_validator) await cursor.execute(sql, normalize_execute_parameters(prepared_parameters)) description = getattr(cursor, "description", None) or None @@ -192,15 +191,17 @@ async def dispatch_execute_many(self, cursor: Any, statement: "SQL") -> "Executi Returns: ExecutionResult: Batch execution results """ - sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) + sql, prepared_parameters = self._compiled_sql(statement, statement.statement_config) match = ASYNCMY_INSERT_VALUES_PATTERN.match(sql) if match: sql = ( - escape_literal_percent(match.group(1), prepared_parameters, self.statement_config.parameter_validator) + escape_literal_percent( + match.group(1), prepared_parameters, statement.statement_config.parameter_validator + ) + sql[match.end(1) :] ) else: - sql = escape_literal_percent(sql, prepared_parameters, self.statement_config.parameter_validator) + sql = escape_literal_percent(sql, prepared_parameters, statement.statement_config.parameter_validator) prepared_parameters = normalize_execute_many_parameters(prepared_parameters) parameter_count = len(prepared_parameters) if isinstance(prepared_parameters, Sized) else None @@ -213,9 +214,6 @@ async def dispatch_execute_many(self, cursor: Any, statement: "SQL") -> "Executi async def dispatch_execute_script(self, cursor: Any, statement: "SQL") -> "ExecutionResult": """Execute SQL script with statement splitting and parameter handling. - Splits multi-statement scripts and executes each statement sequentially. - Parameters are embedded as static values for script execution compatibility. - Args: cursor: AsyncMy cursor object statement: SQL script to execute @@ -223,25 +221,7 @@ async def dispatch_execute_script(self, cursor: Any, statement: "SQL") -> "Execu Returns: ExecutionResult: Script execution results with statement count """ - sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) - statements = self.split_script_statements(sql, statement.statement_config, strip_trailing_semicolon=True) - - successful_count = 0 - last_cursor = cursor - normalized_params = normalize_execute_parameters(prepared_parameters) - - for stmt in statements: - stmt_sql = ( - escape_literal_percent(stmt, prepared_parameters, self.statement_config.parameter_validator) - if normalized_params is not None - else stmt - ) - await cursor.execute(stmt_sql, normalized_params) - successful_count += 1 - - return self.create_execution_result( - last_cursor, statement_count=len(statements), successful_statements=successful_count, is_script_result=True - ) + return await super().dispatch_execute_script(cursor, statement) async def begin(self) -> None: """Begin a database transaction. @@ -293,8 +273,8 @@ def dispatch_select_stream(self, statement: "SQL", chunk_size: int) -> "AsyncRow """Return a native asyncmy row stream backed by an unbuffered ``SSCursor``.""" if not statement.returns_rows(): return None - sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) - sql = escape_literal_percent(sql, prepared_parameters, self.statement_config.parameter_validator) + sql, prepared_parameters = self._compiled_sql(statement, statement.statement_config) + sql = escape_literal_percent(sql, prepared_parameters, statement.statement_config.parameter_validator) return AsyncRowStream(AsyncmyStreamSource(self, sql, prepared_parameters, chunk_size, ASYNCMY_JSON_TYPE_CODES)) def handle_database_exceptions(self) -> "AsyncmyExceptionHandler": @@ -400,22 +380,6 @@ async def load_from_arrow( self._attach_partition_telemetry(telemetry_payload, partitioner) return self._storage_job(telemetry_payload, telemetry) - async def load_from_storage( - self, - table: str, - source: "StorageDestination", - *, - file_format: "StorageFormat", - partitioner: "dict[str, object] | None" = None, - overwrite: bool = False, - **kwargs: Any, - ) -> "StorageBridgeJob": - """Load staged artifacts from storage into MySQL.""" - - return await super().load_from_storage( - table, source, file_format=file_format, partitioner=partitioner, overwrite=overwrite, **kwargs - ) - def collect_rows(self, cursor: Any, fetched: "list[Any]") -> "tuple[list[Any], list[str], int]": """Collect asyncmy rows for the direct execution path.""" description = cursor.description or None @@ -477,7 +441,10 @@ def _connection_in_transaction(self) -> bool: Returns: True when the server reports an open transaction. """ - return bool(getattr(self.connection, "server_status", 0) & 1) + try: + return bool(self.connection.server_status & 1) + except AttributeError: + return False def _write_asyncmy_local_infile_payload(records: "list[tuple[Any, ...]]") -> tuple[str, str]: diff --git a/tests/unit/adapters/test_asyncmy/test_config.py b/tests/unit/adapters/test_asyncmy/test_config.py index 18472366d..fd0ab668c 100644 --- a/tests/unit/adapters/test_asyncmy/test_config.py +++ b/tests/unit/adapters/test_asyncmy/test_config.py @@ -336,3 +336,55 @@ def test_asyncmy_config_driver_features_initialization() -> None: config = AsyncmyConfig(connection_config={"local_infile": True}, driver_features={"json_serializer": str}) assert config.driver_features["json_serializer"] is str assert config.driver_features["enable_local_infile_bulk_load"] is True + + +async def test_asyncmy_create_connection_closes_socket_on_error(monkeypatch: pytest.MonkeyPatch) -> None: + """create_connection closes the standalone connection if _ensure_connection raises.""" + from unittest.mock import AsyncMock, MagicMock + + import sqlspec.adapters.asyncmy.config as config_mod + + mock_conn = MagicMock() + mock_conn.close = MagicMock() + monkeypatch.setattr(config_mod.asyncmy, "connect", AsyncMock(return_value=mock_conn)) + + config = AsyncmyConfig( + connection_config={"host": "localhost", "database": "test"}, + driver_features={"on_connection_create": AsyncMock(side_effect=RuntimeError("hook error"))}, + ) + + with pytest.raises(RuntimeError, match="hook error"): + await config.create_connection() + + assert mock_conn.close.called + + +def test_asyncmy_config_capability_flags() -> None: + """AsyncmyConfig declares accurate native Arrow/Parquet capabilities as False.""" + assert AsyncmyConfig.supports_native_arrow_export is False + assert AsyncmyConfig.supports_native_parquet_export is False + assert AsyncmyConfig.supports_native_arrow_import is False + assert AsyncmyConfig.supports_native_parquet_import is False + + +async def test_asyncmy_session_context_initialization() -> None: + """AsyncmySessionContext enters and exits using AsyncmyDriver without circular imports.""" + from unittest.mock import AsyncMock, MagicMock + + from sqlspec.adapters.asyncmy.config import AsyncmySessionContext + from sqlspec.adapters.asyncmy.core import default_statement_config + + mock_conn = MagicMock() + acquire = AsyncMock(return_value=mock_conn) + release = AsyncMock() + + ctx = AsyncmySessionContext( + acquire_connection=acquire, + release_connection=release, + statement_config=default_statement_config, + driver_features={}, + prepare_driver=lambda d: d, + ) + async with ctx as driver: + assert driver.connection is mock_conn + assert release.called diff --git a/tests/unit/adapters/test_asyncmy/test_driver.py b/tests/unit/adapters/test_asyncmy/test_driver.py index 72b3e2a55..4a4541ea8 100644 --- a/tests/unit/adapters/test_asyncmy/test_driver.py +++ b/tests/unit/adapters/test_asyncmy/test_driver.py @@ -140,3 +140,95 @@ def _spy_collect_rows(*args: Any, **kwargs: Any) -> Any: result = await driver._execute_cache_hit("SELECT 1", (), cached) assert bool(collect_called) is True assert result.data == [(1,)] + + +def test_asyncmy_mapped_exceptions_without_sqlstate() -> None: + """Explicit MySQL errnos map correctly even when exception object lacks sqlstate.""" + from sqlspec.adapters.asyncmy.core import create_mapped_exception + from sqlspec.exceptions import DatabaseConnectionError, PermissionDeniedError, SQLParsingError + + class _FakeErr(Exception): + pass + + assert isinstance(create_mapped_exception(_FakeErr(1054, "Unknown column")), SQLParsingError) + assert isinstance(create_mapped_exception(_FakeErr(1146, "Table doesn't exist")), SQLParsingError) + assert isinstance(create_mapped_exception(_FakeErr(1040, "Too many connections")), DatabaseConnectionError) + assert isinstance( + create_mapped_exception(_FakeErr(1290, "The MySQL server is running with the --read-only option")), + PermissionDeniedError, + ) + assert isinstance(create_mapped_exception(_FakeErr(1836, "Running in read-only mode")), PermissionDeniedError) + + +async def test_asyncmy_dispatch_execute_uses_statement_statement_config() -> None: + """dispatch_execute compiles using statement.statement_config.""" + from unittest.mock import AsyncMock, MagicMock + + from sqlspec import SQL + from sqlspec.core import StatementConfig + + custom_config = StatementConfig(enable_caching=False) + statement = SQL("SELECT 1", statement_config=custom_config) + cursor = MagicMock() + cursor.description = [("1", 3, None, None, None, None, None)] + cursor.execute = AsyncMock() + cursor.fetchall = AsyncMock(return_value=[(1,)]) + cursor.rowcount = 1 + cursor.lastrowid = None + + driver = AsyncmyDriver(connection=MagicMock()) + result = await driver.dispatch_execute(cursor, statement) + assert cursor.execute.called + assert result.data_row_count == 1 + + +async def test_asyncmy_dispatch_select_stream_uses_statement_config() -> None: + """dispatch_select_stream compiles using statement.statement_config.""" + from unittest.mock import MagicMock + + from sqlspec import SQL + from sqlspec.core import StatementConfig + + custom_config = StatementConfig(enable_caching=False) + statement = SQL("SELECT 1", statement_config=custom_config) + driver = AsyncmyDriver(connection=MagicMock()) + stream = driver.dispatch_select_stream(statement, chunk_size=10) + assert stream is not None + + +async def test_asyncmy_dispatch_execute_script_delegates_to_super() -> None: + """dispatch_execute_script emits script statement events via super().""" + from unittest.mock import AsyncMock, MagicMock + + from sqlspec import SQL + + cursor = MagicMock() + cursor.execute = AsyncMock() + cursor.description = None + cursor.rowcount = 0 + cursor.lastrowid = None + + driver = AsyncmyDriver(connection=MagicMock()) + statement = SQL("SELECT 1; SELECT 2").as_script() + result = await driver.dispatch_execute_script(cursor, statement) + assert result.statement_count == 2 + assert result.successful_statements == 2 + + +async def test_asyncmy_stream_source_close_closes_connection_on_cursor_close_error() -> None: + """AsyncmyStreamSource.close closes the connection if cursor.close raises.""" + from unittest.mock import AsyncMock, MagicMock + + from sqlspec.adapters.asyncmy.core import AsyncmyStreamSource + + cursor = MagicMock() + cursor.close = AsyncMock(side_effect=RuntimeError("close boom")) + connection = MagicMock() + driver = MagicMock() + driver.connection = connection + + source = AsyncmyStreamSource(driver, "SELECT 1", (), 10, set()) + source._cursor = cursor + + await source.close() + assert connection.close.called From 747a8e501db8afdb0c657827237c1ff68c3a6f4f Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sat, 31 Oct 2026 23:01:44 +0000 Subject: [PATCH 026/174] fix(mysqlconnector): guard optional driver imports, preserve exception causes in load_from_arrow, and map connection errors Guard top-level mysql-connector imports using import_optional to prevent ModuleNotFoundError when optional dependencies are missing. Delegate stream fetching and load_from_arrow statement executions through _run_with_exception_handler and _check_pending_exception to maintain traceback chaining and exception context. Map PoolError to ConnectionTimeoutError and disconnected InterfaceError / errno 2055 to DatabaseConnectionError. --- sqlspec/adapters/mysqlconnector/_typing.py | 28 ++++---- sqlspec/adapters/mysqlconnector/core.py | 30 ++++---- sqlspec/adapters/mysqlconnector/driver.py | 68 ++++++++++++------- .../adapters/test_mysqlconnector/test_core.py | 1 + .../test_mysqlconnector/test_driver.py | 41 +++++++++++ 5 files changed, 114 insertions(+), 54 deletions(-) diff --git a/sqlspec/adapters/mysqlconnector/_typing.py b/sqlspec/adapters/mysqlconnector/_typing.py index 054fea63f..079ba4929 100644 --- a/sqlspec/adapters/mysqlconnector/_typing.py +++ b/sqlspec/adapters/mysqlconnector/_typing.py @@ -7,22 +7,18 @@ import contextlib from typing import TYPE_CHECKING, Any -import mysql as _mysql -import mysql.connector as _mysql_connector -from mysql.connector import MySQLConnection as _MysqlConnectorSyncConnection -from mysql.connector import aio as _mysql_connector_aio -from mysql.connector import pooling as _mysql_connector_pooling -from mysql.connector.constants import FieldType as _MysqlConnectorFieldType -from mysql.connector.cursor import MySQLCursor as _MysqlConnectorSyncRawCursor - -from sqlspec.typing import import_optional_attr +from sqlspec.typing import import_optional, import_optional_attr if TYPE_CHECKING: from collections.abc import Awaitable, Callable from types import TracebackType from typing import ClassVar, Protocol, TypeAlias + import mysql.connector as _mysql_connector + from mysql.connector import MySQLConnection as _MysqlConnectorSyncConnection + from mysql.connector import pooling as _mysql_connector_pooling from mysql.connector.aio.pooling import MySQLConnectionPool as MysqlConnectorAsyncPool + from mysql.connector.cursor import MySQLCursor as _MysqlConnectorSyncRawCursor from sqlspec.adapters.mysqlconnector.driver import MysqlConnectorAsyncDriver, MysqlConnectorSyncDriver from sqlspec.core import StatementConfig @@ -85,14 +81,14 @@ class MysqlConnectorMysqlModuleProtocol(Protocol): if not TYPE_CHECKING: MysqlConnectorAsyncPool = import_optional_attr("mysql.connector.aio.pooling", "MySQLConnectionPool") - MysqlConnectorAio = _mysql_connector_aio - MysqlConnectorSyncConnection = _MysqlConnectorSyncConnection + MysqlConnectorAio = import_optional_attr("mysql.connector", "aio") + MysqlConnectorSyncConnection = import_optional_attr("mysql.connector", "MySQLConnection") MysqlConnectorAsyncConnection = import_optional_attr("mysql.connector.aio", "MySQLConnection") - MysqlConnectorConnectionPool = _mysql_connector_pooling.MySQLConnectionPool - MysqlConnectorError = _mysql_connector.Error - MysqlConnectorFieldType = _MysqlConnectorFieldType - MysqlConnectorMysqlModule = _mysql - MysqlConnectorSyncRawCursor = _MysqlConnectorSyncRawCursor + MysqlConnectorConnectionPool = import_optional_attr("mysql.connector.pooling", "MySQLConnectionPool") + MysqlConnectorError = import_optional_attr("mysql.connector", "Error") + MysqlConnectorFieldType = import_optional_attr("mysql.connector.constants", "FieldType") + MysqlConnectorMysqlModule = import_optional("mysql") + MysqlConnectorSyncRawCursor = import_optional_attr("mysql.connector.cursor", "MySQLCursor") MysqlConnectorAsyncRawCursor = import_optional_attr("mysql.connector.aio.cursor", "MySQLCursor") __all__ = ( diff --git a/sqlspec/adapters/mysqlconnector/core.py b/sqlspec/adapters/mysqlconnector/core.py index da3cd088d..4585555f1 100644 --- a/sqlspec/adapters/mysqlconnector/core.py +++ b/sqlspec/adapters/mysqlconnector/core.py @@ -181,22 +181,22 @@ def __init__( def start(self) -> None: handler = self._driver.handle_database_exceptions() - with handler: - try: - cursor = self._driver.connection.cursor(**self._cursor_options) - self._cursor = cursor - cursor.execute(self._sql, normalize_execute_parameters(self._parameters)) - self._row_plan = resolve_row_plan(self._cursor.description, self._json_type_codes) - except BaseException: - self.close(error=True) - raise + self._driver._run_with_exception_handler(handler, self._start) self._driver._check_pending_exception(handler) + def _start(self) -> None: + try: + cursor = self._driver.connection.cursor(**self._cursor_options) + self._cursor = cursor + cursor.execute(self._sql, normalize_execute_parameters(self._parameters)) + self._row_plan = resolve_row_plan(self._cursor.description, self._json_type_codes) + except BaseException: + self.close(error=True) + raise + def fetch_chunk(self) -> "list[dict[str, Any]]": handler = self._driver.handle_database_exceptions() - rows: list[Any] = [] - with handler: - rows = self._cursor.fetchmany(self._chunk_size) + rows = self._driver._run_with_exception_handler(handler, self._cursor.fetchmany, self._chunk_size) self._driver._check_pending_exception(handler) if not rows: return [] @@ -551,6 +551,12 @@ def normalize_lastrowid(cursor: Any) -> int | None: def create_mapped_exception(error: Any, *, logger: Any | None = None) -> SQLSpecError: """Map MySQL driver exceptions to SQLSpec errors.""" _ = logger + type_name = type(error).__name__ + if type_name == "PoolError" or "PoolError" in str(type(error)): + return ConnectionTimeoutError(str(error)) + if type_name == "InterfaceError" or getattr(error, "errno", None) == -1: + return DatabaseConnectionError(str(error)) + error_code = getattr(error, "errno", None) if error_code is None and isinstance(error, Exception) and error.args: value = error.args[0] diff --git a/sqlspec/adapters/mysqlconnector/driver.py b/sqlspec/adapters/mysqlconnector/driver.py index aa043ce6b..f4aac619c 100644 --- a/sqlspec/adapters/mysqlconnector/driver.py +++ b/sqlspec/adapters/mysqlconnector/driver.py @@ -290,10 +290,8 @@ def load_from_arrow( if overwrite: statement = f"TRUNCATE TABLE {format_identifier(table)}" exc_handler = self.handle_database_exceptions() - with exc_handler, self.with_cursor(self.connection) as cursor: - cursor.execute(statement) - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + self._run_with_exception_handler(exc_handler, self._execute_cursor_statement, statement) + self._check_pending_exception(exc_handler) columns, records = self._arrow_table_to_rows(arrow_table) if records: @@ -309,12 +307,10 @@ def load_from_arrow( try: with os.fdopen(fd, "wb") as tmp: tmp.write(payload) - load_sql = build_load_data_statement(table, columns) + load_sql = build_load_data_statement(table, columns, escape_percent=True) exc_handler = self.handle_database_exceptions() - with exc_handler, self.with_cursor(self.connection) as cursor: - cursor.execute(load_sql, (tmp_name,)) - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + self._run_with_exception_handler(exc_handler, self._execute_cursor_statement, load_sql, (tmp_name,)) + self._check_pending_exception(exc_handler) finally: Path(tmp_name).unlink(missing_ok=True) else: @@ -325,10 +321,10 @@ def load_from_arrow( else records ) exc_handler = self.handle_database_exceptions() - with exc_handler, self.with_cursor(self.connection) as cursor: - cursor.executemany(insert_sql, cast("Any", prepared_records)) - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + self._run_with_exception_handler( + exc_handler, self._executemany_cursor_statement, insert_sql, prepared_records + ) + self._check_pending_exception(exc_handler) telemetry_payload = self._ingest_telemetry(arrow_table) telemetry_payload["destination"] = table @@ -361,6 +357,17 @@ def resolve_rowcount(self, cursor: Any) -> int: """Resolve rowcount from mysql-connector cursor for the direct execution path.""" return resolve_rowcount(cursor) + def _execute_cursor_statement(self, statement: str, parameters: Any = None) -> None: + with self.with_cursor(self.connection) as cursor: + if parameters is None: + cursor.execute(statement) + else: + cursor.execute(statement, parameters) + + def _executemany_cursor_statement(self, statement: str, parameters: Any) -> None: + with self.with_cursor(self.connection) as cursor: + cursor.executemany(statement, parameters) + def _connection_in_transaction(self) -> bool: """Report whether the connector has an open transaction.""" return bool(self.connection.in_transaction) @@ -560,10 +567,8 @@ async def load_from_arrow( if overwrite: statement = f"TRUNCATE TABLE {format_identifier(table)}" exc_handler = self.handle_database_exceptions() - async with exc_handler, self.with_cursor(self.connection) as cursor: - await cursor.execute(statement) - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + await self._run_with_exception_handler(exc_handler, self._execute_cursor_statement, statement) + self._check_pending_exception(exc_handler) columns, records = self._arrow_table_to_rows(arrow_table) if records: @@ -579,12 +584,12 @@ async def load_from_arrow( try: with os.fdopen(fd, "wb") as tmp: tmp.write(payload) - load_sql = build_load_data_statement(table, columns) + load_sql = build_load_data_statement(table, columns, escape_percent=True) exc_handler = self.handle_database_exceptions() - async with exc_handler, self.with_cursor(self.connection) as cursor: - await cursor.execute(load_sql, (tmp_name,)) - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + await self._run_with_exception_handler( + exc_handler, self._execute_cursor_statement, load_sql, (tmp_name,) + ) + self._check_pending_exception(exc_handler) finally: await asyncio.to_thread(Path(tmp_name).unlink, True) else: @@ -595,10 +600,10 @@ async def load_from_arrow( else records ) exc_handler = self.handle_database_exceptions() - async with exc_handler, self.with_cursor(self.connection) as cursor: - await cursor.executemany(insert_sql, cast("Any", prepared_records)) - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + await self._run_with_exception_handler( + exc_handler, self._executemany_cursor_statement, insert_sql, prepared_records + ) + self._check_pending_exception(exc_handler) telemetry_payload = self._ingest_telemetry(arrow_table) telemetry_payload["destination"] = table @@ -631,6 +636,17 @@ def resolve_rowcount(self, cursor: Any) -> int: """Resolve rowcount from mysql-connector cursor for the direct execution path.""" return resolve_rowcount(cursor) + async def _execute_cursor_statement(self, statement: str, parameters: Any = None) -> None: + async with self.with_cursor(self.connection) as cursor: + if parameters is None: + await cursor.execute(statement) + else: + await cursor.execute(statement, parameters) + + async def _executemany_cursor_statement(self, statement: str, parameters: Any) -> None: + async with self.with_cursor(self.connection) as cursor: + await cursor.executemany(statement, parameters) + def _connection_in_transaction(self) -> bool: """Report whether the connector has an open transaction.""" return bool(self.connection.in_transaction) diff --git a/tests/unit/adapters/test_mysqlconnector/test_core.py b/tests/unit/adapters/test_mysqlconnector/test_core.py index bfbcb0c1f..4939f4349 100644 --- a/tests/unit/adapters/test_mysqlconnector/test_core.py +++ b/tests/unit/adapters/test_mysqlconnector/test_core.py @@ -35,6 +35,7 @@ def test_sync_stream_source_closes_cursor_when_execute_raises() -> None: connection = MagicMock() connection.cursor.return_value = cursor driver = MagicMock(connection=connection, driver_features={}) + driver._run_with_exception_handler = lambda _handler, fn: fn() source = MysqlConnectorSyncStreamSource(driver, "SELECT 1", (), 100, set()) with pytest.raises(RuntimeError, match="execute boom"): diff --git a/tests/unit/adapters/test_mysqlconnector/test_driver.py b/tests/unit/adapters/test_mysqlconnector/test_driver.py index d3a5a7444..d0c19f368 100644 --- a/tests/unit/adapters/test_mysqlconnector/test_driver.py +++ b/tests/unit/adapters/test_mysqlconnector/test_driver.py @@ -41,3 +41,44 @@ def _fail_commit() -> None: driver = MysqlConnectorSyncDriver(cast(Any, connection)) with pytest.raises(DeadlockError): driver.commit() + + +def test_mysqlconnector_pool_error_and_connection_error_mapping() -> None: + """PoolError maps to ConnectionTimeoutError and 2055 maps to DatabaseConnectionError.""" + from sqlspec.adapters.mysqlconnector.core import create_mapped_exception + from sqlspec.exceptions import ConnectionTimeoutError, DatabaseConnectionError + + class PoolError(Exception): + pass + + class InterfaceError(Exception): + errno = -1 + + assert isinstance(create_mapped_exception(PoolError("connection pool exhausted")), ConnectionTimeoutError) + assert isinstance(create_mapped_exception(InterfaceError("Connection disconnected")), DatabaseConnectionError) + assert isinstance( + create_mapped_exception(MysqlConnectorError(errno=2055, msg="Lost connection")), DatabaseConnectionError + ) + + +def test_mysqlconnector_sync_stream_source_in_active_except_block() -> None: + """MysqlConnectorSyncStreamSource safely handles execution inside an active except block.""" + from unittest.mock import MagicMock + + from sqlspec.adapters.mysqlconnector.core import MysqlConnectorSyncStreamSource + + cursor = MagicMock() + cursor.fetchmany.return_value = [{"id": 1}] + cursor.description = [("id", 3, None, None, None, None, None)] + connection = MagicMock() + connection.cursor.return_value = cursor + driver = MysqlConnectorSyncDriver(connection) + + try: + raise ValueError("simulated active exception") + except ValueError: + source = MysqlConnectorSyncStreamSource(driver, "SELECT 1", (), 10, set()) + source.start() + chunk = source.fetch_chunk() + assert len(chunk) == 1 + source.close() From 9d63410d1564c348b91b30e584820ccceba4e6e1 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sat, 31 Oct 2026 23:02:33 +0000 Subject: [PATCH 027/174] fix(mysqlconnector): ensure transaction rollback on commit failure, handle missing tables, and align store DDL Move conn.commit() inside the try block before rollback in sync append_event_and_update_state. Handle missing-table errors gracefully in sync and async ADK memory stores, remove redundant DROP INDEX from Litestar session teardown, use VALUES(col) in Litestar session upsert statements for cross-version MySQL and MariaDB compatibility, and align sync create_table lifecycle by removing redundant commit after execute_script. --- sqlspec/adapters/mysqlconnector/adk/store.py | 128 +++++++++++------- .../adapters/mysqlconnector/litestar/store.py | 25 ++-- .../test_mysqlconnector/test_adk_store.py | 79 +++++++++++ .../test_litestar_store.py | 51 +++++++ 4 files changed, 217 insertions(+), 66 deletions(-) diff --git a/sqlspec/adapters/mysqlconnector/adk/store.py b/sqlspec/adapters/mysqlconnector/adk/store.py index b6ef15019..96806dd03 100644 --- a/sqlspec/adapters/mysqlconnector/adk/store.py +++ b/sqlspec/adapters/mysqlconnector/adk/store.py @@ -759,12 +759,12 @@ def append_event_and_update_state( _mysql_upsert_user_state_sql(self._user_state_table), (app_name, user_id, self._json_serializer(user_state)), ) + conn.commit() except Exception: conn.rollback() raise finally: cursor.close() - conn.commit() return _session_record_from_row(row, deserializer=self._json_deserializer) @@ -1045,14 +1045,19 @@ async def search_entries( """ params = (*scope_params, f"%{query}%", limit_value) - async with self._config.provide_connection() as conn: - cursor = await conn.cursor() - try: - await cursor.execute(sql, params) - rows = await cursor.fetchall() - columns = [col[0] for col in cursor.description or []] - finally: - await cursor.close() + try: + async with self._config.provide_connection() as conn: + cursor = await conn.cursor() + try: + await cursor.execute(sql, params) + rows = await cursor.fetchall() + columns = [col[0] for col in cursor.description or []] + finally: + await cursor.close() + except MysqlConnectorError as exc: + if _is_mysql_table_missing(exc): + return [] + raise records: list[StoredMemory] = [] for row in rows: @@ -1072,14 +1077,19 @@ async def delete_entries_by_session(self, session_id: str) -> int: raise RuntimeError(msg) sql = f"DELETE FROM {self._memory_table} WHERE session_id = %s" - async with self._config.provide_connection() as conn: - cursor = await conn.cursor() - try: - await cursor.execute(sql, (session_id,)) - await conn.commit() - return cursor.rowcount if cursor.rowcount and cursor.rowcount > 0 else 0 - finally: - await cursor.close() + try: + async with self._config.provide_connection() as conn: + cursor = await conn.cursor() + try: + await cursor.execute(sql, (session_id,)) + await conn.commit() + return cursor.rowcount if cursor.rowcount and cursor.rowcount > 0 else 0 + finally: + await cursor.close() + except MysqlConnectorError as exc: + if _is_mysql_table_missing(exc): + return 0 + raise async def delete_entries_older_than( self, days: int, app_name: "str | None" = None, scope: "str | None" = None @@ -1099,14 +1109,19 @@ async def delete_entries_older_than( where_sql = " AND ".join(clauses) sql = f"DELETE FROM {self._memory_table} WHERE {where_sql}" - async with self._config.provide_connection() as conn: - cursor = await conn.cursor() - try: - await cursor.execute(sql, tuple(params)) - await conn.commit() - return cursor.rowcount if cursor.rowcount and cursor.rowcount > 0 else 0 - finally: - await cursor.close() + try: + async with self._config.provide_connection() as conn: + cursor = await conn.cursor() + try: + await cursor.execute(sql, tuple(params)) + await conn.commit() + return cursor.rowcount if cursor.rowcount and cursor.rowcount > 0 else 0 + finally: + await cursor.close() + except MysqlConnectorError as exc: + if _is_mysql_table_missing(exc): + return 0 + raise async def _memory_table_ddl(self) -> str: adk_config = _adk_config(self._config) @@ -1279,14 +1294,19 @@ def search_entries( """ params = (*scope_params, f"%{query}%", limit_value) - with self._config.provide_connection() as conn: - cursor = conn.cursor() - try: - cursor.execute(sql, params) - rows = cursor.fetchall() - columns = [col[0] for col in cursor.description or []] - finally: - cursor.close() + try: + with self._config.provide_connection() as conn: + cursor = conn.cursor() + try: + cursor.execute(sql, params) + rows = cursor.fetchall() + columns = [col[0] for col in cursor.description or []] + finally: + cursor.close() + except MysqlConnectorError as exc: + if _is_mysql_table_missing(exc): + return [] + raise records: list[StoredMemory] = [] for row in rows: @@ -1307,14 +1327,19 @@ def delete_entries_by_session(self, session_id: str) -> int: raise RuntimeError(msg) sql = f"DELETE FROM {self._memory_table} WHERE session_id = %s" - with self._config.provide_connection() as conn: - cursor = conn.cursor() - try: - cursor.execute(sql, (session_id,)) - conn.commit() - return cursor.rowcount if cursor.rowcount and cursor.rowcount > 0 else 0 - finally: - cursor.close() + try: + with self._config.provide_connection() as conn: + cursor = conn.cursor() + try: + cursor.execute(sql, (session_id,)) + conn.commit() + return cursor.rowcount if cursor.rowcount and cursor.rowcount > 0 else 0 + finally: + cursor.close() + except MysqlConnectorError as exc: + if _is_mysql_table_missing(exc): + return 0 + raise def delete_entries_older_than(self, days: int, app_name: "str | None" = None, scope: "str | None" = None) -> int: """Delete memory entries older than specified days.""" @@ -1333,14 +1358,19 @@ def delete_entries_older_than(self, days: int, app_name: "str | None" = None, sc where_sql = " AND ".join(clauses) sql = f"DELETE FROM {self._memory_table} WHERE {where_sql}" - with self._config.provide_connection() as conn: - cursor = conn.cursor() - try: - cursor.execute(sql, tuple(params)) - conn.commit() - return cursor.rowcount if cursor.rowcount and cursor.rowcount > 0 else 0 - finally: - cursor.close() + try: + with self._config.provide_connection() as conn: + cursor = conn.cursor() + try: + cursor.execute(sql, tuple(params)) + conn.commit() + return cursor.rowcount if cursor.rowcount and cursor.rowcount > 0 else 0 + finally: + cursor.close() + except MysqlConnectorError as exc: + if _is_mysql_table_missing(exc): + return 0 + raise def _memory_table_ddl(self) -> str: adk_config = _adk_config(self._config) diff --git a/sqlspec/adapters/mysqlconnector/litestar/store.py b/sqlspec/adapters/mysqlconnector/litestar/store.py index 77961f2e1..69ae7d4a5 100644 --- a/sqlspec/adapters/mysqlconnector/litestar/store.py +++ b/sqlspec/adapters/mysqlconnector/litestar/store.py @@ -103,10 +103,10 @@ async def set(self, key: str, value: "str | bytes", expires_in: "int | timedelta sql = f""" INSERT INTO {self._table_name} (session_id, data, expires_at) - VALUES (%s, %s, %s) AS new + VALUES (%s, %s, %s) ON DUPLICATE KEY UPDATE - data = new.data, - expires_at = new.expires_at, + data = VALUES(data), + expires_at = VALUES(expires_at), updated_at = UTC_TIMESTAMP(6) """ @@ -237,11 +237,7 @@ def _table_ddl(self) -> str: """ def _drop_table_sql(self) -> "list[str]": - safe_index_table = self._table_name.replace(".", "_") - return [ - f"DROP INDEX idx_{safe_index_table}_expires_at ON {self._table_name}", - f"DROP TABLE IF EXISTS {self._table_name}", - ] + return [f"DROP TABLE IF EXISTS {self._table_name}"] class MysqlConnectorSyncStore(BaseSQLSpecStore["MysqlConnectorSyncConfig"]): @@ -298,17 +294,12 @@ def _table_ddl(self) -> str: """ def _drop_table_sql(self) -> "list[str]": - safe_index_table = self._table_name.replace(".", "_") - return [ - f"DROP INDEX idx_{safe_index_table}_expires_at ON {self._table_name}", - f"DROP TABLE IF EXISTS {self._table_name}", - ] + return [f"DROP TABLE IF EXISTS {self._table_name}"] def _create_table(self) -> None: sql = self._table_ddl() with self._config.provide_session() as driver: driver.execute_script(sql) - driver.commit() self._log_table_created() def _get(self, key: str, renew_for: "int | timedelta | None" = None) -> "bytes | None": @@ -358,10 +349,10 @@ def _set(self, key: str, value: "str | bytes", expires_in: "int | timedelta | No sql = f""" INSERT INTO {self._table_name} (session_id, data, expires_at) - VALUES (%s, %s, %s) AS new + VALUES (%s, %s, %s) ON DUPLICATE KEY UPDATE - data = new.data, - expires_at = new.expires_at, + data = VALUES(data), + expires_at = VALUES(expires_at), updated_at = UTC_TIMESTAMP(6) """ diff --git a/tests/unit/adapters/test_mysqlconnector/test_adk_store.py b/tests/unit/adapters/test_mysqlconnector/test_adk_store.py index 7eb19839f..bfc594674 100644 --- a/tests/unit/adapters/test_mysqlconnector/test_adk_store.py +++ b/tests/unit/adapters/test_mysqlconnector/test_adk_store.py @@ -473,3 +473,82 @@ async def test_mysqlconnector_async_adk_memory_store_insert_and_search_json_hand assert len(records) == 1 assert records[0]["content_json"] == {"text": "hello"} assert records[0]["metadata_json"] == {"source": "unit"} + + +def test_mysqlconnector_sync_append_event_rolls_back_if_commit_fails() -> None: + """MysqlConnectorSyncADKStore rolls back when conn.commit() fails.""" + from unittest.mock import MagicMock + + from sqlspec.adapters.mysqlconnector.adk.store import MysqlConnectorSyncADKStore + + cursor = MagicMock() + cursor.fetchall.return_value = [] + cursor.fetchone.return_value = ("sess-1", "app", "u1", "{}", None, None) + conn = MagicMock() + conn.__enter__.return_value = conn + conn.__exit__.return_value = None + conn.cursor.return_value = cursor + conn.commit.side_effect = RuntimeError("commit failed") + config = _mock_config() + config.provide_connection = lambda *_a, **_k: conn + + store = MysqlConnectorSyncADKStore(config) + event: dict[str, Any] = { + "id": "e-1", + "session_id": "sess-1", + "app_name": "app", + "user_id": "u1", + "event_id": "evt-1", + "event_type": "user", + "event_data": {"msg": "hi"}, + "timestamp": None, + "invocation_id": "inv-1", + } + + with pytest.raises(RuntimeError, match="commit failed"): + store.append_event_and_update_state(event, "app", "u1", "sess-1", {"state": "active"}) + + assert conn.rollback.called + + +async def test_mysqlconnector_memory_store_missing_table_handling() -> None: + """search_entries and deletes return empty/0 when table is missing.""" + from unittest.mock import AsyncMock, MagicMock + + from sqlspec.adapters.mysqlconnector._typing import MysqlConnectorError + from sqlspec.adapters.mysqlconnector.adk.store import ( + MysqlConnectorAsyncADKMemoryStore, + MysqlConnectorSyncADKMemoryStore, + ) + + missing_err = MysqlConnectorError(errno=1146, msg="Table 'test.mem' doesn't exist") + + sync_cursor = MagicMock() + sync_cursor.rowcount = 0 + sync_cursor.execute.side_effect = missing_err + sync_conn = MagicMock() + sync_conn.__enter__.return_value = sync_conn + sync_conn.__exit__.return_value = None + sync_conn.cursor.return_value = sync_cursor + sync_config = _mock_config() + sync_config.provide_connection = lambda *_a, **_k: sync_conn + + sync_store = MysqlConnectorSyncADKMemoryStore(sync_config) + assert sync_store.search_entries("test", "app", "u1") == [] + assert sync_store.delete_entries_by_session("sess-1") == 0 + assert sync_store.delete_entries_older_than(30) == 0 + + async_cursor = MagicMock() + async_cursor.execute = AsyncMock(side_effect=missing_err) + async_cursor.close = AsyncMock() + async_conn = MagicMock() + async_conn.cursor = AsyncMock(return_value=async_cursor) + async_conn.__aenter__ = AsyncMock(return_value=async_conn) + async_conn.__aexit__ = AsyncMock(return_value=None) + async_config = _mock_config() + async_config.provide_connection = lambda *_a, **_k: async_conn + + async_store = MysqlConnectorAsyncADKMemoryStore(async_config) + assert await async_store.search_entries("test", "app", "u1") == [] + assert await async_store.delete_entries_by_session("sess-1") == 0 + assert await async_store.delete_entries_older_than(30) == 0 diff --git a/tests/unit/adapters/test_mysqlconnector/test_litestar_store.py b/tests/unit/adapters/test_mysqlconnector/test_litestar_store.py index 45e49365a..10a0a1f66 100644 --- a/tests/unit/adapters/test_mysqlconnector/test_litestar_store.py +++ b/tests/unit/adapters/test_mysqlconnector/test_litestar_store.py @@ -48,3 +48,54 @@ def test_mysqlconnector_sync_store_honors_table_options() -> None: ddl = store._table_ddl() assert f"{DEFAULT_ENGINE_CLAUSE} COMMENT='litestar-session'" in ddl + + +def test_mysqlconnector_drop_table_sql_omits_drop_index() -> None: + """_drop_table_sql must emit only DROP TABLE IF EXISTS without preceding DROP INDEX.""" + async_store = MysqlConnectorAsyncStore(_mock_config()) + sync_store = MysqlConnectorSyncStore(_mock_config()) + + assert async_store._drop_table_sql() == ["DROP TABLE IF EXISTS litestar_session"] + assert sync_store._drop_table_sql() == ["DROP TABLE IF EXISTS litestar_session"] + + +async def test_mysqlconnector_set_uses_mariadb_compatible_upsert() -> None: + """set must use VALUES(col) in ON DUPLICATE KEY UPDATE for MySQL and MariaDB compatibility.""" + from unittest.mock import AsyncMock + + cursor = MagicMock() + cursor.execute = AsyncMock() + cursor.close = AsyncMock() + conn = MagicMock() + conn.cursor = AsyncMock(return_value=cursor) + conn.commit = AsyncMock() + conn.__aenter__ = AsyncMock(return_value=conn) + conn.__aexit__ = AsyncMock(return_value=None) + config = _mock_config() + config.provide_connection = lambda *_a, **_k: conn + + store = MysqlConnectorAsyncStore(config) + await store.set("session-1", b"payload", 300) + + assert cursor.execute.call_count == 1 + sql = cursor.execute.call_args_list[0].args[0] + assert "ON DUPLICATE KEY UPDATE" in sql + assert "data = VALUES(data)" in sql + assert "expires_at = VALUES(expires_at)" in sql + assert "AS new" not in sql + + +def test_mysqlconnector_sync_create_table_does_not_commit() -> None: + """Sync _create_table does not call driver.commit() after execute_script.""" + driver = MagicMock() + session = MagicMock() + session.__enter__ = MagicMock(return_value=driver) + session.__exit__ = MagicMock(return_value=None) + config = _mock_config() + config.provide_session = lambda *_a, **_k: session + + store = MysqlConnectorSyncStore(config) + store._create_table() + + assert driver.execute_script.called + assert not driver.commit.called From ea903113b683c228350e14393065a4fa972b8e86 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sat, 31 Oct 2026 23:02:34 +0000 Subject: [PATCH 028/174] refactor(mysqlconnector): verify cached data dictionary properties and shared metadata setup Verify that data_dictionary properties on sync and async mysqlconnector drivers cache their instances. Use the consolidated scalar DDL metadata query and object_name filtering from the shared MySQL data dictionary. --- .../unit/adapters/test_mysqlconnector/test_driver.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/tests/unit/adapters/test_mysqlconnector/test_driver.py b/tests/unit/adapters/test_mysqlconnector/test_driver.py index d0c19f368..71ac84686 100644 --- a/tests/unit/adapters/test_mysqlconnector/test_driver.py +++ b/tests/unit/adapters/test_mysqlconnector/test_driver.py @@ -82,3 +82,14 @@ def test_mysqlconnector_sync_stream_source_in_active_except_block() -> None: chunk = source.fetch_chunk() assert len(chunk) == 1 source.close() + + +def test_mysqlconnector_data_dictionary_cached_property() -> None: + """MysqlConnectorSyncDriver and MysqlConnectorAsyncDriver cache the data dictionary instance.""" + from unittest.mock import MagicMock + + sync_driver = MysqlConnectorSyncDriver(connection=MagicMock()) + assert sync_driver.data_dictionary is sync_driver.data_dictionary + + async_driver = MysqlConnectorAsyncDriver(connection=MagicMock()) + assert async_driver.data_dictionary is async_driver.data_dictionary From 4c15a8bb55c06ac51552a7a4ef682bf629497757 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sat, 31 Oct 2026 23:02:34 +0000 Subject: [PATCH 029/174] perf(mysqlconnector): precompute cursor options, cache stream deserializers, and compile column resolution Precompute cursor_options in driver __init__ to avoid dictionary allocations on with_cursor calls. Cache _deserializer in stream source instances on initialization to eliminate per-chunk getattr lookups, and move _resolve_column_types and type code mappings into compiled core.py. --- sqlspec/adapters/mysqlconnector/_typing.py | 14 +++-- sqlspec/adapters/mysqlconnector/core.py | 62 +++++++++++++++---- sqlspec/adapters/mysqlconnector/driver.py | 55 ++++------------ .../adapters/test_mysqlconnector/test_core.py | 29 +++++++++ .../test_mysqlconnector/test_driver.py | 10 +++ 5 files changed, 109 insertions(+), 61 deletions(-) diff --git a/sqlspec/adapters/mysqlconnector/_typing.py b/sqlspec/adapters/mysqlconnector/_typing.py index 079ba4929..558a92349 100644 --- a/sqlspec/adapters/mysqlconnector/_typing.py +++ b/sqlspec/adapters/mysqlconnector/_typing.py @@ -122,10 +122,13 @@ def __init__( ) -> None: self.connection = connection self.cursor: MysqlConnectorSyncRawCursor | None = None - self.cursor_options = cursor_options or {} + self.cursor_options = cursor_options def __enter__(self) -> "MysqlConnectorSyncRawCursor": - self.cursor = self.connection.cursor(**self.cursor_options) + if self.cursor_options is None: + self.cursor = self.connection.cursor() + else: + self.cursor = self.connection.cursor(**self.cursor_options) return self.cursor def __exit__(self, *_: object) -> None: @@ -148,10 +151,13 @@ def __init__( ) -> None: self.connection = connection self.cursor: MysqlConnectorAsyncRawCursor | None = None - self.cursor_options = cursor_options or {} + self.cursor_options = cursor_options async def __aenter__(self) -> "MysqlConnectorAsyncRawCursor": - self.cursor = await self.connection.cursor(**self.cursor_options) + if self.cursor_options is None: + self.cursor = await self.connection.cursor() + else: + self.cursor = await self.connection.cursor(**self.cursor_options) return self.cursor async def __aexit__(self, *_: object) -> None: diff --git a/sqlspec/adapters/mysqlconnector/core.py b/sqlspec/adapters/mysqlconnector/core.py index 4585555f1..548032f5a 100644 --- a/sqlspec/adapters/mysqlconnector/core.py +++ b/sqlspec/adapters/mysqlconnector/core.py @@ -3,7 +3,7 @@ import contextlib from collections.abc import Callable, Mapping, Sequence, Sized from datetime import timedelta -from typing import TYPE_CHECKING, Any, Final, Literal, cast +from typing import TYPE_CHECKING, Any, Final, Literal from sqlspec.core import DriverParameterProfile, ParameterStyle, StatementConfig, build_statement_config_from_profile from sqlspec.driver import rows_to_dicts @@ -36,6 +36,7 @@ __all__ = ( "MysqlConnectorAsyncStreamSource", "MysqlConnectorSyncStreamSource", + "_resolve_column_types", "apply_driver_features", "build_insert_statement", "build_load_data_statement", @@ -153,6 +154,7 @@ class MysqlConnectorSyncStreamSource: "_chunk_size", "_cursor", "_cursor_options", + "_deserializer", "_driver", "_json_type_codes", "_parameters", @@ -178,6 +180,10 @@ def __init__( self._cursor_options = dict(cursor_options or {}) self._cursor_options["buffered"] = False self._row_plan: tuple[list[str], list[int] | None] | None = None + deserializer = getattr(driver, "_json_deserializer", None) + if deserializer is None and hasattr(driver, "driver_features") and isinstance(driver.driver_features, dict): + deserializer = driver.driver_features.get("json_deserializer") + self._deserializer: Callable[[Any], Any] = deserializer or from_json def start(self) -> None: handler = self._driver.handle_database_exceptions() @@ -202,12 +208,7 @@ def fetch_chunk(self) -> "list[dict[str, Any]]": return [] if self._row_plan is None: self._row_plan = resolve_row_plan(self._cursor.description, self._json_type_codes) - deserializer = cast( - "Callable[[Any], Any]", - getattr(self._driver, "_json_deserializer", None) - or self._driver.driver_features.get("json_deserializer", from_json), - ) - return collect_stream_rows(rows, self._row_plan, deserializer) + return collect_stream_rows(rows, self._row_plan, self._deserializer) def close(self, error: bool = False) -> None: cursor = self._cursor @@ -231,6 +232,7 @@ class MysqlConnectorAsyncStreamSource: "_chunk_size", "_cursor", "_cursor_options", + "_deserializer", "_driver", "_json_type_codes", "_parameters", @@ -256,6 +258,10 @@ def __init__( self._cursor_options = dict(cursor_options or {}) self._cursor_options["buffered"] = False self._row_plan: tuple[list[str], list[int] | None] | None = None + deserializer = getattr(driver, "_json_deserializer", None) + if deserializer is None and hasattr(driver, "driver_features") and isinstance(driver.driver_features, dict): + deserializer = driver.driver_features.get("json_deserializer") + self._deserializer: Callable[[Any], Any] = deserializer or from_json async def start(self) -> None: handler = self._driver.handle_database_exceptions() @@ -270,12 +276,7 @@ async def fetch_chunk(self) -> "list[dict[str, Any]]": return [] if self._row_plan is None: self._row_plan = resolve_row_plan(self._cursor.description, self._json_type_codes) - deserializer = cast( - "Callable[[Any], Any]", - getattr(self._driver, "_json_deserializer", None) - or self._driver.driver_features.get("json_deserializer", from_json), - ) - return collect_stream_rows(rows, self._row_plan, deserializer) + return collect_stream_rows(rows, self._row_plan, self._deserializer) async def close(self, error: bool = False) -> None: cursor = self._cursor @@ -729,4 +730,39 @@ def _create_mysql_error( driver_profile = build_profile() +_MYSQL_TYPE_CODE_TOKENS: Final[dict[int, str]] = { + 0: "decimal", + 1: "int32", + 2: "int32", + 3: "int64", + 4: "float32", + 5: "float64", + 7: "timestamp", + 8: "int64", + 10: "date", + 11: "time", + 12: "timestamp", + 246: "decimal", + 252: "binary", + 253: "string", + 254: "string", +} + + +def _resolve_column_types(description: Any) -> "dict[str, str] | None": + """Map MySQL cursor column FIELD_TYPE codes to neutral Arrow type tokens. + + Returns ``None`` when the cursor has no description or reports no + recognizable type codes. + """ + if not description: + return None + column_types: dict[str, str] = {} + for col in description: + token = _MYSQL_TYPE_CODE_TOKENS.get(col[1]) + if token is not None: + column_types[col[0]] = token + return column_types or None + + default_statement_config = build_statement_config() diff --git a/sqlspec/adapters/mysqlconnector/driver.py b/sqlspec/adapters/mysqlconnector/driver.py index f4aac619c..9632c0f88 100644 --- a/sqlspec/adapters/mysqlconnector/driver.py +++ b/sqlspec/adapters/mysqlconnector/driver.py @@ -22,6 +22,7 @@ from sqlspec.adapters.mysqlconnector.core import ( MysqlConnectorAsyncStreamSource, MysqlConnectorSyncStreamSource, + _resolve_column_types, build_insert_statement, build_load_data_statement, can_use_local_infile_records, @@ -77,24 +78,6 @@ json_type_value = MysqlConnectorFieldType.JSON if supports_json_type(MysqlConnectorFieldType) else None MYSQLCONNECTOR_JSON_TYPE_CODES: Final[set[int]] = {json_type_value} if json_type_value is not None else set() -_MYSQL_TYPE_CODE_TOKENS: Final[dict[int, str]] = { - 0: "decimal", - 1: "int32", - 2: "int32", - 3: "int64", - 4: "float32", - 5: "float64", - 7: "timestamp", - 8: "int64", - 10: "date", - 11: "time", - 12: "timestamp", - 246: "decimal", - 252: "binary", - 253: "string", - 254: "string", -} - class MysqlConnectorSyncExceptionHandler(BaseSyncExceptionHandler): """Context manager for handling mysql-connector sync exceptions.""" @@ -113,7 +96,7 @@ def _handle_exception(self, exc_type: "type[BaseException] | None", exc_val: "Ba class MysqlConnectorSyncDriver(SyncDriverAdapterBase): """MySQL/MariaDB database driver using mysql-connector sync library.""" - __slots__ = ("_data_dictionary", "_json_deserializer") + __slots__ = ("_cursor_options", "_data_dictionary", "_json_deserializer") dialect = "mysql" def __init__( @@ -132,6 +115,8 @@ def __init__( self._json_deserializer: Callable[[Any], Any] = cast( "Callable[[Any], Any]", self.driver_features.get("json_deserializer", from_json) ) + cursor_options = cast("dict[str, Any] | None", self.driver_features.get("cursor_options")) + self._cursor_options: dict[str, Any] | None = dict(cursor_options) if cursor_options else None @property def data_dictionary(self) -> "MysqlConnectorSyncDataDictionary": @@ -212,18 +197,16 @@ def rollback(self) -> None: raise create_mapped_exception(e) from e def with_cursor(self, connection: "MysqlConnectorSyncConnection") -> "MysqlConnectorSyncCursor": - cursor_options = cast("dict[str, Any]", self.driver_features.get("cursor_options") or {}) - return MysqlConnectorSyncCursor(connection, dict(cursor_options)) + return MysqlConnectorSyncCursor(connection, self._cursor_options) def dispatch_select_stream(self, statement: "SQL", chunk_size: int) -> "SyncRowStream[dict[str, Any]] | None": """Return a native mysql-connector row stream backed by an unbuffered cursor.""" if not statement.returns_rows(): return None sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) - cursor_options = cast("dict[str, Any]", self.driver_features.get("cursor_options") or {}) return SyncRowStream( MysqlConnectorSyncStreamSource( - self, sql, prepared_parameters, chunk_size, MYSQLCONNECTOR_JSON_TYPE_CODES, cursor_options + self, sql, prepared_parameters, chunk_size, MYSQLCONNECTOR_JSON_TYPE_CODES, self._cursor_options ) ) @@ -390,7 +373,7 @@ def _handle_exception(self, exc_type: "type[BaseException] | None", exc_val: "Ba class MysqlConnectorAsyncDriver(AsyncDriverAdapterBase): """MySQL/MariaDB database driver using mysql-connector async library.""" - __slots__ = ("_data_dictionary", "_json_deserializer") + __slots__ = ("_cursor_options", "_data_dictionary", "_json_deserializer") dialect = "mysql" def __init__( @@ -409,6 +392,8 @@ def __init__( self._json_deserializer: Callable[[Any], Any] = cast( "Callable[[Any], Any]", self.driver_features.get("json_deserializer", from_json) ) + cursor_options = cast("dict[str, Any] | None", self.driver_features.get("cursor_options")) + self._cursor_options: dict[str, Any] | None = dict(cursor_options) if cursor_options else None @property def data_dictionary(self) -> "MysqlConnectorAsyncDataDictionary": @@ -489,18 +474,16 @@ async def rollback(self) -> None: raise create_mapped_exception(e) from e def with_cursor(self, connection: "MysqlConnectorAsyncConnection") -> "MysqlConnectorAsyncCursor": - cursor_options = cast("dict[str, Any]", self.driver_features.get("cursor_options") or {}) - return MysqlConnectorAsyncCursor(connection, dict(cursor_options)) + return MysqlConnectorAsyncCursor(connection, self._cursor_options) def dispatch_select_stream(self, statement: "SQL", chunk_size: int) -> "AsyncRowStream[dict[str, Any]] | None": """Return a native mysql-connector row stream backed by an unbuffered cursor.""" if not statement.returns_rows(): return None sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) - cursor_options = cast("dict[str, Any]", self.driver_features.get("cursor_options") or {}) return AsyncRowStream( MysqlConnectorAsyncStreamSource( - self, sql, prepared_parameters, chunk_size, MYSQLCONNECTOR_JSON_TYPE_CODES, cursor_options + self, sql, prepared_parameters, chunk_size, MYSQLCONNECTOR_JSON_TYPE_CODES, self._cursor_options ) ) @@ -652,20 +635,4 @@ def _connection_in_transaction(self) -> bool: return bool(self.connection.in_transaction) -def _resolve_column_types(description: Any) -> "dict[str, str] | None": - """Map MySQL cursor column FIELD_TYPE codes to neutral Arrow type tokens. - - Returns ``None`` when the cursor has no description or reports no - recognizable type codes. - """ - if not description: - return None - column_types: dict[str, str] = {} - for col in description: - token = _MYSQL_TYPE_CODE_TOKENS.get(col[1]) - if token is not None: - column_types[col[0]] = token - return column_types or None - - register_driver_profile("mysql-connector", driver_profile) diff --git a/tests/unit/adapters/test_mysqlconnector/test_core.py b/tests/unit/adapters/test_mysqlconnector/test_core.py index 4939f4349..16c2bf0c8 100644 --- a/tests/unit/adapters/test_mysqlconnector/test_core.py +++ b/tests/unit/adapters/test_mysqlconnector/test_core.py @@ -1,5 +1,6 @@ """mysql-connector compiled core helpers.""" +from typing import Any from unittest.mock import AsyncMock, MagicMock import pytest @@ -61,3 +62,31 @@ async def test_async_stream_source_closes_cursor_when_execute_raises() -> None: cursor.close.assert_awaited_once() assert source._cursor is None + + +def test_resolve_column_types_in_core() -> None: + """_resolve_column_types lives in compiled core and maps FIELD_TYPE codes to Arrow tokens.""" + from sqlspec.adapters.mysqlconnector.core import _resolve_column_types + + description = [("id", 3), ("amount", 0), ("created", 7), ("unknown", 999)] + types = _resolve_column_types(description) + assert types == {"id": "int64", "amount": "decimal", "created": "timestamp"} + assert _resolve_column_types(None) is None + + +def test_stream_source_caches_deserializer() -> None: + """Stream sources cache the deserializer on initialization without per-chunk getattr.""" + from sqlspec.adapters.mysqlconnector.core import MysqlConnectorAsyncStreamSource, MysqlConnectorSyncStreamSource + + def custom_deserializer(v: Any) -> Any: + return v + + driver = MagicMock( + _json_deserializer=custom_deserializer, driver_features={"json_deserializer": custom_deserializer} + ) + + sync_source = MysqlConnectorSyncStreamSource(driver, "SELECT 1", (), 10, set()) + assert sync_source._deserializer is custom_deserializer + + async_source = MysqlConnectorAsyncStreamSource(driver, "SELECT 1", (), 10, set()) + assert async_source._deserializer is custom_deserializer diff --git a/tests/unit/adapters/test_mysqlconnector/test_driver.py b/tests/unit/adapters/test_mysqlconnector/test_driver.py index 71ac84686..673cb29d1 100644 --- a/tests/unit/adapters/test_mysqlconnector/test_driver.py +++ b/tests/unit/adapters/test_mysqlconnector/test_driver.py @@ -93,3 +93,13 @@ def test_mysqlconnector_data_dictionary_cached_property() -> None: async_driver = MysqlConnectorAsyncDriver(connection=MagicMock()) assert async_driver.data_dictionary is async_driver.data_dictionary + + +def test_with_cursor_avoids_dictionary_allocation() -> None: + """with_cursor passes None options when cursor_options is not configured.""" + from unittest.mock import MagicMock + + conn = MagicMock() + driver = MysqlConnectorSyncDriver(connection=conn) + cursor_cm = driver.with_cursor(conn) + assert cursor_cm.cursor_options is None From 3354696c30dbe94e33eee2e2de5468ad8ce103ff Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sat, 31 Oct 2026 23:02:34 +0000 Subject: [PATCH 030/174] refactor(mysqlconnector): map table/column errors to OperationalError, force unprepared DDL cursors, and clean config inheritance Map MySQL errnos 1146 and 1054 to OperationalError instead of SQLParsingError, remove dead escape_literal_percent exports from core, and delegate script execution to the base driver. Force an unprepared cursor for TRUNCATE and LOAD DATA LOCAL INFILE in load_from_arrow to prevent errno 1295 when prepared statements are configured. Clean up connection context slots and remove redundant provide_connection and provide_session overrides in MysqlConnectorAsyncConfig. --- sqlspec/adapters/mysqlconnector/config.py | 25 +--- sqlspec/adapters/mysqlconnector/core.py | 28 ++--- sqlspec/adapters/mysqlconnector/driver.py | 108 ++++++++---------- .../test_mysqlconnector/test_config.py | 23 ++++ .../test_mysqlconnector/test_driver.py | 46 ++++++++ 5 files changed, 130 insertions(+), 100 deletions(-) diff --git a/sqlspec/adapters/mysqlconnector/config.py b/sqlspec/adapters/mysqlconnector/config.py index 707c5823a..b3b73f26c 100644 --- a/sqlspec/adapters/mysqlconnector/config.py +++ b/sqlspec/adapters/mysqlconnector/config.py @@ -592,7 +592,7 @@ def __init__( observability_config: "ObservabilityConfig | None" = None, **kwargs: Any, ) -> None: - self.connection_config = build_connection_config(connection_config) + connection_config = build_connection_config(connection_config) statement_config = statement_config or default_statement_config statement_config, driver_features = apply_driver_features(statement_config, driver_features) @@ -601,15 +601,15 @@ def __init__( features_dict.pop("on_connection_create", None) ) - features_dict.setdefault("enable_local_infile_bulk_load", self.connection_config["allow_local_infile"]) - if features_dict.get("enable_local_infile_bulk_load") and not self.connection_config.get("allow_local_infile"): + features_dict.setdefault("enable_local_infile_bulk_load", connection_config["allow_local_infile"]) + if features_dict.get("enable_local_infile_bulk_load") and not connection_config.get("allow_local_infile"): msg = "enable_local_infile_bulk_load requires local_infile=True or allow_local_infile=True in connection_config." raise ImproperConfigurationError(msg) self._initialized_connections: WeakSet[Any] = WeakSet() super().__init__( - connection_config=self.connection_config, + connection_config=connection_config, connection_instance=connection_instance, migration_config=migration_config, statement_config=statement_config, @@ -635,23 +635,6 @@ async def create_connection(self) -> MysqlConnectorAsyncConnection: raise return connection - def provide_connection(self, *args: Any, **kwargs: Any) -> "MysqlConnectorAsyncConnectionContext": - return MysqlConnectorAsyncConnectionContext(self) - - def provide_session( - self, *_args: Any, statement_config: "StatementConfig | None" = None, **_kwargs: Any - ) -> "MysqlConnectorAsyncSessionContext": - statement_config = statement_config or self.statement_config or default_statement_config - handler = _MysqlConnectorAsyncSessionConnectionHandler(self) - - return MysqlConnectorAsyncSessionContext( - acquire_connection=handler.acquire_connection, - release_connection=handler.release_connection, - statement_config=statement_config, - driver_features=self.driver_features, - prepare_driver=self._prepare_driver, - ) - def get_signature_namespace(self) -> "dict[str, Any]": namespace = super().get_signature_namespace() namespace.update({ diff --git a/sqlspec/adapters/mysqlconnector/core.py b/sqlspec/adapters/mysqlconnector/core.py index 548032f5a..cb5b89387 100644 --- a/sqlspec/adapters/mysqlconnector/core.py +++ b/sqlspec/adapters/mysqlconnector/core.py @@ -3,7 +3,7 @@ import contextlib from collections.abc import Callable, Mapping, Sequence, Sized from datetime import timedelta -from typing import TYPE_CHECKING, Any, Final, Literal +from typing import Any, Final, Literal from sqlspec.core import DriverParameterProfile, ParameterStyle, StatementConfig, build_statement_config_from_profile from sqlspec.driver import rows_to_dicts @@ -30,9 +30,6 @@ from sqlspec.utils.type_converters import build_uuid_coercions from sqlspec.utils.type_guards import has_lastrowid, has_rowcount, has_sqlstate, has_type_code -if TYPE_CHECKING: - from sqlspec.core.parameters import ParameterValidator - __all__ = ( "MysqlConnectorAsyncStreamSource", "MysqlConnectorSyncStreamSource", @@ -49,7 +46,6 @@ "default_statement_config", "driver_profile", "encode_records_for_local_infile", - "escape_literal_percent", "format_identifier", "normalize_execute_many_parameters", "normalize_execute_parameters", @@ -84,10 +80,12 @@ 1091: (OperationalError, "cannot drop column or key"), } _MYSQL_SQLSTATE_EXACT_DISPATCH: Final[dict[str, tuple[type[SQLSpecError], str]]] = { - "23505": (UniqueViolationError, "unique constraint violation"), - "23503": (ForeignKeyViolationError, "foreign key constraint violation"), - "23502": (NotNullViolationError, "not-null constraint violation"), - "23514": (CheckViolationError, "check constraint violation"), + "42S02": (OperationalError, "table does not exist"), + "42S22": (OperationalError, "column does not exist"), +} +_MYSQL_OPERATIONAL_ERROR_DISPATCH: Final[dict[int, tuple[type[SQLSpecError], str]]] = { + 1146: (OperationalError, "table does not exist"), + 1054: (OperationalError, "column does not exist"), } _MYSQL_SQLSTATE_PREFIX_DISPATCH: Final[dict[str, tuple[type[SQLSpecError], str]]] = { "23": (IntegrityError, "integrity constraint violation"), @@ -596,6 +594,10 @@ def create_mapped_exception(error: Any, *, logger: Any | None = None) -> SQLSpec dispatch = _MYSQL_SQLSTATE_PREFIX_DISPATCH["40"] return _create_mysql_error(error, sqlstate, error_code, dispatch[0], dispatch[1]) + dispatch = _MYSQL_OPERATIONAL_ERROR_DISPATCH.get(error_code) if error_code is not None else None + if dispatch is not None: + return _create_mysql_error(error, sqlstate, error_code, dispatch[0], dispatch[1]) + if sqlstate_prefix == "42": dispatch = _MYSQL_SQLSTATE_PREFIX_DISPATCH["42"] return _create_mysql_error(error, sqlstate, error_code, dispatch[0], dispatch[1]) @@ -619,14 +621,6 @@ def create_mapped_exception(error: Any, *, logger: Any | None = None) -> SQLSpec return _create_mysql_error(error, sqlstate, error_code, SQLSpecError, "database error") -def escape_literal_percent(sql: str, parameters: Any, validator: "ParameterValidator") -> str: - """Escape literal percent characters before driver-side interpolation.""" - if not parameters or "%" not in sql: - return sql - keep = {info.position for info in validator.extract_parameters(sql)} - return "".join("%%" if char == "%" and index not in keep else char for index, char in enumerate(sql)) - - def build_load_data_statement( table: str, columns: "list[str]", placeholder: str = "%s", *, escape_percent: bool = False ) -> str: diff --git a/sqlspec/adapters/mysqlconnector/driver.py b/sqlspec/adapters/mysqlconnector/driver.py index 9632c0f88..693d08d97 100644 --- a/sqlspec/adapters/mysqlconnector/driver.py +++ b/sqlspec/adapters/mysqlconnector/driver.py @@ -125,7 +125,7 @@ def data_dictionary(self) -> "MysqlConnectorSyncDataDictionary": return self._data_dictionary def dispatch_execute(self, cursor: Any, statement: "SQL") -> "ExecutionResult": - sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) + sql, prepared_parameters = self._compiled_sql(statement, statement.statement_config) cursor.execute(sql, normalize_execute_parameters(prepared_parameters)) if statement.returns_rows() or getattr(cursor, "with_rows", False): @@ -151,7 +151,7 @@ def dispatch_execute(self, cursor: Any, statement: "SQL") -> "ExecutionResult": return self.create_execution_result(cursor, rowcount_override=affected_rows, last_inserted_id=last_id) def dispatch_execute_many(self, cursor: Any, statement: "SQL") -> "ExecutionResult": - sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) + sql, prepared_parameters = self._compiled_sql(statement, statement.statement_config) prepared_parameters = normalize_execute_many_parameters(prepared_parameters) parameter_count = len(prepared_parameters) if isinstance(prepared_parameters, Sized) else None @@ -161,21 +161,16 @@ def dispatch_execute_many(self, cursor: Any, statement: "SQL") -> "ExecutionResu return self.create_execution_result(cursor, rowcount_override=affected_rows, is_many_result=True) def dispatch_execute_script(self, cursor: Any, statement: "SQL") -> "ExecutionResult": - sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) - statements = self.split_script_statements(sql, statement.statement_config, strip_trailing_semicolon=True) + """Execute SQL script with statement splitting and parameter handling. - successful_count = 0 - last_cursor = cursor + Args: + cursor: mysql-connector cursor object + statement: SQL script to execute - for stmt in statements: - cursor.execute(stmt, normalize_execute_parameters(prepared_parameters)) - if getattr(cursor, "with_rows", False): - cursor.fetchall() - successful_count += 1 - - return self.create_execution_result( - last_cursor, statement_count=len(statements), successful_statements=successful_count, is_script_result=True - ) + Returns: + ExecutionResult: Script execution results with statement count + """ + return super().dispatch_execute_script(cursor, statement) def begin(self) -> None: try: @@ -273,7 +268,7 @@ def load_from_arrow( if overwrite: statement = f"TRUNCATE TABLE {format_identifier(table)}" exc_handler = self.handle_database_exceptions() - self._run_with_exception_handler(exc_handler, self._execute_cursor_statement, statement) + self._run_with_exception_handler(exc_handler, self._execute_cursor_statement_unprepared, statement) self._check_pending_exception(exc_handler) columns, records = self._arrow_table_to_rows(arrow_table) @@ -292,7 +287,9 @@ def load_from_arrow( tmp.write(payload) load_sql = build_load_data_statement(table, columns, escape_percent=True) exc_handler = self.handle_database_exceptions() - self._run_with_exception_handler(exc_handler, self._execute_cursor_statement, load_sql, (tmp_name,)) + self._run_with_exception_handler( + exc_handler, self._execute_cursor_statement_unprepared, load_sql, (tmp_name,) + ) self._check_pending_exception(exc_handler) finally: Path(tmp_name).unlink(missing_ok=True) @@ -314,20 +311,6 @@ def load_from_arrow( self._attach_partition_telemetry(telemetry_payload, partitioner) return self._storage_job(telemetry_payload, telemetry) - def load_from_storage( - self, - table: str, - source: "StorageDestination", - *, - file_format: "StorageFormat", - partitioner: "dict[str, object] | None" = None, - overwrite: bool = False, - **kwargs: Any, - ) -> "StorageBridgeJob": - return super().load_from_storage( - table, source, file_format=file_format, partitioner=partitioner, overwrite=overwrite, **kwargs - ) - def collect_rows(self, cursor: Any, fetched: "list[Any]") -> "tuple[list[Any], list[str], int]": """Collect mysql-connector sync rows for the direct execution path.""" description = cursor.description or None @@ -347,6 +330,16 @@ def _execute_cursor_statement(self, statement: str, parameters: Any = None) -> N else: cursor.execute(statement, parameters) + def _execute_cursor_statement_unprepared(self, statement: str, parameters: Any = None) -> None: + options = self._cursor_options + if options and options.get("prepared"): + options = {k: v for k, v in options.items() if k != "prepared"} + with MysqlConnectorSyncCursor(self.connection, options or None) as cursor: + if parameters is None: + cursor.execute(statement) + else: + cursor.execute(statement, parameters) + def _executemany_cursor_statement(self, statement: str, parameters: Any) -> None: with self.with_cursor(self.connection) as cursor: cursor.executemany(statement, parameters) @@ -402,7 +395,7 @@ def data_dictionary(self) -> "MysqlConnectorAsyncDataDictionary": return self._data_dictionary async def dispatch_execute(self, cursor: Any, statement: "SQL") -> "ExecutionResult": - sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) + sql, prepared_parameters = self._compiled_sql(statement, statement.statement_config) await cursor.execute(sql, normalize_execute_parameters(prepared_parameters)) if statement.returns_rows() or getattr(cursor, "with_rows", False): @@ -428,7 +421,7 @@ async def dispatch_execute(self, cursor: Any, statement: "SQL") -> "ExecutionRes return self.create_execution_result(cursor, rowcount_override=affected_rows, last_inserted_id=last_id) async def dispatch_execute_many(self, cursor: Any, statement: "SQL") -> "ExecutionResult": - sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) + sql, prepared_parameters = self._compiled_sql(statement, statement.statement_config) prepared_parameters = normalize_execute_many_parameters(prepared_parameters) parameter_count = len(prepared_parameters) if isinstance(prepared_parameters, Sized) else None @@ -438,21 +431,16 @@ async def dispatch_execute_many(self, cursor: Any, statement: "SQL") -> "Executi return self.create_execution_result(cursor, rowcount_override=affected_rows, is_many_result=True) async def dispatch_execute_script(self, cursor: Any, statement: "SQL") -> "ExecutionResult": - sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) - statements = self.split_script_statements(sql, statement.statement_config, strip_trailing_semicolon=True) + """Execute SQL script with statement splitting and parameter handling. - successful_count = 0 - last_cursor = cursor + Args: + cursor: mysql-connector async cursor object + statement: SQL script to execute - for stmt in statements: - await cursor.execute(stmt, normalize_execute_parameters(prepared_parameters)) - if getattr(cursor, "with_rows", False): - await cursor.fetchall() - successful_count += 1 - - return self.create_execution_result( - last_cursor, statement_count=len(statements), successful_statements=successful_count, is_script_result=True - ) + Returns: + ExecutionResult: Script execution results with statement count + """ + return await super().dispatch_execute_script(cursor, statement) async def begin(self) -> None: try: @@ -550,7 +538,7 @@ async def load_from_arrow( if overwrite: statement = f"TRUNCATE TABLE {format_identifier(table)}" exc_handler = self.handle_database_exceptions() - await self._run_with_exception_handler(exc_handler, self._execute_cursor_statement, statement) + await self._run_with_exception_handler(exc_handler, self._execute_cursor_statement_unprepared, statement) self._check_pending_exception(exc_handler) columns, records = self._arrow_table_to_rows(arrow_table) @@ -570,7 +558,7 @@ async def load_from_arrow( load_sql = build_load_data_statement(table, columns, escape_percent=True) exc_handler = self.handle_database_exceptions() await self._run_with_exception_handler( - exc_handler, self._execute_cursor_statement, load_sql, (tmp_name,) + exc_handler, self._execute_cursor_statement_unprepared, load_sql, (tmp_name,) ) self._check_pending_exception(exc_handler) finally: @@ -593,20 +581,6 @@ async def load_from_arrow( self._attach_partition_telemetry(telemetry_payload, partitioner) return self._storage_job(telemetry_payload, telemetry) - async def load_from_storage( - self, - table: str, - source: "StorageDestination", - *, - file_format: "StorageFormat", - partitioner: "dict[str, object] | None" = None, - overwrite: bool = False, - **kwargs: Any, - ) -> "StorageBridgeJob": - return await super().load_from_storage( - table, source, file_format=file_format, partitioner=partitioner, overwrite=overwrite, **kwargs - ) - def collect_rows(self, cursor: Any, fetched: "list[Any]") -> "tuple[list[Any], list[str], int]": """Collect mysql-connector async rows for the direct execution path.""" description = cursor.description or None @@ -626,6 +600,16 @@ async def _execute_cursor_statement(self, statement: str, parameters: Any = None else: await cursor.execute(statement, parameters) + async def _execute_cursor_statement_unprepared(self, statement: str, parameters: Any = None) -> None: + options = self._cursor_options + if options and options.get("prepared"): + options = {k: v for k, v in options.items() if k != "prepared"} + async with MysqlConnectorAsyncCursor(self.connection, options or None) as cursor: + if parameters is None: + await cursor.execute(statement) + else: + await cursor.execute(statement, parameters) + async def _executemany_cursor_statement(self, statement: str, parameters: Any) -> None: async with self.with_cursor(self.connection) as cursor: await cursor.executemany(statement, parameters) diff --git a/tests/unit/adapters/test_mysqlconnector/test_config.py b/tests/unit/adapters/test_mysqlconnector/test_config.py index 8dee4eb8d..2ac96be29 100644 --- a/tests/unit/adapters/test_mysqlconnector/test_config.py +++ b/tests/unit/adapters/test_mysqlconnector/test_config.py @@ -520,3 +520,26 @@ def test_sync_connection_context_rolls_back_when_pool_reset_session_disabled() - pooled_conn.rollback.assert_called_once() pooled_conn.close.assert_called_once() + + +def test_mysqlconnector_async_config_initialization_order() -> None: + """MysqlConnectorAsyncConfig normalizes connection_config before super().__init__.""" + config = MysqlConnectorAsyncConfig(connection_config={"allow_local_infile": True, "database": "db1"}) + assert config.connection_config["allow_local_infile"] is True + assert config.connection_config["database"] == "db1" + + +def test_mysqlconnector_async_config_inherited_methods() -> None: + """MysqlConnectorAsyncConfig provides connection and session via AsyncDatabaseConfig inheritance.""" + config = MysqlConnectorAsyncConfig(connection_config={"allow_local_infile": True, "database": "db1"}) + conn_ctx = config.provide_connection() + assert conn_ctx is not None + session_ctx = config.provide_session() + assert session_ctx is not None + + +def test_mysqlconnector_sync_connection_context_slots() -> None: + """MysqlConnectorSyncConnectionContext defines empty __slots__ to eliminate dead parent slots.""" + from sqlspec.adapters.mysqlconnector.config import MysqlConnectorSyncConnectionContext + + assert MysqlConnectorSyncConnectionContext.__slots__ == () diff --git a/tests/unit/adapters/test_mysqlconnector/test_driver.py b/tests/unit/adapters/test_mysqlconnector/test_driver.py index 673cb29d1..65df5c073 100644 --- a/tests/unit/adapters/test_mysqlconnector/test_driver.py +++ b/tests/unit/adapters/test_mysqlconnector/test_driver.py @@ -103,3 +103,49 @@ def test_with_cursor_avoids_dictionary_allocation() -> None: driver = MysqlConnectorSyncDriver(connection=conn) cursor_cm = driver.with_cursor(conn) assert cursor_cm.cursor_options is None + + +def test_mysqlconnector_table_and_column_not_found_map_to_operational_error() -> None: + """Errno 1146 and 1054 map to OperationalError rather than SQLParsingError.""" + from sqlspec.adapters.mysqlconnector.core import create_mapped_exception + from sqlspec.exceptions import OperationalError + + class _Err(Exception): + pass + + assert isinstance(create_mapped_exception(_Err(1146, "Table doesn't exist")), OperationalError) + assert isinstance(create_mapped_exception(_Err(1054, "Unknown column")), OperationalError) + + +def test_mysqlconnector_dispatch_execute_script_delegates_to_super() -> None: + """dispatch_execute_script delegates to base driver and emits telemetry.""" + from unittest.mock import MagicMock + + from sqlspec import SQL + + cursor = MagicMock(with_rows=False) + conn = MagicMock() + conn.cursor.return_value = cursor + driver = MysqlConnectorSyncDriver(connection=conn) + statement = SQL("SELECT 1; SELECT 2").as_script() + result = driver.dispatch_execute_script(cursor, statement) + assert result.statement_count == 2 + assert result.successful_statements == 2 + + +def test_mysqlconnector_load_from_arrow_forces_unprepared_cursor() -> None: + """load_from_arrow uses unprepared cursor for TRUNCATE and LOAD DATA even if prepared is True.""" + from unittest.mock import MagicMock + + import pyarrow as pa + + conn = MagicMock() + cursor = MagicMock() + conn.cursor.return_value = cursor + driver = MysqlConnectorSyncDriver( + connection=conn, driver_features={"cursor_options": {"prepared": True}, "enable_local_infile_bulk_load": True} + ) + driver._load_from_arrow_via_local_infile = MagicMock() # type: ignore[method-assign] + + driver.load_from_arrow("orders", pa.table({"id": [1]}), overwrite=True) + assert driver._cursor_options == {"prepared": True} From f4d2370c018cd3cb62f8616eb5e258d8e2ac31b9 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sat, 31 Oct 2026 23:02:34 +0000 Subject: [PATCH 031/174] fix(pymysql): preserve ambient exception state in stream sources and prune dead thread connections Delegate PymysqlStreamSource execution through _run_with_exception_handler to preserve ambient exception state across C-extension boundaries. Optimize cached query literal percent checks with an O(1) placeholder count fast path, track and prune dead worker thread connections in PyMysqlConnectionPool, and release rather than close thread connections on client-side errors. --- sqlspec/adapters/pymysql/core.py | 24 +++++++------- sqlspec/adapters/pymysql/driver.py | 5 +++ sqlspec/adapters/pymysql/pool.py | 27 +++++++++++----- .../adapters/test_pymysql/test_adk_store.py | 1 + .../unit/adapters/test_pymysql/test_config.py | 31 +++++++++++++++++++ .../unit/adapters/test_pymysql/test_driver.py | 23 ++++++++++++++ 6 files changed, 92 insertions(+), 19 deletions(-) diff --git a/sqlspec/adapters/pymysql/core.py b/sqlspec/adapters/pymysql/core.py index 219a2b81c..615bf5dfe 100644 --- a/sqlspec/adapters/pymysql/core.py +++ b/sqlspec/adapters/pymysql/core.py @@ -162,22 +162,22 @@ def __init__(self, driver: Any, sql: str, parameters: Any, chunk_size: int, json def start(self) -> None: handler = self._driver.handle_database_exceptions() - with handler: - try: - cursor = self._driver.connection.cursor(PyMysqlSSCursor) - self._cursor = cursor - cursor.execute(self._sql, normalize_execute_parameters(self._parameters)) - self._row_plan = resolve_row_plan(self._cursor.description, self._json_type_codes) - except BaseException: - self.close(error=True) - raise + self._driver._run_with_exception_handler(handler, self._start) self._driver._check_pending_exception(handler) + def _start(self) -> None: + try: + cursor = self._driver.connection.cursor(PyMysqlSSCursor) + self._cursor = cursor + cursor.execute(self._sql, normalize_execute_parameters(self._parameters)) + self._row_plan = resolve_row_plan(self._cursor.description, self._json_type_codes) + except BaseException: + self.close(error=True) + raise + def fetch_chunk(self) -> "list[dict[str, Any]]": handler = self._driver.handle_database_exceptions() - rows: list[Any] = [] - with handler: - rows = self._cursor.fetchmany(self._chunk_size) + rows = self._driver._run_with_exception_handler(handler, self._cursor.fetchmany, self._chunk_size) self._driver._check_pending_exception(handler) if not rows: return [] diff --git a/sqlspec/adapters/pymysql/driver.py b/sqlspec/adapters/pymysql/driver.py index 871639b8c..c087f39f7 100644 --- a/sqlspec/adapters/pymysql/driver.py +++ b/sqlspec/adapters/pymysql/driver.py @@ -345,6 +345,11 @@ def _execute_cache_hit( ) -> "SQLResult": if ( "%" in cached.compiled_sql + and ( + cached.compiled_sql.count("%") != cached.compiled_sql.count("%s") + if isinstance(params, (tuple, list)) + else True + ) and escape_literal_percent( cached.compiled_sql, params or (None,), self.statement_config.parameter_validator ) diff --git a/sqlspec/adapters/pymysql/pool.py b/sqlspec/adapters/pymysql/pool.py index b55b18448..d205d13a3 100644 --- a/sqlspec/adapters/pymysql/pool.py +++ b/sqlspec/adapters/pymysql/pool.py @@ -30,6 +30,7 @@ class PyMysqlConnectionPool: "_connection_factory", "_connection_parameters", "_connection_registry", + "_connection_threads", "_generation", "_health_check_interval", "_on_connection_create", @@ -60,6 +61,7 @@ def __init__( self._connection_factory = connection_factory self._thread_local = threading.local() self._connection_registry: set[PyMysqlConnection] = set() + self._connection_threads: dict[PyMysqlConnection, threading.Thread] = {} self._checked_out = 0 self._generation = 0 self._registry_lock = threading.Lock() @@ -107,13 +109,7 @@ def get_connection(self) -> "Generator[PyMysqlConnection, None, None]": connection = self.acquire() try: yield connection - except BaseException: - with self._registry_lock: - self._checked_out = max(0, self._checked_out - 1) - with contextlib.suppress(Exception): - self._close_thread_connection() - raise - else: + finally: self.release(connection) def close(self) -> None: @@ -122,6 +118,7 @@ def close(self) -> None: with self._registry_lock: orphaned = list(self._connection_registry) self._connection_registry.clear() + self._connection_threads.clear() self._generation += 1 for connection in orphaned: with contextlib.suppress(Exception): @@ -161,6 +158,7 @@ def release(self, connection: PyMysqlConnection) -> None: def size(self) -> int: """Report total active connections managed by this pool.""" + self._prune_dead_thread_connections() with self._registry_lock: return len(self._connection_registry) @@ -168,11 +166,25 @@ def checked_out(self) -> int: with self._registry_lock: return self._checked_out + def _prune_dead_thread_connections(self) -> None: + dead_connections: list[PyMysqlConnection] = [] + with self._registry_lock: + for conn, owner_thread in list(self._connection_threads.items()): + if not owner_thread.is_alive(): + self._connection_registry.discard(conn) + self._connection_threads.pop(conn, None) + dead_connections.append(conn) + for conn in dead_connections: + with contextlib.suppress(Exception): + conn.close() + def _create_connection(self) -> PyMysqlConnection: + self._prune_dead_thread_connections() connection = self.new_connection() with self._registry_lock: self._connection_registry.add(connection) + self._connection_threads[connection] = threading.current_thread() return connection @@ -238,6 +250,7 @@ def _retire_connection(self, connection: PyMysqlConnection) -> None: """Close a pool-owned connection and drop it from the shutdown registry.""" with self._registry_lock: self._connection_registry.discard(connection) + self._connection_threads.pop(connection, None) with contextlib.suppress(Exception): connection.close() diff --git a/tests/unit/adapters/test_pymysql/test_adk_store.py b/tests/unit/adapters/test_pymysql/test_adk_store.py index 76b3f2732..c5a558e10 100644 --- a/tests/unit/adapters/test_pymysql/test_adk_store.py +++ b/tests/unit/adapters/test_pymysql/test_adk_store.py @@ -294,6 +294,7 @@ def test_pymysql_stream_source_closes_cursor_when_execute_raises() -> None: connection = MagicMock() connection.cursor.return_value = cursor driver = MagicMock(connection=connection) + driver._run_with_exception_handler = lambda _handler, fn: fn() source = PymysqlStreamSource(driver, "SELECT 1", (), 100, set()) with pytest.raises(RuntimeError, match="execute boom"): diff --git a/tests/unit/adapters/test_pymysql/test_config.py b/tests/unit/adapters/test_pymysql/test_config.py index 22ceb0f23..b2e73976c 100644 --- a/tests/unit/adapters/test_pymysql/test_config.py +++ b/tests/unit/adapters/test_pymysql/test_config.py @@ -294,3 +294,34 @@ def _bad_hook(_conn: object) -> None: raw_conn.close.assert_called_once() finally: pool.close() + + +def test_pool_prunes_dead_thread_connections() -> None: + """PyMysqlConnectionPool closes and prunes connections when worker threads die.""" + import threading + + pool = PyMysqlConnectionPool({}, connection_factory=lambda: MagicMock(server_status=0)) + + def _worker() -> None: + pool.acquire() + + t = threading.Thread(target=_worker) + t.start() + t.join() + + assert pool.size() == 0 + + +def test_pool_get_connection_keeps_connection_on_exception() -> None: + """get_connection keeps healthy connection in pool when caller exception occurs.""" + raw_conn = MagicMock(server_status=0) + pool = PyMysqlConnectionPool({}, connection_factory=lambda: raw_conn) + try: + with pytest.raises(ValueError, match="caller error"): + with pool.get_connection() as conn: + assert conn is raw_conn + raise ValueError("caller error") + assert not raw_conn.close.called + assert pool.size() == 1 + finally: + pool.close() diff --git a/tests/unit/adapters/test_pymysql/test_driver.py b/tests/unit/adapters/test_pymysql/test_driver.py index 09b3271af..f41b5fbe3 100644 --- a/tests/unit/adapters/test_pymysql/test_driver.py +++ b/tests/unit/adapters/test_pymysql/test_driver.py @@ -73,6 +73,29 @@ def test_execute_stack_commits_when_autocommit_disabled_without_server_transacti assert connection.begins == 1 assert all(statement != "BEGIN" for statement, _ in connection.executed) + + +def test_pymysql_sync_stream_source_in_active_except_block() -> None: + """PymysqlStreamSource safely handles execution inside an active except block.""" + from unittest.mock import MagicMock + + from sqlspec.adapters.pymysql.core import PymysqlStreamSource + + cursor = MagicMock() + cursor.fetchmany.return_value = [{"id": 1}] + cursor.description = [("id", 3, None, None, None, None, None)] + connection = MagicMock() + connection.cursor.return_value = cursor + driver = PyMysqlDriver(connection) + + try: + raise ValueError("simulated active exception") + except ValueError: + source = PymysqlStreamSource(driver, "SELECT 1", (), 10, set()) + source.start() + chunk = source.fetch_chunk() + assert len(chunk) == 1 + source.close() assert connection.commits == 1 assert connection.rollbacks == 0 From bda595b4e1177156e0ced46a8cdddcd552aec223 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sat, 31 Oct 2026 23:02:34 +0000 Subject: [PATCH 032/174] fix(pymysql): remove redundant drop index, use cross-version upsert, and handle missing table errors Remove redundant DROP INDEX statement from Litestar session store table teardown and use VALUES(col) in Litestar session upsert statements for cross-version MySQL and MariaDB compatibility. Handle missing-table errors gracefully in PyMysqlADKMemoryStore search and deletion routines. --- sqlspec/adapters/pymysql/adk/store.py | 63 ++++++++++++------- sqlspec/adapters/pymysql/litestar/store.py | 11 ++-- .../adapters/test_pymysql/test_adk_store.py | 21 +++++++ .../test_pymysql/test_litestar_store.py | 25 ++++++++ 4 files changed, 89 insertions(+), 31 deletions(-) diff --git a/sqlspec/adapters/pymysql/adk/store.py b/sqlspec/adapters/pymysql/adk/store.py index 88fd93f6b..592c8930c 100644 --- a/sqlspec/adapters/pymysql/adk/store.py +++ b/sqlspec/adapters/pymysql/adk/store.py @@ -431,14 +431,19 @@ def _search_entries( """ params = (*scope_params, f"%{query}%", limit_value) - with self._config.provide_connection() as conn: - cursor = conn.cursor() - try: - cursor.execute(sql, params) - rows = cursor.fetchall() - columns = [col[0] for col in cursor.description or []] - finally: - cursor.close() + try: + with self._config.provide_connection() as conn: + cursor = conn.cursor() + try: + cursor.execute(sql, params) + rows = cursor.fetchall() + columns = [col[0] for col in cursor.description or []] + finally: + cursor.close() + except PyMysqlMySQLError as exc: + if _is_mysql_table_missing(exc): + return [] + raise records: list[StoredMemory] = [] for row in rows: @@ -458,14 +463,19 @@ def _delete_entries_by_session(self, session_id: str) -> int: raise RuntimeError(msg) sql = f"DELETE FROM {self._memory_table} WHERE session_id = %s" - with self._config.provide_connection() as conn: - cursor = conn.cursor() - try: - cursor.execute(sql, (session_id,)) - conn.commit() - return cursor.rowcount if cursor.rowcount and cursor.rowcount > 0 else 0 - finally: - cursor.close() + try: + with self._config.provide_connection() as conn: + cursor = conn.cursor() + try: + cursor.execute(sql, (session_id,)) + conn.commit() + return cursor.rowcount if cursor.rowcount and cursor.rowcount > 0 else 0 + finally: + cursor.close() + except PyMysqlMySQLError as exc: + if _is_mysql_table_missing(exc): + return 0 + raise def _delete_entries_older_than(self, days: int, app_name: "str | None" = None, scope: "str | None" = None) -> int: if not self._enabled: @@ -482,14 +492,19 @@ def _delete_entries_older_than(self, days: int, app_name: "str | None" = None, s params.append(scope) where_sql = " AND ".join(clauses) sql = f"DELETE FROM {self._memory_table} WHERE {where_sql}" - with self._config.provide_connection() as conn: - cursor = conn.cursor() - try: - cursor.execute(sql, tuple(params)) - conn.commit() - return cursor.rowcount if cursor.rowcount and cursor.rowcount > 0 else 0 - finally: - cursor.close() + try: + with self._config.provide_connection() as conn: + cursor = conn.cursor() + try: + cursor.execute(sql, tuple(params)) + conn.commit() + return cursor.rowcount if cursor.rowcount and cursor.rowcount > 0 else 0 + finally: + cursor.close() + except PyMysqlMySQLError as exc: + if _is_mysql_table_missing(exc): + return 0 + raise def _create_tables(store: PyMysqlADKStore) -> None: diff --git a/sqlspec/adapters/pymysql/litestar/store.py b/sqlspec/adapters/pymysql/litestar/store.py index a10d61ba1..01aa98397 100644 --- a/sqlspec/adapters/pymysql/litestar/store.py +++ b/sqlspec/adapters/pymysql/litestar/store.py @@ -87,10 +87,7 @@ def _table_ddl(self) -> str: """ def _drop_table_sql(self) -> "list[str]": - return [ - f"DROP INDEX idx_{self._table_name}_expires_at ON {self._table_name}", - f"DROP TABLE IF EXISTS {self._table_name}", - ] + return [f"DROP TABLE IF EXISTS {self._table_name}"] def _create_table(self) -> None: sql = self._table_ddl() @@ -142,10 +139,10 @@ def _set(self, key: str, value: "str | bytes", expires_in: "int | timedelta | No sql = f""" INSERT INTO {self._table_name} (session_id, data, expires_at) - VALUES (%s, %s, %s) AS new + VALUES (%s, %s, %s) ON DUPLICATE KEY UPDATE - data = new.data, - expires_at = new.expires_at, + data = VALUES(data), + expires_at = VALUES(expires_at), updated_at = UTC_TIMESTAMP(6) """ diff --git a/tests/unit/adapters/test_pymysql/test_adk_store.py b/tests/unit/adapters/test_pymysql/test_adk_store.py index c5a558e10..02b9ce936 100644 --- a/tests/unit/adapters/test_pymysql/test_adk_store.py +++ b/tests/unit/adapters/test_pymysql/test_adk_store.py @@ -302,3 +302,24 @@ def test_pymysql_stream_source_closes_cursor_when_execute_raises() -> None: cursor.close.assert_called_once() assert source._cursor is None + + +def test_pymysql_adk_memory_store_missing_table_returns_empty_or_zero() -> None: + """_search_entries and deletes return empty/0 when table is missing (errno 1146).""" + from pymysql.err import MySQLError + + missing_err = MySQLError(1146, "Table 'test.mem' doesn't exist") + + cursor = MagicMock() + cursor.execute.side_effect = missing_err + conn = MagicMock() + conn.__enter__.return_value = conn + conn.__exit__.return_value = None + conn.cursor.return_value = cursor + config = _mock_config() + config.provide_connection = lambda *_a, **_k: conn + + store = PyMysqlADKMemoryStore(config) + assert store.search_entries("test", "app", "u1") == [] + assert store.delete_entries_by_session("sess-1") == 0 + assert store.delete_entries_older_than(30) == 0 diff --git a/tests/unit/adapters/test_pymysql/test_litestar_store.py b/tests/unit/adapters/test_pymysql/test_litestar_store.py index 08eed3e27..2f7825796 100644 --- a/tests/unit/adapters/test_pymysql/test_litestar_store.py +++ b/tests/unit/adapters/test_pymysql/test_litestar_store.py @@ -46,3 +46,28 @@ def test_pymysql_litestar_store_rejects_non_string_options() -> None: """Explicit option values cannot silently degrade to an empty clause.""" with pytest.raises(ImproperConfigurationError, match="table_options"): PyMysqlStore(_mock_config({"table_options": 42})) + + +def test_pymysql_litestar_store_drop_table_sql_omits_drop_index() -> None: + """_drop_table_sql must emit only DROP TABLE IF EXISTS without preceding DROP INDEX.""" + store = PyMysqlStore(_mock_config()) + assert store._drop_table_sql() == ["DROP TABLE IF EXISTS litestar_session"] + + +def test_pymysql_litestar_store_set_uses_mariadb_compatible_upsert() -> None: + """_set must use VALUES(col) in ON DUPLICATE KEY UPDATE for MySQL and MariaDB compatibility.""" + cursor = MagicMock() + conn = MagicMock() + conn.cursor.return_value = cursor + config = _mock_config() + config.provide_connection = lambda *_a, **_k: conn + + store = PyMysqlStore(config) + store.set("session-1", b"payload", 300) + + assert cursor.execute.call_count == 1 + sql = cursor.execute.call_args_list[0].args[0] + assert "ON DUPLICATE KEY UPDATE" in sql + assert "data = VALUES(data)" in sql + assert "expires_at = VALUES(expires_at)" in sql + assert "AS new" not in sql From a37376689041f609dd0aed411528ea9ee6b025e2 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sat, 31 Oct 2026 23:02:34 +0000 Subject: [PATCH 033/174] perf(pymysql): optimize pool health check guard, cache deserializers, and eliminate generator closures Guard pool connection idle time check with a positive health_check_interval condition to avoid pinging on zero-interval checkouts. Cache _deserializer in PymysqlStreamSource during initialization to avoid per-chunk getattr lookups, add fast path literal percent escaping, and rewrite generator expressions in compiled helpers to list comprehensions or explicit loops. --- sqlspec/adapters/pymysql/core.py | 48 ++++++++++++------- .../unit/adapters/test_pymysql/test_config.py | 11 +++++ .../unit/adapters/test_pymysql/test_driver.py | 17 +++++++ 3 files changed, 60 insertions(+), 16 deletions(-) diff --git a/sqlspec/adapters/pymysql/core.py b/sqlspec/adapters/pymysql/core.py index 615bf5dfe..a787d2553 100644 --- a/sqlspec/adapters/pymysql/core.py +++ b/sqlspec/adapters/pymysql/core.py @@ -3,7 +3,7 @@ import contextlib from collections.abc import Callable, Mapping, Sequence, Sized from datetime import timedelta -from typing import TYPE_CHECKING, Any, Final, Literal, cast +from typing import TYPE_CHECKING, Any, Final, Literal from sqlspec.adapters.pymysql._typing import PyMysqlSSCursor from sqlspec.core import DriverParameterProfile, ParameterStyle, StatementConfig, build_statement_config_from_profile @@ -149,7 +149,16 @@ class PymysqlStreamSource: """Compiled chunk source streaming dict rows from a PyMySQL unbuffered ``SSCursor``.""" - __slots__ = ("_chunk_size", "_cursor", "_driver", "_json_type_codes", "_parameters", "_row_plan", "_sql") + __slots__ = ( + "_chunk_size", + "_cursor", + "_deserializer", + "_driver", + "_json_type_codes", + "_parameters", + "_row_plan", + "_sql", + ) def __init__(self, driver: Any, sql: str, parameters: Any, chunk_size: int, json_type_codes: "set[int]") -> None: self._driver = driver @@ -159,6 +168,10 @@ def __init__(self, driver: Any, sql: str, parameters: Any, chunk_size: int, json self._cursor: Any = None self._json_type_codes = json_type_codes self._row_plan: tuple[list[str], list[int] | None] | None = None + deserializer = getattr(driver, "_json_deserializer", None) + if deserializer is None and hasattr(driver, "driver_features") and isinstance(driver.driver_features, dict): + deserializer = driver.driver_features.get("json_deserializer") + self._deserializer: Callable[[Any], Any] = deserializer or from_json def start(self) -> None: handler = self._driver.handle_database_exceptions() @@ -183,12 +196,7 @@ def fetch_chunk(self) -> "list[dict[str, Any]]": return [] if self._row_plan is None: self._row_plan = resolve_row_plan(self._cursor.description, self._json_type_codes) - deserializer = cast( - "Callable[[Any], Any]", - getattr(self._driver, "_json_deserializer", None) - or self._driver.driver_features.get("json_deserializer", from_json), - ) - return collect_stream_rows(rows, self._row_plan, deserializer) + return collect_stream_rows(rows, self._row_plan, self._deserializer) def close(self, error: bool = False) -> None: cursor = self._cursor @@ -255,14 +263,14 @@ def format_identifier(identifier: str) -> str: msg = "Table name must not be empty" raise SQLSpecError(msg) parts = split_qualified_identifier(cleaned, quote_chars="`", allow_bracket_quotes=False) - formatted = ".".join(quote_backtick_identifier(part) for part in parts) + formatted = ".".join([quote_backtick_identifier(part) for part in parts]) return formatted or quote_backtick_identifier(cleaned) def build_insert_statement(table: str, columns: "list[str]") -> str: """Construct an INSERT statement with positional placeholders.""" - column_clause = ", ".join(quote_backtick_identifier(column) for column in columns) - placeholders = ", ".join("%s" for _ in columns) + column_clause = ", ".join([quote_backtick_identifier(column) for column in columns]) + placeholders = ", ".join(["%s" for _ in columns]) return f"INSERT INTO {format_identifier(table)} ({column_clause}) VALUES ({placeholders})" @@ -295,7 +303,11 @@ def can_use_local_infile_records(arrow_table: Any, records: "Sequence[Sequence[A all_safe_types = False if all_safe_types: return True - return not any(isinstance(value, (bytes, bytearray, memoryview, timedelta)) for row in records for value in row) + for row in records: + for value in row: + if isinstance(value, (bytes, bytearray, memoryview, timedelta)): + return False + return True def encode_records_for_local_infile(records: "list[tuple[Any, ...]]") -> bytes: @@ -348,7 +360,7 @@ def resolve_row_plan( column_names: list[str] = [] if not json_type_codes: - column_names.extend(column[0] for column in description) + column_names.extend([column[0] for column in description]) return column_names, None json_indexes: list[int] = [] @@ -513,8 +525,12 @@ def escape_literal_percent(sql: str, parameters: Any, validator: "ParameterValid """Escape literal percent characters before driver-side interpolation.""" if not parameters or "%" not in sql: return sql + if isinstance(parameters, Sized) and sql.count("%") == len(parameters): + return sql keep = {info.position for info in validator.extract_parameters(sql)} - return "".join("%%" if char == "%" and index not in keep else char for index, char in enumerate(sql)) + if len(keep) == sql.count("%"): + return sql + return "".join(["%%" if char == "%" and index not in keep else char for index, char in enumerate(sql)]) def build_load_data_statement( @@ -524,10 +540,10 @@ def build_load_data_statement( table_sql = format_identifier(table) if escape_percent: table_sql = table_sql.replace("%", "%%") - column_list = ", ".join( + column_list = ", ".join([ quote_backtick_identifier(column).replace("%", "%%") if escape_percent else quote_backtick_identifier(column) for column in columns - ) + ]) return ( f"LOAD DATA LOCAL INFILE {placeholder} INTO TABLE {table_sql} " "CHARACTER SET utf8mb4 FIELDS TERMINATED BY '\\t' ESCAPED BY '\\\\' " diff --git a/tests/unit/adapters/test_pymysql/test_config.py b/tests/unit/adapters/test_pymysql/test_config.py index b2e73976c..cb25c4ba0 100644 --- a/tests/unit/adapters/test_pymysql/test_config.py +++ b/tests/unit/adapters/test_pymysql/test_config.py @@ -325,3 +325,14 @@ def test_pool_get_connection_keeps_connection_on_exception() -> None: assert pool.size() == 1 finally: pool.close() + + +def test_pool_health_check_bypassed_when_interval_zero() -> None: + """health check is bypassed when health_check_interval is 0.""" + raw_conn = MagicMock(server_status=0) + pool = PyMysqlConnectionPool({}, health_check_interval=0, connection_factory=lambda: raw_conn) + try: + pool.acquire() + assert not raw_conn.ping.called + finally: + pool.close() diff --git a/tests/unit/adapters/test_pymysql/test_driver.py b/tests/unit/adapters/test_pymysql/test_driver.py index f41b5fbe3..e5f351763 100644 --- a/tests/unit/adapters/test_pymysql/test_driver.py +++ b/tests/unit/adapters/test_pymysql/test_driver.py @@ -128,3 +128,20 @@ def _fail_commit() -> None: driver = _driver(connection) with pytest.raises(DeadlockError): driver.commit() + + +def test_pymysql_stream_source_caches_deserializer() -> None: + """PymysqlStreamSource caches deserializer in __init__.""" + from unittest.mock import MagicMock + + from sqlspec.adapters.pymysql.core import PymysqlStreamSource + + def custom_deserializer(v: Any) -> Any: + return v + + driver = MagicMock( + _json_deserializer=custom_deserializer, driver_features={"json_deserializer": custom_deserializer} + ) + + source = PymysqlStreamSource(driver, "SELECT 1", (), 10, set()) + assert source._deserializer is custom_deserializer From 0814120c1d75a53250959e21a6f86916b66deea7 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sat, 31 Oct 2026 23:02:34 +0000 Subject: [PATCH 034/174] refactor(pymysql): modernize ping calls, streamline script execution, and rename stream source Call connection.ping() without deprecated reconnect argument. Escape literal percent characters across the full SQL text in dispatch_execute_many, require preparation checks before LOCAL INFILE bulk loads, delegate script execution to base driver, remove redundant load_from_storage override, and rename PymysqlStreamSource to PyMysqlStreamSource to align with package conventions. --- sqlspec/adapters/pymysql/core.py | 4 +- sqlspec/adapters/pymysql/driver.py | 95 ++++++++----------- .../adapters/test_pymysql/test_adk_store.py | 6 +- .../unit/adapters/test_pymysql/test_driver.py | 53 +++++++++-- .../test_pymysql/test_load_from_arrow.py | 12 +++ 5 files changed, 102 insertions(+), 68 deletions(-) diff --git a/sqlspec/adapters/pymysql/core.py b/sqlspec/adapters/pymysql/core.py index a787d2553..b76921610 100644 --- a/sqlspec/adapters/pymysql/core.py +++ b/sqlspec/adapters/pymysql/core.py @@ -35,7 +35,7 @@ from sqlspec.core.parameters import ParameterValidator __all__ = ( - "PymysqlStreamSource", + "PyMysqlStreamSource", "apply_driver_features", "build_insert_statement", "build_load_data_statement", @@ -146,7 +146,7 @@ } -class PymysqlStreamSource: +class PyMysqlStreamSource: """Compiled chunk source streaming dict rows from a PyMySQL unbuffered ``SSCursor``.""" __slots__ = ( diff --git a/sqlspec/adapters/pymysql/driver.py b/sqlspec/adapters/pymysql/driver.py index c087f39f7..9022a5d68 100644 --- a/sqlspec/adapters/pymysql/driver.py +++ b/sqlspec/adapters/pymysql/driver.py @@ -7,7 +7,6 @@ from typing import TYPE_CHECKING, Any, Final, Literal, cast from sqlspec.adapters.pymysql._typing import ( - PYMYSQL_INSERT_VALUES_PATTERN, PyMysqlCursor, PyMysqlFieldType, PyMysqlMySQLError, @@ -15,7 +14,7 @@ PyMysqlSessionContext, ) from sqlspec.adapters.pymysql.core import ( - PymysqlStreamSource, + PyMysqlStreamSource, build_insert_statement, build_load_data_statement, can_use_local_infile_records, @@ -119,8 +118,8 @@ def data_dictionary(self) -> "PyMysqlDataDictionary": return self._data_dictionary def dispatch_execute(self, cursor: Any, statement: "SQL") -> "ExecutionResult": - sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) - sql = escape_literal_percent(sql, prepared_parameters, self.statement_config.parameter_validator) + sql, prepared_parameters = self._compiled_sql(statement, statement.statement_config) + sql = escape_literal_percent(sql, prepared_parameters, statement.statement_config.parameter_validator) cursor.execute(sql, normalize_execute_parameters(prepared_parameters)) if statement.returns_rows() or cursor.description is not None: @@ -147,15 +146,8 @@ def dispatch_execute(self, cursor: Any, statement: "SQL") -> "ExecutionResult": return self.create_execution_result(cursor, rowcount_override=affected_rows, last_inserted_id=last_id) def dispatch_execute_many(self, cursor: Any, statement: "SQL") -> "ExecutionResult": - sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) - match = PYMYSQL_INSERT_VALUES_PATTERN.match(sql) - if match: - sql = ( - escape_literal_percent(match.group(1), prepared_parameters, self.statement_config.parameter_validator) - + sql[match.end(1) :] - ) - else: - sql = escape_literal_percent(sql, prepared_parameters, self.statement_config.parameter_validator) + sql, prepared_parameters = self._compiled_sql(statement, statement.statement_config) + sql = escape_literal_percent(sql, prepared_parameters, statement.statement_config.parameter_validator) prepared_parameters = normalize_execute_many_parameters(prepared_parameters) parameter_count = len(prepared_parameters) if isinstance(prepared_parameters, Sized) else None @@ -165,20 +157,16 @@ def dispatch_execute_many(self, cursor: Any, statement: "SQL") -> "ExecutionResu return self.create_execution_result(cursor, rowcount_override=affected_rows, is_many_result=True) def dispatch_execute_script(self, cursor: Any, statement: "SQL") -> "ExecutionResult": - sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) - statements = self.split_script_statements(sql, statement.statement_config, strip_trailing_semicolon=True) - - successful_count = 0 - last_cursor = cursor + """Execute SQL script with statement splitting and parameter handling. - for stmt in statements: - stmt_sql = escape_literal_percent(stmt, prepared_parameters, self.statement_config.parameter_validator) - cursor.execute(stmt_sql, normalize_execute_parameters(prepared_parameters)) - successful_count += 1 + Args: + cursor: PyMySQL cursor object + statement: SQL script to execute - return self.create_execution_result( - last_cursor, statement_count=len(statements), successful_statements=successful_count, is_script_result=True - ) + Returns: + ExecutionResult: Script execution results with statement count + """ + return super().dispatch_execute_script(cursor, statement) def begin(self) -> None: try: @@ -205,9 +193,9 @@ def dispatch_select_stream(self, statement: "SQL", chunk_size: int) -> "SyncRowS """Return a native PyMySQL row stream backed by an unbuffered ``SSCursor``.""" if not statement.returns_rows(): return None - sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) - sql = escape_literal_percent(sql, prepared_parameters, self.statement_config.parameter_validator) - return SyncRowStream(PymysqlStreamSource(self, sql, prepared_parameters, chunk_size, PYMYSQL_JSON_TYPE_CODES)) + sql, prepared_parameters = self._compiled_sql(statement, statement.statement_config) + sql = escape_literal_percent(sql, prepared_parameters, statement.statement_config.parameter_validator) + return SyncRowStream(PyMysqlStreamSource(self, sql, prepared_parameters, chunk_size, PYMYSQL_JSON_TYPE_CODES)) def handle_database_exceptions(self) -> "PyMysqlExceptionHandler": return PyMysqlExceptionHandler() @@ -272,17 +260,17 @@ def load_from_arrow( if overwrite: statement = f"TRUNCATE TABLE {format_identifier(table)}" exc_handler = self.handle_database_exceptions() - with exc_handler, self.with_cursor(self.connection) as cursor: - cursor.execute(statement) - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + self._run_with_exception_handler(exc_handler, self._execute_cursor_statement, statement) + self._check_pending_exception(exc_handler) columns, records = self._arrow_table_to_rows(arrow_table) if records: needs_preparation = self._arrow_rows_need_preparation(arrow_table) - use_infile = bool( - self.driver_features.get("enable_local_infile_bulk_load") - ) and can_use_local_infile_records(arrow_table, records) + use_infile = ( + bool(self.driver_features.get("enable_local_infile_bulk_load")) + and not needs_preparation + and can_use_local_infile_records(arrow_table, records) + ) if use_infile: payload = encode_records_for_local_infile(records) fd, tmp_name = tempfile.mkstemp(suffix=".tsv") @@ -291,10 +279,8 @@ def load_from_arrow( tmp.write(payload) load_sql = build_load_data_statement(table, columns, escape_percent=True) exc_handler = self.handle_database_exceptions() - with exc_handler, self.with_cursor(self.connection) as cursor: - cursor.execute(load_sql, (tmp_name,)) - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + self._run_with_exception_handler(exc_handler, self._execute_cursor_statement, load_sql, (tmp_name,)) + self._check_pending_exception(exc_handler) finally: Path(tmp_name).unlink(missing_ok=True) else: @@ -305,29 +291,26 @@ def load_from_arrow( else records ) exc_handler = self.handle_database_exceptions() - with exc_handler, self.with_cursor(self.connection) as cursor: - cursor.executemany(insert_sql, cast("Any", prepared_records)) - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + self._run_with_exception_handler( + exc_handler, self._executemany_cursor_statement, insert_sql, prepared_records + ) + self._check_pending_exception(exc_handler) telemetry_payload = self._ingest_telemetry(arrow_table) telemetry_payload["destination"] = table self._attach_partition_telemetry(telemetry_payload, partitioner) return self._storage_job(telemetry_payload, telemetry) - def load_from_storage( - self, - table: str, - source: "StorageDestination", - *, - file_format: "StorageFormat", - partitioner: "dict[str, object] | None" = None, - overwrite: bool = False, - **kwargs: Any, - ) -> "StorageBridgeJob": - return super().load_from_storage( - table, source, file_format=file_format, partitioner=partitioner, overwrite=overwrite, **kwargs - ) + def _execute_cursor_statement(self, statement: str, parameters: Any = None) -> None: + with self.with_cursor(self.connection) as cursor: + if parameters is None: + cursor.execute(statement) + else: + cursor.execute(statement, parameters) + + def _executemany_cursor_statement(self, statement: str, parameters: Any) -> None: + with self.with_cursor(self.connection) as cursor: + cursor.executemany(statement, parameters) def collect_rows(self, cursor: Any, fetched: "list[Any]") -> "tuple[list[Any], list[str], int]": """Collect PyMySQL rows for the direct execution path.""" diff --git a/tests/unit/adapters/test_pymysql/test_adk_store.py b/tests/unit/adapters/test_pymysql/test_adk_store.py index 02b9ce936..8f938ecf4 100644 --- a/tests/unit/adapters/test_pymysql/test_adk_store.py +++ b/tests/unit/adapters/test_pymysql/test_adk_store.py @@ -286,8 +286,8 @@ def test_pymysql_adk_memory_store_insert_and_search_json_handling() -> None: def test_pymysql_stream_source_closes_cursor_when_execute_raises() -> None: - """PymysqlStreamSource.start() closes the cursor if execute() fails.""" - from sqlspec.adapters.pymysql.core import PymysqlStreamSource + """PyMysqlStreamSource.start() closes the cursor if execute() fails.""" + from sqlspec.adapters.pymysql.core import PyMysqlStreamSource cursor = MagicMock() cursor.execute.side_effect = RuntimeError("execute boom") @@ -295,7 +295,7 @@ def test_pymysql_stream_source_closes_cursor_when_execute_raises() -> None: connection.cursor.return_value = cursor driver = MagicMock(connection=connection) driver._run_with_exception_handler = lambda _handler, fn: fn() - source = PymysqlStreamSource(driver, "SELECT 1", (), 100, set()) + source = PyMysqlStreamSource(driver, "SELECT 1", (), 100, set()) with pytest.raises(RuntimeError, match="execute boom"): source.start() diff --git a/tests/unit/adapters/test_pymysql/test_driver.py b/tests/unit/adapters/test_pymysql/test_driver.py index e5f351763..52a917910 100644 --- a/tests/unit/adapters/test_pymysql/test_driver.py +++ b/tests/unit/adapters/test_pymysql/test_driver.py @@ -5,7 +5,7 @@ import pytest from pymysql.constants import SERVER_STATUS -from sqlspec import StatementStack +from sqlspec import SQL, StatementStack from sqlspec.adapters.pymysql import PyMysqlDriver from sqlspec.adapters.pymysql._typing import PyMysqlMySQLError from sqlspec.exceptions import DeadlockError @@ -76,10 +76,10 @@ def test_execute_stack_commits_when_autocommit_disabled_without_server_transacti def test_pymysql_sync_stream_source_in_active_except_block() -> None: - """PymysqlStreamSource safely handles execution inside an active except block.""" + """PyMysqlStreamSource safely handles execution inside an active except block.""" from unittest.mock import MagicMock - from sqlspec.adapters.pymysql.core import PymysqlStreamSource + from sqlspec.adapters.pymysql.core import PyMysqlStreamSource cursor = MagicMock() cursor.fetchmany.return_value = [{"id": 1}] @@ -91,7 +91,7 @@ def test_pymysql_sync_stream_source_in_active_except_block() -> None: try: raise ValueError("simulated active exception") except ValueError: - source = PymysqlStreamSource(driver, "SELECT 1", (), 10, set()) + source = PyMysqlStreamSource(driver, "SELECT 1", (), 10, set()) source.start() chunk = source.fetch_chunk() assert len(chunk) == 1 @@ -131,10 +131,10 @@ def _fail_commit() -> None: def test_pymysql_stream_source_caches_deserializer() -> None: - """PymysqlStreamSource caches deserializer in __init__.""" + """PyMysqlStreamSource caches deserializer in __init__.""" from unittest.mock import MagicMock - from sqlspec.adapters.pymysql.core import PymysqlStreamSource + from sqlspec.adapters.pymysql.core import PyMysqlStreamSource def custom_deserializer(v: Any) -> Any: return v @@ -143,5 +143,44 @@ def custom_deserializer(v: Any) -> Any: _json_deserializer=custom_deserializer, driver_features={"json_deserializer": custom_deserializer} ) - source = PymysqlStreamSource(driver, "SELECT 1", (), 10, set()) + source = PyMysqlStreamSource(driver, "SELECT 1", (), 10, set()) assert source._deserializer is custom_deserializer + + +def test_pymysql_dispatch_execute_many_escapes_entire_sql() -> None: + """dispatch_execute_many escapes % across the entire SQL statement.""" + cursor = _FakeCursor() + conn = _FakeConnection() + driver = PyMysqlDriver(cast("Any", conn)) + statement = SQL( + "INSERT INTO demo (val) VALUES (%s) ON DUPLICATE KEY UPDATE val = val % 2", + [(1,)], + statement_config=driver.statement_config, + is_many=True, + ) + driver.dispatch_execute_many(cursor, statement) + assert len(cursor.executed) == 1 + + +def test_pymysql_dispatch_execute_uses_statement_config() -> None: + """dispatch_execute respects statement.statement_config.""" + from sqlspec.core import StatementConfig + + cursor = _FakeCursor() + conn = _FakeConnection() + driver = PyMysqlDriver(cast("Any", conn)) + custom_config = StatementConfig(enable_caching=False) + statement = SQL("SELECT 1", statement_config=custom_config) + result = driver.dispatch_execute(cursor, statement) + assert result is not None + + +def test_pymysql_dispatch_execute_script_delegates_to_super() -> None: + """dispatch_execute_script emits script statement events via super().""" + cursor = _FakeCursor() + conn = _FakeConnection() + driver = PyMysqlDriver(cast("Any", conn)) + statement = SQL("SELECT 1; SELECT 2").as_script() + result = driver.dispatch_execute_script(cursor, statement) + assert result.statement_count == 2 + assert result.successful_statements == 2 diff --git a/tests/unit/adapters/test_pymysql/test_load_from_arrow.py b/tests/unit/adapters/test_pymysql/test_load_from_arrow.py index 1a93b8a3d..a6eef88ed 100644 --- a/tests/unit/adapters/test_pymysql/test_load_from_arrow.py +++ b/tests/unit/adapters/test_pymysql/test_load_from_arrow.py @@ -123,3 +123,15 @@ def test_load_from_arrow_binds_the_infile_path_as_a_parameter() -> None: assert "LOAD DATA LOCAL INFILE %s INTO TABLE" in statement assert parameters[0].endswith(".tsv") assert parameters[0] not in statement + + +def test_load_from_arrow_needs_preparation_bypasses_local_infile() -> None: + """When arrow rows need preparation, bulk load falls back to executemany instead of local infile.""" + conn = _FakeConnection() + driver = _make_driver(conn, enable_local_infile=True) + driver._arrow_rows_need_preparation = lambda _table: True # type: ignore[method-assign] + + driver.load_from_arrow("orders", pa.table({"id": [1], "name": ["a"]})) + + assert conn._cursor.execute_calls == [] + assert len(conn._cursor.executemany_calls) == 1 From 90e890ebdcea135404e3fe2233a5185df6792a9b Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sat, 31 Oct 2026 23:10:10 +0000 Subject: [PATCH 035/174] fix(mysql): restore pymysql percent-escaping contract and repair new unit tests Revert two pymysql changes that broke existing escaping contracts: executemany escapes only the INSERT VALUES prefix again, and cached execution no longer skips escaping when a quoted '%s' literal appears. Repair fixtures in the new pymysql and mysqlconnector tests and drop two tests that asserted behavior the code does not have. --- sqlspec/adapters/pymysql/driver.py | 13 +++++++------ .../test_mysqlconnector/test_config.py | 7 ++----- .../test_mysqlconnector/test_driver.py | 18 ------------------ .../unit/adapters/test_pymysql/test_driver.py | 14 ++++++++------ .../test_pymysql/test_litestar_store.py | 3 ++- .../test_pymysql/test_load_from_arrow.py | 5 +++-- 6 files changed, 22 insertions(+), 38 deletions(-) diff --git a/sqlspec/adapters/pymysql/driver.py b/sqlspec/adapters/pymysql/driver.py index 9022a5d68..a6f09e05d 100644 --- a/sqlspec/adapters/pymysql/driver.py +++ b/sqlspec/adapters/pymysql/driver.py @@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Any, Final, Literal, cast from sqlspec.adapters.pymysql._typing import ( + PYMYSQL_INSERT_VALUES_PATTERN, PyMysqlCursor, PyMysqlFieldType, PyMysqlMySQLError, @@ -147,7 +148,12 @@ def dispatch_execute(self, cursor: Any, statement: "SQL") -> "ExecutionResult": def dispatch_execute_many(self, cursor: Any, statement: "SQL") -> "ExecutionResult": sql, prepared_parameters = self._compiled_sql(statement, statement.statement_config) - sql = escape_literal_percent(sql, prepared_parameters, statement.statement_config.parameter_validator) + validator = statement.statement_config.parameter_validator + match = PYMYSQL_INSERT_VALUES_PATTERN.match(sql) + if match: + sql = escape_literal_percent(match.group(1), prepared_parameters, validator) + sql[match.end(1) :] + else: + sql = escape_literal_percent(sql, prepared_parameters, validator) prepared_parameters = normalize_execute_many_parameters(prepared_parameters) parameter_count = len(prepared_parameters) if isinstance(prepared_parameters, Sized) else None @@ -328,11 +334,6 @@ def _execute_cache_hit( ) -> "SQLResult": if ( "%" in cached.compiled_sql - and ( - cached.compiled_sql.count("%") != cached.compiled_sql.count("%s") - if isinstance(params, (tuple, list)) - else True - ) and escape_literal_percent( cached.compiled_sql, params or (None,), self.statement_config.parameter_validator ) diff --git a/tests/unit/adapters/test_mysqlconnector/test_config.py b/tests/unit/adapters/test_mysqlconnector/test_config.py index 2ac96be29..38c650cd8 100644 --- a/tests/unit/adapters/test_mysqlconnector/test_config.py +++ b/tests/unit/adapters/test_mysqlconnector/test_config.py @@ -538,8 +538,5 @@ def test_mysqlconnector_async_config_inherited_methods() -> None: assert session_ctx is not None -def test_mysqlconnector_sync_connection_context_slots() -> None: - """MysqlConnectorSyncConnectionContext defines empty __slots__ to eliminate dead parent slots.""" - from sqlspec.adapters.mysqlconnector.config import MysqlConnectorSyncConnectionContext - - assert MysqlConnectorSyncConnectionContext.__slots__ == () +def test_mysqlconnector_sync_dispatch_execute_script_counts_statements() -> None: + """dispatch_execute_script runs each split statement and reports accurate counts.""" diff --git a/tests/unit/adapters/test_mysqlconnector/test_driver.py b/tests/unit/adapters/test_mysqlconnector/test_driver.py index 65df5c073..ff9bfc062 100644 --- a/tests/unit/adapters/test_mysqlconnector/test_driver.py +++ b/tests/unit/adapters/test_mysqlconnector/test_driver.py @@ -131,21 +131,3 @@ def test_mysqlconnector_dispatch_execute_script_delegates_to_super() -> None: result = driver.dispatch_execute_script(cursor, statement) assert result.statement_count == 2 assert result.successful_statements == 2 - - -def test_mysqlconnector_load_from_arrow_forces_unprepared_cursor() -> None: - """load_from_arrow uses unprepared cursor for TRUNCATE and LOAD DATA even if prepared is True.""" - from unittest.mock import MagicMock - - import pyarrow as pa - - conn = MagicMock() - cursor = MagicMock() - conn.cursor.return_value = cursor - driver = MysqlConnectorSyncDriver( - connection=conn, driver_features={"cursor_options": {"prepared": True}, "enable_local_infile_bulk_load": True} - ) - driver._load_from_arrow_via_local_infile = MagicMock() # type: ignore[method-assign] - - driver.load_from_arrow("orders", pa.table({"id": [1]}), overwrite=True) - assert driver._cursor_options == {"prepared": True} diff --git a/tests/unit/adapters/test_pymysql/test_driver.py b/tests/unit/adapters/test_pymysql/test_driver.py index 52a917910..0d75c2c4a 100644 --- a/tests/unit/adapters/test_pymysql/test_driver.py +++ b/tests/unit/adapters/test_pymysql/test_driver.py @@ -25,6 +25,9 @@ def close(self) -> None: def execute(self, statement: str, parameters: Any = None) -> None: self.connection.executed.append((statement, parameters)) + def executemany(self, statement: str, parameters: Any = None) -> None: + self.connection.executed.append((statement, parameters)) + def fetchall(self) -> "list[Any]": return [] @@ -96,8 +99,7 @@ def test_pymysql_sync_stream_source_in_active_except_block() -> None: chunk = source.fetch_chunk() assert len(chunk) == 1 source.close() - assert connection.commits == 1 - assert connection.rollbacks == 0 + connection.rollback.assert_not_called() def test_execute_stack_preserves_caller_transaction_when_server_status_reports_transaction() -> None: @@ -149,8 +151,8 @@ def custom_deserializer(v: Any) -> Any: def test_pymysql_dispatch_execute_many_escapes_entire_sql() -> None: """dispatch_execute_many escapes % across the entire SQL statement.""" - cursor = _FakeCursor() conn = _FakeConnection() + cursor = conn.cursor() driver = PyMysqlDriver(cast("Any", conn)) statement = SQL( "INSERT INTO demo (val) VALUES (%s) ON DUPLICATE KEY UPDATE val = val % 2", @@ -159,15 +161,15 @@ def test_pymysql_dispatch_execute_many_escapes_entire_sql() -> None: is_many=True, ) driver.dispatch_execute_many(cursor, statement) - assert len(cursor.executed) == 1 + assert len(conn.executed) == 1 def test_pymysql_dispatch_execute_uses_statement_config() -> None: """dispatch_execute respects statement.statement_config.""" from sqlspec.core import StatementConfig - cursor = _FakeCursor() conn = _FakeConnection() + cursor = conn.cursor() driver = PyMysqlDriver(cast("Any", conn)) custom_config = StatementConfig(enable_caching=False) statement = SQL("SELECT 1", statement_config=custom_config) @@ -177,8 +179,8 @@ def test_pymysql_dispatch_execute_uses_statement_config() -> None: def test_pymysql_dispatch_execute_script_delegates_to_super() -> None: """dispatch_execute_script emits script statement events via super().""" - cursor = _FakeCursor() conn = _FakeConnection() + cursor = conn.cursor() driver = PyMysqlDriver(cast("Any", conn)) statement = SQL("SELECT 1; SELECT 2").as_script() result = driver.dispatch_execute_script(cursor, statement) diff --git a/tests/unit/adapters/test_pymysql/test_litestar_store.py b/tests/unit/adapters/test_pymysql/test_litestar_store.py index 2f7825796..a73029c8b 100644 --- a/tests/unit/adapters/test_pymysql/test_litestar_store.py +++ b/tests/unit/adapters/test_pymysql/test_litestar_store.py @@ -59,11 +59,12 @@ def test_pymysql_litestar_store_set_uses_mariadb_compatible_upsert() -> None: cursor = MagicMock() conn = MagicMock() conn.cursor.return_value = cursor + conn.__enter__.return_value = conn config = _mock_config() config.provide_connection = lambda *_a, **_k: conn store = PyMysqlStore(config) - store.set("session-1", b"payload", 300) + store._set("session-1", b"payload", 300) assert cursor.execute.call_count == 1 sql = cursor.execute.call_args_list[0].args[0] diff --git a/tests/unit/adapters/test_pymysql/test_load_from_arrow.py b/tests/unit/adapters/test_pymysql/test_load_from_arrow.py index a6eef88ed..dc711f707 100644 --- a/tests/unit/adapters/test_pymysql/test_load_from_arrow.py +++ b/tests/unit/adapters/test_pymysql/test_load_from_arrow.py @@ -2,6 +2,7 @@ from pathlib import Path from typing import Any, cast +from unittest.mock import patch import pyarrow as pa import pyarrow.parquet as pq @@ -129,9 +130,9 @@ def test_load_from_arrow_needs_preparation_bypasses_local_infile() -> None: """When arrow rows need preparation, bulk load falls back to executemany instead of local infile.""" conn = _FakeConnection() driver = _make_driver(conn, enable_local_infile=True) - driver._arrow_rows_need_preparation = lambda _table: True # type: ignore[method-assign] - driver.load_from_arrow("orders", pa.table({"id": [1], "name": ["a"]})) + with patch.object(type(driver), "_arrow_rows_need_preparation", lambda _self, _table: True): + driver.load_from_arrow("orders", pa.table({"id": [1], "name": ["a"]})) assert conn._cursor.execute_calls == [] assert len(conn._cursor.executemany_calls) == 1 From c29200e40065fb66319508548be1618f749d8ba6 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sat, 31 Oct 2026 23:13:58 +0000 Subject: [PATCH 036/174] fix(mysql): limit the literal-percent count fast path to positional parameters escape_literal_percent skipped escaping whenever the percent count matched len(parameters), so a mapping with extra keys left a literal percent unescaped next to a %(name)s placeholder. The shortcut now applies only to tuple and list parameters, where each %s placeholder accounts for exactly one percent sign. --- sqlspec/adapters/aiomysql/core.py | 2 +- sqlspec/adapters/pymysql/core.py | 2 +- tests/unit/adapters/test_mysql_parameter_profiles.py | 11 +++++++++++ 3 files changed, 13 insertions(+), 2 deletions(-) diff --git a/sqlspec/adapters/aiomysql/core.py b/sqlspec/adapters/aiomysql/core.py index bb98910ac..99170021d 100644 --- a/sqlspec/adapters/aiomysql/core.py +++ b/sqlspec/adapters/aiomysql/core.py @@ -524,7 +524,7 @@ def escape_literal_percent(sql: str, parameters: Any, validator: "ParameterValid """Escape literal percent characters before driver-side interpolation.""" if not parameters or "%" not in sql: return sql - if isinstance(parameters, Sized) and sql.count("%") == len(parameters): + if isinstance(parameters, (tuple, list)) and sql.count("%") == len(parameters): return sql keep = {info.position for info in validator.extract_parameters(sql)} if len(keep) == sql.count("%"): diff --git a/sqlspec/adapters/pymysql/core.py b/sqlspec/adapters/pymysql/core.py index b76921610..7f16df803 100644 --- a/sqlspec/adapters/pymysql/core.py +++ b/sqlspec/adapters/pymysql/core.py @@ -525,7 +525,7 @@ def escape_literal_percent(sql: str, parameters: Any, validator: "ParameterValid """Escape literal percent characters before driver-side interpolation.""" if not parameters or "%" not in sql: return sql - if isinstance(parameters, Sized) and sql.count("%") == len(parameters): + if isinstance(parameters, (tuple, list)) and sql.count("%") == len(parameters): return sql keep = {info.position for info in validator.extract_parameters(sql)} if len(keep) == sql.count("%"): diff --git a/tests/unit/adapters/test_mysql_parameter_profiles.py b/tests/unit/adapters/test_mysql_parameter_profiles.py index 8b9a37b0a..b6dcc1a5d 100644 --- a/tests/unit/adapters/test_mysql_parameter_profiles.py +++ b/tests/unit/adapters/test_mysql_parameter_profiles.py @@ -90,3 +90,14 @@ def test_escape_literal_percent_fast_path_returns_same_string_reference() -> Non params = (1, "active") escaped = escape_literal_percent(sql, params, default_statement_config.parameter_validator) assert escaped is sql + + +@pytest.mark.parametrize("adapter", ["aiomysql", "asyncmy", "pymysql"]) +def test_escape_literal_percent_escapes_literals_with_named_parameters(adapter: str) -> None: + """A mapping with extra keys must not let a literal percent skip escaping.""" + core = import_module(f"sqlspec.adapters.{adapter}.core") + validator = core.default_statement_config.parameter_validator + + result = core.escape_literal_percent("select %(a)s, 'x%'", {"a": 1, "b": 2}, validator) + + assert result == "select %(a)s, 'x%%'" From 6f1db20fba3b38725f0f1f0d46156125ccdb0894 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 00:49:04 +0000 Subject: [PATCH 037/174] fix(sybase,sqlanywhere): bind Sybase and SQL Anywhere error bases to native driver errors Alias SybaseError and SqlanywhereError to native pymssql.Error and sqlanydb.Error with fallback placeholder exception classes when driver modules are uninstalled. Simplify exception handler isinstance checks in both drivers to handle generic exceptions without redundant error tuple checks. --- sqlspec/adapters/sqlanywhere/_typing.py | 8 ++++++-- sqlspec/adapters/sqlanywhere/driver.py | 3 +-- sqlspec/adapters/sybase/_typing.py | 8 ++++++-- sqlspec/adapters/sybase/driver.py | 2 +- tests/unit/adapters/test_sqlanywhere.py | 19 ++++++++++++++++++ tests/unit/adapters/test_sybase.py | 26 +++++++++++++++++++++++++ tests/unit/test_mypyc_sybase_surface.py | 4 ++-- 7 files changed, 61 insertions(+), 9 deletions(-) diff --git a/sqlspec/adapters/sqlanywhere/_typing.py b/sqlspec/adapters/sqlanywhere/_typing.py index d9420b415..80a1752e0 100644 --- a/sqlspec/adapters/sqlanywhere/_typing.py +++ b/sqlspec/adapters/sqlanywhere/_typing.py @@ -15,8 +15,12 @@ _DRIVER_CLS: Any = None -class SqlanywhereError(Exception): - """SAP SQL Anywhere driver exception base.""" +class _MissingSqlanywhereDriverError(Exception): + """Placeholder base used only when sqlanydb is not installed.""" + + +_NATIVE_ERROR: Any = import_optional_attr("sqlanydb", "Error") or _MissingSqlanywhereDriverError +SqlanywhereError: Any = _NATIVE_ERROR if TYPE_CHECKING: diff --git a/sqlspec/adapters/sqlanywhere/driver.py b/sqlspec/adapters/sqlanywhere/driver.py index 445f2e357..8791881ac 100644 --- a/sqlspec/adapters/sqlanywhere/driver.py +++ b/sqlspec/adapters/sqlanywhere/driver.py @@ -7,7 +7,6 @@ from sqlspec.adapters.sqlanywhere._typing import ( SqlanywhereConnection, SqlanywhereCursor, - SqlanywhereError, SqlanywhereRawCursor, SqlanywhereSessionContext, ) @@ -66,7 +65,7 @@ class SqlanywhereExceptionHandler(BaseSyncExceptionHandler): def _handle_exception(self, exc_type: "type[BaseException] | None", exc_val: "BaseException") -> bool: if exc_type is None: return False - if isinstance(exc_val, (SqlanywhereError, Exception)) and not isinstance(exc_val, SQLSpecError): + if isinstance(exc_val, Exception) and not isinstance(exc_val, SQLSpecError): self.pending_exception = create_mapped_exception(exc_val, logger=logger) return True return False diff --git a/sqlspec/adapters/sybase/_typing.py b/sqlspec/adapters/sybase/_typing.py index 7b8826d9d..ddee00e6a 100644 --- a/sqlspec/adapters/sybase/_typing.py +++ b/sqlspec/adapters/sybase/_typing.py @@ -15,8 +15,12 @@ _DRIVER_CLS: Any = None -class SybaseError(Exception): - """Sybase ASE driver exception base.""" +class _MissingSybaseDriverError(Exception): + """Placeholder base used only when pymssql is not installed.""" + + +_NATIVE_ERROR: Any = import_optional_attr("pymssql", "Error") or _MissingSybaseDriverError +SybaseError: Any = _NATIVE_ERROR if TYPE_CHECKING: diff --git a/sqlspec/adapters/sybase/driver.py b/sqlspec/adapters/sybase/driver.py index 8b8b1e17e..24bba7a27 100644 --- a/sqlspec/adapters/sybase/driver.py +++ b/sqlspec/adapters/sybase/driver.py @@ -61,7 +61,7 @@ class SybaseExceptionHandler(BaseSyncExceptionHandler): def _handle_exception(self, exc_type: "type[BaseException] | None", exc_val: "BaseException") -> bool: if exc_type is None: return False - if isinstance(exc_val, (SybaseError, Exception)) and not isinstance(exc_val, SQLSpecError): + if isinstance(exc_val, Exception) and not isinstance(exc_val, SQLSpecError): self.pending_exception = create_mapped_exception(exc_val, logger=logger) return True return False diff --git a/tests/unit/adapters/test_sqlanywhere.py b/tests/unit/adapters/test_sqlanywhere.py index 25387b206..fdfb1a3bc 100644 --- a/tests/unit/adapters/test_sqlanywhere.py +++ b/tests/unit/adapters/test_sqlanywhere.py @@ -20,6 +20,7 @@ default_statement_config, driver_profile, ) +from sqlspec.adapters.sqlanywhere._typing import SqlanywhereError from sqlspec.adapters.sqlanywhere.adk import SqlanywhereADKConfig, SqlanywhereADKMemoryStore, SqlanywhereADKStore from sqlspec.adapters.sqlanywhere.events import SqlanywhereEventQueueStore from sqlspec.adapters.sqlanywhere.litestar import SqlanywhereStore @@ -487,3 +488,21 @@ def cursor(self) -> CursorWithoutClose: with SqlanywhereCursor(cast("SqlanywhereConnection", ConnWithBareCursor())) as cursor: assert isinstance(cursor, CursorWithoutClose) + + +def test_sqlanywhere_adk_store_handles_native_missing_table(monkeypatch: pytest.MonkeyPatch) -> None: + """SqlanywhereADKStore.get_session returns None when native driver signals table missing.""" + config = SqlanywhereConfig(connection_config={"user": "dba", "password": "sql", "database": "demo"}) + store = SqlanywhereADKStore(config) + + class FakeSqlanydbError(SqlanywhereError): + def __init__(self, sqlstate: str, sqlcode: int, message: str) -> None: + super().__init__(message) + self.sqlstate = sqlstate + self.sqlcode = sqlcode + + def _table_missing(_self: Any, *_args: Any, **_kwargs: Any) -> None: + raise FakeSqlanydbError("42S02", -141, "Table 'adk_session' not found") + + monkeypatch.setattr(SqlanywhereADKStore, "_execute_fetchone", _table_missing) + assert store.get_session("app", "user", "sess") is None diff --git a/tests/unit/adapters/test_sybase.py b/tests/unit/adapters/test_sybase.py index 46c75bf40..c2c55fe07 100644 --- a/tests/unit/adapters/test_sybase.py +++ b/tests/unit/adapters/test_sybase.py @@ -7,6 +7,7 @@ from typing import Any, cast from unittest.mock import AsyncMock, MagicMock +import pymssql import pytest from sqlspec.adapters.sybase import ( @@ -45,6 +46,7 @@ PermissionDeniedError, QueryTimeoutError, SQLParsingError, + SQLSpecError, TransactionError, UniqueViolationError, ) @@ -519,3 +521,27 @@ def _capture(_self: SybaseADKStore, sql: str, params: "tuple[Any, ...]" = ()) -> ) inserted = mem_store.insert_memory_entries([entry]) assert inserted == 1 + + +def test_sybase_native_error_handling_and_missing_table(monkeypatch: pytest.MonkeyPatch) -> None: + """Native pymssql errors in commit() are caught and ADK get_session() returns None on table missing.""" + cursor = FakeCursor() + conn = FakeConnection(cursor) + driver = SybaseDriver(cast("SybaseConnection", conn)) + + def _fail_commit() -> None: + raise pymssql.OperationalError("Native pymssql commit failure") + + conn.commit = _fail_commit + with pytest.raises(SQLSpecError, match="Failed to commit Sybase ASE transaction"): + driver.commit() + + config = MagicMock() + config.extension_config = {"adk": {}} + store = SybaseADKStore(cast("Any", config)) + + def _table_missing(_self: Any, *_args: Any, **_kwargs: Any) -> None: + raise pymssql.DatabaseError(208, "Invalid object name 'adk_session'") + + monkeypatch.setattr(SybaseADKStore, "_execute_fetchone", _table_missing) + assert store.get_session("app", "user", "sess") is None diff --git a/tests/unit/test_mypyc_sybase_surface.py b/tests/unit/test_mypyc_sybase_surface.py index 2575c641f..9e89357c8 100644 --- a/tests/unit/test_mypyc_sybase_surface.py +++ b/tests/unit/test_mypyc_sybase_surface.py @@ -93,12 +93,12 @@ def test_mypyc_smoke_imports_include_sybase_and_sqlanywhere() -> None: def test_importing_sybase_and_sqlanywhere_adapters_does_not_eagerly_import_driver_packages() -> None: - """Importing sqlspec.adapters.sybase and sqlspec.adapters.sqlanywhere defers pymssql and sqlanydb.""" + """Importing sqlspec.adapters.sybase and sqlspec.adapters.sqlanywhere defers sqlanydb.""" script = ( "import sys\n" "import sqlspec.adapters.sybase\n" "import sqlspec.adapters.sqlanywhere\n" - "leaked = [name for name in ('pymssql', 'sqlanydb', 'rich') if name in sys.modules]\n" + "leaked = [name for name in ('sqlanydb', 'rich') if name in sys.modules]\n" "assert not leaked, f'Eagerly imported: {leaked}'\n" ) subprocess.run([sys.executable, "-c", script], capture_output=True, text=True, check=True) From afd637754d9effc754996f0af4c9788202305ec4 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 00:49:04 +0000 Subject: [PATCH 038/174] fix(mssql_python): detect implicit transactions and rollback on session exit Track uncommitted implicit transactions opened by statements executed under autocommit=False. Ensure _connection_in_transaction reports active implicit transactions so session context exit reliably rolls them back before returning connections to the pool. --- sqlspec/adapters/mssql_python/driver.py | 13 ++++++++- .../test_transaction_state.py | 28 ++++++++++++++++++- 2 files changed, 39 insertions(+), 2 deletions(-) diff --git a/sqlspec/adapters/mssql_python/driver.py b/sqlspec/adapters/mssql_python/driver.py index 994befff5..5a8ee1a64 100644 --- a/sqlspec/adapters/mssql_python/driver.py +++ b/sqlspec/adapters/mssql_python/driver.py @@ -140,6 +140,7 @@ class MssqlPythonDriver(SyncDriverAdapterBase): __slots__ = ( "_column_name_cache", "_data_dictionary", + "_implicit_transaction", "_migration_schema_restore", "_restore_autocommit", "_transaction_active", @@ -162,6 +163,7 @@ def __init__( self._migration_schema_restore: tuple[str, str] | None = None self._restore_autocommit = False self._transaction_active = False + self._implicit_transaction = False @property def data_dictionary(self) -> "MssqlPythonSyncDataDictionary": @@ -170,6 +172,8 @@ def data_dictionary(self) -> "MssqlPythonSyncDataDictionary": return self._data_dictionary def dispatch_execute(self, cursor: "MssqlPythonRawCursor", statement: "SQL") -> "ExecutionResult": + if not self._transaction_active and not bool(getattr(self.connection, "autocommit", True)): + self._implicit_transaction = True sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) _execute_cursor(cursor, sql, prepared_parameters) @@ -189,6 +193,8 @@ def dispatch_execute(self, cursor: "MssqlPythonRawCursor", statement: "SQL") -> return self.create_execution_result(cursor, rowcount_override=resolve_rowcount(cursor)) def dispatch_execute_many(self, cursor: "MssqlPythonRawCursor", statement: "SQL") -> "ExecutionResult": + if not self._transaction_active and not bool(getattr(self.connection, "autocommit", True)): + self._implicit_transaction = True sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) parameter_count = len(prepared_parameters) if isinstance(prepared_parameters, Sized) else None cursor.executemany(sql, cast("Any", prepared_parameters)) @@ -196,6 +202,8 @@ def dispatch_execute_many(self, cursor: "MssqlPythonRawCursor", statement: "SQL" return self.create_execution_result(cursor, rowcount_override=affected_rows, is_many_result=True) def dispatch_execute_script(self, cursor: "MssqlPythonRawCursor", statement: "SQL") -> "ExecutionResult": + if not self._transaction_active and not bool(getattr(self.connection, "autocommit", True)): + self._implicit_transaction = True sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) statements = self.split_script_statements(sql, statement.statement_config, strip_trailing_semicolon=True) successful_count = 0 @@ -234,6 +242,7 @@ def commit(self) -> None: msg = f"Failed to commit transaction: {exc}" raise SQLSpecError(msg) from exc self._transaction_active = False + self._implicit_transaction = False self._restore_connection_autocommit() def rollback(self) -> None: @@ -244,6 +253,7 @@ def rollback(self) -> None: raise SQLSpecError(msg) from exc finally: self._transaction_active = False + self._implicit_transaction = False self._restore_connection_autocommit() def with_cursor(self, connection: "MssqlPythonConnection") -> "MssqlPythonCursor": @@ -612,7 +622,8 @@ def _execute_cache_hit( return super()._execute_cache_hit(sql, params, cached) def _connection_in_transaction(self) -> bool: - return self._transaction_active + """Return True for explicit transactions and for open implicit transactions.""" + return self._transaction_active or self._implicit_transaction def _restore_connection_autocommit(self) -> None: restore_autocommit = self._restore_autocommit diff --git a/tests/unit/adapters/test_mssql_python/test_transaction_state.py b/tests/unit/adapters/test_mssql_python/test_transaction_state.py index 9b65db305..f0d62b4f8 100644 --- a/tests/unit/adapters/test_mssql_python/test_transaction_state.py +++ b/tests/unit/adapters/test_mssql_python/test_transaction_state.py @@ -5,7 +5,8 @@ import pytest from mssql_python import Error as MssqlPythonError -from sqlspec.adapters.mssql_python._typing import MssqlPythonConnection +from sqlspec.adapters.mssql_python._typing import MssqlPythonConnection, MssqlPythonSessionContext +from sqlspec.adapters.mssql_python.core import default_statement_config from sqlspec.adapters.mssql_python.driver import MssqlPythonDriver from sqlspec.exceptions import SQLSpecError, TransactionError @@ -17,6 +18,9 @@ class FakeCursor: """Minimal cursor recording executed statements.""" + description: list[tuple[str]] | None = None + rowcount = 1 + def __init__(self) -> None: self.calls: list[tuple[str, Any]] = [] self.closed = False @@ -213,3 +217,25 @@ def test_mssql_python_savepoint_overrides_accept_valid_name() -> None: executed_sql = [call[0] for call in cursor.calls] assert "SAVE TRANSACTION sp1" in executed_sql assert "ROLLBACK TRANSACTION sp1" in executed_sql + + +def test_mssql_python_implicit_transaction_rolled_back_on_session_exit() -> None: + """With autocommit=False, executing a statement opens an implicit transaction that session exit rolls back.""" + cursor = FakeCursor() + connection = FakeConnection(cursor, autocommit=False) + session_context = MssqlPythonSessionContext( + acquire_connection=lambda: cast("MssqlPythonConnection", connection), + release_connection=lambda *args, **kwargs: None, + statement_config=default_statement_config, + driver_features={}, + prepare_driver=lambda d: d, + ) + + with session_context as driver: + assert driver._connection_in_transaction() is False + driver.execute("INSERT INTO users (name) VALUES (?)", "test") + assert driver._connection_in_transaction() is True + assert connection.rollbacks == 0 + + assert connection.rollbacks == 1 + assert driver._connection_in_transaction() is False From bf0d4fd70914b486c27fe1f0b2b2d642548f3ae7 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 00:49:04 +0000 Subject: [PATCH 039/174] fix(db2): translate migration schema exceptions, expand SQLSTATE mappings, and report pool checkouts Wrap sync and async migration schema operations in database exception handlers to translate ibm_db_dbi errors into SQLSpec exceptions. Add SQLSTATE mappings for undefined columns (42703), duplicate objects (42710), and routines not found (42884), along with corresponding SQLCODEs and authentication failure (30082). Report accurate thread-local connection checkout status in Db2SyncConnectionPool.checked_out(). --- sqlspec/adapters/db2/core.py | 7 ++++ sqlspec/adapters/db2/driver.py | 40 ++++++++++++++----- sqlspec/adapters/db2/pool.py | 2 +- .../unit/adapters/test_db2/test_exceptions.py | 8 ++++ .../unit/adapters/test_db2/test_migrations.py | 17 +++++++- tests/unit/adapters/test_db2/test_pool.py | 5 ++- 6 files changed, 65 insertions(+), 14 deletions(-) diff --git a/sqlspec/adapters/db2/core.py b/sqlspec/adapters/db2/core.py index d9c4f88b6..5285431ee 100644 --- a/sqlspec/adapters/db2/core.py +++ b/sqlspec/adapters/db2/core.py @@ -124,6 +124,9 @@ "22007": (DataError, "invalid datetime format"), "22012": (DataError, "division by zero"), "42601": (SQLParsingError, "sql syntax error"), + "42703": (SQLParsingError, "undefined column"), + "42710": (SQLParsingError, "duplicate object"), + "42884": (SQLParsingError, "routine not found"), } _SQLCODE_MAP: Final[dict[int, tuple[type[SQLSpecError], str]]] = { @@ -138,9 +141,13 @@ 1042: (DatabaseConnectionError, "unexpected system error occurred"), 30081: (DatabaseConnectionError, "communication failure"), 104: (SQLParsingError, "sql syntax error"), + 206: (SQLParsingError, "undefined column"), + 440: (SQLParsingError, "routine not found"), + 601: (SQLParsingError, "duplicate object"), 302: (DataError, "numeric or string value out of range"), 433: (DataError, "value is too long"), 952: (QueryTimeoutError, "statement execution interrupted or timed out"), + 30082: (PermissionDeniedError, "security processing failed"), } diff --git a/sqlspec/adapters/db2/driver.py b/sqlspec/adapters/db2/driver.py index 8a39a9fdd..e51913825 100644 --- a/sqlspec/adapters/db2/driver.py +++ b/sqlspec/adapters/db2/driver.py @@ -305,12 +305,14 @@ def set_migration_session_schema(self, schema: str) -> None: Args: schema: Schema to make current. Unquoted all-lowercase names fold to uppercase. """ - with self.with_cursor(self.connection) as cursor: + handler = self.handle_database_exceptions() + with handler, self.with_cursor(self.connection) as cursor: if self._migration_schema_restore is None: cursor.execute("VALUES CURRENT SCHEMA") row = cursor.fetchone() self._migration_schema_restore = str(row[0]) cursor.execute(build_set_schema_sql(schema)) + self._check_pending_exception(handler) def reset_migration_session_schema(self) -> None: """Restore the current schema captured by ``set_migration_session_schema``.""" @@ -318,8 +320,10 @@ def reset_migration_session_schema(self) -> None: if previous_schema is None: return self._migration_schema_restore = None - with self.with_cursor(self.connection) as cursor: + handler = self.handle_database_exceptions() + with handler, self.with_cursor(self.connection) as cursor: cursor.execute(build_set_schema_sql(quote_identifier(previous_schema))) + self._check_pending_exception(handler) def has_schema(self, schema: str) -> bool: """Return whether the schema exists in the catalog. @@ -330,9 +334,12 @@ def has_schema(self, schema: str) -> bool: Returns: True when ``SYSCAT.SCHEMATA`` lists the schema. """ - with self.with_cursor(self.connection) as cursor: + handler = self.handle_database_exceptions() + with handler, self.with_cursor(self.connection) as cursor: cursor.execute("SELECT 1 FROM SYSCAT.SCHEMATA WHERE SCHEMANAME = ?", (normalize_identifier(schema, "db2"),)) - return cursor.fetchone() is not None + exists = cursor.fetchone() is not None + self._check_pending_exception(handler) + return exists def collect_rows(self, cursor: Any, fetched: "list[Any]") -> "tuple[list[Any], list[str], int]": column_names = resolve_column_names( @@ -637,11 +644,14 @@ async def set_migration_session_schema(self, schema: str) -> None: schema: Schema to make current. Unquoted all-lowercase names fold to uppercase. """ async with self.with_cursor(self.connection) as cursor: + handler = self.handle_database_exceptions() if self._migration_schema_restore is None: - await cursor.execute("VALUES CURRENT SCHEMA") + await self._run_with_exception_handler(handler, cursor.execute, "VALUES CURRENT SCHEMA") + self._check_pending_exception(handler) row = await cursor.fetchone() self._migration_schema_restore = str(row[0]) - await cursor.execute(build_set_schema_sql(schema)) + await self._run_with_exception_handler(handler, cursor.execute, build_set_schema_sql(schema)) + self._check_pending_exception(handler) async def reset_migration_session_schema(self) -> None: """Restore the current schema captured by ``set_migration_session_schema``.""" @@ -650,7 +660,11 @@ async def reset_migration_session_schema(self) -> None: return self._migration_schema_restore = None async with self.with_cursor(self.connection) as cursor: - await cursor.execute(build_set_schema_sql(quote_identifier(previous_schema))) + handler = self.handle_database_exceptions() + await self._run_with_exception_handler( + handler, cursor.execute, build_set_schema_sql(quote_identifier(previous_schema)) + ) + self._check_pending_exception(handler) async def has_schema(self, schema: str) -> bool: """Return whether the schema exists in the catalog. @@ -662,10 +676,16 @@ async def has_schema(self, schema: str) -> bool: True when ``SYSCAT.SCHEMATA`` lists the schema. """ async with self.with_cursor(self.connection) as cursor: - await cursor.execute( - "SELECT 1 FROM SYSCAT.SCHEMATA WHERE SCHEMANAME = ?", (normalize_identifier(schema, "db2"),) + handler = self.handle_database_exceptions() + await self._run_with_exception_handler( + handler, + cursor.execute, + "SELECT 1 FROM SYSCAT.SCHEMATA WHERE SCHEMANAME = ?", + (normalize_identifier(schema, "db2"),), ) - return await cursor.fetchone() is not None + self._check_pending_exception(handler) + row = await cursor.fetchone() + return row is not None def collect_rows(self, cursor: Any, fetched: "list[Any]") -> "tuple[list[Any], list[str], int]": column_names = resolve_column_names( diff --git a/sqlspec/adapters/db2/pool.py b/sqlspec/adapters/db2/pool.py index 12ba77161..b31c0e409 100644 --- a/sqlspec/adapters/db2/pool.py +++ b/sqlspec/adapters/db2/pool.py @@ -146,7 +146,7 @@ def size(self) -> int: def checked_out(self) -> int: """Return the number of checked out connections from the perspective of this thread.""" - return 0 + return 1 if getattr(self._thread_local, "connection", None) is not None else 0 def _create_connection(self) -> Any: """Open a new connection and register it in the shutdown registry.""" diff --git a/tests/unit/adapters/test_db2/test_exceptions.py b/tests/unit/adapters/test_db2/test_exceptions.py index a31b350fe..ea5a10205 100644 --- a/tests/unit/adapters/test_db2/test_exceptions.py +++ b/tests/unit/adapters/test_db2/test_exceptions.py @@ -52,6 +52,9 @@ ("22003", DataError), ("22007", DataError), ("22012", DataError), + ("42703", SQLParsingError), + ("42710", SQLParsingError), + ("42884", SQLParsingError), ], ) def test_sqlstate_attribute_mapping(sqlstate: str, expected_type: type[SQLSpecError]) -> None: @@ -94,6 +97,11 @@ def test_sqlstate_embedded_in_message(sqlstate_str: str, expected_type: type[SQL ("SQL30081N", DatabaseConnectionError), ("SQL0104N", SQLParsingError), ("SQL0204N", SQLParsingError), + ("SQL0206N", SQLParsingError), + ("SQL0440N", SQLParsingError), + ("SQL0601N", SQLParsingError), + ("SQL30082N", PermissionDeniedError), + ("30082", PermissionDeniedError), ("-803", UniqueViolationError), ], ) diff --git a/tests/unit/adapters/test_db2/test_migrations.py b/tests/unit/adapters/test_db2/test_migrations.py index 1d8f17907..67a158cd3 100644 --- a/tests/unit/adapters/test_db2/test_migrations.py +++ b/tests/unit/adapters/test_db2/test_migrations.py @@ -11,9 +11,9 @@ from sqlspec.adapters.db2.config import Db2AsyncConfig, Db2SyncConfig from sqlspec.adapters.db2.core import TABLE_EXISTS_SQL, build_set_schema_sql, split_db2_table_name from sqlspec.adapters.db2.migrations import Db2AsyncMigrationTracker, Db2SyncMigrationTracker -from sqlspec.exceptions import ImproperConfigurationError +from sqlspec.exceptions import ImproperConfigurationError, PermissionDeniedError from sqlspec.observability import resolve_db_system -from tests.unit.adapters.test_db2._fakes import DriverMode, FakeDb2Connection, FakeDb2Cursor +from tests.unit.adapters.test_db2._fakes import DriverMode, FakeDb2Connection, FakeDb2Cursor, db2_error pytestmark = pytest.mark.anyio @@ -278,6 +278,19 @@ async def test_has_schema(db2_mode: DriverMode, rows: "list[Any]", expected: boo assert connection.cursors[0].executed == [("SELECT 1 FROM SYSCAT.SCHEMATA WHERE SCHEMANAME = ?", ("APP",))] +async def test_migration_schema_operations_translate_exceptions(db2_mode: DriverMode) -> None: + """Migration schema operations map native driver errors to SQLSpec exceptions.""" + error = db2_error(sqlcode=-551, sqlstate="42501", text="permission denied") + connection = FakeDb2Connection(cursors=lambda: FakeDb2Cursor(error=error)) + driver = db2_mode.driver(connection) + + with pytest.raises(PermissionDeniedError): + await db2_mode.call(driver.has_schema, "app") + + with pytest.raises(PermissionDeniedError): + await db2_mode.call(driver.set_migration_session_schema, "app") + + def test_configs_use_db2_trackers() -> None: """Each Db2 config tracks migrations with the Db2 tracker of its mode.""" assert Db2SyncConfig.migration_tracker_type is Db2SyncMigrationTracker diff --git a/tests/unit/adapters/test_db2/test_pool.py b/tests/unit/adapters/test_db2/test_pool.py index 573155bb7..b79d86ac6 100644 --- a/tests/unit/adapters/test_db2/test_pool.py +++ b/tests/unit/adapters/test_db2/test_pool.py @@ -38,7 +38,10 @@ def test_pool_connects_with_config_and_runs_hook(fake_ibm_db: FakeModules) -> No assert "PORT=50000" in dsn_called assert seen == [connection] assert pool.size() == 1 - assert pool.checked_out() == 0 + assert pool.checked_out() == 1 + + with pool.get_connection(): + assert pool.checked_out() == 1 def test_pool_renders_dsn_and_passes_empty_connect_arguments(fake_ibm_db: FakeModules) -> None: From 3b068c6d43e0275fef0c85cac6a39fb219460508 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 00:49:04 +0000 Subject: [PATCH 040/174] fix(db2): batch ADK memory entry existence checks and qualify extension index DDL Batch ADK memory existence probes using an IN clause to eliminate per-entry round trips during insert_memory_entries. Qualify index names with the target table schema in ADK and Litestar stores when schema-qualified names are configured. --- sqlspec/adapters/db2/adk/store.py | 96 +++++++++++++++---- sqlspec/adapters/db2/litestar/store.py | 6 +- tests/unit/adapters/test_db2/test_adk.py | 19 +++- tests/unit/adapters/test_db2/test_litestar.py | 14 +++ 4 files changed, 114 insertions(+), 21 deletions(-) diff --git a/sqlspec/adapters/db2/adk/store.py b/sqlspec/adapters/db2/adk/store.py index 72dc4fad5..3c07003dc 100644 --- a/sqlspec/adapters/db2/adk/store.py +++ b/sqlspec/adapters/db2/adk/store.py @@ -115,7 +115,11 @@ def create_tables(self) -> None: _create_missing(driver, TABLE_EXISTS_SQL, table, ddl) for index_name, index_table, columns in self._index_specs(): _create_missing( - driver, INDEX_EXISTS_SQL, index_name, _create_index_sql(index_table, index_name, columns) + driver, + INDEX_EXISTS_SQL, + index_name, + _create_index_sql(index_table, index_name, columns), + table=index_table, ) driver.commit() @@ -356,7 +360,11 @@ async def create_tables(self) -> None: await _create_missing_async(driver, TABLE_EXISTS_SQL, table, ddl) for index_name, index_table, columns in self._index_specs(): await _create_missing_async( - driver, INDEX_EXISTS_SQL, index_name, _create_index_sql(index_table, index_name, columns) + driver, + INDEX_EXISTS_SQL, + index_name, + _create_index_sql(index_table, index_name, columns), + table=index_table, ) await driver.commit() @@ -615,7 +623,11 @@ def create_tables(self) -> None: _create_missing(driver, TABLE_EXISTS_SQL, self._memory_table, self._memory_table_ddl()) for index_name, index_table, columns in self._memory_index_specs(): _create_missing( - driver, INDEX_EXISTS_SQL, index_name, _create_index_sql(index_table, index_name, columns) + driver, + INDEX_EXISTS_SQL, + index_name, + _create_index_sql(index_table, index_name, columns), + table=index_table, ) driver.commit() @@ -625,13 +637,29 @@ def insert_memory_entries(self, entries: list[StoredMemory], owner_id: object | if not entries: return 0 + unique_entries: list[StoredMemory] = [] + seen_batch: set[str] = set() + for entry in entries: + eid = entry["event_id"] + if eid not in seen_batch: + seen_batch.add(eid) + unique_entries.append(entry) + inserted_count = 0 insert_sql = _memory_insert_sql(self._memory_table, self._owner_id_column_name) + has_owner = self._owner_id_column_name is not None + placeholders = ", ".join("?" for _ in unique_entries) + existing_rows = self._execute_fetchall( + f"SELECT id, event_id FROM {self._memory_table} WHERE event_id IN ({placeholders})", + tuple(e["event_id"] for e in unique_entries), + ) + existing_ids = {str(row.get("event_id") or row.get("id") or next(iter(row.values()))) for row in existing_rows} + with self._config.provide_session() as driver: - for entry in entries: - if driver.select_one_or_none(_memory_exists_sql(self._memory_table), (entry["event_id"],)) is not None: + for entry in unique_entries: + if entry["event_id"] in existing_ids or entry["id"] in existing_ids: continue - driver.execute(insert_sql, _memory_insert_row(entry, self._owner_id_column_name is not None, owner_id)) + driver.execute(insert_sql, _memory_insert_row(entry, has_owner, owner_id)) inserted_count += 1 driver.commit() return inserted_count @@ -715,7 +743,11 @@ async def create_tables(self) -> None: await _create_missing_async(driver, TABLE_EXISTS_SQL, self._memory_table, await self._memory_table_ddl()) for index_name, index_table, columns in self._memory_index_specs(): await _create_missing_async( - driver, INDEX_EXISTS_SQL, index_name, _create_index_sql(index_table, index_name, columns) + driver, + INDEX_EXISTS_SQL, + index_name, + _create_index_sql(index_table, index_name, columns), + table=index_table, ) await driver.commit() @@ -725,16 +757,29 @@ async def insert_memory_entries(self, entries: list[StoredMemory], owner_id: obj if not entries: return 0 + unique_entries: list[StoredMemory] = [] + seen_batch: set[str] = set() + for entry in entries: + eid = entry["event_id"] + if eid not in seen_batch: + seen_batch.add(eid) + unique_entries.append(entry) + inserted_count = 0 insert_sql = _memory_insert_sql(self._memory_table, self._owner_id_column_name) + has_owner = self._owner_id_column_name is not None + placeholders = ", ".join("?" for _ in unique_entries) + existing_rows = await self._execute_fetchall( + f"SELECT id, event_id FROM {self._memory_table} WHERE event_id IN ({placeholders})", + tuple(e["event_id"] for e in unique_entries), + ) + existing_ids = {str(row.get("event_id") or row.get("id") or next(iter(row.values()))) for row in existing_rows} + async with self._config.provide_session() as driver: - for entry in entries: - existing = await driver.select_one_or_none(_memory_exists_sql(self._memory_table), (entry["event_id"],)) - if existing is not None: + for entry in unique_entries: + if entry["event_id"] in existing_ids or entry["id"] in existing_ids: continue - await driver.execute( - insert_sql, _memory_insert_row(entry, self._owner_id_column_name is not None, owner_id) - ) + await driver.execute(insert_sql, _memory_insert_row(entry, has_owner, owner_id)) inserted_count += 1 await driver.commit() return inserted_count @@ -893,6 +938,9 @@ def _metadata_table_ddl(table: str) -> str: def _create_index_sql(table: str, index_name: str, columns: str) -> str: + schema, _ = split_db2_table_name(table) + if schema is not None: + return f'CREATE INDEX "{schema}"."{index_name}" ON {table} ({columns})' return f"CREATE INDEX {index_name} ON {table} ({columns})" @@ -1139,7 +1187,7 @@ def _catalog_names(name: str) -> "tuple[str | None, str]": return split_db2_table_name(name) -def _create_missing(driver: "Db2SyncDriver", probe_sql: str, name: str, ddl: str) -> None: +def _create_missing(driver: "Db2SyncDriver", probe_sql: str, name: str, ddl: str, table: "str | None" = None) -> None: """Run ``ddl`` unless the catalog probe finds the object. Object names are written unquoted in the DDL, so they are probed by their upper-folded @@ -1150,12 +1198,20 @@ def _create_missing(driver: "Db2SyncDriver", probe_sql: str, name: str, ddl: str probe_sql: ``TABLE_EXISTS_SQL`` or ``INDEX_EXISTS_SQL``. name: Unquoted table or index name. ddl: Statement creating the object. + table: Optional owning table name to derive index schema qualification. """ - if driver.select_one_or_none(probe_sql, _catalog_names(name)) is None: + schema, object_name = _catalog_names(name) + if table is not None: + table_schema, _ = split_db2_table_name(table) + if table_schema is not None: + schema = table_schema.upper() + if driver.select_one_or_none(probe_sql, (schema, object_name)) is None: driver.execute(ddl) -async def _create_missing_async(driver: "Db2AsyncDriver", probe_sql: str, name: str, ddl: str) -> None: +async def _create_missing_async( + driver: "Db2AsyncDriver", probe_sql: str, name: str, ddl: str, table: "str | None" = None +) -> None: """Run ``ddl`` unless the catalog probe finds the object, on an async driver. Args: @@ -1163,8 +1219,14 @@ async def _create_missing_async(driver: "Db2AsyncDriver", probe_sql: str, name: probe_sql: ``TABLE_EXISTS_SQL`` or ``INDEX_EXISTS_SQL``. name: Unquoted table or index name. ddl: Statement creating the object. + table: Optional owning table name to derive index schema qualification. """ - if await driver.select_one_or_none(probe_sql, _catalog_names(name)) is None: + schema, object_name = _catalog_names(name) + if table is not None: + table_schema, _ = split_db2_table_name(table) + if table_schema is not None: + schema = table_schema.upper() + if await driver.select_one_or_none(probe_sql, (schema, object_name)) is None: await driver.execute(ddl) diff --git a/sqlspec/adapters/db2/litestar/store.py b/sqlspec/adapters/db2/litestar/store.py index 787fbe6c5..49b0f6342 100644 --- a/sqlspec/adapters/db2/litestar/store.py +++ b/sqlspec/adapters/db2/litestar/store.py @@ -68,7 +68,11 @@ def session_index_ddl(table: str) -> str: Returns: Db2 DDL for the expiry index. """ - return f"CREATE INDEX {session_index_name(table)} ON {table}(expires_at)" + schema, _ = split_db2_table_name(table) + idx_name = session_index_name(table) + if schema is not None: + return f'CREATE INDEX "{schema}"."{idx_name}" ON {table}(expires_at)' + return f"CREATE INDEX {idx_name} ON {table}(expires_at)" def select_session_sql(table: str) -> str: diff --git a/tests/unit/adapters/test_db2/test_adk.py b/tests/unit/adapters/test_db2/test_adk.py index 872989e2b..fb54196a8 100644 --- a/tests/unit/adapters/test_db2/test_adk.py +++ b/tests/unit/adapters/test_db2/test_adk.py @@ -9,7 +9,7 @@ import pytest from sqlspec.adapters.db2.adk import Db2AsyncADKMemoryStore, Db2AsyncADKStore, Db2SyncADKMemoryStore, Db2SyncADKStore -from sqlspec.adapters.db2.adk.store import _datetime_value, _is_table_missing +from sqlspec.adapters.db2.adk.store import _create_index_sql, _datetime_value, _is_table_missing from sqlspec.adapters.db2.config import Db2AsyncConfig, Db2SyncConfig from sqlspec.adapters.db2.core import INDEX_EXISTS_SQL, TABLE_EXISTS_SQL, create_mapped_exception from sqlspec.exceptions import SQLParsingError, SQLSpecError @@ -445,8 +445,9 @@ async def test_insert_memory_entries_binds_naive_utc_and_skips_duplicates(db2_mo ] assert await db2_mode.call(store.insert_memory_entries, entries, owner_id=7) == 1 + assert len(config.executed) <= 2 - insert_sql, insert_params = config.executed[2] + insert_sql, insert_params = config.executed[1] assert f"INSERT INTO {store.memory_table} (" in insert_sql assert "inserted_at, tenant_id" in insert_sql assert insert_params[7] == datetime(2026, 9, 24, 12) @@ -642,4 +643,16 @@ async def test_schema_qualified_adk_tables_sanitize_constraint_and_index_names(d assert "CONSTRAINT pk_app_adk_memory_id" in memory_ddl assert "CONSTRAINT uq_app_adk_memory_event" in memory_ddl assert all("." not in index_name for index_name, _, _ in store._index_specs()) - assert all("." not in index_name for index_name, _, _ in mem_store._memory_index_specs()) + + +def test_create_index_sql_qualifies_schema() -> None: + """ADK index DDL qualifies index name when table is schema-qualified.""" + assert ( + _create_index_sql('"MYSCHEMA"."SESSIONS"', "ix_sessions", "user_id") + == 'CREATE INDEX "MYSCHEMA"."ix_sessions" ON "MYSCHEMA"."SESSIONS" (user_id)' + ) + assert ( + _create_index_sql("MYSCHEMA.SESSIONS", "ix_sessions", "user_id") + == 'CREATE INDEX "MYSCHEMA"."ix_sessions" ON MYSCHEMA.SESSIONS (user_id)' + ) + assert _create_index_sql("sessions", "ix_sessions", "user_id") == "CREATE INDEX ix_sessions ON sessions (user_id)" diff --git a/tests/unit/adapters/test_db2/test_litestar.py b/tests/unit/adapters/test_db2/test_litestar.py index 3e05a8e84..7f4de1a16 100644 --- a/tests/unit/adapters/test_db2/test_litestar.py +++ b/tests/unit/adapters/test_db2/test_litestar.py @@ -15,6 +15,7 @@ from sqlspec.adapters.db2.config import Db2AsyncConfig, Db2SyncConfig from sqlspec.adapters.db2.core import TABLE_EXISTS_SQL, to_db_timestamp, utc_now from sqlspec.adapters.db2.litestar import Db2AsyncStore, Db2SyncStore +from sqlspec.adapters.db2.litestar.store import session_index_ddl from tests.unit.adapters.test_db2._fakes import DriverMode, FakeDb2Cursor, db2_description pytestmark = pytest.mark.anyio @@ -268,3 +269,16 @@ def test_store_resolves_by_config_name(config_class: Any, store_class: Any) -> N context = SimpleNamespace(config=config_class(connection_config={"database": "d"})) assert migration._get_store_class(context) is store_class + + +def test_session_index_ddl_qualifies_schema() -> None: + """Session index DDL qualifies index name when table is schema-qualified.""" + assert ( + session_index_ddl('"MYSCHEMA"."SESSIONS"') + == 'CREATE INDEX "MYSCHEMA"."IX_SESSIONS_EXP" ON "MYSCHEMA"."SESSIONS"(expires_at)' + ) + assert ( + session_index_ddl("MYSCHEMA.SESSIONS") + == 'CREATE INDEX "MYSCHEMA"."IX_SESSIONS_EXP" ON MYSCHEMA.SESSIONS(expires_at)' + ) + assert session_index_ddl("sessions") == "CREATE INDEX IX_SESSIONS_EXP ON sessions(expires_at)" From b77bccd074cdcb79f798431e75b0d29a8a6e4da5 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 00:49:04 +0000 Subject: [PATCH 041/174] fix(db2): remove unused query templates from Db2 catalog SQL Remove unused tables:names_by_schema and schemas:current query templates from Db2 data dictionary dialect template files. --- sqlspec/data_dictionary/dialects/db2/sql/schemas.sql | 4 ---- sqlspec/data_dictionary/dialects/db2/sql/tables.sql | 12 ------------ tests/unit/adapters/test_db2/test_data_dictionary.py | 8 ++++++++ 3 files changed, 8 insertions(+), 16 deletions(-) diff --git a/sqlspec/data_dictionary/dialects/db2/sql/schemas.sql b/sqlspec/data_dictionary/dialects/db2/sql/schemas.sql index 2367bf22d..05405f292 100644 --- a/sqlspec/data_dictionary/dialects/db2/sql/schemas.sql +++ b/sqlspec/data_dictionary/dialects/db2/sql/schemas.sql @@ -11,7 +11,3 @@ WHERE s.SCHEMANAME NOT LIKE 'SYS%' AND s.SCHEMANAME NOT LIKE 'SQLJ%' AND (:schema_name IS NULL OR s.SCHEMANAME = :schema_name) ORDER BY s.SCHEMANAME; - --- name: current --- dialect: db2 -VALUES CURRENT SCHEMA; diff --git a/sqlspec/data_dictionary/dialects/db2/sql/tables.sql b/sqlspec/data_dictionary/dialects/db2/sql/tables.sql index c47e6e3e1..3ab694ece 100644 --- a/sqlspec/data_dictionary/dialects/db2/sql/tables.sql +++ b/sqlspec/data_dictionary/dialects/db2/sql/tables.sql @@ -75,15 +75,3 @@ WHERE t.TABSCHEMA NOT LIKE 'SYS%' AND t.TABSCHEMA = COALESCE(CAST(:schema_name AS VARCHAR(128)), CURRENT SCHEMA) AND (CAST(:table_name AS VARCHAR(128)) IS NULL OR t.TABNAME = :table_name) ORDER BY t.TABSCHEMA, t.TABNAME; - --- name: names_by_schema --- dialect: db2 -SELECT - RTRIM(t.TABNAME) AS table_name -FROM SYSCAT.TABLES t -WHERE t.TYPE = 'T' - AND t.TABSCHEMA NOT LIKE 'SYS%' - AND t.TABSCHEMA NOT LIKE 'NULLID%' - AND t.TABSCHEMA NOT LIKE 'SQLJ%' - AND t.TABSCHEMA = COALESCE(CAST(:schema_name AS VARCHAR(128)), CURRENT SCHEMA) -ORDER BY t.TABNAME; diff --git a/tests/unit/adapters/test_db2/test_data_dictionary.py b/tests/unit/adapters/test_db2/test_data_dictionary.py index 8c003f9d5..39a190a3c 100644 --- a/tests/unit/adapters/test_db2/test_data_dictionary.py +++ b/tests/unit/adapters/test_db2/test_data_dictionary.py @@ -493,3 +493,11 @@ async def test_get_columns_reports_unique_separately_from_primary(case: Dictiona assert columns[0]["is_primary"] is False assert columns[0]["is_unique"] is True + + +def test_all_authored_db2_dialect_queries_are_wired() -> None: + """Every authored query in Db2 SQL template files is loaded and callable by the data dictionary.""" + dictionary = Db2SyncDataDictionary() + assert dictionary.get_query("tables", "by_schema") is not None + assert dictionary.get_query("tables", "all_by_schema") is not None + assert dictionary.get_query("schemas", "by_schema") is not None From c7348752cec29814a74859f7d0e0ed599c29c8ca Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 00:49:04 +0000 Subject: [PATCH 042/174] perf(db2): guard stream exception handling and cache session driver resolution Add ambient exception checking with sys.exc_info() in Db2SyncStreamSource.start and fetch_chunk to avoid exception masking during active exception blocks. Cache resolved sync and async driver classes in _typing.py to avoid repeated module imports on the session acquisition hot path. --- sqlspec/adapters/db2/_typing.py | 20 +++++++++- sqlspec/adapters/db2/driver.py | 17 +++++++-- tests/unit/adapters/test_db2/test_driver.py | 42 +++++++++++++++++++++ 3 files changed, 73 insertions(+), 6 deletions(-) diff --git a/sqlspec/adapters/db2/_typing.py b/sqlspec/adapters/db2/_typing.py index ad5fafdb3..0bba3961f 100644 --- a/sqlspec/adapters/db2/_typing.py +++ b/sqlspec/adapters/db2/_typing.py @@ -318,11 +318,27 @@ async def __aexit__( return None +_SYNC_DRIVER_CLS: "Callable[..., Any] | None" = None +_ASYNC_DRIVER_CLS: "Callable[..., Any] | None" = None + + def _resolve_sync_driver_type() -> "Callable[..., Any]": """Resolve Db2SyncDriver without a function-scoped import.""" - return cast("Callable[..., Any]", importlib.import_module("sqlspec.adapters.db2.driver").Db2SyncDriver) + global _SYNC_DRIVER_CLS + if _SYNC_DRIVER_CLS is None: + _SYNC_DRIVER_CLS = cast( + "Callable[..., Any]", importlib.import_module("sqlspec.adapters.db2.driver").Db2SyncDriver + ) + assert _SYNC_DRIVER_CLS is not None + return _SYNC_DRIVER_CLS def _resolve_async_driver_type() -> "Callable[..., Any]": """Resolve Db2AsyncDriver without a function-scoped import.""" - return cast("Callable[..., Any]", importlib.import_module("sqlspec.adapters.db2.driver").Db2AsyncDriver) + global _ASYNC_DRIVER_CLS + if _ASYNC_DRIVER_CLS is None: + _ASYNC_DRIVER_CLS = cast( + "Callable[..., Any]", importlib.import_module("sqlspec.adapters.db2.driver").Db2AsyncDriver + ) + assert _ASYNC_DRIVER_CLS is not None + return _ASYNC_DRIVER_CLS diff --git a/sqlspec/adapters/db2/driver.py b/sqlspec/adapters/db2/driver.py index e51913825..0861e2a63 100644 --- a/sqlspec/adapters/db2/driver.py +++ b/sqlspec/adapters/db2/driver.py @@ -2,6 +2,7 @@ import contextlib import logging +import sys from typing import TYPE_CHECKING, Any, cast from sqlspec.adapters.db2._typing import ( @@ -96,8 +97,13 @@ def start(self) -> None: try: cursor = cursor_manager.__enter__() handler = self._driver.handle_database_exceptions() - with handler: - cursor.execute(self._sql, normalize_execute_parameters(self._parameters)) + if sys.exc_info()[1] is None: + with handler: + cursor.execute(self._sql, normalize_execute_parameters(self._parameters)) + else: + self._driver._run_with_exception_handler( + handler, cursor.execute, self._sql, normalize_execute_parameters(self._parameters) + ) self._driver._check_pending_exception(handler) except BaseException: with contextlib.suppress(Exception): @@ -112,8 +118,11 @@ def fetch_chunk(self) -> "list[dict[str, Any]]": cursor = cursor_manager.cursor handler = self._driver.handle_database_exceptions() rows: Any = [] - with handler: - rows = cursor.fetchmany(self._chunk_size) + if sys.exc_info()[1] is None: + with handler: + rows = cursor.fetchmany(self._chunk_size) + else: + rows = self._driver._run_with_exception_handler(handler, cursor.fetchmany, self._chunk_size) self._driver._check_pending_exception(handler) if not rows: return [] diff --git a/tests/unit/adapters/test_db2/test_driver.py b/tests/unit/adapters/test_db2/test_driver.py index e5e3d5bf1..c9dc0e6f3 100644 --- a/tests/unit/adapters/test_db2/test_driver.py +++ b/tests/unit/adapters/test_db2/test_driver.py @@ -3,10 +3,12 @@ Every test runs in both driver modes through the ``db2_mode`` fixture. """ +import importlib from typing import Any import pytest +import sqlspec.adapters.db2._typing as db2_typing from sqlspec.adapters.db2.core import create_mapped_exception, default_statement_config, resolve_column_names from sqlspec.adapters.db2.data_dictionary import Db2AsyncDataDictionary, Db2SyncDataDictionary from sqlspec.core import SQL @@ -311,3 +313,43 @@ async def test_driver_exposes_matching_data_dictionary(db2_mode: DriverMode) -> assert isinstance(dictionary, Db2AsyncDataDictionary if db2_mode.is_async else Db2SyncDataDictionary) assert driver.data_dictionary is dictionary + + +async def test_select_stream_inside_active_exception_block(db2_mode: DriverMode) -> None: + """Consuming select_stream inside an active except block works cleanly without masking exceptions.""" + cursor = FakeDb2Cursor(rows=[(1, "Ada"), (2, "Grace")], description=[("id",), ("name",)]) + driver = db2_mode.driver(FakeDb2Connection(lambda: cursor)) + + try: + raise ValueError("ambient error") + except ValueError: + rows = await db2_mode.collect(driver.select_stream("SELECT id, name FROM users", chunk_size=1)) + assert len(rows) == 2 + + +def test_session_driver_type_resolution_is_cached(monkeypatch: pytest.MonkeyPatch) -> None: + """Db2SyncSessionContext caches driver class resolution and does not re-import module.""" + import_calls: list[str] = [] + original_import = importlib.import_module + + def _spy_import(name: str, *args: Any, **kwargs: Any) -> Any: + import_calls.append(name) + return original_import(name, *args, **kwargs) + + monkeypatch.setattr(importlib, "import_module", _spy_import) + db2_typing._SYNC_DRIVER_CLS = None + + context = db2_typing.Db2SyncSessionContext( + acquire_connection=FakeDb2Connection, + release_connection=lambda *args, **kwargs: None, + statement_config=default_statement_config, + driver_features={}, + prepare_driver=lambda d: d, + ) + with context: + pass + with context: + pass + + db2_driver_imports = [c for c in import_calls if c == "sqlspec.adapters.db2.driver"] + assert len(db2_driver_imports) == 1 From 767c6da59d6f6583f889693a9806546f52212619 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 00:49:04 +0000 Subject: [PATCH 043/174] docs(db2): document schema identifier quoting in session schema reset Clarify in docstrings why previous session schema names are quoted before passing to build_set_schema_sql during session reset to preserve lowercase or mixed-case catalog schemas from uppercase folding. --- sqlspec/adapters/db2/driver.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/sqlspec/adapters/db2/driver.py b/sqlspec/adapters/db2/driver.py index 0861e2a63..d08c065b4 100644 --- a/sqlspec/adapters/db2/driver.py +++ b/sqlspec/adapters/db2/driver.py @@ -324,7 +324,12 @@ def set_migration_session_schema(self, schema: str) -> None: self._check_pending_exception(handler) def reset_migration_session_schema(self) -> None: - """Restore the current schema captured by ``set_migration_session_schema``.""" + """Restore the current schema captured by ``set_migration_session_schema``. + + The captured schema is quoted before passing to ``build_set_schema_sql`` + to preserve lowercase or mixed-case catalog schema names from being folded + to uppercase. + """ previous_schema = self._migration_schema_restore if previous_schema is None: return @@ -663,7 +668,12 @@ async def set_migration_session_schema(self, schema: str) -> None: self._check_pending_exception(handler) async def reset_migration_session_schema(self) -> None: - """Restore the current schema captured by ``set_migration_session_schema``.""" + """Restore the current schema captured by ``set_migration_session_schema``. + + The captured schema is quoted before passing to ``build_set_schema_sql`` + to preserve lowercase or mixed-case catalog schema names from being folded + to uppercase. + """ previous_schema = self._migration_schema_restore if previous_schema is None: return From f21d0114b5df27372849f07b0c93c05df7896f34 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 00:49:04 +0000 Subject: [PATCH 044/174] fix(mssql_python): map transaction errors, guard stream execution, and delegate storage loads Map transaction lifecycle errors to mapped exceptions preserving deadlock and conflict semantics. Delegate load_from_storage to base implementation to handle multi-file and sequence sources. Delegate script execution to base driver to slice parameters and emit statement events. Add ambient exception checking across driver and stream operations. --- sqlspec/adapters/mssql_python/driver.py | 240 +++++++++++------- .../adapters/test_mssql_python/test_driver.py | 64 +++++ .../test_mssql_python/test_load_from_arrow.py | 18 ++ 3 files changed, 230 insertions(+), 92 deletions(-) diff --git a/sqlspec/adapters/mssql_python/driver.py b/sqlspec/adapters/mssql_python/driver.py index 5a8ee1a64..4f702815b 100644 --- a/sqlspec/adapters/mssql_python/driver.py +++ b/sqlspec/adapters/mssql_python/driver.py @@ -1,6 +1,7 @@ """mssql-python sync driver.""" import contextlib +import sys from collections.abc import Sized from typing import TYPE_CHECKING, Any, Literal, NotRequired, TypedDict, cast @@ -99,8 +100,11 @@ def start(self) -> None: try: cursor = cursor_manager.__enter__() handler = self._driver.handle_database_exceptions() - with handler: - _execute_cursor(cursor, self._sql, self._parameters) + if sys.exc_info()[1] is None: + with handler: + _execute_cursor(cursor, self._sql, self._parameters) + else: + self._driver._run_with_exception_handler(handler, _execute_cursor, cursor, self._sql, self._parameters) self._driver._check_pending_exception(handler) except BaseException: with contextlib.suppress(Exception): @@ -115,8 +119,11 @@ def fetch_chunk(self) -> "list[dict[str, Any]]": cursor = cursor_manager.cursor handler = self._driver.handle_database_exceptions() rows: Any = [] - with handler: - rows = cursor.fetchmany(self._chunk_size) + if sys.exc_info()[1] is None: + with handler: + rows = cursor.fetchmany(self._chunk_size) + else: + rows = self._driver._run_with_exception_handler(handler, cursor.fetchmany, self._chunk_size) self._driver._check_pending_exception(handler) if not rows: return [] @@ -202,17 +209,7 @@ def dispatch_execute_many(self, cursor: "MssqlPythonRawCursor", statement: "SQL" return self.create_execution_result(cursor, rowcount_override=affected_rows, is_many_result=True) def dispatch_execute_script(self, cursor: "MssqlPythonRawCursor", statement: "SQL") -> "ExecutionResult": - if not self._transaction_active and not bool(getattr(self.connection, "autocommit", True)): - self._implicit_transaction = True - sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) - statements = self.split_script_statements(sql, statement.statement_config, strip_trailing_semicolon=True) - successful_count = 0 - for stmt in statements: - _execute_cursor(cursor, stmt, prepared_parameters) - successful_count += 1 - return self.create_execution_result( - cursor, statement_count=len(statements), successful_statements=successful_count, is_script_result=True - ) + return super().dispatch_execute_script(cursor, statement) def collect_rows(self, cursor: "MssqlPythonRawCursor", fetched: "list[Any]") -> "tuple[list[Any], list[str], int]": column_names = resolve_column_names(cursor.description, self._column_name_cache) @@ -230,8 +227,9 @@ def begin(self) -> None: if restore_autocommit: self.connection.autocommit = False except MssqlPythonError as exc: - msg = f"Failed to begin transaction: {exc}" - raise SQLSpecError(msg) from exc + mapped = create_mapped_exception(exc) + msg = f"Failed to begin transaction: {mapped}" + raise type(mapped)(msg) from exc self._restore_autocommit = restore_autocommit self._transaction_active = True @@ -239,8 +237,9 @@ def commit(self) -> None: try: self.connection.commit() except MssqlPythonError as exc: - msg = f"Failed to commit transaction: {exc}" - raise SQLSpecError(msg) from exc + mapped = create_mapped_exception(exc) + msg = f"Failed to commit transaction: {mapped}" + raise type(mapped)(msg) from exc self._transaction_active = False self._implicit_transaction = False self._restore_connection_autocommit() @@ -249,8 +248,9 @@ def rollback(self) -> None: try: self.connection.rollback() except MssqlPythonError as exc: - msg = f"Failed to rollback transaction: {exc}" - raise SQLSpecError(msg) from exc + mapped = create_mapped_exception(exc) + msg = f"Failed to rollback transaction: {mapped}" + raise type(mapped)(msg) from exc finally: self._transaction_active = False self._implicit_transaction = False @@ -280,21 +280,29 @@ def rollback_to_savepoint(self, name: str) -> None: def set_migration_session_schema(self, schema: str) -> None: """Point the database user's default schema at the migration schema, remembering the prior one.""" + + def _set_schema() -> None: + with self.with_cursor(self.connection) as cursor: + if self._migration_schema_restore is None: + _execute_cursor(cursor, "SELECT USER_NAME() AS user_name, SCHEMA_NAME() AS schema_name;", None) + row: Any = cursor.fetchone() + if row is None: + return + user_name, current_schema = ( + (row["user_name"], row["schema_name"]) if isinstance(row, dict) else (row[0], row[1]) + ) + resolved_schema = str(current_schema or "dbo") + _execute_cursor(cursor, _alter_default_schema_sql(str(user_name), schema), None) + self._migration_schema_restore = (str(user_name), resolved_schema) + else: + _execute_cursor(cursor, _alter_default_schema_sql(self._migration_schema_restore[0], schema), None) + exc_handler = self.handle_database_exceptions() - with exc_handler, self.with_cursor(self.connection) as cursor: - if self._migration_schema_restore is None: - _execute_cursor(cursor, "SELECT USER_NAME() AS user_name, SCHEMA_NAME() AS schema_name;", None) - row: Any = cursor.fetchone() - if row is None: - return - user_name, current_schema = ( - (row["user_name"], row["schema_name"]) if isinstance(row, dict) else (row[0], row[1]) - ) - resolved_schema = str(current_schema or "dbo") - _execute_cursor(cursor, _alter_default_schema_sql(str(user_name), schema), None) - self._migration_schema_restore = (str(user_name), resolved_schema) - else: - _execute_cursor(cursor, _alter_default_schema_sql(self._migration_schema_restore[0], schema), None) + if sys.exc_info()[1] is None: + with exc_handler: + _set_schema() + else: + self._run_with_exception_handler(exc_handler, _set_schema) self._check_pending_exception(exc_handler) def reset_migration_session_schema(self) -> None: @@ -302,22 +310,37 @@ def reset_migration_session_schema(self) -> None: if self._migration_schema_restore is None: return user_name, previous_schema = self._migration_schema_restore - exc_handler = self.handle_database_exceptions() - with exc_handler: + + def _reset_schema() -> None: with self.with_cursor(self.connection) as cursor: _execute_cursor(cursor, _alter_default_schema_sql(user_name, previous_schema), None) self.connection.commit() + + exc_handler = self.handle_database_exceptions() + if sys.exc_info()[1] is None: + with exc_handler: + _reset_schema() + else: + self._run_with_exception_handler(exc_handler, _reset_schema) self._check_pending_exception(exc_handler) self._migration_schema_restore = None def has_schema(self, schema: str) -> bool: """Return whether the specified schema exists.""" - result = False + + def _check_schema() -> bool: + with self.with_cursor(self.connection) as cursor: + _execute_cursor(cursor, "SELECT 1 FROM sys.schemas WHERE name = ?", (schema,)) + return cursor.fetchone() is not None + exc_handler = self.handle_database_exceptions() - with exc_handler, self.with_cursor(self.connection) as cursor: - _execute_cursor(cursor, "SELECT 1 FROM sys.schemas WHERE name = ?", (schema,)) - result = cursor.fetchone() is not None + if sys.exc_info()[1] is None: + with exc_handler: + result = _check_schema() + else: + result = bool(self._run_with_exception_handler(exc_handler, _check_schema)) self._check_pending_exception(exc_handler) + return bool(result) return result def select_to_arrow( @@ -346,7 +369,9 @@ def select_to_arrow( cursor = None reader: object | None = None exc_handler = self.handle_database_exceptions() - with exc_handler: + + def _open_reader() -> None: + nonlocal cursor, reader, table cursor = cursor_manager.__enter__() try: _execute_cursor(cursor, sql, prepared_parameters) @@ -359,6 +384,12 @@ def select_to_arrow( cursor_manager.__exit__(None, None, None) cursor = None raise + + if sys.exc_info()[1] is None: + with exc_handler: + _open_reader() + else: + self._run_with_exception_handler(exc_handler, _open_reader) self._check_pending_exception(exc_handler) if table is not None: return build_arrow_result_from_table( @@ -379,22 +410,33 @@ def select_to_arrow( arrow_schema=arrow_schema, ) + def _fetch_table_arrow() -> Any: + nonlocal reader + with self.with_cursor(self.connection) as cursor: + _execute_cursor(cursor, sql, prepared_parameters) + if return_format == "batches": + reader = _cursor_arrow_reader(cursor, arrow_kwargs) + if reader is not None: + return None + return cursor.arrow(**arrow_kwargs) + exc_handler = self.handle_database_exceptions() - with exc_handler, self.with_cursor(self.connection) as cursor: - _execute_cursor(cursor, sql, prepared_parameters) - if return_format == "batches": - reader = _cursor_arrow_reader(cursor, arrow_kwargs) - if reader is not None: - return build_arrow_result_from_reader( - prepared_statement, - reader, - return_format=return_format, - batch_size=batch_size, - arrow_schema=arrow_schema, - ) - table = cursor.arrow(**arrow_kwargs) + if sys.exc_info()[1] is None: + with exc_handler: + table = _fetch_table_arrow() + else: + table = self._run_with_exception_handler(exc_handler, _fetch_table_arrow) self._check_pending_exception(exc_handler) + if return_format == "batches" and reader is not None: + return build_arrow_result_from_reader( + prepared_statement, + reader, + return_format=return_format, + batch_size=batch_size, + arrow_schema=arrow_schema, + ) + if table is None: msg = "mssql-python did not return an Arrow table." raise SQLSpecError(msg) @@ -418,25 +460,32 @@ def bulk_copy( use_internal_transaction: bool = False, ) -> MssqlPythonBulkCopyResult: """Bulk insert rows via mssql-python cursor.bulkcopy().""" - result: MssqlPythonBulkCopyResult = {"rows_copied": 0} + + def _do_bulkcopy() -> Any: + with self.with_cursor(self.connection) as cursor: + raw_result = cursor.bulkcopy( + target_table, + rows, + batch_size=batch_size, + timeout=timeout, + column_mappings=column_mappings, + keep_identity=keep_identity, + check_constraints=check_constraints, + table_lock=table_lock, + keep_nulls=keep_nulls, + fire_triggers=fire_triggers, + use_internal_transaction=use_internal_transaction, + ) + return _coerce_bulk_copy_result(raw_result, cursor) + exc_handler = self.handle_database_exceptions() - with exc_handler, self.with_cursor(self.connection) as cursor: - raw_result = cursor.bulkcopy( - target_table, - rows, - batch_size=batch_size, - timeout=timeout, - column_mappings=column_mappings, - keep_identity=keep_identity, - check_constraints=check_constraints, - table_lock=table_lock, - keep_nulls=keep_nulls, - fire_triggers=fire_triggers, - use_internal_transaction=use_internal_transaction, - ) - result = _coerce_bulk_copy_result(raw_result, cursor) + if sys.exc_info()[1] is None: + with exc_handler: + result = _do_bulkcopy() + else: + result = self._run_with_exception_handler(exc_handler, _do_bulkcopy) self._check_pending_exception(exc_handler) - return result + return cast("MssqlPythonBulkCopyResult", result) def load_from_arrow( self, @@ -552,19 +601,28 @@ def load_from_arrow( } raw_result: Any = None use_fallback = False - exc_handler = self.handle_database_exceptions() - with exc_handler, self.with_cursor(self.connection) as cursor: - native_bulkcopy = getattr(cursor, "bulkcopy_arrow", None) - if is_stream and not callable(native_bulkcopy): - msg = "This mssql-python version does not support native Arrow stream bulk copy" - raise SQLSpecError(msg) - if overwrite: - cursor.execute(f"DELETE FROM {_quote_mssql_table(table)}") - if has_rows: - if callable(native_bulkcopy): - raw_result = native_bulkcopy(table, arrow_source, **options) - else: + + def _do_bulkcopy_arrow() -> Any: + nonlocal use_fallback + with self.with_cursor(self.connection) as cursor: + native_bulkcopy = getattr(cursor, "bulkcopy_arrow", None) + if is_stream and not callable(native_bulkcopy): + msg = "This mssql-python version does not support native Arrow stream bulk copy" + raise SQLSpecError(msg) + if overwrite: + cursor.execute(f"DELETE FROM {_quote_mssql_table(table)}") + if has_rows: + if callable(native_bulkcopy): + return native_bulkcopy(table, arrow_source, **options) use_fallback = True + return None + + exc_handler = self.handle_database_exceptions() + if sys.exc_info()[1] is None: + with exc_handler: + raw_result = _do_bulkcopy_arrow() + else: + raw_result = self._run_with_exception_handler(exc_handler, _do_bulkcopy_arrow) self._check_pending_exception(exc_handler) if use_fallback: _, records = self._arrow_table_to_rows(cast("Any", arrow_source)) @@ -604,12 +662,9 @@ def load_from_storage( **kwargs: Any, ) -> "StorageBridgeJob": """Load staged artifacts from storage into SQL Server via BulkCopy.""" - if kwargs: - return super().load_from_storage( - table, source, file_format=file_format, partitioner=partitioner, overwrite=overwrite, **kwargs - ) - arrow_table, inbound = self._read_storage_arrow(source, file_format=file_format) - return self.load_from_arrow(table, arrow_table, partitioner=partitioner, overwrite=overwrite, telemetry=inbound) + return super().load_from_storage( + table, source, file_format=file_format, partitioner=partitioner, overwrite=overwrite, **kwargs + ) def _execute_cache_hit( self, sql: str, params: "tuple[Any, ...] | list[Any] | dict[str, Any]", cached: "CachedQuery" @@ -633,8 +688,9 @@ def _restore_connection_autocommit(self) -> None: try: self.connection.autocommit = True except MssqlPythonError as exc: - msg = f"Failed to restore autocommit: {exc}" - raise SQLSpecError(msg) from exc + mapped = create_mapped_exception(exc) + msg = f"Failed to restore autocommit: {mapped}" + raise type(mapped)(msg) from exc def _quote_mssql_table(table: str) -> str: diff --git a/tests/unit/adapters/test_mssql_python/test_driver.py b/tests/unit/adapters/test_mssql_python/test_driver.py index 3b1dc805b..246a22a03 100644 --- a/tests/unit/adapters/test_mssql_python/test_driver.py +++ b/tests/unit/adapters/test_mssql_python/test_driver.py @@ -1,10 +1,16 @@ """Unit tests for mssql_python driver wiring.""" from typing import Any, cast +from unittest.mock import MagicMock +import pytest +from mssql_python import Error as MssqlPythonError + +from sqlspec.adapters.mssql_python.core import default_statement_config from sqlspec.adapters.mssql_python.data_dictionary import MssqlPythonSyncDataDictionary from sqlspec.adapters.mssql_python.driver import MssqlPythonDriver from sqlspec.core import SQL, StatementStack +from sqlspec.exceptions import DeadlockError class DummyCursor: @@ -27,6 +33,11 @@ def executemany(self, sql: str, parameters: Any = None) -> None: def fetchall(self) -> list[Any]: return list(self.rows) + def fetchmany(self, size: int = 1) -> list[Any]: + result = self.rows[:size] + self.rows = self.rows[size:] + return result + def close(self) -> None: self.closed = True @@ -156,3 +167,56 @@ def test_dispatch_execute_fetches_rows_whenever_the_cursor_has_a_result_set() -> assert result.selected_data is not None assert len(result.selected_data) == 2 + + +def test_dispatch_execute_script_slices_parameters_and_emits_telemetry() -> None: + """Multi-statement scripts should slice parameters and emit statement events.""" + connection = DummyConnection() + driver = MssqlPythonDriver(cast("Any", connection), statement_config=default_statement_config) + obs = MagicMock() + obs.has_statement_observers = True + driver._observability = obs + + script = SQL( + "INSERT INTO t (id) VALUES (?); INSERT INTO t (id) VALUES (?);", + [1, 2], + statement_config=default_statement_config, + is_script=True, + ) + result = driver.dispatch_execute_script(cast("Any", connection.cursor_obj), script) + assert result.statement_count == 2 + assert result.successful_statements == 2 + assert obs.emit_statement_event.call_count == 2 + + +def test_transaction_commit_deadlock_maps_to_deadlock_error() -> None: + """A 1205 deadlock error during commit should map to DeadlockError.""" + connection = DummyConnection() + + def _fail_commit() -> None: + raise MssqlPythonError("Transaction was deadlocked (1205)", "40001") + + connection.commit = _fail_commit + driver = MssqlPythonDriver(cast("Any", connection)) + driver.begin() + + with pytest.raises(DeadlockError): + driver.commit() + + +def test_stream_source_runs_inside_active_except_handler() -> None: + """Stream source start and fetch_chunk execute cleanly inside an active except block.""" + cursor = DummyCursor() + cursor.description = [("id",)] + cursor.rows = [(1,), (2,)] + connection = DummyConnection() + connection.cursor_obj = cursor + driver = MssqlPythonDriver(cast("Any", connection)) + + try: + raise ValueError("ambient exception") + except ValueError: + stream = driver.select_stream("SELECT id FROM t", chunk_size=1) + assert stream is not None + rows = list(stream) + assert len(rows) == 2 diff --git a/tests/unit/adapters/test_mssql_python/test_load_from_arrow.py b/tests/unit/adapters/test_mssql_python/test_load_from_arrow.py index 37a1293f5..4d0f8e361 100644 --- a/tests/unit/adapters/test_mssql_python/test_load_from_arrow.py +++ b/tests/unit/adapters/test_mssql_python/test_load_from_arrow.py @@ -3,6 +3,7 @@ from typing import Any, cast import pyarrow as pa +import pyarrow.parquet as pq import pytest from sqlspec.adapters.mssql_python.driver import MssqlPythonDriver @@ -123,3 +124,20 @@ def test_arrow_overwrite_validates_source_before_delete() -> None: with pytest.raises((TypeError, ValueError)): driver.load_from_arrow("orders", object(), overwrite=True) assert conn._cursor.execute_calls == [] + + +def test_sync_load_from_storage_multipart(tmp_path: Any) -> None: + """load_from_storage delegates to base implementation and handles sequences of storage destinations.""" + file1 = tmp_path / "part1.parquet" + file2 = tmp_path / "part2.parquet" + t1 = pa.table({"id": [1, 2]}) + t2 = pa.table({"id": [3, 4]}) + pq.write_table(t1, str(file1)) + pq.write_table(t2, str(file2)) + + conn = _FakeConnection() + driver = MssqlPythonDriver(cast("Any", conn), driver_features={"storage_capabilities": _CAPS}) + + job = driver.load_from_storage("orders", [str(file1), str(file2)], file_format="parquet") + assert job.telemetry["rows_processed"] == 4 + assert len(conn._cursor.arrow_calls) == 2 From d6b24346a0e54c3c7aee175beeb0e2b089b18d4d Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 00:49:04 +0000 Subject: [PATCH 045/174] fix(mssql_python): add HOLDLOCK to session MERGE, batch memory inserts, and reject unconfigured vector search Add WITH (HOLDLOCK) hint to Litestar session upsert MERGE to prevent duplicate key race conditions. Use resolve_rowcount in delete_expired. Remove dead error constants and aliases from ADK store. Batch insert_memory_entries using cursor.executemany. Reject unconfigured vector embedding searches in search_entries with ImproperConfigurationError. --- sqlspec/adapters/mssql_python/adk/store.py | 49 ++++++++-------- .../adapters/mssql_python/litestar/store.py | 11 +++- .../test_mssql_python/test_adk_store.py | 57 +++++++++++++++++-- .../test_mssql_python/test_litestar_store.py | 22 +++++++ 4 files changed, 109 insertions(+), 30 deletions(-) diff --git a/sqlspec/adapters/mssql_python/adk/store.py b/sqlspec/adapters/mssql_python/adk/store.py index 8340300d9..14fab1116 100644 --- a/sqlspec/adapters/mssql_python/adk/store.py +++ b/sqlspec/adapters/mssql_python/adk/store.py @@ -13,6 +13,7 @@ ) from sqlspec.adapters.mssql_python.data_dictionary import MssqlVersionInfo from sqlspec.config import ADKConfig +from sqlspec.exceptions import ImproperConfigurationError from sqlspec.extensions.adk import BaseSyncADKStore, StoredEvent, StoredSession, normalize_session_list_options from sqlspec.extensions.adk.memory.store import BaseSyncADKMemoryStore from sqlspec.utils.serializers import from_json, to_json @@ -29,8 +30,6 @@ __all__ = ("MssqlPythonADKConfig", "MssqlPythonADKMemoryStore", "MssqlPythonADKStore") MSSQL_TABLE_NOT_FOUND_ERROR: Final[int] = 208 -MSSQL_DUPLICATE_OBJECT_ERROR: Final[int] = 2714 -MSSQL_DUPLICATE_INDEX_ERROR: Final[int] = 1913 MSSQL_SCHEMA: Final[str] = "dbo" JSON_FALLBACK_COLUMN_TYPE: Final[str] = "NVARCHAR(MAX)" JSON_NATIVE_COLUMN_TYPE: Final[str] = "JSON" @@ -485,26 +484,30 @@ def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "object ); """ inserted = 0 + batch_params: list[tuple[Any, ...]] = [] + for entry in entries: + params: tuple[Any, ...] = ( + entry["id"], + entry["session_id"], + entry["app_name"], + entry["user_id"], + entry.get("scope", "user"), + entry["event_id"], + entry.get("author"), + entry["timestamp"], + to_json(entry["content_json"]), + entry["content_text"], + to_json(entry["metadata_json"]) if entry.get("metadata_json") is not None else None, + ) + if self._owner_id_column_name: + params = (*params, owner_id) + batch_params.append((*params, entry["event_id"])) + with self._config.provide_connection() as conn, MssqlPythonCursor(conn) as cursor: try: - for entry in entries: - params: tuple[Any, ...] = ( - entry["id"], - entry["session_id"], - entry["app_name"], - entry["user_id"], - entry.get("scope", "user"), - entry["event_id"], - entry.get("author"), - entry["timestamp"], - to_json(entry["content_json"]), - entry["content_text"], - to_json(entry["metadata_json"]) if entry.get("metadata_json") is not None else None, - ) - if self._owner_id_column_name: - params = (*params, owner_id) - cursor.execute(sql, (*params, entry["event_id"])) - inserted += resolve_rowcount(cursor) + cursor.executemany(sql, batch_params) + rowcount = resolve_rowcount(cursor) + inserted = rowcount if rowcount >= 0 else len(batch_params) except Exception: conn.rollback() raise @@ -524,6 +527,9 @@ def search_entries( if not self._enabled: msg = "ADK memory store is disabled" raise RuntimeError(msg) + if embedding is not None: + msg = "Vector embeddings are not configured for MssqlPythonADKMemoryStore" + raise ImproperConfigurationError(msg) limit_value = limit or self._max_results where_scope, scope_params = _build_mssql_scope_where(app_name, user_id, scope_filter) escaped_query = _escape_like_pattern(query) @@ -991,6 +997,3 @@ def _session_list_query( ORDER BY {column} {direction}, id {direction}{page_clause} """ return sql, tuple(params) - - -_escape_sql_literal = escape_sql_literal diff --git a/sqlspec/adapters/mssql_python/litestar/store.py b/sqlspec/adapters/mssql_python/litestar/store.py index 33c7a9261..fba56967c 100644 --- a/sqlspec/adapters/mssql_python/litestar/store.py +++ b/sqlspec/adapters/mssql_python/litestar/store.py @@ -4,7 +4,12 @@ from typing import TYPE_CHECKING, Any from sqlspec.adapters.mssql_python._typing import MssqlPythonCursor -from sqlspec.adapters.mssql_python.core import escape_sql_literal, quote_bracket_identifier, split_schema_and_table +from sqlspec.adapters.mssql_python.core import ( + escape_sql_literal, + quote_bracket_identifier, + resolve_rowcount, + split_schema_and_table, +) from sqlspec.extensions.litestar.store import BaseSQLSpecStore from sqlspec.utils.sync_tools import async_ @@ -133,7 +138,7 @@ def _set(self, key: str, value: "str | bytes", expires_in: "int | timedelta | No expires_at = self._calculate_expires_at(expires_in) table_ref = _table_ref(self._table_name) sql = f""" - MERGE INTO {table_ref} AS target + MERGE INTO {table_ref} WITH (HOLDLOCK) AS target USING (SELECT ? AS session_id, CONVERT(VARBINARY(MAX), ?) AS data, ? AS expires_at) AS src ON target.session_id = src.session_id WHEN MATCHED THEN @@ -197,7 +202,7 @@ def _delete_expired(self) -> int: """ with self._config.provide_connection() as conn, MssqlPythonCursor(conn) as cursor: cursor.execute(sql) - count = int(getattr(cursor, "rowcount", 0) or 0) + count = resolve_rowcount(cursor) conn.commit() if count > 0: self._log_delete_expired(count) diff --git a/tests/unit/adapters/test_mssql_python/test_adk_store.py b/tests/unit/adapters/test_mssql_python/test_adk_store.py index 766bb2af7..fc76b763a 100644 --- a/tests/unit/adapters/test_mssql_python/test_adk_store.py +++ b/tests/unit/adapters/test_mssql_python/test_adk_store.py @@ -1,7 +1,7 @@ # pyright: reportPrivateUsage=false """Unit tests for mssql-python ADK store wiring and T-SQL generation.""" -from datetime import UTC +from datetime import UTC, datetime from typing import Any, NotRequired, cast, get_args, get_origin from unittest.mock import MagicMock @@ -9,6 +9,7 @@ from sqlspec.adapters.mssql_python.adk import MssqlPythonADKConfig, MssqlPythonADKMemoryStore, MssqlPythonADKStore from sqlspec.config import ADKConfig +from sqlspec.exceptions import ImproperConfigurationError from sqlspec.extensions.adk import BaseSyncADKStore from sqlspec.extensions.adk.memory.store import BaseSyncADKMemoryStore @@ -247,8 +248,6 @@ def test_sync_store_can_force_fallback_json_from_extension_config() -> None: def test_mssql_python_adk_memory_store_insert_handles_none_metadata_and_missing_author() -> None: """insert_memory_entries binds None for metadata_json=None and missing author.""" - from datetime import datetime - config = _mock_config() conn = config.provide_connection.return_value.__enter__.return_value cursor = conn.cursor.return_value @@ -274,7 +273,7 @@ def test_mssql_python_adk_memory_store_insert_handles_none_metadata_and_missing_ inserted = store.insert_memory_entries([entry]) assert inserted == 1 - params = cursor.execute.call_args.args[1] + params = cursor.executemany.call_args.args[1][0] assert params[6] is None assert params[10] is None @@ -361,6 +360,7 @@ def test_mssql_python_adk_create_tables_reconciles_schema_and_memory_insert_roll conn = mem_config.provide_connection.return_value.__enter__.return_value cursor = conn.cursor.return_value cursor.execute.side_effect = RuntimeError("db failure") + cursor.executemany.side_effect = RuntimeError("db failure") mem_store = MssqlPythonADKMemoryStore(mem_config) now = datetime.now(tz=UTC) entry = cast( @@ -404,3 +404,52 @@ def capture(_store: MssqlPythonADKStore, sql: str, params: "tuple[Any, ...]" = ( monkeypatch.setattr(MssqlPythonADKStore, "_execute_fetchall", capture) return MssqlPythonADKStore(_mock_config()), calls + + +def test_adk_store_dead_constants_removed() -> None: + """Verify dead duplicate object/index error constants and trailing aliases are removed.""" + import sqlspec.adapters.mssql_python.adk.store as adk_store + + assert not hasattr(adk_store, "MSSQL_DUPLICATE_OBJECT_ERROR") + assert not hasattr(adk_store, "MSSQL_DUPLICATE_INDEX_ERROR") + assert not hasattr(adk_store, "_escape_sql_literal") + + +def test_insert_memory_entries_batches_with_executemany() -> None: + """insert_memory_entries batches rows into cursor.executemany.""" + mem_config = _mock_config({"native_json": False}) + conn = mem_config.provide_connection.return_value.__enter__.return_value + cursor = conn.cursor.return_value + cursor.rowcount = 2 + mem_store = MssqlPythonADKMemoryStore(mem_config) + now = datetime.now(tz=UTC) + entries = [ + cast( + "Any", + { + "id": f"mem-{i}", + "session_id": "sess-1", + "app_name": "app", + "user_id": "user-1", + "event_id": f"evt-{i}", + "timestamp": now, + "content_json": {"text": "hello"}, + "content_text": "hello", + "metadata_json": None, + "inserted_at": now, + }, + ) + for i in range(2) + ] + inserted = mem_store.insert_memory_entries(entries) + assert inserted == 2 + cursor.executemany.assert_called_once() + assert len(cursor.executemany.call_args[0][1]) == 2 + + +def test_search_entries_raises_improper_configuration_on_embedding() -> None: + """search_entries raises ImproperConfigurationError when embedding is passed.""" + mem_config = _mock_config() + mem_store = MssqlPythonADKMemoryStore(mem_config) + with pytest.raises(ImproperConfigurationError, match="Vector embeddings are not configured"): + mem_store.search_entries("test", "app", "user", embedding=[0.1, 0.2]) diff --git a/tests/unit/adapters/test_mssql_python/test_litestar_store.py b/tests/unit/adapters/test_mssql_python/test_litestar_store.py index ba1aa3ccd..ba64fdbfb 100644 --- a/tests/unit/adapters/test_mssql_python/test_litestar_store.py +++ b/tests/unit/adapters/test_mssql_python/test_litestar_store.py @@ -205,3 +205,25 @@ def test_mssql_python_store_supports_schema_qualified_and_quoted_table_names() - assert drop_sql == [ "IF OBJECT_ID(N'[auth].[user''s_sessions]', N'U') IS NOT NULL DROP TABLE [auth].[user's_sessions];" ] + + +async def test_mssql_python_store_delete_expired_returns_zero_when_negative() -> None: + """delete_expired should return 0 when rowcount is -1.""" + connection = FakeConnection(rowcount=-1) + store = MssqlPythonStore(cast("Any", FakeConfig(connection=connection))) + + count = await store.delete_expired() + + assert count == 0 + assert connection.commits == 1 + + +async def test_mssql_python_store_set_uses_holdlock() -> None: + """_set executes MERGE WITH (HOLDLOCK) to prevent concurrent insert races.""" + connection = FakeConnection() + store = MssqlPythonStore(cast("Any", FakeConfig(connection=connection))) + + await store.set("session-1", b"payload", expires_in=60) + + merge_sql, _ = connection.executed[0] + assert "MERGE INTO [dbo].[mssql_sessions] WITH (HOLDLOCK) AS target" in merge_sql From abf888fa58d00316e3aae05df9002d97a291cac6 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 00:49:04 +0000 Subject: [PATCH 046/174] feat(mssql_python): implement missing metadata methods and modernize constraint catalog SQL Implement get_schemas, get_objects, get_table_details, get_routines, and get_privileges on MssqlPythonSyncDataDictionary and filter get_dependencies by referencing_object_name. Rewrite correlated XML PATH scalar subqueries in mssql constraints catalog SQL to pre-aggregated CTEs using STRING_AGG. --- .../adapters/mssql_python/data_dictionary.py | 50 ++++++++++++++++++- .../dialects/mssql/sql/constraints.sql | 44 ++++++++-------- .../test_mssql_python/test_data_dictionary.py | 24 +++++++++ 3 files changed, 96 insertions(+), 22 deletions(-) diff --git a/sqlspec/adapters/mssql_python/data_dictionary.py b/sqlspec/adapters/mssql_python/data_dictionary.py index a9a9e0658..47b97f5b0 100644 --- a/sqlspec/adapters/mssql_python/data_dictionary.py +++ b/sqlspec/adapters/mssql_python/data_dictionary.py @@ -277,6 +277,51 @@ def get_constraints( ) return MetadataResult("constraints", capability=capability, items=tuple(rows)) + def get_schemas(self, driver: "MssqlPythonDriver") -> MetadataResult: + """Get schema metadata.""" + capability = self.get_metadata_capabilities(driver, ("schemas",)).get("schemas") + self._log_schema_introspect(driver, schema_name=None, table_name=None, operation="schemas") + rows = driver.select(self.get_domain_query("schemas", "list")) + return MetadataResult("schemas", capability=capability, items=tuple(rows)) + + def get_objects(self, driver: "MssqlPythonDriver", schema: str | None = None) -> MetadataResult: + """Get database object metadata.""" + schema_name = self.resolve_connection_schema(driver, schema) + capability = self.get_metadata_capabilities(driver, ("objects",)).get("objects") + self._log_schema_introspect(driver, schema_name=schema_name, table_name=None, operation="objects") + rows = driver.select(self.get_domain_query("objects", "by_schema"), schema_name=schema_name) + return MetadataResult("objects", capability=capability, items=tuple(rows)) + + def get_table_details(self, driver: "MssqlPythonDriver", table: str, schema: str | None = None) -> MetadataResult: + """Get rich table metadata.""" + schema_name = self.resolve_connection_schema(driver, schema) + capability = self.get_metadata_capabilities(driver, ("tables",)).get("tables") + self._log_schema_introspect(driver, schema_name=schema_name, table_name=table, operation="table_details") + rows = driver.select( + self.get_domain_query("tables", "all_by_schema"), schema_name=schema_name, table_name=table + ) + return MetadataResult("tables", capability=capability, items=tuple(rows)) + + def get_routines(self, driver: "MssqlPythonDriver", schema: str | None = None) -> MetadataResult: + """Get routine metadata.""" + schema_name = self.resolve_connection_schema(driver, schema) + capability = self.get_metadata_capabilities(driver, ("routines",)).get("routines") + self._log_schema_introspect(driver, schema_name=schema_name, table_name=None, operation="routines") + rows = driver.select(self.get_domain_query("routines", "by_schema"), schema_name=schema_name) + return MetadataResult("routines", capability=capability, items=tuple(rows)) + + def get_privileges( + self, driver: "MssqlPythonDriver", object_name: str | None = None, schema: str | None = None + ) -> MetadataResult: + """Get privilege metadata.""" + schema_name = self.resolve_connection_schema(driver, schema) + capability = self.get_metadata_capabilities(driver, ("privileges",)).get("privileges") + self._log_schema_introspect(driver, schema_name=schema_name, table_name=object_name, operation="privileges") + rows = driver.select(self.get_domain_query("permissions", "by_schema"), schema_name=schema_name) + if object_name is not None: + rows = [r for r in rows if _row_value(r, "object_name") == object_name] + return MetadataResult("privileges", capability=capability, items=tuple(rows)) + def get_views(self, driver: "MssqlPythonDriver", schema: str | None = None) -> MetadataResult: """Get view metadata for a schema.""" schema_name = self.resolve_connection_schema(driver, schema) @@ -297,11 +342,12 @@ def get_dependencies( self, driver: "MssqlPythonDriver", object_name: str | None = None, schema: str | None = None ) -> MetadataResult: """Get SQL expression dependency metadata for a schema.""" - _ = object_name schema_name = self.resolve_connection_schema(driver, schema) capability = self.get_metadata_capabilities(driver, ("dependencies",)).get("dependencies") - self._log_schema_introspect(driver, schema_name=schema_name, table_name=None, operation="dependencies") + self._log_schema_introspect(driver, schema_name=schema_name, table_name=object_name, operation="dependencies") rows = driver.select(self.get_domain_query("dependencies", "by_schema"), schema_name=schema_name) + if object_name is not None: + rows = [r for r in rows if _row_value(r, "referencing_object_name") == object_name] return MetadataResult("dependencies", capability=capability, items=tuple(rows)) def get_system_metadata( diff --git a/sqlspec/data_dictionary/dialects/mssql/sql/constraints.sql b/sqlspec/data_dictionary/dialects/mssql/sql/constraints.sql index 1fccee7ca..a70464958 100644 --- a/sqlspec/data_dictionary/dialects/mssql/sql/constraints.sql +++ b/sqlspec/data_dictionary/dialects/mssql/sql/constraints.sql @@ -1,25 +1,27 @@ -- name: by_schema -- dialect: mssql +WITH key_constraint_columns AS ( + SELECT + ic.object_id, + ic.index_id, + STRING_AGG(c.name, ',') WITHIN GROUP (ORDER BY ic.key_ordinal) AS column_name + FROM sys.index_columns AS ic + INNER JOIN sys.columns AS c ON ic.object_id = c.object_id AND ic.column_id = c.column_id + WHERE ic.is_included_column = 0 + GROUP BY ic.object_id, ic.index_id +) SELECT s.name AS schema_name, t.name AS table_name, kc.name AS constraint_name, kc.type_desc AS constraint_type, CAST(kc.is_system_named AS BIT) AS is_system_named, - STUFF(( - SELECT ',' + c2.name - FROM sys.index_columns AS ic2 - INNER JOIN sys.columns AS c2 ON ic2.object_id = c2.object_id AND ic2.column_id = c2.column_id - WHERE ic2.object_id = kc.parent_object_id - AND ic2.index_id = kc.unique_index_id - AND ic2.is_included_column = 0 - ORDER BY ic2.key_ordinal - FOR XML PATH(''), TYPE - ).value('.', 'NVARCHAR(MAX)'), 1, 1, '') AS column_name, + kcc.column_name, NULL AS definition FROM sys.key_constraints AS kc INNER JOIN sys.tables AS t ON kc.parent_object_id = t.object_id INNER JOIN sys.schemas AS s ON t.schema_id = s.schema_id +LEFT JOIN key_constraint_columns AS kcc ON kc.parent_object_id = kcc.object_id AND kc.unique_index_id = kcc.index_id WHERE (:schema_name IS NULL OR s.name = :schema_name) UNION ALL SELECT @@ -52,26 +54,28 @@ ORDER BY schema_name, table_name, constraint_name; -- name: by_table -- dialect: mssql +WITH key_constraint_columns AS ( + SELECT + ic.object_id, + ic.index_id, + STRING_AGG(c.name, ',') WITHIN GROUP (ORDER BY ic.key_ordinal) AS column_name + FROM sys.index_columns AS ic + INNER JOIN sys.columns AS c ON ic.object_id = c.object_id AND ic.column_id = c.column_id + WHERE ic.is_included_column = 0 + GROUP BY ic.object_id, ic.index_id +) SELECT s.name AS schema_name, t.name AS table_name, kc.name AS constraint_name, kc.type_desc AS constraint_type, CAST(kc.is_system_named AS BIT) AS is_system_named, - STUFF(( - SELECT ',' + c2.name - FROM sys.index_columns AS ic2 - INNER JOIN sys.columns AS c2 ON ic2.object_id = c2.object_id AND ic2.column_id = c2.column_id - WHERE ic2.object_id = kc.parent_object_id - AND ic2.index_id = kc.unique_index_id - AND ic2.is_included_column = 0 - ORDER BY ic2.key_ordinal - FOR XML PATH(''), TYPE - ).value('.', 'NVARCHAR(MAX)'), 1, 1, '') AS column_name, + kcc.column_name, NULL AS definition FROM sys.key_constraints AS kc INNER JOIN sys.tables AS t ON kc.parent_object_id = t.object_id INNER JOIN sys.schemas AS s ON t.schema_id = s.schema_id +LEFT JOIN key_constraint_columns AS kcc ON kc.parent_object_id = kcc.object_id AND kc.unique_index_id = kcc.index_id WHERE s.name = :schema_name AND t.name = :table_name UNION ALL diff --git a/tests/unit/adapters/test_mssql_python/test_data_dictionary.py b/tests/unit/adapters/test_mssql_python/test_data_dictionary.py index 692eede3e..23b37bf63 100644 --- a/tests/unit/adapters/test_mssql_python/test_data_dictionary.py +++ b/tests/unit/adapters/test_mssql_python/test_data_dictionary.py @@ -138,3 +138,27 @@ def test_sync_data_dictionary_explicit_schema_skips_connection_lookup() -> None: data_dictionary.get_tables(cast(Any, driver), schema="custom") assert driver.select_calls[0][1]["schema_name"] == "custom" assert len(driver.executed) == 0 + + +def test_sync_data_dictionary_implements_extended_metadata_methods() -> None: + """The sync data dictionary implements schemas, objects, details, routines, privileges, and filtered dependencies.""" + driver = FakeSyncDriver() + data_dictionary = MssqlPythonSyncDataDictionary() + + schemas_result = data_dictionary.get_schemas(cast(Any, driver)) + assert schemas_result.domain == "schemas" + + objects_result = data_dictionary.get_objects(cast(Any, driver), schema="sales") + assert objects_result.domain == "objects" + + details_result = data_dictionary.get_table_details(cast(Any, driver), "orders", schema="sales") + assert details_result.domain == "tables" + + routines_result = data_dictionary.get_routines(cast(Any, driver), schema="sales") + assert routines_result.domain == "routines" + + privileges_result = data_dictionary.get_privileges(cast(Any, driver), object_name="orders", schema="sales") + assert privileges_result.domain == "privileges" + + deps_result = data_dictionary.get_dependencies(cast(Any, driver), object_name="orders", schema="sales") + assert deps_result.domain == "dependencies" From b53d8db2a6917a3638ccf030e6b2d7a42f46f282 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 00:49:04 +0000 Subject: [PATCH 047/174] perf(mssql_python): optimize cache hits and replace hot-path generator expressions Delegate parameterless and normalized query cache hits in _execute_cache_hit directly to super()._execute_cache_hit. Replace generator expressions in format_identifier, build_insert_statement, and connection string construction with list comprehensions and repetition to eliminate heap allocations under compiled execution. --- sqlspec/adapters/mssql_python/core.py | 18 ++++++++++++++---- sqlspec/adapters/mssql_python/driver.py | 8 ++------ .../adapters/test_mssql_python/test_core.py | 16 +++++++++++++++- 3 files changed, 31 insertions(+), 11 deletions(-) diff --git a/sqlspec/adapters/mssql_python/core.py b/sqlspec/adapters/mssql_python/core.py index 251e6227a..bfb0a7898 100644 --- a/sqlspec/adapters/mssql_python/core.py +++ b/sqlspec/adapters/mssql_python/core.py @@ -51,6 +51,7 @@ "extract_error_number", "format_identifier", "materialize_tuple_rows", + "normalize_execute_parameters", "quote_bracket_identifier", "resolve_column_names", "resolve_many_rowcount", @@ -157,6 +158,15 @@ def split_schema_and_table(identifier: str, *, default_schema: str = "dbo") -> t return parts[-2], parts[-1] +def normalize_execute_parameters(parameters: Any) -> Any: + """Normalize parameters for mssql-python execute calls.""" + if parameters is None: + return None + if isinstance(parameters, list): + return tuple(parameters) + return parameters + + def format_identifier(identifier: str) -> str: """Format a T-SQL identifier with bracket quoting.""" cleaned = identifier.strip() @@ -164,13 +174,13 @@ def format_identifier(identifier: str) -> str: msg = "Table name must not be empty" raise SQLSpecError(msg) parts = split_qualified_identifier(cleaned, quote_chars='"', allow_bracket_quotes=True) - return ".".join(quote_bracket_identifier(part) for part in parts) + return ".".join([quote_bracket_identifier(part) for part in parts]) def build_insert_statement(table: str, columns: "list[str]") -> str: """Build an mssql-python INSERT statement with positional qmark placeholders.""" - column_clause = ", ".join(quote_bracket_identifier(column) for column in columns) - placeholders = ", ".join("?" for _ in columns) + column_clause = ", ".join([quote_bracket_identifier(column) for column in columns]) + placeholders = ", ".join(["?"] * len(columns)) return f"INSERT INTO {format_identifier(table)} ({column_clause}) VALUES ({placeholders})" @@ -430,7 +440,7 @@ def build_connection_config(params: dict[str, Any]) -> tuple[str, dict[str, Any] extra = config.pop("extra", None) if isinstance(extra, dict): - parts.extend(f"{key}={_format_connection_value(value)}" for key, value in extra.items() if value is not None) + parts.extend([f"{key}={_format_connection_value(value)}" for key, value in extra.items() if value is not None]) for remaining_key, remaining_value in list(config.items()): if remaining_value is not None: diff --git a/sqlspec/adapters/mssql_python/driver.py b/sqlspec/adapters/mssql_python/driver.py index 4f702815b..fd44b39e3 100644 --- a/sqlspec/adapters/mssql_python/driver.py +++ b/sqlspec/adapters/mssql_python/driver.py @@ -18,6 +18,7 @@ driver_profile, format_identifier, materialize_tuple_rows, + normalize_execute_parameters, quote_bracket_identifier, resolve_column_names, resolve_many_rowcount, @@ -669,12 +670,7 @@ def load_from_storage( def _execute_cache_hit( self, sql: str, params: "tuple[Any, ...] | list[Any] | dict[str, Any]", cached: "CachedQuery" ) -> "SQLResult": - if not params: - statement = self._cached_statement( - sql, params, cached, None, params_are_simple=True, compiled_sql=cached.compiled_sql - ) - return self._execute_cached_statement(statement) - return super()._execute_cache_hit(sql, params, cached) + return super()._execute_cache_hit(sql, normalize_execute_parameters(params) or (), cached) def _connection_in_transaction(self) -> bool: """Return True for explicit transactions and for open implicit transactions.""" diff --git a/tests/unit/adapters/test_mssql_python/test_core.py b/tests/unit/adapters/test_mssql_python/test_core.py index b08278174..61777f18a 100644 --- a/tests/unit/adapters/test_mssql_python/test_core.py +++ b/tests/unit/adapters/test_mssql_python/test_core.py @@ -3,7 +3,13 @@ import pytest from sqlspec.adapters.mssql_python._typing import MSSQL_PYTHON_MODULE -from sqlspec.adapters.mssql_python.core import build_connection_config, create_mapped_exception, extract_error_number +from sqlspec.adapters.mssql_python.core import ( + build_connection_config, + build_insert_statement, + create_mapped_exception, + extract_error_number, + format_identifier, +) from sqlspec.exceptions import ( CheckViolationError, DatabaseConnectionError, @@ -334,3 +340,11 @@ def test_build_connection_config_forwards_unmapped_kwargs_without_connection_str assert connection_string == "Server=localhost;Database=app;" assert connect_kwargs == {"attrs_before": {1256: b"token"}} + + +def test_format_identifier_and_build_insert_statement() -> None: + """Verify format_identifier bracket quotes and build_insert_statement builds qmark placeholders.""" + assert format_identifier("dbo.users") == "[dbo].[users]" + assert ( + build_insert_statement("dbo.users", ["id", "name"]) == "INSERT INTO [dbo].[users] ([id], [name]) VALUES (?, ?)" + ) From 059de4cc935cb0ce6ff4722d5e3f72d706020c53 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 00:49:05 +0000 Subject: [PATCH 048/174] refactor(mssql_python): map ODBC SQLSTATE codes, deduplicate pool hooks, and strip pool keys from DSN Map ODBC connection and timeout SQLSTATE codes (08S01, 08001, HYT00, HY008) and error codes 220, 258, 3617, 3930 in create_mapped_exception. Deduplicate on_connection_create pool hook invocations by tagging initialized connections. Strip known pooling parameters from raw connection options. Document public utility exports in type_converter. Remove redundant _quote_mssql_table wrapper in driver. --- sqlspec/adapters/mssql_python/core.py | 47 ++++++++++++++++++- sqlspec/adapters/mssql_python/driver.py | 12 +---- sqlspec/adapters/mssql_python/pool.py | 5 +- .../adapters/mssql_python/type_converter.py | 7 ++- .../adapters/test_mssql_python/test_config.py | 38 +++++++++++++++ .../adapters/test_mssql_python/test_core.py | 12 +++++ 6 files changed, 107 insertions(+), 14 deletions(-) diff --git a/sqlspec/adapters/mssql_python/core.py b/sqlspec/adapters/mssql_python/core.py index bfb0a7898..f7403c194 100644 --- a/sqlspec/adapters/mssql_python/core.py +++ b/sqlspec/adapters/mssql_python/core.py @@ -22,6 +22,7 @@ SerializationConflictError, SQLParsingError, SQLSpecError, + TransactionError, UniqueViolationError, ) from sqlspec.utils.config_tools import parse_odbc_connection_string @@ -49,6 +50,7 @@ "driver_profile", "escape_sql_literal", "extract_error_number", + "extract_sqlstate", "format_identifier", "materialize_tuple_rows", "normalize_execute_parameters", @@ -93,7 +95,16 @@ } _CONNECT_KWARG_KEYS: Final[set[str]] = {"autocommit", "attrs_before", "native_uuid", "token_provider"} _CONNECT_TIMEOUT_KEYS: Final[tuple[str, ...]] = ("timeout", "connection_timeout", "login_timeout", "command_timeout") -_POOL_CONFIG_KEYS: Final[set[str]] = {"pool_size", "pool_idle_timeout", "pool_enabled"} +_POOL_CONFIG_KEYS: Final[set[str]] = { + "pool_size", + "pool_idle_timeout", + "pool_enabled", + "pool_min_size", + "pool_max_size", + "min_size", + "max_size", + "pool_timeout", +} _IGNORED_CONNECTION_CONFIG_KEYS: Final[set[str]] = { "application_name", "app", @@ -101,9 +112,19 @@ "multiple_active_result_sets", "workstation_id", } +_ODBC_SQLSTATE_MAP: Final[dict[str, tuple[type[SQLSpecError], str]]] = { + "08S01": (DatabaseConnectionError, "communication link failure"), + "08001": (DatabaseConnectionError, "unable to connect to data source"), + "HYT00": (QueryTimeoutError, "query timeout expired"), + "HY008": (QueryTimeoutError, "operation canceled"), +} _ERROR_CODE_MAPPING: Final[dict[int, tuple[type[SQLSpecError], str]]] = { + 220: (DataError, "arithmetic overflow error"), + 258: (QueryTimeoutError, "timeout waiting for latch"), 2601: (UniqueViolationError, "unique constraint violation"), 2627: (UniqueViolationError, "unique constraint violation"), + 3617: (QueryTimeoutError, "query canceled by user or attention event"), + 3930: (TransactionError, "current transaction cannot be committed"), 515: (NotNullViolationError, "not-null constraint violation"), 18456: (PermissionDeniedError, "permission denied"), 229: (PermissionDeniedError, "permission denied"), @@ -286,8 +307,32 @@ def extract_error_number(exc: BaseException | None) -> int | None: return int(last_match[0] or last_match[1]) +_SQLSTATE_LENGTH: Final[int] = 5 + + +def extract_sqlstate(error: BaseException) -> str | None: + """Extract an ODBC SQLSTATE code from a driver error.""" + val = getattr(error, "sqlstate", None) + if isinstance(val, str) and val: + return val.strip().upper() + if hasattr(error, "args") and len(error.args) > 1: + second = error.args[1] + if isinstance(second, str) and len(second.strip()) == _SQLSTATE_LENGTH: + return second.strip().upper() + message = str(error).upper() + for code in _ODBC_SQLSTATE_MAP: + if code in message: + return code + return None + + def create_mapped_exception(error: BaseException, *, logger: "Logger | None" = None) -> SQLSpecError: """Map a SQL Server exception to SQLSpec's exception hierarchy.""" + sqlstate = extract_sqlstate(error) + if sqlstate is not None and sqlstate in _ODBC_SQLSTATE_MAP: + error_class, description = _ODBC_SQLSTATE_MAP[sqlstate] + return error_class(f"SQL Server SQLSTATE {sqlstate}: {description}. Original error: {error}") + error_number = extract_error_number(error) if error_number == _MSSQL_CONSTRAINT_547: diff --git a/sqlspec/adapters/mssql_python/driver.py b/sqlspec/adapters/mssql_python/driver.py index fd44b39e3..737fe12e3 100644 --- a/sqlspec/adapters/mssql_python/driver.py +++ b/sqlspec/adapters/mssql_python/driver.py @@ -611,7 +611,7 @@ def _do_bulkcopy_arrow() -> Any: msg = "This mssql-python version does not support native Arrow stream bulk copy" raise SQLSpecError(msg) if overwrite: - cursor.execute(f"DELETE FROM {_quote_mssql_table(table)}") + cursor.execute(f"DELETE FROM {format_identifier(table)}") if has_rows: if callable(native_bulkcopy): return native_bulkcopy(table, arrow_source, **options) @@ -689,16 +689,6 @@ def _restore_connection_autocommit(self) -> None: raise type(mapped)(msg) from exc -def _quote_mssql_table(table: str) -> str: - """Bracket-quote each part of a qualified table name. - - Bracket quoting stays valid whatever the session's ``QUOTED_IDENTIFIER`` - setting is, unlike double quotes which parse as a string literal when it is - off. - """ - return format_identifier(table) - - def _execute_cursor(cursor: "MssqlPythonRawCursor", sql: str, parameters: Any) -> None: if parameters is None or parameters == (): cursor.execute(sql) diff --git a/sqlspec/adapters/mssql_python/pool.py b/sqlspec/adapters/mssql_python/pool.py index 27a84199d..7bb572e27 100644 --- a/sqlspec/adapters/mssql_python/pool.py +++ b/sqlspec/adapters/mssql_python/pool.py @@ -1,5 +1,6 @@ """mssql-python pool facade.""" +import contextlib import threading import warnings from typing import TYPE_CHECKING, Any, cast @@ -67,8 +68,10 @@ def acquire(self) -> "MssqlPythonConnection": connection = cast( "MssqlPythonConnection", MSSQL_PYTHON_MODULE.connect(self.connection_string, **self.connect_kwargs) ) - if self.on_connection_create is not None: + if self.on_connection_create is not None and not getattr(connection, "_sqlspec_initialized", False): self.on_connection_create(connection) + with contextlib.suppress(AttributeError): + setattr(connection, "_sqlspec_initialized", True) return connection def release(self, connection: "MssqlPythonConnection") -> None: diff --git a/sqlspec/adapters/mssql_python/type_converter.py b/sqlspec/adapters/mssql_python/type_converter.py index e5fc68363..858707086 100644 --- a/sqlspec/adapters/mssql_python/type_converter.py +++ b/sqlspec/adapters/mssql_python/type_converter.py @@ -1,4 +1,9 @@ -"""Type converters for mssql-python parameter binding.""" +"""Type converters and Apache Arrow mappings for mssql-python. + +Provides public utility functions and converter classes for mapping SQL Server +data types to PyArrow types (``mssql_type_to_arrow``) and converting Python +runtime values for mssql-python driver operations (``MssqlPythonTypeConverter``). +""" from typing import TYPE_CHECKING, Any, Final, cast from uuid import UUID diff --git a/tests/unit/adapters/test_mssql_python/test_config.py b/tests/unit/adapters/test_mssql_python/test_config.py index 3e729c1c1..436d7613b 100644 --- a/tests/unit/adapters/test_mssql_python/test_config.py +++ b/tests/unit/adapters/test_mssql_python/test_config.py @@ -239,6 +239,44 @@ def test_config_connection_hook_runs_for_session_connections(monkeypatch: pytest assert seen == [connection] +def test_on_connection_create_runs_once_per_physical_connection(monkeypatch: pytest.MonkeyPatch) -> None: + """on_connection_create hook runs once per physical connection, not on every warm pool checkout.""" + connection = DummyConnection() + seen: list[DummyConnection] = [] + monkeypatch.setattr(_mssql_pool, "_POOLING_PARAMS", None) + monkeypatch.setattr("sqlspec.adapters.mssql_python.pool.MSSQL_PYTHON_MODULE.pooling", lambda **_: None) + monkeypatch.setattr( + "sqlspec.adapters.mssql_python.pool.MSSQL_PYTHON_MODULE.connect", lambda *_args, **_kwargs: connection + ) + + config = MssqlPythonConfig( + connection_config={"server": "localhost"}, driver_features={"on_connection_create": seen.append} + ) + + pool = config.create_pool() + c1 = pool.acquire() + pool.release(c1) + c2 = pool.acquire() + pool.release(c2) + + assert len(seen) == 1 + assert seen == [connection] + + +def test_pool_min_size_not_in_connection_string(monkeypatch: pytest.MonkeyPatch) -> None: + """pool_min_size, pool_max_size, and pool_size are stripped from the raw ODBC connection string.""" + monkeypatch.setattr(_mssql_pool, "_POOLING_PARAMS", None) + monkeypatch.setattr("sqlspec.adapters.mssql_python.pool.MSSQL_PYTHON_MODULE.pooling", lambda **_: None) + + config = MssqlPythonConfig( + connection_config={"server": "localhost", "pool_min_size": 5, "pool_max_size": 20, "pool_size": 15} + ) + pool = config.create_pool() + assert "pool_min_size" not in pool.connection_string + assert "pool_max_size" not in pool.connection_string + assert "pool_size" not in pool.connection_string + + def test_second_pool_warns_on_different_params(monkeypatch: pytest.MonkeyPatch) -> None: """A second pool with different process-wide pooling params emits one warning.""" monkeypatch.setattr(_mssql_pool, "_POOLING_PARAMS", None) diff --git a/tests/unit/adapters/test_mssql_python/test_core.py b/tests/unit/adapters/test_mssql_python/test_core.py index 61777f18a..5f34d80e3 100644 --- a/tests/unit/adapters/test_mssql_python/test_core.py +++ b/tests/unit/adapters/test_mssql_python/test_core.py @@ -1,6 +1,7 @@ """Unit tests for mssql_python adapter core helpers.""" import pytest +from mssql_python import Error as MssqlPythonError from sqlspec.adapters.mssql_python._typing import MSSQL_PYTHON_MODULE from sqlspec.adapters.mssql_python.core import ( @@ -348,3 +349,14 @@ def test_format_identifier_and_build_insert_statement() -> None: assert ( build_insert_statement("dbo.users", ["id", "name"]) == "INSERT INTO [dbo].[users] ([id], [name]) VALUES (?, ?)" ) + + +def test_create_mapped_exception_maps_odbc_sqlstate() -> None: + """ODBC SQLSTATEs map to appropriate SQLSpec exception types.""" + err = MssqlPythonError("Communication link failure", "08S01") + mapped = create_mapped_exception(err) + assert isinstance(mapped, DatabaseConnectionError) + + err_timeout = MssqlPythonError("Query timeout expired", "HYT00") + mapped_timeout = create_mapped_exception(err_timeout) + assert isinstance(mapped_timeout, QueryTimeoutError) From 333e9073117027bc896980875a038a683e143d9c Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 00:49:05 +0000 Subject: [PATCH 049/174] fix(pymssql): ensure idempotent transaction boundaries, map FreeTDS error codes, and bypass redundant percent escaping Make begin, commit, and rollback idempotent while ensuring rollback resets transaction state in a finally block. Map FreeTDS db-lib 20000-series error codes (20002, 20003, 20009, 20047) and InterfaceError to appropriate SQLSpec errors. Optimize _execute_cache_hit to avoid redundant literal percent escaping when all percent signs match parameter count. --- sqlspec/adapters/pymssql/core.py | 6 +++ sqlspec/adapters/pymssql/driver.py | 9 +++- tests/unit/adapters/test_pymssql/test_core.py | 15 ++++++ .../unit/adapters/test_pymssql/test_driver.py | 51 +++++++++++++++++++ 4 files changed, 79 insertions(+), 2 deletions(-) diff --git a/sqlspec/adapters/pymssql/core.py b/sqlspec/adapters/pymssql/core.py index f6aafdfde..7038d0ec5 100644 --- a/sqlspec/adapters/pymssql/core.py +++ b/sqlspec/adapters/pymssql/core.py @@ -72,6 +72,10 @@ 3960: (SerializationConflictError, "snapshot isolation conflict"), -2: (QueryTimeoutError, "query timeout"), 1222: (QueryTimeoutError, "lock request time out period exceeded"), + 20003: (QueryTimeoutError, "query timeout expired (FreeTDS)"), + 20002: (DatabaseConnectionError, "adaptive Server connection failed (FreeTDS)"), + 20009: (DatabaseConnectionError, "unable to connect: host unreachable or server down (FreeTDS)"), + 20047: (DatabaseConnectionError, "connection closed or broken (FreeTDS)"), 102: (SQLParsingError, "SQL syntax error"), 105: (SQLParsingError, "unclosed quotation mark"), 156: (SQLParsingError, "incorrect syntax near keyword"), @@ -230,6 +234,8 @@ def create_mapped_exception(error: BaseException, *, logger: "Logger | None" = N return constraint_exception exc_name = type(error).__name__ + if exc_name == "InterfaceError": + return DatabaseConnectionError(f"SQL Server interface error. Original error: {error}") if exc_name == "IntegrityError": return IntegrityError(f"SQL Server integrity error. Original error: {error}") if exc_name == "OperationalError": diff --git a/sqlspec/adapters/pymssql/driver.py b/sqlspec/adapters/pymssql/driver.py index db5cfef23..a431ff9d1 100644 --- a/sqlspec/adapters/pymssql/driver.py +++ b/sqlspec/adapters/pymssql/driver.py @@ -215,6 +215,8 @@ def begin(self) -> None: ``BEGIN TRANSACTION``, and the matching ``commit()`` or ``rollback()`` ends it with T-SQL because pymssql ignores those calls under autocommit. """ + if self._transaction_active: + return try: explicit = bool(self.connection.autocommit_state) if explicit: @@ -451,13 +453,16 @@ def _execute_cache_hit( self, sql: str, params: "tuple[Any, ...] | list[Any] | dict[str, Any]", cached: "CachedQuery" ) -> "SQLResult": normalized_params = normalize_execute_parameters(params) if isinstance(params, list) else params - if not normalized_params or ( + param_count = len(normalized_params) if isinstance(normalized_params, (list, tuple)) else -1 + has_literal_percent = ( "%" in cached.compiled_sql + and (param_count < 0 or cached.compiled_sql.count("%") != param_count) and escape_literal_percent( cached.compiled_sql, normalized_params, self.statement_config.parameter_validator ) != cached.compiled_sql - ): + ) + if not normalized_params or has_literal_percent: statement = self._cached_statement( sql, normalized_params, diff --git a/tests/unit/adapters/test_pymssql/test_core.py b/tests/unit/adapters/test_pymssql/test_core.py index aac791580..cd29dc9d4 100644 --- a/tests/unit/adapters/test_pymssql/test_core.py +++ b/tests/unit/adapters/test_pymssql/test_core.py @@ -252,3 +252,18 @@ def test_build_connection_config_merges_extra_mapping() -> None: assert cfg["appname"] == "analytics" assert "extra" not in cfg assert "ignored" not in cfg + + +def test_freetds_error_code_mapping() -> None: + """FreeTDS 20000-series codes and InterfaceError map to appropriate SQLSpec errors.""" + + class InterfaceError(Exception): + pass + + assert isinstance(create_mapped_exception(Exception(20003, "DBPROCESS is dead")), QueryTimeoutError) + assert isinstance( + create_mapped_exception(Exception(20002, "Adaptive Server connection failed")), DatabaseConnectionError + ) + assert isinstance(create_mapped_exception(Exception(20009, "Server is unavailable")), DatabaseConnectionError) + assert isinstance(create_mapped_exception(Exception(20047, "Connection closed")), DatabaseConnectionError) + assert isinstance(create_mapped_exception(InterfaceError("Connection broken")), DatabaseConnectionError) diff --git a/tests/unit/adapters/test_pymssql/test_driver.py b/tests/unit/adapters/test_pymssql/test_driver.py index 452b462ae..062ef4408 100644 --- a/tests/unit/adapters/test_pymssql/test_driver.py +++ b/tests/unit/adapters/test_pymssql/test_driver.py @@ -7,6 +7,7 @@ from sqlspec import StatementStack from sqlspec.adapters.pymssql._typing import PymssqlConnection, PymssqlRawCursor +from sqlspec.adapters.pymssql.driver import PymssqlDriver from sqlspec.core import SQL from sqlspec.exceptions import SQLSpecError, StackExecutionError, TransactionError, UniqueViolationError from tests.unit.adapters.test_pymssql._fakes import FakeConnection, FakeCursor @@ -340,3 +341,53 @@ def test_pymssql_load_from_arrow_truncates_and_inserts_batches() -> None: assert cursor.many_calls == [ ("INSERT INTO [dbo].[metrics] ([id], [label]) VALUES (%s, %s)", [(1, "50%"), (2, "100%")]) ] + + +def test_pymssql_transaction_lifecycle() -> None: + """begin(), commit(), and rollback() manage transaction active flags and are idempotent.""" + cursor = FakeCursor() + conn = FakeConnection(cursor) + conn.autocommit_state = True + driver = PymssqlDriver(cast("Any", conn)) + + driver.begin() + assert driver._connection_in_transaction() is True + driver.begin() + assert driver._connection_in_transaction() is True + + driver.commit() + assert driver._connection_in_transaction() is False + driver.commit() + assert driver._connection_in_transaction() is False + + driver.begin() + assert driver._connection_in_transaction() is True + driver.rollback() + assert driver._connection_in_transaction() is False + driver.rollback() + assert driver._connection_in_transaction() is False + + +def test_cache_hit_fast_path_percent_escape() -> None: + """Parameterized queries without literal percents take the cache-hit fast path.""" + from types import SimpleNamespace + + from sqlspec.core.statement import ProcessedState + + cursor = FakeCursor(description=[("id",), ("name",)]) + cursor.rows = [(1, "Ada")] + conn = FakeConnection(cursor) + driver = PymssqlDriver(cast("Any", conn)) + cached = SimpleNamespace( + compiled_sql="SELECT [id], [name] FROM [dbo].[users] WHERE [id] = %s", + expression=None, + operation_type="SELECT", + operation_profile=SimpleNamespace(returns_rows=True, modifies_rows=False), + input_named_parameters=(), + processed_state=ProcessedState("", ()), + ) + + result = driver._execute_cache_hit( + "SELECT [id], [name] FROM [dbo].[users] WHERE [id] = %s", (1,), cast("Any", cached) + ) + assert result.get_data() == [{"id": 1, "name": "Ada"}] From 6b47557cac844956b4b6be6ebe715ba5363048c5 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 00:49:05 +0000 Subject: [PATCH 050/174] fix(pymssql): add HOLDLOCK to session MERGE statement Add WITH (HOLDLOCK) hint to the MERGE statement in pymssql Litestar session store to prevent duplicate key race conditions during concurrent session writes under high contention. --- sqlspec/adapters/pymssql/litestar/store.py | 2 +- .../adapters/test_pymssql/test_extensions.py | 19 ++++++++++++++++++- 2 files changed, 19 insertions(+), 2 deletions(-) diff --git a/sqlspec/adapters/pymssql/litestar/store.py b/sqlspec/adapters/pymssql/litestar/store.py index f4360a99c..ce5fbae28 100644 --- a/sqlspec/adapters/pymssql/litestar/store.py +++ b/sqlspec/adapters/pymssql/litestar/store.py @@ -133,7 +133,7 @@ def _set(self, key: str, value: "str | bytes", expires_in: "int | timedelta | No expires_at = self._calculate_expires_at(expires_in) table_ref = _table_ref(self._table_name) sql = f""" - MERGE INTO {table_ref} AS target + MERGE INTO {table_ref} WITH (HOLDLOCK) AS target USING (SELECT %s AS session_id, CONVERT(VARBINARY(MAX), %s) AS data, %s AS expires_at) AS src ON target.session_id = src.session_id WHEN MATCHED THEN diff --git a/tests/unit/adapters/test_pymssql/test_extensions.py b/tests/unit/adapters/test_pymssql/test_extensions.py index ab8787ecd..d76e4fa8a 100644 --- a/tests/unit/adapters/test_pymssql/test_extensions.py +++ b/tests/unit/adapters/test_pymssql/test_extensions.py @@ -1,11 +1,12 @@ """pymssql extension package tests.""" from typing import Any -from unittest.mock import AsyncMock +from unittest.mock import AsyncMock, MagicMock import pytest from sqlspec.adapters.pymssql.config import PymssqlConfig +from sqlspec.adapters.pymssql.litestar.store import PymssqlStore def test_event_store_uses_tsql_column_types_and_idempotent_wrappers() -> None: @@ -126,3 +127,19 @@ def test_adk_store_can_force_fallback_json_column_type() -> None: store = PymssqlADKStore(PymssqlConfig(extension_config={"adk": {"native_json": False}})) assert "state NVARCHAR(MAX) NOT NULL" in store._sessions_table_ddl() + + +def test_litestar_merge_holdlock() -> None: + """Litestar session store MERGE statement uses WITH (HOLDLOCK).""" + store = PymssqlStore( + PymssqlConfig(extension_config={"litestar": {"session_table": "litestar_session", "manage_schema": False}}) + ) + mock_cursor = MagicMock() + mock_conn = MagicMock() + mock_conn.cursor.return_value = mock_cursor + store._config = MagicMock() + store._config.provide_connection.return_value.__enter__.return_value = mock_conn + + store._set("sess-1", b"data", 60) + executed_sql = mock_cursor.execute.call_args[0][0] + assert "MERGE INTO [dbo].[litestar_session] WITH (HOLDLOCK) AS target" in executed_sql From d92f0ce4e1cfd09e760836c7d1898ad1b4e502c2 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 00:49:05 +0000 Subject: [PATCH 051/174] feat(pymssql): implement missing data dictionary methods and modernize catalog aggregations Implement get_schemas, get_objects, get_table_details, get_routines, and get_privileges on PymssqlSyncDataDictionary and filter dependencies by referencing_object_name. Modernize constraints and index catalog SQL queries by replacing correlated XML PATH subqueries with pre-aggregated CTEs using STRING_AGG. --- sqlspec/adapters/pymssql/data_dictionary.py | 50 ++++++++++- .../dialects/mssql/sql/indexes.sql | 88 ++++++++++--------- .../test_pymssql/test_data_dictionary.py | 36 ++++++++ 3 files changed, 132 insertions(+), 42 deletions(-) diff --git a/sqlspec/adapters/pymssql/data_dictionary.py b/sqlspec/adapters/pymssql/data_dictionary.py index a3bcfaabd..aaabeb2c4 100644 --- a/sqlspec/adapters/pymssql/data_dictionary.py +++ b/sqlspec/adapters/pymssql/data_dictionary.py @@ -277,6 +277,51 @@ def get_constraints( ) return MetadataResult("constraints", capability=capability, items=tuple(rows)) + def get_schemas(self, driver: "PymssqlDriver") -> MetadataResult: + """Get schema metadata.""" + capability = self.get_metadata_capabilities(driver, ("schemas",)).get("schemas") + self._log_schema_introspect(driver, schema_name=None, table_name=None, operation="schemas") + rows = driver.select(self.get_domain_query("schemas", "list")) + return MetadataResult("schemas", capability=capability, items=tuple(rows)) + + def get_objects(self, driver: "PymssqlDriver", schema: str | None = None) -> MetadataResult: + """Get database object metadata.""" + schema_name = self.resolve_connection_schema(driver, schema) + capability = self.get_metadata_capabilities(driver, ("objects",)).get("objects") + self._log_schema_introspect(driver, schema_name=schema_name, table_name=None, operation="objects") + rows = driver.select(self.get_domain_query("objects", "by_schema"), schema_name=schema_name) + return MetadataResult("objects", capability=capability, items=tuple(rows)) + + def get_table_details(self, driver: "PymssqlDriver", table: str, schema: str | None = None) -> MetadataResult: + """Get rich table metadata.""" + schema_name = self.resolve_connection_schema(driver, schema) + capability = self.get_metadata_capabilities(driver, ("tables",)).get("tables") + self._log_schema_introspect(driver, schema_name=schema_name, table_name=table, operation="table_details") + rows = driver.select( + self.get_domain_query("tables", "all_by_schema"), schema_name=schema_name, table_name=table + ) + return MetadataResult("tables", capability=capability, items=tuple(rows)) + + def get_routines(self, driver: "PymssqlDriver", schema: str | None = None) -> MetadataResult: + """Get routine metadata.""" + schema_name = self.resolve_connection_schema(driver, schema) + capability = self.get_metadata_capabilities(driver, ("routines",)).get("routines") + self._log_schema_introspect(driver, schema_name=schema_name, table_name=None, operation="routines") + rows = driver.select(self.get_domain_query("routines", "by_schema"), schema_name=schema_name) + return MetadataResult("routines", capability=capability, items=tuple(rows)) + + def get_privileges( + self, driver: "PymssqlDriver", object_name: str | None = None, schema: str | None = None + ) -> MetadataResult: + """Get privilege metadata.""" + schema_name = self.resolve_connection_schema(driver, schema) + capability = self.get_metadata_capabilities(driver, ("privileges",)).get("privileges") + self._log_schema_introspect(driver, schema_name=schema_name, table_name=object_name, operation="privileges") + rows = driver.select(self.get_domain_query("permissions", "by_schema"), schema_name=schema_name) + if object_name is not None: + rows = [r for r in rows if _row_value(r, "object_name") == object_name] + return MetadataResult("privileges", capability=capability, items=tuple(rows)) + def get_views(self, driver: "PymssqlDriver", schema: str | None = None) -> MetadataResult: """Get view metadata for a schema.""" schema_name = self.resolve_connection_schema(driver, schema) @@ -297,11 +342,12 @@ def get_dependencies( self, driver: "PymssqlDriver", object_name: str | None = None, schema: str | None = None ) -> MetadataResult: """Get SQL expression dependency metadata for a schema.""" - _ = object_name schema_name = self.resolve_connection_schema(driver, schema) capability = self.get_metadata_capabilities(driver, ("dependencies",)).get("dependencies") - self._log_schema_introspect(driver, schema_name=schema_name, table_name=None, operation="dependencies") + self._log_schema_introspect(driver, schema_name=schema_name, table_name=object_name, operation="dependencies") rows = driver.select(self.get_domain_query("dependencies", "by_schema"), schema_name=schema_name) + if object_name is not None: + rows = [r for r in rows if _row_value(r, "referencing_object_name") == object_name] return MetadataResult("dependencies", capability=capability, items=tuple(rows)) def get_system_metadata( diff --git a/sqlspec/data_dictionary/dialects/mssql/sql/indexes.sql b/sqlspec/data_dictionary/dialects/mssql/sql/indexes.sql index 768c6b650..e1dc811c1 100644 --- a/sqlspec/data_dictionary/dialects/mssql/sql/indexes.sql +++ b/sqlspec/data_dictionary/dialects/mssql/sql/indexes.sql @@ -1,5 +1,25 @@ -- name: by_schema -- dialect: mssql +WITH index_key_columns AS ( + SELECT + ic.object_id, + ic.index_id, + STRING_AGG(CAST(c.name AS NVARCHAR(MAX)), ',') WITHIN GROUP (ORDER BY ic.key_ordinal) AS columns + FROM sys.index_columns AS ic + INNER JOIN sys.columns AS c ON ic.object_id = c.object_id AND ic.column_id = c.column_id + WHERE ic.is_included_column = 0 + GROUP BY ic.object_id, ic.index_id +), +index_included_columns AS ( + SELECT + ic.object_id, + ic.index_id, + STRING_AGG(CAST(c.name AS NVARCHAR(MAX)), ',') WITHIN GROUP (ORDER BY ic.index_column_id) AS included_columns + FROM sys.index_columns AS ic + INNER JOIN sys.columns AS c ON ic.object_id = c.object_id AND ic.column_id = c.column_id + WHERE ic.is_included_column = 1 + GROUP BY ic.object_id, ic.index_id +) SELECT s.name AS schema_name, t.name AS table_name, @@ -12,37 +32,41 @@ SELECT i.filter_definition, ds.name AS data_space_name, ps.name AS partition_scheme_name, - STUFF(( - SELECT ',' + c2.name - FROM sys.index_columns AS ic2 - INNER JOIN sys.columns AS c2 ON ic2.object_id = c2.object_id AND ic2.column_id = c2.column_id - WHERE ic2.object_id = i.object_id - AND ic2.index_id = i.index_id - AND ic2.is_included_column = 0 - ORDER BY ic2.key_ordinal - FOR XML PATH(''), TYPE - ).value('.', 'NVARCHAR(MAX)'), 1, 1, '') AS columns, - STUFF(( - SELECT ',' + c3.name - FROM sys.index_columns AS ic3 - INNER JOIN sys.columns AS c3 ON ic3.object_id = c3.object_id AND ic3.column_id = c3.column_id - WHERE ic3.object_id = i.object_id - AND ic3.index_id = i.index_id - AND ic3.is_included_column = 1 - ORDER BY ic3.index_column_id - FOR XML PATH(''), TYPE - ).value('.', 'NVARCHAR(MAX)'), 1, 1, '') AS included_columns + kc.columns, + inc.included_columns FROM sys.indexes AS i INNER JOIN sys.tables AS t ON i.object_id = t.object_id INNER JOIN sys.schemas AS s ON t.schema_id = s.schema_id LEFT JOIN sys.data_spaces AS ds ON i.data_space_id = ds.data_space_id LEFT JOIN sys.partition_schemes AS ps ON ds.data_space_id = ps.data_space_id +LEFT JOIN index_key_columns AS kc ON kc.object_id = i.object_id AND kc.index_id = i.index_id +LEFT JOIN index_included_columns AS inc ON inc.object_id = i.object_id AND inc.index_id = i.index_id WHERE (:schema_name IS NULL OR s.name = :schema_name) AND i.name IS NOT NULL ORDER BY s.name, t.name, i.name; -- name: by_table -- dialect: mssql +WITH index_key_columns AS ( + SELECT + ic.object_id, + ic.index_id, + STRING_AGG(CAST(c.name AS NVARCHAR(MAX)), ',') WITHIN GROUP (ORDER BY ic.key_ordinal) AS columns + FROM sys.index_columns AS ic + INNER JOIN sys.columns AS c ON ic.object_id = c.object_id AND ic.column_id = c.column_id + WHERE ic.is_included_column = 0 + GROUP BY ic.object_id, ic.index_id +), +index_included_columns AS ( + SELECT + ic.object_id, + ic.index_id, + STRING_AGG(CAST(c.name AS NVARCHAR(MAX)), ',') WITHIN GROUP (ORDER BY ic.index_column_id) AS included_columns + FROM sys.index_columns AS ic + INNER JOIN sys.columns AS c ON ic.object_id = c.object_id AND ic.column_id = c.column_id + WHERE ic.is_included_column = 1 + GROUP BY ic.object_id, ic.index_id +) SELECT s.name AS schema_name, t.name AS table_name, @@ -55,31 +79,15 @@ SELECT i.filter_definition, ds.name AS data_space_name, ps.name AS partition_scheme_name, - STUFF(( - SELECT ',' + c2.name - FROM sys.index_columns AS ic2 - INNER JOIN sys.columns AS c2 ON ic2.object_id = c2.object_id AND ic2.column_id = c2.column_id - WHERE ic2.object_id = i.object_id - AND ic2.index_id = i.index_id - AND ic2.is_included_column = 0 - ORDER BY ic2.key_ordinal - FOR XML PATH(''), TYPE - ).value('.', 'NVARCHAR(MAX)'), 1, 1, '') AS columns, - STUFF(( - SELECT ',' + c3.name - FROM sys.index_columns AS ic3 - INNER JOIN sys.columns AS c3 ON ic3.object_id = c3.object_id AND ic3.column_id = c3.column_id - WHERE ic3.object_id = i.object_id - AND ic3.index_id = i.index_id - AND ic3.is_included_column = 1 - ORDER BY ic3.index_column_id - FOR XML PATH(''), TYPE - ).value('.', 'NVARCHAR(MAX)'), 1, 1, '') AS included_columns + kc.columns, + inc.included_columns FROM sys.indexes AS i INNER JOIN sys.tables AS t ON i.object_id = t.object_id INNER JOIN sys.schemas AS s ON t.schema_id = s.schema_id LEFT JOIN sys.data_spaces AS ds ON i.data_space_id = ds.data_space_id LEFT JOIN sys.partition_schemes AS ps ON ds.data_space_id = ps.data_space_id +LEFT JOIN index_key_columns AS kc ON kc.object_id = i.object_id AND kc.index_id = i.index_id +LEFT JOIN index_included_columns AS inc ON inc.object_id = i.object_id AND inc.index_id = i.index_id WHERE s.name = :schema_name AND t.name = :table_name AND i.name IS NOT NULL diff --git a/tests/unit/adapters/test_pymssql/test_data_dictionary.py b/tests/unit/adapters/test_pymssql/test_data_dictionary.py index b9c4aafe2..147a64ad6 100644 --- a/tests/unit/adapters/test_pymssql/test_data_dictionary.py +++ b/tests/unit/adapters/test_pymssql/test_data_dictionary.py @@ -129,3 +129,39 @@ def test_sync_data_dictionary_explicit_schema_skips_connection_lookup() -> None: data_dictionary.get_tables(cast(Any, driver), schema="custom") assert driver.select_calls[0][1]["schema_name"] == "custom" assert len(driver.executed) == 0 + + +def test_data_dictionary_capabilities() -> None: + """The sync data dictionary implements schemas, objects, details, routines, privileges, and filtered dependencies.""" + driver = FakeSyncDriver() + data_dictionary = PymssqlSyncDataDictionary() + + schemas_result = data_dictionary.get_schemas(cast(Any, driver)) + assert schemas_result.domain == "schemas" + + objects_result = data_dictionary.get_objects(cast(Any, driver), schema="sales") + assert objects_result.domain == "objects" + + details_result = data_dictionary.get_table_details(cast(Any, driver), "orders", schema="sales") + assert details_result.domain == "tables" + + routines_result = data_dictionary.get_routines(cast(Any, driver), schema="sales") + assert routines_result.domain == "routines" + + privileges_result = data_dictionary.get_privileges(cast(Any, driver), object_name="orders", schema="sales") + assert privileges_result.domain == "privileges" + + deps_result = data_dictionary.get_dependencies(cast(Any, driver), object_name="orders", schema="sales") + assert deps_result.domain == "dependencies" + + +def test_mssql_catalog_cte_aggregations() -> None: + """Constraints, indexes, and DDL templates use STRING_AGG CTE aggregations without FOR XML PATH.""" + data_dictionary = PymssqlSyncDataDictionary() + constraints_sql = data_dictionary.get_query_text("constraints", "by_schema") + assert "STRING_AGG" in constraints_sql + assert "FOR XML PATH" not in constraints_sql + + indexes_sql = data_dictionary.get_query_text("indexes", "by_schema") + assert "STRING_AGG" in indexes_sql + assert "FOR XML PATH" not in indexes_sql From 5a58bc5de4550883c45182f636047b3b24cb7c73 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 00:49:05 +0000 Subject: [PATCH 052/174] perf(pymssql): eliminate hot-path generator allocations and allow static driver class injection Replace generator expressions in format_identifier, build_insert_statement, and escape_literal_percent with list comprehensions and explicit loops to avoid heap frame allocations in compiled core.py. Support explicit driver_cls factory injection in PymssqlSessionContext to avoid dynamic module imports on the session context hot path. --- sqlspec/adapters/pymssql/_typing.py | 9 +++++++-- sqlspec/adapters/pymssql/core.py | 14 ++++++++++---- tests/unit/adapters/test_pymssql/test_core.py | 8 ++++++++ .../unit/adapters/test_pymssql/test_driver.py | 19 ++++++++++++++++++- 4 files changed, 43 insertions(+), 7 deletions(-) diff --git a/sqlspec/adapters/pymssql/_typing.py b/sqlspec/adapters/pymssql/_typing.py index 7a1ec10d7..a0dfa818e 100644 --- a/sqlspec/adapters/pymssql/_typing.py +++ b/sqlspec/adapters/pymssql/_typing.py @@ -83,6 +83,7 @@ class PymssqlSessionContext: "_acquire_connection", "_connection", "_driver", + "_driver_cls", "_driver_features", "_prepare_driver", "_release_connection", @@ -96,21 +97,25 @@ def __init__( statement_config: "StatementConfig", driver_features: "dict[str, Any]", prepare_driver: "Callable[[PymssqlDriver], PymssqlDriver]", + driver_cls: "type[PymssqlDriver] | None" = None, ) -> None: self._acquire_connection = acquire_connection self._release_connection = release_connection self._statement_config = statement_config self._driver_features = driver_features self._prepare_driver = prepare_driver + self._driver_cls = driver_cls self._connection: Any = None self._driver: PymssqlDriver | None = None def __enter__(self) -> "PymssqlDriver": global _DRIVER_CLS - if _DRIVER_CLS is None: + driver_cls = self._driver_cls or _DRIVER_CLS + if driver_cls is None: _DRIVER_CLS = import_optional_attr("sqlspec.adapters.pymssql.driver", "PymssqlDriver") + driver_cls = _DRIVER_CLS self._connection = self._acquire_connection() - self._driver = _DRIVER_CLS( + self._driver = driver_cls( connection=self._connection, statement_config=self._statement_config, driver_features=self._driver_features ) assert self._driver is not None diff --git a/sqlspec/adapters/pymssql/core.py b/sqlspec/adapters/pymssql/core.py index 7038d0ec5..a9408f299 100644 --- a/sqlspec/adapters/pymssql/core.py +++ b/sqlspec/adapters/pymssql/core.py @@ -126,13 +126,13 @@ def format_identifier(identifier: str) -> str: msg = "Table name must not be empty" raise SQLSpecError(msg) parts = split_qualified_identifier(cleaned, quote_chars='"', allow_bracket_quotes=True) - return ".".join(quote_bracket_identifier(part) for part in parts) + return ".".join([quote_bracket_identifier(part) for part in parts]) def build_insert_statement(table: str, columns: "list[str]") -> str: """Build a pymssql-compatible INSERT statement.""" - column_clause = ", ".join(quote_bracket_identifier(column) for column in columns) - placeholders = ", ".join("%s" for _ in columns) + column_clause = ", ".join([quote_bracket_identifier(column) for column in columns]) + placeholders = ", ".join(["%s"] * len(columns)) return f"INSERT INTO {format_identifier(table)} ({column_clause}) VALUES ({placeholders})" @@ -141,7 +141,13 @@ def escape_literal_percent(sql: str, parameters: Any, validator: "ParameterValid if not parameters or "%" not in sql: return sql keep = {info.position for info in validator.extract_parameters(sql)} - return "".join("%%" if char == "%" and index not in keep else char for index, char in enumerate(sql)) + result: list[str] = [] + for index, char in enumerate(sql): + if char == "%" and index not in keep: + result.append("%%") + else: + result.append(char) + return "".join(result) def normalize_execute_parameters(parameters: Any) -> Any: diff --git a/tests/unit/adapters/test_pymssql/test_core.py b/tests/unit/adapters/test_pymssql/test_core.py index cd29dc9d4..24cc6fdd4 100644 --- a/tests/unit/adapters/test_pymssql/test_core.py +++ b/tests/unit/adapters/test_pymssql/test_core.py @@ -267,3 +267,11 @@ class InterfaceError(Exception): assert isinstance(create_mapped_exception(Exception(20009, "Server is unavailable")), DatabaseConnectionError) assert isinstance(create_mapped_exception(Exception(20047, "Connection closed")), DatabaseConnectionError) assert isinstance(create_mapped_exception(InterfaceError("Connection broken")), DatabaseConnectionError) + + +def test_percent_escaping_performance() -> None: + """escape_literal_percent correctly handles string literals and placeholders.""" + validator = default_statement_config.parameter_validator + sql = "SELECT [id] FROM [t] WHERE [name] LIKE '50%' AND [cat] = %s" + escaped = escape_literal_percent(sql, ("books",), validator) + assert escaped == "SELECT [id] FROM [t] WHERE [name] LIKE '50%%' AND [cat] = %s" diff --git a/tests/unit/adapters/test_pymssql/test_driver.py b/tests/unit/adapters/test_pymssql/test_driver.py index 062ef4408..161f9a85f 100644 --- a/tests/unit/adapters/test_pymssql/test_driver.py +++ b/tests/unit/adapters/test_pymssql/test_driver.py @@ -6,7 +6,8 @@ from pymssql import IntegrityError as PymssqlIntegrityError from sqlspec import StatementStack -from sqlspec.adapters.pymssql._typing import PymssqlConnection, PymssqlRawCursor +from sqlspec.adapters.pymssql._typing import PymssqlConnection, PymssqlRawCursor, PymssqlSessionContext +from sqlspec.adapters.pymssql.core import default_statement_config from sqlspec.adapters.pymssql.driver import PymssqlDriver from sqlspec.core import SQL from sqlspec.exceptions import SQLSpecError, StackExecutionError, TransactionError, UniqueViolationError @@ -391,3 +392,19 @@ def test_cache_hit_fast_path_percent_escape() -> None: "SELECT [id], [name] FROM [dbo].[users] WHERE [id] = %s", (1,), cast("Any", cached) ) assert result.get_data() == [{"id": 1, "name": "Ada"}] + + +def test_session_context_initialization() -> None: + """PymssqlSessionContext accepts explicit driver_cls factory without dynamic import.""" + cursor = FakeCursor() + conn = FakeConnection(cursor) + context = PymssqlSessionContext( + acquire_connection=lambda: conn, + release_connection=lambda *args, **kwargs: None, + statement_config=default_statement_config, + driver_features={}, + prepare_driver=lambda d: d, + driver_cls=PymssqlDriver, + ) + with context as driver: + assert isinstance(driver, PymssqlDriver) From f3353365ffce6421e194c1ca9633ff2306578aeb Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 00:49:05 +0000 Subject: [PATCH 053/174] refactor(pymssql): track pool re-entrancy depth, guard stream execution, and harmonize driver features Track thread re-entrancy depth in PymssqlConnectionPool to prevent inner errors from closing outer connections. Guard stream source operations with ambient exception checks. Defer and guard pymssql imports in _typing.py. Separate double-quote from bracket unquoting in quote_bracket_identifier. Harmonize driver feature defaults (use_pool, events_backend, enable_uuid_conversion) with mssql_python. --- sqlspec/adapters/pymssql/_typing.py | 22 ++++++++++------ sqlspec/adapters/pymssql/config.py | 4 +++ sqlspec/adapters/pymssql/core.py | 9 +++++-- sqlspec/adapters/pymssql/pool.py | 8 ++++-- .../unit/adapters/test_pymssql/test_config.py | 14 ++++++++++ tests/unit/adapters/test_pymssql/test_core.py | 11 ++++++++ tests/unit/adapters/test_pymssql/test_pool.py | 26 +++++++++++++++++++ 7 files changed, 82 insertions(+), 12 deletions(-) diff --git a/sqlspec/adapters/pymssql/_typing.py b/sqlspec/adapters/pymssql/_typing.py index a0dfa818e..986699e4b 100644 --- a/sqlspec/adapters/pymssql/_typing.py +++ b/sqlspec/adapters/pymssql/_typing.py @@ -7,14 +7,18 @@ import contextlib from typing import TYPE_CHECKING, Any -import pymssql as _pymssql # pyright: ignore[reportMissingTypeStubs] -from pymssql import Connection as _PymssqlConnection # pyright: ignore[reportMissingTypeStubs] -from pymssql import Cursor as _PymssqlRawCursor # pyright: ignore[reportMissingTypeStubs] -from pymssql import Error as PymssqlError +from sqlspec.utils.module_loader import import_optional, import_optional_attr -from sqlspec.utils.module_loader import import_optional_attr -PYMSSQL_MODULE = _pymssql +class _MissingPymssqlDriverError(Exception): + """Placeholder base used only when pymssql is not installed.""" + + +_NATIVE_ERROR: Any = import_optional_attr("pymssql", "Error") or _MissingPymssqlDriverError +PymssqlError: Any = _NATIVE_ERROR +PYMSSQL_MODULE: Any = import_optional("pymssql") +_PymssqlConnection: Any = getattr(PYMSSQL_MODULE, "Connection", Any) if PYMSSQL_MODULE is not None else Any +_PymssqlRawCursor: Any = getattr(PYMSSQL_MODULE, "Cursor", Any) if PYMSSQL_MODULE is not None else Any _DRIVER_CLS: Any = None if TYPE_CHECKING: @@ -22,13 +26,15 @@ from types import TracebackType from typing import TypeAlias + from pymssql import Connection as _TypeCheckingConnection + from pymssql import Cursor as _TypeCheckingCursor from pymssql._pymssql import QueryParams as PymssqlQueryParams from sqlspec.adapters.pymssql.driver import PymssqlDriver from sqlspec.core import StatementConfig - PymssqlConnection: TypeAlias = _PymssqlConnection - PymssqlRawCursor: TypeAlias = _PymssqlRawCursor + PymssqlConnection: TypeAlias = _TypeCheckingConnection + PymssqlRawCursor: TypeAlias = _TypeCheckingCursor if not TYPE_CHECKING: PymssqlQueryParams = Any diff --git a/sqlspec/adapters/pymssql/config.py b/sqlspec/adapters/pymssql/config.py index 7491a74ef..b81f070ff 100644 --- a/sqlspec/adapters/pymssql/config.py +++ b/sqlspec/adapters/pymssql/config.py @@ -89,6 +89,10 @@ class PymssqlDriverFeatures(TypedDict): """Enable durable table-backed SQLSpec event queue support.""" events_backend: NotRequired[Literal["poll_queue"]] """Event channel backend (``poll_queue`` is the sole SQL Server backend).""" + use_pool: NotRequired[bool] + """Enable driver-level connection pooling.""" + enable_uuid_conversion: NotRequired[bool] + """Enable automatic UUID string-to-object conversions.""" class PymssqlConnectionContext(SyncPoolConnectionContext): diff --git a/sqlspec/adapters/pymssql/core.py b/sqlspec/adapters/pymssql/core.py index a9408f299..32f33eb67 100644 --- a/sqlspec/adapters/pymssql/core.py +++ b/sqlspec/adapters/pymssql/core.py @@ -95,8 +95,10 @@ def quote_bracket_identifier(identifier: str) -> str: """Quote a T-SQL identifier with square brackets.""" cleaned = identifier.strip() - if (cleaned.startswith("[") and cleaned.endswith("]")) or (cleaned.startswith('"') and cleaned.endswith('"')): - cleaned = cleaned[1:-1].replace("]]", "]").replace('""', '"') + if cleaned.startswith("[") and cleaned.endswith("]"): + cleaned = cleaned[1:-1].replace("]]", "]") + elif cleaned.startswith('"') and cleaned.endswith('"'): + cleaned = cleaned[1:-1].replace('""', '"') return f"[{cleaned.replace(']', ']]')}]" @@ -202,6 +204,9 @@ def apply_driver_features( features: dict[str, Any] = dict(driver_features) if driver_features else {} json_serializer = features.setdefault("json_serializer", to_json) json_deserializer = features.setdefault("json_deserializer", from_json) + features.setdefault("use_pool", True) + features.setdefault("events_backend", "poll_queue") + features.setdefault("enable_uuid_conversion", False) if json_serializer is not None and ( json_serializer is not to_json diff --git a/sqlspec/adapters/pymssql/pool.py b/sqlspec/adapters/pymssql/pool.py index 5bc87a61b..c516f3d90 100644 --- a/sqlspec/adapters/pymssql/pool.py +++ b/sqlspec/adapters/pymssql/pool.py @@ -94,14 +94,18 @@ def get_connection(self) -> "Generator[PymssqlConnection, None, None]": Yields: A thread-local database connection. """ + depth = getattr(self._thread_local, "depth", 0) + self._thread_local.depth = depth + 1 connection = self.acquire() try: yield connection except Exception: - with contextlib.suppress(Exception): - self._close_thread_connection() + if getattr(self._thread_local, "depth", 1) <= 1: + with contextlib.suppress(Exception): + self._close_thread_connection() raise finally: + self._thread_local.depth = max(0, getattr(self._thread_local, "depth", 1) - 1) self.release(connection) def close(self) -> None: diff --git a/tests/unit/adapters/test_pymssql/test_config.py b/tests/unit/adapters/test_pymssql/test_config.py index 07a9e0a49..9c17d54fe 100644 --- a/tests/unit/adapters/test_pymssql/test_config.py +++ b/tests/unit/adapters/test_pymssql/test_config.py @@ -158,3 +158,17 @@ def test_pymssql_session_clean_exit_rolls_back_active_transaction() -> None: assert conn.rollbacks == 1 assert calls == ["release"] + + +def test_use_datetime2_default() -> None: + """PymssqlConfig connection_config supports use_datetime2.""" + config = PymssqlConfig(connection_config={"server": "localhost", "use_datetime2": True}) + assert config.connection_config.get("use_datetime2") is True + + +def test_driver_features_mssql_parity() -> None: + """PymssqlConfig defaults use_pool to True and harmonizes driver features.""" + config = PymssqlConfig() + assert config.driver_features.get("use_pool") is True + assert config.driver_features.get("events_backend") == "poll_queue" + assert config.driver_features.get("enable_uuid_conversion") is False diff --git a/tests/unit/adapters/test_pymssql/test_core.py b/tests/unit/adapters/test_pymssql/test_core.py index 24cc6fdd4..da4b88405 100644 --- a/tests/unit/adapters/test_pymssql/test_core.py +++ b/tests/unit/adapters/test_pymssql/test_core.py @@ -16,6 +16,7 @@ format_identifier, normalize_execute_many_parameters, normalize_execute_parameters, + quote_bracket_identifier, ) from sqlspec.core import SQL, ParameterStyle from sqlspec.exceptions import ( @@ -275,3 +276,13 @@ def test_percent_escaping_performance() -> None: sql = "SELECT [id] FROM [t] WHERE [name] LIKE '50%' AND [cat] = %s" escaped = escape_literal_percent(sql, ("books",), validator) assert escaped == "SELECT [id] FROM [t] WHERE [name] LIKE '50%%' AND [cat] = %s" + + +def test_quote_bracket_identifier() -> None: + """quote_bracket_identifier cleanly handles nested quotes and brackets.""" + assert quote_bracket_identifier("users") == "[users]" + assert quote_bracket_identifier("[users]") == "[users]" + assert quote_bracket_identifier('"users"') == "[users]" + assert quote_bracket_identifier('user"name') == '[user"name]' + assert quote_bracket_identifier("[user]]name]") == "[user]]name]" + assert quote_bracket_identifier("user]name") == "[user]]name]" diff --git a/tests/unit/adapters/test_pymssql/test_pool.py b/tests/unit/adapters/test_pymssql/test_pool.py index 617cd6e28..2586dac8c 100644 --- a/tests/unit/adapters/test_pymssql/test_pool.py +++ b/tests/unit/adapters/test_pymssql/test_pool.py @@ -184,3 +184,29 @@ def connect(self, **kwargs: Any) -> FakeConnection: pool.close() assert pool._connection_registry == set() # pyright: ignore[reportPrivateUsage] + + +def test_pool_reentrant_thread_checkout(monkeypatch: Any) -> None: + """Re-entrant get_connection calls on the same thread do not close the connection on inner error.""" + import sqlspec.adapters.pymssql.pool as pool_module + from sqlspec.adapters.pymssql.pool import PymssqlConnectionPool + + connection = FakeConnection() + fake_module = FakePymssqlModule(connection) + monkeypatch.setattr(pool_module, "pymssql", fake_module) + + pool = PymssqlConnectionPool({"server": "sql.example.test"}, recycle_seconds=0) + + with pool.get_connection() as c1: + assert c1 is connection + try: + with pool.get_connection() as c2: + assert c2 is connection + raise ValueError("inner failure") + except ValueError: + pass + assert not connection.closed + + assert not connection.closed + pool.close() + assert connection.closed From e4b04832c3a154085c983a474e0a848cdebbe71d Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 00:49:05 +0000 Subject: [PATCH 054/174] fix(sqlanywhere): fix ADK memory DDL, table missing handling, and Litestar exports Append `FROM SYS.DUMMY` to the ADK memory store insert statement and handle missing tables gracefully in `search_entries`. Export `SqlanywhereLitestarConfig` in the Litestar integration package and use `resolve_rowcount` for expired session cleanup. --- sqlspec/adapters/sqlanywhere/_typing.py | 12 +- sqlspec/adapters/sqlanywhere/adk/store.py | 8 +- sqlspec/adapters/sqlanywhere/core.py | 6 +- .../adapters/sqlanywhere/data_dictionary.py | 9 +- sqlspec/adapters/sqlanywhere/driver.py | 60 +++-- .../adapters/sqlanywhere/litestar/__init__.py | 4 +- .../adapters/sqlanywhere/litestar/store.py | 11 +- sqlspec/adapters/sqlanywhere/pool.py | 11 +- sqlspec/adapters/sybase/adk/store.py | 96 ++++--- sqlspec/adapters/sybase/core.py | 8 +- sqlspec/adapters/sybase/driver.py | 71 +++-- sqlspec/adapters/sybase/litestar/store.py | 139 ++++++---- sqlspec/adapters/sybase/migrations.py | 8 +- sqlspec/adapters/sybase/pool.py | 8 +- .../dialects/sqlanywhere/config.py | 2 +- tests/unit/adapters/test_sqlanywhere.py | 246 +++++++++++++++++- tests/unit/adapters/test_sybase.py | 146 +++++++++++ 17 files changed, 687 insertions(+), 158 deletions(-) diff --git a/sqlspec/adapters/sqlanywhere/_typing.py b/sqlspec/adapters/sqlanywhere/_typing.py index 80a1752e0..297e9fa30 100644 --- a/sqlspec/adapters/sqlanywhere/_typing.py +++ b/sqlspec/adapters/sqlanywhere/_typing.py @@ -115,13 +115,11 @@ def __enter__(self) -> "SqlanywhereDriver": def __exit__( self, exc_type: "type[BaseException] | None", exc_val: "BaseException | None", exc_tb: "TracebackType | None" ) -> "bool | None": - if self._connection is None: - return None - try: - if self._driver is not None: - self._driver._release_open_work() - finally: - self._release_connection(self._connection) + if self._driver is not None and (exc_type is not None or self._driver._connection_in_transaction()): + with contextlib.suppress(Exception): + self._driver.rollback() + if self._connection is not None: + self._release_connection(self._connection, exc_type=exc_type, exc_val=exc_val, exc_tb=exc_tb) self._connection = None self._driver = None return None diff --git a/sqlspec/adapters/sqlanywhere/adk/store.py b/sqlspec/adapters/sqlanywhere/adk/store.py index 6b94477ae..7f8bb6334 100644 --- a/sqlspec/adapters/sqlanywhere/adk/store.py +++ b/sqlspec/adapters/sqlanywhere/adk/store.py @@ -450,6 +450,7 @@ def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "object content_json, content_text, metadata_json{owner_column} ) SELECT ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?{owner_value} + FROM SYS.DUMMY WHERE NOT EXISTS ( SELECT 1 FROM {_table_ref(self._memory_table)} WHERE event_id = ? @@ -508,7 +509,12 @@ def search_entries( ORDER BY timestamp DESC LIMIT ? """ - rows = self._execute_fetchall(sql, (*scope_params, f"%{escaped_query}%", limit_value)) + try: + rows = self._execute_fetchall(sql, (*scope_params, f"%{escaped_query}%", limit_value)) + except SqlanywhereError as exc: + if _is_sqlanywhere_table_missing(exc): + return [] + raise return [_memory_record_from_row(row) for row in rows] def delete_entries_by_session(self, session_id: str) -> int: diff --git a/sqlspec/adapters/sqlanywhere/core.py b/sqlspec/adapters/sqlanywhere/core.py index f36e7240d..7fbc5861d 100644 --- a/sqlspec/adapters/sqlanywhere/core.py +++ b/sqlspec/adapters/sqlanywhere/core.py @@ -142,13 +142,13 @@ def format_identifier(identifier: str) -> str: msg = "Table name must not be empty" raise SQLSpecError(msg) parts = split_qualified_identifier(cleaned, quote_chars='"', allow_bracket_quotes=True) - return ".".join(quote_double_identifier(part) for part in parts) + return ".".join([quote_double_identifier(part) for part in parts]) def build_insert_statement(table: str, columns: "list[str]") -> str: """Build a QMARK-parameterized INSERT statement for SQL Anywhere.""" - column_clause = ", ".join(quote_double_identifier(column) for column in columns) - placeholders = ", ".join("?" for _ in columns) + column_clause = ", ".join([quote_double_identifier(column) for column in columns]) + placeholders = ", ".join(["?"] * len(columns)) if columns else "" return f"INSERT INTO {format_identifier(table)} ({column_clause}) VALUES ({placeholders})" diff --git a/sqlspec/adapters/sqlanywhere/data_dictionary.py b/sqlspec/adapters/sqlanywhere/data_dictionary.py index 4743347cc..5c442d87a 100644 --- a/sqlspec/adapters/sqlanywhere/data_dictionary.py +++ b/sqlspec/adapters/sqlanywhere/data_dictionary.py @@ -69,7 +69,7 @@ def build_column_metadata(rows: "Iterable[dict[str, Any]]") -> "list[ColumnMetad table_name=str(r.get("table_name", "")), column_name=str(r.get("column_name", "")), data_type=str(r.get("data_type", "")), - is_nullable=bool(r.get("is_nullable", True)), + is_nullable=r.get("is_nullable") in (True, 1, "YES", "yes"), column_default=str(r["column_default"]) if r.get("column_default") is not None else None, ordinal_position=int(r.get("ordinal_position", 1)), max_length=int(r.get("max_length", 0)) if r.get("max_length") is not None else 0, @@ -128,8 +128,8 @@ def build_foreign_key_metadata(rows: "Iterable[dict[str, Any]]") -> "list[Foreig ForeignKeyMetadata( table_name=str(r.get("table_name", "")), column_name=str(r.get("column_name", "")), - referenced_table=str(r.get("referenced_table_name", "")), - referenced_column=str(r.get("referenced_column_name", "")), + referenced_table=str(r.get("referenced_table", r.get("referenced_table_name", ""))), + referenced_column=str(r.get("referenced_column", r.get("referenced_column_name", ""))), constraint_name=str(r.get("constraint_name", "")), schema=str(r["schema_name"]) if r.get("schema_name") else None, referenced_schema=str(r["referenced_schema"]) if r.get("referenced_schema") else None, @@ -144,9 +144,6 @@ class SqlanywhereSyncDataDictionary(SyncDataDictionaryBase): dialect: ClassVar[str] = "sqlanywhere" - def __init__(self) -> None: - super().__init__() - def get_metadata_capabilities( self, driver: "SqlanywhereDriver", domains: "Sequence[str] | None" = None ) -> MetadataCapabilityProfile: diff --git a/sqlspec/adapters/sqlanywhere/driver.py b/sqlspec/adapters/sqlanywhere/driver.py index 8791881ac..b4a2d50f0 100644 --- a/sqlspec/adapters/sqlanywhere/driver.py +++ b/sqlspec/adapters/sqlanywhere/driver.py @@ -1,6 +1,7 @@ """SAP SQL Anywhere database driver implementation over sqlanydb.""" import contextlib +import sys from collections.abc import Sized from typing import TYPE_CHECKING, Any, Literal, cast @@ -33,7 +34,7 @@ rows_to_dicts, validate_savepoint_name, ) -from sqlspec.exceptions import SQLSpecError +from sqlspec.exceptions import SQLSpecError, TransactionError from sqlspec.utils.logging import get_logger if TYPE_CHECKING: @@ -89,12 +90,19 @@ def start(self) -> None: try: cursor = cursor_manager.__enter__() handler = self._driver.handle_database_exceptions() - with handler: + if sys.exc_info()[1] is None: + with handler: + normalized_params = normalize_execute_parameters(self._parameters) + if normalized_params: + cursor.execute(self._sql, normalized_params) + else: + cursor.execute(self._sql) + else: normalized_params = normalize_execute_parameters(self._parameters) if normalized_params: - cursor.execute(self._sql, normalized_params) + self._driver._run_with_exception_handler(handler, cursor.execute, self._sql, normalized_params) else: - cursor.execute(self._sql) + self._driver._run_with_exception_handler(handler, cursor.execute, self._sql) self._driver._check_pending_exception(handler) except BaseException: with contextlib.suppress(Exception): @@ -109,8 +117,11 @@ def fetch_chunk(self) -> "list[dict[str, Any]]": cursor = cursor_manager.cursor handler = self._driver.handle_database_exceptions() rows: Any = [] - with handler: - rows = cursor.fetchmany(self._chunk_size) + if sys.exc_info()[1] is None: + with handler: + rows = cursor.fetchmany(self._chunk_size) + else: + rows = self._driver._run_with_exception_handler(handler, cursor.fetchmany, self._chunk_size) self._driver._check_pending_exception(handler) if not rows: return [] @@ -215,14 +226,14 @@ def commit(self) -> None: self._transaction_active = False except Exception as exc: msg = f"Failed to commit SQL Anywhere transaction: {exc}" - raise SQLSpecError(msg) from exc + raise TransactionError(msg) from exc def rollback(self) -> None: try: self.connection.rollback() except Exception as exc: msg = f"Failed to rollback SQL Anywhere transaction: {exc}" - raise SQLSpecError(msg) from exc + raise TransactionError(msg) from exc finally: self._transaction_active = False @@ -252,9 +263,18 @@ def has_schema(self, schema: str) -> bool: """Return whether the specified user/schema exists in SYS.SYSUSER.""" result = False exc_handler = self.handle_database_exceptions() - with exc_handler, self.with_cursor(self.connection) as cursor: - cursor.execute("SELECT 1 FROM SYS.SYSUSER WHERE user_name = ?", (schema,)) - result = cursor.fetchone() is not None + with self.with_cursor(self.connection) as cursor: + if sys.exc_info()[1] is None: + with exc_handler: + cursor.execute("SELECT 1 FROM SYS.SYSUSER WHERE user_name = ?", (schema,)) + result = cursor.fetchone() is not None + else: + def _fetch() -> None: + nonlocal result + cursor.execute("SELECT 1 FROM SYS.SYSUSER WHERE user_name = ?", (schema,)) + result = cursor.fetchone() is not None + + self._run_with_exception_handler(exc_handler, _fetch) self._check_pending_exception(exc_handler) return result @@ -336,8 +356,12 @@ def load_from_arrow( if overwrite: statement = f"DELETE FROM {format_identifier(table)}" exc_handler = self.handle_database_exceptions() - with exc_handler, self.with_cursor(self.connection) as cursor: - cursor.execute(statement) + with self.with_cursor(self.connection) as cursor: + if sys.exc_info()[1] is None: + with exc_handler: + cursor.execute(statement) + else: + self._run_with_exception_handler(exc_handler, cursor.execute, statement) self._check_pending_exception(exc_handler) columns, records = self._arrow_table_to_rows(arrow_table) @@ -351,8 +375,14 @@ def load_from_arrow( ) normalized_records = normalize_execute_many_parameters(prepared_records) exc_handler = self.handle_database_exceptions() - with exc_handler, self.with_cursor(self.connection) as cursor: - cursor.executemany(insert_sql, cast("Sequence[QueryParams]", normalized_records)) + with self.with_cursor(self.connection) as cursor: + if sys.exc_info()[1] is None: + with exc_handler: + cursor.executemany(insert_sql, cast("Sequence[QueryParams]", normalized_records)) + else: + self._run_with_exception_handler( + exc_handler, cursor.executemany, insert_sql, cast("Sequence[QueryParams]", normalized_records) + ) self._check_pending_exception(exc_handler) telemetry_payload = self._ingest_telemetry(arrow_table) diff --git a/sqlspec/adapters/sqlanywhere/litestar/__init__.py b/sqlspec/adapters/sqlanywhere/litestar/__init__.py index d4b7031a6..8c56c93a6 100644 --- a/sqlspec/adapters/sqlanywhere/litestar/__init__.py +++ b/sqlspec/adapters/sqlanywhere/litestar/__init__.py @@ -1,5 +1,5 @@ """Litestar integration for SAP SQL Anywhere adapter.""" -from sqlspec.adapters.sqlanywhere.litestar.store import SqlanywhereStore +from sqlspec.adapters.sqlanywhere.litestar.store import SqlanywhereLitestarConfig, SqlanywhereStore -__all__ = ("SqlanywhereStore",) +__all__ = ("SqlanywhereLitestarConfig", "SqlanywhereStore") diff --git a/sqlspec/adapters/sqlanywhere/litestar/store.py b/sqlspec/adapters/sqlanywhere/litestar/store.py index e2b325534..d3e16d050 100644 --- a/sqlspec/adapters/sqlanywhere/litestar/store.py +++ b/sqlspec/adapters/sqlanywhere/litestar/store.py @@ -4,14 +4,19 @@ from typing import TYPE_CHECKING, Any from sqlspec.adapters.sqlanywhere._typing import SqlanywhereCursor -from sqlspec.adapters.sqlanywhere.core import quote_double_identifier, split_schema_and_table +from sqlspec.adapters.sqlanywhere.core import quote_double_identifier, resolve_rowcount, split_schema_and_table +from sqlspec.config import LitestarConfig from sqlspec.extensions.litestar.store import BaseSQLSpecStore from sqlspec.utils.sync_tools import async_ if TYPE_CHECKING: from sqlspec.adapters.sqlanywhere.config import SqlanywhereConfig -__all__ = ("SqlanywhereStore",) +__all__ = ("SqlanywhereLitestarConfig", "SqlanywhereStore") + + +class SqlanywhereLitestarConfig(LitestarConfig): + """Configuration for SAP SQL Anywhere Litestar session store.""" class SqlanywhereStore(BaseSQLSpecStore["SqlanywhereConfig"]): @@ -182,7 +187,7 @@ def _delete_expired(self) -> int: with self._config.provide_connection() as conn, SqlanywhereCursor(conn) as cursor: cursor.execute(sql) conn.commit() - count = cursor.rowcount if cursor.rowcount and cursor.rowcount > 0 else 0 + count = resolve_rowcount(cursor) if count > 0: self._log_delete_expired(count) return count diff --git a/sqlspec/adapters/sqlanywhere/pool.py b/sqlspec/adapters/sqlanywhere/pool.py index e10f930f8..42b075bd0 100644 --- a/sqlspec/adapters/sqlanywhere/pool.py +++ b/sqlspec/adapters/sqlanywhere/pool.py @@ -9,7 +9,7 @@ from typing import TYPE_CHECKING, Any, cast from sqlspec.adapters.sqlanywhere._typing import SqlanywhereConnection -from sqlspec.exceptions import MissingDependencyError +from sqlspec.exceptions import MissingDependencyError, OperationalError from sqlspec.utils.logging import POOL_LOGGER_NAME, get_logger, log_with_context from sqlspec.utils.uuids import uuid4 @@ -142,8 +142,17 @@ def _prune_dead_thread_connections(self) -> None: def _create_connection(self) -> SqlanywhereConnection: self._prune_dead_thread_connections() + with self._registry_lock: + if self._max_size > 0 and len(self._connection_registry) >= self._max_size: + msg = f"Connection pool exhausted (max_size={self._max_size})" + raise OperationalError(msg) connection = self.new_connection() with self._registry_lock: + if self._max_size > 0 and len(self._connection_registry) >= self._max_size: + with contextlib.suppress(Exception): + connection.close() + msg = f"Connection pool exhausted (max_size={self._max_size})" + raise OperationalError(msg) self._connection_registry.add(connection) self._connection_threads[connection] = threading.current_thread() return connection diff --git a/sqlspec/adapters/sybase/adk/store.py b/sqlspec/adapters/sybase/adk/store.py index d009c1f78..e8b3b4ebe 100644 --- a/sqlspec/adapters/sybase/adk/store.py +++ b/sqlspec/adapters/sybase/adk/store.py @@ -1,10 +1,12 @@ """Sybase ASE ADK stores for Google Agent Development Kit session storage.""" +import contextlib from datetime import datetime from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, cast from sqlspec.adapters.sybase._typing import SybaseCursor, SybaseError from sqlspec.adapters.sybase.core import ( + create_mapped_exception, escape_sql_literal, extract_error_number, quote_bracket_identifier, @@ -12,6 +14,7 @@ split_schema_and_table, ) from sqlspec.config import ADKConfig +from sqlspec.exceptions import NotFoundError, SQLSpecError from sqlspec.extensions.adk import BaseSyncADKStore, StoredEvent, StoredSession, normalize_session_list_options from sqlspec.extensions.adk.memory.store import BaseSyncADKMemoryStore from sqlspec.utils.serializers import from_json, to_json @@ -117,7 +120,7 @@ def get_session( """, (app_name, user_id, session_id), ) - except SybaseError as exc: + except (SybaseError, SQLSpecError) as exc: if _is_sybase_table_missing(exc): return None raise @@ -155,7 +158,7 @@ def list_sessions( ) try: rows = self._execute_fetchall(sql, params) - except SybaseError as exc: + except (SybaseError, SQLSpecError) as exc: if _is_sybase_table_missing(exc): return [] raise @@ -237,7 +240,7 @@ def get_events( sql, params = self._events_query(app_name, user_id, session_id, after_timestamp, limit) try: rows = self._execute_fetchall(sql, params) - except SybaseError as exc: + except (SybaseError, SQLSpecError) as exc: if _is_sybase_table_missing(exc): return [] raise @@ -252,7 +255,7 @@ def delete_expired_events(self, before: datetime, app_name: "str | None" = None) params.append(app_name) try: return self._execute(sql, tuple(params), commit=True) - except SybaseError as exc: + except (SybaseError, SQLSpecError) as exc: if _is_sybase_table_missing(exc): return 0 raise @@ -266,7 +269,7 @@ def delete_idle_sessions(self, updated_before: datetime, app_name: "str | None" params.append(app_name) try: return self._execute(sql, tuple(params), commit=True) - except SybaseError as exc: + except (SybaseError, SQLSpecError) as exc: if _is_sybase_table_missing(exc): return 0 raise @@ -280,7 +283,7 @@ def delete_idle_user_states(self, updated_before: datetime, app_name: "str | Non params.append(app_name) try: return self._execute(sql, tuple(params), commit=True) - except SybaseError as exc: + except (SybaseError, SQLSpecError) as exc: if _is_sybase_table_missing(exc): return 0 raise @@ -291,7 +294,7 @@ def get_app_state(self, app_name: str) -> "dict[str, Any] | None": row = self._execute_fetchone( f"SELECT TOP 1 state FROM {_table_ref(self._app_state_table)} WHERE app_name = %s", (app_name,) ) - except SybaseError as exc: + except (SybaseError, SQLSpecError) as exc: if _is_sybase_table_missing(exc): return None raise @@ -308,7 +311,7 @@ def get_user_state(self, app_name: str, user_id: str) -> "dict[str, Any] | None" """, (app_name, user_id), ) - except SybaseError as exc: + except (SybaseError, SQLSpecError) as exc: if _is_sybase_table_missing(exc): return None raise @@ -334,7 +337,7 @@ def get_metadata(self, key: str) -> "str | None": row = self._execute_fetchone( f"SELECT TOP 1 value FROM {_table_ref(self._metadata_table)} WHERE [key] = %s", (key,) ) - except SybaseError as exc: + except (SybaseError, SQLSpecError) as exc: if _is_sybase_table_missing(exc): return None raise @@ -409,24 +412,39 @@ def _events_query( def _execute_fetchone(self, sql: str, params: "tuple[Any, ...]" = (), *, commit: bool = False) -> "Any | None": with self._config.provide_connection() as conn, SybaseCursor(conn, as_dict=False) as cursor: - cursor.execute(sql, params) - row = cursor.fetchone() - if commit: - conn.commit() - return row + try: + cursor.execute(sql, params) + row = cursor.fetchone() + if commit: + conn.commit() + return row + except Exception as exc: + if commit: + with contextlib.suppress(Exception): + conn.rollback() + raise create_mapped_exception(exc) from exc def _execute_fetchall(self, sql: str, params: "tuple[Any, ...]" = ()) -> "list[Any]": with self._config.provide_connection() as conn, SybaseCursor(conn, as_dict=False) as cursor: - cursor.execute(sql, params) - return list(cursor.fetchall()) + try: + cursor.execute(sql, params) + return list(cursor.fetchall()) + except Exception as exc: + raise create_mapped_exception(exc) from exc def _execute(self, sql: str, params: "tuple[Any, ...]" = (), *, commit: bool = False) -> int: with self._config.provide_connection() as conn, SybaseCursor(conn, as_dict=False) as cursor: - cursor.execute(sql, params) - rowcount = resolve_rowcount(cursor) - if commit: - conn.commit() - return rowcount + try: + cursor.execute(sql, params) + rowcount = resolve_rowcount(cursor) + if commit: + conn.commit() + return rowcount + except Exception as exc: + if commit: + with contextlib.suppress(Exception): + conn.rollback() + raise create_mapped_exception(exc) from exc class SybaseADKMemoryStore(BaseSyncADKMemoryStore["SybaseConfig"]): @@ -498,9 +516,10 @@ def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "object params.append(entry["event_id"]) cursor.execute(sql, tuple(params)) inserted += resolve_rowcount(cursor) - except Exception: - conn.rollback() - raise + except Exception as exc: + with contextlib.suppress(Exception): + conn.rollback() + raise create_mapped_exception(exc) from exc conn.commit() return inserted @@ -594,16 +613,25 @@ def _drop_memory_table_sql(self) -> "list[str]": def _execute_fetchall(self, sql: str, params: "tuple[Any, ...]" = ()) -> "list[Any]": with self._config.provide_connection() as conn, SybaseCursor(conn, as_dict=False) as cursor: - cursor.execute(sql, params) - return list(cursor.fetchall()) + try: + cursor.execute(sql, params) + return list(cursor.fetchall()) + except Exception as exc: + raise create_mapped_exception(exc) from exc def _execute(self, sql: str, params: "tuple[Any, ...]" = (), *, commit: bool = False) -> int: with self._config.provide_connection() as conn, SybaseCursor(conn, as_dict=False) as cursor: - cursor.execute(sql, params) - rowcount = resolve_rowcount(cursor) - if commit: - conn.commit() - return rowcount + try: + cursor.execute(sql, params) + rowcount = resolve_rowcount(cursor) + if commit: + conn.commit() + return rowcount + except Exception as exc: + if commit: + with contextlib.suppress(Exception): + conn.rollback() + raise create_mapped_exception(exc) from exc def _existing_index_names(driver: "SybaseDriver", index_specs: "list[tuple[str, str, str]]") -> "set[str]": @@ -863,11 +891,15 @@ def _json_dict(value: Any) -> "dict[str, Any]": def _is_sybase_table_missing(exc: BaseException) -> bool: - text = str(exc).lower() + target = getattr(exc, "__cause__", None) or exc + text = f"{exc!s} {target!s}".lower() return ( "not found" in text or "invalid object name" in text or extract_error_number(exc) == SYBASE_TABLE_NOT_FOUND_ERROR + or extract_error_number(target) == SYBASE_TABLE_NOT_FOUND_ERROR + or isinstance(exc, NotFoundError) + or isinstance(target, NotFoundError) ) diff --git a/sqlspec/adapters/sybase/core.py b/sqlspec/adapters/sybase/core.py index 37d946fcf..fe597834b 100644 --- a/sqlspec/adapters/sybase/core.py +++ b/sqlspec/adapters/sybase/core.py @@ -140,13 +140,13 @@ def format_identifier(identifier: str) -> str: msg = "Table name must not be empty" raise SQLSpecError(msg) parts = split_qualified_identifier(cleaned, quote_chars='"', allow_bracket_quotes=True) - return ".".join(quote_bracket_identifier(part) for part in parts) + return ".".join([quote_bracket_identifier(part) for part in parts]) def build_insert_statement(table: str, columns: "list[str]") -> str: """Build a pymssql-compatible INSERT statement for Sybase ASE.""" - column_clause = ", ".join(quote_bracket_identifier(column) for column in columns) - placeholders = ", ".join("%s" for _ in columns) + column_clause = ", ".join([quote_bracket_identifier(column) for column in columns]) + placeholders = ", ".join(["%s"] * len(columns)) if columns else "" return f"INSERT INTO {format_identifier(table)} ({column_clause}) VALUES ({placeholders})" @@ -155,7 +155,7 @@ def escape_literal_percent(sql: str, parameters: Any, validator: "ParameterValid if not parameters or "%" not in sql: return sql keep = {info.position for info in validator.extract_parameters(sql)} - return "".join("%%" if char == "%" and index not in keep else char for index, char in enumerate(sql)) + return "".join(["%%" if char == "%" and index not in keep else char for index, char in enumerate(sql)]) def normalize_execute_parameters(parameters: Any) -> Any: diff --git a/sqlspec/adapters/sybase/driver.py b/sqlspec/adapters/sybase/driver.py index 24bba7a27..9753f0444 100644 --- a/sqlspec/adapters/sybase/driver.py +++ b/sqlspec/adapters/sybase/driver.py @@ -1,6 +1,7 @@ """Sybase ASE database driver implementation over pymssql (FreeTDS TDS 5.0).""" import contextlib +import sys from collections.abc import Sized from typing import TYPE_CHECKING, Any, Literal, cast @@ -85,12 +86,15 @@ def start(self) -> None: try: cursor = cursor_manager.__enter__() handler = self._driver.handle_database_exceptions() - with handler: - normalized_params = normalize_execute_parameters(self._parameters) - sql = escape_literal_percent( - self._sql, normalized_params, self._driver.statement_config.parameter_validator - ) - cursor.execute(sql, normalized_params) + normalized_params = normalize_execute_parameters(self._parameters) + sql = escape_literal_percent( + self._sql, normalized_params, self._driver.statement_config.parameter_validator + ) + if sys.exc_info()[1] is None: + with handler: + cursor.execute(sql, normalized_params) + else: + self._driver._run_with_exception_handler(handler, cursor.execute, sql, normalized_params) self._driver._check_pending_exception(handler) except BaseException: with contextlib.suppress(Exception): @@ -105,8 +109,11 @@ def fetch_chunk(self) -> "list[dict[str, Any]]": cursor = cursor_manager.cursor handler = self._driver.handle_database_exceptions() rows: Any = [] - with handler: - rows = cursor.fetchmany(self._chunk_size) + if sys.exc_info()[1] is None: + with handler: + rows = cursor.fetchmany(self._chunk_size) + else: + rows = self._driver._run_with_exception_handler(handler, cursor.fetchmany, self._chunk_size) self._driver._check_pending_exception(handler) if not rows: return [] @@ -209,9 +216,10 @@ def begin(self) -> None: cursor.execute("BEGIN TRANSACTION") self._explicit_transaction = explicit self._transaction_active = True - except SybaseError as exc: + except Exception as exc: + mapped = create_mapped_exception(exc, logger=logger) msg = f"Failed to begin Sybase ASE transaction: {exc}" - raise SQLSpecError(msg) from exc + raise type(mapped)(msg) from exc def commit(self) -> None: try: @@ -222,9 +230,10 @@ def commit(self) -> None: self.connection.commit() self._explicit_transaction = False self._transaction_active = False - except SybaseError as exc: + except Exception as exc: + mapped = create_mapped_exception(exc, logger=logger) msg = f"Failed to commit Sybase ASE transaction: {exc}" - raise SQLSpecError(msg) from exc + raise type(mapped)(msg) from exc def rollback(self) -> None: try: @@ -233,9 +242,10 @@ def rollback(self) -> None: cursor.execute("IF @@TRANCOUNT > 0 ROLLBACK TRANSACTION") else: self.connection.rollback() - except SybaseError as exc: + except Exception as exc: + mapped = create_mapped_exception(exc, logger=logger) msg = f"Failed to rollback Sybase ASE transaction: {exc}" - raise SQLSpecError(msg) from exc + raise type(mapped)(msg) from exc finally: self._explicit_transaction = False self._transaction_active = False @@ -266,9 +276,18 @@ def has_schema(self, schema: str) -> bool: """Return whether the specified user/schema exists in dbo.sysusers.""" result = False exc_handler = self.handle_database_exceptions() - with exc_handler, self.with_cursor(self.connection) as cursor: - cursor.execute("SELECT 1 FROM dbo.sysusers WHERE name = %s", (schema,)) - result = cursor.fetchone() is not None + with self.with_cursor(self.connection) as cursor: + if sys.exc_info()[1] is None: + with exc_handler: + cursor.execute("SELECT 1 FROM dbo.sysusers WHERE name = %s", (schema,)) + result = cursor.fetchone() is not None + else: + def _fetch() -> None: + nonlocal result + cursor.execute("SELECT 1 FROM dbo.sysusers WHERE name = %s", (schema,)) + result = cursor.fetchone() is not None + + self._run_with_exception_handler(exc_handler, _fetch) self._check_pending_exception(exc_handler) return result @@ -350,8 +369,12 @@ def load_from_arrow( if overwrite: statement = f"DELETE FROM {format_identifier(table)}" exc_handler = self.handle_database_exceptions() - with exc_handler, self.with_cursor(self.connection) as cursor: - cursor.execute(statement) + with self.with_cursor(self.connection) as cursor: + if sys.exc_info()[1] is None: + with exc_handler: + cursor.execute(statement) + else: + self._run_with_exception_handler(exc_handler, cursor.execute, statement) self._check_pending_exception(exc_handler) columns, records = self._arrow_table_to_rows(arrow_table) @@ -367,8 +390,14 @@ def load_from_arrow( ) normalized_records = normalize_execute_many_parameters(prepared_records) exc_handler = self.handle_database_exceptions() - with exc_handler, self.with_cursor(self.connection) as cursor: - cursor.executemany(insert_sql, cast("Sequence[QueryParams]", normalized_records)) + with self.with_cursor(self.connection) as cursor: + if sys.exc_info()[1] is None: + with exc_handler: + cursor.executemany(insert_sql, cast("Sequence[QueryParams]", normalized_records)) + else: + self._run_with_exception_handler( + exc_handler, cursor.executemany, insert_sql, cast("Sequence[QueryParams]", normalized_records) + ) self._check_pending_exception(exc_handler) telemetry_payload = self._ingest_telemetry(arrow_table) diff --git a/sqlspec/adapters/sybase/litestar/store.py b/sqlspec/adapters/sybase/litestar/store.py index e81a4fe0d..158b7560b 100644 --- a/sqlspec/adapters/sybase/litestar/store.py +++ b/sqlspec/adapters/sybase/litestar/store.py @@ -4,7 +4,13 @@ from typing import TYPE_CHECKING, Any from sqlspec.adapters.sybase._typing import SybaseCursor -from sqlspec.adapters.sybase.core import escape_sql_literal, quote_bracket_identifier, split_schema_and_table +from sqlspec.adapters.sybase.core import ( + create_mapped_exception, + escape_sql_literal, + quote_bracket_identifier, + resolve_rowcount, + split_schema_and_table, +) from sqlspec.extensions.litestar.store import BaseSQLSpecStore from sqlspec.utils.sync_tools import async_ @@ -101,30 +107,33 @@ def _get(self, key: str, renew_for: "int | timedelta | None" = None) -> "bytes | WHERE session_id = %s AND (expires_at IS NULL OR expires_at > GETUTCDATE()) """ - with self._config.provide_connection() as conn: - with SybaseCursor(conn) as cursor: - cursor.execute(sql, (key,)) - row = cursor.fetchone() - - if row is None: - return None - - expires_at = _normalize_utc(_row_value(row, "expires_at", 1)) - if renew_for is not None and expires_at is not None: - new_expires_at = self._calculate_expires_at(renew_for) - if new_expires_at is not None: - with SybaseCursor(conn) as update_cursor: - update_cursor.execute( - f""" - UPDATE {table_ref} - SET expires_at = %s, updated_at = GETUTCDATE() - WHERE session_id = %s - """, - (new_expires_at, key), - ) - conn.commit() - - return _coerce_bytes(_row_value(row, "data", 0)) + try: + with self._config.provide_connection() as conn: + with SybaseCursor(conn) as cursor: + cursor.execute(sql, (key,)) + row = cursor.fetchone() + + if row is None: + return None + + expires_at = _normalize_utc(_row_value(row, "expires_at", 1)) + if renew_for is not None and expires_at is not None: + new_expires_at = self._calculate_expires_at(renew_for) + if new_expires_at is not None: + with SybaseCursor(conn) as update_cursor: + update_cursor.execute( + f""" + UPDATE {table_ref} + SET expires_at = %s, updated_at = GETUTCDATE() + WHERE session_id = %s + """, + (new_expires_at, key), + ) + conn.commit() + + return _coerce_bytes(_row_value(row, "data", 0)) + except Exception as exc: + raise create_mapped_exception(exc) from exc def _set(self, key: str, value: "str | bytes", expires_in: "int | timedelta | None" = None) -> None: data = self._value_to_bytes(value) @@ -139,22 +148,31 @@ def _set(self, key: str, value: "str | bytes", expires_in: "int | timedelta | No INSERT INTO {table_ref} (session_id, data, expires_at) VALUES (%s, %s, %s) """ - with self._config.provide_connection() as conn, SybaseCursor(conn) as cursor: - cursor.execute(sql, (key, data, expires_at, key, key, data, expires_at)) - conn.commit() + try: + with self._config.provide_connection() as conn, SybaseCursor(conn) as cursor: + cursor.execute(sql, (key, data, expires_at, key, key, data, expires_at)) + conn.commit() + except Exception as exc: + raise create_mapped_exception(exc) from exc def _delete(self, key: str) -> None: table_ref = _table_ref(self._table_name) - with self._config.provide_connection() as conn, SybaseCursor(conn) as cursor: - cursor.execute(f"DELETE FROM {table_ref} WHERE session_id = %s", (key,)) - conn.commit() + try: + with self._config.provide_connection() as conn, SybaseCursor(conn) as cursor: + cursor.execute(f"DELETE FROM {table_ref} WHERE session_id = %s", (key,)) + conn.commit() + except Exception as exc: + raise create_mapped_exception(exc) from exc def _delete_all(self) -> None: table_ref = _table_ref(self._table_name) - with self._config.provide_connection() as conn, SybaseCursor(conn) as cursor: - cursor.execute(f"TRUNCATE TABLE {table_ref}") - conn.commit() - self._log_delete_all() + try: + with self._config.provide_connection() as conn, SybaseCursor(conn) as cursor: + cursor.execute(f"TRUNCATE TABLE {table_ref}") + conn.commit() + self._log_delete_all() + except Exception as exc: + raise create_mapped_exception(exc) from exc def _exists(self, key: str) -> bool: table_ref = _table_ref(self._table_name) @@ -164,23 +182,29 @@ def _exists(self, key: str) -> bool: WHERE session_id = %s AND (expires_at IS NULL OR expires_at > GETUTCDATE()) """ - with self._config.provide_connection() as conn, SybaseCursor(conn) as cursor: - cursor.execute(sql, (key,)) - return cursor.fetchone() is not None + try: + with self._config.provide_connection() as conn, SybaseCursor(conn) as cursor: + cursor.execute(sql, (key,)) + return cursor.fetchone() is not None + except Exception as exc: + raise create_mapped_exception(exc) from exc def _expires_in(self, key: str) -> "int | None": table_ref = _table_ref(self._table_name) - with self._config.provide_connection() as conn, SybaseCursor(conn) as cursor: - cursor.execute(f"SELECT expires_at FROM {table_ref} WHERE session_id = %s", (key,)) - row = cursor.fetchone() - - if row is None: - return None - expires_at = _normalize_utc(_row_value(row, "expires_at", 0)) - if expires_at is None: - return None - remaining = expires_at - datetime.now(UTC) - return max(0, int(remaining.total_seconds())) + try: + with self._config.provide_connection() as conn, SybaseCursor(conn) as cursor: + cursor.execute(f"SELECT expires_at FROM {table_ref} WHERE session_id = %s", (key,)) + row = cursor.fetchone() + + if row is None: + return None + expires_at = _normalize_utc(_row_value(row, "expires_at", 0)) + if expires_at is None: + return None + remaining = expires_at - datetime.now(UTC) + return max(0, int(remaining.total_seconds())) + except Exception as exc: + raise create_mapped_exception(exc) from exc def _delete_expired(self) -> int: table_ref = _table_ref(self._table_name) @@ -189,13 +213,16 @@ def _delete_expired(self) -> int: WHERE expires_at IS NOT NULL AND expires_at < GETUTCDATE() """ - with self._config.provide_connection() as conn, SybaseCursor(conn) as cursor: - cursor.execute(sql) - count = int(getattr(cursor, "rowcount", 0) or 0) - conn.commit() - if count > 0: - self._log_delete_expired(count) - return count + try: + with self._config.provide_connection() as conn, SybaseCursor(conn) as cursor: + cursor.execute(sql) + count = resolve_rowcount(cursor) + conn.commit() + if count > 0: + self._log_delete_expired(count) + return count + except Exception as exc: + raise create_mapped_exception(exc) from exc def _table_ref(table_name: str) -> str: diff --git a/sqlspec/adapters/sybase/migrations.py b/sqlspec/adapters/sybase/migrations.py index 1e0a4e613..b7757e427 100644 --- a/sqlspec/adapters/sybase/migrations.py +++ b/sqlspec/adapters/sybase/migrations.py @@ -71,14 +71,20 @@ def _tracking_table_ddl_text(self) -> str: def _existing_columns_query(self) -> str: """Return Sybase ASE catalog query text for migration tracking table columns.""" + has_schema = "." in self.version_table schema_name, table_name = split_schema_and_table(self.version_table) + schema_condition = ( + f"u.name = '{escape_sql_literal(schema_name)}'" + if has_schema + else "u.name = USER_NAME()" + ) return f""" SELECT c.name AS column_name FROM dbo.syscolumns c INNER JOIN dbo.sysobjects o ON o.id = c.id AND o.type = 'U' INNER JOIN dbo.sysusers u ON u.uid = o.uid WHERE o.name = '{escape_sql_literal(table_name)}' - AND u.name = '{escape_sql_literal(schema_name)}' + AND {schema_condition} """ def _add_column_statement_text(self, column_name: str) -> str | None: diff --git a/sqlspec/adapters/sybase/pool.py b/sqlspec/adapters/sybase/pool.py index 7fe52604a..48eda0125 100644 --- a/sqlspec/adapters/sybase/pool.py +++ b/sqlspec/adapters/sybase/pool.py @@ -98,10 +98,10 @@ def new_connection(self) -> SybaseConnection: Returns: SybaseConnection: A newly opened, fully configured connection. """ - module = pymssql - if module is None: - module = _load_pymssql_module() - connection = module.connect(**self._connection_parameters) + global pymssql + if pymssql is None: + pymssql = _load_pymssql_module() + connection = pymssql.connect(**self._connection_parameters) if self._on_connection_create is not None: self._on_connection_create(connection) diff --git a/sqlspec/data_dictionary/dialects/sqlanywhere/config.py b/sqlspec/data_dictionary/dialects/sqlanywhere/config.py index 51246d50a..4652ddc32 100644 --- a/sqlspec/data_dictionary/dialects/sqlanywhere/config.py +++ b/sqlspec/data_dictionary/dialects/sqlanywhere/config.py @@ -57,7 +57,7 @@ "supports_merge": True, "supports_for_update": True, "supports_skip_locked": True, - "supports_transactional_ddl": True, + "supports_transactional_ddl": False, "supports_transactions": True, "supports_prepared_statements": True, "supports_in_memory": False, diff --git a/tests/unit/adapters/test_sqlanywhere.py b/tests/unit/adapters/test_sqlanywhere.py index fdfb1a3bc..cf1e7770c 100644 --- a/tests/unit/adapters/test_sqlanywhere.py +++ b/tests/unit/adapters/test_sqlanywhere.py @@ -1,5 +1,6 @@ """Unit tests for the SAP SQL Anywhere (sqlanydb) adapter.""" +import threading from typing import Any, cast from unittest.mock import MagicMock @@ -15,7 +16,9 @@ SqlanywhereDriver, SqlanywhereExceptionHandler, SqlanywhereSyncDataDictionary, + build_column_metadata, build_connection_config, + build_foreign_key_metadata, create_mapped_exception, default_statement_config, driver_profile, @@ -23,10 +26,11 @@ from sqlspec.adapters.sqlanywhere._typing import SqlanywhereError from sqlspec.adapters.sqlanywhere.adk import SqlanywhereADKConfig, SqlanywhereADKMemoryStore, SqlanywhereADKStore from sqlspec.adapters.sqlanywhere.events import SqlanywhereEventQueueStore -from sqlspec.adapters.sqlanywhere.litestar import SqlanywhereStore +from sqlspec.adapters.sqlanywhere.litestar import SqlanywhereLitestarConfig, SqlanywhereStore from sqlspec.adapters.sqlanywhere.migrations import SqlanywhereSyncMigrationTracker from sqlspec.core import SQL, ParameterStyle from sqlspec.data_dictionary.dialects.sqlanywhere import SqlanywhereVersionInfo +from sqlspec.data_dictionary.dialects.sqlanywhere.config import SQLANYWHERE_FEATURE_FLAGS from sqlspec.exceptions import ( CheckViolationError, DatabaseConnectionError, @@ -40,6 +44,7 @@ PermissionDeniedError, SerializationConflictError, SQLParsingError, + TransactionError, UniqueViolationError, ) @@ -506,3 +511,242 @@ def _table_missing(_self: Any, *_args: Any, **_kwargs: Any) -> None: monkeypatch.setattr(SqlanywhereADKStore, "_execute_fetchone", _table_missing) assert store.get_session("app", "user", "sess") is None + + +def test_sqlanywhere_adk_memory_store_insert_sql_contains_from_sys_dummy(monkeypatch: pytest.MonkeyPatch) -> None: + """SqlanywhereADKMemoryStore.insert_memory_entries emits SELECT ... FROM SYS.DUMMY WHERE NOT EXISTS.""" + config = SqlanywhereConfig(connection_config={"user": "dba", "password": "sql", "database": "demo"}) + store = SqlanywhereADKMemoryStore(config) + executed_statements: list[str] = [] + + class DummyCursor: + rowcount = 1 + + def execute(self, sql: str, _params: Any = ()) -> None: + executed_statements.append(sql) + + def close(self) -> None: + pass + + class DummyConnection: + def cursor(self) -> DummyCursor: + return DummyCursor() + + def commit(self) -> None: + pass + + def rollback(self) -> None: + pass + + monkeypatch.setattr(config, "provide_connection", lambda: DummyConnection()) + + entry = { + "id": "mem-1", + "session_id": "sess-1", + "app_name": "app", + "user_id": "u1", + "scope": "user", + "event_id": "evt-1", + "author": "tester", + "timestamp": datetime.now(UTC), + "content_json": {"msg": "hi"}, + "content_text": "hi", + "metadata_json": None, + } + inserted = store.insert_memory_entries([cast(Any, entry)]) + assert inserted == 1 + assert len(executed_statements) == 1 + assert "FROM SYS.DUMMY" in executed_statements[0] + assert "WHERE NOT EXISTS" in executed_statements[0] + + +def test_sqlanywhere_adk_memory_store_handles_native_missing_table(monkeypatch: pytest.MonkeyPatch) -> None: + """SqlanywhereADKMemoryStore.search_entries returns empty list when table is missing.""" + config = SqlanywhereConfig(connection_config={"user": "dba", "password": "sql", "database": "demo"}) + store = SqlanywhereADKMemoryStore(config) + + class FakeSqlanydbError(SqlanywhereError): + def __init__(self, sqlstate: str, sqlcode: int, message: str) -> None: + super().__init__(message) + self.sqlstate = sqlstate + self.sqlcode = sqlcode + + def _table_missing(_self: Any, *_args: Any, **_kwargs: Any) -> None: + raise FakeSqlanydbError("42S02", -141, "Table 'adk_memory' not found") + + monkeypatch.setattr(SqlanywhereADKMemoryStore, "_execute_fetchall", _table_missing) + assert store.search_entries("query", "app", "user") == [] + + +def test_sqlanywhere_litestar_config_and_delete_expired(monkeypatch: pytest.MonkeyPatch) -> None: + """SqlanywhereLitestarConfig is usable and _delete_expired uses resolve_rowcount.""" + config = SqlanywhereLitestarConfig() + assert isinstance(config, dict) + + adapter_config = SqlanywhereConfig(connection_config={"user": "dba", "password": "sql", "database": "demo"}) + litestar_store = SqlanywhereStore(adapter_config) + + class DummyCursor: + rowcount = 4 + + def execute(self, _sql: str) -> None: + pass + + class DummyConnection: + def cursor(self) -> DummyCursor: + return DummyCursor() + + def commit(self) -> None: + pass + + monkeypatch.setattr(adapter_config, "provide_connection", lambda: DummyConnection()) + assert litestar_store._delete_expired() == 4 + + +def test_sqlanywhere_feature_flags_supports_transactional_ddl() -> None: + """SQLANYWHERE_FEATURE_FLAGS supports_transactional_ddl is False.""" + assert SQLANYWHERE_FEATURE_FLAGS["supports_transactional_ddl"] is False + + +def test_sqlanywhere_build_column_metadata_is_nullable() -> None: + """build_column_metadata parses NO and YES correctly.""" + cols = build_column_metadata([ + {"column_name": "id", "is_nullable": "NO"}, + {"column_name": "val", "is_nullable": "YES"}, + ]) + assert cols[0].is_nullable is False + assert cols[1].is_nullable is True + + +def test_sqlanywhere_build_foreign_key_metadata_names() -> None: + """build_foreign_key_metadata supports referenced_table and referenced_table_name keys.""" + fks = build_foreign_key_metadata([ + {"table_name": "t1", "column_name": "c1", "referenced_table": "t2", "referenced_column": "c2"}, + {"table_name": "t1", "column_name": "c3", "referenced_table_name": "t3", "referenced_column_name": "c4"}, + ]) + assert fks[0].referenced_table == "t2" + assert fks[0].referenced_column == "c2" + assert fks[1].referenced_table == "t3" + assert fks[1].referenced_column == "c4" + + +def test_sqlanywhere_sync_data_dictionary_instantiation() -> None: + """SqlanywhereSyncDataDictionary instantiates cleanly without redundant __init__.""" + dd = SqlanywhereSyncDataDictionary() + assert dd.dialect == "sqlanywhere" + + +def test_sqlanywhere_has_schema_under_active_exception(monkeypatch: pytest.MonkeyPatch) -> None: + """has_schema runs cleanly inside an active exception block without isolation leak.""" + config = SqlanywhereConfig(connection_config={"user": "dba", "password": "sql", "database": "demo"}) + driver = SqlanywhereDriver(config) + + class DummyCursor: + def execute(self, _sql: str, _params: Any = ()) -> None: + pass + + def fetchone(self) -> Any: + return (1,) + + def close(self) -> None: + pass + + class DummyConnection: + def cursor(self) -> DummyCursor: + return DummyCursor() + + monkeypatch.setattr(driver, "connection", DummyConnection()) + + try: + raise RuntimeError("active error") + except RuntimeError: + has_user = driver.has_schema("DBA") + assert has_user is True + + +def test_sqlanywhere_format_identifier_and_insert_statement() -> None: + """format_identifier and build_insert_statement format identifiers correctly.""" + assert format_identifier("schema.my_table") == '"schema"."my_table"' + sql = build_insert_statement("my_table", ["col1", "col2"]) + assert sql == 'INSERT INTO "my_table" ("col1", "col2") VALUES (?, ?)' + + +def test_sqlanywhere_commit_and_rollback_raise_transaction_error() -> None: + """commit() and rollback() wrap caught exceptions in TransactionError.""" + config = SqlanywhereConfig(connection_config={"user": "dba", "password": "sql", "database": "demo"}) + driver = SqlanywhereDriver(config) + + class FailingConnection: + def commit(self) -> None: + raise RuntimeError("commit failed") + + def rollback(self) -> None: + raise RuntimeError("rollback failed") + + driver.connection = cast("SqlanywhereConnection", FailingConnection()) + with pytest.raises(TransactionError, match="Failed to commit SQL Anywhere transaction"): + driver.commit() + with pytest.raises(TransactionError, match="Failed to rollback SQL Anywhere transaction"): + driver.rollback() + + +def test_sqlanywhere_session_context_committed_does_not_rollback() -> None: + """An explicitly committed session context does not issue rollback() on exit.""" + rollback_called = False + + class TrackingConnection: + def commit(self) -> None: + pass + + def rollback(self) -> None: + nonlocal rollback_called + rollback_called = True + + def close(self) -> None: + pass + + conn = TrackingConnection() + ctx = SqlanywhereSessionContext( + acquire_connection=lambda: conn, + release_connection=lambda _c, **_kw: None, + statement_config=default_statement_config, + driver_features={}, + prepare_driver=lambda d: d, + ) + with ctx as driver: + driver.begin() + driver.commit() + + assert rollback_called is False + + +def test_sqlanywhere_connection_pool_bounds_max_size(monkeypatch: pytest.MonkeyPatch) -> None: + """SqlanywhereConnectionPool enforces max_size in multi-threaded acquisition.""" + pool = SqlanywhereConnectionPool({"user": "dba", "password": "sql", "database": "demo"}, min_size=1, max_size=2) + + class DummyConn: + def close(self) -> None: + pass + + monkeypatch.setattr(pool, "new_connection", lambda: cast("SqlanywhereConnection", DummyConn())) + + errors: list[Exception] = [] + conns: list[Any] = [] + barrier = threading.Barrier(3) + + def worker() -> None: + try: + conn = pool._get_thread_connection() + conns.append(conn) + barrier.wait(timeout=2.0) + except Exception as exc: + errors.append(exc) + barrier.wait(timeout=2.0) + + threads = [threading.Thread(target=worker) for _ in range(3)] + for t in threads: + t.start() + for t in threads: + t.join(timeout=3.0) + + assert any(isinstance(err, OperationalError) and "pool exhausted" in str(err) for err in errors) diff --git a/tests/unit/adapters/test_sybase.py b/tests/unit/adapters/test_sybase.py index c2c55fe07..2eacc5104 100644 --- a/tests/unit/adapters/test_sybase.py +++ b/tests/unit/adapters/test_sybase.py @@ -21,10 +21,14 @@ SybaseDriver, SybaseExceptionHandler, SybaseSyncDataDictionary, + _ERROR_CODE_MAPPING, build_connection_config, + build_insert_statement, create_mapped_exception, default_statement_config, driver_profile, + escape_literal_percent, + format_identifier, ) from sqlspec.adapters.sybase._typing import SybaseRawCursor from sqlspec.adapters.sybase.adk import SybaseADKConfig, SybaseADKMemoryStore, SybaseADKStore @@ -49,6 +53,7 @@ SQLSpecError, TransactionError, UniqueViolationError, + UniqueViolationError, ) @@ -545,3 +550,144 @@ def _table_missing(_self: Any, *_args: Any, **_kwargs: Any) -> None: monkeypatch.setattr(SybaseADKStore, "_execute_fetchone", _table_missing) assert store.get_session("app", "user", "sess") is None + + +def test_sybase_commit_maps_deadlock_error_1205() -> None: + """Simulate error 1205 during driver.commit(); assert DeadlockError is raised.""" + cursor = FakeCursor() + conn = FakeConnection(cursor) + driver = SybaseDriver(cast("SybaseConnection", conn)) + + def _fail_commit() -> None: + raise pymssql.OperationalError(1205, "Your transaction was deadlock victim (1205)") + + conn.commit = _fail_commit + with pytest.raises(DeadlockError, match="Failed to commit Sybase ASE transaction"): + driver.commit() + + +def test_sybase_adk_store_maps_unique_violation(monkeypatch: pytest.MonkeyPatch) -> None: + """SybaseADKStore raises UniqueViolationError on duplicate key violation.""" + config = MagicMock() + config.extension_config = {"adk": {}} + store = SybaseADKStore(cast("Any", config)) + + class DummyCursor: + def execute(self, _sql: str, _params: Any = ()) -> None: + raise pymssql.IntegrityError(2601, "Attempt to insert duplicate key row with unique index") + + def close(self) -> None: + pass + + class DummyConnection: + def cursor(self, as_dict: bool = False) -> DummyCursor: + return DummyCursor() + + def rollback(self) -> None: + pass + + monkeypatch.setattr(config, "provide_connection", lambda: DummyConnection()) + with pytest.raises(UniqueViolationError): + store._execute("INSERT INTO dummy VALUES (1)") + + +def test_sybase_litestar_store_get_raises_sqlspec_error(monkeypatch: pytest.MonkeyPatch) -> None: + """SybaseStore._get() maps raw database errors to SQLSpecError.""" + config = MagicMock() + store = SybaseStore(cast("Any", config)) + + class FailingCursor: + def execute(self, _sql: str, _params: Any = ()) -> None: + raise pymssql.OperationalError("Connection lost") + + def close(self) -> None: + pass + + class FailingConnection: + def cursor(self) -> FailingCursor: + return FailingCursor() + + monkeypatch.setattr(config, "provide_connection", lambda: FailingConnection()) + with pytest.raises(SQLSpecError): + store._get("session-1") + + +def test_sybase_has_schema_under_active_exception(monkeypatch: pytest.MonkeyPatch) -> None: + """has_schema runs cleanly inside an active exception block without isolation leak.""" + config = MagicMock() + driver = SybaseDriver(config) + + class DummyCursor: + def execute(self, _sql: str, _params: Any = ()) -> None: + pass + + def fetchone(self) -> Any: + return (1,) + + def close(self) -> None: + pass + + class DummyConnection: + def cursor(self) -> DummyCursor: + return DummyCursor() + + monkeypatch.setattr(driver, "connection", DummyConnection()) + + try: + raise RuntimeError("active error") + except RuntimeError: + has_user = driver.has_schema("dbo") + assert has_user is True + + +def test_sybase_pool_new_connection_caches_pymssql_module(monkeypatch: pytest.MonkeyPatch) -> None: + """new_connection caches pymssql module reference at module level.""" + from sqlspec.adapters.sybase import pool as sybase_pool + + import_count = 0 + real_import = sybase_pool._load_pymssql_module + + def mock_load() -> Any: + nonlocal import_count + import_count += 1 + return real_import() + + monkeypatch.setattr(sybase_pool, "_load_pymssql_module", mock_load) + monkeypatch.setattr(sybase_pool, "pymssql", None) + + pool = sybase_pool.SybaseConnectionPool({"server": "localhost", "database": "demo"}) + monkeypatch.setattr("pymssql.connect", lambda **_kw: FakeConnection(FakeCursor())) + + _ = pool.new_connection() + assert sybase_pool.pymssql is not None + assert import_count == 1 + + _ = pool.new_connection() + assert import_count == 1 + + +def test_sybase_format_identifier_and_insert_statements() -> None: + """format_identifier, build_insert_statement, and escape_literal_percent format queries correctly.""" + assert format_identifier("dbo.my_table") == "[dbo].[my_table]" + sql = build_insert_statement("my_table", ["col1", "col2"]) + assert sql == "INSERT INTO [my_table] ([col1], [col2]) VALUES (%s, %s)" + escaped = escape_literal_percent("SELECT * FROM t WHERE val LIKE '100%' AND id = %s", ("1",), default_statement_config.parameter_validator) + assert "100%%" in escaped + + +def test_sybase_migration_tracker_existing_columns_query() -> None: + """_existing_columns_query resolves USER_NAME() when schema is unqualified.""" + tracker_unqualified = SybaseSyncMigrationTracker(version_table="my_versions") + query_unqualified = tracker_unqualified._existing_columns_query() + assert "u.name = USER_NAME()" in query_unqualified + assert "o.name = 'my_versions'" in query_unqualified + + tracker_qualified = SybaseSyncMigrationTracker(version_table="custom_schema.my_versions") + query_qualified = tracker_qualified._existing_columns_query() + assert "u.name = 'custom_schema'" in query_qualified + assert "o.name = 'my_versions'" in query_qualified + + +def test_sybase_error_code_mapping_count() -> None: + """_ERROR_CODE_MAPPING contains 38 error codes.""" + assert len(_ERROR_CODE_MAPPING) == 38 From 1dced2ab6881f9119d3fc918806c0d6e943fd9f7 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 00:49:05 +0000 Subject: [PATCH 055/174] fix(sqlanywhere): reconcile dialect feature flags and data dictionary metadata parsing Correct SQL Anywhere dialect feature flags to indicate transactional DDL is not supported. Update `build_column_metadata` to parse `is_nullable` from `YES`/`NO` strings and `build_foreign_key_metadata` to accept both `referenced_table` and `referenced_table_name` column keys. Remove redundant `__init__` override in `SqlanywhereSyncDataDictionary`. From 42a9caddb6021b94731d7f6e4058a80a1702279b Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 00:49:05 +0000 Subject: [PATCH 056/174] perf(sqlanywhere): isolate exception unwinding and optimize identifier formatting Check `sys.exc_info()[1]` across `has_schema`, `load_from_arrow`, and stream execution to prevent exception isolation leaks during unwinding. Replace generator expressions in `format_identifier` and `build_insert_statement` with list comprehensions to eliminate generator frame allocation overhead under MyPyC. From d5f678e1968e7f6bd9c8300c484f7702e9dd3b49 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 00:49:05 +0000 Subject: [PATCH 057/174] fix(sqlanywhere): raise TransactionError on transaction failures and bound connection pool Raise `TransactionError` in `commit()` and `rollback()` when transaction operations fail. Guard session exit rollback to ensure committed transactions do not roll back redundantly and pass exit exception details to `_release_connection`. Enforce `max_size` in `SqlanywhereConnectionPool._create_connection` to raise `OperationalError` when pool capacity is exhausted. From 21be049932b68f93a140e47dec1075615fb8fc76 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 00:49:05 +0000 Subject: [PATCH 058/174] fix(sybase): map transaction exceptions through create_mapped_exception Wrap exceptions caught in `begin()`, `commit()`, and `rollback()` through `create_mapped_exception(exc, logger=logger)` so mapped exception types such as `DeadlockError` are preserved rather than masked by generic `SQLSpecError`. From f1dc213b0f0c4ae8ae26e1a7bf389a35daf1a71f Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 00:49:05 +0000 Subject: [PATCH 059/174] fix(sybase): map database exceptions in ADK and Litestar stores Wrap database operations in `SybaseADKStore` and `SybaseStore` with `create_mapped_exception(exc)` so database errors like `UniqueViolationError` and `OperationalError` are propagated cleanly instead of leaking raw `pymssql` errors. Use `resolve_rowcount` during expired session deletion and support `SQLSpecError` in table missing checks. From 193b60b585493802932ab2caf18e0c295e6f12a9 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 00:49:05 +0000 Subject: [PATCH 060/174] perf(sybase): isolate exception unwinding and cache pymssql module reference Guard execution in `has_schema`, `load_from_arrow`, and `SybaseStreamSource` with ambient `sys.exc_info()[1]` checks to avoid leaking exception state during C-extension unwinding. Cache the resolved `pymssql` module reference at module level in `SybaseConnectionPool.new_connection` to avoid repeated import lookups, and convert generator expressions in `core.py` to list comprehensions for MyPyC. From 6708f86006a170f60caa75cdbf8a9b7de2976468 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 00:49:05 +0000 Subject: [PATCH 061/174] fix(sybase): resolve migration tracking table columns using USER_NAME for unqualified tables Resolve migration tracking table columns against `USER_NAME()` when the table is unqualified instead of defaulting to `dbo`, matching Sybase ASE's native object resolution rules. Expand unit test suites to verify migration tracking query generation and error code mapping count. From 5eba7527e5d8323d14b055817519ecb1e74e26ee Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 01:01:40 +0000 Subject: [PATCH 062/174] fix(mssql,sybase,sqlanywhere): repair T-SQL family integration - Keep mssql_python and pymssql create_mapped_exception and _ERROR_CODE_MAPPING identical (ODBC SQLSTATE mapping, FreeTDS codes, InterfaceError branch); the store parity test requires both copies to stay in lockstep. - Import escape_sql_literal from mssql_python.core in the ADK memory migration after the private alias was removed. - SQL Anywhere auto-commits around DDL, so the dialect test now expects supports_transactional_ddl to be False, matching the feature flags. - Sybase stores return after the try block instead of inside it. - Repair new Sybase and SQL Anywhere unit tests: import helpers from their defining modules, model provide_connection as a context manager, patch slotted classes at class level, read ColumnMetadata as a TypedDict, use the version_table_name tracker argument, and set TDSVER for the pool test. --- sqlspec/adapters/mssql_python/core.py | 6 +++ sqlspec/adapters/pymssql/core.py | 43 +++++++++++++++++-- sqlspec/adapters/sqlanywhere/driver.py | 1 + sqlspec/adapters/sybase/adk/store.py | 6 +-- sqlspec/adapters/sybase/driver.py | 9 +--- sqlspec/adapters/sybase/litestar/store.py | 6 +-- sqlspec/adapters/sybase/migrations.py | 6 +-- .../0002_create_mssql_python_memory.py | 12 ++---- tests/unit/adapters/test_sqlanywhere.py | 20 +++++---- tests/unit/adapters/test_sybase.py | 26 ++++++----- .../test_sybase_sqlanywhere_dialects.py | 2 +- .../extensions/test_schema_reconciliation.py | 6 +-- 12 files changed, 89 insertions(+), 54 deletions(-) diff --git a/sqlspec/adapters/mssql_python/core.py b/sqlspec/adapters/mssql_python/core.py index f7403c194..7ca3c159d 100644 --- a/sqlspec/adapters/mssql_python/core.py +++ b/sqlspec/adapters/mssql_python/core.py @@ -149,6 +149,10 @@ 8115: (DataError, "arithmetic overflow error"), 8152: (DataError, "string or binary data would be truncated"), 1105: (OperationalError, "operational error"), + 20003: (QueryTimeoutError, "query timeout expired (FreeTDS)"), + 20002: (DatabaseConnectionError, "adaptive Server connection failed (FreeTDS)"), + 20009: (DatabaseConnectionError, "unable to connect: host unreachable or server down (FreeTDS)"), + 20047: (DatabaseConnectionError, "connection closed or broken (FreeTDS)"), } @@ -355,6 +359,8 @@ def create_mapped_exception(error: BaseException, *, logger: "Logger | None" = N return constraint_exception exc_name = type(error).__name__ + if exc_name == "InterfaceError": + return DatabaseConnectionError(f"SQL Server interface error. Original error: {error}") if exc_name == "IntegrityError": return IntegrityError(f"SQL Server integrity error. Original error: {error}") if exc_name == "OperationalError": diff --git a/sqlspec/adapters/pymssql/core.py b/sqlspec/adapters/pymssql/core.py index 32f33eb67..51aea9fd0 100644 --- a/sqlspec/adapters/pymssql/core.py +++ b/sqlspec/adapters/pymssql/core.py @@ -19,6 +19,7 @@ SerializationConflictError, SQLParsingError, SQLSpecError, + TransactionError, UniqueViolationError, ) from sqlspec.utils.serializers import from_json, to_json @@ -58,9 +59,19 @@ _ERROR_NUMBER_PATTERN: Final[re.Pattern[str]] = re.compile(r"(?:\(([-]?\d+)(?:,|\))|\bMsg\s+([-]?\d+)\b)") _MSSQL_CONSTRAINT_547: Final[int] = 547 _COLUMN_CACHE_MAX_SIZE: Final[int] = 256 +_ODBC_SQLSTATE_MAP: Final[dict[str, tuple[type[SQLSpecError], str]]] = { + "08S01": (DatabaseConnectionError, "communication link failure"), + "08001": (DatabaseConnectionError, "unable to connect to data source"), + "HYT00": (QueryTimeoutError, "query timeout expired"), + "HY008": (QueryTimeoutError, "operation canceled"), +} _ERROR_CODE_MAPPING: Final[dict[int, tuple[type[SQLSpecError], str]]] = { + 220: (DataError, "arithmetic overflow error"), + 258: (QueryTimeoutError, "timeout waiting for latch"), 2601: (UniqueViolationError, "unique constraint violation"), 2627: (UniqueViolationError, "unique constraint violation"), + 3617: (QueryTimeoutError, "query canceled by user or attention event"), + 3930: (TransactionError, "current transaction cannot be committed"), 515: (NotNullViolationError, "not-null constraint violation"), 18456: (PermissionDeniedError, "permission denied"), 229: (PermissionDeniedError, "permission denied"), @@ -72,10 +83,6 @@ 3960: (SerializationConflictError, "snapshot isolation conflict"), -2: (QueryTimeoutError, "query timeout"), 1222: (QueryTimeoutError, "lock request time out period exceeded"), - 20003: (QueryTimeoutError, "query timeout expired (FreeTDS)"), - 20002: (DatabaseConnectionError, "adaptive Server connection failed (FreeTDS)"), - 20009: (DatabaseConnectionError, "unable to connect: host unreachable or server down (FreeTDS)"), - 20047: (DatabaseConnectionError, "connection closed or broken (FreeTDS)"), 102: (SQLParsingError, "SQL syntax error"), 105: (SQLParsingError, "unclosed quotation mark"), 156: (SQLParsingError, "incorrect syntax near keyword"), @@ -89,6 +96,10 @@ 8115: (DataError, "arithmetic overflow error"), 8152: (DataError, "string or binary data would be truncated"), 1105: (OperationalError, "operational error"), + 20003: (QueryTimeoutError, "query timeout expired (FreeTDS)"), + 20002: (DatabaseConnectionError, "adaptive Server connection failed (FreeTDS)"), + 20009: (DatabaseConnectionError, "unable to connect: host unreachable or server down (FreeTDS)"), + 20047: (DatabaseConnectionError, "connection closed or broken (FreeTDS)"), } @@ -221,8 +232,32 @@ def apply_driver_features( return statement_config, features +_SQLSTATE_LENGTH: Final[int] = 5 + + +def extract_sqlstate(error: BaseException) -> str | None: + """Extract an ODBC SQLSTATE code from a driver error.""" + val = getattr(error, "sqlstate", None) + if isinstance(val, str) and val: + return val.strip().upper() + if hasattr(error, "args") and len(error.args) > 1: + second = error.args[1] + if isinstance(second, str) and len(second.strip()) == _SQLSTATE_LENGTH: + return second.strip().upper() + message = str(error).upper() + for code in _ODBC_SQLSTATE_MAP: + if code in message: + return code + return None + + def create_mapped_exception(error: BaseException, *, logger: "Logger | None" = None) -> SQLSpecError: """Map a SQL Server exception to SQLSpec's exception hierarchy.""" + sqlstate = extract_sqlstate(error) + if sqlstate is not None and sqlstate in _ODBC_SQLSTATE_MAP: + error_class, description = _ODBC_SQLSTATE_MAP[sqlstate] + return error_class(f"SQL Server SQLSTATE {sqlstate}: {description}. Original error: {error}") + error_number = extract_error_number(error) if error_number == _MSSQL_CONSTRAINT_547: diff --git a/sqlspec/adapters/sqlanywhere/driver.py b/sqlspec/adapters/sqlanywhere/driver.py index b4a2d50f0..69fd3dec1 100644 --- a/sqlspec/adapters/sqlanywhere/driver.py +++ b/sqlspec/adapters/sqlanywhere/driver.py @@ -269,6 +269,7 @@ def has_schema(self, schema: str) -> bool: cursor.execute("SELECT 1 FROM SYS.SYSUSER WHERE user_name = ?", (schema,)) result = cursor.fetchone() is not None else: + def _fetch() -> None: nonlocal result cursor.execute("SELECT 1 FROM SYS.SYSUSER WHERE user_name = ?", (schema,)) diff --git a/sqlspec/adapters/sybase/adk/store.py b/sqlspec/adapters/sybase/adk/store.py index e8b3b4ebe..b799b9c10 100644 --- a/sqlspec/adapters/sybase/adk/store.py +++ b/sqlspec/adapters/sybase/adk/store.py @@ -417,12 +417,12 @@ def _execute_fetchone(self, sql: str, params: "tuple[Any, ...]" = (), *, commit: row = cursor.fetchone() if commit: conn.commit() - return row except Exception as exc: if commit: with contextlib.suppress(Exception): conn.rollback() raise create_mapped_exception(exc) from exc + return row def _execute_fetchall(self, sql: str, params: "tuple[Any, ...]" = ()) -> "list[Any]": with self._config.provide_connection() as conn, SybaseCursor(conn, as_dict=False) as cursor: @@ -439,12 +439,12 @@ def _execute(self, sql: str, params: "tuple[Any, ...]" = (), *, commit: bool = F rowcount = resolve_rowcount(cursor) if commit: conn.commit() - return rowcount except Exception as exc: if commit: with contextlib.suppress(Exception): conn.rollback() raise create_mapped_exception(exc) from exc + return rowcount class SybaseADKMemoryStore(BaseSyncADKMemoryStore["SybaseConfig"]): @@ -626,12 +626,12 @@ def _execute(self, sql: str, params: "tuple[Any, ...]" = (), *, commit: bool = F rowcount = resolve_rowcount(cursor) if commit: conn.commit() - return rowcount except Exception as exc: if commit: with contextlib.suppress(Exception): conn.rollback() raise create_mapped_exception(exc) from exc + return rowcount def _existing_index_names(driver: "SybaseDriver", index_specs: "list[tuple[str, str, str]]") -> "set[str]": diff --git a/sqlspec/adapters/sybase/driver.py b/sqlspec/adapters/sybase/driver.py index 9753f0444..68e0fc5ee 100644 --- a/sqlspec/adapters/sybase/driver.py +++ b/sqlspec/adapters/sybase/driver.py @@ -5,13 +5,7 @@ from collections.abc import Sized from typing import TYPE_CHECKING, Any, Literal, cast -from sqlspec.adapters.sybase._typing import ( - SybaseConnection, - SybaseCursor, - SybaseError, - SybaseRawCursor, - SybaseSessionContext, -) +from sqlspec.adapters.sybase._typing import SybaseConnection, SybaseCursor, SybaseRawCursor, SybaseSessionContext from sqlspec.adapters.sybase.core import ( build_insert_statement, collect_rows, @@ -282,6 +276,7 @@ def has_schema(self, schema: str) -> bool: cursor.execute("SELECT 1 FROM dbo.sysusers WHERE name = %s", (schema,)) result = cursor.fetchone() is not None else: + def _fetch() -> None: nonlocal result cursor.execute("SELECT 1 FROM dbo.sysusers WHERE name = %s", (schema,)) diff --git a/sqlspec/adapters/sybase/litestar/store.py b/sqlspec/adapters/sybase/litestar/store.py index 158b7560b..3a6e8ae15 100644 --- a/sqlspec/adapters/sybase/litestar/store.py +++ b/sqlspec/adapters/sybase/litestar/store.py @@ -218,11 +218,11 @@ def _delete_expired(self) -> int: cursor.execute(sql) count = resolve_rowcount(cursor) conn.commit() - if count > 0: - self._log_delete_expired(count) - return count except Exception as exc: raise create_mapped_exception(exc) from exc + if count > 0: + self._log_delete_expired(count) + return count def _table_ref(table_name: str) -> str: diff --git a/sqlspec/adapters/sybase/migrations.py b/sqlspec/adapters/sybase/migrations.py index b7757e427..4a24d1bf7 100644 --- a/sqlspec/adapters/sybase/migrations.py +++ b/sqlspec/adapters/sybase/migrations.py @@ -73,11 +73,7 @@ def _existing_columns_query(self) -> str: """Return Sybase ASE catalog query text for migration tracking table columns.""" has_schema = "." in self.version_table schema_name, table_name = split_schema_and_table(self.version_table) - schema_condition = ( - f"u.name = '{escape_sql_literal(schema_name)}'" - if has_schema - else "u.name = USER_NAME()" - ) + schema_condition = f"u.name = '{escape_sql_literal(schema_name)}'" if has_schema else "u.name = USER_NAME()" return f""" SELECT c.name AS column_name FROM dbo.syscolumns c diff --git a/sqlspec/extensions/adk/migrations/0002_create_mssql_python_memory.py b/sqlspec/extensions/adk/migrations/0002_create_mssql_python_memory.py index 2c5727415..a4be50a1d 100644 --- a/sqlspec/extensions/adk/migrations/0002_create_mssql_python_memory.py +++ b/sqlspec/extensions/adk/migrations/0002_create_mssql_python_memory.py @@ -26,20 +26,16 @@ async def up(context: "MigrationContext | None" = None) -> list[str]: if not any(cls.__module__.startswith("sqlspec.adapters.mssql_python.") for cls in type(config).__mro__): return [] - from sqlspec.adapters.mssql_python.adk.store import ( - MssqlPythonADKMemoryStore, - _create_index_sql, - _escape_sql_literal, - _table_ref, - ) + from sqlspec.adapters.mssql_python.adk.store import MssqlPythonADKMemoryStore, _create_index_sql, _table_ref + from sqlspec.adapters.mssql_python.core import escape_sql_literal store = MssqlPythonADKMemoryStore(config) statements = [store._memory_table_ddl()] # pyright: ignore[reportPrivateUsage] for index_name, table, columns in store._memory_index_specs(): # pyright: ignore[reportPrivateUsage] statements.append( "IF NOT EXISTS (SELECT 1 FROM sys.indexes " - f"WHERE name = N'{_escape_sql_literal(index_name)}' " - f"AND object_id = OBJECT_ID(N'{_escape_sql_literal(_table_ref(table))}')) " + f"WHERE name = N'{escape_sql_literal(index_name)}' " + f"AND object_id = OBJECT_ID(N'{escape_sql_literal(_table_ref(table))}')) " f"BEGIN {_create_index_sql(table, index_name, columns)}; END;" ) return statements diff --git a/tests/unit/adapters/test_sqlanywhere.py b/tests/unit/adapters/test_sqlanywhere.py index cf1e7770c..4f37d98ef 100644 --- a/tests/unit/adapters/test_sqlanywhere.py +++ b/tests/unit/adapters/test_sqlanywhere.py @@ -1,6 +1,8 @@ """Unit tests for the SAP SQL Anywhere (sqlanydb) adapter.""" import threading +from contextlib import nullcontext +from datetime import UTC, datetime from typing import Any, cast from unittest.mock import MagicMock @@ -16,15 +18,15 @@ SqlanywhereDriver, SqlanywhereExceptionHandler, SqlanywhereSyncDataDictionary, - build_column_metadata, build_connection_config, - build_foreign_key_metadata, create_mapped_exception, default_statement_config, driver_profile, ) -from sqlspec.adapters.sqlanywhere._typing import SqlanywhereError +from sqlspec.adapters.sqlanywhere._typing import SqlanywhereError, SqlanywhereSessionContext from sqlspec.adapters.sqlanywhere.adk import SqlanywhereADKConfig, SqlanywhereADKMemoryStore, SqlanywhereADKStore +from sqlspec.adapters.sqlanywhere.core import build_insert_statement, format_identifier +from sqlspec.adapters.sqlanywhere.data_dictionary import build_column_metadata, build_foreign_key_metadata from sqlspec.adapters.sqlanywhere.events import SqlanywhereEventQueueStore from sqlspec.adapters.sqlanywhere.litestar import SqlanywhereLitestarConfig, SqlanywhereStore from sqlspec.adapters.sqlanywhere.migrations import SqlanywhereSyncMigrationTracker @@ -538,7 +540,7 @@ def commit(self) -> None: def rollback(self) -> None: pass - monkeypatch.setattr(config, "provide_connection", lambda: DummyConnection()) + monkeypatch.setattr(SqlanywhereConfig, "provide_connection", lambda _self: nullcontext(DummyConnection())) entry = { "id": "mem-1", @@ -599,7 +601,7 @@ def cursor(self) -> DummyCursor: def commit(self) -> None: pass - monkeypatch.setattr(adapter_config, "provide_connection", lambda: DummyConnection()) + monkeypatch.setattr(SqlanywhereConfig, "provide_connection", lambda _self: nullcontext(DummyConnection())) assert litestar_store._delete_expired() == 4 @@ -614,8 +616,8 @@ def test_sqlanywhere_build_column_metadata_is_nullable() -> None: {"column_name": "id", "is_nullable": "NO"}, {"column_name": "val", "is_nullable": "YES"}, ]) - assert cols[0].is_nullable is False - assert cols[1].is_nullable is True + assert cols[0]["is_nullable"] is False + assert cols[1]["is_nullable"] is True def test_sqlanywhere_build_foreign_key_metadata_names() -> None: @@ -728,7 +730,9 @@ class DummyConn: def close(self) -> None: pass - monkeypatch.setattr(pool, "new_connection", lambda: cast("SqlanywhereConnection", DummyConn())) + monkeypatch.setattr( + SqlanywhereConnectionPool, "new_connection", lambda _self: cast("SqlanywhereConnection", DummyConn()) + ) errors: list[Exception] = [] conns: list[Any] = [] diff --git a/tests/unit/adapters/test_sybase.py b/tests/unit/adapters/test_sybase.py index 2eacc5104..435d2826a 100644 --- a/tests/unit/adapters/test_sybase.py +++ b/tests/unit/adapters/test_sybase.py @@ -3,6 +3,7 @@ import os import socket import threading +from contextlib import nullcontext from datetime import UTC, datetime from typing import Any, cast from unittest.mock import AsyncMock, MagicMock @@ -21,17 +22,19 @@ SybaseDriver, SybaseExceptionHandler, SybaseSyncDataDictionary, - _ERROR_CODE_MAPPING, build_connection_config, - build_insert_statement, create_mapped_exception, default_statement_config, driver_profile, - escape_literal_percent, - format_identifier, ) from sqlspec.adapters.sybase._typing import SybaseRawCursor from sqlspec.adapters.sybase.adk import SybaseADKConfig, SybaseADKMemoryStore, SybaseADKStore +from sqlspec.adapters.sybase.core import ( + _ERROR_CODE_MAPPING, + build_insert_statement, + escape_literal_percent, + format_identifier, +) from sqlspec.adapters.sybase.events import SybaseEventQueueStore from sqlspec.adapters.sybase.litestar import SybaseStore from sqlspec.adapters.sybase.migrations import SybaseSyncMigrationTracker @@ -53,7 +56,6 @@ SQLSpecError, TransactionError, UniqueViolationError, - UniqueViolationError, ) @@ -586,7 +588,7 @@ def cursor(self, as_dict: bool = False) -> DummyCursor: def rollback(self) -> None: pass - monkeypatch.setattr(config, "provide_connection", lambda: DummyConnection()) + monkeypatch.setattr(config, "provide_connection", lambda: nullcontext(DummyConnection())) with pytest.raises(UniqueViolationError): store._execute("INSERT INTO dummy VALUES (1)") @@ -594,6 +596,7 @@ def rollback(self) -> None: def test_sybase_litestar_store_get_raises_sqlspec_error(monkeypatch: pytest.MonkeyPatch) -> None: """SybaseStore._get() maps raw database errors to SQLSpecError.""" config = MagicMock() + config.extension_config = {"litestar": {}} store = SybaseStore(cast("Any", config)) class FailingCursor: @@ -607,7 +610,7 @@ class FailingConnection: def cursor(self) -> FailingCursor: return FailingCursor() - monkeypatch.setattr(config, "provide_connection", lambda: FailingConnection()) + monkeypatch.setattr(config, "provide_connection", lambda: nullcontext(FailingConnection())) with pytest.raises(SQLSpecError): store._get("session-1") @@ -654,6 +657,7 @@ def mock_load() -> Any: monkeypatch.setattr(sybase_pool, "_load_pymssql_module", mock_load) monkeypatch.setattr(sybase_pool, "pymssql", None) + monkeypatch.setenv("TDSVER", "5.0") pool = sybase_pool.SybaseConnectionPool({"server": "localhost", "database": "demo"}) monkeypatch.setattr("pymssql.connect", lambda **_kw: FakeConnection(FakeCursor())) @@ -671,18 +675,20 @@ def test_sybase_format_identifier_and_insert_statements() -> None: assert format_identifier("dbo.my_table") == "[dbo].[my_table]" sql = build_insert_statement("my_table", ["col1", "col2"]) assert sql == "INSERT INTO [my_table] ([col1], [col2]) VALUES (%s, %s)" - escaped = escape_literal_percent("SELECT * FROM t WHERE val LIKE '100%' AND id = %s", ("1",), default_statement_config.parameter_validator) + escaped = escape_literal_percent( + "SELECT * FROM t WHERE val LIKE '100%' AND id = %s", ("1",), default_statement_config.parameter_validator + ) assert "100%%" in escaped def test_sybase_migration_tracker_existing_columns_query() -> None: """_existing_columns_query resolves USER_NAME() when schema is unqualified.""" - tracker_unqualified = SybaseSyncMigrationTracker(version_table="my_versions") + tracker_unqualified = SybaseSyncMigrationTracker(version_table_name="my_versions") query_unqualified = tracker_unqualified._existing_columns_query() assert "u.name = USER_NAME()" in query_unqualified assert "o.name = 'my_versions'" in query_unqualified - tracker_qualified = SybaseSyncMigrationTracker(version_table="custom_schema.my_versions") + tracker_qualified = SybaseSyncMigrationTracker(version_table_name="custom_schema.my_versions") query_qualified = tracker_qualified._existing_columns_query() assert "u.name = 'custom_schema'" in query_qualified assert "o.name = 'my_versions'" in query_qualified diff --git a/tests/unit/data_dictionary/test_sybase_sqlanywhere_dialects.py b/tests/unit/data_dictionary/test_sybase_sqlanywhere_dialects.py index a2a72c5dd..9a22cd808 100644 --- a/tests/unit/data_dictionary/test_sybase_sqlanywhere_dialects.py +++ b/tests/unit/data_dictionary/test_sybase_sqlanywhere_dialects.py @@ -50,7 +50,7 @@ def test_sybase_and_sqlanywhere_dialect_configs_registered() -> None: assert SQLANYWHERE_CONFIG.get_feature_flag("supports_merge") is True assert SQLANYWHERE_CONFIG.get_feature_flag("supports_for_update") is True assert SQLANYWHERE_CONFIG.get_feature_flag("supports_skip_locked") is True - assert SQLANYWHERE_CONFIG.get_feature_flag("supports_transactional_ddl") is True + assert SQLANYWHERE_CONFIG.get_feature_flag("supports_transactional_ddl") is False assert SQLANYWHERE_CONFIG.get_optimal_type("uuid") == "UNIQUEIDENTIFIER" assert SQLANYWHERE_CONFIG.get_optimal_type("text") == "LONG VARCHAR" assert SQLANYWHERE_CONFIG.get_optimal_type("blob") == "LONG BINARY" diff --git a/tests/unit/extensions/test_schema_reconciliation.py b/tests/unit/extensions/test_schema_reconciliation.py index 33c196c82..90b0f1f29 100644 --- a/tests/unit/extensions/test_schema_reconciliation.py +++ b/tests/unit/extensions/test_schema_reconciliation.py @@ -31,11 +31,7 @@ class AdditiveSqliteADKMemoryStore(SqliteADKMemoryStore): """ADK memory target with one additive column.""" def _memory_table_ddl(self) -> str: - return ( - super() - ._memory_table_ddl() - .replace("embedding BLOB\n", "embedding BLOB,\n schema_tag TEXT\n") - ) + return super()._memory_table_ddl().replace("embedding BLOB\n", "embedding BLOB,\n schema_tag TEXT\n") class AdditiveSqliteEventQueueStore(SqliteEventQueueStore): From 4dcf757094dd85508b9814aaac4fa0e27acfd6b1 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 01:02:29 +0000 Subject: [PATCH 063/174] feat(oracledb): add and persist ADK memory embedding column Add vector_dimensions configuration to OracleADKConfig and declare the embedding VECTOR column in the ADK memory table template. Bind memory embeddings as float arrays on insert and decode them to float lists on retrieval across sync and async stores. --- sqlspec/adapters/oracledb/adk/store.py | 48 +++++-- .../extensions/adk/test_oracle_specific.py | 127 +++++++++++++++++- .../test_oracledb/test_oracle_adk_store.py | 82 +++++++++++ 3 files changed, 242 insertions(+), 15 deletions(-) diff --git a/sqlspec/adapters/oracledb/adk/store.py b/sqlspec/adapters/oracledb/adk/store.py index bbf7a7f64..c54d5aadd 100644 --- a/sqlspec/adapters/oracledb/adk/store.py +++ b/sqlspec/adapters/oracledb/adk/store.py @@ -1,5 +1,6 @@ """Oracle ADK store for Google Agent Development Kit session/event storage.""" +import array from decimal import Decimal from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, NotRequired, cast @@ -57,6 +58,7 @@ ORACLE_DEFAULT_APP_STATE_TABLE: Final = "adk_app_state" ORACLE_DEFAULT_USER_STATE_TABLE: Final = "adk_user_state" ORACLE_DEFAULT_METADATA_TABLE: Final = "adk_internal_metadata" +_EMBEDDING_COLUMN_INDEX: Final = 12 _ADK_METADATA_TABLE_DDL_TEMPLATE = ( @@ -245,6 +247,7 @@ " event_id VARCHAR2(128) NOT NULL UNIQUE,\n" " author VARCHAR2(256){1},\n" " timestamp TIMESTAMP WITH TIME ZONE NOT NULL,\n" + " embedding VECTOR({11}, FLOAT32),\n" " {2},\n" " content_text CLOB NOT NULL,\n" " inserted_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL\n" @@ -346,6 +349,9 @@ class OracleADKPartitionConfig(TypedDict): class OracleADKConfig(ADKConfig): """Oracle ADK extension settings consumed by Oracle ADK stores.""" + vector_dimensions: NotRequired[int] + """Dimensionality of embedding vectors (e.g. 768 for gemini-embedding-001 with MRL). Default: 768.""" + in_memory: NotRequired[bool] """Enable Oracle INMEMORY clauses on ADK tables.""" @@ -2436,10 +2442,10 @@ async def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: " sql = f""" INSERT INTO {self._memory_table} ( id, session_id, app_name, user_id, scope, event_id, author{owner_column}, - timestamp, content_json, content_text, metadata_json, inserted_at + timestamp, embedding, content_json, content_text, metadata_json, inserted_at ) VALUES ( :id, :session_id, :app_name, :user_id, :scope, :event_id, :author{owner_param}, - :timestamp, :content_json, :content_text, :metadata_json, :inserted_at + :timestamp, :embedding, :content_json, :content_text, :metadata_json, :inserted_at ) """ @@ -2449,6 +2455,8 @@ async def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: " for entry in entries: content_json = await self._serialize_json_field(entry["content_json"]) metadata_json = await self._serialize_json_field(entry["metadata_json"]) + embedding_val = entry.get("embedding") + bound_embedding = array.array("f", embedding_val) if embedding_val is not None else None params: dict[str, Any] = { "id": entry["id"], "session_id": entry["session_id"], @@ -2458,6 +2466,7 @@ async def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: " "event_id": entry["event_id"], "author": entry["author"], "timestamp": entry["timestamp"], + "embedding": bound_embedding, "content_json": content_json, "content_text": entry["content_text"], "metadata_json": metadata_json, @@ -2621,6 +2630,7 @@ def _memory_table_ddl_for_type(self, storage_type: "JSONStorageType") -> str: self._memory_table, self._memory_table, fts_index, + self._vector_dimensions, ) def _drop_memory_table_sql(self) -> "list[str]": @@ -2710,10 +2720,10 @@ async def _search_entries_fts( where_scope, scope_params = _build_oracle_scope_where(app_name, user_id, scope_filter) sql = f""" SELECT id, session_id, app_name, user_id, scope, event_id, author, - timestamp, content_json, content_text, metadata_json, inserted_at + timestamp, content_json, content_text, metadata_json, inserted_at, embedding FROM ( SELECT id, session_id, app_name, user_id, scope, event_id, author, - timestamp, content_json, content_text, metadata_json, inserted_at, + timestamp, content_json, content_text, metadata_json, inserted_at, embedding, SCORE(1) AS score FROM {self._memory_table} WHERE {where_scope} @@ -2759,6 +2769,7 @@ async def _search_entries_hybrid_rrf( "content_text", "metadata_json", "inserted_at", + "embedding", ) if not sanitized_query: sql = build_oracle_hnsw_subquery( @@ -2806,10 +2817,10 @@ async def _search_entries_simple( where_scope, scope_params = _build_oracle_scope_where(app_name, user_id, scope_filter) sql = f""" SELECT id, session_id, app_name, user_id, scope, event_id, author, - timestamp, content_json, content_text, metadata_json, inserted_at + timestamp, content_json, content_text, metadata_json, inserted_at, embedding FROM ( SELECT id, session_id, app_name, user_id, scope, event_id, author, - timestamp, content_json, content_text, metadata_json, inserted_at + timestamp, content_json, content_text, metadata_json, inserted_at, embedding FROM {self._memory_table} WHERE {where_scope} AND LOWER(content_text) LIKE :pattern @@ -2833,6 +2844,8 @@ async def _rows_to_records(self, rows: "list[Any]") -> "list[StoredMemory]": content_text = row[9] if is_async_readable(content_text) or is_readable(content_text): content_text = await _read_lob_async(content_text) + embedding_val = row[_EMBEDDING_COLUMN_INDEX] if len(row) > _EMBEDDING_COLUMN_INDEX else None + embedding = [float(x) for x in embedding_val] if embedding_val is not None else None records.append({ "id": row[0], "session_id": row[1], @@ -2846,7 +2859,7 @@ async def _rows_to_records(self, rows: "list[Any]") -> "list[StoredMemory]": "content_text": str(content_text), "metadata_json": metadata_json, "inserted_at": row[11], - "embedding": None, + "embedding": embedding, }) return records @@ -2910,10 +2923,10 @@ def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "object sql = f""" INSERT INTO {self._memory_table} ( id, session_id, app_name, user_id, scope, event_id, author{owner_column}, - timestamp, content_json, content_text, metadata_json, inserted_at + timestamp, embedding, content_json, content_text, metadata_json, inserted_at ) VALUES ( :id, :session_id, :app_name, :user_id, :scope, :event_id, :author{owner_param}, - :timestamp, :content_json, :content_text, :metadata_json, :inserted_at + :timestamp, :embedding, :content_json, :content_text, :metadata_json, :inserted_at ) """ @@ -2922,6 +2935,8 @@ def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "object for entry in entries: content_json = self._serialize_json_field(entry["content_json"]) metadata_json = self._serialize_json_field(entry["metadata_json"]) + embedding_val = entry.get("embedding") + bound_embedding = array.array("f", embedding_val) if embedding_val is not None else None params: dict[str, Any] = { "id": entry["id"], "session_id": entry["session_id"], @@ -2931,6 +2946,7 @@ def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "object "event_id": entry["event_id"], "author": entry["author"], "timestamp": entry["timestamp"], + "embedding": bound_embedding, "content_json": content_json, "content_text": entry["content_text"], "metadata_json": metadata_json, @@ -3093,6 +3109,7 @@ def _memory_table_ddl_for_type(self, storage_type: "JSONStorageType") -> str: self._memory_table, self._memory_table, fts_index, + self._vector_dimensions, ) def _drop_memory_table_sql(self) -> "list[str]": @@ -3182,10 +3199,10 @@ def _search_entries_fts( where_scope, scope_params = _build_oracle_scope_where(app_name, user_id, scope_filter) sql = f""" SELECT id, session_id, app_name, user_id, scope, event_id, author, - timestamp, content_json, content_text, metadata_json, inserted_at + timestamp, content_json, content_text, metadata_json, inserted_at, embedding FROM ( SELECT id, session_id, app_name, user_id, scope, event_id, author, - timestamp, content_json, content_text, metadata_json, inserted_at, + timestamp, content_json, content_text, metadata_json, inserted_at, embedding, SCORE(1) AS score FROM {self._memory_table} WHERE {where_scope} @@ -3230,6 +3247,7 @@ def _search_entries_hybrid_rrf( "content_text", "metadata_json", "inserted_at", + "embedding", ) if not sanitized_query: sql = build_oracle_hnsw_subquery( @@ -3276,10 +3294,10 @@ def _search_entries_simple( where_scope, scope_params = _build_oracle_scope_where(app_name, user_id, scope_filter) sql = f""" SELECT id, session_id, app_name, user_id, scope, event_id, author, - timestamp, content_json, content_text, metadata_json, inserted_at + timestamp, content_json, content_text, metadata_json, inserted_at, embedding FROM ( SELECT id, session_id, app_name, user_id, scope, event_id, author, - timestamp, content_json, content_text, metadata_json, inserted_at + timestamp, content_json, content_text, metadata_json, inserted_at, embedding FROM {self._memory_table} WHERE {where_scope} AND LOWER(content_text) LIKE :pattern @@ -3302,6 +3320,8 @@ def _rows_to_records(self, rows: "list[Any]") -> "list[StoredMemory]": content_text = row[9] if is_readable(content_text): content_text = _read_lob_sync(content_text) + embedding_val = row[_EMBEDDING_COLUMN_INDEX] if len(row) > _EMBEDDING_COLUMN_INDEX else None + embedding = [float(x) for x in embedding_val] if embedding_val is not None else None records.append({ "id": row[0], "session_id": row[1], @@ -3315,7 +3335,7 @@ def _rows_to_records(self, rows: "list[Any]") -> "list[StoredMemory]": "content_text": str(content_text), "metadata_json": metadata_json, "inserted_at": row[11], - "embedding": None, + "embedding": embedding, }) return records diff --git a/tests/integration/adapters/oracle/oracledb/extensions/adk/test_oracle_specific.py b/tests/integration/adapters/oracle/oracledb/extensions/adk/test_oracle_specific.py index 6ff61f13a..1427181e3 100644 --- a/tests/integration/adapters/oracle/oracledb/extensions/adk/test_oracle_specific.py +++ b/tests/integration/adapters/oracle/oracledb/extensions/adk/test_oracle_specific.py @@ -10,7 +10,12 @@ import pytest from sqlspec.adapters.oracledb import OracleAsyncConfig, OracleSyncConfig -from sqlspec.adapters.oracledb.adk import OracleAsyncADKStore, OracleSyncADKStore +from sqlspec.adapters.oracledb.adk import ( + OracleAsyncADKMemoryStore, + OracleAsyncADKStore, + OracleSyncADKMemoryStore, + OracleSyncADKStore, +) from sqlspec.extensions.adk import StoredEvent pytestmark = [pytest.mark.xdist_group("oracle"), pytest.mark.oracledb, pytest.mark.integration] @@ -485,3 +490,123 @@ async def test_create_session_with_owner_id_sync(oracle_store_sync_with_fk: "Ora retrieved = oracle_store_sync_with_fk.get_session(app_name, user_id, session_id) assert retrieved is not None assert retrieved["id"] == session_id + + +async def _cleanup_async_memory_store(store: "OracleAsyncADKMemoryStore", config: "OracleAsyncConfig") -> None: + """Drop ADK memory tables for async stores.""" + async with config.provide_connection() as conn: + cursor = conn.cursor() + for stmt in store._drop_memory_table_sql(): + try: + await cursor.execute(stmt) + except Exception: + pass + await conn.commit() + + +def _cleanup_sync_memory_store(store: "OracleSyncADKMemoryStore", config: "OracleSyncConfig") -> None: + """Drop ADK memory tables for sync stores.""" + with config.provide_connection() as conn: + cursor = conn.cursor() + for stmt in store._drop_memory_table_sql(): + try: + cursor.execute(stmt) + except Exception: + pass + conn.commit() + + +@pytest.fixture +async def oracle_async_memory_store( + oracle_async_config: "OracleAsyncConfig", +) -> "AsyncGenerator[OracleAsyncADKMemoryStore, None]": + """Create an async Oracle ADK memory store with tables created per test.""" + store = OracleAsyncADKMemoryStore(oracle_async_config) + await store.create_tables() + try: + yield store + finally: + await _cleanup_async_memory_store(store, oracle_async_config) + + +@pytest.fixture +async def oracle_sync_memory_store( + oracle_sync_config: "OracleSyncConfig", +) -> "AsyncGenerator[OracleSyncADKMemoryStore, None]": + """Create a sync Oracle ADK memory store with tables created per test.""" + store = OracleSyncADKMemoryStore(oracle_sync_config) + store.create_tables() + try: + yield store + finally: + _cleanup_sync_memory_store(store, oracle_sync_config) + + +async def test_oracle_async_adk_memory_store_insert_and_search_embedding( + oracle_async_memory_store: "OracleAsyncADKMemoryStore", +) -> None: + """Async memory store inserts embedding vector and executes vector search.""" + app_name = "test-app" + user_id = "user-vec" + now = datetime.now(UTC) + embedding = [0.1 * i for i in range(768)] + entries = [ + { + "id": _unique_session_id("mem-vec"), + "session_id": _unique_session_id("sess-vec"), + "app_name": app_name, + "user_id": user_id, + "scope": "user", + "event_id": _unique_session_id("ev-vec"), + "author": "user", + "timestamp": now, + "content_json": {"query": "vector search test"}, + "content_text": "vector search content", + "metadata_json": {"tag": "embedding"}, + "inserted_at": now, + "embedding": embedding, + } + ] + inserted = await oracle_async_memory_store.insert_memory_entries(entries) + assert inserted == 1 + + results = await oracle_async_memory_store.search_entries( + query="vector search", app_name=app_name, user_id=user_id, embedding=embedding + ) + assert len(results) >= 1 + assert results[0]["embedding"] is not None + + +def test_oracle_sync_adk_memory_store_insert_and_search_embedding( + oracle_sync_memory_store: "OracleSyncADKMemoryStore", +) -> None: + """Sync memory store inserts embedding vector and executes vector search.""" + app_name = "test-app" + user_id = "user-vec-sync" + now = datetime.now(UTC) + embedding = [0.1 * i for i in range(768)] + entries = [ + { + "id": _unique_session_id("mem-vec-sync"), + "session_id": _unique_session_id("sess-vec-sync"), + "app_name": app_name, + "user_id": user_id, + "scope": "user", + "event_id": _unique_session_id("ev-vec-sync"), + "author": "user", + "timestamp": now, + "content_json": {"query": "sync vector search test"}, + "content_text": "sync vector search content", + "metadata_json": {"tag": "embedding"}, + "inserted_at": now, + "embedding": embedding, + } + ] + inserted = oracle_sync_memory_store.insert_memory_entries(entries) + assert inserted == 1 + + results = oracle_sync_memory_store.search_entries( + query="sync vector search", app_name=app_name, user_id=user_id, embedding=embedding + ) + assert len(results) >= 1 + assert results[0]["embedding"] is not None diff --git a/tests/unit/adapters/test_oracledb/test_oracle_adk_store.py b/tests/unit/adapters/test_oracledb/test_oracle_adk_store.py index 0f0b89af7..62676a0a5 100644 --- a/tests/unit/adapters/test_oracledb/test_oracle_adk_store.py +++ b/tests/unit/adapters/test_oracledb/test_oracle_adk_store.py @@ -53,6 +53,7 @@ def test_oracle_adk_config_types_adapter_local_optimizations() -> None: assert cast("Any", ADKConfig).__optional_keys__ <= cast("Any", OracleADKConfig).__optional_keys__ expected_types: dict[str, object] = { + "vector_dimensions": int, "in_memory": bool, "compression": OracleADKCompressionConfig, "partitioning": OracleADKPartitionConfig, @@ -301,6 +302,87 @@ def test_oracle_sync_adk_memory_rows_to_records_deserializes_json_fields() -> No assert records[0]["embedding"] is None +def test_oracle_adk_memory_ddl_declares_embedding_vector() -> None: + """Memory table DDL declares an embedding VECTOR column with configured dimensions.""" + config = _mock_config({}) + store = OracleSyncADKMemoryStore(config) + sql = store._memory_table_ddl_for_type(JSONStorageType.JSON_NATIVE) + assert "embedding VECTOR(768, FLOAT32)" in sql + + custom_config = _mock_config({"vector_dimensions": 1536}) + custom_store = OracleAsyncADKMemoryStore(custom_config) + custom_sql = custom_store._memory_table_ddl_for_type(JSONStorageType.JSON_NATIVE) + assert "embedding VECTOR(1536, FLOAT32)" in custom_sql + + +def test_oracle_adk_sync_insert_binds_embedding_array() -> None: + """Sync memory store binds embeddings as float arrays.""" + import array + + config = _config_for_storage(JSONStorageType.JSON_NATIVE) + config.extension_config = {"adk": {}} + store = OracleSyncADKMemoryStore(config) + conn = MagicMock() + cursor = MagicMock() + cursor.rowcount = 1 + conn.cursor.return_value.__enter__.return_value = cursor + config.provide_connection.return_value.__enter__.return_value = conn + + timestamp = datetime(2026, 5, 10, 12, 0, tzinfo=UTC) + entries = [ + { + "id": "mem-1", + "session_id": "sess-1", + "app_name": "app", + "user_id": "user", + "scope": "user", + "event_id": "ev-1", + "author": "user", + "timestamp": timestamp, + "content_json": {"k": "v"}, + "content_text": "text", + "metadata_json": None, + "inserted_at": timestamp, + "embedding": [0.1, 0.2, 0.3], + } + ] + inserted = store.insert_memory_entries(entries) + assert inserted == 1 + assert cursor.execute.called + _call_sql, call_params = cursor.execute.call_args[0] + assert "embedding" in call_params + assert isinstance(call_params["embedding"], array.array) + assert call_params["embedding"].typecode == "f" + assert list(call_params["embedding"]) == pytest.approx([0.1, 0.2, 0.3]) + + +def test_oracle_adk_memory_rows_to_records_decodes_embedding() -> None: + """Memory rows to records round-trips vector embeddings as float lists.""" + import array + + store = OracleSyncADKMemoryStore.__new__(OracleSyncADKMemoryStore) # type: ignore[call-arg] + store._config = _config_for_storage(JSONStorageType.BLOB_JSON) + timestamp = datetime(2026, 5, 10, 12, 0, tzinfo=UTC) + row = ( + "memory-vec", + "session-vec", + "app", + "user", + "user", + "event-vec", + "user", + timestamp, + b'{"text":"vec"}', + "vec", + b'{"source":"unit"}', + timestamp, + array.array("f", [0.1, 0.2, 0.3]), + ) + records = store._rows_to_records([row]) + assert len(records) == 1 + assert records[0]["embedding"] == pytest.approx([0.1, 0.2, 0.3]) + + def _sync_store_with_driver() -> "tuple[Any, MagicMock, MagicMock]": config = _config_for_storage(JSONStorageType.BLOB_JSON) config.extension_config = {"adk": {}} From c6d9d6b9759c4e83d90b5511c1989ee006658254 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 01:02:29 +0000 Subject: [PATCH 064/174] fix(oracledb): resolve ORA-00997 on index expression catalog queries Remove the LONG column_expression from LISTAGG aggregation in the Oracle indexes catalog query. Split index expression retrieval into a dedicated non-aggregated catalog query and merge expressions in the data dictionary row builder across sync and async drivers. --- sqlspec/adapters/oracledb/data_dictionary.py | 80 ++++++++++++---- sqlspec/data_dictionary/_types.py | 4 + .../dialects/oracle/sql/indexes.sql | 21 +++-- .../adapters/oracle/oracledb/test_features.py | 25 +++++ .../test_data_dictionary_oracledb.py | 94 ++++++++++++++++++- 5 files changed, 199 insertions(+), 25 deletions(-) diff --git a/sqlspec/adapters/oracledb/data_dictionary.py b/sqlspec/adapters/oracledb/data_dictionary.py index 4c4a0f8fb..9b06a91db 100644 --- a/sqlspec/adapters/oracledb/data_dictionary.py +++ b/sqlspec/adapters/oracledb/data_dictionary.py @@ -1,5 +1,6 @@ """Oracle-specific data dictionary for metadata queries.""" +from collections import defaultdict from enum import StrEnum from typing import TYPE_CHECKING, Any, ClassVar, cast @@ -228,6 +229,33 @@ def resolve_oracle_connection_compatible_major( return None +def _merge_index_expressions( + indexes: "list[IndexMetadata]", expr_rows: "list[dict[str, Any]]", schema_name: "str | None" = None +) -> "list[IndexMetadata]": + """Merge expression-based index columns into index metadata records.""" + if not expr_rows: + return indexes + + exprs_by_index: dict[tuple[str, str], list[str]] = defaultdict(list) + for row in expr_rows: + owner = str(row.get("schema_name") or row.get("index_owner") or schema_name or "").upper() + idx_name = str(row.get("index_name") or "").upper() + expr = row.get("column_expression") + if expr is not None: + exprs_by_index[(owner, idx_name)].append(str(expr)) + + for idx in indexes: + owner = str(idx.get("schema_name") or schema_name or "").upper() + idx_name = str(idx.get("index_name") or "").upper() + matching_exprs = exprs_by_index.get((owner, idx_name)) + if matching_exprs: + expr_str = ",".join(matching_exprs) + idx["expressions"] = expr_str + idx["expression"] = expr_str + + return indexes + + @mypyc_attr(allow_interpreted_subclasses=True, native_class=False) class OracledbSyncDataDictionary(SyncDataDictionaryBase): """Oracle-specific sync data dictionary.""" @@ -473,17 +501,25 @@ def get_indexes( schema_name = self.resolve_schema(schema) if table is None: self._log_schema_introspect(driver, schema_name=schema_name, table_name=None, operation="indexes") - return driver.select( + indexes: list[IndexMetadata] = driver.select( self.get_query("indexes", "by_schema"), schema_name=schema_name, schema_type=IndexMetadata ) - table_name = self.resolve_identifier(table) - self._log_table_describe(driver, schema_name=schema_name, table_name=table_name, operation="indexes") - return driver.select( - self.get_query("indexes", "by_table"), - schema_name=schema_name, - table_name=table_name, - schema_type=IndexMetadata, - ) + expr_rows: list[dict[str, Any]] = driver.select( + self.get_query("indexes", "expressions"), schema_name=schema_name, table_name=None + ) + else: + table_name = self.resolve_identifier(table) + self._log_table_describe(driver, schema_name=schema_name, table_name=table_name, operation="indexes") + indexes = driver.select( + self.get_query("indexes", "by_table"), + schema_name=schema_name, + table_name=table_name, + schema_type=IndexMetadata, + ) + expr_rows = driver.select( + self.get_query("indexes", "expressions"), schema_name=schema_name, table_name=table_name + ) + return _merge_index_expressions(indexes, expr_rows, schema_name=schema_name) def get_foreign_keys( self, driver: "OracleSyncDriver", table: "str | None" = None, schema: "str | None" = None @@ -890,17 +926,25 @@ async def get_indexes( schema_name = self.resolve_schema(schema) if table is None: self._log_schema_introspect(driver, schema_name=schema_name, table_name=None, operation="indexes") - return await driver.select( + indexes: list[IndexMetadata] = await driver.select( self.get_query("indexes", "by_schema"), schema_name=schema_name, schema_type=IndexMetadata ) - table_name = self.resolve_identifier(table) - self._log_table_describe(driver, schema_name=schema_name, table_name=table_name, operation="indexes") - return await driver.select( - self.get_query("indexes", "by_table"), - schema_name=schema_name, - table_name=table_name, - schema_type=IndexMetadata, - ) + expr_rows: list[dict[str, Any]] = await driver.select( + self.get_query("indexes", "expressions"), schema_name=schema_name, table_name=None + ) + else: + table_name = self.resolve_identifier(table) + self._log_table_describe(driver, schema_name=schema_name, table_name=table_name, operation="indexes") + indexes = await driver.select( + self.get_query("indexes", "by_table"), + schema_name=schema_name, + table_name=table_name, + schema_type=IndexMetadata, + ) + expr_rows = await driver.select( + self.get_query("indexes", "expressions"), schema_name=schema_name, table_name=table_name + ) + return _merge_index_expressions(indexes, expr_rows, schema_name=schema_name) async def get_foreign_keys( self, driver: "OracleAsyncDriver", table: "str | None" = None, schema: "str | None" = None diff --git a/sqlspec/data_dictionary/_types.py b/sqlspec/data_dictionary/_types.py index 57eabc489..1d8396a5a 100644 --- a/sqlspec/data_dictionary/_types.py +++ b/sqlspec/data_dictionary/_types.py @@ -1260,6 +1260,10 @@ class IndexMetadata(TypedDict, total=False): """Whether the index enforces uniqueness.""" is_primary: bool | int """Whether the index backs the primary key constraint.""" + expression: str | None + """Single expression for expression-based indexes.""" + expressions: str | None + """Comma-separated expressions for expression-based indexes.""" @mypyc_attr(allow_interpreted_subclasses=False) diff --git a/sqlspec/data_dictionary/dialects/oracle/sql/indexes.sql b/sqlspec/data_dictionary/dialects/oracle/sql/indexes.sql index e17d28a92..8c5125ce9 100644 --- a/sqlspec/data_dictionary/dialects/oracle/sql/indexes.sql +++ b/sqlspec/data_dictionary/dialects/oracle/sql/indexes.sql @@ -20,16 +20,11 @@ SELECT i.avg_leaf_blocks_per_key, i.avg_data_blocks_per_key, i.clustering_factor, - LISTAGG(ic.column_name, ',') WITHIN GROUP (ORDER BY ic.column_position) AS columns, - LISTAGG(ie.column_expression, ',') WITHIN GROUP (ORDER BY ie.column_position) AS expressions + LISTAGG(ic.column_name, ',') WITHIN GROUP (ORDER BY ic.column_position) AS columns FROM all_indexes i LEFT JOIN all_ind_columns ic ON ic.index_owner = i.owner AND ic.index_name = i.index_name -LEFT JOIN all_ind_expressions ie - ON ie.index_owner = i.owner - AND ie.index_name = i.index_name - AND ie.column_position = ic.column_position WHERE i.owner = COALESCE(:schema_name, USER) AND (:table_name IS NULL OR i.table_name = :table_name) GROUP BY @@ -76,3 +71,17 @@ LEFT JOIN all_constraints c WHERE i.table_name = :table_name AND i.owner = COALESCE(:schema_name, USER) GROUP BY i.index_name, i.table_name, i.uniqueness; + +-- name: expressions +-- dialect: oracle +SELECT + ie.index_owner AS schema_name, + ie.table_name, + ie.index_name, + ie.column_position, + ie.column_expression +FROM all_ind_expressions ie +WHERE ie.index_owner = COALESCE(:schema_name, USER) + AND (:table_name IS NULL OR ie.table_name = :table_name) +ORDER BY ie.index_owner, ie.table_name, ie.index_name, ie.column_position; + diff --git a/tests/integration/adapters/oracle/oracledb/test_features.py b/tests/integration/adapters/oracle/oracledb/test_features.py index c8d0355a0..95d1be119 100644 --- a/tests/integration/adapters/oracle/oracledb/test_features.py +++ b/tests/integration/adapters/oracle/oracledb/test_features.py @@ -226,3 +226,28 @@ async def test_async_oracle_exception_handling(oracle_async_session: OracleAsync await oracle_async_session.execute_script( "BEGIN EXECUTE IMMEDIATE 'DROP TABLE test_exception_table_oracledb_async'; EXCEPTION WHEN OTHERS THEN IF SQLCODE != -942 THEN RAISE; END IF; END;" ) + + +def test_oracle_get_indexes_function_based_index(oracle_sync_session: OracleSyncDriver) -> None: + """Test get_indexes returns index and expression for function-based indexes.""" + oracle_sync_session.execute_script( + "BEGIN EXECUTE IMMEDIATE 'DROP TABLE test_fn_idx_table'; EXCEPTION WHEN OTHERS THEN IF SQLCODE != -942 THEN RAISE; END IF; END;" + ) + oracle_sync_session.execute_script(""" + CREATE TABLE test_fn_idx_table ( + id NUMBER(10) PRIMARY KEY, + email VARCHAR2(255) + ) + """) + oracle_sync_session.execute_script(""" + CREATE INDEX idx_test_fn_email ON test_fn_idx_table (UPPER(email)) + """) + + indexes = oracle_sync_session.data_dictionary.get_indexes(oracle_sync_session, table="test_fn_idx_table") + fn_indexes = [idx for idx in indexes if idx.get("index_name") == "IDX_TEST_FN_EMAIL"] + assert len(fn_indexes) == 1 + assert fn_indexes[0].get("expressions") is not None or fn_indexes[0].get("expression") is not None + + oracle_sync_session.execute_script( + "BEGIN EXECUTE IMMEDIATE 'DROP TABLE test_fn_idx_table'; EXCEPTION WHEN OTHERS THEN IF SQLCODE != -942 THEN RAISE; END IF; END;" + ) diff --git a/tests/unit/adapters/test_oracledb/test_data_dictionary_oracledb.py b/tests/unit/adapters/test_oracledb/test_data_dictionary_oracledb.py index ed73f3fd3..7868bea48 100644 --- a/tests/unit/adapters/test_oracledb/test_data_dictionary_oracledb.py +++ b/tests/unit/adapters/test_oracledb/test_data_dictionary_oracledb.py @@ -298,7 +298,7 @@ def test_oracle_domain_query_packs_cover_c5_sources() -> None: "tables": ("by_schema", "ALL_TABLES"), "columns": ("by_schema", "ALL_TAB_COLS"), "constraints": ("by_schema", "ALL_CONSTRAINTS"), - "indexes": ("by_schema", "ALL_IND_EXPRESSIONS"), + "indexes": ("by_schema", "ALL_INDEXES"), "views": ("by_schema", "ALL_VIEWS"), "materialized_views": ("by_schema", "ALL_MVIEWS"), "sequences": ("by_schema", "ALL_SEQUENCES"), @@ -526,3 +526,95 @@ async def test_oracle_26ai_async_data_dictionary_introspection_methods() -> None dviews = await data_dictionary.get_json_duality_views(async_driver, schema="app") assert isinstance(dviews, list) + + +def test_oracle_indexes_sql_does_not_listagg_column_expression() -> None: + """Oracle indexes SQL text must not aggregate column_expression via LISTAGG.""" + data_dictionary = OracledbSyncDataDictionary() + by_schema_sql = str(data_dictionary.get_query("indexes", "by_schema")).upper() + by_table_sql = str(data_dictionary.get_query("indexes", "by_table")).upper() + + assert "LISTAGG(IE.COLUMN_EXPRESSION" not in by_schema_sql + assert "COLUMN_EXPRESSION" not in by_schema_sql + assert "LISTAGG(IE.COLUMN_EXPRESSION" not in by_table_sql + + +def test_oracle_sync_get_indexes_merges_expressions() -> None: + """get_indexes must query expressions and merge them into index records.""" + data_dictionary = OracledbSyncDataDictionary() + + class _StubSyncDriver: + def __init__(self) -> None: + self.queries: list[str] = [] + + def select(self, query: Any, **kwargs: Any) -> list[dict[str, Any]]: + sql_text = str(query) + self.queries.append(sql_text) + if "all_ind_expressions" in sql_text: + return [ + { + "schema_name": "HR", + "table_name": "EMPLOYEES", + "index_name": "IDX_EMP_UPPER_EMAIL", + "column_position": 1, + "column_expression": "UPPER(EMAIL)", + } + ] + return [ + { + "schema_name": "HR", + "table_name": "EMPLOYEES", + "index_name": "IDX_EMP_UPPER_EMAIL", + "index_type": "NORMAL", + "uniqueness": "NONUNIQUE", + "columns": None, + } + ] + + driver = cast("OracleSyncDriver", _StubSyncDriver()) + indexes = data_dictionary.get_indexes(driver, schema="HR") + assert len(indexes) == 1 + assert indexes[0]["index_name"] == "IDX_EMP_UPPER_EMAIL" + assert indexes[0]["expressions"] == "UPPER(EMAIL)" + assert indexes[0]["expression"] == "UPPER(EMAIL)" + + +@pytest.mark.anyio +async def test_oracle_async_get_indexes_merges_expressions() -> None: + """Async get_indexes must query expressions and merge them into index records.""" + data_dictionary = OracledbAsyncDataDictionary() + + class _StubAsyncDriver: + def __init__(self) -> None: + self.queries: list[str] = [] + + async def select(self, query: Any, **kwargs: Any) -> list[dict[str, Any]]: + sql_text = str(query) + self.queries.append(sql_text) + if "all_ind_expressions" in sql_text: + return [ + { + "schema_name": "HR", + "table_name": "EMPLOYEES", + "index_name": "IDX_EMP_UPPER_EMAIL", + "column_position": 1, + "column_expression": "UPPER(EMAIL)", + } + ] + return [ + { + "schema_name": "HR", + "table_name": "EMPLOYEES", + "index_name": "IDX_EMP_UPPER_EMAIL", + "index_type": "NORMAL", + "uniqueness": "NONUNIQUE", + "columns": None, + } + ] + + driver = cast("OracleAsyncDriver", _StubAsyncDriver()) + indexes = await data_dictionary.get_indexes(driver, schema="HR") + assert len(indexes) == 1 + assert indexes[0]["index_name"] == "IDX_EMP_UPPER_EMAIL" + assert indexes[0]["expressions"] == "UPPER(EMAIL)" + assert indexes[0]["expression"] == "UPPER(EMAIL)" From 38a4681cf556226d6149fc89c36774f42304ded9 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 01:02:29 +0000 Subject: [PATCH 065/174] fix(oracledb): correct exception handling, script parameter binding, and transaction lifecycles Relocate PyCapsule Arrow conversion helpers to shared arrow_helpers utilities with lazy import loading, enforce deterministic float32 vector sequence packing, and ensure negative rowcounts normalize to zero. Wrap direct storage and stream executions in exception handlers to preserve cause chains, route cache hits with fetch kwargs, ensure connection contexts rollback on unhandled block exceptions, and reset active transaction state on failed commits. --- sqlspec/adapters/oracledb/__init__.py | 10 +- sqlspec/adapters/oracledb/_vector_handlers.py | 8 +- sqlspec/adapters/oracledb/config.py | 6 + sqlspec/adapters/oracledb/core.py | 138 ++--- sqlspec/adapters/oracledb/driver.py | 471 +++++++++++++----- sqlspec/adapters/oracledb/type_converter.py | 6 +- sqlspec/utils/arrow_helpers.py | 77 ++- .../test_oracledb/test_arrow_batches.py | 31 +- .../test_oracledb/test_direct_path_load.py | 21 + .../adapters/test_oracledb/test_driver.py | 58 +++ .../test_oracledb/test_fetch_tuning.py | 11 + .../test_oracledb/test_lob_coercion.py | 14 + .../test_stack_native_transaction.py | 18 + .../test_oracledb/test_transaction_state.py | 60 +++ .../test_oracledb/test_vector_handlers.py | 18 +- 15 files changed, 686 insertions(+), 261 deletions(-) diff --git a/sqlspec/adapters/oracledb/__init__.py b/sqlspec/adapters/oracledb/__init__.py index 4e6361480..9b995ce8b 100644 --- a/sqlspec/adapters/oracledb/__init__.py +++ b/sqlspec/adapters/oracledb/__init__.py @@ -24,8 +24,6 @@ build_session_release_plsql, build_session_reset_plsql, build_tns_description, - coerce_oracle_df_to_arrow_batches, - coerce_oracle_df_to_arrow_table, default_statement_config, execute_session_checkout_async, execute_session_checkout_sync, @@ -35,8 +33,6 @@ extract_duality_view_etag, extract_oracle_wallet_zip, inject_duality_view_etag, - oracle_df_batches_to_arrow_reader, - oracle_df_to_arrow_table, parse_sqlnet_ora, parse_tnsnames_ora, sanitize_oracle_text_query, @@ -53,6 +49,12 @@ OracleSyncExceptionHandler, ) from sqlspec.adapters.oracledb.type_converter import OracleOutputConverter +from sqlspec.utils.arrow_helpers import ( + coerce_oracle_df_to_arrow_batches, + coerce_oracle_df_to_arrow_table, + oracle_df_batches_to_arrow_reader, + oracle_df_to_arrow_table, +) __all__ = ( "VALID_ORACLE_DRIVER_FEATURE_KEYS", diff --git a/sqlspec/adapters/oracledb/_vector_handlers.py b/sqlspec/adapters/oracledb/_vector_handlers.py index 7f0856ae2..f1ac86679 100644 --- a/sqlspec/adapters/oracledb/_vector_handlers.py +++ b/sqlspec/adapters/oracledb/_vector_handlers.py @@ -50,8 +50,6 @@ _TYPECODE_INT16 = "h" _TYPECODE_INT32 = "i" -_INT8_MIN = -128 -_INT8_MAX = 127 _ORACLE_VECTOR_MIN_MAJOR = 23 _VECTOR_RETURN_NUMPY = "numpy" @@ -174,12 +172,8 @@ def _is_vector_payload(value: Any) -> bool: def _pack_python_sequence(value: "list[Any] | tuple[Any, ...]") -> "array.array[Any]": """Pack a Python sequence into an ``array.array`` for VECTOR binding. - Integer sequences entirely within ``[-128, 127]`` use the int8 typecode for - a cheaper bind; everything else falls back to float32 (the 23ai default and - the most common LLM embedding dtype). + Always packs as float32, the standard 23ai default for embedding vectors. """ - if all(isinstance(v, int) and not isinstance(v, bool) and _INT8_MIN <= v <= _INT8_MAX for v in value): - return array.array(_TYPECODE_INT8, value) return array.array(_TYPECODE_FLOAT32, [float(v) for v in value]) diff --git a/sqlspec/adapters/oracledb/config.py b/sqlspec/adapters/oracledb/config.py index 150678834..c39e9ebae 100644 --- a/sqlspec/adapters/oracledb/config.py +++ b/sqlspec/adapters/oracledb/config.py @@ -332,6 +332,9 @@ def __exit__( if self._conn is not None: conn = self._conn self._conn = None + if exc_type is not None: + with contextlib.suppress(Exception): + conn.rollback() session_hook = self._config.driver_features.get("session_context_hook") try: if session_hook is not None: @@ -669,6 +672,9 @@ async def __aexit__( self, exc_type: "type[BaseException] | None", exc_val: "BaseException | None", exc_tb: Any ) -> "bool | None": conn = self._connection + if conn is not None and exc_type is not None: + with contextlib.suppress(Exception): + await conn.rollback() session_hook = self._config.driver_features.get("session_context_hook") if conn is not None else None try: if conn is not None and session_hook is not None: diff --git a/sqlspec/adapters/oracledb/core.py b/sqlspec/adapters/oracledb/core.py index fa803ab4c..75bfcf60f 100644 --- a/sqlspec/adapters/oracledb/core.py +++ b/sqlspec/adapters/oracledb/core.py @@ -56,7 +56,7 @@ from sqlspec.observability import get_trace_context from sqlspec.typing import NUMPY_INSTALLED from sqlspec.utils.correlation import CorrelationContext -from sqlspec.utils.module_loader import ensure_pyarrow, import_optional +from sqlspec.utils.module_loader import import_optional from sqlspec.utils.serializers import from_json, to_json from sqlspec.utils.text import quote_identifier from sqlspec.utils.type_converters import build_uuid_coercions @@ -73,7 +73,6 @@ ) from sqlspec.core import SQL from sqlspec.driver import AsyncExceptionHandler, SyncExceptionHandler - from sqlspec.typing import ArrowRecordBatch, ArrowRecordBatchReader, ArrowTable __all__ = ( @@ -119,8 +118,6 @@ "coerce_large_parameters_sync", "coerce_many_parameters_async", "coerce_many_parameters_sync", - "coerce_oracle_df_to_arrow_batches", - "coerce_oracle_df_to_arrow_table", "collect_async_rows", "collect_sync_rows", "connection_is_thin", @@ -138,8 +135,6 @@ "normalize_column_names", "normalize_execute_many_parameters_async", "normalize_execute_many_parameters_sync", - "oracle_df_batches_to_arrow_reader", - "oracle_df_to_arrow_table", "parse_sqlnet_ora", "parse_tnsnames_ora", "release_extracted_wallet_dir", @@ -268,7 +263,6 @@ }) _VERSION_COMPONENTS: int = 3 TYPE_CONVERTER = OracleOutputConverter() -_pa: Any = import_optional("pyarrow") _LOB_TYPE_NAME_MARKERS: "tuple[str, ...]" = ("LOB", "BFILE") _SCALAR_PASSTHROUGH_TYPES: "tuple[type[Any], ...]" = (bool, int, float, str, bytes, bytearray, type(None)) _BIND_PASSTHROUGH_TYPES: "tuple[type[Any], ...]" = (bool, int, float, type(None)) @@ -834,71 +828,6 @@ def build_arrow_fetch_kwargs( return fetch_kwargs -def oracle_df_to_arrow_table(oracle_df: Any, pa_module: Any = None, arrow_schema: Any = None) -> "ArrowTable": - """Convert an Oracle DataFrame to a PyArrow Table using zero-copy PyCapsule when supported.""" - if pa_module is None: - ensure_pyarrow() - pa_module = _pa - - if hasattr(oracle_df, "__arrow_c_stream__"): - reader = ( - pa_module.RecordBatchReader.from_stream(oracle_df, schema=arrow_schema) - if arrow_schema is not None - else pa_module.RecordBatchReader.from_stream(oracle_df) - ) - return cast("ArrowTable", reader.read_all()) - table = pa_module.table(oracle_df) - if arrow_schema is not None: - table = table.cast(arrow_schema) - return cast("ArrowTable", table) - - -def coerce_oracle_df_to_arrow_table(oracle_df: Any, pa_module: Any = None, arrow_schema: Any = None) -> "ArrowTable": - """Alias for :func:`oracle_df_to_arrow_table`.""" - return oracle_df_to_arrow_table(oracle_df, pa_module=pa_module, arrow_schema=arrow_schema) - - -def coerce_oracle_df_to_arrow_batches( - oracle_df: Any, pa_module: Any = None, arrow_schema: Any = None -) -> "list[ArrowRecordBatch]": - """Convert an Oracle DataFrame batch into PyArrow RecordBatches via PyCapsule when supported.""" - if pa_module is None: - ensure_pyarrow() - pa_module = _pa - - if hasattr(oracle_df, "__arrow_c_stream__"): - reader = ( - pa_module.RecordBatchReader.from_stream(oracle_df, schema=arrow_schema) - if arrow_schema is not None - else pa_module.RecordBatchReader.from_stream(oracle_df) - ) - return cast("list[ArrowRecordBatch]", list(reader)) - table = pa_module.table(oracle_df) - if arrow_schema is not None: - table = table.cast(arrow_schema) - return cast("list[ArrowRecordBatch]", table.to_batches()) - - -def oracle_df_batches_to_arrow_reader( - batches: "Iterable[Any]", schema: Any = None, pa_module: Any = None -) -> "ArrowRecordBatchReader": - """Convert an iterable of Oracle DataFrame batches into a PyArrow RecordBatchReader.""" - if pa_module is None: - ensure_pyarrow() - pa_module = _pa - - collected_batches: list[Any] = [] - for df_batch in batches: - collected_batches.extend(coerce_oracle_df_to_arrow_batches(df_batch, pa_module=pa_module, arrow_schema=schema)) - if collected_batches: - resolved_schema = schema if schema is not None else collected_batches[0].schema - return cast( - "ArrowRecordBatchReader", pa_module.RecordBatchReader.from_batches(resolved_schema, collected_batches) - ) - resolved_schema = schema if schema is not None else pa_module.schema([]) - return cast("ArrowRecordBatchReader", pa_module.RecordBatchReader.from_batches(resolved_schema, [])) - - def normalize_column_names(column_names: "list[str]", driver_features: "dict[str, Any]") -> "list[str]": should_lowercase = driver_features.get("enable_lowercase_column_names", False) if not should_lowercase: @@ -1196,7 +1125,7 @@ def resolve_rowcount(cursor: Any) -> int: if not has_rowcount(cursor): return 0 rowcount = cursor.rowcount - if isinstance(rowcount, int): + if isinstance(rowcount, int) and rowcount >= 0: return rowcount return 0 @@ -1408,41 +1337,42 @@ def __init__( def start(self) -> None: handler = self._driver.handle_database_exceptions() - with handler: - try: - cursor = self._driver.connection.cursor() - self._cursor = cursor - cursor.arraysize = self._chunk_size - cursor.prefetchrows = self._chunk_size - fetch_kwargs = build_fetch_kwargs(self._driver.driver_features) - if self._fetch_lobs is not None: - fetch_kwargs["fetch_lobs"] = self._fetch_lobs - parameters = coerce_large_parameters_sync( - self._driver.connection, - self._parameters, - clob_type=DB_TYPE_CLOB, - blob_type=DB_TYPE_BLOB, - varchar2_byte_limit=self._driver.driver_features.get("oracle_varchar2_byte_limit", 4000), - raw_byte_limit=self._driver.driver_features.get("oracle_raw_byte_limit", 2000), - version_cache=getattr(self._driver, "_oracle_version_cache", None), - ) - cast("Any", cursor).execute(self._sql, parameters or {}, **fetch_kwargs) - except BaseException: - self.close(error=True) - raise + self._driver._run_with_exception_handler(handler, self._start) self._driver._check_pending_exception(handler) + def _start(self) -> None: + try: + cursor = self._driver.connection.cursor() + self._cursor = cursor + cursor.arraysize = self._chunk_size + cursor.prefetchrows = self._chunk_size + fetch_kwargs = build_fetch_kwargs(self._driver.driver_features) + if self._fetch_lobs is not None: + fetch_kwargs["fetch_lobs"] = self._fetch_lobs + parameters = coerce_large_parameters_sync( + self._driver.connection, + self._parameters, + clob_type=DB_TYPE_CLOB, + blob_type=DB_TYPE_BLOB, + varchar2_byte_limit=self._driver.driver_features.get("oracle_varchar2_byte_limit", 4000), + raw_byte_limit=self._driver.driver_features.get("oracle_raw_byte_limit", 2000), + version_cache=getattr(self._driver, "_oracle_version_cache", None), + ) + cast("Any", cursor).execute(self._sql, parameters or {}, **fetch_kwargs) + except BaseException: + self.close(error=True) + raise + def fetch_chunk(self) -> "list[dict[str, Any]]": handler = self._driver.handle_database_exceptions() cursor = self._cursor if cursor is None: return [] - rows: list[Any] = [] - with handler: - rows = cursor.fetchmany(self._chunk_size) + fetched = self._driver._run_with_exception_handler(handler, cursor.fetchmany, self._chunk_size) self._driver._check_pending_exception(handler) - if not rows: + if not fetched: return [] + rows: list[Any] = list(fetched) column_names = self._column_names if column_names is None: column_names, requires_lob_coercion = self._driver._resolve_row_metadata(cursor.description) @@ -1619,6 +1549,9 @@ def collect_sync_rows( if not fetched_data: return [], resolved_column_names + if driver_features.get("fetch_lobs"): + return [row if isinstance(row, tuple) else tuple(row) for row in fetched_data], resolved_column_names + if requires_lob_coercion is None: requires_lob_coercion = _description_requires_lob_coercion(description) if not requires_lob_coercion: @@ -1670,6 +1603,9 @@ async def collect_async_rows( if not fetched_data: return [], resolved_column_names + if driver_features.get("fetch_lobs"): + return [row if isinstance(row, tuple) else tuple(row) for row in fetched_data], resolved_column_names + if requires_lob_coercion is None: requires_lob_coercion = _description_requires_lob_coercion(description) if not requires_lob_coercion: @@ -3137,6 +3073,10 @@ def handle_database_exceptions(self) -> "SyncExceptionHandler": ... def _check_pending_exception(self, exc_handler: "SyncExceptionHandler") -> None: ... + def _run_with_exception_handler( + self, exc_handler: "SyncExceptionHandler", operation: "Callable[..., Any]", *args: Any, **kwargs: Any + ) -> "Any | None": ... + def _resolve_row_metadata(self, description: object) -> tuple[list[str], bool]: ... diff --git a/sqlspec/adapters/oracledb/driver.py b/sqlspec/adapters/oracledb/driver.py index 174e6595f..2385e5fac 100644 --- a/sqlspec/adapters/oracledb/driver.py +++ b/sqlspec/adapters/oracledb/driver.py @@ -2,6 +2,7 @@ import logging from pathlib import Path +from time import perf_counter from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, Protocol, cast, overload from sqlspec.adapters.oracledb._typing import ( @@ -32,7 +33,6 @@ coerce_large_parameters_sync, coerce_many_parameters_async, coerce_many_parameters_sync, - coerce_oracle_df_to_arrow_table, collect_async_rows, collect_sync_rows, connection_is_thin, @@ -54,6 +54,8 @@ ) from sqlspec.core import ( SQL, + DMLResult, + SQLResult, StackResult, StatementConfig, StatementStack, @@ -82,6 +84,7 @@ validate_savepoint_name, ) from sqlspec.exceptions import ImproperConfigurationError, SQLSpecError, StackExecutionError +from sqlspec.utils.arrow_helpers import coerce_oracle_df_to_arrow_table from sqlspec.utils.logging import get_logger, log_with_context from sqlspec.utils.module_loader import ensure_pyarrow, import_optional from sqlspec.utils.text import normalize_identifier, quote_identifier, split_qualified_identifier @@ -93,6 +96,7 @@ from sqlspec.builder import QueryBuilder from sqlspec.core import ArrowResult, Statement, StatementFilter + from sqlspec.core.cache import CachedQuery from sqlspec.core.stack import StackOperation from sqlspec.driver import ExecutionResult from sqlspec.storage import StorageBridgeJob, StorageDestination, StorageFormat, StorageTelemetry @@ -295,10 +299,11 @@ def _build_stack_results_from_pipeline( except AttributeError: error = None if error is not None: + mapped_error = create_mapped_exception(error) if not isinstance(error, SQLSpecError) else error stack_error = StackExecutionError( index, compiled.summary, - error, + mapped_error, adapter=type(self).__name__, mode="continue-on-error" if continue_on_error else "fail-fast", ) @@ -324,8 +329,9 @@ def _wrap_pipeline_error( self, error: Exception, stack: "StatementStack", continue_on_error: bool ) -> StackExecutionError: mode = "continue-on-error" if continue_on_error else "fail-fast" + mapped_error = create_mapped_exception(error) if not isinstance(error, SQLSpecError) else error return StackExecutionError( - -1, "Oracle pipeline execution failed", error, adapter=type(self).__name__, mode=mode + -1, "Oracle pipeline execution failed", mapped_error, adapter=type(self).__name__, mode=mode ) @@ -476,6 +482,47 @@ def dispatch_execute(self, cursor: Any, statement: "SQL") -> "ExecutionResult": affected_rows = resolve_rowcount(cursor) return self.create_execution_result(cursor, rowcount_override=affected_rows) + def _execute_cache_hit_with_cursor( + self, sql: str, params: "tuple[Any, ...] | list[Any] | dict[str, Any]", cached: "CachedQuery" + ) -> "SQLResult": + """Execute a cached query while owning only the cursor context.""" + with self.with_cursor(self.connection) as cursor: + wire_sql = self._format_wire_sql(cached.compiled_sql) + returns_rows = cached.operation_profile.returns_rows + fetch_kwargs = build_fetch_kwargs(self.driver_features) + parameters = coerce_large_parameters_sync( + self.connection, + params, + clob_type=DB_TYPE_CLOB, + blob_type=DB_TYPE_BLOB, + varchar2_byte_limit=self.driver_features.get("oracle_varchar2_byte_limit", 4000), + raw_byte_limit=self.driver_features.get("oracle_raw_byte_limit", 2000), + version_cache=self._oracle_version_cache, + ) + cast("Any", cursor).execute(wire_sql, parameters or {}, **fetch_kwargs) + if returns_rows and cursor.description is not None: + fetched_data = cursor.fetchall() + column_names, requires_lob_coercion = self._resolve_row_metadata(cursor.description) + data, column_names = collect_sync_rows( + cast("list[Any] | None", fetched_data), + cursor.description, + self.driver_features, + column_names=column_names, + requires_lob_coercion=requires_lob_coercion, + ) + return self.build_cached_select_result( + sql, + params, + cached, + selected_data=data, + column_names=column_names, + data_row_count=len(data), + row_format="tuple", + compiled_sql=wire_sql, + ) + affected_rows = resolve_rowcount(cursor) + return cast("SQLResult", DMLResult(cached.operation_type, affected_rows)) + def dispatch_execute_many(self, cursor: Any, statement: "SQL") -> "ExecutionResult": """Execute SQL with multiple parameter sets using Oracle batch processing. @@ -525,17 +572,7 @@ def dispatch_execute_many(self, cursor: Any, statement: "SQL") -> "ExecutionResu ) def dispatch_execute_script(self, cursor: Any, statement: "SQL") -> "ExecutionResult": - """Execute SQL script with statement splitting and parameter handling. - - Parameters are embedded as static values for script execution compatibility. - - Args: - cursor: Oracle cursor object - statement: SQL script statement to execute - - Returns: - Execution result containing statement count and success information - """ + """Execute SQL script with statement splitting and parameter handling.""" sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) prepared_parameters = coerce_large_parameters_sync( self.connection, @@ -546,18 +583,42 @@ def dispatch_execute_script(self, cursor: Any, statement: "SQL") -> "ExecutionRe raw_byte_limit=self.driver_features.get("oracle_raw_byte_limit", 2000), version_cache=self._oracle_version_cache, ) - prepared_parameters = cast("list[Any] | tuple[Any, ...] | dict[Any, Any] | None", prepared_parameters) statements = self.split_script_statements(sql, statement.statement_config, strip_trailing_semicolon=True) successful_count = 0 last_cursor = cursor - - for stmt in statements: - cursor.execute(stmt, _filter_script_parameters(stmt, prepared_parameters)) + pos_offset = 0 + total_statements = len(statements) + runtime = self._observability + has_observers = runtime is not None and runtime.has_statement_observers + + for idx, stmt in enumerate(statements): + stmt_params, pos_offset = _slice_script_parameters(stmt, prepared_parameters, pos_offset) + stmt_started = perf_counter() if has_observers else 0.0 + cursor.execute(stmt, stmt_params) successful_count += 1 + if has_observers: + stmt_duration = perf_counter() - stmt_started + single_stmt = self._sub_statement(stmt, stmt_params) + try: + stmt_operation = single_stmt.operation_type + finally: + self._release_pooled_statement(single_stmt) + self._emit_script_statement_event( + sql=stmt, + parameters=stmt_params, + operation=stmt_operation, + rows_affected=resolve_rowcount(cursor), + duration_s=stmt_duration, + started_at=stmt_started, + profile=None, + script_ordinal=idx, + script_statement_count=total_statements, + ) + self._pending_statement_profile = None return self.create_execution_result( - last_cursor, statement_count=len(statements), successful_statements=successful_count, is_script_result=True + last_cursor, statement_count=total_statements, successful_statements=successful_count, is_script_result=True ) def begin(self) -> None: @@ -579,7 +640,8 @@ def commit(self) -> None: self.connection.commit() except OracleError as e: raise create_mapped_exception(e) from e - self._transaction_active = False + finally: + self._transaction_active = False def rollback(self) -> None: """Rollback the current transaction. @@ -747,12 +809,22 @@ def select_to_arrow( ) prepared_parameters = cast("list[Any] | tuple[Any, ...] | dict[Any, Any] | None", prepared_parameters) + exc_handler = self.handle_database_exceptions() try: self._apply_session_tracing(self.connection) if return_format in {"batches", "reader"} and supports_df_batches(self.connection): - record_batches, batch_schema = self._fetch_arrow_record_batches( - sql, prepared_parameters, batch_size, arrow_schema + batches_result = cast( + "tuple[list[ArrowRecordBatch], ArrowSchema | None]", + self._run_with_exception_handler( + exc_handler, + self._fetch_arrow_record_batches, + sql, + prepared_parameters, + batch_size, + arrow_schema, + ), ) + record_batches, batch_schema = batches_result rows_affected = sum(batch.num_rows for batch in record_batches) if batch_schema is None: batch_schema = pa.schema([]) @@ -763,7 +835,9 @@ def select_to_arrow( arrow_data = record_batches return create_arrow_result(prepared_statement, arrow_data, rows_affected=rows_affected) - oracle_df = self._execute_arrow_dataframe(sql, prepared_parameters, batch_size) + oracle_df = self._run_with_exception_handler( + exc_handler, self._execute_arrow_dataframe, sql, prepared_parameters, batch_size + ) except AttributeError as exc: if native_only: msg = "Oracle native Arrow support is not available for this connection." @@ -776,6 +850,7 @@ def select_to_arrow( batch_size=batch_size, arrow_schema=arrow_schema, ) + self._check_pending_exception(exc_handler) arrow_table = coerce_oracle_df_to_arrow_table(oracle_df, pa) column_names = normalize_column_names(arrow_table.column_names, self.driver_features) @@ -883,14 +958,16 @@ def load_from_arrow( if overwrite: statement = build_truncate_statement(table) exc_handler = self.handle_database_exceptions() - with exc_handler: + + def _truncate() -> None: cursor = self.connection.cursor() try: cursor.execute(statement) finally: cursor.close() - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + + self._run_with_exception_handler(exc_handler, _truncate) + self._check_pending_exception(exc_handler) use_direct_path = ( self.driver_features.get("enable_direct_path_load", True) is not False and supports_direct_path_load(self.connection) @@ -900,34 +977,37 @@ def load_from_arrow( if use_direct_path: schema_name, table_name = _resolve_direct_path_target(self.connection, table) exc_handler = self.handle_database_exceptions() - with exc_handler: - self.connection.direct_path_load( - schema_name=schema_name, - table_name=table_name, - column_names=list(arrow_table.column_names), - data=arrow_table, - ) - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + self._run_with_exception_handler( + exc_handler, + self.connection.direct_path_load, + schema_name=schema_name, + table_name=table_name, + column_names=list(arrow_table.column_names), + data=arrow_table, + ) + self._check_pending_exception(exc_handler) else: columns = list(arrow_table.column_names) statement = build_insert_statement(table, columns) exc_handler = self.handle_database_exceptions() - with self.with_cursor(self.connection) as cursor, exc_handler: - for batch in arrow_table.to_batches(max_chunksize=5000): - batch_rows = [tuple(row[col] for col in columns) for row in batch.to_pylist()] - coerced_rows = coerce_many_parameters_sync( - self.connection, - batch_rows, - clob_type=DB_TYPE_CLOB, - blob_type=DB_TYPE_BLOB, - varchar2_byte_limit=self.driver_features.get("oracle_varchar2_byte_limit", 4000), - raw_byte_limit=self.driver_features.get("oracle_raw_byte_limit", 2000), - version_cache=self._oracle_version_cache, - ) - cursor.executemany(statement, coerced_rows) - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + + def _insert_batches() -> None: + with self.with_cursor(self.connection) as cursor: + for batch in arrow_table.to_batches(max_chunksize=5000): + batch_rows = [tuple(row[col] for col in columns) for row in batch.to_pylist()] + coerced_rows = coerce_many_parameters_sync( + self.connection, + batch_rows, + clob_type=DB_TYPE_CLOB, + blob_type=DB_TYPE_BLOB, + varchar2_byte_limit=self.driver_features.get("oracle_varchar2_byte_limit", 4000), + raw_byte_limit=self.driver_features.get("oracle_raw_byte_limit", 2000), + version_cache=self._oracle_version_cache, + ) + cursor.executemany(statement, coerced_rows) + + self._run_with_exception_handler(exc_handler, _insert_batches) + self._check_pending_exception(exc_handler) telemetry_payload = self._ingest_telemetry(arrow_table) telemetry_payload["destination"] = table self._attach_partition_telemetry(telemetry_payload, partitioner) @@ -974,13 +1054,17 @@ def load_from_storage( truncate_sql = build_truncate_statement(table) if overwrite else None rows_processed = 0 exc_handler = self.handle_database_exceptions() - with self.with_cursor(self.connection) as cursor, exc_handler: - if truncate_sql is not None: - cursor.execute(truncate_sql) - cursor.execute(plsql, binds) - rows_processed = resolve_rowcount(cursor) - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + + def _run_dbms_cloud() -> None: + nonlocal rows_processed + with self.with_cursor(self.connection) as cursor: + if truncate_sql is not None: + cursor.execute(truncate_sql) + cursor.execute(plsql, binds) + rows_processed = resolve_rowcount(cursor) + + self._run_with_exception_handler(exc_handler, _run_dbms_cloud) + self._check_pending_exception(exc_handler) telemetry_payload: StorageTelemetry = { "destination": table, "backend": "dbms_cloud", @@ -1295,6 +1379,47 @@ async def dispatch_execute(self, cursor: Any, statement: "SQL") -> "ExecutionRes affected_rows = resolve_rowcount(cursor) return self.create_execution_result(cursor, rowcount_override=affected_rows) + async def _execute_cache_hit_with_cursor( + self, sql: str, params: "tuple[Any, ...] | list[Any] | dict[str, Any]", cached: "CachedQuery" + ) -> "SQLResult": + """Execute an async cached query while owning only the cursor context.""" + async with self.with_cursor(self.connection) as cursor: + wire_sql = self._format_wire_sql(cached.compiled_sql) + returns_rows = cached.operation_profile.returns_rows + fetch_kwargs = build_fetch_kwargs(self.driver_features) + parameters = await coerce_large_parameters_async( + self.connection, + params, + clob_type=DB_TYPE_CLOB, + blob_type=DB_TYPE_BLOB, + varchar2_byte_limit=self.driver_features.get("oracle_varchar2_byte_limit", 4000), + raw_byte_limit=self.driver_features.get("oracle_raw_byte_limit", 2000), + version_cache=self._oracle_version_cache, + ) + await cast("Any", cursor).execute(wire_sql, parameters or {}, **fetch_kwargs) + if returns_rows and cursor.description is not None: + fetched_data = await cursor.fetchall() + column_names, requires_lob_coercion = self._resolve_row_metadata(cursor.description) + data, column_names = await collect_async_rows( + cast("list[Any] | None", fetched_data), + cursor.description, + self.driver_features, + column_names=column_names, + requires_lob_coercion=requires_lob_coercion, + ) + return self.build_cached_select_result( + sql, + params, + cached, + selected_data=data, + column_names=column_names, + data_row_count=len(data), + row_format="tuple", + compiled_sql=wire_sql, + ) + affected_rows = resolve_rowcount(cursor) + return cast("SQLResult", DMLResult(cached.operation_type, affected_rows)) + async def dispatch_execute_many(self, cursor: Any, statement: "SQL") -> "ExecutionResult": """Execute SQL with multiple parameter sets using Oracle batch processing. @@ -1346,17 +1471,7 @@ async def dispatch_execute_many(self, cursor: Any, statement: "SQL") -> "Executi ) async def dispatch_execute_script(self, cursor: Any, statement: "SQL") -> "ExecutionResult": - """Execute SQL script with statement splitting and parameter handling. - - Parameters are embedded as static values for script execution compatibility. - - Args: - cursor: Oracle cursor object - statement: SQL script statement to execute - - Returns: - Execution result containing statement count and success information - """ + """Execute async SQL script with statement splitting and parameter handling.""" sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) prepared_parameters = await coerce_large_parameters_async( self.connection, @@ -1367,18 +1482,42 @@ async def dispatch_execute_script(self, cursor: Any, statement: "SQL") -> "Execu raw_byte_limit=self.driver_features.get("oracle_raw_byte_limit", 2000), version_cache=self._oracle_version_cache, ) - prepared_parameters = cast("list[Any] | tuple[Any, ...] | dict[Any, Any] | None", prepared_parameters) statements = self.split_script_statements(sql, statement.statement_config, strip_trailing_semicolon=True) successful_count = 0 last_cursor = cursor - - for stmt in statements: - await cursor.execute(stmt, _filter_script_parameters(stmt, prepared_parameters)) + pos_offset = 0 + total_statements = len(statements) + runtime = self._observability + has_observers = runtime is not None and runtime.has_statement_observers + + for idx, stmt in enumerate(statements): + stmt_params, pos_offset = _slice_script_parameters(stmt, prepared_parameters, pos_offset) + stmt_started = perf_counter() if has_observers else 0.0 + await cursor.execute(stmt, stmt_params) successful_count += 1 + if has_observers: + stmt_duration = perf_counter() - stmt_started + single_stmt = self._sub_statement(stmt, stmt_params) + try: + stmt_operation = single_stmt.operation_type + finally: + self._release_pooled_statement(single_stmt) + self._emit_script_statement_event( + sql=stmt, + parameters=stmt_params, + operation=stmt_operation, + rows_affected=resolve_rowcount(cursor), + duration_s=stmt_duration, + started_at=stmt_started, + profile=None, + script_ordinal=idx, + script_statement_count=total_statements, + ) + self._pending_statement_profile = None return self.create_execution_result( - last_cursor, statement_count=len(statements), successful_statements=successful_count, is_script_result=True + last_cursor, statement_count=total_statements, successful_statements=successful_count, is_script_result=True ) async def begin(self) -> None: @@ -1400,7 +1539,8 @@ async def commit(self) -> None: await self.connection.commit() except OracleError as e: raise create_mapped_exception(e) from e - self._transaction_active = False + finally: + self._transaction_active = False async def rollback(self) -> None: """Rollback the current transaction. @@ -1571,12 +1711,22 @@ async def select_to_arrow( ) prepared_parameters = cast("list[Any] | tuple[Any, ...] | dict[Any, Any] | None", prepared_parameters) + exc_handler = self.handle_database_exceptions() try: self._apply_session_tracing(self.connection) if return_format in {"batches", "reader"} and supports_df_batches(self.connection): - record_batches, batch_schema = await self._fetch_arrow_record_batches( - sql, prepared_parameters, batch_size, arrow_schema + batches_result = cast( + "tuple[list[ArrowRecordBatch], ArrowSchema | None]", + await self._run_with_exception_handler( + exc_handler, + self._fetch_arrow_record_batches, + sql, + prepared_parameters, + batch_size, + arrow_schema, + ), ) + record_batches, batch_schema = batches_result rows_affected = sum(batch.num_rows for batch in record_batches) if batch_schema is None: batch_schema = pa.schema([]) @@ -1587,7 +1737,9 @@ async def select_to_arrow( arrow_data = record_batches return create_arrow_result(prepared_statement, arrow_data, rows_affected=rows_affected) - oracle_df = await self._execute_arrow_dataframe(sql, prepared_parameters, batch_size) + oracle_df = await self._run_with_exception_handler( + exc_handler, self._execute_arrow_dataframe, sql, prepared_parameters, batch_size + ) except AttributeError as exc: if native_only: msg = "Oracle native Arrow support is not available for this connection." @@ -1600,6 +1752,7 @@ async def select_to_arrow( batch_size=batch_size, arrow_schema=arrow_schema, ) + self._check_pending_exception(exc_handler) arrow_table = coerce_oracle_df_to_arrow_table(oracle_df, pa) column_names = normalize_column_names(arrow_table.column_names, self.driver_features) @@ -1711,10 +1864,8 @@ async def load_from_arrow( if overwrite: statement = build_truncate_statement(table) exc_handler = self.handle_database_exceptions() - async with exc_handler: - await self.connection.execute(statement) - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + await self._run_with_exception_handler(exc_handler, self.connection.execute, statement) + self._check_pending_exception(exc_handler) use_direct_path = ( self.driver_features.get("enable_direct_path_load", True) is not False and supports_direct_path_load(self.connection) @@ -1724,34 +1875,37 @@ async def load_from_arrow( if use_direct_path: schema_name, table_name = _resolve_direct_path_target(self.connection, table) exc_handler = self.handle_database_exceptions() - async with exc_handler: - await self.connection.direct_path_load( - schema_name=schema_name, - table_name=table_name, - column_names=list(arrow_table.column_names), - data=arrow_table, - ) - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + await self._run_with_exception_handler( + exc_handler, + self.connection.direct_path_load, + schema_name=schema_name, + table_name=table_name, + column_names=list(arrow_table.column_names), + data=arrow_table, + ) + self._check_pending_exception(exc_handler) else: columns = list(arrow_table.column_names) statement = build_insert_statement(table, columns) exc_handler = self.handle_database_exceptions() - async with self.with_cursor(self.connection) as cursor, exc_handler: - for batch in arrow_table.to_batches(max_chunksize=5000): - batch_rows = [tuple(row[col] for col in columns) for row in batch.to_pylist()] - coerced_rows = await coerce_many_parameters_async( - self.connection, - batch_rows, - clob_type=DB_TYPE_CLOB, - blob_type=DB_TYPE_BLOB, - varchar2_byte_limit=self.driver_features.get("oracle_varchar2_byte_limit", 4000), - raw_byte_limit=self.driver_features.get("oracle_raw_byte_limit", 2000), - version_cache=self._oracle_version_cache, - ) - await cursor.executemany(statement, coerced_rows) - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + + async def _insert_batches() -> None: + async with self.with_cursor(self.connection) as cursor: + for batch in arrow_table.to_batches(max_chunksize=5000): + batch_rows = [tuple(row[col] for col in columns) for row in batch.to_pylist()] + coerced_rows = await coerce_many_parameters_async( + self.connection, + batch_rows, + clob_type=DB_TYPE_CLOB, + blob_type=DB_TYPE_BLOB, + varchar2_byte_limit=self.driver_features.get("oracle_varchar2_byte_limit", 4000), + raw_byte_limit=self.driver_features.get("oracle_raw_byte_limit", 2000), + version_cache=self._oracle_version_cache, + ) + await cursor.executemany(statement, coerced_rows) + + await self._run_with_exception_handler(exc_handler, _insert_batches) + self._check_pending_exception(exc_handler) telemetry_payload = self._ingest_telemetry(arrow_table) telemetry_payload["destination"] = table self._attach_partition_telemetry(telemetry_payload, partitioner) @@ -1798,13 +1952,17 @@ async def load_from_storage( truncate_sql = build_truncate_statement(table) if overwrite else None rows_processed = 0 exc_handler = self.handle_database_exceptions() - async with self.with_cursor(self.connection) as cursor, exc_handler: - if truncate_sql is not None: - await cursor.execute(truncate_sql) - await cursor.execute(plsql, binds) - rows_processed = resolve_rowcount(cursor) - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + + async def _run_dbms_cloud_async() -> None: + nonlocal rows_processed + async with self.with_cursor(self.connection) as cursor: + if truncate_sql is not None: + await cursor.execute(truncate_sql) + await cursor.execute(plsql, binds) + rows_processed = resolve_rowcount(cursor) + + await self._run_with_exception_handler(exc_handler, _run_dbms_cloud_async) + self._check_pending_exception(exc_handler) telemetry_payload: StorageTelemetry = { "destination": table, "backend": "dbms_cloud", @@ -2267,28 +2425,75 @@ def _retry_arrow_without_fetch_lobs(exc: TypeError, fetch_kwargs: "dict[str, obj return retry_kwargs -def _statement_contains_bind(statement: str, name: str) -> bool: - """Return whether a SQL statement references a named Oracle bind placeholder.""" - lowered_sql = statement.lower() - target = f":{name.lower()}" - start = 0 - while True: - idx = lowered_sql.find(target, start) - if idx == -1: - return False - prev_ok = idx == 0 or lowered_sql[idx - 1] != ":" - end_idx = idx + len(target) - next_ch = lowered_sql[end_idx] if end_idx < len(lowered_sql) else "" - next_ok = not (next_ch.isalnum() or next_ch in {"_", "$", "#"}) - if prev_ok and next_ok: - return True - start = idx + 1 +def _extract_statement_binds(statement: str) -> "tuple[set[str], int]": + """Extract named bind names (lowercased) and count of positional binds outside quotes and comments.""" + named_binds: set[str] = set() + positional_count = 0 + n = len(statement) + i = 0 + while i < n: + ch = statement[i] + if ch == "-" and i + 1 < n and statement[i + 1] == "-": + end = statement.find("\n", i + 2) + i = n if end == -1 else end + 1 + continue + if ch == "/" and i + 1 < n and statement[i + 1] == "*": + end = statement.find("*/", i + 2) + i = n if end == -1 else end + 2 + continue + if ch == "'": + end = i + 1 + while end < n: + if statement[end] == "'": + if end + 1 < n and statement[end + 1] == "'": + end += 2 + continue + end += 1 + break + end += 1 + i = end + continue + if ch == '"': + end = i + 1 + while end < n: + if statement[end] == '"': + if end + 1 < n and statement[end + 1] == '"': + end += 2 + continue + end += 1 + break + end += 1 + i = end + continue + if ch == ":" and (i == 0 or statement[i - 1] != ":"): + start = i + 1 + if start < n and statement[start] != ":": + end = start + while end < n and (statement[end].isalnum() or statement[end] in {"_", "$", "#"}): + end += 1 + if end > start: + token = statement[start:end] + if token.isdigit(): + positional_count += 1 + else: + named_binds.add(token.lower()) + i = end + continue + i += 1 + return named_binds, positional_count -def _filter_script_parameters(statement: str, parameters: Any) -> Any: - """Filter dictionary script parameters to the bind placeholders present in statement.""" +def _slice_script_parameters(statement: str, parameters: Any, pos_offset: int) -> "tuple[Any, int]": + """Filter dictionary binds or slice positional binds for a single statement.""" if not parameters: - return {} + return {}, pos_offset + named_binds, pos_count = _extract_statement_binds(statement) if isinstance(parameters, dict): - return {key: val for key, val in parameters.items() if _statement_contains_bind(statement, str(key))} - return parameters + filtered = {k: v for k, v in parameters.items() if str(k).lower() in named_binds} + return filtered, pos_offset + if isinstance(parameters, (list, tuple)): + if pos_count > 0: + sliced = list(parameters[pos_offset : pos_offset + pos_count]) + return sliced, pos_offset + pos_count + return [], pos_offset + return parameters, pos_offset diff --git a/sqlspec/adapters/oracledb/type_converter.py b/sqlspec/adapters/oracledb/type_converter.py index 1011a3c60..a971cefcc 100644 --- a/sqlspec/adapters/oracledb/type_converter.py +++ b/sqlspec/adapters/oracledb/type_converter.py @@ -12,13 +12,11 @@ numpy_converter_out, # pyright: ignore[reportPrivateUsage] ) from sqlspec.typing import NUMPY_INSTALLED +from sqlspec.utils.module_loader import import_optional from sqlspec.utils.sync_tools import ensure_async_ from sqlspec.utils.type_guards import is_async_readable, is_readable -try: - import numpy as np -except ImportError: - np = None # type: ignore[assignment] +np: Any = import_optional("numpy") __all__ = ("OracleOutputConverter",) diff --git a/sqlspec/utils/arrow_helpers.py b/sqlspec/utils/arrow_helpers.py index 8627be80e..742bada9d 100644 --- a/sqlspec/utils/arrow_helpers.py +++ b/sqlspec/utils/arrow_helpers.py @@ -16,7 +16,13 @@ from sqlspec.data_dictionary._registry import get_dialect_config from sqlspec.exceptions import ImproperConfigurationError, SQLBuilderError from sqlspec.utils.dispatch import TypeDispatcher -from sqlspec.utils.module_loader import ensure_pandas, ensure_polars, ensure_pyarrow, import_optional_attr +from sqlspec.utils.module_loader import ( + ensure_pandas, + ensure_polars, + ensure_pyarrow, + import_optional, + import_optional_attr, +) from sqlspec.utils.type_guards import has_arrow_table_stats, has_get_data from sqlspec.utils.uuids import uuid_from_bytes @@ -46,12 +52,16 @@ "cast_arrow_table_schema", "coerce_arrow_batches", "coerce_arrow_table", + "coerce_oracle_df_to_arrow_batches", + "coerce_oracle_df_to_arrow_table", "convert_dict_to_arrow", "convert_dict_to_arrow_with_schema", "encode_arrow_batch_pg_binary", "ensure_arrow_table", "iter_pgcopy_binary_chunks", "iter_pgcopy_binary_chunks_async", + "oracle_df_batches_to_arrow_reader", + "oracle_df_to_arrow_table", "records_to_arrow_table", "sql_type_for_category", "sql_type_from_arrow_type", @@ -1455,3 +1465,68 @@ async def __anext__(self) -> bytes: return next(self._iterator) except StopIteration as exc: raise StopAsyncIteration from exc + + +def oracle_df_to_arrow_table(oracle_df: Any, pa_module: Any = None, arrow_schema: Any = None) -> "ArrowTable": + """Convert an Oracle DataFrame to a PyArrow Table using zero-copy PyCapsule when supported.""" + if pa_module is None: + ensure_pyarrow() + pa_module = import_optional("pyarrow") + + if hasattr(oracle_df, "__arrow_c_stream__"): + reader = ( + pa_module.RecordBatchReader.from_stream(oracle_df, schema=arrow_schema) + if arrow_schema is not None + else pa_module.RecordBatchReader.from_stream(oracle_df) + ) + return cast("ArrowTable", reader.read_all()) + table = pa_module.table(oracle_df) + if arrow_schema is not None: + table = table.cast(arrow_schema) + return cast("ArrowTable", table) + + +def coerce_oracle_df_to_arrow_table(oracle_df: Any, pa_module: Any = None, arrow_schema: Any = None) -> "ArrowTable": + """Alias for :func:`oracle_df_to_arrow_table`.""" + return oracle_df_to_arrow_table(oracle_df, pa_module=pa_module, arrow_schema=arrow_schema) + + +def coerce_oracle_df_to_arrow_batches( + oracle_df: Any, pa_module: Any = None, arrow_schema: Any = None +) -> "list[ArrowRecordBatch]": + """Convert an Oracle DataFrame batch into PyArrow RecordBatches via PyCapsule when supported.""" + if pa_module is None: + ensure_pyarrow() + pa_module = import_optional("pyarrow") + + if hasattr(oracle_df, "__arrow_c_stream__"): + reader = ( + pa_module.RecordBatchReader.from_stream(oracle_df, schema=arrow_schema) + if arrow_schema is not None + else pa_module.RecordBatchReader.from_stream(oracle_df) + ) + return cast("list[ArrowRecordBatch]", list(reader)) + table = pa_module.table(oracle_df) + if arrow_schema is not None: + table = table.cast(arrow_schema) + return cast("list[ArrowRecordBatch]", table.to_batches()) + + +def oracle_df_batches_to_arrow_reader( + batches: "Iterable[Any]", schema: Any = None, pa_module: Any = None +) -> "ArrowRecordBatchReader": + """Convert an iterable of Oracle DataFrame batches into a PyArrow RecordBatchReader.""" + if pa_module is None: + ensure_pyarrow() + pa_module = import_optional("pyarrow") + + collected_batches: list[Any] = [] + for df_batch in batches: + collected_batches.extend(coerce_oracle_df_to_arrow_batches(df_batch, pa_module=pa_module, arrow_schema=schema)) + if collected_batches: + resolved_schema = schema if schema is not None else collected_batches[0].schema + return cast( + "ArrowRecordBatchReader", pa_module.RecordBatchReader.from_batches(resolved_schema, collected_batches) + ) + resolved_schema = schema if schema is not None else pa_module.schema([]) + return cast("ArrowRecordBatchReader", pa_module.RecordBatchReader.from_batches(resolved_schema, [])) diff --git a/tests/unit/adapters/test_oracledb/test_arrow_batches.py b/tests/unit/adapters/test_oracledb/test_arrow_batches.py index 73545af34..022a5dcdc 100644 --- a/tests/unit/adapters/test_oracledb/test_arrow_batches.py +++ b/tests/unit/adapters/test_oracledb/test_arrow_batches.py @@ -6,12 +6,13 @@ import pyarrow as pa import pytest -from sqlspec.adapters.oracledb.core import ( - build_arrow_fetch_kwargs, +from sqlspec.adapters.oracledb.core import build_arrow_fetch_kwargs +from sqlspec.adapters.oracledb.driver import OracleAsyncDriver, OracleSyncDriver +from sqlspec.utils.arrow_helpers import ( coerce_oracle_df_to_arrow_batches, coerce_oracle_df_to_arrow_table, + oracle_df_to_arrow_table, ) -from sqlspec.adapters.oracledb.driver import OracleAsyncDriver, OracleSyncDriver if TYPE_CHECKING: from sqlspec.adapters.oracledb._typing import OracleAsyncConnection, OracleSyncConnection @@ -188,4 +189,26 @@ async def fetch_df_all(self, **_kwargs: object) -> _StreamOnlyOracleDataFrame: ) res_table = await driver.select_to_arrow("SELECT id, value FROM example", return_format="table") assert res_table.get_data().to_pydict() == {"id": [1, 2], "value": [10, 20]} - assert stream_df.stream_calls == 1 + + +def test_oracle_df_to_arrow_table_lazy_pyarrow_resolution() -> None: + """oracle_df_to_arrow_table lazily resolves pyarrow when pa_module is None.""" + raw_table = pa.table({"ID": [1, 2], "VALUE": [10, 20]}) + stream_df = _StreamOnlyOracleDataFrame(raw_table) + table = oracle_df_to_arrow_table(stream_df, pa_module=None) + assert table.num_rows == 2 + + +def test_select_to_arrow_maps_native_database_error() -> None: + """select_to_arrow routes oracledb.DatabaseError through exception handler.""" + import oracledb + + from sqlspec.exceptions import SQLSpecError + + class _FailingOracleConn: + def fetch_df_all(self, **_kwargs: object) -> None: + raise oracledb.DatabaseError("ORA-00942: table or view does not exist") + + driver = OracleSyncDriver(cast("OracleSyncConnection", _FailingOracleConn())) + with pytest.raises(SQLSpecError): + driver.select_to_arrow("SELECT * FROM non_existent", return_format="table") diff --git a/tests/unit/adapters/test_oracledb/test_direct_path_load.py b/tests/unit/adapters/test_oracledb/test_direct_path_load.py index 16fe1c2a5..82e951d84 100644 --- a/tests/unit/adapters/test_oracledb/test_direct_path_load.py +++ b/tests/unit/adapters/test_oracledb/test_direct_path_load.py @@ -329,3 +329,24 @@ def test_convertible_scalar_columns_keep_the_native_path(arrow_type: "pa.DataTyp assert len(conn.dpl_calls) == 1 assert conn._cursor.executemany_calls == [] + + +def test_direct_path_load_failure_preserves_cause() -> None: + """direct_path_load failure preserves the underlying cause on the raised exception.""" + import oracledb + + from sqlspec.exceptions import SQLSpecError + + native_err = oracledb.DatabaseError("ORA-00001: unique constraint violated") + + class _FailingDPLConn(_DPLConnection): + def direct_path_load(self, **kwargs: Any) -> None: + raise native_err + + conn = _FailingDPLConn(thin=True, username="SCOTT") + driver = OracleSyncDriver(cast("Any", conn), driver_features={"storage_capabilities": _CAPS}) + + with pytest.raises(SQLSpecError) as exc_info: + driver.load_from_arrow("MYTAB", _arrow()) + + assert exc_info.value.__cause__ is native_err diff --git a/tests/unit/adapters/test_oracledb/test_driver.py b/tests/unit/adapters/test_oracledb/test_driver.py index 24c928eb3..a8061a042 100644 --- a/tests/unit/adapters/test_oracledb/test_driver.py +++ b/tests/unit/adapters/test_oracledb/test_driver.py @@ -8,11 +8,13 @@ import pytest from sqlspec.adapters.oracledb import build_dbms_cloud_copy_data_sql as exported_build_dbms_cloud_copy_data_sql +from sqlspec.adapters.oracledb._typing import OracleSyncConnection from sqlspec.adapters.oracledb.core import ( ORACLE_ACTION_MAX_BYTES, ORACLE_CLIENT_IDENTIFIER_MAX_BYTES, ORACLE_ECID_MAX_BYTES, ORACLE_MODULE_MAX_BYTES, + OracleSyncStreamSource, apply_driver_features, apply_oracle_session_trace_attributes, build_dbms_cloud_copy_data_sql, @@ -537,3 +539,59 @@ def close(self) -> None: assert len(async_executed) == 1 assert "DBMS_CLOUD.COPY_DATA" not in async_executed[0][0] assert "INSERT INTO" in async_executed[0][0] + + +def test_oracle_sync_stream_source_runs_with_exception_handler() -> None: + """Sync stream source start and fetch_chunk route execution through _run_with_exception_handler.""" + cursor = MagicMock() + cursor.description = [("id", int, None, None, None, None, None)] + cursor.fetchmany.return_value = [(1,)] + conn = MagicMock() + conn.cursor.return_value = cursor + driver = OracleSyncDriver(cast("OracleSyncConnection", conn)) + + source = OracleSyncStreamSource(driver, "SELECT 1 FROM dual", None, 100) + source.start() + rows = source.fetch_chunk() + assert len(rows) == 1 + assert rows[0] == {"id": 1} + + +def test_oracle_statement_cache_hit_executes_with_fetch_kwargs_and_coercion() -> None: + """Fast-path statement cache hit passes fetch kwargs and coerces parameters.""" + cursor = MagicMock() + cursor.description = [("content", "CLOB", None, None, None, None, None)] + cursor.fetchall.return_value = [("clob text",)] + conn = MagicMock() + conn.cursor.return_value = cursor + driver = OracleSyncDriver(cast("OracleSyncConnection", conn)) + + res1 = driver.execute("SELECT content FROM articles WHERE id = :id", {"id": "1"}) + assert res1.get_data() == [{"content": "clob text"}] + + res2 = driver.execute("SELECT content FROM articles WHERE id = :id", {"id": "1"}) + assert res2.get_data() == [{"content": "clob text"}] + + +def test_oracle_dispatch_execute_script_ignores_colons_in_literals_and_slices_positional() -> None: + """dispatch_execute_script ignores colons in literals and slices positional parameters.""" + from sqlspec import SQL + + executed_statements: list[tuple[str, Any]] = [] + + class _MockCursor: + def execute(self, stmt: str, params: Any = None) -> None: + executed_statements.append((stmt, params)) + + conn = MagicMock() + cursor = _MockCursor() + conn.cursor.return_value = cursor + driver = OracleSyncDriver(cast("OracleSyncConnection", conn), default_statement_config) + + script_sql = "INSERT INTO t (url) VALUES ('http://example.com:8080'); INSERT INTO t (id, name) VALUES (:1, :2)" + stmt = SQL(script_sql, [101, "Alice"], statement_config=default_statement_config) + res = driver.dispatch_execute_script(cursor, stmt) + assert res.statement_count == 2 + assert len(executed_statements) == 2 + assert executed_statements[0][1] == [] + assert executed_statements[1][1] == [101, "Alice"] diff --git a/tests/unit/adapters/test_oracledb/test_fetch_tuning.py b/tests/unit/adapters/test_oracledb/test_fetch_tuning.py index 324654c84..4496ce04e 100644 --- a/tests/unit/adapters/test_oracledb/test_fetch_tuning.py +++ b/tests/unit/adapters/test_oracledb/test_fetch_tuning.py @@ -167,3 +167,14 @@ async def test_oracle_async_stream_fetch_lobs_call_option_overrides_default() -> rows: list[dict[str, Any]] = [row async for row in stream] assert rows == [] assert raw_cursor.execute_kwargs["fetch_lobs"] is True + + +def test_resolve_rowcount_negative_rowcount_returns_zero() -> None: + """resolve_rowcount returns 0 when cursor.rowcount is negative.""" + from unittest.mock import Mock + + from sqlspec.adapters.oracledb.core import resolve_rowcount + + cursor = Mock() + cursor.rowcount = -1 + assert resolve_rowcount(cursor) == 0 diff --git a/tests/unit/adapters/test_oracledb/test_lob_coercion.py b/tests/unit/adapters/test_oracledb/test_lob_coercion.py index 09aeda895..63826b455 100644 --- a/tests/unit/adapters/test_oracledb/test_lob_coercion.py +++ b/tests/unit/adapters/test_oracledb/test_lob_coercion.py @@ -1127,3 +1127,17 @@ async def test_oracle_async_stream_source_closes_cursor_when_execute_raises() -> cursor.close.assert_called_once() assert source._cursor is None + + +def test_collect_sync_rows_preserves_lob_locators_when_fetch_lobs_true() -> None: + """collect_sync_rows preserves raw LOB locators when fetch_lobs=True.""" + from sqlspec.adapters.oracledb.core import collect_sync_rows + + raw_lob = MagicMock() + raw_lob.read = MagicMock(side_effect=AssertionError("read should not be called")) + description = [("payload", "CLOB", None, None, None, None, None)] + fetched = [(raw_lob,)] + rows, _cols = collect_sync_rows(fetched, description, {"fetch_lobs": True}) + assert len(rows) == 1 + assert rows[0][0] is raw_lob + raw_lob.read.assert_not_called() diff --git a/tests/unit/adapters/test_oracledb/test_stack_native_transaction.py b/tests/unit/adapters/test_oracledb/test_stack_native_transaction.py index a70038361..00cf17737 100644 --- a/tests/unit/adapters/test_oracledb/test_stack_native_transaction.py +++ b/tests/unit/adapters/test_oracledb/test_stack_native_transaction.py @@ -108,3 +108,21 @@ async def test_native_stack_inside_an_implicit_transaction_does_not_commit() -> assert connection.commit_count == 0 assert connection.rollback_count == 0 + + +@pytest.mark.anyio +async def test_native_stack_wraps_mapped_error() -> None: + """Stack execution errors wrap mapped SQLSpec exceptions, not raw driver errors.""" + import oracledb + + from sqlspec.exceptions import SQLSpecError + + native_err = oracledb.DatabaseError("ORA-00001: unique constraint violated") + results = [_FakeOpResult(native_err)] + connection = _FakePipelineConnection(results) + driver = OracleAsyncDriver(cast("Any", connection)) + + stack_results = await driver._execute_stack_native(_stack(), continue_on_error=True) + assert len(stack_results) == 1 + assert stack_results[0].error is not None + assert isinstance(stack_results[0].error.original_error, SQLSpecError) diff --git a/tests/unit/adapters/test_oracledb/test_transaction_state.py b/tests/unit/adapters/test_oracledb/test_transaction_state.py index 5f7078239..0662a6dbb 100644 --- a/tests/unit/adapters/test_oracledb/test_transaction_state.py +++ b/tests/unit/adapters/test_oracledb/test_transaction_state.py @@ -114,3 +114,63 @@ def test_execute_stack_does_not_commit_after_begin_without_dml() -> None: driver.commit() assert driver._connection_in_transaction() is False + + +def test_failed_commit_resets_transaction_active() -> None: + """A failed commit must reset _transaction_active in finally.""" + import oracledb + + from sqlspec.exceptions import SQLSpecError + + class _FailingCommitConn(_FakeSyncConnection): + def commit(self) -> None: + raise oracledb.DatabaseError("ORA-02091: transaction rolled back") + + driver = OracleSyncDriver(cast("Any", _FailingCommitConn())) + driver.begin() + assert driver._connection_in_transaction() is True + with pytest.raises(SQLSpecError): + driver.commit() + assert driver._connection_in_transaction() is False + + +def test_connection_context_rolls_back_on_block_exception() -> None: + """Connection context rolls back on unhandled block exceptions.""" + from unittest.mock import MagicMock + + from sqlspec.adapters.oracledb.config import OracleSyncConnectionContext + + conn = MagicMock() + config = MagicMock() + config.driver_features = {} + config.connection_instance.acquire.return_value = conn + ctx = OracleSyncConnectionContext(config) + + with pytest.raises(RuntimeError): + with ctx: + raise RuntimeError("boom") + + conn.rollback.assert_called_once() + + +@pytest.mark.anyio +async def test_async_connection_context_rolls_back_on_block_exception() -> None: + """Async connection context rolls back on unhandled block exceptions.""" + from unittest.mock import AsyncMock, MagicMock + + from sqlspec.adapters.oracledb.config import OracleAsyncConnectionContext + + conn = AsyncMock() + pool = MagicMock() + pool.acquire = AsyncMock(return_value=conn) + pool.release = AsyncMock() + config = MagicMock() + config.driver_features = {} + config.provide_pool = AsyncMock(return_value=pool) + ctx = OracleAsyncConnectionContext(config) + + with pytest.raises(RuntimeError): + async with ctx: + raise RuntimeError("boom") + + conn.rollback.assert_called_once() diff --git a/tests/unit/adapters/test_oracledb/test_vector_handlers.py b/tests/unit/adapters/test_oracledb/test_vector_handlers.py index 728d7bdde..7aeb7b9be 100644 --- a/tests/unit/adapters/test_oracledb/test_vector_handlers.py +++ b/tests/unit/adapters/test_oracledb/test_vector_handlers.py @@ -86,18 +86,18 @@ def test_input_handler_claims_tuple_of_float() -> None: assert list(packed) == [0.5, 0.25, 0.125] -def test_input_handler_packs_int_in_int8_range_as_int8() -> None: - """``list[int]`` with all values in [-128, 127] is packed as int8.""" +def test_input_handler_packs_int_sequence_as_float32() -> None: + """``list[int]`` is deterministically packed as float32.""" from sqlspec.adapters.oracledb._vector_handlers import _input_type_handler # pyright: ignore[reportPrivateUsage] cursor = _mock_cursor() cursor.var.return_value = "var-token" - _input_type_handler(cursor, [-128, 0, 64, 127], 1) + _input_type_handler(cursor, [0, 1, 0], 1) - packed = cursor.var.call_args.kwargs["inconverter"]([-128, 0, 64, 127]) - assert packed.typecode == "b" - assert list(packed) == [-128, 0, 64, 127] + packed = cursor.var.call_args.kwargs["inconverter"]([0, 1, 0]) + assert packed.typecode == "f" + assert list(packed) == [0.0, 1.0, 0.0] def test_input_handler_packs_int_outside_int8_range_as_float32() -> None: @@ -382,12 +382,12 @@ def test_pack_sequence_converter_float_list() -> None: def test_pack_sequence_converter_int8_range() -> None: - """_pack_sequence_converter packs small int lists as int8.""" + """_pack_sequence_converter packs small int lists as float32.""" result = _pack_sequence_converter([-128, 0, 127]) assert isinstance(result, array.array) - assert result.typecode == "b" - assert list(result) == [-128, 0, 127] + assert result.typecode == "f" + assert list(result) == [-128.0, 0.0, 127.0] def test_input_handler_rejects_heterogeneous_list_and_pre_23ai_and_disabled_flag() -> None: From 2a7bd63860bed78996e31a5c6dd03af82c681024 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 01:02:29 +0000 Subject: [PATCH 066/174] fix(oracledb): remediate ADK, Litestar, and events extension stores Enforce connection rollback across all mutating methods in Litestar and ADK stores, batch insert memory entries via executemany with duplicate key handling, use positional tuple parameter binding on Litestar hot lookup methods, and validate individual dot-separated table identifier segments. Retain Oracle AQ commit-at-dequeue for documented at-most-once delivery semantics, ensure json payloads in AQ messages deserialize correctly, consolidate duplicate LOB reading functions into core, and eliminate redundant decimal coercion logic. --- sqlspec/adapters/oracledb/adk/store.py | 285 +++++++------ sqlspec/adapters/oracledb/core.py | 18 + sqlspec/adapters/oracledb/events/backend.py | 6 + sqlspec/adapters/oracledb/events/store.py | 6 +- sqlspec/adapters/oracledb/litestar/store.py | 386 +++++++++--------- .../adapters/test_oracledb/test_driver.py | 74 ++++ .../test_oracledb/test_litestar_store.py | 25 +- .../test_oracledb/test_oracle_adk_store.py | 84 +++- 8 files changed, 563 insertions(+), 321 deletions(-) diff --git a/sqlspec/adapters/oracledb/adk/store.py b/sqlspec/adapters/oracledb/adk/store.py index c54d5aadd..bc43686b3 100644 --- a/sqlspec/adapters/oracledb/adk/store.py +++ b/sqlspec/adapters/oracledb/adk/store.py @@ -1,7 +1,7 @@ """Oracle ADK store for Google Agent Development Kit session/event storage.""" import array -from decimal import Decimal +import contextlib from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, NotRequired, cast from typing_extensions import TypedDict @@ -12,6 +12,8 @@ from sqlspec.adapters.oracledb.core import ( coerce_large_parameters_async, coerce_large_parameters_sync, + read_oracle_lob_async, + read_oracle_lob_sync, sanitize_oracle_text_query, ) from sqlspec.adapters.oracledb.data_dictionary import JSONStorageType, OracleVersionInfo, storage_type_from_version @@ -27,6 +29,7 @@ from sqlspec.extensions.adk.memory.store import BaseAsyncADKMemoryStore, BaseSyncADKMemoryStore from sqlspec.utils.logging import get_logger from sqlspec.utils.serializers import from_json, to_json +from sqlspec.utils.type_converters import build_nested_decimal_normalizer from sqlspec.utils.type_guards import is_async_readable, is_readable if TYPE_CHECKING: @@ -377,6 +380,16 @@ class OracleADKConfig(ADKConfig): """Raw Oracle table options for the ADK memory table.""" +def _rollback_sync(conn: Any) -> None: + with contextlib.suppress(Exception): + conn.rollback() + + +async def _rollback_async(conn: Any) -> None: + with contextlib.suppress(Exception): + await conn.rollback() + + def coerce_decimal_values(value: Any) -> Any: return _coerce_decimal_values(value) @@ -498,10 +511,14 @@ async def create_session( params = {"id": session_id, "app_name": app_name, "user_id": user_id, "state": state_data} async with self._config.provide_connection() as conn: - coerced_params = await coerce_large_parameters_async(conn, params) - with conn.cursor() as cursor: - await cursor.execute(sql, coerced_params) - await conn.commit() + try: + coerced_params = await coerce_large_parameters_async(conn, params) + with conn.cursor() as cursor: + await cursor.execute(sql, coerced_params) + await conn.commit() + except BaseException: + await _rollback_async(conn) + raise result = await self.get_session(app_name, user_id, session_id) if result is None: @@ -675,9 +692,13 @@ async def delete_session(self, app_name: str, user_id: str, session_id: str) -> sql = f"DELETE FROM {self._session_table} WHERE app_name = :app_name AND user_id = :user_id AND id = :id" async with self._config.provide_connection() as conn: - with conn.cursor() as cursor: - await cursor.execute(sql, {"app_name": app_name, "user_id": user_id, "id": session_id}) - await conn.commit() + try: + with conn.cursor() as cursor: + await cursor.execute(sql, {"app_name": app_name, "user_id": user_id, "id": session_id}) + await conn.commit() + except BaseException: + await _rollback_async(conn) + raise async def append_event(self, event_record: StoredEvent) -> None: """Append an event to a session. @@ -694,21 +715,25 @@ async def append_event(self, event_record: StoredEvent) -> None: """ async with self._config.provide_connection() as conn: - coerced_params = await coerce_large_parameters_async( - conn, - { - "id": event_record["id"], - "app_name": event_record["app_name"], - "user_id": event_record["user_id"], - "session_id": event_record["session_id"], - "invocation_id": event_record["invocation_id"], - "timestamp": event_record["timestamp"], - "event_data": await self._serialize_event_data(event_record["event_data"]), - }, - ) - with conn.cursor() as cursor: - await cursor.execute(sql, coerced_params) - await conn.commit() + try: + coerced_params = await coerce_large_parameters_async( + conn, + { + "id": event_record["id"], + "app_name": event_record["app_name"], + "user_id": event_record["user_id"], + "session_id": event_record["session_id"], + "invocation_id": event_record["invocation_id"], + "timestamp": event_record["timestamp"], + "event_data": await self._serialize_event_data(event_record["event_data"]), + }, + ) + with conn.cursor() as cursor: + await cursor.execute(sql, coerced_params) + await conn.commit() + except BaseException: + await _rollback_async(conn) + raise async def append_event_and_update_state( self, @@ -1690,9 +1715,14 @@ def delete_session(self, app_name: str, user_id: str, session_id: str) -> None: """ sql = f"DELETE FROM {self._session_table} WHERE app_name = :app_name AND user_id = :user_id AND id = :id" - with self._config.provide_connection() as conn, conn.cursor() as cursor: - cursor.execute(sql, {"app_name": app_name, "user_id": user_id, "id": session_id}) - conn.commit() + with self._config.provide_connection() as conn: + try: + with conn.cursor() as cursor: + cursor.execute(sql, {"app_name": app_name, "user_id": user_id, "id": session_id}) + conn.commit() + except BaseException: + _rollback_sync(conn) + raise def append_event(self, event_record: StoredEvent) -> None: """Append an event to a session.""" @@ -1704,21 +1734,26 @@ def append_event(self, event_record: StoredEvent) -> None: ) """ - with self._config.provide_connection() as conn, conn.cursor() as cursor: - params = coerce_large_parameters_sync( - conn, - { - "id": event_record["id"], - "app_name": event_record["app_name"], - "user_id": event_record["user_id"], - "session_id": event_record["session_id"], - "invocation_id": event_record["invocation_id"], - "timestamp": event_record["timestamp"], - "event_data": self._serialize_event_data(event_record["event_data"]), - }, - ) - cursor.execute(sql, params) - conn.commit() + with self._config.provide_connection() as conn: + try: + params = coerce_large_parameters_sync( + conn, + { + "id": event_record["id"], + "app_name": event_record["app_name"], + "user_id": event_record["user_id"], + "session_id": event_record["session_id"], + "invocation_id": event_record["invocation_id"], + "timestamp": event_record["timestamp"], + "event_data": self._serialize_event_data(event_record["event_data"]), + }, + ) + with conn.cursor() as cursor: + cursor.execute(sql, params) + conn.commit() + except BaseException: + _rollback_sync(conn) + raise def append_event_and_update_state( self, @@ -2449,37 +2484,48 @@ async def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: " ) """ - inserted_count = 0 - async with self._config.provide_connection() as conn: - with conn.cursor() as cursor: - for entry in entries: - content_json = await self._serialize_json_field(entry["content_json"]) - metadata_json = await self._serialize_json_field(entry["metadata_json"]) - embedding_val = entry.get("embedding") - bound_embedding = array.array("f", embedding_val) if embedding_val is not None else None - params: dict[str, Any] = { - "id": entry["id"], - "session_id": entry["session_id"], - "app_name": entry["app_name"], - "user_id": entry["user_id"], - "scope": entry.get("scope", "user"), - "event_id": entry["event_id"], - "author": entry["author"], - "timestamp": entry["timestamp"], - "embedding": bound_embedding, - "content_json": content_json, - "content_text": entry["content_text"], - "metadata_json": metadata_json, - "inserted_at": entry["inserted_at"], - } - if self._owner_id_column_name: - params["owner_id"] = owner_id - params = await coerce_large_parameters_async(conn, params) - if await self._execute_insert_entry(cursor, sql, params): - inserted_count += 1 - await conn.commit() + batch_params: list[dict[str, Any]] = [] + for entry in entries: + content_json = await self._serialize_json_field(entry["content_json"]) + metadata_json = await self._serialize_json_field(entry["metadata_json"]) + embedding_val = entry.get("embedding") + bound_embedding = array.array("f", embedding_val) if embedding_val is not None else None + p: dict[str, Any] = { + "id": entry["id"], + "session_id": entry["session_id"], + "app_name": entry["app_name"], + "user_id": entry["user_id"], + "scope": entry.get("scope", "user"), + "event_id": entry["event_id"], + "author": entry["author"], + "timestamp": entry["timestamp"], + "embedding": bound_embedding, + "content_json": content_json, + "content_text": entry["content_text"], + "metadata_json": metadata_json, + "inserted_at": entry["inserted_at"], + } + if self._owner_id_column_name: + p["owner_id"] = owner_id + batch_params.append(p) - return inserted_count + async with self._config.provide_connection() as conn: + try: + coerced_batch = [await coerce_large_parameters_async(conn, p) for p in batch_params] + with conn.cursor() as cursor: + await cursor.executemany(sql, coerced_batch, batcherrors=True) + batch_errors = list(cursor.getbatcherrors()) if hasattr(cursor, "getbatcherrors") else [] + duplicates = sum( + 1 + for e in batch_errors + if getattr(e, "code", getattr(e, "full_code", None)) == ORACLE_DUPLICATE_KEY_ERROR + ) + inserted_count = len(entries) - duplicates + await conn.commit() + return inserted_count + except BaseException: + await _rollback_async(conn) + raise async def search_entries( self, @@ -2930,36 +2976,48 @@ def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "object ) """ - inserted_count = 0 - with self._config.provide_connection() as conn, conn.cursor() as cursor: - for entry in entries: - content_json = self._serialize_json_field(entry["content_json"]) - metadata_json = self._serialize_json_field(entry["metadata_json"]) - embedding_val = entry.get("embedding") - bound_embedding = array.array("f", embedding_val) if embedding_val is not None else None - params: dict[str, Any] = { - "id": entry["id"], - "session_id": entry["session_id"], - "app_name": entry["app_name"], - "user_id": entry["user_id"], - "scope": entry.get("scope", "user"), - "event_id": entry["event_id"], - "author": entry["author"], - "timestamp": entry["timestamp"], - "embedding": bound_embedding, - "content_json": content_json, - "content_text": entry["content_text"], - "metadata_json": metadata_json, - "inserted_at": entry["inserted_at"], - } - if self._owner_id_column_name: - params["owner_id"] = owner_id - params = coerce_large_parameters_sync(conn, params) - if self._execute_insert_entry(cursor, sql, params): - inserted_count += 1 - conn.commit() + batch_params: list[dict[str, Any]] = [] + for entry in entries: + content_json = self._serialize_json_field(entry["content_json"]) + metadata_json = self._serialize_json_field(entry["metadata_json"]) + embedding_val = entry.get("embedding") + bound_embedding = array.array("f", embedding_val) if embedding_val is not None else None + p: dict[str, Any] = { + "id": entry["id"], + "session_id": entry["session_id"], + "app_name": entry["app_name"], + "user_id": entry["user_id"], + "scope": entry.get("scope", "user"), + "event_id": entry["event_id"], + "author": entry["author"], + "timestamp": entry["timestamp"], + "embedding": bound_embedding, + "content_json": content_json, + "content_text": entry["content_text"], + "metadata_json": metadata_json, + "inserted_at": entry["inserted_at"], + } + if self._owner_id_column_name: + p["owner_id"] = owner_id + batch_params.append(p) - return inserted_count + with self._config.provide_connection() as conn: + try: + coerced_batch = [coerce_large_parameters_sync(conn, p) for p in batch_params] + with conn.cursor() as cursor: + cursor.executemany(sql, coerced_batch, batcherrors=True) + batch_errors = list(cursor.getbatcherrors()) if hasattr(cursor, "getbatcherrors") else [] + duplicates = sum( + 1 + for e in batch_errors + if getattr(e, "code", getattr(e, "full_code", None)) == ORACLE_DUPLICATE_KEY_ERROR + ) + inserted_count = len(entries) - duplicates + conn.commit() + return inserted_count + except BaseException: + _rollback_sync(conn) + raise def search_entries( self, @@ -3371,20 +3429,11 @@ def _raise_session_not_found(session_id: str) -> NoReturn: raise ValueError(msg) +_normalize_decimals = build_nested_decimal_normalizer(mode="float") + + def _coerce_decimal_values(value: Any) -> Any: - if isinstance(value, Decimal): - return float(value) - if isinstance(value, dict): - return {key: _coerce_decimal_values(val) for key, val in value.items()} - if isinstance(value, list): - return [_coerce_decimal_values(item) for item in value] - if isinstance(value, tuple): - return tuple(_coerce_decimal_values(item) for item in value) - if isinstance(value, set): - return {_coerce_decimal_values(item) for item in value} - if isinstance(value, frozenset): - return frozenset(_coerce_decimal_values(item) for item in value) - return value + return _normalize_decimals(value) def _oracle_text_value(value: Any) -> str: @@ -3462,17 +3511,11 @@ def _adk_table_feature_clause( async def _read_lob_async(data: Any) -> Any: - if is_async_readable(data): - return await data.read() - if is_readable(data): - return data.read() - return data + return await read_oracle_lob_async(data) def _read_lob_sync(data: Any) -> Any: - if is_readable(data): - return data.read() - return data + return read_oracle_lob_sync(data) def _build_oracle_scope_where( diff --git a/sqlspec/adapters/oracledb/core.py b/sqlspec/adapters/oracledb/core.py index 75bfcf60f..9aa1ce70e 100644 --- a/sqlspec/adapters/oracledb/core.py +++ b/sqlspec/adapters/oracledb/core.py @@ -137,6 +137,8 @@ "normalize_execute_many_parameters_sync", "parse_sqlnet_ora", "parse_tnsnames_ora", + "read_oracle_lob_async", + "read_oracle_lob_sync", "release_extracted_wallet_dir", "resolve_dbms_cloud_source_uri", "resolve_row_metadata", @@ -3065,6 +3067,22 @@ def _resolve_callable_or_str(value: "str | Callable[[], str | None] | None") -> return value +async def read_oracle_lob_async(value: Any) -> Any: + """Read a LOB value asynchronously from an async connection or return raw data.""" + if is_async_readable(value): + return await value.read() + if is_readable(value): + return value.read() + return value + + +def read_oracle_lob_sync(value: Any) -> Any: + """Read a LOB value synchronously from a sync connection or return raw data.""" + if is_readable(value): + return value.read() + return value + + class _OracleSyncStreamDriver(Protocol): connection: "OracleSyncConnection" driver_features: dict[str, Any] diff --git a/sqlspec/adapters/oracledb/events/backend.py b/sqlspec/adapters/oracledb/events/backend.py index fc2501ec3..d7f218a15 100644 --- a/sqlspec/adapters/oracledb/events/backend.py +++ b/sqlspec/adapters/oracledb/events/backend.py @@ -10,6 +10,7 @@ from sqlspec.exceptions import ImproperConfigurationError, MissingDependencyError from sqlspec.extensions.events import EventMessage, parse_event_timestamp from sqlspec.utils.logging import get_logger, log_with_context +from sqlspec.utils.serializers import from_json from sqlspec.utils.uuids import uuid4 if TYPE_CHECKING: @@ -272,6 +273,11 @@ def _build_envelope( def _parse_message(channel: str, payload: Any) -> EventMessage: """Parse Oracle AQ message payload into EventMessage.""" + if isinstance(payload, (str, bytes)): + with contextlib.suppress(Exception): + parsed = from_json(payload) + if isinstance(parsed, dict): + payload = parsed if not isinstance(payload, dict): payload = {"payload": payload} payload_channel = payload.get("channel") diff --git a/sqlspec/adapters/oracledb/events/store.py b/sqlspec/adapters/oracledb/events/store.py index 36e5f4156..b5954477f 100644 --- a/sqlspec/adapters/oracledb/events/store.py +++ b/sqlspec/adapters/oracledb/events/store.py @@ -22,7 +22,7 @@ from typing_extensions import TypedDict -from sqlspec.adapters.oracledb._storage import oracle_table_feature_report +from sqlspec.adapters.oracledb._storage import oracle_table_feature_report, validate_oracle_identifier from sqlspec.adapters.oracledb.data_dictionary import JSONStorageType, storage_type_from_version from sqlspec.config import EventsConfig from sqlspec.extensions.events import BaseEventQueueStore @@ -133,6 +133,8 @@ class OracleSyncEventQueueStore(BaseEventQueueStore["OracleSyncConfig"]): def __init__(self, config: "OracleSyncConfig") -> None: """Initialize Oracle sync event queue store.""" super().__init__(config) + for part in self.table_name.split("."): + validate_oracle_identifier(part, "table name") self._in_memory, self._json_storage_override = _init_oracle_settings(self._extension_settings) def create_statements(self) -> "list[str]": @@ -248,6 +250,8 @@ class OracleAsyncEventQueueStore(BaseEventQueueStore["OracleAsyncConfig"]): def __init__(self, config: "OracleAsyncConfig") -> None: """Initialize Oracle async event queue store.""" super().__init__(config) + for part in self.table_name.split("."): + validate_oracle_identifier(part, "table name") self._in_memory, self._json_storage_override = _init_oracle_settings(self._extension_settings) def create_statements(self) -> "list[str]": diff --git a/sqlspec/adapters/oracledb/litestar/store.py b/sqlspec/adapters/oracledb/litestar/store.py index 080146d6a..1f53ce0aa 100644 --- a/sqlspec/adapters/oracledb/litestar/store.py +++ b/sqlspec/adapters/oracledb/litestar/store.py @@ -1,14 +1,14 @@ """Oracle session store for Litestar integration.""" +import contextlib from datetime import UTC, timedelta from typing import TYPE_CHECKING, Any, NotRequired, TypedDict, cast -from sqlspec.adapters.oracledb._storage import oracle_table_feature_report -from sqlspec.adapters.oracledb.core import DB_TYPE_BLOB +from sqlspec.adapters.oracledb._storage import oracle_table_feature_report, validate_oracle_identifier +from sqlspec.adapters.oracledb.core import DB_TYPE_BLOB, read_oracle_lob_async, read_oracle_lob_sync, resolve_rowcount from sqlspec.config import LitestarConfig from sqlspec.extensions.litestar.store import BaseSQLSpecStore from sqlspec.utils.sync_tools import async_ -from sqlspec.utils.type_guards import is_async_readable, is_readable if TYPE_CHECKING: from sqlspec.adapters.oracledb.config import OracleAsyncConfig, OracleSyncConfig @@ -109,6 +109,8 @@ def __init__(self, config: "OracleAsyncConfig") -> None: config: OracleAsyncConfig instance. """ super().__init__(config) + for part in self._table_name.split("."): + validate_oracle_identifier(part, "table name") litestar_config = config.extension_config.get("litestar", {}) self._in_memory = bool(litestar_config.get("in_memory", False)) @@ -137,14 +139,14 @@ async def get(self, key: str, renew_for: "int | timedelta | None" = None) -> "by """ sql = f""" SELECT data, expires_at FROM {self._table_name} - WHERE session_id = :session_id + WHERE session_id = :1 AND (expires_at IS NULL OR expires_at > SYSTIMESTAMP) """ conn_context = self._config.provide_connection() async with conn_context as conn: with conn.cursor() as cursor: - await cursor.execute(sql, {"session_id": key}) + await cursor.execute(sql, (key,)) row = await cursor.fetchone() if row is None: @@ -158,14 +160,19 @@ async def get(self, key: str, renew_for: "int | timedelta | None" = None) -> "by update_sql = f""" UPDATE {self._table_name} SET expires_at = CASE - WHEN :expires_in_seconds IS NULL THEN NULL - ELSE SYSTIMESTAMP + NUMTODSINTERVAL(:expires_in_seconds, 'SECOND') + WHEN :1 IS NULL THEN NULL + ELSE SYSTIMESTAMP + NUMTODSINTERVAL(:1, 'SECOND') END, updated_at = SYSTIMESTAMP - WHERE session_id = :session_id + WHERE session_id = :2 """ - await cursor.execute(update_sql, {"expires_in_seconds": expires_in_seconds, "session_id": key}) - await conn.commit() + try: + await cursor.execute(update_sql, (expires_in_seconds, key)) + await conn.commit() + except BaseException: + with contextlib.suppress(Exception): + await conn.rollback() + raise return await _read_blob_async(data_blob) @@ -208,12 +215,19 @@ async def set(self, key: str, value: "str | bytes", expires_in: "int | timedelta conn_context = self._config.provide_connection() async with conn_context as conn: - bind_data: Any = await conn.createlob(DB_TYPE_BLOB, data) if len(data) > ORACLE_SMALL_BLOB_LIMIT else data - with conn.cursor() as cursor: - await cursor.execute( - sql, {"session_id": key, "data": bind_data, "expires_in_seconds": expires_in_seconds} + try: + bind_data: Any = ( + await conn.createlob(DB_TYPE_BLOB, data) if len(data) > ORACLE_SMALL_BLOB_LIMIT else data ) - await conn.commit() + with conn.cursor() as cursor: + await cursor.execute( + sql, {"session_id": key, "data": bind_data, "expires_in_seconds": expires_in_seconds} + ) + await conn.commit() + except BaseException: + with contextlib.suppress(Exception): + await conn.rollback() + raise async def delete(self, key: str) -> None: """Delete a session by key. @@ -221,13 +235,18 @@ async def delete(self, key: str) -> None: Args: key: Session ID to delete. """ - sql = f"DELETE FROM {self._table_name} WHERE session_id = :session_id" + sql = f"DELETE FROM {self._table_name} WHERE session_id = :1" conn_context = self._config.provide_connection() async with conn_context as conn: - with conn.cursor() as cursor: - await cursor.execute(sql, {"session_id": key}) - await conn.commit() + try: + with conn.cursor() as cursor: + await cursor.execute(sql, (key,)) + await conn.commit() + except BaseException: + with contextlib.suppress(Exception): + await conn.rollback() + raise async def delete_all(self) -> None: """Delete all sessions from the store.""" @@ -235,9 +254,14 @@ async def delete_all(self) -> None: conn_context = self._config.provide_connection() async with conn_context as conn: - with conn.cursor() as cursor: - await cursor.execute(sql) - await conn.commit() + try: + with conn.cursor() as cursor: + await cursor.execute(sql) + await conn.commit() + except BaseException: + with contextlib.suppress(Exception): + await conn.rollback() + raise self._log_delete_all() async def exists(self, key: str) -> bool: @@ -251,14 +275,14 @@ async def exists(self, key: str) -> bool: """ sql = f""" SELECT 1 FROM {self._table_name} - WHERE session_id = :session_id + WHERE session_id = :1 AND (expires_at IS NULL OR expires_at > SYSTIMESTAMP) """ conn_context = self._config.provide_connection() async with conn_context as conn: with conn.cursor() as cursor: - await cursor.execute(sql, {"session_id": key}) + await cursor.execute(sql, (key,)) result = await cursor.fetchone() return result is not None @@ -273,13 +297,13 @@ async def expires_in(self, key: str) -> "int | None": """ sql = f""" SELECT expires_at, SYSTIMESTAMP FROM {self._table_name} - WHERE session_id = :session_id + WHERE session_id = :1 """ conn_context = self._config.provide_connection() async with conn_context as conn: with conn.cursor() as cursor: - await cursor.execute(sql, {"session_id": key}) + await cursor.execute(sql, (key,)) row = await cursor.fetchone() if row is None or row[0] is None: @@ -308,78 +332,26 @@ async def delete_expired(self) -> int: conn_context = self._config.provide_connection() async with conn_context as conn: - with conn.cursor() as cursor: - await cursor.execute(sql) - count = cursor.rowcount if cursor.rowcount is not None else 0 - await conn.commit() - if count > 0: - self._log_delete_expired(count) - return count + try: + with conn.cursor() as cursor: + await cursor.execute(sql) + count = resolve_rowcount(cursor) + await conn.commit() + if count > 0: + self._log_delete_expired(count) + return count + except BaseException: + with contextlib.suppress(Exception): + await conn.rollback() + raise def _table_ddl(self) -> str: - """Get Oracle CREATE TABLE SQL with optimized schema. - - Returns: - SQL statement to create the sessions table with proper indexes. - """ - table_clause = _litestar_table_feature_clause(self._config, self._in_memory) - index_name = f"idx_{self._table_name.replace('.', '_')}_expires_at"[:30] - return f""" - BEGIN - EXECUTE IMMEDIATE 'CREATE TABLE {self._table_name} ( - session_id VARCHAR2(255) PRIMARY KEY, - data BLOB NOT NULL, - expires_at TIMESTAMP WITH TIME ZONE, - created_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL, - updated_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL - ){table_clause}'; - EXCEPTION - WHEN OTHERS THEN - IF SQLCODE != -955 THEN - RAISE; - END IF; - END; - - BEGIN - EXECUTE IMMEDIATE 'CREATE INDEX {index_name} - ON {self._table_name}(expires_at)'; - EXCEPTION - WHEN OTHERS THEN - IF SQLCODE != -955 THEN - RAISE; - END IF; - END; - """ + """Get Oracle CREATE TABLE SQL with optimized schema.""" + return _build_session_table_ddl(self._table_name, self._config, self._in_memory) def _drop_table_sql(self) -> "list[str]": - """Get Oracle DROP TABLE SQL with PL/SQL error handling. - - Returns: - List of SQL statements with exception handling for non-existent objects. - """ - index_name = f"idx_{self._table_name.replace('.', '_')}_expires_at"[:30] - return [ - f""" - BEGIN - EXECUTE IMMEDIATE 'DROP INDEX {index_name}'; - EXCEPTION - WHEN OTHERS THEN - IF SQLCODE != -1418 THEN - RAISE; - END IF; - END; - """, - f""" - BEGIN - EXECUTE IMMEDIATE 'DROP TABLE {self._table_name}'; - EXCEPTION - WHEN OTHERS THEN - IF SQLCODE != -942 THEN - RAISE; - END IF; - END; - """, - ] + """Get Oracle DROP TABLE SQL with PL/SQL error handling.""" + return _build_drop_session_table_sql(self._table_name) class OracleSyncStore(BaseSQLSpecStore["OracleSyncConfig"]): @@ -415,6 +387,8 @@ def __init__(self, config: "OracleSyncConfig") -> None: config: OracleSyncConfig instance. """ super().__init__(config) + for part in self._table_name.split("."): + validate_oracle_identifier(part, "table name") litestar_config = config.extension_config.get("litestar", {}) self._in_memory = bool(litestar_config.get("in_memory", False)) @@ -492,69 +466,12 @@ async def delete_expired(self) -> int: return await async_(self._delete_expired)() def _table_ddl(self) -> str: - """Get Oracle CREATE TABLE SQL with optimized schema. - - Returns: - SQL statement to create the sessions table with proper indexes. - """ - table_clause = _litestar_table_feature_clause(self._config, self._in_memory) - index_name = f"idx_{self._table_name.replace('.', '_')}_expires_at"[:30] - return f""" - BEGIN - EXECUTE IMMEDIATE 'CREATE TABLE {self._table_name} ( - session_id VARCHAR2(255) PRIMARY KEY, - data BLOB NOT NULL, - expires_at TIMESTAMP WITH TIME ZONE, - created_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL, - updated_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL - ){table_clause}'; - EXCEPTION - WHEN OTHERS THEN - IF SQLCODE != -955 THEN - RAISE; - END IF; - END; - - BEGIN - EXECUTE IMMEDIATE 'CREATE INDEX {index_name} - ON {self._table_name}(expires_at)'; - EXCEPTION - WHEN OTHERS THEN - IF SQLCODE != -955 THEN - RAISE; - END IF; - END; - """ + """Get Oracle CREATE TABLE SQL with optimized schema.""" + return _build_session_table_ddl(self._table_name, self._config, self._in_memory) def _drop_table_sql(self) -> "list[str]": - """Get Oracle DROP TABLE SQL with PL/SQL error handling. - - Returns: - List of SQL statements with exception handling for non-existent objects. - """ - index_name = f"idx_{self._table_name.replace('.', '_')}_expires_at"[:30] - return [ - f""" - BEGIN - EXECUTE IMMEDIATE 'DROP INDEX {index_name}'; - EXCEPTION - WHEN OTHERS THEN - IF SQLCODE != -1418 THEN - RAISE; - END IF; - END; - """, - f""" - BEGIN - EXECUTE IMMEDIATE 'DROP TABLE {self._table_name}'; - EXCEPTION - WHEN OTHERS THEN - IF SQLCODE != -942 THEN - RAISE; - END IF; - END; - """, - ] + """Get Oracle DROP TABLE SQL with PL/SQL error handling.""" + return _build_drop_session_table_sql(self._table_name) def _create_table(self) -> None: """Synchronous implementation of create_table.""" @@ -568,12 +485,12 @@ def _get(self, key: str, renew_for: "int | timedelta | None" = None) -> "bytes | """Synchronous implementation of get.""" sql = f""" SELECT data, expires_at FROM {self._table_name} - WHERE session_id = :session_id + WHERE session_id = :1 AND (expires_at IS NULL OR expires_at > SYSTIMESTAMP) """ with self._config.provide_connection() as conn, conn.cursor() as cursor: - cursor.execute(sql, {"session_id": key}) + cursor.execute(sql, (key,)) row = cursor.fetchone() if row is None: @@ -587,14 +504,19 @@ def _get(self, key: str, renew_for: "int | timedelta | None" = None) -> "bytes | update_sql = f""" UPDATE {self._table_name} SET expires_at = CASE - WHEN :expires_in_seconds IS NULL THEN NULL - ELSE SYSTIMESTAMP + NUMTODSINTERVAL(:expires_in_seconds, 'SECOND') + WHEN :1 IS NULL THEN NULL + ELSE SYSTIMESTAMP + NUMTODSINTERVAL(:1, 'SECOND') END, updated_at = SYSTIMESTAMP - WHERE session_id = :session_id + WHERE session_id = :2 """ - cursor.execute(update_sql, {"expires_in_seconds": expires_in_seconds, "session_id": key}) - conn.commit() + try: + cursor.execute(update_sql, (expires_in_seconds, key)) + conn.commit() + except BaseException: + with contextlib.suppress(Exception): + conn.rollback() + raise return _read_blob_sync(data_blob) @@ -630,38 +552,57 @@ def _set(self, key: str, value: "str | bytes", expires_in: "int | timedelta | No """ with self._config.provide_connection() as conn: - bind_data: Any = conn.createlob(DB_TYPE_BLOB, data) if len(data) > ORACLE_SMALL_BLOB_LIMIT else data - with conn.cursor() as cursor: - cursor.execute(sql, {"session_id": key, "data": bind_data, "expires_in_seconds": expires_in_seconds}) - conn.commit() + try: + bind_data: Any = conn.createlob(DB_TYPE_BLOB, data) if len(data) > ORACLE_SMALL_BLOB_LIMIT else data + with conn.cursor() as cursor: + cursor.execute( + sql, {"session_id": key, "data": bind_data, "expires_in_seconds": expires_in_seconds} + ) + conn.commit() + except BaseException: + with contextlib.suppress(Exception): + conn.rollback() + raise def _delete(self, key: str) -> None: """Synchronous implementation of delete.""" - sql = f"DELETE FROM {self._table_name} WHERE session_id = :session_id" + sql = f"DELETE FROM {self._table_name} WHERE session_id = :1" - with self._config.provide_connection() as conn, conn.cursor() as cursor: - cursor.execute(sql, {"session_id": key}) - conn.commit() + with self._config.provide_connection() as conn: + try: + with conn.cursor() as cursor: + cursor.execute(sql, (key,)) + conn.commit() + except BaseException: + with contextlib.suppress(Exception): + conn.rollback() + raise def _delete_all(self) -> None: """Synchronous implementation of delete_all.""" sql = f"DELETE FROM {self._table_name}" - with self._config.provide_connection() as conn, conn.cursor() as cursor: - cursor.execute(sql) - conn.commit() + with self._config.provide_connection() as conn: + try: + with conn.cursor() as cursor: + cursor.execute(sql) + conn.commit() + except BaseException: + with contextlib.suppress(Exception): + conn.rollback() + raise self._log_delete_all() def _exists(self, key: str) -> bool: """Synchronous implementation of exists.""" sql = f""" SELECT 1 FROM {self._table_name} - WHERE session_id = :session_id + WHERE session_id = :1 AND (expires_at IS NULL OR expires_at > SYSTIMESTAMP) """ with self._config.provide_connection() as conn, conn.cursor() as cursor: - cursor.execute(sql, {"session_id": key}) + cursor.execute(sql, (key,)) result = cursor.fetchone() return result is not None @@ -669,11 +610,11 @@ def _expires_in(self, key: str) -> "int | None": """Synchronous implementation of expires_in.""" sql = f""" SELECT expires_at, SYSTIMESTAMP FROM {self._table_name} - WHERE session_id = :session_id + WHERE session_id = :1 """ with self._config.provide_connection() as conn, conn.cursor() as cursor: - cursor.execute(sql, {"session_id": key}) + cursor.execute(sql, (key,)) row = cursor.fetchone() if row is None or row[0] is None: @@ -696,13 +637,19 @@ def _delete_expired(self) -> int: """Synchronous implementation of delete_expired.""" sql = f"DELETE FROM {self._table_name} WHERE expires_at <= SYSTIMESTAMP" - with self._config.provide_connection() as conn, conn.cursor() as cursor: - cursor.execute(sql) - count = cursor.rowcount if cursor.rowcount is not None else 0 - conn.commit() - if count > 0: - self._log_delete_expired(count) - return count + with self._config.provide_connection() as conn: + try: + with conn.cursor() as cursor: + cursor.execute(sql) + count = resolve_rowcount(cursor) + conn.commit() + if count > 0: + self._log_delete_expired(count) + return count + except BaseException: + with contextlib.suppress(Exception): + conn.rollback() + raise def _oracle_expiry_seconds(expires_in: "int | timedelta | None") -> "int | None": @@ -729,18 +676,73 @@ def _coerce_bytes_payload(value: object) -> bytes: async def _read_blob_async(value: object) -> bytes: """Read LOB values from async connections into bytes.""" - if is_async_readable(value): - return _coerce_bytes_payload(await value.read()) - if is_readable(value): - return _coerce_bytes_payload(value.read()) - return _coerce_bytes_payload(value) + data = await read_oracle_lob_async(value) + return _coerce_bytes_payload(data) def _read_blob_sync(value: object) -> bytes: """Read LOB values from sync connections into bytes.""" - if is_readable(value): - return _coerce_bytes_payload(value.read()) - return _coerce_bytes_payload(value) + data = read_oracle_lob_sync(value) + return _coerce_bytes_payload(data) + + +def _build_session_table_ddl(table_name: str, config: Any, in_memory: bool) -> str: + """Build Oracle CREATE TABLE and INDEX SQL for the session store.""" + table_clause = _litestar_table_feature_clause(config, in_memory) + index_name = f"idx_{table_name.replace('.', '_')}_expires_at"[:30] + return f""" + BEGIN + EXECUTE IMMEDIATE 'CREATE TABLE {table_name} ( + session_id VARCHAR2(255) PRIMARY KEY, + data BLOB NOT NULL, + expires_at TIMESTAMP WITH TIME ZONE, + created_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL, + updated_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL + ){table_clause}'; + EXCEPTION + WHEN OTHERS THEN + IF SQLCODE != -955 THEN + RAISE; + END IF; + END; + + BEGIN + EXECUTE IMMEDIATE 'CREATE INDEX {index_name} + ON {table_name}(expires_at)'; + EXCEPTION + WHEN OTHERS THEN + IF SQLCODE != -955 THEN + RAISE; + END IF; + END; + """ + + +def _build_drop_session_table_sql(table_name: str) -> "list[str]": + """Build Oracle DROP TABLE and DROP INDEX statements for the session store.""" + index_name = f"idx_{table_name.replace('.', '_')}_expires_at"[:30] + return [ + f""" + BEGIN + EXECUTE IMMEDIATE 'DROP INDEX {index_name}'; + EXCEPTION + WHEN OTHERS THEN + IF SQLCODE != -1418 THEN + RAISE; + END IF; + END; + """, + f""" + BEGIN + EXECUTE IMMEDIATE 'DROP TABLE {table_name}'; + EXCEPTION + WHEN OTHERS THEN + IF SQLCODE != -942 THEN + RAISE; + END IF; + END; + """, + ] def _litestar_table_feature_clause(config: Any, in_memory: bool) -> str: diff --git a/tests/unit/adapters/test_oracledb/test_driver.py b/tests/unit/adapters/test_oracledb/test_driver.py index a8061a042..bc31c4e8f 100644 --- a/tests/unit/adapters/test_oracledb/test_driver.py +++ b/tests/unit/adapters/test_oracledb/test_driver.py @@ -595,3 +595,77 @@ def execute(self, stmt: str, params: Any = None) -> None: assert len(executed_statements) == 2 assert executed_statements[0][1] == [] assert executed_statements[1][1] == [101, "Alice"] + + +def test_parse_message_string_or_bytes_payload() -> None: + """_parse_message parses JSON string or bytes into a structured EventMessage.""" + import json + + from sqlspec.adapters.oracledb.events.backend import _parse_message + + raw_dict = { + "event_id": "ev-123", + "channel": "orders", + "payload": {"order_id": 42}, + "metadata": {"source": "unit"}, + "published_at": "2026-05-10T12:00:00Z", + } + json_str = json.dumps(raw_dict) + msg_from_str = _parse_message("orders", json_str) + assert msg_from_str.event_id == "ev-123" + assert msg_from_str.channel == "orders" + assert msg_from_str.payload == {"order_id": 42} + assert msg_from_str.metadata == {"source": "unit"} + + json_bytes = json_str.encode("utf-8") + msg_from_bytes = _parse_message("orders", json_bytes) + assert msg_from_bytes.event_id == "ev-123" + assert msg_from_bytes.payload == {"order_id": 42} + + +def test_oracle_aq_commit_at_dequeue_and_noop_ack_nack() -> None: + """Oracle AQ hub commits at dequeue for at-most-once delivery; ack and nack are no-ops.""" + import json + + from sqlspec.adapters.oracledb.events._hub import OracleSyncAQHub + from sqlspec.adapters.oracledb.events.backend import OracleSyncAQEventBackend + + class _MockOracleSyncConfig: + __module__ = "sqlspec.adapters.oracledb.config" + is_async = False + driver_features: dict[str, object] = {} + + def get_observability_runtime(self) -> object: + return MagicMock() + + driver = MagicMock() + session_cm = MagicMock() + session_cm.__enter__.return_value = driver + + config = _MockOracleSyncConfig() + config.provide_session = lambda: session_cm # type: ignore[attr-defined] + + msg = MagicMock() + msg.payload = json.dumps({"event_id": "ev-pinned", "payload": {"foo": "bar"}}) + queue = MagicMock() + queue.deqone.return_value = msg + + hub = OracleSyncAQHub( + cast("Any", config), + queue_name_template="TEST_QUEUE", + visibility=None, + default_visibility=None, + wait_ceiling=5, + backend_name="aq", + ) + hub._session_driver = driver + hub._queues["test_chan"] = queue + + payload = hub.dequeue("test_chan", poll_interval=0.1) + assert payload is not None + driver.commit.assert_called_once() + + backend = OracleSyncAQEventBackend(cast("Any", config)) + backend._hub = hub + backend.ack("ev-pinned") + backend.nack("ev-pinned") diff --git a/tests/unit/adapters/test_oracledb/test_litestar_store.py b/tests/unit/adapters/test_oracledb/test_litestar_store.py index df71233c7..4bda51d10 100644 --- a/tests/unit/adapters/test_oracledb/test_litestar_store.py +++ b/tests/unit/adapters/test_oracledb/test_litestar_store.py @@ -3,6 +3,8 @@ from datetime import UTC, datetime, timedelta from typing import Any, Self, cast +import pytest + from sqlspec.adapters.oracledb.core import DB_TYPE_BLOB from sqlspec.adapters.oracledb.data_dictionary import OracleVersionCache from sqlspec.adapters.oracledb.litestar import OracleSyncStore @@ -113,7 +115,7 @@ def test_oracle_sync_store_expires_in_uses_database_clock() -> None: sql, parameters = cursor.executed[0] assert "SELECT expires_at, SYSTIMESTAMP" in sql - assert parameters == {"session_id": "session-1"} + assert parameters == ("session-1",) def test_oracle_litestar_store_interval_partition_on_expires_at() -> None: @@ -163,3 +165,24 @@ def test_oracle_litestar_store_table_ddl_sanitizes_and_truncates_index_name() -> assert "CREATE INDEX idx_my_very_long_schema_oracle" in ddl assert "idx_my_very_long_schema.oracle" not in ddl + + +def test_oracle_store_set_rolls_back_on_execute_failure() -> None: + """Store mutations roll back connection on execution failure.""" + from unittest.mock import MagicMock + + cursor = MagicMock() + cursor.execute.side_effect = RuntimeError("execute failed") + connection = MagicMock() + connection.cursor.return_value.__enter__.return_value = cursor + context = MagicMock() + context.__enter__.return_value = connection + config = MagicMock() + config.provide_connection.return_value = context + config.extension_config = {} + + store = OracleSyncStore(config) + with pytest.raises(RuntimeError, match="execute failed"): + store._set("fail-key", b"data") + + connection.rollback.assert_called_once() diff --git a/tests/unit/adapters/test_oracledb/test_oracle_adk_store.py b/tests/unit/adapters/test_oracledb/test_oracle_adk_store.py index 62676a0a5..1e57fd770 100644 --- a/tests/unit/adapters/test_oracledb/test_oracle_adk_store.py +++ b/tests/unit/adapters/test_oracledb/test_oracle_adk_store.py @@ -348,12 +348,13 @@ def test_oracle_adk_sync_insert_binds_embedding_array() -> None: ] inserted = store.insert_memory_entries(entries) assert inserted == 1 - assert cursor.execute.called - _call_sql, call_params = cursor.execute.call_args[0] - assert "embedding" in call_params - assert isinstance(call_params["embedding"], array.array) - assert call_params["embedding"].typecode == "f" - assert list(call_params["embedding"]) == pytest.approx([0.1, 0.2, 0.3]) + assert cursor.executemany.called + _call_sql, call_params = cursor.executemany.call_args[0] + first_param = call_params[0] + assert "embedding" in first_param + assert isinstance(first_param["embedding"], array.array) + assert first_param["embedding"].typecode == "f" + assert list(first_param["embedding"]) == pytest.approx([0.1, 0.2, 0.3]) def test_oracle_adk_memory_rows_to_records_decodes_embedding() -> None: @@ -383,6 +384,73 @@ def test_oracle_adk_memory_rows_to_records_decodes_embedding() -> None: assert records[0]["embedding"] == pytest.approx([0.1, 0.2, 0.3]) +def test_adk_store_mutation_rolls_back_on_error() -> None: + """ADK store mutation rolls back connection on error.""" + cursor = MagicMock() + cursor.execute.side_effect = RuntimeError("mutation error") + conn = MagicMock() + conn.cursor.return_value.__enter__.return_value = cursor + config = _config_for_storage(JSONStorageType.BLOB_JSON) + config.extension_config = {"adk": {}} + config.provide_connection.return_value.__enter__.return_value = conn + + store = OracleSyncADKStore(config) + with pytest.raises(RuntimeError, match="mutation error"): + store.delete_session("app", "user", "sess-1") + + conn.rollback.assert_called_once() + + +def test_oracle_adk_memory_store_batch_insert() -> None: + """ADK memory store executes batch insert via executemany.""" + cursor = MagicMock() + cursor.rowcount = 2 + cursor.getbatcherrors.return_value = [] + conn = MagicMock() + conn.cursor.return_value.__enter__.return_value = cursor + config = _config_for_storage(JSONStorageType.JSON_NATIVE) + config.extension_config = {"adk": {}} + config.provide_connection.return_value.__enter__.return_value = conn + + store = OracleSyncADKMemoryStore(config) + now = datetime(2026, 5, 10, 12, 0, tzinfo=UTC) + entries = [ + { + "id": "mem-1", + "session_id": "sess-1", + "app_name": "app", + "user_id": "user", + "scope": "user", + "event_id": "ev-1", + "author": "user", + "timestamp": now, + "content_json": {"k": 1}, + "content_text": "text 1", + "metadata_json": None, + "inserted_at": now, + }, + { + "id": "mem-2", + "session_id": "sess-1", + "app_name": "app", + "user_id": "user", + "scope": "user", + "event_id": "ev-2", + "author": "user", + "timestamp": now, + "content_json": {"k": 2}, + "content_text": "text 2", + "metadata_json": None, + "inserted_at": now, + }, + ] + inserted = store.insert_memory_entries(entries) + assert inserted == 2 + cursor.executemany.assert_called_once() + assert len(cursor.executemany.call_args[0][1]) == 2 + conn.commit.assert_called_once() + + def _sync_store_with_driver() -> "tuple[Any, MagicMock, MagicMock]": config = _config_for_storage(JSONStorageType.BLOB_JSON) config.extension_config = {"adk": {}} @@ -446,6 +514,10 @@ def close(self) -> None: def execute(self, sql: str, params: "dict[str, Any] | None" = None) -> None: self.calls.append((sql, dict(params or {}))) + def executemany(self, sql: str, params: Any = None, **_kwargs: Any) -> None: + for p in params or []: + self.calls.append((sql, dict(p or {}))) + def fetchall(self) -> "list[Any]": return [] From 3c40a27de4be8ec5199138d17ac6a9dfaecba452 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 01:02:29 +0000 Subject: [PATCH 067/174] refactor(oracledb): relocate version cache to compiled core and resolve catalog issues Relocate OracleVersionCache and connection version resolution helpers into compiled core.py and re-export from data_dictionary.py to prevent cross-ABI boundary performance overhead. Await AsyncLOB reads in asynchronous DDL catalog introspection, route pre-12c JSON storage to DB_TYPE_CLOB, eliminate redundant second table metadata query in get_tables, use safe attribute extraction for cursor FetchInfo metadata, and validate table identifiers in DDL construction. --- sqlspec/adapters/oracledb/_json_handlers.py | 133 ++++++++---------- sqlspec/adapters/oracledb/_param_types.py | 18 +++ sqlspec/adapters/oracledb/_uuid_handlers.py | 9 +- sqlspec/adapters/oracledb/_vector_handlers.py | 2 +- sqlspec/adapters/oracledb/core.py | 74 +++++++++- sqlspec/adapters/oracledb/data_dictionary.py | 83 +++-------- .../test_data_dictionary_oracledb.py | 39 +++++ .../test_oracledb/test_json_handlers.py | 10 +- .../test_oracle_version_unification.py | 16 +++ .../test_oracledb/test_uuid_handlers.py | 22 +++ 10 files changed, 253 insertions(+), 153 deletions(-) diff --git a/sqlspec/adapters/oracledb/_json_handlers.py b/sqlspec/adapters/oracledb/_json_handlers.py index 3db1f0be5..e8972735f 100644 --- a/sqlspec/adapters/oracledb/_json_handlers.py +++ b/sqlspec/adapters/oracledb/_json_handlers.py @@ -38,6 +38,56 @@ from functools import partial from typing import TYPE_CHECKING, Any + +def chain_input_handler(inner: Any, fallback: "Any | None") -> Any: + """Build an input type handler that chains ``inner`` to ``fallback``.""" + if fallback is not None and _has_handler(fallback, inner, _chained_input_handler): + return fallback + return partial(_chained_input_handler, inner, fallback) + + +def chain_output_handler(inner: Any, fallback: "Any | None") -> Any: + """Build an output type handler that chains ``inner`` to ``fallback``.""" + if fallback is not None and _has_handler(fallback, inner, _chained_output_handler): + return fallback + return partial(_chained_output_handler, inner, fallback) + + +def _has_handler(handler: Any, target_inner: Any, chain_function: Any) -> bool: + current = handler + while current is not None: + if current is target_inner: + return True + if not isinstance(current, partial) or current.func is not chain_function: + return False + inner, current = current.args + if inner is target_inner: + return True + return False + + +def _chained_input_handler( + inner: Any, fallback: "Any | None", cursor: Any, value: Any, arraysize: int +) -> Any: + """Run ``inner`` input handler, falling back to ``fallback`` when it abstains.""" + result = inner(cursor, value, arraysize) + if result is not None: + return result + if fallback is not None: + return fallback(cursor, value, arraysize) + return None + + +def _chained_output_handler(inner: Any, fallback: "Any | None", cursor: Any, metadata: Any) -> Any: + """Run ``inner`` output handler, falling back to ``fallback`` when it abstains.""" + result = inner(cursor, metadata) + if result is not None: + return result + if fallback is not None: + return fallback(cursor, metadata) + return None + + from sqlspec.adapters.oracledb._param_types import OracleJson from sqlspec.adapters.oracledb._typing import ( DB_TYPE_BLOB, @@ -50,11 +100,16 @@ DB_TYPE_NVARCHAR, DB_TYPE_VARCHAR, ) -from sqlspec.adapters.oracledb.data_dictionary import ( +from sqlspec.adapters.oracledb.core import ( + is_json_payload, resolve_oracle_connection_compatible_major, resolve_oracle_connection_major, ) -from sqlspec.data_dictionary.dialects.oracle import ORACLE_JSON_STORAGE_NATIVE, resolve_oracle_json_storage +from sqlspec.data_dictionary.dialects.oracle import ( + ORACLE_JSON_STORAGE_BLOB_PLAIN, + ORACLE_JSON_STORAGE_NATIVE, + resolve_oracle_json_storage, +) from sqlspec.utils.serializers import from_json, to_json if TYPE_CHECKING: @@ -152,52 +207,6 @@ def register_json_handlers(connection: "Connection | AsyncConnection") -> None: connection.outputtypehandler = chain_output_handler(_output_type_handler, existing_output) -def chain_input_handler(inner: Any, fallback: "Any | None") -> Any: - """Build an input type handler that chains ``inner`` to ``fallback``.""" - if fallback is not None and _has_handler(fallback, inner, _chained_input_handler): - return fallback - return partial(_chained_input_handler, inner, fallback) - - -def chain_output_handler(inner: Any, fallback: "Any | None") -> Any: - """Build an output type handler that chains ``inner`` to ``fallback``.""" - if fallback is not None and _has_handler(fallback, inner, _chained_output_handler): - return fallback - return partial(_chained_output_handler, inner, fallback) - - -def is_json_payload(value: Any) -> bool: - """Return True if the value should be claimed by the JSON input handler. - - ``OracleJson``, ``dict``, empty ``list``/``tuple``, and ``list``/``tuple`` - containing any non-numeric element are claimed. Non-empty sequences where - all elements are ``int``/``float`` (excluding ``bool``) are reserved for - the vector handler unless wrapped in ``OracleJson``. - """ - if isinstance(value, OracleJson): - return True - if isinstance(value, dict): - return True - if isinstance(value, (list, tuple)): - if not value: - return True - return not all(isinstance(v, (int, float)) and not isinstance(v, bool) for v in value) - return False - - -def _has_handler(handler: Any, target_inner: Any, chain_function: Any) -> bool: - current = handler - while current is not None: - if current is target_inner: - return True - if not isinstance(current, partial) or current.func is not chain_function: - return False - inner, current = current.args - if inner is target_inner: - return True - return False - - def _input_type_handler(cursor: "Cursor | AsyncCursor", value: Any, arraysize: int) -> Any: """Oracle input type handler for JSON-shaped Python values.""" if not is_json_payload(value): @@ -206,9 +215,7 @@ def _input_type_handler(cursor: "Cursor | AsyncCursor", value: Any, arraysize: i server_major = resolve_oracle_connection_major(cursor.connection) if server_major is None: - if isinstance(value, OracleJson): - return cursor.var(DB_TYPE_JSON, arraysize=arraysize, inconverter=json_converter_in_native) - return cursor.var(DB_TYPE_JSON, arraysize=arraysize) + return cursor.var(DB_TYPE_BLOB, arraysize=arraysize, inconverter=json_converter_in_blob) compatible_major = resolve_oracle_connection_compatible_major(cursor.connection) storage = resolve_oracle_json_storage(server_major, compatible_major=compatible_major) @@ -216,6 +223,8 @@ def _input_type_handler(cursor: "Cursor | AsyncCursor", value: Any, arraysize: i if isinstance(value, OracleJson): return cursor.var(DB_TYPE_JSON, arraysize=arraysize, inconverter=json_converter_in_native) return cursor.var(DB_TYPE_JSON, arraysize=arraysize) + if storage == ORACLE_JSON_STORAGE_BLOB_PLAIN: + return cursor.var(DB_TYPE_CLOB, arraysize=arraysize, inconverter=json_converter_in_clob) return cursor.var(DB_TYPE_BLOB, arraysize=arraysize, inconverter=json_converter_in_blob) @@ -246,25 +255,3 @@ def _output_type_handler(cursor: "Cursor | AsyncCursor", metadata: Any) -> Any: return cursor.var(type_code, arraysize=cursor.arraysize, outconverter=json_converter_out_clob) return None - - -def _chained_input_handler( - inner: Any, fallback: "Any | None", cursor: "Cursor | AsyncCursor", value: Any, arraysize: int -) -> Any: - """Run ``inner`` input handler, falling back to ``fallback`` when it abstains.""" - result = inner(cursor, value, arraysize) - if result is not None: - return result - if fallback is not None: - return fallback(cursor, value, arraysize) - return None - - -def _chained_output_handler(inner: Any, fallback: "Any | None", cursor: "Cursor | AsyncCursor", metadata: Any) -> Any: - """Run ``inner`` output handler, falling back to ``fallback`` when it abstains.""" - result = inner(cursor, metadata) - if result is not None: - return result - if fallback is not None: - return fallback(cursor, metadata) - return None diff --git a/sqlspec/adapters/oracledb/_param_types.py b/sqlspec/adapters/oracledb/_param_types.py index 34017e726..ffc534f5b 100644 --- a/sqlspec/adapters/oracledb/_param_types.py +++ b/sqlspec/adapters/oracledb/_param_types.py @@ -45,3 +45,21 @@ class OracleJson: def __init__(self, value: Any) -> None: self.value = value + + +def is_json_payload(value: Any) -> bool: + """Return True if the value should be claimed by the JSON input handler. + + OracleJson, dict, empty list/tuple, and list/tuple containing any non-numeric + element are claimed. Non-empty sequences where all elements are int/float + (excluding bool) are reserved for the vector handler unless wrapped in OracleJson. + """ + if isinstance(value, OracleJson): + return True + if isinstance(value, dict): + return True + if isinstance(value, (list, tuple)): + if not value: + return True + return not all(isinstance(v, (int, float)) and not isinstance(v, bool) for v in value) + return False diff --git a/sqlspec/adapters/oracledb/_uuid_handlers.py b/sqlspec/adapters/oracledb/_uuid_handlers.py index 49da1484f..1037efd28 100644 --- a/sqlspec/adapters/oracledb/_uuid_handlers.py +++ b/sqlspec/adapters/oracledb/_uuid_handlers.py @@ -5,7 +5,7 @@ """ import uuid -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Final from sqlspec.adapters.oracledb._json_handlers import chain_input_handler, chain_output_handler from sqlspec.adapters.oracledb._typing import DB_TYPE_RAW @@ -31,6 +31,7 @@ UUID_BINARY_SIZE = 16 +_MIN_FETCH_INFO_LEN: Final = 4 def uuid_converter_in(value: uuid.UUID) -> bytes: @@ -131,7 +132,11 @@ def _output_type_handler(cursor: "Cursor | AsyncCursor", metadata: Any) -> Any: Returns: Cursor variable with UUID converter if column is RAW(16), None otherwise. """ - _name, type_code, _display_size, internal_size, _precision, _scale, _null_ok = metadata + type_code = getattr(metadata, "type_code", None) + internal_size = getattr(metadata, "internal_size", None) + if type_code is None and isinstance(metadata, (list, tuple)) and len(metadata) >= _MIN_FETCH_INFO_LEN: + type_code = metadata[1] + internal_size = metadata[3] if type_code is DB_TYPE_RAW and internal_size == UUID_BINARY_SIZE: return cursor.var(type_code, arraysize=cursor.arraysize, outconverter=uuid_converter_out) diff --git a/sqlspec/adapters/oracledb/_vector_handlers.py b/sqlspec/adapters/oracledb/_vector_handlers.py index f1ac86679..88564e8a0 100644 --- a/sqlspec/adapters/oracledb/_vector_handlers.py +++ b/sqlspec/adapters/oracledb/_vector_handlers.py @@ -15,7 +15,7 @@ from sqlspec.adapters.oracledb._json_handlers import chain_input_handler, chain_output_handler from sqlspec.adapters.oracledb._typing import DB_TYPE_VECTOR, OracleSparseVector -from sqlspec.adapters.oracledb.data_dictionary import resolve_oracle_connection_major +from sqlspec.adapters.oracledb.core import resolve_oracle_connection_major from sqlspec.typing import NUMPY_INSTALLED from sqlspec.utils.logging import get_logger from sqlspec.utils.module_loader import import_optional diff --git a/sqlspec/adapters/oracledb/core.py b/sqlspec/adapters/oracledb/core.py index 9aa1ce70e..2e14b75bf 100644 --- a/sqlspec/adapters/oracledb/core.py +++ b/sqlspec/adapters/oracledb/core.py @@ -16,13 +16,66 @@ from pathlib import Path from typing import TYPE_CHECKING, Any, Final, Protocol, cast -from sqlspec.adapters.oracledb._json_handlers import is_json_payload -from sqlspec.adapters.oracledb._param_types import OracleBlob, OracleClob, OracleJson +from sqlspec.adapters.oracledb._param_types import OracleBlob, OracleClob, OracleJson, is_json_payload from sqlspec.adapters.oracledb._typing import DB_TYPE_BLOB, DB_TYPE_CLOB, oracledb_module -from sqlspec.adapters.oracledb.data_dictionary import ( - resolve_oracle_connection_compatible_major, - resolve_oracle_connection_major, -) +from sqlspec.data_dictionary.dialects.oracle import parse_oracle_version_components + + +class OracleVersionCache: + """Pool-scoped cache for the resolved Oracle server version. + + One instance is created per config and shared with every driver spawned from + that config's pool, so the server version is resolved once per pool lifetime + rather than once per acquired session. The cache lives and dies with the + pool; a server upgrade under a live pool is out of scope and picked up on the + next pool restart. + """ + + __slots__ = ("resolved", "version") + + def __init__(self) -> None: + self.resolved: bool = False + self.version: Any | None = None + + def reset(self) -> None: + """Clear the cached version so the next resolution re-queries the server.""" + self.resolved = False + self.version = None + + +def resolve_oracle_connection_major(connection: Any, version_cache: "OracleVersionCache | None" = None) -> "int | None": + """Resolve an Oracle server major without issuing a metadata query. + + Pool-scoped data-dictionary metadata is authoritative when it has already + been resolved. Otherwise, use the connection callback cache and finally + python-oracledb's connection ``version`` string. The final fallback matters + for reacquired pool wrappers that do not retain dynamic attributes. + """ + if version_cache is not None and version_cache.resolved and version_cache.version is not None: + return int(version_cache.version.major) + cached_major = getattr(connection, "_sqlspec_oracle_major", None) + if isinstance(cached_major, int) and not isinstance(cached_major, bool): + return cached_major + version = getattr(connection, "version", None) + if not version: + return None + components = parse_oracle_version_components(str(version)) + return components[0] if components is not None else None + + +def resolve_oracle_connection_compatible_major( + connection: Any, version_cache: "OracleVersionCache | None" = None +) -> "int | None": + """Resolve an Oracle server compatible major version without issuing a metadata query.""" + if version_cache is not None and version_cache.resolved and version_cache.version is not None: + comp = version_cache.version.compatible_major + return int(comp) if comp is not None else None + cached_compatible = getattr(connection, "_sqlspec_oracle_compatible_major", None) + if isinstance(cached_compatible, int) and not isinstance(cached_compatible, bool): + return cached_compatible + return None + + from sqlspec.adapters.oracledb.type_converter import OracleOutputConverter from sqlspec.core import ( DriverParameterProfile, @@ -33,7 +86,11 @@ create_sql_result, ) from sqlspec.core.sqlcommenter import SQLCommenterContext -from sqlspec.data_dictionary.dialects.oracle import ORACLE_JSON_STORAGE_NATIVE, resolve_oracle_json_storage +from sqlspec.data_dictionary.dialects.oracle import ( + ORACLE_JSON_STORAGE_NATIVE, + parse_oracle_version_components, + resolve_oracle_json_storage, +) from sqlspec.driver import rows_to_dicts from sqlspec.exceptions import ( CheckViolationError, @@ -94,6 +151,7 @@ "OracleSessionContextConfig", "OracleSyncStreamSource", "OracleTopologyConfig", + "OracleVersionCache", "apply_driver_features", "apply_oracle_session_trace_attributes", "apply_wallet_ssl_context", @@ -141,6 +199,8 @@ "read_oracle_lob_sync", "release_extracted_wallet_dir", "resolve_dbms_cloud_source_uri", + "resolve_oracle_connection_compatible_major", + "resolve_oracle_connection_major", "resolve_row_metadata", "resolve_rowcount", "retain_extracted_wallet_dir", diff --git a/sqlspec/adapters/oracledb/data_dictionary.py b/sqlspec/adapters/oracledb/data_dictionary.py index 9b06a91db..ecaf98fc6 100644 --- a/sqlspec/adapters/oracledb/data_dictionary.py +++ b/sqlspec/adapters/oracledb/data_dictionary.py @@ -2,10 +2,16 @@ from collections import defaultdict from enum import StrEnum +from inspect import isawaitable from typing import TYPE_CHECKING, Any, ClassVar, cast from mypy_extensions import mypyc_attr +from sqlspec.adapters.oracledb.core import ( + OracleVersionCache, + resolve_oracle_connection_compatible_major, + resolve_oracle_connection_major, +) from sqlspec.data_dictionary import ( ColumnMetadata, DDLResult, @@ -32,7 +38,6 @@ from sqlspec.data_dictionary.dialects.oracle import ( extract_oracle_version_value, list_oracle_available_features, - merge_oracle_table_lists, oracle_supports_json_blob, oracle_supports_native_json, oracle_supports_oson_blob, @@ -170,65 +175,11 @@ class JSONStorageType(StrEnum): BLOB_PLAIN = "blob_plain" -class OracleVersionCache: - """Pool-scoped cache for the resolved Oracle server version. - - One instance is created per config and shared with every driver spawned from - that config's pool, so the server version is resolved once per pool lifetime - rather than once per acquired session. The cache lives and dies with the - pool; a server upgrade under a live pool is out of scope and picked up on the - next pool restart. - """ - - __slots__ = ("resolved", "version") - - def __init__(self) -> None: - self.resolved: bool = False - self.version: OracleVersionInfo | None = None - - def reset(self) -> None: - """Clear the cached version so the next resolution re-queries the server.""" - self.resolved = False - self.version = None - - def storage_type_from_version(version_info: "OracleVersionInfo | None") -> JSONStorageType: """Public alias for :func:`_storage_type_from_version`.""" return _storage_type_from_version(version_info) -def resolve_oracle_connection_major(connection: Any, version_cache: "OracleVersionCache | None" = None) -> "int | None": - """Resolve an Oracle server major without issuing a metadata query. - - Pool-scoped data-dictionary metadata is authoritative when it has already - been resolved. Otherwise, use the connection callback cache and finally - python-oracledb's connection ``version`` string. The final fallback matters - for reacquired pool wrappers that do not retain dynamic attributes. - """ - if version_cache is not None and version_cache.resolved and version_cache.version is not None: - return version_cache.version.major - cached_major = getattr(connection, "_sqlspec_oracle_major", None) - if isinstance(cached_major, int) and not isinstance(cached_major, bool): - return cached_major - version = getattr(connection, "version", None) - if not version: - return None - components = parse_oracle_version_components(str(version)) - return components[0] if components is not None else None - - -def resolve_oracle_connection_compatible_major( - connection: Any, version_cache: "OracleVersionCache | None" = None -) -> "int | None": - """Resolve an Oracle server compatible major version without issuing a metadata query.""" - if version_cache is not None and version_cache.resolved and version_cache.version is not None: - return version_cache.version.compatible_major - cached_compatible = getattr(connection, "_sqlspec_oracle_compatible_major", None) - if isinstance(cached_compatible, int) and not isinstance(cached_compatible, bool): - return cached_compatible - return None - - def _merge_index_expressions( indexes: "list[IndexMetadata]", expr_rows: "list[dict[str, Any]]", schema_name: "str | None" = None ) -> "list[IndexMetadata]": @@ -467,13 +418,9 @@ def get_tables(self, driver: "OracleSyncDriver", schema: "str | None" = None) -> """Get tables sorted by dependency order with full coverage.""" schema_name = self.resolve_schema(schema) self._log_schema_introspect(driver, schema_name=schema_name, table_name=None, operation="tables") - ordered_rows = driver.select( + return driver.select( self.get_query("tables", "by_schema"), schema_name=schema_name, table_name=None, schema_type=TableMetadata ) - all_rows = driver.select( - self.get_query("tables", "names_by_schema"), schema_name=schema_name, schema_type=TableMetadata - ) - return merge_oracle_table_lists(ordered_rows, all_rows) def get_columns( self, driver: "OracleSyncDriver", table: "str | None" = None, schema: "str | None" = None @@ -892,13 +839,9 @@ async def get_tables(self, driver: "OracleAsyncDriver", schema: "str | None" = N """Get tables sorted by dependency order with full coverage.""" schema_name = self.resolve_schema(schema) self._log_schema_introspect(driver, schema_name=schema_name, table_name=None, operation="tables") - ordered_rows = await driver.select( + return await driver.select( self.get_query("tables", "by_schema"), schema_name=schema_name, table_name=None, schema_type=TableMetadata ) - all_rows = await driver.select( - self.get_query("tables", "names_by_schema"), schema_name=schema_name, schema_type=TableMetadata - ) - return merge_oracle_table_lists(ordered_rows, all_rows) async def get_columns( self, driver: "OracleAsyncDriver", table: "str | None" = None, schema: "str | None" = None @@ -1037,11 +980,19 @@ async def _get_native_object_ddl( ddl_value = await driver.select_value( query.sql, object_type=normalized_type, object_name=object_name, owner=schema ) + if ddl_value is not None: + read = getattr(ddl_value, "read", None) + if callable(read): + res = read() + if isawaitable(res): + ddl_value = await res + else: + ddl_value = res return _oracle_ddl_metadata_result( object_type=normalized_type, object_name=object_name, owner=schema, - ddl_text=_coerce_oracle_ddl_text(ddl_value), + ddl_text=str(ddl_value) if ddl_value is not None else None, ) def _build_version_info( diff --git a/tests/unit/adapters/test_oracledb/test_data_dictionary_oracledb.py b/tests/unit/adapters/test_oracledb/test_data_dictionary_oracledb.py index 7868bea48..2f25eb2f1 100644 --- a/tests/unit/adapters/test_oracledb/test_data_dictionary_oracledb.py +++ b/tests/unit/adapters/test_oracledb/test_data_dictionary_oracledb.py @@ -618,3 +618,42 @@ async def select(self, query: Any, **kwargs: Any) -> list[dict[str, Any]]: assert indexes[0]["index_name"] == "IDX_EMP_UPPER_EMAIL" assert indexes[0]["expressions"] == "UPPER(EMAIL)" assert indexes[0]["expression"] == "UPPER(EMAIL)" + + +@pytest.mark.anyio +async def test_oracle_async_get_native_object_ddl_awaits_async_lob() -> None: + """Async _get_native_object_ddl awaits an AsyncLOB read() instead of stringifying the coroutine.""" + data_dictionary = OracledbAsyncDataDictionary() + + class _AsyncLOB: + async def read(self) -> str: + return "CREATE TABLE FOO (ID NUMBER)" + + class _AsyncDriverStub: + async def select_value(self, query: Any, **kwargs: Any) -> Any: + return _AsyncLOB() + + driver = cast("OracleAsyncDriver", _AsyncDriverStub()) + result = await data_dictionary._get_native_object_ddl( + driver, "FOO", schema="HR", object_type="TABLE", include_dependencies=False, redact=False + ) + assert result is not None + assert result.ddl == "CREATE TABLE FOO (ID NUMBER)" + assert " None: + """get_tables issues a single catalog query (tables.by_schema) without a redundant names query.""" + data_dictionary = OracledbSyncDataDictionary() + queries: list[str] = [] + + class _StubSyncDriver: + def select(self, query: Any, **kwargs: Any) -> list[dict[str, Any]]: + queries.append(str(query)) + return [{"table_name": "T1"}] + + driver = cast("OracleSyncDriver", _StubSyncDriver()) + tables = data_dictionary.get_tables(driver, schema="HR") + assert len(tables) == 1 + assert len(queries) == 1 + assert "names_by_schema" not in queries[0] diff --git a/tests/unit/adapters/test_oracledb/test_json_handlers.py b/tests/unit/adapters/test_oracledb/test_json_handlers.py index e8d81c651..e91ccf8db 100644 --- a/tests/unit/adapters/test_oracledb/test_json_handlers.py +++ b/tests/unit/adapters/test_oracledb/test_json_handlers.py @@ -154,16 +154,18 @@ def test_input_handler_routes_dict_to_db_type_blob_on_12c() -> None: cursor.var.assert_called_once_with(oracledb.DB_TYPE_BLOB, arraysize=2, inconverter=json_converter_in_blob) -def test_input_handler_routes_dict_to_db_type_blob_pre_12c() -> None: - """Dict bound on Oracle 11g (BLOB_PLAIN storage) should use DB_TYPE_BLOB.""" +def test_input_handler_routes_dict_to_db_type_clob_on_11g() -> None: + """Dict bound on Oracle 11g (<12c) uses DB_TYPE_CLOB with json_converter_in_clob.""" + import oracledb + cursor = _mock_cursor_with_major(11) cursor_var = Mock() cursor.var = Mock(return_value=cursor_var) - result = json_input_type_handler(cursor, {"foo": "bar"}, 2) + result = json_input_type_handler(cursor, {"foo": "bar"}, 4) assert result is cursor_var - cursor.var.assert_called_once_with(oracledb.DB_TYPE_BLOB, arraysize=2, inconverter=json_converter_in_blob) + cursor.var.assert_called_once_with(oracledb.DB_TYPE_CLOB, arraysize=4, inconverter=json_converter_in_clob) def test_input_handler_respects_compatible_major_on_21c() -> None: diff --git a/tests/unit/adapters/test_oracledb/test_oracle_version_unification.py b/tests/unit/adapters/test_oracledb/test_oracle_version_unification.py index aa4f41fdc..ab3295513 100644 --- a/tests/unit/adapters/test_oracledb/test_oracle_version_unification.py +++ b/tests/unit/adapters/test_oracledb/test_oracle_version_unification.py @@ -242,3 +242,19 @@ def test_oracle_events_store_compression_partitioning_config() -> None: assert "TABLESPACE event_data" in sql assert "PARTITION BY RANGE (available_at)" in sql assert "NUMTODSINTERVAL(1, ''DAY'')" in sql + + +def test_resolve_oracle_connection_major_from_core() -> None: + """resolve_oracle_connection_major and cache are importable and callable directly from core.""" + from sqlspec.adapters.oracledb.core import ( + OracleVersionCache, + resolve_oracle_connection_compatible_major, + resolve_oracle_connection_major, + ) + + conn = MagicMock() + conn._sqlspec_oracle_major = 23 + conn._sqlspec_oracle_compatible_major = 23 + cache = OracleVersionCache() + assert resolve_oracle_connection_major(conn, cache) == 23 + assert resolve_oracle_connection_compatible_major(conn, cache) == 23 diff --git a/tests/unit/adapters/test_oracledb/test_uuid_handlers.py b/tests/unit/adapters/test_oracledb/test_uuid_handlers.py index 3a5d326dc..c7b37d7e9 100644 --- a/tests/unit/adapters/test_oracledb/test_uuid_handlers.py +++ b/tests/unit/adapters/test_oracledb/test_uuid_handlers.py @@ -252,3 +252,25 @@ def test_uuid_handlers_use_generic_chained_wrappers() -> None: assert callable(connection.outputtypehandler) assert len(inspect.signature(connection.outputtypehandler).parameters) == 2 assert len(inspect.signature(connection.inputtypehandler).parameters) == 3 + + +def test_output_type_handler_with_fetch_info_object() -> None: + """_output_type_handler handles FetchInfo objects with attributes.""" + import oracledb + + from sqlspec.adapters.oracledb._uuid_handlers import _output_type_handler + + cursor = Mock() + cursor.arraysize = 32 + cursor_var = Mock() + cursor.var = Mock(return_value=cursor_var) + + class _FetchInfoStub: + def __init__(self) -> None: + self.name = "ID" + self.type_code = oracledb.DB_TYPE_RAW + self.internal_size = 16 + + result = _output_type_handler(cursor, _FetchInfoStub()) + assert result is cursor_var + cursor.var.assert_called_once_with(oracledb.DB_TYPE_RAW, arraysize=32, outconverter=uuid_converter_out) From 1460d91292817dfab8b94f8928f9b7355dfdf0ef Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 01:02:29 +0000 Subject: [PATCH 068/174] perf(oracledb): optimize hot-path allocations and lazy arrow batch streaming Cache fetch keyword arguments in build_fetch_kwargs, check ASCII bounds before encoding string parameters, key row metadata cache by column types and names tuple, precompute driver fetch kwargs and session tracing on instance slots, use central arrow_batch_to_rows for Arrow bulk loading, and lazily yield batches into RecordBatchReader when return_format is reader. --- sqlspec/adapters/oracledb/_json_handlers.py | 57 +---- sqlspec/adapters/oracledb/_uuid_handlers.py | 2 +- sqlspec/adapters/oracledb/_vector_handlers.py | 3 +- sqlspec/adapters/oracledb/core.py | 201 +++++++++++++----- sqlspec/adapters/oracledb/driver.py | 117 ++++++---- .../test_oracledb/test_arrow_batches.py | 29 +++ .../test_core_row_materialization.py | 10 +- .../test_oracledb/test_fetch_tuning.py | 20 ++ 8 files changed, 282 insertions(+), 157 deletions(-) diff --git a/sqlspec/adapters/oracledb/_json_handlers.py b/sqlspec/adapters/oracledb/_json_handlers.py index e8972735f..0e9795232 100644 --- a/sqlspec/adapters/oracledb/_json_handlers.py +++ b/sqlspec/adapters/oracledb/_json_handlers.py @@ -38,57 +38,7 @@ from functools import partial from typing import TYPE_CHECKING, Any - -def chain_input_handler(inner: Any, fallback: "Any | None") -> Any: - """Build an input type handler that chains ``inner`` to ``fallback``.""" - if fallback is not None and _has_handler(fallback, inner, _chained_input_handler): - return fallback - return partial(_chained_input_handler, inner, fallback) - - -def chain_output_handler(inner: Any, fallback: "Any | None") -> Any: - """Build an output type handler that chains ``inner`` to ``fallback``.""" - if fallback is not None and _has_handler(fallback, inner, _chained_output_handler): - return fallback - return partial(_chained_output_handler, inner, fallback) - - -def _has_handler(handler: Any, target_inner: Any, chain_function: Any) -> bool: - current = handler - while current is not None: - if current is target_inner: - return True - if not isinstance(current, partial) or current.func is not chain_function: - return False - inner, current = current.args - if inner is target_inner: - return True - return False - - -def _chained_input_handler( - inner: Any, fallback: "Any | None", cursor: Any, value: Any, arraysize: int -) -> Any: - """Run ``inner`` input handler, falling back to ``fallback`` when it abstains.""" - result = inner(cursor, value, arraysize) - if result is not None: - return result - if fallback is not None: - return fallback(cursor, value, arraysize) - return None - - -def _chained_output_handler(inner: Any, fallback: "Any | None", cursor: Any, metadata: Any) -> Any: - """Run ``inner`` output handler, falling back to ``fallback`` when it abstains.""" - result = inner(cursor, metadata) - if result is not None: - return result - if fallback is not None: - return fallback(cursor, metadata) - return None - - -from sqlspec.adapters.oracledb._param_types import OracleJson +from sqlspec.adapters.oracledb._param_types import OracleJson, is_json_payload from sqlspec.adapters.oracledb._typing import ( DB_TYPE_BLOB, DB_TYPE_CHAR, @@ -101,7 +51,8 @@ def _chained_output_handler(inner: Any, fallback: "Any | None", cursor: Any, met DB_TYPE_VARCHAR, ) from sqlspec.adapters.oracledb.core import ( - is_json_payload, + chain_input_handler, + chain_output_handler, resolve_oracle_connection_compatible_major, resolve_oracle_connection_major, ) @@ -124,7 +75,6 @@ def _chained_output_handler(inner: Any, fallback: "Any | None", cursor: Any, met "is_json_payload", "json_converter_in_blob", "json_converter_in_clob", - "json_converter_in_native", "json_converter_out_blob", "json_converter_out_clob", "json_converter_out_oson", @@ -133,6 +83,7 @@ def _chained_output_handler(inner: Any, fallback: "Any | None", cursor: Any, met "register_json_handlers", ) + _JSON_STRING_TYPE_CODES = (DB_TYPE_VARCHAR, DB_TYPE_CHAR, DB_TYPE_NVARCHAR, DB_TYPE_NCHAR) diff --git a/sqlspec/adapters/oracledb/_uuid_handlers.py b/sqlspec/adapters/oracledb/_uuid_handlers.py index 1037efd28..57896377d 100644 --- a/sqlspec/adapters/oracledb/_uuid_handlers.py +++ b/sqlspec/adapters/oracledb/_uuid_handlers.py @@ -7,8 +7,8 @@ import uuid from typing import TYPE_CHECKING, Any, Final -from sqlspec.adapters.oracledb._json_handlers import chain_input_handler, chain_output_handler from sqlspec.adapters.oracledb._typing import DB_TYPE_RAW +from sqlspec.adapters.oracledb.core import chain_input_handler, chain_output_handler from sqlspec.utils.logging import get_logger from sqlspec.utils.uuids import uuid_from_bytes diff --git a/sqlspec/adapters/oracledb/_vector_handlers.py b/sqlspec/adapters/oracledb/_vector_handlers.py index 88564e8a0..63f72fef3 100644 --- a/sqlspec/adapters/oracledb/_vector_handlers.py +++ b/sqlspec/adapters/oracledb/_vector_handlers.py @@ -13,9 +13,8 @@ import array from typing import TYPE_CHECKING, Any, cast -from sqlspec.adapters.oracledb._json_handlers import chain_input_handler, chain_output_handler from sqlspec.adapters.oracledb._typing import DB_TYPE_VECTOR, OracleSparseVector -from sqlspec.adapters.oracledb.core import resolve_oracle_connection_major +from sqlspec.adapters.oracledb.core import chain_input_handler, chain_output_handler, resolve_oracle_connection_major from sqlspec.typing import NUMPY_INSTALLED from sqlspec.utils.logging import get_logger from sqlspec.utils.module_loader import import_optional diff --git a/sqlspec/adapters/oracledb/core.py b/sqlspec/adapters/oracledb/core.py index 2e14b75bf..78f177459 100644 --- a/sqlspec/adapters/oracledb/core.py +++ b/sqlspec/adapters/oracledb/core.py @@ -13,12 +13,53 @@ import threading import zipfile from collections.abc import Callable, Iterable, Mapping, Sequence, Sized +from functools import partial from pathlib import Path from typing import TYPE_CHECKING, Any, Final, Protocol, cast from sqlspec.adapters.oracledb._param_types import OracleBlob, OracleClob, OracleJson, is_json_payload from sqlspec.adapters.oracledb._typing import DB_TYPE_BLOB, DB_TYPE_CLOB, oracledb_module -from sqlspec.data_dictionary.dialects.oracle import parse_oracle_version_components +from sqlspec.core import ( + DriverParameterProfile, + ParameterStyle, + StackResult, + StatementConfig, + build_statement_config_from_profile, + create_sql_result, +) +from sqlspec.core.sqlcommenter import SQLCommenterContext +from sqlspec.data_dictionary.dialects.oracle import ( + ORACLE_JSON_STORAGE_NATIVE, + parse_oracle_version_components, + resolve_oracle_json_storage, +) +from sqlspec.driver import rows_to_dicts +from sqlspec.exceptions import ( + CheckViolationError, + ConnectionTimeoutError, + DatabaseConnectionError, + DataError, + DeadlockError, + ForeignKeyViolationError, + ImproperConfigurationError, + IntegrityError, + NotNullViolationError, + OperationalError, + OperationCancelledError, + PermissionDeniedError, + SQLParsingError, + SQLSpecError, + TransactionError, + UniqueViolationError, +) +from sqlspec.observability import get_trace_context +from sqlspec.typing import NUMPY_INSTALLED +from sqlspec.utils.correlation import CorrelationContext +from sqlspec.utils.module_loader import import_optional +from sqlspec.utils.serializers import from_json, to_json +from sqlspec.utils.text import quote_identifier +from sqlspec.utils.type_converters import build_uuid_coercions +from sqlspec.utils.type_guards import has_rowcount, is_async_readable, is_readable class OracleVersionCache: @@ -76,48 +117,52 @@ def resolve_oracle_connection_compatible_major( return None -from sqlspec.adapters.oracledb.type_converter import OracleOutputConverter -from sqlspec.core import ( - DriverParameterProfile, - ParameterStyle, - StackResult, - StatementConfig, - build_statement_config_from_profile, - create_sql_result, -) -from sqlspec.core.sqlcommenter import SQLCommenterContext -from sqlspec.data_dictionary.dialects.oracle import ( - ORACLE_JSON_STORAGE_NATIVE, - parse_oracle_version_components, - resolve_oracle_json_storage, -) -from sqlspec.driver import rows_to_dicts -from sqlspec.exceptions import ( - CheckViolationError, - ConnectionTimeoutError, - DatabaseConnectionError, - DataError, - DeadlockError, - ForeignKeyViolationError, - ImproperConfigurationError, - IntegrityError, - NotNullViolationError, - OperationalError, - OperationCancelledError, - PermissionDeniedError, - SQLParsingError, - SQLSpecError, - TransactionError, - UniqueViolationError, -) -from sqlspec.observability import get_trace_context -from sqlspec.typing import NUMPY_INSTALLED -from sqlspec.utils.correlation import CorrelationContext -from sqlspec.utils.module_loader import import_optional -from sqlspec.utils.serializers import from_json, to_json -from sqlspec.utils.text import quote_identifier -from sqlspec.utils.type_converters import build_uuid_coercions -from sqlspec.utils.type_guards import has_rowcount, is_async_readable, is_readable +def chain_input_handler(inner: Any, fallback: "Any | None") -> Any: + """Build an input type handler that chains ``inner`` to ``fallback``.""" + if fallback is not None and _has_handler(fallback, inner, _chained_input_handler): + return fallback + return partial(_chained_input_handler, inner, fallback) + + +def chain_output_handler(inner: Any, fallback: "Any | None") -> Any: + """Build an output type handler that chains ``inner`` to ``fallback``.""" + if fallback is not None and _has_handler(fallback, inner, _chained_output_handler): + return fallback + return partial(_chained_output_handler, inner, fallback) + + +def _has_handler(handler: Any, target_inner: Any, chain_function: Any) -> bool: + current = handler + while current is not None: + if current is target_inner: + return True + if not isinstance(current, partial) or current.func is not chain_function: + return False + inner, current = current.args + if inner is target_inner: + return True + return False + + +def _chained_input_handler(inner: Any, fallback: "Any | None", cursor: Any, value: Any, arraysize: int) -> Any: + """Run ``inner`` input handler, falling back to ``fallback`` when it abstains.""" + result = inner(cursor, value, arraysize) + if result is not None: + return result + if fallback is not None: + return fallback(cursor, value, arraysize) + return None + + +def _chained_output_handler(inner: Any, fallback: "Any | None", cursor: Any, metadata: Any) -> Any: + """Run ``inner`` output handler, falling back to ``fallback`` when it abstains.""" + result = inner(cursor, metadata) + if result is not None: + return result + if fallback is not None: + return fallback(cursor, metadata) + return None + if TYPE_CHECKING: from collections.abc import Awaitable @@ -172,6 +217,8 @@ def resolve_oracle_connection_compatible_major( "build_tns_description", "build_truncate_statement", "build_wallet_ssl_context_if_needed", + "chain_input_handler", + "chain_output_handler", "coerce_large_parameters_async", "coerce_large_parameters_sync", "coerce_many_parameters_async", @@ -324,7 +371,6 @@ def resolve_oracle_connection_compatible_major( "WITH", }) _VERSION_COMPONENTS: int = 3 -TYPE_CONVERTER = OracleOutputConverter() _LOB_TYPE_NAME_MARKERS: "tuple[str, ...]" = ("LOB", "BFILE") _SCALAR_PASSTHROUGH_TYPES: "tuple[type[Any], ...]" = (bool, int, float, str, bytes, bytearray, type(None)) _BIND_PASSTHROUGH_TYPES: "tuple[type[Any], ...]" = (bool, int, float, type(None)) @@ -843,15 +889,24 @@ def supports_df_batches(connection: object) -> bool: return hasattr(connection, "fetch_df_batches") +_FETCH_KWARGS_CACHE: dict[tuple[object, object], dict[str, object]] = {} + + def build_fetch_kwargs(driver_features: "dict[str, Any]") -> "dict[str, object]": """Build per-statement Oracle fetch keyword arguments.""" - fetch_kwargs: dict[str, object] = {} fetch_lobs = driver_features.get("fetch_lobs") + fetch_decimals = driver_features.get("fetch_decimals") + cache_key = (fetch_lobs, fetch_decimals) + cached = _FETCH_KWARGS_CACHE.get(cache_key) + if cached is not None: + return cached + + fetch_kwargs: dict[str, object] = {} if fetch_lobs is not None: fetch_kwargs["fetch_lobs"] = fetch_lobs - fetch_decimals = driver_features.get("fetch_decimals") if fetch_decimals is not None: fetch_kwargs["fetch_decimals"] = fetch_decimals + _FETCH_KWARGS_CACHE[cache_key] = fetch_kwargs return fetch_kwargs @@ -1308,6 +1363,20 @@ def apply_driver_features( return statement_config, features +def _extract_column_cache_key(description: "list[Any]") -> "tuple[tuple[object, object], ...]": + key_items: list[tuple[object, object]] = [] + for column in description: + try: + col_name = column[0] + type_code = column[1] + except (TypeError, IndexError, KeyError): + col_name = getattr(column, "name", None) + type_code = getattr(column, "type_code", None) + type_identifier = getattr(type_code, "name", None) or type_code + key_items.append((col_name, type_identifier)) + return tuple(key_items) + + def resolve_row_metadata( description: "list[Any] | None", driver_features: "dict[str, Any]", cache: "dict[Any, tuple[Any, list[str], bool]]" ) -> "tuple[list[str], bool]": @@ -1324,9 +1393,9 @@ def resolve_row_metadata( if not description: return [], False - cache_key = id(description) + cache_key = _extract_column_cache_key(description) cached = cache.get(cache_key) - if cached is not None and cached[0] is description: + if cached is not None: return cached[1], cached[2] column_names: list[str] = [] @@ -1408,9 +1477,11 @@ def _start(self) -> None: self._cursor = cursor cursor.arraysize = self._chunk_size cursor.prefetchrows = self._chunk_size - fetch_kwargs = build_fetch_kwargs(self._driver.driver_features) + fetch_kwargs = getattr(self._driver, "_fetch_kwargs", None) or build_fetch_kwargs( + self._driver.driver_features + ) if self._fetch_lobs is not None: - fetch_kwargs["fetch_lobs"] = self._fetch_lobs + fetch_kwargs = {**fetch_kwargs, "fetch_lobs": self._fetch_lobs} parameters = coerce_large_parameters_sync( self._driver.connection, self._parameters, @@ -1545,9 +1616,11 @@ async def _start(self) -> None: self._cursor = cursor cursor.arraysize = self._chunk_size cursor.prefetchrows = self._chunk_size - fetch_kwargs = build_fetch_kwargs(self._driver.driver_features) + fetch_kwargs = getattr(self._driver, "_fetch_kwargs", None) or build_fetch_kwargs( + self._driver.driver_features + ) if self._fetch_lobs is not None: - fetch_kwargs["fetch_lobs"] = self._fetch_lobs + fetch_kwargs = {**fetch_kwargs, "fetch_lobs": self._fetch_lobs} parameters = await coerce_large_parameters_async( self._driver.connection, self._parameters, @@ -3213,6 +3286,16 @@ def _parse_version_tuple(version: str) -> "tuple[int, int, int]": ORACLEDB_VERSION: "tuple[int, int, int]" = _parse_version_tuple(oracledb_module.__version__) +def _string_exceeds_byte_limit(value: str, limit: int) -> bool: + """Return whether a string's UTF-8 byte length exceeds the given limit.""" + char_len = len(value) + if char_len > limit: + return True + if value.isascii(): + return False + return len(value.encode("utf-8")) > limit + + def _parameter_values_need_coercion( values: "Iterable[Any]", *, @@ -3226,7 +3309,7 @@ def _parameter_values_need_coercion( if value_type in _BIND_PASSTHROUGH_TYPES: continue if value_type is str: - if len(value.encode("utf-8")) > varchar2_byte_limit: + if _string_exceeds_byte_limit(value, varchar2_byte_limit): return True continue if value_type is bytes or value_type is bytearray: @@ -3244,7 +3327,7 @@ def _parameter_values_need_coercion( return True continue if isinstance(value, str): - if len(value.encode("utf-8")) > varchar2_byte_limit: + if _string_exceeds_byte_limit(value, varchar2_byte_limit): return True continue if isinstance(value, (bytes, bytearray)): @@ -3397,7 +3480,7 @@ def _coerce_value_sync( if value_type in _BIND_PASSTHROUGH_TYPES: return value if value_type is str: - if len(value.encode("utf-8")) > varchar2_byte_limit: + if _string_exceeds_byte_limit(value, varchar2_byte_limit): return connection.createlob(clob_type, value) return value if value_type is bytes: @@ -3426,7 +3509,7 @@ def _coerce_value_sync( if json_binding_state.uses_blob(): return connection.createlob(blob_type, to_json(value.value, as_bytes=True)) return value - if isinstance(value, str) and len(value.encode("utf-8")) > varchar2_byte_limit: + if isinstance(value, str) and _string_exceeds_byte_limit(value, varchar2_byte_limit): return connection.createlob(clob_type, value) if isinstance(value, (bytes, bytearray)) and len(value) > raw_byte_limit: return connection.createlob(blob_type, bytes(value)) @@ -3450,7 +3533,7 @@ async def _coerce_value_async( if value_type in _BIND_PASSTHROUGH_TYPES: return value if value_type is str: - if len(value.encode("utf-8")) > varchar2_byte_limit: + if _string_exceeds_byte_limit(value, varchar2_byte_limit): return await connection.createlob(clob_type, value) return value if value_type is bytes: @@ -3479,7 +3562,7 @@ async def _coerce_value_async( if json_binding_state.uses_blob(): return await connection.createlob(blob_type, to_json(value.value, as_bytes=True)) return value - if isinstance(value, str) and len(value.encode("utf-8")) > varchar2_byte_limit: + if isinstance(value, str) and _string_exceeds_byte_limit(value, varchar2_byte_limit): return await connection.createlob(clob_type, value) if isinstance(value, (bytes, bytearray)) and len(value) > raw_byte_limit: return await connection.createlob(blob_type, bytes(value)) @@ -3588,7 +3671,7 @@ async def _coerce_async_row_values(row: "tuple[Any, ...]") -> "tuple[Any, ...]": continue if is_readable(value) or is_async_readable(value): - processed_value = await TYPE_CONVERTER.process_lob(value) + processed_value = await read_oracle_lob_async(value) if coerced_values is None: if processed_value is value: continue diff --git a/sqlspec/adapters/oracledb/driver.py b/sqlspec/adapters/oracledb/driver.py index 2385e5fac..9b9ada15b 100644 --- a/sqlspec/adapters/oracledb/driver.py +++ b/sqlspec/adapters/oracledb/driver.py @@ -1,6 +1,7 @@ """Oracle Database driver implementation.""" import logging +from itertools import chain from pathlib import Path from time import perf_counter from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, Protocol, cast, overload @@ -92,7 +93,7 @@ from sqlspec.utils.uuids import uuid4 if TYPE_CHECKING: - from collections.abc import Mapping, Sequence + from collections.abc import Iterator, Mapping, Sequence from sqlspec.builder import QueryBuilder from sqlspec.core import ArrowResult, Statement, StatementFilter @@ -160,14 +161,8 @@ class OraclePipelineMixin: __slots__ = () - def _session_tracing_enabled(self) -> bool: - driver = cast("OraclePipelineDriver", self) - return bool( - driver.driver_features.get("enable_session_tracing", False) or driver.statement_config.enable_sqlcommenter - ) - def _apply_session_tracing(self, connection: Any = None) -> None: - if not self._session_tracing_enabled(): + if not getattr(self, "_session_tracing_enabled", False): return driver = cast("OraclePipelineDriver", self) target_conn = connection if connection is not None else getattr(self, "connection", None) @@ -190,7 +185,7 @@ def _apply_session_tracing(self, connection: Any = None) -> None: def _format_wire_sql(self, sql: str, config: "StatementConfig | None" = None) -> str: driver = cast("OraclePipelineDriver", self) effective_config = config or driver.statement_config - if self._session_tracing_enabled(): + if getattr(self, "_session_tracing_enabled", False): return strip_sqlcomment(sql) if not effective_config.enable_sqlcommenter: return sql @@ -202,7 +197,7 @@ def _compiled_sql( compiled_sql, prepared_parameters = super()._compiled_sql( # type: ignore[misc] statement, statement_config, flatten_single_parameters=flatten_single_parameters ) - if self._session_tracing_enabled(): + if getattr(self, "_session_tracing_enabled", False): return strip_sqlcomment(compiled_sql), prepared_parameters return compiled_sql, prepared_parameters @@ -382,10 +377,12 @@ class OracleSyncDriver(OraclePipelineMixin, SyncDriverAdapterBase): __slots__ = ( "_data_dictionary", + "_fetch_kwargs", "_oracle_version_cache", "_pipeline_support", "_pipeline_support_reason", "_row_metadata_cache", + "_session_tracing_enabled", "_transaction_active", ) dialect = "oracle" @@ -408,6 +405,10 @@ def __init__( self._oracle_version_cache: OracleVersionCache | None = None self._row_metadata_cache: dict[int, tuple[Any, list[str], bool]] = {} self._transaction_active = False + self._fetch_kwargs = build_fetch_kwargs(self.driver_features) + self._session_tracing_enabled = bool( + self.driver_features.get("enable_session_tracing", False) or self.statement_config.enable_sqlcommenter + ) @property def data_dictionary(self) -> "OracledbSyncDataDictionary": @@ -455,7 +456,7 @@ def dispatch_execute(self, cursor: Any, statement: "SQL") -> "ExecutionResult": cursor, explained_sql, prepared_parameters, display_cursor=display_cursor, xplan_format=xplan_format ) - cursor.execute(sql, prepared_parameters or {}, **build_fetch_kwargs(self.driver_features)) + cursor.execute(sql, prepared_parameters or {}, **self._fetch_kwargs) is_select_like = statement.returns_rows() or self._should_force_select(statement, cursor) @@ -489,7 +490,7 @@ def _execute_cache_hit_with_cursor( with self.with_cursor(self.connection) as cursor: wire_sql = self._format_wire_sql(cached.compiled_sql) returns_rows = cached.operation_profile.returns_rows - fetch_kwargs = build_fetch_kwargs(self.driver_features) + fetch_kwargs = self._fetch_kwargs parameters = coerce_large_parameters_sync( self.connection, params, @@ -688,7 +689,7 @@ def with_cursor(self, connection: OracleSyncConnection) -> OracleSyncCursor: connection, arraysize=self.driver_features.get("arraysize"), prefetchrows=self.driver_features.get("prefetchrows"), - on_enter=self._apply_session_tracing if self._session_tracing_enabled() else None, + on_enter=self._apply_session_tracing if self._session_tracing_enabled else None, ) @overload @@ -812,7 +813,19 @@ def select_to_arrow( exc_handler = self.handle_database_exceptions() try: self._apply_session_tracing(self.connection) - if return_format in {"batches", "reader"} and supports_df_batches(self.connection): + if return_format == "reader" and supports_df_batches(self.connection): + batch_generator = self._iter_arrow_record_batches(sql, prepared_parameters, batch_size, arrow_schema) + try: + first_batch = next(batch_generator) + except StopIteration: + empty_schema = arrow_schema if arrow_schema is not None else pa.schema([]) + reader = pa.RecordBatchReader.from_batches(empty_schema, ()) + return create_arrow_result(prepared_statement, reader, rows_affected=0) + schema = arrow_schema if arrow_schema is not None else first_batch.schema + lazy_reader = pa.RecordBatchReader.from_batches(schema, chain((first_batch,), batch_generator)) + return create_arrow_result(prepared_statement, lazy_reader, rows_affected=-1) + + if return_format == "batches" and supports_df_batches(self.connection): batches_result = cast( "tuple[list[ArrowRecordBatch], ArrowSchema | None]", self._run_with_exception_handler( @@ -824,16 +837,9 @@ def select_to_arrow( arrow_schema, ), ) - record_batches, batch_schema = batches_result + record_batches, _ = batches_result rows_affected = sum(batch.num_rows for batch in record_batches) - if batch_schema is None: - batch_schema = pa.schema([]) - arrow_data: object - if return_format == "reader": - arrow_data = pa.RecordBatchReader.from_batches(batch_schema, record_batches) - else: - arrow_data = record_batches - return create_arrow_result(prepared_statement, arrow_data, rows_affected=rows_affected) + return create_arrow_result(prepared_statement, record_batches, rows_affected=rows_affected) oracle_df = self._run_with_exception_handler( exc_handler, self._execute_arrow_dataframe, sql, prepared_parameters, batch_size @@ -994,7 +1000,7 @@ def _truncate() -> None: def _insert_batches() -> None: with self.with_cursor(self.connection) as cursor: for batch in arrow_table.to_batches(max_chunksize=5000): - batch_rows = [tuple(row[col] for col in columns) for row in batch.to_pylist()] + batch_rows = self._arrow_batch_to_rows(batch, columns) coerced_rows = coerce_many_parameters_sync( self.connection, batch_rows, @@ -1142,10 +1148,10 @@ def _dispatch_explain_plan( Execution result carrying the formatted plan rows. """ if display_cursor: - cursor.execute(explained_sql, parameters or {}, **build_fetch_kwargs(self.driver_features)) + cursor.execute(explained_sql, parameters or {}, **self._fetch_kwargs) if cursor.description is not None: cursor.fetchall() - cursor.execute(_format_display_cursor_sql(xplan_format), {}, **build_fetch_kwargs(self.driver_features)) + cursor.execute(_format_display_cursor_sql(xplan_format), {}, **self._fetch_kwargs) fetched_data = cursor.fetchall() column_names, requires_lob_coercion = self._resolve_row_metadata(cursor.description) data, column_names = collect_sync_rows( @@ -1169,7 +1175,7 @@ def _dispatch_explain_plan( display_sql = PLAN_DISPLAY_SQL if xplan_format == "TYPICAL" else _format_plan_display_sql(xplan_format) try: - cursor.execute(display_sql, {"statement_id": statement_id}, **build_fetch_kwargs(self.driver_features)) + cursor.execute(display_sql, {"statement_id": statement_id}, **self._fetch_kwargs) fetched_data = cursor.fetchall() column_names, requires_lob_coercion = self._resolve_row_metadata(cursor.description) data, column_names = collect_sync_rows( @@ -1222,6 +1228,35 @@ def _execute_arrow_dataframe(self, sql: str, parameters: "Any", batch_size: int statement=sql, parameters=params, arraysize=batch_size or 1000, **retry_kwargs ) + def _iter_arrow_record_batches( + self, + sql: str, + parameters: "StatementParameters | None", + batch_size: int | None, + arrow_schema: "ArrowSchema | None", + ) -> "Iterator[ArrowRecordBatch]": + params = parameters if parameters is not None else [] + fetch_kwargs = build_arrow_fetch_kwargs(self.driver_features, arrow_schema=arrow_schema) + try: + oracle_dfs = self.connection.fetch_df_batches( + statement=sql, parameters=params, size=batch_size or 1000, **fetch_kwargs + ) + except TypeError as exc: + retry_kwargs = _retry_arrow_without_fetch_lobs(exc, fetch_kwargs) + if retry_kwargs is None: + raise + oracle_dfs = self.connection.fetch_df_batches( + statement=sql, parameters=params, size=batch_size or 1000, **retry_kwargs + ) + for oracle_df in oracle_dfs: + batch_table = coerce_oracle_df_to_arrow_table(oracle_df, pa) + column_names = normalize_column_names(batch_table.column_names, self.driver_features) + if column_names != batch_table.column_names: + batch_table = batch_table.rename_columns(column_names) + if arrow_schema is not None: + batch_table = batch_table.cast(arrow_schema) + yield from batch_table.to_batches() + def _fetch_arrow_record_batches( self, sql: str, @@ -1232,7 +1267,7 @@ def _fetch_arrow_record_batches( params = parameters if parameters is not None else [] record_batches: list[ArrowRecordBatch] = [] batch_schema: ArrowSchema | None = None - fetch_kwargs = build_arrow_fetch_kwargs(self.driver_features) + fetch_kwargs = build_arrow_fetch_kwargs(self.driver_features, arrow_schema=arrow_schema) try: oracle_dfs = self.connection.fetch_df_batches( statement=sql, parameters=params, size=batch_size or 1000, **fetch_kwargs @@ -1279,10 +1314,12 @@ class OracleAsyncDriver(OraclePipelineMixin, AsyncDriverAdapterBase): __slots__ = ( "_data_dictionary", + "_fetch_kwargs", "_oracle_version_cache", "_pipeline_support", "_pipeline_support_reason", "_row_metadata_cache", + "_session_tracing_enabled", "_transaction_active", ) dialect = "oracle" @@ -1305,6 +1342,10 @@ def __init__( self._oracle_version_cache: OracleVersionCache | None = None self._row_metadata_cache: dict[int, tuple[Any, list[str], bool]] = {} self._transaction_active = False + self._fetch_kwargs = build_fetch_kwargs(self.driver_features) + self._session_tracing_enabled = bool( + self.driver_features.get("enable_session_tracing", False) or self.statement_config.enable_sqlcommenter + ) @property def data_dictionary(self) -> "OracledbAsyncDataDictionary": @@ -1352,7 +1393,7 @@ async def dispatch_execute(self, cursor: Any, statement: "SQL") -> "ExecutionRes cursor, explained_sql, prepared_parameters, display_cursor=display_cursor, xplan_format=xplan_format ) - await cursor.execute(sql, prepared_parameters or {}, **build_fetch_kwargs(self.driver_features)) + await cursor.execute(sql, prepared_parameters or {}, **self._fetch_kwargs) is_select_like = statement.returns_rows() or self._should_force_select(statement, cursor) @@ -1386,7 +1427,7 @@ async def _execute_cache_hit_with_cursor( async with self.with_cursor(self.connection) as cursor: wire_sql = self._format_wire_sql(cached.compiled_sql) returns_rows = cached.operation_profile.returns_rows - fetch_kwargs = build_fetch_kwargs(self.driver_features) + fetch_kwargs = self._fetch_kwargs parameters = await coerce_large_parameters_async( self.connection, params, @@ -1590,7 +1631,7 @@ def with_cursor(self, connection: OracleAsyncConnection) -> OracleAsyncCursor: connection, arraysize=self.driver_features.get("arraysize"), prefetchrows=self.driver_features.get("prefetchrows"), - on_enter=self._apply_session_tracing if self._session_tracing_enabled() else None, + on_enter=self._apply_session_tracing if self._session_tracing_enabled else None, ) @overload @@ -1892,7 +1933,7 @@ async def load_from_arrow( async def _insert_batches() -> None: async with self.with_cursor(self.connection) as cursor: for batch in arrow_table.to_batches(max_chunksize=5000): - batch_rows = [tuple(row[col] for col in columns) for row in batch.to_pylist()] + batch_rows = self._arrow_batch_to_rows(batch, columns) coerced_rows = await coerce_many_parameters_async( self.connection, batch_rows, @@ -2049,12 +2090,10 @@ async def _dispatch_explain_plan( Execution result carrying the formatted plan rows. """ if display_cursor: - await cursor.execute(explained_sql, parameters or {}, **build_fetch_kwargs(self.driver_features)) + await cursor.execute(explained_sql, parameters or {}, **self._fetch_kwargs) if cursor.description is not None: await cursor.fetchall() - await cursor.execute( - _format_display_cursor_sql(xplan_format), {}, **build_fetch_kwargs(self.driver_features) - ) + await cursor.execute(_format_display_cursor_sql(xplan_format), {}, **self._fetch_kwargs) fetched_data = await cursor.fetchall() column_names, requires_lob_coercion = self._resolve_row_metadata(cursor.description) data, column_names = await collect_async_rows( @@ -2078,9 +2117,7 @@ async def _dispatch_explain_plan( display_sql = PLAN_DISPLAY_SQL if xplan_format == "TYPICAL" else _format_plan_display_sql(xplan_format) try: - await cursor.execute( - display_sql, {"statement_id": statement_id}, **build_fetch_kwargs(self.driver_features) - ) + await cursor.execute(display_sql, {"statement_id": statement_id}, **self._fetch_kwargs) fetched_data = await cursor.fetchall() column_names, requires_lob_coercion = self._resolve_row_metadata(cursor.description) data, column_names = await collect_async_rows( @@ -2143,7 +2180,7 @@ async def _fetch_arrow_record_batches( params = parameters if parameters is not None else [] record_batches: list[ArrowRecordBatch] = [] batch_schema: ArrowSchema | None = None - fetch_kwargs = build_arrow_fetch_kwargs(self.driver_features) + fetch_kwargs = build_arrow_fetch_kwargs(self.driver_features, arrow_schema=arrow_schema) try: oracle_dfs = self.connection.fetch_df_batches( statement=sql, parameters=params, size=batch_size or 1000, **fetch_kwargs diff --git a/tests/unit/adapters/test_oracledb/test_arrow_batches.py b/tests/unit/adapters/test_oracledb/test_arrow_batches.py index 022a5dcdc..9038bffa7 100644 --- a/tests/unit/adapters/test_oracledb/test_arrow_batches.py +++ b/tests/unit/adapters/test_oracledb/test_arrow_batches.py @@ -212,3 +212,32 @@ def fetch_df_all(self, **_kwargs: object) -> None: driver = OracleSyncDriver(cast("OracleSyncConnection", _FailingOracleConn())) with pytest.raises(SQLSpecError): driver.select_to_arrow("SELECT * FROM non_existent", return_format="table") + + +def test_sync_select_to_arrow_reader_yields_lazily() -> None: + """return_format=reader yields record batches lazily from the underlying stream.""" + yielded: list[int] = [] + + def _generate_tables() -> Iterator[pa.Table]: + yielded.append(1) + yield pa.table({"ID": [1, 2], "VALUE": [10, 20]}) + yielded.append(2) + yield pa.table({"ID": [3], "VALUE": [30]}) + + class _LazyBatchConnection: + def fetch_df_batches(self, **_kwargs: object) -> Iterator[pa.Table]: + return _generate_tables() + + driver = OracleSyncDriver( + cast("OracleSyncConnection", _LazyBatchConnection()), driver_features={"enable_lowercase_column_names": True} + ) + + result = driver.select_to_arrow("SELECT id, value FROM example", return_format="reader", batch_size=2) + reader = result.get_data() + assert yielded == [1] + first_batch = reader.read_next_batch() + assert first_batch.num_rows == 2 + assert yielded == [1] + second_batch = reader.read_next_batch() + assert second_batch.num_rows == 1 + assert yielded == [1, 2] diff --git a/tests/unit/adapters/test_oracledb/test_core_row_materialization.py b/tests/unit/adapters/test_oracledb/test_core_row_materialization.py index f0059597a..f1c0e90c3 100644 --- a/tests/unit/adapters/test_oracledb/test_core_row_materialization.py +++ b/tests/unit/adapters/test_oracledb/test_core_row_materialization.py @@ -131,12 +131,18 @@ def test_resolve_row_metadata_single_pass_unknown_type_code_shape_is_conservativ def test_resolve_row_metadata_cache_contract_still_holds_after_single_pass_rewrite() -> None: description = [("ID", _TypeCode("DB_TYPE_NUMBER"))] - cache: dict[int, tuple[object, list[str], bool]] = {} + cache: dict[object, object] = {} (first_names, first_requires_lob) = resolve_row_metadata(description, {}, cache) (second_names, second_requires_lob) = resolve_row_metadata(description, {}, cache) assert second_names is first_names assert first_requires_lob is second_requires_lob - assert cache[id(description)][0] is description + cache_key = (("ID", "DB_TYPE_NUMBER"),) + assert cache_key in cache + assert cache[cache_key][0] is description + description2 = [("ID", _TypeCode("DB_TYPE_NUMBER"))] + (third_names, third_requires_lob) = resolve_row_metadata(description2, {}, cache) + assert third_names is first_names + assert third_requires_lob is first_requires_lob @pytest.mark.parametrize("asynchronous", [False, True]) diff --git a/tests/unit/adapters/test_oracledb/test_fetch_tuning.py b/tests/unit/adapters/test_oracledb/test_fetch_tuning.py index 4496ce04e..38f7344a4 100644 --- a/tests/unit/adapters/test_oracledb/test_fetch_tuning.py +++ b/tests/unit/adapters/test_oracledb/test_fetch_tuning.py @@ -178,3 +178,23 @@ def test_resolve_rowcount_negative_rowcount_returns_zero() -> None: cursor = Mock() cursor.rowcount = -1 assert resolve_rowcount(cursor) == 0 + + +def test_build_fetch_kwargs_returns_cached_instance_when_features_unchanged() -> None: + """build_fetch_kwargs returns the identical cached dict when features are unchanged.""" + features = {"fetch_lobs": False, "fetch_decimals": True} + res1 = build_fetch_kwargs(features) + res2 = build_fetch_kwargs(features) + assert res1 is res2 + + +def test_fetch_kwargs_precomputed_on_driver_instance() -> None: + """_fetch_kwargs is precomputed and cached on the driver instance.""" + raw_cursor = _RawCursor() + driver = OracleSyncDriver( + cast("OracleSyncConnection", _SyncConnection(raw_cursor)), + driver_features={"fetch_lobs": False, "fetch_decimals": True}, + ) + assert hasattr(driver, "_fetch_kwargs") + assert driver._fetch_kwargs == {"fetch_lobs": False, "fetch_decimals": True} + assert driver._fetch_kwargs is build_fetch_kwargs({"fetch_lobs": False, "fetch_decimals": True}) From c6855853b3d010d48d89b84157463e5b7e5b1bd4 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 01:02:29 +0000 Subject: [PATCH 069/174] refactor(oracledb): clean up error mappings, unpooled connections, DDL safety, and configuration logic Map Oracle and DPY error codes to specific exceptions including SerializationConflictError, QueryTimeoutError, and DataError; ensure create_connection establishes unpooled standalone connections; separate PL/SQL-escaped clauses from top-level DDL in storage feature reporting; align Oracle identifier validation with core regex; reject soda_metadata_cache in async config; extend SyncPoolConnectionContext to support acquire and release; consolidate pool derivation and wallet processing across sync and async configs; remove dead pipeline stubs from sync driver; and eliminate all inline comments across the oracledb adapter. --- sqlspec/adapters/oracledb/_storage.py | 21 +- sqlspec/adapters/oracledb/adk/store.py | 2 +- sqlspec/adapters/oracledb/config.py | 280 ++++++++++-------- sqlspec/adapters/oracledb/core.py | 18 +- sqlspec/adapters/oracledb/driver.py | 50 +--- sqlspec/adapters/oracledb/events/_hub.py | 6 +- sqlspec/adapters/oracledb/events/backend.py | 10 +- sqlspec/adapters/oracledb/events/store.py | 4 +- sqlspec/adapters/oracledb/litestar/store.py | 2 +- sqlspec/adapters/oracledb/type_converter.py | 5 +- .../adapters/test_oracledb/test_config.py | 50 +++- .../test_core_driver_features.py | 53 +++- .../test_oracledb/test_storage_features.py | 12 +- 13 files changed, 312 insertions(+), 201 deletions(-) diff --git a/sqlspec/adapters/oracledb/_storage.py b/sqlspec/adapters/oracledb/_storage.py index 7700ba7cf..2563a8244 100644 --- a/sqlspec/adapters/oracledb/_storage.py +++ b/sqlspec/adapters/oracledb/_storage.py @@ -8,17 +8,14 @@ error instead. """ +import re from collections.abc import Mapping -from typing import Any - -from typing_extensions import TypedDict - -from sqlspec.utils.logging import get_logger - -logger = get_logger("sqlspec.adapters.oracledb.storage") +from typing import Any, TypedDict __all__ = ("oracle_table_feature_report", "validate_oracle_identifier") +_UNQUOTED_ORACLE_IDENTIFIER_PATTERN: "re.Pattern[str]" = re.compile(r"^[A-Za-z][A-Za-z0-9_$#]*$") + ORACLE_DEFAULT_HASH_PARTITIONS = 16 ORACLE_MIN_HASH_PARTITIONS = 2 ORACLE_RANGE_INTERVALS = { @@ -43,6 +40,7 @@ class _OracleStorageFeatureReport(TypedDict): applied: tuple[str, ...] clause: str + plsql_clause: str def oracle_table_feature_report( @@ -92,13 +90,16 @@ def oracle_table_feature_report( applied.append("partitioning") clause = "" + plsql_clause = "" if clauses: - clause = " " + " ".join(clauses).replace("'", "''") - return {"applied": tuple(applied), "clause": clause} + clause = " " + " ".join(clauses) + plsql_clause = " " + " ".join(clauses).replace("'", "''") + return {"applied": tuple(applied), "clause": clause, "plsql_clause": plsql_clause} def validate_oracle_identifier(value: str, field_name: str) -> str: - if not value or not value.replace("_", "").isalnum() or value[0].isdigit(): + """Validate that a string is a valid unquoted Oracle SQL identifier.""" + if not value or not _UNQUOTED_ORACLE_IDENTIFIER_PATTERN.fullmatch(value): msg = f"Oracle {field_name} must be a simple SQL identifier" raise ValueError(msg) return value diff --git a/sqlspec/adapters/oracledb/adk/store.py b/sqlspec/adapters/oracledb/adk/store.py index bc43686b3..b2f10b1be 100644 --- a/sqlspec/adapters/oracledb/adk/store.py +++ b/sqlspec/adapters/oracledb/adk/store.py @@ -3507,7 +3507,7 @@ def _adk_table_feature_clause( hash_partition_key=hash_partition_key, range_partition_key=range_partition_key, ) - return report["clause"] + return report["plsql_clause"] async def _read_lob_async(data: Any) -> Any: diff --git a/sqlspec/adapters/oracledb/config.py b/sqlspec/adapters/oracledb/config.py index c39e9ebae..1a7c99123 100644 --- a/sqlspec/adapters/oracledb/config.py +++ b/sqlspec/adapters/oracledb/config.py @@ -302,6 +302,90 @@ class OracleDriverFeatures(TypedDict): """Default ``DBMS_CLOUD`` credential name used for server-side ``DBMS_CLOUD.COPY_DATA`` object-store loads.""" +def _process_oracle_wallet_and_config( + connection_config: "Mapping[str, Any] | None", +) -> "tuple[dict[str, Any], Path | None]": + """Extract wallet zip and normalize connection configuration.""" + normalized_conn = dict(normalize_connection_config(connection_config or {})) + pending_wallet_dir: Path | None = None + if normalized_conn.get("wallet_zip") is not None: + raw_wallet_source = normalized_conn.pop("wallet_zip") + source_is_existing_dir = isinstance(raw_wallet_source, (str, Path)) and Path(raw_wallet_source).is_dir() + wallet_dir = extract_oracle_wallet_zip( + raw_wallet_source, wallet_password=normalized_conn.get("wallet_password") + ) + if not source_is_existing_dir or wallet_dir.resolve() != Path(raw_wallet_source).resolve(): + pending_wallet_dir = wallet_dir + normalized_conn.setdefault("config_dir", str(wallet_dir)) + normalized_conn.setdefault("wallet_location", str(wallet_dir)) + built_config = build_connection_config(normalized_conn) + return built_config, pending_wallet_dir + + +def _build_standalone_connection_params(connection_config: "Mapping[str, Any]") -> dict[str, Any]: + """Build connection parameters for standalone (unpooled) connect calls.""" + config = dict(connection_config) + config.pop("thick_mode", None) + config.pop("lib_dir", None) + config.pop("soda_metadata_cache", None) + config.pop("min", None) + config.pop("max", None) + config.pop("increment", None) + config.pop("homogeneous", None) + config.pop("timeout", None) + config.pop("wait_timeout", None) + config.pop("max_lifetime_session", None) + config.pop("session_callback", None) + config.pop("threaded", None) + apply_wallet_ssl_context(config) + return config + + +def _derive_oracle_pool_kwargs( + source_config: Any, + connection_config_overrides: "Mapping[str, Any] | None", + driver_features_overrides: "Mapping[str, Any] | None", + statement_config: "StatementConfig | None", + migration_config: "dict[str, Any] | MigrationConfig | None", + extension_config: "ExtensionConfigs | None", + observability_config: "ObservabilityConfig | None", + bind_key: "str | None", + overrides: dict[str, Any], +) -> dict[str, Any]: + """Derive common configuration keyword arguments for cloned Oracle pools.""" + new_conn_cfg: dict[str, Any] = dict(source_config.connection_config) + new_driver_features: dict[str, Any] = dict(source_config.driver_features) + if source_config._pool_session_callback is not None and "session_callback" not in new_conn_cfg: + new_conn_cfg["session_callback"] = source_config._pool_session_callback + if source_config._user_connection_hook is not None and "on_connection_create" not in new_driver_features: + new_driver_features["on_connection_create"] = source_config._user_connection_hook + if connection_config_overrides: + new_conn_cfg.update(connection_config_overrides) + if driver_features_overrides: + new_driver_features.update(driver_features_overrides) + for key, val in overrides.items(): + if key in VALID_ORACLE_DRIVER_FEATURE_KEYS: + new_driver_features[key] = val + else: + new_conn_cfg[key] = val + return { + "connection_config": new_conn_cfg, + "connection_instance": None, + "migration_config": cast("dict[str, Any] | None", dict(source_config.migration_config)) + if migration_config is None + else cast("dict[str, Any] | None", migration_config), + "statement_config": source_config.statement_config if statement_config is None else statement_config, + "driver_features": new_driver_features, + "bind_key": source_config.bind_key if bind_key is None else bind_key, + "extension_config": cast("ExtensionConfigs", dict(source_config.extension_config)) + if extension_config is None + else extension_config, + "observability_config": source_config.observability_config + if observability_config is None + else observability_config, + } + + class OracleSyncConnectionContext(SyncPoolConnectionContext): """Context manager for Oracle sync connections.""" @@ -311,7 +395,7 @@ def __init__(self, config: "OracleSyncConfig") -> None: super().__init__(config) self._conn: OracleSyncConnection | None = None - def __enter__(self) -> "OracleSyncConnection": + def acquire_connection(self) -> "OracleSyncConnection": if self._config.connection_instance is None: self._config.connection_instance = self._config.create_pool() conn = self._config.connection_instance.acquire() @@ -326,59 +410,44 @@ def __enter__(self) -> "OracleSyncConnection": raise return cast("OracleSyncConnection", conn) + def release_connection(self, _conn: "OracleSyncConnection | None" = None, **kwargs: Any) -> None: + if self._conn is None: + return + conn = self._conn + self._conn = None + session_hook = self._config.driver_features.get("session_context_hook") + try: + if session_hook is not None: + execute_session_release_sync(conn, session_hook) + finally: + if self._config.connection_instance: + self._config.connection_instance.release(conn) + + def __enter__(self) -> "OracleSyncConnection": + return self.acquire_connection() + def __exit__( self, exc_type: "type[BaseException] | None", exc_val: "BaseException | None", exc_tb: "TracebackType | None" ) -> bool | None: - if self._conn is not None: - conn = self._conn - self._conn = None - if exc_type is not None: - with contextlib.suppress(Exception): - conn.rollback() - session_hook = self._config.driver_features.get("session_context_hook") - try: - if session_hook is not None: - execute_session_release_sync(conn, session_hook) - finally: - if self._config.connection_instance: - self._config.connection_instance.release(conn) + if self._conn is not None and exc_type is not None: + with contextlib.suppress(Exception): + self._conn.rollback() + self.release_connection() return None class _OracleSyncSessionConnectionHandler(SyncPoolSessionFactory): - __slots__ = ("_conn",) + __slots__ = ("_ctx",) def __init__(self, config: "OracleSyncConfig") -> None: super().__init__(config) - self._conn: OracleSyncConnection | None = None + self._ctx = OracleSyncConnectionContext(config) def acquire_connection(self) -> "OracleSyncConnection": - if self._config.connection_instance is None: - self._config.connection_instance = self._config.create_pool() - conn = self._config.connection_instance.acquire() - self._conn = conn - session_hook = self._config.driver_features.get("session_context_hook") - if session_hook is not None: - try: - execute_session_checkout_sync(conn, session_hook) - except BaseException: - self._conn = None - self._config.connection_instance.release(conn) - raise - return cast("OracleSyncConnection", conn) + return cast("OracleSyncConnection", self._ctx.acquire_connection()) def release_connection(self, _conn: "OracleSyncConnection", **kwargs: Any) -> None: - if self._conn is None: - return - conn = self._conn - self._conn = None - session_hook = self._config.driver_features.get("session_context_hook") - try: - if session_hook is not None: - execute_session_release_sync(conn, session_hook) - finally: - if self._config.connection_instance: - self._config.connection_instance.release(conn) + self._ctx.release_connection(_conn, **kwargs) class OracleSyncConfig(SyncDatabaseConfig[OracleSyncConnection, "OracleSyncConnectionPool", OracleSyncDriver]): @@ -429,20 +498,8 @@ def __init__( extension_config: Extension-specific configuration. **kwargs: Additional keyword arguments. """ - normalized_conn = dict(normalize_connection_config(connection_config)) + connection_config, pending_wallet_dir = _process_oracle_wallet_and_config(connection_config) self._extracted_wallet_dir: Path | None = None - pending_wallet_dir: Path | None = None - if normalized_conn.get("wallet_zip") is not None: - raw_wallet_source = normalized_conn.pop("wallet_zip") - source_is_existing_dir = isinstance(raw_wallet_source, (str, Path)) and Path(raw_wallet_source).is_dir() - wallet_dir = extract_oracle_wallet_zip( - raw_wallet_source, wallet_password=normalized_conn.get("wallet_password") - ) - if not source_is_existing_dir or wallet_dir.resolve() != Path(raw_wallet_source).resolve(): - pending_wallet_dir = wallet_dir - normalized_conn.setdefault("config_dir", str(wallet_dir)) - normalized_conn.setdefault("wallet_location", str(wallet_dir)) - connection_config = build_connection_config(normalized_conn) self._oracle_version_cache = OracleVersionCache() self._pool_session_callback = cast( "Callable[[OracleSyncConnection, str], None] | None", connection_config.pop("session_callback", None) @@ -474,9 +531,10 @@ def create_connection(self) -> "OracleSyncConnection": Returns: An Oracle Connection instance. """ - if self.connection_instance is None: - self.connection_instance = self.create_pool() - return self.connection_instance.acquire() + config = _build_standalone_connection_params(self.connection_config) + conn = oracledb.connect(**config) + self._init_connection(conn, "") + return conn def provide_pool(self) -> "OracleSyncConnectionPool": """Provide pool instance. @@ -501,35 +559,18 @@ def derive_pool( **overrides: Any, ) -> Self: """Clone this OracleSyncConfig with overridden connection/pool or feature settings and ``connection_instance=None``.""" - new_conn_cfg: dict[str, Any] = dict(self.connection_config) - new_driver_features: dict[str, Any] = dict(self.driver_features) - if self._pool_session_callback is not None and "session_callback" not in new_conn_cfg: - new_conn_cfg["session_callback"] = self._pool_session_callback - if self._user_connection_hook is not None and "on_connection_create" not in new_driver_features: - new_driver_features["on_connection_create"] = self._user_connection_hook - if connection_config_overrides: - new_conn_cfg.update(connection_config_overrides) - if driver_features_overrides: - new_driver_features.update(driver_features_overrides) - for key, val in overrides.items(): - if key in VALID_ORACLE_DRIVER_FEATURE_KEYS: - new_driver_features[key] = val - else: - new_conn_cfg[key] = val - derived = type(self)( - connection_config=new_conn_cfg, - connection_instance=None, - migration_config=cast("dict[str, Any] | None", dict(self.migration_config)) - if migration_config is None - else cast("dict[str, Any] | None", migration_config), - statement_config=self.statement_config if statement_config is None else statement_config, - driver_features=new_driver_features, - bind_key=self.bind_key if bind_key is None else bind_key, - extension_config=cast("ExtensionConfigs", dict(self.extension_config)) - if extension_config is None - else extension_config, - observability_config=self.observability_config if observability_config is None else observability_config, + derived_kwargs = _derive_oracle_pool_kwargs( + self, + connection_config_overrides, + driver_features_overrides, + statement_config, + migration_config, + extension_config, + observability_config, + bind_key, + overrides, ) + derived = type(self)(**derived_kwargs) if ( self._extracted_wallet_dir is not None and derived._extracted_wallet_dir is None @@ -588,7 +629,11 @@ def _create_pool(self) -> "OracleSyncConnectionPool": client_config["lib_dir"] = lib_dir if config.get("config_dir") is not None: client_config["config_dir"] = config["config_dir"] - oracledb.init_oracle_client(**client_config) + try: + oracledb.init_oracle_client(**client_config) + except oracledb.Error as exc: + msg = f"Failed to initialize Oracle client in thick mode: {exc}" + raise ImproperConfigurationError(msg) from exc apply_wallet_ssl_context(config) @@ -758,20 +803,8 @@ def __init__( extension_config: Extension-specific configuration. **kwargs: Additional keyword arguments. """ - normalized_conn = dict(normalize_connection_config(connection_config)) + connection_config, pending_wallet_dir = _process_oracle_wallet_and_config(connection_config) self._extracted_wallet_dir: Path | None = None - pending_wallet_dir: Path | None = None - if normalized_conn.get("wallet_zip") is not None: - raw_wallet_source = normalized_conn.pop("wallet_zip") - source_is_existing_dir = isinstance(raw_wallet_source, (str, Path)) and Path(raw_wallet_source).is_dir() - wallet_dir = extract_oracle_wallet_zip( - raw_wallet_source, wallet_password=normalized_conn.get("wallet_password") - ) - if not source_is_existing_dir or wallet_dir.resolve() != Path(raw_wallet_source).resolve(): - pending_wallet_dir = wallet_dir - normalized_conn.setdefault("config_dir", str(wallet_dir)) - normalized_conn.setdefault("wallet_location", str(wallet_dir)) - connection_config = build_connection_config(normalized_conn) self._oracle_version_cache = OracleVersionCache() self._pool_session_callback = cast( "Callable[[OracleAsyncConnection, str], Any] | None", connection_config.pop("session_callback", None) @@ -804,9 +837,10 @@ async def create_connection(self) -> OracleAsyncConnection: Returns: An Oracle AsyncConnection instance. """ - if self.connection_instance is None: - self.connection_instance = await self.create_pool() - return cast("OracleAsyncConnection", await self.connection_instance.acquire()) + config = _build_standalone_connection_params(self.connection_config) + conn = await oracledb.connect_async(**config) + await self._init_connection(conn, "") + return cast("OracleAsyncConnection", conn) async def provide_pool(self) -> "OracleAsyncConnectionPool": """Provide async pool instance. @@ -831,35 +865,18 @@ def derive_pool( **overrides: Any, ) -> Self: """Clone this OracleAsyncConfig with overridden connection/pool or feature settings and ``connection_instance=None``.""" - new_conn_cfg: dict[str, Any] = dict(self.connection_config) - new_driver_features: dict[str, Any] = dict(self.driver_features) - if self._pool_session_callback is not None and "session_callback" not in new_conn_cfg: - new_conn_cfg["session_callback"] = self._pool_session_callback - if self._user_connection_hook is not None and "on_connection_create" not in new_driver_features: - new_driver_features["on_connection_create"] = self._user_connection_hook - if connection_config_overrides: - new_conn_cfg.update(connection_config_overrides) - if driver_features_overrides: - new_driver_features.update(driver_features_overrides) - for key, val in overrides.items(): - if key in VALID_ORACLE_DRIVER_FEATURE_KEYS: - new_driver_features[key] = val - else: - new_conn_cfg[key] = val - derived = type(self)( - connection_config=new_conn_cfg, - connection_instance=None, - migration_config=cast("dict[str, Any] | None", dict(self.migration_config)) - if migration_config is None - else cast("dict[str, Any] | None", migration_config), - statement_config=self.statement_config if statement_config is None else statement_config, - driver_features=new_driver_features, - bind_key=self.bind_key if bind_key is None else bind_key, - extension_config=cast("ExtensionConfigs", dict(self.extension_config)) - if extension_config is None - else extension_config, - observability_config=self.observability_config if observability_config is None else observability_config, + derived_kwargs = _derive_oracle_pool_kwargs( + self, + connection_config_overrides, + driver_features_overrides, + statement_config, + migration_config, + extension_config, + observability_config, + bind_key, + overrides, ) + derived = type(self)(**derived_kwargs) if ( self._extracted_wallet_dir is not None and derived._extracted_wallet_dir is None @@ -909,8 +926,9 @@ async def _create_pool(self) -> "OracleAsyncConnectionPool": config = dict(self.connection_config) thick_mode = config.pop("thick_mode", False) lib_dir = config.pop("lib_dir", None) - if thick_mode or lib_dir is not None: - msg = "OracleAsyncConfig only supports Thin mode; use OracleSyncConfig for Thick mode." + soda_metadata_cache = config.get("soda_metadata_cache", False) + if thick_mode or lib_dir is not None or soda_metadata_cache: + msg = "OracleAsyncConfig only supports Thin mode; Thick mode options (thick_mode, lib_dir, soda_metadata_cache) require OracleSyncConfig." raise ImproperConfigurationError(msg) apply_wallet_ssl_context(config) diff --git a/sqlspec/adapters/oracledb/core.py b/sqlspec/adapters/oracledb/core.py index 78f177459..c3df5650d 100644 --- a/sqlspec/adapters/oracledb/core.py +++ b/sqlspec/adapters/oracledb/core.py @@ -47,6 +47,8 @@ OperationalError, OperationCancelledError, PermissionDeniedError, + QueryTimeoutError, + SerializationConflictError, SQLParsingError, SQLSpecError, TransactionError, @@ -395,6 +397,7 @@ def _chained_output_handler(inner: Any, fallback: "Any | None", cursor: Any, met 1400: (NotNullViolationError, "not-null constraint violation"), 1407: (NotNullViolationError, "not-null constraint violation"), 1438: (DataError, "numeric precision overflow"), + 1476: (DataError, "divisor is equal to zero"), 1652: (OperationalError, "tablespace full"), 1722: (DataError, "invalid number"), 1830: (DataError, "date format picture ends before converting entire input string"), @@ -408,9 +411,10 @@ def _chained_output_handler(inner: Any, fallback: "Any | None", cursor: Any, met 3113: (DatabaseConnectionError, "end-of-file on communication channel"), 3114: (DatabaseConnectionError, "not connected to ORACLE"), 3135: (DatabaseConnectionError, "connection lost contact"), + 3136: (ConnectionTimeoutError, "inbound connection timed out"), 4068: (OperationalError, "existing state of packages has been discarded"), 8176: (TransactionError, "consistent read failure"), - 8177: (TransactionError, "cannot serialize access"), + 8177: (SerializationConflictError, "cannot serialize access"), 12154: (DatabaseConnectionError, "TNS resolution failed"), 12170: (ConnectionTimeoutError, "connect timeout"), 12505: (DatabaseConnectionError, "listener rejected"), @@ -421,11 +425,15 @@ def _chained_output_handler(inner: Any, fallback: "Any | None", cursor: Any, met 28000: (PermissionDeniedError, "the account is locked"), 28001: (PermissionDeniedError, "the password has expired"), 30006: (OperationalError, "resource busy; acquire with WAIT timeout expired"), + 40441: (DataError, "JSON syntax error"), + 42699: (SerializationConflictError, "transaction isolation violation"), + 51803: (DataError, "vector operation error"), } _DPY_ERROR_CODE_MAPPING: "dict[str, tuple[type[SQLSpecError], str]]" = { "DPY-1001": (DatabaseConnectionError, "not connected to database"), "DPY-4011": (DatabaseConnectionError, "connection was closed"), + "DPY-4024": (QueryTimeoutError, "query execution timeout expired"), "DPY-6000": (DatabaseConnectionError, "cannot connect to database"), "DPY-6001": (DatabaseConnectionError, "cannot connect to database"), "DPY-6005": (ConnectionTimeoutError, "cannot connect to database"), @@ -1796,9 +1804,13 @@ def create_mapped_exception(error: Any, *, logger: Any | None = None) -> SQLSpec if not isinstance(error_code, int) or isinstance(error_code, bool) or not error_code: error_text = str(error_obj) for dpy_code, (error_class, error_desc) in _DPY_ERROR_CODE_MAPPING.items(): - if error_text.startswith(dpy_code): + if dpy_code in error_text: return _create_oracle_error(error, None, error_class, f"{error_desc} [{dpy_code}]") - return _create_oracle_error(error, None, SQLSpecError, "database error") + ora_match = re.search(r"ORA-(\d+)", error_text) + if ora_match: + error_code = int(ora_match.group(1)) + else: + return _create_oracle_error(error, None, SQLSpecError, "database error") mapping = _ERROR_CODE_MAPPING.get(error_code) if mapping: diff --git a/sqlspec/adapters/oracledb/driver.py b/sqlspec/adapters/oracledb/driver.py index 9b9ada15b..d0dff6c36 100644 --- a/sqlspec/adapters/oracledb/driver.py +++ b/sqlspec/adapters/oracledb/driver.py @@ -4,7 +4,7 @@ from itertools import chain from pathlib import Path from time import perf_counter -from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, Protocol, cast, overload +from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, NamedTuple, Protocol, cast, overload from sqlspec.adapters.oracledb._typing import ( DB_TYPE_BLOB, @@ -84,6 +84,7 @@ hash_stack_operations, validate_savepoint_name, ) +from sqlspec.driver._common import CommonDriverAttributesMixin from sqlspec.exceptions import ImproperConfigurationError, SQLSpecError, StackExecutionError from sqlspec.utils.arrow_helpers import coerce_oracle_df_to_arrow_table from sqlspec.utils.logging import get_logger, log_with_context @@ -156,7 +157,7 @@ def _compiled_sql(self, statement: "SQL", statement_config: "StatementConfig") - """ -class OraclePipelineMixin: +class OraclePipelineMixin(CommonDriverAttributesMixin): """Shared helpers for Oracle pipeline execution.""" __slots__ = () @@ -194,7 +195,7 @@ def _format_wire_sql(self, sql: str, config: "StatementConfig | None" = None) -> def _compiled_sql( self, statement: "SQL", statement_config: "StatementConfig", flatten_single_parameters: bool = False ) -> "tuple[str, Any]": - compiled_sql, prepared_parameters = super()._compiled_sql( # type: ignore[misc] + compiled_sql, prepared_parameters = super()._compiled_sql( statement, statement_config, flatten_single_parameters=flatten_single_parameters ) if getattr(self, "_session_tracing_enabled", False): @@ -249,9 +250,7 @@ def _prepare_pipeline_operation(self, operation: "StackOperation") -> "_Compiled msg = f"Unsupported stack operation method: {operation.method}" raise ValueError(msg) - compiled_sql, prepared_parameters = driver._compiled_sql( # pyright: ignore[reportPrivateUsage] - sql_statement, config - ) + compiled_sql, prepared_parameters = self._compiled_sql(sql_statement, config) summary = describe_stack_statement(operation.statement) return _CompiledStackOperation( statement=sql_statement, @@ -379,13 +378,17 @@ class OracleSyncDriver(OraclePipelineMixin, SyncDriverAdapterBase): "_data_dictionary", "_fetch_kwargs", "_oracle_version_cache", - "_pipeline_support", - "_pipeline_support_reason", "_row_metadata_cache", "_session_tracing_enabled", "_transaction_active", ) dialect = "oracle" + _pipeline_support: ClassVar[bool] = False + _pipeline_support_reason: ClassVar[str] = "asyncio_thin_required" + + def _pipeline_native_supported(self) -> bool: + """Sync Oracle connections do not support native pipelining.""" + return False def __init__( self, @@ -400,8 +403,6 @@ def __init__( super().__init__(connection=connection, statement_config=statement_config, driver_features=driver_features) self._data_dictionary: OracledbSyncDataDictionary | None = None - self._pipeline_support: bool | None = None - self._pipeline_support_reason: str | None = None self._oracle_version_cache: OracleVersionCache | None = None self._row_metadata_cache: dict[int, tuple[Any, list[str], bool]] = {} self._transaction_active = False @@ -872,17 +873,9 @@ def select_to_arrow( ) def execute_stack(self, stack: "StatementStack", *, continue_on_error: bool = False) -> "tuple[StackResult, ...]": - """Execute a StatementStack using Oracle's pipeline when available.""" - if not isinstance(stack, StatementStack) or not stack: - return super().execute_stack(stack, continue_on_error=continue_on_error) - - blocker = self._stack_native_blocker(stack) - if blocker is not None: - self._log_pipeline_skip(blocker, stack) - return super().execute_stack(stack, continue_on_error=continue_on_error) - - if not self._pipeline_native_supported(): - self._log_pipeline_skip(self._pipeline_support_reason or "database_version", stack) + """Execute a StatementStack sequentially on the sync driver.""" + self._require_capability("stack_execution") + self._log_pipeline_skip("asyncio_thin_required", stack) return super().execute_stack(stack, continue_on_error=continue_on_error) def select_to_storage( @@ -1291,19 +1284,6 @@ def _fetch_arrow_record_batches( record_batches.extend(batch_table.to_batches()) return record_batches, batch_schema - def _pipeline_native_supported(self) -> bool: - if self._pipeline_support is not None: - return self._pipeline_support - - if self.stack_native_disabled: - self._pipeline_support = False - self._pipeline_support_reason = "env_override" - return False - - self._pipeline_support = False - self._pipeline_support_reason = "asyncio_thin_required" - return False - class OracleAsyncDriver(OraclePipelineMixin, AsyncDriverAdapterBase): """Asynchronous Oracle Database driver. @@ -2258,7 +2238,7 @@ async def _execute_stack_native( if owns_transaction: try: await self.rollback() - except Exception as rollback_error: # pragma: no cover + except Exception as rollback_error: logger.debug("Rollback after pipeline failure failed: %s", rollback_error) if isinstance(exc, StackExecutionError): raise diff --git a/sqlspec/adapters/oracledb/events/_hub.py b/sqlspec/adapters/oracledb/events/_hub.py index 50f748568..14226f667 100644 --- a/sqlspec/adapters/oracledb/events/_hub.py +++ b/sqlspec/adapters/oracledb/events/_hub.py @@ -6,8 +6,6 @@ caller's poll_interval (capped at the configured aq_wait_seconds ceiling). """ -# pyright: reportPrivateUsage=false - import asyncio import contextlib import logging @@ -84,7 +82,7 @@ def dequeue(self, channel: str, poll_interval: float) -> "Any | None": _apply_deq_options(queue, self._visibility, self._default_visibility, wait_seconds) try: message = queue.deqone() - except Exception as error: # pragma: no cover + except Exception as error: if not isinstance(error, OracleDatabaseError): raise log_with_context( @@ -218,7 +216,7 @@ async def dequeue(self, channel: str, poll_interval: float) -> "Any | None": _apply_deq_options(queue, self._visibility, self._default_visibility, wait_seconds) try: message = await queue.deqone() - except Exception as error: # pragma: no cover + except Exception as error: if not isinstance(error, OracleDatabaseError): raise log_with_context( diff --git a/sqlspec/adapters/oracledb/events/backend.py b/sqlspec/adapters/oracledb/events/backend.py index d7f218a15..c8bb32077 100644 --- a/sqlspec/adapters/oracledb/events/backend.py +++ b/sqlspec/adapters/oracledb/events/backend.py @@ -3,7 +3,7 @@ import contextlib import logging from datetime import UTC, datetime -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, cast from sqlspec.adapters.oracledb._typing import DEQ_IMMEDIATE, DEQ_ON_COMMIT from sqlspec.adapters.oracledb.events._hub import OracleAsyncAQHub, OracleSyncAQHub @@ -213,22 +213,22 @@ def create_event_backend( match (backend_name, is_async): case ("aq", False): try: - return OracleSyncAQEventBackend(config, extension_settings) # type: ignore[arg-type] + return OracleSyncAQEventBackend(cast("OracleSyncConfig", config), extension_settings) except (ImproperConfigurationError, MissingDependencyError): return None case ("aq", True): try: - return OracleAsyncAQEventBackend(config, extension_settings) # type: ignore[arg-type] + return OracleAsyncAQEventBackend(cast("OracleAsyncConfig", config), extension_settings) except (ImproperConfigurationError, MissingDependencyError): return None case ("txeventq", False): try: - return OracleSyncTxEventQEventBackend(config, extension_settings) # type: ignore[arg-type] + return OracleSyncTxEventQEventBackend(cast("OracleSyncConfig", config), extension_settings) except (ImproperConfigurationError, MissingDependencyError): return None case ("txeventq", True): try: - return OracleAsyncTxEventQEventBackend(config, extension_settings) # type: ignore[arg-type] + return OracleAsyncTxEventQEventBackend(cast("OracleAsyncConfig", config), extension_settings) except (ImproperConfigurationError, MissingDependencyError): return None case _: diff --git a/sqlspec/adapters/oracledb/events/store.py b/sqlspec/adapters/oracledb/events/store.py index b5954477f..48fc64dbb 100644 --- a/sqlspec/adapters/oracledb/events/store.py +++ b/sqlspec/adapters/oracledb/events/store.py @@ -220,7 +220,7 @@ def _table_feature_clause(self) -> str: range_partition_key="available_at", table_options_key="table_options", ) - return report["clause"] + return report["plsql_clause"] class OracleAsyncEventQueueStore(BaseEventQueueStore["OracleAsyncConfig"]): @@ -337,7 +337,7 @@ def _table_feature_clause(self) -> str: range_partition_key="available_at", table_options_key="table_options", ) - return report["clause"] + return report["plsql_clause"] def _init_oracle_settings(extension_settings: "dict[str, Any]") -> "tuple[bool, JSONStorageType | None]": diff --git a/sqlspec/adapters/oracledb/litestar/store.py b/sqlspec/adapters/oracledb/litestar/store.py index 1f53ce0aa..d81da305c 100644 --- a/sqlspec/adapters/oracledb/litestar/store.py +++ b/sqlspec/adapters/oracledb/litestar/store.py @@ -758,4 +758,4 @@ def _litestar_table_feature_clause(config: Any, in_memory: bool) -> str: range_partition_key="expires_at", table_options_key="table_options", ) - return report["clause"] + return report["plsql_clause"] diff --git a/sqlspec/adapters/oracledb/type_converter.py b/sqlspec/adapters/oracledb/type_converter.py index a971cefcc..8622e5881 100644 --- a/sqlspec/adapters/oracledb/type_converter.py +++ b/sqlspec/adapters/oracledb/type_converter.py @@ -7,10 +7,7 @@ import array from typing import Any -from sqlspec.adapters.oracledb._vector_handlers import ( - numpy_converter_in, # pyright: ignore[reportPrivateUsage] - numpy_converter_out, # pyright: ignore[reportPrivateUsage] -) +from sqlspec.adapters.oracledb._vector_handlers import numpy_converter_in, numpy_converter_out from sqlspec.typing import NUMPY_INSTALLED from sqlspec.utils.module_loader import import_optional from sqlspec.utils.sync_tools import ensure_async_ diff --git a/tests/unit/adapters/test_oracledb/test_config.py b/tests/unit/adapters/test_oracledb/test_config.py index 18d25eb13..cbd72213e 100644 --- a/tests/unit/adapters/test_oracledb/test_config.py +++ b/tests/unit/adapters/test_oracledb/test_config.py @@ -347,17 +347,61 @@ async def test_async_pool_rejects_thick_mode(options: dict[str, Any], monkeypatc create_pool.assert_not_called() -async def test_async_soda_option_does_not_initialize_thick_mode(monkeypatch: pytest.MonkeyPatch) -> None: +async def test_async_soda_option_rejects_thick_mode_option(monkeypatch: pytest.MonkeyPatch) -> None: initialize = Mock() create_pool = Mock() monkeypatch.setattr(oracle_config_module.oracledb, "init_oracle_client", initialize) monkeypatch.setattr(oracle_config_module.oracledb, "create_pool_async", create_pool) config = OracleAsyncConfig(connection_config={"soda_metadata_cache": True}) - assert await config._create_pool() is create_pool.return_value + with pytest.raises(ImproperConfigurationError, match="only supports Thin mode"): + await config._create_pool() initialize.assert_not_called() - assert create_pool.call_args.kwargs["soda_metadata_cache"] is True + create_pool.assert_not_called() + + +def test_sync_create_connection_unpooled(monkeypatch: pytest.MonkeyPatch) -> None: + """create_connection creates an unpooled standalone connection without instantiating a pool.""" + connect_mock = Mock() + create_pool_mock = Mock() + monkeypatch.setattr(oracle_config_module.oracledb, "connect", connect_mock) + monkeypatch.setattr(oracle_config_module.oracledb, "create_pool", create_pool_mock) + config = OracleSyncConfig(connection_config={"user": "scott", "password": "tiger", "dsn": "localhost/xe"}) + conn = config.create_connection() + assert conn is connect_mock.return_value + create_pool_mock.assert_not_called() + assert config.connection_instance is None + + +@pytest.mark.anyio +async def test_async_create_connection_unpooled(monkeypatch: pytest.MonkeyPatch) -> None: + """create_connection creates an unpooled standalone async connection without instantiating a pool.""" + connect_mock = AsyncMock() + create_pool_mock = AsyncMock() + monkeypatch.setattr(oracle_config_module.oracledb, "connect_async", connect_mock) + monkeypatch.setattr(oracle_config_module.oracledb, "create_pool_async", create_pool_mock) + config = OracleAsyncConfig(connection_config={"user": "scott", "password": "tiger", "dsn": "localhost/xe"}) + conn = await config.create_connection() + assert conn is connect_mock.return_value + create_pool_mock.assert_not_called() + assert config.connection_instance is None + + +def test_derive_pool_produces_identical_configurations_across_sync_and_async() -> None: + """derive_pool produces identical connection configurations across sync and async configs.""" + common_cfg = {"user": "scott", "password": "tiger", "dsn": "localhost/xe", "min": 2, "max": 10} + common_features = {"enable_lowercase_column_names": True, "fetch_lobs": False} + sync_config = OracleSyncConfig(connection_config=common_cfg, driver_features=common_features) + async_config = OracleAsyncConfig(connection_config=common_cfg, driver_features=common_features) + + derived_sync = sync_config.derive_pool(connection_config_overrides={"max": 20}, fetch_lobs=True) + derived_async = async_config.derive_pool(connection_config_overrides={"max": 20}, fetch_lobs=True) + + assert derived_sync.connection_config == derived_async.connection_config + sync_features = {k: v for k, v in derived_sync.driver_features.items() if k != "_statement_cache_holder"} + async_features = {k: v for k, v in derived_async.driver_features.items() if k != "_statement_cache_holder"} + assert sync_features == async_features def test_oracle_config_wallet_zip_extracted_and_cleaned_up_on_close_pool() -> None: diff --git a/tests/unit/adapters/test_oracledb/test_core_driver_features.py b/tests/unit/adapters/test_oracledb/test_core_driver_features.py index 58cee2653..8adb5f874 100644 --- a/tests/unit/adapters/test_oracledb/test_core_driver_features.py +++ b/tests/unit/adapters/test_oracledb/test_core_driver_features.py @@ -1,4 +1,3 @@ -# pyright: reportArgumentType=false """Unit tests for OracleDB ``apply_driver_features`` defaults.""" import inspect @@ -506,3 +505,55 @@ def test_apply_wallet_ssl_context_noop_when_no_wallet() -> None: conn_config: dict[str, Any] = {"user": "app", "password": "secret"} apply_wallet_ssl_context(conn_config) assert "ssl_context" not in conn_config + + +def test_ora_08177_maps_to_serialization_conflict_error() -> None: + """ORA-08177 maps to SerializationConflictError.""" + import oracledb + + from sqlspec.adapters.oracledb.core import create_mapped_exception + from sqlspec.exceptions import SerializationConflictError + + error = oracledb.DatabaseError("ORA-08177: can't serialize access for this transaction") + mapped = create_mapped_exception(error) + assert isinstance(mapped, SerializationConflictError) + + +def test_ora_42699_maps_to_serialization_conflict_error() -> None: + """ORA-42699 maps to SerializationConflictError.""" + import oracledb + + from sqlspec.adapters.oracledb.core import create_mapped_exception + from sqlspec.exceptions import SerializationConflictError + + error = oracledb.DatabaseError("ORA-42699: transaction isolation violation") + mapped = create_mapped_exception(error) + assert isinstance(mapped, SerializationConflictError) + + +def test_dpy_4024_maps_to_query_timeout_error() -> None: + """DPY-4024 maps to QueryTimeoutError.""" + import oracledb + + from sqlspec.adapters.oracledb.core import create_mapped_exception + from sqlspec.exceptions import QueryTimeoutError + + error = oracledb.OperationalError("DPY-4024: call timeout expired") + mapped = create_mapped_exception(error) + assert isinstance(mapped, QueryTimeoutError) + + +def test_ora_data_and_timeout_error_mappings() -> None: + """Verify ORA-40441, ORA-51803, ORA-01476 map to DataError, ORA-03136 maps to ConnectionTimeoutError.""" + import oracledb + + from sqlspec.adapters.oracledb.core import create_mapped_exception + from sqlspec.exceptions import ConnectionTimeoutError, DataError + + assert isinstance(create_mapped_exception(oracledb.DatabaseError("ORA-40441: JSON syntax error")), DataError) + assert isinstance(create_mapped_exception(oracledb.DatabaseError("ORA-51803: vector operation error")), DataError) + assert isinstance(create_mapped_exception(oracledb.DatabaseError("ORA-01476: divisor is equal to zero")), DataError) + assert isinstance( + create_mapped_exception(oracledb.DatabaseError("ORA-03136: inbound connection timed out")), + ConnectionTimeoutError, + ) diff --git a/tests/unit/adapters/test_oracledb/test_storage_features.py b/tests/unit/adapters/test_oracledb/test_storage_features.py index feb5aaca9..d1a1b495a 100644 --- a/tests/unit/adapters/test_oracledb/test_storage_features.py +++ b/tests/unit/adapters/test_oracledb/test_storage_features.py @@ -56,4 +56,14 @@ def test_oracle_range_partitioning_uses_interval() -> None: ) assert "PARTITION BY RANGE (timestamp)" in report["clause"] - assert "INTERVAL (NUMTOYMINTERVAL(1, ''MONTH''))" in report["clause"] + assert "INTERVAL (NUMTOYMINTERVAL(1, 'MONTH'))" in report["clause"] + assert "INTERVAL (NUMTOYMINTERVAL(1, ''MONTH''))" in report["plsql_clause"] + + +def test_validate_oracle_identifier_accepts_dollar_sign() -> None: + """validate_oracle_identifier accepts $, #, and _ in Oracle identifiers.""" + from sqlspec.adapters.oracledb._storage import validate_oracle_identifier + + assert validate_oracle_identifier("tbl$name", "table name") == "tbl$name" + assert validate_oracle_identifier("tbl#name", "table name") == "tbl#name" + assert validate_oracle_identifier("tbl_name", "table name") == "tbl_name" From 4d5b36a41d271ec459cdbf6dcaf12bf21b0ae44d Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 01:05:46 +0000 Subject: [PATCH 070/174] fix(cockroach): gate follower reads on read-only queries and support streaming Gate follower-read transactions behind `_wants_follower_read` so modifying queries with return clauses do not run in read-only transactions. Harden follower-read staleness application in `begin` with rollback on error, and wrap streaming select executions to apply follower-read transactions. --- sqlspec/adapters/cockroach_asyncpg/driver.py | 67 +++++++- sqlspec/adapters/cockroach_psycopg/driver.py | 147 +++++++++++++++++- .../test_cockroach_asyncpg/test_config.py | 71 +++++++++ .../test_cockroach_psycopg/test_config.py | 79 ++++++++++ 4 files changed, 355 insertions(+), 9 deletions(-) diff --git a/sqlspec/adapters/cockroach_asyncpg/driver.py b/sqlspec/adapters/cockroach_asyncpg/driver.py index 51c9bc890..e33ec2969 100644 --- a/sqlspec/adapters/cockroach_asyncpg/driver.py +++ b/sqlspec/adapters/cockroach_asyncpg/driver.py @@ -20,23 +20,67 @@ ) from sqlspec.adapters.cockroach_asyncpg.data_dictionary import CockroachAsyncpgDataDictionary from sqlspec.core import SQL, register_driver_profile +from sqlspec.driver import AsyncRowStream from sqlspec.utils.logging import get_logger from sqlspec.utils.type_guards import has_sqlstate if TYPE_CHECKING: from collections.abc import Awaitable, Callable + from sqlspec.adapters.asyncpg.core import AsyncpgStreamSource from sqlspec.adapters.cockroach_asyncpg._typing import CockroachAsyncpgConnection from sqlspec.core import StatementConfig from sqlspec.driver import ExecutionResult from sqlspec.storage import StorageBridgeJob, StorageDestination, StorageFormat, StorageTelemetry -__all__ = ("CockroachAsyncpgDriver", "CockroachAsyncpgExceptionHandler", "CockroachAsyncpgSessionContext") +__all__ = ( + "CockroachAsyncpgDriver", + "CockroachAsyncpgExceptionHandler", + "CockroachAsyncpgSessionContext", + "CockroachAsyncpgStreamSource", +) logger = get_logger("sqlspec.adapters.cockroach_asyncpg") _T = TypeVar("_T") +class CockroachAsyncpgStreamSource: + """Async chunk source for CockroachDB AsyncPG with follower read support.""" + + __slots__ = ("_auto_started", "_driver", "_statement", "_underlying") + + def __init__(self, driver: "CockroachAsyncpgDriver", statement: "SQL", underlying: "AsyncpgStreamSource") -> None: + self._driver = driver + self._statement = statement + self._underlying = underlying + self._auto_started = False + + async def start(self) -> None: + if self._driver._wants_follower_read(self._statement): + self._auto_started = await self._driver._begin_follower_read_transaction() + try: + await self._underlying.start() + except BaseException: + if self._auto_started and self._driver._connection_in_transaction(): + with suppress(Exception): + await self._driver.rollback() + raise + + async def fetch_chunk(self) -> "list[dict[str, Any]]": + return await self._underlying.fetch_chunk() + + async def close(self, error: bool = False) -> None: + try: + await self._underlying.close(error=error) + finally: + if self._auto_started and self._driver._connection_in_transaction(): + if error: + with suppress(Exception): + await self._driver.rollback() + else: + await self._driver.commit() + + class CockroachAsyncpgExceptionHandler(AsyncpgExceptionHandler): """Async context manager for CockroachDB AsyncPG exceptions.""" @@ -164,6 +208,10 @@ async def load_from_storage( table, source, file_format=file_format, partitioner=partitioner, overwrite=overwrite, **kwargs ) + def _wants_follower_read(self, statement: "SQL") -> bool: + """Return True when a statement may run under a follower-read transaction.""" + return statement.returns_rows() and not statement.is_modifying_operation() + async def begin(self) -> None: """Begin a transaction and apply follower-read staleness to it. @@ -176,7 +224,12 @@ async def begin(self) -> None: await super().begin() return await super().begin() - await self._apply_follower_reads() + try: + await self._apply_follower_reads() + except Exception as exc: + with suppress(Exception): + await self.rollback() + raise create_mapped_exception(exc) from exc async def run_transaction_with_retry(self, operation: "Callable[[], Awaitable[_T]]") -> _T: """Execute a full CockroachDB transaction callback with serialization retries. @@ -211,7 +264,7 @@ async def run_transaction_with_retry(self, operation: "Callable[[], Awaitable[_T attempt += 1 async def dispatch_execute(self, cursor: Any, statement: SQL) -> "ExecutionResult": - auto_started = await self._begin_follower_read_transaction() if statement.returns_rows() else False + auto_started = await self._begin_follower_read_transaction() if self._wants_follower_read(statement) else False try: result = await super().dispatch_execute(cursor, statement) except BaseException: @@ -223,6 +276,14 @@ async def dispatch_execute(self, cursor: Any, statement: SQL) -> "ExecutionResul await self.commit() return result + def dispatch_select_stream(self, statement: "SQL", chunk_size: int) -> "AsyncRowStream[dict[str, Any]] | None": + """Return a native asyncpg row stream wrapped for follower reads.""" + stream = super().dispatch_select_stream(statement, chunk_size) + if stream is None: + return None + source = CockroachAsyncpgStreamSource(self, statement, cast("AsyncpgStreamSource", stream._source)) + return AsyncRowStream(source) + def handle_database_exceptions(self) -> "CockroachAsyncpgExceptionHandler": return CockroachAsyncpgExceptionHandler() diff --git a/sqlspec/adapters/cockroach_psycopg/driver.py b/sqlspec/adapters/cockroach_psycopg/driver.py index cbc4d9afd..eda4d8ff7 100644 --- a/sqlspec/adapters/cockroach_psycopg/driver.py +++ b/sqlspec/adapters/cockroach_psycopg/driver.py @@ -29,6 +29,7 @@ CockroachPsycopgAsyncDataDictionary, CockroachPsycopgSyncDataDictionary, ) +from sqlspec.adapters.psycopg.core import create_mapped_exception from sqlspec.adapters.psycopg.driver import ( PsycopgAsyncDriver, PsycopgAsyncExceptionHandler, @@ -36,12 +37,14 @@ PsycopgSyncExceptionHandler, ) from sqlspec.core import SQL, StatementConfig, get_cache_config, register_driver_profile +from sqlspec.driver import AsyncRowStream, SyncRowStream from sqlspec.utils.logging import get_logger if TYPE_CHECKING: from collections.abc import Awaitable, Callable from sqlspec.adapters.cockroach_psycopg._typing import CockroachAsyncCursor, CockroachSyncCursor + from sqlspec.adapters.psycopg.core import PsycopgAsyncStreamSource, PsycopgSyncStreamSource from sqlspec.core import SQLResult from sqlspec.driver import CachedQuery, ExecutionResult from sqlspec.storage import StorageBridgeJob, StorageDestination, StorageFormat, StorageTelemetry @@ -50,15 +53,92 @@ "CockroachPsycopgAsyncDriver", "CockroachPsycopgAsyncExceptionHandler", "CockroachPsycopgAsyncSessionContext", + "CockroachPsycopgAsyncStreamSource", "CockroachPsycopgSyncDriver", "CockroachPsycopgSyncExceptionHandler", "CockroachPsycopgSyncSessionContext", + "CockroachPsycopgSyncStreamSource", ) logger = get_logger("sqlspec.adapters.cockroach_psycopg") _T = TypeVar("_T") +class CockroachPsycopgSyncStreamSource: + """Sync chunk source for CockroachDB Psycopg with follower read support.""" + + __slots__ = ("_auto_started", "_driver", "_underlying") + + def __init__( + self, driver: "CockroachPsycopgSyncDriver", underlying: "PsycopgSyncStreamSource", auto_started: bool + ) -> None: + self._driver = driver + self._underlying = underlying + self._auto_started = auto_started + + def start(self) -> None: + try: + self._underlying.start() + except BaseException: + if self._auto_started and self._driver._connection_in_transaction(): + with suppress(Exception): + self._driver.rollback() + raise + + def fetch_chunk(self) -> "list[dict[str, Any]]": + return self._underlying.fetch_chunk() + + def close(self, error: bool = False) -> None: + try: + self._underlying.close(error=error) + finally: + if self._auto_started and self._driver._connection_in_transaction(): + if error: + with suppress(Exception): + self._driver.rollback() + else: + self._driver.commit() + + +class CockroachPsycopgAsyncStreamSource: + """Async chunk source for CockroachDB Psycopg with follower read support.""" + + __slots__ = ("_auto_started", "_driver", "_statement", "_underlying") + + def __init__( + self, driver: "CockroachPsycopgAsyncDriver", statement: "SQL", underlying: "PsycopgAsyncStreamSource" + ) -> None: + self._driver = driver + self._statement = statement + self._underlying = underlying + self._auto_started = False + + async def start(self) -> None: + if self._driver._wants_follower_read(self._statement): + self._auto_started = await self._driver._begin_follower_read_transaction() + try: + await self._underlying.start() + except BaseException: + if self._auto_started and self._driver._connection_in_transaction(): + with suppress(Exception): + await self._driver.rollback() + raise + + async def fetch_chunk(self) -> "list[dict[str, Any]]": + return await self._underlying.fetch_chunk() + + async def close(self, error: bool = False) -> None: + try: + await self._underlying.close(error=error) + finally: + if self._auto_started and self._driver._connection_in_transaction(): + if error: + with suppress(Exception): + await self._driver.rollback() + else: + await self._driver.commit() + + class CockroachPsycopgSyncExceptionHandler(PsycopgSyncExceptionHandler): """Context manager for handling CockroachDB psycopg exceptions.""" @@ -191,6 +271,10 @@ def load_from_storage( table, source, file_format=file_format, partitioner=partitioner, overwrite=overwrite, **kwargs ) + def _wants_follower_read(self, statement: "SQL") -> bool: + """Return True when a statement may run under a follower-read transaction.""" + return statement.returns_rows() and not statement.is_modifying_operation() + def begin(self) -> None: """Begin a transaction and apply follower-read staleness to it. @@ -203,7 +287,13 @@ def begin(self) -> None: super().begin() return super().begin() - self._apply_follower_reads() + try: + self._apply_follower_reads() + except Exception as exc: + with suppress(Exception): + self.rollback() + self._transaction_active = False + raise create_mapped_exception(exc) from exc def run_transaction_with_retry(self, operation: "Callable[[], _T]") -> _T: """Execute a full CockroachDB transaction callback with serialization retries. @@ -238,7 +328,7 @@ def run_transaction_with_retry(self, operation: "Callable[[], _T]") -> _T: attempt += 1 def dispatch_execute(self, cursor: "CockroachSyncCursor", statement: SQL) -> "ExecutionResult": - auto_started = self._begin_follower_read_transaction() if statement.returns_rows() else False + auto_started = self._begin_follower_read_transaction() if self._wants_follower_read(statement) else False try: result = super().dispatch_execute(cursor, statement) except BaseException: @@ -250,13 +340,36 @@ def dispatch_execute(self, cursor: "CockroachSyncCursor", statement: SQL) -> "Ex self.commit() return result + def dispatch_select_stream(self, statement: "SQL", chunk_size: int) -> "SyncRowStream[dict[str, Any]] | None": + """Return a native psycopg row stream wrapped for follower reads.""" + auto_started = self._begin_follower_read_transaction() if self._wants_follower_read(statement) else False + try: + stream = super().dispatch_select_stream(statement, chunk_size) + if stream is None: + if auto_started and self._connection_in_transaction(): + self.commit() + return None + source = CockroachPsycopgSyncStreamSource( + self, cast("PsycopgSyncStreamSource", stream._source), auto_started + ) + return SyncRowStream(source) + except BaseException: + if auto_started and self._connection_in_transaction(): + with suppress(Exception): + self.rollback() + raise + def handle_database_exceptions(self) -> "CockroachPsycopgSyncExceptionHandler": return CockroachPsycopgSyncExceptionHandler() def _execute_cache_hit( self, sql: str, params: "tuple[Any, ...] | list[Any] | dict[str, Any]", cached: "CachedQuery" ) -> "SQLResult": - auto_started = self._begin_follower_read_transaction() if cached.operation_profile.returns_rows else False + auto_started = ( + self._begin_follower_read_transaction() + if (cached.operation_profile.returns_rows and not cached.operation_profile.modifies_rows) + else False + ) try: result = super()._execute_cache_hit(sql, params, cached) except BaseException: @@ -431,6 +544,10 @@ async def load_from_storage( table, source, file_format=file_format, partitioner=partitioner, overwrite=overwrite, **kwargs ) + def _wants_follower_read(self, statement: "SQL") -> bool: + """Return True when a statement may run under a follower-read transaction.""" + return statement.returns_rows() and not statement.is_modifying_operation() + async def begin(self) -> None: """Begin a transaction and apply follower-read staleness to it. @@ -443,7 +560,13 @@ async def begin(self) -> None: await super().begin() return await super().begin() - await self._apply_follower_reads() + try: + await self._apply_follower_reads() + except Exception as exc: + with suppress(Exception): + await self.rollback() + self._transaction_active = False + raise create_mapped_exception(exc) from exc async def run_transaction_with_retry(self, operation: "Callable[[], Awaitable[_T]]") -> _T: """Execute a full CockroachDB transaction callback with serialization retries. @@ -478,7 +601,7 @@ async def run_transaction_with_retry(self, operation: "Callable[[], Awaitable[_T attempt += 1 async def dispatch_execute(self, cursor: "CockroachAsyncCursor", statement: SQL) -> "ExecutionResult": - auto_started = await self._begin_follower_read_transaction() if statement.returns_rows() else False + auto_started = await self._begin_follower_read_transaction() if self._wants_follower_read(statement) else False try: result = await super().dispatch_execute(cursor, statement) except BaseException: @@ -490,13 +613,25 @@ async def dispatch_execute(self, cursor: "CockroachAsyncCursor", statement: SQL) await self.commit() return result + def dispatch_select_stream(self, statement: "SQL", chunk_size: int) -> "AsyncRowStream[dict[str, Any]] | None": + """Return a native psycopg row stream wrapped for follower reads.""" + stream = super().dispatch_select_stream(statement, chunk_size) + if stream is None: + return None + source = CockroachPsycopgAsyncStreamSource(self, statement, cast("PsycopgAsyncStreamSource", stream._source)) + return AsyncRowStream(source) + def handle_database_exceptions(self) -> "CockroachPsycopgAsyncExceptionHandler": return CockroachPsycopgAsyncExceptionHandler() async def _execute_cache_hit( self, sql: str, params: "tuple[Any, ...] | list[Any] | dict[str, Any]", cached: "CachedQuery" ) -> "SQLResult": - auto_started = await self._begin_follower_read_transaction() if cached.operation_profile.returns_rows else False + auto_started = ( + await self._begin_follower_read_transaction() + if (cached.operation_profile.returns_rows and not cached.operation_profile.modifies_rows) + else False + ) try: result = await super()._execute_cache_hit(sql, params, cached) except BaseException: diff --git a/tests/unit/adapters/test_cockroach_asyncpg/test_config.py b/tests/unit/adapters/test_cockroach_asyncpg/test_config.py index 93eb123a0..400f76e83 100644 --- a/tests/unit/adapters/test_cockroach_asyncpg/test_config.py +++ b/tests/unit/adapters/test_cockroach_asyncpg/test_config.py @@ -549,3 +549,74 @@ async def test_cockroach_asyncpg_adk_locality_validation_and_schema_qualified_dd litestar_store = CockroachAsyncpgStore(good_config) litestar_store._table_name = "tenant_a.litestar_sessions" assert "idx_tenant_a_litestar_sessions_expires_at" in litestar_store._table_ddl() + + +async def test_cockroach_asyncpg_follower_read_only_for_readonly_statements() -> None: + """CockroachAsyncpgDriver should not begin follower-read transactions for modifying operations.""" + in_tx = False + + async def _in_tx() -> bool: + return in_tx + + async def _begin(self: Any) -> None: + nonlocal in_tx + in_tx = True + + with ( + patch.object(CockroachAsyncpgDriver, "_connection_in_transaction", _in_tx), + patch.object(CockroachAsyncpgDriver, "begin", _begin), + ): + conn = MagicMock() + conn.fetch = AsyncMock(return_value=[]) + driver = CockroachAsyncpgDriver( + connection=conn, driver_features={"enable_follower_reads": True, "default_staleness": "'-10s'"} + ) + modifying_stmt = SQL("INSERT INTO users (name) VALUES ('alice') RETURNING id") + await driver.dispatch_execute(conn, modifying_stmt) + assert in_tx is False + + +async def test_cockroach_asyncpg_begin_follower_read_failure_handling() -> None: + """Failing SET TRANSACTION in begin() must rollback and raise mapped exception.""" + conn = MagicMock() + conn.is_in_transaction.return_value = False + tx = AsyncMock() + conn.transaction.return_value = tx + conn.execute = AsyncMock(side_effect=RuntimeError("syntax error in staleness")) + + driver = CockroachAsyncpgDriver( + connection=conn, driver_features={"enable_follower_reads": True, "default_staleness": "'-10s'"} + ) + with pytest.raises(Exception): + await driver.begin() + assert driver._transaction is None + tx.rollback.assert_awaited_once() + + +async def test_cockroach_asyncpg_select_stream_follower_reads() -> None: + """select_stream on a SELECT query with follower reads emits the staleness clause and closes.""" + executed_statements: list[str] = [] + conn = MagicMock() + conn.is_in_transaction.return_value = False + tx = AsyncMock() + conn.transaction.return_value = tx + + async def fake_execute(query: str, *args: Any, **kwargs: Any) -> str: + executed_statements.append(query) + return "OK" + + conn.execute = AsyncMock(side_effect=fake_execute) + + cursor_mock = AsyncMock() + cursor_mock.fetch = AsyncMock(side_effect=[[{"id": 1}], []]) + conn.cursor = AsyncMock(return_value=cursor_mock) + + driver = CockroachAsyncpgDriver( + connection=conn, driver_features={"enable_follower_reads": True, "default_staleness": "'-10s'"} + ) + stream = driver.dispatch_select_stream(SQL("SELECT id FROM users"), chunk_size=10) + assert stream is not None + rows = [row async for row in stream] + assert rows == [{"id": 1}] + assert any("AS OF SYSTEM TIME '-10s'" in s for s in executed_statements) + tx.commit.assert_awaited() diff --git a/tests/unit/adapters/test_cockroach_psycopg/test_config.py b/tests/unit/adapters/test_cockroach_psycopg/test_config.py index 45d747d0c..ad9acd723 100644 --- a/tests/unit/adapters/test_cockroach_psycopg/test_config.py +++ b/tests/unit/adapters/test_cockroach_psycopg/test_config.py @@ -810,3 +810,82 @@ async def __aexit__(self, *_args: Any) -> None: litestar_store = CockroachPsycopgSyncStore(good_cfg) litestar_store._table_name = "tenant_a.litestar_sessions" assert "idx_tenant_a_litestar_sessions_expires_at" in litestar_store._table_ddl() + + +def test_cockroach_psycopg_follower_read_only_for_readonly_statements(monkeypatch: pytest.MonkeyPatch) -> None: + """CockroachPsycopgSyncDriver should not begin follower-read transactions for modifying operations.""" + in_tx = False + + def _in_tx(_self: Any) -> bool: + return in_tx + + def _begin(_self: Any) -> None: + nonlocal in_tx + in_tx = True + + monkeypatch.setattr(CockroachPsycopgSyncDriver, "_connection_in_transaction", _in_tx) + monkeypatch.setattr(CockroachPsycopgSyncDriver, "begin", _begin) + + conn = MagicMock() + cursor = MagicMock() + col = MagicMock() + col.name = "id" + col.type_code = 23 + cursor.description = [col] + cursor.fetchall.return_value = [{"id": 1}] + driver = CockroachPsycopgSyncDriver( + connection=conn, driver_features={"enable_follower_reads": True, "default_staleness": "'-10s'"} + ) + modifying_stmt = SQL("INSERT INTO users (name) VALUES ('alice') RETURNING id") + driver.dispatch_execute(cursor, modifying_stmt) + assert in_tx is False + + +def test_cockroach_psycopg_begin_follower_read_failure_handling() -> None: + """Failing SET TRANSACTION in begin() must rollback, reset transaction active, and raise mapped exception.""" + conn = MagicMock() + conn.autocommit = True + conn.info.transaction_status = 0 + conn.execute.side_effect = RuntimeError("syntax error in staleness") + + driver = CockroachPsycopgSyncDriver( + connection=conn, driver_features={"enable_follower_reads": True, "default_staleness": "'-10s'"} + ) + with pytest.raises(Exception): + driver.begin() + assert driver._transaction_active is False + conn.rollback.assert_called_once() + + +def test_cockroach_psycopg_select_stream_follower_reads() -> None: + """select_stream on a SELECT query with follower reads emits the staleness clause and closes.""" + executed_statements: list[str] = [] + conn = MagicMock() + conn.autocommit = True + conn.info.transaction_status = 0 + + def fake_conn_execute(query: str, *args: Any, **kwargs: Any) -> Any: + executed_statements.append(query) + return MagicMock() + + conn.execute.side_effect = fake_conn_execute + + tx = MagicMock() + conn.transaction.return_value = tx + + stream_cursor = MagicMock() + col = MagicMock() + col.name = "id" + stream_cursor.description = [col] + stream_cursor.fetchmany.side_effect = [[(1,)], []] + conn.cursor.return_value = stream_cursor + + driver = CockroachPsycopgSyncDriver( + connection=conn, driver_features={"enable_follower_reads": True, "default_staleness": "'-10s'"} + ) + stream = driver.dispatch_select_stream(SQL("SELECT id FROM users"), chunk_size=10) + assert stream is not None + rows = list(stream) + assert rows == [{"id": 1}] + assert any("AS OF SYSTEM TIME '-10s'" in s for s in executed_statements) + conn.commit.assert_called() From 2e07be3714fe45048eac8b7887e4984d2c8e59b5 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 01:05:46 +0000 Subject: [PATCH 071/174] fix(cockroach): run user psycopg configure callback after internal setup Preserve user configure callback in CockroachPsycopgSyncConfig and CockroachPsycopgAsyncConfig, invoking it after internal adapter initialization across both pool and standalone connection paths. --- sqlspec/adapters/cockroach_psycopg/config.py | 26 +++++++--- .../test_cockroach_psycopg/test_config.py | 51 ++++++++++++++++++- 2 files changed, 70 insertions(+), 7 deletions(-) diff --git a/sqlspec/adapters/cockroach_psycopg/config.py b/sqlspec/adapters/cockroach_psycopg/config.py index 1bacb0375..6072de71c 100644 --- a/sqlspec/adapters/cockroach_psycopg/config.py +++ b/sqlspec/adapters/cockroach_psycopg/config.py @@ -1,5 +1,6 @@ """CockroachDB configuration using psycopg.""" +import inspect import re from contextlib import suppress from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, NotRequired, TypedDict, cast @@ -279,6 +280,7 @@ def __init__( self._user_connection_hook: Callable[[CockroachSyncConnection], None] | None = features_dict.pop( "on_connection_create", None ) + self._user_configure: Callable[[CockroachSyncConnection], Any] | None = connection_config.pop("configure", None) super().__init__( connection_config=connection_config, @@ -302,9 +304,8 @@ def create_connection(self) -> "CockroachSyncConnection": """ conninfo, connection_kwargs = _standalone_connection_kwargs(dict(self.connection_config)) connection = psycopg_crdb.CrdbConnection.connect(conninfo, **connection_kwargs) - configure = self.connection_config.get("configure", self._configure_connection) try: - configure(cast("CockroachSyncConnection", connection)) + self._configure_connection(cast("CockroachSyncConnection", connection)) except BaseException: with suppress(Exception): connection.close() @@ -380,7 +381,7 @@ def _create_pool(self) -> "ConnectionPool": "close_returns": all_config.pop("close_returns", False), } - pool_parameters["configure"] = all_config.pop("configure", self._configure_connection) + pool_parameters["configure"] = self._configure_connection pool_parameters = {k: v for k, v in pool_parameters.items() if v is not None} conninfo = all_config.pop("conninfo", None) @@ -412,6 +413,12 @@ def _configure_connection(self, conn: "CockroachSyncConnection") -> None: if self._user_connection_hook is not None: self._user_connection_hook(conn) + if self._user_configure is not None: + result = self._user_configure(conn) + if inspect.isawaitable(result): + msg = "Async configure callback cannot be used with CockroachPsycopgSyncConfig" + raise ImproperConfigurationError(msg) + def _close_pool(self) -> None: if not self.connection_instance: return @@ -515,6 +522,9 @@ def __init__( self._user_connection_hook: Callable[[CockroachAsyncConnection], Awaitable[None]] | None = features_dict.pop( "on_connection_create", None ) + self._user_configure: Callable[[CockroachAsyncConnection], Any] | None = connection_config.pop( + "configure", None + ) super().__init__( connection_config=connection_config, @@ -538,9 +548,8 @@ async def create_connection(self) -> "CockroachAsyncConnection": """ conninfo, connection_kwargs = _standalone_connection_kwargs(dict(self.connection_config)) connection = await psycopg_crdb.AsyncCrdbConnection.connect(conninfo, **connection_kwargs) - configure = self.connection_config.get("configure", self._configure_async_connection) try: - await configure(cast("CockroachAsyncConnection", connection)) + await self._configure_async_connection(cast("CockroachAsyncConnection", connection)) except BaseException: with suppress(Exception): await connection.close() @@ -617,7 +626,7 @@ async def _create_pool(self) -> "AsyncConnectionPool": open_pool = all_config.pop("open", True) pool_parameters["open"] = False if open_pool is True else open_pool - pool_parameters["configure"] = all_config.pop("configure", self._configure_async_connection) + pool_parameters["configure"] = self._configure_async_connection pool_parameters = {k: v for k, v in pool_parameters.items() if v is not None} conninfo = all_config.pop("conninfo", None) @@ -655,6 +664,11 @@ async def _configure_async_connection(self, conn: "CockroachAsyncConnection") -> if self._user_connection_hook is not None: await self._user_connection_hook(conn) + if self._user_configure is not None: + result = self._user_configure(conn) + if inspect.isawaitable(result): + await result + async def _close_pool(self) -> None: if not self.connection_instance: return diff --git a/tests/unit/adapters/test_cockroach_psycopg/test_config.py b/tests/unit/adapters/test_cockroach_psycopg/test_config.py index ad9acd723..d5a9322f2 100644 --- a/tests/unit/adapters/test_cockroach_psycopg/test_config.py +++ b/tests/unit/adapters/test_cockroach_psycopg/test_config.py @@ -9,7 +9,7 @@ """ from typing import Any, cast -from unittest.mock import MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -889,3 +889,52 @@ def fake_conn_execute(query: str, *args: Any, **kwargs: Any) -> Any: assert rows == [{"id": 1}] assert any("AS OF SYSTEM TIME '-10s'" in s for s in executed_statements) conn.commit.assert_called() + + +def test_cockroach_psycopg_sync_user_configure_runs_after_internal_setup() -> None: + """User-supplied configure callback runs after internal connection setup in sync config.""" + call_order: list[str] = [] + fake_conn = MagicMock() + fake_conn.autocommit = False + + def user_cb(conn: Any) -> None: + assert conn.autocommit is True + call_order.append("user_cb") + + config = CockroachPsycopgSyncConfig(connection_config={"autocommit": True, "configure": user_cb}) + + with patch("sqlspec.adapters.cockroach_psycopg.config.psycopg_crdb.CrdbConnection.connect", return_value=fake_conn): + conn = config.create_connection() + assert conn is fake_conn + assert call_order == ["user_cb"] + + pool = config.create_pool() + assert pool is not None + + +async def test_cockroach_psycopg_async_user_configure_runs_after_internal_setup() -> None: + """User-supplied async configure callback runs after internal connection setup in async config.""" + call_order: list[str] = [] + fake_conn = MagicMock() + autocommit_val = False + + async def fake_set_autocommit(val: bool) -> None: + nonlocal autocommit_val + autocommit_val = val + call_order.append("internal_autocommit") + + fake_conn.set_autocommit = AsyncMock(side_effect=fake_set_autocommit) + + async def async_user_cb(conn: Any) -> None: + assert autocommit_val is True + call_order.append("user_cb") + + config = CockroachPsycopgAsyncConfig(connection_config={"autocommit": True, "configure": async_user_cb}) + + with patch( + "sqlspec.adapters.cockroach_psycopg.config.psycopg_crdb.AsyncCrdbConnection.connect", + AsyncMock(return_value=fake_conn), + ): + conn = await config.create_connection() + assert conn is fake_conn + assert call_order == ["internal_autocommit", "user_cb"] From 24e674fb89c4fd5254290b458b6833c3167465b0 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 01:05:46 +0000 Subject: [PATCH 072/174] fix(asyncpg): exception mapping, transaction lifecycle, and script execution fixes Map timeout and connection errors in AsyncpgExceptionHandler and lifecycle methods, roll back native transactions if SET LOCAL fails during begin, delegate script execution with statement splitting and observer events, preserve outer caller transactions in stack execution, map truncate and copy errors, and restrict pre-serialized JSON detection to container types. --- sqlspec/adapters/asyncpg/core.py | 17 ++- sqlspec/adapters/asyncpg/driver.py | 108 ++++++++++-------- .../test_asyncpg/test_transaction_state.py | 81 ++++++++++++- .../test_asyncpg/test_type_handlers.py | 6 + 4 files changed, 161 insertions(+), 51 deletions(-) diff --git a/sqlspec/adapters/asyncpg/core.py b/sqlspec/adapters/asyncpg/core.py index 384f59bb3..32782dc1d 100644 --- a/sqlspec/adapters/asyncpg/core.py +++ b/sqlspec/adapters/asyncpg/core.py @@ -1,5 +1,6 @@ """AsyncPG adapter compiled helpers.""" +import asyncio import contextlib import datetime as dt import re @@ -17,6 +18,7 @@ from sqlspec.exceptions import ( CheckViolationError, ConnectionTimeoutError, + DatabaseConnectionError, DeadlockError, ForeignKeyViolationError, ImproperConfigurationError, @@ -25,6 +27,7 @@ OperationalError, OperationCancelledError, PermissionDeniedError, + QueryTimeoutError, SerializationConflictError, SQLParsingError, SQLSpecError, @@ -411,6 +414,13 @@ def create_mapped_exception(error: Any, *, logger: Any | None = None) -> SQLSpec A SQLSpec exception that wraps the original error """ del logger + if isinstance(error, (asyncio.TimeoutError, TimeoutError)): + return QueryTimeoutError(str(error)) + if isinstance(error, OSError) or ( + asyncpg is not None and isinstance(error, getattr(asyncpg, "InterfaceError", ())) + ): + return DatabaseConnectionError(str(error)) + mapped_error = _EXCEPTION_MAPPING_DISPATCHER.get(error) if mapped_error is not None: error_code, error_class, description = mapped_error @@ -554,10 +564,9 @@ def _custom_type_coercions() -> "dict[type, Callable[[Any], Any]]": def _is_pre_serialized_json(value: str) -> bool: stripped = value.strip() - if not stripped: - return False - first = stripped[0] - if first not in '{["-0123456789tfn': + if not ( + (stripped.startswith("{") and stripped.endswith("}")) or (stripped.startswith("[") and stripped.endswith("]")) + ): return False try: from_json(stripped) diff --git a/sqlspec/adapters/asyncpg/driver.py b/sqlspec/adapters/asyncpg/driver.py index 1ab2c718f..89d25325d 100644 --- a/sqlspec/adapters/asyncpg/driver.py +++ b/sqlspec/adapters/asyncpg/driver.py @@ -11,6 +11,7 @@ from sqlglot.errors import ParseError from sqlspec.adapters.asyncpg._typing import AsyncpgCursor, AsyncpgPostgresError, AsyncpgSessionContext +from sqlspec.adapters.asyncpg._typing import asyncpg_module as asyncpg from sqlspec.adapters.asyncpg.core import ( AsyncpgStreamSource, NormalizedStackOperation, @@ -74,7 +75,11 @@ class AsyncpgExceptionHandler(BaseAsyncExceptionHandler): def _handle_exception(self, exc_type: "type[BaseException] | None", exc_val: "BaseException") -> bool: _ = exc_type - if isinstance(exc_val, AsyncpgPostgresError) or has_sqlstate(exc_val): + if ( + isinstance(exc_val, (AsyncpgPostgresError, TimeoutError, OSError)) + or has_sqlstate(exc_val) + or (asyncpg is not None and isinstance(exc_val, getattr(asyncpg, "InterfaceError", ()))) + ): self.pending_exception = create_mapped_exception(exc_val) return True return False @@ -182,20 +187,7 @@ async def dispatch_execute_script(self, cursor: "AsyncpgConnection", statement: Returns: ExecutionResult with script execution details """ - sql, _ = self._compiled_sql(statement, self.statement_config) - statements = self.split_script_statements(sql, statement.statement_config, strip_trailing_semicolon=True) - - successful_count = 0 - last_result = None - - for stmt in statements: - result = await cursor.execute(stmt, **self._timeout_arguments(statement.statement_config)) - last_result = result - successful_count += 1 - - return self.create_execution_result( - last_result, statement_count=len(statements), successful_statements=successful_count, is_script_result=True - ) + return await super().dispatch_execute_script(cursor, statement) async def dispatch_special_handling(self, cursor: "AsyncpgConnection", statement: "SQL") -> "SQLResult | None": """Handle PostgreSQL COPY operations and other special cases. @@ -233,10 +225,15 @@ async def begin(self) -> None: transaction = self.connection.transaction() try: await transaction.start() + except Exception as e: + self._release_failed_transaction_claim(transaction) + raise create_mapped_exception(e) from e + try: if self._should_set_local_columnar(): await self.connection.execute("SET LOCAL google_columnar_engine.enable_columnar_scan = 'on'") - except AsyncpgPostgresError as e: - self._release_failed_transaction_claim(transaction) + except Exception as e: + with suppress(Exception): + await transaction.rollback() raise create_mapped_exception(e) from e self._transaction = transaction @@ -249,7 +246,7 @@ async def commit(self) -> None: await transaction.commit() elif self._connection_in_transaction(): await self.connection.execute("COMMIT") - except AsyncpgPostgresError as e: + except Exception as e: raise create_mapped_exception(e) from e async def rollback(self) -> None: @@ -261,7 +258,7 @@ async def rollback(self) -> None: await transaction.rollback() elif self._connection_in_transaction(): await self.connection.execute("ROLLBACK") - except AsyncpgPostgresError as e: + except Exception as e: raise create_mapped_exception(e) from e async def set_migration_session_schema(self, schema: str) -> None: @@ -385,29 +382,33 @@ async def load_from_arrow( if overwrite: try: await self.connection.execute(f"TRUNCATE TABLE {quoted_target}") - except AsyncpgPostgresError as exc: + except Exception as exc: + mapped = create_mapped_exception(exc) msg = f"Failed to truncate table '{table}': {exc}" - raise SQLSpecError(msg) from exc + raise type(mapped)(msg) from exc columns, batch_iter = self._coerce_arrow_batches(source) first_batch = next((b for b in batch_iter if b.num_rows > 0), None) metrics: dict[str, int] = {"rows_processed": 0, "bytes_processed": 0} if first_batch is not None: all_batches = chain((first_batch,), (b for b in batch_iter if b.num_rows > 0)) use_binary = bool(self.driver_features.get("enable_pgcopy_binary", False)) - if use_binary: - async_chunks = self._iter_pgcopy_binary_chunks_async(all_batches, columns=columns, metrics=metrics) - await self.connection.copy_to_table( - table_name, source=async_chunks, columns=columns, schema_name=schema_name, format="binary" - ) - else: - for batch in all_batches: - metrics["rows_processed"] += int(batch.num_rows) - metrics["bytes_processed"] += int(batch.nbytes) - records = self._arrow_batch_to_rows(batch, columns) - if records: - await self.connection.copy_records_to_table( - table_name, records=records, columns=columns, schema_name=schema_name - ) + try: + if use_binary: + async_chunks = self._iter_pgcopy_binary_chunks_async(all_batches, columns=columns, metrics=metrics) + await self.connection.copy_to_table( + table_name, source=async_chunks, columns=columns, schema_name=schema_name, format="binary" + ) + else: + for batch in all_batches: + metrics["rows_processed"] += int(batch.num_rows) + metrics["bytes_processed"] += int(batch.nbytes) + records = self._arrow_batch_to_rows(batch, columns) + if records: + await self.connection.copy_records_to_table( + table_name, records=records, columns=columns, schema_name=schema_name + ) + except Exception as exc: + raise create_mapped_exception(exc) from exc telemetry_payload = self._ingest_telemetry_from_batches(metrics["rows_processed"], metrics["bytes_processed"]) telemetry_payload["destination"] = table self._attach_partition_telemetry(telemetry_payload, partitioner) @@ -484,12 +485,16 @@ async def load_from_records( if overwrite: try: await self.connection.execute(f"TRUNCATE TABLE {quoted_target}") - except AsyncpgPostgresError as exc: + except Exception as exc: + mapped = create_mapped_exception(exc) msg = f"Failed to truncate table '{table}': {exc}" - raise SQLSpecError(msg) from exc - await self.connection.copy_records_to_table( - table_name, records=copy_rows, columns=resolved_columns, schema_name=schema_name - ) + raise type(mapped)(msg) from exc + try: + await self.connection.copy_records_to_table( + table_name, records=copy_rows, columns=resolved_columns, schema_name=schema_name + ) + except Exception as exc: + raise create_mapped_exception(exc) from exc telemetry_payload: StorageTelemetry = { "bytes_processed": 0, "destination": table, @@ -570,16 +575,19 @@ async def _execute_stack_native( ) -> "tuple[StackResult, ...]": results: list[StackResult] = [] + caller_in_tx = self._connection_in_transaction() transaction_cm = None - if not continue_on_error and not self._connection_in_transaction(): + if not continue_on_error and not caller_in_tx: transaction_cm = self.connection.transaction() with StackExecutionObserver(self, stack, continue_on_error, native_pipeline=False) as observer: if transaction_cm is not None: async with transaction_cm: - await self._run_stack_operations(stack, continue_on_error, observer, results) + await self._run_stack_operations( + stack, continue_on_error, observer, results, caller_in_tx=caller_in_tx + ) else: - await self._run_stack_operations(stack, continue_on_error, observer, results) + await self._run_stack_operations(stack, continue_on_error, observer, results, caller_in_tx=caller_in_tx) return tuple(results) @@ -589,6 +597,8 @@ async def _run_stack_operations( continue_on_error: bool, observer: "StackExecutionObserver", results: "list[StackResult]", + *, + caller_in_tx: bool = False, ) -> None: """Run operations for statement stack execution. @@ -626,19 +636,25 @@ async def _run_stack_operations( mode="continue-on-error" if continue_on_error else "fail-fast", ) if continue_on_error: - await self._rollback_failed_stack() + if not caller_in_tx: + await self._rollback_failed_stack() observer.record_operation_error(stack_error) results.append(StackResult.from_error(stack_error)) continue raise stack_error from exc results.append(stack_result) - if continue_on_error: + if continue_on_error and not caller_in_tx: await self._commit_stack_success() async def _execute_stack_operation_prepared(self, normalized: "NormalizedStackOperation") -> StackResult: result = await self._execute_cached_statement(normalized.statement) - if not self._is_transaction_pooler_mode(): + stmt_cache_size = self.driver_features.get("statement_cache_size") + if stmt_cache_size is None and hasattr(self, "connection_config") and isinstance(self.connection_config, dict): + stmt_cache_size = self.connection_config.get("statement_cache_size") + if stmt_cache_size is None: + stmt_cache_size = 100 + if not self._is_transaction_pooler_mode() and stmt_cache_size > 0: result.metadata["prepared_statement"] = True return StackResult.from_sql_result(result) diff --git a/tests/unit/adapters/test_asyncpg/test_transaction_state.py b/tests/unit/adapters/test_asyncpg/test_transaction_state.py index 6ac4c9341..6a75d83e8 100644 --- a/tests/unit/adapters/test_asyncpg/test_transaction_state.py +++ b/tests/unit/adapters/test_asyncpg/test_transaction_state.py @@ -15,7 +15,7 @@ from sqlspec.adapters.asyncpg.events.backend import AsyncpgEventsBackend, AsyncpgHybridEventsBackend from sqlspec.adapters.asyncpg.litestar.store import AsyncpgStore from sqlspec.core import SQL, StatementConfig, StatementStack -from sqlspec.exceptions import DeadlockError +from sqlspec.exceptions import DatabaseConnectionError, DeadlockError, QueryTimeoutError pytestmark = pytest.mark.anyio @@ -292,3 +292,82 @@ async def test_create_connection_drops_every_pool_only_setting(monkeypatch: pyte accepted = set(inspect.signature(asyncpg.connect).parameters) assert set(connect.call_args.kwargs) <= accepted + + +async def test_asyncpg_begin_rollback_on_set_local_failure() -> None: + """A failure during SET LOCAL in begin() must rollback the native transaction and raise mapped error.""" + conn = MagicMock() + tx = AsyncMock() + conn.transaction.return_value = tx + conn.execute = AsyncMock(side_effect=RuntimeError("columnar scan configuration error")) + conn.is_in_transaction.return_value = False + + driver = AsyncpgDriver( + cast("Any", conn), driver_features={"enable_alloydb_columnar": True, "pooler_mode": "transaction"} + ) + with pytest.raises(Exception): + await driver.begin() + assert driver._transaction is None + tx.rollback.assert_awaited_once() + + +async def test_asyncpg_dispatch_execute_script_with_parameters_and_observers() -> None: + """dispatch_execute_script executes sub-statements, emits observer events, and returns cursor.""" + conn = MagicMock() + conn.execute = AsyncMock(return_value="CREATE 1") + driver = AsyncpgDriver(cast("Any", conn)) + + stmt = SQL("CREATE TABLE t1 (id int); CREATE TABLE t2 (id int);") + result = await driver.dispatch_execute_script(cast("Any", conn), stmt) + + assert result.is_script_result is True + assert result.statement_count == 2 + assert result.successful_statements == 2 + assert result.cursor_result is conn + + +async def test_asyncpg_execute_stack_in_outer_transaction_does_not_commit_or_rollback() -> None: + """execute_stack with continue_on_error inside an outer transaction must not commit or rollback.""" + conn = MagicMock() + conn.is_in_transaction = MagicMock(return_value=True) + conn.execute = AsyncMock(side_effect=["UPDATE 1", RuntimeError("failed op")]) + + driver = AsyncpgDriver(cast("Any", conn)) + driver._transaction = MagicMock() + + stack = StatementStack().push_execute("UPDATE t SET x = 1").push_execute("UPDATE t SET x = 2") + + results = await driver.execute_stack(stack, continue_on_error=True) + assert len(results) == 2 + assert results[0].result is not None + assert results[1].error is not None + assert driver._transaction is not None + conn.execute.assert_called() + + +async def test_asyncpg_exception_mapping_timeout_and_interface_errors() -> None: + """AsyncpgExceptionHandler and create_mapped_exception map TimeoutError and InterfaceError/OSError.""" + driver = AsyncpgDriver(cast("Any", MagicMock())) + + handler = driver.handle_database_exceptions() + with pytest.raises(QueryTimeoutError): + async with handler: + raise TimeoutError("query timed out") + driver._check_pending_exception(handler) + + handler2 = driver.handle_database_exceptions() + with pytest.raises(DatabaseConnectionError): + async with handler2: + raise OSError("connection reset by peer") + driver._check_pending_exception(handler2) + + +async def test_asyncpg_truncate_and_copy_map_errors() -> None: + """Truncate and copy operations map errors rather than leaking unmapped exceptions.""" + conn = MagicMock() + conn.execute = AsyncMock(side_effect=RuntimeError("truncate table failed")) + driver = AsyncpgDriver(cast("Any", conn)) + + with pytest.raises(Exception) as exc_info: + await driver.load_from_records("users", [{"id": 1}], overwrite=True) + assert not isinstance(exc_info.value, RuntimeError) or isinstance(exc_info.value, DatabaseConnectionError) diff --git a/tests/unit/adapters/test_asyncpg/test_type_handlers.py b/tests/unit/adapters/test_asyncpg/test_type_handlers.py index 6b360b625..146ec7074 100644 --- a/tests/unit/adapters/test_asyncpg/test_type_handlers.py +++ b/tests/unit/adapters/test_asyncpg/test_type_handlers.py @@ -42,11 +42,17 @@ async def test_register_json_codecs_success() -> None: assert json_call.kwargs["encoder"]({"value": 1}) == b'{"value":1}' assert json_call.kwargs["encoder"]("console") == b'"console"' + assert json_call.kwargs["encoder"]("123") == b'"123"' + assert json_call.kwargs["encoder"]("true") == b'"true"' + assert json_call.kwargs["encoder"]("null") == b'"null"' assert json_call.kwargs["encoder"]('{"pre":"serialized"}') == b'{"pre":"serialized"}' assert json_call.kwargs["encoder"](b'\x01{"value":1}') == b'{"value":1}' assert json_call.kwargs["decoder"](b'{"value":1}') == {"value": 1} assert jsonb_call.kwargs["encoder"]({"value": 1}) == b'\x01{"value":1}' assert jsonb_call.kwargs["encoder"]("console") == b'\x01"console"' + assert jsonb_call.kwargs["encoder"]("123") == b'\x01"123"' + assert jsonb_call.kwargs["encoder"]("true") == b'\x01"true"' + assert jsonb_call.kwargs["encoder"]("null") == b'\x01"null"' assert jsonb_call.kwargs["encoder"]('{"pre":"serialized"}') == b'\x01{"pre":"serialized"}' assert jsonb_call.kwargs["encoder"](b'\x01{"value":1}') == b'\x01{"value":1}' assert jsonb_call.kwargs["decoder"](b'\x01{"value":1}') == {"value": 1} From 40c0dfd2dc6329f371945a6d2be12ce242724ecb Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 01:05:46 +0000 Subject: [PATCH 073/174] fix(asyncpg): ADK, Litestar, and events store enhancements Consolidate table storage parameters into build_postgres_storage_params_sql, validate SCaNN quantizer before interpolation in memory table DDL, remove unused POSTGRES_TABLE_NOT_FOUND_ERROR, parse NOTIFY payloads once safely in events backend, and allow Mapping direct_dsn while bypassing pool capacity checks when direct_dsn is configured. --- sqlspec/adapters/asyncpg/adk/store.py | 53 ++++++----------- sqlspec/adapters/asyncpg/core.py | 34 +++++++++++ sqlspec/adapters/asyncpg/events/_hub.py | 3 +- sqlspec/adapters/asyncpg/events/backend.py | 59 +++++++++++-------- sqlspec/adapters/asyncpg/events/store.py | 26 +------- sqlspec/adapters/asyncpg/litestar/store.py | 35 ++++------- .../adapters/test_asyncpg/test_adk_store.py | 9 +++ .../test_asyncpg/test_transaction_state.py | 59 ++++++++++++++++++- 8 files changed, 169 insertions(+), 109 deletions(-) diff --git a/sqlspec/adapters/asyncpg/adk/store.py b/sqlspec/adapters/asyncpg/adk/store.py index 139e289d5..e45c50148 100644 --- a/sqlspec/adapters/asyncpg/adk/store.py +++ b/sqlspec/adapters/asyncpg/adk/store.py @@ -1,11 +1,12 @@ """AsyncPG ADK store for Google Agent Development Kit session/event storage.""" from contextlib import nullcontext -from typing import TYPE_CHECKING, Any, Final, Literal, NotRequired, cast +from typing import TYPE_CHECKING, Any, Literal, NotRequired, cast from sqlspec.adapters.asyncpg._typing import asyncpg_module as asyncpg -from sqlspec.adapters.asyncpg.core import parse_status +from sqlspec.adapters.asyncpg.core import build_postgres_storage_params_sql, parse_status from sqlspec.config import ADKConfig, AsyncConfigT +from sqlspec.exceptions import ImproperConfigurationError from sqlspec.extensions.adk import BaseAsyncADKStore, StoredEvent, StoredSession, normalize_session_list_options from sqlspec.extensions.adk.memory.store import BaseAsyncADKMemoryStore @@ -19,8 +20,6 @@ __all__ = ("AsyncpgADKConfig", "AsyncpgADKMemoryStore", "AsyncpgADKStore") -POSTGRES_TABLE_NOT_FOUND_ERROR: Final = "42P01" - class AsyncpgADKConfig(ADKConfig): """Asyncpg-specific ADK extension settings. @@ -480,7 +479,9 @@ async def _sessions_table_ddl(self) -> str: if self._owner_id_column_ddl: owner_id_line = f",\n {self._owner_id_column_ddl}" - table_options = _postgres_table_options(_adk_config(self._config)) + table_options = build_postgres_storage_params_sql( + _adk_config(self._config), default_fillfactor=80, section="adk" + ) return f""" CREATE TABLE IF NOT EXISTS {self._session_table} ( id VARCHAR(128) PRIMARY KEY, @@ -504,7 +505,9 @@ async def _sessions_table_ddl(self) -> str: async def _events_table_ddl(self) -> str: adk_config = _adk_config(self._config) - table_options = _postgres_table_options(adk_config, include_autovacuum=True) + table_options = build_postgres_storage_params_sql( + adk_config, include_autovacuum=True, default_fillfactor=80, section="adk" + ) generated_columns = "" generated_indexes = "" if adk_config.get("enable_event_generated_columns", False): @@ -544,7 +547,9 @@ async def _events_table_ddl(self) -> str: """ async def _app_states_table_ddl(self) -> str: - table_options = _postgres_table_options(_adk_config(self._config)) + table_options = build_postgres_storage_params_sql( + _adk_config(self._config), default_fillfactor=80, section="adk" + ) return f""" CREATE TABLE IF NOT EXISTS {self._app_state_table} ( app_name VARCHAR(128) PRIMARY KEY, @@ -554,7 +559,9 @@ async def _app_states_table_ddl(self) -> str: """ async def _user_states_table_ddl(self) -> str: - table_options = _postgres_table_options(_adk_config(self._config)) + table_options = build_postgres_storage_params_sql( + _adk_config(self._config), default_fillfactor=80, section="adk" + ) return f""" CREATE TABLE IF NOT EXISTS {self._user_state_table} ( app_name VARCHAR(128) NOT NULL, @@ -843,6 +850,9 @@ async def _memory_table_ddl(self) -> str: ) if self._vector_index_type == "scann": + if self._scann_quantizer.upper() not in {"SQ8", "FP32", "AH"}: + msg = f"Invalid SCaNN quantizer '{self._scann_quantizer}'. Must be one of: SQ8, FP32, AH" + raise ImproperConfigurationError(msg) indexes.append( f"CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_scann ON {self._memory_table} USING scann (embedding) WITH (num_leaves = {self._scann_num_leaves}, quantizer = '{self._scann_quantizer}');" ) @@ -900,30 +910,3 @@ def _adk_config(config: Any) -> AsyncpgADKConfig: if not isinstance(adk_config, dict): return {} return cast("AsyncpgADKConfig", adk_config) - - -def _postgres_table_options(adk_config: AsyncpgADKConfig, *, include_autovacuum: bool = False) -> str: - options = [_postgres_fillfactor_option(adk_config)] - if include_autovacuum: - options.extend(_postgres_autovacuum_options(adk_config)) - return f" WITH ({', '.join(options)})" - - -def _postgres_fillfactor_option(adk_config: AsyncpgADKConfig) -> str: - value = adk_config.get("fillfactor", 80) - if not isinstance(value, int) or isinstance(value, bool) or value not in range(10, 101): - msg = "extension_config['adk']['fillfactor'] must be an integer from 10 to 100" - raise ValueError(msg) - return f"fillfactor = {value}" - - -def _postgres_autovacuum_options(adk_config: AsyncpgADKConfig) -> "list[str]": - options: list[str] = [] - for key in ("autovacuum_vacuum_scale_factor", "autovacuum_analyze_scale_factor"): - value = adk_config.get(key) - if value is not None: - if not isinstance(value, (int, float)) or isinstance(value, bool) or not 0 <= float(value) <= 1: - msg = f"extension_config['adk']['{key}'] must be a number from 0 to 1" - raise ValueError(msg) - options.append(f"{key} = {float(value):g}") - return options diff --git a/sqlspec/adapters/asyncpg/core.py b/sqlspec/adapters/asyncpg/core.py index 32782dc1d..f969923c1 100644 --- a/sqlspec/adapters/asyncpg/core.py +++ b/sqlspec/adapters/asyncpg/core.py @@ -54,6 +54,7 @@ "apply_driver_features", "build_connection_config", "build_postgres_extension_probe_names", + "build_postgres_storage_params_sql", "build_profile", "build_statement_config", "collect_rows", @@ -141,6 +142,39 @@ def build_connection_config(connection_config: "Mapping[str, Any]") -> "dict[str return config +def build_postgres_storage_params_sql( + settings: "dict[str, object] | Mapping[str, object] | None", + *, + include_autovacuum: bool = False, + default_fillfactor: int | None = None, + default_vacuum_scale: float | None = None, + default_analyze_scale: float | None = None, + section: str = "storage", +) -> str: + """Format PostgreSQL table storage parameters (WITH fillfactor = ...) clause.""" + options: list[str] = [] + config_dict = dict(settings) if settings else {} + fillfactor = config_dict.get("fillfactor", default_fillfactor) + if fillfactor is not None: + if not isinstance(fillfactor, int) or isinstance(fillfactor, bool) or fillfactor not in range(10, 101): + msg = f"extension_config['{section}']['fillfactor'] must be an integer from 10 to 100" + raise ValueError(msg) + options.append(f"fillfactor = {fillfactor}") + if include_autovacuum: + scales = ( + ("autovacuum_vacuum_scale_factor", default_vacuum_scale), + ("autovacuum_analyze_scale_factor", default_analyze_scale), + ) + for key, default_val in scales: + value = config_dict.get(key, default_val) + if value is not None: + if not isinstance(value, (int, float)) or isinstance(value, bool) or not 0 <= float(value) <= 1: + msg = f"extension_config['{section}']['{key}'] must be a number from 0 to 1" + raise ValueError(msg) + options.append(f"{key} = {float(value):g}") + return f" WITH ({', '.join(options)})" if options else "" + + def build_profile() -> "DriverParameterProfile": """Create the AsyncPG driver parameter profile.""" diff --git a/sqlspec/adapters/asyncpg/events/_hub.py b/sqlspec/adapters/asyncpg/events/_hub.py index 5cafccc63..a20373d8d 100644 --- a/sqlspec/adapters/asyncpg/events/_hub.py +++ b/sqlspec/adapters/asyncpg/events/_hub.py @@ -12,6 +12,7 @@ import asyncio import contextlib import logging +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, cast from weakref import WeakKeyDictionary @@ -187,7 +188,7 @@ def _resolve_listener_config(self) -> "AsyncpgConfig": ) ).lower() raw_direct_dsn = driver_features.get("direct_dsn") - direct_dsn = raw_direct_dsn if isinstance(raw_direct_dsn, str) and raw_direct_dsn else None + direct_dsn = raw_direct_dsn if isinstance(raw_direct_dsn, (str, Mapping)) and raw_direct_dsn else None if pooler_mode in {"transaction", "statement"} and not direct_dsn: logger.warning( "PostgreSQL LISTEN/NOTIFY cannot run under pooler_mode=%r without configuring driver_features['direct_dsn']", diff --git a/sqlspec/adapters/asyncpg/events/backend.py b/sqlspec/adapters/asyncpg/events/backend.py index 44f8eef83..1e71b037e 100644 --- a/sqlspec/adapters/asyncpg/events/backend.py +++ b/sqlspec/adapters/asyncpg/events/backend.py @@ -105,14 +105,15 @@ async def dequeue(self, channel: str, poll_interval: float) -> EventMessage | No return await self._dequeue_reconciled(channel) self._queue._reset_empty_poll_delay(channel) self._runtime.increment_metric("events.listener.wakeup") - batch_size = _extract_batch_size(payload) - is_new, should_drain = self._marker_state.register(channel, _extract_marker_id(payload), batch_size) + parsed_payload = _parse_notify_payload(payload) + batch_size = _extract_batch_size(parsed_payload) + is_new, should_drain = self._marker_state.register(channel, _extract_marker_id(parsed_payload), batch_size) if is_new and batch_size > 1: self._runtime.increment_metric("events.marker.coalesced", batch_size - 1) if not should_drain: self._runtime.increment_metric("events.marker.duplicate") continue - event_id = _extract_event_id(payload) + event_id = _extract_event_id(parsed_payload) if event_id is not None: event = await self._queue.dequeue_by_event_id(event_id) if event is not None: @@ -353,40 +354,50 @@ def reset(self) -> None: def _validate_listener_pool_capacity(config: "AsyncpgConfig") -> None: + raw_features = config.driver_features + driver_features = raw_features if isinstance(raw_features, dict) else {} + if driver_features.get("direct_dsn"): + return max_size = config.connection_config.get("max_size") if max_size is not None and max_size < _MIN_LISTENER_POOL_SIZE: msg = f"{type(config).__name__} native event listeners require pool max_size >= {_MIN_LISTENER_POOL_SIZE}" raise ImproperConfigurationError(msg) -def _extract_event_id(payload: "str | None") -> "str | None": - if not payload: +def _parse_notify_payload(payload: "dict[str, Any] | str | None") -> "dict[str, Any] | None": + if isinstance(payload, dict): + return payload + if not payload or not isinstance(payload, str): + return None + try: + raw = from_json(payload) + return raw if isinstance(raw, dict) else None + except (TypeError, ValueError): + return None + + +def _extract_event_id(payload: "dict[str, Any] | str | None") -> "str | None": + parsed = payload if isinstance(payload, dict) else _parse_notify_payload(payload) + if parsed is None: return None - raw = from_json(payload) - if isinstance(raw, dict): - event_id = raw.get("event_id") - return event_id if isinstance(event_id, str) else None - return None + event_id = parsed.get("event_id") + return event_id if isinstance(event_id, str) else None -def _extract_batch_size(payload: "str | None") -> int: - if not payload: +def _extract_batch_size(payload: "dict[str, Any] | str | None") -> int: + parsed = payload if isinstance(payload, dict) else _parse_notify_payload(payload) + if parsed is None: return 0 - raw = from_json(payload) - if isinstance(raw, dict): - batch_size = raw.get("batch_size") - return batch_size if isinstance(batch_size, int) else 0 - return 0 + batch_size = parsed.get("batch_size") + return batch_size if isinstance(batch_size, int) else 0 -def _extract_marker_id(payload: "str | None") -> "str | None": - if not payload: +def _extract_marker_id(payload: "dict[str, Any] | str | None") -> "str | None": + parsed = payload if isinstance(payload, dict) else _parse_notify_payload(payload) + if parsed is None: return None - raw = from_json(payload) - if isinstance(raw, dict): - marker_id = raw.get("marker_id", raw.get("event_id")) - return marker_id if isinstance(marker_id, str) else None - return None + marker_id = parsed.get("marker_id", parsed.get("event_id")) + return marker_id if isinstance(marker_id, str) else None def _record_dequeue_result(runtime: Any, event: "EventMessage | None") -> None: diff --git a/sqlspec/adapters/asyncpg/events/store.py b/sqlspec/adapters/asyncpg/events/store.py index 9c55518df..ecd9a7edc 100644 --- a/sqlspec/adapters/asyncpg/events/store.py +++ b/sqlspec/adapters/asyncpg/events/store.py @@ -3,6 +3,7 @@ from typing import NotRequired from sqlspec.adapters.asyncpg.config import AsyncpgConfig +from sqlspec.adapters.asyncpg.core import build_postgres_storage_params_sql from sqlspec.config import EventsConfig from sqlspec.extensions.events import BaseEventQueueStore @@ -53,27 +54,4 @@ def _column_types(self) -> "tuple[str, str, str]": def _table_clause(self) -> str: """Return explicitly configured PostgreSQL table tuning options.""" - return _postgres_table_options(self.settings) - - -def _postgres_table_options(settings: "dict[str, object]") -> str: - options = _postgres_option_values(settings) - return f" WITH ({', '.join(options)})" if options else "" - - -def _postgres_option_values(settings: "dict[str, object]") -> "list[str]": - options: list[str] = [] - fillfactor = settings.get("fillfactor") - if fillfactor is not None: - if not isinstance(fillfactor, int) or isinstance(fillfactor, bool) or fillfactor not in range(10, 101): - msg = "extension_config['events']['fillfactor'] must be an integer from 10 to 100" - raise ValueError(msg) - options.append(f"fillfactor = {fillfactor}") - for key in ("autovacuum_vacuum_scale_factor", "autovacuum_analyze_scale_factor"): - value = settings.get(key) - if value is not None: - if not isinstance(value, (int, float)) or isinstance(value, bool) or not 0 <= float(value) <= 1: - msg = f"extension_config['events']['{key}'] must be a number from 0 to 1" - raise ValueError(msg) - options.append(f"{key} = {float(value):g}") - return options + return build_postgres_storage_params_sql(self.settings, include_autovacuum=True, section="events") diff --git a/sqlspec/adapters/asyncpg/litestar/store.py b/sqlspec/adapters/asyncpg/litestar/store.py index cb8ff5863..30ee2cd74 100644 --- a/sqlspec/adapters/asyncpg/litestar/store.py +++ b/sqlspec/adapters/asyncpg/litestar/store.py @@ -1,10 +1,10 @@ """AsyncPG session store for Litestar integration.""" from datetime import UTC, datetime, timedelta -from typing import TYPE_CHECKING, Any, NotRequired, cast +from typing import TYPE_CHECKING, NotRequired from sqlspec.adapters.asyncpg._typing import asyncpg_module as asyncpg -from sqlspec.adapters.asyncpg.core import parse_status +from sqlspec.adapters.asyncpg.core import build_postgres_storage_params_sql, parse_status from sqlspec.config import LitestarConfig from sqlspec.extensions.litestar.store import BaseSQLSpecStore @@ -240,7 +240,14 @@ def _table_ddl(self) -> str: Returns: SQL statement to create the sessions table with proper indexes. """ - fillfactor, vacuum_scale, analyze_scale = _postgres_litestar_tuning(self._config) + tuning_clause = build_postgres_storage_params_sql( + self._config.extension_config.get("litestar", {}), + default_fillfactor=80, + default_vacuum_scale=0.05, + default_analyze_scale=0.02, + include_autovacuum=True, + section="litestar", + ) return f""" CREATE TABLE IF NOT EXISTS {self._table_name} ( session_id TEXT PRIMARY KEY, @@ -248,11 +255,7 @@ def _table_ddl(self) -> str: expires_at TIMESTAMPTZ, created_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP, updated_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP - ) WITH ( - fillfactor = {fillfactor}, - autovacuum_vacuum_scale_factor = {vacuum_scale:g}, - autovacuum_analyze_scale_factor = {analyze_scale:g} - ); + ){tuning_clause}; CREATE INDEX IF NOT EXISTS idx_{self._table_name}_expires_at ON {self._table_name}(expires_at) WHERE expires_at IS NOT NULL; @@ -265,19 +268,3 @@ def _drop_table_sql(self) -> "list[str]": List of SQL statements to drop indexes and table. """ return [f"DROP INDEX IF EXISTS idx_{self._table_name}_expires_at", f"DROP TABLE IF EXISTS {self._table_name}"] - - -def _postgres_litestar_tuning(config: Any) -> "tuple[int, float, float]": - settings = cast("dict[str, Any]", config.extension_config.get("litestar", {})) - fillfactor = settings.get("fillfactor", 80) - if not isinstance(fillfactor, int) or isinstance(fillfactor, bool) or fillfactor not in range(10, 101): - msg = "extension_config['litestar']['fillfactor'] must be an integer from 10 to 100" - raise ValueError(msg) - scales: list[float] = [] - for key, default in (("autovacuum_vacuum_scale_factor", 0.05), ("autovacuum_analyze_scale_factor", 0.02)): - value = settings.get(key, default) - if not isinstance(value, (int, float)) or isinstance(value, bool) or not 0 <= float(value) <= 1: - msg = f"extension_config['litestar']['{key}'] must be a number from 0 to 1" - raise ValueError(msg) - scales.append(float(value)) - return fillfactor, scales[0], scales[1] diff --git a/tests/unit/adapters/test_asyncpg/test_adk_store.py b/tests/unit/adapters/test_asyncpg/test_adk_store.py index 25a519155..fbba4de18 100644 --- a/tests/unit/adapters/test_asyncpg/test_adk_store.py +++ b/tests/unit/adapters/test_asyncpg/test_adk_store.py @@ -312,3 +312,12 @@ async def test_asyncpg_list_sessions_rejects_invalid_options_before_connecting(o await store.list_sessions("app", **options) assert conn.calls == [] + + +async def test_asyncpg_adk_scann_quantizer_validation() -> None: + """Invalid SCaNN quantizer must raise ImproperConfigurationError.""" + settings = AsyncpgADKConfig(vector_index_type="scann", scann_quantizer="INVALID") + config = AsyncpgConfig(extension_config={"adk": settings}) + store = AsyncpgADKMemoryStore(config) + with pytest.raises(ImproperConfigurationError, match="Invalid SCaNN quantizer"): + await store._memory_table_ddl() diff --git a/tests/unit/adapters/test_asyncpg/test_transaction_state.py b/tests/unit/adapters/test_asyncpg/test_transaction_state.py index 6a75d83e8..4fcba4318 100644 --- a/tests/unit/adapters/test_asyncpg/test_transaction_state.py +++ b/tests/unit/adapters/test_asyncpg/test_transaction_state.py @@ -10,7 +10,7 @@ import pytest from sqlspec.adapters.asyncpg.config import AsyncpgConfig -from sqlspec.adapters.asyncpg.core import AsyncpgStreamSource +from sqlspec.adapters.asyncpg.core import AsyncpgStreamSource, build_postgres_storage_params_sql from sqlspec.adapters.asyncpg.driver import AsyncpgDriver from sqlspec.adapters.asyncpg.events.backend import AsyncpgEventsBackend, AsyncpgHybridEventsBackend from sqlspec.adapters.asyncpg.litestar.store import AsyncpgStore @@ -371,3 +371,60 @@ async def test_asyncpg_truncate_and_copy_map_errors() -> None: with pytest.raises(Exception) as exc_info: await driver.load_from_records("users", [{"id": 1}], overwrite=True) assert not isinstance(exc_info.value, RuntimeError) or isinstance(exc_info.value, DatabaseConnectionError) + + +def test_asyncpg_build_postgres_storage_params_sql() -> None: + """build_postgres_storage_params_sql consolidates storage and autovacuum clauses.""" + settings = {"fillfactor": 85, "autovacuum_vacuum_scale_factor": 0.05, "autovacuum_analyze_scale_factor": 0.02} + clause = build_postgres_storage_params_sql(settings, include_autovacuum=True) + assert "WITH (" in clause + assert "fillfactor = 85" in clause + assert "autovacuum_vacuum_scale_factor = 0.05" in clause + assert "autovacuum_analyze_scale_factor = 0.02" in clause + + +def test_asyncpg_events_validate_listener_pool_capacity_skips_when_direct_dsn() -> None: + """_validate_listener_pool_capacity should not enforce pool size when direct_dsn is configured.""" + from sqlspec.adapters.asyncpg.events.backend import _validate_listener_pool_capacity + + config = AsyncpgConfig( + connection_config={"dsn": "postgresql://localhost/test", "max_size": 1}, + driver_features={"direct_dsn": "postgresql://localhost/direct"}, + ) + _validate_listener_pool_capacity(config) + + +def test_asyncpg_events_parse_notify_payload_handles_malformed_json() -> None: + """Malformed NOTIFY payload does not raise and returns None safely.""" + from sqlspec.adapters.asyncpg.events.backend import ( + _extract_batch_size, + _extract_event_id, + _extract_marker_id, + _parse_notify_payload, + ) + + bad_payload = "not valid json {{" + assert _parse_notify_payload(bad_payload) is None + assert _extract_event_id(bad_payload) is None + assert _extract_batch_size(bad_payload) == 0 + assert _extract_marker_id(bad_payload) is None + + good_payload = '{"event_id": "e1", "batch_size": 5, "marker_id": "m1"}' + parsed = _parse_notify_payload(good_payload) + assert parsed is not None + assert _extract_event_id(parsed) == "e1" + assert _extract_batch_size(parsed) == 5 + assert _extract_marker_id(parsed) == "m1" + + +def test_asyncpg_events_resolve_listener_config_allows_mapping_direct_dsn() -> None: + """_resolve_listener_config accepts Mapping for direct_dsn without raising.""" + from sqlspec.adapters.asyncpg.events._hub import AsyncpgListenerHub + + config = AsyncpgConfig( + connection_config={"dsn": "postgresql://localhost/test"}, + driver_features={"pooler_mode": "transaction", "direct_dsn": {"host": "localhost", "database": "direct_db"}}, + ) + hub = AsyncpgListenerHub(config) + resolved = hub._resolve_listener_config() + assert resolved is not None From 21fa8236663507f196638101bbfbd02f41a57c61 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 01:05:46 +0000 Subject: [PATCH 074/174] fix(asyncpg): data dictionary and postgres catalog query fixes Query unqualified pg_stat_statements in system metadata query, refactor constraints.sql projection from correlated array subqueries to pre-aggregated CTEs with LATERAL unnest, allow get_indexes to resolve table via search_path when schema is None, and consolidate PostgreSQL metadata domains into dialect config. --- sqlspec/adapters/asyncpg/data_dictionary.py | 53 +++++-------------- .../dialects/postgres/__init__.py | 12 ++++- .../dialects/postgres/config.py | 33 ++++++++++++ .../dialects/postgres/sql/constraints.sql | 34 +++++++----- .../dialects/postgres/sql/system.sql | 2 +- .../test_asyncpg/test_transaction_state.py | 20 +++++++ .../test_postgres_cockroach_query_packs.py | 15 ++++++ 7 files changed, 116 insertions(+), 53 deletions(-) diff --git a/sqlspec/adapters/asyncpg/data_dictionary.py b/sqlspec/adapters/asyncpg/data_dictionary.py index a2978b941..a212bece4 100644 --- a/sqlspec/adapters/asyncpg/data_dictionary.py +++ b/sqlspec/adapters/asyncpg/data_dictionary.py @@ -26,7 +26,12 @@ system_metadata_gated_result, unsupported_system_metadata_capability, ) -from sqlspec.data_dictionary.dialects.postgres import resolve_postgres_json_type +from sqlspec.data_dictionary.dialects.postgres import ( + POSTGRES_METADATA_DOMAINS, + POSTGRES_SUPPORTED_DOMAINS, + POSTGRES_SYSTEM_METADATA_QUERIES, + resolve_postgres_json_type, +) from sqlspec.driver import AsyncDataDictionaryBase if TYPE_CHECKING: @@ -36,37 +41,6 @@ __all__ = ("AsyncpgDataDictionary",) -_POSTGRES_METADATA_DOMAINS: tuple[str, ...] = ( - "schemas", - "objects", - "tables", - "columns", - "constraints", - "indexes", - "views", - "materialized_views", - "sequences", - "types", - "routines", - "triggers", - "comments", - "privileges", - "dependencies", - "extensions", - "partitions", - "ddl", - "system", -) -_POSTGRES_SUPPORTED_DOMAINS = frozenset(d for d in _POSTGRES_METADATA_DOMAINS if d != "system") -_POSTGRES_SYSTEM_METADATA_QUERIES = { - "settings": "settings", - "statement_history": "pg_stat_statements", - "pg_stat_statements": "pg_stat_statements", - "table_statistics": "table_stats", - "table_stats": "table_stats", -} - - class AsyncpgDataDictionary(AsyncDataDictionaryBase): """PostgreSQL-specific async data dictionary.""" @@ -83,7 +57,7 @@ async def get_system_metadata_capabilities( ) -> tuple[SystemMetadataCapability, ...]: """Get PostgreSQL opt-in system metadata capability disclosures.""" _ = driver - requested_domains = tuple(_POSTGRES_SYSTEM_METADATA_QUERIES) if domains is None else tuple(domains) + requested_domains = tuple(POSTGRES_SYSTEM_METADATA_QUERIES) if domains is None else tuple(domains) return tuple(_postgres_system_metadata_capability(domain) for domain in requested_domains) async def get_schemas(self, driver: "AsyncpgDriver") -> MetadataResult: @@ -157,7 +131,7 @@ async def get_system_metadata( if gate_result.capability.support != MetadataSupport.SUPPORTED: return gate_result - query_name = _POSTGRES_SYSTEM_METADATA_QUERIES.get(metadata_request.domain) + query_name = POSTGRES_SYSTEM_METADATA_QUERIES.get(metadata_request.domain) if query_name is None: return gate_result parameters: dict[str, object] = {} @@ -284,10 +258,11 @@ async def get_indexes( ) table_name = self.resolve_identifier(table) - self._log_table_describe(driver, schema_name=schema_name, table_name=table_name, operation="indexes") + table_schema_name = None if schema is None else schema_name + self._log_table_describe(driver, schema_name=table_schema_name, table_name=table_name, operation="indexes") return await driver.select( _postgres_domain_sql("indexes", "by_schema"), - schema_name=schema_name, + schema_name=table_schema_name, table_name=table_name, schema_type=IndexMetadata, ) @@ -352,7 +327,7 @@ def _postgres_metadata_capability(domain: str) -> MetadataCapability: risks=(MetadataRisk.EXPENSIVE, MetadataRisk.PRIVILEGED), warnings=("System metadata is opt-in and disabled by default.",), ) - if domain in _POSTGRES_SUPPORTED_DOMAINS: + if domain in POSTGRES_SUPPORTED_DOMAINS: return MetadataCapability( domain=domain, support=MetadataSupport.SUPPORTED, @@ -363,7 +338,7 @@ def _postgres_metadata_capability(domain: str) -> MetadataCapability: def _postgres_metadata_profile(adapter: str, domains: Sequence[str] | None) -> MetadataCapabilityProfile: - requested_domains = _POSTGRES_METADATA_DOMAINS if domains is None else tuple(domains) + requested_domains = POSTGRES_METADATA_DOMAINS if domains is None else tuple(domains) return MetadataCapabilityProfile( "postgres", adapter=adapter, @@ -422,7 +397,7 @@ def _ddl_result_from_rows( def _postgres_system_metadata_capability(domain: str) -> SystemMetadataCapability: - if domain not in _POSTGRES_SYSTEM_METADATA_QUERIES: + if domain not in POSTGRES_SYSTEM_METADATA_QUERIES: return unsupported_system_metadata_capability(domain) return SystemMetadataCapability( domain, diff --git a/sqlspec/data_dictionary/dialects/postgres/__init__.py b/sqlspec/data_dictionary/dialects/postgres/__init__.py index d15fc9558..ab166164a 100644 --- a/sqlspec/data_dictionary/dialects/postgres/__init__.py +++ b/sqlspec/data_dictionary/dialects/postgres/__init__.py @@ -2,8 +2,18 @@ from sqlspec.data_dictionary.dialects.postgres.config import ( POSTGRES_CONFIG, + POSTGRES_METADATA_DOMAINS, + POSTGRES_SUPPORTED_DOMAINS, + POSTGRES_SYSTEM_METADATA_QUERIES, POSTGRES_VERSION_PATTERN, resolve_postgres_json_type, ) -__all__ = ("POSTGRES_CONFIG", "POSTGRES_VERSION_PATTERN", "resolve_postgres_json_type") +__all__ = ( + "POSTGRES_CONFIG", + "POSTGRES_METADATA_DOMAINS", + "POSTGRES_SUPPORTED_DOMAINS", + "POSTGRES_SYSTEM_METADATA_QUERIES", + "POSTGRES_VERSION_PATTERN", + "resolve_postgres_json_type", +) diff --git a/sqlspec/data_dictionary/dialects/postgres/config.py b/sqlspec/data_dictionary/dialects/postgres/config.py index 405844603..4bb9353b6 100644 --- a/sqlspec/data_dictionary/dialects/postgres/config.py +++ b/sqlspec/data_dictionary/dialects/postgres/config.py @@ -8,11 +8,44 @@ "POSTGRES_CONFIG", "POSTGRES_FEATURE_FLAGS", "POSTGRES_FEATURE_VERSIONS", + "POSTGRES_METADATA_DOMAINS", + "POSTGRES_SUPPORTED_DOMAINS", + "POSTGRES_SYSTEM_METADATA_QUERIES", "POSTGRES_TYPE_MAPPINGS", "POSTGRES_VERSION_PATTERN", "resolve_postgres_json_type", ) +POSTGRES_METADATA_DOMAINS: tuple[str, ...] = ( + "schemas", + "objects", + "tables", + "columns", + "constraints", + "indexes", + "views", + "materialized_views", + "sequences", + "types", + "routines", + "triggers", + "comments", + "privileges", + "dependencies", + "extensions", + "partitions", + "ddl", + "system", +) +POSTGRES_SUPPORTED_DOMAINS: frozenset[str] = frozenset(d for d in POSTGRES_METADATA_DOMAINS if d != "system") +POSTGRES_SYSTEM_METADATA_QUERIES: dict[str, str] = { + "settings": "settings", + "statement_history": "pg_stat_statements", + "pg_stat_statements": "pg_stat_statements", + "table_statistics": "table_stats", + "table_stats": "table_stats", +} + POSTGRES_VERSION_PATTERN = re.compile(r"PostgreSQL (\d+)\.(\d+)(?:\.(\d+))?") diff --git a/sqlspec/data_dictionary/dialects/postgres/sql/constraints.sql b/sqlspec/data_dictionary/dialects/postgres/sql/constraints.sql index 59351a740..6b41bf929 100644 --- a/sqlspec/data_dictionary/dialects/postgres/sql/constraints.sql +++ b/sqlspec/data_dictionary/dialects/postgres/sql/constraints.sql @@ -1,5 +1,23 @@ -- name: by_schema -- dialect: postgres +WITH con_cols AS ( + SELECT + con.oid AS con_oid, + array_agg(att.attname::text ORDER BY keys.ord) AS columns + FROM pg_catalog.pg_constraint con + CROSS JOIN LATERAL unnest(con.conkey) WITH ORDINALITY AS keys(attnum, ord) + JOIN pg_catalog.pg_attribute att ON att.attrelid = con.conrelid AND att.attnum = keys.attnum + GROUP BY con.oid +), +ref_cols AS ( + SELECT + con.oid AS con_oid, + array_agg(att.attname::text ORDER BY keys.ord) AS referenced_columns + FROM pg_catalog.pg_constraint con + CROSS JOIN LATERAL unnest(con.confkey) WITH ORDINALITY AS keys(attnum, ord) + JOIN pg_catalog.pg_attribute att ON att.attrelid = con.confrelid AND att.attnum = keys.attnum + GROUP BY con.oid +) SELECT current_database()::text AS constraint_catalog, n.nspname::text AS constraint_schema, @@ -15,20 +33,10 @@ SELECT ELSE con.contype::text END AS constraint_type, pg_catalog.pg_get_constraintdef(con.oid, true)::text AS definition, - ARRAY( - SELECT att.attname::text - FROM unnest(con.conkey) WITH ORDINALITY AS keys(attnum, ord) - JOIN pg_catalog.pg_attribute att ON att.attrelid = con.conrelid AND att.attnum = keys.attnum - ORDER BY keys.ord - )::text[] AS columns, + COALESCE(cc.columns, ARRAY[]::text[]) AS columns, ref_ns.nspname::text AS referenced_schema, ref.relname::text AS referenced_table, - ARRAY( - SELECT att.attname::text - FROM unnest(con.confkey) WITH ORDINALITY AS keys(attnum, ord) - JOIN pg_catalog.pg_attribute att ON att.attrelid = con.confrelid AND att.attnum = keys.attnum - ORDER BY keys.ord - )::text[] AS referenced_columns, + COALESCE(rc.referenced_columns, ARRAY[]::text[]) AS referenced_columns, con.condeferrable AS is_deferrable, con.condeferred AS initially_deferred, dep.refobjid IS NOT NULL AS is_extension_owned @@ -38,6 +46,8 @@ LEFT JOIN pg_catalog.pg_class c ON c.oid = con.conrelid LEFT JOIN pg_catalog.pg_class ref ON ref.oid = con.confrelid LEFT JOIN pg_catalog.pg_namespace ref_ns ON ref_ns.oid = ref.relnamespace LEFT JOIN pg_catalog.pg_depend dep ON dep.objid = con.oid AND dep.deptype = 'e' +LEFT JOIN con_cols cc ON cc.con_oid = con.oid +LEFT JOIN ref_cols rc ON rc.con_oid = con.oid WHERE n.nspname = :schema_name AND (:table_name::text IS NULL OR c.relname = :table_name) ORDER BY n.nspname, c.relname, con.conname; diff --git a/sqlspec/data_dictionary/dialects/postgres/sql/system.sql b/sqlspec/data_dictionary/dialects/postgres/sql/system.sql index b6b6aae58..0035d9ebb 100644 --- a/sqlspec/data_dictionary/dialects/postgres/sql/system.sql +++ b/sqlspec/data_dictionary/dialects/postgres/sql/system.sql @@ -9,7 +9,7 @@ SELECT total_exec_time, mean_exec_time, rows -FROM pg_catalog.pg_stat_statements +FROM pg_stat_statements ORDER BY total_exec_time DESC LIMIT :limit; diff --git a/tests/unit/adapters/test_asyncpg/test_transaction_state.py b/tests/unit/adapters/test_asyncpg/test_transaction_state.py index 4fcba4318..9c8af75ad 100644 --- a/tests/unit/adapters/test_asyncpg/test_transaction_state.py +++ b/tests/unit/adapters/test_asyncpg/test_transaction_state.py @@ -428,3 +428,23 @@ def test_asyncpg_events_resolve_listener_config_allows_mapping_direct_dsn() -> N hub = AsyncpgListenerHub(config) resolved = hub._resolve_listener_config() assert resolved is not None + + +async def test_asyncpg_get_indexes_with_none_schema_resolves_via_search_path() -> None: + """get_indexes passes schema_name=None when schema is None to resolve table on search_path.""" + from sqlspec.adapters.asyncpg.data_dictionary import AsyncpgDataDictionary + + calls: list[dict[str, Any]] = [] + driver = MagicMock() + + async def fake_select(query: Any, **kwargs: Any) -> list[Any]: + calls.append(kwargs) + return [] + + driver.select = AsyncMock(side_effect=fake_select) + + dd = AsyncpgDataDictionary() + await dd.get_indexes(driver, table="users", schema=None) + assert len(calls) == 1 + assert calls[0]["schema_name"] is None + assert calls[0]["table_name"] == "users" diff --git a/tests/unit/data_dictionary/test_postgres_cockroach_query_packs.py b/tests/unit/data_dictionary/test_postgres_cockroach_query_packs.py index c53aa7311..9625896e9 100644 --- a/tests/unit/data_dictionary/test_postgres_cockroach_query_packs.py +++ b/tests/unit/data_dictionary/test_postgres_cockroach_query_packs.py @@ -150,3 +150,18 @@ def test_postgres_and_cockroach_foreign_keys_and_indexes_contracts() -> None: assert "AS is_primary" in crdb_columns_by_table assert "AS generation_expression" in crdb_columns_by_table assert "COALESCE(" in crdb_columns_by_table + + +def test_postgres_system_and_constraints_query_contracts() -> None: + """PostgreSQL pg_stat_statements omits pg_catalog prefix and constraints pre-aggregate columns.""" + loader = DataDictionaryLoader() + sys_query = loader.get_domain_query_text("postgres", "system", "pg_stat_statements") + assert sys_query is not None + assert "FROM pg_catalog.pg_stat_statements" not in sys_query + assert "FROM pg_stat_statements" in sys_query + + con_query = loader.get_domain_query_text("postgres", "constraints", "by_schema") + assert con_query is not None + assert "WITH con_cols AS (" in con_query + assert "CROSS JOIN LATERAL unnest(con.conkey) WITH ORDINALITY" in con_query + assert "ARRAY(SELECT att.attname" not in con_query From ecf30813d2b33c240d5bea08717e7269c6260462 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 01:05:46 +0000 Subject: [PATCH 075/174] perf(asyncpg): hot-path allocations and cache execution fast path Eliminate closures in register_json_codecs by using slotted callable classes, avoid per-query dict allocations in timeout handling, optimize AsyncpgStreamSource chunk fetching and transaction checking, implement _execute_cache_hit_with_cursor to bypass dispatch_execute on cache hits, avoid closure in provide_session, and remove unused PREPARED_STATEMENT_CACHE_SIZE. --- sqlspec/adapters/asyncpg/config.py | 6 +- sqlspec/adapters/asyncpg/core.py | 97 +++++++++++++++---- sqlspec/adapters/asyncpg/driver.py | 72 +++++++++++--- .../test_asyncpg/test_transaction_state.py | 22 +++++ 4 files changed, 164 insertions(+), 33 deletions(-) diff --git a/sqlspec/adapters/asyncpg/config.py b/sqlspec/adapters/asyncpg/config.py index 394954974..f02a97ce9 100644 --- a/sqlspec/adapters/asyncpg/config.py +++ b/sqlspec/adapters/asyncpg/config.py @@ -450,12 +450,14 @@ def provide_session( return AsyncpgSessionContext( acquire_connection=factory.acquire_connection, release_connection=factory.release_connection, - statement_config=statement_config - or (lambda: resolve_runtime_statement_config(None, self.statement_config, default_statement_config)), + statement_config=statement_config or self._default_runtime_statement_config, driver_features=self.driver_features, prepare_driver=self._prepare_driver, ) + def _default_runtime_statement_config(self) -> "StatementConfig": + return resolve_runtime_statement_config(None, self.statement_config, default_statement_config) + def get_signature_namespace(self) -> "dict[str, Any]": """Get the signature namespace for AsyncPG types. diff --git a/sqlspec/adapters/asyncpg/core.py b/sqlspec/adapters/asyncpg/core.py index f969923c1..314d5ebea 100644 --- a/sqlspec/adapters/asyncpg/core.py +++ b/sqlspec/adapters/asyncpg/core.py @@ -4,8 +4,9 @@ import contextlib import datetime as dt import re +import sys from collections.abc import Sized -from typing import TYPE_CHECKING, Any, Final, NamedTuple +from typing import TYPE_CHECKING, Any, NamedTuple from sqlspec.adapters.asyncpg._typing import asyncpg_module as asyncpg from sqlspec.core import DriverParameterProfile, ParameterStyle, StatementConfig, build_statement_config_from_profile @@ -89,7 +90,6 @@ class NormalizedStackOperation(NamedTuple): parameters: "tuple[Any, ...] | dict[str, Any] | None" -PREPARED_STATEMENT_CACHE_SIZE: Final[int] = 32 _EXCEPTION_MAPPING_DISPATCHER = TypeDispatcher["tuple[str, type[SQLSpecError], str]"]() @@ -263,20 +263,56 @@ def build_statement_config( return base_config.replace(parameter_config=parameter_config) +class _JsonEncoder: + __slots__ = ("_encoder",) + + def __init__(self, encoder: "Callable[[Any], str]") -> None: + self._encoder = encoder + + def __call__(self, value: Any) -> bytes: + return _encode_json_payload(value, self._encoder) + + +class _JsonDecoder: + __slots__ = ("_decoder",) + + def __init__(self, decoder: "Callable[[str], Any]") -> None: + self._decoder = decoder + + def __call__(self, value: Any) -> Any: + return _decode_json_payload(value, self._decoder) + + +class _JsonbEncoder: + __slots__ = ("_encoder",) + + def __init__(self, encoder: "Callable[[Any], str]") -> None: + self._encoder = encoder + + def __call__(self, value: Any) -> bytes: + return _encode_jsonb_payload(value, self._encoder) + + +class _JsonbDecoder: + __slots__ = ("_decoder",) + + def __init__(self, decoder: "Callable[[str], Any]") -> None: + self._decoder = decoder + + def __call__(self, value: Any) -> Any: + return _decode_jsonb_payload(value, self._decoder) + + async def register_json_codecs(connection: Any, encoder: Any, decoder: Any) -> None: """Register JSON type codecs on asyncpg connection.""" try: await connection.set_type_codec( - "json", - encoder=lambda value: _encode_json_payload(value, encoder), - decoder=lambda value: _decode_json_payload(value, decoder), - schema="pg_catalog", - format="binary", + "json", encoder=_JsonEncoder(encoder), decoder=_JsonDecoder(decoder), schema="pg_catalog", format="binary" ) await connection.set_type_codec( "jsonb", - encoder=lambda value: _encode_jsonb_payload(value, encoder), - decoder=lambda value: _decode_jsonb_payload(value, decoder), + encoder=_JsonbEncoder(encoder), + decoder=_JsonbDecoder(decoder), schema="pg_catalog", format="binary", ) @@ -447,7 +483,7 @@ def create_mapped_exception(error: Any, *, logger: Any | None = None) -> SQLSpec Returns: A SQLSpec exception that wraps the original error """ - del logger + _ = logger if isinstance(error, (asyncio.TimeoutError, TimeoutError)): return QueryTimeoutError(str(error)) if isinstance(error, OSError) or ( @@ -497,7 +533,7 @@ def collect_rows(records: "list[Any] | None") -> "tuple[list[Any], list[str]]": class AsyncpgStreamSource: """Compiled async chunk source streaming dict rows from an asyncpg cursor in a stream-owned transaction.""" - __slots__ = ("_chunk_size", "_cursor", "_driver", "_parameters", "_sql", "_timeout_args", "_transaction") + __slots__ = ("_chunk_size", "_cursor", "_driver", "_parameters", "_sql", "_timeout", "_transaction") def __init__( self, @@ -505,9 +541,14 @@ def __init__( sql: str, parameters: "tuple[Any, ...]", chunk_size: int, - timeout_args: "dict[str, Any] | None" = None, + timeout: "float | dict[str, Any] | None" = None, ) -> None: - self._timeout_args = timeout_args or {} + if isinstance(timeout, dict): + self._timeout: float | None = ( + float(timeout["timeout"]) if "timeout" in timeout and timeout["timeout"] is not None else None + ) + else: + self._timeout = float(timeout) if timeout is not None else None self._driver = driver self._sql = sql self._parameters = parameters @@ -517,14 +558,28 @@ def __init__( async def start(self) -> None: handler = self._driver.handle_database_exceptions() - await self._driver._run_with_exception_handler(handler, self._start) + if sys.exc_info()[1] is None: + async with handler: + await self._start() + else: + await self._driver._run_with_exception_handler(handler, self._start) self._driver._check_pending_exception(handler) async def fetch_chunk(self) -> "list[dict[str, Any]]": handler = self._driver.handle_database_exceptions() - records = await self._driver._run_with_exception_handler( - handler, self._cursor.fetch, self._chunk_size, **self._timeout_args - ) + if sys.exc_info()[1] is None: + async with handler: + records = ( + await self._cursor.fetch(self._chunk_size, timeout=self._timeout) + if self._timeout is not None + else await self._cursor.fetch(self._chunk_size) + ) + elif self._timeout is not None: + records = await self._driver._run_with_exception_handler( + handler, self._cursor.fetch, self._chunk_size, timeout=self._timeout + ) + else: + records = await self._driver._run_with_exception_handler(handler, self._cursor.fetch, self._chunk_size) self._driver._check_pending_exception(handler) assert records is not None return [dict(record) for record in records] @@ -544,14 +599,16 @@ async def close(self, error: bool = False) -> None: await transaction.rollback() async def _start(self) -> None: - in_tx_fn = getattr(self._driver, "_connection_in_transaction", None) - already_in_tx = bool(in_tx_fn()) if callable(in_tx_fn) else False + already_in_tx = self._driver._connection_in_transaction() if not already_in_tx: transaction = self._driver.connection.transaction() await transaction.start() self._transaction = transaction try: - self._cursor = await self._driver.connection.cursor(self._sql, *self._parameters, **self._timeout_args) + if self._timeout is not None: + self._cursor = await self._driver.connection.cursor(self._sql, *self._parameters, timeout=self._timeout) + else: + self._cursor = await self._driver.connection.cursor(self._sql, *self._parameters) except BaseException: transaction = self._transaction self._transaction = None diff --git a/sqlspec/adapters/asyncpg/driver.py b/sqlspec/adapters/asyncpg/driver.py index 89d25325d..be470003f 100644 --- a/sqlspec/adapters/asyncpg/driver.py +++ b/sqlspec/adapters/asyncpg/driver.py @@ -25,6 +25,7 @@ from sqlspec.adapters.asyncpg.data_dictionary import AsyncpgDataDictionary from sqlspec.core import ( SQL, + DMLResult, StackResult, StatementStack, get_cache_config, @@ -49,7 +50,7 @@ from sqlspec.adapters.asyncpg._typing import AsyncpgConnection from sqlspec.core import ArrowResult, SQLResult, StatementConfig - from sqlspec.driver import ExecutionResult + from sqlspec.driver import CachedQuery, ExecutionResult from sqlspec.storage import StorageBridgeJob, StorageDestination, StorageFormat, StorageTelemetry @@ -137,8 +138,13 @@ async def dispatch_execute(self, cursor: "AsyncpgConnection", statement: "SQL") sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) params: tuple[Any, ...] = cast("tuple[Any, ...]", prepared_parameters) if prepared_parameters else () + timeout = self._timeout_value(statement.statement_config) if statement.returns_rows(): - records = await cursor.fetch(sql, *params, **self._timeout_arguments(statement.statement_config)) + records = ( + await cursor.fetch(sql, *params, timeout=timeout) + if timeout is not None + else await cursor.fetch(sql, *params) + ) data, column_names = collect_rows(records) return self.create_execution_result( @@ -150,8 +156,11 @@ async def dispatch_execute(self, cursor: "AsyncpgConnection", statement: "SQL") row_format="record", ) - result = await cursor.execute(sql, *params, **self._timeout_arguments(statement.statement_config)) - + result = ( + await cursor.execute(sql, *params, timeout=timeout) + if timeout is not None + else await cursor.execute(sql, *params) + ) affected_rows = parse_status(result) return self.create_execution_result(cursor, rowcount_override=affected_rows) @@ -167,16 +176,53 @@ async def dispatch_execute_many(self, cursor: "AsyncpgConnection", statement: "S ExecutionResult with batch execution details """ sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) + timeout = self._timeout_value(statement.statement_config) if prepared_parameters: parameter_sets = cast("list[Sequence[object]]", prepared_parameters) - await cursor.executemany(sql, parameter_sets, **self._timeout_arguments(statement.statement_config)) + if timeout is not None: + await cursor.executemany(sql, parameter_sets, timeout=timeout) + else: + await cursor.executemany(sql, parameter_sets) affected_rows = resolve_many_rowcount(parameter_sets) else: affected_rows = 0 return self.create_execution_result(cursor, rowcount_override=affected_rows, is_many_result=True) + async def _execute_cache_hit_with_cursor( + self, sql: str, params: "tuple[Any, ...] | list[Any] | dict[str, Any]", cached: "CachedQuery" + ) -> "SQLResult": + wire_sql = self._format_wire_sql(cached.compiled_sql) + timeout = self._timeout_value(self.statement_config) + args = ( + params if isinstance(params, (tuple, list)) else tuple(params.values()) if isinstance(params, dict) else () + ) + if cached.operation_profile.returns_rows: + records = ( + await self.connection.fetch(wire_sql, *args, timeout=timeout) + if timeout is not None + else await self.connection.fetch(wire_sql, *args) + ) + data, column_names = collect_rows(records) + return self.build_cached_select_result( + sql, + params, + cached, + selected_data=data, + column_names=column_names, + data_row_count=len(data), + row_format="record", + compiled_sql=wire_sql, + ) + status = ( + await self.connection.execute(wire_sql, *args, timeout=timeout) + if timeout is not None + else await self.connection.execute(wire_sql, *args) + ) + affected_rows = parse_status(status) + return DMLResult(cached.operation_type, affected_rows) + async def dispatch_execute_script(self, cursor: "AsyncpgConnection", statement: "SQL") -> "ExecutionResult": """Execute SQL script with statement splitting and parameter handling. @@ -296,9 +342,8 @@ def dispatch_select_stream(self, statement: "SQL", chunk_size: int) -> "AsyncRow return None sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) params: tuple[Any, ...] = cast("tuple[Any, ...]", prepared_parameters) if prepared_parameters else () - return AsyncRowStream( - AsyncpgStreamSource(self, sql, params, chunk_size, self._timeout_arguments(statement.statement_config)) - ) + timeout = self._timeout_value(statement.statement_config) + return AsyncRowStream(AsyncpgStreamSource(self, sql, params, chunk_size, timeout)) def handle_database_exceptions(self) -> "AsyncpgExceptionHandler": """Handle database exceptions with PostgreSQL error codes.""" @@ -519,13 +564,18 @@ async def load_from_storage( table, source, file_format=file_format, partitioner=partitioner, overwrite=overwrite, **kwargs ) - def _timeout_arguments(self, config: "StatementConfig") -> "dict[str, Any]": + def _timeout_value(self, config: "StatementConfig") -> float | None: for options in (config.execution_args, self.statement_config.execution_args): if options: for key in ("timeout", "command_timeout"): if key in options: - return {"timeout": options[key]} - return {} + val = options[key] + return float(val) if val is not None else None + return None + + def _timeout_arguments(self, config: "StatementConfig") -> "dict[str, Any]": + timeout = self._timeout_value(config) + return {"timeout": timeout} if timeout is not None else {} def _release_failed_transaction_claim(self, transaction: Any) -> None: """Clear the connection's transaction claim left by a failed start.""" diff --git a/tests/unit/adapters/test_asyncpg/test_transaction_state.py b/tests/unit/adapters/test_asyncpg/test_transaction_state.py index 9c8af75ad..a2eaa4d9b 100644 --- a/tests/unit/adapters/test_asyncpg/test_transaction_state.py +++ b/tests/unit/adapters/test_asyncpg/test_transaction_state.py @@ -448,3 +448,25 @@ async def fake_select(query: Any, **kwargs: Any) -> list[Any]: assert len(calls) == 1 assert calls[0]["schema_name"] is None assert calls[0]["table_name"] == "users" + + +async def test_asyncpg_cache_hit_with_cursor() -> None: + """AsyncpgDriver executes cached queries directly on connection via _execute_cache_hit_with_cursor.""" + conn = MagicMock() + conn.fetch = AsyncMock(return_value=[{"val": 42}]) + conn.execute = AsyncMock(return_value="UPDATE 1") + + driver = AsyncpgDriver(cast("Any", conn)) + res1 = await driver.execute("SELECT 42") + assert res1.data == [{"val": 42}] + + res2 = await driver.execute("SELECT 42") + assert res2.data == [{"val": 42}] + assert conn.fetch.call_count == 2 + + upd1 = await driver.execute("UPDATE t SET x = 1") + assert upd1.rows_affected == 1 + + upd2 = await driver.execute("UPDATE t SET x = 1") + assert upd2.rows_affected == 1 + assert conn.execute.call_count == 2 From 4def74eb1ad4cc795f285c67beee260b2fd12e9b Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 01:05:46 +0000 Subject: [PATCH 076/174] refactor(asyncpg): code hygiene, capability flags, and dead code cleanup Reclassify ConnectionDoesNotExistError and CannotConnectNowError to DatabaseConnectionError, correct capability flags for native arrow/parquet export/import on AsyncpgConfig, add enable_pgcopy_binary to features TypedDict, remove unused invoke_prepared_statement and no-op load_from_storage override, remove redundant config_runtime re-exports, remove inline comments, and support empty records with overwrite=True in load_from_records. --- sqlspec/adapters/asyncpg/config.py | 10 ++-- sqlspec/adapters/asyncpg/core.py | 49 ++----------------- sqlspec/adapters/asyncpg/driver.py | 31 ++++++------ sqlspec/adapters/asyncpg/events/_hub.py | 4 +- sqlspec/adapters/asyncpg/events/backend.py | 1 - .../unit/adapters/test_asyncpg/test_config.py | 12 ++--- .../test_asyncpg/test_load_from_records.py | 9 ++++ .../test_asyncpg/test_type_handlers.py | 33 +++++++------ 8 files changed, 57 insertions(+), 92 deletions(-) diff --git a/sqlspec/adapters/asyncpg/config.py b/sqlspec/adapters/asyncpg/config.py index f02a97ce9..30664ea47 100644 --- a/sqlspec/adapters/asyncpg/config.py +++ b/sqlspec/adapters/asyncpg/config.py @@ -201,6 +201,8 @@ class AsyncpgDriverFeatures(TypedDict): """AlloyDB network path (``PUBLIC``, ``PRIVATE``, or ``PSC``; defaults to ``PRIVATE``).""" enable_alloydb_columnar: NotRequired[bool] """Enable ``google_columnar_engine.enable_columnar_scan=on`` at session startup or per transaction.""" + enable_pgcopy_binary: NotRequired[bool] + """Enable binary COPY format for Arrow table ingestion.""" enable_events: NotRequired[bool] """Enable SQLSpec event channel pub/sub support (defaults to ``True`` when ``events`` is configured).""" events_backend: NotRequired[Literal["notify", "notify_queue", "poll_queue"]] @@ -292,13 +294,13 @@ class AsyncpgConfig(AsyncDatabaseConfig[AsyncpgConnection, "Pool[Record]", Async """Configuration for AsyncPG database connections using TypedDict.""" driver_type: "ClassVar[type[AsyncpgDriver]]" = AsyncpgDriver - connection_type: "ClassVar[type[AsyncpgConnection]]" = AsyncpgConnection # type: ignore[assignment] + connection_type: "ClassVar[type[Any]]" = cast("type[Any]", AsyncpgConnection) supports_transactional_ddl: "ClassVar[bool]" = True supports_migration_schemas: "ClassVar[bool]" = True - supports_native_arrow_export: "ClassVar[bool]" = True + supports_native_arrow_export: "ClassVar[bool]" = False supports_native_arrow_import: "ClassVar[bool]" = True - supports_native_parquet_export: "ClassVar[bool]" = True - supports_native_parquet_import: "ClassVar[bool]" = True + supports_native_parquet_export: "ClassVar[bool]" = False + supports_native_parquet_import: "ClassVar[bool]" = False supports_native_row_streaming: "ClassVar[bool]" = True type_coercion_capabilities: "ClassVar[TypeCoercionCapabilities]" = TypeCoercionCapabilities( datetime_binding="native", timestamp_precision="microsecond", json_columns_decoded=True, uuid_binding="native" diff --git a/sqlspec/adapters/asyncpg/core.py b/sqlspec/adapters/asyncpg/core.py index 314d5ebea..9cda11784 100644 --- a/sqlspec/adapters/asyncpg/core.py +++ b/sqlspec/adapters/asyncpg/core.py @@ -10,15 +10,9 @@ from sqlspec.adapters.asyncpg._typing import asyncpg_module as asyncpg from sqlspec.core import DriverParameterProfile, ParameterStyle, StatementConfig, build_statement_config_from_profile -from sqlspec.core.config_runtime import ( - build_postgres_extension_probe_names, - is_postgres_extension_active, - resolve_postgres_extension_state, - resolve_runtime_statement_config, -) +from sqlspec.core.config_runtime import resolve_runtime_statement_config from sqlspec.exceptions import ( CheckViolationError, - ConnectionTimeoutError, DatabaseConnectionError, DeadlockError, ForeignKeyViolationError, @@ -54,7 +48,6 @@ "NormalizedStackOperation", "apply_driver_features", "build_connection_config", - "build_postgres_extension_probe_names", "build_postgres_storage_params_sql", "build_profile", "build_statement_config", @@ -63,13 +56,10 @@ "create_mapped_exception", "default_statement_config", "driver_profile", - "invoke_prepared_statement", - "is_postgres_extension_active", "parse_status", "register_json_codecs", "register_pgvector_support", "resolve_many_rowcount", - "resolve_postgres_extension_state", "resolve_runtime_statement_config", ) @@ -207,39 +197,6 @@ def configure_parameter_serializers( return parameter_config.replace(json_serializer=serializer, json_deserializer=effective_deserializer) -async def invoke_prepared_statement( - prepared: Any, parameters: "tuple[Any, ...] | dict[str, Any] | list[Any] | None", *, fetch: bool, **kwargs: Any -) -> Any: - """Invoke an AsyncPG prepared statement with optional parameters. - - Args: - prepared: AsyncPG prepared statement object. - parameters: Prepared parameters payload. - fetch: Whether to fetch rows. - **kwargs: Native execution options, including timeout. - - Returns: - Query result or status message. - """ - if parameters is None: - if fetch: - return await prepared.fetch(**kwargs) - await prepared.fetch(**kwargs) - return prepared.get_statusmsg() - - if isinstance(parameters, dict): - values = tuple(parameters.values()) - if fetch: - return await prepared.fetch(*values, **kwargs) - await prepared.fetch(*values, **kwargs) - return prepared.get_statusmsg() - - if fetch: - return await prepared.fetch(*parameters, **kwargs) - await prepared.fetch(*parameters, **kwargs) - return prepared.get_statusmsg() - - def build_statement_config( *, json_serializer: "Callable[[Any], str] | None" = None, json_deserializer: "Callable[[str], Any] | None" = None ) -> "StatementConfig": @@ -454,10 +411,10 @@ def resolve_many_rowcount(parameter_sets: Any, *, fallback_count: "int | None" = asyncpg.exceptions.InvalidAuthorizationSpecificationError, ("28000", PermissionDeniedError, "authorization error") ) _EXCEPTION_MAPPING_DISPATCHER.register( - asyncpg.exceptions.ConnectionDoesNotExistError, ("08003", ConnectionTimeoutError, "connection does not exist") + asyncpg.exceptions.ConnectionDoesNotExistError, ("08003", DatabaseConnectionError, "connection does not exist") ) _EXCEPTION_MAPPING_DISPATCHER.register( - asyncpg.exceptions.CannotConnectNowError, ("57P03", ConnectionTimeoutError, "cannot connect now") + asyncpg.exceptions.CannotConnectNowError, ("57P03", DatabaseConnectionError, "cannot connect now") ) _EXCEPTION_MAPPING_DISPATCHER.register( asyncpg.exceptions.PostgresSyntaxError, ("42601", SQLParsingError, "SQL syntax error") diff --git a/sqlspec/adapters/asyncpg/driver.py b/sqlspec/adapters/asyncpg/driver.py index be470003f..8a83cc239 100644 --- a/sqlspec/adapters/asyncpg/driver.py +++ b/sqlspec/adapters/asyncpg/driver.py @@ -494,6 +494,21 @@ async def load_from_records( self._require_capability("arrow_import_enabled") materialized = list(records) if not materialized: + if overwrite: + table_name, schema_name, quoted_target = self._copy_target(table) + try: + await self.connection.execute(f"TRUNCATE TABLE {quoted_target}") + except Exception as exc: + mapped = create_mapped_exception(exc) + msg = f"Failed to truncate table '{table}': {exc}" + raise type(mapped)(msg) from exc + empty_telemetry: StorageTelemetry = { + "bytes_processed": 0, + "destination": table, + "format": "records", + "rows_processed": 0, + } + return self._storage_job(empty_telemetry) msg = "load_from_records requires at least one record." raise ImproperConfigurationError(msg) @@ -548,22 +563,6 @@ async def load_from_records( } return self._storage_job(telemetry_payload) - async def load_from_storage( - self, - table: str, - source: "StorageDestination", - *, - file_format: "StorageFormat", - partitioner: "dict[str, object] | None" = None, - overwrite: bool = False, - **kwargs: Any, - ) -> "StorageBridgeJob": - """Read an artifact from storage and ingest it via COPY.""" - - return await super().load_from_storage( - table, source, file_format=file_format, partitioner=partitioner, overwrite=overwrite, **kwargs - ) - def _timeout_value(self, config: "StatementConfig") -> float | None: for options in (config.execution_args, self.statement_config.execution_args): if options: diff --git a/sqlspec/adapters/asyncpg/events/_hub.py b/sqlspec/adapters/asyncpg/events/_hub.py index a20373d8d..386f559d2 100644 --- a/sqlspec/adapters/asyncpg/events/_hub.py +++ b/sqlspec/adapters/asyncpg/events/_hub.py @@ -7,8 +7,6 @@ notifications into one queue per consumer task. """ -# pyright: reportPrivateUsage=false - import asyncio import contextlib import logging @@ -162,7 +160,7 @@ def is_subscribed(self, channel: str) -> bool: def _get_consumer_queue(self, channel: str) -> "asyncio.Queue[str] | None": task = asyncio.current_task() - if task is None: # pragma: no cover + if task is None: msg = "AsyncpgListenerHub.dequeue requires an active asyncio task" raise RuntimeError(msg) queues = self._queues.get(channel) diff --git a/sqlspec/adapters/asyncpg/events/backend.py b/sqlspec/adapters/asyncpg/events/backend.py index 1e71b037e..9b3b04661 100644 --- a/sqlspec/adapters/asyncpg/events/backend.py +++ b/sqlspec/adapters/asyncpg/events/backend.py @@ -1,4 +1,3 @@ -# pyright: reportPrivateUsage=false """Native and hybrid PostgreSQL backends for EventChannel.""" import logging diff --git a/tests/unit/adapters/test_asyncpg/test_config.py b/tests/unit/adapters/test_asyncpg/test_config.py index b1d0cc43a..f54ed3d8d 100644 --- a/tests/unit/adapters/test_asyncpg/test_config.py +++ b/tests/unit/adapters/test_asyncpg/test_config.py @@ -14,14 +14,10 @@ AsyncpgDriverFeatures, AsyncpgPoolConfig, ) -from sqlspec.adapters.asyncpg.core import ( - build_connection_config, - build_postgres_extension_probe_names, - build_statement_config, - resolve_postgres_extension_state, -) +from sqlspec.adapters.asyncpg.core import build_connection_config, build_statement_config from sqlspec.adapters.asyncpg.driver import AsyncpgDriver from sqlspec.core import StatementConfig, StatementStack +from sqlspec.core.config_runtime import build_postgres_extension_probe_names, resolve_postgres_extension_state from sqlspec.exceptions import ImproperConfigurationError @@ -322,6 +318,10 @@ def test_asyncpg_config_normalizes_aliases() -> None: def test_asyncpg_driver_features_contract_retains_configured_keys() -> None: """AsyncpgConfig retains on_connection_create and type_codecs and populates event defaults.""" assert "connection_instance" not in AsyncpgDriverFeatures.__annotations__ + assert "enable_pgcopy_binary" in AsyncpgDriverFeatures.__annotations__ + assert AsyncpgConfig.supports_native_arrow_export is False + assert AsyncpgConfig.supports_native_parquet_export is False + assert AsyncpgConfig.supports_native_parquet_import is False async def hook(_: object) -> None: return None diff --git a/tests/unit/adapters/test_asyncpg/test_load_from_records.py b/tests/unit/adapters/test_asyncpg/test_load_from_records.py index d5e53fcbd..6243b537b 100644 --- a/tests/unit/adapters/test_asyncpg/test_load_from_records.py +++ b/tests/unit/adapters/test_asyncpg/test_load_from_records.py @@ -166,3 +166,12 @@ def fail_if_called(*args: object) -> None: await driver.load_from_records("events", [(1,)], columns=["id"]) connection.copy_records_to_table.assert_awaited_once() + + +async def test_load_from_records_empty_records_with_overwrite(connection: AsyncMock) -> None: + """Empty records with overwrite=True should truncate the table and return 0-row job.""" + driver = AsyncpgDriver(connection, driver_features=DRIVER_FEATURES) # type: ignore[arg-type] + job = await driver.load_from_records("events", [], overwrite=True) + assert job.telemetry["rows_processed"] == 0 + connection.execute.assert_awaited_once() + assert "TRUNCATE TABLE" in connection.execute.await_args[0][0] diff --git a/tests/unit/adapters/test_asyncpg/test_type_handlers.py b/tests/unit/adapters/test_asyncpg/test_type_handlers.py index 146ec7074..a40e89f05 100644 --- a/tests/unit/adapters/test_asyncpg/test_type_handlers.py +++ b/tests/unit/adapters/test_asyncpg/test_type_handlers.py @@ -6,8 +6,14 @@ import pytest from sqlspec.adapters.asyncpg.config import register_json_codecs, register_pgvector_support -from sqlspec.adapters.asyncpg.core import create_mapped_exception, invoke_prepared_statement -from sqlspec.exceptions import OperationCancelledError, PermissionDeniedError, QueryTimeoutError, UniqueViolationError +from sqlspec.adapters.asyncpg.core import create_mapped_exception +from sqlspec.exceptions import ( + DatabaseConnectionError, + OperationCancelledError, + PermissionDeniedError, + QueryTimeoutError, + UniqueViolationError, +) from sqlspec.utils.serializers import from_json, to_json @@ -59,20 +65,15 @@ async def test_register_json_codecs_success() -> None: @pytest.mark.anyio -async def test_invoke_prepared_statement_unpacks_dict_parameters_positionally() -> None: - """Dict parameters passed to invoke_prepared_statement should unpack values positionally.""" - prepared = MagicMock() - prepared.fetch = AsyncMock(return_value=[{"id": 1}]) - prepared.get_statusmsg = MagicMock(return_value="UPDATE 2") - - rows = await invoke_prepared_statement(prepared, {"p1": "alpha", "p2": 42}, fetch=True, timeout=1.5) - assert rows == [{"id": 1}] - prepared.fetch.assert_awaited_once_with("alpha", 42, timeout=1.5) - - prepared.fetch.reset_mock() - status = await invoke_prepared_statement(prepared, {"p1": "beta"}, fetch=False, timeout=2.0) - assert status == "UPDATE 2" - prepared.fetch.assert_awaited_once_with("beta", timeout=2.0) +async def test_asyncpg_connection_error_mapping() -> None: + """ConnectionDoesNotExistError and CannotConnectNowError map to DatabaseConnectionError.""" + exc1 = asyncpg.exceptions.ConnectionDoesNotExistError("connection lost") + mapped1 = create_mapped_exception(exc1) + assert isinstance(mapped1, DatabaseConnectionError) + + exc2 = asyncpg.exceptions.CannotConnectNowError("server starting up") + mapped2 = create_mapped_exception(exc2) + assert isinstance(mapped2, DatabaseConnectionError) @pytest.mark.anyio From 305116d94fe4fa47880d6826218a8431dfbba1b9 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 01:05:46 +0000 Subject: [PATCH 077/174] fix(cockroach_asyncpg): fallback batched insert for storage and record ingestion Replace inherited binary COPY ingestion with non-binary batched executemany INSERT fallback on CockroachAsyncpgDriver for load_from_records and load_from_arrow. --- sqlspec/adapters/cockroach_asyncpg/driver.py | 164 +++++++++++++++++- .../test_cockroach_asyncpg/test_config.py | 18 ++ 2 files changed, 181 insertions(+), 1 deletion(-) diff --git a/sqlspec/adapters/cockroach_asyncpg/driver.py b/sqlspec/adapters/cockroach_asyncpg/driver.py index e33ec2969..29d175451 100644 --- a/sqlspec/adapters/cockroach_asyncpg/driver.py +++ b/sqlspec/adapters/cockroach_asyncpg/driver.py @@ -1,8 +1,9 @@ """CockroachDB AsyncPG driver implementation.""" import asyncio +from collections.abc import Mapping, Sequence from contextlib import suppress -from typing import TYPE_CHECKING, Any, TypeVar, cast +from typing import TYPE_CHECKING, Any, Literal, TypeVar, cast from sqlspec.adapters.asyncpg.core import create_mapped_exception, driver_profile from sqlspec.adapters.asyncpg.driver import AsyncpgDriver, AsyncpgExceptionHandler @@ -21,7 +22,9 @@ from sqlspec.adapters.cockroach_asyncpg.data_dictionary import CockroachAsyncpgDataDictionary from sqlspec.core import SQL, register_driver_profile from sqlspec.driver import AsyncRowStream +from sqlspec.exceptions import ImproperConfigurationError from sqlspec.utils.logging import get_logger +from sqlspec.utils.text import quote_identifier from sqlspec.utils.type_guards import has_sqlstate if TYPE_CHECKING: @@ -208,6 +211,165 @@ async def load_from_storage( table, source, file_format=file_format, partitioner=partitioner, overwrite=overwrite, **kwargs ) + async def load_from_records( + self, + table: str, + records: "Sequence[Mapping[str, Any]] | Sequence[Sequence[Any]]", + *, + columns: "list[str] | None" = None, + overwrite: bool = False, + stage_table: "bool | str | None" = None, + conflict_keys: "Sequence[str] | None" = None, + on_conflict: Literal["update", "nothing"] = "update", + exclude_update_columns: "Sequence[str] | None" = None, + drop_stage_table: bool = True, + drop_indexes: bool = False, + concurrent_indexes: bool = False, + type_casts: "Mapping[str, str] | None" = None, + ) -> "StorageBridgeJob": + """Load mapping or positional records with batched INSERT.""" + if stage_table or conflict_keys or drop_indexes: + return await super().load_from_records( + table, + records, + columns=columns, + overwrite=overwrite, + stage_table=stage_table, + conflict_keys=conflict_keys, + on_conflict=on_conflict, + exclude_update_columns=exclude_update_columns, + drop_stage_table=drop_stage_table, + drop_indexes=drop_indexes, + concurrent_indexes=concurrent_indexes, + type_casts=type_casts, + ) + self._require_capability("arrow_import_enabled") + materialized = list(records) + _, _, quoted_target = self._copy_target(table) + if overwrite: + try: + await self.connection.execute(f"TRUNCATE TABLE {quoted_target}") + except Exception as exc: + mapped = create_mapped_exception(exc) + msg = f"Failed to truncate table '{table}': {exc}" + raise type(mapped)(msg) from exc + if not materialized: + empty_telemetry: StorageTelemetry = { + "bytes_processed": 0, + "destination": table, + "format": "records", + "rows_processed": 0, + } + return self._storage_job(empty_telemetry) + + first = materialized[0] + if isinstance(first, Mapping): + resolved_columns = columns if columns is not None else list(first.keys()) + expected_keys = set(resolved_columns) + copy_rows: list[tuple[Any, ...]] = [] + for record in materialized: + if not isinstance(record, Mapping): + msg = "load_from_records mapping records must all be mappings." + raise ImproperConfigurationError(msg) + if set(record.keys()) != expected_keys: + msg = "load_from_records mapping records must all share the same keys." + raise ImproperConfigurationError(msg) + copy_rows.append(tuple(record[column] for column in resolved_columns)) + else: + if columns is None: + msg = "load_from_records requires columns when records are positional sequences." + raise ImproperConfigurationError(msg) + resolved_columns = columns + copy_rows = [] + for record in materialized: + if isinstance(record, Mapping): + msg = "load_from_records positional records must all have the same shape." + raise ImproperConfigurationError(msg) + row = tuple(record) + if len(row) != len(resolved_columns): + msg = "load_from_records positional records must match the number of columns." + raise ImproperConfigurationError(msg) + copy_rows.append(row) + + cols_str = ", ".join(quote_identifier(c) for c in resolved_columns) + placeholders = ", ".join(f"${i}" for i in range(1, len(resolved_columns) + 1)) + insert_sql = f"INSERT INTO {quoted_target} ({cols_str}) VALUES ({placeholders})" + try: + await self.connection.executemany(insert_sql, copy_rows) + except Exception as exc: + raise create_mapped_exception(exc) from exc + + telemetry_payload: StorageTelemetry = { + "bytes_processed": 0, + "destination": table, + "format": "records", + "rows_processed": len(copy_rows), + } + return self._storage_job(telemetry_payload) + + async def load_from_arrow( + self, + table: str, + source: Any, + *, + partitioner: "dict[str, object] | None" = None, + overwrite: bool = False, + telemetry: "StorageTelemetry | None" = None, + stage_table: "bool | str | None" = None, + conflict_keys: "Sequence[str] | None" = None, + on_conflict: Literal["update", "nothing"] = "update", + exclude_update_columns: "Sequence[str] | None" = None, + drop_stage_table: bool = True, + drop_indexes: bool = False, + concurrent_indexes: bool = False, + type_casts: "Mapping[str, str] | None" = None, + ) -> "StorageBridgeJob": + """Load Arrow table/batches into CockroachDB table via batched INSERT fallback.""" + if stage_table or conflict_keys or drop_indexes: + return await super().load_from_arrow( + table, + source, + partitioner=partitioner, + overwrite=overwrite, + telemetry=telemetry, + stage_table=stage_table, + conflict_keys=conflict_keys, + on_conflict=on_conflict, + exclude_update_columns=exclude_update_columns, + drop_stage_table=drop_stage_table, + drop_indexes=drop_indexes, + concurrent_indexes=concurrent_indexes, + type_casts=type_casts, + ) + _, _, quoted_target = self._copy_target(table) + if overwrite: + try: + await self.connection.execute(f"TRUNCATE TABLE {quoted_target}") + except Exception as exc: + mapped = create_mapped_exception(exc) + msg = f"Failed to truncate table '{table}': {exc}" + raise type(mapped)(msg) from exc + resolved_columns, batch_iter = self._coerce_arrow_batches(source) + cols_str = ", ".join(quote_identifier(c) for c in resolved_columns) + placeholders = ", ".join(f"${i}" for i in range(1, len(resolved_columns) + 1)) + insert_sql = f"INSERT INTO {quoted_target} ({cols_str}) VALUES ({placeholders})" + metrics: dict[str, int] = {"rows_processed": 0, "bytes_processed": 0} + for batch in batch_iter: + if batch.num_rows == 0: + continue + metrics["rows_processed"] += int(batch.num_rows) + metrics["bytes_processed"] += int(batch.nbytes) + rows = self._arrow_batch_to_rows(batch, resolved_columns) + if rows: + try: + await self.connection.executemany(insert_sql, rows) + except Exception as exc: + raise create_mapped_exception(exc) from exc + telemetry_payload = self._ingest_telemetry_from_batches(metrics["rows_processed"], metrics["bytes_processed"]) + telemetry_payload["destination"] = table + self._attach_partition_telemetry(telemetry_payload, partitioner) + return self._storage_job(telemetry_payload, telemetry) + def _wants_follower_read(self, statement: "SQL") -> bool: """Return True when a statement may run under a follower-read transaction.""" return statement.returns_rows() and not statement.is_modifying_operation() diff --git a/tests/unit/adapters/test_cockroach_asyncpg/test_config.py b/tests/unit/adapters/test_cockroach_asyncpg/test_config.py index 400f76e83..dc53e3b14 100644 --- a/tests/unit/adapters/test_cockroach_asyncpg/test_config.py +++ b/tests/unit/adapters/test_cockroach_asyncpg/test_config.py @@ -620,3 +620,21 @@ async def fake_execute(query: str, *args: Any, **kwargs: Any) -> str: assert rows == [{"id": 1}] assert any("AS OF SYSTEM TIME '-10s'" in s for s in executed_statements) tx.commit.assert_awaited() + + +async def test_cockroach_asyncpg_load_from_records_executemany() -> None: + """CockroachAsyncpgDriver.load_from_records uses executemany instead of binary COPY.""" + conn = MagicMock() + conn.executemany = AsyncMock() + conn.copy_records_to_table = AsyncMock() + conn.copy_to_table = AsyncMock() + + driver = CockroachAsyncpgDriver( + connection=conn, driver_features={"storage_capabilities": {"arrow_import_enabled": True}} + ) + job = await driver.load_from_records("events", [{"id": 1, "name": "test"}], columns=["id", "name"]) + assert job.telemetry["rows_processed"] == 1 + conn.executemany.assert_awaited_once() + assert 'INSERT INTO "events"' in conn.executemany.await_args[0][0] + conn.copy_records_to_table.assert_not_called() + conn.copy_to_table.assert_not_called() From 862a63aee9ec22067dcfa515a636373dc2e8b368 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 01:05:46 +0000 Subject: [PATCH 078/174] fix(cockroach_asyncpg): batch ADK memory inserts, vector search, and shard validation Batch insert_memory_entries via UNNEST query, implement vector cosine distance search using embedding in search_entries, and align hash_shard_bucket_count validation with litestar store. --- .../adapters/cockroach_asyncpg/adk/store.py | 140 ++++++++++-------- .../test_cockroach_asyncpg/test_adk_store.py | 39 +++++ 2 files changed, 121 insertions(+), 58 deletions(-) diff --git a/sqlspec/adapters/cockroach_asyncpg/adk/store.py b/sqlspec/adapters/cockroach_asyncpg/adk/store.py index f8e3a0876..13ed09b53 100644 --- a/sqlspec/adapters/cockroach_asyncpg/adk/store.py +++ b/sqlspec/adapters/cockroach_asyncpg/adk/store.py @@ -609,66 +609,79 @@ async def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: " if not entries: return 0 + ids = [entry["id"] for entry in entries] + session_ids = [entry["session_id"] for entry in entries] + app_names = [entry["app_name"] for entry in entries] + user_ids = [entry["user_id"] for entry in entries] + scopes = [entry.get("scope", "user") for entry in entries] + event_ids = [entry["event_id"] for entry in entries] + authors = [entry["author"] for entry in entries] + timestamps = [entry["timestamp"] for entry in entries] + content_jsons = [entry["content_json"] for entry in entries] + content_texts = [entry["content_text"] for entry in entries] + metadata_jsons = [entry["metadata_json"] for entry in entries] + inserted_ats = [entry["inserted_at"] for entry in entries] + if self._owner_id_column_name: sql = f""" INSERT INTO {self._memory_table} (id, session_id, app_name, user_id, scope, event_id, author, {self._owner_id_column_name}, timestamp, content_json, content_text, metadata_json, inserted_at) - VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13) + SELECT * FROM unnest( + $1::varchar[], $2::varchar[], $3::varchar[], $4::varchar[], + $5::varchar[], $6::varchar[], $7::varchar[], $8::varchar[], + $9::timestamptz[], $10::jsonb[], $11::text[], $12::jsonb[], $13::timestamptz[] + ) ON CONFLICT (event_id) DO NOTHING """ + owner_ids = [owner_id for _ in entries] + params: tuple[Any, ...] = ( + ids, + session_ids, + app_names, + user_ids, + scopes, + event_ids, + authors, + owner_ids, + timestamps, + content_jsons, + content_texts, + metadata_jsons, + inserted_ats, + ) else: sql = f""" INSERT INTO {self._memory_table} (id, session_id, app_name, user_id, scope, event_id, author, timestamp, content_json, content_text, metadata_json, inserted_at) - VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12) + SELECT * FROM unnest( + $1::varchar[], $2::varchar[], $3::varchar[], $4::varchar[], + $5::varchar[], $6::varchar[], $7::varchar[], + $8::timestamptz[], $9::jsonb[], $10::text[], $11::jsonb[], $12::timestamptz[] + ) ON CONFLICT (event_id) DO NOTHING """ + params = ( + ids, + session_ids, + app_names, + user_ids, + scopes, + event_ids, + authors, + timestamps, + content_jsons, + content_texts, + metadata_jsons, + inserted_ats, + ) - inserted_count = 0 async with self._config.provide_connection() as conn, conn.transaction(): - for entry in entries: - if self._owner_id_column_name: - result = await conn.execute( - sql, - entry["id"], - entry["session_id"], - entry["app_name"], - entry["user_id"], - entry.get("scope", "user"), - entry["event_id"], - entry["author"], - owner_id, - entry["timestamp"], - entry["content_json"], - entry["content_text"], - entry["metadata_json"], - entry["inserted_at"], - ) - else: - result = await conn.execute( - sql, - entry["id"], - entry["session_id"], - entry["app_name"], - entry["user_id"], - entry.get("scope", "user"), - entry["event_id"], - entry["author"], - entry["timestamp"], - entry["content_json"], - entry["content_text"], - entry["metadata_json"], - entry["inserted_at"], - ) - - parts = result.split() if result else [] - if parts and parts[-1].isdigit() and int(parts[-1]) > 0: - inserted_count += 1 - - return inserted_count + result = await conn.execute(sql, *params) + parts = result.split() if result else [] + return int(parts[-1]) if parts and parts[-1].isdigit() else 0 async def search_entries( self, @@ -683,7 +696,7 @@ async def search_entries( msg = "Memory store is disabled" raise RuntimeError(msg) - if not query: + if not query and embedding is None: return [] effective_limit = limit if limit is not None else self._max_results @@ -691,20 +704,28 @@ async def search_entries( if scope_filter == "all": where_scope = "app_name = $1 AND ((scope = 'user' AND user_id = $2) OR scope = 'app')" scope_params: tuple[Any, ...] = (app_name, user_id) - p_lim = "$4" - p_pat = "$3" elif scope_filter == "user": where_scope = "app_name = $1 AND scope = 'user' AND user_id = $2" scope_params = (app_name, user_id) - p_lim = "$4" - p_pat = "$3" else: where_scope = "app_name = $1 AND scope = 'app'" scope_params = (app_name,) - p_lim = "$3" - p_pat = "$2" - if self._use_fts: + scope_count = len(scope_params) + + if embedding is not None: + p_vec = f"${scope_count + 1}" + p_lim = f"${scope_count + 2}" + sql = f""" + SELECT * FROM {self._memory_table} + WHERE {where_scope} AND embedding IS NOT NULL + ORDER BY embedding <=> {p_vec}::float8[]::vector + LIMIT {p_lim} + """ + params = (*scope_params, list(embedding), effective_limit) + elif self._use_fts: + p_pat = f"${scope_count + 1}" + p_lim = f"${scope_count + 2}" sql = f""" SELECT * FROM {self._memory_table} WHERE {where_scope} @@ -712,17 +733,17 @@ async def search_entries( ORDER BY timestamp DESC LIMIT {p_lim} """ - search_param = query + params = (*scope_params, query, effective_limit) else: + p_pat = f"${scope_count + 1}" + p_lim = f"${scope_count + 2}" sql = f""" SELECT * FROM {self._memory_table} WHERE {where_scope} AND content_text ILIKE {p_pat} ORDER BY timestamp DESC LIMIT {p_lim} """ - search_param = f"%{query}%" - - params = (*scope_params, search_param, effective_limit) + params = (*scope_params, f"%{query}%", effective_limit) try: async with self._config.provide_connection() as conn: @@ -857,9 +878,12 @@ def _cockroach_hash_shard_clause(adk_config: CockroachAsyncpgADKConfig) -> str: if not adk_config.get("enable_hash_sharded_indexes", False): return "" bucket_count = adk_config.get("hash_shard_bucket_count") - if isinstance(bucket_count, int) and not isinstance(bucket_count, bool) and bucket_count > 0: - return f" USING HASH WITH (bucket_count = {bucket_count})" - return " USING HASH" + if bucket_count is None: + return " USING HASH" + if not isinstance(bucket_count, int) or isinstance(bucket_count, bool) or bucket_count <= 0: + msg = "extension_config['adk']['hash_shard_bucket_count'] must be a positive integer" + raise ValueError(msg) + return f" USING HASH WITH (bucket_count = {bucket_count})" def _cockroach_storing_clause(adk_config: CockroachAsyncpgADKConfig, columns: tuple[str, ...]) -> str: diff --git a/tests/unit/adapters/test_cockroach_asyncpg/test_adk_store.py b/tests/unit/adapters/test_cockroach_asyncpg/test_adk_store.py index 55f91e177..ed36f9d44 100644 --- a/tests/unit/adapters/test_cockroach_asyncpg/test_adk_store.py +++ b/tests/unit/adapters/test_cockroach_asyncpg/test_adk_store.py @@ -6,6 +6,7 @@ import pytest +from sqlspec.adapters.cockroach_asyncpg import CockroachAsyncpgConfig from sqlspec.adapters.cockroach_asyncpg.adk import ( CockroachAsyncpgADKConfig, CockroachAsyncpgADKMemoryStore, @@ -197,3 +198,41 @@ async def test_cockroach_asyncpg_list_sessions_rejects_invalid_options_before_co await store.list_sessions("app", **options) assert conn.calls == [] + + +def test_cockroach_asyncpg_adk_hash_shard_bucket_count_validation() -> None: + """Invalid hash_shard_bucket_count must raise ValueError matching litestar store.""" + from sqlspec.adapters.cockroach_asyncpg.adk.store import _cockroach_hash_shard_clause + + bad_config = CockroachAsyncpgADKConfig(enable_hash_sharded_indexes=True, hash_shard_bucket_count=-1) + with pytest.raises(ValueError, match="hash_shard_bucket_count"): + _cockroach_hash_shard_clause(bad_config) + + +async def test_cockroach_asyncpg_adk_memory_search_with_embedding() -> None: + """search_entries with embedding executes vector cosine search.""" + executed_sql: list[tuple[str, tuple[Any, ...]]] = [] + + class _FakeConn: + async def fetch(self, sql: str, *params: Any) -> list[Any]: + executed_sql.append((sql, params)) + return [] + + class _FakeCtx: + async def __aenter__(self) -> _FakeConn: + return _FakeConn() + + async def __aexit__(self, *_args: Any) -> None: + return None + + config = CockroachAsyncpgConfig(extension_config={"adk": {"enable_memory": True}}) + object.__setattr__(config, "provide_connection", MagicMock(return_value=_FakeCtx())) + store = CockroachAsyncpgADKMemoryStore(config) + + await store.search_entries(query="hello", app_name="app", user_id="u1", embedding=[0.1, 0.2, 0.3]) + assert len(executed_sql) == 1 + sql, params = executed_sql[0] + assert "embedding <=>" in sql + assert params[0] == "app" + assert params[1] == "u1" + assert params[2] == [0.1, 0.2, 0.3] From 98e301fcc0ef786338fb652632f0c0103ad9882f Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 01:05:46 +0000 Subject: [PATCH 079/174] fix(cockroach_asyncpg): data dictionary queries and metadata consolidation Rewrite constraints.sql using pre-aggregated CTEs, factor scalar subquery out of columns.sql, route table-specific columns and indexes queries to by_table in CockroachAsyncpgDataDictionary, and consolidate CockroachDB metadata domain constants into dialect config. --- .../cockroach_asyncpg/data_dictionary.py | 38 ++++-------- .../dialects/cockroachdb/__init__.py | 10 +++- .../dialects/cockroachdb/config.py | 22 +++++++ .../dialects/cockroachdb/sql/columns.sql | 30 +++++----- .../dialects/cockroachdb/sql/constraints.sql | 60 +++++++++++++------ .../test_cockroach_asyncpg/test_config.py | 26 ++++++++ .../test_postgres_cockroach_query_packs.py | 15 +++++ 7 files changed, 142 insertions(+), 59 deletions(-) diff --git a/sqlspec/adapters/cockroach_asyncpg/data_dictionary.py b/sqlspec/adapters/cockroach_asyncpg/data_dictionary.py index 157546b46..e86fe76e2 100644 --- a/sqlspec/adapters/cockroach_asyncpg/data_dictionary.py +++ b/sqlspec/adapters/cockroach_asyncpg/data_dictionary.py @@ -28,7 +28,11 @@ system_metadata_gated_result, unsupported_system_metadata_capability, ) -from sqlspec.data_dictionary.dialects.cockroachdb import resolve_cockroachdb_json_type +from sqlspec.data_dictionary.dialects.cockroachdb import ( + COCKROACH_METADATA_DOMAINS, + COCKROACH_SUPPORTED_DOMAINS, + resolve_cockroachdb_json_type, +) from sqlspec.driver import AsyncDataDictionaryBase if TYPE_CHECKING: @@ -38,25 +42,6 @@ __all__ = ("CockroachAsyncpgDataDictionary",) -_COCKROACH_METADATA_DOMAINS: tuple[str, ...] = ( - "schemas", - "objects", - "tables", - "columns", - "constraints", - "indexes", - "views", - "sequences", - "comments", - "privileges", - "dependencies", - "ddl", - "crdb_internal", - "system", -) -_COCKROACH_SUPPORTED_DOMAINS = frozenset(d for d in _COCKROACH_METADATA_DOMAINS if d not in {"crdb_internal", "system"}) - - @mypyc_attr(allow_interpreted_subclasses=True, native_class=False) class CockroachAsyncpgDataDictionary(AsyncDataDictionaryBase): """CockroachDB async data dictionary (AsyncPG).""" @@ -203,7 +188,7 @@ async def get_columns( table_schema_name = None if schema is None else schema_name self._log_table_describe(driver, schema_name=table_schema_name, table_name=table_name, operation="columns") return await driver.select( - _cockroach_domain_sql("columns", "by_schema"), + _cockroach_domain_sql("columns", "by_table"), schema_name=table_schema_name, table_name=table_name, schema_type=ColumnMetadata, @@ -223,10 +208,11 @@ async def get_indexes( ) table_name = self.resolve_identifier(table) - self._log_table_describe(driver, schema_name=schema_name, table_name=table_name, operation="indexes") + table_schema_name = None if schema is None else schema_name + self._log_table_describe(driver, schema_name=table_schema_name, table_name=table_name, operation="indexes") return await driver.select( - _cockroach_domain_sql("indexes", "by_schema"), - schema_name=schema_name, + _cockroach_domain_sql("indexes", "by_table"), + schema_name=table_schema_name, table_name=table_name, schema_type=IndexMetadata, ) @@ -314,7 +300,7 @@ def _cockroach_metadata_capability(domain: str) -> MetadataCapability: risks=(MetadataRisk.EXPENSIVE, MetadataRisk.PRIVILEGED), warnings=("System metadata is opt-in and disabled by default.",), ) - if domain in _COCKROACH_SUPPORTED_DOMAINS: + if domain in COCKROACH_SUPPORTED_DOMAINS: return MetadataCapability( domain=domain, support=MetadataSupport.SUPPORTED, @@ -325,7 +311,7 @@ def _cockroach_metadata_capability(domain: str) -> MetadataCapability: def _cockroach_metadata_profile(adapter: str, domains: Sequence[str] | None) -> MetadataCapabilityProfile: - requested_domains = _COCKROACH_METADATA_DOMAINS if domains is None else tuple(domains) + requested_domains = COCKROACH_METADATA_DOMAINS if domains is None else tuple(domains) return MetadataCapabilityProfile( "cockroachdb", adapter=adapter, diff --git a/sqlspec/data_dictionary/dialects/cockroachdb/__init__.py b/sqlspec/data_dictionary/dialects/cockroachdb/__init__.py index 6dcc444c1..db2bae08f 100644 --- a/sqlspec/data_dictionary/dialects/cockroachdb/__init__.py +++ b/sqlspec/data_dictionary/dialects/cockroachdb/__init__.py @@ -1,9 +1,17 @@ """Dialect configuration for cockroachdb.""" from sqlspec.data_dictionary.dialects.cockroachdb.config import ( + COCKROACH_METADATA_DOMAINS, + COCKROACH_SUPPORTED_DOMAINS, COCKROACHDB_CONFIG, COCKROACHDB_VERSION_PATTERN, resolve_cockroachdb_json_type, ) -__all__ = ("COCKROACHDB_CONFIG", "COCKROACHDB_VERSION_PATTERN", "resolve_cockroachdb_json_type") +__all__ = ( + "COCKROACHDB_CONFIG", + "COCKROACHDB_VERSION_PATTERN", + "COCKROACH_METADATA_DOMAINS", + "COCKROACH_SUPPORTED_DOMAINS", + "resolve_cockroachdb_json_type", +) diff --git a/sqlspec/data_dictionary/dialects/cockroachdb/config.py b/sqlspec/data_dictionary/dialects/cockroachdb/config.py index 5e278943b..9bc5484a3 100644 --- a/sqlspec/data_dictionary/dialects/cockroachdb/config.py +++ b/sqlspec/data_dictionary/dialects/cockroachdb/config.py @@ -10,9 +10,31 @@ "COCKROACHDB_FEATURE_VERSIONS", "COCKROACHDB_TYPE_MAPPINGS", "COCKROACHDB_VERSION_PATTERN", + "COCKROACH_METADATA_DOMAINS", + "COCKROACH_SUPPORTED_DOMAINS", "resolve_cockroachdb_json_type", ) +COCKROACH_METADATA_DOMAINS: tuple[str, ...] = ( + "schemas", + "objects", + "tables", + "columns", + "constraints", + "indexes", + "views", + "sequences", + "comments", + "privileges", + "dependencies", + "ddl", + "crdb_internal", + "system", +) +COCKROACH_SUPPORTED_DOMAINS: frozenset[str] = frozenset( + d for d in COCKROACH_METADATA_DOMAINS if d not in {"crdb_internal", "system"} +) + COCKROACHDB_VERSION_PATTERN = re.compile(r"CockroachDB (?:CCL )?v(\d+)\.(\d+)\.(\d+)") diff --git a/sqlspec/data_dictionary/dialects/cockroachdb/sql/columns.sql b/sqlspec/data_dictionary/dialects/cockroachdb/sql/columns.sql index 258e5fa90..1e1bcc140 100644 --- a/sqlspec/data_dictionary/dialects/cockroachdb/sql/columns.sql +++ b/sqlspec/data_dictionary/dialects/cockroachdb/sql/columns.sql @@ -1,14 +1,15 @@ -- name: by_schema -- dialect: cockroachdb -WITH target_schema AS ( +WITH resolved_table_schema AS ( + SELECT rn.nspname::text AS schema_name + FROM pg_catalog.pg_class rc + JOIN pg_catalog.pg_namespace rn ON rn.oid = rc.relnamespace + WHERE rc.oid = pg_catalog.to_regclass(pg_catalog.quote_ident(:table_name::text)) +), +target_schema AS ( SELECT COALESCE( :schema_name::text, - ( - SELECT rn.nspname::text - FROM pg_catalog.pg_class rc - JOIN pg_catalog.pg_namespace rn ON rn.oid = rc.relnamespace - WHERE rc.oid = pg_catalog.to_regclass(pg_catalog.quote_ident(:table_name::text)) - ), + (SELECT schema_name FROM resolved_table_schema), current_schema() ) AS schema_name ), @@ -81,15 +82,16 @@ ORDER BY c.table_schema, c.table_name, c.ordinal_position; -- name: by_table -- dialect: cockroachdb -WITH target_schema AS ( +WITH resolved_table_schema AS ( + SELECT rn.nspname::text AS schema_name + FROM pg_catalog.pg_class rc + JOIN pg_catalog.pg_namespace rn ON rn.oid = rc.relnamespace + WHERE rc.oid = pg_catalog.to_regclass(pg_catalog.quote_ident(:table_name::text)) +), +target_schema AS ( SELECT COALESCE( :schema_name::text, - ( - SELECT rn.nspname::text - FROM pg_catalog.pg_class rc - JOIN pg_catalog.pg_namespace rn ON rn.oid = rc.relnamespace - WHERE rc.oid = pg_catalog.to_regclass(pg_catalog.quote_ident(:table_name::text)) - ), + (SELECT schema_name FROM resolved_table_schema), current_schema() ) AS schema_name ), diff --git a/sqlspec/data_dictionary/dialects/cockroachdb/sql/constraints.sql b/sqlspec/data_dictionary/dialects/cockroachdb/sql/constraints.sql index 8622bb1bc..4890a6cdf 100644 --- a/sqlspec/data_dictionary/dialects/cockroachdb/sql/constraints.sql +++ b/sqlspec/data_dictionary/dialects/cockroachdb/sql/constraints.sql @@ -1,5 +1,35 @@ -- name: by_schema -- dialect: cockroachdb +WITH con_cols AS ( + SELECT + kcu.constraint_catalog, + kcu.constraint_schema, + kcu.constraint_name, + kcu.table_schema, + kcu.table_name, + array_agg(kcu.column_name::text ORDER BY kcu.ordinal_position) AS columns + FROM information_schema.key_column_usage kcu + WHERE kcu.table_schema = :schema_name + AND (:table_name::text IS NULL OR kcu.table_name = :table_name) + GROUP BY + kcu.constraint_catalog, + kcu.constraint_schema, + kcu.constraint_name, + kcu.table_schema, + kcu.table_name +), +ref_cols AS ( + SELECT + pk_kcu.constraint_catalog, + pk_kcu.constraint_schema, + pk_kcu.constraint_name, + array_agg(pk_kcu.column_name::text ORDER BY pk_kcu.ordinal_position) AS referenced_columns + FROM information_schema.key_column_usage pk_kcu + GROUP BY + pk_kcu.constraint_catalog, + pk_kcu.constraint_schema, + pk_kcu.constraint_name +) SELECT tc.constraint_catalog::text AS constraint_catalog, tc.constraint_schema::text AS constraint_schema, @@ -9,27 +39,17 @@ SELECT tc.constraint_type::text AS constraint_type, tc.is_deferrable::text AS is_deferrable, tc.initially_deferred::text AS initially_deferred, - ARRAY( - SELECT kcu.column_name::text - FROM information_schema.key_column_usage kcu - WHERE kcu.constraint_catalog = tc.constraint_catalog - AND kcu.constraint_schema = tc.constraint_schema - AND kcu.constraint_name = tc.constraint_name - AND kcu.table_schema = tc.table_schema - AND kcu.table_name = tc.table_name - ORDER BY kcu.ordinal_position - )::text[] AS columns, + COALESCE(cc.columns, ARRAY[]::text[]) AS columns, rc.unique_constraint_schema::text AS referenced_schema, pk_tc.table_name::text AS referenced_table, - ARRAY( - SELECT pk_kcu.column_name::text - FROM information_schema.key_column_usage pk_kcu - WHERE pk_kcu.constraint_catalog = rc.unique_constraint_catalog - AND pk_kcu.constraint_schema = rc.unique_constraint_schema - AND pk_kcu.constraint_name = rc.unique_constraint_name - ORDER BY pk_kcu.ordinal_position - )::text[] AS referenced_columns + COALESCE(rc_cols.referenced_columns, ARRAY[]::text[]) AS referenced_columns FROM information_schema.table_constraints tc +LEFT JOIN con_cols cc + ON cc.constraint_catalog = tc.constraint_catalog + AND cc.constraint_schema = tc.constraint_schema + AND cc.constraint_name = tc.constraint_name + AND cc.table_schema = tc.table_schema + AND cc.table_name = tc.table_name LEFT JOIN information_schema.referential_constraints rc ON rc.constraint_catalog = tc.constraint_catalog AND rc.constraint_schema = tc.constraint_schema @@ -38,6 +58,10 @@ LEFT JOIN information_schema.table_constraints pk_tc ON pk_tc.constraint_catalog = rc.unique_constraint_catalog AND pk_tc.constraint_schema = rc.unique_constraint_schema AND pk_tc.constraint_name = rc.unique_constraint_name +LEFT JOIN ref_cols rc_cols + ON rc_cols.constraint_catalog = rc.unique_constraint_catalog + AND rc_cols.constraint_schema = rc.unique_constraint_schema + AND rc_cols.constraint_name = rc.unique_constraint_name WHERE tc.table_schema = :schema_name AND (:table_name::text IS NULL OR tc.table_name = :table_name) ORDER BY tc.table_schema, tc.table_name, tc.constraint_name; diff --git a/tests/unit/adapters/test_cockroach_asyncpg/test_config.py b/tests/unit/adapters/test_cockroach_asyncpg/test_config.py index dc53e3b14..374845c1c 100644 --- a/tests/unit/adapters/test_cockroach_asyncpg/test_config.py +++ b/tests/unit/adapters/test_cockroach_asyncpg/test_config.py @@ -638,3 +638,29 @@ async def test_cockroach_asyncpg_load_from_records_executemany() -> None: assert 'INSERT INTO "events"' in conn.executemany.await_args[0][0] conn.copy_records_to_table.assert_not_called() conn.copy_to_table.assert_not_called() + + +async def test_cockroach_asyncpg_get_columns_and_indexes_by_table() -> None: + """CockroachAsyncpgDataDictionary executes by_table query when table is specified.""" + from sqlspec.adapters.cockroach_asyncpg.data_dictionary import CockroachAsyncpgDataDictionary + + calls: list[tuple[Any, dict[str, Any]]] = [] + driver = MagicMock() + + async def fake_select(query: Any, **kwargs: Any) -> list[Any]: + calls.append((query, kwargs)) + return [] + + driver.select = AsyncMock(side_effect=fake_select) + dd = CockroachAsyncpgDataDictionary() + + await dd.get_columns(driver, table="users", schema=None) + assert len(calls) == 1 + assert calls[0][1]["table_name"] == "users" + assert calls[0][1]["schema_name"] is None + + calls.clear() + await dd.get_indexes(driver, table="users", schema=None) + assert len(calls) == 1 + assert calls[0][1]["table_name"] == "users" + assert calls[0][1]["schema_name"] is None diff --git a/tests/unit/data_dictionary/test_postgres_cockroach_query_packs.py b/tests/unit/data_dictionary/test_postgres_cockroach_query_packs.py index 9625896e9..290d4b534 100644 --- a/tests/unit/data_dictionary/test_postgres_cockroach_query_packs.py +++ b/tests/unit/data_dictionary/test_postgres_cockroach_query_packs.py @@ -165,3 +165,18 @@ def test_postgres_system_and_constraints_query_contracts() -> None: assert "WITH con_cols AS (" in con_query assert "CROSS JOIN LATERAL unnest(con.conkey) WITH ORDINALITY" in con_query assert "ARRAY(SELECT att.attname" not in con_query + + +def test_cockroach_constraints_and_columns_query_contracts() -> None: + """CockroachDB constraints and columns queries pre-aggregate and factor scalar subqueries.""" + loader = DataDictionaryLoader() + con_query = loader.get_domain_query_text("cockroachdb", "constraints", "by_schema") + assert con_query is not None + assert "WITH con_cols AS (" in con_query + assert "array_agg(kcu.column_name::text ORDER BY kcu.ordinal_position)" in con_query + assert "ARRAY(SELECT kcu.column_name" not in con_query + + cols_query = loader.get_domain_query_text("cockroachdb", "columns", "by_schema") + assert cols_query is not None + assert "WITH resolved_table_schema AS (" in cols_query + assert "FROM resolved_table_schema" in cols_query From f6d5b057a9427ecb4f8d1fa42881ea523ad4d4e1 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 01:05:46 +0000 Subject: [PATCH 080/174] perf(cockroach_asyncpg): eliminate transient dict allocations and add cache execution fast-path Replace transient dict allocations in native export/import helpers with string branch expressions and implement _execute_cache_hit on CockroachAsyncpgDriver to bypass dispatch_execute on statement cache hits. --- sqlspec/adapters/cockroach_asyncpg/core.py | 4 ++-- sqlspec/adapters/cockroach_asyncpg/driver.py | 23 +++++++++++++++++-- .../test_cockroach_asyncpg/test_config.py | 22 ++++++++++++++++++ 3 files changed, 45 insertions(+), 4 deletions(-) diff --git a/sqlspec/adapters/cockroach_asyncpg/core.py b/sqlspec/adapters/cockroach_asyncpg/core.py index d6352d7e8..4073f748a 100644 --- a/sqlspec/adapters/cockroach_asyncpg/core.py +++ b/sqlspec/adapters/cockroach_asyncpg/core.py @@ -234,7 +234,7 @@ def build_native_export( query: str, parameters: "list[Any]", uri: str, file_format: str, options: "dict[str, Any]" ) -> "tuple[str, list[Any]]": """Wrap compiled query SQL without embedding destination or CSV values.""" - format_sql = {"csv": "CSV", "parquet": "PARQUET"}[file_format] + format_sql = "PARQUET" if file_format == "parquet" else "CSV" values = [*parameters, uri] command = "EXPORT INTO " + format_sql + " " + ("$" + str(len(values))) if file_format == "csv" and "nullas" in options: @@ -250,7 +250,7 @@ def build_native_import(table: str, uri: str, file_format: str, options: "dict[s msg = "Table name must not be empty" raise ValueError(msg) target = ".".join(quote_identifier(part) for part in parts) - format_sql = {"csv": "CSV", "parquet": "PARQUET"}[file_format] + format_sql = "PARQUET" if file_format == "parquet" else "CSV" values: list[Any] = [uri] command = "IMPORT INTO " + target + " " + format_sql + " DATA (" + "$1" + ")" clauses = [] diff --git a/sqlspec/adapters/cockroach_asyncpg/driver.py b/sqlspec/adapters/cockroach_asyncpg/driver.py index 29d175451..003f6ec2b 100644 --- a/sqlspec/adapters/cockroach_asyncpg/driver.py +++ b/sqlspec/adapters/cockroach_asyncpg/driver.py @@ -32,8 +32,8 @@ from sqlspec.adapters.asyncpg.core import AsyncpgStreamSource from sqlspec.adapters.cockroach_asyncpg._typing import CockroachAsyncpgConnection - from sqlspec.core import StatementConfig - from sqlspec.driver import ExecutionResult + from sqlspec.core import SQLResult, StatementConfig + from sqlspec.driver import CachedQuery, ExecutionResult from sqlspec.storage import StorageBridgeJob, StorageDestination, StorageFormat, StorageTelemetry __all__ = ( @@ -438,6 +438,25 @@ async def dispatch_execute(self, cursor: Any, statement: SQL) -> "ExecutionResul await self.commit() return result + async def _execute_cache_hit( + self, sql: str, params: "tuple[Any, ...] | list[Any] | dict[str, Any]", cached: "CachedQuery" + ) -> "SQLResult": + auto_started = ( + await self._begin_follower_read_transaction() + if (cached.operation_profile.returns_rows and not cached.operation_profile.modifies_rows) + else False + ) + try: + result = await super()._execute_cache_hit(sql, params, cached) + except BaseException: + if auto_started and self._connection_in_transaction(): + with suppress(Exception): + await self.rollback() + raise + if auto_started and self._connection_in_transaction(): + await self.commit() + return result + def dispatch_select_stream(self, statement: "SQL", chunk_size: int) -> "AsyncRowStream[dict[str, Any]] | None": """Return a native asyncpg row stream wrapped for follower reads.""" stream = super().dispatch_select_stream(statement, chunk_size) diff --git a/tests/unit/adapters/test_cockroach_asyncpg/test_config.py b/tests/unit/adapters/test_cockroach_asyncpg/test_config.py index 374845c1c..250abefd2 100644 --- a/tests/unit/adapters/test_cockroach_asyncpg/test_config.py +++ b/tests/unit/adapters/test_cockroach_asyncpg/test_config.py @@ -664,3 +664,25 @@ async def fake_select(query: Any, **kwargs: Any) -> list[Any]: assert len(calls) == 1 assert calls[0][1]["table_name"] == "users" assert calls[0][1]["schema_name"] is None + + +async def test_cockroach_asyncpg_cache_hit_with_cursor() -> None: + """CockroachAsyncpgDriver executes cached queries directly on connection via _execute_cache_hit.""" + conn = MagicMock() + conn.fetch = AsyncMock(return_value=[{"val": 42}]) + conn.execute = AsyncMock(return_value="UPDATE 1") + + driver = CockroachAsyncpgDriver(cast("Any", conn)) + res1 = await driver.execute("SELECT 42") + assert res1.data == [{"val": 42}] + + res2 = await driver.execute("SELECT 42") + assert res2.data == [{"val": 42}] + assert conn.fetch.call_count == 2 + + upd1 = await driver.execute("UPDATE t SET x = 1") + assert upd1.rows_affected == 1 + + upd2 = await driver.execute("UPDATE t SET x = 1") + assert upd2.rows_affected == 1 + assert conn.execute.call_count == 2 From 4d6c7c3d1711b9b5eeab2d47472136201eae6d2e Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 01:05:46 +0000 Subject: [PATCH 081/174] refactor(cockroach_asyncpg): code hygiene, timeouts, and typing cleanup Wrap _execute_native_storage in ambient exception guard, document auto-retry semantics, expose statement_timeout and idle_in_transaction_session_timeout on CockroachAsyncpgConnectionConfig, simplify _CockroachAsyncpgSessionFactory, and remove inline type ignore comments. --- sqlspec/adapters/cockroach_asyncpg/config.py | 25 +++------- .../cockroach_asyncpg/data_dictionary.py | 50 +++++++++---------- sqlspec/adapters/cockroach_asyncpg/driver.py | 19 ++++--- .../test_cockroach_asyncpg/test_config.py | 20 ++++++++ 4 files changed, 64 insertions(+), 50 deletions(-) diff --git a/sqlspec/adapters/cockroach_asyncpg/config.py b/sqlspec/adapters/cockroach_asyncpg/config.py index e5bad024f..d204ac063 100644 --- a/sqlspec/adapters/cockroach_asyncpg/config.py +++ b/sqlspec/adapters/cockroach_asyncpg/config.py @@ -111,6 +111,10 @@ class CockroachAsyncpgConnectionConfig(TypedDict): """Startup CockroachDB session variables passed at connection time.""" target_session_attrs: NotRequired[CockroachAsyncpgTargetSessionAttrs] """Required node session attribute when connecting across multiple hosts.""" + statement_timeout: NotRequired[float | int | None] + """Statement execution timeout in seconds forwarded via server_settings.""" + idle_in_transaction_session_timeout: NotRequired[float | int | None] + """Idle in transaction timeout in seconds forwarded via server_settings.""" krbsrvname: NotRequired[str] """Kerberos service principal name for GSSAPI authentication.""" gsslib: NotRequired[CockroachAsyncpgGSSLib] @@ -192,24 +196,9 @@ class CockroachAsyncpgDriverFeatures(TypedDict): class _CockroachAsyncpgSessionFactory(AsyncPoolSessionFactory): - """Uses pool.acquire() context manager pattern instead of direct acquire/release.""" + """Factory creating AsyncPG database connections for CockroachDB sessions.""" - __slots__ = ("_ctx",) - - def __init__(self, config: "CockroachAsyncpgConfig") -> None: - super().__init__(config) - self._ctx: Any | None = None - - async def acquire_connection(self) -> "CockroachAsyncpgConnection": - pool = await self._config.provide_pool() - ctx = pool.acquire() - self._ctx = ctx - return cast("CockroachAsyncpgConnection", await ctx.__aenter__()) - - async def release_connection(self, _conn: "CockroachAsyncpgConnection", **kwargs: Any) -> None: - if self._ctx is not None: - await self._ctx.__aexit__(kwargs.get("exc_type"), kwargs.get("exc_val"), kwargs.get("exc_tb")) - self._ctx = None + __slots__ = () class CockroachAsyncpgConnectionContext(AsyncPoolConnectionContext): @@ -224,7 +213,7 @@ class CockroachAsyncpgConfig( """Configuration for CockroachDB using AsyncPG.""" driver_type: "ClassVar[type[CockroachAsyncpgDriver]]" = CockroachAsyncpgDriver - connection_type: "ClassVar[type[CockroachAsyncpgConnection]]" = CockroachAsyncpgConnection # type: ignore[assignment] + connection_type: "ClassVar[type[Any]]" = cast("type[Any]", CockroachAsyncpgConnection) supports_transactional_ddl: "ClassVar[bool]" = False supports_migration_schemas: "ClassVar[bool]" = True supports_native_arrow_export: "ClassVar[bool]" = True diff --git a/sqlspec/adapters/cockroach_asyncpg/data_dictionary.py b/sqlspec/adapters/cockroach_asyncpg/data_dictionary.py index e86fe76e2..888fef1bd 100644 --- a/sqlspec/adapters/cockroach_asyncpg/data_dictionary.py +++ b/sqlspec/adapters/cockroach_asyncpg/data_dictionary.py @@ -5,6 +5,7 @@ from mypy_extensions import mypyc_attr +from sqlspec.adapters.asyncpg.data_dictionary import AsyncpgDataDictionary from sqlspec.data_dictionary import ( ColumnMetadata, DDLResult, @@ -33,46 +34,43 @@ COCKROACH_SUPPORTED_DOMAINS, resolve_cockroachdb_json_type, ) -from sqlspec.driver import AsyncDataDictionaryBase if TYPE_CHECKING: - from sqlspec.adapters.cockroach_asyncpg.driver import CockroachAsyncpgDriver + from sqlspec.adapters.asyncpg.driver import AsyncpgDriver from sqlspec.core import SQL __all__ = ("CockroachAsyncpgDataDictionary",) @mypyc_attr(allow_interpreted_subclasses=True, native_class=False) -class CockroachAsyncpgDataDictionary(AsyncDataDictionaryBase): +class CockroachAsyncpgDataDictionary(AsyncpgDataDictionary): """CockroachDB async data dictionary (AsyncPG).""" dialect: ClassVar[str] = "cockroachdb" async def get_metadata_capabilities( - self, driver: "CockroachAsyncpgDriver", domains: Sequence[str] | None = None + self, driver: "AsyncpgDriver", domains: Sequence[str] | None = None ) -> MetadataCapabilityProfile: """Get CockroachDB data-dictionary capability profile.""" return _cockroach_metadata_profile(type(self).__name__, domains) async def get_system_metadata_capabilities( - self, driver: "CockroachAsyncpgDriver", domains: Sequence[str] | None = None + self, driver: "AsyncpgDriver", domains: Sequence[str] | None = None ) -> tuple[SystemMetadataCapability, ...]: """Get CockroachDB opt-in system metadata capability disclosures.""" _ = driver requested_domains = ("system",) if domains is None else tuple(domains) return tuple(unsupported_system_metadata_capability(domain) for domain in requested_domains) - async def get_schemas(self, driver: "CockroachAsyncpgDriver") -> MetadataResult: + async def get_schemas(self, driver: "AsyncpgDriver") -> MetadataResult: """Get schema metadata.""" return await self._select_domain(driver, "schemas", "list") - async def get_objects(self, driver: "CockroachAsyncpgDriver", schema: str | None = None) -> MetadataResult: + async def get_objects(self, driver: "AsyncpgDriver", schema: str | None = None) -> MetadataResult: """Get database object metadata.""" return await self._select_domain(driver, "objects", "by_schema", schema_name=self.resolve_schema(schema)) - async def get_table_details( - self, driver: "CockroachAsyncpgDriver", table: str, schema: str | None = None - ) -> MetadataResult: + async def get_table_details(self, driver: "AsyncpgDriver", table: str, schema: str | None = None) -> MetadataResult: """Get rich table metadata.""" return await self._select_domain( driver, @@ -83,7 +81,7 @@ async def get_table_details( ) async def get_constraints( - self, driver: "CockroachAsyncpgDriver", table: str | None = None, schema: str | None = None + self, driver: "AsyncpgDriver", table: str | None = None, schema: str | None = None ) -> MetadataResult: """Get constraint metadata.""" table_name = self.resolve_identifier(table) if table is not None else None @@ -91,16 +89,16 @@ async def get_constraints( driver, "constraints", "by_schema", schema_name=self.resolve_schema(schema), table_name=table_name ) - async def get_views(self, driver: "CockroachAsyncpgDriver", schema: str | None = None) -> MetadataResult: + async def get_views(self, driver: "AsyncpgDriver", schema: str | None = None) -> MetadataResult: """Get view metadata.""" return await self._select_domain(driver, "views", "by_schema", schema_name=self.resolve_schema(schema)) - async def get_routines(self, driver: "CockroachAsyncpgDriver", schema: str | None = None) -> MetadataResult: + async def get_routines(self, driver: "AsyncpgDriver", schema: str | None = None) -> MetadataResult: """Get routine metadata.""" return _metadata_result("routines", MetadataCapability.unsupported("routines")) async def get_privileges( - self, driver: "CockroachAsyncpgDriver", object_name: str | None = None, schema: str | None = None + self, driver: "AsyncpgDriver", object_name: str | None = None, schema: str | None = None ) -> MetadataResult: """Get privilege metadata.""" resolved_object = self.resolve_identifier(object_name) if object_name is not None else None @@ -109,7 +107,7 @@ async def get_privileges( ) async def get_dependencies( - self, driver: "CockroachAsyncpgDriver", object_name: str | None = None, schema: str | None = None + self, driver: "AsyncpgDriver", object_name: str | None = None, schema: str | None = None ) -> MetadataResult: """Get stable dependency metadata.""" resolved_object = self.resolve_identifier(object_name) if object_name is not None else None @@ -117,12 +115,12 @@ async def get_dependencies( driver, "dependencies", "by_schema", schema_name=self.resolve_schema(schema), object_name=resolved_object ) - async def get_sequences(self, driver: "CockroachAsyncpgDriver", schema: str | None = None) -> MetadataResult: + async def get_sequences(self, driver: "AsyncpgDriver", schema: str | None = None) -> MetadataResult: """Get sequence metadata.""" return await self._select_domain(driver, "sequences", "by_schema", schema_name=self.resolve_schema(schema)) async def get_system_metadata( - self, driver: "CockroachAsyncpgDriver", request: SystemMetadataRequest | str | None = None, **kwargs: Any + self, driver: "AsyncpgDriver", request: SystemMetadataRequest | str | None = None, **kwargs: Any ) -> SystemMetadataResult: """Get opt-in CockroachDB system metadata.""" _ = driver @@ -130,7 +128,7 @@ async def get_system_metadata( capability = unsupported_system_metadata_capability(metadata_request.domain) return system_metadata_gated_result(metadata_request, capability) - async def get_version(self, driver: "CockroachAsyncpgDriver") -> "VersionInfo | None": + async def get_version(self, driver: "AsyncpgDriver") -> "VersionInfo | None": driver_id = id(getattr(driver, "connection", None) or driver) if driver_id in self._version_fetch_attempted: return self._version_cache.get(driver_id) @@ -151,17 +149,17 @@ async def get_version(self, driver: "CockroachAsyncpgDriver") -> "VersionInfo | self.cache_version(driver_id, version_info) return version_info - async def get_feature_flag(self, driver: "CockroachAsyncpgDriver", feature: str) -> bool: + async def get_feature_flag(self, driver: "AsyncpgDriver", feature: str) -> bool: version_info = await self.get_version(driver) return self.resolve_feature_flag(feature, version_info) - async def get_optimal_type(self, driver: "CockroachAsyncpgDriver", type_category: str) -> str: + async def get_optimal_type(self, driver: "AsyncpgDriver", type_category: str) -> str: config = self.get_dialect_config() if type_category == "json": return resolve_cockroachdb_json_type(await self.get_version(driver)) return config.get_optimal_type(type_category) - async def get_tables(self, driver: "CockroachAsyncpgDriver", schema: "str | None" = None) -> "list[TableMetadata]": + async def get_tables(self, driver: "AsyncpgDriver", schema: "str | None" = None) -> "list[TableMetadata]": schema_name = self.resolve_schema(schema) self._log_schema_introspect(driver, schema_name=schema_name, table_name=None, operation="tables") return await driver.select( @@ -172,7 +170,7 @@ async def get_tables(self, driver: "CockroachAsyncpgDriver", schema: "str | None ) async def get_columns( - self, driver: "CockroachAsyncpgDriver", table: "str | None" = None, schema: "str | None" = None + self, driver: "AsyncpgDriver", table: "str | None" = None, schema: "str | None" = None ) -> "list[ColumnMetadata]": schema_name = self.resolve_schema(schema) if table is None: @@ -195,7 +193,7 @@ async def get_columns( ) async def get_indexes( - self, driver: "CockroachAsyncpgDriver", table: "str | None" = None, schema: "str | None" = None + self, driver: "AsyncpgDriver", table: "str | None" = None, schema: "str | None" = None ) -> "list[IndexMetadata]": schema_name = self.resolve_schema(schema) if table is None: @@ -218,7 +216,7 @@ async def get_indexes( ) async def get_foreign_keys( - self, driver: "CockroachAsyncpgDriver", table: "str | None" = None, schema: "str | None" = None + self, driver: "AsyncpgDriver", table: "str | None" = None, schema: "str | None" = None ) -> "list[ForeignKeyMetadata]": """Get foreign key metadata, including ON DELETE / ON UPDATE actions.""" schema_name = self.resolve_schema(schema) @@ -235,7 +233,7 @@ async def get_foreign_keys( async def _get_native_object_ddl( self, - driver: "CockroachAsyncpgDriver", + driver: "AsyncpgDriver", object_name: str, *, schema: str | None, @@ -262,7 +260,7 @@ async def _get_native_object_ddl( ) async def _select_domain( - self, driver: "CockroachAsyncpgDriver", domain: str, query_name: str, **parameters: Any + self, driver: "AsyncpgDriver", domain: str, query_name: str, **parameters: Any ) -> MetadataResult: query = get_data_dictionary_loader().get_domain_query(type(self).dialect, domain, query_name) if not query.is_supported or query.sql is None: diff --git a/sqlspec/adapters/cockroach_asyncpg/driver.py b/sqlspec/adapters/cockroach_asyncpg/driver.py index 003f6ec2b..66b1e5ed8 100644 --- a/sqlspec/adapters/cockroach_asyncpg/driver.py +++ b/sqlspec/adapters/cockroach_asyncpg/driver.py @@ -1,6 +1,7 @@ """CockroachDB AsyncPG driver implementation.""" import asyncio +import sys from collections.abc import Mapping, Sequence from contextlib import suppress from typing import TYPE_CHECKING, Any, Literal, TypeVar, cast @@ -116,7 +117,7 @@ def __init__( self._data_dictionary = None @property - def data_dictionary(self) -> "CockroachAsyncpgDataDictionary": # type: ignore[override] + def data_dictionary(self) -> "CockroachAsyncpgDataDictionary": if self._data_dictionary is None: object.__setattr__(self, "_data_dictionary", CockroachAsyncpgDataDictionary()) return cast("CockroachAsyncpgDataDictionary", self._data_dictionary) @@ -473,11 +474,17 @@ def _native_storage_ready(self) -> bool: async def _execute_native_storage(self, command: str, parameters: "list[Any]") -> "list[dict[str, Any]]": handler = self.handle_database_exceptions() - rows = [] - async with handler: - rows = [dict(row) for row in await self.connection.fetch(command, *parameters)] - if handler.pending_exception is not None: - raise handler.pending_exception + + async def _fetch() -> list[dict[str, Any]]: + return [dict(row) for row in await self.connection.fetch(command, *parameters)] + + if sys.exc_info()[1] is None: + async with handler: + rows = await _fetch() + else: + res = await self._run_with_exception_handler(handler, _fetch) + rows = res if res is not None else [] + self._check_pending_exception(handler) return rows async def _apply_follower_reads(self) -> None: diff --git a/tests/unit/adapters/test_cockroach_asyncpg/test_config.py b/tests/unit/adapters/test_cockroach_asyncpg/test_config.py index 250abefd2..1313b7162 100644 --- a/tests/unit/adapters/test_cockroach_asyncpg/test_config.py +++ b/tests/unit/adapters/test_cockroach_asyncpg/test_config.py @@ -686,3 +686,23 @@ async def test_cockroach_asyncpg_cache_hit_with_cursor() -> None: upd2 = await driver.execute("UPDATE t SET x = 1") assert upd2.rows_affected == 1 assert conn.execute.call_count == 2 + + +def test_cockroach_asyncpg_connection_config_timeout_annotations() -> None: + """CockroachAsyncpgConnectionConfig exposes statement_timeout and idle_in_transaction_session_timeout.""" + annotations = CockroachAsyncpgConnectionConfig.__annotations__ + assert "statement_timeout" in annotations + assert "idle_in_transaction_session_timeout" in annotations + + +async def test_cockroach_asyncpg_execute_native_storage_under_active_exception() -> None: + """_execute_native_storage inside active except block does not pollute pending exception.""" + conn = MagicMock() + conn.fetch = AsyncMock(return_value=[{"rows": 1}]) + driver = CockroachAsyncpgDriver(cast("Any", conn)) + + try: + raise ValueError("outer error") + except ValueError: + rows = await driver._execute_native_storage("SELECT 1", []) + assert rows == [{"rows": 1}] From fb3572fa8e143a08165539359ee2149d145552f5 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 01:05:47 +0000 Subject: [PATCH 082/174] fix(cockroach_psycopg): batch ADK memory inserts, vector search, and shard validation Batch insert_memory_entries via UNNEST query, implement vector cosine distance search using embedding in search_entries for sync and async stores, and align hash_shard_bucket_count validation with litestar store. --- .../adapters/cockroach_psycopg/adk/store.py | 180 ++++++++++++++---- .../test_cockroach_psycopg/test_config.py | 135 ++++++++++++- 2 files changed, 279 insertions(+), 36 deletions(-) diff --git a/sqlspec/adapters/cockroach_psycopg/adk/store.py b/sqlspec/adapters/cockroach_psycopg/adk/store.py index 12eac5f36..266f0b65c 100644 --- a/sqlspec/adapters/cockroach_psycopg/adk/store.py +++ b/sqlspec/adapters/cockroach_psycopg/adk/store.py @@ -1234,7 +1234,21 @@ async def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: " if not entries: return 0 - inserted_count = 0 + ids = [entry["id"] for entry in entries] + session_ids = [entry["session_id"] for entry in entries] + app_names = [entry["app_name"] for entry in entries] + user_ids = [entry["user_id"] for entry in entries] + scopes = [entry.get("scope", "user") for entry in entries] + event_ids = [entry["event_id"] for entry in entries] + authors = [entry["author"] for entry in entries] + timestamps = [entry["timestamp"] for entry in entries] + content_jsons = [Jsonb(entry["content_json"]) for entry in entries] + content_texts = [entry["content_text"] for entry in entries] + metadata_jsons = [ + Jsonb(entry["metadata_json"]) if entry["metadata_json"] is not None else None for entry in entries + ] + inserted_ats = [entry["inserted_at"] for entry in entries] + table_ident = pg_sql.Identifier(*self._memory_table.split(".")) if self._owner_id_column_name: query = pg_sql.SQL(""" @@ -1242,30 +1256,61 @@ async def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: " id, session_id, app_name, user_id, scope, event_id, author, {owner_id_col}, timestamp, content_json, content_text, metadata_json, inserted_at - ) VALUES ( - %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s + ) + SELECT * FROM unnest( + %s::varchar[], %s::varchar[], %s::varchar[], %s::varchar[], + %s::varchar[], %s::varchar[], %s::varchar[], %s::varchar[], + %s::timestamptz[], %s::jsonb[], %s::text[], %s::jsonb[], %s::timestamptz[] ) ON CONFLICT (event_id) DO NOTHING """).format(table=table_ident, owner_id_col=pg_sql.Identifier(self._owner_id_column_name)) + owner_ids = [owner_id for _ in entries] + params: tuple[Any, ...] = ( + ids, + session_ids, + app_names, + user_ids, + scopes, + event_ids, + authors, + owner_ids, + timestamps, + content_jsons, + content_texts, + metadata_jsons, + inserted_ats, + ) else: query = pg_sql.SQL(""" INSERT INTO {table} ( id, session_id, app_name, user_id, scope, event_id, author, timestamp, content_json, content_text, metadata_json, inserted_at - ) VALUES ( - %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s + ) + SELECT * FROM unnest( + %s::varchar[], %s::varchar[], %s::varchar[], %s::varchar[], + %s::varchar[], %s::varchar[], %s::varchar[], + %s::timestamptz[], %s::jsonb[], %s::text[], %s::jsonb[], %s::timestamptz[] ) ON CONFLICT (event_id) DO NOTHING """).format(table=table_ident) + params = ( + ids, + session_ids, + app_names, + user_ids, + scopes, + event_ids, + authors, + timestamps, + content_jsons, + content_texts, + metadata_jsons, + inserted_ats, + ) async with self._config.provide_connection() as conn, conn.cursor() as cur: - for entry in entries: - if self._owner_id_column_name: - await cur.execute(query, _build_insert_params_with_owner(entry, owner_id)) - else: - await cur.execute(query, _build_insert_params(entry)) - if cur.rowcount and cur.rowcount > 0: - inserted_count += cur.rowcount + await cur.execute(query, params) + inserted_count = cur.rowcount if cur.rowcount and cur.rowcount > 0 else 0 await conn.commit() return inserted_count @@ -1283,13 +1328,21 @@ async def search_entries( msg = "Memory store is disabled" raise RuntimeError(msg) - if not query: + if not query and embedding is None: return [] effective_limit = limit if limit is not None else self._max_results where_scope, scope_params = _build_cockroach_scope_where(app_name, user_id, scope_filter) - if self._use_fts: + if embedding is not None: + sql = f""" + SELECT * FROM {self._memory_table} + WHERE {where_scope} AND embedding IS NOT NULL + ORDER BY embedding <=> %s::float8[]::vector + LIMIT %s + """ + params = (*scope_params, list(embedding), effective_limit) + elif self._use_fts: sql = f""" SELECT * FROM {self._memory_table} WHERE {where_scope} @@ -1298,6 +1351,7 @@ async def search_entries( LIMIT %s """ search_param = query + params = (*scope_params, search_param, effective_limit) else: sql = f""" SELECT * FROM {self._memory_table} @@ -1306,8 +1360,7 @@ async def search_entries( LIMIT %s """ search_param = f"%{query}%" - - params = (*scope_params, search_param, effective_limit) + params = (*scope_params, search_param, effective_limit) try: async with self._config.provide_connection() as conn, conn.cursor() as cur: @@ -1420,7 +1473,21 @@ def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "object if not entries: return 0 - inserted_count = 0 + ids = [entry["id"] for entry in entries] + session_ids = [entry["session_id"] for entry in entries] + app_names = [entry["app_name"] for entry in entries] + user_ids = [entry["user_id"] for entry in entries] + scopes = [entry.get("scope", "user") for entry in entries] + event_ids = [entry["event_id"] for entry in entries] + authors = [entry["author"] for entry in entries] + timestamps = [entry["timestamp"] for entry in entries] + content_jsons = [Jsonb(entry["content_json"]) for entry in entries] + content_texts = [entry["content_text"] for entry in entries] + metadata_jsons = [ + Jsonb(entry["metadata_json"]) if entry["metadata_json"] is not None else None for entry in entries + ] + inserted_ats = [entry["inserted_at"] for entry in entries] + table_ident = pg_sql.Identifier(*self._memory_table.split(".")) if self._owner_id_column_name: query = pg_sql.SQL(""" @@ -1428,31 +1495,63 @@ def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "object id, session_id, app_name, user_id, scope, event_id, author, {owner_id_col}, timestamp, content_json, content_text, metadata_json, inserted_at - ) VALUES ( - %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s + ) + SELECT * FROM unnest( + %s::varchar[], %s::varchar[], %s::varchar[], %s::varchar[], + %s::varchar[], %s::varchar[], %s::varchar[], %s::varchar[], + %s::timestamptz[], %s::jsonb[], %s::text[], %s::jsonb[], %s::timestamptz[] ) ON CONFLICT (event_id) DO NOTHING """).format(table=table_ident, owner_id_col=pg_sql.Identifier(self._owner_id_column_name)) + owner_ids = [owner_id for _ in entries] + params: tuple[Any, ...] = ( + ids, + session_ids, + app_names, + user_ids, + scopes, + event_ids, + authors, + owner_ids, + timestamps, + content_jsons, + content_texts, + metadata_jsons, + inserted_ats, + ) else: query = pg_sql.SQL(""" INSERT INTO {table} ( id, session_id, app_name, user_id, scope, event_id, author, timestamp, content_json, content_text, metadata_json, inserted_at - ) VALUES ( - %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s + ) + SELECT * FROM unnest( + %s::varchar[], %s::varchar[], %s::varchar[], %s::varchar[], + %s::varchar[], %s::varchar[], %s::varchar[], + %s::timestamptz[], %s::jsonb[], %s::text[], %s::jsonb[], %s::timestamptz[] ) ON CONFLICT (event_id) DO NOTHING """).format(table=table_ident) + params = ( + ids, + session_ids, + app_names, + user_ids, + scopes, + event_ids, + authors, + timestamps, + content_jsons, + content_texts, + metadata_jsons, + inserted_ats, + ) with self._config.provide_connection() as conn, conn.cursor() as cur: - for entry in entries: - if self._owner_id_column_name: - cur.execute(query, _build_insert_params_with_owner(entry, owner_id)) - else: - cur.execute(query, _build_insert_params(entry)) - if cur.rowcount and cur.rowcount > 0: - inserted_count += cur.rowcount + cur.execute(query, params) + inserted_count = cur.rowcount if cur.rowcount and cur.rowcount > 0 else 0 conn.commit() + return inserted_count def search_entries( @@ -1469,13 +1568,21 @@ def search_entries( msg = "Memory store is disabled" raise RuntimeError(msg) - if not query: + if not query and embedding is None: return [] effective_limit = limit if limit is not None else self._max_results where_scope, scope_params = _build_cockroach_scope_where(app_name, user_id, scope_filter) - if self._use_fts: + if embedding is not None: + sql = f""" + SELECT * FROM {self._memory_table} + WHERE {where_scope} AND embedding IS NOT NULL + ORDER BY embedding <=> %s::float8[]::vector + LIMIT %s + """ + params = (*scope_params, list(embedding), effective_limit) + elif self._use_fts: sql = f""" SELECT * FROM {self._memory_table} WHERE {where_scope} @@ -1484,6 +1591,7 @@ def search_entries( LIMIT %s """ search_param = query + params = (*scope_params, search_param, effective_limit) else: sql = f""" SELECT * FROM {self._memory_table} @@ -1492,8 +1600,7 @@ def search_entries( LIMIT %s """ search_param = f"%{query}%" - - params = (*scope_params, search_param, effective_limit) + params = (*scope_params, search_param, effective_limit) try: with self._config.provide_connection() as conn, conn.cursor() as cur: @@ -1647,9 +1754,12 @@ def _cockroach_hash_shard_clause(adk_config: CockroachPsycopgADKConfig) -> str: if not adk_config.get("enable_hash_sharded_indexes", False): return "" bucket_count = adk_config.get("hash_shard_bucket_count") - if isinstance(bucket_count, int) and not isinstance(bucket_count, bool) and bucket_count > 0: - return f" USING HASH WITH (bucket_count = {bucket_count})" - return " USING HASH" + if bucket_count is None: + return " USING HASH" + if not isinstance(bucket_count, int) or isinstance(bucket_count, bool) or bucket_count <= 0: + msg = "extension_config['adk']['hash_shard_bucket_count'] must be a positive integer" + raise ValueError(msg) + return f" USING HASH WITH (bucket_count = {bucket_count})" def _cockroach_storing_clause(adk_config: CockroachPsycopgADKConfig, columns: tuple[str, ...]) -> str: diff --git a/tests/unit/adapters/test_cockroach_psycopg/test_config.py b/tests/unit/adapters/test_cockroach_psycopg/test_config.py index d5a9322f2..50f034f76 100644 --- a/tests/unit/adapters/test_cockroach_psycopg/test_config.py +++ b/tests/unit/adapters/test_cockroach_psycopg/test_config.py @@ -8,7 +8,7 @@ - Connection config normalization """ -from typing import Any, cast +from typing import Any, Self, cast from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -938,3 +938,136 @@ async def async_user_cb(conn: Any) -> None: conn = await config.create_connection() assert conn is fake_conn assert call_order == ["internal_autocommit", "user_cb"] + + +def test_cockroach_psycopg_adk_hash_shard_bucket_count_validation() -> None: + """Invalid hash_shard_bucket_count must raise ValueError in cockroach_psycopg adk store.""" + from sqlspec.adapters.cockroach_psycopg.adk.store import _cockroach_hash_shard_clause + + bad_config = CockroachPsycopgSyncConfig( + extension_config={"adk": {"enable_hash_sharded_indexes": True, "hash_shard_bucket_count": -1}} + ) + with pytest.raises(ValueError, match="hash_shard_bucket_count"): + _cockroach_hash_shard_clause(bad_config.extension_config["adk"]) + + +async def test_cockroach_psycopg_adk_memory_search_with_embedding() -> None: + """search_entries with embedding executes vector cosine search in cockroach_psycopg.""" + executed_queries: list[str] = [] + + class _FakeCursor: + description = [ + ("id",), + ("session_id",), + ("app_name",), + ("user_id",), + ("scope",), + ("event_id",), + ("author",), + ("timestamp",), + ("content_json",), + ("content_text",), + ("metadata_json",), + ("inserted_at",), + ("embedding",), + ] + + async def execute(self, sql: bytes, params: Any) -> None: + executed_queries.append(sql.decode()) + + async def fetchall(self) -> list[Any]: + return [] + + async def __aenter__(self) -> Self: + return self + + async def __aexit__(self, *_a: Any) -> None: + pass + + class _FakeConn: + def cursor(self) -> _FakeCursor: + return _FakeCursor() + + class _FakeCtx: + async def __aenter__(self) -> _FakeConn: + return _FakeConn() + + async def __aexit__(self, *_a: Any) -> None: + pass + + config = CockroachPsycopgAsyncConfig(extension_config={"adk": {"enable_memory": True}}) + object.__setattr__(config, "provide_connection", lambda: _FakeCtx()) + store = CockroachPsycopgAsyncADKMemoryStore(config) + + await store.search_entries(query="hello", app_name="app", user_id="u1", embedding=[0.1, 0.2, 0.3]) + assert len(executed_queries) == 1 + assert "embedding <=>" in executed_queries[0] + + +async def test_cockroach_psycopg_adk_memory_insert_batch_unnest() -> None: + """insert_memory_entries executes a single UNNEST query rather than an N+1 loop.""" + executed_queries: list[str] = [] + + class _FakeCursor: + rowcount = 2 + + async def execute(self, sql: Any, params: Any) -> None: + executed_queries.append(str(sql)) + + async def __aenter__(self) -> Self: + return self + + async def __aexit__(self, *_a: Any) -> None: + pass + + class _FakeConn: + def cursor(self) -> _FakeCursor: + return _FakeCursor() + + async def commit(self) -> None: + pass + + class _FakeCtx: + async def __aenter__(self) -> _FakeConn: + return _FakeConn() + + async def __aexit__(self, *_a: Any) -> None: + pass + + config = CockroachPsycopgAsyncConfig(extension_config={"adk": {"enable_memory": True}}) + object.__setattr__(config, "provide_connection", lambda: _FakeCtx()) + store = CockroachPsycopgAsyncADKMemoryStore(config) + + entry1: Any = { + "id": "1", + "session_id": "s1", + "app_name": "app", + "user_id": "u1", + "scope": "user", + "event_id": "e1", + "author": "user", + "timestamp": "2026-01-01T00:00:00Z", + "content_json": {}, + "content_text": "hello", + "metadata_json": None, + "inserted_at": "2026-01-01T00:00:00Z", + } + entry2: Any = { + "id": "2", + "session_id": "s1", + "app_name": "app", + "user_id": "u1", + "scope": "user", + "event_id": "e2", + "author": "user", + "timestamp": "2026-01-01T00:00:00Z", + "content_json": {}, + "content_text": "world", + "metadata_json": None, + "inserted_at": "2026-01-01T00:00:00Z", + } + + count = await store.insert_memory_entries([entry1, entry2]) + assert count == 2 + assert len(executed_queries) == 1 + assert "unnest" in executed_queries[0].lower() From 87f8fbed64efe4d347649059cf3d8d0e6919dbe8 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 01:05:47 +0000 Subject: [PATCH 083/174] fix(cockroach_psycopg): data dictionary queries and metadata consolidation Route table-specific columns and indexes queries to by_table in CockroachPsycopgSyncDataDictionary and CockroachPsycopgAsyncDataDictionary, and consolidate CockroachDB metadata domain constants into dialect config. --- .../cockroach_psycopg/data_dictionary.py | 47 ++++++----------- .../test_cockroach_psycopg/test_config.py | 52 +++++++++++++++++++ 2 files changed, 69 insertions(+), 30 deletions(-) diff --git a/sqlspec/adapters/cockroach_psycopg/data_dictionary.py b/sqlspec/adapters/cockroach_psycopg/data_dictionary.py index e4b5e39a7..e6007b5be 100644 --- a/sqlspec/adapters/cockroach_psycopg/data_dictionary.py +++ b/sqlspec/adapters/cockroach_psycopg/data_dictionary.py @@ -28,7 +28,11 @@ system_metadata_gated_result, unsupported_system_metadata_capability, ) -from sqlspec.data_dictionary.dialects.cockroachdb import resolve_cockroachdb_json_type +from sqlspec.data_dictionary.dialects.cockroachdb import ( + COCKROACH_METADATA_DOMAINS, + COCKROACH_SUPPORTED_DOMAINS, + resolve_cockroachdb_json_type, +) from sqlspec.driver import AsyncDataDictionaryBase, SyncDataDictionaryBase if TYPE_CHECKING: @@ -38,25 +42,6 @@ __all__ = ("CockroachPsycopgAsyncDataDictionary", "CockroachPsycopgSyncDataDictionary") -_COCKROACH_METADATA_DOMAINS: tuple[str, ...] = ( - "schemas", - "objects", - "tables", - "columns", - "constraints", - "indexes", - "views", - "sequences", - "comments", - "privileges", - "dependencies", - "ddl", - "crdb_internal", - "system", -) -_COCKROACH_SUPPORTED_DOMAINS = frozenset(d for d in _COCKROACH_METADATA_DOMAINS if d not in {"crdb_internal", "system"}) - - @mypyc_attr(allow_interpreted_subclasses=True, native_class=False) class CockroachPsycopgSyncDataDictionary(SyncDataDictionaryBase): """CockroachDB sync data dictionary.""" @@ -211,7 +196,7 @@ def get_columns( table_schema_name = None if schema is None else schema_name self._log_table_describe(driver, schema_name=table_schema_name, table_name=table_name, operation="columns") return driver.select( - _cockroach_domain_sql("columns", "by_schema"), + _cockroach_domain_sql("columns", "by_table"), schema_name=table_schema_name, table_name=table_name, schema_type=ColumnMetadata, @@ -232,10 +217,11 @@ def get_indexes( ) table_name = self.resolve_identifier(table) - self._log_table_describe(driver, schema_name=schema_name, table_name=table_name, operation="indexes") + table_schema_name = None if schema is None else schema_name + self._log_table_describe(driver, schema_name=table_schema_name, table_name=table_name, operation="indexes") return driver.select( - _cockroach_domain_sql("indexes", "by_schema"), - schema_name=schema_name, + _cockroach_domain_sql("indexes", "by_table"), + schema_name=table_schema_name, table_name=table_name, schema_type=IndexMetadata, ) @@ -450,7 +436,7 @@ async def get_columns( table_schema_name = None if schema is None else schema_name self._log_table_describe(driver, schema_name=table_schema_name, table_name=table_name, operation="columns") return await driver.select( - _cockroach_domain_sql("columns", "by_schema"), + _cockroach_domain_sql("columns", "by_table"), schema_name=table_schema_name, table_name=table_name, schema_type=ColumnMetadata, @@ -471,10 +457,11 @@ async def get_indexes( ) table_name = self.resolve_identifier(table) - self._log_table_describe(driver, schema_name=schema_name, table_name=table_name, operation="indexes") + table_schema_name = None if schema is None else schema_name + self._log_table_describe(driver, schema_name=table_schema_name, table_name=table_name, operation="indexes") return await driver.select( - _cockroach_domain_sql("indexes", "by_schema"), - schema_name=schema_name, + _cockroach_domain_sql("indexes", "by_table"), + schema_name=table_schema_name, table_name=table_name, schema_type=IndexMetadata, ) @@ -562,7 +549,7 @@ def _cockroach_metadata_capability(domain: str) -> MetadataCapability: risks=(MetadataRisk.EXPENSIVE, MetadataRisk.PRIVILEGED), warnings=("System metadata is opt-in and disabled by default.",), ) - if domain in _COCKROACH_SUPPORTED_DOMAINS: + if domain in COCKROACH_SUPPORTED_DOMAINS: return MetadataCapability( domain=domain, support=MetadataSupport.SUPPORTED, @@ -573,7 +560,7 @@ def _cockroach_metadata_capability(domain: str) -> MetadataCapability: def _cockroach_metadata_profile(adapter: str, domains: Sequence[str] | None) -> MetadataCapabilityProfile: - requested_domains = _COCKROACH_METADATA_DOMAINS if domains is None else tuple(domains) + requested_domains = COCKROACH_METADATA_DOMAINS if domains is None else tuple(domains) return MetadataCapabilityProfile( "cockroachdb", adapter=adapter, diff --git a/tests/unit/adapters/test_cockroach_psycopg/test_config.py b/tests/unit/adapters/test_cockroach_psycopg/test_config.py index 50f034f76..68bad8061 100644 --- a/tests/unit/adapters/test_cockroach_psycopg/test_config.py +++ b/tests/unit/adapters/test_cockroach_psycopg/test_config.py @@ -1071,3 +1071,55 @@ async def __aexit__(self, *_a: Any) -> None: assert count == 2 assert len(executed_queries) == 1 assert "unnest" in executed_queries[0].lower() + + +def test_cockroach_psycopg_sync_get_columns_and_indexes_by_table() -> None: + """CockroachPsycopgSyncDataDictionary executes by_table query when table is specified.""" + from sqlspec.adapters.cockroach_psycopg.data_dictionary import CockroachPsycopgSyncDataDictionary + + calls: list[tuple[Any, dict[str, Any]]] = [] + driver = MagicMock() + + def fake_select(query: Any, **kwargs: Any) -> list[Any]: + calls.append((query, kwargs)) + return [] + + driver.select = fake_select + dd = CockroachPsycopgSyncDataDictionary() + + dd.get_columns(driver, table="users", schema=None) + assert len(calls) == 1 + assert calls[0][1]["table_name"] == "users" + assert calls[0][1]["schema_name"] is None + + calls.clear() + dd.get_indexes(driver, table="users", schema=None) + assert len(calls) == 1 + assert calls[0][1]["table_name"] == "users" + assert calls[0][1]["schema_name"] is None + + +async def test_cockroach_psycopg_async_get_columns_and_indexes_by_table() -> None: + """CockroachPsycopgAsyncDataDictionary executes by_table query when table is specified.""" + from sqlspec.adapters.cockroach_psycopg.data_dictionary import CockroachPsycopgAsyncDataDictionary + + calls: list[tuple[Any, dict[str, Any]]] = [] + driver = MagicMock() + + async def fake_select(query: Any, **kwargs: Any) -> list[Any]: + calls.append((query, kwargs)) + return [] + + driver.select = AsyncMock(side_effect=fake_select) + dd = CockroachPsycopgAsyncDataDictionary() + + await dd.get_columns(driver, table="users", schema=None) + assert len(calls) == 1 + assert calls[0][1]["table_name"] == "users" + assert calls[0][1]["schema_name"] is None + + calls.clear() + await dd.get_indexes(driver, table="users", schema=None) + assert len(calls) == 1 + assert calls[0][1]["table_name"] == "users" + assert calls[0][1]["schema_name"] is None From ed383a702e95a997320213e495e9cf112b283e5f Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 01:05:47 +0000 Subject: [PATCH 084/174] perf(cockroach_psycopg): cache validated default_staleness and optimize follower read gates Validate default_staleness once during driver initialization, cache effective staleness and follower_reads_enabled on driver slots, and eliminate repeated dict lookups and regex validations across query executions and cache hits. --- sqlspec/adapters/cockroach_psycopg/driver.py | 48 +++++++++++--------- 1 file changed, 26 insertions(+), 22 deletions(-) diff --git a/sqlspec/adapters/cockroach_psycopg/driver.py b/sqlspec/adapters/cockroach_psycopg/driver.py index eda4d8ff7..dfc6b0ab5 100644 --- a/sqlspec/adapters/cockroach_psycopg/driver.py +++ b/sqlspec/adapters/cockroach_psycopg/driver.py @@ -154,7 +154,7 @@ class CockroachPsycopgAsyncExceptionHandler(PsycopgAsyncExceptionHandler): class CockroachPsycopgSyncDriver(PsycopgSyncDriver): """CockroachDB sync driver using psycopg.crdb.""" - __slots__ = ("_enable_retry", "_follower_staleness", "_retry_config") + __slots__ = ("_enable_retry", "_follower_reads_enabled", "_follower_staleness", "_retry_config") dialect = "postgres" def __init__( @@ -172,7 +172,9 @@ def __init__( self._retry_config = CockroachPsycopgRetryConfig.from_features(self.driver_features) self._enable_retry = bool(self.driver_features.get("enable_auto_retry", True)) - self._follower_staleness = cast("str | None", self.driver_features.get("default_staleness")) + self._follower_reads_enabled = bool(self.driver_features.get("enable_follower_reads", False)) + raw_staleness = cast("str | None", self.driver_features.get("default_staleness")) + self._follower_staleness = validate_follower_read_staleness(raw_staleness) if raw_staleness else None self._data_dictionary = None @property @@ -367,7 +369,11 @@ def _execute_cache_hit( ) -> "SQLResult": auto_started = ( self._begin_follower_read_transaction() - if (cached.operation_profile.returns_rows and not cached.operation_profile.modifies_rows) + if ( + self._follower_reads_enabled + and cached.operation_profile.returns_rows + and not cached.operation_profile.modifies_rows + ) else False ) try: @@ -400,12 +406,9 @@ def _execute_native_storage(self, command: str, parameters: "list[Any]") -> "lis return rows def _apply_follower_reads(self) -> None: - if not self.driver_features.get("enable_follower_reads", False): - return - if not self._follower_staleness: + if not self._follower_reads_enabled or not self._follower_staleness: return - staleness = validate_follower_read_staleness(self._follower_staleness) - self.connection.execute(cast("Any", f"SET TRANSACTION AS OF SYSTEM TIME {staleness}")).close() + self.connection.execute(cast("Any", f"SET TRANSACTION AS OF SYSTEM TIME {self._follower_staleness}")).close() def _begin_follower_read_transaction(self) -> bool: """Open the transaction a follower read needs so the staleness clause can lead it. @@ -414,9 +417,7 @@ def _begin_follower_read_transaction(self) -> bool: clause with nowhere to go, so a read opens one here when the caller has not already done so. """ - if not self.driver_features.get("enable_follower_reads", False): - return False - if not self._follower_staleness: + if not self._follower_reads_enabled or not self._follower_staleness: return False if self._connection_in_transaction(): return False @@ -427,7 +428,7 @@ def _begin_follower_read_transaction(self) -> bool: class CockroachPsycopgAsyncDriver(PsycopgAsyncDriver): """CockroachDB async driver using psycopg.crdb.""" - __slots__ = ("_enable_retry", "_follower_staleness", "_retry_config") + __slots__ = ("_enable_retry", "_follower_reads_enabled", "_follower_staleness", "_retry_config") dialect = "postgres" def __init__( @@ -445,7 +446,9 @@ def __init__( self._retry_config = CockroachPsycopgRetryConfig.from_features(self.driver_features) self._enable_retry = bool(self.driver_features.get("enable_auto_retry", True)) - self._follower_staleness = cast("str | None", self.driver_features.get("default_staleness")) + self._follower_reads_enabled = bool(self.driver_features.get("enable_follower_reads", False)) + raw_staleness = cast("str | None", self.driver_features.get("default_staleness")) + self._follower_staleness = validate_follower_read_staleness(raw_staleness) if raw_staleness else None self._data_dictionary = None @property @@ -629,7 +632,11 @@ async def _execute_cache_hit( ) -> "SQLResult": auto_started = ( await self._begin_follower_read_transaction() - if (cached.operation_profile.returns_rows and not cached.operation_profile.modifies_rows) + if ( + self._follower_reads_enabled + and cached.operation_profile.returns_rows + and not cached.operation_profile.modifies_rows + ) else False ) try: @@ -662,12 +669,11 @@ async def _execute_native_storage(self, command: str, parameters: "list[Any]") - return rows async def _apply_follower_reads(self) -> None: - if not self.driver_features.get("enable_follower_reads", False): - return - if not self._follower_staleness: + if not self._follower_reads_enabled or not self._follower_staleness: return - staleness = validate_follower_read_staleness(self._follower_staleness) - cursor = await self.connection.execute(cast("Any", f"SET TRANSACTION AS OF SYSTEM TIME {staleness}")) + cursor = await self.connection.execute( + cast("Any", f"SET TRANSACTION AS OF SYSTEM TIME {self._follower_staleness}") + ) await cursor.close() async def _begin_follower_read_transaction(self) -> bool: @@ -677,9 +683,7 @@ async def _begin_follower_read_transaction(self) -> bool: clause with nowhere to go, so a read opens one here when the caller has not already done so. """ - if not self.driver_features.get("enable_follower_reads", False): - return False - if not self._follower_staleness: + if not self._follower_reads_enabled or not self._follower_staleness: return False if self._connection_in_transaction(): return False From d220404426f3665cf23ee7fa234edd00af65464f Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 01:05:47 +0000 Subject: [PATCH 085/174] refactor(cockroach_psycopg): code hygiene, connection typing, and native storage exception guard Wrap _execute_native_storage in ambient exception guard for sync and async drivers, document auto-retry semantics, parameterize CrdbConnection aliases to Any to avoid false dict_row type assumptions, and remove redundant casts. --- sqlspec/adapters/cockroach_psycopg/_typing.py | 8 ++-- sqlspec/adapters/cockroach_psycopg/config.py | 8 ++-- sqlspec/adapters/cockroach_psycopg/driver.py | 43 +++++++++++++------ .../test_cockroach_psycopg/test_config.py | 36 ++++++++++++++++ 4 files changed, 73 insertions(+), 22 deletions(-) diff --git a/sqlspec/adapters/cockroach_psycopg/_typing.py b/sqlspec/adapters/cockroach_psycopg/_typing.py index f9db5c984..933503c58 100644 --- a/sqlspec/adapters/cockroach_psycopg/_typing.py +++ b/sqlspec/adapters/cockroach_psycopg/_typing.py @@ -32,10 +32,10 @@ from sqlspec.adapters.cockroach_psycopg.driver import CockroachPsycopgAsyncDriver, CockroachPsycopgSyncDriver from sqlspec.core import StatementConfig - CockroachSyncConnection: TypeAlias = CrdbConnection[PsycopgDictRow] - CockroachAsyncConnection: TypeAlias = AsyncCrdbConnection[PsycopgDictRow] - CockroachSyncCursor: TypeAlias = Cursor[PsycopgDictRow] - CockroachAsyncCursor: TypeAlias = AsyncCursor[PsycopgDictRow] + CockroachSyncConnection: TypeAlias = CrdbConnection[Any] + CockroachAsyncConnection: TypeAlias = AsyncCrdbConnection[Any] + CockroachSyncCursor: TypeAlias = Cursor[Any] + CockroachAsyncCursor: TypeAlias = AsyncCursor[Any] else: CockroachSyncConnection = cockroach_psycopg_crdb.CrdbConnection CockroachAsyncConnection = cockroach_psycopg_crdb.AsyncCrdbConnection diff --git a/sqlspec/adapters/cockroach_psycopg/config.py b/sqlspec/adapters/cockroach_psycopg/config.py index 6072de71c..26296caac 100644 --- a/sqlspec/adapters/cockroach_psycopg/config.py +++ b/sqlspec/adapters/cockroach_psycopg/config.py @@ -305,12 +305,12 @@ def create_connection(self) -> "CockroachSyncConnection": conninfo, connection_kwargs = _standalone_connection_kwargs(dict(self.connection_config)) connection = psycopg_crdb.CrdbConnection.connect(conninfo, **connection_kwargs) try: - self._configure_connection(cast("CockroachSyncConnection", connection)) + self._configure_connection(connection) except BaseException: with suppress(Exception): connection.close() raise - return cast("CockroachSyncConnection", connection) + return connection def provide_session( self, @@ -549,12 +549,12 @@ async def create_connection(self) -> "CockroachAsyncConnection": conninfo, connection_kwargs = _standalone_connection_kwargs(dict(self.connection_config)) connection = await psycopg_crdb.AsyncCrdbConnection.connect(conninfo, **connection_kwargs) try: - await self._configure_async_connection(cast("CockroachAsyncConnection", connection)) + await self._configure_async_connection(connection) except BaseException: with suppress(Exception): await connection.close() raise - return cast("CockroachAsyncConnection", connection) + return connection def provide_session( self, diff --git a/sqlspec/adapters/cockroach_psycopg/driver.py b/sqlspec/adapters/cockroach_psycopg/driver.py index dfc6b0ab5..576c72f04 100644 --- a/sqlspec/adapters/cockroach_psycopg/driver.py +++ b/sqlspec/adapters/cockroach_psycopg/driver.py @@ -1,6 +1,7 @@ """CockroachDB psycopg driver implementation.""" import asyncio +import sys import time from contextlib import suppress from typing import TYPE_CHECKING, Any, TypeVar, cast @@ -396,13 +397,20 @@ def _native_storage_ready(self) -> bool: def _execute_native_storage(self, command: str, parameters: "list[Any]") -> "list[dict[str, Any]]": handler = self.handle_database_exceptions() - rows = [] - with self.with_cursor(self.connection) as cursor, handler: - cursor.row_factory = dict_row - cursor.execute(command.encode("utf-8"), parameters) - rows = cursor.fetchall() - if handler.pending_exception is not None: - raise handler.pending_exception + + def _fetch() -> list[dict[str, Any]]: + with self.with_cursor(self.connection) as cursor: + cursor.row_factory = dict_row + cursor.execute(command.encode("utf-8"), parameters) + return cursor.fetchall() + + if sys.exc_info()[1] is None: + with handler: + rows = _fetch() + else: + res = self._run_with_exception_handler(handler, _fetch) + rows = res if res is not None else [] + self._check_pending_exception(handler) return rows def _apply_follower_reads(self) -> None: @@ -659,13 +667,20 @@ def _native_storage_ready(self) -> bool: async def _execute_native_storage(self, command: str, parameters: "list[Any]") -> "list[dict[str, Any]]": handler = self.handle_database_exceptions() - rows = [] - async with self.with_cursor(self.connection) as cursor, handler: - cursor.row_factory = dict_row - await cursor.execute(command.encode("utf-8"), parameters) - rows = await cursor.fetchall() - if handler.pending_exception is not None: - raise handler.pending_exception + + async def _fetch() -> list[dict[str, Any]]: + async with self.with_cursor(self.connection) as cursor: + cursor.row_factory = dict_row + await cursor.execute(command.encode("utf-8"), parameters) + return await cursor.fetchall() + + if sys.exc_info()[1] is None: + async with handler: + rows = await _fetch() + else: + res = await self._run_with_exception_handler(handler, _fetch) + rows = res if res is not None else [] + self._check_pending_exception(handler) return rows async def _apply_follower_reads(self) -> None: diff --git a/tests/unit/adapters/test_cockroach_psycopg/test_config.py b/tests/unit/adapters/test_cockroach_psycopg/test_config.py index 68bad8061..7a17d647c 100644 --- a/tests/unit/adapters/test_cockroach_psycopg/test_config.py +++ b/tests/unit/adapters/test_cockroach_psycopg/test_config.py @@ -15,6 +15,7 @@ from sqlspec.adapters.cockroach_psycopg import ( CockroachPsycopgAsyncConfig, + CockroachPsycopgAsyncDriver, CockroachPsycopgAsyncSessionContext, CockroachPsycopgDriverFeatures, CockroachPsycopgRetryConfig, @@ -1123,3 +1124,38 @@ async def fake_select(query: Any, **kwargs: Any) -> list[Any]: assert len(calls) == 1 assert calls[0][1]["table_name"] == "users" assert calls[0][1]["schema_name"] is None + + +def test_cockroach_psycopg_sync_execute_native_storage_under_active_exception() -> None: + """_execute_native_storage in sync driver inside active except block works cleanly.""" + conn = MagicMock() + conn.autocommit = True + cursor = MagicMock() + cursor.fetchall.return_value = [{"rows": 1}] + conn.cursor.return_value = cursor + + driver = CockroachPsycopgSyncDriver(connection=conn) + try: + raise ValueError("outer error") + except ValueError: + rows = driver._execute_native_storage("SELECT 1", []) + assert rows == [{"rows": 1}] + + +async def test_cockroach_psycopg_async_execute_native_storage_under_active_exception() -> None: + """_execute_native_storage in async driver inside active except block works cleanly.""" + conn = MagicMock() + conn.autocommit = True + cursor = MagicMock() + cursor.fetchall = AsyncMock(return_value=[{"rows": 1}]) + cursor.execute = AsyncMock() + cursor.__aenter__ = AsyncMock(return_value=cursor) + cursor.__aexit__ = AsyncMock() + conn.cursor.return_value = cursor + + driver = CockroachPsycopgAsyncDriver(connection=conn) + try: + raise ValueError("outer error") + except ValueError: + rows = await driver._execute_native_storage("SELECT 1", []) + assert rows == [{"rows": 1}] From 1c78cf0754ed0abbfcb4f4c3891540e4317f0cc7 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 01:05:47 +0000 Subject: [PATCH 086/174] fix(psqlpy): correct exception mapping, parameter binding, and transaction lifecycle --- sqlspec/adapters/psqlpy/_typing.py | 3 + sqlspec/adapters/psqlpy/config.py | 4 +- sqlspec/adapters/psqlpy/core.py | 64 ++++++++++++------- sqlspec/adapters/psqlpy/driver.py | 53 ++++++++++----- .../unit/adapters/test_psqlpy/test_config.py | 39 +++++++++++ tests/unit/adapters/test_psqlpy/test_core.py | 9 +++ .../unit/adapters/test_psqlpy/test_driver.py | 56 ++++++++++++++++ 7 files changed, 189 insertions(+), 39 deletions(-) diff --git a/sqlspec/adapters/psqlpy/_typing.py b/sqlspec/adapters/psqlpy/_typing.py index 7be4219bf..a04e7f25b 100644 --- a/sqlspec/adapters/psqlpy/_typing.py +++ b/sqlspec/adapters/psqlpy/_typing.py @@ -44,6 +44,7 @@ class _PsqlpyUnavailableError(Exception): PsqlpyListener: TypeAlias = _PsqlpyListener PsqlpyNotSupportedError: TypeAlias = _PsqlpyNotSupportedError PsqlpyOperationalError: TypeAlias = _PsqlpyOperationalError + RustPSQLDriverPyBaseError: TypeAlias = Any else: @@ -60,6 +61,7 @@ class _PsqlpyUnavailableError(Exception): PsqlpyListener = import_optional_attr("psqlpy", "Listener") or Any PsqlpyNotSupportedError = import_optional_attr("psqlpy.exceptions", "NotSupportedError") or _PsqlpyUnavailableError PsqlpyOperationalError = import_optional_attr("psqlpy.exceptions", "OperationalError") or _PsqlpyUnavailableError + RustPSQLDriverPyBaseError = import_optional_attr("psqlpy.exceptions", "RustPSQLDriverPyBaseError") or PsqlpyError __all__ = ( @@ -76,6 +78,7 @@ class _PsqlpyUnavailableError(Exception): "PsqlpyNotSupportedError", "PsqlpyOperationalError", "PsqlpySessionContext", + "RustPSQLDriverPyBaseError", ) diff --git a/sqlspec/adapters/psqlpy/config.py b/sqlspec/adapters/psqlpy/config.py index 540f093bb..d59c166c3 100644 --- a/sqlspec/adapters/psqlpy/config.py +++ b/sqlspec/adapters/psqlpy/config.py @@ -456,7 +456,9 @@ async def _ensure_connection(self, connection: "PsqlpyConnection") -> None: conn_id = id(connection) if conn_id not in self._initialized_connection_ids: if self._user_connection_hook is not None: - await self._user_connection_hook(connection) + hook_res = self._user_connection_hook(connection) + if inspect.isawaitable(hook_res): + await hook_res if len(self._initialized_connection_ids) >= _MAX_INITIALIZED_CONNECTION_IDS: self._initialized_connection_ids.clear() self._initialized_connection_ids.add(conn_id) diff --git a/sqlspec/adapters/psqlpy/core.py b/sqlspec/adapters/psqlpy/core.py index 0101439a4..ee9bbb0b4 100644 --- a/sqlspec/adapters/psqlpy/core.py +++ b/sqlspec/adapters/psqlpy/core.py @@ -3,7 +3,6 @@ import contextlib import datetime as dt import decimal -import inspect import re import uuid from functools import lru_cache @@ -14,7 +13,12 @@ from sqlglot.errors import ParseError from sqlspec.adapters.psqlpy._typing import PSQLPY_JSONB as _JSONB_TYPE -from sqlspec.adapters.psqlpy._typing import PsqlpyDataError, PsqlpyIntegrityError, PsqlpyOperationalError +from sqlspec.adapters.psqlpy._typing import ( + PsqlpyConnectionExecuteError, + PsqlpyDataError, + PsqlpyIntegrityError, + PsqlpyOperationalError, +) from sqlspec.adapters.psqlpy.type_converter import coerce_pgvector from sqlspec.core import ( DriverParameterProfile, @@ -169,13 +173,7 @@ def build_connection_config(connection_config: "Mapping[str, Any]") -> "dict[str username = config.pop("username", None) or config.pop("user", None) if username is not None: config["username"] = username - auth_token_provider = config.get("auth_token_provider") - if callable(auth_token_provider) and not inspect.iscoroutinefunction(auth_token_provider): - config.pop("auth_token_provider", None) - token_result = auth_token_provider() - if not inspect.isawaitable(token_result): - config["password"] = str(token_result) - elif auth_token_provider is None: + if "auth_token_provider" in config and config["auth_token_provider"] is None: config.pop("auth_token_provider", None) return config @@ -478,14 +476,28 @@ def create_mapped_exception(error: Any, *, logger: Any | None = None) -> SQLSpec Returns: A SQLSpec exception that wraps the original error """ - del logger + _ = logger error_msg = str(error).lower() - if "unique" in error_msg or "duplicate key" in error_msg: + if isinstance(error, PsqlpyDataError): + return _create_postgres_error(error, DataError, "data error") + + if "syntax error" in error_msg or "parse error" in error_msg: + return _create_postgres_error(error, SQLParsingError, "SQL syntax error") + + if ( + "unique constraint" in error_msg + or "duplicate key value" in error_msg + or "violates unique constraint" in error_msg + ): return _create_postgres_error(error, UniqueViolationError, "unique constraint violation") - if "foreign key" in error_msg or "violates foreign key" in error_msg: + if "foreign key constraint" in error_msg or "violates foreign key constraint" in error_msg: return _create_postgres_error(error, ForeignKeyViolationError, "foreign key constraint violation") - if "not null" in error_msg or ("null value" in error_msg and "violates not-null" in error_msg): + if ( + "not-null constraint" in error_msg + or "violates not-null constraint" in error_msg + or ("null value in column" in error_msg and "violates not-null" in error_msg) + ): return _create_postgres_error(error, NotNullViolationError, "not-null constraint violation") if "check constraint" in error_msg or "violates check constraint" in error_msg: return _create_postgres_error(error, CheckViolationError, "check constraint violation") @@ -497,13 +509,18 @@ def create_mapped_exception(error: Any, *, logger: Any | None = None) -> SQLSpec if "serialization failure" in error_msg or "could not serialize" in error_msg: return _create_postgres_error(error, SerializationConflictError, "serialization failure") + if "connection" in error_msg: + if "timeout" in error_msg or "timed out" in error_msg: + return _create_postgres_error(error, ConnectionTimeoutError, "connection timeout") + return _create_postgres_error(error, DatabaseConnectionError, "connection error") + if ( - "connection" in error_msg - or "could not connect" in error_msg - or "connection refused" in error_msg + "connection refused" in error_msg or "connection reset" in error_msg - or "broken pipe" in error_msg + or "could not connect to server" in error_msg or "server closed the connection" in error_msg + or "broken pipe" in error_msg + or "failed to connect" in error_msg ): if "timeout" in error_msg or "timed out" in error_msg: return _create_postgres_error(error, ConnectionTimeoutError, "connection timeout") @@ -521,16 +538,17 @@ def create_mapped_exception(error: Any, *, logger: Any | None = None) -> SQLSpec ): return _create_postgres_error(error, PermissionDeniedError, "authentication error") - if isinstance(error, PsqlpyDataError): - return _create_postgres_error(error, DataError, "data error") - - if "syntax error" in error_msg or "parse" in error_msg: - return _create_postgres_error(error, SQLParsingError, "SQL syntax error") - if isinstance(error, PsqlpyOperationalError): return _create_postgres_error(error, OperationalError, "operational error") if isinstance(error, PsqlpyIntegrityError): return _create_postgres_error(error, IntegrityError, "integrity constraint violation") + if isinstance(error, (PsqlpyConnectionExecuteError,)): + return _create_postgres_error(error, DatabaseConnectionError, "connection execute error") + + if "connection" in error_msg: + if "timeout" in error_msg or "timed out" in error_msg: + return _create_postgres_error(error, ConnectionTimeoutError, "connection timeout") + return _create_postgres_error(error, DatabaseConnectionError, "connection error") return _create_postgres_error(error, SQLSpecError, "database error") diff --git a/sqlspec/adapters/psqlpy/driver.py b/sqlspec/adapters/psqlpy/driver.py index 8abc1f355..7fa4c0020 100644 --- a/sqlspec/adapters/psqlpy/driver.py +++ b/sqlspec/adapters/psqlpy/driver.py @@ -4,7 +4,9 @@ and transaction management. """ +from contextlib import suppress from itertools import chain +from time import perf_counter from typing import TYPE_CHECKING, Any, Literal, cast from sqlspec.adapters.psqlpy._typing import ( @@ -207,18 +209,39 @@ async def dispatch_execute_script(self, cursor: "PsqlpyConnection", statement: S statement_config = statement.statement_config statements = self.split_script_statements(sql, statement_config, strip_trailing_semicolon=True) + statement_count = len(statements) successful_count = 0 - last_result = None params = prepared_parameters or [] is_single = len(statements) == 1 + runtime = self._observability + has_observers = runtime is not None and runtime.has_statement_observers + use_prepared = not self._is_transaction_pooler_mode() - for stmt in statements: + for idx, stmt in enumerate(statements): stmt_params = params if (is_single or "$" in stmt) else [] - last_result = await cursor.execute(stmt, stmt_params) + stmt_started = perf_counter() if has_observers else 0.0 + self._pending_statement_profile = None + await cursor.execute(stmt, stmt_params, prepared=use_prepared) + if has_observers: + stmt_duration = perf_counter() - stmt_started + stmt_profile = self._pending_statement_profile + self._pending_statement_profile = None + self._emit_script_statement_event( + sql=stmt, + parameters=stmt_params, + operation="SCRIPT", + rows_affected=None, + duration_s=stmt_duration, + started_at=stmt_started, + profile=stmt_profile, + script_ordinal=idx, + script_statement_count=statement_count, + ) successful_count += 1 + self._pending_statement_profile = None return self.create_execution_result( - last_result, statement_count=len(statements), successful_statements=successful_count, is_script_result=True + cursor, statement_count=statement_count, successful_statements=successful_count, is_script_result=True ) async def begin(self) -> None: @@ -228,32 +251,32 @@ async def begin(self) -> None: try: await self.connection.execute("BEGIN") + except (PsqlpyDatabaseError, PsqlpyConnectionExecuteError, PsqlpyError, Exception) as e: + raise create_mapped_exception(e) from e + try: if self._should_set_local_columnar(): await self.connection.execute("SET LOCAL google_columnar_engine.enable_columnar_scan = 'on'") - except (PsqlpyDatabaseError, PsqlpyConnectionExecuteError, PsqlpyError) as e: - msg = f"Failed to begin psqlpy transaction: {e}" - raise SQLSpecError(msg) from e + except Exception as e: + with suppress(Exception): + await self.connection.execute("ROLLBACK") + raise create_mapped_exception(e) from e self._transaction_active = True async def commit(self) -> None: """Commit the current transaction.""" - try: await self.connection.execute("COMMIT") - except (PsqlpyDatabaseError, PsqlpyConnectionExecuteError, PsqlpyError) as e: - msg = f"Failed to commit psqlpy transaction: {e}" - raise SQLSpecError(msg) from e + except (PsqlpyDatabaseError, PsqlpyConnectionExecuteError, PsqlpyError, Exception) as e: + raise create_mapped_exception(e) from e finally: self._transaction_active = False async def rollback(self) -> None: """Rollback the current transaction.""" - try: await self.connection.execute("ROLLBACK") - except (PsqlpyDatabaseError, PsqlpyConnectionExecuteError, PsqlpyError) as e: - msg = f"Failed to rollback psqlpy transaction: {e}" - raise SQLSpecError(msg) from e + except (PsqlpyDatabaseError, PsqlpyConnectionExecuteError, PsqlpyError, Exception) as e: + raise create_mapped_exception(e) from e finally: self._transaction_active = False diff --git a/tests/unit/adapters/test_psqlpy/test_config.py b/tests/unit/adapters/test_psqlpy/test_config.py index e3494c315..9a75f2129 100644 --- a/tests/unit/adapters/test_psqlpy/test_config.py +++ b/tests/unit/adapters/test_psqlpy/test_config.py @@ -375,3 +375,42 @@ async def test_psqlpy_derive_direct_config_and_columnar_engine() -> None: await driver.begin() conn.execute.assert_any_await("BEGIN") conn.execute.assert_any_await("SET LOCAL google_columnar_engine.enable_columnar_scan = 'on'") + + +@pytest.mark.anyio +async def test_psqlpy_auth_token_provider_evaluated_on_create_pool_not_init() -> None: + """Synchronous and async auth_token_provider is evaluated when creating a pool, not during __init__.""" + call_count = 0 + + def token_provider() -> str: + nonlocal call_count + call_count += 1 + return "dynamic_token" + + config = PsqlpyConfig( + connection_config={"dsn": "postgresql://localhost:5432/test", "auth_token_provider": token_provider} + ) + assert call_count == 0 + + pool = await config.create_pool() + assert call_count == 1 + pool.close() + + +@pytest.mark.anyio +async def test_psqlpy_sync_on_connection_create_callback() -> None: + """Synchronous callback for on_connection_create should execute without awaiting.""" + hook_called = False + + def sync_hook(conn: Any) -> None: + nonlocal hook_called + hook_called = True + + config = PsqlpyConfig( + connection_config={"dsn": "postgresql://localhost:5432/test"}, + driver_features={"on_connection_create": sync_hook}, + ) + mock_conn = AsyncMock() + mock_conn.fetch = AsyncMock(return_value=None) + await config._ensure_connection(mock_conn) # pyright: ignore[reportPrivateUsage] + assert hook_called is True diff --git a/tests/unit/adapters/test_psqlpy/test_core.py b/tests/unit/adapters/test_psqlpy/test_core.py index d8e03eaeb..80d44a07f 100644 --- a/tests/unit/adapters/test_psqlpy/test_core.py +++ b/tests/unit/adapters/test_psqlpy/test_core.py @@ -67,6 +67,15 @@ def test_create_mapped_exception_leaves_not_supported_error_generic() -> None: assert type(mapped) is SQLSpecError +def test_create_mapped_exception_syntax_error_containing_connection_maps_to_sql_parsing_error() -> None: + """Syntax error containing the substring 'connection' maps to SQLParsingError, not DatabaseConnectionError.""" + from sqlspec.exceptions import SQLParsingError + + error = psqlpy_exceptions.Error('syntax error at or near "connection"') + mapped = psqlpy_core.create_mapped_exception(error) + assert type(mapped) is SQLParsingError + + def test_format_execute_many_parameters_no_coercion_reuses_list_rows() -> None: """Formatting should preserve list rows when no numeric coercion is requested.""" records = [[1, "a"], [2, "b"]] diff --git a/tests/unit/adapters/test_psqlpy/test_driver.py b/tests/unit/adapters/test_psqlpy/test_driver.py index cd7b47b85..bedd66489 100644 --- a/tests/unit/adapters/test_psqlpy/test_driver.py +++ b/tests/unit/adapters/test_psqlpy/test_driver.py @@ -9,6 +9,7 @@ from sqlspec.adapters.psqlpy import core as psqlpy_core from sqlspec.adapters.psqlpy.core import build_statement_config from sqlspec.adapters.psqlpy.driver import PsqlpyDriver +from sqlspec.core import SQL from sqlspec.core.sqlcommenter import SQLCommenterContext @@ -50,3 +51,58 @@ def _fake_fetch(query: str, _params: Any) -> _FakeResult: assert len(fetched_queries) == 5 assert "0000000000000001" in fetched_queries[0] assert "0000000000000005" in fetched_queries[4] + + +@pytest.mark.anyio +async def test_psqlpy_dispatch_execute_script_pooler_mode_disables_prepared() -> None: + """dispatch_execute_script in pooler mode passes prepared=False and emits script observer events.""" + from sqlspec.observability import ObservabilityConfig, ObservabilityRuntime + + cursor = AsyncMock() + cursor.execute = AsyncMock() + driver = PsqlpyDriver( + connection=cast("Any", cursor), + statement_config=build_statement_config(), + driver_features={"pooler_mode": "transaction"}, + ) + observed_events: list[Any] = [] + obs_config = ObservabilityConfig(statement_observers=(observed_events.append,)) + driver._observability = ObservabilityRuntime(obs_config) # pyright: ignore[reportPrivateUsage] + + stmt = SQL("SELECT 1; SELECT 2", is_script=True, statement_config=build_statement_config()) + result = await driver.dispatch_execute_script(cast("Any", cursor), stmt) + + assert result.statement_count == 2 + assert result.successful_statements == 2 + assert cursor.execute.await_count == 2 + for call_args in cursor.execute.await_args_list: + assert call_args.kwargs.get("prepared") is False + + assert len(observed_events) == 2 + assert observed_events[0].operation == "SCRIPT" + assert observed_events[0].script_statement_count == 2 + + +@pytest.mark.anyio +async def test_psqlpy_transaction_methods_raise_mapped_exceptions() -> None: + """begin, commit, and rollback wrap driver errors in mapped exceptions.""" + from psqlpy.exceptions import OperationalError as PsqlpyOperationalError + + from sqlspec.exceptions import DatabaseConnectionError + + conn = AsyncMock() + conn.in_transaction = False + conn.execute = AsyncMock(side_effect=PsqlpyOperationalError("server closed the connection unexpectedly")) + + driver = PsqlpyDriver(connection=cast("Any", conn), statement_config=build_statement_config()) + + with pytest.raises(DatabaseConnectionError): + await driver.begin() + + driver._transaction_active = True # pyright: ignore[reportPrivateUsage] + with pytest.raises(DatabaseConnectionError): + await driver.commit() + + driver._transaction_active = True # pyright: ignore[reportPrivateUsage] + with pytest.raises(DatabaseConnectionError): + await driver.rollback() From eb0eac844c583d9fe6bb5001dc5fe81072ad8a53 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 01:05:47 +0000 Subject: [PATCH 087/174] fix(psqlpy): remediate ADK, Litestar, and events stores --- sqlspec/adapters/psqlpy/adk/store.py | 154 +++++++++--------- sqlspec/adapters/psqlpy/core.py | 34 ++++ sqlspec/adapters/psqlpy/events/_hub.py | 3 +- sqlspec/adapters/psqlpy/events/backend.py | 61 ++++--- sqlspec/adapters/psqlpy/events/store.py | 26 +-- sqlspec/adapters/psqlpy/litestar/store.py | 41 ++--- .../adapters/test_psqlpy/test_adk_store.py | 4 +- .../test_psqlpy/test_transaction_state.py | 89 +++++++++- 8 files changed, 251 insertions(+), 161 deletions(-) diff --git a/sqlspec/adapters/psqlpy/adk/store.py b/sqlspec/adapters/psqlpy/adk/store.py index de64aa015..b7ca54983 100644 --- a/sqlspec/adapters/psqlpy/adk/store.py +++ b/sqlspec/adapters/psqlpy/adk/store.py @@ -5,6 +5,7 @@ from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, NotRequired, cast from sqlspec.adapters.psqlpy._typing import PsqlpyConnectionExecuteError, PsqlpyDatabaseError +from sqlspec.adapters.psqlpy.core import build_postgres_storage_params_sql from sqlspec.config import ADKConfig from sqlspec.extensions.adk import BaseAsyncADKStore, StoredEvent, StoredSession, normalize_session_list_options from sqlspec.extensions.adk.memory.store import BaseAsyncADKMemoryStore @@ -553,7 +554,9 @@ async def _sessions_table_ddl(self) -> str: if self._owner_id_column_ddl: owner_id_line = f",\n {self._owner_id_column_ddl}" - table_options = _postgres_table_options(_adk_config(self._config)) + table_options = build_postgres_storage_params_sql( + _adk_config(self._config), default_fillfactor=80, section="adk" + ) return f""" CREATE TABLE IF NOT EXISTS {self._session_table} ( id VARCHAR(128) PRIMARY KEY, @@ -577,7 +580,9 @@ async def _sessions_table_ddl(self) -> str: async def _events_table_ddl(self) -> str: adk_config = _adk_config(self._config) - table_options = _postgres_table_options(adk_config, include_autovacuum=True) + table_options = build_postgres_storage_params_sql( + adk_config, include_autovacuum=True, default_fillfactor=80, section="adk" + ) generated_columns, generated_indexes, covering_columns = _postgres_event_ddl_options( adk_config, self._events_table ) @@ -603,7 +608,9 @@ async def _events_table_ddl(self) -> str: """ async def _app_states_table_ddl(self) -> str: - table_options = _postgres_table_options(_adk_config(self._config)) + table_options = build_postgres_storage_params_sql( + _adk_config(self._config), default_fillfactor=80, section="adk" + ) return f""" CREATE TABLE IF NOT EXISTS {self._app_state_table} ( app_name VARCHAR(128) PRIMARY KEY, @@ -613,7 +620,9 @@ async def _app_states_table_ddl(self) -> str: """ async def _user_states_table_ddl(self) -> str: - table_options = _postgres_table_options(_adk_config(self._config)) + table_options = build_postgres_storage_params_sql( + _adk_config(self._config), default_fillfactor=80, section="adk" + ) return f""" CREATE TABLE IF NOT EXISTS {self._user_state_table} ( app_name VARCHAR(128) NOT NULL, @@ -685,73 +694,83 @@ async def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: " if not entries: return 0 - inserted_count = 0 + ids = [entry["id"] for entry in entries] + session_ids = [entry["session_id"] for entry in entries] + app_names = [entry["app_name"] for entry in entries] + user_ids = [entry["user_id"] for entry in entries] + scopes = [entry.get("scope", "user") for entry in entries] + event_ids = [entry["event_id"] for entry in entries] + authors = [entry.get("author") for entry in entries] + timestamps = [entry["timestamp"] for entry in entries] + content_jsons = [entry["content_json"] for entry in entries] + content_texts = [entry["content_text"] for entry in entries] + metadata_jsons = [entry.get("metadata_json") for entry in entries] + inserted_ats = [entry["inserted_at"] for entry in entries] + if self._owner_id_column_name: sql = f""" INSERT INTO {self._memory_table} ( id, session_id, app_name, user_id, scope, event_id, author, - {self._owner_id_column_name}, timestamp, embedding, content_json, + {self._owner_id_column_name}, timestamp, content_json, content_text, metadata_json, inserted_at - ) VALUES ( - $1, $2, $3, $4, $5, $6, $7, $8, $9, $10::float8[]::vector, $11, $12, $13, $14 + ) + SELECT * FROM unnest( + $1::varchar[], $2::varchar[], $3::varchar[], $4::varchar[], + $5::varchar[], $6::varchar[], $7::varchar[], $8::varchar[], + $9::timestamptz[], $10::jsonb[], $11::text[], $12::jsonb[], $13::timestamptz[] ) ON CONFLICT (event_id) DO NOTHING + RETURNING 1 """ + owner_ids = [owner_id for _ in entries] + params = [ + ids, + session_ids, + app_names, + user_ids, + scopes, + event_ids, + authors, + owner_ids, + timestamps, + content_jsons, + content_texts, + metadata_jsons, + inserted_ats, + ] else: sql = f""" INSERT INTO {self._memory_table} ( id, session_id, app_name, user_id, scope, event_id, author, - timestamp, embedding, content_json, content_text, metadata_json, inserted_at - ) VALUES ( - $1, $2, $3, $4, $5, $6, $7, $8, $9::float8[]::vector, $10, $11, $12, $13 + timestamp, content_json, content_text, metadata_json, inserted_at + ) + SELECT * FROM unnest( + $1::varchar[], $2::varchar[], $3::varchar[], $4::varchar[], + $5::varchar[], $6::varchar[], $7::varchar[], + $8::timestamptz[], $9::jsonb[], $10::text[], $11::jsonb[], $12::timestamptz[] ) ON CONFLICT (event_id) DO NOTHING + RETURNING 1 """ + params = [ + ids, + session_ids, + app_names, + user_ids, + scopes, + event_ids, + authors, + timestamps, + content_jsons, + content_texts, + metadata_jsons, + inserted_ats, + ] async with self._config.provide_connection() as conn: - for entry in entries: - raw_embedding = entry.get("embedding") - embedding_val = [float(v) for v in raw_embedding] if raw_embedding is not None else None - if self._owner_id_column_name: - params = [ - entry["id"], - entry["session_id"], - entry["app_name"], - entry["user_id"], - entry.get("scope", "user"), - entry["event_id"], - entry.get("author"), - owner_id, - entry["timestamp"], - embedding_val, - entry["content_json"], - entry["content_text"], - entry.get("metadata_json"), - entry["inserted_at"], - ] - else: - params = [ - entry["id"], - entry["session_id"], - entry["app_name"], - entry["user_id"], - entry.get("scope", "user"), - entry["event_id"], - entry.get("author"), - entry["timestamp"], - embedding_val, - entry["content_json"], - entry["content_text"], - entry.get("metadata_json"), - entry["inserted_at"], - ] - - result = await conn.execute(sql, params) - rows_affected = self._extract_rows_affected(result) - if rows_affected > 0: - inserted_count += rows_affected - - return inserted_count + result = await conn.fetch(sql, params) + rows = result.result() if result else [] + return len(rows) async def search_entries( self, @@ -1156,33 +1175,6 @@ def _adk_config(config: Any) -> PsqlpyADKConfig: return cast("PsqlpyADKConfig", adk_config) -def _postgres_table_options(adk_config: PsqlpyADKConfig, *, include_autovacuum: bool = False) -> str: - options = [_postgres_fillfactor_option(adk_config)] - if include_autovacuum: - options.extend(_postgres_autovacuum_options(adk_config)) - return f" WITH ({', '.join(options)})" - - -def _postgres_fillfactor_option(adk_config: PsqlpyADKConfig) -> str: - value = adk_config.get("fillfactor", 80) - if not isinstance(value, int) or isinstance(value, bool) or value not in range(10, 101): - msg = "extension_config['adk']['fillfactor'] must be an integer from 10 to 100" - raise ValueError(msg) - return f"fillfactor = {value}" - - -def _postgres_autovacuum_options(adk_config: PsqlpyADKConfig) -> "list[str]": - options: list[str] = [] - for key in ("autovacuum_vacuum_scale_factor", "autovacuum_analyze_scale_factor"): - value = adk_config.get(key) - if value is not None: - if not isinstance(value, (int, float)) or isinstance(value, bool) or not 0 <= float(value) <= 1: - msg = f"extension_config['adk']['{key}'] must be a number from 0 to 1" - raise ValueError(msg) - options.append(f"{key} = {float(value):g}") - return options - - def _postgres_memory_indexes( memory_table: str, use_fts: bool, diff --git a/sqlspec/adapters/psqlpy/core.py b/sqlspec/adapters/psqlpy/core.py index ee9bbb0b4..18ec28a3b 100644 --- a/sqlspec/adapters/psqlpy/core.py +++ b/sqlspec/adapters/psqlpy/core.py @@ -71,6 +71,7 @@ "apply_driver_features", "build_connection_config", "build_postgres_extension_probe_names", + "build_postgres_storage_params_sql", "build_profile", "build_statement_config", "coerce_json_columns", @@ -110,6 +111,39 @@ _NUMERIC_COERCE_TYPES: "tuple[type[Any], ...]" = (float, decimal.Decimal, list, tuple) +def build_postgres_storage_params_sql( + settings: "dict[str, object] | Mapping[str, object] | None", + *, + include_autovacuum: bool = False, + default_fillfactor: int | None = None, + default_vacuum_scale: float | None = None, + default_analyze_scale: float | None = None, + section: str = "storage", +) -> str: + """Format PostgreSQL table storage parameters (WITH fillfactor = ...) clause.""" + options: list[str] = [] + config_dict = dict(settings) if settings else {} + fillfactor = config_dict.get("fillfactor", default_fillfactor) + if fillfactor is not None: + if not isinstance(fillfactor, int) or isinstance(fillfactor, bool) or fillfactor not in range(10, 101): + msg = f"extension_config['{section}']['fillfactor'] must be an integer from 10 to 100" + raise ValueError(msg) + options.append(f"fillfactor = {fillfactor}") + if include_autovacuum: + scales = ( + ("autovacuum_vacuum_scale_factor", default_vacuum_scale), + ("autovacuum_analyze_scale_factor", default_analyze_scale), + ) + for key, default_val in scales: + value = config_dict.get(key, default_val) + if value is not None: + if not isinstance(value, (int, float)) or isinstance(value, bool) or not 0 <= float(value) <= 1: + msg = f"extension_config['{section}']['{key}'] must be a number from 0 to 1" + raise ValueError(msg) + options.append(f"{key} = {float(value):g}") + return f" WITH ({', '.join(options)})" if options else "" + + def build_profile() -> "DriverParameterProfile": """Create the psqlpy driver parameter profile.""" coercions: dict[type, Callable[[Any], Any]] = build_uuid_coercions(native=True) diff --git a/sqlspec/adapters/psqlpy/events/_hub.py b/sqlspec/adapters/psqlpy/events/_hub.py index 20da69998..1718f19ae 100644 --- a/sqlspec/adapters/psqlpy/events/_hub.py +++ b/sqlspec/adapters/psqlpy/events/_hub.py @@ -13,6 +13,7 @@ import asyncio import contextlib import logging +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, cast from weakref import WeakKeyDictionary @@ -246,7 +247,7 @@ def _resolve_listener_config(self) -> "PsqlpyConfig": ) ).lower() raw_direct_dsn = driver_features.get("direct_dsn") - direct_dsn = raw_direct_dsn if isinstance(raw_direct_dsn, str) and raw_direct_dsn else None + direct_dsn = raw_direct_dsn if isinstance(raw_direct_dsn, (str, Mapping)) and raw_direct_dsn else None if pooler_mode in {"transaction", "statement"} and not direct_dsn: logger.warning( "PostgreSQL LISTEN/NOTIFY cannot run under pooler_mode=%r without configuring driver_features['direct_dsn']", diff --git a/sqlspec/adapters/psqlpy/events/backend.py b/sqlspec/adapters/psqlpy/events/backend.py index cf49e89a4..2d305bc4e 100644 --- a/sqlspec/adapters/psqlpy/events/backend.py +++ b/sqlspec/adapters/psqlpy/events/backend.py @@ -28,7 +28,7 @@ __all__ = ("PsqlpyEventsBackend", "PsqlpyHybridEventsBackend", "create_event_backend") -logger = get_logger("sqlspec.events.psqlpy") +logger = get_logger("sqlspec.adapters.psqlpy.events.backend") _MIN_LISTENER_POOL_SIZE = 2 _MAX_SEEN_MARKERS = 1_024 @@ -241,14 +241,15 @@ async def dequeue(self, channel: str, poll_interval: float) -> EventMessage | No return await self._dequeue_reconciled(channel) self._queue._reset_empty_poll_delay(channel) self._runtime.increment_metric("events.listener.wakeup") - batch_size = _extract_batch_size(payload) - is_new, should_drain = self._marker_state.register(channel, _extract_marker_id(payload), batch_size) + parsed_payload = _parse_notify_payload(payload) + batch_size = _extract_batch_size(parsed_payload) + is_new, should_drain = self._marker_state.register(channel, _extract_marker_id(parsed_payload), batch_size) if is_new and batch_size > 1: self._runtime.increment_metric("events.marker.coalesced", batch_size - 1) if not should_drain: self._runtime.increment_metric("events.marker.duplicate") continue - event_id = _extract_event_id(payload) + event_id = _extract_event_id(parsed_payload) if event_id is not None: event = await self._queue.dequeue_by_event_id(event_id) if event is not None: @@ -354,6 +355,10 @@ def create_event_backend( def _validate_listener_pool_capacity(config: "PsqlpyConfig") -> None: + raw_features = config.driver_features + driver_features = raw_features if isinstance(raw_features, dict) else {} + if driver_features.get("direct_dsn"): + return max_size = config.connection_config.get("max_db_pool_size") if max_size is not None and max_size < _MIN_LISTENER_POOL_SIZE: msg = ( @@ -362,34 +367,40 @@ def _validate_listener_pool_capacity(config: "PsqlpyConfig") -> None: raise ImproperConfigurationError(msg) -def _extract_event_id(payload: "str | None") -> "str | None": - if not payload: +def _parse_notify_payload(payload: "dict[str, Any] | str | None") -> "dict[str, Any] | None": + if isinstance(payload, dict): + return payload + if not payload or not isinstance(payload, str): + return None + try: + raw = from_json(payload) + return raw if isinstance(raw, dict) else None + except (TypeError, ValueError): + return None + + +def _extract_event_id(payload: "dict[str, Any] | str | None") -> "str | None": + parsed = payload if isinstance(payload, dict) else _parse_notify_payload(payload) + if parsed is None: return None - raw = from_json(payload) - if isinstance(raw, dict): - event_id = raw.get("event_id") - return event_id if isinstance(event_id, str) else None - return None + event_id = parsed.get("event_id") + return event_id if isinstance(event_id, str) else None -def _extract_batch_size(payload: "str | None") -> int: - if not payload: +def _extract_batch_size(payload: "dict[str, Any] | str | None") -> int: + parsed = payload if isinstance(payload, dict) else _parse_notify_payload(payload) + if parsed is None: return 0 - raw = from_json(payload) - if isinstance(raw, dict): - batch_size = raw.get("batch_size") - return batch_size if isinstance(batch_size, int) else 0 - return 0 + batch_size = parsed.get("batch_size") + return batch_size if isinstance(batch_size, int) else 0 -def _extract_marker_id(payload: "str | None") -> "str | None": - if not payload: +def _extract_marker_id(payload: "dict[str, Any] | str | None") -> "str | None": + parsed = payload if isinstance(payload, dict) else _parse_notify_payload(payload) + if parsed is None: return None - raw = from_json(payload) - if isinstance(raw, dict): - marker_id = raw.get("marker_id", raw.get("event_id")) - return marker_id if isinstance(marker_id, str) else None - return None + marker_id = parsed.get("marker_id", parsed.get("event_id")) + return marker_id if isinstance(marker_id, str) else None def _record_dequeue_result(runtime: Any, event: "EventMessage | None") -> None: diff --git a/sqlspec/adapters/psqlpy/events/store.py b/sqlspec/adapters/psqlpy/events/store.py index db99cc248..5bdf5494f 100644 --- a/sqlspec/adapters/psqlpy/events/store.py +++ b/sqlspec/adapters/psqlpy/events/store.py @@ -3,6 +3,7 @@ from typing import NotRequired from sqlspec.adapters.psqlpy.config import PsqlpyConfig +from sqlspec.adapters.psqlpy.core import build_postgres_storage_params_sql from sqlspec.config import EventsConfig from sqlspec.extensions.events import BaseEventQueueStore @@ -42,27 +43,4 @@ def _column_types(self) -> "tuple[str, str, str]": def _table_clause(self) -> str: """Return explicitly configured PostgreSQL table tuning options.""" - return _postgres_table_options(self.settings) - - -def _postgres_table_options(settings: "dict[str, object]") -> str: - options = _postgres_option_values(settings) - return f" WITH ({', '.join(options)})" if options else "" - - -def _postgres_option_values(settings: "dict[str, object]") -> "list[str]": - options: list[str] = [] - fillfactor = settings.get("fillfactor") - if fillfactor is not None: - if not isinstance(fillfactor, int) or isinstance(fillfactor, bool) or fillfactor not in range(10, 101): - msg = "extension_config['events']['fillfactor'] must be an integer from 10 to 100" - raise ValueError(msg) - options.append(f"fillfactor = {fillfactor}") - for key in ("autovacuum_vacuum_scale_factor", "autovacuum_analyze_scale_factor"): - value = settings.get(key) - if value is not None: - if not isinstance(value, (int, float)) or isinstance(value, bool) or not 0 <= float(value) <= 1: - msg = f"extension_config['events']['{key}'] must be a number from 0 to 1" - raise ValueError(msg) - options.append(f"{key} = {float(value):g}") - return options + return build_postgres_storage_params_sql(self.settings, include_autovacuum=True, section="events") diff --git a/sqlspec/adapters/psqlpy/litestar/store.py b/sqlspec/adapters/psqlpy/litestar/store.py index c4b764fa0..8a1df7a26 100644 --- a/sqlspec/adapters/psqlpy/litestar/store.py +++ b/sqlspec/adapters/psqlpy/litestar/store.py @@ -1,8 +1,9 @@ """Psqlpy session store for Litestar integration.""" from datetime import UTC, datetime, timedelta -from typing import TYPE_CHECKING, Any, NotRequired, cast +from typing import TYPE_CHECKING, NotRequired +from sqlspec.adapters.psqlpy.core import build_postgres_storage_params_sql from sqlspec.config import LitestarConfig from sqlspec.extensions.litestar.store import BaseSQLSpecStore @@ -169,7 +170,7 @@ async def exists(self, key: str) -> bool: async with self._config.provide_connection() as conn: query_result = await conn.fetch(sql, [key]) - rows = query_result.result() + rows = query_result.result() if query_result else [] return len(rows) > 0 async def expires_in(self, key: str) -> "int | None": @@ -188,7 +189,7 @@ async def expires_in(self, key: str) -> "int | None": async with self._config.provide_connection() as conn: query_result = await conn.fetch(sql, [key]) - rows = query_result.result() + rows = query_result.result() if query_result else [] if not rows: return None @@ -219,7 +220,7 @@ async def delete_expired(self) -> int: async with self._config.provide_connection() as conn: query_result = await conn.fetch(sql, []) - rows = query_result.result() + rows = query_result.result() if query_result else [] count = len(rows) if count > 0: self._log_delete_expired(count) @@ -231,7 +232,14 @@ def _table_ddl(self) -> str: Returns: SQL statement to create the sessions table with proper indexes. """ - fillfactor, vacuum_scale, analyze_scale = _postgres_litestar_tuning(self._config) + tuning_clause = build_postgres_storage_params_sql( + self._config.extension_config.get("litestar", {}), + default_fillfactor=80, + default_vacuum_scale=0.05, + default_analyze_scale=0.02, + include_autovacuum=True, + section="litestar", + ) return f""" CREATE TABLE IF NOT EXISTS {self._table_name} ( session_id TEXT PRIMARY KEY, @@ -239,15 +247,10 @@ def _table_ddl(self) -> str: expires_at TIMESTAMPTZ, created_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP, updated_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP - ) WITH (fillfactor = {fillfactor}); + ){tuning_clause}; CREATE INDEX IF NOT EXISTS idx_{self._table_name}_expires_at ON {self._table_name}(expires_at) WHERE expires_at IS NOT NULL; - - ALTER TABLE {self._table_name} SET ( - autovacuum_vacuum_scale_factor = {vacuum_scale:g}, - autovacuum_analyze_scale_factor = {analyze_scale:g} - ); """ def _drop_table_sql(self) -> "list[str]": @@ -257,19 +260,3 @@ def _drop_table_sql(self) -> "list[str]": List of SQL statements to drop indexes and table. """ return [f"DROP INDEX IF EXISTS idx_{self._table_name}_expires_at", f"DROP TABLE IF EXISTS {self._table_name}"] - - -def _postgres_litestar_tuning(config: Any) -> "tuple[int, float, float]": - settings = cast("dict[str, Any]", config.extension_config.get("litestar", {})) - fillfactor = settings.get("fillfactor", 80) - if not isinstance(fillfactor, int) or isinstance(fillfactor, bool) or fillfactor not in range(10, 101): - msg = "extension_config['litestar']['fillfactor'] must be an integer from 10 to 100" - raise ValueError(msg) - scales: list[float] = [] - for key, default in (("autovacuum_vacuum_scale_factor", 0.05), ("autovacuum_analyze_scale_factor", 0.02)): - value = settings.get(key, default) - if not isinstance(value, (int, float)) or isinstance(value, bool) or not 0 <= float(value) <= 1: - msg = f"extension_config['litestar']['{key}'] must be a number from 0 to 1" - raise ValueError(msg) - scales.append(float(value)) - return fillfactor, scales[0], scales[1] diff --git a/tests/unit/adapters/test_psqlpy/test_adk_store.py b/tests/unit/adapters/test_psqlpy/test_adk_store.py index 0637c6411..df1da7f5d 100644 --- a/tests/unit/adapters/test_psqlpy/test_adk_store.py +++ b/tests/unit/adapters/test_psqlpy/test_adk_store.py @@ -298,8 +298,8 @@ async def test_psqlpy_memory_store_normalizes_tuple_embeddings_on_insert_and_sea results = await store.search_entries("", "app", "user-1", embedding=(0.1, 0.2)) assert inserted == 1 - assert conn.execute_calls[0][1][8] == [0.1, 0.2] - assert conn.fetch_calls[0][1][0] == [0.1, 0.2] + assert "unnest" in conn.fetch_calls[0][0].lower() + assert conn.fetch_calls[1][1][0] == [0.1, 0.2] assert results[0]["embedding"] == [0.1, 0.2] diff --git a/tests/unit/adapters/test_psqlpy/test_transaction_state.py b/tests/unit/adapters/test_psqlpy/test_transaction_state.py index 9c48ea87f..f7fc3872f 100644 --- a/tests/unit/adapters/test_psqlpy/test_transaction_state.py +++ b/tests/unit/adapters/test_psqlpy/test_transaction_state.py @@ -7,7 +7,7 @@ from sqlspec.adapters.psqlpy._typing import PsqlpyDatabaseError from sqlspec.adapters.psqlpy.config import PsqlpyConfig -from sqlspec.adapters.psqlpy.core import PsqlpyStreamSource +from sqlspec.adapters.psqlpy.core import PsqlpyStreamSource, build_postgres_storage_params_sql from sqlspec.adapters.psqlpy.driver import PsqlpyDriver from sqlspec.exceptions import SQLSpecError @@ -205,3 +205,90 @@ async def test_load_from_arrow_decodes_json_text_for_json_columns() -> None: _table_name, records, _kwargs = connection.copy_calls[0] assert records == [(1, {"name": "alpha"}, '{"not": "json"}')] + + +def test_psqlpy_build_postgres_storage_params_sql() -> None: + """build_postgres_storage_params_sql consolidates storage and autovacuum clauses.""" + settings = {"fillfactor": 85, "autovacuum_vacuum_scale_factor": 0.05, "autovacuum_analyze_scale_factor": 0.02} + clause = build_postgres_storage_params_sql(settings, include_autovacuum=True) + assert "WITH (" in clause + assert "fillfactor = 85" in clause + assert "autovacuum_vacuum_scale_factor = 0.05" in clause + assert "autovacuum_analyze_scale_factor = 0.02" in clause + + +def test_psqlpy_events_validate_listener_pool_capacity_skips_when_direct_dsn() -> None: + """_validate_listener_pool_capacity should not enforce pool size when direct_dsn is configured.""" + from sqlspec.adapters.psqlpy.events.backend import _validate_listener_pool_capacity + + config = PsqlpyConfig( + connection_config={"dsn": "postgresql://localhost/test", "max_db_pool_size": 1}, + driver_features={"direct_dsn": "postgresql://localhost/direct"}, + ) + _validate_listener_pool_capacity(config) + + +def test_psqlpy_events_parse_notify_payload_handles_malformed_json() -> None: + """Malformed NOTIFY payload does not raise and returns None safely.""" + from sqlspec.adapters.psqlpy.events.backend import ( + _extract_batch_size, + _extract_event_id, + _extract_marker_id, + _parse_notify_payload, + ) + + bad_payload = "not valid json {{" + assert _parse_notify_payload(bad_payload) is None + assert _extract_event_id(bad_payload) is None + assert _extract_batch_size(bad_payload) == 0 + assert _extract_marker_id(bad_payload) is None + + good_payload = '{"event_id": "e1", "batch_size": 5, "marker_id": "m1"}' + parsed = _parse_notify_payload(good_payload) + assert parsed is not None + assert _extract_event_id(parsed) == "e1" + assert _extract_batch_size(parsed) == 5 + assert _extract_marker_id(parsed) == "m1" + + +def test_psqlpy_events_resolve_listener_config_allows_mapping_direct_dsn() -> None: + """_resolve_listener_config accepts Mapping for direct_dsn without raising.""" + from sqlspec.adapters.psqlpy.events._hub import PsqlpyListenerHub + + config = PsqlpyConfig( + connection_config={"dsn": "postgresql://localhost/test"}, + driver_features={"pooler_mode": "transaction", "direct_dsn": {"host": "localhost", "database": "direct_db"}}, + ) + hub = PsqlpyListenerHub(config) + resolved = hub._resolve_listener_config() + assert resolved is not None + + +def test_psqlpy_events_backend_logger_namespace() -> None: + """Psqlpy events backend logger namespace follows standard naming.""" + from sqlspec.adapters.psqlpy.events.backend import logger + + assert logger.name == "sqlspec.adapters.psqlpy.events.backend" + + +async def test_psqlpy_litestar_store_operations_handle_none_query_result() -> None: + """exists, delete, delete_expired handle None query result safely.""" + from unittest.mock import AsyncMock + + from sqlspec.adapters.psqlpy.litestar.store import PsqlpyStore + + mock_conn = AsyncMock() + mock_conn.fetch = AsyncMock(return_value=None) + mock_conn.execute = AsyncMock(return_value=None) + + config = PsqlpyConfig(connection_config={"dsn": "postgresql://localhost/test"}) + store = PsqlpyStore(config) + object.__setattr__( + config, + "provide_connection", + lambda: AsyncMock(__aenter__=AsyncMock(return_value=mock_conn), __aexit__=AsyncMock()), + ) + + assert await store.exists("sess_1") is False + assert await store.expires_in("sess_1") is None + assert await store.delete_expired() == 0 From c1c7623fae35eab2bf78fd3d972fed1f4b7b5891 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 01:05:47 +0000 Subject: [PATCH 088/174] fix(psqlpy): correct data dictionary and catalog queries --- sqlspec/adapters/psqlpy/data_dictionary.py | 53 +++++-------------- .../test_psqlpy/test_transaction_state.py | 22 ++++++++ 2 files changed, 36 insertions(+), 39 deletions(-) diff --git a/sqlspec/adapters/psqlpy/data_dictionary.py b/sqlspec/adapters/psqlpy/data_dictionary.py index a45f39843..a4706a434 100644 --- a/sqlspec/adapters/psqlpy/data_dictionary.py +++ b/sqlspec/adapters/psqlpy/data_dictionary.py @@ -28,7 +28,12 @@ system_metadata_gated_result, unsupported_system_metadata_capability, ) -from sqlspec.data_dictionary.dialects.postgres import resolve_postgres_json_type +from sqlspec.data_dictionary.dialects.postgres import ( + POSTGRES_METADATA_DOMAINS, + POSTGRES_SUPPORTED_DOMAINS, + POSTGRES_SYSTEM_METADATA_QUERIES, + resolve_postgres_json_type, +) from sqlspec.driver import AsyncDataDictionaryBase if TYPE_CHECKING: @@ -38,37 +43,6 @@ __all__ = ("PsqlpyDataDictionary",) -_POSTGRES_METADATA_DOMAINS: tuple[str, ...] = ( - "schemas", - "objects", - "tables", - "columns", - "constraints", - "indexes", - "views", - "materialized_views", - "sequences", - "types", - "routines", - "triggers", - "comments", - "privileges", - "dependencies", - "extensions", - "partitions", - "ddl", - "system", -) -_POSTGRES_SUPPORTED_DOMAINS = frozenset(d for d in _POSTGRES_METADATA_DOMAINS if d != "system") -_POSTGRES_SYSTEM_METADATA_QUERIES = { - "settings": "settings", - "statement_history": "pg_stat_statements", - "pg_stat_statements": "pg_stat_statements", - "table_statistics": "table_stats", - "table_stats": "table_stats", -} - - @mypyc_attr(allow_interpreted_subclasses=True, native_class=False) class PsqlpyDataDictionary(AsyncDataDictionaryBase): """PostgreSQL-specific async data dictionary via psqlpy.""" @@ -89,7 +63,7 @@ async def get_system_metadata_capabilities( ) -> tuple[SystemMetadataCapability, ...]: """Get PostgreSQL opt-in system metadata capability disclosures.""" _ = driver - requested_domains = tuple(_POSTGRES_SYSTEM_METADATA_QUERIES) if domains is None else tuple(domains) + requested_domains = tuple(POSTGRES_SYSTEM_METADATA_QUERIES) if domains is None else tuple(domains) return tuple(_postgres_system_metadata_capability(domain) for domain in requested_domains) async def get_schemas(self, driver: "PsqlpyDriver") -> MetadataResult: @@ -163,7 +137,7 @@ async def get_system_metadata( if gate_result.capability.support != MetadataSupport.SUPPORTED: return gate_result - query_name = _POSTGRES_SYSTEM_METADATA_QUERIES.get(metadata_request.domain) + query_name = POSTGRES_SYSTEM_METADATA_QUERIES.get(metadata_request.domain) if query_name is None: return gate_result parameters: dict[str, object] = {} @@ -266,10 +240,11 @@ async def get_indexes( ) table_name = self.resolve_identifier(table) - self._log_table_describe(driver, schema_name=schema_name, table_name=table_name, operation="indexes") + table_schema_name = None if schema is None else schema_name + self._log_table_describe(driver, schema_name=table_schema_name, table_name=table_name, operation="indexes") return await driver.select( _postgres_domain_sql("indexes", "by_schema"), - schema_name=schema_name, + schema_name=table_schema_name, table_name=table_name, schema_type=IndexMetadata, ) @@ -334,7 +309,7 @@ def _postgres_metadata_capability(domain: str) -> MetadataCapability: risks=(MetadataRisk.EXPENSIVE, MetadataRisk.PRIVILEGED), warnings=("System metadata is opt-in and disabled by default.",), ) - if domain in _POSTGRES_SUPPORTED_DOMAINS: + if domain in POSTGRES_SUPPORTED_DOMAINS: return MetadataCapability( domain=domain, support=MetadataSupport.SUPPORTED, @@ -345,7 +320,7 @@ def _postgres_metadata_capability(domain: str) -> MetadataCapability: def _postgres_metadata_profile(adapter: str, domains: Sequence[str] | None) -> MetadataCapabilityProfile: - requested_domains = _POSTGRES_METADATA_DOMAINS if domains is None else tuple(domains) + requested_domains = POSTGRES_METADATA_DOMAINS if domains is None else tuple(domains) return MetadataCapabilityProfile( "postgres", adapter=adapter, @@ -404,7 +379,7 @@ def _ddl_result_from_rows( def _postgres_system_metadata_capability(domain: str) -> SystemMetadataCapability: - if domain not in _POSTGRES_SYSTEM_METADATA_QUERIES: + if domain not in POSTGRES_SYSTEM_METADATA_QUERIES: return unsupported_system_metadata_capability(domain) return SystemMetadataCapability( domain, diff --git a/tests/unit/adapters/test_psqlpy/test_transaction_state.py b/tests/unit/adapters/test_psqlpy/test_transaction_state.py index f7fc3872f..6faedab9f 100644 --- a/tests/unit/adapters/test_psqlpy/test_transaction_state.py +++ b/tests/unit/adapters/test_psqlpy/test_transaction_state.py @@ -292,3 +292,25 @@ async def test_psqlpy_litestar_store_operations_handle_none_query_result() -> No assert await store.exists("sess_1") is False assert await store.expires_in("sess_1") is None assert await store.delete_expired() == 0 + + +async def test_psqlpy_get_indexes_with_none_schema_resolves_via_search_path() -> None: + """get_indexes passes schema_name=None when schema is None to resolve table on search_path.""" + from unittest.mock import AsyncMock, MagicMock + + from sqlspec.adapters.psqlpy.data_dictionary import PsqlpyDataDictionary + + calls: list[dict[str, Any]] = [] + driver = MagicMock() + + async def fake_select(query: Any, **kwargs: Any) -> list[Any]: + calls.append(kwargs) + return [] + + driver.select = AsyncMock(side_effect=fake_select) + + dd = PsqlpyDataDictionary() + await dd.get_indexes(driver, table="users", schema=None) + assert len(calls) == 1 + assert calls[0]["schema_name"] is None + assert calls[0]["table_name"] == "users" From 93ff6ece67ac8c9d0b3ec5531a8019c441a5f268 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 01:05:47 +0000 Subject: [PATCH 089/174] perf(psqlpy): reduce hot-path allocations and add cache-hit fast path --- sqlspec/adapters/psqlpy/_typing.py | 10 ++- sqlspec/adapters/psqlpy/core.py | 31 ++++++--- sqlspec/adapters/psqlpy/driver.py | 68 +++++++++++++------ .../unit/adapters/test_psqlpy/test_driver.py | 61 +++++++++++++++++ 4 files changed, 138 insertions(+), 32 deletions(-) diff --git a/sqlspec/adapters/psqlpy/_typing.py b/sqlspec/adapters/psqlpy/_typing.py index a04e7f25b..6ba1186b4 100644 --- a/sqlspec/adapters/psqlpy/_typing.py +++ b/sqlspec/adapters/psqlpy/_typing.py @@ -171,6 +171,12 @@ async def __aexit__( return None +_CACHED_PSQLPY_DRIVER_CLS: "type[PsqlpyDriver] | None" = None + + def _resolve_psqlpy_driver() -> "type[PsqlpyDriver]": - module = importlib.import_module("sqlspec.adapters.psqlpy.driver") - return module.PsqlpyDriver # type: ignore[no-any-return] + global _CACHED_PSQLPY_DRIVER_CLS + if _CACHED_PSQLPY_DRIVER_CLS is None: + module = importlib.import_module("sqlspec.adapters.psqlpy.driver") + _CACHED_PSQLPY_DRIVER_CLS = module.PsqlpyDriver + return _CACHED_PSQLPY_DRIVER_CLS diff --git a/sqlspec/adapters/psqlpy/core.py b/sqlspec/adapters/psqlpy/core.py index 18ec28a3b..e25e9eb4e 100644 --- a/sqlspec/adapters/psqlpy/core.py +++ b/sqlspec/adapters/psqlpy/core.py @@ -5,7 +5,9 @@ import decimal import re import uuid +from collections.abc import Mapping from functools import lru_cache +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, cast import sqlglot @@ -33,7 +35,7 @@ resolve_postgres_extension_state, resolve_runtime_statement_config, ) -from sqlspec.core.parameters import type_coercion_dispatcher +from sqlspec.core.parameters import type_coercion_dispatcher, type_coercion_fallbacks from sqlspec.driver import CachedQuery from sqlspec.exceptions import ( CheckViolationError, @@ -62,7 +64,7 @@ from sqlspec.utils.uuids import uuid_from_string if TYPE_CHECKING: - from collections.abc import Callable, Mapping + from collections.abc import Callable from sqlspec.core import SQL @@ -449,23 +451,28 @@ def extract_rows_affected(result: Any) -> int: return 0 -def get_parameter_casts(statement: "SQL | CachedQuery") -> "dict[int, str]": +_EMPTY_PARAMETER_CASTS: Final[Mapping[int, str]] = MappingProxyType({}) + + +def get_parameter_casts(statement: "SQL | CachedQuery") -> "Mapping[int, str]": """Get parameter cast metadata from a SQLSpec statement.""" if isinstance(statement, CachedQuery): return statement.parameter_casts state = statement.get_processed_state() - return {} if state is Empty else state.parameter_casts or {} + return _EMPTY_PARAMETER_CASTS if state is Empty or not state.parameter_casts else state.parameter_casts def prepare_parameters_with_casts( - parameters: Any, parameter_casts: "dict[int, str]", statement_config: "StatementConfig" + parameters: Any, parameter_casts: "Mapping[int, str]", statement_config: "StatementConfig" ) -> Any: """Prepare parameters with cast-aware type coercion.""" if isinstance(parameters, (list, tuple)): - result: list[Any] = [] serializer = statement_config.parameter_config.json_serializer or to_json type_map = statement_config.parameter_config.type_coercion_map - dispatcher = type_coercion_dispatcher(tuple(type_map.items())) if type_map else None + fallbacks = type_coercion_fallbacks(type_map) if type_map else () + dispatcher = type_coercion_dispatcher(fallbacks) if fallbacks else None + coerced: list[Any] | None = None + for idx, param in enumerate(parameters, start=1): cast_type = parameter_casts.get(idx, "") prepared_value = param @@ -479,8 +486,14 @@ def prepare_parameters_with_casts( prepared_value = fallback_converter(prepared_value) if cast_type: prepared_value = _coerce_parameter_for_cast(prepared_value, cast_type, serializer) - result.append(prepared_value) - return tuple(result) if isinstance(parameters, tuple) else result + if prepared_value is not param and coerced is None: + coerced = list(parameters[: idx - 1]) + if coerced is not None: + coerced.append(prepared_value) + + if coerced is not None: + return tuple(coerced) if isinstance(parameters, tuple) else coerced + return parameters return parameters diff --git a/sqlspec/adapters/psqlpy/driver.py b/sqlspec/adapters/psqlpy/driver.py index 7fa4c0020..41932c78c 100644 --- a/sqlspec/adapters/psqlpy/driver.py +++ b/sqlspec/adapters/psqlpy/driver.py @@ -34,7 +34,7 @@ split_schema_and_table, ) from sqlspec.adapters.psqlpy.data_dictionary import PsqlpyDataDictionary -from sqlspec.core import SQL, StatementConfig, get_cache_config, register_driver_profile +from sqlspec.core import SQL, DMLResult, StatementConfig, get_cache_config, register_driver_profile from sqlspec.core.sqlcommenter import split_sqlcomment from sqlspec.driver import AsyncDriverAdapterBase, AsyncRowStream, BaseAsyncExceptionHandler from sqlspec.exceptions import SQLSpecError @@ -80,7 +80,7 @@ class PsqlpyDriver(AsyncDriverAdapterBase): and transaction management. """ - __slots__ = ("_data_dictionary", "_transaction_active") + __slots__ = ("_data_dictionary", "_pooler_unprepared", "_transaction_active") dialect = "postgres" def __init__( @@ -97,6 +97,10 @@ def __init__( super().__init__(connection=connection, statement_config=statement_config, driver_features=driver_features) self._data_dictionary: PsqlpyDataDictionary | None = None self._transaction_active = False + pooler_mode = self.driver_features.get("pooler_mode", "session") + self._pooler_unprepared = pooler_mode in {"transaction", "statement"} or bool( + self.driver_features.get("pgbouncer") or self.driver_features.get("pgbouncer_mode") + ) @property def data_dictionary(self) -> "PsqlpyDataDictionary": @@ -396,9 +400,8 @@ async def load_from_arrow( qualified = format_table_identifier(table) exc_handler = self.handle_database_exceptions() async with exc_handler, self.with_cursor(self.connection) as cursor: - await cursor.execute(f"TRUNCATE TABLE {qualified}") - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + await self._run_with_exception_handler(exc_handler, cursor.execute, f"TRUNCATE TABLE {qualified}") + self._check_pending_exception(exc_handler) columns, batch_iter = self._coerce_arrow_batches(source) first_batch = next((b for b in batch_iter if b.num_rows > 0), None) @@ -419,7 +422,9 @@ async def load_from_arrow( payload = self._encode_arrow_batch_pg_binary( batch, columns=columns, include_header=True, include_trailer=True ) - await cursor.binary_copy_to_table(payload, table_name, **copy_kwargs) + await self._run_with_exception_handler( + exc_handler, cursor.binary_copy_to_table, payload, table_name, **copy_kwargs + ) else: json_columns = await self._resolve_json_columns(schema_name, table_name) deserializer = cast( @@ -430,13 +435,14 @@ async def load_from_arrow( metrics["bytes_processed"] += int(batch.nbytes) records = self._arrow_batch_to_rows(batch, columns) if records: - await cursor.copy_records_to_table( + await self._run_with_exception_handler( + exc_handler, + cursor.copy_records_to_table, table_name, coerce_json_columns(records, columns, json_columns, deserializer=deserializer), **copy_kwargs, ) - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + self._check_pending_exception(exc_handler) telemetry_payload = self._ingest_telemetry_from_batches(metrics["rows_processed"], metrics["bytes_processed"]) telemetry_payload["destination"] = table @@ -554,15 +560,38 @@ async def _execute_cache_hit( ) -> "SQLResult": """Execute cached psqlpy queries with cast-aware parameter preparation.""" prepared_params = self.prepare_driver_parameters(params, self.statement_config, prepared_statement=cached) - direct_statement = self._cached_statement( - sql, - params, - cached, - cast("tuple[Any, ...] | list[Any] | dict[str, Any]", prepared_params), - params_are_simple=True, - compiled_sql=cached.compiled_sql, + args = ( + prepared_params + if isinstance(prepared_params, (tuple, list)) + else tuple(prepared_params.values()) + if isinstance(prepared_params, dict) + else () ) - return await self._execute_cached_statement(direct_statement) + wire_sql = self._format_wire_sql(cached.compiled_sql) + if cached.operation_profile.returns_rows: + query_result = ( + await self.connection.fetch(wire_sql, args, prepared=False) + if self._pooler_unprepared + else await self.connection.fetch(wire_sql, args) + ) + dict_rows, column_names = collect_rows(query_result) + return self.build_cached_select_result( + sql, + params, + cached, + selected_data=dict_rows, + column_names=column_names, + data_row_count=len(dict_rows), + row_format="dict", + compiled_sql=wire_sql, + ) + result = ( + await self.connection.execute(wire_sql, args, prepared=False) + if self._pooler_unprepared + else await self.connection.execute(wire_sql, args) + ) + affected_rows = extract_rows_affected(result) + return DMLResult(cached.operation_type, affected_rows) def _connection_in_transaction(self) -> bool: """Check if connection is in transaction. @@ -573,10 +602,7 @@ def _connection_in_transaction(self) -> bool: return self._transaction_active def _is_transaction_pooler_mode(self) -> bool: - pooler_mode = self.driver_features.get("pooler_mode", "session") - return pooler_mode in {"transaction", "statement"} or bool( - self.driver_features.get("pgbouncer") or self.driver_features.get("pgbouncer_mode") - ) + return self._pooler_unprepared def _should_set_local_columnar(self) -> bool: return bool(self.driver_features.get("enable_alloydb_columnar", False)) and self._is_transaction_pooler_mode() diff --git a/tests/unit/adapters/test_psqlpy/test_driver.py b/tests/unit/adapters/test_psqlpy/test_driver.py index bedd66489..a956a374b 100644 --- a/tests/unit/adapters/test_psqlpy/test_driver.py +++ b/tests/unit/adapters/test_psqlpy/test_driver.py @@ -106,3 +106,64 @@ async def test_psqlpy_transaction_methods_raise_mapped_exceptions() -> None: driver._transaction_active = True # pyright: ignore[reportPrivateUsage] with pytest.raises(DatabaseConnectionError): await driver.rollback() + + +@pytest.mark.anyio +async def test_psqlpy_prepare_parameters_with_casts_zero_allocation() -> None: + """prepare_parameters_with_casts returns original sequence when no coercion or cast is needed.""" + from sqlspec.adapters.psqlpy.core import _EMPTY_PARAMETER_CASTS, prepare_parameters_with_casts + + params = (1, "two", 3.0) + cfg = build_statement_config() + result = prepare_parameters_with_casts(params, _EMPTY_PARAMETER_CASTS, cfg) + assert result is params + + +@pytest.mark.anyio +async def test_psqlpy_execute_cache_hit_direct_fast_path() -> None: + """_execute_cache_hit executes directly on connection without calling _execute_cached_statement.""" + from sqlspec.core.cache import CachedQuery, StatementOperationProfile + + conn = AsyncMock() + conn.fetch = AsyncMock(return_value=[{"col": 1}]) + driver = PsqlpyDriver(connection=cast("Any", conn), statement_config=build_statement_config()) + + cached = CachedQuery( + compiled_sql="SELECT 1 AS col", + operation_profile=StatementOperationProfile(returns_rows=True), + operation_type="SELECT", + ) + result = await driver._execute_cache_hit("SELECT 1 AS col", (), cached) + assert result.data == [{"col": 1}] + assert conn.fetch.await_count == 1 + + +@pytest.mark.anyio +async def test_psqlpy_load_from_arrow_preserves_cause_on_failure() -> None: + """load_from_arrow preserves underlying __cause__ on execution failure.""" + import pyarrow as pa + from psqlpy.exceptions import ConnectionExecuteError + + conn = AsyncMock() + conn.binary_copy_to_table = AsyncMock(side_effect=ConnectionExecuteError("binary copy network drop")) + driver = PsqlpyDriver( + connection=cast("Any", conn), + statement_config=build_statement_config(), + driver_features={"enable_pgcopy_binary": True, "arrow_import_enabled": True}, + ) + tbl = pa.Table.from_arrays([pa.array([1, 2])], names=["id"]) + + with pytest.raises(Exception) as exc_info: + await driver.load_from_arrow("items", tbl) + + assert exc_info.value.__cause__ is not None or "binary copy network drop" in str(exc_info.value) + + +def test_psqlpy_driver_class_cached_in_resolve() -> None: + """_resolve_psqlpy_driver caches the resolved driver class.""" + from sqlspec.adapters.psqlpy._typing import _resolve_psqlpy_driver + + cls1 = _resolve_psqlpy_driver() + cls2 = _resolve_psqlpy_driver() + assert cls1 is cls2 + assert cls1 is PsqlpyDriver From 20d3a6718a49920a6731319f513295659946a96a Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 01:05:47 +0000 Subject: [PATCH 090/174] refactor(psqlpy): clean up structure, naming, and dead code --- sqlspec/adapters/psqlpy/adk/store.py | 1 - sqlspec/adapters/psqlpy/config.py | 8 +++-- sqlspec/adapters/psqlpy/core.py | 36 ------------------- sqlspec/adapters/psqlpy/type_converter.py | 29 +++------------ .../adapters/test_psqlpy/test_adk_store.py | 12 ------- 5 files changed, 9 insertions(+), 77 deletions(-) diff --git a/sqlspec/adapters/psqlpy/adk/store.py b/sqlspec/adapters/psqlpy/adk/store.py index b7ca54983..f8e43f265 100644 --- a/sqlspec/adapters/psqlpy/adk/store.py +++ b/sqlspec/adapters/psqlpy/adk/store.py @@ -24,7 +24,6 @@ logger = get_logger("sqlspec.adapters.psqlpy.adk.store") -POSTGRES_TABLE_NOT_FOUND_SQLSTATE: Final = "42P01" PSQLPY_STATUS_REGEX: Final[re.Pattern[str]] = re.compile(r"^([A-Z]+)(?:\s+(\d+))?\s+(\d+)$", re.IGNORECASE) diff --git a/sqlspec/adapters/psqlpy/config.py b/sqlspec/adapters/psqlpy/config.py index d59c166c3..c9a5abc4c 100644 --- a/sqlspec/adapters/psqlpy/config.py +++ b/sqlspec/adapters/psqlpy/config.py @@ -162,6 +162,8 @@ class PsqlpyDriverFeatures(TypedDict): """Async callback invoked once per physical ``psqlpy`` connection acquired from the pool.""" enable_events: NotRequired[bool] """Enable SQLSpec event channel pub/sub support (defaults to ``True`` when ``events`` is configured).""" + enable_pgcopy_binary: NotRequired[bool] + """Enable binary wire format for PostgreSQL COPY operations.""" events_backend: NotRequired[Literal["notify", "notify_queue", "poll_queue"]] """Event transport backend (``notify`` for LISTEN/NOTIFY, ``notify_queue``, or ``poll_queue``).""" @@ -240,10 +242,10 @@ class PsqlpyConfig(AsyncDatabaseConfig[PsqlpyConnection, "ConnectionPool", Psqlp connection_type: "ClassVar[type[PsqlpyConnection]]" = PsqlpyConnection supports_transactional_ddl: "ClassVar[bool]" = True supports_migration_schemas: "ClassVar[bool]" = True - supports_native_arrow_export: ClassVar[bool] = True + supports_native_arrow_export: ClassVar[bool] = False supports_native_arrow_import: ClassVar[bool] = True - supports_native_parquet_export: ClassVar[bool] = True - supports_native_parquet_import: ClassVar[bool] = True + supports_native_parquet_export: ClassVar[bool] = False + supports_native_parquet_import: ClassVar[bool] = False supports_native_row_streaming: ClassVar[bool] = True type_coercion_capabilities: "ClassVar[TypeCoercionCapabilities]" = TypeCoercionCapabilities( datetime_binding="native", timestamp_precision="microsecond", json_columns_decoded=True, uuid_binding="native" diff --git a/sqlspec/adapters/psqlpy/core.py b/sqlspec/adapters/psqlpy/core.py index e25e9eb4e..9bead644d 100644 --- a/sqlspec/adapters/psqlpy/core.py +++ b/sqlspec/adapters/psqlpy/core.py @@ -81,7 +81,6 @@ "collect_rows", "create_mapped_exception", "default_statement_config", - "detect_postgres_extensions", "driver_profile", "extract_rows_affected", "format_execute_many_parameters", @@ -637,41 +636,6 @@ def format_execute_many_parameters(parameters: Any, *, coerce_numeric: bool) -> return [_format_execute_many_param_set(parameters, coerce_numeric=coerce_numeric)] -async def detect_postgres_extensions( - connection: Any, *, enable_pgvector: bool = False, enable_paradedb: bool = False, enable_pg_textsearch: bool = False -) -> "tuple[bool, bool, bool]": - """Detect whether pgvector, paradedb, and pg_textsearch extensions are installed on PostgreSQL. - - Args: - connection: Active psqlpy connection. - enable_pgvector: Whether to probe for pgvector extension. - enable_paradedb: Whether to probe for paradedb extension. - enable_pg_textsearch: Whether to probe for pg_textsearch extension. - - Returns: - Tuple of ``(pgvector_available, paradedb_available, pg_textsearch_available)``. - """ - driver_features: dict[str, Any] = { - "enable_pgvector": enable_pgvector, - "enable_paradedb": enable_paradedb, - "enable_pg_textsearch": enable_pg_textsearch, - } - extensions_to_check = build_postgres_extension_probe_names(driver_features) - if not extensions_to_check: - return False, False, False - - try: - result = await connection.fetch( - "SELECT extname FROM pg_extension WHERE extname = ANY($1::text[])", [extensions_to_check] - ) - rows = result.result() if result else [] - detected = {cast("str", row["extname"] if isinstance(row, dict) else row[0]) for row in rows} - except Exception: - logger.debug("PostgreSQL extension detection failed for psqlpy connection") - return False, False, False - return "vector" in detected, "pg_search" in detected, "pg_textsearch" in detected - - def _coerce_json_parameter(value: Any, cast_type: str, serializer: "Callable[[Any], str]") -> Any: """Serialize JSON parameters according to the detected cast type.""" diff --git a/sqlspec/adapters/psqlpy/type_converter.py b/sqlspec/adapters/psqlpy/type_converter.py index 7d3941837..25ac54664 100644 --- a/sqlspec/adapters/psqlpy/type_converter.py +++ b/sqlspec/adapters/psqlpy/type_converter.py @@ -1,17 +1,10 @@ -"""PostgreSQL-specific helpers for the psqlpy adapter. +"""PostgreSQL-specific helpers for the psqlpy adapter.""" -This module preserves the ``register_pgvector`` placeholder used by the -driver configuration layer. -""" +from typing import Any -from typing import TYPE_CHECKING, Any +from sqlspec.typing import import_optional_attr -from sqlspec.typing import PGVECTOR_INSTALLED, import_optional_attr - -if TYPE_CHECKING: - from sqlspec.adapters.psqlpy._typing import PsqlpyConnection as Connection - -__all__ = ("coerce_pgvector", "register_pgvector") +__all__ = ("coerce_pgvector",) _PGVECTOR_TYPE = import_optional_attr("psqlpy.extra_types", "PgVector") @@ -26,17 +19,3 @@ def coerce_pgvector(value: Any) -> Any: if callable(tolist): return _PGVECTOR_TYPE(tolist()) return value - - -def register_pgvector(connection: "Connection") -> None: - """Register pgvector type handlers on psqlpy connection. - - Currently a placeholder for future implementation. The psqlpy library - does not yet expose a type handler registration API compatible with - pgvector's automatic conversion system. - - Args: - connection: Psqlpy connection instance. - """ - if not PGVECTOR_INSTALLED: - return diff --git a/tests/unit/adapters/test_psqlpy/test_adk_store.py b/tests/unit/adapters/test_psqlpy/test_adk_store.py index df1da7f5d..b7ba74ca7 100644 --- a/tests/unit/adapters/test_psqlpy/test_adk_store.py +++ b/tests/unit/adapters/test_psqlpy/test_adk_store.py @@ -9,7 +9,6 @@ from sqlspec.adapters.psqlpy.adk import PsqlpyADKConfig, PsqlpyADKMemoryStore, PsqlpyADKStore from sqlspec.adapters.psqlpy.adk.store import _is_psqlpy_database_error, _is_table_missing_error, decode_postgres_vector -from sqlspec.adapters.psqlpy.core import detect_postgres_extensions from sqlspec.config import ADKConfig @@ -301,14 +300,3 @@ async def test_psqlpy_memory_store_normalizes_tuple_embeddings_on_insert_and_sea assert "unnest" in conn.fetch_calls[0][0].lower() assert conn.fetch_calls[1][1][0] == [0.1, 0.2] assert results[0]["embedding"] == [0.1, 0.2] - - -async def test_psqlpy_detect_postgres_extensions_sets_pg_textsearch_available() -> None: - """detect_postgres_extensions detects vector, pg_search, and pg_textsearch for psqlpy connections.""" - conn = _FakeMemoryConnection([{"extname": "vector"}, {"extname": "pg_search"}, {"extname": "pg_textsearch"}]) - pgvector, paradedb, pg_textsearch = await detect_postgres_extensions( - conn, enable_pgvector=True, enable_paradedb=True, enable_pg_textsearch=True - ) - assert pgvector is True - assert paradedb is True - assert pg_textsearch is True From 8a393e90b8637b0136c2b7945cafc78f13ff6b9b Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 01:05:47 +0000 Subject: [PATCH 091/174] fix(psycopg): correct exception mapping, parameter binding, and transaction lifecycle --- sqlspec/adapters/psycopg/_typing.py | 10 +- sqlspec/adapters/psycopg/config.py | 128 ++++++++++-------- sqlspec/adapters/psycopg/core.py | 12 ++ sqlspec/adapters/psycopg/driver.py | 28 +++- .../unit/adapters/test_psycopg/test_config.py | 8 +- .../unit/adapters/test_psycopg/test_driver.py | 58 +++++++- .../test_psycopg/test_type_handlers.py | 12 ++ 7 files changed, 184 insertions(+), 72 deletions(-) diff --git a/sqlspec/adapters/psycopg/_typing.py b/sqlspec/adapters/psycopg/_typing.py index 9a74fe490..2e4a8cdbd 100644 --- a/sqlspec/adapters/psycopg/_typing.py +++ b/sqlspec/adapters/psycopg/_typing.py @@ -215,7 +215,10 @@ def __exit__( self, exc_type: "type[BaseException] | None", exc_val: "BaseException | None", exc_tb: "TracebackType | None" ) -> "bool | None": try: - if self._driver is not None and getattr(self._driver, "_transaction_active", False): + if self._driver is not None and ( + getattr(self._driver, "_transaction_active", False) + or getattr(self._driver, "_connection_in_transaction", lambda: False)() + ): with contextlib.suppress(Exception): self._driver.rollback() finally: @@ -276,7 +279,10 @@ async def __aexit__( self, exc_type: "type[BaseException] | None", exc_val: "BaseException | None", exc_tb: "TracebackType | None" ) -> "bool | None": try: - if self._driver is not None and getattr(self._driver, "_transaction_active", False): + if self._driver is not None and ( + getattr(self._driver, "_transaction_active", False) + or getattr(self._driver, "_connection_in_transaction", lambda: False)() + ): with contextlib.suppress(Exception): await self._driver.rollback() finally: diff --git a/sqlspec/adapters/psycopg/config.py b/sqlspec/adapters/psycopg/config.py index 80813c2cd..a2d5c8231 100644 --- a/sqlspec/adapters/psycopg/config.py +++ b/sqlspec/adapters/psycopg/config.py @@ -1,6 +1,8 @@ """Psycopg database configuration with direct field-based configuration.""" +import asyncio import inspect +import threading from collections.abc import Mapping from contextlib import suppress from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, NotRequired, Self, TypedDict, cast @@ -347,8 +349,8 @@ def __init__( statement_config, driver_features = apply_driver_features(statement_config, driver_features) features_dict = dict(driver_features) if driver_features else {} - self._user_connection_hook: Callable[[PsycopgSyncConnection], None] | None = features_dict.get( - "on_connection_create" + self._user_connection_hook: Callable[[PsycopgSyncConnection], None] | None = features_dict.pop( + "on_connection_create", None ) features_dict.setdefault("enable_alloydb", False) features_dict.setdefault("enable_alloydb_iam_auth", False) @@ -368,6 +370,7 @@ def __init__( connection_config["options"] = columnar_opt elif columnar_opt not in str(existing_options): connection_config["options"] = f"{existing_options} {columnar_opt}" + self._probe_lock: threading.Lock = threading.Lock() self._pgvector_available: bool | None = None self._paradedb_available: bool | None = None self._pg_textsearch_available: bool | None = None @@ -636,29 +639,37 @@ def _configure_connection(self, conn: "PsycopgSyncConnection") -> None: set_json_loads(deserializer, conn) if self._pgvector_available is None: - detected_extensions: set[str] = set() - extensions = build_postgres_extension_probe_names(self.driver_features) - probe_succeeded = True - if extensions: - try: - cursor = conn.execute( - "SELECT extname FROM pg_extension WHERE extname = ANY(%s::text[])", (extensions,) - ) - results = cursor.fetchall() - detected_extensions = { - cast("str", r["extname"] if isinstance(r, Mapping) else r[0]) for r in results - } - except Exception as exc: - probe_succeeded = False - self._pg_textsearch_probe_error = exc + with self._probe_lock: + if self._pgvector_available is None: + detected_extensions: set[str] = set() + extensions = build_postgres_extension_probe_names(self.driver_features) + probe_succeeded = True + if extensions: + try: + cursor = conn.execute( + "SELECT extname FROM pg_extension WHERE extname = ANY(%s::text[])", (extensions,) + ) + results = cursor.fetchall() + detected_extensions = { + cast("str", r["extname"] if isinstance(r, Mapping) else r[0]) for r in results + } + except Exception as exc: + probe_succeeded = False + self._pg_textsearch_probe_error = exc + if not conn.autocommit: + with suppress(Exception): + conn.rollback() + if probe_succeeded: + self.statement_config, self._pgvector_available, self._paradedb_available = ( + resolve_postgres_extension_state( + self.statement_config, self.driver_features, detected_extensions + ) + ) + self._pg_textsearch_available = is_postgres_extension_active( + self.driver_features, "pg_textsearch" + ) if not conn.autocommit: - with suppress(Exception): - conn.rollback() - if probe_succeeded: - self.statement_config, self._pgvector_available, self._paradedb_available = ( - resolve_postgres_extension_state(self.statement_config, self.driver_features, detected_extensions) - ) - self._pg_textsearch_available = is_postgres_extension_active(self.driver_features, "pg_textsearch") + conn.rollback() if self._pgvector_available: register_pgvector_sync(conn) @@ -666,9 +677,6 @@ def _configure_connection(self, conn: "PsycopgSyncConnection") -> None: if self._user_connection_hook is not None: self._user_connection_hook(conn) - if not conn.autocommit: - conn.rollback() - def _ensure_pg_textsearch_available(self) -> None: if self._pg_textsearch_available: return @@ -811,8 +819,8 @@ def __init__( if features_dict.get("enable_alloydb", False): msg = "Google AlloyDB connector integration (enable_alloydb) is only supported on PsycopgSyncConfig" raise ImproperConfigurationError(msg) - self._user_connection_hook: Callable[[PsycopgAsyncConnection], Awaitable[None]] | None = features_dict.get( - "on_connection_create" + self._user_connection_hook: Callable[[PsycopgAsyncConnection], Awaitable[None]] | None = features_dict.pop( + "on_connection_create", None ) pooler_mode = features_dict.get("pooler_mode", "session") is_transaction_pooler = bool( @@ -829,6 +837,7 @@ def __init__( connection_config["options"] = columnar_opt elif columnar_opt not in str(existing_options): connection_config["options"] = f"{existing_options} {columnar_opt}" + self._probe_lock: asyncio.Lock | None = None self._pgvector_available: bool | None = None self._paradedb_available: bool | None = None self._pg_textsearch_available: bool | None = None @@ -1045,38 +1054,47 @@ async def _configure_async_connection(self, conn: "PsycopgAsyncConnection") -> N set_json_loads(deserializer, conn) if self._pgvector_available is None: - detected_extensions: set[str] = set() - extensions = build_postgres_extension_probe_names(self.driver_features) - probe_succeeded = True - if extensions: - try: - cursor = await conn.execute( - "SELECT extname FROM pg_extension WHERE extname = ANY(%s::text[])", (extensions,) - ) - results = await cursor.fetchall() - detected_extensions = { - cast("str", r["extname"] if isinstance(r, Mapping) else r[0]) for r in results - } - except Exception as exc: - probe_succeeded = False - self._pg_textsearch_probe_error = exc + if self._probe_lock is None: + self._probe_lock = asyncio.Lock() + async with self._probe_lock: + if self._pgvector_available is None: + detected_extensions: set[str] = set() + extensions = build_postgres_extension_probe_names(self.driver_features) + probe_succeeded = True + if extensions: + try: + cursor = await conn.execute( + "SELECT extname FROM pg_extension WHERE extname = ANY(%s::text[])", (extensions,) + ) + results = await cursor.fetchall() + detected_extensions = { + cast("str", r["extname"] if isinstance(r, Mapping) else r[0]) for r in results + } + except Exception as exc: + probe_succeeded = False + self._pg_textsearch_probe_error = exc + if not conn.autocommit: + with suppress(Exception): + await conn.rollback() + if probe_succeeded: + self.statement_config, self._pgvector_available, self._paradedb_available = ( + resolve_postgres_extension_state( + self.statement_config, self.driver_features, detected_extensions + ) + ) + self._pg_textsearch_available = is_postgres_extension_active( + self.driver_features, "pg_textsearch" + ) if not conn.autocommit: - with suppress(Exception): - await conn.rollback() - if probe_succeeded: - self.statement_config, self._pgvector_available, self._paradedb_available = ( - resolve_postgres_extension_state(self.statement_config, self.driver_features, detected_extensions) - ) - self._pg_textsearch_available = is_postgres_extension_active(self.driver_features, "pg_textsearch") + await conn.rollback() if self._pgvector_available: await register_pgvector_async(conn) if self._user_connection_hook is not None: - await self._user_connection_hook(conn) - - if not conn.autocommit: - await conn.rollback() + hook_res = self._user_connection_hook(conn) + if inspect.isawaitable(hook_res): + await hook_res def _ensure_pg_textsearch_available(self) -> None: if self._pg_textsearch_available: diff --git a/sqlspec/adapters/psycopg/core.py b/sqlspec/adapters/psycopg/core.py index 2de7d9a12..9e2c58724 100644 --- a/sqlspec/adapters/psycopg/core.py +++ b/sqlspec/adapters/psycopg/core.py @@ -26,6 +26,8 @@ from sqlspec.exceptions import ( CheckViolationError, ConnectionTimeoutError, + DatabaseConnectionError, + DataError, DeadlockError, ForeignKeyViolationError, ImproperConfigurationError, @@ -614,6 +616,12 @@ def escape_literal_percent(sql: str, parameters: Any, validator: "ParameterValid def _register_exception_mappings() -> None: + from sqlspec.typing import import_optional_attr + + pool_timeout = import_optional_attr("psycopg_pool", "PoolTimeout") + if pool_timeout is not None: + _EXCEPTION_MAPPING[pool_timeout] = ("", ConnectionTimeoutError, "connection pool timeout") + _EXCEPTION_MAPPING.update({ pg_errors.UniqueViolation: ("23505", UniqueViolationError, "unique constraint violation"), pg_errors.ForeignKeyViolation: ("23503", ForeignKeyViolationError, "foreign key constraint violation"), @@ -627,6 +635,10 @@ def _register_exception_mappings() -> None: pg_errors.AdminShutdown: ("57P01", ConnectionTimeoutError, "admin shutdown"), pg_errors.CannotConnectNow: ("57P03", ConnectionTimeoutError, "cannot connect now"), pg_errors.SyntaxError: ("42601", SQLParsingError, "SQL syntax error"), + psycopg.OperationalError: ("", DatabaseConnectionError, "operational error"), + psycopg.DataError: ("", DataError, "data error"), + psycopg.NotSupportedError: ("", SQLSpecError, "not supported"), + psycopg.InterfaceError: ("", DatabaseConnectionError, "interface error"), }) diff --git a/sqlspec/adapters/psycopg/driver.py b/sqlspec/adapters/psycopg/driver.py index 3ddebc4eb..83cfe42ce 100644 --- a/sqlspec/adapters/psycopg/driver.py +++ b/sqlspec/adapters/psycopg/driver.py @@ -1,7 +1,7 @@ """PostgreSQL psycopg driver implementation.""" from collections.abc import Sized -from contextlib import AsyncExitStack, ExitStack +from contextlib import AsyncExitStack, ExitStack, suppress from itertools import chain from typing import TYPE_CHECKING, Any, Literal, LiteralString, cast @@ -380,17 +380,21 @@ def dispatch_special_handling(self, cursor: Any, statement: "SQL") -> "SQLResult def begin(self) -> None: """Begin a database transaction on the current connection.""" + if getattr(getattr(self.connection, "info", None), "transaction_status", None) == TRANSACTION_STATUS_INERROR: + self.connection.rollback() if self._connection_in_transaction(): return + restore_autocommit = False try: - if self.connection.info.transaction_status == TRANSACTION_STATUS_INERROR: - self.connection.rollback() restore_autocommit = self.connection.autocommit if restore_autocommit: self.connection.autocommit = False if self._should_set_local_columnar(): self.connection.execute("SET LOCAL google_columnar_engine.enable_columnar_scan = 'on'") except psycopg.Error as e: + if restore_autocommit: + with suppress(Exception): + self.connection.autocommit = True mapped = create_mapped_exception(e) if type(mapped) is not SQLSpecError: raise mapped from e @@ -711,9 +715,12 @@ def _resolve_column_names(self, description: Any) -> list[str]: def _connection_in_transaction(self) -> bool: """Check if connection is in transaction.""" - return self._transaction_active or self.connection.info.transaction_status in { + info = getattr(self.connection, "info", None) + status = getattr(info, "transaction_status", None) + return self._transaction_active or status in { TRANSACTION_STATUS_ACTIVE, TRANSACTION_STATUS_INTRANS, + TRANSACTION_STATUS_INERROR, } def _restore_original_autocommit(self) -> None: @@ -946,17 +953,21 @@ async def dispatch_special_handling(self, cursor: Any, statement: "SQL") -> "SQL async def begin(self) -> None: """Begin a database transaction on the current connection.""" + if getattr(getattr(self.connection, "info", None), "transaction_status", None) == TRANSACTION_STATUS_INERROR: + await self.connection.rollback() if self._connection_in_transaction(): return + restore_autocommit = False try: - if self.connection.info.transaction_status == TRANSACTION_STATUS_INERROR: - await self.connection.rollback() restore_autocommit = self.connection.autocommit if restore_autocommit: await self.connection.set_autocommit(False) if self._should_set_local_columnar(): await self.connection.execute("SET LOCAL google_columnar_engine.enable_columnar_scan = 'on'") except psycopg.Error as e: + if restore_autocommit: + with suppress(Exception): + await self.connection.set_autocommit(True) mapped = create_mapped_exception(e) if type(mapped) is not SQLSpecError: raise mapped from e @@ -1282,9 +1293,12 @@ def _resolve_column_names(self, description: Any) -> list[str]: def _connection_in_transaction(self) -> bool: """Check if connection is in transaction.""" - return self._transaction_active or self.connection.info.transaction_status in { + info = getattr(self.connection, "info", None) + status = getattr(info, "transaction_status", None) + return self._transaction_active or status in { TRANSACTION_STATUS_ACTIVE, TRANSACTION_STATUS_INTRANS, + TRANSACTION_STATUS_INERROR, } async def _restore_original_autocommit(self) -> None: diff --git a/tests/unit/adapters/test_psycopg/test_config.py b/tests/unit/adapters/test_psycopg/test_config.py index af34e66a6..5529adea5 100644 --- a/tests/unit/adapters/test_psycopg/test_config.py +++ b/tests/unit/adapters/test_psycopg/test_config.py @@ -632,7 +632,7 @@ def test_async_config_normalizes_aliases() -> None: def test_psycopg_sync_configure_connection_supports_dict_row_and_rolls_back_after_hook() -> None: - """_configure_connection should handle dict_row probe results and roll back after user hooks when autocommit=False.""" + """_configure_connection should handle dict_row probe results and roll back before user hooks when autocommit=False.""" events: list[str] = [] connection = MagicMock() connection.autocommit = False @@ -652,12 +652,12 @@ def test_psycopg_sync_configure_connection_supports_dict_row_and_rolls_back_afte assert config.pg_textsearch_available is True assert config.statement_config.dialect == "pg_textsearch" - assert events == ["hook", "rollback"] + assert events == ["rollback", "hook"] @pytest.mark.anyio async def test_psycopg_async_configure_connection_supports_dict_row_and_rolls_back_after_hook() -> None: - """_configure_async_connection should handle dict_row probe results and roll back after user hooks when autocommit=False.""" + """_configure_async_connection should handle dict_row probe results and roll back before user hooks when autocommit=False.""" events: list[str] = [] connection = AsyncMock() connection.autocommit = False @@ -686,7 +686,7 @@ async def record_hook(_: object) -> None: assert config.pg_textsearch_available is True assert config.statement_config.dialect == "pg_textsearch" - assert events == ["hook", "rollback"] + assert events == ["rollback", "hook"] def test_psycopg_pooler_mode_transaction_sets_prepare_threshold_none_and_alloydb_class( diff --git a/tests/unit/adapters/test_psycopg/test_driver.py b/tests/unit/adapters/test_psycopg/test_driver.py index 4cdf238c7..1955594a9 100644 --- a/tests/unit/adapters/test_psycopg/test_driver.py +++ b/tests/unit/adapters/test_psycopg/test_driver.py @@ -119,6 +119,7 @@ def commit(self) -> None: def rollback(self) -> None: self.rollback_calls += 1 + self.info.transaction_status = 0 if self.rollback_error is not None: raise self.rollback_error @@ -158,6 +159,7 @@ async def commit(self) -> None: async def rollback(self) -> None: self.rollback_calls += 1 + self.info.transaction_status = 0 if self.rollback_error is not None: raise self.rollback_error @@ -409,11 +411,11 @@ async def test_async_logical_transaction_is_reported_while_libpq_is_idle() -> No def test_sync_inerror_status_is_not_active_transaction_and_rolls_back_on_begin() -> None: - """TRANSACTION_STATUS_INERROR (3) should not report an active transaction and should roll back on begin().""" + """TRANSACTION_STATUS_INERROR (3) is reported in transaction state and rolls back on begin().""" connection = _SyncTransactionConnection(transaction_status=3) driver = _sync_driver(connection) - assert driver._connection_in_transaction() is False # pyright: ignore[reportPrivateUsage] + assert driver._connection_in_transaction() is True # pyright: ignore[reportPrivateUsage] driver.begin() assert connection.rollback_calls == 1 @@ -422,11 +424,11 @@ def test_sync_inerror_status_is_not_active_transaction_and_rolls_back_on_begin() async def test_async_inerror_status_is_not_active_transaction_and_rolls_back_on_begin() -> None: - """TRANSACTION_STATUS_INERROR (3) should not report an active transaction and should roll back on async begin().""" + """TRANSACTION_STATUS_INERROR (3) is reported in transaction state and rolls back on async begin().""" connection = _AsyncTransactionConnection(transaction_status=3) driver = _async_driver(connection) - assert driver._connection_in_transaction() is False # pyright: ignore[reportPrivateUsage] + assert driver._connection_in_transaction() is True # pyright: ignore[reportPrivateUsage] await driver.begin() assert connection.rollback_calls == 1 @@ -460,3 +462,51 @@ def test_sync_copy_to_stdout_decodes_memoryview_chunks_and_reports_rowcount() -> assert result is not None assert result.get_data() == [{"copy_output": "1\talice\n2\tbob\n"}] assert result.rows_affected == 2 + + +def test_sync_session_exit_rolls_back_uncommitted_implicit_transaction() -> None: + """Session exit rolls back uncommitted implicit autocommit=False transactions.""" + connection = _SyncTransactionConnection(transaction_status=2) + config = PsycopgSyncConfig() + context = config.provide_session() + object.__setattr__(context, "_acquire_connection", lambda: connection) + with context: + pass + assert connection.rollback_calls == 1 + + +@pytest.mark.anyio +async def test_async_session_exit_rolls_back_uncommitted_implicit_transaction() -> None: + """Async session exit rolls back uncommitted implicit autocommit=False transactions.""" + connection = _AsyncTransactionConnection(transaction_status=2) + config = PsycopgAsyncConfig() + context = config.provide_session() + object.__setattr__(context, "_acquire_connection", AsyncMock(return_value=connection)) + async with context: + pass + assert connection.rollback_calls == 1 + + +def test_sync_begin_restores_autocommit_on_set_local_failure() -> None: + """begin() restores autocommit=True if SET LOCAL columnar fails.""" + connection = _SyncTransactionConnection(autocommit=True) + connection.execute = MagicMock(side_effect=psycopg.Error("columnar extension missing")) + driver = PsycopgSyncDriver( + connection=connection, driver_features={"enable_alloydb_columnar": True, "pooler_mode": "transaction"} + ) + with pytest.raises(SQLSpecError): + driver.begin() + assert connection.autocommit is True + + +@pytest.mark.anyio +async def test_async_begin_restores_autocommit_on_set_local_failure() -> None: + """async begin() restores autocommit=True if SET LOCAL columnar fails.""" + connection = _AsyncTransactionConnection(autocommit=True) + connection.execute = AsyncMock(side_effect=psycopg.Error("columnar extension missing")) + driver = PsycopgAsyncDriver( + connection=connection, driver_features={"enable_alloydb_columnar": True, "pooler_mode": "transaction"} + ) + with pytest.raises(SQLSpecError): + await driver.begin() + assert connection.autocommit is True diff --git a/tests/unit/adapters/test_psycopg/test_type_handlers.py b/tests/unit/adapters/test_psycopg/test_type_handlers.py index 9d959e557..1ac27970e 100644 --- a/tests/unit/adapters/test_psycopg/test_type_handlers.py +++ b/tests/unit/adapters/test_psycopg/test_type_handlers.py @@ -93,3 +93,15 @@ def test_pgvector_module_reference_is_cached() -> None: assert not hasattr(type_converter, "importlib") assert hasattr(type_converter, "_pgvector_psycopg") + + +def test_operational_error_with_none_sqlstate_maps_to_database_connection_error() -> None: + """OperationalError with sqlstate=None maps to DatabaseConnectionError.""" + import psycopg + + from sqlspec.adapters.psycopg.core import create_mapped_exception + from sqlspec.exceptions import DatabaseConnectionError + + err = psycopg.OperationalError("connection failed") + mapped = create_mapped_exception(err) + assert type(mapped) is DatabaseConnectionError From a79400dfbd68656702078774bfcf7f7d99bbdc32 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 01:05:47 +0000 Subject: [PATCH 092/174] fix(psycopg): remediate ADK, Litestar, and events stores --- sqlspec/adapters/psycopg/adk/store.py | 52 ++-- sqlspec/adapters/psycopg/core.py | 34 +++ sqlspec/adapters/psycopg/events/_hub.py | 5 +- sqlspec/adapters/psycopg/events/backend.py | 64 +++-- sqlspec/adapters/psycopg/events/store.py | 28 +- sqlspec/adapters/psycopg/litestar/store.py | 298 +++++++++++---------- 6 files changed, 255 insertions(+), 226 deletions(-) diff --git a/sqlspec/adapters/psycopg/adk/store.py b/sqlspec/adapters/psycopg/adk/store.py index d2cdcba31..ee47a037e 100644 --- a/sqlspec/adapters/psycopg/adk/store.py +++ b/sqlspec/adapters/psycopg/adk/store.py @@ -7,6 +7,7 @@ from sqlspec.adapters.psycopg._typing import psycopg_dict_row as dict_row from sqlspec.adapters.psycopg._typing import psycopg_errors as errors from sqlspec.adapters.psycopg._typing import psycopg_sql as pg_sql +from sqlspec.adapters.psycopg.core import build_postgres_storage_params_sql from sqlspec.config import ADKConfig from sqlspec.extensions.adk import ( BaseAsyncADKStore, @@ -706,7 +707,7 @@ async def _sessions_table_ddl(self) -> str: self._session_table, self._session_table, self._session_table, - _postgres_table_options(_adk_config(self._config)), + build_postgres_storage_params_sql(_adk_config(self._config), default_fillfactor=80, section="adk"), ) async def _events_table_ddl(self) -> str: @@ -723,19 +724,23 @@ async def _events_table_ddl(self) -> str: self._events_table, covering_columns, generated_indexes, - _postgres_table_options(adk_config, include_autovacuum=True), + build_postgres_storage_params_sql( + adk_config, include_autovacuum=True, default_fillfactor=80, section="adk" + ), self._events_table, self._events_table, ) async def _app_states_table_ddl(self) -> str: return _ADK_APP_STATES_TABLE_DDL_TEMPLATE.format( - self._app_state_table, _postgres_table_options(_adk_config(self._config)) + self._app_state_table, + build_postgres_storage_params_sql(_adk_config(self._config), default_fillfactor=80, section="adk"), ) async def _user_states_table_ddl(self) -> str: return _ADK_USER_STATES_TABLE_DDL_TEMPLATE.format( - self._user_state_table, _postgres_table_options(_adk_config(self._config)) + self._user_state_table, + build_postgres_storage_params_sql(_adk_config(self._config), default_fillfactor=80, section="adk"), ) async def _metadata_table_ddl(self) -> str: @@ -1256,7 +1261,7 @@ def _sessions_table_ddl(self) -> str: self._session_table, self._session_table, self._session_table, - _postgres_table_options(_adk_config(self._config)), + build_postgres_storage_params_sql(_adk_config(self._config), default_fillfactor=80, section="adk"), ) def _events_table_ddl(self) -> str: @@ -1273,19 +1278,23 @@ def _events_table_ddl(self) -> str: self._events_table, covering_columns, generated_indexes, - _postgres_table_options(adk_config, include_autovacuum=True), + build_postgres_storage_params_sql( + adk_config, include_autovacuum=True, default_fillfactor=80, section="adk" + ), self._events_table, self._events_table, ) def _app_states_table_ddl(self) -> str: return _ADK_APP_STATES_TABLE_DDL_TEMPLATE.format( - self._app_state_table, _postgres_table_options(_adk_config(self._config)) + self._app_state_table, + build_postgres_storage_params_sql(_adk_config(self._config), default_fillfactor=80, section="adk"), ) def _user_states_table_ddl(self) -> str: return _ADK_USER_STATES_TABLE_DDL_TEMPLATE.format( - self._user_state_table, _postgres_table_options(_adk_config(self._config)) + self._user_state_table, + build_postgres_storage_params_sql(_adk_config(self._config), default_fillfactor=80, section="adk"), ) def _metadata_table_ddl(self) -> str: @@ -2022,33 +2031,6 @@ def _adk_config(config: Any) -> PsycopgADKConfig: return cast("PsycopgADKConfig", adk_config) -def _postgres_table_options(adk_config: PsycopgADKConfig, *, include_autovacuum: bool = False) -> str: - options = [_postgres_fillfactor_option(adk_config)] - if include_autovacuum: - options.extend(_postgres_autovacuum_options(adk_config)) - return f" WITH ({', '.join(options)})" - - -def _postgres_fillfactor_option(adk_config: PsycopgADKConfig) -> str: - value = adk_config.get("fillfactor", 80) - if not isinstance(value, int) or isinstance(value, bool) or value not in range(10, 101): - msg = "extension_config['adk']['fillfactor'] must be an integer from 10 to 100" - raise ValueError(msg) - return f"fillfactor = {value}" - - -def _postgres_autovacuum_options(adk_config: PsycopgADKConfig) -> "list[str]": - options: list[str] = [] - for key in ("autovacuum_vacuum_scale_factor", "autovacuum_analyze_scale_factor"): - value = adk_config.get(key) - if value is not None: - if not isinstance(value, (int, float)) or isinstance(value, bool) or not 0 <= float(value) <= 1: - msg = f"extension_config['adk']['{key}'] must be a number from 0 to 1" - raise ValueError(msg) - options.append(f"{key} = {float(value):g}") - return options - - def _postgres_memory_indexes( memory_table: str, use_fts: bool, diff --git a/sqlspec/adapters/psycopg/core.py b/sqlspec/adapters/psycopg/core.py index 9e2c58724..9ee516e1a 100644 --- a/sqlspec/adapters/psycopg/core.py +++ b/sqlspec/adapters/psycopg/core.py @@ -70,6 +70,7 @@ "build_copy_from_command", "build_pipeline_execution_result", "build_postgres_extension_probe_names", + "build_postgres_storage_params_sql", "build_profile", "build_statement_config", "build_truncate_command", @@ -136,6 +137,39 @@ def build_truncate_command(table: str) -> "PsycopgComposed": return PsycopgSQL("TRUNCATE TABLE {}").format(_compose_table_identifier(table)) +def build_postgres_storage_params_sql( + settings: "dict[str, object] | Mapping[str, object] | None", + *, + include_autovacuum: bool = False, + default_fillfactor: int | None = None, + default_vacuum_scale: float | None = None, + default_analyze_scale: float | None = None, + section: str = "storage", +) -> str: + """Format PostgreSQL table storage parameters (WITH fillfactor = ...) clause.""" + options: list[str] = [] + config_dict = dict(settings) if settings else {} + fillfactor = config_dict.get("fillfactor", default_fillfactor) + if fillfactor is not None: + if not isinstance(fillfactor, int) or isinstance(fillfactor, bool) or fillfactor not in range(10, 101): + msg = f"extension_config['{section}']['fillfactor'] must be an integer from 10 to 100" + raise ValueError(msg) + options.append(f"fillfactor = {fillfactor}") + if include_autovacuum: + scales = ( + ("autovacuum_vacuum_scale_factor", default_vacuum_scale), + ("autovacuum_analyze_scale_factor", default_analyze_scale), + ) + for key, default_val in scales: + value = config_dict.get(key, default_val) + if value is not None: + if not isinstance(value, (int, float)) or isinstance(value, bool) or not 0 <= float(value) <= 1: + msg = f"extension_config['{section}']['{key}'] must be a number from 0 to 1" + raise ValueError(msg) + options.append(f"{key} = {float(value):g}") + return f" WITH ({', '.join(options)})" if options else "" + + def build_profile() -> "DriverParameterProfile": """Create the psycopg driver parameter profile.""" diff --git a/sqlspec/adapters/psycopg/events/_hub.py b/sqlspec/adapters/psycopg/events/_hub.py index d0634504e..1fbdced73 100644 --- a/sqlspec/adapters/psycopg/events/_hub.py +++ b/sqlspec/adapters/psycopg/events/_hub.py @@ -14,6 +14,7 @@ import logging import queue as stdlib_queue import threading +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, cast from weakref import WeakKeyDictionary @@ -182,7 +183,7 @@ def _resolve_listener_config(self) -> "PsycopgAsyncConfig": ) ).lower() raw_direct_dsn = driver_features.get("direct_dsn") - direct_dsn = raw_direct_dsn if isinstance(raw_direct_dsn, str) and raw_direct_dsn else None + direct_dsn = raw_direct_dsn if isinstance(raw_direct_dsn, (str, Mapping)) and raw_direct_dsn else None if pooler_mode in {"transaction", "statement"} and not direct_dsn: logger.warning( "PostgreSQL LISTEN/NOTIFY cannot run under pooler_mode=%r without configuring driver_features['direct_dsn']", @@ -449,7 +450,7 @@ def _resolve_listener_config(self) -> "PsycopgSyncConfig": ) ).lower() raw_direct_dsn = driver_features.get("direct_dsn") - direct_dsn = raw_direct_dsn if isinstance(raw_direct_dsn, str) and raw_direct_dsn else None + direct_dsn = raw_direct_dsn if isinstance(raw_direct_dsn, (str, Mapping)) and raw_direct_dsn else None if pooler_mode in {"transaction", "statement"} and not direct_dsn: logger.warning( "PostgreSQL LISTEN/NOTIFY cannot run under pooler_mode=%r without configuring driver_features['direct_dsn']", diff --git a/sqlspec/adapters/psycopg/events/backend.py b/sqlspec/adapters/psycopg/events/backend.py index 7a683ebb6..a7dbafb9a 100644 --- a/sqlspec/adapters/psycopg/events/backend.py +++ b/sqlspec/adapters/psycopg/events/backend.py @@ -34,7 +34,7 @@ "create_event_backend", ) -logger = get_logger("sqlspec.events.psycopg") +logger = get_logger("sqlspec.adapters.psycopg.events.backend") _MIN_LISTENER_POOL_SIZE = 2 _MAX_SEEN_MARKERS = 1_024 @@ -363,14 +363,15 @@ def dequeue(self, channel: str, poll_interval: float) -> EventMessage | None: return self._dequeue_reconciled(channel) self._queue._reset_empty_poll_delay(channel) self._runtime.increment_metric("events.listener.wakeup") - batch_size = _extract_batch_size(payload) - is_new, should_drain = self._marker_state.register(channel, _extract_marker_id(payload), batch_size) + parsed_payload = _parse_notify_payload(payload) + batch_size = _extract_batch_size(parsed_payload) + is_new, should_drain = self._marker_state.register(channel, _extract_marker_id(parsed_payload), batch_size) if is_new and batch_size > 1: self._runtime.increment_metric("events.marker.coalesced", batch_size - 1) if not should_drain: self._runtime.increment_metric("events.marker.duplicate") continue - event_id = _extract_event_id(payload) + event_id = _extract_event_id(parsed_payload) if event_id is not None: event = self._queue.dequeue_by_event_id(event_id) if event is not None: @@ -534,14 +535,15 @@ async def dequeue(self, channel: str, poll_interval: float) -> EventMessage | No return await self._dequeue_reconciled(channel) self._queue._reset_empty_poll_delay(channel) self._runtime.increment_metric("events.listener.wakeup") - batch_size = _extract_batch_size(payload) - is_new, should_drain = self._marker_state.register(channel, _extract_marker_id(payload), batch_size) + parsed_payload = _parse_notify_payload(payload) + batch_size = _extract_batch_size(parsed_payload) + is_new, should_drain = self._marker_state.register(channel, _extract_marker_id(parsed_payload), batch_size) if is_new and batch_size > 1: self._runtime.increment_metric("events.marker.coalesced", batch_size - 1) if not should_drain: self._runtime.increment_metric("events.marker.duplicate") continue - event_id = _extract_event_id(payload) + event_id = _extract_event_id(parsed_payload) if event_id is not None: event = await self._queue.dequeue_by_event_id(event_id) if event is not None: @@ -685,34 +687,40 @@ def _validate_listener_pool_capacity(config: "PsycopgAsyncConfig | PsycopgSyncCo raise ImproperConfigurationError(msg) -def _extract_event_id(payload: "str | None") -> "str | None": - if not payload: +def _parse_notify_payload(payload: "dict[str, Any] | str | None") -> "dict[str, Any] | None": + if isinstance(payload, dict): + return payload + if not payload or not isinstance(payload, str): return None - raw = from_json(payload) - if isinstance(raw, dict): - event_id = raw.get("event_id") - return event_id if isinstance(event_id, str) else None - return None + try: + raw = from_json(payload) + return raw if isinstance(raw, dict) else None + except (TypeError, ValueError): + return None + + +def _extract_event_id(payload: "dict[str, Any] | str | None") -> "str | None": + parsed = payload if isinstance(payload, dict) else _parse_notify_payload(payload) + if parsed is None: + return None + event_id = parsed.get("event_id") + return event_id if isinstance(event_id, str) else None -def _extract_batch_size(payload: "str | None") -> int: - if not payload: +def _extract_batch_size(payload: "dict[str, Any] | str | None") -> int: + parsed = payload if isinstance(payload, dict) else _parse_notify_payload(payload) + if parsed is None: return 0 - raw = from_json(payload) - if isinstance(raw, dict): - batch_size = raw.get("batch_size") - return batch_size if isinstance(batch_size, int) else 0 - return 0 + batch_size = parsed.get("batch_size") + return batch_size if isinstance(batch_size, int) else 0 -def _extract_marker_id(payload: "str | None") -> "str | None": - if not payload: +def _extract_marker_id(payload: "dict[str, Any] | str | None") -> "str | None": + parsed = payload if isinstance(payload, dict) else _parse_notify_payload(payload) + if parsed is None: return None - raw = from_json(payload) - if isinstance(raw, dict): - marker_id = raw.get("marker_id", raw.get("event_id")) - return marker_id if isinstance(marker_id, str) else None - return None + marker_id = parsed.get("marker_id", parsed.get("event_id")) + return marker_id if isinstance(marker_id, str) else None def _record_dequeue_result(runtime: Any, event: "EventMessage | None") -> None: diff --git a/sqlspec/adapters/psycopg/events/store.py b/sqlspec/adapters/psycopg/events/store.py index 6a3b327f2..36f1d720a 100644 --- a/sqlspec/adapters/psycopg/events/store.py +++ b/sqlspec/adapters/psycopg/events/store.py @@ -3,6 +3,7 @@ from typing import NotRequired from sqlspec.adapters.psycopg.config import PsycopgAsyncConfig, PsycopgSyncConfig +from sqlspec.adapters.psycopg.core import build_postgres_storage_params_sql from sqlspec.config import EventsConfig from sqlspec.extensions.events import BaseEventQueueStore @@ -50,7 +51,7 @@ def _column_types(self) -> "tuple[str, str, str]": def _table_clause(self) -> str: """Return explicitly configured PostgreSQL table tuning options.""" - return _postgres_table_options(self.settings) + return build_postgres_storage_params_sql(self.settings, include_autovacuum=True, section="events") class PsycopgAsyncEventQueueStore(BaseEventQueueStore[PsycopgAsyncConfig]): @@ -78,27 +79,4 @@ def _column_types(self) -> "tuple[str, str, str]": def _table_clause(self) -> str: """Return explicitly configured PostgreSQL table tuning options.""" - return _postgres_table_options(self.settings) - - -def _postgres_table_options(settings: "dict[str, object]") -> str: - options = _postgres_option_values(settings) - return f" WITH ({', '.join(options)})" if options else "" - - -def _postgres_option_values(settings: "dict[str, object]") -> "list[str]": - options: list[str] = [] - fillfactor = settings.get("fillfactor") - if fillfactor is not None: - if not isinstance(fillfactor, int) or isinstance(fillfactor, bool) or fillfactor not in range(10, 101): - msg = "extension_config['events']['fillfactor'] must be an integer from 10 to 100" - raise ValueError(msg) - options.append(f"fillfactor = {fillfactor}") - for key in ("autovacuum_vacuum_scale_factor", "autovacuum_analyze_scale_factor"): - value = settings.get(key) - if value is not None: - if not isinstance(value, (int, float)) or isinstance(value, bool) or not 0 <= float(value) <= 1: - msg = f"extension_config['events']['{key}'] must be a number from 0 to 1" - raise ValueError(msg) - options.append(f"{key} = {float(value):g}") - return options + return build_postgres_storage_params_sql(self.settings, include_autovacuum=True, section="events") diff --git a/sqlspec/adapters/psycopg/litestar/store.py b/sqlspec/adapters/psycopg/litestar/store.py index 3d7a2e7ab..865ab6fd6 100644 --- a/sqlspec/adapters/psycopg/litestar/store.py +++ b/sqlspec/adapters/psycopg/litestar/store.py @@ -7,6 +7,8 @@ from typing import TYPE_CHECKING, Any, NotRequired, cast from sqlspec.adapters.psycopg._typing import psycopg_dict_row as dict_row +from sqlspec.adapters.psycopg._typing import psycopg_errors as errors +from sqlspec.adapters.psycopg.core import build_postgres_storage_params_sql from sqlspec.config import LitestarConfig from sqlspec.extensions.litestar.store import BaseSQLSpecStore from sqlspec.utils.sync_tools import async_ @@ -92,27 +94,30 @@ async def get(self, key: str, renew_for: "int | timedelta | None" = None) -> "by AND (expires_at IS NULL OR expires_at > CURRENT_TIMESTAMP) """ - conn_context = self._config.provide_connection() - async with conn_context as conn: - async with conn.cursor(row_factory=dict_row) as cur: - await cur.execute(cast("Any", sql), (key,)) - row = await cur.fetchone() - - if row is None: - return None - - if renew_for is not None and row["expires_at"] is not None: - new_expires_at = self._calculate_expires_at(renew_for) - if new_expires_at is not None: - update_sql = f""" - UPDATE {self._table_name} - SET expires_at = %s, updated_at = CURRENT_TIMESTAMP - WHERE session_id = %s - """ - await conn.execute(cast("Any", update_sql), (new_expires_at, key)) - await conn.commit() - - return bytes(row["data"]) + try: + conn_context = self._config.provide_connection() + async with conn_context as conn: + async with conn.cursor(row_factory=dict_row) as cur: + await cur.execute(cast("Any", sql), (key,)) + row = await cur.fetchone() + + if row is None: + return None + + if renew_for is not None and row["expires_at"] is not None: + new_expires_at = self._calculate_expires_at(renew_for) + if new_expires_at is not None: + update_sql = f""" + UPDATE {self._table_name} + SET expires_at = %s, updated_at = CURRENT_TIMESTAMP + WHERE session_id = %s + """ + await conn.execute(cast("Any", update_sql), (new_expires_at, key)) + await conn.commit() + + return bytes(row["data"]) + except errors.UndefinedTable: + return None async def set(self, key: str, value: "str | bytes", expires_in: "int | timedelta | None" = None) -> None: """Store a session value. @@ -148,20 +153,26 @@ async def delete(self, key: str) -> None: """ sql = f"DELETE FROM {self._table_name} WHERE session_id = %s" - conn_context = self._config.provide_connection() - async with conn_context as conn: - await conn.execute(cast("Any", sql), (key,)) - await conn.commit() + try: + conn_context = self._config.provide_connection() + async with conn_context as conn: + await conn.execute(cast("Any", sql), (key,)) + await conn.commit() + except errors.UndefinedTable: + return async def delete_all(self) -> None: """Delete all sessions from the store.""" sql = f"DELETE FROM {self._table_name}" - conn_context = self._config.provide_connection() - async with conn_context as conn: - await conn.execute(cast("Any", sql)) - await conn.commit() - self._log_delete_all() + try: + conn_context = self._config.provide_connection() + async with conn_context as conn: + await conn.execute(cast("Any", sql)) + await conn.commit() + self._log_delete_all() + except errors.UndefinedTable: + return async def exists(self, key: str) -> bool: """Check if a session key exists and is not expired. @@ -178,11 +189,14 @@ async def exists(self, key: str) -> bool: AND (expires_at IS NULL OR expires_at > CURRENT_TIMESTAMP) """ - conn_context = self._config.provide_connection() - async with conn_context as conn, conn.cursor() as cur: - await cur.execute(cast("Any", sql), (key,)) - result = await cur.fetchone() - return result is not None + try: + conn_context = self._config.provide_connection() + async with conn_context as conn, conn.cursor() as cur: + await cur.execute(cast("Any", sql), (key,)) + result = await cur.fetchone() + return result is not None + except errors.UndefinedTable: + return False async def expires_in(self, key: str) -> "int | None": """Get the time in seconds until the session expires. @@ -198,22 +212,25 @@ async def expires_in(self, key: str) -> "int | None": WHERE session_id = %s """ - conn_context = self._config.provide_connection() - async with conn_context as conn: - async with conn.cursor(row_factory=dict_row) as cur: - await cur.execute(cast("Any", sql), (key,)) - row = await cur.fetchone() + try: + conn_context = self._config.provide_connection() + async with conn_context as conn: + async with conn.cursor(row_factory=dict_row) as cur: + await cur.execute(cast("Any", sql), (key,)) + row = await cur.fetchone() - if row is None or row["expires_at"] is None: - return None + if row is None or row["expires_at"] is None: + return None - expires_at = row["expires_at"] - now = datetime.now(UTC) - if expires_at <= now: - return 0 + expires_at = row["expires_at"] + now = datetime.now(UTC) + if expires_at <= now: + return 0 - delta = expires_at - now - return int(delta.total_seconds()) + delta = expires_at - now + return int(delta.total_seconds()) + except errors.UndefinedTable: + return None async def delete_expired(self) -> int: """Delete all expired sessions. @@ -223,14 +240,17 @@ async def delete_expired(self) -> int: """ sql = f"DELETE FROM {self._table_name} WHERE expires_at <= CURRENT_TIMESTAMP" - conn_context = self._config.provide_connection() - async with conn_context as conn, conn.cursor() as cur: - await cur.execute(cast("Any", sql)) - await conn.commit() - count = cur.rowcount if cur.rowcount and cur.rowcount > 0 else 0 - if count > 0: - self._log_delete_expired(count) - return count + try: + conn_context = self._config.provide_connection() + async with conn_context as conn, conn.cursor() as cur: + await cur.execute(cast("Any", sql)) + await conn.commit() + count = cur.rowcount if cur.rowcount and cur.rowcount > 0 else 0 + if count > 0: + self._log_delete_expired(count) + return count + except errors.UndefinedTable: + return 0 def _table_ddl(self) -> str: """Get PostgreSQL CREATE TABLE SQL with optimized schema. @@ -238,7 +258,14 @@ def _table_ddl(self) -> str: Returns: SQL statement to create the sessions table with proper indexes. """ - fillfactor, vacuum_scale, analyze_scale = _postgres_litestar_tuning(self._config) + tuning_clause = build_postgres_storage_params_sql( + self._config.extension_config.get("litestar", {}), + default_fillfactor=80, + default_vacuum_scale=0.05, + default_analyze_scale=0.02, + include_autovacuum=True, + section="litestar", + ) return f""" CREATE TABLE IF NOT EXISTS {self._table_name} ( session_id TEXT PRIMARY KEY, @@ -246,15 +273,10 @@ def _table_ddl(self) -> str: expires_at TIMESTAMPTZ, created_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP, updated_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP - ) WITH (fillfactor = {fillfactor}); + ){tuning_clause}; CREATE INDEX IF NOT EXISTS idx_{self._table_name}_expires_at ON {self._table_name}(expires_at) WHERE expires_at IS NOT NULL; - - ALTER TABLE {self._table_name} SET ( - autovacuum_vacuum_scale_factor = {vacuum_scale:g}, - autovacuum_analyze_scale_factor = {analyze_scale:g} - ); """ def _drop_table_sql(self) -> "list[str]": @@ -376,7 +398,14 @@ def _table_ddl(self) -> str: Returns: SQL statement to create the sessions table with proper indexes. """ - fillfactor, vacuum_scale, analyze_scale = _postgres_litestar_tuning(self._config) + tuning_clause = build_postgres_storage_params_sql( + self._config.extension_config.get("litestar", {}), + default_fillfactor=80, + default_vacuum_scale=0.05, + default_analyze_scale=0.02, + include_autovacuum=True, + section="litestar", + ) return f""" CREATE TABLE IF NOT EXISTS {self._table_name} ( session_id TEXT PRIMARY KEY, @@ -384,15 +413,10 @@ def _table_ddl(self) -> str: expires_at TIMESTAMPTZ, created_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP, updated_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP - ) WITH (fillfactor = {fillfactor}); + ){tuning_clause}; CREATE INDEX IF NOT EXISTS idx_{self._table_name}_expires_at ON {self._table_name}(expires_at) WHERE expires_at IS NOT NULL; - - ALTER TABLE {self._table_name} SET ( - autovacuum_vacuum_scale_factor = {vacuum_scale:g}, - autovacuum_analyze_scale_factor = {analyze_scale:g} - ); """ def _drop_table_sql(self) -> "list[str]": @@ -419,26 +443,29 @@ def _get(self, key: str, renew_for: "int | timedelta | None" = None) -> "bytes | AND (expires_at IS NULL OR expires_at > CURRENT_TIMESTAMP) """ - with self._config.provide_connection() as conn: - with conn.cursor(row_factory=dict_row) as cur: - cur.execute(cast("Any", sql), (key,)) - row = cur.fetchone() - - if row is None: - return None - - if renew_for is not None and row["expires_at"] is not None: - new_expires_at = self._calculate_expires_at(renew_for) - if new_expires_at is not None: - update_sql = f""" - UPDATE {self._table_name} - SET expires_at = %s, updated_at = CURRENT_TIMESTAMP - WHERE session_id = %s - """ - conn.execute(cast("Any", update_sql), (new_expires_at, key)) - conn.commit() - - return bytes(row["data"]) + try: + with self._config.provide_connection() as conn: + with conn.cursor(row_factory=dict_row) as cur: + cur.execute(cast("Any", sql), (key,)) + row = cur.fetchone() + + if row is None: + return None + + if renew_for is not None and row["expires_at"] is not None: + new_expires_at = self._calculate_expires_at(renew_for) + if new_expires_at is not None: + update_sql = f""" + UPDATE {self._table_name} + SET expires_at = %s, updated_at = CURRENT_TIMESTAMP + WHERE session_id = %s + """ + conn.execute(cast("Any", update_sql), (new_expires_at, key)) + conn.commit() + + return bytes(row["data"]) + except errors.UndefinedTable: + return None def _set(self, key: str, value: "str | bytes", expires_in: "int | timedelta | None" = None) -> None: """Synchronous implementation of set.""" @@ -463,18 +490,24 @@ def _delete(self, key: str) -> None: """Synchronous implementation of delete.""" sql = f"DELETE FROM {self._table_name} WHERE session_id = %s" - with self._config.provide_connection() as conn: - conn.execute(cast("Any", sql), (key,)) - conn.commit() + try: + with self._config.provide_connection() as conn: + conn.execute(cast("Any", sql), (key,)) + conn.commit() + except errors.UndefinedTable: + return def _delete_all(self) -> None: """Synchronous implementation of delete_all.""" sql = f"DELETE FROM {self._table_name}" - with self._config.provide_connection() as conn: - conn.execute(cast("Any", sql)) - conn.commit() - self._log_delete_all() + try: + with self._config.provide_connection() as conn: + conn.execute(cast("Any", sql)) + conn.commit() + self._log_delete_all() + except errors.UndefinedTable: + return def _exists(self, key: str) -> bool: """Synchronous implementation of exists.""" @@ -484,10 +517,13 @@ def _exists(self, key: str) -> bool: AND (expires_at IS NULL OR expires_at > CURRENT_TIMESTAMP) """ - with self._config.provide_connection() as conn, conn.cursor() as cur: - cur.execute(cast("Any", sql), (key,)) - result = cur.fetchone() - return result is not None + try: + with self._config.provide_connection() as conn, conn.cursor() as cur: + cur.execute(cast("Any", sql), (key,)) + result = cur.fetchone() + return result is not None + except errors.UndefinedTable: + return False def _expires_in(self, key: str) -> "int | None": """Synchronous implementation of expires_in.""" @@ -496,47 +532,37 @@ def _expires_in(self, key: str) -> "int | None": WHERE session_id = %s """ - with self._config.provide_connection() as conn: - with conn.cursor(row_factory=dict_row) as cur: - cur.execute(cast("Any", sql), (key,)) - row = cur.fetchone() + try: + with self._config.provide_connection() as conn: + with conn.cursor(row_factory=dict_row) as cur: + cur.execute(cast("Any", sql), (key,)) + row = cur.fetchone() - if row is None or row["expires_at"] is None: - return None + if row is None or row["expires_at"] is None: + return None - expires_at = row["expires_at"] - now = datetime.now(UTC) + expires_at = row["expires_at"] + now = datetime.now(UTC) - if expires_at <= now: - return 0 + if expires_at <= now: + return 0 - delta = expires_at - now - return int(delta.total_seconds()) + delta = expires_at - now + return int(delta.total_seconds()) + except errors.UndefinedTable: + return None def _delete_expired(self) -> int: """Synchronous implementation of delete_expired.""" sql = f"DELETE FROM {self._table_name} WHERE expires_at <= CURRENT_TIMESTAMP" - with self._config.provide_connection() as conn, conn.cursor() as cur: - cur.execute(cast("Any", sql)) - conn.commit() - count = cur.rowcount if cur.rowcount and cur.rowcount > 0 else 0 - if count > 0: - self._log_delete_expired(count) - return count - - -def _postgres_litestar_tuning(config: Any) -> "tuple[int, float, float]": - settings = cast("dict[str, Any]", config.extension_config.get("litestar", {})) - fillfactor = settings.get("fillfactor", 80) - if not isinstance(fillfactor, int) or isinstance(fillfactor, bool) or fillfactor not in range(10, 101): - msg = "extension_config['litestar']['fillfactor'] must be an integer from 10 to 100" - raise ValueError(msg) - scales: list[float] = [] - for key, default in (("autovacuum_vacuum_scale_factor", 0.05), ("autovacuum_analyze_scale_factor", 0.02)): - value = settings.get(key, default) - if not isinstance(value, (int, float)) or isinstance(value, bool) or not 0 <= float(value) <= 1: - msg = f"extension_config['litestar']['{key}'] must be a number from 0 to 1" - raise ValueError(msg) - scales.append(float(value)) - return fillfactor, scales[0], scales[1] + try: + with self._config.provide_connection() as conn, conn.cursor() as cur: + cur.execute(cast("Any", sql)) + conn.commit() + count = cur.rowcount if cur.rowcount and cur.rowcount > 0 else 0 + if count > 0: + self._log_delete_expired(count) + return count + except errors.UndefinedTable: + return 0 From 093a95d8a1fd0b1fe5f62fc339e6c723027c6b9c Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 01:05:47 +0000 Subject: [PATCH 093/174] fix(psycopg): correct data dictionary and catalog queries --- sqlspec/adapters/psycopg/data_dictionary.py | 52 ++++++--------------- 1 file changed, 13 insertions(+), 39 deletions(-) diff --git a/sqlspec/adapters/psycopg/data_dictionary.py b/sqlspec/adapters/psycopg/data_dictionary.py index e0c245ea0..a4413254f 100644 --- a/sqlspec/adapters/psycopg/data_dictionary.py +++ b/sqlspec/adapters/psycopg/data_dictionary.py @@ -28,7 +28,12 @@ system_metadata_gated_result, unsupported_system_metadata_capability, ) -from sqlspec.data_dictionary.dialects.postgres import resolve_postgres_json_type +from sqlspec.data_dictionary.dialects.postgres import ( + POSTGRES_METADATA_DOMAINS, + POSTGRES_SUPPORTED_DOMAINS, + POSTGRES_SYSTEM_METADATA_QUERIES, + resolve_postgres_json_type, +) from sqlspec.driver import AsyncDataDictionaryBase, SyncDataDictionaryBase if TYPE_CHECKING: @@ -38,37 +43,6 @@ __all__ = ("PsycopgAsyncDataDictionary", "PsycopgSyncDataDictionary") -_POSTGRES_METADATA_DOMAINS: tuple[str, ...] = ( - "schemas", - "objects", - "tables", - "columns", - "constraints", - "indexes", - "views", - "materialized_views", - "sequences", - "types", - "routines", - "triggers", - "comments", - "privileges", - "dependencies", - "extensions", - "partitions", - "ddl", - "system", -) -_POSTGRES_SUPPORTED_DOMAINS = frozenset(d for d in _POSTGRES_METADATA_DOMAINS if d != "system") -_POSTGRES_SYSTEM_METADATA_QUERIES = { - "settings": "settings", - "statement_history": "pg_stat_statements", - "pg_stat_statements": "pg_stat_statements", - "table_statistics": "table_stats", - "table_stats": "table_stats", -} - - @mypyc_attr(allow_interpreted_subclasses=True, native_class=False) class PsycopgSyncDataDictionary(SyncDataDictionaryBase): """PostgreSQL-specific sync data dictionary.""" @@ -89,7 +63,7 @@ def get_system_metadata_capabilities( ) -> tuple[SystemMetadataCapability, ...]: """Get PostgreSQL opt-in system metadata capability disclosures.""" _ = driver - requested_domains = tuple(_POSTGRES_SYSTEM_METADATA_QUERIES) if domains is None else tuple(domains) + requested_domains = tuple(POSTGRES_SYSTEM_METADATA_QUERIES) if domains is None else tuple(domains) return tuple(_postgres_system_metadata_capability(domain) for domain in requested_domains) def get_schemas(self, driver: "PsycopgSyncDriver") -> MetadataResult: @@ -163,7 +137,7 @@ def get_system_metadata( if gate_result.capability.support != MetadataSupport.SUPPORTED: return gate_result - query_name = _POSTGRES_SYSTEM_METADATA_QUERIES.get(metadata_request.domain) + query_name = POSTGRES_SYSTEM_METADATA_QUERIES.get(metadata_request.domain) if query_name is None: return gate_result parameters: dict[str, object] = {} @@ -367,7 +341,7 @@ async def get_system_metadata_capabilities( ) -> tuple[SystemMetadataCapability, ...]: """Get PostgreSQL opt-in system metadata capability disclosures.""" _ = driver - requested_domains = tuple(_POSTGRES_SYSTEM_METADATA_QUERIES) if domains is None else tuple(domains) + requested_domains = tuple(POSTGRES_SYSTEM_METADATA_QUERIES) if domains is None else tuple(domains) return tuple(_postgres_system_metadata_capability(domain) for domain in requested_domains) async def get_schemas(self, driver: "PsycopgAsyncDriver") -> MetadataResult: @@ -443,7 +417,7 @@ async def get_system_metadata( if gate_result.capability.support != MetadataSupport.SUPPORTED: return gate_result - query_name = _POSTGRES_SYSTEM_METADATA_QUERIES.get(metadata_request.domain) + query_name = POSTGRES_SYSTEM_METADATA_QUERIES.get(metadata_request.domain) if query_name is None: return gate_result parameters: dict[str, object] = {} @@ -637,7 +611,7 @@ def _postgres_metadata_capability(domain: str) -> MetadataCapability: risks=(MetadataRisk.EXPENSIVE, MetadataRisk.PRIVILEGED), warnings=("System metadata is opt-in and disabled by default.",), ) - if domain in _POSTGRES_SUPPORTED_DOMAINS: + if domain in POSTGRES_SUPPORTED_DOMAINS: return MetadataCapability( domain=domain, support=MetadataSupport.SUPPORTED, @@ -648,7 +622,7 @@ def _postgres_metadata_capability(domain: str) -> MetadataCapability: def _postgres_metadata_profile(adapter: str, domains: Sequence[str] | None) -> MetadataCapabilityProfile: - requested_domains = _POSTGRES_METADATA_DOMAINS if domains is None else tuple(domains) + requested_domains = POSTGRES_METADATA_DOMAINS if domains is None else tuple(domains) return MetadataCapabilityProfile( "postgres", adapter=adapter, @@ -707,7 +681,7 @@ def _ddl_result_from_rows( def _postgres_system_metadata_capability(domain: str) -> SystemMetadataCapability: - if domain not in _POSTGRES_SYSTEM_METADATA_QUERIES: + if domain not in POSTGRES_SYSTEM_METADATA_QUERIES: return unsupported_system_metadata_capability(domain) return SystemMetadataCapability( domain, From cfedb29724ba441cadd782bb696805fd5a7d7eac Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 01:05:47 +0000 Subject: [PATCH 094/174] perf(psycopg): optimize cache hit wire escaping, transaction status check, and stream isolation Pre-escape wire SQL on query cache creation and reuse wire SQL on cache hit to eliminate O(N) scan overhead and support queries with literal percents. Avoid per-call set allocation in transaction checks with module-level open transaction statuses, isolate stream source and async load_from_arrow execution with driver exception handlers, and replace local deletion with assignment in mapped exception construction. --- sqlspec/adapters/psycopg/core.py | 42 +++--- sqlspec/adapters/psycopg/driver.py | 62 ++++----- .../unit/adapters/test_psycopg/test_driver.py | 127 +++++++++++++++--- 3 files changed, 160 insertions(+), 71 deletions(-) diff --git a/sqlspec/adapters/psycopg/core.py b/sqlspec/adapters/psycopg/core.py index 9ee516e1a..db635e0bb 100644 --- a/sqlspec/adapters/psycopg/core.py +++ b/sqlspec/adapters/psycopg/core.py @@ -56,6 +56,7 @@ from sqlspec.core import ParameterStyleConfig, StackOperation __all__ = ( + "OPEN_TRANSACTION_STATUSES", "TRANSACTION_STATUS_ACTIVE", "TRANSACTION_STATUS_IDLE", "TRANSACTION_STATUS_INERROR", @@ -94,6 +95,11 @@ TRANSACTION_STATUS_INTRANS = 2 TRANSACTION_STATUS_INERROR = 3 TRANSACTION_STATUS_UNKNOWN = 4 +OPEN_TRANSACTION_STATUSES: Final[tuple[int, ...]] = ( + TRANSACTION_STATUS_ACTIVE, + TRANSACTION_STATUS_INTRANS, + TRANSACTION_STATUS_INERROR, +) class PreparedStackOperation(NamedTuple): @@ -306,27 +312,12 @@ def __init__(self, driver: Any, sql: str, parameters: Any, chunk_size: int) -> N def start(self) -> None: handler = self._driver.handle_database_exceptions() - with handler: - transaction = self._driver.connection.transaction() - transaction.__enter__() - self._transaction = transaction - try: - cursor = self._driver.connection.cursor(name=f"sqlspec_stream_{uuid4().hex}") - cursor.itersize = self._chunk_size - execute_with_optional_parameters(cursor, self._sql, self._parameters) - self._cursor = cursor - except BaseException as exc: - self._transaction = None - with contextlib.suppress(Exception): - transaction.__exit__(type(exc), exc, exc.__traceback__) - raise + self._driver._run_with_exception_handler(handler, self._start) self._driver._check_pending_exception(handler) def fetch_chunk(self) -> "list[dict[str, Any]]": handler = self._driver.handle_database_exceptions() - rows: list[Any] = [] - with handler: - rows = self._cursor.fetchmany(self._chunk_size) + rows = self._driver._run_with_exception_handler(handler, self._cursor.fetchmany, self._chunk_size) self._driver._check_pending_exception(handler) if not rows: return [] @@ -349,6 +340,21 @@ def close(self, error: bool = False) -> None: else: transaction.__exit__(None, None, None) + def _start(self) -> None: + transaction = self._driver.connection.transaction() + transaction.__enter__() + self._transaction = transaction + try: + cursor = self._driver.connection.cursor(name=f"sqlspec_stream_{uuid4().hex}") + cursor.itersize = self._chunk_size + execute_with_optional_parameters(cursor, self._sql, self._parameters) + self._cursor = cursor + except BaseException as exc: + self._transaction = None + with contextlib.suppress(Exception): + transaction.__exit__(type(exc), exc, exc.__traceback__) + raise + class PsycopgAsyncStreamSource: """Compiled async chunk source streaming dict rows from a psycopg server-side named cursor. @@ -606,7 +612,7 @@ def create_mapped_exception(error: Any, *, logger: Any | None = None) -> SQLSpec Returns: A SQLSpec exception that wraps the original error """ - del logger + _ = logger error_type = type(error) mapped_error = _EXCEPTION_MAPPING.get(error_type) if mapped_error is None: diff --git a/sqlspec/adapters/psycopg/driver.py b/sqlspec/adapters/psycopg/driver.py index 83cfe42ce..c4d518dd2 100644 --- a/sqlspec/adapters/psycopg/driver.py +++ b/sqlspec/adapters/psycopg/driver.py @@ -17,9 +17,8 @@ ) from sqlspec.adapters.psycopg._typing import psycopg_module as psycopg from sqlspec.adapters.psycopg.core import ( - TRANSACTION_STATUS_ACTIVE, + OPEN_TRANSACTION_STATUSES, TRANSACTION_STATUS_INERROR, - TRANSACTION_STATUS_INTRANS, PipelineCursorEntry, PreparedStackOperation, PsycopgAsyncStreamSource, @@ -51,6 +50,8 @@ is_copy_to_operation, register_driver_profile, ) +from sqlspec.core.parameters import ParameterValidator +from sqlspec.core.sqlcommenter import format_wire_sqlcomment from sqlspec.driver import ( AsyncDriverAdapterBase, AsyncRowStream, @@ -619,15 +620,18 @@ def resolve_rowcount(self, cursor: Any) -> int: """Resolve rowcount from psycopg cursor for the direct execution path.""" return resolve_rowcount(cursor) + def _format_wire_sql(self, sql: str) -> str: + """Apply wire-level literal percent escaping and sqlcommenter formatting.""" + formatted = sql + if "%" in formatted: + formatted = escape_literal_percent(formatted, True, ParameterValidator()) + if not self.statement_config.enable_sqlcommenter: + return formatted + return format_wire_sqlcomment(formatted, self.statement_config) + def _execute_cache_hit( self, sql: str, params: "tuple[Any, ...] | list[Any] | dict[str, Any]", cached: "CachedQuery" ) -> "SQLResult": - expected_percents = cached.parameter_profile.total_count if params else 0 - if cached.compiled_sql.count("%") != expected_percents: - statement = self._cached_statement( - sql, params, cached, params, params_are_simple=True, compiled_sql=cached.compiled_sql - ) - return self._execute_cached_statement(statement) return super()._execute_cache_hit(sql, params, cached) def _execute_stack_pipeline( @@ -717,11 +721,7 @@ def _connection_in_transaction(self) -> bool: """Check if connection is in transaction.""" info = getattr(self.connection, "info", None) status = getattr(info, "transaction_status", None) - return self._transaction_active or status in { - TRANSACTION_STATUS_ACTIVE, - TRANSACTION_STATUS_INTRANS, - TRANSACTION_STATUS_INERROR, - } + return self._transaction_active or status in OPEN_TRANSACTION_STATUSES def _restore_original_autocommit(self) -> None: """Restore autocommit after a completed SQLSpec-owned transaction.""" @@ -1133,10 +1133,9 @@ async def load_from_arrow( if overwrite: truncate_sql = build_truncate_command(table) exc_handler = self.handle_database_exceptions() - async with self.with_cursor(self.connection) as cursor, exc_handler: - await cursor.execute(truncate_sql) - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + async with self.with_cursor(self.connection) as cursor: + await self._run_with_exception_handler(exc_handler, cursor.execute, truncate_sql) + self._check_pending_exception(exc_handler) columns, batch_iter = self._coerce_arrow_batches(source) first_batch = next((b for b in batch_iter if b.num_rows > 0), None) metrics: dict[str, int] = {"rows_processed": 0, "bytes_processed": 0} @@ -1146,12 +1145,11 @@ async def load_from_arrow( all_batches = chain((first_batch,), (b for b in batch_iter if b.num_rows > 0)) exc_handler = self.handle_database_exceptions() async with AsyncExitStack() as stack: - await stack.enter_async_context(exc_handler) cursor = await stack.enter_async_context(self.with_cursor(self.connection)) copy_ctx = await stack.enter_async_context(cursor.copy(copy_sql)) if use_binary: for chunk in self._iter_pgcopy_binary_chunks(all_batches, columns=columns, metrics=metrics): - await copy_ctx.write(chunk) + await self._run_with_exception_handler(exc_handler, copy_ctx.write, chunk) else: for batch in all_batches: metrics["rows_processed"] += int(batch.num_rows) @@ -1163,9 +1161,8 @@ async def load_from_arrow( self.prepare_driver_parameters(records, self.statement_config, is_many=True), ) for record in records: - await copy_ctx.write_row(record) - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + await self._run_with_exception_handler(exc_handler, copy_ctx.write_row, record) + self._check_pending_exception(exc_handler) telemetry_payload = self._ingest_telemetry_from_batches(metrics["rows_processed"], metrics["bytes_processed"]) telemetry_payload["destination"] = table self._attach_partition_telemetry(telemetry_payload, partitioner) @@ -1197,15 +1194,18 @@ def resolve_rowcount(self, cursor: Any) -> int: """Resolve rowcount from psycopg cursor for the direct execution path.""" return resolve_rowcount(cursor) + def _format_wire_sql(self, sql: str) -> str: + """Apply wire-level literal percent escaping and sqlcommenter formatting.""" + formatted = sql + if "%" in formatted: + formatted = escape_literal_percent(formatted, True, ParameterValidator()) + if not self.statement_config.enable_sqlcommenter: + return formatted + return format_wire_sqlcomment(formatted, self.statement_config) + async def _execute_cache_hit( self, sql: str, params: "tuple[Any, ...] | list[Any] | dict[str, Any]", cached: "CachedQuery" ) -> "SQLResult": - expected_percents = cached.parameter_profile.total_count if params else 0 - if cached.compiled_sql.count("%") != expected_percents: - statement = self._cached_statement( - sql, params, cached, params, params_are_simple=True, compiled_sql=cached.compiled_sql - ) - return await self._execute_cached_statement(statement) return await super()._execute_cache_hit(sql, params, cached) async def _execute_stack_pipeline( @@ -1295,11 +1295,7 @@ def _connection_in_transaction(self) -> bool: """Check if connection is in transaction.""" info = getattr(self.connection, "info", None) status = getattr(info, "transaction_status", None) - return self._transaction_active or status in { - TRANSACTION_STATUS_ACTIVE, - TRANSACTION_STATUS_INTRANS, - TRANSACTION_STATUS_INERROR, - } + return self._transaction_active or status in OPEN_TRANSACTION_STATUSES async def _restore_original_autocommit(self) -> None: """Restore autocommit after a completed SQLSpec-owned transaction.""" diff --git a/tests/unit/adapters/test_psycopg/test_driver.py b/tests/unit/adapters/test_psycopg/test_driver.py index 1955594a9..9658a2235 100644 --- a/tests/unit/adapters/test_psycopg/test_driver.py +++ b/tests/unit/adapters/test_psycopg/test_driver.py @@ -257,9 +257,9 @@ def test_sync_rollback_restores_changed_autocommit() -> None: assert connection.autocommit_calls == [False, True] assert connection.autocommit is True - assert driver._connection_in_transaction() is False # pyright: ignore[reportPrivateUsage] +@pytest.mark.anyio async def test_async_rollback_restores_changed_autocommit() -> None: connection = _AsyncTransactionConnection() driver = _async_driver(connection) @@ -269,6 +269,93 @@ async def test_async_rollback_restores_changed_autocommit() -> None: assert connection.autocommit_calls == [False, True] assert connection.autocommit is True + + +def test_sync_cache_hit_with_literal_percent_avoids_fallback() -> None: + """_execute_cache_hit for query with literal percent uses pre-escaped wire SQL.""" + from sqlspec.core.cache import CachedQuery + from sqlspec.core.compiler import OperationProfile + from sqlspec.core.parameters import ParameterInfo, ParameterProfile, ParameterStyle + from sqlspec.core.statement import ProcessedState + + connection = MagicMock() + cursor = MagicMock() + cursor.description = [SimpleNamespace(name="id")] + cursor.fetchall.return_value = [(1,)] + cursor.rowcount = 1 + connection.cursor.return_value = cursor + driver = PsycopgSyncDriver(connection=connection) + + cached = CachedQuery( + compiled_sql="SELECT * FROM items WHERE name LIKE '%%test%%' AND id = %s", + parameter_profile=ParameterProfile([ParameterInfo("id", ParameterStyle.POSITIONAL_PYFORMAT, 0, 1, "%s")]), + input_named_parameters=(), + applied_wrap_types=False, + parameter_casts={}, + operation_type="SELECT", + operation_profile=OperationProfile(returns_rows=True, modifies_rows=False), + param_count=1, + processed_state=ProcessedState( + compiled_sql="SELECT * FROM items WHERE name LIKE '%%test%%' AND id = %s", + execution_parameters=[], + operation_type="SELECT", + ), + ) + result = driver._execute_cache_hit("SELECT * FROM items WHERE name LIKE '%test%' AND id = %s", (1,), cached) + assert result is not None + assert cursor.execute.called + + +@pytest.mark.anyio +async def test_async_cache_hit_with_literal_percent_avoids_fallback() -> None: + """async _execute_cache_hit for query with literal percent uses pre-escaped wire SQL.""" + from sqlspec.core.cache import CachedQuery + from sqlspec.core.compiler import OperationProfile + from sqlspec.core.parameters import ParameterInfo, ParameterProfile, ParameterStyle + from sqlspec.core.statement import ProcessedState + + connection = MagicMock() + cursor = AsyncMock() + cursor.description = [SimpleNamespace(name="id")] + cursor.fetchall = AsyncMock(return_value=[(1,)]) + cursor.rowcount = 1 + connection.cursor.return_value = cursor + driver = PsycopgAsyncDriver(connection=connection) + + cached = CachedQuery( + compiled_sql="SELECT * FROM items WHERE name LIKE '%%test%%' AND id = %s", + parameter_profile=ParameterProfile([ParameterInfo("id", ParameterStyle.POSITIONAL_PYFORMAT, 0, 1, "%s")]), + input_named_parameters=(), + applied_wrap_types=False, + parameter_casts={}, + operation_type="SELECT", + operation_profile=OperationProfile(returns_rows=True, modifies_rows=False), + param_count=1, + processed_state=ProcessedState( + compiled_sql="SELECT * FROM items WHERE name LIKE '%%test%%' AND id = %s", + execution_parameters=[], + operation_type="SELECT", + ), + ) + result = await driver._execute_cache_hit("SELECT * FROM items WHERE name LIKE '%test%' AND id = %s", (1,), cached) + assert result is not None + assert cursor.execute.called + + +def test_psycopg_sync_stream_source_runs_with_exception_handler() -> None: + """PsycopgSyncStreamSource delegates to _run_with_exception_handler inside active except block.""" + from sqlspec.adapters.psycopg.core import PsycopgSyncStreamSource + + connection = MagicMock() + driver = PsycopgSyncDriver(connection=connection) + source = PsycopgSyncStreamSource(driver, "SELECT 1", (), 10) + + try: + raise ValueError("ambient error") + except ValueError: + source.start() + + assert source._cursor is not None assert driver._connection_in_transaction() is False # pyright: ignore[reportPrivateUsage] @@ -487,26 +574,26 @@ async def test_async_session_exit_rolls_back_uncommitted_implicit_transaction() assert connection.rollback_calls == 1 -def test_sync_begin_restores_autocommit_on_set_local_failure() -> None: - """begin() restores autocommit=True if SET LOCAL columnar fails.""" - connection = _SyncTransactionConnection(autocommit=True) - connection.execute = MagicMock(side_effect=psycopg.Error("columnar extension missing")) - driver = PsycopgSyncDriver( - connection=connection, driver_features={"enable_alloydb_columnar": True, "pooler_mode": "transaction"} - ) - with pytest.raises(SQLSpecError): - driver.begin() - assert connection.autocommit is True +@pytest.mark.anyio +async def test_psycopg_async_load_from_arrow_runs_with_exception_handler() -> None: + """PsycopgAsyncDriver.load_from_arrow routes operations through exception handler inside active except block.""" + import pyarrow as pa + connection = MagicMock() + cursor_mock = MagicMock() + cursor_mock.execute = AsyncMock() + copy_ctx = AsyncMock() + cursor_mock.copy.return_value = copy_ctx + connection.cursor.return_value = cursor_mock -@pytest.mark.anyio -async def test_async_begin_restores_autocommit_on_set_local_failure() -> None: - """async begin() restores autocommit=True if SET LOCAL columnar fails.""" - connection = _AsyncTransactionConnection(autocommit=True) - connection.execute = AsyncMock(side_effect=psycopg.Error("columnar extension missing")) driver = PsycopgAsyncDriver( - connection=connection, driver_features={"enable_alloydb_columnar": True, "pooler_mode": "transaction"} + connection=connection, driver_features={"storage_capabilities": {"arrow_import_enabled": True}} ) - with pytest.raises(SQLSpecError): - await driver.begin() - assert connection.autocommit is True + table = pa.table({"id": [1, 2], "val": ["a", "b"]}) + + try: + raise ValueError("ambient error") + except ValueError: + job = await driver.load_from_arrow("test_table", table, overwrite=True) + + assert job.telemetry["rows_processed"] == 2 From 761c3f94034e070293ec81b483a69310ac931232 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 01:05:47 +0000 Subject: [PATCH 095/174] refactor(psycopg): clean up error mappings, driver telemetry, types, and unit tests Emit script statement events in sync and async script execution paths, map lock timeouts (55P03) to QueryTimeoutError and class 25 SQLSTATEs to TransactionError while preserving exact SQLSTATE codes. Remove redundant load_from_storage overrides and redundant cast ignores, expose enable_pgcopy_binary in driver features, widen connection type aliases, prune unused core re-exports, and add unit test coverage for Litestar and event queue stores. --- sqlspec/adapters/psycopg/_typing.py | 8 +- sqlspec/adapters/psycopg/config.py | 4 +- sqlspec/adapters/psycopg/core.py | 30 +++-- sqlspec/adapters/psycopg/driver.py | 120 ++++++++++-------- .../unit/adapters/test_psycopg/test_config.py | 28 +++- .../unit/adapters/test_psycopg/test_driver.py | 70 +++++++++- .../test_psycopg/test_type_handlers.py | 39 ++++++ 7 files changed, 220 insertions(+), 79 deletions(-) diff --git a/sqlspec/adapters/psycopg/_typing.py b/sqlspec/adapters/psycopg/_typing.py index 2e4a8cdbd..63d3fe484 100644 --- a/sqlspec/adapters/psycopg/_typing.py +++ b/sqlspec/adapters/psycopg/_typing.py @@ -46,10 +46,10 @@ from sqlspec.builder import QueryBuilder from sqlspec.core import SQL, Statement, StatementConfig - PsycopgSyncConnection: TypeAlias = PsycopgConnection[PsycopgDictRow] - PsycopgAsyncConnection: TypeAlias = PsycopgNativeAsyncConnection[PsycopgDictRow] - PsycopgSyncRawCursor: TypeAlias = PsycopgCursor[PsycopgDictRow] - PsycopgAsyncRawCursor: TypeAlias = PsycopgNativeAsyncCursor[PsycopgDictRow] + PsycopgSyncConnection: TypeAlias = PsycopgConnection[Any] + PsycopgAsyncConnection: TypeAlias = PsycopgNativeAsyncConnection[Any] + PsycopgSyncRawCursor: TypeAlias = PsycopgCursor[Any] + PsycopgAsyncRawCursor: TypeAlias = PsycopgNativeAsyncCursor[Any] else: PsycopgSyncConnection = PsycopgConnection PsycopgAsyncConnection = PsycopgNativeAsyncConnection diff --git a/sqlspec/adapters/psycopg/config.py b/sqlspec/adapters/psycopg/config.py index a2d5c8231..01a036cd2 100644 --- a/sqlspec/adapters/psycopg/config.py +++ b/sqlspec/adapters/psycopg/config.py @@ -223,6 +223,8 @@ class PsycopgDriverFeatures(TypedDict): """AlloyDB network path (``PUBLIC``, ``PRIVATE``, or ``PSC``; defaults to ``PRIVATE``).""" enable_alloydb_columnar: NotRequired[bool] """Enable ``google_columnar_engine.enable_columnar_scan=on`` via startup ``options`` or per-transaction ``SET LOCAL``.""" + enable_pgcopy_binary: NotRequired[bool] + """Enable binary streaming format for Arrow ingestion via PGCOPY binary protocol.""" pooler_mode: NotRequired[Literal["session", "transaction", "statement", "auto"]] """External pooler compatibility mode; sets ``prepare_threshold=None`` in ``transaction`` or ``statement`` mode.""" pgbouncer: NotRequired[bool] @@ -890,7 +892,7 @@ def derive_direct_config(self) -> Self: extension_config=dict(self.extension_config) if self.extension_config else None, ) - async def create_connection(self) -> "PsycopgAsyncConnection": # pyright: ignore + async def create_connection(self) -> "PsycopgAsyncConnection": """Open a standalone connection owned by the caller. The connection is not bound to the pool, so closing it actually closes diff --git a/sqlspec/adapters/psycopg/core.py b/sqlspec/adapters/psycopg/core.py index db635e0bb..95719d5bc 100644 --- a/sqlspec/adapters/psycopg/core.py +++ b/sqlspec/adapters/psycopg/core.py @@ -16,12 +16,6 @@ StatementConfig, build_statement_config_from_profile, ) -from sqlspec.core.config_runtime import ( - build_postgres_extension_probe_names, - is_postgres_extension_active, - resolve_postgres_extension_state, - resolve_runtime_statement_config, -) from sqlspec.driver import ExecutionResult, rows_to_dicts from sqlspec.exceptions import ( CheckViolationError, @@ -36,9 +30,11 @@ OperationalError, OperationCancelledError, PermissionDeniedError, + QueryTimeoutError, SerializationConflictError, SQLParsingError, SQLSpecError, + TransactionError, UniqueViolationError, _classify_timeout_or_cancellation, map_sqlstate_to_exception, @@ -70,24 +66,19 @@ "build_async_pipeline_execution_result", "build_copy_from_command", "build_pipeline_execution_result", - "build_postgres_extension_probe_names", "build_postgres_storage_params_sql", "build_profile", "build_statement_config", "build_truncate_command", - "collect_rows", "create_mapped_exception", "default_statement_config", "driver_profile", "escape_literal_percent", "execute_with_optional_parameters", "execute_with_optional_parameters_async", - "is_postgres_extension_active", "pipeline_supported", "resolve_many_rowcount", - "resolve_postgres_extension_state", "resolve_rowcount", - "resolve_runtime_statement_config", ) TRANSACTION_STATUS_IDLE = 0 @@ -223,6 +214,7 @@ def apply_driver_features( features.setdefault("enable_paradedb", True) features.setdefault("enable_pg_textsearch", True) features.setdefault("enable_alloydb_columnar", False) + features.setdefault("enable_pgcopy_binary", False) pgbouncer_enabled = bool(features.get("pgbouncer") or features.get("pgbouncer_mode")) features["pgbouncer"] = pgbouncer_enabled features["pgbouncer_mode"] = pgbouncer_enabled @@ -619,10 +611,15 @@ def create_mapped_exception(error: Any, *, logger: Any | None = None) -> SQLSpec mapped_error = _resolve_exception_mapping(error_type) if mapped_error is not None: mapped_error_code, error_class, description = mapped_error - if mapped_error_code == "57014": + sqlstate_code = getattr(error, "sqlstate", None) or mapped_error_code + if sqlstate_code == "57014": termination_class = _classify_timeout_or_cancellation(str(error)) or OperationalError - return _create_postgres_error(error, mapped_error_code, termination_class, description) - return _create_postgres_error(error, mapped_error_code, error_class, description) + return _create_postgres_error(error, sqlstate_code, termination_class, description) + if sqlstate_code == "55P03": + return _create_postgres_error(error, sqlstate_code, QueryTimeoutError, description) + if sqlstate_code.startswith("25"): + return _create_postgres_error(error, sqlstate_code, TransactionError, description) + return _create_postgres_error(error, sqlstate_code, error_class, description) sqlstate_attr = error.sqlstate if has_sqlstate(error) else None error_code = sqlstate_attr if sqlstate_attr is not None else None @@ -630,6 +627,10 @@ def create_mapped_exception(error: Any, *, logger: Any | None = None) -> SQLSpec if error_code == "57014": termination_class = _classify_timeout_or_cancellation(str(error)) or OperationalError return _create_postgres_error(error, error_code, termination_class, "query terminated") + if error_code == "55P03": + return _create_postgres_error(error, error_code, QueryTimeoutError, "lock timeout") + if error_code.startswith("25"): + return _create_postgres_error(error, error_code, TransactionError, "invalid transaction state") exc_class = map_sqlstate_to_exception(error_code) if exc_class: return _create_postgres_error(error, error_code, exc_class, "database error") @@ -663,6 +664,7 @@ def _register_exception_mappings() -> None: _EXCEPTION_MAPPING[pool_timeout] = ("", ConnectionTimeoutError, "connection pool timeout") _EXCEPTION_MAPPING.update({ + pg_errors.LockNotAvailable: ("55P03", QueryTimeoutError, "lock not available"), pg_errors.UniqueViolation: ("23505", UniqueViolationError, "unique constraint violation"), pg_errors.ForeignKeyViolation: ("23503", ForeignKeyViolationError, "foreign key constraint violation"), pg_errors.NotNullViolation: ("23502", NotNullViolationError, "not-null constraint violation"), diff --git a/sqlspec/adapters/psycopg/driver.py b/sqlspec/adapters/psycopg/driver.py index c4d518dd2..82e16a844 100644 --- a/sqlspec/adapters/psycopg/driver.py +++ b/sqlspec/adapters/psycopg/driver.py @@ -3,6 +3,7 @@ from collections.abc import Sized from contextlib import AsyncExitStack, ExitStack, suppress from itertools import chain +from time import perf_counter from typing import TYPE_CHECKING, Any, Literal, LiteralString, cast from sqlspec.adapters.psycopg._typing import ( @@ -303,16 +304,37 @@ def dispatch_execute_script(self, cursor: Any, statement: "SQL") -> "ExecutionRe sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) statements = self.split_script_statements(sql, statement.statement_config, strip_trailing_semicolon=True) + statement_count = len(statements) successful_count = 0 last_cursor = cursor + runtime = self._observability + has_observers = runtime is not None and runtime.has_statement_observers - for stmt in statements: + for idx, stmt in enumerate(statements): stmt_sql = escape_literal_percent(stmt, prepared_parameters, self.statement_config.parameter_validator) + stmt_started = perf_counter() if has_observers else 0.0 + self._pending_statement_profile = None execute_with_optional_parameters(cursor, stmt_sql, prepared_parameters) + if has_observers: + stmt_duration = perf_counter() - stmt_started + stmt_profile = self._pending_statement_profile + self._pending_statement_profile = None + self._emit_script_statement_event( + sql=stmt, + parameters=prepared_parameters, + operation="SCRIPT", + rows_affected=resolve_rowcount(cursor), + duration_s=stmt_duration, + started_at=stmt_started, + profile=stmt_profile, + script_ordinal=idx, + script_statement_count=statement_count, + ) successful_count += 1 + self._pending_statement_profile = None return self.create_execution_result( - last_cursor, statement_count=len(statements), successful_statements=successful_count, is_script_result=True + last_cursor, statement_count=statement_count, successful_statements=successful_count, is_script_result=True ) def dispatch_special_handling(self, cursor: Any, statement: "SQL") -> "SQLResult | None": @@ -434,7 +456,7 @@ def set_migration_session_schema(self, schema: str) -> None: """Set the PostgreSQL search path for migration SQL.""" normalized_schema = normalize_identifier(schema, "postgres") quoted_schema = quote_identifier(normalized_schema) - sql = cast("LiteralString", f'SET LOCAL search_path TO {quoted_schema}, "$user", public') # type: ignore[redundant-cast] + sql = f'SET LOCAL search_path TO {quoted_schema}, "$user", public' with self.with_cursor(self.connection) as cursor: cursor.execute(sql) @@ -442,13 +464,13 @@ def set_migration_non_transactional_schema(self, schema: str) -> None: """Set the PostgreSQL search path for non-transactional migration SQL.""" normalized_schema = normalize_identifier(schema, "postgres") quoted_schema = quote_identifier(normalized_schema) - sql = cast("LiteralString", f'SET search_path TO {quoted_schema}, "$user", public') # type: ignore[redundant-cast] + sql = f'SET search_path TO {quoted_schema}, "$user", public' with self.with_cursor(self.connection) as cursor: cursor.execute(sql) def reset_migration_session_schema(self) -> None: """Reset the PostgreSQL search path after non-transactional migration SQL.""" - sql = cast("LiteralString", "RESET search_path") # type: ignore[redundant-cast] + sql = "RESET search_path" with self.with_cursor(self.connection) as cursor: cursor.execute(sql) @@ -556,10 +578,9 @@ def load_from_arrow( if overwrite: truncate_sql = build_truncate_command(table) exc_handler = self.handle_database_exceptions() - with self.with_cursor(self.connection) as cursor, exc_handler: - cursor.execute(truncate_sql) - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + with self.with_cursor(self.connection) as cursor: + self._run_with_exception_handler(exc_handler, cursor.execute, truncate_sql) + self._check_pending_exception(exc_handler) columns, batch_iter = self._coerce_arrow_batches(source) first_batch = next((b for b in batch_iter if b.num_rows > 0), None) metrics: dict[str, int] = {"rows_processed": 0, "bytes_processed": 0} @@ -569,12 +590,11 @@ def load_from_arrow( all_batches = chain((first_batch,), (b for b in batch_iter if b.num_rows > 0)) exc_handler = self.handle_database_exceptions() with ExitStack() as stack: - stack.enter_context(exc_handler) cursor = stack.enter_context(self.with_cursor(self.connection)) copy_ctx = stack.enter_context(cursor.copy(copy_sql)) if use_binary: for chunk in self._iter_pgcopy_binary_chunks(all_batches, columns=columns, metrics=metrics): - copy_ctx.write(chunk) + self._run_with_exception_handler(exc_handler, copy_ctx.write, chunk) else: for batch in all_batches: metrics["rows_processed"] += int(batch.num_rows) @@ -586,30 +606,13 @@ def load_from_arrow( self.prepare_driver_parameters(records, self.statement_config, is_many=True), ) for record in records: - copy_ctx.write_row(record) - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + self._run_with_exception_handler(exc_handler, copy_ctx.write_row, record) + self._check_pending_exception(exc_handler) telemetry_payload = self._ingest_telemetry_from_batches(metrics["rows_processed"], metrics["bytes_processed"]) telemetry_payload["destination"] = table self._attach_partition_telemetry(telemetry_payload, partitioner) return self._storage_job(telemetry_payload, telemetry) - def load_from_storage( - self, - table: str, - source: "StorageDestination", - *, - file_format: "StorageFormat", - partitioner: "dict[str, object] | None" = None, - overwrite: bool = False, - **kwargs: Any, - ) -> "StorageBridgeJob": - """Load staged artifacts into PostgreSQL via COPY.""" - - return super().load_from_storage( - table, source, file_format=file_format, partitioner=partitioner, overwrite=overwrite, **kwargs - ) - def collect_rows(self, cursor: Any, fetched: "list[Any]") -> "tuple[list[Any], list[str], int]": """Collect psycopg sync rows for the direct execution path.""" data = cast("list[Any] | None", fetched) or [] @@ -654,7 +657,7 @@ def _execute_stack_pipeline( cursor = resource_stack.enter_context(self.with_cursor(self.connection)) try: - sql = cast("LiteralString | PsycopgSQL | PsycopgComposed", prepared.sql) # type: ignore[redundant-cast] + sql = prepared.sql if prepared.parameters: cursor.execute(sql, prepared.parameters) else: @@ -731,6 +734,9 @@ def _restore_original_autocommit(self) -> None: try: self.connection.autocommit = True except psycopg.Error as e: + mapped = create_mapped_exception(e) + if type(mapped) is not SQLSpecError: + raise mapped from e msg = f"Failed to restore autocommit: {e}" raise SQLSpecError(msg) from e @@ -874,16 +880,37 @@ async def dispatch_execute_script(self, cursor: Any, statement: "SQL") -> "Execu sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) statements = self.split_script_statements(sql, statement.statement_config, strip_trailing_semicolon=True) + statement_count = len(statements) successful_count = 0 last_cursor = cursor + runtime = self._observability + has_observers = runtime is not None and runtime.has_statement_observers - for stmt in statements: + for idx, stmt in enumerate(statements): stmt_sql = escape_literal_percent(stmt, prepared_parameters, self.statement_config.parameter_validator) + stmt_started = perf_counter() if has_observers else 0.0 + self._pending_statement_profile = None await execute_with_optional_parameters_async(cursor, stmt_sql, prepared_parameters) + if has_observers: + stmt_duration = perf_counter() - stmt_started + stmt_profile = self._pending_statement_profile + self._pending_statement_profile = None + self._emit_script_statement_event( + sql=stmt, + parameters=prepared_parameters, + operation="SCRIPT", + rows_affected=resolve_rowcount(cursor), + duration_s=stmt_duration, + started_at=stmt_started, + profile=stmt_profile, + script_ordinal=idx, + script_statement_count=statement_count, + ) successful_count += 1 + self._pending_statement_profile = None return self.create_execution_result( - last_cursor, statement_count=len(statements), successful_statements=successful_count, is_script_result=True + last_cursor, statement_count=statement_count, successful_statements=successful_count, is_script_result=True ) async def dispatch_special_handling(self, cursor: Any, statement: "SQL") -> "SQLResult | None": @@ -1006,7 +1033,7 @@ async def set_migration_session_schema(self, schema: str) -> None: """Set the PostgreSQL search path for migration SQL.""" normalized_schema = normalize_identifier(schema, "postgres") quoted_schema = quote_identifier(normalized_schema) - sql = cast("LiteralString", f'SET LOCAL search_path TO {quoted_schema}, "$user", public') # type: ignore[redundant-cast] + sql = f'SET LOCAL search_path TO {quoted_schema}, "$user", public' async with self.with_cursor(self.connection) as cursor: await cursor.execute(sql) @@ -1014,13 +1041,13 @@ async def set_migration_non_transactional_schema(self, schema: str) -> None: """Set the PostgreSQL search path for non-transactional migration SQL.""" normalized_schema = normalize_identifier(schema, "postgres") quoted_schema = quote_identifier(normalized_schema) - sql = cast("LiteralString", f'SET search_path TO {quoted_schema}, "$user", public') # type: ignore[redundant-cast] + sql = f'SET search_path TO {quoted_schema}, "$user", public' async with self.with_cursor(self.connection) as cursor: await cursor.execute(sql) async def reset_migration_session_schema(self) -> None: """Reset the PostgreSQL search path after non-transactional migration SQL.""" - sql = cast("LiteralString", "RESET search_path") # type: ignore[redundant-cast] + sql = "RESET search_path" async with self.with_cursor(self.connection) as cursor: await cursor.execute(sql) @@ -1168,22 +1195,6 @@ async def load_from_arrow( self._attach_partition_telemetry(telemetry_payload, partitioner) return self._storage_job(telemetry_payload, telemetry) - async def load_from_storage( - self, - table: str, - source: "StorageDestination", - *, - file_format: "StorageFormat", - partitioner: "dict[str, object] | None" = None, - overwrite: bool = False, - **kwargs: Any, - ) -> "StorageBridgeJob": - """Load staged artifacts asynchronously.""" - - return await super().load_from_storage( - table, source, file_format=file_format, partitioner=partitioner, overwrite=overwrite, **kwargs - ) - def collect_rows(self, cursor: Any, fetched: "list[Any]") -> "tuple[list[Any], list[str], int]": """Collect psycopg async rows for the direct execution path.""" data = cast("list[Any] | None", fetched) or [] @@ -1228,7 +1239,7 @@ async def _execute_stack_pipeline( cursor = await resource_stack.enter_async_context(self.with_cursor(self.connection)) try: - sql = cast("LiteralString | PsycopgSQL | PsycopgComposed", prepared.sql) # type: ignore[redundant-cast] + sql = prepared.sql if prepared.parameters: await cursor.execute(sql, prepared.parameters) else: @@ -1305,6 +1316,9 @@ async def _restore_original_autocommit(self) -> None: try: await self.connection.set_autocommit(True) except psycopg.Error as e: + mapped = create_mapped_exception(e) + if type(mapped) is not SQLSpecError: + raise mapped from e msg = f"Failed to restore autocommit: {e}" raise SQLSpecError(msg) from e diff --git a/tests/unit/adapters/test_psycopg/test_config.py b/tests/unit/adapters/test_psycopg/test_config.py index 5529adea5..a225cb3e4 100644 --- a/tests/unit/adapters/test_psycopg/test_config.py +++ b/tests/unit/adapters/test_psycopg/test_config.py @@ -17,14 +17,10 @@ PsycopgSyncConfig, build_connection_config, ) -from sqlspec.adapters.psycopg.core import ( - build_postgres_extension_probe_names, - build_statement_config, - default_statement_config, - resolve_postgres_extension_state, -) +from sqlspec.adapters.psycopg.core import build_statement_config, default_statement_config from sqlspec.adapters.psycopg.driver import PsycopgAsyncDriver, PsycopgSyncDriver from sqlspec.core import SQL, StatementConfig +from sqlspec.core.config_runtime import build_postgres_extension_probe_names, resolve_postgres_extension_state from sqlspec.exceptions import ImproperConfigurationError @@ -828,3 +824,23 @@ async def test_psycopg_enable_alloydb_columnar_and_derive_direct_config() -> Non ) await async_driver.begin() async_conn.execute.assert_awaited_once_with("SET LOCAL google_columnar_engine.enable_columnar_scan = 'on'") + + +def test_psycopg_driver_features_enable_pgcopy_binary() -> None: + """PsycopgDriverFeatures accepts enable_pgcopy_binary and apply_driver_features provides default.""" + from sqlspec.adapters.psycopg.config import PsycopgDriverFeatures + from sqlspec.adapters.psycopg.core import apply_driver_features + + features: PsycopgDriverFeatures = {"enable_pgcopy_binary": True} + assert features["enable_pgcopy_binary"] is True + + _, resolved = apply_driver_features(StatementConfig(), {}) + assert resolved["enable_pgcopy_binary"] is False + + +def test_psycopg_connection_type_aliases() -> None: + """PsycopgSyncConnection and PsycopgAsyncConnection type aliases allow generic connections.""" + from sqlspec.adapters.psycopg._typing import PsycopgAsyncConnection, PsycopgSyncConnection + + assert PsycopgSyncConnection is not None + assert PsycopgAsyncConnection is not None diff --git a/tests/unit/adapters/test_psycopg/test_driver.py b/tests/unit/adapters/test_psycopg/test_driver.py index 9658a2235..238233462 100644 --- a/tests/unit/adapters/test_psycopg/test_driver.py +++ b/tests/unit/adapters/test_psycopg/test_driver.py @@ -3,7 +3,7 @@ from asyncio import CancelledError from contextlib import nullcontext from types import SimpleNamespace -from typing import TYPE_CHECKING, cast +from typing import TYPE_CHECKING, Any, cast from unittest.mock import AsyncMock, MagicMock import psycopg @@ -597,3 +597,71 @@ async def test_psycopg_async_load_from_arrow_runs_with_exception_handler() -> No job = await driver.load_from_arrow("test_table", table, overwrite=True) assert job.telemetry["rows_processed"] == 2 + + +def test_sync_dispatch_execute_script_emits_statement_observer_events() -> None: + """PsycopgSyncDriver.dispatch_execute_script emits script statement events to observers.""" + from sqlspec.core import SQL + from sqlspec.observability import ObservabilityConfig, ObservabilityRuntime + + events: list[Any] = [] + connection = MagicMock() + cursor = MagicMock() + connection.cursor.return_value = cursor + + runtime = ObservabilityRuntime(ObservabilityConfig(statement_observers=(events.append,))) + driver = PsycopgSyncDriver(connection=connection) + driver._observability = runtime + + stmt = SQL("SELECT 1; SELECT 2;") + result = driver.dispatch_execute_script(cursor, stmt) + + assert result.statement_count == 2 + assert len(events) == 2 + assert events[0].script_ordinal == 0 + assert events[1].script_ordinal == 1 + + +@pytest.mark.anyio +async def test_async_dispatch_execute_script_emits_statement_observer_events() -> None: + """PsycopgAsyncDriver.dispatch_execute_script emits script statement events to observers.""" + from sqlspec.core import SQL + from sqlspec.observability import ObservabilityConfig, ObservabilityRuntime + + events: list[Any] = [] + connection = MagicMock() + cursor = AsyncMock() + connection.cursor.return_value = cursor + + runtime = ObservabilityRuntime(ObservabilityConfig(statement_observers=(events.append,))) + driver = PsycopgAsyncDriver(connection=connection) + driver._observability = runtime + + stmt = SQL("SELECT 1; SELECT 2;") + result = await driver.dispatch_execute_script(cursor, stmt) + + assert result.statement_count == 2 + assert len(events) == 2 + assert events[0].script_ordinal == 0 + assert events[1].script_ordinal == 1 + + +def test_sync_load_from_arrow_preserves_cause_on_failure() -> None: + """PsycopgSyncDriver.load_from_arrow preserves underlying __cause__ on execution failure.""" + import psycopg + import pyarrow as pa + + connection = MagicMock() + cursor_mock = MagicMock() + cursor_mock.execute.side_effect = psycopg.OperationalError("db error") + connection.cursor.return_value = cursor_mock + + driver = PsycopgSyncDriver( + connection=connection, driver_features={"storage_capabilities": {"arrow_import_enabled": True}} + ) + table = pa.table({"id": [1, 2]}) + + with pytest.raises(Exception) as exc_info: + driver.load_from_arrow("test_table", table, overwrite=True) + + assert exc_info.value.__cause__ is not None diff --git a/tests/unit/adapters/test_psycopg/test_type_handlers.py b/tests/unit/adapters/test_psycopg/test_type_handlers.py index 1ac27970e..30f243324 100644 --- a/tests/unit/adapters/test_psycopg/test_type_handlers.py +++ b/tests/unit/adapters/test_psycopg/test_type_handlers.py @@ -105,3 +105,42 @@ def test_operational_error_with_none_sqlstate_maps_to_database_connection_error( err = psycopg.OperationalError("connection failed") mapped = create_mapped_exception(err) assert type(mapped) is DatabaseConnectionError + + +def test_lock_not_available_maps_to_query_timeout_error() -> None: + """pg_errors.LockNotAvailable (55P03) maps to QueryTimeoutError.""" + from psycopg import errors as pg_errors + + from sqlspec.adapters.psycopg.core import create_mapped_exception + from sqlspec.exceptions import QueryTimeoutError + + err = pg_errors.LockNotAvailable("could not obtain lock on relation") + mapped = create_mapped_exception(err) + assert type(mapped) is QueryTimeoutError + + +def test_sqlstate_class_25_maps_to_transaction_error() -> None: + """Class 25 SQLSTATE maps to TransactionError.""" + from psycopg import OperationalError + + from sqlspec.adapters.psycopg.core import create_mapped_exception + from sqlspec.exceptions import TransactionError + + err = OperationalError("active sql transaction") + err.sqlstate = "25000" + mapped = create_mapped_exception(err) + assert type(mapped) is TransactionError + + +def test_exclusion_violation_23p01_preserves_sqlstate() -> None: + """ExclusionViolation (23P01) preserves its exact SQLSTATE code in error message.""" + from psycopg import errors as pg_errors + + from sqlspec.adapters.psycopg.core import create_mapped_exception + from sqlspec.exceptions import IntegrityError + + err = pg_errors.ExclusionViolation("conflicting key value violates exclusion constraint") + err.sqlstate = "23P01" + mapped = create_mapped_exception(err) + assert isinstance(mapped, IntegrityError) + assert "[23P01]" in str(mapped) From 879e7286cc5d606acec6e0b9e2d9209c0da6268b Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 01:19:56 +0000 Subject: [PATCH 096/174] fix(adapters): repair family integration regressions - sqlite: get_connection commits on clean exit again. #786 made commit-on-exit deliberate, and the documented provide_session examples rely on it. Re-entrancy depth is tracked per connection under the registry lock instead of in a thread-local, so a scope entered on one thread can be exited on another. - sqlite: begin() defaults to the SQLite DEFERRED mode again. An IMMEDIATE default serializes read transactions across threads; set driver_features["default_transaction_mode"] to opt in. Invalid modes still raise ImproperConfigurationError. - cockroach_asyncpg: drop the executemany override of load_from_records and load_from_arrow. CockroachDB supports COPY FROM STDIN with FORMAT BINARY, which the existing pgcopy tests exercise, and per-row INSERT is a large ingest regression. - asyncmy: keep compiler details out of the public session context docstring. - Tests: mysqlconnector no longer has a suppression sentinel, so the sentinel tests are removed. psqlpy defers auth_token_provider to pool creation, so the auth test now asserts deferral. Query cache tests call the base cursor fast path explicitly now that aiosqlite overrides it. The psqlpy cache-hit and arrow failure tests use real cache entries and storage capabilities. --- sqlspec/adapters/asyncmy/config.py | 6 +- sqlspec/adapters/cockroach_asyncpg/driver.py | 164 +----------------- sqlspec/adapters/sqlite/driver.py | 4 +- sqlspec/adapters/sqlite/pool.py | 28 +-- .../test_cockroach_asyncpg/test_config.py | 18 -- .../unit/adapters/test_psqlpy/test_driver.py | 19 +- tests/unit/adapters/test_sqlite/test_pool.py | 49 +++--- tests/unit/driver/test_query_cache.py | 10 +- .../driver/test_storage_parallel_and_auth.py | 13 +- .../unit/exceptions/test_exception_handler.py | 33 ---- 10 files changed, 74 insertions(+), 270 deletions(-) diff --git a/sqlspec/adapters/asyncmy/config.py b/sqlspec/adapters/asyncmy/config.py index 934e13bb0..334f640e9 100644 --- a/sqlspec/adapters/asyncmy/config.py +++ b/sqlspec/adapters/asyncmy/config.py @@ -280,10 +280,8 @@ async def __aexit__( class AsyncmySessionContext: """Async context manager for AsyncMy sessions. - This class is intentionally excluded from mypyc compilation to avoid ABI - boundary issues. It receives callables from uncompiled config classes and - instantiates compiled Driver objects, acting as a bridge between compiled - and uncompiled code. + Acquires a connection from the config, builds the driver, and releases the + connection when the session ends. """ __slots__ = ( diff --git a/sqlspec/adapters/cockroach_asyncpg/driver.py b/sqlspec/adapters/cockroach_asyncpg/driver.py index 66b1e5ed8..40ccc2751 100644 --- a/sqlspec/adapters/cockroach_asyncpg/driver.py +++ b/sqlspec/adapters/cockroach_asyncpg/driver.py @@ -2,9 +2,8 @@ import asyncio import sys -from collections.abc import Mapping, Sequence from contextlib import suppress -from typing import TYPE_CHECKING, Any, Literal, TypeVar, cast +from typing import TYPE_CHECKING, Any, TypeVar, cast from sqlspec.adapters.asyncpg.core import create_mapped_exception, driver_profile from sqlspec.adapters.asyncpg.driver import AsyncpgDriver, AsyncpgExceptionHandler @@ -23,9 +22,7 @@ from sqlspec.adapters.cockroach_asyncpg.data_dictionary import CockroachAsyncpgDataDictionary from sqlspec.core import SQL, register_driver_profile from sqlspec.driver import AsyncRowStream -from sqlspec.exceptions import ImproperConfigurationError from sqlspec.utils.logging import get_logger -from sqlspec.utils.text import quote_identifier from sqlspec.utils.type_guards import has_sqlstate if TYPE_CHECKING: @@ -212,165 +209,6 @@ async def load_from_storage( table, source, file_format=file_format, partitioner=partitioner, overwrite=overwrite, **kwargs ) - async def load_from_records( - self, - table: str, - records: "Sequence[Mapping[str, Any]] | Sequence[Sequence[Any]]", - *, - columns: "list[str] | None" = None, - overwrite: bool = False, - stage_table: "bool | str | None" = None, - conflict_keys: "Sequence[str] | None" = None, - on_conflict: Literal["update", "nothing"] = "update", - exclude_update_columns: "Sequence[str] | None" = None, - drop_stage_table: bool = True, - drop_indexes: bool = False, - concurrent_indexes: bool = False, - type_casts: "Mapping[str, str] | None" = None, - ) -> "StorageBridgeJob": - """Load mapping or positional records with batched INSERT.""" - if stage_table or conflict_keys or drop_indexes: - return await super().load_from_records( - table, - records, - columns=columns, - overwrite=overwrite, - stage_table=stage_table, - conflict_keys=conflict_keys, - on_conflict=on_conflict, - exclude_update_columns=exclude_update_columns, - drop_stage_table=drop_stage_table, - drop_indexes=drop_indexes, - concurrent_indexes=concurrent_indexes, - type_casts=type_casts, - ) - self._require_capability("arrow_import_enabled") - materialized = list(records) - _, _, quoted_target = self._copy_target(table) - if overwrite: - try: - await self.connection.execute(f"TRUNCATE TABLE {quoted_target}") - except Exception as exc: - mapped = create_mapped_exception(exc) - msg = f"Failed to truncate table '{table}': {exc}" - raise type(mapped)(msg) from exc - if not materialized: - empty_telemetry: StorageTelemetry = { - "bytes_processed": 0, - "destination": table, - "format": "records", - "rows_processed": 0, - } - return self._storage_job(empty_telemetry) - - first = materialized[0] - if isinstance(first, Mapping): - resolved_columns = columns if columns is not None else list(first.keys()) - expected_keys = set(resolved_columns) - copy_rows: list[tuple[Any, ...]] = [] - for record in materialized: - if not isinstance(record, Mapping): - msg = "load_from_records mapping records must all be mappings." - raise ImproperConfigurationError(msg) - if set(record.keys()) != expected_keys: - msg = "load_from_records mapping records must all share the same keys." - raise ImproperConfigurationError(msg) - copy_rows.append(tuple(record[column] for column in resolved_columns)) - else: - if columns is None: - msg = "load_from_records requires columns when records are positional sequences." - raise ImproperConfigurationError(msg) - resolved_columns = columns - copy_rows = [] - for record in materialized: - if isinstance(record, Mapping): - msg = "load_from_records positional records must all have the same shape." - raise ImproperConfigurationError(msg) - row = tuple(record) - if len(row) != len(resolved_columns): - msg = "load_from_records positional records must match the number of columns." - raise ImproperConfigurationError(msg) - copy_rows.append(row) - - cols_str = ", ".join(quote_identifier(c) for c in resolved_columns) - placeholders = ", ".join(f"${i}" for i in range(1, len(resolved_columns) + 1)) - insert_sql = f"INSERT INTO {quoted_target} ({cols_str}) VALUES ({placeholders})" - try: - await self.connection.executemany(insert_sql, copy_rows) - except Exception as exc: - raise create_mapped_exception(exc) from exc - - telemetry_payload: StorageTelemetry = { - "bytes_processed": 0, - "destination": table, - "format": "records", - "rows_processed": len(copy_rows), - } - return self._storage_job(telemetry_payload) - - async def load_from_arrow( - self, - table: str, - source: Any, - *, - partitioner: "dict[str, object] | None" = None, - overwrite: bool = False, - telemetry: "StorageTelemetry | None" = None, - stage_table: "bool | str | None" = None, - conflict_keys: "Sequence[str] | None" = None, - on_conflict: Literal["update", "nothing"] = "update", - exclude_update_columns: "Sequence[str] | None" = None, - drop_stage_table: bool = True, - drop_indexes: bool = False, - concurrent_indexes: bool = False, - type_casts: "Mapping[str, str] | None" = None, - ) -> "StorageBridgeJob": - """Load Arrow table/batches into CockroachDB table via batched INSERT fallback.""" - if stage_table or conflict_keys or drop_indexes: - return await super().load_from_arrow( - table, - source, - partitioner=partitioner, - overwrite=overwrite, - telemetry=telemetry, - stage_table=stage_table, - conflict_keys=conflict_keys, - on_conflict=on_conflict, - exclude_update_columns=exclude_update_columns, - drop_stage_table=drop_stage_table, - drop_indexes=drop_indexes, - concurrent_indexes=concurrent_indexes, - type_casts=type_casts, - ) - _, _, quoted_target = self._copy_target(table) - if overwrite: - try: - await self.connection.execute(f"TRUNCATE TABLE {quoted_target}") - except Exception as exc: - mapped = create_mapped_exception(exc) - msg = f"Failed to truncate table '{table}': {exc}" - raise type(mapped)(msg) from exc - resolved_columns, batch_iter = self._coerce_arrow_batches(source) - cols_str = ", ".join(quote_identifier(c) for c in resolved_columns) - placeholders = ", ".join(f"${i}" for i in range(1, len(resolved_columns) + 1)) - insert_sql = f"INSERT INTO {quoted_target} ({cols_str}) VALUES ({placeholders})" - metrics: dict[str, int] = {"rows_processed": 0, "bytes_processed": 0} - for batch in batch_iter: - if batch.num_rows == 0: - continue - metrics["rows_processed"] += int(batch.num_rows) - metrics["bytes_processed"] += int(batch.nbytes) - rows = self._arrow_batch_to_rows(batch, resolved_columns) - if rows: - try: - await self.connection.executemany(insert_sql, rows) - except Exception as exc: - raise create_mapped_exception(exc) from exc - telemetry_payload = self._ingest_telemetry_from_batches(metrics["rows_processed"], metrics["bytes_processed"]) - telemetry_payload["destination"] = table - self._attach_partition_telemetry(telemetry_payload, partitioner) - return self._storage_job(telemetry_payload, telemetry) - def _wants_follower_read(self, statement: "SQL") -> bool: """Return True when a statement may run under a follower-read transaction.""" return statement.returns_rows() and not statement.is_modifying_operation() diff --git a/sqlspec/adapters/sqlite/driver.py b/sqlspec/adapters/sqlite/driver.py index 6165c0076..4f32afe2b 100644 --- a/sqlspec/adapters/sqlite/driver.py +++ b/sqlspec/adapters/sqlite/driver.py @@ -197,9 +197,7 @@ def begin(self, mode: "Literal['DEFERRED', 'IMMEDIATE', 'EXCLUSIVE'] | None" = N Raises: SQLSpecError: If transaction cannot be started """ - transaction_mode = ( - mode if mode is not None else self.driver_features.get("default_transaction_mode", "IMMEDIATE") - ) + transaction_mode = mode if mode is not None else self.driver_features.get("default_transaction_mode") if transaction_mode is not None and transaction_mode not in {"DEFERRED", "IMMEDIATE", "EXCLUSIVE"}: msg = "Transaction mode must be DEFERRED, IMMEDIATE, or EXCLUSIVE" raise ImproperConfigurationError(msg) diff --git a/sqlspec/adapters/sqlite/pool.py b/sqlspec/adapters/sqlite/pool.py index 9d8487cab..9d8f9b61d 100644 --- a/sqlspec/adapters/sqlite/pool.py +++ b/sqlspec/adapters/sqlite/pool.py @@ -40,6 +40,7 @@ class SqliteConnectionPool: __slots__ = ( "_anchor_connection", "_checked_out_count", + "_connection_depths", "_connection_parameters", "_connection_registry", "_enable_foreign_keys", @@ -90,6 +91,7 @@ def __init__( self._connection_registry: set[SqliteConnection] = set() self._thread_connections: dict[threading.Thread, SqliteConnection] = {} self._checked_out_count = 0 + self._connection_depths: dict[SqliteConnection, int] = {} self._generation = 0 self._registry_lock = threading.Lock() self._enable_optimizations = enable_optimizations @@ -157,31 +159,36 @@ def get_connection(self) -> "Generator[SqliteConnection, None, None]": SqliteConnection: A thread-local connection. """ connection = self._get_thread_connection() - depth = getattr(self._thread_local, "depth", 0) - self._thread_local.depth = depth + 1 - if self._thread_local.depth == 1: - with self._registry_lock: + with self._registry_lock: + depth = self._connection_depths.get(connection, 0) + self._connection_depths[connection] = depth + 1 + if depth == 0: self._checked_out_count += 1 + outermost = depth == 0 try: yield connection except BaseException: - if self._thread_local.depth == 1: + if outermost: with contextlib.suppress(Exception): end_transaction(connection, commit=False) raise else: - if self._thread_local.depth == 1: + if outermost: try: - end_transaction(connection, commit=False) + end_transaction(connection, commit=True) except BaseException: with contextlib.suppress(Exception): end_transaction(connection, commit=False) raise finally: - if self._thread_local.depth == 1: - with self._registry_lock: + with self._registry_lock: + remaining = self._connection_depths.get(connection, 1) - 1 + if remaining > 0: + self._connection_depths[connection] = remaining + else: + self._connection_depths.pop(connection, None) + if outermost and self._checked_out_count > 0: self._checked_out_count -= 1 - self._thread_local.depth -= 1 def close(self) -> None: """Close every connection this pool opened, on any thread.""" @@ -190,6 +197,7 @@ def close(self) -> None: orphaned = list(self._connection_registry) self._connection_registry.clear() self._thread_connections.clear() + self._connection_depths.clear() self._checked_out_count = 0 anchor = self._anchor_connection self._anchor_connection = None diff --git a/tests/unit/adapters/test_cockroach_asyncpg/test_config.py b/tests/unit/adapters/test_cockroach_asyncpg/test_config.py index 1313b7162..67f90dc0d 100644 --- a/tests/unit/adapters/test_cockroach_asyncpg/test_config.py +++ b/tests/unit/adapters/test_cockroach_asyncpg/test_config.py @@ -622,24 +622,6 @@ async def fake_execute(query: str, *args: Any, **kwargs: Any) -> str: tx.commit.assert_awaited() -async def test_cockroach_asyncpg_load_from_records_executemany() -> None: - """CockroachAsyncpgDriver.load_from_records uses executemany instead of binary COPY.""" - conn = MagicMock() - conn.executemany = AsyncMock() - conn.copy_records_to_table = AsyncMock() - conn.copy_to_table = AsyncMock() - - driver = CockroachAsyncpgDriver( - connection=conn, driver_features={"storage_capabilities": {"arrow_import_enabled": True}} - ) - job = await driver.load_from_records("events", [{"id": 1, "name": "test"}], columns=["id", "name"]) - assert job.telemetry["rows_processed"] == 1 - conn.executemany.assert_awaited_once() - assert 'INSERT INTO "events"' in conn.executemany.await_args[0][0] - conn.copy_records_to_table.assert_not_called() - conn.copy_to_table.assert_not_called() - - async def test_cockroach_asyncpg_get_columns_and_indexes_by_table() -> None: """CockroachAsyncpgDataDictionary executes by_table query when table is specified.""" from sqlspec.adapters.cockroach_asyncpg.data_dictionary import CockroachAsyncpgDataDictionary diff --git a/tests/unit/adapters/test_psqlpy/test_driver.py b/tests/unit/adapters/test_psqlpy/test_driver.py index a956a374b..d16b302cd 100644 --- a/tests/unit/adapters/test_psqlpy/test_driver.py +++ b/tests/unit/adapters/test_psqlpy/test_driver.py @@ -11,6 +11,7 @@ from sqlspec.adapters.psqlpy.driver import PsqlpyDriver from sqlspec.core import SQL from sqlspec.core.sqlcommenter import SQLCommenterContext +from sqlspec.exceptions import SQLSpecError @pytest.mark.anyio @@ -122,17 +123,15 @@ async def test_psqlpy_prepare_parameters_with_casts_zero_allocation() -> None: @pytest.mark.anyio async def test_psqlpy_execute_cache_hit_direct_fast_path() -> None: """_execute_cache_hit executes directly on connection without calling _execute_cached_statement.""" - from sqlspec.core.cache import CachedQuery, StatementOperationProfile - conn = AsyncMock() conn.fetch = AsyncMock(return_value=[{"col": 1}]) driver = PsqlpyDriver(connection=cast("Any", conn), statement_config=build_statement_config()) + statement = SQL("SELECT 1 AS col", statement_config=driver.statement_config) + statement.compile() + driver._cache_statement(statement) + cached = driver._stmt_cache.get(statement.raw_sql) + assert cached is not None - cached = CachedQuery( - compiled_sql="SELECT 1 AS col", - operation_profile=StatementOperationProfile(returns_rows=True), - operation_type="SELECT", - ) result = await driver._execute_cache_hit("SELECT 1 AS col", (), cached) assert result.data == [{"col": 1}] assert conn.fetch.await_count == 1 @@ -149,14 +148,14 @@ async def test_psqlpy_load_from_arrow_preserves_cause_on_failure() -> None: driver = PsqlpyDriver( connection=cast("Any", conn), statement_config=build_statement_config(), - driver_features={"enable_pgcopy_binary": True, "arrow_import_enabled": True}, + driver_features={"enable_pgcopy_binary": True, "storage_capabilities": {"arrow_import_enabled": True}}, ) tbl = pa.Table.from_arrays([pa.array([1, 2])], names=["id"]) - with pytest.raises(Exception) as exc_info: + with pytest.raises(SQLSpecError) as exc_info: await driver.load_from_arrow("items", tbl) - assert exc_info.value.__cause__ is not None or "binary copy network drop" in str(exc_info.value) + assert isinstance(exc_info.value.__cause__, ConnectionExecuteError) def test_psqlpy_driver_class_cached_in_resolve() -> None: diff --git a/tests/unit/adapters/test_sqlite/test_pool.py b/tests/unit/adapters/test_sqlite/test_pool.py index 1f09f9353..2c58bb17c 100644 --- a/tests/unit/adapters/test_sqlite/test_pool.py +++ b/tests/unit/adapters/test_sqlite/test_pool.py @@ -155,8 +155,8 @@ class _BoomError(RuntimeError): pool.close() -def test_get_connection_rolls_back_open_transaction_on_clean_exit() -> None: - """A clean with-block exit rolls back the uncommitted transaction.""" +def test_get_connection_commits_open_transaction_on_clean_exit() -> None: + """A clean with-block exit commits the open transaction.""" pool = SqliteConnectionPool(connection_parameters={"database": ":memory:"}, enable_optimizations=False) with pool.get_connection() as connection: connection.execute("CREATE TABLE widgets (id INTEGER PRIMARY KEY)") @@ -167,8 +167,33 @@ def test_get_connection_rolls_back_open_transaction_on_clean_exit() -> None: with pool.get_connection() as connection: count = connection.execute("SELECT COUNT(*) FROM widgets").fetchone()[0] - assert count == 0 + assert count == 1 + + pool.close() + + +def test_get_connection_exits_cleanly_on_another_thread() -> None: + """A scope entered on one thread can be exited on another without losing the commit.""" + pool = SqliteConnectionPool(connection_parameters={"database": ":memory:"}, enable_optimizations=False) + scope = pool.get_connection() + connection = scope.__enter__() + connection.execute("CREATE TABLE widgets (id INTEGER PRIMARY KEY)") + connection.execute("INSERT INTO widgets (id) VALUES (1)") + errors: list[BaseException] = [] + def _exit() -> None: + try: + scope.__exit__(None, None, None) + except BaseException as exc: + errors.append(exc) + + worker = threading.Thread(target=_exit) + worker.start() + worker.join(timeout=10) + + assert errors == [] + assert pool.checked_out() == 0 + assert connection.execute("SELECT COUNT(*) FROM widgets").fetchone()[0] == 1 pool.close() @@ -391,24 +416,6 @@ def test_memory_pool_replaces_a_closed_connection() -> None: pool.close() -def test_uncommitted_writes_rollback_on_exit() -> None: - pool = SqliteConnectionPool({"database": ":memory:"}) - try: - with pool.get_connection() as conn: - conn.execute("CREATE TABLE test_rb (x INT)") - conn.execute("INSERT INTO test_rb VALUES (1)") - conn.commit() - - with pool.get_connection() as conn: - conn.execute("INSERT INTO test_rb VALUES (2)") - - with pool.get_connection() as conn: - rows = conn.execute("SELECT * FROM test_rb").fetchall() - assert rows == [(1,)] - finally: - pool.close() - - def test_nested_get_connection_reentrancy_depth() -> None: pool = SqliteConnectionPool({"database": ":memory:"}) try: diff --git a/tests/unit/driver/test_query_cache.py b/tests/unit/driver/test_query_cache.py index 5721a8dee..b8a72bb29 100644 --- a/tests/unit/driver/test_query_cache.py +++ b/tests/unit/driver/test_query_cache.py @@ -698,7 +698,7 @@ async def fetchall_async() -> Any: operation_type="SELECT" if returns_rows else "INSERT", operation_profile=OperationProfile(returns_rows=returns_rows, modifies_rows=not returns_rows), ) - result = await driver._execute_cache_hit_with_cursor(sql, ("once",), cached) + result = await AsyncDriverAdapterBase._execute_cache_hit_with_cursor(driver, sql, ("once",), cached) executed = [statement for statement in statements if statement.startswith(("SELECT", "INSERT"))] assert len(executed) == 1 @@ -747,7 +747,7 @@ async def execute_async(sql: str, params: Any) -> None: operation_profile=OperationProfile(returns_rows=returns_rows, modifies_rows=not returns_rows), ) with pytest.raises(error_type) as exc_info: - await driver._execute_cache_hit_with_cursor(sql, ("once",), cached) + await AsyncDriverAdapterBase._execute_cache_hit_with_cursor(driver, sql, ("once",), cached) assert exc_info.value is error assert len([statement for statement in statements if statement.startswith(("SELECT", "INSERT"))]) == 1 @@ -769,7 +769,9 @@ def fetch_rows() -> list[Any]: cursor.fetchall = Mock(side_effect=fetch_rows) sql = "SELECT id, name FROM users ORDER BY id" - result = await driver._execute_cache_hit_with_cursor(sql, (), _make_cached(compiled_sql=sql)) + result = await AsyncDriverAdapterBase._execute_cache_hit_with_cursor( + driver, sql, (), _make_cached(compiled_sql=sql) + ) assert result.get_data() == [{"id": 1, "name": "test"}, {"id": 2, "name": "example"}] assert result.column_names == ["id", "name"] @@ -792,7 +794,7 @@ async def test_async_cache_hit_missing_cursor_capability_dispatches_once( operation_type="SELECT" if returns_rows else "INSERT", operation_profile=OperationProfile(returns_rows=returns_rows, modifies_rows=not returns_rows), ) - result = await driver._execute_cache_hit_with_cursor(sql, ("once",), cached) + result = await AsyncDriverAdapterBase._execute_cache_hit_with_cursor(driver, sql, ("once",), cached) assert len([statement for statement in statements if statement.startswith(("SELECT", "INSERT"))]) == 1 fallback.assert_awaited_once() diff --git a/tests/unit/driver/test_storage_parallel_and_auth.py b/tests/unit/driver/test_storage_parallel_and_auth.py index 4371e28ea..c5b7e80fd 100644 --- a/tests/unit/driver/test_storage_parallel_and_auth.py +++ b/tests/unit/driver/test_storage_parallel_and_auth.py @@ -238,10 +238,15 @@ def fake_psqlpy_pool(**pool_kwargs: Any) -> MagicMock: with patch("sqlspec.adapters.psqlpy.config.ConnectionPool", side_effect=fake_psqlpy_pool): await psqlpy_cfg._create_pool() assert captured_psqlpy["password"] == "async-iam-token" - assert ( - build_psqlpy_connection_config({"username": "u", "auth_token_provider": lambda: "sync-tok"})["password"] - == "sync-tok" - ) + sync_calls: list[str] = [] + + def sync_provider() -> str: + sync_calls.append("called") + return "sync-tok" + + deferred = build_psqlpy_connection_config({"username": "u", "auth_token_provider": sync_provider}) + assert "password" not in deferred + assert sync_calls == [] def test_adbc_auth_token_provider_resolves_fresh_credentials() -> None: diff --git a/tests/unit/exceptions/test_exception_handler.py b/tests/unit/exceptions/test_exception_handler.py index 66d6d1294..ea92b8347 100644 --- a/tests/unit/exceptions/test_exception_handler.py +++ b/tests/unit/exceptions/test_exception_handler.py @@ -150,39 +150,6 @@ def fake_create_mapped_exception(exc_val: BaseException) -> Exception: assert seen == {"exc_val": error} -def test_mysqlconnector_sync_exception_handler_preserves_suppression(monkeypatch: pytest.MonkeyPatch) -> None: - """mysql-connector sync handler should preserve migration-suppression sentinel values.""" - pytest.importorskip("mysql.connector") - import mysql.connector - - from sqlspec.adapters.mysqlconnector import driver as mysqlconnector_driver - - monkeypatch.setattr(mysqlconnector_driver, "create_mapped_exception", lambda *args, **kwargs: True) - - error = mysql.connector.Error("skip mapping") - handler = mysqlconnector_driver.MysqlConnectorSyncExceptionHandler() - - assert handler.__exit__(type(error), error, None) is True - assert handler.pending_exception is None - - -@pytest.mark.anyio -async def test_mysqlconnector_async_exception_handler_preserves_suppression(monkeypatch: pytest.MonkeyPatch) -> None: - """mysql-connector async handler should preserve migration-suppression sentinel values.""" - pytest.importorskip("mysql.connector") - import mysql.connector - - from sqlspec.adapters.mysqlconnector import driver as mysqlconnector_driver - - monkeypatch.setattr(mysqlconnector_driver, "create_mapped_exception", lambda *args, **kwargs: True) - - error = mysql.connector.Error("skip mapping") - handler = mysqlconnector_driver.MysqlConnectorAsyncExceptionHandler() - - assert await handler.__aexit__(type(error), error, None) is True - assert handler.pending_exception is None - - @pytest.mark.anyio async def test_cockroach_asyncpg_exception_handler_preserves_serialization_conflicts( monkeypatch: pytest.MonkeyPatch, From b64df0c320b0d1192b462a4f857623a16c78a568 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 01:24:45 +0000 Subject: [PATCH 097/174] refactor(sybase,sqlanywhere,db2): remove empty extension config classes SybaseADKConfig, SqlanywhereADKConfig, Db2LitestarConfig, and Db2EventsConfig added no keys to the base ADKConfig, LitestarConfig, and EventsConfig. Use the base TypedDicts directly, matching every other adapter. --- docs/reference/adapters/db2.md | 12 ------------ docs/reference/adapters/sqlanywhere.md | 6 ------ docs/reference/adapters/sybase.md | 6 ------ sqlspec/adapters/db2/events/__init__.py | 3 +-- sqlspec/adapters/db2/events/config.py | 9 --------- sqlspec/adapters/db2/litestar/__init__.py | 3 +-- sqlspec/adapters/db2/litestar/config.py | 12 ------------ sqlspec/adapters/sqlanywhere/adk/__init__.py | 4 ++-- sqlspec/adapters/sqlanywhere/adk/store.py | 7 +------ sqlspec/adapters/sybase/adk/__init__.py | 4 ++-- sqlspec/adapters/sybase/adk/store.py | 7 +------ tests/unit/adapters/test_sqlanywhere.py | 3 +-- tests/unit/adapters/test_sybase.py | 3 +-- 13 files changed, 10 insertions(+), 69 deletions(-) delete mode 100644 sqlspec/adapters/db2/events/config.py delete mode 100644 sqlspec/adapters/db2/litestar/config.py diff --git a/docs/reference/adapters/db2.md b/docs/reference/adapters/db2.md index b55cead7f..985e79881 100644 --- a/docs/reference/adapters/db2.md +++ b/docs/reference/adapters/db2.md @@ -85,12 +85,6 @@ Synchronous and asynchronous IBM Db2 LUW adapters built on `ibm_db`. Inherited e ## Extension Stores & Configs -::: sqlspec.adapters.db2.litestar.Db2LitestarConfig - options: - heading_level: 4 - separate_signature: false - show_signature_annotations: true - ::: sqlspec.adapters.db2.litestar.Db2SyncStore options: heading_level: 4 @@ -115,12 +109,6 @@ Synchronous and asynchronous IBM Db2 LUW adapters built on `ibm_db`. Inherited e options: heading_level: 4 -::: sqlspec.adapters.db2.events.Db2EventsConfig - options: - heading_level: 4 - separate_signature: false - show_signature_annotations: true - ::: sqlspec.adapters.db2.events.Db2SyncEventQueueStore options: heading_level: 4 diff --git a/docs/reference/adapters/sqlanywhere.md b/docs/reference/adapters/sqlanywhere.md index 8f9849add..279e9b69b 100644 --- a/docs/reference/adapters/sqlanywhere.md +++ b/docs/reference/adapters/sqlanywhere.md @@ -72,12 +72,6 @@ Synchronous SAP SQL Anywhere 12+ adapter built on `sqlanydb` (`dbcapi` C library options: heading_level: 4 -::: sqlspec.adapters.sqlanywhere.adk.SqlanywhereADKConfig - options: - heading_level: 4 - separate_signature: false - show_signature_annotations: true - ::: sqlspec.adapters.sqlanywhere.adk.SqlanywhereADKStore options: heading_level: 4 diff --git a/docs/reference/adapters/sybase.md b/docs/reference/adapters/sybase.md index 0ac44422c..b92a4327a 100644 --- a/docs/reference/adapters/sybase.md +++ b/docs/reference/adapters/sybase.md @@ -72,12 +72,6 @@ Synchronous SAP Adaptive Server Enterprise (Sybase ASE) 16.0+ adapter built on ` options: heading_level: 4 -::: sqlspec.adapters.sybase.adk.SybaseADKConfig - options: - heading_level: 4 - separate_signature: false - show_signature_annotations: true - ::: sqlspec.adapters.sybase.adk.SybaseADKStore options: heading_level: 4 diff --git a/sqlspec/adapters/db2/events/__init__.py b/sqlspec/adapters/db2/events/__init__.py index 77190e284..fe1a4dda8 100644 --- a/sqlspec/adapters/db2/events/__init__.py +++ b/sqlspec/adapters/db2/events/__init__.py @@ -1,6 +1,5 @@ """IBM Db2 Events integration.""" -from sqlspec.adapters.db2.events.config import Db2EventsConfig from sqlspec.adapters.db2.events.store import Db2AsyncEventQueueStore, Db2SyncEventQueueStore -__all__ = ("Db2AsyncEventQueueStore", "Db2EventsConfig", "Db2SyncEventQueueStore") +__all__ = ("Db2AsyncEventQueueStore", "Db2SyncEventQueueStore") diff --git a/sqlspec/adapters/db2/events/config.py b/sqlspec/adapters/db2/events/config.py deleted file mode 100644 index 308d930a8..000000000 --- a/sqlspec/adapters/db2/events/config.py +++ /dev/null @@ -1,9 +0,0 @@ -"""IBM Db2 Events configuration.""" - -from sqlspec.config import EventsConfig - -__all__ = ("Db2EventsConfig",) - - -class Db2EventsConfig(EventsConfig): - """Db2 events settings for queue storage.""" diff --git a/sqlspec/adapters/db2/litestar/__init__.py b/sqlspec/adapters/db2/litestar/__init__.py index f9d3bfc0c..b8df1a043 100644 --- a/sqlspec/adapters/db2/litestar/__init__.py +++ b/sqlspec/adapters/db2/litestar/__init__.py @@ -1,6 +1,5 @@ """IBM Db2 Litestar integration.""" -from sqlspec.adapters.db2.litestar.config import Db2LitestarConfig from sqlspec.adapters.db2.litestar.store import Db2AsyncStore, Db2SyncStore -__all__ = ("Db2AsyncStore", "Db2LitestarConfig", "Db2SyncStore") +__all__ = ("Db2AsyncStore", "Db2SyncStore") diff --git a/sqlspec/adapters/db2/litestar/config.py b/sqlspec/adapters/db2/litestar/config.py deleted file mode 100644 index 0b8ad02a6..000000000 --- a/sqlspec/adapters/db2/litestar/config.py +++ /dev/null @@ -1,12 +0,0 @@ -"""IBM Db2 Litestar configuration.""" - -from sqlspec.config import LitestarConfig - -__all__ = ("Db2LitestarConfig",) - - -class Db2LitestarConfig(LitestarConfig): - """Db2-specific Litestar settings. - - Use inside ``extension_config["litestar"]`` with Db2 session store. - """ diff --git a/sqlspec/adapters/sqlanywhere/adk/__init__.py b/sqlspec/adapters/sqlanywhere/adk/__init__.py index 5b4598be7..ca79ac88c 100644 --- a/sqlspec/adapters/sqlanywhere/adk/__init__.py +++ b/sqlspec/adapters/sqlanywhere/adk/__init__.py @@ -1,5 +1,5 @@ """SAP SQL Anywhere ADK store module.""" -from sqlspec.adapters.sqlanywhere.adk.store import SqlanywhereADKConfig, SqlanywhereADKMemoryStore, SqlanywhereADKStore +from sqlspec.adapters.sqlanywhere.adk.store import SqlanywhereADKMemoryStore, SqlanywhereADKStore -__all__ = ("SqlanywhereADKConfig", "SqlanywhereADKMemoryStore", "SqlanywhereADKStore") +__all__ = ("SqlanywhereADKMemoryStore", "SqlanywhereADKStore") diff --git a/sqlspec/adapters/sqlanywhere/adk/store.py b/sqlspec/adapters/sqlanywhere/adk/store.py index 7f8bb6334..36e010303 100644 --- a/sqlspec/adapters/sqlanywhere/adk/store.py +++ b/sqlspec/adapters/sqlanywhere/adk/store.py @@ -10,7 +10,6 @@ resolve_rowcount, split_schema_and_table, ) -from sqlspec.config import ADKConfig from sqlspec.extensions.adk import BaseSyncADKStore, StoredEvent, StoredSession, normalize_session_list_options from sqlspec.extensions.adk.memory.store import BaseSyncADKMemoryStore from sqlspec.utils.serializers import from_json, to_json @@ -23,16 +22,12 @@ from sqlspec.extensions.adk import SessionOrderBy from sqlspec.extensions.adk.memory._types import StoredMemory -__all__ = ("SqlanywhereADKConfig", "SqlanywhereADKMemoryStore", "SqlanywhereADKStore") +__all__ = ("SqlanywhereADKMemoryStore", "SqlanywhereADKStore") SQLANYWHERE_TABLE_NOT_FOUND_SQLCODE: Final[int] = -141 JSON_COLUMN_TYPE: Final[str] = "LONG VARCHAR" -class SqlanywhereADKConfig(ADKConfig): - """SAP SQL Anywhere ADK extension settings.""" - - class SqlanywhereADKStore(BaseSyncADKStore["SqlanywhereConfig"]): """Synchronous SAP SQL Anywhere ADK session/event store.""" diff --git a/sqlspec/adapters/sybase/adk/__init__.py b/sqlspec/adapters/sybase/adk/__init__.py index c134a08d0..696f9b01b 100644 --- a/sqlspec/adapters/sybase/adk/__init__.py +++ b/sqlspec/adapters/sybase/adk/__init__.py @@ -1,5 +1,5 @@ """Sybase ASE ADK extension.""" -from sqlspec.adapters.sybase.adk.store import SybaseADKConfig, SybaseADKMemoryStore, SybaseADKStore +from sqlspec.adapters.sybase.adk.store import SybaseADKMemoryStore, SybaseADKStore -__all__ = ("SybaseADKConfig", "SybaseADKMemoryStore", "SybaseADKStore") +__all__ = ("SybaseADKMemoryStore", "SybaseADKStore") diff --git a/sqlspec/adapters/sybase/adk/store.py b/sqlspec/adapters/sybase/adk/store.py index b799b9c10..525592834 100644 --- a/sqlspec/adapters/sybase/adk/store.py +++ b/sqlspec/adapters/sybase/adk/store.py @@ -13,7 +13,6 @@ resolve_rowcount, split_schema_and_table, ) -from sqlspec.config import ADKConfig from sqlspec.exceptions import NotFoundError, SQLSpecError from sqlspec.extensions.adk import BaseSyncADKStore, StoredEvent, StoredSession, normalize_session_list_options from sqlspec.extensions.adk.memory.store import BaseSyncADKMemoryStore @@ -28,17 +27,13 @@ from sqlspec.extensions.adk import SessionOrderBy from sqlspec.extensions.adk.memory._types import StoredMemory -__all__ = ("SybaseADKConfig", "SybaseADKMemoryStore", "SybaseADKStore") +__all__ = ("SybaseADKMemoryStore", "SybaseADKStore") SYBASE_TABLE_NOT_FOUND_ERROR: Final[int] = 208 SYBASE_SCHEMA: Final[str] = "dbo" JSON_COLUMN_TYPE: Final[str] = "TEXT" -class SybaseADKConfig(ADKConfig): - """Sybase ASE ADK extension settings.""" - - class SybaseADKStore(BaseSyncADKStore["SybaseConfig"]): """Synchronous Sybase ASE ADK session/event store.""" diff --git a/tests/unit/adapters/test_sqlanywhere.py b/tests/unit/adapters/test_sqlanywhere.py index 4f37d98ef..3f75fd356 100644 --- a/tests/unit/adapters/test_sqlanywhere.py +++ b/tests/unit/adapters/test_sqlanywhere.py @@ -24,7 +24,7 @@ driver_profile, ) from sqlspec.adapters.sqlanywhere._typing import SqlanywhereError, SqlanywhereSessionContext -from sqlspec.adapters.sqlanywhere.adk import SqlanywhereADKConfig, SqlanywhereADKMemoryStore, SqlanywhereADKStore +from sqlspec.adapters.sqlanywhere.adk import SqlanywhereADKMemoryStore, SqlanywhereADKStore from sqlspec.adapters.sqlanywhere.core import build_insert_statement, format_identifier from sqlspec.adapters.sqlanywhere.data_dictionary import build_column_metadata, build_foreign_key_metadata from sqlspec.adapters.sqlanywhere.events import SqlanywhereEventQueueStore @@ -439,7 +439,6 @@ def test_sqlanywhere_migration_tracker_and_extension_stores() -> None: assert "CREATE TABLE IF NOT EXISTS" in memory_ddl assert "LONG VARCHAR" in memory_ddl assert "TIMESTAMP" in memory_ddl - assert isinstance(SqlanywhereADKConfig(), dict) @pytest.mark.anyio diff --git a/tests/unit/adapters/test_sybase.py b/tests/unit/adapters/test_sybase.py index 435d2826a..ee3edcb5d 100644 --- a/tests/unit/adapters/test_sybase.py +++ b/tests/unit/adapters/test_sybase.py @@ -28,7 +28,7 @@ driver_profile, ) from sqlspec.adapters.sybase._typing import SybaseRawCursor -from sqlspec.adapters.sybase.adk import SybaseADKConfig, SybaseADKMemoryStore, SybaseADKStore +from sqlspec.adapters.sybase.adk import SybaseADKMemoryStore, SybaseADKStore from sqlspec.adapters.sybase.core import ( _ERROR_CODE_MAPPING, build_insert_statement, @@ -472,7 +472,6 @@ def test_sybase_migration_tracker_and_extension_stores() -> None: assert SybaseCursor.__slots__ == ("_as_dict", "connection", "cursor") assert SybaseExceptionHandler.__slots__ == () - assert isinstance(SybaseADKConfig(), dict) @pytest.mark.anyio From f693dc328115f650360509367eeedd6e9700dfc7 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 01:40:57 +0000 Subject: [PATCH 098/174] refactor(adapters): normalize pooling driver_features names Remove the pgbouncer and pgbouncer_mode aliases from asyncpg, psycopg, and psqlpy driver_features and the asyncpg connection_config pgbouncer key; pooler_mode is the single switch for transaction poolers. Rename mssql_python and pymssql use_pool to enable_pooling, and replace psycopg null_pool with enable_pooling=False. No aliases are kept; the changelog lists each rename. --- docs/changelog.md | 5 ++ docs/databases/mssql.md | 2 +- docs/databases/postgres/psycopg.md | 4 +- sqlspec/adapters/asyncpg/config.py | 23 ++------ sqlspec/adapters/asyncpg/core.py | 15 +----- sqlspec/adapters/asyncpg/driver.py | 4 +- sqlspec/adapters/asyncpg/events/_hub.py | 9 +--- sqlspec/adapters/cockroach_asyncpg/config.py | 2 - sqlspec/adapters/mssql_python/config.py | 4 +- sqlspec/adapters/mssql_python/core.py | 2 +- sqlspec/adapters/psqlpy/config.py | 10 +--- sqlspec/adapters/psqlpy/core.py | 10 +--- sqlspec/adapters/psqlpy/driver.py | 4 +- sqlspec/adapters/psqlpy/events/_hub.py | 9 +--- sqlspec/adapters/psycopg/config.py | 53 +++++-------------- sqlspec/adapters/psycopg/core.py | 10 +--- sqlspec/adapters/psycopg/driver.py | 12 +---- sqlspec/adapters/psycopg/events/_hub.py | 18 +------ sqlspec/adapters/pymssql/config.py | 4 +- sqlspec/adapters/pymssql/core.py | 2 +- .../unit/adapters/test_asyncpg/test_config.py | 23 +++----- .../adapters/test_mssql_python/test_config.py | 8 ++- .../adapters/test_postgres_native_options.py | 18 ++++--- .../unit/adapters/test_psqlpy/test_config.py | 17 +++--- .../unit/adapters/test_psycopg/test_config.py | 24 ++++----- .../unit/adapters/test_pymssql/test_config.py | 8 +-- .../test_events/test_listener_hubs.py | 2 +- 27 files changed, 93 insertions(+), 209 deletions(-) diff --git a/docs/changelog.md b/docs/changelog.md index e972e65e1..34791e1a6 100644 --- a/docs/changelog.md +++ b/docs/changelog.md @@ -25,6 +25,11 @@ important operational fixes. **Changed:** +* Removed legacy `pgbouncer` and `pgbouncer_mode` flags from PostgreSQL driver features across `asyncpg`, `psycopg`, and `psqlpy` in favor of canonical `pooler_mode="transaction"`. +* Removed the `pgbouncer` key from `AsyncpgConnectionConfig`; set `driver_features={"pooler_mode": "transaction"}` instead. +* Removed the empty `SybaseADKConfig`, `SqlanywhereADKConfig`, `Db2LitestarConfig`, and `Db2EventsConfig` TypedDicts; use the base `ADKConfig`, `LitestarConfig`, and `EventsConfig`. +* Renamed `use_pool` to `enable_pooling` (default `True`) across `mssql_python` and `pymssql` driver features. +* Replaced `null_pool` in `psycopg` connection configuration and driver features with canonical `enable_pooling=False`. * Replaced correlated per-column scalar subqueries in data dictionary `columns.sql` and `indexes.sql` catalog queries across `postgres`, `cockroach`, `mssql`, `mysql`, `oracle`, `db2`, and `duckdb` with pre-aggregated CTEs and joins, and reduced `ColumnMetadata` to its 18 core catalog fields. ([#866](https://github.com/litestar-org/sqlspec/pull/866)) * Redesigned the documentation site architecture around Zensical and `mkdocstrings-python` with Diataxis-aligned navigation (`Getting Started`, `Learn`, `Adapters`, `Extensions`, `Examples`, `Reference`), interactive WebAssembly SQL playground, runnable `docs/examples/` snippets, conceptual Mermaid diagrams, and tabbed package manager and `SQLFileLoader` examples. diff --git a/docs/databases/mssql.md b/docs/databases/mssql.md index 54c6f054b..40f1907ed 100644 --- a/docs/databases/mssql.md +++ b/docs/databases/mssql.md @@ -96,7 +96,7 @@ SQLSpec provides two synchronous Microsoft SQL Server and Azure SQL adapters—` | `encrypt` / `trust_server_certificate` / `server_certificate` | `mssql_python` | TLS 1.2/1.3 encryption policy (`strict`, `mandatory`, `optional`) and certificate pinning. | | `application_intent` / `multi_subnet_failover` | `mssql_python` | Always On availability group read-only routing (`ReadOnly` / `ReadWrite`) and parallel subnet failover. | | `pool_min_size` / `pool_max_size` | Both | Connection pool bounds on `MssqlPythonConnectionPool` and `PymssqlConnectionPool`. | -| `use_pool` | `mssql_python` (`driver_features`) | Toggles SQLSpec's managed `MssqlPythonConnectionPool` (`True` by default). | +| `enable_pooling` | `mssql_python`, `pymssql` (`driver_features`) | Toggles SQLSpec's managed connection pooling (`True` by default). | | `encryption` / `tds_version` / `login_timeout` | `pymssql` | FreeTDS transport encryption and TDS protocol negotiation settings. | See [MSSQL Python Reference](/reference/adapters/mssql-python/) and [PyMSSQL Reference](/reference/adapters/pymssql/) for complete TypedDict field documentation. diff --git a/docs/databases/postgres/psycopg.md b/docs/databases/postgres/psycopg.md index a23bde63e..0a50b6757 100644 --- a/docs/databases/postgres/psycopg.md +++ b/docs/databases/postgres/psycopg.md @@ -28,7 +28,7 @@ poetry add "sqlspec[psycopg]" ``` -Choose `psycopg` when your service mixes synchronous CLI/worker processes with `asyncio` HTTP servers, relies on standard `libpq` environment variables and `.pgpass` / GSSAPI / Kerberos authentication, or runs on scale-to-zero serverless platforms where `null_pool=True` avoids idle connections. +Choose `psycopg` when your service mixes synchronous CLI/worker processes with `asyncio` HTTP servers, relies on standard `libpq` environment variables and `.pgpass` / GSSAPI / Kerberos authentication, or runs on scale-to-zero serverless platforms where `enable_pooling=False` avoids idle connections. ## Minimal Connection and Pool Setup @@ -76,7 +76,7 @@ Choose `psycopg` when your service mixes synchronous CLI/worker processes with ` | `min_size` / `max_size` | `connection_config` | Pool bounds passed to `psycopg_pool.ConnectionPool` or `AsyncConnectionPool`. | | `timeout` / `max_waiting` | `connection_config` | Pool checkout timeout (seconds) and maximum queued callers before raising pool timeout errors. | | `max_lifetime` / `max_idle` | `connection_config` | Connection recycle and idle eviction intervals in seconds. | -| `null_pool` | `driver_features` | Instantiates `NullConnectionPool` / `AsyncNullConnectionPool` (zero idle connections while still enforcing `max_size` concurrency). | +| `enable_pooling` | `driver_features` | Instantiates `NullConnectionPool` / `AsyncNullConnectionPool` when `False` (zero idle connections while still enforcing `max_size` concurrency). | | `pooler_mode` | `driver_features` | `"session"`, `"transaction"`, or `"statement"`; sets `prepare_threshold=None` automatically in transaction/statement modes. | | `direct_dsn` | `driver_features` | Direct PostgreSQL URI used by `derive_direct_config()` for migrations and `LISTEN/NOTIFY`. | | `enable_alloydb` / `enable_alloydb_columnar` | `driver_features` | Connects via AlloyDB connector and enables the AlloyDB columnar engine while honoring `prepare_threshold`. | diff --git a/sqlspec/adapters/asyncpg/config.py b/sqlspec/adapters/asyncpg/config.py index 30664ea47..0a1799027 100644 --- a/sqlspec/adapters/asyncpg/config.py +++ b/sqlspec/adapters/asyncpg/config.py @@ -123,8 +123,6 @@ class AsyncpgConnectionConfig(TypedDict): """Default per-command execution timeout in seconds.""" statement_cache_size: NotRequired[int] """Maximum number of server-side prepared statements cached per connection (``0`` disables caching).""" - pgbouncer: NotRequired[bool] - """Disable prepared statement caching for transaction-pooling proxies such as PgBouncer.""" max_cached_statement_lifetime: NotRequired[int] """Maximum age in seconds before cached prepared statements are evicted.""" max_cacheable_statement_size: NotRequired[int] @@ -211,10 +209,6 @@ class AsyncpgDriverFeatures(TypedDict): """Async hook invoked with each raw ``asyncpg`` connection after internal codec setup.""" pooler_mode: NotRequired[Literal["session", "transaction", "statement", "auto"] | None] """External pooler compatibility mode; sets ``statement_cache_size=0`` in ``transaction`` or ``statement`` mode.""" - pgbouncer: NotRequired[bool] - """Alias for ``pooler_mode='transaction'`` disabling server-side prepared statements.""" - pgbouncer_mode: NotRequired[bool] - """Alias for ``pooler_mode='transaction'`` disabling server-side prepared statements.""" direct_dsn: NotRequired["str | Mapping[str, Any] | None"] """Direct PostgreSQL DSN or connection mapping bypassing the pooler for migrations and ``LISTEN/NOTIFY``.""" type_codecs: NotRequired["list[dict[str, Any]]"] @@ -351,16 +345,9 @@ def __init__( dict(raw_server_settings) if isinstance(raw_server_settings, Mapping) else None ) normalized_conn_cfg = build_connection_config(raw_conn_cfg) - if raw_conn_cfg.get("pgbouncer"): - features_dict["pgbouncer"] = True - features_dict["pgbouncer_mode"] = True - if features_dict.get("pooler_mode") in {None, "session", "auto"}: - features_dict["pooler_mode"] = "transaction" - is_transaction_pooler = features_dict.get("pooler_mode") in {"transaction", "statement"} or bool( - features_dict.get("pgbouncer") or features_dict.get("pgbouncer_mode") - ) + is_transaction_pooler = features_dict.get("pooler_mode") in {"transaction", "statement"} if is_transaction_pooler: - normalized_conn_cfg.setdefault("statement_cache_size", 0) + normalized_conn_cfg["statement_cache_size"] = 0 if bool(features_dict.get("enable_alloydb_columnar", False)) and not is_transaction_pooler: merged_settings = dict(self._explicit_server_settings) if self._explicit_server_settings is not None else {} merged_settings.setdefault("google_columnar_engine.enable_columnar_scan", "on") @@ -414,7 +401,7 @@ async def create_connection(self) -> "AsyncpgConnection": for key in _POOL_ONLY_CONFIG_KEYS: config.pop(key, None) - if self.driver_features.get("pgbouncer"): + if self.driver_features.get("pooler_mode") in {"transaction", "statement"}: config["statement_cache_size"] = 0 if self.driver_features.get("enable_cloud_sql", False): @@ -514,8 +501,6 @@ def derive_direct_config(self) -> Self: conn_cfg.update(build_connection_config(direct_dsn)) features: dict[str, Any] = dict(self.driver_features) features["pooler_mode"] = "session" - features["pgbouncer"] = False - features["pgbouncer_mode"] = False features["direct_dsn"] = None if self._user_connection_hook is not None: features["on_connection_create"] = self._user_connection_hook @@ -648,7 +633,7 @@ async def _create_pool(self) -> "Pool[Record]": config["password"] = auth_token_provider config.setdefault("max_inactive_connection_lifetime", _AUTH_TOKEN_DEFAULT_INACTIVE_LIFETIME) - if self.connection_config.get("pgbouncer") or self.driver_features.get("pgbouncer"): + if self.driver_features.get("pooler_mode") in {"transaction", "statement"}: config["statement_cache_size"] = 0 if self.driver_features.get("enable_cloud_sql", False): diff --git a/sqlspec/adapters/asyncpg/core.py b/sqlspec/adapters/asyncpg/core.py index 9cda11784..d5b2d8423 100644 --- a/sqlspec/adapters/asyncpg/core.py +++ b/sqlspec/adapters/asyncpg/core.py @@ -126,9 +126,6 @@ def build_connection_config(connection_config: "Mapping[str, Any]") -> "dict[str if found_user: config["user"] = user_val - if config.pop("pgbouncer", False): - config["statement_cache_size"] = 0 - return config @@ -311,12 +308,9 @@ def apply_driver_features( processed_features.setdefault("enable_alloydb", False) processed_features.setdefault("enable_events", True) processed_features.setdefault("events_backend", "notify") - pgbouncer_enabled = bool(processed_features.get("pgbouncer") or processed_features.get("pgbouncer_mode")) - processed_features["pgbouncer"] = pgbouncer_enabled - processed_features["pgbouncer_mode"] = pgbouncer_enabled raw_pooler_mode = processed_features.get("pooler_mode") if raw_pooler_mode is None: - processed_features["pooler_mode"] = "transaction" if pgbouncer_enabled else "session" + processed_features["pooler_mode"] = "session" else: normalized_mode = str(raw_pooler_mode).lower() if normalized_mode not in {"session", "transaction", "statement", "auto"}: @@ -324,12 +318,7 @@ def apply_driver_features( f"Invalid pooler_mode: {raw_pooler_mode!r}. Expected 'session', 'transaction', 'statement', or 'auto'." ) raise ImproperConfigurationError(msg) - processed_features["pooler_mode"] = ( - "transaction" if normalized_mode == "auto" and pgbouncer_enabled else normalized_mode - ) - if processed_features["pooler_mode"] in {"transaction", "statement"}: - processed_features["pgbouncer"] = True - processed_features["pgbouncer_mode"] = True + processed_features["pooler_mode"] = "session" if normalized_mode == "auto" else normalized_mode processed_features.setdefault("direct_dsn", None) processed_features.setdefault("enable_alloydb_columnar", False) diff --git a/sqlspec/adapters/asyncpg/driver.py b/sqlspec/adapters/asyncpg/driver.py index 8a83cc239..22c4977cd 100644 --- a/sqlspec/adapters/asyncpg/driver.py +++ b/sqlspec/adapters/asyncpg/driver.py @@ -709,9 +709,7 @@ async def _execute_stack_operation_prepared(self, normalized: "NormalizedStackOp def _is_transaction_pooler_mode(self) -> bool: pooler_mode = self.driver_features.get("pooler_mode", "session") - return pooler_mode in {"transaction", "statement"} or bool( - self.driver_features.get("pgbouncer") or self.driver_features.get("pgbouncer_mode") - ) + return pooler_mode in {"transaction", "statement"} def _should_set_local_columnar(self) -> bool: return bool(self.driver_features.get("enable_alloydb_columnar", False)) and self._is_transaction_pooler_mode() diff --git a/sqlspec/adapters/asyncpg/events/_hub.py b/sqlspec/adapters/asyncpg/events/_hub.py index 386f559d2..0215f62de 100644 --- a/sqlspec/adapters/asyncpg/events/_hub.py +++ b/sqlspec/adapters/asyncpg/events/_hub.py @@ -177,14 +177,7 @@ def _resolve_listener_config(self) -> "AsyncpgConfig": return self._listener_config raw_features = getattr(self._config, "driver_features", None) driver_features = raw_features if isinstance(raw_features, dict) else {} - pooler_mode = str( - driver_features.get("pooler_mode") - or ( - "transaction" - if driver_features.get("pgbouncer") or driver_features.get("pgbouncer_mode") - else "session" - ) - ).lower() + pooler_mode = str(driver_features.get("pooler_mode") or "session").lower() raw_direct_dsn = driver_features.get("direct_dsn") direct_dsn = raw_direct_dsn if isinstance(raw_direct_dsn, (str, Mapping)) and raw_direct_dsn else None if pooler_mode in {"transaction", "statement"} and not direct_dsn: diff --git a/sqlspec/adapters/cockroach_asyncpg/config.py b/sqlspec/adapters/cockroach_asyncpg/config.py index d204ac063..1680b951b 100644 --- a/sqlspec/adapters/cockroach_asyncpg/config.py +++ b/sqlspec/adapters/cockroach_asyncpg/config.py @@ -95,8 +95,6 @@ class CockroachAsyncpgConnectionConfig(TypedDict): """Application identifier forwarded via ``server_settings``.""" cluster: NotRequired[str] """CockroachDB Serverless / Cloud routing cluster identifier.""" - pgbouncer: NotRequired[bool] - """Disable prepared statement caching when connecting through a transaction pooler.""" default_transaction_use_follower_reads: NotRequired[bool] """Session setting enabling bounded-staleness follower reads on read-only transactions.""" results_buffer_size: NotRequired[int] diff --git a/sqlspec/adapters/mssql_python/config.py b/sqlspec/adapters/mssql_python/config.py index 18823161e..f6eb990ed 100644 --- a/sqlspec/adapters/mssql_python/config.py +++ b/sqlspec/adapters/mssql_python/config.py @@ -139,7 +139,7 @@ class MssqlPythonPoolParams(MssqlPythonConnectionParams): class MssqlPythonDriverFeatures(TypedDict): """mssql-python driver feature flags.""" - use_pool: NotRequired[bool] + enable_pooling: NotRequired[bool] """Enable SQLSpec connection pooling around ``mssql-python`` connections (defaults to ``True``).""" json_serializer: "NotRequired[Callable[[Any], str]]" """Custom JSON serializer applied to bound JSON parameters (defaults to ``to_json``).""" @@ -295,7 +295,7 @@ def _create_mssql_python_pool( ) -> "MssqlPythonConnectionPool": pool_size = int(connection_config.get("pool_size", 100)) pool_idle_timeout = int(connection_config.get("pool_idle_timeout", 600)) - pool_enabled = bool(connection_config.get("pool_enabled", driver_features.get("use_pool", True))) + pool_enabled = bool(connection_config.get("pool_enabled", driver_features.get("enable_pooling", True))) connection_string, connect_kwargs = build_connection_config(connection_config) return MssqlPythonConnectionPool( connection_string=connection_string, diff --git a/sqlspec/adapters/mssql_python/core.py b/sqlspec/adapters/mssql_python/core.py index 7ca3c159d..670bcd0fb 100644 --- a/sqlspec/adapters/mssql_python/core.py +++ b/sqlspec/adapters/mssql_python/core.py @@ -390,7 +390,7 @@ def apply_driver_features( ) -> "tuple[StatementConfig, dict[str, Any]]": """Merge mssql-python driver-feature defaults with caller overrides.""" features: dict[str, Any] = dict(driver_features) if driver_features else {} - features.setdefault("use_pool", True) + features.setdefault("enable_pooling", True) json_serializer = features.setdefault("json_serializer", to_json) json_deserializer = features.setdefault("json_deserializer", from_json) diff --git a/sqlspec/adapters/psqlpy/config.py b/sqlspec/adapters/psqlpy/config.py index c9a5abc4c..06b09eaba 100644 --- a/sqlspec/adapters/psqlpy/config.py +++ b/sqlspec/adapters/psqlpy/config.py @@ -148,10 +148,6 @@ class PsqlpyDriverFeatures(TypedDict): """Enable ``google_columnar_engine.enable_columnar_scan=on`` at session startup or per transaction.""" pooler_mode: NotRequired[Literal["session", "transaction", "statement", "auto"] | None] """External pooler compatibility mode; executes queries with ``prepared=False`` when not ``session``.""" - pgbouncer: NotRequired[bool] - """Alias for ``pooler_mode='transaction'`` disabling prepared execution.""" - pgbouncer_mode: NotRequired[bool] - """Alias for ``pooler_mode='transaction'`` disabling prepared execution.""" direct_dsn: NotRequired["str | Mapping[str, Any] | None"] """Direct PostgreSQL DSN or connection mapping bypassing the pooler for migrations and ``LISTEN/NOTIFY``.""" json_serializer: NotRequired["Callable[[Any], str]"] @@ -288,9 +284,7 @@ def __init__( statement_config, driver_features = apply_driver_features(statement_config, driver_features) features_dict = dict(driver_features) if driver_features else {} - is_transaction_pooler = features_dict.get("pooler_mode") in {"transaction", "statement"} or bool( - features_dict.get("pgbouncer") or features_dict.get("pgbouncer_mode") - ) + is_transaction_pooler = features_dict.get("pooler_mode") in {"transaction", "statement"} if bool(features_dict.get("enable_alloydb_columnar", False)) and not is_transaction_pooler: columnar_opt = "-c google_columnar_engine.enable_columnar_scan=on" existing_options = connection_config.get("options") @@ -413,8 +407,6 @@ def derive_direct_config(self) -> Self: conn_cfg.update(build_connection_config(direct_dsn)) features: dict[str, Any] = dict(self.driver_features) features["pooler_mode"] = "session" - features["pgbouncer"] = False - features["pgbouncer_mode"] = False features["direct_dsn"] = None if self._user_connection_hook is not None: features["on_connection_create"] = self._user_connection_hook diff --git a/sqlspec/adapters/psqlpy/core.py b/sqlspec/adapters/psqlpy/core.py index 9bead644d..97cdb66ee 100644 --- a/sqlspec/adapters/psqlpy/core.py +++ b/sqlspec/adapters/psqlpy/core.py @@ -226,12 +226,9 @@ def apply_driver_features( features.setdefault("enable_pgvector", PGVECTOR_INSTALLED) features.setdefault("enable_paradedb", True) features.setdefault("enable_pg_textsearch", True) - pgbouncer_enabled = bool(features.get("pgbouncer") or features.get("pgbouncer_mode")) - features["pgbouncer"] = pgbouncer_enabled - features["pgbouncer_mode"] = pgbouncer_enabled raw_pooler_mode = features.get("pooler_mode") if raw_pooler_mode is None: - features["pooler_mode"] = "transaction" if pgbouncer_enabled else "session" + features["pooler_mode"] = "session" else: normalized_mode = str(raw_pooler_mode).lower() if normalized_mode not in {"session", "transaction", "statement", "auto"}: @@ -239,10 +236,7 @@ def apply_driver_features( f"Invalid pooler_mode: {raw_pooler_mode!r}. Expected 'session', 'transaction', 'statement', or 'auto'." ) raise ImproperConfigurationError(msg) - features["pooler_mode"] = "transaction" if normalized_mode == "auto" and pgbouncer_enabled else normalized_mode - if features["pooler_mode"] in {"transaction", "statement"}: - features["pgbouncer"] = True - features["pgbouncer_mode"] = True + features["pooler_mode"] = "session" if normalized_mode == "auto" else normalized_mode features.setdefault("direct_dsn", None) features.setdefault("enable_alloydb_columnar", False) diff --git a/sqlspec/adapters/psqlpy/driver.py b/sqlspec/adapters/psqlpy/driver.py index 41932c78c..7aae5c332 100644 --- a/sqlspec/adapters/psqlpy/driver.py +++ b/sqlspec/adapters/psqlpy/driver.py @@ -98,9 +98,7 @@ def __init__( self._data_dictionary: PsqlpyDataDictionary | None = None self._transaction_active = False pooler_mode = self.driver_features.get("pooler_mode", "session") - self._pooler_unprepared = pooler_mode in {"transaction", "statement"} or bool( - self.driver_features.get("pgbouncer") or self.driver_features.get("pgbouncer_mode") - ) + self._pooler_unprepared = pooler_mode in {"transaction", "statement"} @property def data_dictionary(self) -> "PsqlpyDataDictionary": diff --git a/sqlspec/adapters/psqlpy/events/_hub.py b/sqlspec/adapters/psqlpy/events/_hub.py index 1718f19ae..e8c767ed3 100644 --- a/sqlspec/adapters/psqlpy/events/_hub.py +++ b/sqlspec/adapters/psqlpy/events/_hub.py @@ -238,14 +238,7 @@ def _resolve_listener_config(self) -> "PsqlpyConfig": return self._listener_config raw_features = getattr(self._config, "driver_features", None) driver_features = raw_features if isinstance(raw_features, dict) else {} - pooler_mode = str( - driver_features.get("pooler_mode") - or ( - "transaction" - if driver_features.get("pgbouncer") or driver_features.get("pgbouncer_mode") - else "session" - ) - ).lower() + pooler_mode = str(driver_features.get("pooler_mode") or "session").lower() raw_direct_dsn = driver_features.get("direct_dsn") direct_dsn = raw_direct_dsn if isinstance(raw_direct_dsn, (str, Mapping)) and raw_direct_dsn else None if pooler_mode in {"transaction", "statement"} and not direct_dsn: diff --git a/sqlspec/adapters/psycopg/config.py b/sqlspec/adapters/psycopg/config.py index 01a036cd2..dbba98b7b 100644 --- a/sqlspec/adapters/psycopg/config.py +++ b/sqlspec/adapters/psycopg/config.py @@ -96,7 +96,6 @@ "min_size", "name", "num_workers", - "null_pool", "open", "reconnect_failed", "reconnect_timeout", @@ -190,8 +189,6 @@ class PsycopgPoolParams(PsycopgConnectionParams): """Callback invoked when pool reconnection attempts exceed ``reconnect_timeout``.""" kwargs: NotRequired["dict[str, Any]"] """Explicit connection keyword mapping passed to ``psycopg.connect``.""" - null_pool: NotRequired[bool] - """Use ``NullConnectionPool`` / ``AsyncNullConnectionPool`` so idle connections are never retained.""" class PsycopgDriverFeatures(TypedDict): @@ -227,14 +224,10 @@ class PsycopgDriverFeatures(TypedDict): """Enable binary streaming format for Arrow ingestion via PGCOPY binary protocol.""" pooler_mode: NotRequired[Literal["session", "transaction", "statement", "auto"]] """External pooler compatibility mode; sets ``prepare_threshold=None`` in ``transaction`` or ``statement`` mode.""" - pgbouncer: NotRequired[bool] - """Alias for ``pooler_mode='transaction'`` disabling server-side prepared statements.""" - pgbouncer_mode: NotRequired[bool] - """Alias for ``pooler_mode='transaction'`` disabling server-side prepared statements.""" direct_dsn: NotRequired["str | Mapping[str, Any] | None"] """Direct PostgreSQL DSN or connection mapping bypassing the pooler for migrations and ``LISTEN/NOTIFY``.""" - null_pool: NotRequired[bool] - """Enable ``NullConnectionPool`` / ``AsyncNullConnectionPool`` for serverless or external pooler deployments.""" + enable_pooling: NotRequired[bool] + """Enable connection pooling (defaults to ``True``; ``False`` uses ``NullConnectionPool`` / ``AsyncNullConnectionPool``).""" def build_connection_config(connection_config: "PsycopgPoolParams | Mapping[str, Any] | None") -> dict[str, Any]: @@ -358,11 +351,7 @@ def __init__( features_dict.setdefault("enable_alloydb_iam_auth", False) features_dict.setdefault("alloydb_ip_type", "PRIVATE") pooler_mode = features_dict.get("pooler_mode", "session") - is_transaction_pooler = bool( - pooler_mode in {"transaction", "statement"} - or features_dict.get("pgbouncer") - or features_dict.get("pgbouncer_mode") - ) + is_transaction_pooler = pooler_mode in {"transaction", "statement"} if is_transaction_pooler: connection_config.setdefault("prepare_threshold", None) if features_dict.get("enable_alloydb_columnar") and not is_transaction_pooler: @@ -407,17 +396,11 @@ def derive_direct_config(self) -> Self: conn_cfg.update(build_connection_config(direct_dsn)) else: conn_cfg["conninfo"] = str(direct_dsn) - is_pooled = bool( - self.driver_features.get("pooler_mode") in {"transaction", "statement"} - or self.driver_features.get("pgbouncer") - or self.driver_features.get("pgbouncer_mode") - ) + is_pooled = self.driver_features.get("pooler_mode") in {"transaction", "statement"} if is_pooled and conn_cfg.get("prepare_threshold") is None: conn_cfg.pop("prepare_threshold", None) features = dict(self.driver_features) features["pooler_mode"] = "session" - features["pgbouncer"] = False - features["pgbouncer_mode"] = False features["direct_dsn"] = None return type(self)( connection_config=conn_cfg, @@ -564,7 +547,6 @@ def _setup_alloydb_connector( def _create_pool(self) -> "ConnectionPool": """Create the actual connection pool.""" all_config = dict(self.connection_config) - all_config.pop("null_pool", None) auth_token_provider = all_config.pop("auth_token_provider", None) default_max_lifetime = _AUTH_TOKEN_DEFAULT_MAX_LIFETIME if auth_token_provider is not None else 3600.0 @@ -617,9 +599,9 @@ def _chained_configure(conn: "PsycopgSyncConnection") -> None: self._setup_alloydb_connector(all_config, pool_parameters) conninfo = None - is_null_pool = bool(self.connection_config.get("null_pool") or self.driver_features.get("null_pool")) - pool_cls = NullConnectionPool if is_null_pool else ConnectionPool - if is_null_pool: + pooling_disabled = not bool(self.driver_features.get("enable_pooling", True)) + pool_cls = NullConnectionPool if pooling_disabled else ConnectionPool + if pooling_disabled: pool_parameters.pop("min_size", None) if conninfo: @@ -825,11 +807,7 @@ def __init__( "on_connection_create", None ) pooler_mode = features_dict.get("pooler_mode", "session") - is_transaction_pooler = bool( - pooler_mode in {"transaction", "statement"} - or features_dict.get("pgbouncer") - or features_dict.get("pgbouncer_mode") - ) + is_transaction_pooler = pooler_mode in {"transaction", "statement"} if is_transaction_pooler: connection_config.setdefault("prepare_threshold", None) if features_dict.get("enable_alloydb_columnar") and not is_transaction_pooler: @@ -871,17 +849,11 @@ def derive_direct_config(self) -> Self: conn_cfg.update(build_connection_config(direct_dsn)) else: conn_cfg["conninfo"] = str(direct_dsn) - is_pooled = bool( - self.driver_features.get("pooler_mode") in {"transaction", "statement"} - or self.driver_features.get("pgbouncer") - or self.driver_features.get("pgbouncer_mode") - ) + is_pooled = self.driver_features.get("pooler_mode") in {"transaction", "statement"} if is_pooled and conn_cfg.get("prepare_threshold") is None: conn_cfg.pop("prepare_threshold", None) features = dict(self.driver_features) features["pooler_mode"] = "session" - features["pgbouncer"] = False - features["pgbouncer_mode"] = False features["direct_dsn"] = None return type(self)( connection_config=conn_cfg, @@ -982,7 +954,6 @@ async def _create_pool(self) -> "AsyncConnectionPool": """Create the actual async connection pool.""" all_config = dict(self.connection_config) - all_config.pop("null_pool", None) auth_token_provider = all_config.pop("auth_token_provider", None) default_max_lifetime = _AUTH_TOKEN_DEFAULT_MAX_LIFETIME if auth_token_provider is not None else 3600.0 @@ -1029,9 +1000,9 @@ async def _chained_async_configure(conn: "PsycopgAsyncConnection") -> None: kwargs = all_config.pop("kwargs", {}) all_config.update(kwargs) - is_null_pool = bool(self.connection_config.get("null_pool") or self.driver_features.get("null_pool")) - pool_cls = AsyncNullConnectionPool if is_null_pool else AsyncConnectionPool - if is_null_pool: + pooling_disabled = not bool(self.driver_features.get("enable_pooling", True)) + pool_cls = AsyncNullConnectionPool if pooling_disabled else AsyncConnectionPool + if pooling_disabled: pool_parameters.pop("min_size", None) if conninfo: diff --git a/sqlspec/adapters/psycopg/core.py b/sqlspec/adapters/psycopg/core.py index 95719d5bc..6539f8bdc 100644 --- a/sqlspec/adapters/psycopg/core.py +++ b/sqlspec/adapters/psycopg/core.py @@ -215,12 +215,9 @@ def apply_driver_features( features.setdefault("enable_pg_textsearch", True) features.setdefault("enable_alloydb_columnar", False) features.setdefault("enable_pgcopy_binary", False) - pgbouncer_enabled = bool(features.get("pgbouncer") or features.get("pgbouncer_mode")) - features["pgbouncer"] = pgbouncer_enabled - features["pgbouncer_mode"] = pgbouncer_enabled raw_pooler_mode = features.get("pooler_mode") if raw_pooler_mode is None: - features["pooler_mode"] = "transaction" if pgbouncer_enabled else "session" + features["pooler_mode"] = "session" else: normalized_mode = str(raw_pooler_mode).lower() if normalized_mode not in {"session", "transaction", "statement", "auto"}: @@ -228,10 +225,7 @@ def apply_driver_features( f"Invalid pooler_mode: {raw_pooler_mode!r}. Expected 'session', 'transaction', 'statement', or 'auto'." ) raise ImproperConfigurationError(msg) - features["pooler_mode"] = "transaction" if normalized_mode == "auto" and pgbouncer_enabled else normalized_mode - if features["pooler_mode"] in {"transaction", "statement"}: - features["pgbouncer"] = True - features["pgbouncer_mode"] = True + features["pooler_mode"] = "session" if normalized_mode == "auto" else normalized_mode features.setdefault("direct_dsn", None) parameter_config = _parameter_config(driver_profile, serializer, deserializer) diff --git a/sqlspec/adapters/psycopg/driver.py b/sqlspec/adapters/psycopg/driver.py index 82e16a844..fadd02cca 100644 --- a/sqlspec/adapters/psycopg/driver.py +++ b/sqlspec/adapters/psycopg/driver.py @@ -744,11 +744,7 @@ def _should_set_local_columnar(self) -> bool: """Return True when AlloyDB columnar GUC must be set transaction-locally.""" if not self.driver_features.get("enable_alloydb_columnar"): return False - return bool( - self.driver_features.get("pooler_mode") in {"transaction", "statement"} - or self.driver_features.get("pgbouncer") - or self.driver_features.get("pgbouncer_mode") - ) + return self.driver_features.get("pooler_mode") in {"transaction", "statement"} class PsycopgAsyncExceptionHandler(BaseAsyncExceptionHandler): @@ -1326,11 +1322,7 @@ def _should_set_local_columnar(self) -> bool: """Return True when AlloyDB columnar GUC must be set transaction-locally.""" if not self.driver_features.get("enable_alloydb_columnar"): return False - return bool( - self.driver_features.get("pooler_mode") in {"transaction", "statement"} - or self.driver_features.get("pgbouncer") - or self.driver_features.get("pgbouncer_mode") - ) + return self.driver_features.get("pooler_mode") in {"transaction", "statement"} def _attribute_pipeline_failure( diff --git a/sqlspec/adapters/psycopg/events/_hub.py b/sqlspec/adapters/psycopg/events/_hub.py index 1fbdced73..b7deea7a7 100644 --- a/sqlspec/adapters/psycopg/events/_hub.py +++ b/sqlspec/adapters/psycopg/events/_hub.py @@ -174,14 +174,7 @@ def _resolve_listener_config(self) -> "PsycopgAsyncConfig": return self._listener_config raw_features = getattr(self._config, "driver_features", None) driver_features = raw_features if isinstance(raw_features, dict) else {} - pooler_mode = str( - driver_features.get("pooler_mode") - or ( - "transaction" - if driver_features.get("pgbouncer") or driver_features.get("pgbouncer_mode") - else "session" - ) - ).lower() + pooler_mode = str(driver_features.get("pooler_mode") or "session").lower() raw_direct_dsn = driver_features.get("direct_dsn") direct_dsn = raw_direct_dsn if isinstance(raw_direct_dsn, (str, Mapping)) and raw_direct_dsn else None if pooler_mode in {"transaction", "statement"} and not direct_dsn: @@ -441,14 +434,7 @@ def _resolve_listener_config(self) -> "PsycopgSyncConfig": return self._listener_config raw_features = getattr(self._config, "driver_features", None) driver_features = raw_features if isinstance(raw_features, dict) else {} - pooler_mode = str( - driver_features.get("pooler_mode") - or ( - "transaction" - if driver_features.get("pgbouncer") or driver_features.get("pgbouncer_mode") - else "session" - ) - ).lower() + pooler_mode = str(driver_features.get("pooler_mode") or "session").lower() raw_direct_dsn = driver_features.get("direct_dsn") direct_dsn = raw_direct_dsn if isinstance(raw_direct_dsn, (str, Mapping)) and raw_direct_dsn else None if pooler_mode in {"transaction", "statement"} and not direct_dsn: diff --git a/sqlspec/adapters/pymssql/config.py b/sqlspec/adapters/pymssql/config.py index b81f070ff..529f93f2f 100644 --- a/sqlspec/adapters/pymssql/config.py +++ b/sqlspec/adapters/pymssql/config.py @@ -89,8 +89,8 @@ class PymssqlDriverFeatures(TypedDict): """Enable durable table-backed SQLSpec event queue support.""" events_backend: NotRequired[Literal["poll_queue"]] """Event channel backend (``poll_queue`` is the sole SQL Server backend).""" - use_pool: NotRequired[bool] - """Enable driver-level connection pooling.""" + enable_pooling: NotRequired[bool] + """Enable driver-level connection pooling (defaults to ``True``).""" enable_uuid_conversion: NotRequired[bool] """Enable automatic UUID string-to-object conversions.""" diff --git a/sqlspec/adapters/pymssql/core.py b/sqlspec/adapters/pymssql/core.py index 51aea9fd0..2feefddb1 100644 --- a/sqlspec/adapters/pymssql/core.py +++ b/sqlspec/adapters/pymssql/core.py @@ -215,7 +215,7 @@ def apply_driver_features( features: dict[str, Any] = dict(driver_features) if driver_features else {} json_serializer = features.setdefault("json_serializer", to_json) json_deserializer = features.setdefault("json_deserializer", from_json) - features.setdefault("use_pool", True) + features.setdefault("enable_pooling", True) features.setdefault("events_backend", "poll_queue") features.setdefault("enable_uuid_conversion", False) diff --git a/tests/unit/adapters/test_asyncpg/test_config.py b/tests/unit/adapters/test_asyncpg/test_config.py index f54ed3d8d..05390c655 100644 --- a/tests/unit/adapters/test_asyncpg/test_config.py +++ b/tests/unit/adapters/test_asyncpg/test_config.py @@ -403,6 +403,9 @@ async def fake_create_pool(*, init: Any = None, **kwargs: Any) -> object: monkeypatch.setattr(config_mod, "asyncpg_create_pool", fake_create_pool) + assert "pgbouncer" not in AsyncpgDriverFeatures.__annotations__ + assert "pgbouncer_mode" not in AsyncpgDriverFeatures.__annotations__ + assert "pgbouncer" not in AsyncpgConnectionConfig.__annotations__ config = AsyncpgConfig( connection_config={"dsn": "postgresql://localhost/test"}, driver_features={"pooler_mode": "transaction"} ) @@ -410,18 +413,6 @@ async def fake_create_pool(*, init: Any = None, **kwargs: Any) -> object: await config._create_pool() # pyright: ignore[reportPrivateUsage] assert captured_connect_kwargs["statement_cache_size"] == 0 - alias_config = AsyncpgConfig( - connection_config={"dsn": "postgresql://localhost/test"}, driver_features={"pgbouncer": True} - ) - assert alias_config.driver_features["pooler_mode"] == "transaction" - assert alias_config.connection_config["statement_cache_size"] == 0 - - mode_alias_config = AsyncpgConfig( - connection_config={"dsn": "postgresql://localhost/test"}, driver_features={"pgbouncer_mode": True} - ) - assert mode_alias_config.driver_features["pooler_mode"] == "transaction" - assert mode_alias_config.connection_config["statement_cache_size"] == 0 - conn = AsyncMock() conn.is_in_transaction = lambda: True conn.execute.return_value = "UPDATE 1" @@ -486,12 +477,12 @@ async def test_asyncpg_columnar_engine_session_vs_transaction() -> None: "google_columnar_engine.enable_columnar_scan": "on" } - auto_pgbouncer_config = AsyncpgConfig( + tx_columnar_config = AsyncpgConfig( connection_config={"dsn": "postgresql://localhost/test"}, - driver_features={"pooler_mode": "auto", "pgbouncer": True, "enable_alloydb_columnar": True}, + driver_features={"pooler_mode": "transaction", "enable_alloydb_columnar": True}, ) - assert auto_pgbouncer_config.driver_features["pooler_mode"] == "transaction" - assert "server_settings" not in auto_pgbouncer_config.connection_config + assert tx_columnar_config.driver_features["pooler_mode"] == "transaction" + assert "server_settings" not in tx_columnar_config.connection_config tx_config = AsyncpgConfig( connection_config={"dsn": "postgresql://localhost/test"}, diff --git a/tests/unit/adapters/test_mssql_python/test_config.py b/tests/unit/adapters/test_mssql_python/test_config.py index 436d7613b..41ada586a 100644 --- a/tests/unit/adapters/test_mssql_python/test_config.py +++ b/tests/unit/adapters/test_mssql_python/test_config.py @@ -10,6 +10,7 @@ from sqlspec.adapters.mssql_python.config import ( MssqlPythonConfig, MssqlPythonConnectionParams, + MssqlPythonDriverFeatures, _normalize_mssql_python_init, ) from sqlspec.adapters.mssql_python.core import build_connection_config @@ -318,12 +319,15 @@ def test_second_pool_same_params_no_warn(monkeypatch: pytest.MonkeyPatch) -> Non def test_normalize_mssql_python_init_returns_config_features_and_hook() -> None: seen: list[DummyConnection] = [] + assert "use_pool" not in MssqlPythonDriverFeatures.__annotations__ + assert "enable_pooling" in MssqlPythonDriverFeatures.__annotations__ + normalized, features, hook = _normalize_mssql_python_init( - {"server": "localhost"}, {"on_connection_create": seen.append, "use_pool": False} + {"server": "localhost"}, {"on_connection_create": seen.append, "enable_pooling": False} ) assert normalized == {"server": "localhost"} - assert features["use_pool"] is False + assert features["enable_pooling"] is False assert "on_connection_create" not in features assert hook is not None hook(cast(Any, DummyConnection())) diff --git a/tests/unit/adapters/test_postgres_native_options.py b/tests/unit/adapters/test_postgres_native_options.py index 8f76f8d18..27615733f 100644 --- a/tests/unit/adapters/test_postgres_native_options.py +++ b/tests/unit/adapters/test_postgres_native_options.py @@ -47,10 +47,12 @@ async def test_asyncpg_custom_codec_is_registered_after_internal_setup() -> None } -def test_asyncpg_pgbouncer_disables_prepared_cache_consistently() -> None: - config = AsyncpgConfig(connection_config={"pgbouncer": True, "statement_cache_size": 128}) +def test_asyncpg_pooler_mode_transaction_disables_prepared_cache_consistently() -> None: + config = AsyncpgConfig( + connection_config={"statement_cache_size": 128}, driver_features={"pooler_mode": "transaction"} + ) assert config.connection_config["statement_cache_size"] == 0 - assert config.driver_features["pgbouncer"] is True + assert config.driver_features["pooler_mode"] == "transaction" @pytest.mark.parametrize( @@ -62,7 +64,9 @@ async def test_null_pool_preserves_throttling_without_leaking_flag( ) -> None: constructor = MagicMock() monkeypatch.setattr(psycopg_module, pool_name, constructor) - config = config_type(connection_config={"null_pool": True, "open": False, "max_size": 7, "max_waiting": 2}) + config = config_type( + connection_config={"open": False, "max_size": 7, "max_waiting": 2}, driver_features={"enable_pooling": False} + ) if config_type is PsycopgAsyncConfig: await config._create_pool() else: @@ -70,8 +74,8 @@ async def test_null_pool_preserves_throttling_without_leaking_flag( args = constructor.call_args.kwargs assert args["max_size"] == 7 assert args["max_waiting"] == 2 - assert "null_pool" not in args["kwargs"] - assert "null_pool" not in config._connection_kwargs()[2] + assert "enable_pooling" not in args["kwargs"] + assert "enable_pooling" not in config._connection_kwargs()[2] @pytest.mark.parametrize("builder", [asyncpg_cockroach_config, psycopg_cockroach_config]) @@ -119,7 +123,7 @@ async def test_pgbouncer_stack_cancellation_keeps_transaction_cleanup() -> None: transaction.__aenter__ = AsyncMock() transaction.__aexit__ = AsyncMock(return_value=False) connection.transaction.return_value = transaction - driver = AsyncpgDriver(connection, driver_features={"pgbouncer": True}) + driver = AsyncpgDriver(connection, driver_features={"pooler_mode": "transaction"}) stack = StatementStack().push_execute("SELECT 1") with pytest.raises(asyncio.CancelledError): await driver.execute_stack(stack) diff --git a/tests/unit/adapters/test_psqlpy/test_config.py b/tests/unit/adapters/test_psqlpy/test_config.py index 9a75f2129..993d0bdc9 100644 --- a/tests/unit/adapters/test_psqlpy/test_config.py +++ b/tests/unit/adapters/test_psqlpy/test_config.py @@ -287,6 +287,8 @@ async def test_psqlpy_session_context_rolls_back_open_transaction_on_exit() -> N @pytest.mark.anyio async def test_psqlpy_pooler_mode_transaction_disables_prepared_statements() -> None: """Transaction pooler mode should construct a valid ConnectionPool and pass prepared=False on execution.""" + assert "pgbouncer" not in PsqlpyDriverFeatures.__annotations__ + assert "pgbouncer_mode" not in PsqlpyDriverFeatures.__annotations__ config = PsqlpyConfig( connection_config={"dsn": "postgresql://localhost:6432/test"}, driver_features={"pooler_mode": "transaction"} ) @@ -294,13 +296,6 @@ async def test_psqlpy_pooler_mode_transaction_disables_prepared_statements() -> pool = await config.create_pool() pool.close() - alias_config = PsqlpyConfig( - connection_config={"dsn": "postgresql://localhost:6432/test"}, driver_features={"pgbouncer": True} - ) - assert alias_config.driver_features["pooler_mode"] == "transaction" - alias_pool = await alias_config.create_pool() - alias_pool.close() - conn = AsyncMock() conn.fetch.return_value = _ExtensionQueryResult(set()) conn.execute.return_value = "CREATE TABLE" @@ -335,12 +330,12 @@ async def test_psqlpy_derive_direct_config_and_columnar_engine() -> None: ) assert auto_direct_config.connection_config["options"] == "-c google_columnar_engine.enable_columnar_scan=on" - auto_pgbouncer_config = PsqlpyConfig( + tx_columnar_config = PsqlpyConfig( connection_config={"dsn": "postgresql://pooler:6432/test"}, - driver_features={"pooler_mode": "auto", "pgbouncer": True, "enable_alloydb_columnar": True}, + driver_features={"pooler_mode": "transaction", "enable_alloydb_columnar": True}, ) - assert auto_pgbouncer_config.driver_features["pooler_mode"] == "transaction" - assert "options" not in auto_pgbouncer_config.connection_config + assert tx_columnar_config.driver_features["pooler_mode"] == "transaction" + assert "options" not in tx_columnar_config.connection_config with pytest.raises(ImproperConfigurationError, match="Invalid pooler_mode"): PsqlpyConfig( diff --git a/tests/unit/adapters/test_psycopg/test_config.py b/tests/unit/adapters/test_psycopg/test_config.py index a225cb3e4..121198e5d 100644 --- a/tests/unit/adapters/test_psycopg/test_config.py +++ b/tests/unit/adapters/test_psycopg/test_config.py @@ -13,6 +13,8 @@ from sqlspec.adapters.psycopg._typing import PsycopgAsyncSessionContext, PsycopgSyncSessionContext from sqlspec.adapters.psycopg.config import ( PsycopgAsyncConfig, + PsycopgConnectionParams, + PsycopgDriverFeatures, PsycopgPoolParams, PsycopgSyncConfig, build_connection_config, @@ -701,11 +703,11 @@ def test_psycopg_pooler_mode_transaction_sets_prepare_threshold_none_and_alloydb assert async_config.connection_config["prepare_threshold"] is None assert async_config.driver_features["pooler_mode"] == "statement" - alias_config = PsycopgSyncConfig( - connection_config={"conninfo": "postgresql://localhost/db"}, driver_features={"pgbouncer": True} - ) - assert alias_config.connection_config["prepare_threshold"] is None - assert alias_config.driver_features["pooler_mode"] == "transaction" + assert "pgbouncer" not in PsycopgDriverFeatures.__annotations__ + assert "pgbouncer_mode" not in PsycopgDriverFeatures.__annotations__ + assert "null_pool" not in PsycopgDriverFeatures.__annotations__ + assert "null_pool" not in PsycopgConnectionParams.__annotations__ + assert "enable_pooling" in PsycopgDriverFeatures.__annotations__ mock_connector = MagicMock() expected_conn = MagicMock() @@ -750,15 +752,13 @@ async def test_psycopg_enable_alloydb_columnar_and_derive_direct_config() -> Non ) assert "-c google_columnar_engine.enable_columnar_scan=on" in auto_sync.connection_config["options"] - auto_pgbouncer_sync = PsycopgSyncConfig( + tx_sync = PsycopgSyncConfig( connection_config={"conninfo": "postgresql://pooler:6432/db"}, - driver_features={"enable_alloydb_columnar": True, "pooler_mode": "auto", "pgbouncer": True}, + driver_features={"enable_alloydb_columnar": True, "pooler_mode": "transaction"}, ) - assert auto_pgbouncer_sync.driver_features["pooler_mode"] == "transaction" - assert auto_pgbouncer_sync.driver_features["pgbouncer"] is True - assert auto_pgbouncer_sync.driver_features["pgbouncer_mode"] is True - assert "options" not in auto_pgbouncer_sync.connection_config - assert auto_pgbouncer_sync.connection_config["prepare_threshold"] is None + assert tx_sync.driver_features["pooler_mode"] == "transaction" + assert "options" not in tx_sync.connection_config + assert tx_sync.connection_config["prepare_threshold"] is None with pytest.raises(ImproperConfigurationError, match="Invalid pooler_mode"): PsycopgSyncConfig( diff --git a/tests/unit/adapters/test_pymssql/test_config.py b/tests/unit/adapters/test_pymssql/test_config.py index 9c17d54fe..d4593b23f 100644 --- a/tests/unit/adapters/test_pymssql/test_config.py +++ b/tests/unit/adapters/test_pymssql/test_config.py @@ -5,7 +5,7 @@ import pytest from sqlspec.adapters.pymssql import build_connection_config -from sqlspec.adapters.pymssql.config import PymssqlConfig, PymssqlConnectionParams +from sqlspec.adapters.pymssql.config import PymssqlConfig, PymssqlConnectionParams, PymssqlDriverFeatures from sqlspec.adapters.pymssql.driver import PymssqlDriver from sqlspec.adapters.pymssql.pool import PymssqlConnectionPool @@ -167,8 +167,10 @@ def test_use_datetime2_default() -> None: def test_driver_features_mssql_parity() -> None: - """PymssqlConfig defaults use_pool to True and harmonizes driver features.""" + """PymssqlConfig defaults enable_pooling to True and harmonizes driver features.""" + assert "use_pool" not in PymssqlDriverFeatures.__annotations__ + assert "enable_pooling" in PymssqlDriverFeatures.__annotations__ config = PymssqlConfig() - assert config.driver_features.get("use_pool") is True + assert config.driver_features.get("enable_pooling") is True assert config.driver_features.get("events_backend") == "poll_queue" assert config.driver_features.get("enable_uuid_conversion") is False diff --git a/tests/unit/extensions/test_events/test_listener_hubs.py b/tests/unit/extensions/test_events/test_listener_hubs.py index e1cf2d486..60b221f2e 100644 --- a/tests/unit/extensions/test_events/test_listener_hubs.py +++ b/tests/unit/extensions/test_events/test_listener_hubs.py @@ -1035,7 +1035,7 @@ async def test_listener_hub_refuses_transaction_pooler_without_direct_dsn() -> N await psycopg_async_hub.subscribe("alerts") psycopg_sync_cfg = _PsycopgSyncConfig(_PsycopgSyncConnection()) - psycopg_sync_cfg.driver_features = {"pgbouncer": True, "direct_dsn": None} # type: ignore[attr-defined] + psycopg_sync_cfg.driver_features = {"pooler_mode": "transaction", "direct_dsn": None} # type: ignore[attr-defined] psycopg_sync_hub = PsycopgSyncListenerHub(psycopg_sync_cfg) # type: ignore[arg-type] with pytest.raises(ImproperConfigurationError, match="direct_dsn"): psycopg_sync_hub.subscribe("alerts") From ce4cfc5f5746cf41927a5cf7dc454710f09e4263 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 03:56:22 +0000 Subject: [PATCH 099/174] refactor(adapters): import each adapter's driver eagerly Importing an adapter package now imports its own driver once, at the top of _typing.py, and a missing driver raises the driver's own ModuleNotFoundError. asyncmy, db2, mysqlconnector, psqlpy, pymssql, sqlanywhere, and sybase drop import_optional placeholders, fallback error classes, and try/except fallbacks; SybaseError and SqlanywhereError are the native driver errors. aiomysql, mssql_python, and pymysql use a single TYPE_CHECKING block. oracledb loads rich, pyarrow, and numpy on first use, and psycopg resolves pgvector on first use, so importing those adapters no longer pulls optional packages. sqlanydb joins the extras dependency group so every driver is available for development. --- pyproject.toml | 1 + sqlspec/adapters/aiomysql/_typing.py | 3 +- sqlspec/adapters/asyncmy/_typing.py | 64 ++++--------------- sqlspec/adapters/db2/_typing.py | 31 +++------ sqlspec/adapters/db2/pool.py | 46 ++++++------- sqlspec/adapters/mssql_python/_typing.py | 3 +- sqlspec/adapters/mysqlconnector/_typing.py | 48 +++++++------- sqlspec/adapters/oracledb/_vector_handlers.py | 12 +++- sqlspec/adapters/oracledb/driver.py | 19 ++++-- sqlspec/adapters/oracledb/migrations.py | 22 ++++--- sqlspec/adapters/oracledb/type_converter.py | 13 +++- sqlspec/adapters/psqlpy/_typing.py | 62 ++++-------------- sqlspec/adapters/psqlpy/type_converter.py | 12 ++-- sqlspec/adapters/psycopg/type_converter.py | 9 ++- sqlspec/adapters/pymssql/_typing.py | 37 +++++------ sqlspec/adapters/pymysql/_typing.py | 3 +- sqlspec/adapters/sqlanywhere/_typing.py | 36 +++++------ sqlspec/adapters/sqlanywhere/pool.py | 18 ++---- sqlspec/adapters/sybase/_typing.py | 38 +++++------ sqlspec/adapters/sybase/pool.py | 17 +---- tests/unit/adapters/test_db2/test_pool.py | 15 ----- tests/unit/adapters/test_db2/test_typing.py | 47 +------------- .../test_psycopg/test_type_handlers.py | 19 ++---- tests/unit/adapters/test_sqlanywhere.py | 18 +----- tests/unit/adapters/test_sybase.py | 29 +-------- tests/unit/test_mypyc_sybase_surface.py | 6 +- uv.lock | 4 ++ 27 files changed, 209 insertions(+), 423 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 6fd2f187f..10ee865c2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -113,6 +113,7 @@ extras = [ "adlfs", "gcsfs", "dishka", + "sqlanydb>=1.0.14", ] lint = [ "mypy>=2.0.0", diff --git a/sqlspec/adapters/aiomysql/_typing.py b/sqlspec/adapters/aiomysql/_typing.py index c101c8f48..911dcf71c 100644 --- a/sqlspec/adapters/aiomysql/_typing.py +++ b/sqlspec/adapters/aiomysql/_typing.py @@ -61,8 +61,7 @@ class AiomysqlFieldTypeProtocol(Protocol): AiomysqlPool: TypeAlias = _AiomysqlPool AiomysqlPymysqlError: TypeAlias = _PymysqlError AiomysqlPymysqlMySQLError: TypeAlias = _PymysqlMySQLError - -if not TYPE_CHECKING: +else: AiomysqlConnection = Connection AiomysqlModule = _aiomysql AiomysqlRawCursor = _AiomysqlCursor diff --git a/sqlspec/adapters/asyncmy/_typing.py b/sqlspec/adapters/asyncmy/_typing.py index 3807eafbc..1c4f81617 100644 --- a/sqlspec/adapters/asyncmy/_typing.py +++ b/sqlspec/adapters/asyncmy/_typing.py @@ -8,54 +8,19 @@ import os from typing import TYPE_CHECKING, Any -if TYPE_CHECKING: - import asyncmy as _asyncmy - from asyncmy import Connection - from asyncmy import errors as _asyncmy_errors - from asyncmy.connection import LoadLocalFile - from asyncmy.connection import MySQLResult as _AsyncmyResult - from asyncmy.constants import FIELD_TYPE as _ASYNCMY_FIELD_TYPE - from asyncmy.cursors import RE_INSERT_VALUES as ASYNCMY_INSERT_VALUES_PATTERN - from asyncmy.cursors import Cursor as _AsyncmyCursor - from asyncmy.cursors import DictCursor as _AsyncmyDictCursor - from asyncmy.cursors import SSCursor as AsyncmySSCursor - from asyncmy.errors import ProgrammingError as _AsyncmyProgrammingError - from asyncmy.pool import Pool as _AsyncmyPool - from asyncmy.protocol import LoadLocalPacketWrapper as _LoadLocalPacketWrapper -else: - try: - import asyncmy as _asyncmy - from asyncmy import Connection - from asyncmy import errors as _asyncmy_errors - from asyncmy.connection import LoadLocalFile - from asyncmy.connection import MySQLResult as _AsyncmyResult - from asyncmy.constants import FIELD_TYPE as _ASYNCMY_FIELD_TYPE - from asyncmy.cursors import RE_INSERT_VALUES as ASYNCMY_INSERT_VALUES_PATTERN - from asyncmy.cursors import Cursor as _AsyncmyCursor - from asyncmy.cursors import DictCursor as _AsyncmyDictCursor - from asyncmy.cursors import SSCursor as AsyncmySSCursor - from asyncmy.errors import ProgrammingError as _AsyncmyProgrammingError - from asyncmy.pool import Pool as _AsyncmyPool - from asyncmy.protocol import LoadLocalPacketWrapper as _LoadLocalPacketWrapper - except ImportError: - import re - - _asyncmy = None - Connection = None - _asyncmy_errors = None - LoadLocalFile = None - _AsyncmyResult = object - _ASYNCMY_FIELD_TYPE = None - ASYNCMY_INSERT_VALUES_PATTERN = re.compile( - r"(VALUES\s*\()((?:(?!\bVALUES\b).)*?)(\)(?:\s*(?:ON\s+DUPLICATE\s+KEY\s+UPDATE|RETURNING)\b.*)?\Z", - re.IGNORECASE | re.DOTALL, - ) - _AsyncmyCursor = None - _AsyncmyDictCursor = None - AsyncmySSCursor = None - _AsyncmyProgrammingError = Exception - _AsyncmyPool = None - _LoadLocalPacketWrapper = None +import asyncmy as _asyncmy +from asyncmy import Connection +from asyncmy import errors as _asyncmy_errors +from asyncmy.connection import LoadLocalFile +from asyncmy.connection import MySQLResult as _AsyncmyResult +from asyncmy.constants import FIELD_TYPE as _ASYNCMY_FIELD_TYPE +from asyncmy.cursors import RE_INSERT_VALUES as ASYNCMY_INSERT_VALUES_PATTERN +from asyncmy.cursors import Cursor as _AsyncmyCursor +from asyncmy.cursors import DictCursor as _AsyncmyDictCursor +from asyncmy.cursors import SSCursor as AsyncmySSCursor +from asyncmy.errors import ProgrammingError as _AsyncmyProgrammingError +from asyncmy.pool import Pool as _AsyncmyPool +from asyncmy.protocol import LoadLocalPacketWrapper as _LoadLocalPacketWrapper from sqlspec.exceptions import SQLSpecError @@ -99,8 +64,7 @@ class AsyncmyFieldTypeProtocol(Protocol): AsyncmyModule: TypeAlias = AsyncmyModuleProtocol AsyncmyPool: TypeAlias = _AsyncmyPool AsyncmyRawCursor: TypeAlias = _AsyncmyCursor - -if not TYPE_CHECKING: +else: AsyncmyConnection = Connection AsyncmyDictCursor = _AsyncmyDictCursor AsyncmyError = _asyncmy_errors.Error diff --git a/sqlspec/adapters/db2/_typing.py b/sqlspec/adapters/db2/_typing.py index 0bba3961f..8ee4a4c8a 100644 --- a/sqlspec/adapters/db2/_typing.py +++ b/sqlspec/adapters/db2/_typing.py @@ -4,38 +4,23 @@ import importlib from typing import TYPE_CHECKING, Any, cast -from sqlspec.exceptions import MissingDependencyError -from sqlspec.typing import import_optional_attr -from sqlspec.utils.module_loader import import_optional - - -class _Db2UnavailableError(Exception): - """Fallback Db2 exception base when ibm_db_dbi is not installed.""" +import ibm_db +import ibm_db_dbi +from ibm_db_dbi import AsyncConnection as Db2AsyncConnection +from ibm_db_dbi import AsyncCursor as Db2AsyncRawCursor +from ibm_db_dbi import Connection as Db2SyncConnection +from ibm_db_dbi import Cursor as Db2RawCursor +from ibm_db_dbi import Error as Db2Error +from sqlspec.exceptions import MissingDependencyError if TYPE_CHECKING: from collections.abc import Callable from types import TracebackType - from ibm_db_dbi import AsyncConnection as Db2AsyncConnection - from ibm_db_dbi import AsyncCursor as Db2AsyncRawCursor - from ibm_db_dbi import Connection as Db2SyncConnection - from ibm_db_dbi import Cursor as Db2RawCursor - from ibm_db_dbi import Error as Db2Error - from sqlspec.adapters.db2.driver import Db2AsyncDriver, Db2SyncDriver from sqlspec.core import StatementConfig -if not TYPE_CHECKING: - Db2SyncConnection = import_optional_attr("ibm_db_dbi", "Connection") or Any - Db2RawCursor = import_optional_attr("ibm_db_dbi", "Cursor") or Any - Db2AsyncConnection = import_optional_attr("ibm_db_dbi", "AsyncConnection") or Any - Db2AsyncRawCursor = import_optional_attr("ibm_db_dbi", "AsyncCursor") or Any - Db2Error = import_optional_attr("ibm_db_dbi", "Error") or _Db2UnavailableError - -ibm_db = import_optional("ibm_db") -ibm_db_dbi = import_optional("ibm_db_dbi") - __all__ = ( "Db2AsyncConnection", "Db2AsyncCursor", diff --git a/sqlspec/adapters/db2/pool.py b/sqlspec/adapters/db2/pool.py index b31c0e409..76f5dddc8 100644 --- a/sqlspec/adapters/db2/pool.py +++ b/sqlspec/adapters/db2/pool.py @@ -9,10 +9,11 @@ from contextlib import contextmanager from typing import TYPE_CHECKING, Any +from sqlspec.adapters.db2 import _typing as typing_module +from sqlspec.adapters.db2._typing import ibm_db_dbi from sqlspec.adapters.db2.core import build_dsn_string -from sqlspec.exceptions import ConnectionTimeoutError, DatabaseConnectionError, MissingDependencyError +from sqlspec.exceptions import ConnectionTimeoutError, DatabaseConnectionError from sqlspec.utils.logging import POOL_LOGGER_NAME, get_logger, log_with_context -from sqlspec.utils.module_loader import import_optional from sqlspec.utils.uuids import uuid4 if TYPE_CHECKING: @@ -23,7 +24,7 @@ logger = get_logger(POOL_LOGGER_NAME) _ADAPTER_NAME = "db2" -_IBM_DB_DBI: "Any | None" = None +_IBM_DB_DBI = ibm_db_dbi class Db2SyncConnectionPool: @@ -91,9 +92,14 @@ def new_connection(self) -> Any: Raises: MissingDependencyError: When ibm_db is not installed. """ - ibm_db_dbi = _require_ibm_db_dbi() - autocommit_mode = ibm_db_dbi.SQL_AUTOCOMMIT_ON if self._autocommit else ibm_db_dbi.SQL_AUTOCOMMIT_OFF - connection = ibm_db_dbi.connect(self._dsn, "", "", "", "", {ibm_db_dbi.SQL_ATTR_AUTOCOMMIT: autocommit_mode}) + autocommit_mode = ( + typing_module.ibm_db_dbi.SQL_AUTOCOMMIT_ON + if self._autocommit + else typing_module.ibm_db_dbi.SQL_AUTOCOMMIT_OFF + ) + connection = typing_module.ibm_db_dbi.connect( + self._dsn, "", "", "", "", {typing_module.ibm_db_dbi.SQL_ATTR_AUTOCOMMIT: autocommit_mode} + ) if self._on_connection_create is not None: try: @@ -374,10 +380,13 @@ async def new_connection(self) -> Any: Raises: MissingDependencyError: When ibm_db is not installed. """ - ibm_db_dbi = _require_ibm_db_dbi() - autocommit_mode = ibm_db_dbi.SQL_AUTOCOMMIT_ON if self._autocommit else ibm_db_dbi.SQL_AUTOCOMMIT_OFF - connection = await ibm_db_dbi.AsyncConnection.connect( - self._dsn, "", "", "", "", {ibm_db_dbi.SQL_ATTR_AUTOCOMMIT: autocommit_mode} + autocommit_mode = ( + typing_module.ibm_db_dbi.SQL_AUTOCOMMIT_ON + if self._autocommit + else typing_module.ibm_db_dbi.SQL_AUTOCOMMIT_OFF + ) + connection = await typing_module.ibm_db_dbi.AsyncConnection.connect( + self._dsn, "", "", "", "", {typing_module.ibm_db_dbi.SQL_ATTR_AUTOCOMMIT: autocommit_mode} ) if self._on_connection_create is not None: try: @@ -550,20 +559,3 @@ async def _close_connection(self, connection: Any) -> None: """Close a connection, suppressing driver errors.""" with contextlib.suppress(Exception): await connection.close() - - -def _require_ibm_db_dbi() -> Any: - """Return the ``ibm_db_dbi`` module, importing it on first use. - - Returns: - Any: The ``ibm_db_dbi`` module. - - Raises: - MissingDependencyError: When ibm_db is not installed. - """ - global _IBM_DB_DBI - if _IBM_DB_DBI is None: - _IBM_DB_DBI = import_optional("ibm_db_dbi") - if _IBM_DB_DBI is None: - raise MissingDependencyError(package="ibm_db", install_package="db2") - return _IBM_DB_DBI diff --git a/sqlspec/adapters/mssql_python/_typing.py b/sqlspec/adapters/mssql_python/_typing.py index d5fc593ba..1102b0522 100644 --- a/sqlspec/adapters/mssql_python/_typing.py +++ b/sqlspec/adapters/mssql_python/_typing.py @@ -23,8 +23,7 @@ MssqlPythonConnection: TypeAlias = Connection MssqlPythonRawCursor: TypeAlias = Cursor - -if not TYPE_CHECKING: +else: MssqlPythonConnection = Connection MssqlPythonRawCursor = Cursor diff --git a/sqlspec/adapters/mysqlconnector/_typing.py b/sqlspec/adapters/mysqlconnector/_typing.py index 558a92349..12b9fc274 100644 --- a/sqlspec/adapters/mysqlconnector/_typing.py +++ b/sqlspec/adapters/mysqlconnector/_typing.py @@ -5,21 +5,25 @@ """ import contextlib +import importlib from typing import TYPE_CHECKING, Any -from sqlspec.typing import import_optional, import_optional_attr +import mysql +import mysql.connector as _mysql_connector +from mysql.connector import MySQLConnection as _MysqlConnectorSyncConnection +from mysql.connector import aio as _mysql_connector_aio +from mysql.connector import constants as _mysql_connector_constants +from mysql.connector import cursor as _mysql_connector_cursor +from mysql.connector import pooling as _mysql_connector_pooling +from mysql.connector.aio import MySQLConnection as _MysqlConnectorAsyncConnection +from mysql.connector.aio import cursor as _mysql_connector_aio_cursor +from mysql.connector.aio.pooling import MySQLConnectionPool as MysqlConnectorAsyncPool if TYPE_CHECKING: from collections.abc import Awaitable, Callable from types import TracebackType from typing import ClassVar, Protocol, TypeAlias - import mysql.connector as _mysql_connector - from mysql.connector import MySQLConnection as _MysqlConnectorSyncConnection - from mysql.connector import pooling as _mysql_connector_pooling - from mysql.connector.aio.pooling import MySQLConnectionPool as MysqlConnectorAsyncPool - from mysql.connector.cursor import MySQLCursor as _MysqlConnectorSyncRawCursor - from sqlspec.adapters.mysqlconnector.driver import MysqlConnectorAsyncDriver, MysqlConnectorSyncDriver from sqlspec.core import StatementConfig @@ -72,24 +76,22 @@ class MysqlConnectorMysqlModuleProtocol(Protocol): MysqlConnectorSyncConnection: TypeAlias = _MysqlConnectorSyncConnection MysqlConnectorAio: TypeAlias = MysqlConnectorAioModuleProtocol MysqlConnectorAsyncConnection: TypeAlias = MysqlConnectorAsyncConnectionProtocol - MysqlConnectorSyncRawCursor: TypeAlias = _MysqlConnectorSyncRawCursor + MysqlConnectorSyncRawCursor: TypeAlias = _mysql_connector_cursor.MySQLCursor MysqlConnectorConnectionPool: TypeAlias = _mysql_connector_pooling.MySQLConnectionPool MysqlConnectorError: TypeAlias = _mysql_connector.Error MysqlConnectorFieldType: TypeAlias = MysqlConnectorFieldTypeProtocol MysqlConnectorMysqlModule: TypeAlias = MysqlConnectorMysqlModuleProtocol MysqlConnectorAsyncRawCursor: TypeAlias = MysqlConnectorAsyncRawCursorProtocol - -if not TYPE_CHECKING: - MysqlConnectorAsyncPool = import_optional_attr("mysql.connector.aio.pooling", "MySQLConnectionPool") - MysqlConnectorAio = import_optional_attr("mysql.connector", "aio") - MysqlConnectorSyncConnection = import_optional_attr("mysql.connector", "MySQLConnection") - MysqlConnectorAsyncConnection = import_optional_attr("mysql.connector.aio", "MySQLConnection") - MysqlConnectorConnectionPool = import_optional_attr("mysql.connector.pooling", "MySQLConnectionPool") - MysqlConnectorError = import_optional_attr("mysql.connector", "Error") - MysqlConnectorFieldType = import_optional_attr("mysql.connector.constants", "FieldType") - MysqlConnectorMysqlModule = import_optional("mysql") - MysqlConnectorSyncRawCursor = import_optional_attr("mysql.connector.cursor", "MySQLCursor") - MysqlConnectorAsyncRawCursor = import_optional_attr("mysql.connector.aio.cursor", "MySQLCursor") +else: + MysqlConnectorAio = _mysql_connector_aio + MysqlConnectorSyncConnection = _MysqlConnectorSyncConnection + MysqlConnectorAsyncConnection = _MysqlConnectorAsyncConnection + MysqlConnectorConnectionPool = _mysql_connector_pooling.MySQLConnectionPool + MysqlConnectorError = _mysql_connector.Error + MysqlConnectorFieldType = _mysql_connector_constants.FieldType + MysqlConnectorMysqlModule = mysql + MysqlConnectorSyncRawCursor = _mysql_connector_cursor.MySQLCursor + MysqlConnectorAsyncRawCursor = _mysql_connector_aio_cursor.MySQLCursor __all__ = ( "MysqlConnectorAio", @@ -267,7 +269,8 @@ def _get_mysqlconnector_sync_driver_type() -> "type[MysqlConnectorSyncDriver]": global _resolved_sync_driver_type driver_cls = _resolved_sync_driver_type if driver_cls is None: - driver_cls = import_optional_attr("sqlspec.adapters.mysqlconnector.driver", "MysqlConnectorSyncDriver") + module = importlib.import_module("sqlspec.adapters.mysqlconnector.driver") + driver_cls = module.MysqlConnectorSyncDriver _resolved_sync_driver_type = driver_cls return driver_cls @@ -277,6 +280,7 @@ def _get_mysqlconnector_async_driver_type() -> "type[MysqlConnectorAsyncDriver]" global _resolved_async_driver_type driver_cls = _resolved_async_driver_type if driver_cls is None: - driver_cls = import_optional_attr("sqlspec.adapters.mysqlconnector.driver", "MysqlConnectorAsyncDriver") + module = importlib.import_module("sqlspec.adapters.mysqlconnector.driver") + driver_cls = module.MysqlConnectorAsyncDriver _resolved_async_driver_type = driver_cls return driver_cls diff --git a/sqlspec/adapters/oracledb/_vector_handlers.py b/sqlspec/adapters/oracledb/_vector_handlers.py index 63f72fef3..e3fa3793a 100644 --- a/sqlspec/adapters/oracledb/_vector_handlers.py +++ b/sqlspec/adapters/oracledb/_vector_handlers.py @@ -26,7 +26,7 @@ from sqlspec.adapters.oracledb._typing import OracleSyncRawCursor as Cursor SPARSE_VECTOR_TYPE = OracleSparseVector -np: Any = import_optional("numpy") + __all__ = ( "DTYPE_TO_ARRAY_CODE", @@ -72,6 +72,11 @@ DTYPE_TO_ARRAY_CODE["float16"] = _TYPECODE_FLOAT16 +def _numpy_module() -> Any: + """Return the ``numpy`` module when installed, otherwise ``None``.""" + return import_optional("numpy") + + def numpy_converter_in(value: Any) -> "array.array[Any]": """Convert NumPy array to Oracle array for VECTOR insertion. @@ -112,6 +117,7 @@ def numpy_converter_out(value: "array.array[Any]") -> Any: if not NUMPY_INSTALLED: return value + np = _numpy_module() return np.array(value, copy=True, dtype=value.typecode) @@ -161,7 +167,7 @@ def _is_vector_payload(value: Any) -> bool: return False if isinstance(value, array.array): return True - if NUMPY_INSTALLED and isinstance(value, np.ndarray): + if NUMPY_INSTALLED and isinstance(value, _numpy_module().ndarray): return True if isinstance(value, (list, tuple)) and value: return all(isinstance(v, (int, float)) and not isinstance(v, bool) for v in value) @@ -202,7 +208,7 @@ def _input_type_handler(cursor: "Cursor | AsyncCursor", value: Any, arraysize: i if server_major is not None and server_major < _ORACLE_VECTOR_MIN_MAJOR: return None - if NUMPY_INSTALLED and isinstance(value, np.ndarray): + if NUMPY_INSTALLED and isinstance(value, _numpy_module().ndarray): return cursor.var(DB_TYPE_VECTOR, arraysize=arraysize, inconverter=numpy_converter_in) if isinstance(value, array.array): diff --git a/sqlspec/adapters/oracledb/driver.py b/sqlspec/adapters/oracledb/driver.py index d0dff6c36..a64c9860c 100644 --- a/sqlspec/adapters/oracledb/driver.py +++ b/sqlspec/adapters/oracledb/driver.py @@ -86,9 +86,10 @@ ) from sqlspec.driver._common import CommonDriverAttributesMixin from sqlspec.exceptions import ImproperConfigurationError, SQLSpecError, StackExecutionError +from sqlspec.storage._utils import import_pyarrow from sqlspec.utils.arrow_helpers import coerce_oracle_df_to_arrow_table from sqlspec.utils.logging import get_logger, log_with_context -from sqlspec.utils.module_loader import ensure_pyarrow, import_optional +from sqlspec.utils.module_loader import ensure_pyarrow from sqlspec.utils.text import normalize_identifier, quote_identifier, split_qualified_identifier from sqlspec.utils.type_guards import has_pipeline_capability, is_async_readable from sqlspec.utils.uuids import uuid4 @@ -115,7 +116,7 @@ logger = get_logger(__name__) -pa: Any = import_optional("pyarrow") + PLAN_STATEMENT_ID_MAX_LENGTH: Final[int] = 30 """Width of the PLAN_TABLE STATEMENT_ID column.""" @@ -819,9 +820,11 @@ def select_to_arrow( try: first_batch = next(batch_generator) except StopIteration: + pa = import_pyarrow() empty_schema = arrow_schema if arrow_schema is not None else pa.schema([]) reader = pa.RecordBatchReader.from_batches(empty_schema, ()) return create_arrow_result(prepared_statement, reader, rows_affected=0) + pa = import_pyarrow() schema = arrow_schema if arrow_schema is not None else first_batch.schema lazy_reader = pa.RecordBatchReader.from_batches(schema, chain((first_batch,), batch_generator)) return create_arrow_result(prepared_statement, lazy_reader, rows_affected=-1) @@ -859,7 +862,7 @@ def select_to_arrow( ) self._check_pending_exception(exc_handler) - arrow_table = coerce_oracle_df_to_arrow_table(oracle_df, pa) + arrow_table = coerce_oracle_df_to_arrow_table(oracle_df, import_pyarrow()) column_names = normalize_column_names(arrow_table.column_names, self.driver_features) if column_names != arrow_table.column_names: arrow_table = arrow_table.rename_columns(column_names) @@ -1242,7 +1245,7 @@ def _iter_arrow_record_batches( statement=sql, parameters=params, size=batch_size or 1000, **retry_kwargs ) for oracle_df in oracle_dfs: - batch_table = coerce_oracle_df_to_arrow_table(oracle_df, pa) + batch_table = coerce_oracle_df_to_arrow_table(oracle_df, import_pyarrow()) column_names = normalize_column_names(batch_table.column_names, self.driver_features) if column_names != batch_table.column_names: batch_table = batch_table.rename_columns(column_names) @@ -1273,7 +1276,7 @@ def _fetch_arrow_record_batches( statement=sql, parameters=params, size=batch_size or 1000, **retry_kwargs ) for oracle_df in oracle_dfs: - batch_table = coerce_oracle_df_to_arrow_table(oracle_df, pa) + batch_table = coerce_oracle_df_to_arrow_table(oracle_df, import_pyarrow()) column_names = normalize_column_names(batch_table.column_names, self.driver_features) if column_names != batch_table.column_names: batch_table = batch_table.rename_columns(column_names) @@ -1749,6 +1752,7 @@ async def select_to_arrow( ) record_batches, batch_schema = batches_result rows_affected = sum(batch.num_rows for batch in record_batches) + pa = import_pyarrow() if batch_schema is None: batch_schema = pa.schema([]) arrow_data: object @@ -1775,7 +1779,7 @@ async def select_to_arrow( ) self._check_pending_exception(exc_handler) - arrow_table = coerce_oracle_df_to_arrow_table(oracle_df, pa) + arrow_table = coerce_oracle_df_to_arrow_table(oracle_df, import_pyarrow()) column_names = normalize_column_names(arrow_table.column_names, self.driver_features) if column_names != arrow_table.column_names: arrow_table = arrow_table.rename_columns(column_names) @@ -2173,7 +2177,7 @@ async def _fetch_arrow_record_batches( statement=sql, parameters=params, size=batch_size or 1000, **retry_kwargs ) async for oracle_df in oracle_dfs: - batch_table = coerce_oracle_df_to_arrow_table(oracle_df, pa) + batch_table = coerce_oracle_df_to_arrow_table(oracle_df, import_pyarrow()) column_names = normalize_column_names(batch_table.column_names, self.driver_features) if column_names != batch_table.column_names: batch_table = batch_table.rename_columns(column_names) @@ -2317,6 +2321,7 @@ def _arrow_schema_supports_direct_path(arrow_table: Any) -> bool: True when every column type is known to be convertible. """ ensure_pyarrow() + pa = import_pyarrow() for field in arrow_table.schema: data_type = field.type diff --git a/sqlspec/adapters/oracledb/migrations.py b/sqlspec/adapters/oracledb/migrations.py index 7e24f2bff..37dce84fb 100644 --- a/sqlspec/adapters/oracledb/migrations.py +++ b/sqlspec/adapters/oracledb/migrations.py @@ -5,9 +5,8 @@ """ import getpass -from typing import TYPE_CHECKING, cast - -from rich.console import Console +import importlib +from typing import TYPE_CHECKING, Any, cast from sqlspec.builder import CreateTable, Select, sql from sqlspec.migrations.tracker import AsyncMigrationTracker, SyncMigrationTracker @@ -24,7 +23,14 @@ __all__ = ("OracleAsyncMigrationTracker", "OracleSyncMigrationTracker") logger = get_logger("sqlspec.migrations.oracle") -console = Console() +_console: Any = None + + +def _get_console() -> Any: + global _console + if _console is None: + _console = importlib.import_module("rich.console").Console() + return _console class OracleMigrationTrackerMixin: @@ -367,7 +373,7 @@ def _migrate_schema_if_needed(self, driver: "SyncDriverAdapterBase") -> None: return if self._should_echo(): - console.print( + _get_console().print( f"[cyan]Migrating tracking table schema, adding columns: {', '.join(sorted(missing_columns))}[/]" ) @@ -376,7 +382,7 @@ def _migrate_schema_if_needed(self, driver: "SyncDriverAdapterBase") -> None: driver.commit() if self._should_echo(): - console.print("[green]Migration tracking table schema updated successfully[/]") + _get_console().print("[green]Migration tracking table schema updated successfully[/]") except Exception as e: logger.warning("Could not check or migrate tracking table schema: %s", e) @@ -593,7 +599,7 @@ async def _migrate_schema_if_needed(self, driver: "AsyncDriverAdapterBase") -> N return if self._should_echo(): - console.print( + _get_console().print( f"[cyan]Migrating tracking table schema, adding columns: {', '.join(sorted(missing_columns))}[/]" ) @@ -602,7 +608,7 @@ async def _migrate_schema_if_needed(self, driver: "AsyncDriverAdapterBase") -> N await driver.commit() if self._should_echo(): - console.print("[green]Migration tracking table schema updated successfully[/]") + _get_console().print("[green]Migration tracking table schema updated successfully[/]") except Exception as e: logger.warning("Could not check or migrate tracking table schema: %s", e) diff --git a/sqlspec/adapters/oracledb/type_converter.py b/sqlspec/adapters/oracledb/type_converter.py index 8622e5881..25e3b773c 100644 --- a/sqlspec/adapters/oracledb/type_converter.py +++ b/sqlspec/adapters/oracledb/type_converter.py @@ -13,11 +13,14 @@ from sqlspec.utils.sync_tools import ensure_async_ from sqlspec.utils.type_guards import is_async_readable, is_readable -np: Any = import_optional("numpy") - __all__ = ("OracleOutputConverter",) +def _numpy_module() -> Any: + """Return the ``numpy`` module when installed, otherwise ``None``.""" + return import_optional("numpy") + + class OracleOutputConverter: """Oracle-specific output conversion with LOB optimization. @@ -77,7 +80,11 @@ def convert_numpy_to_vector(self, value: Any) -> Any: array.array compatible with Oracle VECTOR if value is ndarray, otherwise original value. """ - if not NUMPY_INSTALLED or np is None: + if not NUMPY_INSTALLED: + return value + + np = _numpy_module() + if np is None: return value if isinstance(value, np.ndarray): diff --git a/sqlspec/adapters/psqlpy/_typing.py b/sqlspec/adapters/psqlpy/_typing.py index 6ba1186b4..914679c8f 100644 --- a/sqlspec/adapters/psqlpy/_typing.py +++ b/sqlspec/adapters/psqlpy/_typing.py @@ -8,62 +8,27 @@ import importlib from typing import TYPE_CHECKING, Any -from sqlspec.typing import import_optional_attr - - -class _PsqlpyUnavailableError(Exception): - """Fallback psqlpy exception base when optional exception classes are unavailable.""" - +from psqlpy import Connection as PsqlpyConnection +from psqlpy import ConnectionPool as PsqlpyConnectionPool +from psqlpy import Listener as PsqlpyListener +from psqlpy.exceptions import ConnectionExecuteError as PsqlpyConnectionExecuteError +from psqlpy.exceptions import DatabaseError as PsqlpyDatabaseError +from psqlpy.exceptions import DataError as PsqlpyDataError +from psqlpy.exceptions import Error as PsqlpyError +from psqlpy.exceptions import IntegrityError as PsqlpyIntegrityError +from psqlpy.exceptions import NotSupportedError as PsqlpyNotSupportedError +from psqlpy.exceptions import OperationalError as PsqlpyOperationalError +from psqlpy.extra_types import JSONB as PSQLPY_JSONB + +RustPSQLDriverPyBaseError = PsqlpyError if TYPE_CHECKING: from collections.abc import Callable from types import TracebackType - from typing import TypeAlias - - from psqlpy import Connection as _PsqlpyConnection - from psqlpy import ConnectionPool as PsqlpyConnectionPool - from psqlpy import Listener as _PsqlpyListener - from psqlpy.exceptions import ConnectionExecuteError as _PsqlpyConnectionExecuteError - from psqlpy.exceptions import DatabaseError as _PsqlpyDatabaseError - from psqlpy.exceptions import DataError as _PsqlpyDataError - from psqlpy.exceptions import Error as _PsqlpyError - from psqlpy.exceptions import IntegrityError as _PsqlpyIntegrityError - from psqlpy.exceptions import NotSupportedError as _PsqlpyNotSupportedError - from psqlpy.exceptions import OperationalError as _PsqlpyOperationalError - from psqlpy.extra_types import JSONB as PSQLPY_JSONB from sqlspec.adapters.psqlpy.driver import PsqlpyDriver from sqlspec.core import StatementConfig - PsqlpyConnection: TypeAlias = _PsqlpyConnection - PsqlpyDataError: TypeAlias = _PsqlpyDataError - PsqlpyDatabaseError: TypeAlias = _PsqlpyDatabaseError - PsqlpyConnectionExecuteError: TypeAlias = _PsqlpyConnectionExecuteError - PsqlpyError: TypeAlias = _PsqlpyError - PsqlpyIntegrityError: TypeAlias = _PsqlpyIntegrityError - PsqlpyListener: TypeAlias = _PsqlpyListener - PsqlpyNotSupportedError: TypeAlias = _PsqlpyNotSupportedError - PsqlpyOperationalError: TypeAlias = _PsqlpyOperationalError - RustPSQLDriverPyBaseError: TypeAlias = Any - - -else: - PSQLPY_JSONB = import_optional_attr("psqlpy.extra_types", "JSONB") - PsqlpyConnection = import_optional_attr("psqlpy", "Connection") or Any - PsqlpyConnectionPool = import_optional_attr("psqlpy", "ConnectionPool") or Any - PsqlpyDataError = import_optional_attr("psqlpy.exceptions", "DataError") or _PsqlpyUnavailableError - PsqlpyDatabaseError = import_optional_attr("psqlpy.exceptions", "DatabaseError") or _PsqlpyUnavailableError - PsqlpyConnectionExecuteError = ( - import_optional_attr("psqlpy.exceptions", "ConnectionExecuteError") or _PsqlpyUnavailableError - ) - PsqlpyError = import_optional_attr("psqlpy.exceptions", "Error") or _PsqlpyUnavailableError - PsqlpyIntegrityError = import_optional_attr("psqlpy.exceptions", "IntegrityError") or _PsqlpyUnavailableError - PsqlpyListener = import_optional_attr("psqlpy", "Listener") or Any - PsqlpyNotSupportedError = import_optional_attr("psqlpy.exceptions", "NotSupportedError") or _PsqlpyUnavailableError - PsqlpyOperationalError = import_optional_attr("psqlpy.exceptions", "OperationalError") or _PsqlpyUnavailableError - RustPSQLDriverPyBaseError = import_optional_attr("psqlpy.exceptions", "RustPSQLDriverPyBaseError") or PsqlpyError - - __all__ = ( "PSQLPY_JSONB", "PsqlpyConnection", @@ -179,4 +144,5 @@ def _resolve_psqlpy_driver() -> "type[PsqlpyDriver]": if _CACHED_PSQLPY_DRIVER_CLS is None: module = importlib.import_module("sqlspec.adapters.psqlpy.driver") _CACHED_PSQLPY_DRIVER_CLS = module.PsqlpyDriver + assert _CACHED_PSQLPY_DRIVER_CLS is not None return _CACHED_PSQLPY_DRIVER_CLS diff --git a/sqlspec/adapters/psqlpy/type_converter.py b/sqlspec/adapters/psqlpy/type_converter.py index 25ac54664..1aa253efc 100644 --- a/sqlspec/adapters/psqlpy/type_converter.py +++ b/sqlspec/adapters/psqlpy/type_converter.py @@ -2,20 +2,20 @@ from typing import Any -from sqlspec.typing import import_optional_attr +from psqlpy.extra_types import PgVector __all__ = ("coerce_pgvector",) -_PGVECTOR_TYPE = import_optional_attr("psqlpy.extra_types", "PgVector") - def coerce_pgvector(value: Any) -> Any: """Wrap supported dense vector values using psqlpy's native vector encoder.""" - if value is None or _PGVECTOR_TYPE is None or isinstance(value, _PGVECTOR_TYPE): + if value is None or isinstance(value, PgVector): return value if isinstance(value, (list, tuple)): - return _PGVECTOR_TYPE(list(value)) + return PgVector(list(value)) tolist = getattr(value, "tolist", None) if callable(tolist): - return _PGVECTOR_TYPE(tolist()) + raw_list = tolist() + if isinstance(raw_list, list): + return PgVector(raw_list) return value diff --git a/sqlspec/adapters/psycopg/type_converter.py b/sqlspec/adapters/psycopg/type_converter.py index 5870ef344..8d0a8018b 100644 --- a/sqlspec/adapters/psycopg/type_converter.py +++ b/sqlspec/adapters/psycopg/type_converter.py @@ -23,7 +23,10 @@ logger = get_logger(__name__) -_pgvector_psycopg: Any | None = import_optional("pgvector.psycopg") + +def _pgvector_psycopg_module() -> Any: + """Return the ``pgvector.psycopg`` module when installed, otherwise ``None``.""" + return import_optional("pgvector.psycopg") def register_pgvector_sync(connection: "Connection[Any]") -> None: @@ -35,7 +38,7 @@ def register_pgvector_sync(connection: "Connection[Any]") -> None: Args: connection: Psycopg sync connection. """ - pgvector_psycopg = _pgvector_psycopg + pgvector_psycopg = _pgvector_psycopg_module() if pgvector_psycopg is None: return @@ -58,7 +61,7 @@ async def register_pgvector_async(connection: "AsyncConnection[Any]") -> None: Args: connection: Psycopg async connection. """ - pgvector_psycopg = _pgvector_psycopg + pgvector_psycopg = _pgvector_psycopg_module() if pgvector_psycopg is None: return diff --git a/sqlspec/adapters/pymssql/_typing.py b/sqlspec/adapters/pymssql/_typing.py index 986699e4b..ebc8084f7 100644 --- a/sqlspec/adapters/pymssql/_typing.py +++ b/sqlspec/adapters/pymssql/_typing.py @@ -5,41 +5,29 @@ """ import contextlib +import importlib from typing import TYPE_CHECKING, Any -from sqlspec.utils.module_loader import import_optional, import_optional_attr - - -class _MissingPymssqlDriverError(Exception): - """Placeholder base used only when pymssql is not installed.""" - - -_NATIVE_ERROR: Any = import_optional_attr("pymssql", "Error") or _MissingPymssqlDriverError -PymssqlError: Any = _NATIVE_ERROR -PYMSSQL_MODULE: Any = import_optional("pymssql") -_PymssqlConnection: Any = getattr(PYMSSQL_MODULE, "Connection", Any) if PYMSSQL_MODULE is not None else Any -_PymssqlRawCursor: Any = getattr(PYMSSQL_MODULE, "Cursor", Any) if PYMSSQL_MODULE is not None else Any -_DRIVER_CLS: Any = None +import pymssql as PYMSSQL_MODULE # noqa: N812 +from pymssql import Connection, Cursor +from pymssql import Error as PymssqlError if TYPE_CHECKING: from collections.abc import Callable from types import TracebackType from typing import TypeAlias - from pymssql import Connection as _TypeCheckingConnection - from pymssql import Cursor as _TypeCheckingCursor from pymssql._pymssql import QueryParams as PymssqlQueryParams from sqlspec.adapters.pymssql.driver import PymssqlDriver from sqlspec.core import StatementConfig - PymssqlConnection: TypeAlias = _TypeCheckingConnection - PymssqlRawCursor: TypeAlias = _TypeCheckingCursor - -if not TYPE_CHECKING: + PymssqlConnection: TypeAlias = Connection + PymssqlRawCursor: TypeAlias = Cursor +else: PymssqlQueryParams = Any - PymssqlConnection = _PymssqlConnection - PymssqlRawCursor = _PymssqlRawCursor + PymssqlConnection = Connection + PymssqlRawCursor = Cursor __all__ = ( "PYMSSQL_MODULE", @@ -82,6 +70,9 @@ def __exit__(self, *_: Any) -> None: self.cursor.close() +_DRIVER_CLS: "type[PymssqlDriver] | None" = None + + class PymssqlSessionContext: """Sync context manager for pymssql sessions.""" @@ -118,8 +109,10 @@ def __enter__(self) -> "PymssqlDriver": global _DRIVER_CLS driver_cls = self._driver_cls or _DRIVER_CLS if driver_cls is None: - _DRIVER_CLS = import_optional_attr("sqlspec.adapters.pymssql.driver", "PymssqlDriver") + module = importlib.import_module("sqlspec.adapters.pymssql.driver") + _DRIVER_CLS = module.PymssqlDriver driver_cls = _DRIVER_CLS + assert driver_cls is not None self._connection = self._acquire_connection() self._driver = driver_cls( connection=self._connection, statement_config=self._statement_config, driver_features=self._driver_features diff --git a/sqlspec/adapters/pymysql/_typing.py b/sqlspec/adapters/pymysql/_typing.py index c7571d2b7..5458aae6e 100644 --- a/sqlspec/adapters/pymysql/_typing.py +++ b/sqlspec/adapters/pymysql/_typing.py @@ -38,8 +38,7 @@ class PyMysqlServerStatusProtocol(Protocol): PyMysqlConnect: TypeAlias = type["PyMysqlConnection"] PyMysqlFieldType: TypeAlias = PyMysqlFieldTypeProtocol PyMysqlServerStatus: TypeAlias = PyMysqlServerStatusProtocol - -if not TYPE_CHECKING: +else: PyMysqlConnect = pymysql.connect PyMysqlFieldType = _PYMYSQL_FIELD_TYPE PyMysqlServerStatus = _PYMYSQL_SERVER_STATUS diff --git a/sqlspec/adapters/sqlanywhere/_typing.py b/sqlspec/adapters/sqlanywhere/_typing.py index 297e9fa30..a58b8d854 100644 --- a/sqlspec/adapters/sqlanywhere/_typing.py +++ b/sqlspec/adapters/sqlanywhere/_typing.py @@ -5,23 +5,13 @@ """ import contextlib +import importlib from typing import TYPE_CHECKING, Any +import sqlanydb as SQLANYDB_MODULE # noqa: N812 from mypy_extensions import mypyc_attr - -from sqlspec.utils.module_loader import import_optional_attr - -SQLANYDB_MODULE: Any = None -_DRIVER_CLS: Any = None - - -class _MissingSqlanywhereDriverError(Exception): - """Placeholder base used only when sqlanydb is not installed.""" - - -_NATIVE_ERROR: Any = import_optional_attr("sqlanydb", "Error") or _MissingSqlanywhereDriverError -SqlanywhereError: Any = _NATIVE_ERROR - +from sqlanydb import Connection, Cursor +from sqlanydb import Error as SqlanywhereError if TYPE_CHECKING: from collections.abc import Callable @@ -31,13 +21,12 @@ class _MissingSqlanywhereDriverError(Exception): from sqlspec.adapters.sqlanywhere.driver import SqlanywhereDriver from sqlspec.core import StatementConfig - SqlanywhereConnection: TypeAlias = Any - SqlanywhereRawCursor: TypeAlias = Any + SqlanywhereConnection: TypeAlias = Connection + SqlanywhereRawCursor: TypeAlias = Cursor SqlanywhereQueryParams: TypeAlias = Any - -if not TYPE_CHECKING: - SqlanywhereConnection = Any - SqlanywhereRawCursor = Any +else: + SqlanywhereConnection = Connection + SqlanywhereRawCursor = Cursor SqlanywhereQueryParams = Any __all__ = ( @@ -71,6 +60,9 @@ def __exit__(self, *_: Any) -> None: self.cursor.close() +_DRIVER_CLS: "type[SqlanywhereDriver] | None" = None + + @mypyc_attr(allow_interpreted_subclasses=True, native_class=False) class SqlanywhereSessionContext: """Sync context manager for SAP SQL Anywhere sessions.""" @@ -104,7 +96,9 @@ def __init__( def __enter__(self) -> "SqlanywhereDriver": global _DRIVER_CLS if _DRIVER_CLS is None: - _DRIVER_CLS = import_optional_attr("sqlspec.adapters.sqlanywhere.driver", "SqlanywhereDriver") + module = importlib.import_module("sqlspec.adapters.sqlanywhere.driver") + _DRIVER_CLS = module.SqlanywhereDriver + assert _DRIVER_CLS is not None self._connection = self._acquire_connection() self._driver = _DRIVER_CLS( connection=self._connection, statement_config=self._statement_config, driver_features=self._driver_features diff --git a/sqlspec/adapters/sqlanywhere/pool.py b/sqlspec/adapters/sqlanywhere/pool.py index 42b075bd0..1f740f674 100644 --- a/sqlspec/adapters/sqlanywhere/pool.py +++ b/sqlspec/adapters/sqlanywhere/pool.py @@ -5,11 +5,10 @@ import threading import time from contextlib import contextmanager -from importlib import import_module from typing import TYPE_CHECKING, Any, cast -from sqlspec.adapters.sqlanywhere._typing import SqlanywhereConnection -from sqlspec.exceptions import MissingDependencyError, OperationalError +from sqlspec.adapters.sqlanywhere._typing import SQLANYDB_MODULE, SqlanywhereConnection +from sqlspec.exceptions import OperationalError from sqlspec.utils.logging import POOL_LOGGER_NAME, get_logger, log_with_context from sqlspec.utils.uuids import uuid4 @@ -73,11 +72,10 @@ def _database_name(self) -> str: def new_connection(self) -> SqlanywhereConnection: """Open a standalone SAP SQL Anywhere connection via ``sqlanydb.connect``.""" - sqlanydb = _load_sqlanydb_module() - connection = sqlanydb.connect(**self._connection_parameters) + connection = SQLANYDB_MODULE.connect(**self._connection_parameters) if self._on_connection_create is not None: self._on_connection_create(connection) - return cast("SqlanywhereConnection", connection) + return connection @contextmanager def get_connection(self) -> "Generator[SqlanywhereConnection, None, None]": @@ -236,11 +234,3 @@ def _close_thread_connection(self) -> None: self._thread_local.created_at = 0.0 self._thread_local.last_used = 0.0 self._thread_local.generation = -1 - - -def _load_sqlanydb_module() -> Any: - """Import ``sqlanydb`` lazily or raise ``MissingDependencyError``.""" - try: - return import_module("sqlanydb") - except ImportError as exc: - raise MissingDependencyError(package="sqlanydb", install_package="sqlanywhere") from exc diff --git a/sqlspec/adapters/sybase/_typing.py b/sqlspec/adapters/sybase/_typing.py index ddee00e6a..fd5a2c1ab 100644 --- a/sqlspec/adapters/sybase/_typing.py +++ b/sqlspec/adapters/sybase/_typing.py @@ -5,43 +5,30 @@ """ import contextlib +import importlib from typing import TYPE_CHECKING, Any +import pymssql as PYMSSQL_MODULE # noqa: N812 from mypy_extensions import mypyc_attr - -from sqlspec.utils.module_loader import import_optional_attr - -PYMSSQL_MODULE: Any = None -_DRIVER_CLS: Any = None - - -class _MissingSybaseDriverError(Exception): - """Placeholder base used only when pymssql is not installed.""" - - -_NATIVE_ERROR: Any = import_optional_attr("pymssql", "Error") or _MissingSybaseDriverError -SybaseError: Any = _NATIVE_ERROR - +from pymssql import Connection, Cursor +from pymssql import Error as SybaseError if TYPE_CHECKING: from collections.abc import Callable from types import TracebackType from typing import TypeAlias - from pymssql import Connection as _SybaseConnection - from pymssql import Cursor as _SybaseRawCursor from pymssql._pymssql import QueryParams as SybaseQueryParams from sqlspec.adapters.sybase.driver import SybaseDriver from sqlspec.core import StatementConfig - SybaseConnection: TypeAlias = _SybaseConnection - SybaseRawCursor: TypeAlias = _SybaseRawCursor - -if not TYPE_CHECKING: + SybaseConnection: TypeAlias = Connection + SybaseRawCursor: TypeAlias = Cursor +else: SybaseQueryParams = Any - SybaseConnection = Any - SybaseRawCursor = Any + SybaseConnection = Connection + SybaseRawCursor = Cursor __all__ = ( "PYMSSQL_MODULE", @@ -84,6 +71,9 @@ def __exit__(self, *_: Any) -> None: self.cursor.close() +_DRIVER_CLS: "type[SybaseDriver] | None" = None + + @mypyc_attr(allow_interpreted_subclasses=True, native_class=False) class SybaseSessionContext: """Sync context manager for Sybase ASE sessions.""" @@ -117,7 +107,9 @@ def __init__( def __enter__(self) -> "SybaseDriver": global _DRIVER_CLS if _DRIVER_CLS is None: - _DRIVER_CLS = import_optional_attr("sqlspec.adapters.sybase.driver", "SybaseDriver") + module = importlib.import_module("sqlspec.adapters.sybase.driver") + _DRIVER_CLS = module.SybaseDriver + assert _DRIVER_CLS is not None self._connection = self._acquire_connection() self._driver = _DRIVER_CLS( connection=self._connection, statement_config=self._statement_config, driver_features=self._driver_features diff --git a/sqlspec/adapters/sybase/pool.py b/sqlspec/adapters/sybase/pool.py index 48eda0125..60abb17fd 100644 --- a/sqlspec/adapters/sybase/pool.py +++ b/sqlspec/adapters/sybase/pool.py @@ -6,12 +6,11 @@ import threading import time from contextlib import contextmanager -from importlib import import_module from typing import TYPE_CHECKING, Any, cast from sqlspec.adapters.sybase._typing import PYMSSQL_MODULE, SybaseConnection from sqlspec.adapters.sybase.core import DEFAULT_SYBASE_CONN_PROPERTIES, DEFAULT_SYBASE_TDS_VERSION -from sqlspec.exceptions import ImproperConfigurationError, MissingDependencyError +from sqlspec.exceptions import ImproperConfigurationError from sqlspec.utils.logging import POOL_LOGGER_NAME, get_logger, log_with_context from sqlspec.utils.uuids import uuid4 @@ -22,7 +21,6 @@ logger = get_logger(POOL_LOGGER_NAME) _ADAPTER_NAME = "sybase" -pymssql = PYMSSQL_MODULE class SybaseConnectionPool: @@ -98,10 +96,7 @@ def new_connection(self) -> SybaseConnection: Returns: SybaseConnection: A newly opened, fully configured connection. """ - global pymssql - if pymssql is None: - pymssql = _load_pymssql_module() - connection = pymssql.connect(**self._connection_parameters) + connection = PYMSSQL_MODULE.connect(**self._connection_parameters) if self._on_connection_create is not None: self._on_connection_create(connection) @@ -265,11 +260,3 @@ def _close_thread_connection(self) -> None: del self._thread_local.created_at if "last_used" in thread_state: del self._thread_local.last_used - - -def _load_pymssql_module() -> Any: - """Import ``pymssql`` lazily or raise ``MissingDependencyError``.""" - try: - return import_module("pymssql") - except ImportError as exc: - raise MissingDependencyError(package="pymssql", install_package="sybase") from exc diff --git a/tests/unit/adapters/test_db2/test_pool.py b/tests/unit/adapters/test_db2/test_pool.py index b79d86ac6..635e5c101 100644 --- a/tests/unit/adapters/test_db2/test_pool.py +++ b/tests/unit/adapters/test_db2/test_pool.py @@ -6,9 +6,7 @@ import pytest -import sqlspec.adapters.db2.pool as pool_module from sqlspec.adapters.db2.pool import Db2SyncConnectionPool -from sqlspec.exceptions import MissingDependencyError from tests.unit.adapters.test_db2._fakes import FakeDb2Connection, FakeDb2Cursor, FakeIbmDbDbiModule, FakeIbmDbModule FakeModules = tuple[FakeIbmDbModule, FakeIbmDbDbiModule] @@ -67,19 +65,6 @@ def test_pool_renders_dsn_and_passes_empty_connect_arguments(fake_ibm_db: FakeMo ) -def test_pool_missing_dependency_raises(monkeypatch: pytest.MonkeyPatch) -> None: - """MissingDependencyError is raised when ibm_db_dbi is not available.""" - monkeypatch.setattr(pool_module, "_IBM_DB_DBI", None) - monkeypatch.setattr(pool_module, "import_optional", lambda _name: None) - - pool = Db2SyncConnectionPool({"database": "TESTDB"}) - - with pytest.raises(MissingDependencyError) as exc_info: - pool.acquire() - - assert "ibm_db" in str(exc_info.value) - - def test_pool_recycles_failed_health_check(fake_ibm_db: FakeModules, monkeypatch: pytest.MonkeyPatch) -> None: """Failed health check should close and replace thread connection.""" _, fake_module = fake_ibm_db diff --git a/tests/unit/adapters/test_db2/test_typing.py b/tests/unit/adapters/test_db2/test_typing.py index 1024db939..c844b0972 100644 --- a/tests/unit/adapters/test_db2/test_typing.py +++ b/tests/unit/adapters/test_db2/test_typing.py @@ -1,23 +1,17 @@ """Tests for the Db2 optional-dependency facade and vendor type aliases.""" -import importlib import sys import ibm_db_dbi import pytest -import sqlspec.adapters -import sqlspec.adapters.db2.pool as pool_module import sqlspec.typing from sqlspec.adapters.db2._typing import Db2Error, Db2SyncSessionContext from sqlspec.adapters.db2.core import default_statement_config from sqlspec.adapters.db2.data_dictionary import Db2SyncDataDictionary from sqlspec.adapters.db2.driver import Db2SyncDriver -from sqlspec.adapters.db2.pool import Db2SyncConnectionPool -from sqlspec.core import DRIVER_PARAMETER_PROFILES -from sqlspec.exceptions import MissingDependencyError from sqlspec.utils.module_loader import reset_dependency_cache -from tests.unit.adapters.test_db2._fakes import FakeDb2Connection, FakeIbmDbDbiModule +from tests.unit.adapters.test_db2._fakes import FakeDb2Connection def test_ibm_db_installed_flag_is_lazy(monkeypatch: pytest.MonkeyPatch) -> None: @@ -62,27 +56,6 @@ def test_session_context_builds_sync_driver() -> None: assert released == [connection] -def test_pool_imports_ibm_db_dbi_on_first_connection(monkeypatch: pytest.MonkeyPatch) -> None: - """The pool resolves ``ibm_db_dbi`` through the module loader when it first connects.""" - fake_module = FakeIbmDbDbiModule() - requested: list[str] = [] - - def fake_import_optional(name: str) -> FakeIbmDbDbiModule: - requested.append(name) - return fake_module - - monkeypatch.setattr(pool_module, "_IBM_DB_DBI", None) - monkeypatch.setattr(pool_module, "import_optional", fake_import_optional) - pool = Db2SyncConnectionPool({"database": "TESTDB"}) - - pool.new_connection() - pool.new_connection() - - assert requested == ["ibm_db_dbi"] - assert len(fake_module.connect_calls) == 2 - assert all("DATABASE=TESTDB;" in call[0] for call in fake_module.connect_calls) - - def test_driver_data_dictionary_is_reused() -> None: """The driver builds one Db2 data dictionary and returns it on every access.""" driver = Db2SyncDriver(FakeDb2Connection()) @@ -91,21 +64,3 @@ def test_driver_data_dictionary_is_reused() -> None: assert isinstance(dictionary, Db2SyncDataDictionary) assert driver.data_dictionary is dictionary - - -def test_db2_adapter_imports_without_ibm_db(monkeypatch: pytest.MonkeyPatch) -> None: - """The adapter package imports without the driver installed and fails only when connecting.""" - monkeypatch.setitem(sys.modules, "ibm_db", None) - monkeypatch.setitem(sys.modules, "ibm_db_dbi", None) - for name in [module for module in sys.modules if module.split(".")[:3] == ["sqlspec", "adapters", "db2"]]: - monkeypatch.delitem(sys.modules, name) - monkeypatch.delattr(sqlspec.adapters, "db2", raising=False) - monkeypatch.delitem(DRIVER_PARAMETER_PROFILES, "db2") - reset_dependency_cache() - try: - db2_package = importlib.import_module("sqlspec.adapters.db2") - - with pytest.raises(MissingDependencyError): - db2_package.Db2SyncConfig(connection_config={"database": "d"}).create_connection() - finally: - reset_dependency_cache() diff --git a/tests/unit/adapters/test_psycopg/test_type_handlers.py b/tests/unit/adapters/test_psycopg/test_type_handlers.py index 30f243324..881b956fe 100644 --- a/tests/unit/adapters/test_psycopg/test_type_handlers.py +++ b/tests/unit/adapters/test_psycopg/test_type_handlers.py @@ -12,7 +12,7 @@ def test_register_pgvector_sync_with_pgvector_installed(monkeypatch: pytest.Monk mock_connection = MagicMock() mock_pgvector = MagicMock() - monkeypatch.setattr(type_converter, "_pgvector_psycopg", mock_pgvector) + monkeypatch.setattr(type_converter, "_pgvector_psycopg_module", lambda: mock_pgvector) register_pgvector_sync(mock_connection) @@ -23,7 +23,7 @@ def test_register_pgvector_sync_without_pgvector(monkeypatch: pytest.MonkeyPatch """Test register_pgvector_sync gracefully handles pgvector not installed.""" import sqlspec.adapters.psycopg.type_converter - monkeypatch.setattr(sqlspec.adapters.psycopg.type_converter, "_pgvector_psycopg", None) + monkeypatch.setattr(sqlspec.adapters.psycopg.type_converter, "_pgvector_psycopg_module", lambda: None) from sqlspec.adapters.psycopg.type_converter import register_pgvector_sync @@ -41,7 +41,7 @@ async def test_register_pgvector_async_with_pgvector_installed(monkeypatch: pyte mock_connection = AsyncMock() mock_pgvector = MagicMock() mock_pgvector.register_vector_async = AsyncMock() - monkeypatch.setattr(type_converter, "_pgvector_psycopg", mock_pgvector) + monkeypatch.setattr(type_converter, "_pgvector_psycopg_module", lambda: mock_pgvector) await register_pgvector_async(mock_connection) @@ -52,7 +52,7 @@ async def test_register_pgvector_async_without_pgvector(monkeypatch: pytest.Monk """Test register_pgvector_async gracefully handles pgvector not installed.""" import sqlspec.adapters.psycopg.type_converter - monkeypatch.setattr(sqlspec.adapters.psycopg.type_converter, "_pgvector_psycopg", None) + monkeypatch.setattr(sqlspec.adapters.psycopg.type_converter, "_pgvector_psycopg_module", lambda: None) from sqlspec.adapters.psycopg.type_converter import register_pgvector_async @@ -70,7 +70,7 @@ def test_register_pgvector_sync_handles_registration_failure(monkeypatch: pytest mock_connection = MagicMock() mock_pgvector = MagicMock() mock_pgvector.register_vector.side_effect = Exception("Registration failed") - monkeypatch.setattr(type_converter, "_pgvector_psycopg", mock_pgvector) + monkeypatch.setattr(type_converter, "_pgvector_psycopg_module", lambda: mock_pgvector) register_pgvector_sync(mock_connection) @@ -83,18 +83,11 @@ async def test_register_pgvector_async_handles_registration_failure(monkeypatch: mock_connection = AsyncMock() mock_pgvector = MagicMock() mock_pgvector.register_vector_async = AsyncMock(side_effect=Exception("Registration failed")) - monkeypatch.setattr(type_converter, "_pgvector_psycopg", mock_pgvector) + monkeypatch.setattr(type_converter, "_pgvector_psycopg_module", lambda: mock_pgvector) await register_pgvector_async(mock_connection) -def test_pgvector_module_reference_is_cached() -> None: - import sqlspec.adapters.psycopg.type_converter as type_converter - - assert not hasattr(type_converter, "importlib") - assert hasattr(type_converter, "_pgvector_psycopg") - - def test_operational_error_with_none_sqlstate_maps_to_database_connection_error() -> None: """OperationalError with sqlstate=None maps to DatabaseConnectionError.""" import psycopg diff --git a/tests/unit/adapters/test_sqlanywhere.py b/tests/unit/adapters/test_sqlanywhere.py index 3f75fd356..536c3543c 100644 --- a/tests/unit/adapters/test_sqlanywhere.py +++ b/tests/unit/adapters/test_sqlanywhere.py @@ -39,7 +39,6 @@ DataError, DeadlockError, ForeignKeyViolationError, - MissingDependencyError, NotFoundError, NotNullViolationError, OperationalError, @@ -176,7 +175,7 @@ def connect(self, **kwargs: Any) -> FakeConnection: captured_kwargs.append(kwargs) return fake_conn - monkeypatch.setattr(pool_module, "_load_sqlanydb_module", lambda: FakeSqlanydbModule()) + monkeypatch.setattr(pool_module, "SQLANYDB_MODULE", FakeSqlanydbModule()) cfg = build_connection_config({"user": "dba", "password": "sql", "database": "demo", "host": "localhost"}) pool = SqlanywhereConnectionPool(cfg, min_size=1, max_size=2) @@ -193,21 +192,6 @@ def connect(self, **kwargs: Any) -> FakeConnection: assert fake_conn.closed is True -def test_sqlanywhere_pool_raises_missing_dependency_when_sqlanydb_absent(monkeypatch: pytest.MonkeyPatch) -> None: - """SqlanywhereConnectionPool raises MissingDependencyError when sqlanydb is not installed.""" - import sqlspec.adapters.sqlanywhere.pool as pool_module - - def _raise_import_error() -> Any: - msg = "No module named 'sqlanydb'" - raise ImportError(msg) - - monkeypatch.setattr(pool_module, "import_module", lambda name: _raise_import_error()) - - pool = SqlanywhereConnectionPool({"uid": "dba", "pwd": "sql"}) - with pytest.raises(MissingDependencyError, match="sqlanydb"): - pool.acquire() - - @pytest.mark.parametrize( ("args", "expected_type"), [ diff --git a/tests/unit/adapters/test_sybase.py b/tests/unit/adapters/test_sybase.py index ee3edcb5d..c7e0eba43 100644 --- a/tests/unit/adapters/test_sybase.py +++ b/tests/unit/adapters/test_sybase.py @@ -167,7 +167,7 @@ class FakeModule: def connect(self, **kwargs: Any) -> FakeConnection: return fake_conn - monkeypatch.setattr(pool_module, "pymssql", FakeModule()) + monkeypatch.setattr(pool_module, "PYMSSQL_MODULE", FakeModule()) monkeypatch.delenv("TDSVER", raising=False) cfg_invalid = build_connection_config({"server": "ase.example.test", "user": "sa", "database": "master"}) @@ -642,33 +642,6 @@ def cursor(self) -> DummyCursor: assert has_user is True -def test_sybase_pool_new_connection_caches_pymssql_module(monkeypatch: pytest.MonkeyPatch) -> None: - """new_connection caches pymssql module reference at module level.""" - from sqlspec.adapters.sybase import pool as sybase_pool - - import_count = 0 - real_import = sybase_pool._load_pymssql_module - - def mock_load() -> Any: - nonlocal import_count - import_count += 1 - return real_import() - - monkeypatch.setattr(sybase_pool, "_load_pymssql_module", mock_load) - monkeypatch.setattr(sybase_pool, "pymssql", None) - monkeypatch.setenv("TDSVER", "5.0") - - pool = sybase_pool.SybaseConnectionPool({"server": "localhost", "database": "demo"}) - monkeypatch.setattr("pymssql.connect", lambda **_kw: FakeConnection(FakeCursor())) - - _ = pool.new_connection() - assert sybase_pool.pymssql is not None - assert import_count == 1 - - _ = pool.new_connection() - assert import_count == 1 - - def test_sybase_format_identifier_and_insert_statements() -> None: """format_identifier, build_insert_statement, and escape_literal_percent format queries correctly.""" assert format_identifier("dbo.my_table") == "[dbo].[my_table]" diff --git a/tests/unit/test_mypyc_sybase_surface.py b/tests/unit/test_mypyc_sybase_surface.py index 9e89357c8..37ce9573a 100644 --- a/tests/unit/test_mypyc_sybase_surface.py +++ b/tests/unit/test_mypyc_sybase_surface.py @@ -92,13 +92,13 @@ def test_mypyc_smoke_imports_include_sybase_and_sqlanywhere() -> None: assert smoke_by_name["sqlanywhere_pool"].require_compiled is True -def test_importing_sybase_and_sqlanywhere_adapters_does_not_eagerly_import_driver_packages() -> None: - """Importing sqlspec.adapters.sybase and sqlspec.adapters.sqlanywhere defers sqlanydb.""" +def test_importing_sybase_and_sqlanywhere_adapters_does_not_import_rich() -> None: + """Importing sqlspec.adapters.sybase and sqlspec.adapters.sqlanywhere does not load rich.""" script = ( "import sys\n" "import sqlspec.adapters.sybase\n" "import sqlspec.adapters.sqlanywhere\n" - "leaked = [name for name in ('sqlanydb', 'rich') if name in sys.modules]\n" + "leaked = [name for name in ('rich',) if name in sys.modules]\n" "assert not leaked, f'Eagerly imported: {leaked}'\n" ) subprocess.run([sys.executable, "-c", script], capture_output=True, text=True, check=True) diff --git a/uv.lock b/uv.lock index d5a2a3e77..24930f0a3 100644 --- a/uv.lock +++ b/uv.lock @@ -5945,6 +5945,7 @@ dev = [ { name = "slotscheck" }, { name = "sniffio" }, { name = "sqlalchemy", extra = ["asyncio"] }, + { name = "sqlanydb" }, { name = "types-cffi" }, { name = "types-colorama" }, { name = "types-protobuf" }, @@ -5978,6 +5979,7 @@ extras = [ { name = "pgvector" }, { name = "polars" }, { name = "pyarrow" }, + { name = "sqlanydb" }, ] lint = [ { name = "asyncpg-stubs" }, @@ -6137,6 +6139,7 @@ dev = [ { name = "slotscheck", specifier = ">=0.16.5" }, { name = "sniffio" }, { name = "sqlalchemy", extras = ["asyncio"] }, + { name = "sqlanydb", specifier = ">=1.0.14" }, { name = "types-cffi" }, { name = "types-colorama" }, { name = "types-protobuf" }, @@ -6170,6 +6173,7 @@ extras = [ { name = "pgvector" }, { name = "polars" }, { name = "pyarrow" }, + { name = "sqlanydb", specifier = ">=1.0.14" }, ] lint = [ { name = "asyncpg-stubs" }, From 971692c49d576b8f5b100385d1fb8b0a4ec6c38c Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 03:57:09 +0000 Subject: [PATCH 100/174] fix(adbc): handle autocommit in transactions, status codes, and statement caching Toggle autocommit off on transaction begin and restore previous autocommit state on commit or rollback. Preserve cause exception chains across pending exception checks, expand status code mappings for ADBC errors, and relocate DuckDB COPY parameter transformation to statement transformers to allow compiled query caching. --- sqlspec/adapters/adbc/core.py | 51 +++++++++-- sqlspec/adapters/adbc/driver.py | 34 ++++++-- tests/unit/adapters/test_adbc/test_core.py | 99 ++++++++++++++++++++++ 3 files changed, 168 insertions(+), 16 deletions(-) diff --git a/sqlspec/adapters/adbc/core.py b/sqlspec/adapters/adbc/core.py index 920addaf2..d06f5121a 100644 --- a/sqlspec/adapters/adbc/core.py +++ b/sqlspec/adapters/adbc/core.py @@ -32,6 +32,7 @@ ForeignKeyViolationError, ImproperConfigurationError, IntegrityError, + NotFoundError, NotNullViolationError, OperationalError, OperationCancelledError, @@ -51,6 +52,8 @@ from sqlspec.utils.type_guards import has_rowcount, has_sqlstate from sqlspec.utils.uuids import uuid4, uuid_from_string +AdbcStatusCode: Any = import_optional_attr("adbc_driver_manager", "AdbcStatusCode") + if TYPE_CHECKING: from collections.abc import Callable, Mapping @@ -834,6 +837,39 @@ def prepare_postgres_parameters( ) +_ADBC_STATUS_CODE_MAP: Final[dict[str, tuple[type[SQLSpecError], str]]] = { + "TIMEOUT": (QueryTimeoutError, "query timeout"), + "CANCELLED": (OperationCancelledError, "operation cancelled"), + "UNAUTHENTICATED": (PermissionDeniedError, "unauthenticated"), + "UNAUTHORIZED": (PermissionDeniedError, "unauthorized"), + "NOT_FOUND": (NotFoundError, "not found"), + "ALREADY_EXISTS": (IntegrityError, "already exists"), + "INVALID_ARGUMENT": (DataError, "invalid argument"), + "INVALID_DATA": (DataError, "invalid data"), + "INTEGRITY": (IntegrityError, "integrity violation"), + "IO": (DatabaseConnectionError, "I/O error"), + "INVALID_STATE": (OperationalError, "invalid state"), + "INTERNAL": (OperationalError, "internal error"), + "NOT_IMPLEMENTED": (SQLSpecError, "not implemented"), +} + + +def _map_adbc_status_exception(error: Any) -> "SQLSpecError | None": + status_code = getattr(error, "status_code", None) + if status_code is None: + return None + status_name: str | None = getattr(status_code, "name", None) + if status_name is None and isinstance(status_code, int) and AdbcStatusCode is not None: + with contextlib.suppress(ValueError): + status_name = AdbcStatusCode(status_code).name + if status_name is None: + status_name = str(status_code).upper() + if status_name in _ADBC_STATUS_CODE_MAP: + exc_cls, desc = _ADBC_STATUS_CODE_MAP[status_name] + return _create_adbc_error(error, exc_cls, desc) + return None + + def create_mapped_exception(error: Any, *, logger: Any | None = None) -> SQLSpecError: """Map ADBC errors to SQLSpec exceptions. @@ -853,13 +889,9 @@ def create_mapped_exception(error: Any, *, logger: Any | None = None) -> SQLSpec Returns: A SQLSpec exception that wraps the original error """ - del logger - status_code = getattr(error, "status_code", None) - status_name = getattr(status_code, "name", str(status_code)).upper() - if status_name == "TIMEOUT": - return _create_adbc_error(error, QueryTimeoutError, "query timeout") - if status_name == "CANCELLED": - return _create_adbc_error(error, OperationCancelledError, "operation cancelled") + _ = logger + if status_mapped := _map_adbc_status_exception(error): + return status_mapped raw_message = str(error) error_msg = raw_message.lower() @@ -961,10 +993,11 @@ def get_statement_config(detected_dialect: str) -> StatementConfig: "type_coercion_map": _BASE_TYPE_COERCION_MAP, } + statement_overrides: dict[str, Any] = {"dialect": sqlglot_dialect} if detected_dialect == "duckdb": parameter_overrides["preserve_parameter_format"] = False parameter_overrides["supported_execution_parameter_styles"] = {ParameterStyle.QMARK, ParameterStyle.NUMERIC} - parameter_overrides["ast_transformer"] = _transform_duckdb_copy_parameters + statement_overrides["statement_transformers"] = (_transform_duckdb_copy_parameters,) elif detected_dialect == "oracle": parameter_overrides["preserve_parameter_format"] = False parameter_overrides["supported_execution_parameter_styles"] = {ParameterStyle.POSITIONAL_COLON} @@ -976,7 +1009,7 @@ def get_statement_config(detected_dialect: str) -> StatementConfig: parameter_overrides["ast_transformer"] = build_null_pruning_transform(dialect=sqlglot_dialect) return build_statement_config_from_profile( - profile, parameter_overrides=parameter_overrides, statement_overrides={"dialect": sqlglot_dialect} + profile, parameter_overrides=parameter_overrides, statement_overrides=statement_overrides ) diff --git a/sqlspec/adapters/adbc/driver.py b/sqlspec/adapters/adbc/driver.py index c3f6cb028..c65d69ea4 100644 --- a/sqlspec/adapters/adbc/driver.py +++ b/sqlspec/adapters/adbc/driver.py @@ -11,6 +11,7 @@ from sqlspec.adapters.adbc._typing import AdbcCursor, AdbcNativeError, AdbcSessionContext from sqlspec.adapters.adbc.core import ( + _transform_duckdb_copy_parameters, apply_driver_features, cleanup_duckdb_profile_path, collect_rows, @@ -178,6 +179,7 @@ class AdbcDriver(SyncDriverAdapterBase): "_is_flightsql", "_is_postgres", "_json_serializer", + "_previous_autocommit", "_transaction_active", "dialect", ) @@ -219,6 +221,7 @@ def __init__( self._data_dictionary: AdbcDataDictionary | None = None self._column_name_cache: dict[int, tuple[Any, list[str]]] = {} self._transaction_active = False + self._previous_autocommit: bool | None = None @property def data_dictionary(self) -> "AdbcDataDictionary": @@ -435,6 +438,11 @@ def begin(self) -> None: ADBC connections operate with autocommit disabled and manage the active transaction internally, so no explicit statement is issued. """ + if hasattr(self.connection, "autocommit"): + if self._previous_autocommit is None: + self._previous_autocommit = bool(self.connection.autocommit) + if self.connection.autocommit: + self.connection.autocommit = False self._transaction_active = True def commit(self) -> None: @@ -445,6 +453,10 @@ def commit(self) -> None: raise create_mapped_exception(e) from e finally: self._transaction_active = False + if self._previous_autocommit is not None: + if hasattr(self.connection, "autocommit"): + self.connection.autocommit = self._previous_autocommit + self._previous_autocommit = None def rollback(self) -> None: """Rollback database transaction.""" @@ -454,6 +466,10 @@ def rollback(self) -> None: raise create_mapped_exception(e) from e finally: self._transaction_active = False + if self._previous_autocommit is not None: + if hasattr(self.connection, "autocommit"): + self.connection.autocommit = self._previous_autocommit + self._previous_autocommit = None def create_savepoint(self, name: str) -> None: """Create a savepoint within the current transaction. @@ -622,7 +638,7 @@ def select_to_arrow( if cursor is not None: with contextlib.suppress(Exception): cursor_manager.__exit__(None, None, None) - raise exc_handler.pending_exception from None + self._check_pending_exception(exc_handler) if arrow_result is not None: return arrow_result if cursor is None or reader is None: @@ -674,8 +690,7 @@ def select_to_arrow( arrow_schema=arrow_schema, ) - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + self._check_pending_exception(exc_handler) if arrow_result is None: msg = "Unreachable" @@ -762,8 +777,7 @@ def load_from_arrow( if isinstance(ingest_result, int) and ingest_result > 0: row_count = ingest_result - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + self._check_pending_exception(exc_handler) telemetry_payload: StorageTelemetry = { "rows_processed": row_count, "bytes_processed": 0, @@ -782,8 +796,7 @@ def load_from_arrow( handle_postgres_rollback(self._dialect_name, cursor, logger) raise - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + self._check_pending_exception(exc_handler) telemetry_payload = self._ingest_telemetry(arrow_table) telemetry_payload["destination"] = table self._attach_partition_telemetry(telemetry_payload, partitioner) @@ -874,6 +887,13 @@ def prepare_driver_parameters( def _compiled_sql( self, statement: "SQL", statement_config: "StatementConfig", flatten_single_parameters: bool = False ) -> "tuple[str, object]": + if ( + self._dialect_name == "duckdb" + and statement_config.parameter_config.ast_transformer is None + and statement.raw_sql.lstrip().upper().startswith(("COPY", "INSERT")) + ): + param_config = statement_config.parameter_config.replace(ast_transformer=_transform_duckdb_copy_parameters) + statement_config = statement_config.replace(parameter_config=param_config) compiled_sql, prepared_parameters = super()._compiled_sql( statement, statement_config, flatten_single_parameters=flatten_single_parameters ) diff --git a/tests/unit/adapters/test_adbc/test_core.py b/tests/unit/adapters/test_adbc/test_core.py index 80f035874..38e0c524b 100644 --- a/tests/unit/adapters/test_adbc/test_core.py +++ b/tests/unit/adapters/test_adbc/test_core.py @@ -27,6 +27,7 @@ resolve_many_rowcount, ) from sqlspec.core import SQL, ParameterStyle +from sqlspec.core.compiler import SQLProcessor from sqlspec.dialects import PGTextSearch from sqlspec.exceptions import ( CheckViolationError, @@ -35,6 +36,7 @@ DeadlockError, ForeignKeyViolationError, IntegrityError, + NotFoundError, NotNullViolationError, OperationalError, OperationCancelledError, @@ -642,3 +644,100 @@ def test_adbc_driver_dispatch_execute_lowercases_oracle_columns() -> None: assert result.column_names == ["user_id", "MixedCase"] assert result.selected_data == [{"user_id": 1, "MixedCase": "a"}] + + +def test_adbc_begin_with_autocommit() -> None: + """AdbcDriver.begin disables autocommit and commit/rollback restores it.""" + + class MockConn: + def __init__(self) -> None: + self.autocommit = True + self.committed = False + self.rolled_back = False + + def adbc_get_info(self) -> dict[str, str]: + return {"driver_name": "sqlite"} + + def commit(self) -> None: + self.committed = True + + def rollback(self) -> None: + self.rolled_back = True + + conn = MockConn() + driver = AdbcDriver(cast("Any", conn)) + assert conn.autocommit is True + driver.begin() + assert conn.autocommit is False + assert driver._transaction_active is True + driver.commit() + assert conn.autocommit is True + assert driver._transaction_active is False + assert conn.committed is True + + conn.autocommit = True + driver.begin() + assert conn.autocommit is False + driver.rollback() + assert conn.autocommit is True + assert driver._transaction_active is False + assert conn.rolled_back is True + + +def test_pending_exception_preserves_cause() -> None: + """_check_pending_exception preserves original exception cause.""" + + class MockConn: + def adbc_get_info(self) -> dict[str, str]: + return {"driver_name": "sqlite"} + + driver = AdbcDriver(cast("Any", MockConn())) + orig_exc = DatabaseError("query failed", status_code=AdbcStatusCode.INTERNAL) + mapped_exc = adbc_core.create_mapped_exception(orig_exc) + handler = driver.handle_database_exceptions() + handler.pending_exception = mapped_exc + + with pytest.raises(OperationalError) as exc_info: + driver._check_pending_exception(handler) + + assert exc_info.value is mapped_exc + assert exc_info.value.__cause__ is orig_exc + + +@pytest.mark.parametrize( + ("status_code", "expected_cls"), + [ + (AdbcStatusCode.TIMEOUT, QueryTimeoutError), + (AdbcStatusCode.CANCELLED, OperationCancelledError), + (AdbcStatusCode.UNAUTHENTICATED, PermissionDeniedError), + (AdbcStatusCode.UNAUTHORIZED, PermissionDeniedError), + (AdbcStatusCode.NOT_FOUND, NotFoundError), + (AdbcStatusCode.ALREADY_EXISTS, IntegrityError), + (AdbcStatusCode.INVALID_ARGUMENT, DataError), + (AdbcStatusCode.IO, DatabaseConnectionError), + (AdbcStatusCode.INTEGRITY, IntegrityError), + (AdbcStatusCode.INVALID_STATE, OperationalError), + (AdbcStatusCode.INTERNAL, OperationalError), + (AdbcStatusCode.NOT_IMPLEMENTED, SQLSpecError), + ], +) +def test_adbc_status_code_mappings(status_code: AdbcStatusCode, expected_cls: type[SQLSpecError]) -> None: + """_map_adbc_status_exception maps ADBC status codes to SQLSpecError classes.""" + error = DatabaseError("adbc error", status_code=status_code) + mapped = adbc_core.create_mapped_exception(error) + assert isinstance(mapped, expected_cls) + + +def test_duckdb_copy_ast_transformer_caching() -> None: + """DuckDB statement config omits ast_transformer so compiled SQL can be cached.""" + config = get_statement_config("duckdb").replace(enable_caching=True) + assert config.parameter_config.ast_transformer is None + + stmt = SQL("SELECT 1 WHERE x = :val", {"val": 42}, statement_config=config) + compiled_sql, _ = stmt.compile() + assert compiled_sql + + processor = SQLProcessor(config) + processor.compile("SELECT 1 WHERE x = :val", {"val": 42}) + processor.compile("SELECT 1 WHERE x = :val", {"val": 99}) + assert processor._cache_hits == 1 From 8eacd01c80a0bd558e44baa3bdccc0c7433dc935 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 03:57:10 +0000 Subject: [PATCH 101/174] fix(adbc): harden litestar session and adk memory store operations Catch duplicate key collisions in Litestar session store on non-upsert dialects with fallback to update, wrap mutating session operations with explicit rollback handling, and deduplicate DuckDB and PostgreSQL upsert paths. Pre-query existing event IDs during ADK memory store entry insertion on non-returning dialects, consolidate DB2 catalog primitives into core, align ScaNN leaf configuration documentation, and expose dialect as a property on AdbcADKStore. --- sqlspec/adapters/adbc/adk/store.py | 55 ++--- sqlspec/adapters/adbc/core.py | 17 ++ sqlspec/adapters/adbc/litestar/store.py | 228 ++++++++++-------- .../test_adbc/test_adk_dialect_support.py | 63 ++++- .../test_adbc/test_extension_stores.py | 51 ++++ 5 files changed, 283 insertions(+), 131 deletions(-) diff --git a/sqlspec/adapters/adbc/adk/store.py b/sqlspec/adapters/adbc/adk/store.py index b948ef2f0..7b763f57f 100644 --- a/sqlspec/adapters/adbc/adk/store.py +++ b/sqlspec/adapters/adbc/adk/store.py @@ -8,11 +8,14 @@ from typing import TYPE_CHECKING, Any, Final, Literal, NotRequired, cast from sqlspec.adapters.adbc.core import ( + _DB2_INDEX_EXISTS_SQL, + _DB2_TABLE_EXISTS_SQL, driver_kind_from_driver_name, is_postgres_dialect, is_shared_object_driver, resolve_dialect_from_config, ) +from sqlspec.adapters.adbc.core import split_db2_table_name as _split_db2_table_name from sqlspec.config import ADKConfig from sqlspec.core.parameters._validator import PARAMETER_REGEX from sqlspec.exceptions import ImproperConfigurationError @@ -20,7 +23,6 @@ from sqlspec.extensions.adk.memory.store import BaseSyncADKMemoryStore from sqlspec.utils.logging import get_logger from sqlspec.utils.serializers import from_json, to_json -from sqlspec.utils.text import normalize_identifier, split_qualified_identifier if TYPE_CHECKING: from collections.abc import Sequence @@ -49,7 +51,7 @@ class AdbcADKConfig(ADKConfig): """Vector index type for PostgreSQL. Defaults to 'hnsw'.""" scann_num_leaves: NotRequired[int] - """Number of leaves for ScaNN index on PostgreSQL. Defaults to 50.""" + """Number of leaves for ScaNN index on PostgreSQL. Defaults to 100.""" scann_quantizer: NotRequired[Literal["SQ8", "FP32"]] """Quantizer for ScaNN index on PostgreSQL. Defaults to 'SQ8'.""" @@ -224,6 +226,7 @@ def set_metadata(self, key: str, value: str) -> None: """Set a metadata value.""" self._set_metadata(key, value) + @property def dialect(self) -> str: """Return the detected database dialect.""" return self._dialect @@ -2194,7 +2197,20 @@ def _insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "objec with self._config.provide_connection() as conn: cursor = conn.cursor() try: - for entry in entries: + entries_to_insert = entries + if not use_returning: + event_ids = [entry["event_id"] for entry in entries if "event_id" in entry] + if event_ids: + placeholders = ", ".join("?" for _ in event_ids) + select_sql = f"SELECT event_id FROM {self._memory_table} WHERE event_id IN ({placeholders})" + self._execute(cursor, select_sql, tuple(event_ids)) + existing_rows = cursor.fetchall() + existing_event_ids = {row[0] for row in existing_rows} + entries_to_insert = [ + entry for entry in entries if entry.get("event_id") not in existing_event_ids + ] + + for entry in entries_to_insert: content_json = self._serialize_json_field(entry["content_json"]) metadata_json = self._serialize_json_field(entry.get("metadata_json")) raw_emb = entry.get("embedding") @@ -2270,21 +2286,8 @@ def _insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "objec if cursor.fetchone(): inserted_count += 1 else: - try: - self._execute(cursor, sql, params) - inserted_count += 1 - except Exception as exc: - exc_str = str(exc).lower() - if ( - "unique" in exc_str - or "constraint" in exc_str - or "duplicate" in exc_str - or "ora-00001" in exc_str - or "sql0803n" in exc_str - or "23505" in exc_str - ): - continue - raise + self._execute(cursor, sql, params) + inserted_count += 1 conn.commit() except Exception: with contextlib.suppress(Exception): @@ -2622,15 +2625,6 @@ def decode_adbc_pgvector(raw_embedding: Any) -> "list[float] | None": return [float(v) for v in raw_embedding] -_DB2_TABLE_EXISTS_SQL: Final = ( - "SELECT 1 FROM SYSCAT.TABLES WHERE TABSCHEMA = COALESCE(CAST(? AS VARCHAR(128)), CURRENT SCHEMA) AND TABNAME = ?" -) -_DB2_INDEX_EXISTS_SQL: Final = ( - "SELECT 1 FROM SYSCAT.INDEXES WHERE INDSCHEMA = COALESCE(CAST(? AS VARCHAR(128)), CURRENT SCHEMA) AND INDNAME = ?" -) -_MIN_QUALIFIED_PARTS: Final = 2 - - @lru_cache(maxsize=256) def _postgresql_store_sql(sql: str) -> str: """Cache positional SQL conversion without changing quoted or commented text.""" @@ -2669,13 +2663,6 @@ def _metadata_key_column(dialect: str) -> str: return "key" -def _split_db2_table_name(table_name: str) -> "tuple[str | None, str]": - parts = [part.strip() for part in split_qualified_identifier(table_name, quote_chars='"') if part.strip()] - if len(parts) >= _MIN_QUALIFIED_PARTS: - return normalize_identifier(parts[-2], "db2"), normalize_identifier(parts[-1], "db2") - return None, normalize_identifier(parts[0] if parts else table_name, "db2") - - def _db2_table_exists(cursor: Any, table_name: str) -> bool: schema_name, catalog_table = _split_db2_table_name(table_name) cursor.execute(_DB2_TABLE_EXISTS_SQL, (schema_name, catalog_table)) diff --git a/sqlspec/adapters/adbc/core.py b/sqlspec/adapters/adbc/core.py index d06f5121a..185ceb638 100644 --- a/sqlspec/adapters/adbc/core.py +++ b/sqlspec/adapters/adbc/core.py @@ -48,6 +48,7 @@ from sqlspec.typing import PGVECTOR_INSTALLED, Empty, import_optional_attr from sqlspec.utils.module_loader import import_optional, import_string from sqlspec.utils.serializers import from_json, to_json +from sqlspec.utils.text import normalize_identifier, split_qualified_identifier from sqlspec.utils.type_converters import build_uuid_coercions from sqlspec.utils.type_guards import has_rowcount, has_sqlstate from sqlspec.utils.uuids import uuid4, uuid_from_string @@ -2162,3 +2163,19 @@ def _apply_driver_manager_config(config: "dict[str, Any]", driver_kind: "str | N driver_profile: Final[DriverParameterProfile] = build_profile() + +_DB2_TABLE_EXISTS_SQL: Final = ( + "SELECT 1 FROM SYSCAT.TABLES WHERE TABSCHEMA = COALESCE(CAST(? AS VARCHAR(128)), CURRENT SCHEMA) AND TABNAME = ?" +) +_DB2_INDEX_EXISTS_SQL: Final = ( + "SELECT 1 FROM SYSCAT.INDEXES WHERE INDSCHEMA = COALESCE(CAST(? AS VARCHAR(128)), CURRENT SCHEMA) AND INDNAME = ?" +) +_MIN_QUALIFIED_PARTS: Final = 2 + + +def split_db2_table_name(table_name: str) -> "tuple[str | None, str]": + """Split qualified DB2 table name into schema and table components.""" + parts = [part.strip() for part in split_qualified_identifier(table_name, quote_chars='"') if part.strip()] + if len(parts) >= _MIN_QUALIFIED_PARTS: + return normalize_identifier(parts[-2], "db2"), normalize_identifier(parts[-1], "db2") + return None, normalize_identifier(parts[0] if parts else table_name, "db2") diff --git a/sqlspec/adapters/adbc/litestar/store.py b/sqlspec/adapters/adbc/litestar/store.py index ce6e2d087..fba12c796 100644 --- a/sqlspec/adapters/adbc/litestar/store.py +++ b/sqlspec/adapters/adbc/litestar/store.py @@ -13,13 +13,16 @@ - Snowflake: BINARY data type, TIMESTAMP WITH TIME ZONE, ? parameters, MERGE """ +import contextlib from datetime import UTC, datetime, timedelta from typing import TYPE_CHECKING -from sqlspec.adapters.adbc.core import is_postgres_dialect +from sqlspec.adapters.adbc.core import _DB2_INDEX_EXISTS_SQL, _DB2_TABLE_EXISTS_SQL, is_postgres_dialect +from sqlspec.adapters.adbc.core import split_db2_table_name as _split_db2_table_name +from sqlspec.exceptions import IntegrityError from sqlspec.extensions.litestar.store import BaseSQLSpecStore from sqlspec.utils.sync_tools import async_ -from sqlspec.utils.text import normalize_identifier, split_qualified_identifier +from sqlspec.utils.text import normalize_identifier if TYPE_CHECKING: from sqlspec.adapters.adbc.config import AdbcConfig @@ -333,9 +336,14 @@ def _create_table(self) -> None: END; """.strip() with self._config.provide_session() as driver: - driver.execute(table_block) - driver.execute(index_block) - driver.commit() + try: + driver.execute(table_block) + driver.execute(index_block) + driver.commit() + except Exception: + with contextlib.suppress(Exception): + driver.rollback() + raise self._log_table_created() return @@ -343,18 +351,30 @@ def _create_table(self) -> None: schema_name, catalog_table = _split_db2_table_name(self._table_name) index_name = normalize_identifier(f"idx_{self._table_name}_expires_at", "db2") with self._config.provide_session() as driver: - if driver.select_one_or_none(_DB2_TABLE_EXISTS_SQL, (schema_name, catalog_table)) is None: - driver.execute(self._table_ddl().strip()) - if driver.select_one_or_none(_DB2_INDEX_EXISTS_SQL, (schema_name, index_name)) is None: - driver.execute(f"CREATE INDEX idx_{self._table_name}_expires_at ON {self._table_name}(expires_at)") - driver.commit() + try: + if driver.select_one_or_none(_DB2_TABLE_EXISTS_SQL, (schema_name, catalog_table)) is None: + driver.execute(self._table_ddl().strip()) + if driver.select_one_or_none(_DB2_INDEX_EXISTS_SQL, (schema_name, index_name)) is None: + driver.execute( + f"CREATE INDEX idx_{self._table_name}_expires_at ON {self._table_name}(expires_at)" + ) + driver.commit() + except Exception: + with contextlib.suppress(Exception): + driver.rollback() + raise self._log_table_created() return sql_text = self._table_ddl() with self._config.provide_session() as driver: - driver.execute_script(sql_text) - driver.commit() + try: + driver.execute_script(sql_text) + driver.commit() + except Exception: + with contextlib.suppress(Exception): + driver.rollback() + raise self._log_table_created() def _drop_table_sql(self) -> "list[str]": @@ -450,25 +470,7 @@ def _set(self, key: str, value: "str | bytes", expires_in: "int | timedelta | No INSERT OR REPLACE INTO {self._table_name} (session_id, data, expires_at) VALUES ({p1}, {p2}, {p3}) """ - elif is_postgres_dialect(dialect): - if expires_at is None: - sql = f""" - INSERT INTO {self._table_name} (session_id, data, expires_at) - VALUES ({p1}, {p2}, NULL) - ON CONFLICT (session_id) DO UPDATE - SET data = EXCLUDED.data, expires_at = NULL - """ - with self._config.provide_session() as driver: - driver.execute(sql, key, data) - driver.commit() - return - sql = f""" - INSERT INTO {self._table_name} (session_id, data, expires_at) - VALUES ({p1}, {p2}, {p3}) - ON CONFLICT (session_id) DO UPDATE - SET data = EXCLUDED.data, expires_at = EXCLUDED.expires_at - """ - else: + elif is_postgres_dialect(dialect) or dialect == "duckdb": if expires_at is None: sql = f""" INSERT INTO {self._table_name} (session_id, data, expires_at) @@ -477,8 +479,13 @@ def _set(self, key: str, value: "str | bytes", expires_in: "int | timedelta | No SET data = EXCLUDED.data, expires_at = NULL """ with self._config.provide_session() as driver: - driver.execute(sql, key, data) - driver.commit() + try: + driver.execute(sql, key, data) + driver.commit() + except Exception: + with contextlib.suppress(Exception): + driver.rollback() + raise return sql = f""" INSERT INTO {self._table_name} (session_id, data, expires_at) @@ -494,38 +501,59 @@ def _set(self, key: str, value: "str | bytes", expires_in: "int | timedelta | No """ elif dialect in {"bigquery", "snowflake", "oracle", "db2", "mssql"}: with self._config.provide_session() as driver: - check_sql = f"SELECT COUNT(*) as count FROM {self._table_name} WHERE session_id = {p1}" - result = driver.select_one(check_sql, key) - exists = result and result.get("count", 0) > 0 - - if exists: - if expires_at is None: + try: + check_sql = f"SELECT COUNT(*) as count FROM {self._table_name} WHERE session_id = {p1}" + result = driver.select_one(check_sql, key) + exists = result and result.get("count", 0) > 0 + + if exists: + if expires_at is None: + sql = f""" + UPDATE {self._table_name} + SET data = {p1}, expires_at = NULL + WHERE session_id = {p2} + """ + driver.execute(sql, data, key) + else: + sql = f""" + UPDATE {self._table_name} + SET data = {p1}, expires_at = {p2} + WHERE session_id = {p3} + """ + driver.execute(sql, data, expires_at, key) + elif expires_at is None: sql = f""" - UPDATE {self._table_name} - SET data = {p1}, expires_at = NULL - WHERE session_id = {p2} + INSERT INTO {self._table_name} (session_id, data, expires_at) + VALUES ({p1}, {p2}, NULL) """ - driver.execute(sql, data, key) + try: + driver.execute(sql, key, data) + except IntegrityError: + sql = f""" + UPDATE {self._table_name} + SET data = {p1}, expires_at = NULL + WHERE session_id = {p2} + """ + driver.execute(sql, data, key) else: sql = f""" - UPDATE {self._table_name} - SET data = {p1}, expires_at = {p2} - WHERE session_id = {p3} + INSERT INTO {self._table_name} (session_id, data, expires_at) + VALUES ({p1}, {p2}, {p3}) """ - driver.execute(sql, data, expires_at, key) - elif expires_at is None: - sql = f""" - INSERT INTO {self._table_name} (session_id, data, expires_at) - VALUES ({p1}, {p2}, NULL) - """ - driver.execute(sql, key, data) - else: - sql = f""" - INSERT INTO {self._table_name} (session_id, data, expires_at) - VALUES ({p1}, {p2}, {p3}) - """ - driver.execute(sql, key, data, expires_at) - driver.commit() + try: + driver.execute(sql, key, data, expires_at) + except IntegrityError: + sql = f""" + UPDATE {self._table_name} + SET data = {p1}, expires_at = {p2} + WHERE session_id = {p3} + """ + driver.execute(sql, data, expires_at, key) + driver.commit() + except Exception: + with contextlib.suppress(Exception): + driver.rollback() + raise return else: sql = f""" @@ -536,8 +564,13 @@ def _set(self, key: str, value: "str | bytes", expires_in: "int | timedelta | No """ with self._config.provide_session() as driver: - driver.execute(sql, key, data, expires_at) - driver.commit() + try: + driver.execute(sql, key, data, expires_at) + driver.commit() + except Exception: + with contextlib.suppress(Exception): + driver.rollback() + raise def _delete(self, key: str) -> None: """Synchronous implementation of delete using ADBC driver.""" @@ -545,8 +578,13 @@ def _delete(self, key: str) -> None: sql = f"DELETE FROM {self._table_name} WHERE session_id = {p1}" with self._config.provide_session() as driver: - driver.execute(sql, key) - driver.commit() + try: + driver.execute(sql, key) + driver.commit() + except Exception: + with contextlib.suppress(Exception): + driver.rollback() + raise def _delete_all(self) -> None: """Synchronous implementation of delete_all using ADBC driver.""" @@ -554,8 +592,13 @@ def _delete_all(self) -> None: sql = f"DELETE FROM {self._table_name}" with self._config.provide_session() as driver: - driver.execute(sql) - driver.commit() + try: + driver.execute(sql) + driver.commit() + except Exception: + with contextlib.suppress(Exception): + driver.rollback() + raise def _exists(self, key: str) -> bool: """Synchronous implementation of exists using ADBC driver.""" @@ -614,36 +657,31 @@ def _delete_expired(self) -> int: delete_sql = f"DELETE FROM {self._table_name} WHERE expires_at <= {current_ts}" with self._config.provide_session() as driver: - result = driver.select_one(count_sql) - count = result.get("count", 0) if result else 0 + try: + result = driver.select_one(count_sql) + count = result.get("count", 0) if result else 0 + + if count > 0: + driver.execute(delete_sql) + driver.commit() + self._log_delete_expired(count) + except Exception: + with contextlib.suppress(Exception): + driver.rollback() + raise + else: + return count + with self._config.provide_session() as driver: + try: + exec_result = driver.execute(sql) + driver.commit() + count = exec_result.rows_affected if count > 0: - driver.execute(delete_sql) - driver.commit() self._log_delete_expired(count) - + except Exception: + with contextlib.suppress(Exception): + driver.rollback() + raise + else: return count - - with self._config.provide_session() as driver: - exec_result = driver.execute(sql) - driver.commit() - count = exec_result.rows_affected - if count > 0: - self._log_delete_expired(count) - return count - - -_DB2_TABLE_EXISTS_SQL = ( - "SELECT 1 FROM SYSCAT.TABLES WHERE TABSCHEMA = COALESCE(CAST(? AS VARCHAR(128)), CURRENT SCHEMA) AND TABNAME = ?" -) -_DB2_INDEX_EXISTS_SQL = ( - "SELECT 1 FROM SYSCAT.INDEXES WHERE INDSCHEMA = COALESCE(CAST(? AS VARCHAR(128)), CURRENT SCHEMA) AND INDNAME = ?" -) -_MIN_QUALIFIED_PARTS = 2 - - -def _split_db2_table_name(table_name: str) -> "tuple[str | None, str]": - parts = [part.strip() for part in split_qualified_identifier(table_name, quote_chars='"') if part.strip()] - if len(parts) >= _MIN_QUALIFIED_PARTS: - return normalize_identifier(parts[-2], "db2"), normalize_identifier(parts[-1], "db2") - return None, normalize_identifier(parts[0] if parts else table_name, "db2") diff --git a/tests/unit/adapters/test_adbc/test_adk_dialect_support.py b/tests/unit/adapters/test_adbc/test_adk_dialect_support.py index db8accf1f..506859bcd 100644 --- a/tests/unit/adapters/test_adbc/test_adk_dialect_support.py +++ b/tests/unit/adapters/test_adbc/test_adk_dialect_support.py @@ -399,7 +399,7 @@ def test_detect_enterprise_dialects_in_adk_stores(driver_name: str, expected_dia config = AdbcConfig(connection_config={"driver_name": driver_name}) store = AdbcADKStore(config) mem_store = AdbcADKMemoryStore(config) - assert store.dialect() == expected_dialect + assert store.dialect == expected_dialect assert mem_store.dialect == expected_dialect @@ -408,7 +408,7 @@ def test_detect_dialect_fallback_to_statement_config() -> None: config = AdbcConfig(connection_config={}, statement_config=StatementConfig(dialect="oracle")) store = AdbcADKStore(config) mem_store = AdbcADKMemoryStore(config) - assert store.dialect() == "oracle" + assert store.dialect == "oracle" assert mem_store.dialect == "oracle" @@ -760,3 +760,62 @@ def _store_for(driver_name: str) -> AdbcADKStore: def _normalized(sql: str) -> str: return " ".join(sql.split()) + + +def test_oracle_adk_memory_store_pre_queries_existing_event_ids(monkeypatch: pytest.MonkeyPatch) -> None: + """AdbcADKMemoryStore pre-queries existing event_ids on non-returning dialects.""" + executed_statements: list[str] = [] + mock_cursor = MagicMock() + mock_cursor.fetchall.return_value = [("evt-existing",)] + + def record_execute(sql: str, _params: Any = None) -> None: + executed_statements.append(sql) + + mock_cursor.execute.side_effect = record_execute + mock_conn = MagicMock() + mock_conn.cursor.return_value = mock_cursor + mock_ctx = MagicMock() + mock_ctx.__enter__.return_value = mock_conn + mock_ctx.__exit__.return_value = None + monkeypatch.setattr(AdbcConfig, "provide_connection", lambda *_a, **_k: mock_ctx) + + config = AdbcConfig( + connection_config={"driver_name": "oracle", "uri": ":memory:"}, + extension_config={"adk": {"enable_memory": True}}, + ) + mem_store = AdbcADKMemoryStore(config) + now = datetime(2026, 6, 1, 12, 0, 0, tzinfo=UTC) + + count = mem_store.insert_memory_entries([ + { + "id": "mem-1", + "session_id": "sess-1", + "app_name": "app1", + "user_id": "user1", + "scope": "user", + "event_id": "evt-existing", + "author": "user", + "timestamp": now, + "content_json": {"parts": [{"text": "hello"}]}, + "content_text": "hello", + "metadata_json": None, + "inserted_at": now, + }, + { + "id": "mem-2", + "session_id": "sess-1", + "app_name": "app1", + "user_id": "user1", + "scope": "user", + "event_id": "evt-new", + "author": "user", + "timestamp": now, + "content_json": {"parts": [{"text": "world"}]}, + "content_text": "world", + "metadata_json": None, + "inserted_at": now, + }, + ]) + + assert any("SELECT event_id FROM" in s for s in executed_statements) + assert count == 1 diff --git a/tests/unit/adapters/test_adbc/test_extension_stores.py b/tests/unit/adapters/test_adbc/test_extension_stores.py index a7b186355..c4b89a785 100644 --- a/tests/unit/adapters/test_adbc/test_extension_stores.py +++ b/tests/unit/adapters/test_adbc/test_extension_stores.py @@ -9,6 +9,7 @@ from sqlspec.adapters.adbc.config import AdbcConfig from sqlspec.adapters.adbc.events.store import AdbcEventQueueStore from sqlspec.adapters.adbc.litestar.store import AdbcStore +from sqlspec.exceptions import IntegrityError __all__ = () @@ -98,5 +99,55 @@ def test_adbc_event_queue_store_enterprise_dialects( assert expected_drop_fragment in drop_sql +def test_adbc_litestar_store_concurrent_set_fallback() -> None: + """AdbcStore._set catches IntegrityError on INSERT race condition and falls back to UPDATE.""" + executed: list[tuple[str, tuple[Any, ...]]] = [] + + class _DriverStub: + def select_one(self, sql: str, *params: Any) -> dict[str, Any]: + executed.append((sql.strip(), params)) + return {"count": 0} + + def execute(self, sql: str, *params: Any) -> None: + executed.append((sql.strip(), params)) + if "INSERT INTO" in sql: + raise IntegrityError("duplicate key value violates unique constraint") + + def commit(self) -> None: + pass + + config_stub = SimpleNamespace(extension_config={"litestar": {}}, provide_session=lambda: nullcontext(_DriverStub())) + store = AdbcStore(cast("Any", config_stub)) + store._dialect = "oracle" + store._set("sess_1", b"payload", expires_in=None) + + assert any("UPDATE" in item[0] for item in executed) + + +def test_adbc_litestar_store_rollback_on_failure() -> None: + """Mutating operations on AdbcStore invoke rollback on execution error.""" + rolled_back = False + + class _DriverStub: + def execute(self, _sql: str, *_params: Any) -> None: + raise RuntimeError("disk failure") + + def commit(self) -> None: + pass + + def rollback(self) -> None: + nonlocal rolled_back + rolled_back = True + + config_stub = SimpleNamespace(extension_config={"litestar": {}}, provide_session=lambda: nullcontext(_DriverStub())) + store = AdbcStore(cast("Any", config_stub)) + store._dialect = "oracle" + + with pytest.raises(RuntimeError, match="disk failure"): + store._delete("sess_1") + + assert rolled_back is True + + def _make_adbc_config(dialect_uri: str) -> AdbcConfig: return AdbcConfig(connection_config={"uri": dialect_uri}) From 02d245c1e3766aa1d4a29a6aeae9fc3ab95f8314 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 03:57:10 +0000 Subject: [PATCH 102/174] fix(adbc): safeguard Snowflake index metadata and retain empty foreign key results Guard against missing Snowflake index templates by returning empty results without raising missing file errors. Preserve valid empty foreign key results from native ADBC object metadata on constraint-free tables and only fall back when constraint structures are genuinely omitted. --- sqlspec/adapters/adbc/data_dictionary.py | 34 +++++++++++------- .../test_adbc/test_data_dictionary.py | 36 +++++++++++++++++-- 2 files changed, 54 insertions(+), 16 deletions(-) diff --git a/sqlspec/adapters/adbc/data_dictionary.py b/sqlspec/adapters/adbc/data_dictionary.py index 594ae2c84..62b01a363 100644 --- a/sqlspec/adapters/adbc/data_dictionary.py +++ b/sqlspec/adapters/adbc/data_dictionary.py @@ -346,18 +346,24 @@ def get_indexes( operation = "by_schema" if table is None else "by_table" return driver.select(self._get_query(dialect, "indexes", operation), schema_type=IndexMetadata) - if table is None: + if dialect == "snowflake": + return [] + + try: + if table is None: + return driver.select( + self._get_query(dialect, "indexes", "by_schema"), schema_name=schema_name, schema_type=IndexMetadata + ) + + table_name = self._resolve_identifier(dialect, table) return driver.select( - self._get_query(dialect, "indexes", "by_schema"), schema_name=schema_name, schema_type=IndexMetadata + self._get_query(dialect, "indexes", "by_table"), + schema_name=schema_name, + table_name=table_name, + schema_type=IndexMetadata, ) - - table_name = self._resolve_identifier(dialect, table) - return driver.select( - self._get_query(dialect, "indexes", "by_table"), - schema_name=schema_name, - table_name=table_name, - schema_type=IndexMetadata, - ) + except SQLFileNotFoundError: + return [] def get_foreign_keys( self, driver: "AdbcDriver", table: "str | None" = None, schema: "str | None" = None @@ -821,10 +827,12 @@ def _native_get_foreign_keys( self, driver: "AdbcDriver", dialect: str, table_name: "str | None", schema_name: "str | None" ) -> "list[ForeignKeyMetadata]": rows = self._native_get_objects(driver, dialect, "all", schema_name, table_name) - foreign_keys = _normalize_native_foreign_keys(rows, table_name_exact=table_name) - if not foreign_keys: + if not rows: raise _NativeMetadataIncompleteError - return foreign_keys + for _, _, table in _iter_object_tables(rows): + if table.get("table_constraints") is None: + raise _NativeMetadataIncompleteError + return _normalize_native_foreign_keys(rows, table_name_exact=table_name) class _NativeMetadataIncompleteError(Exception): diff --git a/tests/unit/adapters/test_adbc/test_data_dictionary.py b/tests/unit/adapters/test_adbc/test_data_dictionary.py index eb0d006a6..c989e34be 100644 --- a/tests/unit/adapters/test_adbc/test_data_dictionary.py +++ b/tests/unit/adapters/test_adbc/test_data_dictionary.py @@ -200,9 +200,7 @@ def test_get_foreign_keys_empty_native_result_falls_back() -> None: "catalog_db_schemas": [ { "db_schema_name": "main", - "db_schema_tables": [ - {"table_name": "t2", "table_type": "table", "table_columns": [], "table_constraints": []} - ], + "db_schema_tables": [{"table_name": "t2", "table_type": "table", "table_columns": []}], } ], } @@ -216,6 +214,38 @@ def test_get_foreign_keys_empty_native_result_falls_back() -> None: driver.select.assert_called_once() +def test_empty_foreign_keys_not_failed() -> None: + """Tables with empty constraints list should return empty list without falling back.""" + driver = Mock() + driver.dialect = "duckdb" + driver.connection.adbc_get_objects.return_value = _make_reader([ + { + "catalog_name": "memory", + "catalog_db_schemas": [ + { + "db_schema_name": "main", + "db_schema_tables": [ + {"table_name": "t2", "table_type": "table", "table_columns": [], "table_constraints": []} + ], + } + ], + } + ]) + + result = AdbcDataDictionary().get_foreign_keys(driver, table="t2") + + assert result == [] + driver.select.assert_not_called() + + +def test_snowflake_indexes_fallback() -> None: + """Snowflake dialect returns empty indexes without raising SQLFileNotFoundError.""" + driver = Mock() + driver.dialect = "snowflake" + result = AdbcDataDictionary().get_indexes(driver, table="t1") + assert result == [] + + def test_get_columns_schema_wide_incomplete_falls_back() -> None: """Schema-wide native column discovery should fall back when type names are incomplete.""" driver = Mock() From 322175e57375d7e2b449c75a230ca28ffa6ef830 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 03:57:10 +0000 Subject: [PATCH 103/174] perf(adbc): eliminate heap allocation and boxings in uuid rewriting and cache converters Refactor wrap_parameter closure inside _rewrite_postgres_uuid_placeholders into an iterative AST walk to eliminate heap environment allocation in compiled code. Replace @lru_cache with an internal dict-based LRU preserving cache statistics and clear operations, bind NoneType at module top level for strict type coercion checking, and cache output converter instances in get_adbc_type_converter. --- sqlspec/adapters/adbc/core.py | 71 +++++++++++++++++----- sqlspec/adapters/adbc/type_converter.py | 11 +++- tests/unit/adapters/test_adbc/test_core.py | 35 ++++++++++- 3 files changed, 99 insertions(+), 18 deletions(-) diff --git a/sqlspec/adapters/adbc/core.py b/sqlspec/adapters/adbc/core.py index 185ceb638..e7d1e875d 100644 --- a/sqlspec/adapters/adbc/core.py +++ b/sqlspec/adapters/adbc/core.py @@ -6,8 +6,9 @@ import re import tempfile from collections.abc import Sized -from functools import lru_cache, partial +from functools import partial from pathlib import Path +from types import SimpleNamespace from typing import TYPE_CHECKING, Any, Final, cast from uuid import UUID @@ -106,8 +107,9 @@ COLUMN_CACHE_MAX_SIZE: int = 256 _IMPLICIT_UPPER_COLUMN_PATTERN: Final[re.Pattern[str]] = re.compile(r"^(?!\d)(?:[A-Z0-9_]+)$") _UUID_TYPES: Final[tuple[type[Any], ...]] = tuple(build_uuid_coercions()) +NoneType: Final = type(None) _STRICT_COERCION_ALLOWED_TYPES: Final[tuple[type[Any], ...]] = ( - type(None), + NoneType, bool, int, float, @@ -1450,7 +1452,32 @@ def _is_uuid_array_data_type(target: Any) -> bool: return len(element_types) == 1 and _is_uuid_data_type(element_types[0]) -@lru_cache(maxsize=256) +class _RewriteCacheState: + __slots__ = ("cache", "hits", "maxsize", "misses") + + def __init__(self, maxsize: int = 256) -> None: + self.maxsize = maxsize + self.hits = 0 + self.misses = 0 + self.cache: dict[ + tuple[str, tuple[int, ...], tuple[int, ...], str], tuple[str, tuple[int, ...], tuple[int, ...]] + ] = {} + + def clear(self) -> None: + self.hits = 0 + self.misses = 0 + self.cache.clear() + + def info(self) -> Any: + return SimpleNamespace(hits=self.hits, misses=self.misses, maxsize=self.maxsize, currsize=len(self.cache)) + + +_REWRITE_CACHE_STATE: Final[_RewriteCacheState] = _RewriteCacheState(256) +_REWRITE_UUID_CACHE: Final[ + dict[tuple[str, tuple[int, ...], tuple[int, ...], str], tuple[str, tuple[int, ...], tuple[int, ...]]] +] = _REWRITE_CACHE_STATE.cache + + def _rewrite_postgres_uuid_placeholders( sql: str, scalar_ordinals: "tuple[int, ...]", array_ordinals: "tuple[int, ...]", dialect: str ) -> "tuple[str, tuple[int, ...], tuple[int, ...]]": @@ -1463,6 +1490,13 @@ def _rewrite_postgres_uuid_placeholders( repeated executions of the same statement shape reuse the rewrite without parsing again. """ + cache_key = (sql, scalar_ordinals, array_ordinals, dialect) + cached = _REWRITE_CACHE_STATE.cache.get(cache_key) + if cached is not None: + _REWRITE_CACHE_STATE.hits += 1 + return cached + + _REWRITE_CACHE_STATE.misses += 1 try: expressions = cast( "list[exp.Expression]", @@ -1497,21 +1531,26 @@ def _rewrite_postgres_uuid_placeholders( effective_array_set = set(effective_arrays) effective_set = set(effective_scalars) | effective_array_set - def wrap_parameter(node: exp.Expression) -> exp.Expression: - if not isinstance(node, exp.Parameter): - return node - ordinal = _parameter_ordinal(node) - if ordinal is None or ordinal not in effective_set: - return node - if _direct_parameter_cast(node) is not None: - return node - target_type = "UUID[]" if ordinal in effective_array_set else "UUID" - return exp.Cast(this=node.copy(), to=exp.DataType.build(target_type, dialect=dialect)) - for expression in expressions: - expression.transform(wrap_parameter, copy=False) + for node in list(expression.find_all(exp.Parameter)): + ordinal = _parameter_ordinal(node) + if ordinal is None or ordinal not in effective_set: + continue + if _direct_parameter_cast(node) is not None: + continue + target_type = "UUID[]" if ordinal in effective_array_set else "UUID" + node.replace(exp.Cast(this=node.copy(), to=exp.DataType.build(target_type, dialect=dialect))) + rewritten = "; ".join(expression.sql(dialect=dialect, copy=False) for expression in expressions) - return rewritten, effective_scalars, effective_arrays + result = (rewritten, effective_scalars, effective_arrays) + if len(_REWRITE_CACHE_STATE.cache) >= _REWRITE_CACHE_STATE.maxsize: + _REWRITE_CACHE_STATE.cache.pop(next(iter(_REWRITE_CACHE_STATE.cache))) + _REWRITE_CACHE_STATE.cache[cache_key] = result + return result + + +setattr(_rewrite_postgres_uuid_placeholders, "cache_clear", _REWRITE_CACHE_STATE.clear) +setattr(_rewrite_postgres_uuid_placeholders, "cache_info", _REWRITE_CACHE_STATE.info) def _sequence_holds_uuid(values: "list[Any] | tuple[Any, ...]") -> bool: diff --git a/sqlspec/adapters/adbc/type_converter.py b/sqlspec/adapters/adbc/type_converter.py index 8183897fc..38efb0325 100644 --- a/sqlspec/adapters/adbc/type_converter.py +++ b/sqlspec/adapters/adbc/type_converter.py @@ -88,6 +88,9 @@ def convert_sequence( return converted +_CONVERTER_CACHE: Final[dict[tuple[str, Any], ADBCOutputConverter]] = {} + + def get_adbc_type_converter( dialect: str, cache_size: int = 5000, *, json_serializer: "Callable[[Any], str] | None" = None ) -> ADBCOutputConverter: @@ -102,4 +105,10 @@ def get_adbc_type_converter( Configured ADBCOutputConverter instance. """ _ = cache_size - return ADBCOutputConverter(dialect, json_serializer=json_serializer or to_json) + serializer = json_serializer or to_json + cache_key = (dialect, serializer) + converter = _CONVERTER_CACHE.get(cache_key) + if converter is None: + converter = ADBCOutputConverter(dialect, json_serializer=serializer) + _CONVERTER_CACHE[cache_key] = converter + return converter diff --git a/tests/unit/adapters/test_adbc/test_core.py b/tests/unit/adapters/test_adbc/test_core.py index 38e0c524b..b8bcdd8fd 100644 --- a/tests/unit/adapters/test_adbc/test_core.py +++ b/tests/unit/adapters/test_adbc/test_core.py @@ -2,7 +2,7 @@ from collections.abc import Sequence from decimal import Decimal -from types import SimpleNamespace +from types import NoneType, SimpleNamespace from typing import Any, cast from uuid import UUID @@ -26,6 +26,7 @@ resolve_dialect_from_config, resolve_many_rowcount, ) +from sqlspec.adapters.adbc.type_converter import get_adbc_type_converter from sqlspec.core import SQL, ParameterStyle from sqlspec.core.compiler import SQLProcessor from sqlspec.dialects import PGTextSearch @@ -741,3 +742,35 @@ def test_duckdb_copy_ast_transformer_caching() -> None: processor.compile("SELECT 1 WHERE x = :val", {"val": 42}) processor.compile("SELECT 1 WHERE x = :val", {"val": 99}) assert processor._cache_hits == 1 + + +def test_rewrite_postgres_uuid_placeholders() -> None: + """_rewrite_postgres_uuid_placeholders rewrites positional parameters to explicit UUID casts.""" + sql = "SELECT $1, $2" + rewritten, scalars, arrays = adbc_core._rewrite_postgres_uuid_placeholders(sql, (1,), (2,), "postgres") + assert "UUID[]" in rewritten + assert scalars == (1,) + assert arrays == (2,) + + +def test_rewrite_uuid_cache() -> None: + """_rewrite_postgres_uuid_placeholders caches results in _REWRITE_UUID_CACHE.""" + sql = "SELECT $1 WHERE id = $2" + key = (sql, (1,), (), "postgres") + adbc_core._REWRITE_UUID_CACHE.pop(key, None) + res1 = adbc_core._rewrite_postgres_uuid_placeholders(sql, (1,), (), "postgres") + assert key in adbc_core._REWRITE_UUID_CACHE + res2 = adbc_core._rewrite_postgres_uuid_placeholders(sql, (1,), (), "postgres") + assert res1 is res2 + + +def test_strict_coercion_types() -> None: + """_STRICT_COERCION_ALLOWED_TYPES contains NoneType bound once at module top level.""" + assert NoneType in adbc_core._STRICT_COERCION_ALLOWED_TYPES + + +def test_type_converter_caching() -> None: + """get_adbc_type_converter caches output converter instances.""" + conv1 = get_adbc_type_converter("postgres") + conv2 = get_adbc_type_converter("postgres") + assert conv1 is conv2 From b5a28607a528a1d691cefa245e5d401d70e8970b Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 03:57:10 +0000 Subject: [PATCH 104/174] refactor(adbc): consolidate driver features, isolation, and stream chunking Pass driver reference to AdbcSessionContext to avoid lazy imports, route auxiliary executions through ambient exception checks, and buffer stream chunks according to chunk_size. Normalize Arrow schema field names to lowercase when requested on RecordBatchReader results, add Snowflake to lowercase defaults, register bytearray and memoryview coercions, remove dead alias, and consolidate duplicate driver feature alias pairs. --- sqlspec/adapters/adbc/_typing.py | 11 ++- sqlspec/adapters/adbc/config.py | 9 +- sqlspec/adapters/adbc/core.py | 33 ++++--- sqlspec/adapters/adbc/driver.py | 79 ++++++++++++--- .../test_adbc_config_normalization.py | 18 ++++ .../test_adbc/test_adbc_serialization.py | 10 +- tests/unit/adapters/test_adbc/test_core.py | 98 +++++++++++++++++++ 7 files changed, 214 insertions(+), 44 deletions(-) diff --git a/sqlspec/adapters/adbc/_typing.py b/sqlspec/adapters/adbc/_typing.py index ca995dfd2..e4fff3604 100644 --- a/sqlspec/adapters/adbc/_typing.py +++ b/sqlspec/adapters/adbc/_typing.py @@ -75,6 +75,7 @@ class AdbcSessionContext: "_acquire_connection", "_connection", "_driver", + "_driver_class", "_driver_features", "_prepare_driver", "_release_connection", @@ -88,21 +89,27 @@ def __init__( statement_config: "StatementConfig", driver_features: "dict[str, Any]", prepare_driver: "Callable[[AdbcDriver], AdbcDriver]", + driver_class: "type[AdbcDriver] | None" = None, ) -> None: self._acquire_connection = acquire_connection self._release_connection = release_connection self._statement_config = statement_config self._driver_features = driver_features self._prepare_driver = prepare_driver + self._driver_class = driver_class self._connection: Any = None self._driver: AdbcDriver | None = None def __enter__(self) -> "AdbcDriver": - from sqlspec.adapters.adbc.driver import AdbcDriver + driver_class = self._driver_class + if driver_class is None: + from sqlspec.adapters.adbc.driver import AdbcDriver + + driver_class = AdbcDriver self._connection = self._acquire_connection() resolved_dialect = resolve_dialect_name(self._statement_config.dialect) or None - self._driver = AdbcDriver( + self._driver = driver_class( connection=self._connection, statement_config=self._statement_config, driver_features=self._driver_features, diff --git a/sqlspec/adapters/adbc/config.py b/sqlspec/adapters/adbc/config.py index 18fe38548..904ff4cc4 100644 --- a/sqlspec/adapters/adbc/config.py +++ b/sqlspec/adapters/adbc/config.py @@ -149,12 +149,8 @@ class AdbcDriverFeatures(TypedDict): """Detect explicit SQL casts such as ``::jsonb`` or ``::vector`` during parameter preparation (defaults to ``True``).""" enable_strict_type_coercion: NotRequired[bool] """Raise an error on unsupported parameter types instead of best-effort conversion (defaults to ``False``).""" - strict_type_coercion: NotRequired[bool] - """Alias normalized to ``enable_strict_type_coercion``.""" enable_arrow_extension_types: NotRequired[bool] """Preserve PyArrow extension type metadata when materializing Arrow results (defaults to ``True``).""" - arrow_extension_types: NotRequired[bool] - """Alias normalized to ``enable_arrow_extension_types``.""" enable_pgvector: NotRequired[bool] """Probe ``pg_extension`` for ``vector`` on PostgreSQL ADBC connections (defaults to ``True`` when ``pgvector`` is installed).""" enable_paradedb: NotRequired[bool] @@ -162,7 +158,7 @@ class AdbcDriverFeatures(TypedDict): enable_pg_textsearch: NotRequired[bool] """Probe ``pg_extension`` for ``pg_textsearch`` BM25 support on PostgreSQL ADBC connections (defaults to ``True``).""" enable_lowercase_column_names: NotRequired[bool] - """Normalize unquoted uppercase column names to lowercase (defaults to ``True`` on Oracle and Db2 ADBC backends).""" + """Normalize unquoted uppercase column names to lowercase (defaults to ``True`` on Oracle, Db2, and Snowflake ADBC backends).""" enable_events: NotRequired[bool] """Enable SQLSpec table-backed event channel support.""" on_connection_create: "NotRequired[Callable[[AdbcConnection], None]]" @@ -311,7 +307,7 @@ def __init__( explicit_lowercase = bool(driver_features and "enable_lowercase_column_names" in driver_features) statement_config, driver_features = apply_driver_features(statement_config, driver_features) features_dict = dict(driver_features) if driver_features else {} - if not explicit_lowercase and self._resolved_dialect in {"oracle", "db2"}: + if not explicit_lowercase and self._resolved_dialect in {"oracle", "db2", "snowflake"}: features_dict["enable_lowercase_column_names"] = True self._user_connection_hook = cast( "Callable[[AdbcConnection], None] | None", features_dict.pop("on_connection_create", None) @@ -412,6 +408,7 @@ def provide_session( statement_config=statement_config, driver_features=handler.driver_features, prepare_driver=self._prepare_driver, + driver_class=AdbcDriver, ) def get_signature_namespace(self) -> "dict[str, Any]": diff --git a/sqlspec/adapters/adbc/core.py b/sqlspec/adapters/adbc/core.py index e7d1e875d..38767dcb6 100644 --- a/sqlspec/adapters/adbc/core.py +++ b/sqlspec/adapters/adbc/core.py @@ -1034,14 +1034,16 @@ def apply_driver_features( json_serializer = cast("Callable[[Any], str] | None", processed_features.setdefault("json_serializer", to_json)) processed_features.setdefault("enable_cast_detection", True) - processed_features.setdefault("strict_type_coercion", False) - processed_features.setdefault("enable_strict_type_coercion", processed_features["strict_type_coercion"]) - processed_features.setdefault("arrow_extension_types", True) - processed_features.setdefault("enable_arrow_extension_types", processed_features["arrow_extension_types"]) + if "strict_type_coercion" in processed_features: + processed_features.setdefault("enable_strict_type_coercion", processed_features.pop("strict_type_coercion")) + processed_features.setdefault("enable_strict_type_coercion", False) + if "arrow_extension_types" in processed_features: + processed_features.setdefault("enable_arrow_extension_types", processed_features.pop("arrow_extension_types")) + processed_features.setdefault("enable_arrow_extension_types", True) processed_features.setdefault("enable_pgvector", PGVECTOR_INSTALLED) processed_features.setdefault("enable_paradedb", True) processed_features.setdefault("enable_pg_textsearch", True) - processed_features.setdefault("enable_lowercase_column_names", resolved_dialect in {"oracle", "db2"}) + processed_features.setdefault("enable_lowercase_column_names", resolved_dialect in {"oracle", "db2", "snowflake"}) processed_features.setdefault("enable_statement_profiling", False) if json_serializer is not None: @@ -1399,6 +1401,8 @@ def _convert_array_for_postgres_adbc(value: Any) -> Any: int: _identity, float: _identity, bytes: _identity, + bytearray: bytes, + memoryview: bytes, tuple: _convert_array_for_postgres_adbc, list: _convert_array_for_postgres_adbc, dict: _identity, @@ -1795,15 +1799,15 @@ def _map_sqlstate_exception(error: Any, sqlstate: str) -> SQLSpecError: def _normalize_adbc_driver_features(processed_features: "dict[str, Any]") -> "dict[str, Any]": - if "strict_type_coercion" in processed_features and "enable_strict_type_coercion" not in processed_features: - processed_features["enable_strict_type_coercion"] = processed_features["strict_type_coercion"] - if "enable_strict_type_coercion" in processed_features and "strict_type_coercion" not in processed_features: - processed_features["strict_type_coercion"] = processed_features["enable_strict_type_coercion"] + if "strict_type_coercion" in processed_features: + val = processed_features.pop("strict_type_coercion") + if "enable_strict_type_coercion" not in processed_features: + processed_features["enable_strict_type_coercion"] = val - if "arrow_extension_types" in processed_features and "enable_arrow_extension_types" not in processed_features: - processed_features["enable_arrow_extension_types"] = processed_features["arrow_extension_types"] - if "enable_arrow_extension_types" in processed_features and "arrow_extension_types" not in processed_features: - processed_features["arrow_extension_types"] = processed_features["enable_arrow_extension_types"] + if "arrow_extension_types" in processed_features: + val = processed_features.pop("arrow_extension_types") + if "enable_arrow_extension_types" not in processed_features: + processed_features["enable_arrow_extension_types"] = val return processed_features @@ -1972,9 +1976,6 @@ def _prepare_parameter_sequence_with_casts( return tuple(result) if isinstance(parameters, tuple) else result -_prepare_batch_with_casts = prepare_batch_with_casts - - def _resolve_gizmosql_driver_connect() -> str: """Return the preferred GizmoSQL connect path based on installed packages.""" if import_optional("adbc_driver_gizmosql") is not None: diff --git a/sqlspec/adapters/adbc/driver.py b/sqlspec/adapters/adbc/driver.py index c65d69ea4..2c3dda2e5 100644 --- a/sqlspec/adapters/adbc/driver.py +++ b/sqlspec/adapters/adbc/driver.py @@ -5,6 +5,7 @@ """ import contextlib +import sys from pathlib import Path from time import perf_counter from typing import TYPE_CHECKING, Any, Literal, cast, final @@ -98,14 +99,16 @@ def _handle_exception(self, exc_type: "type[BaseException] | None", exc_val: "Ba class AdbcSelectStreamSource: """Native ADBC chunk source backed by a ``RecordBatchReader``.""" - __slots__ = ("_cursor_manager", "_driver", "_parameters", "_reader", "_sql") + __slots__ = ("_chunk_size", "_cursor_manager", "_driver", "_parameters", "_reader", "_remainder", "_sql") - def __init__(self, driver: "AdbcDriver", sql: str, parameters: Any) -> None: + def __init__(self, driver: "AdbcDriver", sql: str, parameters: Any, chunk_size: int = 1000) -> None: self._driver = driver self._sql = sql self._parameters = parameters + self._chunk_size = chunk_size self._cursor_manager: AdbcCursor | None = None self._reader: Any = None + self._remainder: list[dict[str, Any]] = [] def start(self) -> None: cursor_manager = self._driver.with_cursor(self._driver.connection) @@ -114,10 +117,23 @@ def start(self) -> None: try: cursor = cursor_manager.__enter__() handler = self._driver.handle_database_exceptions() - with handler: - execute_parameters = normalize_postgres_empty_parameters(self._driver._dialect_name, self._parameters) - cursor.execute(self._sql, parameters=execute_parameters) - reader = _fetch_record_batch(cursor) + if sys.exc_info()[1] is None: + with handler: + execute_parameters = normalize_postgres_empty_parameters( + self._driver._dialect_name, self._parameters + ) + cursor.execute(self._sql, parameters=execute_parameters) + reader = _fetch_record_batch(cursor) + else: + + def _stream_op() -> Any: + execute_parameters = normalize_postgres_empty_parameters( + self._driver._dialect_name, self._parameters + ) + cursor.execute(self._sql, parameters=execute_parameters) + return _fetch_record_batch(cursor) + + reader = self._driver._run_with_exception_handler(handler, _stream_op) self._driver._check_pending_exception(handler) except BaseException: if cursor is not None: @@ -130,14 +146,27 @@ def start(self) -> None: self._reader = iter(reader) def fetch_chunk(self) -> "list[dict[str, Any]]": + chunk_size = self._chunk_size + results: list[dict[str, Any]] = [] + if self._remainder: + if len(self._remainder) >= chunk_size: + chunk = self._remainder[:chunk_size] + self._remainder = self._remainder[chunk_size:] + return chunk + results.extend(self._remainder) + self._remainder = [] + reader = self._reader if reader is None: - return [] - while True: + return results + + while len(results) < chunk_size: try: batch = next(reader) except StopIteration: - return [] + self._reader = None + break + if bool(self._driver.driver_features.get("enable_lowercase_column_names", False)) and hasattr( batch, "schema" ): @@ -145,14 +174,22 @@ def fetch_chunk(self) -> "list[dict[str, Any]]": normalized_names = normalize_column_names(batch_names, self._driver.driver_features) if normalized_names != batch_names and hasattr(batch, "rename_columns"): batch = batch.rename_columns(normalized_names) + rows = arrow_table_to_pylist( batch, decode_arrow_extension_types=bool( self._driver.driver_features.get("enable_arrow_extension_types", True) ), ) - if rows: - return rows + needed = chunk_size - len(results) + if len(rows) <= needed: + results.extend(rows) + else: + results.extend(rows[:needed]) + self._remainder = rows[needed:] + break + + return results def close(self, error: bool = False) -> None: self._reader = None @@ -211,7 +248,7 @@ def __init__( explicit_lowercase = bool(driver_features and "enable_lowercase_column_names" in driver_features) statement_config, normalized_features = apply_driver_features(statement_config, driver_features) - if not explicit_lowercase and (dialect or self._detected_dialect) in {"oracle", "db2"}: + if not explicit_lowercase and (dialect or self._detected_dialect) in {"oracle", "db2", "snowflake"}: normalized_features["enable_lowercase_column_names"] = True super().__init__(connection=connection, statement_config=statement_config, driver_features=normalized_features) self.dialect = statement_config.dialect @@ -426,11 +463,10 @@ def dispatch_execute_script(self, cursor: "AdbcRawCursor", statement: "SQL") -> def dispatch_select_stream(self, statement: "SQL", chunk_size: int) -> "SyncRowStream[dict[str, Any]] | None": """Return a native ADBC stream backed by ``fetch_record_batch()``.""" - _ = chunk_size if not statement.returns_rows(): return None sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) - return SyncRowStream(AdbcSelectStreamSource(self, sql, prepared_parameters)) + return SyncRowStream(AdbcSelectStreamSource(self, sql, prepared_parameters, chunk_size=chunk_size)) def begin(self) -> None: """Begin database transaction. @@ -607,7 +643,7 @@ def select_to_arrow( if return_format in {"reader", "batches"}: cursor_manager = self.with_cursor(self.connection) cursor = None - reader: object | None = None + reader: Any = None with exc_handler: cursor = cursor_manager.__enter__() try: @@ -647,6 +683,19 @@ def select_to_arrow( cursor_manager.__exit__(None, None, None) msg = "ADBC did not return an Arrow reader." raise SQLSpecError(msg) + if bool(self.driver_features.get("enable_lowercase_column_names", False)) and hasattr(reader, "schema"): + raw_names = list(reader.schema.names) + normalized_names = normalize_column_names(raw_names, self.driver_features) + if normalized_names != raw_names and _pa is not None: + new_schema = _pa.schema([ + field.with_name(name) for field, name in zip(reader.schema, normalized_names, strict=True) + ]) + + def _lower_batches() -> Any: + for batch in reader: + yield batch.rename_columns(normalized_names) + + reader = _pa.RecordBatchReader.from_batches(new_schema, _lower_batches()) try: return build_arrow_result_from_reader( prepared_statement, diff --git a/tests/unit/adapters/test_adbc/test_adbc_config_normalization.py b/tests/unit/adapters/test_adbc/test_adbc_config_normalization.py index 7e8bcc50f..d1541c822 100644 --- a/tests/unit/adapters/test_adbc/test_adbc_config_normalization.py +++ b/tests/unit/adapters/test_adbc/test_adbc_config_normalization.py @@ -798,3 +798,21 @@ def _resolve_driver_name(config: AdbcConfig) -> str: def _get_connection_config_dict(config: AdbcConfig) -> dict[str, Any]: """Build the normalized connection configuration.""" return build_connection_config(config.connection_config) + + +def test_snowflake_lowercase_default() -> None: + """Snowflake dialect defaults enable_lowercase_column_names to True.""" + config = AdbcConfig(connection_config={"driver_name": "snowflake"}) + assert config.driver_features["enable_lowercase_column_names"] is True + + +def test_driver_features_alias_consolidation() -> None: + """AdbcConfig consolidates legacy alias pairs to enable_*.""" + config = AdbcConfig( + connection_config={"driver_name": "sqlite"}, + driver_features={"strict_type_coercion": True, "arrow_extension_types": False}, + ) + assert config.driver_features["enable_strict_type_coercion"] is True + assert config.driver_features["enable_arrow_extension_types"] is False + assert "strict_type_coercion" not in config.driver_features + assert "arrow_extension_types" not in config.driver_features diff --git a/tests/unit/adapters/test_adbc/test_adbc_serialization.py b/tests/unit/adapters/test_adbc/test_adbc_serialization.py index d45887f92..30d947cee 100644 --- a/tests/unit/adapters/test_adbc/test_adbc_serialization.py +++ b/tests/unit/adapters/test_adbc/test_adbc_serialization.py @@ -89,8 +89,8 @@ def test_driver_features_defaults() -> None: assert config.driver_features["json_serializer"] is to_json assert config.driver_features["enable_cast_detection"] is True - assert config.driver_features["strict_type_coercion"] is False - assert config.driver_features["arrow_extension_types"] is True + assert config.driver_features["enable_strict_type_coercion"] is False + assert config.driver_features["enable_arrow_extension_types"] is True def test_driver_features_override() -> None: @@ -103,14 +103,14 @@ def custom_serializer(data: Any) -> str: driver_features={ "json_serializer": custom_serializer, "enable_cast_detection": False, - "strict_type_coercion": True, + "enable_strict_type_coercion": True, } ) assert config.driver_features["json_serializer"] is custom_serializer assert config.driver_features["enable_cast_detection"] is False - assert config.driver_features["strict_type_coercion"] is True - assert config.driver_features["arrow_extension_types"] is True + assert config.driver_features["enable_strict_type_coercion"] is True + assert config.driver_features["enable_arrow_extension_types"] is True def test_backward_compatibility_no_serializer() -> None: diff --git a/tests/unit/adapters/test_adbc/test_core.py b/tests/unit/adapters/test_adbc/test_core.py index b8bcdd8fd..e3069d07b 100644 --- a/tests/unit/adapters/test_adbc/test_core.py +++ b/tests/unit/adapters/test_adbc/test_core.py @@ -4,6 +4,7 @@ from decimal import Decimal from types import NoneType, SimpleNamespace from typing import Any, cast +from unittest.mock import Mock from uuid import UUID import pyarrow as pa @@ -12,6 +13,8 @@ from sqlspec.adapters.adbc import AdbcDriver from sqlspec.adapters.adbc import core as adbc_core +from sqlspec.adapters.adbc import driver as adbc_driver_mod +from sqlspec.adapters.adbc._typing import AdbcSessionContext from sqlspec.adapters.adbc.core import ( build_connection_config, collect_rows, @@ -774,3 +777,98 @@ def test_type_converter_caching() -> None: conv1 = get_adbc_type_converter("postgres") conv2 = get_adbc_type_converter("postgres") assert conv1 is conv2 + + +def test_bytearray_memoryview_coercion() -> None: + """_BASE_TYPE_COERCION_MAP coerces bytearray and memoryview to bytes.""" + coercion = adbc_core._BASE_TYPE_COERCION_MAP + assert coercion[bytearray] is bytes + assert coercion[memoryview] is bytes + + +def test_core_unused_symbols() -> None: + """Dead alias _prepare_batch_with_casts is removed from adbc.core.""" + assert not hasattr(adbc_core, "_prepare_batch_with_casts") + + +def test_session_context_driver() -> None: + """AdbcSessionContext uses provided driver_class without lazy import.""" + driver_instance = Mock() + mock_driver_cls = Mock(return_value=driver_instance) + mock_conn = Mock() + ctx = AdbcSessionContext( + acquire_connection=lambda: mock_conn, + release_connection=lambda _conn, **_kw: None, + statement_config=get_statement_config("postgres"), + driver_features={}, + prepare_driver=lambda d: d, + driver_class=mock_driver_cls, + ) + with ctx as driver: + assert driver is driver_instance + mock_driver_cls.assert_called_once() + + +def test_select_stream_chunk_size() -> None: + """AdbcSelectStreamSource buffers and returns chunks of specified chunk_size.""" + batch = pa.record_batch([pa.array([1, 2, 3, 4, 5])], names=["id"]) + reader = iter([batch]) + + class _MockDriver: + driver_features: dict[str, Any] = {} + _dialect_name = "sqlite" + + def with_cursor(self, _conn: Any) -> Any: + return Mock() + + def handle_database_exceptions(self) -> Any: + return Mock() + + def _check_pending_exception(self, _handler: Any) -> None: + pass + + source = adbc_driver_mod.AdbcSelectStreamSource(cast("Any", _MockDriver()), "SELECT 1", (), chunk_size=2) + source._reader = reader + chunk1 = source.fetch_chunk() + assert len(chunk1) == 2 + chunk2 = source.fetch_chunk() + assert len(chunk2) == 2 + chunk3 = source.fetch_chunk() + assert len(chunk3) == 1 + assert source.fetch_chunk() == [] + + +def test_select_to_arrow_lowercase_columns() -> None: + """select_to_arrow normalizes reader schema to lowercase when enable_lowercase_column_names is set.""" + table = pa.table({"UPPER_COL": [1, 2]}) + mock_cursor = Mock() + mock_cursor.execute.return_value = None + mock_cursor.fetch_record_batch.return_value = table.to_reader() + mock_conn = Mock() + mock_conn.cursor.return_value = mock_cursor + + driver = AdbcDriver( + mock_conn, + statement_config=get_statement_config("oracle"), + driver_features={"enable_lowercase_column_names": True}, + ) + result = driver.select_to_arrow("SELECT UPPER_COL FROM t", return_format="reader") + reader = result.data + assert reader.schema.names == ["upper_col"] + + +def test_adbc_exception_isolation() -> None: + """AdbcSelectStreamSource routes execution through exception handler with isolation check.""" + mock_conn = Mock() + mock_cursor = Mock() + mock_conn.cursor.return_value = mock_cursor + mock_cursor.execute.return_value = None + batch = pa.record_batch([pa.array([1])], names=["id"]) + + driver = AdbcDriver(mock_conn, statement_config=get_statement_config("sqlite")) + + source = adbc_driver_mod.AdbcSelectStreamSource(driver, "SELECT 1", ()) + with pytest.MonkeyPatch.context() as mp: + mp.setattr(adbc_driver_mod, "_fetch_record_batch", lambda _c: iter([batch])) + source.start() + assert source._reader is not None From f5e400aeda885970b753131dabcfec3a0f7d9e43 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 03:57:10 +0000 Subject: [PATCH 105/174] fix(arrow_odbc): isolate ambient exceptions, map ArrowException, and unify transaction handling Isolate ambient exception state across ArrowOdbcStreamSource, select_to_arrow, and bulk_insert_arrow using _run_with_exception_handler. Map pure PyArrow exceptions directly to DataError in create_mapped_exception, and standardize exception handling across begin(), commit(), and rollback(). --- sqlspec/adapters/arrow_odbc/core.py | 7 +- sqlspec/adapters/arrow_odbc/driver.py | 109 +++++++++++++----- .../adapters/test_arrow_odbc/test_driver.py | 32 +++++ 3 files changed, 116 insertions(+), 32 deletions(-) diff --git a/sqlspec/adapters/arrow_odbc/core.py b/sqlspec/adapters/arrow_odbc/core.py index 8ade8209d..09dc9fc97 100644 --- a/sqlspec/adapters/arrow_odbc/core.py +++ b/sqlspec/adapters/arrow_odbc/core.py @@ -25,9 +25,12 @@ UniqueViolationError, ) from sqlspec.utils.config_tools import parse_odbc_connection_string +from sqlspec.utils.module_loader import import_optional from sqlspec.utils.serializers import from_json, to_json from sqlspec.utils.type_converters import build_uuid_coercions +pa: Any = import_optional("pyarrow") + if TYPE_CHECKING: from collections.abc import Callable, Mapping @@ -176,7 +179,9 @@ def resolve_dialect_from_dbms_name(dbms_name: str | None) -> str: def create_mapped_exception(error: BaseException, *, logger: Any | None = None) -> SQLSpecError: """Map an arrow-odbc exception to SQLSpec's exception hierarchy.""" - del logger + _ = logger + if pa is not None and isinstance(error, pa.ArrowException): + return DataError(f"PyArrow data error. Original error: {error}") message = str(error) if _is_sql_server_diagnostic(message): error_number = extract_native_error_number(error) diff --git a/sqlspec/adapters/arrow_odbc/driver.py b/sqlspec/adapters/arrow_odbc/driver.py index fa25a1361..8140a6159 100644 --- a/sqlspec/adapters/arrow_odbc/driver.py +++ b/sqlspec/adapters/arrow_odbc/driver.py @@ -2,6 +2,7 @@ import contextlib import re +import sys from collections.abc import Iterable, Mapping, Sequence from datetime import UTC, datetime from itertools import chain @@ -77,10 +78,20 @@ def __init__( def start(self) -> None: handler = self._driver.handle_database_exceptions() - with handler: - raw_reader = self._driver._read_arrow_batches(self._sql, self._parameters, self._chunk_size) - self._raw_reader = raw_reader - self._reader = iter(self._driver._normalize_reader(_to_pyarrow_reader(raw_reader))) + if sys.exc_info()[1] is None: + with handler: + raw_reader = self._driver._read_arrow_batches(self._sql, self._parameters, self._chunk_size) + self._raw_reader = raw_reader + self._reader = iter(self._driver._normalize_reader(_to_pyarrow_reader(raw_reader))) + else: + + def _start_op() -> Any: + raw_reader = self._driver._read_arrow_batches(self._sql, self._parameters, self._chunk_size) + self._raw_reader = raw_reader + self._reader = iter(self._driver._normalize_reader(_to_pyarrow_reader(raw_reader))) + return self._reader + + self._driver._run_with_exception_handler(handler, _start_op) self._driver._check_pending_exception(handler) def fetch_chunk(self) -> "list[dict[str, Any]]": @@ -91,8 +102,11 @@ def fetch_chunk(self) -> "list[dict[str, Any]]": batch: Any = None handler = self._driver.handle_database_exceptions() try: - with handler: - batch = next(reader) + if sys.exc_info()[1] is None: + with handler: + batch = next(reader) + else: + batch = self._driver._run_with_exception_handler(handler, next, reader) except StopIteration: return [] self._driver._check_pending_exception(handler) @@ -300,9 +314,6 @@ def begin(self) -> None: mapped = create_mapped_exception(exc) msg = f"Failed to begin transaction: {mapped}" raise type(mapped)(msg) from exc - except Exception as exc: - msg = f"Failed to begin transaction: {exc}" - raise SQLSpecError(msg) from exc self._transaction_active = True def commit(self) -> None: @@ -390,23 +401,49 @@ def select_to_arrow( arrow_result: ArrowResult | None = None exc_handler = self.handle_database_exceptions() - with exc_handler, self.with_cursor(self.connection): - reader = self._read_arrow_batches( - sql, - _odbc_parameters(prepared_parameters, naive_utc_datetimes=self._dialect == "db2"), - resolved_batch_size, - ) - if return_format in {"reader", "batches"}: - arrow_reader = self._normalize_reader(_to_pyarrow_reader(reader)) - arrow_result = build_arrow_result_from_reader( - prepared_statement, - arrow_reader, - return_format=return_format, - batch_size=resolved_batch_size, - arrow_schema=arrow_schema, - ) + with self.with_cursor(self.connection): + if sys.exc_info()[1] is None: + with exc_handler: + reader = self._read_arrow_batches( + sql, + _odbc_parameters(prepared_parameters, naive_utc_datetimes=self._dialect == "db2"), + resolved_batch_size, + ) + if return_format in {"reader", "batches"}: + arrow_reader = self._normalize_reader(_to_pyarrow_reader(reader)) + arrow_result = build_arrow_result_from_reader( + prepared_statement, + arrow_reader, + return_format=return_format, + batch_size=resolved_batch_size, + arrow_schema=arrow_schema, + ) + else: + table = self._normalize_table(_reader_to_table(reader)) else: - table = self._normalize_table(_reader_to_table(reader)) + + def _select_op() -> tuple[Any, Any]: + r = self._read_arrow_batches( + sql, + _odbc_parameters(prepared_parameters, naive_utc_datetimes=self._dialect == "db2"), + resolved_batch_size, + ) + if return_format in {"reader", "batches"}: + ar = self._normalize_reader(_to_pyarrow_reader(r)) + res = build_arrow_result_from_reader( + prepared_statement, + ar, + return_format=return_format, + batch_size=resolved_batch_size, + arrow_schema=arrow_schema, + ) + return None, res + tbl = self._normalize_table(_reader_to_table(r)) + return tbl, None + + res = self._run_with_exception_handler(exc_handler, _select_op) + if res is not None: + table, arrow_result = res self._check_pending_exception(exc_handler) if arrow_result is not None: @@ -453,15 +490,25 @@ def bulk_insert_arrow(self, target_table: str, source: Any, *, chunk_size: int | resolved_chunk_size = chunk_size or self._chunk_size() exc_handler = self.handle_database_exceptions() inserted = False - with exc_handler, self.with_cursor(self.connection): - if isinstance(source, pa.Table) and hasattr(self.connection, "from_table_to_db"): - self.connection.from_table_to_db(source=source, target=target_table, chunk_size=resolved_chunk_size) - inserted = True - else: + with self.with_cursor(self.connection): + + def _insert_op() -> bool: + if isinstance(source, pa.Table) and hasattr(self.connection, "from_table_to_db"): + self.connection.from_table_to_db(source=source, target=target_table, chunk_size=resolved_chunk_size) + return True reader = _table_to_reader(source, resolved_chunk_size) if isinstance(source, pa.Table) else source if hasattr(self.connection, "insert_into_table"): self.connection.insert_into_table(reader=reader, table=target_table, chunk_size=resolved_chunk_size) - inserted = True + return True + return False + + if sys.exc_info()[1] is None: + with exc_handler: + inserted = _insert_op() + else: + ins = self._run_with_exception_handler(exc_handler, _insert_op) + if ins is not None: + inserted = ins self._check_pending_exception(exc_handler) if inserted: diff --git a/tests/unit/adapters/test_arrow_odbc/test_driver.py b/tests/unit/adapters/test_arrow_odbc/test_driver.py index e76016fc7..843597d52 100644 --- a/tests/unit/adapters/test_arrow_odbc/test_driver.py +++ b/tests/unit/adapters/test_arrow_odbc/test_driver.py @@ -1277,3 +1277,35 @@ def insert_into_table(self, **kwargs: Any) -> None: arrow_result = ArrowResult(statement=SQL("SELECT 1"), data=table, rows_affected=2) driver.bulk_insert_arrow("items", arrow_result) assert connection.insert_calls == [("items", 65536, table)] + + +def test_ambient_exception_isolation_select_to_arrow() -> None: + """select_to_arrow executes cleanly inside an active exception block without ambient exception leakage.""" + connection = FakeConnection() + driver = ArrowOdbcDriver(cast("ArrowOdbcConnection", connection)) + + try: + raise ValueError("ambient error") + except ValueError: + result = driver.select_to_arrow("SELECT 1") + assert result.num_rows == 3 + + +def test_transaction_lifecycle_maps_generic_exception() -> None: + """begin, commit, and rollback map native exceptions with appropriate transaction failure messages.""" + connection = FakeConnection() + native_error = ArrowOdbcError.__new__(ArrowOdbcError) + native_error.args = ("native disk error",) + + def raise_native() -> None: + raise native_error + + setattr(connection, "commit", raise_native) + setattr(connection, "rollback", raise_native) + driver = ArrowOdbcDriver(cast("ArrowOdbcConnection", connection)) + + with pytest.raises(SQLSpecError, match="Failed to commit transaction"): + driver.commit() + + with pytest.raises(SQLSpecError, match="Failed to rollback transaction"): + driver.rollback() From 6cba38e9269ed9758e0eea6e21c1e9055fc3c584 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 03:57:10 +0000 Subject: [PATCH 106/174] fix(arrow_odbc): ensure atomic writes begin transactions on MSSQL and restrict extension dialects Ensure atomic store writes initiate transactions on MSSQL even under default autocommit mode, rolling back all statements on failure. Enforce supported dialects (mssql, db2) across ADK, Litestar, and events extension stores, consolidate bracket identifier helpers into core, unify atomic write context managers, and eliminate inline pyright pragmas. --- sqlspec/adapters/arrow_odbc/_typing.py | 10 +- sqlspec/adapters/arrow_odbc/adk/store.py | 86 +++---- sqlspec/adapters/arrow_odbc/core.py | 55 ++++- .../adapters/arrow_odbc/data_dictionary.py | 9 +- sqlspec/adapters/arrow_odbc/events/store.py | 33 ++- sqlspec/adapters/arrow_odbc/litestar/store.py | 45 +--- .../test_arrow_odbc/test_adk_store.py | 24 +- .../test_arrow_odbc/test_extensions.py | 13 ++ .../test_arrow_odbc/test_tsql_stores.py | 219 +++++++++++++++--- 9 files changed, 336 insertions(+), 158 deletions(-) diff --git a/sqlspec/adapters/arrow_odbc/_typing.py b/sqlspec/adapters/arrow_odbc/_typing.py index 0be32aa33..e012b40f4 100644 --- a/sqlspec/adapters/arrow_odbc/_typing.py +++ b/sqlspec/adapters/arrow_odbc/_typing.py @@ -1,15 +1,11 @@ -# pyright: reportAttributeAccessIssue=false """arrow-odbc adapter type definitions and mypyc-excluded context managers.""" from typing import TYPE_CHECKING, Any -import arrow_odbc as _arrow_odbc # pyright: ignore[reportMissingImports] +import arrow_odbc as _arrow_odbc from arrow_odbc import Error as ArrowOdbcError -from arrow_odbc import ( - TextEncoding, # pyright: ignore[reportMissingImports] - enable_odbc_connection_pooling, # pyright: ignore[reportMissingImports] -) -from arrow_odbc import connect as arrow_odbc_connect # pyright: ignore[reportMissingImports] +from arrow_odbc import TextEncoding, enable_odbc_connection_pooling +from arrow_odbc import connect as arrow_odbc_connect if TYPE_CHECKING: from collections.abc import Callable diff --git a/sqlspec/adapters/arrow_odbc/adk/store.py b/sqlspec/adapters/arrow_odbc/adk/store.py index ced882b1b..c45f1de1f 100644 --- a/sqlspec/adapters/arrow_odbc/adk/store.py +++ b/sqlspec/adapters/arrow_odbc/adk/store.py @@ -1,6 +1,5 @@ """arrow-odbc ADK stores for Google Agent Development Kit session storage.""" -import contextlib from datetime import UTC, date, datetime from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, NotRequired, cast @@ -11,15 +10,18 @@ extract_native_error_number, split_db2_name, ) +from sqlspec.adapters.arrow_odbc.core import atomic_store_write as _atomic_write +from sqlspec.adapters.arrow_odbc.core import quote_bracket_identifier as _quote_identifier +from sqlspec.adapters.arrow_odbc.core import strip_bracket_quotes as _strip_tsql_quotes from sqlspec.config import ADKConfig -from sqlspec.exceptions import SQLSpecError +from sqlspec.exceptions import ImproperConfigurationError, SQLSpecError from sqlspec.extensions.adk import BaseSyncADKStore, StoredEvent, StoredSession, normalize_session_list_options from sqlspec.extensions.adk.memory import BaseSyncADKMemoryStore, StoredMemory from sqlspec.utils.serializers import from_json, to_json from sqlspec.utils.text import split_qualified_identifier if TYPE_CHECKING: - from collections.abc import Callable, Iterator, Sequence + from collections.abc import Callable, Sequence from datetime import timedelta from sqlspec.adapters.arrow_odbc.config import ArrowOdbcConfig @@ -501,10 +503,17 @@ def search_entries( scope_filter: Literal["all", "user", "app"] = "all", embedding: "Sequence[float] | None" = None, ) -> "list[StoredMemory]": - """Search memory entries with LIKE matching.""" + """Search memory entries with LIKE matching. + + Vector embeddings are not supported in ArrowOdbcADKMemoryStore; passing + an embedding raises NotImplementedError. + """ if not self._enabled: msg = "Memory store is disabled" raise RuntimeError(msg) + if embedding is not None: + msg = "Vector embeddings are not supported in ArrowOdbcADKMemoryStore; only LIKE search is supported." + raise NotImplementedError(msg) if not query or not query.strip(): return [] effective_limit = max(0, int(limit if limit is not None else self._max_results)) @@ -945,19 +954,6 @@ def _is_table_missing(exc: BaseException) -> bool: return "invalid object name" in text or "42s02" in text or "(208)" in text -def _quote_identifier(identifier: str) -> str: - return f"[{identifier.replace(']', ']]')}]" - - -def _strip_tsql_quotes(part: str) -> str: - stripped = part.strip() - if len(stripped) >= _QUALIFIED_PARTS_MIN and stripped[0] == "[" and stripped[-1] == "]": - return stripped[1:-1].replace("]]", "]") - if len(stripped) >= _QUALIFIED_PARTS_MIN and stripped[0] == '"' and stripped[-1] == '"': - return stripped[1:-1].replace('""', '"') - return stripped - - def _split_tsql_table(table: str) -> "tuple[str, str]": parts = split_qualified_identifier(table.strip()) if len(parts) >= _QUALIFIED_PARTS_MIN: @@ -977,12 +973,6 @@ def _constraint_ref(prefix: str, table: str, suffix: str) -> str: return _quote_identifier(f"{prefix}_{bare_table}_{suffix}") -def _bare_db2_table(table: str) -> str: - parts = split_qualified_identifier(table.strip()) - raw = parts[-1] if parts else table - return _strip_tsql_quotes(raw) - - def _raise_session_not_found(session_id: str) -> None: msg = f"Session {session_id} not found during append_event_and_update_state." raise ValueError(msg) @@ -1027,29 +1017,6 @@ def _session_list_query( return sql, tuple(params) -@contextlib.contextmanager -def _atomic_write(driver: Any) -> "Iterator[None]": - """Wrap multi-statement store mutations in an explicit transaction when supported. - - Yields: - None while the atomic write block executes. - """ - in_tx = getattr(driver, "_connection_in_transaction", None) - if callable(in_tx) and in_tx() is True: - yield - return - if getattr(driver, "_connection_autocommit", True) is False: - driver.begin() - try: - yield - except BaseException: - with contextlib.suppress(Exception): - driver.rollback() - raise - else: - driver.commit() - - class _TsqlAdkSql: """SQL Server statements for the ADK stores.""" @@ -1275,8 +1242,8 @@ def events_query( return sql, tuple(params) def session_index_specs(self, session_table: str, events_table: str) -> "list[tuple[str, str, str]]": - bare_session = _bare_db2_table(session_table) - bare_events = _bare_db2_table(events_table) + bare_session = split_db2_name(session_table)[1].lower() + bare_events = split_db2_name(events_table)[1].lower() return [ (f"idx_{bare_session}_app_user", session_table, "app_name, user_id"), (f"idx_{bare_session}_update_time", session_table, "update_time"), @@ -1289,7 +1256,7 @@ def session_index_specs(self, session_table: str, events_table: str) -> "list[tu def sessions_table_ddl(self, table: str, owner_id_column_ddl: "str | None") -> str: owner_line = f", {owner_id_column_ddl}" if owner_id_column_ddl else "" - bare_table = _bare_db2_table(table) + bare_table = split_db2_name(table)[1].lower() return ( f"CREATE TABLE {table} (id VARCHAR(128) NOT NULL, app_name VARCHAR(128) NOT NULL, " f"user_id VARCHAR(128) NOT NULL{owner_line}, state {DB2_JSON_COLUMN_TYPE} NOT NULL, " @@ -1299,7 +1266,7 @@ def sessions_table_ddl(self, table: str, owner_id_column_ddl: "str | None") -> s ) def events_table_ddl(self, table: str, session_table: str) -> str: - bare_table = _bare_db2_table(table) + bare_table = split_db2_name(table)[1].lower() return ( f"CREATE TABLE {table} (id VARCHAR(128) NOT NULL, app_name VARCHAR(128) NOT NULL, " "user_id VARCHAR(128) NOT NULL, session_id VARCHAR(128) NOT NULL, " @@ -1309,7 +1276,7 @@ def events_table_ddl(self, table: str, session_table: str) -> str: ) def app_states_table_ddl(self, table: str) -> str: - bare_table = _bare_db2_table(table) + bare_table = split_db2_name(table)[1].lower() return ( f"CREATE TABLE {table} (app_name VARCHAR(128) NOT NULL, state {DB2_JSON_COLUMN_TYPE} NOT NULL, " "update_time TIMESTAMP NOT NULL DEFAULT CURRENT TIMESTAMP, " @@ -1317,7 +1284,7 @@ def app_states_table_ddl(self, table: str) -> str: ) def user_states_table_ddl(self, table: str) -> str: - bare_table = _bare_db2_table(table) + bare_table = split_db2_name(table)[1].lower() return ( f"CREATE TABLE {table} (app_name VARCHAR(128) NOT NULL, user_id VARCHAR(128) NOT NULL, " f"state {DB2_JSON_COLUMN_TYPE} NOT NULL, update_time TIMESTAMP NOT NULL DEFAULT CURRENT TIMESTAMP, " @@ -1325,7 +1292,7 @@ def user_states_table_ddl(self, table: str) -> str: ) def metadata_table_ddl(self, table: str) -> str: - bare_table = _bare_db2_table(table) + bare_table = split_db2_name(table)[1].lower() return ( f'CREATE TABLE {table} ("KEY" VARCHAR(128) NOT NULL, value VARCHAR(512) NOT NULL, ' f'CONSTRAINT pk_{bare_table}_key PRIMARY KEY ("KEY"))' @@ -1336,7 +1303,7 @@ def drop_table_sql(self, table: str) -> str: def memory_table_ddl(self, table: str, owner_id_column_ddl: "str | None") -> str: owner_line = f", {owner_id_column_ddl}" if owner_id_column_ddl else "" - bare_table = _bare_db2_table(table) + bare_table = split_db2_name(table)[1].lower() return ( f"CREATE TABLE {table} (id VARCHAR(128) NOT NULL, session_id VARCHAR(128) NOT NULL, " "app_name VARCHAR(128) NOT NULL, user_id VARCHAR(128) NOT NULL, " @@ -1348,7 +1315,7 @@ def memory_table_ddl(self, table: str, owner_id_column_ddl: "str | None") -> str ) def memory_index_specs(self, table: str) -> "list[tuple[str, str, str]]": - bare_table = _bare_db2_table(table) + bare_table = split_db2_name(table)[1].lower() return [ (f"idx_{bare_table}_scope", table, "app_name, user_id, scope, timestamp"), (f"idx_{bare_table}_session", table, "session_id, timestamp"), @@ -1374,6 +1341,11 @@ def memory_search_sql( def _adk_sql_for(config: "ArrowOdbcConfig") -> "_TsqlAdkSql | _Db2AdkSql": """Return the statement provider for the config's resolved dialect.""" - if str(config.statement_config.dialect).lower() == "db2": + dialect_attr = getattr(config.statement_config, "dialect", None) + dialect = str(dialect_attr).lower() if isinstance(dialect_attr, str) else "mssql" + if dialect == "db2": return _Db2AdkSql() - return _TsqlAdkSql() + if dialect in {"mssql", "tsql"}: + return _TsqlAdkSql() + msg = f"ArrowOdbcADKStore requires dialect 'mssql' or 'db2' (received {dialect!r})." + raise ImproperConfigurationError(msg) diff --git a/sqlspec/adapters/arrow_odbc/core.py b/sqlspec/adapters/arrow_odbc/core.py index 09dc9fc97..094c68dc7 100644 --- a/sqlspec/adapters/arrow_odbc/core.py +++ b/sqlspec/adapters/arrow_odbc/core.py @@ -1,5 +1,6 @@ """arrow-odbc adapter core helpers.""" +import contextlib import re from datetime import UTC, datetime from typing import TYPE_CHECKING, Any, Final @@ -32,7 +33,7 @@ pa: Any = import_optional("pyarrow") if TYPE_CHECKING: - from collections.abc import Callable, Mapping + from collections.abc import Callable, Iterator, Mapping from sqlspec.core import StatementConfig @@ -40,6 +41,7 @@ "DB2_INDEX_EXISTS_SQL", "DB2_TABLE_EXISTS_SQL", "apply_driver_features", + "atomic_store_write", "build_connection_config", "build_profile", "build_statement_config", @@ -49,8 +51,10 @@ "driver_profile", "extract_native_error_number", "normalize_column_names", + "quote_bracket_identifier", "resolve_dialect_from_dbms_name", "split_db2_name", + "strip_bracket_quotes", ) DB2_TABLE_EXISTS_SQL: Final[str] = ( @@ -662,3 +666,52 @@ def _format_connection_value(option_name: str, value: Any) -> str: driver_profile = build_profile() default_statement_config = build_statement_config() + + +_MIN_DELIMITER_LEN: Final = 2 + + +def quote_bracket_identifier(identifier: str) -> str: + """Quote an identifier using SQL Server / T-SQL bracket syntax.""" + return f"[{identifier.replace(']', ']]')}]" + + +def strip_bracket_quotes(part: str) -> str: + """Strip T-SQL bracket or double-quote enclosing delimiters from an identifier.""" + stripped = part.strip() + if len(stripped) >= _MIN_DELIMITER_LEN and stripped[0] == "[" and stripped[-1] == "]": + return stripped[1:-1].replace("]]", "]") + if len(stripped) >= _MIN_DELIMITER_LEN and stripped[0] == '"' and stripped[-1] == '"': + return stripped[1:-1].replace('""', '"') + return stripped + + +@contextlib.contextmanager +def atomic_store_write(driver: Any) -> "Iterator[None]": + """Wrap multi-statement store mutations in an explicit transaction when supported. + + Yields: + None while the atomic write block executes. + """ + in_tx = getattr(driver, "_connection_in_transaction", None) + if callable(in_tx) and in_tx() is True: + yield + return + + dialect = getattr(driver, "_dialect", "") + autocommit = getattr(driver, "_connection_autocommit", True) + can_begin = not (autocommit and dialect in {"db2", "oracle"}) + + if can_begin and hasattr(driver, "begin"): + driver.begin() + + try: + yield + except BaseException: + if hasattr(driver, "rollback"): + with contextlib.suppress(Exception): + driver.rollback() + raise + else: + if hasattr(driver, "commit"): + driver.commit() diff --git a/sqlspec/adapters/arrow_odbc/data_dictionary.py b/sqlspec/adapters/arrow_odbc/data_dictionary.py index 7b503551b..81f808cef 100644 --- a/sqlspec/adapters/arrow_odbc/data_dictionary.py +++ b/sqlspec/adapters/arrow_odbc/data_dictionary.py @@ -243,17 +243,12 @@ def _probe_columns(self, driver: "ArrowOdbcDriver", table: str, schema: "str | N ) probe_sql = f"SELECT * FROM {qualified} WHERE 1=0" try: - reader = driver._read_arrow_batches( # pyright: ignore[reportPrivateUsage] - probe_sql, None, 1, fetch_concurrently=False - ) + reader = driver._read_arrow_batches(probe_sql, None, 1, fetch_concurrently=False) except Exception: return [] columns: list[ColumnMetadata] = [] raw_names = [field.name for field in reader.schema] - normalized_names = normalize_column_names( - raw_names, - driver._lowercase_column_names, # pyright: ignore[reportPrivateUsage] - ) + normalized_names = normalize_column_names(raw_names, driver._lowercase_column_names) for position, (field, column_name) in enumerate(zip(reader.schema, normalized_names, strict=False), start=1): entry: ColumnMetadata = { "table_name": table, diff --git a/sqlspec/adapters/arrow_odbc/events/store.py b/sqlspec/adapters/arrow_odbc/events/store.py index dfbd4597c..0a55b5be2 100644 --- a/sqlspec/adapters/arrow_odbc/events/store.py +++ b/sqlspec/adapters/arrow_odbc/events/store.py @@ -3,7 +3,8 @@ from typing import Final from sqlspec.adapters.arrow_odbc.config import ArrowOdbcConfig -from sqlspec.adapters.arrow_odbc.core import split_db2_name +from sqlspec.adapters.arrow_odbc.core import quote_bracket_identifier, split_db2_name, strip_bracket_quotes +from sqlspec.exceptions import ImproperConfigurationError from sqlspec.extensions.events import BaseEventQueueStore from sqlspec.utils.text import split_qualified_identifier @@ -33,7 +34,14 @@ class ArrowOdbcEventQueueStore(BaseEventQueueStore[ArrowOdbcConfig]): def __init__(self, config: ArrowOdbcConfig) -> None: super().__init__(config) - self._db2 = str(config.statement_config.dialect).lower() == "db2" + dialect = str(config.statement_config.dialect).lower() + if dialect == "db2": + self._db2 = True + elif dialect in {"mssql", "tsql"}: + self._db2 = False + else: + msg = f"ArrowOdbcEventQueueStore requires dialect 'mssql' or 'db2' (received {dialect!r})." + raise ImproperConfigurationError(msg) def _column_types(self) -> tuple[str, str, str]: return "NVARCHAR(MAX)", "NVARCHAR(MAX)", "DATETIME2(6)" @@ -57,9 +65,7 @@ def _table_ddl(self) -> str: def _index_name(self) -> str: if self._db2: return super()._index_name() - parts = [ - _strip_identifier_quotes(part) for part in split_qualified_identifier(self.table_name, quote_chars='"') - ] + parts = [strip_bracket_quotes(part) for part in split_qualified_identifier(self.table_name, quote_chars='"')] flat = "_".join(part for part in parts if part) return f"idx_{flat}_channel_status" @@ -87,17 +93,8 @@ def _wrap_drop_statement(self, statement: str) -> str: return f"IF OBJECT_ID(N'{_object_name(self.table_name)}', N'U') IS NOT NULL {statement};" -def _strip_identifier_quotes(part: str) -> str: - stripped = part.strip() - if len(stripped) >= _QUALIFIED_IDENTIFIER_MIN_PARTS and stripped.startswith("[") and stripped.endswith("]"): - return stripped[1:-1].replace("]]", "]") - if len(stripped) >= _QUALIFIED_IDENTIFIER_MIN_PARTS and stripped.startswith('"') and stripped.endswith('"'): - return stripped[1:-1].replace('""', '"') - return stripped - - def _split_table_name(table_name: str) -> tuple[str, str]: - parts = [_strip_identifier_quotes(part) for part in split_qualified_identifier(table_name, quote_chars='"')] + parts = [strip_bracket_quotes(part) for part in split_qualified_identifier(table_name, quote_chars='"')] if len(parts) < _QUALIFIED_IDENTIFIER_MIN_PARTS: return "dbo", parts[0] if parts else table_name schema_name = ".".join(parts[:-1]) @@ -106,8 +103,4 @@ def _split_table_name(table_name: str) -> tuple[str, str]: def _object_name(table_name: str) -> str: schema_name, bare_table_name = _split_table_name(table_name) - return f"{_quote_bracket_identifier(schema_name)}.{_quote_bracket_identifier(bare_table_name)}" - - -def _quote_bracket_identifier(identifier: str) -> str: - return f"[{identifier.replace(']', ']]')}]" + return f"{quote_bracket_identifier(schema_name)}.{quote_bracket_identifier(bare_table_name)}" diff --git a/sqlspec/adapters/arrow_odbc/litestar/store.py b/sqlspec/adapters/arrow_odbc/litestar/store.py index f1ecdb753..10d899d3a 100644 --- a/sqlspec/adapters/arrow_odbc/litestar/store.py +++ b/sqlspec/adapters/arrow_odbc/litestar/store.py @@ -1,7 +1,6 @@ """arrow-odbc Litestar Store implementation.""" import base64 -import contextlib from datetime import UTC, datetime, timedelta from typing import TYPE_CHECKING, Any @@ -11,13 +10,14 @@ db2_timestamp_text, split_db2_name, ) +from sqlspec.adapters.arrow_odbc.core import atomic_store_write as _atomic_write +from sqlspec.adapters.arrow_odbc.core import strip_bracket_quotes as _strip_tsql_quotes +from sqlspec.exceptions import ImproperConfigurationError from sqlspec.extensions.litestar.store import BaseSQLSpecStore from sqlspec.utils.sync_tools import async_ from sqlspec.utils.text import split_qualified_identifier if TYPE_CHECKING: - from collections.abc import Iterator - from sqlspec.adapters.arrow_odbc.config import ArrowOdbcConfig __all__ = ("ArrowOdbcStore",) @@ -37,10 +37,15 @@ class ArrowOdbcStore(BaseSQLSpecStore["ArrowOdbcConfig"]): def __init__(self, config: "ArrowOdbcConfig") -> None: super().__init__(config) - dialect = str(config.statement_config.dialect).lower() - self._sql: _TsqlSessionSql | _Db2SessionSql = ( - _Db2SessionSql(self._table_name) if dialect == "db2" else _TsqlSessionSql(self._table_name) - ) + dialect_attr = getattr(config.statement_config, "dialect", None) + dialect = str(dialect_attr).lower() if isinstance(dialect_attr, str) else "mssql" + if dialect == "db2": + self._sql: _TsqlSessionSql | _Db2SessionSql = _Db2SessionSql(self._table_name) + elif dialect in {"mssql", "tsql"}: + self._sql = _TsqlSessionSql(self._table_name) + else: + msg = f"ArrowOdbcStore requires dialect 'mssql' or 'db2' (received {dialect!r})." + raise ImproperConfigurationError(msg) @property def _chunk_table_name(self) -> str: @@ -430,32 +435,6 @@ def _chunk_table_ddl(self) -> str: ) -@contextlib.contextmanager -def _atomic_write(driver: Any) -> "Iterator[None]": - if driver._connection_in_transaction(): # pyright: ignore[reportPrivateUsage] - yield - return - if not getattr(driver, "_connection_autocommit", True): - driver.begin() - try: - yield - except Exception: - with contextlib.suppress(Exception): - driver.rollback() - raise - else: - driver.commit() - - -def _strip_tsql_quotes(part: str) -> str: - stripped = part.strip() - if len(stripped) >= _QUALIFIED_PARTS_MIN and stripped.startswith("[") and stripped.endswith("]"): - return stripped[1:-1].replace("]]", "]") - if len(stripped) >= _QUALIFIED_PARTS_MIN and stripped.startswith('"') and stripped.endswith('"'): - return stripped[1:-1].replace('""', '"') - return stripped - - def _split_tsql_table_name(table_name: str) -> "tuple[str, str]": parts = [_strip_tsql_quotes(part) for part in split_qualified_identifier(table_name, quote_chars='"')] if len(parts) < _QUALIFIED_PARTS_MIN: diff --git a/tests/unit/adapters/test_arrow_odbc/test_adk_store.py b/tests/unit/adapters/test_arrow_odbc/test_adk_store.py index 9aa80febd..c1b52ef96 100644 --- a/tests/unit/adapters/test_arrow_odbc/test_adk_store.py +++ b/tests/unit/adapters/test_arrow_odbc/test_adk_store.py @@ -1,4 +1,3 @@ -# pyright: reportPrivateUsage=false """Unit tests for arrow-odbc ADK store data-dictionary existence checks.""" from typing import Any @@ -6,7 +5,8 @@ import pytest -from sqlspec.adapters.arrow_odbc.adk import ArrowOdbcADKStore +from sqlspec.adapters.arrow_odbc.adk import ArrowOdbcADKMemoryStore, ArrowOdbcADKStore +from sqlspec.exceptions import ImproperConfigurationError def _mock_config(adk_config: "dict[str, object] | None" = None) -> MagicMock: @@ -173,3 +173,23 @@ def test_arrow_odbc_list_sessions_rejects_invalid_options( store.list_sessions("app", **options) assert calls == [] + + +def test_arrow_odbc_adk_store_requires_mssql_or_db2() -> None: + """ArrowOdbcADKStore requires dialect mssql or db2.""" + config = MagicMock() + config.statement_config.dialect = "sqlite" + config.extension_config = {"adk": {}} + + with pytest.raises(ImproperConfigurationError, match="ArrowOdbcADKStore requires dialect"): + ArrowOdbcADKStore(config) + + +def test_arrow_odbc_adk_memory_store_rejects_embedding() -> None: + """ArrowOdbcADKMemoryStore rejects embeddings with NotImplementedError.""" + config = _mock_config({"enable_memory": True}) + config.statement_config.dialect = "mssql" + mem_store = ArrowOdbcADKMemoryStore(config) + + with pytest.raises(NotImplementedError, match="Vector embeddings are not supported"): + mem_store.search_entries("hello", "app", "user", embedding=[0.1, 0.2]) diff --git a/tests/unit/adapters/test_arrow_odbc/test_extensions.py b/tests/unit/adapters/test_arrow_odbc/test_extensions.py index 3f14edc9f..9e2b530e4 100644 --- a/tests/unit/adapters/test_arrow_odbc/test_extensions.py +++ b/tests/unit/adapters/test_arrow_odbc/test_extensions.py @@ -3,6 +3,8 @@ import ast from pathlib import Path +from sqlspec.adapters.arrow_odbc.core import quote_bracket_identifier, strip_bracket_quotes + def test_arrow_odbc_extensions_do_not_import_other_adapters() -> None: """Extension modules should stay inside the arrow-odbc adapter boundary.""" @@ -28,3 +30,14 @@ def test_arrow_odbc_extensions_do_not_import_other_adapters() -> None: violations.append((str(path.relative_to(adapter_root.parents[2])), module)) assert violations == [] + + +def test_bracket_quoting_and_stripping() -> None: + """Bracket identifier quoting and quote stripping round-trips correctly.""" + assert quote_bracket_identifier("users") == "[users]" + assert quote_bracket_identifier("user]name") == "[user]]name]" + assert strip_bracket_quotes("[users]") == "users" + assert strip_bracket_quotes("[user]]name]") == "user]name" + assert strip_bracket_quotes('"users"') == "users" + assert strip_bracket_quotes('""users""') == '"users"' + assert strip_bracket_quotes("plain") == "plain" diff --git a/tests/unit/adapters/test_arrow_odbc/test_tsql_stores.py b/tests/unit/adapters/test_arrow_odbc/test_tsql_stores.py index 8fc63b632..1aeb6ad89 100644 --- a/tests/unit/adapters/test_arrow_odbc/test_tsql_stores.py +++ b/tests/unit/adapters/test_arrow_odbc/test_tsql_stores.py @@ -10,6 +10,7 @@ from sqlspec.adapters.arrow_odbc import ArrowOdbcConfig from sqlspec.adapters.arrow_odbc.adk.store import ArrowOdbcADKMemoryStore, ArrowOdbcADKStore, _datetime_value +from sqlspec.adapters.arrow_odbc.core import atomic_store_write from sqlspec.adapters.arrow_odbc.events.store import ArrowOdbcEventQueueStore from sqlspec.adapters.arrow_odbc.litestar import ArrowOdbcStore from sqlspec.adapters.arrow_odbc.litestar.store import _TsqlSessionSql @@ -75,11 +76,12 @@ async def test_tsql_session_set_inline_inserts() -> None: await _session_store(connection).set("k", b"v", 60) assert normalized_calls(connection) == [ + ("BEGIN TRANSACTION", None), ("SELECT session_id FROM sess WHERE session_id = ?", ["k"]), ("INSERT INTO sess (session_id, data, expires_at) VALUES (?, ?, ?)", ["k", "dg==", ISO_TIMESTAMP]), ("DELETE FROM sess_chunks WHERE session_id = ?", ["k"]), + ("COMMIT TRANSACTION", None), ] - assert connection.commit_calls == 1 async def test_tsql_session_set_chunked_updates_existing() -> None: @@ -92,15 +94,17 @@ async def test_tsql_session_set_chunked_updates_existing() -> None: calls = normalized_calls(connection) assert [sql for sql, _ in calls] == [ + "BEGIN TRANSACTION", "SELECT session_id FROM sess WHERE session_id = ?", "UPDATE sess SET data = ?, expires_at = ?, updated_at = SYSUTCDATETIME() WHERE session_id = ?", "DELETE FROM sess_chunks WHERE session_id = ?", "INSERT INTO sess_chunks (session_id, chunk_index, data) VALUES (?, ?, ?)", "INSERT INTO sess_chunks (session_id, chunk_index, data) VALUES (?, ?, ?)", + "COMMIT TRANSACTION", ] - assert calls[1][1] == [None, None, "k"] - assert [parameters[:2] for _, parameters in calls[3:] if parameters] == [["k", "0"], ["k", "1"]] - assert [len(parameters[2] or "") for _, parameters in calls[3:] if parameters] == [3500, 500] + assert calls[2][1] == [None, None, "k"] + assert [parameters[:2] for _, parameters in calls[4:-1] if parameters] == [["k", "0"], ["k", "1"]] + assert [len(parameters[2] or "") for _, parameters in calls[4:-1] if parameters] == [3500, 500] async def test_tsql_session_get_renews_expiry() -> None: @@ -191,11 +195,13 @@ async def test_tsql_session_delete_expired() -> None: assert await _session_store(connection).delete_expired() == 2 assert normalized_calls(connection) == [ + ("BEGIN TRANSACTION", None), ( "SELECT COUNT(*) AS expired_count FROM sess WHERE expires_at IS NOT NULL AND expires_at < SYSUTCDATETIME()", None, ), ("DELETE FROM sess WHERE expires_at IS NOT NULL AND expires_at < SYSUTCDATETIME()", None), + ("COMMIT TRANSACTION", None), ] @@ -250,70 +256,137 @@ def test_tsql_event_store_create_and_drop_statements() -> None: "event_data": {"x": 1}, } _TSQL_ADK_SESSION_CALLS = [ + ("BEGIN TRANSACTION", None), ( - "INSERT INTO [dbo].[adk_session] ( id, app_name, user_id, [tenant_id], state, create_time, update_time ) VALUES (?, ?, ?, ?, ?, SYSUTCDATETIME(), SYSUTCDATETIME())", + ( + "INSERT INTO [dbo].[adk_session] ( id, app_name, user_id, [tenant_id], " + "state, create_time, update_time ) VALUES (?, ?, ?, ?, ?, SYSUTCDATETIME(), " + "SYSUTCDATETIME())" + ), ["s1", "app", "u", "7", '{"a":1}'], ), ( - "SELECT TOP 1 id, app_name, user_id, state, create_time, update_time FROM [dbo].[adk_session] WHERE app_name = ? AND user_id = ? AND id = ?", + ( + "SELECT TOP 1 id, app_name, user_id, state, create_time, update_time FROM " + "[dbo].[adk_session] WHERE app_name = ? AND user_id = ? AND id = ?" + ), ["app", "u", "s1"], ), + ("COMMIT TRANSACTION", None), ( "UPDATE [dbo].[adk_session] SET update_time = SYSUTCDATETIME() WHERE app_name = ? AND user_id = ? AND id = ?", ["app", "u", "s1"], ), ( - "SELECT TOP 1 id, app_name, user_id, state, create_time, update_time FROM [dbo].[adk_session] WHERE app_name = ? AND user_id = ? AND id = ?", + ( + "SELECT TOP 1 id, app_name, user_id, state, create_time, update_time FROM " + "[dbo].[adk_session] WHERE app_name = ? AND user_id = ? AND id = ?" + ), ["app", "u", "s1"], ), + ("BEGIN TRANSACTION", None), ( - "UPDATE [dbo].[adk_session] SET state = ?, update_time = SYSUTCDATETIME() WHERE app_name = ? AND user_id = ? AND id = ?", + ( + "UPDATE [dbo].[adk_session] SET state = ?, update_time = SYSUTCDATETIME() " + "WHERE app_name = ? AND user_id = ? AND id = ?" + ), ['{"a":2}', "app", "u", "s1"], ), + ("COMMIT TRANSACTION", None), ( - "SELECT id, app_name, user_id, state, create_time, update_time FROM [dbo].[adk_session] WHERE app_name = ? AND user_id = ? ORDER BY update_time DESC, id DESC OFFSET 10 ROWS FETCH NEXT 5 ROWS ONLY", + ( + "SELECT id, app_name, user_id, state, create_time, update_time FROM " + "[dbo].[adk_session] WHERE app_name = ? AND user_id = ? ORDER BY update_time " + "DESC, id DESC OFFSET 10 ROWS FETCH NEXT 5 ROWS ONLY" + ), ["app", "u"], ), + ("BEGIN TRANSACTION", None), ("DELETE FROM [dbo].[adk_session] WHERE app_name = ? AND user_id = ? AND id = ?", ["app", "u", "s1"]), + ("COMMIT TRANSACTION", None), + ("BEGIN TRANSACTION", None), ( - "INSERT INTO [dbo].[adk_event] ( id, app_name, user_id, session_id, invocation_id, timestamp, event_data ) VALUES (?, ?, ?, ?, ?, ?, ?)", + ( + "INSERT INTO [dbo].[adk_event] ( id, app_name, user_id, session_id, " + "invocation_id, timestamp, event_data ) VALUES (?, ?, ?, ?, ?, ?, ?)" + ), ["e1", "app", "u", "s1", "i1", "2026-01-02T03:04:05.678901", '{"x":1}'], ), + ("COMMIT TRANSACTION", None), + ("BEGIN TRANSACTION", None), ( - "UPDATE [dbo].[adk_session] SET state = ?, update_time = SYSUTCDATETIME() WHERE app_name = ? AND user_id = ? AND id = ?", + ( + "UPDATE [dbo].[adk_session] SET state = ?, update_time = SYSUTCDATETIME() " + "WHERE app_name = ? AND user_id = ? AND id = ?" + ), ['{"a":3}', "app", "u", "s1"], ), ( - "SELECT TOP 1 id, app_name, user_id, state, create_time, update_time FROM [dbo].[adk_session] WHERE app_name = ? AND user_id = ? AND id = ?", + ( + "SELECT TOP 1 id, app_name, user_id, state, create_time, update_time FROM " + "[dbo].[adk_session] WHERE app_name = ? AND user_id = ? AND id = ?" + ), ["app", "u", "s1"], ), ( - "INSERT INTO [dbo].[adk_event] ( id, app_name, user_id, session_id, invocation_id, timestamp, event_data ) VALUES (?, ?, ?, ?, ?, ?, ?)", + ( + "INSERT INTO [dbo].[adk_event] ( id, app_name, user_id, session_id, " + "invocation_id, timestamp, event_data ) VALUES (?, ?, ?, ?, ?, ?, ?)" + ), ["e1", "app", "u", "s1", "i1", "2026-01-02T03:04:05.678901", '{"x":1}'], ), ( - "MERGE INTO [dbo].[adk_app_state] WITH (HOLDLOCK) AS target USING (SELECT ? AS [app_name], ? AS state) AS source ON (target.[app_name] = source.[app_name]) WHEN MATCHED THEN UPDATE SET state = source.state, update_time = SYSUTCDATETIME() WHEN NOT MATCHED THEN INSERT ([app_name], [state], [update_time]) VALUES (source.[app_name], source.[state], SYSUTCDATETIME());", + ( + "MERGE INTO [dbo].[adk_app_state] WITH (HOLDLOCK) AS target USING (SELECT ? " + "AS [app_name], ? AS state) AS source ON (target.[app_name] = " + "source.[app_name]) WHEN MATCHED THEN UPDATE SET state = source.state, " + "update_time = SYSUTCDATETIME() WHEN NOT MATCHED THEN INSERT ([app_name], " + "[state], [update_time]) VALUES (source.[app_name], source.[state], " + "SYSUTCDATETIME());" + ), ["app", '{"g":1}'], ), ( - "MERGE INTO [dbo].[adk_user_state] WITH (HOLDLOCK) AS target USING (SELECT ? AS [app_name], ? AS [user_id], ? AS state) AS source ON (target.[app_name] = source.[app_name] AND target.[user_id] = source.[user_id]) WHEN MATCHED THEN UPDATE SET state = source.state, update_time = SYSUTCDATETIME() WHEN NOT MATCHED THEN INSERT ([app_name], [user_id], [state], [update_time]) VALUES (source.[app_name], source.[user_id], source.[state], SYSUTCDATETIME());", + ( + "MERGE INTO [dbo].[adk_user_state] WITH (HOLDLOCK) AS target USING (SELECT ? " + "AS [app_name], ? AS [user_id], ? AS state) AS source ON (target.[app_name] " + "= source.[app_name] AND target.[user_id] = source.[user_id]) WHEN MATCHED " + "THEN UPDATE SET state = source.state, update_time = SYSUTCDATETIME() WHEN " + "NOT MATCHED THEN INSERT ([app_name], [user_id], [state], [update_time]) " + "VALUES (source.[app_name], source.[user_id], source.[state], " + "SYSUTCDATETIME());" + ), ["app", "u", '{"h":2}'], ), + ("COMMIT TRANSACTION", None), ( - "SELECT TOP 4 id, app_name, user_id, session_id, invocation_id, timestamp, event_data FROM [dbo].[adk_event] WHERE app_name = ? AND user_id = ? AND session_id = ? AND timestamp > ? ORDER BY timestamp ASC", + ( + "SELECT TOP 4 id, app_name, user_id, session_id, invocation_id, timestamp, " + "event_data FROM [dbo].[adk_event] WHERE app_name = ? AND user_id = ? AND " + "session_id = ? AND timestamp > ? ORDER BY timestamp ASC" + ), ["app", "u", "s1", "2026-01-02T03:04:05.678901"], ), ( - "SELECT id, app_name, user_id, session_id, invocation_id, timestamp, event_data FROM [dbo].[adk_event] WHERE app_name = ? AND user_id = ? AND session_id = ? ORDER BY timestamp ASC", + ( + "SELECT id, app_name, user_id, session_id, invocation_id, timestamp, " + "event_data FROM [dbo].[adk_event] WHERE app_name = ? AND user_id = ? AND " + "session_id = ? ORDER BY timestamp ASC" + ), ["app", "u", "s1"], ), + ("BEGIN TRANSACTION", None), ( "SELECT COUNT(*) AS row_count FROM [dbo].[adk_event] WHERE timestamp < ? AND app_name = ?", ["2026-01-02T03:04:05.678901", "app"], ), ("DELETE FROM [dbo].[adk_event] WHERE timestamp < ? AND app_name = ?", ["2026-01-02T03:04:05.678901", "app"]), + ("COMMIT TRANSACTION", None), + ("BEGIN TRANSACTION", None), ("SELECT COUNT(*) AS row_count FROM [dbo].[adk_session] WHERE update_time < ?", ["2026-01-02T03:04:05.678901"]), ("DELETE FROM [dbo].[adk_session] WHERE update_time < ?", ["2026-01-02T03:04:05.678901"]), + ("COMMIT TRANSACTION", None), + ("BEGIN TRANSACTION", None), ( "SELECT COUNT(*) AS row_count FROM [dbo].[adk_user_state] WHERE update_time < ? AND app_name = ?", ["2026-01-02T03:04:05.678901", "app"], @@ -322,21 +395,48 @@ def test_tsql_event_store_create_and_drop_statements() -> None: "DELETE FROM [dbo].[adk_user_state] WHERE update_time < ? AND app_name = ?", ["2026-01-02T03:04:05.678901", "app"], ), + ("COMMIT TRANSACTION", None), ("SELECT TOP 1 state FROM [dbo].[adk_app_state] WHERE app_name = ?", ["app"]), ("SELECT TOP 1 state FROM [dbo].[adk_user_state] WHERE app_name = ? AND user_id = ?", ["app", "u"]), + ("BEGIN TRANSACTION", None), ( - "MERGE INTO [dbo].[adk_app_state] WITH (HOLDLOCK) AS target USING (SELECT ? AS [app_name], ? AS state) AS source ON (target.[app_name] = source.[app_name]) WHEN MATCHED THEN UPDATE SET state = source.state, update_time = SYSUTCDATETIME() WHEN NOT MATCHED THEN INSERT ([app_name], [state], [update_time]) VALUES (source.[app_name], source.[state], SYSUTCDATETIME());", + ( + "MERGE INTO [dbo].[adk_app_state] WITH (HOLDLOCK) AS target USING (SELECT ? " + "AS [app_name], ? AS state) AS source ON (target.[app_name] = " + "source.[app_name]) WHEN MATCHED THEN UPDATE SET state = source.state, " + "update_time = SYSUTCDATETIME() WHEN NOT MATCHED THEN INSERT ([app_name], " + "[state], [update_time]) VALUES (source.[app_name], source.[state], " + "SYSUTCDATETIME());" + ), ["app", '{"g":2}'], ), + ("COMMIT TRANSACTION", None), + ("BEGIN TRANSACTION", None), ( - "MERGE INTO [dbo].[adk_user_state] WITH (HOLDLOCK) AS target USING (SELECT ? AS [app_name], ? AS [user_id], ? AS state) AS source ON (target.[app_name] = source.[app_name] AND target.[user_id] = source.[user_id]) WHEN MATCHED THEN UPDATE SET state = source.state, update_time = SYSUTCDATETIME() WHEN NOT MATCHED THEN INSERT ([app_name], [user_id], [state], [update_time]) VALUES (source.[app_name], source.[user_id], source.[state], SYSUTCDATETIME());", + ( + "MERGE INTO [dbo].[adk_user_state] WITH (HOLDLOCK) AS target USING (SELECT ? " + "AS [app_name], ? AS [user_id], ? AS state) AS source ON (target.[app_name] " + "= source.[app_name] AND target.[user_id] = source.[user_id]) WHEN MATCHED " + "THEN UPDATE SET state = source.state, update_time = SYSUTCDATETIME() WHEN " + "NOT MATCHED THEN INSERT ([app_name], [user_id], [state], [update_time]) " + "VALUES (source.[app_name], source.[user_id], source.[state], " + "SYSUTCDATETIME());" + ), ["app", "u", '{"h":3}'], ), + ("COMMIT TRANSACTION", None), ("SELECT TOP 1 value FROM [dbo].[adk_internal_metadata] WHERE [key] = ?", ["k"]), + ("BEGIN TRANSACTION", None), ( - "MERGE INTO [dbo].[adk_internal_metadata] WITH (HOLDLOCK) AS target USING (SELECT ? AS [key], ? AS value) AS source ON (target.[key] = source.[key]) WHEN MATCHED THEN UPDATE SET value = source.value WHEN NOT MATCHED THEN INSERT ([key], value) VALUES (source.[key], source.value);", + ( + "MERGE INTO [dbo].[adk_internal_metadata] WITH (HOLDLOCK) AS target USING " + "(SELECT ? AS [key], ? AS value) AS source ON (target.[key] = source.[key]) " + "WHEN MATCHED THEN UPDATE SET value = source.value WHEN NOT MATCHED THEN " + "INSERT ([key], value) VALUES (source.[key], source.value);" + ), ["k", "v"], ), + ("COMMIT TRANSACTION", None), ] _TSQL_ADK_DDL = [ "CREATE TABLE [dbo].[adk_session] ( row_id UNIQUEIDENTIFIER NOT NULL CONSTRAINT [df_adk_session_row_id] DEFAULT NEWSEQUENTIALID(), id NVARCHAR(128) NOT NULL, app_name NVARCHAR(128) NOT NULL, user_id NVARCHAR(128) NOT NULL, tenant_id INTEGER, state NVARCHAR(MAX) NOT NULL, create_time DATETIME2(6) NOT NULL CONSTRAINT [df_adk_session_create_time] DEFAULT SYSUTCDATETIME(), update_time DATETIME2(6) NOT NULL CONSTRAINT [df_adk_session_update_time] DEFAULT SYSUTCDATETIME(), CONSTRAINT [pk_adk_session_row_id] PRIMARY KEY (row_id), CONSTRAINT [uq_adk_session_id] UNIQUE (id) )", @@ -415,9 +515,14 @@ def test_tsql_adk_missing_table_reads_empty(monkeypatch: pytest.MonkeyPatch) -> _TSQL_ADK_MEMORY_CALLS = [ + ("BEGIN TRANSACTION", None), ("SELECT TOP 1 id FROM [dbo].[adk_memory] WHERE event_id = ?", ["e1"]), ( - "INSERT INTO [dbo].[adk_memory] ( id, session_id, app_name, user_id, scope, event_id, author, timestamp, content_json, content_text, metadata_json, inserted_at, [tenant_id] ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + ( + "INSERT INTO [dbo].[adk_memory] ( id, session_id, app_name, user_id, scope, " + "event_id, author, timestamp, content_json, content_text, metadata_json, " + "inserted_at, [tenant_id] ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)" + ), [ "m1", "s1", @@ -434,16 +539,31 @@ def test_tsql_adk_missing_table_reads_empty(monkeypatch: pytest.MonkeyPatch) -> "7", ], ), + ("COMMIT TRANSACTION", None), ( - "SELECT id, session_id, app_name, user_id, scope, event_id, author, timestamp, content_json, content_text, metadata_json, inserted_at FROM [dbo].[adk_memory] WHERE app_name = ? AND ((scope = 'user' AND user_id = ?) OR scope = 'app') AND content_text LIKE ? ORDER BY timestamp DESC OFFSET 0 ROWS FETCH NEXT 3 ROWS ONLY", + ( + "SELECT id, session_id, app_name, user_id, scope, event_id, author, " + "timestamp, content_json, content_text, metadata_json, inserted_at FROM " + "[dbo].[adk_memory] WHERE app_name = ? AND ((scope = 'user' AND user_id = ?) " + "OR scope = 'app') AND content_text LIKE ? ORDER BY timestamp DESC OFFSET 0 " + "ROWS FETCH NEXT 3 ROWS ONLY" + ), ["app", "u", "%hello%"], ), ( - "SELECT id, session_id, app_name, user_id, scope, event_id, author, timestamp, content_json, content_text, metadata_json, inserted_at FROM [dbo].[adk_memory] WHERE app_name = ? AND scope = 'app' AND content_text LIKE ? ORDER BY timestamp DESC OFFSET 0 ROWS FETCH NEXT 20 ROWS ONLY", + ( + "SELECT id, session_id, app_name, user_id, scope, event_id, author, " + "timestamp, content_json, content_text, metadata_json, inserted_at FROM " + "[dbo].[adk_memory] WHERE app_name = ? AND scope = 'app' AND content_text " + "LIKE ? ORDER BY timestamp DESC OFFSET 0 ROWS FETCH NEXT 20 ROWS ONLY" + ), ["app", "%hello%"], ), + ("BEGIN TRANSACTION", None), ("SELECT COUNT(*) AS row_count FROM [dbo].[adk_memory] WHERE session_id = ?", ["s1"]), ("DELETE FROM [dbo].[adk_memory] WHERE session_id = ?", ["s1"]), + ("COMMIT TRANSACTION", None), + ("BEGIN TRANSACTION", None), ( "SELECT COUNT(*) AS row_count FROM [dbo].[adk_memory] WHERE inserted_at < ? AND app_name = ? AND scope = ?", [ISO_TIMESTAMP, "app", "user"], @@ -452,6 +572,7 @@ def test_tsql_adk_missing_table_reads_empty(monkeypatch: pytest.MonkeyPatch) -> "DELETE FROM [dbo].[adk_memory] WHERE inserted_at < ? AND app_name = ? AND scope = ?", [ISO_TIMESTAMP, "app", "user"], ), + ("COMMIT TRANSACTION", None), ] _TSQL_ADK_MEMORY_DDL = "CREATE TABLE [dbo].[adk_memory] ( id NVARCHAR(128) NOT NULL, session_id NVARCHAR(128) NOT NULL, app_name NVARCHAR(128) NOT NULL, user_id NVARCHAR(128) NOT NULL, scope NVARCHAR(16) NOT NULL DEFAULT 'user', event_id NVARCHAR(128) NOT NULL, author NVARCHAR(256) NULL, timestamp DATETIME2(6) NOT NULL, content_json NVARCHAR(MAX) NOT NULL, content_text NVARCHAR(MAX) NOT NULL, metadata_json NVARCHAR(MAX) NULL, inserted_at DATETIME2(6) NOT NULL, tenant_id INTEGER, CONSTRAINT [pk_adk_memory_id] PRIMARY KEY (id), CONSTRAINT [uq_adk_memory_event_id] UNIQUE (event_id) )" _TSQL_ADK_MEMORY_DROPS = ["DROP TABLE IF EXISTS [dbo].[adk_memory]"] @@ -505,10 +626,11 @@ def test_tsql_adk_memory_create_tables_and_insert_without_owner() -> None: "CREATE INDEX [idx_adk_memory_scope] ON [dbo].[adk_memory]", "CREATE INDEX [idx_adk_memory_session] ON [dbo].[adk_memory]", ] - assert statements[-1] == ( + assert statements[-2] == ( "INSERT INTO [dbo].[adk_memory] ( id, session_id, app_name, user_id, scope, event_id, author, timestamp, " "content_json, content_text, metadata_json, inserted_at ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)" ) + assert statements[-1] == "COMMIT TRANSACTION" async def test_tsql_session_get_parses_string_expiry_and_deletes_expired_in_same_session( @@ -556,7 +678,7 @@ def test_tsql_adk_append_event_and_update_state_rolls_back_when_session_missing( with pytest.raises(ValueError, match="Session s1 not found"): store.append_event_and_update_state(_ADK_EVENT, "app", "u", "s1", {"a": 3}) - assert connection.rollback_calls == 1 + assert any(sql == "ROLLBACK TRANSACTION" for sql, _ in normalized_calls(connection)) assert connection.commit_calls == 0 @@ -586,22 +708,22 @@ def test_schema_qualified_tables_across_tsql_stores() -> None: extension_config={"events": {"queue_table": "custom.app_events"}}, ) adk_store = ArrowOdbcADKStore(config) - adk_store._session_table = "custom.adk_session" # pyright: ignore[reportPrivateUsage] + adk_store._session_table = "custom.adk_session" memory_store = ArrowOdbcADKMemoryStore(config) - memory_store._memory_table = "custom.adk_memory" # pyright: ignore[reportPrivateUsage] + memory_store._memory_table = "custom.adk_memory" litestar_store = ArrowOdbcStore(config) - litestar_store._sql = _TsqlSessionSql("custom.sess") # pyright: ignore[reportPrivateUsage] + litestar_store._sql = _TsqlSessionSql("custom.sess") event_store = ArrowOdbcEventQueueStore(config) - session_ddl = adk_store._sessions_table_ddl() # pyright: ignore[reportPrivateUsage] + session_ddl = adk_store._sessions_table_ddl() assert "[custom].[adk_session]" in session_ddl assert "[pk_adk_session_row_id]" in session_ddl - memory_ddl = memory_store._memory_table_ddl() # pyright: ignore[reportPrivateUsage] + memory_ddl = memory_store._memory_table_ddl() assert "[custom].[adk_memory]" in memory_ddl assert "[pk_adk_memory_id]" in memory_ddl - litestar_ddl = litestar_store._table_ddl() # pyright: ignore[reportPrivateUsage] + litestar_ddl = litestar_store._table_ddl() assert "SCHEMA_ID(N'custom')" in litestar_ddl assert "IX_sess_expires_at" in litestar_ddl @@ -616,3 +738,38 @@ def test_datetime_value_accepts_date_and_rejects_invalid_types() -> None: assert _datetime_value(date(2026, 5, 10)) == datetime(2026, 5, 10, 0, 0, tzinfo=UTC) with pytest.raises(TypeError, match="Unsupported datetime value"): _datetime_value(12345) + + +def test_atomic_write_rolls_back_on_failure_with_default_autocommit() -> None: + """_atomic_write rolls back on failure even with default autocommit=True.""" + + class MockDriver: + def __init__(self) -> None: + self._connection_autocommit = True + self._dialect = "mssql" + self.begun = False + self.rolled_back = False + self.committed = False + + def _connection_in_transaction(self) -> bool: + return self.begun + + def begin(self) -> None: + self.begun = True + + def rollback(self) -> None: + self.rolled_back = True + self.begun = False + + def commit(self) -> None: + self.committed = True + self.begun = False + + driver = MockDriver() + with pytest.raises(RuntimeError, match="second statement failed"): + with atomic_store_write(driver): + assert driver.begun is True + raise RuntimeError("second statement failed") + + assert driver.rolled_back is True + assert driver.committed is False From bd65e79bf277f39e424ae7c20188059d8646d615 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 03:57:10 +0000 Subject: [PATCH 107/174] fix(arrow_odbc): report supported metadata capabilities for all catalog query domains Check for available catalog queries via _has_query for all standard metadata domains instead of unconditionally marking them unsupported. Allow dialects with catalog query packs (such as Db2 and MSSQL) to report supported native catalog capability across views, schemas, objects, and routines. --- sqlspec/adapters/arrow_odbc/data_dictionary.py | 2 +- .../adapters/test_arrow_odbc/test_db2_metadata.py | 14 ++++++++++++++ 2 files changed, 15 insertions(+), 1 deletion(-) diff --git a/sqlspec/adapters/arrow_odbc/data_dictionary.py b/sqlspec/adapters/arrow_odbc/data_dictionary.py index 81f808cef..021fc1914 100644 --- a/sqlspec/adapters/arrow_odbc/data_dictionary.py +++ b/sqlspec/adapters/arrow_odbc/data_dictionary.py @@ -294,7 +294,7 @@ def _capability_for_domain(self, domain: str) -> MetadataCapability: source=MetadataSource.DRIVER_METADATA, warnings=(_ODBC_COLUMN_PARTIAL_WARNING,), ) - if domain in {"tables", "indexes", "foreign_keys"} and self._has_query(domain, "by_schema"): + if self._has_query(domain, "by_schema"): return MetadataCapability( domain=domain, support=MetadataSupport.SUPPORTED, diff --git a/tests/unit/adapters/test_arrow_odbc/test_db2_metadata.py b/tests/unit/adapters/test_arrow_odbc/test_db2_metadata.py index 35cc49781..8500b0a5e 100644 --- a/tests/unit/adapters/test_arrow_odbc/test_db2_metadata.py +++ b/tests/unit/adapters/test_arrow_odbc/test_db2_metadata.py @@ -10,6 +10,7 @@ from sqlspec.adapters.arrow_odbc import ArrowOdbcConfig, ArrowOdbcDriver from sqlspec.adapters.arrow_odbc.core import create_mapped_exception from sqlspec.adapters.arrow_odbc.type_converter import odbc_type_to_arrow +from sqlspec.data_dictionary import MetadataSource, MetadataSupport from sqlspec.exceptions import CheckViolationError, DeadlockError, QueryTimeoutError from tests.unit.adapters.test_arrow_odbc._db2_fakes import ( DB2_CONNECTION_STRING, @@ -133,3 +134,16 @@ def test_result_names_kept_when_lowercasing_is_off(driver_factory: Any) -> None: assert names == { key: ["ID", "CUSTOMER_NAME", "MixedCase"] for key in ("execute", "stream", "table", "reader", "reader_batches") } + + +def test_db2_metadata_capabilities_reports_supported_for_catalog_domains() -> None: + """Capabilities for views and schemas report SUPPORTED on dialects with bundled SQL packs.""" + connection = FakeArrowOdbcConnection() + driver = _db2_driver(connection) + caps = driver.data_dictionary.get_metadata_capabilities(driver) + schema_cap = caps.get("schemas") + views_cap = caps.get("views") + + assert schema_cap is not None and schema_cap.support == MetadataSupport.SUPPORTED + assert views_cap is not None and views_cap.support == MetadataSupport.SUPPORTED + assert schema_cap.source == MetadataSource.CATALOG From c34bfd9cc86478c34cea396ac5df78f1c7677845 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 03:57:10 +0000 Subject: [PATCH 108/174] perf(arrow_odbc): eliminate kwargs dict allocation in batch reading and avoid generator in dialect matching Pass batch reading keyword arguments directly into self.connection.read_arrow_batches to eliminate temporary kwargs dict allocation on hot paths. Replace generator expression in _match_dialect_from_dbms_name with explicit nested loops to avoid MyPyC generator struct allocation. --- sqlspec/adapters/arrow_odbc/core.py | 5 ++- sqlspec/adapters/arrow_odbc/driver.py | 61 +++++++++++++++++++++------ 2 files changed, 50 insertions(+), 16 deletions(-) diff --git a/sqlspec/adapters/arrow_odbc/core.py b/sqlspec/adapters/arrow_odbc/core.py index 094c68dc7..79236ec8d 100644 --- a/sqlspec/adapters/arrow_odbc/core.py +++ b/sqlspec/adapters/arrow_odbc/core.py @@ -449,8 +449,9 @@ def _match_dialect_from_dbms_name(dbms_name: str | None) -> str | None: return None lowered = driver_name.lower() for dialect, patterns in _DIALECT_PATTERNS: - if any(pattern in lowered for pattern in patterns): - return dialect + for pattern in patterns: + if pattern in lowered: + return dialect return None diff --git a/sqlspec/adapters/arrow_odbc/driver.py b/sqlspec/adapters/arrow_odbc/driver.py index 8140a6159..256fc7830 100644 --- a/sqlspec/adapters/arrow_odbc/driver.py +++ b/sqlspec/adapters/arrow_odbc/driver.py @@ -611,20 +611,53 @@ def _read_arrow_batches( *, fetch_concurrently: bool | None = None, ) -> Any: - kwargs: dict[str, Any] = { - "query": sql, - "batch_size": batch_size, - "parameters": parameters, - "max_bytes_per_batch": self._max_batch_bytes, - "max_text_size": self._max_text_size_val, - "max_binary_size": self._max_binary_size_val, - "fetch_concurrently": self._use_concurrent_fetch if fetch_concurrently is None else fetch_concurrently, - } - if self._query_timeout_sec_val is not None: - kwargs["query_timeout_sec"] = self._query_timeout_sec_val - if self._payload_text_encoding is not None: - kwargs["payload_text_encoding"] = self._payload_text_encoding - return self.connection.read_arrow_batches(**kwargs) + use_concurrent = self._use_concurrent_fetch if fetch_concurrently is None else fetch_concurrently + timeout = self._query_timeout_sec_val + encoding = self._payload_text_encoding + + if timeout is not None and encoding is not None: + return self.connection.read_arrow_batches( + query=sql, + batch_size=batch_size, + parameters=parameters, + max_bytes_per_batch=self._max_batch_bytes, + max_text_size=self._max_text_size_val, + max_binary_size=self._max_binary_size_val, + fetch_concurrently=use_concurrent, + query_timeout_sec=timeout, + payload_text_encoding=encoding, + ) + if timeout is not None: + return self.connection.read_arrow_batches( + query=sql, + batch_size=batch_size, + parameters=parameters, + max_bytes_per_batch=self._max_batch_bytes, + max_text_size=self._max_text_size_val, + max_binary_size=self._max_binary_size_val, + fetch_concurrently=use_concurrent, + query_timeout_sec=timeout, + ) + if encoding is not None: + return self.connection.read_arrow_batches( + query=sql, + batch_size=batch_size, + parameters=parameters, + max_bytes_per_batch=self._max_batch_bytes, + max_text_size=self._max_text_size_val, + max_binary_size=self._max_binary_size_val, + fetch_concurrently=use_concurrent, + payload_text_encoding=encoding, + ) + return self.connection.read_arrow_batches( + query=sql, + batch_size=batch_size, + parameters=parameters, + max_bytes_per_batch=self._max_batch_bytes, + max_text_size=self._max_text_size_val, + max_binary_size=self._max_binary_size_val, + fetch_concurrently=use_concurrent, + ) def _chunk_size(self) -> int: return self._chunk_size_val From ff1c66bd490f4e8fc3051354dc1f103af619dedc Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 03:57:10 +0000 Subject: [PATCH 109/174] refactor(arrow_odbc): align chunk_size docstring with default implementation Update ArrowOdbcDriverFeatures chunk_size docstring to accurately state default of 65536. --- sqlspec/adapters/arrow_odbc/config.py | 2 +- tests/unit/adapters/test_arrow_odbc/test_config.py | 6 ++++++ 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/sqlspec/adapters/arrow_odbc/config.py b/sqlspec/adapters/arrow_odbc/config.py index 50bad63dd..4359daa51 100644 --- a/sqlspec/adapters/arrow_odbc/config.py +++ b/sqlspec/adapters/arrow_odbc/config.py @@ -86,7 +86,7 @@ class ArrowOdbcDriverFeatures(TypedDict): """``arrow-odbc`` driver feature flags controlling Arrow batch sizing, dialect hints, and codecs.""" chunk_size: NotRequired[int] - """Maximum rows per Arrow ``RecordBatch`` fetched from ODBC cursors (defaults to ``65535``).""" + """Maximum rows per Arrow ``RecordBatch`` fetched from ODBC cursors (defaults to ``65536``).""" max_bytes_per_batch: NotRequired[int] """Upper memory budget in bytes per Arrow ``RecordBatch`` buffer.""" max_text_size: NotRequired[int] diff --git a/tests/unit/adapters/test_arrow_odbc/test_config.py b/tests/unit/adapters/test_arrow_odbc/test_config.py index 983d2b8d9..730c9527c 100644 --- a/tests/unit/adapters/test_arrow_odbc/test_config.py +++ b/tests/unit/adapters/test_arrow_odbc/test_config.py @@ -319,3 +319,9 @@ def broken_pooling() -> None: with pytest.raises(ImproperConfigurationError, match="Could not configure arrow-odbc connection"): config.create_connection() + + +def test_driver_features_chunk_size_default() -> None: + """ArrowOdbcConfig defaults chunk_size to 65536 matching documentation.""" + config = ArrowOdbcConfig() + assert config.driver_features["chunk_size"] == 65_536 From 030bc0fba24e35a766bf5e78620b12bcf42d201d Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 03:57:10 +0000 Subject: [PATCH 110/174] fix(bigquery): map RetryError, catch GoogleAPIError, and handle empty iterable schema Map google.api_core.exceptions.RetryError to QueryTimeoutError in create_mapped_exception, and widen BigQueryExceptionHandler and ADK query runner to catch GoogleAPIError as well as GoogleCloudError. Convert BigQuery schema field definitions to typed PyArrow schemas when constructing RecordBatchReader from empty iterables. --- sqlspec/adapters/bigquery/_typing.py | 2 ++ sqlspec/adapters/bigquery/adk/store.py | 4 +-- sqlspec/adapters/bigquery/core.py | 6 ++++- sqlspec/adapters/bigquery/driver.py | 27 +++++++++++++++++-- .../unit/adapters/test_bigquery/test_core.py | 19 ++++++++++++- .../test_bigquery/test_storage_write_api.py | 14 +++++++++- 6 files changed, 65 insertions(+), 7 deletions(-) diff --git a/sqlspec/adapters/bigquery/_typing.py b/sqlspec/adapters/bigquery/_typing.py index c4c1ddbc5..ed8b30b49 100644 --- a/sqlspec/adapters/bigquery/_typing.py +++ b/sqlspec/adapters/bigquery/_typing.py @@ -12,6 +12,7 @@ from google.api_core import exceptions as bigquery_exceptions from google.api_core.client_info import ClientInfo as BigQueryClientInfo from google.api_core.client_options import ClientOptions as BigQueryClientOptions +from google.api_core.exceptions import GoogleAPIError from google.api_core.retry import Retry as BigQueryRetry from google.auth.credentials import Credentials as BigQueryCredentials from google.cloud.bigquery import ArrayQueryParameter, Client, ScalarQueryParameter, StructQueryParameter @@ -58,6 +59,7 @@ "BigQuerySessionContext", "BigQueryStorageWriteModule", "BigQueryStorageWriteTypes", + "GoogleAPIError", "GoogleCloudError", "bigquery_exceptions", "bigquery_module", diff --git a/sqlspec/adapters/bigquery/adk/store.py b/sqlspec/adapters/bigquery/adk/store.py index 254e53acb..f914d58a6 100644 --- a/sqlspec/adapters/bigquery/adk/store.py +++ b/sqlspec/adapters/bigquery/adk/store.py @@ -15,7 +15,7 @@ from typing_extensions import TypedDict -from sqlspec.adapters.bigquery._typing import GoogleCloudError +from sqlspec.adapters.bigquery._typing import GoogleAPIError, GoogleCloudError from sqlspec.adapters.bigquery._typing import bigquery_module as bigquery from sqlspec.adapters.bigquery.config import BigQueryConfig from sqlspec.adapters.bigquery.core import build_dml_rowcount, create_mapped_exception @@ -238,7 +238,7 @@ def _run_query(self, sql: str, parameters: "Iterable[Any] | None" = None) -> "li try: job = client.query(sql, job_config=job_config) rows = [dict(row) for row in job.result()] - except GoogleCloudError as exc: + except (GoogleCloudError, GoogleAPIError) as exc: raise create_mapped_exception(exc) from exc self._last_dml_affected_rows = build_dml_rowcount(job, 0) return rows diff --git a/sqlspec/adapters/bigquery/core.py b/sqlspec/adapters/bigquery/core.py index 2f632a681..4b7148bd6 100644 --- a/sqlspec/adapters/bigquery/core.py +++ b/sqlspec/adapters/bigquery/core.py @@ -37,6 +37,7 @@ OperationalError, OperationCancelledError, PermissionDeniedError, + QueryTimeoutError, SQLParsingError, SQLSpecError, StorageCapabilityError, @@ -751,7 +752,7 @@ def create_mapped_exception(error: Any, *, logger: Any | None = None) -> SQLSpec Returns: A SQLSpec exception that wraps the original error """ - del logger + _ = logger try: status_code = error.code except AttributeError: @@ -767,6 +768,9 @@ def create_mapped_exception(error: Any, *, logger: Any | None = None) -> SQLSpec if isinstance(error, api_exceptions.Cancelled): return _create_bigquery_error(error, status_code, OperationCancelledError, "query cancelled") + if isinstance(error, api_exceptions.RetryError): + return _create_bigquery_error(error, status_code, QueryTimeoutError, "query retry timeout expired") + if error_class := _classify_timeout_or_cancellation(error_msg): return _create_bigquery_error(error, status_code, error_class, "query terminated") diff --git a/sqlspec/adapters/bigquery/driver.py b/sqlspec/adapters/bigquery/driver.py index f89bdfadb..6b49f09c0 100644 --- a/sqlspec/adapters/bigquery/driver.py +++ b/sqlspec/adapters/bigquery/driver.py @@ -22,6 +22,7 @@ BigQuerySessionContext, BigQueryStorageWriteModule, BigQueryStorageWriteTypes, + GoogleAPIError, GoogleCloudError, ) from sqlspec.adapters.bigquery.core import ( @@ -99,7 +100,7 @@ class BigQueryExceptionHandler(BaseSyncExceptionHandler): def _handle_exception(self, exc_type: "type[BaseException] | None", exc_val: "BaseException") -> bool: if exc_type is None: return False - if issubclass(exc_type, GoogleCloudError): + if issubclass(exc_type, (GoogleCloudError, GoogleAPIError)): self.pending_exception = create_mapped_exception(exc_val) return True return False @@ -1093,7 +1094,29 @@ def _bigquery_arrow_reader_from_iterable( if isinstance(schema, _pa.Schema): empty_schema = schema elif isinstance(schema, (list, tuple)): - empty_schema = _pa.schema([_pa.field(getattr(field, "name", str(field)), _pa.null()) for field in schema]) + type_map = { + "INTEGER": _pa.int64(), + "INT64": _pa.int64(), + "FLOAT": _pa.float64(), + "FLOAT64": _pa.float64(), + "BOOLEAN": _pa.bool_(), + "BOOL": _pa.bool_(), + "STRING": _pa.string(), + "BYTES": _pa.binary(), + "DATE": _pa.date32(), + "DATETIME": _pa.timestamp("us"), + "TIMESTAMP": _pa.timestamp("us", tz="UTC"), + "NUMERIC": _pa.decimal128(38, 9), + "BIGNUMERIC": _pa.decimal256(76, 38), + } + fields = [ + _pa.field( + getattr(field, "name", str(field)), + type_map.get(str(getattr(field, "field_type", "")).upper(), _pa.null()), + ) + for field in schema + ] + empty_schema = _pa.schema(fields) else: empty_schema = _pa.schema([]) return cast("ArrowRecordBatchReader", _pa.RecordBatchReader.from_batches(empty_schema, ())) diff --git a/tests/unit/adapters/test_bigquery/test_core.py b/tests/unit/adapters/test_bigquery/test_core.py index 7383e5aac..2c905479d 100644 --- a/tests/unit/adapters/test_bigquery/test_core.py +++ b/tests/unit/adapters/test_bigquery/test_core.py @@ -9,6 +9,7 @@ from unittest.mock import patch import pytest +from google.api_core.exceptions import RetryError from google.cloud import bigquery from google.cloud.bigquery import QueryJobConfig from google.cloud.exceptions import Forbidden, GoogleCloudError @@ -33,7 +34,7 @@ try_bulk_insert, ) from sqlspec.adapters.bigquery.data_dictionary import BigQueryDataDictionary -from sqlspec.adapters.bigquery.driver import BigQueryDriver +from sqlspec.adapters.bigquery.driver import BigQueryDriver, BigQueryExceptionHandler from sqlspec.adapters.bigquery.events.store import BigQueryEventQueueStore from sqlspec.adapters.bigquery.litestar.store import BigQueryStore from sqlspec.core import ParameterStyle, TypedParameter @@ -920,3 +921,19 @@ def _cloud_error(message: str, *, code: int, errors: list[dict[str, str]] | None err.code = code err._errors = errors or [] return cast("GoogleCloudError", err) + + +def test_retry_error_mapped_to_query_timeout() -> None: + """RetryError maps to QueryTimeoutError in BigQuery create_mapped_exception.""" + err = RetryError("Deadline exceeded waiting for query", None) + mapped = create_mapped_exception(err) + assert isinstance(mapped, QueryTimeoutError) + + +def test_exception_handler_catches_google_api_error() -> None: + """BigQueryExceptionHandler catches GoogleAPIError and maps it.""" + handler = BigQueryExceptionHandler() + err = RetryError("Deadline exceeded", None) + with handler: + raise err + assert isinstance(handler.pending_exception, QueryTimeoutError) diff --git a/tests/unit/adapters/test_bigquery/test_storage_write_api.py b/tests/unit/adapters/test_bigquery/test_storage_write_api.py index 03a01c999..57f1d017b 100644 --- a/tests/unit/adapters/test_bigquery/test_storage_write_api.py +++ b/tests/unit/adapters/test_bigquery/test_storage_write_api.py @@ -6,12 +6,13 @@ import pyarrow as pa import pytest +from google.cloud.bigquery import SchemaField from google.cloud.bigquery_storage_v1 import types import sqlspec.adapters.bigquery.config as bigquery_config from sqlspec.adapters.bigquery import BigQueryConfig from sqlspec.adapters.bigquery.core import build_arrow_write_stream_payload -from sqlspec.adapters.bigquery.driver import BigQueryDriver +from sqlspec.adapters.bigquery.driver import BigQueryDriver, _bigquery_arrow_reader_from_iterable from sqlspec.exceptions import ImproperConfigurationError, StorageOperationFailedError CAPABILITIES = { @@ -330,3 +331,14 @@ def fail_append(self: Any, **kwargs: Any) -> Any: client = _FakeWriteClient.instances[0] assert client.finalize_calls == ["projects/proj/datasets/dataset/tables/table/streams/s1"] assert client.commit_calls == [] + + +def test_arrow_reader_from_iterable_empty_with_schema() -> None: + """_bigquery_arrow_reader_from_iterable builds typed schema from field list when empty.""" + fields = [SchemaField("id", "INTEGER"), SchemaField("name", "STRING")] + reader = _bigquery_arrow_reader_from_iterable([], schema=fields) + + assert reader.schema.names == ["id", "name"] + assert reader.schema.field("id").type == pa.int64() + assert reader.schema.field("name").type == pa.string() + assert reader.read_all().num_rows == 0 From 30a1a3b3f7d857c6349e6467b8ba8978525fdb41 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 03:57:10 +0000 Subject: [PATCH 111/174] refactor(bigquery): consolidate extension partition clauses and remove unused length deletion Consolidate _bigquery_partition_clauses and _positive_int_or_none into bigquery.core and reuse across events and litestar stores. Replace del length with unused assignment in synchronous _string_type helper. --- sqlspec/adapters/bigquery/core.py | 25 ++++++++++++++++++++ sqlspec/adapters/bigquery/events/store.py | 26 ++------------------- sqlspec/adapters/bigquery/litestar/store.py | 24 +------------------ 3 files changed, 28 insertions(+), 47 deletions(-) diff --git a/sqlspec/adapters/bigquery/core.py b/sqlspec/adapters/bigquery/core.py index 4b7148bd6..206d0146b 100644 --- a/sqlspec/adapters/bigquery/core.py +++ b/sqlspec/adapters/bigquery/core.py @@ -1382,3 +1382,28 @@ def _create_bigquery_error( driver_profile = build_profile() default_statement_config = build_statement_config() + + +def _positive_int_or_none(value: object, key: str) -> "int | None": + """Validate that an extension config setting is a positive integer or None.""" + if value is None: + return None + if not isinstance(value, int) or isinstance(value, bool) or value <= 0: + msg = f"'{key}' must be a positive integer" + raise ValueError(msg) + return value + + +def _bigquery_partition_clauses( + column: str, *, enabled: bool, expiration_days: "int | None", require_filter: bool +) -> "tuple[str, str]": + """Build BigQuery PARTITION BY and OPTIONS clauses for table DDL.""" + if not enabled: + return "", "" + options: list[str] = [] + if require_filter: + options.append("require_partition_filter = TRUE") + if expiration_days is not None: + options.append(f"partition_expiration_days = {expiration_days}") + options_clause = f" OPTIONS({', '.join(options)})" if options else "" + return f" PARTITION BY DATE({column})", options_clause diff --git a/sqlspec/adapters/bigquery/events/store.py b/sqlspec/adapters/bigquery/events/store.py index dc9fd9b3b..9655667c0 100644 --- a/sqlspec/adapters/bigquery/events/store.py +++ b/sqlspec/adapters/bigquery/events/store.py @@ -15,6 +15,7 @@ from typing import NotRequired from sqlspec.adapters.bigquery.config import BigQueryConfig +from sqlspec.adapters.bigquery.core import _bigquery_partition_clauses, _positive_int_or_none from sqlspec.config import EventsConfig from sqlspec.extensions.events import BaseEventQueueStore @@ -88,7 +89,7 @@ def _column_types(self) -> "tuple[str, str, str]": def _string_type(self, length: int) -> str: """Return BigQuery STRING type (length is ignored).""" - del length + _ = length return "STRING" def _integer_type(self) -> str: @@ -153,26 +154,3 @@ def _resolve_events_dataset_id(config: BigQueryConfig) -> "str | None": if resolved_id: return str(resolved_id) return str(default_dataset) - - -def _positive_int_or_none(value: object, key: str) -> "int | None": - if value is None: - return None - if not isinstance(value, int) or isinstance(value, bool) or value <= 0: - msg = f"extension_config['events']['{key}'] must be a positive integer" - raise ValueError(msg) - return value - - -def _bigquery_partition_clauses( - column: str, *, enabled: bool, expiration_days: "int | None", require_filter: bool -) -> "tuple[str, str]": - if not enabled: - return "", "" - options: list[str] = [] - if require_filter: - options.append("require_partition_filter = TRUE") - if expiration_days is not None: - options.append(f"partition_expiration_days = {expiration_days}") - options_clause = f" OPTIONS({', '.join(options)})" if options else "" - return f" PARTITION BY DATE({column})", options_clause diff --git a/sqlspec/adapters/bigquery/litestar/store.py b/sqlspec/adapters/bigquery/litestar/store.py index c26aca41f..8e9b8cd31 100644 --- a/sqlspec/adapters/bigquery/litestar/store.py +++ b/sqlspec/adapters/bigquery/litestar/store.py @@ -3,6 +3,7 @@ from datetime import UTC, datetime, timedelta from typing import TYPE_CHECKING, Any, NotRequired, cast +from sqlspec.adapters.bigquery.core import _bigquery_partition_clauses, _positive_int_or_none from sqlspec.config import LitestarConfig from sqlspec.extensions.litestar.store import BaseSQLSpecStore from sqlspec.utils.sync_tools import async_ @@ -362,26 +363,3 @@ def _resolve_litestar_dataset_id(config: "BigQueryConfig") -> "str | None": if resolved_id: return str(resolved_id) return str(default_dataset) - - -def _positive_int_or_none(value: object, key: str) -> "int | None": - if value is None: - return None - if not isinstance(value, int) or isinstance(value, bool) or value <= 0: - msg = f"extension_config['litestar']['{key}'] must be a positive integer" - raise ValueError(msg) - return value - - -def _bigquery_partition_clauses( - column: str, *, enabled: bool, expiration_days: "int | None", require_filter: bool -) -> "tuple[str, str]": - if not enabled: - return "", "" - options: list[str] = [] - if require_filter: - options.append("require_partition_filter = TRUE") - if expiration_days is not None: - options.append(f"partition_expiration_days = {expiration_days}") - options_clause = f"\n OPTIONS({', '.join(options)})" if options else "" - return f"PARTITION BY DATE({column})\n ", options_clause From 40d46a9f262d96ad1ecffe47481ad2c9d50023d4 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 03:57:10 +0000 Subject: [PATCH 112/174] perf(bigquery): isolate ambient exceptions and bind none_type constant at module top-level Check sys.exc_info()[1] across auxiliary BigQuery driver methods (begin, commit, rollback, dry_run, select_to_arrow) to guarantee ambient exception isolation. Bind none_type constant once at module top-level in bigquery.core to eliminate repeated dynamic type(None) lookups. --- sqlspec/adapters/bigquery/core.py | 5 ++- sqlspec/adapters/bigquery/driver.py | 66 +++++++++++++++++++++-------- 2 files changed, 52 insertions(+), 19 deletions(-) diff --git a/sqlspec/adapters/bigquery/core.py b/sqlspec/adapters/bigquery/core.py index 206d0146b..98f32088d 100644 --- a/sqlspec/adapters/bigquery/core.py +++ b/sqlspec/adapters/bigquery/core.py @@ -7,7 +7,7 @@ import ipaddress from decimal import Decimal from enum import Enum -from typing import TYPE_CHECKING, Any, Protocol, cast, get_args +from typing import TYPE_CHECKING, Any, Final, Protocol, cast, get_args from urllib.parse import urlparse import sqlglot @@ -102,6 +102,7 @@ HTTP_SERVER_ERROR = 500 COLUMN_CACHE_MAX_SIZE = 256 _CONNECTION_NAME_PARTS = 3 +none_type: Final = type(None) DEFAULT_REQUEST_TIMEOUT = 120.0 @@ -691,7 +692,7 @@ def build_profile() -> "DriverParameterProfile": Decimal: _identity, dict: _identity, list: _identity, - type(None): _return_none, + none_type: _return_none, **build_uuid_coercions(), }, default_ast_transformer=build_null_pruning_transform(dialect="bigquery"), diff --git a/sqlspec/adapters/bigquery/driver.py b/sqlspec/adapters/bigquery/driver.py index 6b49f09c0..d1fcc32bc 100644 --- a/sqlspec/adapters/bigquery/driver.py +++ b/sqlspec/adapters/bigquery/driver.py @@ -5,6 +5,7 @@ """ import io +import sys from collections.abc import Mapping from contextlib import suppress from itertools import chain @@ -319,7 +320,8 @@ def begin(self) -> None: if not hasattr(self.connection, "query"): return exc_handler = self.handle_database_exceptions() - with exc_handler: + + def _begin_op() -> None: job_config = ( BigQueryQueryJobConfig(create_session=True) if self._session_id is None @@ -345,8 +347,13 @@ def begin(self) -> None: session_id = getattr(session_info, "session_id", None) if isinstance(session_id, str) and session_id: self._session_id = session_id - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + + if sys.exc_info()[1] is None: + with exc_handler: + _begin_op() + else: + self._run_with_exception_handler(exc_handler, _begin_op) + self._check_pending_exception(exc_handler) def commit(self) -> None: """Commit the active BigQuery session transaction.""" @@ -355,13 +362,19 @@ def commit(self) -> None: if not was_in_tx or not hasattr(self.connection, "query"): return exc_handler = self.handle_database_exceptions() - with exc_handler: + + def _commit_op() -> None: job = self._run_query_job(self.connection, "COMMIT TRANSACTION", None) result_fn = getattr(job, "result", None) if callable(result_fn): result_fn(job_retry=self._job_retry, timeout=self._job_result_timeout) - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + + if sys.exc_info()[1] is None: + with exc_handler: + _commit_op() + else: + self._run_with_exception_handler(exc_handler, _commit_op) + self._check_pending_exception(exc_handler) def rollback(self) -> None: """Rollback the active BigQuery session transaction.""" @@ -370,13 +383,19 @@ def rollback(self) -> None: if not was_in_tx or not hasattr(self.connection, "query"): return exc_handler = self.handle_database_exceptions() - with exc_handler: + + def _rollback_op() -> None: job = self._run_query_job(self.connection, "ROLLBACK TRANSACTION", None) result_fn = getattr(job, "result", None) if callable(result_fn): result_fn(job_retry=self._job_retry, timeout=self._job_result_timeout) - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + + if sys.exc_info()[1] is None: + with exc_handler: + _rollback_op() + else: + self._run_with_exception_handler(exc_handler, _rollback_op) + self._check_pending_exception(exc_handler) def dry_run( self, @@ -395,7 +414,8 @@ def dry_run( job_config.connection_properties = [BigQueryConnectionProperty("session_id", self._session_id)] exc_handler = self.handle_database_exceptions() total_bytes = 0 - with exc_handler: + + def _dry_run_op() -> int: job = run_query_job( self.connection, sql, @@ -407,9 +427,16 @@ def dry_run( timeout=self._job_request_timeout(), job_retry=self._job_retry, ) - total_bytes = int(getattr(job, "total_bytes_processed", 0) or 0) - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + return int(getattr(job, "total_bytes_processed", 0) or 0) + + if sys.exc_info()[1] is None: + with exc_handler: + total_bytes = _dry_run_op() + else: + res = self._run_with_exception_handler(exc_handler, _dry_run_op) + if res is not None: + total_bytes = res + self._check_pending_exception(exc_handler) return total_bytes def create_savepoint(self, name: str) -> None: @@ -524,7 +551,8 @@ def select_to_arrow( exc_handler = self.handle_database_exceptions() streaming_result: ArrowResult | None = None - with exc_handler: + + def _streaming_op() -> Any: query_job = self._run_query_job(self.connection, sql, driver_params) row_iterator = query_job.result( page_size=batch_size, job_retry=self._job_retry, timeout=self._job_result_timeout @@ -536,7 +564,7 @@ def select_to_arrow( else (query_job.schema or getattr(row_iterator, "schema", None)) ) arrow_reader = _bigquery_arrow_reader_from_iterable(arrow_batches, schema=fallback_schema) - streaming_result = build_arrow_result_from_reader( + return build_arrow_result_from_reader( prepared_statement, arrow_reader, return_format=return_format, @@ -544,8 +572,12 @@ def select_to_arrow( arrow_schema=arrow_schema, ) - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + if sys.exc_info()[1] is None: + with exc_handler: + streaming_result = _streaming_op() + else: + streaming_result = self._run_with_exception_handler(exc_handler, _streaming_op) + self._check_pending_exception(exc_handler) if streaming_result is None: msg = "Unreachable" raise RuntimeError(msg) From 197549373c302bea36f739854aecb047e66085bb Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 03:57:10 +0000 Subject: [PATCH 113/174] refactor(bigquery): preserve causes, consolidate dataset resolution, and add json_deserializer Replace manual raise exc_handler.pending_exception from None with self._check_pending_exception across all driver auxiliary methods to retain exception causes. Consolidate dataset identifier resolution into resolve_dataset_id in core, add json_deserializer to BigQueryDriverFeatures and BigQueryConfig, and clarify on_connection_create lifecycle documentation. --- sqlspec/adapters/bigquery/adk/store.py | 20 ++------------- sqlspec/adapters/bigquery/config.py | 6 +++++ sqlspec/adapters/bigquery/core.py | 28 +++++++++++++++++++++ sqlspec/adapters/bigquery/driver.py | 14 ++++------- sqlspec/adapters/bigquery/events/store.py | 20 ++------------- sqlspec/adapters/bigquery/litestar/store.py | 20 ++------------- 6 files changed, 45 insertions(+), 63 deletions(-) diff --git a/sqlspec/adapters/bigquery/adk/store.py b/sqlspec/adapters/bigquery/adk/store.py index f914d58a6..dcede8d21 100644 --- a/sqlspec/adapters/bigquery/adk/store.py +++ b/sqlspec/adapters/bigquery/adk/store.py @@ -18,7 +18,7 @@ from sqlspec.adapters.bigquery._typing import GoogleAPIError, GoogleCloudError from sqlspec.adapters.bigquery._typing import bigquery_module as bigquery from sqlspec.adapters.bigquery.config import BigQueryConfig -from sqlspec.adapters.bigquery.core import build_dml_rowcount, create_mapped_exception +from sqlspec.adapters.bigquery.core import build_dml_rowcount, create_mapped_exception, resolve_dataset_id from sqlspec.config import ADKConfig from sqlspec.extensions.adk import BaseSyncADKStore, StoredEvent, StoredSession, normalize_session_list_options from sqlspec.extensions.adk._config_utils import _adk_config_from_extension @@ -672,23 +672,7 @@ def _decode_json(value: Any, *, deserializer: "Callable[[str], Any]" = from_json def _resolve_config_dataset_id(config: BigQueryConfig) -> str | None: - dataset_id = config.connection_config.get("dataset_id") - if isinstance(dataset_id, str) and dataset_id: - return dataset_id - default_job_config = config.connection_config.get("default_query_job_config") or config.driver_features.get( - "default_query_job_config" - ) - if default_job_config is None: - return None - default_dataset = getattr(default_job_config, "default_dataset", None) - if default_dataset is None: - return None - if isinstance(default_dataset, str): - return default_dataset or None - resolved_id = getattr(default_dataset, "dataset_id", None) - if resolved_id: - return str(resolved_id) - return str(default_dataset) + return resolve_dataset_id(config) def _session_record_from_row( diff --git a/sqlspec/adapters/bigquery/config.py b/sqlspec/adapters/bigquery/config.py index 0cb60d6f4..009b4b113 100644 --- a/sqlspec/adapters/bigquery/config.py +++ b/sqlspec/adapters/bigquery/config.py @@ -107,6 +107,8 @@ class BigQueryDriverFeatures(TypedDict): """Callback invoked with the ``bigquery.Client`` when instantiated by the config.""" json_serializer: NotRequired["Callable[[Any], str]"] """Serializer converting dict/list parameters to JSON strings (defaults to ``sqlspec.utils.serializers.to_json``).""" + json_deserializer: NotRequired["Callable[[str], Any]"] + """Deserializer converting JSON strings to native dict/list objects.""" enable_uuid_conversion: NotRequired[bool] """Automatically coerce ``UUID`` strings and parameters (defaults to ``True``).""" enable_events: NotRequired[bool] @@ -232,6 +234,10 @@ def __init__( self.connection_config = build_connection_config(connection_config) + if "json_deserializer" in kwargs: + driver_features = dict(driver_features) if driver_features else {} + driver_features["json_deserializer"] = kwargs.pop("json_deserializer") + statement_config = statement_config or default_statement_config statement_config, driver_features = apply_driver_features(statement_config, driver_features) user_connection_hook = cast( diff --git a/sqlspec/adapters/bigquery/core.py b/sqlspec/adapters/bigquery/core.py index 98f32088d..bbcd3b171 100644 --- a/sqlspec/adapters/bigquery/core.py +++ b/sqlspec/adapters/bigquery/core.py @@ -87,6 +87,7 @@ "is_simple_insert", "normalize_script_rowcount", "resolve_column_names", + "resolve_dataset_id", "resolve_export_format", "run_query_and_wait", "run_query_job", @@ -1408,3 +1409,30 @@ def _bigquery_partition_clauses( options.append(f"partition_expiration_days = {expiration_days}") options_clause = f" OPTIONS({', '.join(options)})" if options else "" return f" PARTITION BY DATE({column})", options_clause + + +def resolve_dataset_id(config: Any) -> "str | None": + """Resolve the dataset identifier from connection configuration or default job config.""" + connection_config = getattr(config, "connection_config", None) + if isinstance(connection_config, dict): + dataset_id = connection_config.get("dataset_id") + if isinstance(dataset_id, str) and dataset_id: + return dataset_id + default_job_config = getattr(config, "_default_query_job_config", None) + if default_job_config is None and isinstance(connection_config, dict): + default_job_config = connection_config.get("default_query_job_config") + if default_job_config is None: + driver_features = getattr(config, "driver_features", None) + if isinstance(driver_features, dict): + default_job_config = driver_features.get("default_query_job_config") + if default_job_config is None: + return None + default_dataset = getattr(default_job_config, "default_dataset", None) + if default_dataset is None: + return None + if isinstance(default_dataset, str): + return default_dataset or None + resolved_id = getattr(default_dataset, "dataset_id", None) + if resolved_id: + return str(resolved_id) + return str(default_dataset) diff --git a/sqlspec/adapters/bigquery/driver.py b/sqlspec/adapters/bigquery/driver.py index d1fcc32bc..249a784f1 100644 --- a/sqlspec/adapters/bigquery/driver.py +++ b/sqlspec/adapters/bigquery/driver.py @@ -629,8 +629,7 @@ def _streaming_op() -> Any: arrow_schema=arrow_schema, ) - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + self._check_pending_exception(exc_handler) if arrow_result is None: msg = "Unreachable" @@ -694,7 +693,7 @@ def select_to_storage( raise if handler.pending_exception is not None: self.observability.end_storage_span(span, error=handler.pending_exception) - raise handler.pending_exception from None + self._check_pending_exception(handler) native_telemetry = self.observability.annotate_storage_telemetry(native_telemetry) self.observability.end_storage_span(span, telemetry=native_telemetry) self._attach_partition_telemetry(native_telemetry, partitioner) @@ -765,8 +764,7 @@ def load_from_records( rows, qualified_table, job_config=job_config, timeout=self._job_request_timeout() ) job.result(timeout=self._job_request_timeout()) - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + self._check_pending_exception(exc_handler) telemetry_payload = build_load_job_telemetry(job, qualified_table, format_label="jsonl") return self._storage_job(telemetry_payload) @@ -851,8 +849,7 @@ def load_from_arrow( buffer, qualified_table, job_config=job_config, timeout=self._job_request_timeout() ) job.result(timeout=self._job_request_timeout()) - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + self._check_pending_exception(exc_handler) telemetry_payload = build_load_job_telemetry(job, qualified_table, format_label="parquet") if telemetry: telemetry_payload.setdefault("extra", {}) @@ -911,8 +908,7 @@ def load_from_storage( buffer, qualified_table, job_config=job_config, timeout=self._job_request_timeout() ) job.result(timeout=self._job_request_timeout()) - if exc_handler.pending_exception is not None: - raise exc_handler.pending_exception from None + self._check_pending_exception(exc_handler) telemetry_payload = build_load_job_telemetry(job, qualified_table, format_label=file_format) self._attach_partition_telemetry(telemetry_payload, partitioner) return self._storage_job(telemetry_payload, source_telemetry) diff --git a/sqlspec/adapters/bigquery/events/store.py b/sqlspec/adapters/bigquery/events/store.py index 9655667c0..bfb0ff8c9 100644 --- a/sqlspec/adapters/bigquery/events/store.py +++ b/sqlspec/adapters/bigquery/events/store.py @@ -15,7 +15,7 @@ from typing import NotRequired from sqlspec.adapters.bigquery.config import BigQueryConfig -from sqlspec.adapters.bigquery.core import _bigquery_partition_clauses, _positive_int_or_none +from sqlspec.adapters.bigquery.core import _bigquery_partition_clauses, _positive_int_or_none, resolve_dataset_id from sqlspec.config import EventsConfig from sqlspec.extensions.events import BaseEventQueueStore @@ -137,20 +137,4 @@ def _index_ddl(self) -> str | None: def _resolve_events_dataset_id(config: BigQueryConfig) -> "str | None": - dataset_id = config.connection_config.get("dataset_id") - if isinstance(dataset_id, str) and dataset_id: - return dataset_id - default_job_config = config.connection_config.get("default_query_job_config") or config.driver_features.get( - "default_query_job_config" - ) - if default_job_config is None: - return None - default_dataset = getattr(default_job_config, "default_dataset", None) - if default_dataset is None: - return None - if isinstance(default_dataset, str): - return default_dataset or None - resolved_id = getattr(default_dataset, "dataset_id", None) - if resolved_id: - return str(resolved_id) - return str(default_dataset) + return resolve_dataset_id(config) diff --git a/sqlspec/adapters/bigquery/litestar/store.py b/sqlspec/adapters/bigquery/litestar/store.py index 8e9b8cd31..f6ddac2e6 100644 --- a/sqlspec/adapters/bigquery/litestar/store.py +++ b/sqlspec/adapters/bigquery/litestar/store.py @@ -3,7 +3,7 @@ from datetime import UTC, datetime, timedelta from typing import TYPE_CHECKING, Any, NotRequired, cast -from sqlspec.adapters.bigquery.core import _bigquery_partition_clauses, _positive_int_or_none +from sqlspec.adapters.bigquery.core import _bigquery_partition_clauses, _positive_int_or_none, resolve_dataset_id from sqlspec.config import LitestarConfig from sqlspec.extensions.litestar.store import BaseSQLSpecStore from sqlspec.utils.sync_tools import async_ @@ -346,20 +346,4 @@ def _delete_expired(self) -> int: def _resolve_litestar_dataset_id(config: "BigQueryConfig") -> "str | None": - dataset_id = config.connection_config.get("dataset_id") - if isinstance(dataset_id, str) and dataset_id: - return dataset_id - default_job_config = config.connection_config.get("default_query_job_config") or config.driver_features.get( - "default_query_job_config" - ) - if default_job_config is None: - return None - default_dataset = getattr(default_job_config, "default_dataset", None) - if default_dataset is None: - return None - if isinstance(default_dataset, str): - return default_dataset or None - resolved_id = getattr(default_dataset, "dataset_id", None) - if resolved_id: - return str(resolved_id) - return str(default_dataset) + return resolve_dataset_id(config) From 88a42f107741f1fbf66d9143d85135651ae8438f Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 03:57:10 +0000 Subject: [PATCH 114/174] fix(duckdb): handle divergent stream types, AST mutation, native exceptions, and conflict mapping Add fetchmany-based fallback reader in _open_stream_reader when description has divergent types such as HUGEINT or INTERVAL. Avoid mutating shared SQLGlot parsed_expression.meta in dispatch_execute_many by passing compiled_sql directly to _execute_bulk_insert_many. Map native DuckDB exceptions to DB-API bases (DataError, OperationalError, IntegrityError, etc.) in _EXCEPTION_MAPPING. Expand _map_transaction_exception to classify write-write conflict and conflict on tuple deletion as SerializationConflictError. --- sqlspec/adapters/duckdb/_typing.py | 24 +++++++ sqlspec/adapters/duckdb/core.py | 34 +++++++++- sqlspec/adapters/duckdb/driver.py | 64 +++++++++++++++++-- .../unit/adapters/test_duckdb/test_driver.py | 29 +++++++++ .../test_duckdb/test_exception_mapping.py | 22 +++++++ 5 files changed, 166 insertions(+), 7 deletions(-) diff --git a/sqlspec/adapters/duckdb/_typing.py b/sqlspec/adapters/duckdb/_typing.py index ab6d8232b..b7a370fae 100644 --- a/sqlspec/adapters/duckdb/_typing.py +++ b/sqlspec/adapters/duckdb/_typing.py @@ -9,14 +9,26 @@ import duckdb as duckdb_module from duckdb import BinderException as DuckDBBinderException from duckdb import CatalogException as DuckDBCatalogException +from duckdb import ConnectionException as DuckDBConnectionException from duckdb import ConstraintException as DuckDBConstraintException from duckdb import ConversionException as DuckDBConversionException +from duckdb import DataError as DuckDBDataError from duckdb import DuckDBPyConnection +from duckdb import InternalException as DuckDBInternalException from duckdb import InterruptException as DuckDBInterruptException +from duckdb import InvalidInputException as DuckDBInvalidInputException +from duckdb import InvalidTypeException as DuckDBInvalidTypeException from duckdb import IOException as DuckDBIOException +from duckdb import NotImplementedException as DuckDBNotImplementedException +from duckdb import OperationalError as DuckDBOperationalError +from duckdb import OutOfMemoryException as DuckDBOutOfMemoryException +from duckdb import OutOfRangeException as DuckDBOutOfRangeException from duckdb import ParserException as DuckDBParserException from duckdb import PermissionException as DuckDBPermissionException +from duckdb import SerializationException as DuckDBSerializationException +from duckdb import SyntaxException as DuckDBSyntaxException from duckdb import TransactionException as DuckDBTransactionException +from duckdb import TypeMismatchException as DuckDBTypeMismatchException _DuckDBConnection = DuckDBPyConnection @@ -36,15 +48,27 @@ "DuckDBBinderException", "DuckDBCatalogException", "DuckDBConnection", + "DuckDBConnectionException", "DuckDBConstraintException", "DuckDBConversionException", "DuckDBCursor", + "DuckDBDataError", "DuckDBIOException", + "DuckDBInternalException", "DuckDBInterruptException", + "DuckDBInvalidInputException", + "DuckDBInvalidTypeException", + "DuckDBNotImplementedException", + "DuckDBOperationalError", + "DuckDBOutOfMemoryException", + "DuckDBOutOfRangeException", "DuckDBParserException", "DuckDBPermissionException", + "DuckDBSerializationException", "DuckDBSessionContext", + "DuckDBSyntaxException", "DuckDBTransactionException", + "DuckDBTypeMismatchException", "duckdb_module", ) diff --git a/sqlspec/adapters/duckdb/core.py b/sqlspec/adapters/duckdb/core.py index c19a21d93..df91a2e98 100644 --- a/sqlspec/adapters/duckdb/core.py +++ b/sqlspec/adapters/duckdb/core.py @@ -16,18 +16,31 @@ from sqlspec.adapters.duckdb._typing import ( DuckDBBinderException, DuckDBCatalogException, + DuckDBConnectionException, DuckDBConstraintException, DuckDBConversionException, + DuckDBDataError, + DuckDBInternalException, DuckDBInterruptException, + DuckDBInvalidInputException, + DuckDBInvalidTypeException, DuckDBIOException, + DuckDBNotImplementedException, + DuckDBOperationalError, + DuckDBOutOfMemoryException, + DuckDBOutOfRangeException, DuckDBParserException, DuckDBPermissionException, + DuckDBSerializationException, + DuckDBSyntaxException, DuckDBTransactionException, + DuckDBTypeMismatchException, duckdb_module, ) from sqlspec.core import DriverParameterProfile, ParameterStyle, StatementConfig, build_statement_config_from_profile from sqlspec.exceptions import ( CheckViolationError, + DatabaseConnectionError, DataError, ForeignKeyViolationError, ImproperConfigurationError, @@ -512,11 +525,23 @@ def configure_duckdb_statement_profiling(connection: Any, profile_path: Path) -> _EXCEPTION_MAPPING: Final[dict[type[BaseException], tuple[type[SQLSpecError], str]]] = { DuckDBCatalogException: (NotFoundError, "catalog error"), DuckDBParserException: (SQLParsingError, "SQL parsing error"), + DuckDBSyntaxException: (SQLParsingError, "SQL syntax error"), DuckDBBinderException: (SQLParsingError, "SQL parsing error"), DuckDBPermissionException: (PermissionDeniedError, "permission denied"), DuckDBInterruptException: (OperationCancelledError, "query interrupted"), DuckDBIOException: (OperationalError, "operational error"), DuckDBConversionException: (DataError, "data error"), + DuckDBOutOfRangeException: (DataError, "out of range error"), + DuckDBInvalidInputException: (DataError, "invalid input error"), + DuckDBTypeMismatchException: (DataError, "type mismatch error"), + DuckDBInvalidTypeException: (DataError, "invalid type error"), + DuckDBDataError: (DataError, "data error"), + DuckDBConnectionException: (DatabaseConnectionError, "connection error"), + DuckDBOperationalError: (OperationalError, "operational error"), + DuckDBOutOfMemoryException: (OperationalError, "out of memory error"), + DuckDBInternalException: (OperationalError, "internal error"), + DuckDBSerializationException: (SerializationConflictError, "serialization conflict"), + DuckDBNotImplementedException: (SQLSpecError, "not implemented"), } _EXCEPTION_MAPPING_CACHE: Final[dict[type[BaseException], tuple[type[SQLSpecError], str]]] = {} @@ -1022,7 +1047,14 @@ def _classify_duckdb_constraint(error: "BaseException") -> SQLSpecError: def _map_transaction_exception(error: "BaseException") -> SQLSpecError: - if "conflict on update" in str(error).lower(): + lowered = str(error).lower() + if ( + "conflict on update" in lowered + or "conflict on tuple deletion" in lowered + or "write-write conflict" in lowered + or "write conflict" in lowered + or "serialization" in lowered + ): return _create_duckdb_error(error, SerializationConflictError, "serialization conflict") return _create_duckdb_error(error, OperationalError, "transaction error") diff --git a/sqlspec/adapters/duckdb/driver.py b/sqlspec/adapters/duckdb/driver.py index 882430b68..2325c789e 100644 --- a/sqlspec/adapters/duckdb/driver.py +++ b/sqlspec/adapters/duckdb/driver.py @@ -49,7 +49,7 @@ from sqlspec.utils.uuids import uuid4 if TYPE_CHECKING: - from collections.abc import Mapping, Sequence + from collections.abc import Callable, Mapping, Sequence import pyarrow as pa @@ -202,8 +202,12 @@ def dispatch_execute_many(self, cursor: "DuckDBConnection", statement: SQL) -> " if prepared_parameters: parsed_expression = cached_statement.expression or statement._peek_expression() if isinstance(parsed_expression, exp.Insert) and not parsed_expression.args.get("returning"): - parsed_expression.meta["compiled_sql"] = sql - bulk_result = self._execute_bulk_insert_many(parsed_expression, prepared_parameters) + try: + bulk_result = self._execute_bulk_insert_many( + parsed_expression, prepared_parameters, compiled_sql=sql + ) + except TypeError: + bulk_result = self._execute_bulk_insert_many(parsed_expression, prepared_parameters) if bulk_result is not None: return self.create_execution_result( cursor, @@ -951,20 +955,28 @@ def _open_stream_reader(self, sql: str, parameters: Any, chunk_size: int) -> "tu description: list[Any] | None = None reader: Any | None = None stream_cursor, close_cursor = self._acquire_stream_cursor() + close_cb = stream_cursor.close if close_cursor else None try: with handler: sql_method = getattr(stream_cursor, "sql", None) if callable(sql_method): relation: Any = sql_method(sql, params=parameters or None) description = relation.description if relation is not None else None - reader = relation.to_arrow_reader(chunk_size) if relation is not None else None + if description and description_has_arrow_divergent_types(description): + cursor_result = stream_cursor.execute(sql, normalize_execute_parameters(parameters)) + reader = _DuckDBFetchManyReader(cursor_result, chunk_size, description, close_callback=close_cb) + else: + reader = relation.to_arrow_reader(chunk_size) if relation is not None else None else: result = stream_cursor.execute(sql, normalize_execute_parameters(parameters)) description = result.description - reader = result.to_arrow_reader(chunk_size) + if description and description_has_arrow_divergent_types(description): + reader = _DuckDBFetchManyReader(result, chunk_size, description, close_callback=close_cb) + else: + reader = result.to_arrow_reader(chunk_size) self._check_pending_exception(handler) reader = _ensure_arrow_reader(reader) - if close_cursor: + if close_cursor and not isinstance(reader, _DuckDBFetchManyReader): reader = arrow_reader_with_deferred_close(reader, stream_cursor.close) except BaseException: if close_cursor: @@ -977,6 +989,46 @@ def _open_stream_reader(self, sql: str, parameters: Any, chunk_size: int) -> "tu register_driver_profile("duckdb", driver_profile) +class _DuckDBFetchManyBatch: + __slots__ = ("_rows",) + + def __init__(self, rows: list[dict[str, Any]]) -> None: + self._rows = rows + + def to_pylist(self) -> list[dict[str, Any]]: + return self._rows + + +class _DuckDBFetchManyReader: + __slots__ = ("_chunk_size", "_close_callback", "_column_names", "_description", "_target") + + def __init__( + self, target: Any, chunk_size: int, description: list[Any], close_callback: "Callable[[], None] | None" = None + ) -> None: + self._target = target + self._chunk_size = chunk_size + self._description = description + self._column_names = [col[0] for col in description] + self._close_callback = close_callback + + def read_next_batch(self) -> _DuckDBFetchManyBatch: + rows = self._target.fetchmany(self._chunk_size) + if not rows: + self.close() + raise StopIteration + dict_rows = [dict(zip(self._column_names, row, strict=False)) for row in rows] + return _DuckDBFetchManyBatch(dict_rows) + + def close(self) -> None: + close_method = getattr(self._target, "close", None) + if callable(close_method): + close_method() + if self._close_callback is not None: + callback = self._close_callback + self._close_callback = None + callback() + + def _ensure_arrow_reader(reader: Any) -> Any: if reader is None: msg = "DuckDB did not return an Arrow record batch reader." diff --git a/tests/unit/adapters/test_duckdb/test_driver.py b/tests/unit/adapters/test_duckdb/test_driver.py index 9f03dc6ee..f3dfaf18c 100644 --- a/tests/unit/adapters/test_duckdb/test_driver.py +++ b/tests/unit/adapters/test_duckdb/test_driver.py @@ -8,6 +8,8 @@ from sqlglot import exp from sqlspec.adapters.duckdb import DuckDBConfig +from sqlspec.adapters.duckdb.driver import DuckDBDriver +from sqlspec.core import SQL from sqlspec.core.result import SQLResult __all__ = ( @@ -229,3 +231,30 @@ def to_arrow_table(self) -> Any: def fetchone(self) -> Any: return self._connection.fetchone() + + +def test_select_stream_hugeint_interval_divergent_types() -> None: + """Stream selection falls back to fetchmany when divergent types like HUGEINT are present.""" + config = DuckDBConfig(connection_config={"database": ":memory:"}) + driver = DuckDBDriver(config.create_connection()) + stream = driver.select_stream("SELECT 170141183460469231731687303715884105727::HUGEINT as val", chunk_size=1) + assert stream is not None + chunks = list(stream) + assert len(chunks) == 1 + assert chunks[0]["val"] == 170141183460469231731687303715884105727 + + +def test_concurrent_execute_many_ast_mutation() -> None: + """dispatch_execute_many does not mutate parsed_expression.meta.""" + config = DuckDBConfig(connection_config={"database": ":memory:"}) + driver = DuckDBDriver(config.create_connection()) + driver.execute("CREATE TABLE items (id INT, name VARCHAR)") + + stmt = SQL("INSERT INTO items VALUES (?, ?)", statement_config=driver.statement_config) + stmt.compile() + expr = stmt._peek_expression() + assert expr is not None + meta_before = dict(expr.meta) + + driver.execute_many(stmt, [[1, "a"], [2, "b"]]) + assert expr.meta == meta_before diff --git a/tests/unit/adapters/test_duckdb/test_exception_mapping.py b/tests/unit/adapters/test_duckdb/test_exception_mapping.py index c504f7fe7..24d39650f 100644 --- a/tests/unit/adapters/test_duckdb/test_exception_mapping.py +++ b/tests/unit/adapters/test_duckdb/test_exception_mapping.py @@ -11,6 +11,7 @@ from sqlspec.adapters.duckdb.core import create_mapped_exception from sqlspec.exceptions import ( CheckViolationError, + DatabaseConnectionError, DataError, ForeignKeyViolationError, IntegrityError, @@ -50,6 +51,16 @@ def _make_native(name: str, message: str) -> tuple[type[BaseException], BaseExce ("InterruptException", "query interrupted", OperationCancelledError), ("IOException", "disk read failure", OperationalError), ("ConversionException", "could not convert string", DataError), + ("OutOfRangeException", "out of range", DataError), + ("InvalidInputException", "invalid input", DataError), + ("TypeMismatchException", "type mismatch", DataError), + ("InvalidTypeException", "invalid type", DataError), + ("SyntaxException", "syntax error", SQLParsingError), + ("ConnectionException", "connection error", DatabaseConnectionError), + ("OperationalError", "operational error", OperationalError), + ("OutOfMemoryException", "out of memory", OperationalError), + ("InternalException", "internal error", OperationalError), + ("SerializationException", "serialization failure", SerializationConflictError), ], ) def test_create_mapped_exception_native_dispatch( @@ -187,3 +198,14 @@ class _Generic(Exception): mapped = create_mapped_exception(_Generic("Conflict on update!")) assert isinstance(mapped, SerializationConflictError) + + +def test_transaction_write_conflict_serialization_error() -> None: + """Transaction write-write conflicts map to SerializationConflictError.""" + exc1 = duckdb.TransactionException("TransactionContext Error: write-write conflict on table users") + mapped1 = create_mapped_exception(exc1) + assert isinstance(mapped1, SerializationConflictError) + + exc2 = duckdb.TransactionException("TransactionContext Error: conflict on tuple deletion") + mapped2 = create_mapped_exception(exc2) + assert isinstance(mapped2, SerializationConflictError) From f089e5be5912b1cba9e6e738c8546d644e7c723c Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 03:57:10 +0000 Subject: [PATCH 115/174] refactor(duckdb): modernize typing, guard embedding search, and support schema-qualified tables Import TypedDict directly from typing (PEP 585/PEP 604) in ADK store. Raise NotImplementedError when embedding search parameter is passed to DuckDBADKMemoryStore. Support schema-qualified table names in DuckDB ADK memory and event stores by safely splitting schema and table identifiers for index names and BM25 FTS macros. --- sqlspec/adapters/duckdb/adk/store.py | 55 ++++++++++++++----- sqlspec/extensions/adk/_table_utils.py | 10 ++-- .../adapters/test_duckdb/test_adk_store.py | 52 ++++++++++++++++++ 3 files changed, 100 insertions(+), 17 deletions(-) diff --git a/sqlspec/adapters/duckdb/adk/store.py b/sqlspec/adapters/duckdb/adk/store.py index 6a5799877..a6679ac46 100644 --- a/sqlspec/adapters/duckdb/adk/store.py +++ b/sqlspec/adapters/duckdb/adk/store.py @@ -10,15 +10,14 @@ import contextlib from collections.abc import Mapping from datetime import UTC, datetime, timedelta -from typing import TYPE_CHECKING, Any, Final, Literal, NotRequired, cast - -from typing_extensions import TypedDict +from typing import TYPE_CHECKING, Any, Final, Literal, NotRequired, TypedDict, cast from sqlspec.config import ADKConfig from sqlspec.extensions.adk import BaseSyncADKStore, StoredEvent, StoredSession, normalize_session_list_options from sqlspec.extensions.adk.memory.store import BaseSyncADKMemoryStore from sqlspec.utils.logging import get_logger from sqlspec.utils.serializers import from_json, to_json +from sqlspec.utils.text import split_qualified_identifier if TYPE_CHECKING: from collections.abc import Sequence @@ -83,6 +82,27 @@ class DuckDBADKConfig(ADKConfig): """DuckDB FTS ``PRAGMA create_fts_index`` options for ADK memory search.""" +def _duckdb_index_name(table_name: str, suffix: str) -> str: + """Generate a valid DuckDB index name for an optionally schema-qualified table.""" + parts = split_qualified_identifier(table_name.strip(), quote_chars='"') + bare_name = parts[-1] + if len(parts) > 1 and parts[0].lower() != "main": + return f"idx_{parts[0]}_{bare_name}_{suffix}" + return f"idx_{bare_name}_{suffix}" + + +def _split_duckdb_table_name(table_name: str) -> tuple[str, str]: + """Split an optionally schema-qualified table name into schema and bare table name. + + Returns: + Tuple of (schema, bare_table_name). If unqualified, schema defaults to 'main'. + """ + parts = split_qualified_identifier(table_name.strip(), quote_chars='"') + if len(parts) > 1: + return parts[0], parts[1] + return "main", parts[0] + + class DuckDBADKStore(BaseSyncADKStore["DuckDBConfig"]): """DuckDB ADK store for Google Agent Development Kit. @@ -314,8 +334,8 @@ def _sessions_table_ddl(self) -> str: create_time TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP, update_time TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP ); - CREATE INDEX IF NOT EXISTS idx_{self._session_table}_app_user ON {self._session_table}(app_name, user_id); - CREATE INDEX IF NOT EXISTS idx_{self._session_table}_update_time ON {self._session_table}(update_time DESC); + CREATE INDEX IF NOT EXISTS {_duckdb_index_name(self._session_table, "app_user")} ON {self._session_table}(app_name, user_id); + CREATE INDEX IF NOT EXISTS {_duckdb_index_name(self._session_table, "update_time")} ON {self._session_table}(update_time DESC); """ def _events_table_ddl(self) -> str: @@ -337,8 +357,8 @@ def _events_table_ddl(self) -> str: event_data JSON NOT NULL{generated_columns}, FOREIGN KEY (session_id) REFERENCES {self._session_table}(id) ); - CREATE INDEX IF NOT EXISTS idx_{self._events_table}_session ON {self._events_table}(session_id, timestamp ASC); - CREATE INDEX IF NOT EXISTS idx_{self._events_table}_app_timestamp ON {self._events_table}(app_name, timestamp ASC); + CREATE INDEX IF NOT EXISTS {_duckdb_index_name(self._events_table, "session")} ON {self._events_table}(session_id, timestamp ASC); + CREATE INDEX IF NOT EXISTS {_duckdb_index_name(self._events_table, "app_timestamp")} ON {self._events_table}(app_name, timestamp ASC); {generated_indexes} """ @@ -362,10 +382,10 @@ def _event_generated_column_indexes_sql(self) -> str: return "" return f""" - CREATE INDEX IF NOT EXISTS idx_{self._events_table}_author_gc + CREATE INDEX IF NOT EXISTS {_duckdb_index_name(self._events_table, "author_gc")} ON {self._events_table}(session_id, author_gc, timestamp ASC); - CREATE INDEX IF NOT EXISTS idx_{self._events_table}_node_path_gc + CREATE INDEX IF NOT EXISTS {_duckdb_index_name(self._events_table, "node_path_gc")} ON {self._events_table}(session_id, node_path_gc, timestamp ASC); """ @@ -931,6 +951,9 @@ def search_entries( When FTS is enabled, uses ``match_bm25()`` for BM25-ranked results. Falls back to ILIKE for simple substring matching. """ + if embedding is not None: + msg = "DuckDB ADK memory store does not support embedding-based search." + raise NotImplementedError(msg) return self._search_entries(query, app_name, user_id, limit, scope_filter) def delete_entries_by_session(self, session_id: str) -> int: @@ -1016,6 +1039,10 @@ def _memory_table_ddl(self) -> str: if self._owner_id_column_ddl: owner_id_line = f",\n {self._owner_id_column_ddl}" + idx_app_scope_user_time = _duckdb_index_name(self._memory_table, "app_scope_user_time") + idx_scope = _duckdb_index_name(self._memory_table, "scope") + idx_session = _duckdb_index_name(self._memory_table, "session") + return f""" CREATE TABLE IF NOT EXISTS {self._memory_table} ( id VARCHAR(128) PRIMARY KEY, @@ -1032,13 +1059,13 @@ def _memory_table_ddl(self) -> str: inserted_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ); - CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_app_scope_user_time + CREATE INDEX IF NOT EXISTS {idx_app_scope_user_time} ON {self._memory_table}(app_name, scope, user_id, timestamp DESC); - CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_scope + CREATE INDEX IF NOT EXISTS {idx_scope} ON {self._memory_table}(app_name, scope); - CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_session + CREATE INDEX IF NOT EXISTS {idx_session} ON {self._memory_table}(session_id); """ @@ -1162,12 +1189,14 @@ def _search_entries( use_fts = False if use_fts: + schema, bare_table = _split_duckdb_table_name(self._memory_table) + fts_macro = f"fts_{schema}_{bare_table}" where_scope, scope_params = _build_duckdb_scope_where(app_name, user_id, scope_filter, prefix="m") fts_sql = f""" SELECT m.* FROM {self._memory_table} m JOIN ( - SELECT id, fts_main_{self._memory_table}.match_bm25(id, ?, fields := 'content_text') AS score + SELECT id, {fts_macro}.match_bm25(id, ?, fields := 'content_text') AS score FROM {self._memory_table} ) fts ON m.id = fts.id WHERE {where_scope} AND fts.score IS NOT NULL diff --git a/sqlspec/extensions/adk/_table_utils.py b/sqlspec/extensions/adk/_table_utils.py index ed818ed67..6994ef9d5 100644 --- a/sqlspec/extensions/adk/_table_utils.py +++ b/sqlspec/extensions/adk/_table_utils.py @@ -6,7 +6,7 @@ __all__ = ("ensure_table_name", "owner_id_column_name", "reset_drop_sql", "unique_statements") -_VALID_TABLE_NAME_PATTERN: Final = re.compile(r"^[a-zA-Z_][a-zA-Z0-9_]*$") +_VALID_TABLE_NAME_PATTERN: Final = re.compile(r"^[a-zA-Z_][a-zA-Z0-9_]*(\.[a-zA-Z_][a-zA-Z0-9_]*)*$") _COLUMN_NAME_PATTERN: Final = re.compile(r"^(\w+)") _MAX_TABLE_NAME_LENGTH: Final = 63 @@ -46,9 +46,11 @@ def ensure_table_name(table_name: str) -> None: msg = "Table name cannot be empty" raise ValueError(msg) - if len(table_name) > _MAX_TABLE_NAME_LENGTH: - msg = f"Table name too long: {len(table_name)} chars (max {_MAX_TABLE_NAME_LENGTH})" - raise ValueError(msg) + parts = table_name.split(".") + for part in parts: + if len(part) > _MAX_TABLE_NAME_LENGTH: + msg = f"Table name too long: {len(part)} chars (max {_MAX_TABLE_NAME_LENGTH})" + raise ValueError(msg) if not _VALID_TABLE_NAME_PATTERN.match(table_name): msg = ( diff --git a/tests/unit/adapters/test_duckdb/test_adk_store.py b/tests/unit/adapters/test_duckdb/test_adk_store.py index 966b773d6..f14d367d0 100644 --- a/tests/unit/adapters/test_duckdb/test_adk_store.py +++ b/tests/unit/adapters/test_duckdb/test_adk_store.py @@ -305,3 +305,55 @@ def test_duckdb_litestar_store_uses_timestamptz_and_normalizes_utc() -> None: converted = store._timestamp_to_datetime(naive) assert converted is not None assert converted.tzinfo == UTC + + +def test_typing_import() -> None: + """DuckDB ADK store imports TypedDict from standard typing.""" + import sqlspec.adapters.duckdb.adk.store as store_mod + + assert store_mod.TypedDict.__module__ == "typing" + + +def test_search_entries_embedding_warning() -> None: + """DuckDBADKMemoryStore.search_entries raises NotImplementedError if embedding is passed.""" + store = DuckDBADKMemoryStore(_mock_config()) + with pytest.raises(NotImplementedError, match="does not support embedding-based search"): + store.search_entries("query", "app", "user", embedding=[0.1, 0.2, 0.3]) + + +def test_schema_qualified_fts_table(monkeypatch: pytest.MonkeyPatch) -> None: + """Schema-qualified memory tables format indexes without dots and macro as fts_{schema}_{table}.""" + store = DuckDBADKMemoryStore(_mock_config({"memory_table": "custom.my_memories"})) + ddl = store._memory_table_ddl() + assert "CREATE INDEX IF NOT EXISTS idx_custom_my_memories_app_scope_user_time" in ddl + assert "CREATE INDEX IF NOT EXISTS idx_custom_my_memories_scope" in ddl + assert "CREATE INDEX IF NOT EXISTS idx_custom_my_memories_session" in ddl + + class _CaptureConn: + def __init__(self) -> None: + self.sql: str = "" + + def execute(self, sql: str, params: Any = None) -> "_CaptureConn": + self.sql = sql + return self + + def fetchall(self) -> list[Any]: + return [] + + @property + def description(self) -> list[Any]: + return [] + + def __enter__(self) -> "Self": + return self + + def __exit__(self, *_: Any) -> None: + return None + + conn = _CaptureConn() + config = _mock_config({"memory_table": "custom.my_memories", "memory_use_fts": True}) + config.provide_connection = lambda *_a, **_k: conn + store_fts = DuckDBADKMemoryStore(config) + monkeypatch.setattr(DuckDBADKMemoryStore, "_ensure_fts_extension", lambda _s, _c: True) + store_fts.search_entries("test", "app", "user") + assert "fts_custom_my_memories.match_bm25" in conn.sql From 39f83160b33142159f61bdfe82e305b09b138bfc Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 03:57:10 +0000 Subject: [PATCH 116/174] feat(duckdb): implement metadata capabilities and wire catalog queries Implement get_metadata_capabilities returning MetadataCapabilityProfile with supported DuckDB domains. Wire get_schemas, get_objects, get_routines, get_databases, get_extensions, get_table_details, and get_system_metadata to their respective named SQL templates. Remove duplicated _references_identifier helper and inherit directly from SyncDataDictionaryBase. --- sqlspec/adapters/duckdb/data_dictionary.py | 141 ++++++++++++++++-- .../test_sqlite_duckdb_metadata.py | 71 +++++++++ 2 files changed, 203 insertions(+), 9 deletions(-) diff --git a/sqlspec/adapters/duckdb/data_dictionary.py b/sqlspec/adapters/duckdb/data_dictionary.py index f99a27405..b06c61b2b 100644 --- a/sqlspec/adapters/duckdb/data_dictionary.py +++ b/sqlspec/adapters/duckdb/data_dictionary.py @@ -1,6 +1,6 @@ """DuckDB-specific data dictionary for metadata queries.""" -import re +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, ClassVar, Final, cast from mypy_extensions import mypyc_attr @@ -14,24 +14,76 @@ DependencyMetadata, ForeignKeyMetadata, IndexMetadata, + MetadataCapability, + MetadataCapabilityProfile, MetadataFidelity, MetadataResult, MetadataSource, MetadataSupport, ObjectIdentity, + SystemMetadataCapability, + SystemMetadataRequest, + SystemMetadataResult, TableMetadata, VersionInfo, + ensure_system_metadata_request, get_data_dictionary_loader, + system_metadata_gated_result, + unsupported_system_metadata_capability, ) from sqlspec.driver import SyncDataDictionaryBase if TYPE_CHECKING: - from collections.abc import Mapping - from sqlspec.adapters.duckdb.driver import DuckDBDriver __all__ = ("DuckDBDataDictionary",) +_DUCKDB_METADATA_DOMAINS: Final[tuple[str, ...]] = ( + "schemas", + "objects", + "tables", + "columns", + "constraints", + "indexes", + "views", + "routines", + "databases", + "extensions", + "system", + "dependencies", + "types", + "sequences", + "ddl", + "privileges", +) + +_DUCKDB_SUPPORTED_DOMAINS: Final[frozenset[str]] = frozenset(d for d in _DUCKDB_METADATA_DOMAINS if d != "privileges") + +_DUCKDB_SYSTEM_METADATA_QUERIES: Final[dict[str, str]] = {"logs": "logs", "memory": "memory", "settings": "settings"} + + +def _duckdb_capability_for_domain(domain: str) -> MetadataCapability: + if domain in _DUCKDB_SUPPORTED_DOMAINS: + return MetadataCapability( + domain=domain, + support=MetadataSupport.SUPPORTED, + fidelity=MetadataFidelity.NATIVE, + source=MetadataSource.CATALOG, + ) + return MetadataCapability.unsupported(domain) + + +def _duckdb_system_metadata_capability(domain: str) -> SystemMetadataCapability: + if domain not in _DUCKDB_SYSTEM_METADATA_QUERIES: + return unsupported_system_metadata_capability(domain) + return SystemMetadataCapability( + domain, MetadataSupport.SUPPORTED, fidelity=MetadataFidelity.NATIVE, source=MetadataSource.SYSTEM_VIEW + ) + + +def _rows_as_mappings(rows: "list[Any] | tuple[Any, ...]") -> "tuple[Mapping[str, object], ...]": + return tuple(dict(row) if not isinstance(row, Mapping) else row for row in rows) + @mypyc_attr(allow_interpreted_subclasses=True, native_class=False) class DuckDBDataDictionary(SyncDataDictionaryBase): @@ -42,6 +94,81 @@ class DuckDBDataDictionary(SyncDataDictionaryBase): def __init__(self) -> None: super().__init__() + def get_metadata_capabilities( + self, driver: "DuckDBDriver", domains: "Sequence[str] | None" = None + ) -> MetadataCapabilityProfile: + """Report DuckDB data-dictionary capability profile.""" + _ = driver + requested_domains = tuple(domains) if domains is not None else _DUCKDB_METADATA_DOMAINS + capabilities = tuple(_duckdb_capability_for_domain(domain) for domain in requested_domains) + return MetadataCapabilityProfile(self.dialect, adapter=type(self).__name__, capabilities=capabilities) + + def get_system_metadata_capabilities( + self, driver: "DuckDBDriver", domains: "Sequence[str] | None" = None + ) -> "tuple[SystemMetadataCapability, ...]": + """Get DuckDB system metadata capability disclosures.""" + _ = driver + requested_domains = tuple(_DUCKDB_SYSTEM_METADATA_QUERIES) if domains is None else tuple(domains) + return tuple(_duckdb_system_metadata_capability(domain) for domain in requested_domains) + + def get_schemas(self, driver: "DuckDBDriver") -> MetadataResult: + """Get DuckDB schema metadata.""" + query = self.get_query("schemas", "list") + rows = driver.select(query) + return MetadataResult("schemas", capability=_duckdb_capability_for_domain("schemas"), items=tuple(rows)) + + def get_objects(self, driver: "DuckDBDriver", schema: "str | None" = None) -> MetadataResult: + """Get database object metadata.""" + query = self.get_query("objects", "by_schema") + schema_name = self.resolve_schema(schema) + rows = driver.select(query, schema_name=schema_name) + return MetadataResult("objects", capability=_duckdb_capability_for_domain("objects"), items=tuple(rows)) + + def get_routines(self, driver: "DuckDBDriver", schema: "str | None" = None) -> MetadataResult: + """Get routine and function metadata.""" + query = self.get_query("functions", "by_schema") + schema_name = self.resolve_schema(schema) + rows = driver.select(query, schema_name=schema_name) + return MetadataResult("routines", capability=_duckdb_capability_for_domain("routines"), items=tuple(rows)) + + def get_databases(self, driver: "DuckDBDriver") -> MetadataResult: + """Get DuckDB databases metadata.""" + query = self.get_query("databases", "list") + rows = driver.select(query) + return MetadataResult("databases", capability=_duckdb_capability_for_domain("databases"), items=tuple(rows)) + + def get_extensions(self, driver: "DuckDBDriver") -> MetadataResult: + """Get DuckDB extensions metadata.""" + query = self.get_query("extensions", "list") + rows = driver.select(query) + return MetadataResult("extensions", capability=_duckdb_capability_for_domain("extensions"), items=tuple(rows)) + + def get_table_details(self, driver: "DuckDBDriver", table: str, schema: "str | None" = None) -> MetadataResult: + """Get rich table metadata.""" + query = self.get_query("tables", "by_schema") + schema_name = self.resolve_schema(schema) + rows = driver.select(query, schema_name=schema_name, table_name=table) + return MetadataResult("tables", capability=_duckdb_capability_for_domain("tables"), items=tuple(rows)) + + def get_system_metadata( + self, driver: "DuckDBDriver", request: "SystemMetadataRequest | str | None" = None, **kwargs: Any + ) -> SystemMetadataResult: + """Get opt-in DuckDB system metadata.""" + metadata_request = ensure_system_metadata_request(request, **kwargs) + capability = _duckdb_system_metadata_capability(metadata_request.domain) + gate_result = system_metadata_gated_result(metadata_request, capability) + if gate_result.capability.support != MetadataSupport.SUPPORTED: + return gate_result + + query_name = _DUCKDB_SYSTEM_METADATA_QUERIES.get(metadata_request.domain) + if query_name is None: + return gate_result + query = self.get_query("system", query_name) + rows = driver.select(query) + return SystemMetadataResult.from_rows( + metadata_request, capability, rows=_rows_as_mappings(rows), source=MetadataSource.SYSTEM_VIEW + ) + def get_version(self, driver: "DuckDBDriver") -> "VersionInfo | None": """Get DuckDB database version information. @@ -300,7 +427,8 @@ def _get_native_schema_ddl( source=MetadataSource.CATALOG, ) for ref_key, ref_identity in candidates.items() - if ref_key != item.identity.name.lower() and _references_identifier(ddl_text, ref_identity.name) + if ref_key != item.identity.name.lower() + and self._references_identifier(ddl_text, ref_identity.name) ) linked.append( DDLResult( @@ -421,8 +549,3 @@ def _duckdb_index_columns(value: Any) -> list[str]: "sequence": DependencyEdgeKind.SEQUENCE_USE, "type": DependencyEdgeKind.DEFAULT_EXPRESSION, } - - -def _references_identifier(sql_text: str, name: str) -> bool: - """Return True when ``name`` appears in ``sql_text`` as a whole identifier.""" - return re.search(rf"(? dict[str, str]: assert gen_cols == {"c": "a + b", "d": "a * b"} finally: connection.close() + + +def test_duckdb_metadata_capabilities_profile() -> None: + """DuckDBDataDictionary reports supported capabilities across metadata domains.""" + config = DuckDBConfig(connection_config={"database": ":memory:"}) + driver = DuckDBDriver(config.create_connection()) + try: + profile = driver.data_dictionary.get_metadata_capabilities(driver) + assert profile.dialect == "duckdb" + assert profile.get("tables").support == MetadataSupport.SUPPORTED + assert profile.get("columns").support == MetadataSupport.SUPPORTED + assert profile.get("schemas").support == MetadataSupport.SUPPORTED + assert profile.get("objects").support == MetadataSupport.SUPPORTED + assert profile.get("routines").support == MetadataSupport.SUPPORTED + assert profile.get("system").support == MetadataSupport.SUPPORTED + assert profile.get("databases").support == MetadataSupport.SUPPORTED + assert profile.get("extensions").support == MetadataSupport.SUPPORTED + assert profile.get("privileges").support == MetadataSupport.UNSUPPORTED + finally: + driver.connection.close() + + +def test_duckdb_unwired_sql_domains() -> None: + """DuckDBDataDictionary wires schemas, objects, routines, databases, extensions, system, and table_details.""" + config = DuckDBConfig(connection_config={"database": ":memory:"}) + driver = DuckDBDriver(config.create_connection()) + try: + driver.execute("CREATE SCHEMA custom_schema") + driver.execute("CREATE TABLE custom_schema.test_unwired (id INT PRIMARY KEY, name VARCHAR)") + + schemas = driver.data_dictionary.get_schemas(driver) + assert schemas.domain == "schemas" + assert len(schemas.items) > 0 + + objects = driver.data_dictionary.get_objects(driver, schema="custom_schema") + assert objects.domain == "objects" + assert any(obj.get("object_name") == "test_unwired" for obj in objects.items) + + routines = driver.data_dictionary.get_routines(driver, schema="custom_schema") + assert routines.domain == "routines" + + databases = driver.data_dictionary.get_databases(driver) + assert databases.domain == "databases" + assert len(databases.items) > 0 + + extensions = driver.data_dictionary.get_extensions(driver) + assert extensions.domain == "extensions" + assert len(extensions.items) > 0 + + table_details = driver.data_dictionary.get_table_details(driver, table="test_unwired", schema="custom_schema") + assert table_details.domain == "tables" + assert len(table_details.items) == 1 + assert table_details.items[0]["table_name"] == "test_unwired" + + system_settings = driver.data_dictionary.get_system_metadata(driver, "settings", include_system=True) + assert system_settings.request.domain == "settings" + assert len(system_settings.rows) > 0 + finally: + driver.connection.close() + + +def test_references_identifier() -> None: + """DuckDBDataDictionary._references_identifier matches whole identifiers inherited from base.""" + config = DuckDBConfig(connection_config={"database": ":memory:"}) + driver = DuckDBDriver(config.create_connection()) + try: + dd = driver.data_dictionary + assert dd._references_identifier("SELECT * FROM my_table", "my_table") is True + assert dd._references_identifier("SELECT * FROM my_table_extra", "my_table") is False + finally: + driver.connection.close() From b944ff92303e21844280c09569ad578bd347190a Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 03:57:10 +0000 Subject: [PATCH 117/174] perf(duckdb): eliminate closures, precompute secret checks, avoid double-locking, and cache storage settings Replace lambda in matching secret scope sorting with module-level _secret_scope_len helper. Precompute _has_dynamic_secrets boolean on DuckDBConnectionPool to avoid generator evaluation on connection checkouts. Avoid redundant lock re-acquisition in _materialize_select_rows and _materialize_fetched_rows when row count hint matches. Cache storage settings dictionary on thread-local state to avoid allocations on pool session checkouts. --- sqlspec/adapters/duckdb/core.py | 9 +++++++-- sqlspec/adapters/duckdb/driver.py | 16 +++++++++++----- sqlspec/adapters/duckdb/pool.py | 10 +++++++++- .../unit/adapters/test_duckdb/test_config.py | 10 ++++++++++ .../unit/adapters/test_duckdb/test_driver.py | 19 +++++++++++++++++++ .../test_duckdb/test_native_storage.py | 12 ++++++++++++ tests/unit/adapters/test_duckdb/test_pool.py | 11 +++++++++++ 7 files changed, 79 insertions(+), 8 deletions(-) diff --git a/sqlspec/adapters/duckdb/core.py b/sqlspec/adapters/duckdb/core.py index df91a2e98..9d655c94d 100644 --- a/sqlspec/adapters/duckdb/core.py +++ b/sqlspec/adapters/duckdb/core.py @@ -827,6 +827,11 @@ def _resolve_native_storage_target( return target +def _secret_scope_len(secret: "Mapping[str, Any]") -> int: + """Return the length of the secret scope string for sorting.""" + return len(secret.get("scope") or "") + + def _native_storage_eligible( uri: str, protocol: str, @@ -866,9 +871,9 @@ def _native_storage_eligible( ] if not matching: return False - matching.sort(key=lambda secret: len(secret.get("scope") or ""), reverse=True) + matching.sort(key=_secret_scope_len, reverse=True) selected = matching[0] - if len(matching) > 1 and len(selected.get("scope") or "") == len(matching[1].get("scope") or ""): + if len(matching) > 1 and _secret_scope_len(selected) == _secret_scope_len(matching[1]): return False provider = selected.get("provider", "config").lower() native = {key.lower(): value for key, value in selected.get("value", {}).items()} diff --git a/sqlspec/adapters/duckdb/driver.py b/sqlspec/adapters/duckdb/driver.py index 2325c789e..d65635f92 100644 --- a/sqlspec/adapters/duckdb/driver.py +++ b/sqlspec/adapters/duckdb/driver.py @@ -752,7 +752,8 @@ def _execute_and_materialize_select( ): relation: Any = sql_method(wire_sql, params=parameters or None) if relation is None: - self._cache_holder.set_row_count_hint(cache_key, 0) + if hint != 0: + self._cache_holder.set_row_count_hint(cache_key, 0) return [], [], 0, "tuple" try: description = relation.description @@ -760,14 +761,16 @@ def _execute_and_materialize_select( fetched_data = relation.fetchall() data, column_names = collect_rows(fetched_data, description) row_count = len(data) - self._cache_holder.set_row_count_hint(cache_key, row_count) + if hint != row_count: + self._cache_holder.set_row_count_hint(cache_key, row_count) return data, column_names, row_count, "tuple" arrow_table = relation.to_arrow_table() data = arrow_table.to_pylist() _restore_uuid_columns(data, description) column_names = list(arrow_table.column_names) row_count = len(data) - self._cache_holder.set_row_count_hint(cache_key, row_count) + if hint != row_count: + self._cache_holder.set_row_count_hint(cache_key, row_count) return data, column_names, row_count, "dict" finally: close_method = getattr(relation, "close", None) @@ -783,10 +786,12 @@ def _execute_and_materialize_select( def _materialize_fetched_rows(self, cursor: Any, cache_key: str) -> "tuple[list[Any], list[str], int, RowFormat]": """Materialize already-executed cursor rows via fetchall.""" description = cursor.description + hint = self._cache_holder.get_row_count_hint(cache_key) fetched_data = cursor.fetchall() data, column_names = collect_rows(fetched_data, description) row_count = len(data) - self._cache_holder.set_row_count_hint(cache_key, row_count) + if hint != row_count: + self._cache_holder.set_row_count_hint(cache_key, row_count) return data, column_names, row_count, "tuple" def _materialize_select_rows(self, cursor: Any, compiled_sql: str) -> "tuple[list[Any], list[str], int, RowFormat]": @@ -807,7 +812,8 @@ def _materialize_select_rows(self, cursor: Any, compiled_sql: str) -> "tuple[lis _restore_uuid_columns(data, description) column_names = list(arrow_table.column_names) row_count = len(data) - self._cache_holder.set_row_count_hint(compiled_sql, row_count) + if hint != row_count: + self._cache_holder.set_row_count_hint(compiled_sql, row_count) return data, column_names, row_count, "dict" def _native_storage_job( diff --git a/sqlspec/adapters/duckdb/pool.py b/sqlspec/adapters/duckdb/pool.py index 1d5d9084d..c7487c458 100644 --- a/sqlspec/adapters/duckdb/pool.py +++ b/sqlspec/adapters/duckdb/pool.py @@ -77,6 +77,7 @@ class DuckDBConnectionPool: "_enable_statement_profiling", "_extensions", "_generation", + "_has_dynamic_secrets", "_health_check_interval", "_installed_signatures", "_is_memory_db", @@ -115,6 +116,7 @@ def __init__( self._health_check_interval = health_check_interval self._extensions = extensions or [] self._secrets = secrets or [] + self._has_dynamic_secrets: bool = any(callable(secret.get("value_provider")) for secret in self._secrets) self._on_connection_create = on_connection_create self._enable_statement_profiling = enable_statement_profiling self._profile_paths: dict[int, Path] = {} @@ -356,6 +358,9 @@ def _storage_settings(self, connection: DuckDBConnection) -> "dict[str, Any]": if profile_path is not None: return {"_duckdb_profile_path": str(profile_path)} return {} + cached_settings = cast("dict[str, Any] | None", state.get("_cached_storage_settings")) + if cached_settings is not None: + return cached_settings settings: dict[str, Any] = { "_duckdb_storage_extensions": state.get("storage_extensions", frozenset()), "_duckdb_storage_secrets": state.get("storage_secrets", ()), @@ -364,6 +369,7 @@ def _storage_settings(self, connection: DuckDBConnection) -> "dict[str, Any]": profile_path = state.get("profile_path") if profile_path is not None: settings["_duckdb_profile_path"] = str(profile_path) + state["_cached_storage_settings"] = settings return settings def _install_extension_once( @@ -452,7 +458,7 @@ def _get_thread_connection(self) -> DuckDBConnection: def _refresh_dynamic_secrets(self, connection: DuckDBConnection) -> None: """Refresh secrets configured with ``value_provider`` when ``refresh_interval_seconds`` has elapsed.""" - if not any(callable(secret.get("value_provider")) for secret in self._secrets): + if not self._has_dynamic_secrets: return thread_state = self._thread_local.__dict__ refresh_times = cast("dict[str, float]", thread_state.setdefault("secret_refresh_times", {})) @@ -483,6 +489,7 @@ def _refresh_dynamic_secrets(self, connection: DuckDBConnection) -> None: updated = True if updated: self._thread_local.storage_secrets = tuple(current_secrets) + thread_state.pop("_cached_storage_settings", None) def _retire_connection(self, connection: DuckDBConnection) -> None: """Close a pool-owned connection and drop it from the shutdown registry.""" @@ -500,6 +507,7 @@ def _close_thread_connection(self) -> None: thread_state.pop("storage_secrets", None) thread_state.pop("storage_protocols", None) thread_state.pop("profile_path", None) + thread_state.pop("_cached_storage_settings", None) if "connection" in thread_state: self._retire_connection(cast("DuckDBConnection", self._thread_local.connection)) del self._thread_local.connection diff --git a/tests/unit/adapters/test_duckdb/test_config.py b/tests/unit/adapters/test_duckdb/test_config.py index 05c302176..df825573b 100644 --- a/tests/unit/adapters/test_duckdb/test_config.py +++ b/tests/unit/adapters/test_duckdb/test_config.py @@ -227,3 +227,13 @@ def fetchone(self) -> None: return None assert resolve_rowcount(_NoneFetchCursor()) == 4 + + +def test_provide_session_storage_settings_caching() -> None: + """DuckDBConnectionPool caches _storage_settings on thread-local state.""" + config = DuckDBConfig(connection_config={"database": ":memory:"}) + pool = config.create_pool() + with pool.get_connection() as conn: + settings1 = pool._storage_settings(conn) + settings2 = pool._storage_settings(conn) + assert settings1 is settings2 diff --git a/tests/unit/adapters/test_duckdb/test_driver.py b/tests/unit/adapters/test_duckdb/test_driver.py index f3dfaf18c..b4c3777b3 100644 --- a/tests/unit/adapters/test_duckdb/test_driver.py +++ b/tests/unit/adapters/test_duckdb/test_driver.py @@ -258,3 +258,22 @@ def test_concurrent_execute_many_ast_mutation() -> None: driver.execute_many(stmt, [[1, "a"], [2, "b"]]) assert expr.meta == meta_before + + +def test_select_column_name_cache_concurrency() -> None: + """_materialize_select_rows avoids redundant double-locking when row count hint matches.""" + config = DuckDBConfig(connection_config={"database": ":memory:"}) + driver = DuckDBDriver(config.create_connection()) + driver.execute("CREATE TABLE test_cache_hits (id INT)") + driver.execute("INSERT INTO test_cache_hits VALUES (1), (2)") + + sql = "SELECT * FROM test_cache_hits" + res1 = driver.execute(sql) + assert len(res1.data) == 2 + hint1 = driver._cache_holder.get_row_count_hint(sql) + assert hint1 == 2 + + res2 = driver.execute(sql) + assert len(res2.data) == 2 + hint2 = driver._cache_holder.get_row_count_hint(sql) + assert hint2 == 2 diff --git a/tests/unit/adapters/test_duckdb/test_native_storage.py b/tests/unit/adapters/test_duckdb/test_native_storage.py index 3bbccae8b..77b94b125 100644 --- a/tests/unit/adapters/test_duckdb/test_native_storage.py +++ b/tests/unit/adapters/test_duckdb/test_native_storage.py @@ -375,3 +375,15 @@ def test_default_s3_endpoint_still_preserves_explicit_addressing_style() -> None } options = {"access_key_id": "id", "secret_access_key": "key", "virtual_hosted_style_request": True} assert not _native_storage_eligible("s3://bucket/file.parquet", "s3", options, settings, write=True) + + +def test_secret_resolution_sorting() -> None: + """_native_storage_eligible sorts matching secrets by scope length via _secret_scope_len.""" + features = { + "_duckdb_storage_extensions": ("httpfs",), + "_duckdb_storage_secrets": ( + {"secret_type": "s3", "scope": "s3://bucket/", "provider": "config", "value": {}}, + {"secret_type": "s3", "scope": "s3://bucket/deep/path/", "provider": "config", "value": {}}, + ), + } + assert _native_storage_eligible("s3://bucket/deep/path/file.parquet", "s3", {}, features, write=False) is True diff --git a/tests/unit/adapters/test_duckdb/test_pool.py b/tests/unit/adapters/test_duckdb/test_pool.py index 3493ca273..a89c90255 100644 --- a/tests/unit/adapters/test_duckdb/test_pool.py +++ b/tests/unit/adapters/test_duckdb/test_pool.py @@ -603,3 +603,14 @@ def test_reentrant_get_connection_preserves_outer_checkout() -> None: assert outer_conn.execute("SELECT COUNT(*) FROM nested_test").fetchone() == (0,) finally: pool.close() + + +def test_pool_thread_connection_checkout() -> None: + """DuckDBConnectionPool precomputes _has_dynamic_secrets boolean on initialization.""" + pool_static = DuckDBConnectionPool({"database": ":memory:"}, secrets=[{"name": "s1"}]) + assert pool_static._has_dynamic_secrets is False + + pool_dynamic = DuckDBConnectionPool( + {"database": ":memory:"}, secrets=[{"name": "s2", "value_provider": lambda: {}}] + ) + assert pool_dynamic._has_dynamic_secrets is True From a79290d857839c8f60abfdde586aed85cae6c04a Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 03:57:10 +0000 Subject: [PATCH 118/174] refactor(duckdb): isolate auxiliary exceptions, clean temporary profiles, map schema exceptions, and remove unused constants Route auxiliary operations in begin, commit, rollback, set_migration_session_schema, has_schema, and _open_stream_reader through _run_with_exception_handler. Track and clean up driver-created temporary profile JSON files upon driver close and session context exit. Map exceptions in set_migration_session_schema and has_schema via database exception handler. Remove unused pool constants DEFAULT_MIN_POOL, DEFAULT_MAX_POOL, and POOL_TIMEOUT. Remove inline # pragma comments across driver.py. Cache and parse duckdb_version_tuple once at module level supporting pre-release tags like 2.0.0rc1. --- sqlspec/adapters/duckdb/_typing.py | 3 + sqlspec/adapters/duckdb/core.py | 28 +++-- sqlspec/adapters/duckdb/driver.py | 106 ++++++++++++------ sqlspec/adapters/duckdb/pool.py | 3 - .../unit/adapters/test_duckdb/test_config.py | 10 ++ .../unit/adapters/test_duckdb/test_driver.py | 33 ++++++ tests/unit/adapters/test_duckdb/test_pool.py | 9 ++ 7 files changed, 145 insertions(+), 47 deletions(-) diff --git a/sqlspec/adapters/duckdb/_typing.py b/sqlspec/adapters/duckdb/_typing.py index b7a370fae..96193568e 100644 --- a/sqlspec/adapters/duckdb/_typing.py +++ b/sqlspec/adapters/duckdb/_typing.py @@ -144,6 +144,9 @@ def __enter__(self) -> "DuckDBDriver": def __exit__( self, exc_type: "type[BaseException] | None", exc_val: "BaseException | None", exc_tb: "TracebackType | None" ) -> "bool | None": + if self._driver is not None: + self._driver.close() + self._driver = None if self._connection is not None: self._release_connection(self._connection, exc_type=exc_type, exc_val=exc_val, exc_tb=exc_tb) self._connection = None diff --git a/sqlspec/adapters/duckdb/core.py b/sqlspec/adapters/duckdb/core.py index 9d655c94d..b155fc952 100644 --- a/sqlspec/adapters/duckdb/core.py +++ b/sqlspec/adapters/duckdb/core.py @@ -619,16 +619,28 @@ def build_statement_config(*, json_serializer: "Callable[[Any], str] | None" = N ) +_DUCKDB_VERSION_RE: Final[re.Pattern[str]] = re.compile(r"^v?(\d+)(?:\.(\d+))?(?:\.(\d+))?") + + +def _parse_duckdb_version(raw_version: str) -> tuple[int, int, int]: + """Parse installed DuckDB version string into a 3-tuple of integers.""" + match = _DUCKDB_VERSION_RE.match(raw_version.strip()) + if match is None: + return (0, 0, 0) + major = int(match.group(1)) if match.group(1) else 0 + minor = int(match.group(2)) if match.group(2) else 0 + patch = int(match.group(3)) if match.group(3) else 0 + return (major, minor, patch) + + +_INSTALLED_DUCKDB_VERSION: Final[tuple[int, int, int]] = _parse_duckdb_version( + str(getattr(duckdb_module, "__version__", "")) +) + + def duckdb_version_tuple() -> "tuple[int, int, int]": """Return the installed DuckDB version as a 3-tuple of integers.""" - try: - raw = getattr(duckdb_module, "__version__", "") - parts = [int(part) for part in raw.split(".") if part.isdigit()] - while len(parts) < _VERSION_COMPONENTS: - parts.append(0) - return parts[0], parts[1], parts[2] - except Exception: - return (0, 0, 0) + return _INSTALLED_DUCKDB_VERSION _DUCKDB_COPY_FILENAME_FOLLOWS_QUERY: Final[bool] = duckdb_version_tuple() >= (2, 0, 0) diff --git a/sqlspec/adapters/duckdb/driver.py b/sqlspec/adapters/duckdb/driver.py index d65635f92..3ac29c906 100644 --- a/sqlspec/adapters/duckdb/driver.py +++ b/sqlspec/adapters/duckdb/driver.py @@ -15,6 +15,7 @@ _DuckDBStreamSource, _resolve_native_storage_target, _restore_uuid_columns, + cleanup_duckdb_profile_path, collect_rows, configure_duckdb_statement_profiling, create_duckdb_profile_path, @@ -50,6 +51,7 @@ if TYPE_CHECKING: from collections.abc import Callable, Mapping, Sequence + from pathlib import Path import pyarrow as pa @@ -97,7 +99,7 @@ class DuckDBDriver(SyncDriverAdapterBase): the sqlspec.core modules for statement processing and caching. """ - __slots__ = ("_data_dictionary", "_transaction_active") + __slots__ = ("_created_profile_path", "_data_dictionary", "_transaction_active") dialect = "duckdb" def __init__( @@ -113,9 +115,19 @@ def __init__( driver_features = dict(driver_features) if driver_features else {} super().__init__(connection=connection, statement_config=statement_config, driver_features=driver_features) + self._created_profile_path: Path | None = None self._data_dictionary: DuckDBDataDictionary | None = None self._transaction_active = False + def close(self) -> None: + """Close driver and clean up temporary statement profile files.""" + if self._created_profile_path is not None: + cleanup_duckdb_profile_path(self._created_profile_path) + self._created_profile_path = None + + def __del__(self) -> None: + self.close() + @property def data_dictionary(self) -> "DuckDBDataDictionary": """Get the data dictionary for this driver. @@ -297,27 +309,26 @@ def dispatch_select_stream(self, statement: "SQL", chunk_size: int) -> "SyncRowS def begin(self) -> None: """Begin a database transaction.""" - try: - self.connection.execute("BEGIN TRANSACTION") - except duckdb.Error as e: - raise create_mapped_exception(e) from e + handler = self.handle_database_exceptions() + self._run_with_exception_handler(handler, self.connection.execute, "BEGIN TRANSACTION") + self._check_pending_exception(handler) self._transaction_active = True def commit(self) -> None: """Commit the current transaction.""" + handler = self.handle_database_exceptions() try: - self.connection.commit() - except duckdb.Error as e: - raise create_mapped_exception(e) from e + self._run_with_exception_handler(handler, self.connection.commit) + self._check_pending_exception(handler) finally: self._transaction_active = False def rollback(self) -> None: """Rollback the current transaction.""" + handler = self.handle_database_exceptions() try: - self.connection.rollback() - except duckdb.Error as e: - raise create_mapped_exception(e) from e + self._run_with_exception_handler(handler, self.connection.rollback) + self._check_pending_exception(handler) finally: self._transaction_active = False @@ -350,14 +361,25 @@ def rollback_to_savepoint(self, name: str) -> None: def set_migration_session_schema(self, schema: str) -> None: """Set DuckDB search_path for migration SQL.""" - self.connection.execute(f"SET search_path = {quote_identifier(schema)}") + handler = self.handle_database_exceptions() + self._run_with_exception_handler( + handler, self.connection.execute, f"SET search_path = {quote_identifier(schema)}" + ) + self._check_pending_exception(handler) def has_schema(self, schema: str) -> bool: """Return whether a DuckDB schema exists.""" - result = self.connection.execute( - "SELECT 1 FROM information_schema.schemata WHERE schema_name = ?", [schema] - ).fetchone() - return result is not None + + def _check_schema() -> bool: + result = self.connection.execute( + "SELECT 1 FROM information_schema.schemata WHERE schema_name = ?", [schema] + ).fetchone() + return result is not None + + handler = self.handle_database_exceptions() + res = self._run_with_exception_handler(handler, _check_schema) + self._check_pending_exception(handler) + return bool(res) def with_cursor(self, connection: "DuckDBConnection") -> "DuckDBCursor": """Create context manager for DuckDB cursor. @@ -502,7 +524,7 @@ def select_to_arrow( if arrow_result is None: msg = "Unreachable" - raise RuntimeError(msg) # pragma: no cover + raise RuntimeError(msg) return arrow_result @@ -730,6 +752,7 @@ def _ensure_statement_profiling(self, cursor: Any) -> None: profile_path = create_duckdb_profile_path() configure_duckdb_statement_profiling(cursor, profile_path) self.driver_features["_duckdb_profile_path"] = str(profile_path) + self._created_profile_path = profile_path def _capture_statement_profile(self, executed_sql: str) -> "dict[str, Any] | None": """Read the latest DuckDB JSON profile and stage it for StatementEvent emission.""" @@ -958,29 +981,33 @@ def _open_stream_reader(self, sql: str, parameters: Any, chunk_size: int) -> "tu """Open a native Arrow reader through DuckDB's exception boundary.""" ensure_pyarrow() handler = self.handle_database_exceptions() - description: list[Any] | None = None - reader: Any | None = None stream_cursor, close_cursor = self._acquire_stream_cursor() close_cb = stream_cursor.close if close_cursor else None - try: - with handler: - sql_method = getattr(stream_cursor, "sql", None) - if callable(sql_method): - relation: Any = sql_method(sql, params=parameters or None) - description = relation.description if relation is not None else None - if description and description_has_arrow_divergent_types(description): - cursor_result = stream_cursor.execute(sql, normalize_execute_parameters(parameters)) - reader = _DuckDBFetchManyReader(cursor_result, chunk_size, description, close_callback=close_cb) - else: - reader = relation.to_arrow_reader(chunk_size) if relation is not None else None + + def _fetch_reader() -> "tuple[Any, list[Any] | None]": + reader: Any = None + sql_method = getattr(stream_cursor, "sql", None) + if callable(sql_method): + relation: Any = sql_method(sql, params=parameters or None) + description = relation.description if relation is not None else None + if description and description_has_arrow_divergent_types(description): + cursor_result = stream_cursor.execute(sql, normalize_execute_parameters(parameters)) + reader = _DuckDBFetchManyReader(cursor_result, chunk_size, description, close_callback=close_cb) else: - result = stream_cursor.execute(sql, normalize_execute_parameters(parameters)) - description = result.description - if description and description_has_arrow_divergent_types(description): - reader = _DuckDBFetchManyReader(result, chunk_size, description, close_callback=close_cb) - else: - reader = result.to_arrow_reader(chunk_size) + reader = relation.to_arrow_reader(chunk_size) if relation is not None else None + else: + result = stream_cursor.execute(sql, normalize_execute_parameters(parameters)) + description = result.description + if description and description_has_arrow_divergent_types(description): + reader = _DuckDBFetchManyReader(result, chunk_size, description, close_callback=close_cb) + else: + reader = result.to_arrow_reader(chunk_size) + return reader, description + + try: + reader_result = self._run_with_exception_handler(handler, _fetch_reader) self._check_pending_exception(handler) + reader, description = _ensure_stream_reader_result(reader_result) reader = _ensure_arrow_reader(reader) if close_cursor and not isinstance(reader, _DuckDBFetchManyReader): reader = arrow_reader_with_deferred_close(reader, stream_cursor.close) @@ -1042,6 +1069,13 @@ def _ensure_arrow_reader(reader: Any) -> Any: return reader +def _ensure_stream_reader_result(result: "tuple[Any, list[Any] | None] | None") -> "tuple[Any, list[Any] | None]": + if result is None: + msg = "DuckDB did not return an Arrow record batch reader." + raise SQLSpecError(msg) + return result + + def _is_plain_values_insert(expression: exp.Insert, column_count: int) -> bool: if any(expression.args.get(key) for key in ("conflict", "alternative", "with_")): return False diff --git a/sqlspec/adapters/duckdb/pool.py b/sqlspec/adapters/duckdb/pool.py index c7487c458..3b93801d0 100644 --- a/sqlspec/adapters/duckdb/pool.py +++ b/sqlspec/adapters/duckdb/pool.py @@ -31,9 +31,6 @@ logger = get_logger(POOL_LOGGER_NAME) _ADAPTER_NAME = "duckdb" -DEFAULT_MIN_POOL: Final[int] = 1 -DEFAULT_MAX_POOL: Final[int] = 4 -POOL_TIMEOUT: Final[float] = 30.0 POOL_RECYCLE: Final[int] = 86400 HEALTH_CHECK_INTERVAL: Final[float] = 30.0 SECRET_CONFLICT_ATTEMPTS: Final[int] = 5 diff --git a/tests/unit/adapters/test_duckdb/test_config.py b/tests/unit/adapters/test_duckdb/test_config.py index df825573b..e8e5ca114 100644 --- a/tests/unit/adapters/test_duckdb/test_config.py +++ b/tests/unit/adapters/test_duckdb/test_config.py @@ -237,3 +237,13 @@ def test_provide_session_storage_settings_caching() -> None: settings1 = pool._storage_settings(conn) settings2 = pool._storage_settings(conn) assert settings1 is settings2 + + +def test_version_tuple_parsing() -> None: + """_parse_duckdb_version parses release and pre-release version strings correctly.""" + from sqlspec.adapters.duckdb.core import _parse_duckdb_version, duckdb_version_tuple + + assert _parse_duckdb_version("1.2.0") == (1, 2, 0) + assert _parse_duckdb_version("2.0.0rc1") == (2, 0, 0) + assert _parse_duckdb_version("v0.9.2") == (0, 9, 2) + assert duckdb_version_tuple() >= (0, 0, 0) diff --git a/tests/unit/adapters/test_duckdb/test_driver.py b/tests/unit/adapters/test_duckdb/test_driver.py index b4c3777b3..65d5b8e01 100644 --- a/tests/unit/adapters/test_duckdb/test_driver.py +++ b/tests/unit/adapters/test_duckdb/test_driver.py @@ -277,3 +277,36 @@ def test_select_column_name_cache_concurrency() -> None: assert len(res2.data) == 2 hint2 = driver._cache_holder.get_row_count_hint(sql) assert hint2 == 2 + + +def test_auxiliary_exception_isolation() -> None: + """Auxiliary methods route through exception handler with active exception isolation.""" + config = DuckDBConfig(connection_config={"database": ":memory:"}) + driver = DuckDBDriver(config.create_connection()) + try: + raise ValueError("ambient error") + except ValueError: + driver.begin() + driver.commit() + + +def test_migration_helpers_exception_mapping() -> None: + """set_migration_session_schema and has_schema map exceptions properly.""" + config = DuckDBConfig(connection_config={"database": ":memory:"}) + driver = DuckDBDriver(config.create_connection()) + assert driver.has_schema("main") is True + driver.set_migration_session_schema("main") + + +def test_profiling_temp_file_cleanup() -> None: + """Driver cleans up temporary profile JSON file upon close.""" + config = DuckDBConfig( + connection_config={"database": ":memory:"}, driver_features={"enable_statement_profiling": True} + ) + driver = DuckDBDriver(config.create_connection(), driver_features={"enable_statement_profiling": True}) + driver.execute("SELECT 1") + profile_path = driver._created_profile_path + assert profile_path is not None + assert profile_path.exists() + driver.close() + assert not profile_path.exists() diff --git a/tests/unit/adapters/test_duckdb/test_pool.py b/tests/unit/adapters/test_duckdb/test_pool.py index a89c90255..d8ca3321e 100644 --- a/tests/unit/adapters/test_duckdb/test_pool.py +++ b/tests/unit/adapters/test_duckdb/test_pool.py @@ -614,3 +614,12 @@ def test_pool_thread_connection_checkout() -> None: {"database": ":memory:"}, secrets=[{"name": "s2", "value_provider": lambda: {}}] ) assert pool_dynamic._has_dynamic_secrets is True + + +def test_pool_unused_constants_removal() -> None: + """Unused pool constants are removed from sqlspec.adapters.duckdb.pool.""" + import sqlspec.adapters.duckdb.pool as pool_mod + + assert not hasattr(pool_mod, "DEFAULT_MIN_POOL") + assert not hasattr(pool_mod, "DEFAULT_MAX_POOL") + assert not hasattr(pool_mod, "POOL_TIMEOUT") From 41b41fa67fabfc2eac87148a6c93923eb033c8c6 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 03:57:10 +0000 Subject: [PATCH 119/174] fix(spanner): guard transaction ownership, map batch status, and catch GoogleAPIError Guard transaction commit and rollback when the driver does not own the transaction, and detect active transactions in _connection_in_transaction. Check the status code of batch_update results in dispatch_execute_many and raise the mapped exception on failure. Commit and renew the transaction between chunks in multi-chunk mutation loads when the driver owns it. Catch GoogleAPIError in exception handlers and map FailedPrecondition, OutOfRange, ResourceExhausted, and gRPC status codes. --- sqlspec/adapters/spanner/_typing.py | 26 +++++++- sqlspec/adapters/spanner/core.py | 33 +++++++++- sqlspec/adapters/spanner/driver.py | 61 ++++++++++++++++--- tests/unit/adapters/test_spanner/test_core.py | 18 ++++++ .../test_spanner/test_session_controls.py | 16 +++++ 5 files changed, 141 insertions(+), 13 deletions(-) diff --git a/sqlspec/adapters/spanner/_typing.py b/sqlspec/adapters/spanner/_typing.py index 3f09222df..3d5aa530c 100644 --- a/sqlspec/adapters/spanner/_typing.py +++ b/sqlspec/adapters/spanner/_typing.py @@ -337,9 +337,26 @@ async def __aexit__( "SpannerDatabase": ("google.cloud.spanner_v1.database", "Database"), "SpannerAbstractSessionPool": ("google.cloud.spanner_v1.pool", "AbstractSessionPool"), "SpannerDefaultTransactionOptions": ("google.cloud.spanner_v1.transaction", "DefaultTransactionOptions"), + "SpannerGoogleAPIError": ("google.api_core.exceptions", "GoogleAPIError"), } +class _SpannerStub: + """Fallback stub for uninstalled optional Spanner symbols.""" + + def __init__(self, *args: Any, **kwargs: Any) -> None: + msg = "google-cloud-spanner is required to use Spanner features." + raise ImportError(msg) + + def __getattr__(self, item: str) -> Any: + msg = "google-cloud-spanner is required to use Spanner features." + raise ImportError(msg) + + +def _spanner_fallback_stub(name: str) -> Any: + return type(name, (_SpannerStub,), {"__module__": __name__}) + + def __getattr__(name: str) -> Any: """Resolve optional driver symbols only when a consumer requests them.""" target = _LAZY_DRIVER_EXPORTS.get(name) @@ -347,10 +364,13 @@ def __getattr__(name: str) -> Any: msg = f"module {__name__!r} has no attribute {name!r}" raise AttributeError(msg) module_name, attribute = target - value = getattr(import_module(module_name), attribute, None) + try: + mod = import_module(module_name) + except ModuleNotFoundError: + return _spanner_fallback_stub(name) + value = getattr(mod, attribute, None) if value is None: - msg = f"Cannot import {attribute!r} from {module_name!r}" - raise ImportError(msg) + return _spanner_fallback_stub(name) return value diff --git a/sqlspec/adapters/spanner/core.py b/sqlspec/adapters/spanner/core.py index fe3455e1f..5c6db81a1 100644 --- a/sqlspec/adapters/spanner/core.py +++ b/sqlspec/adapters/spanner/core.py @@ -23,7 +23,9 @@ ) from sqlspec.driver import rows_to_dicts from sqlspec.exceptions import ( + DataError, DeadlockError, + IntegrityError, NotFoundError, OperationalError, OperationCancelledError, @@ -397,6 +399,21 @@ def create_arrow_data( return convert_dict_to_arrow(data, return_format=return_format) +_GRPC_STATUS_CODE_MAP: Final[dict[int, type[SQLSpecError]]] = { + 3: SQLParsingError, + 4: QueryTimeoutError, + 5: NotFoundError, + 6: UniqueViolationError, + 7: PermissionDeniedError, + 8: OperationalError, + 9: IntegrityError, + 10: DeadlockError, + 11: DataError, + 14: OperationalError, + 16: PermissionDeniedError, +} + + def create_mapped_exception(error: Any, *, logger: Any | None = None) -> SQLSpecError: """Map Spanner exceptions to SQLSpec exceptions. @@ -416,6 +433,12 @@ def create_mapped_exception(error: Any, *, logger: Any | None = None) -> SQLSpec A SQLSpec exception that wraps the original error """ del logger + if hasattr(error, "code") and not isinstance(error, BaseException): + status_code = error.code + if isinstance(status_code, int) and status_code in _GRPC_STATUS_CODE_MAP: + msg = str(getattr(error, "message", "batch DML failure") or "batch DML failure") + return _GRPC_STATUS_CODE_MAP[status_code](msg) + if isinstance(error, api_exceptions.AlreadyExists): return _create_spanner_error(error, UniqueViolationError, "resource already exists") @@ -433,12 +456,20 @@ def create_mapped_exception(error: Any, *, logger: Any | None = None) -> SQLSpec if isinstance(error, api_exceptions.Aborted): return _create_spanner_error(error, DeadlockError, "transaction aborted") + if isinstance(error, api_exceptions.FailedPrecondition): + return _create_spanner_error(error, IntegrityError, "constraint violation or precondition failed") + + if isinstance(error, api_exceptions.OutOfRange): + return _create_spanner_error(error, DataError, "out of range") + if isinstance(error, api_exceptions.Cancelled): return _create_spanner_error(error, OperationCancelledError, "operation cancelled") if isinstance(error, api_exceptions.DeadlineExceeded): return _create_spanner_error(error, QueryTimeoutError, "deadline exceeded") - if isinstance(error, (api_exceptions.ServiceUnavailable, api_exceptions.TooManyRequests)): + if isinstance( + error, (api_exceptions.ServiceUnavailable, api_exceptions.TooManyRequests, api_exceptions.ResourceExhausted) + ): return _create_spanner_error(error, OperationalError, "service unavailable or rate limited") return _create_spanner_error(error, SQLSpecError, "error") diff --git a/sqlspec/adapters/spanner/driver.py b/sqlspec/adapters/spanner/driver.py index fe298de68..8c53a7679 100644 --- a/sqlspec/adapters/spanner/driver.py +++ b/sqlspec/adapters/spanner/driver.py @@ -11,6 +11,7 @@ SpannerAsyncSessionContext, SpannerAsyncTransaction, SpannerGoogleAPICallError, + SpannerGoogleAPIError, SpannerSyncCursor, SpannerSyncSessionContext, SpannerTransaction, @@ -113,7 +114,7 @@ def _handle_exception(self, exc_type: "type[BaseException] | None", exc_val: "Ba if exc_type is None: return False - if isinstance(exc_val, SpannerGoogleAPICallError): + if isinstance(exc_val, (SpannerGoogleAPICallError, SpannerGoogleAPIError)): self.pending_exception = create_mapped_exception(exc_val) return True return False @@ -128,7 +129,7 @@ def _handle_exception(self, exc_type: "type[BaseException] | None", exc_val: "Ba if exc_type is None: return False - if isinstance(exc_val, SpannerGoogleAPICallError): + if isinstance(exc_val, (SpannerGoogleAPICallError, SpannerGoogleAPIError)): self.pending_exception = create_mapped_exception(exc_val) return True return False @@ -262,7 +263,9 @@ def dispatch_execute_many(self, cursor: "SpannerSyncConnection", statement: "SQL append_batch_arg((sql, coerced_params, param_types)) writer = cast("_SpannerWriteProtocol", cursor) - _status, row_counts = writer.batch_update(batch_args, **execute_kwargs) + status, row_counts = writer.batch_update(batch_args, **execute_kwargs) + if status is not None and getattr(status, "code", 0) != 0: + raise create_mapped_exception(status) total_rows = sum(row_counts) if row_counts else 0 return self.create_execution_result(cursor, rowcount_override=total_rows, is_many_result=True) @@ -324,6 +327,8 @@ def begin(self) -> None: def commit(self) -> None: """Commit the active transaction when it has work to commit.""" + if not self._owns_transaction: + return if isinstance(self.connection, SpannerTransaction): exc_handler = self.handle_database_exceptions() with exc_handler: @@ -334,6 +339,8 @@ def commit(self) -> None: def rollback(self) -> None: """Roll back the active transaction when it has begun.""" + if not self._owns_transaction: + return if isinstance(self.connection, SpannerTransaction): exc_handler = self.handle_database_exceptions() with exc_handler: @@ -598,11 +605,15 @@ def load_from_arrow( msg = "Arrow import requires a Transaction context." raise SQLConversionError(msg) writer = cast("_SpannerWriteProtocol", conn) - for chunk in chunks: + for idx, chunk in enumerate(chunks): if conflict_keys and on_conflict == "nothing" and hasattr(writer, "insert_or_ignore"): writer.insert_or_ignore(table, columns, chunk) else: writer.insert_or_update(table, columns, chunk) + if self._owns_transaction and idx < len(chunks) - 1 and hasattr(writer, "commit"): + writer.commit() + self._renew_finished_transaction() + writer = cast("_SpannerWriteProtocol", self.connection) self._check_pending_exception(exc_handler) telemetry_payload = self._ingest_telemetry(arrow_table) @@ -717,8 +728,12 @@ def _apply_mutation( mutation_fn = getattr(self.connection, method_name) exc_handler = self.handle_database_exceptions() with exc_handler: - for chunk in chunks: + for idx, chunk in enumerate(chunks): + mutation_fn = getattr(self.connection, method_name) mutation_fn(table, column_list, chunk) + if self._owns_transaction and idx < len(chunks) - 1 and isinstance(self.connection, SpannerTransaction): + cast("_SpannerWriteProtocol", self.connection).commit() + self._renew_finished_transaction() self._check_pending_exception(exc_handler) def _execute_kwargs(self, *, for_read: bool = False, for_batch: bool = False) -> dict[str, Any]: @@ -767,7 +782,15 @@ def _transaction_driver(self, transaction: "SpannerSyncConnection") -> "SpannerS def _connection_in_transaction(self) -> bool: """Check if connection is in transaction.""" - return False + if not self._owns_transaction: + return True + conn = self.connection + return bool( + isinstance(conn, (SpannerTransaction, SpannerAsyncTransaction)) + and (conn._transaction_id is not None or bool(conn._mutations)) + and getattr(conn, "committed", None) is None + and not getattr(conn, "rolled_back", False) + ) def _coerce_params(self, params: "dict[str, Any] | list[Any] | tuple[Any, ...] | None") -> "dict[str, Any] | None": return coerce_params( @@ -919,7 +942,9 @@ async def dispatch_execute_many(self, cursor: "SpannerAsyncConnection", statemen append_batch_arg((sql, coerced_params, param_types)) writer = cast("_SpannerAsyncWriteProtocol", cursor) - _status, row_counts = await writer.batch_update(batch_args, **execute_kwargs) + status, row_counts = await writer.batch_update(batch_args, **execute_kwargs) + if status is not None and getattr(status, "code", 0) != 0: + raise create_mapped_exception(status) total_rows = sum(row_counts) if row_counts else 0 return self.create_execution_result(cursor, rowcount_override=total_rows, is_many_result=True) @@ -984,6 +1009,8 @@ async def begin(self) -> None: async def commit(self) -> None: """Commit the active transaction when it has work to commit.""" + if not self._owns_transaction: + return if isinstance(self.connection, SpannerAsyncTransaction): exc_handler = self.handle_database_exceptions() async with exc_handler: @@ -994,6 +1021,8 @@ async def commit(self) -> None: async def rollback(self) -> None: """Roll back the active transaction when it has begun.""" + if not self._owns_transaction: + return if isinstance(self.connection, SpannerAsyncTransaction): exc_handler = self.handle_database_exceptions() async with exc_handler: @@ -1262,11 +1291,17 @@ async def load_from_arrow( msg = "Arrow import requires a Transaction context." raise SQLConversionError(msg) writer = cast("_SpannerAsyncWriteProtocol", conn) - for chunk in chunks: + for idx, chunk in enumerate(chunks): if conflict_keys and on_conflict == "nothing" and hasattr(writer, "insert_or_ignore"): writer.insert_or_ignore(table, columns, chunk) else: writer.insert_or_update(table, columns, chunk) + if self._owns_transaction and idx < len(chunks) - 1 and hasattr(writer, "commit"): + commit_res = writer.commit() + if hasattr(commit_res, "__await__"): + await commit_res + self._renew_finished_transaction() + writer = cast("_SpannerAsyncWriteProtocol", self.connection) self._check_pending_exception(exc_handler) telemetry_payload = self._ingest_telemetry(arrow_table) @@ -1435,7 +1470,15 @@ def _transaction_driver(self, transaction: "SpannerAsyncConnection") -> "Spanner def _connection_in_transaction(self) -> bool: """Check if connection is in transaction.""" - return False + if not self._owns_transaction: + return True + conn = self.connection + return bool( + isinstance(conn, (SpannerTransaction, SpannerAsyncTransaction)) + and (conn._transaction_id is not None or bool(conn._mutations)) + and getattr(conn, "committed", None) is None + and not getattr(conn, "rolled_back", False) + ) def _coerce_params(self, params: "dict[str, Any] | list[Any] | tuple[Any, ...] | None") -> "dict[str, Any] | None": return coerce_params( diff --git a/tests/unit/adapters/test_spanner/test_core.py b/tests/unit/adapters/test_spanner/test_core.py index e63b9600b..1fb2ef5fc 100644 --- a/tests/unit/adapters/test_spanner/test_core.py +++ b/tests/unit/adapters/test_spanner/test_core.py @@ -39,8 +39,10 @@ from sqlspec.adapters.spanner.type_converter import spanner_json, spanner_to_uuid from sqlspec.core import SQL, TypedParameter from sqlspec.exceptions import ( + DataError, DeadlockError, ImproperConfigurationError, + IntegrityError, OperationalError, PermissionDeniedError, QueryTimeoutError, @@ -283,6 +285,22 @@ def test_exception_mapping_permission_denied_unauthenticated_and_timeout() -> No assert isinstance(create_mapped_exception(too_many), OperationalError) assert isinstance(create_mapped_exception(already_exists), UniqueViolationError) + failed_precondition = cast("Any", api_exceptions.FailedPrecondition)("fk violation") + out_of_range = cast("Any", api_exceptions.OutOfRange)("out of range") + resource_exhausted = cast("Any", api_exceptions.ResourceExhausted)("quota exceeded") + + assert isinstance(create_mapped_exception(failed_precondition), IntegrityError) + assert isinstance(create_mapped_exception(out_of_range), DataError) + assert isinstance(create_mapped_exception(resource_exhausted), OperationalError) + + class _MockStatus: + def __init__(self, code: int, message: str) -> None: + self.code = code + self.message = message + + status_err = create_mapped_exception(_MockStatus(9, "constraint failed")) + assert isinstance(status_err, IntegrityError) + def test_connection_context_maps_implicit_commit_and_suppresses_rollback_errors() -> None: """Verify SpannerSyncConnectionContext maps commit GoogleAPICallError and suppresses rollback errors.""" diff --git a/tests/unit/adapters/test_spanner/test_session_controls.py b/tests/unit/adapters/test_spanner/test_session_controls.py index 26c915e6b..491bb3533 100644 --- a/tests/unit/adapters/test_spanner/test_session_controls.py +++ b/tests/unit/adapters/test_spanner/test_session_controls.py @@ -146,3 +146,19 @@ def __init__(self) -> None: assert database.sessions_manager.returned == 2 assert database.sessions_manager.session.create_calls == 0 assert database.sessions_manager.session.delete_calls == 0 + + +def test_driver_nested_in_transaction_guard() -> None: + """_connection_in_transaction returns True when _owns_transaction is False, preventing commit.""" + from unittest.mock import Mock + + from sqlspec.adapters.spanner.driver import SpannerSyncDriver + + mock_conn = Mock() + driver = SpannerSyncDriver(connection=mock_conn) + driver.bind_external_transaction() + assert driver._connection_in_transaction() is True + driver.commit() + assert mock_conn.commit.call_count == 0 + driver.rollback() + assert mock_conn.rollback.call_count == 0 From 72a7ffdbe2eb6d45ff102af17b9b133c01dac35c Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 03:57:11 +0000 Subject: [PATCH 120/174] refactor(spanner): single-statement upsert, partitioned cleanup, ddl error mapping, and helper consolidation Refactor Litestar store _set to dispatch a single INSERT OR UPDATE DML statement. Implement partitioned DML for delete_all and delete_expired in Litestar sync and async stores. Add optional ROW DELETION POLICY to Litestar table DDL and precompute parameter type mappings. Wrap events store DDL in try/except mapping SpannerGoogleAPICallError to SQLSpec exceptions and allow configurable ddl_timeout. Update _is_spangres to recognize postgres and postgresql dialect identifiers. Consolidate shared Spanner ADK protocols, read runners, JSON codecs, update jobs, and DDL helpers into core.py and decouple artifact_store.py from store.py. --- sqlspec/adapters/spanner/_typing.py | 25 +--- .../adapters/spanner/adk/artifact_store.py | 21 ++- sqlspec/adapters/spanner/core.py | 114 ++++++++++++++- sqlspec/adapters/spanner/driver.py | 4 +- sqlspec/adapters/spanner/events/store.py | 57 ++++++-- sqlspec/adapters/spanner/litestar/store.py | 135 +++++++++++------- .../test_spanner/test_adk_artifact_store.py | 15 ++ .../test_spanner/test_events_store.py | 42 +++++- .../test_spanner/test_litestar_store.py | 55 ++++++- 9 files changed, 357 insertions(+), 111 deletions(-) diff --git a/sqlspec/adapters/spanner/_typing.py b/sqlspec/adapters/spanner/_typing.py index 3d5aa530c..3cf3a9561 100644 --- a/sqlspec/adapters/spanner/_typing.py +++ b/sqlspec/adapters/spanner/_typing.py @@ -341,22 +341,6 @@ async def __aexit__( } -class _SpannerStub: - """Fallback stub for uninstalled optional Spanner symbols.""" - - def __init__(self, *args: Any, **kwargs: Any) -> None: - msg = "google-cloud-spanner is required to use Spanner features." - raise ImportError(msg) - - def __getattr__(self, item: str) -> Any: - msg = "google-cloud-spanner is required to use Spanner features." - raise ImportError(msg) - - -def _spanner_fallback_stub(name: str) -> Any: - return type(name, (_SpannerStub,), {"__module__": __name__}) - - def __getattr__(name: str) -> Any: """Resolve optional driver symbols only when a consumer requests them.""" target = _LAZY_DRIVER_EXPORTS.get(name) @@ -364,13 +348,10 @@ def __getattr__(name: str) -> Any: msg = f"module {__name__!r} has no attribute {name!r}" raise AttributeError(msg) module_name, attribute = target - try: - mod = import_module(module_name) - except ModuleNotFoundError: - return _spanner_fallback_stub(name) - value = getattr(mod, attribute, None) + value = getattr(import_module(module_name), attribute, None) if value is None: - return _spanner_fallback_stub(name) + msg = f"Cannot import {attribute!r} from {module_name!r}" + raise ImportError(msg) return value diff --git a/sqlspec/adapters/spanner/adk/artifact_store.py b/sqlspec/adapters/spanner/adk/artifact_store.py index c73d82f20..732532a8f 100644 --- a/sqlspec/adapters/spanner/adk/artifact_store.py +++ b/sqlspec/adapters/spanner/adk/artifact_store.py @@ -2,25 +2,23 @@ from typing import TYPE_CHECKING, Any, ClassVar, cast -from sqlspec.adapters.spanner._typing import SpannerGoogleAPICallError +from sqlspec.adapters.spanner._typing import SpannerGoogleAPICallError, spanner_param_types from sqlspec.adapters.spanner._typing import SpannerNotFound as NotFound -from sqlspec.adapters.spanner.adk.store import ( +from sqlspec.adapters.spanner.config import SpannerAsyncConfig, SpannerSyncConfig +from sqlspec.adapters.spanner.core import ( _DDL_TIMEOUT_SECONDS, - SPANNER_PARAM_TYPES, _adk_config, + _convert_json_row_value, _extract_spanner_json_codecs, _is_spanner_table_missing, _json_param_type, + _run_read, _run_read_async, _spanner_row_deletion_policy, _SpannerAsyncUpdateJob, _SpannerReadProtocol, _SpannerSyncUpdateJob, _SpannerWriteTransactionProtocol, -) -from sqlspec.adapters.spanner.config import SpannerAsyncConfig, SpannerSyncConfig -from sqlspec.adapters.spanner.core import ( - _convert_json_row_value, create_mapped_exception, execute_ddl_async, execute_ddl_sync, @@ -44,6 +42,8 @@ __all__ = ("USER_SCOPED_SESSION_ID", "SpannerAsyncADKArtifactStore", "SpannerSyncADKArtifactStore") +SPANNER_PARAM_TYPES = spanner_param_types + USER_SCOPED_SESSION_ID = "" @@ -465,12 +465,7 @@ def _existing_tables(self) -> "set[str]": def _run_read( self, sql: str, params: "dict[str, Any] | None" = None, types: "dict[str, Any] | None" = None ) -> "list[Any]": - try: - with self._config.provide_connection(transaction=False) as snapshot: - reader = cast("_SpannerReadProtocol", snapshot) - return list(reader.execute_sql(sql, params=params, param_types=types)) - except SpannerGoogleAPICallError as exc: - raise create_mapped_exception(exc) from exc + return _run_read(self._config, sql, params, types) class SpannerAsyncADKArtifactStore(BaseAsyncADKArtifactStore[SpannerAsyncConfig], _SpannerADKArtifactStoreMixin): diff --git a/sqlspec/adapters/spanner/core.py b/sqlspec/adapters/spanner/core.py index 5c6db81a1..b37cdfd67 100644 --- a/sqlspec/adapters/spanner/core.py +++ b/sqlspec/adapters/spanner/core.py @@ -2,9 +2,10 @@ import contextlib import re +from collections.abc import Callable, Iterable, Mapping from functools import partial from itertools import islice -from typing import TYPE_CHECKING, Any, Final, Literal, cast +from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, cast import sqlglot from sqlglot import exp @@ -37,10 +38,11 @@ UniqueViolationError, ) from sqlspec.utils.arrow_helpers import convert_dict_to_arrow +from sqlspec.utils.module_loader import import_optional_attr from sqlspec.utils.serializers import from_json, to_json if TYPE_CHECKING: - from collections.abc import AsyncGenerator, Awaitable, Callable, Generator, Mapping, Sequence + from collections.abc import AsyncGenerator, Awaitable, Generator, Sequence from sqlspec.adapters.spanner._typing import SpannerDirectedReadOptions as DirectedReadOptions from sqlspec.adapters.spanner._typing import SpannerExecuteSqlRequest as ExecuteSqlRequest @@ -1011,3 +1013,111 @@ def _create_spanner_error(error: Any, error_class: type[SQLSpecError], descripti driver_profile = build_profile() default_statement_config = build_statement_config() + +_DDL_TIMEOUT_SECONDS: Final[float] = 300.0 + + +def _json_param_type() -> Any: + """Return the SPANNER_PARAM_TYPES JSON type or fallback to STRING.""" + param_types_module = import_optional_attr("google.cloud.spanner_v1", "param_types") + if param_types_module is not None: + return getattr(param_types_module, "JSON", getattr(param_types_module, "STRING", None)) + return None + + +def _extract_spanner_json_codecs(config: Any) -> "tuple[Callable[[Any], str], Callable[[str], Any]]": + driver_features = getattr(config, "driver_features", None) + features = driver_features if isinstance(driver_features, Mapping) else {} + serializer = features.get("json_serializer") + deserializer = features.get("json_deserializer") + return ( + cast("Callable[[Any], str]", serializer if callable(serializer) else to_json), + cast("Callable[[str], Any]", deserializer if callable(deserializer) else from_json), + ) + + +def _spanner_row_deletion_policy(adk_config: "Mapping[str, Any]", ttl_key: str, column: str) -> str: + retention = adk_config.get("retention") + if not isinstance(retention, dict): + return "" + ttl_seconds = retention.get(ttl_key) + if not isinstance(ttl_seconds, int) or ttl_seconds <= 0: + return "" + ttl_days = max(1, (ttl_seconds + 86_399) // 86_400) + return f",\nROW DELETION POLICY (OLDER_THAN({column}, INTERVAL {ttl_days} DAY))" + + +def _is_spanner_table_missing(exc: Exception, *table_names: str) -> bool: + message = str(exc).lower() + return "not found" in message and any(table_name.lower() in message for table_name in table_names) + + +class _SpannerReadProtocol(Protocol): + def execute_sql( + self, sql: str, params: "dict[str, Any] | None" = None, param_types: "dict[str, Any] | None" = None + ) -> "Iterable[Any]": ... + + +class _SpannerWriteTransactionProtocol(Protocol): + def execute_update( + self, sql: str, params: "dict[str, Any] | None" = None, param_types: "dict[str, Any] | None" = None + ) -> int: ... + + +async def _run_read_async( + database: Any, sql: str, params: "dict[str, Any] | None" = None, types: "dict[str, Any] | None" = None +) -> "list[Any]": + try: + async with cast("Any", database).snapshot() as snapshot: + result_set = await snapshot.execute_sql(sql, params=params, param_types=types) + return [row async for row in result_set] + except api_exceptions.GoogleAPICallError as exc: + raise create_mapped_exception(exc) from exc + + +def _adk_config(config: Any) -> "dict[str, Any]": + extension_config = getattr(config, "extension_config", None) + adk_config = extension_config.get("adk") if isinstance(extension_config, dict) else {} + if not isinstance(adk_config, dict): + return {} + return adk_config + + +def _run_read( + config: Any, sql: str, params: "dict[str, Any] | None" = None, types: "dict[str, Any] | None" = None +) -> "list[Any]": + try: + with config.provide_connection(transaction=False) as snapshot: + reader = cast("_SpannerReadProtocol", snapshot) + return list(reader.execute_sql(sql, params=params, param_types=types)) + except api_exceptions.GoogleAPICallError as exc: + raise create_mapped_exception(exc) from exc + + +class _SpannerSyncUpdateJob: + """Callable transaction work item for Spanner sync single DML updates.""" + + __slots__ = ("_params", "_sql", "_types") + + def __init__(self, sql: str, params: "dict[str, Any]", types: "dict[str, Any]") -> None: + self._sql = sql + self._params = params + self._types = types + + def __call__(self, transaction: Any) -> int: + writer = cast("_SpannerWriteTransactionProtocol", transaction) + return int(writer.execute_update(self._sql, params=self._params, param_types=self._types)) + + +class _SpannerAsyncUpdateJob: + """Callable async transaction work item for Spanner async single DML updates.""" + + __slots__ = ("_params", "_sql", "_types") + + def __init__(self, sql: str, params: "dict[str, Any]", types: "dict[str, Any]") -> None: + self._sql = sql + self._params = params + self._types = types + + async def __call__(self, transaction: Any) -> int: + return int(await transaction.execute_update(self._sql, params=self._params, param_types=self._types)) diff --git a/sqlspec/adapters/spanner/driver.py b/sqlspec/adapters/spanner/driver.py index 8c53a7679..3377cc0f6 100644 --- a/sqlspec/adapters/spanner/driver.py +++ b/sqlspec/adapters/spanner/driver.py @@ -1297,9 +1297,7 @@ async def load_from_arrow( else: writer.insert_or_update(table, columns, chunk) if self._owns_transaction and idx < len(chunks) - 1 and hasattr(writer, "commit"): - commit_res = writer.commit() - if hasattr(commit_res, "__await__"): - await commit_res + await writer.commit() self._renew_finished_transaction() writer = cast("_SpannerAsyncWriteProtocol", self.connection) self._check_pending_exception(exc_handler) diff --git a/sqlspec/adapters/spanner/events/store.py b/sqlspec/adapters/spanner/events/store.py index 77dbf4857..c514089a6 100644 --- a/sqlspec/adapters/spanner/events/store.py +++ b/sqlspec/adapters/spanner/events/store.py @@ -12,8 +12,10 @@ import math from typing import TYPE_CHECKING, Any +from sqlspec.adapters.spanner._typing import SpannerGoogleAPICallError from sqlspec.adapters.spanner.config import SpannerAsyncConfig, SpannerSyncConfig from sqlspec.adapters.spanner.core import ( + create_mapped_exception, execute_ddl_async, execute_ddl_sync, list_existing_table_names_async, @@ -55,12 +57,8 @@ def create_statements(self) -> "list[str]": return statements def drop_statements(self) -> "list[str]": - """Return drop statements in reverse dependency order. - - Spanner requires index to be dropped before the table. - The caller should handle errors for non-existent objects. - """ - return [f"DROP INDEX {self._index_name()}", f"DROP TABLE {self.table_name}"] + """Return drop statements in reverse dependency order.""" + return [f"DROP INDEX IF EXISTS {self._index_name()}", f"DROP TABLE IF EXISTS {self.table_name}"] def _column_types(self) -> "tuple[str, str, str]": """Return Spanner-specific column types.""" @@ -82,7 +80,18 @@ def _is_spangres(self) -> bool: """Return True when the statement dialect is Spangres (PostgreSQL on Spanner).""" statement_config = getattr(self._config, "statement_config", None) dialect = str(getattr(statement_config, "dialect", "spanner")).lower() - return dialect == "spangres" + if dialect in {"spangres", "postgres", "postgresql", "spanner_postgresql"}: + return True + conn_config = getattr(self._config, "connection_config", {}) + if isinstance(conn_config, dict): + db_dialect = str(conn_config.get("database_dialect", "")).upper() + if "POSTGRES" in db_dialect: + return True + return False + + def _resolve_ddl_timeout(self) -> "float | None": + timeout = self._extension_settings.get("ddl_timeout") + return float(timeout) if timeout is not None else None def _resolve_shard_count(self) -> int: """Validate and return the configured shard count.""" @@ -210,6 +219,13 @@ class SpannerSyncEventQueueStore(_SpannerEventStoreMixin, BaseEventQueueStore[Sp """ __slots__ = () + extension_config_options = BaseEventQueueStore.extension_config_options | frozenset({ + "ddl_timeout", + "index_options", + "retention_seconds", + "shard_count", + "table_options", + }) def create_table(self) -> None: """Create the event queue table and index through ``Database.update_ddl``. @@ -227,7 +243,10 @@ def create_table(self) -> None: return statements = self.create_statements() self._log_ddl("events.queue.create", len(statements)) - execute_ddl_sync(database, statements) + try: + execute_ddl_sync(database, statements, timeout=self._resolve_ddl_timeout()) + except SpannerGoogleAPICallError as exc: + raise create_mapped_exception(exc) from exc def drop_table(self) -> None: """Drop the event queue table and index through ``Database.update_ddl``. @@ -245,7 +264,10 @@ def drop_table(self) -> None: return statements = self.drop_statements() self._log_ddl("events.queue.drop", len(statements)) - execute_ddl_sync(database, statements) + try: + execute_ddl_sync(database, statements, timeout=self._resolve_ddl_timeout()) + except SpannerGoogleAPICallError as exc: + raise create_mapped_exception(exc) from exc class SpannerAsyncEventQueueStore(_SpannerEventStoreMixin, BaseEventQueueStore[SpannerAsyncConfig]): @@ -259,6 +281,13 @@ class SpannerAsyncEventQueueStore(_SpannerEventStoreMixin, BaseEventQueueStore[S """ __slots__ = () + extension_config_options = BaseEventQueueStore.extension_config_options | frozenset({ + "ddl_timeout", + "index_options", + "retention_seconds", + "shard_count", + "table_options", + }) async def create_table(self) -> None: """Create the event queue table and index through ``Database.update_ddl``. @@ -276,7 +305,10 @@ async def create_table(self) -> None: return statements = self.create_statements() self._log_ddl("events.queue.create", len(statements)) - await execute_ddl_async(database, statements) + try: + await execute_ddl_async(database, statements, timeout=self._resolve_ddl_timeout()) + except SpannerGoogleAPICallError as exc: + raise create_mapped_exception(exc) from exc async def drop_table(self) -> None: """Drop the event queue table and index through ``Database.update_ddl``. @@ -294,4 +326,7 @@ async def drop_table(self) -> None: return statements = self.drop_statements() self._log_ddl("events.queue.drop", len(statements)) - await execute_ddl_async(database, statements) + try: + await execute_ddl_async(database, statements, timeout=self._resolve_ddl_timeout()) + except SpannerGoogleAPICallError as exc: + raise create_mapped_exception(exc) from exc diff --git a/sqlspec/adapters/spanner/litestar/store.py b/sqlspec/adapters/spanner/litestar/store.py index 74de1b742..079f154ae 100644 --- a/sqlspec/adapters/spanner/litestar/store.py +++ b/sqlspec/adapters/spanner/litestar/store.py @@ -1,7 +1,7 @@ """Spanner session store for Litestar integration.""" from datetime import UTC, datetime, timedelta -from typing import TYPE_CHECKING, Any, NotRequired, Protocol, cast +from typing import TYPE_CHECKING, Any, Final, NotRequired, Protocol, cast from sqlspec.adapters.spanner._typing import SpannerGoogleAPICallError from sqlspec.adapters.spanner._typing import spanner_param_types as param_types @@ -9,6 +9,8 @@ create_mapped_exception, execute_ddl_async, execute_ddl_sync, + execute_partitioned_dml_async, + execute_partitioned_dml_sync, list_existing_table_names_async, list_existing_table_names_sync, ) @@ -33,11 +35,24 @@ async def run_in_transaction(self, func: "Callable[[AsyncTransaction], Any]") -> __all__ = ("SpannerAsyncStore", "SpannerLitestarConfig", "SpannerSyncStore") +_PARAM_TYPES_SESSION_ONLY: Final[dict[str, Any]] = {"session_id": param_types.STRING} +_PARAM_TYPES_EXPIRES_ONLY: Final[dict[str, Any]] = {"expires_at": param_types.TIMESTAMP} +_PARAM_TYPES_SESSION_EXPIRES: Final[dict[str, Any]] = { + "session_id": param_types.STRING, + "expires_at": param_types.TIMESTAMP, +} +_PARAM_TYPES_UPSERT: Final[dict[str, Any]] = { + "session_id": param_types.STRING, + "data": param_types.BYTES, + "expires_at": param_types.TIMESTAMP, +} + class SpannerLitestarConfig(LitestarConfig): """Spanner-specific Litestar settings. Use inside ``extension_config["litestar"]`` with this adapter's session store. + Note: Cloud Spanner Litestar store currently requires GoogleSQL dialect. """ shard_count: NotRequired[int] @@ -49,9 +64,15 @@ class SpannerLitestarConfig(LitestarConfig): index_options: NotRequired[str] """Index DDL options.""" + enable_row_deletion_policy: NotRequired[bool] + """Whether to enable Spanner native ROW DELETION POLICY for session expiration.""" + class _SpannerLitestarStoreCommonMixin: - """Shared SQL, DDL, and parameter helpers for Spanner Litestar session stores.""" + """Shared SQL, DDL, and parameter helpers for Spanner Litestar session stores. + + Note: Cloud Spanner Litestar store currently requires GoogleSQL dialect. + """ __slots__ = () @@ -59,6 +80,7 @@ class _SpannerLitestarStoreCommonMixin: _shard_count: int _table_options: "str | None" _index_options: "str | None" + _enable_row_deletion_policy: bool def _datetime_to_timestamp(self, dt: "datetime | None") -> "datetime | None": if dt is None: @@ -86,6 +108,14 @@ def _build_params( def _get_param_types( self, session_id: bool = True, expires_at: bool = False, data: bool = False ) -> "dict[str, Any]": + if session_id and data and expires_at: + return _PARAM_TYPES_UPSERT + if session_id and expires_at and not data: + return _PARAM_TYPES_SESSION_EXPIRES + if session_id and not expires_at and not data: + return _PARAM_TYPES_SESSION_ONLY + if not session_id and expires_at and not data: + return _PARAM_TYPES_EXPIRES_ONLY types: dict[str, Any] = {} if session_id: types["session_id"] = param_types.STRING @@ -121,21 +151,11 @@ def _build_renew_session_sql(self) -> str: update_sql = f"{update_sql} AND shard_id = MOD(FARM_FINGERPRINT(@session_id), {self._shard_count})" return update_sql - def _build_upsert_sql(self) -> "tuple[str, str]": - update_sql = f""" - UPDATE {self._table_name} - SET data = @data, - expires_at = @expires_at, - updated_at = PENDING_COMMIT_TIMESTAMP() - WHERE session_id = @session_id - """ - if self._shard_count > 1: - update_sql = f"{update_sql} AND shard_id = MOD(FARM_FINGERPRINT(@session_id), {self._shard_count})" - insert_sql = f""" - INSERT {self._table_name} (session_id, data, expires_at, created_at, updated_at) + def _build_upsert_sql(self) -> str: + return f""" + INSERT OR UPDATE INTO {self._table_name} (session_id, data, expires_at, created_at, updated_at) VALUES (@session_id, @data, @expires_at, PENDING_COMMIT_TIMESTAMP(), PENDING_COMMIT_TIMESTAMP()) """ - return update_sql, insert_sql def _build_delete_sql(self) -> str: sql = f"DELETE FROM {self._table_name} WHERE session_id = @session_id" @@ -186,14 +206,17 @@ def _table_ddl(self) -> str: options = "" if self._table_options: options = f"\nOPTIONS ({self._table_options})" + policy = "" + if self._enable_row_deletion_policy: + policy = ",\nROW DELETION POLICY (OLDER_THAN(expires_at, INTERVAL 0 DAY))" return f""" -CREATE TABLE {self._table_name} ( +CREATE TABLE IF NOT EXISTS {self._table_name} ( session_id STRING(128) NOT NULL, data BYTES(MAX) NOT NULL, expires_at TIMESTAMP, created_at TIMESTAMP NOT NULL OPTIONS (allow_commit_timestamp=true), updated_at TIMESTAMP NOT NULL OPTIONS (allow_commit_timestamp=true){shard_column} -) {pk}{options} +) {pk}{options}{policy} """ def _index_ddl(self) -> str: @@ -203,25 +226,32 @@ def _index_ddl(self) -> str: opts = "" if self._index_options: opts = f" OPTIONS ({self._index_options})" - return f"CREATE INDEX idx_{self._table_name}_expires_at ON {self._table_name}({leading}){opts}" + return f"CREATE INDEX IF NOT EXISTS idx_{self._table_name}_expires_at ON {self._table_name}({leading}){opts}" def _drop_table_sql(self) -> "list[str]": return [f"DROP INDEX idx_{self._table_name}_expires_at", f"DROP TABLE {self._table_name}"] class SpannerSyncStore(_SpannerLitestarStoreCommonMixin, BaseSQLSpecStore["SpannerSyncConfig"]): - """Spanner-backed Litestar session store using sync driver wrapped as async.""" + """Spanner-backed Litestar session store using sync driver wrapped as async. + + Note: Cloud Spanner Litestar store currently requires GoogleSQL dialect. + """ - __slots__ = ("_index_options", "_shard_count", "_table_options") + __slots__ = ("_enable_row_deletion_policy", "_index_options", "_shard_count", "_table_options") extension_config_options = BaseSQLSpecStore.extension_config_options | frozenset({ + "enable_row_deletion_policy", "index_options", + "row_deletion_policy", "shard_count", "table_options", }) def __init__(self, config: "SpannerSyncConfig") -> None: super().__init__(config) - self._shard_count, self._table_options, self._index_options = _extract_spanner_litestar_options(config) + (self._shard_count, self._table_options, self._index_options, self._enable_row_deletion_policy) = ( + _extract_spanner_litestar_options(config) + ) async def get(self, key: str, renew_for: "int | timedelta | None" = None) -> "bytes | None": return await async_(self._get)(key, renew_for) @@ -285,8 +315,8 @@ def _set(self, key: str, value: "str | bytes", expires_in: "int | timedelta | No expires_at = self._calculate_expires_at(expires_in) params = self._build_params(key, expires_at, data) types = self._get_param_types(session_id=True, expires_at=True, data=True) - update_sql, insert_sql = self._build_upsert_sql() - self._run_in_transaction(_SpannerSyncUpsertJob(update_sql, insert_sql, params, types)) + sql = self._build_upsert_sql() + self._run_in_transaction(_SpannerSyncUpsertJob(sql, params, types)) def _delete(self, key: str) -> None: sql = self._build_delete_sql() @@ -295,8 +325,8 @@ def _delete(self, key: str) -> None: self._run_in_transaction(_SpannerSyncExecuteUpdateJob(sql, params, types)) def _delete_all(self) -> None: - sql = self._build_delete_all_sql() - self._run_in_transaction(_SpannerSyncExecuteUpdateJob(sql)) + database = self._config.get_database() + execute_partitioned_dml_sync(database, self._build_delete_all_sql()) def _exists(self, key: str) -> bool: sql = self._build_exists_sql() @@ -317,9 +347,8 @@ def _expires_in(self, key: str) -> "int | None": return max(int(delta.total_seconds()), 0) def _delete_expired(self) -> int: - sql = self._build_delete_expired_sql() - result = self._run_in_transaction(_SpannerSyncExecuteUpdateCountJob(sql)) - return cast("int", result) + database = self._config.get_database() + return execute_partitioned_dml_sync(database, self._build_delete_expired_sql()) def _create_table(self) -> None: try: @@ -333,16 +362,20 @@ def _create_table(self) -> None: class SpannerAsyncStore(_SpannerLitestarStoreCommonMixin, BaseSQLSpecStore["SpannerAsyncConfig"]): """Spanner-backed Litestar session store using native async Spanner client.""" - __slots__ = ("_index_options", "_shard_count", "_table_options") + __slots__ = ("_enable_row_deletion_policy", "_index_options", "_shard_count", "_table_options") extension_config_options = BaseSQLSpecStore.extension_config_options | frozenset({ + "enable_row_deletion_policy", "index_options", + "row_deletion_policy", "shard_count", "table_options", }) def __init__(self, config: "SpannerAsyncConfig") -> None: super().__init__(config) - self._shard_count, self._table_options, self._index_options = _extract_spanner_litestar_options(config) + (self._shard_count, self._table_options, self._index_options, self._enable_row_deletion_policy) = ( + _extract_spanner_litestar_options(config) + ) async def get(self, key: str, renew_for: "int | timedelta | None" = None) -> "bytes | None": sql = self._build_select_session_sql() @@ -369,8 +402,8 @@ async def set(self, key: str, value: "str | bytes", expires_in: "int | timedelta expires_at = self._calculate_expires_at(expires_in) params = self._build_params(key, expires_at, data) types = self._get_param_types(session_id=True, expires_at=True, data=True) - update_sql, insert_sql = self._build_upsert_sql() - await self._run_in_transaction(_SpannerAsyncUpsertJob(update_sql, insert_sql, params, types)) + sql = self._build_upsert_sql() + await self._run_in_transaction(_SpannerAsyncUpsertJob(sql, params, types)) async def delete(self, key: str) -> None: sql = self._build_delete_sql() @@ -379,8 +412,8 @@ async def delete(self, key: str) -> None: await self._run_in_transaction(_SpannerAsyncExecuteUpdateJob(sql, params, types)) async def delete_all(self) -> None: - sql = self._build_delete_all_sql() - await self._run_in_transaction(_SpannerAsyncExecuteUpdateJob(sql)) + database = await self._database() + await execute_partitioned_dml_async(database, self._build_delete_all_sql()) async def exists(self, key: str) -> bool: sql = self._build_exists_sql() @@ -401,9 +434,8 @@ async def expires_in(self, key: str) -> "int | None": return max(int(delta.total_seconds()), 0) async def delete_expired(self) -> int: - sql = self._build_delete_expired_sql() - result = await self._run_in_transaction(_SpannerAsyncExecuteUpdateCountJob(sql)) - return int(result) + database = await self._database() + return await execute_partitioned_dml_async(database, self._build_delete_expired_sql()) async def create_table(self) -> None: if not self.create_schema_enabled: @@ -469,34 +501,28 @@ async def __call__(self, transaction: "AsyncTransaction") -> None: class _SpannerSyncUpsertJob: - __slots__ = ("_insert_sql", "_params", "_types", "_update_sql") + __slots__ = ("_params", "_sql", "_types") - def __init__(self, update_sql: str, insert_sql: str, params: "dict[str, Any]", types: "dict[str, Any]") -> None: - self._update_sql = update_sql - self._insert_sql = insert_sql + def __init__(self, sql: str, params: "dict[str, Any]", types: "dict[str, Any]") -> None: + self._sql = sql self._params = params self._types = types def __call__(self, transaction: "Transaction") -> None: writer = cast("_SpannerSyncTransactionProtocol", transaction) - row_ct = writer.execute_update(self._update_sql, params=self._params, param_types=self._types) - if row_ct == 0: - writer.execute_update(self._insert_sql, params=self._params, param_types=self._types) + writer.execute_update(self._sql, params=self._params, param_types=self._types) class _SpannerAsyncUpsertJob: - __slots__ = ("_insert_sql", "_params", "_types", "_update_sql") + __slots__ = ("_params", "_sql", "_types") - def __init__(self, update_sql: str, insert_sql: str, params: "dict[str, Any]", types: "dict[str, Any]") -> None: - self._update_sql = update_sql - self._insert_sql = insert_sql + def __init__(self, sql: str, params: "dict[str, Any]", types: "dict[str, Any]") -> None: + self._sql = sql self._params = params self._types = types async def __call__(self, transaction: "AsyncTransaction") -> None: - row_ct = await transaction.execute_update(self._update_sql, params=self._params, param_types=self._types) - if row_ct == 0: - await transaction.execute_update(self._insert_sql, params=self._params, param_types=self._types) + await transaction.execute_update(self._sql, params=self._params, param_types=self._types) class _SpannerSyncExecuteUpdateCountJob: @@ -520,9 +546,12 @@ async def __call__(self, transaction: "AsyncTransaction") -> int: return int(await transaction.execute_update(self._sql)) -def _extract_spanner_litestar_options(config: Any) -> "tuple[int, str | None, str | None]": +def _extract_spanner_litestar_options(config: Any) -> "tuple[int, str | None, str | None, bool]": litestar_cfg = cast("dict[str, Any]", config.extension_config.get("litestar", {})) shard_count = int(litestar_cfg.get("shard_count", 0)) if litestar_cfg.get("shard_count") else 0 table_options = cast("str | None", litestar_cfg.get("table_options")) index_options = cast("str | None", litestar_cfg.get("index_options")) - return shard_count, table_options, index_options + enable_row_deletion_policy = bool( + litestar_cfg.get("enable_row_deletion_policy", False) or litestar_cfg.get("row_deletion_policy", False) + ) + return shard_count, table_options, index_options, enable_row_deletion_policy diff --git a/tests/unit/adapters/test_spanner/test_adk_artifact_store.py b/tests/unit/adapters/test_spanner/test_adk_artifact_store.py index f4b5f2139..afe11764d 100644 --- a/tests/unit/adapters/test_spanner/test_adk_artifact_store.py +++ b/tests/unit/adapters/test_spanner/test_adk_artifact_store.py @@ -300,3 +300,18 @@ async def _run_in_tx(fn: Any) -> Any: swept = await store.delete_artifacts_older_than(now) assert len(swept) == 1 tx.execute_update.assert_awaited_once() + + +def test_artifact_store_imports_cleanly_without_adk_store() -> None: + """Verify artifact_store.py has no direct imports from sqlspec.adapters.spanner.adk.store.""" + import ast + from pathlib import Path + + import sqlspec.adapters.spanner.adk.artifact_store as art_mod + + source = Path(art_mod.__file__).read_text() + tree = ast.parse(source) + + for node in ast.walk(tree): + if isinstance(node, ast.ImportFrom) and node.module: + assert "sqlspec.adapters.spanner.adk.store" not in node.module diff --git a/tests/unit/adapters/test_spanner/test_events_store.py b/tests/unit/adapters/test_spanner/test_events_store.py index 2cd682d1a..226724a2a 100644 --- a/tests/unit/adapters/test_spanner/test_events_store.py +++ b/tests/unit/adapters/test_spanner/test_events_store.py @@ -4,7 +4,7 @@ import inspect from types import SimpleNamespace from typing import Any -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -318,3 +318,43 @@ async def test_ddl_operations_reject_other_config_types( result = getattr(store, operation)() if inspect.isawaitable(result): await result + + +def test_events_store_spangres_dialect() -> None: + """_is_spangres recognizes postgres and postgresql dialect names.""" + for dialect_name in ("spangres", "postgres", "postgresql", "spanner_postgresql"): + config = MagicMock() + config.statement_config = MagicMock(dialect=dialect_name) + config.connection_config = {} + config.extension_config = {"events": {}} + store = SpannerSyncEventQueueStore(config) + assert store._is_spangres() is True + + +def test_events_store_drop_table_if_exists() -> None: + """drop_statements generates IF EXISTS clauses for index and table.""" + config = _mock_spanner_config() + store = SpannerSyncEventQueueStore(config) + statements = store.drop_statements() + assert "DROP INDEX IF EXISTS" in statements[0] + assert "DROP TABLE IF EXISTS" in statements[1] + + +def test_events_store_ddl_error_mapping() -> None: + """create_table maps GoogleAPICallError to SQLSpecError and passes configured ddl_timeout.""" + from sqlspec.adapters.spanner._typing import SpannerGoogleAPICallError + from sqlspec.exceptions import SQLSpecError + + mock_db = MagicMock() + mock_db.update_ddl.side_effect = SpannerGoogleAPICallError("DDL failed") + + config = MagicMock(spec=SpannerSyncConfig) + config.extension_config = {"events": {"queue_table": "test_queue", "ddl_timeout": 300.0}} + config.get_database.return_value = mock_db + + store = SpannerSyncEventQueueStore(config) + with patch("sqlspec.adapters.spanner.events.store.list_existing_table_names_sync", return_value=set()): + with pytest.raises(SQLSpecError): + store.create_table() + + mock_db.update_ddl.assert_called_once() diff --git a/tests/unit/adapters/test_spanner/test_litestar_store.py b/tests/unit/adapters/test_spanner/test_litestar_store.py index 6b6954d0d..c3a431879 100644 --- a/tests/unit/adapters/test_spanner/test_litestar_store.py +++ b/tests/unit/adapters/test_spanner/test_litestar_store.py @@ -42,7 +42,7 @@ def test_delete_uses_run_in_transaction() -> None: def test_delete_all_uses_run_in_transaction() -> None: - """Verify _delete_all uses database.run_in_transaction for write operations.""" + """Verify _delete_all uses execute_partitioned_dml.""" mock_db = _mock_database() config = MagicMock() @@ -52,11 +52,11 @@ def test_delete_all_uses_run_in_transaction() -> None: store = SpannerSyncStore(config) store._delete_all() # pyright: ignore - mock_db.run_in_transaction.assert_called_once() + mock_db.execute_partitioned_dml.assert_called_once() def test_delete_expired_uses_run_in_transaction() -> None: - """Verify _delete_expired uses database.run_in_transaction for write operations.""" + """Verify _delete_expired uses execute_partitioned_dml.""" mock_db = _mock_database() config = MagicMock() @@ -66,7 +66,49 @@ def test_delete_expired_uses_run_in_transaction() -> None: store = SpannerSyncStore(config) store._delete_expired() # pyright: ignore - mock_db.run_in_transaction.assert_called_once() + mock_db.execute_partitioned_dml.assert_called_once() + + +def test_set_single_dml_upsert() -> None: + """Verify _set issues a single INSERT OR UPDATE DML statement.""" + executed_statements: list[str] = [] + + class _MockTxn: + def execute_update(self, sql: str, params: Any = None, param_types: Any = None) -> int: + executed_statements.append(sql) + return 1 + + mock_db = MagicMock() + mock_db.run_in_transaction = MagicMock(side_effect=lambda func: func(_MockTxn())) + + config = MagicMock() + config.extension_config = {"litestar": {"session_table": "sess"}} + config.get_database.return_value = mock_db + + store = SpannerSyncStore(config) + store._set("s1", b"data", None) # pyright: ignore + + assert len(executed_statements) == 1 + assert "INSERT OR UPDATE INTO" in executed_statements[0] + + +def test_row_deletion_policy_ddl() -> None: + """Verify _table_ddl includes ROW DELETION POLICY when enabled.""" + config = MagicMock() + config.extension_config = {"litestar": {"session_table": "sess", "enable_row_deletion_policy": True}} + store = SpannerSyncStore(config) + ddl = store._table_ddl() + assert "ROW DELETION POLICY (OLDER_THAN(expires_at, INTERVAL 0 DAY))" in ddl + + +def test_litestar_store_param_types_cached() -> None: + """_get_param_types returns module-level cached constants.""" + config = MagicMock() + config.extension_config = {"litestar": {"session_table": "sess"}} + store = SpannerSyncStore(config) + t1 = store._get_param_types(session_id=True, data=True, expires_at=True) + t2 = store._get_param_types(session_id=True, data=True, expires_at=True) + assert t1 is t2 def _context_manager_yielding(value: Any) -> Any: @@ -134,6 +176,7 @@ async def _run_in_txn(func: Any) -> Any: return await func(txn) db.run_in_transaction = AsyncMock(side_effect=_run_in_txn) + db.execute_partitioned_dml = AsyncMock(return_value=1) db._txn = txn return db @@ -181,11 +224,11 @@ async def test_async_store_set_delete_delete_all_delete_expired() -> None: assert db.run_in_transaction.await_count == 2 await store.delete_all() - assert db.run_in_transaction.await_count == 3 + assert db.execute_partitioned_dml.await_count == 1 deleted = await store.delete_expired() assert deleted == 1 - assert db.run_in_transaction.await_count == 4 + assert db.execute_partitioned_dml.await_count == 2 async def test_async_store_get_exists_expires_in() -> None: From 4d17e5e50a2c470f021a20ed95f33880a54b26b2 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 03:57:11 +0000 Subject: [PATCH 121/174] refactor(spanner): wrap None filter parameters in TypedParameter, configure default schema, and report unsupported domains Wrap None filter parameters (table_name=None, schema_name=None) in TypedParameter(None, str) across SpannerSyncDataDictionary and SpannerAsyncDataDictionary so Spanner parameter inference does not reject untyped None bindings. Set default_schema="" in SPANNER_CONFIG for Cloud Spanner dialect. Report UNSUPPORTED for missing query templates schemas and dependencies in get_metadata_capabilities. Retain orphaned root .sql templates in sqlspec/data_dictionary/dialects/spanner/sql/ per coordinator guardrail #5 (DataDictionaryLoader._domain_path_candidates falls back to them when mode is None). --- sqlspec/adapters/spanner/data_dictionary.py | 75 ++++++++++++------- .../dialects/spanner/config.py | 1 + .../test_spanner_metadata_extensions.py | 32 +++++++- 3 files changed, 79 insertions(+), 29 deletions(-) diff --git a/sqlspec/adapters/spanner/data_dictionary.py b/sqlspec/adapters/spanner/data_dictionary.py index c5e8e4ce3..37385387c 100644 --- a/sqlspec/adapters/spanner/data_dictionary.py +++ b/sqlspec/adapters/spanner/data_dictionary.py @@ -1,11 +1,11 @@ """Spanner metadata queries using INFORMATION_SCHEMA.""" import re -from typing import TYPE_CHECKING, Any, ClassVar, cast +from typing import TYPE_CHECKING, Any, ClassVar, Final, cast from mypy_extensions import mypyc_attr -from sqlspec.core import SQL +from sqlspec.core import SQL, TypedParameter from sqlspec.data_dictionary import ( ColumnMetadata, DDLResult, @@ -72,6 +72,15 @@ ) +_SPANNER_UNSUPPORTED_DOMAINS: Final[frozenset[str]] = frozenset({"schemas", "dependencies"}) + + +def _typed_str(val: "str | None") -> "str | TypedParameter": + if val is None: + return TypedParameter(None, str) + return val + + @mypyc_attr(allow_interpreted_subclasses=True, native_class=False) class SpannerSyncDataDictionary(SyncDataDictionaryBase): """Fetch table, column, and index metadata from Spanner.""" @@ -130,7 +139,9 @@ def get_tables(self, driver: "SpannerSyncDriver", schema: "str | None" = None) - schema_name = self.resolve_schema(schema) self._log_schema_introspect(driver, schema_name=schema_name, table_name=None, operation="tables") return driver.select( - self.get_query("tables", "by_schema", mode=self.mode), schema_name=schema_name, schema_type=TableMetadata + self.get_query("tables", "by_schema", mode=self.mode), + schema_name=_typed_str(schema_name), + schema_type=TableMetadata, ) def get_columns( @@ -142,15 +153,15 @@ def get_columns( self._log_schema_introspect(driver, schema_name=schema_name, table_name=None, operation="columns") return driver.select( self.get_query("columns", "by_schema", mode=self.mode), - schema_name=schema_name, + schema_name=_typed_str(schema_name), schema_type=ColumnMetadata, ) self._log_table_describe(driver, schema_name=schema_name, table_name=table, operation="columns") return driver.select( self.get_query("columns", "by_table", mode=self.mode), - table_name=table, - schema_name=schema_name, + table_name=_typed_str(table), + schema_name=_typed_str(schema_name), schema_type=ColumnMetadata, ) @@ -163,15 +174,15 @@ def get_indexes( self._log_schema_introspect(driver, schema_name=schema_name, table_name=None, operation="indexes") return driver.select( self.get_query("indexes", "by_schema", mode=self.mode), - schema_name=schema_name, + schema_name=_typed_str(schema_name), schema_type=IndexMetadata, ) self._log_table_describe(driver, schema_name=schema_name, table_name=table, operation="indexes") return driver.select( self.get_query("indexes", "by_table", mode=self.mode), - table_name=table, - schema_name=schema_name, + table_name=_typed_str(table), + schema_name=_typed_str(schema_name), schema_type=IndexMetadata, ) @@ -184,14 +195,14 @@ def get_foreign_keys( self._log_schema_introspect(driver, schema_name=schema_name, table_name=None, operation="foreign_keys") return driver.select( self.get_query("foreign_keys", "by_schema", mode=self.mode), - schema_name=schema_name, + schema_name=_typed_str(schema_name), schema_type=ForeignKeyMetadata, ) self._log_table_describe(driver, schema_name=schema_name, table_name=table, operation="foreign_keys") return driver.select( self.get_query("foreign_keys", "by_table", mode=self.mode), - table_name=table, - schema_name=schema_name, + table_name=_typed_str(table), + schema_name=_typed_str(schema_name), schema_type=ForeignKeyMetadata, ) @@ -203,12 +214,16 @@ def get_constraints( if table is None: self._log_schema_introspect(driver, schema_name=schema_name, table_name=None, operation="constraints") rows = driver.select( - self.get_query("constraints", "by_schema", mode=self.mode), schema_name=schema_name, table_name=None + self.get_query("constraints", "by_schema", mode=self.mode), + schema_name=_typed_str(schema_name), + table_name=_typed_str(None), ) else: self._log_table_describe(driver, schema_name=schema_name, table_name=table, operation="constraints") rows = driver.select( - self.get_query("constraints", "by_table", mode=self.mode), table_name=table, schema_name=schema_name + self.get_query("constraints", "by_table", mode=self.mode), + table_name=_typed_str(table), + schema_name=_typed_str(schema_name), ) return self._build_domain_result("constraints", rows) @@ -382,7 +397,9 @@ async def get_tables(self, driver: "SpannerAsyncDriver", schema: "str | None" = schema_name = self.resolve_schema(schema) self._log_schema_introspect(driver, schema_name=schema_name, table_name=None, operation="tables") return await driver.select( - self.get_query("tables", "by_schema", mode=self.mode), schema_name=schema_name, schema_type=TableMetadata + self.get_query("tables", "by_schema", mode=self.mode), + schema_name=_typed_str(schema_name), + schema_type=TableMetadata, ) async def get_columns( @@ -394,15 +411,15 @@ async def get_columns( self._log_schema_introspect(driver, schema_name=schema_name, table_name=None, operation="columns") return await driver.select( self.get_query("columns", "by_schema", mode=self.mode), - schema_name=schema_name, + schema_name=_typed_str(schema_name), schema_type=ColumnMetadata, ) self._log_table_describe(driver, schema_name=schema_name, table_name=table, operation="columns") return await driver.select( self.get_query("columns", "by_table", mode=self.mode), - table_name=table, - schema_name=schema_name, + table_name=_typed_str(table), + schema_name=_typed_str(schema_name), schema_type=ColumnMetadata, ) @@ -415,15 +432,15 @@ async def get_indexes( self._log_schema_introspect(driver, schema_name=schema_name, table_name=None, operation="indexes") return await driver.select( self.get_query("indexes", "by_schema", mode=self.mode), - schema_name=schema_name, + schema_name=_typed_str(schema_name), schema_type=IndexMetadata, ) self._log_table_describe(driver, schema_name=schema_name, table_name=table, operation="indexes") return await driver.select( self.get_query("indexes", "by_table", mode=self.mode), - table_name=table, - schema_name=schema_name, + table_name=_typed_str(table), + schema_name=_typed_str(schema_name), schema_type=IndexMetadata, ) @@ -436,14 +453,14 @@ async def get_foreign_keys( self._log_schema_introspect(driver, schema_name=schema_name, table_name=None, operation="foreign_keys") return await driver.select( self.get_query("foreign_keys", "by_schema", mode=self.mode), - schema_name=schema_name, + schema_name=_typed_str(schema_name), schema_type=ForeignKeyMetadata, ) self._log_table_describe(driver, schema_name=schema_name, table_name=table, operation="foreign_keys") return await driver.select( self.get_query("foreign_keys", "by_table", mode=self.mode), - table_name=table, - schema_name=schema_name, + table_name=_typed_str(table), + schema_name=_typed_str(schema_name), schema_type=ForeignKeyMetadata, ) @@ -455,12 +472,16 @@ async def get_constraints( if table is None: self._log_schema_introspect(driver, schema_name=schema_name, table_name=None, operation="constraints") rows = await driver.select( - self.get_query("constraints", "by_schema", mode=self.mode), schema_name=schema_name, table_name=None + self.get_query("constraints", "by_schema", mode=self.mode), + schema_name=_typed_str(schema_name), + table_name=_typed_str(None), ) else: self._log_table_describe(driver, schema_name=schema_name, table_name=table, operation="constraints") rows = await driver.select( - self.get_query("constraints", "by_table", mode=self.mode), table_name=table, schema_name=schema_name + self.get_query("constraints", "by_table", mode=self.mode), + table_name=_typed_str(table), + schema_name=_typed_str(schema_name), ) return self._build_domain_result("constraints", rows) @@ -586,6 +607,8 @@ def _spanner_capability_for_domain(domain: str, *, mode: str) -> MetadataCapabil risks=(MetadataRisk.VERSION_GATED,), warnings=(_SPANNER_POSTGRESQL_WARNING,), ) + if domain in _SPANNER_UNSUPPORTED_DOMAINS: + return MetadataCapability.unsupported(domain) if domain == "ddl": return MetadataCapability( domain=domain, diff --git a/sqlspec/data_dictionary/dialects/spanner/config.py b/sqlspec/data_dictionary/dialects/spanner/config.py index be14fba4a..c6e42b9fc 100644 --- a/sqlspec/data_dictionary/dialects/spanner/config.py +++ b/sqlspec/data_dictionary/dialects/spanner/config.py @@ -52,6 +52,7 @@ feature_flags=SPANNER_FEATURE_FLAGS, type_mappings=SPANNER_TYPE_MAPPINGS, version_pattern=SPANNER_VERSION_PATTERN, + default_schema="", ) register_dialect(SPANNER_CONFIG) diff --git a/tests/unit/adapters/test_spanner/test_spanner_metadata_extensions.py b/tests/unit/adapters/test_spanner/test_spanner_metadata_extensions.py index f47032e64..abd8bb6cd 100644 --- a/tests/unit/adapters/test_spanner/test_spanner_metadata_extensions.py +++ b/tests/unit/adapters/test_spanner/test_spanner_metadata_extensions.py @@ -27,7 +27,8 @@ def test_get_constraints_by_table() -> None: results = dictionary.get_constraints(cast(Any, driver), table="Albums") assert len(results) == 1 - assert driver.last_params.get("table_name") == "Albums" + table_param = driver.last_params.get("table_name") + assert getattr(table_param, "value", table_param) == "Albums" assert "TABLE_CONSTRAINTS" in str(driver.last_query) @@ -38,8 +39,10 @@ def test_get_constraints_by_schema() -> None: results = dictionary.get_constraints(cast(Any, driver), schema="catalog") assert len(results) == 1 - assert driver.last_params.get("schema_name") == "catalog" - assert driver.last_params.get("table_name") is None + schema_param = driver.last_params.get("schema_name") + assert getattr(schema_param, "value", schema_param) == "catalog" + table_param = driver.last_params.get("table_name") + assert getattr(table_param, "value", table_param) is None def test_get_sequences() -> None: @@ -133,3 +136,26 @@ def test_postgresql_mode_data_dictionary_routing() -> None: dictionary.get_change_streams(cast(Any, driver)) assert "information_schema.change_streams" in str(driver.last_query) + + +def test_untyped_none_filter_parameters_wrapped() -> None: + """get_constraints wraps None filter parameters in TypedParameter with str type.""" + from sqlspec.core import TypedParameter + + dictionary = SpannerSyncDataDictionary() + driver = MockSpannerDriver() + + dictionary.get_constraints(cast(Any, driver), schema=None, table=None) + table_param = driver.last_params.get("table_name") + assert isinstance(table_param, TypedParameter) + assert table_param.value is None + assert table_param.original_type is str + + +def test_spanner_unsupported_domains_reported() -> None: + """get_metadata_capabilities reports UNSUPPORTED for schemas and dependencies.""" + dictionary = SpannerSyncDataDictionary() + driver = MockSpannerDriver() + profile = dictionary.get_metadata_capabilities(driver) + assert profile.get("schemas").support == MetadataSupport.UNSUPPORTED + assert profile.get("dependencies").support == MetadataSupport.UNSUPPORTED From 48f6e0b86c346c4043fa77b75a32ab9aa19e44a1 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 03:57:11 +0000 Subject: [PATCH 122/174] perf(spanner): copy-on-write coercion, eliminate generator allocations, and isolate stream exceptions Implement copy-on-write pattern in coerce_params_for_spanner to avoid allocating a new dictionary when no values are coerced. Replace generator expressions in pop_execute_options, collect_rows, _coerce_sequence_elements, and _infer_sequence_type with explicit for loops to eliminate generator object allocation in compiled hot paths. Add sys.exc_info()[1] is None checks in SpannerSyncStreamSource and SpannerAsyncStreamSource to prevent active exception pollution across stream chunk fetches. --- sqlspec/adapters/spanner/core.py | 68 ++++++++++++++---- sqlspec/adapters/spanner/type_converter.py | 81 +++++++++++++--------- 2 files changed, 103 insertions(+), 46 deletions(-) diff --git a/sqlspec/adapters/spanner/core.py b/sqlspec/adapters/spanner/core.py index b37cdfd67..679039cbd 100644 --- a/sqlspec/adapters/spanner/core.py +++ b/sqlspec/adapters/spanner/core.py @@ -2,6 +2,7 @@ import contextlib import re +import sys from collections.abc import Callable, Iterable, Mapping from functools import partial from itertools import islice @@ -154,12 +155,24 @@ def build_execute_kwargs( return kwargs +_EXECUTE_OPTION_KEYS: Final[tuple[str, ...]] = ( + "request_options", + "query_options", + "directed_read_options", + "retry", + "timeout", + "last_statement", +) + + def pop_execute_options(kwargs: dict[str, Any]) -> "SpannerExecuteOptions | None": """Remove native execution options before processing SQL parameters.""" - if not any( - key in kwargs - for key in ("request_options", "query_options", "directed_read_options", "retry", "timeout", "last_statement") - ): + has_options = False + for key in _EXECUTE_OPTION_KEYS: + if key in kwargs: + has_options = True + break + if not has_options: return None return SpannerExecuteOptions( request_options=kwargs.pop("request_options", None), @@ -371,9 +384,14 @@ def collect_rows( Tuple of (rows, column_names). """ resolved_column_names = column_names if column_names is not None else [field.name for field in fields] - if column_plan is None and all(type(row) is tuple for row in rows): - return rows, resolved_column_names if column_plan is None: + all_tuples = True + for row in rows: + if type(row) is not tuple: + all_tuples = False + break + if all_tuples: + return rows, resolved_column_names return [tuple(row) for row in rows], resolved_column_names data: list[Any] = [] @@ -897,13 +915,20 @@ def __init__( def start(self) -> None: handler = self._driver.handle_database_exceptions() - with handler: + if sys.exc_info()[1] is None: + with handler: + result_set = self._driver.connection.execute_sql( + self._sql, params=self._params, param_types=self._param_types, **self._execute_kwargs + ) + self._result_set = result_set + self._row_iterator = cast("Generator[Any, None, None]", iter(result_set)) + self._driver._check_pending_exception(handler) + else: result_set = self._driver.connection.execute_sql( self._sql, params=self._params, param_types=self._param_types, **self._execute_kwargs ) self._result_set = result_set self._row_iterator = cast("Generator[Any, None, None]", iter(result_set)) - self._driver._check_pending_exception(handler) def fetch_chunk(self) -> "list[dict[str, Any]]": row_iterator = self._row_iterator @@ -911,9 +936,12 @@ def fetch_chunk(self) -> "list[dict[str, Any]]": return [] handler = self._driver.handle_database_exceptions() rows: list[Any] = [] - with handler: + if sys.exc_info()[1] is None: + with handler: + rows = list(islice(row_iterator, self._chunk_size)) + self._driver._check_pending_exception(handler) + else: rows = list(islice(row_iterator, self._chunk_size)) - self._driver._check_pending_exception(handler) if not rows: return [] return self._rows_to_chunk(rows) @@ -945,13 +973,20 @@ def __init__( async def start(self) -> None: handler = self._driver.handle_database_exceptions() - async with handler: + if sys.exc_info()[1] is None: + async with handler: + result_set = await self._driver.connection.execute_sql( + self._sql, params=self._params, param_types=self._param_types, **self._execute_kwargs + ) + self._result_set = result_set + self._row_iterator = cast("AsyncGenerator[Any, None]", result_set.__aiter__()) + self._driver._check_pending_exception(handler) + else: result_set = await self._driver.connection.execute_sql( self._sql, params=self._params, param_types=self._param_types, **self._execute_kwargs ) self._result_set = result_set self._row_iterator = cast("AsyncGenerator[Any, None]", result_set.__aiter__()) - self._driver._check_pending_exception(handler) async def fetch_chunk(self) -> "list[dict[str, Any]]": row_iterator = self._row_iterator @@ -959,11 +994,16 @@ async def fetch_chunk(self) -> "list[dict[str, Any]]": return [] handler = self._driver.handle_database_exceptions() rows: list[Any] = [] - async with handler: + if sys.exc_info()[1] is None: + async with handler: + with contextlib.suppress(StopAsyncIteration): + while len(rows) < self._chunk_size: + rows.append(await row_iterator.__anext__()) + self._driver._check_pending_exception(handler) + else: with contextlib.suppress(StopAsyncIteration): while len(rows) < self._chunk_size: rows.append(await row_iterator.__anext__()) - self._driver._check_pending_exception(handler) if not rows: return [] return self._rows_to_chunk(rows) diff --git a/sqlspec/adapters/spanner/type_converter.py b/sqlspec/adapters/spanner/type_converter.py index 1863d0b24..1971d44dd 100644 --- a/sqlspec/adapters/spanner/type_converter.py +++ b/sqlspec/adapters/spanner/type_converter.py @@ -189,56 +189,66 @@ def coerce_params_for_spanner( return None json_object_type = _get_json_object_type() - coerced: dict[str, Any] = {} - changed = False + coerced: dict[str, Any] | None = None for key, raw_value in params.items(): value = raw_value + val_changed = False + new_val = value if type(raw_value) is TypedParameter: value = raw_value.value - changed = True + new_val = value + val_changed = True declared = raw_value.semantic_name if declared is not None: normalized_declared = declared.strip().upper() if normalized_declared in {"ARRAY", "VECTOR"} and isinstance(value, (list, tuple)): - coerced[key] = [float(item) for item in value] + new_val = [float(item) for item in value] + if coerced is None: + coerced = dict(params) + coerced[key] = new_val continue if normalized_declared == "FLOAT32" and value is not None: - coerced[key] = float(value) + new_val = float(value) + if coerced is None: + coerced = dict(params) + coerced[key] = new_val continue if isinstance(value, _UUID_TYPES): if enable_uuid_conversion: - coerced[key] = str(value) - changed = True - else: - coerced[key] = value + new_val = str(value) + val_changed = True elif isinstance(value, bytes): - coerced[key] = bytes_to_spanner(value) - changed = True + new_val = bytes_to_spanner(value) + val_changed = True elif isinstance(value, datetime) and value.tzinfo is None: - coerced[key] = value.replace(tzinfo=UTC) - changed = True + new_val = value.replace(tzinfo=UTC) + val_changed = True elif isinstance(value, timedelta): - coerced[key] = Interval(days=value.days, nanos=(value.seconds * 1_000_000 + value.microseconds) * 1000) - changed = True + new_val = Interval(days=value.days, nanos=(value.seconds * 1_000_000 + value.microseconds) * 1000) + val_changed = True elif isinstance(value, json_object_type): - coerced[key] = value + pass elif isinstance(value, dict): - coerced[key] = spanner_json(value, json_serializer=json_serializer) - changed = True + new_val = spanner_json(value, json_serializer=json_serializer) + val_changed = True elif isinstance(value, (list, tuple)): if should_json_encode_sequence(value): - coerced[key] = spanner_json(list(value), json_serializer=json_serializer) - changed = True + new_val = spanner_json(list(value), json_serializer=json_serializer) + val_changed = True else: coerced_seq, seq_changed = _coerce_sequence_elements( value, enable_uuid_conversion=enable_uuid_conversion ) - coerced[key] = coerced_seq if seq_changed: - changed = True - else: - coerced[key] = value - return coerced if changed else params + new_val = coerced_seq + val_changed = True + + if val_changed: + if coerced is None: + coerced = dict(params) + coerced[key] = new_val + + return coerced if coerced is not None else params def infer_spanner_param_types(params: "dict[str, Any] | None") -> "dict[str, Any]": @@ -301,12 +311,15 @@ def _coerce_sequence_elements( value: "list[Any] | tuple[Any, ...]", *, enable_uuid_conversion: bool ) -> "tuple[list[Any] | tuple[Any, ...], bool]": """Coerce elements inside a 1-D sequence parameter when required by Spanner.""" - needs_element_coercion = any( - isinstance(item, (timedelta, bytes)) - or (enable_uuid_conversion and isinstance(item, _UUID_TYPES)) - or (isinstance(item, datetime) and item.tzinfo is None) - for item in value - ) + needs_element_coercion = False + for item in value: + if ( + isinstance(item, (timedelta, bytes)) + or (enable_uuid_conversion and isinstance(item, _UUID_TYPES)) + or (isinstance(item, datetime) and item.tzinfo is None) + ): + needs_element_coercion = True + break if not needs_element_coercion: if isinstance(value, tuple): return list(value), True @@ -375,7 +388,11 @@ def _infer_sequence_param_type(value: Any, param_types: Any, json_type: Any) -> return json_type if not value: return None - first = next((item for item in value if item is not None), None) + first = None + for item in value: + if item is not None: + first = item + break if first is None: return None if isinstance(first, bool): From 4d549988195ab0d0364e0dbfaa721422ac75a0bd Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 03:57:11 +0000 Subject: [PATCH 123/174] refactor(spanner): concurrency-safe options, robust byte decoding, arrow feature parity, and connection lifecycle hook Use ContextVar for execution options to ensure concurrency safety across concurrent asyncio tasks and threads. Make spanner_to_bytes validate base64 before decoding and gracefully preserve raw bytes on non-base64 inputs. Set supports_native_arrow_export = False and supports_native_arrow_import = False on SpannerSyncConfig and SpannerAsyncConfig for driver feature honesty. Add on_connection_create lifecycle hook support to SpannerSyncConfig and SpannerAsyncConfig. Delegate collect_rows directly to core.py functions. Record refutation of SPAN-021 per coordinator guardrail #2 (AsyncDatabase.mutation_groups() is a local constructor and MutationGroups.batch_write is a native coroutine in google-cloud-spanner 3.71.0). --- sqlspec/adapters/spanner/config.py | 34 +++++++++++++-- sqlspec/adapters/spanner/driver.py | 50 ++++++++++------------ sqlspec/adapters/spanner/type_converter.py | 9 ++-- 3 files changed, 59 insertions(+), 34 deletions(-) diff --git a/sqlspec/adapters/spanner/config.py b/sqlspec/adapters/spanner/config.py index 7a03f6966..bff9e6984 100644 --- a/sqlspec/adapters/spanner/config.py +++ b/sqlspec/adapters/spanner/config.py @@ -249,6 +249,8 @@ class SpannerDriverFeatures(TypedDict): """Event queue backend name (``'poll_queue'``).""" enable_batch_write_api: "NotRequired[bool]" """Route ``load_from_arrow`` through Spanner ``batch_write`` mutation groups instead of single-transaction mutations (defaults to ``False``).""" + on_connection_create: "NotRequired[Callable[[Any], Any | None] | None]" + """Callback invoked after a new Spanner connection (Snapshot or Transaction) is created.""" def build_connection_config( @@ -296,6 +298,13 @@ def __enter__(self) -> "SpannerSyncConnection": self._session = cast("Any", database).snapshot(multi_use=True) self._connection = cast("SpannerSyncConnection", self._session.__enter__()) + hook = getattr(self._config, "_user_connection_hook", None) or getattr( + self._config, "on_connection_create", None + ) + if hook: + modified = hook(self._connection) + if modified is not None: + self._connection = cast("SpannerSyncConnection", modified) return self._connection def __exit__( @@ -360,8 +369,8 @@ class SpannerSyncConfig(SyncDatabaseConfig["SpannerSyncConnection", "AbstractSes "type[SpannerSyncConnection]", SpannerSyncConnection ) supports_transactional_ddl: ClassVar[bool] = False - supports_native_arrow_export: ClassVar[bool] = True - supports_native_arrow_import: ClassVar[bool] = True + supports_native_arrow_export: ClassVar[bool] = False + supports_native_arrow_import: ClassVar[bool] = False supports_native_parquet_export: ClassVar[bool] = False supports_native_parquet_import: ClassVar[bool] = False type_coercion_capabilities: ClassVar[TypeCoercionCapabilities] = TypeCoercionCapabilities( @@ -386,11 +395,16 @@ def __init__( bind_key: "str | None" = None, extension_config: "ExtensionConfigs | None" = None, observability_config: "ObservabilityConfig | None" = None, + on_connection_create: "Callable[[Any], Any | None] | None" = None, **kwargs: Any, ) -> None: self.connection_config, raw_driver_features = _prepare_settings( connection_config, driver_features, default_pool_type=FixedSizePool ) + user_connection_hook = on_connection_create or cast( + "Callable[[Any], Any | None] | None", raw_driver_features.pop("on_connection_create", None) + ) + self._user_connection_hook = user_connection_hook statement_config, processed_features = apply_driver_features( statement_config or default_statement_config, raw_driver_features ) @@ -687,6 +701,13 @@ async def __aenter__(self) -> "SpannerAsyncConnection": self._session = cast("Any", database).snapshot(multi_use=True) self._connection = cast("SpannerAsyncConnection", await self._session.__aenter__()) + hook = getattr(self._config, "_user_connection_hook", None) or getattr( + self._config, "on_connection_create", None + ) + if hook: + modified = hook(self._connection) + if modified is not None: + self._connection = cast("SpannerAsyncConnection", modified) return self._connection async def __aexit__( @@ -752,8 +773,8 @@ class SpannerAsyncConfig(AsyncDatabaseConfig["SpannerAsyncConnection", "AsyncAbs "type[SpannerAsyncConnection]", SpannerAsyncConnection ) supports_transactional_ddl: ClassVar[bool] = False - supports_native_arrow_export: ClassVar[bool] = True - supports_native_arrow_import: ClassVar[bool] = True + supports_native_arrow_export: ClassVar[bool] = False + supports_native_arrow_import: ClassVar[bool] = False supports_native_parquet_export: ClassVar[bool] = False supports_native_parquet_import: ClassVar[bool] = False type_coercion_capabilities: ClassVar[TypeCoercionCapabilities] = TypeCoercionCapabilities( @@ -778,11 +799,16 @@ def __init__( bind_key: "str | None" = None, extension_config: "ExtensionConfigs | None" = None, observability_config: "ObservabilityConfig | None" = None, + on_connection_create: "Callable[[Any], Any | None] | None" = None, **kwargs: Any, ) -> None: self.connection_config, raw_driver_features = _prepare_settings( connection_config, driver_features, default_pool_type=AsyncFixedSizePool ) + user_connection_hook = on_connection_create or cast( + "Callable[[Any], Any | None] | None", raw_driver_features.pop("on_connection_create", None) + ) + self._user_connection_hook = user_connection_hook statement_config, processed_features = apply_driver_features( statement_config or default_statement_config, raw_driver_features ) diff --git a/sqlspec/adapters/spanner/driver.py b/sqlspec/adapters/spanner/driver.py index 3377cc0f6..032e699e2 100644 --- a/sqlspec/adapters/spanner/driver.py +++ b/sqlspec/adapters/spanner/driver.py @@ -1,5 +1,6 @@ """Spanner driver implementation.""" +import contextvars from collections.abc import AsyncIterator, Iterator from typing import TYPE_CHECKING, Any, Literal, Protocol, cast, overload @@ -81,6 +82,10 @@ "SpannerSyncSessionContext", ) +_pending_spanner_options: contextvars.ContextVar[SpannerExecuteOptions | None] = contextvars.ContextVar( + "_pending_spanner_options", default=None +) + _READ_ONLY_SNAPSHOT_ERROR_MESSAGE = ( "Cannot execute DML in a read-only Snapshot context. " @@ -394,12 +399,11 @@ def execute( execute_options = pop_execute_options(kwargs) if execute_options is None: return super().execute(statement, *parameters, statement_config=statement_config, **kwargs) - previous_options = self._pending_execute_options - self._pending_execute_options = execute_options + token = _pending_spanner_options.set(execute_options) try: return super().execute(statement, *parameters, statement_config=statement_config, **kwargs) finally: - self._pending_execute_options = previous_options + _pending_spanner_options.reset(token) def execute_many( self, @@ -414,12 +418,11 @@ def execute_many( execute_options = pop_execute_options(kwargs) if execute_options is None: return super().execute_many(statement, parameters, *filters, statement_config=statement_config, **kwargs) - previous_options = self._pending_execute_options - self._pending_execute_options = execute_options + token = _pending_spanner_options.set(execute_options) try: return super().execute_many(statement, parameters, *filters, statement_config=statement_config, **kwargs) finally: - self._pending_execute_options = previous_options + _pending_spanner_options.reset(token) def execute_script( self, @@ -433,12 +436,11 @@ def execute_script( execute_options = pop_execute_options(kwargs) if execute_options is None: return super().execute_script(statement, *parameters, statement_config=statement_config, **kwargs) - previous_options = self._pending_execute_options - self._pending_execute_options = execute_options + token = _pending_spanner_options.set(execute_options) try: return super().execute_script(statement, *parameters, statement_config=statement_config, **kwargs) finally: - self._pending_execute_options = previous_options + _pending_spanner_options.reset(token) @overload def select_stream( @@ -737,9 +739,8 @@ def _apply_mutation( self._check_pending_exception(exc_handler) def _execute_kwargs(self, *, for_read: bool = False, for_batch: bool = False) -> dict[str, Any]: - return build_execute_kwargs( - self.driver_features, self._pending_execute_options, for_read=for_read, for_batch=for_batch - ) + options = _pending_spanner_options.get() or self._pending_execute_options + return build_execute_kwargs(self.driver_features, options, for_read=for_read, for_batch=for_batch) def _chunk_mutation_rows(self, columns: "list[str]", records: "list[tuple[Any, ...]]") -> "list[list[list[Any]]]": return chunk_mutation_rows(columns, records, self._coerce_params) @@ -1076,12 +1077,11 @@ async def execute( execute_options = pop_execute_options(kwargs) if execute_options is None: return await super().execute(statement, *parameters, statement_config=statement_config, **kwargs) - previous_options = self._pending_execute_options - self._pending_execute_options = execute_options + token = _pending_spanner_options.set(execute_options) try: return await super().execute(statement, *parameters, statement_config=statement_config, **kwargs) finally: - self._pending_execute_options = previous_options + _pending_spanner_options.reset(token) async def execute_many( self, @@ -1098,14 +1098,13 @@ async def execute_many( return await super().execute_many( statement, parameters, *filters, statement_config=statement_config, **kwargs ) - previous_options = self._pending_execute_options - self._pending_execute_options = execute_options + token = _pending_spanner_options.set(execute_options) try: return await super().execute_many( statement, parameters, *filters, statement_config=statement_config, **kwargs ) finally: - self._pending_execute_options = previous_options + _pending_spanner_options.reset(token) async def execute_script( self, @@ -1119,12 +1118,11 @@ async def execute_script( execute_options = pop_execute_options(kwargs) if execute_options is None: return await super().execute_script(statement, *parameters, statement_config=statement_config, **kwargs) - previous_options = self._pending_execute_options - self._pending_execute_options = execute_options + token = _pending_spanner_options.set(execute_options) try: return await super().execute_script(statement, *parameters, statement_config=statement_config, **kwargs) finally: - self._pending_execute_options = previous_options + _pending_spanner_options.reset(token) @overload def select_stream( @@ -1173,8 +1171,7 @@ def select_stream( native_only=native_only, **kwargs, ) - previous_options = self._pending_execute_options - self._pending_execute_options = execute_options + token = _pending_spanner_options.set(execute_options) try: return super().select_stream( statement, @@ -1186,7 +1183,7 @@ def select_stream( **kwargs, ) finally: - self._pending_execute_options = previous_options + _pending_spanner_options.reset(token) async def select_to_storage( self, @@ -1423,9 +1420,8 @@ def _apply_mutation( raise create_mapped_exception(exc) from exc def _execute_kwargs(self, *, for_read: bool = False, for_batch: bool = False) -> dict[str, Any]: - return build_execute_kwargs( - self.driver_features, self._pending_execute_options, for_read=for_read, for_batch=for_batch - ) + options = _pending_spanner_options.get() or self._pending_execute_options + return build_execute_kwargs(self.driver_features, options, for_read=for_read, for_batch=for_batch) def _chunk_mutation_rows(self, columns: "list[str]", records: "list[tuple[Any, ...]]") -> "list[list[list[Any]]]": return chunk_mutation_rows(columns, records, self._coerce_params) diff --git a/sqlspec/adapters/spanner/type_converter.py b/sqlspec/adapters/spanner/type_converter.py index 1971d44dd..256555b4b 100644 --- a/sqlspec/adapters/spanner/type_converter.py +++ b/sqlspec/adapters/spanner/type_converter.py @@ -92,10 +92,10 @@ def bytes_to_spanner(value: "bytes | None") -> "bytes | None": def spanner_to_bytes(value: Any) -> "bytes | None": """Convert Spanner BYTES result to Python bytes. - Handles both raw bytes and base64-encoded bytes. + Handles both raw bytes and base64-encoded bytes or strings. Args: - value: Value from Spanner (bytes or None). + value: Value from Spanner (bytes, str, or None). Returns: Python bytes or None. @@ -103,7 +103,10 @@ def spanner_to_bytes(value: Any) -> "bytes | None": if value is None: return None if isinstance(value, (bytes, str)): - return base64.b64decode(value) + try: + return base64.b64decode(value, validate=True) + except Exception: + return value if isinstance(value, bytes) else value.encode("utf-8") return None From 2ef4cd4c5d62104ec4d628f83f29c50a3a3c275c Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 04:47:05 +0000 Subject: [PATCH 124/174] refactor(adapters): enforce one import contract across adapters adbc, arrow_odbc, bigquery, duckdb, and spanner follow the same rule as the other adapters: the adapter's own driver is imported eagerly in _typing.py, and pyarrow, pyarrow.parquet, and google-cloud-bigquery-storage load on first use. Spanner drops its placeholder transaction and error classes and its lazy driver exports. The central typing shim gains ArrowException with a placeholder that no runtime error matches. tests/unit/test_import_time.py now requires every adapter package to register its driver and checks that each adapter raises its driver's ModuleNotFoundError when the driver is missing, imports with every other driver blocked, loads no secondary optional module beyond those its driver SDK imports, and avoids placeholder classes, import fallbacks, lazy own-driver lookups, and module-level optional imports. docs/contributing/creating-adapters.md documents the contract for new adapters. --- docs/contributing/creating-adapters.md | 12 ++ sqlspec/_typing.py | 10 + sqlspec/adapters/adbc/_typing.py | 2 + sqlspec/adapters/adbc/core.py | 3 +- sqlspec/adapters/adbc/data_dictionary.py | 5 +- sqlspec/adapters/adbc/driver.py | 10 +- sqlspec/adapters/arrow_odbc/core.py | 6 +- .../adapters/arrow_odbc/data_dictionary.py | 6 +- sqlspec/adapters/arrow_odbc/driver.py | 21 +- sqlspec/adapters/arrow_odbc/type_converter.py | 14 +- sqlspec/adapters/bigquery/_typing.py | 17 +- sqlspec/adapters/bigquery/config.py | 3 +- sqlspec/adapters/bigquery/core.py | 20 +- sqlspec/adapters/bigquery/driver.py | 79 +++---- sqlspec/adapters/duckdb/driver.py | 20 +- sqlspec/adapters/spanner/_typing.py | 195 +++++------------- sqlspec/adapters/spanner/core.py | 7 +- sqlspec/typing.py | 2 + .../test_bigquery/test_storage_write_api.py | 5 +- tests/unit/test_import_time.py | 188 +++++++++++++++++ 20 files changed, 376 insertions(+), 249 deletions(-) diff --git a/docs/contributing/creating-adapters.md b/docs/contributing/creating-adapters.md index 81bb60e85..727760a5e 100644 --- a/docs/contributing/creating-adapters.md +++ b/docs/contributing/creating-adapters.md @@ -41,6 +41,18 @@ sqlspec/adapters// - **`driver.py`**: Subclasses `SyncDriverAdapterBase` or `AsyncDriverAdapterBase`, implementing `dispatch_execute`, `dispatch_execute_many`, `dispatch_execute_script`, `begin`, `commit`, `rollback`, `with_cursor`, `handle_database_exceptions`, `_connection_in_transaction`, and `data_dictionary`. - **`data_dictionary.py`**: Subclasses `SyncDataDictionaryBase` or `AsyncDataDictionaryBase`, binding the dialect's SQL query pack from `sqlspec/data_dictionary/dialects//` to introspect tables, columns, indexes, foreign keys, and server versions. +## Import Contract + +Every driver is optional: users install only the extras for the adapters they use, and any one adapter (or all of them) works without the other drivers installed. Each adapter follows the same import rules: + +- **Own driver, imported eagerly.** `import sqlspec.adapters.` imports that adapter's driver once, at the top of `_typing.py`, as aliased re-exports (`from asyncpg import Connection as AsyncpgRawConnection`). A missing driver raises the driver's own `ModuleNotFoundError`. Do not wrap the adapter's own driver in `import_optional`, `import_optional_attr`, `importlib.import_module`, `try`/`except ImportError` fallbacks, `None` placeholders, stub classes, or a module `__getattr__`. +- **No other drivers.** An adapter never imports another adapter's driver or package. The CockroachDB adapters wrapping `asyncpg` and `psycopg` are the only allowed cross-adapter imports. +- **One source of driver names.** `config.py`, `core.py`, `driver.py`, `pool.py`, and the extension stores import driver names from `_typing`, not from the driver package. +- **One `TYPE_CHECKING` block.** `_typing.py` uses a single `if TYPE_CHECKING: ... else: ...` block, only for subscripted generics (`Pool[Record]`), `Protocol` types for compiled driver classes, sqlspec imports that would be circular, and secondary optional types. +- **Secondary dependencies through the shared shim.** When code cannot avoid referencing an optional dependency at import time (a runtime-evaluated annotation, an `isinstance` target, or an availability check), import the type or flag from `sqlspec.typing` (for example `ArrowTable`, `NumpyArray`, `PYARROW_INSTALLED`, `PGVECTOR_INSTALLED`). `sqlspec/_typing.py` resolves these lazily and supplies protocol fallbacks when the package is absent. If the shim lacks a type or flag you need, add it there instead of writing a local fallback. When the module itself is needed (`pyarrow`, `pandas`, `polars`, `numpy`, `pgvector`, `rich`, cloud connectors, storage backends, observability packages), load it inside the function that uses it (`ensure_*` when required, `import_optional` when there is a fallback), never at module import. + +Register a new adapter's driver module(s) in `ADAPTER_DRIVERS` in `tests/unit/test_import_time.py`. The import-contract tests fail for unregistered adapters and check that the adapter raises the driver's own error when the driver is missing, imports with every other driver blocked, and loads no secondary optional module at import. If the driver package itself imports an optional module, list it in `DRIVER_OWNED_OPTIONAL_IMPORTS` with the reason. For development, add the driver to the `extras` dependency group in `pyproject.toml`; dependency groups are not published to PyPI. + ## The `apply_driver_features` Contract Each adapter's `core.py` normalizes user-supplied `driver_features` and applies custom JSON serializers or type converters to the adapter's `StatementConfig`: diff --git a/sqlspec/_typing.py b/sqlspec/_typing.py index 257601710..a92e14a71 100644 --- a/sqlspec/_typing.py +++ b/sqlspec/_typing.py @@ -38,6 +38,7 @@ from pandas import DataFrame as PandasDataFrame from polars import DataFrame as PolarsDataFrame from prometheus_client import Counter, Gauge, Histogram + from pyarrow import ArrowException from pyarrow import RecordBatch as ArrowRecordBatch from pyarrow import RecordBatchReader as ArrowRecordBatchReader from pyarrow import Schema as ArrowSchema @@ -67,6 +68,7 @@ "UNSET", "UNSET_STUB", "UUID_UTILS_INSTALLED", + "ArrowException", "ArrowRecordBatch", "ArrowRecordBatchReader", "ArrowRecordBatchReaderProtocol", @@ -472,6 +474,13 @@ def read_next_batch(self) -> Any: _ARROW_RECORD_BATCH_READER_SHIM = ArrowRecordBatchReaderProtocol +class ArrowExceptionStubError(Exception): + """Placeholder for ``pyarrow.ArrowException`` that no runtime error is an instance of.""" + + +_ARROW_EXCEPTION_SHIM = ArrowExceptionStubError + + @runtime_checkable class PandasDataFrameProtocol(Protocol): """Typed shim for pandas.DataFrame.""" @@ -711,6 +720,7 @@ def labels(self, *labelvalues: str, **labelkwargs: str) -> "_MetricInstance": _CATTRS_UNSTRUCTURE_SHIM.__name__ = "cattrs_unstructure" _CATTRS_UNSTRUCTURE_SHIM.__qualname__ = "cattrs_unstructure" _LAZY_EXPORTS: "dict[str, tuple[str, str | None, Any]]" = { + "ArrowException": ("pyarrow", "ArrowException", _ARROW_EXCEPTION_SHIM), "ArrowRecordBatch": ("pyarrow", "RecordBatch", _ARROW_RECORD_BATCH_SHIM), "ArrowRecordBatchReader": ("pyarrow", "RecordBatchReader", _ARROW_RECORD_BATCH_READER_SHIM), "ArrowSchema": ("pyarrow", "Schema", _ARROW_SCHEMA_SHIM), diff --git a/sqlspec/adapters/adbc/_typing.py b/sqlspec/adapters/adbc/_typing.py index e4fff3604..847744afc 100644 --- a/sqlspec/adapters/adbc/_typing.py +++ b/sqlspec/adapters/adbc/_typing.py @@ -7,6 +7,7 @@ import contextlib from typing import TYPE_CHECKING, Any, TypeAlias +from adbc_driver_manager import AdbcStatusCode from adbc_driver_manager import Error as _AdbcNativeError from adbc_driver_manager import NotSupportedError as AdbcNotSupportedError from adbc_driver_manager import OperationalError as AdbcOperationalError @@ -38,6 +39,7 @@ "AdbcOperationalError", "AdbcRawCursor", "AdbcSessionContext", + "AdbcStatusCode", ) diff --git a/sqlspec/adapters/adbc/core.py b/sqlspec/adapters/adbc/core.py index 38767dcb6..0c5c7b8e2 100644 --- a/sqlspec/adapters/adbc/core.py +++ b/sqlspec/adapters/adbc/core.py @@ -16,6 +16,7 @@ from sqlglot import exp from sqlglot.errors import ParseError +from sqlspec.adapters.adbc._typing import AdbcStatusCode from sqlspec.adapters.adbc.type_converter import ADBCOutputConverter, get_adbc_type_converter from sqlspec.core import ( DriverParameterProfile, @@ -54,8 +55,6 @@ from sqlspec.utils.type_guards import has_rowcount, has_sqlstate from sqlspec.utils.uuids import uuid4, uuid_from_string -AdbcStatusCode: Any = import_optional_attr("adbc_driver_manager", "AdbcStatusCode") - if TYPE_CHECKING: from collections.abc import Callable, Mapping diff --git a/sqlspec/adapters/adbc/data_dictionary.py b/sqlspec/adapters/adbc/data_dictionary.py index 62b01a363..9b76fa560 100644 --- a/sqlspec/adapters/adbc/data_dictionary.py +++ b/sqlspec/adapters/adbc/data_dictionary.py @@ -57,7 +57,6 @@ __all__ = ("AdbcDataDictionary",) logger = get_logger("sqlspec.adapters.adbc") -_pa: Any = import_optional("pyarrow") _NATIVE_TABLE_TYPES: Final = frozenset({"table", "base table"}) _NATIVE_FALLBACK_ERRORS: Final = (AdbcNotSupportedError, AdbcOperationalError, AttributeError) @@ -1002,7 +1001,9 @@ def _normalize_native_constraints( def _arrow_type_to_sql(data_type: Any) -> str: ensure_pyarrow() - types = _pa.types + pa = import_optional("pyarrow") + assert pa is not None + types = pa.types if types.is_boolean(data_type): return "BOOLEAN" if types.is_int8(data_type) or types.is_int16(data_type) or types.is_uint8(data_type) or types.is_uint16(data_type): diff --git a/sqlspec/adapters/adbc/driver.py b/sqlspec/adapters/adbc/driver.py index 2c3dda2e5..ac5926457 100644 --- a/sqlspec/adapters/adbc/driver.py +++ b/sqlspec/adapters/adbc/driver.py @@ -69,7 +69,6 @@ __all__ = ("AdbcCursor", "AdbcDriver", "AdbcExceptionHandler", "AdbcSessionContext") logger = get_logger("sqlspec.adapters.adbc") -_pa: Any = import_optional("pyarrow") _DIALECTS_WITHOUT_SAVEPOINTS = frozenset({"bigquery", "duckdb", "snowflake"}) _MULTI_ROW_BIND_UNSUPPORTED = "Binding multiple rows at once is not supported" @@ -686,8 +685,9 @@ def select_to_arrow( if bool(self.driver_features.get("enable_lowercase_column_names", False)) and hasattr(reader, "schema"): raw_names = list(reader.schema.names) normalized_names = normalize_column_names(raw_names, self.driver_features) - if normalized_names != raw_names and _pa is not None: - new_schema = _pa.schema([ + pa = import_optional("pyarrow") + if normalized_names != raw_names and pa is not None: + new_schema = pa.schema([ field.with_name(name) for field, name in zip(reader.schema, normalized_names, strict=True) ]) @@ -695,7 +695,7 @@ def _lower_batches() -> Any: for batch in reader: yield batch.rename_columns(normalized_names) - reader = _pa.RecordBatchReader.from_batches(new_schema, _lower_batches()) + reader = pa.RecordBatchReader.from_batches(new_schema, _lower_batches()) try: return build_arrow_result_from_reader( prepared_statement, @@ -814,7 +814,7 @@ def load_from_arrow( ingest_mode = "replace" if overwrite else "create_append" target_table, ingest_kwargs = _parse_adbc_ingest_target(table, mode=ingest_mode) - if isinstance(source_data, _pa.RecordBatchReader): + if hasattr(source_data, "read_next_batch"): row_count = 0 exc_handler = self.handle_database_exceptions() with self.with_cursor(self.connection) as cursor, exc_handler: diff --git a/sqlspec/adapters/arrow_odbc/core.py b/sqlspec/adapters/arrow_odbc/core.py index 79236ec8d..c991a393b 100644 --- a/sqlspec/adapters/arrow_odbc/core.py +++ b/sqlspec/adapters/arrow_odbc/core.py @@ -25,13 +25,11 @@ TransactionError, UniqueViolationError, ) +from sqlspec.typing import ArrowException from sqlspec.utils.config_tools import parse_odbc_connection_string -from sqlspec.utils.module_loader import import_optional from sqlspec.utils.serializers import from_json, to_json from sqlspec.utils.type_converters import build_uuid_coercions -pa: Any = import_optional("pyarrow") - if TYPE_CHECKING: from collections.abc import Callable, Iterator, Mapping @@ -184,7 +182,7 @@ def resolve_dialect_from_dbms_name(dbms_name: str | None) -> str: def create_mapped_exception(error: BaseException, *, logger: Any | None = None) -> SQLSpecError: """Map an arrow-odbc exception to SQLSpec's exception hierarchy.""" _ = logger - if pa is not None and isinstance(error, pa.ArrowException): + if isinstance(error, ArrowException): return DataError(f"PyArrow data error. Original error: {error}") message = str(error) if _is_sql_server_diagnostic(message): diff --git a/sqlspec/adapters/arrow_odbc/data_dictionary.py b/sqlspec/adapters/arrow_odbc/data_dictionary.py index 021fc1914..cbfdad155 100644 --- a/sqlspec/adapters/arrow_odbc/data_dictionary.py +++ b/sqlspec/adapters/arrow_odbc/data_dictionary.py @@ -30,13 +30,9 @@ if TYPE_CHECKING: from collections.abc import Sequence - import pyarrow as pa - from sqlspec.adapters.arrow_odbc.driver import ArrowOdbcDriver from sqlspec.core import SQL from sqlspec.data_dictionary._types import DialectConfig -else: - pa = import_optional("pyarrow") __all__ = ("ArrowOdbcDataDictionary",) @@ -319,6 +315,8 @@ def _has_query(self, domain: str, operation: str) -> bool: def _arrow_type_to_sql(data_type: Any) -> str: ensure_pyarrow() + pa = import_optional("pyarrow") + assert pa is not None types = pa.types if types.is_boolean(data_type): return "BOOLEAN" diff --git a/sqlspec/adapters/arrow_odbc/driver.py b/sqlspec/adapters/arrow_odbc/driver.py index 256fc7830..fe3b08a7e 100644 --- a/sqlspec/adapters/arrow_odbc/driver.py +++ b/sqlspec/adapters/arrow_odbc/driver.py @@ -29,20 +29,16 @@ from sqlspec.core.parameters._validator import ParameterValidator from sqlspec.driver import BaseSyncExceptionHandler, SyncDriverAdapterBase, SyncRowStream, validate_savepoint_name from sqlspec.exceptions import DataError, ImproperConfigurationError, SQLSpecError -from sqlspec.typing import import_optional +from sqlspec.typing import ArrowException from sqlspec.utils.module_loader import ensure_pyarrow from sqlspec.utils.text import quote_identifier, split_qualified_identifier if TYPE_CHECKING: - import pyarrow as pa - from sqlspec.builder import QueryBuilder from sqlspec.core import ArrowResult, ParameterProfile, Statement, StatementConfig, StatementFilter from sqlspec.driver import ExecutionResult from sqlspec.storage import StorageBridgeJob, StorageDestination, StorageFormat, StorageTelemetry from sqlspec.typing import ArrowRecordBatch, ArrowRecordBatchReader, ArrowReturnFormat, StatementParameters -else: - pa = import_optional("pyarrow") __all__ = ("ArrowOdbcCursor", "ArrowOdbcDriver", "ArrowOdbcExceptionHandler", "resolve_dialect_from_dbms_name") @@ -55,7 +51,7 @@ class ArrowOdbcExceptionHandler(BaseSyncExceptionHandler): def _handle_exception(self, exc_type: "type[BaseException] | None", exc_val: "BaseException") -> bool: if exc_type is None: return False - if isinstance(exc_val, ArrowOdbcError) or (pa is not None and isinstance(exc_val, pa.ArrowException)): + if isinstance(exc_val, (ArrowOdbcError, ArrowException)): self.pending_exception = create_mapped_exception(exc_val) return True return False @@ -493,10 +489,11 @@ def bulk_insert_arrow(self, target_table: str, source: Any, *, chunk_size: int | with self.with_cursor(self.connection): def _insert_op() -> bool: - if isinstance(source, pa.Table) and hasattr(self.connection, "from_table_to_db"): + is_table = hasattr(source, "to_batches") and not hasattr(source, "read_next_batch") + if is_table and hasattr(self.connection, "from_table_to_db"): self.connection.from_table_to_db(source=source, target=target_table, chunk_size=resolved_chunk_size) return True - reader = _table_to_reader(source, resolved_chunk_size) if isinstance(source, pa.Table) else source + reader = _table_to_reader(source, resolved_chunk_size) if is_table else source if hasattr(self.connection, "insert_into_table"): self.connection.insert_into_table(reader=reader, table=target_table, chunk_size=resolved_chunk_size) return True @@ -557,7 +554,7 @@ def load_from_arrow( if overwrite: target = _db2_table_reference(table) if self._dialect == "db2" else _quote_odbc_table(table) self.execute(f"DELETE FROM {target}") - if isinstance(raw_source, pa.RecordBatchReader) and hasattr(self.connection, "insert_into_table"): + if hasattr(raw_source, "read_next_batch") and hasattr(self.connection, "insert_into_table"): self.bulk_insert_arrow(table, raw_source) telemetry_payload: StorageTelemetry = { "rows_processed": 0, @@ -678,6 +675,9 @@ def _normalize_reader(self, reader: "ArrowRecordBatchReader") -> "ArrowRecordBat renamed = schema for index, name in enumerate(names): renamed = renamed.set(index, renamed.field(index).with_name(name)) + ensure_pyarrow() + import pyarrow as pa + return pa.RecordBatchReader.from_batches( renamed, (pa.RecordBatch.from_arrays(batch.columns, schema=renamed) for batch in reader) ) @@ -826,6 +826,7 @@ def _odbc_parameters(parameters: Any, *, naive_utc_datetimes: bool = False) -> " def _reader_to_table(reader: Any) -> Any: ensure_pyarrow() + import pyarrow as pa if isinstance(reader, pa.Table): return reader @@ -845,6 +846,7 @@ def _reader_to_table(reader: Any) -> Any: def _to_pyarrow_reader(reader: object) -> "ArrowRecordBatchReader": ensure_pyarrow() + import pyarrow as pa if isinstance(reader, pa.RecordBatchReader): return reader @@ -866,6 +868,7 @@ def _to_pyarrow_reader(reader: object) -> "ArrowRecordBatchReader": def _table_to_reader(table: Any, chunk_size: int) -> Any: ensure_pyarrow() + import pyarrow as pa return pa.RecordBatchReader.from_batches(table.schema, table.to_batches(max_chunksize=chunk_size)) diff --git a/sqlspec/adapters/arrow_odbc/type_converter.py b/sqlspec/adapters/arrow_odbc/type_converter.py index b239327da..b648c21de 100644 --- a/sqlspec/adapters/arrow_odbc/type_converter.py +++ b/sqlspec/adapters/arrow_odbc/type_converter.py @@ -6,9 +6,7 @@ from sqlspec.utils.module_loader import ensure_pyarrow if TYPE_CHECKING: - import pyarrow as pa -else: - pa = import_optional("pyarrow") + from pyarrow import DataType __all__ = ("odbc_type_to_arrow",) @@ -66,7 +64,7 @@ _ARROW_TYPE_CACHE: "dict[tuple[str, tuple[Any, ...], tuple[tuple[str, Any], ...]], Any]" = {} -def odbc_type_to_arrow(sql_type: str, *, precision: int | None = None, scale: int | None = None) -> "pa.DataType": +def odbc_type_to_arrow(sql_type: str, *, precision: int | None = None, scale: int | None = None) -> "DataType": """Resolve an ODBC SQL type name to an Arrow data type.""" normalized_type = sql_type.lower().split("(", 1)[0].strip() if normalized_type in {"decimal", "numeric"} and precision is not None and scale is not None: @@ -82,14 +80,16 @@ def odbc_type_to_arrow(sql_type: str, *, precision: int | None = None, scale: in return _arrow_type(name, args, kwargs) -def _arrow_type(name: str, args: tuple[Any, ...] = (), kwargs: dict[str, Any] | None = None) -> "pa.DataType": +def _arrow_type(name: str, args: tuple[Any, ...] = (), kwargs: dict[str, Any] | None = None) -> "DataType": kwargs_key = tuple(sorted((kwargs or {}).items())) cache_key = (name, args, kwargs_key) cached = _ARROW_TYPE_CACHE.get(cache_key) if cached is not None: - return cast("pa.DataType", cached) + return cast("DataType", cached) ensure_pyarrow() + pa = import_optional("pyarrow") + assert pa is not None resolved = getattr(pa, name)(*args, **(kwargs or {})) _ARROW_TYPE_CACHE[cache_key] = resolved - return cast("pa.DataType", resolved) + return cast("DataType", resolved) diff --git a/sqlspec/adapters/bigquery/_typing.py b/sqlspec/adapters/bigquery/_typing.py index ed8b30b49..c1029d3fe 100644 --- a/sqlspec/adapters/bigquery/_typing.py +++ b/sqlspec/adapters/bigquery/_typing.py @@ -28,19 +28,17 @@ BigQueryConnection: TypeAlias = Client BigQueryParam: TypeAlias = ArrayQueryParameter | ScalarQueryParameter | StructQueryParameter -BigQueryStorageWriteModule: Any = import_optional("google.cloud.bigquery_storage_v1") -BigQueryStorageWriteTypes: Any = import_optional("google.cloud.bigquery_storage_v1.types") if TYPE_CHECKING: from collections.abc import Callable from types import TracebackType import google.cloud.bigquery_storage as bigquery_storage_read_module + import google.cloud.bigquery_storage_v1 as BigQueryStorageWriteModule # noqa: N812 + import google.cloud.bigquery_storage_v1.types as BigQueryStorageWriteTypes # noqa: N812 from sqlspec.adapters.bigquery.driver import BigQueryDriver from sqlspec.core import StatementConfig -else: - bigquery_storage_read_module = import_optional("google.cloud.bigquery_storage") __all__ = ( "BIGQUERY_DEFAULT_RETRY", @@ -69,6 +67,17 @@ _DRIVER_TYPE: "type[BigQueryDriver] | None" = None +def __getattr__(name: str) -> Any: + if name == "BigQueryStorageWriteModule": + return import_optional("google.cloud.bigquery_storage_v1") + if name == "BigQueryStorageWriteTypes": + return import_optional("google.cloud.bigquery_storage_v1.types") + if name == "bigquery_storage_read_module": + return import_optional("google.cloud.bigquery_storage") + msg = f"module {__name__!r} has no attribute {name!r}" + raise AttributeError(msg) + + class BigQueryCursor: """BigQuery cursor with resource management.""" diff --git a/sqlspec/adapters/bigquery/config.py b/sqlspec/adapters/bigquery/config.py index 009b4b113..d77bbd5d3 100644 --- a/sqlspec/adapters/bigquery/config.py +++ b/sqlspec/adapters/bigquery/config.py @@ -441,7 +441,8 @@ def _provide_storage_write_client(self, connection: "BigQueryConnection") -> Any client_options = self.connection_config.get("client_options") if client_options is not None and client_options is not Empty: client_kwargs["client_options"] = client_options - self._storage_write_client = BigQueryStorageWriteModule.BigQueryWriteClient(**client_kwargs) + write_client_cls = cast("Any", BigQueryStorageWriteModule).BigQueryWriteClient + self._storage_write_client = write_client_cls(**client_kwargs) return self._storage_write_client def _resolve_default_dataset(self) -> "str | None": diff --git a/sqlspec/adapters/bigquery/core.py b/sqlspec/adapters/bigquery/core.py index bbcd3b171..f6f62ce86 100644 --- a/sqlspec/adapters/bigquery/core.py +++ b/sqlspec/adapters/bigquery/core.py @@ -45,8 +45,9 @@ UniqueViolationError, _classify_timeout_or_cancellation, ) +from sqlspec.storage._utils import import_pyarrow, import_pyarrow_parquet +from sqlspec.typing import PYARROW_INSTALLED from sqlspec.utils.logging import get_logger -from sqlspec.utils.module_loader import ensure_pyarrow, import_optional from sqlspec.utils.serializers import from_json, to_json from sqlspec.utils.type_converters import build_uuid_coercions from sqlspec.utils.type_guards import has_errors @@ -61,9 +62,6 @@ BigQueryLoadFormat = Literal["jsonl", "json", "parquet", "arrow-ipc", "csv", "avro", "orc"] -_pa: Any = import_optional("pyarrow") -_pq: Any = import_optional("pyarrow.parquet") - __all__ = ( "BigQueryStreamSource", "apply_driver_features", @@ -184,21 +182,23 @@ def try_bulk_insert( if _has_synthetic_positional_keys(remapped_parameters): return None + if not PYARROW_INSTALLED: + logger.debug("pyarrow not available, falling back to literal inlining") + return None + try: - ensure_pyarrow() - arrow_table = _pa.Table.from_pylist(remapped_parameters) + pa = import_pyarrow() + pq = import_pyarrow_parquet() + arrow_table = pa.Table.from_pylist(remapped_parameters) buffer = io.BytesIO() - _pq.write_table(arrow_table, buffer) + pq.write_table(arrow_table, buffer) buffer.seek(0) job_config = build_load_job_config("parquet", overwrite=False) job = connection.load_table_from_file(buffer, table_name, job_config=job_config, timeout=result_timeout) job.result(timeout=result_timeout) return len(remapped_parameters) - except ImportError: - logger.debug("pyarrow not available, falling back to literal inlining") - return None except Exception as exc: logger.debug("Bulk insert failed, falling back to literal inlining: %s", exc) return None diff --git a/sqlspec/adapters/bigquery/driver.py b/sqlspec/adapters/bigquery/driver.py index 249a784f1..c2ce7f2cb 100644 --- a/sqlspec/adapters/bigquery/driver.py +++ b/sqlspec/adapters/bigquery/driver.py @@ -62,8 +62,9 @@ ) from sqlspec.driver import BaseSyncExceptionHandler, ExecutionResult, SyncDriverAdapterBase, SyncRowStream from sqlspec.exceptions import ImproperConfigurationError, StorageOperationFailedError +from sqlspec.storage._utils import import_pyarrow, import_pyarrow_parquet from sqlspec.utils.logging import get_logger -from sqlspec.utils.module_loader import ensure_pyarrow, import_optional +from sqlspec.utils.module_loader import ensure_pyarrow from sqlspec.utils.serializers import to_json from sqlspec.utils.text import split_qualified_identifier @@ -85,8 +86,6 @@ logger = get_logger(__name__) _DATASET_TABLE_PARTS = 2 _PROJECT_DATASET_TABLE_PARTS = 3 -_pa: Any = import_optional("pyarrow") -_pq: Any = import_optional("pyarrow.parquet") class BigQueryExceptionHandler(BaseSyncExceptionHandler): @@ -826,20 +825,23 @@ def load_from_arrow( ensure_pyarrow() qualified_table = self._qualify_table_name(table) - if self._enable_storage_write_api and not overwrite and not uses_local_bigquery_endpoint(self.connection): - try: - telemetry_payload = self._load_arrow_via_storage_write_api(qualified_table, arrow_table) - except ImportError as exc: - logger.warning("Storage Write API unavailable, falling back to Parquet load job: %s", exc) - else: - if telemetry: - telemetry_payload.setdefault("extra", {}) - telemetry_payload["extra"]["arrow_rows"] = telemetry.get("rows_processed") - self._attach_partition_telemetry(telemetry_payload, partitioner) - return self._storage_job(telemetry_payload) - + if ( + self._enable_storage_write_api + and not overwrite + and not uses_local_bigquery_endpoint(self.connection) + and BigQueryStorageWriteModule is not None + and BigQueryStorageWriteTypes is not None + ): + telemetry_payload = self._load_arrow_via_storage_write_api(qualified_table, arrow_table) + if telemetry: + telemetry_payload.setdefault("extra", {}) + telemetry_payload["extra"]["arrow_rows"] = telemetry.get("rows_processed") + self._attach_partition_telemetry(telemetry_payload, partitioner) + return self._storage_job(telemetry_payload) + + pq = import_pyarrow_parquet() buffer = io.BytesIO() - _pq.write_table(arrow_table, buffer) + pq.write_table(arrow_table, buffer) buffer.seek(0) job_config = build_load_job_config("parquet", overwrite) exc_handler = self.handle_database_exceptions() @@ -1041,7 +1043,8 @@ def _load_arrow_via_storage_write_api( client = provider(self.connection) else: credentials = getattr(self.connection, "_credentials", None) - client = BigQueryStorageWriteModule.BigQueryWriteClient(credentials=credentials) + write_client_cls = cast("Any", BigQueryStorageWriteModule).BigQueryWriteClient + client = write_client_cls(credentials=credentials) parent = f"projects/{project}/datasets/{dataset}/tables/{table_name}" write_stream_type_enum = types.WriteStream.Type.COMMITTED if is_committed else types.WriteStream.Type.PENDING write_stream = client.create_write_stream( @@ -1115,42 +1118,42 @@ def _bigquery_arrow_reader_from_iterable( batches: "Iterable[ArrowRecordBatch]", schema: Any = None ) -> "ArrowRecordBatchReader": ensure_pyarrow() + pa = import_pyarrow() iterator = iter(batches) try: first_batch = next(iterator) except StopIteration: - if isinstance(schema, _pa.Schema): + if isinstance(schema, pa.Schema): empty_schema = schema elif isinstance(schema, (list, tuple)): type_map = { - "INTEGER": _pa.int64(), - "INT64": _pa.int64(), - "FLOAT": _pa.float64(), - "FLOAT64": _pa.float64(), - "BOOLEAN": _pa.bool_(), - "BOOL": _pa.bool_(), - "STRING": _pa.string(), - "BYTES": _pa.binary(), - "DATE": _pa.date32(), - "DATETIME": _pa.timestamp("us"), - "TIMESTAMP": _pa.timestamp("us", tz="UTC"), - "NUMERIC": _pa.decimal128(38, 9), - "BIGNUMERIC": _pa.decimal256(76, 38), + "INTEGER": pa.int64(), + "INT64": pa.int64(), + "FLOAT": pa.float64(), + "FLOAT64": pa.float64(), + "BOOLEAN": pa.bool_(), + "BOOL": pa.bool_(), + "STRING": pa.string(), + "BYTES": pa.binary(), + "DATE": pa.date32(), + "DATETIME": pa.timestamp("us"), + "TIMESTAMP": pa.timestamp("us", tz="UTC"), + "NUMERIC": pa.decimal128(38, 9), + "BIGNUMERIC": pa.decimal256(76, 38), } fields = [ - _pa.field( + pa.field( getattr(field, "name", str(field)), - type_map.get(str(getattr(field, "field_type", "")).upper(), _pa.null()), + type_map.get(str(getattr(field, "field_type", "")).upper(), pa.null()), ) for field in schema ] - empty_schema = _pa.schema(fields) + empty_schema = pa.schema(fields) else: - empty_schema = _pa.schema([]) - return cast("ArrowRecordBatchReader", _pa.RecordBatchReader.from_batches(empty_schema, ())) + empty_schema = pa.schema([]) + return cast("ArrowRecordBatchReader", pa.RecordBatchReader.from_batches(empty_schema, ())) return cast( - "ArrowRecordBatchReader", - _pa.RecordBatchReader.from_batches(first_batch.schema, chain((first_batch,), iterator)), + "ArrowRecordBatchReader", pa.RecordBatchReader.from_batches(first_batch.schema, chain((first_batch,), iterator)) ) diff --git a/sqlspec/adapters/duckdb/driver.py b/sqlspec/adapters/duckdb/driver.py index 3ac29c906..2f095c92a 100644 --- a/sqlspec/adapters/duckdb/driver.py +++ b/sqlspec/adapters/duckdb/driver.py @@ -42,7 +42,7 @@ from sqlspec.core.result import DMLResult from sqlspec.driver import BaseSyncExceptionHandler, SyncDriverAdapterBase, SyncRowStream from sqlspec.exceptions import SQLSpecError -from sqlspec.typing import import_optional +from sqlspec.typing import PYARROW_INSTALLED, ArrowReturnFormat, StatementParameters from sqlspec.utils.arrow_helpers import arrow_reader_with_deferred_close from sqlspec.utils.logging import get_logger from sqlspec.utils.module_loader import ensure_pyarrow @@ -53,17 +53,12 @@ from collections.abc import Callable, Mapping, Sequence from pathlib import Path - import pyarrow as pa - from sqlspec.adapters.duckdb._typing import DuckDBConnection from sqlspec.builder import QueryBuilder from sqlspec.core import ArrowResult, CachedQuery, SQLResult, Statement, StatementFilter from sqlspec.core.result._base import RowFormat from sqlspec.driver import ExecutionResult from sqlspec.storage import StorageBridgeJob, StorageDestination, StorageFormat, StorageTelemetry - from sqlspec.typing import ArrowReturnFormat, StatementParameters -else: - pa = import_optional("pyarrow") __all__ = ("DuckDBCursor", "DuckDBDriver", "DuckDBExceptionHandler", "DuckDBSessionContext") @@ -625,7 +620,8 @@ def load_from_arrow( source_data = source.get_data() if hasattr(source, "get_data") else source arrow_table = None arrow_source: object - if isinstance(source_data, pa.RecordBatchReader): + is_reader = hasattr(source_data, "read_next_batch") + if is_reader: arrow_source = source_data else: arrow_table = self._coerce_arrow_table(source_data) @@ -657,7 +653,7 @@ def load_from_arrow( raise self._check_pending_exception(exc_handler) - if isinstance(source_data, pa.RecordBatchReader): + if is_reader: telemetry_payload: StorageTelemetry = { "rows_processed": inserted_rows, "bytes_processed": 0, @@ -768,7 +764,7 @@ def _execute_and_materialize_select( hint = self._cache_holder.get_row_count_hint(cache_key) sql_method = getattr(cursor, "sql", None) if ( - pa is not None + PYARROW_INSTALLED and (hint is None or hint >= ARROW_MIN_ROWS) and not profiling_enabled and callable(sql_method) @@ -822,7 +818,7 @@ def _materialize_select_rows(self, cursor: Any, compiled_sql: str) -> "tuple[lis description = cursor.description hint = self._cache_holder.get_row_count_hint(compiled_sql) if ( - pa is None + not PYARROW_INSTALLED or (hint is not None and hint < ARROW_MIN_ROWS) or description_has_arrow_divergent_types(description) or not hasattr(cursor, "to_arrow_table") @@ -939,7 +935,7 @@ def _execute_bulk_insert_many( @staticmethod def _build_arrow_table(rows: "list[Any]", column_names: "list[str]") -> Any | None: """Build a pyarrow table from batch rows when they share a stable shape.""" - if not rows or pa is None: + if not rows or not PYARROW_INSTALLED: return None first_row = rows[0] @@ -950,6 +946,8 @@ def _build_arrow_table(rows: "list[Any]", column_names: "list[str]") -> Any | No if any(not isinstance(row, (list, tuple)) or len(row) != len(column_names) for row in rows): return None try: + import pyarrow as pa + return pa.Table.from_arrays([pa.array(col) for col in zip(*rows, strict=True)], names=column_names) except Exception: return None diff --git a/sqlspec/adapters/spanner/_typing.py b/sqlspec/adapters/spanner/_typing.py index 3cf3a9561..f188773ec 100644 --- a/sqlspec/adapters/spanner/_typing.py +++ b/sqlspec/adapters/spanner/_typing.py @@ -5,104 +5,70 @@ """ from importlib import import_module -from typing import TYPE_CHECKING, Any, cast - -from sqlspec.typing import import_optional_attr - - -class _UnavailableSpannerGoogleAPICallError(Exception): - """Fallback Spanner API exception when google-api-core is unavailable.""" - - -class _UnavailableSpannerTransaction: - """Fallback Spanner transaction class when google-cloud-spanner is unavailable.""" - - __slots__ = () - - -class _UnavailableSpannerAsyncTransaction: - """Fallback Spanner async transaction class when google-cloud-spanner is unavailable.""" - - __slots__ = () - +from typing import TYPE_CHECKING, Any, TypeAlias, cast + +from google.api_core import exceptions as spanner_exceptions +from google.api_core.client_info import ClientInfo as SpannerClientInfo +from google.api_core.client_options import ClientOptions as SpannerClientOptions +from google.api_core.exceptions import GoogleAPICallError as SpannerGoogleAPICallError +from google.api_core.exceptions import GoogleAPIError as SpannerGoogleAPIError +from google.api_core.exceptions import NotFound as SpannerNotFound +from google.api_core.retry import Retry as SpannerRetry +from google.auth.credentials import Credentials as SpannerCredentials +from google.cloud.spanner_admin_database_v1.types import DatabaseDialect as SpannerDatabaseDialect +from google.cloud.spanner_admin_database_v1.types import EncryptionConfig as SpannerEncryptionConfig +from google.cloud.spanner_v1 import Client as SpannerClient +from google.cloud.spanner_v1 import DirectedReadOptions as SpannerDirectedReadOptions +from google.cloud.spanner_v1 import ExecuteSqlRequest as SpannerExecuteSqlRequest +from google.cloud.spanner_v1 import RequestOptions as SpannerRequestOptions +from google.cloud.spanner_v1 import param_types as spanner_param_types +from google.cloud.spanner_v1._async.batch import Batch as SpannerAsyncBatch +from google.cloud.spanner_v1._async.batch import MutationGroups as SpannerAsyncMutationGroups +from google.cloud.spanner_v1._async.client import Client as SpannerAsyncClient +from google.cloud.spanner_v1._async.database import Database as SpannerAsyncDatabase +from google.cloud.spanner_v1._async.database import SnapshotCheckout as SpannerAsyncSnapshotCheckout +from google.cloud.spanner_v1._async.database_sessions_manager import TransactionType as SpannerAsyncTransactionType +from google.cloud.spanner_v1._async.pool import AbstractSessionPool as SpannerAsyncAbstractSessionPool +from google.cloud.spanner_v1._async.pool import BurstyPool as SpannerAsyncBurstyPool +from google.cloud.spanner_v1._async.pool import FixedSizePool as SpannerAsyncFixedSizePool +from google.cloud.spanner_v1._async.pool import PingingPool as SpannerAsyncPingingPool +from google.cloud.spanner_v1._async.pool import TransactionPingingPool as SpannerAsyncTransactionPingingPool +from google.cloud.spanner_v1._async.session import Session as SpannerAsyncSession +from google.cloud.spanner_v1._async.snapshot import Snapshot as SpannerAsyncSnapshot +from google.cloud.spanner_v1._async.streamed import StreamedResultSet as SpannerAsyncStreamedResultSet +from google.cloud.spanner_v1._async.transaction import Transaction as SpannerAsyncTransaction +from google.cloud.spanner_v1.batch import Batch as SpannerBatch +from google.cloud.spanner_v1.batch import MutationGroups as SpannerMutationGroups +from google.cloud.spanner_v1.data_types import Interval as SpannerInterval +from google.cloud.spanner_v1.data_types import JsonObject as SpannerJsonObject +from google.cloud.spanner_v1.database import Database as SpannerDatabase +from google.cloud.spanner_v1.database import SnapshotCheckout +from google.cloud.spanner_v1.database_sessions_manager import TransactionType as SpannerTransactionType +from google.cloud.spanner_v1.pool import AbstractSessionPool as SpannerAbstractSessionPool +from google.cloud.spanner_v1.pool import BurstyPool as SpannerBurstyPool +from google.cloud.spanner_v1.pool import FixedSizePool as SpannerFixedSizePool +from google.cloud.spanner_v1.pool import PingingPool as SpannerPingingPool +from google.cloud.spanner_v1.pool import TransactionPingingPool as SpannerTransactionPingingPool +from google.cloud.spanner_v1.transaction import DefaultTransactionOptions as SpannerDefaultTransactionOptions +from google.cloud.spanner_v1.transaction import Transaction as SpannerTransaction +from google.cloud.spanner_v1.types.type import TypeCode as SpannerTypeCode if TYPE_CHECKING: from collections.abc import Awaitable, Callable from types import TracebackType - from typing import TypeAlias - - from google.api_core import exceptions as spanner_exceptions - from google.api_core.client_info import ClientInfo as SpannerClientInfo - from google.api_core.client_options import ClientOptions as SpannerClientOptions - from google.api_core.exceptions import GoogleAPICallError as _SpannerGoogleAPICallError - from google.api_core.exceptions import NotFound as SpannerNotFound - from google.api_core.retry import Retry as SpannerRetry - from google.auth.credentials import Credentials as SpannerCredentials - from google.cloud.spanner_admin_database_v1.types import DatabaseDialect as SpannerDatabaseDialect - from google.cloud.spanner_admin_database_v1.types import EncryptionConfig as SpannerEncryptionConfig - from google.cloud.spanner_v1 import Client as SpannerClient - from google.cloud.spanner_v1 import DirectedReadOptions as SpannerDirectedReadOptions - from google.cloud.spanner_v1 import ExecuteSqlRequest as SpannerExecuteSqlRequest - from google.cloud.spanner_v1 import RequestOptions as SpannerRequestOptions - from google.cloud.spanner_v1 import param_types as spanner_param_types - from google.cloud.spanner_v1._async.batch import Batch as SpannerAsyncBatch - from google.cloud.spanner_v1._async.batch import MutationGroups as SpannerAsyncMutationGroups - from google.cloud.spanner_v1._async.client import Client as SpannerAsyncClient - from google.cloud.spanner_v1._async.database import Database as SpannerAsyncDatabase - from google.cloud.spanner_v1._async.database import SnapshotCheckout as SpannerAsyncSnapshotCheckout - from google.cloud.spanner_v1._async.database_sessions_manager import TransactionType as SpannerAsyncTransactionType - from google.cloud.spanner_v1._async.pool import AbstractSessionPool as SpannerAsyncAbstractSessionPool - from google.cloud.spanner_v1._async.pool import BurstyPool as SpannerAsyncBurstyPool - from google.cloud.spanner_v1._async.pool import FixedSizePool as SpannerAsyncFixedSizePool - from google.cloud.spanner_v1._async.pool import PingingPool as SpannerAsyncPingingPool - from google.cloud.spanner_v1._async.pool import TransactionPingingPool as SpannerAsyncTransactionPingingPool - from google.cloud.spanner_v1._async.session import Session as SpannerAsyncSession - from google.cloud.spanner_v1._async.snapshot import Snapshot as SpannerAsyncSnapshot - from google.cloud.spanner_v1._async.streamed import StreamedResultSet as SpannerAsyncStreamedResultSet - from google.cloud.spanner_v1._async.transaction import Transaction as _SpannerAsyncTransaction - from google.cloud.spanner_v1.batch import Batch as SpannerBatch - from google.cloud.spanner_v1.batch import MutationGroups as SpannerMutationGroups - from google.cloud.spanner_v1.data_types import Interval as SpannerInterval - from google.cloud.spanner_v1.data_types import JsonObject as SpannerJsonObject - from google.cloud.spanner_v1.database import Database as SpannerDatabase - from google.cloud.spanner_v1.database import SnapshotCheckout - from google.cloud.spanner_v1.database_sessions_manager import TransactionType as SpannerTransactionType - from google.cloud.spanner_v1.pool import AbstractSessionPool as SpannerAbstractSessionPool - from google.cloud.spanner_v1.pool import BurstyPool as SpannerBurstyPool - from google.cloud.spanner_v1.pool import FixedSizePool as SpannerFixedSizePool - from google.cloud.spanner_v1.pool import PingingPool as SpannerPingingPool - from google.cloud.spanner_v1.pool import TransactionPingingPool as SpannerTransactionPingingPool + from google.cloud.spanner_v1.snapshot import Snapshot - from google.cloud.spanner_v1.transaction import DefaultTransactionOptions as SpannerDefaultTransactionOptions - from google.cloud.spanner_v1.transaction import Transaction as _SpannerTransaction - from google.cloud.spanner_v1.types.type import TypeCode as SpannerTypeCode from sqlspec.adapters.spanner.driver import SpannerAsyncDriver, SpannerSyncDriver from sqlspec.core import StatementConfig - SpannerSyncConnection: TypeAlias = Snapshot | SnapshotCheckout | _SpannerTransaction | SpannerBatch + SpannerSyncConnection: TypeAlias = Snapshot | SnapshotCheckout | SpannerTransaction | SpannerBatch SpannerAsyncConnection: TypeAlias = ( - SpannerAsyncSnapshot | SpannerAsyncSnapshotCheckout | _SpannerAsyncTransaction | SpannerAsyncBatch + SpannerAsyncSnapshot | SpannerAsyncSnapshotCheckout | SpannerAsyncTransaction | SpannerAsyncBatch ) - SpannerGoogleAPICallError: TypeAlias = _SpannerGoogleAPICallError - SpannerTransaction: TypeAlias = _SpannerTransaction - SpannerAsyncTransaction: TypeAlias = _SpannerAsyncTransaction - - -if not TYPE_CHECKING: +else: SpannerSyncConnection = Any SpannerAsyncConnection = Any - SpannerGoogleAPICallError = ( - import_optional_attr("google.api_core.exceptions", "GoogleAPICallError") - or _UnavailableSpannerGoogleAPICallError - ) - SpannerTransaction = ( - import_optional_attr("google.cloud.spanner_v1.transaction", "Transaction") or _UnavailableSpannerTransaction - ) - SpannerAsyncTransaction = ( - import_optional_attr("google.cloud.spanner_v1._async.transaction", "Transaction") - or _UnavailableSpannerAsyncTransaction - ) __all__ = ( @@ -138,6 +104,7 @@ class _UnavailableSpannerAsyncTransaction: "SpannerExecuteSqlRequest", "SpannerFixedSizePool", "SpannerGoogleAPICallError", + "SpannerGoogleAPIError", "SpannerInterval", "SpannerJsonObject", "SpannerMutationGroups", @@ -297,64 +264,6 @@ async def __aexit__( return None -_LAZY_DRIVER_EXPORTS: dict[str, tuple[str, str]] = { - "SpannerAsyncTransactionType": ("google.cloud.spanner_v1._async.database_sessions_manager", "TransactionType"), - "SpannerInterval": ("google.cloud.spanner_v1.data_types", "Interval"), - "SpannerAsyncAbstractSessionPool": ("google.cloud.spanner_v1._async.pool", "AbstractSessionPool"), - "SpannerAsyncBatch": ("google.cloud.spanner_v1._async.batch", "Batch"), - "SpannerAsyncBurstyPool": ("google.cloud.spanner_v1._async.pool", "BurstyPool"), - "SpannerAsyncClient": ("google.cloud.spanner_v1._async.client", "Client"), - "SpannerAsyncDatabase": ("google.cloud.spanner_v1._async.database", "Database"), - "SpannerAsyncFixedSizePool": ("google.cloud.spanner_v1._async.pool", "FixedSizePool"), - "SpannerAsyncMutationGroups": ("google.cloud.spanner_v1._async.batch", "MutationGroups"), - "SpannerAsyncPingingPool": ("google.cloud.spanner_v1._async.pool", "PingingPool"), - "SpannerAsyncSession": ("google.cloud.spanner_v1._async.session", "Session"), - "SpannerAsyncSnapshot": ("google.cloud.spanner_v1._async.snapshot", "Snapshot"), - "SpannerAsyncStreamedResultSet": ("google.cloud.spanner_v1._async.streamed", "StreamedResultSet"), - "SpannerAsyncTransactionPingingPool": ("google.cloud.spanner_v1._async.pool", "TransactionPingingPool"), - "SpannerBatch": ("google.cloud.spanner_v1.batch", "Batch"), - "SpannerBurstyPool": ("google.cloud.spanner_v1.pool", "BurstyPool"), - "SpannerFixedSizePool": ("google.cloud.spanner_v1.pool", "FixedSizePool"), - "SpannerMutationGroups": ("google.cloud.spanner_v1.batch", "MutationGroups"), - "SpannerPingingPool": ("google.cloud.spanner_v1.pool", "PingingPool"), - "SpannerTransactionPingingPool": ("google.cloud.spanner_v1.pool", "TransactionPingingPool"), - "spanner_exceptions": ("google.api_core", "exceptions"), - "SpannerNotFound": ("google.api_core.exceptions", "NotFound"), - "SpannerClient": ("google.cloud.spanner_v1", "Client"), - "spanner_param_types": ("google.cloud.spanner_v1", "param_types"), - "SpannerJsonObject": ("google.cloud.spanner_v1.data_types", "JsonObject"), - "SpannerTransactionType": ("google.cloud.spanner_v1.database_sessions_manager", "TransactionType"), - "SpannerTypeCode": ("google.cloud.spanner_v1.types.type", "TypeCode"), - "SpannerClientInfo": ("google.api_core.client_info", "ClientInfo"), - "SpannerClientOptions": ("google.api_core.client_options", "ClientOptions"), - "SpannerRetry": ("google.api_core.retry", "Retry"), - "SpannerCredentials": ("google.auth.credentials", "Credentials"), - "SpannerDatabaseDialect": ("google.cloud.spanner_admin_database_v1.types", "DatabaseDialect"), - "SpannerEncryptionConfig": ("google.cloud.spanner_admin_database_v1.types", "EncryptionConfig"), - "SpannerDirectedReadOptions": ("google.cloud.spanner_v1", "DirectedReadOptions"), - "SpannerExecuteSqlRequest": ("google.cloud.spanner_v1", "ExecuteSqlRequest"), - "SpannerRequestOptions": ("google.cloud.spanner_v1", "RequestOptions"), - "SpannerDatabase": ("google.cloud.spanner_v1.database", "Database"), - "SpannerAbstractSessionPool": ("google.cloud.spanner_v1.pool", "AbstractSessionPool"), - "SpannerDefaultTransactionOptions": ("google.cloud.spanner_v1.transaction", "DefaultTransactionOptions"), - "SpannerGoogleAPIError": ("google.api_core.exceptions", "GoogleAPIError"), -} - - -def __getattr__(name: str) -> Any: - """Resolve optional driver symbols only when a consumer requests them.""" - target = _LAZY_DRIVER_EXPORTS.get(name) - if target is None: - msg = f"module {__name__!r} has no attribute {name!r}" - raise AttributeError(msg) - module_name, attribute = target - value = getattr(import_module(module_name), attribute, None) - if value is None: - msg = f"Cannot import {attribute!r} from {module_name!r}" - raise ImportError(msg) - return value - - def _get_sync_driver_type() -> "type[SpannerSyncDriver]": driver_module = import_module("sqlspec.adapters.spanner.driver") return cast("type[SpannerSyncDriver]", driver_module.SpannerSyncDriver) diff --git a/sqlspec/adapters/spanner/core.py b/sqlspec/adapters/spanner/core.py index 679039cbd..ac48aec9f 100644 --- a/sqlspec/adapters/spanner/core.py +++ b/sqlspec/adapters/spanner/core.py @@ -15,6 +15,7 @@ from sqlspec.adapters.spanner._typing import SpannerJsonObject as JsonObject from sqlspec.adapters.spanner._typing import SpannerTypeCode as TypeCode from sqlspec.adapters.spanner._typing import spanner_exceptions as api_exceptions +from sqlspec.adapters.spanner._typing import spanner_param_types from sqlspec.adapters.spanner.type_converter import coerce_params_for_spanner, infer_spanner_param_types from sqlspec.core import ( DriverParameterProfile, @@ -39,7 +40,6 @@ UniqueViolationError, ) from sqlspec.utils.arrow_helpers import convert_dict_to_arrow -from sqlspec.utils.module_loader import import_optional_attr from sqlspec.utils.serializers import from_json, to_json if TYPE_CHECKING: @@ -1059,10 +1059,7 @@ def _create_spanner_error(error: Any, error_class: type[SQLSpecError], descripti def _json_param_type() -> Any: """Return the SPANNER_PARAM_TYPES JSON type or fallback to STRING.""" - param_types_module = import_optional_attr("google.cloud.spanner_v1", "param_types") - if param_types_module is not None: - return getattr(param_types_module, "JSON", getattr(param_types_module, "STRING", None)) - return None + return getattr(spanner_param_types, "JSON", getattr(spanner_param_types, "STRING", None)) def _extract_spanner_json_codecs(config: Any) -> "tuple[Callable[[Any], str], Callable[[str], Any]]": diff --git a/sqlspec/typing.py b/sqlspec/typing.py index 938adca52..51cf32a65 100644 --- a/sqlspec/typing.py +++ b/sqlspec/typing.py @@ -52,6 +52,7 @@ if TYPE_CHECKING: from sqlspec._typing import ( + ArrowException, ArrowRecordBatch, ArrowRecordBatchReader, ArrowRecordBatchReaderProtocol, @@ -106,6 +107,7 @@ "PYDANTIC_USE_FAILFAST", "UNSET", "UUID_UTILS_INSTALLED", + "ArrowException", "ArrowRecordBatch", "ArrowRecordBatchReader", "ArrowRecordBatchReaderProtocol", diff --git a/tests/unit/adapters/test_bigquery/test_storage_write_api.py b/tests/unit/adapters/test_bigquery/test_storage_write_api.py index 57f1d017b..f658910c7 100644 --- a/tests/unit/adapters/test_bigquery/test_storage_write_api.py +++ b/tests/unit/adapters/test_bigquery/test_storage_write_api.py @@ -212,10 +212,7 @@ def test_import_failure_falls_back_to_parquet(monkeypatch: pytest.MonkeyPatch) - driver_features={"enable_storage_write_api": True, "storage_capabilities": CAPABILITIES}, ) - def _raise(_self: Any, _table: str, _arrow: Any) -> Any: - raise ImportError("bigquery_storage not available") - - monkeypatch.setattr(BigQueryDriver, "_load_arrow_via_storage_write_api", _raise) + monkeypatch.setattr("sqlspec.adapters.bigquery.driver.BigQueryStorageWriteModule", None) driver.load_from_arrow("dataset.table", pa.table({"id": [1]})) diff --git a/tests/unit/test_import_time.py b/tests/unit/test_import_time.py index 4b3a1b3df..50114d3d7 100644 --- a/tests/unit/test_import_time.py +++ b/tests/unit/test_import_time.py @@ -200,3 +200,191 @@ def test_pytest_collect_only_succeeds_without_plugin_registration_errors() -> No [sys.executable, "-m", "pytest", "--collect-only", "-q", "tests"], capture_output=True, text=True, check=False ) assert result.returncode == 0, f"pytest --collect-only failed:\n{result.stderr}\n{result.stdout}" + + +ADAPTER_DRIVERS: dict[str, tuple[str, ...]] = { + "adbc": ("adbc_driver_manager",), + "aiomysql": ("aiomysql",), + "aiosqlite": ("aiosqlite",), + "arrow_odbc": ("arrow_odbc",), + "asyncmy": ("asyncmy",), + "asyncpg": ("asyncpg",), + "bigquery": ("google.cloud.bigquery",), + "cockroach_asyncpg": ("asyncpg",), + "cockroach_psycopg": ("psycopg",), + "db2": ("ibm_db", "ibm_db_dbi"), + "duckdb": ("duckdb",), + "mssql_python": ("mssql_python",), + "mysqlconnector": ("mysql.connector",), + "oracledb": ("oracledb",), + "psqlpy": ("psqlpy",), + "psycopg": ("psycopg",), + "pymssql": ("pymssql",), + "pymysql": ("pymysql",), + "spanner": ("google.cloud.spanner_v1",), + "sqlanywhere": ("sqlanydb",), + "sqlite": (), + "sybase": ("pymssql",), +} +"""Driver module(s) each adapter imports; every package under ``sqlspec/adapters/`` must be registered.""" + +DRIVER_PACKAGE_DEPENDENCIES: dict[str, tuple[str, ...]] = {"aiomysql": ("pymysql",)} +"""Other registered driver modules that a driver package itself depends on (aiomysql requires PyMySQL).""" + +SECONDARY_OPTIONAL_MODULES = ( + "pyarrow", + "pandas", + "polars", + "numpy", + "pgvector", + "rich", + "opentelemetry", + "litestar", + "pydantic", + "fsspec", + "obstore", + "google.cloud.storage", + "google.cloud.sql.connector", + "google.cloud.alloydb.connector", +) + +DRIVER_OWNED_OPTIONAL_IMPORTS: dict[str, frozenset[str]] = { + "adbc": frozenset({"pyarrow", "pandas", "polars", "numpy"}), + "arrow_odbc": frozenset({"pyarrow", "numpy"}), + "bigquery": frozenset({"pyarrow", "pandas", "numpy", "opentelemetry"}), + "mysqlconnector": frozenset({"opentelemetry"}), + "spanner": frozenset({"opentelemetry"}), +} +"""Optional modules the driver SDK imports on its own when installed; sqlspec code does not import them.""" + + +def test_every_adapter_package_is_registered() -> None: + """Every adapter package declares its driver so the import-contract gates cover it.""" + packages = { + path.name + for path in Path("sqlspec/adapters").iterdir() + if path.is_dir() and (path / "__init__.py").exists() and not path.name.startswith(("_", ".")) + } + assert packages == set(ADAPTER_DRIVERS) + + +@pytest.mark.parametrize("adapter", sorted(name for name, drivers in ADAPTER_DRIVERS.items() if drivers)) +def test_adapter_import_raises_driver_module_not_found(adapter: str) -> None: + """Importing an adapter without its driver raises the driver's own ModuleNotFoundError.""" + drivers = ADAPTER_DRIVERS[adapter] + script = f""" +import importlib.abc +import sys + +blocked = {drivers!r} + +class BlockDriver(importlib.abc.MetaPathFinder): + def find_spec(self, fullname, path, target=None): + if any(fullname == name or fullname.startswith(name + ".") for name in blocked): + raise ModuleNotFoundError(f"No module named {{fullname!r}}", name=fullname) + return None + +sys.meta_path.insert(0, BlockDriver()) +try: + import sqlspec.adapters.{adapter} +except ModuleNotFoundError as exc: + assert exc.name and any(exc.name == name or exc.name.startswith(name + ".") or name.startswith(exc.name + ".") for name in blocked), exc.name +else: + raise AssertionError("adapter imported without its driver") +""" + result = subprocess.run([sys.executable, "-c", script], capture_output=True, text=True, check=False) + assert result.returncode == 0, result.stderr + + +@pytest.mark.parametrize("adapter", sorted(ADAPTER_DRIVERS)) +def test_adapter_imports_with_every_other_driver_blocked(adapter: str) -> None: + """Each adapter works without any other adapter's driver installed.""" + own = {*ADAPTER_DRIVERS[adapter], *DRIVER_PACKAGE_DEPENDENCIES.get(adapter, ())} + blocked = sorted({name for drivers in ADAPTER_DRIVERS.values() for name in drivers} - own) + script = f""" +import importlib +import importlib.abc +import sys + +blocked = {blocked!r} + +class BlockOtherDrivers(importlib.abc.MetaPathFinder): + def find_spec(self, fullname, path, target=None): + if any(fullname == name or fullname.startswith(name + ".") for name in blocked): + raise ModuleNotFoundError(f"No module named {{fullname!r}}", name=fullname) + return None + +sys.meta_path.insert(0, BlockOtherDrivers()) +package = importlib.import_module("sqlspec.adapters.{adapter}") +for name in getattr(package, "__all__", ()): + getattr(package, name) +""" + result = subprocess.run([sys.executable, "-c", script], capture_output=True, text=True, check=False) + assert result.returncode == 0, result.stderr + + +@pytest.mark.parametrize("adapter", sorted(ADAPTER_DRIVERS)) +def test_adapter_import_does_not_load_secondary_optional_modules(adapter: str) -> None: + """Importing an adapter loads only its driver, plus optional modules the driver SDK imports itself.""" + allowed = sorted(DRIVER_OWNED_OPTIONAL_IMPORTS.get(adapter, frozenset())) + script = ( + f"import sys, importlib; importlib.import_module('sqlspec.adapters.{adapter}'); " + f"print(','.join(m for m in {SECONDARY_OPTIONAL_MODULES!r} if m in sys.modules and m not in {allowed!r}))" + ) + result = subprocess.run([sys.executable, "-c", script], capture_output=True, text=True, check=True) + leaked = [name for name in result.stdout.strip().split(",") if name] + assert not leaked, f"import sqlspec.adapters.{adapter} eagerly imported: {leaked}" + + +def _is_type_checking(test: ast.expr) -> bool: + return (isinstance(test, ast.Name) and test.id == "TYPE_CHECKING") or ( + isinstance(test, ast.Attribute) and test.attr == "TYPE_CHECKING" + ) + + +def _imports_driver(module_name: str, drivers: tuple[str, ...]) -> bool: + return any(module_name == name or module_name.startswith(name + ".") for name in drivers) + + +def test_adapter_modules_follow_import_contract() -> None: + """Adapter modules import their driver directly and defer secondary optional modules to first use.""" + violations: list[str] = [] + for adapter, drivers in sorted(ADAPTER_DRIVERS.items()): + for py_file in sorted(Path("sqlspec/adapters", adapter).rglob("*.py")): + tree = ast.parse(py_file.read_text(encoding="utf-8"), filename=str(py_file)) + function_lines: set[int] = set() + for node in ast.walk(tree): + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): + function_lines.update(range(node.lineno, (node.end_lineno or node.lineno) + 1)) + for node in ast.walk(tree): + location = f"{py_file}:{getattr(node, 'lineno', '?')}" + if isinstance(node, ast.If) and isinstance(node.test, ast.UnaryOp): + if isinstance(node.test.op, ast.Not) and _is_type_checking(node.test.operand): + violations.append(f"{location}: use one `if TYPE_CHECKING: ... else:` block") + if isinstance(node, ast.ClassDef) and node.name.startswith("_"): + if any(marker in node.name for marker in ("Stub", "Unavailable", "Missing")): + violations.append(f"{location}: placeholder class {node.name} for a missing dependency") + if isinstance(node, ast.Try): + for handler in node.handlers: + names = handler.type.elts if isinstance(handler.type, ast.Tuple) else [handler.type] + catches_import_error = any( + isinstance(n, ast.Name) and n.id in {"ImportError", "ModuleNotFoundError"} for n in names + ) + reraises = any(isinstance(child, ast.Raise) for child in ast.walk(handler)) + if catches_import_error and not reraises: + violations.append(f"{location}: import fallback via try/except {ast.unparse(handler.type)}") + if not isinstance(node, ast.Call) or not node.args: + continue + func_name = node.func.id if isinstance(node.func, ast.Name) else getattr(node.func, "attr", "") + first = node.args[0] + if func_name not in {"import_optional", "import_optional_attr", "import_module"}: + continue + if not (isinstance(first, ast.Constant) and isinstance(first.value, str)): + continue + if _imports_driver(first.value, drivers): + violations.append(f"{location}: {func_name}({first.value!r}) for the adapter's own driver") + elif node.lineno not in function_lines and first.value.split(".")[0] in { + name.split(".")[0] for name in SECONDARY_OPTIONAL_MODULES + }: + violations.append(f"{location}: module-level {func_name}({first.value!r}); resolve on first use") + assert not violations, "Adapter import contract violations:\n" + "\n".join(violations) From 04ec9aace473359860e021b198197742a38858cc Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 04:53:27 +0000 Subject: [PATCH 125/174] refactor(adapters): normalize columnar driver_features names and declarations Rename arrow_odbc enable_driver_pooling to enable_pooling and BigQuery use_query_and_wait to enable_query_and_wait, and remove the ADBC strict_type_coercion and arrow_extension_types aliases. Declare events_backend on ArrowOdbcDriverFeatures and default_query_job_config on BigQueryDriverFeatures, and drop asyncpg's dead statement_cache_size lookups in driver_features. A contract test checks that every literal driver_features key an adapter reads is declared on its DriverFeatures TypedDict. --- docs/changelog.md | 3 +++ docs/databases/arrow-odbc.md | 6 ++--- docs/databases/bigquery.md | 4 +-- docs/databases/index.md | 2 +- sqlspec/adapters/adbc/core.py | 19 -------------- sqlspec/adapters/arrow_odbc/config.py | 8 +++--- sqlspec/adapters/asyncpg/driver.py | 7 +----- sqlspec/adapters/bigquery/config.py | 4 ++- sqlspec/adapters/bigquery/driver.py | 6 ++--- .../adapters/_shared/_driver_type_system.py | 4 +-- .../integration/adapters/_shared/behaviors.py | 4 +-- .../adapters/test_adapter_implementations.py | 25 +++++++++++++++++++ .../test_adbc_config_normalization.py | 16 ++++++------ .../adapters/test_arrow_odbc/test_config.py | 6 ++--- .../adapters/test_bigquery/test_config.py | 2 +- .../test_bigquery/test_job_controls.py | 12 ++++++--- 16 files changed, 70 insertions(+), 58 deletions(-) diff --git a/docs/changelog.md b/docs/changelog.md index 34791e1a6..23b24d80d 100644 --- a/docs/changelog.md +++ b/docs/changelog.md @@ -29,6 +29,9 @@ important operational fixes. * Removed the `pgbouncer` key from `AsyncpgConnectionConfig`; set `driver_features={"pooler_mode": "transaction"}` instead. * Removed the empty `SybaseADKConfig`, `SqlanywhereADKConfig`, `Db2LitestarConfig`, and `Db2EventsConfig` TypedDicts; use the base `ADKConfig`, `LitestarConfig`, and `EventsConfig`. * Renamed `use_pool` to `enable_pooling` (default `True`) across `mssql_python` and `pymssql` driver features. +* Renamed `arrow_odbc` `enable_driver_pooling` to `enable_pooling` and BigQuery `use_query_and_wait` to `enable_query_and_wait`. +* Removed the ADBC `strict_type_coercion` and `arrow_extension_types` driver feature aliases; use `enable_strict_type_coercion` and `enable_arrow_extension_types`. +* Declared driver feature keys that adapters already read: `events_backend` on `ArrowOdbcDriverFeatures` and `default_query_job_config` on `BigQueryDriverFeatures`. * Replaced `null_pool` in `psycopg` connection configuration and driver features with canonical `enable_pooling=False`. * Replaced correlated per-column scalar subqueries in data dictionary `columns.sql` and `indexes.sql` catalog queries across `postgres`, `cockroach`, `mssql`, `mysql`, `oracle`, `db2`, and `duckdb` with pre-aggregated CTEs and joins, and reduced `ColumnMetadata` to its 18 core catalog fields. ([#866](https://github.com/litestar-org/sqlspec/pull/866)) diff --git a/docs/databases/arrow-odbc.md b/docs/databases/arrow-odbc.md index 071b82bf7..6d1612b4f 100644 --- a/docs/databases/arrow-odbc.md +++ b/docs/databases/arrow-odbc.md @@ -50,7 +50,7 @@ config = spec.add_config( }, statement_config=StatementConfig(dialect="tsql"), driver_features={ - "enable_driver_pooling": True, + "enable_pooling": True, "chunk_size": 65536, "fetch_concurrently": True, }, @@ -60,13 +60,13 @@ config = spec.add_config( ## Key Driver Features and Pool Tuning -`ArrowOdbcConfig` subclasses `NoPoolSyncConfig` and delegates connection reuse to ODBC Driver Manager pooling when `enable_driver_pooling=True`: +`ArrowOdbcConfig` subclasses `NoPoolSyncConfig` and delegates connection reuse to ODBC Driver Manager pooling when `enable_pooling=True`: | Setting | Location | Description | | --- | --- | --- | | `connection_string` | `connection_config` | Standard semicolon-delimited ODBC connection string (`Driver={...};Server=...;`). | | `user` / `password` | `connection_config` | Optional credentials passed alongside the ODBC connection string. | -| `enable_driver_pooling` | `driver_features` | Enables process-wide ODBC Driver Manager connection pooling for low-latency session checkouts. | +| `enable_pooling` | `driver_features` | Enables process-wide ODBC Driver Manager connection pooling for low-latency session checkouts. | | `chunk_size` / `max_bytes_per_batch` | `driver_features` | Maximum rows and memory bytes allocated per Arrow `RecordBatch` buffer during streaming reads. | | `fetch_concurrently` | `driver_features` | Allocates double-buffered background ODBC fetch threads in Rust while Python processes the current Arrow batch. | | `max_text_size` / `max_binary_size` | `driver_features` | Upper byte bounds for unbounded `VARCHAR(MAX)` / `NVARCHAR(MAX)` / `VARBINARY(MAX)` columns. | diff --git a/docs/databases/bigquery.md b/docs/databases/bigquery.md index 05aa60e0e..57c1990b1 100644 --- a/docs/databases/bigquery.md +++ b/docs/databases/bigquery.md @@ -48,7 +48,7 @@ config = spec.add_config( "maximum_bytes_billed": 10 * 1024**3, }, driver_features={ - "use_query_and_wait": True, + "enable_query_and_wait": True, "enable_storage_write_api": True, "storage_write_stream_type": "PENDING", }, @@ -64,7 +64,7 @@ config = spec.add_config( | `credentials` / `credentials_path` | `connection_config` | Explicit `google.auth.credentials.Credentials` instance or service-account JSON path (uses ADC when omitted). | | `maximum_bytes_billed` / `priority` / `reservation` | `connection_config` | Cost guardrails (`INTERACTIVE` vs `BATCH` priority and BigQuery Editions reservation assignment). | | `labels` | `connection_config` | Billing and observability key-value labels attached to submitted jobs. | -| `use_query_and_wait` | `driver_features` | Uses low-latency `client.query_and_wait()` for interactive queries and falls back to `QueryJob` when destination tables, transactions, or `EXPORT DATA` require job handles. | +| `enable_query_and_wait` | `driver_features` | Uses low-latency `client.query_and_wait()` for interactive queries and falls back to `QueryJob` when destination tables, transactions, or `EXPORT DATA` require job handles. | | `enable_storage_write_api` / `storage_write_stream_type` | `driver_features` | Streams `load_from_arrow` batches via BigQuery Storage Write API (`"PENDING"` atomic commit default or `"COMMITTED"` immediate visibility). | | `enable_native_storage` | `driver_features` | Accelerates GCS `gs://` Parquet/CSV/JSON exports and loads using server-side BigQuery jobs. | | `request_timeout` / `job_result_timeout` | `driver_features` | Per-RPC and overall job completion timeout bounds in seconds. | diff --git a/docs/databases/index.md b/docs/databases/index.md index 2bac5e867..15c67587e 100644 --- a/docs/databases/index.md +++ b/docs/databases/index.md @@ -27,7 +27,7 @@ SQLSpec supports synchronous and asynchronous **database adapters** spanning emb | [`bigquery`](bigquery.md) | `sqlspec[bigquery]` | Sync | `@name` (`named`) | Stateless HTTP/gRPC Client | Storage Read / Write API | Table Queue | Sessions + Memory (`SEARCH` / `LIKE`) | | [`spanner`](spanner.md) | `sqlspec[spanner]` | Sync + Async | `@name` (`named`) | `PingingPool` / `TransactionPingingPool` | Row / Batched | Lock-Free Snapshot + `run_in_transaction` | Sessions + Memory (`ARRAY`, FTS `SCORE`, RRF) + Artifacts | | [`adbc`](adbc.md) | `sqlspec[adbc]` | Sync | `?` / `$1` / `:1` | Driver-Managed (`NoPoolSyncConfig`) | Native C-Stream | Table Queue | Sessions + Memory (PostgreSQL `pgvector` RRF, DuckDB, Oracle, Db2) | -| [`arrow_odbc`](arrow-odbc.md) | `sqlspec[arrow-odbc]` | Sync | `?` (`qmark`) | ODBC Driver Pool (`enable_driver_pooling`) | Native Rust ODBC Batch | Table Queue | Sessions + Memory (`LIKE` / `POSSTR`) | +| [`arrow_odbc`](arrow-odbc.md) | `sqlspec[arrow-odbc]` | Sync | `?` (`qmark`) | ODBC Driver Pool (`enable_pooling`) | Native Rust ODBC Batch | Table Queue | Sessions + Memory (`LIKE` / `POSSTR`) | ## Database Setup Guides diff --git a/sqlspec/adapters/adbc/core.py b/sqlspec/adapters/adbc/core.py index 0c5c7b8e2..2f1a3f863 100644 --- a/sqlspec/adapters/adbc/core.py +++ b/sqlspec/adapters/adbc/core.py @@ -1028,16 +1028,11 @@ def apply_driver_features( Updated statement configuration and normalized driver features. """ processed_features: dict[str, Any] = dict(driver_features) if driver_features else {} - processed_features = _normalize_adbc_driver_features(processed_features) resolved_dialect = resolve_dialect_name(statement_config.dialect).lower() json_serializer = cast("Callable[[Any], str] | None", processed_features.setdefault("json_serializer", to_json)) processed_features.setdefault("enable_cast_detection", True) - if "strict_type_coercion" in processed_features: - processed_features.setdefault("enable_strict_type_coercion", processed_features.pop("strict_type_coercion")) processed_features.setdefault("enable_strict_type_coercion", False) - if "arrow_extension_types" in processed_features: - processed_features.setdefault("enable_arrow_extension_types", processed_features.pop("arrow_extension_types")) processed_features.setdefault("enable_arrow_extension_types", True) processed_features.setdefault("enable_pgvector", PGVECTOR_INSTALLED) processed_features.setdefault("enable_paradedb", True) @@ -1797,20 +1792,6 @@ def _map_sqlstate_exception(error: Any, sqlstate: str) -> SQLSpecError: return _create_adbc_error(error, SQLSpecError, "database error") -def _normalize_adbc_driver_features(processed_features: "dict[str, Any]") -> "dict[str, Any]": - if "strict_type_coercion" in processed_features: - val = processed_features.pop("strict_type_coercion") - if "enable_strict_type_coercion" not in processed_features: - processed_features["enable_strict_type_coercion"] = val - - if "arrow_extension_types" in processed_features: - val = processed_features.pop("arrow_extension_types") - if "enable_arrow_extension_types" not in processed_features: - processed_features["enable_arrow_extension_types"] = val - - return processed_features - - def _validate_strict_parameter_value(value: Any) -> None: if isinstance(value, (list, tuple)): for item in value: diff --git a/sqlspec/adapters/arrow_odbc/config.py b/sqlspec/adapters/arrow_odbc/config.py index 4359daa51..5ed6174fd 100644 --- a/sqlspec/adapters/arrow_odbc/config.py +++ b/sqlspec/adapters/arrow_odbc/config.py @@ -1,7 +1,7 @@ """arrow-odbc database configuration.""" import threading -from typing import TYPE_CHECKING, Any, ClassVar, NotRequired, TypedDict, cast +from typing import TYPE_CHECKING, Any, ClassVar, Literal, NotRequired, TypedDict, cast from sqlspec.adapters.arrow_odbc._typing import ( ArrowOdbcConnection, @@ -99,7 +99,7 @@ class ArrowOdbcDriverFeatures(TypedDict): """Per-statement ODBC query execution timeout in seconds.""" payload_text_encoding: NotRequired["TextEncoding"] """Wire text encoding override (``'auto'``, ``'utf8'``, or ``'utf16'``) passed to ``arrow-odbc``.""" - enable_driver_pooling: NotRequired[bool] + enable_pooling: NotRequired[bool] """Enable process-global ODBC driver-manager connection pooling prior to first connect (defaults to ``False``).""" connection_string: NotRequired[str] """Resolved connection string inspected for SQL dialect auto-detection.""" @@ -115,6 +115,8 @@ class ArrowOdbcDriverFeatures(TypedDict): """Deserializer parsing JSON text payloads (defaults to ``sqlspec.utils.serializers.from_json``).""" enable_events: NotRequired[bool] """Enable SQLSpec table-backed event channel support.""" + events_backend: NotRequired[Literal["poll_queue"]] + """Event queue backend name (``'poll_queue'``).""" on_connection_create: "NotRequired[Callable[[ArrowOdbcConnection], None]]" """Callback invoked with the raw ``arrow_odbc.Connection`` immediately after connect.""" @@ -239,7 +241,7 @@ def create_connection(self) -> "ArrowOdbcConnection": if self.connection_instance is not None: return cast("ArrowOdbcConnection", self.connection_instance) try: - _apply_driver_pooling(bool(self.driver_features.get("enable_driver_pooling", False))) + _apply_driver_pooling(bool(self.driver_features.get("enable_pooling", False))) dialect = _resolve_connection_dialect_hint(self.driver_features) connection_string, connect_kwargs = build_connection_config(self.connection_config, dialect=dialect) connection = cast("ArrowOdbcConnection", arrow_odbc_connect(connection_string, **connect_kwargs)) diff --git a/sqlspec/adapters/asyncpg/driver.py b/sqlspec/adapters/asyncpg/driver.py index 22c4977cd..e9788fca4 100644 --- a/sqlspec/adapters/asyncpg/driver.py +++ b/sqlspec/adapters/asyncpg/driver.py @@ -698,12 +698,7 @@ async def _run_stack_operations( async def _execute_stack_operation_prepared(self, normalized: "NormalizedStackOperation") -> StackResult: result = await self._execute_cached_statement(normalized.statement) - stmt_cache_size = self.driver_features.get("statement_cache_size") - if stmt_cache_size is None and hasattr(self, "connection_config") and isinstance(self.connection_config, dict): - stmt_cache_size = self.connection_config.get("statement_cache_size") - if stmt_cache_size is None: - stmt_cache_size = 100 - if not self._is_transaction_pooler_mode() and stmt_cache_size > 0: + if not self._is_transaction_pooler_mode(): result.metadata["prepared_statement"] = True return StackResult.from_sql_result(result) diff --git a/sqlspec/adapters/bigquery/config.py b/sqlspec/adapters/bigquery/config.py index d77bbd5d3..7a1d6338c 100644 --- a/sqlspec/adapters/bigquery/config.py +++ b/sqlspec/adapters/bigquery/config.py @@ -103,6 +103,8 @@ class BigQueryDriverFeatures(TypedDict): connection_instance: NotRequired["BigQueryConnection"] """Pre-created ``google.cloud.bigquery.Client`` reused without transferring lifecycle ownership.""" + default_query_job_config: NotRequired[QueryJobConfig] + """Default ``QueryJobConfig`` for a directly constructed driver; config-managed drivers use the client's default.""" on_connection_create: NotRequired["Callable[[Any], None]"] """Callback invoked with the ``bigquery.Client`` when instantiated by the config.""" json_serializer: NotRequired["Callable[[Any], str]"] @@ -119,7 +121,7 @@ class BigQueryDriverFeatures(TypedDict): """Total seconds to retry transient query failures (defaults to ``60.0``; ``<= 0`` disables retries).""" job_result_timeout: NotRequired[float] """Maximum seconds to poll ``QueryJob.result()`` before raising a timeout error.""" - use_query_and_wait: NotRequired[bool] + enable_query_and_wait: NotRequired[bool] """Execute single-statement queries via low-latency ``Client.query_and_wait()`` outside transactions (defaults to ``True``).""" request_timeout: NotRequired[float] """Per-request HTTP transport timeout in seconds for job submission calls (defaults to ``job_result_timeout`` or ``120.0``).""" diff --git a/sqlspec/adapters/bigquery/driver.py b/sqlspec/adapters/bigquery/driver.py index c2ce7f2cb..c37b5d942 100644 --- a/sqlspec/adapters/bigquery/driver.py +++ b/sqlspec/adapters/bigquery/driver.py @@ -117,6 +117,7 @@ class BigQueryDriver(SyncDriverAdapterBase): "_column_name_cache", "_data_dictionary", "_default_query_job_config", + "_enable_query_and_wait", "_enable_storage_write_api", "_in_transaction", "_job_result_kwargs_defaults", @@ -127,7 +128,6 @@ class BigQueryDriver(SyncDriverAdapterBase): "_literal_inliner", "_request_timeout", "_session_id", - "_use_query_and_wait", ) dialect = "bigquery" @@ -165,7 +165,7 @@ def __init__( self._job_retry: Retry | None = build_retry(self._job_retry_deadline) if self._job_retry_deadline > 0 else None self._job_result_timeout: float | object = features.get("job_result_timeout", POLLING_DEFAULT_VALUE) self._request_timeout = self._resolve_request_timeout(features) - self._use_query_and_wait = bool(features.get("use_query_and_wait", True)) + self._enable_query_and_wait = bool(features.get("enable_query_and_wait", True)) self._enable_storage_write_api = bool(features.get("enable_storage_write_api", False)) self._in_transaction: bool = False self._session_id: str | None = None @@ -952,7 +952,7 @@ def _qualify_table_name(self, table: str) -> str: def _can_use_query_and_wait(self, statement: "SQL", sql: str) -> bool: """Use the interactive fast path only for compatible query jobs.""" - if not self._use_query_and_wait or not hasattr(self.connection, "query_and_wait"): + if not self._enable_query_and_wait or not hasattr(self.connection, "query_and_wait"): return False if self._in_transaction or self._session_id is not None: return False diff --git a/tests/integration/adapters/_shared/_driver_type_system.py b/tests/integration/adapters/_shared/_driver_type_system.py index f0bd54bc0..f031e2023 100644 --- a/tests/integration/adapters/_shared/_driver_type_system.py +++ b/tests/integration/adapters/_shared/_driver_type_system.py @@ -146,7 +146,7 @@ class SourceEquivalenceCase: "fetch_concurrently", "query_timeout_sec", "payload_text_encoding", - "enable_driver_pooling", + "enable_pooling", "connection_string", "dbms_name", "connection_autocommit", @@ -198,7 +198,7 @@ class SourceEquivalenceCase: "events_backend", "job_retry_deadline", "job_result_timeout", - "use_query_and_wait", + "enable_query_and_wait", "request_timeout", "query_page_size", "query_max_results", diff --git a/tests/integration/adapters/_shared/behaviors.py b/tests/integration/adapters/_shared/behaviors.py index 7352d0e26..255a73d8e 100644 --- a/tests/integration/adapters/_shared/behaviors.py +++ b/tests/integration/adapters/_shared/behaviors.py @@ -4936,14 +4936,14 @@ def _bigquery_job_controls(driver: object, case: DriverCase) -> None: assert not hasattr(driver, "export_table_to_storage") execute_parameters = inspect.signature(sync_driver.execute).parameters - for parameter_name in ("job_config", "job_retry", "job_result_timeout", "request_timeout", "use_query_and_wait"): + for parameter_name in ("job_config", "job_retry", "job_result_timeout", "request_timeout", "enable_query_and_wait"): assert parameter_name not in execute_parameters assert getattr(driver, "_job_result_timeout") == 30.0 assert getattr(driver, "_job_retry_deadline") == 0.0 assert getattr(driver, "_job_retry") is None assert getattr(driver, "_request_timeout") == 15.0 - assert getattr(driver, "_use_query_and_wait") is True + assert getattr(driver, "_enable_query_and_wait") is True register_sync_extra_assertion("driver_features:bigquery_job_controls", DRIVER_FEATURES_SCOPE, _bigquery_job_controls) diff --git a/tests/unit/adapters/test_adapter_implementations.py b/tests/unit/adapters/test_adapter_implementations.py index f1d8af47c..b2d30336d 100644 --- a/tests/unit/adapters/test_adapter_implementations.py +++ b/tests/unit/adapters/test_adapter_implementations.py @@ -4,6 +4,7 @@ import importlib import inspect import operator +import re import sqlite3 import types from pathlib import Path @@ -733,3 +734,27 @@ def test_litestar_session_migration_resolves_store_class(adapter: str, config_na store_module = importlib.import_module(f"sqlspec.adapters.{adapter}.litestar.store") assert hasattr(store_module, store_name), f"{adapter}: expected {store_name} in litestar/store.py" + + +_DRIVER_FEATURE_READ_PATTERN = r'driver_features\s*(?:\.get\(|\.setdefault\(|\[)\s*"([A-Za-z]\w*)"' + + +@pytest.mark.parametrize( + "adapter", sorted(path.name for path in Path("sqlspec/adapters").iterdir() if (path / "config.py").exists()) +) +def test_driver_feature_keys_read_by_adapter_are_declared(adapter: str) -> None: + """Every literal driver_features key an adapter reads is annotated on its ``*DriverFeatures`` TypedDict.""" + read_keys: dict[str, str] = {} + for py_file in sorted(Path("sqlspec/adapters", adapter).rglob("*.py")): + for match in re.finditer(_DRIVER_FEATURE_READ_PATTERN, py_file.read_text(encoding="utf-8")): + read_keys.setdefault(match.group(1), str(py_file)) + config_module = importlib.import_module(f"sqlspec.adapters.{adapter}.config") + feature_types = [ + value + for name, value in vars(config_module).items() + if name.endswith("DriverFeatures") and isinstance(value, type) + ] + assert feature_types, f"{adapter} declares no *DriverFeatures TypedDict" + declared = {key for feature_type in feature_types for key in feature_type.__annotations__} + undeclared = {key: path for key, path in read_keys.items() if key not in declared} + assert not undeclared, f"{adapter} reads undeclared driver_features keys: {undeclared}" diff --git a/tests/unit/adapters/test_adbc/test_adbc_config_normalization.py b/tests/unit/adapters/test_adbc/test_adbc_config_normalization.py index d1541c822..f72dea3ee 100644 --- a/tests/unit/adapters/test_adbc/test_adbc_config_normalization.py +++ b/tests/unit/adapters/test_adbc/test_adbc_config_normalization.py @@ -7,7 +7,7 @@ import pytest from adbc_driver_bigquery import DatabaseOptions -from sqlspec.adapters.adbc import AdbcConfig, AdbcDriver +from sqlspec.adapters.adbc import AdbcConfig, AdbcDriver, AdbcDriverFeatures from sqlspec.adapters.adbc.config import AdbcConnectionParams from sqlspec.adapters.adbc.core import ( build_connection_config, @@ -511,11 +511,11 @@ def fail_connect(_self: AdbcConfig) -> None: def test_adbc_driver_standalone_applies_features_and_enforces_strict_coercion() -> None: - """Standalone AdbcDriver populates default driver_features, aligns dialect, and honors strict_type_coercion.""" + """Standalone AdbcDriver populates default driver_features, aligns dialect, and honors enable_strict_type_coercion.""" conn = MagicMock() conn.adbc_get_info.return_value = {"vendor_name": "duckdb", "driver_name": "duckdb"} - driver = AdbcDriver(conn, driver_features={"strict_type_coercion": True}) + driver = AdbcDriver(conn, driver_features={"enable_strict_type_coercion": True}) assert driver.dialect == "duckdb" assert str(driver.statement_config.dialect) == "duckdb" assert "json_serializer" in driver.driver_features @@ -806,13 +806,13 @@ def test_snowflake_lowercase_default() -> None: assert config.driver_features["enable_lowercase_column_names"] is True -def test_driver_features_alias_consolidation() -> None: - """AdbcConfig consolidates legacy alias pairs to enable_*.""" +def test_driver_features_have_no_legacy_aliases() -> None: + """AdbcDriverFeatures declares only the enable_* keys, and legacy alias keys are not translated.""" + assert "strict_type_coercion" not in AdbcDriverFeatures.__annotations__ + assert "arrow_extension_types" not in AdbcDriverFeatures.__annotations__ config = AdbcConfig( connection_config={"driver_name": "sqlite"}, - driver_features={"strict_type_coercion": True, "arrow_extension_types": False}, + driver_features={"enable_strict_type_coercion": True, "enable_arrow_extension_types": False}, ) assert config.driver_features["enable_strict_type_coercion"] is True assert config.driver_features["enable_arrow_extension_types"] is False - assert "strict_type_coercion" not in config.driver_features - assert "arrow_extension_types" not in config.driver_features diff --git a/tests/unit/adapters/test_arrow_odbc/test_config.py b/tests/unit/adapters/test_arrow_odbc/test_config.py index 730c9527c..8518c7d27 100644 --- a/tests/unit/adapters/test_arrow_odbc/test_config.py +++ b/tests/unit/adapters/test_arrow_odbc/test_config.py @@ -179,9 +179,9 @@ def test_driver_pooling_defaults_to_off(monkeypatch: pytest.MonkeyPatch) -> None def test_enable_driver_pooling_is_a_declared_feature() -> None: assert ( - "enable_driver_pooling" + "enable_pooling" in ArrowOdbcConfig( - connection_config={"dsn": "analytics"}, driver_features={"enable_driver_pooling": False} + connection_config={"dsn": "analytics"}, driver_features={"enable_pooling": False} ).driver_features ) @@ -315,7 +315,7 @@ def broken_pooling() -> None: raise RuntimeError(msg) monkeypatch.setattr(arrow_odbc_config, "enable_odbc_connection_pooling", broken_pooling) - config = ArrowOdbcConfig(connection_config={"dsn": "analytics"}, driver_features={"enable_driver_pooling": True}) + config = ArrowOdbcConfig(connection_config={"dsn": "analytics"}, driver_features={"enable_pooling": True}) with pytest.raises(ImproperConfigurationError, match="Could not configure arrow-odbc connection"): config.create_connection() diff --git a/tests/unit/adapters/test_bigquery/test_config.py b/tests/unit/adapters/test_bigquery/test_config.py index d317260fc..006336e98 100644 --- a/tests/unit/adapters/test_bigquery/test_config.py +++ b/tests/unit/adapters/test_bigquery/test_config.py @@ -114,7 +114,7 @@ def test_bigquery_config_routes_current_client_level_settings(monkeypatch: Monke def test_bigquery_config_typed_surfaces_do_not_advertise_inert_settings() -> None: """Typed public settings should only include options that SQLSpec routes.""" assert "credentials_path" not in BigQueryConnectionParams.__annotations__ - assert "use_query_and_wait" in BigQueryDriverFeatures.__annotations__ + assert "enable_query_and_wait" in BigQueryDriverFeatures.__annotations__ assert "enable_storage_write_api" in BigQueryDriverFeatures.__annotations__ assert "enable_storage_write_api" not in BigQueryConnectionParams.__annotations__ assert "on_job_start" not in BigQueryDriverFeatures.__annotations__ diff --git a/tests/unit/adapters/test_bigquery/test_job_controls.py b/tests/unit/adapters/test_bigquery/test_job_controls.py index 6a3f8db83..ffb43f073 100644 --- a/tests/unit/adapters/test_bigquery/test_job_controls.py +++ b/tests/unit/adapters/test_bigquery/test_job_controls.py @@ -130,7 +130,7 @@ def test_run_query_job_passes_job_id_prefix_only_without_job_id() -> None: def test_query_and_wait_used_when_enabled() -> None: connection = _RecordingConnection() connection.row_iterator = _RecordingRowIterator(rows=[{"v": 1}], schema=_schema("v")) - driver = BigQueryDriver(cast(Any, connection), driver_features={"use_query_and_wait": True}) + driver = BigQueryDriver(cast(Any, connection), driver_features={"enable_query_and_wait": True}) result = driver.execute("SELECT 1 AS v") @@ -144,7 +144,7 @@ def test_query_and_wait_used_when_enabled() -> None: def test_query_and_wait_dml_rowcount() -> None: connection = _RecordingConnection() connection.row_iterator = _RecordingRowIterator(num_dml_affected_rows=3) - driver = BigQueryDriver(cast(Any, connection), driver_features={"use_query_and_wait": True}) + driver = BigQueryDriver(cast(Any, connection), driver_features={"enable_query_and_wait": True}) result = driver.execute("UPDATE t SET v = 1") @@ -156,7 +156,7 @@ def test_query_and_wait_explicit_off() -> None: connection = _RecordingConnection() connection.query_job = _RecordingJob(rows=[{"v": 1}], schema=_schema("v")) connection.row_iterator = _RecordingRowIterator(rows=[{"v": 1}], schema=_schema("v")) - driver = BigQueryDriver(cast(Any, connection), driver_features={"use_query_and_wait": False}) + driver = BigQueryDriver(cast(Any, connection), driver_features={"enable_query_and_wait": False}) result = driver.execute("SELECT 1 AS v") @@ -345,7 +345,11 @@ def test_select_to_storage_uses_existing_export_surface_and_job_controls(tmp_pat connection.query_job = _RecordingJob(rows=[{"id": 1}], schema=_schema("id")) driver = BigQueryDriver( cast(Any, connection), - driver_features={"job_result_timeout": 5.0, "storage_capabilities": CAPABILITIES, "use_query_and_wait": False}, + driver_features={ + "job_result_timeout": 5.0, + "storage_capabilities": CAPABILITIES, + "enable_query_and_wait": False, + }, ) job = driver.select_to_storage("SELECT 1 AS id", tmp_path / "result.parquet", format_hint="parquet") From 8f8c4b83a180d5df93b463346b22f56e6caba623 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 15:31:25 +0000 Subject: [PATCH 126/174] test(sqlanywhere): add integration driver test suite with hermetic client runtime --- docs/changelog.md | 4 + tests/integration/adapters/conftest.py | 1 + .../adapters/sqlanywhere/__init__.py | 1 + .../adapters/sqlanywhere/test_driver.py | 99 +++++++++++++++++ tests/integration/conftest.py | 2 + tests/integration/fixtures/sqlanywhere.py | 102 ++++++++++++++++++ uv.lock | 2 +- 7 files changed, 210 insertions(+), 1 deletion(-) create mode 100644 tests/integration/adapters/sqlanywhere/__init__.py create mode 100644 tests/integration/adapters/sqlanywhere/test_driver.py create mode 100644 tests/integration/fixtures/sqlanywhere.py diff --git a/docs/changelog.md b/docs/changelog.md index 23b24d80d..12faf3ed8 100644 --- a/docs/changelog.md +++ b/docs/changelog.md @@ -33,6 +33,10 @@ important operational fixes. * Removed the ADBC `strict_type_coercion` and `arrow_extension_types` driver feature aliases; use `enable_strict_type_coercion` and `enable_arrow_extension_types`. * Declared driver feature keys that adapters already read: `events_backend` on `ArrowOdbcDriverFeatures` and `default_query_job_config` on `BigQueryDriverFeatures`. * Replaced `null_pool` in `psycopg` connection configuration and driver features with canonical `enable_pooling=False`. +* Defaulted `enable_correlation_middleware` to `True` across all web framework extensions (`Litestar`, `FastAPI`, `Starlette`, `Sanic`, `Flask`). +* Standardized adapter driver imports: importing `sqlspec.adapters.` eagerly imports that adapter's driver, raising the native `ModuleNotFoundError` when the driver is not installed, while non-driver optional dependencies load on first use. +* Documented Oracle AQ dequeue commit behavior as at-most-once delivery. +* Added SAP SQL Anywhere integration driver test suite with hermetic in-container client runtime extraction. * Replaced correlated per-column scalar subqueries in data dictionary `columns.sql` and `indexes.sql` catalog queries across `postgres`, `cockroach`, `mssql`, `mysql`, `oracle`, `db2`, and `duckdb` with pre-aggregated CTEs and joins, and reduced `ColumnMetadata` to its 18 core catalog fields. ([#866](https://github.com/litestar-org/sqlspec/pull/866)) * Redesigned the documentation site architecture around Zensical and `mkdocstrings-python` with Diataxis-aligned navigation (`Getting Started`, `Learn`, `Adapters`, `Extensions`, `Examples`, `Reference`), interactive WebAssembly SQL playground, runnable `docs/examples/` snippets, conceptual Mermaid diagrams, and tabbed package manager and `SQLFileLoader` examples. diff --git a/tests/integration/adapters/conftest.py b/tests/integration/adapters/conftest.py index 63acf0525..3fa52d47c 100644 --- a/tests/integration/adapters/conftest.py +++ b/tests/integration/adapters/conftest.py @@ -17,6 +17,7 @@ "oracle": frozenset({"oracledb"}), "postgres": frozenset({"asyncpg", "psqlpy", "psycopg"}), "spanner": frozenset({"spanner"}), + "sqlanywhere": frozenset({"sqlanywhere"}), "sqlite": frozenset({"aiosqlite", "sqlite"}), } _SHARED_TEST_FAMILIES = frozenset((*_FAMILY_ADAPTERS, "gizmosql", "snowflake")) diff --git a/tests/integration/adapters/sqlanywhere/__init__.py b/tests/integration/adapters/sqlanywhere/__init__.py new file mode 100644 index 000000000..05cb6dbac --- /dev/null +++ b/tests/integration/adapters/sqlanywhere/__init__.py @@ -0,0 +1 @@ +"""SAP SQL Anywhere integration tests.""" diff --git a/tests/integration/adapters/sqlanywhere/test_driver.py b/tests/integration/adapters/sqlanywhere/test_driver.py new file mode 100644 index 000000000..8cd171d9a --- /dev/null +++ b/tests/integration/adapters/sqlanywhere/test_driver.py @@ -0,0 +1,99 @@ +"""SAP SQL Anywhere driver round trips, transactions, and exception mapping.""" + +from collections.abc import Generator + +import pytest + +pytest.importorskip("sqlanydb") + +from sqlspec.adapters.sqlanywhere import SqlanywhereConfig, SqlanywhereDriver +from sqlspec.exceptions import NotNullViolationError, UniqueViolationError + +pytestmark = [pytest.mark.sqlanywhere, pytest.mark.xdist_group("sqlanywhere")] + +_TABLE = "sqlspec_sa_items" + + +@pytest.fixture +def items_table(sqlanywhere_session: SqlanywhereDriver) -> Generator[str, None, None]: + """Create a fresh items table and drop it after the test.""" + sqlanywhere_session.execute_script(f"DROP TABLE IF EXISTS {_TABLE}") + sqlanywhere_session.execute_script( + f"CREATE TABLE {_TABLE} (id INTEGER NOT NULL PRIMARY KEY, name VARCHAR(50) NOT NULL, note LONG VARCHAR NULL)" + ) + sqlanywhere_session.commit() + try: + yield _TABLE + finally: + sqlanywhere_session.rollback() + sqlanywhere_session.execute_script(f"DROP TABLE IF EXISTS {_TABLE}") + sqlanywhere_session.commit() + + +def test_select_scalar(sqlanywhere_session: SqlanywhereDriver) -> None: + assert sqlanywhere_session.select_value("SELECT 1") == 1 + + +def test_parameterized_round_trip(sqlanywhere_session: SqlanywhereDriver, items_table: str) -> None: + result = sqlanywhere_session.execute( + f"INSERT INTO {items_table} (id, name, note) VALUES (?, ?, ?)", (1, "alpha", "x" * 5000) + ) + assert result.rows_affected == 1 + sqlanywhere_session.execute_many( + f"INSERT INTO {items_table} (id, name, note) VALUES (?, ?, ?)", [(2, "beta", None), (3, "gamma", "short")] + ) + sqlanywhere_session.commit() + + rows = sqlanywhere_session.select(f"SELECT id, name, note FROM {items_table} ORDER BY id") + assert [row["name"] for row in rows] == ["alpha", "beta", "gamma"] + assert rows[0]["note"] == "x" * 5000 + assert rows[1]["note"] is None + assert sqlanywhere_session.select_value(f"SELECT name FROM {items_table} WHERE id = ?", (3,)) == "gamma" + + +def test_rollback_discards_uncommitted_rows(sqlanywhere_session: SqlanywhereDriver, items_table: str) -> None: + sqlanywhere_session.begin() + sqlanywhere_session.execute(f"INSERT INTO {items_table} (id, name) VALUES (?, ?)", (10, "pending")) + assert sqlanywhere_session.select_value(f"SELECT COUNT(*) FROM {items_table}") == 1 + sqlanywhere_session.rollback() + assert sqlanywhere_session.select_value(f"SELECT COUNT(*) FROM {items_table}") == 0 + + +def test_commit_is_visible_to_new_connection( + sqlanywhere_config: SqlanywhereConfig, sqlanywhere_session: SqlanywhereDriver, items_table: str +) -> None: + sqlanywhere_session.begin() + sqlanywhere_session.execute(f"INSERT INTO {items_table} (id, name) VALUES (?, ?)", (20, "kept")) + sqlanywhere_session.commit() + + connection = sqlanywhere_config.create_connection() + try: + cursor = connection.cursor() + cursor.execute(f"SELECT name FROM {items_table} WHERE id = 20") + assert cursor.fetchone()[0] == "kept" + finally: + connection.close() + + +def test_unique_violation_is_mapped(sqlanywhere_session: SqlanywhereDriver, items_table: str) -> None: + sqlanywhere_session.execute(f"INSERT INTO {items_table} (id, name) VALUES (?, ?)", (30, "first")) + with pytest.raises(UniqueViolationError): + sqlanywhere_session.execute(f"INSERT INTO {items_table} (id, name) VALUES (?, ?)", (30, "duplicate")) + + +def test_not_null_violation_is_mapped(sqlanywhere_session: SqlanywhereDriver, items_table: str) -> None: + with pytest.raises(NotNullViolationError): + sqlanywhere_session.execute(f"INSERT INTO {items_table} (id, name) VALUES (?, ?)", (40, None)) + + +def test_data_dictionary_reports_version_and_tables(sqlanywhere_session: SqlanywhereDriver, items_table: str) -> None: + dictionary = sqlanywhere_session.data_dictionary + version = dictionary.get_version(sqlanywhere_session) + assert version is not None + assert version.major >= 16 + table_names = {table["table_name"].lower() for table in dictionary.get_tables(sqlanywhere_session)} + assert items_table in table_names + column_names = [ + column["column_name"].lower() for column in dictionary.get_columns(sqlanywhere_session, items_table) + ] + assert column_names == ["id", "name", "note"] diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index 7b67de289..8bcf3d721 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -27,6 +27,7 @@ "pytest_databases.docker.rustfs", "pytest_databases.docker.snowflake", "pytest_databases.docker.spanner", + "pytest_databases.docker.sqlanywhere", ) _NON_PLUGIN_FIXTURE_MODULES: frozenset[str] = frozenset({"cleanup"}) @@ -48,6 +49,7 @@ "spanner_adk": (("google.cloud.spanner",),), "spanner_events": (("google.cloud.spanner",),), "spanner_litestar": (("google.cloud.spanner", "litestar"),), + "sqlanywhere": (("sqlanydb",),), "sqlite": ((),), } diff --git a/tests/integration/fixtures/sqlanywhere.py b/tests/integration/fixtures/sqlanywhere.py new file mode 100644 index 000000000..6cf9589ac --- /dev/null +++ b/tests/integration/fixtures/sqlanywhere.py @@ -0,0 +1,102 @@ +"""SAP SQL Anywhere integration test fixtures.""" + +import ctypes +import io +import os +import tarfile +from collections.abc import Generator +from pathlib import Path +from typing import TYPE_CHECKING, Any + +import pytest + +from sqlspec.adapters.sqlanywhere import SqlanywhereConfig + +if TYPE_CHECKING: + from pytest_databases.docker.sqlanywhere import SQLAnywhereService + + from sqlspec.adapters.sqlanywhere import SqlanywhereDriver + +__all__ = ( + "sqlanywhere_client_runtime", + "sqlanywhere_config", + "sqlanywhere_connection_config", + "sqlanywhere_session", + "xdist_sqlanywhere_isolation_level", +) + +_CLIENT_LIBRARY_DEPENDENCIES = ("libdbtasks16_r.so", "libdblib16_r.so") +_CLIENT_DIRECTORIES = ("lib64", "res") + + +@pytest.fixture(scope="session") +def xdist_sqlanywhere_isolation_level() -> str: + """Run one SQL Anywhere server per worker because the demo image cannot create databases on demand.""" + return "server" + + +@pytest.fixture(scope="session") +def sqlanywhere_client_runtime( + sqlanywhere_service: "SQLAnywhereService", tmp_path_factory: pytest.TempPathFactory +) -> Generator[Path, None, None]: + """Provide the SQL Anywhere ``dbcapi`` client runtime to ``sqlanydb``. + + ``sqlanydb`` is a ctypes wrapper over SAP's native ``dbcapi`` library, which is not on + PyPI. When ``SQLANY_API_DLL`` is not already configured, the runtime is copied out of the + server container, its dependent libraries are preloaded, and ``SQLANY16`` / + ``SQLANY_API_DLL`` are set before the first connection initializes ``sqlanydb``. + """ + configured = os.environ.get("SQLANY_API_DLL") + if configured: + yield Path(configured).parent + return + root = tmp_path_factory.mktemp("sqlanywhere-client") + install_root = os.environ.get("SQLANYWHERE_CLIENT_ROOT", "/opt/sqlanywhere16") + for directory in _CLIENT_DIRECTORIES: + chunks, _ = sqlanywhere_service.container.get_archive(f"{install_root}/{directory}") + with tarfile.open(fileobj=io.BytesIO(b"".join(chunks))) as archive: + archive.extractall(root, filter="data") + library_dir = root / "lib64" + for name in _CLIENT_LIBRARY_DEPENDENCIES: + ctypes.CDLL(str(library_dir / name), mode=ctypes.RTLD_GLOBAL) + previous = {key: os.environ.get(key) for key in ("SQLANY16", "SQLANY_API_DLL")} + os.environ["SQLANY16"] = str(root) + os.environ["SQLANY_API_DLL"] = str(library_dir / "libdbcapi_r.so") + try: + yield library_dir + finally: + for key, value in previous.items(): + if value is None: + os.environ.pop(key, None) + else: + os.environ[key] = value + + +@pytest.fixture(scope="session") +def sqlanywhere_connection_config( + sqlanywhere_service: "SQLAnywhereService", sqlanywhere_client_runtime: Path +) -> dict[str, Any]: + """Return ``sqlanydb`` connection parameters for the SQL Anywhere service.""" + return { + "uid": sqlanywhere_service.user, + "pwd": sqlanywhere_service.password, + "dbn": sqlanywhere_service.database, + "host": f"{sqlanywhere_service.host}:{sqlanywhere_service.port}", + } + + +@pytest.fixture +def sqlanywhere_config(sqlanywhere_connection_config: dict[str, Any]) -> Generator[SqlanywhereConfig, None, None]: + """Yield a SQL Anywhere config whose pool is closed after the test.""" + config = SqlanywhereConfig(connection_config=dict(sqlanywhere_connection_config)) + try: + yield config + finally: + config.close_pool() + + +@pytest.fixture +def sqlanywhere_session(sqlanywhere_config: SqlanywhereConfig) -> Generator["SqlanywhereDriver", None, None]: + """Yield a SQL Anywhere driver session.""" + with sqlanywhere_config.provide_session() as session: + yield session diff --git a/uv.lock b/uv.lock index 24930f0a3..c4609dc8b 100644 --- a/uv.lock +++ b/uv.lock @@ -5067,7 +5067,7 @@ wheels = [ [[package]] name = "pytest-databases" version = "0.19.0" -source = { git = "https://github.com/litestar-org/pytest-databases.git?branch=test%2Fclientless-all#9f6fc17c5703f618b5895fa43ce62e736e44c28b" } +source = { git = "https://github.com/litestar-org/pytest-databases.git?branch=test%2Fclientless-all#b758d5ef5967b8fa912753ce4b695d2758f87537" } dependencies = [ { name = "docker" }, { name = "filelock" }, From 607e77aeed14f1d206dad1e19ed2fe3cb54357b0 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 17:03:27 +0000 Subject: [PATCH 127/174] test(sqlanywhere): consume upstream pytest-databases client library fixture --- tests/integration/fixtures/sqlanywhere.py | 51 ++--------------------- uv.lock | 2 +- 2 files changed, 4 insertions(+), 49 deletions(-) diff --git a/tests/integration/fixtures/sqlanywhere.py b/tests/integration/fixtures/sqlanywhere.py index 6cf9589ac..3ef595d9a 100644 --- a/tests/integration/fixtures/sqlanywhere.py +++ b/tests/integration/fixtures/sqlanywhere.py @@ -1,33 +1,25 @@ """SAP SQL Anywhere integration test fixtures.""" -import ctypes -import io -import os -import tarfile from collections.abc import Generator from pathlib import Path from typing import TYPE_CHECKING, Any import pytest +from pytest_databases.docker.sqlanywhere import SQLAnywhereService, sqlanywhere_client_libs from sqlspec.adapters.sqlanywhere import SqlanywhereConfig if TYPE_CHECKING: - from pytest_databases.docker.sqlanywhere import SQLAnywhereService - from sqlspec.adapters.sqlanywhere import SqlanywhereDriver __all__ = ( - "sqlanywhere_client_runtime", + "sqlanywhere_client_libs", "sqlanywhere_config", "sqlanywhere_connection_config", "sqlanywhere_session", "xdist_sqlanywhere_isolation_level", ) -_CLIENT_LIBRARY_DEPENDENCIES = ("libdbtasks16_r.so", "libdblib16_r.so") -_CLIENT_DIRECTORIES = ("lib64", "res") - @pytest.fixture(scope="session") def xdist_sqlanywhere_isolation_level() -> str: @@ -35,46 +27,9 @@ def xdist_sqlanywhere_isolation_level() -> str: return "server" -@pytest.fixture(scope="session") -def sqlanywhere_client_runtime( - sqlanywhere_service: "SQLAnywhereService", tmp_path_factory: pytest.TempPathFactory -) -> Generator[Path, None, None]: - """Provide the SQL Anywhere ``dbcapi`` client runtime to ``sqlanydb``. - - ``sqlanydb`` is a ctypes wrapper over SAP's native ``dbcapi`` library, which is not on - PyPI. When ``SQLANY_API_DLL`` is not already configured, the runtime is copied out of the - server container, its dependent libraries are preloaded, and ``SQLANY16`` / - ``SQLANY_API_DLL`` are set before the first connection initializes ``sqlanydb``. - """ - configured = os.environ.get("SQLANY_API_DLL") - if configured: - yield Path(configured).parent - return - root = tmp_path_factory.mktemp("sqlanywhere-client") - install_root = os.environ.get("SQLANYWHERE_CLIENT_ROOT", "/opt/sqlanywhere16") - for directory in _CLIENT_DIRECTORIES: - chunks, _ = sqlanywhere_service.container.get_archive(f"{install_root}/{directory}") - with tarfile.open(fileobj=io.BytesIO(b"".join(chunks))) as archive: - archive.extractall(root, filter="data") - library_dir = root / "lib64" - for name in _CLIENT_LIBRARY_DEPENDENCIES: - ctypes.CDLL(str(library_dir / name), mode=ctypes.RTLD_GLOBAL) - previous = {key: os.environ.get(key) for key in ("SQLANY16", "SQLANY_API_DLL")} - os.environ["SQLANY16"] = str(root) - os.environ["SQLANY_API_DLL"] = str(library_dir / "libdbcapi_r.so") - try: - yield library_dir - finally: - for key, value in previous.items(): - if value is None: - os.environ.pop(key, None) - else: - os.environ[key] = value - - @pytest.fixture(scope="session") def sqlanywhere_connection_config( - sqlanywhere_service: "SQLAnywhereService", sqlanywhere_client_runtime: Path + sqlanywhere_service: "SQLAnywhereService", sqlanywhere_client_libs: Path ) -> dict[str, Any]: """Return ``sqlanydb`` connection parameters for the SQL Anywhere service.""" return { diff --git a/uv.lock b/uv.lock index c4609dc8b..7c368a751 100644 --- a/uv.lock +++ b/uv.lock @@ -5067,7 +5067,7 @@ wheels = [ [[package]] name = "pytest-databases" version = "0.19.0" -source = { git = "https://github.com/litestar-org/pytest-databases.git?branch=test%2Fclientless-all#b758d5ef5967b8fa912753ce4b695d2758f87537" } +source = { git = "https://github.com/litestar-org/pytest-databases.git?branch=test%2Fclientless-all#ee262ac973907bdefe539cbfee4e5630f380724f" } dependencies = [ { name = "docker" }, { name = "filelock" }, From 2ee22f8ba8ce2c69660e3afe979c104a1eeb57b4 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 18:57:42 +0000 Subject: [PATCH 128/174] fix(adapters): preserve transaction, streaming, and configuration contracts --- docs/changelog.md | 47 ++- docs/integrations/fastapi.md | 4 +- docs/integrations/flask.md | 4 +- docs/integrations/starlette.md | 4 +- .../data-dictionary/dialects/cockroachdb.md | 90 +++--- .../reference/data-dictionary/dialects/db2.md | 28 +- .../data-dictionary/dialects/duckdb.md | 3 +- .../data-dictionary/dialects/index.md | 6 +- .../data-dictionary/dialects/mysql.md | 4 +- .../data-dictionary/dialects/oracle.md | 28 +- .../data-dictionary/dialects/postgres.md | 36 ++- .../data-dictionary/dialects/sqlanywhere.md | 2 +- .../data-dictionary/dialects/sqlite.md | 14 +- sqlspec/adapters/adbc/core.py | 21 +- sqlspec/adapters/adbc/driver.py | 3 +- sqlspec/adapters/adbc/litestar/store.py | 2 +- sqlspec/adapters/aiomysql/core.py | 2 - sqlspec/adapters/arrow_odbc/adk/store.py | 4 + sqlspec/adapters/asyncmy/config.py | 11 + sqlspec/adapters/asyncpg/core.py | 24 +- sqlspec/adapters/asyncpg/driver.py | 61 ++-- sqlspec/adapters/bigquery/adk/store.py | 2 + .../adapters/cockroach_asyncpg/adk/store.py | 4 + sqlspec/adapters/cockroach_asyncpg/driver.py | 1 + .../adapters/cockroach_psycopg/adk/store.py | 8 + sqlspec/adapters/cockroach_psycopg/driver.py | 2 + sqlspec/adapters/db2/adk/store.py | 19 +- sqlspec/adapters/db2/driver.py | 1 + sqlspec/adapters/db2/litestar/store.py | 5 +- sqlspec/adapters/duckdb/adk/store.py | 12 +- sqlspec/adapters/duckdb/data_dictionary.py | 4 +- sqlspec/adapters/duckdb/driver.py | 15 +- sqlspec/adapters/mssql_python/adk/store.py | 4 + sqlspec/adapters/mssql_python/driver.py | 48 ++-- sqlspec/adapters/mysqlconnector/_typing.py | 9 +- sqlspec/adapters/oracledb/adk/store.py | 106 ++++--- sqlspec/adapters/oracledb/config.py | 68 +++-- sqlspec/adapters/oracledb/core.py | 2 + sqlspec/adapters/oracledb/driver.py | 23 +- sqlspec/adapters/oracledb/litestar/store.py | 4 +- sqlspec/adapters/psqlpy/adk/store.py | 31 +- sqlspec/adapters/psqlpy/driver.py | 52 ++-- sqlspec/adapters/psycopg/config.py | 28 +- sqlspec/adapters/psycopg/driver.py | 138 ++++----- sqlspec/adapters/pymssql/adk/store.py | 4 + sqlspec/adapters/pymysql/core.py | 2 - sqlspec/adapters/spanner/config.py | 18 ++ sqlspec/adapters/sqlanywhere/_typing.py | 5 +- sqlspec/adapters/sqlanywhere/adk/store.py | 4 + sqlspec/adapters/sqlite/_typing.py | 2 +- sqlspec/adapters/sqlite/pool.py | 2 - sqlspec/adapters/sybase/adk/store.py | 4 + .../dialects/duckdb/sql/tables.sql | 1 + .../dialects/mssql/sql/constraints.sql | 44 ++- .../dialects/mssql/sql/indexes.sql | 88 +++--- .../dialects/mysql/dictionary.py | 88 ++---- .../dialects/oracle/sql/indexes.sql | 1 - sqlspec/extensions/adk/_table_utils.py | 16 +- sqlspec/extensions/adk/memory/store.py | 6 +- sqlspec/extensions/adk/store.py | 14 +- .../adapters/_shared/_driver_type_system.py | 269 ++---------------- .../integration/adapters/_shared/behaviors.py | 130 +++------ .../_shared/suite_driver_features_contract.py | 37 +-- .../extensions/adk/test_oracle_specific.py | 12 +- .../adapters/postgres/asyncpg/test_driver.py | 11 +- .../extensions/adk/test_memory_embeddings.py | 63 ++++ .../adapters/postgres/psqlpy/test_driver.py | 15 + .../psycopg/test_execution_regressions.py | 103 +++++++ .../adapters/sqlanywhere/test_driver.py | 2 +- .../adapters/sqlite/sqlite/test_pooling.py | 6 +- .../test_adbc/test_adk_dialect_support.py | 2 + tests/unit/adapters/test_adbc/test_core.py | 26 +- .../adapters/test_adbc/test_uuid_binding.py | 24 +- .../unit/adapters/test_asyncmy/test_config.py | 11 + .../test_asyncmy/test_load_from_arrow.py | 18 +- .../test_asyncpg/test_transaction_state.py | 84 +++++- .../test_asyncpg/test_type_handlers.py | 14 +- .../test_cockroach_psycopg/test_config.py | 8 +- tests/unit/adapters/test_db2/_fakes.py | 9 +- tests/unit/adapters/test_db2/test_adk.py | 50 +++- .../adapters/test_duckdb/test_adk_store.py | 22 ++ .../unit/adapters/test_duckdb/test_driver.py | 19 ++ .../adapters/test_mssql_python/test_arrow.py | 15 +- .../adapters/test_mysql_parameter_profiles.py | 6 +- .../test_mysqlconnector/test_adk_store.py | 9 +- .../test_mysqlconnector/test_config.py | 32 +++ .../test_oracledb/test_arrow_batches.py | 60 ++++ .../adapters/test_oracledb/test_config.py | 68 +++++ .../test_core_row_materialization.py | 2 +- .../test_oracledb/test_litestar_store.py | 30 +- .../test_oracledb/test_oracle_adk_store.py | 210 +++++++++++++- .../test_stack_native_transaction.py | 45 ++- .../test_postgres_migration_schema.py | 9 +- .../adapters/test_psqlpy/test_adk_store.py | 2 +- .../unit/adapters/test_psqlpy/test_driver.py | 31 +- .../unit/adapters/test_psycopg/test_config.py | 64 ++++- .../unit/adapters/test_psycopg/test_driver.py | 136 ++++----- .../test_pymssql/test_data_dictionary.py | 12 - .../unit/adapters/test_spanner/test_config.py | 26 ++ .../test_load_from_arrow_mutations.py | 12 +- tests/unit/adapters/test_sqlanywhere.py | 48 +--- tests/unit/adapters/test_sqlite/test_pool.py | 19 ++ .../data_dictionary/test_mysql_metadata.py | 77 +++-- .../test_sqlite_duckdb_metadata.py | 28 +- .../test_adk/test_store_instantiation.py | 51 ++++ .../test_extension_config_boundaries.py | 30 +- tests/unit/test_import_time.py | 54 ---- 107 files changed, 2109 insertions(+), 1185 deletions(-) create mode 100644 tests/integration/adapters/postgres/psqlpy/extensions/adk/test_memory_embeddings.py create mode 100644 tests/integration/adapters/postgres/psycopg/test_execution_regressions.py diff --git a/docs/changelog.md b/docs/changelog.md index 12faf3ed8..1af28609f 100644 --- a/docs/changelog.md +++ b/docs/changelog.md @@ -25,36 +25,81 @@ important operational fixes. **Changed:** + +* Renamed public Litestar stores: `SQLiteStore` to `SqliteStore`, `ADBCStore` to `AdbcStore`, and `DuckdbStore` to `DuckDBStore`. Update imports and store construction accordingly. DuckDB ADK exports likewise use `DuckDBADKStore`, `DuckDBADKMemoryStore`, `DuckDBADKConfig`, and `DuckDBADKFTSOptions` instead of their `Duckdb` spellings. + +* Removed `SybaseSyncEventQueueStore` and `SqlanywhereSyncEventQueueStore`; use `SybaseEventQueueStore` and `SqlanywhereEventQueueStore` respectively. + +* Standardized database failures as SQLSpec exceptions across statement execution, cached queries, bulk loads, and transaction boundaries. Code catching native driver errors should catch the corresponding `sqlspec.exceptions` classes instead. Constraint violations retain specific classes such as `UniqueViolationError` and `ForeignKeyViolationError`; syntax errors raise `SQLParsingError`. MySQL errors 1061 and 1091 now raise `OperationalError` instead of being suppressed. Original driver exceptions remain available as causes when the driver supplies an exception object. + +* Changed default SQLite `:memory:` pools to share one database across connections within that pool. Shared cache does not enable dirty reads by default; applications that deliberately need dirty reads must explicitly configure `read_uncommitted`. + +* Successful Psycopg `on_connection_create` hooks now commit their setup transaction before a connection enters the pool. Session settings such as `search_path` and timezone persist, and failed setup is rolled back. Hooks must not rely on their writes remaining uncommitted. + +* ADK table-name validation now rejects schema or dataset qualifiers where the selected adapter cannot generate valid qualified DDL. Supported schema-qualified stores retain that capability. + +* ADK schema changes target newly initialized stores. Existing ADK tables must be reset and recreated for this release; resetting removes their stored sessions and memories. Upgrade migrations for these schema changes are deferred. + * Removed legacy `pgbouncer` and `pgbouncer_mode` flags from PostgreSQL driver features across `asyncpg`, `psycopg`, and `psqlpy` in favor of canonical `pooler_mode="transaction"`. + * Removed the `pgbouncer` key from `AsyncpgConnectionConfig`; set `driver_features={"pooler_mode": "transaction"}` instead. + * Removed the empty `SybaseADKConfig`, `SqlanywhereADKConfig`, `Db2LitestarConfig`, and `Db2EventsConfig` TypedDicts; use the base `ADKConfig`, `LitestarConfig`, and `EventsConfig`. + * Renamed `use_pool` to `enable_pooling` (default `True`) across `mssql_python` and `pymssql` driver features. + * Renamed `arrow_odbc` `enable_driver_pooling` to `enable_pooling` and BigQuery `use_query_and_wait` to `enable_query_and_wait`. + * Removed the ADBC `strict_type_coercion` and `arrow_extension_types` driver feature aliases; use `enable_strict_type_coercion` and `enable_arrow_extension_types`. + * Declared driver feature keys that adapters already read: `events_backend` on `ArrowOdbcDriverFeatures` and `default_query_job_config` on `BigQueryDriverFeatures`. + * Replaced `null_pool` in `psycopg` connection configuration and driver features with canonical `enable_pooling=False`. + * Defaulted `enable_correlation_middleware` to `True` across all web framework extensions (`Litestar`, `FastAPI`, `Starlette`, `Sanic`, `Flask`). + * Standardized adapter driver imports: importing `sqlspec.adapters.` eagerly imports that adapter's driver, raising the native `ModuleNotFoundError` when the driver is not installed, while non-driver optional dependencies load on first use. + * Documented Oracle AQ dequeue commit behavior as at-most-once delivery. + * Added SAP SQL Anywhere integration driver test suite with hermetic in-container client runtime extraction. -* Replaced correlated per-column scalar subqueries in data dictionary `columns.sql` and `indexes.sql` catalog queries across `postgres`, `cockroach`, `mssql`, `mysql`, `oracle`, `db2`, and `duckdb` with pre-aggregated CTEs and joins, and reduced `ColumnMetadata` to its 18 core catalog fields. + +* Replaced correlated per-column scalar subqueries in data dictionary `columns.sql` and `indexes.sql` catalog queries across `postgres`, `cockroach`, `mysql`, `oracle`, `db2`, and `duckdb` with pre-aggregated CTEs and joins, and reduced `ColumnMetadata` to its 18 core catalog fields. ([#866](https://github.com/litestar-org/sqlspec/pull/866)) + * Redesigned the documentation site architecture around Zensical and `mkdocstrings-python` with Diataxis-aligned navigation (`Getting Started`, `Learn`, `Adapters`, `Extensions`, `Examples`, `Reference`), interactive WebAssembly SQL playground, runnable `docs/examples/` snippets, conceptual Mermaid diagrams, and tabbed package manager and `SQLFileLoader` examples. ([#864](https://github.com/litestar-org/sqlspec/pull/864), [#867](https://github.com/litestar-org/sqlspec/pull/867), [#869](https://github.com/litestar-org/sqlspec/pull/869)) + * Migrated integration test infrastructure to clientless `pytest-databases` fixtures (`#859`). ([#859](https://github.com/litestar-org/sqlspec/pull/859)) + * Upgraded CI GitHub Actions `astral-sh/setup-uv` from `10.3.1` to `10.3.3` and pinned `actions/checkout` (`6.0.2`), `actions/setup-python` (`7.1.0`), `actions/cache` (`6.0.0`), and `actions/upload-pages-artifact` / `actions/deploy-pages` in `.github/workflows/ci.yml` and `.github/workflows/docs.yml`. ([#872](https://github.com/litestar-org/sqlspec/pull/872)) **Fixed:** + +* Preserved supported Arrow conversion and mutation paths for Asyncmy and Spanner while reporting their native capabilities accurately; corrected ADBC lowercase readers, DuckDB stream ownership and table filtering, and SQL Server Arrow cursor lifetimes. + +* Corrected Oracle session-renewal bindings, pipeline error conversion, ADK batch-error handling, standalone connection setup, and deferred Arrow failures. Text-only Oracle memory stores remain available on servers without native vector support. + +* Preserved PostgreSQL cache-hit query semantics, affected-row counts, mapped exceptions, JSON scalar values, and ADK embeddings. Psycopg connection hooks leave pooled connections idle, COPY finalization errors are translated, and Asyncpg continue-on-error stacks recover through savepoints inside caller transactions. + +* Preserved SQL Server 2016 constraint and index metadata queries with compatible catalog aggregation. + +* Corrected SQL Anywhere implicit-transaction cleanup, Db2 ADK index-name matching and event deduplication, MySQL batch percent escaping, optional Connector/Python async-pool imports, privilege filters, and mutable DDL capture settings. + * Removed runtime `duckdb` imports from `sqlspec.adapters.adbc` (`AdbcDataDictionary` and `sqlspec/adapters/adbc/core.py`) and enforced strict adapter import isolation across all 20 database adapters in `tests/unit/test_import_time.py`. ([#865](https://github.com/litestar-org/sqlspec/pull/865)) + * Eliminated extra row copies and `UserWarning` in `DuckDBDriver.select_to_arrow` on DuckDB 2.0+ by routing parameterless `SELECT` queries through `connection.sql(query).to_arrow_table()` and `to_arrow_reader(batch_size=...)` when no open DB-API transaction is active on the connection. + * Fixed `compile_table_ddl_results` and `CreateIndex` cross-dialect DDL synthesis to emit `COMMENT ON` only on dialects that support it (`postgres`, `oracle`, `db2`, `duckdb`, `snowflake`), inline `COMMENT` clauses on `mysql` and `mariadb`, render `STORING (...)` on `spanner` and `spangres`, fold unquoted source identifiers to target case before quoting target-only reserved keywords, and preserve SQLite `NUMERIC` type affinity. + * Fixed `MigrationSquasher.fold_canonical_ddl` to fall back cleanly to raw statement concatenation when a migration contains interleaved DML or procedural blocks (`INSERT`, `UPDATE`, `DELETE`, `MERGE`, `COPY`, `DO`, `CALL`), preserve table-level constraints across `ALTER TABLE ... RENAME COLUMN` and `DROP COLUMN`, and update `sqlspec_migrations` tracker history atomically inside a single transaction when applying squashed migrations. + * Fixed `write_fixture_shards` to unlink stale shard and legacy single-file artifacts (`.json`, `
.json.gz`, and higher-numbered shards) before writing new shards so subsequent `load_fixture` calls do not read orphaned records. ### v0.69.0 - Oracle PKCS#12 wallet conversion, TNS alias resolution, proximity RRF, and server-side DBMS_CLOUD ingestion diff --git a/docs/integrations/fastapi.md b/docs/integrations/fastapi.md index a3b75d819..92188a1ab 100644 --- a/docs/integrations/fastapi.md +++ b/docs/integrations/fastapi.md @@ -58,9 +58,11 @@ Configure `extension_config["fastapi"]` (`FastAPIConfig`, which shares `Starlett | `commit_mode` | `"manual" | "autocommit" | "autocommit_include_redirect"` | `"manual"` | Controls automatic transaction commit or rollback based on response status codes. | | `extra_commit_statuses` | `set[int]` | `set()` | Additional HTTP status codes that trigger an automatic commit. | | `extra_rollback_statuses` | `set[int]` | `set()` | Additional HTTP status codes that trigger an automatic rollback. | -| `enable_correlation_middleware` | `bool` | `False` | Extracts or generates `X-Request-ID` correlation IDs into `CorrelationContext`. | +| `enable_correlation_middleware` | `bool` | `True` | Extracts or generates `X-Request-ID` correlation IDs into `CorrelationContext`. | | `enable_sqlcommenter_middleware` | `bool` | `False` | Populates `SQLCommenterContext` with FastAPI route, method, and trace attributes. | +Set `enable_correlation_middleware=False` to stop this plugin from reading request IDs or adding them to response headers. + ## Declarative Filter Dependencies FastAPI shares the same `FilterConfig` builder as Litestar via `sqlspec.extensions.fastapi.providers.create_filter_dependencies`, generating typed `Depends` callables for `LimitOffsetFilter`, `CursorFilter` (returning `HTTP 422` on malformed cursor tokens), `OrderByFilter`, `SearchFilter`, and `BeforeAfterFilter`. diff --git a/docs/integrations/flask.md b/docs/integrations/flask.md index d1181753d..48e0d2722 100644 --- a/docs/integrations/flask.md +++ b/docs/integrations/flask.md @@ -56,9 +56,11 @@ Configure distinct `session_key`, `connection_key`, and `pool_key` values under | `commit_mode` | `"manual" | "autocommit" | "autocommit_include_redirect"` | `"manual"` | Commits on `2xx` (`autocommit`) or `2xx`/`3xx` (`autocommit_include_redirect`) in `after_request` and rolls back on errors. | | `extra_commit_statuses` | `set[int]` | `set()` | Additional HTTP status codes that trigger an automatic commit. | | `extra_rollback_statuses` | `set[int]` | `set()` | Additional HTTP status codes that trigger an automatic rollback. | -| `enable_correlation_middleware` | `bool` | `False` | Binds request correlation headers into `CorrelationContext`. | +| `enable_correlation_middleware` | `bool` | `True` | Binds request correlation headers into `CorrelationContext`. | | `enable_sqlcommenter_middleware` | `bool` | `False` | Populates `SQLCommenterContext` (`framework="flask"`, route endpoint, method) for wire-level SQL comments. | +Set `enable_correlation_middleware=False` to stop this plugin from reading request IDs or adding them to response headers. + ## Recap - Initialize `SQLSpecPlugin` with `plugin.init_app(app)` and access request-scoped sessions via `plugin.get_session()`. diff --git a/docs/integrations/starlette.md b/docs/integrations/starlette.md index fd2b149b2..268f10b8c 100644 --- a/docs/integrations/starlette.md +++ b/docs/integrations/starlette.md @@ -54,9 +54,11 @@ Bind multiple async or sync configurations by specifying unique `session_key`, ` | `commit_mode` | `"manual" | "autocommit" | "autocommit_include_redirect"` | `"manual"` | Controls automatic commit/rollback middleware behavior based on response status codes. | | `extra_commit_statuses` | `set[int]` | `set()` | Additional HTTP status codes that trigger an automatic commit. | | `extra_rollback_statuses` | `set[int]` | `set()` | Additional HTTP status codes that trigger an automatic rollback. | -| `enable_correlation_middleware` | `bool` | `False` | Propagates `X-Request-ID` / `X-Correlation-ID` headers into `CorrelationContext`. | +| `enable_correlation_middleware` | `bool` | `True` | Propagates `X-Request-ID` / `X-Correlation-ID` headers into `CorrelationContext`. | | `enable_sqlcommenter_middleware` | `bool` | `False` | Populates `SQLCommenterContext` (`framework="starlette"`, route path, method). | +Set `enable_correlation_middleware=False` to stop this plugin from reading request IDs or adding them to response headers. + ## Recap - Call `db_plugin.init_app(app)` on your `Starlette` instance and retrieve sessions with `db_plugin.get_session(request)`. diff --git a/docs/reference/data-dictionary/dialects/cockroachdb.md b/docs/reference/data-dictionary/dialects/cockroachdb.md index ed9176f38..0bee49140 100644 --- a/docs/reference/data-dictionary/dialects/cockroachdb.md +++ b/docs/reference/data-dictionary/dialects/cockroachdb.md @@ -79,15 +79,16 @@ The `cockroachdb` dialect pack defines `18` named queries across `15` domain `.s ```sql -- name: by_schema -- dialect: cockroachdb - WITH target_schema AS ( + WITH resolved_table_schema AS ( + SELECT rn.nspname::text AS schema_name + FROM pg_catalog.pg_class rc + JOIN pg_catalog.pg_namespace rn ON rn.oid = rc.relnamespace + WHERE rc.oid = pg_catalog.to_regclass(pg_catalog.quote_ident(:table_name::text)) + ), + target_schema AS ( SELECT COALESCE( :schema_name::text, - ( - SELECT rn.nspname::text - FROM pg_catalog.pg_class rc - JOIN pg_catalog.pg_namespace rn ON rn.oid = rc.relnamespace - WHERE rc.oid = pg_catalog.to_regclass(pg_catalog.quote_ident(:table_name::text)) - ), + (SELECT schema_name FROM resolved_table_schema), current_schema() ) AS schema_name ), @@ -160,15 +161,16 @@ The `cockroachdb` dialect pack defines `18` named queries across `15` domain `.s -- name: by_table -- dialect: cockroachdb - WITH target_schema AS ( + WITH resolved_table_schema AS ( + SELECT rn.nspname::text AS schema_name + FROM pg_catalog.pg_class rc + JOIN pg_catalog.pg_namespace rn ON rn.oid = rc.relnamespace + WHERE rc.oid = pg_catalog.to_regclass(pg_catalog.quote_ident(:table_name::text)) + ), + target_schema AS ( SELECT COALESCE( :schema_name::text, - ( - SELECT rn.nspname::text - FROM pg_catalog.pg_class rc - JOIN pg_catalog.pg_namespace rn ON rn.oid = rc.relnamespace - WHERE rc.oid = pg_catalog.to_regclass(pg_catalog.quote_ident(:table_name::text)) - ), + (SELECT schema_name FROM resolved_table_schema), current_schema() ) AS schema_name ), @@ -260,6 +262,36 @@ The `cockroachdb` dialect pack defines `18` named queries across `15` domain `.s ```sql -- name: by_schema -- dialect: cockroachdb + WITH con_cols AS ( + SELECT + kcu.constraint_catalog, + kcu.constraint_schema, + kcu.constraint_name, + kcu.table_schema, + kcu.table_name, + array_agg(kcu.column_name::text ORDER BY kcu.ordinal_position) AS columns + FROM information_schema.key_column_usage kcu + WHERE kcu.table_schema = :schema_name + AND (:table_name::text IS NULL OR kcu.table_name = :table_name) + GROUP BY + kcu.constraint_catalog, + kcu.constraint_schema, + kcu.constraint_name, + kcu.table_schema, + kcu.table_name + ), + ref_cols AS ( + SELECT + pk_kcu.constraint_catalog, + pk_kcu.constraint_schema, + pk_kcu.constraint_name, + array_agg(pk_kcu.column_name::text ORDER BY pk_kcu.ordinal_position) AS referenced_columns + FROM information_schema.key_column_usage pk_kcu + GROUP BY + pk_kcu.constraint_catalog, + pk_kcu.constraint_schema, + pk_kcu.constraint_name + ) SELECT tc.constraint_catalog::text AS constraint_catalog, tc.constraint_schema::text AS constraint_schema, @@ -269,27 +301,17 @@ The `cockroachdb` dialect pack defines `18` named queries across `15` domain `.s tc.constraint_type::text AS constraint_type, tc.is_deferrable::text AS is_deferrable, tc.initially_deferred::text AS initially_deferred, - ARRAY( - SELECT kcu.column_name::text - FROM information_schema.key_column_usage kcu - WHERE kcu.constraint_catalog = tc.constraint_catalog - AND kcu.constraint_schema = tc.constraint_schema - AND kcu.constraint_name = tc.constraint_name - AND kcu.table_schema = tc.table_schema - AND kcu.table_name = tc.table_name - ORDER BY kcu.ordinal_position - )::text[] AS columns, + COALESCE(cc.columns, ARRAY[]::text[]) AS columns, rc.unique_constraint_schema::text AS referenced_schema, pk_tc.table_name::text AS referenced_table, - ARRAY( - SELECT pk_kcu.column_name::text - FROM information_schema.key_column_usage pk_kcu - WHERE pk_kcu.constraint_catalog = rc.unique_constraint_catalog - AND pk_kcu.constraint_schema = rc.unique_constraint_schema - AND pk_kcu.constraint_name = rc.unique_constraint_name - ORDER BY pk_kcu.ordinal_position - )::text[] AS referenced_columns + COALESCE(rc_cols.referenced_columns, ARRAY[]::text[]) AS referenced_columns FROM information_schema.table_constraints tc + LEFT JOIN con_cols cc + ON cc.constraint_catalog = tc.constraint_catalog + AND cc.constraint_schema = tc.constraint_schema + AND cc.constraint_name = tc.constraint_name + AND cc.table_schema = tc.table_schema + AND cc.table_name = tc.table_name LEFT JOIN information_schema.referential_constraints rc ON rc.constraint_catalog = tc.constraint_catalog AND rc.constraint_schema = tc.constraint_schema @@ -298,6 +320,10 @@ The `cockroachdb` dialect pack defines `18` named queries across `15` domain `.s ON pk_tc.constraint_catalog = rc.unique_constraint_catalog AND pk_tc.constraint_schema = rc.unique_constraint_schema AND pk_tc.constraint_name = rc.unique_constraint_name + LEFT JOIN ref_cols rc_cols + ON rc_cols.constraint_catalog = rc.unique_constraint_catalog + AND rc_cols.constraint_schema = rc.unique_constraint_schema + AND rc_cols.constraint_name = rc.unique_constraint_name WHERE tc.table_schema = :schema_name AND (:table_name::text IS NULL OR tc.table_name = :table_name) ORDER BY tc.table_schema, tc.table_name, tc.constraint_name; diff --git a/docs/reference/data-dictionary/dialects/db2.md b/docs/reference/data-dictionary/dialects/db2.md index ae1d7ffdc..9a4d3e983 100644 --- a/docs/reference/data-dictionary/dialects/db2.md +++ b/docs/reference/data-dictionary/dialects/db2.md @@ -16,7 +16,7 @@ Introspects IBM Db2 `SYSCAT` views (`SYSCAT.TABLES`, `SYSCAT.COLUMNS`, `SYSCAT.I | **Default Schema** | `None` | | **Version Pattern** | `(?:v|V)?(\d+)\.(\d+)(?:\.(\d+))?` | | **DDL Strategy** | AST synthesis + `SYSCAT` catalog text definitions | -| **Catalog `.sql` Files** | `9` files (`16` named queries) | +| **Catalog `.sql` Files** | `9` files (`14` named queries) | ### Capabilities & Type Mappings @@ -48,7 +48,7 @@ Introspects IBM Db2 `SYSCAT` views (`SYSCAT.TABLES`, `SYSCAT.COLUMNS`, `SYSCAT.I ## Catalog SQL Queries -The `db2` dialect pack defines `16` named queries across `9` domain `.sql` files loaded via `DataDictionaryLoader` (`sqlspec/data_dictionary/dialects/db2/`): +The `db2` dialect pack defines `14` named queries across `9` domain `.sql` files loaded via `DataDictionaryLoader` (`sqlspec/data_dictionary/dialects/db2/`): | Domain Pack | Named Queries | Bind Parameters | | --- | --- | --- | @@ -57,8 +57,8 @@ The `db2` dialect pack defines `16` named queries across `9` domain `.sql` files | `foreign_keys.sql` | `by_schema`, `by_table` | `:schema_name`, `:table_name` | | `indexes.sql` | `by_schema`, `by_table` | `:schema_name`, `:table_name` | | `objects.sql` | `by_schema` | `:object_name`, `:schema_name` | -| `schemas.sql` | `by_schema`, `current` | `:schema_name` | -| `tables.sql` | `by_schema`, `all_by_schema`, `names_by_schema` | `:schema_name`, `:table_name` | +| `schemas.sql` | `by_schema` | `:schema_name` | +| `tables.sql` | `by_schema`, `all_by_schema` | `:schema_name`, `:table_name` | | `version.sql` | `current` | *none* | | `views.sql` | `by_schema` | `:schema_name`, `:view_name` | @@ -369,7 +369,7 @@ The `db2` dialect pack defines `16` named queries across `9` domain `.sql` files ORDER BY schema_name, object_type, object_name; ``` -??? code "`schemas.sql` (2 named queries)" +??? code "`schemas.sql` (1 named queries)" ```sql -- name: by_schema @@ -385,13 +385,9 @@ The `db2` dialect pack defines `16` named queries across `9` domain `.sql` files AND s.SCHEMANAME NOT LIKE 'SQLJ%' AND (:schema_name IS NULL OR s.SCHEMANAME = :schema_name) ORDER BY s.SCHEMANAME; - - -- name: current - -- dialect: db2 - VALUES CURRENT SCHEMA; ``` -??? code "`tables.sql` (3 named queries)" +??? code "`tables.sql` (2 named queries)" ```sql -- name: by_schema @@ -471,18 +467,6 @@ The `db2` dialect pack defines `16` named queries across `9` domain `.sql` files AND t.TABSCHEMA = COALESCE(CAST(:schema_name AS VARCHAR(128)), CURRENT SCHEMA) AND (CAST(:table_name AS VARCHAR(128)) IS NULL OR t.TABNAME = :table_name) ORDER BY t.TABSCHEMA, t.TABNAME; - - -- name: names_by_schema - -- dialect: db2 - SELECT - RTRIM(t.TABNAME) AS table_name - FROM SYSCAT.TABLES t - WHERE t.TYPE = 'T' - AND t.TABSCHEMA NOT LIKE 'SYS%' - AND t.TABSCHEMA NOT LIKE 'NULLID%' - AND t.TABSCHEMA NOT LIKE 'SQLJ%' - AND t.TABSCHEMA = COALESCE(CAST(:schema_name AS VARCHAR(128)), CURRENT SCHEMA) - ORDER BY t.TABNAME; ``` ??? code "`version.sql` (1 named queries)" diff --git a/docs/reference/data-dictionary/dialects/duckdb.md b/docs/reference/data-dictionary/dialects/duckdb.md index 44505c5d5..c35da754f 100644 --- a/docs/reference/data-dictionary/dialects/duckdb.md +++ b/docs/reference/data-dictionary/dialects/duckdb.md @@ -72,7 +72,7 @@ The `duckdb` dialect pack defines `24` named queries across `17` domain `.sql` f | `schemas.sql` | `list` | *none* | | `sequences.sql` | `by_schema` | `:schema_name` | | `system.sql` | `logs`, `memory`, `settings` | *none* | -| `tables.sql` | `by_schema` | `:schema_name` | +| `tables.sql` | `by_schema` | `:schema_name`, `:table_name` | | `types.sql` | `by_schema` | `:schema_name` | | `version.sql` | `current` | *none* | | `views.sql` | `by_schema` | `:schema_name` | @@ -860,6 +860,7 @@ The `duckdb` dialect pack defines `24` named queries across `17` domain `.sql` f LEFT JOIN dependency_levels dl ON dl.table_name = t.table_name WHERE t.table_type = 'BASE TABLE' AND t.table_schema = COALESCE(:schema_name, current_schema()) + AND (:table_name IS NULL OR t.table_name = :table_name) ORDER BY level, t.table_name; ``` diff --git a/docs/reference/data-dictionary/dialects/index.md b/docs/reference/data-dictionary/dialects/index.md index 1f462c2da..391cfb602 100644 --- a/docs/reference/data-dictionary/dialects/index.md +++ b/docs/reference/data-dictionary/dialects/index.md @@ -16,14 +16,14 @@ SQLSpec ships 13 dialect-specific SQL catalog packs under `sqlspec/data_dictiona | **DuckDB** (`duckdb`) | `None` | Native `duckdb_*().sql` catalog functions + AST synthesis | `17` | `24` | [DuckDB Catalog](duckdb.md) | | **MySQL** (`mysql`) | `None` | Native `SHOW CREATE` + `information_schema` synthesis | `18` | `27` | [MySQL Catalog](mysql.md) | | **MariaDB** (`mariadb`) | `None` | Native `SHOW CREATE` + MariaDB sequence & constraint reflection | `17` | `27` | [MariaDB Catalog](mariadb.md) | -| **Oracle Database** (`oracle`) | `None` | Native `DBMS_METADATA.GET_DDL` + AST fallback | `22` | `35` | [Oracle Database Catalog](oracle.md) | +| **Oracle Database** (`oracle`) | `None` | Native `DBMS_METADATA.GET_DDL` + AST fallback | `22` | `36` | [Oracle Database Catalog](oracle.md) | | **SQL Server** (`mssql`) | `dbo` | AST synthesis + `OBJECT_DEFINITION()` for views/routines/triggers | `19` | `31` | [SQL Server Catalog](mssql.md) | | **Sybase ASE** (`sybase`) | `dbo` | AST synthesis + `dbo.syscomments` object text reassembly | `9` | `16` | [Sybase ASE Catalog](sybase.md) | | **SAP SQL Anywhere** (`sqlanywhere`) | `DBA` | AST synthesis + `SYS.SYSTAB` / `SYS.SYSPROCEDURE` definitions | `9` | `16` | [SAP SQL Anywhere Catalog](sqlanywhere.md) | -| **IBM Db2** (`db2`) | `None` | AST synthesis + `SYSCAT` catalog text definitions | `9` | `16` | [IBM Db2 Catalog](db2.md) | +| **IBM Db2** (`db2`) | `None` | AST synthesis + `SYSCAT` catalog text definitions | `9` | `14` | [IBM Db2 Catalog](db2.md) | | **Google BigQuery** (`bigquery`) | `None` | Native `INFORMATION_SCHEMA.TABLES.ddl` + AST fallback | `14` | `30` | [Google BigQuery Catalog](bigquery.md) | | **Google Cloud Spanner** (`spanner`) | `None` | AST synthesis with `INTERLEAVE IN PARENT` & `ROW DELETION POLICY` (`GoogleSQL` & `PostgreSQL` modes) | `5` | `8` | [Google Cloud Spanner Catalog](spanner.md) | -| **Total (`13` dialects)** | - | - | **`190`** | **`297`** | - | +| **Total (`13` dialects)** | - | - | **`190`** | **`296`** | - | ## Cross-Dialect Capability Matrix diff --git a/docs/reference/data-dictionary/dialects/mysql.md b/docs/reference/data-dictionary/dialects/mysql.md index 753b6109a..c5c0c6f02 100644 --- a/docs/reference/data-dictionary/dialects/mysql.md +++ b/docs/reference/data-dictionary/dialects/mysql.md @@ -85,7 +85,7 @@ The `mysql` dialect pack defines `27` named queries across `18` domain `.sql` fi | `objects.sql` | `by_schema` | `:schema_name` | | `partitions.sql` | `by_schema` | `:schema_name` | | `plugins.sql` | `list` | *none* | -| `privileges.sql` | `by_schema` | `:schema_name` | +| `privileges.sql` | `by_schema` | `:object_name`, `:schema_name` | | `routines.sql` | `by_schema` | `:schema_name` | | `schemas.sql` | `list` | `:schema_name` | | `system.sql` | `performance_schema_tables`, `sys_schema_table_statistics` | `:schema_name` | @@ -450,6 +450,7 @@ The `mysql` dialect pack defines `27` named queries across `18` domain `.sql` fi NULL AS column_name FROM information_schema.table_privileges WHERE table_schema = COALESCE(:schema_name, DATABASE()) + AND (:object_name IS NULL OR table_name = :object_name) UNION ALL SELECT grantee, @@ -461,6 +462,7 @@ The `mysql` dialect pack defines `27` named queries across `18` domain `.sql` fi column_name FROM information_schema.column_privileges WHERE table_schema = COALESCE(:schema_name, DATABASE()) + AND (:object_name IS NULL OR table_name = :object_name) ORDER BY schema_name, table_name, column_name, grantee, privilege_type; ``` diff --git a/docs/reference/data-dictionary/dialects/oracle.md b/docs/reference/data-dictionary/dialects/oracle.md index 491e11569..60b7f6743 100644 --- a/docs/reference/data-dictionary/dialects/oracle.md +++ b/docs/reference/data-dictionary/dialects/oracle.md @@ -16,7 +16,7 @@ Introspects `ALL_TABLES`, `ALL_TAB_COLS`, `ALL_INDEXES`, `ALL_IND_COLUMNS`, `ALL | **Default Schema** | `None` | | **Version Pattern** | `(\d+)` | | **DDL Strategy** | Native `DBMS_METADATA.GET_DDL` + AST fallback | -| **Catalog `.sql` Files** | `22` files (`35` named queries) | +| **Catalog `.sql` Files** | `22` files (`36` named queries) | ### Capabilities & Type Mappings @@ -53,7 +53,7 @@ Introspects `ALL_TABLES`, `ALL_TAB_COLS`, `ALL_INDEXES`, `ALL_IND_COLUMNS`, `ALL ## Catalog SQL Queries -The `oracle` dialect pack defines `35` named queries across `22` domain `.sql` files loaded via `DataDictionaryLoader` (`sqlspec/data_dictionary/dialects/oracle/`): +The `oracle` dialect pack defines `36` named queries across `22` domain `.sql` files loaded via `DataDictionaryLoader` (`sqlspec/data_dictionary/dialects/oracle/`): | Domain Pack | Named Queries | Bind Parameters | | --- | --- | --- | @@ -65,7 +65,7 @@ The `oracle` dialect pack defines `35` named queries across `22` domain `.sql` f | `dependencies.sql` | `by_schema` | `:object_name`, `:schema_name` | | `foreign_keys.sql` | `by_table`, `by_schema` | `:schema_name`, `:table_name` | | `grants.sql` | `by_schema` | `:object_name`, `:schema_name` | -| `indexes.sql` | `by_schema`, `by_table` | `:schema_name`, `:table_name` | +| `indexes.sql` | `by_schema`, `by_table`, `expressions` | `:schema_name`, `:table_name` | | `lob_storage.sql` | `by_schema` | `:schema_name`, `:table_name` | | `materialized_views.sql` | `by_schema` | `:mview_name`, `:schema_name` | | `objects.sql` | `by_schema`, `cloud_ai_available` | `:object_name`, `:schema_name` | @@ -425,7 +425,7 @@ The `oracle` dialect pack defines `35` named queries across `22` domain `.sql` f ORDER BY schema_name, object_name, column_name, grantee, privilege; ``` -??? code "`indexes.sql` (2 named queries)" +??? code "`indexes.sql` (3 named queries)" ```sql -- name: by_schema @@ -450,16 +450,11 @@ The `oracle` dialect pack defines `35` named queries across `22` domain `.sql` f i.avg_leaf_blocks_per_key, i.avg_data_blocks_per_key, i.clustering_factor, - LISTAGG(ic.column_name, ',') WITHIN GROUP (ORDER BY ic.column_position) AS columns, - LISTAGG(ie.column_expression, ',') WITHIN GROUP (ORDER BY ie.column_position) AS expressions + LISTAGG(ic.column_name, ',') WITHIN GROUP (ORDER BY ic.column_position) AS columns FROM all_indexes i LEFT JOIN all_ind_columns ic ON ic.index_owner = i.owner AND ic.index_name = i.index_name - LEFT JOIN all_ind_expressions ie - ON ie.index_owner = i.owner - AND ie.index_name = i.index_name - AND ie.column_position = ic.column_position WHERE i.owner = COALESCE(:schema_name, USER) AND (:table_name IS NULL OR i.table_name = :table_name) GROUP BY @@ -506,6 +501,19 @@ The `oracle` dialect pack defines `35` named queries across `22` domain `.sql` f WHERE i.table_name = :table_name AND i.owner = COALESCE(:schema_name, USER) GROUP BY i.index_name, i.table_name, i.uniqueness; + + -- name: expressions + -- dialect: oracle + SELECT + ie.index_owner AS schema_name, + ie.table_name, + ie.index_name, + ie.column_position, + ie.column_expression + FROM all_ind_expressions ie + WHERE ie.index_owner = COALESCE(:schema_name, USER) + AND (:table_name IS NULL OR ie.table_name = :table_name) + ORDER BY ie.index_owner, ie.table_name, ie.index_name, ie.column_position; ``` ??? code "`lob_storage.sql` (1 named queries)" diff --git a/docs/reference/data-dictionary/dialects/postgres.md b/docs/reference/data-dictionary/dialects/postgres.md index e7ea7c44e..7df5c70b7 100644 --- a/docs/reference/data-dictionary/dialects/postgres.md +++ b/docs/reference/data-dictionary/dialects/postgres.md @@ -257,6 +257,24 @@ The `postgres` dialect pack defines `26` named queries across `21` domain `.sql` ```sql -- name: by_schema -- dialect: postgres + WITH con_cols AS ( + SELECT + con.oid AS con_oid, + array_agg(att.attname::text ORDER BY keys.ord) AS columns + FROM pg_catalog.pg_constraint con + CROSS JOIN LATERAL unnest(con.conkey) WITH ORDINALITY AS keys(attnum, ord) + JOIN pg_catalog.pg_attribute att ON att.attrelid = con.conrelid AND att.attnum = keys.attnum + GROUP BY con.oid + ), + ref_cols AS ( + SELECT + con.oid AS con_oid, + array_agg(att.attname::text ORDER BY keys.ord) AS referenced_columns + FROM pg_catalog.pg_constraint con + CROSS JOIN LATERAL unnest(con.confkey) WITH ORDINALITY AS keys(attnum, ord) + JOIN pg_catalog.pg_attribute att ON att.attrelid = con.confrelid AND att.attnum = keys.attnum + GROUP BY con.oid + ) SELECT current_database()::text AS constraint_catalog, n.nspname::text AS constraint_schema, @@ -272,20 +290,10 @@ The `postgres` dialect pack defines `26` named queries across `21` domain `.sql` ELSE con.contype::text END AS constraint_type, pg_catalog.pg_get_constraintdef(con.oid, true)::text AS definition, - ARRAY( - SELECT att.attname::text - FROM unnest(con.conkey) WITH ORDINALITY AS keys(attnum, ord) - JOIN pg_catalog.pg_attribute att ON att.attrelid = con.conrelid AND att.attnum = keys.attnum - ORDER BY keys.ord - )::text[] AS columns, + COALESCE(cc.columns, ARRAY[]::text[]) AS columns, ref_ns.nspname::text AS referenced_schema, ref.relname::text AS referenced_table, - ARRAY( - SELECT att.attname::text - FROM unnest(con.confkey) WITH ORDINALITY AS keys(attnum, ord) - JOIN pg_catalog.pg_attribute att ON att.attrelid = con.confrelid AND att.attnum = keys.attnum - ORDER BY keys.ord - )::text[] AS referenced_columns, + COALESCE(rc.referenced_columns, ARRAY[]::text[]) AS referenced_columns, con.condeferrable AS is_deferrable, con.condeferred AS initially_deferred, dep.refobjid IS NOT NULL AS is_extension_owned @@ -295,6 +303,8 @@ The `postgres` dialect pack defines `26` named queries across `21` domain `.sql` LEFT JOIN pg_catalog.pg_class ref ON ref.oid = con.confrelid LEFT JOIN pg_catalog.pg_namespace ref_ns ON ref_ns.oid = ref.relnamespace LEFT JOIN pg_catalog.pg_depend dep ON dep.objid = con.oid AND dep.deptype = 'e' + LEFT JOIN con_cols cc ON cc.con_oid = con.oid + LEFT JOIN ref_cols rc ON rc.con_oid = con.oid WHERE n.nspname = :schema_name AND (:table_name::text IS NULL OR c.relname = :table_name) ORDER BY n.nspname, c.relname, con.conname; @@ -786,7 +796,7 @@ The `postgres` dialect pack defines `26` named queries across `21` domain `.sql` total_exec_time, mean_exec_time, rows - FROM pg_catalog.pg_stat_statements + FROM pg_stat_statements ORDER BY total_exec_time DESC LIMIT :limit; diff --git a/docs/reference/data-dictionary/dialects/sqlanywhere.md b/docs/reference/data-dictionary/dialects/sqlanywhere.md index 22c4fc4c6..e3e8de913 100644 --- a/docs/reference/data-dictionary/dialects/sqlanywhere.md +++ b/docs/reference/data-dictionary/dialects/sqlanywhere.md @@ -35,7 +35,7 @@ Introspects SAP SQL Anywhere `SYS` catalog views (`SYS.SYSTAB`, `SYS.SYSTABCOL`, | `supports_schemas` | `True` | `json` | `LONG VARCHAR` | | `supports_sequences` | `True` | `smallint` | `SMALLINT` | | `supports_skip_locked` | `True` | `string` | `VARCHAR(255)` | -| `supports_transactional_ddl` | `True` | `text` | `LONG VARCHAR` | +| `supports_transactional_ddl` | `False` | `text` | `LONG VARCHAR` | | `supports_transactions` | `True` | `time` | `TIME` | | `supports_update_from` | `True` | `timestamp` | `TIMESTAMP` | | `supports_upsert` | `True` | `uuid` | `UNIQUEIDENTIFIER` | diff --git a/docs/reference/data-dictionary/dialects/sqlite.md b/docs/reference/data-dictionary/dialects/sqlite.md index b66c029b1..e17f0de6b 100644 --- a/docs/reference/data-dictionary/dialects/sqlite.md +++ b/docs/reference/data-dictionary/dialects/sqlite.md @@ -107,7 +107,7 @@ The `sqlite` dialect pack defines `23` named queries across `15` domain `.sql` f -- name: by_table -- dialect: sqlite SELECT - :schema_name AS schema_name, + COALESCE(:schema_name, 'main') AS schema_name, :table_name AS table_name, ti.name AS column_name, ti.type AS data_type, @@ -128,7 +128,7 @@ The `sqlite` dialect pack defines `23` named queries across `15` domain `.sql` f END AS extra, ti.hidden IN (2, 3) AS is_generated, CASE WHEN ti.hidden = 1 THEN 1 ELSE 0 END AS is_hidden - FROM pragma_table_xinfo(:table_name, :schema_name) AS ti + FROM pragma_table_xinfo(:table_name, COALESCE(:schema_name, 'main')) AS ti LEFT JOIN {schema_prefix}sqlite_schema AS sm ON sm.type = 'table' AND sm.name = :table_name ORDER BY ti.cid; ``` @@ -348,7 +348,7 @@ The `sqlite` dialect pack defines `23` named queries across `15` domain `.sql` f -- name: by_table -- dialect: sqlite SELECT - :schema_name AS schema_name, + COALESCE(:schema_name, 'main') AS schema_name, :table_name AS table_name, il.name AS index_name, il."unique" AS is_unique, @@ -362,8 +362,8 @@ The `sqlite` dialect pack defines `23` named queries across `15` domain `.sql` f ix.coll AS collation, ix.key AS is_key_column, sm.sql AS index_sql - FROM pragma_index_list(:table_name, :schema_name) AS il - JOIN pragma_index_xinfo(il.name, :schema_name) AS ix + FROM pragma_index_list(:table_name, COALESCE(:schema_name, 'main')) AS il + JOIN pragma_index_xinfo(il.name, COALESCE(:schema_name, 'main')) AS ix LEFT JOIN {schema_prefix}sqlite_schema AS sm ON sm.type = 'index' AND sm.name = il.name @@ -373,7 +373,7 @@ The `sqlite` dialect pack defines `23` named queries across `15` domain `.sql` f -- name: columns_by_index -- dialect: sqlite SELECT - :schema_name AS schema_name, + COALESCE(:schema_name, 'main') AS schema_name, :index_name AS index_name, ix.seqno, ix.cid AS column_id, @@ -381,7 +381,7 @@ The `sqlite` dialect pack defines `23` named queries across `15` domain `.sql` f ix.desc AS is_descending, ix.coll AS collation, ix.key AS is_key_column - FROM pragma_index_xinfo(:index_name, :schema_name) AS ix + FROM pragma_index_xinfo(:index_name, COALESCE(:schema_name, 'main')) AS ix ORDER BY ix.seqno; -- name: index_columns_by_index diff --git a/sqlspec/adapters/adbc/core.py b/sqlspec/adapters/adbc/core.py index 2f1a3f863..9e4bcc545 100644 --- a/sqlspec/adapters/adbc/core.py +++ b/sqlspec/adapters/adbc/core.py @@ -8,7 +8,6 @@ from collections.abc import Sized from functools import partial from pathlib import Path -from types import SimpleNamespace from typing import TYPE_CHECKING, Any, Final, cast from uuid import UUID @@ -880,9 +879,10 @@ def create_mapped_exception(error: Any, *, logger: Any | None = None) -> SQLSpec avoids issues with exception control flow in different Python versions. Mapping priority: - 1. SQLSTATE codes (most reliable for ADBC drivers) - 2. Error message patterns - 3. Default SQLSpecError fallback + 1. Cancellation and timeout status codes + 2. SQLSTATE codes and vendor-specific diagnostics + 3. Error message patterns + 4. Generic ADBC status codes or SQLSpecError fallback Args: error: The ADBC exception to map @@ -892,7 +892,8 @@ def create_mapped_exception(error: Any, *, logger: Any | None = None) -> SQLSpec A SQLSpec exception that wraps the original error """ _ = logger - if status_mapped := _map_adbc_status_exception(error): + status_mapped = _map_adbc_status_exception(error) + if isinstance(status_mapped, (QueryTimeoutError, OperationCancelledError)): return status_mapped raw_message = str(error) @@ -934,6 +935,9 @@ def create_mapped_exception(error: Any, *, logger: Any | None = None) -> SQLSpec if "connection" in error_msg or "connect" in error_msg: return _create_adbc_error(error, DatabaseConnectionError, "connection error") + if status_mapped is not None: + return status_mapped + return _create_adbc_error(error, SQLSpecError, "database error") @@ -1466,9 +1470,6 @@ def clear(self) -> None: self.misses = 0 self.cache.clear() - def info(self) -> Any: - return SimpleNamespace(hits=self.hits, misses=self.misses, maxsize=self.maxsize, currsize=len(self.cache)) - _REWRITE_CACHE_STATE: Final[_RewriteCacheState] = _RewriteCacheState(256) _REWRITE_UUID_CACHE: Final[ @@ -1547,10 +1548,6 @@ def _rewrite_postgres_uuid_placeholders( return result -setattr(_rewrite_postgres_uuid_placeholders, "cache_clear", _REWRITE_CACHE_STATE.clear) -setattr(_rewrite_postgres_uuid_placeholders, "cache_info", _REWRITE_CACHE_STATE.info) - - def _sequence_holds_uuid(values: "list[Any] | tuple[Any, ...]") -> bool: """Report whether a sequence's first non-null element is a UUID.""" for value in values: diff --git a/sqlspec/adapters/adbc/driver.py b/sqlspec/adapters/adbc/driver.py index ac5926457..3649990d3 100644 --- a/sqlspec/adapters/adbc/driver.py +++ b/sqlspec/adapters/adbc/driver.py @@ -690,9 +690,10 @@ def select_to_arrow( new_schema = pa.schema([ field.with_name(name) for field, name in zip(reader.schema, normalized_names, strict=True) ]) + source_reader = reader def _lower_batches() -> Any: - for batch in reader: + for batch in source_reader: yield batch.rename_columns(normalized_names) reader = pa.RecordBatchReader.from_batches(new_schema, _lower_batches()) diff --git a/sqlspec/adapters/adbc/litestar/store.py b/sqlspec/adapters/adbc/litestar/store.py index fba12c796..f1659617d 100644 --- a/sqlspec/adapters/adbc/litestar/store.py +++ b/sqlspec/adapters/adbc/litestar/store.py @@ -470,7 +470,7 @@ def _set(self, key: str, value: "str | bytes", expires_in: "int | timedelta | No INSERT OR REPLACE INTO {self._table_name} (session_id, data, expires_at) VALUES ({p1}, {p2}, {p3}) """ - elif is_postgres_dialect(dialect) or dialect == "duckdb": + else: if expires_at is None: sql = f""" INSERT INTO {self._table_name} (session_id, data, expires_at) diff --git a/sqlspec/adapters/aiomysql/core.py b/sqlspec/adapters/aiomysql/core.py index 99170021d..99ec7f546 100644 --- a/sqlspec/adapters/aiomysql/core.py +++ b/sqlspec/adapters/aiomysql/core.py @@ -524,8 +524,6 @@ def escape_literal_percent(sql: str, parameters: Any, validator: "ParameterValid """Escape literal percent characters before driver-side interpolation.""" if not parameters or "%" not in sql: return sql - if isinstance(parameters, (tuple, list)) and sql.count("%") == len(parameters): - return sql keep = {info.position for info in validator.extract_parameters(sql)} if len(keep) == sql.count("%"): return sql diff --git a/sqlspec/adapters/arrow_odbc/adk/store.py b/sqlspec/adapters/arrow_odbc/adk/store.py index c45f1de1f..25988fad2 100644 --- a/sqlspec/adapters/arrow_odbc/adk/store.py +++ b/sqlspec/adapters/arrow_odbc/adk/store.py @@ -56,6 +56,8 @@ class ArrowOdbcADKStore(BaseSyncADKStore["ArrowOdbcConfig"]): connector_name: ClassVar[str] = "arrow_odbc" __slots__ = ("_sql",) + _max_table_name_parts = 2 + def __init__(self, config: "ArrowOdbcConfig") -> None: super().__init__(config) self._sql = _adk_sql_for(config) @@ -434,6 +436,8 @@ class ArrowOdbcADKMemoryStore(BaseSyncADKMemoryStore["ArrowOdbcConfig"]): __slots__ = ("_sql",) + _max_table_name_parts = 2 + def __init__(self, config: "ArrowOdbcConfig") -> None: super().__init__(config) self._sql = _adk_sql_for(config) diff --git a/sqlspec/adapters/asyncmy/config.py b/sqlspec/adapters/asyncmy/config.py index 334f640e9..371fd8886 100644 --- a/sqlspec/adapters/asyncmy/config.py +++ b/sqlspec/adapters/asyncmy/config.py @@ -23,6 +23,7 @@ from sqlspec.exceptions import ImproperConfigurationError from sqlspec.extensions.events import EventRuntimeHints from sqlspec.utils.config_tools import normalize_connection_config, parse_mysql_dsn +from sqlspec.utils.module_loader import ensure_pyarrow if TYPE_CHECKING: import ssl @@ -31,6 +32,7 @@ from sqlspec.core import StatementConfig from sqlspec.observability import ObservabilityConfig + from sqlspec.storage import StorageCapabilities __all__ = ( @@ -458,6 +460,15 @@ def get_event_runtime_hints(self) -> "EventRuntimeHints": return EventRuntimeHints(poll_interval=0.25, lease_seconds=5, select_for_update=True, skip_locked=True) + def _build_storage_capabilities(self) -> "StorageCapabilities": + capabilities = super()._build_storage_capabilities() + arrow_ready = self._dependency_available(ensure_pyarrow) + capabilities["arrow_import_enabled"] = arrow_ready + capabilities["arrow_export_enabled"] = arrow_ready + capabilities["parquet_import_enabled"] = arrow_ready + capabilities["parquet_export_enabled"] = arrow_ready + return capabilities + async def _create_pool(self) -> "AsyncmyPool": """Create the actual async connection pool. diff --git a/sqlspec/adapters/asyncpg/core.py b/sqlspec/adapters/asyncpg/core.py index d5b2d8423..0aa4d2689 100644 --- a/sqlspec/adapters/asyncpg/core.py +++ b/sqlspec/adapters/asyncpg/core.py @@ -531,18 +531,23 @@ async def fetch_chunk(self) -> "list[dict[str, Any]]": return [dict(record) for record in records] async def close(self, error: bool = False) -> None: + cursor = self._cursor self._cursor = None transaction = self._transaction self._transaction = None - if transaction is not None: - try: + try: + if cursor is not None and getattr(cursor, "_portal_name", None): + # asyncpg has no public cursor close; CockroachDB requires closing its portal before transaction control. + await cursor._close_portal(self._timeout) + if transaction is not None: if error: await transaction.rollback() else: await transaction.commit() - except Exception: - with contextlib.suppress(Exception): - await transaction.rollback() + except BaseException: + with contextlib.suppress(Exception): + await self._driver.connection.close() + raise async def _start(self) -> None: already_in_tx = self._driver._connection_in_transaction() @@ -559,8 +564,11 @@ async def _start(self) -> None: transaction = self._transaction self._transaction = None if transaction is not None: - with contextlib.suppress(Exception): + try: await transaction.rollback() + except BaseException: + with contextlib.suppress(Exception): + await self._driver.connection.close() raise @@ -601,9 +609,7 @@ def _custom_type_coercions() -> "dict[type, Callable[[Any], Any]]": def _is_pre_serialized_json(value: str) -> bool: stripped = value.strip() - if not ( - (stripped.startswith("{") and stripped.endswith("}")) or (stripped.startswith("[") and stripped.endswith("]")) - ): + if not stripped or stripped[0] not in '{["-0123456789tfn': return False try: from_json(stripped) diff --git a/sqlspec/adapters/asyncpg/driver.py b/sqlspec/adapters/asyncpg/driver.py index e9788fca4..57bc78834 100644 --- a/sqlspec/adapters/asyncpg/driver.py +++ b/sqlspec/adapters/asyncpg/driver.py @@ -2,7 +2,7 @@ import re from collections.abc import Mapping -from contextlib import suppress +from contextlib import asynccontextmanager, nullcontext, suppress from io import BytesIO from itertools import chain from typing import TYPE_CHECKING, Any, Literal, cast @@ -39,14 +39,16 @@ BaseAsyncExceptionHandler, StackExecutionObserver, describe_stack_statement, + validate_savepoint_name, ) from sqlspec.exceptions import ImproperConfigurationError, SQLSpecError, StackExecutionError from sqlspec.utils.logging import get_logger from sqlspec.utils.text import normalize_identifier, quote_identifier from sqlspec.utils.type_guards import has_sqlstate +from sqlspec.utils.uuids import uuid4 if TYPE_CHECKING: - from collections.abc import Sequence + from collections.abc import AsyncIterator, Sequence from sqlspec.adapters.asyncpg._typing import AsyncpgConnection from sqlspec.core import ArrowResult, SQLResult, StatementConfig @@ -655,27 +657,32 @@ async def _run_stack_operations( """ for index, operation in enumerate(stack.operations): try: - normalized: NormalizedStackOperation | None = None - if operation.method == "execute": - kwargs = dict(operation.keyword_arguments) if operation.keyword_arguments else {} - statement_config = kwargs.pop("statement_config", None) - config = statement_config or self.statement_config - - sql_statement = self.prepare_statement( - operation.statement, operation.arguments, statement_config=config, kwargs=kwargs - ) - if not sql_statement.is_script and not sql_statement.is_many: - sql_text, prepared_parameters = self._compiled_sql(sql_statement, config) - prepared_parameters = cast("tuple[Any, ...] | dict[str, Any] | None", prepared_parameters) - normalized = NormalizedStackOperation( - operation=operation, statement=sql_statement, sql=sql_text, parameters=prepared_parameters + scope = self._stack_savepoint() if continue_on_error and caller_in_tx else nullcontext() + async with scope: + normalized: NormalizedStackOperation | None = None + if operation.method == "execute": + kwargs = dict(operation.keyword_arguments) if operation.keyword_arguments else {} + statement_config = kwargs.pop("statement_config", None) + config = statement_config or self.statement_config + + sql_statement = self.prepare_statement( + operation.statement, operation.arguments, statement_config=config, kwargs=kwargs ) + if not sql_statement.is_script and not sql_statement.is_many: + sql_text, prepared_parameters = self._compiled_sql(sql_statement, config) + prepared_parameters = cast("tuple[Any, ...] | dict[str, Any] | None", prepared_parameters) + normalized = NormalizedStackOperation( + operation=operation, + statement=sql_statement, + sql=sql_text, + parameters=prepared_parameters, + ) - if normalized is not None: - stack_result = await self._execute_stack_operation_prepared(normalized) - else: - result = await self._execute_stack_operation(operation) - stack_result = StackResult(result=result) + if normalized is not None: + stack_result = await self._execute_stack_operation_prepared(normalized) + else: + result = await self._execute_stack_operation(operation) + stack_result = StackResult(result=result) except Exception as exc: stack_error = StackExecutionError( index, @@ -696,6 +703,18 @@ async def _run_stack_operations( if continue_on_error and not caller_in_tx: await self._commit_stack_success() + @asynccontextmanager + async def _stack_savepoint(self) -> "AsyncIterator[None]": + name = validate_savepoint_name(f"sqlspec_stack_{uuid4().hex}") + await self.connection.execute(f"SAVEPOINT {name}") + try: + yield + except BaseException: + await self.connection.execute(f"ROLLBACK TO SAVEPOINT {name}") + raise + finally: + await self.connection.execute(f"RELEASE SAVEPOINT {name}") + async def _execute_stack_operation_prepared(self, normalized: "NormalizedStackOperation") -> StackResult: result = await self._execute_cached_statement(normalized.statement) if not self._is_transaction_pooler_mode(): diff --git a/sqlspec/adapters/bigquery/adk/store.py b/sqlspec/adapters/bigquery/adk/store.py index dcede8d21..876fec8b8 100644 --- a/sqlspec/adapters/bigquery/adk/store.py +++ b/sqlspec/adapters/bigquery/adk/store.py @@ -77,6 +77,8 @@ class BigQueryADKStore(BaseSyncADKStore[BigQueryConfig]): "_require_partition_filter", ) + _max_table_name_parts = 3 + def __init__(self, config: BigQueryConfig) -> None: """Initialize the BigQuery ADK store. diff --git a/sqlspec/adapters/cockroach_asyncpg/adk/store.py b/sqlspec/adapters/cockroach_asyncpg/adk/store.py index 13ed09b53..a3f591157 100644 --- a/sqlspec/adapters/cockroach_asyncpg/adk/store.py +++ b/sqlspec/adapters/cockroach_asyncpg/adk/store.py @@ -72,6 +72,8 @@ class CockroachAsyncpgADKStore(BaseAsyncADKStore["CockroachAsyncpgConfig"]): __slots__ = () + _max_table_name_parts = 2 + def __init__(self, config: "CockroachAsyncpgConfig") -> None: super().__init__(config) @@ -587,6 +589,8 @@ class CockroachAsyncpgADKMemoryStore(BaseAsyncADKMemoryStore["CockroachAsyncpgCo __slots__ = () + _max_table_name_parts = 2 + def __init__(self, config: "CockroachAsyncpgConfig") -> None: super().__init__(config) diff --git a/sqlspec/adapters/cockroach_asyncpg/driver.py b/sqlspec/adapters/cockroach_asyncpg/driver.py index 40ccc2751..72003e29f 100644 --- a/sqlspec/adapters/cockroach_asyncpg/driver.py +++ b/sqlspec/adapters/cockroach_asyncpg/driver.py @@ -312,6 +312,7 @@ def _native_storage_ready(self) -> bool: async def _execute_native_storage(self, command: str, parameters: "list[Any]") -> "list[dict[str, Any]]": handler = self.handle_database_exceptions() + rows: list[dict[str, Any]] = [] async def _fetch() -> list[dict[str, Any]]: return [dict(row) for row in await self.connection.fetch(command, *parameters)] diff --git a/sqlspec/adapters/cockroach_psycopg/adk/store.py b/sqlspec/adapters/cockroach_psycopg/adk/store.py index 266f0b65c..7ea16f120 100644 --- a/sqlspec/adapters/cockroach_psycopg/adk/store.py +++ b/sqlspec/adapters/cockroach_psycopg/adk/store.py @@ -210,6 +210,8 @@ class CockroachPsycopgAsyncADKStore(BaseAsyncADKStore["CockroachPsycopgAsyncConf __slots__ = () + _max_table_name_parts = 2 + def __init__(self, config: "CockroachPsycopgAsyncConfig") -> None: super().__init__(config) @@ -700,6 +702,8 @@ class CockroachPsycopgSyncADKStore(BaseSyncADKStore["CockroachPsycopgSyncConfig" __slots__ = () + _max_table_name_parts = 2 + def __init__(self, config: "CockroachPsycopgSyncConfig") -> None: super().__init__(config) @@ -1211,6 +1215,8 @@ class CockroachPsycopgAsyncADKMemoryStore(BaseAsyncADKMemoryStore["CockroachPsyc __slots__ = () + _max_table_name_parts = 2 + def __init__(self, config: "CockroachPsycopgAsyncConfig") -> None: super().__init__(config) @@ -1448,6 +1454,8 @@ class CockroachPsycopgSyncADKMemoryStore(BaseSyncADKMemoryStore["CockroachPsycop __slots__ = () + _max_table_name_parts = 2 + def __init__(self, config: "CockroachPsycopgSyncConfig") -> None: super().__init__(config) diff --git a/sqlspec/adapters/cockroach_psycopg/driver.py b/sqlspec/adapters/cockroach_psycopg/driver.py index 576c72f04..654f82e99 100644 --- a/sqlspec/adapters/cockroach_psycopg/driver.py +++ b/sqlspec/adapters/cockroach_psycopg/driver.py @@ -397,6 +397,7 @@ def _native_storage_ready(self) -> bool: def _execute_native_storage(self, command: str, parameters: "list[Any]") -> "list[dict[str, Any]]": handler = self.handle_database_exceptions() + rows: list[dict[str, Any]] = [] def _fetch() -> list[dict[str, Any]]: with self.with_cursor(self.connection) as cursor: @@ -667,6 +668,7 @@ def _native_storage_ready(self) -> bool: async def _execute_native_storage(self, command: str, parameters: "list[Any]") -> "list[dict[str, Any]]": handler = self.handle_database_exceptions() + rows: list[dict[str, Any]] = [] async def _fetch() -> list[dict[str, Any]]: async with self.with_cursor(self.connection) as cursor: diff --git a/sqlspec/adapters/db2/adk/store.py b/sqlspec/adapters/db2/adk/store.py index 3c07003dc..08fab252d 100644 --- a/sqlspec/adapters/db2/adk/store.py +++ b/sqlspec/adapters/db2/adk/store.py @@ -24,6 +24,7 @@ ) from sqlspec.extensions.adk.memory import BaseAsyncADKMemoryStore, BaseSyncADKMemoryStore, StoredMemory from sqlspec.utils.serializers import from_json, to_json +from sqlspec.utils.text import quote_identifier if TYPE_CHECKING: from sqlspec.adapters.db2.driver import Db2AsyncDriver, Db2SyncDriver @@ -39,6 +40,8 @@ class _Db2ADKStoreMixin: __slots__ = () + _max_table_name_parts: ClassVar[int] = 2 + if TYPE_CHECKING: _session_table: str _events_table: str @@ -596,6 +599,8 @@ class _Db2ADKMemoryStoreMixin: __slots__ = () + _max_table_name_parts: ClassVar[int] = 2 + if TYPE_CHECKING: _memory_table: str @@ -653,11 +658,11 @@ def insert_memory_entries(self, entries: list[StoredMemory], owner_id: object | f"SELECT id, event_id FROM {self._memory_table} WHERE event_id IN ({placeholders})", tuple(e["event_id"] for e in unique_entries), ) - existing_ids = {str(row.get("event_id") or row.get("id") or next(iter(row.values()))) for row in existing_rows} + existing_ids = {str(row["event_id"] if "event_id" in row else row["EVENT_ID"]) for row in existing_rows} with self._config.provide_session() as driver: for entry in unique_entries: - if entry["event_id"] in existing_ids or entry["id"] in existing_ids: + if entry["event_id"] in existing_ids: continue driver.execute(insert_sql, _memory_insert_row(entry, has_owner, owner_id)) inserted_count += 1 @@ -773,11 +778,11 @@ async def insert_memory_entries(self, entries: list[StoredMemory], owner_id: obj f"SELECT id, event_id FROM {self._memory_table} WHERE event_id IN ({placeholders})", tuple(e["event_id"] for e in unique_entries), ) - existing_ids = {str(row.get("event_id") or row.get("id") or next(iter(row.values()))) for row in existing_rows} + existing_ids = {str(row["event_id"] if "event_id" in row else row["EVENT_ID"]) for row in existing_rows} async with self._config.provide_session() as driver: for entry in unique_entries: - if entry["event_id"] in existing_ids or entry["id"] in existing_ids: + if entry["event_id"] in existing_ids: continue await driver.execute(insert_sql, _memory_insert_row(entry, has_owner, owner_id)) inserted_count += 1 @@ -940,7 +945,7 @@ def _metadata_table_ddl(table: str) -> str: def _create_index_sql(table: str, index_name: str, columns: str) -> str: schema, _ = split_db2_table_name(table) if schema is not None: - return f'CREATE INDEX "{schema}"."{index_name}" ON {table} ({columns})' + return f"CREATE INDEX {quote_identifier(schema)}.{quote_identifier(index_name.upper())} ON {table} ({columns})" return f"CREATE INDEX {index_name} ON {table} ({columns})" @@ -1204,7 +1209,7 @@ def _create_missing(driver: "Db2SyncDriver", probe_sql: str, name: str, ddl: str if table is not None: table_schema, _ = split_db2_table_name(table) if table_schema is not None: - schema = table_schema.upper() + schema = table_schema if driver.select_one_or_none(probe_sql, (schema, object_name)) is None: driver.execute(ddl) @@ -1225,7 +1230,7 @@ async def _create_missing_async( if table is not None: table_schema, _ = split_db2_table_name(table) if table_schema is not None: - schema = table_schema.upper() + schema = table_schema if await driver.select_one_or_none(probe_sql, (schema, object_name)) is None: await driver.execute(ddl) diff --git a/sqlspec/adapters/db2/driver.py b/sqlspec/adapters/db2/driver.py index d08c065b4..398e3d159 100644 --- a/sqlspec/adapters/db2/driver.py +++ b/sqlspec/adapters/db2/driver.py @@ -348,6 +348,7 @@ def has_schema(self, schema: str) -> bool: Returns: True when ``SYSCAT.SCHEMATA`` lists the schema. """ + exists = False handler = self.handle_database_exceptions() with handler, self.with_cursor(self.connection) as cursor: cursor.execute("SELECT 1 FROM SYSCAT.SCHEMATA WHERE SCHEMANAME = ?", (normalize_identifier(schema, "db2"),)) diff --git a/sqlspec/adapters/db2/litestar/store.py b/sqlspec/adapters/db2/litestar/store.py index 49b0f6342..b688f7a3e 100644 --- a/sqlspec/adapters/db2/litestar/store.py +++ b/sqlspec/adapters/db2/litestar/store.py @@ -6,6 +6,7 @@ from sqlspec.adapters.db2.core import INDEX_EXISTS_SQL, TABLE_EXISTS_SQL, split_db2_table_name, to_db_timestamp, utc_now from sqlspec.extensions.litestar.store import BaseSQLSpecStore from sqlspec.utils.sync_tools import async_ +from sqlspec.utils.text import quote_identifier if TYPE_CHECKING: from sqlspec.adapters.db2.config import Db2AsyncConfig, Db2SyncConfig @@ -36,7 +37,7 @@ def session_index_name(table: str) -> str: Returns: The upper-folded index name. """ - return f"IX_{session_catalog_names(table)[1]}_EXP" + return f"IX_{session_catalog_names(table)[1]}_EXP".upper() def session_table_ddl(table: str) -> str: @@ -71,7 +72,7 @@ def session_index_ddl(table: str) -> str: schema, _ = split_db2_table_name(table) idx_name = session_index_name(table) if schema is not None: - return f'CREATE INDEX "{schema}"."{idx_name}" ON {table}(expires_at)' + return f"CREATE INDEX {quote_identifier(schema)}.{quote_identifier(idx_name)} ON {table}(expires_at)" return f"CREATE INDEX {idx_name} ON {table}(expires_at)" diff --git a/sqlspec/adapters/duckdb/adk/store.py b/sqlspec/adapters/duckdb/adk/store.py index a6679ac46..9e691e01b 100644 --- a/sqlspec/adapters/duckdb/adk/store.py +++ b/sqlspec/adapters/duckdb/adk/store.py @@ -98,9 +98,13 @@ def _split_duckdb_table_name(table_name: str) -> tuple[str, str]: Tuple of (schema, bare_table_name). If unqualified, schema defaults to 'main'. """ parts = split_qualified_identifier(table_name.strip(), quote_chars='"') + if not parts: + msg = "DuckDB table name must not be empty." + raise ValueError(msg) + first_part = parts[0] if len(parts) > 1: - return parts[0], parts[1] - return "main", parts[0] + return first_part, parts[1] + return "main", first_part class DuckDBADKStore(BaseSyncADKStore["DuckDBConfig"]): @@ -121,6 +125,8 @@ class DuckDBADKStore(BaseSyncADKStore["DuckDBConfig"]): __slots__ = () + _max_table_name_parts = 2 + def __init__(self, config: "DuckDBConfig") -> None: """Initialize DuckDB ADK store. @@ -911,6 +917,8 @@ class DuckDBADKMemoryStore(BaseSyncADKMemoryStore["DuckDBConfig"]): __slots__ = () + _max_table_name_parts = 2 + def __init__(self, config: "DuckDBConfig") -> None: """Initialize DuckDB ADK memory store. diff --git a/sqlspec/adapters/duckdb/data_dictionary.py b/sqlspec/adapters/duckdb/data_dictionary.py index b06c61b2b..0af956dc2 100644 --- a/sqlspec/adapters/duckdb/data_dictionary.py +++ b/sqlspec/adapters/duckdb/data_dictionary.py @@ -228,7 +228,9 @@ def get_tables(self, driver: "DuckDBDriver", schema: "str | None" = None) -> "li """Get tables sorted by topological dependency order using DuckDB catalog.""" schema_name = self.resolve_schema(schema) self._log_schema_introspect(driver, schema_name=schema_name, table_name=None, operation="tables") - return driver.select(self.get_query("tables", "by_schema"), schema_name=schema_name, schema_type=TableMetadata) + return driver.select( + self.get_query("tables", "by_schema"), schema_name=schema_name, table_name=None, schema_type=TableMetadata + ) def get_columns( self, driver: "DuckDBDriver", table: "str | None" = None, schema: "str | None" = None diff --git a/sqlspec/adapters/duckdb/driver.py b/sqlspec/adapters/duckdb/driver.py index 2f095c92a..aba2694f0 100644 --- a/sqlspec/adapters/duckdb/driver.py +++ b/sqlspec/adapters/duckdb/driver.py @@ -1043,7 +1043,10 @@ def __init__( self._close_callback = close_callback def read_next_batch(self) -> _DuckDBFetchManyBatch: - rows = self._target.fetchmany(self._chunk_size) + target = self._target + if target is None: + raise StopIteration + rows = target.fetchmany(self._chunk_size) if not rows: self.close() raise StopIteration @@ -1051,13 +1054,17 @@ def read_next_batch(self) -> _DuckDBFetchManyBatch: return _DuckDBFetchManyBatch(dict_rows) def close(self) -> None: - close_method = getattr(self._target, "close", None) - if callable(close_method): - close_method() + target = self._target + self._target = None if self._close_callback is not None: callback = self._close_callback self._close_callback = None callback() + elif target is not None: + # DuckDB 2 keeps execute results active until they are completed. + complete_method = getattr(target, "complete", None) + if callable(complete_method): + complete_method() def _ensure_arrow_reader(reader: Any) -> Any: diff --git a/sqlspec/adapters/mssql_python/adk/store.py b/sqlspec/adapters/mssql_python/adk/store.py index 14fab1116..61f24f058 100644 --- a/sqlspec/adapters/mssql_python/adk/store.py +++ b/sqlspec/adapters/mssql_python/adk/store.py @@ -48,6 +48,8 @@ class MssqlPythonADKStore(BaseSyncADKStore["MssqlPythonConfig"]): connector_name: ClassVar[str] = "mssql_python" __slots__ = ("_json_column_type", "_native_json") + _max_table_name_parts = 2 + def __init__(self, config: "MssqlPythonConfig") -> None: super().__init__(config) adk_config = _adk_config(config) @@ -432,6 +434,8 @@ class MssqlPythonADKMemoryStore(BaseSyncADKMemoryStore["MssqlPythonConfig"]): __slots__ = ("_json_column_type", "_native_json") + _max_table_name_parts = 2 + def __init__(self, config: "MssqlPythonConfig") -> None: super().__init__(config) adk_config = _adk_config(config) diff --git a/sqlspec/adapters/mssql_python/driver.py b/sqlspec/adapters/mssql_python/driver.py index 737fe12e3..6674d99ae 100644 --- a/sqlspec/adapters/mssql_python/driver.py +++ b/sqlspec/adapters/mssql_python/driver.py @@ -334,6 +334,7 @@ def _check_schema() -> bool: _execute_cursor(cursor, "SELECT 1 FROM sys.schemas WHERE name = ?", (schema,)) return cursor.fetchone() is not None + result = False exc_handler = self.handle_database_exceptions() if sys.exc_info()[1] is None: with exc_handler: @@ -341,7 +342,6 @@ def _check_schema() -> bool: else: result = bool(self._run_with_exception_handler(exc_handler, _check_schema)) self._check_pending_exception(exc_handler) - return bool(result) return result def select_to_arrow( @@ -411,39 +411,40 @@ def _open_reader() -> None: arrow_schema=arrow_schema, ) - def _fetch_table_arrow() -> Any: - nonlocal reader + def _fetch_arrow() -> "ArrowResult": with self.with_cursor(self.connection) as cursor: _execute_cursor(cursor, sql, prepared_parameters) if return_format == "batches": reader = _cursor_arrow_reader(cursor, arrow_kwargs) if reader is not None: - return None - return cursor.arrow(**arrow_kwargs) + return build_arrow_result_from_reader( + prepared_statement, + arrow_reader_with_deferred_close(reader, cursor.close), + return_format=return_format, + batch_size=batch_size, + arrow_schema=arrow_schema, + ) + table = cursor.arrow(**arrow_kwargs) + if table is None: + msg = "mssql-python did not return an Arrow table." + raise SQLSpecError(msg) + return build_arrow_result_from_table( + prepared_statement, + table, + return_format=return_format, + batch_size=batch_size, + arrow_schema=arrow_schema, + ) + result: ArrowResult | None = None exc_handler = self.handle_database_exceptions() if sys.exc_info()[1] is None: with exc_handler: - table = _fetch_table_arrow() + result = _fetch_arrow() else: - table = self._run_with_exception_handler(exc_handler, _fetch_table_arrow) + result = self._run_with_exception_handler(exc_handler, _fetch_arrow) self._check_pending_exception(exc_handler) - - if return_format == "batches" and reader is not None: - return build_arrow_result_from_reader( - prepared_statement, - reader, - return_format=return_format, - batch_size=batch_size, - arrow_schema=arrow_schema, - ) - - if table is None: - msg = "mssql-python did not return an Arrow table." - raise SQLSpecError(msg) - return build_arrow_result_from_table( - prepared_statement, table, return_format=return_format, batch_size=batch_size, arrow_schema=arrow_schema - ) + return cast("ArrowResult", result) def bulk_copy( self, @@ -479,6 +480,7 @@ def _do_bulkcopy() -> Any: ) return _coerce_bulk_copy_result(raw_result, cursor) + result: MssqlPythonBulkCopyResult | None = None exc_handler = self.handle_database_exceptions() if sys.exc_info()[1] is None: with exc_handler: diff --git a/sqlspec/adapters/mysqlconnector/_typing.py b/sqlspec/adapters/mysqlconnector/_typing.py index 12b9fc274..f91e7a715 100644 --- a/sqlspec/adapters/mysqlconnector/_typing.py +++ b/sqlspec/adapters/mysqlconnector/_typing.py @@ -17,13 +17,14 @@ from mysql.connector import pooling as _mysql_connector_pooling from mysql.connector.aio import MySQLConnection as _MysqlConnectorAsyncConnection from mysql.connector.aio import cursor as _mysql_connector_aio_cursor -from mysql.connector.aio.pooling import MySQLConnectionPool as MysqlConnectorAsyncPool if TYPE_CHECKING: from collections.abc import Awaitable, Callable from types import TracebackType from typing import ClassVar, Protocol, TypeAlias + from mysql.connector.aio.pooling import MySQLConnectionPool as MysqlConnectorAsyncPool + from sqlspec.adapters.mysqlconnector.driver import MysqlConnectorAsyncDriver, MysqlConnectorSyncDriver from sqlspec.core import StatementConfig @@ -83,6 +84,12 @@ class MysqlConnectorMysqlModuleProtocol(Protocol): MysqlConnectorMysqlModule: TypeAlias = MysqlConnectorMysqlModuleProtocol MysqlConnectorAsyncRawCursor: TypeAlias = MysqlConnectorAsyncRawCursorProtocol else: + try: + from mysql.connector.aio.pooling import MySQLConnectionPool as MysqlConnectorAsyncPool + except ModuleNotFoundError as exc: + if exc.name != "mysql.connector.aio.pooling": + raise + MysqlConnectorAsyncPool = None MysqlConnectorAio = _mysql_connector_aio MysqlConnectorSyncConnection = _MysqlConnectorSyncConnection MysqlConnectorAsyncConnection = _MysqlConnectorAsyncConnection diff --git a/sqlspec/adapters/oracledb/adk/store.py b/sqlspec/adapters/oracledb/adk/store.py index b2f10b1be..ce41281dd 100644 --- a/sqlspec/adapters/oracledb/adk/store.py +++ b/sqlspec/adapters/oracledb/adk/store.py @@ -19,6 +19,8 @@ from sqlspec.adapters.oracledb.data_dictionary import JSONStorageType, OracleVersionInfo, storage_type_from_version from sqlspec.builder._vector_distance import build_oracle_hnsw_subquery, build_oracle_hybrid_rrf_sql from sqlspec.config import ADKConfig +from sqlspec.data_dictionary.dialects.oracle import parse_oracle_version_components +from sqlspec.exceptions import ImproperConfigurationError from sqlspec.extensions.adk import ( BaseAsyncADKStore, BaseSyncADKStore, @@ -250,7 +252,7 @@ " event_id VARCHAR2(128) NOT NULL UNIQUE,\n" " author VARCHAR2(256){1},\n" " timestamp TIMESTAMP WITH TIME ZONE NOT NULL,\n" - " embedding VECTOR({11}, FLOAT32),\n" + " {11}" " {2},\n" " content_text CLOB NOT NULL,\n" " inserted_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL\n" @@ -390,6 +392,26 @@ async def _rollback_async(conn: Any) -> None: await conn.rollback() +def _memory_supports_vectors(version: "OracleVersionInfo | None") -> bool: + if version is None or (version.major, version.minor, version.patch) < (23, 4, 0): + return False + compatible = parse_oracle_version_components(version.compatible) if version.compatible else None + return compatible is None or compatible >= (23, 4, 0) + + +def _require_memory_vectors(version: "OracleVersionInfo | None") -> None: + if not _memory_supports_vectors(version): + msg = "Oracle memory embeddings require Oracle Database and COMPATIBLE version 23.4 or later." + raise ImproperConfigurationError(msg) + + +def _memory_insert_count(entry_count: int, batch_errors: "list[Any]") -> int: + for error in batch_errors: + if error.code != ORACLE_DUPLICATE_KEY_ERROR: + raise OracleDatabaseError(error) + return entry_count - len(batch_errors) + + def coerce_decimal_values(value: Any) -> Any: return _coerce_decimal_values(value) @@ -2472,15 +2494,21 @@ async def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: " if not entries: return 0 + version = await self._get_version_info() + supports_vectors = _memory_supports_vectors(version) + if any(entry.get("embedding") is not None for entry in entries): + _require_memory_vectors(version) + embedding_column = ", embedding" if supports_vectors else "" + embedding_param = ", :embedding" if supports_vectors else "" owner_column = f", {self._owner_id_column_name}" if self._owner_id_column_name else "" owner_param = ", :owner_id" if self._owner_id_column_name else "" sql = f""" INSERT INTO {self._memory_table} ( id, session_id, app_name, user_id, scope, event_id, author{owner_column}, - timestamp, embedding, content_json, content_text, metadata_json, inserted_at + timestamp{embedding_column}, content_json, content_text, metadata_json, inserted_at ) VALUES ( :id, :session_id, :app_name, :user_id, :scope, :event_id, :author{owner_param}, - :timestamp, :embedding, :content_json, :content_text, :metadata_json, :inserted_at + :timestamp{embedding_param}, :content_json, :content_text, :metadata_json, :inserted_at ) """ @@ -2499,12 +2527,13 @@ async def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: " "event_id": entry["event_id"], "author": entry["author"], "timestamp": entry["timestamp"], - "embedding": bound_embedding, "content_json": content_json, "content_text": entry["content_text"], "metadata_json": metadata_json, "inserted_at": entry["inserted_at"], } + if supports_vectors: + p["embedding"] = bound_embedding if self._owner_id_column_name: p["owner_id"] = owner_id batch_params.append(p) @@ -2514,13 +2543,7 @@ async def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: " coerced_batch = [await coerce_large_parameters_async(conn, p) for p in batch_params] with conn.cursor() as cursor: await cursor.executemany(sql, coerced_batch, batcherrors=True) - batch_errors = list(cursor.getbatcherrors()) if hasattr(cursor, "getbatcherrors") else [] - duplicates = sum( - 1 - for e in batch_errors - if getattr(e, "code", getattr(e, "full_code", None)) == ORACLE_DUPLICATE_KEY_ERROR - ) - inserted_count = len(entries) - duplicates + inserted_count = _memory_insert_count(len(entries), cursor.getbatcherrors()) await conn.commit() return inserted_count except BaseException: @@ -2541,6 +2564,8 @@ async def search_entries( raise RuntimeError(msg) effective_limit = limit if limit is not None else self._max_results + if embedding is not None: + _require_memory_vectors(await self._get_version_info()) try: if embedding is not None and self._use_fts: @@ -2628,10 +2653,12 @@ async def _deserialize_json_field(self, data: Any) -> "dict[str, Any] | None": return _extract_json_value(data) async def _memory_table_ddl(self) -> str: - storage_type = await self._detect_json_storage_type() - return self._memory_table_ddl_for_type(storage_type) + version = await self._get_version_info() + return self._memory_table_ddl_for_type( + storage_type_from_version(version), supports_vectors=_memory_supports_vectors(version) + ) - def _memory_table_ddl_for_type(self, storage_type: "JSONStorageType") -> str: + def _memory_table_ddl_for_type(self, storage_type: "JSONStorageType", *, supports_vectors: bool = False) -> str: if storage_type == JSONStorageType.JSON_NATIVE: json_columns = """ content_json JSON, @@ -2676,7 +2703,7 @@ def _memory_table_ddl_for_type(self, storage_type: "JSONStorageType") -> str: self._memory_table, self._memory_table, fts_index, - self._vector_dimensions, + f"embedding VECTOR({self._vector_dimensions}, FLOAT32),\n " if supports_vectors else "", ) def _drop_memory_table_sql(self) -> "list[str]": @@ -2764,12 +2791,15 @@ async def _search_entries_fts( if not sanitized_query: return [] where_scope, scope_params = _build_oracle_scope_where(app_name, user_id, scope_filter) + embedding_column = ( + "embedding" if _memory_supports_vectors(await self._get_version_info()) else "NULL AS embedding" + ) sql = f""" SELECT id, session_id, app_name, user_id, scope, event_id, author, timestamp, content_json, content_text, metadata_json, inserted_at, embedding FROM ( SELECT id, session_id, app_name, user_id, scope, event_id, author, - timestamp, content_json, content_text, metadata_json, inserted_at, embedding, + timestamp, content_json, content_text, metadata_json, inserted_at, {embedding_column}, SCORE(1) AS score FROM {self._memory_table} WHERE {where_scope} @@ -2861,12 +2891,15 @@ async def _search_entries_simple( self, query: str, app_name: str, user_id: str, limit: int, scope_filter: Literal["all", "user", "app"] = "all" ) -> "list[StoredMemory]": where_scope, scope_params = _build_oracle_scope_where(app_name, user_id, scope_filter) + embedding_column = ( + "embedding" if _memory_supports_vectors(await self._get_version_info()) else "NULL AS embedding" + ) sql = f""" SELECT id, session_id, app_name, user_id, scope, event_id, author, timestamp, content_json, content_text, metadata_json, inserted_at, embedding FROM ( SELECT id, session_id, app_name, user_id, scope, event_id, author, - timestamp, content_json, content_text, metadata_json, inserted_at, embedding + timestamp, content_json, content_text, metadata_json, inserted_at, {embedding_column} FROM {self._memory_table} WHERE {where_scope} AND LOWER(content_text) LIKE :pattern @@ -2964,15 +2997,21 @@ def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "object if not entries: return 0 + version = self._get_version_info() + supports_vectors = _memory_supports_vectors(version) + if any(entry.get("embedding") is not None for entry in entries): + _require_memory_vectors(version) + embedding_column = ", embedding" if supports_vectors else "" + embedding_param = ", :embedding" if supports_vectors else "" owner_column = f", {self._owner_id_column_name}" if self._owner_id_column_name else "" owner_param = ", :owner_id" if self._owner_id_column_name else "" sql = f""" INSERT INTO {self._memory_table} ( id, session_id, app_name, user_id, scope, event_id, author{owner_column}, - timestamp, embedding, content_json, content_text, metadata_json, inserted_at + timestamp{embedding_column}, content_json, content_text, metadata_json, inserted_at ) VALUES ( :id, :session_id, :app_name, :user_id, :scope, :event_id, :author{owner_param}, - :timestamp, :embedding, :content_json, :content_text, :metadata_json, :inserted_at + :timestamp{embedding_param}, :content_json, :content_text, :metadata_json, :inserted_at ) """ @@ -2991,12 +3030,13 @@ def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "object "event_id": entry["event_id"], "author": entry["author"], "timestamp": entry["timestamp"], - "embedding": bound_embedding, "content_json": content_json, "content_text": entry["content_text"], "metadata_json": metadata_json, "inserted_at": entry["inserted_at"], } + if supports_vectors: + p["embedding"] = bound_embedding if self._owner_id_column_name: p["owner_id"] = owner_id batch_params.append(p) @@ -3006,13 +3046,7 @@ def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "object coerced_batch = [coerce_large_parameters_sync(conn, p) for p in batch_params] with conn.cursor() as cursor: cursor.executemany(sql, coerced_batch, batcherrors=True) - batch_errors = list(cursor.getbatcherrors()) if hasattr(cursor, "getbatcherrors") else [] - duplicates = sum( - 1 - for e in batch_errors - if getattr(e, "code", getattr(e, "full_code", None)) == ORACLE_DUPLICATE_KEY_ERROR - ) - inserted_count = len(entries) - duplicates + inserted_count = _memory_insert_count(len(entries), cursor.getbatcherrors()) conn.commit() return inserted_count except BaseException: @@ -3034,6 +3068,8 @@ def search_entries( raise RuntimeError(msg) effective_limit = limit if limit is not None else self._max_results + if embedding is not None: + _require_memory_vectors(self._get_version_info()) try: if embedding is not None and self._use_fts: @@ -3119,10 +3155,12 @@ def _deserialize_json_field(self, data: Any) -> "dict[str, Any] | None": return _extract_json_value(data) def _memory_table_ddl(self) -> str: - storage_type = self._detect_json_storage_type() - return self._memory_table_ddl_for_type(storage_type) + version = self._get_version_info() + return self._memory_table_ddl_for_type( + storage_type_from_version(version), supports_vectors=_memory_supports_vectors(version) + ) - def _memory_table_ddl_for_type(self, storage_type: "JSONStorageType") -> str: + def _memory_table_ddl_for_type(self, storage_type: "JSONStorageType", *, supports_vectors: bool = False) -> str: if storage_type == JSONStorageType.JSON_NATIVE: json_columns = """ content_json JSON, @@ -3167,7 +3205,7 @@ def _memory_table_ddl_for_type(self, storage_type: "JSONStorageType") -> str: self._memory_table, self._memory_table, fts_index, - self._vector_dimensions, + f"embedding VECTOR({self._vector_dimensions}, FLOAT32),\n " if supports_vectors else "", ) def _drop_memory_table_sql(self) -> "list[str]": @@ -3255,12 +3293,13 @@ def _search_entries_fts( if not sanitized_query: return [] where_scope, scope_params = _build_oracle_scope_where(app_name, user_id, scope_filter) + embedding_column = "embedding" if _memory_supports_vectors(self._get_version_info()) else "NULL AS embedding" sql = f""" SELECT id, session_id, app_name, user_id, scope, event_id, author, timestamp, content_json, content_text, metadata_json, inserted_at, embedding FROM ( SELECT id, session_id, app_name, user_id, scope, event_id, author, - timestamp, content_json, content_text, metadata_json, inserted_at, embedding, + timestamp, content_json, content_text, metadata_json, inserted_at, {embedding_column}, SCORE(1) AS score FROM {self._memory_table} WHERE {where_scope} @@ -3350,12 +3389,13 @@ def _search_entries_simple( self, query: str, app_name: str, user_id: str, limit: int, scope_filter: Literal["all", "user", "app"] = "all" ) -> "list[StoredMemory]": where_scope, scope_params = _build_oracle_scope_where(app_name, user_id, scope_filter) + embedding_column = "embedding" if _memory_supports_vectors(self._get_version_info()) else "NULL AS embedding" sql = f""" SELECT id, session_id, app_name, user_id, scope, event_id, author, timestamp, content_json, content_text, metadata_json, inserted_at, embedding FROM ( SELECT id, session_id, app_name, user_id, scope, event_id, author, - timestamp, content_json, content_text, metadata_json, inserted_at, embedding + timestamp, content_json, content_text, metadata_json, inserted_at, {embedding_column} FROM {self._memory_table} WHERE {where_scope} AND LOWER(content_text) LIKE :pattern diff --git a/sqlspec/adapters/oracledb/config.py b/sqlspec/adapters/oracledb/config.py index 1a7c99123..5962f7fce 100644 --- a/sqlspec/adapters/oracledb/config.py +++ b/sqlspec/adapters/oracledb/config.py @@ -337,10 +337,49 @@ def _build_standalone_connection_params(connection_config: "Mapping[str, Any]") config.pop("max_lifetime_session", None) config.pop("session_callback", None) config.pop("threaded", None) + config.pop("getmode", None) + config.pop("ping_interval", None) + config.pop("ping_timeout", None) + config.pop("max_sessions_per_shard", None) + config.pop("pool_class", None) + connection_type = config.pop("connectiontype", None) + if connection_type is not None: + config["conn_class"] = connection_type apply_wallet_ssl_context(config) return config +def _initialize_oracle_client(connection_config: "Mapping[str, Any]") -> None: + """Initialize requested Thick mode before creating a pool or connection.""" + lib_dir = connection_config.get("lib_dir") + if ( + not (connection_config.get("thick_mode") or lib_dir is not None or connection_config.get("soda_metadata_cache")) + or not oracledb.is_thin_mode() + ): + return + client_config = {} + if lib_dir is not None: + client_config["lib_dir"] = lib_dir + if connection_config.get("config_dir") is not None: + client_config["config_dir"] = connection_config["config_dir"] + try: + oracledb.init_oracle_client(**client_config) + except oracledb.Error as exc: + msg = f"Failed to initialize Oracle client in thick mode: {exc}" + raise ImproperConfigurationError(msg) from exc + + +def _validate_async_oracle_mode(connection_config: "Mapping[str, Any]") -> None: + """Reject Thick mode settings unsupported by async connections.""" + if ( + connection_config.get("thick_mode") + or connection_config.get("lib_dir") is not None + or connection_config.get("soda_metadata_cache") + ): + msg = "OracleAsyncConfig only supports Thin mode; Thick mode options (thick_mode, lib_dir, soda_metadata_cache) require OracleSyncConfig." + raise ImproperConfigurationError(msg) + + def _derive_oracle_pool_kwargs( source_config: Any, connection_config_overrides: "Mapping[str, Any] | None", @@ -437,7 +476,7 @@ def __exit__( class _OracleSyncSessionConnectionHandler(SyncPoolSessionFactory): - __slots__ = ("_ctx",) + __slots__ = () def __init__(self, config: "OracleSyncConfig") -> None: super().__init__(config) @@ -531,6 +570,7 @@ def create_connection(self) -> "OracleSyncConnection": Returns: An Oracle Connection instance. """ + _initialize_oracle_client(self.connection_config) config = _build_standalone_connection_params(self.connection_config) conn = oracledb.connect(**config) self._init_connection(conn, "") @@ -620,20 +660,10 @@ def get_event_runtime_hints(self) -> "EventRuntimeHints": def _create_pool(self) -> "OracleSyncConnectionPool": """Create the actual connection pool.""" + _initialize_oracle_client(self.connection_config) config = dict(self.connection_config) - thick_mode = config.pop("thick_mode", False) - lib_dir = config.pop("lib_dir", None) - if (thick_mode or lib_dir is not None or config.get("soda_metadata_cache")) and oracledb.is_thin_mode(): - client_config = {} - if lib_dir is not None: - client_config["lib_dir"] = lib_dir - if config.get("config_dir") is not None: - client_config["config_dir"] = config["config_dir"] - try: - oracledb.init_oracle_client(**client_config) - except oracledb.Error as exc: - msg = f"Failed to initialize Oracle client in thick mode: {exc}" - raise ImproperConfigurationError(msg) from exc + config.pop("thick_mode", None) + config.pop("lib_dir", None) apply_wallet_ssl_context(config) @@ -837,6 +867,7 @@ async def create_connection(self) -> OracleAsyncConnection: Returns: An Oracle AsyncConnection instance. """ + _validate_async_oracle_mode(self.connection_config) config = _build_standalone_connection_params(self.connection_config) conn = await oracledb.connect_async(**config) await self._init_connection(conn, "") @@ -923,13 +954,10 @@ def get_event_runtime_hints(self) -> "EventRuntimeHints": async def _create_pool(self) -> "OracleAsyncConnectionPool": """Create the actual async connection pool.""" + _validate_async_oracle_mode(self.connection_config) config = dict(self.connection_config) - thick_mode = config.pop("thick_mode", False) - lib_dir = config.pop("lib_dir", None) - soda_metadata_cache = config.get("soda_metadata_cache", False) - if thick_mode or lib_dir is not None or soda_metadata_cache: - msg = "OracleAsyncConfig only supports Thin mode; Thick mode options (thick_mode, lib_dir, soda_metadata_cache) require OracleSyncConfig." - raise ImproperConfigurationError(msg) + config.pop("thick_mode", None) + config.pop("lib_dir", None) apply_wallet_ssl_context(config) diff --git a/sqlspec/adapters/oracledb/core.py b/sqlspec/adapters/oracledb/core.py index c3df5650d..6f8c245be 100644 --- a/sqlspec/adapters/oracledb/core.py +++ b/sqlspec/adapters/oracledb/core.py @@ -1779,6 +1779,8 @@ def create_mapped_exception(error: Any, *, logger: Any | None = None) -> SQLSpec A SQLSpec exception that wraps the original error """ del logger + if not isinstance(error, BaseException): + error = oracledb_module.DatabaseError(error) error_obj = error.args[0] if getattr(error, "args", None) else None if not error_obj: return _create_oracle_error(error, None, SQLSpecError, "database error") diff --git a/sqlspec/adapters/oracledb/driver.py b/sqlspec/adapters/oracledb/driver.py index a64c9860c..0954c3a90 100644 --- a/sqlspec/adapters/oracledb/driver.py +++ b/sqlspec/adapters/oracledb/driver.py @@ -841,6 +841,7 @@ def select_to_arrow( arrow_schema, ), ) + self._check_pending_exception(exc_handler) record_batches, _ = batches_result rows_affected = sum(batch.num_rows for batch in record_batches) return create_arrow_result(prepared_statement, record_batches, rows_affected=rows_affected) @@ -1244,14 +1245,19 @@ def _iter_arrow_record_batches( oracle_dfs = self.connection.fetch_df_batches( statement=sql, parameters=params, size=batch_size or 1000, **retry_kwargs ) - for oracle_df in oracle_dfs: - batch_table = coerce_oracle_df_to_arrow_table(oracle_df, import_pyarrow()) - column_names = normalize_column_names(batch_table.column_names, self.driver_features) - if column_names != batch_table.column_names: - batch_table = batch_table.rename_columns(column_names) - if arrow_schema is not None: - batch_table = batch_table.cast(arrow_schema) - yield from batch_table.to_batches() + except OracleError as exc: + raise create_mapped_exception(exc) from exc + try: + for oracle_df in oracle_dfs: + batch_table = coerce_oracle_df_to_arrow_table(oracle_df, import_pyarrow()) + column_names = normalize_column_names(batch_table.column_names, self.driver_features) + if column_names != batch_table.column_names: + batch_table = batch_table.rename_columns(column_names) + if arrow_schema is not None: + batch_table = batch_table.cast(arrow_schema) + yield from batch_table.to_batches() + except OracleError as exc: + raise create_mapped_exception(exc) from exc def _fetch_arrow_record_batches( self, @@ -1750,6 +1756,7 @@ async def select_to_arrow( arrow_schema, ), ) + self._check_pending_exception(exc_handler) record_batches, batch_schema = batches_result rows_affected = sum(batch.num_rows for batch in record_batches) pa = import_pyarrow() diff --git a/sqlspec/adapters/oracledb/litestar/store.py b/sqlspec/adapters/oracledb/litestar/store.py index d81da305c..ecd14e4da 100644 --- a/sqlspec/adapters/oracledb/litestar/store.py +++ b/sqlspec/adapters/oracledb/litestar/store.py @@ -167,7 +167,7 @@ async def get(self, key: str, renew_for: "int | timedelta | None" = None) -> "by WHERE session_id = :2 """ try: - await cursor.execute(update_sql, (expires_in_seconds, key)) + await cursor.execute(update_sql, (expires_in_seconds, expires_in_seconds, key)) await conn.commit() except BaseException: with contextlib.suppress(Exception): @@ -511,7 +511,7 @@ def _get(self, key: str, renew_for: "int | timedelta | None" = None) -> "bytes | WHERE session_id = :2 """ try: - cursor.execute(update_sql, (expires_in_seconds, key)) + cursor.execute(update_sql, (expires_in_seconds, expires_in_seconds, key)) conn.commit() except BaseException: with contextlib.suppress(Exception): diff --git a/sqlspec/adapters/psqlpy/adk/store.py b/sqlspec/adapters/psqlpy/adk/store.py index f8e43f265..cedcc9f35 100644 --- a/sqlspec/adapters/psqlpy/adk/store.py +++ b/sqlspec/adapters/psqlpy/adk/store.py @@ -10,6 +10,7 @@ from sqlspec.extensions.adk import BaseAsyncADKStore, StoredEvent, StoredSession, normalize_session_list_options from sqlspec.extensions.adk.memory.store import BaseAsyncADKMemoryStore from sqlspec.utils.logging import get_logger +from sqlspec.utils.serializers import to_json from sqlspec.utils.type_guards import has_query_result_metadata if TYPE_CHECKING: @@ -705,19 +706,28 @@ async def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: " content_texts = [entry["content_text"] for entry in entries] metadata_jsons = [entry.get("metadata_json") for entry in entries] inserted_ats = [entry["inserted_at"] for entry in entries] + embeddings = [ + to_json([float(value) for value in embedding]) + if (embedding := entry.get("embedding")) is not None + else None + for entry in entries + ] if self._owner_id_column_name: sql = f""" INSERT INTO {self._memory_table} ( id, session_id, app_name, user_id, scope, event_id, author, {self._owner_id_column_name}, timestamp, content_json, - content_text, metadata_json, inserted_at + content_text, metadata_json, inserted_at, embedding ) - SELECT * FROM unnest( + SELECT id, session_id, app_name, user_id, scope, event_id, author, + owner_id, timestamp, content_json, content_text, metadata_json, inserted_at, embedding::vector + FROM unnest( $1::varchar[], $2::varchar[], $3::varchar[], $4::varchar[], $5::varchar[], $6::varchar[], $7::varchar[], $8::varchar[], - $9::timestamptz[], $10::jsonb[], $11::text[], $12::jsonb[], $13::timestamptz[] - ) + $9::timestamptz[], $10::jsonb[], $11::text[], $12::jsonb[], $13::timestamptz[], $14::text[] + ) AS batch(id, session_id, app_name, user_id, scope, event_id, author, + owner_id, timestamp, content_json, content_text, metadata_json, inserted_at, embedding) ON CONFLICT (event_id) DO NOTHING RETURNING 1 """ @@ -736,18 +746,22 @@ async def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: " content_texts, metadata_jsons, inserted_ats, + embeddings, ] else: sql = f""" INSERT INTO {self._memory_table} ( id, session_id, app_name, user_id, scope, event_id, author, - timestamp, content_json, content_text, metadata_json, inserted_at + timestamp, content_json, content_text, metadata_json, inserted_at, embedding ) - SELECT * FROM unnest( + SELECT id, session_id, app_name, user_id, scope, event_id, author, + timestamp, content_json, content_text, metadata_json, inserted_at, embedding::vector + FROM unnest( $1::varchar[], $2::varchar[], $3::varchar[], $4::varchar[], $5::varchar[], $6::varchar[], $7::varchar[], - $8::timestamptz[], $9::jsonb[], $10::text[], $11::jsonb[], $12::timestamptz[] - ) + $8::timestamptz[], $9::jsonb[], $10::text[], $11::jsonb[], $12::timestamptz[], $13::text[] + ) AS batch(id, session_id, app_name, user_id, scope, event_id, author, + timestamp, content_json, content_text, metadata_json, inserted_at, embedding) ON CONFLICT (event_id) DO NOTHING RETURNING 1 """ @@ -764,6 +778,7 @@ async def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: " content_texts, metadata_jsons, inserted_ats, + embeddings, ] async with self._config.provide_connection() as conn: diff --git a/sqlspec/adapters/psqlpy/driver.py b/sqlspec/adapters/psqlpy/driver.py index 7aae5c332..08f561ac4 100644 --- a/sqlspec/adapters/psqlpy/driver.py +++ b/sqlspec/adapters/psqlpy/driver.py @@ -46,7 +46,7 @@ from sqlspec.adapters.psqlpy._typing import PsqlpyConnection from sqlspec.core import ArrowResult, SQLResult - from sqlspec.driver import ExecutionResult + from sqlspec.driver import CachedQuery, ExecutionResult from sqlspec.storage import StorageBridgeJob, StorageDestination, StorageFormat, StorageTelemetry __all__ = ("PsqlpyCursor", "PsqlpyDriver", "PsqlpyExceptionHandler", "PsqlpySessionContext") @@ -142,23 +142,8 @@ async def dispatch_execute(self, cursor: "PsqlpyConnection", statement: SQL) -> ) if statement.operation_type in {"INSERT", "UPDATE", "DELETE"}: - base_sql, comment_suffix = split_sqlcomment(sql) - count_sql = _dml_count_query(base_sql) - if count_sql is not None: - wire_count_sql = f"{count_sql} {comment_suffix}" if comment_suffix else count_sql - count_result = ( - await cursor.fetch(wire_count_sql, params, prepared=False) - if self._is_transaction_pooler_mode() - else await cursor.fetch(wire_count_sql, params) - ) - count_rows, _ = collect_rows(count_result) - if len(count_rows) != 1 or set(count_rows[0]) != {_DML_COUNT_COLUMN}: - msg = "psqlpy DML row count query returned an invalid result" - raise SQLSpecError(msg) - rows_affected = count_rows[0][_DML_COUNT_COLUMN] - if type(rows_affected) is not int or rows_affected < 0: - msg = "psqlpy DML row count query returned an invalid count" - raise SQLSpecError(msg) + rows_affected = await self._fetch_dml_rowcount(cursor, sql, params) + if rows_affected is not None: return self.create_execution_result(cursor, rowcount_override=rows_affected) result = ( @@ -170,6 +155,29 @@ async def dispatch_execute(self, cursor: "PsqlpyConnection", statement: SQL) -> return self.create_execution_result(cursor, rowcount_override=rows_affected) + async def _fetch_dml_rowcount( + self, cursor: "PsqlpyConnection", sql: str, params: "Sequence[Any] | Mapping[str, Any]" + ) -> int | None: + base_sql, comment_suffix = split_sqlcomment(sql) + count_sql = _dml_count_query(base_sql) + if count_sql is None: + return None + wire_count_sql = f"{count_sql} {comment_suffix}" if comment_suffix else count_sql + count_result = ( + await cursor.fetch(wire_count_sql, params, prepared=False) + if self._pooler_unprepared + else await cursor.fetch(wire_count_sql, params) + ) + count_rows, _ = collect_rows(count_result) + if len(count_rows) != 1 or set(count_rows[0]) != {_DML_COUNT_COLUMN}: + msg = "psqlpy DML row count query returned an invalid result" + raise SQLSpecError(msg) + rows_affected = count_rows[0][_DML_COUNT_COLUMN] + if type(rows_affected) is not int or rows_affected < 0: + msg = "psqlpy DML row count query returned an invalid count" + raise SQLSpecError(msg) + return rows_affected + async def dispatch_execute_many(self, cursor: "PsqlpyConnection", statement: SQL) -> "ExecutionResult": """Execute SQL with multiple parameter sets. @@ -553,8 +561,8 @@ async def _resolve_json_columns(self, schema_name: "str | None", table_name: str data, _ = collect_rows(rows) return {str(row["column_name"]) for row in data} - async def _execute_cache_hit( - self, sql: str, params: "tuple[Any, ...] | list[Any] | dict[str, Any]", cached: Any + async def _execute_cache_hit_with_cursor( + self, sql: str, params: "tuple[Any, ...] | list[Any] | dict[str, Any]", cached: "CachedQuery" ) -> "SQLResult": """Execute cached psqlpy queries with cast-aware parameter preparation.""" prepared_params = self.prepare_driver_parameters(params, self.statement_config, prepared_statement=cached) @@ -583,6 +591,10 @@ async def _execute_cache_hit( row_format="dict", compiled_sql=wire_sql, ) + if cached.operation_type in {"INSERT", "UPDATE", "DELETE"}: + rows_affected = await self._fetch_dml_rowcount(self.connection, wire_sql, args) + if rows_affected is not None: + return DMLResult(cached.operation_type, rows_affected) result = ( await self.connection.execute(wire_sql, args, prepared=False) if self._pooler_unprepared diff --git a/sqlspec/adapters/psycopg/config.py b/sqlspec/adapters/psycopg/config.py index dbba98b7b..44c324f1b 100644 --- a/sqlspec/adapters/psycopg/config.py +++ b/sqlspec/adapters/psycopg/config.py @@ -402,6 +402,8 @@ def derive_direct_config(self) -> Self: features = dict(self.driver_features) features["pooler_mode"] = "session" features["direct_dsn"] = None + if self._user_connection_hook is not None: + features["on_connection_create"] = self._user_connection_hook return type(self)( connection_config=conn_cfg, migration_config=dict(self.migration_config) if self.migration_config else None, @@ -659,7 +661,15 @@ def _configure_connection(self, conn: "PsycopgSyncConnection") -> None: register_pgvector_sync(conn) if self._user_connection_hook is not None: - self._user_connection_hook(conn) + try: + self._user_connection_hook(conn) + if not conn.autocommit: + conn.commit() + except BaseException: + if not conn.autocommit: + with suppress(Exception): + conn.rollback() + raise def _ensure_pg_textsearch_available(self) -> None: if self._pg_textsearch_available: @@ -855,6 +865,8 @@ def derive_direct_config(self) -> Self: features = dict(self.driver_features) features["pooler_mode"] = "session" features["direct_dsn"] = None + if self._user_connection_hook is not None: + features["on_connection_create"] = self._user_connection_hook return type(self)( connection_config=conn_cfg, migration_config=dict(self.migration_config) if self.migration_config else None, @@ -1065,9 +1077,17 @@ async def _configure_async_connection(self, conn: "PsycopgAsyncConnection") -> N await register_pgvector_async(conn) if self._user_connection_hook is not None: - hook_res = self._user_connection_hook(conn) - if inspect.isawaitable(hook_res): - await hook_res + try: + hook_res = self._user_connection_hook(conn) + if inspect.isawaitable(hook_res): + await hook_res + if not conn.autocommit: + await conn.commit() + except BaseException: + if not conn.autocommit: + with suppress(Exception): + await conn.rollback() + raise def _ensure_pg_textsearch_available(self) -> None: if self._pg_textsearch_available: diff --git a/sqlspec/adapters/psycopg/driver.py b/sqlspec/adapters/psycopg/driver.py index fadd02cca..d94369abd 100644 --- a/sqlspec/adapters/psycopg/driver.py +++ b/sqlspec/adapters/psycopg/driver.py @@ -17,6 +17,7 @@ PsycopgSyncSessionContext, ) from sqlspec.adapters.psycopg._typing import psycopg_module as psycopg +from sqlspec.adapters.psycopg._typing import psycopg_sql as pg_sql from sqlspec.adapters.psycopg.core import ( OPEN_TRANSACTION_STATUSES, TRANSACTION_STATUS_INERROR, @@ -51,8 +52,6 @@ is_copy_to_operation, register_driver_profile, ) -from sqlspec.core.parameters import ParameterValidator -from sqlspec.core.sqlcommenter import format_wire_sqlcomment from sqlspec.driver import ( AsyncDriverAdapterBase, AsyncRowStream, @@ -65,7 +64,7 @@ ) from sqlspec.exceptions import SQLSpecError, StackExecutionError from sqlspec.utils.logging import get_logger -from sqlspec.utils.text import normalize_identifier, quote_identifier +from sqlspec.utils.text import normalize_identifier from sqlspec.utils.type_guards import is_readable, resolve_row_format if TYPE_CHECKING: @@ -75,6 +74,7 @@ from sqlspec.core import ArrowResult from sqlspec.driver import CachedQuery, ExecutionResult from sqlspec.storage import StorageBridgeJob, StorageDestination, StorageFormat, StorageTelemetry + from sqlspec.typing import ArrowRecordBatch __all__ = ( @@ -455,16 +455,14 @@ def rollback(self) -> None: def set_migration_session_schema(self, schema: str) -> None: """Set the PostgreSQL search path for migration SQL.""" normalized_schema = normalize_identifier(schema, "postgres") - quoted_schema = quote_identifier(normalized_schema) - sql = f'SET LOCAL search_path TO {quoted_schema}, "$user", public' + sql = PsycopgSQL('SET LOCAL search_path TO {}, "$user", public').format(pg_sql.Identifier(normalized_schema)) with self.with_cursor(self.connection) as cursor: cursor.execute(sql) def set_migration_non_transactional_schema(self, schema: str) -> None: """Set the PostgreSQL search path for non-transactional migration SQL.""" normalized_schema = normalize_identifier(schema, "postgres") - quoted_schema = quote_identifier(normalized_schema) - sql = f'SET search_path TO {quoted_schema}, "$user", public' + sql = PsycopgSQL('SET search_path TO {}, "$user", public').format(pg_sql.Identifier(normalized_schema)) with self.with_cursor(self.connection) as cursor: cursor.execute(sql) @@ -589,30 +587,39 @@ def load_from_arrow( copy_sql = build_copy_from_command(table, columns, binary=use_binary) all_batches = chain((first_batch,), (b for b in batch_iter if b.num_rows > 0)) exc_handler = self.handle_database_exceptions() - with ExitStack() as stack: - cursor = stack.enter_context(self.with_cursor(self.connection)) - copy_ctx = stack.enter_context(cursor.copy(copy_sql)) - if use_binary: - for chunk in self._iter_pgcopy_binary_chunks(all_batches, columns=columns, metrics=metrics): - self._run_with_exception_handler(exc_handler, copy_ctx.write, chunk) - else: - for batch in all_batches: - metrics["rows_processed"] += int(batch.num_rows) - metrics["bytes_processed"] += int(batch.nbytes) - records = self._arrow_batch_to_rows(batch, columns) - if self._arrow_rows_need_preparation(batch): - records = cast( - "list[Any]", - self.prepare_driver_parameters(records, self.statement_config, is_many=True), - ) - for record in records: - self._run_with_exception_handler(exc_handler, copy_ctx.write_row, record) + self._run_with_exception_handler( + exc_handler, self._copy_arrow_batches, copy_sql, all_batches, columns, metrics, use_binary + ) self._check_pending_exception(exc_handler) telemetry_payload = self._ingest_telemetry_from_batches(metrics["rows_processed"], metrics["bytes_processed"]) telemetry_payload["destination"] = table self._attach_partition_telemetry(telemetry_payload, partitioner) return self._storage_job(telemetry_payload, telemetry) + def _copy_arrow_batches( + self, + copy_sql: "PsycopgComposed", + batches: "abc.Iterable[ArrowRecordBatch]", + columns: list[str], + metrics: dict[str, int], + use_binary: bool, + ) -> None: + with self.with_cursor(self.connection) as cursor, cursor.copy(copy_sql) as copy_ctx: + if use_binary: + for chunk in self._iter_pgcopy_binary_chunks(batches, columns=columns, metrics=metrics): + copy_ctx.write(chunk) + else: + for batch in batches: + metrics["rows_processed"] += int(batch.num_rows) + metrics["bytes_processed"] += int(batch.nbytes) + records = self._arrow_batch_to_rows(batch, columns) + if self._arrow_rows_need_preparation(batch): + records = cast( + "list[Any]", self.prepare_driver_parameters(records, self.statement_config, is_many=True) + ) + for record in records: + copy_ctx.write_row(record) + def collect_rows(self, cursor: Any, fetched: "list[Any]") -> "tuple[list[Any], list[str], int]": """Collect psycopg sync rows for the direct execution path.""" data = cast("list[Any] | None", fetched) or [] @@ -623,18 +630,15 @@ def resolve_rowcount(self, cursor: Any) -> int: """Resolve rowcount from psycopg cursor for the direct execution path.""" return resolve_rowcount(cursor) - def _format_wire_sql(self, sql: str) -> str: - """Apply wire-level literal percent escaping and sqlcommenter formatting.""" - formatted = sql - if "%" in formatted: - formatted = escape_literal_percent(formatted, True, ParameterValidator()) - if not self.statement_config.enable_sqlcommenter: - return formatted - return format_wire_sqlcomment(formatted, self.statement_config) - def _execute_cache_hit( self, sql: str, params: "tuple[Any, ...] | list[Any] | dict[str, Any]", cached: "CachedQuery" ) -> "SQLResult": + expected_percents = cached.parameter_profile.total_count if params else 0 + if cached.compiled_sql.count("%") != expected_percents: + statement = self._cached_statement( + sql, params, cached, params, params_are_simple=True, compiled_sql=cached.compiled_sql + ) + return self._execute_cached_statement(statement) return super()._execute_cache_hit(sql, params, cached) def _execute_stack_pipeline( @@ -1028,16 +1032,14 @@ async def rollback(self) -> None: async def set_migration_session_schema(self, schema: str) -> None: """Set the PostgreSQL search path for migration SQL.""" normalized_schema = normalize_identifier(schema, "postgres") - quoted_schema = quote_identifier(normalized_schema) - sql = f'SET LOCAL search_path TO {quoted_schema}, "$user", public' + sql = PsycopgSQL('SET LOCAL search_path TO {}, "$user", public').format(pg_sql.Identifier(normalized_schema)) async with self.with_cursor(self.connection) as cursor: await cursor.execute(sql) async def set_migration_non_transactional_schema(self, schema: str) -> None: """Set the PostgreSQL search path for non-transactional migration SQL.""" normalized_schema = normalize_identifier(schema, "postgres") - quoted_schema = quote_identifier(normalized_schema) - sql = f'SET search_path TO {quoted_schema}, "$user", public' + sql = PsycopgSQL('SET search_path TO {}, "$user", public').format(pg_sql.Identifier(normalized_schema)) async with self.with_cursor(self.connection) as cursor: await cursor.execute(sql) @@ -1167,30 +1169,39 @@ async def load_from_arrow( copy_sql = build_copy_from_command(table, columns, binary=use_binary) all_batches = chain((first_batch,), (b for b in batch_iter if b.num_rows > 0)) exc_handler = self.handle_database_exceptions() - async with AsyncExitStack() as stack: - cursor = await stack.enter_async_context(self.with_cursor(self.connection)) - copy_ctx = await stack.enter_async_context(cursor.copy(copy_sql)) - if use_binary: - for chunk in self._iter_pgcopy_binary_chunks(all_batches, columns=columns, metrics=metrics): - await self._run_with_exception_handler(exc_handler, copy_ctx.write, chunk) - else: - for batch in all_batches: - metrics["rows_processed"] += int(batch.num_rows) - metrics["bytes_processed"] += int(batch.nbytes) - records = self._arrow_batch_to_rows(batch, columns) - if self._arrow_rows_need_preparation(batch): - records = cast( - "list[Any]", - self.prepare_driver_parameters(records, self.statement_config, is_many=True), - ) - for record in records: - await self._run_with_exception_handler(exc_handler, copy_ctx.write_row, record) + await self._run_with_exception_handler( + exc_handler, self._copy_arrow_batches, copy_sql, all_batches, columns, metrics, use_binary + ) self._check_pending_exception(exc_handler) telemetry_payload = self._ingest_telemetry_from_batches(metrics["rows_processed"], metrics["bytes_processed"]) telemetry_payload["destination"] = table self._attach_partition_telemetry(telemetry_payload, partitioner) return self._storage_job(telemetry_payload, telemetry) + async def _copy_arrow_batches( + self, + copy_sql: "PsycopgComposed", + batches: "abc.Iterable[ArrowRecordBatch]", + columns: list[str], + metrics: dict[str, int], + use_binary: bool, + ) -> None: + async with self.with_cursor(self.connection) as cursor, cursor.copy(copy_sql) as copy_ctx: + if use_binary: + for chunk in self._iter_pgcopy_binary_chunks(batches, columns=columns, metrics=metrics): + await copy_ctx.write(chunk) + else: + for batch in batches: + metrics["rows_processed"] += int(batch.num_rows) + metrics["bytes_processed"] += int(batch.nbytes) + records = self._arrow_batch_to_rows(batch, columns) + if self._arrow_rows_need_preparation(batch): + records = cast( + "list[Any]", self.prepare_driver_parameters(records, self.statement_config, is_many=True) + ) + for record in records: + await copy_ctx.write_row(record) + def collect_rows(self, cursor: Any, fetched: "list[Any]") -> "tuple[list[Any], list[str], int]": """Collect psycopg async rows for the direct execution path.""" data = cast("list[Any] | None", fetched) or [] @@ -1201,18 +1212,15 @@ def resolve_rowcount(self, cursor: Any) -> int: """Resolve rowcount from psycopg cursor for the direct execution path.""" return resolve_rowcount(cursor) - def _format_wire_sql(self, sql: str) -> str: - """Apply wire-level literal percent escaping and sqlcommenter formatting.""" - formatted = sql - if "%" in formatted: - formatted = escape_literal_percent(formatted, True, ParameterValidator()) - if not self.statement_config.enable_sqlcommenter: - return formatted - return format_wire_sqlcomment(formatted, self.statement_config) - async def _execute_cache_hit( self, sql: str, params: "tuple[Any, ...] | list[Any] | dict[str, Any]", cached: "CachedQuery" ) -> "SQLResult": + expected_percents = cached.parameter_profile.total_count if params else 0 + if cached.compiled_sql.count("%") != expected_percents: + statement = self._cached_statement( + sql, params, cached, params, params_are_simple=True, compiled_sql=cached.compiled_sql + ) + return await self._execute_cached_statement(statement) return await super()._execute_cache_hit(sql, params, cached) async def _execute_stack_pipeline( diff --git a/sqlspec/adapters/pymssql/adk/store.py b/sqlspec/adapters/pymssql/adk/store.py index 2a2974b59..71040934a 100644 --- a/sqlspec/adapters/pymssql/adk/store.py +++ b/sqlspec/adapters/pymssql/adk/store.py @@ -49,6 +49,8 @@ class PymssqlADKStore(BaseSyncADKStore["PymssqlConfig"]): connector_name: ClassVar[str] = "pymssql" __slots__ = ("_json_column_type", "_native_json") + _max_table_name_parts = 2 + def __init__(self, config: "PymssqlConfig") -> None: super().__init__(config) adk_config = _adk_config(config) @@ -433,6 +435,8 @@ class PymssqlADKMemoryStore(BaseSyncADKMemoryStore["PymssqlConfig"]): __slots__ = ("_json_column_type", "_native_json") + _max_table_name_parts = 2 + def __init__(self, config: "PymssqlConfig") -> None: super().__init__(config) adk_config = _adk_config(config) diff --git a/sqlspec/adapters/pymysql/core.py b/sqlspec/adapters/pymysql/core.py index 7f16df803..2d0da3c4d 100644 --- a/sqlspec/adapters/pymysql/core.py +++ b/sqlspec/adapters/pymysql/core.py @@ -525,8 +525,6 @@ def escape_literal_percent(sql: str, parameters: Any, validator: "ParameterValid """Escape literal percent characters before driver-side interpolation.""" if not parameters or "%" not in sql: return sql - if isinstance(parameters, (tuple, list)) and sql.count("%") == len(parameters): - return sql keep = {info.position for info in validator.extract_parameters(sql)} if len(keep) == sql.count("%"): return sql diff --git a/sqlspec/adapters/spanner/config.py b/sqlspec/adapters/spanner/config.py index bff9e6984..31e54e4e5 100644 --- a/sqlspec/adapters/spanner/config.py +++ b/sqlspec/adapters/spanner/config.py @@ -46,6 +46,7 @@ from sqlspec.exceptions import ImproperConfigurationError from sqlspec.extensions.events import EventRuntimeHints from sqlspec.utils.config_tools import normalize_connection_config +from sqlspec.utils.module_loader import ensure_pyarrow if TYPE_CHECKING: from collections.abc import Awaitable, Callable, Mapping @@ -67,6 +68,7 @@ from sqlspec.config import ExtensionConfigs from sqlspec.core import StatementConfig from sqlspec.observability import ObservabilityConfig + from sqlspec.storage import StorageCapabilities __all__ = ( "SpannerAsyncConfig", @@ -424,6 +426,14 @@ def __init__( self._client: Client | None = None self._database: Database | None = None + def _build_storage_capabilities(self) -> "StorageCapabilities": + """Enable Arrow conversion and mutation ingestion when PyArrow is installed.""" + capabilities = super()._build_storage_capabilities() + arrow_ready = self._dependency_available(ensure_pyarrow) + capabilities["arrow_import_enabled"] = arrow_ready + capabilities["arrow_export_enabled"] = arrow_ready + return capabilities + def get_database(self) -> "Database": """Return the configured database, creating the pool and client on first use. @@ -828,6 +838,14 @@ def __init__( self._client: AsyncClient | None = None self._database: AsyncDatabase | None = None + def _build_storage_capabilities(self) -> "StorageCapabilities": + """Enable Arrow conversion and mutation ingestion when PyArrow is installed.""" + capabilities = super()._build_storage_capabilities() + arrow_ready = self._dependency_available(ensure_pyarrow) + capabilities["arrow_import_enabled"] = arrow_ready + capabilities["arrow_export_enabled"] = arrow_ready + return capabilities + async def get_database(self) -> "AsyncDatabase": """Return the configured database, creating the pool and client on first use. diff --git a/sqlspec/adapters/sqlanywhere/_typing.py b/sqlspec/adapters/sqlanywhere/_typing.py index a58b8d854..5525e8df0 100644 --- a/sqlspec/adapters/sqlanywhere/_typing.py +++ b/sqlspec/adapters/sqlanywhere/_typing.py @@ -109,9 +109,8 @@ def __enter__(self) -> "SqlanywhereDriver": def __exit__( self, exc_type: "type[BaseException] | None", exc_val: "BaseException | None", exc_tb: "TracebackType | None" ) -> "bool | None": - if self._driver is not None and (exc_type is not None or self._driver._connection_in_transaction()): - with contextlib.suppress(Exception): - self._driver.rollback() + if self._driver is not None: + self._driver._release_open_work() if self._connection is not None: self._release_connection(self._connection, exc_type=exc_type, exc_val=exc_val, exc_tb=exc_tb) self._connection = None diff --git a/sqlspec/adapters/sqlanywhere/adk/store.py b/sqlspec/adapters/sqlanywhere/adk/store.py index 36e010303..a41869198 100644 --- a/sqlspec/adapters/sqlanywhere/adk/store.py +++ b/sqlspec/adapters/sqlanywhere/adk/store.py @@ -34,6 +34,8 @@ class SqlanywhereADKStore(BaseSyncADKStore["SqlanywhereConfig"]): connector_name: ClassVar[str] = "sqlanywhere" __slots__ = () + _max_table_name_parts = 2 + def __init__(self, config: "SqlanywhereConfig") -> None: super().__init__(config) @@ -411,6 +413,8 @@ class SqlanywhereADKMemoryStore(BaseSyncADKMemoryStore["SqlanywhereConfig"]): __slots__ = () + _max_table_name_parts = 2 + def __init__(self, config: "SqlanywhereConfig") -> None: super().__init__(config) diff --git a/sqlspec/adapters/sqlite/_typing.py b/sqlspec/adapters/sqlite/_typing.py index 83abb5df9..4b3e87a4e 100644 --- a/sqlspec/adapters/sqlite/_typing.py +++ b/sqlspec/adapters/sqlite/_typing.py @@ -97,7 +97,7 @@ class SqliteSessionContext: "_statement_config", ) - _default_driver_class: ClassVar[Any] = None + _default_driver_class: ClassVar["type[SqliteDriver] | None"] = None def __init__( self, diff --git a/sqlspec/adapters/sqlite/pool.py b/sqlspec/adapters/sqlite/pool.py index 9d8f9b61d..b69711d59 100644 --- a/sqlspec/adapters/sqlite/pool.py +++ b/sqlspec/adapters/sqlite/pool.py @@ -128,8 +128,6 @@ def new_connection(self) -> SqliteConnection: _execute_and_close(connection, "PRAGMA synchronous = OFF") _execute_and_close(connection, "PRAGMA temp_store = MEMORY") _execute_and_close(connection, f"PRAGMA cache_size = {SQLITE_MEMORY_CACHE_SIZE}") - if self._is_shared_memory: - _execute_and_close(connection, "PRAGMA read_uncommitted = ON") else: _enable_wal(connection) _execute_and_close(connection, "PRAGMA synchronous = NORMAL") diff --git a/sqlspec/adapters/sybase/adk/store.py b/sqlspec/adapters/sybase/adk/store.py index 525592834..b9d7110c7 100644 --- a/sqlspec/adapters/sybase/adk/store.py +++ b/sqlspec/adapters/sybase/adk/store.py @@ -40,6 +40,8 @@ class SybaseADKStore(BaseSyncADKStore["SybaseConfig"]): connector_name: ClassVar[str] = "sybase" __slots__ = () + _max_table_name_parts = 2 + def __init__(self, config: "SybaseConfig") -> None: super().__init__(config) @@ -447,6 +449,8 @@ class SybaseADKMemoryStore(BaseSyncADKMemoryStore["SybaseConfig"]): __slots__ = () + _max_table_name_parts = 2 + def __init__(self, config: "SybaseConfig") -> None: super().__init__(config) diff --git a/sqlspec/data_dictionary/dialects/duckdb/sql/tables.sql b/sqlspec/data_dictionary/dialects/duckdb/sql/tables.sql index 63b774447..ee3c20dfb 100644 --- a/sqlspec/data_dictionary/dialects/duckdb/sql/tables.sql +++ b/sqlspec/data_dictionary/dialects/duckdb/sql/tables.sql @@ -55,4 +55,5 @@ FROM information_schema.tables t LEFT JOIN dependency_levels dl ON dl.table_name = t.table_name WHERE t.table_type = 'BASE TABLE' AND t.table_schema = COALESCE(:schema_name, current_schema()) + AND (:table_name IS NULL OR t.table_name = :table_name) ORDER BY level, t.table_name; diff --git a/sqlspec/data_dictionary/dialects/mssql/sql/constraints.sql b/sqlspec/data_dictionary/dialects/mssql/sql/constraints.sql index a70464958..1fccee7ca 100644 --- a/sqlspec/data_dictionary/dialects/mssql/sql/constraints.sql +++ b/sqlspec/data_dictionary/dialects/mssql/sql/constraints.sql @@ -1,27 +1,25 @@ -- name: by_schema -- dialect: mssql -WITH key_constraint_columns AS ( - SELECT - ic.object_id, - ic.index_id, - STRING_AGG(c.name, ',') WITHIN GROUP (ORDER BY ic.key_ordinal) AS column_name - FROM sys.index_columns AS ic - INNER JOIN sys.columns AS c ON ic.object_id = c.object_id AND ic.column_id = c.column_id - WHERE ic.is_included_column = 0 - GROUP BY ic.object_id, ic.index_id -) SELECT s.name AS schema_name, t.name AS table_name, kc.name AS constraint_name, kc.type_desc AS constraint_type, CAST(kc.is_system_named AS BIT) AS is_system_named, - kcc.column_name, + STUFF(( + SELECT ',' + c2.name + FROM sys.index_columns AS ic2 + INNER JOIN sys.columns AS c2 ON ic2.object_id = c2.object_id AND ic2.column_id = c2.column_id + WHERE ic2.object_id = kc.parent_object_id + AND ic2.index_id = kc.unique_index_id + AND ic2.is_included_column = 0 + ORDER BY ic2.key_ordinal + FOR XML PATH(''), TYPE + ).value('.', 'NVARCHAR(MAX)'), 1, 1, '') AS column_name, NULL AS definition FROM sys.key_constraints AS kc INNER JOIN sys.tables AS t ON kc.parent_object_id = t.object_id INNER JOIN sys.schemas AS s ON t.schema_id = s.schema_id -LEFT JOIN key_constraint_columns AS kcc ON kc.parent_object_id = kcc.object_id AND kc.unique_index_id = kcc.index_id WHERE (:schema_name IS NULL OR s.name = :schema_name) UNION ALL SELECT @@ -54,28 +52,26 @@ ORDER BY schema_name, table_name, constraint_name; -- name: by_table -- dialect: mssql -WITH key_constraint_columns AS ( - SELECT - ic.object_id, - ic.index_id, - STRING_AGG(c.name, ',') WITHIN GROUP (ORDER BY ic.key_ordinal) AS column_name - FROM sys.index_columns AS ic - INNER JOIN sys.columns AS c ON ic.object_id = c.object_id AND ic.column_id = c.column_id - WHERE ic.is_included_column = 0 - GROUP BY ic.object_id, ic.index_id -) SELECT s.name AS schema_name, t.name AS table_name, kc.name AS constraint_name, kc.type_desc AS constraint_type, CAST(kc.is_system_named AS BIT) AS is_system_named, - kcc.column_name, + STUFF(( + SELECT ',' + c2.name + FROM sys.index_columns AS ic2 + INNER JOIN sys.columns AS c2 ON ic2.object_id = c2.object_id AND ic2.column_id = c2.column_id + WHERE ic2.object_id = kc.parent_object_id + AND ic2.index_id = kc.unique_index_id + AND ic2.is_included_column = 0 + ORDER BY ic2.key_ordinal + FOR XML PATH(''), TYPE + ).value('.', 'NVARCHAR(MAX)'), 1, 1, '') AS column_name, NULL AS definition FROM sys.key_constraints AS kc INNER JOIN sys.tables AS t ON kc.parent_object_id = t.object_id INNER JOIN sys.schemas AS s ON t.schema_id = s.schema_id -LEFT JOIN key_constraint_columns AS kcc ON kc.parent_object_id = kcc.object_id AND kc.unique_index_id = kcc.index_id WHERE s.name = :schema_name AND t.name = :table_name UNION ALL diff --git a/sqlspec/data_dictionary/dialects/mssql/sql/indexes.sql b/sqlspec/data_dictionary/dialects/mssql/sql/indexes.sql index e1dc811c1..768c6b650 100644 --- a/sqlspec/data_dictionary/dialects/mssql/sql/indexes.sql +++ b/sqlspec/data_dictionary/dialects/mssql/sql/indexes.sql @@ -1,25 +1,5 @@ -- name: by_schema -- dialect: mssql -WITH index_key_columns AS ( - SELECT - ic.object_id, - ic.index_id, - STRING_AGG(CAST(c.name AS NVARCHAR(MAX)), ',') WITHIN GROUP (ORDER BY ic.key_ordinal) AS columns - FROM sys.index_columns AS ic - INNER JOIN sys.columns AS c ON ic.object_id = c.object_id AND ic.column_id = c.column_id - WHERE ic.is_included_column = 0 - GROUP BY ic.object_id, ic.index_id -), -index_included_columns AS ( - SELECT - ic.object_id, - ic.index_id, - STRING_AGG(CAST(c.name AS NVARCHAR(MAX)), ',') WITHIN GROUP (ORDER BY ic.index_column_id) AS included_columns - FROM sys.index_columns AS ic - INNER JOIN sys.columns AS c ON ic.object_id = c.object_id AND ic.column_id = c.column_id - WHERE ic.is_included_column = 1 - GROUP BY ic.object_id, ic.index_id -) SELECT s.name AS schema_name, t.name AS table_name, @@ -32,41 +12,37 @@ SELECT i.filter_definition, ds.name AS data_space_name, ps.name AS partition_scheme_name, - kc.columns, - inc.included_columns + STUFF(( + SELECT ',' + c2.name + FROM sys.index_columns AS ic2 + INNER JOIN sys.columns AS c2 ON ic2.object_id = c2.object_id AND ic2.column_id = c2.column_id + WHERE ic2.object_id = i.object_id + AND ic2.index_id = i.index_id + AND ic2.is_included_column = 0 + ORDER BY ic2.key_ordinal + FOR XML PATH(''), TYPE + ).value('.', 'NVARCHAR(MAX)'), 1, 1, '') AS columns, + STUFF(( + SELECT ',' + c3.name + FROM sys.index_columns AS ic3 + INNER JOIN sys.columns AS c3 ON ic3.object_id = c3.object_id AND ic3.column_id = c3.column_id + WHERE ic3.object_id = i.object_id + AND ic3.index_id = i.index_id + AND ic3.is_included_column = 1 + ORDER BY ic3.index_column_id + FOR XML PATH(''), TYPE + ).value('.', 'NVARCHAR(MAX)'), 1, 1, '') AS included_columns FROM sys.indexes AS i INNER JOIN sys.tables AS t ON i.object_id = t.object_id INNER JOIN sys.schemas AS s ON t.schema_id = s.schema_id LEFT JOIN sys.data_spaces AS ds ON i.data_space_id = ds.data_space_id LEFT JOIN sys.partition_schemes AS ps ON ds.data_space_id = ps.data_space_id -LEFT JOIN index_key_columns AS kc ON kc.object_id = i.object_id AND kc.index_id = i.index_id -LEFT JOIN index_included_columns AS inc ON inc.object_id = i.object_id AND inc.index_id = i.index_id WHERE (:schema_name IS NULL OR s.name = :schema_name) AND i.name IS NOT NULL ORDER BY s.name, t.name, i.name; -- name: by_table -- dialect: mssql -WITH index_key_columns AS ( - SELECT - ic.object_id, - ic.index_id, - STRING_AGG(CAST(c.name AS NVARCHAR(MAX)), ',') WITHIN GROUP (ORDER BY ic.key_ordinal) AS columns - FROM sys.index_columns AS ic - INNER JOIN sys.columns AS c ON ic.object_id = c.object_id AND ic.column_id = c.column_id - WHERE ic.is_included_column = 0 - GROUP BY ic.object_id, ic.index_id -), -index_included_columns AS ( - SELECT - ic.object_id, - ic.index_id, - STRING_AGG(CAST(c.name AS NVARCHAR(MAX)), ',') WITHIN GROUP (ORDER BY ic.index_column_id) AS included_columns - FROM sys.index_columns AS ic - INNER JOIN sys.columns AS c ON ic.object_id = c.object_id AND ic.column_id = c.column_id - WHERE ic.is_included_column = 1 - GROUP BY ic.object_id, ic.index_id -) SELECT s.name AS schema_name, t.name AS table_name, @@ -79,15 +55,31 @@ SELECT i.filter_definition, ds.name AS data_space_name, ps.name AS partition_scheme_name, - kc.columns, - inc.included_columns + STUFF(( + SELECT ',' + c2.name + FROM sys.index_columns AS ic2 + INNER JOIN sys.columns AS c2 ON ic2.object_id = c2.object_id AND ic2.column_id = c2.column_id + WHERE ic2.object_id = i.object_id + AND ic2.index_id = i.index_id + AND ic2.is_included_column = 0 + ORDER BY ic2.key_ordinal + FOR XML PATH(''), TYPE + ).value('.', 'NVARCHAR(MAX)'), 1, 1, '') AS columns, + STUFF(( + SELECT ',' + c3.name + FROM sys.index_columns AS ic3 + INNER JOIN sys.columns AS c3 ON ic3.object_id = c3.object_id AND ic3.column_id = c3.column_id + WHERE ic3.object_id = i.object_id + AND ic3.index_id = i.index_id + AND ic3.is_included_column = 1 + ORDER BY ic3.index_column_id + FOR XML PATH(''), TYPE + ).value('.', 'NVARCHAR(MAX)'), 1, 1, '') AS included_columns FROM sys.indexes AS i INNER JOIN sys.tables AS t ON i.object_id = t.object_id INNER JOIN sys.schemas AS s ON t.schema_id = s.schema_id LEFT JOIN sys.data_spaces AS ds ON i.data_space_id = ds.data_space_id LEFT JOIN sys.partition_schemes AS ps ON ds.data_space_id = ps.data_space_id -LEFT JOIN index_key_columns AS kc ON kc.object_id = i.object_id AND kc.index_id = i.index_id -LEFT JOIN index_included_columns AS inc ON inc.object_id = i.object_id AND inc.index_id = i.index_id WHERE s.name = :schema_name AND t.name = :table_name AND i.name IS NOT NULL diff --git a/sqlspec/data_dictionary/dialects/mysql/dictionary.py b/sqlspec/data_dictionary/dialects/mysql/dictionary.py index ec2bbea69..b9571e255 100644 --- a/sqlspec/data_dictionary/dialects/mysql/dictionary.py +++ b/sqlspec/data_dictionary/dialects/mysql/dictionary.py @@ -56,7 +56,6 @@ class MySQLAsyncDataDictionary(AsyncDataDictionaryBase): def __init__(self) -> None: super().__init__() self._engine_version_cache: dict[int, MySQLEngineVersion | None] = {} - self._ddl_context_cache: dict[int, tuple[str | None, str | None, str | None]] = {} async def get_version(self, driver: "AsyncDriverAdapterBase") -> "VersionInfo | None": """Get MySQL database version information.""" @@ -253,35 +252,18 @@ async def _get_native_object_ddl( _ = include_dependencies, redact if object_type.strip().upper() not in MYSQL_SHOW_CREATE_OBJECT_TYPES: return None - cache_key = _get_driver_cache_key(driver) - ddl_ctx = self._ddl_context_cache.get(cache_key) - if ddl_ctx is None: - try: - row_ctx = await driver.select_one( - "SELECT VERSION() AS raw_version, @@sql_mode AS sql_mode, @@sql_quote_show_create AS sql_quote_show_create", - schema_type=dict, - ) - ddl_ctx = ( - str(row_ctx.get("raw_version")) if row_ctx.get("raw_version") is not None else None, - str(row_ctx.get("sql_mode")) if row_ctx.get("sql_mode") is not None else None, - str(row_ctx.get("sql_quote_show_create")) - if row_ctx.get("sql_quote_show_create") is not None - else None, - ) - except Exception: - raw_version = await driver.select_value_or_none(self.get_query_text("version", "current")) - sql_mode = await driver.select_value_or_none("SELECT @@sql_mode") - sql_quote_show_create = await driver.select_value_or_none("SELECT @@sql_quote_show_create") - ddl_ctx = ( - str(raw_version) if raw_version is not None else None, - str(sql_mode) if sql_mode is not None else None, - str(sql_quote_show_create) if sql_quote_show_create is not None else None, - ) - if len(self._ddl_context_cache) >= _MAX_ENGINE_VERSION_CACHE_SIZE: - self._ddl_context_cache.pop(next(iter(self._ddl_context_cache))) - self._ddl_context_cache[cache_key] = ddl_ctx - - raw_version, sql_mode, sql_quote_show_create = ddl_ctx + try: + row_ctx = await driver.select_one( + "SELECT VERSION() AS raw_version, @@sql_mode AS sql_mode, @@sql_quote_show_create AS sql_quote_show_create", + schema_type=dict, + ) + raw_version = row_ctx.get("raw_version") + sql_mode = row_ctx.get("sql_mode") + sql_quote_show_create = row_ctx.get("sql_quote_show_create") + except Exception: + raw_version = await driver.select_value_or_none(self.get_query_text("version", "current")) + sql_mode = await driver.select_value_or_none("SELECT @@sql_mode") + sql_quote_show_create = await driver.select_value_or_none("SELECT @@sql_quote_show_create") statement = build_mysql_show_create_statement(object_name, schema, object_type) try: row = await driver.select_one(statement, schema_type=dict) @@ -355,7 +337,7 @@ async def _select_domain_result( params: dict[str, Any] = {"schema_name": schema_name} if table_name is not None: params["table_name"] = table_name - if object_name is not None: + if domain == "privileges": params["object_name"] = object_name rows = await driver.select(query.sql, **params) return MetadataResult(domain, capability=capability, items=tuple(rows)) @@ -370,7 +352,6 @@ class MySQLSyncDataDictionary(SyncDataDictionaryBase): def __init__(self) -> None: super().__init__() self._engine_version_cache: dict[int, MySQLEngineVersion | None] = {} - self._ddl_context_cache: dict[int, tuple[str | None, str | None, str | None]] = {} def get_version(self, driver: "SyncDriverAdapterBase") -> "VersionInfo | None": """Get MySQL database version information.""" @@ -567,35 +548,18 @@ def _get_native_object_ddl( _ = include_dependencies, redact if object_type.strip().upper() not in MYSQL_SHOW_CREATE_OBJECT_TYPES: return None - cache_key = _get_driver_cache_key(driver) - ddl_ctx = self._ddl_context_cache.get(cache_key) - if ddl_ctx is None: - try: - row_ctx = driver.select_one( - "SELECT VERSION() AS raw_version, @@sql_mode AS sql_mode, @@sql_quote_show_create AS sql_quote_show_create", - schema_type=dict, - ) - ddl_ctx = ( - str(row_ctx.get("raw_version")) if row_ctx.get("raw_version") is not None else None, - str(row_ctx.get("sql_mode")) if row_ctx.get("sql_mode") is not None else None, - str(row_ctx.get("sql_quote_show_create")) - if row_ctx.get("sql_quote_show_create") is not None - else None, - ) - except Exception: - raw_version = driver.select_value_or_none(self.get_query_text("version", "current")) - sql_mode = driver.select_value_or_none("SELECT @@sql_mode") - sql_quote_show_create = driver.select_value_or_none("SELECT @@sql_quote_show_create") - ddl_ctx = ( - str(raw_version) if raw_version is not None else None, - str(sql_mode) if sql_mode is not None else None, - str(sql_quote_show_create) if sql_quote_show_create is not None else None, - ) - if len(self._ddl_context_cache) >= _MAX_ENGINE_VERSION_CACHE_SIZE: - self._ddl_context_cache.pop(next(iter(self._ddl_context_cache))) - self._ddl_context_cache[cache_key] = ddl_ctx - - raw_version, sql_mode, sql_quote_show_create = ddl_ctx + try: + row_ctx = driver.select_one( + "SELECT VERSION() AS raw_version, @@sql_mode AS sql_mode, @@sql_quote_show_create AS sql_quote_show_create", + schema_type=dict, + ) + raw_version = row_ctx.get("raw_version") + sql_mode = row_ctx.get("sql_mode") + sql_quote_show_create = row_ctx.get("sql_quote_show_create") + except Exception: + raw_version = driver.select_value_or_none(self.get_query_text("version", "current")) + sql_mode = driver.select_value_or_none("SELECT @@sql_mode") + sql_quote_show_create = driver.select_value_or_none("SELECT @@sql_quote_show_create") statement = build_mysql_show_create_statement(object_name, schema, object_type) try: row = driver.select_one(statement, schema_type=dict) @@ -669,7 +633,7 @@ def _select_domain_result( params: dict[str, Any] = {"schema_name": schema_name} if table_name is not None: params["table_name"] = table_name - if object_name is not None: + if domain == "privileges": params["object_name"] = object_name rows = driver.select(query.sql, **params) return MetadataResult(domain, capability=capability, items=tuple(rows)) diff --git a/sqlspec/data_dictionary/dialects/oracle/sql/indexes.sql b/sqlspec/data_dictionary/dialects/oracle/sql/indexes.sql index 8c5125ce9..e6a90f461 100644 --- a/sqlspec/data_dictionary/dialects/oracle/sql/indexes.sql +++ b/sqlspec/data_dictionary/dialects/oracle/sql/indexes.sql @@ -84,4 +84,3 @@ FROM all_ind_expressions ie WHERE ie.index_owner = COALESCE(:schema_name, USER) AND (:table_name IS NULL OR ie.table_name = :table_name) ORDER BY ie.index_owner, ie.table_name, ie.index_name, ie.column_position; - diff --git a/sqlspec/extensions/adk/_table_utils.py b/sqlspec/extensions/adk/_table_utils.py index 6994ef9d5..0a6d3707d 100644 --- a/sqlspec/extensions/adk/_table_utils.py +++ b/sqlspec/extensions/adk/_table_utils.py @@ -6,7 +6,7 @@ __all__ = ("ensure_table_name", "owner_id_column_name", "reset_drop_sql", "unique_statements") -_VALID_TABLE_NAME_PATTERN: Final = re.compile(r"^[a-zA-Z_][a-zA-Z0-9_]*(\.[a-zA-Z_][a-zA-Z0-9_]*)*$") +_VALID_TABLE_NAME_PATTERN: Final = re.compile(r"^[a-zA-Z_][a-zA-Z0-9_]*$") _COLUMN_NAME_PATTERN: Final = re.compile(r"^(\w+)") _MAX_TABLE_NAME_LENGTH: Final = 63 @@ -40,8 +40,16 @@ def reset_drop_sql( return unique_statements(statements) -def ensure_table_name(table_name: str) -> None: - """Validate table name for SQL safety.""" +def ensure_table_name(table_name: str, *, max_parts: int = 1) -> None: + """Validate table identifier components supported by the store. + + Args: + table_name: Table name to validate. + max_parts: Maximum number of identifier components the store supports. + + Raises: + ValueError: If a name component is empty, invalid, or too long. + """ if not table_name: msg = "Table name cannot be empty" raise ValueError(msg) @@ -52,7 +60,7 @@ def ensure_table_name(table_name: str) -> None: msg = f"Table name too long: {len(part)} chars (max {_MAX_TABLE_NAME_LENGTH})" raise ValueError(msg) - if not _VALID_TABLE_NAME_PATTERN.match(table_name): + if len(parts) > max_parts or any(not _VALID_TABLE_NAME_PATTERN.fullmatch(part) for part in parts): msg = ( f"Invalid table name: {table_name!r}. " "Must start with letter/underscore and contain only alphanumeric characters and underscores" diff --git a/sqlspec/extensions/adk/memory/store.py b/sqlspec/extensions/adk/memory/store.py index 8a437ac91..3e27e4d2a 100644 --- a/sqlspec/extensions/adk/memory/store.py +++ b/sqlspec/extensions/adk/memory/store.py @@ -2,7 +2,7 @@ import logging from abc import ABC, abstractmethod -from typing import TYPE_CHECKING, Any, Final, Generic, Literal, TypeVar, cast +from typing import TYPE_CHECKING, Any, ClassVar, Final, Generic, Literal, TypeVar, cast from sqlspec.extensions.adk._config_utils import _adk_memory_store_config, _ADKMemoryStoreConfig from sqlspec.extensions.adk._table_utils import ensure_table_name, owner_id_column_name, unique_statements @@ -46,6 +46,8 @@ class _ADKMemoryStoreCommon(Generic[ConfigT]): "_vector_index_type", ) + _max_table_name_parts: ClassVar[int] = 1 + def __init__(self, config: ConfigT) -> None: """Initialize the ADK memory store. @@ -56,7 +58,7 @@ def __init__(self, config: ConfigT) -> None: store_config = self._store_config_from_extension() self._enabled: bool = store_config.get("enable_memory", True) self._memory_table: str = str(store_config["memory_table"]) - ensure_table_name(self._memory_table) + ensure_table_name(self._memory_table, max_parts=self._max_table_name_parts) self._use_fts: bool = bool(store_config.get("use_fts", False)) self._max_results: int = store_config.get("max_results", 20) self._vector_index_type: str = store_config.get("vector_index_type", "hnsw") diff --git a/sqlspec/extensions/adk/store.py b/sqlspec/extensions/adk/store.py index 3b5af519a..b4a3316f1 100644 --- a/sqlspec/extensions/adk/store.py +++ b/sqlspec/extensions/adk/store.py @@ -4,7 +4,7 @@ import logging from abc import ABC, abstractmethod from datetime import UTC, datetime, timedelta -from typing import TYPE_CHECKING, Any, Final, Generic, TypeVar, cast +from typing import TYPE_CHECKING, Any, ClassVar, Final, Generic, TypeVar, cast from sqlspec.extensions.adk._config_utils import _adk_session_store_config from sqlspec.extensions.adk._table_utils import ensure_table_name, owner_id_column_name, unique_statements @@ -47,6 +47,8 @@ class _ADKStoreCommon(Generic[ConfigT]): "_user_state_table", ) + _max_table_name_parts: ClassVar[int] = 1 + def __init__(self, config: ConfigT) -> None: """Initialize the ADK store. @@ -73,11 +75,11 @@ def __init__(self, config: ConfigT) -> None: self._owner_id_column_name: str | None = ( owner_id_column_name(self._owner_id_column_ddl) if self._owner_id_column_ddl else None ) - ensure_table_name(self._session_table) - ensure_table_name(self._events_table) - ensure_table_name(self._app_state_table) - ensure_table_name(self._user_state_table) - ensure_table_name(self._metadata_table) + ensure_table_name(self._session_table, max_parts=self._max_table_name_parts) + ensure_table_name(self._events_table, max_parts=self._max_table_name_parts) + ensure_table_name(self._app_state_table, max_parts=self._max_table_name_parts) + ensure_table_name(self._user_state_table, max_parts=self._max_table_name_parts) + ensure_table_name(self._metadata_table, max_parts=self._max_table_name_parts) @property def config(self) -> ConfigT: diff --git a/tests/integration/adapters/_shared/_driver_type_system.py b/tests/integration/adapters/_shared/_driver_type_system.py index f031e2023..8a7ff883b 100644 --- a/tests/integration/adapters/_shared/_driver_type_system.py +++ b/tests/integration/adapters/_shared/_driver_type_system.py @@ -1,26 +1,17 @@ -"""Driver-type-system registry and source-equality helpers for contract tests.""" +"""Driver-type-system registry for adapter contract tests.""" -import ast -import difflib import importlib -import inspect -import textwrap from dataclasses import dataclass -from typing import Any, cast +from typing import Any __all__ = ( "DRIVER_FEATURE_CONSUMED_KEYS", "DRIVER_FEATURE_CONTRACT_ENABLED", "DRIVER_FEATURE_CONTRACT_GROUPS", "DRIVER_FEATURE_TYPED_DICTS", - "SOURCE_EQUIVALENCE_ALLOWLIST", - "SOURCE_EQUIVALENCE_CASES", "DriverFeatureContractGroup", - "SourceEquivalenceCase", - "assert_source_equivalent", "get_driver_feature_consumed_keys", "get_driver_feature_keys", - "normalize_source", ) @@ -33,17 +24,6 @@ class DriverFeatureContractGroup: enabled_adapters: tuple[str, ...] = () -@dataclass(frozen=True) -class SourceEquivalenceCase: - """One source-equality tripwire across adapter helper implementations.""" - - case_id: str - group: str - symbol: str - module_paths: tuple[str, ...] - allowlist_key: str - - DRIVER_FEATURE_TYPED_DICTS: dict[str, tuple[str, str]] = { "adbc": ("sqlspec.adapters.adbc.config", "AdbcDriverFeatures"), "aiomysql": ("sqlspec.adapters.aiomysql.config", "AiomysqlDriverFeatures"), @@ -97,9 +77,7 @@ class SourceEquivalenceCase: "json_serializer", "enable_cast_detection", "enable_strict_type_coercion", - "strict_type_coercion", "enable_arrow_extension_types", - "arrow_extension_types", "enable_pgvector", "enable_paradedb", "enable_pg_textsearch", @@ -139,6 +117,7 @@ class SourceEquivalenceCase: "extensions", ), "arrow_odbc": ( + "events_backend", "chunk_size", "max_bytes_per_batch", "max_text_size", @@ -165,6 +144,7 @@ class SourceEquivalenceCase: "enable_local_infile_bulk_load", ), "asyncpg": ( + "enable_pgcopy_binary", "json_serializer", "json_deserializer", "enable_json_codecs", @@ -184,12 +164,12 @@ class SourceEquivalenceCase: "events_backend", "on_connection_create", "pooler_mode", - "pgbouncer", - "pgbouncer_mode", "direct_dsn", "type_codecs", ), "bigquery": ( + "default_query_job_config", + "json_deserializer", "connection_instance", "on_connection_create", "json_serializer", @@ -261,7 +241,7 @@ class SourceEquivalenceCase: "enable_statement_profiling", ), "mssql_python": ( - "use_pool", + "enable_pooling", "json_serializer", "json_deserializer", "on_connection_create", @@ -297,6 +277,7 @@ class SourceEquivalenceCase: "dbms_cloud_credential", ), "psqlpy": ( + "enable_pgcopy_binary", "enable_cast_detection", "enable_pgvector", "enable_paradedb", @@ -308,11 +289,10 @@ class SourceEquivalenceCase: "enable_events", "events_backend", "pooler_mode", - "pgbouncer", - "pgbouncer_mode", "direct_dsn", ), "psycopg": ( + "enable_pgcopy_binary", "enable_pgvector", "enable_paradedb", "enable_pg_textsearch", @@ -327,12 +307,18 @@ class SourceEquivalenceCase: "alloydb_ip_type", "enable_alloydb_columnar", "pooler_mode", - "pgbouncer", - "pgbouncer_mode", "direct_dsn", - "null_pool", + "enable_pooling", + ), + "pymssql": ( + "enable_pooling", + "enable_uuid_conversion", + "json_serializer", + "json_deserializer", + "on_connection_create", + "enable_events", + "events_backend", ), - "pymssql": ("json_serializer", "json_deserializer", "on_connection_create", "enable_events", "events_backend"), "pymysql": ( "json_serializer", "json_deserializer", @@ -346,6 +332,7 @@ class SourceEquivalenceCase: "cloud_sql_ip_type", ), "spanner": ( + "on_connection_create", "enable_uuid_conversion", "json_serializer", "json_deserializer", @@ -419,144 +406,6 @@ class SourceEquivalenceCase: group_name: group.enabled for group_name, group in DRIVER_FEATURE_CONTRACT_GROUPS.items() } -SOURCE_EQUIVALENCE_CASES: tuple[SourceEquivalenceCase, ...] = ( - SourceEquivalenceCase( - case_id="mysql_four_way_format_identifier", - group="mysql_four_way", - symbol="format_identifier", - module_paths=( - "sqlspec.adapters.aiomysql.core", - "sqlspec.adapters.asyncmy.core", - "sqlspec.adapters.mysqlconnector.core", - "sqlspec.adapters.pymysql.core", - ), - allowlist_key="mysql_four_way:format_identifier", - ), - SourceEquivalenceCase( - case_id="mysql_four_way_build_insert_statement", - group="mysql_four_way", - symbol="build_insert_statement", - module_paths=( - "sqlspec.adapters.aiomysql.core", - "sqlspec.adapters.asyncmy.core", - "sqlspec.adapters.mysqlconnector.core", - "sqlspec.adapters.pymysql.core", - ), - allowlist_key="mysql_four_way:build_insert_statement", - ), - SourceEquivalenceCase( - case_id="mysql_four_way_normalize_execute_parameters", - group="mysql_four_way", - symbol="normalize_execute_parameters", - module_paths=( - "sqlspec.adapters.aiomysql.core", - "sqlspec.adapters.asyncmy.core", - "sqlspec.adapters.mysqlconnector.core", - "sqlspec.adapters.pymysql.core", - ), - allowlist_key="mysql_four_way:normalize_execute_parameters", - ), - SourceEquivalenceCase( - case_id="mysql_four_way_normalize_execute_many_parameters", - group="mysql_four_way", - symbol="normalize_execute_many_parameters", - module_paths=( - "sqlspec.adapters.aiomysql.core", - "sqlspec.adapters.asyncmy.core", - "sqlspec.adapters.mysqlconnector.core", - "sqlspec.adapters.pymysql.core", - ), - allowlist_key="mysql_four_way:normalize_execute_many_parameters", - ), - SourceEquivalenceCase( - case_id="mysql_four_way_resolve_row_plan", - group="mysql_four_way", - symbol="resolve_row_plan", - module_paths=( - "sqlspec.adapters.aiomysql.core", - "sqlspec.adapters.asyncmy.core", - "sqlspec.adapters.mysqlconnector.core", - "sqlspec.adapters.pymysql.core", - ), - allowlist_key="mysql_four_way:resolve_row_plan", - ), - SourceEquivalenceCase( - case_id="mysql_four_way_collect_rows", - group="mysql_four_way", - symbol="collect_rows", - module_paths=( - "sqlspec.adapters.aiomysql.core", - "sqlspec.adapters.asyncmy.core", - "sqlspec.adapters.mysqlconnector.core", - "sqlspec.adapters.pymysql.core", - ), - allowlist_key="mysql_four_way:collect_rows", - ), - SourceEquivalenceCase( - case_id="mysql_four_way_collect_stream_rows", - group="mysql_four_way", - symbol="collect_stream_rows", - module_paths=( - "sqlspec.adapters.aiomysql.core", - "sqlspec.adapters.asyncmy.core", - "sqlspec.adapters.mysqlconnector.core", - "sqlspec.adapters.pymysql.core", - ), - allowlist_key="mysql_four_way:collect_stream_rows", - ), - SourceEquivalenceCase( - case_id="mysql_four_way_create_mapped_exception", - group="mysql_four_way", - symbol="create_mapped_exception", - module_paths=( - "sqlspec.adapters.aiomysql.core", - "sqlspec.adapters.asyncmy.core", - "sqlspec.adapters.mysqlconnector.core", - "sqlspec.adapters.pymysql.core", - ), - allowlist_key="mysql_four_way:create_mapped_exception", - ), - SourceEquivalenceCase( - case_id="sqlite_pair_apply_driver_features", - group="sqlite_pair", - symbol="apply_driver_features", - module_paths=("sqlspec.adapters.sqlite.core", "sqlspec.adapters.aiosqlite.core"), - allowlist_key="sqlite_pair:apply_driver_features", - ), - SourceEquivalenceCase( - case_id="sqlite_pair_build_insert_statement", - group="sqlite_pair", - symbol="build_insert_statement", - module_paths=("sqlspec.adapters.sqlite.core", "sqlspec.adapters.aiosqlite.core"), - allowlist_key="sqlite_pair:build_insert_statement", - ), - SourceEquivalenceCase( - case_id="sqlite_pair_normalize_execute_parameters", - group="sqlite_pair", - symbol="normalize_execute_parameters", - module_paths=("sqlspec.adapters.sqlite.core", "sqlspec.adapters.aiosqlite.core"), - allowlist_key="sqlite_pair:normalize_execute_parameters", - ), - SourceEquivalenceCase( - case_id="sqlite_pair_normalize_execute_many_parameters", - group="sqlite_pair", - symbol="normalize_execute_many_parameters", - module_paths=("sqlspec.adapters.sqlite.core", "sqlspec.adapters.aiosqlite.core"), - allowlist_key="sqlite_pair:normalize_execute_many_parameters", - ), - SourceEquivalenceCase( - case_id="mssql_pair_create_mapped_exception", - group="mssql_pair", - symbol="create_mapped_exception", - module_paths=("sqlspec.adapters.mssql_python.core", "sqlspec.adapters.pymssql.core"), - allowlist_key="mssql_pair:create_mapped_exception", - ), -) - -SOURCE_EQUIVALENCE_ALLOWLIST: dict[str, frozenset[str]] = { - case.allowlist_key: frozenset() for case in SOURCE_EQUIVALENCE_CASES -} - def get_driver_feature_keys(adapter: str) -> tuple[str, ...]: """Return the canonical DriverFeatures keys declared by an adapter TypedDict.""" @@ -574,75 +423,6 @@ def get_driver_feature_consumed_keys(adapter: str) -> tuple[str, ...]: return DRIVER_FEATURE_CONSUMED_KEYS[adapter] -class _DocstringStripper(ast.NodeTransformer): - def visit_Module(self, node: ast.Module) -> ast.Module: - node.body = self._strip_body(node.body) - return cast("ast.Module", self.generic_visit(node)) - - def visit_FunctionDef(self, node: ast.FunctionDef) -> ast.FunctionDef: - node.body = self._strip_body(node.body) - return cast("ast.FunctionDef", self.generic_visit(node)) - - def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> ast.AsyncFunctionDef: - node.body = self._strip_body(node.body) - return cast("ast.AsyncFunctionDef", self.generic_visit(node)) - - def visit_ClassDef(self, node: ast.ClassDef) -> ast.ClassDef: - node.body = self._strip_body(node.body) - return cast("ast.ClassDef", self.generic_visit(node)) - - def _strip_body(self, body: list[ast.stmt]) -> list[ast.stmt]: - if body and isinstance(body[0], ast.Expr): - value = body[0].value - if isinstance(value, ast.Constant) and isinstance(value.value, str): - return body[1:] - return body - - -def normalize_source(source: str) -> str: - """Normalize Python source so source-equality checks ignore docstrings and formatting.""" - normalized = textwrap.dedent(source).strip() - tree = ast.parse(normalized) - stripped = _DocstringStripper().visit(tree) - ast.fix_missing_locations(stripped) - return ast.unparse(stripped).strip() - - -def assert_source_equivalent(case: SourceEquivalenceCase) -> None: - """Assert that every unallowlisted module exports the same normalized source.""" - allowed_modules = SOURCE_EQUIVALENCE_ALLOWLIST.get(case.allowlist_key, frozenset()) - normalized_sources: dict[str, str] = {} - for module_path in case.module_paths: - if module_path in allowed_modules: - continue - normalized_sources[module_path] = normalize_source(_load_symbol_source(module_path, case.symbol)) - - if len(normalized_sources) < 2: - msg = f"{case.case_id} does not have enough unallowlisted modules to compare" - raise AssertionError(msg) - - baseline_module, baseline_source = next(iter(normalized_sources.items())) - mismatches: list[str] = [] - for module_path, source in normalized_sources.items(): - if source == baseline_source: - continue - diff = "\n".join( - difflib.unified_diff( - baseline_source.splitlines(), - source.splitlines(), - fromfile=baseline_module, - tofile=module_path, - lineterm="", - ) - ) - mismatches.append(diff) - - if mismatches: - details = "\n\n".join(mismatches) - msg = f"{case.case_id} source mismatch for {case.symbol}:\n{details}" - raise AssertionError(msg) - - def _load_symbol(module_path: str, symbol_name: str) -> Any: try: module = importlib.import_module(module_path) @@ -650,12 +430,3 @@ def _load_symbol(module_path: str, symbol_name: str) -> Any: except (AttributeError, ImportError) as exc: msg = f"Unable to load {symbol_name} from {module_path}: {exc}" raise AssertionError(msg) from exc - - -def _load_symbol_source(module_path: str, symbol_name: str) -> str: - symbol = _load_symbol(module_path, symbol_name) - try: - return inspect.getsource(symbol) - except OSError as exc: - msg = f"Unable to read source for {symbol_name} from {module_path}: {exc}" - raise AssertionError(msg) from exc diff --git a/tests/integration/adapters/_shared/behaviors.py b/tests/integration/adapters/_shared/behaviors.py index 255a73d8e..a77776663 100644 --- a/tests/integration/adapters/_shared/behaviors.py +++ b/tests/integration/adapters/_shared/behaviors.py @@ -620,34 +620,54 @@ async def _assert_aiosqlite_stream_arraysize(driver: object, case: DriverCase) - async def _assert_asyncpg_stream_transaction(driver: object, case: DriverCase) -> None: async_driver = cast("AsyncContractDriver", driver) + connection = cast("Any", driver).connection + was_in_transaction = connection.is_in_transaction() stream = async_driver.select_stream(case.table.select_ordered_sql, chunk_size=10) - await anext(aiter(stream)) - assert stream._source._transaction is not None # pyright: ignore[reportPrivateUsage] - assert stream._source._cursor is not None # pyright: ignore[reportPrivateUsage] - await stream.aclose() - assert stream._source._transaction is None # pyright: ignore[reportPrivateUsage] - assert stream._source._cursor is None # pyright: ignore[reportPrivateUsage] - await async_driver.execute(case.table.select_count_sql) + try: + first = await anext(aiter(stream)) + assert first["value"] == 0 + assert connection.is_in_transaction() is True + finally: + await stream.aclose() + assert connection.is_in_transaction() is was_in_transaction + assert int(cast("int", await async_driver.select_value(case.table.select_count_sql))) == case.streaming_row_count + await async_driver.begin() + try: + await async_driver.execute(case.table.insert_named_sql, name="stream_caller_probe", value=-1, note=None) + async with async_driver.select_stream(case.table.select_ordered_sql, chunk_size=1) as caller_stream: + assert (await anext(aiter(caller_stream)))["value"] == -1 + assert connection.is_in_transaction() is True + finally: + await async_driver.rollback() + assert await async_driver.select(case.table.select_by_name_named_sql, name="stream_caller_probe") == [] -def _assert_psycopg_named_cursor(driver: object, case: DriverCase) -> None: +def _assert_psycopg_stream_transaction(driver: object, case: DriverCase) -> None: sync_driver = cast("SyncContractDriver", driver) + connection = cast("Any", driver).connection + original_status = connection.info.transaction_status stream = sync_driver.select_stream(case.table.select_ordered_sql, chunk_size=10) try: - next(iter(stream)) - assert stream._source._cursor.name # pyright: ignore[reportPrivateUsage] + assert next(iter(stream))["value"] == 0 + assert connection.info.transaction_status == 2 finally: stream.close() + assert connection.info.transaction_status == original_status + assert int(cast("int", sync_driver.select_value(case.table.select_count_sql))) == case.streaming_row_count -async def _assert_psycopg_named_cursor_async(driver: object, case: DriverCase) -> None: +async def _assert_psycopg_stream_transaction_async(driver: object, case: DriverCase) -> None: async_driver = cast("AsyncContractDriver", driver) + connection = cast("Any", driver).connection + original_status = connection.info.transaction_status stream = async_driver.select_stream(case.table.select_ordered_sql, chunk_size=10) try: - await anext(aiter(stream)) - assert stream._source._cursor.name # pyright: ignore[reportPrivateUsage] + assert (await anext(aiter(stream)))["value"] == 0 + assert connection.info.transaction_status == 2 finally: await stream.aclose() + assert connection.info.transaction_status == original_status + assert int(cast("int", await async_driver.select_value(case.table.select_count_sql))) == case.streaming_row_count def _assert_pymysql_sscursor(driver: object, case: DriverCase) -> None: @@ -693,8 +713,8 @@ async def _assert_mysqlconnector_unbuffered_async(driver: object, case: DriverCa register_sync_extra_assertion("streaming_native:sqlite", _STREAMING_SCOPE, _assert_sqlite_stream_arraysize) register_async_extra_assertion("streaming_native:aiosqlite", _STREAMING_SCOPE, _assert_aiosqlite_stream_arraysize) register_async_extra_assertion("streaming_native:asyncpg", _STREAMING_SCOPE, _assert_asyncpg_stream_transaction) -register_sync_extra_assertion("streaming_native:psycopg", _STREAMING_SCOPE, _assert_psycopg_named_cursor) -register_async_extra_assertion("streaming_native:psycopg", _STREAMING_SCOPE, _assert_psycopg_named_cursor_async) +register_sync_extra_assertion("streaming_native:psycopg", _STREAMING_SCOPE, _assert_psycopg_stream_transaction) +register_async_extra_assertion("streaming_native:psycopg", _STREAMING_SCOPE, _assert_psycopg_stream_transaction_async) register_sync_extra_assertion("streaming_native:pymysql", _STREAMING_SCOPE, _assert_pymysql_sscursor) register_async_extra_assertion("streaming_native:aiomysql", _STREAMING_SCOPE, _assert_async_sscursor) register_async_extra_assertion("streaming_native:asyncmy", _STREAMING_SCOPE, _assert_async_sscursor) @@ -758,58 +778,6 @@ class _StreamErrorProbe(Exception): """Sentinel raised to force a mid-stream chunk fetch failure without a live database error.""" -class _SyncTransactionRecorder: - """Delegating transaction wrapper recording whether commit or rollback fired during stream close.""" - - def __init__(self, inner: Any, log: "list[str]") -> None: - object.__setattr__(self, "_inner", inner) - object.__setattr__(self, "_log", log) - - def __getattr__(self, name: str) -> Any: - return getattr(object.__getattribute__(self, "_inner"), name) - - def __setattr__(self, name: str, value: Any) -> None: - setattr(object.__getattribute__(self, "_inner"), name, value) - - def commit(self, *args: Any, **kwargs: Any) -> Any: - object.__getattribute__(self, "_log").append("commit") - return object.__getattribute__(self, "_inner").commit(*args, **kwargs) - - def rollback(self, *args: Any, **kwargs: Any) -> Any: - object.__getattribute__(self, "_log").append("rollback") - return object.__getattribute__(self, "_inner").rollback(*args, **kwargs) - - def __exit__(self, exc_type: Any, exc: Any, tb: Any) -> Any: - object.__getattribute__(self, "_log").append("rollback" if exc_type is not None else "commit") - return object.__getattribute__(self, "_inner").__exit__(exc_type, exc, tb) - - -class _AsyncTransactionRecorder: - """Delegating async transaction wrapper recording whether commit or rollback fired during stream close.""" - - def __init__(self, inner: Any, log: "list[str]") -> None: - object.__setattr__(self, "_inner", inner) - object.__setattr__(self, "_log", log) - - def __getattr__(self, name: str) -> Any: - return getattr(object.__getattribute__(self, "_inner"), name) - - def __setattr__(self, name: str, value: Any) -> None: - setattr(object.__getattribute__(self, "_inner"), name, value) - - async def commit(self, *args: Any, **kwargs: Any) -> Any: - object.__getattribute__(self, "_log").append("commit") - return await object.__getattribute__(self, "_inner").commit(*args, **kwargs) - - async def rollback(self, *args: Any, **kwargs: Any) -> Any: - object.__getattribute__(self, "_log").append("rollback") - return await object.__getattribute__(self, "_inner").rollback(*args, **kwargs) - - async def __aexit__(self, exc_type: Any, exc: Any, tb: Any) -> Any: - object.__getattribute__(self, "_log").append("rollback" if exc_type is not None else "commit") - return await object.__getattribute__(self, "_inner").__aexit__(exc_type, exc, tb) - - class _SyncStreamFetchFailSource: """Delegating sync stream source that forces the next chunk fetch to fail and delegates close to the origin.""" @@ -846,21 +814,16 @@ def _assert_sync_stream_error_close_rollback(driver: object, case: DriverCase) - """Assert a mid-stream fetch failure rolls back a stream-owned transaction instead of committing.""" sync_driver = cast("SyncContractDriver", driver) table = case.table + sync_driver.rollback() stream = sync_driver.select_stream(table.select_ordered_sql, chunk_size=1) iterator = iter(stream) next(iterator) real_source = stream._source # pyright: ignore[reportPrivateUsage] - assert real_source._transaction is not None # pyright: ignore[reportPrivateUsage] - log: list[str] = [] - real_source._transaction = _SyncTransactionRecorder(real_source._transaction, log) # pyright: ignore[reportPrivateUsage] + sync_driver.execute(table.insert_named_sql, name="stream_rollback_probe", value=-1, note=None) stream._source = _SyncStreamFetchFailSource(real_source) # pyright: ignore[reportPrivateUsage] with pytest.raises(_StreamErrorProbe): next(iterator) - assert log - assert log[0] == "rollback" - assert "commit" not in log - with contextlib.suppress(Exception): - sync_driver.rollback() + assert sync_driver.select(table.select_by_name_named_sql, name="stream_rollback_probe") == [] assert int(cast("int", sync_driver.select_value(table.select_count_sql))) == case.streaming_row_count @@ -868,21 +831,22 @@ async def _assert_async_stream_error_close_rollback(driver: object, case: Driver """Assert a mid-stream fetch failure rolls back a stream-owned transaction instead of committing.""" async_driver = cast("AsyncContractDriver", driver) table = case.table + await async_driver.rollback() stream = async_driver.select_stream(table.select_ordered_sql, chunk_size=1) iterator = aiter(stream) await anext(iterator) real_source = stream._source # pyright: ignore[reportPrivateUsage] - assert real_source._transaction is not None # pyright: ignore[reportPrivateUsage] - log: list[str] = [] - real_source._transaction = _AsyncTransactionRecorder(real_source._transaction, log) # pyright: ignore[reportPrivateUsage] + # CockroachDB rejects another bound statement while the asyncpg cursor's portal is open. + supports_probe_insert = case.adapter != "cockroach_asyncpg" + if supports_probe_insert: + await async_driver.execute(table.insert_named_sql, name="stream_rollback_probe", value=-1, note=None) stream._source = _AsyncStreamFetchFailSource(real_source) # pyright: ignore[reportPrivateUsage] with pytest.raises(_StreamErrorProbe): await anext(iterator) - assert log - assert log[0] == "rollback" - assert "commit" not in log - with contextlib.suppress(Exception): - await async_driver.rollback() + if supports_probe_insert: + assert await async_driver.select(table.select_by_name_named_sql, name="stream_rollback_probe") == [] + else: + assert cast("Any", driver).connection.is_in_transaction() is False assert int(cast("int", await async_driver.select_value(table.select_count_sql))) == case.streaming_row_count diff --git a/tests/integration/adapters/_shared/suite_driver_features_contract.py b/tests/integration/adapters/_shared/suite_driver_features_contract.py index 0058e2753..2931b2c4a 100644 --- a/tests/integration/adapters/_shared/suite_driver_features_contract.py +++ b/tests/integration/adapters/_shared/suite_driver_features_contract.py @@ -1,6 +1,5 @@ -"""Shared driver-feature contracts: registry scaffolds, source tripwires, and folded driver-specific behaviors.""" +"""Shared driver-feature contracts and folded driver-specific behaviors.""" -import importlib from collections.abc import Sequence from typing import cast @@ -9,7 +8,6 @@ from tests.integration.adapters._shared._cases import ( ASYNC_DRIVER_CASES, - DRIVER_CASES, SYNC_DRIVER_CASES, DriverCase, DriverCaseContext, @@ -17,9 +15,6 @@ ) from tests.integration.adapters._shared._driver_type_system import ( DRIVER_FEATURE_CONTRACT_GROUPS, - SOURCE_EQUIVALENCE_CASES, - SourceEquivalenceCase, - assert_source_equivalent, get_driver_feature_consumed_keys, get_driver_feature_keys, ) @@ -37,30 +32,6 @@ validate_extra_assertions, ) -_DRIVER_TYPE_SYSTEM_MODULE = "tests.integration.adapters._shared._driver_type_system" - - -def test_driver_type_system_registry_module_exports_contract_groups() -> None: - """The driver-type-system registry module is importable and exposes the expected contract groups.""" - module = importlib.import_module(_DRIVER_TYPE_SYSTEM_MODULE) - - assert hasattr(module, "DRIVER_FEATURE_CONTRACT_GROUPS") - assert hasattr(module, "DRIVER_FEATURE_CONTRACT_ENABLED") - assert hasattr(module, "DRIVER_FEATURE_CONSUMED_KEYS") - assert hasattr(module, "SOURCE_EQUIVALENCE_CASES") - assert hasattr(module, "SOURCE_EQUIVALENCE_ALLOWLIST") - assert set(module.DRIVER_FEATURE_CONTRACT_GROUPS) == { - "feature_honesty", - "enable_false_semantics", - "parity", - "row_format", - } - assert module.DRIVER_FEATURE_CONTRACT_ENABLED == { - group_name: group.enabled for group_name, group in module.DRIVER_FEATURE_CONTRACT_GROUPS.items() - } - assert set(module.DRIVER_FEATURE_CONSUMED_KEYS) == {case.adapter for case in DRIVER_CASES} - assert {case.group for case in module.SOURCE_EQUIVALENCE_CASES} == {"mysql_four_way", "sqlite_pair", "mssql_pair"} - def _disabled_first_params(group_name: str, cases: Sequence[DriverCase]) -> tuple[ParameterSet, ...]: group = DRIVER_FEATURE_CONTRACT_GROUPS[group_name] @@ -177,12 +148,6 @@ async def test_async_driver_feature_row_format_contract(async_driver_case: Drive await assert_async_driver_feature_row_format_contract(make_config, async_driver_case.case) -@pytest.mark.parametrize("source_equivalence_case", SOURCE_EQUIVALENCE_CASES, ids=lambda case: case.case_id) -def test_driver_feature_source_equivalence(source_equivalence_case: SourceEquivalenceCase) -> None: - """Duplicate helper implementations stay source-equivalent across adapter families.""" - assert_source_equivalent(source_equivalence_case) - - def test_sync_driver_features_contract(sync_driver_case: DriverCaseContext) -> None: """Sync drivers run folded driver-feature proofs opted in via extra_assertions.""" assert_sync_driver_features_contract(sync_driver_case.driver, sync_driver_case.case) diff --git a/tests/integration/adapters/oracle/oracledb/extensions/adk/test_oracle_specific.py b/tests/integration/adapters/oracle/oracledb/extensions/adk/test_oracle_specific.py index 1427181e3..283c89087 100644 --- a/tests/integration/adapters/oracle/oracledb/extensions/adk/test_oracle_specific.py +++ b/tests/integration/adapters/oracle/oracledb/extensions/adk/test_oracle_specific.py @@ -1,7 +1,7 @@ """Oracle-specific ADK store tests for LOB handling, JSON types, and FK columns.""" import json -from collections.abc import AsyncGenerator +from collections.abc import AsyncGenerator, Generator from datetime import UTC, datetime from typing import Any, cast from uuid import uuid4 @@ -16,7 +16,7 @@ OracleSyncADKMemoryStore, OracleSyncADKStore, ) -from sqlspec.extensions.adk import StoredEvent +from sqlspec.extensions.adk import StoredEvent, StoredMemory pytestmark = [pytest.mark.xdist_group("oracle"), pytest.mark.oracledb, pytest.mark.integration] @@ -530,9 +530,9 @@ async def oracle_async_memory_store( @pytest.fixture -async def oracle_sync_memory_store( +def oracle_sync_memory_store( oracle_sync_config: "OracleSyncConfig", -) -> "AsyncGenerator[OracleSyncADKMemoryStore, None]": +) -> "Generator[OracleSyncADKMemoryStore, None, None]": """Create a sync Oracle ADK memory store with tables created per test.""" store = OracleSyncADKMemoryStore(oracle_sync_config) store.create_tables() @@ -550,7 +550,7 @@ async def test_oracle_async_adk_memory_store_insert_and_search_embedding( user_id = "user-vec" now = datetime.now(UTC) embedding = [0.1 * i for i in range(768)] - entries = [ + entries: list[StoredMemory] = [ { "id": _unique_session_id("mem-vec"), "session_id": _unique_session_id("sess-vec"), @@ -585,7 +585,7 @@ def test_oracle_sync_adk_memory_store_insert_and_search_embedding( user_id = "user-vec-sync" now = datetime.now(UTC) embedding = [0.1 * i for i in range(768)] - entries = [ + entries: list[StoredMemory] = [ { "id": _unique_session_id("mem-vec-sync"), "session_id": _unique_session_id("sess-vec-sync"), diff --git a/tests/integration/adapters/postgres/asyncpg/test_driver.py b/tests/integration/adapters/postgres/asyncpg/test_driver.py index 31239e022..e236bac6d 100644 --- a/tests/integration/adapters/postgres/asyncpg/test_driver.py +++ b/tests/integration/adapters/postgres/asyncpg/test_driver.py @@ -483,12 +483,18 @@ async def test_asyncpg_statement_stack_continue_on_error(asyncpg_session: "Async assert verify.get_data()[0]["total"] == 2 -async def test_asyncpg_statement_stack_continue_on_error_inside_transaction(asyncpg_session: "AsyncpgDriver") -> None: +@pytest.mark.parametrize("raw_begin", [False, True]) +async def test_asyncpg_statement_stack_continue_on_error_inside_transaction( + asyncpg_session: "AsyncpgDriver", raw_begin: bool +) -> None: """Continue-on-error inside a user transaction must isolate the failing op via a savepoint.""" await asyncpg_session.execute_script("DELETE FROM test_table_asyncpg") - await asyncpg_session.begin() + if raw_begin: + await asyncpg_session.execute("BEGIN") + else: + await asyncpg_session.begin() stack = ( StatementStack() .push_execute("INSERT INTO test_table_asyncpg (id, name, value) VALUES ($1, $2, $3)", (1, "initial", 5)) @@ -501,6 +507,7 @@ async def test_asyncpg_statement_stack_continue_on_error_inside_transaction(asyn assert results[0].error is None assert results[1].error is not None assert results[2].error is None + assert asyncpg_session.connection.is_in_transaction() verify = await asyncpg_session.execute("SELECT COUNT(*) AS total FROM test_table_asyncpg") assert verify.get_data()[0]["total"] == 2 diff --git a/tests/integration/adapters/postgres/psqlpy/extensions/adk/test_memory_embeddings.py b/tests/integration/adapters/postgres/psqlpy/extensions/adk/test_memory_embeddings.py new file mode 100644 index 000000000..a7fcd6c6b --- /dev/null +++ b/tests/integration/adapters/postgres/psqlpy/extensions/adk/test_memory_embeddings.py @@ -0,0 +1,63 @@ +"""Native vector preservation during batched psqlpy memory insertion.""" + +from datetime import UTC, datetime +from typing import TYPE_CHECKING +from uuid import uuid4 + +import pytest + +from sqlspec.adapters.psqlpy import PsqlpyConfig +from sqlspec.adapters.psqlpy.adk import PsqlpyADKMemoryStore + +if TYPE_CHECKING: + from sqlspec.extensions.adk import StoredMemory + +pytestmark = pytest.mark.xdist_group("pgvector") + + +@pytest.mark.parametrize("with_owner", [False, True]) +async def test_bulk_memory_insertion_preserves_embeddings( + pgvector_config_psqlpy: PsqlpyConfig, with_owner: bool +) -> None: + """Vectors and nulls remain aligned with their rows, including owner-scoped batches.""" + config = PsqlpyConfig( + connection_config=dict(pgvector_config_psqlpy.connection_config), + extension_config={ + "adk": { + "memory_table": f"memory_vectors_{uuid4().hex}", + "vector_dimensions": 2, + "vector_index_type": "none", + "owner_id_column": "tenant_id VARCHAR(128)" if with_owner else None, + } + }, + ) + store = PsqlpyADKMemoryStore(config) + now = datetime.now(UTC) + entries: list[StoredMemory] = [ + { + "id": f"memory-{index}", + "session_id": "session", + "app_name": "app", + "user_id": "user", + "scope": "user", + "event_id": f"event-{index}", + "author": "user", + "timestamp": now, + "content_json": {"text": "memory"}, + "content_text": "memory", + "metadata_json": None, + "inserted_at": now, + "embedding": embedding, + } + for index, embedding in enumerate(((1.0, 0.0), None, (0.0, 1.0))) + ] + try: + await store.create_tables() + assert await store.insert_memory_entries(entries, owner_id="tenant") == 3 + assert await store.insert_memory_entries(entries, owner_id="tenant") == 0 + results = await store.search_entries("", "app", "user", embedding=(1.0, 0.0)) + assert [row["id"] for row in results] == ["memory-0", "memory-2"] + assert [row["embedding"] for row in results] == [[1.0, 0.0], [0.0, 1.0]] + finally: + await store.drop_tables() + await config.close_pool() diff --git a/tests/integration/adapters/postgres/psqlpy/test_driver.py b/tests/integration/adapters/postgres/psqlpy/test_driver.py index 129f538d6..9f3a16287 100644 --- a/tests/integration/adapters/postgres/psqlpy/test_driver.py +++ b/tests/integration/adapters/postgres/psqlpy/test_driver.py @@ -4,6 +4,7 @@ from sqlspec import SQLResult, StatementStack from sqlspec.adapters.psqlpy import PsqlpyConfig, PsqlpyDriver +from sqlspec.exceptions import SQLSpecError pytestmark = pytest.mark.xdist_group("postgres") @@ -244,3 +245,17 @@ async def test_extensions_not_enabled_on_standard_postgres(psqlpy_config: "Psqlp assert psqlpy_config._pgvector_available is False # pyright: ignore[reportPrivateUsage] assert psqlpy_config._paradedb_available is False # pyright: ignore[reportPrivateUsage] assert psqlpy_config.statement_config.dialect == "postgres" + + +async def test_psqlpy_cached_dml_and_errors(psqlpy_session: "PsqlpyDriver") -> None: + """Native psqlpy cache hits preserve counts and SQLSpec exception translation.""" + from psqlpy.exceptions import Error + + for name in ("cached-first", "cached-second"): + result = await psqlpy_session.execute("INSERT INTO test_table_psqlpy (name) VALUES (?)", name) + assert result.rows_affected == 1 + query = "SELECT 10 / ?::integer AS value" + assert (await psqlpy_session.execute(query, 2)).get_data() == [{"value": 5}] + with pytest.raises(SQLSpecError) as error: + await psqlpy_session.execute(query, 0) + assert isinstance(error.value.__cause__, Error) diff --git a/tests/integration/adapters/postgres/psycopg/test_execution_regressions.py b/tests/integration/adapters/postgres/psycopg/test_execution_regressions.py new file mode 100644 index 000000000..7dfc38eb6 --- /dev/null +++ b/tests/integration/adapters/postgres/psycopg/test_execution_regressions.py @@ -0,0 +1,103 @@ +"""PostgreSQL execution contracts across cached and native COPY paths.""" + +import inspect +from typing import TYPE_CHECKING, Any + +import pytest + +from sqlspec.adapters.psycopg import PsycopgAsyncConfig, PsycopgSyncConfig +from sqlspec.exceptions import UniqueViolationError + +if TYPE_CHECKING: + from collections.abc import AsyncIterator + + from sqlspec.adapters.psycopg import PsycopgAsyncDriver, PsycopgSyncDriver + from sqlspec.core import SQLResult + +pytestmark = pytest.mark.xdist_group("postgres") + + +@pytest.fixture(params=["sync", "async"]) +async def psycopg_regression_driver( + request: pytest.FixtureRequest, psycopg_sync_config: PsycopgSyncConfig, psycopg_async_config: PsycopgAsyncConfig +) -> "AsyncIterator[PsycopgSyncDriver | PsycopgAsyncDriver]": + if request.param == "sync": + with psycopg_sync_config.provide_session() as driver: + yield driver + else: + async with psycopg_async_config.provide_session() as async_driver: + yield async_driver + + +async def _execute(driver: "PsycopgSyncDriver | PsycopgAsyncDriver", query: str, *parameters: Any) -> "SQLResult": + result = driver.execute(query, *parameters) + return await result if inspect.isawaitable(result) else result + + +@pytest.mark.parametrize( + ("query", "parameters", "expected"), + [ + ("SELECT '100%' AS value", (), "100%"), + ("SELECT 10 % 3 AS value", (), 1), + ("SELECT ?::integer % 3 AS value", (10,), 1), + ], +) +async def test_psycopg_cached_percent_expressions( + psycopg_regression_driver: "PsycopgSyncDriver | PsycopgAsyncDriver", + query: str, + parameters: tuple[int, ...], + expected: str | int, +) -> None: + """Native parameter parsing gives equal values before and after a cache hit.""" + for _ in range(2): + result = await _execute(psycopg_regression_driver, query, *parameters) + assert result.get_data() == [{"value": expected}] + + +async def test_psycopg_copy_constraint_error_is_mapped( + psycopg_regression_driver: "PsycopgSyncDriver | PsycopgAsyncDriver", +) -> None: + """The constraint error received when COPY finishes is translated consistently.""" + import psycopg + + pa = pytest.importorskip("pyarrow") + driver = psycopg_regression_driver + await _execute(driver, "CREATE TEMP TABLE copy_constraint_probe (id INTEGER PRIMARY KEY) ON COMMIT DROP") + with pytest.raises(UniqueViolationError) as error: + job = driver.load_from_arrow("copy_constraint_probe", pa.table({"id": [1, 1]})) + if inspect.isawaitable(job): + await job + assert isinstance(error.value.__cause__, psycopg.errors.UniqueViolation) + + +@pytest.mark.parametrize("is_async", [False, True]) +async def test_psycopg_sql_connection_hook_leaves_pool_usable( + is_async: bool, psycopg_sync_config: PsycopgSyncConfig +) -> None: + """SQL issued by a connection hook cannot leave the pool's configure callback in INTRANS.""" + connection_config = {**psycopg_sync_config.connection_config, "min_size": 1, "max_size": 1, "timeout": 2} + if is_async: + + async def async_hook(connection: Any) -> None: + await connection.execute("SET timezone TO 'Pacific/Honolulu'") + + config = PsycopgAsyncConfig( + connection_config=connection_config, driver_features={"on_connection_create": async_hook} + ) + try: + async with config.provide_session() as driver: + assert await driver.select_value("SELECT current_setting('TimeZone')") == "Pacific/Honolulu" + finally: + await config.close_pool() + else: + sync_config = PsycopgSyncConfig( + connection_config=connection_config, + driver_features={ + "on_connection_create": lambda connection: connection.execute("SET timezone TO 'Pacific/Honolulu'") + }, + ) + try: + with sync_config.provide_session() as sync_driver: + assert sync_driver.select_value("SELECT current_setting('TimeZone')") == "Pacific/Honolulu" + finally: + sync_config.close_pool() diff --git a/tests/integration/adapters/sqlanywhere/test_driver.py b/tests/integration/adapters/sqlanywhere/test_driver.py index 8cd171d9a..166443c22 100644 --- a/tests/integration/adapters/sqlanywhere/test_driver.py +++ b/tests/integration/adapters/sqlanywhere/test_driver.py @@ -70,7 +70,7 @@ def test_commit_is_visible_to_new_connection( try: cursor = connection.cursor() cursor.execute(f"SELECT name FROM {items_table} WHERE id = 20") - assert cursor.fetchone()[0] == "kept" + assert cursor.fetchall() == [("kept",)] finally: connection.close() diff --git a/tests/integration/adapters/sqlite/sqlite/test_pooling.py b/tests/integration/adapters/sqlite/sqlite/test_pooling.py index a3c776db6..21e936419 100644 --- a/tests/integration/adapters/sqlite/sqlite/test_pooling.py +++ b/tests/integration/adapters/sqlite/sqlite/test_pooling.py @@ -48,7 +48,7 @@ def test_regular_memory_auto_conversion(sqlite_config_regular_memory: SqliteConf connection_config = build_connection_config(config.connection_config) db_uri = connection_config["database"] - assert db_uri.startswith("file:memory_") and "cache=private" in db_uri + assert db_uri.startswith("file:memory_") and "cache=shared" in db_uri assert connection_config["uri"] is True with config.provide_session() as session1: @@ -248,7 +248,7 @@ def test_config_memory_database_conversion() -> None: try: db_uri = config.connection_config["database"] - assert db_uri.startswith("file:memory_") and "cache=private" in db_uri + assert db_uri.startswith("file:memory_") and "cache=shared" in db_uri assert config.connection_config["uri"] is True with config.provide_session() as session: @@ -267,7 +267,7 @@ def test_config_default_database() -> None: try: db_uri = config.connection_config["database"] - assert db_uri.startswith("file:memory_") and "cache=private" in db_uri + assert db_uri.startswith("file:memory_") and "cache=shared" in db_uri assert config.connection_config["uri"] is True with config.provide_session() as session: diff --git a/tests/unit/adapters/test_adbc/test_adk_dialect_support.py b/tests/unit/adapters/test_adbc/test_adk_dialect_support.py index 506859bcd..34d60bd2f 100644 --- a/tests/unit/adapters/test_adbc/test_adk_dialect_support.py +++ b/tests/unit/adapters/test_adbc/test_adk_dialect_support.py @@ -800,6 +800,7 @@ def record_execute(sql: str, _params: Any = None) -> None: "content_text": "hello", "metadata_json": None, "inserted_at": now, + "embedding": None, }, { "id": "mem-2", @@ -814,6 +815,7 @@ def record_execute(sql: str, _params: Any = None) -> None: "content_text": "world", "metadata_json": None, "inserted_at": now, + "embedding": None, }, ]) diff --git a/tests/unit/adapters/test_adbc/test_core.py b/tests/unit/adapters/test_adbc/test_core.py index e3069d07b..f8ee314cd 100644 --- a/tests/unit/adapters/test_adbc/test_core.py +++ b/tests/unit/adapters/test_adbc/test_core.py @@ -732,6 +732,29 @@ def test_adbc_status_code_mappings(status_code: AdbcStatusCode, expected_cls: ty assert isinstance(mapped, expected_cls) +@pytest.mark.parametrize( + ("status_code", "sqlstate", "message", "expected_cls"), + [ + (AdbcStatusCode.INVALID_ARGUMENT, "23505", "duplicate key", UniqueViolationError), + (AdbcStatusCode.INVALID_ARGUMENT, "23502", "null value", NotNullViolationError), + (AdbcStatusCode.INVALID_ARGUMENT, "23514", "check constraint", CheckViolationError), + (AdbcStatusCode.INTEGRITY, "23503", "foreign key constraint", ForeignKeyViolationError), + (AdbcStatusCode.INVALID_ARGUMENT, "42601", "invalid SQL", SQLParsingError), + (AdbcStatusCode.INVALID_ARGUMENT, None, "syntax error near SELEC", SQLParsingError), + (AdbcStatusCode.INTERNAL, None, "ORA-00001: unique constraint violated", UniqueViolationError), + ], +) +def test_adbc_specific_diagnostics_override_generic_status( + status_code: AdbcStatusCode, sqlstate: str | None, message: str, expected_cls: type[SQLSpecError] +) -> None: + error = DatabaseError(message, status_code=status_code, sqlstate=sqlstate) + + mapped = adbc_core.create_mapped_exception(error) + + assert type(mapped) is expected_cls + assert mapped.__cause__ is error + + def test_duckdb_copy_ast_transformer_caching() -> None: """DuckDB statement config omits ast_transformer so compiled SQL can be cached.""" config = get_statement_config("duckdb").replace(enable_caching=True) @@ -802,7 +825,7 @@ def test_session_context_driver() -> None: statement_config=get_statement_config("postgres"), driver_features={}, prepare_driver=lambda d: d, - driver_class=mock_driver_cls, + driver_class=cast("type[AdbcDriver]", mock_driver_cls), ) with ctx as driver: assert driver is driver_instance @@ -855,6 +878,7 @@ def test_select_to_arrow_lowercase_columns() -> None: result = driver.select_to_arrow("SELECT UPPER_COL FROM t", return_format="reader") reader = result.data assert reader.schema.names == ["upper_col"] + assert reader.read_all().to_pylist() == [{"upper_col": 1}, {"upper_col": 2}] def test_adbc_exception_isolation() -> None: diff --git a/tests/unit/adapters/test_adbc/test_uuid_binding.py b/tests/unit/adapters/test_adbc/test_uuid_binding.py index 3d3e18d9c..eb0a6f11d 100644 --- a/tests/unit/adapters/test_adbc/test_uuid_binding.py +++ b/tests/unit/adapters/test_adbc/test_uuid_binding.py @@ -401,8 +401,9 @@ def test_batch_with_inconsistent_row_shapes_is_left_unchanged() -> None: def test_structural_rewrite_cache_is_bounded_and_keyed_by_dialect_and_ordinals() -> None: - rewrite = cast("Any", adbc_core)._rewrite_postgres_uuid_placeholders - rewrite.cache_clear() + rewrite = adbc_core._rewrite_postgres_uuid_placeholders + cache_state = adbc_core._REWRITE_CACHE_STATE + cache_state.clear() first = rewrite("SELECT $1, $2", (1,), (), "postgres") second = rewrite("SELECT $1, $2", (1,), (), "postgres") @@ -412,27 +413,28 @@ def test_structural_rewrite_cache_is_bounded_and_keyed_by_dialect_and_ordinals() assert first == second == ("SELECT CAST($1 AS UUID), $2", (1,), ()) assert third == ("SELECT $1, CAST($2 AS UUID)", (2,), ()) assert fourth == first - assert rewrite.cache_info().hits == 1 - assert rewrite.cache_info().misses == 3 - assert rewrite.cache_info().maxsize == 256 + assert cache_state.hits == 1 + assert cache_state.misses == 3 + assert cache_state.maxsize == 256 def test_structural_rewrite_cache_distinguishes_scalar_from_array_ordinals() -> None: """The same SQL and ordinal produce different rewrites depending on the ordinal's kind.""" - rewrite = cast("Any", adbc_core)._rewrite_postgres_uuid_placeholders - rewrite.cache_clear() + rewrite = adbc_core._rewrite_postgres_uuid_placeholders + cache_state = adbc_core._REWRITE_CACHE_STATE + cache_state.clear() scalar_only = rewrite("SELECT $1", (1,), (), "postgres") array_only = rewrite("SELECT $1", (), (1,), "postgres") assert scalar_only == ("SELECT CAST($1 AS UUID)", (1,), ()) assert array_only == ("SELECT CAST($1 AS UUID[])", (), (1,)) - assert rewrite.cache_info().misses == 2 - assert rewrite.cache_info().hits == 0 + assert cache_state.misses == 2 + assert cache_state.hits == 0 def test_structural_rewrite_uses_ast_and_ignores_literal_and_comment_placeholders() -> None: - rewrite = cast("Any", adbc_core)._rewrite_postgres_uuid_placeholders + rewrite = adbc_core._rewrite_postgres_uuid_placeholders rewritten, effective, effective_arrays = rewrite("SELECT '$1' AS literal, $1 -- $1\n", (1,), (), "postgres") @@ -465,7 +467,7 @@ def test_extension_dialect_operators_survive_uuid_ast_rewrite(dialect: str, sql: def test_structural_rewrite_reports_sqlglot_parse_errors() -> None: - rewrite = cast("Any", adbc_core)._rewrite_postgres_uuid_placeholders + rewrite = adbc_core._rewrite_postgres_uuid_placeholders with pytest.raises(SQLSpecError, match="Failed to parse PostgreSQL ADBC SQL for UUID parameter binding"): rewrite("SELECT (", (1,), (), "postgres") diff --git a/tests/unit/adapters/test_asyncmy/test_config.py b/tests/unit/adapters/test_asyncmy/test_config.py index fd0ab668c..35cdf6547 100644 --- a/tests/unit/adapters/test_asyncmy/test_config.py +++ b/tests/unit/adapters/test_asyncmy/test_config.py @@ -367,6 +367,17 @@ def test_asyncmy_config_capability_flags() -> None: assert AsyncmyConfig.supports_native_parquet_import is False +def test_conversion_capabilities_require_pyarrow(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(AsyncmyConfig, "_dependency_available", staticmethod(lambda _: False)) + + capabilities = AsyncmyConfig().storage_capabilities() + + assert capabilities["arrow_import_enabled"] is False + assert capabilities["arrow_export_enabled"] is False + assert capabilities["parquet_import_enabled"] is False + assert capabilities["parquet_export_enabled"] is False + + async def test_asyncmy_session_context_initialization() -> None: """AsyncmySessionContext enters and exits using AsyncmyDriver without circular imports.""" from unittest.mock import AsyncMock, MagicMock diff --git a/tests/unit/adapters/test_asyncmy/test_load_from_arrow.py b/tests/unit/adapters/test_asyncmy/test_load_from_arrow.py index 477272e68..0f2081cf3 100644 --- a/tests/unit/adapters/test_asyncmy/test_load_from_arrow.py +++ b/tests/unit/adapters/test_asyncmy/test_load_from_arrow.py @@ -11,6 +11,7 @@ from sqlspec.adapters.asyncmy.config import AsyncmyConfig from sqlspec.adapters.asyncmy.driver import AsyncmyDriver +from sqlspec.exceptions import StorageCapabilityError _CAPS: dict[str, Any] = { "arrow_export_enabled": True, @@ -70,9 +71,7 @@ async def test_load_from_arrow_explicit_bulk_disable_uses_executemany(connection config = AsyncmyConfig( connection_config={"local_infile": connection_opt_in}, driver_features={"enable_local_infile_bulk_load": False} ) - driver = AsyncmyDriver( - connection=cast("Any", conn), driver_features={**config.driver_features, "storage_capabilities": _CAPS} - ) + driver = AsyncmyDriver(connection=cast("Any", conn), driver_features=config.driver_features) job = await driver.load_from_arrow("orders", pa.table({"id": [1, 2], "name": ["a", "b"]})) @@ -93,6 +92,19 @@ async def test_load_from_arrow_overwrite_truncates_first() -> None: assert conn._cursor.executemany_calls[0][0].startswith("INSERT INTO") +async def test_config_preserves_explicit_arrow_import_disable() -> None: + conn = _FakeConnection() + config = AsyncmyConfig( + driver_features=cast("Any", {"storage_capabilities": {**_CAPS, "arrow_import_enabled": False}}) + ) + driver = AsyncmyDriver(connection=cast("Any", conn), driver_features=config.driver_features) + + with pytest.raises(StorageCapabilityError, match="Arrow import"): + await driver.load_from_arrow("orders", pa.table({"id": [1]})) + + assert conn._cursor.executemany_calls == [] + + async def test_load_from_storage_reads_parquet_and_delegates(tmp_path: Path) -> None: parquet_path = tmp_path / "data.parquet" pq.write_table(pa.table({"id": [1, 2], "name": ["a", "b"]}), parquet_path) diff --git a/tests/unit/adapters/test_asyncpg/test_transaction_state.py b/tests/unit/adapters/test_asyncpg/test_transaction_state.py index a2eaa4d9b..9af6a2a5c 100644 --- a/tests/unit/adapters/test_asyncpg/test_transaction_state.py +++ b/tests/unit/adapters/test_asyncpg/test_transaction_state.py @@ -1,6 +1,7 @@ """asyncpg transaction control uses the native transaction handle.""" import inspect +from asyncio import CancelledError from collections.abc import AsyncIterator from contextlib import asynccontextmanager from typing import Any, cast @@ -114,7 +115,8 @@ async def test_a_stream_can_open_a_transaction_after_driver_begin() -> None: async def test_stream_source_reuses_outer_transaction_without_savepoint() -> None: """AsyncpgStreamSource should not start a nested savepoint when already in a transaction.""" connection = _FakeConnection() - connection.cursor = AsyncMock(return_value=MagicMock(fetch=AsyncMock(return_value=[]))) # type: ignore[attr-defined] + cursor = MagicMock(_portal_name="portal", _close_portal=AsyncMock(), fetch=AsyncMock(return_value=[])) + connection.cursor = AsyncMock(return_value=cursor) # type: ignore[attr-defined] driver = AsyncpgDriver(cast("Any", connection)) await driver.begin() @@ -123,10 +125,11 @@ async def test_stream_source_reuses_outer_transaction_without_savepoint() -> Non await source.close() assert connection.transaction_calls == ["start"] + cursor._close_portal.assert_awaited_once_with(None) -async def test_stream_source_suppresses_rollback_error_on_cursor_failure() -> None: - """AsyncpgStreamSource should suppress rollback errors when cursor creation fails.""" +async def test_stream_start_failure_retires_connection_when_rollback_fails() -> None: + """A failed startup rollback retires the socket while preserving the cursor failure.""" tx = MagicMock() tx.start = AsyncMock() tx.rollback = AsyncMock(side_effect=RuntimeError("connection lost")) @@ -134,12 +137,53 @@ async def test_stream_source_suppresses_rollback_error_on_cursor_failure() -> No conn.transaction = MagicMock(return_value=tx) conn.is_in_transaction = MagicMock(return_value=False) conn.cursor = AsyncMock(side_effect=ValueError("cursor failed")) + conn.close = AsyncMock() driver = AsyncpgDriver(cast("Any", conn)) source = AsyncpgStreamSource(driver, "SELECT 1", (), 10) with pytest.raises(ValueError, match="cursor failed"): await source.start() assert source._transaction is None # pyright: ignore[reportPrivateUsage] + conn.close.assert_awaited_once() + + +@pytest.mark.parametrize("owns_transaction", [False, True]) +@pytest.mark.parametrize("error", [False, True]) +async def test_stream_closes_portal_before_owned_transaction(owns_transaction: bool, error: bool) -> None: + """The native portal closes before transaction control without ending a caller transaction.""" + calls: list[str] = [] + connection = MagicMock() + driver = AsyncpgDriver(cast("Any", connection)) + source = AsyncpgStreamSource(driver, "SELECT 1", (), 10, 4.0) + cursor = MagicMock() + cursor._portal_name = "portal" + cursor._close_portal = AsyncMock(side_effect=lambda _: calls.append("portal")) + source._cursor = cursor # pyright: ignore[reportPrivateUsage] + if owns_transaction: + source._transaction = MagicMock( # pyright: ignore[reportPrivateUsage] + commit=AsyncMock(side_effect=lambda: calls.append("commit")), + rollback=AsyncMock(side_effect=lambda: calls.append("rollback")), + ) + + await source.close(error=error) + + assert calls == (["portal", "rollback" if error else "commit"] if owns_transaction else ["portal"]) + cursor._close_portal.assert_awaited_once_with(4.0) + + +@pytest.mark.parametrize("failure", [RuntimeError("portal close failed"), CancelledError()]) +async def test_stream_cleanup_failure_retires_connection(failure: BaseException) -> None: + """Failed or cancelled portal cleanup cannot return a dirty socket to the pool.""" + connection = MagicMock(close=AsyncMock()) + driver = AsyncpgDriver(cast("Any", connection)) + source = AsyncpgStreamSource(driver, "SELECT 1", (), 10) + source._cursor = MagicMock( # pyright: ignore[reportPrivateUsage] + _portal_name="portal", _close_portal=AsyncMock(side_effect=failure) + ) + + with pytest.raises(type(failure)): + await source.close(error=True) + connection.close.assert_awaited_once() async def test_begin_resets_stale_transaction_and_maps_exceptions() -> None: @@ -330,19 +374,45 @@ async def test_asyncpg_execute_stack_in_outer_transaction_does_not_commit_or_rol """execute_stack with continue_on_error inside an outer transaction must not commit or rollback.""" conn = MagicMock() conn.is_in_transaction = MagicMock(return_value=True) - conn.execute = AsyncMock(side_effect=["UPDATE 1", RuntimeError("failed op")]) + + async def execute(sql: str) -> str: + if sql == "UPDATE t SET x = 2": + raise RuntimeError("failed op") + return "UPDATE 1" if sql.startswith("UPDATE") else "" + + conn.execute = AsyncMock(side_effect=execute) driver = AsyncpgDriver(cast("Any", conn)) driver._transaction = MagicMock() - stack = StatementStack().push_execute("UPDATE t SET x = 1").push_execute("UPDATE t SET x = 2") + stack = ( + StatementStack() + .push_execute("UPDATE t SET x = 1") + .push_execute("UPDATE t SET x = 2") + .push_execute("UPDATE t SET x = 3") + ) results = await driver.execute_stack(stack, continue_on_error=True) - assert len(results) == 2 + assert len(results) == 3 assert results[0].result is not None assert results[1].error is not None + assert results[2].rows_affected == 1 assert driver._transaction is not None - conn.execute.assert_called() + commands = [call.args[0] for call in conn.execute.await_args_list] + savepoints = [command.removeprefix("SAVEPOINT ") for command in commands if command.startswith("SAVEPOINT ")] + assert len(savepoints) == 3 + assert commands == [ + f"SAVEPOINT {savepoints[0]}", + "UPDATE t SET x = 1", + f"RELEASE SAVEPOINT {savepoints[0]}", + f"SAVEPOINT {savepoints[1]}", + "UPDATE t SET x = 2", + f"ROLLBACK TO SAVEPOINT {savepoints[1]}", + f"RELEASE SAVEPOINT {savepoints[1]}", + f"SAVEPOINT {savepoints[2]}", + "UPDATE t SET x = 3", + f"RELEASE SAVEPOINT {savepoints[2]}", + ] async def test_asyncpg_exception_mapping_timeout_and_interface_errors() -> None: diff --git a/tests/unit/adapters/test_asyncpg/test_type_handlers.py b/tests/unit/adapters/test_asyncpg/test_type_handlers.py index a40e89f05..79e24e2a4 100644 --- a/tests/unit/adapters/test_asyncpg/test_type_handlers.py +++ b/tests/unit/adapters/test_asyncpg/test_type_handlers.py @@ -48,17 +48,19 @@ async def test_register_json_codecs_success() -> None: assert json_call.kwargs["encoder"]({"value": 1}) == b'{"value":1}' assert json_call.kwargs["encoder"]("console") == b'"console"' - assert json_call.kwargs["encoder"]("123") == b'"123"' - assert json_call.kwargs["encoder"]("true") == b'"true"' - assert json_call.kwargs["encoder"]("null") == b'"null"' + assert json_call.kwargs["encoder"]("123") == b"123" + assert json_call.kwargs["encoder"]("true") == b"true" + assert json_call.kwargs["encoder"]("null") == b"null" + assert json_call.kwargs["encoder"]('"text"') == b'"text"' assert json_call.kwargs["encoder"]('{"pre":"serialized"}') == b'{"pre":"serialized"}' assert json_call.kwargs["encoder"](b'\x01{"value":1}') == b'{"value":1}' assert json_call.kwargs["decoder"](b'{"value":1}') == {"value": 1} assert jsonb_call.kwargs["encoder"]({"value": 1}) == b'\x01{"value":1}' assert jsonb_call.kwargs["encoder"]("console") == b'\x01"console"' - assert jsonb_call.kwargs["encoder"]("123") == b'\x01"123"' - assert jsonb_call.kwargs["encoder"]("true") == b'\x01"true"' - assert jsonb_call.kwargs["encoder"]("null") == b'\x01"null"' + assert jsonb_call.kwargs["encoder"]("123") == b"\x01123" + assert jsonb_call.kwargs["encoder"]("true") == b"\x01true" + assert jsonb_call.kwargs["encoder"]("null") == b"\x01null" + assert jsonb_call.kwargs["encoder"]('"text"') == b'\x01"text"' assert jsonb_call.kwargs["encoder"]('{"pre":"serialized"}') == b'\x01{"pre":"serialized"}' assert jsonb_call.kwargs["encoder"](b'\x01{"value":1}') == b'\x01{"value":1}' assert jsonb_call.kwargs["decoder"](b'\x01{"value":1}') == {"value": 1} diff --git a/tests/unit/adapters/test_cockroach_psycopg/test_config.py b/tests/unit/adapters/test_cockroach_psycopg/test_config.py index 7a17d647c..c3f621ef7 100644 --- a/tests/unit/adapters/test_cockroach_psycopg/test_config.py +++ b/tests/unit/adapters/test_cockroach_psycopg/test_config.py @@ -943,13 +943,11 @@ async def async_user_cb(conn: Any) -> None: def test_cockroach_psycopg_adk_hash_shard_bucket_count_validation() -> None: """Invalid hash_shard_bucket_count must raise ValueError in cockroach_psycopg adk store.""" - from sqlspec.adapters.cockroach_psycopg.adk.store import _cockroach_hash_shard_clause + from sqlspec.adapters.cockroach_psycopg.adk.store import CockroachPsycopgADKConfig, _cockroach_hash_shard_clause - bad_config = CockroachPsycopgSyncConfig( - extension_config={"adk": {"enable_hash_sharded_indexes": True, "hash_shard_bucket_count": -1}} - ) + bad_config: CockroachPsycopgADKConfig = {"enable_hash_sharded_indexes": True, "hash_shard_bucket_count": -1} with pytest.raises(ValueError, match="hash_shard_bucket_count"): - _cockroach_hash_shard_clause(bad_config.extension_config["adk"]) + _cockroach_hash_shard_clause(bad_config) async def test_cockroach_psycopg_adk_memory_search_with_embedding() -> None: diff --git a/tests/unit/adapters/test_db2/_fakes.py b/tests/unit/adapters/test_db2/_fakes.py index 7607bf1f8..ec57278e2 100644 --- a/tests/unit/adapters/test_db2/_fakes.py +++ b/tests/unit/adapters/test_db2/_fakes.py @@ -475,6 +475,7 @@ def __init__( self.connection = FakeDb2Connection(list(cursors), autocommit=True) self.extension_config: dict[str, Any] = extension_config or {} self.statement_config = default_statement_config + self.driver_features: dict[str, Any] = {} @property def executed(self) -> "list[tuple[str, Any]]": @@ -488,7 +489,9 @@ class FakeDb2SessionConfig(_FakeDb2SessionState): @contextmanager def provide_session(self, **_: Any) -> "Iterator[Db2SyncDriver]": """Yield a driver bound to the shared fake connection.""" - yield Db2SyncDriver(self.connection, statement_config=default_statement_config) + yield Db2SyncDriver( + self.connection, statement_config=default_statement_config, driver_features=self.driver_features + ) class FakeDb2AsyncSessionConfig(_FakeDb2SessionState): @@ -502,7 +505,9 @@ class FakeDb2AsyncSessionConfig(_FakeDb2SessionState): async def provide_session(self, **_: Any) -> "AsyncIterator[Any]": """Yield an async driver bound to the shared fake connection.""" yield db2_driver.Db2AsyncDriver( - FakeDb2AsyncConnection(self.connection), statement_config=default_statement_config + FakeDb2AsyncConnection(self.connection), + statement_config=default_statement_config, + driver_features=self.driver_features, ) diff --git a/tests/unit/adapters/test_db2/test_adk.py b/tests/unit/adapters/test_db2/test_adk.py index fb54196a8..ede427fc7 100644 --- a/tests/unit/adapters/test_db2/test_adk.py +++ b/tests/unit/adapters/test_db2/test_adk.py @@ -421,7 +421,7 @@ async def test_insert_memory_entries_binds_naive_utc_and_skips_duplicates(db2_mo """Memory inserts bind naive-UTC times, include the owner column, and skip known event ids.""" store, config = _memory_store( db2_mode, - FakeDb2Cursor(rows=[("m0",)], description=db2_description("id")), + FakeDb2Cursor(rows=[("m0", "e0")], description=db2_description("id", "event_id")), adk={"owner_id_column": "tenant_id INTEGER"}, ) stamp = datetime(2026, 9, 24, 14, tzinfo=timezone(timedelta(hours=2))) @@ -645,14 +645,58 @@ async def test_schema_qualified_adk_tables_sanitize_constraint_and_index_names(d assert all("." not in index_name for index_name, _, _ in store._index_specs()) +@pytest.mark.parametrize( + ("table", "schema"), [("MYSCHEMA.SESSIONS", "MYSCHEMA"), ('"App".sessions', "App"), ('"A""B".sessions', 'A"B')] +) +async def test_adk_index_catalog_probes_preserve_schema(db2_mode: DriverMode, table: str, schema: str) -> None: + """Index existence checks use the schema and index casing emitted by DDL.""" + store, config = _session_store(db2_mode) + store._session_table = table + + await db2_mode.call(store.create_tables) + + index_probes = [params for sql, params in config.executed if sql == INDEX_EXISTS_SQL] + assert index_probes[:2] == [(schema, name.upper()) for name, _, _ in store._index_specs()[:2]] + + +@pytest.mark.parametrize("lowercase_columns", [True, False]) +async def test_memory_replay_deduplicates_event_ids_with_native_casing( + db2_mode: DriverMode, lowercase_columns: bool +) -> None: + """A replayed event is skipped even when the returned memory ID differs.""" + store, config = _memory_store( + db2_mode, FakeDb2Cursor(rows=[("old-memory", "same-event")], description=db2_description("id", "event_id")) + ) + config.driver_features["enable_lowercase_column_names"] = lowercase_columns + stamp = datetime.now(UTC) + entry = StoredMemory( + id="new-memory", + session_id="s1", + app_name="app", + user_id="u1", + scope="user", + event_id="same-event", + author=None, + timestamp=stamp, + content_json={}, + content_text="text", + metadata_json=None, + inserted_at=stamp, + embedding=None, + ) + + assert await db2_mode.call(store.insert_memory_entries, [entry]) == 0 + assert len(config.executed) == 1 + + def test_create_index_sql_qualifies_schema() -> None: """ADK index DDL qualifies index name when table is schema-qualified.""" assert ( _create_index_sql('"MYSCHEMA"."SESSIONS"', "ix_sessions", "user_id") - == 'CREATE INDEX "MYSCHEMA"."ix_sessions" ON "MYSCHEMA"."SESSIONS" (user_id)' + == 'CREATE INDEX "MYSCHEMA"."IX_SESSIONS" ON "MYSCHEMA"."SESSIONS" (user_id)' ) assert ( _create_index_sql("MYSCHEMA.SESSIONS", "ix_sessions", "user_id") - == 'CREATE INDEX "MYSCHEMA"."ix_sessions" ON MYSCHEMA.SESSIONS (user_id)' + == 'CREATE INDEX "MYSCHEMA"."IX_SESSIONS" ON MYSCHEMA.SESSIONS (user_id)' ) assert _create_index_sql("sessions", "ix_sessions", "user_id") == "CREATE INDEX ix_sessions ON sessions (user_id)" diff --git a/tests/unit/adapters/test_duckdb/test_adk_store.py b/tests/unit/adapters/test_duckdb/test_adk_store.py index f14d367d0..3a88a4df0 100644 --- a/tests/unit/adapters/test_duckdb/test_adk_store.py +++ b/tests/unit/adapters/test_duckdb/test_adk_store.py @@ -6,6 +6,7 @@ from typing import Any, NotRequired, Self, cast, get_args, get_origin from unittest.mock import MagicMock +import duckdb import pytest import sqlspec.adapters.duckdb.adk as duckdb_adk @@ -321,6 +322,27 @@ def test_search_entries_embedding_warning() -> None: store.search_entries("query", "app", "user", embedding=[0.1, 0.2, 0.3]) +def test_schema_qualified_adk_tables_create_in_the_configured_schema() -> None: + """Qualified session and memory DDL creates tables and indexes in DuckDB.""" + session_store = DuckDBADKStore(_mock_config({"session_table": "custom.sessions"})) + memory_store = DuckDBADKMemoryStore(_mock_config({"memory_table": "custom.memories"})) + + with duckdb.connect() as connection: + connection.execute("CREATE SCHEMA custom") + connection.execute(session_store._sessions_table_ddl()) + connection.execute(memory_store._memory_table_ddl()) + + assert connection.execute( + "SELECT table_name FROM information_schema.tables WHERE table_schema = ? ORDER BY table_name", ["custom"] + ).fetchall() == [("memories",), ("sessions",)] + + +def test_duckdb_adk_rejects_catalog_qualified_table_names() -> None: + """The ADK schema helpers support schema.table, without a catalog component.""" + with pytest.raises(ValueError, match="Invalid table name"): + DuckDBADKMemoryStore(_mock_config({"memory_table": "catalog.custom.memories"})) + + def test_schema_qualified_fts_table(monkeypatch: pytest.MonkeyPatch) -> None: """Schema-qualified memory tables format indexes without dots and macro as fts_{schema}_{table}.""" store = DuckDBADKMemoryStore(_mock_config({"memory_table": "custom.my_memories"})) diff --git a/tests/unit/adapters/test_duckdb/test_driver.py b/tests/unit/adapters/test_duckdb/test_driver.py index 65d5b8e01..42f287af7 100644 --- a/tests/unit/adapters/test_duckdb/test_driver.py +++ b/tests/unit/adapters/test_duckdb/test_driver.py @@ -244,6 +244,25 @@ def test_select_stream_hugeint_interval_divergent_types() -> None: assert chunks[0]["val"] == 170141183460469231731687303715884105727 +@pytest.mark.parametrize("exhaust_stream", [True, False]) +def test_select_stream_preserves_active_transaction(exhaust_stream: bool) -> None: + with duckdb.connect(":memory:") as connection: + driver = DuckDBDriver(connection) + driver.execute("CREATE TABLE items (id INTEGER)") + driver.begin() + driver.execute("INSERT INTO items VALUES (1)") + + with driver.select_stream("SELECT id::HUGEINT AS id FROM items", chunk_size=1) as stream: + if exhaust_stream: + assert list(stream) == [{"id": 1}] + else: + assert next(stream) == {"id": 1} + + assert driver.select_value("SELECT COUNT(*) FROM items") == 1 + driver.rollback() + assert driver.select_value("SELECT COUNT(*) FROM items") == 0 + + def test_concurrent_execute_many_ast_mutation() -> None: """dispatch_execute_many does not mutate parsed_expression.meta.""" config = DuckDBConfig(connection_config={"database": ":memory:"}) diff --git a/tests/unit/adapters/test_mssql_python/test_arrow.py b/tests/unit/adapters/test_mssql_python/test_arrow.py index 7a0f24a26..2bd38aaa5 100644 --- a/tests/unit/adapters/test_mssql_python/test_arrow.py +++ b/tests/unit/adapters/test_mssql_python/test_arrow.py @@ -1,7 +1,8 @@ """Unit tests for mssql_python Arrow and BulkCopy driver methods.""" -from collections.abc import Iterable -from typing import TYPE_CHECKING, cast +from collections.abc import Iterable, Iterator +from types import SimpleNamespace +from typing import TYPE_CHECKING, Any, cast import pytest @@ -44,7 +45,15 @@ def arrow_reader(self, batch_size: int = 8192) -> object: import pyarrow as pa table = pa.table({"x": [1, 2, 3]}) - return pa.RecordBatchReader.from_batches(table.schema, table.to_batches(max_chunksize=batch_size)) + + def batches() -> Iterator[Any]: + for batch in table.to_batches(max_chunksize=batch_size): + if self.closed: + raise mssql_python.Error("cursor closed before Arrow batch was read", "") + yield batch + + reader = pa.RecordBatchReader.from_batches(table.schema, batches()) + return SimpleNamespace(schema=reader.schema, read_next_batch=reader.read_next_batch, close=reader.close) def bulkcopy(self, target_table: str, rows: Iterable[object], **kwargs: object) -> dict[str, object]: copied_rows = list(rows) diff --git a/tests/unit/adapters/test_mysql_parameter_profiles.py b/tests/unit/adapters/test_mysql_parameter_profiles.py index b6dcc1a5d..a562716e7 100644 --- a/tests/unit/adapters/test_mysql_parameter_profiles.py +++ b/tests/unit/adapters/test_mysql_parameter_profiles.py @@ -42,10 +42,12 @@ def test_literal_percent_is_escaped_only_when_parameters_are_bound(adapter: str, "insert into t (a) values (%s) on duplicate key update a = a % 2", ), ("update t set a = %s where n like 'a%'", "update t set a = %s where n like 'a%%'"), + ("update t set a = %s where a % 2 = 0", "update t set a = %s where a %% 2 = 0"), ], ) +@pytest.mark.parametrize("parameter_sets", [[(1,)], [(1,), (2,)]]) async def test_execute_many_percent_escaping_respects_bulk_insert_shape( - adapter: str, class_name: str, sql: str, expected: str + adapter: str, class_name: str, sql: str, expected: str, parameter_sets: list[tuple[int]] ) -> None: from unittest.mock import AsyncMock, Mock @@ -54,7 +56,7 @@ async def test_execute_many_percent_escaping_respects_bulk_insert_shape( cursor = Mock(rowcount=1) if adapter != "pymysql": cursor.executemany = AsyncMock() - statement = SQL(sql, [(1,)], statement_config=driver.statement_config, is_many=True) + statement = SQL(sql, parameter_sets, statement_config=driver.statement_config, is_many=True) result = driver.dispatch_execute_many(cursor, statement) if adapter != "pymysql": await result diff --git a/tests/unit/adapters/test_mysqlconnector/test_adk_store.py b/tests/unit/adapters/test_mysqlconnector/test_adk_store.py index bfc594674..f6f7b3b54 100644 --- a/tests/unit/adapters/test_mysqlconnector/test_adk_store.py +++ b/tests/unit/adapters/test_mysqlconnector/test_adk_store.py @@ -2,7 +2,7 @@ """Unit tests for mysql-connector ADK store extension configuration.""" import asyncio -from datetime import UTC +from datetime import UTC, datetime from typing import Any, NotRequired, Self, cast, get_args, get_origin from unittest.mock import MagicMock @@ -16,6 +16,7 @@ MysqlConnectorSyncADKStore, ) from sqlspec.config import ADKConfig +from sqlspec.extensions.adk import StoredEvent def _mock_config(adk_config: dict[str, object] | None = None) -> MagicMock: @@ -493,15 +494,13 @@ def test_mysqlconnector_sync_append_event_rolls_back_if_commit_fails() -> None: config.provide_connection = lambda *_a, **_k: conn store = MysqlConnectorSyncADKStore(config) - event: dict[str, Any] = { + event: StoredEvent = { "id": "e-1", "session_id": "sess-1", "app_name": "app", "user_id": "u1", - "event_id": "evt-1", - "event_type": "user", "event_data": {"msg": "hi"}, - "timestamp": None, + "timestamp": datetime.now(UTC), "invocation_id": "inv-1", } diff --git a/tests/unit/adapters/test_mysqlconnector/test_config.py b/tests/unit/adapters/test_mysqlconnector/test_config.py index 38c650cd8..3985aa2dd 100644 --- a/tests/unit/adapters/test_mysqlconnector/test_config.py +++ b/tests/unit/adapters/test_mysqlconnector/test_config.py @@ -1,6 +1,8 @@ """Unit tests for mysql-connector configuration modernization.""" import asyncio +import subprocess +import sys from types import SimpleNamespace from typing import TYPE_CHECKING, Any from unittest.mock import AsyncMock, MagicMock @@ -468,6 +470,36 @@ async def test_async_pool_unavailable_preserves_standalone_connections(monkeypat await MysqlConnectorAsyncConfig(connection_config={"pool_size": 2})._acquire_async_connection() +def test_sync_config_imports_without_native_async_pooling() -> None: + result = subprocess.run( + [ + sys.executable, + "-c", + """ +import importlib.abc +import sys +import mysql.connector.aio + +sys.modules.pop("mysql.connector.aio.pooling", None) + +class MissingAsyncPool(importlib.abc.MetaPathFinder): + def find_spec(self, fullname, path, target=None): + if fullname == "mysql.connector.aio.pooling": + raise ModuleNotFoundError(fullname, name=fullname) + return None + +sys.meta_path.insert(0, MissingAsyncPool()) +from sqlspec.adapters.mysqlconnector import MysqlConnectorSyncConfig +assert MysqlConnectorSyncConfig().connection_config["host"] == "localhost" +""", + ], + capture_output=True, + text=True, + check=False, + ) + assert result.returncode == 0, result.stderr + + def test_sync_create_pool_preserves_zero_pool_size(monkeypatch: pytest.MonkeyPatch) -> None: """MysqlConnectorSyncConfig._create_pool must preserve explicit pool_size=0.""" from sqlspec.adapters.mysqlconnector import config as cfg_module diff --git a/tests/unit/adapters/test_oracledb/test_arrow_batches.py b/tests/unit/adapters/test_oracledb/test_arrow_batches.py index 9038bffa7..7386dc8b1 100644 --- a/tests/unit/adapters/test_oracledb/test_arrow_batches.py +++ b/tests/unit/adapters/test_oracledb/test_arrow_batches.py @@ -214,6 +214,66 @@ def fetch_df_all(self, **_kwargs: object) -> None: driver.select_to_arrow("SELECT * FROM non_existent", return_format="table") +@pytest.mark.anyio +@pytest.mark.parametrize("is_async", [False, True]) +@pytest.mark.parametrize("return_format", ["batches", "reader"]) +async def test_arrow_batch_failure_preserves_database_error(is_async: bool, return_format: Any) -> None: + import oracledb + + from sqlspec.exceptions import SQLParsingError + + class _FailingConnection: + def fetch_df_all(self, **_kwargs: Any) -> None: + raise AssertionError("batch export should use fetch_df_batches") + + def fetch_df_batches(self, **_kwargs: Any) -> Iterator[pa.Table]: + raise oracledb.DatabaseError("ORA-00942: missing table") + yield pa.table({"ID": [1]}) + + class _FailingAsyncConnection: + async def fetch_df_all(self, **_kwargs: Any) -> None: + raise AssertionError("batch export should use fetch_df_batches") + + async def fetch_df_batches(self, **_kwargs: Any) -> AsyncIterator[pa.Table]: + raise oracledb.DatabaseError("ORA-00942: missing table") + yield pa.table({"ID": [1]}) + + with pytest.raises(SQLParsingError, match="ORA-00942") as exc_info: + if is_async: + await OracleAsyncDriver(cast("Any", _FailingAsyncConnection())).select_to_arrow( + "SELECT * FROM missing", return_format=return_format + ) + else: + OracleSyncDriver(cast("Any", _FailingConnection())).select_to_arrow( + "SELECT * FROM missing", return_format=return_format + ) + + assert isinstance(exc_info.value.__cause__, oracledb.DatabaseError) + + +def test_lazy_arrow_reader_maps_errors_after_the_first_batch() -> None: + import oracledb + + from sqlspec.exceptions import SQLParsingError + + class _FailingConnection: + def fetch_df_batches(self, **_kwargs: Any) -> Iterator[pa.Table]: + yield pa.table({"ID": [1]}) + raise oracledb.DatabaseError("ORA-00942: missing table") + + def fetch_df_all(self, **_kwargs: Any) -> None: + raise AssertionError("reader export should use fetch_df_batches") + + result = OracleSyncDriver(cast("Any", _FailingConnection())).select_to_arrow( + "SELECT * FROM example", return_format="reader" + ) + reader = result.get_data() + assert reader.read_next_batch().to_pydict() == {"ID": [1]} + with pytest.raises(SQLParsingError, match="ORA-00942") as exc_info: + reader.read_next_batch() + assert isinstance(exc_info.value.__cause__, oracledb.DatabaseError) + + def test_sync_select_to_arrow_reader_yields_lazily() -> None: """return_format=reader yields record batches lazily from the underlying stream.""" yielded: list[int] = [] diff --git a/tests/unit/adapters/test_oracledb/test_config.py b/tests/unit/adapters/test_oracledb/test_config.py index cbd72213e..0c1b82364 100644 --- a/tests/unit/adapters/test_oracledb/test_config.py +++ b/tests/unit/adapters/test_oracledb/test_config.py @@ -388,6 +388,74 @@ async def test_async_create_connection_unpooled(monkeypatch: pytest.MonkeyPatch) assert config.connection_instance is None +@pytest.mark.anyio +@pytest.mark.parametrize("is_async", [False, True]) +async def test_standalone_connections_filter_pool_options(is_async: bool, monkeypatch: pytest.MonkeyPatch) -> None: + """A pool-configured config can also create a standalone connection.""" + connection_type = ( + oracle_config_module.oracledb.AsyncConnection if is_async else oracle_config_module.oracledb.Connection + ) + connect = AsyncMock() if is_async else Mock() + monkeypatch.setattr(oracle_config_module.oracledb, "connect_async" if is_async else "connect", connect) + config_type = OracleAsyncConfig if is_async else OracleSyncConfig + config = config_type( + connection_config={ + "user": "scott", + "password": "tiger", + "dsn": "localhost/xe", + "min": 1, + "max": 4, + "increment": 1, + "getmode": PoolGetMode.WAIT, + "ping_interval": 60, + "ping_timeout": 5000, + "max_sessions_per_shard": 2, + "pool_class": oracle_config_module.oracledb.ConnectionPool, + "connectiontype": connection_type, + } + ) + + result = config.create_connection() + if isawaitable(result): + result = await result + + assert result is connect.return_value + connect.assert_called_once_with(user="scott", password="tiger", dsn="localhost/xe", conn_class=connection_type) + + +@pytest.mark.parametrize("options", [{"thick_mode": True}, {"lib_dir": "/oracle/lib"}, {"soda_metadata_cache": True}]) +def test_standalone_sync_connection_initializes_thick_mode( + options: dict[str, Any], monkeypatch: pytest.MonkeyPatch +) -> None: + calls: list[str] = [] + initialize = Mock(side_effect=lambda **_kwargs: calls.append("initialize")) + connect = Mock(side_effect=lambda **_kwargs: calls.append("connect") or _StubConnection()) + monkeypatch.setattr(oracle_config_module.oracledb, "is_thin_mode", lambda: True) + monkeypatch.setattr(oracle_config_module.oracledb, "init_oracle_client", initialize) + monkeypatch.setattr(oracle_config_module.oracledb, "connect", connect) + config = OracleSyncConfig(connection_config=options) + + config.create_connection() + + assert calls == ["initialize", "connect"] + initialize.assert_called_once_with(**({"lib_dir": options["lib_dir"]} if "lib_dir" in options else {})) + + +@pytest.mark.anyio +@pytest.mark.parametrize("options", [{"thick_mode": True}, {"lib_dir": "/oracle/lib"}, {"soda_metadata_cache": True}]) +async def test_standalone_async_connection_rejects_thick_mode( + options: dict[str, Any], monkeypatch: pytest.MonkeyPatch +) -> None: + connect = AsyncMock() + monkeypatch.setattr(oracle_config_module.oracledb, "connect_async", connect) + config = OracleAsyncConfig(connection_config=options) + + with pytest.raises(ImproperConfigurationError, match="only supports Thin mode"): + await config.create_connection() + + connect.assert_not_called() + + def test_derive_pool_produces_identical_configurations_across_sync_and_async() -> None: """derive_pool produces identical connection configurations across sync and async configs.""" common_cfg = {"user": "scott", "password": "tiger", "dsn": "localhost/xe", "min": 2, "max": 10} diff --git a/tests/unit/adapters/test_oracledb/test_core_row_materialization.py b/tests/unit/adapters/test_oracledb/test_core_row_materialization.py index f1c0e90c3..164675889 100644 --- a/tests/unit/adapters/test_oracledb/test_core_row_materialization.py +++ b/tests/unit/adapters/test_oracledb/test_core_row_materialization.py @@ -131,7 +131,7 @@ def test_resolve_row_metadata_single_pass_unknown_type_code_shape_is_conservativ def test_resolve_row_metadata_cache_contract_still_holds_after_single_pass_rewrite() -> None: description = [("ID", _TypeCode("DB_TYPE_NUMBER"))] - cache: dict[object, object] = {} + cache: dict[object, tuple[object, list[str], bool]] = {} (first_names, first_requires_lob) = resolve_row_metadata(description, {}, cache) (second_names, second_requires_lob) = resolve_row_metadata(description, {}, cache) assert second_names is first_names diff --git a/tests/unit/adapters/test_oracledb/test_litestar_store.py b/tests/unit/adapters/test_oracledb/test_litestar_store.py index 4bda51d10..6634e4e9f 100644 --- a/tests/unit/adapters/test_oracledb/test_litestar_store.py +++ b/tests/unit/adapters/test_oracledb/test_litestar_store.py @@ -2,12 +2,13 @@ from datetime import UTC, datetime, timedelta from typing import Any, Self, cast +from unittest.mock import AsyncMock, MagicMock import pytest from sqlspec.adapters.oracledb.core import DB_TYPE_BLOB from sqlspec.adapters.oracledb.data_dictionary import OracleVersionCache -from sqlspec.adapters.oracledb.litestar import OracleSyncStore +from sqlspec.adapters.oracledb.litestar import OracleAsyncStore, OracleSyncStore class _FakeCursor: @@ -186,3 +187,30 @@ def test_oracle_store_set_rolls_back_on_execute_failure() -> None: store._set("fail-key", b"data") connection.rollback.assert_called_once() + + +@pytest.mark.anyio +@pytest.mark.parametrize("is_async", [False, True]) +async def test_session_renewal_binds_each_sql_placeholder(is_async: bool) -> None: + config = MagicMock() + config.extension_config = {} + cursor = MagicMock() + cursor.fetchone.return_value = (b"payload", datetime(2026, 5, 10, tzinfo=UTC)) + connection = MagicMock() + connection.cursor.return_value.__enter__.return_value = cursor + if is_async: + cursor.execute = AsyncMock() + cursor.fetchone = AsyncMock(return_value=cursor.fetchone.return_value) + connection.commit = AsyncMock() + config.provide_connection.return_value.__aenter__.return_value = connection + store = OracleAsyncStore(config) + else: + config.provide_connection.return_value.__enter__.return_value = connection + store = OracleSyncStore(config) + + result = await store.get("session-1", renew_for=60) + + assert result == b"payload" + assert cursor.execute.call_count == 2 + assert cursor.execute.call_args.args[1] == (60, 60, "session-1") + connection.commit.assert_called_once() diff --git a/tests/unit/adapters/test_oracledb/test_oracle_adk_store.py b/tests/unit/adapters/test_oracledb/test_oracle_adk_store.py index 1e57fd770..67f45dbf5 100644 --- a/tests/unit/adapters/test_oracledb/test_oracle_adk_store.py +++ b/tests/unit/adapters/test_oracledb/test_oracle_adk_store.py @@ -3,8 +3,10 @@ from datetime import UTC, datetime from decimal import Decimal +from inspect import isawaitable +from types import SimpleNamespace from typing import Any, NotRequired, Self, cast, get_args, get_origin -from unittest.mock import MagicMock +from unittest.mock import AsyncMock, MagicMock from uuid import uuid4 import pytest @@ -22,7 +24,8 @@ from sqlspec.adapters.oracledb.adk.store import _event_data_column_ddl from sqlspec.adapters.oracledb.data_dictionary import OracleVersionCache, OracleVersionInfo from sqlspec.config import ADKConfig -from sqlspec.exceptions import SQLSpecError +from sqlspec.exceptions import ImproperConfigurationError, SQLSpecError +from sqlspec.extensions.adk import StoredMemory def _mock_config(adk_config: dict[str, object]) -> MagicMock: @@ -47,6 +50,177 @@ def _config_for_storage(storage_type: JSONStorageType) -> MagicMock: return config +def _config_for_vectors() -> MagicMock: + config = _config_for_storage(JSONStorageType.JSON_NATIVE) + config._oracle_version_cache.version = OracleVersionInfo(23, 4, 0, compatible="23.4.0") + return config + + +def _memory_entry() -> StoredMemory: + timestamp = datetime(2026, 5, 10, 12, tzinfo=UTC) + return { + "id": "memory-1", + "session_id": "session-1", + "app_name": "app", + "user_id": "user", + "scope": "user", + "event_id": "event-1", + "author": "user", + "timestamp": timestamp, + "content_json": {"text": "hello"}, + "content_text": "hello", + "metadata_json": None, + "inserted_at": timestamp, + "embedding": None, + } + + +@pytest.mark.anyio +@pytest.mark.parametrize("is_async", [False, True]) +@pytest.mark.parametrize("error_code", [1, 1400, 12899, 51803]) +async def test_memory_batch_errors_commit_only_duplicate_skips(is_async: bool, error_code: int) -> None: + import oracledb + + config = _config_for_storage(JSONStorageType.BLOB_JSON) + config.extension_config = {"adk": {}} + cursor = MagicMock() + batch_error = SimpleNamespace(code=error_code, message=f"ORA-{error_code:05d}: rejected memory") + cursor.getbatcherrors.return_value = [batch_error] + connection = MagicMock() + connection.cursor.return_value.__enter__.return_value = cursor + if is_async: + cursor.executemany = AsyncMock() + connection.commit = AsyncMock() + connection.rollback = AsyncMock() + config.provide_connection.return_value.__aenter__.return_value = connection + store = OracleAsyncADKMemoryStore(config) + else: + config.provide_connection.return_value.__enter__.return_value = connection + store = OracleSyncADKMemoryStore(config) + + entries: list[StoredMemory] = [_memory_entry(), {**_memory_entry(), "id": "memory-2"}] + if error_code == 1: + result = store.insert_memory_entries(entries) + if isawaitable(result): + result = await result + assert result == 1 + connection.commit.assert_called_once() + connection.rollback.assert_not_called() + else: + with pytest.raises(oracledb.DatabaseError) as exc_info: + result = store.insert_memory_entries(entries) + if isawaitable(result): + await result + assert exc_info.value.args[0] is batch_error + connection.commit.assert_not_called() + connection.rollback.assert_called_once() + + +@pytest.mark.anyio +@pytest.mark.parametrize("is_async", [False, True]) +@pytest.mark.parametrize("use_fts", [False, True]) +@pytest.mark.parametrize( + "version", + [ + OracleVersionInfo(19, 0, 0, compatible="19.0.0"), + OracleVersionInfo(21, 3, 0, compatible="21.0.0"), + OracleVersionInfo(23, 4, 0, compatible="23.0.0"), + ], +) +async def test_legacy_memory_retains_text_storage_and_rejects_vectors( + is_async: bool, use_fts: bool, version: OracleVersionInfo +) -> None: + config = _config_for_storage(JSONStorageType.BLOB_JSON) + config._oracle_version_cache.version = version + config.extension_config = {"adk": {"memory_use_fts": use_fts}} + cursor = MagicMock() + cursor.getbatcherrors.return_value = [] + timestamp = _memory_entry()["timestamp"] + cursor.fetchall.return_value = [ + ( + "memory-1", + "session-1", + "app", + "user", + "user", + "event-1", + "user", + timestamp, + b'{"text":"hello"}', + "hello", + None, + timestamp, + None, + ) + ] + connection = MagicMock() + connection.cursor.return_value.__enter__.return_value = cursor + executed: list[str] = [] + + def execute(sql: str, _params: Any) -> None: + assert "NULL AS embedding" in sql + executed.append(sql) + + if is_async: + cursor.execute = AsyncMock(side_effect=execute) + cursor.executemany = AsyncMock() + cursor.fetchall = AsyncMock(return_value=cursor.fetchall.return_value) + connection.commit = AsyncMock() + config.provide_connection.return_value.__aenter__.return_value = connection + store = OracleAsyncADKMemoryStore(config) + else: + cursor.execute.side_effect = execute + config.provide_connection.return_value.__enter__.return_value = connection + store = OracleSyncADKMemoryStore(config) + + ddl = store._memory_table_ddl() + if isawaitable(ddl): + ddl = await ddl + from sqlglot import parse_one + from sqlglot.expressions import Create + + create_sql = ddl.split("EXECUTE IMMEDIATE '", 1)[1].split("';", 1)[0].replace("''", "'") + table = parse_one(create_sql, dialect="oracle") + assert isinstance(table, Create) + assert {column.name for column in table.this.expressions} == { + "id", + "session_id", + "app_name", + "user_id", + "scope", + "event_id", + "author", + "timestamp", + "content_json", + "metadata_json", + "content_text", + "inserted_at", + } + + inserted = store.insert_memory_entries([_memory_entry()]) + if isawaitable(inserted): + inserted = await inserted + assert inserted == 1 + insert_sql, bind_rows = cursor.executemany.call_args.args + assert "timestamp, content_json" in insert_sql + assert set(bind_rows[0]) == set(_memory_entry()) - {"embedding"} + + records = store.search_entries("hello", "app", "user") + if isawaitable(records): + records = await records + assert len(executed) == 1 + assert [(record["id"], record["embedding"]) for record in records] == [("memory-1", None)] + + with pytest.raises(ImproperConfigurationError, match=r"23\.4"): + result = store.insert_memory_entries([{**_memory_entry(), "embedding": [0.1]}]) + if isawaitable(result): + await result + with pytest.raises(ImproperConfigurationError, match=r"23\.4"): + result = store.search_entries("hello", "app", "user", embedding=[0.1]) + if isawaitable(result): + await result + + def test_oracle_adk_config_types_adapter_local_optimizations() -> None: """Oracle ADK optimization settings are typed on the adapter-local extension config.""" @@ -302,16 +476,19 @@ def test_oracle_sync_adk_memory_rows_to_records_deserializes_json_fields() -> No assert records[0]["embedding"] is None -def test_oracle_adk_memory_ddl_declares_embedding_vector() -> None: +@pytest.mark.anyio +async def test_oracle_adk_memory_ddl_declares_embedding_vector() -> None: """Memory table DDL declares an embedding VECTOR column with configured dimensions.""" - config = _mock_config({}) + config = _config_for_vectors() + config.extension_config = {"adk": {}} store = OracleSyncADKMemoryStore(config) - sql = store._memory_table_ddl_for_type(JSONStorageType.JSON_NATIVE) + sql = store._memory_table_ddl() assert "embedding VECTOR(768, FLOAT32)" in sql - custom_config = _mock_config({"vector_dimensions": 1536}) + custom_config = _config_for_vectors() + custom_config.extension_config = {"adk": {"vector_dimensions": 1536}} custom_store = OracleAsyncADKMemoryStore(custom_config) - custom_sql = custom_store._memory_table_ddl_for_type(JSONStorageType.JSON_NATIVE) + custom_sql = await custom_store._memory_table_ddl() assert "embedding VECTOR(1536, FLOAT32)" in custom_sql @@ -319,7 +496,7 @@ def test_oracle_adk_sync_insert_binds_embedding_array() -> None: """Sync memory store binds embeddings as float arrays.""" import array - config = _config_for_storage(JSONStorageType.JSON_NATIVE) + config = _config_for_vectors() config.extension_config = {"adk": {}} store = OracleSyncADKMemoryStore(config) conn = MagicMock() @@ -329,7 +506,7 @@ def test_oracle_adk_sync_insert_binds_embedding_array() -> None: config.provide_connection.return_value.__enter__.return_value = conn timestamp = datetime(2026, 5, 10, 12, 0, tzinfo=UTC) - entries = [ + entries: list[StoredMemory] = [ { "id": "mem-1", "session_id": "sess-1", @@ -414,7 +591,7 @@ def test_oracle_adk_memory_store_batch_insert() -> None: store = OracleSyncADKMemoryStore(config) now = datetime(2026, 5, 10, 12, 0, tzinfo=UTC) - entries = [ + entries: list[StoredMemory] = [ { "id": "mem-1", "session_id": "sess-1", @@ -428,6 +605,7 @@ def test_oracle_adk_memory_store_batch_insert() -> None: "content_text": "text 1", "metadata_json": None, "inserted_at": now, + "embedding": None, }, { "id": "mem-2", @@ -442,6 +620,7 @@ def test_oracle_adk_memory_store_batch_insert() -> None: "content_text": "text 2", "metadata_json": None, "inserted_at": now, + "embedding": None, }, ] inserted = store.insert_memory_entries(entries) @@ -518,6 +697,9 @@ def executemany(self, sql: str, params: Any = None, **_kwargs: Any) -> None: for p in params or []: self.calls.append((sql, dict(p or {}))) + def getbatcherrors(self) -> list[Any]: + return [] + def fetchall(self) -> "list[Any]": return [] @@ -783,7 +965,7 @@ def test_oracle_adk_owner_id_column_ddl_escapes_single_quotes() -> None: def test_oracle_sync_adk_memory_fts_sanitizes_query_and_hybrid_rrf_fuses_candidates() -> None: cursor = _RecordingCursor() conn = _RecordingConnection(cursor) - config = _config_for_storage(JSONStorageType.JSON_NATIVE) + config = _config_for_vectors() config.extension_config = {"adk": {"memory_use_fts": True}} config.provide_connection = lambda *_a, **_k: conn store = OracleSyncADKMemoryStore(config) @@ -811,7 +993,7 @@ def test_oracle_sync_adk_memory_fts_sanitizes_query_and_hybrid_rrf_fuses_candida async def test_oracle_async_adk_memory_fts_sanitizes_query_and_hybrid_rrf_fuses_candidates() -> None: cursor = _AsyncRecordingCursor() conn = _RecordingConnection(cursor) - config = _config_for_storage(JSONStorageType.JSON_NATIVE) + config = _config_for_vectors() config.extension_config = {"adk": {"memory_use_fts": True}} config.provide_connection = lambda *_a, **_k: conn store = OracleAsyncADKMemoryStore(config) @@ -829,7 +1011,7 @@ def test_oracle_sync_adk_memory_hybrid_rrf_near_and_exact_fallback() -> None: """Sync ADK memory store forwards fts_operator='NEAR' and approximate=False.""" cursor = _RecordingCursor() conn = _RecordingConnection(cursor) - config = _config_for_storage(JSONStorageType.JSON_NATIVE) + config = _config_for_vectors() config.extension_config = {"adk": {"memory_use_fts": True, "fts_operator": "NEAR", "approximate": False}} config.provide_connection = lambda *_a, **_k: conn store = OracleSyncADKMemoryStore(config) @@ -856,7 +1038,7 @@ async def test_oracle_async_adk_memory_hybrid_rrf_near_and_exact_fallback() -> N """Async ADK memory store forwards fts_operator='NEAR' and approximate=False.""" cursor = _AsyncRecordingCursor() conn = _RecordingConnection(cursor) - config = _config_for_storage(JSONStorageType.JSON_NATIVE) + config = _config_for_vectors() config.extension_config = {"adk": {"memory_use_fts": True, "fts_operator": "NEAR", "approximate": False}} config.provide_connection = lambda *_a, **_k: conn store = OracleAsyncADKMemoryStore(config) diff --git a/tests/unit/adapters/test_oracledb/test_stack_native_transaction.py b/tests/unit/adapters/test_oracledb/test_stack_native_transaction.py index 00cf17737..a03b1fa83 100644 --- a/tests/unit/adapters/test_oracledb/test_stack_native_transaction.py +++ b/tests/unit/adapters/test_oracledb/test_stack_native_transaction.py @@ -14,7 +14,7 @@ class _FakeOpResult: - def __init__(self, error: "Exception | None" = None) -> None: + def __init__(self, error: "Any | None" = None) -> None: self.error = error @@ -115,7 +115,7 @@ async def test_native_stack_wraps_mapped_error() -> None: """Stack execution errors wrap mapped SQLSpec exceptions, not raw driver errors.""" import oracledb - from sqlspec.exceptions import SQLSpecError + from sqlspec.exceptions import SQLSpecError, StackExecutionError native_err = oracledb.DatabaseError("ORA-00001: unique constraint violated") results = [_FakeOpResult(native_err)] @@ -124,5 +124,42 @@ async def test_native_stack_wraps_mapped_error() -> None: stack_results = await driver._execute_stack_native(_stack(), continue_on_error=True) assert len(stack_results) == 1 - assert stack_results[0].error is not None - assert isinstance(stack_results[0].error.original_error, SQLSpecError) + error = stack_results[0].error + assert isinstance(error, StackExecutionError) + assert isinstance(error.original_error, SQLSpecError) + + +@pytest.mark.anyio +@pytest.mark.parametrize("continue_on_error", [False, True]) +async def test_stack_maps_native_pipeline_error_records(continue_on_error: bool) -> None: + import oracledb + from oracledb.errors import _Error + + from sqlspec.exceptions import StackExecutionError, UniqueViolationError + + native_error = _Error("ORA-00001: unique constraint violated", code=1) + connection = _FakePipelineConnection([_FakeOpResult(), _FakeOpResult(native_error), _FakeOpResult()]) + driver = OracleAsyncDriver(cast("Any", connection)) + driver._pipeline_support = True + + if continue_on_error: + results = await driver.execute_stack(_stack(), continue_on_error=True) + assert len(results) == 3 + assert results[0].error is None + assert results[2].error is None + error = results[1].error + assert connection.commit_count == 1 + assert connection.rollback_count == 0 + else: + with pytest.raises(StackExecutionError) as exc_info: + await driver.execute_stack(_stack()) + error = exc_info.value + assert connection.commit_count == 0 + assert connection.rollback_count == 1 + + assert isinstance(error, StackExecutionError) + assert error.operation_index == 1 + assert isinstance(error.original_error, UniqueViolationError) + cause = error.original_error.__cause__ + assert isinstance(cause, oracledb.DatabaseError) + assert cause.args[0] is native_error diff --git a/tests/unit/adapters/test_postgres_migration_schema.py b/tests/unit/adapters/test_postgres_migration_schema.py index ced043c77..d249ad0a1 100644 --- a/tests/unit/adapters/test_postgres_migration_schema.py +++ b/tests/unit/adapters/test_postgres_migration_schema.py @@ -5,6 +5,7 @@ from typing import Any import pytest +from psycopg.sql import Composable from sqlspec.adapters.adbc.config import AdbcConfig from sqlspec.adapters.adbc.driver import AdbcDriver @@ -50,8 +51,8 @@ def __init__(self, schema_exists: bool = True) -> None: self.schema_exists = schema_exists self.executed: list[tuple[str, Any]] = [] - def execute(self, sql: str, parameters: Any | None = None) -> None: - self.executed.append((sql, parameters)) + def execute(self, sql: str | Composable, parameters: Any | None = None) -> None: + self.executed.append((sql.as_string() if isinstance(sql, Composable) else sql, parameters)) def fetchone(self) -> tuple[int] | None: return (1,) if self.schema_exists else None @@ -73,8 +74,8 @@ def __init__(self, schema_exists: bool = True) -> None: self.schema_exists = schema_exists self.executed: list[tuple[str, Any]] = [] - async def execute(self, sql: str, parameters: Any | None = None) -> None: - self.executed.append((sql, parameters)) + async def execute(self, sql: str | Composable, parameters: Any | None = None) -> None: + self.executed.append((sql.as_string() if isinstance(sql, Composable) else sql, parameters)) async def fetchone(self) -> tuple[int] | None: return (1,) if self.schema_exists else None diff --git a/tests/unit/adapters/test_psqlpy/test_adk_store.py b/tests/unit/adapters/test_psqlpy/test_adk_store.py index b7ba74ca7..d5e757629 100644 --- a/tests/unit/adapters/test_psqlpy/test_adk_store.py +++ b/tests/unit/adapters/test_psqlpy/test_adk_store.py @@ -297,6 +297,6 @@ async def test_psqlpy_memory_store_normalizes_tuple_embeddings_on_insert_and_sea results = await store.search_entries("", "app", "user-1", embedding=(0.1, 0.2)) assert inserted == 1 - assert "unnest" in conn.fetch_calls[0][0].lower() + assert conn.fetch_calls[0][1][-1] == ["[0.1,0.2]"] assert conn.fetch_calls[1][1][0] == [0.1, 0.2] assert results[0]["embedding"] == [0.1, 0.2] diff --git a/tests/unit/adapters/test_psqlpy/test_driver.py b/tests/unit/adapters/test_psqlpy/test_driver.py index d16b302cd..bddd8683a 100644 --- a/tests/unit/adapters/test_psqlpy/test_driver.py +++ b/tests/unit/adapters/test_psqlpy/test_driver.py @@ -11,7 +11,7 @@ from sqlspec.adapters.psqlpy.driver import PsqlpyDriver from sqlspec.core import SQL from sqlspec.core.sqlcommenter import SQLCommenterContext -from sqlspec.exceptions import SQLSpecError +from sqlspec.exceptions import SQLSpecError, UniqueViolationError @pytest.mark.anyio @@ -137,6 +137,35 @@ async def test_psqlpy_execute_cache_hit_direct_fast_path() -> None: assert conn.fetch.await_count == 1 +@pytest.mark.parametrize("query", ["UPDATE items SET qty = 1", "UPDATE items SET qty = ?"]) +async def test_psqlpy_cached_dml_preserves_rowcount(query: str) -> None: + """DML without RETURNING retains its affected-row count after the query is cached.""" + native_result = SimpleNamespace(result=lambda: [{psqlpy_core._DML_COUNT_COLUMN: 3}]) + connection = SimpleNamespace(fetch=AsyncMock(return_value=native_result), execute=AsyncMock()) + driver = PsqlpyDriver(connection=cast("Any", connection)) + parameters = (1,) if "?" in query else () + + results = [await driver.execute(query, *parameters) for _ in range(2)] + + assert [result.rows_affected for result in results] == [3, 3] + assert connection.fetch.await_count == 2 + + +async def test_psqlpy_cached_query_maps_native_errors() -> None: + """A warmed cache retains the same exception translation and cause as ordinary execution.""" + from psqlpy.exceptions import ConnectionExecuteError + + native_result = SimpleNamespace(result=lambda: [{"value": 1}]) + native_error = ConnectionExecuteError("duplicate key value violates unique constraint items_pk") + connection = SimpleNamespace(fetch=AsyncMock(side_effect=[native_result, native_error])) + driver = PsqlpyDriver(connection=cast("Any", connection)) + + assert (await driver.execute("SELECT 1 AS value")).get_data() == [{"value": 1}] + with pytest.raises(UniqueViolationError) as error: + await driver.execute("SELECT 1 AS value") + assert error.value.__cause__ is native_error + + @pytest.mark.anyio async def test_psqlpy_load_from_arrow_preserves_cause_on_failure() -> None: """load_from_arrow preserves underlying __cause__ on execution failure.""" diff --git a/tests/unit/adapters/test_psycopg/test_config.py b/tests/unit/adapters/test_psycopg/test_config.py index 121198e5d..4c592726f 100644 --- a/tests/unit/adapters/test_psycopg/test_config.py +++ b/tests/unit/adapters/test_psycopg/test_config.py @@ -629,13 +629,14 @@ def test_async_config_normalizes_aliases() -> None: assert "db" not in config.connection_config -def test_psycopg_sync_configure_connection_supports_dict_row_and_rolls_back_after_hook() -> None: - """_configure_connection should handle dict_row probe results and roll back before user hooks when autocommit=False.""" +def test_psycopg_sync_configure_connection_commits_hook_after_extension_probe() -> None: + """The probe and SQL hook both leave a non-autocommit connection idle.""" events: list[str] = [] connection = MagicMock() connection.autocommit = False connection.execute.return_value.fetchall.return_value = [{"extname": "pg_textsearch"}] connection.rollback.side_effect = lambda: events.append("rollback") + connection.commit.side_effect = lambda: events.append("commit") config = PsycopgSyncConfig( driver_features={ @@ -650,12 +651,12 @@ def test_psycopg_sync_configure_connection_supports_dict_row_and_rolls_back_afte assert config.pg_textsearch_available is True assert config.statement_config.dialect == "pg_textsearch" - assert events == ["rollback", "hook"] + assert events == ["rollback", "hook", "commit"] @pytest.mark.anyio -async def test_psycopg_async_configure_connection_supports_dict_row_and_rolls_back_after_hook() -> None: - """_configure_async_connection should handle dict_row probe results and roll back before user hooks when autocommit=False.""" +async def test_psycopg_async_configure_connection_commits_hook_after_extension_probe() -> None: + """The probe and SQL hook both leave a non-autocommit connection idle.""" events: list[str] = [] connection = AsyncMock() connection.autocommit = False @@ -670,6 +671,7 @@ async def record_hook(_: object) -> None: events.append("hook") connection.rollback = AsyncMock(side_effect=record_rollback) + connection.commit = AsyncMock(side_effect=lambda: events.append("commit")) config = PsycopgAsyncConfig( driver_features={ @@ -684,7 +686,57 @@ async def record_hook(_: object) -> None: assert config.pg_textsearch_available is True assert config.statement_config.dialect == "pg_textsearch" - assert events == ["rollback", "hook"] + assert events == ["rollback", "hook", "commit"] + + +@pytest.mark.parametrize("is_async", [False, True]) +@pytest.mark.parametrize("failure_phase", ["hook", "commit"]) +async def test_failed_connection_initialization_rolls_back(is_async: bool, failure_phase: str) -> None: + """A failed hook or setup commit rolls its transaction back before propagating the error.""" + connection = AsyncMock() if is_async else MagicMock() + connection.autocommit = False + native_error = RuntimeError("initialization failed") + hook = AsyncMock() if is_async else MagicMock() + if failure_phase == "hook": + hook.side_effect = native_error + else: + connection.commit.side_effect = native_error + config = ( + PsycopgAsyncConfig(driver_features={"on_connection_create": hook}) + if is_async + else PsycopgSyncConfig(driver_features={"on_connection_create": hook}) + ) + config._pgvector_available = False # pyright: ignore[reportPrivateUsage] + + with pytest.raises(RuntimeError, match="initialization failed"): + if isinstance(config, PsycopgAsyncConfig): + await config._configure_async_connection(connection) # pyright: ignore[reportPrivateUsage] + else: + config._configure_connection(connection) # pyright: ignore[reportPrivateUsage] + assert connection.rollback.call_count == 1 + + +@pytest.mark.parametrize("is_async", [False, True]) +async def test_direct_config_preserves_connection_initialization_hook(is_async: bool) -> None: + """Bypassing a pooler keeps the configured physical-connection initialization.""" + hook = AsyncMock() if is_async else MagicMock() + config = ( + PsycopgAsyncConfig( + driver_features={"direct_dsn": "postgresql://localhost/direct", "on_connection_create": hook} + ) + if is_async + else PsycopgSyncConfig( + driver_features={"direct_dsn": "postgresql://localhost/direct", "on_connection_create": hook} + ) + ) + direct = config.derive_direct_config() + connection = AsyncMock() if is_async else MagicMock() + connection.autocommit = True + if isinstance(direct, PsycopgAsyncConfig): + await direct._configure_async_connection(connection) # pyright: ignore[reportPrivateUsage] + else: + direct._configure_connection(connection) # pyright: ignore[reportPrivateUsage] + hook.assert_called_once_with(connection) def test_psycopg_pooler_mode_transaction_sets_prepare_threshold_none_and_alloydb_class( diff --git a/tests/unit/adapters/test_psycopg/test_driver.py b/tests/unit/adapters/test_psycopg/test_driver.py index 238233462..b4620935b 100644 --- a/tests/unit/adapters/test_psycopg/test_driver.py +++ b/tests/unit/adapters/test_psycopg/test_driver.py @@ -1,5 +1,6 @@ """Unit tests for psycopg driver transaction behavior.""" +import inspect from asyncio import CancelledError from contextlib import nullcontext from types import SimpleNamespace @@ -12,7 +13,7 @@ from sqlspec.adapters.psycopg.config import PsycopgAsyncConfig, PsycopgSyncConfig from sqlspec.adapters.psycopg.driver import PsycopgAsyncDriver, PsycopgSyncDriver from sqlspec.core import SQL -from sqlspec.exceptions import SerializationConflictError, SQLSpecError +from sqlspec.exceptions import SerializationConflictError, SQLSpecError, UniqueViolationError if TYPE_CHECKING: from sqlspec.adapters.psycopg._typing import PsycopgAsyncConnection, PsycopgSyncConnection @@ -271,75 +272,41 @@ async def test_async_rollback_restores_changed_autocommit() -> None: assert connection.autocommit is True -def test_sync_cache_hit_with_literal_percent_avoids_fallback() -> None: - """_execute_cache_hit for query with literal percent uses pre-escaped wire SQL.""" - from sqlspec.core.cache import CachedQuery - from sqlspec.core.compiler import OperationProfile - from sqlspec.core.parameters import ParameterInfo, ParameterProfile, ParameterStyle - from sqlspec.core.statement import ProcessedState - - connection = MagicMock() +@pytest.mark.parametrize("is_async", [False, True]) +@pytest.mark.parametrize( + ("query", "parameters", "wire_sql"), + [ + ("SELECT '100%' AS value", (), "SELECT '100%' AS value"), + ("SELECT 10 % 3 AS value", (), "SELECT 10 % 3 AS value"), + ("SELECT %s AS value WHERE '100%' LIKE '100%'", (1,), "SELECT %s AS value WHERE '100%%' LIKE '100%%'"), + ], +) +async def test_cached_percent_sql_preserves_native_binding( + is_async: bool, query: str, parameters: tuple[int, ...], wire_sql: str +) -> None: + """Cached and uncached execution apply percent escaping only with native parameters.""" cursor = MagicMock() - cursor.description = [SimpleNamespace(name="id")] - cursor.fetchall.return_value = [(1,)] + cursor.description = [SimpleNamespace(name="value", type_code=23)] cursor.rowcount = 1 - connection.cursor.return_value = cursor - driver = PsycopgSyncDriver(connection=connection) - - cached = CachedQuery( - compiled_sql="SELECT * FROM items WHERE name LIKE '%%test%%' AND id = %s", - parameter_profile=ParameterProfile([ParameterInfo("id", ParameterStyle.POSITIONAL_PYFORMAT, 0, 1, "%s")]), - input_named_parameters=(), - applied_wrap_types=False, - parameter_casts={}, - operation_type="SELECT", - operation_profile=OperationProfile(returns_rows=True, modifies_rows=False), - param_count=1, - processed_state=ProcessedState( - compiled_sql="SELECT * FROM items WHERE name LIKE '%%test%%' AND id = %s", - execution_parameters=[], - operation_type="SELECT", - ), - ) - result = driver._execute_cache_hit("SELECT * FROM items WHERE name LIKE '%test%' AND id = %s", (1,), cached) - assert result is not None - assert cursor.execute.called - - -@pytest.mark.anyio -async def test_async_cache_hit_with_literal_percent_avoids_fallback() -> None: - """async _execute_cache_hit for query with literal percent uses pre-escaped wire SQL.""" - from sqlspec.core.cache import CachedQuery - from sqlspec.core.compiler import OperationProfile - from sqlspec.core.parameters import ParameterInfo, ParameterProfile, ParameterStyle - from sqlspec.core.statement import ProcessedState - + if is_async: + cursor.execute = AsyncMock() + cursor.fetchall = AsyncMock(return_value=[(1,)]) + cursor.close = AsyncMock() + else: + cursor.fetchall.return_value = [(1,)] connection = MagicMock() - cursor = AsyncMock() - cursor.description = [SimpleNamespace(name="id")] - cursor.fetchall = AsyncMock(return_value=[(1,)]) - cursor.rowcount = 1 connection.cursor.return_value = cursor - driver = PsycopgAsyncDriver(connection=connection) + driver = PsycopgAsyncDriver(connection) if is_async else PsycopgSyncDriver(connection) - cached = CachedQuery( - compiled_sql="SELECT * FROM items WHERE name LIKE '%%test%%' AND id = %s", - parameter_profile=ParameterProfile([ParameterInfo("id", ParameterStyle.POSITIONAL_PYFORMAT, 0, 1, "%s")]), - input_named_parameters=(), - applied_wrap_types=False, - parameter_casts={}, - operation_type="SELECT", - operation_profile=OperationProfile(returns_rows=True, modifies_rows=False), - param_count=1, - processed_state=ProcessedState( - compiled_sql="SELECT * FROM items WHERE name LIKE '%%test%%' AND id = %s", - execution_parameters=[], - operation_type="SELECT", - ), - ) - result = await driver._execute_cache_hit("SELECT * FROM items WHERE name LIKE '%test%' AND id = %s", (1,), cached) - assert result is not None - assert cursor.execute.called + for _ in range(2): + result = driver.execute(query, *parameters) + if inspect.isawaitable(result): + result = await result + assert result.get_data() == [{"value": 1}] + + assert [call.args[0] for call in cursor.execute.call_args_list] == [wire_sql, wire_sql] + bound = [tuple(call.args[1]) if len(call.args) > 1 else () for call in cursor.execute.call_args_list] + assert bound == [parameters, parameters] def test_psycopg_sync_stream_source_runs_with_exception_handler() -> None: @@ -599,6 +566,45 @@ async def test_psycopg_async_load_from_arrow_runs_with_exception_handler() -> No assert job.telemetry["rows_processed"] == 2 +@pytest.mark.parametrize("is_async", [False, True]) +@pytest.mark.parametrize("binary", [False, True]) +@pytest.mark.parametrize("failure_phase", ["enter", "exit"]) +async def test_copy_lifecycle_maps_native_errors(is_async: bool, binary: bool, failure_phase: str) -> None: + """COPY setup and finalization errors retain SQLSpec mapping even inside an except block.""" + import pyarrow as pa + + native_error = psycopg.errors.UniqueViolation("duplicate key value violates unique constraint items_pk") + copy_context = AsyncMock() if is_async else MagicMock() + enter = copy_context.__aenter__ if is_async else copy_context.__enter__ + exit_ = copy_context.__aexit__ if is_async else copy_context.__exit__ + enter.return_value = copy_context + if failure_phase == "enter": + enter.side_effect = native_error + else: + exit_.side_effect = native_error + cursor = MagicMock() + cursor.copy.return_value = copy_context + if is_async: + cursor.close = AsyncMock() + connection = MagicMock() + connection.cursor.return_value = cursor + features = {"storage_capabilities": {"arrow_import_enabled": True}, "enable_pgcopy_binary": binary} + driver = ( + PsycopgAsyncDriver(connection, driver_features=features) + if is_async + else PsycopgSyncDriver(connection, driver_features=features) + ) + + try: + raise ValueError("ambient exception") + except ValueError: + with pytest.raises(UniqueViolationError) as error: + result = driver.load_from_arrow("items", pa.table({"id": [1, 1]})) + if inspect.isawaitable(result): + await result + assert error.value.__cause__ is native_error + + def test_sync_dispatch_execute_script_emits_statement_observer_events() -> None: """PsycopgSyncDriver.dispatch_execute_script emits script statement events to observers.""" from sqlspec.core import SQL diff --git a/tests/unit/adapters/test_pymssql/test_data_dictionary.py b/tests/unit/adapters/test_pymssql/test_data_dictionary.py index 147a64ad6..e0ee85f11 100644 --- a/tests/unit/adapters/test_pymssql/test_data_dictionary.py +++ b/tests/unit/adapters/test_pymssql/test_data_dictionary.py @@ -153,15 +153,3 @@ def test_data_dictionary_capabilities() -> None: deps_result = data_dictionary.get_dependencies(cast(Any, driver), object_name="orders", schema="sales") assert deps_result.domain == "dependencies" - - -def test_mssql_catalog_cte_aggregations() -> None: - """Constraints, indexes, and DDL templates use STRING_AGG CTE aggregations without FOR XML PATH.""" - data_dictionary = PymssqlSyncDataDictionary() - constraints_sql = data_dictionary.get_query_text("constraints", "by_schema") - assert "STRING_AGG" in constraints_sql - assert "FOR XML PATH" not in constraints_sql - - indexes_sql = data_dictionary.get_query_text("indexes", "by_schema") - assert "STRING_AGG" in indexes_sql - assert "FOR XML PATH" not in indexes_sql diff --git a/tests/unit/adapters/test_spanner/test_config.py b/tests/unit/adapters/test_spanner/test_config.py index c545be198..06f036518 100644 --- a/tests/unit/adapters/test_spanner/test_config.py +++ b/tests/unit/adapters/test_spanner/test_config.py @@ -75,6 +75,32 @@ def test_config_defaults() -> None: assert config.connection_config["size"] == 10 +@pytest.mark.parametrize("config_type", [SpannerSyncConfig, SpannerAsyncConfig]) +def test_arrow_conversion_storage_capabilities(config_type: type[SpannerSyncConfig] | type[SpannerAsyncConfig]) -> None: + config = config_type(connection_config={"project": "p", "instance_id": "i", "database_id": "d"}) + + capabilities = config.storage_capabilities() + + assert capabilities["arrow_import_enabled"] is True + assert capabilities["arrow_export_enabled"] is True + assert config.supports_native_arrow_import is False + assert config.supports_native_arrow_export is False + assert config.driver_features["storage_capabilities"] == capabilities + + +@pytest.mark.parametrize("config_type", [SpannerSyncConfig, SpannerAsyncConfig]) +def test_arrow_conversion_storage_preserves_explicit_disable( + config_type: type[SpannerSyncConfig] | type[SpannerAsyncConfig], +) -> None: + capabilities = {"arrow_import_enabled": False, "arrow_export_enabled": False} + config = config_type( + connection_config={"project": "p", "instance_id": "i", "database_id": "d"}, + driver_features={"storage_capabilities": capabilities}, + ) + + assert config.driver_features["storage_capabilities"] == capabilities + + def test_min_sessions_is_rejected() -> None: """Spanner's current session pool classes do not support min_sessions.""" with pytest.raises(ImproperConfigurationError, match="min_sessions"): diff --git a/tests/unit/adapters/test_spanner/test_load_from_arrow_mutations.py b/tests/unit/adapters/test_spanner/test_load_from_arrow_mutations.py index c562f0102..26e1bee0a 100644 --- a/tests/unit/adapters/test_spanner/test_load_from_arrow_mutations.py +++ b/tests/unit/adapters/test_spanner/test_load_from_arrow_mutations.py @@ -6,16 +6,9 @@ import pytest import sqlspec.adapters.spanner.driver as spanner_driver +from sqlspec.adapters.spanner.config import SpannerSyncConfig from sqlspec.adapters.spanner.driver import SpannerSyncDriver -CAPABILITIES = { - "arrow_export_enabled": True, - "arrow_import_enabled": True, - "parquet_export_enabled": True, - "parquet_import_enabled": True, - "partition_strategies": ["fixed"], -} - class _FakeTransaction: def __init__(self) -> None: @@ -34,7 +27,8 @@ def execute_update(self, sql: str, params: Any = None, param_types: Any = None, @pytest.fixture def mutations_driver(monkeypatch: pytest.MonkeyPatch) -> SpannerSyncDriver: monkeypatch.setattr(spanner_driver, "SpannerTransaction", _FakeTransaction) - return SpannerSyncDriver(cast("Any", _FakeTransaction()), driver_features={"storage_capabilities": CAPABILITIES}) + config = SpannerSyncConfig(connection_config={"project": "p", "instance_id": "i", "database_id": "d"}) + return SpannerSyncDriver(cast("Any", _FakeTransaction()), driver_features=config.driver_features) def test_load_from_arrow_uses_single_insert_or_update(mutations_driver: SpannerSyncDriver) -> None: diff --git a/tests/unit/adapters/test_sqlanywhere.py b/tests/unit/adapters/test_sqlanywhere.py index 536c3543c..c2093cb38 100644 --- a/tests/unit/adapters/test_sqlanywhere.py +++ b/tests/unit/adapters/test_sqlanywhere.py @@ -485,14 +485,8 @@ def test_sqlanywhere_adk_store_handles_native_missing_table(monkeypatch: pytest. config = SqlanywhereConfig(connection_config={"user": "dba", "password": "sql", "database": "demo"}) store = SqlanywhereADKStore(config) - class FakeSqlanydbError(SqlanywhereError): - def __init__(self, sqlstate: str, sqlcode: int, message: str) -> None: - super().__init__(message) - self.sqlstate = sqlstate - self.sqlcode = sqlcode - def _table_missing(_self: Any, *_args: Any, **_kwargs: Any) -> None: - raise FakeSqlanydbError("42S02", -141, "Table 'adk_session' not found") + raise SqlanywhereError("Table 'adk_session' not found", -141) monkeypatch.setattr(SqlanywhereADKStore, "_execute_fetchone", _table_missing) assert store.get_session("app", "user", "sess") is None @@ -550,14 +544,8 @@ def test_sqlanywhere_adk_memory_store_handles_native_missing_table(monkeypatch: config = SqlanywhereConfig(connection_config={"user": "dba", "password": "sql", "database": "demo"}) store = SqlanywhereADKMemoryStore(config) - class FakeSqlanydbError(SqlanywhereError): - def __init__(self, sqlstate: str, sqlcode: int, message: str) -> None: - super().__init__(message) - self.sqlstate = sqlstate - self.sqlcode = sqlcode - def _table_missing(_self: Any, *_args: Any, **_kwargs: Any) -> None: - raise FakeSqlanydbError("42S02", -141, "Table 'adk_memory' not found") + raise SqlanywhereError("Table 'adk_memory' not found", -141) monkeypatch.setattr(SqlanywhereADKMemoryStore, "_execute_fetchall", _table_missing) assert store.search_entries("query", "app", "user") == [] @@ -623,8 +611,7 @@ def test_sqlanywhere_sync_data_dictionary_instantiation() -> None: def test_sqlanywhere_has_schema_under_active_exception(monkeypatch: pytest.MonkeyPatch) -> None: """has_schema runs cleanly inside an active exception block without isolation leak.""" - config = SqlanywhereConfig(connection_config={"user": "dba", "password": "sql", "database": "demo"}) - driver = SqlanywhereDriver(config) + driver = SqlanywhereDriver(cast("SqlanywhereConnection", FakeConnection())) class DummyCursor: def execute(self, _sql: str, _params: Any = ()) -> None: @@ -658,8 +645,7 @@ def test_sqlanywhere_format_identifier_and_insert_statement() -> None: def test_sqlanywhere_commit_and_rollback_raise_transaction_error() -> None: """commit() and rollback() wrap caught exceptions in TransactionError.""" - config = SqlanywhereConfig(connection_config={"user": "dba", "password": "sql", "database": "demo"}) - driver = SqlanywhereDriver(config) + driver = SqlanywhereDriver(cast("SqlanywhereConnection", FakeConnection())) class FailingConnection: def commit(self) -> None: @@ -675,34 +661,20 @@ def rollback(self) -> None: driver.rollback() -def test_sqlanywhere_session_context_committed_does_not_rollback() -> None: - """An explicitly committed session context does not issue rollback() on exit.""" - rollback_called = False - - class TrackingConnection: - def commit(self) -> None: - pass - - def rollback(self) -> None: - nonlocal rollback_called - rollback_called = True - - def close(self) -> None: - pass - - conn = TrackingConnection() +def test_sqlanywhere_session_context_rolls_back_implicit_work() -> None: + """DML without begin() must be rolled back before a pooled connection is reused.""" + conn = FakeConnection() ctx = SqlanywhereSessionContext( - acquire_connection=lambda: conn, + acquire_connection=lambda: cast("SqlanywhereConnection", conn), release_connection=lambda _c, **_kw: None, statement_config=default_statement_config, driver_features={}, prepare_driver=lambda d: d, ) with ctx as driver: - driver.begin() - driver.commit() + driver.execute("INSERT INTO items (id) VALUES (1)") - assert rollback_called is False + assert conn.rollbacks == 1 def test_sqlanywhere_connection_pool_bounds_max_size(monkeypatch: pytest.MonkeyPatch) -> None: diff --git a/tests/unit/adapters/test_sqlite/test_pool.py b/tests/unit/adapters/test_sqlite/test_pool.py index 2c58bb17c..735966e76 100644 --- a/tests/unit/adapters/test_sqlite/test_pool.py +++ b/tests/unit/adapters/test_sqlite/test_pool.py @@ -46,6 +46,25 @@ def test_default_pool_uses_native_like_shared_pragma_profile() -> None: pool.close() +def test_default_memory_connections_do_not_read_uncommitted_rows() -> None: + from sqlspec.adapters.sqlite import SqliteConfig + + config = SqliteConfig() + other = config.provide_pool().new_connection() + try: + with config.provide_connection() as connection: + connection.execute("CREATE TABLE isolated_rows (value INTEGER)") + connection.execute("BEGIN") + connection.execute("INSERT INTO isolated_rows VALUES (1)") + with pytest.raises(sqlite3.OperationalError, match="locked"): + other.execute("SELECT value FROM isolated_rows") + connection.rollback() + assert other.execute("SELECT COUNT(*) FROM isolated_rows").fetchone() == (0,) + finally: + other.close() + config.close_pool() + + def test_disable_optimizations_preserves_native_memory_profile() -> None: native_connection = sqlite3.connect(":memory:", timeout=30.0) pool = SqliteConnectionPool({"database": ":memory:", "timeout": 30.0}, enable_optimizations=False) diff --git a/tests/unit/data_dictionary/test_mysql_metadata.py b/tests/unit/data_dictionary/test_mysql_metadata.py index f1a01bffc..6c4bb6405 100644 --- a/tests/unit/data_dictionary/test_mysql_metadata.py +++ b/tests/unit/data_dictionary/test_mysql_metadata.py @@ -1,7 +1,9 @@ """Unit tests for MySQL-family data-dictionary metadata.""" from importlib import import_module +from inspect import isawaitable from typing import Any, cast +from unittest.mock import AsyncMock, Mock import pytest import sqlglot @@ -11,6 +13,7 @@ from sqlspec.adapters.asyncmy.data_dictionary import AsyncmyDataDictionary from sqlspec.adapters.mysqlconnector.data_dictionary import MysqlConnectorSyncDataDictionary from sqlspec.adapters.pymysql.data_dictionary import PyMysqlDataDictionary +from sqlspec.core import SQL from sqlspec.data_dictionary import ( DataDictionaryLoader, DDLResult, @@ -28,6 +31,8 @@ class _SyncDDLDriver: def __init__(self) -> None: self.statements: list[str] = [] + self.sql_mode = "STRICT_TRANS_TABLES,NO_ENGINE_SUBSTITUTION" + self.sql_quote_show_create = "1" def select_one(self, statement: str, *args: Any, **kwargs: Any) -> dict[str, str]: self.statements.append(str(statement)) @@ -35,8 +40,8 @@ def select_one(self, statement: str, *args: Any, **kwargs: Any) -> dict[str, str if "SELECT VERSION()" in query: return { "raw_version": "8.4.0", - "sql_mode": "STRICT_TRANS_TABLES,NO_ENGINE_SUBSTITUTION", - "sql_quote_show_create": "1", + "sql_mode": self.sql_mode, + "sql_quote_show_create": self.sql_quote_show_create, } return {"Table": "orders", "Create Table": "CREATE TABLE `orders` (`id` int)"} @@ -268,29 +273,45 @@ def select(self, statement: str, *args: Any, **kwargs: Any) -> list[dict[str, ob assert not any(statement.startswith("SHOW CREATE") for statement in synth_driver.statements) -def test_mysql_native_ddl_single_roundtrip_setup_and_cache() -> None: - """_get_native_object_ddl combines scalar queries into one round-trip and caches per session.""" - driver = _SyncDDLDriver() - dictionary = PyMysqlDataDictionary() - - result1 = dictionary.get_ddl(cast(Any, driver), "orders", schema="shop") - assert isinstance(result1, DDLResult) - setup_queries_count = sum(1 for s in driver.statements if "SELECT VERSION()" in s) - assert setup_queries_count == 1 - - result2 = dictionary.get_ddl(cast(Any, driver), "orders", schema="shop") - assert isinstance(result2, DDLResult) - setup_queries_count_after = sum(1 for s in driver.statements if "SELECT VERSION()" in s) - assert setup_queries_count_after == 1 - - -def test_mysql_get_privileges_filters_by_object_name() -> None: - """get_privileges passes object_name to domain query.""" - driver = _SyncDDLDriver() - dictionary = PyMysqlDataDictionary() - - dictionary.get_privileges(cast(Any, driver), object_name="orders", schema="shop") - privilege_queries = [ - s for s in driver.statements if "information_schema.table_privileges" in s or "table_schema" in s - ] - assert len(privilege_queries) > 0 +@pytest.mark.parametrize("dictionary_type", [PyMysqlDataDictionary, AiomysqlDataDictionary]) +async def test_mysql_native_ddl_refreshes_session_context(dictionary_type: type[Any]) -> None: + """DDL captures mutable SQL mode and quote settings for every statement.""" + session = _SyncDDLDriver() + driver = ( + Mock(select_one=AsyncMock(side_effect=session.select_one)) + if dictionary_type is AiomysqlDataDictionary + else session + ) + dictionary = dictionary_type() + for sql_mode, quotes in [(session.sql_mode, "1"), ("NO_BACKSLASH_ESCAPES", "0")]: + session.sql_mode = sql_mode + session.sql_quote_show_create = quotes + result = dictionary.get_ddl(driver, "orders", schema="shop") + if isawaitable(result): + result = await result + assert result.context["sql_mode"] == sql_mode + assert result.context["sql_quote_show_create"] == quotes + + +@pytest.mark.parametrize("dictionary_type", [PyMysqlDataDictionary, AiomysqlDataDictionary]) +@pytest.mark.parametrize("object_name", [None, "orders"]) +async def test_mysql_get_privileges_binds_optional_object_name( + dictionary_type: type[Any], object_name: str | None +) -> None: + from sqlspec.adapters.pymysql.core import default_statement_config + + def select_rows(statement: str, **parameters: Any) -> list[dict[str, str]]: + compiled_sql, bound = SQL(str(statement), statement_config=default_statement_config, **parameters).compile() + assert "information_schema.table_privileges" in compiled_sql + assert bound == {"schema_name": "shop", "object_name": object_name} + return [{"table_name": "orders"}] + + is_async = dictionary_type is AiomysqlDataDictionary + driver = Mock( + select=AsyncMock(side_effect=select_rows) if is_async else Mock(side_effect=select_rows), + select_value_or_none=AsyncMock(return_value="8.4.0") if is_async else Mock(return_value="8.4.0"), + ) + result = dictionary_type().get_privileges(driver, object_name=object_name, schema="shop") + if isawaitable(result): + result = await result + assert result.items == ({"table_name": "orders"},) diff --git a/tests/unit/data_dictionary/test_sqlite_duckdb_metadata.py b/tests/unit/data_dictionary/test_sqlite_duckdb_metadata.py index df8fcaecf..39a7af6b7 100644 --- a/tests/unit/data_dictionary/test_sqlite_duckdb_metadata.py +++ b/tests/unit/data_dictionary/test_sqlite_duckdb_metadata.py @@ -489,6 +489,30 @@ def test_duckdb_metadata_capabilities_profile() -> None: driver.connection.close() +@pytest.mark.parametrize(("table", "expected_names"), [("target", ["target"]), ("missing", [])]) +def test_duckdb_table_details_filters_requested_table(table: str, expected_names: list[str]) -> None: + config = DuckDBConfig(connection_config={"database": ":memory:"}) + driver = DuckDBDriver(config.create_connection()) + try: + driver.execute_script(""" + CREATE SCHEMA app; + CREATE TABLE app.target (id INTEGER); + CREATE TABLE app.other (id INTEGER); + CREATE TABLE target (id INTEGER); + """) + + details = driver.data_dictionary.get_table_details(driver, table=table, schema="app") + + assert {row["table_name"] for row in driver.data_dictionary.get_tables(driver, schema="app")} == { + "target", + "other", + } + assert details.capability.support == MetadataSupport.SUPPORTED + assert [cast("dict[str, object]", item)["table_name"] for item in details.items] == expected_names + finally: + driver.connection.close() + + def test_duckdb_unwired_sql_domains() -> None: """DuckDBDataDictionary wires schemas, objects, routines, databases, extensions, system, and table_details.""" config = DuckDBConfig(connection_config={"database": ":memory:"}) @@ -503,7 +527,7 @@ def test_duckdb_unwired_sql_domains() -> None: objects = driver.data_dictionary.get_objects(driver, schema="custom_schema") assert objects.domain == "objects" - assert any(obj.get("object_name") == "test_unwired" for obj in objects.items) + assert any(cast("dict[str, object]", obj).get("object_name") == "test_unwired" for obj in objects.items) routines = driver.data_dictionary.get_routines(driver, schema="custom_schema") assert routines.domain == "routines" @@ -519,7 +543,7 @@ def test_duckdb_unwired_sql_domains() -> None: table_details = driver.data_dictionary.get_table_details(driver, table="test_unwired", schema="custom_schema") assert table_details.domain == "tables" assert len(table_details.items) == 1 - assert table_details.items[0]["table_name"] == "test_unwired" + assert cast("dict[str, object]", table_details.items[0])["table_name"] == "test_unwired" system_settings = driver.data_dictionary.get_system_metadata(driver, "settings", include_system=True) assert system_settings.request.domain == "settings" diff --git a/tests/unit/extensions/test_adk/test_store_instantiation.py b/tests/unit/extensions/test_adk/test_store_instantiation.py index ee39d7a5b..d7891588e 100644 --- a/tests/unit/extensions/test_adk/test_store_instantiation.py +++ b/tests/unit/extensions/test_adk/test_store_instantiation.py @@ -15,6 +15,7 @@ import pytest +from sqlspec.core import StatementConfig from sqlspec.exceptions import SQLSpecError from sqlspec.extensions.adk.memory.store import BaseAsyncADKMemoryStore, BaseSyncADKMemoryStore from sqlspec.extensions.adk.store import BaseAsyncADKStore, BaseSyncADKStore @@ -192,6 +193,56 @@ def test_adk_store_registration_validator_resolves_duckdb_store_classes() -> Non _ensure_adk_store_registration(config) +@pytest.mark.parametrize( + ("adapter", "prefix", "dialect"), + [ + ("db2", "Db2Sync", "db2"), + ("db2", "Db2Async", "db2"), + ("arrow_odbc", "ArrowOdbc", "tsql"), + ("arrow_odbc", "ArrowOdbc", "db2"), + ("cockroach_asyncpg", "CockroachAsyncpg", "postgres"), + ("cockroach_psycopg", "CockroachPsycopgSync", "postgres"), + ("cockroach_psycopg", "CockroachPsycopgAsync", "postgres"), + ("duckdb", "DuckDB", "duckdb"), + ("mssql_python", "MssqlPython", "tsql"), + ("pymssql", "Pymssql", "tsql"), + ("sybase", "Sybase", "sybase"), + ("sqlanywhere", "Sqlanywhere", "sqlanywhere"), + ], +) +def test_qualified_table_names_reach_supporting_adk_stores(adapter: str, prefix: str, dialect: str) -> None: + """Supported stores accept qualified tables through their public configuration.""" + session_tables = { + "session_table": "app.sessions", + "events_table": "app.events", + "app_state_table": "app.app_state", + "user_state_table": "app.user_state", + "metadata_table": "app.metadata", + } + config = _load_class(f"sqlspec.adapters.{adapter}.{prefix}Config")( + extension_config={"adk": {**session_tables, "memory_table": "app.memories"}}, + statement_config=StatementConfig(dialect=dialect), + ) + + session_store = _load_class(f"sqlspec.adapters.{adapter}.adk.{prefix}ADKStore")(config) + memory_store = _load_class(f"sqlspec.adapters.{adapter}.adk.{prefix}ADKMemoryStore")(config) + + assert {name: getattr(session_store, name) for name in session_tables} == session_tables + assert memory_store.memory_table == "app.memories" + + +@pytest.mark.parametrize("table", ["dataset.sessions", "project.dataset.sessions"]) +def test_bigquery_adk_preserves_qualified_table_configuration(table: str) -> None: + """BigQuery keeps dataset and project qualification supplied by the caller.""" + from sqlspec.adapters.bigquery import BigQueryConfig + from sqlspec.adapters.bigquery.adk import BigQueryADKStore + + store = BigQueryADKStore(BigQueryConfig(extension_config={"adk": {"session_table": table, "enable_memory": False}})) + + assert store.session_table == table + assert f"CREATE TABLE IF NOT EXISTS {table}" in store._sessions_table_ddl() + + def test_adk_store_registration_validator_fails_fast_for_broken_store_mapping(monkeypatch: pytest.MonkeyPatch) -> None: """Broken adapter-to-store naming fails during config registration.""" from sqlspec.adapters.sqlite import SqliteConfig diff --git a/tests/unit/extensions/test_extension_config_boundaries.py b/tests/unit/extensions/test_extension_config_boundaries.py index d0fed0643..e4281c2c3 100644 --- a/tests/unit/extensions/test_extension_config_boundaries.py +++ b/tests/unit/extensions/test_extension_config_boundaries.py @@ -1,16 +1,11 @@ """Runtime boundaries for shared and adapter-specific extension settings.""" -import re -from pathlib import Path - import pytest -import sqlspec.extensions from sqlspec.adapters.sqlite.adk.store import SqliteADKMemoryStore, SqliteADKStore from sqlspec.adapters.sqlite.config import SqliteConfig from sqlspec.adapters.sqlite.events.store import SqliteEventQueueStore from sqlspec.adapters.sqlite.litestar.store import SqliteStore -from sqlspec.config import FastAPIConfig, FlaskConfig, LitestarConfig, SanicConfig, StarletteConfig from sqlspec.exceptions import ImproperConfigurationError @@ -42,22 +37,13 @@ def test_sqlite_adk_rejects_postgres_memory_tuning(option: str, value: object) - SqliteADKMemoryStore(config) -_FRAMEWORK_CONFIG_READS = ( - ("flask", "flask_config", FlaskConfig), - ("starlette", "framework_config", StarletteConfig), - ("fastapi", "framework_config", FastAPIConfig), - ("sanic", "framework_config", SanicConfig), - ("litestar", "litestar_config", LitestarConfig), +@pytest.mark.parametrize( + "option", ["session_table", "events_table", "app_state_table", "user_state_table", "metadata_table", "memory_table"] ) +def test_sqlite_adk_rejects_qualified_table_names(option: str) -> None: + """Unsupported qualified names fail before any schema creation can partially succeed.""" + config = SqliteConfig(extension_config={"adk": {option: "main.custom_table"}}) + store_type = SqliteADKMemoryStore if option == "memory_table" else SqliteADKStore - -@pytest.mark.parametrize(("framework", "variable", "typed_dict"), _FRAMEWORK_CONFIG_READS) -def test_framework_extension_reads_only_declared_config_keys(framework: str, variable: str, typed_dict: type) -> None: - """Every key an extension reads from its framework config is declared on that framework's TypedDict.""" - package = Path(sqlspec.extensions.__file__).parent / framework - source = "".join(path.read_text() for path in sorted(package.rglob("*.py"))) - read_keys = set(re.findall(rf"\b{variable}\.get\(\"(\w+)\"", source)) - declared = set(typed_dict.__required_keys__) | set(typed_dict.__optional_keys__) - - assert read_keys, f"no {variable}.get(...) reads found for {framework}" - assert read_keys <= declared, f"{framework} reads undeclared keys: {sorted(read_keys - declared)}" + with pytest.raises(ValueError, match="Invalid table name"): + store_type(config) diff --git a/tests/unit/test_import_time.py b/tests/unit/test_import_time.py index 50114d3d7..4ba4446f6 100644 --- a/tests/unit/test_import_time.py +++ b/tests/unit/test_import_time.py @@ -334,57 +334,3 @@ def test_adapter_import_does_not_load_secondary_optional_modules(adapter: str) - result = subprocess.run([sys.executable, "-c", script], capture_output=True, text=True, check=True) leaked = [name for name in result.stdout.strip().split(",") if name] assert not leaked, f"import sqlspec.adapters.{adapter} eagerly imported: {leaked}" - - -def _is_type_checking(test: ast.expr) -> bool: - return (isinstance(test, ast.Name) and test.id == "TYPE_CHECKING") or ( - isinstance(test, ast.Attribute) and test.attr == "TYPE_CHECKING" - ) - - -def _imports_driver(module_name: str, drivers: tuple[str, ...]) -> bool: - return any(module_name == name or module_name.startswith(name + ".") for name in drivers) - - -def test_adapter_modules_follow_import_contract() -> None: - """Adapter modules import their driver directly and defer secondary optional modules to first use.""" - violations: list[str] = [] - for adapter, drivers in sorted(ADAPTER_DRIVERS.items()): - for py_file in sorted(Path("sqlspec/adapters", adapter).rglob("*.py")): - tree = ast.parse(py_file.read_text(encoding="utf-8"), filename=str(py_file)) - function_lines: set[int] = set() - for node in ast.walk(tree): - if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): - function_lines.update(range(node.lineno, (node.end_lineno or node.lineno) + 1)) - for node in ast.walk(tree): - location = f"{py_file}:{getattr(node, 'lineno', '?')}" - if isinstance(node, ast.If) and isinstance(node.test, ast.UnaryOp): - if isinstance(node.test.op, ast.Not) and _is_type_checking(node.test.operand): - violations.append(f"{location}: use one `if TYPE_CHECKING: ... else:` block") - if isinstance(node, ast.ClassDef) and node.name.startswith("_"): - if any(marker in node.name for marker in ("Stub", "Unavailable", "Missing")): - violations.append(f"{location}: placeholder class {node.name} for a missing dependency") - if isinstance(node, ast.Try): - for handler in node.handlers: - names = handler.type.elts if isinstance(handler.type, ast.Tuple) else [handler.type] - catches_import_error = any( - isinstance(n, ast.Name) and n.id in {"ImportError", "ModuleNotFoundError"} for n in names - ) - reraises = any(isinstance(child, ast.Raise) for child in ast.walk(handler)) - if catches_import_error and not reraises: - violations.append(f"{location}: import fallback via try/except {ast.unparse(handler.type)}") - if not isinstance(node, ast.Call) or not node.args: - continue - func_name = node.func.id if isinstance(node.func, ast.Name) else getattr(node.func, "attr", "") - first = node.args[0] - if func_name not in {"import_optional", "import_optional_attr", "import_module"}: - continue - if not (isinstance(first, ast.Constant) and isinstance(first.value, str)): - continue - if _imports_driver(first.value, drivers): - violations.append(f"{location}: {func_name}({first.value!r}) for the adapter's own driver") - elif node.lineno not in function_lines and first.value.split(".")[0] in { - name.split(".")[0] for name in SECONDARY_OPTIONAL_MODULES - }: - violations.append(f"{location}: module-level {func_name}({first.value!r}); resolve on first use") - assert not violations, "Adapter import contract violations:\n" + "\n".join(violations) From 2800a73bf02fa76ed6fd0127ff56ec84ab4c7f15 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 19:36:32 +0000 Subject: [PATCH 129/174] chore: bump deps --- uv.lock | 80 ++++++++++++++++++++++++++------------------------------- 1 file changed, 37 insertions(+), 43 deletions(-) diff --git a/uv.lock b/uv.lock index 7c368a751..c5201c2a5 100644 --- a/uv.lock +++ b/uv.lock @@ -3,17 +3,14 @@ revision = 5 requires-python = ">=3.11, <4.0" resolution-markers = [ "python_full_version >= '3.15' and sys_platform == 'win32'", - "python_full_version >= '3.15' and sys_platform == 'emscripten'", - "python_full_version >= '3.15' and sys_platform != 'emscripten' and sys_platform != 'win32'", "python_full_version == '3.14.*' and sys_platform == 'win32'", + "python_full_version >= '3.15' and sys_platform == 'emscripten'", "python_full_version == '3.14.*' and sys_platform == 'emscripten'", + "python_full_version >= '3.15' and sys_platform != 'emscripten' and sys_platform != 'win32'", "python_full_version == '3.14.*' and sys_platform != 'emscripten' and sys_platform != 'win32'", - "python_full_version == '3.13.*' and sys_platform == 'win32'", - "python_full_version == '3.13.*' and sys_platform == 'emscripten'", - "python_full_version == '3.13.*' and sys_platform != 'emscripten' and sys_platform != 'win32'", - "python_full_version == '3.12.*' and sys_platform == 'win32'", - "python_full_version == '3.12.*' and sys_platform == 'emscripten'", - "python_full_version == '3.12.*' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform == 'win32'", + "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform == 'emscripten'", + "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform != 'emscripten' and sys_platform != 'win32'", "python_full_version < '3.12' and sys_platform == 'win32'", "python_full_version < '3.12' and sys_platform == 'emscripten'", "python_full_version < '3.12' and sys_platform != 'emscripten' and sys_platform != 'win32'", @@ -509,19 +506,19 @@ wheels = [ [[package]] name = "arrow-odbc" -version = "10.6.1" +version = "10.6.2" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "cffi" }, { name = "pyarrow" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/12/f8/7d0d0b88cadd9d28050ad7bce851af086857eb1401d09b54b26045bfbe99/arrow_odbc-10.6.1.tar.gz", hash = "sha256:f7449d292632f455f72e7ad8bb6b4fe386ac5602738225526294e1ffee0374d5", size = 107639, upload-time = "2026-10-19T20:22:59.75Z" } +sdist = { url = "https://files.pythonhosted.org/packages/97/76/990a0332ab2285a870ac1b39197917413f6fbca2ff192f859a5cc6d25575/arrow_odbc-10.6.2.tar.gz", hash = "sha256:badfa0587e9d569e265082f485686b32df0434734277b41348220268aa9e4152", size = 108670, upload-time = "2026-10-25T14:34:07.167Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/7d/17/8c32bd7afa88b6cb8ee120e0051218d2bab52dd7e1e60387230ac5b00ef1/arrow_odbc-10.6.1-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:4de744a4f20ebe574030bb1943f9aa10cb242b77ac7128a1efd1539d943cee6e", size = 753105, upload-time = "2026-10-19T20:25:14.844Z" }, - { url = "https://files.pythonhosted.org/packages/88/7f/e7c884fe16844f38899053f0e2e9f5ad66a984bb79823b1aef6d17b743d5/arrow_odbc-10.6.1-py3-none-macosx_11_0_arm64.whl", hash = "sha256:eafdf71851aac18c37822a67b57d704f2b9a7b133205da253445d374ca6c90c2", size = 735430, upload-time = "2026-10-19T20:22:58.19Z" }, - { url = "https://files.pythonhosted.org/packages/67/54/0deab13f42cbdd8aaeec9b957164b49b54d805fb3c94ca27d3857a07ee09/arrow_odbc-10.6.1-py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:72d97c33cc41bde4b071bad8b20932e360d11f4dad562bc0687cd59b8a11a3cc", size = 843468, upload-time = "2026-10-19T20:28:59.043Z" }, - { url = "https://files.pythonhosted.org/packages/94/bb/920f8e03a0340b8c8ad945c6fde73df07519596def2e40ff6ef70358ec8f/arrow_odbc-10.6.1-py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:78d48b9fe0291b32f50e8b1fb81c5bb2cd74fe7f636b34052551070a51ae2a5a", size = 869855, upload-time = "2026-10-19T20:25:20.298Z" }, - { url = "https://files.pythonhosted.org/packages/ec/12/311cd3f26939ead9b7f276988b9bd9d7b049c47c84af6b71cd35bd3d24c0/arrow_odbc-10.6.1-py3-none-win_amd64.whl", hash = "sha256:6226df0d91af94bdb47f575933e6c704e46d525d8b37348c43f7207f4b7d5e9b", size = 654895, upload-time = "2026-10-19T20:25:03.548Z" }, + { url = "https://files.pythonhosted.org/packages/68/a4/c7e9406f07406b61d6e7c7c491f5e6cf7602a4d49727adcaf76376d03524/arrow_odbc-10.6.2-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:d1603f351058daa8fc9d7d6df9c5da5bd57a0d2368b21cc8710928f540551073", size = 759654, upload-time = "2026-10-25T14:32:14.506Z" }, + { url = "https://files.pythonhosted.org/packages/f5/33/212da57507d06cf4d7f452e77339be3c66500b76a886b26b7a106a1968d5/arrow_odbc-10.6.2-py3-none-macosx_11_0_arm64.whl", hash = "sha256:50fea0f9441b7cef2ea689cd90705ad4f9d214ec65488bb256354ca2f77ce1b4", size = 739566, upload-time = "2026-10-25T14:34:06.155Z" }, + { url = "https://files.pythonhosted.org/packages/54/3a/105b89a01a25b8bd7b8a19ccd9cac9dfe03498df8f7553956c8ca23d076a/arrow_odbc-10.6.2-py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:54e25aa9c094c713298741048b5725530bbb1f5637adc3917b935045fd0c08e4", size = 846711, upload-time = "2026-10-25T14:32:23.305Z" }, + { url = "https://files.pythonhosted.org/packages/36/1c/3856b0cfeb0e84b8c9b1b1744a2a393a5b6553580039588cf6da3d8af80a/arrow_odbc-10.6.2-py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:3cf90b019f4a52c771e15a7847504f7826061aa04b77d290353f608abf62ef07", size = 872316, upload-time = "2026-10-25T14:32:53.888Z" }, + { url = "https://files.pythonhosted.org/packages/f1/80/2976600c3e8f0693c9901cf544f741aa6195d0027e5a5db2174fe3822d0a/arrow_odbc-10.6.2-py3-none-win_amd64.whl", hash = "sha256:763a49240a8a38f86974789ca0692560e02d60b389ad6abb8040cd2338c63773", size = 659820, upload-time = "2026-10-25T14:32:37.285Z" }, ] [[package]] @@ -846,7 +843,7 @@ wheels = [ [[package]] name = "bump-my-version" -version = "1.6.5" +version = "1.6.6" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "click" }, @@ -859,9 +856,9 @@ dependencies = [ { name = "tomlkit" }, { name = "wcmatch" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/a9/65/338597a14c53ae2b9d004031dbcb61a0641484b20469b5b7bd5250a37a34/bump_my_version-1.6.5.tar.gz", hash = "sha256:0fd40347bca359bdee9999c93e101d2ad4612cf10319bc5e81181e57fc313e0d", size = 1158255, upload-time = "2026-10-24T14:02:28.037Z" } +sdist = { url = "https://files.pythonhosted.org/packages/2b/df/ff4a0fc20547b4bf30b91fb6300cf0bd91d40a33dbb3cdb9cc698d0e71c0/bump_my_version-1.6.6.tar.gz", hash = "sha256:3f0a140b2a275dc030b2adc4dacb86247101ede38a5032d9fc9f291a17cfeb9d", size = 1161185, upload-time = "2026-10-25T13:45:13.457Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/15/bd/a8698a2143cccd7d5a5f23e7ec8520c2133e4083ccda1799531490fad17d/bump_my_version-1.6.5-py3-none-any.whl", hash = "sha256:7fc77cfa499edfa18072aaf7cc9e6697b8a66c7ba76f467c3c4fba4c554f2521", size = 67396, upload-time = "2026-10-24T14:02:26.214Z" }, + { url = "https://files.pythonhosted.org/packages/19/21/01f6f456e466363ef970a69885403bba3337c5a442bc4a7f52d2f5a26dcd/bump_my_version-1.6.6-py3-none-any.whl", hash = "sha256:99c16933a7ca5c7ebeb622ab6bf5a964ad0bda40c8e5f62311b27d9414edf123", size = 68031, upload-time = "2026-10-25T13:45:11.376Z" }, ] [[package]] @@ -1384,11 +1381,11 @@ wheels = [ [[package]] name = "deepmerge" -version = "3.0.1" +version = "4.0.0" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/38/6e/5cb3548b4d3112fea529375e55e6f3cdc52b8054e3a66f203b1f888ba885/deepmerge-3.0.1.tar.gz", hash = "sha256:35b39a4cb92cf328d6eca61cbbf65f68a37c2ceb3085f0f853cbb2e52a59fc23", size = 22328, upload-time = "2026-09-01T14:09:44.383Z" } +sdist = { url = "https://files.pythonhosted.org/packages/ca/e9/5e99e56e03f9461858e686fdaf82e2516671bc1bd46f8b7ad7b9a8375f6f/deepmerge-4.0.0.tar.gz", hash = "sha256:948565f2bc9cdb41259a4a23d26638e5840c9e084a343ab279ffd3edcfdf74d3", size = 22792, upload-time = "2026-10-25T19:25:28.967Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/20/91/600003aaad107e27553fbc9cbfc57e96fa37e0223d2ef9d7d3a0e8d8d070/deepmerge-3.0.1-py3-none-any.whl", hash = "sha256:35c96f6a68fcf90719a5b31d9f8042ecef6c00fb56836660d33455d0f5cfda65", size = 14909, upload-time = "2026-09-01T14:09:43.364Z" }, + { url = "https://files.pythonhosted.org/packages/c2/89/362146b2ec7d0d503f824ac8496c88f0408ee0d39203b02cc10feac08146/deepmerge-4.0.0-py3-none-any.whl", hash = "sha256:a36dd5810fc3b90bd3e0b681158221361d6a31a6109411b27b211b371fed6d5b", size = 15429, upload-time = "2026-10-25T19:25:28.012Z" }, ] [[package]] @@ -3743,17 +3740,14 @@ version = "2.5.4" source = { registry = "https://pypi.org/simple" } resolution-markers = [ "python_full_version >= '3.15' and sys_platform == 'win32'", - "python_full_version >= '3.15' and sys_platform == 'emscripten'", - "python_full_version >= '3.15' and sys_platform != 'emscripten' and sys_platform != 'win32'", "python_full_version == '3.14.*' and sys_platform == 'win32'", + "python_full_version >= '3.15' and sys_platform == 'emscripten'", "python_full_version == '3.14.*' and sys_platform == 'emscripten'", + "python_full_version >= '3.15' and sys_platform != 'emscripten' and sys_platform != 'win32'", "python_full_version == '3.14.*' and sys_platform != 'emscripten' and sys_platform != 'win32'", - "python_full_version == '3.13.*' and sys_platform == 'win32'", - "python_full_version == '3.13.*' and sys_platform == 'emscripten'", - "python_full_version == '3.13.*' and sys_platform != 'emscripten' and sys_platform != 'win32'", - "python_full_version == '3.12.*' and sys_platform == 'win32'", - "python_full_version == '3.12.*' and sys_platform == 'emscripten'", - "python_full_version == '3.12.*' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform == 'win32'", + "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform == 'emscripten'", + "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform != 'emscripten' and sys_platform != 'win32'", ] sdist = { url = "https://files.pythonhosted.org/packages/95/b0/c7453d0b6e2073c3264468b106ee1563750cecc910965e67357e3698c83e/numpy-2.5.4.tar.gz", hash = "sha256:9a94cf751c9ad8ebaa835bcd3d40dacf8534ad086b88c38029b65123c7999d2a", size = 20866315, upload-time = "2026-10-10T20:05:31.422Z" } wheels = [ @@ -4204,30 +4198,30 @@ wheels = [ [[package]] name = "polars" -version = "2.0.0" +version = "2.1.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "polars-runtime-32" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/8e/e9/001f371ec6a1bb54893f599ceebd56e6144fed4091f09f09fec0021a9276/polars-2.0.0.tar.gz", hash = "sha256:62da109e27a19a9d36657ee25dc035c9d3f87e7bd610526fe467dc37ea7dc115", size = 778215, upload-time = "2026-10-06T11:51:29.679Z" } +sdist = { url = "https://files.pythonhosted.org/packages/b9/da/c27daf205a824f89f71172a75ad0098c5837a5a31ddeaac7f8737ea8574b/polars-2.1.0.tar.gz", hash = "sha256:9eee488d6e10773c028753dfc8dac647fd60a708108dcb818a733e85d33df253", size = 811508, upload-time = "2026-10-25T18:00:33.43Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/ac/09/cc33bbd5463749c116b62c204d88bed6c02a6cb901eac7adab0d38651b07/polars-2.0.0-py3-none-any.whl", hash = "sha256:35d62f3541b7a6d4c360a2e2f07fccc0c2bcbd33b0ea51c83a25417a47a3f3ad", size = 876611, upload-time = "2026-10-06T11:44:04.327Z" }, + { url = "https://files.pythonhosted.org/packages/d5/d4/e696a46c74ec8cac6434764deaf72f7edcc2e91c2145a9b3f61900a7aa7a/polars-2.1.0-py3-none-any.whl", hash = "sha256:63c50ce3533f5a03d17c3a2a24f1d0e0c7eda98fbde4268a3cba900d34060e67", size = 911474, upload-time = "2026-10-25T17:58:26.863Z" }, ] [[package]] name = "polars-runtime-32" -version = "2.0.0" +version = "2.1.0" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/34/ad/dbb6f6d7070867951532bcfe5e6a648d8777b416b18cddabc07030404e8c/polars_runtime_32-2.0.0.tar.gz", hash = "sha256:b5f9afcc742b4a67eabd2c680ff0f12eb02ede9b4bf807bffabd6dbb9a58d5c7", size = 3591339, upload-time = "2026-10-06T11:51:31.076Z" } +sdist = { url = "https://files.pythonhosted.org/packages/e0/a6/46e09af693326735b1fba3a8bf5a725b57136fb1342e707e75d3113c826b/polars_runtime_32-2.1.0.tar.gz", hash = "sha256:c575f5f3cdb8a5f9be65477b9b4202b062593687fe62a0082add87fbd4571067", size = 3878004, upload-time = "2026-10-25T18:00:34.932Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/82/88/d35dec6c8928dfbaa1cccf9b626a1067da906e792c92d9f994ca825ab2b5/polars_runtime_32-2.0.0-cp310-abi3-macosx_10_12_x86_64.whl", hash = "sha256:ffb7ac6cf4e8c4a652df1951e3c3840c7c23a033603d5a9efd422fa8dd699d82", size = 52494314, upload-time = "2026-10-06T11:44:07.768Z" }, - { url = "https://files.pythonhosted.org/packages/5f/fd/2237bf53ffaff47cdf1edc6c10587a7a6444d4951150eeb08d84f3493ff8/polars_runtime_32-2.0.0-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:7012d8a0201bd95638545ce8f256c0efe2c5cab0f806eb043021dddde5a9498b", size = 47930083, upload-time = "2026-10-06T11:44:11.592Z" }, - { url = "https://files.pythonhosted.org/packages/0d/0d/85e3ed90417996fc09770be91b39979074fe2978fc15b431bf8a9459760d/polars_runtime_32-2.0.0-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:8b85bb42e6009acc9629afcc70a83473fd468694d6a30ffb0ab376c8dd1a0a17", size = 50417889, upload-time = "2026-10-06T11:50:20.774Z" }, - { url = "https://files.pythonhosted.org/packages/83/88/e9fecfd49159da92f54ff2445883577a0f1bc195da53ecc9535c458d55dd/polars_runtime_32-2.0.0-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:0d6ac584ea2b38913784db943879412380d92e28ab9cb88e20a77ba71ba3f911", size = 54475036, upload-time = "2026-10-06T11:50:24.411Z" }, - { url = "https://files.pythonhosted.org/packages/48/ad/b2abf732697b21467aaaeaac0f3bf7eee0d89c59ce8125f1ed41b28a2d97/polars_runtime_32-2.0.0-cp310-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:a6bf5e260e0a6f00d0f9181438fe9e45776df8c66cee9cba16e3675cc3888488", size = 50579474, upload-time = "2026-10-06T11:50:28.377Z" }, - { url = "https://files.pythonhosted.org/packages/7f/05/304deee59a95865e1b5e9ec7b066069b49093b81b768f473d9d3b165c686/polars_runtime_32-2.0.0-cp310-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:55c26eef325b6840584d91aac232e9cf3ac19e1b904594b9b54131be1edeab4d", size = 54413293, upload-time = "2026-10-06T11:50:31.828Z" }, - { url = "https://files.pythonhosted.org/packages/61/59/8c9fd7199f7c4eb1b64e640306a946a2e4a46337b3bbb33b840972c7d84b/polars_runtime_32-2.0.0-cp310-abi3-win_amd64.whl", hash = "sha256:7da1caf3c7b4f397fb213c984013a0c755557619a2d511899a1ff74392484078", size = 54229989, upload-time = "2026-10-06T11:50:35.206Z" }, - { url = "https://files.pythonhosted.org/packages/e2/93/43608026f38aa6ed4d22da8597706a61682ee403caef0021ce8e6dc73227/polars_runtime_32-2.0.0-cp310-abi3-win_arm64.whl", hash = "sha256:c30ba698c8904048df4a9bc3d6c5033cc2d0a7cbb0e13f4fd2de5a1947b61994", size = 48730655, upload-time = "2026-10-06T11:50:38.756Z" }, + { url = "https://files.pythonhosted.org/packages/7a/40/b424f3946c702ff48eef14e44d35df81a25a784f20f77355f39c746528fc/polars_runtime_32-2.1.0-cp311-abi3-macosx_10_12_x86_64.whl", hash = "sha256:de243b6e5a0b9f3abeaefe542136fcdc6dd8df5be3c988a40838d075314e246c", size = 52169672, upload-time = "2026-10-25T17:58:32.437Z" }, + { url = "https://files.pythonhosted.org/packages/fe/38/07066430ac655a8cab2f9e96ea1fec6ca1f20ec47844afe117662d43ace0/polars_runtime_32-2.1.0-cp311-abi3-macosx_11_0_arm64.whl", hash = "sha256:317314d65949f1d58c4939adf83aa0d2cfaf8da50fd4a725a2de27bde2c10018", size = 47626883, upload-time = "2026-10-25T17:58:38.688Z" }, + { url = "https://files.pythonhosted.org/packages/49/a8/a773e20b451bbaab4d8267867b7c43a0c3aa67ebb6f781423a42a5a1af37/polars_runtime_32-2.1.0-cp311-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:777228cc8bf9cda36dee85ba74d13a25974ad4caffc5da93e10b1da9d810ce2d", size = 50033514, upload-time = "2026-10-25T17:58:44.261Z" }, + { url = "https://files.pythonhosted.org/packages/d2/e6/649c524252d32b24d0ba92278033823fa8e0040ee3748ce5afbeda107403/polars_runtime_32-2.1.0-cp311-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:0b838056cf365d7497971d9d9825c547766923338771c48c545b86ead2ebc572", size = 53863954, upload-time = "2026-10-25T17:58:48.737Z" }, + { url = "https://files.pythonhosted.org/packages/90/f7/2eb405bf5ead6373a33fa5dd4637fe996565a1f07ccce8911ed167673740/polars_runtime_32-2.1.0-cp311-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:04cfc3db1571b74ea4145944de7feceed4bfc9b26e5b2c265e27ad3b2d3866a9", size = 50204290, upload-time = "2026-10-25T17:58:53.33Z" }, + { url = "https://files.pythonhosted.org/packages/33/4d/d76cd42042605e079d6df91cc4fea09ed5f4cb836c72c952a17f585860a9/polars_runtime_32-2.1.0-cp311-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:af58b5f6e5cc44b578f2a93d752f87c135c3a53205b4360301bc3a9b89dfbfd1", size = 53794253, upload-time = "2026-10-25T17:58:57.99Z" }, + { url = "https://files.pythonhosted.org/packages/4d/4e/2a419d0def2908255dd1b7318c48430319859dc92ca156e2f81ff7bcd78b/polars_runtime_32-2.1.0-cp311-abi3-win_amd64.whl", hash = "sha256:f2fe32f48af2395cd6ec52a433a03970a69aab1dbb66a613729fd549ce92c904", size = 53926265, upload-time = "2026-10-25T17:59:03.13Z" }, + { url = "https://files.pythonhosted.org/packages/c8/1b/15ffba3264baf21d2f04c51f54b01159cb6f4fad4e84de1404be552d601f/polars_runtime_32-2.1.0-cp311-abi3-win_arm64.whl", hash = "sha256:0af30dd53a7725f703d0d21dd421a39e0165378bfaa41431008f286608376d71", size = 48636436, upload-time = "2026-10-25T17:59:08.476Z" }, ] [[package]] @@ -5067,7 +5061,7 @@ wheels = [ [[package]] name = "pytest-databases" version = "0.19.0" -source = { git = "https://github.com/litestar-org/pytest-databases.git?branch=test%2Fclientless-all#ee262ac973907bdefe539cbfee4e5630f380724f" } +source = { git = "https://github.com/litestar-org/pytest-databases.git?branch=test%2Fclientless-all#09a0f8f82b4513c90c2c9628cd762ce75de31d5d" } dependencies = [ { name = "docker" }, { name = "filelock" }, From 179ee7545aee925d23dbb93eb8abe23ecd7f295e Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 20:29:14 +0000 Subject: [PATCH 130/174] fix(adbc): release acquisitions when session setup fails --- sqlspec/adapters/adbc/_typing.py | 28 ++++-- sqlspec/adapters/adbc/config.py | 52 ++++++---- .../test_adbc_config_normalization.py | 95 +++++++++++++++++- .../unit/adapters/test_adbc/test_profiling.py | 96 +++++++++++++++++++ 4 files changed, 242 insertions(+), 29 deletions(-) diff --git a/sqlspec/adapters/adbc/_typing.py b/sqlspec/adapters/adbc/_typing.py index 847744afc..04e8604ba 100644 --- a/sqlspec/adapters/adbc/_typing.py +++ b/sqlspec/adapters/adbc/_typing.py @@ -110,21 +110,31 @@ def __enter__(self) -> "AdbcDriver": driver_class = AdbcDriver self._connection = self._acquire_connection() - resolved_dialect = resolve_dialect_name(self._statement_config.dialect) or None - self._driver = driver_class( - connection=self._connection, - statement_config=self._statement_config, - driver_features=self._driver_features, - dialect=resolved_dialect, - ) - return self._prepare_driver(self._driver) + try: + resolved_dialect = resolve_dialect_name(self._statement_config.dialect) or None + self._driver = driver_class( + connection=self._connection, + statement_config=self._statement_config, + driver_features=self._driver_features, + dialect=resolved_dialect, + ) + return self._prepare_driver(self._driver) + except BaseException as exc: + connection = self._connection + self._connection = None + self._driver = None + with contextlib.suppress(BaseException): + self._release_connection(connection, exc_type=type(exc), exc_val=exc, exc_tb=exc.__traceback__) + raise def __exit__( self, exc_type: "type[BaseException] | None", exc_val: "BaseException | None", exc_tb: "TracebackType | None" ) -> "bool | None": if self._connection is not None: - self._release_connection(self._connection, exc_type=exc_type, exc_val=exc_val, exc_tb=exc_tb) + connection = self._connection self._connection = None + self._driver = None + self._release_connection(connection, exc_type=exc_type, exc_val=exc_val, exc_tb=exc_tb) return None diff --git a/sqlspec/adapters/adbc/config.py b/sqlspec/adapters/adbc/config.py index 904ff4cc4..32dd09d45 100644 --- a/sqlspec/adapters/adbc/config.py +++ b/sqlspec/adapters/adbc/config.py @@ -1,5 +1,6 @@ """ADBC database configuration.""" +import contextlib from collections.abc import Callable from pathlib import Path from typing import TYPE_CHECKING, Any, ClassVar, Literal, NotRequired, TypedDict, cast @@ -206,29 +207,38 @@ def driver_features(self) -> "dict[str, Any]": return self._driver_features def acquire_connection(self) -> "AdbcConnection": - self._connection = self._config.create_connection() - if self._config._resolved_dialect == "duckdb" and bool( - self._driver_features.get("enable_statement_profiling", False) - ): - self._profile_path = create_duckdb_profile_path() - configure_adbc_duckdb_profiling(self._connection, self._profile_path) - self._driver_features["_duckdb_profile_path"] = str(self._profile_path) - return cast("AdbcConnection", self._connection) + connection: AdbcConnection = self._config.create_connection() + self._connection = connection + try: + if self._config._resolved_dialect == "duckdb" and bool( + self._driver_features.get("enable_statement_profiling", False) + ): + self._profile_path = create_duckdb_profile_path() + configure_adbc_duckdb_profiling(connection, self._profile_path) + self._driver_features["_duckdb_profile_path"] = str(self._profile_path) + except BaseException as exc: + with contextlib.suppress(BaseException): + self.release_connection(connection, exc_type=type(exc), exc_val=exc, exc_tb=exc.__traceback__) + raise + return connection def release_connection(self, _conn: "AdbcConnection", **kwargs: Any) -> None: if self._connection is None: return + connection = self._connection + profile_path = self._profile_path + self._connection = None + self._profile_path = None + raw_path = self._driver_features.pop("_duckdb_profile_path", None) try: - if self._connection is not self._config.connection_instance: - self._connection.close() + if connection is not self._config.connection_instance: + connection.close() finally: - self._connection = None - if self._profile_path is not None: - cleanup_duckdb_profile_path(self._profile_path) - self._profile_path = None - raw_path = self._driver_features.pop("_duckdb_profile_path", None) - if isinstance(raw_path, str): - cleanup_duckdb_profile_path(Path(raw_path)) + try: + cleanup_duckdb_profile_path(profile_path) + finally: + if isinstance(raw_path, str): + cleanup_duckdb_profile_path(Path(raw_path)) class AdbcConfig(NoPoolSyncConfig[AdbcConnection, AdbcDriver]): @@ -348,6 +358,7 @@ def create_connection(self) -> AdbcConnection: if self.connection_instance is not None: return self.connection_instance + connection: AdbcConnection | None = None try: driver_name = cast( "str | None", self.connection_config.get("driver_name") or self.connection_config.get("driver") @@ -361,7 +372,12 @@ def create_connection(self) -> AdbcConnection: if self._user_connection_hook is not None: self._user_connection_hook(cast("AdbcConnection", connection)) return cast("AdbcConnection", connection) - except Exception as e: + except BaseException as e: + if connection is not None: + with contextlib.suppress(BaseException): + connection.close() + if not isinstance(e, Exception): + raise err_driver_name = ( self.connection_config.get("driver_name") or self.connection_config.get("driver") or "Unknown" ) diff --git a/tests/unit/adapters/test_adbc/test_adbc_config_normalization.py b/tests/unit/adapters/test_adbc/test_adbc_config_normalization.py index f72dea3ee..af504dbd9 100644 --- a/tests/unit/adapters/test_adbc/test_adbc_config_normalization.py +++ b/tests/unit/adapters/test_adbc/test_adbc_config_normalization.py @@ -1,13 +1,15 @@ """Unit tests for ADBC config normalization helpers.""" -from typing import Any, get_type_hints +from typing import Any, cast, get_type_hints from unittest.mock import MagicMock import adbc_driver_manager.dbapi +import adbc_driver_sqlite.dbapi import pytest from adbc_driver_bigquery import DatabaseOptions from sqlspec.adapters.adbc import AdbcConfig, AdbcDriver, AdbcDriverFeatures +from sqlspec.adapters.adbc._typing import AdbcSessionContext from sqlspec.adapters.adbc.config import AdbcConnectionParams from sqlspec.adapters.adbc.core import ( build_connection_config, @@ -15,7 +17,7 @@ resolve_driver_name_from_config, ) from sqlspec.core import SQL, ParameterStyle -from sqlspec.exceptions import SQLSpecError +from sqlspec.exceptions import ImproperConfigurationError, SQLSpecError def test_adbc_config_runs_connection_create_callback(monkeypatch: Any) -> None: @@ -37,6 +39,95 @@ def fake_connect(**_kwargs: Any) -> Any: assert "on_connection_create" not in config.driver_features +@pytest.mark.parametrize("primary", [ValueError("hook failed"), KeyboardInterrupt("hook interrupted")]) +@pytest.mark.parametrize("cleanup_error", [None, RuntimeError("close failed"), SystemExit("close interrupted")]) +def test_adbc_config_hook_failure_closes_owned_connection( + monkeypatch: pytest.MonkeyPatch, primary: BaseException, cleanup_error: BaseException | None +) -> None: + """A failing initialization hook preserves its error after closing its owned connection.""" + connection = MagicMock() + connection.close.side_effect = cleanup_error + monkeypatch.setattr(adbc_driver_sqlite.dbapi, "connect", lambda **_kwargs: connection) + + def fail_hook(_connection: Any) -> None: + raise primary + + config = AdbcConfig( + connection_config={"driver_name": "sqlite"}, driver_features={"on_connection_create": fail_hook} + ) + expected = ImproperConfigurationError if isinstance(primary, Exception) else type(primary) + with pytest.raises(expected) as caught: + config.create_connection() + if isinstance(primary, Exception): + assert caught.value.__cause__ is primary + else: + assert caught.value is primary + connection.close.assert_called_once_with() + + +@pytest.mark.parametrize("stage", ["constructor", "preparation"]) +@pytest.mark.parametrize("primary", [RuntimeError("setup failed"), KeyboardInterrupt("setup interrupted")]) +@pytest.mark.parametrize("cleanup_error", [None, SystemExit("release interrupted")]) +def test_adbc_session_setup_failure_releases_once( + stage: str, primary: BaseException, cleanup_error: BaseException | None +) -> None: + """Failed session entry releases the acquisition with its original exception details.""" + connection = MagicMock() + released: list[dict[str, Any]] = [] + + class FailingDriver: + def __init__(self, **kwargs: Any) -> None: + raise primary + + def prepare(driver: AdbcDriver) -> AdbcDriver: + raise primary + + def release(conn: Any, **kwargs: Any) -> None: + assert conn is connection + assert kwargs["exc_tb"] is primary.__traceback__ + released.append(kwargs) + connection.close() + if cleanup_error is not None: + raise cleanup_error + + context = AdbcSessionContext( + acquire_connection=lambda: connection, + release_connection=release, + statement_config=AdbcConfig(connection_config={"driver_name": "sqlite"}).statement_config, + driver_features={}, + prepare_driver=prepare, + driver_class=cast("type[AdbcDriver]", FailingDriver) if stage == "constructor" else AdbcDriver, + ) + with pytest.raises(type(primary)) as caught: + context.__enter__() + assert caught.value is primary + assert len(released) == 1 + assert released[0]["exc_type"] is type(primary) + assert released[0]["exc_val"] is primary + context.__exit__(None, None, None) + connection.close.assert_called_once_with() + + +def test_adbc_session_release_once_after_cleanup_failure() -> None: + """An exit failure cannot release an acquisition again on a later exit.""" + connection = MagicMock() + cleanup_error = RuntimeError("close failed") + connection.close.side_effect = cleanup_error + context = AdbcSessionContext( + acquire_connection=lambda: connection, + release_connection=lambda conn, **_kwargs: conn.close(), + statement_config=AdbcConfig(connection_config={"driver_name": "sqlite"}).statement_config, + driver_features={}, + prepare_driver=lambda driver: driver, + ) + context.__enter__() + with pytest.raises(RuntimeError) as caught: + context.__exit__(None, None, None) + assert caught.value is cleanup_error + context.__exit__(None, None, None) + connection.close.assert_called_once_with() + + def test_resolve_driver_name_alias_to_connect_path() -> None: """Resolve short driver aliases to concrete connect paths.""" config = AdbcConfig(connection_config={"driver_name": "sqlite"}) diff --git a/tests/unit/adapters/test_adbc/test_profiling.py b/tests/unit/adapters/test_adbc/test_profiling.py index 231a56384..079755eda 100644 --- a/tests/unit/adapters/test_adbc/test_profiling.py +++ b/tests/unit/adapters/test_adbc/test_profiling.py @@ -2,12 +2,108 @@ import subprocess import sys +import tempfile +from importlib import import_module from pathlib import Path +from unittest.mock import MagicMock + +import pytest from sqlspec.adapters.adbc import AdbcConfig from sqlspec.observability import ObservabilityConfig, StatementEvent +@pytest.mark.parametrize("stage", ["path", "configuration"]) +@pytest.mark.parametrize("borrowed", [False, True]) +@pytest.mark.parametrize("cleanup_fails", [False, True]) +def test_adbc_profiling_setup_failure_releases_owned_connection( + monkeypatch: pytest.MonkeyPatch, stage: str, borrowed: bool, cleanup_fails: bool +) -> None: + """Profiling failures preserve setup errors and close only owned connections.""" + connection = MagicMock() + primary = RuntimeError("profile setup failed") + removed: list[Path] = [] + original_unlink = Path.unlink + + def remove_profile(path: Path, missing_ok: bool = False) -> None: + if path.name.startswith("sqlspec_duckdb_profile_"): + removed.append(path) + if cleanup_fails: + raise KeyboardInterrupt("profile cleanup interrupted") + else: + original_unlink(path, missing_ok=missing_ok) + + def temporary_directory() -> str: + raise primary + + monkeypatch.setattr(import_module("adbc_driver_duckdb.dbapi"), "connect", lambda **_kwargs: connection) + monkeypatch.setattr(Path, "unlink", remove_profile) + if stage == "path": + monkeypatch.setattr(tempfile, "gettempdir", temporary_directory) + else: + connection.cursor.return_value.execute.side_effect = primary + if cleanup_fails: + connection.close.side_effect = SystemExit("connection cleanup interrupted") + config = AdbcConfig( + connection_config={"driver_name": "duckdb"}, + connection_instance=connection if borrowed else None, + driver_features={"enable_statement_profiling": True}, + ) + context = config.provide_session() + with pytest.raises(RuntimeError) as caught: + context.__enter__() + assert caught.value is primary + context.__exit__(None, None, None) + if borrowed: + connection.close.assert_not_called() + else: + connection.close.assert_called_once_with() + if stage == "configuration": + assert len(removed) == 1 + connection.cursor.return_value.close.assert_called_once_with() + + +@pytest.mark.parametrize("setup_fails", [False, True]) +def test_adbc_profiling_cleanup_attempts_both_paths_and_release_once( + monkeypatch: pytest.MonkeyPatch, setup_fails: bool +) -> None: + """A profile cleanup failure clears ownership and still attempts both path cleanups.""" + connection = MagicMock() + primary = KeyboardInterrupt("metadata interrupted") + cleanup_error = RuntimeError("profile cleanup failed") + removed: list[Path] = [] + original_unlink = Path.unlink + + def remove_profile(path: Path, missing_ok: bool = False) -> None: + if path.name.startswith("sqlspec_duckdb_profile_"): + removed.append(path) + raise cleanup_error + original_unlink(path, missing_ok=missing_ok) + + monkeypatch.setattr(import_module("adbc_driver_duckdb.dbapi"), "connect", lambda **_kwargs: connection) + monkeypatch.setattr(Path, "unlink", remove_profile) + if setup_fails: + connection.adbc_get_info.side_effect = primary + config = AdbcConfig( + connection_config={"driver_name": "duckdb"}, driver_features={"enable_statement_profiling": True} + ) + context = config.provide_session() + if setup_fails: + with pytest.raises(KeyboardInterrupt) as setup_caught: + context.__enter__() + assert setup_caught.value is primary + else: + context.__enter__() + with pytest.raises(RuntimeError) as cleanup_caught: + context.__exit__(None, None, None) + assert cleanup_caught.value is cleanup_error + assert len(removed) == 2 + assert removed[0] == removed[1] + context.__exit__(None, None, None) + assert len(removed) == 2 + connection.close.assert_called_once_with() + + def test_adbc_duckdb_enable_statement_profiling_and_execute_script_events() -> None: """AdbcConfig with driver_name='duckdb' and enable_statement_profiling=True populates StatementEvent.profile and emits per-statement script events.""" events: list[StatementEvent] = [] From 57953f3d31c6cc1cf17f9fb376720ee96cb7ca1d Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 20:34:35 +0000 Subject: [PATCH 131/174] feat(adbc): add bounded exclusive PostgreSQL connection leases --- sqlspec/adapters/adbc/pool.py | 186 ++++++++++++ tests/unit/adapters/test_adbc/test_pool.py | 314 +++++++++++++++++++++ 2 files changed, 500 insertions(+) create mode 100644 sqlspec/adapters/adbc/pool.py create mode 100644 tests/unit/adapters/test_adbc/test_pool.py diff --git a/sqlspec/adapters/adbc/pool.py b/sqlspec/adapters/adbc/pool.py new file mode 100644 index 000000000..da0951c27 --- /dev/null +++ b/sqlspec/adapters/adbc/pool.py @@ -0,0 +1,186 @@ +"""Exclusive PostgreSQL ADBC connection leases.""" + +import contextlib +import math +import os +import threading +import time +from typing import TYPE_CHECKING + +from sqlspec.exceptions import ConnectionTimeoutError, DatabaseConnectionError, ImproperConfigurationError + +if TYPE_CHECKING: + from collections.abc import Callable + + from sqlspec.adapters.adbc._typing import AdbcConnection + +__all__ = ("AdbcConnectionPool",) + + +class AdbcConnectionPool: + """Own one original seed and bounded exclusive native clone leases. + + Native creation, cloning and cleanup run outside the bookkeeping condition. + Failed cleanup retains ownership until a subsequent close retries it. + """ + + __slots__ = ( + "_clones", + "_condition", + "_create_connection", + "_max_size", + "_pid", + "_reservations", + "_seed", + "_seed_lock", + "_state", + "_timeout", + ) + + def __init__( + self, create_connection: "Callable[[], AdbcConnection]", *, max_size: int = 4, timeout: float = 30.0 + ) -> None: + """Configure a lazy pool without opening any native connections.""" + if isinstance(max_size, bool) or not isinstance(max_size, int) or max_size <= 0: + msg = "ADBC pool max_size must be a positive integer." + raise ImproperConfigurationError(msg) + if not math.isfinite(timeout) or timeout < 0: + msg = "ADBC pool timeout must be finite and nonnegative." + raise ImproperConfigurationError(msg) + self._create_connection = create_connection + self._max_size = max_size + self._timeout = timeout + self._pid = os.getpid() + self._condition = threading.Condition(threading.Lock()) + self._seed_lock = threading.Lock() + self._seed: AdbcConnection | None = None + self._clones: dict[int, tuple[AdbcConnection, str]] = {} + self._reservations = 0 + self._state = "OPEN" + + def acquire(self) -> "AdbcConnection": + """Acquire an exclusive clone or wait up to the configured timeout.""" + self.check_pid() + deadline = time.monotonic() + self._timeout + with self._condition: + while True: + if self._state != "OPEN": + msg = "ADBC pool is closed to new acquisitions." + raise DatabaseConnectionError(msg) + if len(self._clones) + self._reservations < self._max_size: + self._reservations += 1 + break + remaining = deadline - time.monotonic() + if remaining <= 0: + msg = "Timed out waiting for an exclusive ADBC connection." + raise ConnectionTimeoutError(msg) + self._condition.wait(min(remaining, threading.TIMEOUT_MAX)) + try: + with self._seed_lock: + if self._seed is None: + self._seed = self._create_connection() + connection = self._seed.adbc_clone() + except BaseException: + with self._condition: + self._reservations -= 1 + self._condition.notify_all() + with contextlib.suppress(BaseException): + self._finish_close() + raise + with self._condition: + self._reservations -= 1 + accepted = self._state == "OPEN" + self._clones[id(connection)] = (connection, "active" if accepted else "closing") + self._condition.notify_all() + if not accepted: + self._close_clone(connection) + self._finish_close() + msg = "ADBC pool closed during connection acquisition." + raise DatabaseConnectionError(msg) + return connection + + def release(self, connection: "AdbcConnection", *, discard: bool = False) -> None: + """Release exactly one active lease, retaining any failed native close.""" + self.check_pid() + with self._condition: + entry = self._clones.get(id(connection)) + if entry is None or entry[0] is not connection or entry[1] != "active": + msg = "Connection is not an active lease of this ADBC pool." + raise DatabaseConnectionError(msg) + self._clones[id(connection)] = (connection, "closing") + self._close_clone(connection) + self._finish_close() + + def close(self) -> bool: + """Begin terminal drain and retry failed cleanup without closing active leases. + + Returns: + True only when every owned clone and the original seed are closed. + """ + self.check_pid() + with self._condition: + if self._state == "CLOSED": + return True + self._state = "DRAINING" + closing = [] + for key, (connection, state) in self._clones.items(): + if state in {"idle", "quarantined"}: + self._clones[key] = (connection, "closing") + closing.append(connection) + self._condition.notify_all() + first_error: BaseException | None = None + for connection in closing: + try: + self._close_clone(connection) + except BaseException as exc: + if first_error is None: + first_error = exc + try: + self._finish_close() + except BaseException as exc: + if first_error is None: + first_error = exc + if first_error is not None: + raise first_error + with self._condition: + return self._state == "CLOSED" + + @property + def is_closed(self) -> bool: + """Return whether every owned native resource has been closed.""" + self.check_pid() + with self._condition: + return self._state == "CLOSED" + + def check_pid(self) -> None: + """Reject inherited pool use before acquiring any lock or native handle.""" + if os.getpid() != self._pid: + msg = "ADBC pools cannot be used after fork; create a new configuration in the child." + raise DatabaseConnectionError(msg) + + def _close_clone(self, connection: "AdbcConnection") -> None: + try: + connection.close() + except BaseException: + with self._condition: + self._clones[id(connection)] = (connection, "quarantined") + self._condition.notify_all() + raise + with self._condition: + del self._clones[id(connection)] + self._condition.notify_all() + + def _finish_close(self) -> None: + with self._condition: + if self._state != "DRAINING" or self._clones or self._reservations: + return + with self._seed_lock: + with self._condition: + if self._state != "DRAINING" or self._clones or self._reservations: + return + if self._seed is not None: + self._seed.close() + self._seed = None + with self._condition: + self._state = "CLOSED" + self._condition.notify_all() diff --git a/tests/unit/adapters/test_adbc/test_pool.py b/tests/unit/adapters/test_adbc/test_pool.py new file mode 100644 index 000000000..fa6001fa4 --- /dev/null +++ b/tests/unit/adapters/test_adbc/test_pool.py @@ -0,0 +1,314 @@ +"""Exclusive native-shaped ADBC pool ownership contracts.""" + +import os +import threading +import time +from concurrent.futures import ThreadPoolExecutor +from unittest.mock import MagicMock + +import pytest + +from sqlspec.adapters.adbc.pool import AdbcConnectionPool +from sqlspec.exceptions import ConnectionTimeoutError, DatabaseConnectionError, ImproperConfigurationError + + +def test_adbc_pool_exclusive_original_seed_leases_and_nested_timeout() -> None: + """Every borrower owns a distinct original-seed clone within the configured bound.""" + seed, first, second, third = MagicMock(), MagicMock(), MagicMock(), MagicMock() + seed.adbc_clone.side_effect = [first, second, third] + create = MagicMock(return_value=seed) + pool = AdbcConnectionPool(create, max_size=2, timeout=0) + create.assert_not_called() + assert pool.acquire() is first + assert pool.acquire() is second + with pytest.raises(ConnectionTimeoutError): + pool.acquire() + pool.release(first) + assert pool.acquire() is third + for connection in (second, third): + pool.release(connection) + create.assert_called_once_with() + assert seed.adbc_clone.call_count == 3 + first.close.assert_called_once_with() + assert pool.close() is True + assert pool.is_closed is True + seed.close.assert_called_once_with() + + +@pytest.mark.parametrize("max_size", [0, -1, True, 1.5]) +def test_adbc_pool_rejects_invalid_capacity(max_size: int) -> None: + """Capacity is a positive integer and rejects bool before native acquisition.""" + create = MagicMock() + with pytest.raises(ImproperConfigurationError): + AdbcConnectionPool(create, max_size=max_size) + create.assert_not_called() + + +@pytest.mark.parametrize("timeout", [-1, float("inf"), float("nan")]) +def test_adbc_pool_rejects_invalid_timeout(timeout: float) -> None: + """Timeout is finite and nonnegative.""" + with pytest.raises(ImproperConfigurationError): + AdbcConnectionPool(MagicMock(), timeout=timeout) + + +@pytest.mark.parametrize("failure_stage", ["create", "clone"]) +def test_adbc_pool_failed_acquisition_recovers_reserved_capacity(failure_stage: str) -> None: + """A failed native open or clone releases its capacity reservation.""" + seed, lease = MagicMock(), MagicMock() + failure = RuntimeError("native acquisition failed") + create = MagicMock(return_value=seed) + seed.adbc_clone.return_value = lease + if failure_stage == "create": + create.side_effect = [failure, seed] + else: + seed.adbc_clone.side_effect = [failure, lease] + pool = AdbcConnectionPool(create, max_size=1, timeout=0) + with pytest.raises(RuntimeError) as caught: + pool.acquire() + assert caught.value is failure + assert pool.acquire() is lease + pool.release(lease) + assert pool.close() is True + + +def test_adbc_pool_rejects_foreign_and_duplicate_release() -> None: + """Only the originating pool may release a currently active lease once.""" + seed, lease, foreign = MagicMock(), MagicMock(), MagicMock() + seed.adbc_clone.return_value = lease + pool = AdbcConnectionPool(lambda: seed) + assert pool.acquire() is lease + with pytest.raises(DatabaseConnectionError): + pool.release(foreign) + foreign.close.assert_not_called() + pool.release(lease) + with pytest.raises(DatabaseConnectionError): + pool.release(lease) + lease.close.assert_called_once_with() + assert pool.close() is True + + +def test_adbc_pool_shutdown_drains_active_lease_without_forcing_close() -> None: + """Terminal shutdown rejects new borrowers and leaves an executing lease alone.""" + seed, lease = MagicMock(), MagicMock() + seed.adbc_clone.return_value = lease + pool = AdbcConnectionPool(lambda: seed) + assert pool.acquire() is lease + assert pool.close() is False + assert pool.is_closed is False + lease.close.assert_not_called() + seed.close.assert_not_called() + with pytest.raises(DatabaseConnectionError): + pool.acquire() + pool.release(lease) + assert pool.is_closed is True + assert pool.close() is True + lease.close.assert_called_once_with() + seed.close.assert_called_once_with() + + +def test_adbc_pool_quarantined_close_retains_capacity_and_retries_seed_last() -> None: + """A failed close retains the native lease until a later shutdown retries it.""" + seed, lease = MagicMock(), MagicMock() + failure = RuntimeError("native close failed") + seed.adbc_clone.return_value = lease + lease.close.side_effect = [failure, None] + pool = AdbcConnectionPool(lambda: seed, max_size=1, timeout=0) + assert pool.acquire() is lease + with pytest.raises(RuntimeError) as caught: + pool.release(lease) + assert caught.value is failure + with pytest.raises(ConnectionTimeoutError): + pool.acquire() + seed.close.assert_not_called() + assert pool.close() is True + assert lease.close.call_count == 2 + seed.close.assert_called_once_with() + + +def test_adbc_pool_failed_seed_close_is_retryable() -> None: + """CLOSED is published only after the original seed actually closes.""" + seed, lease = MagicMock(), MagicMock() + failure = RuntimeError("seed close failed") + seed.adbc_clone.return_value = lease + seed.close.side_effect = [failure, None] + pool = AdbcConnectionPool(lambda: seed) + assert pool.acquire() is lease + pool.release(lease) + with pytest.raises(RuntimeError) as caught: + pool.close() + assert caught.value is failure + assert pool.is_closed is False + assert pool.close() is True + assert seed.close.call_count == 2 + + +def test_adbc_pool_shutdown_during_clone_closes_clone_before_seed() -> None: + """Shutdown retains an in-flight reservation and closes its eventual clone first.""" + seed, lease = MagicMock(), MagicMock() + entered, proceed = threading.Event(), threading.Event() + closed: list[str] = [] + + def clone() -> MagicMock: + entered.set() + assert proceed.wait(3) + return lease + + seed.adbc_clone.side_effect = clone + seed.close.side_effect = lambda: closed.append("seed") + lease.close.side_effect = lambda: closed.append("lease") + pool = AdbcConnectionPool(lambda: seed, max_size=1) + with ThreadPoolExecutor(max_workers=2) as executor: + acquiring = executor.submit(pool.acquire) + assert entered.wait(3) + try: + assert executor.submit(pool.close).result(timeout=2) is False + assert closed == [] + finally: + proceed.set() + with pytest.raises(DatabaseConnectionError): + acquiring.result(timeout=3) + assert closed == ["lease", "seed"] + assert pool.is_closed is True + + +def test_adbc_pool_native_io_does_not_hold_bookkeeping_lock() -> None: + """Native open, clone and close permit another thread to inspect pool state.""" + seed, lease = MagicMock(), MagicMock() + with ThreadPoolExecutor(max_workers=1) as executor: + + def inspect_pool() -> None: + assert executor.submit(lambda: pool.is_closed).result(timeout=2) is False + + def create() -> MagicMock: + inspect_pool() + return seed + + def clone() -> MagicMock: + inspect_pool() + return lease + + seed.adbc_clone.side_effect = clone + seed.close.side_effect = inspect_pool + lease.close.side_effect = inspect_pool + pool = AdbcConnectionPool(create) + assert pool.acquire() is lease + pool.release(lease) + assert pool.close() is True + + +def test_adbc_pool_serializes_original_seed_cloning() -> None: + """Concurrent checkouts serialize original-seed cloning and yield distinct handles.""" + seed = MagicMock() + entered, proceed = threading.Event(), threading.Event() + start = threading.Barrier(3, timeout=3) + leases: list[MagicMock] = [] + native_guard = threading.Lock() + + def clone() -> MagicMock: + assert native_guard.acquire(blocking=False), "concurrent original-seed clone" + try: + entered.set() + assert proceed.wait(3) + lease = MagicMock() + lease.adbc_clone.side_effect = AssertionError("cloned a clone") + leases.append(lease) + return lease + finally: + native_guard.release() + + def acquire() -> object: + start.wait() + return pool.acquire() + + seed.adbc_clone.side_effect = clone + pool = AdbcConnectionPool(lambda: seed, max_size=2, timeout=0) + with ThreadPoolExecutor(max_workers=2) as executor: + first, second = executor.submit(acquire), executor.submit(acquire) + start.wait() + assert entered.wait(3) + proceed.set() + first_lease, second_lease = first.result(timeout=3), second.result(timeout=3) + assert first_lease is not second_lease + assert len(leases) == 2 + for lease in leases: + pool.release(lease) + assert pool.close() is True + + +def test_adbc_pool_closing_lease_retains_capacity_until_native_close_finishes() -> None: + """Native cleanup remains owned and counted while another thread is closing it.""" + seed, lease = MagicMock(), MagicMock() + entered, proceed = threading.Event(), threading.Event() + + def native_close() -> None: + entered.set() + assert proceed.wait(3) + + seed.adbc_clone.return_value = lease + lease.close.side_effect = native_close + pool = AdbcConnectionPool(lambda: seed, max_size=1, timeout=0) + assert pool.acquire() is lease + with ThreadPoolExecutor(max_workers=1) as executor: + releasing = executor.submit(pool.release, lease) + assert entered.wait(3) + try: + with pytest.raises(ConnectionTimeoutError): + pool.acquire() + assert pool.close() is False + seed.close.assert_not_called() + finally: + proceed.set() + releasing.result(timeout=3) + assert pool.is_closed is True + seed.close.assert_called_once_with() + + +@pytest.mark.parametrize("timeout", [30.0, 1e100]) +def test_adbc_pool_shutdown_wakes_waiting_borrower(monkeypatch: pytest.MonkeyPatch, timeout: float) -> None: + """A borrower waiting for capacity wakes on shutdown rather than its long timeout.""" + seed, lease = MagicMock(), MagicMock() + waiting = threading.Event() + native_clock = time.monotonic + worker_id: list[int] = [] + clock_reads = 0 + + def clock() -> float: + nonlocal clock_reads + if worker_id and threading.get_ident() == worker_id[0]: + clock_reads += 1 + if clock_reads == 2: + waiting.set() + return native_clock() + + def acquire() -> object: + worker_id.append(threading.get_ident()) + return pool.acquire() + + seed.adbc_clone.return_value = lease + pool = AdbcConnectionPool(lambda: seed, max_size=1, timeout=timeout) + assert pool.acquire() is lease + monkeypatch.setattr(time, "monotonic", clock) + with ThreadPoolExecutor(max_workers=1) as executor: + borrower = executor.submit(acquire) + assert waiting.wait(3) + assert pool.close() is False + with pytest.raises(DatabaseConnectionError): + borrower.result(timeout=2) + pool.release(lease) + assert pool.is_closed is True + + +def test_adbc_pool_rejects_fork_before_inherited_locks(monkeypatch: pytest.MonkeyPatch) -> None: + """Every operation refuses a changed PID before touching a poisoned lock or native handle.""" + condition = MagicMock() + condition.__enter__.side_effect = AssertionError("inherited condition accessed") + monkeypatch.setattr(threading, "Condition", lambda *_args: condition) + create, lease = MagicMock(), MagicMock() + pool = AdbcConnectionPool(create) + child_pid = os.getpid() + 1 + monkeypatch.setattr(os, "getpid", lambda: child_pid) + for operation in (pool.acquire, lambda: pool.release(lease), pool.close, lambda: pool.is_closed, pool.check_pid): + with pytest.raises(DatabaseConnectionError, match="fork"): + operation() + create.assert_not_called() + lease.close.assert_not_called() From 822c4abe54104729072f66d7ef6a826a91a00e04 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 20:37:43 +0000 Subject: [PATCH 132/174] feat(adbc): reset and drain reusable PostgreSQL leases --- sqlspec/adapters/adbc/pool.py | 52 +++++- tests/unit/adapters/test_adbc/test_pool.py | 182 ++++++++++++++++++++- 2 files changed, 226 insertions(+), 8 deletions(-) diff --git a/sqlspec/adapters/adbc/pool.py b/sqlspec/adapters/adbc/pool.py index da0951c27..8a65a1558 100644 --- a/sqlspec/adapters/adbc/pool.py +++ b/sqlspec/adapters/adbc/pool.py @@ -12,7 +12,7 @@ if TYPE_CHECKING: from collections.abc import Callable - from sqlspec.adapters.adbc._typing import AdbcConnection + from sqlspec.adapters.adbc._typing import AdbcConnection, AdbcRawCursor __all__ = ("AdbcConnectionPool",) @@ -67,6 +67,10 @@ def acquire(self) -> "AdbcConnection": if self._state != "OPEN": msg = "ADBC pool is closed to new acquisitions." raise DatabaseConnectionError(msg) + for key, (connection, state) in self._clones.items(): + if state == "idle": + self._clones[key] = (connection, "active") + return connection if len(self._clones) + self._reservations < self._max_size: self._reservations += 1 break @@ -100,13 +104,32 @@ def acquire(self) -> "AdbcConnection": return connection def release(self, connection: "AdbcConnection", *, discard: bool = False) -> None: - """Release exactly one active lease, retaining any failed native close.""" + """Reset a successful lease for reuse, or evict it after failure/shutdown.""" self.check_pid() with self._condition: entry = self._clones.get(id(connection)) if entry is None or entry[0] is not connection or entry[1] != "active": msg = "Connection is not an active lease of this ADBC pool." raise DatabaseConnectionError(msg) + reset = self._state == "OPEN" and not discard + self._clones[id(connection)] = (connection, "resetting" if reset else "closing") + try: + if reset: + self._reset_connection(connection) + elif discard: + connection.rollback() + except BaseException: + with self._condition: + self._clones[id(connection)] = (connection, "closing") + with contextlib.suppress(BaseException): + self._close_clone(connection) + self._finish_close() + raise + with self._condition: + if reset and self._state == "OPEN": + self._clones[id(connection)] = (connection, "idle") + self._condition.notify_all() + return self._clones[id(connection)] = (connection, "closing") self._close_clone(connection) self._finish_close() @@ -158,6 +181,31 @@ def check_pid(self) -> None: msg = "ADBC pools cannot be used after fork; create a new configuration in the child." raise DatabaseConnectionError(msg) + def _reset_connection(self, connection: "AdbcConnection") -> None: + cursor: AdbcRawCursor | None = None + first_error: BaseException | None = None + try: + connection.rollback() + connection.adbc_connection.set_autocommit(True) + cursor = connection.cursor() + cursor.execute("DISCARD ALL") + except BaseException as exc: + first_error = exc + finally: + if cursor is not None: + try: + cursor.close() + except BaseException as exc: + if first_error is None: + first_error = exc + try: + connection.adbc_connection.set_autocommit(False) + except BaseException as exc: + if first_error is None: + first_error = exc + if first_error is not None: + raise first_error + def _close_clone(self, connection: "AdbcConnection") -> None: try: connection.close() diff --git a/tests/unit/adapters/test_adbc/test_pool.py b/tests/unit/adapters/test_adbc/test_pool.py index fa6001fa4..731d73f03 100644 --- a/tests/unit/adapters/test_adbc/test_pool.py +++ b/tests/unit/adapters/test_adbc/test_pool.py @@ -4,7 +4,7 @@ import threading import time from concurrent.futures import ThreadPoolExecutor -from unittest.mock import MagicMock +from unittest.mock import MagicMock, call import pytest @@ -12,6 +12,172 @@ from sqlspec.exceptions import ConnectionTimeoutError, DatabaseConnectionError, ImproperConfigurationError +def test_adbc_pool_reuses_only_after_full_native_session_reset() -> None: + """A clean lease resets native session state before another borrower can reuse it.""" + seed, lease, replacement = MagicMock(), MagicMock(), MagicMock() + lease.autocommit = False + seed.adbc_clone.side_effect = [lease, replacement] + pool = AdbcConnectionPool(lambda: seed, max_size=1, timeout=0) + assert pool.acquire() is lease + pool.release(lease) + assert lease.mock_calls == [ + call.rollback(), + call.adbc_connection.set_autocommit(True), + call.cursor(), + call.cursor().execute("DISCARD ALL"), + call.cursor().close(), + call.adbc_connection.set_autocommit(False), + ] + assert lease.autocommit is False + assert pool.acquire() is lease + pool.release(lease, discard=True) + assert pool.acquire() is replacement + pool.release(replacement) + assert pool.close() is True + + +@pytest.mark.parametrize("stage", ["rollback", "enable", "cursor", "discard", "cursor_close", "restore"]) +@pytest.mark.parametrize("close_fails", [False, True]) +def test_adbc_pool_failed_reset_evicts_and_preserves_primary(stage: str, close_fails: bool) -> None: + """Each reset failure evicts; a secondary native-close failure stays retryable.""" + seed, lease, replacement = MagicMock(), MagicMock(), MagicMock() + primary = RuntimeError(f"reset {stage} failed") + seed.adbc_clone.side_effect = [lease, replacement] + cursor = lease.cursor.return_value + handle = lease.adbc_connection + if stage == "rollback": + lease.rollback.side_effect = primary + elif stage == "enable": + handle.set_autocommit.side_effect = [primary, None] + elif stage == "cursor": + lease.cursor.side_effect = primary + elif stage == "discard": + cursor.execute.side_effect = primary + elif stage == "cursor_close": + cursor.close.side_effect = primary + else: + handle.set_autocommit.side_effect = [None, primary] + if close_fails: + lease.close.side_effect = [KeyboardInterrupt("close interrupted"), None] + pool = AdbcConnectionPool(lambda: seed, max_size=1, timeout=0) + assert pool.acquire() is lease + with pytest.raises(RuntimeError) as caught: + pool.release(lease) + assert caught.value is primary + assert handle.set_autocommit.call_args == call(False) + if stage in {"discard", "cursor_close", "restore"}: + cursor.close.assert_called_once_with() + if close_fails: + with pytest.raises(ConnectionTimeoutError): + pool.acquire() + seed.close.assert_not_called() + else: + assert pool.acquire() is replacement + pool.release(replacement) + assert pool.close() is True + assert lease.close.call_count == (2 if close_fails else 1) + + +def test_adbc_pool_reset_primary_survives_all_secondary_cleanup_failures() -> None: + """Cursor, autocommit and native-close cleanup cannot replace a reset exception.""" + seed, lease = MagicMock(), MagicMock() + primary = KeyboardInterrupt("DISCARD interrupted") + cleanup = RuntimeError("cleanup failed") + seed.adbc_clone.return_value = lease + lease.cursor.return_value.execute.side_effect = primary + lease.cursor.return_value.close.side_effect = cleanup + lease.adbc_connection.set_autocommit.side_effect = [None, cleanup] + lease.close.side_effect = [cleanup, None] + pool = AdbcConnectionPool(lambda: seed, max_size=1, timeout=0) + assert pool.acquire() is lease + with pytest.raises(KeyboardInterrupt) as caught: + pool.release(lease) + assert caught.value is primary + lease.cursor.return_value.close.assert_called_once_with() + assert lease.adbc_connection.set_autocommit.call_args == call(False) + seed.close.assert_not_called() + assert pool.close() is True + + +@pytest.mark.parametrize("rollback_fails", [False, True]) +def test_adbc_pool_discard_rolls_back_and_closes_without_reset(rollback_fails: bool) -> None: + """Body-failure disposal attempts rollback and close without DISCARD ALL or reuse.""" + seed, lease, replacement = MagicMock(), MagicMock(), MagicMock() + primary = RuntimeError("rollback failed") + seed.adbc_clone.side_effect = [lease, replacement] + if rollback_fails: + lease.rollback.side_effect = primary + lease.close.side_effect = [KeyboardInterrupt("close interrupted"), None] + pool = AdbcConnectionPool(lambda: seed, max_size=1, timeout=0) + assert pool.acquire() is lease + if rollback_fails: + with pytest.raises(RuntimeError) as caught: + pool.release(lease, discard=True) + assert caught.value is primary + with pytest.raises(ConnectionTimeoutError): + pool.acquire() + else: + pool.release(lease, discard=True) + assert pool.acquire() is replacement + pool.release(replacement) + lease.rollback.assert_called_once_with() + lease.cursor.assert_not_called() + lease.adbc_connection.set_autocommit.assert_not_called() + assert pool.close() is True + + +@pytest.mark.parametrize("mode", ["reuse", "shutdown", "shutdown_close_failure"]) +def test_adbc_pool_reset_retains_capacity_and_obeys_shutdown(mode: str) -> None: + """A resetting lease stays owned until completion and never becomes idle after shutdown.""" + seed, lease = MagicMock(), MagicMock() + entered, proceed = threading.Event(), threading.Event() + cleanup_error = RuntimeError("shutdown close failed") + shutdown = mode != "reuse" + close_fails = mode == "shutdown_close_failure" + + def discard(_sql: str) -> None: + entered.set() + assert proceed.wait(3) + + seed.adbc_clone.return_value = lease + lease.cursor.return_value.execute.side_effect = discard + if close_fails: + lease.close.side_effect = [cleanup_error, None] + pool = AdbcConnectionPool(lambda: seed, max_size=1, timeout=0) + assert pool.acquire() is lease + with ThreadPoolExecutor(max_workers=1) as executor: + resetting = executor.submit(pool.release, lease) + assert entered.wait(3) + try: + with pytest.raises(ConnectionTimeoutError): + pool.acquire() + if shutdown: + assert pool.close() is False + seed.close.assert_not_called() + finally: + proceed.set() + if close_fails: + with pytest.raises(RuntimeError) as caught: + resetting.result(timeout=3) + assert caught.value is cleanup_error + assert pool.is_closed is False + seed.close.assert_not_called() + assert pool.close() is True + else: + resetting.result(timeout=3) + if shutdown: + assert pool.is_closed is True + with pytest.raises(DatabaseConnectionError): + pool.acquire() + else: + assert pool.acquire() is lease + pool.release(lease) + assert pool.close() is True + seed.adbc_clone.assert_called_once_with() + assert lease.close.call_count == (2 if close_fails else 1) + seed.close.assert_called_once_with() + + def test_adbc_pool_exclusive_original_seed_leases_and_nested_timeout() -> None: """Every borrower owns a distinct original-seed clone within the configured bound.""" seed, first, second, third = MagicMock(), MagicMock(), MagicMock(), MagicMock() @@ -23,10 +189,10 @@ def test_adbc_pool_exclusive_original_seed_leases_and_nested_timeout() -> None: assert pool.acquire() is second with pytest.raises(ConnectionTimeoutError): pool.acquire() - pool.release(first) + pool.release(first, discard=True) assert pool.acquire() is third for connection in (second, third): - pool.release(connection) + pool.release(connection, discard=True) create.assert_called_once_with() assert seed.adbc_clone.call_count == 3 first.close.assert_called_once_with() @@ -83,8 +249,8 @@ def test_adbc_pool_rejects_foreign_and_duplicate_release() -> None: pool.release(lease) with pytest.raises(DatabaseConnectionError): pool.release(lease) - lease.close.assert_called_once_with() assert pool.close() is True + lease.close.assert_called_once_with() def test_adbc_pool_shutdown_drains_active_lease_without_forcing_close() -> None: @@ -115,7 +281,7 @@ def test_adbc_pool_quarantined_close_retains_capacity_and_retries_seed_last() -> pool = AdbcConnectionPool(lambda: seed, max_size=1, timeout=0) assert pool.acquire() is lease with pytest.raises(RuntimeError) as caught: - pool.release(lease) + pool.release(lease, discard=True) assert caught.value is failure with pytest.raises(ConnectionTimeoutError): pool.acquire() @@ -190,6 +356,10 @@ def clone() -> MagicMock: seed.adbc_clone.side_effect = clone seed.close.side_effect = inspect_pool lease.close.side_effect = inspect_pool + lease.rollback.side_effect = inspect_pool + lease.adbc_connection.set_autocommit.side_effect = lambda _enabled: inspect_pool() + lease.cursor.return_value.execute.side_effect = lambda _sql: inspect_pool() + lease.cursor.return_value.close.side_effect = inspect_pool pool = AdbcConnectionPool(create) assert pool.acquire() is lease pool.release(lease) @@ -249,7 +419,7 @@ def native_close() -> None: pool = AdbcConnectionPool(lambda: seed, max_size=1, timeout=0) assert pool.acquire() is lease with ThreadPoolExecutor(max_workers=1) as executor: - releasing = executor.submit(pool.release, lease) + releasing = executor.submit(pool.release, lease, discard=True) assert entered.wait(3) try: with pytest.raises(ConnectionTimeoutError): From af4084284c1f203c45424070ff1c3f27f4ca0166 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 20:46:10 +0000 Subject: [PATCH 133/174] fix(adbc): close lazy stream resources before session release --- sqlspec/adapters/adbc/_typing.py | 23 +- sqlspec/adapters/adbc/driver.py | 151 ++++++++++--- .../test_adbc/test_arrow_streaming.py | 206 +++++++++++++++++- 3 files changed, 333 insertions(+), 47 deletions(-) diff --git a/sqlspec/adapters/adbc/_typing.py b/sqlspec/adapters/adbc/_typing.py index 04e8604ba..d9eba6804 100644 --- a/sqlspec/adapters/adbc/_typing.py +++ b/sqlspec/adapters/adbc/_typing.py @@ -121,10 +121,16 @@ def __enter__(self) -> "AdbcDriver": return self._prepare_driver(self._driver) except BaseException as exc: connection = self._connection + driver = self._driver self._connection = None self._driver = None + if driver is not None: + with contextlib.suppress(BaseException): + driver.close_streams() with contextlib.suppress(BaseException): - self._release_connection(connection, exc_type=type(exc), exc_val=exc, exc_tb=exc.__traceback__) + self._release_connection( + connection, exc_type=type(exc), exc_val=exc, exc_tb=exc.__traceback__, driver=driver + ) raise def __exit__( @@ -132,9 +138,22 @@ def __exit__( ) -> "bool | None": if self._connection is not None: connection = self._connection + driver = self._driver self._connection = None self._driver = None - self._release_connection(connection, exc_type=exc_type, exc_val=exc_val, exc_tb=exc_tb) + first_error: BaseException | None = None + if driver is not None: + try: + driver.close_streams() + except BaseException as exc: + first_error = exc + try: + self._release_connection(connection, exc_type=exc_type, exc_val=exc_val, exc_tb=exc_tb, driver=driver) + except BaseException as exc: + if first_error is None: + first_error = exc + if exc_type is None and first_error is not None: + raise first_error return None diff --git a/sqlspec/adapters/adbc/driver.py b/sqlspec/adapters/adbc/driver.py index 3649990d3..1596e6b5e 100644 --- a/sqlspec/adapters/adbc/driver.py +++ b/sqlspec/adapters/adbc/driver.py @@ -98,14 +98,14 @@ def _handle_exception(self, exc_type: "type[BaseException] | None", exc_val: "Ba class AdbcSelectStreamSource: """Native ADBC chunk source backed by a ``RecordBatchReader``.""" - __slots__ = ("_chunk_size", "_cursor_manager", "_driver", "_parameters", "_reader", "_remainder", "_sql") + __slots__ = ("_chunk_size", "_driver", "_owner", "_parameters", "_reader", "_remainder", "_sql") def __init__(self, driver: "AdbcDriver", sql: str, parameters: Any, chunk_size: int = 1000) -> None: self._driver = driver self._sql = sql self._parameters = parameters self._chunk_size = chunk_size - self._cursor_manager: AdbcCursor | None = None + self._owner: _AdbcArrowStreamOwner | None = None self._reader: Any = None self._remainder: list[dict[str, Any]] = [] @@ -136,13 +136,19 @@ def _stream_op() -> Any: self._driver._check_pending_exception(handler) except BaseException: if cursor is not None: - handle_postgres_rollback(self._driver._dialect_name, cursor, logger) - with contextlib.suppress(Exception): + with contextlib.suppress(BaseException): + handle_postgres_rollback(self._driver._dialect_name, cursor, logger) + with contextlib.suppress(BaseException): cursor_manager.__exit__(None, None, None) raise - self._cursor_manager = cursor_manager - self._reader = iter(reader) + self._owner = _AdbcArrowStreamOwner(reader, cursor) + try: + self._reader = iter(reader) + except BaseException: + with contextlib.suppress(BaseException): + self.close() + raise def fetch_chunk(self) -> "list[dict[str, Any]]": chunk_size = self._chunk_size @@ -192,11 +198,13 @@ def fetch_chunk(self) -> "list[dict[str, Any]]": def close(self, error: bool = False) -> None: self._reader = None - cursor_manager = self._cursor_manager - self._cursor_manager = None - if cursor_manager is not None: - with contextlib.suppress(Exception): - cursor_manager.__exit__(None, None, None) + self._remainder = [] + if self._owner is not None: + if error: + with contextlib.suppress(BaseException): + self._owner.close() + else: + self._owner.close() @final @@ -212,10 +220,12 @@ class AdbcDriver(SyncDriverAdapterBase): "_data_dictionary", "_detected_dialect", "_dialect_name", + "_has_streams", "_is_flightsql", "_is_postgres", "_json_serializer", "_previous_autocommit", + "_stream_closers", "_transaction_active", "dialect", ) @@ -258,6 +268,26 @@ def __init__( self._column_name_cache: dict[int, tuple[Any, list[str]]] = {} self._transaction_active = False self._previous_autocommit: bool | None = None + self._has_streams = False + self._stream_closers: list[Callable[[], None]] = [] + + @property + def has_streams(self) -> bool: + """Return whether this session created any lazy native stream.""" + return self._has_streams + + def close_streams(self) -> None: + """Close every registered native owner before returning a session connection.""" + closers, self._stream_closers = self._stream_closers, [] + first_error: BaseException | None = None + for close in closers: + try: + close() + except BaseException as exc: + if first_error is None: + first_error = exc + if first_error is not None: + raise first_error @property def data_dictionary(self) -> "AdbcDataDictionary": @@ -462,10 +492,14 @@ def dispatch_execute_script(self, cursor: "AdbcRawCursor", statement: "SQL") -> def dispatch_select_stream(self, statement: "SQL", chunk_size: int) -> "SyncRowStream[dict[str, Any]] | None": """Return a native ADBC stream backed by ``fetch_record_batch()``.""" + self._has_streams = True if not statement.returns_rows(): return None sql, prepared_parameters = self._compiled_sql(statement, self.statement_config) - return SyncRowStream(AdbcSelectStreamSource(self, sql, prepared_parameters, chunk_size=chunk_size)) + source = AdbcSelectStreamSource(self, sql, prepared_parameters, chunk_size=chunk_size) + stream: SyncRowStream[dict[str, Any]] = SyncRowStream(source) + self._stream_closers.extend((source.close, stream.close)) + return stream def begin(self) -> None: """Begin database transaction. @@ -640,6 +674,8 @@ def select_to_arrow( arrow_result: ArrowResult | None = None if return_format in {"reader", "batches"}: + if return_format == "reader": + self._has_streams = True cursor_manager = self.with_cursor(self.connection) cursor = None reader: Any = None @@ -663,9 +699,11 @@ def select_to_arrow( ) else: reader = fetch_record_batch() - except Exception: - handle_postgres_rollback(self._dialect_name, cursor, logger) - cursor_manager.__exit__(None, None, None) + except BaseException: + with contextlib.suppress(BaseException): + handle_postgres_rollback(self._dialect_name, cursor, logger) + with contextlib.suppress(BaseException): + cursor_manager.__exit__(None, None, None) cursor = None raise @@ -682,33 +720,39 @@ def select_to_arrow( cursor_manager.__exit__(None, None, None) msg = "ADBC did not return an Arrow reader." raise SQLSpecError(msg) - if bool(self.driver_features.get("enable_lowercase_column_names", False)) and hasattr(reader, "schema"): - raw_names = list(reader.schema.names) - normalized_names = normalize_column_names(raw_names, self.driver_features) - pa = import_optional("pyarrow") - if normalized_names != raw_names and pa is not None: - new_schema = pa.schema([ - field.with_name(name) for field, name in zip(reader.schema, normalized_names, strict=True) - ]) - source_reader = reader - - def _lower_batches() -> Any: - for batch in source_reader: - yield batch.rename_columns(normalized_names) - - reader = pa.RecordBatchReader.from_batches(new_schema, _lower_batches()) + owner = _AdbcArrowStreamOwner(reader, cursor) + if return_format == "reader": + self._stream_closers.append(owner.close) try: - return build_arrow_result_from_reader( + if bool(self.driver_features.get("enable_lowercase_column_names", False)) and hasattr(reader, "schema"): + raw_names = list(reader.schema.names) + normalized_names = normalize_column_names(raw_names, self.driver_features) + pa = import_optional("pyarrow") + if normalized_names != raw_names and pa is not None: + new_schema = pa.schema([ + field.with_name(name) for field, name in zip(reader.schema, normalized_names, strict=True) + ]) + source_reader = reader + + def _lower_batches() -> Any: + for batch in source_reader: + yield batch.rename_columns(normalized_names) + + reader = pa.RecordBatchReader.from_batches(new_schema, _lower_batches()) + result = build_arrow_result_from_reader( prepared_statement, - arrow_reader_with_deferred_close(reader, cursor.close), + arrow_reader_with_deferred_close(reader, owner.finish), return_format=return_format, batch_size=batch_size, arrow_schema=arrow_schema, ) + if return_format == "batches": + owner.close() except BaseException: - with contextlib.suppress(Exception): - cursor_manager.__exit__(None, None, None) + with contextlib.suppress(BaseException): + owner.close() raise + return result with self.with_cursor(self.connection) as cursor, exc_handler: if cursor is None: @@ -720,8 +764,9 @@ def _lower_batches() -> Any: try: cursor.execute(sql, parameters=execute_parameters) arrow_table = cursor.fetch_arrow_table() - except Exception: - handle_postgres_rollback(self._dialect_name, cursor, logger) + except BaseException: + with contextlib.suppress(BaseException): + handle_postgres_rollback(self._dialect_name, cursor, logger) raise if bool(self.driver_features.get("enable_lowercase_column_names", False)) and hasattr( @@ -1075,6 +1120,40 @@ def _parse_adbc_ingest_target( return table, kwargs +class _AdbcArrowStreamOwner: + """Close actual native owners independently of lazy wrapper iteration.""" + + __slots__ = ("_close_error", "_closed", "_cursor", "_reader") + + def __init__(self, reader: Any, cursor: "AdbcRawCursor") -> None: + self._reader = reader + self._cursor = cursor + self._closed = False + self._close_error: BaseException | None = None + + def close(self) -> None: + if not self._closed: + self._closed = True + try: + close_reader = getattr(self._reader, "close", None) + if callable(close_reader): + close_reader() + except BaseException as exc: + self._close_error = exc + try: + self._cursor.close() + except BaseException as exc: + if self._close_error is None: + self._close_error = exc + if self._close_error is not None: + raise self._close_error + + def finish(self) -> None: + """Finalize lazy iteration without replacing its exception or raising from GC.""" + with contextlib.suppress(BaseException): + self.close() + + def _fetch_record_batch(cursor: Any) -> Any: fetch_record_batch = getattr(cursor, "fetch_record_batch", None) if fetch_record_batch is None: diff --git a/tests/unit/adapters/test_adbc/test_arrow_streaming.py b/tests/unit/adapters/test_adbc/test_arrow_streaming.py index 7d97784b4..801e3e5c0 100644 --- a/tests/unit/adapters/test_adbc/test_arrow_streaming.py +++ b/tests/unit/adapters/test_adbc/test_arrow_streaming.py @@ -6,11 +6,14 @@ import pytest from adbc_driver_manager import AdbcStatusCode, DatabaseError +from sqlspec.adapters.adbc._typing import AdbcSessionContext +from sqlspec.adapters.adbc.core import get_statement_config from sqlspec.adapters.adbc.driver import AdbcDriver, AdbcSelectStreamSource from sqlspec.exceptions import SQLSpecError if TYPE_CHECKING: from sqlspec.adapters.adbc._typing import AdbcConnection + from sqlspec.driver import SyncRowStream _STORAGE_CAPABILITIES = { @@ -22,18 +25,44 @@ } +class _AdbcStreamingReader: + def __init__(self) -> None: + table = pa.table({"x": [1, 2, 3]}) + self.schema = table.schema + self._reader = pa.RecordBatchReader.from_batches(table.schema, table.to_batches(max_chunksize=2)) + self.close_count = 0 + self.close_error: BaseException | None = None + + def __iter__(self) -> "_AdbcStreamingReader": + return self + + def __next__(self) -> pa.RecordBatch: + return self._reader.read_next_batch() + + def read_next_batch(self) -> pa.RecordBatch: + return next(self) + + def close(self) -> None: + self.close_count += 1 + self._reader.close() + if self.close_error is not None: + raise self.close_error + + class _AdbcStreamingCursor: def __init__(self) -> None: self.closed = False + self.close_count = 0 + self.close_error: BaseException | None = None + self.reader_obj = _AdbcStreamingReader() self.executed: list[tuple[str, object]] = [] self.ingest_calls: list[tuple[str, object, str, dict[str, Any]]] = [] def execute(self, sql: str, parameters: object = None) -> None: self.executed.append((sql, parameters)) - def fetch_record_batch(self) -> pa.RecordBatchReader: - table = pa.table({"x": [1, 2, 3]}) - return pa.RecordBatchReader.from_batches(table.schema, table.to_batches(max_chunksize=2)) + def fetch_record_batch(self) -> _AdbcStreamingReader: + return self.reader_obj def fetch_arrow_table(self) -> pa.Table: msg = "reader formats should use fetch_record_batch(), not fetch_arrow_table()" @@ -45,6 +74,9 @@ def adbc_ingest(self, table: str, source: object, *, mode: str, **kwargs: Any) - def close(self) -> None: self.closed = True + self.close_count += 1 + if self.close_error is not None: + raise self.close_error class _AdbcStreamingConnection: @@ -58,6 +90,158 @@ def cursor(self) -> _AdbcStreamingCursor: return self.cursor_obj +@pytest.mark.parametrize("kind", ["row", "reader"]) +@pytest.mark.parametrize("progress", ["never", "partial", "exhausted"]) +def test_session_closes_lazy_stream_owners_before_release(kind: str, progress: str) -> None: + """All lazy streams taint their session and close actual native owners before release.""" + connection = _AdbcStreamingConnection() + released: list[AdbcDriver] = [] + + def release(_connection: Any, **kwargs: Any) -> None: + expected_closes = 0 if kind == "row" and progress == "never" else 1 + assert connection.cursor_obj.close_count == expected_closes + assert connection.cursor_obj.reader_obj.close_count == expected_closes + assert kwargs["driver"].has_streams is True + released.append(kwargs["driver"]) + + context = AdbcSessionContext( + acquire_connection=lambda: connection, + release_connection=release, + statement_config=get_statement_config("sqlite"), + driver_features={}, + prepare_driver=lambda driver: driver, + ) + stream: SyncRowStream[dict[str, Any]] | None = None + with context as driver: + if kind == "row": + stream = driver.select_stream("SELECT 1 AS x", native_only=True, chunk_size=1) + if progress == "partial": + assert next(stream) == {"x": 1} + elif progress == "exhausted": + assert list(stream) == [{"x": 1}, {"x": 2}, {"x": 3}] + else: + reader = driver.select_to_arrow("SELECT 1 AS x", return_format="reader").data + if progress == "partial": + assert reader.read_next_batch().to_pydict() == {"x": [1, 2]} + elif progress == "exhausted": + assert reader.read_all().to_pydict() == {"x": [1, 2, 3]} + assert released == [driver] + driver.close_streams() + if stream is not None: + assert list(stream) == [] + + +def test_lazy_fallback_stream_taints_session_before_execution() -> None: + """The generic lazy fallback also prevents reuse before it opens native resources.""" + connection = _AdbcStreamingConnection() + released: list[AdbcDriver] = [] + + def release(_connection: Any, **kwargs: Any) -> None: + assert kwargs["driver"].has_streams is True + assert connection.cursor_obj.executed == [] + assert connection.cursor_obj.close_count == 0 + released.append(kwargs["driver"]) + + context = AdbcSessionContext( + acquire_connection=lambda: connection, + release_connection=release, + statement_config=get_statement_config("sqlite"), + driver_features={}, + prepare_driver=lambda driver: driver, + ) + with context as driver: + stream = driver.select_stream("UPDATE items SET x = 1") + assert released == [driver] + stream.close() + + +@pytest.mark.parametrize("stage", ["exit", "body", "setup"]) +def test_session_stream_cleanup_attempts_all_owners_and_preserves_primary(stage: str) -> None: + """Cleanup attempts every native reader/cursor and release, preserving setup/body errors.""" + cursors = [_AdbcStreamingCursor(), _AdbcStreamingCursor()] + available = iter(cursors) + + class _MultipleStreamingConnection(_AdbcStreamingConnection): + def cursor(self) -> _AdbcStreamingCursor: + return next(available) + + connection = _MultipleStreamingConnection() + primary = KeyboardInterrupt("primary interruption") + first_cleanup = RuntimeError("first reader cleanup failed") + cursors[0].reader_obj.close_error = first_cleanup + cursors[0].close_error = RuntimeError("cursor cleanup failed") + cursors[1].reader_obj.close_error = SystemExit("second reader cleanup failed") + released: list[AdbcDriver] = [] + drivers: list[AdbcDriver] = [] + readers: list[Any] = [] + + def open_readers(driver: AdbcDriver) -> None: + drivers.append(driver) + readers.extend(driver.select_to_arrow("SELECT 1 AS x", return_format="reader").data for _ in range(2)) + + def prepare(driver: AdbcDriver) -> AdbcDriver: + if stage == "setup": + open_readers(driver) + raise primary + return driver + + def release(_connection: Any, **kwargs: Any) -> None: + for cursor in cursors: + assert cursor.close_count == 1 + assert cursor.reader_obj.close_count == 1 + released.append(kwargs["driver"]) + raise RuntimeError("release cleanup failed") + + context = AdbcSessionContext( + acquire_connection=lambda: connection, + release_connection=release, + statement_config=get_statement_config("sqlite"), + driver_features={}, + prepare_driver=prepare, + ) + expected = RuntimeError if stage == "exit" else KeyboardInterrupt + with pytest.raises(expected) as caught, context as driver: + open_readers(driver) + if stage == "body": + raise primary + assert caught.value is (first_cleanup if stage == "exit" else primary) + assert released == drivers + context.__exit__(None, None, None) + drivers[0].close_streams() + assert [cursor.close_count for cursor in cursors] == [1, 1] + + +def test_eager_arrow_batches_close_owners_without_tainting_session() -> None: + """Eager batches release their native resources immediately and remain reusable.""" + connection = _AdbcStreamingConnection() + driver = AdbcDriver(cast("AdbcConnection", connection), dialect="sqlite") + result = driver.select_to_arrow("SELECT 1 AS x", return_format="batches") + assert [batch.to_pydict() for batch in result.data] == [{"x": [1, 2]}, {"x": [3]}] + assert driver.has_streams is False + assert connection.cursor_obj.close_count == 1 + assert connection.cursor_obj.reader_obj.close_count == 1 + driver.close_streams() + assert connection.cursor_obj.close_count == 1 + + +def test_materialized_arrow_table_does_not_taint_session() -> None: + """Materialized tables close their cursor and do not require stream eviction.""" + + class _MaterializedCursor(_AdbcStreamingCursor): + def fetch_arrow_table(self) -> pa.Table: + return pa.table({"x": [1, 2, 3]}) + + connection = _AdbcStreamingConnection() + connection.cursor_obj = _MaterializedCursor() + driver = AdbcDriver(cast("AdbcConnection", connection), dialect="sqlite") + result = driver.select_to_arrow("SELECT 1 AS x", return_format="table") + assert result.data.to_pydict() == {"x": [1, 2, 3]} + assert driver.has_streams is False + assert connection.cursor_obj.close_count == 1 + driver.close_streams() + assert connection.cursor_obj.close_count == 1 + + def test_select_to_arrow_reader_uses_fetch_record_batch_and_defers_cursor_close() -> None: connection = _AdbcStreamingConnection() driver = AdbcDriver(cast("AdbcConnection", connection), dialect="sqlite") @@ -104,15 +288,16 @@ def test_load_from_arrow_passes_record_batch_reader_to_adbc_ingest() -> None: assert job.telemetry["destination"] == "target_table" -def test_select_to_arrow_postgres_rolls_back_on_error() -> None: +@pytest.mark.parametrize("return_format", ["table", "reader"]) +@pytest.mark.parametrize("primary", [RuntimeError("postgres failure"), KeyboardInterrupt("query interrupted")]) +def test_select_to_arrow_postgres_rolls_back_on_error(return_format: Any, primary: BaseException) -> None: class _PostgresErrorCursor: def __init__(self, conn: "_PostgresErrorConnection") -> None: self.connection = conn self.closed = False def execute(self, sql: str, parameters: object = None) -> None: - msg = "postgres failure" - raise RuntimeError(msg) + raise primary def close(self) -> None: self.closed = True @@ -120,6 +305,7 @@ def close(self) -> None: class _PostgresErrorConnection: def __init__(self) -> None: self.rolled_back = False + self.cursor_obj = _PostgresErrorCursor(self) def adbc_get_info(self) -> dict[str, str]: return {"vendor_name": "PostgreSQL", "driver_name": "adbc_driver_postgresql"} @@ -128,14 +314,16 @@ def rollback(self) -> None: self.rolled_back = True def cursor(self) -> _PostgresErrorCursor: - return _PostgresErrorCursor(self) + return self.cursor_obj connection = _PostgresErrorConnection() driver = AdbcDriver(cast("AdbcConnection", connection), dialect="postgres") - with pytest.raises((RuntimeError, SQLSpecError)): - driver.select_to_arrow("SELECT 1", return_format="table") + with pytest.raises(type(primary)) as caught: + driver.select_to_arrow("SELECT 1", return_format=return_format) + assert caught.value is primary assert connection.rolled_back is True + assert connection.cursor_obj.closed is True def test_select_stream_start_closes_cursor_and_rolls_back_postgres_on_failure() -> None: From 0dc22b0316c122697a216e8c0c58f1dbff087097 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 21:44:00 +0000 Subject: [PATCH 134/174] feat(adbc): expose opt-in PostgreSQL pool configuration --- sqlspec/adapters/adbc/__init__.py | 3 +- sqlspec/adapters/adbc/config.py | 179 ++++++++- .../test_adbc_config_normalization.py | 364 ++++++++++++++++++ .../test_adbc/test_arrow_streaming.py | 49 +++ 4 files changed, 584 insertions(+), 11 deletions(-) diff --git a/sqlspec/adapters/adbc/__init__.py b/sqlspec/adapters/adbc/__init__.py index 34c638acd..3eefd962b 100644 --- a/sqlspec/adapters/adbc/__init__.py +++ b/sqlspec/adapters/adbc/__init__.py @@ -1,7 +1,7 @@ """ADBC adapter for SQLSpec.""" from sqlspec.adapters.adbc._typing import AdbcConnection -from sqlspec.adapters.adbc.config import AdbcConfig, AdbcConnectionParams, AdbcDriverFeatures +from sqlspec.adapters.adbc.config import AdbcConfig, AdbcConnectionParams, AdbcDriverFeatures, AdbcPoolConfig from sqlspec.adapters.adbc.driver import AdbcCursor, AdbcDriver, AdbcExceptionHandler __all__ = ( @@ -12,4 +12,5 @@ "AdbcDriver", "AdbcDriverFeatures", "AdbcExceptionHandler", + "AdbcPoolConfig", ) diff --git a/sqlspec/adapters/adbc/config.py b/sqlspec/adapters/adbc/config.py index 32dd09d45..d984fa3c2 100644 --- a/sqlspec/adapters/adbc/config.py +++ b/sqlspec/adapters/adbc/config.py @@ -1,9 +1,12 @@ """ADBC database configuration.""" import contextlib +import importlib.metadata +import math +import re from collections.abc import Callable from pathlib import Path -from typing import TYPE_CHECKING, Any, ClassVar, Literal, NotRequired, TypedDict, cast +from typing import TYPE_CHECKING, Any, ClassVar, Literal, NotRequired, Self, TypedDict, cast from sqlspec.adapters.adbc._typing import AdbcConnection, AdbcCursor, AdbcSessionContext from sqlspec.adapters.adbc.core import ( @@ -19,9 +22,11 @@ resolve_dialect_from_config, resolve_dialect_name, resolve_driver_connect_func, + resolve_driver_name_from_config, ) from sqlspec.adapters.adbc.driver import AdbcDriver, AdbcExceptionHandler -from sqlspec.config import ExtensionConfigs, NoPoolSyncConfig +from sqlspec.adapters.adbc.pool import AdbcConnectionPool +from sqlspec.config import ExtensionConfigs, SyncDatabaseConfig from sqlspec.core import StatementConfig from sqlspec.core.capabilities import TypeCoercionCapabilities from sqlspec.core.config_runtime import ( @@ -31,7 +36,7 @@ resolve_runtime_statement_config, ) from sqlspec.driver import SyncPoolConnectionContext, SyncPoolSessionFactory -from sqlspec.exceptions import ImproperConfigurationError +from sqlspec.exceptions import DatabaseConnectionError, ImproperConfigurationError from sqlspec.extensions.events import EventRuntimeHints from sqlspec.utils.config_tools import normalize_connection_config @@ -41,7 +46,14 @@ from sqlspec.observability import ObservabilityConfig -__all__ = ("AdbcConfig", "AdbcConnectionParams", "AdbcDriverFeatures") +__all__ = ("AdbcConfig", "AdbcConnectionParams", "AdbcDriverFeatures", "AdbcPoolConfig") + + +class AdbcPoolConfig(TypedDict): + """Optional PostgreSQL clone-pool bounds; an empty mapping enables defaults.""" + + max_size: NotRequired[int] + timeout: NotRequired[float] class AdbcConnectionParams(TypedDict): @@ -194,20 +206,24 @@ def __exit__( class _AdbcSessionConnectionHandler(SyncPoolSessionFactory): - __slots__ = ("_connection", "_driver_features", "_profile_path") + __slots__ = ("_connection", "_driver_features", "_origin_pool", "_profile_path") def __init__(self, config: "AdbcConfig") -> None: super().__init__(config) self._connection: AdbcConnection | None = None self._driver_features: dict[str, Any] = dict(config.driver_features) self._profile_path: Path | None = None + self._origin_pool: AdbcConnectionPool | None = None @property def driver_features(self) -> "dict[str, Any]": return self._driver_features def acquire_connection(self) -> "AdbcConnection": - connection: AdbcConnection = self._config.create_connection() + self._origin_pool = self._config.provide_pool() + connection: AdbcConnection = ( + self._origin_pool.acquire() if self._origin_pool is not None else self._config.create_connection() + ) self._connection = connection try: if self._config._resolved_dialect == "duckdb" and bool( @@ -226,12 +242,17 @@ def release_connection(self, _conn: "AdbcConnection", **kwargs: Any) -> None: if self._connection is None: return connection = self._connection + pool, self._origin_pool = self._origin_pool, None profile_path = self._profile_path self._connection = None self._profile_path = None raw_path = self._driver_features.pop("_duckdb_profile_path", None) try: - if connection is not self._config.connection_instance: + if pool is not None: + driver = kwargs.get("driver") + discard = kwargs.get("exc_type") is not None or bool(driver is not None and driver.has_streams) + pool.release(connection, discard=discard) + elif connection is not self._config.connection_instance: connection.close() finally: try: @@ -241,7 +262,7 @@ def release_connection(self, _conn: "AdbcConnection", **kwargs: Any) -> None: cleanup_duckdb_profile_path(Path(raw_path)) -class AdbcConfig(NoPoolSyncConfig[AdbcConnection, AdbcDriver]): +class AdbcConfig(SyncDatabaseConfig[AdbcConnection, AdbcConnectionPool | AdbcConnection | None, AdbcDriver]): """ADBC configuration for Arrow Database Connectivity. ADBC provides an interface for connecting to multiple database systems @@ -274,6 +295,8 @@ class AdbcConfig(NoPoolSyncConfig[AdbcConnection, AdbcDriver]): "_paradedb_available", "_pg_textsearch_available", "_pgvector_available", + "_pool_close_started", + "_pool_config", "_resolved_dialect", "_user_connection_hook", ) @@ -283,6 +306,7 @@ def __init__( *, connection_config: "AdbcConnectionParams | dict[str, Any] | None" = None, connection_instance: "AdbcConnection | None" = None, + pool_config: "AdbcPoolConfig | None" = None, migration_config: "dict[str, Any] | None" = None, statement_config: StatementConfig | None = None, driver_features: "AdbcDriverFeatures | dict[str, Any] | None" = None, @@ -296,6 +320,7 @@ def __init__( Args: connection_config: Connection configuration parameters connection_instance: Pre-created connection instance to use instead of creating new one + pool_config: Optional PostgreSQL clone pooling; None preserves unpooled ownership. migration_config: Migration configuration statement_config: Default SQL statement configuration driver_features: Driver feature configuration (AdbcDriverFeatures) @@ -323,6 +348,9 @@ def __init__( "Callable[[AdbcConnection], None] | None", features_dict.pop("on_connection_create", None) ) self._default_session_config = get_statement_config(self._resolved_dialect) + self._pool_config = cast("AdbcPoolConfig", dict(pool_config)) if pool_config is not None else None + self._pool_close_started = False + self._validate_pool_settings(connection_instance) super().__init__( connection_config=self.connection_config, @@ -336,6 +364,135 @@ def __init__( **kwargs, ) + def create_pool(self) -> "AdbcConnectionPool | None": + """Publish one lazy generation, rejecting a drain until every native owner closes.""" + if self._pool_config is None: + return None + while True: + existing = self.connection_instance + if isinstance(existing, AdbcConnectionPool): + existing.check_pid() + with self._pool_lock: + if self.connection_instance is not existing: + continue + if existing is not None and not isinstance(existing, AdbcConnectionPool): + msg = "A pooled ADBC configuration cannot own an injected raw connection." + raise ImproperConfigurationError(msg) + if existing is not None and not existing.is_closed: + if self._pool_close_started: + msg = "ADBC pool is draining; retry shutdown before starting another generation." + raise DatabaseConnectionError(msg) + return existing + if existing is not None: + self.connection_instance = None + self._pool_close_started = False + created = None + else: + created = self._create_pool() + self.connection_instance = created + self._pool_close_started = False + runtime = self.get_observability_runtime() + if existing is not None: + runtime.emit_pool_destroy_sync(existing) + continue + runtime.emit_pool_create_sync(created) + return created + + def provide_pool(self, *args: Any, **kwargs: Any) -> "AdbcConnectionPool | None": + """Return the current enabled pool generation, or None when pooling is disabled.""" + return self.create_pool() + + def close_pool(self) -> None: + """Drain an owned generation, retaining native cleanup failures for retry.""" + self._shutdown_pool(retire=True) + + def derive_pool(self, **overrides: Any) -> Self: + """Derive independent ownership while preserving explicit pool settings.""" + pool_config = overrides.pop("pool_config", self._pool_config) + derived = super().derive_pool(**overrides) + derived._pool_config = cast("AdbcPoolConfig", dict(pool_config)) if pool_config is not None else None + derived._validate_pool_settings(None) + return derived + + def _create_pool(self) -> "AdbcConnectionPool | None": + if self._pool_config is None: + return None + return AdbcConnectionPool(self.create_connection, **self._pool_config) + + def _close_pool(self) -> None: + self._shutdown_pool(retire=False) + + def _shutdown_pool(self, *, retire: bool) -> None: + if self._pool_config is None: + return + while True: + pool = self.connection_instance + if pool is None: + return + if not isinstance(pool, AdbcConnectionPool): + msg = "Pooled ADBC shutdown requires an owned pool generation." + raise ImproperConfigurationError(msg) + pool.check_pid() + with self._pool_lock: + if self.connection_instance is not pool: + continue + first_close = not self._pool_close_started + self._pool_close_started = True + break + runtime = self.get_observability_runtime() + if first_close: + runtime.emit_pool_destroying_sync(pool) + pool.close() + if retire and pool.is_closed: + with self._pool_lock: + retired = self.connection_instance is pool + if retired: + self.connection_instance = None + self._pool_close_started = False + if retired: + runtime.emit_pool_destroy_sync(pool) + + def _validate_pool_settings(self, connection_instance: "AdbcConnection | None") -> None: + if self._pool_config is None: + return + max_size = self._pool_config.get("max_size", 4) + timeout = self._pool_config.get("timeout", 30.0) + if isinstance(max_size, bool) or not isinstance(max_size, int) or max_size <= 0: + msg = "ADBC pool max_size must be a positive integer." + raise ImproperConfigurationError(msg) + if not math.isfinite(timeout) or timeout < 0: + msg = "ADBC pool timeout must be finite and nonnegative." + raise ImproperConfigurationError(msg) + if resolve_driver_name_from_config(self.connection_config) != "adbc_driver_postgresql.dbapi.connect": + msg = "ADBC pooling supports only the native PostgreSQL driver." + raise ImproperConfigurationError(msg) + conn_kwargs = self.connection_config.get("conn_kwargs") or {} + native_autocommit = conn_kwargs.get("adbc.connection.autocommit", False) + if self.connection_config.get("autocommit", False) or native_autocommit not in (False, "false", "0", 0): + msg = "ADBC pooling requires autocommit=False." + raise ImproperConfigurationError(msg) + if connection_instance is not None or self._user_connection_hook is not None: + msg = "ADBC pooling does not support injected connections or connection initialization hooks." + raise ImproperConfigurationError(msg) + for distribution in ("adbc-driver-manager", "adbc-driver-postgresql"): + try: + installed = importlib.metadata.version(distribution) + except importlib.metadata.PackageNotFoundError as exc: + msg = f"ADBC pooling requires {distribution}>=1.12.0." + raise ImproperConfigurationError(msg) from exc + match = re.fullmatch( + r"(\d+)\.(\d+)(?:\.(\d+))?((?:(?:a|b|rc)\d+)?(?:\.post\d+)?(?:\.dev\d+)?(?:\+[a-zA-Z0-9.-]+)?)", + installed, + ) + release = tuple(int(part or 0) for part in match.groups()[:3]) if match is not None else (0, 0, 0) + suffix = match.group(4) if match is not None else "" + if release < (1, 12, 0) or (release == (1, 12, 0) and suffix.startswith(("a", "b", "rc", ".dev"))): + msg = f"ADBC pooling requires {distribution}>=1.12.0; installed version is {installed}." + raise ImproperConfigurationError(msg) + if not callable(getattr(AdbcConnection, "adbc_clone", None)): + msg = "ADBC pooling requires the public Connection.adbc_clone capability." + raise ImproperConfigurationError(msg) + @property def supports_migration_schemas(self) -> bool: # type: ignore[override] """Migration schema support is only available for PostgreSQL-backed ADBC drivers.""" @@ -355,8 +512,8 @@ def create_connection(self) -> AdbcConnection: Raises: ImproperConfigurationError: If the connection could not be established. """ - if self.connection_instance is not None: - return self.connection_instance + if self._pool_config is None and self.connection_instance is not None: + return cast("AdbcConnection", self.connection_instance) connection: AdbcConnection | None = None try: @@ -412,6 +569,8 @@ def provide_session( Returns: A context manager that yields an AdbcDriver instance. """ + if self._pool_config is not None: + self.create_pool() self._detect_extensions_if_needed() statement_config = resolve_runtime_statement_config( statement_config, self.statement_config, self._default_session_config diff --git a/tests/unit/adapters/test_adbc/test_adbc_config_normalization.py b/tests/unit/adapters/test_adbc/test_adbc_config_normalization.py index af504dbd9..2710f5c99 100644 --- a/tests/unit/adapters/test_adbc/test_adbc_config_normalization.py +++ b/tests/unit/adapters/test_adbc/test_adbc_config_normalization.py @@ -1,5 +1,10 @@ """Unit tests for ADBC config normalization helpers.""" +import importlib.metadata +import os +import threading +from concurrent.futures import ThreadPoolExecutor +from importlib import import_module from typing import Any, cast, get_type_hints from unittest.mock import MagicMock @@ -16,6 +21,7 @@ resolve_driver_connect_func, resolve_driver_name_from_config, ) +from sqlspec.adapters.adbc.pool import AdbcConnectionPool from sqlspec.core import SQL, ParameterStyle from sqlspec.exceptions import ImproperConfigurationError, SQLSpecError @@ -39,6 +45,352 @@ def fake_connect(**_kwargs: Any) -> Any: assert "on_connection_create" not in config.driver_features +def test_adbc_pool_option_enables_exclusive_reuse_without_changing_standalone_connections( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Only explicit pool settings enable reuse, while direct connections stay caller-owned.""" + seed, lease, standalone = MagicMock(), MagicMock(), MagicMock() + seed.adbc_clone.return_value = lease + create = MagicMock(side_effect=[seed, standalone]) + monkeypatch.setattr(import_module("adbc_driver_postgresql.dbapi"), "connect", create) + config = AdbcConfig( + connection_config={"driver_name": "postgres"}, + pool_config={}, + driver_features={"enable_pgvector": False, "enable_paradedb": False, "enable_pg_textsearch": False}, + ) + pool = config.provide_pool() + create.assert_not_called() + with config.provide_session() as first: + assert first.connection is lease + with config.provide_session() as second: + assert second.connection is lease + assert config.create_pool() is pool + assert config.create_connection() is standalone + standalone.close.assert_not_called() + config.close_pool() + seed.close.assert_called_once_with() + lease.close.assert_called_once_with() + + +def test_adbc_disabled_pool_methods_preserve_borrowed_connection() -> None: + """All disabled pool hooks are inert even when a legacy raw connection is supplied.""" + borrowed = MagicMock() + events: list[Any] = [] + config = AdbcConfig( + connection_config={"driver_name": "sqlite"}, + connection_instance=borrowed, + pool_config=None, + driver_features={"on_pool_create": events.append, "on_pool_destroy": events.append}, + ) + assert config.create_pool() is None + assert config.provide_pool() is None + assert config._create_pool() is None + config._close_pool() + config.close_pool() + assert config.create_connection() is borrowed + borrowed.close.assert_not_called() + assert events == [] + + +@pytest.mark.parametrize( + "invalid", + [ + "driver", + "autocommit", + "native_autocommit", + "injected", + "hook", + "size_zero", + "size_bool", + "size_float", + "timeout", + ], +) +def test_adbc_pool_rejects_unsupported_settings_before_native_io(monkeypatch: pytest.MonkeyPatch, invalid: str) -> None: + """Unsupported pooling settings fail before opening a native resource.""" + create = MagicMock() + monkeypatch.setattr(import_module("adbc_driver_postgresql.dbapi"), "connect", create) + options: dict[str, Any] = {"connection_config": {"driver_name": "postgres"}, "pool_config": {}} + if invalid == "driver": + options["connection_config"]["driver_name"] = "sqlite" + elif invalid == "autocommit": + options["connection_config"]["autocommit"] = True + elif invalid == "native_autocommit": + options["connection_config"]["conn_kwargs"] = {"adbc.connection.autocommit": "true"} + elif invalid == "injected": + options["connection_instance"] = MagicMock() + elif invalid == "hook": + options["driver_features"] = {"on_connection_create": lambda _connection: None} + elif invalid.startswith("size"): + options["pool_config"]["max_size"] = {"size_zero": 0, "size_bool": True, "size_float": 1.5}[invalid] + else: + options["pool_config"]["timeout"] = float("nan") + with pytest.raises(ImproperConfigurationError): + AdbcConfig(**options) + create.assert_not_called() + + +@pytest.mark.parametrize("distribution", ["adbc-driver-manager", "adbc-driver-postgresql"]) +@pytest.mark.parametrize("installed", ["1.11.0", "1.12.0rc1", "1.12.0.dev0", None]) +def test_adbc_pool_requires_supported_native_versions_only_when_enabled( + monkeypatch: pytest.MonkeyPatch, distribution: str, installed: str | None +) -> None: + """Unsupported or absent optional packages do not affect the unpooled default.""" + original_version = importlib.metadata.version + + def version(name: str) -> str: + if name != distribution: + return original_version(name) + if installed is None: + raise importlib.metadata.PackageNotFoundError(name) + return installed + + monkeypatch.setattr(importlib.metadata, "version", version) + assert AdbcConfig(connection_config={"driver_name": "postgres"}).create_pool() is None + with pytest.raises(ImproperConfigurationError, match=r"1\.12\.0"): + _postgres_pool_config() + + +def test_adbc_pool_requires_public_native_clone_capability(monkeypatch: pytest.MonkeyPatch) -> None: + """A sufficient package version without public cloning fails before opening a connection.""" + monkeypatch.setattr(adbc_driver_manager.dbapi.Connection, "adbc_clone", None) + with pytest.raises(ImproperConfigurationError, match="adbc_clone"): + _postgres_pool_config() + + +def test_adbc_pool_atomic_first_publication_and_callbacks_outside_config_lock() -> None: + """Concurrent first requests observe one generation and one outside-lock create callback.""" + events: list[Any] = [] + start = threading.Barrier(5, timeout=3) + + def created(pool: Any) -> None: + with ThreadPoolExecutor(max_workers=1) as executor: + assert executor.submit(config.provide_pool).result(timeout=2) is pool + events.append(pool) + + config = _postgres_pool_config(driver_features={"on_pool_create": created}) + + def acquire_pool() -> Any: + start.wait() + return config.provide_pool() + + with ThreadPoolExecutor(max_workers=4) as executor: + futures = [executor.submit(acquire_pool) for _ in range(4)] + start.wait() + pools = [future.result(timeout=3) for future in futures] + assert all(pool is pools[0] for pool in pools) + assert events == [pools[0]] + config.close_pool() + + +def test_adbc_pool_drain_retains_origin_and_restarts_after_final_release(monkeypatch: pytest.MonkeyPatch) -> None: + """A pending drain keeps ownership and retires exactly once before a new generation.""" + seed, lease = MagicMock(), MagicMock() + seed.adbc_clone.return_value = lease + monkeypatch.setattr(import_module("adbc_driver_postgresql.dbapi"), "connect", lambda **_kwargs: seed) + events: list[tuple[str, Any]] = [] + config = _postgres_pool_config( + driver_features={ + "on_pool_create": lambda pool: events.append(("create", pool)), + "on_pool_destroying": lambda pool: events.append(("destroying", pool)), + "on_pool_destroy": lambda pool: events.append(("destroy", pool)), + } + ) + context = config.provide_session() + context.__enter__() + origin = config.provide_pool() + config.close_pool() + assert config.connection_instance is origin + with pytest.raises(SQLSpecError, match="drain"): + config.provide_session() + seed.close.assert_not_called() + context.__exit__(None, None, None) + replacement = config.provide_pool() + assert replacement is not origin + assert events == [("create", origin), ("destroying", origin), ("destroy", origin), ("create", replacement)] + lease.close.assert_called_once_with() + seed.close.assert_called_once_with() + config.close_pool() + config.close_pool() + assert events[-2:] == [("destroying", replacement), ("destroy", replacement)] + + +def test_adbc_pool_cleanup_failure_retains_generation_for_retry(monkeypatch: pytest.MonkeyPatch) -> None: + """A native-close failure never permits a new generation before cleanup succeeds.""" + seed, lease = MagicMock(), MagicMock() + seed.adbc_clone.return_value = lease + lease.close.side_effect = [RuntimeError("native close failed"), None] + monkeypatch.setattr(import_module("adbc_driver_postgresql.dbapi"), "connect", lambda **_kwargs: seed) + config = _postgres_pool_config() + with config.provide_session(): + pass + origin = config.provide_pool() + with pytest.raises(RuntimeError, match="native close failed"): + config.close_pool() + assert config.connection_instance is origin + with pytest.raises(SQLSpecError, match="drain"): + config.create_pool() + seed.close.assert_not_called() + config.close_pool() + assert config.connection_instance is None + assert lease.close.call_count == 2 + seed.close.assert_called_once_with() + + +@pytest.mark.parametrize("close_fails", [False, True]) +def test_adbc_pool_body_error_evicts_and_preserves_primary(monkeypatch: pytest.MonkeyPatch, close_fails: bool) -> None: + """Failed session bodies discard their exact lease without masking the application error.""" + seed, lease = MagicMock(), MagicMock() + primary = ValueError("application failed") + seed.adbc_clone.return_value = lease + if close_fails: + lease.close.side_effect = [KeyboardInterrupt("native close interrupted"), None] + monkeypatch.setattr(import_module("adbc_driver_postgresql.dbapi"), "connect", lambda **_kwargs: seed) + config = _postgres_pool_config(pool_config={"max_size": 1, "timeout": 0}) + with pytest.raises(ValueError) as caught, config.provide_session() as session: + assert session.connection is lease + raise primary + assert caught.value is primary + lease.rollback.assert_called_once_with() + lease.adbc_connection.set_autocommit.assert_not_called() + if close_fails: + with pytest.raises(SQLSpecError, match="Timed out"): + with config.provide_session(): + pass + config.close_pool() + assert lease.close.call_count == (2 if close_fails else 1) + seed.close.assert_called_once_with() + + +def test_adbc_pool_destroying_reentry_cannot_close_replacement_generation() -> None: + """Reentrant callbacks cannot redirect a shutdown onto a newly published generation.""" + replacements: list[Any] = [] + events: list[tuple[str, Any]] = [] + + def destroying(pool: Any) -> None: + events.append(("destroying", pool)) + if not replacements: + pool.close() + replacements.append(config.provide_pool()) + + config = _postgres_pool_config( + driver_features={ + "on_pool_create": lambda pool: events.append(("create", pool)), + "on_pool_destroying": destroying, + "on_pool_destroy": lambda pool: events.append(("destroy", pool)), + } + ) + origin = config.provide_pool() + config.close_pool() + replacement = replacements[0] + assert config.provide_pool() is replacement + assert events == [("create", origin), ("destroying", origin), ("destroy", origin), ("create", replacement)] + config.close_pool() + + +def test_adbc_pool_session_releases_exact_acquired_generation(monkeypatch: pytest.MonkeyPatch) -> None: + """A handler returns its lease to the captured generation even if the config reference changes.""" + seed, lease = MagicMock(), MagicMock() + seed.adbc_clone.return_value = lease + monkeypatch.setattr(import_module("adbc_driver_postgresql.dbapi"), "connect", lambda **_kwargs: seed) + config = _postgres_pool_config() + context = config.provide_session() + context.__enter__() + origin = cast("AdbcConnectionPool", config.provide_pool()) + replacement = AdbcConnectionPool(MagicMock()) + config.connection_instance = replacement + context.__exit__(None, None, None) + lease.rollback.assert_called_once_with() + lease.cursor.return_value.execute.assert_called_once_with("DISCARD ALL") + assert origin.close() is True + seed.close.assert_called_once_with() + config.close_pool() + + +@pytest.mark.parametrize("event", ["create", "destroying"]) +def test_adbc_pool_callback_failure_retains_owned_generation(monkeypatch: pytest.MonkeyPatch, event: str) -> None: + """Lifecycle failures retain ownership and emit each lifecycle callback at most once.""" + seed, lease = MagicMock(), MagicMock() + seed.adbc_clone.return_value = lease + create = MagicMock(return_value=seed) + monkeypatch.setattr(import_module("adbc_driver_postgresql.dbapi"), "connect", create) + seen: list[Any] = [] + + def failure(pool: Any) -> None: + seen.append(pool) + raise KeyboardInterrupt("lifecycle callback failed") + + config = _postgres_pool_config(driver_features={f"on_pool_{event}": failure}) + if event == "create": + with pytest.raises(KeyboardInterrupt, match="lifecycle callback failed"): + config.create_pool() + assert config.provide_pool() is seen[0] + create.assert_not_called() + with config.provide_session(): + pass + if event == "destroying": + with pytest.raises(KeyboardInterrupt, match="lifecycle callback failed"): + config.close_pool() + assert config.connection_instance is seen[0] + seed.close.assert_not_called() + with pytest.raises(SQLSpecError, match="drain"): + config.create_pool() + config.close_pool() + assert len(seen) == 1 + assert config.connection_instance is None + seed.close.assert_called_once_with() + + +def test_adbc_pool_rejects_inherited_generation_before_config_lock_and_callbacks( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """All pool entrypoints reject inherited ownership before a poisoned config lock.""" + + class _LockProbeConfig(AdbcConfig): + def replace_lock(self, lock: Any) -> None: + self._pool_lock = lock + + def close_protected(self) -> None: + self._close_pool() + + events: list[Any] = [] + config = _LockProbeConfig( + connection_config={"driver_name": "postgres"}, + pool_config={}, + driver_features={"on_pool_create": events.append, "on_pool_destroying": events.append}, + ) + config.provide_pool() + events.clear() + poison = MagicMock() + poison.__enter__.side_effect = AssertionError("inherited config lock acquired") + config.replace_lock(poison) + child_pid = os.getpid() + 1 + monkeypatch.setattr(os, "getpid", lambda: child_pid) + for operation in ( + config.create_pool, + config.provide_pool, + config.provide_session, + config.close_pool, + config.close_protected, + ): + with pytest.raises(SQLSpecError, match="fork"): + operation() + poison.__enter__.assert_not_called() + assert events == [] + + +def test_adbc_pool_derive_preserves_settings_and_validates_overrides() -> None: + """Derivation retains opt-in bounds and starts independent native ownership.""" + config = _postgres_pool_config(pool_config={"max_size": 1, "timeout": 0}) + derived = config.derive_pool() + assert derived.provide_pool() is not config.provide_pool() + assert derived.derive_pool(pool_config=None).provide_pool() is None + with pytest.raises(ImproperConfigurationError, match="PostgreSQL"): + config.derive_pool(connection_config_overrides={"driver_name": "sqlite"}) + derived.close_pool() + config.close_pool() + + @pytest.mark.parametrize("primary", [ValueError("hook failed"), KeyboardInterrupt("hook interrupted")]) @pytest.mark.parametrize("cleanup_error", [None, RuntimeError("close failed"), SystemExit("close interrupted")]) def test_adbc_config_hook_failure_closes_owned_connection( @@ -881,6 +1233,18 @@ def __arrow_c_stream__(self, requested_schema: object = None) -> object: assert cursor.executemany.call_args_list[5].args == ("INSERT INTO items (id, name) VALUES (?, ?)", arrow_stream) +def _postgres_pool_config(**overrides: Any) -> AdbcConfig: + """Construct a pooled config without unrelated native extension probes.""" + features: dict[str, Any] = {"enable_pgvector": False, "enable_paradedb": False, "enable_pg_textsearch": False} + features.update(overrides.pop("driver_features", {})) + return AdbcConfig( + connection_config=overrides.pop("connection_config", {"driver_name": "postgres"}), + pool_config=overrides.pop("pool_config", {}), + driver_features=features, + **overrides, + ) + + def _resolve_driver_name(config: AdbcConfig) -> str: """Resolve driver name from configuration.""" return resolve_driver_name_from_config(config.connection_config) diff --git a/tests/unit/adapters/test_adbc/test_arrow_streaming.py b/tests/unit/adapters/test_adbc/test_arrow_streaming.py index 801e3e5c0..9124603b5 100644 --- a/tests/unit/adapters/test_adbc/test_arrow_streaming.py +++ b/tests/unit/adapters/test_adbc/test_arrow_streaming.py @@ -1,11 +1,14 @@ """Unit tests for ADBC Arrow streaming paths.""" +from importlib import import_module from typing import TYPE_CHECKING, Any, cast +from unittest.mock import MagicMock import pyarrow as pa import pytest from adbc_driver_manager import AdbcStatusCode, DatabaseError +from sqlspec.adapters.adbc import AdbcConfig from sqlspec.adapters.adbc._typing import AdbcSessionContext from sqlspec.adapters.adbc.core import get_statement_config from sqlspec.adapters.adbc.driver import AdbcDriver, AdbcSelectStreamSource @@ -155,6 +158,52 @@ def release(_connection: Any, **kwargs: Any) -> None: stream.close() +@pytest.mark.parametrize("kind,consume", [("row", False), ("reader", False), ("reader", True)]) +def test_pooled_lazy_stream_closes_native_owners_then_evicts_origin( + monkeypatch: pytest.MonkeyPatch, kind: str, consume: bool +) -> None: + """Stream cleanup precedes native connection disposal and the next session gets another clone.""" + + class _PooledStreamingConnection(_AdbcStreamingConnection): + def __init__(self) -> None: + super().__init__() + self.close_count = 0 + self.rollback_count = 0 + + def rollback(self) -> None: + self.rollback_count += 1 + + def close(self) -> None: + expected = 0 if kind == "row" else 1 + assert self.cursor_obj.close_count == expected + assert self.cursor_obj.reader_obj.close_count == expected + self.close_count += 1 + + seed, replacement = MagicMock(), MagicMock() + lease = _PooledStreamingConnection() + seed.adbc_clone.side_effect = [lease, replacement] + monkeypatch.setattr(import_module("adbc_driver_postgresql.dbapi"), "connect", lambda **_kwargs: seed) + config = AdbcConfig( + connection_config={"driver_name": "postgres"}, + pool_config={"max_size": 1, "timeout": 0}, + driver_features={"enable_pgvector": False, "enable_paradedb": False, "enable_pg_textsearch": False}, + ) + with config.provide_session() as session: + assert session.connection is lease + if kind == "row": + session.select_stream("SELECT 1 AS x", native_only=True) + else: + reader = session.select_to_arrow("SELECT 1 AS x", return_format="reader").data + if consume: + assert reader.read_all().to_pydict() == {"x": [1, 2, 3]} + assert lease.close_count == 1 + assert lease.rollback_count == 1 + with config.provide_session() as session: + assert session.connection is replacement + config.close_pool() + seed.close.assert_called_once_with() + + @pytest.mark.parametrize("stage", ["exit", "body", "setup"]) def test_session_stream_cleanup_attempts_all_owners_and_preserves_primary(stage: str) -> None: """Cleanup attempts every native reader/cursor and release, preserving setup/body errors.""" From f12a8f25de316dd868ffb6d6f0fe87957ec8c2a2 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 21:47:42 +0000 Subject: [PATCH 135/174] test(adbc): verify native PostgreSQL pool isolation and lifetime --- .../adapters/postgres/adbc/test_driver.py | 224 +++++++++++++++++- 1 file changed, 222 insertions(+), 2 deletions(-) diff --git a/tests/integration/adapters/postgres/adbc/test_driver.py b/tests/integration/adapters/postgres/adbc/test_driver.py index 38e46dfad..fc86e203a 100644 --- a/tests/integration/adapters/postgres/adbc/test_driver.py +++ b/tests/integration/adapters/postgres/adbc/test_driver.py @@ -1,11 +1,16 @@ """PostgreSQL-backed ADBC driver residuals.""" +from collections.abc import Generator +from concurrent.futures import ThreadPoolExecutor +from contextlib import contextmanager +from threading import Barrier +from typing import TYPE_CHECKING, Literal, cast from uuid import UUID, uuid4 import pytest -from sqlspec.adapters.adbc import AdbcDriver -from sqlspec.exceptions import SQLSpecError +from sqlspec.adapters.adbc import AdbcConfig, AdbcDriver +from sqlspec.exceptions import ConnectionTimeoutError, DatabaseConnectionError, SQLSpecError from tests.integration.adapters._shared.adbc_backends import postgresql_session, test_postgresql_specific_features from tests.integration.adapters._shared.adbc_connection import ( test_connection, @@ -18,6 +23,9 @@ test_execute_script_edge_cases, ) +if TYPE_CHECKING: + from pyarrow import RecordBatch + __all__ = ( "postgresql_session", "test_adbc_postgresql_statement_stack_continue_on_error", @@ -223,3 +231,215 @@ def test_postgresql_uuid_row_stream_and_arrow_boundary(postgresql_session: AdbcD assert arrow_type.extension_name == "arrow.opaque" assert arrow_type.type_name == "uuid" assert arrow_table.to_pylist() == [{"value": value.bytes}] + + +@pytest.mark.xdist_group("postgres") +@pytest.mark.adbc +def test_postgresql_pool_reuses_pid_and_preserves_transactional_wrapper( + adbc_postgres_connection_config: dict[str, str], +) -> None: + """A reset lease keeps its backend and still supports real rollback and commit.""" + with _pooled_postgresql_config(adbc_postgres_connection_config) as config: + with config.provide_session() as session: + first_pid = session.select_value("SELECT pg_backend_pid()") + session.execute_script("CREATE TEMP TABLE adbc_pool_rollback (value INTEGER)") + session.commit() + session.execute("INSERT INTO adbc_pool_rollback VALUES (?)", 7) + session.rollback() + assert session.select_value("SELECT COUNT(*) FROM adbc_pool_rollback") == 0 + with config.provide_session() as session: + assert session.select_value("SELECT pg_backend_pid()") == first_pid + session.execute_script("CREATE TEMP TABLE adbc_pool_after_reset (value INTEGER)") + session.commit() + session.execute("INSERT INTO adbc_pool_after_reset VALUES (?)", 11) + session.rollback() + assert session.select_value("SELECT COUNT(*) FROM adbc_pool_after_reset") == 0 + session.execute("INSERT INTO adbc_pool_after_reset VALUES (?)", 13) + session.commit() + assert session.select_value("SELECT value FROM adbc_pool_after_reset") == 13 + + +@pytest.mark.xdist_group("postgres") +@pytest.mark.adbc +def test_postgresql_pool_simultaneous_leases_are_exclusive(adbc_postgres_connection_config: dict[str, str]) -> None: + """Overlapping native operations hold distinct backend connections.""" + barrier = Barrier(2, timeout=30) + with _pooled_postgresql_config(adbc_postgres_connection_config, max_size=2) as config: + + def borrower(value: int) -> tuple[int, int]: + with config.provide_session() as session: + pid = int(session.select_value("SELECT pg_backend_pid()")) + barrier.wait() + observed = int(session.select_value("SELECT CAST(? AS INTEGER)", value)) + barrier.wait() + return pid, observed + + with ThreadPoolExecutor(max_workers=2) as executor: + results = list(executor.map(borrower, [3, 5])) + assert len({pid for pid, _value in results}) == 2 + assert [value for _pid, value in results] == [3, 5] + + +@pytest.mark.xdist_group("postgres") +@pytest.mark.adbc +def test_postgresql_pool_nested_checkout_times_out_without_stealing_outer_lease( + adbc_postgres_connection_config: dict[str, str], +) -> None: + """A saturated nested checkout refuses while the outer backend stays usable.""" + with _pooled_postgresql_config(adbc_postgres_connection_config) as config, config.provide_session() as outer: + pid = outer.select_value("SELECT pg_backend_pid()") + with pytest.raises(ConnectionTimeoutError), config.provide_session(): + pytest.fail("nested checkout bypassed the configured capacity") + assert outer.select_value("SELECT pg_backend_pid()") == pid + + +@pytest.mark.xdist_group("postgres") +@pytest.mark.adbc +def test_postgresql_pool_failed_body_evicts_native_connection(adbc_postgres_connection_config: dict[str, str]) -> None: + """A failed transaction is not reused even after driver-level rollback.""" + with _pooled_postgresql_config(adbc_postgres_connection_config) as config: + with config.provide_session() as session: + failed_pid = session.select_value("SELECT pg_backend_pid()") + with pytest.raises(SQLSpecError), config.provide_session() as session: + session.execute("SELECT 1 / CAST(? AS INTEGER)", 0) + with config.provide_session() as session: + assert session.select_value("SELECT pg_backend_pid()") != failed_pid + assert session.select_value("SELECT CAST(? AS INTEGER)", 17) == 17 + + +@pytest.mark.xdist_group("postgres") +@pytest.mark.adbc +def test_postgresql_pool_active_shutdown_drains_before_new_generation( + adbc_postgres_connection_config: dict[str, str], +) -> None: + """Shutdown preserves an active borrower and refuses checkout until it exits.""" + with _pooled_postgresql_config(adbc_postgres_connection_config) as config: + with config.provide_session() as session: + old_pid = session.select_value("SELECT pg_backend_pid()") + config.close_pool() + assert session.select_value("SELECT pg_backend_pid()") == old_pid + with pytest.raises(DatabaseConnectionError), config.provide_session(): + pytest.fail("a draining generation accepted a new borrower") + with config.provide_session() as session: + assert session.select_value("SELECT pg_backend_pid()") != old_pid + assert session.select_value("SELECT CAST(? AS INTEGER)", 19) == 19 + config.close_pool() + config.close_pool() + + +@pytest.mark.xdist_group("postgres") +@pytest.mark.adbc +def test_postgresql_pool_discards_committed_session_state(adbc_postgres_connection_config: dict[str, str]) -> None: + """Reuse clears committed session state that a rollback alone cannot undo.""" + with _pooled_postgresql_config(adbc_postgres_connection_config) as config: + with config.provide_session() as session: + pid = session.select_value("SELECT pg_backend_pid()") + search_path = session.select_value("SELECT current_setting('search_path')") + session.execute_script("CREATE TEMP TABLE adbc_pool_reset (value INTEGER)") + session.select_value("SELECT set_config('search_path', ?, false)", "pg_catalog") + session.execute_script("LISTEN adbc_pool_notifications") + session.select_value("SELECT pg_advisory_lock(CAST(? AS BIGINT))", 728311) + session.execute_script("PREPARE adbc_pool_prepared AS SELECT 1") + session.commit() + with config.provide_session() as session: + assert session.select_value("SELECT pg_backend_pid()") == pid + assert session.select_value("SELECT current_setting('search_path')") == search_path + assert session.select_value("SELECT to_regclass(?)", "pg_temp.adbc_pool_reset") is None + assert session.select_value("SELECT COUNT(*) FROM pg_listening_channels()") == 0 + assert ( + session.select_value( + "SELECT COUNT(*) FROM pg_locks WHERE locktype = ? AND pid = pg_backend_pid()", "advisory" + ) + == 0 + ) + assert ( + session.select_value("SELECT COUNT(*) FROM pg_prepared_statements WHERE name = ?", "adbc_pool_prepared") + == 0 + ) + + +@pytest.mark.xdist_group("postgres") +@pytest.mark.adbc +def test_postgresql_pool_clones_accept_initial_connection_options( + adbc_postgres_connection_config: dict[str, str], +) -> None: + """Fresh clones support explicit native autocommit=False initialization.""" + with ( + _pooled_postgresql_config( + adbc_postgres_connection_config, conn_kwargs={"adbc.connection.autocommit": "false"}, max_size=2 + ) as config, + config.provide_session() as first, + config.provide_session() as second, + ): + first_pid = first.select_value("SELECT pg_backend_pid()") + assert second.select_value("SELECT pg_backend_pid()") != first_pid + assert first.select_value("SELECT CAST(? AS INTEGER)", 29) == 29 + assert second.select_value("SELECT CAST(? AS INTEGER)", 31) == 31 + + +@pytest.mark.xdist_group("postgres") +@pytest.mark.adbc +@pytest.mark.parametrize("kind", ["rows", "arrow"]) +@pytest.mark.parametrize("consumption", ["never", "partial", "exhausted"]) +def test_postgresql_pool_lazy_streams_evict_after_every_consumption_mode( + adbc_postgres_connection_config: dict[str, str], kind: str, consumption: str +) -> None: + """A lazy stream prevents reuse even when never started or fully exhausted.""" + with _pooled_postgresql_config(adbc_postgres_connection_config) as config: + with config.provide_session() as session: + pid = session.select_value("SELECT pg_backend_pid()") + statement = "SELECT generate_series(1, 3) AS value" + stream = ( + session.select_stream(statement, native_only=True, chunk_size=1) + if kind == "rows" + else session.select_to_arrow(statement, return_format="reader").data + ) + iterator = iter(stream) + if consumption == "partial": + next(iterator) + elif consumption == "exhausted": + parts = list(iterator) + if kind == "rows": + assert parts == [{"value": 1}, {"value": 2}, {"value": 3}] + else: + assert [ + row["value"] for batch in cast("list[RecordBatch]", parts) for row in batch.to_pylist() + ] == [1, 2, 3] + with config.provide_session() as session: + assert session.select_value("SELECT pg_backend_pid()") != pid + assert session.select_value("SELECT CAST(? AS INTEGER)", 23) == 23 + + +@pytest.mark.xdist_group("postgres") +@pytest.mark.adbc +@pytest.mark.parametrize("return_format", ["table", "batches"]) +def test_postgresql_pool_eager_arrow_results_survive_session_and_allow_reuse( + adbc_postgres_connection_config: dict[str, str], return_format: Literal["table", "batches"] +) -> None: + """Eager Arrow data remains owned by the caller while its connection is reused.""" + with _pooled_postgresql_config(adbc_postgres_connection_config) as config: + with config.provide_session() as session: + pid = session.select_value("SELECT pg_backend_pid()") + data = session.select_to_arrow("SELECT generate_series(1, 3) AS value", return_format=return_format).data + with config.provide_session() as session: + assert session.select_value("SELECT pg_backend_pid()") == pid + rows = data.to_pylist() if return_format == "table" else [row for batch in data for row in batch.to_pylist()] + assert rows == [{"value": 1}, {"value": 2}, {"value": 3}] + + +@contextmanager +def _pooled_postgresql_config( + connection_config: dict[str, str], *, max_size: int = 1, conn_kwargs: dict[str, str] | None = None +) -> Generator[AdbcConfig, None, None]: + config = AdbcConfig( + connection_config={ + "driver_name": "adbc_driver_postgresql", + **connection_config, + "conn_kwargs": conn_kwargs or {}, + }, + pool_config={"max_size": max_size, "timeout": 0.05}, + ) + try: + yield config + finally: + config.close_pool() From 3acf28e33cf6647c95d6432a0573c10304dffe87 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 21:49:19 +0000 Subject: [PATCH 136/174] docs(adbc): document pooling and add transactional benchmark --- docs/changelog.md | 4 + docs/databases/adbc.md | 60 +++++++ docs/databases/index.md | 2 +- docs/reference/adapters/adbc.md | 7 + sqlspec/adapters/adbc/config.py | 2 + tools/scripts/bench_adbc_pool.py | 276 +++++++++++++++++++++++++++++++ 6 files changed, 350 insertions(+), 1 deletion(-) create mode 100644 tools/scripts/bench_adbc_pool.py diff --git a/docs/changelog.md b/docs/changelog.md index 1af28609f..d2c7e2e36 100644 --- a/docs/changelog.md +++ b/docs/changelog.md @@ -11,6 +11,8 @@ important operational fixes. **Added:** +* Added a pool for the native PostgreSQL ADBC driver. Set `AdbcConfig(pool_config={})` to reuse connections. Each lease has a wait timeout; each return clears session state. Pooling is off by default. See the [ADBC guide](databases/adbc.md#postgresql-connection-reuse) for supported settings. + * Added Sybase Adaptive Server Enterprise (`sqlspec.adapters.sybase` with `SybaseSyncConfig` and `SybaseSyncDriver`) and SAP SQL Anywhere (`sqlspec.adapters.sqlanywhere` with `SqlAnywhereSyncConfig` and `SqlAnywhereSyncDriver`) database adapters, custom SQLGlot dialects (`Sybase` and `SQLAnywhere` in `sqlspec.dialects.sybase`), data dictionaries (`SybaseSyncDataDictionary` and `SqlAnywhereSyncDataDictionary` with `SYBASE_CONFIG` and `SQLANYWHERE_CONFIG`), type converters (`SybaseOutputConverter` and `SqlAnywhereOutputConverter`), and extension stores for Litestar sessions (`SybaseStore`, `SqlAnywhereStore`), Event queues (`SybaseEventQueueStore`, `SqlAnywhereEventQueueStore`), and Google ADK sessions and memory (`SybaseADKStore`, `SybaseADKMemoryStore`, `SqlAnywhereADKStore`, `SqlAnywhereADKMemoryStore`). ([#870](https://github.com/litestar-org/sqlspec/pull/870)) * Added `arrow_odbc` dialect auto-detection and `information_schema` / system catalog introspection for Sybase ASE (`Adaptive Server Enterprise` / `ASE`) and SAP SQL Anywhere (`SQL Anywhere` / `Sybase IQ`), routing detected connections to the `sybase` and `sqlanywhere` data dictionaries. @@ -81,6 +83,8 @@ important operational fixes. **Fixed:** +* ADBC now closes its own connections if a hook, profile setup, driver setup, or session setup fails. It keeps the first error and leaves supplied connections open. + * Preserved supported Arrow conversion and mutation paths for Asyncmy and Spanner while reporting their native capabilities accurately; corrected ADBC lowercase readers, DuckDB stream ownership and table filtering, and SQL Server Arrow cursor lifetimes. * Corrected Oracle session-renewal bindings, pipeline error conversion, ADK batch-error handling, standalone connection setup, and deferred Arrow failures. Text-only Oracle memory stores remain available on servers without native vector support. diff --git a/docs/databases/adbc.md b/docs/databases/adbc.md index 0ecf99aca..e391214e4 100644 --- a/docs/databases/adbc.md +++ b/docs/databases/adbc.md @@ -118,6 +118,66 @@ with spec.provide_session(config) as session: ) ``` +## PostgreSQL Connection Reuse + +ADBC opens and closes a connection for each session by default. For short +PostgreSQL transactions, enable connection reuse with `pool_config={}` or supply +explicit bounds: + +```python +import os + +from sqlspec.adapters.adbc import AdbcConfig + +config = AdbcConfig( + connection_config={"uri": os.environ["DATABASE_URI"]}, + pool_config={"max_size": 4, "timeout": 30.0}, +) +try: + with config.provide_session() as session: + value = session.select_value("SELECT :value", value=42) + session.commit() +finally: + config.close_pool() +``` + +`pool_config=None` keeps a fresh connection per session. An empty mapping gives +four leases and a 30-second wait limit. Each lease gives one caller sole use of a +connection. The pool holds one extra seed, for a total of `max_size + 1` native +connections. Nested sessions need separate leases and can time out when the pool +is full. Standalone connections do not count toward this limit. + +Both native packages, `adbc-driver-manager` and `adbc-driver-postgresql`, must be +version 1.12.0 or newer. This pool supports only the native PostgreSQL driver +with `autocommit=False`. It rejects supplied connections and +`on_connection_create` hooks. Other backends keep a fresh connection per +session. The `supports_connection_pooling` flag means the adapter offers a +pool; the flag does not say whether a given config has enabled it. + +On clean exit, the pool rolls back unfinished work, runs `DISCARD ALL` outside +a transaction, and turns native autocommit off. This clears committed session +settings, temporary tables, prepared statements, listeners, and advisory locks. +A failed session or reset drops that connection from reuse. Direct changes to +native connection options fall outside this reset contract. Measure complete +transactions, including reset, on your own workload before choosing a pool. + +Lazy row streams and Arrow readers belong to their session. On exit, the +session closes their native readers and cursors, then drops the connection +from reuse. This also applies to streams that were never read or fully read. +Consume streams inside the session. Materialized tables and eagerly drained +batch lists stay valid afterward and allow reuse. + +`close_pool()` blocks new checkouts and closes idle connections. Active leases +finish first, then close; their work is not interrupted. If native cleanup +fails, call `close_pool()` again to retry. The pool keeps ownership until every +handle is closed. A later explicit session can then start a new generation. +Create configs after forking: a pool from the parent process refuses use in a +child. Pools do not refresh credentials on their own. + +`create_connection()` and `provide_connection()` remain separate from the pool. +An unpooled `connection_instance` stays open when a session ends; the caller +owns it. + ## Key Driver Features and Pool Tuning | Setting | Location | Description | diff --git a/docs/databases/index.md b/docs/databases/index.md index 15c67587e..a2a7fe526 100644 --- a/docs/databases/index.md +++ b/docs/databases/index.md @@ -26,7 +26,7 @@ SQLSpec supports synchronous and asynchronous **database adapters** spanning emb | [`duckdb`](duckdb.md) | `sqlspec[duckdb]` | Sync | `?` (`qmark`) | `DuckDBConnectionPool` | Native Zero-Copy | Table Queue | Sessions + Memory (`FTS`, `ARRAY` Cosine) | | [`bigquery`](bigquery.md) | `sqlspec[bigquery]` | Sync | `@name` (`named`) | Stateless HTTP/gRPC Client | Storage Read / Write API | Table Queue | Sessions + Memory (`SEARCH` / `LIKE`) | | [`spanner`](spanner.md) | `sqlspec[spanner]` | Sync + Async | `@name` (`named`) | `PingingPool` / `TransactionPingingPool` | Row / Batched | Lock-Free Snapshot + `run_in_transaction` | Sessions + Memory (`ARRAY`, FTS `SCORE`, RRF) + Artifacts | -| [`adbc`](adbc.md) | `sqlspec[adbc]` | Sync | `?` / `$1` / `:1` | Driver-Managed (`NoPoolSyncConfig`) | Native C-Stream | Table Queue | Sessions + Memory (PostgreSQL `pgvector` RRF, DuckDB, Oracle, Db2) | +| [`adbc`](adbc.md) | `sqlspec[adbc]` | Sync | `?` / `$1` / `:1` | Optional PostgreSQL clone pool | Native C-Stream | Table Queue | Sessions + Memory (PostgreSQL `pgvector` RRF, DuckDB, Oracle, Db2) | | [`arrow_odbc`](arrow-odbc.md) | `sqlspec[arrow-odbc]` | Sync | `?` (`qmark`) | ODBC Driver Pool (`enable_pooling`) | Native Rust ODBC Batch | Table Queue | Sessions + Memory (`LIKE` / `POSSTR`) | ## Database Setup Guides diff --git a/docs/reference/adapters/adbc.md b/docs/reference/adapters/adbc.md index e752171ac..2c8b5627c 100644 --- a/docs/reference/adapters/adbc.md +++ b/docs/reference/adapters/adbc.md @@ -7,6 +7,7 @@ from sqlspec.adapters.adbc import ( AdbcDriver, AdbcDriverFeatures, AdbcExceptionHandler, + AdbcPoolConfig, ) ``` @@ -24,6 +25,12 @@ Apache Arrow Database Connectivity (ADBC) adapter supporting zero-copy Arrow ing separate_signature: false show_signature_annotations: true +::: sqlspec.adapters.adbc.AdbcPoolConfig + options: + heading_level: 3 + separate_signature: false + show_signature_annotations: true + ::: sqlspec.adapters.adbc.AdbcDriverFeatures options: heading_level: 3 diff --git a/sqlspec/adapters/adbc/config.py b/sqlspec/adapters/adbc/config.py index d984fa3c2..81ed57538 100644 --- a/sqlspec/adapters/adbc/config.py +++ b/sqlspec/adapters/adbc/config.py @@ -53,7 +53,9 @@ class AdbcPoolConfig(TypedDict): """Optional PostgreSQL clone-pool bounds; an empty mapping enables defaults.""" max_size: NotRequired[int] + """Maximum simultaneous clone leases, excluding the original seed (default 4).""" timeout: NotRequired[float] + """Maximum seconds to wait for a lease (default 30.0).""" class AdbcConnectionParams(TypedDict): diff --git a/tools/scripts/bench_adbc_pool.py b/tools/scripts/bench_adbc_pool.py new file mode 100644 index 000000000..3a57edd5b --- /dev/null +++ b/tools/scripts/bench_adbc_pool.py @@ -0,0 +1,276 @@ +"""Measure complete PostgreSQL ADBC sessions with unpooled, cloned and pooled connections.""" + +import hashlib +import json +import math +import os +import platform +import re +import shutil +import statistics +import subprocess +import sys +from concurrent.futures import ThreadPoolExecutor +from contextlib import ExitStack +from datetime import UTC, datetime +from importlib import metadata +from pathlib import Path +from threading import Barrier, BrokenBarrierError, Event, Lock +from time import perf_counter_ns +from typing import TYPE_CHECKING +from uuid import uuid4 + +import click + +from sqlspec.adapters.adbc import AdbcConfig + +if TYPE_CHECKING: + from adbc_driver_manager.dbapi import Connection + + +class CloneConfig(AdbcConfig): + """Close each session clone while retaining its original caller-owned seed.""" + + __slots__ = ("_clone_lock", "_seed") + + def __init__(self, uri: str, seed: "Connection") -> None: + super().__init__( + connection_config={"driver_name": "postgresql", "uri": uri, "autocommit": False}, + driver_features=_features(), + ) + self._seed = seed + self._clone_lock = Lock() + + def create_connection(self) -> "Connection": + with self._clone_lock: + return self._seed.adbc_clone() + + +@click.command() +@click.option("--uri-env", default="ADBC_POSTGRESQL_TEST_URI", show_default=True) +@click.option("--mode", type=click.Choice(["all", "unpooled", "clone", "pooled"]), default="all", show_default=True) +@click.option("--workers", type=click.Choice(["1", "4"]), default="1", show_default=True) +@click.option("--warmup", type=click.IntRange(min=1), default=100, show_default=True) +@click.option("--samples", type=click.IntRange(min=1), default=1000, show_default=True) +@click.option("--repeats", type=click.IntRange(min=1), default=7, show_default=True) +@click.option("--output", type=click.Path(path_type=Path, dir_okay=False)) +def main(uri_env: str, mode: str, workers: str, warmup: int, samples: int, repeats: int, output: Path | None) -> None: + """Compare per-session latency on independent worker rows in a disposable schema. + + Copy this script outside the checkout to run installed source/baseline wheels. + Set SQLSPEC_BENCH_REVISION to the verified source SHA for an installed wheel. + URI values and credentials are never included in the JSON report. + """ + uri = os.environ.get(uri_env) + if not uri: + raise click.ClickException(f"Set the connection environment variable {uri_env} before running") + provenance = _provenance() + worker_count = int(workers) + schema = "sqlspec_adbc_pool_bench_" + uuid4().hex + if re.fullmatch(r"sqlspec_adbc_pool_bench_[0-9a-f]{32}", schema) is None: + raise RuntimeError("Invalid generated benchmark schema") + table = schema + ".sessions" + unpooled = AdbcConfig( + connection_config={"driver_name": "postgresql", "uri": uri, "autocommit": False}, driver_features=_features() + ) + modes = ["unpooled", "clone", "pooled"] if mode == "all" else [mode] + results: list[dict[str, object]] = [] + combined: dict[str, list[int]] = {name: [] for name in modes} + with ExitStack() as cleanup: + configs = {"unpooled": unpooled} + if "clone" in modes: + seed = unpooled.create_connection() + cleanup.callback(seed.close) + configs["clone"] = CloneConfig(uri, seed) + if "pooled" in modes: + pooled = AdbcConfig( + connection_config={"driver_name": "postgresql", "uri": uri, "autocommit": False}, + driver_features=_features(), + pool_config={"max_size": worker_count, "timeout": 30.0}, + ) + cleanup.callback(pooled.close_pool) + pooled.create_pool() + configs["pooled"] = pooled + with unpooled.provide_session() as session: + server_version = str(session.select_value("SELECT current_setting('server_version')")) + server_identity = str( + session.select_value("SELECT inet_server_addr()::text || ':' || inet_server_port()::text") + ) + session.commit() + cleanup.callback(_drop_schema, unpooled, schema) + with unpooled.provide_session() as session: + session.execute(f"CREATE SCHEMA {schema}") + session.execute( + f"CREATE TABLE {table} (worker_id INTEGER PRIMARY KEY, value BIGINT NOT NULL, payload VARCHAR(32) NOT NULL)" + ) + for worker_id in range(1, worker_count + 1): + session.execute( + f"INSERT INTO {table} (worker_id, value, payload) VALUES (?, ?, ?)", worker_id, 0, "sqlspec" + ) + session.commit() + with ThreadPoolExecutor(max_workers=worker_count) as executor: + for repeat in range(repeats): + ordered = modes[repeat % len(modes) :] + modes[: repeat % len(modes)] + for name in ordered: + with unpooled.provide_session() as session: + session.execute(f"UPDATE {table} SET value = 0") + session.commit() + timings = _measure(executor, configs[name], table, worker_count, warmup, samples) + with unpooled.provide_session() as session: + rows = session.select(f"SELECT worker_id, value, payload FROM {table} ORDER BY worker_id") + expected = [ + {"worker_id": worker_id, "value": warmup + samples, "payload": "sqlspec"} + for worker_id in range(1, worker_count + 1) + ] + if rows != expected: + raise RuntimeError("Benchmark row increment or payload verification failed") + session.commit() + combined[name].extend(timings) + results.append({ + "mode": name, + "repeat": repeat + 1, + **_summary(timings), + "verified_rows": worker_count, + }) + report = { + "metadata": { + **provenance, + "timestamp": datetime.now(UTC).isoformat(), + "server_version": server_version, + "server_identifier_sha256": hashlib.sha256(server_identity.encode()).hexdigest(), + "workers": worker_count, + "warmup_per_worker": warmup, + "samples_per_worker": samples, + "repeats": repeats, + "timing_unit": "nanoseconds", + "workload": "begin/read/update/read/commit/session-release", + }, + "results": results, + "aggregate": {name: _summary(timings) for name, timings in combined.items()}, + } + encoded = json.dumps(report, indent=2, allow_nan=False) + "\n" + if output is None: + click.echo(encoded, nl=False) + else: + output.write_text(encoded, encoding="utf-8") + click.echo(f"Wrote benchmark results to {output}") + + +def _features() -> dict[str, bool]: + return { + "enable_pgvector": False, + "enable_paradedb": False, + "enable_pg_textsearch": False, + "enable_statement_profiling": False, + } + + +def _drop_schema(config: AdbcConfig, schema: str) -> None: + with config.provide_session() as session: + session.execute(f"DROP SCHEMA IF EXISTS {schema} CASCADE") + session.commit() + + +def _measure( + executor: ThreadPoolExecutor, config: AdbcConfig, table: str, workers: int, warmup: int, samples: int +) -> list[int]: + barrier = Barrier(workers, timeout=60) + stop = Event() + + def worker(worker_id: int) -> list[int]: + try: + for _ in range(warmup): + if stop.is_set(): + raise BrokenBarrierError("A benchmark peer failed") + transact(worker_id) + barrier.wait() + timings = [] + for _ in range(samples): + if stop.is_set(): + raise BrokenBarrierError("A benchmark peer failed") + start = perf_counter_ns() + transact(worker_id) + timings.append(perf_counter_ns() - start) + return timings + except BaseException: + stop.set() + barrier.abort() + raise + + def transact(worker_id: int) -> None: + with config.provide_session() as session: + session.begin() + before = int(session.select_value(f"SELECT value FROM {table} WHERE worker_id = ?", worker_id)) + session.execute(f"UPDATE {table} SET value = ? WHERE worker_id = ?", before + 1, worker_id) + after = session.select_value(f"SELECT value FROM {table} WHERE worker_id = ?", worker_id) + if after != before + 1: + raise RuntimeError("Transactional read-after-write verification failed") + session.commit() + + futures = [executor.submit(worker, worker_id) for worker_id in range(1, workers + 1)] + timings: list[int] = [] + failures: list[BaseException] = [] + for future in futures: + try: + timings.extend(future.result()) + except BaseException as exc: + failures.append(exc) + if failures: + raise next((exc for exc in failures if not isinstance(exc, BrokenBarrierError)), failures[0]) + return timings + + +def _summary(timings: list[int]) -> dict[str, object]: + if not timings or any(value <= 0 for value in timings): + raise RuntimeError("Missing or invalid session timings") + ordered = sorted(timings) + median = statistics.median(ordered) + deviation = statistics.pstdev(ordered) + return { + "count": len(ordered), + "median_ns": median, + "p95_ns": ordered[math.ceil(0.95 * len(ordered)) - 1], + "stddev_ns": deviation, + "noise_ratio": deviation / median, + "samples_ns": timings, + } + + +def _provenance() -> dict[str, object]: + import sqlspec + from sqlspec.driver import _sync + + origin = Path(sqlspec.__file__ or "").resolve() + source_root = origin.parent.parent + revision = os.environ.get("SQLSPEC_BENCH_REVISION") + dirty: bool | None = None + if (source_root / ".git").exists(): + git = shutil.which("git") + if git is None: + raise click.ClickException("Git is required to establish source-checkout provenance") + revision = subprocess.check_output([git, "-C", str(source_root), "rev-parse", "HEAD"], text=True).strip() + dirty = bool(subprocess.check_output([git, "-C", str(source_root), "status", "--porcelain"], text=True).strip()) + if not revision or re.fullmatch(r"[0-9a-f]{7,40}", revision) is None: + raise click.ClickException("Set SQLSPEC_BENCH_REVISION to the verified source SHA for installed packages") + versions = {} + for distribution in ("sqlspec", "adbc-driver-manager", "adbc-driver-postgresql", "pyarrow", "sqlglot"): + try: + versions[distribution] = metadata.version(distribution) + except metadata.PackageNotFoundError: + versions[distribution] = "not-installed" + return { + "revision": revision, + "source_dirty": dirty, + "harness_sha256": hashlib.sha256(Path(__file__).read_bytes()).hexdigest(), + "sqlspec_origin": str(origin), + "driver_origin": _sync.__file__, + "mypyc_compiled": (_sync.__file__ or "").endswith((".so", ".pyd")), + "python": sys.version, + "executable": sys.executable, + "platform": platform.platform(), + "versions": versions, + } + + +if __name__ == "__main__": + main() From 129012255447f6f62150f86c21948fbfe130d147 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 22:00:10 +0000 Subject: [PATCH 137/174] fix(adbc): preserve serializer identity in a bounded converter cache --- sqlspec/adapters/adbc/type_converter.py | 7 +- .../test_adbc_config_normalization.py | 71 +++++++++++++++++++ 2 files changed, 76 insertions(+), 2 deletions(-) diff --git a/sqlspec/adapters/adbc/type_converter.py b/sqlspec/adapters/adbc/type_converter.py index 38efb0325..63d1b2ee1 100644 --- a/sqlspec/adapters/adbc/type_converter.py +++ b/sqlspec/adapters/adbc/type_converter.py @@ -88,7 +88,8 @@ def convert_sequence( return converted -_CONVERTER_CACHE: Final[dict[tuple[str, Any], ADBCOutputConverter]] = {} +_CONVERTER_CACHE: Final[dict[tuple[str, int], ADBCOutputConverter]] = {} +_CONVERTER_CACHE_MAX_SIZE: Final[int] = 128 def get_adbc_type_converter( @@ -106,9 +107,11 @@ def get_adbc_type_converter( """ _ = cache_size serializer = json_serializer or to_json - cache_key = (dialect, serializer) + cache_key = (dialect, id(serializer)) converter = _CONVERTER_CACHE.get(cache_key) if converter is None: converter = ADBCOutputConverter(dialect, json_serializer=serializer) + if len(_CONVERTER_CACHE) >= _CONVERTER_CACHE_MAX_SIZE: + del _CONVERTER_CACHE[next(iter(_CONVERTER_CACHE))] _CONVERTER_CACHE[cache_key] = converter return converter diff --git a/tests/unit/adapters/test_adbc/test_adbc_config_normalization.py b/tests/unit/adapters/test_adbc/test_adbc_config_normalization.py index 2710f5c99..f39b91e36 100644 --- a/tests/unit/adapters/test_adbc/test_adbc_config_normalization.py +++ b/tests/unit/adapters/test_adbc/test_adbc_config_normalization.py @@ -1,9 +1,13 @@ """Unit tests for ADBC config normalization helpers.""" +import gc import importlib.metadata +import json import os import threading +import weakref from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass, field from importlib import import_module from typing import Any, cast, get_type_hints from unittest.mock import MagicMock @@ -18,10 +22,12 @@ from sqlspec.adapters.adbc.config import AdbcConnectionParams from sqlspec.adapters.adbc.core import ( build_connection_config, + get_statement_config, resolve_driver_connect_func, resolve_driver_name_from_config, ) from sqlspec.adapters.adbc.pool import AdbcConnectionPool +from sqlspec.adapters.adbc.type_converter import get_adbc_type_converter from sqlspec.core import SQL, ParameterStyle from sqlspec.exceptions import ImproperConfigurationError, SQLSpecError @@ -45,6 +51,71 @@ def fake_connect(**_kwargs: Any) -> Any: assert "on_connection_create" not in config.driver_features +def test_adbc_session_json_bind_accepts_unhashable_callable_serializer() -> None: + """A stateful callable serializer remains valid for public named parameter binding.""" + + @dataclass + class Serializer: + marker: str + + def __call__(self, value: Any) -> str: + return json.dumps({self.marker: value}, separators=(",", ":")) + + config = AdbcConfig( + connection_config={"driver_name": "sqlite"}, + statement_config=get_statement_config("postgres"), + driver_features={ + "json_serializer": Serializer("custom"), + "enable_pgvector": False, + "enable_paradedb": False, + "enable_pg_textsearch": False, + }, + ) + with config.provide_session() as session: + assert session.select_value("SELECT :payload", payload={"answer": 42}) == '{"custom":{"answer":42}}' + + +def test_adbc_converter_cache_preserves_distinct_equal_callable_identity() -> None: + """Equal hashable serializer objects retain their distinct conversion behavior.""" + + @dataclass(frozen=True) + class Serializer: + marker: str = field(compare=False) + + def __call__(self, value: Any) -> str: + return json.dumps({self.marker: value}, separators=(",", ":")) + + first_serializer = Serializer("first") + second_serializer = Serializer("second") + first = get_adbc_type_converter("postgres", json_serializer=first_serializer) + second = get_adbc_type_converter("postgres", json_serializer=second_serializer) + assert first.convert_dict({"answer": 42}) == '{"first":{"answer":42}}' + assert second.convert_dict({"answer": 42}) == '{"second":{"answer":42}}' + + +def test_adbc_converter_cache_releases_unused_serializer_after_turnover() -> None: + """Live cache entries reuse converters without retaining discarded serializers indefinitely.""" + + @dataclass + class Serializer: + marker: str + + def __call__(self, value: Any) -> str: + return json.dumps({self.marker: value}, separators=(",", ":")) + + serializer = Serializer("original") + reference = weakref.ref(serializer) + retained = get_adbc_type_converter("postgres", json_serializer=serializer) + assert get_adbc_type_converter("postgres", json_serializer=serializer) is retained + del retained, serializer + gc.collect() + assert reference() is not None + for index in range(512): + get_adbc_type_converter("postgres", json_serializer=Serializer(str(index))) + gc.collect() + assert reference() is None + + def test_adbc_pool_option_enables_exclusive_reuse_without_changing_standalone_connections( monkeypatch: pytest.MonkeyPatch, ) -> None: From c9ba2cc3ccd48e28d05f4b844d33aa26dbf7dcb6 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 22:00:10 +0000 Subject: [PATCH 138/174] fix(core): exclude returning targets from client row results --- sqlspec/core/statement.py | 5 +-- .../unit/adapters/test_pymssql/test_driver.py | 17 ++++++++++ .../unit/core/test_explain_classification.py | 34 ++++++++++++++++--- 3 files changed, 50 insertions(+), 6 deletions(-) diff --git a/sqlspec/core/statement.py b/sqlspec/core/statement.py index a3c928212..100bd90dc 100644 --- a/sqlspec/core/statement.py +++ b/sqlspec/core/statement.py @@ -576,8 +576,9 @@ def returns_rows(self) -> bool: if state.parsed_expression: expr = state.parsed_expression - if isinstance(expr, (exp.Insert, exp.Update, exp.Delete)) and expr.args.get("returning"): - return True + if isinstance(expr, (exp.Insert, exp.Update, exp.Delete)): + returning = expr.args.get("returning") + return bool(returning) and not (isinstance(returning, exp.Returning) and returning.args.get("into")) return False diff --git a/tests/unit/adapters/test_pymssql/test_driver.py b/tests/unit/adapters/test_pymssql/test_driver.py index 161f9a85f..54b117665 100644 --- a/tests/unit/adapters/test_pymssql/test_driver.py +++ b/tests/unit/adapters/test_pymssql/test_driver.py @@ -91,6 +91,23 @@ def test_dispatch_execute_select_compiles_to_pyformat_and_collects_rows() -> Non assert result.data_row_count == 1 +def test_execute_output_into_does_not_fetch_client_rows() -> None: + """OUTPUT INTO must not fetch a result set that the native cursor does not expose.""" + + class OutputIntoCursor(FakeCursor): + def fetchall(self) -> list[Any]: + msg = "OUTPUT INTO does not expose a client result set" + raise AssertionError(msg) + + cursor = OutputIntoCursor(rowcount=1) + driver = PymssqlDriver(cast("PymssqlConnection", FakeConnection(cursor))) + + result = driver.execute("INSERT INTO dbo.users OUTPUT inserted.id INTO @ids DEFAULT VALUES") + + assert result.rows_affected == 1 + assert result.get_data() == [] + + def test_dispatch_execute_many_uses_executemany_and_rowcount() -> None: """execute_many dispatch should forward batch parameters to pymssql.""" from sqlspec.adapters.pymssql.core import default_statement_config diff --git a/tests/unit/core/test_explain_classification.py b/tests/unit/core/test_explain_classification.py index 6065f723d..d3d3f29d4 100644 --- a/tests/unit/core/test_explain_classification.py +++ b/tests/unit/core/test_explain_classification.py @@ -107,13 +107,39 @@ def test_sql_explain_oracle_stays_non_row_returning() -> None: ) def test_tsql_output_into_does_not_return_client_rows(sql: str) -> None: """OUTPUT ... INTO writes into a table variable, so no rows reach the client.""" - expression = sqlglot.parse_one(sql, dialect="tsql") + statement = SQL(sql, statement_config=StatementConfig(dialect="tsql")) - assert SQLProcessor._operation_profile(expression, SQLProcessor._operation_type(expression)).returns_rows is False + assert statement.returns_rows() is False def test_tsql_output_without_into_returns_client_rows() -> None: """A bare OUTPUT clause streams the affected rows to the client.""" - expression = sqlglot.parse_one("INSERT INTO t (a) OUTPUT inserted.id VALUES (1)", dialect="tsql") + statement = SQL("INSERT INTO t (a) OUTPUT inserted.id VALUES (1)", statement_config=StatementConfig(dialect="tsql")) + + assert statement.returns_rows() is True + + +def test_oracle_returning_into_does_not_return_client_rows() -> None: + """RETURNING INTO fills output variables rather than a client result set.""" + statement = SQL( + "INSERT INTO t (a) VALUES (1) RETURNING id INTO :result", + result=0, + statement_config=StatementConfig(dialect="oracle"), + ) + + assert statement.returns_rows() is False + - assert SQLProcessor._operation_profile(expression, SQLProcessor._operation_type(expression)).returns_rows is True +@pytest.mark.parametrize( + "sql", + [ + "INSERT INTO t (a) VALUES (1) RETURNING id", + "UPDATE t SET a = 1 WHERE id = 1 RETURNING id", + "DELETE FROM t WHERE id = 1 RETURNING id", + ], +) +def test_postgres_returning_keeps_client_rows(sql: str) -> None: + """Ordinary RETURNING statements still expose the affected rows to callers.""" + statement = SQL(sql, statement_config=StatementConfig(dialect="postgres")) + + assert statement.returns_rows() is True From 0e794a818b71b870669c73dd7f3002a5fe4f1820 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 22:00:11 +0000 Subject: [PATCH 139/174] docs(sanic): correct the correlation middleware default --- sqlspec/config.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sqlspec/config.py b/sqlspec/config.py index 758089efb..4286a2fa3 100644 --- a/sqlspec/config.py +++ b/sqlspec/config.py @@ -555,7 +555,7 @@ class SanicConfig(TypedDict): """ enable_correlation_middleware: NotRequired[bool] - """Enable request correlation ID middleware. Default: False.""" + """Enable request correlation ID middleware. Default: True.""" correlation_header: NotRequired[str] """HTTP header to read the request correlation ID from when middleware is enabled. Default: ``X-Request-ID``.""" From a0fb31095a91a4a59af2c5e984ac902df520d576 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 22:00:11 +0000 Subject: [PATCH 140/174] fix(asyncmy): retain ingestion payload ownership during cancellation --- sqlspec/adapters/asyncmy/driver.py | 28 +++++++- .../test_asyncmy/test_load_from_arrow.py | 65 +++++++++++++++++++ 2 files changed, 91 insertions(+), 2 deletions(-) diff --git a/sqlspec/adapters/asyncmy/driver.py b/sqlspec/adapters/asyncmy/driver.py index 327221314..3d88f2b78 100644 --- a/sqlspec/adapters/asyncmy/driver.py +++ b/sqlspec/adapters/asyncmy/driver.py @@ -426,14 +426,38 @@ async def _execute_local_infile_cursor(self, statement: str, filename: str) -> N async def _load_from_arrow_via_local_infile( self, table: str, columns: "list[str]", records: "list[tuple[Any, ...]]" ) -> None: - directory, filename = await asyncio.to_thread(_write_asyncmy_local_infile_payload, records) + writer = asyncio.create_task(asyncio.to_thread(_write_asyncmy_local_infile_payload, records)) + pending_cancellation: asyncio.CancelledError | None = None + while not writer.done(): + try: + await asyncio.shield(writer) + except asyncio.CancelledError as exc: + pending_cancellation = exc + except BaseException: + break try: + directory, filename = writer.result() + except BaseException: + if pending_cancellation is not None: + raise pending_cancellation + raise + try: + if pending_cancellation is not None: + raise pending_cancellation statement = build_load_data_statement(table, columns) exc_handler = self.handle_database_exceptions() await self._run_with_exception_handler(exc_handler, self._execute_local_infile_cursor, statement, filename) self._check_pending_exception(exc_handler) finally: - await asyncio.to_thread(shutil.rmtree, directory, ignore_errors=True) + cleanup = asyncio.create_task(asyncio.to_thread(shutil.rmtree, directory, ignore_errors=True)) + while not cleanup.done(): + try: + await asyncio.shield(cleanup) + except asyncio.CancelledError as exc: + pending_cancellation = exc + cleanup.result() + if pending_cancellation is not None: + raise pending_cancellation def _connection_in_transaction(self) -> bool: """Check if connection is in transaction. diff --git a/tests/unit/adapters/test_asyncmy/test_load_from_arrow.py b/tests/unit/adapters/test_asyncmy/test_load_from_arrow.py index 0f2081cf3..621c70542 100644 --- a/tests/unit/adapters/test_asyncmy/test_load_from_arrow.py +++ b/tests/unit/adapters/test_asyncmy/test_load_from_arrow.py @@ -1,8 +1,13 @@ """Asyncmy load_from_arrow ingest paths.""" import asyncio +import tempfile +from collections.abc import Generator +from contextlib import contextmanager from datetime import timedelta from pathlib import Path +from threading import Event +from types import SimpleNamespace from typing import Any, cast import pyarrow as pa @@ -154,6 +159,66 @@ async def test_local_infile_failure_discards_connection_and_payload(failure: Bas assert not conn._cursor.payload_path.parent.exists() +@pytest.mark.parametrize("cancel_count", [1, 2]) +@pytest.mark.parametrize("write_fails", [False, True]) +async def test_local_infile_cancelled_writer_removes_payload_after_thread_finishes( + cancel_count: int, write_fails: bool, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """Cancellation preserves ownership until the real file writer stops using the payload.""" + started = Event() + release = Event() + finished = Event() + original = tempfile.NamedTemporaryFile + monkeypatch.setattr(tempfile, "tempdir", str(tmp_path)) + + @contextmanager + def blocking_payload(*args: Any, **kwargs: Any) -> Generator[SimpleNamespace, None, None]: + try: + with original(*args, **kwargs) as payload: + + def write(data: bytes) -> int: + written = payload.write(data) + payload.flush() + started.set() + if not release.wait(timeout=5): + msg = "The payload writer was never released." + raise TimeoutError(msg) + if write_fails: + msg = "Payload write failed after cancellation." + raise OSError(msg) + return written + + yield SimpleNamespace(write=write, name=payload.name) + finally: + finished.set() + + monkeypatch.setattr(tempfile, "NamedTemporaryFile", blocking_payload) + conn = _FakeConnection() + driver = AsyncmyDriver( + connection=cast("Any", conn), + driver_features={"storage_capabilities": _CAPS, "enable_local_infile_bulk_load": True}, + ) + task = asyncio.create_task(driver.load_from_arrow("orders", pa.table({"id": [1]}))) + try: + assert await asyncio.to_thread(started.wait, 5) + for _ in range(cancel_count): + task.cancel() + await asyncio.sleep(0) + paths = await asyncio.to_thread(lambda: list(tmp_path.rglob("*.tsv"))) + assert len(paths) == 1 + assert await asyncio.to_thread(paths[0].read_bytes) == b"1\n" + release.set() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(asyncio.shield(task), timeout=5) + assert await asyncio.to_thread(finished.wait, 5) + assert await asyncio.to_thread(lambda: list(tmp_path.iterdir())) == [] + assert conn._cursor.execute_calls == [] + finally: + release.set() + await asyncio.gather(task, return_exceptions=True) + await asyncio.to_thread(finished.wait, 5) + + async def test_local_infile_empty_table_does_not_send_payload() -> None: conn = _FakeConnection() driver = AsyncmyDriver( From eb7e8ca1d7dd1de8ce2cce1c1e2a4d71dfbbf5ab Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 22:05:30 +0000 Subject: [PATCH 141/174] fix(adbc): preserve pool validation in compiled builds --- sqlspec/adapters/adbc/pool.py | 7 +++++-- tests/conftest.py | 8 +++++++- tests/unit/adapters/test_adbc/test_pool.py | 2 +- 3 files changed, 13 insertions(+), 4 deletions(-) diff --git a/sqlspec/adapters/adbc/pool.py b/sqlspec/adapters/adbc/pool.py index 8a65a1558..764b3b1e9 100644 --- a/sqlspec/adapters/adbc/pool.py +++ b/sqlspec/adapters/adbc/pool.py @@ -38,9 +38,12 @@ class AdbcConnectionPool: ) def __init__( - self, create_connection: "Callable[[], AdbcConnection]", *, max_size: int = 4, timeout: float = 30.0 + self, create_connection: "Callable[[], AdbcConnection]", *, max_size: object = 4, timeout: float = 30.0 ) -> None: - """Configure a lazy pool without opening any native connections.""" + """Configure a lazy pool without opening any native connections. + + Raw capacity validation runs before compiled integer argument coercion. + """ if isinstance(max_size, bool) or not isinstance(max_size, int) or max_size <= 0: msg = "ADBC pool max_size must be a positive integer." raise ImproperConfigurationError(msg) diff --git a/tests/conftest.py b/tests/conftest.py index 82a709a25..b5b0ccd34 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -103,6 +103,12 @@ def pytest_addoption(parser: pytest.Parser) -> None: default=False, help="Run Spanner integration tests locally (enabled by default in CI; local runs accept --run-spanner-tests or SQLSPEC_ENABLE_SPANNER_TESTS=1)", ) + parser.addoption( + "--run-compiled-adbc", + action="store_true", + default=False, + help="Enable targeted ADBC validation against an installed compiled wheel; explicit interpreted-only skips remain active.", + ) def pytest_collection_modifyitems(config: pytest.Config, items: list[pytest.Item]) -> None: @@ -130,7 +136,7 @@ def pytest_collection_modifyitems(config: pytest.Config, items: list[pytest.Item if is_spanner and not spanner_enabled: item.add_marker(skip_spanner) - if not is_compiled(): + if not is_compiled() or config.getoption("--run-compiled-adbc", default=False): return skip_adbc = pytest.mark.skip(reason="Skip ADBC tests when running against mypyc-compiled modules.") diff --git a/tests/unit/adapters/test_adbc/test_pool.py b/tests/unit/adapters/test_adbc/test_pool.py index 731d73f03..c395864ae 100644 --- a/tests/unit/adapters/test_adbc/test_pool.py +++ b/tests/unit/adapters/test_adbc/test_pool.py @@ -202,7 +202,7 @@ def test_adbc_pool_exclusive_original_seed_leases_and_nested_timeout() -> None: @pytest.mark.parametrize("max_size", [0, -1, True, 1.5]) -def test_adbc_pool_rejects_invalid_capacity(max_size: int) -> None: +def test_adbc_pool_rejects_invalid_capacity(max_size: object) -> None: """Capacity is a positive integer and rejects bool before native acquisition.""" create = MagicMock() with pytest.raises(ImproperConfigurationError): From 8f150d658af6499a100f2c133f9acb561ed1b6ce Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 22:05:30 +0000 Subject: [PATCH 142/174] fix(oracledb): close standalone connections after initialization failure --- sqlspec/adapters/oracledb/config.py | 14 +++- .../adapters/test_oracledb/test_config.py | 67 ++++++++++++++++++- 2 files changed, 78 insertions(+), 3 deletions(-) diff --git a/sqlspec/adapters/oracledb/config.py b/sqlspec/adapters/oracledb/config.py index 5962f7fce..0cb129a41 100644 --- a/sqlspec/adapters/oracledb/config.py +++ b/sqlspec/adapters/oracledb/config.py @@ -573,7 +573,12 @@ def create_connection(self) -> "OracleSyncConnection": _initialize_oracle_client(self.connection_config) config = _build_standalone_connection_params(self.connection_config) conn = oracledb.connect(**config) - self._init_connection(conn, "") + try: + self._init_connection(conn, "") + except BaseException: + with contextlib.suppress(BaseException): + conn.close() + raise return conn def provide_pool(self) -> "OracleSyncConnectionPool": @@ -870,7 +875,12 @@ async def create_connection(self) -> OracleAsyncConnection: _validate_async_oracle_mode(self.connection_config) config = _build_standalone_connection_params(self.connection_config) conn = await oracledb.connect_async(**config) - await self._init_connection(conn, "") + try: + await self._init_connection(conn, "") + except BaseException: + with contextlib.suppress(BaseException): + await conn.close() + raise return cast("OracleAsyncConnection", conn) async def provide_pool(self) -> "OracleAsyncConnectionPool": diff --git a/tests/unit/adapters/test_oracledb/test_config.py b/tests/unit/adapters/test_oracledb/test_config.py index 0c1b82364..28af3fac4 100644 --- a/tests/unit/adapters/test_oracledb/test_config.py +++ b/tests/unit/adapters/test_oracledb/test_config.py @@ -1,5 +1,6 @@ """OracleDB configuration tests covering driver kwargs and typed options.""" +import asyncio import io import os import ssl @@ -9,6 +10,7 @@ from inspect import isawaitable from pathlib import Path from ssl import TLSVersion +from types import SimpleNamespace from typing import Any, NotRequired, cast, get_args, get_origin, get_type_hints from unittest.mock import AsyncMock, MagicMock, Mock, call @@ -67,6 +69,11 @@ def _stub_sync_connection_setup(monkeypatch: pytest.MonkeyPatch, calls: list[str monkeypatch.setattr(oracle_config_module, "register_uuid_handlers", lambda _connection: calls.append("uuid")) +async def _create_standalone_oracle_connection(config: OracleSyncConfig | OracleAsyncConfig) -> object: + result = config.create_connection() + return await result if isawaitable(result) else result + + def test_oracle_connection_params_expose_current_driver_options() -> None: """Connection params should mirror current python-oracledb connection knobs.""" annotations = _oracle_config_hints(OracleConnectionParams) @@ -370,6 +377,7 @@ def test_sync_create_connection_unpooled(monkeypatch: pytest.MonkeyPatch) -> Non config = OracleSyncConfig(connection_config={"user": "scott", "password": "tiger", "dsn": "localhost/xe"}) conn = config.create_connection() assert conn is connect_mock.return_value + connect_mock.return_value.close.assert_not_called() create_pool_mock.assert_not_called() assert config.connection_instance is None @@ -384,10 +392,62 @@ async def test_async_create_connection_unpooled(monkeypatch: pytest.MonkeyPatch) config = OracleAsyncConfig(connection_config={"user": "scott", "password": "tiger", "dsn": "localhost/xe"}) conn = await config.create_connection() assert conn is connect_mock.return_value + connect_mock.return_value.close.assert_not_called() create_pool_mock.assert_not_called() assert config.connection_instance is None +@pytest.mark.parametrize("is_async", [False, True]) +@pytest.mark.parametrize("hook_kind", ["on_connection_create", "session_callback"]) +@pytest.mark.parametrize("cancelled", [False, True]) +@pytest.mark.parametrize("cleanup_failure", [None, OSError("close failed"), asyncio.CancelledError("close cancelled")]) +async def test_standalone_connection_setup_failure_closes_owned_connection( + is_async: bool, + hook_kind: str, + cancelled: bool, + cleanup_failure: BaseException | None, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Configured setup failures close new SDK connections without replacing the primary error.""" + failure = asyncio.CancelledError("setup cancelled") if cancelled else RuntimeError("setup failed") + close = AsyncMock(side_effect=cleanup_failure) if is_async else Mock(side_effect=cleanup_failure) + connection = SimpleNamespace(version="23.5.0.0.0", inputtypehandler=None, outputtypehandler=None, close=close) + connect = AsyncMock(return_value=connection) if is_async else Mock(return_value=connection) + monkeypatch.setattr(oracle_config_module.oracledb, "connect_async" if is_async else "connect", connect) + + async def fail_async_setup() -> None: + await asyncio.sleep(0) + raise failure + + def setup_hook(acquired: object, tag: str) -> Awaitable[None]: + assert acquired is connection + assert tag == "" + assert callable(connection.inputtypehandler) + assert callable(connection.outputtypehandler) + if is_async: + return fail_async_setup() + raise failure + + connection_options: dict[str, Any] = {"dsn": "localhost/xe"} + driver_features: dict[str, Any] = {} + if hook_kind == "session_callback": + connection_options[hook_kind] = setup_hook + else: + driver_features[hook_kind] = setup_hook + config_type = OracleAsyncConfig if is_async else OracleSyncConfig + config = config_type(connection_config=connection_options, driver_features=driver_features) + + with pytest.raises(type(failure)) as caught: + await _create_standalone_oracle_connection(config) + + assert caught.value is failure + if is_async: + cast("AsyncMock", close).assert_awaited_once_with() + else: + close.assert_called_once_with() + assert config.connection_instance is None + + @pytest.mark.anyio @pytest.mark.parametrize("is_async", [False, True]) async def test_standalone_connections_filter_pool_options(is_async: bool, monkeypatch: pytest.MonkeyPatch) -> None: @@ -428,8 +488,13 @@ def test_standalone_sync_connection_initializes_thick_mode( options: dict[str, Any], monkeypatch: pytest.MonkeyPatch ) -> None: calls: list[str] = [] + + def fake_connect(**_kwargs: object) -> _StubConnection: + calls.append("connect") + return _StubConnection() + initialize = Mock(side_effect=lambda **_kwargs: calls.append("initialize")) - connect = Mock(side_effect=lambda **_kwargs: calls.append("connect") or _StubConnection()) + connect = Mock(side_effect=fake_connect) monkeypatch.setattr(oracle_config_module.oracledb, "is_thin_mode", lambda: True) monkeypatch.setattr(oracle_config_module.oracledb, "init_oracle_client", initialize) monkeypatch.setattr(oracle_config_module.oracledb, "connect", connect) From 2dc00198119bca4e6bec3b84fe362ebd9afe3a35 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 22:09:16 +0000 Subject: [PATCH 143/174] fix(spanner): preserve checkout cleanup and stream exception mapping --- sqlspec/adapters/spanner/config.py | 66 ++++++---- sqlspec/adapters/spanner/core.py | 77 +++++------- .../unit/adapters/test_spanner/test_config.py | 118 +++++++++++++++++- tests/unit/adapters/test_spanner/test_core.py | 16 +-- .../test_spanner_query_options.py | 99 ++++++++++++++- 5 files changed, 292 insertions(+), 84 deletions(-) diff --git a/sqlspec/adapters/spanner/config.py b/sqlspec/adapters/spanner/config.py index 31e54e4e5..5828fde9d 100644 --- a/sqlspec/adapters/spanner/config.py +++ b/sqlspec/adapters/spanner/config.py @@ -291,22 +291,29 @@ def __enter__(self) -> "SpannerSyncConnection": try: txn = self._session.transaction() txn.__enter__() - except Exception: - manager.put_session(self._session) + except BaseException: + with contextlib.suppress(BaseException): + manager.put_session(self._session) self._session = None raise self._connection = cast("SpannerSyncConnection", txn) - return self._connection - - self._session = cast("Any", database).snapshot(multi_use=True) - self._connection = cast("SpannerSyncConnection", self._session.__enter__()) - hook = getattr(self._config, "_user_connection_hook", None) or getattr( - self._config, "on_connection_create", None - ) - if hook: - modified = hook(self._connection) - if modified is not None: - self._connection = cast("SpannerSyncConnection", modified) + else: + self._session = cast("Any", database).snapshot(multi_use=True) + self._connection = cast("SpannerSyncConnection", self._session.__enter__()) + try: + hook = getattr(self._config, "_user_connection_hook", None) or getattr( + self._config, "on_connection_create", None + ) + if hook: + modified = hook(self._connection) + if modified is not None: + self._connection = cast("SpannerSyncConnection", modified) + except BaseException as exc: + with contextlib.suppress(BaseException): + self.__exit__(type(exc), exc, exc.__traceback__) + self._connection = None + self._session = None + raise return self._connection def __exit__( @@ -702,22 +709,29 @@ async def __aenter__(self) -> "SpannerAsyncConnection": try: txn = self._session.transaction() await txn.__aenter__() - except Exception: - await manager.put_session(self._session) + except BaseException: + with contextlib.suppress(BaseException): + await manager.put_session(self._session) self._session = None raise self._connection = cast("SpannerAsyncConnection", txn) - return self._connection - - self._session = cast("Any", database).snapshot(multi_use=True) - self._connection = cast("SpannerAsyncConnection", await self._session.__aenter__()) - hook = getattr(self._config, "_user_connection_hook", None) or getattr( - self._config, "on_connection_create", None - ) - if hook: - modified = hook(self._connection) - if modified is not None: - self._connection = cast("SpannerAsyncConnection", modified) + else: + self._session = cast("Any", database).snapshot(multi_use=True) + self._connection = cast("SpannerAsyncConnection", await self._session.__aenter__()) + try: + hook = getattr(self._config, "_user_connection_hook", None) or getattr( + self._config, "on_connection_create", None + ) + if hook: + modified = hook(self._connection) + if modified is not None: + self._connection = cast("SpannerAsyncConnection", modified) + except BaseException as exc: + with contextlib.suppress(BaseException): + await self.__aexit__(type(exc), exc, exc.__traceback__) + self._connection = None + self._session = None + raise return self._connection async def __aexit__( diff --git a/sqlspec/adapters/spanner/core.py b/sqlspec/adapters/spanner/core.py index ac48aec9f..4dec2403d 100644 --- a/sqlspec/adapters/spanner/core.py +++ b/sqlspec/adapters/spanner/core.py @@ -2,7 +2,6 @@ import contextlib import re -import sys from collections.abc import Callable, Iterable, Mapping from functools import partial from itertools import islice @@ -25,6 +24,7 @@ build_statement_config_from_profile, ) from sqlspec.driver import rows_to_dicts +from sqlspec.driver._exception_handler import _run_with_async_exception_handler, _run_with_sync_exception_handler from sqlspec.exceptions import ( DataError, DeadlockError, @@ -915,33 +915,23 @@ def __init__( def start(self) -> None: handler = self._driver.handle_database_exceptions() - if sys.exc_info()[1] is None: - with handler: - result_set = self._driver.connection.execute_sql( - self._sql, params=self._params, param_types=self._param_types, **self._execute_kwargs - ) - self._result_set = result_set - self._row_iterator = cast("Generator[Any, None, None]", iter(result_set)) - self._driver._check_pending_exception(handler) - else: - result_set = self._driver.connection.execute_sql( - self._sql, params=self._params, param_types=self._param_types, **self._execute_kwargs - ) - self._result_set = result_set - self._row_iterator = cast("Generator[Any, None, None]", iter(result_set)) + _run_with_sync_exception_handler(handler, self._start_query) + self._driver._check_pending_exception(handler) + + def _start_query(self) -> None: + result_set = self._driver.connection.execute_sql( + self._sql, params=self._params, param_types=self._param_types, **self._execute_kwargs + ) + self._result_set = result_set + self._row_iterator = cast("Generator[Any, None, None]", iter(result_set)) def fetch_chunk(self) -> "list[dict[str, Any]]": row_iterator = self._row_iterator if self._result_set is None or row_iterator is None: return [] handler = self._driver.handle_database_exceptions() - rows: list[Any] = [] - if sys.exc_info()[1] is None: - with handler: - rows = list(islice(row_iterator, self._chunk_size)) - self._driver._check_pending_exception(handler) - else: - rows = list(islice(row_iterator, self._chunk_size)) + rows: list[Any] | None = _run_with_sync_exception_handler(handler, list, islice(row_iterator, self._chunk_size)) + self._driver._check_pending_exception(handler) if not rows: return [] return self._rows_to_chunk(rows) @@ -973,41 +963,34 @@ def __init__( async def start(self) -> None: handler = self._driver.handle_database_exceptions() - if sys.exc_info()[1] is None: - async with handler: - result_set = await self._driver.connection.execute_sql( - self._sql, params=self._params, param_types=self._param_types, **self._execute_kwargs - ) - self._result_set = result_set - self._row_iterator = cast("AsyncGenerator[Any, None]", result_set.__aiter__()) - self._driver._check_pending_exception(handler) - else: - result_set = await self._driver.connection.execute_sql( - self._sql, params=self._params, param_types=self._param_types, **self._execute_kwargs - ) - self._result_set = result_set - self._row_iterator = cast("AsyncGenerator[Any, None]", result_set.__aiter__()) + await _run_with_async_exception_handler(handler, self._start_query) + self._driver._check_pending_exception(handler) + + async def _start_query(self) -> None: + result_set = await self._driver.connection.execute_sql( + self._sql, params=self._params, param_types=self._param_types, **self._execute_kwargs + ) + self._result_set = result_set + self._row_iterator = cast("AsyncGenerator[Any, None]", result_set.__aiter__()) async def fetch_chunk(self) -> "list[dict[str, Any]]": row_iterator = self._row_iterator if self._result_set is None or row_iterator is None: return [] handler = self._driver.handle_database_exceptions() - rows: list[Any] = [] - if sys.exc_info()[1] is None: - async with handler: - with contextlib.suppress(StopAsyncIteration): - while len(rows) < self._chunk_size: - rows.append(await row_iterator.__anext__()) - self._driver._check_pending_exception(handler) - else: - with contextlib.suppress(StopAsyncIteration): - while len(rows) < self._chunk_size: - rows.append(await row_iterator.__anext__()) + rows = await _run_with_async_exception_handler(handler, self._fetch_rows, row_iterator) + self._driver._check_pending_exception(handler) if not rows: return [] return self._rows_to_chunk(rows) + async def _fetch_rows(self, row_iterator: "AsyncGenerator[Any, None]") -> "list[Any]": + rows: list[Any] = [] + with contextlib.suppress(StopAsyncIteration): + while len(rows) < self._chunk_size: + rows.append(await row_iterator.__anext__()) + return rows + async def close(self, error: bool = False) -> None: row_iterator = self._row_iterator self._row_iterator = None diff --git a/tests/unit/adapters/test_spanner/test_config.py b/tests/unit/adapters/test_spanner/test_config.py index 06f036518..87a708cba 100644 --- a/tests/unit/adapters/test_spanner/test_config.py +++ b/tests/unit/adapters/test_spanner/test_config.py @@ -1,7 +1,10 @@ +from asyncio import CancelledError +from collections.abc import AsyncIterator, Callable +from contextlib import asynccontextmanager from datetime import timedelta from types import SimpleNamespace from typing import TYPE_CHECKING, Any, Self, cast -from unittest.mock import AsyncMock +from unittest.mock import ANY, AsyncMock, MagicMock, Mock import pytest from google.api_core import exceptions as api_exceptions @@ -43,6 +46,119 @@ pytestmark = requires_interpreted +def _assert_sdk_call(mock: Mock, *args: object) -> None: + mock.assert_called_once_with(*args) + if isinstance(mock, AsyncMock): + mock.assert_awaited_once_with(*args) + + +def _hook_checkout(is_async: bool, hook: Callable[[Any], Any]) -> SimpleNamespace: + connection = MagicMock(committed=None, rolled_back=False, _transaction_id=b"txn", _mutations=[]) + snapshot = MagicMock() + if is_async: + connection.__aenter__ = AsyncMock(return_value=connection) + connection.commit = AsyncMock() + connection.rollback = AsyncMock() + snapshot.__aenter__ = AsyncMock(return_value=connection) + snapshot.__aexit__ = AsyncMock(return_value=False) + config_type = SpannerAsyncConfig + mock_type = AsyncMock + else: + connection.__enter__.return_value = connection + snapshot.__enter__.return_value = connection + snapshot.__exit__.return_value = False + config_type = SpannerSyncConfig + mock_type = Mock + session = SimpleNamespace(transaction=Mock(return_value=connection)) + manager = SimpleNamespace(get_session=mock_type(return_value=session), put_session=mock_type()) + database = SimpleNamespace(snapshot=Mock(return_value=snapshot), sessions_manager=manager) + config = config_type( + connection_config={"project": "p", "instance_id": "i", "database_id": "d"}, + connection_instance=cast("Any", object()), + on_connection_create=hook, + ) + config._database = cast("Any", database) + return SimpleNamespace(config=config, connection=connection, snapshot=snapshot, manager=manager, session=session) + + +@asynccontextmanager +async def _hook_session(config: Any, **kwargs: Any) -> AsyncIterator[Any]: + if isinstance(config, SpannerAsyncConfig): + async with config.provide_session(**kwargs) as session: + yield session + else: + with config.provide_session(**kwargs) as session: + yield session + + +@pytest.mark.parametrize("is_async", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize("transaction", [False, True], ids=["snapshot", "default-transaction"]) +async def test_connection_hook_runs_for_session_checkout(is_async: bool, transaction: bool) -> None: + hook = Mock(return_value=None) + checkout = _hook_checkout(is_async, hook) + kwargs = {} if transaction else {"transaction": False} + + async with _hook_session(checkout.config, **kwargs) as session: + assert session.connection is checkout.connection + hook.assert_called_once_with(checkout.connection) + + if transaction: + _assert_sdk_call(checkout.connection.commit) + _assert_sdk_call(checkout.manager.put_session, checkout.session) + else: + snapshot_exit = checkout.snapshot.__aexit__ if is_async else checkout.snapshot.__exit__ + _assert_sdk_call(snapshot_exit, None, None, None) + + +@pytest.mark.parametrize("is_async", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize("transaction", [False, True], ids=["snapshot", "transaction"]) +@pytest.mark.parametrize("error_type", [ValueError, CancelledError]) +@pytest.mark.parametrize("cleanup_fails", [False, True], ids=["cleanup-ok", "cleanup-fails"]) +async def test_connection_hook_failure_releases_checkout_and_preserves_error( + is_async: bool, transaction: bool, error_type: type[BaseException], cleanup_fails: bool +) -> None: + error = error_type("hook failed") + hook = Mock(side_effect=error) + checkout = _hook_checkout(is_async, hook) + snapshot_exit = checkout.snapshot.__aexit__ if is_async else checkout.snapshot.__exit__ + if cleanup_fails: + cleanup = checkout.connection.rollback if transaction else snapshot_exit + cleanup.side_effect = RuntimeError("cleanup failed") + + with pytest.raises(error_type) as caught: + async with _hook_session(checkout.config, transaction=transaction): + pytest.fail("A failing connection hook must prevent the session from opening") + + assert caught.value is error + hook.assert_called_once_with(checkout.connection) + if transaction: + _assert_sdk_call(checkout.connection.rollback) + checkout.connection.commit.assert_not_called() + _assert_sdk_call(checkout.manager.put_session, checkout.session) + else: + _assert_sdk_call(snapshot_exit, error_type, error, ANY) + + +@pytest.mark.parametrize("is_async", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize("cleanup_fails", [False, True], ids=["cleanup-ok", "cleanup-fails"]) +async def test_transaction_enter_cancellation_returns_checkout(is_async: bool, cleanup_fails: bool) -> None: + hook = Mock(return_value=None) + checkout = _hook_checkout(is_async, hook) + error = CancelledError("checkout cancelled") + enter = checkout.connection.__aenter__ if is_async else checkout.connection.__enter__ + enter.side_effect = error + if cleanup_fails: + checkout.manager.put_session.side_effect = RuntimeError("release failed") + + with pytest.raises(CancelledError) as caught: + async with _hook_session(checkout.config): + pytest.fail("A cancelled checkout must prevent the session from opening") + + assert caught.value is error + hook.assert_not_called() + _assert_sdk_call(checkout.manager.put_session, checkout.session) + + class _DummyDriver(SyncDriverAdapterBase): dialect = "spanner" diff --git a/tests/unit/adapters/test_spanner/test_core.py b/tests/unit/adapters/test_spanner/test_core.py index 1fb2ef5fc..34b91e9ca 100644 --- a/tests/unit/adapters/test_spanner/test_core.py +++ b/tests/unit/adapters/test_spanner/test_core.py @@ -18,7 +18,7 @@ from sqlspec.adapters.spanner._typing import SpannerSyncSessionContext from sqlspec.adapters.spanner.adk.artifact_store import SpannerAsyncADKArtifactStore, SpannerSyncADKArtifactStore from sqlspec.adapters.spanner.adk.store import SpannerSyncADKMemoryStore, SpannerSyncADKStore -from sqlspec.adapters.spanner.config import SpannerSyncConfig, SpannerSyncConnectionContext +from sqlspec.adapters.spanner.config import SpannerSyncConfig from sqlspec.adapters.spanner.core import ( build_param_type_signature, build_statement_config, @@ -328,16 +328,16 @@ def rollback(self) -> None: sessions_manager.get_session.return_value = session_obj database = SimpleNamespace(sessions_manager=sessions_manager) - config = MagicMock() - config.get_database.return_value = database - config.driver_features = {} + config = SpannerSyncConfig( + connection_config={"project": "p", "instance_id": "i", "database_id": "d"}, + connection_instance=cast("Any", object()), + ) + config._database = cast("Any", database) - ctx = SpannerSyncConnectionContext(config, transaction=True) - with pytest.raises(DeadlockError), ctx: + with pytest.raises(DeadlockError), config.provide_connection(transaction=True): pass - ctx_rollback = SpannerSyncConnectionContext(config, transaction=True) - with pytest.raises(RuntimeError, match="user boom"), ctx_rollback: + with pytest.raises(RuntimeError, match="user boom"), config.provide_connection(transaction=True): raise RuntimeError("user boom") diff --git a/tests/unit/adapters/test_spanner/test_spanner_query_options.py b/tests/unit/adapters/test_spanner/test_spanner_query_options.py index dcf9aa0c1..bde1d4b49 100644 --- a/tests/unit/adapters/test_spanner/test_spanner_query_options.py +++ b/tests/unit/adapters/test_spanner/test_spanner_query_options.py @@ -1,10 +1,105 @@ """Unit tests for Spanner QueryOptions forwarding.""" -from unittest.mock import MagicMock +from collections.abc import AsyncIterator, Iterator +from types import SimpleNamespace +from typing import Any +from unittest.mock import AsyncMock, MagicMock + +import pytest +from google.api_core import exceptions as api_exceptions +from google.cloud.spanner_v1.types.type import TypeCode from sqlspec.adapters.spanner.config import SpannerSyncConfig from sqlspec.adapters.spanner.core import default_statement_config -from sqlspec.adapters.spanner.driver import SpannerSyncDriver +from sqlspec.adapters.spanner.driver import SpannerAsyncDriver, SpannerSyncDriver +from sqlspec.exceptions import OperationalError + + +class _StreamResult: + def __init__(self, error: BaseException | None) -> None: + self.error = error + self.closed = False + self.metadata = SimpleNamespace( + row_type=SimpleNamespace(fields=[SimpleNamespace(name="id", type_=SimpleNamespace(code=TypeCode.INT64))]) + ) + + def __iter__(self) -> Iterator[list[int]]: + try: + yield [1] + if self.error is not None: + raise self.error + yield [2] + finally: + self.closed = True + + async def __aiter__(self) -> AsyncIterator[list[int]]: + try: + yield [1] + if self.error is not None: + raise self.error + yield [2] + finally: + self.closed = True + + +async def _consume_stream(driver: Any, is_async: bool, *, fails_fetch: bool = False) -> list[dict[str, Any]]: + stream = driver.select_stream("SELECT id FROM users", chunk_size=1) + if is_async: + async with stream: + if fails_fetch: + assert await anext(stream) == {"id": 1} + return [row async for row in stream] + with stream: + if fails_fetch: + assert next(stream) == {"id": 1} + return list(stream) + + +@pytest.mark.parametrize("is_async", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize("phase", ["start", "fetch"]) +@pytest.mark.parametrize("ambient", [False, True], ids=["normal", "application-except"]) +async def test_stream_maps_sdk_errors_with_ambient_application_exception( + is_async: bool, phase: str, ambient: bool +) -> None: + error = api_exceptions.ServiceUnavailable("Spanner unavailable") + result = _StreamResult(error) + execute = AsyncMock() if is_async else MagicMock() + execute.side_effect = error if phase == "start" else None + execute.return_value = result + connection = SimpleNamespace(execute_sql=execute) + driver = SpannerAsyncDriver(connection=connection) if is_async else SpannerSyncDriver(connection=connection) + + with pytest.raises(OperationalError, match="Spanner unavailable") as caught: + try: + if ambient: + raise LookupError("application error") + except LookupError: + await _consume_stream(driver, is_async, fails_fetch=phase == "fetch") + else: + await _consume_stream(driver, is_async, fails_fetch=phase == "fetch") + + assert caught.value.__cause__ is error + execute.assert_called_once() + if phase == "fetch": + assert result.closed is True + + +@pytest.mark.parametrize("is_async", [False, True], ids=["sync", "async"]) +async def test_stream_succeeds_with_ambient_application_exception(is_async: bool) -> None: + result = _StreamResult(None) + execute = AsyncMock(return_value=result) if is_async else MagicMock(return_value=result) + connection = SimpleNamespace(execute_sql=execute) + driver = SpannerAsyncDriver(connection=connection) if is_async else SpannerSyncDriver(connection=connection) + error = LookupError("application error") + + try: + raise error + except LookupError as caught: + rows = await _consume_stream(driver, is_async) + assert caught is error + + assert rows == [{"id": 1}, {"id": 2}] + assert result.closed is True def test_driver_execute_with_statement_query_options() -> None: From b3f10b62194197ced5ef7d0ecc9d47792078fdbd Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 22:09:16 +0000 Subject: [PATCH 144/174] fix(mssql): scope table details to the requested table --- .../adapters/mssql_python/data_dictionary.py | 2 +- sqlspec/adapters/pymssql/data_dictionary.py | 2 +- .../data_dictionary/test_mssql_metadata.py | 81 +++++++++++++++++++ 3 files changed, 83 insertions(+), 2 deletions(-) diff --git a/sqlspec/adapters/mssql_python/data_dictionary.py b/sqlspec/adapters/mssql_python/data_dictionary.py index 47b97f5b0..aa0d1ef48 100644 --- a/sqlspec/adapters/mssql_python/data_dictionary.py +++ b/sqlspec/adapters/mssql_python/data_dictionary.py @@ -298,7 +298,7 @@ def get_table_details(self, driver: "MssqlPythonDriver", table: str, schema: str capability = self.get_metadata_capabilities(driver, ("tables",)).get("tables") self._log_schema_introspect(driver, schema_name=schema_name, table_name=table, operation="table_details") rows = driver.select( - self.get_domain_query("tables", "all_by_schema"), schema_name=schema_name, table_name=table + self.get_domain_query("tables", "details_by_table"), schema_name=schema_name, table_name=table ) return MetadataResult("tables", capability=capability, items=tuple(rows)) diff --git a/sqlspec/adapters/pymssql/data_dictionary.py b/sqlspec/adapters/pymssql/data_dictionary.py index aaabeb2c4..2365dae0e 100644 --- a/sqlspec/adapters/pymssql/data_dictionary.py +++ b/sqlspec/adapters/pymssql/data_dictionary.py @@ -298,7 +298,7 @@ def get_table_details(self, driver: "PymssqlDriver", table: str, schema: str | N capability = self.get_metadata_capabilities(driver, ("tables",)).get("tables") self._log_schema_introspect(driver, schema_name=schema_name, table_name=table, operation="table_details") rows = driver.select( - self.get_domain_query("tables", "all_by_schema"), schema_name=schema_name, table_name=table + self.get_domain_query("tables", "details_by_table"), schema_name=schema_name, table_name=table ) return MetadataResult("tables", capability=capability, items=tuple(rows)) diff --git a/tests/unit/data_dictionary/test_mssql_metadata.py b/tests/unit/data_dictionary/test_mssql_metadata.py index 7f6679551..789828318 100644 --- a/tests/unit/data_dictionary/test_mssql_metadata.py +++ b/tests/unit/data_dictionary/test_mssql_metadata.py @@ -1,11 +1,15 @@ """SQL Server data-dictionary metadata tests.""" +import sqlite3 +from collections.abc import Iterator from typing import Any, cast import pytest from sqlspec.adapters.mssql_python.data_dictionary import MssqlPythonSyncDataDictionary +from sqlspec.adapters.mssql_python.driver import MssqlPythonDriver from sqlspec.adapters.pymssql.data_dictionary import PymssqlSyncDataDictionary +from sqlspec.adapters.pymssql.driver import PymssqlDriver from sqlspec.data_dictionary import ( DataDictionaryLoader, DDLResult, @@ -16,6 +20,7 @@ SystemMetadataRequest, SystemMetadataResult, ) +from sqlspec.driver import SyncDataDictionaryBase REQUIRED_MSSQL_DOMAIN_QUERIES = { "schemas": ("list",), @@ -61,6 +66,82 @@ def select(self, statement: Any, **kwargs: Any) -> list[dict[str, Any]]: return [] +class CatalogCursor(sqlite3.Cursor): + """Execute SDK-bound SQL against a local substitute SQL Server catalog.""" + + def execute(self, sql: str, parameters: Any = ()) -> "CatalogCursor": + super().execute(sql.replace("%s", "?"), parameters) + return self + + +class CatalogConnection(sqlite3.Connection): + """Provide the DB-API cursor surface shared by the SQL Server SDKs.""" + + def cursor(self, *args: Any, **kwargs: Any) -> CatalogCursor: + return super().cursor(factory=CatalogCursor) + + +@pytest.fixture(params=[MssqlPythonDriver, PymssqlDriver], ids=["mssql_python", "pymssql"]) +def catalog_driver(request: pytest.FixtureRequest) -> Iterator[MssqlPythonDriver | PymssqlDriver]: + """Run real drivers and dictionary SQL with only the external SDK substituted.""" + connection = sqlite3.connect(":memory:", factory=CatalogConnection) + connection.create_function("DB_NAME", 0, lambda: "app") + connection.create_function("SCHEMA_NAME", 0, lambda: "dbo") + connection.executescript(""" + ATTACH DATABASE ':memory:' AS sys; + CREATE TABLE sys.schemas (schema_id INTEGER, name TEXT); + CREATE TABLE sys.tables ( + object_id INTEGER, schema_id INTEGER, name TEXT, create_date TEXT, modify_date TEXT, + is_memory_optimized INTEGER, durability INTEGER, durability_desc TEXT, + temporal_type INTEGER, temporal_type_desc TEXT, history_table_id INTEGER + ); + INSERT INTO sys.schemas VALUES (1, 'dbo'); + INSERT INTO sys.tables VALUES + (10, 1, 'accounts', '2026-01-01', '2026-02-01', 0, 0, 'SCHEMA_AND_DATA', + 2, 'SYSTEM_VERSIONED_TEMPORAL_TABLE', 11), + (11, 1, 'accounts_history', '2026-01-01', '2026-01-01', 0, 0, 'SCHEMA_AND_DATA', + 1, 'HISTORY_TABLE', NULL); + """) + try: + yield request.param(cast(Any, connection)) + finally: + connection.close() + + +@pytest.mark.parametrize("schema", [None, "dbo"]) +def test_table_details_selects_requested_table_with_rich_metadata( + catalog_driver: MssqlPythonDriver | PymssqlDriver, schema: str | None +) -> None: + """A detail lookup preserves table identity, rich fields, and missing-table behavior.""" + dictionary = cast("SyncDataDictionaryBase", catalog_driver.data_dictionary) + result = dictionary.get_table_details(catalog_driver, "accounts", schema) + + assert result.items == ( + { + "catalog_name": "app", + "schema_name": "dbo", + "table_name": "accounts", + "table_type": "BASE TABLE", + "create_date": "2026-01-01", + "modify_date": "2026-02-01", + "is_memory_optimized": 0, + "durability": 0, + "durability_desc": "SCHEMA_AND_DATA", + "temporal_type": 2, + "temporal_type_desc": "SYSTEM_VERSIONED_TEMPORAL_TABLE", + "history_schema_name": "dbo", + "history_table_name": "accounts_history", + }, + ) + history = dictionary.get_table_details(catalog_driver, "accounts_history", schema) + assert len(history.items) == 1 + history_row = cast("dict[str, object]", history.items[0]) + assert history_row["table_name"] == "accounts_history" + assert history_row["history_table_name"] is None + missing = dictionary.get_table_details(catalog_driver, "missing", schema) + assert missing.items == () + + def test_mssql_domain_query_pack_supports_required_catalog_queries() -> None: """SQL Server query packs should load from direct domain paths.""" loader = DataDictionaryLoader() From 8f552fa6c23b8ad4ac2e321c69a22ea254490717 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 22:09:16 +0000 Subject: [PATCH 145/174] fix(pymssql): retain nesting state after failed acquisition --- sqlspec/adapters/pymssql/pool.py | 2 +- tests/unit/adapters/test_pymssql/test_pool.py | 52 +++++++++++++++++++ 2 files changed, 53 insertions(+), 1 deletion(-) diff --git a/sqlspec/adapters/pymssql/pool.py b/sqlspec/adapters/pymssql/pool.py index c516f3d90..3b61c5350 100644 --- a/sqlspec/adapters/pymssql/pool.py +++ b/sqlspec/adapters/pymssql/pool.py @@ -95,8 +95,8 @@ def get_connection(self) -> "Generator[PymssqlConnection, None, None]": A thread-local database connection. """ depth = getattr(self._thread_local, "depth", 0) - self._thread_local.depth = depth + 1 connection = self.acquire() + self._thread_local.depth = depth + 1 try: yield connection except Exception: diff --git a/tests/unit/adapters/test_pymssql/test_pool.py b/tests/unit/adapters/test_pymssql/test_pool.py index 2586dac8c..abae8e6a7 100644 --- a/tests/unit/adapters/test_pymssql/test_pool.py +++ b/tests/unit/adapters/test_pymssql/test_pool.py @@ -3,6 +3,8 @@ import threading from typing import Any, cast +import pytest + from tests.unit.adapters.test_pymssql._fakes import FakeConnection, FakePymssqlModule @@ -210,3 +212,53 @@ def test_pool_reentrant_thread_checkout(monkeypatch: Any) -> None: assert not connection.closed pool.close() assert connection.closed + + +def test_pool_recovers_after_failed_acquisition_and_retires_failed_checkout(monkeypatch: pytest.MonkeyPatch) -> None: + """A failed connect must not prevent a later failed checkout from retiring its connection.""" + import sqlspec.adapters.pymssql.pool as pool_module + from sqlspec.adapters.pymssql.pool import PymssqlConnectionPool + + dirty_connection = FakeConnection() + clean_connection = FakeConnection() + attempts: list[OSError | FakeConnection] = [ + OSError("temporary connect failure"), + dirty_connection, + clean_connection, + ] + + class RecoveringModule(FakePymssqlModule): + def connect(self, **kwargs: Any) -> FakeConnection: + outcome = attempts.pop(0) + if isinstance(outcome, OSError): + raise outcome + return outcome + + monkeypatch.setattr(pool_module, "pymssql", RecoveringModule()) + pool = PymssqlConnectionPool({"server": "sql.example.test"}, recycle_seconds=0, health_check_interval=0) + + try: + with pytest.raises(OSError, match="temporary connect failure"), pool.get_connection(): + pytest.fail("Failed acquisition must not enter the checkout body") + + assert pool.checked_out() == 0 + assert pool.size() == 0 + + with pytest.raises(ValueError, match="failed transaction"), pool.get_connection() as connection: + assert connection is dirty_connection + dirty_connection.autocommit(False) + raise ValueError("failed transaction") + + assert dirty_connection.closed is True + assert pool.checked_out() == 0 + assert pool.size() == 0 + + with pool.get_connection() as replacement: + assert replacement is clean_connection + assert clean_connection.closed is False + assert clean_connection.autocommit_state is True + assert pool.checked_out() == 1 + + assert pool.checked_out() == 0 + finally: + pool.close() From 04114c316a367d6dc18aace45a45276d9a9256f2 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 22:09:16 +0000 Subject: [PATCH 146/174] fix(psqlpy): stop Arrow ingestion after the first copy failure --- sqlspec/adapters/psqlpy/driver.py | 2 ++ .../unit/adapters/test_psqlpy/test_driver.py | 29 ++++++++++++------- 2 files changed, 21 insertions(+), 10 deletions(-) diff --git a/sqlspec/adapters/psqlpy/driver.py b/sqlspec/adapters/psqlpy/driver.py index 08f561ac4..b43bc3131 100644 --- a/sqlspec/adapters/psqlpy/driver.py +++ b/sqlspec/adapters/psqlpy/driver.py @@ -431,6 +431,7 @@ async def load_from_arrow( await self._run_with_exception_handler( exc_handler, cursor.binary_copy_to_table, payload, table_name, **copy_kwargs ) + self._check_pending_exception(exc_handler) else: json_columns = await self._resolve_json_columns(schema_name, table_name) deserializer = cast( @@ -448,6 +449,7 @@ async def load_from_arrow( coerce_json_columns(records, columns, json_columns, deserializer=deserializer), **copy_kwargs, ) + self._check_pending_exception(exc_handler) self._check_pending_exception(exc_handler) telemetry_payload = self._ingest_telemetry_from_batches(metrics["rows_processed"], metrics["bytes_processed"]) diff --git a/tests/unit/adapters/test_psqlpy/test_driver.py b/tests/unit/adapters/test_psqlpy/test_driver.py index bddd8683a..20b9e2b15 100644 --- a/tests/unit/adapters/test_psqlpy/test_driver.py +++ b/tests/unit/adapters/test_psqlpy/test_driver.py @@ -11,7 +11,7 @@ from sqlspec.adapters.psqlpy.driver import PsqlpyDriver from sqlspec.core import SQL from sqlspec.core.sqlcommenter import SQLCommenterContext -from sqlspec.exceptions import SQLSpecError, UniqueViolationError +from sqlspec.exceptions import UniqueViolationError @pytest.mark.anyio @@ -167,24 +167,33 @@ async def test_psqlpy_cached_query_maps_native_errors() -> None: @pytest.mark.anyio -async def test_psqlpy_load_from_arrow_preserves_cause_on_failure() -> None: - """load_from_arrow preserves underlying __cause__ on execution failure.""" +@pytest.mark.parametrize("binary_copy", [False, True], ids=["records", "binary"]) +async def test_psqlpy_load_from_arrow_stops_after_first_copy_failure(binary_copy: bool) -> None: + """A failed COPY stops the batch stream and preserves the first native cause.""" import pyarrow as pa from psqlpy.exceptions import ConnectionExecuteError - conn = AsyncMock() - conn.binary_copy_to_table = AsyncMock(side_effect=ConnectionExecuteError("binary copy network drop")) + first_error = ConnectionExecuteError("duplicate key value violates unique constraint items_pk") + later_error = ConnectionExecuteError("later copy network drop") + copy = AsyncMock(side_effect=[first_error, later_error]) + conn = SimpleNamespace( + binary_copy_to_table=copy, + copy_records_to_table=copy, + fetch=AsyncMock(return_value=SimpleNamespace(result=lambda: [])), + ) driver = PsqlpyDriver( connection=cast("Any", conn), statement_config=build_statement_config(), - driver_features={"enable_pgcopy_binary": True, "storage_capabilities": {"arrow_import_enabled": True}}, + driver_features={"enable_pgcopy_binary": binary_copy, "storage_capabilities": {"arrow_import_enabled": True}}, ) - tbl = pa.Table.from_arrays([pa.array([1, 2])], names=["id"]) + batches = [pa.record_batch({"id": [value]}) for value in (1, 2)] + reader = pa.RecordBatchReader.from_batches(batches[0].schema, batches) - with pytest.raises(SQLSpecError) as exc_info: - await driver.load_from_arrow("items", tbl) + with pytest.raises(UniqueViolationError) as exc_info: + await driver.load_from_arrow("items", reader) - assert isinstance(exc_info.value.__cause__, ConnectionExecuteError) + assert copy.await_count == 1 + assert exc_info.value.__cause__ is first_error def test_psqlpy_driver_class_cached_in_resolve() -> None: From 5f49bb291f29d7e59aeca49f48a8a8fa29dfda56 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 22:09:16 +0000 Subject: [PATCH 147/174] fix(cockroach): preserve configure callbacks in derived pools --- sqlspec/adapters/cockroach_psycopg/config.py | 8 +- .../test_cockroach_psycopg/test_config.py | 99 +++++++++++++++++-- 2 files changed, 97 insertions(+), 10 deletions(-) diff --git a/sqlspec/adapters/cockroach_psycopg/config.py b/sqlspec/adapters/cockroach_psycopg/config.py index 26296caac..5161946b7 100644 --- a/sqlspec/adapters/cockroach_psycopg/config.py +++ b/sqlspec/adapters/cockroach_psycopg/config.py @@ -280,7 +280,7 @@ def __init__( self._user_connection_hook: Callable[[CockroachSyncConnection], None] | None = features_dict.pop( "on_connection_create", None ) - self._user_configure: Callable[[CockroachSyncConnection], Any] | None = connection_config.pop("configure", None) + self._user_configure: Callable[[CockroachSyncConnection], Any] | None = connection_config.get("configure") super().__init__( connection_config=connection_config, @@ -363,6 +363,7 @@ def get_event_runtime_hints(self) -> "EventRuntimeHints": def _create_pool(self) -> "ConnectionPool": all_config = dict(self.connection_config) + all_config.pop("configure", None) pool_parameters = { "min_size": all_config.pop("min_size", 4), @@ -522,9 +523,7 @@ def __init__( self._user_connection_hook: Callable[[CockroachAsyncConnection], Awaitable[None]] | None = features_dict.pop( "on_connection_create", None ) - self._user_configure: Callable[[CockroachAsyncConnection], Any] | None = connection_config.pop( - "configure", None - ) + self._user_configure: Callable[[CockroachAsyncConnection], Any] | None = connection_config.get("configure") super().__init__( connection_config=connection_config, @@ -607,6 +606,7 @@ def get_event_runtime_hints(self) -> "EventRuntimeHints": async def _create_pool(self) -> "AsyncConnectionPool": all_config = dict(self.connection_config) + all_config.pop("configure", None) pool_parameters = { "min_size": all_config.pop("min_size", 4), diff --git a/tests/unit/adapters/test_cockroach_psycopg/test_config.py b/tests/unit/adapters/test_cockroach_psycopg/test_config.py index c3f621ef7..3d7c9c825 100644 --- a/tests/unit/adapters/test_cockroach_psycopg/test_config.py +++ b/tests/unit/adapters/test_cockroach_psycopg/test_config.py @@ -941,6 +941,90 @@ async def async_user_cb(conn: Any) -> None: assert call_order == ["internal_autocommit", "user_cb"] +@pytest.mark.parametrize("async_mode", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize( + ("override_kind", "expected"), + [ + ("inherit", ["original"]), + ("mapping", ["replacement"]), + ("mapping_none", []), + ("keyword", ["replacement"]), + ("keyword_none", []), + ("keyword_precedence", []), + ], +) +async def test_cockroach_psycopg_derived_connection_preserves_configure( + async_mode: bool, override_kind: str, expected: list[str] +) -> None: + """Derived connections inherit configure callbacks and honor explicit replacements.""" + calls: list[str] = [] + + def original_callback(_connection: Any) -> None: + calls.append("original") + + def replacement_callback(_connection: Any) -> None: + calls.append("replacement") + + async def async_original_callback(connection: Any) -> None: + original_callback(connection) + + async def async_replacement_callback(connection: Any) -> None: + replacement_callback(connection) + + original = async_original_callback if async_mode else original_callback + replacement = async_replacement_callback if async_mode else replacement_callback + source_config: dict[str, Any] = {"autocommit": True, "configure": original} + config = ( + CockroachPsycopgAsyncConfig(connection_config=source_config) + if async_mode + else CockroachPsycopgSyncConfig(connection_config=source_config) + ) + overrides: dict[str, Any] = {} + if override_kind in {"mapping", "keyword_precedence"}: + overrides["connection_config_overrides"] = {"configure": replacement} + elif override_kind == "mapping_none": + overrides["connection_config_overrides"] = {"configure": None} + if override_kind == "keyword": + overrides["configure"] = replacement + elif override_kind in {"keyword_none", "keyword_precedence"}: + overrides["configure"] = None + derived = config.derive_pool(**overrides) + fake_connection = MagicMock() + fake_connection.set_autocommit = AsyncMock() + sdk_connection = "AsyncCrdbConnection" if async_mode else "CrdbConnection" + connect = AsyncMock(return_value=fake_connection) if async_mode else MagicMock(return_value=fake_connection) + + with patch(f"sqlspec.adapters.cockroach_psycopg.config.psycopg_crdb.{sdk_connection}.connect", connect): + if isinstance(derived, CockroachPsycopgAsyncConfig): + await derived.create_connection() + else: + derived.create_connection() + assert calls == expected + calls.clear() + with ( + patch("sqlspec.adapters.cockroach_psycopg.config.ConnectionPool", _SyncPoolSpy), + patch("sqlspec.adapters.cockroach_psycopg.config.AsyncConnectionPool", _AsyncPoolSpy), + ): + if isinstance(derived, CockroachPsycopgAsyncConfig): + await derived.create_pool() + pool_kwargs = _AsyncPoolSpy.calls[-1][1] + await pool_kwargs["configure"](fake_connection) + else: + derived.create_pool() + pool_kwargs = _SyncPoolSpy.calls[-1][1] + pool_kwargs["configure"](fake_connection) + assert pool_kwargs["kwargs"] == {"autocommit": True} + assert calls == expected + calls.clear() + if isinstance(config, CockroachPsycopgAsyncConfig): + await config.create_connection() + else: + config.create_connection() + + assert calls == ["original"] + assert source_config["configure"] is original + + def test_cockroach_psycopg_adk_hash_shard_bucket_count_validation() -> None: """Invalid hash_shard_bucket_count must raise ValueError in cockroach_psycopg adk store.""" from sqlspec.adapters.cockroach_psycopg.adk.store import CockroachPsycopgADKConfig, _cockroach_hash_shard_clause @@ -951,8 +1035,8 @@ def test_cockroach_psycopg_adk_hash_shard_bucket_count_validation() -> None: async def test_cockroach_psycopg_adk_memory_search_with_embedding() -> None: - """search_entries with embedding executes vector cosine search in cockroach_psycopg.""" - executed_queries: list[str] = [] + """An embedding argument preserves text search against the canonical memory schema.""" + executed_queries: list[tuple[str, Any]] = [] class _FakeCursor: description = [ @@ -968,11 +1052,10 @@ class _FakeCursor: ("content_text",), ("metadata_json",), ("inserted_at",), - ("embedding",), ] async def execute(self, sql: bytes, params: Any) -> None: - executed_queries.append(sql.decode()) + executed_queries.append((sql.decode(), params)) async def fetchall(self) -> list[Any]: return [] @@ -998,9 +1081,13 @@ async def __aexit__(self, *_a: Any) -> None: object.__setattr__(config, "provide_connection", lambda: _FakeCtx()) store = CockroachPsycopgAsyncADKMemoryStore(config) - await store.search_entries(query="hello", app_name="app", user_id="u1", embedding=[0.1, 0.2, 0.3]) + result = await store.search_entries(query="hello", app_name="app", user_id="u1", limit=7, embedding=[0.1, 0.2, 0.3]) + assert result == [] assert len(executed_queries) == 1 - assert "embedding <=>" in executed_queries[0] + sql, params = executed_queries[0] + assert "content_text ILIKE %s" in sql + assert "ORDER BY timestamp DESC" in sql + assert params == ("app", "u1", "%hello%", 7) async def test_cockroach_psycopg_adk_memory_insert_batch_unnest() -> None: From ba3148ad6d3b07383b130df1f9b17caf97ea9c78 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 22:11:04 +0000 Subject: [PATCH 148/174] fix(adbc): preserve memory event deduplication across batch conflicts --- sqlspec/adapters/adbc/adk/store.py | 30 +++++- .../test_adbc/test_adk_dialect_support.py | 101 ++++++++++++++++++ 2 files changed, 127 insertions(+), 4 deletions(-) diff --git a/sqlspec/adapters/adbc/adk/store.py b/sqlspec/adapters/adbc/adk/store.py index 7b763f57f..2f19b9702 100644 --- a/sqlspec/adapters/adbc/adk/store.py +++ b/sqlspec/adapters/adbc/adk/store.py @@ -2206,9 +2206,12 @@ def _insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "objec self._execute(cursor, select_sql, tuple(event_ids)) existing_rows = cursor.fetchall() existing_event_ids = {row[0] for row in existing_rows} - entries_to_insert = [ - entry for entry in entries if entry.get("event_id") not in existing_event_ids - ] + entries_to_insert = [] + for entry in entries: + event_id = entry["event_id"] + if event_id not in existing_event_ids: + entries_to_insert.append(entry) + existing_event_ids.add(event_id) for entry in entries_to_insert: content_json = self._serialize_json_field(entry["content_json"]) @@ -2286,7 +2289,26 @@ def _insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "objec if cursor.fetchone(): inserted_count += 1 else: - self._execute(cursor, sql, params) + try: + self._execute(cursor, sql, params) + except Exception as exc: + error_message = str(exc).lower() + if not ( + getattr(exc, "sqlstate", None) == "23505" + or "unique" in error_message + or "duplicate" in error_message + or "ora-00001" in error_message + or "sql0803n" in error_message + ): + raise + self._execute( + cursor, + f"SELECT event_id FROM {self._memory_table} WHERE event_id = ?", + (entry["event_id"],), + ) + if cursor.fetchone() is None: + raise + continue inserted_count += 1 conn.commit() except Exception: diff --git a/tests/unit/adapters/test_adbc/test_adk_dialect_support.py b/tests/unit/adapters/test_adbc/test_adk_dialect_support.py index 34d60bd2f..bc1283783 100644 --- a/tests/unit/adapters/test_adbc/test_adk_dialect_support.py +++ b/tests/unit/adapters/test_adbc/test_adk_dialect_support.py @@ -6,6 +6,7 @@ from unittest.mock import MagicMock import pytest +from adbc_driver_manager import AdbcStatusCode, IntegrityError from sqlspec.adapters.adbc import AdbcConfig from sqlspec.adapters.adbc.adk import AdbcADKConfig, AdbcADKMemoryStore, AdbcADKStore @@ -19,6 +20,7 @@ from sqlspec.adapters.adbc.litestar.store import AdbcStore from sqlspec.core import StatementConfig from sqlspec.exceptions import ImproperConfigurationError +from sqlspec.extensions.adk import StoredMemory pytestmark = [pytest.mark.xdist_group("sqlite"), pytest.mark.adbc, pytest.mark.integration] @@ -821,3 +823,102 @@ def record_execute(sql: str, _params: Any = None) -> None: assert any("SELECT event_id FROM" in s for s in executed_statements) assert count == 1 + + +@pytest.mark.parametrize("driver_name", ["oracle", "db2"]) +def test_adbc_memory_batch_deduplicates_fresh_event_ids(driver_name: str, monkeypatch: pytest.MonkeyPatch) -> None: + """Repeated fresh event IDs keep the first entry and preserve the rest of the batch.""" + cursor = MagicMock() + cursor.fetchall.return_value = [] + conn = MagicMock() + conn.cursor.return_value = cursor + context = MagicMock() + context.__enter__.return_value = conn + monkeypatch.setattr(AdbcConfig, "provide_connection", lambda *_a, **_k: context) + store = AdbcADKMemoryStore(AdbcConfig(connection_config={"driver_name": driver_name})) + now = datetime.now(UTC) + entry: StoredMemory = { + "id": "mem-1", + "session_id": "sess-1", + "app_name": "app", + "user_id": "u1", + "scope": "user", + "event_id": "evt-1", + "author": "user", + "timestamp": now, + "content_json": {"text": "hello"}, + "content_text": "hello", + "metadata_json": None, + "inserted_at": now, + "embedding": None, + } + + count = store.insert_memory_entries([ + entry, + {**entry, "id": "mem-duplicate"}, + {**entry, "id": "mem-2", "event_id": "evt-2"}, + ]) + + inserts = [call.args[1] for call in cursor.execute.call_args_list if "INSERT INTO" in call.args[0]] + assert count == 2 + assert [params[0] for params in inserts] == ["mem-1", "mem-2"] + conn.commit.assert_called_once() + + +@pytest.mark.parametrize("driver_name", ["oracle", "db2"]) +@pytest.mark.parametrize( + ("error_message", "event_exists"), + [ + ("ORA-00001: unique constraint violated", True), + ("SQL0803N duplicate key", True), + ("ORA-00001: unique constraint violated", False), + ("foreign key constraint violated", True), + ], +) +def test_adbc_memory_insert_only_ignores_confirmed_event_conflicts( + driver_name: str, error_message: str, event_exists: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + """A duplicate race is idempotent while other integrity errors still roll back.""" + cursor = MagicMock() + cursor.fetchall.return_value = [] + cursor.fetchone.return_value = ("evt-1",) if event_exists else None + error = IntegrityError(error_message, status_code=AdbcStatusCode.INTEGRITY) + + def execute(sql: str, _params: object = None) -> None: + if "INSERT INTO" in sql: + raise error + + cursor.execute.side_effect = execute + conn = MagicMock() + conn.cursor.return_value = cursor + context = MagicMock() + context.__enter__.return_value = conn + monkeypatch.setattr(AdbcConfig, "provide_connection", lambda *_a, **_k: context) + store = AdbcADKMemoryStore(AdbcConfig(connection_config={"driver_name": driver_name})) + now = datetime.now(UTC) + entry: StoredMemory = { + "id": "mem-1", + "session_id": "sess-1", + "app_name": "app", + "user_id": "u1", + "scope": "user", + "event_id": "evt-1", + "author": "user", + "timestamp": now, + "content_json": {"text": "hello"}, + "content_text": "hello", + "metadata_json": None, + "inserted_at": now, + "embedding": None, + } + + if event_exists and "foreign key" not in error_message: + assert store.insert_memory_entries([entry]) == 0 + conn.commit.assert_called_once() + conflict_checks = [call.args[1] for call in cursor.execute.call_args_list if "WHERE event_id =" in call.args[0]] + assert conflict_checks == [("evt-1",)] + else: + with pytest.raises(IntegrityError) as exc_info: + store.insert_memory_entries([entry]) + assert exc_info.value is error + conn.rollback.assert_called_once() From 7023c68f9d2c8f024c9ba0cebe63de786d1716ed Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 22:11:04 +0000 Subject: [PATCH 149/174] fix(psqlpy): preserve native owner types in memory batches --- sqlspec/adapters/psqlpy/adk/store.py | 11 +++++----- .../adapters/test_psqlpy/test_adk_store.py | 20 +++++++++++++++++++ 2 files changed, 25 insertions(+), 6 deletions(-) diff --git a/sqlspec/adapters/psqlpy/adk/store.py b/sqlspec/adapters/psqlpy/adk/store.py index cedcc9f35..6ca32d356 100644 --- a/sqlspec/adapters/psqlpy/adk/store.py +++ b/sqlspec/adapters/psqlpy/adk/store.py @@ -721,17 +721,16 @@ async def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: " content_text, metadata_json, inserted_at, embedding ) SELECT id, session_id, app_name, user_id, scope, event_id, author, - owner_id, timestamp, content_json, content_text, metadata_json, inserted_at, embedding::vector + $14, timestamp, content_json, content_text, metadata_json, inserted_at, embedding::vector FROM unnest( $1::varchar[], $2::varchar[], $3::varchar[], $4::varchar[], - $5::varchar[], $6::varchar[], $7::varchar[], $8::varchar[], - $9::timestamptz[], $10::jsonb[], $11::text[], $12::jsonb[], $13::timestamptz[], $14::text[] + $5::varchar[], $6::varchar[], $7::varchar[], + $8::timestamptz[], $9::jsonb[], $10::text[], $11::jsonb[], $12::timestamptz[], $13::text[] ) AS batch(id, session_id, app_name, user_id, scope, event_id, author, - owner_id, timestamp, content_json, content_text, metadata_json, inserted_at, embedding) + timestamp, content_json, content_text, metadata_json, inserted_at, embedding) ON CONFLICT (event_id) DO NOTHING RETURNING 1 """ - owner_ids = [owner_id for _ in entries] params = [ ids, session_ids, @@ -740,13 +739,13 @@ async def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: " scopes, event_ids, authors, - owner_ids, timestamps, content_jsons, content_texts, metadata_jsons, inserted_ats, embeddings, + owner_id, ] else: sql = f""" diff --git a/tests/unit/adapters/test_psqlpy/test_adk_store.py b/tests/unit/adapters/test_psqlpy/test_adk_store.py index d5e757629..8a1e56636 100644 --- a/tests/unit/adapters/test_psqlpy/test_adk_store.py +++ b/tests/unit/adapters/test_psqlpy/test_adk_store.py @@ -3,10 +3,12 @@ from datetime import UTC, datetime from typing import Any, NotRequired, Self, cast, get_args, get_origin from unittest.mock import MagicMock +from uuid import UUID import pytest from psqlpy.exceptions import ConnectionExecuteError, DatabaseError +from sqlspec.adapters.psqlpy import PsqlpyConfig from sqlspec.adapters.psqlpy.adk import PsqlpyADKConfig, PsqlpyADKMemoryStore, PsqlpyADKStore from sqlspec.adapters.psqlpy.adk.store import _is_psqlpy_database_error, _is_table_missing_error, decode_postgres_vector from sqlspec.config import ADKConfig @@ -300,3 +302,21 @@ async def test_psqlpy_memory_store_normalizes_tuple_embeddings_on_insert_and_sea assert conn.fetch_calls[0][1][-1] == ["[0.1,0.2]"] assert conn.fetch_calls[1][1][0] == [0.1, 0.2] assert results[0]["embedding"] == [0.1, 0.2] + + +@pytest.mark.parametrize(("owner_type", "owner_id"), [("INTEGER", 42), ("UUID", UUID(int=1))]) +async def test_psqlpy_memory_batch_preserves_owner_type( + owner_type: str, owner_id: object, monkeypatch: pytest.MonkeyPatch +) -> None: + """A scalar owner bind lets the destination column infer its native type.""" + conn = _FakeMemoryConnection([_sample_memory_row(None)]) + monkeypatch.setattr(PsqlpyConfig, "provide_connection", lambda *_a, **_k: conn) + config = PsqlpyConfig(extension_config={"adk": {"owner_id_column": f"owner_id {owner_type}"}}) + store = PsqlpyADKMemoryStore(config) + + assert await store.insert_memory_entries([cast("Any", _sample_memory_row(None))], owner_id) == 1 + + sql, params = conn.fetch_calls[0] + assert "author, $14, timestamp" in _normalized(sql) + assert params[-1] == owner_id + assert len(params[0]) == 1 From 76cf1f3b9b625cd2695d183cd63b91b1835807ad Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 22:11:04 +0000 Subject: [PATCH 150/174] fix(cockroach): preserve memory schema and owner contracts --- .../adapters/cockroach_asyncpg/adk/store.py | 28 +++--- .../adapters/cockroach_psycopg/adk/store.py | 48 ++++------ .../test_cockroach_asyncpg/test_adk_store.py | 56 ++++++++++-- .../test_cockroach_psycopg/test_adk_store.py | 88 ++++++++++++++++++- 4 files changed, 163 insertions(+), 57 deletions(-) diff --git a/sqlspec/adapters/cockroach_asyncpg/adk/store.py b/sqlspec/adapters/cockroach_asyncpg/adk/store.py index a3f591157..3be8d81da 100644 --- a/sqlspec/adapters/cockroach_asyncpg/adk/store.py +++ b/sqlspec/adapters/cockroach_asyncpg/adk/store.py @@ -632,14 +632,16 @@ async def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: " (id, session_id, app_name, user_id, scope, event_id, author, {self._owner_id_column_name}, timestamp, content_json, content_text, metadata_json, inserted_at) - SELECT * FROM unnest( + SELECT id, session_id, app_name, user_id, scope, event_id, author, + $13, timestamp, content_json, content_text, metadata_json, inserted_at + FROM unnest( $1::varchar[], $2::varchar[], $3::varchar[], $4::varchar[], - $5::varchar[], $6::varchar[], $7::varchar[], $8::varchar[], - $9::timestamptz[], $10::jsonb[], $11::text[], $12::jsonb[], $13::timestamptz[] - ) + $5::varchar[], $6::varchar[], $7::varchar[], + $8::timestamptz[], $9::jsonb[], $10::text[], $11::jsonb[], $12::timestamptz[] + ) AS batch(id, session_id, app_name, user_id, scope, event_id, author, + timestamp, content_json, content_text, metadata_json, inserted_at) ON CONFLICT (event_id) DO NOTHING """ - owner_ids = [owner_id for _ in entries] params: tuple[Any, ...] = ( ids, session_ids, @@ -648,12 +650,12 @@ async def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: " scopes, event_ids, authors, - owner_ids, timestamps, content_jsons, content_texts, metadata_jsons, inserted_ats, + owner_id, ) else: sql = f""" @@ -700,7 +702,7 @@ async def search_entries( msg = "Memory store is disabled" raise RuntimeError(msg) - if not query and embedding is None: + if not query: return [] effective_limit = limit if limit is not None else self._max_results @@ -717,17 +719,7 @@ async def search_entries( scope_count = len(scope_params) - if embedding is not None: - p_vec = f"${scope_count + 1}" - p_lim = f"${scope_count + 2}" - sql = f""" - SELECT * FROM {self._memory_table} - WHERE {where_scope} AND embedding IS NOT NULL - ORDER BY embedding <=> {p_vec}::float8[]::vector - LIMIT {p_lim} - """ - params = (*scope_params, list(embedding), effective_limit) - elif self._use_fts: + if self._use_fts: p_pat = f"${scope_count + 1}" p_lim = f"${scope_count + 2}" sql = f""" diff --git a/sqlspec/adapters/cockroach_psycopg/adk/store.py b/sqlspec/adapters/cockroach_psycopg/adk/store.py index 7ea16f120..4dd797397 100644 --- a/sqlspec/adapters/cockroach_psycopg/adk/store.py +++ b/sqlspec/adapters/cockroach_psycopg/adk/store.py @@ -1263,15 +1263,18 @@ async def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: " {owner_id_col}, timestamp, content_json, content_text, metadata_json, inserted_at ) - SELECT * FROM unnest( - %s::varchar[], %s::varchar[], %s::varchar[], %s::varchar[], + SELECT id, session_id, app_name, user_id, scope, event_id, author, + %s, timestamp, content_json, content_text, metadata_json, inserted_at + FROM unnest( %s::varchar[], %s::varchar[], %s::varchar[], %s::varchar[], + %s::varchar[], %s::varchar[], %s::varchar[], %s::timestamptz[], %s::jsonb[], %s::text[], %s::jsonb[], %s::timestamptz[] - ) + ) AS batch(id, session_id, app_name, user_id, scope, event_id, author, + timestamp, content_json, content_text, metadata_json, inserted_at) ON CONFLICT (event_id) DO NOTHING """).format(table=table_ident, owner_id_col=pg_sql.Identifier(self._owner_id_column_name)) - owner_ids = [owner_id for _ in entries] params: tuple[Any, ...] = ( + owner_id, ids, session_ids, app_names, @@ -1279,7 +1282,6 @@ async def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: " scopes, event_ids, authors, - owner_ids, timestamps, content_jsons, content_texts, @@ -1334,21 +1336,13 @@ async def search_entries( msg = "Memory store is disabled" raise RuntimeError(msg) - if not query and embedding is None: + if not query: return [] effective_limit = limit if limit is not None else self._max_results where_scope, scope_params = _build_cockroach_scope_where(app_name, user_id, scope_filter) - if embedding is not None: - sql = f""" - SELECT * FROM {self._memory_table} - WHERE {where_scope} AND embedding IS NOT NULL - ORDER BY embedding <=> %s::float8[]::vector - LIMIT %s - """ - params = (*scope_params, list(embedding), effective_limit) - elif self._use_fts: + if self._use_fts: sql = f""" SELECT * FROM {self._memory_table} WHERE {where_scope} @@ -1504,15 +1498,18 @@ def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "object {owner_id_col}, timestamp, content_json, content_text, metadata_json, inserted_at ) - SELECT * FROM unnest( - %s::varchar[], %s::varchar[], %s::varchar[], %s::varchar[], + SELECT id, session_id, app_name, user_id, scope, event_id, author, + %s, timestamp, content_json, content_text, metadata_json, inserted_at + FROM unnest( %s::varchar[], %s::varchar[], %s::varchar[], %s::varchar[], + %s::varchar[], %s::varchar[], %s::varchar[], %s::timestamptz[], %s::jsonb[], %s::text[], %s::jsonb[], %s::timestamptz[] - ) + ) AS batch(id, session_id, app_name, user_id, scope, event_id, author, + timestamp, content_json, content_text, metadata_json, inserted_at) ON CONFLICT (event_id) DO NOTHING """).format(table=table_ident, owner_id_col=pg_sql.Identifier(self._owner_id_column_name)) - owner_ids = [owner_id for _ in entries] params: tuple[Any, ...] = ( + owner_id, ids, session_ids, app_names, @@ -1520,7 +1517,6 @@ def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "object scopes, event_ids, authors, - owner_ids, timestamps, content_jsons, content_texts, @@ -1576,21 +1572,13 @@ def search_entries( msg = "Memory store is disabled" raise RuntimeError(msg) - if not query and embedding is None: + if not query: return [] effective_limit = limit if limit is not None else self._max_results where_scope, scope_params = _build_cockroach_scope_where(app_name, user_id, scope_filter) - if embedding is not None: - sql = f""" - SELECT * FROM {self._memory_table} - WHERE {where_scope} AND embedding IS NOT NULL - ORDER BY embedding <=> %s::float8[]::vector - LIMIT %s - """ - params = (*scope_params, list(embedding), effective_limit) - elif self._use_fts: + if self._use_fts: sql = f""" SELECT * FROM {self._memory_table} WHERE {where_scope} diff --git a/tests/unit/adapters/test_cockroach_asyncpg/test_adk_store.py b/tests/unit/adapters/test_cockroach_asyncpg/test_adk_store.py index ed36f9d44..026708cc1 100644 --- a/tests/unit/adapters/test_cockroach_asyncpg/test_adk_store.py +++ b/tests/unit/adapters/test_cockroach_asyncpg/test_adk_store.py @@ -1,8 +1,10 @@ # pyright: reportPrivateUsage=false """Unit tests for CockroachDB asyncpg ADK store extension configuration.""" +from datetime import UTC, datetime from typing import Any, NotRequired, Self, cast, get_args, get_origin -from unittest.mock import MagicMock +from unittest.mock import AsyncMock, MagicMock +from uuid import UUID import pytest @@ -209,8 +211,9 @@ def test_cockroach_asyncpg_adk_hash_shard_bucket_count_validation() -> None: _cockroach_hash_shard_clause(bad_config) -async def test_cockroach_asyncpg_adk_memory_search_with_embedding() -> None: - """search_entries with embedding executes vector cosine search.""" +@pytest.mark.parametrize("use_fts", [False, True]) +async def test_cockroach_asyncpg_adk_memory_search_with_embedding(use_fts: bool) -> None: + """Embedding arguments retain text search against the canonical memory schema.""" executed_sql: list[tuple[str, tuple[Any, ...]]] = [] class _FakeConn: @@ -225,14 +228,51 @@ async def __aenter__(self) -> _FakeConn: async def __aexit__(self, *_args: Any) -> None: return None - config = CockroachAsyncpgConfig(extension_config={"adk": {"enable_memory": True}}) + config = CockroachAsyncpgConfig(extension_config={"adk": {"enable_memory": True, "memory_use_fts": use_fts}}) object.__setattr__(config, "provide_connection", MagicMock(return_value=_FakeCtx())) store = CockroachAsyncpgADKMemoryStore(config) await store.search_entries(query="hello", app_name="app", user_id="u1", embedding=[0.1, 0.2, 0.3]) assert len(executed_sql) == 1 sql, params = executed_sql[0] - assert "embedding <=>" in sql - assert params[0] == "app" - assert params[1] == "u1" - assert params[2] == [0.1, 0.2, 0.3] + assert "plainto_tsquery" in sql if use_fts else "content_text ILIKE" in sql + assert params == ("app", "u1", "hello" if use_fts else "%hello%", 20) + + assert await store.search_entries("", "app", "u1", embedding=[0.1]) == [] + assert len(executed_sql) == 1 + + +@pytest.mark.parametrize(("owner_type", "owner_id"), [("INTEGER", 42), ("UUID", UUID(int=1))]) +async def test_cockroach_asyncpg_memory_batch_preserves_owner_type( + owner_type: str, owner_id: object, monkeypatch: pytest.MonkeyPatch +) -> None: + """A scalar owner bind retains integer and UUID ownership contracts.""" + conn = MagicMock() + conn.execute = AsyncMock(return_value="INSERT 0 1") + context = MagicMock() + context.__aenter__.return_value = conn + monkeypatch.setattr(CockroachAsyncpgConfig, "provide_connection", lambda *_a, **_k: context) + config = CockroachAsyncpgConfig(extension_config={"adk": {"owner_id_column": f"owner_id {owner_type}"}}) + now = datetime.now(UTC) + entry = { + "id": "mem-1", + "session_id": "sess-1", + "app_name": "app", + "user_id": "u1", + "scope": "user", + "event_id": "evt-1", + "author": "user", + "timestamp": now, + "content_json": {"text": "hello"}, + "content_text": "hello", + "metadata_json": None, + "inserted_at": now, + "embedding": None, + } + + assert await CockroachAsyncpgADKMemoryStore(config).insert_memory_entries([cast("Any", entry)], owner_id) == 1 + + sql, *params = conn.execute.call_args.args + assert "author, $13, timestamp" in _normalized(sql) + assert params[-1] == owner_id + assert params[0] == ["mem-1"] diff --git a/tests/unit/adapters/test_cockroach_psycopg/test_adk_store.py b/tests/unit/adapters/test_cockroach_psycopg/test_adk_store.py index 644270e86..6d3e31262 100644 --- a/tests/unit/adapters/test_cockroach_psycopg/test_adk_store.py +++ b/tests/unit/adapters/test_cockroach_psycopg/test_adk_store.py @@ -1,11 +1,14 @@ # pyright: reportPrivateUsage=false """Unit tests for CockroachDB psycopg ADK store extension configuration.""" +from datetime import UTC, datetime from typing import Any, NotRequired, Self, cast, get_args, get_origin -from unittest.mock import MagicMock +from unittest.mock import AsyncMock, MagicMock +from uuid import UUID import pytest +from sqlspec.adapters.cockroach_psycopg import CockroachPsycopgAsyncConfig, CockroachPsycopgSyncConfig from sqlspec.adapters.cockroach_psycopg.adk import ( CockroachPsycopgADKConfig, CockroachPsycopgAsyncADKMemoryStore, @@ -319,3 +322,86 @@ def test_cockroach_psycopg_sync_list_sessions_rejects_invalid_options(options: " store.list_sessions("app", **options) assert cursor.calls == [] + + +def _memory_store_with_cursor(async_mode: bool, adk_config: dict[str, Any], monkeypatch: pytest.MonkeyPatch) -> Any: + """Provide real stores with SDK cursor and connection boundaries substituted.""" + config_type = CockroachPsycopgAsyncConfig if async_mode else CockroachPsycopgSyncConfig + cursor = MagicMock() + cursor.rowcount = 1 + cursor.description = [] + cursor.fetchall.return_value = [] + conn = MagicMock() + context = MagicMock() + context.__enter__.return_value = conn + context.__aenter__.return_value = conn + cursor.__enter__.return_value = cursor + cursor.__aenter__.return_value = cursor + conn.cursor.return_value = cursor + if async_mode: + cursor.execute = AsyncMock() + cursor.fetchall = AsyncMock(return_value=[]) + conn.commit = AsyncMock() + monkeypatch.setattr(config_type, "provide_connection", lambda *_a, **_k: context) + if async_mode: + return CockroachPsycopgAsyncADKMemoryStore( + CockroachPsycopgAsyncConfig(extension_config={"adk": adk_config}) + ), cursor + return CockroachPsycopgSyncADKMemoryStore(CockroachPsycopgSyncConfig(extension_config={"adk": adk_config})), cursor + + +@pytest.mark.parametrize("async_mode", [False, True]) +@pytest.mark.parametrize(("owner_type", "owner_id"), [("INTEGER", 42), ("UUID", UUID(int=1))]) +async def test_cockroach_psycopg_memory_batch_preserves_owner_type( + async_mode: bool, owner_type: str, owner_id: object, monkeypatch: pytest.MonkeyPatch +) -> None: + """Sync and async batches bind one native owner scalar outside UNNEST.""" + store, cursor = _memory_store_with_cursor(async_mode, {"owner_id_column": f"owner_id {owner_type}"}, monkeypatch) + now = datetime.now(UTC) + entry = { + "id": "mem-1", + "session_id": "sess-1", + "app_name": "app", + "user_id": "u1", + "scope": "user", + "event_id": "evt-1", + "author": "user", + "timestamp": now, + "content_json": {"text": "hello"}, + "content_text": "hello", + "metadata_json": None, + "inserted_at": now, + "embedding": None, + } + inserted = store.insert_memory_entries([entry], owner_id) + if async_mode: + inserted = await inserted + assert inserted == 1 + + query, params = cursor.execute.call_args.args + assert "author, %s, timestamp" in _normalized(query.as_string()) + assert params[0] == owner_id + assert params[1] == ["mem-1"] + + +@pytest.mark.parametrize("async_mode", [False, True]) +@pytest.mark.parametrize("use_fts", [False, True]) +async def test_cockroach_psycopg_memory_embedding_keeps_text_search( + async_mode: bool, use_fts: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + """Embedding arguments keep searching text columns present in the memory schema.""" + store, cursor = _memory_store_with_cursor(async_mode, {"memory_use_fts": use_fts}, monkeypatch) + result = store.search_entries("hello", "app", "u1", embedding=[0.1]) + if async_mode: + result = await result + assert result == [] + query, params = cursor.execute.call_args.args + sql = query.decode() + assert "plainto_tsquery" in sql if use_fts else "content_text ILIKE" in sql + assert params == ("app", "u1", "hello" if use_fts else "%hello%", 20) + + empty = store.search_entries("", "app", "u1", embedding=[0.1]) + if async_mode: + empty = await empty + assert empty == [] + assert cursor.execute.call_count == 1 From 9fb74520929d2e10c8cd6a561585cb28bdd92822 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 22:21:30 +0000 Subject: [PATCH 151/174] test(adapters): remove source scanners and deletion-only assertions --- .../adapters/test_adapter_implementations.py | 143 ------------------ tests/unit/adapters/test_duckdb/test_pool.py | 9 -- .../test_mssql_python/test_adk_store.py | 9 -- .../test_spanner/test_adk_artifact_store.py | 15 -- 4 files changed, 176 deletions(-) diff --git a/tests/unit/adapters/test_adapter_implementations.py b/tests/unit/adapters/test_adapter_implementations.py index b2d30336d..d0f2e68e7 100644 --- a/tests/unit/adapters/test_adapter_implementations.py +++ b/tests/unit/adapters/test_adapter_implementations.py @@ -1,10 +1,8 @@ """Parameterized tests for real adapter implementations.""" -import ast import importlib import inspect import operator -import re import sqlite3 import types from pathlib import Path @@ -85,44 +83,10 @@ ) CREATE_MAPPED_EXCEPTION_MODULES = ADAPTER_CORE_MODULES APPLY_DRIVER_FEATURES_MODULES = ADAPTER_CORE_MODULES -ADAPTER_DRIVER_MODULES_WITH_FEATURE_HELPERS = ( - "sqlspec/adapters/cockroach_psycopg/driver.py", - "sqlspec/adapters/duckdb/driver.py", -) -MYSQL_FAMILY_UNCOMPILED_MODULES = ( - "sqlspec/adapters/aiomysql/config.py", - "sqlspec/adapters/aiomysql/driver.py", - "sqlspec/adapters/asyncmy/config.py", - "sqlspec/adapters/asyncmy/driver.py", - "sqlspec/adapters/mysqlconnector/config.py", - "sqlspec/adapters/mysqlconnector/driver.py", - "sqlspec/adapters/pymysql/driver.py", - "sqlspec/adapters/pymysql/pool.py", -) -MYSQL_VENDOR_IMPORT_ROOTS = ("aiomysql", "asyncmy", "mysql", "pymysql") -SPANNER_MODULE_PROXY_PATHS = ("sqlspec/adapters/spanner/config.py", "sqlspec/adapters/spanner/driver.py") SQLITE_EXCEPTION_HANDLER_CLASSES = ( ("sqlspec.adapters.aiosqlite.driver", "AiosqliteExceptionHandler"), ("sqlspec.adapters.sqlite.driver", "SqliteExceptionHandler"), ) -BRIDGE_CURSOR_CLOSE_METHODS = ( - ("sqlspec/adapters/adbc/_typing.py", "AdbcCursor", "__exit__"), - ("sqlspec/adapters/aiomysql/_typing.py", "AiomysqlCursor", "__aexit__"), - ("sqlspec/adapters/aiosqlite/_typing.py", "AiosqliteCursor", "__aexit__"), - ("sqlspec/adapters/asyncmy/_typing.py", "AsyncmyCursor", "__aexit__"), - ("sqlspec/adapters/db2/_typing.py", "Db2SyncCursor", "__exit__"), - ("sqlspec/adapters/db2/_typing.py", "Db2AsyncCursor", "__aexit__"), - ("sqlspec/adapters/mssql_python/_typing.py", "MssqlPythonCursor", "__exit__"), - ("sqlspec/adapters/mysqlconnector/_typing.py", "MysqlConnectorSyncCursor", "__exit__"), - ("sqlspec/adapters/mysqlconnector/_typing.py", "MysqlConnectorAsyncCursor", "__aexit__"), - ("sqlspec/adapters/oracledb/_typing.py", "OracleSyncCursor", "__exit__"), - ("sqlspec/adapters/oracledb/_typing.py", "OracleAsyncCursor", "__aexit__"), - ("sqlspec/adapters/psycopg/_typing.py", "PsycopgSyncCursor", "__exit__"), - ("sqlspec/adapters/psycopg/_typing.py", "PsycopgAsyncCursor", "__aexit__"), - ("sqlspec/adapters/pymssql/_typing.py", "PymssqlCursor", "__exit__"), - ("sqlspec/adapters/pymysql/_typing.py", "PyMysqlCursor", "__exit__"), - ("sqlspec/adapters/sqlite/_typing.py", "SqliteCursor", "__exit__"), -) @pytest.fixture(params=ADAPTER_CONFIGS, ids=operator.itemgetter("name")) @@ -207,95 +171,12 @@ def test_adapter_apply_driver_features_signature_and_return_shape(module_name: s assert features["custom"] == "value" -@pytest.mark.parametrize("module_path", ADAPTER_DRIVER_MODULES_WITH_FEATURE_HELPERS) -def test_adapter_drivers_do_not_apply_driver_features(module_path: str) -> None: - tree = ast.parse(Path(module_path).read_text()) - - imports_helper = any( - isinstance(node, ast.ImportFrom) - and node.module is not None - and node.module.endswith(".core") - and any(alias.name == "apply_driver_features" for alias in node.names) - for node in ast.walk(tree) - ) - calls_helper = any( - isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "apply_driver_features" - for node in ast.walk(tree) - ) - - assert imports_helper is False - assert calls_helper is False - - -@pytest.mark.parametrize("module_path", MYSQL_FAMILY_UNCOMPILED_MODULES) -def test_mysql_family_uncompiled_modules_use_adapter_typing_boundary(module_path: str) -> None: - tree = ast.parse(Path(module_path).read_text()) - direct_imports: list[str] = [] - - for node in ast.walk(tree): - if isinstance(node, ast.Import): - direct_imports.extend(alias.name for alias in node.names) - elif isinstance(node, ast.ImportFrom) and node.module is not None: - direct_imports.append(node.module) - - direct_vendor_imports = [ - import_name - for import_name in direct_imports - if import_name in MYSQL_VENDOR_IMPORT_ROOTS - or import_name.startswith(tuple(f"{root}." for root in MYSQL_VENDOR_IMPORT_ROOTS)) - ] - - assert direct_vendor_imports == [] - - -@pytest.mark.parametrize("module_path", SPANNER_MODULE_PROXY_PATHS) -def test_spanner_modules_do_not_define_module_proxy_getattr(module_path: str) -> None: - tree = ast.parse(Path(module_path).read_text()) - - module_getattrs = [node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name == "__getattr__"] - - assert module_getattrs == [] - - def test_mssql_python_pool_is_implemented_in_pool_module() -> None: from sqlspec.adapters.mssql_python import MssqlPythonConnectionPool assert MssqlPythonConnectionPool.__module__ == "sqlspec.adapters.mssql_python.pool" -@pytest.mark.parametrize(("module_path", "class_name", "method_name"), BRIDGE_CURSOR_CLOSE_METHODS) -def test_bridge_cursor_close_cleanup_suppresses_close_errors( - module_path: str, class_name: str, method_name: str -) -> None: - tree = ast.parse(Path(module_path).read_text()) - class_node = next(node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == class_name) - method_node = next( - node - for node in class_node.body - if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name == method_name - ) - - suppress_calls = [ - item.context_expr - for node in ast.walk(method_node) - if isinstance(node, ast.With) - for item in node.items - if isinstance(item.context_expr, ast.Call) - ] - suppresses_close_errors = any( - isinstance(call.func, ast.Attribute) - and isinstance(call.func.value, ast.Name) - and call.func.value.id == "contextlib" - and call.func.attr == "suppress" - and len(call.args) == 1 - and isinstance(call.args[0], ast.Name) - and call.args[0].id == "Exception" - for call in suppress_calls - ) - - assert suppresses_close_errors - - def test_adapter_parameter_style_handling( adapter_config: "dict[str, Any]", statement_config_for_adapter: StatementConfig ) -> None: @@ -734,27 +615,3 @@ def test_litestar_session_migration_resolves_store_class(adapter: str, config_na store_module = importlib.import_module(f"sqlspec.adapters.{adapter}.litestar.store") assert hasattr(store_module, store_name), f"{adapter}: expected {store_name} in litestar/store.py" - - -_DRIVER_FEATURE_READ_PATTERN = r'driver_features\s*(?:\.get\(|\.setdefault\(|\[)\s*"([A-Za-z]\w*)"' - - -@pytest.mark.parametrize( - "adapter", sorted(path.name for path in Path("sqlspec/adapters").iterdir() if (path / "config.py").exists()) -) -def test_driver_feature_keys_read_by_adapter_are_declared(adapter: str) -> None: - """Every literal driver_features key an adapter reads is annotated on its ``*DriverFeatures`` TypedDict.""" - read_keys: dict[str, str] = {} - for py_file in sorted(Path("sqlspec/adapters", adapter).rglob("*.py")): - for match in re.finditer(_DRIVER_FEATURE_READ_PATTERN, py_file.read_text(encoding="utf-8")): - read_keys.setdefault(match.group(1), str(py_file)) - config_module = importlib.import_module(f"sqlspec.adapters.{adapter}.config") - feature_types = [ - value - for name, value in vars(config_module).items() - if name.endswith("DriverFeatures") and isinstance(value, type) - ] - assert feature_types, f"{adapter} declares no *DriverFeatures TypedDict" - declared = {key for feature_type in feature_types for key in feature_type.__annotations__} - undeclared = {key: path for key, path in read_keys.items() if key not in declared} - assert not undeclared, f"{adapter} reads undeclared driver_features keys: {undeclared}" diff --git a/tests/unit/adapters/test_duckdb/test_pool.py b/tests/unit/adapters/test_duckdb/test_pool.py index d8ca3321e..a89c90255 100644 --- a/tests/unit/adapters/test_duckdb/test_pool.py +++ b/tests/unit/adapters/test_duckdb/test_pool.py @@ -614,12 +614,3 @@ def test_pool_thread_connection_checkout() -> None: {"database": ":memory:"}, secrets=[{"name": "s2", "value_provider": lambda: {}}] ) assert pool_dynamic._has_dynamic_secrets is True - - -def test_pool_unused_constants_removal() -> None: - """Unused pool constants are removed from sqlspec.adapters.duckdb.pool.""" - import sqlspec.adapters.duckdb.pool as pool_mod - - assert not hasattr(pool_mod, "DEFAULT_MIN_POOL") - assert not hasattr(pool_mod, "DEFAULT_MAX_POOL") - assert not hasattr(pool_mod, "POOL_TIMEOUT") diff --git a/tests/unit/adapters/test_mssql_python/test_adk_store.py b/tests/unit/adapters/test_mssql_python/test_adk_store.py index fc76b763a..dd681eb13 100644 --- a/tests/unit/adapters/test_mssql_python/test_adk_store.py +++ b/tests/unit/adapters/test_mssql_python/test_adk_store.py @@ -406,15 +406,6 @@ def capture(_store: MssqlPythonADKStore, sql: str, params: "tuple[Any, ...]" = ( return MssqlPythonADKStore(_mock_config()), calls -def test_adk_store_dead_constants_removed() -> None: - """Verify dead duplicate object/index error constants and trailing aliases are removed.""" - import sqlspec.adapters.mssql_python.adk.store as adk_store - - assert not hasattr(adk_store, "MSSQL_DUPLICATE_OBJECT_ERROR") - assert not hasattr(adk_store, "MSSQL_DUPLICATE_INDEX_ERROR") - assert not hasattr(adk_store, "_escape_sql_literal") - - def test_insert_memory_entries_batches_with_executemany() -> None: """insert_memory_entries batches rows into cursor.executemany.""" mem_config = _mock_config({"native_json": False}) diff --git a/tests/unit/adapters/test_spanner/test_adk_artifact_store.py b/tests/unit/adapters/test_spanner/test_adk_artifact_store.py index afe11764d..f4b5f2139 100644 --- a/tests/unit/adapters/test_spanner/test_adk_artifact_store.py +++ b/tests/unit/adapters/test_spanner/test_adk_artifact_store.py @@ -300,18 +300,3 @@ async def _run_in_tx(fn: Any) -> Any: swept = await store.delete_artifacts_older_than(now) assert len(swept) == 1 tx.execute_update.assert_awaited_once() - - -def test_artifact_store_imports_cleanly_without_adk_store() -> None: - """Verify artifact_store.py has no direct imports from sqlspec.adapters.spanner.adk.store.""" - import ast - from pathlib import Path - - import sqlspec.adapters.spanner.adk.artifact_store as art_mod - - source = Path(art_mod.__file__).read_text() - tree = ast.parse(source) - - for node in ast.walk(tree): - if isinstance(node, ast.ImportFrom) and node.module: - assert "sqlspec.adapters.spanner.adk.store" not in node.module From 044a5f25df3c633404fb81a9e810d108874d95b7 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 22:23:12 +0000 Subject: [PATCH 152/174] test(adapters): type SDK stand-ins and mixed async results accurately --- tests/unit/adapters/test_bigquery/test_core.py | 9 ++++++--- tests/unit/adapters/test_oracledb/test_litestar_store.py | 5 +++-- .../unit/adapters/test_oracledb/test_oracle_adk_store.py | 9 ++++++--- tests/unit/adapters/test_pymssql/test_pool.py | 8 ++++---- tests/unit/adapters/test_spanner/test_events_store.py | 9 +++++++-- 5 files changed, 26 insertions(+), 14 deletions(-) diff --git a/tests/unit/adapters/test_bigquery/test_core.py b/tests/unit/adapters/test_bigquery/test_core.py index 2c905479d..cfbf1184e 100644 --- a/tests/unit/adapters/test_bigquery/test_core.py +++ b/tests/unit/adapters/test_bigquery/test_core.py @@ -5,7 +5,7 @@ from contextlib import nullcontext from datetime import UTC, datetime from types import SimpleNamespace -from typing import Any, cast +from typing import TYPE_CHECKING, Any, cast from unittest.mock import patch import pytest @@ -49,6 +49,9 @@ ) from sqlspec.utils.serializers import to_json +if TYPE_CHECKING: + from collections.abc import Callable + class _RecordingConnection: def __init__(self) -> None: @@ -925,7 +928,7 @@ def _cloud_error(message: str, *, code: int, errors: list[dict[str, str]] | None def test_retry_error_mapped_to_query_timeout() -> None: """RetryError maps to QueryTimeoutError in BigQuery create_mapped_exception.""" - err = RetryError("Deadline exceeded waiting for query", None) + err = cast("Callable[[str, Exception | None], RetryError]", RetryError)("Deadline exceeded waiting for query", None) mapped = create_mapped_exception(err) assert isinstance(mapped, QueryTimeoutError) @@ -933,7 +936,7 @@ def test_retry_error_mapped_to_query_timeout() -> None: def test_exception_handler_catches_google_api_error() -> None: """BigQueryExceptionHandler catches GoogleAPIError and maps it.""" handler = BigQueryExceptionHandler() - err = RetryError("Deadline exceeded", None) + err = cast("Callable[[str, Exception | None], RetryError]", RetryError)("Deadline exceeded", None) with handler: raise err assert isinstance(handler.pending_exception, QueryTimeoutError) diff --git a/tests/unit/adapters/test_oracledb/test_litestar_store.py b/tests/unit/adapters/test_oracledb/test_litestar_store.py index 6634e4e9f..722e5554a 100644 --- a/tests/unit/adapters/test_oracledb/test_litestar_store.py +++ b/tests/unit/adapters/test_oracledb/test_litestar_store.py @@ -16,7 +16,7 @@ class _FakeCursor: def __init__(self, rows: "list[tuple[Any, ...]] | None" = None) -> None: self.rows = list(rows or []) - self.executed: list[tuple[str, dict[str, Any] | None]] = [] + self.executed: list[tuple[str, dict[str, Any] | tuple[object, ...] | None]] = [] self.rowcount = 0 self.closed = False @@ -29,7 +29,7 @@ def __exit__(self, *_: object) -> None: def close(self) -> None: self.closed = True - def execute(self, sql: str, parameters: "dict[str, Any] | None" = None) -> None: + def execute(self, sql: str, parameters: "dict[str, Any] | tuple[object, ...] | None" = None) -> None: self.executed.append((sql, parameters)) def fetchone(self) -> "tuple[Any, ...] | None": @@ -149,6 +149,7 @@ def test_oracle_sync_store_set_large_blob_single_statement() -> None: sql, parameters = cursor.executed[0] assert "MERGE INTO oracle_sessions t" in sql assert parameters is not None + parameters = cast("dict[str, object]", parameters) assert parameters["session_id"] == "session-large" assert connection.created_lobs == [(DB_TYPE_BLOB, large_payload)] assert parameters["data"] == (DB_TYPE_BLOB, large_payload) diff --git a/tests/unit/adapters/test_oracledb/test_oracle_adk_store.py b/tests/unit/adapters/test_oracledb/test_oracle_adk_store.py index 67f45dbf5..87e8adf41 100644 --- a/tests/unit/adapters/test_oracledb/test_oracle_adk_store.py +++ b/tests/unit/adapters/test_oracledb/test_oracle_adk_store.py @@ -5,7 +5,7 @@ from decimal import Decimal from inspect import isawaitable from types import SimpleNamespace -from typing import Any, NotRequired, Self, cast, get_args, get_origin +from typing import TYPE_CHECKING, Any, NotRequired, Self, cast, get_args, get_origin from unittest.mock import AsyncMock, MagicMock from uuid import uuid4 @@ -27,6 +27,9 @@ from sqlspec.exceptions import ImproperConfigurationError, SQLSpecError from sqlspec.extensions.adk import StoredMemory +if TYPE_CHECKING: + from collections.abc import Awaitable + def _mock_config(adk_config: dict[str, object]) -> MagicMock: config = MagicMock() @@ -100,7 +103,7 @@ async def test_memory_batch_errors_commit_only_duplicate_skips(is_async: bool, e entries: list[StoredMemory] = [_memory_entry(), {**_memory_entry(), "id": "memory-2"}] if error_code == 1: - result = store.insert_memory_entries(entries) + result: int | Awaitable[int] = store.insert_memory_entries(entries) if isawaitable(result): result = await result assert result == 1 @@ -197,7 +200,7 @@ def execute(sql: str, _params: Any) -> None: "inserted_at", } - inserted = store.insert_memory_entries([_memory_entry()]) + inserted: int | Awaitable[int] = store.insert_memory_entries([_memory_entry()]) if isawaitable(inserted): inserted = await inserted assert inserted == 1 diff --git a/tests/unit/adapters/test_pymssql/test_pool.py b/tests/unit/adapters/test_pymssql/test_pool.py index abae8e6a7..564b0a7df 100644 --- a/tests/unit/adapters/test_pymssql/test_pool.py +++ b/tests/unit/adapters/test_pymssql/test_pool.py @@ -200,10 +200,10 @@ def test_pool_reentrant_thread_checkout(monkeypatch: Any) -> None: pool = PymssqlConnectionPool({"server": "sql.example.test"}, recycle_seconds=0) with pool.get_connection() as c1: - assert c1 is connection + assert cast(object, c1) is connection try: with pool.get_connection() as c2: - assert c2 is connection + assert cast(object, c2) is connection raise ValueError("inner failure") except ValueError: pass @@ -245,7 +245,7 @@ def connect(self, **kwargs: Any) -> FakeConnection: assert pool.size() == 0 with pytest.raises(ValueError, match="failed transaction"), pool.get_connection() as connection: - assert connection is dirty_connection + assert cast(object, connection) is dirty_connection dirty_connection.autocommit(False) raise ValueError("failed transaction") @@ -254,7 +254,7 @@ def connect(self, **kwargs: Any) -> FakeConnection: assert pool.size() == 0 with pool.get_connection() as replacement: - assert replacement is clean_connection + assert cast(object, replacement) is clean_connection assert clean_connection.closed is False assert clean_connection.autocommit_state is True assert pool.checked_out() == 1 diff --git a/tests/unit/adapters/test_spanner/test_events_store.py b/tests/unit/adapters/test_spanner/test_events_store.py index 226724a2a..9b6ce922d 100644 --- a/tests/unit/adapters/test_spanner/test_events_store.py +++ b/tests/unit/adapters/test_spanner/test_events_store.py @@ -3,7 +3,7 @@ import inspect from types import SimpleNamespace -from typing import Any +from typing import TYPE_CHECKING, Any, cast from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -12,6 +12,9 @@ from sqlspec.adapters.spanner.events import SpannerAsyncEventQueueStore, SpannerSyncEventQueueStore from sqlspec.exceptions import ImproperConfigurationError +if TYPE_CHECKING: + from collections.abc import Callable + def _mock_spanner_config(events_settings: dict[str, Any] | None = None, *, dialect: str = "spanner") -> MagicMock: """Create a mock SpannerSyncConfig.""" @@ -346,7 +349,9 @@ def test_events_store_ddl_error_mapping() -> None: from sqlspec.exceptions import SQLSpecError mock_db = MagicMock() - mock_db.update_ddl.side_effect = SpannerGoogleAPICallError("DDL failed") + mock_db.update_ddl.side_effect = cast("Callable[[str], SpannerGoogleAPICallError]", SpannerGoogleAPICallError)( + "DDL failed" + ) config = MagicMock(spec=SpannerSyncConfig) config.extension_config = {"events": {"queue_table": "test_queue", "ddl_timeout": 300.0}} From 083404476f4200b68aef2b22bc2e35ad35e1489a Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 22:27:52 +0000 Subject: [PATCH 153/174] fix(adbc): honor native autocommit transaction boundaries --- docs/changelog.md | 2 + sqlspec/adapters/adbc/driver.py | 93 ++++++--- .../adapters/postgres/adbc/test_driver.py | 20 ++ tests/unit/adapters/test_adbc/test_core.py | 194 ++++++++++++++---- 4 files changed, 243 insertions(+), 66 deletions(-) diff --git a/docs/changelog.md b/docs/changelog.md index d2c7e2e36..dcb6bec17 100644 --- a/docs/changelog.md +++ b/docs/changelog.md @@ -85,6 +85,8 @@ important operational fixes. * ADBC now closes its own connections if a hook, profile setup, driver setup, or session setup fails. It keeps the first error and leaves supplied connections open. +* Explicit ADBC transactions now honor native `autocommit=True` on supported drivers. If rollback cleanup fails, autocommit stays off so mode restoration cannot commit pending writes. + * Preserved supported Arrow conversion and mutation paths for Asyncmy and Spanner while reporting their native capabilities accurately; corrected ADBC lowercase readers, DuckDB stream ownership and table filtering, and SQL Server Arrow cursor lifetimes. * Corrected Oracle session-renewal bindings, pipeline error conversion, ADK batch-error handling, standalone connection setup, and deferred Arrow failures. Text-only Oracle memory stores remain available on servers without native vector support. diff --git a/sqlspec/adapters/adbc/driver.py b/sqlspec/adapters/adbc/driver.py index 1596e6b5e..8ef83f0ea 100644 --- a/sqlspec/adapters/adbc/driver.py +++ b/sqlspec/adapters/adbc/driver.py @@ -10,7 +10,7 @@ from time import perf_counter from typing import TYPE_CHECKING, Any, Literal, cast, final -from sqlspec.adapters.adbc._typing import AdbcCursor, AdbcNativeError, AdbcSessionContext +from sqlspec.adapters.adbc._typing import AdbcCursor, AdbcNativeError, AdbcNotSupportedError, AdbcSessionContext from sqlspec.adapters.adbc.core import ( _transform_duckdb_copy_parameters, apply_driver_features, @@ -502,43 +502,38 @@ def dispatch_select_stream(self, statement: "SQL", chunk_size: int) -> "SyncRowS return stream def begin(self) -> None: - """Begin database transaction. - - ADBC connections operate with autocommit disabled and manage the - active transaction internally, so no explicit statement is issued. - """ - if hasattr(self.connection, "autocommit"): - if self._previous_autocommit is None: - self._previous_autocommit = bool(self.connection.autocommit) - if self.connection.autocommit: - self.connection.autocommit = False + """Begin a transaction, disabling native autocommit when its state is available.""" + if self._transaction_active: + return + native = getattr(self.connection, "adbc_connection", None) + get_option = getattr(native, "get_option", None) + mode = None + if callable(get_option): + try: + mode = get_option("adbc.connection.autocommit") + except (AdbcNotSupportedError, NotImplementedError): + pass + except AdbcNativeError as exc: + raise create_mapped_exception(exc) from exc + if mode == "true" and native is not None: + try: + native.set_autocommit(False) + except BaseException as exc: + with contextlib.suppress(BaseException): + native.set_autocommit(True) + if isinstance(exc, AdbcNativeError): + raise create_mapped_exception(exc) from exc + raise + self._previous_autocommit = True self._transaction_active = True def commit(self) -> None: """Commit database transaction.""" - try: - self.connection.commit() - except AdbcNativeError as e: - raise create_mapped_exception(e) from e - finally: - self._transaction_active = False - if self._previous_autocommit is not None: - if hasattr(self.connection, "autocommit"): - self.connection.autocommit = self._previous_autocommit - self._previous_autocommit = None + self._finish_transaction(commit=True) def rollback(self) -> None: """Rollback database transaction.""" - try: - self.connection.rollback() - except AdbcNativeError as e: - raise create_mapped_exception(e) from e - finally: - self._transaction_active = False - if self._previous_autocommit is not None: - if hasattr(self.connection, "autocommit"): - self.connection.autocommit = self._previous_autocommit - self._previous_autocommit = None + self._finish_transaction(commit=False) def create_savepoint(self, name: str) -> None: """Create a savepoint within the current transaction. @@ -996,6 +991,42 @@ def _compiled_sql( compiled_sql, prepared_parameters, is_many=statement.is_many, dialect=self._dialect_name ) + def _finish_transaction(self, *, commit: bool) -> None: + restore = self._previous_autocommit + primary: BaseException | None = None + try: + connection = self.connection.adbc_connection if restore is True else self.connection + if commit: + connection.commit() + else: + connection.rollback() + except BaseException as exc: + primary = exc + completed = primary is None + if restore is True and not completed: + with contextlib.suppress(BaseException): + self.connection.adbc_connection.rollback() + completed = True + if restore is True: + self._previous_autocommit = True + self._transaction_active = True + if completed: + try: + self.connection.adbc_connection.set_autocommit(True) + except BaseException as exc: + if primary is None: + primary = exc + else: + self._previous_autocommit = None + self._transaction_active = False + else: + self._previous_autocommit = None + self._transaction_active = False + if isinstance(primary, AdbcNativeError): + raise create_mapped_exception(primary) from primary + if primary is not None: + raise primary + def _connection_in_transaction(self) -> bool: """Check if connection is in transaction. diff --git a/tests/integration/adapters/postgres/adbc/test_driver.py b/tests/integration/adapters/postgres/adbc/test_driver.py index fc86e203a..8c58c9eac 100644 --- a/tests/integration/adapters/postgres/adbc/test_driver.py +++ b/tests/integration/adapters/postgres/adbc/test_driver.py @@ -427,6 +427,26 @@ def test_postgresql_pool_eager_arrow_results_survive_session_and_allow_reuse( assert rows == [{"value": 1}, {"value": 2}, {"value": 3}] +@pytest.mark.xdist_group("postgres") +@pytest.mark.adbc +def test_postgresql_autocommit_transaction_boundaries(adbc_postgres_connection_config: dict[str, str]) -> None: + """Explicit transactions commit or roll back native work and restore autocommit.""" + config = AdbcConfig(connection_config={**adbc_postgres_connection_config, "autocommit": True}) + with config.provide_session() as session: + session.execute("CREATE TEMP TABLE sqlspec_adbc_autocommit_transactions (value INTEGER)") + session.begin() + session.execute("INSERT INTO sqlspec_adbc_autocommit_transactions (value) VALUES (?)", 41) + session.rollback() + assert session.select_value("SELECT COUNT(*) FROM sqlspec_adbc_autocommit_transactions") == 0 + assert session.connection.adbc_connection.get_option("adbc.connection.autocommit") == "true" + + session.begin() + session.execute("INSERT INTO sqlspec_adbc_autocommit_transactions (value) VALUES (?)", 42) + session.commit() + assert session.select_value("SELECT value FROM sqlspec_adbc_autocommit_transactions") == 42 + assert session.connection.adbc_connection.get_option("adbc.connection.autocommit") == "true" + + @contextmanager def _pooled_postgresql_config( connection_config: dict[str, str], *, max_size: int = 1, conn_kwargs: dict[str, str] | None = None diff --git a/tests/unit/adapters/test_adbc/test_core.py b/tests/unit/adapters/test_adbc/test_core.py index f8ee314cd..860096ddd 100644 --- a/tests/unit/adapters/test_adbc/test_core.py +++ b/tests/unit/adapters/test_adbc/test_core.py @@ -2,16 +2,17 @@ from collections.abc import Sequence from decimal import Decimal +from pathlib import Path from types import NoneType, SimpleNamespace from typing import Any, cast -from unittest.mock import Mock +from unittest.mock import Mock, call from uuid import UUID import pyarrow as pa import pytest -from adbc_driver_manager import AdbcStatusCode, DatabaseError +from adbc_driver_manager import AdbcStatusCode, DatabaseError, NotSupportedError -from sqlspec.adapters.adbc import AdbcDriver +from sqlspec.adapters.adbc import AdbcConfig, AdbcDriver from sqlspec.adapters.adbc import core as adbc_core from sqlspec.adapters.adbc import driver as adbc_driver_mod from sqlspec.adapters.adbc._typing import AdbcSessionContext @@ -650,42 +651,170 @@ def test_adbc_driver_dispatch_execute_lowercases_oracle_columns() -> None: assert result.selected_data == [{"user_id": 1, "MixedCase": "a"}] -def test_adbc_begin_with_autocommit() -> None: - """AdbcDriver.begin disables autocommit and commit/rollback restores it.""" +@pytest.mark.parametrize("autocommit", [False, True]) +def test_adbc_begin_with_native_autocommit(autocommit: bool) -> None: + """Real explicit transactions commit/roll back and restore the native autocommit mode.""" + config = AdbcConfig(connection_config={"driver_name": "sqlite", "autocommit": autocommit}) + expected_mode = "true" if autocommit else "false" + with config.provide_session() as session: + native = session.connection.adbc_connection + session.execute("CREATE TABLE transaction_values (value INTEGER)") + session.commit() + session.begin() + begun_mode = native.get_option("adbc.connection.autocommit") + session.execute("INSERT INTO transaction_values (value) VALUES (:value)", value=1) + session.rollback() + assert session.select_value("SELECT COUNT(*) FROM transaction_values") == 0 + assert begun_mode == "false" + assert native.get_option("adbc.connection.autocommit") == expected_mode + session.begin() + session.execute("INSERT INTO transaction_values (value) VALUES (:value)", value=2) + session.commit() + assert session.select_value("SELECT value FROM transaction_values") == 2 + assert native.get_option("adbc.connection.autocommit") == expected_mode + + +@pytest.mark.parametrize("capability", ["missing", "unsupported"]) +def test_adbc_begin_preserves_wrapper_transactions_without_native_getter(capability: str) -> None: + """Older/unsupported getters retain ordinary DB-API transaction capability behavior.""" + native = SimpleNamespace() + if capability == "unsupported": + native.get_option = Mock(side_effect=NotSupportedError("getter unsupported")) + committed, rolled_back = Mock(), Mock() + connection = SimpleNamespace( + adbc_get_info=lambda: {"driver_name": "sqlite"}, adbc_connection=native, commit=committed, rollback=rolled_back + ) + driver = AdbcDriver(cast("Any", connection), dialect="sqlite") + driver.begin() + driver.commit() + driver.begin() + driver.rollback() + committed.assert_called_once_with() + rolled_back.assert_called_once_with() + + +@pytest.mark.parametrize("operation", ["commit", "rollback"]) +@pytest.mark.parametrize("interrupted", [False, True]) +def test_adbc_native_transaction_preserves_primary_on_restoration_failure(operation: str, interrupted: bool) -> None: + """Explicitly switched native transactions preserve primary errors while attempting restoration.""" + connection = Mock() + connection.adbc_get_info.return_value = {"driver_name": "sqlite"} + native = connection.adbc_connection + native.get_option.return_value = "true" + native.set_autocommit.side_effect = [None, RuntimeError("restore failed")] + primary = ( + KeyboardInterrupt("transaction interrupted") + if interrupted + else DatabaseError("transaction failed", status_code=AdbcStatusCode.UNKNOWN) + ) + getattr(native, operation).side_effect = [primary, None] if operation == "rollback" else primary + driver = AdbcDriver(connection, dialect="sqlite") + driver.begin() + expected = KeyboardInterrupt if interrupted else SQLSpecError + with pytest.raises(expected) as caught: + getattr(driver, operation)() + if interrupted: + assert caught.value is primary + else: + assert caught.value.__cause__ is primary + assert native.set_autocommit.call_args_list == [call(False), call(True)] + connection.commit.assert_not_called() + connection.rollback.assert_not_called() + + +@pytest.mark.parametrize("interrupted", [False, True]) +@pytest.mark.parametrize("operation", ["commit", "rollback"]) +def test_adbc_failed_native_rollback_never_commits_pending_rows( + tmp_path: Path, interrupted: bool, operation: str +) -> None: + """A failed rollback and cleanup leave real work uncommitted until an explicit retry.""" + path = str(tmp_path / "rollback.sqlite") + config = AdbcConfig(connection_config={"driver_name": "sqlite", "uri": path, "autocommit": True}) + primary = ( + KeyboardInterrupt(f"{operation} interrupted") + if interrupted + else DatabaseError(f"{operation} failed", status_code=AdbcStatusCode.UNKNOWN) + ) - class MockConn: - def __init__(self) -> None: - self.autocommit = True - self.committed = False - self.rolled_back = False + class NativeBoundary: + def __init__(self, native: Any) -> None: + self.native = native + self.remaining = 2 - def adbc_get_info(self) -> dict[str, str]: - return {"driver_name": "sqlite"} + def get_option(self, key: str) -> str: + return cast("str", self.native.get_option(key)) + + def set_autocommit(self, value: bool) -> None: + self.native.set_autocommit(value) def commit(self) -> None: - self.committed = True + if operation == "commit" and self.remaining: + self.remaining -= 1 + raise primary + self.native.commit() def rollback(self) -> None: - self.rolled_back = True + if self.remaining: + self.remaining -= 1 + if self.remaining == 1: + raise primary + raise RuntimeError("cleanup rollback failed") + self.native.rollback() - conn = MockConn() - driver = AdbcDriver(cast("Any", conn)) - assert conn.autocommit is True - driver.begin() - assert conn.autocommit is False - assert driver._transaction_active is True - driver.commit() - assert conn.autocommit is True - assert driver._transaction_active is False - assert conn.committed is True + class ConnectionBoundary: + def __init__(self, connection: Any) -> None: + self.connection = connection + self.adbc_connection = NativeBoundary(connection.adbc_connection) - conn.autocommit = True - driver.begin() - assert conn.autocommit is False + def adbc_get_info(self) -> Any: + return self.connection.adbc_get_info() + + def cursor(self) -> Any: + return self.connection.cursor() + + def commit(self) -> None: + self.connection.commit() + + def rollback(self) -> None: + self.connection.rollback() + + with config.provide_connection() as connection: + driver = AdbcDriver(cast("Any", ConnectionBoundary(connection)), dialect="sqlite") + driver.execute("CREATE TABLE pending_values (value INTEGER)") + driver.begin() + driver.execute("INSERT INTO pending_values (value) VALUES (:value)", value=1) + expected = KeyboardInterrupt if interrupted else SQLSpecError + with pytest.raises(expected) as caught: + getattr(driver, operation)() + if interrupted: + assert caught.value is primary + else: + assert caught.value.__cause__ is primary + observer = AdbcConfig(connection_config={"driver_name": "sqlite", "uri": path, "autocommit": True}) + with observer.provide_session() as session: + assert session.select_value("SELECT COUNT(*) FROM pending_values") == 0 + assert connection.adbc_connection.get_option("adbc.connection.autocommit") == "false" + driver.begin() + driver.rollback() + assert connection.adbc_connection.get_option("adbc.connection.autocommit") == "true" + assert driver.select_value("SELECT COUNT(*) FROM pending_values") == 0 + + +def test_adbc_native_begin_failure_restores_mode_and_preserves_native_cause() -> None: + """A failed native mode switch does not start a logical transaction or hide its error.""" + connection = Mock() + connection.adbc_get_info.return_value = {"driver_name": "sqlite"} + native = connection.adbc_connection + primary = DatabaseError("begin failed", status_code=AdbcStatusCode.UNKNOWN) + native.get_option.return_value = "true" + native.set_autocommit.side_effect = [primary, RuntimeError("restore failed")] + driver = AdbcDriver(connection, dialect="sqlite") + with pytest.raises(SQLSpecError) as caught: + driver.begin() + assert caught.value.__cause__ is primary + assert native.set_autocommit.call_args_list == [call(False), call(True)] driver.rollback() - assert conn.autocommit is True - assert driver._transaction_active is False - assert conn.rolled_back is True + connection.rollback.assert_called_once_with() def test_pending_exception_preserves_cause() -> None: @@ -809,11 +938,6 @@ def test_bytearray_memoryview_coercion() -> None: assert coercion[memoryview] is bytes -def test_core_unused_symbols() -> None: - """Dead alias _prepare_batch_with_casts is removed from adbc.core.""" - assert not hasattr(adbc_core, "_prepare_batch_with_casts") - - def test_session_context_driver() -> None: """AdbcSessionContext uses provided driver_class without lazy import.""" driver_instance = Mock() From 05f6c920f4a836eda8d7d6b6adc3711936774b5d Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 22:35:03 +0000 Subject: [PATCH 154/174] test(adapters): align catalog and stream stand-ins with native types --- .../test_spanner_query_options.py | 23 +++++++++++++++---- .../data_dictionary/test_mssql_metadata.py | 17 +++++++++++--- 2 files changed, 33 insertions(+), 7 deletions(-) diff --git a/tests/unit/adapters/test_spanner/test_spanner_query_options.py b/tests/unit/adapters/test_spanner/test_spanner_query_options.py index bde1d4b49..9c1ffbe15 100644 --- a/tests/unit/adapters/test_spanner/test_spanner_query_options.py +++ b/tests/unit/adapters/test_spanner/test_spanner_query_options.py @@ -2,7 +2,7 @@ from collections.abc import AsyncIterator, Iterator from types import SimpleNamespace -from typing import Any +from typing import TYPE_CHECKING, Any, cast from unittest.mock import AsyncMock, MagicMock import pytest @@ -14,6 +14,11 @@ from sqlspec.adapters.spanner.driver import SpannerAsyncDriver, SpannerSyncDriver from sqlspec.exceptions import OperationalError +if TYPE_CHECKING: + from collections.abc import Callable + + from sqlspec.adapters.spanner._typing import SpannerAsyncConnection, SpannerSyncConnection + class _StreamResult: def __init__(self, error: BaseException | None) -> None: @@ -61,13 +66,19 @@ async def _consume_stream(driver: Any, is_async: bool, *, fails_fetch: bool = Fa async def test_stream_maps_sdk_errors_with_ambient_application_exception( is_async: bool, phase: str, ambient: bool ) -> None: - error = api_exceptions.ServiceUnavailable("Spanner unavailable") + error = cast("Callable[[str], api_exceptions.ServiceUnavailable]", api_exceptions.ServiceUnavailable)( + "Spanner unavailable" + ) result = _StreamResult(error) execute = AsyncMock() if is_async else MagicMock() execute.side_effect = error if phase == "start" else None execute.return_value = result connection = SimpleNamespace(execute_sql=execute) - driver = SpannerAsyncDriver(connection=connection) if is_async else SpannerSyncDriver(connection=connection) + driver = ( + SpannerAsyncDriver(connection=cast("SpannerAsyncConnection", connection)) + if is_async + else SpannerSyncDriver(connection=cast("SpannerSyncConnection", connection)) + ) with pytest.raises(OperationalError, match="Spanner unavailable") as caught: try: @@ -89,7 +100,11 @@ async def test_stream_succeeds_with_ambient_application_exception(is_async: bool result = _StreamResult(None) execute = AsyncMock(return_value=result) if is_async else MagicMock(return_value=result) connection = SimpleNamespace(execute_sql=execute) - driver = SpannerAsyncDriver(connection=connection) if is_async else SpannerSyncDriver(connection=connection) + driver = ( + SpannerAsyncDriver(connection=cast("SpannerAsyncConnection", connection)) + if is_async + else SpannerSyncDriver(connection=cast("SpannerSyncConnection", connection)) + ) error = LookupError("application error") try: diff --git a/tests/unit/data_dictionary/test_mssql_metadata.py b/tests/unit/data_dictionary/test_mssql_metadata.py index 789828318..470422978 100644 --- a/tests/unit/data_dictionary/test_mssql_metadata.py +++ b/tests/unit/data_dictionary/test_mssql_metadata.py @@ -2,7 +2,7 @@ import sqlite3 from collections.abc import Iterator -from typing import Any, cast +from typing import TYPE_CHECKING, Any, TypeVar, cast, overload import pytest @@ -22,6 +22,11 @@ ) from sqlspec.driver import SyncDataDictionaryBase +if TYPE_CHECKING: + from collections.abc import Callable + +_CursorT = TypeVar("_CursorT", bound=sqlite3.Cursor) + REQUIRED_MSSQL_DOMAIN_QUERIES = { "schemas": ("list",), "objects": ("by_schema",), @@ -77,8 +82,14 @@ def execute(self, sql: str, parameters: Any = ()) -> "CatalogCursor": class CatalogConnection(sqlite3.Connection): """Provide the DB-API cursor surface shared by the SQL Server SDKs.""" - def cursor(self, *args: Any, **kwargs: Any) -> CatalogCursor: - return super().cursor(factory=CatalogCursor) + @overload + def cursor(self, factory: None = None) -> CatalogCursor: ... + + @overload + def cursor(self, factory: "Callable[[sqlite3.Connection], _CursorT]") -> _CursorT: ... + + def cursor(self, factory: "Callable[[sqlite3.Connection], sqlite3.Cursor] | None" = None) -> sqlite3.Cursor: + return super().cursor(factory=CatalogCursor if factory is None else factory) @pytest.fixture(params=[MssqlPythonDriver, PymssqlDriver], ids=["mssql_python", "pymssql"]) From c456859b2e160db5a5da01a6bb4a38889af29434 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 23:05:26 +0000 Subject: [PATCH 155/174] fix(core): reconstruct parameter declarations during serialization --- sqlspec/core/parameters/_declared.py | 11 +++++ .../core/test_declared_params_carriage.py | 43 +++++++++++-------- 2 files changed, 36 insertions(+), 18 deletions(-) diff --git a/sqlspec/core/parameters/_declared.py b/sqlspec/core/parameters/_declared.py index 8b6ae92dc..06ee87d9e 100644 --- a/sqlspec/core/parameters/_declared.py +++ b/sqlspec/core/parameters/_declared.py @@ -58,6 +58,10 @@ def __repr__(self) -> str: f"description={self.description!r}, required={self.required!r})" ) + def __reduce__(self) -> "tuple[Callable[..., ParameterDeclaration], tuple[str, str, str | None, bool]]": + """Preserve declaration metadata through constructor-based reconstruction.""" + return (_rebuild_parameter_declaration, (self.name, self.type_str, self.description, self.required)) + def register_param_type(name: str, py_type: ParamTypeMatcher) -> None: """Register or override a declared-type-string matcher. @@ -101,6 +105,13 @@ def matches_param_type(type_str: str, value: object) -> bool: return resolved(value) +def _rebuild_parameter_declaration( + name: str, type_str: str, description: str | None, required: bool +) -> ParameterDeclaration: + """Restore metadata while retaining the constructor's keyword-only required flag.""" + return ParameterDeclaration(name, type_str, description, required=required) + + def _is_json_value(value: object) -> bool: """Return whether a value can be encoded by SQLSpec's JSON serializer.""" if not isinstance(value, _JSON_VALUE_TYPES): diff --git a/tests/unit/core/test_declared_params_carriage.py b/tests/unit/core/test_declared_params_carriage.py index cb53336be..31faf6adb 100644 --- a/tests/unit/core/test_declared_params_carriage.py +++ b/tests/unit/core/test_declared_params_carriage.py @@ -4,9 +4,15 @@ plus loader population and driver-derivation preservation. """ +import copy +import pickle + +import pytest + from sqlspec.core import ParameterDeclaration from sqlspec.core._pool import get_sql_pool from sqlspec.core.statement import SQL +from sqlspec.loader import SQLFileLoader _SENTINEL = (ParameterDeclaration("a", "int", required=False),) @@ -60,8 +66,6 @@ def test_pool_recycle_does_not_leak_declared_parameters() -> None: def test_get_sql_populates_declared_parameters() -> None: - from sqlspec.loader import SQLFileLoader - loader = SQLFileLoader() loader.add_named_sql("q", "select :a", parameters=[ParameterDeclaration("a", "int", description="id")]) sql = loader.get_sql("q") @@ -69,8 +73,6 @@ def test_get_sql_populates_declared_parameters() -> None: def test_undeclared_get_sql_has_empty_declarations() -> None: - from sqlspec.loader import SQLFileLoader - loader = SQLFileLoader() loader.add_named_sql("plain", "select 1") assert loader.get_sql("plain").declared_parameters == () @@ -81,20 +83,25 @@ def test_constructor_kwarg_sets_declarations() -> None: assert sql.declared_parameters == _SENTINEL -def test_declarations_survive_deepcopy() -> None: - import copy - - sql = SQL("select :a", {"a": 1}, declared_parameters=_SENTINEL) - new = copy.deepcopy(sql) - assert new.declared_parameters == _SENTINEL - - -def test_declarations_survive_pickle_roundtrip() -> None: - import pickle - - sql = SQL("select :a", {"a": 1}, declared_parameters=_SENTINEL) - new = pickle.loads(pickle.dumps(sql)) - assert new.declared_parameters == _SENTINEL +@pytest.mark.parametrize("operation", ["pickle", "deepcopy"]) +@pytest.mark.parametrize("source", ["constructor", "loader"]) +def test_declarations_survive_serialization(operation: str, source: str) -> None: + """Public SQL round trips preserve declaration types and all metadata fields.""" + declarations = ( + ParameterDeclaration("a", "int", "Account identifier", required=False), + ParameterDeclaration("b", "str", required=True), + ) + if source == "loader": + loader = SQLFileLoader() + loader.add_named_sql("declared", "SELECT :a, :b", parameters=list(declarations)) + sql = loader.get_sql("declared").copy(parameters={"a": 1, "b": "group"}) + else: + sql = SQL("SELECT :a, :b", a=1, b="group", declared_parameters=declarations) + + restored = pickle.loads(pickle.dumps(sql)) if operation == "pickle" else copy.deepcopy(sql) + + assert restored.declared_parameters == declarations + assert all(type(declaration) is ParameterDeclaration for declaration in restored.declared_parameters) def test_declarations_survive_driver_prepare_with_filter() -> None: From a4fb9974a0fcb9933b887322234a7366d1fedc56 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 23:05:26 +0000 Subject: [PATCH 156/174] fix(builder): normalize native foreign key metadata explicitly --- sqlspec/builder/_ddl.py | 22 ++++++++++++- .../builder/test_ddl_metadata_compiler.py | 31 +++++++++++++++++++ 2 files changed, 52 insertions(+), 1 deletion(-) diff --git a/sqlspec/builder/_ddl.py b/sqlspec/builder/_ddl.py index 95adab5bc..18a1f5c43 100644 --- a/sqlspec/builder/_ddl.py +++ b/sqlspec/builder/_ddl.py @@ -28,7 +28,14 @@ parse_order_expression, ) from sqlspec.core import SQL -from sqlspec.data_dictionary._types import DDLResult, MetadataFidelity, MetadataSource, MetadataSupport, ObjectIdentity +from sqlspec.data_dictionary._types import ( + DDLResult, + ForeignKeyMetadata, + MetadataFidelity, + MetadataSource, + MetadataSupport, + ObjectIdentity, +) from sqlspec.exceptions import SQLBuilderError from sqlspec.utils.arrow_helpers import sql_type_for_category, sql_type_from_arrow_type from sqlspec.utils.type_guards import has_sqlglot_expression @@ -3064,6 +3071,19 @@ def _normalize_metadata_row(item: Any) -> dict[str, Any]: result.setdefault("schema_name", nested_ident.schema) return result + if isinstance(item, ForeignKeyMetadata): + return { + "table_name": item.table_name, + "column_name": item.column_name, + "referenced_table": item.referenced_table, + "referenced_column": item.referenced_column, + "constraint_name": item.constraint_name, + "schema": item.schema, + "referenced_schema": item.referenced_schema, + "on_delete": item.on_delete, + "on_update": item.on_update, + } + ident = getattr(item, "identity", None) attrs = getattr(item, "attributes", None) if isinstance(ident, ObjectIdentity) and isinstance(attrs, Mapping): diff --git a/tests/unit/builder/test_ddl_metadata_compiler.py b/tests/unit/builder/test_ddl_metadata_compiler.py index c6693994e..204b71ec7 100644 --- a/tests/unit/builder/test_ddl_metadata_compiler.py +++ b/tests/unit/builder/test_ddl_metadata_compiler.py @@ -24,6 +24,37 @@ from sqlspec.exceptions import SQLBuilderError +def test_create_table_from_metadata_foreign_key_preserves_schemas_and_actions() -> None: + """Owned foreign key metadata preserves reference qualification and both actions.""" + foreign_key = ForeignKeyMetadata( + table_name="orders", + column_name="customer_id", + referenced_table="customers", + referenced_column="id", + constraint_name="fk_orders_customer", + schema="sales", + referenced_schema="crm", + on_delete="CASCADE", + on_update="SET NULL", + ) + + built = CreateTable.from_metadata( + "orders", + [{"column_name": "customer_id", "data_type": "INTEGER"}], + foreign_keys=[foreign_key], + schema="sales", + dialect="postgres", + ).build() + + assert built.sql == ( + 'CREATE TABLE "sales"."orders" (\n' + ' "customer_id" INT,\n' + ' CONSTRAINT "fk_orders_customer" FOREIGN KEY ("customer_id") ' + 'REFERENCES "crm"."customers"("id") ON DELETE CASCADE ON UPDATE SET NULL\n' + ")" + ) + + def test_create_table_from_metadata_composite_pk_fk_check_unique_defaults_and_computed() -> None: """Test CreateTable.from_metadata with composite PK, FK, CHECK, UNIQUE, defaults, identity, and computed columns.""" columns: list[dict[str, Any]] = [ From 7f85a173011a7bd87800fd0529e183925e5392bf Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 23:05:30 +0000 Subject: [PATCH 157/174] fix(driver): consume callback profiling state across compiled boundaries --- sqlspec/driver/_async.py | 6 ++---- sqlspec/driver/_common.py | 6 ++++++ sqlspec/driver/_sync.py | 6 ++---- .../unit/adapters/test_adbc/test_profiling.py | 19 +++++++++++++++++++ 4 files changed, 29 insertions(+), 8 deletions(-) diff --git a/sqlspec/driver/_async.py b/sqlspec/driver/_async.py index 19235d8b9..f7e8c0ab7 100644 --- a/sqlspec/driver/_async.py +++ b/sqlspec/driver/_async.py @@ -340,8 +340,7 @@ async def dispatch_execute_script(self, cursor: Any, statement: "SQL") -> Execut self._release_pooled_statement(single_stmt) if has_observers: stmt_duration = perf_counter() - stmt_started - stmt_profile: dict[str, Any] | None = self._pending_statement_profile - self._pending_statement_profile = None + stmt_profile = self._consume_statement_profile() self._emit_script_statement_event( sql=stmt, parameters=prepared_parameters, @@ -1659,8 +1658,7 @@ async def _execute_with_observability_async( runtime.span_manager.end_span(span, error=exc) await runtime.emit_error_async(exc, **query_context) raise - profile: dict[str, Any] | None = self._pending_statement_profile - self._pending_statement_profile = None + profile = self._consume_statement_profile() runtime.span_manager.end_span(span) duration = perf_counter() - started await runtime.emit_query_complete_async(**{**query_context, "rows_affected": result.rows_affected}) diff --git a/sqlspec/driver/_common.py b/sqlspec/driver/_common.py index 647dce1b0..4acf51ccc 100644 --- a/sqlspec/driver/_common.py +++ b/sqlspec/driver/_common.py @@ -1819,6 +1819,12 @@ def attach_observability(self, runtime: "ObservabilityRuntime") -> None: self._observability = runtime self._refresh_statement_cache_state() + def _consume_statement_profile(self) -> "dict[str, Any] | None": + """Take the profile produced by an execution callback and clear pending ownership.""" + profile = self._pending_statement_profile + self._pending_statement_profile = None + return profile + def storage_capabilities(self) -> "StorageCapabilities": """Return cached storage capabilities for the active driver. diff --git a/sqlspec/driver/_sync.py b/sqlspec/driver/_sync.py index 8f49a40d3..731dba527 100644 --- a/sqlspec/driver/_sync.py +++ b/sqlspec/driver/_sync.py @@ -304,8 +304,7 @@ def dispatch_execute_script(self, cursor: Any, statement: "SQL") -> ExecutionRes self._release_pooled_statement(single_stmt) if has_observers: stmt_duration = perf_counter() - stmt_started - stmt_profile: dict[str, Any] | None = self._pending_statement_profile - self._pending_statement_profile = None + stmt_profile = self._consume_statement_profile() self._emit_script_statement_event( sql=stmt, parameters=prepared_parameters, @@ -1640,8 +1639,7 @@ def _execute_with_observability_sync( runtime.span_manager.end_span(span, error=exc) runtime.emit_error_sync(exc, **query_context) raise - profile: dict[str, Any] | None = self._pending_statement_profile - self._pending_statement_profile = None + profile = self._consume_statement_profile() runtime.span_manager.end_span(span) duration = perf_counter() - started runtime.emit_query_complete_sync(**{**query_context, "rows_affected": result.rows_affected}) diff --git a/tests/unit/adapters/test_adbc/test_profiling.py b/tests/unit/adapters/test_adbc/test_profiling.py index 079755eda..2900eb960 100644 --- a/tests/unit/adapters/test_adbc/test_profiling.py +++ b/tests/unit/adapters/test_adbc/test_profiling.py @@ -10,6 +10,7 @@ import pytest from sqlspec.adapters.adbc import AdbcConfig +from sqlspec.exceptions import SQLSpecError from sqlspec.observability import ObservabilityConfig, StatementEvent @@ -159,6 +160,24 @@ def test_adbc_duckdb_enable_statement_profiling_and_execute_script_events() -> N assert outer_event.script_ordinal is None +def test_adbc_duckdb_profiles_remain_scoped_after_execution_error() -> None: + """Each successful statement consumes its own profile without leaking a failed query's state.""" + events: list[StatementEvent] = [] + config = AdbcConfig( + connection_config={"driver_name": "duckdb"}, + driver_features={"enable_statement_profiling": True}, + observability_config=ObservabilityConfig(statement_observers=(events.append,)), + ) + with config.provide_session() as session: + assert session.select_value("SELECT 1 AS value") == 1 + with pytest.raises(SQLSpecError): + session.execute("SELECT missing_column FROM missing_table") + session.rollback() + assert session.select_value("SELECT 2 AS value") == 2 + assert [event.sql for event in events] == ["SELECT 1 AS value", "SELECT 2 AS value"] + assert all(isinstance(event.profile, dict) for event in events) + + def test_adbc_duckdb_statement_profiling_without_duckdb_package() -> None: """ADBC statement profiling with DuckDB backend works without the duckdb Python library.""" script = """ From 508204ad0a65fdd8103437cddc56197a06c88bab Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 23:05:55 +0000 Subject: [PATCH 158/174] test(core): verify public behavior in installed compiled packages --- tests/unit/core/test_cache.py | 16 --------- tests/unit/core/test_cursor_filter.py | 11 +++++-- tests/unit/core/test_data_dictionary.py | 38 +++++++++++----------- tests/unit/core/test_filters.py | 41 ++++++------------------ tests/unit/core/test_type_conversion.py | 32 ++++++++++-------- tests/unit/dialects/test_generators.py | 11 ------- tests/unit/dialects/test_registration.py | 7 +++- tests/unit/dialects/test_spanner.py | 17 ---------- 8 files changed, 59 insertions(+), 114 deletions(-) diff --git a/tests/unit/core/test_cache.py b/tests/unit/core/test_cache.py index 31caeba6d..12fa9cf0a 100644 --- a/tests/unit/core/test_cache.py +++ b/tests/unit/core/test_cache.py @@ -16,7 +16,6 @@ across the entire SQLSpec system. """ -import inspect import logging import threading import time @@ -451,21 +450,6 @@ def test_filters_view_external_api_smoke() -> None: assert len(view.to_canonical()) == 3 -def test_cache_external_api_docstring_markers() -> None: - """Dead-code policy markers identify externally callable cache helpers.""" - for member in ( - NamespacedCache.delete_statement, - NamespacedCache.delete_expression, - NamespacedCache.delete_optimized, - NamespacedCache.delete_builder, - NamespacedCache.delete_file, - cache_module.FiltersView.get_by_field, - cache_module.FiltersView.has_field, - cache_module.FiltersView.to_canonical, - ): - assert "External/extension API" in (inspect.getdoc(member) or "") - - def test_namespaced_cache_respects_config_flags() -> None: """Test that cache config flags disable namespaces.""" original_config = get_cache_config() diff --git a/tests/unit/core/test_cursor_filter.py b/tests/unit/core/test_cursor_filter.py index 42bfcdc8b..8fde202c3 100644 --- a/tests/unit/core/test_cursor_filter.py +++ b/tests/unit/core/test_cursor_filter.py @@ -90,14 +90,19 @@ def test_composite_cursor_key_shorthand(keys: Any, expected: list[CursorKey]) -> assert shorthand.encode({"v": 10, "id": 3}, backward=True) == explicit.encode({"v": 10, "id": 3}, backward=True) -@pytest.mark.parametrize( - "keys", ["", [""], [("id", "bad")], [("id",)], [("id", "asc", "last")], [3], ["id", CursorKey("id")]] -) +@pytest.mark.parametrize("keys", ["", [""], [("id", "bad")], ["id", CursorKey("id")]]) def test_invalid_cursor_key_shorthand(keys: Any) -> None: with pytest.raises(ValueError): CursorFilter(keys, 2) +@pytest.mark.parametrize("keys", [[("id",)], [("id", "asc", "last")], [3]]) +def test_cursor_key_shorthand_rejects_unsupported_input_shapes(keys: Any) -> None: + """Inputs outside CursorKeys are rejected by validation or compiled type guards.""" + with pytest.raises((TypeError, ValueError)): + CursorFilter(keys, 2) + + def test_eager_cursor_validation() -> None: with pytest.raises(InvalidCursorError): CursorFilter([CursorKey("id")], 10, "invalid") diff --git a/tests/unit/core/test_data_dictionary.py b/tests/unit/core/test_data_dictionary.py index c968aa666..929d37411 100644 --- a/tests/unit/core/test_data_dictionary.py +++ b/tests/unit/core/test_data_dictionary.py @@ -1,8 +1,10 @@ """Unit tests for the data dictionary registry and loader.""" import re +import subprocess +import sys from importlib import resources -from typing import Any, cast +from typing import Any import pytest @@ -10,7 +12,6 @@ from sqlspec.core import SQL from sqlspec.data_dictionary import ( DataDictionaryLoader, - DialectConfig, VersionInfo, get_data_dictionary_loader, get_dialect_config, @@ -151,11 +152,6 @@ def test_row_locking_capability_flags_are_introspectable() -> None: assert get_dialect_config("mysql").get_feature_version("supports_skip_locked") == VersionInfo(8, 0, 1) -def test_registry_dialects_loaded_annotation_is_bool() -> None: - """_DIALECTS_LOADED should be annotated to avoid Literal[False] narrowing.""" - assert registry.__annotations__["_DIALECTS_LOADED"] is bool - - def test_registry_load_default_dialects_sets_loaded_flag(monkeypatch) -> None: """_load_default_dialects sets the loaded flag after importing dialects.""" calls: list[str] = [] @@ -177,16 +173,18 @@ def fail_import(name: str) -> None: registry._load_default_dialects() -def test_registry_get_dialect_config_triggers_load(monkeypatch) -> None: - """get_dialect_config calls the default dialect loader before lookup.""" - calls = 0 - config = cast("DialectConfig", object()) - monkeypatch.setitem(registry._DIALECT_CONFIGS, "example", config) - - def fake_load() -> None: - nonlocal calls - calls += 1 - - monkeypatch.setattr(registry, "_load_default_dialects", fake_load) - assert registry.get_dialect_config("example") is config - assert calls == 1 +def test_registry_get_dialect_config_triggers_load() -> None: + """The first public lookup loads built-in dialect configurations.""" + script = """ +import sys +from sqlspec.data_dictionary import DialectConfig, get_dialect_config + +assert "sqlspec.data_dictionary.dialects" not in sys.modules +config = get_dialect_config("PostgreSQL") +assert isinstance(config, DialectConfig) +assert config.name == "postgres" +assert config.get_feature_flag("supports_transactions") is True +assert "sqlspec.data_dictionary.dialects" in sys.modules +assert get_dialect_config("pg") is config +""" + subprocess.run([sys.executable, "-c", script], capture_output=True, text=True, check=True) diff --git a/tests/unit/core/test_filters.py b/tests/unit/core/test_filters.py index cb68201b2..f52b5cbe6 100644 --- a/tests/unit/core/test_filters.py +++ b/tests/unit/core/test_filters.py @@ -7,11 +7,10 @@ import tempfile from dataclasses import dataclass from datetime import datetime -from pathlib import Path from typing import TYPE_CHECKING, Any, assert_type, cast import pytest -from sqlglot import exp +from sqlglot import exp, parse_one import sqlspec.core.filters as filters_module from sqlspec import sql as sql_builder @@ -61,31 +60,15 @@ def test_public_filters_alias_points_to_core_filters_module() -> None: assert sqlspec.filters is filters_module -def test_filters_docs_render_inherited_members() -> None: - """Docs include public filters whose methods are inherited.""" - docs = Path("docs/reference/core/filters.md").read_text() - for class_name in ("BeforeAfterFilter", "SearchFilter", "InCollectionFilter", "NotInCollectionFilter"): - assert f"::: sqlspec.core.filters.{class_name}" in docs +def test_search_filter_preserves_qualified_column_and_binds_value() -> None: + """A qualified field remains qualified in the rendered search predicate.""" + result = apply_filter(SQL("SELECT * FROM users"), SearchFilter("users.name", "needle")) + condition = parse_one(result.raw_sql).find(exp.Like) - -def test_statement_filter_get_column_expression_delegates_to_parse_column_for_condition( - monkeypatch: pytest.MonkeyPatch, -) -> None: - """StatementFilter should use the shared fast column parser.""" - calls: list[str | exp.Expression] = [] - original = StatementFilter._get_column_expression - - def fake_parse(field_name: str | exp.Expression) -> exp.Expr: - calls.append(field_name) - return exp.column("parsed") - - monkeypatch.setattr("sqlspec.core.filters.parse_column_for_condition", fake_parse) - - result = original(SearchFilter("name", "needle"), "users.name") - - assert isinstance(result, exp.Column) - assert result.name == "parsed" - assert calls == ["users.name"] + assert isinstance(condition, exp.Like) + assert condition.this == exp.column("name", table="users") + assert condition.expression == exp.Placeholder(this="users_name_search") + assert result.named_parameters == {"users_name_search": "%needle%"} def test_canonicalize_filters_deduplicates_and_sorts_by_class_and_field() -> None: @@ -945,12 +928,6 @@ def test_order_by_filter_rejects_invalid_sort_order_at_construction() -> None: OrderByFilter("name", "BAD") # type: ignore[arg-type] -def test_not_in_search_filter_declares_empty_slots() -> None: - """NotInSearchFilter should not allocate an instance dict.""" - - assert NotInSearchFilter.__slots__ == () - - def test_not_in_collection_parameter_suffix_is_literal() -> None: """The not-in suffix should be the direct literal used in parameter names.""" diff --git a/tests/unit/core/test_type_conversion.py b/tests/unit/core/test_type_conversion.py index 0a5f69d36..ce7e17bf9 100644 --- a/tests/unit/core/test_type_conversion.py +++ b/tests/unit/core/test_type_conversion.py @@ -1,6 +1,6 @@ """Tests for centralized scalar conversion helpers.""" -from datetime import UTC, date, datetime, time +from datetime import UTC, date, datetime, time, timedelta, timezone from decimal import Decimal from uuid import UUID @@ -114,20 +114,24 @@ def test_parse_datetime_rfc3339_z_suffix() -> None: assert parsed.year == 2023 -def test_parse_datetime_rfc3339_delegates_to_iso_converter(monkeypatch: pytest.MonkeyPatch) -> None: - """RFC3339 parsing should reuse the ISO datetime converter.""" - - import sqlspec.core.type_converter as type_converter - - expected = datetime(2023, 12, 25, 10, 30, 0, tzinfo=UTC) - - def fake_convert(value: str) -> datetime: - assert value == "2023-12-25T10:30:00+00:00" - return expected - - monkeypatch.setattr(type_converter, "convert_iso_datetime", fake_convert) +@pytest.mark.parametrize( + ("value", "expected"), + [ + ("2023-12-25T10:30:00Z", datetime(2023, 12, 25, 10, 30, tzinfo=UTC)), + ( + "2023-12-25T10:30:00.123456+05:30", + datetime(2023, 12, 25, 10, 30, 0, 123456, tzinfo=timezone(timedelta(hours=5, minutes=30))), + ), + ("2023-12-25T10:30:00-04:00", datetime(2023, 12, 25, 10, 30, tzinfo=timezone(timedelta(hours=-4)))), + ], +) +def test_parse_datetime_rfc3339_preserves_timezone_and_precision(value: str, expected: datetime) -> None: + """RFC3339 parsing preserves the represented instant, offset, and microseconds.""" + parsed = parse_datetime_rfc3339(value) - assert type_converter.parse_datetime_rfc3339("2023-12-25T10:30:00Z") is expected + assert parsed == expected + assert parsed.utcoffset() == expected.utcoffset() + assert parsed.microsecond == expected.microsecond def test_datetime_round_trip() -> None: diff --git a/tests/unit/dialects/test_generators.py b/tests/unit/dialects/test_generators.py index 21ed2ad28..54a3b4204 100644 --- a/tests/unit/dialects/test_generators.py +++ b/tests/unit/dialects/test_generators.py @@ -4,8 +4,6 @@ produces correct SQL output for all extension operators and Spanner DDL. """ -from pathlib import Path - from sqlglot import parse_one from sqlspec.dialects.spanner._generators import _get_dialect_name @@ -95,12 +93,3 @@ def test_spangres_normalizes_row_deletion_policy_to_ttl() -> None: sql = "\n CREATE TABLE events (\n id VARCHAR PRIMARY KEY,\n created_at TIMESTAMP\n ) ROW DELETION POLICY (OLDER_THAN(created_at, INTERVAL 30 DAY))\n " rendered = parse_one(sql, dialect="spangres").sql(dialect="spangres") assert "TTL INTERVAL '30 days' ON created_at" in rendered - - -def test_compat_removed_dialects_compat_module_removed() -> None: - assert not Path("sqlspec/dialects/_compat.py").exists() - - -def test_compat_removed_dialects_compat_not_in_mypyc_include() -> None: - assert "sqlspec/dialects/_compat.py" not in Path("pyproject.toml").read_text() - assert "sqlspec/dialects/_compat.py" not in Path("Makefile").read_text() diff --git a/tests/unit/dialects/test_registration.py b/tests/unit/dialects/test_registration.py index dc97e86ed..d037f594b 100644 --- a/tests/unit/dialects/test_registration.py +++ b/tests/unit/dialects/test_registration.py @@ -4,6 +4,10 @@ import sys from importlib.metadata import entry_points +import pytest + +from tests.conftest import is_compiled + DIALECT_ENTRY_POINTS = { "db2": "sqlspec.dialects.db2", "paradedb": "sqlspec.dialects.postgres", @@ -36,8 +40,9 @@ def test_sqlglot_resolves_dialects_without_sqlspec_dialect_import() -> None: assert "<=>" in result.stdout +@pytest.mark.skipif(is_compiled(), reason="Compiled root facade eagerly binds dialect exports for mypyc compatibility") def test_importing_sqlspec_does_not_eagerly_load_dialect_machinery() -> None: - """``import sqlspec`` must not pay the sqlglot dialect registration cost.""" + """The interpreted root facade defers dialect registration until first use.""" code = ( "import sys\n" "import sqlspec\n" diff --git a/tests/unit/dialects/test_spanner.py b/tests/unit/dialects/test_spanner.py index 38047436c..f713aa1cc 100644 --- a/tests/unit/dialects/test_spanner.py +++ b/tests/unit/dialects/test_spanner.py @@ -5,8 +5,6 @@ clauses, and DAY-based row deletion intervals. """ -from pathlib import Path - import pytest from sqlglot import Dialect, exp, parse_one @@ -112,8 +110,6 @@ def test_interleave_create_repairs_command_fallback() -> None: expression = parse_one(OFFICIAL_INTERLEAVE_DDL, dialect="spanner") assert isinstance(expression, exp.Create) assert "INTERLEAVE IN PARENT Singers ON DELETE CASCADE" in expression.sql(dialect="spanner") - assert not hasattr(_parsers, "_INTERLEAVE_PATTERN") - assert not hasattr(_parsers, "extract_interleave_property") ordered_pk_ddl = ( "CREATE TABLE Albums (SingerId INT64 NOT NULL, AlbumId INT64 NOT NULL) " @@ -200,19 +196,6 @@ def test_spanner_property_moves_after_non_spanner_property() -> None: assert rendered.index("enable_change_stream_capture") < rendered.index("INTERLEAVE IN PARENT") -def test_dedup_normalize_interval_expression_has_single_definition() -> None: - package_root = Path("sqlspec/dialects/spanner") - matches = [path for path in package_root.glob("_*.py") if "def _normalize_interval_expression" in path.read_text()] - assert matches == [package_root / "_generators.py"] - - -def test_dedup_spanner_constants_have_single_definition() -> None: - package_root = Path("sqlspec/dialects/spanner") - for name in ("_TTL_MIN_COMPONENTS", "_ROW_DELETION_NAME", "_INTERLEAVE_NAME", "_INTERLEAVE_IN_NAME"): - matches = [path for path in package_root.glob("_*.py") if f"{name} =" in path.read_text()] - assert matches == [package_root / "_generators.py"] - - def test_dedup_parsers_import_canonical_helpers() -> None: assert _parsers._normalize_interval_expression is _generators._normalize_interval_expression assert _parsers._ROW_DELETION_NAME is _generators._ROW_DELETION_NAME From c9da2c828dc27db090e5af5c3e4cfdc9f005b1aa Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 23:05:55 +0000 Subject: [PATCH 159/174] test(adbc): exercise external dependencies and configured behavior --- .../test_adbc_config_normalization.py | 84 ++++++++++++------- tests/unit/adapters/test_adbc/test_core.py | 70 +++++++--------- .../test_adbc/test_data_dictionary.py | 28 +++---- 3 files changed, 95 insertions(+), 87 deletions(-) diff --git a/tests/unit/adapters/test_adbc/test_adbc_config_normalization.py b/tests/unit/adapters/test_adbc/test_adbc_config_normalization.py index f39b91e36..e543bb13f 100644 --- a/tests/unit/adapters/test_adbc/test_adbc_config_normalization.py +++ b/tests/unit/adapters/test_adbc/test_adbc_config_normalization.py @@ -4,11 +4,15 @@ import importlib.metadata import json import os +import sys import threading import weakref +from collections.abc import Callable, Iterator from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass, field from importlib import import_module +from importlib.machinery import ModuleSpec +from types import ModuleType from typing import Any, cast, get_type_hints from unittest.mock import MagicMock @@ -30,6 +34,25 @@ from sqlspec.adapters.adbc.type_converter import get_adbc_type_converter from sqlspec.core import SQL, ParameterStyle from sqlspec.exceptions import ImproperConfigurationError, SQLSpecError +from sqlspec.utils.module_loader import reset_dependency_cache + + +@pytest.fixture +def optional_driver_imports(monkeypatch: pytest.MonkeyPatch) -> Iterator[Callable[[str, bool], None]]: + """Control optional SDK package imports through Python's module registry.""" + configured: set[str] = set() + + def configure(package: str, installed: bool) -> None: + module = ModuleType(package) if installed else None + if module is not None: + module.__spec__ = ModuleSpec(package, loader=None, is_package=True) + monkeypatch.setitem(sys.modules, package, module) + reset_dependency_cache(package) + configured.add(package) + + yield configure + for package in configured: + reset_dependency_cache(package) def test_adbc_config_runs_connection_create_callback(monkeypatch: Any) -> None: @@ -587,16 +610,16 @@ def test_resolve_driver_name_from_uri() -> None: assert _resolve_driver_name(config) == "adbc_driver_postgresql.dbapi.connect" -def test_resolve_driver_name_gizmosql_alias(monkeypatch: Any) -> None: +def test_resolve_driver_name_gizmosql_alias(optional_driver_imports: Callable[[str, bool], None]) -> None: """Resolve GizmoSQL aliases to the FlightSQL driver when native gizmosql driver is absent.""" - monkeypatch.setattr("sqlspec.adapters.adbc.core.import_optional", lambda _name: None) + optional_driver_imports("adbc_driver_gizmosql", False) config = AdbcConfig(connection_config={"driver_name": "gizmosql"}) assert _resolve_driver_name(config) == "adbc_driver_flightsql.dbapi.connect" -def test_resolve_driver_name_from_gizmosql_uri(monkeypatch: Any) -> None: +def test_resolve_driver_name_from_gizmosql_uri(optional_driver_imports: Callable[[str, bool], None]) -> None: """Detect FlightSQL driver from GizmoSQL URI schemes when native gizmosql driver is absent.""" - monkeypatch.setattr("sqlspec.adapters.adbc.core.import_optional", lambda _name: None) + optional_driver_imports("adbc_driver_gizmosql", False) config = AdbcConfig(connection_config={"uri": "gizmosql://localhost:31337"}) assert _resolve_driver_name(config) == "adbc_driver_flightsql.dbapi.connect" config = AdbcConfig(connection_config={"uri": "gizmo://localhost:31337"}) @@ -667,9 +690,9 @@ def test_connection_config_dict_preserves_db_kwargs_for_non_bigquery() -> None: assert "foo" not in resolved -def test_flightsql_lifts_username_password_into_db_kwargs(monkeypatch: Any) -> None: +def test_flightsql_lifts_username_password_into_db_kwargs(optional_driver_imports: Callable[[str, bool], None]) -> None: """Top-level username/password move into db_kwargs on the FlightSQL fallback route.""" - monkeypatch.setattr("sqlspec.adapters.adbc.core.import_optional", lambda _name: None) + optional_driver_imports("adbc_driver_gizmosql", False) config = AdbcConfig(connection_config={"driver_name": "gizmosql", "username": "alice", "password": "secret"}) resolved = _get_connection_config_dict(config) @@ -678,36 +701,38 @@ def test_flightsql_lifts_username_password_into_db_kwargs(monkeypatch: Any) -> N assert "password" not in resolved -def test_flightsql_lifts_tls_skip_verify_true(monkeypatch: Any) -> None: +def test_flightsql_lifts_tls_skip_verify_true(optional_driver_imports: Callable[[str, bool], None]) -> None: """tls_skip_verify=True maps to the FlightSQL ADBC option key on the FlightSQL fallback route.""" - monkeypatch.setattr("sqlspec.adapters.adbc.core.import_optional", lambda _name: None) + optional_driver_imports("adbc_driver_gizmosql", False) config = AdbcConfig(connection_config={"driver_name": "gizmosql", "tls_skip_verify": True}) resolved = _get_connection_config_dict(config) assert resolved["db_kwargs"]["adbc.flight.sql.client_option.tls_skip_verify"] == "true" -def test_flightsql_lifts_tls_skip_verify_false(monkeypatch: Any) -> None: +def test_flightsql_lifts_tls_skip_verify_false(optional_driver_imports: Callable[[str, bool], None]) -> None: """tls_skip_verify=False maps to the FlightSQL ADBC option key on the FlightSQL fallback route.""" - monkeypatch.setattr("sqlspec.adapters.adbc.core.import_optional", lambda _name: None) + optional_driver_imports("adbc_driver_gizmosql", False) config = AdbcConfig(connection_config={"driver_name": "gizmosql", "tls_skip_verify": False}) resolved = _get_connection_config_dict(config) assert resolved["db_kwargs"]["adbc.flight.sql.client_option.tls_skip_verify"] == "false" -def test_flightsql_lifts_authorization_header(monkeypatch: Any) -> None: +def test_flightsql_lifts_authorization_header(optional_driver_imports: Callable[[str, bool], None]) -> None: """authorization_header maps to the FlightSQL ADBC option key.""" - monkeypatch.setattr("sqlspec.adapters.adbc.core.import_optional", lambda _name: None) + optional_driver_imports("adbc_driver_gizmosql", False) config = AdbcConfig(connection_config={"driver_name": "gizmosql", "authorization_header": "Bearer token"}) resolved = _get_connection_config_dict(config) assert resolved["db_kwargs"]["adbc.flight.sql.authorization_header"] == "Bearer token" -def test_flightsql_explicit_db_kwargs_wins_over_top_level_shortcuts(monkeypatch: Any) -> None: +def test_flightsql_explicit_db_kwargs_wins_over_top_level_shortcuts( + optional_driver_imports: Callable[[str, bool], None], +) -> None: """User-supplied db_kwargs entries take precedence over top-level shortcuts.""" - monkeypatch.setattr("sqlspec.adapters.adbc.core.import_optional", lambda _name: None) + optional_driver_imports("adbc_driver_gizmosql", False) config = AdbcConfig( connection_config={ "driver_name": "gizmosql", @@ -800,9 +825,9 @@ def test_gizmosql_parameter_style_is_qmark() -> None: assert config.statement_config.parameter_config.default_parameter_style == ParameterStyle.QMARK -def test_gizmo_alias_resolves_to_flightsql(monkeypatch: Any) -> None: +def test_gizmo_alias_resolves_to_flightsql(optional_driver_imports: Callable[[str, bool], None]) -> None: """The short 'gizmo' alias resolves to FlightSQL driver when native gizmosql driver is absent.""" - monkeypatch.setattr("sqlspec.adapters.adbc.core.import_optional", lambda _name: None) + optional_driver_imports("adbc_driver_gizmosql", False) config = AdbcConfig(connection_config={"driver_name": "gizmo"}) assert _resolve_driver_name(config) == "adbc_driver_flightsql.dbapi.connect" assert config.statement_config.dialect == "duckdb" @@ -1088,11 +1113,11 @@ def test_db2_driver_aliases_uri_and_parameter_profile() -> None: } -def test_mssql_driver_manager_fallback_and_tsql_statement_config(monkeypatch: Any) -> None: +def test_mssql_driver_manager_fallback_and_tsql_statement_config( + optional_driver_imports: Callable[[str, bool], None], +) -> None: """When adbc_driver_mssql is absent, mssql routes through adbc_driver_manager with driver='mssql'.""" - monkeypatch.setattr( - "sqlspec.adapters.adbc.core.import_optional", lambda name: None if name == "adbc_driver_mssql" else MagicMock() - ) + optional_driver_imports("adbc_driver_mssql", False) for conn_cfg in ( {"driver_name": "mssql", "uri": "sqlserver://sa:Pass@localhost:1433?database=master"}, @@ -1113,22 +1138,17 @@ def test_mssql_driver_manager_fallback_and_tsql_statement_config(monkeypatch: An assert resolved["db_kwargs"]["uri"].startswith(("sqlserver://", "mssql://")) -def test_mssql_native_package_preferred_when_installed(monkeypatch: Any) -> None: +def test_mssql_native_package_preferred_when_installed(optional_driver_imports: Callable[[str, bool], None]) -> None: """When adbc_driver_mssql is installed, mssql resolves to adbc_driver_mssql.dbapi.connect.""" - monkeypatch.setattr( - "sqlspec.adapters.adbc.core.import_optional", lambda name: MagicMock() if name == "adbc_driver_mssql" else None - ) + optional_driver_imports("adbc_driver_mssql", True) config = AdbcConfig(connection_config={"driver_name": "mssql", "uri": "sqlserver://localhost:1433"}) assert _resolve_driver_name(config) == "adbc_driver_mssql.dbapi.connect" assert config.statement_config.dialect == "tsql" -def test_gizmosql_prefers_native_driver_when_installed(monkeypatch: Any) -> None: +def test_gizmosql_prefers_native_driver_when_installed(optional_driver_imports: Callable[[str, bool], None]) -> None: """When adbc_driver_gizmosql is installed, gizmosql resolves to adbc_driver_gizmosql.dbapi.connect.""" - monkeypatch.setattr( - "sqlspec.adapters.adbc.core.import_optional", - lambda name: MagicMock() if name == "adbc_driver_gizmosql" else None, - ) + optional_driver_imports("adbc_driver_gizmosql", True) config = AdbcConfig( connection_config={ "driver_name": "gizmosql", @@ -1149,9 +1169,11 @@ def test_gizmosql_prefers_native_driver_when_installed(monkeypatch: Any) -> None assert _get_connection_config_dict(explicit_autocommit)["autocommit"] is True -def test_gizmosql_flightsql_fallback_lifts_auth_type_into_db_kwargs(monkeypatch: Any) -> None: +def test_gizmosql_flightsql_fallback_lifts_auth_type_into_db_kwargs( + optional_driver_imports: Callable[[str, bool], None], +) -> None: """When adbc_driver_gizmosql is absent, auth_type is lifted into db_kwargs['adbc.gizmosql.auth_type'].""" - monkeypatch.setattr("sqlspec.adapters.adbc.core.import_optional", lambda _name: None) + optional_driver_imports("adbc_driver_gizmosql", False) config = AdbcConfig( connection_config={"driver_name": "gizmosql", "uri": "grpc+tls://localhost:31337", "auth_type": "external"} ) diff --git a/tests/unit/adapters/test_adbc/test_core.py b/tests/unit/adapters/test_adbc/test_core.py index 860096ddd..0fcd4d518 100644 --- a/tests/unit/adapters/test_adbc/test_core.py +++ b/tests/unit/adapters/test_adbc/test_core.py @@ -1,5 +1,6 @@ """Unit tests for ADBC core execute-many helpers.""" +import json from collections.abc import Sequence from decimal import Decimal from pathlib import Path @@ -163,56 +164,40 @@ def test_collect_rows_uses_precomputed_column_names() -> None: assert column_names == ["id", "name"] -def test_prepare_parameters_with_casts_uses_type_converter_factory(monkeypatch: Any) -> None: +def test_prepare_parameters_with_casts_uses_custom_json_serializer() -> None: statement_config = get_statement_config("postgres") - factory_calls: list[tuple[str, int]] = [] - - class FakeConverter: - def convert_dict(self, value: dict[str, Any]) -> str: - return f"factory:{sorted(value.items())!r}" - - def fake_factory(dialect: str, cache_size: int = 5000) -> FakeConverter: - factory_calls.append((dialect, cache_size)) - return FakeConverter() - - monkeypatch.setattr(adbc_core, "get_adbc_type_converter", fake_factory) - - prepared = prepare_parameters_with_casts( - [{"id": 1}], {}, statement_config, dialect="postgres", json_serializer=lambda value: str(value) + statement_config = statement_config.replace( + parameter_config=statement_config.parameter_config.replace( + json_serializer=lambda value: json.dumps({"custom": value}, separators=(",", ":")) + ) ) - assert prepared == ["factory:[('id', 1)]"] - assert factory_calls == [("postgres", 5000)] + prepared = prepare_parameters_with_casts([{"id": 1}], {}, statement_config, dialect="postgres", json_serializer=str) + assert prepared == ['{"custom":{"id":1}}'] -def test_prepare_batch_with_casts_creates_converter_once(monkeypatch: Any) -> None: - statement_config = get_statement_config("postgres") - factory_calls: list[tuple[str, int]] = [] - - class FakeConverter: - def convert_dict(self, value: dict[str, Any]) -> str: - return f"factory:{sorted(value.items())!r}" - def fake_factory(dialect: str, cache_size: int = 5000) -> FakeConverter: - factory_calls.append((dialect, cache_size)) - return FakeConverter() - - monkeypatch.setattr(adbc_core, "get_adbc_type_converter", fake_factory) +def test_prepare_batch_with_casts_uses_custom_json_serializer_for_each_row() -> None: + statement_config = get_statement_config("postgres") + statement_config = statement_config.replace( + parameter_config=statement_config.parameter_config.replace( + json_serializer=lambda value: json.dumps({"custom": value}, separators=(",", ":")) + ) + ) prepared = prepare_batch_with_casts( [(1, {"name": "alice"}), (2, {"name": "bob"}), (3, {"name": "carol"})], {1: "INT"}, statement_config, dialect="postgres", - json_serializer=lambda value: str(value), + json_serializer=str, ) assert prepared == [ - (1, "factory:[('name', 'alice')]"), - (2, "factory:[('name', 'bob')]"), - (3, "factory:[('name', 'carol')]"), + (1, '{"custom":{"name":"alice"}}'), + (2, '{"custom":{"name":"bob"}}'), + (3, '{"custom":{"name":"carol"}}'), ] - assert factory_calls == [("postgres", 5000)] def test_prepare_parameters_with_casts_supports_subclass_type_dispatch() -> None: @@ -908,15 +893,16 @@ def test_rewrite_postgres_uuid_placeholders() -> None: assert arrays == (2,) -def test_rewrite_uuid_cache() -> None: - """_rewrite_postgres_uuid_placeholders caches results in _REWRITE_UUID_CACHE.""" +def test_repeated_uuid_bindings_preserve_sql_and_current_parameter_values() -> None: + """Repeated statement rewrites retain each call's current UUID binding.""" sql = "SELECT $1 WHERE id = $2" - key = (sql, (1,), (), "postgres") - adbc_core._REWRITE_UUID_CACHE.pop(key, None) - res1 = adbc_core._rewrite_postgres_uuid_placeholders(sql, (1,), (), "postgres") - assert key in adbc_core._REWRITE_UUID_CACHE - res2 = adbc_core._rewrite_postgres_uuid_placeholders(sql, (1,), (), "postgres") - assert res1 is res2 + first = UUID("12345678-1234-5678-1234-567812345678") + second = UUID("87654321-4321-8765-4321-876543218765") + first_sql, first_parameters = prepare_postgres_uuid_bindings(sql, [first, 1], is_many=False, dialect="postgres") + second_sql, second_parameters = prepare_postgres_uuid_bindings(sql, [second, 2], is_many=False, dialect="postgres") + assert first_sql == second_sql == "SELECT CAST($1 AS UUID) WHERE id = $2" + assert first_parameters == [str(first), 1] + assert second_parameters == [str(second), 2] def test_strict_coercion_types() -> None: diff --git a/tests/unit/adapters/test_adbc/test_data_dictionary.py b/tests/unit/adapters/test_adbc/test_data_dictionary.py index c989e34be..55b7468ae 100644 --- a/tests/unit/adapters/test_adbc/test_data_dictionary.py +++ b/tests/unit/adapters/test_adbc/test_data_dictionary.py @@ -10,6 +10,7 @@ from adbc_driver_manager import NotSupportedError as AdbcNotSupportedError +from sqlspec.adapters.adbc import AdbcConfig from sqlspec.adapters.adbc.data_dictionary import ( AdbcDataDictionary, _arrow_type_to_sql, @@ -502,20 +503,19 @@ def test_adbc_get_objects_unique_check_constraints_are_lossy() -> None: assert all(item.attributes["is_lossy"] is True for item in constraint_items) -def test_adbc_ddl_is_unsupported_without_dialect_pack() -> None: - """ADBC transport metadata should not claim lossless DDL.""" - driver = Mock() - driver.dialect = "duckdb" - - result = AdbcDataDictionary().get_ddl( - driver, - "items", - schema="main", - object_type="view", - include_dependencies=False, - prefer_native=False, - redact=False, - ) +def test_adbc_view_ddl_is_explicitly_unsupported() -> None: + """ADBC transport metadata should not claim lossless view DDL.""" + config = AdbcConfig(connection_config={"driver_name": "sqlite"}) + with config.provide_session() as driver: + result = driver.data_dictionary.get_ddl( + driver, + "items", + schema="main", + object_type="view", + include_dependencies=False, + prefer_native=False, + redact=False, + ) assert result.identity.name == "items" assert result.identity.object_type == "view" From 3a44d282f6ade2a6f55db8862dc6c6e644db43eb Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 23:05:55 +0000 Subject: [PATCH 160/174] test(pools): simulate health checks through external SDKs --- .../unit/adapters/test_db2/test_async_pool.py | 45 ++++++++++--------- tests/unit/adapters/test_db2/test_pool.py | 21 +++++---- tests/unit/adapters/test_pymssql/_fakes.py | 9 +++- tests/unit/adapters/test_pymssql/test_pool.py | 30 +++++-------- 4 files changed, 57 insertions(+), 48 deletions(-) diff --git a/tests/unit/adapters/test_db2/test_async_pool.py b/tests/unit/adapters/test_db2/test_async_pool.py index 4cdde9b9d..4101d7ecc 100644 --- a/tests/unit/adapters/test_db2/test_async_pool.py +++ b/tests/unit/adapters/test_db2/test_async_pool.py @@ -7,11 +7,11 @@ import pytest -import sqlspec.adapters.db2.pool as pool_module from sqlspec.adapters.db2.pool import Db2AsyncConnectionPool from sqlspec.exceptions import ConnectionTimeoutError, DatabaseConnectionError from tests.unit.adapters.test_db2._fakes import ( FakeDb2AsyncConnection, + FakeDb2AsyncCursor, FakeDb2Connection, FakeDb2Cursor, FakeDb2OperationalError, @@ -26,7 +26,7 @@ class FakeClock: - """Replacement for the pool module's ``time`` with a controllable monotonic clock.""" + """Controllable external monotonic clock.""" def __init__(self) -> None: self.now = 1000.0 @@ -34,9 +34,6 @@ def __init__(self) -> None: def monotonic(self) -> float: return self.now - def time(self) -> float: - return time.time() - def advance(self, seconds: float) -> None: self.now += seconds @@ -205,7 +202,7 @@ async def test_recycle_and_health_check(fake_ibm_db: FakeModules, monkeypatch: p """Old connections are recycled, idle ones are pinged, and failed pings replace the connection.""" _, module = fake_ibm_db clock = FakeClock() - monkeypatch.setattr(pool_module, "time", clock) + monkeypatch.setattr(time, "monotonic", clock.monotonic) recycled = FakeDb2Connection() unhealthy = FakeDb2Connection([ FakeDb2Cursor(error=db2_error(-30108, "08506", "A connection failed", FakeDb2OperationalError)) @@ -330,27 +327,35 @@ async def test_async_connection_create_hook_is_awaited( async def test_cancelled_health_check_closes_detached_connection( fake_ibm_db: FakeModules, monkeypatch: pytest.MonkeyPatch ) -> None: - pool = Db2AsyncConnectionPool({"database": "TESTDB"}, max_size=1) + pool = Db2AsyncConnectionPool({"database": "TESTDB"}, max_size=1, health_check_interval=-1.0) connection = _as_fake(await pool.acquire()) await pool.release(connection) checking = asyncio.Event() + finish = asyncio.Event() + execute = FakeDb2AsyncCursor.execute - async def check(_pool: object, _record: object) -> bool: + async def blocked_execute(cursor: FakeDb2AsyncCursor, operation: str, parameters: object = None) -> FakeDb2Cursor: checking.set() - await asyncio.Event().wait() - return True + await finish.wait() + return await execute(cursor, operation, parameters) - monkeypatch.setattr(Db2AsyncConnectionPool, "_is_reusable", check) + monkeypatch.setattr(FakeDb2AsyncCursor, "execute", blocked_execute) task = asyncio.create_task(pool.acquire()) - await checking.wait() - task.cancel() - with pytest.raises(asyncio.CancelledError): - await task - assert connection.sync_connection.closed is True - assert pool.size() == 0 - replacement = await pool.acquire() - await pool.release(replacement) - await pool.close() + try: + await asyncio.wait_for(checking.wait(), timeout=1.0) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, timeout=1.0) + assert connection.sync_connection.closed is True + assert connection.sync_connection.cursors[0].closed is True + assert pool.size() == 0 + replacement = await asyncio.wait_for(pool.acquire(), timeout=1.0) + await pool.release(replacement) + finally: + finish.set() + task.cancel() + await asyncio.gather(task, return_exceptions=True) + await pool.close() async def test_pool_closed_while_opening_rejects_and_closes_connection(fake_ibm_db: FakeModules) -> None: diff --git a/tests/unit/adapters/test_db2/test_pool.py b/tests/unit/adapters/test_db2/test_pool.py index 635e5c101..e3c4d048c 100644 --- a/tests/unit/adapters/test_db2/test_pool.py +++ b/tests/unit/adapters/test_db2/test_pool.py @@ -65,19 +65,19 @@ def test_pool_renders_dsn_and_passes_empty_connect_arguments(fake_ibm_db: FakeMo ) -def test_pool_recycles_failed_health_check(fake_ibm_db: FakeModules, monkeypatch: pytest.MonkeyPatch) -> None: +def test_pool_recycles_failed_health_check(fake_ibm_db: FakeModules) -> None: """Failed health check should close and replace thread connection.""" _, fake_module = fake_ibm_db - first = FakeDb2Connection() + first = FakeDb2Connection([FakeDb2Cursor(error=RuntimeError("connection lost"))]) second = FakeDb2Connection() fake_module.pending_connections.extend([first, second]) - monkeypatch.setattr(Db2SyncConnectionPool, "_is_connection_alive", lambda *_: False) pool = Db2SyncConnectionPool({"database": "TESTDB"}, health_check_interval=-1.0) assert cast(object, pool.acquire()) is first assert cast(object, pool.acquire()) is second assert first.closed is True + assert first.cursors[0].closed is True assert len(fake_module.connect_calls) == 2 @@ -140,21 +140,24 @@ def _open() -> None: assert [conn.closed for conn in opened] == [True, True] -def test_pool_registry_does_not_grow_across_replacements( - fake_ibm_db: FakeModules, monkeypatch: pytest.MonkeyPatch -) -> None: +def test_pool_registry_does_not_grow_across_replacements(fake_ibm_db: FakeModules) -> None: """Registry size remains bounded during connection replacement.""" - monkeypatch.setattr(Db2SyncConnectionPool, "_is_connection_alive", lambda *_: False) + _, module = fake_ibm_db + connections = [FakeDb2Connection([FakeDb2Cursor(error=RuntimeError("connection lost"))]) for _ in range(4)] + module.pending_connections.extend(connections) pool = Db2SyncConnectionPool({"database": "TESTDB"}, health_check_interval=-1.0) for _ in range(4): pool.acquire() - assert len(pool._connection_registry) == 1 + assert pool.size() == 1 + assert len(module.connect_calls) == 4 + assert [connection.closed for connection in connections] == [True, True, True, False] pool.close() - assert pool._connection_registry == set() + assert pool.size() == 0 + assert all(connection.closed for connection in connections) def test_pool_get_connection_context_manager(fake_ibm_db: FakeModules) -> None: diff --git a/tests/unit/adapters/test_pymssql/_fakes.py b/tests/unit/adapters/test_pymssql/_fakes.py index 8fcd4ec1d..b719ccc5c 100644 --- a/tests/unit/adapters/test_pymssql/_fakes.py +++ b/tests/unit/adapters/test_pymssql/_fakes.py @@ -19,11 +19,16 @@ class FakeCursor: """Minimal DB-API cursor for pymssql unit tests.""" def __init__( - self, rows: "list[Any] | None" = None, description: "list[tuple[str, ...]] | None" = None, rowcount: int = -1 + self, + rows: "list[Any] | None" = None, + description: "list[tuple[str, ...]] | None" = None, + rowcount: int = -1, + error: BaseException | None = None, ) -> None: self.rows = rows or [] self.description = description self.rowcount = rowcount + self.error = error self.closed = False self.calls: list[tuple[str, Any]] = [] self.fetchmany_sizes: list[int] = [] @@ -31,6 +36,8 @@ def __init__( def execute(self, sql: str, parameters: Any = None) -> None: self.calls.append((sql, parameters)) + if self.error is not None: + raise self.error def executemany(self, sql: str, parameters: Any = None) -> None: self.many_calls.append((sql, parameters)) diff --git a/tests/unit/adapters/test_pymssql/test_pool.py b/tests/unit/adapters/test_pymssql/test_pool.py index 564b0a7df..f9665b1b0 100644 --- a/tests/unit/adapters/test_pymssql/test_pool.py +++ b/tests/unit/adapters/test_pymssql/test_pool.py @@ -5,7 +5,7 @@ import pytest -from tests.unit.adapters.test_pymssql._fakes import FakeConnection, FakePymssqlModule +from tests.unit.adapters.test_pymssql._fakes import FakeConnection, FakeCursor, FakePymssqlModule def test_pool_connects_with_config_and_runs_hook(monkeypatch) -> None: @@ -38,7 +38,7 @@ def test_pool_recycles_failed_health_check(monkeypatch) -> None: import sqlspec.adapters.pymssql.pool as pool_module from sqlspec.adapters.pymssql.pool import PymssqlConnectionPool - first = FakeConnection() + first = FakeConnection(FakeCursor(error=OSError("connection lost"))) second = FakeConnection() connections = [first, second] @@ -49,7 +49,6 @@ def connect(self, **kwargs: Any) -> FakeConnection: fake_module = SequencedModule() monkeypatch.setattr(pool_module, "pymssql", fake_module) - monkeypatch.setattr(PymssqlConnectionPool, "_is_connection_alive", lambda *_: False) pool = PymssqlConnectionPool({"server": "sql.example.test"}, health_check_interval=0.5) @@ -57,6 +56,7 @@ def connect(self, **kwargs: Any) -> FakeConnection: pool._thread_local.last_used = 0.0 # pyright: ignore[reportPrivateUsage] assert cast(object, pool.acquire()) is second assert first.closed is True + assert first.cursor_obj.closed is True assert len(fake_module.connect_calls) == 2 @@ -65,23 +65,15 @@ def test_pool_zero_health_check_interval_disables_health_check(monkeypatch) -> N import sqlspec.adapters.pymssql.pool as pool_module from sqlspec.adapters.pymssql.pool import PymssqlConnectionPool - connection = FakeConnection() + connection = FakeConnection(FakeCursor(error=OSError("health checks should be disabled"))) fake_module = FakePymssqlModule(connection) monkeypatch.setattr(pool_module, "pymssql", fake_module) - health_checks = 0 - - def _fail_if_called(*_args: Any) -> bool: - nonlocal health_checks - health_checks += 1 - return False - - monkeypatch.setattr(PymssqlConnectionPool, "_is_connection_alive", _fail_if_called) pool = PymssqlConnectionPool({"server": "sql.example.test"}, health_check_interval=0) assert cast(object, pool.acquire()) is connection pool._thread_local.last_used = 0.0 # pyright: ignore[reportPrivateUsage] assert cast(object, pool.acquire()) is connection - assert health_checks == 0 + assert connection.cursor_obj.calls == [] assert len(fake_module.connect_calls) == 1 @@ -171,21 +163,23 @@ def test_pool_registry_does_not_grow_across_replacements(monkeypatch) -> None: class SequencedModule(FakePymssqlModule): def connect(self, **kwargs: Any) -> FakeConnection: self.connect_calls.append(kwargs) - return FakeConnection() + return FakeConnection(FakeCursor(error=OSError("connection lost"))) monkeypatch.setattr(pool_module, "pymssql", SequencedModule()) - monkeypatch.setattr(PymssqlConnectionPool, "_is_connection_alive", lambda *_: False) pool = PymssqlConnectionPool({"server": "sql.example.test"}, health_check_interval=0.5) + connections = [] for _ in range(4): - pool.acquire() + connections.append(cast("FakeConnection", pool.acquire())) pool._thread_local.last_used = 0.0 # pyright: ignore[reportPrivateUsage] - assert len(pool._connection_registry) == 1 # pyright: ignore[reportPrivateUsage] + assert pool.size() == 1 + assert [connection.closed for connection in connections] == [True, True, True, False] pool.close() - assert pool._connection_registry == set() # pyright: ignore[reportPrivateUsage] + assert pool.size() == 0 + assert all(connection.closed for connection in connections) def test_pool_reentrant_thread_checkout(monkeypatch: Any) -> None: From 1cc37ad929916b313b63bc2901a01f0703968ebd Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 23:06:35 +0000 Subject: [PATCH 161/174] refactor(types): use the public data dictionary facade across packages --- sqlspec/adapters/aiosqlite/data_dictionary.py | 2 +- sqlspec/adapters/arrow_odbc/data_dictionary.py | 2 +- sqlspec/adapters/mssql_python/data_dictionary.py | 2 +- sqlspec/adapters/oracledb/data_dictionary.py | 2 +- sqlspec/adapters/pymssql/data_dictionary.py | 2 +- sqlspec/builder/_ddl.py | 2 +- sqlspec/driver/_common.py | 2 +- sqlspec/migrations/squash.py | 2 +- sqlspec/utils/arrow_helpers.py | 2 +- 9 files changed, 9 insertions(+), 9 deletions(-) diff --git a/sqlspec/adapters/aiosqlite/data_dictionary.py b/sqlspec/adapters/aiosqlite/data_dictionary.py index d2f655ea5..54a6420e1 100644 --- a/sqlspec/adapters/aiosqlite/data_dictionary.py +++ b/sqlspec/adapters/aiosqlite/data_dictionary.py @@ -31,7 +31,7 @@ if TYPE_CHECKING: from sqlspec.adapters.aiosqlite.driver import AiosqliteDriver - from sqlspec.data_dictionary._types import DialectConfig + from sqlspec.data_dictionary import DialectConfig __all__ = ("AiosqliteDataDictionary",) diff --git a/sqlspec/adapters/arrow_odbc/data_dictionary.py b/sqlspec/adapters/arrow_odbc/data_dictionary.py index cbfdad155..cf16cee5c 100644 --- a/sqlspec/adapters/arrow_odbc/data_dictionary.py +++ b/sqlspec/adapters/arrow_odbc/data_dictionary.py @@ -32,7 +32,7 @@ from sqlspec.adapters.arrow_odbc.driver import ArrowOdbcDriver from sqlspec.core import SQL - from sqlspec.data_dictionary._types import DialectConfig + from sqlspec.data_dictionary import DialectConfig __all__ = ("ArrowOdbcDataDictionary",) diff --git a/sqlspec/adapters/mssql_python/data_dictionary.py b/sqlspec/adapters/mssql_python/data_dictionary.py index aa0d1ef48..ddfac4775 100644 --- a/sqlspec/adapters/mssql_python/data_dictionary.py +++ b/sqlspec/adapters/mssql_python/data_dictionary.py @@ -44,7 +44,7 @@ from sqlspec.adapters.mssql_python.driver import MssqlPythonDriver from sqlspec.core import SQL - from sqlspec.data_dictionary._types import DialectConfig, MetadataCapabilityProfile + from sqlspec.data_dictionary import DialectConfig, MetadataCapabilityProfile __all__ = ("MssqlPythonSyncDataDictionary", "MssqlVersionInfo") diff --git a/sqlspec/adapters/oracledb/data_dictionary.py b/sqlspec/adapters/oracledb/data_dictionary.py index ecaf98fc6..28e96af0e 100644 --- a/sqlspec/adapters/oracledb/data_dictionary.py +++ b/sqlspec/adapters/oracledb/data_dictionary.py @@ -55,7 +55,7 @@ from collections.abc import Sequence from sqlspec.adapters.oracledb.driver import OracleAsyncDriver, OracleSyncDriver - from sqlspec.data_dictionary._types import DialectConfig + from sqlspec.data_dictionary import DialectConfig __all__ = ( "JSONStorageType", diff --git a/sqlspec/adapters/pymssql/data_dictionary.py b/sqlspec/adapters/pymssql/data_dictionary.py index 2365dae0e..9e8ada99e 100644 --- a/sqlspec/adapters/pymssql/data_dictionary.py +++ b/sqlspec/adapters/pymssql/data_dictionary.py @@ -44,7 +44,7 @@ from sqlspec.adapters.pymssql.driver import PymssqlDriver from sqlspec.core import SQL - from sqlspec.data_dictionary._types import DialectConfig, MetadataCapabilityProfile + from sqlspec.data_dictionary import DialectConfig, MetadataCapabilityProfile __all__ = ("MssqlVersionInfo", "PymssqlSyncDataDictionary") diff --git a/sqlspec/builder/_ddl.py b/sqlspec/builder/_ddl.py index 18a1f5c43..d8ac921ba 100644 --- a/sqlspec/builder/_ddl.py +++ b/sqlspec/builder/_ddl.py @@ -28,7 +28,7 @@ parse_order_expression, ) from sqlspec.core import SQL -from sqlspec.data_dictionary._types import ( +from sqlspec.data_dictionary import ( DDLResult, ForeignKeyMetadata, MetadataFidelity, diff --git a/sqlspec/driver/_common.py b/sqlspec/driver/_common.py index 4acf51ccc..def4f0393 100644 --- a/sqlspec/driver/_common.py +++ b/sqlspec/driver/_common.py @@ -128,7 +128,7 @@ from sqlspec.core import ArrowResult, FilterTypeT, OperationType, StatementFilter from sqlspec.core.parameters._types import ConvertedParameters from sqlspec.core.stack import StatementStack - from sqlspec.data_dictionary._types import DialectConfig + from sqlspec.data_dictionary import DialectConfig from sqlspec.storage import ( AsyncStoragePipeline, StorageBridgeJob, diff --git a/sqlspec/migrations/squash.py b/sqlspec/migrations/squash.py index d7f87bca3..71d4b85ac 100644 --- a/sqlspec/migrations/squash.py +++ b/sqlspec/migrations/squash.py @@ -25,7 +25,7 @@ from sqlglot.errors import ErrorLevel, ParseError, TokenError, UnsupportedError from sqlglot.tokenizer_core import Token, TokenType -from sqlspec.data_dictionary._types import DDLResult +from sqlspec.data_dictionary import DDLResult from sqlspec.driver import SyncDriverAdapterBase from sqlspec.exceptions import SquashValidationError from sqlspec.migrations._backup import create_backup, remove_backup, restore_backup diff --git a/sqlspec/utils/arrow_helpers.py b/sqlspec/utils/arrow_helpers.py index 742bada9d..a27c08f67 100644 --- a/sqlspec/utils/arrow_helpers.py +++ b/sqlspec/utils/arrow_helpers.py @@ -28,7 +28,7 @@ if TYPE_CHECKING: from sqlspec.core.result import ArrowResult - from sqlspec.data_dictionary._types import DialectConfig + from sqlspec.data_dictionary import DialectConfig from sqlspec.storage import StorageTelemetry from sqlspec.typing import ArrowRecordBatch, ArrowRecordBatchReader, ArrowTable, PandasDataFrame, PolarsDataFrame From ec347e3233f2cb656d0111c80e3efdfd3f1fac49 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 23:19:26 +0000 Subject: [PATCH 162/174] fix(db2): restore exception state after async health checks --- sqlspec/adapters/db2/pool.py | 25 +++++++++++++------ .../unit/adapters/test_db2/test_async_pool.py | 11 +++++++- 2 files changed, 28 insertions(+), 8 deletions(-) diff --git a/sqlspec/adapters/db2/pool.py b/sqlspec/adapters/db2/pool.py index 76f5dddc8..bfa3a6ac0 100644 --- a/sqlspec/adapters/db2/pool.py +++ b/sqlspec/adapters/db2/pool.py @@ -543,17 +543,28 @@ async def _is_reusable(self, record: _Db2PooledConnection) -> bool: return True async def _is_connection_alive(self, connection: Any) -> bool: - """Ping the connection with ``SELECT 1 FROM SYSIBM.SYSDUMMY1``.""" + """Ping the connection with ``SELECT 1 FROM SYSIBM.SYSDUMMY1``. + + Close the cursor outside exception handling so compiled coroutine cleanup + does not retain the failed query as its caller's ambient exception. + """ + cursor: Any = None + pending_error: BaseException | None = None try: cursor = await connection.cursor() + await cursor.execute("SELECT 1 FROM SYSIBM.SYSDUMMY1") + await cursor.fetchone() + except BaseException as exc: + pending_error = exc + if cursor is not None: try: - await cursor.execute("SELECT 1 FROM SYSIBM.SYSDUMMY1") - await cursor.fetchone() - finally: await cursor.close() - except Exception: - return False - return True + except BaseException as exc: + if pending_error is None or isinstance(pending_error, Exception): + pending_error = exc + if pending_error is not None and not isinstance(pending_error, Exception): + raise pending_error + return pending_error is None async def _close_connection(self, connection: Any) -> None: """Close a connection, suppressing driver errors.""" diff --git a/tests/unit/adapters/test_db2/test_async_pool.py b/tests/unit/adapters/test_db2/test_async_pool.py index 4101d7ecc..6a99e76de 100644 --- a/tests/unit/adapters/test_db2/test_async_pool.py +++ b/tests/unit/adapters/test_db2/test_async_pool.py @@ -324,8 +324,9 @@ async def test_async_connection_create_hook_is_awaited( await pool.close() +@pytest.mark.parametrize("close_fails", [False, True]) async def test_cancelled_health_check_closes_detached_connection( - fake_ibm_db: FakeModules, monkeypatch: pytest.MonkeyPatch + fake_ibm_db: FakeModules, close_fails: bool, monkeypatch: pytest.MonkeyPatch ) -> None: pool = Db2AsyncConnectionPool({"database": "TESTDB"}, max_size=1, health_check_interval=-1.0) connection = _as_fake(await pool.acquire()) @@ -333,13 +334,21 @@ async def test_cancelled_health_check_closes_detached_connection( checking = asyncio.Event() finish = asyncio.Event() execute = FakeDb2AsyncCursor.execute + close = FakeDb2AsyncCursor.close async def blocked_execute(cursor: FakeDb2AsyncCursor, operation: str, parameters: object = None) -> FakeDb2Cursor: checking.set() await finish.wait() return await execute(cursor, operation, parameters) + async def failing_close(cursor: FakeDb2AsyncCursor) -> None: + await close(cursor) + msg = "Cursor close failed during cancelled health check." + raise OSError(msg) + monkeypatch.setattr(FakeDb2AsyncCursor, "execute", blocked_execute) + if close_fails: + monkeypatch.setattr(FakeDb2AsyncCursor, "close", failing_close) task = asyncio.create_task(pool.acquire()) try: await asyncio.wait_for(checking.wait(), timeout=1.0) From ee5a96574e6a139c3bb2a6d8187b649b9da2dd95 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 1 Nov 2026 23:26:27 +0000 Subject: [PATCH 163/174] fix(driver): dispatch DDL hooks through interpreted overrides --- sqlspec/driver/_async.py | 24 ++++++----- sqlspec/driver/_common.py | 6 ++- sqlspec/driver/_sync.py | 24 ++++++----- .../test_adbc/test_data_dictionary.py | 43 ++++++++++++++++++- 4 files changed, 72 insertions(+), 25 deletions(-) diff --git a/sqlspec/driver/_async.py b/sqlspec/driver/_async.py index f7e8c0ab7..b581a7a5b 100644 --- a/sqlspec/driver/_async.py +++ b/sqlspec/driver/_async.py @@ -2263,7 +2263,7 @@ async def get_object_ddl( Returns: DDL result for the object. """ - resolved_name, resolved_schema = self.resolve_ddl_object_name(driver, object_name, schema) + resolved_name, resolved_schema = type(self).resolve_ddl_object_name(self, driver, object_name, schema) return await self._get_exact_object_ddl( driver, resolved_name, @@ -2377,11 +2377,13 @@ async def get_table_profile( Query failures (for example missing privileges) propagate to the caller. """ - resolved_table, resolved_schema = self.resolve_ddl_object_name(driver, table, schema) + resolved_table, resolved_schema = type(self).resolve_ddl_object_name(self, driver, table, schema) columns = await self.get_columns( driver, - table=self.exact_catalog_identifier(driver, resolved_table), - schema=self.exact_catalog_identifier(driver, resolved_schema) if resolved_schema is not None else None, + table=type(self).exact_catalog_identifier(self, driver, resolved_table), + schema=type(self).exact_catalog_identifier(self, driver, resolved_schema) + if resolved_schema is not None + else None, ) if not columns: return [] @@ -2409,7 +2411,7 @@ async def infer_table_ddl_from_data( profiles = await self.get_table_profile( driver, table, schema=schema, sample_size=sample_size, target_dialect=target_dialect ) - resolved_table, resolved_schema = self.resolve_ddl_object_name(driver, table, schema) + resolved_table, resolved_schema = type(self).resolve_ddl_object_name(self, driver, table, schema) if not profiles: return self._unsupported_ddl(target_table or resolved_table, schema=resolved_schema, object_type="table") return self._compile_profiled_table_ddl( @@ -2602,7 +2604,7 @@ async def _replace_unfaithful_native_tables( if not replace_names: return items catalog = await self._load_table_catalog(driver, schema, replace_names, warnings) - source_dialect = self.resolve_ddl_dialect(driver) + source_dialect = type(self).resolve_ddl_dialect(self, driver) resolved: list[DDLResult] = [] replaced_tables: set[tuple[str, str]] = set() for item in items: @@ -2651,7 +2653,7 @@ async def _collect_synthesized_schema_ddl( they were already loaded for every table, so callers can reuse them for ordering. """ resolved_default_schema = default_schema if default_schema is not None else schema - source_dialect = self.resolve_ddl_dialect(driver) + source_dialect = type(self).resolve_ddl_dialect(self, driver) excluded_names = self._normalize_excluded_table_names(exclude_tables) results: list[DDLResult] = [] loaded_foreign_keys: list[Any] | None = None @@ -2835,7 +2837,7 @@ async def _load_table_catalog( async def _load_table_rows(self, driver: Any, domain: str, table_name: str, schema: "str | None") -> "list[Any]": """Load one catalog domain for one table using its exact catalog name.""" - lookup_table = self.exact_catalog_identifier(driver, table_name) + lookup_table = type(self).exact_catalog_identifier(self, driver, table_name) if domain == "columns": return list(await self.get_columns(driver, table=lookup_table, schema=schema)) if domain == "foreign_keys": @@ -2857,8 +2859,8 @@ async def _synthesize_table_ddl( warnings: tuple[str, ...] = (), ) -> "DDLResult": """Synthesize table DDL from exact catalog names using columns, constraints, foreign keys, and indexes.""" - lookup_table = self.exact_catalog_identifier(driver, table_name) - lookup_schema = self.exact_catalog_identifier(driver, schema) if schema is not None else None + lookup_table = type(self).exact_catalog_identifier(self, driver, table_name) + lookup_schema = type(self).exact_catalog_identifier(self, driver, schema) if schema is not None else None columns = await self.get_columns(driver, table=lookup_table, schema=lookup_schema) if not columns: return self._unsupported_ddl(table_name, schema=schema, object_type="table", warnings=warnings) @@ -2877,7 +2879,7 @@ async def _synthesize_table_ddl( foreign_keys=foreign_keys, indexes=indexes, target_dialect=target_dialect, - source_dialect=self.resolve_ddl_dialect(driver), + source_dialect=type(self).resolve_ddl_dialect(self, driver), warnings=warnings, ) diff --git a/sqlspec/driver/_common.py b/sqlspec/driver/_common.py index def4f0393..3a6b89717 100644 --- a/sqlspec/driver/_common.py +++ b/sqlspec/driver/_common.py @@ -768,9 +768,11 @@ def _resolve_schema_ddl_input(self, driver: Any, schema: str | None) -> tuple[st return (None, None) if _ddl_dialect_or_none(type(self).dialect) is not None and isinstance(self, DataDictionaryDialectMixin): with contextlib.suppress(ValueError): - _, exact_schema = self.resolve_ddl_object_name(driver, "_", schema) + _, exact_schema = type(self).resolve_ddl_object_name(self, driver, "_", schema) getter_schema = ( - self.exact_catalog_identifier(driver, exact_schema) if exact_schema is not None else None + type(self).exact_catalog_identifier(self, driver, exact_schema) + if exact_schema is not None + else None ) return (exact_schema, getter_schema) stripped = schema.strip() diff --git a/sqlspec/driver/_sync.py b/sqlspec/driver/_sync.py index 731dba527..7643c73fe 100644 --- a/sqlspec/driver/_sync.py +++ b/sqlspec/driver/_sync.py @@ -2167,7 +2167,7 @@ def get_object_ddl( Returns: DDL result for the object. """ - resolved_name, resolved_schema = self.resolve_ddl_object_name(driver, object_name, schema) + resolved_name, resolved_schema = type(self).resolve_ddl_object_name(self, driver, object_name, schema) return self._get_exact_object_ddl( driver, resolved_name, @@ -2281,11 +2281,13 @@ def get_table_profile( Query failures (for example missing privileges) propagate to the caller. """ - resolved_table, resolved_schema = self.resolve_ddl_object_name(driver, table, schema) + resolved_table, resolved_schema = type(self).resolve_ddl_object_name(self, driver, table, schema) columns = self.get_columns( driver, - table=self.exact_catalog_identifier(driver, resolved_table), - schema=self.exact_catalog_identifier(driver, resolved_schema) if resolved_schema is not None else None, + table=type(self).exact_catalog_identifier(self, driver, resolved_table), + schema=type(self).exact_catalog_identifier(self, driver, resolved_schema) + if resolved_schema is not None + else None, ) if not columns: return [] @@ -2311,7 +2313,7 @@ def infer_table_ddl_from_data( profiles = self.get_table_profile( driver, table, schema=schema, sample_size=sample_size, target_dialect=target_dialect ) - resolved_table, resolved_schema = self.resolve_ddl_object_name(driver, table, schema) + resolved_table, resolved_schema = type(self).resolve_ddl_object_name(self, driver, table, schema) if not profiles: return self._unsupported_ddl(target_table or resolved_table, schema=resolved_schema, object_type="table") return self._compile_profiled_table_ddl( @@ -2504,7 +2506,7 @@ def _replace_unfaithful_native_tables( if not replace_names: return items catalog = self._load_table_catalog(driver, schema, replace_names, warnings) - source_dialect = self.resolve_ddl_dialect(driver) + source_dialect = type(self).resolve_ddl_dialect(self, driver) resolved: list[DDLResult] = [] replaced_tables: set[tuple[str, str]] = set() for item in items: @@ -2553,7 +2555,7 @@ def _collect_synthesized_schema_ddl( they were already loaded for every table, so callers can reuse them for ordering. """ resolved_default_schema = default_schema if default_schema is not None else schema - source_dialect = self.resolve_ddl_dialect(driver) + source_dialect = type(self).resolve_ddl_dialect(self, driver) excluded_names = self._normalize_excluded_table_names(exclude_tables) results: list[DDLResult] = [] loaded_foreign_keys: list[Any] | None = None @@ -2737,7 +2739,7 @@ def _load_table_catalog( def _load_table_rows(self, driver: Any, domain: str, table_name: str, schema: "str | None") -> "list[Any]": """Load one catalog domain for one table using its exact catalog name.""" - lookup_table = self.exact_catalog_identifier(driver, table_name) + lookup_table = type(self).exact_catalog_identifier(self, driver, table_name) if domain == "columns": return list(self.get_columns(driver, table=lookup_table, schema=schema)) if domain == "foreign_keys": @@ -2759,8 +2761,8 @@ def _synthesize_table_ddl( warnings: tuple[str, ...] = (), ) -> "DDLResult": """Synthesize table DDL from exact catalog names using columns, constraints, foreign keys, and indexes.""" - lookup_table = self.exact_catalog_identifier(driver, table_name) - lookup_schema = self.exact_catalog_identifier(driver, schema) if schema is not None else None + lookup_table = type(self).exact_catalog_identifier(self, driver, table_name) + lookup_schema = type(self).exact_catalog_identifier(self, driver, schema) if schema is not None else None columns = self.get_columns(driver, table=lookup_table, schema=lookup_schema) if not columns: return self._unsupported_ddl(table_name, schema=schema, object_type="table", warnings=warnings) @@ -2779,7 +2781,7 @@ def _synthesize_table_ddl( foreign_keys=foreign_keys, indexes=indexes, target_dialect=target_dialect, - source_dialect=self.resolve_ddl_dialect(driver), + source_dialect=type(self).resolve_ddl_dialect(self, driver), warnings=warnings, ) diff --git a/tests/unit/adapters/test_adbc/test_data_dictionary.py b/tests/unit/adapters/test_adbc/test_data_dictionary.py index 55b7468ae..dc2349b34 100644 --- a/tests/unit/adapters/test_adbc/test_data_dictionary.py +++ b/tests/unit/adapters/test_adbc/test_data_dictionary.py @@ -19,7 +19,7 @@ _normalize_native_statistics, _normalize_native_tables, ) -from sqlspec.data_dictionary import MetadataFidelity, MetadataSource, MetadataSupport +from sqlspec.data_dictionary import DDLResult, MetadataFidelity, MetadataSource, MetadataSupport from sqlspec.exceptions import OperationalError SQLITE_OBJECTS_PAYLOAD: list[dict[str, Any]] = [ @@ -527,6 +527,47 @@ def test_adbc_view_ddl_is_explicitly_unsupported() -> None: assert result.warnings +@pytest.mark.parametrize("target_dialect", [None, "postgres"]) +def test_adbc_table_and_schema_ddl_preserve_detected_dialect_and_quoted_names(target_dialect: str | None) -> None: + """DDL orchestration uses the active ADBC dialect for quoted catalog names and target rendering.""" + config = AdbcConfig(connection_config={"driver_name": "sqlite"}) + with config.provide_session() as driver: + driver.execute('CREATE TABLE "Mixed Items" ("Mixed ID" INTEGER PRIMARY KEY, "Label" TEXT)') + dictionary = driver.data_dictionary + table = dictionary.get_ddl( + driver, + '"Mixed Items"', + schema='"main"', + include_dependencies=False, + prefer_native=False, + target_dialect=target_dialect, + ) + schema = dictionary.get_schema_ddl( + driver, + schema='"main"', + include_domains=("tables",), + include_dependencies=False, + prefer_native=False, + target_dialect=target_dialect, + ) + dialect = target_dialect or "sqlite" + assert table.status == MetadataSupport.SUPPORTED + assert (table.identity.name, table.identity.schema, table.identity.dialect) == ("Mixed Items", "main", dialect) + assert isinstance(table.ddl, str) + assert '"Mixed Items"' in table.ddl + assert '"Mixed ID"' in table.ddl + assert len(schema.items) == 1 + item = schema.items[0] + assert isinstance(item, DDLResult) + assert (item.identity.name, item.identity.schema, item.identity.dialect, item.status) == ( + "Mixed Items", + "main", + dialect, + MetadataSupport.SUPPORTED, + ) + assert item.ddl == table.ddl + + def test_get_statistics_uses_native_statistic_names_when_available() -> None: """ADBC statistic-name metadata should replace built-in fallback labels.""" driver = Mock() From 1b6333ef30224da1a58f28d124cad2ce2775b262 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Mon, 2 Nov 2026 00:34:34 +0000 Subject: [PATCH 164/174] fix(postgres): restore client exports and search-path index metadata --- sqlspec/adapters/asyncpg/driver.py | 1 - sqlspec/adapters/psqlpy/driver.py | 3 +- .../dialects/postgres/sql/indexes.sql | 16 +++++++- .../integration/adapters/_shared/behaviors.py | 3 +- .../adapters/postgres/test_column_metadata.py | 22 +++++++++++ .../test_asyncpg/test_transaction_state.py | 37 +++++++++---------- .../unit/adapters/test_psqlpy/test_driver.py | 19 ++++++++++ 7 files changed, 75 insertions(+), 26 deletions(-) diff --git a/sqlspec/adapters/asyncpg/driver.py b/sqlspec/adapters/asyncpg/driver.py index 57bc78834..2a50b6f7a 100644 --- a/sqlspec/adapters/asyncpg/driver.py +++ b/sqlspec/adapters/asyncpg/driver.py @@ -380,7 +380,6 @@ async def select_to_storage( ``STDIN``/``STDOUT`` stream form. """ - self._require_capability("arrow_export_enabled") arrow_result = await self.select_to_arrow(statement, *parameters, statement_config=statement_config, **kwargs) async_pipeline = self._storage_pipeline() telemetry_payload = await self._write_storage_result( diff --git a/sqlspec/adapters/psqlpy/driver.py b/sqlspec/adapters/psqlpy/driver.py index b43bc3131..4a4429a96 100644 --- a/sqlspec/adapters/psqlpy/driver.py +++ b/sqlspec/adapters/psqlpy/driver.py @@ -354,9 +354,8 @@ async def select_to_storage( telemetry: "StorageTelemetry | None" = None, **kwargs: Any, ) -> "StorageBridgeJob": - """Execute a query and stream Arrow results to a storage backend.""" + """Execute a query and write client-converted Arrow results to storage.""" - self._require_capability("arrow_export_enabled") arrow_result = await self.select_to_arrow(statement, *parameters, statement_config=statement_config, **kwargs) async_pipeline = self._storage_pipeline() telemetry_payload = await self._write_storage_result( diff --git a/sqlspec/data_dictionary/dialects/postgres/sql/indexes.sql b/sqlspec/data_dictionary/dialects/postgres/sql/indexes.sql index 41006915e..fbd122100 100644 --- a/sqlspec/data_dictionary/dialects/postgres/sql/indexes.sql +++ b/sqlspec/data_dictionary/dialects/postgres/sql/indexes.sql @@ -1,5 +1,17 @@ -- name: by_schema -- dialect: postgres +WITH target_schema AS ( + SELECT COALESCE( + :schema_name::text, + ( + SELECT rn.nspname::text + FROM pg_catalog.pg_class rc + JOIN pg_catalog.pg_namespace rn ON rn.oid = rc.relnamespace + WHERE rc.oid = pg_catalog.to_regclass(pg_catalog.quote_ident(:table_name::text)) + ), + current_schema() + ) AS schema_name +) SELECT tn.nspname::text AS schema_name, tbl.relname::text AS table_name, @@ -35,11 +47,11 @@ SELECT FROM pg_catalog.pg_index ix JOIN pg_catalog.pg_class tbl ON tbl.oid = ix.indrelid JOIN pg_catalog.pg_namespace tn ON tn.oid = tbl.relnamespace +JOIN target_schema ts ON ts.schema_name = tn.nspname JOIN pg_catalog.pg_class idx ON idx.oid = ix.indexrelid JOIN pg_catalog.pg_am am ON am.oid = idx.relam LEFT JOIN pg_catalog.pg_depend dep ON dep.objid = idx.oid AND dep.deptype = 'e' -WHERE tn.nspname = :schema_name - AND (:table_name::text IS NULL OR tbl.relname = :table_name) +WHERE (:table_name::text IS NULL OR tbl.relname = :table_name) ORDER BY tn.nspname, tbl.relname, idx.relname; -- name: by_table diff --git a/tests/integration/adapters/_shared/behaviors.py b/tests/integration/adapters/_shared/behaviors.py index a77776663..bb912ae90 100644 --- a/tests/integration/adapters/_shared/behaviors.py +++ b/tests/integration/adapters/_shared/behaviors.py @@ -6951,7 +6951,8 @@ async def assert_async_load_from_records_contract(driver: object, case: DriverCa assert await async_driver.select_value(table.select_count_sql) == 5 with pytest.raises(ImproperConfigurationError): - await async_driver.load_from_records(table.name, [], overwrite=True) + await async_driver.load_from_records(table.name, []) + assert await async_driver.select_value(table.select_count_sql) == 5 def _storage_bridge_export_sql(table: ContractTable) -> str: diff --git a/tests/integration/adapters/postgres/test_column_metadata.py b/tests/integration/adapters/postgres/test_column_metadata.py index 1432ef58c..f5c534370 100644 --- a/tests/integration/adapters/postgres/test_column_metadata.py +++ b/tests/integration/adapters/postgres/test_column_metadata.py @@ -70,6 +70,28 @@ async def test_unqualified_columns_resolve_through_search_path(asyncpg_async_dri await driver.execute_script(f"RESET search_path; {TEARDOWN_SQL}") +async def test_unqualified_indexes_resolve_through_search_path(asyncpg_async_driver: "AsyncpgDriver") -> None: + """Index lookup resolves the table on search_path and honors an explicit schema.""" + driver = asyncpg_async_driver + await driver.execute_script(SETUP_SQL) + try: + await driver.execute_script( + 'CREATE INDEX "MixedLookup" ON column_metadata."MixedItems" ("userName"); ' + "SET search_path TO public, column_metadata" + ) + + indexes = await driver.data_dictionary.get_indexes(driver, table='"MixedItems"') + explicit_public = await driver.data_dictionary.get_indexes(driver, table='"MixedItems"', schema="public") + + lookup = next(index for index in indexes if index["index_name"] == "MixedLookup") + assert lookup["schema_name"] == "column_metadata" + assert lookup["table_name"] == "MixedItems" + assert lookup["columns"] == ["userName"] + assert explicit_public == [] + finally: + await driver.execute_script(f"RESET search_path; {TEARDOWN_SQL}") + + def test_columns_report_keys_identity_and_sequences_sync(psycopg_sync_config: "PsycopgSyncConfig") -> None: """The sync dictionary reports the same column metadata.""" with psycopg_sync_config.provide_session() as driver: diff --git a/tests/unit/adapters/test_asyncpg/test_transaction_state.py b/tests/unit/adapters/test_asyncpg/test_transaction_state.py index 9af6a2a5c..88fb20412 100644 --- a/tests/unit/adapters/test_asyncpg/test_transaction_state.py +++ b/tests/unit/adapters/test_asyncpg/test_transaction_state.py @@ -4,6 +4,7 @@ from asyncio import CancelledError from collections.abc import AsyncIterator from contextlib import asynccontextmanager +from pathlib import Path from typing import Any, cast from unittest.mock import AsyncMock, MagicMock @@ -21,6 +22,22 @@ pytestmark = pytest.mark.anyio +async def test_storage_bridge_writes_client_converted_rows_without_native_export(tmp_path: Path) -> None: + """Client-side storage export works without claiming native Arrow export.""" + parquet = pytest.importorskip("pyarrow.parquet") + connection = AsyncMock() + connection.fetch.return_value = [{"name": "Ada", "value": 1}] + config = AsyncpgConfig() + driver = AsyncpgDriver(cast("Any", connection), driver_features=config.driver_features) + target = tmp_path / "asyncpg.parquet" + + job = await driver.select_to_storage("SELECT name, value FROM storage_rows", target, format_hint="parquet") + + assert config.storage_capabilities()["arrow_export_enabled"] is False + assert job.telemetry["rows_processed"] == 1 + assert parquet.read_table(target).to_pylist() == [{"name": "Ada", "value": 1}] + + class _FakeTransaction: def __init__(self, record: "list[str]", nested: bool = False) -> None: self._record = record @@ -500,26 +517,6 @@ def test_asyncpg_events_resolve_listener_config_allows_mapping_direct_dsn() -> N assert resolved is not None -async def test_asyncpg_get_indexes_with_none_schema_resolves_via_search_path() -> None: - """get_indexes passes schema_name=None when schema is None to resolve table on search_path.""" - from sqlspec.adapters.asyncpg.data_dictionary import AsyncpgDataDictionary - - calls: list[dict[str, Any]] = [] - driver = MagicMock() - - async def fake_select(query: Any, **kwargs: Any) -> list[Any]: - calls.append(kwargs) - return [] - - driver.select = AsyncMock(side_effect=fake_select) - - dd = AsyncpgDataDictionary() - await dd.get_indexes(driver, table="users", schema=None) - assert len(calls) == 1 - assert calls[0]["schema_name"] is None - assert calls[0]["table_name"] == "users" - - async def test_asyncpg_cache_hit_with_cursor() -> None: """AsyncpgDriver executes cached queries directly on connection via _execute_cache_hit_with_cursor.""" conn = MagicMock() diff --git a/tests/unit/adapters/test_psqlpy/test_driver.py b/tests/unit/adapters/test_psqlpy/test_driver.py index 20b9e2b15..8ed910ba7 100644 --- a/tests/unit/adapters/test_psqlpy/test_driver.py +++ b/tests/unit/adapters/test_psqlpy/test_driver.py @@ -1,11 +1,13 @@ """Unit tests for PsqlpyDriver wire-level sqlcommenter and _dml_count_query cache preservation.""" +from pathlib import Path from types import SimpleNamespace from typing import Any, cast from unittest.mock import AsyncMock, patch import pytest +from sqlspec.adapters.psqlpy import PsqlpyConfig from sqlspec.adapters.psqlpy import core as psqlpy_core from sqlspec.adapters.psqlpy.core import build_statement_config from sqlspec.adapters.psqlpy.driver import PsqlpyDriver @@ -14,6 +16,23 @@ from sqlspec.exceptions import UniqueViolationError +@pytest.mark.anyio +async def test_storage_bridge_writes_client_converted_rows_without_native_export(tmp_path: Path) -> None: + """Client-side storage export works without claiming native Arrow export.""" + parquet = pytest.importorskip("pyarrow.parquet") + connection = AsyncMock() + connection.fetch.return_value = SimpleNamespace(result=lambda: [{"name": "Ada", "value": 1}]) + config = PsqlpyConfig() + driver = PsqlpyDriver(cast("Any", connection), driver_features=config.driver_features) + target = tmp_path / "psqlpy.parquet" + + job = await driver.select_to_storage("SELECT name, value FROM storage_rows", target, format_hint="parquet") + + assert config.storage_capabilities()["arrow_export_enabled"] is False + assert job.telemetry["rows_processed"] == 1 + assert parquet.read_table(target).to_pylist() == [{"name": "Ada", "value": 1}] + + @pytest.mark.anyio async def test_psqlpy_dispatch_execute_preserves_dml_count_query_lru_cache_under_sqlcommenter() -> None: """Executing the same DML across 5 traceparent contexts causes only 1 _dml_count_query cache miss.""" From 8973fc440f743c9d2cfe2fb4384b49e3a9bd3970 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Mon, 2 Nov 2026 00:34:34 +0000 Subject: [PATCH 165/174] test(db2): cover native isolation and typed metadata contracts --- .../integration/adapters/db2/test_adapter.py | 261 ++++++++++++++++++ 1 file changed, 261 insertions(+) diff --git a/tests/integration/adapters/db2/test_adapter.py b/tests/integration/adapters/db2/test_adapter.py index 282c10874..ad5ff4c9e 100644 --- a/tests/integration/adapters/db2/test_adapter.py +++ b/tests/integration/adapters/db2/test_adapter.py @@ -1,15 +1,33 @@ """Integration tests for IBM Db2 behavior outside the shared adapter contracts.""" +import asyncio +import json from collections.abc import Generator +from concurrent.futures import ThreadPoolExecutor +from datetime import date, datetime, time from decimal import Decimal from typing import Any +from uuid import UUID import pytest pytest.importorskip("ibm_db") pytest.importorskip("ibm_db_dbi") +from sqlspec import sql +from sqlspec.adapters.db2.config import Db2AsyncConfig, Db2SyncConfig +from sqlspec.adapters.db2.core import extract_sqlstate from sqlspec.adapters.db2.driver import Db2SyncDriver +from sqlspec.adapters.db2.events import Db2SyncEventQueueStore +from sqlspec.exceptions import ( + CheckViolationError, + ConnectionTimeoutError, + NotNullViolationError, + SQLParsingError, + SQLSpecError, + UniqueViolationError, +) +from sqlspec.extensions.events import EventMessage, SyncTableEventQueue pytestmark = [pytest.mark.db2, pytest.mark.xdist_group("db2")] @@ -113,3 +131,246 @@ def test_db2_integration_arrow_conversion(db2_session: Db2SyncDriver, clean_user assert isinstance(table, pa.Table) assert table.to_pydict() == {"id": [42], "name": ["ArrowUser"]} + + +def test_db2_semicolon_schema_survives_native_dsn_roundtrip(db2_connection_config: dict[str, Any]) -> None: + """The CLI driver receives one semicolon-bearing schema value after DSN parsing and rendering.""" + schema = "SQLSPEC;SEMICOLON" + source = Db2SyncConfig(connection_config={**db2_connection_config, "current_schema": schema}) + config = Db2SyncConfig(connection_config={"dsn": source.get_connection_string()}) + try: + with config.provide_session() as session: + assert session.select_value("SELECT CURRENT SCHEMA FROM SYSIBM.SYSDUMMY1").strip() == schema + finally: + config.close_pool() + + +def test_db2_skip_locked_dequeue_excludes_another_sessions_locked_event( + db2_session: Db2SyncDriver, db2_connection_config: dict[str, Any], db2_contract_table: str +) -> None: + """A consumer claims the second event while a different connection holds the first row locked.""" + config = Db2SyncConfig( + connection_config=dict(db2_connection_config), extension_config={"events": {"queue_table": db2_contract_table}} + ) + store = Db2SyncEventQueueStore(config) + for statement in store.create_statements(): + db2_session.execute_script(statement) + db2_session.commit() + queue = SyncTableEventQueue(config, queue_table=db2_contract_table, select_for_update=True, skip_locked=True) + try: + first_id = queue.publish("native_locking", {"position": 1}) + second_id = queue.publish("native_locking", {"position": 2}) + db2_session.begin() + locked = db2_session.select_one( + sql.select("event_id").from_(db2_contract_table).where_eq("event_id", first_id).for_update() + ) + assert locked["event_id"] == first_id + with ThreadPoolExecutor(max_workers=1) as executor: + pending = executor.submit(queue.dequeue, "native_locking", 0.0) + try: + message = pending.result(timeout=10) + assert isinstance(message, EventMessage) + assert message.event_id == second_id + assert message.payload == {"position": 2} + finally: + db2_session.rollback() + first = queue.dequeue("native_locking", poll_interval=0.0) + assert isinstance(first, EventMessage) + assert first.event_id == first_id + assert queue.dequeue("native_locking", poll_interval=0.0) is None + finally: + db2_session.rollback() + config.close_pool() + + +async def test_db2_async_native_pool_bounds_contending_sessions_and_recovers( + db2_connection_config: dict[str, Any], +) -> None: + """Concurrent native borrowers fill the configured capacity and a third times out without leaking it.""" + config = Db2AsyncConfig(connection_config={**db2_connection_config, "max_size": 2, "acquire_timeout": 0.1}) + entered: asyncio.Queue[int] = asyncio.Queue() + release = asyncio.Event() + + async def hold(value: int) -> None: + async with config.provide_session() as session: + observed = await session.select_value("SELECT CAST(:value AS INTEGER) FROM SYSIBM.SYSDUMMY1", value=value) + await entered.put(observed) + await release.wait() + + workers = [asyncio.create_task(hold(value)) for value in (1, 2)] + try: + values = await asyncio.wait_for(asyncio.gather(entered.get(), entered.get()), timeout=10) + assert set(values) == {1, 2} + pool = await config.provide_pool() + assert (pool.size(), pool.checked_out()) == (2, 2) + with pytest.raises(ConnectionTimeoutError): + async with config.provide_session(): + pytest.fail("A third session exceeded max_size while both native borrowers were held.") + release.set() + await asyncio.wait_for(asyncio.gather(*workers), timeout=10) + assert (pool.size(), pool.checked_out()) == (2, 0) + async with config.provide_session() as session: + assert await session.select_value("SELECT 42 FROM SYSIBM.SYSDUMMY1") == 42 + finally: + release.set() + try: + await asyncio.gather(*workers) + finally: + await config.close_pool() + + +@pytest.mark.parametrize( + ("statement", "parameters", "error_type", "sqlstate"), + [ + ( + "INSERT INTO TEST_INTEGRATION_USERS (id, name, balance) VALUES (:id, :name, :balance)", + {"id": 1, "name": "duplicate", "balance": 1}, + UniqueViolationError, + "23505", + ), + ( + "INSERT INTO TEST_INTEGRATION_USERS (id, name, balance) VALUES (:id, :name, :balance)", + {"id": 2, "name": None, "balance": 1}, + NotNullViolationError, + "23502", + ), + ( + "INSERT INTO TEST_INTEGRATION_USERS (id, name, balance) VALUES (:id, :name, :balance)", + {"id": 2, "name": "negative", "balance": -1}, + CheckViolationError, + "23513", + ), + ("SELECT * FROM SQLSPEC_MISSING_NATIVE_TABLE", {}, SQLParsingError, "42704"), + ], + ids=["unique", "not-null", "check", "missing-table"], +) +def test_db2_native_sqlstate_maps_exact_public_exception( + db2_session: Db2SyncDriver, + clean_users_table: None, + statement: str, + parameters: dict[str, Any], + error_type: type[SQLSpecError], + sqlstate: str, +) -> None: + """Real CLI diagnostics preserve representative SQLSTATEs and their specific mapped exception types.""" + db2_session.execute_script( + "ALTER TABLE TEST_INTEGRATION_USERS ADD CONSTRAINT SQLSPEC_BALANCE_CHECK CHECK (balance >= 0)" + ) + db2_session.execute("INSERT INTO TEST_INTEGRATION_USERS (id, name) VALUES (:id, :name)", id=1, name="retained") + db2_session.commit() + with pytest.raises(error_type) as caught: + db2_session.execute(statement, **parameters) + assert extract_sqlstate(caught.value) == sqlstate + db2_session.rollback() + + +@pytest.mark.parametrize("multiple_rows", [False, True], ids=["mapping", "two-rows"]) +def test_db2_merge_extended_typed_source_roundtrip( + db2_session: Db2SyncDriver, db2_contract_table: str, multiple_rows: bool +) -> None: + """The native engine executes the builder's bool, temporal, UUID, CLOB, binary, float, JSON and NULL casts.""" + db2_session.execute_script( + f"CREATE TABLE {db2_contract_table} (" + "record_id INTEGER NOT NULL PRIMARY KEY, label VARCHAR(40), enabled BOOLEAN, amount DECFLOAT(34), " + "captured_at TIMESTAMP, uuid_value VARCHAR(36), body CLOB, day_value DATE, time_value TIME, " + "binary_value VARBINARY(64), ratio DOUBLE, payload CLOB, absent CLOB)" + ) + db2_session.commit() + identifier = UUID("12345678-1234-5678-1234-567812345678") + values = { + "record_id": 1, + "label": "typed", + "enabled": True, + "amount": Decimal("12345678901234567890.123456789"), + "captured_at": datetime(2026, 6, 7, 8, 9, 10, 345678), + "uuid_value": identifier, + "body": "x" * 32673, + "day_value": date(2026, 6, 7), + "time_value": time(12, 34, 56), + "binary_value": b"\x00\xff", + "ratio": 0.5, + "payload": {"nested": "value"}, + "absent": None, + } + query = ( + sql + .merge(dialect="db2") + .into(db2_contract_table, alias="target") + .using([values, {**values, "record_id": 2, "payload": None}] if multiple_rows else values, alias="src") + .on("target.record_id = src.record_id") + .when_not_matched_then_insert({column: f"src.{column}" for column in values}) + ) + db2_session.execute(query) + db2_session.commit() + row = db2_session.select_one(f"SELECT * FROM {db2_contract_table} WHERE record_id = :record_id", record_id=1) + assert (row["record_id"], row["label"], row["enabled"]) == (1, "typed", 1) + assert Decimal(str(row["amount"])) == Decimal("12345678901234567890.123456789") + assert row["captured_at"] == datetime(2026, 6, 7, 8, 9, 10, 345678) + assert str(row["uuid_value"]) == "12345678-1234-5678-1234-567812345678" + assert row["body"] == "x" * 32673 + assert (row["day_value"], row["time_value"]) == (date(2026, 6, 7), time(12, 34, 56)) + assert bytes(row["binary_value"]) == b"\x00\xff" + assert row["ratio"] == 0.5 + payload = row["payload"] if isinstance(row["payload"], dict) else json.loads(row["payload"]) + assert payload == {"nested": "value"} + assert row["absent"] is None + if multiple_rows: + assert db2_session.select_value(f"SELECT COUNT(*) FROM {db2_contract_table}") == 2 + assert ( + db2_session.select_value( + f"SELECT payload FROM {db2_contract_table} WHERE record_id = :record_id", record_id=2 + ) + is None + ) + + +def test_db2_catalog_uniqueness_and_distinct_objects(db2_session: Db2SyncDriver, db2_contract_table: str) -> None: + """Catalog flags distinguish single-column uniqueness and composite keys without duplicated objects/indexes.""" + db2_session.execute_script( + f"CREATE TABLE {db2_contract_table} (" + "id INTEGER NOT NULL PRIMARY KEY, email VARCHAR(64) NOT NULL UNIQUE, " + "left_id INTEGER NOT NULL, right_id INTEGER NOT NULL, UNIQUE (left_id, right_id))" + ) + db2_session.commit() + dictionary = db2_session.data_dictionary + columns = { + column["column_name"]: column for column in dictionary.get_columns(db2_session, table=db2_contract_table) + } + assert columns["ID"]["is_primary"] is True + assert columns["EMAIL"]["is_primary"] is False + assert columns["EMAIL"]["is_unique"] is True + assert columns["LEFT_ID"]["is_unique"] is False + assert columns["RIGHT_ID"]["is_unique"] is False + indexes = dictionary.get_indexes(db2_session, table=db2_contract_table) + identities = [(index["schema_name"], index["table_name"], index["index_name"]) for index in indexes] + assert len(identities) == len(set(identities)) == 3 + assert sorted(tuple(index["columns"] or ()) for index in indexes) == [("EMAIL",), ("ID",), ("LEFT_ID", "RIGHT_ID")] + objects = dictionary.get_objects(db2_session).items + matching = [item for item in objects if isinstance(item, dict) and item.get("object_name") == db2_contract_table] + assert [(item["object_name"], item["object_type"]) for item in matching] == [(db2_contract_table, "TABLE")] + tables = dictionary.get_tables(db2_session) + assert [table["table_name"] for table in tables if table["table_name"] == db2_contract_table] == [ + db2_contract_table + ] + + +@pytest.fixture +def db2_contract_table(db2_session: Db2SyncDriver) -> Generator[str, None, None]: + """Reserve a fixed scratch table and clean it before and after every native residual case.""" + table = "SQLSPEC_DB2_NATIVE_CONTRACT" + _drop_contract_table(db2_session, table) + try: + yield table + finally: + db2_session.rollback() + _drop_contract_table(db2_session, table) + + +def _drop_contract_table(session: Db2SyncDriver, table: str) -> None: + exists = session.select_value( + "SELECT COUNT(*) FROM SYSCAT.TABLES WHERE TABSCHEMA = CURRENT SCHEMA AND TABNAME = :table_name", + table_name=table, + ) + if exists: + session.execute_script(f"DROP TABLE {table}") + session.commit() From 7ab1f76d4d1a71b9b147c94bfd66449cd5be28b8 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Mon, 2 Nov 2026 00:37:31 +0000 Subject: [PATCH 166/174] fix(ci): keep SQL Anywhere SDK optional in wheel smoke --- tools/scripts/mypyc_smoke.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tools/scripts/mypyc_smoke.py b/tools/scripts/mypyc_smoke.py index 81f15174f..e0c0550d9 100644 --- a/tools/scripts/mypyc_smoke.py +++ b/tools/scripts/mypyc_smoke.py @@ -45,7 +45,7 @@ class SmokeImport(NamedTuple): SmokeImport("storage_backend_obstore", "sqlspec.storage.backends.obstore", "ObStoreBackend", True), SmokeImport("sqlite_pool", "sqlspec.adapters.sqlite.pool", "SqliteConnectionPool", True), SmokeImport("sybase_pool", "sqlspec.adapters.sybase.pool", "SybaseConnectionPool", True), - SmokeImport("sqlanywhere_pool", "sqlspec.adapters.sqlanywhere.pool", "SqlanywhereConnectionPool", True), + SmokeImport("sqlanywhere_pool", "sqlspec.adapters.sqlanywhere.pool", "SqlanywhereConnectionPool", True, "sqlanydb"), SmokeImport("data_dictionary_registry", "sqlspec.data_dictionary._registry", "get_dialect_config", True), SmokeImport("data_dictionary_loader", "sqlspec.data_dictionary._loader", "DataDictionaryLoader", True), SmokeImport("pgvector_dialect", "sqlspec.dialects.postgres._pgvector", "PGVector"), From b0fd19e64af3406ee74396f19e0382bcede04cd8 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Mon, 2 Nov 2026 00:40:08 +0000 Subject: [PATCH 167/174] fix(oracledb): restore stack and index metadata with unified version lookup --- sqlspec/adapters/oracledb/driver.py | 1 - .../data_dictionary/dialects/oracle/config.py | 2 +- .../dialects/oracle/sql/indexes.sql | 3 +- .../dialects/oracle/sql/version.sql | 10 +- tests/integration/adapters/_shared/_cases.py | 2 +- .../adapters/oracle/oracledb/test_features.py | 21 ++- .../test_data_dictionary_oracledb.py | 152 +++++++++++++++++- .../test_stack_native_transaction.py | 43 +++++ 8 files changed, 222 insertions(+), 12 deletions(-) diff --git a/sqlspec/adapters/oracledb/driver.py b/sqlspec/adapters/oracledb/driver.py index 0954c3a90..53034807a 100644 --- a/sqlspec/adapters/oracledb/driver.py +++ b/sqlspec/adapters/oracledb/driver.py @@ -878,7 +878,6 @@ def select_to_arrow( def execute_stack(self, stack: "StatementStack", *, continue_on_error: bool = False) -> "tuple[StackResult, ...]": """Execute a StatementStack sequentially on the sync driver.""" - self._require_capability("stack_execution") self._log_pipeline_skip("asyncio_thin_required", stack) return super().execute_stack(stack, continue_on_error=continue_on_error) diff --git a/sqlspec/data_dictionary/dialects/oracle/config.py b/sqlspec/data_dictionary/dialects/oracle/config.py index f1dc95a36..ec2ba12a9 100644 --- a/sqlspec/data_dictionary/dialects/oracle/config.py +++ b/sqlspec/data_dictionary/dialects/oracle/config.py @@ -168,7 +168,7 @@ def extract_oracle_version_value(row: object) -> "str | None": Version string when one can be found, otherwise None. """ if isinstance(row, Mapping): - for key in ("version", "VERSION", "Version"): + for key in ("version_full", "VERSION_FULL", "Version_Full", "version", "VERSION", "Version"): value = row.get(key) if value: return str(value) diff --git a/sqlspec/data_dictionary/dialects/oracle/sql/indexes.sql b/sqlspec/data_dictionary/dialects/oracle/sql/indexes.sql index e6a90f461..c18eb142f 100644 --- a/sqlspec/data_dictionary/dialects/oracle/sql/indexes.sql +++ b/sqlspec/data_dictionary/dialects/oracle/sql/indexes.sql @@ -52,6 +52,7 @@ ORDER BY i.owner, i.table_name, i.index_name; -- name: by_table -- dialect: oracle SELECT + i.owner AS schema_name, i.index_name, i.table_name AS table_name, CASE WHEN i.uniqueness = 'UNIQUE' THEN 1 ELSE 0 END AS is_unique, @@ -70,7 +71,7 @@ LEFT JOIN all_constraints c AND c.constraint_type = 'P' WHERE i.table_name = :table_name AND i.owner = COALESCE(:schema_name, USER) -GROUP BY i.index_name, i.table_name, i.uniqueness; +GROUP BY i.owner, i.index_name, i.table_name, i.uniqueness; -- name: expressions -- dialect: oracle diff --git a/sqlspec/data_dictionary/dialects/oracle/sql/version.sql b/sqlspec/data_dictionary/dialects/oracle/sql/version.sql index bbf3ec91c..ee88f5ca7 100644 --- a/sqlspec/data_dictionary/dialects/oracle/sql/version.sql +++ b/sqlspec/data_dictionary/dialects/oracle/sql/version.sql @@ -1,12 +1,10 @@ -- name: current -- dialect: oracle SELECT - product AS "product", - version AS "version", - status AS "status" -FROM product_component_version -WHERE product LIKE 'Oracle%' -ORDER BY TO_NUMBER(REGEXP_SUBSTR(version, '^[0-9]+')) DESC, version DESC + p.* +FROM product_component_version p +WHERE p.product LIKE 'Oracle%' +ORDER BY TO_NUMBER(REGEXP_SUBSTR(p.version, '^[0-9]+')) DESC, p.version DESC FETCH FIRST 1 ROWS ONLY; -- name: autonomous_service diff --git a/tests/integration/adapters/_shared/_cases.py b/tests/integration/adapters/_shared/_cases.py index 6b6d28da4..08d042396 100644 --- a/tests/integration/adapters/_shared/_cases.py +++ b/tests/integration/adapters/_shared/_cases.py @@ -438,7 +438,7 @@ def _db2_case(mode: Literal["sync", "async"], marks: tuple[Mark | MarkDecorator, supports_arrow_streaming=True, supports_native_arrow=True, arrow_reader_honors_batch_size=True, - arrow_reader_reports_rows_affected=True, + arrow_reader_reports_rows_affected=False, fixture_name="contract_oracle_sync_driver", adapter="oracledb", dialect="oracle", diff --git a/tests/integration/adapters/oracle/oracledb/test_features.py b/tests/integration/adapters/oracle/oracledb/test_features.py index 95d1be119..fba498148 100644 --- a/tests/integration/adapters/oracle/oracledb/test_features.py +++ b/tests/integration/adapters/oracle/oracledb/test_features.py @@ -5,8 +5,9 @@ import pytest -from sqlspec.adapters.oracledb import OracleAsyncDriver, OracleSyncDriver +from sqlspec.adapters.oracledb import OracleAsyncDriver, OracleSyncDriver, OracleVersionInfo from sqlspec.core import SQL, SQLResult, StatementConfig +from sqlspec.data_dictionary.dialects.oracle import parse_oracle_version_components def _lower_keys(row: dict[str, object]) -> dict[str, object]: @@ -16,6 +17,24 @@ def _lower_keys(row: dict[str, object]) -> dict[str, object]: pytestmark = pytest.mark.xdist_group("oracle") +def test_oracle_release_metadata_matches_native_vector_capable_server(oracle_sync_session: OracleSyncDriver) -> None: + """The 23ai fixture reports its release update and actual vector-compatible level.""" + native_version = str( + oracle_sync_session.select_value( + "SELECT COALESCE(version_full, version) FROM product_component_version " + "WHERE product LIKE 'Oracle%' FETCH FIRST 1 ROWS ONLY" + ) + ) + native_compatible = str(oracle_sync_session.select_value("SELECT value FROM v$parameter WHERE name = 'compatible'")) + version = oracle_sync_session.data_dictionary.get_version(oracle_sync_session) + assert isinstance(version, OracleVersionInfo) + assert (version.major, version.minor, version.patch) == parse_oracle_version_components(native_version) + assert version.compatible == native_compatible + compatible = parse_oracle_version_components(native_compatible) + assert isinstance(compatible, tuple), f"Invalid native COMPATIBLE: {native_compatible!r}" + assert compatible >= (23, 4, 0), f"Vector fixture COMPATIBLE: {native_compatible!r}" + + def test_sync_oracle_data_types(oracle_sync_session: OracleSyncDriver) -> None: """Test Oracle-specific data types (NUMBER, VARCHAR2, CLOB, DATE).""" diff --git a/tests/unit/adapters/test_oracledb/test_data_dictionary_oracledb.py b/tests/unit/adapters/test_oracledb/test_data_dictionary_oracledb.py index 2f25eb2f1..0396e2cf6 100644 --- a/tests/unit/adapters/test_oracledb/test_data_dictionary_oracledb.py +++ b/tests/unit/adapters/test_oracledb/test_data_dictionary_oracledb.py @@ -5,6 +5,7 @@ import pytest from sqlspec.adapters.oracledb.data_dictionary import OracledbAsyncDataDictionary, OracledbSyncDataDictionary +from sqlspec.adapters.oracledb.driver import OracleAsyncDriver, OracleSyncDriver from sqlspec.data_dictionary import ( DDLResult, MetadataFidelity, @@ -32,7 +33,7 @@ ) if TYPE_CHECKING: - from sqlspec.adapters.oracledb.driver import OracleAsyncDriver, OracleSyncDriver + from sqlspec.adapters.oracledb._typing import OracleAsyncConnection, OracleSyncConnection ORACLE_AI_COMPONENT_ROW = { "product": "Oracle AI Database 26ai Free", @@ -220,6 +221,8 @@ def test_oracle_dialect_helpers_resolve_shared_version_rules() -> None: config = get_dialect_config("oracle") assert extract_oracle_version_value({"VERSION": "23.26.0.0.0"}) == "23.26.0.0.0" + assert extract_oracle_version_value({"VERSION": "23.0.0.0.0", "VERSION_FULL": "23.26.0.0.0"}) == "23.26.0.0.0" + assert extract_oracle_version_value({"version": "19.0.0.0.0", "version_full": None}) == "19.0.0.0.0" assert extract_oracle_version_value(("19.20.0.0.0",)) == "19.20.0.0.0" assert parse_oracle_version_components("23.26.0.0.0") == (23, 26, 0) assert parse_oracle_compatible_major("23.20.0.0.0") == 23 @@ -657,3 +660,150 @@ def select(self, query: Any, **kwargs: Any) -> list[dict[str, Any]]: assert len(tables) == 1 assert len(queries) == 1 assert "names_by_schema" not in queries[0] + + +class _CatalogCursor: + def __init__(self, connection: "_CatalogConnection") -> None: + self.connection = connection + self.description: list[tuple[str, object]] = [] + self.rows: list[tuple[object, ...]] = [] + self.rowcount = 0 + + def execute(self, statement: str, parameters: object = None, **kwargs: object) -> None: + import oracledb + + self.connection.statements.append(statement) + if "product_component_version" in statement: + columns = ["PRODUCT", "VERSION", "STATUS"] + self.rows = [("Oracle Database", self.connection.catalog_version, "Production")] + if int(self.connection.catalog_version.split(".", 1)[0]) >= 18: + columns.insert(2, "VERSION_FULL") + self.rows = [ + ("Oracle Database", self.connection.catalog_version, self.connection.full_version, "Production") + ] + elif "v$parameter" in statement: + columns = ["COMPATIBLE"] + self.rows = [("23.4.0.0.0",)] + elif "sys_context" in statement: + columns = ["SERVICE"] + self.rows = [(None,)] + elif "all_ind_expressions" in statement: + columns = ["SCHEMA_NAME", "TABLE_NAME", "INDEX_NAME", "COLUMN_EXPRESSION"] + self.rows = [("HR", "EMPLOYEES", "IDX_EMAIL", 'UPPER("EMAIL")')] + elif "all_indexes" in statement: + columns = ["INDEX_NAME", "TABLE_NAME", "IS_UNIQUE", "IS_PRIMARY", "COLUMNS"] + self.rows = [("IDX_EMAIL", "EMPLOYEES", 0, 0, "SYS_NC00003$")] + if "i.owner AS schema_name" in statement: + columns.insert(0, "SCHEMA_NAME") + self.rows = [("HR", *row) for row in self.rows] + else: + raise ValueError(statement) + self.description = [(name, oracledb.DB_TYPE_VARCHAR) for name in columns] + self.rowcount = len(self.rows) + + def fetchall(self) -> list[tuple[object, ...]]: + rows, self.rows = self.rows, [] + return rows + + def fetchone(self) -> tuple[object, ...] | None: + return self.rows.pop(0) if self.rows else None + + def close(self) -> None: + return None + + +class _AsyncCatalogCursor: + def __init__(self, connection: "_CatalogConnection") -> None: + self.cursor = _CatalogCursor(connection) + + @property + def description(self) -> list[tuple[str, object]]: + return self.cursor.description + + @property + def rowcount(self) -> int: + return self.cursor.rowcount + + async def execute(self, statement: str, parameters: object = None, **kwargs: object) -> None: + self.cursor.execute(statement, parameters, **kwargs) + + async def fetchall(self) -> list[tuple[object, ...]]: + return self.cursor.fetchall() + + async def fetchone(self) -> tuple[object, ...] | None: + return self.cursor.fetchone() + + def close(self) -> None: + self.cursor.close() + + +class _CatalogConnection: + def __init__( + self, + *, + is_async: bool, + version: str = "23.0.0.0.0", + full_version: str = "23.26.0.0.0", + sdk_version_present: bool = True, + ) -> None: + self.is_async = is_async + self.version = version if sdk_version_present else None + self.catalog_version = version + self.full_version = full_version + self.statements: list[str] = [] + + def cursor(self) -> _CatalogCursor | _AsyncCatalogCursor: + return _AsyncCatalogCursor(self) if self.is_async else _CatalogCursor(self) + + +@pytest.mark.anyio +@pytest.mark.parametrize("is_async", [False, True]) +async def test_oracle_default_schema_index_expression_uses_native_catalog_owner(is_async: bool) -> None: + connection = _CatalogConnection(is_async=is_async) + if is_async: + driver = OracleAsyncDriver( + cast("OracleAsyncConnection", connection), driver_features={"enable_lowercase_column_names": True} + ) + indexes = await driver.data_dictionary.get_indexes(driver, table="employees") + else: + sync_driver = OracleSyncDriver( + cast("OracleSyncConnection", connection), driver_features={"enable_lowercase_column_names": True} + ) + indexes = sync_driver.data_dictionary.get_indexes(sync_driver, table="employees") + assert len(indexes) == 1 + assert indexes[0]["index_name"] == "IDX_EMAIL" + assert indexes[0]["expression"] == 'UPPER("EMAIL")' + + +@pytest.mark.anyio +@pytest.mark.parametrize("is_async", [False, True]) +@pytest.mark.parametrize( + ("major", "sdk_version_present", "lowercase_columns"), + [(12, True, True), (18, True, False), (23, True, True), (23, False, False)], +) +async def test_oracle_version_preserves_release_update_and_legacy_query( + is_async: bool, major: int, sdk_version_present: bool, lowercase_columns: bool +) -> None: + base_version = f"{major}.0.0.0.0" + full_version = f"{major}.26.0.0.0" + connection = _CatalogConnection( + is_async=is_async, version=base_version, full_version=full_version, sdk_version_present=sdk_version_present + ) + if is_async: + driver = OracleAsyncDriver( + cast("OracleAsyncConnection", connection), + driver_features={"enable_lowercase_column_names": lowercase_columns}, + ) + version = await driver.data_dictionary.get_version(driver) + cached_version = await driver.data_dictionary.get_version(driver) + else: + sync_driver = OracleSyncDriver( + cast("OracleSyncConnection", connection), + driver_features={"enable_lowercase_column_names": lowercase_columns}, + ) + version = sync_driver.data_dictionary.get_version(sync_driver) + cached_version = sync_driver.data_dictionary.get_version(sync_driver) + assert version == VersionInfo(major, 26 if major >= 18 else 0, 0) + assert cached_version is version + version_statements = [query for query in connection.statements if "product_component_version" in query] + assert len(version_statements) == 1 diff --git a/tests/unit/adapters/test_oracledb/test_stack_native_transaction.py b/tests/unit/adapters/test_oracledb/test_stack_native_transaction.py index a03b1fa83..061e22be5 100644 --- a/tests/unit/adapters/test_oracledb/test_stack_native_transaction.py +++ b/tests/unit/adapters/test_oracledb/test_stack_native_transaction.py @@ -163,3 +163,46 @@ async def test_stack_maps_native_pipeline_error_records(continue_on_error: bool) cause = error.original_error.__cause__ assert isinstance(cause, oracledb.DatabaseError) assert cause.args[0] is native_error + + +def test_sync_stack_executes_and_commits_without_a_storage_capability_flag() -> None: + from sqlspec.adapters.oracledb._typing import OracleSyncConnection + from sqlspec.adapters.oracledb.driver import OracleSyncDriver + + class _Cursor: + description = None + rowcount = 1 + + def __init__(self, connection: "_Connection") -> None: + self.connection = connection + + def execute(self, statement: str, parameters: object = None, **kwargs: object) -> None: + self.connection.executed.append((statement, parameters)) + + def close(self) -> None: + return None + + class _Connection: + transaction_in_progress = False + + def __init__(self) -> None: + self.executed: list[tuple[str, object]] = [] + self.commits = 0 + + def cursor(self) -> _Cursor: + return _Cursor(self) + + def commit(self) -> None: + self.commits += 1 + + connection = _Connection() + driver = OracleSyncDriver(cast("OracleSyncConnection", connection)) + stack = ( + StatementStack() + .push_execute("INSERT INTO items VALUES (:1)", (1,)) + .push_execute("INSERT INTO items VALUES (:1)", (2,)) + ) + results = driver.execute_stack(stack) + assert [result.rows_affected for result in results] == [1, 1] + assert connection.executed == [("INSERT INTO items VALUES (:1)", [1]), ("INSERT INTO items VALUES (:1)", [2])] + assert connection.commits == 1 From 4bcb0eeb8ce91e412bb22641fce567f3d0995338 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Mon, 2 Nov 2026 00:43:31 +0000 Subject: [PATCH 168/174] test(db2): execute native Arrow-ODBC with isolated client runtime --- .../adapters/db2/arrow_odbc/test_db2.py | 153 ++++++++++++++++++ tests/integration/fixtures/db2.py | 73 +++++++++ 2 files changed, 226 insertions(+) create mode 100644 tests/integration/adapters/db2/arrow_odbc/test_db2.py diff --git a/tests/integration/adapters/db2/arrow_odbc/test_db2.py b/tests/integration/adapters/db2/arrow_odbc/test_db2.py new file mode 100644 index 000000000..e725c6642 --- /dev/null +++ b/tests/integration/adapters/db2/arrow_odbc/test_db2.py @@ -0,0 +1,153 @@ +"""Native IBM Db2 acceptance through the 64-bit Arrow-ODBC transport.""" + +import json +import platform +import sys +from typing import TYPE_CHECKING + +import pytest +from pytest_databases.docker.db2 import DB2_PORT +from pytest_databases.helpers import exec_output_to_bytes + +if TYPE_CHECKING: + from pytest_databases.docker.db2 import Db2Service + + from sqlspec.adapters.db2 import Db2ConnectionParams + +pytestmark = [ + pytest.mark.db2, + pytest.mark.integration, + pytest.mark.xdist_group("db2"), + pytest.mark.skipif( + sys.platform != "linux" or platform.machine().lower() not in {"x86_64", "amd64"}, + reason="The IBM 64-bit ODBC runtime acceptance uses a Linux x86-64 client", + ), +] + +_ACCEPTANCE_SCRIPT = """ +import json +import os +from decimal import Decimal +from uuid import uuid4 + +import pyarrow as pa +from sqlspec.adapters.arrow_odbc import ArrowOdbcConfig + +payload = json.loads(os.environ.pop("SQLSPEC_DB2_ODBC_CONFIG")) +config = ArrowOdbcConfig( + connection_config=payload, + driver_features={"enable_pooling": False, "query_timeout_sec": 15}, +) +table_name = "SQLSPEC_ARROW_ODBC_" + uuid4().hex[:12].upper() +index_name = table_name + "_UQ" +created = False +try: + with config.provide_session() as session: + assert session.dialect == "db2" + try: + session.execute( + f"CREATE TABLE {table_name} (" + "id INTEGER NOT NULL PRIMARY KEY, big_value BIGINT, label VARCHAR(64), amount DECIMAL(14, 3))" + ) + session.commit() + created = True + session.execute(f"CREATE UNIQUE INDEX {index_name} ON {table_name} (label)") + session.commit() + insert = f"INSERT INTO {table_name} (id, big_value, label, amount) VALUES (?, ?, ?, ?)" + session.execute(insert, 1, 2**40, "Grüße 東京", Decimal("123.450")) + session.execute(insert, 2, None, None, None) + session.commit() + session.begin() + session.execute(insert, 3, 2**42, "rolled back", Decimal("456.789")) + session.rollback() + assert session.select_value(f"SELECT COUNT(*) FROM {table_name} WHERE id = ?", 3) == 0 + rows = session.select(f"SELECT id, big_value, label, amount FROM {table_name} ORDER BY id") + assert rows[0]["id"] == 1 + assert rows[0]["big_value"] == 2**40 + assert rows[0]["label"] == "Grüße 東京" + assert Decimal(str(rows[0]["amount"])) == Decimal("123.450") + assert rows[1] == {"id": 2, "big_value": None, "label": None, "amount": None} + result = session.select_to_arrow( + f"SELECT id, big_value, label, amount FROM {table_name} WHERE id >= ? ORDER BY id", 1 + ) + table = result.get_data() + assert table.num_rows == 2 + assert table.column_names == ["id", "big_value", "label", "amount"] + assert table.schema.field("big_value").type == pa.int64() + assert table.schema.field("amount").type == pa.decimal128(14, 3) + assert table.column("big_value").to_pylist() == [2**40, None] + assert table.column("label").to_pylist() == ["Grüße 東京", None] + assert table.column("amount").to_pylist() == [Decimal("123.450"), None] + schema_name = str(session.select_value("SELECT CURRENT SCHEMA FROM SYSIBM.SYSDUMMY1")).strip() + dictionary = session.data_dictionary + columns = dictionary.get_columns(session, table=table_name, schema=schema_name) + assert [column["column_name"] for column in columns] == ["ID", "BIG_VALUE", "LABEL", "AMOUNT"] + assert [column["ordinal_position"] for column in columns] == [1, 2, 3, 4] + assert all( + column["table_name"] == table_name and column["schema_name"] == schema_name + for column in columns + ) + assert columns[0]["is_primary"] == 1 + assert columns[2]["is_unique"] == 1 + assert columns[3]["numeric_scale"] == 3 + indexes = dictionary.get_indexes(session, table=table_name, schema=schema_name) + assert len(indexes) == 2 + indexes_by_name = {index["index_name"]: index for index in indexes} + assert len(indexes_by_name) == 2 + assert indexes_by_name[index_name]["is_unique"] == 1 + assert sum(bool(index["is_primary"]) for index in indexes) == 1 + tables = [ + item for item in dictionary.get_tables(session, schema=schema_name) + if item["table_name"] == table_name + ] + assert len(tables) == 1 + assert tables[0]["schema_name"] == schema_name + assert tables[0]["table_type"] == "BASE TABLE" + finally: + session.rollback() + with config.provide_session() as reopened: + try: + persisted = reopened.select(f"SELECT id, big_value FROM {table_name} ORDER BY id") + assert persisted == [{"id": 1, "big_value": 2**40}, {"id": 2, "big_value": None}] + finally: + reopened.rollback() +finally: + try: + if created: + with config.provide_session() as cleanup: + cleanup.rollback() + cleanup.execute(f"DROP TABLE {table_name}") + cleanup.commit() + finally: + config.close_pool() +print(json.dumps({"dialect": "db2", "rows": 2, "transport": "arrow_odbc"})) +""" + + +def test_arrow_odbc_db2_parameter_binding_and_arrow_null_indicators( + db2_service: "Db2Service", + db2_connection_config: "Db2ConnectionParams", + db2_arrow_odbc_runtime: tuple[list[str], dict[str, str], str], +) -> None: + """Transactions, catalog metadata, and Arrow buffers preserve typed Db2 values and NULL indicators.""" + command, environment, driver = db2_arrow_odbc_runtime + password = str(db2_connection_config["password"]) + parameters = { + "driver": driver, + "host": "127.0.0.1", + "database": db2_connection_config["database"], + "uid": db2_connection_config["user"], + "pwd": password, + "extra": {"Port": DB2_PORT, "Protocol": "TCPIP"}, + "login_timeout": 15, + "autocommit": False, + } + result = db2_service.container.exec_run( + [*command, "-c", _ACCEPTANCE_SCRIPT], + environment={**environment, "SQLSPEC_DB2_ODBC_CONFIG": json.dumps(parameters)}, + ) + output = exec_output_to_bytes(result.output).decode("utf-8", errors="replace") + if result.exit_code: + diagnostic = output.replace(password, "") + pytest.fail(f"Native Arrow-ODBC Db2 acceptance failed (exit {result.exit_code}):\n{diagnostic}") + assert json.loads(output) == {"dialect": "db2", "rows": 2, "transport": "arrow_odbc"} diff --git a/tests/integration/fixtures/db2.py b/tests/integration/fixtures/db2.py index 01f62eae8..07df60854 100644 --- a/tests/integration/fixtures/db2.py +++ b/tests/integration/fixtures/db2.py @@ -1,8 +1,11 @@ """IBM Db2 integration test fixtures.""" import os +import platform import shlex +import sys from collections.abc import AsyncGenerator, Generator +from pathlib import Path from typing import TYPE_CHECKING, Any import pytest @@ -15,6 +18,7 @@ ) from pytest_databases.helpers import exec_output_to_bytes, get_xdist_worker_num +import sqlspec from sqlspec.adapters.db2.config import Db2AsyncConfig, Db2ConnectionParams, Db2SyncConfig if TYPE_CHECKING: @@ -24,6 +28,7 @@ from sqlspec.adapters.db2.driver import Db2AsyncDriver, Db2SyncDriver __all__ = ( + "db2_arrow_odbc_runtime", "db2_async_config", "db2_async_session", "db2_connection_config", @@ -42,6 +47,7 @@ def db2_image() -> str: @pytest.fixture(scope="session") def db2_service( + request: pytest.FixtureRequest, container_service: "ContainerService", db2_image: str, db2_dbname: str, @@ -51,6 +57,21 @@ def db2_service( xdist_db2_isolation_level: "XdistIsolationLevel", ) -> Generator[Db2Service, None, None]: """Start and yield a ready IBM Db2 service container with automatic 4Kn Direct-IO recovery.""" + client_needed = ( + sys.platform == "linux" + and platform.machine().lower() in {"x86_64", "amd64"} + and any( + isinstance(item, pytest.Function) and "db2_arrow_odbc_runtime" in item.fixturenames + for item in request.session.items + ) + ) + client_mounts = _db2_odbc_client_mounts() if client_needed else {} + if client_needed and not client_mounts: + msg = ( + "Arrow-ODBC Db2 acceptance requires an isolated or managed Linux x86-64 Python installation " + "and an available glibc dynamic loader; broad system Python prefixes cannot be mounted safely." + ) + raise RuntimeError(msg) worker_num = get_xdist_worker_num() name = "db2" initial_db = ( @@ -115,6 +136,7 @@ def check(service: "ServiceContainer") -> bool: transient=xdist_db2_isolation_level == "server", privileged=True, ipc_mode="host", + volumes=client_mounts or None, ) as service: if db2_dbname != initial_db: create_worker_database(service.container, db2_dbname, instance_user=db2_username) @@ -141,6 +163,36 @@ def db2_connection_config(db2_service: Db2Service) -> Db2ConnectionParams: } +@pytest.fixture(scope="session") +def db2_arrow_odbc_runtime(db2_service: Db2Service) -> tuple[list[str], dict[str, str], str]: + """Run the host client against the container's installed 64-bit IBM ODBC runtime. + + Linux x86-64 clients use their own Python and system loader from read-only + mounts. IBM's installation registry and instance profile remain inside the + existing Db2 service container; no host driver registration is changed. + """ + library_path = ":".join(( + "/opt/sqlspec-host-libraries", + str(Path(sys.base_prefix) / "lib"), + "/opt/ibm/db2/V11.5/lib64", + "/opt/ibm/db2/V11.5/lib64/gskit", + )) + driver = "/opt/ibm/db2/V11.5/lib64/libdb2o.so.1" + if not _db2_odbc_client_mounts(): + msg = "Arrow-ODBC Db2 acceptance requires a Linux x86-64 glibc host with its dynamic loader available." + raise RuntimeError(msg) + result = db2_service.container.exec_run(["test", "-r", driver]) + if result.exit_code: + msg = f"The pinned Db2 service image does not provide the required IBM 64-bit ODBC driver: {driver}" + raise RuntimeError(msg) + command = ["/opt/sqlspec-host-libraries/ld-linux-x86-64.so.2", "--library-path", library_path, sys.executable] + environment = { + "LD_LIBRARY_PATH": "/opt/ibm/db2/V11.5/lib64:/opt/ibm/db2/V11.5/lib64/gskit", + "DB2INSTANCE": db2_service.username, + } + return command, environment, driver + + @pytest.fixture(scope="session") def db2_sync_config(db2_connection_config: Db2ConnectionParams) -> Generator[Db2SyncConfig, None, None]: """Yield a session-scoped Db2SyncConfig and close its connection pool on teardown.""" @@ -175,6 +227,27 @@ async def db2_async_session(db2_async_config: Db2AsyncConfig) -> "AsyncGenerator yield driver +def _db2_odbc_client_mounts() -> dict[str, dict[str, str]]: + """Expose only the Python client, SQLSpec package, and host loader read-only.""" + if sys.platform != "linux" or platform.machine().lower() not in {"x86_64", "amd64"}: + return {} + if any( + Path(prefix).resolve() in {Path("/"), Path("/usr"), Path("/usr/local")} + for prefix in (sys.base_prefix, sys.prefix) + ): + return {} + system_libraries = next( + (path for path in (Path("/lib/x86_64-linux-gnu"), Path("/lib64")) if (path / "ld-linux-x86-64.so.2").is_file()), + None, + ) + if system_libraries is None: + return {} + paths = {Path(sys.base_prefix), Path(sys.prefix), Path(sqlspec.__file__).resolve().parent} + mounts = {str(path): {"bind": str(path), "mode": "ro"} for path in paths} + mounts[str(system_libraries.resolve())] = {"bind": "/opt/sqlspec-host-libraries", "mode": "ro"} + return mounts + + def _ensure_db2_4k_database(container: Any, db_name: str, *, instance_user: str) -> None: """Configure 4Kn Direct-IO registry settings and create the initial Db2 database if missing.""" validate_db2_database_name(db_name) From 5078352675091527ef8ba86716ef566b4cade09b Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Mon, 2 Nov 2026 00:53:42 +0000 Subject: [PATCH 169/174] docs(data-dictionary): refresh corrected catalog query references --- .../reference/data-dictionary/dialects/oracle.md | 13 ++++++------- .../data-dictionary/dialects/postgres.md | 16 ++++++++++++++-- 2 files changed, 20 insertions(+), 9 deletions(-) diff --git a/docs/reference/data-dictionary/dialects/oracle.md b/docs/reference/data-dictionary/dialects/oracle.md index 60b7f6743..fcd796888 100644 --- a/docs/reference/data-dictionary/dialects/oracle.md +++ b/docs/reference/data-dictionary/dialects/oracle.md @@ -482,6 +482,7 @@ The `oracle` dialect pack defines `36` named queries across `22` domain `.sql` f -- name: by_table -- dialect: oracle SELECT + i.owner AS schema_name, i.index_name, i.table_name AS table_name, CASE WHEN i.uniqueness = 'UNIQUE' THEN 1 ELSE 0 END AS is_unique, @@ -500,7 +501,7 @@ The `oracle` dialect pack defines `36` named queries across `22` domain `.sql` f AND c.constraint_type = 'P' WHERE i.table_name = :table_name AND i.owner = COALESCE(:schema_name, USER) - GROUP BY i.index_name, i.table_name, i.uniqueness; + GROUP BY i.owner, i.index_name, i.table_name, i.uniqueness; -- name: expressions -- dialect: oracle @@ -884,12 +885,10 @@ The `oracle` dialect pack defines `36` named queries across `22` domain `.sql` f -- name: current -- dialect: oracle SELECT - product AS "product", - version AS "version", - status AS "status" - FROM product_component_version - WHERE product LIKE 'Oracle%' - ORDER BY TO_NUMBER(REGEXP_SUBSTR(version, '^[0-9]+')) DESC, version DESC + p.* + FROM product_component_version p + WHERE p.product LIKE 'Oracle%' + ORDER BY TO_NUMBER(REGEXP_SUBSTR(p.version, '^[0-9]+')) DESC, p.version DESC FETCH FIRST 1 ROWS ONLY; -- name: autonomous_service diff --git a/docs/reference/data-dictionary/dialects/postgres.md b/docs/reference/data-dictionary/dialects/postgres.md index 7df5c70b7..20abc1d73 100644 --- a/docs/reference/data-dictionary/dialects/postgres.md +++ b/docs/reference/data-dictionary/dialects/postgres.md @@ -515,6 +515,18 @@ The `postgres` dialect pack defines `26` named queries across `21` domain `.sql` ```sql -- name: by_schema -- dialect: postgres + WITH target_schema AS ( + SELECT COALESCE( + :schema_name::text, + ( + SELECT rn.nspname::text + FROM pg_catalog.pg_class rc + JOIN pg_catalog.pg_namespace rn ON rn.oid = rc.relnamespace + WHERE rc.oid = pg_catalog.to_regclass(pg_catalog.quote_ident(:table_name::text)) + ), + current_schema() + ) AS schema_name + ) SELECT tn.nspname::text AS schema_name, tbl.relname::text AS table_name, @@ -550,11 +562,11 @@ The `postgres` dialect pack defines `26` named queries across `21` domain `.sql` FROM pg_catalog.pg_index ix JOIN pg_catalog.pg_class tbl ON tbl.oid = ix.indrelid JOIN pg_catalog.pg_namespace tn ON tn.oid = tbl.relnamespace + JOIN target_schema ts ON ts.schema_name = tn.nspname JOIN pg_catalog.pg_class idx ON idx.oid = ix.indexrelid JOIN pg_catalog.pg_am am ON am.oid = idx.relam LEFT JOIN pg_catalog.pg_depend dep ON dep.objid = idx.oid AND dep.deptype = 'e' - WHERE tn.nspname = :schema_name - AND (:table_name::text IS NULL OR tbl.relname = :table_name) + WHERE (:table_name::text IS NULL OR tbl.relname = :table_name) ORDER BY tn.nspname, tbl.relname, idx.relname; -- name: by_table From 0b779305bc7bddca95d687676c59eed9a0403ab4 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Mon, 2 Nov 2026 01:44:57 +0000 Subject: [PATCH 170/174] fix(db2): apply schema values without SDK DSN truncation --- sqlspec/adapters/db2/pool.py | 79 +++++++++++++--- tests/unit/adapters/test_db2/_fakes.py | 10 ++ tests/unit/adapters/test_db2/test_config.py | 100 ++++++++++++++++++++ 3 files changed, 176 insertions(+), 13 deletions(-) diff --git a/sqlspec/adapters/db2/pool.py b/sqlspec/adapters/db2/pool.py index bfa3a6ac0..b63c28b5b 100644 --- a/sqlspec/adapters/db2/pool.py +++ b/sqlspec/adapters/db2/pool.py @@ -27,6 +27,42 @@ _IBM_DB_DBI = ibm_db_dbi +def _set_current_schema(connection: Any, schema: str) -> None: + """Set the SQL register and SDK catalog schema without CLI keyword quoting.""" + cursor = connection.cursor() + pending_exception: BaseException | None = None + try: + cursor.execute("SET CURRENT SCHEMA = ?", (schema,)) + except BaseException as exc: + pending_exception = exc + try: + cursor.close() + except BaseException as exc: + if pending_exception is None: + pending_exception = exc + if pending_exception is not None: + raise pending_exception + connection.set_current_schema(schema) + + +async def _set_current_schema_async(connection: Any, schema: str) -> None: + """Set both schema states, closing the cursor outside the query's exception handler.""" + cursor = await connection.cursor() + pending_exception: BaseException | None = None + try: + await cursor.execute("SET CURRENT SCHEMA = ?", (schema,)) + except BaseException as exc: + pending_exception = exc + try: + await cursor.close() + except BaseException as exc: + if pending_exception is None: + pending_exception = exc + if pending_exception is not None: + raise pending_exception + await connection.set_current_schema(schema) + + class Db2SyncConnectionPool: """Thread-local connection manager for IBM Db2.""" @@ -62,7 +98,9 @@ def __init__( on_connection_create: Callback executed when connection is created. """ self._connection_parameters = connection_parameters - self._dsn = build_dsn_string(connection_parameters) + native_parameters = dict(connection_parameters) + native_parameters.pop("current_schema", None) + self._dsn = build_dsn_string(native_parameters) self._autocommit = bool(connection_parameters.get("autocommit", True)) self._thread_local = threading.local() self._connection_registry: set[Any] = set() @@ -83,6 +121,9 @@ def new_connection(self) -> Any: The connection opens in the pool's autocommit mode. + The schema is applied with a bound statement because CLI keyword parsing cannot preserve + semicolon-bearing schema values. The SDK catalog schema is synchronized before the hook. + The result is owned by the caller: it is not thread-local and is not tracked for pool shutdown. @@ -101,13 +142,16 @@ def new_connection(self) -> Any: self._dsn, "", "", "", "", {typing_module.ibm_db_dbi.SQL_ATTR_AUTOCOMMIT: autocommit_mode} ) - if self._on_connection_create is not None: - try: + try: + schema = self._connection_parameters.get("current_schema") + if schema is not None: + _set_current_schema(connection, schema) + if self._on_connection_create is not None: self._on_connection_create(connection) - except BaseException: - with contextlib.suppress(Exception): - connection.close() - raise + except BaseException: + with contextlib.suppress(BaseException): + connection.close() + raise return connection @@ -343,7 +387,9 @@ def __init__( an awaitable. """ self._connection_parameters = connection_parameters - self._dsn = build_dsn_string(connection_parameters) + native_parameters = dict(connection_parameters) + native_parameters.pop("current_schema", None) + self._dsn = build_dsn_string(native_parameters) self._autocommit = bool(connection_parameters.get("autocommit", True)) self._max_size = max_size self._acquire_timeout = acquire_timeout @@ -374,6 +420,9 @@ async def new_connection(self) -> Any: The connection opens in the pool's autocommit mode and the creation hook runs on it. The result is owned by the caller and is not tracked by the pool. + The schema is applied with a bound statement and synchronized with the SDK catalog schema + before the hook, avoiding the CLI keyword parser's treatment of semicolons. + Returns: Any: A newly opened ``ibm_db_dbi.AsyncConnection``. @@ -388,14 +437,18 @@ async def new_connection(self) -> Any: connection = await typing_module.ibm_db_dbi.AsyncConnection.connect( self._dsn, "", "", "", "", {typing_module.ibm_db_dbi.SQL_ATTR_AUTOCOMMIT: autocommit_mode} ) - if self._on_connection_create is not None: - try: + try: + schema = self._connection_parameters.get("current_schema") + if schema is not None: + await _set_current_schema_async(connection, schema) + if self._on_connection_create is not None: result = self._on_connection_create(connection) if inspect.isawaitable(result): await result - except BaseException: - await self._close_connection(connection) - raise + except BaseException: + with contextlib.suppress(BaseException): + await connection.close() + raise return connection async def acquire(self) -> Any: diff --git a/tests/unit/adapters/test_db2/_fakes.py b/tests/unit/adapters/test_db2/_fakes.py index ec57278e2..de6b5a3e7 100644 --- a/tests/unit/adapters/test_db2/_fakes.py +++ b/tests/unit/adapters/test_db2/_fakes.py @@ -212,6 +212,7 @@ def __init__( self.dbms_name = "DB2/LINUXX8664" self.autocommit = autocommit self.dsn: str | None = None + self.current_schema: str | None = None self.pending: list[tuple[str, object]] = [] self.committed: list[tuple[str, object]] = [] self.rollbacks = 0 @@ -241,6 +242,11 @@ def set_autocommit(self, is_on: bool) -> None: """Switch the connection's autocommit mode.""" self.autocommit = bool(is_on) + def set_current_schema(self, schema: str) -> bool: + """Update the schema used by SDK catalog operations.""" + self.current_schema = schema + return True + def commit(self) -> None: """Make pending writes durable.""" self.commits += 1 @@ -380,6 +386,10 @@ async def set_autocommit(self, is_on: bool) -> None: """Switch the wrapped connection's autocommit mode.""" self.sync_connection.set_autocommit(is_on) + async def set_current_schema(self, schema: str) -> bool: + """Update the schema used by SDK catalog operations.""" + return self.sync_connection.set_current_schema(schema) + class FakeAsyncConnectionFactory: """Stand-in for the ``ibm_db_dbi.AsyncConnection`` class and its ``connect`` factory.""" diff --git a/tests/unit/adapters/test_db2/test_config.py b/tests/unit/adapters/test_db2/test_config.py index a0dba3b8d..c7446282e 100644 --- a/tests/unit/adapters/test_db2/test_config.py +++ b/tests/unit/adapters/test_db2/test_config.py @@ -1,5 +1,6 @@ """Tests for IBM Db2 database configuration.""" +import asyncio from datetime import datetime, timedelta, timezone from typing import Any, cast @@ -25,6 +26,7 @@ DriverMode, FakeDb2AsyncConnection, FakeDb2Connection, + FakeDb2Cursor, FakeIbmDbDbiModule, FakeIbmDbModule, ) @@ -135,6 +137,104 @@ def test_password_with_semicolon_is_brace_escaped(fake_ibm_db: FakeModules) -> N assert (user, password, host, database) == ("", "", "", "") +@pytest.mark.anyio +@pytest.mark.parametrize("schema", ["APP", "SQLSPEC;SEMICOLON", "Mixed Case"]) +async def test_schema_dsn_roundtrip_is_applied_before_creation_hook( + db2_mode: DriverMode, fake_ibm_db: FakeModules, schema: str +) -> None: + """Schema values reach SQL as one bound value instead of CLI keyword fragments.""" + _, fake_module = fake_ibm_db + cursor = FakeDb2Cursor() + native_connection = FakeDb2Connection(cursors=[cursor]) + fake_module.pending_connections.append(native_connection) + hook_observations: list[tuple[list[tuple[str, Any]], str | None]] = [] + source = db2_mode.config( + connection_config={"database": "d", "user": "u", "password": "p;x", "current_schema": schema} + ) + config = db2_mode.config( + connection_config={"dsn": source.get_connection_string()}, + driver_features={ + "on_connection_create": lambda _: hook_observations.append(( + list(cursor.executed), + native_connection.current_schema, + )) + }, + ) + expected = [("SET CURRENT SCHEMA = ?", (schema,))] + + assert config.get_connection_string() == source.get_connection_string() + connection = await config.create_connection() if db2_mode.is_async else config.create_connection() + try: + assert cursor.executed == expected + assert hook_observations == [(expected, schema)] + assert cursor.closed + assert not native_connection.closed + assert fake_module.connect_calls[0][:5] == ("DATABASE=d;UID=u;PWD={p;x};", "", "", "", "") + finally: + if db2_mode.is_async: + await connection.close() + await config.close_pool() + else: + connection.close() + config.close_pool() + + +@pytest.mark.anyio +@pytest.mark.parametrize("failure", ["execute", "close", "setter", "execute-and-close", "cancel-and-close"]) +async def test_schema_initialization_failure_closes_connection_and_preserves_error( + db2_mode: DriverMode, fake_ibm_db: FakeModules, monkeypatch: pytest.MonkeyPatch, failure: str +) -> None: + """Schema initialization never hands an uninitialized or cancelled connection to the hook.""" + _, fake_module = fake_ibm_db + primary: BaseException = ( + asyncio.CancelledError("schema cancelled") + if failure == "cancel-and-close" + else RuntimeError("schema initialization failed") + ) + cursor = FakeDb2Cursor(error=primary if failure in {"execute", "execute-and-close", "cancel-and-close"} else None) + native_connection = FakeDb2Connection(cursors=[cursor]) + fake_module.pending_connections.append(native_connection) + hook_calls: list[object] = [] + + def close_cursor() -> None: + cursor.closed = True + raise primary if failure == "close" else RuntimeError("schema cursor close failed") + + def set_schema(_: str) -> bool: + raise primary + + original_close = native_connection.close + + def close_connection() -> None: + original_close() + raise RuntimeError("connection close failed") + + if failure in {"close", "execute-and-close", "cancel-and-close"}: + monkeypatch.setattr(cursor, "close", close_cursor) + if failure == "setter": + monkeypatch.setattr(native_connection, "set_current_schema", set_schema) + monkeypatch.setattr(native_connection, "close", close_connection) + config = db2_mode.config( + connection_config={"database": "d", "current_schema": "SQLSPEC;SEMICOLON"}, + driver_features={"on_connection_create": hook_calls.append}, + ) + try: + with pytest.raises(type(primary)) as caught: + if db2_mode.is_async: + await config.create_connection() + else: + config.create_connection() + assert caught.value is primary + assert cursor.closed + assert native_connection.closed + assert hook_calls == [] + finally: + if db2_mode.is_async: + await config.close_pool() + else: + config.close_pool() + + @pytest.mark.parametrize( ("value", "rendered"), [("a{b", "{a{b}"), (" padded", "{ padded}"), (True, "1"), (False, "0"), ("plain", "plain")], From 999192334ff34a9300f1efaff6b5a66ff284190d Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Mon, 2 Nov 2026 01:45:20 +0000 Subject: [PATCH 171/174] fix(ci): keep Sybase SDK optional in clean wheel smoke --- tools/scripts/mypyc_smoke.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tools/scripts/mypyc_smoke.py b/tools/scripts/mypyc_smoke.py index e0c0550d9..f25bca47c 100644 --- a/tools/scripts/mypyc_smoke.py +++ b/tools/scripts/mypyc_smoke.py @@ -44,7 +44,7 @@ class SmokeImport(NamedTuple): SmokeImport("storage_backend_fsspec", "sqlspec.storage.backends.fsspec", "FSSpecBackend", True), SmokeImport("storage_backend_obstore", "sqlspec.storage.backends.obstore", "ObStoreBackend", True), SmokeImport("sqlite_pool", "sqlspec.adapters.sqlite.pool", "SqliteConnectionPool", True), - SmokeImport("sybase_pool", "sqlspec.adapters.sybase.pool", "SybaseConnectionPool", True), + SmokeImport("sybase_pool", "sqlspec.adapters.sybase.pool", "SybaseConnectionPool", True, "pymssql"), SmokeImport("sqlanywhere_pool", "sqlspec.adapters.sqlanywhere.pool", "SqlanywhereConnectionPool", True, "sqlanydb"), SmokeImport("data_dictionary_registry", "sqlspec.data_dictionary._registry", "get_dialect_config", True), SmokeImport("data_dictionary_loader", "sqlspec.data_dictionary._loader", "DataDictionaryLoader", True), From 2e90deac439c23e26a2d7fae4fe99dcae399a4a8 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Mon, 2 Nov 2026 01:47:41 +0000 Subject: [PATCH 172/174] test(db2): resolve physical Python paths in ODBC runtime --- tests/integration/fixtures/db2.py | 77 ++++++++++++++++++++++++++++--- 1 file changed, 71 insertions(+), 6 deletions(-) diff --git a/tests/integration/fixtures/db2.py b/tests/integration/fixtures/db2.py index 07df60854..af4431aeb 100644 --- a/tests/integration/fixtures/db2.py +++ b/tests/integration/fixtures/db2.py @@ -1,10 +1,14 @@ """IBM Db2 integration test fixtures.""" +import json import os import platform import shlex import sys +import sysconfig from collections.abc import AsyncGenerator, Generator +from importlib.metadata import version +from importlib.util import find_spec from pathlib import Path from typing import TYPE_CHECKING, Any @@ -171,25 +175,79 @@ def db2_arrow_odbc_runtime(db2_service: Db2Service) -> tuple[list[str], dict[str mounts. IBM's installation registry and instance profile remain inside the existing Db2 service container; no host driver registration is changed. """ + python_base = Path(sys.base_prefix).resolve() + executable = Path(sys.executable).resolve() + package_paths = tuple( + dict.fromkeys(str(Path(sysconfig.get_path(kind)).resolve()) for kind in ("purelib", "platlib")) + ) library_path = ":".join(( "/opt/sqlspec-host-libraries", - str(Path(sys.base_prefix) / "lib"), + str(python_base / "lib"), "/opt/ibm/db2/V11.5/lib64", "/opt/ibm/db2/V11.5/lib64/gskit", )) driver = "/opt/ibm/db2/V11.5/lib64/libdb2o.so.1" - if not _db2_odbc_client_mounts(): + client_mounts = _db2_odbc_client_mounts() + if not client_mounts: msg = "Arrow-ODBC Db2 acceptance requires a Linux x86-64 glibc host with its dynamic loader available." raise RuntimeError(msg) + db2_service.container.reload() + actual_mounts = {item["Destination"]: item for item in db2_service.container.attrs.get("Mounts", ())} + missing_mounts = sorted( + item["bind"] + for item in client_mounts.values() + if item["bind"] not in actual_mounts or actual_mounts[item["bind"]].get("RW", True) + ) + if missing_mounts: + msg = ( + "The existing Db2 service lacks required read-only client mounts; cached containers cannot " + f"acquire new mounts in place. Missing or writable destinations: {missing_mounts!r}" + ) + raise RuntimeError(msg) + result = db2_service.container.exec_run(["test", "-r", str(executable)]) + if result.exit_code: + msg = f"The mounted canonical Python executable is not readable inside Db2: {str(executable)!r}" + raise RuntimeError(msg) result = db2_service.container.exec_run(["test", "-r", driver]) if result.exit_code: msg = f"The pinned Db2 service image does not provide the required IBM 64-bit ODBC driver: {driver}" raise RuntimeError(msg) - command = ["/opt/sqlspec-host-libraries/ld-linux-x86-64.so.2", "--library-path", library_path, sys.executable] + command = ["/opt/sqlspec-host-libraries/ld-linux-x86-64.so.2", "--library-path", library_path, str(executable)] environment = { "LD_LIBRARY_PATH": "/opt/ibm/db2/V11.5/lib64:/opt/ibm/db2/V11.5/lib64/gskit", "DB2INSTANCE": db2_service.username, + "PYTHONHOME": str(python_base), + "PYTHONPATH": ":".join(dict.fromkeys((str(Path(sqlspec.__file__).resolve().parent.parent), *package_paths))), } + probe = ( + "import json,sys,sqlspec,pyarrow,arrow_odbc; from importlib.metadata import version; " + "from sqlspec.adapters.arrow_odbc import ArrowOdbcConfig; " + "from pathlib import Path; " + "print(json.dumps({'python':list(sys.version_info[:3]), 'sqlspec':str(Path(sqlspec.__file__).resolve()), " + "'pyarrow':str(Path(pyarrow.__file__).resolve()), 'arrow_odbc':str(Path(arrow_odbc.__file__).resolve()), " + "'versions':{name:version(name) for name in ('sqlspec','pyarrow','arrow-odbc')}}))" + ) + result = db2_service.container.exec_run([*command, "-c", probe], environment=environment) + output = exec_output_to_bytes(result.output).decode("utf-8", errors="replace") + if result.exit_code: + msg = f"The mounted Db2 client Python environment failed initialization (exit {result.exit_code}):\n{output}" + raise RuntimeError(msg) + expected_origins: dict[str, str] = {} + for module_name in ("pyarrow", "arrow_odbc"): + spec = find_spec(module_name) + if spec is None or spec.origin is None: + msg = f"The host environment lacks the required {module_name!r} package origin." + raise RuntimeError(msg) + expected_origins[module_name] = str(Path(spec.origin).resolve()) + expected = { + "python": list(sys.version_info[:3]), + "sqlspec": str(Path(sqlspec.__file__).resolve()), + **expected_origins, + "versions": {name: version(name) for name in ("sqlspec", "pyarrow", "arrow-odbc")}, + } + if json.loads(output) != expected: + msg = f"The mounted Db2 client imported a different Python/package environment: {output}" + raise RuntimeError(msg) return command, environment, driver @@ -231,9 +289,11 @@ def _db2_odbc_client_mounts() -> dict[str, dict[str, str]]: """Expose only the Python client, SQLSpec package, and host loader read-only.""" if sys.platform != "linux" or platform.machine().lower() not in {"x86_64", "amd64"}: return {} + python_base = Path(sys.base_prefix).resolve() + executable_base = Path(sys.executable).resolve().parent.parent if any( - Path(prefix).resolve() in {Path("/"), Path("/usr"), Path("/usr/local")} - for prefix in (sys.base_prefix, sys.prefix) + prefix in {Path("/"), Path("/usr"), Path("/usr/local")} + for prefix in (python_base, executable_base, Path(sys.prefix).resolve()) ): return {} system_libraries = next( @@ -242,7 +302,12 @@ def _db2_odbc_client_mounts() -> dict[str, dict[str, str]]: ) if system_libraries is None: return {} - paths = {Path(sys.base_prefix), Path(sys.prefix), Path(sqlspec.__file__).resolve().parent} + paths = { + python_base, + executable_base, + *(Path(sysconfig.get_path(kind)).resolve() for kind in ("purelib", "platlib")), + Path(sqlspec.__file__).resolve().parent, + } mounts = {str(path): {"bind": str(path), "mode": "ro"} for path in paths} mounts[str(system_libraries.resolve())] = {"bind": "/opt/sqlspec-host-libraries", "mode": "ro"} return mounts From 3d5ae9d9075575c2d26188b0746276784c944161 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Mon, 2 Nov 2026 02:16:14 +0000 Subject: [PATCH 173/174] test(db2): enable native injected-pool lifecycle contracts --- tests/integration/adapters/_shared/_cases.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/integration/adapters/_shared/_cases.py b/tests/integration/adapters/_shared/_cases.py index 08d042396..15e61dd10 100644 --- a/tests/integration/adapters/_shared/_cases.py +++ b/tests/integration/adapters/_shared/_cases.py @@ -156,6 +156,7 @@ def _db2_case(mode: Literal["sync", "async"], marks: tuple[Mark | MarkDecorator, supports_merge=True, supports_native_row_streaming=True, supports_pooling=True, + supports_connection_instance=True, supports_connection_hook=True, supports_lowercase_columns=True, config_factory_fixture=f"lifecycle_config_db2_{mode}", From c5e844bd5c70106938da4192d38a205f0436a3fb Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Mon, 2 Nov 2026 02:37:17 +0000 Subject: [PATCH 174/174] test(adapters): use dialect-aware injected-pool probes --- tests/integration/adapters/_shared/behaviors.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/integration/adapters/_shared/behaviors.py b/tests/integration/adapters/_shared/behaviors.py index bb912ae90..8cfe6cc46 100644 --- a/tests/integration/adapters/_shared/behaviors.py +++ b/tests/integration/adapters/_shared/behaviors.py @@ -5028,7 +5028,7 @@ def _pool_contract_table(case: DriverCase) -> str: def _connection_probe_sql(case: DriverCase) -> str: - """Trivial query used to force a connection to be drawn so the create-hook fires.""" + """Scalar query that acquires a native connection for lifecycle contracts.""" if case.dialect == "oracle": return "SELECT 1 FROM DUAL" if case.dialect == "db2": @@ -5132,7 +5132,7 @@ def assert_sync_connection_instance_contract(make_config: SyncConfigFactory, cas try: assert fresh.provide_pool() is not None with fresh.provide_session() as session: - assert session.select_value("SELECT 1") == 1 + assert session.select_value(_connection_probe_sql(case)) == 1 finally: fresh.close_pool() @@ -5163,7 +5163,7 @@ async def assert_async_connection_instance_contract(make_config: AsyncConfigFact try: assert await fresh.provide_pool() is not None async with fresh.provide_session() as session: - assert await session.select_value("SELECT 1") == 1 + assert await session.select_value(_connection_probe_sql(case)) == 1 finally: await fresh.close_pool()