generalize sql server check for id col to accommodate ORM cases

Fixed issues that prevented the new usage patterns for using DML with ORM
objects presented at :ref:`orm_dml_returning_objects` from working
correctly with the SQL Server pyodbc dialect.

Here we add a step to look in compile_state._dict_values more thoroughly
for the keys we need to determine "identity insert" or not, and also
add a new compiler variable dml_compile_state so that we can skip the
ORM's compile_state if present.

Fixes: #8210
Change-Id: Idbd76bb3eb075c647dc6c1cb78f7315c821e15f7

3 files changed, 43 insertions, 4 deletions

diff --git a/lib/sqlalchemy/sql/coercions.py b/lib/sqlalchemy/sql/coercions.py
index d56035db7..0b250a28e 100644
--- a/lib/sqlalchemy/sql/coercions.py
+++ b/lib/sqlalchemy/sql/coercions.py
@@ -164,6 +164,17 @@ def expect(
 
 @overload
 def expect(
+    role: Type[roles.DMLColumnRole],
+    element: Any,
+    *,
+    as_key: Literal[True] = ...,
+    **kw: Any,
+) -> str:
+    ...
+
+
+@overload
+def expect(
     role: Type[roles.LiteralValueRole],
     element: Any,
     **kw: Any,
@@ -420,9 +431,11 @@ def expect(
         )
 
 
-def expect_as_key(role, element, **kw):
-    kw["as_key"] = True
-    return expect(role, element, **kw)
+def expect_as_key(
+    role: Type[roles.DMLColumnRole], element: Any, **kw: Any
+) -> str:
+    kw.pop("as_key", None)
+    return expect(role, element, as_key=True, **kw)
 
 
 def expect_col_expression_collection(
diff --git a/lib/sqlalchemy/sql/compiler.py b/lib/sqlalchemy/sql/compiler.py
index 60ec09771..a11d83b11 100644
--- a/lib/sqlalchemy/sql/compiler.py
+++ b/lib/sqlalchemy/sql/compiler.py
@@ -529,6 +529,18 @@ class Compiled:
 
     """
 
+    dml_compile_state: Optional[CompileState] = None
+    """Optional :class:`.CompileState` assigned at the same point that
+    .isinsert, .isupdate, or .isdelete is assigned.
+
+    This will normally be the same object as .compile_state, with the
+    exception of cases like the :class:`.ORMFromStatementCompileState`
+    object.
+
+    .. versionadded:: 1.4.40
+
+    """
+
     cache_key: Optional[CacheKey] = None
     """The :class:`.CacheKey` that was generated ahead of creating this
     :class:`.Compiled` object.
@@ -4371,6 +4383,8 @@ class SQLCompiler(Compiled):
 
         if toplevel:
             self.isinsert = True
+            if not self.dml_compile_state:
+                self.dml_compile_state = compile_state
             if not self.compile_state:
                 self.compile_state = compile_state
 
@@ -4548,6 +4562,8 @@ class SQLCompiler(Compiled):
         toplevel = not self.stack
         if toplevel:
             self.isupdate = True
+            if not self.dml_compile_state:
+                self.dml_compile_state = compile_state
             if not self.compile_state:
                 self.compile_state = compile_state
 
@@ -4683,6 +4699,8 @@ class SQLCompiler(Compiled):
         toplevel = not self.stack
         if toplevel:
             self.isdelete = True
+            if not self.dml_compile_state:
+                self.dml_compile_state = compile_state
             if not self.compile_state:
                 self.compile_state = compile_state
 
diff --git a/lib/sqlalchemy/sql/dml.py b/lib/sqlalchemy/sql/dml.py
index 2ed3be9cb..e99f35418 100644
--- a/lib/sqlalchemy/sql/dml.py
+++ b/lib/sqlalchemy/sql/dml.py
@@ -263,7 +263,7 @@ class DMLState(CompileState):
 
     def _process_select_values(self, statement: ValuesBase) -> None:
         assert statement._select_names is not None
-        parameters = {
+        parameters: MutableMapping[_DMLColumnElement, Any] = {
             coercions.expect(roles.DMLColumnRole, name, as_key=True): Null()
             for name in statement._select_names
         }
@@ -312,6 +312,14 @@ class InsertDMLState(DMLState):
         if statement._multi_values:
             self._process_multi_values(statement)
 
+    @util.memoized_property
+    def _insert_col_keys(self) -> List[str]:
+        # this is also done in crud.py -> _key_getters_for_crud_column
+        return [
+            coercions.expect_as_key(roles.DMLColumnRole, col)
+            for col in self._dict_parameters or ()
+        ]
+
 
 @CompileState.plugin_for("default", "update")
 class UpdateDMLState(DMLState):