wip

Author: cofin

sqlspec/adapters/asyncpg/config.py

@@ -212,7 +212,9 @@ def driver_type(self) -> type[AsyncpgDriver]:  # type: ignore[override]
 
     @property
     def connection_config_dict(self) -> dict[str, Any]:
-        """Return the connection configuration as a dict.
+        """Return the connection configuration as a dict for asyncpg.connect().
+
+        This method filters out pool-specific parameters that are not valid for asyncpg.connect().
 
         Raises:
             ImproperConfigurationError: If the configuration is invalid.
@@ -221,19 +223,60 @@ def connection_config_dict(self) -> dict[str, Any]:
         merged_config = {**self.connection_config, **self.pool_config}
         config = {k: v for k, v in merged_config.items() if v is not Empty}
 
+        # Define pool-specific parameters that should not be passed to asyncpg.connect()
+        pool_only_params = {
+            "min_size",
+            "max_size",
+            "max_queries",
+            "max_inactive_connection_lifetime",
+            "setup",
+            "init",
+            "loop",
+            "connection_class",
+            "record_class",
+        }
+
+        # Filter out pool-specific parameters for connection creation
+        connection_config = {k: v for k, v in config.items() if k not in pool_only_params}
+
         # Validate essential connection info
-        has_dsn = config.get("dsn") is not None
-        has_host = config.get("host") is not None
+        has_dsn = connection_config.get("dsn") is not None
+        has_host = connection_config.get("host") is not None
 
         if not (has_dsn or has_host):
-            msg = f"AsyncPG configuration requires either 'dsn' or 'host' in pool_config. Current config: {config}"
+            msg = f"AsyncPG configuration requires either 'dsn' or 'host' in pool_config. Current config: {connection_config}"
             raise ImproperConfigurationError(msg)
 
+        # Set SSL to False by default for non-SSL environments (asyncpg 0.22.0+ defaults to 'prefer')
+        if "ssl" not in connection_config:
+            connection_config["ssl"] = False
+
+        return connection_config
+
+    @property
+    def pool_config_dict(self) -> dict[str, Any]:
+        """Return the full pool configuration as a dict for asyncpg.create_pool().
+
+        Returns:
+            A dictionary containing all pool configuration parameters.
+        """
+        # Merge connection_config into pool_config, with pool_config taking precedence
+        merged_config = {**self.connection_config, **self.pool_config}
+        config = {k: v for k, v in merged_config.items() if v is not Empty}
+
+        # Set SSL to False by default for non-SSL environments
+        if "ssl" not in config:
+            config["ssl"] = False
+
+        # Set reasonable defaults for pool parameters to prevent connection issues
+        if "max_inactive_connection_lifetime" not in config:
+            config["max_inactive_connection_lifetime"] = 300.0  # 5 minutes
+
         return config
 
     async def _create_pool_impl(self) -> "Pool[Record]":
         """Create the actual async connection pool."""
-        pool_args = self.connection_config_dict
+        pool_args = self.pool_config_dict
         return await asyncpg_create_pool(**pool_args)
 
     async def _close_pool_impl(self) -> None:

sqlspec/adapters/asyncpg/driver.py

@@ -61,19 +61,13 @@ def _get_placeholder_style(self) -> ParameterStyle:
     async def _execute_impl(
         self,
         statement: SQL,
-        parameters: Optional[SQLParameterType] = None,
         connection: Optional[AsyncpgConnection] = None,
-        config: Optional[SQLConfig] = None,
-        is_many: bool = False,
-        is_script: bool = False,
         **kwargs: Any,
     ) -> Any:
         async with instrument_operation_async(self, "asyncpg_execute", "database"):
             conn = self._connection(connection)
-            if config is not None and config != statement.config:
-                statement = statement.copy(config=config)
 
-            if is_script:
+            if statement.is_script:
                 final_sql = statement.to_sql(placeholder_style=ParameterStyle.STATIC)
                 if self.instrumentation_config.log_queries:
                     logger.debug("Executing SQL script: %s", final_sql)
