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