fix(ingest/dbt): fix config validator for skip_sources_in_lineage (#1…

…0098)
datahub-project · Mar 21, 2024 · af06f95 · af06f95
1 parent d552106
commit af06f95
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 2 deletions.
diff --git a/metadata-ingestion/src/datahub/ingestion/source/dbt/dbt_common.py b/metadata-ingestion/src/datahub/ingestion/source/dbt/dbt_common.py
@@ -407,8 +407,12 @@ def validate_include_column_lineage(
     def validate_skip_sources_in_lineage(
         cls, skip_sources_in_lineage: bool, values: Dict
     ) -> bool:
-        entites_enabled: DBTEntitiesEnabled = values["entities_enabled"]
-        if skip_sources_in_lineage and entites_enabled.sources == EmitDirective.YES:
+        entites_enabled: Optional[DBTEntitiesEnabled] = values.get("entities_enabled")
+        if (
+            skip_sources_in_lineage
+            and entites_enabled
+            and entites_enabled.sources == EmitDirective.YES
+        ):
             raise ValueError(
                 "When `skip_sources_in_lineage` is enabled, `entities_enabled.sources` must be set to NO."
             )

diff --git a/metadata-ingestion/tests/unit/test_dbt_source.py b/metadata-ingestion/tests/unit/test_dbt_source.py
@@ -197,6 +197,30 @@ def test_dbt_entity_emission_configuration():
     DBTCoreConfig.parse_obj(config_dict)
 
 
+def test_dbt_config_skip_sources_in_lineage():
+    with pytest.raises(
+        ValidationError,
+        match="skip_sources_in_lineage.*entities_enabled.sources.*set to NO",
+    ):
+        config_dict = {
+            "manifest_path": "dummy_path",
+            "catalog_path": "dummy_path",
+            "target_platform": "dummy_platform",
+            "skip_sources_in_lineage": True,
+        }
+        config = DBTCoreConfig.parse_obj(config_dict)
+
+    config_dict = {
+        "manifest_path": "dummy_path",
+        "catalog_path": "dummy_path",
+        "target_platform": "dummy_platform",
+        "skip_sources_in_lineage": True,
+        "entities_enabled": {"sources": "NO"},
+    }
+    config = DBTCoreConfig.parse_obj(config_dict)
+    assert config.skip_sources_in_lineage is True
+
+
 def test_dbt_s3_config():
     # test missing aws config
     config_dict: dict = {