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

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)

metadata/profiler/adaptors/mongodb.py

@@ -0,0 +1,178 @@
+#  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.
+"""
+MongoDB adaptor for the NoSQL profiler.
+"""
+import json
+from enum import Enum
+from typing import TYPE_CHECKING, Dict, List, Optional, Union
+
+from pydantic import BaseModel, Field
+
+from metadata.generated.schema.entity.data.table import Column, Table
+from metadata.profiler.adaptors.nosql_adaptor import NoSQLAdaptor
+from metadata.utils.sqa_like_column import SQALikeColumn
+
+# pylint: disable=invalid-name
+if TYPE_CHECKING:
+    from pymongo import MongoClient
+    from pymongo.command_cursor import CommandCursor
+    from pymongo.cursor import Cursor
+else:
+    MongoClient = None
+    CommandCursor = None
+    Cursor = None
+
+
+class AggregationFunction(Enum):
+    SUM = "$sum"
+    MEAN = "$avg"
+    COUNT = "$count"
+    MAX = "$max"
+    MIN = "$min"
+
+
+class Executable(BaseModel):
+    def to_executable(self, client: MongoClient) -> Union[CommandCursor, Cursor]:
+        raise NotImplementedError
+
+
+class Query(Executable):
+    database: str
+    collection: str
+    filter: dict = Field(default_factory=dict)
+    limit: Optional[int] = None
+
+    def to_executable(self, client: MongoClient) -> Cursor:
+        db = client[self.database]
+        collection = db[self.collection]
+        query = collection.find(self.filter)
+        if self.limit:
+            query = query.limit(self.limit)
+        return query
+
+
+class Aggregation(Executable):
+    database: str
+    collection: str
+    column: str
+    aggregations: List[AggregationFunction]
+
+    def to_executable(self, client: MongoClient) -> CommandCursor:
+        db = client[self.database]
+        collection = db[self.collection]
+        return collection.aggregate(
+            [
+                {
+                    "$group": {
+                        "_id": None,
+                        **{
+                            a.name.lower(): {a.value: f"${self.column}"}
+                            for a in self.aggregations
+                        },
+                    }
+                }
+            ]
+        )
+
+
+class MongoDB(NoSQLAdaptor):
+    """A MongoDB client that serves as an adaptor for profiling data assets on MongoDB"""
+
+    def __init__(self, client: MongoClient):
+        self.client = client
+
+    def item_count(self, table: Table) -> int:
+        db = self.client[table.databaseSchema.name]
+        collection = db[table.name.__root__]
+        return collection.count_documents({})
+
+    def scan(
+        self, table: Table, columns: List[Column], limit: int
+    ) -> List[Dict[str, any]]:
+        return self.execute(
+            Query(
+                database=table.databaseSchema.name,
+                collection=table.name.__root__,
+                limit=limit,
+            )
+        )
+
+    def query(
+        self, table: Table, columns: List[Column], query: any, limit: int
+    ) -> List[Dict[str, any]]:
+        try:
+            json_query = json.loads(query)
+        except json.JSONDecodeError:
+            raise ValueError("Invalid JSON query")
+        return self.execute(
+            Query(
+                database=table.databaseSchema.name,
+                collection=table.name.__root__,
+                filter=json_query,
+            )
+        )
+
+    def get_aggregates(
+        self,
+        table: Table,
+        column: SQALikeColumn,
+        aggregate_functions: List[AggregationFunction],
+    ) -> Dict[str, Union[int, float]]:
+        """
+        Get the aggregate functions for a column in a table
+        Returns:
+            Dict[str, Union[int, float]]: A dictionary of the aggregate functions
+            Example:
+            {
+                "sum": 100,
+                "avg": 50,
+                "count": 2,
+                "max": 75,
+                "min": 25
+            }
+        """
+        row = self.execute(
+            Aggregation(
+                database=table.databaseSchema.name,
+                collection=table.name.__root__,
+                column=column.name,
+                aggregations=aggregate_functions,
+            )
+        )[0]
+        return {k: v for k, v in row.items() if k != "_id"}
+
+    def sum(self, table: Table, column: SQALikeColumn) -> AggregationFunction:
+        return AggregationFunction.SUM
+
+    def mean(self, table: Table, column: SQALikeColumn) -> AggregationFunction:
+        return AggregationFunction.MEAN
+
+    def max(self, table: Table, column: SQALikeColumn) -> AggregationFunction:
+        return AggregationFunction.MAX
+
+    def min(self, table: Table, column: SQALikeColumn) -> AggregationFunction:
+        return AggregationFunction.MIN
+
+    def execute(self, query: Executable) -> List[Dict[str, any]]:
+        records = list(query.to_executable(self.client))
+        result = []
+        for r in records:
+            result.append({c: self._json_safe(r.get(c)) for c in r})
+        return result
+
+    @staticmethod
+    def _json_safe(data: any):
+        try:
+            json.dumps(data)
+            return data
+        except Exception:  # noqa
+            return str(data)

metadata/profiler/adaptors/nosql_adaptor.py

