Skip to content

Commit

Permalink
Add pre-model hook for cleaning up remote temporary table (MotherDuck)
Browse files Browse the repository at this point in the history
  • Loading branch information
guenp committed Apr 11, 2024
1 parent 29160a6 commit 5967ee1
Show file tree
Hide file tree
Showing 2 changed files with 19 additions and 7 deletions.
21 changes: 14 additions & 7 deletions dbt/adapters/duckdb/impl.py
Original file line number Diff line number Diff line change
Expand Up @@ -241,11 +241,23 @@ def render_column_constraint(cls, constraint: ColumnLevelConstraint) -> Optional
else:
return super().render_column_constraint(constraint)

def _clean_up_temp_relation_for_incremental(self, config):
if self.is_motherduck():
if "incremental" == config.model.get_materialization():
temp_relation = self.Relation(
path=self.get_temp_relation_path(config.model), type=RelationType.Table
)
self.drop_relation(temp_relation)

def pre_model_hook(self, config: Any) -> None:
"""A hook for getting the temp schema name from the model config"""
"""A hook for getting the temp schema name from the model config.
Cleans up the remote temporary table on MotherDuck before running
an incremental model.
"""
self._temp_schema_name = config.model.config.meta.get(
TEMP_SCHEMA_NAME, self._temp_schema_name
)
self._clean_up_temp_relation_for_incremental(config)
super().pre_model_hook(config)

@available
Expand All @@ -262,12 +274,7 @@ def post_model_hook(self, config: Any, context: Any) -> None:
"""A hook for cleaning up the remote temporary table on MotherDuck if the
incremental model materialization fails to do so.
"""
if self.is_motherduck():
if "incremental" == config.model.get_materialization():
temp_relation = self.Relation(
path=self.get_temp_relation_path(config.model), type=RelationType.Table
)
self.drop_relation(temp_relation)
self._clean_up_temp_relation_for_incremental(config)
super().post_model_hook(config, context)


Expand Down
5 changes: 5 additions & 0 deletions tests/functional/plugins/test_motherduck.py
Original file line number Diff line number Diff line change
Expand Up @@ -109,6 +109,11 @@ def test_incremental(self, project):
res = project.run_sql("SELECT schema_name FROM information_schema.schemata WHERE catalog_name = 'test'", fetch="all")
assert "dbt_temp_test" in [_r for (_r,) in res]

def test_incremental_temp_table_exists(self, project):
project.run_sql('create or replace table test.dbt_temp_test.summary_of_logs_test as (select 1 from generate_series(1,10) g(x))')
run_dbt()
res = project.run_sql("SELECT count(*) FROM summary_of_logs_test", fetch="one")
assert res == (70,)

@pytest.fixture
def mock_md_plugin():
Expand Down

0 comments on commit 5967ee1

Please sign in to comment.