Merge pull request #333 from motherduckdb/guen/eco-49-add-a-custom_us…

…er_agent-to-dbt-duckdb-integration

Add custom_user_agent to connection config when connecting to MotherDuck

dbt/adapters/duckdb/environments/__init__.py

@@ -16,6 +16,7 @@
 from ..utils import TargetConfig
 from dbt.contracts.connection import AdapterResponse
 from dbt.exceptions import DbtRuntimeError
+from dbt.version import __version__
 
 
 def _ensure_event_loop():
@@ -114,6 +115,12 @@ 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
 
         if creds.retries:
             success, attempt, exc = False, 0, None

tests/functional/plugins/test_motherduck.py

@@ -1,7 +1,11 @@
 import pytest
+from unittest 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.version import __version__
 
 random_logs_sql = """
 {{ config(materialized='table', meta=dict(temp_schema_name='dbt_temp_test')) }}
@@ -100,3 +104,17 @@ 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)
+    with mock.patch("dbt.adapters.duckdb.environments.duckdb.connect") as mock_connect:
+        Environment.initialize_db(creds)
+        if creds.is_motherduck:
+            kwargs = {
+                'read_only': False,
+                'config': {'custom_user_agent': f'dbt/{__version__}'}
+            }
+            mock_connect.assert_called_with(test_path, **kwargs)
+        else:
+            mock_connect.assert_called_with(test_path, read_only=False, config = {})