feat: improve config typing (#221)

Updated extension config type for all drivers

Author: cofin

sqlspec/adapters/adbc/config.py

@@ -9,7 +9,7 @@
 
 from sqlspec.adapters.adbc._types import AdbcConnection
 from sqlspec.adapters.adbc.driver import AdbcCursor, AdbcDriver, get_adbc_statement_config
-from sqlspec.config import NoPoolSyncConfig
+from sqlspec.config import ADKConfig, FastAPIConfig, FlaskConfig, LitestarConfig, NoPoolSyncConfig, StarletteConfig
 from sqlspec.core.statement import StatementConfig
 from sqlspec.exceptions import ImproperConfigurationError
 from sqlspec.utils.module_loader import import_string
@@ -110,7 +110,7 @@ def __init__(
         statement_config: StatementConfig | None = None,
         driver_features: "AdbcDriverFeatures | dict[str, Any] | None" = None,
         bind_key: str | None = None,
-        extension_config: "dict[str, dict[str, Any]] | None" = None,
+        extension_config: "dict[str, dict[str, Any]] | LitestarConfig | FastAPIConfig | StarletteConfig | FlaskConfig | ADKConfig | None" = None,
     ) -> None:
         """Initialize configuration.
 

sqlspec/adapters/aiosqlite/config.py

@@ -14,7 +14,7 @@
     AiosqlitePoolClosedError,
     AiosqlitePoolConnection,
 )
-from sqlspec.config import AsyncDatabaseConfig
+from sqlspec.config import ADKConfig, AsyncDatabaseConfig, FastAPIConfig, FlaskConfig, LitestarConfig, StarletteConfig
 from sqlspec.utils.serializers import from_json, to_json
 
 if TYPE_CHECKING:
@@ -84,7 +84,7 @@ def __init__(
         statement_config: "StatementConfig | None" = None,
         driver_features: "AiosqliteDriverFeatures | dict[str, Any] | None" = None,
         bind_key: "str | None" = None,
-        extension_config: "dict[str, dict[str, Any]] | None" = None,
+        extension_config: "dict[str, dict[str, Any]] | LitestarConfig | FastAPIConfig | StarletteConfig | FlaskConfig | ADKConfig | None" = None,
     ) -> None:
         """Initialize AioSQLite configuration.
 

sqlspec/adapters/asyncmy/config.py

@@ -12,7 +12,7 @@
 
 from sqlspec.adapters.asyncmy._types import AsyncmyConnection
 from sqlspec.adapters.asyncmy.driver import AsyncmyCursor, AsyncmyDriver, asyncmy_statement_config
-from sqlspec.config import AsyncDatabaseConfig
+from sqlspec.config import ADKConfig, AsyncDatabaseConfig, FastAPIConfig, FlaskConfig, LitestarConfig, StarletteConfig
 from sqlspec.utils.serializers import from_json, to_json
 
 if TYPE_CHECKING:
@@ -94,7 +94,7 @@ def __init__(
         statement_config: "StatementConfig | None" = None,
         driver_features: "AsyncmyDriverFeatures | dict[str, Any] | None" = None,
         bind_key: "str | None" = None,
-        extension_config: "dict[str, dict[str, Any]] | None" = None,
+        extension_config: "dict[str, dict[str, Any]] | LitestarConfig | FastAPIConfig | StarletteConfig | FlaskConfig | ADKConfig | None" = None,
     ) -> None:
         """Initialize Asyncmy configuration.
 

sqlspec/adapters/asyncpg/config.py

@@ -13,7 +13,7 @@
 
 from sqlspec.adapters.asyncpg._types import AsyncpgConnection
 from sqlspec.adapters.asyncpg.driver import AsyncpgCursor, AsyncpgDriver, asyncpg_statement_config
-from sqlspec.config import AsyncDatabaseConfig
+from sqlspec.config import ADKConfig, AsyncDatabaseConfig, FastAPIConfig, FlaskConfig, LitestarConfig, StarletteConfig
 from sqlspec.typing import PGVECTOR_INSTALLED
 from sqlspec.utils.serializers import from_json, to_json
 
@@ -107,7 +107,7 @@ def __init__(
         statement_config: "StatementConfig | None" = None,
         driver_features: "AsyncpgDriverFeatures | dict[str, Any] | None" = None,
         bind_key: "str | None" = None,
-        extension_config: "dict[str, dict[str, Any]] | None" = None,
+        extension_config: "dict[str, dict[str, Any]] | LitestarConfig | FastAPIConfig | StarletteConfig | FlaskConfig | ADKConfig | None" = None,
     ) -> None:
         """Initialize AsyncPG configuration.
 

sqlspec/adapters/bigquery/config.py

@@ -9,7 +9,7 @@
 
 from sqlspec.adapters.bigquery._types import BigQueryConnection
 from sqlspec.adapters.bigquery.driver import BigQueryCursor, BigQueryDriver, bigquery_statement_config
-from sqlspec.config import NoPoolSyncConfig
+from sqlspec.config import ADKConfig, FastAPIConfig, FlaskConfig, LitestarConfig, NoPoolSyncConfig, StarletteConfig
 from sqlspec.exceptions import ImproperConfigurationError
 from sqlspec.typing import Empty
 
@@ -109,7 +109,7 @@ def __init__(
         statement_config: "StatementConfig | None" = None,
         driver_features: "BigQueryDriverFeatures | dict[str, Any] | None" = None,
         bind_key: "str | None" = None,
-        extension_config: "dict[str, dict[str, Any]] | None" = None,
+        extension_config: "dict[str, dict[str, Any]] | LitestarConfig | FastAPIConfig | StarletteConfig | FlaskConfig | ADKConfig | None" = None,
     ) -> None:
         """Initialize BigQuery configuration.
 

sqlspec/adapters/duckdb/config.py

@@ -9,7 +9,7 @@
 from sqlspec.adapters.duckdb._types import DuckDBConnection
 from sqlspec.adapters.duckdb.driver import DuckDBCursor, DuckDBDriver, duckdb_statement_config
 from sqlspec.adapters.duckdb.pool import DuckDBConnectionPool
-from sqlspec.config import SyncDatabaseConfig
+from sqlspec.config import ADKConfig, FastAPIConfig, FlaskConfig, LitestarConfig, StarletteConfig, SyncDatabaseConfig
 
 if TYPE_CHECKING:
     from collections.abc import Callable, Generator
@@ -187,7 +187,7 @@ def __init__(
         statement_config: "StatementConfig | None" = None,
         driver_features: "DuckDBDriverFeatures | dict[str, Any] | None" = None,
         bind_key: "str | None" = None,
-        extension_config: "dict[str, dict[str, Any]] | None" = None,
+        extension_config: "dict[str, dict[str, Any]] | LitestarConfig | FastAPIConfig | StarletteConfig | FlaskConfig | ADKConfig | None" = None,
     ) -> None:
         """Initialize DuckDB configuration.
 

sqlspec/adapters/oracledb/config.py

@@ -23,7 +23,15 @@
     oracledb_statement_config,
 )
 from sqlspec.adapters.oracledb.migrations import OracleAsyncMigrationTracker, OracleSyncMigrationTracker
