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.
@@ -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
+ )