Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
39 changes: 30 additions & 9 deletions sqlmesh/core/engine_adapter/duckdb.py
Original file line number Diff line number Diff line change
Expand Up @@ -173,16 +173,8 @@ def _create_table(
track_rows_processed: bool = True,
**kwargs: t.Any,
) -> None:
catalog = self.get_current_catalog()
catalog_type_tuple = self.fetchone(
exp.select("type")
.from_("duckdb_databases()")
.where(exp.column("database_name").eq(catalog))
)
catalog_type = catalog_type_tuple[0] if catalog_type_tuple else None

partitioned_by_exps = None
if catalog_type == "ducklake":
if self._get_catalog_type(self.get_current_catalog()) == "ducklake":
partitioned_by_exps = kwargs.pop("partitioned_by", None)

super()._create_table(
Expand Down Expand Up @@ -215,6 +207,35 @@ def _create_table(
)
self.execute(f"ALTER TABLE {table_name_str} SET PARTITIONED BY ({partitioned_by_str});")

def _drop_object(
self,
name: TableName | SchemaName,
exists: bool = True,
kind: str = "TABLE",
cascade: bool = False,
**drop_args: t.Any,
) -> None:
# DuckLake catalogs do not implement DROP TABLE / DROP VIEW ... CASCADE and raise
# "Cascade Drop not supported in DuckLake". Views in DuckDB are late-binding, so
# dropping the underlying table without CASCADE is safe there.
if cascade and kind.upper() in ("TABLE", "VIEW"):
catalog = exp.to_table(name).catalog or self.get_current_catalog()
if self._get_catalog_type(catalog) == "ducklake":
cascade = False

super()._drop_object(name=name, exists=exists, kind=kind, cascade=cascade, **drop_args)

def _get_catalog_type(self, catalog: t.Optional[str]) -> t.Optional[str]:
"""Returns the type of the given catalog (e.g. 'duckdb', 'ducklake') as reported by duckdb_databases()."""
if not catalog:
return None
catalog_type_tuple = self.fetchone(
exp.select("type")
.from_("duckdb_databases()")
.where(exp.column("database_name").eq(catalog))
)
return catalog_type_tuple[0] if catalog_type_tuple else None

@property
def _is_motherduck(self) -> bool:
return self._extra_config.get("is_motherduck", False)
51 changes: 51 additions & 0 deletions tests/core/engine_adapter/test_duckdb.py
Original file line number Diff line number Diff line change
Expand Up @@ -154,3 +154,54 @@ def test_ducklake_partitioning(adapter: EngineAdapter, duck_conn, tmp_path):
f"SELECT * FROM __ducklake_metadata_{catalog}.main.ducklake_partition_info"
).fetchdf()
assert partition_info.shape[0] == 1


def test_drop_table_ducklake_no_cascade(adapter: EngineAdapter, duck_conn, tmp_path):
# DuckLake does not implement DROP TABLE/VIEW ... CASCADE, so the adapter must
# omit CASCADE for objects in a DuckLake catalog while keeping it for native catalogs.
catalog = "a_ducklake_db"

duck_conn.install_extension("ducklake")
duck_conn.load_extension("ducklake")
duck_conn.execute(
f"ATTACH 'ducklake:{tmp_path}/{catalog}.ducklake' AS {catalog} (DATA_PATH '{tmp_path}');"
)

duck_conn.execute(f"CREATE SCHEMA {catalog}.phys")
duck_conn.execute(f"CREATE SCHEMA {catalog}.virt")
duck_conn.execute(f"CREATE TABLE {catalog}.phys.t (i INTEGER)")
duck_conn.execute(f"CREATE VIEW {catalog}.virt.v AS SELECT * FROM {catalog}.phys.t")

# native catalog, cascade is passed through
duck_conn.execute("CREATE TABLE memory.main.native_t (i INTEGER)")
duck_conn.execute("CREATE VIEW memory.main.native_v AS SELECT * FROM memory.main.native_t")

adapter.drop_table(f"{catalog}.phys.t", cascade=True)
adapter.drop_view(f"{catalog}.virt.v", cascade=True)
adapter.drop_table("memory.main.native_t", cascade=True)
adapter.drop_view("memory.main.native_v", cascade=True)

assert not adapter.table_exists(f"{catalog}.phys.t")
assert not adapter.table_exists(f"{catalog}.virt.v")
assert not adapter.table_exists("memory.main.native_t")
assert not adapter.table_exists("memory.main.native_v")


def test_drop_object_cascade_by_catalog_type(make_mocked_engine_adapter: t.Callable):
adapter = make_mocked_engine_adapter(DuckDBEngineAdapter)
adapter.fetchone = lambda *_args, **_kwargs: ("ducklake",) # type: ignore

adapter.drop_table("lake.phys.t", cascade=True)
adapter.drop_view("lake.virt.v", cascade=True)
# schema cascade is supported by DuckLake and must be preserved
adapter.drop_schema("lake.virt", cascade=True)

adapter.fetchone = lambda *_args, **_kwargs: ("duckdb",) # type: ignore
adapter.drop_table("native.phys.t", cascade=True)

assert to_sql_calls(adapter) == [
'DROP TABLE IF EXISTS "lake"."phys"."t"',
'DROP VIEW IF EXISTS "lake"."virt"."v"',
'DROP SCHEMA IF EXISTS "lake"."virt" CASCADE',
'DROP TABLE IF EXISTS "native"."phys"."t" CASCADE',
]