sqlspec 0.18.0__py3-none-any.whl → 0.20.0__py3-none-any.whl

sqlspec/adapters/adbc/driver.py

@@ -1,8 +1,7 @@
 """ADBC driver implementation for Arrow Database Connectivity.
 
-Provides ADBC driver integration with multi-dialect database connections,
-Arrow-native data handling with type coercion, parameter style conversion
-for different database backends, and transaction management.
+Provides database connectivity through ADBC with support for multiple
+database dialects, parameter style conversion, and transaction management.
 """
 
 import contextlib
@@ -10,6 +9,7 @@ import datetime
 import decimal
 from typing import TYPE_CHECKING, Any, Optional, cast
 
+from adbc_driver_manager.dbapi import DatabaseError, IntegrityError, OperationalError, ProgrammingError
 from sqlglot import exp
 
 from sqlspec.core.cache import get_cache_config
@@ -53,22 +53,88 @@ DIALECT_PARAMETER_STYLES = {
 }
 
 
-def _adbc_ast_transformer(expression: Any, parameters: Any) -> tuple[Any, Any]:
-    """AST transformer for NULL parameter handling.
+def _count_placeholders(expression: Any) -> int:
+    """Count the number of unique parameter placeholders in a SQLGlot expression.
 
-    For PostgreSQL, replaces NULL parameter placeholders with NULL literals
-    in the AST to prevent Arrow from inferring 'na' types which cause binding errors.
+    For PostgreSQL ($1, $2) style: counts highest numbered parameter (e.g., $1, $1, $2 = 2)
+    For QMARK (?) style: counts total occurrences (each ? is a separate parameter)
+    For named (:name) style: counts unique parameter names
 
     Args:
         expression: SQLGlot AST expression
-        parameters: Parameter values that may contain None
 
     Returns:
-        Tuple of (modified_expression, cleaned_parameters)
+        Number of unique parameter placeholders expected
     """
-    if not parameters:
-        return expression, parameters
+    numeric_params = set()  # For $1, $2 style
+    qmark_count = 0  # For ? style
+    named_params = set()  # For :name style
+
+    def count_node(node: Any) -> Any:
+        nonlocal qmark_count
+        if isinstance(node, exp.Parameter):
+            # PostgreSQL style: $1, $2, etc.
+            param_str = str(node)
+            if param_str.startswith("$") and param_str[1:].isdigit():
+                numeric_params.add(int(param_str[1:]))
+            elif ":" in param_str:
+                # Named parameter: :name
+                named_params.add(param_str)
+            else:
+                # Other parameter formats
+                named_params.add(param_str)
+        elif isinstance(node, exp.Placeholder):
+            # QMARK style: ?
+            qmark_count += 1
+        return node
+
+    expression.transform(count_node)
+
+    # Return the appropriate count based on parameter style detected
+    if numeric_params:
+        # PostgreSQL style: return highest numbered parameter
+        return max(numeric_params)
+    if named_params:
+        # Named parameters: return count of unique names
+        return len(named_params)
+    # QMARK style: return total count
+    return qmark_count
+
+
+def _is_execute_many_parameters(parameters: Any) -> bool:
+    """Check if parameters are in execute_many format (list/tuple of lists/tuples)."""
+    return isinstance(parameters, (list, tuple)) and len(parameters) > 0 and isinstance(parameters[0], (list, tuple))
+
 
+def _validate_parameter_counts(expression: Any, parameters: Any, dialect: str) -> None:
+    """Validate parameter count against placeholder count in SQL."""
+    placeholder_count = _count_placeholders(expression)
+    is_execute_many = _is_execute_many_parameters(parameters)
+
+    if is_execute_many:
+        # For execute_many, validate each inner parameter set
+        for i, param_set in enumerate(parameters):
+            param_count = len(param_set) if isinstance(param_set, (list, tuple)) else 0
+            if param_count != placeholder_count:
+                msg = f"Parameter count mismatch in set {i}: {param_count} parameters provided but {placeholder_count} placeholders in SQL (dialect: {dialect})"
+                raise SQLSpecError(msg)
+    else:
+        # For single execution, validate the parameter set directly
+        param_count = (
+            len(parameters)
+            if isinstance(parameters, (list, tuple))
+            else len(parameters)
+            if isinstance(parameters, dict)
+            else 0
+        )
+
+        if param_count != placeholder_count:
+            msg = f"Parameter count mismatch: {param_count} parameters provided but {placeholder_count} placeholders in SQL (dialect: {dialect})"
+            raise SQLSpecError(msg)
+
+
+def _find_null_positions(parameters: Any) -> set[int]:
+    """Find positions of None values in parameters for single execution."""
     null_positions = set()
     if isinstance(parameters, (list, tuple)):
         for i, param in enumerate(parameters):
@@ -83,7 +149,37 @@ def _adbc_ast_transformer(expression: Any, parameters: Any) -> tuple[Any, Any]:
                         null_positions.add(param_num - 1)
                 except ValueError:
                     pass
+    return null_positions
+
+
+def _adbc_ast_transformer(expression: Any, parameters: Any, dialect: str = "postgres") -> tuple[Any, Any]:
+    """Transform AST to handle NULL parameters.
 
