allow joindiff across bigquery projects

Author: Marco Baringer

data_diff/info_tree.py

@@ -13,7 +13,7 @@ class SegmentInfo:
     is_diff: bool = None
     diff_count: int = None
 
-    rowcounts: Dict[int, int] = {}
+    rowcounts: Dict[int, int] = {1: 0, 2: 0}
     max_rows: int = None
 
     def set_diff(self, diff):

data_diff/joindiff_tables.py

@@ -144,7 +144,7 @@ class JoinDiffer(TableDiffer):
     def _diff_tables_root(self, table1: TableSegment, table2: TableSegment, info_tree: InfoTree) -> DiffResult:
         db = table1.database
 
-        if table1.database is not table2.database:
+        if not table1.database.can_join_with(table2.database):
             raise ValueError("Join-diff only works when both tables are in the same database")
 
         table1, table2 = self._threaded_call("with_schema", [table1, table2])
@@ -174,7 +174,7 @@ def _diff_segments(
         segment_index=None,
         segment_count=None,
     ):
-        assert table1.database is table2.database
+        assert table1.database.can_join_with(table2.database)
 
         if segment_index or table1.min_key or max_rows:
             logger.info(
@@ -295,8 +295,7 @@ def _collect_stats(self, i, table_seg: TableSegment, info_tree: InfoTree):
         logger.debug("Done collecting stats for table #%s", i)
 
     def _create_outer_join(self, table1, table2):
-        db = table1.database
-        if db is not table2.database:
+        if not table1.database.can_join_with(table2.database):
             raise ValueError("Joindiff only applies to tables within the same database")
 
         keys1 = table1.key_columns
@@ -319,7 +318,7 @@ def _create_outer_join(self, table1, table2):
         # Order columns as col1_a, col1_b, col2_a, col2_b, etc.
         cols = {k: v for k, v in chain(*zip(a_cols.items(), b_cols.items()))}
 
-        all_rows = _outerjoin(db, a, b, keys1, keys2, {**is_diff_cols, **cols})
+        all_rows = _outerjoin(table1.database, a, b, keys1, keys2, {**is_diff_cols, **cols})
         diff_rows = all_rows.where(or_(this[c] == 1 for c in is_diff_cols))
         return diff_rows, a_cols, b_cols, is_diff_cols, all_rows
 

data_diff/sqeleton/databases/base.py

@@ -453,6 +453,8 @@ def close(self):
         self.is_closed = True
         return super().close()
 
+    def can_join_with(self, other):
+        return self is other
 
 class ThreadedDatabase(Database):
     """Access the database through singleton threads.

data_diff/sqeleton/databases/bigquery.py

@@ -145,21 +145,38 @@ def close(self):
         self._client.close()
 
     def select_table_schema(self, path: DbPath) -> str:
-        schema, name = self._normalize_table_path(path)
-
+        project, schema, name = self._normalize_table_path(path)
         return (
             "SELECT column_name, data_type, 6 as datetime_precision, 38 as numeric_precision, 9 as numeric_scale "
-            f"FROM {schema}.INFORMATION_SCHEMA.COLUMNS "
+            f"FROM `{project}`.`{schema}`.INFORMATION_SCHEMA.COLUMNS "
             f"WHERE table_name = '{name}' AND table_schema = '{schema}'"
         )
 
     def query_table_unique_columns(self, path: DbPath) -> List[str]:
         return []
 
+    def _normalize_table_path(self, path: DbPath) -> DbPath:
+        if len(path) == 0:
+            raise ValueError(f"{self.name}: Bad table path for {self}: ()")
+        elif len(path) == 1:
+            if self.default_schema:
+                return [self.project, self.default_schema, path[0]]
+            else:
+                return path
+        elif len(path) == 2:
+            return [self.project] + path
+        elif len(path) == 3:
+            return path
+        else:
+            raise ValueError(f"{self.name}: Bad table path for {self}: '{'.'.join(path)}'. Expected form: [project.]schema.table")
+
     def parse_table_name(self, name: str) -> DbPath:
         path = parse_table_name(name)
         return tuple(i for i in self._normalize_table_path(path) if i is not None)
 
     @property
     def is_autocommit(self) -> bool:
         return True
+
+    def can_join_with(self, other):
+        return True