ORM bulk insert via execute

* ORM Insert now includes "bulk" mode that will run
  essentially the same process as session.bulk_insert_mappings;
  interprets the given list of values as ORM attributes for
  key names
* ORM UPDATE has a similar feature, without RETURNING support,
  for session.bulk_update_mappings
* Added support for upserts to do RETURNING ORM objects as well
* ORM UPDATE/DELETE with list of parameters + WHERE criteria
  is a not implemented; use connection
* ORM UPDATE/DELETE defaults to "auto" synchronize_session;
  use fetch if RETURNING is present, evaluate if not, as
  "fetch" is much more efficient (no expired object SELECT problem)
  and less error prone if RETURNING is available
  UPDATE: howver this is inefficient!   please continue to
  use evaluate for simple cases, auto can move to fetch
  if criteria not evaluable
* "Evaluate" criteria will now not preemptively
  unexpire and SELECT attributes that were individually
  expired. Instead, if evaluation of the criteria indicates that
  the necessary attrs were expired, we expire the object
  completely (delete) or expire the SET attrs unconditionally
  (update). This keeps the object in the same unloaded state
  where it will refresh those attrs on the next pass, for
  this generally unusual case.  (originally #5664)
* Core change! update/delete rowcount comes from len(rows)
  if RETURNING was used.  SQLite at least otherwise did not
  support this.  adjusted test_rowcount accordingly
* ORM DELETE with a list of parameters at all is also a not
  implemented as this would imply "bulk", and there is no
  bulk_delete_mappings (could be, but we dont have that)
* ORM insert().values() with single or multi-values translates
  key names based on ORM attribute names
* ORM returning() implemented for insert, update, delete;
  explcit returning clauses now interpret rows in an ORM
  context, with support for qualifying loader options as well
* session.bulk_insert_mappings() assigns polymorphic identity
  if not set.
* explicit RETURNING + synchronize_session='fetch' is now
  supported with UPDATE and DELETE.
* expanded return_defaults() to work with DELETE also.
* added support for composite attributes to be present
  in the dictionaries used by bulk_insert_mappings and
  bulk_update_mappings, which is also the new ORM bulk
  insert/update feature, that will expand the composite
  values into their individual mapped attributes the way they'd
  be on a mapped instance.
* bulk UPDATE supports "synchronize_session=evaluate", is the
  default.  this does not apply to session.bulk_update_mappings,
  just the new version
* both bulk UPDATE and bulk INSERT, the latter with or without
  RETURNING, support *heterogenous* parameter sets.
  session.bulk_insert/update_mappings did this, so this feature
  is maintained.  now cursor result can be both horizontally
  and vertically spliced :)

This is now a long story with a lot of options, which in
itself is a problem to be able to document all of this
in some way that makes sense.  raising exceptions for
use cases we haven't supported is pretty important here
too, the tradition of letting unsupported things just not work
is likely not a good idea at this point, though there
are still many cases that aren't easily avoidable

Fixes: #8360
Fixes: #7864
Fixes: #7865
Change-Id: Idf28379f8705e403a3c6a937f6a798a042ef2540

1 files changed, 141 insertions, 32 deletions

diff --git a/lib/sqlalchemy/orm/context.py b/lib/sqlalchemy/orm/context.py
index dc96f8c3c..f8c7ba714 100644
--- a/lib/sqlalchemy/orm/context.py
+++ b/lib/sqlalchemy/orm/context.py
@@ -73,6 +73,7 @@ if TYPE_CHECKING:
     from .query import Query
     from .session import _BindArguments
     from .session import Session
+    from ..engine import Result
     from ..engine.interfaces import _CoreSingleExecuteParams
     from ..engine.interfaces import _ExecuteOptionsParameter
     from ..sql._typing import _ColumnsClauseArgument
