Repository: kayak/pypika
Pull request #285 (2019-08-05 10:00:23)

Add replace table feature
Add feature to replace a specific table with another table in queries. This replaces the previous option to replace all tables in a query.


Patch:
diff --git a/pypika/queries.py b/pypika/queries.py
index d3da3a3..f25d9ba 100644
--- a/pypika/queries.py
+++ b/pypika/queries.py
@@ -462,7 +462,7 @@ class QueryBuilder(Selectable, Term):
     @builder
     def from_(self, selectable):
         """
-        Adds a table to the query.  This function can only be called once and will raise an AttributeError if called a
+        Adds a table to the query. This function can only be called once and will raise an AttributeError if called a
         second time.
 
         :param selectable:
@@ -486,6 +486,41 @@ class QueryBuilder(Selectable, Term):
             selectable.alias = 'sq%d' % sub_query_count
             self._subquery_count = sub_query_count + 1
 
+    @builder
+    def replace_table(self, current_table, new_table):
+        """
+        Replaces all occurrences of the specified table with the new table. Useful when reusing fields across
+        queries.
+
+        :param current_table:
+            The table instance to be replaces.
+        :param new_table:
+            The table instance to replace with.
+        :return:
+            A copy of the query with the tables replaced.
+        """
+        self._from = [new_table if table == current_table else table for table in self._from]
+        self._insert_table = new_table if self._insert_table else None
+        self._update_table = new_table if self._update_table else None
+
+        self._with = [alias_query.replace_table(current_table, new_table) for alias_query in self._with]
+        self._selects = [select.replace_table(current_table, new_table) for select in self._selects]
+        self._columns = [column.replace_table(current_table, new_table) for column in self._columns]
+        self._values = [[value.replace_table(current_table, new_table) for value in value_list]
+                        for value_list
+                        in self._values]
+
+        self._wheres = self._wheres.replace_table(current_table, new_table) if self._wheres else None
+        self._prewheres = self._prewheres.replace_table(current_table, new_table) if self._prewheres else None
+        self._groupbys = [groupby.replace_table(current_table, new_table) for groupby in self._groupbys]
+        self._havings = self._havings.replace_table(current_table, new_table) if self._havings else None
+        self._orderbys = [orderby.replace_table(current_table, new_table) for orderby in self._orderbys]
+        self._joins = [join.replace_table(current_table, new_table) for join in self._joins]
+
+        if current_table in self._select_star_tables:
+            self._select_star_tables.remove(current_table)
+            self._select_star_tables.add(new_table)
+
     @builder
     def with_(self, selectable, name):
         t = AliasedQuery(name, selectable)
@@ -788,7 +823,6 @@ class QueryBuilder(Selectable, Term):
                                  else self.wrap_constant(value)
                                  for value in values])
 
-
     def __str__(self):
         return self.get_sql(quote_char=self.quote_char, dialect=self.dialect)
 
@@ -969,7 +1003,7 @@ class QueryBuilder(Selectable, Term):
             .format(values='),('
                     .join(','
                           .join(term.get_sql(with_alias=True, subquery=True, **kwargs)
-                                   for term in row)
+                                for term in row)
                           for row in self._values))
 
     def _into_sql(self, **kwargs):
@@ -1121,6 +1155,21 @@ class Join:
     def validate(self, _from, _joins):
         pass
 
+    @builder
+    def replace_table(self, current_table, new_table):
+        """
+        Replaces all occurrences of the specified table with the new table. Useful when reusing
+        fields across queries.
+
+        :param current_table:
+            The table to be replaced.
+        :param new_table:
+            The table to replace with.
+        :return:
+            A copy of the join with the tables replaced.
+        """
+        self.item = self.item.replace_table(current_table, new_table)
+
 
 class JoinOn(Join):
     def __init__(self, item, how, criteria):
@@ -1144,6 +1193,22 @@ class JoinOn(Join):
                   tables=', '.join(map(str, missing_tables))
             ))
 
+    @builder
+    def replace_table(self, current_table, new_table):
+        """
+        Replaces all occurrences of the specified table with the new table. Useful when reusing
+        fields across queries.
+
+        :param current_table:
+            The table to be replaced.
+        :param new_table:
+            The table to replace with.
+        :return:
+            A copy of the join with the tables replaced.
+        """
+        self.item = new_table if self.item == current_table else self.item
+        self.criterion = self.criterion.replace_table(current_table, new_table)
+
 
 class JoinUsing(Join):
     def __init__(self, item, how, fields):
@@ -1159,3 +1224,19 @@ class JoinUsing(Join):
 
     def validate(self, _from, _joins):
         pass
+
+    @builder
+    def replace_table(self, current_table, new_table):
+        """
+        Replaces all occurrences of the specified table with the new table. Useful when reusing
+        fields across queries.
+
+        :param current_table:
+            The table to be replaced.
+        :param new_table:
+            The table to replace with.
+        :return:
+            A copy of the join with the tables replaced.
+        """
+        self.item = new_table if self.item == current_table else self.item
+        self.fields = [field.replace_table(current_table, new_table) for field in self.fields]
diff --git a/pypika/terms.py b/pypika/terms.py
index 127338c..e765011 100644
--- a/pypika/terms.py
+++ b/pypika/terms.py
@@ -76,12 +76,14 @@ class Term:
         wrapper_cls = wrapper_cls or ValueWrapper
         return wrapper_cls(val)
 
-    def for_(self, table):
+    def replace_table(self, current_table, new_table):
         """
