sqlspec 0.25.0__py3-none-any.whl → 0.26.0__py3-none-any.whl

sqlspec/adapters/asyncpg/data_dictionary.py

@@ -0,0 +1,134 @@
+"""PostgreSQL-specific data dictionary for metadata queries via asyncpg."""
+
+import re
+from typing import TYPE_CHECKING, Callable, Optional, cast
+
+from sqlspec.driver import AsyncDataDictionaryBase, AsyncDriverAdapterBase, VersionInfo
+from sqlspec.utils.logging import get_logger
+
+if TYPE_CHECKING:
+    from sqlspec.adapters.asyncpg.driver import AsyncpgDriver
+
+logger = get_logger("adapters.asyncpg.data_dictionary")
+
+# Compiled regex patterns
+POSTGRES_VERSION_PATTERN = re.compile(r"PostgreSQL (\d+)\.(\d+)(?:\.(\d+))?")
+
+__all__ = ("PostgresAsyncDataDictionary",)
+
+
+class PostgresAsyncDataDictionary(AsyncDataDictionaryBase):
+    """PostgreSQL-specific async data dictionary."""
+
+    async def get_version(self, driver: AsyncDriverAdapterBase) -> "Optional[VersionInfo]":
+        """Get PostgreSQL database version information.
+
+        Args:
+            driver: Async database driver instance
+
+        Returns:
+            PostgreSQL version information or None if detection fails
+        """
+        asyncpg_driver = cast("AsyncpgDriver", driver)
+        version_str = await asyncpg_driver.select_value("SELECT version()")
+        if not version_str:
+            logger.warning("No PostgreSQL version information found")
+            return None
+
+        # Parse version like "PostgreSQL 15.3 on x86_64-pc-linux-gnu..."
+        version_match = POSTGRES_VERSION_PATTERN.search(str(version_str))
+        if not version_match:
+            logger.warning("Could not parse PostgreSQL version: %s", version_str)
+            return None
+
+        major = int(version_match.group(1))
+        minor = int(version_match.group(2))
+        patch = int(version_match.group(3)) if version_match.group(3) else 0
+
+        version_info = VersionInfo(major, minor, patch)
+        logger.debug("Detected PostgreSQL version: %s", version_info)
+        return version_info
+
+    async def get_feature_flag(self, driver: AsyncDriverAdapterBase, feature: str) -> bool:
+        """Check if PostgreSQL database supports a specific feature.
+
+        Args:
+            driver: Async database driver instance
+            feature: Feature name to check
+
+        Returns:
+            True if feature is supported, False otherwise
+        """
+        version_info = await self.get_version(driver)
+        if not version_info:
+            return False
+
+        feature_checks: dict[str, Callable[..., bool]] = {
+            "supports_json": lambda v: v >= VersionInfo(9, 2, 0),
+            "supports_jsonb": lambda v: v >= VersionInfo(9, 4, 0),
+            "supports_uuid": lambda _: True,  # UUID extension widely available
+            "supports_arrays": lambda _: True,  # PostgreSQL has excellent array support
+            "supports_returning": lambda v: v >= VersionInfo(8, 2, 0),
+            "supports_upsert": lambda v: v >= VersionInfo(9, 5, 0),  # ON CONFLICT
+            "supports_window_functions": lambda v: v >= VersionInfo(8, 4, 0),
+            "supports_cte": lambda v: v >= VersionInfo(8, 4, 0),
+            "supports_transactions": lambda _: True,
+            "supports_prepared_statements": lambda _: True,
+            "supports_schemas": lambda _: True,
+            "supports_partitioning": lambda v: v >= VersionInfo(10, 0, 0),
+        }
+
+        if feature in feature_checks:
+            return bool(feature_checks[feature](version_info))
+
+        return False
+
+    async def get_optimal_type(self, driver: AsyncDriverAdapterBase, type_category: str) -> str:
+        """Get optimal PostgreSQL type for a category.
+
+        Args:
+            driver: Async database driver instance
+            type_category: Type category
+
+        Returns:
+            PostgreSQL-specific type name
+        """
+        version_info = await self.get_version(driver)
+
+        if type_category == "json":
+            if version_info and version_info >= VersionInfo(9, 4, 0):
+                return "JSONB"  # Prefer JSONB over JSON
+            if version_info and version_info >= VersionInfo(9, 2, 0):
+                return "JSON"
+            return "TEXT"
+
+        type_map = {
+            "uuid": "UUID",
+            "boolean": "BOOLEAN",
+            "timestamp": "TIMESTAMP WITH TIME ZONE",
+            "text": "TEXT",
+            "blob": "BYTEA",
+            "array": "ARRAY",
+        }
+        return type_map.get(type_category, "TEXT")
+
+    def list_available_features(self) -> "list[str]":
+        """List available PostgreSQL feature flags.
+
+        Returns:
+            List of supported feature names
+        """
+        return [
+            "supports_json",
+            "supports_jsonb",
+            "supports_uuid",
+            "supports_arrays",
+            "supports_returning",
+            "supports_upsert",
+            "supports_window_functions",
+            "supports_cte",
+            "supports_transactions",
+            "supports_prepared_statements",
+            "supports_schemas",
+            "supports_partitioning",
+        ]