+    Replaces NULL parameter placeholders with NULL literals in the AST
+    to prevent Arrow from inferring 'na' types which cause binding errors.
+    Validates parameter count before transformation.
+
+    Args:
+        expression: SQLGlot AST expression parsed with proper dialect
+        parameters: Parameter values that may contain None
+        dialect: SQLGlot dialect used for parsing (default: "postgres")
+
+    Returns:
+        Tuple of (modified_expression, cleaned_parameters)
+    """
+    if not parameters:
+        return expression, parameters
+
+    # Validate parameter count before transformation
+    _validate_parameter_counts(expression, parameters, dialect)
+
+    # For execute_many operations, skip AST transformation as different parameter
+    # sets may have None values in different positions, making transformation complex
+    if _is_execute_many_parameters(parameters):
+        return expression, parameters
+
+    # Find positions of None values for single execution
+    null_positions = _find_null_positions(parameters)
     if not null_positions:
         return expression, parameters
 
@@ -183,14 +279,28 @@ def get_adbc_statement_config(detected_dialect: str) -> StatementConfig:
 
 
 def _convert_array_for_postgres_adbc(value: Any) -> Any:
-    """Convert array values for PostgreSQL compatibility."""
+    """Convert array values for PostgreSQL compatibility.
+
+    Args:
+        value: Value to convert
+
+    Returns:
+        Converted value (tuples become lists)
+    """
     if isinstance(value, tuple):
         return list(value)
     return value
 
 
 def get_type_coercion_map(dialect: str) -> "dict[type, Any]":
-    """Get type coercion map for Arrow type handling."""
+    """Get type coercion map for Arrow type handling.
+
+    Args:
+        dialect: Database dialect name
+
+    Returns:
+        Mapping of Python types to conversion functions
+    """
     type_map = {
         datetime.datetime: lambda x: x,
         datetime.date: lambda x: x,
@@ -245,8 +355,6 @@ class AdbcExceptionHandler:
             return
 
         try:
-            from adbc_driver_manager.dbapi import DatabaseError, IntegrityError, OperationalError, ProgrammingError
-
             if issubclass(exc_type, IntegrityError):
                 e = exc_val
                 msg = f"Integrity constraint violation: {e}"
@@ -282,9 +390,8 @@ class AdbcExceptionHandler:
 class AdbcDriver(SyncDriverAdapterBase):
     """ADBC driver for Arrow Database Connectivity.
 