-        Replaces the tables of this term for the table parameter provided.  The base implementation returns self
-        because not all terms have a table property.
+        Replaces all occurrences of the specified table with the new table. Useful when reusing fields across queries.
+        The base implementation returns self because not all terms have a table property.
 
-        :param table:
+        :param current_table:
+            The table to be replaced.
+        :param new_table:
             The table to replace with.
         :return:
             Self.
@@ -352,16 +354,18 @@ class Field(Criterion):
         return {self.table}
 
     @builder
-    def for_(self, table):
+    def replace_table(self, current_table, new_table):
         """
-        Replaces the tables of this term for the table parameter provided.  Useful when reusing fields across queries.
+        Replaces all occurrences of the specified table with the new table. Useful when reusing fields across queries.
 
-        :param table:
+        :param current_table:
+            The table to be replaced.
+        :param new_table:
             The table to replace with.
         :return:
-            A copy of the field with it's table value replaced.
+            A copy of the field with the tables replaced. 
         """
-        self.table = table
+        self.table = new_table if self.table == current_table else self.table
 
     def get_sql(self, with_alias=False, with_namespace=False, quote_char=None, **kwargs):
         # Need to add namespace if the table has an alias
@@ -425,6 +429,22 @@ class Tuple(Criterion):
         return all([value.is_aggregate
                     for value in self.values])
 
+    @builder
+    def replace_table(self, current_table, new_table):
+        """
+        Replaces all occurrences of the specified table with the new table. Useful when reusing fields across queries.
+
+        :param current_table:
+            The table to be replaced.
+        :param new_table:
+            The table to replace with.
+        :return:
+            A copy of the field with the tables replaced.
+        """
+        self.values = [[value.replace_table(current_table, new_table) for value in value_list]
+                       for value_list
+                       in self.values]
+
 
 class Array(Tuple):
     def get_sql(self, **kwargs):
@@ -471,8 +491,8 @@ class JSONField(Field):
             name = self.name
 
         return '{quote}{val}{quote}'.format(
-            quote=self.quote_char,
-            val=name,
+              quote=self.quote_char,
+              val=name,
         )
 
     def has_key(self, other):
@@ -508,8 +528,8 @@ class JSONField(Field):
             root_key = JSONField(other[0], quote_char=self.EXPRESSION_QUOTE_CHAR)
             nested_key = JSONField(other[1], quote_char=self.EXPRESSION_QUOTE_CHAR)
             return NestedCriterion(
-                PostgresOperators.GET_JSON_VALUE, PostgresOperators.GET_TEXT_VALUE,
-                self, root_key, nested_key,
+                  PostgresOperators.GET_JSON_VALUE, PostgresOperators.GET_TEXT_VALUE,
+                  self, root_key, nested_key,
             )
         else:
             raise QueryException('Type `{}` does not support for this operation'.format(type(other)))
@@ -549,18 +569,28 @@ class NestedCriterion(Criterion):
         return self.left.tables_ | self.right.tables_ | self.nested.tables_
 
     @builder
