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.
- dirigent_block_sql-0.17.0/LICENSE +18 -0
- dirigent_block_sql-0.17.0/PKG-INFO +28 -0
- dirigent_block_sql-0.17.0/README.md +14 -0
- dirigent_block_sql-0.17.0/pyproject.toml +30 -0
- dirigent_block_sql-0.17.0/pyproject.toml.orig +28 -0
- dirigent_block_sql-0.17.0/src/dirigent_block_sql/__init__.py +34 -0
- dirigent_block_sql-0.17.0/src/dirigent_block_sql/engines.py +321 -0
- dirigent_block_sql-0.17.0/src/dirigent_block_sql/markers.py +18 -0
- dirigent_block_sql-0.17.0/src/dirigent_block_sql/messages.py +65 -0
- dirigent_block_sql-0.17.0/src/dirigent_block_sql/py.typed +0 -0
- dirigent_block_sql-0.17.0/src/dirigent_block_sql/sql.py +480 -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,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
|
+
)
|
|
File without changes
|
|
@@ -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
|