sqlspec/adapters/asyncpg/driver.py

@@ -4,6 +4,7 @@ Provides async PostgreSQL connectivity with parameter processing, resource manag
 PostgreSQL COPY operation support, and transaction management.
 """
 
+import datetime
 import re
 from typing import TYPE_CHECKING, Any, Final, Optional
 
@@ -23,6 +24,7 @@ if TYPE_CHECKING:
     from sqlspec.core.result import SQLResult
     from sqlspec.core.statement import SQL
     from sqlspec.driver import ExecutionResult
+    from sqlspec.driver._async import AsyncDataDictionaryBase
 
 __all__ = ("AsyncpgCursor", "AsyncpgDriver", "AsyncpgExceptionHandler", "asyncpg_statement_config")
 
@@ -36,7 +38,7 @@ asyncpg_statement_config = StatementConfig(
         supported_parameter_styles={ParameterStyle.NUMERIC, ParameterStyle.POSITIONAL_PYFORMAT},
         default_execution_parameter_style=ParameterStyle.NUMERIC,
         supported_execution_parameter_styles={ParameterStyle.NUMERIC},
-        type_coercion_map={},
+        type_coercion_map={datetime.datetime: lambda x: x, datetime.date: lambda x: x, datetime.time: lambda x: x},
         has_native_list_expansion=True,
         needs_static_script_compilation=False,
         preserve_parameter_format=True,
@@ -63,8 +65,7 @@ class AsyncpgCursor:
     async def __aenter__(self) -> "AsyncpgConnection":
         return self.connection
 
-    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
-        _ = (exc_type, exc_val, exc_tb)
+    async def __aexit__(self, *_: Any) -> None: ...
 
 
 class AsyncpgExceptionHandler:
@@ -104,7 +105,7 @@ class AsyncpgDriver(AsyncDriverAdapterBase):
     and caching, and parameter processing with type coercion.
     """
 
-    __slots__ = ()
+    __slots__ = ("_data_dictionary",)
     dialect = "postgres"
 
     def __init__(
@@ -123,6 +124,7 @@ class AsyncpgDriver(AsyncDriverAdapterBase):
             )
 
         super().__init__(connection=connection, statement_config=statement_config, driver_features=driver_features)
+        self._data_dictionary: Optional[AsyncDataDictionaryBase] = None
 
     def with_cursor(self, connection: "AsyncpgConnection") -> "AsyncpgCursor":
         """Create context manager for AsyncPG cursor."""
@@ -311,3 +313,16 @@ class AsyncpgDriver(AsyncDriverAdapterBase):
         except asyncpg.PostgresError as e:
             msg = f"Failed to commit async transaction: {e}"
             raise SQLSpecError(msg) from e
+
+    @property
+    def data_dictionary(self) -> "AsyncDataDictionaryBase":
+        """Get the data dictionary for this driver.
+
+        Returns:
+            Data dictionary instance for metadata queries
+        """
+        if self._data_dictionary is None:
+            from sqlspec.adapters.asyncpg.data_dictionary import PostgresAsyncDataDictionary
+
+            self._data_dictionary = PostgresAsyncDataDictionary()
+        return self._data_dictionary

sqlspec/adapters/bigquery/config.py

@@ -94,6 +94,7 @@ class BigQueryConfig(NoPoolSyncConfig[BigQueryConnection, BigQueryDriver]):
         migration_config: Optional[dict[str, Any]] = None,
         statement_config: "Optional[StatementConfig]" = None,
         driver_features: "Optional[Union[BigQueryDriverFeatures, dict[str, Any]]]" = None,
+        bind_key: "Optional[str]" = None,
     ) -> None:
         """Initialize BigQuery configuration.
 
@@ -102,6 +103,7 @@ class BigQueryConfig(NoPoolSyncConfig[BigQueryConnection, BigQueryDriver]):
             migration_config: Migration configuration
             statement_config: Statement configuration override
             driver_features: BigQuery-specific driver features
+            bind_key: Optional unique identifier for this configuration
         """
 
         self.connection_config: dict[str, Any] = dict(connection_config) if connection_config else {}
@@ -124,6 +126,7 @@ class BigQueryConfig(NoPoolSyncConfig[BigQueryConnection, BigQueryDriver]):
             migration_config=migration_config,
             statement_config=statement_config,
             driver_features=self.driver_features,
+            bind_key=bind_key,
         )
 
     def _setup_default_job_config(self) -> None:

sqlspec/adapters/bigquery/data_dictionary.py

@@ -0,0 +1,109 @@
+"""BigQuery-specific data dictionary for metadata queries."""
+
+from typing import Optional
+
+from sqlspec.driver import SyncDataDictionaryBase, SyncDriverAdapterBase, VersionInfo
+from sqlspec.utils.logging import get_logger
+
+logger = get_logger("adapters.bigquery.data_dictionary")
+
+__all__ = ("BigQuerySyncDataDictionary",)
+
+
+class BigQuerySyncDataDictionary(SyncDataDictionaryBase):
+    """BigQuery-specific sync data dictionary."""
+
+    def get_version(self, driver: SyncDriverAdapterBase) -> "Optional[VersionInfo]":
+        """Get BigQuery version information.
+
+        BigQuery is a cloud service without traditional versioning.
+        Returns a fixed version to indicate feature availability.
+
+        Args:
+            driver: BigQuery driver instance
+
+        Returns:
+            Fixed version info indicating current BigQuery capabilities
+        """
+        # BigQuery is a cloud service - return a fixed version
+        # indicating modern feature support
+        logger.debug("BigQuery cloud service - using fixed version")
+        return VersionInfo(1, 0, 0)
+
+    def get_feature_flag(self, driver: SyncDriverAdapterBase, feature: str) -> bool:
+        """Check if BigQuery supports a specific feature.
+
+        Args:
+            driver: BigQuery driver instance
+            feature: Feature name to check
+
+        Returns:
+            True if feature is supported, False otherwise
+        """
+        # BigQuery feature support based on current capabilities
+        feature_checks = {
+            "supports_json": True,  # Native JSON type
+            "supports_arrays": True,  # ARRAY types
+            "supports_structs": True,  # STRUCT types
+            "supports_geography": True,  # GEOGRAPHY type
+            "supports_returning": False,  # No RETURNING clause
+            "supports_upsert": True,  # MERGE statement
+            "supports_window_functions": True,
+            "supports_cte": True,
+            "supports_transactions": True,  # Multi-statement transactions
+            "supports_prepared_statements": True,
+            "supports_schemas": True,  # Datasets and projects
+            "supports_partitioning": True,  # Table partitioning
+            "supports_clustering": True,  # Table clustering
+            "supports_uuid": False,  # No native UUID, use STRING
+        }
+
+        return feature_checks.get(feature, False)
+
+    def get_optimal_type(self, driver: SyncDriverAdapterBase, type_category: str) -> str:
+        """Get optimal BigQuery type for a category.
+
+        Args:
+            driver: BigQuery driver instance
+            type_category: Type category
+
+        Returns:
+            BigQuery-specific type name
+        """
+        type_map = {
+            "json": "JSON",
+            "uuid": "STRING",
+            "boolean": "BOOL",
+            "timestamp": "TIMESTAMP",
+            "text": "STRING",
+            "blob": "BYTES",
+            "array": "ARRAY",
+            "struct": "STRUCT",
+            "geography": "GEOGRAPHY",
+            "numeric": "NUMERIC",
+            "bignumeric": "BIGNUMERIC",
+        }
+        return type_map.get(type_category, "STRING")
+
+    def list_available_features(self) -> "list[str]":
+        """List available BigQuery feature flags.
+
+        Returns:
+            List of supported feature names
+        """
+        return [
+            "supports_json",
+            "supports_arrays",
+            "supports_structs",
+            "supports_geography",
+            "supports_returning",
+            "supports_upsert",
+            "supports_window_functions",
+            "supports_cte",
+            "supports_transactions",
+            "supports_prepared_statements",
+            "supports_schemas",
+            "supports_partitioning",
+            "supports_clustering",
+            "supports_uuid",
+        ]

sqlspec/adapters/bigquery/driver.py

@@ -15,6 +15,7 @@ from google.cloud.bigquery import ArrayQueryParameter, QueryJob, QueryJobConfig,
 from google.cloud.exceptions import GoogleCloudError
 
 from sqlspec.adapters.bigquery._types import BigQueryConnection
+from sqlspec.adapters.bigquery.type_converter import BigQueryTypeConverter
 from sqlspec.core.cache import get_cache_config
 from sqlspec.core.parameters import ParameterStyle, ParameterStyleConfig
 from sqlspec.core.statement import StatementConfig
@@ -28,11 +29,14 @@ if TYPE_CHECKING:
 
     from sqlspec.core.result import SQLResult
     from sqlspec.core.statement import SQL
+    from sqlspec.driver._sync import SyncDataDictionaryBase
 
 logger = logging.getLogger(__name__)
 
 __all__ = ("BigQueryCursor", "BigQueryDriver", "BigQueryExceptionHandler", "bigquery_statement_config")
 
+_type_converter = BigQueryTypeConverter()
+
 
 _BQ_TYPE_MAP: dict[type, tuple[str, Optional[str]]] = {
     bool: ("BOOL", None),
@@ -134,7 +138,7 @@ bigquery_type_coercion_map = {
     bool: lambda x: x,
     int: lambda x: x,
     float: lambda x: x,
-    str: lambda x: x,
+    str: _type_converter.convert_if_detected,
     bytes: lambda x: x,
     datetime.datetime: lambda x: x,
     datetime.date: lambda x: x,
@@ -177,7 +181,7 @@ class BigQueryCursor:
     def __enter__(self) -> "BigQueryConnection":
         return self.connection
 
-    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+    def __exit__(self, *_: Any) -> None:
         """Clean up cursor resources including active QueryJobs."""
         if self.job is not None:
             try:
@@ -230,7 +234,7 @@ class BigQueryDriver(SyncDriverAdapterBase):
     type coercion, error handling, and query job management.
     """
 
-    __slots__ = ("_default_query_job_config",)
+    __slots__ = ("_data_dictionary", "_default_query_job_config")
     dialect = "bigquery"
 
     def __init__(
@@ -252,6 +256,7 @@ class BigQueryDriver(SyncDriverAdapterBase):
         self._default_query_job_config: Optional[QueryJobConfig] = (driver_features or {}).get(
             "default_query_job_config"
         )
+        self._data_dictionary: Optional[SyncDataDictionaryBase] = None
 
     def with_cursor(self, connection: "BigQueryConnection") -> "BigQueryCursor":
         """Create context manager for cursor management.
@@ -532,3 +537,16 @@ class BigQueryDriver(SyncDriverAdapterBase):
         cursor.job.result()
         affected_rows = cursor.job.num_dml_affected_rows or 0
         return self.create_execution_result(cursor, rowcount_override=affected_rows)
+
+    @property
+    def data_dictionary(self) -> "SyncDataDictionaryBase":
+        """Get the data dictionary for this driver.
+
+        Returns:
+            Data dictionary instance for metadata queries
+        """
+        if self._data_dictionary is None:
+            from sqlspec.adapters.bigquery.data_dictionary import BigQuerySyncDataDictionary
+
+            self._data_dictionary = BigQuerySyncDataDictionary()
+        return self._data_dictionary

sqlspec/adapters/bigquery/type_converter.py

@@ -0,0 +1,93 @@
+"""BigQuery-specific type conversion with UUID support.
+
+Provides specialized type handling for BigQuery, including UUID support
+for the native BigQuery driver.
+"""
+
+from typing import Any, Final, Optional
+from uuid import UUID
+
+from sqlspec.core.type_conversion import BaseTypeConverter, convert_uuid
+
+try:
+    from google.cloud.bigquery import ScalarQueryParameter
+except ImportError:
+    ScalarQueryParameter = None  # type: ignore[assignment,misc]
+
+# Enhanced BigQuery type mapping with UUID support
+BQ_TYPE_MAP: Final[dict[str, str]] = {
+    "str": "STRING",
+    "int": "INT64",
+    "float": "FLOAT64",
+    "bool": "BOOL",
+    "datetime": "DATETIME",
+    "date": "DATE",
+    "time": "TIME",
+    "UUID": "STRING",  # UUID as STRING in BigQuery
+    "uuid": "STRING",
+    "Decimal": "NUMERIC",
+    "bytes": "BYTES",
+    "list": "ARRAY",
+    "dict": "STRUCT",
+}
+
+
+class BigQueryTypeConverter(BaseTypeConverter):
+    """BigQuery-specific type conversion with UUID support.
+
+    Extends the base TypeDetector with BigQuery-specific functionality
+    including UUID parameter handling for the native BigQuery driver.
+    """
+
+    __slots__ = ()
+
+    def create_parameter(self, name: str, value: Any) -> Optional[Any]:
+        """Create BigQuery parameter with proper type mapping.
+
+        Args:
+            name: Parameter name.
+            value: Parameter value.
+
+        Returns:
+            ScalarQueryParameter for native BigQuery driver, None if not available.
+        """
+        if ScalarQueryParameter is None:
+            return None
+
+        if isinstance(value, UUID):
+            return ScalarQueryParameter(name, "STRING", str(value))
+
+        if isinstance(value, str):
+            detected_type = self.detect_type(value)
+            if detected_type == "uuid":
+                uuid_obj = convert_uuid(value)
+                return ScalarQueryParameter(name, "STRING", str(uuid_obj))
+
+        # Handle other types
+        param_type = BQ_TYPE_MAP.get(type(value).__name__, "STRING")
+        return ScalarQueryParameter(name, param_type, value)
+
+    def convert_bigquery_value(self, value: Any, column_type: str) -> Any:
+        """Convert BigQuery value based on column type.
+
+        Args:
+            value: Value to convert.
+            column_type: BigQuery column type.
+
+        Returns:
+            Converted value appropriate for the column type.
+        """
+        if column_type == "STRING" and isinstance(value, str):
+            # Try to detect if this is a special type
+            detected_type = self.detect_type(value)
+            if detected_type:
+                try:
+                    return self.convert_value(value, detected_type)
+                except Exception:
+                    # If conversion fails, return original value
+                    return value
+
+        return value
+
+
+__all__ = ("BQ_TYPE_MAP", "BigQueryTypeConverter")

sqlspec/adapters/duckdb/_types.py

@@ -1,6 +1,6 @@
 from typing import TYPE_CHECKING
 
-from duckdb import DuckDBPyConnection
+from duckdb import DuckDBPyConnection  # type: ignore[import-untyped]
 
 if TYPE_CHECKING:
     from typing_extensions import TypeAlias

sqlspec/adapters/duckdb/config.py

@@ -149,6 +149,7 @@ class DuckDBConfig(SyncDatabaseConfig[DuckDBConnection, DuckDBConnectionPool, Du
         migration_config: Optional[dict[str, Any]] = None,
         statement_config: "Optional[StatementConfig]" = None,
         driver_features: "Optional[Union[DuckDBDriverFeatures, dict[str, Any]]]" = None,
+        bind_key: "Optional[str]" = None,
     ) -> None:
         """Initialize DuckDB configuration."""
         if pool_config is None:
@@ -160,6 +161,7 @@ class DuckDBConfig(SyncDatabaseConfig[DuckDBConnection, DuckDBConnectionPool, Du
             pool_config["database"] = ":memory:shared_db"
 
         super().__init__(
+            bind_key=bind_key,
             pool_config=dict(pool_config),
             pool_instance=pool_instance,
             migration_config=migration_config,

sqlspec/adapters/duckdb/data_dictionary.py

@@ -0,0 +1,124 @@
+"""DuckDB-specific data dictionary for metadata queries."""
+
+import re
+from typing import TYPE_CHECKING, Callable, Optional, cast
+
+from sqlspec.driver import SyncDataDictionaryBase, SyncDriverAdapterBase, VersionInfo
+from sqlspec.utils.logging import get_logger
+
+if TYPE_CHECKING:
+    from sqlspec.adapters.duckdb.driver import DuckDBDriver
+
+logger = get_logger("adapters.duckdb.data_dictionary")
+
+# Compiled regex patterns
+DUCKDB_VERSION_PATTERN = re.compile(r"v?(\d+)\.(\d+)\.(\d+)")
+
+__all__ = ("DuckDBSyncDataDictionary",)
+
+
+class DuckDBSyncDataDictionary(SyncDataDictionaryBase):
+    """DuckDB-specific sync data dictionary."""
+
+    def get_version(self, driver: SyncDriverAdapterBase) -> "Optional[VersionInfo]":
+        """Get DuckDB database version information.
+
+        Args:
+            driver: DuckDB driver instance
+
+        Returns:
+            DuckDB version information or None if detection fails
+        """
+        version_str = cast("DuckDBDriver", driver).select_value("SELECT version()")
+        if not version_str:
+            logger.warning("No DuckDB version information found")
+            return None
+
+        # Parse version like "v0.9.2" or "0.9.2"
+        version_match = DUCKDB_VERSION_PATTERN.search(str(version_str))
+        if not version_match:
+            logger.warning("Could not parse DuckDB version: %s", version_str)
+            return None
+
+        major, minor, patch = map(int, version_match.groups())
+        version_info = VersionInfo(major, minor, patch)
+        logger.debug("Detected DuckDB version: %s", version_info)
+        return version_info
+
+    def get_feature_flag(self, driver: SyncDriverAdapterBase, feature: str) -> bool:
+        """Check if DuckDB database supports a specific feature.
+
+        Args:
+            driver: DuckDB driver instance
+            feature: Feature name to check
+
+        Returns:
+            True if feature is supported, False otherwise
+        """
+        version_info = self.get_version(driver)
+        if not version_info:
+            return False
+
+        feature_checks: dict[str, Callable[..., bool]] = {
+            "supports_json": lambda _: True,  # DuckDB has excellent JSON support
+            "supports_arrays": lambda _: True,  # LIST type
+            "supports_maps": lambda _: True,  # MAP type
+            "supports_structs": lambda _: True,  # STRUCT type
+            "supports_returning": lambda v: v >= VersionInfo(0, 8, 0),
+            "supports_upsert": lambda v: v >= VersionInfo(0, 8, 0),
+            "supports_window_functions": lambda _: True,
+            "supports_cte": lambda _: True,
+            "supports_transactions": lambda _: True,
+            "supports_prepared_statements": lambda _: True,
+            "supports_schemas": lambda _: True,
+            "supports_uuid": lambda _: True,
+        }
+
+        if feature in feature_checks:
+            return bool(feature_checks[feature](version_info))
+
+        return False
+
+    def get_optimal_type(self, driver: SyncDriverAdapterBase, type_category: str) -> str:  # pyright: ignore
+        """Get optimal DuckDB type for a category.
+
+        Args:
+            driver: DuckDB driver instance
+            type_category: Type category
+
+        Returns:
+            DuckDB-specific type name
+        """
+        type_map = {
+            "json": "JSON",
+            "uuid": "UUID",
+            "boolean": "BOOLEAN",
+            "timestamp": "TIMESTAMP",
+            "text": "TEXT",
+            "blob": "BLOB",
+            "array": "LIST",
+            "map": "MAP",
+            "struct": "STRUCT",
+        }
+        return type_map.get(type_category, "VARCHAR")
+
+    def list_available_features(self) -> "list[str]":
+        """List available DuckDB feature flags.
+
+        Returns:
+            List of supported feature names
+        """
+        return [
+            "supports_json",
+            "supports_arrays",
+            "supports_maps",
+            "supports_structs",
+            "supports_returning",
+            "supports_upsert",
+            "supports_window_functions",
+            "supports_cte",
+            "supports_transactions",
+            "supports_prepared_statements",
+            "supports_schemas",
+            "supports_uuid",
+        ]