dirigent-block-duckdb 0.17.2__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- dirigent_block_duckdb-0.17.2/LICENSE +18 -0
- dirigent_block_duckdb-0.17.2/PKG-INFO +60 -0
- dirigent_block_duckdb-0.17.2/README.md +44 -0
- dirigent_block_duckdb-0.17.2/pyproject.toml +35 -0
- dirigent_block_duckdb-0.17.2/pyproject.toml.orig +31 -0
- dirigent_block_duckdb-0.17.2/src/dirigent_block_duckdb/__init__.py +25 -0
- dirigent_block_duckdb-0.17.2/src/dirigent_block_duckdb/engine.py +435 -0
- dirigent_block_duckdb-0.17.2/src/dirigent_block_duckdb/messages.py +54 -0
- dirigent_block_duckdb-0.17.2/src/dirigent_block_duckdb/py.typed +0 -0
|
@@ -0,0 +1,18 @@
|
|
|
1
|
+
Copyright (c) 2026 Morten Olav Hansen <morten@winterop.com>. All rights reserved.
|
|
2
|
+
|
|
3
|
+
This source code and accompanying documentation are the property of
|
|
4
|
+
Morten Olav Hansen. No license, express or implied, is granted to use, copy,
|
|
5
|
+
modify, merge, publish, distribute, sublicense, or sell copies of this
|
|
6
|
+
software or its derivatives.
|
|
7
|
+
|
|
8
|
+
The source is published for reference only. Any use beyond reading
|
|
9
|
+
requires written permission from the copyright holder.
|
|
10
|
+
|
|
11
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
|
|
12
|
+
OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
|
|
13
|
+
MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, AND NONINFRINGEMENT.
|
|
14
|
+
IN NO EVENT SHALL THE COPYRIGHT HOLDER BE LIABLE FOR ANY CLAIM, DAMAGES,
|
|
15
|
+
OR OTHER LIABILITY ARISING FROM THE USE OF THE SOFTWARE.
|
|
16
|
+
|
|
17
|
+
Third-party components redistributed with this software, and the licences they
|
|
18
|
+
carry, are listed in THIRD_PARTY_NOTICES.md.
|
|
@@ -0,0 +1,60 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: dirigent-block-duckdb
|
|
3
|
+
Version: 0.17.2
|
|
4
|
+
Summary: The DuckDB engine of dirigent's sql family: sql.query and sql.execute over files.
|
|
5
|
+
License-Expression: LicenseRef-Proprietary
|
|
6
|
+
License-File: LICENSE
|
|
7
|
+
Classifier: Programming Language :: Python :: 3
|
|
8
|
+
Classifier: Programming Language :: Python :: 3.13
|
|
9
|
+
Requires-Dist: dirigent-block-sql==0.17.2
|
|
10
|
+
Requires-Dist: dirigent-common==0.17.2
|
|
11
|
+
Requires-Dist: dirigent-plugin==0.17.2
|
|
12
|
+
Requires-Dist: duckdb>=1.4.0
|
|
13
|
+
Requires-Dist: duckdb-engine>=0.17.0
|
|
14
|
+
Requires-Python: >=3.13
|
|
15
|
+
Description-Content-Type: text/markdown
|
|
16
|
+
|
|
17
|
+
# dirigent-block-duckdb
|
|
18
|
+
|
|
19
|
+
The DuckDB engine of dirigent's `sql` family. Install it beside `dirigent-block-sql` and a
|
|
20
|
+
`sql` connection whose url is `duckdb:///warehouse.duckdb` or `duckdb:///:memory:` works like
|
|
21
|
+
any other: the same two blocks, the same fields, the same rules.
|
|
22
|
+
|
|
23
|
+
```bash
|
|
24
|
+
uv pip install dirigent-block-duckdb
|
|
25
|
+
```
|
|
26
|
+
|
|
27
|
+
It registers under the `dirigent.sql.engines.v1` entry-point group, so the family finds it with
|
|
28
|
+
no configuration, and a worker that never runs a `duckdb://` url carries none of it -- the
|
|
29
|
+
engine binary is larger than every other driver the family speaks to put together.
|
|
30
|
+
|
|
31
|
+
DuckDB has no async driver at all, so every call runs in a worker thread, and a deadline that
|
|
32
|
+
passes interrupts the engine rather than leaving a statement running behind a step that has
|
|
33
|
+
already failed.
|
|
34
|
+
|
|
35
|
+
## Files, and what holds them in
|
|
36
|
+
|
|
37
|
+
A statement here reads and writes the parquet and csv files a run holds: `read_parquet(:source)`
|
|
38
|
+
and `COPY ... TO :target` take a file the same way a `WHERE` clause takes a value, and a
|
|
39
|
+
`file://` storage URI in `params` arrives as the path duckdb opens.
|
|
40
|
+
|
|
41
|
+
The boundary is duckdb's own, not a check on the parameters. A session is opened, given the
|
|
42
|
+
run's work directory and its local scratch space as its `allowed_directories`, and then closed
|
|
43
|
+
around them: `enable_external_access` goes off, which is what makes those roots the only paths
|
|
44
|
+
the engine will open, and `lock_configuration` goes on, which refuses the `SET` that would give
|
|
45
|
+
any of it back. So a path written straight into the `sql` is refused the same way a parameter
|
|
46
|
+
outside the run is, and a statement cannot `LOAD` or `INSTALL` a further extension to reach past
|
|
47
|
+
the boundary. Spilled intermediates and duckdb's secret store are pointed inside the run's work
|
|
48
|
+
directory for the same reason.
|
|
49
|
+
|
|
50
|
+
Where a parameter or a statement names an `s3://` object, the session loads duckdb's `httpfs`
|
|
51
|
+
extension and is given the endpoint, region, credential and addressing style of the connection
|
|
52
|
+
the `s3` scheme is configured from, and that scheme stays reachable beside the run's own
|
|
53
|
+
directories. The extension is loaded, never installed, at run time; a bare install does it once:
|
|
54
|
+
|
|
55
|
+
```bash
|
|
56
|
+
python -c "import duckdb; duckdb.connect().execute('INSTALL httpfs')"
|
|
57
|
+
```
|
|
58
|
+
|
|
59
|
+
[docs/sql.md](https://github.com/winterop-com/dirigent/blob/main/docs/sql.md) is the family's
|
|
60
|
+
home, and its DuckDB section is the worked example.
|
|
@@ -0,0 +1,44 @@
|
|
|
1
|
+
# dirigent-block-duckdb
|
|
2
|
+
|
|
3
|
+
The DuckDB engine of dirigent's `sql` family. Install it beside `dirigent-block-sql` and a
|
|
4
|
+
`sql` connection whose url is `duckdb:///warehouse.duckdb` or `duckdb:///:memory:` works like
|
|
5
|
+
any other: the same two blocks, the same fields, the same rules.
|
|
6
|
+
|
|
7
|
+
```bash
|
|
8
|
+
uv pip install dirigent-block-duckdb
|
|
9
|
+
```
|
|
10
|
+
|
|
11
|
+
It registers under the `dirigent.sql.engines.v1` entry-point group, so the family finds it with
|
|
12
|
+
no configuration, and a worker that never runs a `duckdb://` url carries none of it -- the
|
|
13
|
+
engine binary is larger than every other driver the family speaks to put together.
|
|
14
|
+
|
|
15
|
+
DuckDB has no async driver at all, so every call runs in a worker thread, and a deadline that
|
|
16
|
+
passes interrupts the engine rather than leaving a statement running behind a step that has
|
|
17
|
+
already failed.
|
|
18
|
+
|
|
19
|
+
## Files, and what holds them in
|
|
20
|
+
|
|
21
|
+
A statement here reads and writes the parquet and csv files a run holds: `read_parquet(:source)`
|
|
22
|
+
and `COPY ... TO :target` take a file the same way a `WHERE` clause takes a value, and a
|
|
23
|
+
`file://` storage URI in `params` arrives as the path duckdb opens.
|
|
24
|
+
|
|
25
|
+
The boundary is duckdb's own, not a check on the parameters. A session is opened, given the
|
|
26
|
+
run's work directory and its local scratch space as its `allowed_directories`, and then closed
|
|
27
|
+
around them: `enable_external_access` goes off, which is what makes those roots the only paths
|
|
28
|
+
the engine will open, and `lock_configuration` goes on, which refuses the `SET` that would give
|
|
29
|
+
any of it back. So a path written straight into the `sql` is refused the same way a parameter
|
|
30
|
+
outside the run is, and a statement cannot `LOAD` or `INSTALL` a further extension to reach past
|
|
31
|
+
the boundary. Spilled intermediates and duckdb's secret store are pointed inside the run's work
|
|
32
|
+
directory for the same reason.
|
|
33
|
+
|
|
34
|
+
Where a parameter or a statement names an `s3://` object, the session loads duckdb's `httpfs`
|
|
35
|
+
extension and is given the endpoint, region, credential and addressing style of the connection
|
|
36
|
+
the `s3` scheme is configured from, and that scheme stays reachable beside the run's own
|
|
37
|
+
directories. The extension is loaded, never installed, at run time; a bare install does it once:
|
|
38
|
+
|
|
39
|
+
```bash
|
|
40
|
+
python -c "import duckdb; duckdb.connect().execute('INSTALL httpfs')"
|
|
41
|
+
```
|
|
42
|
+
|
|
43
|
+
[docs/sql.md](https://github.com/winterop-com/dirigent/blob/main/docs/sql.md) is the family's
|
|
44
|
+
home, and its DuckDB section is the worked example.
|
|
@@ -0,0 +1,35 @@
|
|
|
1
|
+
[project]
|
|
2
|
+
name = "dirigent-block-duckdb"
|
|
3
|
+
version = "0.17.2"
|
|
4
|
+
description = "The DuckDB engine of dirigent's sql family: sql.query and sql.execute over files."
|
|
5
|
+
readme = "README.md"
|
|
6
|
+
requires-python = ">=3.13"
|
|
7
|
+
license = "LicenseRef-Proprietary"
|
|
8
|
+
license-files = ["LICENSE"]
|
|
9
|
+
classifiers = [
|
|
10
|
+
"Programming Language :: Python :: 3",
|
|
11
|
+
"Programming Language :: Python :: 3.13",
|
|
12
|
+
]
|
|
13
|
+
dependencies = [
|
|
14
|
+
"dirigent-block-sql==0.17.2",
|
|
15
|
+
"dirigent-common==0.17.2",
|
|
16
|
+
"dirigent-plugin==0.17.2",
|
|
17
|
+
"duckdb>=1.4.0",
|
|
18
|
+
"duckdb-engine>=0.17.0",
|
|
19
|
+
]
|
|
20
|
+
|
|
21
|
+
[project.entry-points."dirigent.sql.engines.v1"]
|
|
22
|
+
duckdb = "dirigent_block_duckdb:plugin"
|
|
23
|
+
|
|
24
|
+
[build-system]
|
|
25
|
+
requires = ["uv_build>=0.12.0,<0.13.0"]
|
|
26
|
+
build-backend = "uv_build"
|
|
27
|
+
|
|
28
|
+
[tool.uv.sources.dirigent-block-sql]
|
|
29
|
+
workspace = true
|
|
30
|
+
|
|
31
|
+
[tool.uv.sources.dirigent-common]
|
|
32
|
+
workspace = true
|
|
33
|
+
|
|
34
|
+
[tool.uv.sources.dirigent-plugin]
|
|
35
|
+
workspace = true
|
|
@@ -0,0 +1,31 @@
|
|
|
1
|
+
[project]
|
|
2
|
+
name = "dirigent-block-duckdb"
|
|
3
|
+
version = "0.17.2"
|
|
4
|
+
description = "The DuckDB engine of dirigent's sql family: sql.query and sql.execute over files."
|
|
5
|
+
readme = "README.md"
|
|
6
|
+
requires-python = ">=3.13"
|
|
7
|
+
license = "LicenseRef-Proprietary"
|
|
8
|
+
license-files = ["LICENSE"]
|
|
9
|
+
classifiers = [
|
|
10
|
+
"Programming Language :: Python :: 3",
|
|
11
|
+
"Programming Language :: Python :: 3.13",
|
|
12
|
+
]
|
|
13
|
+
dependencies = [
|
|
14
|
+
"dirigent-block-sql==0.17.2",
|
|
15
|
+
"dirigent-common==0.17.2",
|
|
16
|
+
"dirigent-plugin==0.17.2",
|
|
17
|
+
"duckdb>=1.4.0",
|
|
18
|
+
"duckdb-engine>=0.17.0",
|
|
19
|
+
]
|
|
20
|
+
|
|
21
|
+
[project.entry-points."dirigent.sql.engines.v1"]
|
|
22
|
+
duckdb = "dirigent_block_duckdb:plugin"
|
|
23
|
+
|
|
24
|
+
[build-system]
|
|
25
|
+
requires = ["uv_build>=0.12.0,<0.13.0"]
|
|
26
|
+
build-backend = "uv_build"
|
|
27
|
+
|
|
28
|
+
[tool.uv.sources]
|
|
29
|
+
dirigent-block-sql = { workspace = true }
|
|
30
|
+
dirigent-common = { workspace = true }
|
|
31
|
+
dirigent-plugin = { workspace = true }
|
|
@@ -0,0 +1,25 @@
|
|
|
1
|
+
"""The DuckDB engine of dirigent's ``sql`` family, contributed to the family's engine registry."""
|
|
2
|
+
|
|
3
|
+
from collections.abc import Sequence
|
|
4
|
+
|
|
5
|
+
from dirigent_block_duckdb.engine import DuckdbEngine
|
|
6
|
+
from dirigent_block_sql import SqlEngine
|
|
7
|
+
from dirigent_plugin import extension
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
class DuckdbEngines:
|
|
11
|
+
"""The plugin object the family discovers under the dirigent.sql.engines.v1 entry-point group."""
|
|
12
|
+
|
|
13
|
+
@extension
|
|
14
|
+
def engines(self) -> Sequence[SqlEngine]:
|
|
15
|
+
"""Contribute the duckdb engine, which a url naming the duckdb backend is driven through."""
|
|
16
|
+
return [DuckdbEngine()]
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
plugin = DuckdbEngines()
|
|
20
|
+
|
|
21
|
+
__all__ = [
|
|
22
|
+
"DuckdbEngine",
|
|
23
|
+
"DuckdbEngines",
|
|
24
|
+
"plugin",
|
|
25
|
+
]
|
|
@@ -0,0 +1,435 @@
|
|
|
1
|
+
"""The DuckDB engine of the ``sql`` family: the backend whose tables can be files.
|
|
2
|
+
|
|
3
|
+
DuckDB has no async driver at all: neither its own Python API nor ``duckdb_engine`` speaks the
|
|
4
|
+
DBAPI SQLAlchemy's asyncio layer needs, and that layer has no adapter for a synchronous one. So
|
|
5
|
+
every call is made from a worker thread, and a deadline that passes interrupts the engine rather
|
|
6
|
+
than leaving a statement running behind a step that has already failed.
|
|
7
|
+
|
|
8
|
+
A statement here reads and writes the parquet and csv files a run holds, by binding their
|
|
9
|
+
storage URIs as parameters. The session is held to the run's own directories by duckdb's own
|
|
10
|
+
configuration -- ``allowed_directories``, then ``enable_external_access`` off and
|
|
11
|
+
``lock_configuration`` on -- so a path written into the SQL reaches no further than a bound one
|
|
12
|
+
does. Where a step names an ``s3://`` object, the session is given the ``httpfs`` extension and
|
|
13
|
+
the credentials of the connection the ``s3`` scheme is configured from, and that scheme stays
|
|
14
|
+
reachable beside the run's directories.
|
|
15
|
+
"""
|
|
16
|
+
|
|
17
|
+
import asyncio
|
|
18
|
+
from collections.abc import AsyncGenerator, AsyncIterator, Callable, Sequence
|
|
19
|
+
from contextlib import AbstractAsyncContextManager, asynccontextmanager
|
|
20
|
+
from pathlib import Path
|
|
21
|
+
from typing import Any, ClassVar, Final
|
|
22
|
+
from urllib.parse import urlsplit
|
|
23
|
+
|
|
24
|
+
import sqlalchemy
|
|
25
|
+
import sqlalchemy.exc
|
|
26
|
+
from pydantic import SecretStr
|
|
27
|
+
from sqlalchemy.engine import URL, Connection, Engine, create_engine
|
|
28
|
+
from sqlalchemy.pool import NullPool
|
|
29
|
+
from sqlalchemy.sql.elements import TextClause
|
|
30
|
+
|
|
31
|
+
from dirigent_block_duckdb.messages import (
|
|
32
|
+
CONNECT_TIMED_OUT,
|
|
33
|
+
NO_HTTPFS,
|
|
34
|
+
NO_STORAGE_CONNECTION,
|
|
35
|
+
PARAMETER_NAMES_STORAGE,
|
|
36
|
+
PARAMETER_OUTSIDE_THE_RUN,
|
|
37
|
+
READ_ONLY_IN_MEMORY,
|
|
38
|
+
STATEMENT_LOADS_AN_EXTENSION,
|
|
39
|
+
STATEMENT_OUTSIDE_THE_RUN,
|
|
40
|
+
)
|
|
41
|
+
from dirigent_block_sql.engines import (
|
|
42
|
+
CHECK_TIMEOUT_SECONDS,
|
|
43
|
+
OUTSIDE_A_RUN,
|
|
44
|
+
SqlEngine,
|
|
45
|
+
SqlSession,
|
|
46
|
+
reason,
|
|
47
|
+
under_work,
|
|
48
|
+
)
|
|
49
|
+
from dirigent_block_sql.sql import SqlConnectionConfig
|
|
50
|
+
from dirigent_common import BlockModel, HealthReport, JsonMap
|
|
51
|
+
from dirigent_plugin import BlockFailure, ErrorClass, StepContext
|
|
52
|
+
|
|
53
|
+
#: The storage schemes a parameter naming one is refused for: duckdb would open them itself,
|
|
54
|
+
#: and no backend here holds the credentials to give it. ``s3://`` is the one that is wired up.
|
|
55
|
+
REMOTE_STORAGE: Final = frozenset({"gs", "azure"})
|
|
56
|
+
|
|
57
|
+
#: The scheme duckdb reads and writes through its httpfs extension.
|
|
58
|
+
S3: Final = "s3"
|
|
59
|
+
|
|
60
|
+
#: The extension that gives duckdb ``s3://``. It is loaded, never installed, at run time: a
|
|
61
|
+
#: step must not fetch a binary from the internet mid-run, so the image installs it at build
|
|
62
|
+
#: time and a bare install does it once by hand.
|
|
63
|
+
HTTPFS: Final = "httpfs"
|
|
64
|
+
|
|
65
|
+
#: Where a contained session spills a query too large for memory, under the run's work
|
|
66
|
+
#: directory. duckdb's own default is beside the database file or in the worker's cwd, and
|
|
67
|
+
#: neither is inside the roots the session is held to.
|
|
68
|
+
TEMP_DIR: Final = ".duckdb-temp"
|
|
69
|
+
|
|
70
|
+
#: Where a contained session looks for duckdb secrets, under the run's work directory. The
|
|
71
|
+
#: default is the worker's home, which the session may not read, and httpfs reads the secret
|
|
72
|
+
#: store on every remote open.
|
|
73
|
+
SECRET_DIR: Final = ".duckdb-secrets"
|
|
74
|
+
|
|
75
|
+
#: What duckdb says when the containment refused the path a statement named.
|
|
76
|
+
DENIED_PATH: Final = "file system operations are disabled"
|
|
77
|
+
|
|
78
|
+
#: What duckdb says when a statement tried to load an extension the session was not opened with.
|
|
79
|
+
DENIED_EXTENSION: Final = "loading external extensions is disabled"
|
|
80
|
+
|
|
81
|
+
|
|
82
|
+
class DuckdbEngine(SqlEngine):
|
|
83
|
+
"""DuckDB: a worker thread where another engine has an async driver, and files where it has tables."""
|
|
84
|
+
|
|
85
|
+
backend: ClassVar[str] = "duckdb"
|
|
86
|
+
|
|
87
|
+
def validate(self, settings: SqlConnectionConfig, url: URL) -> None:
|
|
88
|
+
"""Refuse ``read_only`` on an in-memory database, which duckdb will not open at all."""
|
|
89
|
+
if settings.read_only and _is_memory(url):
|
|
90
|
+
raise ValueError(READ_ONLY_IN_MEMORY.render())
|
|
91
|
+
|
|
92
|
+
def resolve(self, url: URL, ctx: StepContext) -> URL:
|
|
93
|
+
"""A duckdb database written as a relative path is a file in the run's work directory."""
|
|
94
|
+
return under_work(url, ctx) if _is_run_relative(url) else url
|
|
95
|
+
|
|
96
|
+
def bind(self, params: JsonMap, ctx: StepContext) -> dict[str, Any]:
|
|
97
|
+
"""The parameters as duckdb receives them, storage URIs among them made local paths.
|
|
98
|
+
|
|
99
|
+
DuckDB reads and writes the files a statement names -- ``read_parquet(:source)``,
|
|
100
|
+
``COPY ... TO :target`` -- and a document names a file by its storage URI.
|
|
101
|
+
"""
|
|
102
|
+
return {
|
|
103
|
+
name: _file(name, value, ctx) if isinstance(value, str) and "://" in value else value
|
|
104
|
+
for name, value in params.items()
|
|
105
|
+
}
|
|
106
|
+
|
|
107
|
+
async def check(self, settings: SqlConnectionConfig, url: URL) -> HealthReport:
|
|
108
|
+
"""Open the duckdb file in a thread and ask it for a one, which is what a step does too."""
|
|
109
|
+
if _is_run_relative(url):
|
|
110
|
+
return HealthReport(healthy=False, detail=OUTSIDE_A_RUN)
|
|
111
|
+
engine = _engine(settings, url)
|
|
112
|
+
|
|
113
|
+
def ask() -> None:
|
|
114
|
+
with engine.connect() as connection:
|
|
115
|
+
connection.execute(sqlalchemy.text("SELECT 1"))
|
|
116
|
+
|
|
117
|
+
try:
|
|
118
|
+
async with asyncio.timeout(CHECK_TIMEOUT_SECONDS):
|
|
119
|
+
await asyncio.to_thread(ask)
|
|
120
|
+
except (TimeoutError, sqlalchemy.exc.SQLAlchemyError, OSError) as error:
|
|
121
|
+
return HealthReport(healthy=False, detail=reason(error))
|
|
122
|
+
finally:
|
|
123
|
+
await asyncio.to_thread(engine.dispose)
|
|
124
|
+
return HealthReport(healthy=True, detail="duckdb answered")
|
|
125
|
+
|
|
126
|
+
def session(
|
|
127
|
+
self,
|
|
128
|
+
settings: SqlConnectionConfig,
|
|
129
|
+
url: URL,
|
|
130
|
+
ctx: StepContext,
|
|
131
|
+
*,
|
|
132
|
+
statements: Sequence[str],
|
|
133
|
+
params: JsonMap,
|
|
134
|
+
) -> AbstractAsyncContextManager[SqlSession]:
|
|
135
|
+
"""One duckdb session, given the ``s3`` scheme where anything this step runs names a bucket."""
|
|
136
|
+
return _Session(settings, url, ctx, s3=addresses_s3(params, statements))
|
|
137
|
+
|
|
138
|
+
|
|
139
|
+
class S3StorageSettings(BlockModel):
|
|
140
|
+
"""What the ``s3`` storage connection carries, as duckdb needs to be told it.
|
|
141
|
+
|
|
142
|
+
The field names are the ``s3`` connection kind's own, so the row the instance already
|
|
143
|
+
holds for ``storage_connections`` validates against this without a second connection.
|
|
144
|
+
"""
|
|
145
|
+
|
|
146
|
+
endpoint_url: str | None = None
|
|
147
|
+
region: str = "us-east-1"
|
|
148
|
+
access_key_id: str | None = None
|
|
149
|
+
secret_access_key: SecretStr | None = None
|
|
150
|
+
path_style: bool = False
|
|
151
|
+
verify_tls: bool = True
|
|
152
|
+
bucket: str | None = None
|
|
153
|
+
|
|
154
|
+
|
|
155
|
+
def addresses_s3(values: dict[str, Any], statements: Sequence[str]) -> bool:
|
|
156
|
+
"""Whether anything this step runs names an ``s3://`` object, in a value or in the sql."""
|
|
157
|
+
prefix = f"{S3}://"
|
|
158
|
+
if any(isinstance(value, str) and value.startswith(prefix) for value in values.values()):
|
|
159
|
+
return True
|
|
160
|
+
return any(prefix in statement for statement in statements)
|
|
161
|
+
|
|
162
|
+
|
|
163
|
+
def s3_options(config: S3StorageSettings) -> list[tuple[str, str | bool]]:
|
|
164
|
+
"""The duckdb settings that point httpfs at one endpoint with one credential.
|
|
165
|
+
|
|
166
|
+
duckdb takes the endpoint as host and port with no scheme, and asks separately whether to
|
|
167
|
+
speak TLS to it, so one ``endpoint_url`` becomes two settings.
|
|
168
|
+
"""
|
|
169
|
+
options: list[tuple[str, str | bool]] = [("s3_region", config.region)]
|
|
170
|
+
if config.endpoint_url:
|
|
171
|
+
split = urlsplit(config.endpoint_url)
|
|
172
|
+
options.append(("s3_endpoint", split.netloc or split.path))
|
|
173
|
+
options.append(("s3_use_ssl", split.scheme == "https"))
|
|
174
|
+
else:
|
|
175
|
+
options.append(("s3_use_ssl", config.verify_tls))
|
|
176
|
+
if config.access_key_id:
|
|
177
|
+
options.append(("s3_access_key_id", config.access_key_id))
|
|
178
|
+
if config.secret_access_key is not None:
|
|
179
|
+
options.append(("s3_secret_access_key", config.secret_access_key.get_secret_value()))
|
|
180
|
+
options.append(("s3_url_style", "path" if config.path_style else "vhost"))
|
|
181
|
+
return options
|
|
182
|
+
|
|
183
|
+
|
|
184
|
+
def _is_memory(url: URL) -> bool:
|
|
185
|
+
"""Whether this URL names a database that lives in the process and nowhere else."""
|
|
186
|
+
return (url.database or ":memory:") == ":memory:"
|
|
187
|
+
|
|
188
|
+
|
|
189
|
+
def _is_run_relative(url: URL) -> bool:
|
|
190
|
+
"""Whether this is a duckdb file named by a path relative to the run's work directory."""
|
|
191
|
+
database = url.database or ""
|
|
192
|
+
if _is_memory(url):
|
|
193
|
+
return False
|
|
194
|
+
return bool(database) and not database.startswith("/")
|
|
195
|
+
|
|
196
|
+
|
|
197
|
+
def _engine(settings: SqlConnectionConfig, url: URL) -> Engine:
|
|
198
|
+
"""Build the synchronous duckdb engine, opening the file read-only where the connection says so."""
|
|
199
|
+
return create_engine(
|
|
200
|
+
url,
|
|
201
|
+
poolclass=NullPool,
|
|
202
|
+
connect_args={"read_only": True} if settings.read_only else {},
|
|
203
|
+
)
|
|
204
|
+
|
|
205
|
+
|
|
206
|
+
def _file(name: str, uri: str, ctx: StepContext) -> str:
|
|
207
|
+
"""One storage URI as duckdb receives it: a local path, or a bucket URI it opens itself."""
|
|
208
|
+
split = urlsplit(uri)
|
|
209
|
+
if split.scheme in REMOTE_STORAGE:
|
|
210
|
+
raise BlockFailure(
|
|
211
|
+
PARAMETER_NAMES_STORAGE, error_class=ErrorClass.REJECTED, name=repr(name), scheme=split.scheme
|
|
212
|
+
)
|
|
213
|
+
# An s3:// URI is handed over whole: the session loads httpfs and gives duckdb the
|
|
214
|
+
# scheme's own credentials, so duckdb opens the object rather than a path.
|
|
215
|
+
if split.scheme == S3:
|
|
216
|
+
return uri
|
|
217
|
+
if split.scheme != "file":
|
|
218
|
+
return uri
|
|
219
|
+
roots = _readable(ctx)
|
|
220
|
+
path = Path(f"{split.netloc}{split.path}").absolute()
|
|
221
|
+
if not any(path.is_relative_to(root) for root in roots):
|
|
222
|
+
named = ", ".join(str(root) for root in roots)
|
|
223
|
+
raise BlockFailure(
|
|
224
|
+
PARAMETER_OUTSIDE_THE_RUN,
|
|
225
|
+
error_class=ErrorClass.REJECTED,
|
|
226
|
+
name=repr(name),
|
|
227
|
+
uri=uri,
|
|
228
|
+
named=named,
|
|
229
|
+
)
|
|
230
|
+
# A ``COPY ... TO`` names a file duckdb creates but not a directory it creates.
|
|
231
|
+
path.parent.mkdir(parents=True, exist_ok=True)
|
|
232
|
+
return str(path)
|
|
233
|
+
|
|
234
|
+
|
|
235
|
+
def _readable(ctx: StepContext) -> list[Path]:
|
|
236
|
+
"""The directories a duckdb parameter may name a file inside.
|
|
237
|
+
|
|
238
|
+
The run's work directory always, and its scratch space as well where the artifact root is
|
|
239
|
+
a local one -- which is where a ``file://`` artifact a previous step wrote actually is.
|
|
240
|
+
"""
|
|
241
|
+
roots = [ctx.work]
|
|
242
|
+
split = urlsplit(ctx.scratch)
|
|
243
|
+
if split.scheme == "file":
|
|
244
|
+
roots.append(Path(f"{split.netloc}{split.path}").absolute())
|
|
245
|
+
return roots
|
|
246
|
+
|
|
247
|
+
|
|
248
|
+
def _open_s3(connection: Connection, ctx: StepContext) -> None:
|
|
249
|
+
"""Give this duckdb connection the ``s3`` scheme, on the instance's own credentials.
|
|
250
|
+
|
|
251
|
+
Run in the thread the connection belongs to, before any statement, so a ``read_parquet``
|
|
252
|
+
of an ``s3://`` object opens it rather than looking for a path.
|
|
253
|
+
"""
|
|
254
|
+
config = ctx.storage_connection(S3, S3StorageSettings)
|
|
255
|
+
if config is None:
|
|
256
|
+
raise BlockFailure(NO_STORAGE_CONNECTION, error_class=ErrorClass.REJECTED, scheme=S3)
|
|
257
|
+
try:
|
|
258
|
+
connection.execute(sqlalchemy.text(f"LOAD {HTTPFS}"))
|
|
259
|
+
except sqlalchemy.exc.DatabaseError as error:
|
|
260
|
+
raise BlockFailure(NO_HTTPFS, error_class=ErrorClass.REJECTED, extension=HTTPFS, scheme=S3) from error
|
|
261
|
+
for name, value in s3_options(config):
|
|
262
|
+
_set(connection, name, value)
|
|
263
|
+
# These statements autobegin a transaction, and the step opens its own straight after.
|
|
264
|
+
# They configure the session rather than touching data, so ending this one keeps them.
|
|
265
|
+
connection.commit()
|
|
266
|
+
|
|
267
|
+
|
|
268
|
+
def _set(connection: Connection, name: str, value: Any) -> None:
|
|
269
|
+
"""One duckdb setting, its value bound rather than spliced into the statement.
|
|
270
|
+
|
|
271
|
+
A credential must not become part of a statement, and a path must not be able to end one.
|
|
272
|
+
"""
|
|
273
|
+
connection.execute(sqlalchemy.text(f"SET {name} = :value"), {"value": value})
|
|
274
|
+
|
|
275
|
+
|
|
276
|
+
def _contain(connection: Connection, ctx: StepContext, url: URL, *, s3: bool) -> None:
|
|
277
|
+
"""Hold this duckdb connection to the run's own directories, whatever a statement names.
|
|
278
|
+
|
|
279
|
+
Run in the thread the connection belongs to, after the credentials and before any statement
|
|
280
|
+
the document supplies. The order is the whole mechanism: ``allowed_directories`` only bites
|
|
281
|
+
once ``enable_external_access`` is off, external access cannot be turned back on while the
|
|
282
|
+
database runs, and the lock refuses the ``SET`` that would widen the roots again. With
|
|
283
|
+
external access off duckdb also refuses to load or install any further extension, so the
|
|
284
|
+
scheme this session was opened with is the only one it has.
|
|
285
|
+
"""
|
|
286
|
+
roots = [str(root) for root in _readable(ctx)]
|
|
287
|
+
if s3:
|
|
288
|
+
# The scheme duckdb was just given credentials for stays reachable; no other does.
|
|
289
|
+
roots.append(f"{S3}://")
|
|
290
|
+
_set(connection, "allowed_directories", roots)
|
|
291
|
+
database = _database_path(url)
|
|
292
|
+
if database is not None:
|
|
293
|
+
# The file this connection names may sit anywhere, and duckdb writes a log beside it.
|
|
294
|
+
_set(connection, "allowed_paths", [str(database), f"{database}.wal"])
|
|
295
|
+
work = ctx.work
|
|
296
|
+
_set(connection, "temp_directory", str(work / TEMP_DIR))
|
|
297
|
+
_set(connection, "secret_directory", str(work / SECRET_DIR))
|
|
298
|
+
_set(connection, "enable_external_access", False)
|
|
299
|
+
_set(connection, "lock_configuration", True)
|
|
300
|
+
connection.commit()
|
|
301
|
+
|
|
302
|
+
|
|
303
|
+
def _database_path(url: URL) -> Path | None:
|
|
304
|
+
"""The file this duckdb url opens, or nothing where the database lives only in memory."""
|
|
305
|
+
if _is_memory(url) or not url.database:
|
|
306
|
+
return None
|
|
307
|
+
return Path(url.database).absolute()
|
|
308
|
+
|
|
309
|
+
|
|
310
|
+
def _refusal(error: Exception, roots: list[Path]) -> BlockFailure | None:
|
|
311
|
+
"""What a statement duckdb's containment refused failed for, or nothing for any other error."""
|
|
312
|
+
text = str(error).lower()
|
|
313
|
+
if DENIED_PATH in text:
|
|
314
|
+
named = ", ".join(str(root) for root in roots)
|
|
315
|
+
return BlockFailure(STATEMENT_OUTSIDE_THE_RUN, error_class=ErrorClass.REJECTED, named=named)
|
|
316
|
+
if DENIED_EXTENSION in text:
|
|
317
|
+
return BlockFailure(STATEMENT_LOADS_AN_EXTENSION, error_class=ErrorClass.REJECTED)
|
|
318
|
+
return None
|
|
319
|
+
|
|
320
|
+
|
|
321
|
+
class _Session:
|
|
322
|
+
"""One duckdb connection, opened in a worker thread and closed there too."""
|
|
323
|
+
|
|
324
|
+
def __init__(self, settings: SqlConnectionConfig, url: URL, ctx: StepContext, *, s3: bool = False) -> None:
|
|
325
|
+
"""Hold what opening this step's duckdb file needs."""
|
|
326
|
+
self.settings = settings
|
|
327
|
+
self.ctx = ctx
|
|
328
|
+
self.s3 = s3
|
|
329
|
+
self.url = url
|
|
330
|
+
self.engine = _engine(settings, url)
|
|
331
|
+
self.connection: _Threaded | None = None
|
|
332
|
+
|
|
333
|
+
async def __aenter__(self) -> "_Threaded":
|
|
334
|
+
"""Open the file, read-only where the connection said so, which duckdb enforces itself."""
|
|
335
|
+
try:
|
|
336
|
+
async with asyncio.timeout(self.settings.connect_timeout.total_seconds()):
|
|
337
|
+
opened = await asyncio.to_thread(self.engine.connect)
|
|
338
|
+
except TimeoutError as error:
|
|
339
|
+
await asyncio.to_thread(self.engine.dispose)
|
|
340
|
+
raise BlockFailure(
|
|
341
|
+
CONNECT_TIMED_OUT, error_class=ErrorClass.TRANSIENT, timeout=self.settings.connect_timeout
|
|
342
|
+
) from error
|
|
343
|
+
except BaseException:
|
|
344
|
+
await asyncio.to_thread(self.engine.dispose)
|
|
345
|
+
raise
|
|
346
|
+
try:
|
|
347
|
+
if self.s3:
|
|
348
|
+
await asyncio.to_thread(_open_s3, opened, self.ctx)
|
|
349
|
+
await asyncio.to_thread(_contain, opened, self.ctx, self.url, s3=self.s3)
|
|
350
|
+
except BaseException:
|
|
351
|
+
await asyncio.to_thread(opened.close)
|
|
352
|
+
await asyncio.to_thread(self.engine.dispose)
|
|
353
|
+
raise
|
|
354
|
+
self.connection = _Threaded(opened, _readable(self.ctx))
|
|
355
|
+
return self.connection
|
|
356
|
+
|
|
357
|
+
async def __aexit__(self, *_error: object) -> None:
|
|
358
|
+
"""Close the connection and the engine behind it, however the step left."""
|
|
359
|
+
if self.connection is not None:
|
|
360
|
+
await self.connection.close()
|
|
361
|
+
await asyncio.to_thread(self.engine.dispose)
|
|
362
|
+
|
|
363
|
+
|
|
364
|
+
class _Threaded:
|
|
365
|
+
"""A synchronous connection driven from the event loop, one call in one thread at a time.
|
|
366
|
+
|
|
367
|
+
A deadline that passes cancels the wait but not the thread, so a call cancelled here
|
|
368
|
+
interrupts duckdb, which is what ends the statement rather than leaving it running behind
|
|
369
|
+
a step that has already failed.
|
|
370
|
+
"""
|
|
371
|
+
|
|
372
|
+
def __init__(self, connection: Connection, roots: list[Path]) -> None:
|
|
373
|
+
"""Hold the connection, the handle an interrupt is sent through, and the roots it is held to."""
|
|
374
|
+
self.connection = connection
|
|
375
|
+
self.raw: Any = connection.connection.driver_connection
|
|
376
|
+
self.roots = roots
|
|
377
|
+
|
|
378
|
+
async def execute(self, statement: TextClause, parameters: dict[str, Any] | None = None, /) -> Any:
|
|
379
|
+
"""Run one statement and hand back its result, rowcount included."""
|
|
380
|
+
return await self._call(lambda: self.connection.execute(statement, parameters))
|
|
381
|
+
|
|
382
|
+
async def stream(self, statement: TextClause, parameters: dict[str, Any] | None = None, /) -> "_ThreadedRows":
|
|
383
|
+
"""Run one statement and hand back a cursor read a batch at a time."""
|
|
384
|
+
cursor = self.connection.execution_options(stream_results=True)
|
|
385
|
+
result = await self._call(lambda: cursor.execute(statement, parameters))
|
|
386
|
+
return _ThreadedRows(result, self._call)
|
|
387
|
+
|
|
388
|
+
@asynccontextmanager
|
|
389
|
+
async def begin(self) -> AsyncGenerator[None]:
|
|
390
|
+
"""One transaction, committed when the block leaves cleanly and rolled back otherwise."""
|
|
391
|
+
transaction = await self._call(self.connection.begin)
|
|
392
|
+
try:
|
|
393
|
+
yield
|
|
394
|
+
except BaseException:
|
|
395
|
+
await asyncio.to_thread(transaction.rollback)
|
|
396
|
+
raise
|
|
397
|
+
await asyncio.to_thread(transaction.commit)
|
|
398
|
+
|
|
399
|
+
async def close(self) -> None:
|
|
400
|
+
"""Close the connection in a thread, as everything else on it is done."""
|
|
401
|
+
await asyncio.to_thread(self.connection.close)
|
|
402
|
+
|
|
403
|
+
async def _call[T](self, work: Callable[[], T]) -> T:
|
|
404
|
+
"""Run one blocking call in a thread, interrupting duckdb when the wait is cancelled."""
|
|
405
|
+
try:
|
|
406
|
+
return await asyncio.to_thread(work)
|
|
407
|
+
except asyncio.CancelledError:
|
|
408
|
+
self.raw.interrupt()
|
|
409
|
+
raise
|
|
410
|
+
except sqlalchemy.exc.DatabaseError as error:
|
|
411
|
+
refusal = _refusal(error, self.roots)
|
|
412
|
+
if refusal is None:
|
|
413
|
+
raise
|
|
414
|
+
raise refusal from error
|
|
415
|
+
|
|
416
|
+
|
|
417
|
+
class _ThreadedRows:
|
|
418
|
+
"""The rows of one synchronous result, fetched a batch at a time in a worker thread."""
|
|
419
|
+
|
|
420
|
+
def __init__(self, result: Any, call: Callable[[Callable[[], Any]], Any]) -> None:
|
|
421
|
+
"""Hold the cursor and the thread call every fetch goes through."""
|
|
422
|
+
self.result = result
|
|
423
|
+
self.call = call
|
|
424
|
+
|
|
425
|
+
def keys(self) -> list[str]:
|
|
426
|
+
"""The column names, in the order the statement selected them."""
|
|
427
|
+
return list(self.result.keys())
|
|
428
|
+
|
|
429
|
+
async def partitions(self, size: int) -> AsyncIterator[Sequence[Any]]:
|
|
430
|
+
"""Read the cursor a batch at a time, the way the async result does."""
|
|
431
|
+
while True:
|
|
432
|
+
batch = await self.call(lambda: self.result.fetchmany(size))
|
|
433
|
+
if not batch:
|
|
434
|
+
return
|
|
435
|
+
yield batch
|
|
@@ -0,0 +1,54 @@
|
|
|
1
|
+
"""Every refusal the duckdb engine makes, catalogued under the ``sql.duckdb`` prefix."""
|
|
2
|
+
|
|
3
|
+
from dirigent_common import Catalogue
|
|
4
|
+
|
|
5
|
+
DUCKDB = Catalogue("sql.duckdb")
|
|
6
|
+
|
|
7
|
+
PARAMETER_NAMES_STORAGE = DUCKDB.define(
|
|
8
|
+
"parameter_names_storage",
|
|
9
|
+
"parameter {name} names {scheme}:// storage, and duckdb reads a file through the "
|
|
10
|
+
"worker's own filesystem here; copy it into the run's scratch space with storage.copy first",
|
|
11
|
+
)
|
|
12
|
+
|
|
13
|
+
PARAMETER_OUTSIDE_THE_RUN = DUCKDB.define(
|
|
14
|
+
"parameter_outside_the_run",
|
|
15
|
+
"parameter {name} names {uri}, which is outside this run's own directories ({named}); a "
|
|
16
|
+
"query reads and writes the files of the run it belongs to",
|
|
17
|
+
)
|
|
18
|
+
|
|
19
|
+
NO_STORAGE_CONNECTION = DUCKDB.define(
|
|
20
|
+
"no_storage_connection",
|
|
21
|
+
"this statement names {scheme}:// storage and no connection is bound to the {scheme} scheme, so "
|
|
22
|
+
"duckdb has no endpoint or credential to open it with; set DIRIGENT_STORAGE_CONNECTIONS",
|
|
23
|
+
)
|
|
24
|
+
|
|
25
|
+
NO_HTTPFS = DUCKDB.define(
|
|
26
|
+
"no_httpfs",
|
|
27
|
+
"duckdb could not load its {extension} extension, which is what reads {scheme}:// here; install "
|
|
28
|
+
"it once on this worker with duckdb -c 'INSTALL {extension}'",
|
|
29
|
+
)
|
|
30
|
+
|
|
31
|
+
STATEMENT_OUTSIDE_THE_RUN = DUCKDB.define(
|
|
32
|
+
"statement_outside_the_run",
|
|
33
|
+
"this statement names a file outside the run's own directories ({named}), which is all a "
|
|
34
|
+
"statement reads and writes here; copy it into the run's scratch space with storage.copy first",
|
|
35
|
+
)
|
|
36
|
+
|
|
37
|
+
STATEMENT_LOADS_AN_EXTENSION = DUCKDB.define(
|
|
38
|
+
"statement_loads_an_extension",
|
|
39
|
+
"this statement loads a duckdb extension, and a session carries only the extensions it "
|
|
40
|
+
"was opened with; a bucket a statement names is opened through its storage connection",
|
|
41
|
+
)
|
|
42
|
+
|
|
43
|
+
CONNECT_TIMED_OUT = DUCKDB.define("connect_timed_out", "the database did not answer within {timeout}")
|
|
44
|
+
|
|
45
|
+
|
|
46
|
+
# What a config refuses at validation. Pydantic owns the code a validator's refusal reaches
|
|
47
|
+
# the wire under, so these are rendered into the ``ValueError`` it wraps.
|
|
48
|
+
|
|
49
|
+
READ_ONLY_IN_MEMORY = DUCKDB.define(
|
|
50
|
+
"read_only_in_memory",
|
|
51
|
+
"read_only has no meaning on duckdb:///:memory:, which duckdb refuses to open at all: "
|
|
52
|
+
"an in-memory database starts empty and a read-only one can never be filled, so the "
|
|
53
|
+
"connection would open on nothing; name a duckdb file, or drop read_only",
|
|
54
|
+
)
|
|
File without changes
|