openmetadata-ingestion 1.3.3.0rc2__py3-none-any.whl → 1.4.0.0rc1__py3-none-any.whl

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/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
@@ -183,10 +186,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,9 +217,12 @@ 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)}",
+                    prefix=prefix,
                     creation_date=bucket_response.creation_date.isoformat(),
                     number_of_objects=self._fetch_metric(
                         bucket_name=bucket_name, metric=S3Metric.NUMBER_OF_OBJECTS
@@ -228,6 +235,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),
@@ -336,9 +344,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

@@ -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)
 

metadata/profiler/adaptors/adaptor_factory.py

@@ -0,0 +1,43 @@
+#  Copyright 2024 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.
+"""
+factory for NoSQL adaptors that are used in the NoSQLProfiler.
+"""
+from metadata.generated.schema.entity.services.connections.database.dynamoDBConnection import (
+    DynamoDBConnection,
+)
+from metadata.generated.schema.entity.services.connections.database.mongoDBConnection import (
+    MongoDBConnection,
+)
+from metadata.profiler.adaptors.dynamodb import DynamoDB
+from metadata.profiler.adaptors.mongodb import MongoDB
+from metadata.profiler.factory import Factory
+from metadata.utils.logger import profiler_logger
+
+logger = profiler_logger()
+
+
+class NoSQLAdaptorFactory(Factory):
+    def create(self, interface_type: str, *args, **kwargs) -> any:
+        logger.debug(f"Creating NoSQL client for {interface_type}")
+        client_class = self._interface_type.get(interface_type)
+        if not client_class:
+            raise ValueError(f"Unknown NoSQL source: {interface_type}")
+        logger.debug(f"Using NoSQL client constructor: {client_class.__name__}")
+        return client_class(*args, **kwargs)
+
+
+adaptors = profilers = {
+    MongoDBConnection.__name__: MongoDB,
+    DynamoDBConnection.__name__: DynamoDB,
+}
+factory = NoSQLAdaptorFactory()
+factory.register_many(adaptors)

metadata/profiler/adaptors/dynamodb.py

@@ -0,0 +1,40 @@
+#  Copyright 2024 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.
+"""
+DyanmoDB adaptor for the NoSQL profiler.
+"""
+from typing import TYPE_CHECKING, Dict, List
+
+from metadata.generated.schema.entity.data.table import Column, Table
+from metadata.profiler.adaptors.nosql_adaptor import NoSQLAdaptor
+
+if TYPE_CHECKING:
+    from mypy_boto3_dynamodb.service_resource import DynamoDBServiceResource
+else:
+    DynamoDBServiceResource = None  # pylint: disable=invalid-name
+
+
+class DynamoDB(NoSQLAdaptor):
+    """A MongoDB client that serves as an adaptor for profiling data assets on MongoDB"""
+
+    def __init__(self, client: DynamoDBServiceResource):
+        self.client = client
+
+    def item_count(self, table: Table) -> int:
+        table = self.client.Table(table.name.__root__)
+        return table.item_count
+
+    def scan(
+        self, table: Table, columns: List[Column], limit: int
+    ) -> List[Dict[str, any]]:
+        table = self.client.Table(table.name.__root__)
+        response = table.scan(Limit=limit)
+        return response["Items"]

metadata/profiler/adaptors/factory.py

@@ -0,0 +1,76 @@
+#  Copyright 2024 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.
+"""
+factory for NoSQL adaptors that are used in the NoSQLProfiler.
+"""
+from typing import Callable
+
+from metadata.generated.schema.entity.services.connections.database.dynamoDBConnection import (
+    DynamoDBConnection,
+)
+from metadata.generated.schema.entity.services.connections.database.mongoDBConnection import (
+    MongoDBConnection,
+)
+from metadata.profiler.adaptors.dynamodb import DynamoDB
+from metadata.profiler.adaptors.mongodb import MongoDB
+from metadata.profiler.adaptors.nosql_adaptor import NoSQLAdaptor
+from metadata.profiler.factory import Factory
+
+NoSQLAdaptorConstructor = Callable[[any], NoSQLAdaptor]
+
+
+class NoSQLAdaptorFactory(Factory):
+    """
+    A factory class for creating NoSQL client instances.
+
+    This class maintains a registry of callable constructors for different NoSQL client types.
+    The client types are registered with their corresponding constructors,
+    and can be created using the `construct` method.
+    """
+
+    def register(self, interface_type: str, interface_class: NoSQLAdaptorConstructor):
+        """
+        Register a client type with its constructor.
+
+        Args:
+            source_class_name (str): The class of the source client.
+            target_class (NoSQLClientConstructor): The constructor for the target client.
+
+        Returns:
+            None
+        """
+        self._interface_type[interface_type] = interface_class
+
+    def create(self, interface_type: any, *args, **kwargs) -> NoSQLAdaptor:
+        """
+        Create a client instance of the type of the given source client.
+
+        Args:
+            interface_type (str): The type of the source connection.
+
+        Returns:
+            NoSQLAdaptor: The created client instance.
+
+        Raises:
+            ValueError: If the type of the source client is not registered.
+        """
+        client_class = self._interface_type.get(interface_type)
+        if not client_class:
+            raise ValueError(f"Unknown NoSQL source: {interface_type}")
+        return client_class(*args, **kwargs)
+
+
+factory = NoSQLAdaptorFactory()
+adaptors = {
+    MongoDBConnection.__name__: MongoDB,
+    DynamoDBConnection.__name__: DynamoDB,
+}
+factory.register_many(adaptors)