Skip to content

Commit

Permalink
Add support for configuring community/nightly extensions in the profi…
Browse files Browse the repository at this point in the history
…le directly
  • Loading branch information
jwills committed Oct 3, 2024
1 parent d9da74f commit a823db8
Show file tree
Hide file tree
Showing 4 changed files with 95 additions and 7 deletions.
25 changes: 22 additions & 3 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -59,9 +59,28 @@ option that will be automatically enabled if you are connecting to a MotherDuck

#### DuckDB Extensions, Settings, and Filesystems

You can load any supported [DuckDB extensions](https://duckdb.org/docs/extensions/overview) by listing them in
the `extensions` field in your profile. You can also set any additional [DuckDB configuration options](https://duckdb.org/docs/sql/configuration)
via the `settings` field, including options that are supported in any loaded extensions. To use the [DuckDB Secrets Manager](https://duckdb.org/docs/configuration/secrets_manager.html), you can use the `secrets` field. For example, to be able to connect to S3 and read/write
You can install and load any core [DuckDB extensions](https://duckdb.org/docs/extensions/overview) by listing them in
the `extensions` field in your profile as a string. You can also set any additional [DuckDB configuration options](https://duckdb.org/docs/sql/configuration)
via the `settings` field, including options that are supported in the loaded extensions. You can also configure extensions from outside of the core
extension repository (e.g., a community extension) by configuring the extension as a `name`/`repo` pair:

```
default:
outputs:
dev:
type: duckdb
path: /tmp/dbt.duckdb
extensions:
- httpfs
- parquet
- name: h3
repo: community
- name: uc_catalog
repo: core_nightly
target: dev
```

To use the [DuckDB Secrets Manager](https://duckdb.org/docs/configuration/secrets_manager.html), you can use the `secrets` field. For example, to be able to connect to S3 and read/write
Parquet files using an AWS access key and secret, your profile would look something like this:

```
Expand Down
10 changes: 8 additions & 2 deletions dbt/adapters/duckdb/credentials.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@
from typing import Dict
from typing import List
from typing import Optional
from typing import Tuple
from typing import Union
from urllib.parse import urlparse

from dbt_common.dataclass_schema import dbtClassMixin
Expand Down Expand Up @@ -80,6 +80,12 @@ class Retries(dbtClassMixin):
retryable_exceptions: List[str] = field(default_factory=lambda: ["IOException"])


@dataclass
class Extension(dbtClassMixin):
name: str
repo: str


@dataclass
class DuckDBCredentials(Credentials):
database: str = "main"
Expand All @@ -91,7 +97,7 @@ class DuckDBCredentials(Credentials):
config_options: Optional[Dict[str, Any]] = None

# any DuckDB extensions we want to install and load (httpfs, parquet, etc.)
extensions: Optional[Tuple[str, ...]] = None
extensions: Optional[List[Union[str, Dict[str, str]]]] = None

# any additional pragmas we want to configure on our DuckDB connections;
# a list of the built-in pragmas can be found here:
Expand Down
13 changes: 11 additions & 2 deletions dbt/adapters/duckdb/environments/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
from dbt_common.exceptions import DbtRuntimeError

from ..credentials import DuckDBCredentials
from ..credentials import Extension
from ..plugins import BasePlugin
from ..utils import SourceConfig
from ..utils import TargetConfig
Expand Down Expand Up @@ -167,8 +168,16 @@ def initialize_db(
# install any extensions on the connection
if creds.extensions is not None:
for extension in creds.extensions:
conn.install_extension(extension)
conn.load_extension(extension)
if isinstance(extension, str):
conn.install_extension(extension)
conn.load_extension(extension)
elif isinstance(extension, dict):
try:
ext = Extension(**extension)
except Exception as e:
raise DbtRuntimeError(f"Failed to parse extension: {e}")
conn.execute(f"install {ext.name} from {ext.repo}")
conn.load_extension(ext.name)

# Attach any fsspec filesystems on the database
if creds.filesystems:
Expand Down
54 changes: 54 additions & 0 deletions tests/functional/adapter/test_community_extensions.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,54 @@
import pytest
from dbt.tests.util import (
check_relation_types,
check_relations_equal,
check_result_nodes_by_name,
relation_from_name,
run_dbt,
)

class BaseCommunityExtensions:

@pytest.fixture(scope="class")
def dbt_profile_target(self, dbt_profile_target):
dbt_profile_target["extensions"] = [
{"name": "quack", "repo": "community"},
]
return dbt_profile_target

@pytest.fixture(scope="class")
def models(self):
return {
"quack_model.sql": "select quack('world') as quack_world",
}

@pytest.fixture(scope="class")
def project_config_update(self):
return {
"name": "base",
}

def test_base(self, project):

# run command
results = run_dbt()
# run result length
assert len(results) == 1

# names exist in result nodes
check_result_nodes_by_name(
results,
[
"quack_model",
],
)

# check relation types
expected = {
"quack_model": "view",
}
check_relation_types(project.adapter, expected)

@pytest.mark.skip_profile("buenavista")
class TestCommunityExtensions(BaseCommunityExtensions):
pass

0 comments on commit a823db8

Please sign in to comment.