-    Provides database connectivity through ADBC with multi-database dialect
-    support, Arrow-native data handling with type coercion, parameter style
-    conversion for different backends, and transaction management.
+    Provides database connectivity through ADBC with support for multiple
+    database dialects, parameter style conversion, and transaction management.
     """
 
     __slots__ = ("_detected_dialect", "dialect")
@@ -309,7 +416,11 @@ class AdbcDriver(SyncDriverAdapterBase):
 
     @staticmethod
     def _ensure_pyarrow_installed() -> None:
-        """Ensure PyArrow is installed."""
+        """Ensure PyArrow is installed.
+
+        Raises:
+            MissingDependencyError: If PyArrow is not installed
+        """
         from sqlspec.typing import PYARROW_INSTALLED
 
         if not PYARROW_INSTALLED:
@@ -317,7 +428,14 @@ class AdbcDriver(SyncDriverAdapterBase):
 
     @staticmethod
     def _get_dialect(connection: "AdbcConnection") -> str:
-        """Detect database dialect from connection information."""
+        """Detect database dialect from connection information.
+
+        Args:
+            connection: ADBC connection
+
+        Returns:
+            Detected dialect name (defaults to 'postgres')
+        """
         try:
             driver_info = connection.adbc_get_info()
             vendor_name = driver_info.get("vendor_name", "").lower()
@@ -334,31 +452,53 @@ class AdbcDriver(SyncDriverAdapterBase):
         return "postgres"
 
     def _handle_postgres_rollback(self, cursor: "Cursor") -> None:
-        """Execute rollback for PostgreSQL after transaction failure."""
+        """Execute rollback for PostgreSQL after transaction failure.
+
+        Args:
+            cursor: Database cursor
+        """
         if self.dialect == "postgres":
             with contextlib.suppress(Exception):
                 cursor.execute("ROLLBACK")
                 logger.debug("PostgreSQL rollback executed after transaction failure")
 
     def _handle_postgres_empty_parameters(self, parameters: Any) -> Any:
-        """Process empty parameters for PostgreSQL compatibility."""
+        """Process empty parameters for PostgreSQL compatibility.
+
+        Args:
+            parameters: Parameter values
+
+        Returns:
+            Processed parameters
+        """
         if self.dialect == "postgres" and isinstance(parameters, dict) and not parameters:
             return None
         return parameters
 
     def with_cursor(self, connection: "AdbcConnection") -> "AdbcCursor":
-        """Create context manager for cursor."""
+        """Create context manager for cursor.
+
+        Args:
+            connection: Database connection
+
+        Returns:
+            Cursor context manager
+        """
         return AdbcCursor(connection)
 
     def handle_database_exceptions(self) -> "AbstractContextManager[None]":
-        """Handle database-specific exceptions and wrap them appropriately."""
+        """Handle database-specific exceptions and wrap them appropriately.
+
+        Returns:
+            Exception handler context manager
+        """
         return AdbcExceptionHandler()
 
     def _try_special_handling(self, cursor: "Cursor", statement: SQL) -> "Optional[SQLResult]":
         """Handle special operations.
 
         Args:
-            cursor: Cursor object
+            cursor: Database cursor
             statement: SQL statement to analyze
 
         Returns:
@@ -368,7 +508,15 @@ class AdbcDriver(SyncDriverAdapterBase):
         return None
 
     def _execute_many(self, cursor: "Cursor", statement: SQL) -> "ExecutionResult":
-        """Execute SQL with multiple parameter sets."""
+        """Execute SQL with multiple parameter sets.
+
+        Args:
+            cursor: Database cursor
+            statement: SQL statement to execute
+
+        Returns:
+            Execution result with row counts
+        """
         sql, prepared_parameters = self._get_compiled_sql(statement, self.statement_config)
 
         try:
@@ -398,7 +546,15 @@ class AdbcDriver(SyncDriverAdapterBase):
         return self.create_execution_result(cursor, rowcount_override=row_count, is_many_result=True)
 
     def _execute_statement(self, cursor: "Cursor", statement: SQL) -> "ExecutionResult":
