Merge pull request #643 from datafold/LAB-27_2

add --dbt support for --columns

Author: dlawin

data_diff/__main__.py

@@ -317,6 +317,7 @@ def main(conf, run, **kw):
                 json_output=kw["json_output"],
                 state=state,
                 where_flag=kw["where"],
+                columns_flag=kw["columns"],
             )
         else:
             return _data_diff(

data_diff/dbt.py

@@ -71,6 +71,7 @@ def dbt_diff(
     state: Optional[str] = None,
     log_status_handler: Optional[LogStatusHandler] = None,
     where_flag: Optional[str] = None,
+    columns_flag: Optional[Tuple[str]] = None,
 ) -> None:
     print_version_info()
     diff_threads = []
@@ -110,7 +111,7 @@ def dbt_diff(
             if log_status_handler:
                 log_status_handler.set_prefix(f"Diffing {model.alias} \n")
 
-            diff_vars = _get_diff_vars(dbt_parser, config, model, where_flag)
+            diff_vars = _get_diff_vars(dbt_parser, config, model, where_flag, columns_flag)
 
             # we won't always have a prod path when using state
             # when the model DNE in prod manifest, skip the model diff
@@ -160,7 +161,9 @@ def _get_diff_vars(
     config: TDatadiffConfig,
     model,
     where_flag: Optional[str] = None,
+    columns_flag: Optional[Tuple[str]] = None,
 ) -> TDiffVars:
+    cli_columns = list(columns_flag) if columns_flag else []
     dev_database = model.database
     dev_schema = model.schema_
 
@@ -189,10 +192,10 @@ def _get_diff_vars(
         primary_keys=primary_keys,
         connection=dbt_parser.connection,
         threads=dbt_parser.threads,
-        # --where takes precedence over any model level config
+        # cli flags take precedence over any model level config
         where_filter=where_flag or datadiff_model_config.where_filter,
-        include_columns=datadiff_model_config.include_columns,
-        exclude_columns=datadiff_model_config.exclude_columns,
+        include_columns=cli_columns or datadiff_model_config.include_columns,
+        exclude_columns=[] if cli_columns else datadiff_model_config.exclude_columns,
     )
 
 

tests/test_dbt.py

@@ -30,8 +30,9 @@ class TestDbtDiffer(unittest.TestCase):
     def test_integration_basic_dbt(self):
         artifacts_path = os.getcwd() + "/tests/dbt_artifacts"
         test_project_path = os.environ.get("DATA_DIFF_DBT_PROJ") or artifacts_path
+        test_profiles_path = os.environ.get("DATA_DIFF_DBT_PROFILES") or artifacts_path
         diff = run_datadiff_cli(
-            "--dbt", "--dbt-project-dir", test_project_path, "--dbt-profiles-dir", test_project_path
+            "--dbt", "--dbt-project-dir", test_project_path, "--dbt-profiles-dir", test_profiles_path
         )
 
         # assertions for the diff that exists in tests/dbt_artifacts/jaffle_shop.duckdb
@@ -933,3 +934,37 @@ def test_get_diff_vars_call_get_prod_path_from_manifest(
         mock_prod_path_from_manifest.assert_called_once_with(mock_model, mock_dbt_parser.prod_manifest_obj)
         self.assertEqual(diff_vars.prod_path[0], mock_prod_path_from_manifest.return_value[0])
         self.assertEqual(diff_vars.prod_path[1], mock_prod_path_from_manifest.return_value[1])
+
+    @patch("data_diff.dbt._get_prod_path_from_config")
+    @patch("data_diff.dbt._get_prod_path_from_manifest")
+    def test_get_diff_vars_cli_columns(self, mock_prod_path_from_manifest, mock_prod_path_from_config):
+        config = TDatadiffConfig(prod_database="prod_db")
+        mock_model = Mock()
+        primary_keys = ["a_primary_key"]
+        mock_model.database = "a_dev_db"
+        mock_model.schema_ = "a_schema"
+        mock_model.config.schema_ = None
+        mock_model.config.database = None
+        mock_model.alias = "a_model_name"
+        mock_model.unique_id = "unique_id"
+        mock_tdatadiffmodelconfig = Mock()
+        mock_tdatadiffmodelconfig.where_filter = "where"
+        mock_tdatadiffmodelconfig.include_columns = ["include"]
+        mock_tdatadiffmodelconfig.exclude_columns = ["exclude"]
+        mock_dbt_parser = Mock()
+        mock_dbt_parser.get_datadiff_model_config.return_value = mock_tdatadiffmodelconfig
+        mock_dbt_parser.connection = {}
+        mock_dbt_parser.threads = 0
+        mock_dbt_parser.get_pk_from_model.return_value = primary_keys
+        mock_dbt_parser.requires_upper = False
+        mock_dbt_parser.prod_manifest_obj = None
+        mock_prod_path_from_config.return_value = ("prod_db", "prod_schema")
+        cli_columns = ("col1", "col2")
+
+        diff_vars = _get_diff_vars(mock_dbt_parser, config, mock_model, where_flag=None, columns_flag=cli_columns)
+
+        mock_dbt_parser.get_pk_from_model.assert_called_once()
+        mock_prod_path_from_config.assert_called_once_with(config, mock_model, mock_model.database, mock_model.schema_)
+        mock_prod_path_from_manifest.assert_not_called()
+        self.assertEqual(diff_vars.include_columns, list(cli_columns))
+        self.assertEqual(diff_vars.exclude_columns, [])