@@ -0,0 +1,65 @@
+#  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.
+"""
+NoSQL adaptor for the NoSQL profiler.
+"""
+from abc import ABC, abstractmethod
+from typing import Dict, List, Union
+
+from metadata.generated.schema.entity.data.table import Column, Table
+from metadata.utils.sqa_like_column import SQALikeColumn
+
+
+class NoSQLAdaptor(ABC):
+    """
+    NoSQL adaptor for the NoSQL profiler. This class implememts the required methods for retreiving data from a NoSQL
+    database.
+    """
+
+    @abstractmethod
+    def item_count(self, table: Table) -> int:
+        raise NotImplementedError
+
+    @abstractmethod
+    def scan(
+        self, table: Table, columns: List[Column], limit: int
+    ) -> List[Dict[str, any]]:
+        pass
+
+    def query(
+        self, table: Table, columns: List[Column], query: any, limit: int
+    ) -> List[Dict[str, any]]:
+        raise NotImplementedError
+
+    def get_aggregates(
+        self, table: Table, column: SQALikeColumn, aggregate_functions: List[any]
+    ) -> Dict[str, Union[int, float]]:
+        raise NotImplementedError
+
+    def sum(
+        self, table: Table, column: Column  # pylint: disable=unused-argument
+    ) -> any:
+        return None
+
+    def mean(
+        self, table: Table, column: Column  # pylint: disable=unused-argument
+    ) -> any:
+        return None
+
+    def max(
+        self, table: Table, column: Column  # pylint: disable=unused-argument
+    ) -> any:
+        return None
+
+    def min(
+        self, table: Table, column: Column  # pylint: disable=unused-argument
+    ) -> any:
+        return None

metadata/profiler/api/models.py

@@ -122,7 +122,7 @@ class ProfilerResponse(ConfigModel):
 
 
 class ThreadPoolMetrics(ConfigModel):
-    """thread pool metric"""
+    """A container for all metrics to be computed on the same thread."""
 
     metrics: Union[List[Union[Type[Metric], CustomMetric]], Type[Metric]]
     metric_type: MetricTypes

metadata/profiler/factory.py

@@ -0,0 +1,41 @@
+#  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.
+
+"""
+Factory class for creating profiler interface objects
+"""
+from abc import ABC, abstractmethod
+
+
+class Factory(ABC):
+    """Creational factory for interface objects"""
+
+    def __init__(self):
+        self._interface_type = {}
+
+    def register(self, interface_type: str, interface_class):
+        """Register a new interface"""
+        self._interface_type[interface_type] = interface_class
+
+    def register_many(self, interface_dict):
+        """
+        Registers multiple profiler interfaces at once.
+
+        Args:
+            interface_dict: A dictionary mapping connection class names (strings) to their
+            corresponding profiler interface classes.
+        """
+        for interface_type, interface_class in interface_dict.items():
+            self.register(interface_type, interface_class)
+
+    @abstractmethod
+    def create(self, interface_type: str, *args, **kwargs) -> any:
+        pass

metadata/profiler/interface/nosql/profiler_interface.py

