Repository: kayak/pypika
Pull request #374 (2020-02-11 08:22:51)

Update and Insert improving


Patch:
diff --git a/pypika/dialects.py b/pypika/dialects.py
index 2465c9d..512dca7 100644
--- a/pypika/dialects.py
+++ b/pypika/dialects.py
@@ -10,9 +10,11 @@ from pypika.queries import (
 from pypika.terms import (
     ArithmeticExpression,
     Field,
+    Term,
     Function,
     Star,
     ValueWrapper,
+    EmptyCriterion
 )
 from pypika.utils import (
     QueryException,
@@ -292,28 +294,46 @@ class PostgreQueryBuilder(QueryBuilder):
         super(PostgreQueryBuilder, self).__init__(dialect=Dialects.POSTGRESQL, **kwargs)
         self._returns = []
         self._return_star = False
-        self._on_conflict_field = None
+
+        self._on_conflict = False
+        self._on_conflict_fields = []
         self._on_conflict_do_nothing = False
-        self._on_conflict_updates = []
+        self._on_conflict_do_updates = []
+        self._on_conflict_wheres = None
+        self._on_conflict_do_update_wheres = None
+
+        self._distinct_on = []
 
     def __copy__(self):
         newone = super(PostgreQueryBuilder, self).__copy__()
         newone._returns = copy(self._returns)
-        newone._on_conflict_updates = copy(self._on_conflict_updates)
+        newone._on_conflict_do_updates = copy(self._on_conflict_do_updates)
         return newone
 
     @builder
-    def on_conflict(self, target_field):
+    def distinct_on(self, *fields):
+        for field in fields:
+            if isinstance(field, str):
+                self._distinct_on.append(Field(field))
+            elif isinstance(field, Term):
+                self._distinct_on.append(field)
+
+    @builder
+    def on_conflict(self, *target_fields):
         if not self._insert_table:
             raise QueryException("On conflict only applies to insert query")
-        if isinstance(target_field, str):
-            self._on_conflict_field = self._conflict_field_str(target_field)
-        elif isinstance(target_field, Field):
-            self._on_conflict_field = target_field
+
+        self._on_conflict = True
+
+        for target_field in target_fields:
+            if isinstance(target_field, str):
+                self._on_conflict_fields.append(self._conflict_field_str(target_field))
+            elif isinstance(target_field, Term):
+                self._on_conflict_fields.append(target_field)
 
     @builder
     def do_nothing(self):
-        if len(self._on_conflict_updates) > 0:
+        if len(self._on_conflict_do_updates) > 0:
             raise QueryException("Can not have two conflict handlers")
         self._on_conflict_do_nothing = True
 
@@ -326,43 +346,91 @@ class PostgreQueryBuilder(QueryBuilder):
             field = self._conflict_field_str(update_field)
         elif isinstance(update_field, Field):
             field = update_field
-        self._on_conflict_updates.append((field, ValueWrapper(update_value)))
+        else:
+            raise QueryException("Unsupported update_field")
+
+        self._on_conflict_do_updates.append((field, ValueWrapper(update_value)))
+
+    @builder
+    def where(self, criterion):
+        if not self._on_conflict:
+            return super().where(criterion)
+
+        if isinstance(criterion, EmptyCriterion):
+            return
+
+        if self._on_conflict_do_nothing:
+            raise QueryException('DO NOTHING doest not support WHERE')
+
+        if self._on_conflict_fields and self._on_conflict_do_updates:
+            if self._on_conflict_do_update_wheres:
+                self._on_conflict_do_update_wheres &= criterion
+            else:
+                self._on_conflict_do_update_wheres = criterion
+        elif self._on_conflict_fields:
+            if self._on_conflict_wheres:
+                self._on_conflict_wheres &= criterion
+            else:
+                self._on_conflict_wheres = criterion
+        else:
+            raise QueryException('Can not have fieldless ON CONFLICT WHERE')
+
+    def _distinct_sql(self, **kwargs):
+        if self._distinct_on:
+            return "DISTINCT ON({distinct_on}) ".format(
+                distinct_on=",".join(
+                    term.get_sql(with_alias=True, **kwargs) for term in self._distinct_on
+                )
+            )
+        return super()._distinct_sql(**kwargs)
 
     def _conflict_field_str(self, term):
         if self._insert_table:
             return Field(term, table=self._insert_table)
 
     def _on_conflict_sql(self, **kwargs):
-        if not self._on_conflict_do_nothing and len(self._on_conflict_updates) == 0:
-            if not self._on_conflict_field:
+        if not self._on_conflict_do_nothing and len(self._on_conflict_do_updates) == 0:
+            if not self._on_conflict_fields:
                 return ""
-            else:
-                raise QueryException("No handler defined for on conflict")
-        else:
-            conflict_query = " ON CONFLICT"
-            if self._on_conflict_field:
-                conflict_query += (
-                    " ("
-                    + self._on_conflict_field.get_sql(with_alias=True, **kwargs)
-                    + ")"
-                )
-            if self._on_conflict_do_nothing:
-                conflict_query += " DO NOTHING"
-            elif len(self._on_conflict_updates) > 0:
-                if self._on_conflict_field:
-                    conflict_query += " DO UPDATE SET {updates}".format(
-                        updates=",".join(
-                            "{field}={value}".format(
-                                field=field.get_sql(**kwargs),
-                                value=value.get_sql(**kwargs),
-                            )
-                            for field, value in self._on_conflict_updates
-                        )
+            raise QueryException("No handler defined for on conflict")
+
+        if self._on_conflict_do_updates and not self._on_conflict_fields:
+            raise QueryException("Can not have fieldless on conflict do update")
+
+        conflict_query = " ON CONFLICT"
+        if self._on_conflict_fields:
+            fields = [f.get_sql(with_alias=True, **kwargs)
+                      for f in self._on_conflict_fields]
+            conflict_query += " (" + ', '.join(fields) + ")"
+
+        if self._on_conflict_wheres:
+            conflict_query += " WHERE {where}".format(
+                where=self._on_conflict_wheres.get_sql(subquery=True, **kwargs)
+            )
+
+        return conflict_query
+
+    def _on_conflict_action_sql(self, **kwargs):
+        if self._on_conflict_do_nothing:
+            return " DO NOTHING"
+        elif len(self._on_conflict_do_updates) > 0:
+            action_sql = " DO UPDATE SET {updates}".format(
+                updates=",".join(
+                    "{field}={value}".format(
+                        field=field.get_sql(**kwargs),
+                        value=value.get_sql(with_namespace=True, **kwargs),
                     )
-                else:
-                    raise QueryException("Can not have fieldless on conflict do update")
+                    for field, value in self._on_conflict_do_updates
+                )
+            )
+
+            if self._on_conflict_do_update_wheres:
+                action_sql += " WHERE {where}".format(
+                    where=self._on_conflict_do_update_wheres.get_sql(subquery=True, with_namespace=True, **kwargs)
+                )
+            return action_sql
 
-            return conflict_query
+        return ''
 
     @builder
     def returning(self, *terms):
@@ -433,12 +501,20 @@ class PostgreQueryBuilder(QueryBuilder):
         )
 
     def get_sql(self, with_alias=False, subquery=False, **kwargs):
+        self._set_kwargs_defaults(kwargs)
+
         querystring = super(PostgreQueryBuilder, self).get_sql(
             with_alias, subquery, **kwargs
         )
-        querystring += self._on_conflict_sql()
+        with_namespace = False
+        if self._update_table and self.from_:
+            with_namespace = True
+
+        querystring += self._on_conflict_sql(**kwargs)
+        querystring += self._on_conflict_action_sql(**kwargs)
+
         if self._returns:
-            querystring += self._returning_sql()
+            querystring += self._returning_sql(with_namespace=with_namespace, **kwargs)
         return querystring
 
 
diff --git a/pypika/queries.py b/pypika/queries.py
index 8dbc872..4819e84 100644
--- a/pypika/queries.py
+++ b/pypika/queries.py
@@ -719,9 +719,6 @@ class QueryBuilder(Selectable, Term):
         if self._insert_table is None:
             raise AttributeError("'Query' object has no attribute '%s'" % "insert")
 
-        if isinstance(terms[0], (list, tuple)):
-            terms = terms[0]
-
         for term in terms:
             if isinstance(term, str):
                 term = Field(term, table=self._insert_table)
@@ -1071,6 +1068,7 @@ class QueryBuilder(Selectable, Term):
             self._from[0], QueryBuilder
         )
         has_reference_to_foreign_table = self._foreign_table
+        has_update_from = self._update_table and self._from
 
         kwargs["with_namespace"] = any(
             [
@@ -1078,11 +1076,17 @@ class QueryBuilder(Selectable, Term):
                 has_multiple_from_clauses,
                 has_subquery_from_clause,
                 has_reference_to_foreign_table,
+                has_update_from
             ]
         )
 
         if self._update_table:
-            querystring = self._update_sql(**kwargs)
+            if self._with:
+                querystring = self._with_sql(**kwargs)
+            else:
+                querystring = ""
+
+            querystring += self._update_sql(**kwargs)
 
             if self._joins:
                 querystring += " " + " ".join(
@@ -1091,6 +1095,9 @@ class QueryBuilder(Selectable, Term):
 
             querystring += self._set_sql(**kwargs)
 
+            if self._from:
+                querystring += self._from_sql(**kwargs)
+
             if self._wheres:
                 querystring += self._where_sql(**kwargs)
 
@@ -1103,10 +1110,15 @@ class QueryBuilder(Selectable, Term):
             querystring = self._delete_sql(**kwargs)
 
         elif not self._select_into and self._insert_table:
+            if self._with:
+                querystring = self._with_sql(**kwargs)
+            else:
+                querystring = ""
+
             if self._replace:
-                querystring = self._replace_sql(**kwargs)
+                querystring += self._replace_sql(**kwargs)
             else:
-                querystring = self._insert_sql(**kwargs)
+                querystring += self._insert_sql(**kwargs)
 
             if self._columns:
                 querystring += self._columns_sql(**kwargs)
@@ -1179,9 +1191,17 @@ class QueryBuilder(Selectable, Term):
             for clause in self._with
         )
 
+    def _distinct_sql(self, **kwargs):
+        if self._distinct:
+            distinct = 'DISTINCT '
+        else:
+            distinct = ''
+
+        return distinct
+
     def _select_sql(self, **kwargs):
         return "SELECT {distinct}{select}".format(
-            distinct="DISTINCT " if self._distinct else "",
+            distinct=self._distinct_sql(**kwargs),
             select=",".join(
                 term.get_sql(with_alias=True, subquery=True, **kwargs)
                 for term in self._selects
@@ -1344,7 +1364,7 @@ class QueryBuilder(Selectable, Term):
         return " SET {set}".format(
             set=",".join(
                 "{field}={value}".format(
-                    field=field.get_sql(**kwargs), value=value.get_sql(**kwargs)
+                    field=field.get_sql(**dict(kwargs, with_namespace=False)), value=value.get_sql(**kwargs)
                 )
                 for field, value in self._updates
             )


Test patch:
diff --git a/pypika/tests/dialects/test_postgresql.py b/pypika/tests/dialects/test_postgresql.py
index 75fb868..71884ef 100644
--- a/pypika/tests/dialects/test_postgresql.py
+++ b/pypika/tests/dialects/test_postgresql.py
@@ -207,3 +207,14 @@ class JSONBOperatorsTests(unittest.TestCase):
         self.assertEqual(
             "SELECT * FROM \"abc\" WHERE \"json\"?|ARRAY['dates','imported']", str(q)
         )
+
+
+class DistinctOnTests(unittest.TestCase):
+    table_abc = Table("abc")
+
+    def test_distinct_on(self):
+        q = PostgreSQLQuery.from_(self.table_abc).distinct_on('lname', self.table_abc.fname).select('lname', 'id')
+
+        self.assertEqual(
+            '''SELECT DISTINCT ON("lname","fname") "lname","id" FROM "abc"''', str(q)
+        )
diff --git a/pypika/tests/test_deletes.py b/pypika/tests/test_deletes.py
index 006be1a..626717d 100644
--- a/pypika/tests/test_deletes.py
+++ b/pypika/tests/test_deletes.py
@@ -46,7 +46,7 @@ class PostgresDeleteTests(unittest.TestCase):
             .returning(self.table_abc.id)
         )
 
-        self.assertEqual('DELETE FROM "abc" WHERE "foo"="bar" RETURNING id', str(q1))
+        self.assertEqual('DELETE FROM "abc" WHERE "foo"="bar" RETURNING "id"', str(q1))
 
     def test_delete_returning_str(self):
         q1 = (
@@ -56,4 +56,4 @@ class PostgresDeleteTests(unittest.TestCase):
             .returning("id")
         )
 
-        self.assertEqual('DELETE FROM "abc" WHERE "foo"="bar" RETURNING id', str(q1))
+        self.assertEqual('DELETE FROM "abc" WHERE "foo"="bar" RETURNING "id"', str(q1))
diff --git a/pypika/tests/test_inserts.py b/pypika/tests/test_inserts.py
index 7f76208..987fd19 100644
--- a/pypika/tests/test_inserts.py
+++ b/pypika/tests/test_inserts.py
@@ -1,6 +1,8 @@
 import unittest
 
 from pypika import (
+    AliasedQuery,
+    Case,
     Field as F,
     MySQLQuery,
     PostgreSQLQuery,
@@ -111,28 +113,6 @@ class InsertIntoTests(unittest.TestCase):
             'INSERT INTO "abc" ("foo","bar","buz") VALUES (1,\'a\',true)', str(query)
         )
 
-    def test_insert_selected_columns_type_list(self):
-        query = (
-            Query.into(self.table_abc)
-            .columns([self.table_abc.foo, self.table_abc.bar, self.table_abc.buz])
-            .insert(1, "a", True)
-        )
-
-        self.assertEqual(
-            'INSERT INTO "abc" ("foo","bar","buz") VALUES (1,\'a\',true)', str(query)
-        )
-
-    def test_insert_selected_columns_type_tuple(self):
-        query = (
-            Query.into(self.table_abc)
-            .columns((self.table_abc.foo, self.table_abc.bar, self.table_abc.buz))
-            .insert(1, "a", True)
-        )
-
-        self.assertEqual(
-            'INSERT INTO "abc" ("foo","bar","buz") VALUES (1,\'a\',true)', str(query)
-        )
-
     def test_insert_none_skipped(self):
         query = Query.into(self.table_abc).insert()
 
@@ -158,6 +138,19 @@ class InsertIntoTests(unittest.TestCase):
             "INSERT INTO \"abc\" VALUES (1,'a',true),(2,'b',false)", str(q)
         )
 
+    def test_insert_with_statement(self):
+        sub_query = Query().select(self.table_abc.id).from_(self.table_abc)
+        aliased = AliasedQuery('sub_qs')
+
+        q = (
+            Query().with_(sub_query, 'sub_qs')
+            .into(self.table_abc)
+            .select(aliased.id)
+            .from_(aliased)
+        )
+        self.assertEqual(
+            'WITH sub_qs AS (SELECT "id" FROM "abc") INSERT INTO "abc" SELECT "sub_qs"."id" FROM sub_qs', str(q))
+
 
 class PostgresInsertIntoOnConflictTests(unittest.TestCase):
     table_abc = Table("abc")
@@ -171,7 +164,19 @@ class PostgresInsertIntoOnConflictTests(unittest.TestCase):
         )
 
         self.assertEqual(
-            'INSERT INTO "abc" VALUES (1) ON CONFLICT (id) DO NOTHING', str(query)
+            'INSERT INTO "abc" VALUES (1) ON CONFLICT ("id") DO NOTHING', str(query)
+        )
+
+    def test_insert_on_conflict_do_nothing_multiple_fields(self):
+        query = (
+            PostgreSQLQuery.into(self.table_abc)
+                .insert(1)
+                .on_conflict(self.table_abc.id, self.table_abc.sub_id)
+                .do_nothing()
+        )
+
+        self.assertEqual(
+            'INSERT INTO "abc" VALUES (1) ON CONFLICT ("id", "sub_id") DO NOTHING', str(query)
         )
 
     def test_insert_on_conflict_do_nothing_field_str(self):
@@ -183,7 +188,31 @@ class PostgresInsertIntoOnConflictTests(unittest.TestCase):
         )
 
         self.assertEqual(
-            'INSERT INTO "abc" VALUES (1) ON CONFLICT (id) DO NOTHING', str(query)
+            'INSERT INTO "abc" VALUES (1) ON CONFLICT ("id") DO NOTHING', str(query)
+        )
+
+    def test_insert_on_conflict_do_nothing_multiple_fields_str(self):
+        query = (
+            PostgreSQLQuery.into(self.table_abc)
+            .insert(1)
+            .on_conflict("id", "sub_id")
+            .do_nothing()
+        )
+
+        self.assertEqual(
+            'INSERT INTO "abc" VALUES (1) ON CONFLICT ("id", "sub_id") DO NOTHING', str(query)
+        )
+
+    def test_insert_on_conflict_do_nothing_mixed_fields(self):
+        query = (
+            PostgreSQLQuery.into(self.table_abc)
+            .insert(1)
+            .on_conflict("id", self.table_abc.sub_id)
+            .do_nothing()
+        )
+
+        self.assertEqual(
+            'INSERT INTO "abc" VALUES (1) ON CONFLICT ("id", "sub_id") DO NOTHING', str(query)
         )
 
     def test_insert_on_conflict_do_update_field(self):
