Skip to content

Commit

Permalink
fix(duckdb) only specify custom_user_agent for duckdb > 1.0.0 (#23762)
Browse files Browse the repository at this point in the history
## Summary & Motivation

- Resolves **#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`
  • Loading branch information
cmpadden authored Aug 20, 2024
1 parent f705a4a commit ae95d8f
Show file tree
Hide file tree
Showing 2 changed files with 24 additions and 8 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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,
)
Expand Down
16 changes: 12 additions & 4 deletions python_modules/libraries/dagster-duckdb/dagster_duckdb/resource.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down Expand Up @@ -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,
)
Expand Down

0 comments on commit ae95d8f

Please sign in to comment.