airbyte.caches.duckdb

38class DuckDBCache(DuckDBConfig, CacheBase):
39    """A DuckDB cache."""
40
41    _sql_processor_class: type[DuckDBSqlProcessor] = PrivateAttr(default=DuckDBSqlProcessor)

28class DuckDBConfig(SqlConfig):
29    """Configuration for DuckDB."""
30
31    db_path: Path | str = Field()
32    """Normally db_path is a Path object.
33
34    The database name will be inferred from the file name. For example, given a `db_path` of
35    `/path/to/my/duckdb-file`, the database name is `my_db`.
36    """
37
38    schema_name: str = Field(default="main")
39    """The name of the schema to write to. Defaults to "main"."""
40
41    @overrides
42    def get_sql_alchemy_url(self) -> SecretString:
43        """Return the SQLAlchemy URL to use."""
44        # Suppress warnings from DuckDB about reflection on indices.
45        # https://github.com/Mause/duckdb_engine/issues/905
46        warnings.filterwarnings(
47            "ignore",
48            message="duckdb-engine doesn't yet support reflection on indices",
49            category=DuckDBEngineWarning,
50        )
51        return SecretString(f"duckdb:///{self.db_path!s}")
52
53    @overrides
54    def get_database_name(self) -> str:
55        """Return the name of the database."""
56        if self.db_path == ":memory:":
57            return "memory"
58
59        # Split the path on the appropriate separator ("/" or "\")
60        split_on: Literal["/", "\\"] = "\\" if "\\" in str(self.db_path) else "/"
61
62        # Return the file name without the extension
63        return str(self.db_path).split(sep=split_on)[-1].split(".")[0]
64
65    def _is_file_based_db(self) -> bool:
66        """Return whether the database is file-based."""
67        if isinstance(self.db_path, Path):
68            return True
69
70        db_path_str = str(self.db_path)
71        return (
72            ("/" in db_path_str or "\\" in db_path_str)
73            and db_path_str != ":memory:"
74            and "md:" not in db_path_str
75            and "motherduck:" not in db_path_str
76        )
77
78    @overrides
79    def get_sql_engine(self) -> Engine:
80        """Return the SQL Alchemy engine.
81
82        This method is overridden to ensure that the database parent directory is created if it
83        doesn't exist.
84        """
85        if self._is_file_based_db():
86            Path(self.db_path).parent.mkdir(parents=True, exist_ok=True)
87
88        return super().get_sql_engine()