Handle tag filter array (#835)

grai-io · Dec 13, 2023 · 6874a16 · 6874a16
1 parent 21360d7
commit 6874a16
Show file tree

Hide file tree

Showing 2 changed files with 58 additions and 8 deletions.
diff --git a/grai-server/app/lineage/graph_filter.py b/grai-server/app/lineage/graph_filter.py
@@ -30,22 +30,28 @@ def filter_by_dict(row, query: GraphQuery) -> GraphQuery:
             if row["operator"] == "equals":
                 query.where(f"table.namespace = '{value}'")
             elif row["operator"] == "in":
-                list = "['" + "', '".join(value) + "']"
-                query.where(f"table.namespace IN {list}")
+                value_list = "['" + "', '".join(value) + "']"
+                query.where(f"table.namespace IN {value_list}")
 
         elif row["field"] == "data-source":
             if row["operator"] == "in":
-                list = "['" + "', '".join(value) + "']"
-                query.where(f"any(x IN table.data_sources WHERE x IN {list})")
+                value_list = "['" + "', '".join(value) + "']"
+                query.where(f"any(x IN table.data_sources WHERE x IN {value_list})")
             elif row["operator"] == "not-in":
-                list = "['" + "', '".join(value) + "']"
-                query.where(f"NOT any(x IN table.data_sources WHERE x IN {list})")
+                value_list = "['" + "', '".join(value) + "']"
+                query.where(f"NOT any(x IN table.data_sources WHERE x IN {value_list})")
 
         elif row["field"] == "tag":
             if row["operator"] == "contains":
-                query.where(f"'{value}' IN table.tags")
+                if isinstance(value, list):
+                    query.where(f"'{value[0]}' IN table.tags")
+                else:
+                    query.where(f"'{value}' IN table.tags")
             elif row["operator"] == "not-contains":
-                query.where(f"NOT '{value}' IN table.tags")
+                if isinstance(value, list):
+                    query.where(f"NOT '{value[0]}' IN table.tags")
+                else:
+                    query.where(f"NOT '{value}' IN table.tags")
 
     elif row["type"] == "ancestor":
         if row["field"] == "tag":

diff --git a/grai-server/app/lineage/tests/test_graph_filter.py b/grai-server/app/lineage/tests/test_graph_filter.py
@@ -283,6 +283,28 @@ def test_table_tag_contains():
     assert query.parameters == {}
 
 
+def test_table_tag_contains_array():
+    query = GraphQuery(Match("Table"))
+
+    filter = Filter(
+        name=str(uuid.uuid4()),
+        metadata=[
+            {
+                "type": "table",
+                "field": "tag",
+                "operator": "contains",
+                "value": ["tag1"],
+            }
+        ],
+    )
+
+    filter_by_filter(filter, query)
+
+    assert len(query.clause[0].wheres) == 1
+    assert query.clause[0].wheres[0].where == "'tag1' IN table.tags"
+    assert query.parameters == {}
+
+
 def test_table_tag_doesnt_contain():
     query = GraphQuery(Match("Table"))
 
@@ -305,6 +327,28 @@ def test_table_tag_doesnt_contain():
     assert query.parameters == {}
 
 
+def test_table_tag_doesnt_contain_array():
+    query = GraphQuery(Match("Table"))
+
+    filter = Filter(
+        name=str(uuid.uuid4()),
+        metadata=[
+            {
+                "type": "table",
+                "field": "tag",
+                "operator": "not-contains",
+                "value": ["tag1"],
+            }
+        ],
+    )
+
+    filter_by_filter(filter, query)
+
+    assert len(query.clause[0].wheres) == 1
+    assert query.clause[0].wheres[0].where == "NOT 'tag1' IN table.tags"
+    assert query.parameters == {}
+
+
 def test_no_ancestor():
     query = GraphQuery(Match("Table"))