refactor(ingest): snowflake - update field name for column lineage (d…

…atahub-project#6263)
CybercentreCanada · Nov 18, 2022 · a9262b9 · a9262b9
1 parent 81bf49d
commit a9262b9
Show file tree

Hide file tree

Showing 3 changed files with 196 additions and 4 deletions.
diff --git a/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_lineage.py b/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_lineage.py
@@ -4,6 +4,7 @@
 from dataclasses import dataclass, field
 from typing import Dict, FrozenSet, List, Optional, Set, Tuple
 
+from pydantic import Field
 from pydantic.error_wrappers import ValidationError
 from snowflake.connector import SnowflakeConnection
 
@@ -32,7 +33,13 @@
 
 
 class SnowflakeColumnWithLineage(SnowflakeColumnReference):
-    directSourceColumns: Optional[List[SnowflakeColumnReference]] = None
+    class Config:
+        # This is for backward compatibility and can be removed later
+        allow_population_by_field_name = True
+
+    directSourceColumns: Optional[List[SnowflakeColumnReference]] = Field(
+        default=None, alias="directSources"
+    )
 
 
 @dataclass(frozen=True)
@@ -125,6 +132,7 @@ def update_lineage(self, table: SnowflakeUpstreamTable) -> None:
 
         if table.downstreamColumns:
             for col in table.downstreamColumns:
+
                 if col.directSourceColumns:
                     self.columnLineages[col.columnName].update_column_lineage(
                         col.directSourceColumns