fix(Ingestor/Glue): Implement column upstream lineage between S3 and …

…Glue

metadata-ingestion/src/datahub/ingestion/source/aws/glue.py

@@ -24,6 +24,7 @@
 from pydantic import validator
 from pydantic.fields import Field
 
+import datahub.emitter.mce_builder as builder
 from datahub.configuration.common import AllowDenyPattern
 from datahub.configuration.source_common import DatasetSourceConfigMixin
 from datahub.emitter import mce_builder
@@ -55,7 +56,11 @@
 from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.ingestion.source.aws import s3_util
 from datahub.ingestion.source.aws.aws_common import AwsSourceConfig
-from datahub.ingestion.source.aws.s3_util import is_s3_uri, make_s3_urn
+from datahub.ingestion.source.aws.s3_util import (
+    is_s3_uri,
+    make_s3_urn,
+    make_s3_urn_for_lineage,
+)
 from datahub.ingestion.source.common.subtypes import (
     DatasetContainerSubTypes,
     DatasetSubTypes,
@@ -90,13 +95,17 @@
     DatasetLineageTypeClass,
     DatasetProfileClass,
     DatasetPropertiesClass,
+    FineGrainedLineageClass,
+    FineGrainedLineageDownstreamTypeClass,
+    FineGrainedLineageUpstreamTypeClass,
     GlobalTagsClass,
     MetadataChangeEventClass,
     OwnerClass,
     OwnershipClass,
     OwnershipTypeClass,
     PartitionSpecClass,
     PartitionTypeClass,
+    SchemaMetadataClass,
     TagAssociationClass,
     UpstreamClass,
     UpstreamLineageClass,
@@ -171,6 +180,11 @@ class GlueSourceConfig(
         description="If enabled, delta schemas can be alternatively fetched from table parameters.",
     )
 
+    include_column_lineage: bool = Field(
+        default=True,
+        description="When enabled, column-level lineage will be extracted from the s3.",
+    )
+
     def is_profiling_enabled(self) -> bool:
         return self.profiling is not None and is_profiling_enabled(
             self.profiling.operation_config
@@ -283,6 +297,7 @@ class GlueSource(StatefulIngestionSourceBase):
 
     def __init__(self, config: GlueSourceConfig, ctx: PipelineContext):
         super().__init__(config, ctx)
+        self.ctx = ctx
         self.extract_owners = config.extract_owners
         self.source_config = config
         self.report = GlueSourceReport()
@@ -714,18 +729,43 @@ def get_lineage_if_enabled(
             dataset_properties: Optional[
                 DatasetPropertiesClass
             ] = mce_builder.get_aspect_if_available(mce, DatasetPropertiesClass)
+            # extract dataset schema aspect
+            schema_metadata: Optional[
+                SchemaMetadataClass
+            ] = mce_builder.get_aspect_if_available(mce, SchemaMetadataClass)
+
             if dataset_properties and "Location" in dataset_properties.customProperties:
                 location = dataset_properties.customProperties["Location"]
                 if is_s3_uri(location):
-                    s3_dataset_urn = make_s3_urn(location, self.source_config.env)
+                    s3_dataset_urn = make_s3_urn_for_lineage(
+                        location, self.source_config.env
+                    )
+                    assert self.ctx.graph
+                    schema_metadata_for_s3: Optional[
+                        SchemaMetadataClass
+                    ] = self.ctx.graph.get_schema_metadata(s3_dataset_urn)
+
                     if self.source_config.glue_s3_lineage_direction == "upstream":
+                        fine_grained_lineages = None
+                        if (
+                            self.source_config.include_column_lineage
+                            and schema_metadata
+                            and schema_metadata_for_s3
+                        ):
+                            fine_grained_lineages = self.get_fine_grained_lineages(
+                                mce.proposedSnapshot.urn,
+                                s3_dataset_urn,
+                                schema_metadata,
+                                schema_metadata_for_s3,
+                            )
                         upstream_lineage = UpstreamLineageClass(
                             upstreams=[
                                 UpstreamClass(
                                     dataset=s3_dataset_urn,
                                     type=DatasetLineageTypeClass.COPY,
                                 )
-                            ]
+                            ],
+                            fineGrainedLineages=fine_grained_lineages or None,
                         )
                         return MetadataChangeProposalWrapper(
                             entityUrn=mce.proposedSnapshot.urn,
@@ -747,6 +787,45 @@ def get_lineage_if_enabled(
                         ).as_workunit()
         return None
 
+    def get_fine_grained_lineages(
+        self,
+        dataset_urn: str,
+        s3_dataset_urn: str,
+        schema_metadata: SchemaMetadata,
+        schema_metadata_for_s3: SchemaMetadata,
+    ) -> Optional[List[FineGrainedLineageClass]]:
+        if schema_metadata and schema_metadata_for_s3:
+            fine_grained_lineages: List[FineGrainedLineageClass] = []
+            for field in schema_metadata.fields:
+                matching_s3_field = next(
+                    (
+                        f
+                        for f in schema_metadata_for_s3.fields
+                        if f.fieldPath.split(".")[-1] == field.fieldPath.split(".")[-1]
+                    ),
+                    None,
+                )
+                if matching_s3_field:
+                    fine_grained_lineages.append(
+                        FineGrainedLineageClass(
+                            downstreamType=FineGrainedLineageDownstreamTypeClass.FIELD,
+                            downstreams=[
+                                mce_builder.make_schema_field_urn(
+                                    dataset_urn, field.fieldPath.split(".")[-1]
+                                )
+                            ],
+                            upstreamType=FineGrainedLineageUpstreamTypeClass.FIELD_SET,
+                            upstreams=[
+                                mce_builder.make_schema_field_urn(
+                                    s3_dataset_urn,
+                                    matching_s3_field.fieldPath.split(".")[-1],
+                                )
+                            ],
+                        )
+                    )
+            return fine_grained_lineages
+        return None
+
     def _create_profile_mcp(
         self,
         mce: MetadataChangeEventClass,