acryl-datahub 1.1.1rc4__py3-none-any.whl → 1.3.0.1rc9__py3-none-any.whl

datahub/ingestion/source/excel/source.py

@@ -0,0 +1,662 @@
+import glob
+import io
+import logging
+import os
+import re
+from datetime import datetime, timezone
+from enum import Enum, auto
+from io import BytesIO
+from pathlib import PurePath
+from typing import Any, Dict, Iterable, Iterator, List, Optional, Tuple, Type, Union
+from urllib.parse import urlparse
+
+from datahub.emitter.mce_builder import make_dataset_urn_with_platform_instance
+from datahub.emitter.mcp import MetadataChangeProposalWrapper
+from datahub.ingestion.api.common import PipelineContext
+from datahub.ingestion.api.decorators import (
+    SupportStatus,
+    capability,
+    config_class,
+    platform_name,
+    support_status,
+)
+from datahub.ingestion.api.source import MetadataWorkUnitProcessor, SourceCapability
+from datahub.ingestion.api.workunit import MetadataWorkUnit
+from datahub.ingestion.source.aws.s3_boto_utils import get_s3_tags
+from datahub.ingestion.source.aws.s3_util import (
+    get_bucket_name,
+    get_bucket_relative_path,
+    strip_s3_prefix,
+)
+from datahub.ingestion.source.azure.abs_folder_utils import (
+    get_abs_tags,
+)
+from datahub.ingestion.source.azure.abs_utils import (
+    get_container_relative_path,
+    strip_abs_prefix,
+)
+from datahub.ingestion.source.data_lake_common.data_lake_utils import ContainerWUCreator
+from datahub.ingestion.source.excel.config import ExcelSourceConfig
+from datahub.ingestion.source.excel.excel_file import ExcelFile, ExcelTable
+from datahub.ingestion.source.excel.profiling import ExcelProfiler
+from datahub.ingestion.source.excel.report import ExcelSourceReport
+from datahub.ingestion.source.excel.util import gen_dataset_name
+from datahub.ingestion.source.s3.source import BrowsePath
+from datahub.ingestion.source.state.stale_entity_removal_handler import (
+    StaleEntityRemovalHandler,
+)
+from datahub.ingestion.source.state.stateful_ingestion_base import (
+    StatefulIngestionSourceBase,
+)
+from datahub.metadata.com.linkedin.pegasus2avro.common import TimeStamp
+from datahub.metadata.schema_classes import (
+    BooleanTypeClass,
+    ChangeTypeClass,
+    DatasetPropertiesClass,
+    DateTypeClass,
+    GlobalTagsClass,
+    NullTypeClass,
+    NumberTypeClass,
+    OtherSchemaClass,
+    RecordTypeClass,
+    SchemaFieldClass as SchemaField,
+    SchemaFieldDataTypeClass as SchemaFieldDataType,
+    SchemaMetadataClass as SchemaMetadata,
+    StringTypeClass,
+)
+from datahub.utilities.perf_timer import PerfTimer
+
+logger: logging.Logger = logging.getLogger(__name__)
+
+field_type_mapping: Dict[str, Type] = {
+    "int8": NumberTypeClass,
+    "int16": NumberTypeClass,
+    "int32": NumberTypeClass,
+    "int64": NumberTypeClass,
+    "uint8": NumberTypeClass,
+    "uint16": NumberTypeClass,
+    "uint32": NumberTypeClass,
+    "uint64": NumberTypeClass,
+    "Int8": NumberTypeClass,
+    "Int16": NumberTypeClass,
+    "Int32": NumberTypeClass,
+    "Int64": NumberTypeClass,
+    "UInt8": NumberTypeClass,
+    "UInt16": NumberTypeClass,
+    "UInt32": NumberTypeClass,
+    "UInt64": NumberTypeClass,
+    "intp": NumberTypeClass,
+    "uintp": NumberTypeClass,
+    "float16": NumberTypeClass,
+    "float32": NumberTypeClass,
+    "float64": NumberTypeClass,
+    "float128": NumberTypeClass,
+    "Float32": NumberTypeClass,
+    "Float64": NumberTypeClass,
+    "complex64": NumberTypeClass,
+    "complex128": NumberTypeClass,
+    "complex256": NumberTypeClass,
+    "bool": BooleanTypeClass,
+    "boolean": BooleanTypeClass,
+    "object": StringTypeClass,
+    "string": StringTypeClass,
+    "datetime64": DateTypeClass,
+    "datetime64[ns]": DateTypeClass,
+    "datetime64[ns, tz]": DateTypeClass,
+    "timedelta64": DateTypeClass,
+    "timedelta64[ns]": DateTypeClass,
+    "period": DateTypeClass,
+    "period[D]": DateTypeClass,
+    "period[M]": DateTypeClass,
+    "period[Y]": DateTypeClass,
+    "category": RecordTypeClass,
+    "interval": RecordTypeClass,
+    "sparse": RecordTypeClass,
+    "NA": NullTypeClass,
+}
+
+
+ALLOWED_EXTENSIONS = [".xlsx", ".xlsm", ".xltx", ".xltm"]
+
+
+class UriType(Enum):
+    HTTP = auto()
+    HTTPS = auto()
+    LOCAL_FILE = auto()
+    ABSOLUTE_PATH = auto()
+    RELATIVE_PATH = auto()
+    S3 = auto()
+    S3A = auto()
+    ABS = auto()
+    UNKNOWN = auto()
+
+
+@platform_name("Excel")
+@config_class(ExcelSourceConfig)
+@support_status(SupportStatus.INCUBATING)
+@capability(SourceCapability.CONTAINERS, "Enabled by default")
+@capability(SourceCapability.DATA_PROFILING, "Optionally enabled via configuration")
+@capability(SourceCapability.SCHEMA_METADATA, "Enabled by default")
+@capability(
+    SourceCapability.DELETION_DETECTION,
+    "Optionally enabled via `stateful_ingestion.remove_stale_metadata`",
+    supported=True,
+)
+class ExcelSource(StatefulIngestionSourceBase):
+    config: ExcelSourceConfig
+    report: ExcelSourceReport
+    container_WU_creator: ContainerWUCreator
+    platform: str = "excel"
+
+    def __init__(self, ctx: PipelineContext, config: ExcelSourceConfig):
+        super().__init__(config, ctx)
+        self.ctx = ctx
+        self.config = config
+        self.report: ExcelSourceReport = ExcelSourceReport()
+
+    @classmethod
+    def create(cls, config_dict: dict, ctx: PipelineContext) -> "ExcelSource":
+        config = ExcelSourceConfig.parse_obj(config_dict)
+        return cls(ctx, config)
+
+    def get_workunit_processors(self) -> List[Optional[MetadataWorkUnitProcessor]]:
+        return [
+            *super().get_workunit_processors(),
+            StaleEntityRemovalHandler.create(
+                self, self.config, self.ctx
+            ).workunit_processor,
+        ]
+
+    @staticmethod
+    def uri_type(uri: str) -> Tuple[UriType, str]:
+        if not uri or not isinstance(uri, str):
+            return UriType.UNKNOWN, ""
+
+        uri = uri.strip()
+        parsed = urlparse(uri)
+        scheme = parsed.scheme.lower()
+
+        if scheme == "http":
+            return UriType.HTTP, uri[7:]
+        elif scheme == "https":
+            if parsed.netloc and ".blob.core.windows.net" in parsed.netloc:
+                return UriType.ABS, uri[8:]
+            else:
+                return UriType.HTTPS, uri[8:]
+        elif scheme == "file":
+            if uri.startswith("file:///"):
+                return UriType.LOCAL_FILE, uri[7:]
+
+        if scheme == "s3":
+            return UriType.S3, uri[5:]
+        elif scheme == "s3a":
+            return UriType.S3A, uri[6:]
+
+        if scheme:
+            return UriType.UNKNOWN, uri[len(scheme) + 3 :]
+
+        if os.path.isabs(uri):
+            return UriType.ABSOLUTE_PATH, uri
+        else:
+            return UriType.RELATIVE_PATH, uri
+
+    @staticmethod
+    def is_excel_file(path: str) -> bool:
+        _, ext = os.path.splitext(path)
+        return ext.lower() in ALLOWED_EXTENSIONS
+
+    @staticmethod
+    def local_browser(path_spec: str) -> Iterable[BrowsePath]:
+        matching_paths = glob.glob(path_spec, recursive=True)
+        matching_files = [path for path in matching_paths if os.path.isfile(path)]
+
+        for file in sorted(matching_files):
+            full_path = PurePath(os.path.normpath(file)).as_posix()
+            yield BrowsePath(
+                file=full_path,
+                timestamp=datetime.fromtimestamp(
+                    os.path.getmtime(full_path), timezone.utc
+                ),
+                size=os.path.getsize(full_path),
+                partitions=[],
+            )
+
+    def get_local_file(self, file_path: str) -> Union[BytesIO, None]:
+        try:
+            with open(file_path, "rb") as f:
+                bytes_io = io.BytesIO(f.read())
+                bytes_io.seek(0)
+                return bytes_io
+        except Exception as e:
+            self.report.report_file_dropped(file_path)
+            self.report.warning(
+                message="Error reading local Excel file",
+                context=f"Path={file_path}",
+                exc=e,
+            )
+            return None
+
+    @staticmethod
+    def get_prefix(relative_path: str) -> str:
+        index = re.search(r"[*|{]", relative_path)
+        if index:
+            return relative_path[: index.start()]
+        else:
+            return relative_path
+
+    @staticmethod
+    def create_s3_path(bucket_name: str, key: str) -> str:
+        return f"s3://{bucket_name}/{key}"
+
+    def create_abs_path(self, key: str) -> str:
+        if self.config.azure_config:
+            account_name = self.config.azure_config.account_name
+            container_name = self.config.azure_config.container_name
+            return (
+                f"https://{account_name}.blob.core.windows.net/{container_name}/{key}"
+            )
+        return ""
+
+    @staticmethod
+    def strip_file_prefix(path: str) -> str:
+        if path.startswith("/"):
+            return path[1:]
+        else:
+            return path
+
+    def s3_browser(self, path_spec: str) -> Iterable[BrowsePath]:
+        if self.config.aws_config is None:
+            raise ValueError("aws_config not set. Cannot browse s3")
+        s3 = self.config.aws_config.get_s3_resource(self.config.verify_ssl)
+        bucket_name = get_bucket_name(path_spec)
+        logger.debug(f"Scanning bucket: {bucket_name}")
+        bucket = s3.Bucket(bucket_name)
+        prefix = self.get_prefix(get_bucket_relative_path(path_spec))
+        logger.debug(f"Scanning objects with prefix:{prefix}")
+
+        for obj in bucket.objects.filter(Prefix=prefix).page_size(1000):
+            s3_path = self.create_s3_path(obj.bucket_name, obj.key)
+            logger.debug(f"Path: {s3_path}")
+
+            yield BrowsePath(
+                file=s3_path,
+                timestamp=obj.last_modified,
+                size=obj.size,
+                partitions=[],
+                content_type=None,
+            )
+
+    def get_s3_file(self, path_spec: str) -> Union[BytesIO, None]:
+        if self.config.aws_config is None:
+            raise ValueError("aws_config not set. Cannot browse s3")
+        s3 = self.config.aws_config.get_s3_resource(self.config.verify_ssl)
+        bucket_name = get_bucket_name(path_spec)
+        key = get_bucket_relative_path(path_spec)
+        logger.debug(f"Getting file: {key} from bucket: {bucket_name}")
+        try:
+            obj = s3.Object(bucket_name, key)
+            file_content = obj.get()["Body"].read()
+            binary_stream = io.BytesIO(file_content)
+            binary_stream.seek(0)
+            return binary_stream
+        except Exception as e:
+            self.report.report_file_dropped(path_spec)
+            self.report.warning(
+                message="Error reading Excel file from S3",
+                context=f"Path={path_spec}",
+                exc=e,
+            )
+            return None
+
+    def process_s3_tags(
+        self, path_spec: str, dataset_urn: str
+    ) -> Iterable[MetadataWorkUnit]:
+        bucket_name = get_bucket_name(path_spec)
+        key = get_bucket_relative_path(path_spec)
+
+        s3_tags = get_s3_tags(
+            bucket_name,
+            key,
+            dataset_urn,
+            self.config.aws_config,
+            self.ctx,
+            self.config.use_s3_bucket_tags,
+            self.config.use_s3_object_tags,
+            self.config.verify_ssl,
+        )
+
+        if s3_tags:
+            yield from self.process_global_tags(s3_tags, dataset_urn)
+
+    def abs_browser(self, path_spec: str) -> Iterable[BrowsePath]:
+        if self.config.azure_config is None:
+            raise ValueError("azure_config not set. Cannot browse Azure Blob Storage")
+        abs_blob_service_client = self.config.azure_config.get_blob_service_client()
+        container_client = abs_blob_service_client.get_container_client(
+            self.config.azure_config.container_name
+        )
+
+        container_name = self.config.azure_config.container_name
+        logger.debug(f"Scanning container: {container_name}")
+
+        prefix = self.get_prefix(get_container_relative_path(path_spec))
+        logger.debug(f"Scanning objects with prefix: {prefix}")
+
+        for obj in container_client.list_blobs(
+            name_starts_with=f"{prefix}", results_per_page=1000
+        ):
+            abs_path = self.create_abs_path(obj.name)
+            logger.debug(f"Path: {abs_path}")
+
+            yield BrowsePath(
+                file=abs_path,
+                timestamp=obj.last_modified,
+                size=obj.size,
+                partitions=[],
+                content_type=None,
+            )
+
+    def get_abs_file(self, path_spec: str) -> Union[BytesIO, None]:
+        if self.config.azure_config is None:
+            raise ValueError("azure_config not set. Cannot browse Azure Blob Storage")
+        abs_blob_service_client = self.config.azure_config.get_blob_service_client()
+        container_client = abs_blob_service_client.get_container_client(
+            self.config.azure_config.container_name
+        )
+
+        container_name = self.config.azure_config.container_name
+        blob_path = get_container_relative_path(path_spec)
+        logger.debug(f"Getting file: {blob_path} from container: {container_name}")
+
+        try:
+            blob_client = container_client.get_blob_client(blob_path)
+            download_stream = blob_client.download_blob()
+            file_content = download_stream.readall()
+            binary_stream = io.BytesIO(file_content)
+            binary_stream.seek(0)
+            return binary_stream
+        except Exception as e:
+            self.report.report_file_dropped(path_spec)
+            self.report.warning(
+                message="Error reading Excel file from Azure Blob Storage",
+                context=f"Path={path_spec}",
+                exc=e,
+            )
+            return None
+
+    def process_abs_tags(
+        self, path_spec: str, dataset_urn: str
+    ) -> Iterable[MetadataWorkUnit]:
+        if (
+            self.config.azure_config
+            and self.config.azure_config.container_name is not None
+        ):
+            container_name = self.config.azure_config.container_name
+            blob_path = get_container_relative_path(path_spec)
+
+            abs_tags = get_abs_tags(
+                container_name,
+                blob_path,
+                dataset_urn,
+                self.config.azure_config,
+                self.ctx,
+                self.config.use_abs_blob_tags,
+            )
+
+            if abs_tags:
+                yield from self.process_global_tags(abs_tags, dataset_urn)
+
+    @staticmethod
+    def get_field_type(field_type: str) -> SchemaFieldDataType:
+        type_class = field_type_mapping.get(field_type, NullTypeClass)
+        return SchemaFieldDataType(type=type_class())
+
+    def construct_schema_field(self, f_name: str, f_type: str) -> SchemaField:
+        logger.debug(f"Field: {f_name} Type: {f_type}")
+        return SchemaField(
+            fieldPath=f_name,
+            nativeDataType=f_type,
+            type=self.get_field_type(f_type),
+            description=None,
+            nullable=False,
+            recursive=False,
+        )
+
+    def construct_schema_metadata(
+        self,
+        name: str,
+        dataset: ExcelTable,
+    ) -> SchemaMetadata:
+        canonical_schema: List[SchemaField] = []
+
+        # Get data types for each column
+        data_types = dataset.df.dtypes.to_dict()
+
+        # Convert numpy types to string representation for better readability
+        data_types = {col: str(dtype) for col, dtype in data_types.items()}
+
+        for f_name, f_type in data_types.items():
+            canonical_schema.append(self.construct_schema_field(f_name, f_type))
+
+        return SchemaMetadata(
+            schemaName=name,
+            platform=f"urn:li:dataPlatform:{self.platform}",
+            version=0,
+            hash="",
+            platformSchema=OtherSchemaClass(rawSchema=""),
+            fields=canonical_schema,
+        )
+
+    @staticmethod
+    def get_dataset_attributes(metadata: Dict[str, Any]) -> dict:
+        result = {}
+        for key, value in metadata.items():
+            result[key] = str(value)
+        return result
+
+    @staticmethod
+    def process_global_tags(
+        global_tags: GlobalTagsClass, dataset_urn: str
+    ) -> Iterable[MetadataWorkUnit]:
+        yield MetadataChangeProposalWrapper(
+            entityType="dataset",
+            entityUrn=dataset_urn,
+            aspect=global_tags,
+            changeType=ChangeTypeClass.UPSERT,
+        ).as_workunit()
+
+    def process_dataset(
+        self,
+        relative_path: str,
+        full_path: str,
+        filename: str,
+        table: ExcelTable,
+        source_type: UriType,
+    ) -> Iterable[MetadataWorkUnit]:
+        self.report.report_worksheet_processed()
+        dataset_name = gen_dataset_name(
+            relative_path, table.sheet_name, self.config.convert_urns_to_lowercase
+        )
+        dataset_urn = make_dataset_urn_with_platform_instance(
+            platform=self.platform,
+            name=dataset_name,
+            platform_instance=self.config.platform_instance,
+            env=self.config.env,
+        )
+
+        attributes = self.get_dataset_attributes(table.metadata)
+        created: Optional[datetime] = table.metadata.get("created")
+        modified: Optional[datetime] = table.metadata.get("modified")
+        dataset_properties = DatasetPropertiesClass(
+            tags=[],
+            customProperties=attributes,
+            created=(
+                TimeStamp(time=int(created.timestamp() * 1000)) if created else None
+            ),
+            lastModified=(
+                TimeStamp(time=int(modified.timestamp() * 1000)) if modified else None
+            ),
+        )
+
+        schema_metadata = self.construct_schema_metadata(
+            name=dataset_name,
+            dataset=table,
+        )
+
+        yield MetadataChangeProposalWrapper(
+            entityUrn=dataset_urn,
+            aspect=schema_metadata,
+        ).as_workunit()
+
+        yield MetadataChangeProposalWrapper(
+            entityUrn=dataset_urn,
+            aspect=dataset_properties,
+        ).as_workunit()
+
+        yield from self.container_WU_creator.create_container_hierarchy(
+            relative_path, dataset_urn
+        )
+
+        if source_type == UriType.S3 and (
+            self.config.use_s3_bucket_tags or self.config.use_s3_object_tags
+        ):
+            yield from self.process_s3_tags(full_path, dataset_urn)
+        elif source_type == UriType.ABS and self.config.use_abs_blob_tags:
+            yield from self.process_abs_tags(full_path, dataset_urn)
+
+        if self.config.is_profiling_enabled():
+            profiler = ExcelProfiler(
+                self.config,
+                self.report,
+                table.df,
+                filename,
+                table.sheet_name,
+                dataset_urn,
+                relative_path,
+            )
+            yield from profiler.get_workunits()
+
+    def process_file(
+        self,
+        file_content: BytesIO,
+        relative_path: str,
+        full_path: str,
+        filename: str,
+        source_type: UriType,
+    ) -> Iterable[MetadataWorkUnit]:
+        self.report.report_file_processed()
+        xls = ExcelFile(filename, file_content, self.report)
+        result = xls.load_workbook()
+
+        if result:
+            for table in xls.get_tables(active_only=self.config.active_sheet_only):
+                self.report.report_worksheet_scanned()
+                dataset_name = gen_dataset_name(
+                    relative_path,
+                    table.sheet_name,
+                    self.config.convert_urns_to_lowercase,
+                )
+                if not self.config.worksheet_pattern.allowed(dataset_name):
+                    self.report.report_dropped(dataset_name)
+                    continue
+                yield from self.process_dataset(
+                    relative_path, full_path, filename, table, source_type
+                )
+
+    def check_file_is_valid(self, filename: str) -> bool:
+        self.report.report_file_scanned()
+        if not self.config.path_pattern.allowed(filename):
+            self.report.report_dropped(filename)
+            return False
+        elif not self.is_excel_file(filename):
+            logger.debug(f"File is not an Excel workbook: {filename}")
+            return False
+        return True
+
+    def retrieve_file_data(
+        self, uri_type: UriType, path: str, path_spec: str
+    ) -> Iterator[Tuple[BytesIO, str, str, str]]:
+        if (
+            uri_type == UriType.LOCAL_FILE
+            or uri_type == UriType.ABSOLUTE_PATH
+            or uri_type == UriType.RELATIVE_PATH
+        ):
+            logger.debug(f"Searching local path: {path}")
+            for browse_path in self.local_browser(path):
+                if self.check_file_is_valid(browse_path.file):
+                    basename = os.path.basename(browse_path.file)
+                    file_path = self.strip_file_prefix(browse_path.file)
+                    filename = os.path.splitext(basename)[0]
+
+                    logger.debug(f"Processing {filename}")
+                    with self.report.local_file_get_timer:
+                        file_data = self.get_local_file(browse_path.file)
+
+                    if file_data is not None:
+                        yield file_data, file_path, browse_path.file, filename
+
+        elif uri_type == UriType.S3:
+            logger.debug(f"Searching S3 path: {path}")
+            for browse_path in self.s3_browser(path_spec):
+                if self.check_file_is_valid(browse_path.file):
+                    uri_path = strip_s3_prefix(browse_path.file)
+                    basename = os.path.basename(uri_path)
+                    filename = os.path.splitext(basename)[0]
+
+                    logger.debug(f"Processing {browse_path.file}")
+                    with self.report.s3_file_get_timer:
+                        file_data = self.get_s3_file(browse_path.file)
+
+                    if file_data is not None:
+                        yield file_data, uri_path, browse_path.file, filename
+
+        elif uri_type == UriType.ABS:
+            logger.debug(f"Searching Azure Blob Storage path: {path}")
+            for browse_path in self.abs_browser(path_spec):
+                if self.check_file_is_valid(browse_path.file):
+                    uri_path = strip_abs_prefix(browse_path.file)
+                    basename = os.path.basename(uri_path)
+                    filename = os.path.splitext(basename)[0]
+
+                    logger.debug(f"Processing {browse_path.file}")
+                    with self.report.abs_file_get_timer:
+                        file_data = self.get_abs_file(browse_path.file)
+
+                    if file_data is not None:
+                        yield file_data, uri_path, browse_path.file, filename
+
+        else:
+            self.report.report_file_dropped(path_spec)
+            self.report.warning(
+                message="Unsupported URI Type",
+                context=f"Type={uri_type.name},URI={path_spec}",
+            )
+
+    def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]:
+        self.container_WU_creator = ContainerWUCreator(
+            self.platform,
+            self.config.platform_instance,
+            self.config.env,
+        )
+
+        with PerfTimer() as timer:
+            for path_spec in self.config.path_list:
+                logger.debug(f"Processing path: {path_spec}")
+                uri_type, path = self.uri_type(path_spec)
+                logger.debug(f"URI Type: {uri_type} Path: {path}")
+
+                for (
+                    file_data,
+                    relative_path,
+                    full_path,
+                    filename,
+                ) in self.retrieve_file_data(uri_type, path, path_spec):
+                    yield from self.process_file(
+                        file_data, relative_path, full_path, filename, uri_type
+                    )
+
+        time_taken = timer.elapsed_seconds()
+        logger.info(f"Finished ingestion in {time_taken:.3f} seconds")
+
+    def get_report(self):
+        return self.report

datahub/ingestion/source/excel/util.py

@@ -0,0 +1,18 @@
+import os
+
+
+def gen_dataset_name(path: str, sheet_name: str, lower_case: bool) -> str:
+    sheet_name = sheet_name.strip()
+    directory, filename = os.path.split(path)
+
+    if not directory:
+        excel_path = f"[{filename}]"
+    else:
+        excel_path = os.path.join(directory, f"[{filename}]")
+
+    name = f"{excel_path}{sheet_name}"
+
+    if lower_case:
+        name = name.lower()
+
+    return name