feat(ibis): introduce the relationship validation (#815)

* implement relationship validation

* fix tests for other source

* fix the method argument

* fix test

* fix test

* address comment

Author: goldmedal

ibis-server/app/model/validator.py

@@ -1,28 +1,31 @@
 from __future__ import annotations
 
+import base64
+import json
+
 from app.mdl.rewriter import Rewriter
 from app.model import UnprocessableEntityError
 from app.model.connector import Connector
 
-rules = ["column_is_valid"]
+rules = ["column_is_valid", "relationship_is_valid"]
 
 
 class Validator:
     def __init__(self, connector: Connector, rewriter: Rewriter):
         self.connector = connector
         self.rewriter = rewriter
 
-    def validate(self, rule: str, parameters: dict[str, str]):
+    def validate(self, rule: str, parameters: dict[str, str], manifest_str: str):
         if rule not in rules:
             raise RuleNotFoundError(rule)
         try:
-            getattr(self, f"_validate_{rule}")(parameters)
+            getattr(self, f"_validate_{rule}")(parameters, manifest_str)
         except ValidationError as e:
             raise e
         except Exception as e:
             raise ValidationError(f"Unknown exception: {type(e)}, message: {e!s}")
 
-    def _validate_column_is_valid(self, parameters: dict[str, str]):
+    def _validate_column_is_valid(self, parameters: dict[str, str], manifest_str: str):
         model_name = parameters.get("modelName")
         column_name = parameters.get("columnName")
         if model_name is None:
@@ -37,6 +40,116 @@ def _validate_column_is_valid(self, parameters: dict[str, str]):
         except Exception as e:
             raise ValidationError(f"Exception: {type(e)}, message: {e!s}")
 
+    def _validate_relationship_is_valid(
+        self, parameters: dict[str, str], manifest_str: str
+    ):
+        relationship_name = parameters.get("relationshipName")
+        if relationship_name is None:
+            raise MissingRequiredParameterError("relationship")
+        decoded_manifest = base64.b64decode(manifest_str).decode("utf-8")
+        manifest = json.loads(decoded_manifest)
+
+        relationship = list(
+            filter(lambda r: r["name"] == relationship_name, manifest["relationships"])
+        )
+
+        if len(relationship) == 0:
+            raise ValidationError(
+                f"Relationship {relationship_name} not found in manifest"
+            )
+
+        left_model = self._get_model(manifest, relationship[0]["models"][0])
+        right_model = self._get_model(manifest, relationship[0]["models"][1])
+        relationship_type = relationship[0]["joinType"].lower()
+        condition = relationship[0]["condition"]
+        columns = condition.split("=")
+        left_column = columns[0].strip().split(".")[1]
+        right_column = columns[1].strip().split(".")[1]
+
+        def generate_column_is_unique_sql(model_name, column_name):
+            return f'SELECT count(*) = count(distinct {column_name}) AS result FROM "{model_name}"'
+
+        def generate_is_exist_join_sql(
+            left_model, right_model, left_column, right_column
+        ):
+            return f'SELECT count(*) > 0 AS result FROM "{left_model}" JOIN "{right_model}" ON "{left_model}"."{left_column}" = "{right_model}"."{right_column}"'
+
+        def generate_sql_from_type(
+            relationship_type, left_model, right_model, left_column, right_column
+        ):
+            if relationship_type == "one_to_one":
+                return f"""WITH
+                    lefttable AS ({generate_column_is_unique_sql(left_model, left_column)}),
+                    righttable AS ({generate_column_is_unique_sql(right_model, right_column)}),
+                    joinexist AS ({generate_is_exist_join_sql(left_model, right_model, left_column, right_column)})
+                SELECT lefttable.result AND righttable.result AND joinexist.result result,
+                    lefttable.result left_table_unique,
+                    righttable.result right_table_unique,
+                    joinexist.result is_related
+                FROM lefttable, righttable, joinexist"""
+            elif relationship_type == "many_to_one":
+                return f"""WITH 
+                    righttable AS ({generate_column_is_unique_sql(right_model, right_column)}),
+                    joinexist AS ({generate_is_exist_join_sql(left_model, right_model, left_column, right_column)})
+                SELECT righttable.result AND joinexist.result result,
+                    righttable.result right_table_unique,
+                    joinexist.result is_related
+                FROM righttable, joinexist"""
+            elif relationship_type == "one_to_many":
+                return f"""WITH 
+                    lefttable AS ({generate_column_is_unique_sql(left_model, left_column)}),
+                    joinexist AS ({generate_is_exist_join_sql(left_model, right_model, left_column, right_column)})
+                SELECT lefttable.result AND joinexist.result result,
+                    lefttable.result left_table_unique,
+                    joinexist.result is_related
+                FROM lefttable, joinexist"""
+            elif relationship_type == "many_to_many":
+                return f"""WITH 
+                    joinexist AS ({generate_is_exist_join_sql(left_model, right_model, left_column, right_column)})
+                SELECT joinexist.result result,
+                    joinexist.result is_related
+                FROM joinexist"""
+            else:
+                raise ValidationError(f"Unknown relationship type: {relationship_type}")
+
+        def format_result(result):
+            output = {}
+            output["result"] = str(result.get("result").get(0))
+            output["is_related"] = str(result.get("is_related").get(0))
+            if result.get("left_table_unique") is not None:
+                output["left_table_unique"] = str(
+                    result.get("left_table_unique").get(0)
+                )
+            if result.get("right_table_unique") is not None:
+                output["right_table_unique"] = str(
+                    result.get("right_table_unique").get(0)
+                )
+            return output
+
+        sql = generate_sql_from_type(
+            relationship_type,
+            left_model["name"],
+            right_model["name"],
+            left_column,
+            right_column,
+        )
+        try:
+            rewritten_sql = self.rewriter.rewrite(sql)
+            result = self.connector.query(rewritten_sql, limit=1)
+            if not result.get("result").get(0):
+                raise ValidationError(
+                    f"Relationship {relationship_name} is not valid: {format_result(result)}"
+                )
+
+        except Exception as e:
+            raise ValidationError(f"Exception: {type(e)}, message: {e!s}")
+
+    def _get_model(self, manifest, model_name):
+        models = list(filter(lambda m: m["name"] == model_name, manifest["models"]))
+        if len(models) == 0:
+            raise ValidationError(f"Model {model_name} not found in manifest")
+        return models[0]
+
 
 class ValidationError(UnprocessableEntityError):
     pass

ibis-server/app/routers/v2/connector.py

@@ -91,7 +91,7 @@ def validate(
             non_cached_mdl=True,
         ),
     )
