xtr-orm 1.5.0__py3-none-any.whl

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.
Files changed (57) hide show
  1. xtr_orm/__init__.py +51 -0
  2. xtr_orm/bundle/__init__.py +11 -0
  3. xtr_orm/bundle/connection_config.py +250 -0
  4. xtr_orm/bundle/orm_bundle.py +346 -0
  5. xtr_orm/bundle/orm_config.py +101 -0
  6. xtr_orm/command/__init__.py +56 -0
  7. xtr_orm/command/connection_command.py +134 -0
  8. xtr_orm/command/connections.py +37 -0
  9. xtr_orm/command/database_create_command.py +53 -0
  10. xtr_orm/command/database_drop_command.py +65 -0
  11. xtr_orm/command/migrations_current_command.py +44 -0
  12. xtr_orm/command/migrations_diff_command.py +60 -0
  13. xtr_orm/command/migrations_dump_schema_command.py +56 -0
  14. xtr_orm/command/migrations_execute_command.py +86 -0
  15. xtr_orm/command/migrations_generate_command.py +47 -0
  16. xtr_orm/command/migrations_latest_command.py +44 -0
  17. xtr_orm/command/migrations_list_command.py +71 -0
  18. xtr_orm/command/migrations_migrate_command.py +108 -0
  19. xtr_orm/command/migrations_rollup_command.py +43 -0
  20. xtr_orm/command/migrations_status_command.py +66 -0
  21. xtr_orm/command/migrations_up_to_date_command.py +83 -0
  22. xtr_orm/command/migrations_version_command.py +89 -0
  23. xtr_orm/command/run_sql_command.py +77 -0
  24. xtr_orm/connection_registry.py +192 -0
  25. xtr_orm/database/__init__.py +7 -0
  26. xtr_orm/database/database_manager.py +196 -0
  27. xtr_orm/exception/__init__.py +19 -0
  28. xtr_orm/exception/invalid_argument_error.py +28 -0
  29. xtr_orm/exception/migration_error.py +28 -0
  30. xtr_orm/exception/orm_error.py +14 -0
  31. xtr_orm/exception/session_unavailable_error.py +29 -0
  32. xtr_orm/exception/unknown_connection_error.py +28 -0
  33. xtr_orm/exception/unsupported_database_error.py +29 -0
  34. xtr_orm/messenger/__init__.py +23 -0
  35. xtr_orm/messenger/close_connection_middleware.py +75 -0
  36. xtr_orm/messenger/open_transaction_logger_middleware.py +90 -0
  37. xtr_orm/messenger/transaction_middleware.py +104 -0
  38. xtr_orm/migrations/__init__.py +29 -0
  39. xtr_orm/migrations/_environment.py +188 -0
  40. xtr_orm/migrations/_history.py +178 -0
  41. xtr_orm/migrations/_revisions.py +198 -0
  42. xtr_orm/migrations/_schema.py +66 -0
  43. xtr_orm/migrations/available_migration.py +27 -0
  44. xtr_orm/migrations/direction.py +14 -0
  45. xtr_orm/migrations/executed_migration.py +31 -0
  46. xtr_orm/migrations/execution_result.py +26 -0
  47. xtr_orm/migrations/migration_plan.py +26 -0
  48. xtr_orm/migrations/migration_status.py +48 -0
  49. xtr_orm/migrations/migrations_config.py +162 -0
  50. xtr_orm/migrations/migrator.py +680 -0
  51. xtr_orm/migrations/template/script.py.mako +26 -0
  52. xtr_orm/py.typed +0 -0
  53. xtr_orm-1.5.0.dist-info/METADATA +487 -0
  54. xtr_orm-1.5.0.dist-info/RECORD +57 -0
  55. xtr_orm-1.5.0.dist-info/WHEEL +4 -0
  56. xtr_orm-1.5.0.dist-info/entry_points.txt +3 -0
  57. xtr_orm-1.5.0.dist-info/licenses/LICENSE +21 -0
