Add code to edge metadata

grai-io · Nov 14, 2023 · 5df7e17 · 5df7e17
1 parent 2b1a89f
commit 5df7e17
Show file tree

Hide file tree

Showing 5 changed files with 57 additions and 53 deletions.
diff --git a/grai-integrations/source-dbt/pyproject.toml b/grai-integrations/source-dbt/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "grai_source_dbt"
-version = "0.3.3"
+version = "0.3.4-alpha1"
 description = ""
 authors = ["Ian Eaves <ian@grai.io>", "Edward Louth <edward@grai.io>"]
 license = "Elastic-2.0"

diff --git a/grai-integrations/source-dbt/src/grai_source_dbt/__init__.py b/grai-integrations/source-dbt/src/grai_source_dbt/__init__.py
@@ -9,4 +9,4 @@
 )
 from grai_source_dbt.package_definitions import config
 
-__version__ = "0.3.3-alpha1-alpha1"
+__version__ = "0.3.4-alpha1"
diff --git a/grai-integrations/source-dbt/src/grai_source_dbt/adapters/adapters.py b/grai-integrations/source-dbt/src/grai_source_dbt/adapters/adapters.py
@@ -1,7 +1,7 @@
 from typing import Any, Dict, List, Literal, Sequence, TypeVar, Union
 
-from grai_schemas import config as base_config
 from grai_schemas.v1 import SourcedEdgeV1, SourcedNodeV1, SourceV1
+from grai_schemas.v1.generics import SQL, Code
 from grai_schemas.v1.metadata.edges import (
     ColumnToColumnMetadata,
     EdgeMetadataTypeLabels,
@@ -18,10 +18,9 @@
 from grai_schemas.v1.source import SourceSpec
 from multimethod import multimethod
 
-from grai_source_dbt.loaders import AllDbtNodeInstances, AllDbtNodeTypes
-from grai_source_dbt.models.grai import Column, Constraint, Edge
+from grai_source_dbt.loaders import AllDbtNodeTypes
+from grai_source_dbt.models.grai import Column, Edge
 from grai_source_dbt.package_definitions import config
-from grai_source_dbt.utils import full_name
 
 T = TypeVar("T")
 X = TypeVar("X")
@@ -116,6 +115,11 @@ def build_grai_metadata_from_edge(current: Edge, version: Literal["v1"] = "v1")
         "version": version,
         "edge_type": current.edge_type.value,
         "tags": [config.metadata_id],
+        "code": Code(
+            code=current.definition,
+            language={"language_name": "SQL", "file_extension": ".sql"},
+            has_code=current.definition is not None,
+        ),
     }
     if current.edge_type == EdgeMetadataTypeLabels.table_to_table:
         return TableToTableMetadata(**data)

diff --git a/grai-server/app/poetry.lock b/grai-server/app/poetry.lock
diff --git a/grai-server/app/pyproject.toml b/grai-server/app/pyproject.toml
@@ -48,7 +48,7 @@ grai-schemas = "^0.2.10"
 grai-client = "^0.3.2"
 grai-graph = "^0.2.5"
 grai-source-bigquery = "^0.2.4"
-grai-source-dbt = "^0.3.3"
+grai-source-dbt = "0.3.4a1"
 grai-source-dbt-cloud =  "^0.1.4"
 grai-source-fivetran = "^0.1.2"
 grai-source-flat-file = "^0.1.1"