-    validator.validate(rule_name, dto.parameters)
+    validator.validate(rule_name, dto.parameters, dto.manifest_str)
     return Response(status_code=204)
 
 

ibis-server/app/routers/v3/connector.py

@@ -61,5 +61,5 @@ def validate(data_source: DataSource, rule_name: str, dto: ValidateDTO) -> Respo
         Connector(data_source, dto.connection_info, dto.manifest_str),
         Rewriter(dto.manifest_str, data_source=data_source, experiment=True),
     )
-    validator.validate(rule_name, dto.parameters)
+    validator.validate(rule_name, dto.parameters, dto.manifest_str)
     return Response(status_code=204)

ibis-server/tests/routers/v2/connector/test_bigquery.py

@@ -7,6 +7,7 @@
 
 from app.main import app
 from app.model import HTTPHeaders
+from app.model.validator import rules
 
 pytestmark = pytest.mark.bigquery
 
@@ -247,19 +248,32 @@ def test_query_cast_as_timestamp(manifest_hash):
             "2024-01-01 00:00:00.000000 CST",
         ]
 
-    def test_validate_with_unknown_rule():
+    def test_query_values(manifest_hash):
+        response = client.post(
+            url=f"{base_url}/query",
+            params={"dryRun": True},
+            json={
+                "connectionInfo": connection_info,
+                "manifestStr": manifest_hash,
+                "sql": "SELECT * FROM (VALUES (1, 2), (3, 4))",
+            },
+        )
+
+        assert response.status_code == 204
+
+    def test_validate_with_unknown_rule(manifest_hash):
         response = client.post(
             url=f"{base_url}/validate/unknown_rule",
             json={
                 "connectionInfo": connection_info,
-                "manifestStr": manifest_str,
+                "manifestStr": manifest_hash,
                 "parameters": {"modelName": "Orders", "columnName": "orderkey"},
             },
         )
         assert response.status_code == 422
         assert (
             response.text
-            == "The rule `unknown_rule` is not in the rules, rules: ['column_is_valid']"
+            == f"The rule `unknown_rule` is not in the rules, rules: {rules}"
         )
 
     def test_validate_rule_column_is_valid():

