fix: init consistency (#61)

cofin · web-flow · commit a75d19049ae5 · 2025-08-18T17:31:32.000-05:00
Ensure all configurations have consistent init methods.
diff --git a/sqlspec/adapters/adbc/config.py b/sqlspec/adapters/adbc/config.py
@@ -79,15 +79,17 @@ def __init__(
         self,
         *,
         connection_config: Optional[Union[AdbcConnectionParams, dict[str, Any]]] = None,
-        statement_config: Optional[StatementConfig] = None,
         migration_config: Optional[dict[str, Any]] = None,
+        statement_config: Optional[StatementConfig] = None,
+        driver_features: Optional[dict[str, Any]] = None,
     ) -> None:
         """Initialize ADBC configuration.
 
         Args:
             connection_config: Connection configuration parameters
-            statement_config: Default SQL statement configuration
             migration_config: Migration configuration
+            statement_config: Default SQL statement configuration
+            driver_features: Driver feature configuration
         """
         if connection_config is None:
             connection_config = {}
@@ -106,7 +108,7 @@ def __init__(
             connection_config=self.connection_config,
             migration_config=migration_config,
             statement_config=statement_config,
-            driver_features={},
+            driver_features=driver_features or {},
         )
 
     def _resolve_driver_name(self) -> str:
diff --git a/sqlspec/adapters/aiosqlite/config.py b/sqlspec/adapters/aiosqlite/config.py
@@ -61,7 +61,7 @@ def __init__(
         pool_instance: "Optional[AiosqliteConnectionPool]" = None,
         migration_config: "Optional[dict[str, Any]]" = None,
         statement_config: "Optional[StatementConfig]" = None,
-        **kwargs: Any,
+        driver_features: "Optional[dict[str, Any]]" = None,
     ) -> None:
         """Initialize AioSQLite configuration.
 
@@ -70,10 +70,9 @@ def __init__(
             pool_instance: Optional pre-configured connection pool instance.
             migration_config: Optional migration configuration.
             statement_config: Optional statement configuration.
-            **kwargs: Additional connection parameters that override pool_config.
+            driver_features: Optional driver feature configuration.
         """
         config_dict = dict(pool_config) if pool_config else {}
-        config_dict.update(kwargs)  # Allow kwargs to override pool_config values
 
         # Handle memory database URI conversion - test expectation is different than sqlite pattern
         if "database" not in config_dict or config_dict["database"] == ":memory:":
@@ -85,7 +84,7 @@ def __init__(
             pool_instance=pool_instance,
             migration_config=migration_config,
             statement_config=statement_config or aiosqlite_statement_config,
-            driver_features={},
+            driver_features=driver_features or {},
         )
 
     def _get_pool_config_dict(self) -> "dict[str, Any]":
diff --git a/sqlspec/adapters/asyncmy/config.py b/sqlspec/adapters/asyncmy/config.py
@@ -70,6 +70,7 @@ def __init__(
         pool_instance: "Optional[Pool]" = None,
         migration_config: Optional[dict[str, Any]] = None,
         statement_config: "Optional[StatementConfig]" = None,
+        driver_features: "Optional[dict[str, Any]]" = None,
     ) -> None:
         """Initialize Asyncmy configuration.
 
@@ -78,6 +79,7 @@ def __init__(
             pool_instance: Existing pool instance to use
             migration_config: Migration configuration
             statement_config: Statement configuration override
+            driver_features: Optional driver feature configuration
         """
         processed_pool_config: dict[str, Any] = dict(pool_config) if pool_config else {}
         if "extra" in processed_pool_config:
@@ -97,7 +99,7 @@ def __init__(
             pool_instance=pool_instance,
             migration_config=migration_config,
             statement_config=statement_config,
-            driver_features={},
+            driver_features=driver_features or {},
         )
 
     async def _create_pool(self) -> "Pool":  # pyright: ignore
diff --git a/sqlspec/adapters/bigquery/config.py b/sqlspec/adapters/bigquery/config.py
@@ -71,6 +71,7 @@ class BigQueryDriverFeatures(TypedDict, total=False):
     Only non-standard BigQuery client parameters that are SQLSpec-specific extensions.
     """
 
+    connection_instance: NotRequired["BigQueryConnection"]
     on_job_start: NotRequired["Callable[[str], None]"]
     on_job_complete: NotRequired["Callable[[str, Any], None]"]
     on_connection_create: NotRequired["Callable[[Any], None]"]
@@ -93,7 +94,6 @@ class BigQueryConfig(NoPoolSyncConfig[BigQueryConnection, BigQueryDriver]):
     def __init__(
         self,
         *,
-        connection_instance: "Optional[BigQueryConnection]" = None,
         connection_config: "Optional[Union[BigQueryConnectionParams, dict[str, Any]]]" = None,
         migration_config: Optional[dict[str, Any]] = None,
         statement_config: "Optional[StatementConfig]" = None,
@@ -103,10 +103,10 @@ def __init__(
 
         Args:
             connection_config: Standard connection configuration parameters
-            connection_instance: Existing connection instance to use
             migration_config: Migration configuration
             statement_config: Statement configuration override
-            driver_features: BigQuery-specific driver features and configurations
+            driver_features: BigQuery-specific driver features and configurations.
+                Can include 'connection_instance' to reuse an existing BigQuery connection.
 
         Example:
             >>> # Basic BigQuery connection
@@ -143,19 +143,16 @@ def __init__(
             ... )
         """
 
-        # Store connection instance
-        self._connection_instance = connection_instance
-
-        # Setup configuration following DuckDB pattern
         self.connection_config: dict[str, Any] = dict(connection_config) if connection_config else {}
         if "extra" in self.connection_config:
             extras = self.connection_config.pop("extra")
             self.connection_config.update(extras)
 
-        # Setup driver features
         self.driver_features: dict[str, Any] = dict(driver_features) if driver_features else {}
 
-        # Setup default job config if not provided
+        # Initialize connection instance cache (for performance optimization)
+        self._connection_instance: Optional[BigQueryConnection] = self.driver_features.get("connection_instance")
+
         if "default_query_job_config" not in self.connection_config:
             self._setup_default_job_config()
 
@@ -237,8 +234,8 @@ def create_connection(self) -> BigQueryConnection:
             if on_connection_create:
                 on_connection_create(connection)
 
+            # Cache the connection for reuse (BigQuery connections are expensive)
             self._connection_instance = connection
-
         except Exception as e:
             project = self.connection_config.get("project", "Unknown")
             msg = f"Could not configure BigQuery connection for project '{project}'. Error: {e}"
diff --git a/sqlspec/adapters/duckdb/config.py b/sqlspec/adapters/duckdb/config.py
@@ -149,8 +149,8 @@ def __init__(
         self,
         *,
         pool_config: "Optional[Union[DuckDBPoolParams, dict[str, Any]]]" = None,
-        migration_config: Optional[dict[str, Any]] = None,
         pool_instance: "Optional[DuckDBConnectionPool]" = None,
+        migration_config: Optional[dict[str, Any]] = None,
         statement_config: "Optional[StatementConfig]" = None,
         driver_features: "Optional[Union[DuckDBDriverFeatures, dict[str, Any]]]" = None,
     ) -> None:
diff --git a/sqlspec/adapters/oracledb/config.py b/sqlspec/adapters/oracledb/config.py
@@ -83,18 +83,20 @@ class OracleSyncConfig(SyncDatabaseConfig[OracleSyncConnection, "ConnectionPool"
     def __init__(
         self,
         *,
-        pool_instance: "Optional[ConnectionPool]" = None,
         pool_config: "Optional[Union[OraclePoolParams, dict[str, Any]]]" = None,
-        statement_config: "Optional[StatementConfig]" = None,
+        pool_instance: "Optional[ConnectionPool]" = None,
         migration_config: Optional[dict[str, Any]] = None,
+        statement_config: "Optional[StatementConfig]" = None,
+        driver_features: "Optional[dict[str, Any]]" = None,
     ) -> None:
         """Initialize Oracle synchronous configuration.
 
         Args:
             pool_config: Pool configuration parameters
             pool_instance: Existing pool instance to use
-            statement_config: Default SQL statement configuration
             migration_config: Migration configuration
+            statement_config: Default SQL statement configuration
+            driver_features: Optional driver feature configuration
         """
         # Store the pool config as a dict and extract/merge extras
         processed_pool_config: dict[str, Any] = dict(pool_config) if pool_config else {}
@@ -107,6 +109,7 @@ def __init__(
             pool_instance=pool_instance,
             migration_config=migration_config,
             statement_config=statement_config,
+            driver_features=driver_features or {},
         )
 
     def _create_pool(self) -> "ConnectionPool":
@@ -208,16 +211,18 @@ def __init__(
         *,
         pool_config: "Optional[Union[OraclePoolParams, dict[str, Any]]]" = None,
         pool_instance: "Optional[AsyncConnectionPool]" = None,
-        statement_config: "Optional[StatementConfig]" = None,
         migration_config: Optional[dict[str, Any]] = None,
+        statement_config: "Optional[StatementConfig]" = None,
+        driver_features: "Optional[dict[str, Any]]" = None,
     ) -> None:
         """Initialize Oracle asynchronous configuration.
 
         Args:
             pool_config: Pool configuration parameters
             pool_instance: Existing pool instance to use
-            statement_config: Default SQL statement configuration
             migration_config: Migration configuration
+            statement_config: Default SQL statement configuration
+            driver_features: Optional driver feature configuration
         """
         # Store the pool config as a dict and extract/merge extras
         processed_pool_config: dict[str, Any] = dict(pool_config) if pool_config else {}
@@ -230,6 +235,7 @@ def __init__(
             pool_instance=pool_instance,
             migration_config=migration_config,
             statement_config=statement_config or oracledb_statement_config,
+            driver_features=driver_features or {},
         )
 
     async def _create_pool(self) -> "AsyncConnectionPool":
diff --git a/sqlspec/adapters/psqlpy/config.py b/sqlspec/adapters/psqlpy/config.py
@@ -86,17 +86,19 @@ def __init__(
         self,
         *,
         pool_config: Optional[Union[PsqlpyPoolParams, dict[str, Any]]] = None,
-        statement_config: Optional[StatementConfig] = None,
         pool_instance: Optional[ConnectionPool] = None,
         migration_config: Optional[dict[str, Any]] = None,
+        statement_config: Optional[StatementConfig] = None,
+        driver_features: Optional[dict[str, Any]] = None,
     ) -> None:
         """Initialize Psqlpy asynchronous configuration.
 
         Args:
             pool_config: Pool configuration parameters (TypedDict or dict)
             pool_instance: Existing connection pool instance to use
-            statement_config: Default SQL statement configuration
             migration_config: Migration configuration
+            statement_config: Default SQL statement configuration
+            driver_features: Optional driver feature configuration
         """
         processed_pool_config: dict[str, Any] = dict(pool_config) if pool_config else {}
         if "extra" in processed_pool_config:
@@ -107,6 +109,7 @@ def __init__(
             pool_instance=pool_instance,
             migration_config=migration_config,
             statement_config=statement_config or psqlpy_statement_config,
+            driver_features=driver_features or {},
         )
 
     def _get_pool_config_dict(self) -> dict[str, Any]:
diff --git a/sqlspec/adapters/psycopg/config.py b/sqlspec/adapters/psycopg/config.py
@@ -85,17 +85,18 @@ def __init__(
         *,
         pool_config: "Optional[Union[PsycopgPoolParams, dict[str, Any]]]" = None,
         pool_instance: Optional["ConnectionPool"] = None,
-        statement_config: "Optional[StatementConfig]" = None,
         migration_config: Optional[dict[str, Any]] = None,
+        statement_config: "Optional[StatementConfig]" = None,
+        driver_features: "Optional[dict[str, Any]]" = None,
     ) -> None:
         """Initialize Psycopg synchronous configuration.
 
         Args:
             pool_config: Pool configuration parameters (TypedDict or dict)
             pool_instance: Existing pool instance to use
-            statement_config: Default SQL statement configuration
             migration_config: Migration configuration
-
+            statement_config: Default SQL statement configuration
+            driver_features: Optional driver feature configuration
         """
         processed_pool_config: dict[str, Any] = dict(pool_config) if pool_config else {}
         if "extra" in processed_pool_config:
@@ -107,6 +108,7 @@ def __init__(
             pool_instance=pool_instance,
             migration_config=migration_config,
             statement_config=statement_config or psycopg_statement_config,
+            driver_features=driver_features or {},
         )
 
     def _create_pool(self) -> "ConnectionPool":
@@ -268,14 +270,16 @@ def __init__(
         pool_instance: "Optional[AsyncConnectionPool]" = None,
         migration_config: "Optional[dict[str, Any]]" = None,
         statement_config: "Optional[StatementConfig]" = None,
+        driver_features: "Optional[dict[str, Any]]" = None,
     ) -> None:
         """Initialize Psycopg asynchronous configuration.
 
         Args:
             pool_config: Pool configuration parameters (TypedDict or dict)
             pool_instance: Existing pool instance to use
-            statement_config: Default SQL statement configuration
             migration_config: Migration configuration
+            statement_config: Default SQL statement configuration
+            driver_features: Optional driver feature configuration
         """
         processed_pool_config: dict[str, Any] = dict(pool_config) if pool_config else {}
         if "extra" in processed_pool_config:
@@ -287,6 +291,7 @@ def __init__(
             pool_instance=pool_instance,
             migration_config=migration_config,
             statement_config=statement_config or psycopg_statement_config,
+            driver_features=driver_features or {},
         )
 
     async def _create_pool(self) -> "AsyncConnectionPool":
diff --git a/sqlspec/adapters/sqlite/config.py b/sqlspec/adapters/sqlite/config.py
@@ -44,16 +44,18 @@ def __init__(
         *,
         pool_config: "Optional[Union[SqliteConnectionParams, dict[str, Any]]]" = None,
         pool_instance: "Optional[SqliteConnectionPool]" = None,
-        statement_config: "Optional[StatementConfig]" = None,
         migration_config: "Optional[dict[str, Any]]" = None,
+        statement_config: "Optional[StatementConfig]" = None,
+        driver_features: "Optional[dict[str, Any]]" = None,
     ) -> None:
         """Initialize SQLite configuration.
 
         Args:
             pool_config: Configuration parameters including connection settings
             pool_instance: Pre-created pool instance
-            statement_config: Default SQL statement configuration
             migration_config: Migration configuration
+            statement_config: Default SQL statement configuration
+            driver_features: Optional driver feature configuration
         """
         if pool_config is None:
             pool_config = {}
@@ -66,7 +68,7 @@ def __init__(
             pool_config=cast("dict[str, Any]", pool_config),
             migration_config=migration_config,
             statement_config=statement_config or sqlite_statement_config,
-            driver_features={},
+            driver_features=driver_features or {},
         )
 
     def _get_connection_config_dict(self) -> "dict[str, Any]":
diff --git a/tests/integration/test_adapters/test_aiosqlite/test_connection.py b/tests/integration/test_adapters/test_aiosqlite/test_connection.py
@@ -177,7 +177,9 @@ async def test_config_with_kwargs_override() -> None:
     pool_config = {"database": "base.db", "timeout": 5.0}
 
     unique_db = f"file:override_{uuid4().hex}.db?mode=memory&cache=shared"
-    config = AiosqliteConfig(pool_config=pool_config, database=unique_db, timeout=15.0)
+    # Override pool_config with specific test values
+    test_pool_config = {**pool_config, "database": unique_db, "timeout": 15.0}
+    config = AiosqliteConfig(pool_config=test_pool_config)
 
     try:
         connection_config = config._get_connection_config_dict()