-from sqlspec.config import AsyncDatabaseConfig, SyncDatabaseConfig
+from sqlspec.config import (
+    ADKConfig,
+    AsyncDatabaseConfig,
+    FastAPIConfig,
+    FlaskConfig,
+    LitestarConfig,
+    StarletteConfig,
+    SyncDatabaseConfig,
+)
 from sqlspec.typing import NUMPY_INSTALLED
 
 if TYPE_CHECKING:
@@ -121,7 +129,7 @@ def __init__(
         statement_config: "StatementConfig | None" = None,
         driver_features: "OracleDriverFeatures | dict[str, Any] | None" = None,
         bind_key: "str | None" = None,
-        extension_config: "dict[str, dict[str, Any]] | None" = None,
+        extension_config: "dict[str, dict[str, Any]] | LitestarConfig | FastAPIConfig | StarletteConfig | FlaskConfig | ADKConfig | None" = None,
     ) -> None:
         """Initialize Oracle synchronous configuration.
 
@@ -281,7 +289,7 @@ def __init__(
         statement_config: "StatementConfig | None" = None,
         driver_features: "OracleDriverFeatures | dict[str, Any] | None" = None,
         bind_key: "str | None" = None,
-        extension_config: "dict[str, dict[str, Any]] | None" = None,
+        extension_config: "dict[str, dict[str, Any]] | LitestarConfig | FastAPIConfig | StarletteConfig | FlaskConfig | ADKConfig | None" = None,
     ) -> None:
         """Initialize Oracle asynchronous configuration.
 

sqlspec/adapters/oracledb/driver.py

@@ -202,7 +202,7 @@ async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
                 self.cursor.close()
 
 