-        """Execute single SQL statement."""
+        """Execute single SQL statement.
+
+        Args:
+            cursor: Database cursor
+            statement: SQL statement to execute
+
+        Returns:
+            Execution result with data or row count
+        """
         sql, prepared_parameters = self._get_compiled_sql(statement, self.statement_config)
 
         try:
@@ -430,7 +586,15 @@ class AdbcDriver(SyncDriverAdapterBase):
         return self.create_execution_result(cursor, rowcount_override=row_count)
 
     def _execute_script(self, cursor: "Cursor", statement: "SQL") -> "ExecutionResult":
-        """Execute SQL script."""
+        """Execute SQL script containing multiple statements.
+
+        Args:
+            cursor: Database cursor
+            statement: SQL script to execute
+
+        Returns:
+            Execution result with statement counts
+        """
         if statement.is_script:
             sql = statement._raw_sql
             prepared_parameters: list[Any] = []

sqlspec/adapters/asyncmy/driver.py

@@ -51,7 +51,10 @@ asyncmy_statement_config = StatementConfig(
 
 
 class AsyncmyCursor:
-    """Async context manager for AsyncMy cursor management."""
+    """Context manager for AsyncMy cursor operations.
+
+    Provides automatic cursor acquisition and cleanup for database operations.
+    """
 
     __slots__ = ("connection", "cursor")
 
@@ -70,7 +73,11 @@ class AsyncmyCursor:
 
 
 class AsyncmyExceptionHandler:
-    """Custom async context manager for handling AsyncMy database exceptions."""
+    """Context manager for AsyncMy database exception handling.
+
+    Converts AsyncMy-specific exceptions to SQLSpec exceptions with appropriate
+    error categorization and context preservation.
+    """
 
     __slots__ = ()
 
@@ -116,10 +123,11 @@ class AsyncmyExceptionHandler:
 
 
 class AsyncmyDriver(AsyncDriverAdapterBase):
-    """AsyncMy MySQL/MariaDB driver.
+    """MySQL/MariaDB database driver using AsyncMy client library.
 
-    Provides MySQL/MariaDB connectivity with parameter style conversion,
-    type coercion, error handling, and transaction management.
+    Implements asynchronous database operations for MySQL and MariaDB servers
+    with support for parameter style conversion, type coercion, error handling,
+    and transaction management.
     """
 
     __slots__ = ()
@@ -143,22 +151,33 @@ class AsyncmyDriver(AsyncDriverAdapterBase):
         super().__init__(connection=connection, statement_config=statement_config, driver_features=driver_features)
 
     def with_cursor(self, connection: "AsyncmyConnection") -> "AsyncmyCursor":
-        """Create context manager for AsyncMy cursor."""
+        """Create cursor context manager for the connection.
+
+        Args:
+            connection: AsyncMy database connection
+
+        Returns:
+            AsyncmyCursor: Context manager for cursor operations
+        """
         return AsyncmyCursor(connection)
 
     def handle_database_exceptions(self) -> "AbstractAsyncContextManager[None]":
-        """Handle database-specific exceptions and wrap them appropriately."""
+        """Provide exception handling context manager.
+
+        Returns:
+            AbstractAsyncContextManager[None]: Context manager for AsyncMy exception handling
+        """
         return AsyncmyExceptionHandler()
 
     async def _try_special_handling(self, cursor: Any, statement: "SQL") -> "Optional[SQLResult]":
-        """Hook for AsyncMy-specific special operations.
+        """Handle AsyncMy-specific operations before standard execution.
 
         Args:
             cursor: AsyncMy cursor object
             statement: SQL statement to analyze
 
         Returns:
-            None - always proceeds with standard execution for AsyncMy
+            Optional[SQLResult]: None, always proceeds with standard execution
         """
         _ = (cursor, statement)
         return None
@@ -166,7 +185,15 @@ class AsyncmyDriver(AsyncDriverAdapterBase):
     async def _execute_script(self, cursor: Any, statement: "SQL") -> "ExecutionResult":
         """Execute SQL script with statement splitting and parameter handling.
 
+        Splits multi-statement scripts and executes each statement sequentially.
         Parameters are embedded as static values for script execution compatibility.
+
+        Args:
+            cursor: AsyncMy cursor object
+            statement: SQL script to execute
+
+        Returns:
+            ExecutionResult: Script execution results with statement count
         """
         sql, prepared_parameters = self._get_compiled_sql(statement, self.statement_config)
         statements = self.split_script_statements(sql, statement.statement_config, strip_trailing_semicolon=True)
@@ -183,9 +210,20 @@ class AsyncmyDriver(AsyncDriverAdapterBase):
         )
 
     async def _execute_many(self, cursor: Any, statement: "SQL") -> "ExecutionResult":
