Merge pull request #335 from motherduckdb/guen/eco-57-motherduck_toke…

…n-from-dbt-config-should-override-env-var

Fix: MotherDuck token from plugin config should override env var

dbt/adapters/duckdb/environments/__init__.py

@@ -16,7 +16,6 @@
 from ..utils import TargetConfig
 from dbt.contracts.connection import AdapterResponse
 from dbt.exceptions import DbtRuntimeError
-from dbt.version import __version__
 
 
 def _ensure_event_loop():
@@ -115,12 +114,9 @@ def initialize_db(
         cls, creds: DuckDBCredentials, plugins: Optional[Dict[str, BasePlugin]] = None
     ):
         config = creds.config_options or {}
-        if creds.is_motherduck:
-            user_agent = f"dbt/{__version__}"
-            if "custom_user_agent" in config:
-                user_agent = f"{user_agent} {config['custom_user_agent']}"
-
-            config["custom_user_agent"] = user_agent
+        plugins = plugins or {}
+        for plugin in plugins.values():
+            plugin.update_connection_config(creds, config)
 
         if creds.retries:
             success, attempt, exc = False, 0, None

dbt/adapters/duckdb/plugins/__init__.py

@@ -6,6 +6,7 @@
 
 from duckdb import DuckDBPyConnection
 
+from ..credentials import DuckDBCredentials
 from ..utils import SourceConfig
 from ..utils import TargetConfig
 from dbt.dataclass_schema import dbtClassMixin
@@ -88,6 +89,17 @@ def initialize(self, plugin_config: Dict[str, Any]):
         """
         pass
 
+    def update_connection_config(self, creds: DuckDBCredentials, config: Dict[str, Any]):
+        """
+        This updates the DuckDB connection config if needed.
+        This method should be overridden by subclasses to add any additional
+        config options needed on the connection, such as a connection token or user agent
+
+        :param creds: DuckDB credentials
+        :param config: Config dictionary to be passed to duckdb.connect
+        """
+        pass
+
     def configure_connection(self, conn: DuckDBPyConnection):
         """
         Configure the DuckDB connection with any necessary extensions and/or settings.

dbt/adapters/duckdb/plugins/motherduck.py

@@ -4,6 +4,8 @@
 from duckdb import DuckDBPyConnection
 
 from . import BasePlugin
+from dbt.adapters.duckdb.credentials import DuckDBCredentials
+from dbt.version import __version__
 
 
 class Plugin(BasePlugin):
@@ -12,6 +14,30 @@ def initialize(self, config: Dict[str, Any]):
 
     def configure_connection(self, conn: DuckDBPyConnection):
         conn.load_extension("motherduck")
-        if self._token:
-            connect_stmt = f"SET motherduck_token={self._token}')"
-            conn.execute(connect_stmt)
+
+    @staticmethod
+    def token_from_config(creds: DuckDBCredentials) -> str:
+        """Load the token from the MotherDuck plugin config
+        If not specified, this returns an empty string
+
+        :param str: MotherDuck token
+        """
+        plugins = creds.plugins or []
+        for plugin in plugins:
+            if plugin.config:
+                token = plugin.config.get("token") or ""
+                return str(token)
+        return ""
+
+    def update_connection_config(self, creds: DuckDBCredentials, config: Dict[str, Any]):
+        user_agent = f"dbt/{__version__}"
+        if "custom_user_agent" in config:
+            user_agent = f"{user_agent} {config['custom_user_agent']}"
+
+        config["custom_user_agent"] = user_agent
+
+        # If a user specified the token via the plugin config,
+        # pass it to the config kwarg in duckdb.connect
+        token = self.token_from_config(creds)
+        if token != "":
+            config["motherduck_token"] = token

setup.py

@@ -40,9 +40,7 @@ def _dbt_duckdb_version():
         "dbt-core~=1.7.0",
         "duckdb>=0.7.0",
     ],
-    extras_require={
-        "glue": ["boto3", "mypy-boto3-glue"],
-    },
+    extras_require={"glue": ["boto3", "mypy-boto3-glue"], "md": ["duckdb>=0.7.0,<=0.9.2"]},
     classifiers=[
         "Development Status :: 5 - Production/Stable",
         "License :: OSI Approved :: Apache Software License",

tests/conftest.py

@@ -15,6 +15,9 @@
 # Note: fixtures with session scope need to be local
 pytest_plugins = ["dbt.tests.fixtures.project"]
 
+MOTHERDUCK_TOKEN = "MOTHERDUCK_TOKEN"
+TEST_MOTHERDUCK_TOKEN = "TEST_MOTHERDUCK_TOKEN"
+
 
 def pytest_addoption(parser):
     parser.addoption("--profile", action="store", default="memory", type=str)
@@ -60,10 +63,13 @@ def dbt_profile_target(profile_type, bv_server_process, tmp_path_factory):
         profile["path"] = str(tmp_path_factory.getbasetemp() / "tmp.db")
     elif profile_type == "md":
         # Test against MotherDuck
-        if "MOTHERDUCK_TOKEN" not in os.environ:
-            raise ValueError(
-                "Please set the MOTHERDUCK_TOKEN environment variable to run tests against MotherDuck"
-            )
+        if MOTHERDUCK_TOKEN not in os.environ and MOTHERDUCK_TOKEN.lower() not in os.environ:
+            if TEST_MOTHERDUCK_TOKEN not in os.environ:
+                raise ValueError(
+                    f"Please set the {MOTHERDUCK_TOKEN} or {TEST_MOTHERDUCK_TOKEN} \
+                        environment variable to run tests against MotherDuck"
+                )
+            profile["token"] = os.environ.get(TEST_MOTHERDUCK_TOKEN)
         profile["disable_transactions"] = True
         profile["path"] = "md:test"
     elif profile_type == "memory":

tests/functional/plugins/test_motherduck.py

@@ -1,10 +1,13 @@
 import pytest
 from unittest import mock
+from unittest.mock import Mock
 from dbt.tests.util import (
     run_dbt,
 )
 from dbt.adapters.duckdb.environments import Environment
 from dbt.adapters.duckdb.credentials import DuckDBCredentials
+from dbt.adapters.duckdb.credentials import PluginConfig
+from dbt.adapters.duckdb.plugins.motherduck import Plugin
 from dbt.version import __version__
 
 random_logs_sql = """
@@ -43,7 +46,7 @@
 class TestMDPlugin:
     @pytest.fixture(scope="class")
     def profiles_config_update(self, dbt_profile_target):
-        md_config = {}
+        md_config = {"token": dbt_profile_target.get("token")}
         plugins = [{"module": "motherduck", "config": md_config}]
         return {
             "test": {
@@ -105,16 +108,36 @@ 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_motherduck_user_agent(dbt_profile_target):
-    test_path = dbt_profile_target["path"]
-    creds = DuckDBCredentials(path=test_path)
+
+@pytest.fixture
+def mock_md_plugin():
+    return Plugin.create("motherduck")
+
+
+@pytest.fixture
+def mock_creds(dbt_profile_target):
+    plugin_config = PluginConfig(module="motherduck", config={"token": "quack"})
+    if "md:" in dbt_profile_target["path"]:
+        return DuckDBCredentials(path=dbt_profile_target["path"], plugins=[plugin_config])
+    return DuckDBCredentials(path=dbt_profile_target["path"])
+
+
+@pytest.fixture
+def mock_plugins(mock_creds, mock_md_plugin):
+    plugins = {}
+    if mock_creds.is_motherduck:
+        plugins["motherduck"] = mock_md_plugin
+    return plugins
+
+
+def test_motherduck_user_agent(dbt_profile_target, mock_plugins, mock_creds):
     with mock.patch("dbt.adapters.duckdb.environments.duckdb.connect") as mock_connect:
-        Environment.initialize_db(creds)
-        if creds.is_motherduck:
+        Environment.initialize_db(mock_creds, plugins=mock_plugins)
+        if mock_creds.is_motherduck:
             kwargs = {
                 'read_only': False,
-                'config': {'custom_user_agent': f'dbt/{__version__}'}
+                'config': {'custom_user_agent': f'dbt/{__version__}', 'motherduck_token': 'quack'}
             }
-            mock_connect.assert_called_with(test_path, **kwargs)
+            mock_connect.assert_called_with(dbt_profile_target["path"], **kwargs)
         else:
-            mock_connect.assert_called_with(test_path, read_only=False, config = {})
+            mock_connect.assert_called_with(dbt_profile_target["path"], read_only=False, config = {})

tox.ini

@@ -46,7 +46,7 @@ passenv = *
 commands = {envpython} -m pytest --profile=md --maxfail=2 {posargs} tests/functional/adapter tests/functional/plugins/test_motherduck.py
 deps =
   -rdev-requirements.txt
-  -e.
+  -e.[md]
 
 [testenv:{fsspec,py38,py39,py310,py311,py}]
 description = adapter fsspec testing