Feat: Add BigQuery cache support; Fix: IPython rendering bug when usi…

…ng a TTY (#65)

Co-authored-by: bindipankhudi <[email protected]>
Co-authored-by: Aaron Steers <[email protected]>

airbyte/__init__.py

@@ -7,6 +7,7 @@
 
 from airbyte import caches, datasets, registry, secrets
 from airbyte._factories.connector_factories import get_source
+from airbyte.caches.bigquery import BigQueryCache
 from airbyte.caches.duckdb import DuckDBCache
 from airbyte.caches.factories import get_default_cache, new_local_cache
 from airbyte.datasets import CachedDataset
@@ -29,6 +30,7 @@
     "get_source",
     "new_local_cache",
     # Classes
+    "BigQueryCache",
     "CachedDataset",
     "DuckDBCache",
     "ReadResult",

airbyte/_processors/sql/base.py

@@ -79,7 +79,6 @@ class SqlProcessorBase(RecordProcessor):
 
     # Constructor:
 
-    @final  # We don't want subclasses to have to override the constructor.
     def __init__(
         self,
         cache: CacheBase,
@@ -435,9 +434,6 @@ def _create_table(
         column_definition_str: str,
         primary_keys: list[str] | None = None,
     ) -> None:
-        if DEBUG_MODE:
-            assert table_name not in self._get_tables_list(), f"Table {table_name} already exists."
-
         if primary_keys:
             pk_str = ", ".join(primary_keys)
             column_definition_str += f",\n  PRIMARY KEY ({pk_str})"
@@ -448,11 +444,6 @@ def _create_table(
         )
         """
         _ = self._execute_sql(cmd)
-        if DEBUG_MODE:
-            tables_list = self._get_tables_list()
-            assert (
-                table_name in tables_list
-            ), f"Table {table_name} was not created. Found: {tables_list}"
 
     def _normalize_column_name(
         self,
@@ -804,8 +795,8 @@ def _merge_temp_table_to_final_table(
         columns = {self._quote_identifier(c) for c in self._get_sql_column_definitions(stream_name)}
         pk_columns = {self._quote_identifier(c) for c in self._get_primary_keys(stream_name)}
         non_pk_columns = columns - pk_columns
-        join_clause = "{nl} AND ".join(f"tmp.{pk_col} = final.{pk_col}" for pk_col in pk_columns)
-        set_clause = "{nl}    ".join(f"{col} = tmp.{col}" for col in non_pk_columns)
+        join_clause = f"{nl} AND ".join(f"tmp.{pk_col} = final.{pk_col}" for pk_col in pk_columns)
+        set_clause = f"{nl}  , ".join(f"{col} = tmp.{col}" for col in non_pk_columns)
         self._execute_sql(
             f"""
             MERGE INTO {self._fully_qualified(final_table_name)} final
@@ -908,12 +899,14 @@ def _emulated_merge_temp_table_to_final_table(
             conn.execute(update_stmt)
             conn.execute(insert_new_records_stmt)
 
-    @final
     def _table_exists(
         self,
         table_name: str,
     ) -> bool:
-        """Return true if the given table exists."""
+        """Return true if the given table exists.
+
+        Subclasses may override this method to provide a more efficient implementation.
+        """
         return table_name in self._get_tables_list()
 
     @abc.abstractmethod

airbyte/_processors/sql/bigquery.py

@@ -0,0 +1,204 @@
+# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+"""A BigQuery implementation of the cache."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import TYPE_CHECKING, final
+
+import sqlalchemy
+from google.api_core.exceptions import NotFound
+from google.cloud import bigquery
+from google.oauth2 import service_account
+from overrides import overrides
+
+from airbyte import exceptions as exc
+from airbyte._processors.file.jsonl import JsonlWriter
+from airbyte._processors.sql.base import SqlProcessorBase
+from airbyte.telemetry import CacheTelemetryInfo
+from airbyte.types import SQLTypeConverter
+
+
+if TYPE_CHECKING:
+    from sqlalchemy.engine.reflection import Inspector
+
+    from airbyte._processors.file.base import FileWriterBase
+    from airbyte.caches.base import CacheBase
+    from airbyte.caches.bigquery import BigQueryCache
+
+
+class BigQueryTypeConverter(SQLTypeConverter):
+    """A class to convert types for BigQuery."""
+
+    @overrides
+    def to_sql_type(
+        self,
+        json_schema_property_def: dict[str, str | dict | list],
+    ) -> sqlalchemy.types.TypeEngine:
+        """Convert a value to a SQL type.
+
+        We first call the parent class method to get the type. Then if the type is VARCHAR or
+        BIGINT, we replace it with respective BigQuery types.
+        """
+        sql_type = super().to_sql_type(json_schema_property_def)
+        # to-do: replace hardcoded return types with some sort of snowflake Variant equivalent
+        if isinstance(sql_type, sqlalchemy.types.VARCHAR):
+            return "String"
+        if isinstance(sql_type, sqlalchemy.types.BIGINT):
+            return "INT64"
+
+        return sql_type.__class__.__name__
+
+
+class BigQuerySqlProcessor(SqlProcessorBase):
+    """A BigQuery implementation of the cache."""
+
+    file_writer_class = JsonlWriter
+    type_converter_class = BigQueryTypeConverter
+    supports_merge_insert = True
+
+    cache: BigQueryCache
+
+    def __init__(self, cache: CacheBase, file_writer: FileWriterBase | None = None) -> None:
+        self._credentials: service_account.Credentials | None = None
+        self._schema_exists: bool | None = None
+        super().__init__(cache, file_writer)
+
+    @final
+    @overrides
+    def _fully_qualified(
+        self,
+        table_name: str,
+    ) -> str:
+        """Return the fully qualified name of the given table."""
+        return f"`{self.cache.schema_name}`.`{table_name!s}`"
+
+    @final
+    @overrides
+    def _quote_identifier(self, identifier: str) -> str:
+        """Return the identifier name as is. BigQuery does not require quoting identifiers"""
+        return f"{identifier}"
+
+    @final
+    @overrides
+    def _get_telemetry_info(self) -> CacheTelemetryInfo:
+        return CacheTelemetryInfo("bigquery")
+
+    def _write_files_to_new_table(
+        self,
+        files: list[Path],
+        stream_name: str,
+        batch_id: str,
+    ) -> str:
+        """Write a file(s) to a new table.
+
+        This is a generic implementation, which can be overridden by subclasses
+        to improve performance.
+        """
+        temp_table_name = self._create_table_for_loading(stream_name, batch_id)
+
+        # Specify the table ID (in the format `project_id.dataset_id.table_id`)
+        table_id = f"{self.cache.project_name}.{self.cache.dataset_name}.{temp_table_name}"
+
+        # Initialize a BigQuery client
+        client = bigquery.Client(credentials=self._get_credentials())
+
+        for file_path in files:
+            with Path.open(file_path, "rb") as source_file:
+                load_job = client.load_table_from_file(  # Make an API request
+                    file_obj=source_file,
+                    destination=table_id,
+                    job_config=bigquery.LoadJobConfig(
+                        source_format=bigquery.SourceFormat.NEWLINE_DELIMITED_JSON,
+                        schema=[
+                            bigquery.SchemaField(name, field_type=str(type_))
+                            for name, type_ in self._get_sql_column_definitions(
+                                stream_name=stream_name
+                            ).items()
+                        ],
+                    ),
+                )
+            _ = load_job.result()  # Wait for the job to complete
+
+        return temp_table_name
+
+    def _ensure_schema_exists(
+        self,
+    ) -> None:
+        """Ensure the target schema exists.
+
+        We override the default implementation because BigQuery is very slow at scanning schemas.
+
+        This implementation simply calls "CREATE SCHEMA IF NOT EXISTS" and ignores any errors.
+        """
+        if self._schema_exists:
+            return
+
+        sql = f"CREATE SCHEMA IF NOT EXISTS {self.cache.schema_name}"
+        try:
+            self._execute_sql(sql)
+        except Exception as ex:
+            # Ignore schema exists errors.
+            if "already exists" not in str(ex):
+                raise
+
+        self._schema_exists = True
+
+    def _get_credentials(self) -> service_account.Credentials:
+        """Return the GCP credentials."""
+        if self._credentials is None:
+            self._credentials = service_account.Credentials.from_service_account_file(
+                self.cache.credentials_path
+            )
+
+        return self._credentials
+
+    def _table_exists(
+        self,
+        table_name: str,
+    ) -> bool:
+        """Return true if the given table exists.
+
+        We override the default implementation because BigQuery is very slow at scanning tables.
+        """
+        client = bigquery.Client(credentials=self._get_credentials())
+        table_id = f"{self.cache.project_name}.{self.cache.dataset_name}.{table_name}"
+        try:
+            client.get_table(table_id)
+        except NotFound:
+            return False
+
+        except ValueError as ex:
+            raise exc.AirbyteLibInputError(
+                message="Invalid project name or dataset name.",
+                context={
+                    "table_id": table_id,
+                    "table_name": table_name,
+                    "project_name": self.cache.project_name,
+                    "dataset_name": self.cache.dataset_name,
+                },
+            ) from ex
+
+        return True
+
+    @final
+    @overrides
+    def _get_tables_list(
+        self,
+    ) -> list[str]:
+        """Get the list of available tables in the schema.
+
+        For bigquery, {schema_name}.{table_name} is returned, so we need to
+        strip the schema name in front of the table name, if it exists.
+
+        Warning: This method is slow for BigQuery, as it needs to scan all tables in the dataset.
+        It has been observed to take 30+ seconds in some cases.
+        """
+        with self.get_sql_connection() as conn:
+            inspector: Inspector = sqlalchemy.inspect(conn)
+            tables = inspector.get_table_names(schema=self.cache.schema_name)
+            schema_prefix = f"{self.cache.schema_name}."
+            return [
+                table.replace(schema_prefix, "", 1) if table.startswith(schema_prefix) else table
+                for table in tables
+            ]

airbyte/caches/_catalog_manager.py

@@ -4,12 +4,13 @@
 from __future__ import annotations
 
 import json
+from datetime import datetime
 from typing import TYPE_CHECKING, Callable
 
+from pytz import utc
 from sqlalchemy import Column, DateTime, String
 from sqlalchemy.ext.declarative import declarative_base
 from sqlalchemy.orm import Session
-from sqlalchemy.sql import func
 
 from airbyte_protocol.models import (
     AirbyteStateMessage,
@@ -50,7 +51,9 @@ class StreamState(Base):  # type: ignore[valid-type,misc]
     stream_name = Column(String)
     table_name = Column(String, primary_key=True)
     state_json = Column(String)
-    last_updated = Column(DateTime(timezone=True), onupdate=func.now(), default=func.now())
+    last_updated = Column(
+        DateTime(timezone=True), onupdate=datetime.now(utc), default=datetime.now(utc)
+    )
 
 
 class CatalogManager:

airbyte/caches/bigquery.py

@@ -0,0 +1,38 @@
+# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+"""A BigQuery implementation of the cache."""
+
+from __future__ import annotations
+
+import urllib
+
+from overrides import overrides
+
+from airbyte._processors.sql.bigquery import BigQuerySqlProcessor
+from airbyte.caches.base import (
+    CacheBase,
+)
+
+
+class BigQueryCache(CacheBase):
+    """The BigQuery cache implementation."""
+
+    project_name: str
+    dataset_name: str = "airbyte_raw"
+    credentials_path: str
+
+    _sql_processor_class: type[BigQuerySqlProcessor] = BigQuerySqlProcessor
+
+    def __post_init__(self) -> None:
+        """Initialize the BigQuery cache."""
+        self.schema_name = self.dataset_name
+
+    @overrides
+    def get_database_name(self) -> str:
+        """Return the name of the database. For BigQuery, this is the project name."""
+        return self.project_name
+
+    @overrides
+    def get_sql_alchemy_url(self) -> str:
+        """Return the SQLAlchemy URL to use."""
+        credentials_path_encoded = urllib.parse.quote(self.credentials_path)
+        return f"bigquery://{self.project_name!s}?credentials_path={credentials_path_encoded}"

airbyte/progress.py

@@ -4,6 +4,7 @@
 from __future__ import annotations
 
 import datetime
+import importlib
 import math
 import sys
 import time
@@ -25,10 +26,12 @@
 
 ipy_display: ModuleType | None
 try:
-    IS_NOTEBOOK = True
-    from IPython import display as ipy_display  # type: ignore  # noqa: PGH003
+    # Default to IS_NOTEBOOK=False if a TTY is detected.
+    IS_NOTEBOOK = not sys.stdout.isatty()
+    ipy_display = importlib.import_module("IPython.display")
 
 except ImportError:
+    # If IPython is not installed, then we're definitely not in a notebook.
     ipy_display = None
     IS_NOTEBOOK = False