@@ -203,15 +204,19 @@ _orm_load_exec_options = util.immutabledict(
 
 
 class AbstractORMCompileState(CompileState):
+    is_dml_returning = False
+
     @classmethod
     def create_for_statement(
         cls,
         statement: Union[Select, FromStatement],
         compiler: Optional[SQLCompiler],
         **kw: Any,
-    ) -> ORMCompileState:
+    ) -> AbstractORMCompileState:
         """Create a context for a statement given a :class:`.Compiler`.
+
         This method is always invoked in the context of SQLCompiler.process().
+
         For a Select object, this would be invoked from
         SQLCompiler.visit_select(). For the special FromStatement object used
         by Query to indicate "Query.from_statement()", this is called by
@@ -233,6 +238,28 @@ class AbstractORMCompileState(CompileState):
         raise NotImplementedError()
 
     @classmethod
+    def orm_execute_statement(
+        cls,
+        session,
+        statement,
+        params,
+        execution_options,
+        bind_arguments,
+        conn,
+    ) -> Result:
+        result = conn.execute(
+            statement, params or {}, execution_options=execution_options
+        )
+        return cls.orm_setup_cursor_result(
+            session,
+            statement,
+            params,
+            execution_options,
+            bind_arguments,
+            result,
+        )
+
+    @classmethod
     def orm_setup_cursor_result(
         cls,
         session,
@@ -309,6 +336,17 @@ class ORMCompileState(AbstractORMCompileState):
     def __init__(self, *arg, **kw):
         raise NotImplementedError()
 
+    if TYPE_CHECKING:
+
+        @classmethod
+        def create_for_statement(
+            cls,
+            statement: Union[Select, FromStatement],
+            compiler: Optional[SQLCompiler],
+            **kw: Any,
+        ) -> ORMCompileState:
+            ...
+
     def _append_dedupe_col_collection(self, obj, col_collection):
         dedupe = self.dedupe_columns
         if obj not in dedupe:
@@ -333,26 +371,6 @@ class ORMCompileState(AbstractORMCompileState):
             return SelectState._column_naming_convention(label_style)
 
     @classmethod
-    def create_for_statement(
-        cls,
-        statement: Union[Select, FromStatement],
-        compiler: Optional[SQLCompiler],
-        **kw: Any,
-    ) -> ORMCompileState:
-        """Create a context for a statement given a :class:`.Compiler`.
-
-        This method is always invoked in the context of SQLCompiler.process().
-
-        For a Select object, this would be invoked from
-        SQLCompiler.visit_select(). For the special FromStatement object used
-        by Query to indicate "Query.from_statement()", this is called by
-        FromStatement._compiler_dispatch() that would be called by
-        SQLCompiler.process().
-
-        """
-        raise NotImplementedError()
-
-    @classmethod
     def get_column_descriptions(cls, statement):
         return _column_descriptions(statement)
 
@@ -518,6 +536,49 @@ class ORMCompileState(AbstractORMCompileState):
         )
 
 
+class DMLReturningColFilter:
+    """an adapter used for the DML RETURNING case.
+
+    Has a subset of the interface used by
+    :class:`.ORMAdapter` and is used for :class:`._QueryEntity`
+    instances to set up their columns as used in RETURNING for a
+    DML statement.
+
+    """
+
+    __slots__ = ("mapper", "columns", "__weakref__")
+
+    def __init__(self, target_mapper, immediate_dml_mapper):
+        if (
+            immediate_dml_mapper is not None
+            and target_mapper.local_table
+            is not immediate_dml_mapper.local_table
+        ):
+            # joined inh, or in theory other kinds of multi-table mappings
+            self.mapper = immediate_dml_mapper
+        else:
+            # single inh, normal mappings, etc.
+            self.mapper = target_mapper
+        self.columns = self.columns = util.WeakPopulateDict(
+            self.adapt_check_present  # type: ignore
+        )
+
+    def __call__(self, col, as_filter):
+        for cc in sql_util._find_columns(col):
+            c2 = self.adapt_check_present(cc)
+            if c2 is not None:
+                return col
+        else:
+            return None
+
+    def adapt_check_present(self, col):
+        mapper = self.mapper
+        prop = mapper._columntoproperty.get(col, None)
+        if prop is None:
+            return None
+        return mapper.local_table.c.corresponding_column(col)
+
+
 @sql.base.CompileState.plugin_for("orm", "orm_from_statement")
 class ORMFromStatementCompileState(ORMCompileState):
     _from_obj_alias = None
@@ -525,7 +586,7 @@ class ORMFromStatementCompileState(ORMCompileState):
 
     statement_container: FromStatement
     requested_statement: Union[SelectBase, TextClause, UpdateBase]
-    dml_table: _DMLTableElement
+    dml_table: Optional[_DMLTableElement] = None
 
     _has_orm_entities = False
     multi_row_eager_loaders = False
@@ -541,7 +602,7 @@ class ORMFromStatementCompileState(ORMCompileState):
         statement_container: Union[Select, FromStatement],
         compiler: Optional[SQLCompiler],
         **kw: Any,
-    ) -> ORMCompileState:
+    ) -> ORMFromStatementCompileState:
 
         if compiler is not None:
             toplevel = not compiler.stack
@@ -565,6 +626,7 @@ class ORMFromStatementCompileState(ORMCompileState):
 
         if statement.is_dml:
             self.dml_table = statement.table
+            self.is_dml_returning = True
 
         self._entities = []
         self._polymorphic_adapters = {}
@@ -674,6 +736,18 @@ class ORMFromStatementCompileState(ORMCompileState):
     def _get_current_adapter(self):
         return None
 
+    def setup_dml_returning_compile_state(self, dml_mapper):
+        """used by BulkORMInsert (and Update / Delete?) to set up a handler
+        for RETURNING to return ORM objects and expressions
+
+        """
+        target_mapper = self.statement._propagate_attrs.get(
+            "plugin_subject", None
+        )
+        adapter = DMLReturningColFilter(target_mapper, dml_mapper)
+        for entity in self._entities:
+            entity.setup_dml_returning_compile_state(self, adapter)
+
 
 class FromStatement(GroupedElement, Generative, TypedReturnsRows[_TP]):
     """Core construct that represents a load of ORM objects from various
