Include columns used in on expression in response (#1014)

Author: vchan

web/server/api/endpoints/table_diff.py

@@ -53,4 +53,11 @@ def get_table_diff(
         target_count=_row_diff.target_count,
         count_pct_change=_row_diff.count_pct_change,
     )
-    return TableDiff(schema_diff=schema_diff, row_diff=row_diff)
+    return TableDiff(
+        schema_diff=schema_diff,
+        row_diff=row_diff,
+        on=[
+            (eq.left, eq.right) if eq.left.table == "s" else (eq.right, eq.left)
+            for eq in diff.on.find_all(exp.EQ)
+        ],
+    )

web/server/models.py

@@ -298,3 +298,11 @@ class RowDiff(BaseModel):
 class TableDiff(BaseModel):
     schema_diff: SchemaDiff
     row_diff: RowDiff
+    on: t.List[t.Tuple[str, str]]
+
+    @validator("on", pre=True)
+    def validate_schema(
+        cls,
+        v: list[tuple[str, str]] | list[tuple[exp.Column, exp.Column]],
+    ) -> list[tuple[str, str]]:
+        return [(str(s), str(t)) for s, t in v]