bulk_persistence.py 68 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262126312641265126612671268126912701271127212731274127512761277127812791280128112821283128412851286128712881289129012911292129312941295129612971298129913001301130213031304130513061307130813091310131113121313131413151316131713181319132013211322132313241325132613271328132913301331133213331334133513361337133813391340134113421343134413451346134713481349135013511352135313541355135613571358135913601361136213631364136513661367136813691370137113721373137413751376137713781379138013811382138313841385138613871388138913901391139213931394139513961397139813991400140114021403140414051406140714081409141014111412141314141415141614171418141914201421142214231424142514261427142814291430143114321433143414351436143714381439144014411442144314441445144614471448144914501451145214531454145514561457145814591460146114621463146414651466146714681469147014711472147314741475147614771478147914801481148214831484148514861487148814891490149114921493149414951496149714981499150015011502150315041505150615071508150915101511151215131514151515161517151815191520152115221523152415251526152715281529153015311532153315341535153615371538153915401541154215431544154515461547154815491550155115521553155415551556155715581559156015611562156315641565156615671568156915701571157215731574157515761577157815791580158115821583158415851586158715881589159015911592159315941595159615971598159916001601160216031604160516061607160816091610161116121613161416151616161716181619162016211622162316241625162616271628162916301631163216331634163516361637163816391640164116421643164416451646164716481649165016511652165316541655165616571658165916601661166216631664166516661667166816691670167116721673167416751676167716781679168016811682168316841685168616871688168916901691169216931694169516961697169816991700170117021703170417051706170717081709171017111712171317141715171617171718171917201721172217231724172517261727172817291730173117321733173417351736173717381739174017411742174317441745174617471748174917501751175217531754175517561757175817591760176117621763176417651766176717681769177017711772177317741775177617771778177917801781178217831784178517861787178817891790179117921793179417951796179717981799180018011802180318041805180618071808180918101811181218131814181518161817181818191820182118221823182418251826182718281829183018311832183318341835183618371838183918401841184218431844184518461847184818491850185118521853185418551856185718581859186018611862186318641865186618671868186918701871187218731874187518761877187818791880188118821883188418851886188718881889189018911892189318941895189618971898189919001901190219031904190519061907190819091910191119121913191419151916191719181919192019211922192319241925192619271928192919301931193219331934193519361937193819391940194119421943194419451946194719481949195019511952195319541955195619571958195919601961196219631964196519661967196819691970197119721973197419751976197719781979198019811982198319841985198619871988198919901991199219931994199519961997199819992000200120022003200420052006200720082009201020112012201320142015201620172018201920202021202220232024202520262027202820292030203120322033203420352036203720382039204020412042204320442045204620472048
  1. # orm/bulk_persistence.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. # mypy: ignore-errors
  8. """additional ORM persistence classes related to "bulk" operations,
  9. specifically outside of the flush() process.
  10. """
  11. from __future__ import annotations
  12. from typing import Any
  13. from typing import cast
  14. from typing import Dict
  15. from typing import Iterable
  16. from typing import Optional
  17. from typing import overload
  18. from typing import TYPE_CHECKING
  19. from typing import TypeVar
  20. from typing import Union
  21. from . import attributes
  22. from . import context
  23. from . import evaluator
  24. from . import exc as orm_exc
  25. from . import loading
  26. from . import persistence
  27. from .base import NO_VALUE
  28. from .context import AbstractORMCompileState
  29. from .context import FromStatement
  30. from .context import ORMFromStatementCompileState
  31. from .context import QueryContext
  32. from .. import exc as sa_exc
  33. from .. import util
  34. from ..engine import Dialect
  35. from ..engine import result as _result
  36. from ..sql import coercions
  37. from ..sql import dml
  38. from ..sql import expression
  39. from ..sql import roles
  40. from ..sql import select
  41. from ..sql import sqltypes
  42. from ..sql.base import _entity_namespace_key
  43. from ..sql.base import CompileState
  44. from ..sql.base import Options
  45. from ..sql.dml import DeleteDMLState
  46. from ..sql.dml import InsertDMLState
  47. from ..sql.dml import UpdateDMLState
  48. from ..util import EMPTY_DICT
  49. from ..util.typing import Literal
  50. if TYPE_CHECKING:
  51. from ._typing import DMLStrategyArgument
  52. from ._typing import OrmExecuteOptionsParameter
  53. from ._typing import SynchronizeSessionArgument
  54. from .mapper import Mapper
  55. from .session import _BindArguments
  56. from .session import ORMExecuteState
  57. from .session import Session
  58. from .session import SessionTransaction
  59. from .state import InstanceState
  60. from ..engine import Connection
  61. from ..engine import cursor
  62. from ..engine.interfaces import _CoreAnyExecuteParams
  63. _O = TypeVar("_O", bound=object)
  64. @overload
  65. def _bulk_insert(
  66. mapper: Mapper[_O],
  67. mappings: Union[Iterable[InstanceState[_O]], Iterable[Dict[str, Any]]],
  68. session_transaction: SessionTransaction,
  69. isstates: bool,
  70. return_defaults: bool,
  71. render_nulls: bool,
  72. use_orm_insert_stmt: Literal[None] = ...,
  73. execution_options: Optional[OrmExecuteOptionsParameter] = ...,
  74. ) -> None:
  75. ...
  76. @overload
  77. def _bulk_insert(
  78. mapper: Mapper[_O],
  79. mappings: Union[Iterable[InstanceState[_O]], Iterable[Dict[str, Any]]],
  80. session_transaction: SessionTransaction,
  81. isstates: bool,
  82. return_defaults: bool,
  83. render_nulls: bool,
  84. use_orm_insert_stmt: Optional[dml.Insert] = ...,
  85. execution_options: Optional[OrmExecuteOptionsParameter] = ...,
  86. ) -> cursor.CursorResult[Any]:
  87. ...
  88. def _bulk_insert(
  89. mapper: Mapper[_O],
  90. mappings: Union[Iterable[InstanceState[_O]], Iterable[Dict[str, Any]]],
  91. session_transaction: SessionTransaction,
  92. isstates: bool,
  93. return_defaults: bool,
  94. render_nulls: bool,
  95. use_orm_insert_stmt: Optional[dml.Insert] = None,
  96. execution_options: Optional[OrmExecuteOptionsParameter] = None,
  97. ) -> Optional[cursor.CursorResult[Any]]:
  98. base_mapper = mapper.base_mapper
  99. if session_transaction.session.connection_callable:
  100. raise NotImplementedError(
  101. "connection_callable / per-instance sharding "
  102. "not supported in bulk_insert()"
  103. )
  104. if isstates:
  105. if return_defaults:
  106. states = [(state, state.dict) for state in mappings]
  107. mappings = [dict_ for (state, dict_) in states]
  108. else:
  109. mappings = [state.dict for state in mappings]
  110. else:
  111. mappings = [dict(m) for m in mappings]
  112. _expand_composites(mapper, mappings)
  113. connection = session_transaction.connection(base_mapper)
  114. return_result: Optional[cursor.CursorResult[Any]] = None
  115. mappers_to_run = [
  116. (table, mp)
  117. for table, mp in base_mapper._sorted_tables.items()
  118. if table in mapper._pks_by_table
  119. ]
  120. if return_defaults:
  121. # not used by new-style bulk inserts, only used for legacy
  122. bookkeeping = True
  123. elif len(mappers_to_run) > 1:
  124. # if we have more than one table, mapper to run where we will be
  125. # either horizontally splicing, or copying values between tables,
  126. # we need the "bookkeeping" / deterministic returning order
  127. bookkeeping = True
  128. else:
  129. bookkeeping = False
  130. for table, super_mapper in mappers_to_run:
  131. # find bindparams in the statement. For bulk, we don't really know if
  132. # a key in the params applies to a different table since we are
  133. # potentially inserting for multiple tables here; looking at the
  134. # bindparam() is a lot more direct. in most cases this will
  135. # use _generate_cache_key() which is memoized, although in practice
  136. # the ultimate statement that's executed is probably not the same
  137. # object so that memoization might not matter much.
  138. extra_bp_names = (
  139. [
  140. b.key
  141. for b in use_orm_insert_stmt._get_embedded_bindparams()
  142. if b.key in mappings[0]
  143. ]
  144. if use_orm_insert_stmt is not None
  145. else ()
  146. )
  147. records = (
  148. (
  149. None,
  150. state_dict,
  151. params,
  152. mapper,
  153. connection,
  154. value_params,
  155. has_all_pks,
  156. has_all_defaults,
  157. )
  158. for (
  159. state,
  160. state_dict,
  161. params,
  162. mp,
  163. conn,
  164. value_params,
  165. has_all_pks,
  166. has_all_defaults,
  167. ) in persistence._collect_insert_commands(
  168. table,
  169. ((None, mapping, mapper, connection) for mapping in mappings),
  170. bulk=True,
  171. return_defaults=bookkeeping,
  172. render_nulls=render_nulls,
  173. include_bulk_keys=extra_bp_names,
  174. )
  175. )
  176. result = persistence._emit_insert_statements(
  177. base_mapper,
  178. None,
  179. super_mapper,
  180. table,
  181. records,
  182. bookkeeping=bookkeeping,
  183. use_orm_insert_stmt=use_orm_insert_stmt,
  184. execution_options=execution_options,
  185. )
  186. if use_orm_insert_stmt is not None:
  187. if not use_orm_insert_stmt._returning or return_result is None:
  188. return_result = result
  189. elif result.returns_rows:
  190. assert bookkeeping
  191. return_result = return_result.splice_horizontally(result)
  192. if return_defaults and isstates:
  193. identity_cls = mapper._identity_class
  194. identity_props = [p.key for p in mapper._identity_key_props]
  195. for state, dict_ in states:
  196. state.key = (
  197. identity_cls,
  198. tuple([dict_[key] for key in identity_props]),
  199. )
  200. if use_orm_insert_stmt is not None:
  201. assert return_result is not None
  202. return return_result
  203. @overload
  204. def _bulk_update(
  205. mapper: Mapper[Any],
  206. mappings: Union[Iterable[InstanceState[_O]], Iterable[Dict[str, Any]]],
  207. session_transaction: SessionTransaction,
  208. isstates: bool,
  209. update_changed_only: bool,
  210. use_orm_update_stmt: Literal[None] = ...,
  211. enable_check_rowcount: bool = True,
  212. ) -> None:
  213. ...
  214. @overload
  215. def _bulk_update(
  216. mapper: Mapper[Any],
  217. mappings: Union[Iterable[InstanceState[_O]], Iterable[Dict[str, Any]]],
  218. session_transaction: SessionTransaction,
  219. isstates: bool,
  220. update_changed_only: bool,
  221. use_orm_update_stmt: Optional[dml.Update] = ...,
  222. enable_check_rowcount: bool = True,
  223. ) -> _result.Result[Any]:
  224. ...
  225. def _bulk_update(
  226. mapper: Mapper[Any],
  227. mappings: Union[Iterable[InstanceState[_O]], Iterable[Dict[str, Any]]],
  228. session_transaction: SessionTransaction,
  229. isstates: bool,
  230. update_changed_only: bool,
  231. use_orm_update_stmt: Optional[dml.Update] = None,
  232. enable_check_rowcount: bool = True,
  233. ) -> Optional[_result.Result[Any]]:
  234. base_mapper = mapper.base_mapper
  235. search_keys = mapper._primary_key_propkeys
  236. if mapper._version_id_prop:
  237. search_keys = {mapper._version_id_prop.key}.union(search_keys)
  238. def _changed_dict(mapper, state):
  239. return {
  240. k: v
  241. for k, v in state.dict.items()
  242. if k in state.committed_state or k in search_keys
  243. }
  244. if isstates:
  245. if update_changed_only:
  246. mappings = [_changed_dict(mapper, state) for state in mappings]
  247. else:
  248. mappings = [state.dict for state in mappings]
  249. else:
  250. mappings = [dict(m) for m in mappings]
  251. _expand_composites(mapper, mappings)
  252. if session_transaction.session.connection_callable:
  253. raise NotImplementedError(
  254. "connection_callable / per-instance sharding "
  255. "not supported in bulk_update()"
  256. )
  257. connection = session_transaction.connection(base_mapper)
  258. # find bindparams in the statement. see _bulk_insert for similar
  259. # notes for the insert case
  260. extra_bp_names = (
  261. [
  262. b.key
  263. for b in use_orm_update_stmt._get_embedded_bindparams()
  264. if b.key in mappings[0]
  265. ]
  266. if use_orm_update_stmt is not None
  267. else ()
  268. )
  269. for table, super_mapper in base_mapper._sorted_tables.items():
  270. if not mapper.isa(super_mapper) or table not in mapper._pks_by_table:
  271. continue
  272. records = persistence._collect_update_commands(
  273. None,
  274. table,
  275. (
  276. (
  277. None,
  278. mapping,
  279. mapper,
  280. connection,
  281. (
  282. mapping[mapper._version_id_prop.key]
  283. if mapper._version_id_prop
  284. else None
  285. ),
  286. )
  287. for mapping in mappings
  288. ),
  289. bulk=True,
  290. use_orm_update_stmt=use_orm_update_stmt,
  291. include_bulk_keys=extra_bp_names,
  292. )
  293. persistence._emit_update_statements(
  294. base_mapper,
  295. None,
  296. super_mapper,
  297. table,
  298. records,
  299. bookkeeping=False,
  300. use_orm_update_stmt=use_orm_update_stmt,
  301. enable_check_rowcount=enable_check_rowcount,
  302. )
  303. if use_orm_update_stmt is not None:
  304. return _result.null_result()
  305. def _expand_composites(mapper, mappings):
  306. composite_attrs = mapper.composites
  307. if not composite_attrs:
  308. return
  309. composite_keys = set(composite_attrs.keys())
  310. populators = {
  311. key: composite_attrs[key]._populate_composite_bulk_save_mappings_fn()
  312. for key in composite_keys
  313. }
  314. for mapping in mappings:
  315. for key in composite_keys.intersection(mapping):
  316. populators[key](mapping)
  317. class ORMDMLState(AbstractORMCompileState):
  318. is_dml_returning = True
  319. from_statement_ctx: Optional[ORMFromStatementCompileState] = None
  320. @classmethod
  321. def _get_orm_crud_kv_pairs(
  322. cls, mapper, statement, kv_iterator, needs_to_be_cacheable
  323. ):
  324. core_get_crud_kv_pairs = UpdateDMLState._get_crud_kv_pairs
  325. for k, v in kv_iterator:
  326. k = coercions.expect(roles.DMLColumnRole, k)
  327. if isinstance(k, str):
  328. desc = _entity_namespace_key(mapper, k, default=NO_VALUE)
  329. if desc is NO_VALUE:
  330. yield (
  331. coercions.expect(roles.DMLColumnRole, k),
  332. coercions.expect(
  333. roles.ExpressionElementRole,
  334. v,
  335. type_=sqltypes.NullType(),
  336. is_crud=True,
  337. )
  338. if needs_to_be_cacheable
  339. else v,
  340. )
  341. else:
  342. yield from core_get_crud_kv_pairs(
  343. statement,
  344. desc._bulk_update_tuples(v),
  345. needs_to_be_cacheable,
  346. )
  347. elif "entity_namespace" in k._annotations:
  348. k_anno = k._annotations
  349. attr = _entity_namespace_key(
  350. k_anno["entity_namespace"], k_anno["proxy_key"]
  351. )
  352. yield from core_get_crud_kv_pairs(
  353. statement,
  354. attr._bulk_update_tuples(v),
  355. needs_to_be_cacheable,
  356. )
  357. else:
  358. yield (
  359. k,
  360. v
  361. if not needs_to_be_cacheable
  362. else coercions.expect(
  363. roles.ExpressionElementRole,
  364. v,
  365. type_=sqltypes.NullType(),
  366. is_crud=True,
  367. ),
  368. )
  369. @classmethod
  370. def _get_multi_crud_kv_pairs(cls, statement, kv_iterator):
  371. plugin_subject = statement._propagate_attrs["plugin_subject"]
  372. if not plugin_subject or not plugin_subject.mapper:
  373. return UpdateDMLState._get_multi_crud_kv_pairs(
  374. statement, kv_iterator
  375. )
  376. return [
  377. dict(
  378. cls._get_orm_crud_kv_pairs(
  379. plugin_subject.mapper, statement, value_dict.items(), False
  380. )
  381. )
  382. for value_dict in kv_iterator
  383. ]
  384. @classmethod
  385. def _get_crud_kv_pairs(cls, statement, kv_iterator, needs_to_be_cacheable):
  386. assert (
  387. needs_to_be_cacheable
  388. ), "no test coverage for needs_to_be_cacheable=False"
  389. plugin_subject = statement._propagate_attrs["plugin_subject"]
  390. if not plugin_subject or not plugin_subject.mapper:
  391. return UpdateDMLState._get_crud_kv_pairs(
  392. statement, kv_iterator, needs_to_be_cacheable
  393. )
  394. return list(
  395. cls._get_orm_crud_kv_pairs(
  396. plugin_subject.mapper,
  397. statement,
  398. kv_iterator,
  399. needs_to_be_cacheable,
  400. )
  401. )
  402. @classmethod
  403. def get_entity_description(cls, statement):
  404. ext_info = statement.table._annotations["parententity"]
  405. mapper = ext_info.mapper
  406. if ext_info.is_aliased_class:
  407. _label_name = ext_info.name
  408. else:
  409. _label_name = mapper.class_.__name__
  410. return {
  411. "name": _label_name,
  412. "type": mapper.class_,
  413. "expr": ext_info.entity,
  414. "entity": ext_info.entity,
  415. "table": mapper.local_table,
  416. }
  417. @classmethod
  418. def get_returning_column_descriptions(cls, statement):
  419. def _ent_for_col(c):
  420. return c._annotations.get("parententity", None)
  421. def _attr_for_col(c, ent):
  422. if ent is None:
  423. return c
  424. proxy_key = c._annotations.get("proxy_key", None)
  425. if not proxy_key:
  426. return c
  427. else:
  428. return getattr(ent.entity, proxy_key, c)
  429. return [
  430. {
  431. "name": c.key,
  432. "type": c.type,
  433. "expr": _attr_for_col(c, ent),
  434. "aliased": ent.is_aliased_class,
  435. "entity": ent.entity,
  436. }
  437. for c, ent in [
  438. (c, _ent_for_col(c)) for c in statement._all_selected_columns
  439. ]
  440. ]
  441. def _setup_orm_returning(
  442. self,
  443. compiler,
  444. orm_level_statement,
  445. dml_level_statement,
  446. dml_mapper,
  447. *,
  448. use_supplemental_cols=True,
  449. ):
  450. """establish ORM column handlers for an INSERT, UPDATE, or DELETE
  451. which uses explicit returning().
  452. called within compilation level create_for_statement.
  453. The _return_orm_returning() method then receives the Result
  454. after the statement was executed, and applies ORM loading to the
  455. state that we first established here.
  456. """
  457. if orm_level_statement._returning:
  458. fs = FromStatement(
  459. orm_level_statement._returning,
  460. dml_level_statement,
  461. _adapt_on_names=False,
  462. )
  463. fs = fs.execution_options(**orm_level_statement._execution_options)
  464. fs = fs.options(*orm_level_statement._with_options)
  465. self.select_statement = fs
  466. self.from_statement_ctx = (
  467. fsc
  468. ) = ORMFromStatementCompileState.create_for_statement(fs, compiler)
  469. fsc.setup_dml_returning_compile_state(dml_mapper)
  470. dml_level_statement = dml_level_statement._generate()
  471. dml_level_statement._returning = ()
  472. cols_to_return = [c for c in fsc.primary_columns if c is not None]
  473. # since we are splicing result sets together, make sure there
  474. # are columns of some kind returned in each result set
  475. if not cols_to_return:
  476. cols_to_return.extend(dml_mapper.primary_key)
  477. if use_supplemental_cols:
  478. dml_level_statement = dml_level_statement.return_defaults(
  479. # this is a little weird looking, but by passing
  480. # primary key as the main list of cols, this tells
  481. # return_defaults to omit server-default cols (and
  482. # actually all cols, due to some weird thing we should
  483. # clean up in crud.py).
  484. # Since we have cols_to_return, just return what we asked
  485. # for (plus primary key, which ORM persistence needs since
  486. # we likely set bookkeeping=True here, which is another
  487. # whole thing...). We dont want to clutter the
  488. # statement up with lots of other cols the user didn't
  489. # ask for. see #9685
  490. *dml_mapper.primary_key,
  491. supplemental_cols=cols_to_return,
  492. )
  493. else:
  494. dml_level_statement = dml_level_statement.returning(
  495. *cols_to_return
  496. )
  497. return dml_level_statement
  498. @classmethod
  499. def _return_orm_returning(
  500. cls,
  501. session,
  502. statement,
  503. params,
  504. execution_options,
  505. bind_arguments,
  506. result,
  507. ):
  508. execution_context = result.context
  509. compile_state = execution_context.compiled.compile_state
  510. if (
  511. compile_state.from_statement_ctx
  512. and not compile_state.from_statement_ctx.compile_options._is_star
  513. ):
  514. load_options = execution_options.get(
  515. "_sa_orm_load_options", QueryContext.default_load_options
  516. )
  517. querycontext = QueryContext(
  518. compile_state.from_statement_ctx,
  519. compile_state.select_statement,
  520. params,
  521. session,
  522. load_options,
  523. execution_options,
  524. bind_arguments,
  525. )
  526. return loading.instances(result, querycontext)
  527. else:
  528. return result
  529. class BulkUDCompileState(ORMDMLState):
  530. class default_update_options(Options):
  531. _dml_strategy: DMLStrategyArgument = "auto"
  532. _synchronize_session: SynchronizeSessionArgument = "auto"
  533. _can_use_returning: bool = False
  534. _is_delete_using: bool = False
  535. _is_update_from: bool = False
  536. _autoflush: bool = True
  537. _subject_mapper: Optional[Mapper[Any]] = None
  538. _resolved_values = EMPTY_DICT
  539. _eval_condition = None
  540. _matched_rows = None
  541. _identity_token = None
  542. @classmethod
  543. def can_use_returning(
  544. cls,
  545. dialect: Dialect,
  546. mapper: Mapper[Any],
  547. *,
  548. is_multitable: bool = False,
  549. is_update_from: bool = False,
  550. is_delete_using: bool = False,
  551. is_executemany: bool = False,
  552. ) -> bool:
  553. raise NotImplementedError()
  554. @classmethod
  555. def orm_pre_session_exec(
  556. cls,
  557. session,
  558. statement,
  559. params,
  560. execution_options,
  561. bind_arguments,
  562. is_pre_event,
  563. ):
  564. (
  565. update_options,
  566. execution_options,
  567. ) = BulkUDCompileState.default_update_options.from_execution_options(
  568. "_sa_orm_update_options",
  569. {
  570. "synchronize_session",
  571. "autoflush",
  572. "identity_token",
  573. "is_delete_using",
  574. "is_update_from",
  575. "dml_strategy",
  576. },
  577. execution_options,
  578. statement._execution_options,
  579. )
  580. bind_arguments["clause"] = statement
  581. try:
  582. plugin_subject = statement._propagate_attrs["plugin_subject"]
  583. except KeyError:
  584. assert False, "statement had 'orm' plugin but no plugin_subject"
  585. else:
  586. if plugin_subject:
  587. bind_arguments["mapper"] = plugin_subject.mapper
  588. update_options += {"_subject_mapper": plugin_subject.mapper}
  589. if "parententity" not in statement.table._annotations:
  590. update_options += {"_dml_strategy": "core_only"}
  591. elif not isinstance(params, list):
  592. if update_options._dml_strategy == "auto":
  593. update_options += {"_dml_strategy": "orm"}
  594. elif update_options._dml_strategy == "bulk":
  595. raise sa_exc.InvalidRequestError(
  596. 'Can\'t use "bulk" ORM insert strategy without '
  597. "passing separate parameters"
  598. )
  599. else:
  600. if update_options._dml_strategy == "auto":
  601. update_options += {"_dml_strategy": "bulk"}
  602. sync = update_options._synchronize_session
  603. if sync is not None:
  604. if sync not in ("auto", "evaluate", "fetch", False):
  605. raise sa_exc.ArgumentError(
  606. "Valid strategies for session synchronization "
  607. "are 'auto', 'evaluate', 'fetch', False"
  608. )
  609. if update_options._dml_strategy == "bulk" and sync == "fetch":
  610. raise sa_exc.InvalidRequestError(
  611. "The 'fetch' synchronization strategy is not available "
  612. "for 'bulk' ORM updates (i.e. multiple parameter sets)"
  613. )
  614. if not is_pre_event:
  615. if update_options._autoflush:
  616. session._autoflush()
  617. if update_options._dml_strategy == "orm":
  618. if update_options._synchronize_session == "auto":
  619. update_options = cls._do_pre_synchronize_auto(
  620. session,
  621. statement,
  622. params,
  623. execution_options,
  624. bind_arguments,
  625. update_options,
  626. )
  627. elif update_options._synchronize_session == "evaluate":
  628. update_options = cls._do_pre_synchronize_evaluate(
  629. session,
  630. statement,
  631. params,
  632. execution_options,
  633. bind_arguments,
  634. update_options,
  635. )
  636. elif update_options._synchronize_session == "fetch":
  637. update_options = cls._do_pre_synchronize_fetch(
  638. session,
  639. statement,
  640. params,
  641. execution_options,
  642. bind_arguments,
  643. update_options,
  644. )
  645. elif update_options._dml_strategy == "bulk":
  646. if update_options._synchronize_session == "auto":
  647. update_options += {"_synchronize_session": "evaluate"}
  648. # indicators from the "pre exec" step that are then
  649. # added to the DML statement, which will also be part of the cache
  650. # key. The compile level create_for_statement() method will then
  651. # consume these at compiler time.
  652. statement = statement._annotate(
  653. {
  654. "synchronize_session": update_options._synchronize_session,
  655. "is_delete_using": update_options._is_delete_using,
  656. "is_update_from": update_options._is_update_from,
  657. "dml_strategy": update_options._dml_strategy,
  658. "can_use_returning": update_options._can_use_returning,
  659. }
  660. )
  661. return (
  662. statement,
  663. util.immutabledict(execution_options).union(
  664. {"_sa_orm_update_options": update_options}
  665. ),
  666. )
  667. @classmethod
  668. def orm_setup_cursor_result(
  669. cls,
  670. session,
  671. statement,
  672. params,
  673. execution_options,
  674. bind_arguments,
  675. result,
  676. ):
  677. # this stage of the execution is called after the
  678. # do_orm_execute event hook. meaning for an extension like
  679. # horizontal sharding, this step happens *within* the horizontal
  680. # sharding event handler which calls session.execute() re-entrantly
  681. # and will occur for each backend individually.
  682. # the sharding extension then returns its own merged result from the
  683. # individual ones we return here.
  684. update_options = execution_options["_sa_orm_update_options"]
  685. if update_options._dml_strategy == "orm":
  686. if update_options._synchronize_session == "evaluate":
  687. cls._do_post_synchronize_evaluate(
  688. session, statement, result, update_options
  689. )
  690. elif update_options._synchronize_session == "fetch":
  691. cls._do_post_synchronize_fetch(
  692. session, statement, result, update_options
  693. )
  694. elif update_options._dml_strategy == "bulk":
  695. if update_options._synchronize_session == "evaluate":
  696. cls._do_post_synchronize_bulk_evaluate(
  697. session, params, result, update_options
  698. )
  699. return result
  700. return cls._return_orm_returning(
  701. session,
  702. statement,
  703. params,
  704. execution_options,
  705. bind_arguments,
  706. result,
  707. )
  708. @classmethod
  709. def _adjust_for_extra_criteria(cls, global_attributes, ext_info):
  710. """Apply extra criteria filtering.
  711. For all distinct single-table-inheritance mappers represented in the
  712. table being updated or deleted, produce additional WHERE criteria such
  713. that only the appropriate subtypes are selected from the total results.
  714. Additionally, add WHERE criteria originating from LoaderCriteriaOptions
  715. collected from the statement.
  716. """
  717. return_crit = ()
  718. adapter = ext_info._adapter if ext_info.is_aliased_class else None
  719. if (
  720. "additional_entity_criteria",
  721. ext_info.mapper,
  722. ) in global_attributes:
  723. return_crit += tuple(
  724. ae._resolve_where_criteria(ext_info)
  725. for ae in global_attributes[
  726. ("additional_entity_criteria", ext_info.mapper)
  727. ]
  728. if ae.include_aliases or ae.entity is ext_info
  729. )
  730. if ext_info.mapper._single_table_criterion is not None:
  731. return_crit += (ext_info.mapper._single_table_criterion,)
  732. if adapter:
  733. return_crit = tuple(adapter.traverse(crit) for crit in return_crit)
  734. return return_crit
  735. @classmethod
  736. def _interpret_returning_rows(cls, mapper, rows):
  737. """translate from local inherited table columns to base mapper
  738. primary key columns.
  739. Joined inheritance mappers always establish the primary key in terms of
  740. the base table. When we UPDATE a sub-table, we can only get
  741. RETURNING for the sub-table's columns.
  742. Here, we create a lookup from the local sub table's primary key
  743. columns to the base table PK columns so that we can get identity
  744. key values from RETURNING that's against the joined inheritance
  745. sub-table.
  746. the complexity here is to support more than one level deep of
  747. inheritance, where we have to link columns to each other across
  748. the inheritance hierarchy.
  749. """
  750. if mapper.local_table is not mapper.base_mapper.local_table:
  751. return rows
  752. # this starts as a mapping of
  753. # local_pk_col: local_pk_col.
  754. # we will then iteratively rewrite the "value" of the dict with
  755. # each successive superclass column
  756. local_pk_to_base_pk = {pk: pk for pk in mapper.local_table.primary_key}
  757. for mp in mapper.iterate_to_root():
  758. if mp.inherits is None:
  759. break
  760. elif mp.local_table is mp.inherits.local_table:
  761. continue
  762. t_to_e = dict(mp._table_to_equated[mp.inherits.local_table])
  763. col_to_col = {sub_pk: super_pk for super_pk, sub_pk in t_to_e[mp]}
  764. for pk, super_ in local_pk_to_base_pk.items():
  765. local_pk_to_base_pk[pk] = col_to_col[super_]
  766. lookup = {
  767. local_pk_to_base_pk[lpk]: idx
  768. for idx, lpk in enumerate(mapper.local_table.primary_key)
  769. }
  770. primary_key_convert = [
  771. lookup[bpk] for bpk in mapper.base_mapper.primary_key
  772. ]
  773. return [tuple(row[idx] for idx in primary_key_convert) for row in rows]
  774. @classmethod
  775. def _get_matched_objects_on_criteria(cls, update_options, states):
  776. mapper = update_options._subject_mapper
  777. eval_condition = update_options._eval_condition
  778. raw_data = [
  779. (state.obj(), state, state.dict)
  780. for state in states
  781. if state.mapper.isa(mapper) and not state.expired
  782. ]
  783. identity_token = update_options._identity_token
  784. if identity_token is not None:
  785. raw_data = [
  786. (obj, state, dict_)
  787. for obj, state, dict_ in raw_data
  788. if state.identity_token == identity_token
  789. ]
  790. result = []
  791. for obj, state, dict_ in raw_data:
  792. evaled_condition = eval_condition(obj)
  793. # caution: don't use "in ()" or == here, _EXPIRE_OBJECT
  794. # evaluates as True for all comparisons
  795. if (
  796. evaled_condition is True
  797. or evaled_condition is evaluator._EXPIRED_OBJECT
  798. ):
  799. result.append(
  800. (
  801. obj,
  802. state,
  803. dict_,
  804. evaled_condition is evaluator._EXPIRED_OBJECT,
  805. )
  806. )
  807. return result
  808. @classmethod
  809. def _eval_condition_from_statement(cls, update_options, statement):
  810. mapper = update_options._subject_mapper
  811. target_cls = mapper.class_
  812. evaluator_compiler = evaluator._EvaluatorCompiler(target_cls)
  813. crit = ()
  814. if statement._where_criteria:
  815. crit += statement._where_criteria
  816. global_attributes = {}
  817. for opt in statement._with_options:
  818. if opt._is_criteria_option:
  819. opt.get_global_criteria(global_attributes)
  820. if global_attributes:
  821. crit += cls._adjust_for_extra_criteria(global_attributes, mapper)
  822. if crit:
  823. eval_condition = evaluator_compiler.process(*crit)
  824. else:
  825. # workaround for mypy https://github.com/python/mypy/issues/14027
  826. def _eval_condition(obj):
  827. return True
  828. eval_condition = _eval_condition
  829. return eval_condition
  830. @classmethod
  831. def _do_pre_synchronize_auto(
  832. cls,
  833. session,
  834. statement,
  835. params,
  836. execution_options,
  837. bind_arguments,
  838. update_options,
  839. ):
  840. """setup auto sync strategy
  841. "auto" checks if we can use "evaluate" first, then falls back
  842. to "fetch"
  843. evaluate is vastly more efficient for the common case
  844. where session is empty, only has a few objects, and the UPDATE
  845. statement can potentially match thousands/millions of rows.
  846. OTOH more complex criteria that fails to work with "evaluate"
  847. we would hope usually correlates with fewer net rows.
  848. """
  849. try:
  850. eval_condition = cls._eval_condition_from_statement(
  851. update_options, statement
  852. )
  853. except evaluator.UnevaluatableError:
  854. pass
  855. else:
  856. return update_options + {
  857. "_eval_condition": eval_condition,
  858. "_synchronize_session": "evaluate",
  859. }
  860. update_options += {"_synchronize_session": "fetch"}
  861. return cls._do_pre_synchronize_fetch(
  862. session,
  863. statement,
  864. params,
  865. execution_options,
  866. bind_arguments,
  867. update_options,
  868. )
  869. @classmethod
  870. def _do_pre_synchronize_evaluate(
  871. cls,
  872. session,
  873. statement,
  874. params,
  875. execution_options,
  876. bind_arguments,
  877. update_options,
  878. ):
  879. try:
  880. eval_condition = cls._eval_condition_from_statement(
  881. update_options, statement
  882. )
  883. except evaluator.UnevaluatableError as err:
  884. raise sa_exc.InvalidRequestError(
  885. 'Could not evaluate current criteria in Python: "%s". '
  886. "Specify 'fetch' or False for the "
  887. "synchronize_session execution option." % err
  888. ) from err
  889. return update_options + {
  890. "_eval_condition": eval_condition,
  891. }
  892. @classmethod
  893. def _get_resolved_values(cls, mapper, statement):
  894. if statement._multi_values:
  895. return []
  896. elif statement._ordered_values:
  897. return list(statement._ordered_values)
  898. elif statement._values:
  899. return list(statement._values.items())
  900. else:
  901. return []
  902. @classmethod
  903. def _resolved_keys_as_propnames(cls, mapper, resolved_values):
  904. values = []
  905. for k, v in resolved_values:
  906. if mapper and isinstance(k, expression.ColumnElement):
  907. try:
  908. attr = mapper._columntoproperty[k]
  909. except orm_exc.UnmappedColumnError:
  910. pass
  911. else:
  912. values.append((attr.key, v))
  913. else:
  914. raise sa_exc.InvalidRequestError(
  915. "Attribute name not found, can't be "
  916. "synchronized back to objects: %r" % k
  917. )
  918. return values
  919. @classmethod
  920. def _do_pre_synchronize_fetch(
  921. cls,
  922. session,
  923. statement,
  924. params,
  925. execution_options,
  926. bind_arguments,
  927. update_options,
  928. ):
  929. mapper = update_options._subject_mapper
  930. select_stmt = (
  931. select(*(mapper.primary_key + (mapper.select_identity_token,)))
  932. .select_from(mapper)
  933. .options(*statement._with_options)
  934. )
  935. select_stmt._where_criteria = statement._where_criteria
  936. # conditionally run the SELECT statement for pre-fetch, testing the
  937. # "bind" for if we can use RETURNING or not using the do_orm_execute
  938. # event. If RETURNING is available, the do_orm_execute event
  939. # will cancel the SELECT from being actually run.
  940. #
  941. # The way this is organized seems strange, why don't we just
  942. # call can_use_returning() before invoking the statement and get
  943. # answer?, why does this go through the whole execute phase using an
  944. # event? Answer: because we are integrating with extensions such
  945. # as the horizontal sharding extention that "multiplexes" an individual
  946. # statement run through multiple engines, and it uses
  947. # do_orm_execute() to do that.
  948. can_use_returning = None
  949. def skip_for_returning(orm_context: ORMExecuteState) -> Any:
  950. bind = orm_context.session.get_bind(**orm_context.bind_arguments)
  951. nonlocal can_use_returning
  952. per_bind_result = cls.can_use_returning(
  953. bind.dialect,
  954. mapper,
  955. is_update_from=update_options._is_update_from,
  956. is_delete_using=update_options._is_delete_using,
  957. is_executemany=orm_context.is_executemany,
  958. )
  959. if can_use_returning is not None:
  960. if can_use_returning != per_bind_result:
  961. raise sa_exc.InvalidRequestError(
  962. "For synchronize_session='fetch', can't mix multiple "
  963. "backends where some support RETURNING and others "
  964. "don't"
  965. )
  966. elif orm_context.is_executemany and not per_bind_result:
  967. raise sa_exc.InvalidRequestError(
  968. "For synchronize_session='fetch', can't use multiple "
  969. "parameter sets in ORM mode, which this backend does not "
  970. "support with RETURNING"
  971. )
  972. else:
  973. can_use_returning = per_bind_result
  974. if per_bind_result:
  975. return _result.null_result()
  976. else:
  977. return None
  978. result = session.execute(
  979. select_stmt,
  980. params,
  981. execution_options=execution_options,
  982. bind_arguments=bind_arguments,
  983. _add_event=skip_for_returning,
  984. )
  985. matched_rows = result.fetchall()
  986. return update_options + {
  987. "_matched_rows": matched_rows,
  988. "_can_use_returning": can_use_returning,
  989. }
  990. @CompileState.plugin_for("orm", "insert")
  991. class BulkORMInsert(ORMDMLState, InsertDMLState):
  992. class default_insert_options(Options):
  993. _dml_strategy: DMLStrategyArgument = "auto"
  994. _render_nulls: bool = False
  995. _return_defaults: bool = False
  996. _subject_mapper: Optional[Mapper[Any]] = None
  997. _autoflush: bool = True
  998. _populate_existing: bool = False
  999. select_statement: Optional[FromStatement] = None
  1000. @classmethod
  1001. def orm_pre_session_exec(
  1002. cls,
  1003. session,
  1004. statement,
  1005. params,
  1006. execution_options,
  1007. bind_arguments,
  1008. is_pre_event,
  1009. ):
  1010. (
  1011. insert_options,
  1012. execution_options,
  1013. ) = BulkORMInsert.default_insert_options.from_execution_options(
  1014. "_sa_orm_insert_options",
  1015. {"dml_strategy", "autoflush", "populate_existing", "render_nulls"},
  1016. execution_options,
  1017. statement._execution_options,
  1018. )
  1019. bind_arguments["clause"] = statement
  1020. try:
  1021. plugin_subject = statement._propagate_attrs["plugin_subject"]
  1022. except KeyError:
  1023. assert False, "statement had 'orm' plugin but no plugin_subject"
  1024. else:
  1025. if plugin_subject:
  1026. bind_arguments["mapper"] = plugin_subject.mapper
  1027. insert_options += {"_subject_mapper": plugin_subject.mapper}
  1028. if not params:
  1029. if insert_options._dml_strategy == "auto":
  1030. insert_options += {"_dml_strategy": "orm"}
  1031. elif insert_options._dml_strategy == "bulk":
  1032. raise sa_exc.InvalidRequestError(
  1033. 'Can\'t use "bulk" ORM insert strategy without '
  1034. "passing separate parameters"
  1035. )
  1036. else:
  1037. if insert_options._dml_strategy == "auto":
  1038. insert_options += {"_dml_strategy": "bulk"}
  1039. if insert_options._dml_strategy != "raw":
  1040. # for ORM object loading, like ORMContext, we have to disable
  1041. # result set adapt_to_context, because we will be generating a
  1042. # new statement with specific columns that's cached inside of
  1043. # an ORMFromStatementCompileState, which we will re-use for
  1044. # each result.
  1045. if not execution_options:
  1046. execution_options = context._orm_load_exec_options
  1047. else:
  1048. execution_options = execution_options.union(
  1049. context._orm_load_exec_options
  1050. )
  1051. if not is_pre_event and insert_options._autoflush:
  1052. session._autoflush()
  1053. statement = statement._annotate(
  1054. {"dml_strategy": insert_options._dml_strategy}
  1055. )
  1056. return (
  1057. statement,
  1058. util.immutabledict(execution_options).union(
  1059. {"_sa_orm_insert_options": insert_options}
  1060. ),
  1061. )
  1062. @classmethod
  1063. def orm_execute_statement(
  1064. cls,
  1065. session: Session,
  1066. statement: dml.Insert,
  1067. params: _CoreAnyExecuteParams,
  1068. execution_options: OrmExecuteOptionsParameter,
  1069. bind_arguments: _BindArguments,
  1070. conn: Connection,
  1071. ) -> _result.Result:
  1072. insert_options = execution_options.get(
  1073. "_sa_orm_insert_options", cls.default_insert_options
  1074. )
  1075. if insert_options._dml_strategy not in (
  1076. "raw",
  1077. "bulk",
  1078. "orm",
  1079. "auto",
  1080. ):
  1081. raise sa_exc.ArgumentError(
  1082. "Valid strategies for ORM insert strategy "
  1083. "are 'raw', 'orm', 'bulk', 'auto"
  1084. )
  1085. result: _result.Result[Any]
  1086. if insert_options._dml_strategy == "raw":
  1087. result = conn.execute(
  1088. statement, params or {}, execution_options=execution_options
  1089. )
  1090. return result
  1091. if insert_options._dml_strategy == "bulk":
  1092. mapper = insert_options._subject_mapper
  1093. if (
  1094. statement._post_values_clause is not None
  1095. and mapper._multiple_persistence_tables
  1096. ):
  1097. raise sa_exc.InvalidRequestError(
  1098. "bulk INSERT with a 'post values' clause "
  1099. "(typically upsert) not supported for multi-table "
  1100. f"mapper {mapper}"
  1101. )
  1102. assert mapper is not None
  1103. assert session._transaction is not None
  1104. result = _bulk_insert(
  1105. mapper,
  1106. cast(
  1107. "Iterable[Dict[str, Any]]",
  1108. [params] if isinstance(params, dict) else params,
  1109. ),
  1110. session._transaction,
  1111. isstates=False,
  1112. return_defaults=insert_options._return_defaults,
  1113. render_nulls=insert_options._render_nulls,
  1114. use_orm_insert_stmt=statement,
  1115. execution_options=execution_options,
  1116. )
  1117. elif insert_options._dml_strategy == "orm":
  1118. result = conn.execute(
  1119. statement, params or {}, execution_options=execution_options
  1120. )
  1121. else:
  1122. raise AssertionError()
  1123. if not bool(statement._returning):
  1124. return result
  1125. if insert_options._populate_existing:
  1126. load_options = execution_options.get(
  1127. "_sa_orm_load_options", QueryContext.default_load_options
  1128. )
  1129. load_options += {"_populate_existing": True}
  1130. execution_options = execution_options.union(
  1131. {"_sa_orm_load_options": load_options}
  1132. )
  1133. return cls._return_orm_returning(
  1134. session,
  1135. statement,
  1136. params,
  1137. execution_options,
  1138. bind_arguments,
  1139. result,
  1140. )
  1141. @classmethod
  1142. def create_for_statement(cls, statement, compiler, **kw) -> BulkORMInsert:
  1143. self = cast(
  1144. BulkORMInsert,
  1145. super().create_for_statement(statement, compiler, **kw),
  1146. )
  1147. if compiler is not None:
  1148. toplevel = not compiler.stack
  1149. else:
  1150. toplevel = True
  1151. if not toplevel:
  1152. return self
  1153. mapper = statement._propagate_attrs["plugin_subject"]
  1154. dml_strategy = statement._annotations.get("dml_strategy", "raw")
  1155. if dml_strategy == "bulk":
  1156. self._setup_for_bulk_insert(compiler)
  1157. elif dml_strategy == "orm":
  1158. self._setup_for_orm_insert(compiler, mapper)
  1159. return self
  1160. @classmethod
  1161. def _resolved_keys_as_col_keys(cls, mapper, resolved_value_dict):
  1162. return {
  1163. col.key if col is not None else k: v
  1164. for col, k, v in (
  1165. (mapper.c.get(k), k, v) for k, v in resolved_value_dict.items()
  1166. )
  1167. }
  1168. def _setup_for_orm_insert(self, compiler, mapper):
  1169. statement = orm_level_statement = cast(dml.Insert, self.statement)
  1170. statement = self._setup_orm_returning(
  1171. compiler,
  1172. orm_level_statement,
  1173. statement,
  1174. dml_mapper=mapper,
  1175. use_supplemental_cols=False,
  1176. )
  1177. self.statement = statement
  1178. def _setup_for_bulk_insert(self, compiler):
  1179. """establish an INSERT statement within the context of
  1180. bulk insert.
  1181. This method will be within the "conn.execute()" call that is invoked
  1182. by persistence._emit_insert_statement().
  1183. """
  1184. statement = orm_level_statement = cast(dml.Insert, self.statement)
  1185. an = statement._annotations
  1186. emit_insert_table, emit_insert_mapper = (
  1187. an["_emit_insert_table"],
  1188. an["_emit_insert_mapper"],
  1189. )
  1190. statement = statement._clone()
  1191. statement.table = emit_insert_table
  1192. if self._dict_parameters:
  1193. self._dict_parameters = {
  1194. col: val
  1195. for col, val in self._dict_parameters.items()
  1196. if col.table is emit_insert_table
  1197. }
  1198. statement = self._setup_orm_returning(
  1199. compiler,
  1200. orm_level_statement,
  1201. statement,
  1202. dml_mapper=emit_insert_mapper,
  1203. use_supplemental_cols=True,
  1204. )
  1205. if (
  1206. self.from_statement_ctx is not None
  1207. and self.from_statement_ctx.compile_options._is_star
  1208. ):
  1209. raise sa_exc.CompileError(
  1210. "Can't use RETURNING * with bulk ORM INSERT. "
  1211. "Please use a different INSERT form, such as INSERT..VALUES "
  1212. "or INSERT with a Core Connection"
  1213. )
  1214. self.statement = statement
  1215. @CompileState.plugin_for("orm", "update")
  1216. class BulkORMUpdate(BulkUDCompileState, UpdateDMLState):
  1217. @classmethod
  1218. def create_for_statement(cls, statement, compiler, **kw):
  1219. self = cls.__new__(cls)
  1220. dml_strategy = statement._annotations.get(
  1221. "dml_strategy", "unspecified"
  1222. )
  1223. toplevel = not compiler.stack
  1224. if toplevel and dml_strategy == "bulk":
  1225. self._setup_for_bulk_update(statement, compiler)
  1226. elif (
  1227. dml_strategy == "core_only"
  1228. or dml_strategy == "unspecified"
  1229. and "parententity" not in statement.table._annotations
  1230. ):
  1231. UpdateDMLState.__init__(self, statement, compiler, **kw)
  1232. elif not toplevel or dml_strategy in ("orm", "unspecified"):
  1233. self._setup_for_orm_update(statement, compiler)
  1234. return self
  1235. def _setup_for_orm_update(self, statement, compiler, **kw):
  1236. orm_level_statement = statement
  1237. toplevel = not compiler.stack
  1238. ext_info = statement.table._annotations["parententity"]
  1239. self.mapper = mapper = ext_info.mapper
  1240. self._resolved_values = self._get_resolved_values(mapper, statement)
  1241. self._init_global_attributes(
  1242. statement,
  1243. compiler,
  1244. toplevel=toplevel,
  1245. process_criteria_for_toplevel=toplevel,
  1246. )
  1247. if statement._values:
  1248. self._resolved_values = dict(self._resolved_values)
  1249. new_stmt = statement._clone()
  1250. # note if the statement has _multi_values, these
  1251. # are passed through to the new statement, which will then raise
  1252. # InvalidRequestError because UPDATE doesn't support multi_values
  1253. # right now.
  1254. if statement._ordered_values:
  1255. new_stmt._ordered_values = self._resolved_values
  1256. elif statement._values:
  1257. new_stmt._values = self._resolved_values
  1258. new_crit = self._adjust_for_extra_criteria(
  1259. self.global_attributes, mapper
  1260. )
  1261. if new_crit:
  1262. new_stmt = new_stmt.where(*new_crit)
  1263. # if we are against a lambda statement we might not be the
  1264. # topmost object that received per-execute annotations
  1265. # do this first as we need to determine if there is
  1266. # UPDATE..FROM
  1267. UpdateDMLState.__init__(self, new_stmt, compiler, **kw)
  1268. use_supplemental_cols = False
  1269. if not toplevel:
  1270. synchronize_session = None
  1271. else:
  1272. synchronize_session = compiler._annotations.get(
  1273. "synchronize_session", None
  1274. )
  1275. can_use_returning = compiler._annotations.get(
  1276. "can_use_returning", None
  1277. )
  1278. if can_use_returning is not False:
  1279. # even though pre_exec has determined basic
  1280. # can_use_returning for the dialect, if we are to use
  1281. # RETURNING we need to run can_use_returning() at this level
  1282. # unconditionally because is_delete_using was not known
  1283. # at the pre_exec level
  1284. can_use_returning = (
  1285. synchronize_session == "fetch"
  1286. and self.can_use_returning(
  1287. compiler.dialect, mapper, is_multitable=self.is_multitable
  1288. )
  1289. )
  1290. if synchronize_session == "fetch" and can_use_returning:
  1291. use_supplemental_cols = True
  1292. # NOTE: we might want to RETURNING the actual columns to be
  1293. # synchronized also. however this is complicated and difficult
  1294. # to align against the behavior of "evaluate". Additionally,
  1295. # in a large number (if not the majority) of cases, we have the
  1296. # "evaluate" answer, usually a fixed value, in memory already and
  1297. # there's no need to re-fetch the same value
  1298. # over and over again. so perhaps if it could be RETURNING just
  1299. # the elements that were based on a SQL expression and not
  1300. # a constant. For now it doesn't quite seem worth it
  1301. new_stmt = new_stmt.return_defaults(*new_stmt.table.primary_key)
  1302. if toplevel:
  1303. new_stmt = self._setup_orm_returning(
  1304. compiler,
  1305. orm_level_statement,
  1306. new_stmt,
  1307. dml_mapper=mapper,
  1308. use_supplemental_cols=use_supplemental_cols,
  1309. )
  1310. self.statement = new_stmt
  1311. def _setup_for_bulk_update(self, statement, compiler, **kw):
  1312. """establish an UPDATE statement within the context of
  1313. bulk insert.
  1314. This method will be within the "conn.execute()" call that is invoked
  1315. by persistence._emit_update_statement().
  1316. """
  1317. statement = cast(dml.Update, statement)
  1318. an = statement._annotations
  1319. emit_update_table, _ = (
  1320. an["_emit_update_table"],
  1321. an["_emit_update_mapper"],
  1322. )
  1323. statement = statement._clone()
  1324. statement.table = emit_update_table
  1325. UpdateDMLState.__init__(self, statement, compiler, **kw)
  1326. if self._ordered_values:
  1327. raise sa_exc.InvalidRequestError(
  1328. "bulk ORM UPDATE does not support ordered_values() for "
  1329. "custom UPDATE statements with bulk parameter sets. Use a "
  1330. "non-bulk UPDATE statement or use values()."
  1331. )
  1332. if self._dict_parameters:
  1333. self._dict_parameters = {
  1334. col: val
  1335. for col, val in self._dict_parameters.items()
  1336. if col.table is emit_update_table
  1337. }
  1338. self.statement = statement
  1339. @classmethod
  1340. def orm_execute_statement(
  1341. cls,
  1342. session: Session,
  1343. statement: dml.Update,
  1344. params: _CoreAnyExecuteParams,
  1345. execution_options: OrmExecuteOptionsParameter,
  1346. bind_arguments: _BindArguments,
  1347. conn: Connection,
  1348. ) -> _result.Result:
  1349. update_options = execution_options.get(
  1350. "_sa_orm_update_options", cls.default_update_options
  1351. )
  1352. if update_options._dml_strategy not in (
  1353. "orm",
  1354. "auto",
  1355. "bulk",
  1356. "core_only",
  1357. ):
  1358. raise sa_exc.ArgumentError(
  1359. "Valid strategies for ORM UPDATE strategy "
  1360. "are 'orm', 'auto', 'bulk', 'core_only'"
  1361. )
  1362. result: _result.Result[Any]
  1363. if update_options._dml_strategy == "bulk":
  1364. enable_check_rowcount = not statement._where_criteria
  1365. assert update_options._synchronize_session != "fetch"
  1366. if (
  1367. statement._where_criteria
  1368. and update_options._synchronize_session == "evaluate"
  1369. ):
  1370. raise sa_exc.InvalidRequestError(
  1371. "bulk synchronize of persistent objects not supported "
  1372. "when using bulk update with additional WHERE "
  1373. "criteria right now. add synchronize_session=None "
  1374. "execution option to bypass synchronize of persistent "
  1375. "objects."
  1376. )
  1377. mapper = update_options._subject_mapper
  1378. assert mapper is not None
  1379. assert session._transaction is not None
  1380. result = _bulk_update(
  1381. mapper,
  1382. cast(
  1383. "Iterable[Dict[str, Any]]",
  1384. [params] if isinstance(params, dict) else params,
  1385. ),
  1386. session._transaction,
  1387. isstates=False,
  1388. update_changed_only=False,
  1389. use_orm_update_stmt=statement,
  1390. enable_check_rowcount=enable_check_rowcount,
  1391. )
  1392. return cls.orm_setup_cursor_result(
  1393. session,
  1394. statement,
  1395. params,
  1396. execution_options,
  1397. bind_arguments,
  1398. result,
  1399. )
  1400. else:
  1401. return super().orm_execute_statement(
  1402. session,
  1403. statement,
  1404. params,
  1405. execution_options,
  1406. bind_arguments,
  1407. conn,
  1408. )
  1409. @classmethod
  1410. def can_use_returning(
  1411. cls,
  1412. dialect: Dialect,
  1413. mapper: Mapper[Any],
  1414. *,
  1415. is_multitable: bool = False,
  1416. is_update_from: bool = False,
  1417. is_delete_using: bool = False,
  1418. is_executemany: bool = False,
  1419. ) -> bool:
  1420. # normal answer for "should we use RETURNING" at all.
  1421. normal_answer = (
  1422. dialect.update_returning and mapper.local_table.implicit_returning
  1423. )
  1424. if not normal_answer:
  1425. return False
  1426. if is_executemany:
  1427. return dialect.update_executemany_returning
  1428. # these workarounds are currently hypothetical for UPDATE,
  1429. # unlike DELETE where they impact MariaDB
  1430. if is_update_from:
  1431. return dialect.update_returning_multifrom
  1432. elif is_multitable and not dialect.update_returning_multifrom:
  1433. raise sa_exc.CompileError(
  1434. f'Dialect "{dialect.name}" does not support RETURNING '
  1435. "with UPDATE..FROM; for synchronize_session='fetch', "
  1436. "please add the additional execution option "
  1437. "'is_update_from=True' to the statement to indicate that "
  1438. "a separate SELECT should be used for this backend."
  1439. )
  1440. return True
  1441. @classmethod
  1442. def _do_post_synchronize_bulk_evaluate(
  1443. cls, session, params, result, update_options
  1444. ):
  1445. if not params:
  1446. return
  1447. mapper = update_options._subject_mapper
  1448. pk_keys = [prop.key for prop in mapper._identity_key_props]
  1449. identity_map = session.identity_map
  1450. for param in params:
  1451. identity_key = mapper.identity_key_from_primary_key(
  1452. (param[key] for key in pk_keys),
  1453. update_options._identity_token,
  1454. )
  1455. state = identity_map.fast_get_state(identity_key)
  1456. if not state:
  1457. continue
  1458. evaluated_keys = set(param).difference(pk_keys)
  1459. dict_ = state.dict
  1460. # only evaluate unmodified attributes
  1461. to_evaluate = state.unmodified.intersection(evaluated_keys)
  1462. for key in to_evaluate:
  1463. if key in dict_:
  1464. dict_[key] = param[key]
  1465. state.manager.dispatch.refresh(state, None, to_evaluate)
  1466. state._commit(dict_, list(to_evaluate))
  1467. # attributes that were formerly modified instead get expired.
  1468. # this only gets hit if the session had pending changes
  1469. # and autoflush were set to False.
  1470. to_expire = evaluated_keys.intersection(dict_).difference(
  1471. to_evaluate
  1472. )
  1473. if to_expire:
  1474. state._expire_attributes(dict_, to_expire)
  1475. @classmethod
  1476. def _do_post_synchronize_evaluate(
  1477. cls, session, statement, result, update_options
  1478. ):
  1479. matched_objects = cls._get_matched_objects_on_criteria(
  1480. update_options,
  1481. session.identity_map.all_states(),
  1482. )
  1483. cls._apply_update_set_values_to_objects(
  1484. session,
  1485. update_options,
  1486. statement,
  1487. [(obj, state, dict_) for obj, state, dict_, _ in matched_objects],
  1488. )
  1489. @classmethod
  1490. def _do_post_synchronize_fetch(
  1491. cls, session, statement, result, update_options
  1492. ):
  1493. target_mapper = update_options._subject_mapper
  1494. returned_defaults_rows = result.returned_defaults_rows
  1495. if returned_defaults_rows:
  1496. pk_rows = cls._interpret_returning_rows(
  1497. target_mapper, returned_defaults_rows
  1498. )
  1499. matched_rows = [
  1500. tuple(row) + (update_options._identity_token,)
  1501. for row in pk_rows
  1502. ]
  1503. else:
  1504. matched_rows = update_options._matched_rows
  1505. objs = [
  1506. session.identity_map[identity_key]
  1507. for identity_key in [
  1508. target_mapper.identity_key_from_primary_key(
  1509. list(primary_key),
  1510. identity_token=identity_token,
  1511. )
  1512. for primary_key, identity_token in [
  1513. (row[0:-1], row[-1]) for row in matched_rows
  1514. ]
  1515. if update_options._identity_token is None
  1516. or identity_token == update_options._identity_token
  1517. ]
  1518. if identity_key in session.identity_map
  1519. ]
  1520. if not objs:
  1521. return
  1522. cls._apply_update_set_values_to_objects(
  1523. session,
  1524. update_options,
  1525. statement,
  1526. [
  1527. (
  1528. obj,
  1529. attributes.instance_state(obj),
  1530. attributes.instance_dict(obj),
  1531. )
  1532. for obj in objs
  1533. ],
  1534. )
  1535. @classmethod
  1536. def _apply_update_set_values_to_objects(
  1537. cls, session, update_options, statement, matched_objects
  1538. ):
  1539. """apply values to objects derived from an update statement, e.g.
  1540. UPDATE..SET <values>
  1541. """
  1542. mapper = update_options._subject_mapper
  1543. target_cls = mapper.class_
  1544. evaluator_compiler = evaluator._EvaluatorCompiler(target_cls)
  1545. resolved_values = cls._get_resolved_values(mapper, statement)
  1546. resolved_keys_as_propnames = cls._resolved_keys_as_propnames(
  1547. mapper, resolved_values
  1548. )
  1549. value_evaluators = {}
  1550. for key, value in resolved_keys_as_propnames:
  1551. try:
  1552. _evaluator = evaluator_compiler.process(
  1553. coercions.expect(roles.ExpressionElementRole, value)
  1554. )
  1555. except evaluator.UnevaluatableError:
  1556. pass
  1557. else:
  1558. value_evaluators[key] = _evaluator
  1559. evaluated_keys = list(value_evaluators.keys())
  1560. attrib = {k for k, v in resolved_keys_as_propnames}
  1561. states = set()
  1562. for obj, state, dict_ in matched_objects:
  1563. to_evaluate = state.unmodified.intersection(evaluated_keys)
  1564. for key in to_evaluate:
  1565. if key in dict_:
  1566. # only run eval for attributes that are present.
  1567. dict_[key] = value_evaluators[key](obj)
  1568. state.manager.dispatch.refresh(state, None, to_evaluate)
  1569. state._commit(dict_, list(to_evaluate))
  1570. # attributes that were formerly modified instead get expired.
  1571. # this only gets hit if the session had pending changes
  1572. # and autoflush were set to False.
  1573. to_expire = attrib.intersection(dict_).difference(to_evaluate)
  1574. if to_expire:
  1575. state._expire_attributes(dict_, to_expire)
  1576. states.add(state)
  1577. session._register_altered(states)
  1578. @CompileState.plugin_for("orm", "delete")
  1579. class BulkORMDelete(BulkUDCompileState, DeleteDMLState):
  1580. @classmethod
  1581. def create_for_statement(cls, statement, compiler, **kw):
  1582. self = cls.__new__(cls)
  1583. dml_strategy = statement._annotations.get(
  1584. "dml_strategy", "unspecified"
  1585. )
  1586. if (
  1587. dml_strategy == "core_only"
  1588. or dml_strategy == "unspecified"
  1589. and "parententity" not in statement.table._annotations
  1590. ):
  1591. DeleteDMLState.__init__(self, statement, compiler, **kw)
  1592. return self
  1593. toplevel = not compiler.stack
  1594. orm_level_statement = statement
  1595. ext_info = statement.table._annotations["parententity"]
  1596. self.mapper = mapper = ext_info.mapper
  1597. self._init_global_attributes(
  1598. statement,
  1599. compiler,
  1600. toplevel=toplevel,
  1601. process_criteria_for_toplevel=toplevel,
  1602. )
  1603. new_stmt = statement._clone()
  1604. new_crit = cls._adjust_for_extra_criteria(
  1605. self.global_attributes, mapper
  1606. )
  1607. if new_crit:
  1608. new_stmt = new_stmt.where(*new_crit)
  1609. # do this first as we need to determine if there is
  1610. # DELETE..FROM
  1611. DeleteDMLState.__init__(self, new_stmt, compiler, **kw)
  1612. use_supplemental_cols = False
  1613. if not toplevel:
  1614. synchronize_session = None
  1615. else:
  1616. synchronize_session = compiler._annotations.get(
  1617. "synchronize_session", None
  1618. )
  1619. can_use_returning = compiler._annotations.get(
  1620. "can_use_returning", None
  1621. )
  1622. if can_use_returning is not False:
  1623. # even though pre_exec has determined basic
  1624. # can_use_returning for the dialect, if we are to use
  1625. # RETURNING we need to run can_use_returning() at this level
  1626. # unconditionally because is_delete_using was not known
  1627. # at the pre_exec level
  1628. can_use_returning = (
  1629. synchronize_session == "fetch"
  1630. and self.can_use_returning(
  1631. compiler.dialect,
  1632. mapper,
  1633. is_multitable=self.is_multitable,
  1634. is_delete_using=compiler._annotations.get(
  1635. "is_delete_using", False
  1636. ),
  1637. )
  1638. )
  1639. if can_use_returning:
  1640. use_supplemental_cols = True
  1641. new_stmt = new_stmt.return_defaults(*new_stmt.table.primary_key)
  1642. if toplevel:
  1643. new_stmt = self._setup_orm_returning(
  1644. compiler,
  1645. orm_level_statement,
  1646. new_stmt,
  1647. dml_mapper=mapper,
  1648. use_supplemental_cols=use_supplemental_cols,
  1649. )
  1650. self.statement = new_stmt
  1651. return self
  1652. @classmethod
  1653. def orm_execute_statement(
  1654. cls,
  1655. session: Session,
  1656. statement: dml.Delete,
  1657. params: _CoreAnyExecuteParams,
  1658. execution_options: OrmExecuteOptionsParameter,
  1659. bind_arguments: _BindArguments,
  1660. conn: Connection,
  1661. ) -> _result.Result:
  1662. update_options = execution_options.get(
  1663. "_sa_orm_update_options", cls.default_update_options
  1664. )
  1665. if update_options._dml_strategy == "bulk":
  1666. raise sa_exc.InvalidRequestError(
  1667. "Bulk ORM DELETE not supported right now. "
  1668. "Statement may be invoked at the "
  1669. "Core level using "
  1670. "session.connection().execute(stmt, parameters)"
  1671. )
  1672. if update_options._dml_strategy not in ("orm", "auto", "core_only"):
  1673. raise sa_exc.ArgumentError(
  1674. "Valid strategies for ORM DELETE strategy are 'orm', 'auto', "
  1675. "'core_only'"
  1676. )
  1677. return super().orm_execute_statement(
  1678. session, statement, params, execution_options, bind_arguments, conn
  1679. )
  1680. @classmethod
  1681. def can_use_returning(
  1682. cls,
  1683. dialect: Dialect,
  1684. mapper: Mapper[Any],
  1685. *,
  1686. is_multitable: bool = False,
  1687. is_update_from: bool = False,
  1688. is_delete_using: bool = False,
  1689. is_executemany: bool = False,
  1690. ) -> bool:
  1691. # normal answer for "should we use RETURNING" at all.
  1692. normal_answer = (
  1693. dialect.delete_returning and mapper.local_table.implicit_returning
  1694. )
  1695. if not normal_answer:
  1696. return False
  1697. # now get into special workarounds because MariaDB supports
  1698. # DELETE...RETURNING but not DELETE...USING...RETURNING.
  1699. if is_delete_using:
  1700. # is_delete_using hint was passed. use
  1701. # additional dialect feature (True for PG, False for MariaDB)
  1702. return dialect.delete_returning_multifrom
  1703. elif is_multitable and not dialect.delete_returning_multifrom:
  1704. # is_delete_using hint was not passed, but we determined
  1705. # at compile time that this is in fact a DELETE..USING.
  1706. # it's too late to continue since we did not pre-SELECT.
  1707. # raise that we need that hint up front.
  1708. raise sa_exc.CompileError(
  1709. f'Dialect "{dialect.name}" does not support RETURNING '
  1710. "with DELETE..USING; for synchronize_session='fetch', "
  1711. "please add the additional execution option "
  1712. "'is_delete_using=True' to the statement to indicate that "
  1713. "a separate SELECT should be used for this backend."
  1714. )
  1715. return True
  1716. @classmethod
  1717. def _do_post_synchronize_evaluate(
  1718. cls, session, statement, result, update_options
  1719. ):
  1720. matched_objects = cls._get_matched_objects_on_criteria(
  1721. update_options,
  1722. session.identity_map.all_states(),
  1723. )
  1724. to_delete = []
  1725. for _, state, dict_, is_partially_expired in matched_objects:
  1726. if is_partially_expired:
  1727. state._expire(dict_, session.identity_map._modified)
  1728. else:
  1729. to_delete.append(state)
  1730. if to_delete:
  1731. session._remove_newly_deleted(to_delete)
  1732. @classmethod
  1733. def _do_post_synchronize_fetch(
  1734. cls, session, statement, result, update_options
  1735. ):
  1736. target_mapper = update_options._subject_mapper
  1737. returned_defaults_rows = result.returned_defaults_rows
  1738. if returned_defaults_rows:
  1739. pk_rows = cls._interpret_returning_rows(
  1740. target_mapper, returned_defaults_rows
  1741. )
  1742. matched_rows = [
  1743. tuple(row) + (update_options._identity_token,)
  1744. for row in pk_rows
  1745. ]
  1746. else:
  1747. matched_rows = update_options._matched_rows
  1748. for row in matched_rows:
  1749. primary_key = row[0:-1]
  1750. identity_token = row[-1]
  1751. # TODO: inline this and call remove_newly_deleted
  1752. # once
  1753. identity_key = target_mapper.identity_key_from_primary_key(
  1754. list(primary_key),
  1755. identity_token=identity_token,
  1756. )
  1757. if identity_key in session.identity_map:
  1758. session._remove_newly_deleted(
  1759. [
  1760. attributes.instance_state(
  1761. session.identity_map[identity_key]
  1762. )
  1763. ]
  1764. )