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

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

metadata/profiler/interface/pandas/profiler_interface.py

@@ -22,7 +22,11 @@ from typing import Dict, List, Optional
 
 from sqlalchemy import Column
 
-from metadata.generated.schema.entity.data.table import CustomMetricProfile, TableData
+from metadata.generated.schema.entity.data.table import (
+    CustomMetricProfile,
+    DataType,
+    TableData,
+)
 from metadata.generated.schema.entity.services.connections.database.datalakeConnection import (
     DatalakeConnection,
 )
@@ -32,12 +36,9 @@ from metadata.profiler.api.models import ThreadPoolMetrics
 from metadata.profiler.interface.profiler_interface import ProfilerInterface
 from metadata.profiler.metrics.core import MetricTypes
 from metadata.profiler.metrics.registry import Metrics
-from metadata.readers.dataframe.models import DatalakeTableSchemaWrapper
+from metadata.profiler.processor.metric_filter import MetricFilter
 from metadata.utils.constants import COMPLEX_COLUMN_SEPARATOR, SAMPLE_DATA_DEFAULT_COUNT
-from metadata.utils.datalake.datalake_utils import (
-    GenericDataFrameColumnParser,
-    fetch_dataframe,
-)
+from metadata.utils.datalake.datalake_utils import GenericDataFrameColumnParser
 from metadata.utils.logger import profiler_interface_registry_logger
 from metadata.utils.sqa_like_column import SQALikeColumn
 
@@ -85,29 +86,48 @@ class PandasProfilerInterface(ProfilerInterface, PandasInterfaceMixin):
         )
 
         self.client = self.connection.client
-        self.dfs = self._convert_table_to_list_of_dataframe_objects()
+        self.dfs = self.return_ometa_dataframes_sampled(
+            service_connection_config=self.service_connection_config,
+            client=self.client,
+            table=self.table_entity,
+            profile_sample_config=profile_sample_config,
+        )
         self.sampler = self._get_sampler()
-        self.complex_dataframe_sample = deepcopy(self.sampler.random_sample())
-
-    def _convert_table_to_list_of_dataframe_objects(self):
-        """From a table entity, return the corresponding dataframe object
+        self.complex_dataframe_sample = deepcopy(
+            self.sampler.random_sample(is_sampled=True)
+        )
+        self.complex_df()
 
-        Returns:
-            List[DataFrame]
-        """
-        data = fetch_dataframe(
-            config_source=self.service_connection_config.configSource,
-            client=self.client,
-            file_fqn=DatalakeTableSchemaWrapper(
-                key=self.table_entity.name.__root__,
-                bucket_name=self.table_entity.databaseSchema.name,
-                file_extension=self.table_entity.fileFormat,
-            ),
+    def complex_df(self):
+        """Assign DataTypes to dataframe columns as per the parsed column type"""
+        coltype_mapping_df = []
+        data_formats = (
+            GenericDataFrameColumnParser._data_formats  # pylint: disable=protected-access
         )
+        for index, df in enumerate(self.complex_dataframe_sample):
+            if index == 0:
+                for col in self.table.columns:
+                    coltype = next(
+                        (
+                            key
+                            for key, value in data_formats.items()
+                            if col.dataType == value
+                        ),
+                        None,
+                    )
+                    if coltype and col.dataType not in {DataType.JSON, DataType.ARRAY}:
+                        coltype_mapping_df.append(coltype)
+                    else:
+                        coltype_mapping_df.append("object")
 
-        if not data:
-            raise TypeError(f"Couldn't fetch {self.table_entity.name.__root__}")
-        return data
+            try:
+                self.complex_dataframe_sample[index] = df.astype(
+                    dict(zip(df.keys(), coltype_mapping_df))
+                )
+            except (TypeError, ValueError) as err:
+                self.complex_dataframe_sample[index] = df
+                logger.warning(f"NaN/NoneType found in the Dataframe: {err}")
+                break
 
     def _get_sampler(self):
         """Get dataframe sampler from config"""
@@ -171,19 +191,19 @@ class PandasProfilerInterface(ProfilerInterface, PandasInterfaceMixin):
         """
         import pandas as pd  # pylint: disable=import-outside-toplevel
 
+        row_dict = {}
         try:
-            row_dict = {}
             for metric in metrics:
                 metric_resp = metric(column).df_fn(runner)
                 row_dict[metric.name()] = (
                     None if pd.isnull(metric_resp) else metric_resp
                 )
-            return row_dict
         except Exception as exc:
             logger.debug(
                 f"{traceback.format_exc()}\nError trying to compute profile for {exc}"
             )
             raise RuntimeError(exc)
+        return row_dict
 
     def _compute_query_metrics(
         self,
@@ -283,7 +303,7 @@ class PandasProfilerInterface(ProfilerInterface, PandasInterfaceMixin):
         metric_func: ThreadPoolMetrics,
     ):
         """Run metrics in processor worker"""
-        logger.debug(f"Running profiler for {metric_func.table}")
+        logger.debug(f"Running profiler for {metric_func.table.name.__root__}")
         try:
             row = None
             if self.complex_dataframe_sample:
@@ -360,13 +380,14 @@ class PandasProfilerInterface(ProfilerInterface, PandasInterfaceMixin):
 
     def get_all_metrics(
         self,
-        metric_funcs: list,
+        metric_funcs: List[ThreadPoolMetrics],
     ):
         """get all profiler metrics"""
 
         profile_results = {"table": {}, "columns": defaultdict(dict)}
         metric_list = [
-            self.compute_metrics(metric_func) for metric_func in metric_funcs
+            self.compute_metrics(metric_func)
+            for metric_func in MetricFilter.filter_empty_metrics(metric_funcs)
         ]
         for metric_result in metric_list:
             profile, column, metric_type = metric_result

metadata/profiler/interface/profiler_interface.py

@@ -425,8 +425,12 @@ class ProfilerInterface(ABC):
         runner,
         *args,
         **kwargs,
-    ):
-        """Get metrics"""
+    ) -> Dict[str, Any]:
+        """Get metrics
+        Return:
+            Dict[str, Any]: dict of metrics tio be merged into the final column profile. Keys need to be compatible with
+            the `metadata.generated.schema.entity.data.table.ColumnProfile` schema.
+        """
         raise NotImplementedError
 
     @abstractmethod