Repository: kayak/pypika
Pull request #145 (2018-08-08 10:56:06)

Fixed join queries to work with update and insert queries
Fixed the join function in Query to work with the update table as well as from_ tables


Patch:
diff --git a/pypika/queries.py b/pypika/queries.py
index 53970cd..5cf9ef1 100644
--- a/pypika/queries.py
+++ b/pypika/queries.py
@@ -592,15 +592,17 @@ class QueryBuilder(Selectable, Term):
         return []
 
     def do_join(self, join):
-        join.validate(self._from, self._joins)
+        base_tables = self._from + [self._update_table]
+
+        join.validate(base_tables, self._joins)
 
         if isinstance(join.item, QueryBuilder) and join.item.alias is None:
             self._tag_subquery(join.item)
 
-        table_in_from = any(isinstance(clause, Table)
-                            and join.item in self._from
-                            for clause in self._from)
-        if isinstance(join.item, Table) and join.item.alias is None and table_in_from:
+        table_in_query = any(isinstance(clause, Table)
+                             and join.item in base_tables
+                             for clause in base_tables)
+        if isinstance(join.item, Table) and join.item.alias is None and table_in_query:
             # On the odd chance that we join the same table as the FROM table and don't set an alias
             # FIXME only works once
             join.item.alias = join.item._table_name + '2'
@@ -608,11 +610,13 @@ class QueryBuilder(Selectable, Term):
         self._joins.append(join)
 
     def _validate_term(self, term):
+        base_tables = self._from + [self._update_table]
+
         for field in term.fields():
-            table_in_froms = field.table in self._from
+            table_in_base_tables = field.table in base_tables
             table_in_joins = field.table in [join.item for join in self._joins]
             if field.table is not None \
-                  and not table_in_froms \
+                  and not table_in_base_tables \
                   and not table_in_joins \
                   and field.table != self._update_table:
                 raise JoinException('Table [%s] missing from query.  '


Test patch:
diff --git a/pypika/tests/test_joins.py b/pypika/tests/test_joins.py
index 60bdb53..92235ce 100644
--- a/pypika/tests/test_joins.py
+++ b/pypika/tests/test_joins.py
@@ -17,7 +17,7 @@ __author__ = "Timothy Heys"
 __email__ = "theys@kayak.com"
 
 
-class JoinTypeTests(unittest.TestCase):
+class SelectQueryJoinTests(unittest.TestCase):
     table0, table1, hij = Tables('abc', 'efg', 'hij')
 
     def test_left_join(self):
@@ -482,3 +482,29 @@ class UnionTests(unittest.TestCase):
 
         self.assertEqual('SELECT AVG("sq0"."t") FROM ((SELECT "t" FROM "abc") UNION (SELECT "t" FROM "efg")) "sq0"',
                          str(q))
+
+
+class InsertQueryJoinTests(unittest.TestCase):
+    def test_join_table_on_insert_query(self):
+        a, b, c = Tables('a', 'b', 'c')
+        q = Query.into(c).from_(a) \
+            .join(b).on(a.fkey_id == b.id) \
+            .where(b.foo == 1) \
+            .select('a.*', 'b.*')
+
+        self.assertEqual('INSERT INTO "c" '
+                         'SELECT "a"."a.*","a"."b.*" '
+                         'FROM "a" JOIN "b" '
+                         'ON "a"."fkey_id"="b"."id" '
+                         'WHERE "b"."foo"=1', str(q))
+
+
+class UpdateQueryJoinTests(unittest.TestCase):
+    def test_join_table_on_update_query(self):
+        a, b = Tables('a', 'b')
+        q = Query.update(a) \
+            .join(b).on(a.fkey_id == b.id) \
+            .where(b.foo == 1) \
+            .set('adwords_batch_job_id', 1)
+
+        self.assertEqual('UPDATE "a" SET "adwords_batch_job_id"=1 WHERE "b"."foo"=1', str(q))