-class OracleSyncExceptionHandler:
+class OracleExceptionHandler:
     """Context manager for handling Oracle database exceptions.
 
     Maps Oracle ORA-XXXXX error codes to specific SQLSpec exceptions
@@ -211,16 +211,6 @@ class OracleSyncExceptionHandler:
 
     __slots__ = ()
 
-    def __enter__(self) -> None:
-        return None
-
-    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
-        _ = exc_tb
-        if exc_type is None:
-            return
-        if issubclass(exc_type, oracledb.DatabaseError):
-            self._map_oracle_exception(exc_val)
-
     def _map_oracle_exception(self, e: Any) -> None:
         """Map Oracle exception to SQLSpec exception.
 
@@ -304,106 +294,40 @@ def _raise_generic_error(self, e: Any, code: "int | None") -> None:
         raise SQLSpecError(msg) from e
 
 
-class OracleAsyncExceptionHandler:
-    """Async context manager for handling Oracle database exceptions.
+class OracleSyncExceptionHandler(OracleExceptionHandler):
+    """Sync Context manager for handling Oracle database exceptions.
 
     Maps Oracle ORA-XXXXX error codes to specific SQLSpec exceptions
     for better error handling in application code.
     """
 
-    __slots__ = ()
-
-    async def __aenter__(self) -> None:
+    def __enter__(self) -> None:
         return None
 
-    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
         _ = exc_tb
         if exc_type is None:
             return
         if issubclass(exc_type, oracledb.DatabaseError):
             self._map_oracle_exception(exc_val)
 
-    def _map_oracle_exception(self, e: Any) -> None:
-        """Map Oracle exception to SQLSpec exception.
-
-        Args:
-            e: oracledb.DatabaseError instance
-        """
-        error_obj = e.args[0] if e.args else None
-        if not error_obj:
-            self._raise_generic_error(e, None)
-
-        error_code = getattr(error_obj, "code", None)
-
-        if not error_code:
-            self._raise_generic_error(e, None)
-
-        if error_code == 1:
-            self._raise_unique_violation(e, error_code)
-        elif error_code in {2291, 2292}:
-            self._raise_foreign_key_violation(e, error_code)
-        elif error_code == ORA_CHECK_CONSTRAINT:
-            self._raise_check_violation(e, error_code)
-        elif error_code in {1400, 1407}:
-            self._raise_not_null_violation(e, error_code)
-        elif error_code and ORA_INTEGRITY_RANGE_START <= error_code < ORA_INTEGRITY_RANGE_END:
-            self._raise_integrity_error(e, error_code)
-        elif error_code in {1017, 12154, 12541, 12545, 12514, 12505}:
-            self._raise_connection_error(e, error_code)
-        elif error_code in {60, 8176}:
-            self._raise_transaction_error(e, error_code)
-        elif error_code in {1722, 1858, 1840}:
-            self._raise_data_error(e, error_code)
-        elif error_code and ORA_PARSING_RANGE_START <= error_code < ORA_PARSING_RANGE_END:
-            self._raise_parsing_error(e, error_code)
-        elif error_code == ORA_TABLESPACE_FULL:
-            self._raise_operational_error(e, error_code)
-        else:
-            self._raise_generic_error(e, error_code)
-
-    def _raise_unique_violation(self, e: Any, code: int) -> None:
-        msg = f"Oracle unique constraint violation [ORA-{code:05d}]: {e}"
-        raise UniqueViolationError(msg) from e
-
-    def _raise_foreign_key_violation(self, e: Any, code: int) -> None:
-        msg = f"Oracle foreign key constraint violation [ORA-{code:05d}]: {e}"
-        raise ForeignKeyViolationError(msg) from e
-
-    def _raise_check_violation(self, e: Any, code: int) -> None:
-        msg = f"Oracle check constraint violation [ORA-{code:05d}]: {e}"
-        raise CheckViolationError(msg) from e
 
-    def _raise_not_null_violation(self, e: Any, code: int) -> None:
-        msg = f"Oracle not-null constraint violation [ORA-{code:05d}]: {e}"
-        raise NotNullViolationError(msg) from e
-
-    def _raise_integrity_error(self, e: Any, code: int) -> None:
-        msg = f"Oracle integrity constraint violation [ORA-{code:05d}]: {e}"
-        raise IntegrityError(msg) from e
-
-    def _raise_parsing_error(self, e: Any, code: int) -> None:
-        msg = f"Oracle SQL syntax error [ORA-{code:05d}]: {e}"
-        raise SQLParsingError(msg) from e
-
-    def _raise_connection_error(self, e: Any, code: int) -> None:
-        msg = f"Oracle connection error [ORA-{code:05d}]: {e}"
-        raise DatabaseConnectionError(msg) from e
-
-    def _raise_transaction_error(self, e: Any, code: int) -> None:
-        msg = f"Oracle transaction error [ORA-{code:05d}]: {e}"
-        raise TransactionError(msg) from e
+class OracleAsyncExceptionHandler(OracleExceptionHandler):
+    """Async context manager for handling Oracle database exceptions.
 
