feat: adds a bigquery config and driver (#28)

Implements a `bigquery` config and driver with Export to parquet and `select_arrow`

Author: cofin

pyproject.toml

@@ -24,7 +24,7 @@ msgspec = ["msgspec"]
 nanoid = ["fastnanoid>=0.4.1"]
 oracledb = ["oracledb"]
 orjson = ["orjson"]
-performance = ["sqlglot[rs]"]
+performance = ["sqlglot[rs]", "msgspec"]
 psqlpy = ["psqlpy"]
 psycopg = ["psycopg[binary,pool]"]
 pydantic = ["pydantic", "pydantic-extra-types"]

sqlspec/adapters/bigquery/__init__.py

@@ -0,0 +1,4 @@
+from sqlspec.adapters.bigquery.config import BigQueryConfig, BigQueryConnectionConfig
+from sqlspec.adapters.bigquery.driver import BigQueryConnection, BigQueryDriver
+
+__all__ = ("BigQueryConfig", "BigQueryConnection", "BigQueryConnectionConfig", "BigQueryDriver")

sqlspec/adapters/bigquery/config/__init__.py

@@ -0,0 +1,3 @@
+from sqlspec.adapters.bigquery.config._sync import BigQueryConfig, BigQueryConnectionConfig
+
+__all__ = ("BigQueryConfig", "BigQueryConnectionConfig")

sqlspec/adapters/bigquery/config/_common.py

@@ -0,0 +1,40 @@
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Optional
+
+from google.cloud.bigquery import LoadJobConfig, QueryJobConfig
+
+if TYPE_CHECKING:
+    from google.api_core.client_info import ClientInfo
+    from google.api_core.client_options import ClientOptions
+    from google.auth.credentials import Credentials
+
+__all__ = ("BigQueryConnectionConfigCommon",)
+
+
+@dataclass
+class BigQueryConnectionConfigCommon:
+    """Common configuration options for BigQuery."""
+
+    project: "Optional[str]" = field(default=None)
+    """Google Cloud project ID."""
+    location: "Optional[str]" = field(default=None)
+    """Default geographic location for jobs and datasets."""
+    credentials: "Optional[Credentials]" = field(default=None, hash=False)
+    """Credentials to use for authentication."""
+    dataset_id: "Optional[str]" = field(default=None)
+    """Default dataset ID to use if not specified in queries."""
+    credentials_path: "Optional[str]" = field(default=None)
+    """Path to Google Cloud service account key file (JSON). If None, attempts default authentication."""
+    client_options: "Optional[ClientOptions]" = field(default=None, hash=False)
+    """Client options used to set user options on the client (e.g., api_endpoint)."""
+    default_query_job_config: "Optional[QueryJobConfig]" = field(default=None, hash=False)
+    """Default QueryJobConfig settings."""
+    default_load_job_config: "Optional[LoadJobConfig]" = field(default=None, hash=False)
+    """Default LoadJobConfig settings."""
+    client_info: "Optional[ClientInfo]" = field(default=None, hash=False)
+    """Client info used to send a user-agent string along with API requests."""
+
+    def __post_init__(self) -> None:
+        """Post-initialization hook."""
+        if self.default_query_job_config is None:
+            self.default_query_job_config = QueryJobConfig(default_dataset=self.dataset_id)

sqlspec/adapters/bigquery/config/_sync.py

@@ -0,0 +1,87 @@
+import contextlib
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, Optional
+
+from sqlspec.adapters.bigquery.config._common import BigQueryConnectionConfigCommon
+from sqlspec.adapters.bigquery.driver import BigQueryConnection, BigQueryDriver
+from sqlspec.base import NoPoolSyncConfig
+from sqlspec.typing import dataclass_to_dict
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
+__all__ = ("BigQueryConfig", "BigQueryConnectionConfig")
+
+
+class BigQueryConnectionConfig(BigQueryConnectionConfigCommon):
+    """BigQuery Connection Configuration."""
+
+
+@dataclass
+class BigQueryConfig(NoPoolSyncConfig["BigQueryConnection", "BigQueryDriver"]):
+    """BigQuery Synchronous Driver Configuration."""
+
+    connection_config: "BigQueryConnectionConfig" = field(default_factory=BigQueryConnectionConfig)
+    """BigQuery Connection Configuration."""
+    driver_type: "type[BigQueryDriver]" = field(init=False, repr=False, default=BigQueryDriver)
+    """BigQuery Driver Type."""
+    connection_type: "type[BigQueryConnection]" = field(init=False, repr=False, default=BigQueryConnection)
+    """BigQuery Connection Type."""
+    pool_instance: "None" = field(init=False, repr=False, default=None, hash=False)
+    """This is set to have a init=False since BigQuery does not support pooling."""
+    connection_instance: "Optional[BigQueryConnection]" = field(init=False, repr=False, default=None, hash=False)
+    """BigQuery Connection Instance."""
+
+    @property
+    def connection_config_dict(self) -> "dict[str, Any]":
+        """Return the connection configuration as a dict.
+
+        Returns:
+            A string keyed dict of config kwargs for the BigQueryConnection constructor.
+        """
+        return dataclass_to_dict(
+            self.connection_config,
+            exclude_empty=True,
+            exclude_none=True,
+            exclude={"dataset_id", "credentials_path"},
+        )
+
+    def create_connection(self) -> "BigQueryConnection":
+        """Create a BigQuery Client instance.
+
+        Returns:
+            A BigQuery Client instance.
+        """
+        if self.connection_instance is not None:
+            return self.connection_instance
+
+        self.connection_instance = self.connection_type(**self.connection_config_dict)
+        return self.connection_instance
+
+    @contextlib.contextmanager
+    def provide_connection(self, *args: Any, **kwargs: Any) -> "Iterator[BigQueryConnection]":
+        """Provide a BigQuery client within a context manager.
+
+        Args:
+            *args: Additional arguments to pass to the connection.
+            **kwargs: Additional keyword arguments to pass to the connection.
+
+        Yields:
+            An iterator of BigQuery Client instances.
+        """
+        conn = self.create_connection()
+        yield conn
+
+    @contextlib.contextmanager
+    def provide_session(self, *args: Any, **kwargs: Any) -> "Iterator[BigQueryDriver]":
+        """Provide a BigQuery driver session within a context manager.
+
+        Args:
+            *args: Additional arguments to pass to the driver.
+            **kwargs: Additional keyword arguments to pass to the driver.
+
+        Yields:
+            An iterator of BigQueryDriver instances.
+        """
+        conn = self.create_connection()
+        yield self.driver_type(connection=conn)