feature/gra-754-metabase-column-lineage (#589)

* Feat: Column level lineage for Metabase Questions.

grai-integrations/source-metabase/pyproject.toml

@@ -1,22 +1,23 @@
 [tool.poetry]
 name = "grai-source-metabase"
-version = "0.2.1"
+version = "0.2.2"
 
 description = ""
-authors = ["Elseagle <dowolebolu@gmail.com>"]
+authors = ["Elseagle <dowolebolu@gmail.com>", "Ian Eaves <ian@grai.io", "Edward Louth <edward@grai.io"]
 readme = "README.md"
 packages = [
     { include = "grai_source_metabase", from = "src" },
 ]
 
 [tool.poetry.dependencies]
-python = "^3.10"
+python = "^3.8"
 pydantic = "^1.10.9"
 requests = "^2.31.0"
-grai-schemas = "^0.2.0"
+grai-schemas = "^0.2.1"
 grai-client = "^0.3.0"
 multimethod = "^1.9.1"
 retrying = "^1.3.4"
+httpx = {version="^0.24.1", extras=["http2"]}
 
 [tool.poetry.group.dev.dependencies]
 python-dotenv = "^1.0.0"

grai-integrations/source-metabase/src/grai_source_metabase/__init__.py

@@ -1,11 +1,11 @@
 from grai_source_metabase import (
     adapters,
+    api,
     base,
     loader,
     mock_tools,
     models,
     package_definitions,
 )
 
-__version__ = "0.2.1"
-
+__version__ = "0.2.2"

grai-integrations/source-metabase/src/grai_source_metabase/adapters.py

@@ -1,10 +1,10 @@
 from typing import Any, Dict, List, Literal, Sequence, TypeVar, Union
 
-from grai_schemas import config as base_config
-from grai_schemas.schema import Schema
 from grai_schemas.v1 import EdgeV1, NodeV1, SourcedEdgeV1, SourcedNodeV1, SourceV1
 from grai_schemas.v1.metadata.edges import EdgeMetadataTypeLabels, GenericEdgeMetadataV1
 from grai_schemas.v1.metadata.nodes import (
+    CollectionMetadata,
+    ColumnMetadata,
     GenericNodeMetadataV1,
     NodeMetadataTypeLabels,
     QueryMetadata,
@@ -14,7 +14,14 @@
 from multimethod import multimethod
 from pydantic import BaseModel
 
-from grai_source_metabase.models import Collection, Edge, NodeTypes, Question, Table
+from grai_source_metabase.models import (
+    Collection,
+    Column,
+    Edge,
+    NodeTypes,
+    Question,
+    Table,
+)
 from grai_source_metabase.package_definitions import config
 
 T = TypeVar("T")
@@ -40,6 +47,33 @@ def build_grai_metadata(current: Any, desired: Any) -> None:
     raise NotImplementedError(f"No adapter between {type(current)} and {type(desired)} for value {current}")
 
 
+@build_grai_metadata.register
+def build_grai_metadata_from_table(current: Column, version: Literal["v1"] = "v1") -> ColumnMetadata:
+    """
+    Build grai metadata for a Table object.
+
+    Args:
+        current: The Table object to build grai metadata from.
+        version: The version of grai metadata to build. Defaults to "v1".
+
+    Returns:
+        TableMetadata: grai metadata object for the Table.
+
+    Raises:
+        None.
+
+    """
+
+    data = {
+        "version": version,
+        "node_type": NodeMetadataTypeLabels.column.value,
+        "node_attributes": {},
+        "tags": [config.metadata_id],
+    }
+
+    return ColumnMetadata(**data)
+
+
 @build_grai_metadata.register
 def build_grai_metadata_from_table(current: Table, version: Literal["v1"] = "v1") -> TableMetadata:
     """
@@ -68,7 +102,7 @@ def build_grai_metadata_from_table(current: Table, version: Literal["v1"] = "v1"
 
 
 @build_grai_metadata.register
-def build_grai_metadata_from_question(current: Question, version: Literal["v1"] = "v1") -> GenericNodeMetadataV1:
+def build_grai_metadata_from_question(current: Question, version: Literal["v1"] = "v1") -> QueryMetadata:
     """
     Build grai metadata for a Question object.
 
@@ -95,7 +129,7 @@ def build_grai_metadata_from_question(current: Question, version: Literal["v1"]
 
 
 @build_grai_metadata.register
-def build_grai_metadata_from_collection(current: Collection, version: Literal["v1"] = "v1") -> GenericNodeMetadataV1:
+def build_grai_metadata_from_collection(current: Collection, version: Literal["v1"] = "v1") -> CollectionMetadata:
     """
     Build grai metadata for a Collection object.
 
@@ -110,12 +144,12 @@ def build_grai_metadata_from_collection(current: Collection, version: Literal["v
 
     data = {
         "version": version,
-        "node_type": NodeMetadataTypeLabels.generic.value,
+        "node_type": NodeMetadataTypeLabels.collection.value,
         "node_attributes": {},
         "tags": [config.metadata_id],
     }
 
-    return GenericNodeMetadataV1(**data)
+    return CollectionMetadata(**data)
 
 
 @build_grai_metadata.register
@@ -252,7 +286,9 @@ def adapt_to_client(current: Any, desired: Any):
 
 
 @adapt_to_client.register
-def adapt_table_to_client(current: Table, source: SourceSpec, version: Literal["v1"] = "v1") -> SourcedNodeV1:
+def adapt_table_to_client(
+    current: Union[Table, Column], source: SourceSpec, version: Literal["v1"] = "v1"
+) -> SourcedNodeV1:
     """
     Adapt a Table object to the desired client format.
 
@@ -315,6 +351,7 @@ def adapt_collection_to_client(current: Collection, source: SourceSpec, version:
     Adapt a Collection object to the desired client format.
     Args:
         current:
+        source:
         version:
 
     Returns:

grai-integrations/source-metabase/src/grai_source_metabase/api.py

@@ -0,0 +1,110 @@
+from typing import Dict, List, Optional, Tuple
+
+from pydantic import BaseModel, Field, PrivateAttr
+
+
+class MetabaseDbmsVersion(BaseModel):
+    flavor: str
+    semantic_version: Optional[List]
+    version: str
+
+
+class DB(BaseModel):
+    dbms_version: Optional[MetabaseDbmsVersion]
+    engine: Optional[str]
+    id: int
+    name: str
+    timezone: Optional[str]
+
+
+class QuestionResultMetadata(BaseModel):
+    # id: int
+    base_type: str
+    description: Optional[str]
+    fingerprint: Optional[Dict]
+    display_name: Optional[str]
+
+
+class QuestionDataSetQuery(BaseModel):
+    database: int
+    query: Optional[Dict]
+    type: str
+
+
+class Question(BaseModel):
+    id: int
+    name: str
+    table_id: Optional[int]
+    database_id: Optional[int]
+    collection_id: Optional[int]
+    description: Optional[str]
+    archived: Optional[bool]
+    result_metadata: Optional[List[QuestionResultMetadata]]
+    dataset_query: Optional[QuestionDataSetQuery]
+    # ---- unvalidated ----#
+    creator: Optional[Dict]
+    public_uuid: Optional[str]
+    collection: Optional[Dict]
+
+
+class TableDB(BaseModel):
+    id: int
+    name: str
+    dbms_version: Optional[MetabaseDbmsVersion]
+    engine: str
+
+
+class Table(BaseModel):
+    id: int
+    name: str
+    active: Optional[bool]
+    display_name: Optional[str]
+    table_schema: str = Field(alias="schema")
+    entity_type: Optional[str]
+    db_id: int
+    description: Optional[str]
+    db: TableDB
+
+
+class FingerPrintGlobal(BaseModel):
+    distinct_count: int = Field(..., alias="distinct-count")
+    null_percent: float = Field(..., alias="nil%")
+
+    class Config:
+        allow_population_by_field_name = True
+
+
+class FingerPrint(BaseModel):
+    global_stats: Optional[FingerPrintGlobal] = Field(alias="global")
+    type_stats: Optional[Dict] = Field(alias="type")
+
+
+class TableMetadataField(BaseModel):
+    id: int
+    name: str
+    display_name: str
+    active: bool
+    description: Optional[str]
+    fingerprint: Optional[FingerPrint]
+
+
+class TableMetadata(BaseModel):
+    id: int
+    name: str
+    table_schema: str = Field(..., alias="schema")
+    active: bool
+    db: Optional[TableDB]
+    db_id: int
+    display_name: str
+    fields: Optional[List[TableMetadataField]]
+
+
+class Collection(BaseModel):
+    name: str
+    id: int
+    parent_id: Optional[str]
+    archived: Optional[bool]
+
+    @property
+    def full_name(self):
+        return f"Collection: {self.name}"