@@ -195,7 +224,20 @@ class PostgresInsertIntoOnConflictTests(unittest.TestCase):
         )
 
         self.assertEqual(
-            "INSERT INTO \"abc\" VALUES (1,'m') ON CONFLICT (id) DO UPDATE SET name='m'",
+            "INSERT INTO \"abc\" VALUES (1,'m') ON CONFLICT (\"id\") DO UPDATE SET \"name\"='m'",
+            str(query),
+        )
+
+    def test_insert_on_conflict_do_update_multiple_fields(self):
+        query = (
+            PostgreSQLQuery.into(self.table_abc)
+            .insert(1, "m")
+            .on_conflict(self.table_abc.id, self.table_abc.sub_id)
+            .do_update(self.table_abc.name, "m")
+        )
+
+        self.assertEqual(
+            "INSERT INTO \"abc\" VALUES (1,'m') ON CONFLICT (\"id\", \"sub_id\") DO UPDATE SET \"name\"='m'",
             str(query),
         )
 
@@ -208,7 +250,33 @@ class PostgresInsertIntoOnConflictTests(unittest.TestCase):
         )
 
         self.assertEqual(
-            "INSERT INTO \"abc\" VALUES (1,'m') ON CONFLICT (id) DO UPDATE SET name='m'",
+            "INSERT INTO \"abc\" VALUES (1,'m') ON CONFLICT (\"id\") DO UPDATE SET \"name\"='m'",
+            str(query),
+        )
+
+    def test_insert_on_conflict_do_update_multiple_fields_str(self):
+        query = (
+            PostgreSQLQuery.into(self.table_abc)
+            .insert(1, "m")
+            .on_conflict("id", "sub_id")
+            .do_update("name", "m")
+        )
+
+        self.assertEqual(
+            "INSERT INTO \"abc\" VALUES (1,'m') ON CONFLICT (\"id\", \"sub_id\") DO UPDATE SET \"name\"='m'",
+            str(query),
+        )
+
+    def test_insert_on_conflict_do_update_multiple_mixed_fields(self):
+        query = (
+            PostgreSQLQuery.into(self.table_abc)
+            .insert(1, "m")
+            .on_conflict("id", self.table_abc.sub_id)
+            .do_update("name", "m")
+        )
+
+        self.assertEqual(
+            "INSERT INTO \"abc\" VALUES (1,'m') ON CONFLICT (\"id\", \"sub_id\") DO UPDATE SET \"name\"='m'",
             str(query),
         )
 