@@ -84,37 +78,30 @@ async def _execute_impl(
             if self.instrumentation_config.log_queries:
                 logger.debug("Executing SQL: %s", final_sql)
 
-            if is_many:
-                args_list: list[Sequence[Any]] = []
-                if parameters and isinstance(parameters, Sequence):
-                    for param_set in parameters:
+            if statement.is_many:
+                params_list: list[tuple[Any, ...]] = []
+                if statement.parameters and isinstance(statement.parameters, Sequence):
+                    for param_set in statement.parameters:
                         if isinstance(param_set, (list, tuple)):
-                            args_list.append(param_set)
-                        elif isinstance(param_set, dict):
-                            # Convert dict to ordered args based on statement parameters
-                            ordered_params = statement.get_parameters(style=self._get_placeholder_style())
-                            if isinstance(ordered_params, (list, tuple)):
-                                args_list.append(ordered_params)
-                            else:
-                                args_list.append((ordered_params,) if ordered_params is not None else ())
+                            params_list.append(tuple(param_set))
                         elif param_set is None:
-                            args_list.append(())
+                            params_list.append(())
                         else:
-                            args_list.append((param_set,))
+                            params_list.append((param_set,))
 
-                if self.instrumentation_config.log_parameters and args_list:
-                    logger.debug("Query parameters (batch): %s", args_list)
+                if self.instrumentation_config.log_parameters and params_list:
+                    logger.debug("Query parameters (batch): %s", params_list)
 
-                await conn.executemany(final_sql, args_list)
-                return len(args_list)
+                return await conn.executemany(final_sql, params_list)
             # Single execution
-            ordered_params = statement.get_parameters(style=self._get_placeholder_style())
+            # Use the statement's already-processed parameters directly
+            processed_params = statement._merged_parameters if hasattr(statement, "_merged_parameters") else None
             args: list[Any] = []
 
-            if isinstance(ordered_params, (list, tuple)):
-                args.extend(ordered_params)
-            elif ordered_params is not None:
-                args.append(ordered_params)
+            if isinstance(processed_params, (list, tuple)):
+                args.extend(processed_params)
+            elif processed_params is not None:
+                args.append(processed_params)
 
             if self.instrumentation_config.log_parameters and args:
                 logger.debug("Query parameters: %s", args)
@@ -246,13 +233,14 @@ async def select_to_arrow(
                 raise TypeError(msg)
 
             final_sql = stmt_obj.to_sql(placeholder_style=self._get_placeholder_style())
-            ordered_params = stmt_obj.get_parameters(style=self._get_placeholder_style())
+            # Use the statement's already-processed parameters instead of calling get_parameters()
+            processed_params = stmt_obj._merged_parameters if hasattr(stmt_obj, "_merged_parameters") else None
 
             args: list[Any] = []
-            if isinstance(ordered_params, (list, tuple)):
-                args.extend(ordered_params)
-            elif ordered_params is not None:
-                args.append(ordered_params)
+            if isinstance(processed_params, (list, tuple)):
+                args.extend(processed_params)
+            elif processed_params is not None:
+                args.append(processed_params)
 
             records = await conn.fetch(final_sql, *args)
 

sqlspec/adapters/duckdb/config.py

@@ -262,24 +262,28 @@ def connection_config_dict(self) -> dict[str, Any]:
         # Filter out empty values and prepare config for duckdb.connect()
         config_dict = {k: v for k, v in self.connection_config.items() if v is not Empty}
 
-        # Extract DuckDB-specific settings into the config parameter
-        duckdb_settings = {}
-        connection_params = {}
-
         # Parameters that go directly to duckdb.connect()
-        direct_params = {"database", "read_only", "config"}
+        connection_params = {}
+        duckdb_config_settings = {}
 
+        # Only database and read_only go directly to connect()
+        # Everything else goes into the config dictionary
         for key, value in config_dict.items():
-            if key in direct_params:
-                connection_params[key] = value
+            if key == "database":
+                connection_params["database"] = value
+            elif key == "read_only":
+                connection_params["read_only"] = value
+            elif key == "config":
+                # If user provided a config dict, merge it
+                if isinstance(value, dict):
+                    duckdb_config_settings.update(value)
             else:
                 # All other parameters are DuckDB configuration settings
-                duckdb_settings[key] = value
+                duckdb_config_settings[key] = value
 
-        existing_config: dict[str, Any] = connection_params.get("config", {})
-        if duckdb_settings:
-            existing_config.update(duckdb_settings)
-            connection_params["config"] = existing_config
+        # Add the config dictionary if we have settings
+        if duckdb_config_settings:
+            connection_params["config"] = duckdb_config_settings
 
         return connection_params
 

sqlspec/adapters/duckdb/driver.py

@@ -99,6 +99,14 @@ def _execute_impl(
                     final_exec_params = [list(p) if isinstance(p, (list, tuple)) else [p] for p in parameters]
                 else:
                     final_exec_params = []
+            # Use provided parameters if available, otherwise get from statement
+            elif parameters is not None:
+                if isinstance(parameters, list):
+                    final_exec_params = parameters
+                elif hasattr(parameters, "__iter__") and not isinstance(parameters, (str, bytes)):
+                    final_exec_params = list(parameters)
+                else:
+                    final_exec_params = [parameters]
             else:
                 single_params = statement.get_parameters(style=self._get_placeholder_style())
                 if single_params is not None:
@@ -122,9 +130,17 @@ def _execute_impl(
                     return "SCRIPT EXECUTED"
                 if is_many:
                     cursor.executemany(final_sql, cast("list[list[Any]]", final_exec_params))
-                else:
-                    cursor.execute(final_sql, final_exec_params or [])
-                return cursor
+                    # For executemany, return cursor info for execute result
+                    return {"rowcount": cursor.rowcount if hasattr(cursor, "rowcount") else -1}
+                cursor.execute(final_sql, final_exec_params or [])
+
+                # For SELECT queries, fetch the data immediately since cursor will be closed
+                if self.returns_rows(statement.expression):
+                    fetched_data = cursor.fetchall()
+                    column_names = [col[0] for col in cursor.description or []]
+                    return {"data": fetched_data, "columns": column_names}
+                # For non-SELECT queries, return cursor info
+                return {"rowcount": cursor.rowcount if hasattr(cursor, "rowcount") else -1}
 
     def _wrap_select_result(
         self,
@@ -134,9 +150,23 @@ def _wrap_select_result(
         **kwargs: Any,
     ) -> Union[SelectResult[ModelDTOT], SelectResult[dict[str, Any]]]:
         with instrument_operation(self, "duckdb_wrap_select", "database"):
-            cursor = raw_driver_result
-            fetched_data = cursor.fetchall()
-            column_names = [col[0] for col in cursor.description or []]
+            # Handle the new dictionary format from _execute_impl
+            if isinstance(raw_driver_result, dict) and "data" in raw_driver_result:
+                fetched_data = raw_driver_result["data"]
+                column_names = raw_driver_result["columns"]
+            elif not isinstance(raw_driver_result, dict):
+                # Fallback for backward compatibility (shouldn't happen with new implementation)
+                if hasattr(raw_driver_result, "fetchall") and hasattr(raw_driver_result, "description"):
+                    fetched_data = raw_driver_result.fetchall()
+                    column_names = [col[0] for col in raw_driver_result.description or []]
+                else:
+                    # Should not happen with current implementation
+                    fetched_data = []
+                    column_names = []
+            else:
+                # Should not happen with current implementation
+                fetched_data = []
+                column_names = []
 
             # Convert to list of dicts
             rows_as_dicts: list[dict[str, Any]] = []
@@ -185,8 +215,12 @@ def _wrap_execute_result(
                     operation_type=operation_type or "SCRIPT",
                 )
 
-            cursor = raw_driver_result
-            rows_affected = cursor.rowcount if hasattr(cursor, "rowcount") else -1
+            # Handle the new dictionary format from _execute_impl
+            if isinstance(raw_driver_result, dict) and "rowcount" in raw_driver_result:
+                rows_affected = raw_driver_result["rowcount"]
+            else:
+                # Fallback for backward compatibility
+                rows_affected = getattr(raw_driver_result, "rowcount", -1)
 
             if self.instrumentation_config.log_results_count:
                 logger.debug("Execute operation affected %d rows", rows_affected)

sqlspec/base.py

@@ -5,7 +5,6 @@
 from collections.abc import Awaitable, Coroutine
 from typing import (
     TYPE_CHECKING,
-    Annotated,
     Any,
     Optional,
     Union,
@@ -78,17 +77,17 @@ def _cleanup_pools(self) -> None:
         logger.info("Pool cleanup completed. Cleaned %d pools", cleaned_count)
 
     @overload
-    def add_config(self, config: "SyncConfigT") -> "Annotated[type[SyncConfigT], int]":  # pyright: ignore[reportInvalidTypeVarUse]
+    def add_config(self, config: "SyncConfigT") -> "type[SyncConfigT]":  # pyright: ignore[reportInvalidTypeVarUse]
         ...
 
     @overload
-    def add_config(self, config: "AsyncConfigT") -> "Annotated[type[AsyncConfigT], int]":  # pyright: ignore[reportInvalidTypeVarUse]
+    def add_config(self, config: "AsyncConfigT") -> "type[AsyncConfigT]":  # pyright: ignore[reportInvalidTypeVarUse]
         ...
 
     def add_config(
         self,
         config: "Union[SyncConfigT, AsyncConfigT]",
-    ) -> "Union[Annotated[type[SyncConfigT], int], Annotated[type[AsyncConfigT], int]]":  # pyright: ignore[reportInvalidTypeVarUse]
+    ) -> "type[Union[SyncConfigT, AsyncConfigT]]":  # pyright: ignore[reportInvalidTypeVarUse]
         """Add a configuration instance to the registry.
 
         Args: