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

sqlspec/adapters/bigquery/config.py

@@ -2,7 +2,7 @@
 
 import contextlib
 import logging
-from typing import TYPE_CHECKING, Any, Callable, ClassVar, Optional, TypedDict, Union
+from typing import TYPE_CHECKING, Any, ClassVar, TypedDict
 
 from google.cloud.bigquery import LoadJobConfig, QueryJobConfig
 from typing_extensions import NotRequired
@@ -14,7 +14,7 @@ from sqlspec.exceptions import ImproperConfigurationError
 from sqlspec.typing import Empty
 
 if TYPE_CHECKING:
-    from collections.abc import Generator
+    from collections.abc import Callable, Generator
 
     from google.api_core.client_info import ClientInfo
     from google.api_core.client_options import ClientOptions
@@ -26,7 +26,7 @@ if TYPE_CHECKING:
 logger = logging.getLogger(__name__)
 
 
-class BigQueryConnectionParams(TypedDict, total=False):
+class BigQueryConnectionParams(TypedDict):
     """Standard BigQuery connection parameters.
 
     Includes both official BigQuery client parameters and BigQuery-specific configuration options.
@@ -63,16 +63,29 @@ class BigQueryConnectionParams(TypedDict, total=False):
     extra: NotRequired[dict[str, Any]]
 
 
-class BigQueryDriverFeatures(TypedDict, total=False):
+class BigQueryDriverFeatures(TypedDict):
     """BigQuery driver-specific features configuration.
 
     Only non-standard BigQuery client parameters that are SQLSpec-specific extensions.
+
+    Attributes:
+        connection_instance: Pre-existing BigQuery connection instance to use.
+        on_job_start: Callback invoked when a query job starts.
+        on_job_complete: Callback invoked when a query job completes.
+        on_connection_create: Callback invoked when connection is created.
+        json_serializer: Custom JSON serializer for dict/list parameter conversion.
+            Defaults to sqlspec.utils.serializers.to_json if not provided.
+        enable_uuid_conversion: Enable automatic UUID string conversion.
+            When True (default), UUID strings are automatically converted to UUID objects.
+            When False, UUID strings are treated as regular strings.
     """
 
     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]"]
+    json_serializer: NotRequired["Callable[[Any], str]"]
+    enable_uuid_conversion: NotRequired[bool]
 
 
 __all__ = ("BigQueryConfig", "BigQueryConnectionParams", "BigQueryDriverFeatures")
@@ -86,15 +99,17 @@ class BigQueryConfig(NoPoolSyncConfig[BigQueryConnection, BigQueryDriver]):
 
     driver_type: ClassVar[type[BigQueryDriver]] = BigQueryDriver
     connection_type: "ClassVar[type[BigQueryConnection]]" = BigQueryConnection
+    supports_transactional_ddl: ClassVar[bool] = False
 
     def __init__(
         self,
         *,
-        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,
-        bind_key: "Optional[str]" = None,
+        connection_config: "BigQueryConnectionParams | dict[str, Any] | None" = None,
+        migration_config: dict[str, Any] | None = None,
+        statement_config: "StatementConfig | None" = None,
+        driver_features: "BigQueryDriverFeatures | dict[str, Any] | None" = None,
+        bind_key: "str | None" = None,
+        extension_config: "dict[str, dict[str, Any]] | None" = None,
     ) -> None:
         """Initialize BigQuery configuration.
 
@@ -104,6 +119,7 @@ class BigQueryConfig(NoPoolSyncConfig[BigQueryConnection, BigQueryDriver]):
             statement_config: Statement configuration override
             driver_features: BigQuery-specific driver features
             bind_key: Optional unique identifier for this configuration
+            extension_config: Extension-specific configuration (e.g., Litestar plugin settings)
         """
 
         self.connection_config: dict[str, Any] = dict(connection_config) if connection_config else {}
@@ -113,7 +129,15 @@ class BigQueryConfig(NoPoolSyncConfig[BigQueryConnection, BigQueryDriver]):
 
         self.driver_features: dict[str, Any] = dict(driver_features) if driver_features else {}
 
-        self._connection_instance: Optional[BigQueryConnection] = self.driver_features.get("connection_instance")
+        if "enable_uuid_conversion" not in self.driver_features:
+            self.driver_features["enable_uuid_conversion"] = True
+
+        if "json_serializer" not in self.driver_features:
+            from sqlspec.utils.serializers import to_json
+
+            self.driver_features["json_serializer"] = to_json
+
+        self._connection_instance: BigQueryConnection | None = self.driver_features.get("connection_instance")
 
         if "default_query_job_config" not in self.connection_config:
             self._setup_default_job_config()
@@ -127,6 +151,7 @@ class BigQueryConfig(NoPoolSyncConfig[BigQueryConnection, BigQueryDriver]):
             statement_config=statement_config,
             driver_features=self.driver_features,
             bind_key=bind_key,
+            extension_config=extension_config,
         )
 
     def _setup_default_job_config(self) -> None:
@@ -215,7 +240,7 @@ class BigQueryConfig(NoPoolSyncConfig[BigQueryConnection, BigQueryDriver]):
 
     @contextlib.contextmanager
     def provide_session(
-        self, *_args: Any, statement_config: "Optional[StatementConfig]" = None, **_kwargs: Any
+        self, *_args: Any, statement_config: "StatementConfig | None" = None, **_kwargs: Any
     ) -> "Generator[BigQueryDriver, None, None]":
         """Provide a BigQuery driver session context manager.
 

sqlspec/adapters/bigquery/data_dictionary.py

@@ -1,10 +1,13 @@
 """BigQuery-specific data dictionary for metadata queries."""
 
-from typing import Optional
+from typing import TYPE_CHECKING, Any, cast
 
 from sqlspec.driver import SyncDataDictionaryBase, SyncDriverAdapterBase, VersionInfo
 from sqlspec.utils.logging import get_logger
 
+if TYPE_CHECKING:
+    from sqlspec.adapters.bigquery.driver import BigQueryDriver
+
 logger = get_logger("adapters.bigquery.data_dictionary")
 
 __all__ = ("BigQuerySyncDataDictionary",)
@@ -13,7 +16,7 @@ __all__ = ("BigQuerySyncDataDictionary",)
 class BigQuerySyncDataDictionary(SyncDataDictionaryBase):
     """BigQuery-specific sync data dictionary."""
 
-    def get_version(self, driver: SyncDriverAdapterBase) -> "Optional[VersionInfo]":
+    def get_version(self, driver: SyncDriverAdapterBase) -> "VersionInfo | None":
         """Get BigQuery version information.
 
         BigQuery is a cloud service without traditional versioning.
@@ -85,6 +88,43 @@ class BigQuerySyncDataDictionary(SyncDataDictionaryBase):
         }
         return type_map.get(type_category, "STRING")
 
+    def get_columns(
+        self, driver: SyncDriverAdapterBase, table: str, schema: "str | None" = None
+    ) -> "list[dict[str, Any]]":
+        """Get column information for a table using INFORMATION_SCHEMA.
+
+        Args:
+            driver: BigQuery driver instance
+            table: Table name to query columns for
+            schema: Schema name (dataset name in BigQuery)
+
+        Returns:
+            List of column metadata dictionaries with keys:
+                - column_name: Name of the column
+                - data_type: BigQuery data type
+                - is_nullable: Whether column allows NULL (YES/NO)
+                - column_default: Default value if any
+        """
+        bigquery_driver = cast("BigQueryDriver", driver)
+
+        if schema:
+            sql = f"""
+                SELECT column_name, data_type, is_nullable, column_default
+                FROM `{schema}.INFORMATION_SCHEMA.COLUMNS`
+                WHERE table_name = '{table}'
+                ORDER BY ordinal_position
+            """
+        else:
+            sql = f"""
+                SELECT column_name, data_type, is_nullable, column_default
+                FROM INFORMATION_SCHEMA.COLUMNS
+                WHERE table_name = '{table}'
+                ORDER BY ordinal_position
+            """
+
+        result = bigquery_driver.execute(sql)
+        return result.data or []
+
     def list_available_features(self) -> "list[str]":
         """List available BigQuery feature flags.