Feat: Add MotherDuck cache support 🦆 (#43)

Author: aaronsteers

airbyte/_processors/sql/motherduck.py

@@ -0,0 +1,42 @@
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+"""A MotherDuck implementation of the cache, built on the DuckDB implementation."""
+
+from __future__ import annotations
+
+import warnings
+from typing import TYPE_CHECKING
+
+from overrides import overrides
+
+from airbyte._processors.file import JsonlWriter
+from airbyte._processors.sql.duckdb import DuckDBSqlProcessor
+
+
+if TYPE_CHECKING:
+    from airbyte.caches.motherduck import MotherDuckCache
+
+
+# Suppress warnings from DuckDB about reflection on indices.
+# https://github.com/Mause/duckdb_engine/issues/905
+warnings.filterwarnings(
+    "ignore",
+    message="duckdb-engine doesn't yet support reflection on indices",
+)
+
+
+class MotherDuckSqlProcessor(DuckDBSqlProcessor):
+    """A cache implementation for MotherDuck."""
+
+    supports_merge_insert = False
+    file_writer_class = JsonlWriter
+    cache: MotherDuckCache
+
+    @overrides
+    def _setup(self) -> None:
+        """Do any necessary setup, if applicable.
+
+        Note: The DuckDB parent class requires pre-creation of local directory structure. We
+        don't need to do that here so we override the method be a no-op.
+        """
+        # No setup to do and no need to pre-create local file storage.
+        pass

airbyte/caches/__init__.py

@@ -5,6 +5,7 @@
 from airbyte.caches.base import CacheBase
 from airbyte.caches.duckdb import DuckDBCache
 from airbyte.caches.factories import get_default_cache, new_local_cache
+from airbyte.caches.motherduck import MotherDuckCache
 from airbyte.caches.postgres import PostgresCache
 from airbyte.caches.snowflake import SnowflakeCache
 
@@ -15,8 +16,9 @@
     "get_default_cache",
     "new_local_cache",
     # Classes
+    "CacheBase",
     "DuckDBCache",
+    "MotherDuckCache",
     "PostgresCache",
-    "CacheBase",
     "SnowflakeCache",
 ]

airbyte/caches/motherduck.py

@@ -0,0 +1,31 @@
+"""A cache implementation for the MotherDuck service, built on DuckDB."""
+from __future__ import annotations
+
+from overrides import overrides
+from pydantic import Field
+
+from airbyte._processors.sql.motherduck import MotherDuckSqlProcessor
+from airbyte.caches.duckdb import DuckDBCache
+
+
+class MotherDuckCache(DuckDBCache):
+    """Cache that uses MotherDuck for external persistent storage."""
+
+    db_path: str = Field(default="md:")
+    database: str
+    api_key: str
+
+    _sql_processor_class = MotherDuckSqlProcessor
+
+    @overrides
+    def get_sql_alchemy_url(self) -> str:
+        """Return the SQLAlchemy URL to use."""
+        return (
+            f"duckdb:///md:{self.database}?motherduck_token={self.api_key}"
+            # f"&schema={self.schema_name}"  # TODO: Debug why this doesn't work
+        )
+
+    @overrides
+    def get_database_name(self) -> str:
+        """Return the name of the database."""
+        return self.database

airbyte/progress.py

@@ -9,19 +9,24 @@
 import time
 from contextlib import suppress
 from enum import Enum, auto
-from typing import cast
+from typing import TYPE_CHECKING, cast
 
 from rich.errors import LiveError
 from rich.live import Live as RichLive
 from rich.markdown import Markdown as RichMarkdown
 
 
+if TYPE_CHECKING:
+    from types import ModuleType
+
+
 DEFAULT_REFRESHES_PER_SECOND = 2
 IS_REPL = hasattr(sys, "ps1")  # True if we're in a Python REPL, in which case we can use Rich.
 
+ipy_display: ModuleType | None
 try:
     IS_NOTEBOOK = True
-    from IPython import display as ipy_display
+    from IPython import display as ipy_display  # type: ignore  # noqa: PGH003
 
 except ImportError:
     ipy_display = None
@@ -310,7 +315,7 @@ def update_display(self, *, force_refresh: bool = False) -> None:
 
         status_message = self._get_status_message()
 
-        if self.style == ProgressStyle.IPYTHON:
+        if self.style == ProgressStyle.IPYTHON and ipy_display is not None:
             # We're in a notebook so use the IPython display.
             ipy_display.clear_output(wait=True)
             ipy_display.display(ipy_display.Markdown(status_message))

examples/run_faker_to_motherduck.py

@@ -0,0 +1,41 @@
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+"""A sample execution script which loads data from `source-faker` to a MotherDuck-backed cache.
+
+Usage (from repo root):
+    poetry install
+    poetry run python examples/run_faker_to_motherduck.py
+"""
+
+from __future__ import annotations
+
+import airbyte as ab
+from airbyte.caches import MotherDuckCache
+
+
+MOTHERDUCK_API_KEY = ab.get_secret("MOTHERDUCK_API_KEY")
+"""This is the API key for the MotherDuck service.
+
+It can be auto-detected in env vars and/or a .env file in the root of the project.
+
+If will be prompted (and masked during input) if not found in either location.
+"""
+
+
+source = ab.get_source(
+    "source-faker",
+    config={"count": 10000, "seed": 0, "parallelism": 1, "always_updated": False},
+    install_if_missing=True,
+)
+source.check()
+source.select_all_streams()
+
+cache = MotherDuckCache(
+    database="airbyte_test",
+    schema_name="faker_data",
+    api_key=MOTHERDUCK_API_KEY,
+)
+
+result = source.read(cache=cache)
+
+for name, records in result.streams.items():
+    print(f"Stream {name}: {len(records)} records")