-    def for_(self, table):
-        self.left = self.left.for_(table)
-        self.right = self.right.for_(table)
-        self.nested = self.right.for_(table)
+    def replace_table(self, current_table, new_table):
+        """
+        Replaces all occurrences of the specified table with the new table. Useful when reusing fields across queries.
+
+        :param current_table:
+            The table to be replaced.
+        :param new_table:
+            The table to replace with.
+        :return:
+            A copy of the criterion with the tables replaced. 
+        """
+        self.left = self.left.replace_table(current_table, new_table)
+        self.right = self.right.replace_table(current_table, new_table)
+        self.nested = self.right.replace_table(current_table, new_table)
 
     def get_sql(self, with_alias=False, **kwargs):
         sql = '{left}{comparator}{right}{nested_comparator}{nested}'.format(
-            left=self.left.get_sql(**kwargs),
-            comparator=self.comparator.value,
-            right=self.right.get_sql(**kwargs),
-            nested_comparator=self.nested_comparator.value,
-            nested=self.nested.get_sql(**kwargs)
+              left=self.left.get_sql(**kwargs),
+              comparator=self.comparator.value,
+              right=self.right.get_sql(**kwargs),
+              nested_comparator=self.nested_comparator.value,
+              nested=self.nested.get_sql(**kwargs)
         )
         if with_alias and self.alias:
             return '{sql} "{alias}"'.format(sql=sql, alias=self.alias)
@@ -597,9 +627,19 @@ class BasicCriterion(Criterion):
         return self.left.tables_ | self.right.tables_
 
     @builder
-    def for_(self, table):
-        self.left = self.left.for_(table)
-        self.right = self.right.for_(table)
+    def replace_table(self, current_table, new_table):
+        """
+        Replaces all occurrences of the specified table with the new table. Useful when reusing fields across queries.
+
+        :param current_table:
+            The table to be replaced.
+        :param new_table:
+            The table to replace with.
+        :return:
+            A copy of the criterion with the tables replaced. 
+        """
+        self.left = self.left.replace_table(current_table, new_table)
+        self.right = self.right.replace_table(current_table, new_table)
 
     def fields(self):
         return self.left.fields() + self.right.fields()
@@ -641,6 +681,20 @@ class ContainsCriterion(Criterion):
     def is_aggregate(self):
         return self.term.is_aggregate
 
+    @builder
+    def replace_table(self, current_table, new_table):
+        """
+        Replaces all occurrences of the specified table with the new table. Useful when reusing fields across queries.
+
+        :param current_table:
+            The table to be replaced.
+        :param new_table:
+            The table to replace with.
+        :return:
+            A copy of the criterion with the tables replaced. 
+        """
+        self.term = self.term.replace_table(current_table, new_table)
+
     def fields(self):
         return self.term.fields() if self.term.fields else []
 
@@ -673,8 +727,18 @@ class BetweenCriterion(Criterion):
         return self.term.is_aggregate
 
     @builder
-    def for_(self, table):
-        self.term = self.term.for_(table)
+    def replace_table(self, current_table, new_table):
+        """
+        Replaces all occurrences of the specified table with the new table. Useful when reusing fields across queries.
+
+        :param current_table:
+            The table to be replaced.
+        :param new_table:
+            The table to replace with.
+        :return:
+            A copy of the criterion with the tables replaced. 
+        """
+        self.term = self.term.replace_table(current_table, new_table)
 
     def get_sql(self, **kwargs):
         # FIXME escape
@@ -698,8 +762,18 @@ class NullCriterion(Criterion):
         return self.term.tables_
 
     @builder