-        """Execute SQL with multiple parameter sets using AsyncMy batch processing.
+        """Execute SQL statement with multiple parameter sets.
 
-        Handles MySQL type conversion and parameter processing.
+        Uses AsyncMy's executemany for batch operations with MySQL type conversion
+        and parameter processing.
+
+        Args:
+            cursor: AsyncMy cursor object
+            statement: SQL statement with multiple parameter sets
+
+        Returns:
+            ExecutionResult: Batch execution results
+
+        Raises:
+            ValueError: If no parameters provided for executemany operation
         """
         sql, prepared_parameters = self._get_compiled_sql(statement, self.statement_config)
 
@@ -200,9 +238,17 @@ class AsyncmyDriver(AsyncDriverAdapterBase):
         return self.create_execution_result(cursor, rowcount_override=affected_rows, is_many_result=True)
 
     async def _execute_statement(self, cursor: Any, statement: "SQL") -> "ExecutionResult":
-        """Execute single SQL statement with AsyncMy MySQL data handling.
+        """Execute single SQL statement.
+
+        Handles parameter processing, result fetching, and data transformation
+        for MySQL/MariaDB operations.
+
+        Args:
+            cursor: AsyncMy cursor object
+            statement: SQL statement to execute
 
-        Handles parameter processing and MySQL result processing.
+        Returns:
+            ExecutionResult: Statement execution results with data or row counts
         """
         sql, prepared_parameters = self._get_compiled_sql(statement, self.statement_config)
         await cursor.execute(sql, prepared_parameters or None)
@@ -228,6 +274,9 @@ class AsyncmyDriver(AsyncDriverAdapterBase):
         """Begin a database transaction.
 
         Explicitly starts a MySQL transaction to ensure proper transaction boundaries.
+
+        Raises:
+            SQLSpecError: If transaction initialization fails
         """
         try:
             async with AsyncmyCursor(self.connection) as cursor:
@@ -237,7 +286,11 @@ class AsyncmyDriver(AsyncDriverAdapterBase):
             raise SQLSpecError(msg) from e
 
     async def rollback(self) -> None:
-        """Rollback the current transaction."""
+        """Rollback the current transaction.
+
+        Raises:
+            SQLSpecError: If transaction rollback fails
+        """
         try:
             await self.connection.rollback()
         except asyncmy.errors.MySQLError as e:
@@ -245,7 +298,11 @@ class AsyncmyDriver(AsyncDriverAdapterBase):
             raise SQLSpecError(msg) from e
 
     async def commit(self) -> None:
-        """Commit the current transaction."""
+        """Commit the current transaction.
+
+        Raises:
+            SQLSpecError: If transaction commit fails
+        """
         try:
             await self.connection.commit()
         except asyncmy.errors.MySQLError as e:

sqlspec/adapters/asyncpg/config.py

