Add workaround for temporary tables in remote database when running incremental model

dbt/adapters/duckdb/credentials.py

@@ -146,6 +146,15 @@ class DuckDBCredentials(Credentials):
     # by networking issues)
     retries: Optional[Retries] = None
 
+    @property
+    def is_motherduck(self):
+        parsed = urlparse(self.path)
+        return self._is_motherduck(parsed.scheme)
+
+    @staticmethod
+    def _is_motherduck(scheme: str) -> bool:
+        return scheme in {"md", "motherduck"}
+
     @classmethod
     def __pre_deserialize__(cls, data: Dict[Any, Any]) -> Dict[Any, Any]:
         data = super().__pre_deserialize__(data)
@@ -159,7 +168,7 @@ def __pre_deserialize__(cls, data: Dict[Any, Any]) -> Dict[Any, Any]:
             path_db = os.path.splitext(base_file)[0]
             # For MotherDuck, turn on disable_transactions unless
             # it's explicitly set already by the user
-            if parsed.scheme in {"md", "motherduck"}:
+            if cls._is_motherduck(parsed.scheme):
                 if "disable_transactions" not in data:
                     data["disable_transactions"] = True
                 if path_db == "":

dbt/adapters/duckdb/environments/local.py

@@ -48,10 +48,7 @@ def __init__(self, credentials: credentials.DuckDBCredentials):
         self.handle_count = 0
         self.lock = threading.RLock()
         self._keep_open = (
-            credentials.keep_open
-            or credentials.path == ":memory:"
-            or credentials.path.startswith("md:")
-            or credentials.path.startswith("motherduck:")
+            credentials.keep_open or credentials.path == ":memory:" or credentials.is_motherduck
         )
         self._REGISTERED_DF: dict = {}
 

dbt/adapters/duckdb/impl.py

@@ -1,4 +1,5 @@
 import os
+from typing import Any
 from typing import List
 from typing import Optional
 from typing import Sequence
@@ -19,9 +20,14 @@
 from dbt.contracts.connection import AdapterResponse
 from dbt.contracts.graph.nodes import ColumnLevelConstraint
 from dbt.contracts.graph.nodes import ConstraintType
+from dbt.contracts.relation import Path
+from dbt.contracts.relation import RelationType
 from dbt.exceptions import DbtInternalError
 from dbt.exceptions import DbtRuntimeError
 
+TEMP_SCHEMA_NAME = "temp_schema_name"
+DEFAULT_TEMP_SCHEMA_NAME = "dbt_temp"
+
 
 class DuckDBAdapter(SQLAdapter):
     ConnectionManager = DuckDBConnectionManager
@@ -36,6 +42,9 @@ class DuckDBAdapter(SQLAdapter):
         ConstraintType.foreign_key: ConstraintSupport.ENFORCED,
     }
 
+    # can be overridden via the model config metadata
+    _temp_schema_name = DEFAULT_TEMP_SCHEMA_NAME
+
     @classmethod
     def date_function(cls) -> str:
         return "now()"
@@ -47,6 +56,10 @@ def is_cancelable(cls) -> bool:
     def debug_query(self):
         self.execute("select 1 as id")
 
+    @available
+    def is_motherduck(self):
+        return self.config.credentials.is_motherduck
+
     @available
     def convert_datetimes_to_strs(self, table: agate.Table) -> agate.Table:
         for column in table.columns:
@@ -217,6 +230,35 @@ def render_column_constraint(cls, constraint: ColumnLevelConstraint) -> Optional
         else:
             return super().render_column_constraint(constraint)
 
+    def pre_model_hook(self, config: Any) -> None:
+        """A hook for getting the temp schema name from the model config"""
+        self._temp_schema_name = config.model.config.meta.get(
+            TEMP_SCHEMA_NAME, self._temp_schema_name
+        )
+        super().pre_model_hook(config)
+
+    @available
+    def get_temp_relation_path(self, model: Any):
+        """This is a workaround to enable incremental models on MotherDuck because it
+        currently doesn't support remote temporary tables. Instead we use a regular
+        table that is dropped at the end of the incremental macro or post-model hook.
+        """
+        return Path(
+            schema=self._temp_schema_name, database=model.database, identifier=model.identifier
+        )
+
+    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)
+        super().post_model_hook(config, context)
+
 
 # Change `table_a/b` to `table_aaaaa/bbbbb` to avoid duckdb binding issues when relation_a/b
 # is called "table_a" or "table_b" in some of the dbt tests