@@ -279,6 +347,144 @@ class PostgresInsertIntoOnConflictTests(unittest.TestCase):
                 .do_update(self.table_abc.name, "m")
             )
 
+    def test_on_conflict_from_subquery(self):
+        table_bcd = Table('bcd')
+        query = (
+            PostgreSQLQuery.into(self.table_abc)
+                .insert(self.table_abc.fname, self.table_abc.lname)
+                .select(table_bcd.fname, table_bcd.lname)
+                .from_(table_bcd)
+                .on_conflict(self.table_abc.id, self.table_abc.sub_id)
+                .do_update(self.table_abc.fname, 1)
+                .do_update(self.table_abc.lname, table_bcd.lname)
+                .do_update(self.table_abc.cname,
+                           Case().when(self.table_abc.cname.eq('cname'), 'new_name').else_(self.table_abc.cname))
+        )
+
+        self.assertEqual(
+            'INSERT INTO "abc" VALUES ("fname","lname") '
+            'ON CONFLICT ("id", "sub_id") '
+            'DO UPDATE SET "fname"=1,"lname"="bcd"."lname",'
+            '"cname"=CASE WHEN "abc"."cname"=\'cname\' THEN \'new_name\' ELSE "abc"."cname" END', str(query)
+        )
+
+    def test_on_conflict_where_empty_conflict_fields_do_nothing(self):
+        with self.assertRaises(QueryException):
+            (
+                PostgreSQLQuery.into(self.table_abc)
+                    .insert(1, "m")
+                    .on_conflict()
+                    .where(self.table_abc.abc.eq(0))
+                    .where(self.table_abc.cde.eq(0))
+                    .do_nothing()
+            )
+
+    def test_on_conflict_where_conflict_fields_do_nothing(self):
+        qs = (
+            PostgreSQLQuery.into(self.table_abc)
+                .insert(1, "m")
+                .on_conflict('id')
+                .where(self.table_abc.abc.eq(0))
+                .where(self.table_abc.cde.eq(0))
+                .do_nothing()
+        )
+        self.assertEqual(
+            '''INSERT INTO "abc" VALUES (1,'m') ON CONFLICT ("id") WHERE "abc"=0 AND "cde"=0 DO NOTHING''', str(qs))
+
+    def test_on_conflict_where_empty_conflict_fields_do_update(self):
+        with self.assertRaises(QueryException):
+            (
+                PostgreSQLQuery.into(self.table_abc)
+                    .insert(1, "m")
+                    .on_conflict()
+                    .where(self.table_abc.abc.eq(0))
+                    .where(self.table_abc.cde.eq(0))
+                    .do_update('field', 'val')
+            )
+
+    def test_on_conflict_where_conflict_fields_do_update(self):
+        qs = (
+            PostgreSQLQuery.into(self.table_abc)
+                .insert(1, "m")
+                .on_conflict('id')
+                .where(self.table_abc.abc.eq(0))
+                .where(self.table_abc.cde.eq(0))
+                .do_update('field', 'val')
+        )
+        self.assertEqual(
+            '''INSERT INTO "abc" VALUES (1,'m') ON CONFLICT ("id") WHERE "abc"=0 AND "cde"=0 '''
+            '''DO UPDATE SET "field"='val\'''', str(qs))
+
+    def test_where_and_on_conflict_where(self):
+        table_bcd = Table('bcd')
+
+        qs = (
+            PostgreSQLQuery.into(self.table_abc)
+                .select(table_bcd.abc)
+                .from_(table_bcd)
+                .where(table_bcd.abc.eq('1'))
+                .on_conflict('id')
+                .where(self.table_abc.abc.eq(0))
+                .where(self.table_abc.cde.eq(0))
+                .do_update('field', 'val')
+        )
+
+        self.assertEqual(
+            'INSERT INTO "abc" SELECT "abc" FROM "bcd" WHERE "abc"=\'1\' '
+            'ON CONFLICT ("id") WHERE "abc"=0 AND "cde"=0 DO UPDATE SET "field"=\'val\'', str(qs))
+
+    def test_on_conflict_do_nothing_where(self):
+        with self.assertRaises(QueryException):
+            (
+                PostgreSQLQuery.into(self.table_abc)
+                    .insert(1, "m")
+                    .on_conflict()
+                    .do_nothing()
+                    .where(self.table_abc.abc.eq(1))
+            )
+
+    def test_empty_on_conflict_do_update_where(self):
+        with self.assertRaises(QueryException):
+            (
+                PostgreSQLQuery.into(self.table_abc)
+                    .insert(1, "m")
+                    .on_conflict()
+                    .do_update('abc', 1)
+                    .where(self.table_abc.abc.eq(1))
+            )
+
+    def test_on_conflict_do_update_where(self):
+        qs = (
+            PostgreSQLQuery.into(self.table_abc)
+                .insert(1, "m")
+                .on_conflict("id")
+                .do_update('abc', 1)
+                .where(self.table_abc.abc.eq(1))
+        )
+        self.assertEqual(
+            'INSERT INTO "abc" VALUES (1,\'m\') ON CONFLICT ("id") DO UPDATE SET "abc"=1 WHERE "abc"."abc"=1', str(qs))
+
+    def test_on_conflict_where_complex(self):
+        table_bcd = Table('bcd')
+
+        qs = (
+            PostgreSQLQuery.into(self.table_abc)
+                .select(table_bcd.abc)
+                .from_(table_bcd)
+                .where(table_bcd.abc.eq('1'))
+                .on_conflict('id')
+                .where(self.table_abc.abc.eq(0))
+                .where(self.table_abc.cde.eq(0))
+                .do_update('field', 'val')
+                .where(self.table_abc.id.eq(2))
+                .where(self.table_abc.sub_id.eq(3))
+        )
+        self.assertEqual(
+            'INSERT INTO "abc" SELECT "abc" FROM "bcd" WHERE "abc"=\'1\' '
+            'ON CONFLICT ("id") WHERE "abc"=0 AND "cde"=0 '
+            'DO UPDATE SET "field"=\'val\' WHERE "abc"."id"=2 AND "abc"."sub_id"=3',
+            str(qs))
+
 
 class PostgresInsertIntoReturningTests(unittest.TestCase):
     table_abc = Table("abc")
