sqlspec 0.17.1__py3-none-any.whl → 0.19.0__py3-none-any.whl

sqlspec/adapters/asyncpg/driver.py

@@ -31,7 +31,7 @@ __all__ = ("AsyncpgCursor", "AsyncpgDriver", "AsyncpgExceptionHandler", "asyncpg
 
 logger = get_logger("adapters.asyncpg")
 
-# Enhanced AsyncPG statement configuration using core modules with performance optimizations
+
 asyncpg_statement_config = StatementConfig(
     dialect="postgres",
     parameter_config=ParameterStyleConfig(
@@ -44,20 +44,19 @@ asyncpg_statement_config = StatementConfig(
         needs_static_script_compilation=False,
         preserve_parameter_format=True,
     ),
-    # Core processing features enabled for performance
     enable_parsing=True,
     enable_validation=True,
     enable_caching=True,
     enable_parameter_type_wrapping=True,
 )
 
-# PostgreSQL status parsing constants for row count extraction
+
 ASYNC_PG_STATUS_REGEX: Final[re.Pattern[str]] = re.compile(r"^([A-Z]+)(?:\s+(\d+))?\s+(\d+)$", re.IGNORECASE)
 EXPECTED_REGEX_GROUPS: Final[int] = 3
 
 
 class AsyncpgCursor:
-    """Context manager for AsyncPG cursor management with enhanced error handling."""
+    """Context manager for AsyncPG cursor management."""
 
     __slots__ = ("connection",)
 
@@ -68,12 +67,11 @@ class AsyncpgCursor:
         return self.connection
 
     async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
-        _ = (exc_type, exc_val, exc_tb)  # Mark as intentionally unused
-        # AsyncPG connections don't need explicit cursor cleanup
+        _ = (exc_type, exc_val, exc_tb)
 
 
 class AsyncpgExceptionHandler:
-    """Custom async context manager for handling AsyncPG database exceptions."""
+    """Async context manager for handling AsyncPG database exceptions."""
 
     __slots__ = ()
 
@@ -102,34 +100,15 @@ class AsyncpgExceptionHandler:
 
 
 class AsyncpgDriver(AsyncDriverAdapterBase):
-    """Enhanced AsyncPG PostgreSQL driver with CORE_ROUND_3 architecture integration.
-
-    This driver leverages the complete core module system for maximum performance:
-
-    Performance Improvements:
-    - 5-10x faster SQL compilation through single-pass processing
-    - 40-60% memory reduction through __slots__ optimization
-    - Enhanced caching for repeated statement execution
-    - Zero-copy parameter processing where possible
-    - Async-optimized resource management
-
-    Core Integration Features:
-    - sqlspec.core.statement for enhanced SQL processing
-    - sqlspec.core.parameters for optimized parameter handling
-    - sqlspec.core.cache for unified statement caching
-    - sqlspec.core.config for centralized configuration management
-
-    PostgreSQL Features:
-    - Advanced COPY operation support
-    - Numeric parameter style optimization
-    - PostgreSQL-specific exception handling
-    - Transaction management with async patterns
-
-    Compatibility:
-    - 100% backward compatibility with existing AsyncPG driver interface
-    - All existing async tests pass without modification
-    - Complete StatementConfig API compatibility
-    - Preserved async patterns and exception handling
+    """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
     """
 
     __slots__ = ()
@@ -141,25 +120,23 @@ class AsyncpgDriver(AsyncDriverAdapterBase):
         statement_config: "Optional[StatementConfig]" = None,
         driver_features: "Optional[dict[str, Any]]" = None,
     ) -> None:
-        # Enhanced configuration with global settings integration
         if statement_config is None:
             cache_config = get_cache_config()
-            enhanced_config = asyncpg_statement_config.replace(
+            statement_config = asyncpg_statement_config.replace(
                 enable_caching=cache_config.compiled_cache_enabled,
-                enable_parsing=True,  # Default to enabled
-                enable_validation=True,  # Default to enabled
-                dialect="postgres",  # Use adapter-specific dialect
+                enable_parsing=True,
+                enable_validation=True,
+                dialect="postgres",
             )
-            statement_config = enhanced_config
 
         super().__init__(connection=connection, statement_config=statement_config, driver_features=driver_features)
 
     def with_cursor(self, connection: "AsyncpgConnection") -> "AsyncpgCursor":
-        """Create context manager for AsyncPG cursor with enhanced resource management."""
+        """Create context manager for AsyncPG cursor."""
         return AsyncpgCursor(connection)
 
     def handle_database_exceptions(self) -> "AbstractAsyncContextManager[None]":
-        """Enhanced async exception handling with detailed error categorization."""
+        """Handle database exceptions with PostgreSQL error codes."""
         return AsyncpgExceptionHandler()
 
     async def _try_special_handling(self, cursor: "AsyncpgConnection", statement: "SQL") -> "Optional[SQLResult]":
@@ -179,23 +156,21 @@ class AsyncpgDriver(AsyncDriverAdapterBase):
         return None
 
     async def _handle_copy_operation(self, cursor: "AsyncpgConnection", statement: "SQL") -> None:
-        """Handle PostgreSQL COPY operations with enhanced data processing.
+        """Handle PostgreSQL COPY operations.
 
-        Supports both COPY FROM STDIN and COPY TO STDOUT operations
-        with proper data format handling and error management.
+        Supports both COPY FROM STDIN and COPY TO STDOUT operations.
 
         Args:
             cursor: AsyncPG connection object
             statement: SQL statement with COPY operation
         """
-        # Get metadata for copy operation data if available
+
         metadata: dict[str, Any] = getattr(statement, "metadata", {})
         sql_text = statement.sql
 
         copy_data = metadata.get("postgres_copy_data")
 
         if copy_data:
-            # Process different data formats for COPY operations
             if isinstance(copy_data, dict):
                 data_str = (
                     str(next(iter(copy_data.values())))
@@ -207,25 +182,18 @@ class AsyncpgDriver(AsyncDriverAdapterBase):
             else:
                 data_str = str(copy_data)
 
-            # Handle COPY FROM STDIN operations with binary data support
             if "FROM STDIN" in sql_text.upper():
                 from io import BytesIO
 
                 data_io = BytesIO(data_str.encode("utf-8"))
                 await cursor.copy_from_query(sql_text, output=data_io)
             else:
-                # Standard COPY operation
                 await cursor.execute(sql_text)
         else:
-            # COPY without additional data - execute directly
             await cursor.execute(sql_text)
 
     async def _execute_script(self, cursor: "AsyncpgConnection", statement: "SQL") -> "ExecutionResult":
-        """Execute SQL script using enhanced statement splitting and parameter handling.
-
-        Uses core module optimization for statement parsing and parameter processing.
-        Handles PostgreSQL-specific script execution requirements.
-        """
+        """Execute SQL script with statement splitting and parameter handling."""
         sql, _ = self._get_compiled_sql(statement, self.statement_config)
         statements = self.split_script_statements(sql, statement.statement_config, strip_trailing_semicolon=True)
 
@@ -233,8 +201,6 @@ class AsyncpgDriver(AsyncDriverAdapterBase):
         last_result = None
 
         for stmt in statements:
-            # Execute each statement individually
-            # If parameters were embedded (static style), prepared_parameters will be None/empty
             result = await cursor.execute(stmt)
             last_result = result
             successful_count += 1
@@ -244,38 +210,28 @@ class AsyncpgDriver(AsyncDriverAdapterBase):
         )
 
     async def _execute_many(self, cursor: "AsyncpgConnection", statement: "SQL") -> "ExecutionResult":
-        """Execute SQL with multiple parameter sets using optimized batch processing.
-
-        Leverages AsyncPG's executemany for efficient batch operations with
-        core parameter processing for enhanced type handling and validation.
-        """
+        """Execute SQL with multiple parameter sets using AsyncPG's executemany."""
         sql, prepared_parameters = self._get_compiled_sql(statement, self.statement_config)
 
         if prepared_parameters:
-            # Use AsyncPG's efficient executemany for batch operations
             await cursor.executemany(sql, prepared_parameters)
-            # Calculate affected rows (AsyncPG doesn't provide direct rowcount for executemany)
+
             affected_rows = len(prepared_parameters)
         else:
-            # Handle empty parameter case - no operations to execute
             affected_rows = 0
 
         return self.create_execution_result(cursor, rowcount_override=affected_rows, is_many_result=True)
 
     async def _execute_statement(self, cursor: "AsyncpgConnection", statement: "SQL") -> "ExecutionResult":
-        """Execute single SQL statement with enhanced data handling and performance optimization.
+        """Execute single SQL statement.
 
-        Uses core processing for optimal parameter handling and result processing.
-        Handles both SELECT queries and non-SELECT operations efficiently.
+        Handles both SELECT queries and non-SELECT operations.
         """
         sql, prepared_parameters = self._get_compiled_sql(statement, self.statement_config)
 
-        # Enhanced SELECT result processing
         if statement.returns_rows():
-            # Use AsyncPG's fetch for SELECT operations
             records = await cursor.fetch(sql, *prepared_parameters) if prepared_parameters else await cursor.fetch(sql)
 
-            # Efficient data conversion from asyncpg Records to dicts
             data = [dict(record) for record in records]
             column_names = list(records[0].keys()) if records else []
 
@@ -283,10 +239,8 @@ class AsyncpgDriver(AsyncDriverAdapterBase):
                 cursor, selected_data=data, column_names=column_names, data_row_count=len(data), is_select_result=True
             )
 
-        # Enhanced non-SELECT result processing
         result = await cursor.execute(sql, *prepared_parameters) if prepared_parameters else await cursor.execute(sql)
 
-        # Parse AsyncPG status string for affected rows
         affected_rows = self._parse_asyncpg_status(result) if isinstance(result, str) else 0
 
         return self.create_execution_result(cursor, rowcount_override=affected_rows)
@@ -312,15 +266,14 @@ class AsyncpgDriver(AsyncDriverAdapterBase):
             groups = match.groups()
             if len(groups) >= EXPECTED_REGEX_GROUPS:
                 try:
-                    return int(groups[-1])  # Last group contains the row count
+                    return int(groups[-1])
                 except (ValueError, IndexError):
                     pass
 
         return 0
 
-    # Async transaction management with enhanced error handling
     async def begin(self) -> None:
-        """Begin a database transaction with enhanced error handling."""
+        """Begin a database transaction."""
         try:
             await self.connection.execute("BEGIN")
         except asyncpg.PostgresError as e:
@@ -328,7 +281,7 @@ class AsyncpgDriver(AsyncDriverAdapterBase):
             raise SQLSpecError(msg) from e
 
     async def rollback(self) -> None:
-        """Rollback the current transaction with enhanced error handling."""
+        """Rollback the current transaction."""
         try:
             await self.connection.execute("ROLLBACK")
         except asyncpg.PostgresError as e:
@@ -336,7 +289,7 @@ class AsyncpgDriver(AsyncDriverAdapterBase):
             raise SQLSpecError(msg) from e
 
     async def commit(self) -> None:
-        """Commit the current transaction with enhanced error handling."""
+        """Commit the current transaction."""
         try:
             await self.connection.execute("COMMIT")
         except asyncpg.PostgresError as e:

sqlspec/adapters/bigquery/config.py

@@ -1,4 +1,4 @@
-"""BigQuery database configuration with direct field-based configuration."""
+"""BigQuery database configuration."""
 
 import contextlib
 import logging
@@ -32,14 +32,12 @@ class BigQueryConnectionParams(TypedDict, total=False):
     Includes both official BigQuery client parameters and BigQuery-specific configuration options.
     """
 
-    # Official BigQuery client constructor parameters
     project: NotRequired[str]
     location: NotRequired[str]
     credentials: NotRequired["Credentials"]
     client_options: NotRequired["ClientOptions"]
     client_info: NotRequired["ClientInfo"]
 
-    # BigQuery-specific configuration options
     default_query_job_config: NotRequired[QueryJobConfig]
     default_load_job_config: NotRequired[LoadJobConfig]
     dataset_id: NotRequired[str]
@@ -71,6 +69,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]"]
@@ -80,11 +79,9 @@ __all__ = ("BigQueryConfig", "BigQueryConnectionParams", "BigQueryDriverFeatures
 
 
 class BigQueryConfig(NoPoolSyncConfig[BigQueryConnection, BigQueryDriver]):
-    """Enhanced BigQuery configuration with comprehensive feature support.
+    """BigQuery configuration.
 
-    BigQuery is Google Cloud's serverless, highly scalable data warehouse with
-    advanced analytics, machine learning, and AI capabilities. This configuration
-    supports all BigQuery features including:
+    Configuration for Google Cloud BigQuery connections.
     """
 
     driver_type: ClassVar[type[BigQueryDriver]] = BigQueryDriver
@@ -93,69 +90,29 @@ 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,
         driver_features: "Optional[Union[BigQueryDriverFeatures, dict[str, Any]]]" = None,
     ) -> None:
-        """Initialize BigQuery configuration with comprehensive feature support.
+        """Initialize BigQuery configuration.
 
         Args:
-            connection_config: Standard connection configuration parameters
-            connection_instance: Existing connection instance to use
+            connection_config: Connection configuration parameters
             migration_config: Migration configuration
             statement_config: Statement configuration override
-            driver_features: BigQuery-specific driver features and configurations
-
-        Example:
-            >>> # Basic BigQuery connection
-            >>> config = BigQueryConfig(
-            ...     connection_config={
-            ...         "project": "my-project",
-            ...         "location": "US",
-            ...     }
-            ... )
-
-            >>> # Advanced configuration with ML and AI features
-            >>> config = BigQueryConfig(
-            ...     connection_config={
-            ...         "project": "my-project",
-            ...         "location": "US",
-            ...         "enable_bigquery_ml": True,
-            ...         "enable_gemini_integration": True,
-            ...         "enable_dataframes": True,
-            ...         "enable_vector_search": True,
-            ...         "maximum_bytes_billed": 1000000000,  # 1GB limit
-            ...     }
-            ... )
-
-            >>> # Enterprise configuration with reservations
-            >>> config = BigQueryConfig(
-            ...     connection_config={
-            ...         "project": "my-project",
-            ...         "location": "US",
-            ...         "edition": "Enterprise Plus",
-            ...         "reservation_id": "my-reservation",
-            ...         "enable_continuous_queries": True,
-            ...         "enable_cross_cloud": True,
-            ...     }
-            ... )
+            driver_features: BigQuery-specific driver features
         """
 
-        # 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
+        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()
 
@@ -170,8 +127,8 @@ class BigQueryConfig(NoPoolSyncConfig[BigQueryConnection, BigQueryDriver]):
         )
 
     def _setup_default_job_config(self) -> None:
-        """Set up default job configuration based on connection config."""
-        # Check if already provided in connection_config
+        """Set up default job configuration."""
+
         if self.connection_config.get("default_query_job_config") is not None:
             return
 
@@ -186,14 +143,12 @@ class BigQueryConfig(NoPoolSyncConfig[BigQueryConnection, BigQueryDriver]):
         if use_query_cache is not None:
             job_config.use_query_cache = use_query_cache
         else:
-            job_config.use_query_cache = True  # Default to True
+            job_config.use_query_cache = True
 
-        # Configure cost controls
         maximum_bytes_billed = self.connection_config.get("maximum_bytes_billed")
         if maximum_bytes_billed is not None:
             job_config.maximum_bytes_billed = maximum_bytes_billed
 
-        # Configure timeouts
         query_timeout_ms = self.connection_config.get("query_timeout_ms")
         if query_timeout_ms is not None:
             job_config.job_timeout_ms = query_timeout_ms
@@ -214,7 +169,6 @@ class BigQueryConfig(NoPoolSyncConfig[BigQueryConnection, BigQueryDriver]):
             return self._connection_instance
 
         try:
-            # Filter out extra fields and keep only official BigQuery client constructor fields
             client_fields = {"project", "location", "credentials", "client_options", "client_info"}
             config_dict: dict[str, Any] = {
                 field: value
@@ -223,7 +177,6 @@ class BigQueryConfig(NoPoolSyncConfig[BigQueryConnection, BigQueryDriver]):
             }
             connection = self.connection_type(**config_dict)
 
-            # Store BigQuery-specific config in driver_features for driver access
             default_query_job_config = self.connection_config.get("default_query_job_config")
             if default_query_job_config is not None:
                 self.driver_features["default_query_job_config"] = default_query_job_config
@@ -232,13 +185,11 @@ class BigQueryConfig(NoPoolSyncConfig[BigQueryConnection, BigQueryDriver]):
             if default_load_job_config is not None:
                 self.driver_features["default_load_job_config"] = default_load_job_config
 
-            # Call connection create callback from driver features
             on_connection_create = self.driver_features.get("on_connection_create")
             if on_connection_create:
                 on_connection_create(connection)
 
             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}"
@@ -275,7 +226,6 @@ class BigQueryConfig(NoPoolSyncConfig[BigQueryConnection, BigQueryDriver]):
         """
 
         with self.provide_connection(*_args, **_kwargs) as connection:
-            # Use shared config or user-provided config or instance default
             final_statement_config = statement_config or self.statement_config
 
             driver = self.driver_type(
@@ -286,9 +236,6 @@ class BigQueryConfig(NoPoolSyncConfig[BigQueryConnection, BigQueryDriver]):
     def get_signature_namespace(self) -> "dict[str, type[Any]]":
         """Get the signature namespace for BigQuery types.
 
-        This provides all BigQuery-specific types that Litestar needs to recognize
-        to avoid serialization attempts.
-
         Returns:
             Dictionary mapping type names to types.
         """