dbt/adapters/duckdb/plugins/motherduck.py

@@ -12,7 +12,6 @@ def initialize(self, config: Dict[str, Any]):
 
     def configure_connection(self, conn: DuckDBPyConnection):
         conn.load_extension("motherduck")
-        connect_stmt = "PRAGMA md_connect"
         if self._token:
-            connect_stmt = f"PRAGMA md_connect('token={self._token}')"
-        conn.execute(connect_stmt)
+            connect_stmt = f"SET motherduck_token={self._token}')"
+            conn.execute(connect_stmt)

dbt/include/duckdb/macros/materializations/incremental.sql

@@ -1,6 +1,8 @@
 {% materialization incremental, adapter="duckdb", supported_languages=['sql', 'python'] -%}
 
   {%- set language = model['language'] -%}
+  -- only create temp tables if using local duckdb, as it is not currently supported for remote databases
+  {%- set temporary = not adapter.is_motherduck() -%}
 
   -- relations
   {%- set existing_relation = load_cached_relation(this) -%}
@@ -39,13 +41,20 @@
     {% set build_sql = create_table_as(False, intermediate_relation, compiled_code, language) %}
     {% set need_swap = true %}
   {% else %}
+    {% if not temporary %}
+      -- if not using a temporary table we will update the temp relation to use a different temp schema ("dbt_temp" by default)
+      {% set temp_relation = temp_relation.incorporate(path=adapter.get_temp_relation_path(this)) %}
+      {% do run_query(create_schema(temp_relation)) %}
+      -- then drop the temp relation after we insert the incremental data into the target relation
+      {% set need_drop_temp = True %}
+    {% endif %}
     {% if language == 'python' %}
       {% set build_python = create_table_as(False, temp_relation, compiled_code, language) %}
       {% call statement("pre", language=language) %}
         {{- build_python }}
       {% endcall %}
     {% else %} {# SQL #}
-      {% do run_query(create_table_as(True, temp_relation, compiled_code, language)) %}
+      {% do run_query(create_table_as(temporary, temp_relation, compiled_code, language)) %}
     {% endif %}
     {% do adapter.expand_target_column_types(
              from_relation=temp_relation,
@@ -76,6 +85,10 @@
       {% do to_drop.append(backup_relation) %}
   {% endif %}
 
+  {% if need_drop_temp %}
+      {% do to_drop.append(temp_relation) %}
+  {% endif %}
+
   {% set should_revoke = should_revoke(existing_relation, full_refresh_mode) %}
   {% do apply_grants(target_relation, grant_config, should_revoke=should_revoke) %}
 

tests/functional/plugins/test_motherduck.py

@@ -0,0 +1,102 @@
+import pytest
+from dbt.tests.util import (
+    run_dbt,
+)
+
+random_logs_sql = """
+{{ config(materialized='table', meta=dict(temp_schema_name='dbt_temp_test')) }}
+
+select 
+  uuid()::varchar as log_id,
+ '2023-10-01'::timestamp + interval 1 minute * (random() * 20000)::int as dt ,
+(random() * 4)::int64 as user_id
+from generate_series(1, 10000) g(x)
+"""
+
+summary_of_logs_sql = """
+{{
+    config(
+        materialized='incremental',
+        meta=dict(temp_schema_name='dbt_temp_test'),
+    )
+}}
+
+select dt::date as dt, user_id, count(1) as c
+from {{ ref('random_logs_test') }}
+
+
+{% if is_incremental() %}
+
+  -- this filter will only be applied on an incremental run
+  -- (uses > to include records whose timestamp occurred since the last run of this model)
+  where dt > '2023-10-08'::timestamp
+
+{% endif %}
+group by all
+"""
+
+@pytest.mark.skip_profile("buenavista", "file", "memory")
+class TestMDPlugin:
+    @pytest.fixture(scope="class")
+    def profiles_config_update(self, dbt_profile_target):
+        md_config = {}
+        plugins = [{"module": "motherduck", "config": md_config}]
+        return {
+            "test": {
+                "outputs": {
+                    "dev": {
+                        "type": "duckdb",
+                        "path": dbt_profile_target.get("path", ":memory:"),
+                        "plugins": plugins,
+                    }
+                },
+                "target": "dev",
+            }
+        }
+
+    @pytest.fixture(scope="class")
+    def database_name(self, dbt_profile_target):
+        return dbt_profile_target["path"].replace("md:", "")
+
+    @pytest.fixture(scope="class")
+    def md_sql(self, database_name):
+        # Reads from a MD database in my test account in the cloud
+        return f"""
+            select * FROM {database_name}.main.plugin_table
+        """
+
+    @pytest.fixture(scope="class")
+    def models(self, md_sql):
+        return {
+            "md_table.sql": md_sql,
+            "random_logs_test.sql": random_logs_sql,
+            "summary_of_logs_test.sql": summary_of_logs_sql,
+        }
+
+    @pytest.fixture(autouse=True)
+    def run_dbt_scope(self, project, database_name):
+        project.run_sql(f"CREATE DATABASE IF NOT EXISTS {database_name}")
+        project.run_sql("CREATE OR REPLACE TABLE plugin_table (i integer, j string)")
+        project.run_sql("INSERT INTO plugin_table (i, j) VALUES (1, 'foo')")
+        yield
+        project.run_sql("DROP VIEW md_table")
+        project.run_sql("DROP TABLE random_logs_test")
+        project.run_sql("DROP TABLE summary_of_logs_test")
+        project.run_sql("DROP TABLE plugin_table")
+
+    def test_motherduck(self, project):
+        run_dbt()
+        res = project.run_sql("SELECT * FROM md_table", fetch="one")
+        assert res == (1, "foo")
+
+    def test_incremental(self, project):
+        run_dbt()
+        res = project.run_sql("SELECT count(*) FROM summary_of_logs_test", fetch="one")
+        assert res == (70,)
+
+        run_dbt()
+        res = project.run_sql("SELECT count(*) FROM summary_of_logs_test", fetch="one")
+        assert res == (105,)
+
+        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]

tests/functional/plugins/test_plugins.py

@@ -40,11 +40,6 @@
     select foo() as foo
 """
 
-# Reads from a MD database in my test account in the cloud
-md_sql = """
-    select * FROM plugin_test.main.plugin_table
-"""
-
 
 @pytest.mark.skip_profile("buenavista", "md")
 class TestPlugins:
@@ -133,34 +128,3 @@ def test_plugins(self, project):
 
         res = project.run_sql("SELECT foo FROM foo", fetch="one")
         assert res[0] == 1729
-
-
-@pytest.mark.skip_profile("buenavista", "file", "memory")
-class TestMDPlugin:
-    @pytest.fixture(scope="class")
-    def profiles_config_update(self, dbt_profile_target):
-        md_config = {}
-        plugins = [{"module": "motherduck", "config": md_config}]
-        return {
-            "test": {
-                "outputs": {
-                    "dev": {
-                        "type": "duckdb",
-                        "path": dbt_profile_target.get("path", ":memory:"),
-                        "plugins": plugins,
-                    }
-                },
-                "target": "dev",
-            }
-        }
-
-    @pytest.fixture(scope="class")
-    def models(self):
-        return {
-            "md_table.sql": md_sql,
-        }
-
-    def test_plugins(self, project):
-        run_dbt()
-        res = project.run_sql("SELECT * FROM md_table", fetch="one")
-        assert res == (1, "foo")

tox.ini

@@ -43,7 +43,7 @@ deps =
 description = adapter function testing using MotherDuck
 skip_install = True
 passenv = *
-commands = {envpython} -m pytest --profile=md --maxfail=2 {posargs} tests/functional/adapter
+commands = {envpython} -m pytest --profile=md --maxfail=2 {posargs} tests/functional/adapter tests/functional/plugins/test_motherduck.py
 deps =
   -rdev-requirements.txt
   -e.