Merge pull request #535 from dlawin/issue_518_2

support include/exclude meta config

Author: dlawin

data_diff/cloud/datafold_api.py

@@ -103,6 +103,8 @@ class TCloudApiDataDiff(pydantic.BaseModel):
     pk_columns: List[str]
     filter1: Optional[str] = None
     filter2: Optional[str] = None
+    include_columns: Optional[List[str]]
+    exclude_columns: Optional[List[str]]
 
 
 class TSummaryResultPrimaryKeyStats(pydantic.BaseModel):

data_diff/dbt.py

@@ -1,10 +1,10 @@
 import os
 import time
 import webbrowser
+import pydantic
 import rich
 from rich.prompt import Confirm
 
-from dataclasses import dataclass
 from typing import List, Optional, Dict
 from .utils import dbt_diff_string_template, getLogger
 from pathlib import Path
@@ -32,14 +32,15 @@
 from . import connect_to_table, diff_tables, Algorithm
 
 
-@dataclass
-class DiffVars:
+class TDiffVars(pydantic.BaseModel):
     dev_path: List[str]
     prod_path: List[str]
     primary_keys: List[str]
-    connection: Dict[str, str]
-    threads: Optional[int]
+    connection: Dict[str, Optional[str]]
+    threads: Optional[int] = None
     where_filter: Optional[str] = None
+    include_columns: List[str]
+    exclude_columns: List[str]
 
 
 def dbt_diff(
@@ -122,7 +123,7 @@ def _get_diff_vars(
     config_prod_schema: Optional[str],
     config_prod_custom_schema: Optional[str],
     model,
-) -> DiffVars:
+) -> TDiffVars:
     dev_database = model.database
     dev_schema = model.schema_
 
