dirigent-block-sql 0.17.0__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,28 @@
1
+ Metadata-Version: 2.4
2
+ Name: dirigent-block-sql
3
+ Version: 0.17.0
4
+ Summary: The SQL block family for dirigent: sql.query, sql.execute, and the sql connection kind.
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-common==0.17.0
10
+ Requires-Dist: dirigent-plugin==0.17.0
11
+ Requires-Dist: sqlalchemy[asyncio]>=2.0.52
12
+ Requires-Python: >=3.13
13
+ Description-Content-Type: text/markdown
14
+
15
+ # dirigent-block-sql
16
+
17
+ The SQL block family: `sql.query` reads rows and `sql.execute` writes them, both against the
18
+ database a `sql` connection names.
19
+
20
+ The connection kind holds the URL and its credential, and the driver is SQLAlchemy's, so any
21
+ dialect with an async driver the worker has installed is addressable.
22
+
23
+ Engines are packages. A backend that needs more than a driver implements `SqlEngine` and
24
+ registers under the `dirigent.sql.engines.v1` entry-point group, and the family asks it what a
25
+ valid connection to that backend is, where a database written as a relative path lands, what a
26
+ parameter becomes, how a check reaches it, and what a session on it is.
27
+ `dirigent-block-duckdb` is the first, and `duckdb:///warehouse.duckdb` works once it is
28
+ installed.
@@ -0,0 +1,14 @@
1
+ # dirigent-block-sql
2
+
3
+ The SQL block family: `sql.query` reads rows and `sql.execute` writes them, both against the
4
+ database a `sql` connection names.
5
+
6
+ The connection kind holds the URL and its credential, and the driver is SQLAlchemy's, so any
7
+ dialect with an async driver the worker has installed is addressable.
8
+
9
+ Engines are packages. A backend that needs more than a driver implements `SqlEngine` and
10
+ registers under the `dirigent.sql.engines.v1` entry-point group, and the family asks it what a
11
+ valid connection to that backend is, where a database written as a relative path lands, what a
12
+ parameter becomes, how a check reaches it, and what a session on it is.
13
+ `dirigent-block-duckdb` is the first, and `duckdb:///warehouse.duckdb` works once it is
14
+ installed.
@@ -0,0 +1,30 @@
1
+ [project]
2
+ name = "dirigent-block-sql"
3
+ version = "0.17.0"
4
+ description = "The SQL block family for dirigent: sql.query, sql.execute, and the sql connection kind."
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-common==0.17.0",
15
+ "dirigent-plugin==0.17.0",
16
+ "sqlalchemy[asyncio]>=2.0.52",
17
+ ]
18
+
19
+ [project.entry-points."dirigent.plugins.v1"]
20
+ block-sql = "dirigent_block_sql:plugin"
21
+
22
+ [build-system]
23
+ requires = ["uv_build>=0.12.0,<0.13.0"]
24
+ build-backend = "uv_build"
25
+
26
+ [tool.uv.sources.dirigent-common]
27
+ workspace = true
28
+
29
+ [tool.uv.sources.dirigent-plugin]
30
+ workspace = true
@@ -0,0 +1,28 @@
1
+ [project]
2
+ name = "dirigent-block-sql"
3
+ version = "0.17.0"
4
+ description = "The SQL block family for dirigent: sql.query, sql.execute, and the sql connection kind."
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-common==0.17.0",
15
+ "dirigent-plugin==0.17.0",
16
+ "sqlalchemy[asyncio]>=2.0.52",
17
+ ]
18
+
19
+ [project.entry-points."dirigent.plugins.v1"]
20
+ block-sql = "dirigent_block_sql:plugin"
21
+
22
+ [build-system]
23
+ requires = ["uv_build>=0.12.0,<0.13.0"]
24
+ build-backend = "uv_build"
25
+
26
+ [tool.uv.sources]
27
+ dirigent-common = { workspace = true }
28
+ dirigent-plugin = { workspace = true }
@@ -0,0 +1,34 @@
1
+ """The SQL block family: reading and writing a database through one connection kind."""
2
+
3
+ from dirigent_block_sql.engines import SqlAlchemyEngine, SqlEngine, SqlSession
4
+ from dirigent_block_sql.markers import ENGINES_GROUP
5
+ from dirigent_block_sql.sql import SqlConnectionConfig, SqlConnectionKind, SqlExecuteOperator, SqlQueryOperator
6
+ from dirigent_plugin import Contribution, extension
7
+
8
+
9
+ class SqlBlocks:
10
+ """The plugin object the host discovers under the dirigent.plugins.v1 entry-point group."""
11
+
12
+ @extension
13
+ def contribute(self) -> Contribution:
14
+ """Contribute the two SQL blocks and the connection kind that addresses a database."""
15
+ return Contribution(
16
+ operators=[SqlQueryOperator(), SqlExecuteOperator()],
17
+ connection_kinds=[SqlConnectionKind()],
18
+ )
19
+
20
+
21
+ plugin = SqlBlocks()
22
+
23
+ __all__ = [
24
+ "ENGINES_GROUP",
25
+ "SqlAlchemyEngine",
26
+ "SqlBlocks",
27
+ "SqlConnectionConfig",
28
+ "SqlConnectionKind",
29
+ "SqlEngine",
30
+ "SqlExecuteOperator",
31
+ "SqlQueryOperator",
32
+ "SqlSession",
33
+ "plugin",
34
+ ]
@@ -0,0 +1,321 @@
1
+ """The engine contract the ``sql`` family drives a database through, and the registry of engines.
2
+
3
+ An engine is one backend -- what a URL's scheme names -- and the decisions that differ from one
4
+ backend to the next: what makes a connection to it valid, where a database written as a relative
5
+ path lands, what a parameter has to become before the database sees it, how a check reaches it,
6
+ and what a session on it is. Everything else about ``sql.query`` and ``sql.execute`` is the same
7
+ whichever engine answers.
8
+
9
+ Any dialect with an async driver is the generic path here, and a backend that needs more than a
10
+ driver -- DuckDB, which has no async driver at all -- is a package contributing an engine under
11
+ the ``dirigent.sql.engines.v1`` entry-point group. A block cannot reach the plugin host, so the
12
+ family loads that group itself, once per process.
13
+ """
14
+
15
+ import asyncio
16
+ from abc import ABC, abstractmethod
17
+ from collections.abc import Awaitable, Mapping, Sequence
18
+ from contextlib import AbstractAsyncContextManager
19
+ from functools import cache
20
+ from typing import TYPE_CHECKING, Any, ClassVar, Final, Protocol
21
+
22
+ import sqlalchemy
23
+ import sqlalchemy.exc
24
+ from pluginkit import PluginManager
25
+ from sqlalchemy.engine import URL
26
+ from sqlalchemy.ext.asyncio import AsyncConnection, AsyncEngine, create_async_engine
27
+ from sqlalchemy.pool import NullPool
28
+ from sqlalchemy.sql.elements import TextClause
29
+
30
+ from dirigent_block_sql import markers
31
+ from dirigent_block_sql.markers import ENGINES_GROUP
32
+ from dirigent_block_sql.messages import (
33
+ CONNECT_TIMED_OUT,
34
+ DRIVER_NOT_INSTALLED,
35
+ ENGINE_PACKAGE_MISSING,
36
+ NO_ASYNC_DRIVER,
37
+ READ_ONLY_UNSUPPORTED,
38
+ )
39
+ from dirigent_common import HealthReport, JsonMap
40
+ from dirigent_plugin import PROJECT_NAME, BlockFailure, ErrorClass, StepContext
41
+
42
+ if TYPE_CHECKING:
43
+ from dirigent_block_sql.sql import SqlConnectionConfig
44
+
45
+ #: How long a connection check may spend reaching the database and asking it for a one.
46
+ CHECK_TIMEOUT_SECONDS = 30.0
47
+
48
+ #: What a check answers for a database file that exists only while a run is on the worker.
49
+ OUTSIDE_A_RUN: Final = (
50
+ "a database file relative to a run's work directory exists only inside a run, so there is nothing here to check"
51
+ )
52
+
53
+ #: The dialects whose sessions can be made read-only, and how each one is told.
54
+ READ_ONLY: Final = {
55
+ "postgresql": "SET TRANSACTION READ ONLY",
56
+ "sqlite": "PRAGMA query_only = 1",
57
+ }
58
+
59
+ #: The distribution that provides each async driver a URL can name, for the install line an
60
+ #: absent one is refused with.
61
+ DRIVER_PACKAGE: Final = {
62
+ "aiosqlite": "aiosqlite",
63
+ "asyncpg": "asyncpg",
64
+ "psycopg": "psycopg[binary]",
65
+ "aiomysql": "aiomysql",
66
+ "asyncmy": "asyncmy",
67
+ "aioodbc": "aioodbc",
68
+ "aiooracle": "aiooracle",
69
+ "oracledb": "oracledb",
70
+ }
71
+
72
+ #: The package each backend that is an engine rather than a driver is carried by, for the
73
+ #: install line a URL naming one is refused with when it is not installed.
74
+ KNOWN_ENGINE_PACKAGES: Final = {"duckdb": "dirigent-block-duckdb"}
75
+
76
+
77
+ class DuplicateEngine(Exception):
78
+ """Two plugins claimed the same backend."""
79
+
80
+ def __init__(self, backend: str, first: str, second: str) -> None:
81
+ """Name the backend and both plugins."""
82
+ super().__init__(
83
+ f"backend {backend!r} is contributed by both {first!r} and {second!r}, and a backend is what a "
84
+ f"connection's url dispatches on, so one of the two packages must be uninstalled."
85
+ )
86
+ self.backend = backend
87
+ self.plugins = (first, second)
88
+
89
+
90
+ class SqlSession(Protocol):
91
+ """The connection surface both blocks drive, whichever engine opened it."""
92
+
93
+ async def execute(self, statement: TextClause, parameters: dict[str, Any] | None = None, /) -> Any:
94
+ """Run one statement and hand back its result, rowcount included."""
95
+ ...
96
+
97
+ def stream(self, statement: TextClause, parameters: dict[str, Any] | None = None, /) -> Awaitable[Any]:
98
+ """Run one statement and hand back a cursor read a batch at a time.
99
+
100
+ Awaited rather than declared ``async`` because SQLAlchemy's own is a context manager
101
+ that is also awaitable, and a step awaits it.
102
+ """
103
+ ...
104
+
105
+ def begin(self) -> AbstractAsyncContextManager[Any]:
106
+ """One transaction, committed when the block leaves cleanly and rolled back otherwise."""
107
+ ...
108
+
109
+
110
+ class SqlEngine(ABC):
111
+ """One backend of the ``sql`` family: what a URL naming it means, and how a step drives it."""
112
+
113
+ backend: ClassVar[str]
114
+ """The backend a URL names this engine by, as ``URL.get_backend_name()`` spells it."""
115
+
116
+ @abstractmethod
117
+ def validate(self, settings: "SqlConnectionConfig", url: URL) -> None:
118
+ """Refuse a connection this engine cannot open, saying what about it is refused."""
119
+
120
+ def resolve(self, url: URL, ctx: StepContext) -> URL:
121
+ """The URL a step connects with, for an engine whose database may be a file of the run."""
122
+ return url
123
+
124
+ def bind(self, params: JsonMap, ctx: StepContext) -> dict[str, Any]:
125
+ """The parameters as this engine receives them."""
126
+ return dict(params)
127
+
128
+ @abstractmethod
129
+ async def check(self, settings: "SqlConnectionConfig", url: URL) -> HealthReport:
130
+ """Reach the database and ask it for a one, which is the smallest proof of reach and credential."""
131
+
132
+ @abstractmethod
133
+ def session(
134
+ self,
135
+ settings: "SqlConnectionConfig",
136
+ url: URL,
137
+ ctx: StepContext,
138
+ *,
139
+ statements: Sequence[str],
140
+ params: JsonMap,
141
+ ) -> AbstractAsyncContextManager[SqlSession]:
142
+ """The session a step runs its statements in, opened on entry and closed on exit.
143
+
144
+ The statements and the bound parameters are handed over whole, because what a session
145
+ has to be given -- a scheme's credentials, an extension -- is read off what will run in
146
+ it.
147
+ """
148
+
149
+
150
+ class SqlAlchemyEngine(SqlEngine):
151
+ """Any dialect with an async driver, driven through SQLAlchemy's own asyncio layer."""
152
+
153
+ #: It answers for every backend no installed engine claims, so it is registered under none.
154
+ backend: ClassVar[str] = ""
155
+
156
+ def validate(self, settings: "SqlConnectionConfig", url: URL) -> None:
157
+ """Refuse a URL naming no async driver, and ``read_only`` on a dialect that has no mode for it."""
158
+ backend = url.get_backend_name()
159
+ if "+" not in url.drivername:
160
+ package = KNOWN_ENGINE_PACKAGES.get(backend)
161
+ if package is not None:
162
+ raise ValueError(ENGINE_PACKAGE_MISSING.render(backend=backend, package=package))
163
+ raise ValueError(NO_ASYNC_DRIVER.render(driver=repr(url.drivername), driver_name=url.drivername))
164
+ if settings.read_only and backend not in READ_ONLY:
165
+ supported = ", ".join(sorted(READ_ONLY))
166
+ raise ValueError(READ_ONLY_UNSUPPORTED.render(backend=backend, supported=supported))
167
+
168
+ def resolve(self, url: URL, ctx: StepContext) -> URL:
169
+ """A sqlite database written as a relative path is a file in the run's work directory."""
170
+ return under_work(url, ctx) if _is_run_relative(url) else url
171
+
172
+ async def check(self, settings: "SqlConnectionConfig", url: URL) -> HealthReport:
173
+ """Open a connection and run ``SELECT 1`` over the driver the URL names."""
174
+ if _is_run_relative(url):
175
+ return HealthReport(healthy=False, detail=OUTSIDE_A_RUN)
176
+ try:
177
+ engine = _engine(url)
178
+ except BlockFailure as error:
179
+ return HealthReport(healthy=False, detail=str(error))
180
+ try:
181
+ async with asyncio.timeout(CHECK_TIMEOUT_SECONDS), engine.connect() as connection:
182
+ await connection.execute(sqlalchemy.text("SELECT 1"))
183
+ except (TimeoutError, sqlalchemy.exc.SQLAlchemyError, OSError) as error:
184
+ return HealthReport(healthy=False, detail=reason(error))
185
+ finally:
186
+ await engine.dispose()
187
+ return HealthReport(healthy=True, detail=f"{url.get_backend_name()} answered")
188
+
189
+ def session(
190
+ self,
191
+ settings: "SqlConnectionConfig",
192
+ url: URL,
193
+ ctx: StepContext,
194
+ *,
195
+ statements: Sequence[str],
196
+ params: JsonMap,
197
+ ) -> AbstractAsyncContextManager[SqlSession]:
198
+ """One async connection, put in the mode the connection kind asked for."""
199
+ return _Session(settings, url)
200
+
201
+
202
+ #: The engine a backend no installed package claims is driven through.
203
+ GENERIC: Final = SqlAlchemyEngine()
204
+
205
+
206
+ def under_work(url: URL, ctx: StepContext) -> URL:
207
+ """The same URL with its database file placed under the run's own work directory."""
208
+ database = ctx.work / str(url.database)
209
+ database.parent.mkdir(parents=True, exist_ok=True)
210
+ return url.set(database=str(database))
211
+
212
+
213
+ def reason(error: Exception) -> str:
214
+ """The one sentence a health report gives for a database that did not answer."""
215
+ if isinstance(error, TimeoutError):
216
+ return f"no answer within {CHECK_TIMEOUT_SECONDS:.0f}s"
217
+ text = str(getattr(error, "orig", None) or error).strip().splitlines()
218
+ return text[0] if text else type(error).__name__
219
+
220
+
221
+ def registry(
222
+ *,
223
+ group: str = ENGINES_GROUP,
224
+ extra: Mapping[str, object] | None = None,
225
+ ) -> dict[str, SqlEngine]:
226
+ """Build the registry: the engine every installed plugin contributes, keyed by backend.
227
+
228
+ ``extra`` registers plugin objects that are not installed as distributions.
229
+ """
230
+ found: dict[str, SqlEngine] = {}
231
+ origins: dict[str, str] = {}
232
+ manager = PluginManager(PROJECT_NAME)
233
+ manager.add_extension_points(markers)
234
+ manager.load_entrypoints(group)
235
+ for name, plugin in (extra or {}).items():
236
+ manager.register(plugin, name=name)
237
+ # The hook's return annotation is a declaration, not an enforcement: a plugin may answer
238
+ # with anything, and anything that is not an engine is not registered.
239
+ for plugin_name, contributed in manager.caller(markers.engines).collect_with_plugins():
240
+ for engine in contributed:
241
+ if not isinstance(engine, SqlEngine): # pyright: ignore[reportUnnecessaryIsInstance]
242
+ continue
243
+ owner = origins.get(engine.backend)
244
+ if owner is not None:
245
+ raise DuplicateEngine(engine.backend, owner, plugin_name)
246
+ origins[engine.backend] = plugin_name
247
+ found[engine.backend] = engine
248
+ return found
249
+
250
+
251
+ @cache
252
+ def _installed() -> dict[str, SqlEngine]:
253
+ """The engines this process found, scanned once: a step must not rescan for every statement."""
254
+ return registry()
255
+
256
+
257
+ def engine_for(url: URL) -> SqlEngine:
258
+ """The engine registered for this URL's backend, or the generic async SQLAlchemy one."""
259
+ return _installed().get(url.get_backend_name(), GENERIC)
260
+
261
+
262
+ def _is_run_relative(url: URL) -> bool:
263
+ """Whether this is a sqlite file named by a path relative to the run's work directory."""
264
+ database = url.database or ""
265
+ if url.get_backend_name() != "sqlite":
266
+ return False
267
+ return bool(database) and not database.startswith("/")
268
+
269
+
270
+ def _engine(url: URL) -> AsyncEngine:
271
+ """Build the engine one step uses, naming the package to install when the driver is absent.
272
+
273
+ ``NullPool`` because a step opens one connection and the worker process may run the next
274
+ step against a different database entirely; a pool would outlive what it serves.
275
+ """
276
+ try:
277
+ return create_async_engine(url, poolclass=NullPool)
278
+ except (ModuleNotFoundError, sqlalchemy.exc.NoSuchModuleError) as error:
279
+ driver = url.drivername.partition("+")[2]
280
+ package = DRIVER_PACKAGE.get(driver, driver)
281
+ raise BlockFailure(
282
+ DRIVER_NOT_INSTALLED,
283
+ error_class=ErrorClass.REJECTED,
284
+ driver=repr(driver),
285
+ package=repr(package),
286
+ shipped=", ".join(["asyncpg", "aiosqlite"]),
287
+ ) from error
288
+
289
+
290
+ class _Session:
291
+ """One connection, opened under the connection's timeout and read-only where it says so."""
292
+
293
+ def __init__(self, settings: "SqlConnectionConfig", url: URL) -> None:
294
+ """Hold what opening this step's connection needs."""
295
+ self.settings = settings
296
+ self.url = url
297
+ self.engine = _engine(url)
298
+ self.connection: AsyncConnection | None = None
299
+
300
+ async def __aenter__(self) -> AsyncConnection:
301
+ """Open the connection and put it in the mode the connection kind asked for."""
302
+ try:
303
+ async with asyncio.timeout(self.settings.connect_timeout.total_seconds()):
304
+ self.connection = await self.engine.connect()
305
+ except TimeoutError as error:
306
+ await self.engine.dispose()
307
+ raise BlockFailure(
308
+ CONNECT_TIMED_OUT, error_class=ErrorClass.TRANSIENT, timeout=self.settings.connect_timeout
309
+ ) from error
310
+ except BaseException:
311
+ await self.engine.dispose()
312
+ raise
313
+ if self.settings.read_only:
314
+ await self.connection.execute(sqlalchemy.text(READ_ONLY[self.url.get_backend_name()]))
315
+ return self.connection
316
+
317
+ async def __aexit__(self, *_error: object) -> None:
318
+ """Close the connection and the engine behind it, however the step left."""
319
+ if self.connection is not None:
320
+ await self.connection.close()
321
+ await self.engine.dispose()
@@ -0,0 +1,18 @@
1
+ """The extension point an engine of the ``sql`` family is contributed through."""
2
+
3
+ from collections.abc import Sequence
4
+ from typing import TYPE_CHECKING, Final
5
+
6
+ from dirigent_plugin import extension_point
7
+
8
+ if TYPE_CHECKING:
9
+ from dirigent_block_sql.engines import SqlEngine
10
+
11
+ #: The version is part of the group name, so an incompatible contract ships as a new group.
12
+ ENGINES_GROUP: Final = "dirigent.sql.engines.v1"
13
+
14
+
15
+ @extension_point
16
+ def engines() -> "Sequence[SqlEngine]":
17
+ """Collect the engines a plugin adds to the ``sql`` family, once per process."""
18
+ raise NotImplementedError("an extension point is a declaration; call it via PluginManager.caller(...)")
@@ -0,0 +1,65 @@
1
+ """Every refusal the sql family makes, catalogued under the ``sql`` prefix."""
2
+
3
+ from dirigent_common import Catalogue
4
+
5
+ SQL = Catalogue("sql")
6
+
7
+ DRIVER_NOT_INSTALLED = SQL.define(
8
+ "driver_not_installed",
9
+ "the {driver} driver this url names is not installed on the worker; add the {package} "
10
+ "package to the image, or use a driver that ships with it ({shipped})",
11
+ )
12
+
13
+ CONNECT_TIMED_OUT = SQL.define("connect_timed_out", "the database did not answer within {timeout}")
14
+
15
+ READ_ONLY_CONNECTION = SQL.define(
16
+ "read_only_connection",
17
+ "connection {connection} is read_only, and sql.execute writes; "
18
+ "read it with sql.query, or point this step at a connection that may write",
19
+ )
20
+
21
+ TOO_MANY_ROWS = SQL.define(
22
+ "too_many_rows",
23
+ "the query returned more than max_rows ({maximum}) rows; raise max_rows, or narrow the query",
24
+ )
25
+
26
+ NO_JSON_SPELLING = SQL.define("no_json_spelling", "a column of this result has no JSON spelling: {detail}")
27
+
28
+
29
+ # What a config refuses at validation. Pydantic owns the code a validator's refusal reaches
30
+ # the wire under, so these are rendered into the ``ValueError`` it wraps.
31
+
32
+ ENGINE_PACKAGE_MISSING = SQL.define(
33
+ "engine_package_missing",
34
+ "{backend} needs the engine package: uv pip install {package}",
35
+ )
36
+
37
+ NO_ASYNC_DRIVER = SQL.define(
38
+ "no_async_driver",
39
+ "{driver} names no driver, and these blocks speak to a database over an "
40
+ "async one; write the driver in the url, as in "
41
+ "{driver_name}+asyncpg:// or {driver_name}+aiosqlite://",
42
+ )
43
+
44
+ READ_ONLY_UNSUPPORTED = SQL.define(
45
+ "read_only_unsupported",
46
+ "read_only has no meaning on {backend}: only {supported} can be told to "
47
+ "refuse writes for the length of a session, and a connection that cannot be "
48
+ "is not marked as one that is",
49
+ )
50
+
51
+ INLINE_PASSWORD = SQL.define(
52
+ "inline_password",
53
+ "this url carries a password inline, where it would sit unencrypted in a plain "
54
+ "field; take it out of the url and set the sealed password field instead",
55
+ )
56
+
57
+ EMPTY_STATEMENT = SQL.define("empty_statement", "statement {index} is empty")
58
+
59
+ NOT_A_DATABASE_URL = SQL.define("not_a_database_url", "{url} is not a database url: {detail}")
60
+
61
+ MORE_THAN_ONE_STATEMENT = SQL.define(
62
+ "more_than_one_statement",
63
+ "this is more than one statement: a ';' ends the first and there is more after it. "
64
+ "sql.query runs one statement, and sql.execute takes a list, one statement per entry",
65
+ )
@@ -0,0 +1,480 @@
1
+ """``sql.query`` and ``sql.execute``: read from and write to a database over a ``sql`` connection.
2
+
3
+ Reading a table and writing a table are the two most common things a pipeline does, and until
4
+ these blocks the answer was ``shell.run`` with ``psql`` -- an unsafe block, a credential on a
5
+ command line, and a result that arrives as text somebody has to parse.
6
+
7
+ Both blocks are **ordinary**. They run no command a document supplies and reach nothing but the
8
+ database their connection names, which is a narrower grant than ``shell.run``.
9
+
10
+ Nothing a document writes ever reaches the SQL text. A statement is a constant in the document
11
+ and every value is a named bind parameter, so ``${...}`` resolves into ``params`` and a value
12
+ that looks like SQL stays a value.
13
+
14
+ The engine is whatever the connection's URL names: any dialect with an async driver is driven
15
+ from here, and a backend needing more than a driver is an installed package contributing a
16
+ :class:`~dirigent_block_sql.engines.SqlEngine`, which every decision that differs by backend is
17
+ asked of.
18
+ """
19
+
20
+ import asyncio
21
+ import time
22
+ from collections.abc import AsyncIterator, Sequence
23
+ from datetime import timedelta
24
+ from typing import Annotated, Any, ClassVar, Final
25
+
26
+ import sqlalchemy
27
+ import sqlalchemy.exc
28
+ from pydantic import BaseModel, Field, SecretStr, model_validator
29
+ from sqlalchemy.engine import URL, make_url
30
+
31
+ from dirigent_block_sql.engines import SqlSession, engine_for
32
+ from dirigent_block_sql.messages import (
33
+ EMPTY_STATEMENT,
34
+ INLINE_PASSWORD,
35
+ MORE_THAN_ONE_STATEMENT,
36
+ NO_JSON_SPELLING,
37
+ NOT_A_DATABASE_URL,
38
+ READ_ONLY_CONNECTION,
39
+ TOO_MANY_ROWS,
40
+ )
41
+ from dirigent_common import (
42
+ SQL_MEDIA_TYPE,
43
+ BlockModel,
44
+ Duration,
45
+ HealthReport,
46
+ JsonList,
47
+ JsonMap,
48
+ spelled,
49
+ )
50
+ from dirigent_plugin import (
51
+ BlockFailure,
52
+ ConnectionKind,
53
+ ConnectionRef,
54
+ ErrorClass,
55
+ Operator,
56
+ OperatorSpec,
57
+ RemoteHandle,
58
+ StepContext,
59
+ )
60
+
61
+ #: How many rows are read out of the cursor at a time, inline or on the way to storage.
62
+ BATCH = 1000
63
+
64
+ #: The dialects that can be given a per-statement deadline the server itself enforces.
65
+ STATEMENT_TIMEOUT: Final = {"postgresql": "SET LOCAL statement_timeout = {milliseconds}"}
66
+
67
+ #: What a driver says when the database could not be reached, rather than refusing the request.
68
+ UNREACHABLE: Final = (
69
+ "connection refused",
70
+ "could not connect",
71
+ "could not translate host name",
72
+ "name or service not known",
73
+ "connection reset",
74
+ "server closed the connection",
75
+ "timeout expired",
76
+ "too many clients",
77
+ "the database system is starting up",
78
+ "database is locked",
79
+ )
80
+
81
+
82
+ class SqlConnectionConfig(BlockModel):
83
+ """One database, the sealed password that opens it, and whether it may be written to."""
84
+
85
+ url: str = Field(min_length=1)
86
+ """The database as a SQLAlchemy URL, driver included:
87
+ ``postgresql+asyncpg://user@host:5432/db`` or ``sqlite+aiosqlite:///data.db``.
88
+
89
+ A backend carried by an engine package writes the driver differently or not at all, and
90
+ that engine is what says so; a URL naming one that is not installed is refused with the
91
+ package to install.
92
+
93
+ A URL carrying a password inline is refused: the secret belongs in ``password``, where it
94
+ is encrypted at rest and redacted in every API response, and a plain field is neither.
95
+
96
+ A sqlite database written as a relative path is relative to the run's work directory,
97
+ which is where a database a pipeline builds for itself belongs. It is local to the worker
98
+ that made it, so a later step reading it must be on the same worker."""
99
+
100
+ password: SecretStr | None = None
101
+ """The password, sealed, merged into the URL when a connection is opened and nowhere else."""
102
+
103
+ read_only: bool = False
104
+ """Refuse to write through this connection.
105
+
106
+ Every session it opens is put in the dialect's own read-only mode, so a statement that
107
+ writes is refused by the database rather than by a check here, and ``sql.execute`` refuses
108
+ the connection outright. A dialect with no read-only mode is refused rather than silently
109
+ left writable."""
110
+
111
+ connect_timeout: Duration = timedelta(seconds=10)
112
+ """How long opening a connection may take before the step fails as a transient error."""
113
+
114
+ @model_validator(mode="after")
115
+ def _check_shape(self) -> "SqlConnectionConfig":
116
+ """Refuse a URL that is unparseable, carrying its own password, or one no engine opens."""
117
+ url = _parse(self.url)
118
+ if url.password is not None:
119
+ raise ValueError(INLINE_PASSWORD.render())
120
+ engine_for(url).validate(self, url)
121
+ return self
122
+
123
+
124
+ class SqlConnectionKind(ConnectionKind):
125
+ """The connection kind the ``sql.*`` blocks resolve their database through."""
126
+
127
+ id: ClassVar[str] = "sql"
128
+ config_model: ClassVar[type[BaseModel]] = SqlConnectionConfig
129
+
130
+ async def check(self, config: BaseModel) -> HealthReport:
131
+ """Ask the URL's own engine to reach the database, which is the smallest proof of reach."""
132
+ settings = SqlConnectionConfig.model_validate(config.model_dump())
133
+ url = _with_password(settings)
134
+ return await engine_for(url).check(settings, url)
135
+
136
+
137
+ class SqlQueryConfig(BlockModel):
138
+ """One statement, and the values bound into it."""
139
+
140
+ connection: ConnectionRef
141
+ """The ``sql`` connection naming the database and holding its password."""
142
+
143
+ sql: Annotated[str, Field(min_length=1, json_schema_extra={"contentMediaType": SQL_MEDIA_TYPE})]
144
+ """One statement, and one only. A document that needs two writes two steps, or uses
145
+ ``sql.execute``, which is the block that runs several as one transaction."""
146
+
147
+ params: JsonMap = Field(default_factory=dict[str, Any])
148
+ """Values bound by name, written ``:name`` in the statement.
149
+
150
+ This is where a ``${...}`` reference belongs. A parameter is sent to the database beside
151
+ the statement and never spliced into it, so a value that reads as SQL is still a value.
152
+ Nothing in ``sql`` is substituted, which also means a table or column name cannot come
153
+ from a parameter: only a value can.
154
+
155
+ An engine whose tables can be files is handed a storage URI as the file it opens, within
156
+ the run's own directories; every other engine is handed the value as it was written."""
157
+
158
+ max_rows: int = Field(default=1000, ge=1)
159
+ """How many rows may be carried inline in the step's output.
160
+
161
+ A result past this fails the step rather than being truncated: half an answer is not a
162
+ smaller answer, and a step acting on it would be acting on something the database never
163
+ said. Rows that belong in a file are handed to ``storage.write``."""
164
+
165
+ timeout: Duration = timedelta(minutes=5)
166
+ """How long the statement may run.
167
+
168
+ Set as the database's own statement timeout where the dialect has one, and enforced here
169
+ in every case, so a query that hangs ends the step rather than the deadline."""
170
+
171
+ @model_validator(mode="after")
172
+ def _check_shape(self) -> "SqlQueryConfig":
173
+ """Refuse a config whose ``sql`` is more than the one statement this block runs."""
174
+ _check_single(self.sql)
175
+ return self
176
+
177
+
178
+ class SqlQueryOutput(BlockModel):
179
+ """What one query returned."""
180
+
181
+ rows: JsonList
182
+ """The rows as objects keyed by column name, which a later step reads or writes out."""
183
+
184
+ row_count: int
185
+ """How many rows the query returned."""
186
+
187
+ columns: list[str]
188
+ """The column names, in the order the query selected them."""
189
+
190
+ duration_ms: int
191
+
192
+
193
+ class SqlQueryOperator(Operator[SqlQueryConfig, SqlQueryOutput]):
194
+ """Runs one statement against a database and reports its rows."""
195
+
196
+ spec = OperatorSpec(
197
+ id="sql.query",
198
+ summary="Run one SQL statement and return its rows.",
199
+ idempotent=True,
200
+ )
201
+ config_model: ClassVar[type[BaseModel]] = SqlQueryConfig
202
+ output_model: ClassVar[type[BaseModel]] = SqlQueryOutput
203
+
204
+ async def execute(self, config: SqlQueryConfig, ctx: StepContext) -> SqlQueryOutput | RemoteHandle:
205
+ """Open a session, run the statement, and hand the rows on as the step's output."""
206
+ settings = ctx.connection(config.connection, SqlConnectionConfig)
207
+ url = _resolved(settings, ctx)
208
+ engine = engine_for(url)
209
+ params = engine.bind(config.params, ctx)
210
+ started = time.monotonic()
211
+ async with engine.session(settings, url, ctx, statements=[config.sql], params=params) as session:
212
+ await _limit(session, settings, config.timeout)
213
+ # The deadline covers running the statement and reading the rows out of it, because
214
+ # a query that hangs hangs in either.
215
+ async with asyncio.timeout(config.timeout.total_seconds()):
216
+ result = await session.stream(sqlalchemy.text(config.sql), params)
217
+ columns = list(result.keys())
218
+ rows = await _inline(result, config.max_rows)
219
+ duration = round((time.monotonic() - started) * 1000)
220
+ ctx.log.info(
221
+ "query ran",
222
+ connection=config.connection,
223
+ row_count=len(rows),
224
+ duration_ms=duration,
225
+ )
226
+ return SqlQueryOutput(
227
+ rows=rows,
228
+ row_count=len(rows),
229
+ columns=columns,
230
+ duration_ms=duration,
231
+ )
232
+
233
+ def classify_error(self, error: Exception) -> ErrorClass:
234
+ """A database that could not be reached is transient; one that refused the statement is not."""
235
+ return classify(error)
236
+
237
+
238
+ class SqlExecuteConfig(BlockModel):
239
+ """The statements one transaction runs, and the values bound into all of them."""
240
+
241
+ connection: ConnectionRef
242
+ """The ``sql`` connection naming the database and holding its password. A connection
243
+ marked ``read_only`` is refused: this block writes."""
244
+
245
+ statements: list[str] = Field(min_length=1)
246
+ """The statements, run in order inside one transaction. They all commit or none of them
247
+ does, so a migration, an insert and the index it needs are one step and not three."""
248
+
249
+ params: JsonMap = Field(default_factory=dict[str, Any])
250
+ """Values bound by name, written ``:name``, and shared by every statement.
251
+
252
+ A statement that names no parameter simply binds none of them. An engine whose tables can
253
+ be files is handed a storage URI as the file it opens, within the run's own directories."""
254
+
255
+ timeout: Duration = timedelta(minutes=5)
256
+ """How long the whole transaction may run.
257
+
258
+ Set as the database's own statement timeout where the dialect has one, and enforced here
259
+ in every case, so statements that hang end the step rather than the deadline."""
260
+
261
+ @model_validator(mode="after")
262
+ def _check_shape(self) -> "SqlExecuteConfig":
263
+ """Refuse a list whose entries are empty or hold more than one statement each."""
264
+ for index, statement in enumerate(self.statements):
265
+ if not statement.strip():
266
+ raise ValueError(EMPTY_STATEMENT.render(index=index))
267
+ _check_single(statement)
268
+ return self
269
+
270
+
271
+ class SqlExecuteOutput(BlockModel):
272
+ """What the transaction changed."""
273
+
274
+ row_counts: list[int]
275
+ """Rows affected by each statement, in order; ``-1`` where the driver does not say."""
276
+
277
+ duration_ms: int
278
+
279
+
280
+ class SqlExecuteOperator(Operator[SqlExecuteConfig, SqlExecuteOutput]):
281
+ """Runs several statements against a database as one transaction."""
282
+
283
+ spec = OperatorSpec(
284
+ id="sql.execute",
285
+ summary="Run SQL statements against a database in one transaction.",
286
+ idempotent=False,
287
+ )
288
+ config_model: ClassVar[type[BaseModel]] = SqlExecuteConfig
289
+ output_model: ClassVar[type[BaseModel]] = SqlExecuteOutput
290
+
291
+ async def execute(self, config: SqlExecuteConfig, ctx: StepContext) -> SqlExecuteOutput | RemoteHandle:
292
+ """Run every statement in one transaction and report what each of them touched."""
293
+ settings = ctx.connection(config.connection, SqlConnectionConfig)
294
+ if settings.read_only:
295
+ raise BlockFailure(
296
+ READ_ONLY_CONNECTION, error_class=ErrorClass.REJECTED, connection=repr(config.connection)
297
+ )
298
+ url = _resolved(settings, ctx)
299
+ engine = engine_for(url)
300
+ params = engine.bind(config.params, ctx)
301
+ started = time.monotonic()
302
+ counts: list[int] = []
303
+ opened = engine.session(settings, url, ctx, statements=config.statements, params=params)
304
+ async with opened as session, session.begin():
305
+ await _limit(session, settings, config.timeout)
306
+ # The deadline covers every statement together, because the transaction is what
307
+ # the step commits or loses.
308
+ async with asyncio.timeout(config.timeout.total_seconds()):
309
+ for statement in config.statements:
310
+ result = await session.execute(sqlalchemy.text(statement), params)
311
+ counts.append(result.rowcount)
312
+ duration = round((time.monotonic() - started) * 1000)
313
+ ctx.log.info(
314
+ "statements ran",
315
+ connection=config.connection,
316
+ statements=len(config.statements),
317
+ row_counts=list(counts),
318
+ duration_ms=duration,
319
+ )
320
+ return SqlExecuteOutput(row_counts=counts, duration_ms=duration)
321
+
322
+ def classify_error(self, error: Exception) -> ErrorClass:
323
+ """A database that could not be reached is transient; one that refused a statement is not."""
324
+ return classify(error)
325
+
326
+
327
+ def _parse(url: str) -> URL:
328
+ """Read a SQLAlchemy URL, saying what is wrong with one that is not."""
329
+ try:
330
+ return make_url(url)
331
+ except sqlalchemy.exc.ArgumentError as error:
332
+ raise ValueError(NOT_A_DATABASE_URL.render(url=repr(url), detail=str(error))) from error
333
+
334
+
335
+ def _with_password(settings: SqlConnectionConfig) -> URL:
336
+ """The URL with the sealed password merged in, which is the only place the two meet."""
337
+ url = _parse(settings.url)
338
+ if settings.password is None:
339
+ return url
340
+ return url.set(password=settings.password.get_secret_value())
341
+
342
+
343
+ def _resolved(settings: SqlConnectionConfig, ctx: StepContext) -> URL:
344
+ """The URL a step connects with: the password merged in, and a file of the run made absolute."""
345
+ url = _with_password(settings)
346
+ return engine_for(url).resolve(url, ctx)
347
+
348
+
349
+ async def _limit(session: SqlSession, settings: SqlConnectionConfig, timeout: timedelta) -> None:
350
+ """Ask the database to enforce the statement deadline itself, where the dialect has one."""
351
+ template = STATEMENT_TIMEOUT.get(_parse(settings.url).get_backend_name())
352
+ if template is None:
353
+ return
354
+ await session.execute(sqlalchemy.text(template.format(milliseconds=round(timeout.total_seconds() * 1000))))
355
+
356
+
357
+ async def _inline(result: Any, max_rows: int) -> JsonList:
358
+ """Read the rows into the output, refusing a result the step said it would not hold."""
359
+ rows: JsonList = []
360
+ async for batch in _batches(result):
361
+ rows.extend(spelled_row(one) for one in batch)
362
+ if len(rows) > max_rows:
363
+ raise BlockFailure(TOO_MANY_ROWS, error_class=ErrorClass.REJECTED, maximum=max_rows)
364
+ return rows
365
+
366
+
367
+ async def _batches(result: Any) -> AsyncIterator[Sequence[Any]]:
368
+ """Read the cursor a batch at a time, whatever the driver's own chunking is."""
369
+ async for partition in result.partitions(BATCH):
370
+ yield partition
371
+
372
+
373
+ def spelled_row(row: Any) -> JsonMap:
374
+ """One row as an object keyed by column name, every value in its JSON spelling."""
375
+ mapping: dict[str, Any] = dict(row._mapping)
376
+ try:
377
+ return {name: spelled(value) for name, value in mapping.items()}
378
+ except ValueError as error:
379
+ raise BlockFailure(NO_JSON_SPELLING, error_class=ErrorClass.REJECTED, detail=str(error)) from error
380
+
381
+
382
+ def _check_single(statement: str) -> None:
383
+ """Refuse text holding more than one statement.
384
+
385
+ A ``;`` outside a string literal, an identifier, a comment or a dollar-quoted body ends a
386
+ statement, so anything but whitespace and comments after one means there are two. Two
387
+ statements in one field would run outside ``sql.execute``'s transaction, and a document
388
+ that means to run two says so by listing two.
389
+ """
390
+ rest = _after_first_terminator(statement)
391
+ if rest is not None and _stripped(rest):
392
+ raise ValueError(MORE_THAN_ONE_STATEMENT.render())
393
+
394
+
395
+ def _after_first_terminator(statement: str) -> str | None:
396
+ """The text following the first statement-ending ``;``, or None when there is none."""
397
+ index = 0
398
+ length = len(statement)
399
+ while index < length:
400
+ character = statement[index]
401
+ if character == ";":
402
+ return statement[index + 1 :]
403
+ skipped = _skip(statement, index)
404
+ index = skipped if skipped > index else index + 1
405
+ return None
406
+
407
+
408
+ def _skip(statement: str, index: int) -> int:
409
+ """The index just past whatever quoted or commented run starts here, or ``index`` if none does."""
410
+ character = statement[index]
411
+ if statement.startswith("--", index):
412
+ end = statement.find("\n", index)
413
+ return _length_or(statement, end + 1 if end != -1 else -1)
414
+ if statement.startswith("/*", index):
415
+ end = statement.find("*/", index + 2)
416
+ return _length_or(statement, end + 2 if end != -1 else -1)
417
+ if character in "'\"`":
418
+ return _skip_quoted(statement, index, character)
419
+ if character == "$":
420
+ return _skip_dollar(statement, index)
421
+ return index
422
+
423
+
424
+ def _skip_quoted(statement: str, index: int, quote: str) -> int:
425
+ """Past a quoted run, in which the quote character is doubled to mean itself."""
426
+ cursor = index + 1
427
+ while cursor < len(statement):
428
+ if statement[cursor] == quote:
429
+ if statement.startswith(quote * 2, cursor):
430
+ cursor += 2
431
+ continue
432
+ return cursor + 1
433
+ cursor += 1
434
+ return len(statement)
435
+
436
+
437
+ def _skip_dollar(statement: str, index: int) -> int:
438
+ """Past a ``$tag$ ... $tag$`` body, which is how PostgreSQL writes a function containing ``;``."""
439
+ close = statement.find("$", index + 1)
440
+ if close == -1:
441
+ return index
442
+ tag = statement[index : close + 1]
443
+ if not tag[1:-1].replace("_", "").isalnum() and tag != "$$":
444
+ return index
445
+ end = statement.find(tag, close + 1)
446
+ return _length_or(statement, end + len(tag) if end != -1 else -1)
447
+
448
+
449
+ def _length_or(statement: str, end: int) -> int:
450
+ """An index into the statement, or its end when the run was never closed."""
451
+ return len(statement) if end == -1 else end
452
+
453
+
454
+ def _stripped(text: str) -> str:
455
+ """The text with whitespace and trailing comments removed, so a comment is not a statement."""
456
+ rest = text
457
+ while True:
458
+ rest = rest.strip()
459
+ if rest.startswith("--"):
460
+ _, _, rest = rest.partition("\n")
461
+ continue
462
+ if rest.startswith("/*"):
463
+ _, _, rest = rest.partition("*/")
464
+ continue
465
+ return rest
466
+
467
+
468
+ def classify(error: Exception) -> ErrorClass:
469
+ """A database that could not be reached is transient; one that refused the request is not."""
470
+ if isinstance(error, TimeoutError | sqlalchemy.exc.TimeoutError | sqlalchemy.exc.DisconnectionError):
471
+ return ErrorClass.TRANSIENT
472
+ if isinstance(error, sqlalchemy.exc.DBAPIError):
473
+ if error.connection_invalidated or any(marker in str(error).lower() for marker in UNREACHABLE):
474
+ return ErrorClass.TRANSIENT
475
+ return ErrorClass.REJECTED
476
+ if isinstance(error, sqlalchemy.exc.SQLAlchemyError):
477
+ return ErrorClass.REJECTED
478
+ if isinstance(error, OSError):
479
+ return ErrorClass.TRANSIENT
480
+ return ErrorClass.UNKNOWN