codekingpro/portable-devtools
114k
1# ext/horizontal_shard.py
2# Copyright (C) 2005-2024 the SQLAlchemy authors and contributors
3# <see AUTHORS file>
4#
5# This module is part of SQLAlchemy and is released under
6# the MIT License: https://www.opensource.org/licenses/mit-license.php
7
8"""Horizontal sharding support.
9
10Defines a rudimental 'horizontal sharding' system which allows a Session to
11distribute queries and persistence operations across multiple databases.
12
13For a usage example, see the :ref:`examples_sharding` example included in
14the source distribution.
15
16.. deepalchemy:: The horizontal sharding extension is an advanced feature,
17 involving a complex statement -> database interaction as well as
18 use of semi-public APIs for non-trivial cases. Simpler approaches to
19 refering to multiple database "shards", most commonly using a distinct
20 :class:`_orm.Session` per "shard", should always be considered first
21 before using this more complex and less-production-tested system.
22
23
24
25"""
26from __future__ import annotations
27
28from typing import Any
29from typing import Callable
30from typing import Dict
31from typing import Iterable
32from typing import Optional
33from typing import Tuple
34from typing import Type
35from typing import TYPE_CHECKING
36from typing import TypeVar
37from typing import Union
38
39from .. import event
40from .. import exc
41from .. import inspect
42from .. import util
43from ..orm import PassiveFlag
44from ..orm._typing import OrmExecuteOptionsParameter
45from ..orm.interfaces import ORMOption
46from ..orm.mapper import Mapper
47from ..orm.query import Query
48from ..orm.session import _BindArguments
49from ..orm.session import _PKIdentityArgument
50from ..orm.session import Session
51from ..util.typing import Protocol
52from ..util.typing import Self
53
54if TYPE_CHECKING:
55 from ..engine.base import Connection
56 from ..engine.base import Engine
57 from ..engine.base import OptionEngine
58 from ..engine.result import IteratorResult
59 from ..engine.result import Result
60 from ..orm import LoaderCallableStatus
61 from ..orm._typing import _O
62 from ..orm.bulk_persistence import BulkUDCompileState
63 from ..orm.context import QueryContext
64 from ..orm.session import _EntityBindKey
65 from ..orm.session import _SessionBind
66 from ..orm.session import ORMExecuteState
67 from ..orm.state import InstanceState
68 from ..sql import Executable
69 from ..sql._typing import _TP
70 from ..sql.elements import ClauseElement
71
72__all__ = ["ShardedSession", "ShardedQuery"]
73
74_T = TypeVar("_T", bound=Any)
75
76
77ShardIdentifier = str
78
79
80class ShardChooser(Protocol):
81 def __call__(
82 self,
83 mapper: Optional[Mapper[_T]],
84 instance: Any,
85 clause: Optional[ClauseElement],
86 ) -> Any: ...
87
88
89class IdentityChooser(Protocol):
90 def __call__(
91 self,
92 mapper: Mapper[_T],
93 primary_key: _PKIdentityArgument,
94 *,
95 lazy_loaded_from: Optional[InstanceState[Any]],
96 execution_options: OrmExecuteOptionsParameter,
97 bind_arguments: _BindArguments,
98 **kw: Any,
99 ) -> Any: ...
100
101
102class ShardedQuery(Query[_T]):
103 """Query class used with :class:`.ShardedSession`.
104
105 .. legacy:: The :class:`.ShardedQuery` is a subclass of the legacy
106 :class:`.Query` class. The :class:`.ShardedSession` now supports
107 2.0 style execution via the :meth:`.ShardedSession.execute` method.
108
109 """
110
111 def __init__(self, *args: Any, **kwargs: Any) -> None:
112 super().__init__(*args, **kwargs)
113 assert isinstance(self.session, ShardedSession)
114
115 self.identity_chooser = self.session.identity_chooser
116 self.execute_chooser = self.session.execute_chooser
117 self._shard_id = None
118
119 def set_shard(self, shard_id: ShardIdentifier) -> Self:
120 """Return a new query, limited to a single shard ID.
121
122 All subsequent operations with the returned query will
123 be against the single shard regardless of other state.
124
125 The shard_id can be passed for a 2.0 style execution to the
126 bind_arguments dictionary of :meth:`.Session.execute`::
127
128 results = session.execute(
129 stmt,
130 bind_arguments={"shard_id": "my_shard"}
131 )
132
133 """
134 return self.execution_options(_sa_shard_id=shard_id)
135
136
137class ShardedSession(Session):
138 shard_chooser: ShardChooser
139 identity_chooser: IdentityChooser
140 execute_chooser: Callable[[ORMExecuteState], Iterable[Any]]
141
142 def __init__(
143 self,
144 shard_chooser: ShardChooser,
145 identity_chooser: Optional[IdentityChooser] = None,
146 execute_chooser: Optional[
147 Callable[[ORMExecuteState], Iterable[Any]]
148 ] = None,
149 shards: Optional[Dict[str, Any]] = None,
150 query_cls: Type[Query[_T]] = ShardedQuery,
151 *,
152 id_chooser: Optional[
153 Callable[[Query[_T], Iterable[_T]], Iterable[Any]]
154 ] = None,
155 query_chooser: Optional[Callable[[Executable], Iterable[Any]]] = None,
156 **kwargs: Any,
157 ) -> None:
158 """Construct a ShardedSession.
159
160 :param shard_chooser: A callable which, passed a Mapper, a mapped
161 instance, and possibly a SQL clause, returns a shard ID. This id
162 may be based off of the attributes present within the object, or on
163 some round-robin scheme. If the scheme is based on a selection, it
164 should set whatever state on the instance to mark it in the future as
165 participating in that shard.
166
167 :param identity_chooser: A callable, passed a Mapper and primary key
168 argument, which should return a list of shard ids where this
169 primary key might reside.
170
171 .. versionchanged:: 2.0 The ``identity_chooser`` parameter
172 supersedes the ``id_chooser`` parameter.
173
174 :param execute_chooser: For a given :class:`.ORMExecuteState`,
175 returns the list of shard_ids
176 where the query should be issued. Results from all shards returned
177 will be combined together into a single listing.
178
179 .. versionchanged:: 1.4 The ``execute_chooser`` parameter
180 supersedes the ``query_chooser`` parameter.
181
182 :param shards: A dictionary of string shard names
183 to :class:`~sqlalchemy.engine.Engine` objects.
184
185 """
186 super().__init__(query_cls=query_cls, **kwargs)
187
188 event.listen(
189 self, "do_orm_execute", execute_and_instances, retval=True
190 )
191 self.shard_chooser = shard_chooser
192
193 if id_chooser:
194 _id_chooser = id_chooser
195 util.warn_deprecated(
196 "The ``id_chooser`` parameter is deprecated; "
197 "please use ``identity_chooser``.",
198 "2.0",
199 )
200
201 def _legacy_identity_chooser(
202 mapper: Mapper[_T],
203 primary_key: _PKIdentityArgument,
204 *,
205 lazy_loaded_from: Optional[InstanceState[Any]],
206 execution_options: OrmExecuteOptionsParameter,
207 bind_arguments: _BindArguments,
208 **kw: Any,
209 ) -> Any:
210 q = self.query(mapper)
211 if lazy_loaded_from:
212 q = q._set_lazyload_from(lazy_loaded_from)
213 return _id_chooser(q, primary_key)
214
215 self.identity_chooser = _legacy_identity_chooser
216 elif identity_chooser:
217 self.identity_chooser = identity_chooser
218 else:
219 raise exc.ArgumentError(
220 "identity_chooser or id_chooser is required"
221 )
222
223 if query_chooser:
224 _query_chooser = query_chooser
225 util.warn_deprecated(
226 "The ``query_chooser`` parameter is deprecated; "
227 "please use ``execute_chooser``.",
228 "1.4",
229 )
230 if execute_chooser:
231 raise exc.ArgumentError(
232 "Can't pass query_chooser and execute_chooser "
233 "at the same time."
234 )
235
236 def _default_execute_chooser(
237 orm_context: ORMExecuteState,
238 ) -> Iterable[Any]:
239 return _query_chooser(orm_context.statement)
240
241 if execute_chooser is None:
242 execute_chooser = _default_execute_chooser
243
244 if execute_chooser is None:
245 raise exc.ArgumentError(
246 "execute_chooser or query_chooser is required"
247 )
248 self.execute_chooser = execute_chooser
249 self.__shards: Dict[ShardIdentifier, _SessionBind] = {}
250 if shards is not None:
251 for k in shards:
252 self.bind_shard(k, shards[k])
253
254 def _identity_lookup(
255 self,
256 mapper: Mapper[_O],
257 primary_key_identity: Union[Any, Tuple[Any, ...]],
258 identity_token: Optional[Any] = None,
259 passive: PassiveFlag = PassiveFlag.PASSIVE_OFF,
260 lazy_loaded_from: Optional[InstanceState[Any]] = None,
261 execution_options: OrmExecuteOptionsParameter = util.EMPTY_DICT,
262 bind_arguments: Optional[_BindArguments] = None,
263 **kw: Any,
264 ) -> Union[Optional[_O], LoaderCallableStatus]:
265 """override the default :meth:`.Session._identity_lookup` method so
266 that we search for a given non-token primary key identity across all
267 possible identity tokens (e.g. shard ids).
268
269 .. versionchanged:: 1.4 Moved :meth:`.Session._identity_lookup` from
270 the :class:`_query.Query` object to the :class:`.Session`.
271
272 """
273
274 if identity_token is not None:
275 obj = super()._identity_lookup(
276 mapper,
277 primary_key_identity,
278 identity_token=identity_token,
279 **kw,
280 )
281
282 return obj
283 else:
284 for shard_id in self.identity_chooser(
285 mapper,
286 primary_key_identity,
287 lazy_loaded_from=lazy_loaded_from,
288 execution_options=execution_options,
289 bind_arguments=dict(bind_arguments) if bind_arguments else {},
290 ):
291 obj2 = super()._identity_lookup(
292 mapper,
293 primary_key_identity,
294 identity_token=shard_id,
295 lazy_loaded_from=lazy_loaded_from,
296 **kw,
297 )
298 if obj2 is not None:
299 return obj2
300
301 return None
302
303 def _choose_shard_and_assign(
304 self,
305 mapper: Optional[_EntityBindKey[_O]],
306 instance: Any,
307 **kw: Any,
308 ) -> Any:
309 if instance is not None:
310 state = inspect(instance)
311 if state.key:
312 token = state.key[2]
313 assert token is not None
314 return token
315 elif state.identity_token:
316 return state.identity_token
317
318 assert isinstance(mapper, Mapper)
319 shard_id = self.shard_chooser(mapper, instance, **kw)
320 if instance is not None:
321 state.identity_token = shard_id
322 return shard_id
323
324 def connection_callable( # type: ignore [override]
325 self,
326 mapper: Optional[Mapper[_T]] = None,
327 instance: Optional[Any] = None,
328 shard_id: Optional[ShardIdentifier] = None,
329 **kw: Any,
330 ) -> Connection:
331 """Provide a :class:`_engine.Connection` to use in the unit of work
332 flush process.
333
334 """
335
336 if shard_id is None:
337 shard_id = self._choose_shard_and_assign(mapper, instance)
338
339 if self.in_transaction():
340 trans = self.get_transaction()
341 assert trans is not None
342 return trans.connection(mapper, shard_id=shard_id)
343 else:
344 bind = self.get_bind(
345 mapper=mapper, shard_id=shard_id, instance=instance
346 )
347
348 if isinstance(bind, Engine):
349 return bind.connect(**kw)
350 else:
351 assert isinstance(bind, Connection)
352 return bind
353
354 def get_bind(
355 self,
356 mapper: Optional[_EntityBindKey[_O]] = None,
357 *,
358 shard_id: Optional[ShardIdentifier] = None,
359 instance: Optional[Any] = None,
360 clause: Optional[ClauseElement] = None,
361 **kw: Any,
362 ) -> _SessionBind:
363 if shard_id is None:
364 shard_id = self._choose_shard_and_assign(
365 mapper, instance=instance, clause=clause
366 )
367 assert shard_id is not None
368 return self.__shards[shard_id]
369
370 def bind_shard(
371 self, shard_id: ShardIdentifier, bind: Union[Engine, OptionEngine]
372 ) -> None:
373 self.__shards[shard_id] = bind
374
375
376class set_shard_id(ORMOption):
377 """a loader option for statements to apply a specific shard id to the
378 primary query as well as for additional relationship and column
379 loaders.
380
381 The :class:`_horizontal.set_shard_id` option may be applied using
382 the :meth:`_sql.Executable.options` method of any executable statement::
383
384 stmt = (
385 select(MyObject).
386 where(MyObject.name == 'some name').
387 options(set_shard_id("shard1"))
388 )
389
390 Above, the statement when invoked will limit to the "shard1" shard
391 identifier for the primary query as well as for all relationship and
392 column loading strategies, including eager loaders such as
393 :func:`_orm.selectinload`, deferred column loaders like :func:`_orm.defer`,
394 and the lazy relationship loader :func:`_orm.lazyload`.
395
396 In this way, the :class:`_horizontal.set_shard_id` option has much wider
397 scope than using the "shard_id" argument within the
398 :paramref:`_orm.Session.execute.bind_arguments` dictionary.
399
400
401 .. versionadded:: 2.0.0
402
403 """
404
405 __slots__ = ("shard_id", "propagate_to_loaders")
406
407 def __init__(
408 self, shard_id: ShardIdentifier, propagate_to_loaders: bool = True
409 ):
410 """Construct a :class:`_horizontal.set_shard_id` option.
411
412 :param shard_id: shard identifier
413 :param propagate_to_loaders: if left at its default of ``True``, the
414 shard option will take place for lazy loaders such as
415 :func:`_orm.lazyload` and :func:`_orm.defer`; if False, the option
416 will not be propagated to loaded objects. Note that :func:`_orm.defer`
417 always limits to the shard_id of the parent row in any case, so the
418 parameter only has a net effect on the behavior of the
419 :func:`_orm.lazyload` strategy.
420
421 """
422 self.shard_id = shard_id
423 self.propagate_to_loaders = propagate_to_loaders
424
425
426def execute_and_instances(
427 orm_context: ORMExecuteState,
428) -> Union[Result[_T], IteratorResult[_TP]]:
429 active_options: Union[
430 None,
431 QueryContext.default_load_options,
432 Type[QueryContext.default_load_options],
433 BulkUDCompileState.default_update_options,
434 Type[BulkUDCompileState.default_update_options],
435 ]
436
437 if orm_context.is_select:
438 active_options = orm_context.load_options
439
440 elif orm_context.is_update or orm_context.is_delete:
441 active_options = orm_context.update_delete_options
442 else:
443 active_options = None
444
445 session = orm_context.session
446 assert isinstance(session, ShardedSession)
447
448 def iter_for_shard(
449 shard_id: ShardIdentifier,
450 ) -> Union[Result[_T], IteratorResult[_TP]]:
451 bind_arguments = dict(orm_context.bind_arguments)
452 bind_arguments["shard_id"] = shard_id
453
454 orm_context.update_execution_options(identity_token=shard_id)
455 return orm_context.invoke_statement(bind_arguments=bind_arguments)
456
457 for orm_opt in orm_context._non_compile_orm_options:
458 # TODO: if we had an ORMOption that gets applied at ORM statement
459 # execution time, that would allow this to be more generalized.
460 # for now just iterate and look for our options
461 if isinstance(orm_opt, set_shard_id):
462 shard_id = orm_opt.shard_id
463 break
464 else:
465 if active_options and active_options._identity_token is not None:
466 shard_id = active_options._identity_token
467 elif "_sa_shard_id" in orm_context.execution_options:
468 shard_id = orm_context.execution_options["_sa_shard_id"]
469 elif "shard_id" in orm_context.bind_arguments:
470 shard_id = orm_context.bind_arguments["shard_id"]
471 else:
472 shard_id = None
473
474 if shard_id is not None:
475 return iter_for_shard(shard_id)
476 else:
477 partial = []
478 for shard_id in session.execute_chooser(orm_context):
479 result_ = iter_for_shard(shard_id)
480 partial.append(result_)
481 return partial[0].merge(*partial[1:])
482