Skip to content

Commit af33f6b

Browse files
Tanakpekcursoragent
Tanakpek
and
cursoragent
committed
fix(duckdb): apply extension settings from connector_config before INSTALL/LOAD
Settings that control where DuckDB fetches extensions from (custom_extension_repository, extension_directory, autoinstall_known_extensions, http_proxy, ...) were applied after the extensions had already been installed, so in restricted networks INSTALL reached out to the default repository before the configured one was ever set. These settings are now applied first; all other connector_config settings are still applied after the extensions are loaded. Also quote non-alias `repository` values in `INSTALL ... FROM` so local directories and URLs work, not only aliases such as `community`. Signed-off-by: Tan Akpek <akpektan@gmail.com> Co-authored-by: Cursor <cursoragent@cursor.com>
1 parent d69c262 commit af33f6b

3 files changed

Lines changed: 309 additions & 26 deletions

File tree

‎docs/integrations/engines/duckdb.md‎

Lines changed: 107 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,8 +15,8 @@
1515
| `type` | Engine type name - must be `duckdb` | string | Y |
1616
| `database` | The optional database name. If not specified, the in-memory database is used. Cannot be defined if using `catalogs`. | string | N |
1717
| `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 |
18-
| `extensions` | Extension to load into duckdb. Only autoloadable extensions are supported. | list | N |
19-
| `connector_config` | Configuration to pass into the duckdb connector. | dict | N |
18+
| `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 |
19+
| `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 |
2020
| `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 |
2121
| `filesystems` | Configuration for registering `fsspec` filesystems to the DuckDB connection. | dict | N |
2222

@@ -201,6 +201,111 @@ Example: mounting a SQLite database with the name `sqlite` that has a table `exa
201201
If a connector, like Postgres, requires sensitive information in the path, it might support defining environment variables instead.
202202
[See DuckDB Documentation for more information](https://duckdb.org/docs/extensions/postgres#configuring-via-environment-variables).
203203

204+
#### Extensions
205+
206+
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:
207+
208+
| Key | Description | Required |
209+
|-----------------|----------------------------------------------------------------------------------------------------------------------------------------------------|:--------:|
210+
| `name` | The extension name, e.g. `httpfs`. | Y |
211+
| `repository` | Where to install the extension from. Either a repository alias (`core`, `core_nightly`, `community`), a local directory or a URL (`https://`, `s3://`). | N |
212+
| `force_install` | Reinstall the extension even if it is already installed (`FORCE INSTALL`). | N |
213+
214+
=== "YAML"
215+
216+
```yaml linenums="1"
217+
gateways:
218+
my_gateway:
219+
connection:
220+
type: duckdb
221+
extensions:
222+
- httpfs
223+
- name: h3
224+
repository: community
225+
- name: spatial
226+
repository: /opt/duckdb/extensions
227+
force_install: true
228+
```
229+
230+
=== "Python"
231+
232+
```python linenums="1"
233+
from sqlmesh.core.config import (
234+
Config,
235+
ModelDefaultsConfig,
236+
GatewayConfig,
237+
DuckDBConnectionConfig
238+
)
239+
240+
config = Config(
241+
model_defaults=ModelDefaultsConfig(dialect=<dialect>),
242+
gateways={
243+
"my_gateway": GatewayConfig(
244+
connection=DuckDBConnectionConfig(
245+
extensions=[
246+
"httpfs",
247+
{"name": "h3", "repository": "community"},
248+
{"name": "spatial", "repository": "/opt/duckdb/extensions", "force_install": True},
249+
]
250+
)
251+
),
252+
}
253+
)
254+
```
255+
256+
##### Installing extensions from a custom or offline repository
257+
258+
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.
259+
260+
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.
261+
262+
=== "YAML"
263+
264+
```yaml linenums="1"
265+
gateways:
266+
my_gateway:
267+
connection:
268+
type: duckdb
269+
connector_config:
270+
custom_extension_repository: https://artifactory.example.com/duckdb-extensions
271+
# or, for a directory of pre-downloaded extensions:
272+
# extension_directory: /opt/duckdb/extensions
273+
autoinstall_known_extensions: false
274+
extensions:
275+
- httpfs
276+
```
277+
278+
=== "Python"
279+
280+
```python linenums="1"
281+
from sqlmesh.core.config import (
282+
Config,
283+
ModelDefaultsConfig,
284+
GatewayConfig,
285+
DuckDBConnectionConfig
286+
)
287+
288+
config = Config(
289+
model_defaults=ModelDefaultsConfig(dialect=<dialect>),
290+
gateways={
291+
"my_gateway": GatewayConfig(
292+
connection=DuckDBConnectionConfig(
293+
connector_config={
294+
"custom_extension_repository": "https://artifactory.example.com/duckdb-extensions",
295+
"autoinstall_known_extensions": False,
296+
},
297+
extensions=["httpfs"],
298+
)
299+
),
300+
}
301+
)
302+
```
303+
304+
Alternatively, the repository can be set per extension with the `repository` key, which maps directly to `INSTALL <name> FROM <repository>`.
305+
306+
!!! note "Unsigned extensions"
307+
`allow_unsigned_extensions` can only be set when the DuckDB database is opened and therefore cannot be configured through `connector_config`.
308+
204309
#### Cloud service authentication
205310

