Merge branch 'datahub-project:master' into master

metadata-ingestion/src/datahub/ingestion/source/abs/source.py

@@ -52,13 +52,15 @@
 from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.ingestion.source.abs.config import DataLakeSourceConfig, PathSpec
 from datahub.ingestion.source.abs.report import DataLakeSourceReport
-from datahub.ingestion.source.azure.abs_util import (
+from datahub.ingestion.source.azure.abs_folder_utils import (
     get_abs_properties,
     get_abs_tags,
+    list_folders,
+)
+from datahub.ingestion.source.azure.abs_utils import (
     get_container_name,
     get_container_relative_path,
     get_key_prefix,
-    list_folders,
     strip_abs_prefix,
 )
 from datahub.ingestion.source.data_lake_common.data_lake_utils import ContainerWUCreator

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

@@ -24,6 +24,7 @@
 from pydantic import validator
 from pydantic.fields import Field
 
+from datahub.api.entities.dataset.dataset import Dataset
 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,49 @@ 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]]:
+        def simplify_field_path(field_path):
+            return Dataset._simplify_field_path(field_path)
+
+        if schema_metadata and schema_metadata_for_s3:
+            fine_grained_lineages: List[FineGrainedLineageClass] = []
+            for field in schema_metadata.fields:
+                field_path_v1 = simplify_field_path(field.fieldPath)
+                matching_s3_field = next(
+                    (
+                        f
+                        for f in schema_metadata_for_s3.fields
+                        if simplify_field_path(f.fieldPath) == field_path_v1
+                    ),
+                    None,
+                )
+                if matching_s3_field:
+                    fine_grained_lineages.append(
+                        FineGrainedLineageClass(
+                            downstreamType=FineGrainedLineageDownstreamTypeClass.FIELD,
+                            downstreams=[
+                                mce_builder.make_schema_field_urn(
+                                    dataset_urn, field_path_v1
+                                )
+                            ],
+                            upstreamType=FineGrainedLineageUpstreamTypeClass.FIELD_SET,
+                            upstreams=[
+                                mce_builder.make_schema_field_urn(
+                                    s3_dataset_urn,
+                                    simplify_field_path(matching_s3_field.fieldPath),
+                                )
+                            ],
+                        )
+                    )
+            return fine_grained_lineages
+        return None
+
     def _create_profile_mcp(
         self,
         mce: MetadataChangeEventClass,

metadata-ingestion/src/datahub/ingestion/source/azure/abs_util.py → metadata-ingestion/src/datahub/ingestion/source/azure/abs_folder_utils.py

@@ -1,6 +1,4 @@
 import logging
-import os
-import re
 from typing import Dict, Iterable, List, Optional
 
 from azure.storage.blob import BlobProperties
@@ -10,67 +8,10 @@
 from datahub.ingestion.source.azure.azure_common import AzureConnectionConfig
 from datahub.metadata.schema_classes import GlobalTagsClass, TagAssociationClass
 
-ABS_PREFIXES_REGEX = re.compile(
-    r"(http[s]?://[a-z0-9]{3,24}\.blob\.core\.windows\.net/)"
-)
-
 logging.getLogger("py4j").setLevel(logging.ERROR)
 logger: logging.Logger = logging.getLogger(__name__)
 
 
-def is_abs_uri(uri: str) -> bool:
-    return bool(ABS_PREFIXES_REGEX.match(uri))
-
-
-def get_abs_prefix(abs_uri: str) -> Optional[str]:
-    result = re.search(ABS_PREFIXES_REGEX, abs_uri)
-    if result and result.groups():
-        return result.group(1)
-    return None
-
-
-def strip_abs_prefix(abs_uri: str) -> str:
-    # remove abs prefix https://<storage-account>.blob.core.windows.net
-    abs_prefix = get_abs_prefix(abs_uri)
-    if not abs_prefix:
-        raise ValueError(
-            f"Not an Azure Blob Storage URI. Must match the following regular expression: {str(ABS_PREFIXES_REGEX)}"
-        )
-    length_abs_prefix = len(abs_prefix)
-    return abs_uri[length_abs_prefix:]
-
-
-def make_abs_urn(abs_uri: str, env: str) -> str:
-    abs_name = strip_abs_prefix(abs_uri)
-
-    if abs_name.endswith("/"):
-        abs_name = abs_name[:-1]
-
-    name, extension = os.path.splitext(abs_name)
-
-    if extension != "":
-        extension = extension[1:]  # remove the dot
-        return f"urn:li:dataset:(urn:li:dataPlatform:abs,{name}_{extension},{env})"
-
-    return f"urn:li:dataset:(urn:li:dataPlatform:abs,{abs_name},{env})"
-
-
-def get_container_name(abs_uri: str) -> str:
-    if not is_abs_uri(abs_uri):
-        raise ValueError(
-            f"Not an Azure Blob Storage URI. Must match the following regular expression: {str(ABS_PREFIXES_REGEX)}"
-        )
-    return strip_abs_prefix(abs_uri).split("/")[0]
-
-
-def get_key_prefix(abs_uri: str) -> str:
-    if not is_abs_uri(abs_uri):
-        raise ValueError(
-            f"Not an Azure Blob Storage URI. Must match the following regular expression: {str(ABS_PREFIXES_REGEX)}"
-        )
-    return strip_abs_prefix(abs_uri).split("/", maxsplit=1)[1]
-
-
 def get_abs_properties(
     container_name: str,
     blob_name: Optional[str],
@@ -280,7 +221,3 @@ def list_folders(
             this_dict[folder_name] = folder_name
 
         yield f"{folder_name}"
-
-
-def get_container_relative_path(abs_uri: str) -> str:
-    return "/".join(strip_abs_prefix(abs_uri).split("/")[1:])

metadata-ingestion/src/datahub/ingestion/source/azure/abs_utils.py

@@ -0,0 +1,66 @@
+import os
+import re
+from typing import Optional
+
+# This file should not import any abs spectific modules as we import it in path_spec.py in datat_lake_common.py
+
+ABS_PREFIXES_REGEX = re.compile(
+    r"(http[s]?://[a-z0-9]{3,24}\.blob\.core\.windows\.net/)"
+)
+
+
+def is_abs_uri(uri: str) -> bool:
+    return bool(ABS_PREFIXES_REGEX.match(uri))
+
+
+def get_abs_prefix(abs_uri: str) -> Optional[str]:
+    result = re.search(ABS_PREFIXES_REGEX, abs_uri)
+    if result and result.groups():
+        return result.group(1)
+    return None
+
+
+def strip_abs_prefix(abs_uri: str) -> str:
+    # remove abs prefix https://<storage-account>.blob.core.windows.net
+    abs_prefix = get_abs_prefix(abs_uri)
+    if not abs_prefix:
+        raise ValueError(
+            f"Not an Azure Blob Storage URI. Must match the following regular expression: {str(ABS_PREFIXES_REGEX)}"
+        )
+    length_abs_prefix = len(abs_prefix)
+    return abs_uri[length_abs_prefix:]
+
+
+def make_abs_urn(abs_uri: str, env: str) -> str:
+    abs_name = strip_abs_prefix(abs_uri)
+
+    if abs_name.endswith("/"):
+        abs_name = abs_name[:-1]
+
+    name, extension = os.path.splitext(abs_name)
+
+    if extension != "":
+        extension = extension[1:]  # remove the dot
+        return f"urn:li:dataset:(urn:li:dataPlatform:abs,{name}_{extension},{env})"
+
+    return f"urn:li:dataset:(urn:li:dataPlatform:abs,{abs_name},{env})"
+
+
+def get_container_name(abs_uri: str) -> str:
+    if not is_abs_uri(abs_uri):
+        raise ValueError(
+            f"Not an Azure Blob Storage URI. Must match the following regular expression: {str(ABS_PREFIXES_REGEX)}"
+        )
+    return strip_abs_prefix(abs_uri).split("/")[0]
+
+
+def get_key_prefix(abs_uri: str) -> str:
+    if not is_abs_uri(abs_uri):
+        raise ValueError(
+            f"Not an Azure Blob Storage URI. Must match the following regular expression: {str(ABS_PREFIXES_REGEX)}"
+        )
+    return strip_abs_prefix(abs_uri).split("/", maxsplit=1)[1]
+
+
+def get_container_relative_path(abs_uri: str) -> str:
+    return "/".join(strip_abs_prefix(abs_uri).split("/")[1:])

metadata-ingestion/src/datahub/ingestion/source/data_lake_common/data_lake_utils.py

@@ -16,7 +16,7 @@
     get_s3_prefix,
     is_s3_uri,
 )
-from datahub.ingestion.source.azure.abs_util import (
+from datahub.ingestion.source.azure.abs_utils import (
     get_abs_prefix,
     get_container_name,
     get_container_relative_path,

metadata-ingestion/src/datahub/ingestion/source/data_lake_common/path_spec.py

@@ -11,7 +11,7 @@
 
 from datahub.configuration.common import ConfigModel
 from datahub.ingestion.source.aws.s3_util import is_s3_uri
-from datahub.ingestion.source.azure.abs_util import is_abs_uri
+from datahub.ingestion.source.azure.abs_utils import is_abs_uri
 from datahub.ingestion.source.gcs.gcs_utils import is_gcs_uri
 
 # hide annoying debug errors from py4j