diff --git a/dbt/adapters/duckdb/environments/__init__.py b/dbt/adapters/duckdb/environments/__init__.py index 273b3122..55f6e5e4 100644 --- a/dbt/adapters/duckdb/environments/__init__.py +++ b/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 diff --git a/tests/functional/plugins/test_motherduck.py b/tests/functional/plugins/test_motherduck.py index dfb786c0..3983b5c3 100644 --- a/tests/functional/plugins/test_motherduck.py +++ b/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 = {})