Add Query node_type (#484)

* Add Query node_type * bump --------- Co-authored-by: Ian Eaves <[email protected]>
grai-io · Jul 6, 2023 · 8af2af6 · 8af2af6
1 parent 800ff2e
commit 8af2af6
Show file tree

Hide file tree

Showing 6 changed files with 208 additions and 398 deletions.
diff --git a/grai-schemas/pyproject.toml b/grai-schemas/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "grai_schemas"
-version = "0.1.15"
+version = "0.1.16"
 description = ""
 authors = ["Ian Eaves <[email protected]>", "Edward Louth <[email protected]>"]
 license = "Elastic-2.0"

diff --git a/grai-schemas/src/grai_schemas/__init__.py b/grai-schemas/src/grai_schemas/__init__.py
@@ -10,4 +10,4 @@
 from grai_schemas.base import GraiBaseModel
 from grai_schemas.package_definitions import config
 
-__version__ = "0.1.15"
+__version__ = "0.1.16"
diff --git a/grai-schemas/src/grai_schemas/v1/metadata/nodes.py b/grai-schemas/src/grai_schemas/v1/metadata/nodes.py
@@ -12,9 +12,10 @@ class NodeMetadataTypeLabels(Enum):
     generic: Literal["Generic"] = "Generic"
     table: Literal["Table"] = "Table"
     column: Literal["Column"] = "Column"
+    query: Literal["Query"] = "Query"
 
 
-NodeMetadataTypeLabelLiterals = Literal["Generic", "Table", "Column"]
+NodeMetadataTypeLabelLiterals = Literal["Generic", "Table", "Column", "Query"]
 
 
 class SourceType(Enum):
@@ -74,4 +75,17 @@ class TableMetadata(BaseNodeMetadataV1):
     node_attributes: TableAttributes = TableAttributes()
 
 
-Metadata = Union[ColumnMetadata, TableMetadata, GenericNodeMetadataV1]
+class QueryAttributes(V1Mixin, GenericAttributes):
+    """ """
+
+    pass
+
+
+class QueryMetadata(BaseNodeMetadataV1):
+    """ """
+
+    node_type: Literal["Query"]
+    node_attributes: QueryAttributes = QueryAttributes()
+
+
+Metadata = Union[ColumnMetadata, TableMetadata, QueryMetadata, GenericNodeMetadataV1]
diff --git a/grai-server/app/lineage/graph_cache.py b/grai-server/app/lineage/graph_cache.py
@@ -42,7 +42,7 @@ def clear_cache(self):
     def cache_node(self, node):
         node_type = node.metadata.get("grai", {}).get("node_type")
 
-        if node_type == "Table":
+        if node_type in ["Table", "Query"]:
             self.query(
                 """
                     MERGE (table:Table {id: $id})