Merge pull request #630 from squareapartments/master

Add check if function is aggregate in returning

Author: x8lucas8x

pypika/dialects.py

@@ -557,10 +557,10 @@ def returning(self, *terms: Any) -> "PostgreSQLQueryBuilder":
                 self._return_field(term)
             elif isinstance(term, str):
                 self._return_field_str(term)
-            elif isinstance(term, ArithmeticExpression):
+            elif isinstance(term, (Function, ArithmeticExpression)):
+                if term.is_aggregate:
+                    raise QueryException("Aggregate functions are not allowed in returning")
                 self._return_other(term)
-            elif isinstance(term, Function):
-                raise QueryException("Aggregate functions are not allowed in returning")
             else:
                 self._return_other(self.wrap_constant(term, self._wrapper_cls))
 

pypika/tests/test_inserts.py

@@ -11,7 +11,10 @@
     Tables,
     functions as fn,
 )
-from pypika.functions import Avg
+from pypika.functions import (
+    Avg,
+    Cast,
+)
 from pypika.terms import Values
 from pypika.utils import QueryException
 
@@ -550,6 +553,11 @@ def test_insert_returning_arithmetics(self):
 
         self.assertEqual('INSERT INTO "abc" VALUES (1) RETURNING "f1"+"f2"', str(query))
 
+    def test_insert_returning_functions(self):
+        query = PostgreSQLQuery.into(self.table_abc).insert(1).returning(Cast(self.table_abc.f1, "int"))
+
+        self.assertEqual('INSERT INTO "abc" VALUES (1) RETURNING CAST("f1" AS INT)', str(query))
+
     def test_insert_returning_aggregate(self):
         with self.assertRaises(QueryException):
             PostgreSQLQuery.into(self.table_abc).insert(1).returning(Avg(self.table_abc.views))