@@ -288,12 +494,12 @@ class PostgresInsertIntoReturningTests(unittest.TestCase):
             PostgreSQLQuery.into(self.table_abc).insert(1).returning(self.table_abc.id)
         )
 
-        self.assertEqual('INSERT INTO "abc" VALUES (1) RETURNING id', str(query))
+        self.assertEqual('INSERT INTO "abc" VALUES (1) RETURNING "id"', str(query))
 
     def test_insert_returning_one_field_str(self):
         query = PostgreSQLQuery.into(self.table_abc).insert(1).returning("id")
 
-        self.assertEqual('INSERT INTO "abc" VALUES (1) RETURNING id', str(query))
+        self.assertEqual('INSERT INTO "abc" VALUES (1) RETURNING "id"', str(query))
 
     def test_insert_returning_all_fields(self):
         query = (
@@ -311,7 +517,7 @@ class PostgresInsertIntoReturningTests(unittest.TestCase):
             .returning(self.table_abc.star, self.table_abc.f1 + self.table_abc.f2)
         )
 
-        self.assertEqual('INSERT INTO "abc" VALUES (1) RETURNING *,f1+f2', str(query))
+        self.assertEqual('INSERT INTO "abc" VALUES (1) RETURNING *,"f1"+"f2"', str(query))
 
     def test_insert_all_columns_multi_rows_chained_returning_star(self):
         query = (
@@ -380,7 +586,7 @@ class PostgresInsertIntoReturningTests(unittest.TestCase):
             .returning(self.table_abc.f1 + self.table_abc.f2)
         )
 
