Add test and flake8

Signed-off-by: popcorny <[email protected]>

Author: popcornylu

.flake8

@@ -0,0 +1,15 @@
+[flake8]
+ignore = W605, F401, E501, E125, E126
+max-line-length = 120
+exclude = tests,
+          docs,
+          .git,
+          .venv,
+          venv,
+          debug,
+          build,
+          dist,
+          images,
+          js,
+          *.egg-info
+

Makefile

@@ -0,0 +1,5 @@
+flake8:
+	@flake8
+	@echo "Passed"
+test:
+	@pytest tests

recce/dbt.py

@@ -5,10 +5,15 @@
 from dbt.adapters.factory import get_adapter_by_type
 from dbt.adapters.sql import SQLAdapter
 from dbt.cli.main import dbtRunner
+from dbt.config.profile import Profile
+from dbt.config.project import Project
 from dbt.config.runtime import load_profile, load_project
 from dbt.contracts.graph.manifest import WritableManifest, Manifest
-from dbt.config.project import Project
-from dbt.config.profile import Profile
+from dbt.contracts.results import CatalogArtifact
+
+
+def hello():
+    return 'world'
 
 
 @dataclass
@@ -19,6 +24,8 @@ class DBTContext:
     manifest: Manifest
     curr_manifest: WritableManifest = None
     base_manifest: WritableManifest = None
+    curr_catalog: CatalogArtifact = None
+    base_catalog: CatalogArtifact = None
 
     @classmethod
     def load(cls, target=None):
@@ -36,12 +43,24 @@ def load(cls, target=None):
         curr_manifest = WritableManifest.read_and_check_versions('target/manifest.json')
         base_manifest = WritableManifest.read_and_check_versions('target-base/manifest.json')
 
+        curr_catalog = None
+        if os.path.isfile('target/catalog.json'):
+            curr_catalog = CatalogArtifact.read_and_check_versions('target/catalog.json')
+
+        base_catalog = None
+        if os.path.isfile('target-base/manifest.json'):
+            base_catalog = CatalogArtifact.read_and_check_versions('target-base/manifest.json')
+
+        CatalogArtifact.to_dict()
+
         return cls(profile=profile,
                    project=project,
                    adapter=adapter,
                    manifest=manifest,
                    curr_manifest=curr_manifest,
-                   base_manifest=base_manifest)
+                   base_manifest=base_manifest,
+                   curr_catalog=curr_catalog,
+                   base_catalog=base_catalog)
 
     def find_resource_by_name(self, resource_name, base=False):
 

recce/diff.py

@@ -1,4 +1,5 @@
 import difflib
+
 import pandas as pd
 
 
@@ -27,5 +28,5 @@ def diff_dataframe(before: pd.DataFrame, after: pd.DataFrame):
         return
 
     before_aligned, after_aligned = before.align(after)
-    diff = before_aligned.compare(after_aligned,  result_names=('base', 'current'))
+    diff = before_aligned.compare(after_aligned, result_names=('base', 'current'))
     print(diff.to_string(na_rep='-'))

recce/impact.py

@@ -1,10 +1,11 @@
 import csv
 import io
+
 import agate
 import pandas as pd
+from dbt.contracts.graph.nodes import ModelNode, AnalysisNode
 
 from .dbt import DBTContext
-from dbt.contracts.graph.nodes import ModelNode, AnalysisNode
 
 
 def _dump_result(result: agate.Table):
@@ -35,7 +36,7 @@ def inspect_model_summary(dbt_context: DBTContext, model: ModelNode):
 
         output += f"identity: {model.identifier}\n"
         output += f"rows: {row_count}\n"
-        output += f"columns:\n"
+        output += "columns:\n"
         for column in columns:
             output += f"  {column.name} {column.dtype}\n"
 

recce/server.py

@@ -1,12 +1,13 @@
+from pathlib import Path
 from typing import Optional
 
 from fastapi import FastAPI, Request, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.staticfiles import StaticFiles
-from pathlib import Path
+from pydantic import BaseModel
+
 from .dbt import DBTContext
 from .impact import inspect_sql
-from pydantic import BaseModel
 
 app = FastAPI()
 dbt_context: DBTContext = None
@@ -65,9 +66,8 @@ async def get_lineage(base: Optional[bool] = False):
 
         manifest = dbt_context.curr_manifest if base is False else dbt_context.base_manifest
 
-        parent_map = json.loads(json.dumps(manifest.parent_map, cls=dbt.utils.JSONEncoder))
-
-
+        parent_map: dict = json.loads(json.dumps(manifest.parent_map, cls=dbt.utils.JSONEncoder))
+        parent_map = {k: v for k, v in parent_map.items() if not k.startswith('test.')}
 
         nodes = {}
 
@@ -99,6 +99,21 @@ async def get_lineage(base: Optional[bool] = False):
                 'resource_type': exposure['resource_type'],
                 'package_name': exposure['package_name'],
             }
+        for metric in json.loads(json.dumps(manifest.metrics, cls=dbt.utils.JSONEncoder)).values():
+            nodes[metric['unique_id']] = {
+                'id': metric['unique_id'],
+                'name': metric['name'],
+                'resource_type': metric['resource_type'],
+                'package_name': metric['package_name'],
+            }
+
+        for semantic_models in json.loads(json.dumps(manifest.semantic_models, cls=dbt.utils.JSONEncoder)).values():
+            nodes[semantic_models['unique_id']] = {
+                'id': semantic_models['unique_id'],
+                'name': semantic_models['name'],
+                'resource_type': semantic_models['resource_type'],
+                'package_name': semantic_models['package_name'],
+            }
 
         return dict(parent_map=parent_map, nodes=nodes)
     except Exception as e: