Merge pull request #877 from elefeint/add_user_agent

chore: Add duckdb_engine/sqlalchemy version info to DuckDB user_agent

duckdb_engine/__init__.py

@@ -42,6 +42,7 @@
 sqlalchemy_version = sqlalchemy.__version__
 duckdb_version: str = duckdb.__version__  # type: ignore[attr-defined]
 supports_attach: bool = duckdb_version >= "0.7.0"
+supports_user_agent: bool = duckdb_version >= "0.9.2"
 
 if TYPE_CHECKING:
     from sqlalchemy.base import Connection
@@ -253,6 +254,11 @@ def connect(self, *cargs: Any, **cparams: Any) -> "Connection":
         config.update(cparams.pop("url_config", {}))
 
         ext = {k: config.pop(k) for k in list(config) if k not in core_keys}
+        if supports_user_agent:
+            user_agent = f"duckdb_engine/{__version__}(sqlalchemy/{sqlalchemy_version})"
+            if "custom_user_agent" in config:
+                user_agent = f"{user_agent} {config['custom_user_agent']}"
+            config["custom_user_agent"] = user_agent
 
         conn = duckdb.connect(*cargs, **cparams)
 

duckdb_engine/tests/test_basic.py

@@ -1,5 +1,6 @@
 import logging
 import os
+import re
 import zlib
 from datetime import datetime, timedelta
 from pathlib import Path
@@ -35,7 +36,7 @@
 from sqlalchemy.ext.declarative import declarative_base
 from sqlalchemy.orm import Session, relationship, sessionmaker
 
-from .. import Dialect, supports_attach
+from .. import Dialect, supports_attach, supports_user_agent
 from .._supports import has_comment_support
 
 try:
@@ -508,6 +509,38 @@ def test_url_config_and_dict_config() -> None:
         assert memory_limit in ("500.0MB", "476.8 MiB")
 
 
+@mark.skipif(
+    supports_user_agent is False,
+    reason="custom_user_agent is not supported for DuckDB version < 0.9.2",
+)
+def test_user_agent() -> None:
+    eng = create_engine("duckdb:///:memory:")
+
+    with eng.connect() as conn:
+        res = conn.execute(text("PRAGMA USER_AGENT"))
+        row = res.first()
+        assert row is not None
+        assert re.match(r"duckdb/.*(.*) python duckdb_engine/.*(sqlalchemy/.*)", row[0])
+
+
+@mark.skipif(
+    supports_user_agent is False,
+    reason="custom_user_agent is not supported for DuckDB version < 0.9.2",
+)
+def test_user_agent_with_custom_user_agent() -> None:
+    eng = create_engine(
+        "duckdb:///:memory:", connect_args={"config": {"custom_user_agent": "custom"}}
+    )
+
+    with eng.connect() as conn:
+        res = conn.execute(text("PRAGMA USER_AGENT"))
+        row = res.first()
+        assert row is not None
+        assert re.match(
+            r"duckdb/.*(.*) python duckdb_engine/.*(sqlalchemy/.*) custom", row[0]
+        )
+
+
 def test_do_ping(tmp_path: Path, caplog: LogCaptureFixture) -> None:
     engine = create_engine(
         "duckdb:///" + str(tmp_path / "db"), pool_pre_ping=True, pool_size=1