-        self.assertEqual('INSERT INTO "abc" VALUES (1) RETURNING f1+f2', str(query))
+        self.assertEqual('INSERT INTO "abc" VALUES (1) RETURNING "f1"+"f2"', str(query))
 
     def test_insert_returning_aggregate(self):
         with self.assertRaises(QueryException):
diff --git a/pypika/tests/test_updates.py b/pypika/tests/test_updates.py
index 6fa8a99..ebbceac 100644
--- a/pypika/tests/test_updates.py
+++ b/pypika/tests/test_updates.py
@@ -1,6 +1,6 @@
 import unittest
 
-from pypika import Table, Query, PostgreSQLQuery
+from pypika import Table, Query, PostgreSQLQuery, AliasedQuery
 
 __author__ = "Timothy Heys"
 __email__ = "theys@kayak.com"
@@ -43,7 +43,7 @@ class UpdateTests(unittest.TestCase):
             .set(self.table_abc.lname, self.table_def.lname)
         )
         self.assertEqual(
-            'UPDATE "abc" JOIN "def" ON "def"."abc_id"="abc"."id" SET "abc"."lname"="def"."lname"',
+            'UPDATE "abc" JOIN "def" ON "def"."abc_id"="abc"."id" SET "lname"="def"."lname"',
             str(q),
         )
 
