sqlspec 0.12.1__py3-none-any.whl → 0.13.0__py3-none-any.whl

sqlspec/adapters/psqlpy/config.py

@@ -3,7 +3,6 @@
 import logging
 from collections.abc import AsyncGenerator
 from contextlib import asynccontextmanager
-from dataclasses import replace
 from typing import TYPE_CHECKING, Any, ClassVar, Optional
 
 from psqlpy import ConnectionPool
@@ -302,7 +301,6 @@ class PsqlpyConfig(AsyncDatabaseConfig[PsqlpyConnection, ConnectionPool, PsqlpyD
             if getattr(self, field, None) is not None and getattr(self, field) is not Empty
         }
 
-        # Add connection-specific extras (not pool-specific ones)
         config.update(self.extras)
 
         return config
@@ -359,11 +357,9 @@ class PsqlpyConfig(AsyncDatabaseConfig[PsqlpyConnection, ConnectionPool, PsqlpyD
         Returns:
             A psqlpy Connection instance.
         """
-        # Ensure pool exists
         if not self.pool_instance:
             self.pool_instance = await self._create_pool()
 
-        # Get connection from pool
         return await self.pool_instance.connection()
 
     @asynccontextmanager
@@ -377,7 +373,6 @@ class PsqlpyConfig(AsyncDatabaseConfig[PsqlpyConnection, ConnectionPool, PsqlpyD
         Yields:
             A psqlpy Connection instance.
         """
-        # Ensure pool exists
         if not self.pool_instance:
             self.pool_instance = await self._create_pool()
 
@@ -396,15 +391,16 @@ class PsqlpyConfig(AsyncDatabaseConfig[PsqlpyConnection, ConnectionPool, PsqlpyD
             A PsqlpyDriver instance.
         """
         async with self.provide_connection(*args, **kwargs) as conn:
-            # Create statement config with parameter style info if not already set
             statement_config = self.statement_config
+            # Inject parameter style info if not already set
             if statement_config.allowed_parameter_styles is None:
+                from dataclasses import replace
+
                 statement_config = replace(
                     statement_config,
                     allowed_parameter_styles=self.supported_parameter_styles,
                     target_parameter_style=self.preferred_parameter_style,
                 )
-
             driver = self.driver_type(connection=conn, config=statement_config)
             yield driver
 

sqlspec/adapters/psqlpy/driver.py

@@ -2,11 +2,12 @@
 
 import io
 import logging
-from typing import TYPE_CHECKING, Any, Optional, Union, cast
+from typing import TYPE_CHECKING, Any, Optional, cast
 
 from psqlpy import Connection
 
 from sqlspec.driver import AsyncDriverAdapterProtocol
+from sqlspec.driver.connection import managed_transaction_async
 from sqlspec.driver.mixins import (
     AsyncPipelinedExecutionMixin,
     AsyncStorageMixin,
@@ -14,10 +15,10 @@ from sqlspec.driver.mixins import (
     ToSchemaMixin,
     TypeCoercionMixin,
 )
-from sqlspec.statement.parameters import ParameterStyle
-from sqlspec.statement.result import DMLResultDict, ScriptResultDict, SelectResultDict, SQLResult
+from sqlspec.statement.parameters import ParameterStyle, ParameterValidator
+from sqlspec.statement.result import SQLResult
 from sqlspec.statement.sql import SQL, SQLConfig
-from sqlspec.typing import DictRow, ModelDTOT, RowT
+from sqlspec.typing import DictRow, RowT
 
 if TYPE_CHECKING:
     from sqlglot.dialects.dialect import DialectType
@@ -76,13 +77,36 @@ class PsqlpyDriver(
 
     async def _execute_statement(
         self, statement: SQL, connection: Optional[PsqlpyConnection] = None, **kwargs: Any
-    ) -> Union[SelectResultDict, DMLResultDict, ScriptResultDict]:
+    ) -> SQLResult[RowT]:
         if statement.is_script:
             sql, _ = statement.compile(placeholder_style=ParameterStyle.STATIC)
             return await self._execute_script(sql, connection=connection, **kwargs)
 
-        # Let the SQL object handle parameter style conversion based on dialect support
-        sql, params = statement.compile(placeholder_style=self.default_parameter_style)
+        # Detect parameter styles in the SQL
+        detected_styles = set()
+        sql_str = statement.to_sql(placeholder_style=None)  # Get raw SQL
+        validator = self.config.parameter_validator if self.config else ParameterValidator()
+        param_infos = validator.extract_parameters(sql_str)
+        if param_infos:
+            detected_styles = {p.style for p in param_infos}
+
+        # Determine target style based on what's in the SQL
+        target_style = self.default_parameter_style
+
+        # Check if there are unsupported styles
+        unsupported_styles = detected_styles - set(self.supported_parameter_styles)
+        if unsupported_styles:
+            # Force conversion to default style
+            target_style = self.default_parameter_style
+        elif detected_styles:
+            # Prefer the first supported style found
+            for style in detected_styles:
+                if style in self.supported_parameter_styles:
+                    target_style = style
+                    break
+
+        # Compile with the determined style
+        sql, params = statement.compile(placeholder_style=target_style)
         params = self._process_parameters(params)
 
         if statement.is_many:
@@ -92,43 +116,99 @@ class PsqlpyDriver(
 
     async def _execute(
         self, sql: str, parameters: Any, statement: SQL, connection: Optional[PsqlpyConnection] = None, **kwargs: Any
-    ) -> Union[SelectResultDict, DMLResultDict]:
-        conn = self._connection(connection)
-        if self.returns_rows(statement.expression):
-            query_result = await conn.fetch(sql, parameters=parameters)
-            # Convert query_result to list of dicts
-            dict_rows: list[dict[str, Any]] = []
-            if query_result:
-                # psqlpy QueryResult has a result() method that returns list of dicts
-                dict_rows = query_result.result()
-            column_names = list(dict_rows[0].keys()) if dict_rows else []
-            return {"data": dict_rows, "column_names": column_names, "rows_affected": len(dict_rows)}
-        query_result = await conn.execute(sql, parameters=parameters)
-        # Note: psqlpy doesn't provide rows_affected for DML operations
-        # The QueryResult object only has result(), as_class(), and row_factory() methods
-        # For accurate row counts, use RETURNING clause
-        affected_count = -1  # Unknown, as psqlpy doesn't provide this info
-        return {"rows_affected": affected_count, "status_message": "OK"}
+    ) -> SQLResult[RowT]:
+        # Use provided connection or driver's default connection
+        conn = connection if connection is not None else self._connection(None)
+
+        async with managed_transaction_async(conn, auto_commit=True) as txn_conn:
+            # PSQLPy expects parameters as a list (for $1, $2, etc.) or dict
+            # Ensure we always pass a sequence or mapping, never a scalar
+            final_params: Any
+            if isinstance(parameters, (list, tuple)):
+                final_params = list(parameters)
+            elif isinstance(parameters, dict):
+                final_params = parameters
+            elif parameters is None:
+                final_params = []
+            else:
+                # Single parameter - wrap in list for NUMERIC style ($1)
+                final_params = [parameters]
+
+            if self.returns_rows(statement.expression):
+                query_result = await txn_conn.fetch(sql, parameters=final_params)
+                dict_rows: list[dict[str, Any]] = []
+                if query_result:
+                    # psqlpy QueryResult has a result() method that returns list of dicts
+                    dict_rows = query_result.result()
+                column_names = list(dict_rows[0].keys()) if dict_rows else []
+                return SQLResult(
+                    statement=statement,
+                    data=cast("list[RowT]", dict_rows),
+                    column_names=column_names,
+                    rows_affected=len(dict_rows),
+                    operation_type="SELECT",
+                )
+
+            query_result = await txn_conn.execute(sql, parameters=final_params)
+            # Note: psqlpy doesn't provide rows_affected for DML operations
+            # The QueryResult object only has result(), as_class(), and row_factory() methods
+            affected_count = -1  # Unknown, as psqlpy doesn't provide this info
+            return SQLResult(
+                statement=statement,
+                data=[],
+                rows_affected=affected_count,
+                operation_type=self._determine_operation_type(statement),
+                metadata={"status_message": "OK"},
+            )
 
     async def _execute_many(
         self, sql: str, param_list: Any, connection: Optional[PsqlpyConnection] = None, **kwargs: Any
-    ) -> DMLResultDict:
-        conn = self._connection(connection)
-        await conn.execute_many(sql, param_list or [])
-        # execute_many doesn't return a value with rows_affected
-        affected_count = -1
-        return {"rows_affected": affected_count, "status_message": "OK"}
+    ) -> SQLResult[RowT]:
+        # Use provided connection or driver's default connection
+        conn = connection if connection is not None else self._connection(None)
+
+        async with managed_transaction_async(conn, auto_commit=True) as txn_conn:
+            # PSQLPy expects a list of parameter lists/tuples for execute_many
+            if param_list is None:
+                final_param_list = []
+            elif isinstance(param_list, (list, tuple)):
+                # Ensure each parameter set is a list/tuple
+                final_param_list = [
+                    list(params) if isinstance(params, (list, tuple)) else [params] for params in param_list
+                ]
+            else:
+                # Single parameter set - wrap it
+                final_param_list = [list(param_list) if isinstance(param_list, (list, tuple)) else [param_list]]
+
+            await txn_conn.execute_many(sql, final_param_list)
+            # execute_many doesn't return a value with rows_affected
+            affected_count = -1
+            return SQLResult(
+                statement=SQL(sql, _dialect=self.dialect),
+                data=[],
+                rows_affected=affected_count,
+                operation_type="EXECUTE",
+                metadata={"status_message": "OK"},
+            )
 
     async def _execute_script(
         self, script: str, connection: Optional[PsqlpyConnection] = None, **kwargs: Any
-    ) -> ScriptResultDict:
-        conn = self._connection(connection)
-        # psqlpy can execute multi-statement scripts directly
-        await conn.execute(script)
-        return {
-            "statements_executed": -1,  # Not directly supported, but script is executed
-            "status_message": "SCRIPT EXECUTED",
-        }
+    ) -> SQLResult[RowT]:
+        # Use provided connection or driver's default connection
+        conn = connection if connection is not None else self._connection(None)
+
+        async with managed_transaction_async(conn, auto_commit=True) as txn_conn:
+            # psqlpy can execute multi-statement scripts directly
+            await txn_conn.execute(script)
+            return SQLResult(
+                statement=SQL(script, _dialect=self.dialect).as_script(),
+                data=[],
+                rows_affected=0,
+                operation_type="SCRIPT",
+                metadata={"status_message": "SCRIPT EXECUTED"},
+                total_statements=-1,  # Not directly supported, but script is executed
+                successful_statements=-1,
+            )
 
     async def _ingest_arrow_table(self, table: "Any", table_name: str, mode: str = "append", **options: Any) -> int:
         self._ensure_pyarrow_installed()
@@ -154,61 +234,6 @@ class PsqlpyDriver(
         msg = "Connection does not support COPY operations"
         raise NotImplementedError(msg)
 
-    async def _wrap_select_result(
-        self, statement: SQL, result: SelectResultDict, schema_type: Optional[type[ModelDTOT]] = None, **kwargs: Any
-    ) -> Union[SQLResult[ModelDTOT], SQLResult[RowT]]:
-        dict_rows = result["data"]
-        column_names = result["column_names"]
-        rows_affected = result["rows_affected"]
-
-        if schema_type:
-            converted_data = self.to_schema(data=dict_rows, schema_type=schema_type)
-            return SQLResult[ModelDTOT](
-                statement=statement,
-                data=list(converted_data),
-                column_names=column_names,
-                rows_affected=rows_affected,
-                operation_type="SELECT",
-            )
-        return SQLResult[RowT](
-            statement=statement,
-            data=cast("list[RowT]", dict_rows),
-            column_names=column_names,
-            rows_affected=rows_affected,
-            operation_type="SELECT",
-        )
-
-    async def _wrap_execute_result(
-        self, statement: SQL, result: Union[DMLResultDict, ScriptResultDict], **kwargs: Any
-    ) -> SQLResult[RowT]:
-        operation_type = "UNKNOWN"
-        if statement.expression:
-            operation_type = str(statement.expression.key).upper()
-
-        if "statements_executed" in result:
-            script_result = cast("ScriptResultDict", result)
-            return SQLResult[RowT](
-                statement=statement,
-                data=[],
-                rows_affected=0,
-                operation_type="SCRIPT",
-                metadata={
-                    "status_message": script_result.get("status_message", ""),
-                    "statements_executed": script_result.get("statements_executed", -1),
-                },
-            )
-
-        dml_result = cast("DMLResultDict", result)
-        rows_affected = dml_result.get("rows_affected", -1)
-        status_message = dml_result.get("status_message", "")
-        return SQLResult[RowT](
-            statement=statement,
-            data=[],
-            rows_affected=rows_affected,
-            operation_type=operation_type,
-            metadata={"status_message": status_message},
-        )
-
     def _connection(self, connection: Optional[PsqlpyConnection] = None) -> PsqlpyConnection:
         """Get the connection to use for the operation."""
         return connection or self.connection

sqlspec/adapters/psycopg/config.py

@@ -3,7 +3,6 @@
 import contextlib
 import logging
 from contextlib import asynccontextmanager
-from dataclasses import replace
 from typing import TYPE_CHECKING, Any, ClassVar, Optional, cast
 
 from psycopg.rows import dict_row
@@ -211,7 +210,6 @@ class PsycopgSyncConfig(SyncDatabaseConfig[PsycopgSyncConnection, ConnectionPool
         self.configure = configure
         self.kwargs = kwargs or {}
 
-        # Handle extras and additional kwargs
         self.extras = extras or {}
         self.extras.update(additional_kwargs)
 
@@ -240,7 +238,6 @@ class PsycopgSyncConfig(SyncDatabaseConfig[PsycopgSyncConnection, ConnectionPool
         if self.kwargs:
             config.update(self.kwargs)
 
-        # Set DictRow as the row factory
         config["row_factory"] = dict_row
 
         return config
@@ -263,7 +260,6 @@ class PsycopgSyncConfig(SyncDatabaseConfig[PsycopgSyncConnection, ConnectionPool
         if self.kwargs:
             config.update(self.kwargs)
 
-        # Set DictRow as the row factory
         config["row_factory"] = dict_row
 
         return config
@@ -273,7 +269,6 @@ class PsycopgSyncConfig(SyncDatabaseConfig[PsycopgSyncConnection, ConnectionPool
         logger.info("Creating Psycopg connection pool", extra={"adapter": "psycopg"})
 
         try:
-            # Get all config (creates a new dict)
             all_config = self.pool_config_dict.copy()
 
             # Separate pool-specific parameters that ConnectionPool accepts directly
@@ -289,30 +284,29 @@ class PsycopgSyncConfig(SyncDatabaseConfig[PsycopgSyncConnection, ConnectionPool
                 "num_workers": all_config.pop("num_workers", 3),
             }
 
-            # Create a configure callback to set row_factory
+            # Capture autocommit setting before configuring the pool
+            autocommit_setting = all_config.get("autocommit")
+
             def configure_connection(conn: "PsycopgSyncConnection") -> None:
-                # Set DictRow as the row factory
                 conn.row_factory = dict_row
+                # Apply autocommit setting if specified
+                if autocommit_setting is not None:
+                    conn.autocommit = autocommit_setting
 
             pool_params["configure"] = all_config.pop("configure", configure_connection)
 
-            # Remove None values from pool_params
             pool_params = {k: v for k, v in pool_params.items() if v is not None}
 
-            # Handle conninfo vs individual connection parameters
             conninfo = all_config.pop("conninfo", None)
             if conninfo:
                 # If conninfo is provided, use it directly
                 # Don't pass kwargs when using conninfo string
-                pool = ConnectionPool(conninfo, **pool_params)
+                pool = ConnectionPool(conninfo, open=True, **pool_params)
             else:
-                # Otherwise, pass connection parameters via kwargs
-                # Remove any non-connection parameters
                 # row_factory is already popped out earlier
                 all_config.pop("row_factory", None)
-                # Remove pool-specific settings that may have been left
                 all_config.pop("kwargs", None)
-                pool = ConnectionPool("", kwargs=all_config, **pool_params)
+                pool = ConnectionPool("", kwargs=all_config, open=True, **pool_params)
 
             logger.info("Psycopg connection pool created successfully", extra={"adapter": "psycopg"})
         except Exception as e:
@@ -328,11 +322,17 @@ class PsycopgSyncConfig(SyncDatabaseConfig[PsycopgSyncConnection, ConnectionPool
         logger.info("Closing Psycopg connection pool", extra={"adapter": "psycopg"})
 
         try:
+            # This avoids the "cannot join current thread" error during garbage collection
+            if hasattr(self.pool_instance, "_closed"):
+                self.pool_instance._closed = True
+
             self.pool_instance.close()
             logger.info("Psycopg connection pool closed successfully", extra={"adapter": "psycopg"})
         except Exception as e:
             logger.exception("Failed to close Psycopg connection pool", extra={"adapter": "psycopg", "error": str(e)})
             raise
+        finally:
+            self.pool_instance = None
 
     def create_connection(self) -> "PsycopgSyncConnection":
         """Create a single connection (not from pool).
@@ -377,15 +377,16 @@ class PsycopgSyncConfig(SyncDatabaseConfig[PsycopgSyncConnection, ConnectionPool
             A PsycopgSyncDriver instance.
         """
         with self.provide_connection(*args, **kwargs) as conn:
-            # Create statement config with parameter style info if not already set
             statement_config = self.statement_config
+            # Inject parameter style info if not already set
             if statement_config.allowed_parameter_styles is None:
+                from dataclasses import replace
+
                 statement_config = replace(
                     statement_config,
                     allowed_parameter_styles=self.supported_parameter_styles,
                     target_parameter_style=self.preferred_parameter_style,
                 )
-
             driver = self.driver_type(connection=conn, config=statement_config)
             yield driver
 
@@ -547,7 +548,6 @@ class PsycopgAsyncConfig(AsyncDatabaseConfig[PsycopgAsyncConnection, AsyncConnec
         self.configure = configure
         self.kwargs = kwargs or {}
 
-        # Handle extras and additional kwargs
         self.extras = extras or {}
         self.extras.update(additional_kwargs)
 
@@ -576,7 +576,6 @@ class PsycopgAsyncConfig(AsyncDatabaseConfig[PsycopgAsyncConnection, AsyncConnec
         if self.kwargs:
             config.update(self.kwargs)
 
-        # Set DictRow as the row factory
         config["row_factory"] = dict_row
 
         return config
@@ -599,7 +598,6 @@ class PsycopgAsyncConfig(AsyncDatabaseConfig[PsycopgAsyncConnection, AsyncConnec
         if self.kwargs:
             config.update(self.kwargs)
 
-        # Set DictRow as the row factory
         config["row_factory"] = dict_row
 
         return config
@@ -607,7 +605,6 @@ class PsycopgAsyncConfig(AsyncDatabaseConfig[PsycopgAsyncConnection, AsyncConnec
     async def _create_pool(self) -> "AsyncConnectionPool":
         """Create the actual async connection pool."""
 
-        # Get all config (creates a new dict)
         all_config = self.pool_config_dict.copy()
 
         # Separate pool-specific parameters that AsyncConnectionPool accepts directly
@@ -623,28 +620,27 @@ class PsycopgAsyncConfig(AsyncDatabaseConfig[PsycopgAsyncConnection, AsyncConnec
             "num_workers": all_config.pop("num_workers", 3),
         }
 
-        # Create a configure callback to set row_factory
+        # Capture autocommit setting before configuring the pool
+        autocommit_setting = all_config.get("autocommit")
+
         async def configure_connection(conn: "PsycopgAsyncConnection") -> None:
-            # Set DictRow as the row factory
             conn.row_factory = dict_row
+            # Apply autocommit setting if specified (async version requires await)
+            if autocommit_setting is not None:
+                await conn.set_autocommit(autocommit_setting)
 
         pool_params["configure"] = all_config.pop("configure", configure_connection)
 
-        # Remove None values from pool_params
         pool_params = {k: v for k, v in pool_params.items() if v is not None}
 
-        # Handle conninfo vs individual connection parameters
         conninfo = all_config.pop("conninfo", None)
         if conninfo:
             # If conninfo is provided, use it directly
             # Don't pass kwargs when using conninfo string
             pool = AsyncConnectionPool(conninfo, open=False, **pool_params)
         else:
-            # Otherwise, pass connection parameters via kwargs
-            # Remove any non-connection parameters
             # row_factory is already popped out earlier
             all_config.pop("row_factory", None)
-            # Remove pool-specific settings that may have been left
             all_config.pop("kwargs", None)
             pool = AsyncConnectionPool("", kwargs=all_config, open=False, **pool_params)
 
@@ -657,7 +653,14 @@ class PsycopgAsyncConfig(AsyncDatabaseConfig[PsycopgAsyncConnection, AsyncConnec
         if not self.pool_instance:
             return
 
-        await self.pool_instance.close()
+        try:
+            # This avoids the "cannot join current thread" error during garbage collection
+            if hasattr(self.pool_instance, "_closed"):
+                self.pool_instance._closed = True
+
+            await self.pool_instance.close()
+        finally:
+            self.pool_instance = None
 
     async def create_connection(self) -> "PsycopgAsyncConnection":  # pyright: ignore
         """Create a single async connection (not from pool).
@@ -702,15 +705,16 @@ class PsycopgAsyncConfig(AsyncDatabaseConfig[PsycopgAsyncConnection, AsyncConnec
             A PsycopgAsyncDriver instance.
         """
         async with self.provide_connection(*args, **kwargs) as conn:
-            # Create statement config with parameter style info if not already set
             statement_config = self.statement_config
+            # Inject parameter style info if not already set
             if statement_config.allowed_parameter_styles is None:
+                from dataclasses import replace
+
                 statement_config = replace(
                     statement_config,
                     allowed_parameter_styles=self.supported_parameter_styles,
                     target_parameter_style=self.preferred_parameter_style,
                 )
-
             driver = self.driver_type(connection=conn, config=statement_config)
             yield driver