Repository: kayak/pypika
Pull request #536 (2021-01-04 10:33:21)

Fix PostgreSQLQueryBuilder returning clause code causing Py3.9 test failures
This PR fixes the Python 3.9 test failures in the Github Actions CI. 

This was previously passing due to this statement: `if field.table not in {self._insert_table, self._update_table} and term not in self._from`, specifically the second statement. In my Python 3.6 environment, the Term `__eq__` method was called, always returning a truthy value (as this returns a BasicCriterion object, intended for query building). In Python 3.9, this is not called (I couldn't find anything in the release notes though). Therefore, I have rewritten this to check if the field's table is one of the base and join tables.

FYI @abondar 


Patch:
diff --git a/pypika/dialects.py b/pypika/dialects.py
index 3aa847c..7a12450 100644
--- a/pypika/dialects.py
+++ b/pypika/dialects.py
@@ -1,3 +1,4 @@
+import itertools
 from copy import copy
 from typing import Any, Optional, Union
 
@@ -485,7 +486,12 @@ class PostgreSQLQueryBuilder(QueryBuilder):
         for field in term.fields_():
             if not any([self._insert_table, self._update_table, self._delete_from]):
                 raise QueryException("Returning can't be used in this query")
-            if field.table not in {self._insert_table, self._update_table} and term not in self._from:
+
+            table_is_insert_or_update_table = field.table in {self._insert_table, self._update_table}
+            join_tables = set(itertools.chain.from_iterable([j.criterion.tables_ for j in self._joins]))
+            join_and_base_tables = set(self._from) | join_tables
+            table_not_base_or_join = bool(term.tables_ - join_and_base_tables)
+            if not table_is_insert_or_update_table and table_not_base_or_join:
                 raise QueryException("You can't return from other tables")
 
     def _set_returns_for_star(self) -> None:


Test patch:
diff --git a/pypika/tests/dialects/test_postgresql.py b/pypika/tests/dialects/test_postgresql.py
index e364d4c..34ce021 100644
--- a/pypika/tests/dialects/test_postgresql.py
+++ b/pypika/tests/dialects/test_postgresql.py
@@ -3,7 +3,9 @@ from collections import OrderedDict
 
 from pypika import (
     Array,
+    Field,
     JSON,
+    QueryException,
     Table,
 )
 from pypika.dialects import PostgreSQLQuery
@@ -196,3 +198,56 @@ class ArrayTests(unittest.TestCase):
 
         q = PostgreSQLQuery.from_(tb).select(Array(tb.col).as_("different_name"))
         self.assertEqual(str(q), 'SELECT ARRAY["col"] "different_name" FROM "tb"')
+
+
+class ReturningClauseTests(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls) -> None:
+        super().setUpClass()
+        cls.table_abc = Table('abc')
+
+    def test_returning_from_missing_table_raises_queryexception(self):
+        field_from_diff_table = Field('xyz', table=Table('other'))
+
+        with self.assertRaisesRegex(QueryException, "You can't return from other tables"):
+            (
+                PostgreSQLQuery.from_(self.table_abc)
+                .where(self.table_abc.foo == self.table_abc.bar)
+                .delete()
+                .returning(field_from_diff_table)
+            )
+
+    def test_queryexception_if_returning_used_on_invalid_query(self):
+        with self.assertRaisesRegex(QueryException, "Returning can't be used in this query"):
+            PostgreSQLQuery.from_(self.table_abc).select('abc').returning('abc')
+
+    def test_no_queryexception_if_returning_used_on_valid_query_type(self):
+        # No exceptions for insert, update and delete queries
+        with self.subTest('DELETE'):
+            PostgreSQLQuery.from_(self.table_abc).where(self.table_abc.foo == self.table_abc.bar).delete().returning(
+                "id"
+            )
+        with self.subTest('UPDATE'):
+            PostgreSQLQuery.update(self.table_abc).where(self.table_abc.foo == 0).set("foo", "bar").returning("id")
+        with self.subTest('INSERT'):
+            PostgreSQLQuery.into(self.table_abc).insert('abc').returning('abc')
+
+    def test_return_field_from_join_table(self):
+        new_table = Table('xyz')
+        q = (
+            PostgreSQLQuery.update(self.table_abc)
+            .join(new_table)
+            .on(new_table.id == self.table_abc.xyz)
+            .where(self.table_abc.foo == 0)
+            .set("foo", "bar")
+            .returning(new_table.a)
+        )
+
+        self.assertEqual(
+            'UPDATE "abc" '
+            'JOIN "xyz" ON "xyz"."id"="abc"."xyz" '
+            'SET "foo"=\'bar\' '
+            'WHERE "abc"."foo"=0 '
+            'RETURNING "xyz"."a"',
+            str(q),
+        )
diff --git a/pypika/tests/test_deletes.py b/pypika/tests/test_deletes.py
index 61f1dbe..fadc060 100644
--- a/pypika/tests/test_deletes.py
+++ b/pypika/tests/test_deletes.py
@@ -1,13 +1,18 @@
 import unittest
 
-from pypika import Table, Query, PostgreSQLQuery, SYSTEM_TIME
+from pypika import PostgreSQLQuery, Query, SYSTEM_TIME, Table
 
 __author__ = "Timothy Heys"
 __email__ = "theys@kayak.com"
 
+from pypika.terms import Star
+
 
 class DeleteTests(unittest.TestCase):
-    table_abc = Table("abc")
+    @classmethod
+    def setUpClass(cls) -> None:
+        super().setUpClass()
+        cls.table_abc = Table("abc")
 
     def test_omit_where(self):
         q = Query.from_("abc").delete()
@@ -45,7 +50,10 @@ class DeleteTests(unittest.TestCase):
 
 
 class PostgresDeleteTests(unittest.TestCase):
-    table_abc = Table("abc")
+    @classmethod
+    def setUpClass(cls) -> None:
+        super().setUpClass()
+        cls.table_abc = Table("abc")
 
     def test_delete_returning(self):
         q1 = (
@@ -66,3 +74,13 @@ class PostgresDeleteTests(unittest.TestCase):
         )
 
         self.assertEqual('DELETE FROM "abc" WHERE "foo"="bar" RETURNING "id"', str(q1))
+
+    def test_delete_returning_star(self):
+        q1 = (
+            PostgreSQLQuery.from_(self.table_abc)
+            .where(self.table_abc.foo == self.table_abc.bar)
+            .delete()
+            .returning(Star())
+        )
+
+        self.assertEqual('DELETE FROM "abc" WHERE "foo"="bar" RETURNING *', str(q1))
diff --git a/pypika/tests/test_updates.py b/pypika/tests/test_updates.py
index b4ac2e4..fb9d714 100644
--- a/pypika/tests/test_updates.py
+++ b/pypika/tests/test_updates.py
@@ -1,10 +1,12 @@
 import unittest
 
-from pypika import Table, Query, PostgreSQLQuery, AliasedQuery, SQLLiteQuery, SYSTEM_TIME
+from pypika import AliasedQuery, PostgreSQLQuery, Query, SQLLiteQuery, SYSTEM_TIME, Table
 
 __author__ = "Timothy Heys"
 __email__ = "theys@kayak.com"
 
+from pypika.terms import Star
+
 
 class UpdateTests(unittest.TestCase):
     table_abc = Table("abc")
@@ -145,6 +147,11 @@ class PostgresUpdateTests(unittest.TestCase):
             'UPDATE "abc" SET "lname"="bcd"."long_name" FROM "bcd" RETURNING "abc"."id","bcd"."fname"', str(q)
         )
 
+    def test_update_returning_star(self):
+        q = PostgreSQLQuery.update(self.table_abc).where(self.table_abc.foo == 0).set("foo", "bar").returning(Star())
+
+        self.assertEqual('UPDATE "abc" SET "foo"=\'bar\' WHERE "foo"=0 RETURNING *', str(q))
+
 
 class SQLLiteUpdateTests(unittest.TestCase):
     table_abc = Table("abc")