xtr_orm/__init__.py ADDED
@@ -0,0 +1,51 @@
1
+ """Databases for async applications: engines, sessions and versioned schema migrations.
2
+
3
+ A :class:`~xtr_orm.migrations.Migrator` runs one database's revisions, a
4
+ :class:`~xtr_orm.database.DatabaseManager` creates and drops the database
5
+ itself, and a :class:`ConnectionRegistry` names every connection an
6
+ application has. The bundle in :mod:`xtr_orm.bundle` builds all of them, and
7
+ the engines and sessions, from one configuration.
8
+ """
9
+
10
+ from __future__ import annotations
11
+
12
+ from .connection_registry import DEFAULT_CONNECTION, ConnectionRegistry
13
+ from .database import DatabaseManager
14
+ from .exception import (
15
+ InvalidArgumentError,
16
+ MigrationError,
17
+ OrmError,
18
+ SessionUnavailableError,
19
+ UnknownConnectionError,
20
+ UnsupportedDatabaseError,
21
+ )
22
+ from .migrations import (
23
+ AvailableMigration,
24
+ Direction,
25
+ ExecutedMigration,
26
+ ExecutionResult,
27
+ MigrationPlan,
28
+ MigrationsConfig,
29
+ MigrationStatus,
30
+ Migrator,
31
+ )
32
+
33
+ __all__ = [
34
+ "DEFAULT_CONNECTION",
35
+ "AvailableMigration",
36
+ "ConnectionRegistry",
37
+ "DatabaseManager",
38
+ "Direction",
39
+ "ExecutedMigration",
40
+ "ExecutionResult",
41
+ "InvalidArgumentError",
42
+ "MigrationError",
43
+ "MigrationPlan",
44
+ "MigrationStatus",
45
+ "MigrationsConfig",
46
+ "Migrator",
47
+ "OrmError",
48
+ "SessionUnavailableError",
49
+ "UnknownConnectionError",
50
+ "UnsupportedDatabaseError",
51
+ ]
@@ -0,0 +1,11 @@
1
+ """The xtr-dependency-injection bundle for xtr-orm."""
2
+
3
+ from __future__ import annotations
4
+
5
+ from xtr_orm.migrations import MigrationsConfig
6
+
7
+ from .connection_config import ConnectionConfig
8
+ from .orm_bundle import ORM_CHANNEL, OrmBundle
9
+ from .orm_config import OrmConfig
10
+
11
+ __all__ = ["ORM_CHANNEL", "ConnectionConfig", "MigrationsConfig", "OrmBundle", "OrmConfig"]
@@ -0,0 +1,250 @@
1
+ """How one database connection is built: its URL, engine, sessions and migrations."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import dataclasses
6
+ from collections.abc import Callable, Mapping, Sequence
7
+ from dataclasses import dataclass, field
8
+ from typing import Final
9
+
10
+ from advanced_alchemy.config import AsyncSessionConfig, EngineConfig, SQLAlchemyAsyncConfig
11
+ from sqlalchemy import MetaData, make_url
12
+
13
+ from xtr_orm.exception import InvalidArgumentError
14
+ from xtr_orm.migrations import MigrationsConfig
15
+
16
+ __all__ = ["ConnectionConfig"]
17
+
18
+ _SET_BY_THE_BUNDLE: Final = frozenset(
19
+ {
20
+ "connection_string",
21
+ "engine_config",
22
+ "session_config",
23
+ "bind_key",
24
+ "metadata",
25
+ "engine_instance",
26
+ "session_maker",
27
+ "alembic_config",
28
+ "create_engine_callable",
29
+ "session_maker_class",
30
+ "routing_config",
31
+ },
32
+ )
33
+ """Options of the database layer's own configuration that other fields here already set."""
34
+
35
+
36
+ _NOTHING: Final[frozenset[str]] = frozenset()
37
+
38
+
39
+ def _flag(value: str) -> bool:
40
+ lowered = value.strip().lower()
41
+ if lowered in {"1", "true", "yes", "on"}:
42
+ return True
43
+ if lowered in {"0", "false", "no", "off"}:
44
+ return False
45
+ raise ValueError("a flag: true or false")
46
+
47
+
48
+ def _echo(value: str) -> bool | str:
49
+ return "debug" if value.strip().lower() == "debug" else _flag(value)
50
+
51
+
52
+ def _number(value: str) -> int:
53
+ try:
54
+ return int(value)
55
+ except ValueError:
56
+ raise ValueError("a whole number") from None
57
+
58
+
59
+ def _seconds(value: str) -> float:
60
+ try:
61
+ return float(value)
62
+ except ValueError:
63
+ raise ValueError("a number") from None
64
+
65
+
66
+ def _text(value: str) -> str:
67
+ return value
68
+
69
+
70
+ _ENGINE_URL_OPTIONS: Final[dict[str, Callable[[str], object]]] = {
71
+ "echo": _echo,
72
+ "echo_pool": _echo,
73
+ "enable_from_linting": _flag,
74
+ "hide_parameters": _flag,
75
+ "insertmanyvalues_page_size": _number,
76
+ "isolation_level": _text,
77
+ "label_length": _number,
78
+ "logging_name": _text,
79
+ "max_identifier_length": _number,
80
+ "max_overflow": _number,
81
+ "paramstyle": _text,
82
+ "pool_logging_name": _text,
83
+ "pool_pre_ping": _flag,
84
+ "pool_recycle": _number,
85
+ "pool_size": _number,
86
+ "pool_timeout": _seconds,
87
+ "pool_use_lifo": _flag,
88
+ "query_cache_size": _number,
89
+ "use_insertmanyvalues": _flag,
90
+ }
91
+ """Engine options a URL's query may carry, and how each is read."""
92
+
93
+ _SESSION_URL_OPTIONS: Final[dict[str, Callable[[str], object]]] = {
94
+ "autobegin": _flag,
95
+ "autoflush": _flag,
96
+ "expire_on_commit": _flag,
97
+ "join_transaction_mode": _text,
98
+ "twophase": _flag,
99
+ }
100
+ """Session options a URL's query may carry, and how each is read."""
101
+
102
+
103
+ def _no_options() -> dict[str, object]:
104
+ return {}
105
+
106
+
107
+ def _no_replicas() -> dict[str, str]:
108
+ return {}
109
+
110
+
111
+ @dataclass(frozen=True, slots=True)
112
+ class ConnectionConfig:
113
+ """One database: where it is, its replicas, its engine and sessions, and its migrations.
114
+
115
+ ```python
116
+ ConnectionConfig(
117
+ url=env("DATABASE_URL"),
118
+ replicas={"replica1": env("DATABASE_REPLICA_URL")},
119
+ engine_options={"pool_size": 10, "pool_pre_ping": True},
120
+ session_options={"expire_on_commit": False},
121
+ migrations=MigrationsConfig(render_as_batch=True),
122
+ )
123
+ ```
124
+
125
+ Options may ride on the URL too, which is how one environment variable
126
+ configures a whole connection —
127
+ ``postgresql+asyncpg://app@db/shop?pool_size=20&expire_on_commit=false&ssl=require``.
128
+ An engine or session option, or ``keep_replica``, is taken out of the
129
+ URL and merged over the one configured here: the URL wins. Every other
130
+ option stays on the URL for the driver.
131
+
132
+ Attributes:
133
+ url: The database URL, with an async driver —
134
+ ``postgresql+asyncpg://…``, ``sqlite+aiosqlite:///…``,
135
+ ``mysql+asyncmy://…``. Read from ``DATABASE_URL`` by default, when
136
+ the connection is first used.
137
+ engine_options: Engine options — ``echo``, ``pool_size``,
138
+ ``max_overflow``, ``pool_pre_ping``, ``pool_recycle``,
139
+ ``isolation_level``, ``connect_args``, ``execution_options`` and
140
+ every other the engine takes.
141
+ session_options: Session options — ``expire_on_commit``,
142
+ ``autoflush``, ``info`` and the rest.
143
+ alchemy_options: Any other option of the database layer's own
144
+ configuration, such as ``enable_touch_updated_timestamp_listener``.
145
+ replicas: Read replicas of the database :attr:`url` names, each a URL
146
+ by name. With any, a session reads from a replica picked at
147
+ random, and goes to the primary — :attr:`url` — for a write, a
148
+ transaction's statements, and every read after them. Migrations,
149
+ and the create and drop commands, always use the primary.
150
+ keep_replica: Go back to the replicas once a session commits, rather
151
+ than read from the primary for the rest of its unit of work.
152
+ bind_key: Which table definitions belong to this database when models
153
+ declare several; ``None`` for the default ones.
154
+ metadata: The table definitions a migration diff compares the
155
+ database with; those registered for :attr:`bind_key` when
156
+ ``None``.
157
+ migrations: Where the revisions are and how they are recorded;
158
+ ``migrations`` in the project directory for the default
159
+ connection, ``migrations/<name>`` for any other, when ``None``.
160
+ """
161
+
162
+ url: str = "%env(DATABASE_URL)%"
163
+ engine_options: Mapping[str, object] = field(default_factory=_no_options)
164
+ session_options: Mapping[str, object] = field(default_factory=_no_options)
165
+ alchemy_options: Mapping[str, object] = field(default_factory=_no_options)
166
+ replicas: Mapping[str, str] = field(default_factory=_no_replicas)
167
+ keep_replica: bool = False
168
+ bind_key: str | None = None
169
+ metadata: MetaData | Sequence[MetaData] | None = None
170
+ migrations: MigrationsConfig | None = None
171
+
172
+ def __post_init__(self) -> None:
173
+ """Refuse an empty URL, or an option the engine, a session or the layer does not take.
174
+
175
+ Raises:
176
+ InvalidArgumentError: When the configuration cannot be used.
177
+ """
178
+ # Compared, not tested for truth: the URL may be an environment placeholder,
179
+ # which refuses to decide anything before the variable is read.
180
+ if self.url == "":
181
+ raise InvalidArgumentError("A connection's URL must not be empty.")
182
+ for name, url in self.replicas.items():
183
+ if not isinstance(name, str) or not name: # pyright: ignore[reportUnnecessaryIsInstance] -- configs are written by hand; the annotation is not enforced
184
+ raise InvalidArgumentError(f"A replica needs a non-empty name, got {name!r}.")
185
+ if url == "":
186
+ raise InvalidArgumentError(f'The "{name}" replica\'s URL must not be empty.')
187
+ _check_options("engine", EngineConfig, self.engine_options)
188
+ _check_options("session", AsyncSessionConfig, self.session_options)
189
+ _check_options(
190
+ "alchemy", SQLAlchemyAsyncConfig, self.alchemy_options, exclude=_SET_BY_THE_BUNDLE
191
+ )
192
+
193
+ def with_url_options(self) -> ConnectionConfig:
194
+ """Return this connection with the options its URL carries merged into it.
195
+
196
+ Called once the URL is read — an environment variable is not until a
197
+ service needs it — so a mistake in the URL fails that service.
198
+
199
+ Raises:
200
+ InvalidArgumentError: When an option's value cannot be read.
201
+ """
202
+ url = make_url(self.url)
203
+ driver: dict[str, str | tuple[str, ...]] = {}
204
+ engine = dict(self.engine_options)
205
+ session = dict(self.session_options)
206
+ keep_replica = self.keep_replica
207
+ for name, given in url.query.items():
208
+ value = given if isinstance(given, str) else given[-1]
209
+ if name in _ENGINE_URL_OPTIONS:
210
+ engine[name] = _read(name, value, _ENGINE_URL_OPTIONS[name])
211
+ elif name in _SESSION_URL_OPTIONS:
212
+ session[name] = _read(name, value, _SESSION_URL_OPTIONS[name])
213
+ elif name == "keep_replica":
214
+ keep_replica = bool(_read(name, value, _flag))
215
+ else:
216
+ driver[name] = given
217
+ if len(driver) == len(url.query):
218
+ return self
219
+ return dataclasses.replace(
220
+ self,
221
+ url=url.set(query=driver).render_as_string(hide_password=False),
222
+ engine_options=engine,
223
+ session_options=session,
224
+ keep_replica=keep_replica,
225
+ )
226
+
227
+
228
+ def _read(name: str, value: str, reader: Callable[[str], object]) -> object:
229
+ try:
230
+ return reader(value)
231
+ except ValueError as error:
232
+ raise InvalidArgumentError(
233
+ f'The URL option "{name}" must be {error}, got "{value}".',
234
+ ) from error
235
+
236
+
237
+ def _check_options(
238
+ kind: str,
239
+ target: type[EngineConfig | AsyncSessionConfig | SQLAlchemyAsyncConfig],
240
+ options: Mapping[str, object],
241
+ *,
242
+ exclude: frozenset[str] = _NOTHING,
243
+ ) -> None:
244
+ known = {each.name for each in dataclasses.fields(target) if each.init} - exclude
245
+ for name in options:
246
+ if name not in known:
247
+ listed = ", ".join(sorted(known))
248
+ raise InvalidArgumentError(
249
+ f'Unknown {kind} option "{name}"; the {kind} options are: {listed}.',
250
+ )
@@ -0,0 +1,346 @@
1
+ """The xtr-orm bundle: an engine, sessions and a migrator per configured connection.
2
+
3
+ An application listing :class:`OrmBundle` gets, for every connection of its
4
+ :class:`OrmConfig`, each qualified by the connection's name:
5
+
6
+ - the database layer's :class:`~advanced_alchemy.config.SQLAlchemyAsyncConfig`,
7
+ for its repositories and services;
8
+ - an :class:`~sqlalchemy.ext.asyncio.AsyncEngine` on the primary database;
9
+ - an ``async_sessionmaker[AsyncSession]`` — or, for a connection with
10
+ replicas, a :class:`~advanced_alchemy.routing.RoutingAsyncSessionMaker`
11
+ sending reads to them and writes to the primary;
12
+ - an :class:`~sqlalchemy.ext.asyncio.AsyncSession` per unit of work — a
13
+ request, a message, a command — closed when the unit ends;
14
+ - a :class:`~xtr_orm.migrations.Migrator` and a
15
+ :class:`~xtr_orm.database.DatabaseManager`, both on the primary;
16
+
17
+ and the default connection's without a qualifier too. Every engine is
18
+ disposed when the container closes. When the console bundle is active, the
19
+ ``orm:*`` commands are registered; when the messenger bundle is, the
20
+ ``orm_transaction``, ``orm_close_connection`` and
21
+ ``orm_open_transaction_logger`` middleware are, for a bus to list; when the
22
+ logging bundle is, an ``orm`` channel is added and every migrator logs there.
23
+
24
+ Nothing connects until a service is first used, and a URL given as
25
+ ``env(...)`` is read only then.
26
+ """
27
+
28
+ from __future__ import annotations
29
+
30
+ from collections.abc import AsyncIterator, Awaitable, Callable
31
+ from typing import Final, TypeVar, cast, final
32
+
33
+ from advanced_alchemy.base import metadata_registry
34
+ from advanced_alchemy.config import AsyncSessionConfig, EngineConfig, SQLAlchemyAsyncConfig
35
+ from advanced_alchemy.config.routing import EngineConfig as ReplicaEngine
36
+ from advanced_alchemy.config.routing import RoutingConfig, RoutingStrategy
37
+ from advanced_alchemy.routing import RoutingAsyncSessionMaker, stick_to_primary_var
38
+ from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker
39
+ from typing_extensions import override
40
+ from xtr_dependency_injection import (
41
+ Bundle,
42
+ ContainerBuilder,
43
+ ServiceConfigurator,
44
+ as_bundle,
45
+ bundle_active,
46
+ current_unit_of_work,
47
+ named_factory,
48
+ optional_service,
49
+ required_bundle,
50
+ )
51
+ from xtr_logging_contracts import LoggerInterface
52
+ from xtr_service_contracts import ContainerInterface
53
+
54
+ from xtr_orm.connection_registry import ConnectionRegistry
55
+ from xtr_orm.database import DatabaseManager
56
+ from xtr_orm.exception import SessionUnavailableError
57
+ from xtr_orm.migrations import MigrationsConfig, Migrator
58
+
59
+ from .connection_config import ConnectionConfig
60
+ from .orm_config import OrmConfig
61
+
62
+ __all__ = ["ORM_CHANNEL", "OrmBundle"]
63
+
64
+ _T = TypeVar("_T")
65
+
66
+ ORM_CHANNEL: Final = "orm"
67
+ """The logging channel migrators write to."""
68
+
69
+ _PER_CONNECTION: Final[tuple[type, ...]] = (
70
+ SQLAlchemyAsyncConfig,
71
+ AsyncEngine,
72
+ AsyncSession,
73
+ Migrator,
74
+ DatabaseManager,
75
+ )
76
+
77
+
78
+ @final
79
+ @required_bundle("xtr_logging.bundle:LoggingBundle", ignore_on_invalid=True)
80
+ @required_bundle("xtr_console.bundle:ConsoleBundle", ignore_on_invalid=True)
81
+ @as_bundle("orm", config=OrmConfig)
82
+ class OrmBundle(Bundle[OrmConfig]):
83
+ """Turns an :class:`OrmConfig` into an engine, sessions and a migrator per connection."""
84
+
85
+ @override
86
+ def prepend_extension(self, builder: ContainerBuilder) -> None:
87
+ """When ``logging`` is active, add the ``orm`` channel to its config."""
88
+ if not bundle_active(builder, "logging"):
89
+ return
90
+ # Logging is an optional peer, importable only once it is active.
91
+ from xtr_logging.bundle import LoggingConfig # noqa: PLC0415
92
+
93
+ def add_orm_channel(config: LoggingConfig) -> LoggingConfig:
94
+ return config.with_channels(ORM_CHANNEL)
95
+
96
+ builder.prepend_extension_config(LoggingConfig, add_orm_channel)
97
+
98
+ @override
99
+ def load_extension(
100
+ self,
101
+ config: OrmConfig,
102
+ services: ServiceConfigurator,
103
+ builder: ContainerBuilder,
104
+ ) -> None:
105
+ """Register every connection's services under its name, the registry, the commands."""
106
+ for name, spec in config.connections.items():
107
+ maker_type: type = (
108
+ RoutingAsyncSessionMaker if spec.replicas else async_sessionmaker[AsyncSession]
109
+ )
110
+ factories = (
111
+ (_alchemy_factory(name), "alchemy"),
112
+ (_engine_factory(name), "engine"),
113
+ (
114
+ _routing_maker_factory(name) if spec.replicas else _session_maker_factory(name),
115
+ "session_maker",
116
+ ),
117
+ (_migrator_factory(name), "migrator"),
118
+ (_database_factory(name), "database"),
119
+ )
120
+ for factory, kind in factories:
121
+ _ = services.set(named_factory(factory, f"orm_{kind}_{name}"), qualifier=name)
122
+ _ = services.set(
123
+ named_factory(_session_factory(name), f"orm_session_{name}"),
124
+ qualifier=name,
125
+ lifetime="scoped",
126
+ )
127
+ if name == config.default_connection:
128
+ for service in (*_PER_CONNECTION, maker_type):
129
+ services.alias(service, service, target_qualifier=name)
130
+
131
+ _ = services.set(_connection_registry)
132
+
133
+ if bundle_active(builder, "console"):
134
+ services.load("xtr_orm.command")
135
+ if bundle_active(builder, "messenger"):
136
+ services.load("xtr_orm.messenger")
137
+
138
+
139
+ def _alchemy_factory(
140
+ name: str,
141
+ ) -> Callable[[OrmConfig], AsyncIterator[SQLAlchemyAsyncConfig]]:
142
+ """Build the factory of the connection ``name``'s database layer configuration.
143
+
144
+ It owns every engine of the connection — the primary's, and each
145
+ replica's — and disposes them when the container closes. It injects the
146
+ configuration rather than closing over it: the container hands it a copy
147
+ with every environment placeholder resolved.
148
+ """
149
+
150
+ async def alchemy(config: OrmConfig) -> AsyncIterator[SQLAlchemyAsyncConfig]:
151
+ spec = config.connections[name].with_url_options()
152
+ # The options are forwarded by name, each checked against the target's
153
+ # fields by ConnectionConfig; their values are the application's.
154
+ engine = cast("Callable[..., EngineConfig]", EngineConfig)(**spec.engine_options)
155
+ session = cast("Callable[..., AsyncSessionConfig]", AsyncSessionConfig)(
156
+ **spec.session_options
157
+ )
158
+ built = cast("Callable[..., SQLAlchemyAsyncConfig]", SQLAlchemyAsyncConfig)(
159
+ # With replicas, the primary is given through the routing instead.
160
+ connection_string=None if spec.replicas else spec.url,
161
+ routing_config=_routing(spec) if spec.replicas else None,
162
+ engine_config=engine,
163
+ session_config=session,
164
+ bind_key=spec.bind_key,
165
+ **spec.alchemy_options,
166
+ )
167
+ routing_maker: RoutingAsyncSessionMaker | None = None
168
+ if spec.replicas:
169
+ routing_maker = cast("RoutingAsyncSessionMaker", built.create_session_maker())
170
+ # One pool on the primary: the routing's, which migrations use too.
171
+ built.engine_instance = routing_maker.primary_engine
172
+ try:
173
+ yield built
174
+ finally:
175
+ if routing_maker is not None:
176
+ await routing_maker.close_all()
177
+ elif built.engine_instance is not None:
178
+ await built.engine_instance.dispose()
179
+
180
+ return alchemy
181
+
182
+
183
+ def _routing(spec: ConnectionConfig) -> RoutingConfig:
184
+ """Map the connection's replicas onto the database layer's read/write routing.
185
+
186
+ Reads go to a replica picked at random; a write, a transaction and every
187
+ read after them go to the primary, until a commit when ``keep_replica``.
188
+ Built afresh for each container: that configuration completes itself when
189
+ built.
190
+ """
191
+ return RoutingConfig(
192
+ primary_connection_string=spec.url,
193
+ read_replicas=[
194
+ ReplicaEngine(connection_string=url, name=replica)
195
+ for replica, url in spec.replicas.items()
196
+ ],
197
+ routing_strategy=RoutingStrategy.RANDOM,
198
+ sticky_after_write=True,
199
+ reset_stickiness_on_commit=spec.keep_replica,
200
+ )
201
+
202
+
203
+ def _engine_factory(name: str) -> Callable[[ContainerInterface], Awaitable[AsyncEngine]]:
204
+ """Build the factory of the connection ``name``'s engine, on its primary."""
205
+
206
+ async def engine(container: ContainerInterface) -> AsyncEngine:
207
+ return (await container.get(SQLAlchemyAsyncConfig, name)).get_engine()
208
+
209
+ return engine
210
+
211
+
212
+ def _session_maker_factory(
213
+ name: str,
214
+ ) -> Callable[[ContainerInterface], Awaitable[async_sessionmaker[AsyncSession]]]:
215
+ """Build the factory of the connection ``name``'s session maker, on its engine."""
216
+
217
+ async def session_maker(container: ContainerInterface) -> async_sessionmaker[AsyncSession]:
218
+ alchemy = await container.get(SQLAlchemyAsyncConfig, name)
219
+ return cast("async_sessionmaker[AsyncSession]", alchemy.create_session_maker())
220
+
221
+ return session_maker
222
+
223
+
224
+ def _routing_maker_factory(
225
+ name: str,
226
+ ) -> Callable[[ContainerInterface], Awaitable[RoutingAsyncSessionMaker]]:
227
+ """Build the factory of the connection ``name``'s session maker, routing to its replicas."""
228
+
229
+ async def routing_session_maker(container: ContainerInterface) -> RoutingAsyncSessionMaker:
230
+ alchemy = await container.get(SQLAlchemyAsyncConfig, name)
231
+ return cast("RoutingAsyncSessionMaker", alchemy.create_session_maker())
232
+
233
+ return routing_session_maker
234
+
235
+
236
+ def _session_factory(
237
+ name: str,
238
+ ) -> Callable[[OrmConfig, ContainerInterface], AsyncIterator[AsyncSession]]:
239
+ """Build the factory of the connection ``name``'s session, closed when its unit ends.
240
+
241
+ Closing rolls back what was not committed: committing is the unit's job.
242
+ With replicas, a unit starts reading from them whatever the unit before
243
+ it wrote: sticking to the primary after a write is one unit's state.
244
+ """
245
+
246
+ async def session(
247
+ config: OrmConfig, container: ContainerInterface
248
+ ) -> AsyncIterator[AsyncSession]:
249
+ routed = bool(config.connections[name].replicas)
250
+ maker = (await container.get(SQLAlchemyAsyncConfig, name)).create_session_maker()
251
+ if routed:
252
+ _ = stick_to_primary_var.set(False)
253
+ try:
254
+ async with maker() as opened:
255
+ yield opened
256
+ finally:
257
+ if routed:
258
+ _ = stick_to_primary_var.set(False)
259
+
260
+ return session
261
+
262
+
263
+ def _migrator_factory(name: str) -> Callable[[OrmConfig, ContainerInterface], Awaitable[Migrator]]:
264
+ """Build the factory of the connection ``name``'s migrator, logging to ``orm`` when it can."""
265
+
266
+ async def migrator(config: OrmConfig, container: ContainerInterface) -> Migrator:
267
+ spec = config.connections[name].with_url_options()
268
+ metadata = (
269
+ spec.metadata if spec.metadata is not None else metadata_registry.get(spec.bind_key)
270
+ )
271
+ built = Migrator(
272
+ await container.get(AsyncEngine, name),
273
+ spec.migrations or MigrationsConfig(),
274
+ metadata,
275
+ name=name,
276
+ )
277
+ logger = await optional_service(container, LoggerInterface, ORM_CHANNEL)
278
+ if logger is not None:
279
+ built.set_logger(logger)
280
+ return built
281
+
282
+ return migrator
283
+
284
+
285
+ def _database_factory(name: str) -> Callable[[OrmConfig], DatabaseManager]:
286
+ """Build the factory of the connection ``name``'s database manager."""
287
+
288
+ def database(config: OrmConfig) -> DatabaseManager:
289
+ spec: ConnectionConfig = config.connections[name].with_url_options()
290
+ connect_args = spec.engine_options.get("connect_args")
291
+ return DatabaseManager(
292
+ spec.url,
293
+ connect_args=connect_args if isinstance(connect_args, dict) else None, # pyright: ignore[reportUnknownArgumentType] -- an option mapping, read as given
294
+ )
295
+
296
+ return database
297
+
298
+
299
+ def _connection_registry(config: OrmConfig, container: ContainerInterface) -> ConnectionRegistry:
300
+ """Name every connection for the commands and middleware, each piece built only when needed.
301
+
302
+ A connection's session is the one of the unit of work under way, and
303
+ closing it disposes every engine it has — its replicas' too.
304
+ """
305
+ registry = ConnectionRegistry(config.default_connection)
306
+ for name, spec in config.connections.items():
307
+ registry.register(
308
+ name,
309
+ engine=_provider(container, AsyncEngine, name),
310
+ migrator=_provider(container, Migrator, name),
311
+ database=_provider(container, DatabaseManager, name),
312
+ session=_unit_session(name),
313
+ close=_closer(container, name, routed=bool(spec.replicas)),
314
+ )
315
+ return registry
316
+
317
+
318
+ def _unit_session(name: str) -> Callable[[], Awaitable[AsyncSession]]:
319
+ async def session() -> AsyncSession:
320
+ unit = current_unit_of_work()
321
+ if unit is None:
322
+ raise SessionUnavailableError(name, "no unit of work is under way")
323
+ return await unit.get(AsyncSession, name)
324
+
325
+ return session
326
+
327
+
328
+ def _closer(
329
+ container: ContainerInterface, name: str, *, routed: bool
330
+ ) -> Callable[[], Awaitable[None]]:
331
+ async def close() -> None:
332
+ if routed:
333
+ await (await container.get(RoutingAsyncSessionMaker, name)).close_all()
334
+ else:
335
+ await (await container.get(AsyncEngine, name)).dispose()
336
+
337
+ return close
338
+
339
+
340
+ def _provider(
341
+ container: ContainerInterface, service: type[_T], name: str
342
+ ) -> Callable[[], Awaitable[_T]]:
343
+ async def provide() -> _T:
344
+ return await container.get(service, name)
345
+
346
+ return provide