-    def for_(self, table):
-        self.term = self.term.for_(table)
+    def replace_table(self, current_table, new_table):
+        """
+        Replaces all occurrences of the specified table with the new table. Useful when reusing fields across queries.
+
+        :param current_table:
+            The table to be replaced.
+        :param new_table:
+            The table to replace with.
+        :return:
+            A copy of the criterion with the tables replaced. 
+        """
+        self.term = self.term.replace_table(current_table, new_table)
 
     def get_sql(self, **kwargs):
         return "{term} IS NULL".format(
@@ -772,17 +846,19 @@ class ArithmeticExpression(Term):
         return self.left.tables_ | self.right.tables_
 
     @builder
-    def for_(self, table):
+    def replace_table(self, current_table, new_table):
         """
-        Replaces the tables of this term for the table parameter provided.  Useful when reusing terms across queries.
+        Replaces all occurrences of the specified table with the new table. Useful when reusing fields across queries.
 
-        :param table:
+        :param current_table:
+            The table to be replaced.
+        :param new_table:
             The table to replace with.
         :return:
-            A copy of the term with it's table value replaced.
+            A copy of the term with the tables replaced. 
         """
-        self.left = self.left.for_(table)
-        self.right = self.right.for_(table)
+        self.left = self.left.replace_table(current_table, new_table)
+        self.right = self.right.replace_table(current_table, new_table)
 
     def fields(self):
         return self.left.fields() + self.right.fields()
@@ -821,6 +897,23 @@ class Case(Term):
     def when(self, criterion, term):
         self._cases.append((criterion, self.wrap_constant(term)))
 
+    @builder
+    def replace_table(self, current_table, new_table):
+        """
+        Replaces all occurrences of the specified table with the new table. Useful when reusing fields across queries.
+
+        :param current_table:
+            The table to be replaced.
+        :param new_table:
+            The table to replace with.
+        :return:
+            A copy of the term with the tables replaced. 
+        """
+        self._cases = [[criterion.replace_table(current_table, new_table), term.replace_table(current_table, new_table)]
+                       for criterion, term
+                       in self._cases]
+        self._else = self._else.replace_table(current_table, new_table) if self._else else None
+
     @builder
     def else_(self, term):
         self._else = self.wrap_constant(term)
@@ -906,6 +999,20 @@ class Not(Criterion):
 
         return inner
 
+    @builder
+    def replace_table(self, current_table, new_table):
+        """
+        Replaces all occurrences of the specified table with the new table. Useful when reusing fields across queries.
+
+        :param current_table:
+            The table to be replaced.
+        :param new_table:
+            The table to replace with.
+        :return:
+            A copy of the criterion with the tables replaced. 
+        """
+        self.term = self.term.replace_table(current_table, new_table)
+
     @property
     def tables_(self):
         return self.term.tables_
@@ -942,16 +1049,18 @@ class Function(Criterion):
         return len(self.args) == 1 and self.args[0].is_aggregate
 
     @builder
-    def for_(self, table):
+    def replace_table(self, current_table, new_table):
         """
-        Replaces the tables of this term for the table parameter provided.  Useful when reusing fields across queries.
+        Replaces all occurrences of the specified table with the new table. Useful when reusing fields across queries.
 
-        :param table:
+        :param current_table:
+            The table to be replaced.
+        :param new_table:
             The table to replace with.
         :return:
-            A copy of the field with it's table value replaced.
+            A copy of the criterion with the tables replaced. 
         """
-        self.args = [param.for_(table) for param in self.args]
+        self.args = [param.replace_table(current_table, new_table) for param in self.args]
 
     def get_special_params_sql(self, **kwargs):
         pass


Test patch:
diff --git a/pypika/tests/test_criterions.py b/pypika/tests/test_criterions.py
index b0cddc5..e390725 100644
--- a/pypika/tests/test_criterions.py
+++ b/pypika/tests/test_criterions.py
@@ -247,7 +247,7 @@ class CriterionTests(unittest.TestCase):
 
 
 class NotTests(unittest.TestCase):
-    table_abc = Table('abc', alias='cx0')
+    table_abc, table_efg = Table('abc', alias='cx0'), Table('efg', alias='cx1')
 
     def test_negate(self):
         c1 = Field('foo') >= 1
@@ -283,10 +283,10 @@ class NotTests(unittest.TestCase):
         self.assertEqual('NOT "foo" IS NULL "something"', str(c1))
         self.assertEqual('NOT "cx0"."foo" IS NULL "something"', str(c2))
 
-    def test_notnullcriterion_for_table(self):
-        f = self.table_abc.foo.notnull().for_(self.table_abc)
+    def test_notnullcriterion_replace_table(self):
+        f = self.table_abc.foo.notnull().replace_table(self.table_abc, self.table_efg)
 
-        self.assertEqual('NOT "cx0"."foo" IS NULL', str(f))
+        self.assertEqual('NOT "cx1"."foo" IS NULL', str(f))
 
     def test_not_with_or_criterion(self):
         self.assertEqual('NOT ("foo" OR "bar")', str(~(Field('foo') | Field('bar'))))
@@ -610,44 +610,45 @@ class FieldsAsCriterionTests(unittest.TestCase):
 class CriterionOperationsTests(unittest.TestCase):
     table_abc, table_efg = Table('abc', alias='cx0'), Table('efg', alias='cx1')
 
-    def test_field_for_table(self):
-        f = self.table_abc.foo.for_(self.table_efg)
+    def test_field_replace_table(self):
+        f = self.table_abc.foo.replace_table(self.table_abc, self.table_efg)
 
         self.assertEqual('"cx1"."foo"', str(f))
 
-    def test_arithmeticfunction_for_table(self):
-        f = (self.table_abc.foo + self.table_abc.bar).for_(self.table_efg)
+    def test_arithmeticfunction_replace_table(self):
+        f = (self.table_abc.foo + self.table_abc.bar).replace_table(self.table_abc, self.table_efg)
 
         self.assertEqual('"cx1"."foo"+"cx1"."bar"', str(f))
 
-    def test_criterion_for_table(self):
-        f = (self.table_abc.foo < self.table_abc.bar).for_(self.table_efg)
+    def test_criterion_replace_table(self):
+        f = (self.table_abc.foo < self.table_abc.bar).replace_table(self.table_abc, self.table_efg)
 
         self.assertEqual('"cx1"."foo"<"cx1"."bar"', str(f))
 
-    def test_complexcriterion_for_table(self):
-        f = ((self.table_abc.foo < self.table_abc.bar) & (self.table_abc.fiz > self.table_abc.buz)).for_(self.table_efg)
+    def test_complexcriterion_replace_table(self):
+        f = ((self.table_abc.foo < self.table_abc.bar) & (self.table_abc.fiz > self.table_abc.buz))
+        f = f.replace_table(self.table_abc, self.table_efg)
 
         self.assertEqual('"cx1"."foo"<"cx1"."bar" AND "cx1"."fiz">"cx1"."buz"', str(f))
 
-    def test_function_with_only_fields_for_table(self):
+    def test_function_with_only_fields_replace_table(self):
         f = fn.Sum(self.
-                   table_abc.foo).for_(self.table_efg)
+                   table_abc.foo).replace_table(self.table_abc, self.table_efg)
 
         self.assertEqual('SUM("cx1"."foo")', str(f))
 