ibis-server/tests/routers/v2/connector/test_clickhouse.py

@@ -9,6 +9,7 @@
 
 from app.main import app
 from app.model import HTTPHeaders
+from app.model.validator import rules
 from tests.confest import file_path
 
 pytestmark = pytest.mark.clickhouse
@@ -473,7 +474,7 @@ def test_validate_with_unknown_rule(clickhouse: ClickHouseContainer):
         assert response.status_code == 422
         assert (
             response.text
-            == "The rule `unknown_rule` is not in the rules, rules: ['column_is_valid']"
+            == f"The rule `unknown_rule` is not in the rules, rules: {rules}"
         )
 
     def test_validate_rule_column_is_valid(clickhouse: ClickHouseContainer):

ibis-server/tests/routers/v2/connector/test_mssql.py

@@ -10,6 +10,7 @@
 
 from app.main import app
 from app.model import HTTPHeaders
+from app.model.validator import rules
 from tests.confest import file_path
 
 pytestmark = pytest.mark.mssql
@@ -324,7 +325,7 @@ def test_validate_with_unknown_rule(mssql: SqlServerContainer):
         assert response.status_code == 422
         assert (
             response.text
-            == "The rule `unknown_rule` is not in the rules, rules: ['column_is_valid']"
+            == f"The rule `unknown_rule` is not in the rules, rules: {rules}"
         )
 
     def test_validate_rule_column_is_valid(mssql: SqlServerContainer):

ibis-server/tests/routers/v2/connector/test_mysql.py

@@ -10,6 +10,7 @@
 
 from app.main import app
 from app.model import HTTPHeaders
+from app.model.validator import rules
 from tests.confest import file_path
 
 pytestmark = pytest.mark.mysql
@@ -322,7 +323,7 @@ def test_validate_with_unknown_rule(mysql: MySqlContainer):
         assert response.status_code == 422
         assert (
             response.text
-            == "The rule `unknown_rule` is not in the rules, rules: ['column_is_valid']"
+            == f"The rule `unknown_rule` is not in the rules, rules: {rules}"
         )
 
     def test_validate_rule_column_is_valid(mysql: MySqlContainer):

ibis-server/tests/routers/v2/connector/test_postgres.py

@@ -13,6 +13,7 @@
 
 from app.main import app
 from app.model import VARIABLE_PREFIX, HTTPHeaders
+from app.model.validator import rules
 from tests.confest import file_path
 
 pytestmark = pytest.mark.postgres
@@ -426,7 +427,7 @@ def test_validate_with_unknown_rule(postgres: PostgresContainer):
         assert response.status_code == 422
         assert (
             response.text
-            == "The rule `unknown_rule` is not in the rules, rules: ['column_is_valid']"
+            == f"The rule `unknown_rule` is not in the rules, rules: {rules}"
         )
 
     def test_validate_rule_column_is_valid(postgres: PostgresContainer):

ibis-server/tests/routers/v2/connector/test_snowflake.py

@@ -7,6 +7,7 @@
 
 from app.main import app
 from app.model import HTTPHeaders
+from app.model.validator import rules
 
 pytestmark = pytest.mark.snowflake
 
@@ -255,7 +256,7 @@ def test_validate_with_unknown_rule():
         assert response.status_code == 422
         assert (
             response.text
-            == "The rule `unknown_rule` is not in the rules, rules: ['column_is_valid']"
+            == f"The rule `unknown_rule` is not in the rules, rules: {rules}"
         )
 
     def test_validate_rule_column_is_valid():

ibis-server/tests/routers/v2/connector/test_trino.py

@@ -9,6 +9,7 @@
 
 from app.main import app
 from app.model import HTTPHeaders
+from app.model.validator import rules
 
 pytestmark = pytest.mark.trino
 
@@ -332,7 +333,7 @@ def test_validate_with_unknown_rule(trino: TrinoContainer):
         assert response.status_code == 422
         assert (
             response.text
-            == "The rule `unknown_rule` is not in the rules, rules: ['column_is_valid']"
+            == f"The rule `unknown_rule` is not in the rules, rules: {rules}"
         )
 
     def test_validate_rule_column_is_valid(trino: TrinoContainer):