@@ -813,7 +887,7 @@ class ORMSelectCompileState(ORMCompileState, SelectState):
         statement: Union[Select, FromStatement],
         compiler: Optional[SQLCompiler],
         **kw: Any,
-    ) -> ORMCompileState:
+    ) -> ORMSelectCompileState:
         """compiler hook, we arrive here from compiler.visit_select() only."""
 
         self = cls.__new__(cls)
@@ -2312,6 +2386,13 @@ class _QueryEntity:
     def setup_compile_state(self, compile_state: ORMCompileState) -> None:
         raise NotImplementedError()
 
+    def setup_dml_returning_compile_state(
+        self,
+        compile_state: ORMCompileState,
+        adapter: DMLReturningColFilter,
+    ) -> None:
+        raise NotImplementedError()
+
     def row_processor(self, context, result):
         raise NotImplementedError()
 
@@ -2509,8 +2590,24 @@ class _MapperEntity(_QueryEntity):
 
         return _instance, self._label_name, self._extra_entities
 
-    def setup_compile_state(self, compile_state):
+    def setup_dml_returning_compile_state(
+        self,
+        compile_state: ORMCompileState,
+        adapter: DMLReturningColFilter,
+    ) -> None:
+        loading._setup_entity_query(
+            compile_state,
+            self.mapper,
+            self,
+            self.path,
+            adapter,
+            compile_state.primary_columns,
+            with_polymorphic=self._with_polymorphic_mappers,
+            only_load_props=compile_state.compile_options._only_load_props,
+            polymorphic_discriminator=self._polymorphic_discriminator,
+        )
 
+    def setup_compile_state(self, compile_state):
         adapter = self._get_entity_clauses(compile_state)
 
         single_table_crit = self.mapper._single_table_criterion
@@ -2536,7 +2633,6 @@ class _MapperEntity(_QueryEntity):
             only_load_props=compile_state.compile_options._only_load_props,
             polymorphic_discriminator=self._polymorphic_discriminator,
         )
-
         compile_state._fallback_from_clauses.append(self.selectable)
 
 
@@ -2743,9 +2839,7 @@ class _ColumnEntity(_QueryEntity):
             getter, label_name, extra_entities = self._row_processor
             if self.translate_raw_column:
                 extra_entities += (
-                    result.context.invoked_statement._raw_columns[
-                        self.raw_column_index
-                    ],
+                    context.query._raw_columns[self.raw_column_index],
                 )
 
             return getter, label_name, extra_entities
@@ -2781,9 +2875,7 @@ class _ColumnEntity(_QueryEntity):
 
         if self.translate_raw_column:
             extra_entities = self._extra_entities + (
-                result.context.invoked_statement._raw_columns[
-                    self.raw_column_index
-                ],
+                context.query._raw_columns[self.raw_column_index],
             )
             return getter, self._label_name, extra_entities
         else:
@@ -2843,6 +2935,8 @@ class _RawColumnEntity(_ColumnEntity):
         current_adapter = compile_state._get_current_adapter()
         if current_adapter:
             column = current_adapter(self.column, False)
+            if column is None:
+                return
         else:
             column = self.column
 
@@ -2944,10 +3038,25 @@ class _ORMColumnEntity(_ColumnEntity):
                 self.entity_zero
             ) and entity.common_parent(self.entity_zero)
 
+    def setup_dml_returning_compile_state(
+        self,
+        compile_state: ORMCompileState,
+        adapter: DMLReturningColFilter,
+    ) -> None:
+        self._fetch_column = self.column
+        column = adapter(self.column, False)
+        if column is not None:
+            compile_state.dedupe_columns.add(column)
+            compile_state.primary_columns.append(column)
+
     def setup_compile_state(self, compile_state):
         current_adapter = compile_state._get_current_adapter()
         if current_adapter:
             column = current_adapter(self.column, False)
+            if column is None:
+                assert compile_state.is_dml_returning
+                self._fetch_column = self.column
+                return
         else:
             column = self.column