206311
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)).

‎sqlmesh/core/config/connection.py‎

Lines changed: 72 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,27 @@
6666
}
6767
MOTHERDUCK_TOKEN_REGEX = re.compile(r"(\?|\&)(motherduck_token=)(\S*)")
6868
PASSWORD_REGEX = re.compile(r"(password=)(\S+)")
69+
# DuckDB settings that control how extensions are located, installed and loaded.
70+
# These must be applied before any INSTALL / LOAD statement is executed, otherwise DuckDB
71+
# tries to reach the default extension repository before e.g. a custom repository, a local
72+
# extension directory or a proxy has been configured.
73+
DUCKDB_EXTENSION_SETTINGS = frozenset(
74+
{
75+
"custom_extension_repository",
76+
"autoinstall_extension_repository",
77+
"extension_directory",
78+
"autoinstall_known_extensions",
79+
"autoload_known_extensions",
80+
"allow_community_extensions",
81+
"allow_extensions_metadata_mismatch",
82+
"http_proxy",
83+
"http_proxy_username",
84+
"http_proxy_password",
85+
}
86+
)
87+
# Repository names that DuckDB understands as aliases (e.g. `core`, `core_nightly`, `community`)
88+
# are passed to INSTALL ... FROM unquoted. Anything else (local path, URL) is a string literal.
89+
DUCKDB_EXTENSION_REPOSITORY_ALIAS_REGEX = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")
6990
SUPPORTS_MSSQL_PYTHON_DRIVER = (version_info.major, version_info.minor) >= (3, 10)
7091

7192

@@ -365,14 +386,53 @@ def _cursor_init(self) -> t.Optional[t.Callable[[t.Any], None]]:
365386
import duckdb
366387
from duckdb import BinderException
367388

389+
def apply_settings(cursor: duckdb.DuckDBPyConnection, settings: t.Dict[str, t.Any]) -> None:
390+
if not settings:
391+
return
392+
393+
option_names = list(settings)
394+
in_part = ",".join("?" for _ in range(len(option_names)))
395+
396+
cursor.execute(
397+
f"SELECT name, value FROM duckdb_settings() WHERE name IN ({in_part})",
398+
option_names,
399+
)
400+
401+
existing_values = {field: setting for field, setting in cursor.fetchall()}
402+
403+
# only set connector_config items if the values differ from what is already set
404+
# trying to set options like 'temp_directory' even to the same value can throw errors like:
405+
# Not implemented Error: Cannot switch temporary directory after the current one has been used
406+
for field, setting in settings.items():
407+
if existing_values.get(field) != setting:
408+
try:
409+
cursor.execute(f"SET {field} = '{setting}'")
410+
except Exception as e:
411+
raise ConfigError(
412+
f"Failed to set connector config {field} to {setting}: {e}"
413+
)
414+
368415
def init(cursor: duckdb.DuckDBPyConnection) -> None:
416+
# Settings that control where extensions come from have to be in place before the
417+
# first INSTALL / LOAD, otherwise DuckDB reaches out to the default repository first.
418+
apply_settings(
419+
cursor,
420+
{
421+
field: setting
422+
for field, setting in self.connector_config.items()
423+
if field in DUCKDB_EXTENSION_SETTINGS
424+
},
425+
)
426+
369427
for extension in self.extensions:
370428
extension = extension if isinstance(extension, dict) else {"name": extension}
371429

