Merge pull request #586 from dlawin/issue_483

add support for custom database config

Author: dlawin

data_diff/dbt.py

@@ -138,7 +138,13 @@ def _get_diff_vars(
 
     primary_keys = dbt_parser.get_pk_from_model(model, dbt_parser.unique_columns, "primary-key")
 
-    prod_database = config_prod_database if config_prod_database else dev_database
+    # "custom" dbt config database
+    if model.config.database:
+        prod_database = model.config.database
+    elif config_prod_database:
+        prod_database = config_prod_database
+    else:
+        prod_database = dev_database
 
     # prod schema name differs from dev schema name
     if config_prod_schema:

tests/test_dbt.py

@@ -967,6 +967,7 @@ def test_get_diff_vars_replace_custom_schema(self):
         mock_model.database = "a_dev_db"
         mock_model.schema_ = "a_custom_schema"
         mock_model.config.schema_ = mock_model.schema_
+        mock_model.config.database = None
         mock_model.alias = "a_model_name"
         mock_tdatadiffmodelconfig = Mock()
         mock_tdatadiffmodelconfig.where_filter = "where"
@@ -999,6 +1000,7 @@ def test_get_diff_vars_static_custom_schema(self):
         primary_keys = ["a_primary_key"]
         mock_model.database = "a_dev_db"
         mock_model.schema_ = "a_custom_schema"
+        mock_model.config.database = None
         mock_model.config.schema_ = mock_model.schema_
         mock_model.alias = "a_model_name"
         mock_tdatadiffmodelconfig = Mock()
@@ -1031,6 +1033,7 @@ def test_get_diff_vars_no_custom_schema_on_model(self):
         mock_model.database = "a_dev_db"
         mock_model.schema_ = "a_custom_schema"
         mock_model.config.schema_ = None
+        mock_model.config.database = None
         mock_model.alias = "a_model_name"
         mock_tdatadiffmodelconfig = Mock()
         mock_tdatadiffmodelconfig.where_filter = "where"
@@ -1060,6 +1063,7 @@ def test_get_diff_vars_match_dev_schema(self):
         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_tdatadiffmodelconfig = Mock()
         mock_tdatadiffmodelconfig.where_filter = "where"
@@ -1107,6 +1111,7 @@ def test_get_diff_vars_meta_where(self):
         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_tdatadiffmodelconfig = Mock()
         mock_tdatadiffmodelconfig.where_filter = "where"
@@ -1136,6 +1141,7 @@ def test_get_diff_vars_meta_unrelated(self):
         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_tdatadiffmodelconfig = Mock()
         mock_tdatadiffmodelconfig.where_filter = "where"
@@ -1165,6 +1171,7 @@ def test_get_diff_vars_meta_none(self):
         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_tdatadiffmodelconfig = Mock()
         mock_tdatadiffmodelconfig.where_filter = "where"
@@ -1176,7 +1183,6 @@ def test_get_diff_vars_meta_none(self):
         mock_dbt_parser.threads = 0
         mock_dbt_parser.get_pk_from_model.return_value = primary_keys
         mock_dbt_parser.requires_upper = False
-        where = None
         mock_model.meta = None
 
         diff_vars = _get_diff_vars(mock_dbt_parser, prod_database, None, None, mock_model)
@@ -1188,3 +1194,34 @@ def test_get_diff_vars_meta_none(self):
         assert diff_vars.threads == mock_dbt_parser.threads
         self.assertEqual(diff_vars.where_filter, mock_tdatadiffmodelconfig.where_filter)
         mock_dbt_parser.get_pk_from_model.assert_called_once()
+
+    def test_get_diff_vars_custom_db(self):
+        mock_model = Mock()
+        prod_database = "a_prod_db"
+        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 = "custom_database"
+        mock_model.alias = "a_model_name"
+        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_model.meta = None
+
+        diff_vars = _get_diff_vars(mock_dbt_parser, prod_database, None, None, mock_model)
+
+        assert diff_vars.dev_path == [mock_model.database, mock_model.schema_, mock_model.alias]
+        assert diff_vars.prod_path == [mock_model.config.database, mock_model.schema_, mock_model.alias]
+        assert diff_vars.primary_keys == primary_keys
+        assert diff_vars.connection == mock_dbt_parser.connection
+        assert diff_vars.threads == mock_dbt_parser.threads
+        self.assertEqual(diff_vars.where_filter, mock_tdatadiffmodelconfig.where_filter)
+        mock_dbt_parser.get_pk_from_model.assert_called_once()