@@ -51,6 +51,47 @@ class UpdateTests(unittest.TestCase):
         q = Query.update(self.table_abc).set(self.table_abc.lname, "test").limit(1)
         self.assertEqual('UPDATE "abc" SET "lname"=\'test\' LIMIT 1', str(q))
 
+    def test_update_from(self):
+        from_table = Table('from_table')
+
+        q = Query.update(self.table_abc).set(self.table_abc.lname, from_table.long_name).from_(from_table)
+        self.assertEqual('UPDATE "abc" SET "lname"="from_table"."long_name" FROM "from_table"', str(q))
+
+    def test_update_from_with_where(self):
+        from_table = Table('from_table')
+
+        q = (
+            Query.update(self.table_abc)
+            .set(self.table_abc.lname, from_table.long_name)
+            .from_(from_table)
+            .where(self.table_abc.fname.eq(from_table.full_name))
+        )
+        self.assertEqual(
+            'UPDATE "abc" SET "lname"="from_table"."long_name" FROM "from_table" '
+            'WHERE "abc"."fname"="from_table"."full_name"',
+            str(q)
+        )
+
+    def test_update_with_statement(self):
+        table_efg = Table('efg')
+
+        sub_query = Query.from_(table_efg).select("fizz")
+        an_alias = AliasedQuery("an_alias")
+
+        q = (
+            Query.with_(sub_query, "an_alias")
+            .update(self.table_abc)
+            .from_(an_alias)
+            .set(self.table_abc.lname, an_alias.long_name)
+            .where(self.table_abc.comp.eq(an_alias.alias_comp))
+        )
+        self.assertEqual(
+            'WITH an_alias AS (SELECT "fizz" FROM "efg") '
+            'UPDATE "abc" SET "lname"="an_alias"."long_name" FROM an_alias '
+            'WHERE "abc"."comp"="an_alias"."alias_comp"',
+            str(q)
+        )
+
 
 class PostgresUpdateTests(unittest.TestCase):
     table_abc = Table("abc")
