Update sql_common.py

datahub-project · Dec 2, 2024 · 5ef2404 · 5ef2404
1 parent 847cce9
commit 5ef2404
Showing 1 changed file with 37 additions and 35 deletions.
diff --git a/metadata-ingestion/src/datahub/ingestion/source/sql/sql_common.py b/metadata-ingestion/src/datahub/ingestion/source/sql/sql_common.py
@@ -975,46 +975,48 @@ def get_fine_grained_lineages(
         self,
         dataset_urn: str,
         upstream_dataset_urn: str,
-        schema_fields: Optional[List[SchemaField]],
+        schema_fields: List[SchemaField],
     ) -> Optional[List[FineGrainedLineage]]:
         def simplify_field_path(field_path):
             return Dataset._simplify_field_path(field_path)
 
-        upstream_schema_metadata: Optional[
-            SchemaMetadata
-        ] = self.ctx.graph.get_schema_metadata(upstream_dataset_urn)
-
-        if schema_fields and upstream_schema_metadata:
-            fine_grained_lineages: List[FineGrainedLineage] = []
-            for schema_field in schema_fields:
-                field_path_v1 = simplify_field_path(schema_field.fieldPath)
-                matching_upstream_field = next(
-                    (
-                        f
-                        for f in upstream_schema_metadata.fields
-                        if simplify_field_path(f.fieldPath) == field_path_v1
-                    ),
-                    None,
-                )
-                if matching_upstream_field:
-                    fine_grained_lineages.append(
-                        FineGrainedLineage(
-                            downstreamType=FineGrainedLineageDownstreamType.FIELD,
-                            downstreams=[
-                                make_schema_field_urn(dataset_urn, field_path_v1)
-                            ],
-                            upstreamType=FineGrainedLineageUpstreamType.FIELD_SET,
-                            upstreams=[
-                                make_schema_field_urn(
-                                    upstream_dataset_urn,
-                                    simplify_field_path(
-                                        matching_upstream_field.fieldPath
-                                    ),
-                                )
-                            ],
-                        )
+        if self.ctx.graph:
+            upstream_schema_metadata: Optional[
+                SchemaMetadata
+            ] = self.ctx.graph.get_schema_metadata(upstream_dataset_urn)
+
+            if schema_fields and upstream_schema_metadata:
+                fine_grained_lineages: List[FineGrainedLineage] = []
+                for schema_field in schema_fields:
+                    field_path_v1 = simplify_field_path(schema_field.fieldPath)
+                    matching_upstream_field = next(
+                        (
+                            f
+                            for f in upstream_schema_metadata.fields
+                            if simplify_field_path(f.fieldPath) == field_path_v1
+                        ),
+                        None,
                     )
-            return fine_grained_lineages
+                    if matching_upstream_field:
+                        fine_grained_lineages.append(
+                            FineGrainedLineage(
+                                downstreamType=FineGrainedLineageDownstreamType.FIELD,
+                                downstreams=[
+                                    make_schema_field_urn(dataset_urn, field_path_v1)
+                                ],
+                                upstreamType=FineGrainedLineageUpstreamType.FIELD_SET,
+                                upstreams=[
+                                    make_schema_field_urn(
+                                        upstream_dataset_urn,
+                                        simplify_field_path(
+                                            matching_upstream_field.fieldPath
+                                        ),
+                                    )
+                                ],
+                            )
+                        )
+                if fine_grained_lineages:
+                    return fine_grained_lineages
         return None
 
     def get_schema_fields(