372430
install_command = f"INSTALL {extension['name']}"
373431

374-
if extension.get("repository"):
375-
install_command = f"{install_command} FROM {extension['repository']}"
432+
if repository := extension.get("repository"):
433+
if not DUCKDB_EXTENSION_REPOSITORY_ALIAS_REGEX.match(repository):
434+
repository = exp.Literal.string(repository).sql(dialect="duckdb")
435+
install_command = f"{install_command} FROM {repository}"
376436

377437
if extension.get("force_install"):
378438
install_command = f"FORCE {install_command}"
@@ -383,28 +443,16 @@ def init(cursor: duckdb.DuckDBPyConnection) -> None:
383443
except Exception as e:
384444
raise ConfigError(f"Failed to load extension {extension['name']}: {e}")
385445

386-
if self.connector_config:
387-
option_names = list(self.connector_config)
388-
in_part = ",".join("?" for _ in range(len(option_names)))
389-
390-
cursor.execute(
391-
f"SELECT name, value FROM duckdb_settings() WHERE name IN ({in_part})",
392-
option_names,
393-
)
394-
395-
existing_values = {field: setting for field, setting in cursor.fetchall()}
396-
397-
# only set connector_config items if the values differ from what is already set
398-
# trying to set options like 'temp_directory' even to the same value can throw errors like:
399-
# Not implemented Error: Cannot switch temporary directory after the current one has been used
400-
for field, setting in self.connector_config.items():
401-
if existing_values.get(field) != setting:
402-
try:
403-
cursor.execute(f"SET {field} = '{setting}'")
404-
except Exception as e:
405-
raise ConfigError(
406-
f"Failed to set connector config {field} to {setting}: {e}"
407-
)
446+
# The remaining settings may belong to extensions (e.g. `s3_region` from httpfs),
447+
# so they are applied after the extensions have been loaded.
448+
apply_settings(
449+
cursor,
450+
{
451+
field: setting
452+
for field, setting in self.connector_config.items()
453+
if field not in DUCKDB_EXTENSION_SETTINGS
454+
},
455+
)
408456

409457
if self.secrets:
410458
duckdb_version = duckdb.__version__

‎tests/core/test_connection_config.py‎

Lines changed: 130 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -935,6 +935,136 @@ def test_duckdb_config_json_strings(make_config):
935935
assert config.catalogs.get("test2").path == "test2.duckdb"
936936

937937

