diff --git a/docs/integrations/engines/duckdb.md b/docs/integrations/engines/duckdb.md index aca58615b3..35e90904b1 100644 --- a/docs/integrations/engines/duckdb.md +++ b/docs/integrations/engines/duckdb.md @@ -15,8 +15,8 @@ | `type` | Engine type name - must be `duckdb` | string | Y | | `database` | The optional database name. If not specified, the in-memory database is used. Cannot be defined if using `catalogs`. | string | N | | `catalogs` | Mapping to define multiple catalogs. Can [attach DuckDB catalogs](#duckdb-catalogs-example) or [catalogs for other connections](#other-connection-catalogs-example). First entry is the default catalog. Cannot be defined if using `database`. | dict | N | -| `extensions` | Extension to load into duckdb. Only autoloadable extensions are supported. | list | N | -| `connector_config` | Configuration to pass into the duckdb connector. | dict | N | +| `extensions` | Extensions to install and load into DuckDB. Each entry is either an extension name or a dict with `name`, an optional `repository` (alias such as `community`, a local directory or a URL) and an optional `force_install` flag. See [Extensions](#extensions). | list | N | +| `connector_config` | DuckDB settings applied with `SET` to every connection. Settings that control how extensions are installed (e.g. `custom_extension_repository`, `extension_directory`) are applied before any extension is installed. See [Extensions](#extensions). | dict | N | | `secrets` | Configuration for authenticating external sources (e.g., S3) using DuckDB secrets. Can be a list of secret configurations or a dictionary with custom secret names. | list/dict | N | | `filesystems` | Configuration for registering `fsspec` filesystems to the DuckDB connection. | dict | N | @@ -201,6 +201,111 @@ Example: mounting a SQLite database with the name `sqlite` that has a table `exa If a connector, like Postgres, requires sensitive information in the path, it might support defining environment variables instead. [See DuckDB Documentation for more information](https://duckdb.org/docs/extensions/postgres#configuring-via-environment-variables). +#### Extensions + +The `extensions` option lists the [DuckDB extensions](https://duckdb.org/docs/stable/extensions/overview) that SQLMesh installs and loads on every connection. Each entry is either the extension name or a dictionary with the following keys: + +| Key | Description | Required | +|-----------------|----------------------------------------------------------------------------------------------------------------------------------------------------|:--------:| +| `name` | The extension name, e.g. `httpfs`. | Y | +| `repository` | Where to install the extension from. Either a repository alias (`core`, `core_nightly`, `community`), a local directory or a URL (`https://`, `s3://`). | N | +| `force_install` | Reinstall the extension even if it is already installed (`FORCE INSTALL`). | N | + +=== "YAML" + + ```yaml linenums="1" + gateways: + my_gateway: + connection: + type: duckdb + extensions: + - httpfs + - name: h3 + repository: community + - name: spatial + repository: /opt/duckdb/extensions + force_install: true + ``` + +=== "Python" + + ```python linenums="1" + from sqlmesh.core.config import ( + Config, + ModelDefaultsConfig, + GatewayConfig, + DuckDBConnectionConfig + ) + + config = Config( + model_defaults=ModelDefaultsConfig(dialect=), + gateways={ + "my_gateway": GatewayConfig( + connection=DuckDBConnectionConfig( + extensions=[ + "httpfs", + {"name": "h3", "repository": "community"}, + {"name": "spatial", "repository": "/opt/duckdb/extensions", "force_install": True}, + ] + ) + ), + } + ) + ``` + +##### Installing extensions from a custom or offline repository + +In restricted networks the default DuckDB extension repository is often not reachable, so extensions must be installed from an internal mirror (e.g. Artifactory) or from a directory that was populated ahead of time. + +DuckDB exposes this through settings such as [`custom_extension_repository`](https://duckdb.org/docs/stable/extensions/installing_extensions#custom-repository), `extension_directory` and `autoinstall_known_extensions`. When these are supplied via `connector_config`, SQLMesh applies them **before** any `INSTALL` / `LOAD` statement runs, so DuckDB never contacts the default repository. All other `connector_config` settings are applied after the extensions have been loaded, since some of them (e.g. `s3_region`) are only defined once the corresponding extension is available. + +=== "YAML" + + ```yaml linenums="1" + gateways: + my_gateway: + connection: + type: duckdb + connector_config: + custom_extension_repository: https://artifactory.example.com/duckdb-extensions + # or, for a directory of pre-downloaded extensions: + # extension_directory: /opt/duckdb/extensions + autoinstall_known_extensions: false + extensions: + - httpfs + ``` + +=== "Python" + + ```python linenums="1" + from sqlmesh.core.config import ( + Config, + ModelDefaultsConfig, + GatewayConfig, + DuckDBConnectionConfig + ) + + config = Config( + model_defaults=ModelDefaultsConfig(dialect=), + gateways={ + "my_gateway": GatewayConfig( + connection=DuckDBConnectionConfig( + connector_config={ + "custom_extension_repository": "https://artifactory.example.com/duckdb-extensions", + "autoinstall_known_extensions": False, + }, + extensions=["httpfs"], + ) + ), + } + ) + ``` + +Alternatively, the repository can be set per extension with the `repository` key, which maps directly to `INSTALL FROM `. + +!!! note "Unsigned extensions" + `allow_unsigned_extensions` can only be set when the DuckDB database is opened and therefore cannot be configured through `connector_config`. + #### Cloud service authentication DuckDB can read data directly from cloud services via extensions (e.g., [httpfs](https://duckdb.org/docs/extensions/httpfs/s3api), [azure](https://duckdb.org/docs/extensions/azure)). diff --git a/sqlmesh/core/config/connection.py b/sqlmesh/core/config/connection.py index 29cb1ef467..8741fae8cb 100644 --- a/sqlmesh/core/config/connection.py +++ b/sqlmesh/core/config/connection.py @@ -66,6 +66,27 @@ } MOTHERDUCK_TOKEN_REGEX = re.compile(r"(\?|\&)(motherduck_token=)(\S*)") PASSWORD_REGEX = re.compile(r"(password=)(\S+)") +# DuckDB settings that control how extensions are located, installed and loaded. +# These must be applied before any INSTALL / LOAD statement is executed, otherwise DuckDB +# tries to reach the default extension repository before e.g. a custom repository, a local +# extension directory or a proxy has been configured. +DUCKDB_EXTENSION_SETTINGS = frozenset( + { + "custom_extension_repository", + "autoinstall_extension_repository", + "extension_directory", + "autoinstall_known_extensions", + "autoload_known_extensions", + "allow_community_extensions", + "allow_extensions_metadata_mismatch", + "http_proxy", + "http_proxy_username", + "http_proxy_password", + } +) +# Repository names that DuckDB understands as aliases (e.g. `core`, `core_nightly`, `community`) +# are passed to INSTALL ... FROM unquoted. Anything else (local path, URL) is a string literal. +DUCKDB_EXTENSION_REPOSITORY_ALIAS_REGEX = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") SUPPORTS_MSSQL_PYTHON_DRIVER = (version_info.major, version_info.minor) >= (3, 10) @@ -365,14 +386,53 @@ def _cursor_init(self) -> t.Optional[t.Callable[[t.Any], None]]: import duckdb from duckdb import BinderException + def apply_settings(cursor: duckdb.DuckDBPyConnection, settings: t.Dict[str, t.Any]) -> None: + if not settings: + return + + option_names = list(settings) + in_part = ",".join("?" for _ in range(len(option_names))) + + cursor.execute( + f"SELECT name, value FROM duckdb_settings() WHERE name IN ({in_part})", + option_names, + ) + + existing_values = {field: setting for field, setting in cursor.fetchall()} + + # only set connector_config items if the values differ from what is already set + # trying to set options like 'temp_directory' even to the same value can throw errors like: + # Not implemented Error: Cannot switch temporary directory after the current one has been used + for field, setting in settings.items(): + if existing_values.get(field) != setting: + try: + cursor.execute(f"SET {field} = '{setting}'") + except Exception as e: + raise ConfigError( + f"Failed to set connector config {field} to {setting}: {e}" + ) + def init(cursor: duckdb.DuckDBPyConnection) -> None: + # Settings that control where extensions come from have to be in place before the + # first INSTALL / LOAD, otherwise DuckDB reaches out to the default repository first. + apply_settings( + cursor, + { + field: setting + for field, setting in self.connector_config.items() + if field in DUCKDB_EXTENSION_SETTINGS + }, + ) + for extension in self.extensions: extension = extension if isinstance(extension, dict) else {"name": extension} install_command = f"INSTALL {extension['name']}" - if extension.get("repository"): - install_command = f"{install_command} FROM {extension['repository']}" + if repository := extension.get("repository"): + if not DUCKDB_EXTENSION_REPOSITORY_ALIAS_REGEX.match(repository): + repository = exp.Literal.string(repository).sql(dialect="duckdb") + install_command = f"{install_command} FROM {repository}" if extension.get("force_install"): install_command = f"FORCE {install_command}" @@ -383,28 +443,16 @@ def init(cursor: duckdb.DuckDBPyConnection) -> None: except Exception as e: raise ConfigError(f"Failed to load extension {extension['name']}: {e}") - if self.connector_config: - option_names = list(self.connector_config) - in_part = ",".join("?" for _ in range(len(option_names))) - - cursor.execute( - f"SELECT name, value FROM duckdb_settings() WHERE name IN ({in_part})", - option_names, - ) - - existing_values = {field: setting for field, setting in cursor.fetchall()} - - # only set connector_config items if the values differ from what is already set - # trying to set options like 'temp_directory' even to the same value can throw errors like: - # Not implemented Error: Cannot switch temporary directory after the current one has been used - for field, setting in self.connector_config.items(): - if existing_values.get(field) != setting: - try: - cursor.execute(f"SET {field} = '{setting}'") - except Exception as e: - raise ConfigError( - f"Failed to set connector config {field} to {setting}: {e}" - ) + # The remaining settings may belong to extensions (e.g. `s3_region` from httpfs), + # so they are applied after the extensions have been loaded. + apply_settings( + cursor, + { + field: setting + for field, setting in self.connector_config.items() + if field not in DUCKDB_EXTENSION_SETTINGS + }, + ) if self.secrets: duckdb_version = duckdb.__version__ diff --git a/tests/core/test_connection_config.py b/tests/core/test_connection_config.py index c506d401d9..e53dc16872 100644 --- a/tests/core/test_connection_config.py +++ b/tests/core/test_connection_config.py @@ -935,6 +935,136 @@ def test_duckdb_config_json_strings(make_config): assert config.catalogs.get("test2").path == "test2.duckdb" +@patch("duckdb.connect") +def test_duckdb_extension_settings_applied_before_install(mock_connect, make_config): + """Settings that control where extensions are fetched from must be applied before INSTALL/LOAD, + otherwise DuckDB will try to reach the default repository before e.g. `custom_extension_repository` + or `extension_directory` has been configured. All other connector_config settings keep being applied + after the extensions have been loaded.""" + mock_cursor = MagicMock() + # No settings are currently set, so every connector_config entry should be SET + mock_cursor.fetchall.return_value = [] + mock_connection = MagicMock() + mock_connection.cursor.return_value = mock_cursor + mock_connect.return_value = mock_connection + + config = make_config( + type="duckdb", + extensions=["httpfs"], + connector_config={ + "memory_limit": "1GB", + "custom_extension_repository": "/opt/duckdb/extensions", + "autoinstall_known_extensions": "false", + }, + ) + assert isinstance(config, DuckDBConnectionConfig) + + # Create cursor which triggers _cursor_init + config.create_engine_adapter().cursor + + execute_calls = [call[0][0] for call in mock_cursor.execute.call_args_list] + + set_repository_idx = execute_calls.index( + "SET custom_extension_repository = '/opt/duckdb/extensions'" + ) + set_autoinstall_idx = execute_calls.index("SET autoinstall_known_extensions = 'false'") + install_idx = execute_calls.index("INSTALL httpfs") + load_idx = execute_calls.index("LOAD httpfs") + set_memory_limit_idx = execute_calls.index("SET memory_limit = '1GB'") + + # extension management settings run before the extensions are installed + assert set_repository_idx < install_idx + assert set_autoinstall_idx < install_idx + assert install_idx < load_idx + # the remaining connector_config settings still run after the extensions are loaded + assert load_idx < set_memory_limit_idx + + +@patch("duckdb.connect") +def test_duckdb_connector_config_without_extensions(mock_connect, make_config): + mock_cursor = MagicMock() + mock_cursor.fetchall.return_value = [] + mock_connection = MagicMock() + mock_connection.cursor.return_value = mock_cursor + mock_connect.return_value = mock_connection + + config = make_config( + type="duckdb", + connector_config={ + "custom_extension_repository": "/opt/duckdb/extensions", + "memory_limit": "1GB", + }, + ) + assert isinstance(config, DuckDBConnectionConfig) + + config.create_engine_adapter().cursor + + execute_calls = [call[0][0] for call in mock_cursor.execute.call_args_list] + + assert "SET custom_extension_repository = '/opt/duckdb/extensions'" in execute_calls + assert "SET memory_limit = '1GB'" in execute_calls + assert not any(call.startswith(("INSTALL", "LOAD")) for call in execute_calls) + + +@pytest.mark.parametrize( + "repository, expected_install", + [ + ("community", "INSTALL httpfs FROM community"), + ("core_nightly", "INSTALL httpfs FROM core_nightly"), + ("/opt/duckdb/extensions", "INSTALL httpfs FROM '/opt/duckdb/extensions'"), + ( + "https://artifactory.example.com/duckdb-extensions", + "INSTALL httpfs FROM 'https://artifactory.example.com/duckdb-extensions'", + ), + ("s3://bucket/extensions", "INSTALL httpfs FROM 's3://bucket/extensions'"), + ("/path/with'quote", "INSTALL httpfs FROM '/path/with''quote'"), + ], +) +@patch("duckdb.connect") +def test_duckdb_extension_repository_quoting( + mock_connect, make_config, repository: str, expected_install: str +): + mock_cursor = MagicMock() + mock_cursor.fetchall.return_value = [] + mock_connection = MagicMock() + mock_connection.cursor.return_value = mock_cursor + mock_connect.return_value = mock_connection + + config = make_config( + type="duckdb", + extensions=[{"name": "httpfs", "repository": repository}], + ) + assert isinstance(config, DuckDBConnectionConfig) + + config.create_engine_adapter().cursor + + execute_calls = [call[0][0] for call in mock_cursor.execute.call_args_list] + assert expected_install in execute_calls + assert "LOAD httpfs" in execute_calls + + +@patch("duckdb.connect") +def test_duckdb_extension_force_install_with_repository(mock_connect, make_config): + mock_cursor = MagicMock() + mock_cursor.fetchall.return_value = [] + mock_connection = MagicMock() + mock_connection.cursor.return_value = mock_cursor + mock_connect.return_value = mock_connection + + config = make_config( + type="duckdb", + extensions=[ + {"name": "httpfs", "repository": "/opt/duckdb/extensions", "force_install": True} + ], + ) + assert isinstance(config, DuckDBConnectionConfig) + + config.create_engine_adapter().cursor + + execute_calls = [call[0][0] for call in mock_cursor.execute.call_args_list] + assert "FORCE INSTALL httpfs FROM '/opt/duckdb/extensions'" in execute_calls + + def test_motherduck_attach_catalog(make_config): config = make_config( type="motherduck",