Support Column objects in the SET clause for upsert

Established support for :class:`_schema.Column` objects as well as ORM
instrumented attributes as keys in the ``set_`` dictionary passed to the
:meth:`_postgresql.Insert.on_conflict_do_update` and
:meth:`_sqlite.Insert.on_conflict_do_update` methods, which match to the
:class:`_schema.Column` objects in the ``.c`` collection of the target
:class:`_schema.Table`. Previously,  only string column names were
expected; a column expression would be assumed to be an out-of-table
expression that would render fully along with a warning.

Fixes: #5722
Change-Id: Ice73b501d721c28d978a0277a83cedc6aff756a9

3 files changed, 174 insertions, 0 deletions

diff --git a/test/dialect/postgresql/test_compiler.py b/test/dialect/postgresql/test_compiler.py
index a031c3df9..9651f7bd9 100644
--- a/test/dialect/postgresql/test_compiler.py
+++ b/test/dialect/postgresql/test_compiler.py
@@ -1864,6 +1864,31 @@ class InsertOnConflictTest(fixtures.TestBase, AssertsCompiledSQL):
             },
         )
 
+    def test_do_update_set_clause_column_keys(self):
+        i = insert(self.table_with_metadata).values(myid=1, name="foo")
+        i = i.on_conflict_do_update(
+            index_elements=["myid"],
+            set_=OrderedDict(
+                [
+                    (self.table_with_metadata.c.name, "I'm a name"),
+                    (self.table_with_metadata.c.description, None),
+                ]
+            ),
+        )
+        self.assert_compile(
+            i,
+            "INSERT INTO mytable (myid, name) VALUES "
+            "(%(myid)s, %(name)s) ON CONFLICT (myid) "
+            "DO UPDATE SET name = %(param_1)s, "
+            "description = %(param_2)s",
+            {
+                "myid": 1,
+                "name": "foo",
+                "param_1": "I'm a name",
+                "param_2": None,
+            },
+        )
+
     def test_do_update_set_clause_literal(self):
         i = insert(self.table_with_metadata).values(myid=1, name="foo")
         i = i.on_conflict_do_update(
diff --git a/test/dialect/postgresql/test_on_conflict.py b/test/dialect/postgresql/test_on_conflict.py
index 7a9bfd75d..760487842 100644
--- a/test/dialect/postgresql/test_on_conflict.py
+++ b/test/dialect/postgresql/test_on_conflict.py
@@ -10,6 +10,7 @@ from sqlalchemy import Table
 from sqlalchemy import testing
 from sqlalchemy import types as sqltypes
 from sqlalchemy.dialects.postgresql import insert
+from sqlalchemy.testing import config
 from sqlalchemy.testing import fixtures
 from sqlalchemy.testing.assertions import assert_raises
 from sqlalchemy.testing.assertions import eq_
@@ -30,6 +31,14 @@ class OnConflictTest(fixtures.TablesTest):
             Column("name", String(50)),
         )
 
+        Table(
+            "users_schema",
+            metadata,
+            Column("id", Integer, primary_key=True),
+            Column("name", String(50)),
+            schema=config.test_schema,
+        )
+
         class SpecialType(sqltypes.TypeDecorator):
             impl = String
 
@@ -185,6 +194,99 @@ class OnConflictTest(fixtures.TablesTest):
                 [(1, "name1")],
             )
 