938+
@patch("duckdb.connect")
939+
def test_duckdb_extension_settings_applied_before_install(mock_connect, make_config):
940+
"""Settings that control where extensions are fetched from must be applied before INSTALL/LOAD,
941+
otherwise DuckDB will try to reach the default repository before e.g. `custom_extension_repository`
942+
or `extension_directory` has been configured. All other connector_config settings keep being applied
943+
after the extensions have been loaded."""
944+
mock_cursor = MagicMock()
945+
# No settings are currently set, so every connector_config entry should be SET
946+
mock_cursor.fetchall.return_value = []
947+
mock_connection = MagicMock()
948+
mock_connection.cursor.return_value = mock_cursor
949+
mock_connect.return_value = mock_connection
950+
951+
config = make_config(
952+
type="duckdb",
953+
extensions=["httpfs"],
954+
connector_config={
955+
"memory_limit": "1GB",
956+
"custom_extension_repository": "/opt/duckdb/extensions",
957+
"autoinstall_known_extensions": "false",
958+
},
959+
)
960+
assert isinstance(config, DuckDBConnectionConfig)
961+
962+
# Create cursor which triggers _cursor_init
963+
config.create_engine_adapter().cursor
964+
965+
execute_calls = [call[0][0] for call in mock_cursor.execute.call_args_list]
966+
967+
set_repository_idx = execute_calls.index(
968+
"SET custom_extension_repository = '/opt/duckdb/extensions'"
969+
)
970+
set_autoinstall_idx = execute_calls.index("SET autoinstall_known_extensions = 'false'")
971+
install_idx = execute_calls.index("INSTALL httpfs")
972+
load_idx = execute_calls.index("LOAD httpfs")
973+
set_memory_limit_idx = execute_calls.index("SET memory_limit = '1GB'")
974+
975+
# extension management settings run before the extensions are installed
976+
assert set_repository_idx < install_idx
977+
assert set_autoinstall_idx < install_idx
978+
assert install_idx < load_idx
979+
# the remaining connector_config settings still run after the extensions are loaded
980+
assert load_idx < set_memory_limit_idx
981+
982+
983+
@patch("duckdb.connect")
984+
def test_duckdb_connector_config_without_extensions(mock_connect, make_config):
985+
mock_cursor = MagicMock()
986+
mock_cursor.fetchall.return_value = []
987+
mock_connection = MagicMock()
988+
mock_connection.cursor.return_value = mock_cursor
989+
mock_connect.return_value = mock_connection
990+
991+
config = make_config(
992+
type="duckdb",
993+
connector_config={
994+
"custom_extension_repository": "/opt/duckdb/extensions",
995+
"memory_limit": "1GB",
996+
},
997+
)
998+
assert isinstance(config, DuckDBConnectionConfig)
999+
1000+
config.create_engine_adapter().cursor
1001+
1002+
execute_calls = [call[0][0] for call in mock_cursor.execute.call_args_list]
1003+
1004+
assert "SET custom_extension_repository = '/opt/duckdb/extensions'" in execute_calls
1005+
assert "SET memory_limit = '1GB'" in execute_calls
1006+
assert not any(call.startswith(("INSTALL", "LOAD")) for call in execute_calls)
1007+
1008+
1009+
@pytest.mark.parametrize(
1010+
"repository, expected_install",
1011+
[
1012+
("community", "INSTALL httpfs FROM community"),
1013+
("core_nightly", "INSTALL httpfs FROM core_nightly"),
1014+
("/opt/duckdb/extensions", "INSTALL httpfs FROM '/opt/duckdb/extensions'"),
1015+
(
1016+
"https://artifactory.example.com/duckdb-extensions",
1017+
"INSTALL httpfs FROM 'https://artifactory.example.com/duckdb-extensions'",
1018+
),
1019+
("s3://bucket/extensions", "INSTALL httpfs FROM 's3://bucket/extensions'"),
1020+
("/path/with'quote", "INSTALL httpfs FROM '/path/with''quote'"),
1021+
],
1022+
)
1023+
@patch("duckdb.connect")
1024+
def test_duckdb_extension_repository_quoting(
1025+
mock_connect, make_config, repository: str, expected_install: str
1026+
):
1027+
mock_cursor = MagicMock()
1028+
mock_cursor.fetchall.return_value = []
1029+
mock_connection = MagicMock()
1030+
mock_connection.cursor.return_value = mock_cursor
1031+
mock_connect.return_value = mock_connection
1032+
1033+
config = make_config(
1034+
type="duckdb",
1035+
extensions=[{"name": "httpfs", "repository": repository}],
1036+
)
1037+
assert isinstance(config, DuckDBConnectionConfig)
1038+
1039+
config.create_engine_adapter().cursor
1040+
1041+
execute_calls = [call[0][0] for call in mock_cursor.execute.call_args_list]
1042+
assert expected_install in execute_calls
1043+
assert "LOAD httpfs" in execute_calls
1044+
1045+
1046+
@patch("duckdb.connect")
1047+
def test_duckdb_extension_force_install_with_repository(mock_connect, make_config):
1048+
mock_cursor = MagicMock()
1049+
mock_cursor.fetchall.return_value = []
1050+
mock_connection = MagicMock()
1051+
mock_connection.cursor.return_value = mock_cursor
1052+
mock_connect.return_value = mock_connection
1053+
1054+
config = make_config(
1055+
type="duckdb",
1056+
extensions=[
1057+
{"name": "httpfs", "repository": "/opt/duckdb/extensions", "force_install": True}
1058+
],
1059+
)
1060+
assert isinstance(config, DuckDBConnectionConfig)
1061+
1062+
config.create_engine_adapter().cursor
1063+
1064+
execute_calls = [call[0][0] for call in mock_cursor.execute.call_args_list]
1065+
assert "FORCE INSTALL httpfs FROM '/opt/duckdb/extensions'" in execute_calls
1066+
1067+
9381068
def test_motherduck_attach_catalog(make_config):
9391069
config = make_config(
9401070
type="motherduck",

0 commit comments

Comments
 (0)