@@ -0,0 +1,236 @@
+#  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.
+#  pylint: disable=arguments-differ
+
+"""
+Interfaces with database for all database engine
+supporting sqlalchemy abstraction layer
+"""
+import traceback
+from collections import defaultdict
+from datetime import datetime, timezone
+from typing import Dict, List, Optional, Type
+
+from sqlalchemy import Column
+
+from metadata.generated.schema.entity.data.table import TableData
+from metadata.generated.schema.tests.customMetric import CustomMetric
+from metadata.profiler.adaptors.factory import factory
+from metadata.profiler.adaptors.nosql_adaptor import NoSQLAdaptor
+from metadata.profiler.api.models import ThreadPoolMetrics
+from metadata.profiler.interface.profiler_interface import ProfilerInterface
+from metadata.profiler.metrics.core import Metric, MetricTypes
+from metadata.profiler.metrics.registry import Metrics
+from metadata.profiler.processor.metric_filter import MetricFilter
+from metadata.profiler.processor.sampler.nosql.sampler import NoSQLSampler
+from metadata.utils.logger import profiler_interface_registry_logger
+from metadata.utils.sqa_like_column import SQALikeColumn
+
+logger = profiler_interface_registry_logger()
+
+
+class NoSQLProfilerInterface(ProfilerInterface):
+    """
+    Interface to interact with registry supporting
+    sqlalchemy.
+    """
+
+    # pylint: disable=too-many-arguments
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.sampler = self._get_sampler()
+
+    def _compute_table_metrics(
+        self,
+        metrics: List[Type[Metric]],
+        runner: NoSQLAdaptor,
+        *args,
+        **kwargs,
+    ):
+        result = {}
+        for metric in metrics:
+            try:
+                fn = metric().nosql_fn(runner)
+                result[metric.name()] = fn(self.table)
+            except Exception as exc:
+                logger.debug(
+                    f"{traceback.format_exc()}\n"
+                    f"Error trying to compute metric {metric} for {self.table.fullyQualifiedName}: {exc}"
+                )
+                raise RuntimeError(
+                    f"Error trying to compute metric {metric.name()} for {self.table.fullyQualifiedName}: {exc}"
+                )
+        return result
+
+    def _compute_static_metrics(
+        self,
+        metrics: List[Metrics],
+        runner: NoSQLAdaptor,
+        column: SQALikeColumn,
+        *args,
+        **kwargs,
+    ) -> Dict[str, any]:
+        try:
+            aggs = [metric(column).nosql_fn(runner)(self.table) for metric in metrics]
+            filtered = [agg for agg in aggs if agg is not None]
+            if not filtered:
+                return {}
+            row = runner.get_aggregates(self.table, column, filtered)
+            return dict(row)
+        except Exception as exc:
+            logger.debug(
+                f"{traceback.format_exc()}\n"
+                f"Error trying to compute metrics for {self.table.fullyQualifiedName}: {exc}"
+            )
+            raise RuntimeError(
+                f"Error trying to compute metris for {self.table.fullyQualifiedName}: {exc}"
+            )
+
+    def _compute_query_metrics(
+        self,
+        metric: Metrics,
+        runner,
+        *args,
+        **kwargs,
+    ):
+        return None
+
+    def _compute_window_metrics(
+        self,
+        metrics: List[Metrics],
+        runner,
+        *args,
+        **kwargs,
+    ):
+        return None
+
+    def _compute_system_metrics(
+        self,
+        metrics: Metrics,
+        runner: List,
+        *args,
+        **kwargs,
+    ):
+        return None
+
+    def _compute_custom_metrics(
+        self, metrics: List[CustomMetric], runner, *args, **kwargs
+    ):
+        return None
+
+    def compute_metrics(
+        self,
+        client: NoSQLAdaptor,
+        metric_func: ThreadPoolMetrics,
+    ):
+        """Run metrics in processor worker"""
+        logger.debug(f"Running profiler for {metric_func.table}")
+        try:
+            row = self._get_metric_fn[metric_func.metric_type.value](
+                metric_func.metrics,
+                client,
+                column=metric_func.column,
+            )
+        except Exception as exc:
+            name = f"{metric_func.column if metric_func.column is not None else metric_func.table}"
+            error = f"{name} metric_type.value: {exc}"
+            logger.error(error)
+            self.status.failed_profiler(error, traceback.format_exc())
+            row = None
+        if metric_func.column is not None:
+            column = metric_func.column.name
+            self.status.scanned(f"{metric_func.table.name.__root__}.{column}")
+        else:
+            self.status.scanned(metric_func.table.name.__root__)
+            column = None
+        return row, column, metric_func.metric_type.value
+
+    def fetch_sample_data(self, table, columns: List[SQALikeColumn]) -> TableData:
+        return self.sampler.fetch_sample_data(columns)
+
+    def _get_sampler(self) -> NoSQLSampler:
+        """Get NoSQL sampler from config"""
+        from metadata.profiler.processor.sampler.sampler_factory import (  # pylint: disable=import-outside-toplevel
+            sampler_factory_,
+        )
+
+        return sampler_factory_.create(
+            self.service_connection_config.__class__.__name__,
+            table=self.table,
+            client=factory.create(
+                self.service_connection_config.__class__.__name__,
+                client=self.connection,
+            ),
+            profile_sample_config=self.profile_sample_config,
+            partition_details=self.partition_details,
+            profile_sample_query=self.profile_query,
+        )
+
+    def get_composed_metrics(
+        self, column: Column, metric: Metrics, column_results: Dict
+    ):
+        return None
+
+    def get_hybrid_metrics(
+        self, column: Column, metric: Metrics, column_results: Dict, **kwargs
+    ):
+        return None
+
+    def get_all_metrics(
+        self,
+        metric_funcs: List[ThreadPoolMetrics],
+    ):
+        """get all profiler metrics"""
+        profile_results = {"table": {}, "columns": defaultdict(dict)}
+        runner = factory.create(
+            self.service_connection_config.__class__.__name__,
+            client=self.connection,
+        )
+        metric_list = [
+            self.compute_metrics(runner, metric_func)
+            for metric_func in MetricFilter.filter_empty_metrics(metric_funcs)
+        ]
+        for metric_result in metric_list:
+            profile, column, metric_type = metric_result
+            if profile:
+                if metric_type == MetricTypes.Table.value:
+                    profile_results["table"].update(profile)
+                if metric_type == MetricTypes.System.value:
+                    profile_results["system"] = profile
+                elif metric_type == MetricTypes.Custom.value and column is None:
+                    profile_results["table"].update(profile)
+                else:
+                    profile_results["columns"][column].update(
+                        {
+                            "name": column,
+                            "timestamp": int(
+                                datetime.now(tz=timezone.utc).timestamp() * 1000
+                            ),
+                            **profile,
+                        }
+                    )
+        return profile_results
+
+    @property
+    def table(self):
+        """OM Table entity"""
+        return self.table_entity
+
+    def get_columns(self) -> List[Optional[SQALikeColumn]]:
+        return [
+            SQALikeColumn(name=c.name.__root__, type=c.dataType)
+            for c in self.table.columns
+        ]
+
+    def close(self):
+        if getattr(self.connection, "close", None):
+            self.connection.close()