-    def test_function_with_values_and_fields_for_table(self):
-        f = Mod(self.table_abc.foo, 2).for_(self.table_efg)
+    def test_function_with_values_and_fields_replace_table(self):
+        f = Mod(self.table_abc.foo, 2).replace_table(self.table_abc, self.table_efg)
 
         self.assertEqual('MOD("cx1"."foo",2)', str(f))
 
-    def test_betweencriterion_for_table(self):
-        f = self.table_abc.foo[0:1].for_(self.table_efg)
+    def test_betweencriterion_replace_table(self):
+        f = self.table_abc.foo[0:1].replace_table(self.table_abc, self.table_efg)
 
         self.assertEqual('"cx1"."foo" BETWEEN 0 AND 1', str(f))
 
-    def test_nullcriterion_for_table(self):
-        f = self.table_abc.foo.isnull().for_(self.table_efg)
+    def test_nullcriterion_replace_table(self):
+        f = self.table_abc.foo.isnull().replace_table(self.table_abc, self.table_efg)
 
         self.assertEqual('"cx1"."foo" IS NULL', str(f))
 
diff --git a/pypika/tests/test_internals.py b/pypika/tests/test_internals.py
index 2704b6b..b670925 100644
--- a/pypika/tests/test_internals.py
+++ b/pypika/tests/test_internals.py
@@ -10,10 +10,10 @@ from pypika.terms import Star
 
 class TablesTests(unittest.TestCase):
 
-    def test__criterion_for_with_value(self):
+    def test__criterion_replace_table_with_value(self):
         table = Table('a')
 
