diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml index 35e70a2e0..d2743d630 100644 --- a/.github/workflows/publish.yml +++ b/.github/workflows/publish.yml @@ -132,7 +132,7 @@ jobs: # hatchling versions track the locked dev environment. CIBW_BEFORE_BUILD: "pip install -c {project}/build-constraints.txt hatch-mypyc hatchling" - CIBW_TEST_REQUIRES: "cloud-sql-python-connector google-cloud-alloydb-connector" + CIBW_TEST_REQUIRES: "aiosqlite cloud-sql-python-connector google-cloud-alloydb-connector" CIBW_TEST_COMMAND: >- python {project}/tools/scripts/mypyc_smoke.py --require-compiled diff --git a/.github/workflows/test-build.yml b/.github/workflows/test-build.yml index 76781dc8d..8f23208b1 100644 --- a/.github/workflows/test-build.yml +++ b/.github/workflows/test-build.yml @@ -311,7 +311,7 @@ jobs: CIBW_BEFORE_BUILD: "pip install -c {project}/build-constraints.txt hatch-mypyc hatchling" - CIBW_TEST_REQUIRES: "cloud-sql-python-connector google-cloud-alloydb-connector" + CIBW_TEST_REQUIRES: "aiosqlite cloud-sql-python-connector google-cloud-alloydb-connector" CIBW_TEST_COMMAND: >- python {project}/tools/scripts/mypyc_smoke.py --require-compiled diff --git a/docs/changelog.rst b/docs/changelog.rst index be5e6297b..30bab08a6 100644 --- a/docs/changelog.rst +++ b/docs/changelog.rst @@ -9,6 +9,14 @@ important operational fixes. Recent Updates ============== +Unreleased - Compiled async exception handling +------------------------------------------------------------------------------ + +**Fixed:** + +* Async statement errors from mypyc-compiled drivers are translated into + :class:`~sqlspec.exceptions.SQLSpecError` instead of terminating the process. + v0.58.2 - SQL file parameter diagnostics ------------------------------------------------------------------------------ diff --git a/sqlspec/driver/_async.py b/sqlspec/driver/_async.py index 4f80499c6..b770f73b5 100644 --- a/sqlspec/driver/_async.py +++ b/sqlspec/driver/_async.py @@ -230,24 +230,10 @@ async def dispatch_statement_execution(self, statement: "SQL", connection: "Any" # FAST PATH: Skip all instrumentation if runtime is absent or idle. if runtime is None or runtime.is_idle: exc_handler = self.handle_database_exceptions() - async with exc_handler, self.with_cursor(connection) as cursor: - special_result = await self.dispatch_special_handling(cursor, statement) - if special_result is not None: - result = special_result - elif statement.is_script: - execution_result = await self.dispatch_execute_script(cursor, statement) - result = self.build_statement_result(statement, execution_result) - elif statement.is_many: - if execution_parameters: - execution_result = await self.dispatch_execute_many(cursor, statement) - else: - execution_result = self.create_execution_result( - cursor, rowcount_override=0, is_many_result=True - ) - result = self.build_statement_result(statement, execution_result) - else: - execution_result = await self.dispatch_execute(cursor, statement) - result = self.build_statement_result(statement, execution_result) + async with exc_handler: + result = await self._dispatch_statement_with_cursor( + connection, statement, has_execution_parameters=bool(execution_parameters) + ) self._check_pending_exception(exc_handler) assert result is not None return result @@ -272,24 +258,10 @@ async def dispatch_statement_execution(self, statement: "SQL", connection: "Any" exc_handler = self.handle_database_exceptions() try: - async with exc_handler, self.with_cursor(connection) as cursor: - special_result = await self.dispatch_special_handling(cursor, statement) - if special_result is not None: - result = special_result - elif statement.is_script: - execution_result = await self.dispatch_execute_script(cursor, statement) - result = self.build_statement_result(statement, execution_result) - elif statement.is_many: - if execution_parameters: - execution_result = await self.dispatch_execute_many(cursor, statement) - else: - execution_result = self.create_execution_result( - cursor, rowcount_override=0, is_many_result=True - ) - result = self.build_statement_result(statement, execution_result) - else: - execution_result = await self.dispatch_execute(cursor, statement) - result = self.build_statement_result(statement, execution_result) + async with exc_handler: + result = await self._dispatch_statement_with_cursor( + connection, statement, has_execution_parameters=bool(execution_parameters) + ) except Exception as exc: # pragma: no cover pending_exception = exc_handler.pending_exception if pending_exception is not None: @@ -459,76 +431,13 @@ async def _execute_cache_hit( Returns: SQLResult or DMLResult. """ - direct_statement: SQL | None = None exc_handler = self.handle_database_exceptions() result: SQLResult | None = None - try: - async with exc_handler, self.with_cursor(self.connection) as cursor: - execute = getattr(cursor, "execute", None) - fetchall = getattr(cursor, "fetchall", None) - can_use_cursor_fast_path = execute is not None and ( - (cached.operation_profile.returns_rows and fetchall is not None) - or (not cached.operation_profile.returns_rows and hasattr(cursor, "rowcount")) - ) - if can_use_cursor_fast_path: - assert execute is not None - try: - execute_result = execute(cached.compiled_sql, params) - if isawaitable(execute_result): - await execute_result - if cached.operation_profile.returns_rows: - assert fetchall is not None - fetched_data = fetchall() - if isawaitable(fetched_data): - fetched_data = await fetched_data - data, column_names, row_count = self.collect_rows(cursor, fetched_data) - execution_result = self.create_execution_result( - cursor, - selected_data=data, - column_names=column_names, - data_row_count=row_count, - is_select_result=True, - row_format="tuple", - ) - direct_statement = self._cached_statement( - sql, - params, - cached, - params, - params_are_simple=True, - compiled_sql=cached.compiled_sql, - ) - result = self.build_statement_result(direct_statement, execution_result) - else: - affected_rows = self.resolve_rowcount(cursor) - result = DMLResult(cached.operation_type, affected_rows) - except (AttributeError, NotImplementedError): - pass - - if result is None: - direct_statement = self._cached_statement( - sql, params, cached, params, params_are_simple=True, compiled_sql=cached.compiled_sql - ) - execution_result = await self.dispatch_execute(cursor, direct_statement) - - if cached.operation_profile.returns_rows: - result = self.build_statement_result(direct_statement, execution_result) - else: - # DML path: use DMLResult to bypass full SQLResult construction - affected_rows = ( - execution_result.rowcount_override - if execution_result.rowcount_override is not None - and execution_result.rowcount_override >= 0 - else 0 - ) - result = DMLResult(cached.operation_type, affected_rows) - - self._check_pending_exception(exc_handler) - assert result is not None - return result - finally: - if direct_statement is not None: - self._release_pooled_statement(direct_statement) + async with exc_handler: + result = await self._execute_cache_hit_with_cursor(sql, params, cached) + self._check_pending_exception(exc_handler) + assert result is not None + return result async def _cached_execution( self, statement: str, params: "tuple[Any, ...] | list[Any] | dict[str, Any]" @@ -556,9 +465,8 @@ async def _execute_cached_statement(self, statement: "SQL") -> "SQLResult": exc_handler = self.handle_database_exceptions() result: SQLResult | None = None try: - async with exc_handler, self.with_cursor(self.connection) as cursor: - execution_result = await self.dispatch_execute(cursor, statement) - result = self.build_statement_result(statement, execution_result) + async with exc_handler: + result = await self._execute_cached_statement_with_cursor(statement) self._check_pending_exception(exc_handler) assert result is not None @@ -1761,6 +1669,123 @@ def _connection_in_transaction(self) -> bool: msg = "Adapters must override _connection_in_transaction()" raise NotImplementedError(msg) + async def _dispatch_statement_with_cursor( + self, connection: Any, statement: "SQL", *, has_execution_parameters: bool + ) -> "SQLResult": + """Execute a statement while owning only the cursor context.""" + cursor_manager = self.with_cursor(connection) + cursor_entered = False + exit_suppressed = False + error: Exception | None = None + execution_result: ExecutionResult | None = None + result: SQLResult | None = None + try: + cursor = await cursor_manager.__aenter__() + cursor_entered = True + special_result = await self.dispatch_special_handling(cursor, statement) + if special_result is not None: + result = special_result + elif statement.is_script: + execution_result = await self.dispatch_execute_script(cursor, statement) + elif statement.is_many: + if has_execution_parameters: + execution_result = await self.dispatch_execute_many(cursor, statement) + else: + execution_result = self.create_execution_result(cursor, rowcount_override=0, is_many_result=True) + else: + execution_result = await self.dispatch_execute(cursor, statement) + if special_result is None: + assert execution_result is not None + result = self.build_statement_result(statement, execution_result) + except Exception as exc: + error = exc + finally: + if cursor_entered: + if error is None: + await cursor_manager.__aexit__(None, None, None) + else: + exit_suppressed = bool(await cursor_manager.__aexit__(type(error), error, error.__traceback__)) + if error is not None and not exit_suppressed: + raise error + assert result is not None + return result + + async 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.""" + direct_statement: SQL | None = None + result: SQLResult | None = None + try: + async with self.with_cursor(self.connection) as cursor: + execute = getattr(cursor, "execute", None) + fetchall = getattr(cursor, "fetchall", None) + can_use_cursor_fast_path = execute is not None and ( + (cached.operation_profile.returns_rows and fetchall is not None) + or (not cached.operation_profile.returns_rows and hasattr(cursor, "rowcount")) + ) + if can_use_cursor_fast_path: + assert execute is not None + try: + execute_result = execute(cached.compiled_sql, params) + if isawaitable(execute_result): + await execute_result + if cached.operation_profile.returns_rows: + assert fetchall is not None + fetched_data = fetchall() + if isawaitable(fetched_data): + fetched_data = await fetched_data + data, column_names, row_count = self.collect_rows(cursor, fetched_data) + execution_result = self.create_execution_result( + cursor, + selected_data=data, + column_names=column_names, + data_row_count=row_count, + is_select_result=True, + row_format="tuple", + ) + direct_statement = self._cached_statement( + sql, + params, + cached, + params, + params_are_simple=True, + compiled_sql=cached.compiled_sql, + ) + result = self.build_statement_result(direct_statement, execution_result) + else: + affected_rows = self.resolve_rowcount(cursor) + result = DMLResult(cached.operation_type, affected_rows) + except (AttributeError, NotImplementedError): + pass + + if result is None: + direct_statement = self._cached_statement( + sql, params, cached, params, params_are_simple=True, compiled_sql=cached.compiled_sql + ) + execution_result = await self.dispatch_execute(cursor, direct_statement) + if cached.operation_profile.returns_rows: + result = self.build_statement_result(direct_statement, execution_result) + else: + affected_rows = ( + execution_result.rowcount_override + if execution_result.rowcount_override is not None + and execution_result.rowcount_override >= 0 + else 0 + ) + result = DMLResult(cached.operation_type, affected_rows) + assert result is not None + return result + finally: + if direct_statement is not None: + self._release_pooled_statement(direct_statement) + + async def _execute_cached_statement_with_cursor(self, statement: "SQL") -> "SQLResult": + """Execute a prepared statement while owning only the cursor context.""" + async with self.with_cursor(self.connection) as cursor: + execution_result = await self.dispatch_execute(cursor, statement) + return self.build_statement_result(statement, execution_result) + async def _execute_stack_operation(self, operation: "StackOperation") -> "SQLResult | ArrowResult | None": kwargs = dict(operation.keyword_arguments) if operation.keyword_arguments else {} diff --git a/tests/unit/utils/test_mypyc_smoke.py b/tests/unit/utils/test_mypyc_smoke.py index ffa17f4b9..4c0e1a703 100644 --- a/tests/unit/utils/test_mypyc_smoke.py +++ b/tests/unit/utils/test_mypyc_smoke.py @@ -141,6 +141,7 @@ def test_construction_checks_build_provider_signatures_without_requiring_compila assert all(result["imported"] or result["skipped"] for result in results) assert {result["name"] for result in results} == { + "aiosqlite_exception_mapping", "fastapi_filter_construction", "litestar_filter_construction", "statement_cache_rebind", diff --git a/tools/scripts/mypyc_smoke.py b/tools/scripts/mypyc_smoke.py index b98eb676f..3cf77ae0e 100644 --- a/tools/scripts/mypyc_smoke.py +++ b/tools/scripts/mypyc_smoke.py @@ -3,6 +3,7 @@ import argparse import importlib import json +import subprocess import sys from collections.abc import Sequence from typing import Any, NamedTuple @@ -176,6 +177,87 @@ def _check_statement_cache_rebind(*, require_compiled: bool = False) -> dict[str return result +def _check_aiosqlite_exception_mapping(*, require_compiled: bool = False) -> dict[str, Any]: + """Exercise mapped aiosqlite statement errors in isolated child processes.""" + result = _new_smoke_result( + name="aiosqlite_exception_mapping", + module="sqlspec.driver._async", + attribute="AsyncDriverAdapterBase", + compiled_required=require_compiled, + ) + try: + async_driver_module = importlib.import_module("sqlspec.driver._async") + importlib.import_module("aiosqlite") + except ModuleNotFoundError as exc: + if _is_missing_optional_dependency(exc.name or "", "aiosqlite"): + result["skipped"] = True + result["skip_reason"] = "optional dependency missing: aiosqlite" + return result + result["error"] = f"{type(exc).__name__}: {exc}" + return result + except Exception as exc: + result["error"] = f"{type(exc).__name__}: {exc}" + return result + + result["imported"] = True + result["compiled"] = is_compiled_module(async_driver_module) + if require_compiled and not result["compiled"]: + result["error"] = "module was imported from Python source, not a compiled extension" + return result + + child_script = """ +import asyncio +import sys + +from sqlspec import SQLSpec +from sqlspec.adapters.aiosqlite import AiosqliteConfig +from sqlspec.exceptions import SQLSpecError + + +async def main() -> None: + config = AiosqliteConfig() + spec = SQLSpec() + spec.add_config(config) + try: + async with spec.provide_session(config) as driver: + await driver.execute("CREATE TABLE smoke_items (id INTEGER PRIMARY KEY)") + operation = getattr(driver, sys.argv[1]) + try: + await operation("SELECT missing_column FROM smoke_items") + except SQLSpecError as exc: + if "missing_column" not in str(exc): + raise + else: + raise AssertionError("invalid statement did not raise SQLSpecError") + finally: + await config.close_pool() + + +asyncio.run(main()) +print(f"{sys.argv[1]}:SQLSpecError") +""" + for operation_name in ("select", "execute"): + try: + completed = subprocess.run( + [sys.executable, "-I", "-c", child_script, operation_name], + capture_output=True, + check=False, + text=True, + timeout=30, + ) + except subprocess.TimeoutExpired: + result["error"] = f"{operation_name} exception subprocess timed out" + return result + expected_marker = f"{operation_name}:SQLSpecError" + if completed.returncode != 0 or expected_marker not in completed.stdout: + result["error"] = ( + f"{operation_name} exception subprocess failed with return code {completed.returncode}; " + f"stdout={completed.stdout!r}; stderr={completed.stderr!r}" + ) + return result + return result + + def run_smoke(*, require_compiled: bool = False) -> list[dict[str, Any]]: """Import the compiled-wheel smoke matrix and return per-entry results.""" results: list[dict[str, Any]] = [] @@ -349,6 +431,7 @@ def run_construction_checks(*, require_compiled: bool = False) -> list[dict[str, _check_sqlspec_construction(), _check_statement_sentinel_identity(require_compiled=require_compiled), _check_statement_cache_rebind(require_compiled=require_compiled), + _check_aiosqlite_exception_mapping(require_compiled=require_compiled), _check_fastapi_filter_construction(require_compiled=require_compiled), _check_litestar_filter_construction(require_compiled=require_compiled), ]