-    def _raise_data_error(self, e: Any, code: int) -> None:
-        msg = f"Oracle data error [ORA-{code:05d}]: {e}"
-        raise DataError(msg) from e
+    Maps Oracle ORA-XXXXX error codes to specific SQLSpec exceptions
+    for better error handling in application code.
+    """
 
-    def _raise_operational_error(self, e: Any, code: int) -> None:
-        msg = f"Oracle operational error [ORA-{code:05d}]: {e}"
-        raise OperationalError(msg) from e
+    async def __aenter__(self) -> None:
+        return None
 
-    def _raise_generic_error(self, e: Any, code: "int | None") -> None:
-        msg = f"Oracle database error [ORA-{code:05d}]: {e}" if code else f"Oracle database error: {e}"
-        raise SQLSpecError(msg) from e
+    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        _ = exc_tb
+        if exc_type is None:
+            return
+        if issubclass(exc_type, oracledb.DatabaseError):
+            self._map_oracle_exception(exc_val)
 
 
 class OracleSyncDriver(SyncDriverAdapterBase):

sqlspec/adapters/psqlpy/config.py

@@ -10,7 +10,7 @@
 
 from sqlspec.adapters.psqlpy._types import PsqlpyConnection
 from sqlspec.adapters.psqlpy.driver import PsqlpyCursor, PsqlpyDriver, psqlpy_statement_config
-from sqlspec.config import AsyncDatabaseConfig
+from sqlspec.config import ADKConfig, AsyncDatabaseConfig, FastAPIConfig, FlaskConfig, LitestarConfig, StarletteConfig
 from sqlspec.core.statement import StatementConfig
 from sqlspec.typing import PGVECTOR_INSTALLED
 
@@ -105,7 +105,7 @@ def __init__(
         statement_config: StatementConfig | None = None,
         driver_features: "PsqlpyDriverFeatures | dict[str, Any] | None" = None,
         bind_key: str | None = None,
-        extension_config: "dict[str, dict[str, Any]] | None" = None,
+        extension_config: "dict[str, dict[str, Any]] | LitestarConfig | FastAPIConfig | StarletteConfig | FlaskConfig | ADKConfig | None" = None,
     ) -> None:
         """Initialize Psqlpy configuration.
 

sqlspec/adapters/psycopg/config.py

@@ -17,7 +17,15 @@
     PsycopgSyncDriver,
     psycopg_statement_config,
 )
-from sqlspec.config import AsyncDatabaseConfig, SyncDatabaseConfig
+from sqlspec.config import (
+    ADKConfig,
+    AsyncDatabaseConfig,
+    FastAPIConfig,
+    FlaskConfig,
+    LitestarConfig,
+    StarletteConfig,
+    SyncDatabaseConfig,
+)
 from sqlspec.typing import PGVECTOR_INSTALLED
 
 if TYPE_CHECKING:
@@ -106,7 +114,7 @@ def __init__(
         statement_config: "StatementConfig | None" = None,
         driver_features: "dict[str, Any] | None" = None,
         bind_key: "str | None" = None,
-        extension_config: "dict[str, dict[str, Any]] | None" = None,
+        extension_config: "dict[str, dict[str, Any]] | LitestarConfig | FastAPIConfig | StarletteConfig | FlaskConfig | ADKConfig | None" = None,
     ) -> None:
         """Initialize Psycopg synchronous configuration.
 
@@ -297,7 +305,7 @@ def __init__(
         statement_config: "StatementConfig | None" = None,
         driver_features: "dict[str, Any] | None" = None,
         bind_key: "str | None" = None,
-        extension_config: "dict[str, dict[str, Any]] | None" = None,
+        extension_config: "dict[str, dict[str, Any]] | LitestarConfig | FastAPIConfig | StarletteConfig | FlaskConfig | ADKConfig | None" = None,
     ) -> None:
         """Initialize Psycopg asynchronous configuration.