Repository: kayak/pypika
Pull request #103 (2018-03-16 15:13:04)

RETURNING clause for postgresql
This PR implements postgres `RETURNING` clause.
This PR implements it directly in QueryBuilder, which could be quite misleading, for other sql users, so if you have any propositions regarding how this code should be arranged I can do some refactoring

coveralls: 
[![Coverage Status](https://coveralls.io/builds/16017590/badge)](https://coveralls.io/builds/16017590)

Coverage decreased (-0.5%) to 97.302% when pulling **2ed0d162d38c8a3cd3bc44c0a8b4fa7cd6cf797c on Zeliboba5:master** into **08c8effeba292c28476e98a4048b377a4f3a1341 on kayak:master**.

coveralls: 
[![Coverage Status](https://coveralls.io/builds/16017590/badge)](https://coveralls.io/builds/16017590)

Coverage decreased (-0.5%) to 97.302% when pulling **2ed0d162d38c8a3cd3bc44c0a8b4fa7cd6cf797c on Zeliboba5:master** into **08c8effeba292c28476e98a4048b377a4f3a1341 on kayak:master**.

twheys: If you'd like to add a vendor specific function you can override the querybuilder class and instantiate it in the PostgreSQLQuery class' _builder function. See my last PR for an example: https://github.com/kayak/pypika/pull/109
twheys: Sorry, not sure how I closed this, didn't mean to. Re-opened.
Zeliboba5: Moved everything to postgres QueryBuilder, added some more tests and covered some more clauses, now it all looks neat.
Regarding consolidation of `returning` and `select` I think it is wrong way to go, as it still different functions, which are not fully compatible (`RETURNING` does not support aggregate functions, so I covered it in my new commits)  
twheys: Looks good to go, just need to fix the merge conflicts. I'll merge once that's done, let me know if you need help with it.

Patch:
diff --git a/pypika/dialects.py b/pypika/dialects.py
index c5aa0be..b2ea3b0 100644
--- a/pypika/dialects.py
+++ b/pypika/dialects.py
@@ -1,9 +1,7 @@
-from .utils import builder
-from .enums import Dialects
-from .queries import (
-    Query,
-    QueryBuilder,
-)
+from pypika.terms import ArithmeticExpression, Function, Field, Star
+from pypika.utils import builder, QueryException
+from pypika.enums import Dialects
+from pypika.queries import Query, QueryBuilder
 
 
 class MySQLQuery(Query):
@@ -56,6 +54,86 @@ class OracleQuery(Query):
         return QueryBuilder(dialect=Dialects.ORACLE)
 
 
+class PostgreQueryBuilder(QueryBuilder):
+    def __init__(self):
+        super(PostgreQueryBuilder, self).__init__(dialect=Dialects.POSTGRESQL)
+        self._returns = []
+        self._return_star = False
+
+    @builder
+    def returning(self, *terms):
+        for term in terms:
+            if isinstance(term, Field):
+                self._return_field(term)
+            elif isinstance(term, str):
+                self._return_field_str(term)
+            elif isinstance(term, ArithmeticExpression):
+                self._return_other(term)
+            elif isinstance(term, Function):
+                raise QueryException('Aggregate functions are not allowed in returning')
+            else:
+                self._return_other(self._wrap(term))
+
+    def _validate_returning_term(self, term):
+        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
+            ):
+                raise QueryException('You can\'t return from other tables')
+
+    def _set_returns_for_star(self):
+        self._returns = [returning
+                         for returning in self._returns
+                         if not hasattr(returning, 'table')]
+        self._return_star = True
+
+    def _return_field(self, term):
+        if self._return_star:
+            # Do not add select terms after a star is selected
+            return
+
+        self._validate_returning_term(term)
+
+        if isinstance(term, Star):
+            self._set_returns_for_star()
+
+        self._returns.append(term)
+
+    def _return_field_str(self, term):
+        if term == '*':
+            self._set_returns_for_star()
+            self._returns.append(Star())
+            return
+
+        if self._insert_table:
+            self._return_field(Field(term, table=self._insert_table))
+        elif self._update_table:
+            self._return_field(Field(term, table=self._update_table))
+        elif self._delete_from:
+            self._return_field(Field(term, table=self._from[0]))
+        else:
+            raise QueryException('Returning can\'t be used in this query')
+
+    def _return_other(self, function):
+        self._validate_returning_term(function)
+        self._returns.append(function)
+
+    def _returning_sql(self, **kwargs):
+        return ' RETURNING {returning}'.format(
+            returning=','.join(term.get_sql(with_alias=True, **kwargs)
+                               for term in self._returns),
+        )
+
+    def get_sql(self, with_alias=False, subquery=False, **kwargs):
+        querystring = super(PostgreQueryBuilder, self).get_sql(with_alias, subquery, **kwargs)
+        if self._returns:
+            querystring += self._returning_sql()
+        return querystring
+
+
 class PostgreSQLQuery(Query):
     """
     Defines a query class for use with PostgreSQL.
@@ -63,7 +141,7 @@ class PostgreSQLQuery(Query):
 
     @classmethod
     def _builder(cls):
-        return QueryBuilder(dialect=Dialects.POSTGRESQL)
+        return PostgreQueryBuilder()
 
 
 class RedshiftQuery(Query):
diff --git a/pypika/queries.py b/pypika/queries.py
index 1e31fe3..2883663 100644
--- a/pypika/queries.py
+++ b/pypika/queries.py
@@ -647,6 +647,7 @@ class QueryBuilder(Selectable, Term):
 
             if self._wheres:
                 querystring += self._where_sql(**kwargs)
+
             return querystring
         elif self._delete_from:
             querystring = self._delete_sql(**kwargs)
@@ -657,7 +658,8 @@ class QueryBuilder(Selectable, Term):
                 querystring += self._columns_sql(**kwargs)
 
             if self._values:
-                return querystring + self._values_sql(**kwargs)
+                querystring += self._values_sql(**kwargs)
+                return querystring
             else:
                 querystring += ' ' + self._select_sql(**kwargs)
         else:


Test patch:
diff --git a/pypika/tests/test_deletes.py b/pypika/tests/test_deletes.py
index 4b7f3b0..4196b85 100644
--- a/pypika/tests/test_deletes.py
+++ b/pypika/tests/test_deletes.py
@@ -1,7 +1,7 @@
 # coding: utf8
 import unittest
 
-from pypika import Table, Query
+from pypika import Table, Query, PostgreSQLQuery
 
 __author__ = "Timothy Heys"
 __email__ = "theys@kayak.com"
@@ -26,3 +26,21 @@ class DeleteTests(unittest.TestCase):
 
         self.assertEqual('DELETE FROM "abc" WHERE "foo"="bar"', str(q1))
         self.assertEqual('DELETE FROM "abc" WHERE "foo"="bar"', str(q2))
+
+
+class PostgresDeleteTests(unittest.TestCase):
+    table_abc = Table('abc')
+
+    def test_delete_returning(self):
+        q1 = PostgreSQLQuery.from_(self.table_abc).where(
+            self.table_abc.foo == self.table_abc.bar
+        ).delete().returning(self.table_abc.id)
+
+        self.assertEqual('DELETE FROM "abc" WHERE "foo"="bar" RETURNING id', str(q1))
+
+    def test_delete_returning_str(self):
+        q1 = PostgreSQLQuery.from_(self.table_abc).where(
+            self.table_abc.foo == self.table_abc.bar
+        ).delete().returning('id')
+
+        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 9391561..1712d4a 100644
--- a/pypika/tests/test_inserts.py
+++ b/pypika/tests/test_inserts.py
@@ -1,7 +1,9 @@
 # coding: utf8
 import unittest
 
-from pypika import Table, Tables, Query
+from pypika import Table, Tables, Query, PostgreSQLQuery
+from pypika.functions import Avg
+from pypika.utils import QueryException
 
 __author__ = "Timothy Heys"
 __email__ = "theys@kayak.com"
@@ -88,6 +90,87 @@ class InsertIntoTests(unittest.TestCase):
         self.assertEqual('INSERT INTO "abc" VALUES (NULL)', str(query))
 
 
+class PostgresInsertIntoReturningTests(unittest.TestCase):
+    table_abc = Table('abc')
+
+    def test_insert_returning_one_field(self):
+        query = PostgreSQLQuery.into(self.table_abc).insert(1).returning(self.table_abc.id)
+
+        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))
+
+    def test_insert_returning_all_fields(self):
+        query = PostgreSQLQuery.into(self.table_abc).insert(1).returning(self.table_abc.star)
+
+        self.assertEqual('INSERT INTO "abc" VALUES (1) RETURNING *', str(query))
+
+    def test_insert_returning_all_fields_and_arithmetics(self):
+        query = PostgreSQLQuery.into(self.table_abc).insert(1).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))
+
+    def test_insert_all_columns_multi_rows_chained_returning_star(self):
+        query = PostgreSQLQuery.into(
+            self.table_abc
+        ).insert(1, 'a', True).insert(2, 'b', False).returning(self.table_abc.star)
+
+        self.assertEqual('INSERT INTO "abc" VALUES (1,\'a\',true),(2,\'b\',false) RETURNING *', str(query))
+
+    def test_insert_all_columns_multi_rows_chained_returning_star_and_id(self):
+        query = PostgreSQLQuery.into(
+            self.table_abc
+        ).insert(1, 'a', True).insert(2, 'b', False).returning(
+            self.table_abc.name,
+            self.table_abc.star,
+            self.table_abc.id,
+        )
+
+        self.assertEqual('INSERT INTO "abc" VALUES (1,\'a\',true),(2,\'b\',false) RETURNING *', str(query))
+
+    def test_insert_all_columns_multi_rows_chained_returning_star_str(self):
+        query = PostgreSQLQuery.into(
+            self.table_abc
+        ).insert(1, 'a', True).insert(2, 'b', False).returning('*')
+
+        self.assertEqual('INSERT INTO "abc" VALUES (1,\'a\',true),(2,\'b\',false) RETURNING *', str(query))
+
+    def test_insert_all_columns_single_element_arrays(self):
+        query = PostgreSQLQuery.into(self.table_abc).insert((1, 'a', True)).returning(self.table_abc.star)
+
+        self.assertEqual('INSERT INTO "abc" VALUES (1,\'a\',true) RETURNING *', str(query))
+
+    def test_insert_returning_null(self):
+        query = PostgreSQLQuery.into(self.table_abc).insert(1).returning(None)
+
+        self.assertEqual('INSERT INTO "abc" VALUES (1) RETURNING NULL', str(query))
+
+    def test_insert_returning_tuple(self):
+        query = PostgreSQLQuery.into(self.table_abc).insert(1).returning([1, 2, 3])
+
+        self.assertEqual('INSERT INTO "abc" VALUES (1) RETURNING (1,2,3)', str(query))
+
+    def test_insert_returning_arithmetics(self):
+        query = PostgreSQLQuery.into(self.table_abc).insert(1).returning(self.table_abc.f1 + self.table_abc.f2)
+
+        self.assertEqual('INSERT INTO "abc" VALUES (1) RETURNING f1+f2', 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))
+
+    def test_insert_returning_from_other_table(self):
+        table_cba = Table('cba')
+        with self.assertRaises(QueryException):
+            PostgreSQLQuery.into(self.table_abc).insert(1).returning(table_cba.id)
+
+
 class InsertSelectFromTests(unittest.TestCase):
     table_abc, table_efg, table_hij = Tables('abc', 'efg', 'hij')
 
diff --git a/pypika/tests/test_updates.py b/pypika/tests/test_updates.py
index 50d6c44..87a183e 100644
--- a/pypika/tests/test_updates.py
+++ b/pypika/tests/test_updates.py
@@ -1,7 +1,7 @@
 # coding: utf8
 import unittest
 
-from pypika import Table, Query
+from pypika import Table, Query, PostgreSQLQuery
 
 __author__ = "Timothy Heys"
 __email__ = "theys@kayak.com"
@@ -29,3 +29,25 @@ class UpdateTests(unittest.TestCase):
     def test_update_with_none(self):
         q = Query.update('abc').set('foo', None)
         self.assertEqual('UPDATE "abc" SET "foo"=null', str(q))
+
+
+class PostgesUpdateTests(unittest.TestCase):
+    table_abc = Table('abc')
+
+    def test_update_returning_str(self):
+        q = PostgreSQLQuery.update(self.table_abc).where(
+            self.table_abc.foo == 0
+        ).set('foo', 'bar').returning('id')
+
+        self.assertEqual('UPDATE "abc" SET "foo"=\'bar\' WHERE "foo"=0 RETURNING id', str(q))
+
+
+class PostgresUpdateTests(unittest.TestCase):
+    table_abc = Table('abc')
+
+    def test_update_returning(self):
+        q = PostgreSQLQuery.update(self.table_abc).where(
+            self.table_abc.foo == 0
+        ).set('foo', 'bar').returning(self.table_abc.id)
+
+        self.assertEqual('UPDATE "abc" SET "foo"=\'bar\' WHERE "foo"=0 RETURNING id', str(q))