@@ -64,7 +105,7 @@ class PostgresUpdateTests(unittest.TestCase):
         )
 
         self.assertEqual(
-            'UPDATE "abc" SET "foo"=\'bar\' WHERE "foo"=0 RETURNING id', str(q)
+            'UPDATE "abc" SET "foo"=\'bar\' WHERE "foo"=0 RETURNING "abc"."id"', str(q)
         )
 
     def test_update_returning(self):
@@ -76,5 +117,18 @@ class PostgresUpdateTests(unittest.TestCase):
         )
 
         self.assertEqual(
-            'UPDATE "abc" SET "foo"=\'bar\' WHERE "foo"=0 RETURNING id', str(q)
+            'UPDATE "abc" SET "foo"=\'bar\' WHERE "foo"=0 RETURNING "abc"."id"', str(q)
+        )
+
+    def test_update_returning_from_different_tables(self):
+        table_bcd = Table('bcd')
+
+        q = (
+            PostgreSQLQuery.update(self.table_abc)
+            .from_(table_bcd)
+            .set(self.table_abc.lname, table_bcd.long_name)
+            .returning(self.table_abc.id, table_bcd.fname)
+        )
+        self.assertEqual(
+            'UPDATE "abc" SET "lname"="bcd"."long_name" FROM "bcd" RETURNING "abc"."id","bcd"."fname"', str(q)
         )