@@ -156,19 +157,21 @@ def _get_diff_vars(
         dev_qualified_list = [x for x in [dev_database, dev_schema, model.alias] if x]
         prod_qualified_list = [x for x in [prod_database, prod_schema, model.alias] if x]
 
-    where_filter = None
-    if model.meta:
-        try:
-            where_filter = model.meta["datafold"]["datadiff"]["filter"]
-        except KeyError:
-            pass
-
-    return DiffVars(
-        dev_qualified_list, prod_qualified_list, primary_keys, dbt_parser.connection, dbt_parser.threads, where_filter
+    datadiff_model_config = dbt_parser.get_datadiff_model_config(model.meta)
+
+    return TDiffVars(
+        dev_path=dev_qualified_list,
+        prod_path=prod_qualified_list,
+        primary_keys=primary_keys,
+        connection=dbt_parser.connection,
+        threads=dbt_parser.threads,
+        where_filter=datadiff_model_config.where_filter,
+        include_columns=datadiff_model_config.include_columns,
+        exclude_columns=datadiff_model_config.exclude_columns,
     )
 
 
-def _local_diff(diff_vars: DiffVars) -> None:
+def _local_diff(diff_vars: TDiffVars) -> None:
     column_diffs_str = ""
     dev_qualified_str = ".".join(diff_vars.dev_path)
     prod_qualified_str = ".".join(diff_vars.prod_path)
@@ -189,18 +192,25 @@ def _local_diff(diff_vars: DiffVars) -> None:
         rich.print(diff_output_str)
         return
 
-    mutual_set = set(table1_columns) & set(table2_columns)
-    table1_set_diff = list(set(table1_columns) - set(table2_columns))
-    table2_set_diff = list(set(table2_columns) - set(table1_columns))
+    column_set = set(table1_columns).intersection(table2_columns)
+    table1_diff = set(table1_columns).difference(table2_columns)
+    table2_diff = set(table2_columns).difference(table1_columns)
+
+    if table1_diff:
+        column_diffs_str += f"Column(s) added: {table1_diff}\n"
+
+    if table2_diff:
+        column_diffs_str += f"Column(s) removed: {table2_diff}\n"
+
+    column_set = column_set - set(diff_vars.primary_keys)
 
-    if table1_set_diff:
-        column_diffs_str += "Column(s) added: " + str(table1_set_diff) + "\n"
+    if diff_vars.include_columns:
+        column_set = {x for x in column_set if x.upper() in [y.upper() for y in diff_vars.include_columns]}
 
-    if table2_set_diff:
-        column_diffs_str += "Column(s) removed: " + str(table2_set_diff) + "\n"
+    if diff_vars.exclude_columns:
+        column_set = {x for x in column_set if x.upper() not in [y.upper() for y in diff_vars.exclude_columns]}
 
-    mutual_set = mutual_set - set(diff_vars.primary_keys)
-    extra_columns = tuple(mutual_set)
+    extra_columns = tuple(column_set)
 
     diff = diff_tables(
         table1,
@@ -250,7 +260,7 @@ def _initialize_api() -> Optional[DatafoldAPI]:
     return DatafoldAPI(api_key=api_key, host=datafold_host)
 
 
-def _cloud_diff(diff_vars: DiffVars, datasource_id: int, api: DatafoldAPI) -> None:
+def _cloud_diff(diff_vars: TDiffVars, datasource_id: int, api: DatafoldAPI) -> None:
     diff_output_str = _diff_output_base(".".join(diff_vars.dev_path), ".".join(diff_vars.prod_path))
     payload = TCloudApiDataDiff(
         data_source1_id=datasource_id,
@@ -260,6 +270,8 @@ def _cloud_diff(diff_vars: DiffVars, datasource_id: int, api: DatafoldAPI) -> No
         pk_columns=diff_vars.primary_keys,
         filter1=diff_vars.where_filter,
         filter2=diff_vars.where_filter,
+        include_columns=diff_vars.include_columns,
+        exclude_columns=diff_vars.exclude_columns,
     )
 
     if is_tracking_enabled():

data_diff/dbt_parser.py

@@ -6,6 +6,7 @@
 from typing import List, Dict, Tuple, Set, Optional
 
 from packaging.version import parse as parse_version
+import pydantic
 
 from .utils import getLogger, get_from_dict_with_raise
 from .version import __version__
@@ -68,6 +69,12 @@ def legacy_profiles_dir() -> Path:
     return Path.home() / ".dbt"
 
 
+class TDatadiffModelConfig(pydantic.BaseModel):
+    where_filter: Optional[str] = None
+    include_columns: List[str] = []
+    exclude_columns: List[str] = []
+
+
 class DbtParser:
     def __init__(self, profiles_dir_override: str, project_dir_override: str) -> None:
         (
@@ -79,7 +86,7 @@ def __init__(self, profiles_dir_override: str, project_dir_override: str) -> Non
         ) = import_dbt_dependencies()
         self.profiles_dir = Path(profiles_dir_override or default_profiles_dir())
         self.project_dir = Path(project_dir_override or default_project_dir())
-        self.connection = None
+        self.connection = {}
         self.project_dict = self.get_project_dict()
         self.manifest_obj = self.get_manifest_obj()
         self.dbt_user_id = self.manifest_obj.metadata.user_id
@@ -95,6 +102,21 @@ def get_datadiff_variables(self) -> dict:
         vars = get_from_dict_with_raise(self.project_dict, "vars", error_message)
         return get_from_dict_with_raise(vars, "data_diff", error_message)
 
+    def get_datadiff_model_config(self, model_meta: dict) -> TDatadiffModelConfig:
+        where_filter = None
+        include_columns = []
+        exclude_columns = []
+
+        if "datafold" in model_meta and "datadiff" in model_meta["datafold"]:
+            config = model_meta["datafold"]["datadiff"]
+            where_filter = config.get("filter")
+            include_columns = config.get("include_columns") or []
+            exclude_columns = config.get("exclude_columns") or []
+
+        return TDatadiffModelConfig(
+            where_filter=where_filter, include_columns=include_columns, exclude_columns=exclude_columns
+        )
+
     def get_models(self, dbt_selection: Optional[str] = None):
         dbt_version = parse_version(self.dbt_version)
         if dbt_selection: