openmetadata-ingestion 1.3.4.0__py3-none-any.whl → 1.4.0.0rc2__py3-none-any.whl

metadata/ingestion/source/pipeline/openlineage/models.py

@@ -0,0 +1,88 @@
+#  Copyright 2021 Collate
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  http://www.apache.org/licenses/LICENSE-2.0
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+"""
+Openlineage Source Model module
+"""
+
+from dataclasses import dataclass
+from enum import Enum
+from typing import Any, Dict, List
+
+
+@dataclass
+class OpenLineageEvent:
+    """
+    An object containing data extracted from raw OpenLineage event. Used as a basis for all abstract methods of
+    OpenlineageSource connector.
+    """
+
+    run_facet: Dict
+    job: Dict
+    event_type: str
+    inputs: List[Any]
+    outputs: List[Any]
+
+
+@dataclass
+class TableFQN:
+    """
+    Fully Qualified Name of a Table.
+    """
+
+    value: str
+
+
+@dataclass
+class ColumnFQN:
+    """
+    Fully Qualified Name of a Column.
+    """
+
+    value: str
+
+
+@dataclass
+class LineageNode:
+    """
+    A node being a part of Lineage information.
+    """
+
+    uuid: str
+    fqn: TableFQN
+    node_type: str = "table"
+
+
+@dataclass
+class LineageEdge:
+    """
+    An object describing connection of two nodes in the Lineage information.
+    """
+
+    from_node: LineageNode
+    to_node: LineageNode
+
+
+@dataclass
+class TableDetails:
+    """
+    Minimal table information.
+    """
+
+    schema: str
+    name: str
+
+
+class EventType(str, Enum):
+    """
+    List of used OpenLineage event types.
+    """
+
+    COMPLETE = "COMPLETE"

metadata/ingestion/source/pipeline/openlineage/utils.py

@@ -0,0 +1,59 @@
+"""
+Utils used by OpenlineageSource connector.
+"""
+
+from functools import reduce
+from typing import Dict
+
+from metadata.ingestion.source.pipeline.openlineage.models import OpenLineageEvent
+
+
+def message_to_open_lineage_event(incoming_event: Dict) -> OpenLineageEvent:
+    """
+    Method that takes raw Open Lineage event and parses is to shape into OpenLineageEvent.
+
+    We check whether received event (from Kafka) adheres to expected form and contains all the fields that are required
+    for successful processing by OpenMetadata OpenLineage connector.
+
+    :param incoming_event: raw event received from kafka topic by OpenlineageSource
+    :return: OpenLineageEvent
+    """
+    fields_to_verify = [
+        "run.facets.parent.job.name",
+        "run.facets.parent.job.namespace",
+        "inputs",
+        "outputs",
+        "eventType",
+        "job.name",
+        "job.namespace",
+    ]
+
+    for field in fields_to_verify:
+        try:
+            reduce(lambda x, y: x[y], field.split("."), incoming_event)
+        except KeyError:
+            raise ValueError("Event malformed!")
+
+    run_facet = incoming_event["run"]
+    inputs = incoming_event["inputs"]
+    outputs = incoming_event["outputs"]
+    event_type = incoming_event["eventType"]
+    job = incoming_event["job"]
+
+    result = OpenLineageEvent(
+        run_facet=run_facet,
+        event_type=event_type,
+        job=job,
+        inputs=inputs,
+        outputs=outputs,
+    )
+
+    return result
+
+
+class FQNNotFoundException(Exception):
+    """
+    Error raised when, while searching for an entity (Table, DatabaseSchema) there is no match in OM.
+    """
+
+    pass

metadata/ingestion/source/pipeline/pipeline_service.py

@@ -37,7 +37,7 @@ from metadata.ingestion.models.pipeline_status import OMetaPipelineStatus
 from metadata.ingestion.models.topology import (
     NodeStage,
     ServiceTopology,
-    TopologyContext,
+    TopologyContextManager,
     TopologyNode,
 )
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
@@ -116,7 +116,7 @@ class PipelineServiceSource(TopologyRunnerMixin, Source, ABC):
     service_connection: PipelineConnection.__fields__["config"].type_
 
     topology = PipelineServiceTopology()
-    context = TopologyContext.create(topology)
+    context = TopologyContextManager(topology)
     pipeline_source_state: Set = set()
 
     def __init__(
@@ -228,9 +228,19 @@ class PipelineServiceSource(TopologyRunnerMixin, Source, ABC):
                 entity_type=Pipeline,
                 entity_source_state=self.pipeline_source_state,
                 mark_deleted_entity=self.source_config.markDeletedPipelines,
-                params={"service": self.context.pipeline_service},
+                params={"service": self.context.get().pipeline_service},
             )
 
+    def get_db_service_names(self) -> List[str]:
+        """
+        Get the list of db service names
+        """
+        return (
+            self.source_config.lineageInformation.dbServiceNames or []
+            if self.source_config.lineageInformation
+            else []
+        )
+
     def prepare(self):
         """
         Method to implement any required logic before starting the ingestion process

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

@@ -56,7 +56,9 @@ class SplineSource(PipelineServiceSource):
     """
 
     @classmethod
-    def create(cls, config_dict, metadata: OpenMetadata):
+    def create(
+        cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
+    ):
         config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
         connection: SplineConnection = config.serviceConnection.__root__.config
         if not isinstance(connection, SplineConnection):
@@ -98,7 +100,7 @@ class SplineSource(PipelineServiceSource):
             displayName=pipeline_details.applicationName,
             sourceUrl=connection_url,
             tasks=self.get_connections_jobs(pipeline_details, connection_url),
-            service=self.context.pipeline_service,
+            service=self.context.get().pipeline_service,
         )
         yield Either(right=pipeline_request)
         self.register_record(pipeline_request=pipeline_request)
@@ -113,7 +115,7 @@ class SplineSource(PipelineServiceSource):
     ) -> Optional[Table]:
         if not table_name:
             return None
-        for service_name in self.source_config.dbServiceNames:
+        for service_name in self.get_db_service_names():
             table_fqn = fqn.build(
                 metadata=self.metadata,
                 entity_type=Table,
@@ -163,7 +165,7 @@ class SplineSource(PipelineServiceSource):
         """
         Parse all the executions available and create lineage
         """
-        if not self.source_config.dbServiceNames:
+        if not self.get_db_service_names():
             return
         lineage_details = self.client.get_lineage_details(
             pipeline_details.executionPlanId
@@ -222,8 +224,8 @@ class SplineSource(PipelineServiceSource):
                     pipeline_fqn = fqn.build(
                         metadata=self.metadata,
                         entity_type=Pipeline,
-                        service_name=self.context.pipeline_service,
-                        pipeline_name=self.context.pipeline,
+                        service_name=self.context.get().pipeline_service,
+                        pipeline_name=self.context.get().pipeline,
                     )
                     pipeline_entity = self.metadata.get_by_name(
                         entity=Pipeline, fqn=pipeline_fqn

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

@@ -56,7 +56,9 @@ class ElasticsearchSource(SearchServiceSource):
         self.client: Elasticsearch = self.connection
 
     @classmethod
-    def create(cls, config_dict, metadata: OpenMetadata):
+    def create(
+        cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
+    ):
         config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
         connection: ElasticsearchConnection = config.serviceConnection.__root__.config
         if not isinstance(connection, ElasticsearchConnection):
@@ -96,7 +98,7 @@ class ElasticsearchSource(SearchServiceSource):
                 searchIndexSettings=search_index_details.get(index_name, {}).get(
                     "settings", {}
                 ),
-                service=self.context.search_service,
+                service=self.context.get().search_service,
                 fields=parse_es_index_mapping(
                     search_index_details.get(index_name, {}).get("mappings")
                 ),
@@ -110,9 +112,9 @@ class ElasticsearchSource(SearchServiceSource):
         """
         Method to Get Sample Data of Search Index Entity
         """
-        if self.source_config.includeSampleData and self.context.search_index:
+        if self.source_config.includeSampleData and self.context.get().search_index:
             sample_data = self.client.search(
-                index=self.context.search_index,
+                index=self.context.get().search_index,
                 q=WILDCARD_SEARCH,
                 size=self.source_config.sampleSize,
                 request_timeout=self.service_connection.connectionTimeoutSecs,
@@ -121,8 +123,8 @@ class ElasticsearchSource(SearchServiceSource):
             search_index_fqn = fqn.build(
                 metadata=self.metadata,
                 entity_type=SearchIndex,
-                service_name=self.context.search_service,
-                search_index_name=self.context.search_index,
+                service_name=self.context.get().search_service,
+                search_index_name=self.context.get().search_index,
             )
             search_index_entity = self.metadata.get_by_name(
                 entity=SearchIndex, fqn=search_index_fqn

metadata/ingestion/source/search/search_service.py

@@ -43,7 +43,7 @@ from metadata.ingestion.models.search_index_data import OMetaIndexSampleData
 from metadata.ingestion.models.topology import (
     NodeStage,
     ServiceTopology,
-    TopologyContext,
+    TopologyContextManager,
     TopologyNode,
 )
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
@@ -110,7 +110,7 @@ class SearchServiceSource(TopologyRunnerMixin, Source, ABC):
     service_connection: SearchConnection.__fields__["config"].type_
 
     topology = SearchServiceTopology()
-    context = TopologyContext.create(topology)
+    context = TopologyContextManager(topology)
     index_source_state: Set = set()
 
     def __init__(
@@ -195,7 +195,7 @@ class SearchServiceSource(TopologyRunnerMixin, Source, ABC):
                 entity_type=SearchIndex,
                 entity_source_state=self.index_source_state,
                 mark_deleted_entity=self.source_config.markDeletedSearchIndexes,
-                params={"service": self.context.search_service},
+                params={"service": self.context.get().search_service},
             )
 
     def register_record(self, search_index_request: CreateSearchIndexRequest) -> None:

metadata/ingestion/source/sqa_types.py

@@ -13,7 +13,11 @@ Define custom types as wrappers on top of
 existing SQA types to have a bridge between
 SQA dialects and OM rich type system
 """
+
 from sqlalchemy import types
+from sqlalchemy.sql.sqltypes import TypeDecorator
+
+from metadata.utils.sqlalchemy_utils import convert_numpy_to_list
 
 
 class SQAMap(types.String):
@@ -22,11 +26,26 @@ class SQAMap(types.String):
     """
 
 
-class SQAStruct(types.String):
+class SQAStruct(TypeDecorator):
     """
     Custom Struct type definition
     """
 
+    impl = types.String
+    cache_ok = True
+
+    def process_result_value(self, value, dialect):
+        """This is executed during result retrieval
+
+        Args:
+            value: database record
+            dialect: database dialect
+        Returns:
+            python list conversion of ndarray
+        """
+
+        return convert_numpy_to_list(value)
+
 
 class SQADateTimeRange(types.String):
     """

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

@@ -59,8 +59,17 @@ def test_connection(
     of a metadata workflow or during an Automation Workflow
     """
 
+    def test_buckets(connection: S3Connection, client: S3ObjectStoreClient):
+        if connection.bucketNames:
+            for bucket_name in connection.bucketNames:
+                client.s3_client.list_objects(Bucket=bucket_name)
+            return
+        client.s3_client.list_buckets()
+
     test_fn = {
-        "ListBuckets": client.s3_client.list_buckets,
+        "ListBuckets": partial(
+            test_buckets, client=client, connection=service_connection
+        ),
         "GetMetrics": partial(
             client.cloudwatch_client.list_metrics, Namespace="AWS/S3"
         ),

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

@@ -83,13 +83,13 @@ class S3Source(StorageServiceSource):
         self.s3_reader = get_reader(config_source=S3Config(), client=self.s3_client)
 
     @classmethod
-    def create(cls, config_dict, metadata: OpenMetadata):
+    def create(
+        cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
+    ):
         config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
         connection: S3Connection = config.serviceConnection.__root__.config
         if not isinstance(connection, S3Connection):
-            raise InvalidSourceException(
-                f"Expected S3StoreConnection, but got {connection}"
-            )
+            raise InvalidSourceException(f"Expected S3Connection, but got {connection}")
         return cls(config, metadata)
 
     def get_containers(self) -> Iterable[S3ContainerDetails]:
@@ -103,7 +103,10 @@ class S3Source(StorageServiceSource):
                     bucket_response=bucket_response
                 )
                 container_fqn = fqn._build(  # pylint: disable=protected-access
-                    *(self.context.objectstore_service, self.context.container)
+                    *(
+                        self.context.get().objectstore_service,
+                        self.context.get().container,
+                    )
                 )
                 container_entity = self.metadata.get_by_name(
                     entity=Container, fqn=container_fqn
@@ -123,10 +126,7 @@ class S3Source(StorageServiceSource):
                         # ingest all the relevant valid paths from it
                         yield from self._generate_structured_containers(
                             bucket_response=bucket_response,
-                            entries=self._manifest_entries_to_metadata_entries_by_container(
-                                container_name=bucket_name,
-                                manifest=self.global_manifest,
-                            ),
+                            entries=manifest_entries_for_current_bucket,
                             parent=parent_entity,
                         )
                         # nothing else do to for the current bucket, skipping to the next
@@ -139,23 +139,6 @@ class S3Source(StorageServiceSource):
                         entries=metadata_config.entries,
                         parent=parent_entity,
                     )
-                    for metadata_entry in metadata_config.entries:
-                        logger.info(
-                            f"Extracting metadata from path {metadata_entry.dataPath.strip(KEY_SEPARATOR)} "
-                            f"and generating structured container"
-                        )
-                        structured_container: Optional[
-                            S3ContainerDetails
-                        ] = self._generate_container_details(
-                            bucket_response=bucket_response,
-                            metadata_entry=metadata_entry,
-                            parent=EntityReference(
-                                id=self._bucket_cache[bucket_response.name].id.__root__,
-                                type="container",
-                            ),
-                        )
-                        if structured_container:
-                            yield structured_container
 
             except ValidationError as err:
                 self.status.failed(
@@ -183,10 +166,11 @@ class S3Source(StorageServiceSource):
             numberOfObjects=container_details.number_of_objects,
             size=container_details.size,
             dataModel=container_details.data_model,
-            service=self.context.objectstore_service,
+            service=self.context.get().objectstore_service,
             parent=container_details.parent,
             sourceUrl=container_details.sourceUrl,
             fileFormats=container_details.file_formats,
+            fullPath=container_details.fullPath,
         )
         yield Either(right=container_request)
         self.register_record(container_request=container_request)
@@ -213,10 +197,15 @@ class S3Source(StorageServiceSource):
                 client=self.s3_client,
             )
             if columns:
+                prefix = (
+                    f"{KEY_SEPARATOR}{metadata_entry.dataPath.strip(KEY_SEPARATOR)}"
+                )
                 return S3ContainerDetails(
                     name=metadata_entry.dataPath.strip(KEY_SEPARATOR),
-                    prefix=f"{KEY_SEPARATOR}{metadata_entry.dataPath.strip(KEY_SEPARATOR)}",
-                    creation_date=bucket_response.creation_date.isoformat(),
+                    prefix=prefix,
+                    creation_date=bucket_response.creation_date.isoformat()
+                    if bucket_response.creation_date
+                    else None,
                     number_of_objects=self._fetch_metric(
                         bucket_name=bucket_name, metric=S3Metric.NUMBER_OF_OBJECTS
                     ),
@@ -228,6 +217,7 @@ class S3Source(StorageServiceSource):
                         isPartitioned=metadata_entry.isPartitioned, columns=columns
                     ),
                     parent=parent,
+                    fullPath=self._get_full_path(bucket_name, prefix),
                     sourceUrl=self._get_object_source_url(
                         bucket_name=bucket_name,
                         prefix=metadata_entry.dataPath.strip(KEY_SEPARATOR),
@@ -262,6 +252,11 @@ class S3Source(StorageServiceSource):
     def fetch_buckets(self) -> List[S3BucketResponse]:
         results: List[S3BucketResponse] = []
         try:
+            if self.service_connection.bucketNames:
+                return [
+                    S3BucketResponse(Name=bucket_name)
+                    for bucket_name in self.service_connection.bucketNames
+                ]
             # No pagination required, as there is a hard 1000 limit on nr of buckets per aws account
             for bucket in self.s3_client.list_buckets().get("Buckets") or []:
                 if filter_by_container(
@@ -327,7 +322,9 @@ class S3Source(StorageServiceSource):
         return S3ContainerDetails(
             name=bucket_response.name,
             prefix=KEY_SEPARATOR,
-            creation_date=bucket_response.creation_date.isoformat(),
+            creation_date=bucket_response.creation_date.isoformat()
+            if bucket_response.creation_date
+            else None,
             number_of_objects=self._fetch_metric(
                 bucket_name=bucket_response.name, metric=S3Metric.NUMBER_OF_OBJECTS
             ),
@@ -336,9 +333,27 @@ class S3Source(StorageServiceSource):
             ),
             file_formats=[],
             data_model=None,
+            fullPath=self._get_full_path(bucket_name=bucket_response.name),
             sourceUrl=self._get_bucket_source_url(bucket_name=bucket_response.name),
         )
 
+    def _clean_path(self, path: str) -> str:
+        return path.strip(KEY_SEPARATOR)
+
+    def _get_full_path(self, bucket_name: str, prefix: str = None) -> Optional[str]:
+        """
+        Method to get the full path of the file
+        """
+        if bucket_name is None:
+            return None
+
+        full_path = f"s3://{self._clean_path(bucket_name)}"
+
+        if prefix:
+            full_path += f"/{self._clean_path(prefix)}"
+
+        return full_path
+
     def _get_sample_file_path(
         self, bucket_name: str, metadata_entry: MetadataEntry
     ) -> Optional[str]:

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

@@ -33,8 +33,8 @@ class S3BucketResponse(BaseModel):
         extra = Extra.forbid
 
     name: str = Field(..., description="Bucket name", alias="Name")
-    creation_date: datetime = Field(
-        ...,
+    creation_date: Optional[datetime] = Field(
+        None,
         description="Timestamp of Bucket creation in ISO format",
         alias="CreationDate",
     )
@@ -67,8 +67,8 @@ class S3ContainerDetails(BaseModel):
         ...,
         description="Data Model of the container",
     )
-    creation_date: str = Field(
-        ...,
+    creation_date: Optional[str] = Field(
+        None,
         description="Timestamp of Bucket creation in ISO format",
     )
     parent: Optional[EntityReference] = Field(
@@ -78,3 +78,7 @@ class S3ContainerDetails(BaseModel):
     sourceUrl: Optional[basic.SourceUrl] = Field(
         None, description="Source URL of the container."
     )
+
+    fullPath: Optional[str] = Field(
+        None, description="Full path of the container/file."
+    )

metadata/ingestion/source/storage/storage_service.py

@@ -41,7 +41,7 @@ from metadata.ingestion.models.delete_entity import DeleteEntity
 from metadata.ingestion.models.topology import (
     NodeStage,
     ServiceTopology,
-    TopologyContext,
+    TopologyContextManager,
     TopologyNode,
 )
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
@@ -112,7 +112,7 @@ class StorageServiceSource(TopologyRunnerMixin, Source, ABC):
     service_connection: StorageConnection.__fields__["config"].type_
 
     topology = StorageServiceTopology()
-    context = TopologyContext.create(topology)
+    context = TopologyContextManager(topology)
     container_source_state: Set = set()
 
     global_manifest: Optional[ManifestMetadataConfig]
@@ -191,7 +191,7 @@ class StorageServiceSource(TopologyRunnerMixin, Source, ABC):
         container_fqn = fqn.build(
             self.metadata,
             entity_type=Container,
-            service_name=self.context.objectstore_service,
+            service_name=self.context.get().objectstore_service,
             parent_container=parent_container,
             container_name=container_request.name.__root__,
         )
@@ -210,7 +210,7 @@ class StorageServiceSource(TopologyRunnerMixin, Source, ABC):
                 entity_type=Container,
                 entity_source_state=self.container_source_state,
                 mark_deleted_entity=self.source_config.markDeletedContainers,
-                params={"service": self.context.objectstore_service},
+                params={"service": self.context.get().objectstore_service},
             )
 
     def yield_create_request_objectstore_service(self, config: WorkflowSource):

metadata/ingestion/stage/table_usage.py

@@ -70,7 +70,12 @@ class TableUsageStage(Stage):
         return "Table Usage"
 
     @classmethod
-    def create(cls, config_dict: dict, metadata: OpenMetadata):
+    def create(
+        cls,
+        config_dict: dict,
+        metadata: OpenMetadata,
+        pipeline_name: Optional[str] = None,
+    ):
         config = TableStageConfig.parse_obj(config_dict)
         return cls(config, metadata)
 

metadata/mixins/pandas/pandas_mixin.py

@@ -21,7 +21,7 @@ from metadata.data_quality.validations.table.pandas.tableRowInsertedCountToBeBet
     TableRowInsertedCountToBeBetweenValidator,
 )
 from metadata.generated.schema.entity.data.table import (
-    PartitionIntervalType,
+    PartitionIntervalTypes,
     PartitionProfilerConfig,
     ProfileSampleType,
 )
@@ -47,7 +47,7 @@ class PandasInterfaceMixin:
         partition_field = self.table_partition_config.partitionColumnName
         if (
             self.table_partition_config.partitionIntervalType
-            == PartitionIntervalType.COLUMN_VALUE
+            == PartitionIntervalTypes.COLUMN_VALUE
         ):
             return [
                 df[
@@ -59,7 +59,7 @@ class PandasInterfaceMixin:
             ]
         if (
             self.table_partition_config.partitionIntervalType
-            == PartitionIntervalType.INTEGER_RANGE
+            == PartitionIntervalTypes.INTEGER_RANGE
         ):
             return [
                 df[

metadata/pii/processor.py

@@ -81,7 +81,12 @@ class PIIProcessor(Processor):
         return self._ner_scanner
 
     @classmethod
-    def create(cls, config_dict: dict, metadata: OpenMetadata) -> "Step":
+    def create(
+        cls,
+        config_dict: dict,
+        metadata: OpenMetadata,
+        pipeline_name: Optional[str] = None,
+    ) -> "Step":
         config = parse_workflow_config_gracefully(config_dict)
         return cls(config=config, metadata=metadata)