-        c = (Field('foo') > 1).for_(table)
+        c = (Field('foo') > 1).replace_table(None, table)
         self.assertEqual(c.left, table)
         self.assertEqual(c.tables_, {table})
 
diff --git a/pypika/tests/test_query.py b/pypika/tests/test_query.py
new file mode 100644
index 0000000..3397306
--- /dev/null
+++ b/pypika/tests/test_query.py
@@ -0,0 +1,97 @@
+import unittest
+
+from pypika import (
+    Case,
+    Query,
+    Tables,
+    functions,
+)
+
+
+class QueryTablesTests(unittest.TestCase):
+    table_a, table_b, table_c, table_d = Tables('a', 'b', 'c', 'd')
+
+    def test_replace_table(self):
+        query = Query.from_(self.table_a).select(self.table_a.time)
+        query = query.replace_table(self.table_a, self.table_b)
+
+        self.assertEqual('SELECT "time" FROM "b"', str(query))
+
+    def test_replace_only_specified_table(self):
+        query = Query.from_(self.table_a).select(self.table_a.time)
+        query = query.replace_table(self.table_b, self.table_c)
+
+        self.assertEqual('SELECT "time" FROM "a"', str(query))
+
+    def test_replace_insert_table(self):
+        query = Query.into(self.table_a).insert(1)
+        query = query.replace_table(self.table_a, self.table_b)
+
+        self.assertEqual('INSERT INTO "b" VALUES (1)', str(query))
+
+    def test_replace_update_table(self):
+        query = Query.update(self.table_a).set('foo', 'bar')
+        query = query.replace_table(self.table_a, self.table_b)
+
+        self.assertEqual('UPDATE "b" SET "foo"=\'bar\'', str(query))
+
+    def test_replace_delete_table(self):
+        query = Query.from_(self.table_a).delete()
+        query = query.replace_table(self.table_a, self.table_b)
+
+        self.assertEqual('DELETE FROM "b"', str(query))
+
+    def test_replace_join_tables(self):
+        query = Query \
+            .from_(self.table_a) \
+            .join(self.table_b) \
+            .on(self.table_a.customer_id == self.table_b.id) \
+            .join(self.table_c) \
+            .on(self.table_b.seller_id == self.table_c.id) \
+            .select(self.table_a.star)
+        query = query.replace_table(self.table_a, self.table_d)
+
+        self.assertEqual('SELECT "d".* '
+                         'FROM "d" '
+                         'JOIN "b" ON "d"."customer_id"="b"."id" '
+                         'JOIN "c" ON "b"."seller_id"="c"."id"', str(query))
+
+    def test_replace_filter_tables(self):
+        query = Query \
+            .from_(self.table_a) \
+            .select(self.table_a.name) \
+            .where(self.table_a.name == 'Mustermann')
+        query = query.replace_table(self.table_a, self.table_b)
+
+        self.assertEqual('SELECT "name" '
+                         'FROM "b" '
+                         'WHERE "name"=\'Mustermann\'', str(query))
+
+    def test_replace_having_table(self):
+        query = Query \
+            .from_(self.table_a) \
+            .select(functions.Sum(self.table_a.revenue)) \
+            .groupby(self.table_a.customer) \
+            .having(functions.Sum(self.table_a.revenue) >= 1000)
+        query = query.replace_table(self.table_a, self.table_b)
+
+        self.assertEqual('SELECT SUM("revenue") '
+                         'FROM "b" '
+                         'GROUP BY "customer" '
+                         'HAVING SUM("revenue")>=1000', str(query))
+
+    def test_replace_case_table(self):
+        query = Query \
+            .from_(self.table_a) \
+            .select(Case()
+                    .when(self.table_a.fname == "Tom", "It was Tom")
+                    .when(self.table_a.fname == "John", "It was John")
+                    .else_("It was someone else.")
+                    .as_('who_was_it'))
+        query = query.replace_table(self.table_a, self.table_b)
+
+        self.assertEqual('SELECT CASE '
+                         'WHEN "fname"=\'Tom\' THEN \'It was Tom\' '
+                         'WHEN "fname"=\'John\' THEN \'It was John\' '
+                         'ELSE \'It was someone else.\' END "who_was_it" '
+                         'FROM "b"', str(query))