openmetadata-ingestion 1.2.4.0__py3-none-any.whl → 1.3.2.0rc1__py3-none-any.whl

metadata/ingestion/source/pipeline/pipeline_service.py

@@ -37,8 +37,8 @@ from metadata.ingestion.models.pipeline_status import OMetaPipelineStatus
 from metadata.ingestion.models.topology import (
     NodeStage,
     ServiceTopology,
+    TopologyContext,
     TopologyNode,
-    create_source_context,
 )
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.connections import get_connection, get_test_connection_fn
@@ -116,7 +116,7 @@ class PipelineServiceSource(TopologyRunnerMixin, Source, ABC):
     service_connection: PipelineConnection.__fields__["config"].type_
 
     topology = PipelineServiceTopology()
-    context = create_source_context(topology)
+    context = TopologyContext.create(topology)
     pipeline_source_state: Set = set()
 
     def __init__(
@@ -138,6 +138,10 @@ class PipelineServiceSource(TopologyRunnerMixin, Source, ABC):
         self.client = self.connection
         self.test_connection()
 
+    @property
+    def name(self) -> str:
+        return self.service_connection.type.name
+
     @abstractmethod
     def yield_pipeline(
         self, pipeline_details: Any

metadata/ingestion/source/pipeline/spline/metadata.py

@@ -131,7 +131,6 @@ class SplineSource(PipelineServiceSource):
         return None
 
     def _get_table_from_datasource_name(self, datasource: str) -> Optional[Table]:
-
         if (
             not datasource
             and not datasource.startswith("dbfs")

metadata/ingestion/source/search/elasticsearch/connection.py

@@ -191,9 +191,12 @@ def test_connection(
     of a metadata workflow or during an Automation Workflow
     """
 
+    def test_get_search_indexes():
+        client.indices.get_alias(expand_wildcards="open")
+
     test_fn = {
         "CheckAccess": client.info,
-        "GetSearchIndexes": client.indices.get_alias,
+        "GetSearchIndexes": test_get_search_indexes,
     }
 
     test_connection_steps(

metadata/ingestion/source/search/elasticsearch/metadata.py

@@ -69,7 +69,7 @@ class ElasticsearchSource(SearchServiceSource):
         """
         Get List of all search index
         """
-        index_list = self.client.indices.get_alias() or {}
+        index_list = self.client.indices.get_alias(expand_wildcards="open") or {}
         for index in index_list.keys():
             yield self.client.indices.get(index=str(index))
 
@@ -111,7 +111,6 @@ class ElasticsearchSource(SearchServiceSource):
         Method to Get Sample Data of Search Index Entity
         """
         if self.source_config.includeSampleData and self.context.search_index:
-
             sample_data = self.client.search(
                 index=self.context.search_index,
                 q=WILDCARD_SEARCH,

metadata/ingestion/source/search/search_service.py

@@ -43,8 +43,8 @@ from metadata.ingestion.models.search_index_data import OMetaIndexSampleData
 from metadata.ingestion.models.topology import (
     NodeStage,
     ServiceTopology,
+    TopologyContext,
     TopologyNode,
-    create_source_context,
 )
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.connections import get_connection, get_test_connection_fn
@@ -110,7 +110,7 @@ class SearchServiceSource(TopologyRunnerMixin, Source, ABC):
     service_connection: SearchConnection.__fields__["config"].type_
 
     topology = SearchServiceTopology()
-    context = create_source_context(topology)
+    context = TopologyContext.create(topology)
     index_source_state: Set = set()
 
     def __init__(
@@ -131,6 +131,10 @@ class SearchServiceSource(TopologyRunnerMixin, Source, ABC):
         self.connection_obj = self.connection
         self.test_connection()
 
+    @property
+    def name(self) -> str:
+        return self.service_connection.type.name
+
     @abstractmethod
     def yield_search_index(
         self, search_index_details: Any

metadata/ingestion/source/storage/s3/metadata.py

@@ -30,6 +30,9 @@ from metadata.generated.schema.entity.services.connections.database.datalake.s3C
 from metadata.generated.schema.entity.services.connections.storage.s3Connection import (
     S3Connection,
 )
+from metadata.generated.schema.entity.services.ingestionPipelines.status import (
+    StackTraceError,
+)
 from metadata.generated.schema.metadataIngestion.storage.containerMetadataConfig import (
     MetadataEntry,
     StorageContainerConfig,
@@ -38,7 +41,7 @@ from metadata.generated.schema.metadataIngestion.workflow import (
     Source as WorkflowSource,
 )
 from metadata.generated.schema.type.entityReference import EntityReference
-from metadata.ingestion.api.models import Either, StackTraceError
+from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import InvalidSourceException
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.storage.s3.models import (
@@ -159,7 +162,7 @@ class S3Source(StorageServiceSource):
                     StackTraceError(
                         name=bucket_response.name,
                         error=f"Validation error while creating Container from bucket details - {err}",
-                        stack_trace=traceback.format_exc(),
+                        stackTrace=traceback.format_exc(),
                     )
                 )
             except Exception as err:
@@ -167,26 +170,26 @@ class S3Source(StorageServiceSource):
                     StackTraceError(
                         name=bucket_response.name,
                         error=f"Wild error while creating Container from bucket details - {err}",
-                        stack_trace=traceback.format_exc(),
+                        stackTrace=traceback.format_exc(),
                     )
                 )
 
     def yield_create_container_requests(
         self, container_details: S3ContainerDetails
     ) -> Iterable[Either[CreateContainerRequest]]:
-        yield Either(
-            right=CreateContainerRequest(
-                name=container_details.name,
-                prefix=container_details.prefix,
-                numberOfObjects=container_details.number_of_objects,
-                size=container_details.size,
-                dataModel=container_details.data_model,
-                service=self.context.objectstore_service,
-                parent=container_details.parent,
-                sourceUrl=container_details.sourceUrl,
-                fileFormats=container_details.file_formats,
-            )
+        container_request = CreateContainerRequest(
+            name=container_details.name,
+            prefix=container_details.prefix,
+            numberOfObjects=container_details.number_of_objects,
+            size=container_details.size,
+            dataModel=container_details.data_model,
+            service=self.context.objectstore_service,
+            parent=container_details.parent,
+            sourceUrl=container_details.sourceUrl,
+            fileFormats=container_details.file_formats,
         )
+        yield Either(right=container_request)
+        self.register_record(container_request=container_request)
 
     def _generate_container_details(
         self,
@@ -424,7 +427,7 @@ class S3Source(StorageServiceSource):
         """
         try:
             logger.info(
-                f"Found metadata template file at - s3://{bucket_name}/{OPENMETADATA_TEMPLATE_FILE_NAME}"
+                f"Looking for metadata template file at - s3://{bucket_name}/{OPENMETADATA_TEMPLATE_FILE_NAME}"
             )
             response_object = self.s3_reader.read(
                 path=OPENMETADATA_TEMPLATE_FILE_NAME,
@@ -435,7 +438,9 @@ class S3Source(StorageServiceSource):
             metadata_config = StorageContainerConfig.parse_obj(content)
             return metadata_config
         except ReadException:
-            pass
+            logger.warning(
+                f"No metadata file found at s3://{bucket_name}/{OPENMETADATA_TEMPLATE_FILE_NAME}"
+            )
         except Exception as exc:
             logger.debug(traceback.format_exc())
             logger.warning(

metadata/ingestion/source/storage/storage_service.py

@@ -12,9 +12,7 @@
 Base class for ingesting Object Storage services
 """
 from abc import ABC, abstractmethod
-from typing import Any, Iterable, List, Optional
-
-from pandas import DataFrame
+from typing import Any, Iterable, List, Optional, Set
 
 from metadata.generated.schema.api.data.createContainer import CreateContainerRequest
 from metadata.generated.schema.entity.data.container import Container
@@ -35,14 +33,16 @@ from metadata.generated.schema.metadataIngestion.storageServiceMetadataPipeline
 from metadata.generated.schema.metadataIngestion.workflow import (
     Source as WorkflowSource,
 )
+from metadata.ingestion.api.delete import delete_entity_from_source
 from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import Source
 from metadata.ingestion.api.topology_runner import TopologyRunnerMixin
+from metadata.ingestion.models.delete_entity import DeleteEntity
 from metadata.ingestion.models.topology import (
     NodeStage,
     ServiceTopology,
+    TopologyContext,
     TopologyNode,
-    create_source_context,
 )
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.connections import get_connection, get_test_connection_fn
@@ -50,7 +50,11 @@ from metadata.ingestion.source.database.glue.models import Column
 from metadata.readers.dataframe.models import DatalakeTableSchemaWrapper
 from metadata.readers.dataframe.reader_factory import SupportedTypes
 from metadata.readers.models import ConfigSource
-from metadata.utils.datalake.datalake_utils import fetch_dataframe, get_columns
+from metadata.utils import fqn
+from metadata.utils.datalake.datalake_utils import (
+    DataFrameColumnParser,
+    fetch_dataframe,
+)
 from metadata.utils.logger import ingestion_logger
 from metadata.utils.storage_metadata_config import (
     StorageMetadataConfigException,
@@ -64,7 +68,6 @@ OPENMETADATA_TEMPLATE_FILE_NAME = "openmetadata.json"
 
 
 class StorageServiceTopology(ServiceTopology):
-
     root = TopologyNode(
         producer="get_services",
         stages=[
@@ -78,6 +81,7 @@ class StorageServiceTopology(ServiceTopology):
             ),
         ],
         children=["container"],
+        post_process=["mark_containers_as_deleted"],
     )
 
     container = TopologyNode(
@@ -108,7 +112,8 @@ class StorageServiceSource(TopologyRunnerMixin, Source, ABC):
     service_connection: StorageConnection.__fields__["config"].type_
 
     topology = StorageServiceTopology()
-    context = create_source_context(topology)
+    context = TopologyContext.create(topology)
+    container_source_state: Set = set()
 
     global_manifest: Optional[ManifestMetadataConfig]
 
@@ -135,6 +140,10 @@ class StorageServiceSource(TopologyRunnerMixin, Source, ABC):
             ManifestMetadataConfig
         ] = self.get_manifest_file()
 
+    @property
+    def name(self) -> str:
+        return self.service_connection.type.name
+
     def get_manifest_file(self) -> Optional[ManifestMetadataConfig]:
         if self.source_config.storageMetadataConfigSource and not isinstance(
             self.source_config.storageMetadataConfigSource,
@@ -167,10 +176,43 @@ class StorageServiceSource(TopologyRunnerMixin, Source, ABC):
     def prepare(self):
         """By default, nothing needs to be taken care of when loading the source"""
 
+    def register_record(self, container_request: CreateContainerRequest) -> None:
+        """
+        Mark the container record as scanned and update
+        the storage_source_state
+        """
+        parent_container = (
+            self.metadata.get_by_id(
+                entity=Container, entity_id=container_request.parent.id
+            ).fullyQualifiedName.__root__
+            if container_request.parent
+            else None
+        )
+        container_fqn = fqn.build(
+            self.metadata,
+            entity_type=Container,
+            service_name=self.context.objectstore_service,
+            parent_container=parent_container,
+            container_name=container_request.name.__root__,
+        )
+
+        self.container_source_state.add(container_fqn)
+
     def test_connection(self) -> None:
         test_connection_fn = get_test_connection_fn(self.service_connection)
         test_connection_fn(self.metadata, self.connection_obj, self.service_connection)
 
+    def mark_containers_as_deleted(self) -> Iterable[Either[DeleteEntity]]:
+        """Method to mark the containers as deleted"""
+        if self.source_config.markDeletedContainers:
+            yield from delete_entity_from_source(
+                metadata=self.metadata,
+                entity_type=Container,
+                entity_source_state=self.container_source_state,
+                mark_deleted_entity=self.source_config.markDeletedContainers,
+                params={"service": self.context.objectstore_service},
+            )
+
     def yield_create_request_objectstore_service(self, config: WorkflowSource):
         yield Either(
             right=self.metadata.get_create_service_from_source(
@@ -229,10 +271,10 @@ class StorageServiceSource(TopologyRunnerMixin, Source, ABC):
             ),
         )
         columns = []
-        if isinstance(data_structure_details, DataFrame):
-            columns = get_columns(data_structure_details)
-        if isinstance(data_structure_details, list) and data_structure_details:
-            columns = get_columns(data_structure_details[0])
+        column_parser = DataFrameColumnParser.create(
+            data_structure_details, SupportedTypes(metadata_entry.structureFormat)
+        )
+        columns = column_parser.get_columns()
         return columns
 
     def _get_columns(

metadata/ingestion/stage/table_usage.py

@@ -22,10 +22,13 @@ from typing import Iterable, List, Optional, Tuple
 
 from metadata.config.common import ConfigModel
 from metadata.generated.schema.api.data.createQuery import CreateQueryRequest
+from metadata.generated.schema.entity.services.ingestionPipelines.status import (
+    StackTraceError,
+)
 from metadata.generated.schema.entity.teams.user import User
 from metadata.generated.schema.type.queryParserData import ParsedData, QueryParserData
 from metadata.generated.schema.type.tableUsageCount import TableUsageCount
-from metadata.ingestion.api.models import Either, StackTraceError
+from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import Stage
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.utils.constants import UTF_8
@@ -62,6 +65,10 @@ class TableUsageStage(Stage):
         init_staging_dir(self.config.filename)
         self.wrote_something = False
 
+    @property
+    def name(self) -> str:
+        return "Table Usage"
+
     @classmethod
     def create(cls, config_dict: dict, metadata: OpenMetadata):
         config = TableStageConfig.parse_obj(config_dict)
@@ -153,7 +160,7 @@ class TableUsageStage(Stage):
                 left=StackTraceError(
                     name=table,
                     error=f"Error in staging record [{exc}]",
-                    stack_trace=traceback.format_exc(),
+                    stackTrace=traceback.format_exc(),
                 )
             )
         yield Either(right=table)

metadata/mixins/sqalchemy/sqa_mixin.py

@@ -26,6 +26,9 @@ from metadata.generated.schema.entity.services.connections.database.databricksCo
 from metadata.generated.schema.entity.services.connections.database.snowflakeConnection import (
     SnowflakeType,
 )
+from metadata.generated.schema.entity.services.connections.database.unityCatalogConnection import (
+    UnityCatalogConnection,
+)
 from metadata.ingestion.source.connections import get_connection
 from metadata.ingestion.source.database.snowflake.queries import (
     SNOWFLAKE_SESSION_TAG_QUERY,
@@ -84,17 +87,21 @@ class SQAInterfaceMixin:
             )
 
     def set_catalog(self, session) -> None:
-        """Set catalog for the session. Right now only databricks requires it
+        """Set catalog for the session. Right now only databricks and unity catalog requires it
 
         Args:
             session (Session): sqa session object
         """
-        if isinstance(self.service_connection_config, DatabricksConnection):
-            bind = session.get_bind()
-            bind.execute(
-                "USE CATALOG %(catalog)s;",
-                {"catalog": self.service_connection_config.catalog},
-            ).first()
+        if not isinstance(
+            self.service_connection_config,
+            (UnityCatalogConnection, DatabricksConnection),
+        ):
+            return
+        bind = session.get_bind()
+        bind.execute(
+            "USE CATALOG %(catalog)s;",
+            {"catalog": self.service_connection_config.catalog},
+        ).first()
 
     def close(self):
         """close session"""

metadata/parsers/protobuf_parser.py

@@ -20,12 +20,14 @@ import sys
 import traceback
 from enum import Enum
 from pathlib import Path
-from typing import List, Optional
+from typing import List, Optional, Union
 
 import grpc_tools.protoc
 from pydantic import BaseModel
+from pydantic.main import ModelMetaclass
 
-from metadata.generated.schema.type.schema import FieldModel
+from metadata.generated.schema.entity.data.table import Column, DataType
+from metadata.generated.schema.type.schema import DataTypeTopic, FieldModel
 from metadata.utils.helpers import snake_to_camel
 from metadata.utils.logger import ingestion_logger
 
@@ -37,12 +39,12 @@ class ProtobufDataTypes(Enum):
     Enum for Protobuf Datatypes
     """
 
-    ERROR = 0
+    UNKNOWN = 0
     DOUBLE = 1
     FLOAT = 2
     INT = 3, 4, 5, 13, 17, 18
     FIXED = 6, 7, 15, 16
-    TYPE_BOOL = 8
+    BOOLEAN = 8
     STRING = 9
     UNION = 10
     RECORD = 11
@@ -165,7 +167,9 @@ class ProtobufParser:
             )
         return None
 
-    def parse_protobuf_schema(self) -> Optional[List[FieldModel]]:
+    def parse_protobuf_schema(
+        self, cls: ModelMetaclass = FieldModel
+    ) -> Optional[List[Union[FieldModel, Column]]]:
         """
         Method to parse the protobuf schema
         """
@@ -177,10 +181,12 @@ class ProtobufParser:
             )
 
             field_models = [
-                FieldModel(
+                cls(
                     name=instance.DESCRIPTOR.name,
                     dataType="RECORD",
-                    children=self.get_protobuf_fields(instance.DESCRIPTOR.fields),
+                    children=self.get_protobuf_fields(
+                        instance.DESCRIPTOR.fields, cls=cls
+                    ),
                 )
             ]
 
@@ -196,7 +202,17 @@ class ProtobufParser:
             )
         return None
 
-    def get_protobuf_fields(self, fields) -> Optional[List[FieldModel]]:
+    def _get_field_type(self, type_: int, cls: ModelMetaclass = FieldModel) -> str:
+        if type_ > 18:
+            return DataType.UNKNOWN.value
+        data_type = ProtobufDataTypes(type_).name
+        if cls == Column and data_type == DataTypeTopic.FIXED.value:
+            return DataType.INT.value
+        return data_type
+
+    def get_protobuf_fields(
+        self, fields, cls: ModelMetaclass = FieldModel
+    ) -> Optional[List[Union[FieldModel, Column]]]:
         """
         Recursively convert the parsed schema into required models
         """
@@ -205,10 +221,12 @@ class ProtobufParser:
         for field in fields:
             try:
                 field_models.append(
-                    FieldModel(
+                    cls(
                         name=field.name,
-                        dataType=ProtobufDataTypes(field.type).name,
-                        children=self.get_protobuf_fields(field.message_type.fields)
+                        dataType=self._get_field_type(field.type, cls=cls),
+                        children=self.get_protobuf_fields(
+                            field.message_type.fields, cls=cls
+                        )
                         if field.type == 11
                         else None,
                     )

metadata/pii/processor.py

@@ -16,6 +16,9 @@ import traceback
 from typing import List, Optional, cast
 
 from metadata.generated.schema.entity.data.table import Column, TableData
+from metadata.generated.schema.entity.services.ingestionPipelines.status import (
+    StackTraceError,
+)
 from metadata.generated.schema.metadataIngestion.databaseServiceProfilerPipeline import (
     DatabaseServiceProfilerPipeline,
 )
@@ -28,7 +31,7 @@ from metadata.generated.schema.type.tagLabel import (
     TagLabel,
     TagSource,
 )
-from metadata.ingestion.api.models import Either, StackTraceError
+from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.parser import parse_workflow_config_gracefully
 from metadata.ingestion.api.step import Step
 from metadata.ingestion.api.steps import Processor
@@ -65,6 +68,10 @@ class PIIProcessor(Processor):
         self._ner_scanner = None
         self.confidence_threshold = self.source_config.confidence
 
+    @property
+    def name(self) -> str:
+        return "PII Processor"
+
     @property
     def ner_scanner(self) -> NERScanner:
         """Load the NER Scanner only if called"""
@@ -170,7 +177,7 @@ class PIIProcessor(Processor):
                     StackTraceError(
                         name=record.table.fullyQualifiedName.__root__,
                         error=f"Error computing PII tags for [{column}] - [{err}]",
-                        stack_trace=traceback.format_exc(),
+                        stackTrace=traceback.format_exc(),
                     )
                 )
 

metadata/pii/scanners/ner_scanner.py

@@ -21,7 +21,7 @@ from pydantic import BaseModel
 
 from metadata.generated.schema.entity.classification.tag import Tag
 from metadata.pii.constants import PII, SPACY_EN_MODEL
-from metadata.pii.models import TagAndConfidence, TagType
+from metadata.pii.models import TagAndConfidence
 from metadata.pii.ner import NEREntity
 from metadata.utils import fqn
 from metadata.utils.logger import pii_logger
@@ -119,13 +119,15 @@ class NERScanner:
 
         if entities_score:
             label, score = self.get_highest_score_label(entities_score)
-            tag_type = NEREntity.__members__.get(label, TagType.NONSENSITIVE).value
+            tag_type = NEREntity.__members__.get(label)
+            if not tag_type:
+                return None
             return TagAndConfidence(
                 tag_fqn=fqn.build(
                     metadata=None,
                     entity_type=Tag,
                     classification_name=PII,
-                    tag_name=tag_type,
+                    tag_name=tag_type.value,
                 ),
                 confidence=score,
             )

metadata/profiler/api/models.py

@@ -15,7 +15,10 @@ Return types for Profiler workflow execution.
 We need to define this class as we end up having
 multiple profilers per table and columns.
 """
-from typing import List, Optional, Union
+from typing import List, Optional, Type, Union
+
+from sqlalchemy import Column
+from sqlalchemy.orm import DeclarativeMeta
 
 from metadata.config.common import ConfigModel
 from metadata.generated.schema.api.data.createTableProfile import (
@@ -31,9 +34,12 @@ from metadata.generated.schema.entity.data.table import (
 from metadata.generated.schema.entity.services.connections.connectionBasicType import (
     SampleDataStorageConfig,
 )
+from metadata.generated.schema.tests.customMetric import CustomMetric
 from metadata.generated.schema.type.basic import FullyQualifiedEntityName
 from metadata.ingestion.models.table_metadata import ColumnTag
+from metadata.profiler.metrics.core import Metric, MetricTypes
 from metadata.profiler.processor.models import ProfilerDef
+from metadata.utils.sqa_like_column import SQALikeColumn
 
 
 class ColumnConfig(ConfigModel):
@@ -113,3 +119,15 @@ class ProfilerResponse(ConfigModel):
     def __str__(self):
         """Return the table name being processed"""
         return f"Table [{self.table.name.__root__}]"
+
+
+class ThreadPoolMetrics(ConfigModel):
+    """thread pool metric"""
+
+    metrics: Union[List[Union[Type[Metric], CustomMetric]], Type[Metric]]
+    metric_type: MetricTypes
+    column: Optional[Union[Column, SQALikeColumn]]
+    table: Union[Table, DeclarativeMeta]
+
+    class Config:
+        arbitrary_types_allowed = True