+    def test_on_conflict_do_update_schema(self):
+        users = self.tables.get("%s.users_schema" % config.test_schema)
+
+        with testing.db.connect() as conn:
+            conn.execute(users.insert(), dict(id=1, name="name1"))
+
+            i = insert(users)
+            i = i.on_conflict_do_update(
+                index_elements=[users.c.id], set_=dict(name=i.excluded.name)
+            )
+            result = conn.execute(i, dict(id=1, name="name1"))
+
+            eq_(result.inserted_primary_key, (1,))
+            eq_(result.returned_defaults, None)
+
+            eq_(
+                conn.execute(users.select().where(users.c.id == 1)).fetchall(),
+                [(1, "name1")],
+            )
+
+    def test_on_conflict_do_update_column_as_key_set(self):
+        users = self.tables.users
+
+        with testing.db.connect() as conn:
+            conn.execute(users.insert(), dict(id=1, name="name1"))
+
+            i = insert(users)
+            i = i.on_conflict_do_update(
+                index_elements=[users.c.id],
+                set_={users.c.name: i.excluded.name},
+            )
+            result = conn.execute(i, dict(id=1, name="name1"))
+
+            eq_(result.inserted_primary_key, (1,))
+            eq_(result.returned_defaults, None)
+
+            eq_(
+                conn.execute(users.select().where(users.c.id == 1)).fetchall(),
+                [(1, "name1")],
+            )
+
+    def test_on_conflict_do_update_clauseelem_as_key_set(self):
+        users = self.tables.users
+
+        class MyElem(object):
+            def __init__(self, expr):
+                self.expr = expr
+
+            def __clause_element__(self):
+                return self.expr
+
+        with testing.db.connect() as conn:
+            conn.execute(
+                users.insert(),
+                {"id": 1, "name": "name1"},
+            )
+
+            i = insert(users)
+            i = i.on_conflict_do_update(
+                index_elements=[users.c.id],
+                set_={MyElem(users.c.name): i.excluded.name},
+            ).values({MyElem(users.c.id): 1, MyElem(users.c.name): "name1"})
+            result = conn.execute(i)
+
+            eq_(result.inserted_primary_key, (1,))
+            eq_(result.returned_defaults, None)
+
+            eq_(
+                conn.execute(users.select().where(users.c.id == 1)).fetchall(),
+                [(1, "name1")],
+            )
+
+    def test_on_conflict_do_update_column_as_key_set_schema(self):
+        users = self.tables.get("%s.users_schema" % config.test_schema)
+
+        with testing.db.connect() as conn:
+            conn.execute(users.insert(), dict(id=1, name="name1"))
+
+            i = insert(users)
+            i = i.on_conflict_do_update(
+                index_elements=[users.c.id],
+                set_={users.c.name: i.excluded.name},
+            )
+            result = conn.execute(i, dict(id=1, name="name1"))
+
+            eq_(result.inserted_primary_key, (1,))
+            eq_(result.returned_defaults, None)
+
+            eq_(
+                conn.execute(users.select().where(users.c.id == 1)).fetchall(),
+                [(1, "name1")],
+            )
+
     def test_on_conflict_do_update_two(self):
         users = self.tables.users
 
diff --git a/test/dialect/test_sqlite.py b/test/dialect/test_sqlite.py
index 456bad7bd..f8b50f888 100644
--- a/test/dialect/test_sqlite.py
+++ b/test/dialect/test_sqlite.py
@@ -2921,6 +2921,53 @@ class OnConflictTest(fixtures.TablesTest):
             [(10, "I'm a name")],
         )
 
+    def test_on_conflict_do_update_column_keys(self, connection):
+        users = self.tables.users
+
+        conn = connection
+        conn.execute(users.insert(), dict(id=1, name="name1"))
+
+        i = insert(users)
+        i = i.on_conflict_do_update(
+            index_elements=users.primary_key.columns,
+            set_={users.c.id: 10, users.c.name: "I'm a name"},
+        ).values(id=1, name="name4")
+
+        result = conn.execute(i)
+        eq_(result.inserted_primary_key, (1,))
+
+        eq_(
+            conn.execute(users.select().where(users.c.id == 10)).fetchall(),
+            [(10, "I'm a name")],
+        )
+
+    def test_on_conflict_do_update_clauseelem_keys(self, connection):
+        users = self.tables.users
+
+        class MyElem(object):
+            def __init__(self, expr):
+                self.expr = expr
+
+            def __clause_element__(self):
+                return self.expr
+
+        conn = connection
+        conn.execute(users.insert(), dict(id=1, name="name1"))
+
+        i = insert(users)
+        i = i.on_conflict_do_update(
+            index_elements=users.primary_key.columns,
+            set_={MyElem(users.c.id): 10, MyElem(users.c.name): "I'm a name"},
+        ).values({MyElem(users.c.id): 1, MyElem(users.c.name): "name4"})
+
+        result = conn.execute(i)
+        eq_(result.inserted_primary_key, (1,))
+
+        eq_(
+            conn.execute(users.select().where(users.c.id == 10)).fetchall(),
+            [(10, "I'm a name")],
+        )
+
     def test_on_conflict_do_update_multivalues(self, connection):
         users = self.tables.users