From ae95d8f84305188085054c5a66a1fa2e920116da Mon Sep 17 00:00:00 2001
From: colton <cmpadden@users.noreply.github.com>
Date: Tue, 20 Aug 2024 16:19:34 -0400
Subject: [PATCH] fix(duckdb) only specify custom_user_agent for duckdb > 1.0.0
 (#23762)

## Summary & Motivation

- Resolves **https://github.com/dagster-io/dagster/issues/23761**

Only provide `custom_user_agent` when `duckdb.__version__` is greater
than. `1.0.0`, as that is when this feature was included.

## How I Tested These Changes

- `tox`
---
 .../dagster-duckdb/dagster_duckdb/io_manager.py  | 16 ++++++++++++----
 .../dagster-duckdb/dagster_duckdb/resource.py    | 16 ++++++++++++----
 2 files changed, 24 insertions(+), 8 deletions(-)

diff --git a/python_modules/libraries/dagster-duckdb/dagster_duckdb/io_manager.py b/python_modules/libraries/dagster-duckdb/dagster_duckdb/io_manager.py
index 3e4292c7f1d66..4dd06e1297d97 100644
--- a/python_modules/libraries/dagster-duckdb/dagster_duckdb/io_manager.py
+++ b/python_modules/libraries/dagster-duckdb/dagster_duckdb/io_manager.py
@@ -15,6 +15,7 @@
 )
 from dagster._core.storage.io_manager import dagster_maintained_io_manager
 from dagster._utils.backoff import backoff
+from packaging.version import Version
 from pydantic import Field
 
 DUCKDB_DATETIME_FORMAT = "%Y-%m-%d %H:%M:%S"
@@ -279,16 +280,23 @@ def get_select_statement(table_slice: TableSlice) -> str:
     @staticmethod
     @contextmanager
     def connect(context, _):
+        config = context.resource_config["connection_config"]
+
+        # support for `custom_user_agent` was added in v1.0.0
+        # https://github.com/duckdb/duckdb/commit/0c66b6007b736ed2197bca54d20c9ad9a5eeef46
+        if Version(duckdb.__version__) >= Version("1.0.0"):
+            config = {
+                "custom_user_agent": "dagster",
+                **config,
+            }
+
         conn = backoff(
             fn=duckdb.connect,
             retry_on=(RuntimeError, duckdb.IOException),
             kwargs={
                 "database": context.resource_config["database"],
                 "read_only": False,
-                "config": {
-                    "custom_user_agent": "dagster",
-                    **context.resource_config["connection_config"],
-                },
+                "config": config,
             },
             max_retries=10,
         )
diff --git a/python_modules/libraries/dagster-duckdb/dagster_duckdb/resource.py b/python_modules/libraries/dagster-duckdb/dagster_duckdb/resource.py
index 0e18a1e427009..eacde701002dd 100644
--- a/python_modules/libraries/dagster-duckdb/dagster_duckdb/resource.py
+++ b/python_modules/libraries/dagster-duckdb/dagster_duckdb/resource.py
@@ -4,6 +4,7 @@
 import duckdb
 from dagster import ConfigurableResource
 from dagster._utils.backoff import backoff
+from packaging.version import Version
 from pydantic import Field
 
 
@@ -48,16 +49,23 @@ def _is_dagster_maintained(cls) -> bool:
 
     @contextmanager
     def get_connection(self):
+        config = self.connection_config
+
+        # support for `custom_user_agent` was added in v1.0.0
+        # https://github.com/duckdb/duckdb/commit/0c66b6007b736ed2197bca54d20c9ad9a5eeef46
+        if Version(duckdb.__version__) >= Version("1.0.0"):
+            config = {
+                "custom_user_agent": "dagster",
+                **config,
+            }
+
         conn = backoff(
             fn=duckdb.connect,
             retry_on=(RuntimeError, duckdb.IOException),
             kwargs={
                 "database": self.database,
                 "read_only": False,
-                "config": {
-                    "custom_user_agent": "dagster",
-                    **self.connection_config,
-                },
+                "config": config,
             },
             max_retries=10,
         )