@@ -124,12 +124,32 @@ class AsyncpgConfig(AsyncDatabaseConfig[AsyncpgConnection, "Pool[Record]", Async
         config = self._get_pool_config_dict()
 
         if "init" not in config:
-            config["init"] = self._init_pgvector_connection
+            config["init"] = self._init_connection
 
         return await asyncpg_create_pool(**config)
 
-    async def _init_pgvector_connection(self, connection: "AsyncpgConnection") -> None:
-        """Initialize pgvector support for asyncpg connections."""
+    async def _init_connection(self, connection: "AsyncpgConnection") -> None:
+        """Initialize connection with JSON codecs and pgvector support."""
+
+        try:
+            # Set up JSON type codec
+            await connection.set_type_codec(
+                "json",
+                encoder=self.driver_features.get("json_serializer", to_json),
+                decoder=self.driver_features.get("json_deserializer", from_json),
+                schema="pg_catalog",
+            )
+            # Set up JSONB type codec
+            await connection.set_type_codec(
+                "jsonb",
+                encoder=self.driver_features.get("json_serializer", to_json),
+                decoder=self.driver_features.get("json_deserializer", from_json),
+                schema="pg_catalog",
+            )
+        except Exception as e:
+            logger.debug("Failed to configure JSON type codecs for asyncpg: %s", e)
+
+        # Initialize pgvector support
         try:
             import pgvector.asyncpg
 

sqlspec/adapters/asyncpg/driver.py

@@ -1,10 +1,7 @@
 """AsyncPG PostgreSQL driver implementation for async PostgreSQL operations.
 
-Provides async PostgreSQL connectivity with:
-- Parameter processing with type coercion
-- Resource management
-- PostgreSQL COPY operation support
-- Transaction management
+Provides async PostgreSQL connectivity with parameter processing, resource management,
+PostgreSQL COPY operation support, and transaction management.
 """
 
 import re
@@ -102,13 +99,9 @@ class AsyncpgExceptionHandler:
 class AsyncpgDriver(AsyncDriverAdapterBase):
     """AsyncPG PostgreSQL driver for async database operations.
 
-    Features:
-    - COPY operation support
-    - Numeric parameter style handling
-    - PostgreSQL exception handling
-    - Transaction management
-    - SQL statement compilation and caching
-    - Parameter processing and type coercion
+    Supports COPY operations, numeric parameter style handling, PostgreSQL
+    exception handling, transaction management, SQL statement compilation
+    and caching, and parameter processing with type coercion.
     """
 
     __slots__ = ()
@@ -193,7 +186,15 @@ class AsyncpgDriver(AsyncDriverAdapterBase):
             await cursor.execute(sql_text)
 
     async def _execute_script(self, cursor: "AsyncpgConnection", statement: "SQL") -> "ExecutionResult":
-        """Execute SQL script with statement splitting and parameter handling."""
+        """Execute SQL script with statement splitting and parameter handling.
+
+        Args:
+            cursor: AsyncPG connection object
+            statement: SQL statement containing multiple statements
+
+        Returns:
+            ExecutionResult with script execution details
+        """
         sql, _ = self._get_compiled_sql(statement, self.statement_config)
         statements = self.split_script_statements(sql, statement.statement_config, strip_trailing_semicolon=True)
 
@@ -210,7 +211,15 @@ class AsyncpgDriver(AsyncDriverAdapterBase):
         )
 
     async def _execute_many(self, cursor: "AsyncpgConnection", statement: "SQL") -> "ExecutionResult":
-        """Execute SQL with multiple parameter sets using AsyncPG's executemany."""
+        """Execute SQL with multiple parameter sets using AsyncPG's executemany.
+
+        Args:
+            cursor: AsyncPG connection object
+            statement: SQL statement with multiple parameter sets
+
+        Returns:
+            ExecutionResult with batch execution details
+        """
         sql, prepared_parameters = self._get_compiled_sql(statement, self.statement_config)
 
         if prepared_parameters:
@@ -226,6 +235,13 @@ class AsyncpgDriver(AsyncDriverAdapterBase):
         """Execute single SQL statement.
 
         Handles both SELECT queries and non-SELECT operations.
+
+        Args:
+            cursor: AsyncPG connection object
+            statement: SQL statement to execute
+
+        Returns:
+            ExecutionResult with statement execution details
         """
         sql, prepared_parameters = self._get_compiled_sql(statement, self.statement_config)