diff --git a/.gitignore b/.gitignore index 009c307f..75c102c1 100644 --- a/.gitignore +++ b/.gitignore @@ -79,3 +79,4 @@ target/ .DS_Store .idea/ .vscode/ +.env diff --git a/README.md b/README.md index e4fa9493..2cf2fd37 100644 --- a/README.md +++ b/README.md @@ -61,7 +61,7 @@ option that will be automatically enabled if you are connecting to a MotherDuck 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. For example, to be able to connect to S3 and read/write +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 Parquet files using an AWS access key and secret, your profile would look something like this: ``` @@ -73,10 +73,11 @@ default: extensions: - httpfs - parquet - settings: - s3_region: my-aws-region - s3_access_key_id: "{{ env_var('S3_ACCESS_KEY_ID') }}" - s3_secret_access_key: "{{ env_var('S3_SECRET_ACCESS_KEY') }}" + secrets: + - type: s3 + region: my-aws-region + key_id: "{{ env_var('S3_ACCESS_KEY_ID') }}" + secret: "{{ env_var('S3_SECRET_ACCESS_KEY') }}" target: dev ``` @@ -107,7 +108,23 @@ to load (so `s3`, `gcs`, `abfs`, etc.) and then an arbitrary set of other key-va illustrates the usage of this feature to connect to a Localstack instance running S3 from dbt-duckdb [here](https://github.com/jwills/s3-demo). #### Fetching credentials from context -Instead of specifying the credentials through the settings block, you can also use the use_credential_provider property. If you set this to `aws` (currently the only supported implementation) and you have `boto3` installed in your python environment, we will fetch your AWS credentials using the credential provider chain as described [here](https://boto3.amazonaws.com/v1/documentation/api/latest/guide/credentials.html). This means that you can use any supported mechanism from AWS to obtain credentials (e.g., web identity tokens). + +Instead of specifying the credentials through the settings block, you can also use the `CREDENTIAL_CHAIN` secret provider. This means that you can use any supported mechanism from AWS to obtain credentials (e.g., web identity tokens). You can read more about the secret providers [here](https://duckdb.org/docs/configuration/secrets_manager.html#secret-providers). To use the `CREDENTIAL_CHAIN` provider and automatically fetch credentials from AWS, specify the `provider` in the `secrets` key: + +``` +default: + outputs: + dev: + type: duckdb + path: /tmp/dbt.duckdb + extensions: + - httpfs + - parquet + secrets: + - type: s3 + provider: credential_chain + target: dev +``` #### Attaching Additional Databases diff --git a/dbt/adapters/duckdb/credentials.py b/dbt/adapters/duckdb/credentials.py index 53326cf9..56d3faf3 100644 --- a/dbt/adapters/duckdb/credentials.py +++ b/dbt/adapters/duckdb/credentials.py @@ -1,8 +1,6 @@ import os -import time from dataclasses import dataclass from dataclasses import field -from functools import lru_cache from typing import Any from typing import Dict from typing import List @@ -14,6 +12,8 @@ from dbt_common.exceptions import DbtRuntimeError from dbt.adapters.contracts.connection import Credentials +from dbt.adapters.duckdb.secrets import DEFAULT_SECRET_PREFIX +from dbt.adapters.duckdb.secrets import Secret @dataclass @@ -96,6 +96,10 @@ class DuckDBCredentials(Credentials): # (and extensions may add their own pragmas as well) settings: Optional[Dict[str, Any]] = None + # secrets for connecting to cloud services AWS S3, Azure, Cloudfare R2, + # Google Cloud and Huggingface. + secrets: Optional[List[Dict[str, Any]]] = None + # the root path to use for any external materializations that are specified # in this dbt project; defaults to "." (the current working directory) external_root: str = "." @@ -148,6 +152,10 @@ class DuckDBCredentials(Credentials): retries: Optional[Retries] = None def __post_init__(self): + self.settings = self.settings or {} + self.secrets = self.secrets or [] + self._secrets = [] + # Add MotherDuck plugin if the path is a MotherDuck database # and plugin was not specified in profile.yml if self.is_motherduck: @@ -156,6 +164,29 @@ def __post_init__(self): if "motherduck" not in [plugin.module for plugin in self.plugins]: self.plugins.append(PluginConfig(module="motherduck")) + # For backward compatibility, to be deprecated in the future + if self.use_credential_provider: + if self.use_credential_provider == "aws": + self.secrets.append({"type": "s3", "provider": "credential_chain"}) + else: + raise ValueError( + "Unsupported value for use_credential_provider: " + + self.use_credential_provider + ) + + if self.secrets: + self._secrets = [ + Secret.create( + secret_type=secret.pop("type"), + name=secret.pop("name", f"{DEFAULT_SECRET_PREFIX}{num + 1}"), + **secret, + ) + for num, secret in enumerate(self.secrets) + ] + + def secrets_sql(self) -> List[str]: + return [secret.to_sql() for secret in self._secrets] + @property def is_motherduck(self): parsed = urlparse(self.path) @@ -230,54 +261,3 @@ def _connection_keys(self): "plugins", "disable_transactions", ) - - def load_settings(self) -> Dict[str, str]: - settings = self.settings or {} - if self.use_credential_provider: - if self.use_credential_provider == "aws": - settings.update( - _load_aws_credentials(ttl=_get_ttl_hash(), profile=settings.get("s3_profile")), - ) - else: - raise ValueError( - "Unsupported value for use_credential_provider: " - + self.use_credential_provider - ) - return settings - - -def _get_ttl_hash(seconds=300): - return round(time.time() / seconds) - - -@lru_cache() -def _load_aws_credentials(ttl=None, profile="default") -> Dict[str, Any]: - """ - Load AWS credentials from the environment. - - This function is cached to prevent unnecessary calls to the AWS API. - - :param ttl: Time to live for the cache. If None, the cache will not expire. - :return: A dictionary containing the AWS credentials which can be used to configure DuckDB settings. - """ - del ttl # make mypy happy - import boto3.session - - session = boto3.session.Session(profile_name=profile) - - # use STS to verify that the credentials are valid; we will - # raise a helpful error here if they are not - sts = session.client("sts") - sts.get_caller_identity() - - # now extract/return them - aws_creds = session.get_credentials().get_frozen_credentials() - - credentials = { - "s3_access_key_id": aws_creds.access_key, - "s3_secret_access_key": aws_creds.secret_key, - "s3_session_token": aws_creds.token, - "s3_region": session.region_name, - } - # only return if value is filled - return {k: v for k, v in credentials.items() if v} diff --git a/dbt/adapters/duckdb/environments/__init__.py b/dbt/adapters/duckdb/environments/__init__.py index 3f52a9cf..cab331e6 100644 --- a/dbt/adapters/duckdb/environments/__init__.py +++ b/dbt/adapters/duckdb/environments/__init__.py @@ -201,10 +201,14 @@ def initialize_cursor( plugins: Optional[Dict[str, BasePlugin]] = None, registered_df: dict = {}, ): - for key, value in creds.load_settings().items(): - # Okay to set these as strings because DuckDB will cast them - # to the correct type - cursor.execute(f"SET {key} = '{value}'") + if creds.settings is not None: + for key, value in creds.settings.items(): + # Okay to set these as strings because DuckDB will cast them + # to the correct type + cursor.execute(f"SET {key} = '{value}'") + + for sql in creds.secrets_sql(): + cursor.execute(sql) # update cursor if something is lost in the copy # of the parent connection @@ -229,7 +233,9 @@ def initialize_plugins(cls, creds: DuckDBCredentials) -> Dict[str, BasePlugin]: for plugin_def in creds.plugins or []: config = base_config.copy() config.update(plugin_def.config or {}) - plugin = BasePlugin.create(plugin_def.module, config=config, alias=plugin_def.alias) + plugin = BasePlugin.create( + plugin_def.module, config=config, alias=plugin_def.alias, credentials=creds + ) ret[plugin.name] = plugin return ret diff --git a/dbt/adapters/duckdb/plugins/__init__.py b/dbt/adapters/duckdb/plugins/__init__.py index 0dc463f1..9e610d3a 100644 --- a/dbt/adapters/duckdb/plugins/__init__.py +++ b/dbt/adapters/duckdb/plugins/__init__.py @@ -37,6 +37,7 @@ def create( *, config: Optional[Dict[str, Any]] = None, alias: Optional[str] = None, + credentials: Optional[DuckDBCredentials] = None, ) -> "BasePlugin": """ Create a plugin from a module name and optional configuration. @@ -61,12 +62,22 @@ def create( except ImportError as e: raise ImportError(f"Unable to import module '{module}': {e}") + if config is None and credentials is not None: + config = credentials.settings + if not hasattr(mod, "Plugin"): raise ImportError(f"Module '{module}' does not have a Plugin class.") else: - return mod.Plugin(alias or name, config or {}) - - def __init__(self, name: str, plugin_config: Dict[str, Any]): + return mod.Plugin( + name=alias or name, plugin_config=config or {}, credentials=credentials + ) + + def __init__( + self, + name: str, + plugin_config: Dict[str, Any], + credentials: Optional[DuckDBCredentials] = None, + ): """ Initialize the BasePlugin instance with a name and its configuration. This method should *not* be overriden by subclasses in general; any @@ -74,9 +85,11 @@ def __init__(self, name: str, plugin_config: Dict[str, Any]): defined in the `initialize` method. :param name: A string representing the plugin name. + :param credentials: The DuckDB credentials :param plugin_config: A dictionary representing the plugin configuration. """ self.name = name + self.creds = credentials self.initialize(plugin_config) def initialize(self, plugin_config: Dict[str, Any]): diff --git a/dbt/adapters/duckdb/plugins/glue.py b/dbt/adapters/duckdb/plugins/glue.py index 52c538e2..78fb7581 100644 --- a/dbt/adapters/duckdb/plugins/glue.py +++ b/dbt/adapters/duckdb/plugins/glue.py @@ -16,6 +16,7 @@ from . import BasePlugin from ..utils import TargetConfig from dbt.adapters.base.column import Column +from dbt.adapters.duckdb.secrets import Secret class UnsupportedFormatType(Exception): @@ -263,9 +264,23 @@ def _get_table_def( return table_def -def _get_glue_client(settings: Dict[str, Any]) -> "GlueClient": - if settings: - return boto3.client( +def _get_glue_client( + settings: Dict[str, Any], secrets: Optional[List[Dict[str, Any]]] +) -> "GlueClient": + if secrets is not None: + for secret in secrets: + if isinstance(secret, Secret) and "config" == str(secret.provider).lower(): + secret_kwargs = secret.secret_kwargs or {} + client = boto3.client( + "glue", + aws_access_key_id=secret_kwargs.get("key_id"), + aws_secret_access_key=secret_kwargs.get("secret"), + aws_session_token=secret_kwargs.get("session_token"), + region_name=secret_kwargs.get("region"), + ) + break + elif settings: + client = boto3.client( "glue", aws_access_key_id=settings.get("s3_access_key_id"), aws_secret_access_key=settings.get("s3_secret_access_key"), @@ -273,7 +288,8 @@ def _get_glue_client(settings: Dict[str, Any]) -> "GlueClient": region_name=settings.get("s3_region"), ) else: - return boto3.client("glue") + client = boto3.client("glue") + return client def create_or_update_table( @@ -327,7 +343,9 @@ def create_or_update_table( class Plugin(BasePlugin): def initialize(self, config: Dict[str, Any]): - self.client = _get_glue_client(config) + if self.creds is not None: + secrets = self.creds.secrets + self.client = _get_glue_client(config, secrets) self.database = config.get("glue_database", "default") self.delimiter = config.get("delimiter", ",") diff --git a/dbt/adapters/duckdb/secrets.py b/dbt/adapters/duckdb/secrets.py new file mode 100644 index 00000000..62b1307d --- /dev/null +++ b/dbt/adapters/duckdb/secrets.py @@ -0,0 +1,56 @@ +from dataclasses import dataclass +from typing import Any +from typing import Dict +from typing import Optional + +from dbt_common.dataclass_schema import dbtClassMixin + + +DEFAULT_SECRET_PREFIX = "_dbt_secret_" + + +@dataclass +class Secret(dbtClassMixin): + type: str + persistent: Optional[bool] = False + name: Optional[str] = None + provider: Optional[str] = None + scope: Optional[str] = None + secret_kwargs: Optional[Dict[str, Any]] = None + + @classmethod + def create( + cls, + secret_type: str, + persistent: Optional[bool] = None, + name: Optional[str] = None, + provider: Optional[str] = None, + scope: Optional[str] = None, + **kwargs, + ): + # Create and return Secret + return cls( + type=secret_type, + persistent=persistent, + name=name, + provider=provider, + scope=scope, + secret_kwargs=kwargs, + ) + + def to_sql(self) -> str: + name = f" {self.name}" if self.name else "" + or_replace = " OR REPLACE" if name else "" + persistent = " PERSISTENT" if self.persistent is True else "" + tab = " " + params = self.to_dict(omit_none=True) + params.update(params.pop("secret_kwargs", {})) + params_sql = f",\n{tab}".join( + [ + f"{key} {value}" + for key, value in params.items() + if value is not None and key not in ["name", "persistent"] + ] + ) + sql = f"""CREATE{or_replace}{persistent} SECRET{name} (\n{tab}{params_sql}\n)""" + return sql diff --git a/setup.py b/setup.py index 19f76673..bbd19538 100644 --- a/setup.py +++ b/setup.py @@ -39,7 +39,7 @@ def _dbt_duckdb_version(): install_requires=[ "dbt-common>=1,<2", "dbt-adapters>=1,<2", - "duckdb>=0.7.0,!=0.10.3", + "duckdb>=0.10.0,!=0.10.3", # add dbt-core to ensure backwards compatibility of installation, this is not a functional dependency "dbt-core>=1.8.0", ], diff --git a/tests/unit/test_connections.py b/tests/unit/test_connections.py deleted file mode 100644 index cb33b031..00000000 --- a/tests/unit/test_connections.py +++ /dev/null @@ -1,94 +0,0 @@ -from unittest import mock - -from botocore.credentials import Credentials - -from dbt.adapters.duckdb.credentials import Attachment, DuckDBCredentials - - -def test_load_basic_settings(): - creds = DuckDBCredentials() - creds.settings = { - "s3_access_key_id": "abc", - "s3_secret_access_key": "xyz", - "s3_region": "us-west-2", - } - settings = creds.load_settings() - assert creds.settings == settings - - -@mock.patch("boto3.session.Session") -def test_load_aws_creds(mock_session_class): - mock_session_object = mock.Mock() - mock_client = mock.Mock() - - mock_session_object.get_credentials.return_value = Credentials( - "access_key", "secret_key", "token" - ) - mock_session_object.client.return_value = mock_client - mock_session_class.return_value = mock_session_object - mock_client.get_caller_identity.return_value = {} - - creds = DuckDBCredentials(use_credential_provider="aws") - creds.settings = {"some_other_setting": 1} - - settings = creds.load_settings() - assert settings["s3_access_key_id"] == "access_key" - assert settings["s3_secret_access_key"] == "secret_key" - assert settings["s3_session_token"] == "token" - assert settings["some_other_setting"] == 1 - - -def test_attachments(): - creds = DuckDBCredentials() - creds.attach = [ - {"path": "/tmp/f1234.db"}, - {"path": "/tmp/g1234.db", "alias": "g"}, - {"path": "/tmp/h5678.db", "read_only": 1}, - {"path": "/tmp/i9101.db", "type": "sqlite"}, - {"path": "/tmp/jklm.db", "alias": "jk", "read_only": 1, "type": "sqlite"}, - ] - - expected_sql = [ - "ATTACH '/tmp/f1234.db'", - "ATTACH '/tmp/g1234.db' AS g", - "ATTACH '/tmp/h5678.db' (READ_ONLY)", - "ATTACH '/tmp/i9101.db' (TYPE sqlite)", - "ATTACH '/tmp/jklm.db' AS jk (TYPE sqlite, READ_ONLY)", - ] - - for i, a in enumerate(creds.attach): - attachment = Attachment(**a) - assert expected_sql[i] == attachment.to_sql() - - -def test_infer_database_name_from_path(): - payload = {} - creds = DuckDBCredentials.from_dict(payload) - assert creds.database == "memory" - - payload = {"path": "local.duckdb"} - creds = DuckDBCredentials.from_dict(payload) - assert creds.database == "local" - - payload = {"path": "/tmp/f1234.db"} - creds = DuckDBCredentials.from_dict(payload) - assert creds.database == "f1234" - - payload = {"path": "md:?token=abc123"} - creds = DuckDBCredentials.from_dict(payload) - assert creds.database == "my_db" - - payload = {"path": "md:jaffle_shop?token=abc123"} - creds = DuckDBCredentials.from_dict(payload) - assert creds.database == "jaffle_shop" - - payload = {"database": "memory"} - creds = DuckDBCredentials.from_dict(payload) - assert creds.database == "memory" - - payload = { - "database": "remote", - "remote": {"host": "localhost", "port": 5433, "user": "test"}, - } - creds = DuckDBCredentials.from_dict(payload) - assert creds.database == "remote" diff --git a/tests/unit/test_credentials.py b/tests/unit/test_credentials.py new file mode 100644 index 00000000..72c5b8a3 --- /dev/null +++ b/tests/unit/test_credentials.py @@ -0,0 +1,244 @@ +import duckdb +import pytest +from unittest import mock + +from botocore.credentials import Credentials + +from dbt.adapters.duckdb.credentials import Attachment, DuckDBCredentials + + +def test_load_basic_settings(): + settings = { + "access_mode": "read_only", + "log_query_path": "/path/to/log", + } + creds = DuckDBCredentials(settings=settings) + assert creds.settings == settings + + +def test_add_secret_with_empty_name(): + creds = DuckDBCredentials( + secrets=[ + dict( + type="s3", + name="", + key_id="abc", + secret="xyz", + region="us-west-2" + ) + ] + ) + assert len(creds.secrets) == 1 + assert creds._secrets[0].type == "s3" + assert creds._secrets[0].secret_kwargs.get("key_id") == "abc" + assert creds._secrets[0].secret_kwargs.get("secret") == "xyz" + assert creds._secrets[0].secret_kwargs.get("region") == "us-west-2" + + sql = creds.secrets_sql()[0] + assert sql == \ +"""CREATE SECRET ( + type s3, + key_id abc, + secret xyz, + region us-west-2 +)""" + + +def test_add_secret_with_name(): + creds = DuckDBCredentials.from_dict( + dict(secrets=[ + dict( + type="s3", + name="my_secret", + key_id="abc", + secret="xyz", + region="us-west-2", + scope="s3://my-bucket" + ) + ]) + ) + assert len(creds._secrets) == 1 + assert creds._secrets[0].type == "s3" + assert creds._secrets[0].secret_kwargs.get("key_id") == "abc" + assert creds._secrets[0].secret_kwargs.get("secret") == "xyz" + assert creds._secrets[0].secret_kwargs.get("region") == "us-west-2" + assert creds._secrets[0].scope == "s3://my-bucket" + + sql = creds.secrets_sql()[0] + assert sql == \ +"""CREATE OR REPLACE SECRET my_secret ( + type s3, + scope s3://my-bucket, + key_id abc, + secret xyz, + region us-west-2 +)""" + + +def test_add_unsupported_secret(): + creds = DuckDBCredentials( + secrets=[ + dict( + type="scrooge_mcduck", + name="money" + ) + ] + ) + sql = creds.secrets_sql()[0] + assert sql == \ +"""CREATE OR REPLACE SECRET money ( + type scrooge_mcduck +)""" + with pytest.raises(duckdb.InvalidInputException) as e: + duckdb.sql(sql) + assert "Secret type 'scrooge_mcduck' not found" in str(e) + + +def test_add_unsupported_secret_param(): + creds = DuckDBCredentials( + secrets=[ + dict( + type="s3", + password="secret" + ) + ] + ) + sql = creds.secrets_sql()[0] + assert sql == \ +"""CREATE OR REPLACE SECRET _dbt_secret_1 ( + type s3, + password secret +)""" + with pytest.raises(duckdb.BinderException) as e: + duckdb.sql(sql) + msg = "Unknown parameter 'password' for secret type 's3' with default provider 'config'" + assert msg in str(e) + + +def test_add_azure_secret(): + creds = DuckDBCredentials( + secrets=[ + dict( + type="azure", + name="", + provider="service_principal", + tenant_id="abc", + client_id="xyz", + client_certificate_path="foo\\bar\\baz.pem", + account_name="123" + ) + ] + ) + assert len(creds.secrets) == 1 + assert creds._secrets[0].type == "azure" + assert creds._secrets[0].secret_kwargs.get("tenant_id") == "abc" + assert creds._secrets[0].secret_kwargs.get("client_id") == "xyz" + assert creds._secrets[0].secret_kwargs.get("client_certificate_path") == "foo\\bar\\baz.pem" + assert creds._secrets[0].secret_kwargs.get("account_name") == "123" + + sql = creds.secrets_sql()[0] + assert sql == \ +"""CREATE SECRET ( + type azure, + provider service_principal, + tenant_id abc, + client_id xyz, + client_certificate_path foo\\bar\\baz.pem, + account_name 123 +)""" + + +def test_add_hf_secret(): + creds = DuckDBCredentials( + secrets=[ + dict( + type="huggingface", + name="", + token="abc" + ) + ] + ) + assert len(creds.secrets) == 1 + assert creds._secrets[0].type == "huggingface" + assert creds._secrets[0].secret_kwargs.get("token") == "abc" + + sql = creds.secrets_sql()[0] + assert sql == \ +"""CREATE SECRET ( + type huggingface, + token abc +)""" + + +@mock.patch("boto3.session.Session") +def test_load_aws_creds(mock_session_class): + mock_session_object = mock.Mock() + mock_client = mock.Mock() + + mock_session_object.get_credentials.return_value = Credentials( + "access_key", "secret_key", "token" + ) + mock_session_object.client.return_value = mock_client + mock_session_class.return_value = mock_session_object + mock_client.get_caller_identity.return_value = {} + + creds = DuckDBCredentials(use_credential_provider="aws") + assert len(creds.secrets) == 1 + assert creds._secrets[0].type == "s3" + assert creds._secrets[0].provider == "credential_chain" + + +def test_attachments(): + creds = DuckDBCredentials() + creds.attach = [ + {"path": "/tmp/f1234.db"}, + {"path": "/tmp/g1234.db", "alias": "g"}, + {"path": "/tmp/h5678.db", "read_only": 1}, + {"path": "/tmp/i9101.db", "type": "sqlite"}, + {"path": "/tmp/jklm.db", "alias": "jk", "read_only": 1, "type": "sqlite"}, + ] + + expected_sql = [ + "ATTACH '/tmp/f1234.db'", + "ATTACH '/tmp/g1234.db' AS g", + "ATTACH '/tmp/h5678.db' (READ_ONLY)", + "ATTACH '/tmp/i9101.db' (TYPE sqlite)", + "ATTACH '/tmp/jklm.db' AS jk (TYPE sqlite, READ_ONLY)", + ] + + for i, a in enumerate(creds.attach): + attachment = Attachment(**a) + assert expected_sql[i] == attachment.to_sql() + + +def test_infer_database_name_from_path(): + payload = {} + creds = DuckDBCredentials.from_dict(payload) + assert creds.database == "memory" + + payload = {"path": "local.duckdb"} + creds = DuckDBCredentials.from_dict(payload) + assert creds.database == "local" + + payload = {"path": "/tmp/f1234.db"} + creds = DuckDBCredentials.from_dict(payload) + assert creds.database == "f1234" + + payload = {"path": "md:?token=abc123"} + creds = DuckDBCredentials.from_dict(payload) + assert creds.database == "my_db" + + payload = {"path": "md:jaffle_shop?token=abc123"} + creds = DuckDBCredentials.from_dict(payload) + assert creds.database == "jaffle_shop" + + payload = {"database": "memory"} + creds = DuckDBCredentials.from_dict(payload) + assert creds.database == "memory" + + payload = { + "database": "remote", + "remote": {"host": "localhost", "port": 5433, "user": "test"}, + } + creds = DuckDBCredentials.from_dict(payload) + assert creds.database == "remote" diff --git a/tests/unit/test_duckdb_adapter.py b/tests/unit/test_duckdb_adapter.py index 808bbd90..43fba6c0 100644 --- a/tests/unit/test_duckdb_adapter.py +++ b/tests/unit/test_duckdb_adapter.py @@ -63,3 +63,62 @@ def test_cancel_open_connections_main(self): key = self.adapter.connections.get_thread_identifier() self.adapter.connections.thread_connections[key] = mock_connection("main") self.assertEqual(len(list(self.adapter.cancel_open_connections())), 0) + + +class TestDuckDBAdapterWithSecrets(unittest.TestCase): + def setUp(self): + set_from_args(Namespace(STRICT_MODE=True), {}) + + profile_cfg = { + "outputs": { + "test": { + "type": "duckdb", + "path": ":memory:", + "secrets": [ + { + "type": "s3", + "key_id": "abc", + "secret": "xyz", + "region": "us-west-2" + } + ] + } + }, + "target": "test", + } + + project_cfg = { + "name": "X", + "version": "0.1", + "profile": "test", + "project-root": "/tmp/dbt/does-not-exist", + "quoting": { + "identifier": False, + "schema": True, + }, + "config-version": 2, + } + + self.config = config_from_parts_or_dicts(project_cfg, profile_cfg, cli_vars={}) + self._adapter = None + + @property + def adapter(self): + self.mock_mp_context = mock.MagicMock() + if self._adapter is None: + self._adapter = DuckDBAdapter(self.config, self.mock_mp_context) + return self._adapter + + @mock.patch("dbt.adapters.duckdb.environments.duckdb") + def test_create_secret(self, connector): + connector.__version__ = "0.1.0" # dummy placeholder for semver checks + DuckDBConnectionManager.close_all_connections() + connection = self.adapter.acquire_connection("dummy") + assert connection.handle + connection.handle._cursor._cursor.execute.assert_called_with( +"""CREATE OR REPLACE SECRET _dbt_secret_1 ( + type s3, + key_id abc, + secret xyz, + region us-west-2 +)""") diff --git a/tests/unit/test_glue.py b/tests/unit/test_glue.py index 1086785a..57658b08 100644 --- a/tests/unit/test_glue.py +++ b/tests/unit/test_glue.py @@ -309,4 +309,4 @@ def test_without_update_glue_table(self, mocker, columns): delimiter=",", ) assert len(client.mock_calls) == 1 - client.has_calls([call.get_table(DatabaseName="test", Name="test")]) + client.assert_has_calls([call.get_table(DatabaseName="test", Name="test")])