diff --git a/duckdb_engine/__init__.py b/duckdb_engine/__init__.py index 559ff81d..c658e18e 100644 --- a/duckdb_engine/__init__.py +++ b/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) diff --git a/duckdb_engine/tests/test_basic.py b/duckdb_engine/tests/test_basic.py index cbdc77d5..e323475c 100644 --- a/duckdb_engine/tests/test_basic.py +++ b/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