base.py 72 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405140614071408140914101411141214131414141514161417141814191420142114221423142414251426142714281429143014311432143314341435143614371438143914401441144214431444144514461447144814491450145114521453145414551456145714581459146014611462146314641465146614671468146914701471147214731474147514761477147814791480148114821483148414851486148714881489149014911492149314941495149614971498149915001501150215031504150515061507150815091510151115121513151415151516151715181519152015211522152315241525152615271528152915301531153215331534153515361537153815391540154115421543154415451546154715481549155015511552155315541555155615571558155915601561156215631564156515661567156815691570157115721573157415751576157715781579158015811582158315841585158615871588158915901591159215931594159515961597159815991600160116021603160416051606160716081609161016111612161316141615161616171618161916201621162216231624162516261627162816291630163116321633163416351636163716381639164016411642164316441645164616471648164916501651165216531654165516561657165816591660166116621663166416651666166716681669167016711672167316741675167616771678167916801681168216831684168516861687168816891690169116921693169416951696169716981699170017011702170317041705170617071708170917101711171217131714171517161717171817191720172117221723172417251726172717281729173017311732173317341735173617371738173917401741174217431744174517461747174817491750175117521753175417551756175717581759176017611762176317641765176617671768176917701771177217731774177517761777177817791780178117821783178417851786178717881789179017911792179317941795179617971798179918001801180218031804180518061807180818091810181118121813181418151816181718181819182018211822182318241825182618271828182918301831183218331834183518361837183818391840184118421843184418451846184718481849185018511852185318541855185618571858185918601861186218631864186518661867186818691870187118721873187418751876187718781879188018811882188318841885188618871888188918901891189218931894189518961897189818991900190119021903190419051906190719081909191019111912191319141915191619171918191919201921192219231924192519261927192819291930193119321933193419351936193719381939194019411942194319441945194619471948194919501951195219531954195519561957195819591960196119621963196419651966196719681969197019711972197319741975197619771978197919801981198219831984198519861987198819891990199119921993199419951996199719981999200020012002200320042005200620072008200920102011201220132014201520162017201820192020202120222023202420252026202720282029203020312032203320342035203620372038203920402041204220432044204520462047204820492050205120522053205420552056205720582059206020612062206320642065206620672068206920702071207220732074207520762077207820792080208120822083208420852086208720882089209020912092209320942095209620972098209921002101210221032104210521062107210821092110211121122113211421152116211721182119212021212122212321242125212621272128212921302131213221332134213521362137213821392140214121422143214421452146214721482149215021512152215321542155215621572158215921602161216221632164216521662167216821692170217121722173217421752176217721782179218021812182218321842185218621872188218921902191219221932194219521962197
  1. # sql/base.py
  2. # Copyright (C) 2005-2025 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. # mypy: allow-untyped-defs, allow-untyped-calls
  8. """Foundational utilities common to many sql modules.
  9. """
  10. from __future__ import annotations
  11. import collections
  12. from enum import Enum
  13. import itertools
  14. from itertools import zip_longest
  15. import operator
  16. import re
  17. from typing import Any
  18. from typing import Callable
  19. from typing import cast
  20. from typing import Dict
  21. from typing import FrozenSet
  22. from typing import Generic
  23. from typing import Iterable
  24. from typing import Iterator
  25. from typing import List
  26. from typing import Mapping
  27. from typing import MutableMapping
  28. from typing import NamedTuple
  29. from typing import NoReturn
  30. from typing import Optional
  31. from typing import overload
  32. from typing import Sequence
  33. from typing import Set
  34. from typing import Tuple
  35. from typing import Type
  36. from typing import TYPE_CHECKING
  37. from typing import TypeVar
  38. from typing import Union
  39. from . import roles
  40. from . import visitors
  41. from .cache_key import HasCacheKey # noqa
  42. from .cache_key import MemoizedHasCacheKey # noqa
  43. from .traversals import HasCopyInternals # noqa
  44. from .visitors import ClauseVisitor
  45. from .visitors import ExtendedInternalTraversal
  46. from .visitors import ExternallyTraversible
  47. from .visitors import InternalTraversal
  48. from .. import event
  49. from .. import exc
  50. from .. import util
  51. from ..util import HasMemoized as HasMemoized
  52. from ..util import hybridmethod
  53. from ..util import typing as compat_typing
  54. from ..util.typing import Protocol
  55. from ..util.typing import Self
  56. from ..util.typing import TypeGuard
  57. if TYPE_CHECKING:
  58. from . import coercions
  59. from . import elements
  60. from . import type_api
  61. from ._orm_types import DMLStrategyArgument
  62. from ._orm_types import SynchronizeSessionArgument
  63. from ._typing import _CLE
  64. from .compiler import SQLCompiler
  65. from .elements import BindParameter
  66. from .elements import ClauseList
  67. from .elements import ColumnClause # noqa
  68. from .elements import ColumnElement
  69. from .elements import NamedColumn
  70. from .elements import SQLCoreOperations
  71. from .elements import TextClause
  72. from .schema import Column
  73. from .schema import DefaultGenerator
  74. from .selectable import _JoinTargetElement
  75. from .selectable import _SelectIterable
  76. from .selectable import FromClause
  77. from ..engine import Connection
  78. from ..engine import CursorResult
  79. from ..engine.interfaces import _CoreMultiExecuteParams
  80. from ..engine.interfaces import _ExecuteOptions
  81. from ..engine.interfaces import _ImmutableExecuteOptions
  82. from ..engine.interfaces import CacheStats
  83. from ..engine.interfaces import Compiled
  84. from ..engine.interfaces import CompiledCacheType
  85. from ..engine.interfaces import CoreExecuteOptionsParameter
  86. from ..engine.interfaces import Dialect
  87. from ..engine.interfaces import IsolationLevel
  88. from ..engine.interfaces import SchemaTranslateMapType
  89. from ..event import dispatcher
  90. if not TYPE_CHECKING:
  91. coercions = None # noqa
  92. elements = None # noqa
  93. type_api = None # noqa
  94. class _NoArg(Enum):
  95. NO_ARG = 0
  96. def __repr__(self):
  97. return f"_NoArg.{self.name}"
  98. NO_ARG = _NoArg.NO_ARG
  99. class _NoneName(Enum):
  100. NONE_NAME = 0
  101. """indicate a 'deferred' name that was ultimately the value None."""
  102. _NONE_NAME = _NoneName.NONE_NAME
  103. _T = TypeVar("_T", bound=Any)
  104. _Fn = TypeVar("_Fn", bound=Callable[..., Any])
  105. _AmbiguousTableNameMap = MutableMapping[str, str]
  106. class _DefaultDescriptionTuple(NamedTuple):
  107. arg: Any
  108. is_scalar: Optional[bool]
  109. is_callable: Optional[bool]
  110. is_sentinel: Optional[bool]
  111. @classmethod
  112. def _from_column_default(
  113. cls, default: Optional[DefaultGenerator]
  114. ) -> _DefaultDescriptionTuple:
  115. return (
  116. _DefaultDescriptionTuple(
  117. default.arg, # type: ignore
  118. default.is_scalar,
  119. default.is_callable,
  120. default.is_sentinel,
  121. )
  122. if default
  123. and (
  124. default.has_arg
  125. or (not default.for_update and default.is_sentinel)
  126. )
  127. else _DefaultDescriptionTuple(None, None, None, None)
  128. )
  129. _never_select_column = operator.attrgetter("_omit_from_statements")
  130. class _EntityNamespace(Protocol):
  131. def __getattr__(self, key: str) -> SQLCoreOperations[Any]: ...
  132. class _HasEntityNamespace(Protocol):
  133. @util.ro_non_memoized_property
  134. def entity_namespace(self) -> _EntityNamespace: ...
  135. def _is_has_entity_namespace(element: Any) -> TypeGuard[_HasEntityNamespace]:
  136. return hasattr(element, "entity_namespace")
  137. # Remove when https://github.com/python/mypy/issues/14640 will be fixed
  138. _Self = TypeVar("_Self", bound=Any)
  139. class Immutable:
  140. """mark a ClauseElement as 'immutable' when expressions are cloned.
  141. "immutable" objects refers to the "mutability" of an object in the
  142. context of SQL DQL and DML generation. Such as, in DQL, one can
  143. compose a SELECT or subquery of varied forms, but one cannot modify
  144. the structure of a specific table or column within DQL.
  145. :class:`.Immutable` is mostly intended to follow this concept, and as
  146. such the primary "immutable" objects are :class:`.ColumnClause`,
  147. :class:`.Column`, :class:`.TableClause`, :class:`.Table`.
  148. """
  149. __slots__ = ()
  150. _is_immutable = True
  151. def unique_params(self, *optionaldict, **kwargs):
  152. raise NotImplementedError("Immutable objects do not support copying")
  153. def params(self, *optionaldict, **kwargs):
  154. raise NotImplementedError("Immutable objects do not support copying")
  155. def _clone(self: _Self, **kw: Any) -> _Self:
  156. return self
  157. def _copy_internals(
  158. self, *, omit_attrs: Iterable[str] = (), **kw: Any
  159. ) -> None:
  160. pass
  161. class SingletonConstant(Immutable):
  162. """Represent SQL constants like NULL, TRUE, FALSE"""
  163. _is_singleton_constant = True
  164. _singleton: SingletonConstant
  165. def __new__(cls: _T, *arg: Any, **kw: Any) -> _T:
  166. return cast(_T, cls._singleton)
  167. @util.non_memoized_property
  168. def proxy_set(self) -> FrozenSet[ColumnElement[Any]]:
  169. raise NotImplementedError()
  170. @classmethod
  171. def _create_singleton(cls):
  172. obj = object.__new__(cls)
  173. obj.__init__() # type: ignore
  174. # for a long time this was an empty frozenset, meaning
  175. # a SingletonConstant would never be a "corresponding column" in
  176. # a statement. This referred to #6259. However, in #7154 we see
  177. # that we do in fact need "correspondence" to work when matching cols
  178. # in result sets, so the non-correspondence was moved to a more
  179. # specific level when we are actually adapting expressions for SQL
  180. # render only.
  181. obj.proxy_set = frozenset([obj])
  182. cls._singleton = obj
  183. def _from_objects(
  184. *elements: Union[
  185. ColumnElement[Any], FromClause, TextClause, _JoinTargetElement
  186. ]
  187. ) -> Iterator[FromClause]:
  188. return itertools.chain.from_iterable(
  189. [element._from_objects for element in elements]
  190. )
  191. def _select_iterables(
  192. elements: Iterable[roles.ColumnsClauseRole],
  193. ) -> _SelectIterable:
  194. """expand tables into individual columns in the
  195. given list of column expressions.
  196. """
  197. return itertools.chain.from_iterable(
  198. [c._select_iterable for c in elements]
  199. )
  200. _SelfGenerativeType = TypeVar("_SelfGenerativeType", bound="_GenerativeType")
  201. class _GenerativeType(compat_typing.Protocol):
  202. def _generate(self) -> Self: ...
  203. def _generative(fn: _Fn) -> _Fn:
  204. """non-caching _generative() decorator.
  205. This is basically the legacy decorator that copies the object and
  206. runs a method on the new copy.
  207. """
  208. @util.decorator
  209. def _generative(
  210. fn: _Fn, self: _SelfGenerativeType, *args: Any, **kw: Any
  211. ) -> _SelfGenerativeType:
  212. """Mark a method as generative."""
  213. self = self._generate()
  214. x = fn(self, *args, **kw)
  215. assert x is self, "generative methods must return self"
  216. return self
  217. decorated = _generative(fn)
  218. decorated.non_generative = fn # type: ignore
  219. return decorated
  220. def _exclusive_against(*names: str, **kw: Any) -> Callable[[_Fn], _Fn]:
  221. msgs = kw.pop("msgs", {})
  222. defaults = kw.pop("defaults", {})
  223. getters = [
  224. (name, operator.attrgetter(name), defaults.get(name, None))
  225. for name in names
  226. ]
  227. @util.decorator
  228. def check(fn, *args, **kw):
  229. # make pylance happy by not including "self" in the argument
  230. # list
  231. self = args[0]
  232. args = args[1:]
  233. for name, getter, default_ in getters:
  234. if getter(self) is not default_:
  235. msg = msgs.get(
  236. name,
  237. "Method %s() has already been invoked on this %s construct"
  238. % (fn.__name__, self.__class__),
  239. )
  240. raise exc.InvalidRequestError(msg)
  241. return fn(self, *args, **kw)
  242. return check
  243. def _clone(element, **kw):
  244. return element._clone(**kw)
  245. def _expand_cloned(
  246. elements: Iterable[_CLE],
  247. ) -> Iterable[_CLE]:
  248. """expand the given set of ClauseElements to be the set of all 'cloned'
  249. predecessors.
  250. """
  251. # TODO: cython candidate
  252. return itertools.chain(*[x._cloned_set for x in elements])
  253. def _de_clone(
  254. elements: Iterable[_CLE],
  255. ) -> Iterable[_CLE]:
  256. for x in elements:
  257. while x._is_clone_of is not None:
  258. x = x._is_clone_of
  259. yield x
  260. def _cloned_intersection(a: Iterable[_CLE], b: Iterable[_CLE]) -> Set[_CLE]:
  261. """return the intersection of sets a and b, counting
  262. any overlap between 'cloned' predecessors.
  263. The returned set is in terms of the entities present within 'a'.
  264. """
  265. all_overlap = set(_expand_cloned(a)).intersection(_expand_cloned(b))
  266. return {elem for elem in a if all_overlap.intersection(elem._cloned_set)}
  267. def _cloned_difference(a: Iterable[_CLE], b: Iterable[_CLE]) -> Set[_CLE]:
  268. all_overlap = set(_expand_cloned(a)).intersection(_expand_cloned(b))
  269. return {
  270. elem for elem in a if not all_overlap.intersection(elem._cloned_set)
  271. }
  272. class _DialectArgView(MutableMapping[str, Any]):
  273. """A dictionary view of dialect-level arguments in the form
  274. <dialectname>_<argument_name>.
  275. """
  276. __slots__ = ("obj",)
  277. def __init__(self, obj):
  278. self.obj = obj
  279. def _key(self, key):
  280. try:
  281. dialect, value_key = key.split("_", 1)
  282. except ValueError as err:
  283. raise KeyError(key) from err
  284. else:
  285. return dialect, value_key
  286. def __getitem__(self, key):
  287. dialect, value_key = self._key(key)
  288. try:
  289. opt = self.obj.dialect_options[dialect]
  290. except exc.NoSuchModuleError as err:
  291. raise KeyError(key) from err
  292. else:
  293. return opt[value_key]
  294. def __setitem__(self, key, value):
  295. try:
  296. dialect, value_key = self._key(key)
  297. except KeyError as err:
  298. raise exc.ArgumentError(
  299. "Keys must be of the form <dialectname>_<argname>"
  300. ) from err
  301. else:
  302. self.obj.dialect_options[dialect][value_key] = value
  303. def __delitem__(self, key):
  304. dialect, value_key = self._key(key)
  305. del self.obj.dialect_options[dialect][value_key]
  306. def __len__(self):
  307. return sum(
  308. len(args._non_defaults)
  309. for args in self.obj.dialect_options.values()
  310. )
  311. def __iter__(self):
  312. return (
  313. "%s_%s" % (dialect_name, value_name)
  314. for dialect_name in self.obj.dialect_options
  315. for value_name in self.obj.dialect_options[
  316. dialect_name
  317. ]._non_defaults
  318. )
  319. class _DialectArgDict(MutableMapping[str, Any]):
  320. """A dictionary view of dialect-level arguments for a specific
  321. dialect.
  322. Maintains a separate collection of user-specified arguments
  323. and dialect-specified default arguments.
  324. """
  325. def __init__(self):
  326. self._non_defaults = {}
  327. self._defaults = {}
  328. def __len__(self):
  329. return len(set(self._non_defaults).union(self._defaults))
  330. def __iter__(self):
  331. return iter(set(self._non_defaults).union(self._defaults))
  332. def __getitem__(self, key):
  333. if key in self._non_defaults:
  334. return self._non_defaults[key]
  335. else:
  336. return self._defaults[key]
  337. def __setitem__(self, key, value):
  338. self._non_defaults[key] = value
  339. def __delitem__(self, key):
  340. del self._non_defaults[key]
  341. @util.preload_module("sqlalchemy.dialects")
  342. def _kw_reg_for_dialect(dialect_name):
  343. dialect_cls = util.preloaded.dialects.registry.load(dialect_name)
  344. if dialect_cls.construct_arguments is None:
  345. return None
  346. return dict(dialect_cls.construct_arguments)
  347. class DialectKWArgs:
  348. """Establish the ability for a class to have dialect-specific arguments
  349. with defaults and constructor validation.
  350. The :class:`.DialectKWArgs` interacts with the
  351. :attr:`.DefaultDialect.construct_arguments` present on a dialect.
  352. .. seealso::
  353. :attr:`.DefaultDialect.construct_arguments`
  354. """
  355. __slots__ = ()
  356. _dialect_kwargs_traverse_internals = [
  357. ("dialect_options", InternalTraversal.dp_dialect_options)
  358. ]
  359. @classmethod
  360. def argument_for(cls, dialect_name, argument_name, default):
  361. """Add a new kind of dialect-specific keyword argument for this class.
  362. E.g.::
  363. Index.argument_for("mydialect", "length", None)
  364. some_index = Index("a", "b", mydialect_length=5)
  365. The :meth:`.DialectKWArgs.argument_for` method is a per-argument
  366. way adding extra arguments to the
  367. :attr:`.DefaultDialect.construct_arguments` dictionary. This
  368. dictionary provides a list of argument names accepted by various
  369. schema-level constructs on behalf of a dialect.
  370. New dialects should typically specify this dictionary all at once as a
  371. data member of the dialect class. The use case for ad-hoc addition of
  372. argument names is typically for end-user code that is also using
  373. a custom compilation scheme which consumes the additional arguments.
  374. :param dialect_name: name of a dialect. The dialect must be
  375. locatable, else a :class:`.NoSuchModuleError` is raised. The
  376. dialect must also include an existing
  377. :attr:`.DefaultDialect.construct_arguments` collection, indicating
  378. that it participates in the keyword-argument validation and default
  379. system, else :class:`.ArgumentError` is raised. If the dialect does
  380. not include this collection, then any keyword argument can be
  381. specified on behalf of this dialect already. All dialects packaged
  382. within SQLAlchemy include this collection, however for third party
  383. dialects, support may vary.
  384. :param argument_name: name of the parameter.
  385. :param default: default value of the parameter.
  386. """
  387. construct_arg_dictionary = DialectKWArgs._kw_registry[dialect_name]
  388. if construct_arg_dictionary is None:
  389. raise exc.ArgumentError(
  390. "Dialect '%s' does have keyword-argument "
  391. "validation and defaults enabled configured" % dialect_name
  392. )
  393. if cls not in construct_arg_dictionary:
  394. construct_arg_dictionary[cls] = {}
  395. construct_arg_dictionary[cls][argument_name] = default
  396. @property
  397. def dialect_kwargs(self):
  398. """A collection of keyword arguments specified as dialect-specific
  399. options to this construct.
  400. The arguments are present here in their original ``<dialect>_<kwarg>``
  401. format. Only arguments that were actually passed are included;
  402. unlike the :attr:`.DialectKWArgs.dialect_options` collection, which
  403. contains all options known by this dialect including defaults.
  404. The collection is also writable; keys are accepted of the
  405. form ``<dialect>_<kwarg>`` where the value will be assembled
  406. into the list of options.
  407. .. seealso::
  408. :attr:`.DialectKWArgs.dialect_options` - nested dictionary form
  409. """
  410. return _DialectArgView(self)
  411. @property
  412. def kwargs(self):
  413. """A synonym for :attr:`.DialectKWArgs.dialect_kwargs`."""
  414. return self.dialect_kwargs
  415. _kw_registry = util.PopulateDict(_kw_reg_for_dialect)
  416. @classmethod
  417. def _kw_reg_for_dialect_cls(cls, dialect_name):
  418. construct_arg_dictionary = DialectKWArgs._kw_registry[dialect_name]
  419. d = _DialectArgDict()
  420. if construct_arg_dictionary is None:
  421. d._defaults.update({"*": None})
  422. else:
  423. for cls in reversed(cls.__mro__):
  424. if cls in construct_arg_dictionary:
  425. d._defaults.update(construct_arg_dictionary[cls])
  426. return d
  427. @util.memoized_property
  428. def dialect_options(self):
  429. """A collection of keyword arguments specified as dialect-specific
  430. options to this construct.
  431. This is a two-level nested registry, keyed to ``<dialect_name>``
  432. and ``<argument_name>``. For example, the ``postgresql_where``
  433. argument would be locatable as::
  434. arg = my_object.dialect_options["postgresql"]["where"]
  435. .. versionadded:: 0.9.2
  436. .. seealso::
  437. :attr:`.DialectKWArgs.dialect_kwargs` - flat dictionary form
  438. """
  439. return util.PopulateDict(self._kw_reg_for_dialect_cls)
  440. def _validate_dialect_kwargs(self, kwargs: Dict[str, Any]) -> None:
  441. # validate remaining kwargs that they all specify DB prefixes
  442. if not kwargs:
  443. return
  444. for k in kwargs:
  445. m = re.match("^(.+?)_(.+)$", k)
  446. if not m:
  447. raise TypeError(
  448. "Additional arguments should be "
  449. "named <dialectname>_<argument>, got '%s'" % k
  450. )
  451. dialect_name, arg_name = m.group(1, 2)
  452. try:
  453. construct_arg_dictionary = self.dialect_options[dialect_name]
  454. except exc.NoSuchModuleError:
  455. util.warn(
  456. "Can't validate argument %r; can't "
  457. "locate any SQLAlchemy dialect named %r"
  458. % (k, dialect_name)
  459. )
  460. self.dialect_options[dialect_name] = d = _DialectArgDict()
  461. d._defaults.update({"*": None})
  462. d._non_defaults[arg_name] = kwargs[k]
  463. else:
  464. if (
  465. "*" not in construct_arg_dictionary
  466. and arg_name not in construct_arg_dictionary
  467. ):
  468. raise exc.ArgumentError(
  469. "Argument %r is not accepted by "
  470. "dialect %r on behalf of %r"
  471. % (k, dialect_name, self.__class__)
  472. )
  473. else:
  474. construct_arg_dictionary[arg_name] = kwargs[k]
  475. class CompileState:
  476. """Produces additional object state necessary for a statement to be
  477. compiled.
  478. the :class:`.CompileState` class is at the base of classes that assemble
  479. state for a particular statement object that is then used by the
  480. compiler. This process is essentially an extension of the process that
  481. the SQLCompiler.visit_XYZ() method takes, however there is an emphasis
  482. on converting raw user intent into more organized structures rather than
  483. producing string output. The top-level :class:`.CompileState` for the
  484. statement being executed is also accessible when the execution context
  485. works with invoking the statement and collecting results.
  486. The production of :class:`.CompileState` is specific to the compiler, such
  487. as within the :meth:`.SQLCompiler.visit_insert`,
  488. :meth:`.SQLCompiler.visit_select` etc. methods. These methods are also
  489. responsible for associating the :class:`.CompileState` with the
  490. :class:`.SQLCompiler` itself, if the statement is the "toplevel" statement,
  491. i.e. the outermost SQL statement that's actually being executed.
  492. There can be other :class:`.CompileState` objects that are not the
  493. toplevel, such as when a SELECT subquery or CTE-nested
  494. INSERT/UPDATE/DELETE is generated.
  495. .. versionadded:: 1.4
  496. """
  497. __slots__ = ("statement", "_ambiguous_table_name_map")
  498. plugins: Dict[Tuple[str, str], Type[CompileState]] = {}
  499. _ambiguous_table_name_map: Optional[_AmbiguousTableNameMap]
  500. @classmethod
  501. def create_for_statement(
  502. cls, statement: Executable, compiler: SQLCompiler, **kw: Any
  503. ) -> CompileState:
  504. # factory construction.
  505. if statement._propagate_attrs:
  506. plugin_name = statement._propagate_attrs.get(
  507. "compile_state_plugin", "default"
  508. )
  509. klass = cls.plugins.get(
  510. (plugin_name, statement._effective_plugin_target), None
  511. )
  512. if klass is None:
  513. klass = cls.plugins[
  514. ("default", statement._effective_plugin_target)
  515. ]
  516. else:
  517. klass = cls.plugins[
  518. ("default", statement._effective_plugin_target)
  519. ]
  520. if klass is cls:
  521. return cls(statement, compiler, **kw)
  522. else:
  523. return klass.create_for_statement(statement, compiler, **kw)
  524. def __init__(self, statement, compiler, **kw):
  525. self.statement = statement
  526. @classmethod
  527. def get_plugin_class(
  528. cls, statement: Executable
  529. ) -> Optional[Type[CompileState]]:
  530. plugin_name = statement._propagate_attrs.get(
  531. "compile_state_plugin", None
  532. )
  533. if plugin_name:
  534. key = (plugin_name, statement._effective_plugin_target)
  535. if key in cls.plugins:
  536. return cls.plugins[key]
  537. # there's no case where we call upon get_plugin_class() and want
  538. # to get None back, there should always be a default. return that
  539. # if there was no plugin-specific class (e.g. Insert with "orm"
  540. # plugin)
  541. try:
  542. return cls.plugins[("default", statement._effective_plugin_target)]
  543. except KeyError:
  544. return None
  545. @classmethod
  546. def _get_plugin_class_for_plugin(
  547. cls, statement: Executable, plugin_name: str
  548. ) -> Optional[Type[CompileState]]:
  549. try:
  550. return cls.plugins[
  551. (plugin_name, statement._effective_plugin_target)
  552. ]
  553. except KeyError:
  554. return None
  555. @classmethod
  556. def plugin_for(
  557. cls, plugin_name: str, visit_name: str
  558. ) -> Callable[[_Fn], _Fn]:
  559. def decorate(cls_to_decorate):
  560. cls.plugins[(plugin_name, visit_name)] = cls_to_decorate
  561. return cls_to_decorate
  562. return decorate
  563. class Generative(HasMemoized):
  564. """Provide a method-chaining pattern in conjunction with the
  565. @_generative decorator."""
  566. def _generate(self) -> Self:
  567. skip = self._memoized_keys
  568. cls = self.__class__
  569. s = cls.__new__(cls)
  570. if skip:
  571. # ensure this iteration remains atomic
  572. s.__dict__ = {
  573. k: v for k, v in self.__dict__.copy().items() if k not in skip
  574. }
  575. else:
  576. s.__dict__ = self.__dict__.copy()
  577. return s
  578. class InPlaceGenerative(HasMemoized):
  579. """Provide a method-chaining pattern in conjunction with the
  580. @_generative decorator that mutates in place."""
  581. __slots__ = ()
  582. def _generate(self):
  583. skip = self._memoized_keys
  584. # note __dict__ needs to be in __slots__ if this is used
  585. for k in skip:
  586. self.__dict__.pop(k, None)
  587. return self
  588. class HasCompileState(Generative):
  589. """A class that has a :class:`.CompileState` associated with it."""
  590. _compile_state_plugin: Optional[Type[CompileState]] = None
  591. _attributes: util.immutabledict[str, Any] = util.EMPTY_DICT
  592. _compile_state_factory = CompileState.create_for_statement
  593. class _MetaOptions(type):
  594. """metaclass for the Options class.
  595. This metaclass is actually necessary despite the availability of the
  596. ``__init_subclass__()`` hook as this type also provides custom class-level
  597. behavior for the ``__add__()`` method.
  598. """
  599. _cache_attrs: Tuple[str, ...]
  600. def __add__(self, other):
  601. o1 = self()
  602. if set(other).difference(self._cache_attrs):
  603. raise TypeError(
  604. "dictionary contains attributes not covered by "
  605. "Options class %s: %r"
  606. % (self, set(other).difference(self._cache_attrs))
  607. )
  608. o1.__dict__.update(other)
  609. return o1
  610. if TYPE_CHECKING:
  611. def __getattr__(self, key: str) -> Any: ...
  612. def __setattr__(self, key: str, value: Any) -> None: ...
  613. def __delattr__(self, key: str) -> None: ...
  614. class Options(metaclass=_MetaOptions):
  615. """A cacheable option dictionary with defaults."""
  616. __slots__ = ()
  617. _cache_attrs: Tuple[str, ...]
  618. def __init_subclass__(cls) -> None:
  619. dict_ = cls.__dict__
  620. cls._cache_attrs = tuple(
  621. sorted(
  622. d
  623. for d in dict_
  624. if not d.startswith("__")
  625. and d not in ("_cache_key_traversal",)
  626. )
  627. )
  628. super().__init_subclass__()
  629. def __init__(self, **kw):
  630. self.__dict__.update(kw)
  631. def __add__(self, other):
  632. o1 = self.__class__.__new__(self.__class__)
  633. o1.__dict__.update(self.__dict__)
  634. if set(other).difference(self._cache_attrs):
  635. raise TypeError(
  636. "dictionary contains attributes not covered by "
  637. "Options class %s: %r"
  638. % (self, set(other).difference(self._cache_attrs))
  639. )
  640. o1.__dict__.update(other)
  641. return o1
  642. def __eq__(self, other):
  643. # TODO: very inefficient. This is used only in test suites
  644. # right now.
  645. for a, b in zip_longest(self._cache_attrs, other._cache_attrs):
  646. if getattr(self, a) != getattr(other, b):
  647. return False
  648. return True
  649. def __repr__(self):
  650. # TODO: fairly inefficient, used only in debugging right now.
  651. return "%s(%s)" % (
  652. self.__class__.__name__,
  653. ", ".join(
  654. "%s=%r" % (k, self.__dict__[k])
  655. for k in self._cache_attrs
  656. if k in self.__dict__
  657. ),
  658. )
  659. @classmethod
  660. def isinstance(cls, klass: Type[Any]) -> bool:
  661. return issubclass(cls, klass)
  662. @hybridmethod
  663. def add_to_element(self, name, value):
  664. return self + {name: getattr(self, name) + value}
  665. @hybridmethod
  666. def _state_dict_inst(self) -> Mapping[str, Any]:
  667. return self.__dict__
  668. _state_dict_const: util.immutabledict[str, Any] = util.EMPTY_DICT
  669. @_state_dict_inst.classlevel
  670. def _state_dict(cls) -> Mapping[str, Any]:
  671. return cls._state_dict_const
  672. @classmethod
  673. def safe_merge(cls, other):
  674. d = other._state_dict()
  675. # only support a merge with another object of our class
  676. # and which does not have attrs that we don't. otherwise
  677. # we risk having state that might not be part of our cache
  678. # key strategy
  679. if (
  680. cls is not other.__class__
  681. and other._cache_attrs
  682. and set(other._cache_attrs).difference(cls._cache_attrs)
  683. ):
  684. raise TypeError(
  685. "other element %r is not empty, is not of type %s, "
  686. "and contains attributes not covered here %r"
  687. % (
  688. other,
  689. cls,
  690. set(other._cache_attrs).difference(cls._cache_attrs),
  691. )
  692. )
  693. return cls + d
  694. @classmethod
  695. def from_execution_options(
  696. cls, key, attrs, exec_options, statement_exec_options
  697. ):
  698. """process Options argument in terms of execution options.
  699. e.g.::
  700. (
  701. load_options,
  702. execution_options,
  703. ) = QueryContext.default_load_options.from_execution_options(
  704. "_sa_orm_load_options",
  705. {"populate_existing", "autoflush", "yield_per"},
  706. execution_options,
  707. statement._execution_options,
  708. )
  709. get back the Options and refresh "_sa_orm_load_options" in the
  710. exec options dict w/ the Options as well
  711. """
  712. # common case is that no options we are looking for are
  713. # in either dictionary, so cancel for that first
  714. check_argnames = attrs.intersection(
  715. set(exec_options).union(statement_exec_options)
  716. )
  717. existing_options = exec_options.get(key, cls)
  718. if check_argnames:
  719. result = {}
  720. for argname in check_argnames:
  721. local = "_" + argname
  722. if argname in exec_options:
  723. result[local] = exec_options[argname]
  724. elif argname in statement_exec_options:
  725. result[local] = statement_exec_options[argname]
  726. new_options = existing_options + result
  727. exec_options = util.immutabledict().merge_with(
  728. exec_options, {key: new_options}
  729. )
  730. return new_options, exec_options
  731. else:
  732. return existing_options, exec_options
  733. if TYPE_CHECKING:
  734. def __getattr__(self, key: str) -> Any: ...
  735. def __setattr__(self, key: str, value: Any) -> None: ...
  736. def __delattr__(self, key: str) -> None: ...
  737. class CacheableOptions(Options, HasCacheKey):
  738. __slots__ = ()
  739. @hybridmethod
  740. def _gen_cache_key_inst(self, anon_map, bindparams):
  741. return HasCacheKey._gen_cache_key(self, anon_map, bindparams)
  742. @_gen_cache_key_inst.classlevel
  743. def _gen_cache_key(cls, anon_map, bindparams):
  744. return (cls, ())
  745. @hybridmethod
  746. def _generate_cache_key(self):
  747. return HasCacheKey._generate_cache_key_for_object(self)
  748. class ExecutableOption(HasCopyInternals):
  749. __slots__ = ()
  750. _annotations = util.EMPTY_DICT
  751. __visit_name__ = "executable_option"
  752. _is_has_cache_key = False
  753. _is_core = True
  754. def _clone(self, **kw):
  755. """Create a shallow copy of this ExecutableOption."""
  756. c = self.__class__.__new__(self.__class__)
  757. c.__dict__ = dict(self.__dict__) # type: ignore
  758. return c
  759. class Executable(roles.StatementRole):
  760. """Mark a :class:`_expression.ClauseElement` as supporting execution.
  761. :class:`.Executable` is a superclass for all "statement" types
  762. of objects, including :func:`select`, :func:`delete`, :func:`update`,
  763. :func:`insert`, :func:`text`.
  764. """
  765. supports_execution: bool = True
  766. _execution_options: _ImmutableExecuteOptions = util.EMPTY_DICT
  767. _is_default_generator = False
  768. _with_options: Tuple[ExecutableOption, ...] = ()
  769. _with_context_options: Tuple[
  770. Tuple[Callable[[CompileState], None], Any], ...
  771. ] = ()
  772. _compile_options: Optional[Union[Type[CacheableOptions], CacheableOptions]]
  773. _executable_traverse_internals = [
  774. ("_with_options", InternalTraversal.dp_executable_options),
  775. (
  776. "_with_context_options",
  777. ExtendedInternalTraversal.dp_with_context_options,
  778. ),
  779. ("_propagate_attrs", ExtendedInternalTraversal.dp_propagate_attrs),
  780. ]
  781. is_select = False
  782. is_from_statement = False
  783. is_update = False
  784. is_insert = False
  785. is_text = False
  786. is_delete = False
  787. is_dml = False
  788. if TYPE_CHECKING:
  789. __visit_name__: str
  790. def _compile_w_cache(
  791. self,
  792. dialect: Dialect,
  793. *,
  794. compiled_cache: Optional[CompiledCacheType],
  795. column_keys: List[str],
  796. for_executemany: bool = False,
  797. schema_translate_map: Optional[SchemaTranslateMapType] = None,
  798. **kw: Any,
  799. ) -> Tuple[
  800. Compiled, Optional[Sequence[BindParameter[Any]]], CacheStats
  801. ]: ...
  802. def _execute_on_connection(
  803. self,
  804. connection: Connection,
  805. distilled_params: _CoreMultiExecuteParams,
  806. execution_options: CoreExecuteOptionsParameter,
  807. ) -> CursorResult[Any]: ...
  808. def _execute_on_scalar(
  809. self,
  810. connection: Connection,
  811. distilled_params: _CoreMultiExecuteParams,
  812. execution_options: CoreExecuteOptionsParameter,
  813. ) -> Any: ...
  814. @util.ro_non_memoized_property
  815. def _all_selected_columns(self):
  816. raise NotImplementedError()
  817. @property
  818. def _effective_plugin_target(self) -> str:
  819. return self.__visit_name__
  820. @_generative
  821. def options(self, *options: ExecutableOption) -> Self:
  822. """Apply options to this statement.
  823. In the general sense, options are any kind of Python object
  824. that can be interpreted by the SQL compiler for the statement.
  825. These options can be consumed by specific dialects or specific kinds
  826. of compilers.
  827. The most commonly known kind of option are the ORM level options
  828. that apply "eager load" and other loading behaviors to an ORM
  829. query. However, options can theoretically be used for many other
  830. purposes.
  831. For background on specific kinds of options for specific kinds of
  832. statements, refer to the documentation for those option objects.
  833. .. versionchanged:: 1.4 - added :meth:`.Executable.options` to
  834. Core statement objects towards the goal of allowing unified
  835. Core / ORM querying capabilities.
  836. .. seealso::
  837. :ref:`loading_columns` - refers to options specific to the usage
  838. of ORM queries
  839. :ref:`relationship_loader_options` - refers to options specific
  840. to the usage of ORM queries
  841. """
  842. self._with_options += tuple(
  843. coercions.expect(roles.ExecutableOptionRole, opt)
  844. for opt in options
  845. )
  846. return self
  847. @_generative
  848. def _set_compile_options(self, compile_options: CacheableOptions) -> Self:
  849. """Assign the compile options to a new value.
  850. :param compile_options: appropriate CacheableOptions structure
  851. """
  852. self._compile_options = compile_options
  853. return self
  854. @_generative
  855. def _update_compile_options(self, options: CacheableOptions) -> Self:
  856. """update the _compile_options with new keys."""
  857. assert self._compile_options is not None
  858. self._compile_options += options
  859. return self
  860. @_generative
  861. def _add_context_option(
  862. self,
  863. callable_: Callable[[CompileState], None],
  864. cache_args: Any,
  865. ) -> Self:
  866. """Add a context option to this statement.
  867. These are callable functions that will
  868. be given the CompileState object upon compilation.
  869. A second argument cache_args is required, which will be combined with
  870. the ``__code__`` identity of the function itself in order to produce a
  871. cache key.
  872. """
  873. self._with_context_options += ((callable_, cache_args),)
  874. return self
  875. @overload
  876. def execution_options(
  877. self,
  878. *,
  879. compiled_cache: Optional[CompiledCacheType] = ...,
  880. logging_token: str = ...,
  881. isolation_level: IsolationLevel = ...,
  882. no_parameters: bool = False,
  883. stream_results: bool = False,
  884. max_row_buffer: int = ...,
  885. yield_per: int = ...,
  886. insertmanyvalues_page_size: int = ...,
  887. schema_translate_map: Optional[SchemaTranslateMapType] = ...,
  888. populate_existing: bool = False,
  889. autoflush: bool = False,
  890. synchronize_session: SynchronizeSessionArgument = ...,
  891. dml_strategy: DMLStrategyArgument = ...,
  892. render_nulls: bool = ...,
  893. is_delete_using: bool = ...,
  894. is_update_from: bool = ...,
  895. preserve_rowcount: bool = False,
  896. **opt: Any,
  897. ) -> Self: ...
  898. @overload
  899. def execution_options(self, **opt: Any) -> Self: ...
  900. @_generative
  901. def execution_options(self, **kw: Any) -> Self:
  902. """Set non-SQL options for the statement which take effect during
  903. execution.
  904. Execution options can be set at many scopes, including per-statement,
  905. per-connection, or per execution, using methods such as
  906. :meth:`_engine.Connection.execution_options` and parameters which
  907. accept a dictionary of options such as
  908. :paramref:`_engine.Connection.execute.execution_options` and
  909. :paramref:`_orm.Session.execute.execution_options`.
  910. The primary characteristic of an execution option, as opposed to
  911. other kinds of options such as ORM loader options, is that
  912. **execution options never affect the compiled SQL of a query, only
  913. things that affect how the SQL statement itself is invoked or how
  914. results are fetched**. That is, execution options are not part of
  915. what's accommodated by SQL compilation nor are they considered part of
  916. the cached state of a statement.
  917. The :meth:`_sql.Executable.execution_options` method is
  918. :term:`generative`, as
  919. is the case for the method as applied to the :class:`_engine.Engine`
  920. and :class:`_orm.Query` objects, which means when the method is called,
  921. a copy of the object is returned, which applies the given parameters to
  922. that new copy, but leaves the original unchanged::
  923. statement = select(table.c.x, table.c.y)
  924. new_statement = statement.execution_options(my_option=True)
  925. An exception to this behavior is the :class:`_engine.Connection`
  926. object, where the :meth:`_engine.Connection.execution_options` method
  927. is explicitly **not** generative.
  928. The kinds of options that may be passed to
  929. :meth:`_sql.Executable.execution_options` and other related methods and
  930. parameter dictionaries include parameters that are explicitly consumed
  931. by SQLAlchemy Core or ORM, as well as arbitrary keyword arguments not
  932. defined by SQLAlchemy, which means the methods and/or parameter
  933. dictionaries may be used for user-defined parameters that interact with
  934. custom code, which may access the parameters using methods such as
  935. :meth:`_sql.Executable.get_execution_options` and
  936. :meth:`_engine.Connection.get_execution_options`, or within selected
  937. event hooks using a dedicated ``execution_options`` event parameter
  938. such as
  939. :paramref:`_events.ConnectionEvents.before_execute.execution_options`
  940. or :attr:`_orm.ORMExecuteState.execution_options`, e.g.::
  941. from sqlalchemy import event
  942. @event.listens_for(some_engine, "before_execute")
  943. def _process_opt(conn, statement, multiparams, params, execution_options):
  944. "run a SQL function before invoking a statement"
  945. if execution_options.get("do_special_thing", False):
  946. conn.exec_driver_sql("run_special_function()")
  947. Within the scope of options that are explicitly recognized by
  948. SQLAlchemy, most apply to specific classes of objects and not others.
  949. The most common execution options include:
  950. * :paramref:`_engine.Connection.execution_options.isolation_level` -
  951. sets the isolation level for a connection or a class of connections
  952. via an :class:`_engine.Engine`. This option is accepted only
  953. by :class:`_engine.Connection` or :class:`_engine.Engine`.
  954. * :paramref:`_engine.Connection.execution_options.stream_results` -
  955. indicates results should be fetched using a server side cursor;
  956. this option is accepted by :class:`_engine.Connection`, by the
  957. :paramref:`_engine.Connection.execute.execution_options` parameter
  958. on :meth:`_engine.Connection.execute`, and additionally by
  959. :meth:`_sql.Executable.execution_options` on a SQL statement object,
  960. as well as by ORM constructs like :meth:`_orm.Session.execute`.
  961. * :paramref:`_engine.Connection.execution_options.compiled_cache` -
  962. indicates a dictionary that will serve as the
  963. :ref:`SQL compilation cache <sql_caching>`
  964. for a :class:`_engine.Connection` or :class:`_engine.Engine`, as
  965. well as for ORM methods like :meth:`_orm.Session.execute`.
  966. Can be passed as ``None`` to disable caching for statements.
  967. This option is not accepted by
  968. :meth:`_sql.Executable.execution_options` as it is inadvisable to
  969. carry along a compilation cache within a statement object.
  970. * :paramref:`_engine.Connection.execution_options.schema_translate_map`
  971. - a mapping of schema names used by the
  972. :ref:`Schema Translate Map <schema_translating>` feature, accepted
  973. by :class:`_engine.Connection`, :class:`_engine.Engine`,
  974. :class:`_sql.Executable`, as well as by ORM constructs
  975. like :meth:`_orm.Session.execute`.
  976. .. seealso::
  977. :meth:`_engine.Connection.execution_options`
  978. :paramref:`_engine.Connection.execute.execution_options`
  979. :paramref:`_orm.Session.execute.execution_options`
  980. :ref:`orm_queryguide_execution_options` - documentation on all
  981. ORM-specific execution options
  982. """ # noqa: E501
  983. if "isolation_level" in kw:
  984. raise exc.ArgumentError(
  985. "'isolation_level' execution option may only be specified "
  986. "on Connection.execution_options(), or "
  987. "per-engine using the isolation_level "
  988. "argument to create_engine()."
  989. )
  990. if "compiled_cache" in kw:
  991. raise exc.ArgumentError(
  992. "'compiled_cache' execution option may only be specified "
  993. "on Connection.execution_options(), not per statement."
  994. )
  995. self._execution_options = self._execution_options.union(kw)
  996. return self
  997. def get_execution_options(self) -> _ExecuteOptions:
  998. """Get the non-SQL options which will take effect during execution.
  999. .. versionadded:: 1.3
  1000. .. seealso::
  1001. :meth:`.Executable.execution_options`
  1002. """
  1003. return self._execution_options
  1004. class SchemaEventTarget(event.EventTarget):
  1005. """Base class for elements that are the targets of :class:`.DDLEvents`
  1006. events.
  1007. This includes :class:`.SchemaItem` as well as :class:`.SchemaType`.
  1008. """
  1009. dispatch: dispatcher[SchemaEventTarget]
  1010. def _set_parent(self, parent: SchemaEventTarget, **kw: Any) -> None:
  1011. """Associate with this SchemaEvent's parent object."""
  1012. def _set_parent_with_dispatch(
  1013. self, parent: SchemaEventTarget, **kw: Any
  1014. ) -> None:
  1015. self.dispatch.before_parent_attach(self, parent)
  1016. self._set_parent(parent, **kw)
  1017. self.dispatch.after_parent_attach(self, parent)
  1018. class SchemaVisitable(SchemaEventTarget, visitors.Visitable):
  1019. """Base class for elements that are targets of a :class:`.SchemaVisitor`.
  1020. .. versionadded:: 2.0.41
  1021. """
  1022. class SchemaVisitor(ClauseVisitor):
  1023. """Define the visiting for ``SchemaItem`` and more
  1024. generally ``SchemaVisitable`` objects.
  1025. """
  1026. __traverse_options__ = {"schema_visitor": True}
  1027. class _SentinelDefaultCharacterization(Enum):
  1028. NONE = "none"
  1029. UNKNOWN = "unknown"
  1030. CLIENTSIDE = "clientside"
  1031. SENTINEL_DEFAULT = "sentinel_default"
  1032. SERVERSIDE = "serverside"
  1033. IDENTITY = "identity"
  1034. SEQUENCE = "sequence"
  1035. class _SentinelColumnCharacterization(NamedTuple):
  1036. columns: Optional[Sequence[Column[Any]]] = None
  1037. is_explicit: bool = False
  1038. is_autoinc: bool = False
  1039. default_characterization: _SentinelDefaultCharacterization = (
  1040. _SentinelDefaultCharacterization.NONE
  1041. )
  1042. _COLKEY = TypeVar("_COLKEY", Union[None, str], str)
  1043. _COL_co = TypeVar("_COL_co", bound="ColumnElement[Any]", covariant=True)
  1044. _COL = TypeVar("_COL", bound="ColumnElement[Any]")
  1045. class _ColumnMetrics(Generic[_COL_co]):
  1046. __slots__ = ("column",)
  1047. column: _COL_co
  1048. def __init__(
  1049. self, collection: ColumnCollection[Any, _COL_co], col: _COL_co
  1050. ):
  1051. self.column = col
  1052. # proxy_index being non-empty means it was initialized.
  1053. # so we need to update it
  1054. pi = collection._proxy_index
  1055. if pi:
  1056. for eps_col in col._expanded_proxy_set:
  1057. pi[eps_col].add(self)
  1058. def get_expanded_proxy_set(self):
  1059. return self.column._expanded_proxy_set
  1060. def dispose(self, collection):
  1061. pi = collection._proxy_index
  1062. if not pi:
  1063. return
  1064. for col in self.column._expanded_proxy_set:
  1065. colset = pi.get(col, None)
  1066. if colset:
  1067. colset.discard(self)
  1068. if colset is not None and not colset:
  1069. del pi[col]
  1070. def embedded(
  1071. self,
  1072. target_set: Union[
  1073. Set[ColumnElement[Any]], FrozenSet[ColumnElement[Any]]
  1074. ],
  1075. ) -> bool:
  1076. expanded_proxy_set = self.column._expanded_proxy_set
  1077. for t in target_set.difference(expanded_proxy_set):
  1078. if not expanded_proxy_set.intersection(_expand_cloned([t])):
  1079. return False
  1080. return True
  1081. class ColumnCollection(Generic[_COLKEY, _COL_co]):
  1082. """Collection of :class:`_expression.ColumnElement` instances,
  1083. typically for
  1084. :class:`_sql.FromClause` objects.
  1085. The :class:`_sql.ColumnCollection` object is most commonly available
  1086. as the :attr:`_schema.Table.c` or :attr:`_schema.Table.columns` collection
  1087. on the :class:`_schema.Table` object, introduced at
  1088. :ref:`metadata_tables_and_columns`.
  1089. The :class:`_expression.ColumnCollection` has both mapping- and sequence-
  1090. like behaviors. A :class:`_expression.ColumnCollection` usually stores
  1091. :class:`_schema.Column` objects, which are then accessible both via mapping
  1092. style access as well as attribute access style.
  1093. To access :class:`_schema.Column` objects using ordinary attribute-style
  1094. access, specify the name like any other object attribute, such as below
  1095. a column named ``employee_name`` is accessed::
  1096. >>> employee_table.c.employee_name
  1097. To access columns that have names with special characters or spaces,
  1098. index-style access is used, such as below which illustrates a column named
  1099. ``employee ' payment`` is accessed::
  1100. >>> employee_table.c["employee ' payment"]
  1101. As the :class:`_sql.ColumnCollection` object provides a Python dictionary
  1102. interface, common dictionary method names like
  1103. :meth:`_sql.ColumnCollection.keys`, :meth:`_sql.ColumnCollection.values`,
  1104. and :meth:`_sql.ColumnCollection.items` are available, which means that
  1105. database columns that are keyed under these names also need to use indexed
  1106. access::
  1107. >>> employee_table.c["values"]
  1108. The name for which a :class:`_schema.Column` would be present is normally
  1109. that of the :paramref:`_schema.Column.key` parameter. In some contexts,
  1110. such as a :class:`_sql.Select` object that uses a label style set
  1111. using the :meth:`_sql.Select.set_label_style` method, a column of a certain
  1112. key may instead be represented under a particular label name such
  1113. as ``tablename_columnname``::
  1114. >>> from sqlalchemy import select, column, table
  1115. >>> from sqlalchemy import LABEL_STYLE_TABLENAME_PLUS_COL
  1116. >>> t = table("t", column("c"))
  1117. >>> stmt = select(t).set_label_style(LABEL_STYLE_TABLENAME_PLUS_COL)
  1118. >>> subq = stmt.subquery()
  1119. >>> subq.c.t_c
  1120. <sqlalchemy.sql.elements.ColumnClause at 0x7f59dcf04fa0; t_c>
  1121. :class:`.ColumnCollection` also indexes the columns in order and allows
  1122. them to be accessible by their integer position::
  1123. >>> cc[0]
  1124. Column('x', Integer(), table=None)
  1125. >>> cc[1]
  1126. Column('y', Integer(), table=None)
  1127. .. versionadded:: 1.4 :class:`_expression.ColumnCollection`
  1128. allows integer-based
  1129. index access to the collection.
  1130. Iterating the collection yields the column expressions in order::
  1131. >>> list(cc)
  1132. [Column('x', Integer(), table=None),
  1133. Column('y', Integer(), table=None)]
  1134. The base :class:`_expression.ColumnCollection` object can store
  1135. duplicates, which can
  1136. mean either two columns with the same key, in which case the column
  1137. returned by key access is **arbitrary**::
  1138. >>> x1, x2 = Column("x", Integer), Column("x", Integer)
  1139. >>> cc = ColumnCollection(columns=[(x1.name, x1), (x2.name, x2)])
  1140. >>> list(cc)
  1141. [Column('x', Integer(), table=None),
  1142. Column('x', Integer(), table=None)]
  1143. >>> cc["x"] is x1
  1144. False
  1145. >>> cc["x"] is x2
  1146. True
  1147. Or it can also mean the same column multiple times. These cases are
  1148. supported as :class:`_expression.ColumnCollection`
  1149. is used to represent the columns in
  1150. a SELECT statement which may include duplicates.
  1151. A special subclass :class:`.DedupeColumnCollection` exists which instead
  1152. maintains SQLAlchemy's older behavior of not allowing duplicates; this
  1153. collection is used for schema level objects like :class:`_schema.Table`
  1154. and
  1155. :class:`.PrimaryKeyConstraint` where this deduping is helpful. The
  1156. :class:`.DedupeColumnCollection` class also has additional mutation methods
  1157. as the schema constructs have more use cases that require removal and
  1158. replacement of columns.
  1159. .. versionchanged:: 1.4 :class:`_expression.ColumnCollection`
  1160. now stores duplicate
  1161. column keys as well as the same column in multiple positions. The
  1162. :class:`.DedupeColumnCollection` class is added to maintain the
  1163. former behavior in those cases where deduplication as well as
  1164. additional replace/remove operations are needed.
  1165. """
  1166. __slots__ = "_collection", "_index", "_colset", "_proxy_index"
  1167. _collection: List[Tuple[_COLKEY, _COL_co, _ColumnMetrics[_COL_co]]]
  1168. _index: Dict[Union[None, str, int], Tuple[_COLKEY, _COL_co]]
  1169. _proxy_index: Dict[ColumnElement[Any], Set[_ColumnMetrics[_COL_co]]]
  1170. _colset: Set[_COL_co]
  1171. def __init__(
  1172. self, columns: Optional[Iterable[Tuple[_COLKEY, _COL_co]]] = None
  1173. ):
  1174. object.__setattr__(self, "_colset", set())
  1175. object.__setattr__(self, "_index", {})
  1176. object.__setattr__(
  1177. self, "_proxy_index", collections.defaultdict(util.OrderedSet)
  1178. )
  1179. object.__setattr__(self, "_collection", [])
  1180. if columns:
  1181. self._initial_populate(columns)
  1182. @util.preload_module("sqlalchemy.sql.elements")
  1183. def __clause_element__(self) -> ClauseList:
  1184. elements = util.preloaded.sql_elements
  1185. return elements.ClauseList(
  1186. _literal_as_text_role=roles.ColumnsClauseRole,
  1187. group=False,
  1188. *self._all_columns,
  1189. )
  1190. def _initial_populate(
  1191. self, iter_: Iterable[Tuple[_COLKEY, _COL_co]]
  1192. ) -> None:
  1193. self._populate_separate_keys(iter_)
  1194. @property
  1195. def _all_columns(self) -> List[_COL_co]:
  1196. return [col for (_, col, _) in self._collection]
  1197. def keys(self) -> List[_COLKEY]:
  1198. """Return a sequence of string key names for all columns in this
  1199. collection."""
  1200. return [k for (k, _, _) in self._collection]
  1201. def values(self) -> List[_COL_co]:
  1202. """Return a sequence of :class:`_sql.ColumnClause` or
  1203. :class:`_schema.Column` objects for all columns in this
  1204. collection."""
  1205. return [col for (_, col, _) in self._collection]
  1206. def items(self) -> List[Tuple[_COLKEY, _COL_co]]:
  1207. """Return a sequence of (key, column) tuples for all columns in this
  1208. collection each consisting of a string key name and a
  1209. :class:`_sql.ColumnClause` or
  1210. :class:`_schema.Column` object.
  1211. """
  1212. return [(k, col) for (k, col, _) in self._collection]
  1213. def __bool__(self) -> bool:
  1214. return bool(self._collection)
  1215. def __len__(self) -> int:
  1216. return len(self._collection)
  1217. def __iter__(self) -> Iterator[_COL_co]:
  1218. # turn to a list first to maintain over a course of changes
  1219. return iter([col for _, col, _ in self._collection])
  1220. @overload
  1221. def __getitem__(self, key: Union[str, int]) -> _COL_co: ...
  1222. @overload
  1223. def __getitem__(
  1224. self, key: Tuple[Union[str, int], ...]
  1225. ) -> ReadOnlyColumnCollection[_COLKEY, _COL_co]: ...
  1226. @overload
  1227. def __getitem__(
  1228. self, key: slice
  1229. ) -> ReadOnlyColumnCollection[_COLKEY, _COL_co]: ...
  1230. def __getitem__(
  1231. self, key: Union[str, int, slice, Tuple[Union[str, int], ...]]
  1232. ) -> Union[ReadOnlyColumnCollection[_COLKEY, _COL_co], _COL_co]:
  1233. try:
  1234. if isinstance(key, (tuple, slice)):
  1235. if isinstance(key, slice):
  1236. cols = (
  1237. (sub_key, col)
  1238. for (sub_key, col, _) in self._collection[key]
  1239. )
  1240. else:
  1241. cols = (self._index[sub_key] for sub_key in key)
  1242. return ColumnCollection(cols).as_readonly()
  1243. else:
  1244. return self._index[key][1]
  1245. except KeyError as err:
  1246. if isinstance(err.args[0], int):
  1247. raise IndexError(err.args[0]) from err
  1248. else:
  1249. raise
  1250. def __getattr__(self, key: str) -> _COL_co:
  1251. try:
  1252. return self._index[key][1]
  1253. except KeyError as err:
  1254. raise AttributeError(key) from err
  1255. def __contains__(self, key: str) -> bool:
  1256. if key not in self._index:
  1257. if not isinstance(key, str):
  1258. raise exc.ArgumentError(
  1259. "__contains__ requires a string argument"
  1260. )
  1261. return False
  1262. else:
  1263. return True
  1264. def compare(self, other: ColumnCollection[Any, Any]) -> bool:
  1265. """Compare this :class:`_expression.ColumnCollection` to another
  1266. based on the names of the keys"""
  1267. for l, r in zip_longest(self, other):
  1268. if l is not r:
  1269. return False
  1270. else:
  1271. return True
  1272. def __eq__(self, other: Any) -> bool:
  1273. return self.compare(other)
  1274. @overload
  1275. def get(self, key: str, default: None = None) -> Optional[_COL_co]: ...
  1276. @overload
  1277. def get(self, key: str, default: _COL) -> Union[_COL_co, _COL]: ...
  1278. def get(
  1279. self, key: str, default: Optional[_COL] = None
  1280. ) -> Optional[Union[_COL_co, _COL]]:
  1281. """Get a :class:`_sql.ColumnClause` or :class:`_schema.Column` object
  1282. based on a string key name from this
  1283. :class:`_expression.ColumnCollection`."""
  1284. if key in self._index:
  1285. return self._index[key][1]
  1286. else:
  1287. return default
  1288. def __str__(self) -> str:
  1289. return "%s(%s)" % (
  1290. self.__class__.__name__,
  1291. ", ".join(str(c) for c in self),
  1292. )
  1293. def __setitem__(self, key: str, value: Any) -> NoReturn:
  1294. raise NotImplementedError()
  1295. def __delitem__(self, key: str) -> NoReturn:
  1296. raise NotImplementedError()
  1297. def __setattr__(self, key: str, obj: Any) -> NoReturn:
  1298. raise NotImplementedError()
  1299. def clear(self) -> NoReturn:
  1300. """Dictionary clear() is not implemented for
  1301. :class:`_sql.ColumnCollection`."""
  1302. raise NotImplementedError()
  1303. def remove(self, column: Any) -> None:
  1304. raise NotImplementedError()
  1305. def update(self, iter_: Any) -> NoReturn:
  1306. """Dictionary update() is not implemented for
  1307. :class:`_sql.ColumnCollection`."""
  1308. raise NotImplementedError()
  1309. # https://github.com/python/mypy/issues/4266
  1310. __hash__ = None # type: ignore
  1311. def _populate_separate_keys(
  1312. self, iter_: Iterable[Tuple[_COLKEY, _COL_co]]
  1313. ) -> None:
  1314. """populate from an iterator of (key, column)"""
  1315. self._collection[:] = collection = [
  1316. (k, c, _ColumnMetrics(self, c)) for k, c in iter_
  1317. ]
  1318. self._colset.update(c._deannotate() for _, c, _ in collection)
  1319. self._index.update(
  1320. {idx: (k, c) for idx, (k, c, _) in enumerate(collection)}
  1321. )
  1322. self._index.update({k: (k, col) for k, col, _ in reversed(collection)})
  1323. def add(
  1324. self, column: ColumnElement[Any], key: Optional[_COLKEY] = None
  1325. ) -> None:
  1326. """Add a column to this :class:`_sql.ColumnCollection`.
  1327. .. note::
  1328. This method is **not normally used by user-facing code**, as the
  1329. :class:`_sql.ColumnCollection` is usually part of an existing
  1330. object such as a :class:`_schema.Table`. To add a
  1331. :class:`_schema.Column` to an existing :class:`_schema.Table`
  1332. object, use the :meth:`_schema.Table.append_column` method.
  1333. """
  1334. colkey: _COLKEY
  1335. if key is None:
  1336. colkey = column.key # type: ignore
  1337. else:
  1338. colkey = key
  1339. l = len(self._collection)
  1340. # don't really know how this part is supposed to work w/ the
  1341. # covariant thing
  1342. _column = cast(_COL_co, column)
  1343. self._collection.append(
  1344. (colkey, _column, _ColumnMetrics(self, _column))
  1345. )
  1346. self._colset.add(_column._deannotate())
  1347. self._index[l] = (colkey, _column)
  1348. if colkey not in self._index:
  1349. self._index[colkey] = (colkey, _column)
  1350. def __getstate__(self) -> Dict[str, Any]:
  1351. return {
  1352. "_collection": [(k, c) for k, c, _ in self._collection],
  1353. "_index": self._index,
  1354. }
  1355. def __setstate__(self, state: Dict[str, Any]) -> None:
  1356. object.__setattr__(self, "_index", state["_index"])
  1357. object.__setattr__(
  1358. self, "_proxy_index", collections.defaultdict(util.OrderedSet)
  1359. )
  1360. object.__setattr__(
  1361. self,
  1362. "_collection",
  1363. [
  1364. (k, c, _ColumnMetrics(self, c))
  1365. for (k, c) in state["_collection"]
  1366. ],
  1367. )
  1368. object.__setattr__(
  1369. self, "_colset", {col for k, col, _ in self._collection}
  1370. )
  1371. def contains_column(self, col: ColumnElement[Any]) -> bool:
  1372. """Checks if a column object exists in this collection"""
  1373. if col not in self._colset:
  1374. if isinstance(col, str):
  1375. raise exc.ArgumentError(
  1376. "contains_column cannot be used with string arguments. "
  1377. "Use ``col_name in table.c`` instead."
  1378. )
  1379. return False
  1380. else:
  1381. return True
  1382. def as_readonly(self) -> ReadOnlyColumnCollection[_COLKEY, _COL_co]:
  1383. """Return a "read only" form of this
  1384. :class:`_sql.ColumnCollection`."""
  1385. return ReadOnlyColumnCollection(self)
  1386. def _init_proxy_index(self):
  1387. """populate the "proxy index", if empty.
  1388. proxy index is added in 2.0 to provide more efficient operation
  1389. for the corresponding_column() method.
  1390. For reasons of both time to construct new .c collections as well as
  1391. memory conservation for large numbers of large .c collections, the
  1392. proxy_index is only filled if corresponding_column() is called. once
  1393. filled it stays that way, and new _ColumnMetrics objects created after
  1394. that point will populate it with new data. Note this case would be
  1395. unusual, if not nonexistent, as it means a .c collection is being
  1396. mutated after corresponding_column() were used, however it is tested in
  1397. test/base/test_utils.py.
  1398. """
  1399. pi = self._proxy_index
  1400. if pi:
  1401. return
  1402. for _, _, metrics in self._collection:
  1403. eps = metrics.column._expanded_proxy_set
  1404. for eps_col in eps:
  1405. pi[eps_col].add(metrics)
  1406. def corresponding_column(
  1407. self, column: _COL, require_embedded: bool = False
  1408. ) -> Optional[Union[_COL, _COL_co]]:
  1409. """Given a :class:`_expression.ColumnElement`, return the exported
  1410. :class:`_expression.ColumnElement` object from this
  1411. :class:`_expression.ColumnCollection`
  1412. which corresponds to that original :class:`_expression.ColumnElement`
  1413. via a common
  1414. ancestor column.
  1415. :param column: the target :class:`_expression.ColumnElement`
  1416. to be matched.
  1417. :param require_embedded: only return corresponding columns for
  1418. the given :class:`_expression.ColumnElement`, if the given
  1419. :class:`_expression.ColumnElement`
  1420. is actually present within a sub-element
  1421. of this :class:`_expression.Selectable`.
  1422. Normally the column will match if
  1423. it merely shares a common ancestor with one of the exported
  1424. columns of this :class:`_expression.Selectable`.
  1425. .. seealso::
  1426. :meth:`_expression.Selectable.corresponding_column`
  1427. - invokes this method
  1428. against the collection returned by
  1429. :attr:`_expression.Selectable.exported_columns`.
  1430. .. versionchanged:: 1.4 the implementation for ``corresponding_column``
  1431. was moved onto the :class:`_expression.ColumnCollection` itself.
  1432. """
  1433. # TODO: cython candidate
  1434. # don't dig around if the column is locally present
  1435. if column in self._colset:
  1436. return column
  1437. selected_intersection, selected_metrics = None, None
  1438. target_set = column.proxy_set
  1439. pi = self._proxy_index
  1440. if not pi:
  1441. self._init_proxy_index()
  1442. for current_metrics in (
  1443. mm for ts in target_set if ts in pi for mm in pi[ts]
  1444. ):
  1445. if not require_embedded or current_metrics.embedded(target_set):
  1446. if selected_metrics is None:
  1447. # no corresponding column yet, pick this one.
  1448. selected_metrics = current_metrics
  1449. continue
  1450. current_intersection = target_set.intersection(
  1451. current_metrics.column._expanded_proxy_set
  1452. )
  1453. if selected_intersection is None:
  1454. selected_intersection = target_set.intersection(
  1455. selected_metrics.column._expanded_proxy_set
  1456. )
  1457. if len(current_intersection) > len(selected_intersection):
  1458. # 'current' has a larger field of correspondence than
  1459. # 'selected'. i.e. selectable.c.a1_x->a1.c.x->table.c.x
  1460. # matches a1.c.x->table.c.x better than
  1461. # selectable.c.x->table.c.x does.
  1462. selected_metrics = current_metrics
  1463. selected_intersection = current_intersection
  1464. elif current_intersection == selected_intersection:
  1465. # they have the same field of correspondence. see
  1466. # which proxy_set has fewer columns in it, which
  1467. # indicates a closer relationship with the root
  1468. # column. Also take into account the "weight"
  1469. # attribute which CompoundSelect() uses to give
  1470. # higher precedence to columns based on vertical
  1471. # position in the compound statement, and discard
  1472. # columns that have no reference to the target
  1473. # column (also occurs with CompoundSelect)
  1474. selected_col_distance = sum(
  1475. [
  1476. sc._annotations.get("weight", 1)
  1477. for sc in (
  1478. selected_metrics.column._uncached_proxy_list()
  1479. )
  1480. if sc.shares_lineage(column)
  1481. ],
  1482. )
  1483. current_col_distance = sum(
  1484. [
  1485. sc._annotations.get("weight", 1)
  1486. for sc in (
  1487. current_metrics.column._uncached_proxy_list()
  1488. )
  1489. if sc.shares_lineage(column)
  1490. ],
  1491. )
  1492. if current_col_distance < selected_col_distance:
  1493. selected_metrics = current_metrics
  1494. selected_intersection = current_intersection
  1495. return selected_metrics.column if selected_metrics else None
  1496. _NAMEDCOL = TypeVar("_NAMEDCOL", bound="NamedColumn[Any]")
  1497. class DedupeColumnCollection(ColumnCollection[str, _NAMEDCOL]):
  1498. """A :class:`_expression.ColumnCollection`
  1499. that maintains deduplicating behavior.
  1500. This is useful by schema level objects such as :class:`_schema.Table` and
  1501. :class:`.PrimaryKeyConstraint`. The collection includes more
  1502. sophisticated mutator methods as well to suit schema objects which
  1503. require mutable column collections.
  1504. .. versionadded:: 1.4
  1505. """
  1506. def add( # type: ignore[override]
  1507. self, column: _NAMEDCOL, key: Optional[str] = None
  1508. ) -> None:
  1509. if key is not None and column.key != key:
  1510. raise exc.ArgumentError(
  1511. "DedupeColumnCollection requires columns be under "
  1512. "the same key as their .key"
  1513. )
  1514. key = column.key
  1515. if key is None:
  1516. raise exc.ArgumentError(
  1517. "Can't add unnamed column to column collection"
  1518. )
  1519. if key in self._index:
  1520. existing = self._index[key][1]
  1521. if existing is column:
  1522. return
  1523. self.replace(column)
  1524. # pop out memoized proxy_set as this
  1525. # operation may very well be occurring
  1526. # in a _make_proxy operation
  1527. util.memoized_property.reset(column, "proxy_set")
  1528. else:
  1529. self._append_new_column(key, column)
  1530. def _append_new_column(self, key: str, named_column: _NAMEDCOL) -> None:
  1531. l = len(self._collection)
  1532. self._collection.append(
  1533. (key, named_column, _ColumnMetrics(self, named_column))
  1534. )
  1535. self._colset.add(named_column._deannotate())
  1536. self._index[l] = (key, named_column)
  1537. self._index[key] = (key, named_column)
  1538. def _populate_separate_keys(
  1539. self, iter_: Iterable[Tuple[str, _NAMEDCOL]]
  1540. ) -> None:
  1541. """populate from an iterator of (key, column)"""
  1542. cols = list(iter_)
  1543. replace_col = []
  1544. for k, col in cols:
  1545. if col.key != k:
  1546. raise exc.ArgumentError(
  1547. "DedupeColumnCollection requires columns be under "
  1548. "the same key as their .key"
  1549. )
  1550. if col.name in self._index and col.key != col.name:
  1551. replace_col.append(col)
  1552. elif col.key in self._index:
  1553. replace_col.append(col)
  1554. else:
  1555. self._index[k] = (k, col)
  1556. self._collection.append((k, col, _ColumnMetrics(self, col)))
  1557. self._colset.update(c._deannotate() for (k, c, _) in self._collection)
  1558. self._index.update(
  1559. (idx, (k, c)) for idx, (k, c, _) in enumerate(self._collection)
  1560. )
  1561. for col in replace_col:
  1562. self.replace(col)
  1563. def extend(self, iter_: Iterable[_NAMEDCOL]) -> None:
  1564. self._populate_separate_keys((col.key, col) for col in iter_)
  1565. def remove(self, column: _NAMEDCOL) -> None:
  1566. if column not in self._colset:
  1567. raise ValueError(
  1568. "Can't remove column %r; column is not in this collection"
  1569. % column
  1570. )
  1571. del self._index[column.key]
  1572. self._colset.remove(column)
  1573. self._collection[:] = [
  1574. (k, c, metrics)
  1575. for (k, c, metrics) in self._collection
  1576. if c is not column
  1577. ]
  1578. for metrics in self._proxy_index.get(column, ()):
  1579. metrics.dispose(self)
  1580. self._index.update(
  1581. {idx: (k, col) for idx, (k, col, _) in enumerate(self._collection)}
  1582. )
  1583. # delete higher index
  1584. del self._index[len(self._collection)]
  1585. def replace(
  1586. self,
  1587. column: _NAMEDCOL,
  1588. extra_remove: Optional[Iterable[_NAMEDCOL]] = None,
  1589. ) -> None:
  1590. """add the given column to this collection, removing unaliased
  1591. versions of this column as well as existing columns with the
  1592. same key.
  1593. e.g.::
  1594. t = Table("sometable", metadata, Column("col1", Integer))
  1595. t.columns.replace(Column("col1", Integer, key="columnone"))
  1596. will remove the original 'col1' from the collection, and add
  1597. the new column under the name 'columnname'.
  1598. Used by schema.Column to override columns during table reflection.
  1599. """
  1600. if extra_remove:
  1601. remove_col = set(extra_remove)
  1602. else:
  1603. remove_col = set()
  1604. # remove up to two columns based on matches of name as well as key
  1605. if column.name in self._index and column.key != column.name:
  1606. other = self._index[column.name][1]
  1607. if other.name == other.key:
  1608. remove_col.add(other)
  1609. if column.key in self._index:
  1610. remove_col.add(self._index[column.key][1])
  1611. if not remove_col:
  1612. self._append_new_column(column.key, column)
  1613. return
  1614. new_cols: List[Tuple[str, _NAMEDCOL, _ColumnMetrics[_NAMEDCOL]]] = []
  1615. replaced = False
  1616. for k, col, metrics in self._collection:
  1617. if col in remove_col:
  1618. if not replaced:
  1619. replaced = True
  1620. new_cols.append(
  1621. (column.key, column, _ColumnMetrics(self, column))
  1622. )
  1623. else:
  1624. new_cols.append((k, col, metrics))
  1625. if remove_col:
  1626. self._colset.difference_update(remove_col)
  1627. for rc in remove_col:
  1628. for metrics in self._proxy_index.get(rc, ()):
  1629. metrics.dispose(self)
  1630. if not replaced:
  1631. new_cols.append((column.key, column, _ColumnMetrics(self, column)))
  1632. self._colset.add(column._deannotate())
  1633. self._collection[:] = new_cols
  1634. self._index.clear()
  1635. self._index.update(
  1636. {idx: (k, col) for idx, (k, col, _) in enumerate(self._collection)}
  1637. )
  1638. self._index.update({k: (k, col) for (k, col, _) in self._collection})
  1639. class ReadOnlyColumnCollection(
  1640. util.ReadOnlyContainer, ColumnCollection[_COLKEY, _COL_co]
  1641. ):
  1642. __slots__ = ("_parent",)
  1643. def __init__(self, collection):
  1644. object.__setattr__(self, "_parent", collection)
  1645. object.__setattr__(self, "_colset", collection._colset)
  1646. object.__setattr__(self, "_index", collection._index)
  1647. object.__setattr__(self, "_collection", collection._collection)
  1648. object.__setattr__(self, "_proxy_index", collection._proxy_index)
  1649. def __getstate__(self):
  1650. return {"_parent": self._parent}
  1651. def __setstate__(self, state):
  1652. parent = state["_parent"]
  1653. self.__init__(parent) # type: ignore
  1654. def add(self, column: Any, key: Any = ...) -> Any:
  1655. self._readonly()
  1656. def extend(self, elements: Any) -> NoReturn:
  1657. self._readonly()
  1658. def remove(self, item: Any) -> NoReturn:
  1659. self._readonly()
  1660. class ColumnSet(util.OrderedSet["ColumnClause[Any]"]):
  1661. def contains_column(self, col):
  1662. return col in self
  1663. def extend(self, cols):
  1664. for col in cols:
  1665. self.add(col)
  1666. def __eq__(self, other):
  1667. l = []
  1668. for c in other:
  1669. for local in self:
  1670. if c.shares_lineage(local):
  1671. l.append(c == local)
  1672. return elements.and_(*l)
  1673. def __hash__(self): # type: ignore[override]
  1674. return hash(tuple(x for x in self))
  1675. def _entity_namespace(
  1676. entity: Union[_HasEntityNamespace, ExternallyTraversible]
  1677. ) -> _EntityNamespace:
  1678. """Return the nearest .entity_namespace for the given entity.
  1679. If not immediately available, does an iterate to find a sub-element
  1680. that has one, if any.
  1681. """
  1682. try:
  1683. return cast(_HasEntityNamespace, entity).entity_namespace
  1684. except AttributeError:
  1685. for elem in visitors.iterate(cast(ExternallyTraversible, entity)):
  1686. if _is_has_entity_namespace(elem):
  1687. return elem.entity_namespace
  1688. else:
  1689. raise
  1690. def _entity_namespace_key(
  1691. entity: Union[_HasEntityNamespace, ExternallyTraversible],
  1692. key: str,
  1693. default: Union[SQLCoreOperations[Any], _NoArg] = NO_ARG,
  1694. ) -> SQLCoreOperations[Any]:
  1695. """Return an entry from an entity_namespace.
  1696. Raises :class:`_exc.InvalidRequestError` rather than attribute error
  1697. on not found.
  1698. """
  1699. try:
  1700. ns = _entity_namespace(entity)
  1701. if default is not NO_ARG:
  1702. return getattr(ns, key, default)
  1703. else:
  1704. return getattr(ns, key) # type: ignore
  1705. except AttributeError as err:
  1706. raise exc.InvalidRequestError(
  1707. 'Entity namespace for "%s" has no property "%s"' % (entity, key)
  1708. ) from err