openmetadata-ingestion 1.1.4.0__py3-none-any.whl → 1.1.6.0__py3-none-any.whl

metadata/ingestion/source/database/database_service.py

@@ -112,7 +112,7 @@ class DatabaseServiceTopology(ServiceTopology):
             NodeStage(
                 type_=OMetaTagAndClassification,
                 context="tags",
-                processor="yield_tag_details",
+                processor="yield_database_schema_tag_details",
                 ack_sink=False,
                 nullable=True,
                 cache_all=True,
@@ -130,6 +130,14 @@ class DatabaseServiceTopology(ServiceTopology):
     table = TopologyNode(
         producer="get_tables_name_and_type",
         stages=[
+            NodeStage(
+                type_=OMetaTagAndClassification,
+                context="tags",
+                processor="yield_table_tag_details",
+                ack_sink=False,
+                nullable=True,
+                cache_all=True,
+            ),
             NodeStage(
                 type_=Table,
                 context="table",
@@ -218,7 +226,23 @@ class DatabaseServiceSource(
         From topology. To be run for each schema
         """
 
-    def yield_tag_details(
+    def yield_table_tags(
+        self, table_name_and_type: Tuple[str, TableType]
+    ) -> Iterable[CreateTableRequest]:
+        """
+        From topology. To be run for each table
+        """
+
+    def yield_table_tag_details(
+        self, table_name_and_type: str
+    ) -> Iterable[OMetaTagAndClassification]:
+        """
+        From topology. To be run for each table
+        """
+        if self.source_config.includeTags:
+            yield from self.yield_table_tags(table_name_and_type) or []
+
+    def yield_database_schema_tag_details(
         self, schema_name: str
     ) -> Iterable[OMetaTagAndClassification]:
         """
@@ -267,7 +291,7 @@ class DatabaseServiceSource(
 
         tag_labels = []
         for tag_and_category in self.context.tags or []:
-            if tag_and_category.fqn.__root__ == entity_fqn:
+            if tag_and_category.fqn and tag_and_category.fqn.__root__ == entity_fqn:
                 tag_label = get_tag_label(
                     metadata=self.metadata,
                     tag_name=tag_and_category.tag_request.name.__root__,

metadata/ingestion/source/database/lineage_source.py

@@ -78,25 +78,28 @@ class LineageSource(QueryParserSource, ABC):
         Given an engine, iterate over the query results to
         yield a TableQuery with query parsing info
         """
-        with self.engine.connect() as conn:
-            rows = conn.execute(
-                self.get_sql_statement(
-                    start_time=self.start,
-                    end_time=self.end,
-                )
-            )
-            for row in rows:
-                query_dict = dict(row)
-                try:
-                    yield TableQuery(
-                        query=query_dict["query_text"],
-                        databaseName=self.get_database_name(query_dict),
-                        serviceName=self.config.serviceName,
-                        databaseSchema=self.get_schema_name(query_dict),
+        for engine in self.get_engine():
+            with engine.connect() as conn:
+                rows = conn.execute(
+                    self.get_sql_statement(
+                        start_time=self.start,
+                        end_time=self.end,
                     )
-                except Exception as exc:
-                    logger.debug(traceback.format_exc())
-                    logger.warning(f"Error processing query_dict {query_dict}: {exc}")
+                )
+                for row in rows:
+                    query_dict = dict(row)
+                    try:
+                        yield TableQuery(
+                            query=query_dict["query_text"],
+                            databaseName=self.get_database_name(query_dict),
+                            serviceName=self.config.serviceName,
+                            databaseSchema=self.get_schema_name(query_dict),
+                        )
+                    except Exception as exc:
+                        logger.debug(traceback.format_exc())
+                        logger.warning(
+                            f"Error processing query_dict {query_dict}: {exc}"
+                        )
 
     def next_record(self) -> Iterable[AddLineageRequest]:
         """

metadata/ingestion/source/database/query_parser_source.py

@@ -112,6 +112,9 @@ class QueryParserSource(Source[Union[TableQuery, AddLineageRequest]], ABC):
             return f"{self.filters} AND {self.source_config.filterCondition}"
         return self.filters
 
+    def get_engine(self):
+        yield self.engine
+
     def close(self):
         """
         By default, there is nothing to close

metadata/ingestion/source/database/sample_data.py

@@ -1121,15 +1121,22 @@ class SampleDataSource(
                             rowCount=profile["rowCount"],
                             createDateTime=profile.get("createDateTime"),
                             sizeInByte=profile.get("sizeInByte"),
-                            timestamp=(
-                                datetime.now(tz=timezone.utc) - timedelta(days=days)
-                            ).timestamp(),
+                            timestamp=int(
+                                (
+                                    datetime.now(tz=timezone.utc) - timedelta(days=days)
+                                ).timestamp()
+                                * 1000
+                            ),
                         ),
                         columnProfile=[
                             ColumnProfile(
-                                timestamp=(
-                                    datetime.now(tz=timezone.utc) - timedelta(days=days)
-                                ).timestamp(),
+                                timestamp=int(
+                                    (
+                                        datetime.now(tz=timezone.utc)
+                                        - timedelta(days=days)
+                                    ).timestamp()
+                                    * 1000
+                                ),
                                 **col_profile,
                             )
                             for col_profile in profile["columnProfile"]
@@ -1216,9 +1223,10 @@ class SampleDataSource(
                 for days, result in enumerate(test_case_results["results"]):
                     yield OMetaTestCaseResultsSample(
                         test_case_results=TestCaseResult(
-                            timestamp=(
-                                datetime.now() - timedelta(days=days)
-                            ).timestamp(),
+                            timestamp=int(
+                                (datetime.now() - timedelta(days=days)).timestamp()
+                                * 1000
+                            ),
                             testCaseStatus=result["testCaseStatus"],
                             result=result["result"],
                             testResultValue=[

metadata/ingestion/source/database/usage_source.py

@@ -93,37 +93,38 @@ class UsageSource(QueryParserSource, ABC):
                 f"{(self.start + timedelta(days=days + 1)).date()}"
             )
             try:
-                with self.engine.connect() as conn:
-                    rows = conn.execute(
-                        self.get_sql_statement(
-                            start_time=self.start + timedelta(days=days),
-                            end_time=self.start + timedelta(days=days + 1),
+                for engine in self.get_engine():
+                    with engine.connect() as conn:
+                        rows = conn.execute(
+                            self.get_sql_statement(
+                                start_time=self.start + timedelta(days=days),
+                                end_time=self.start + timedelta(days=days + 1),
+                            )
                         )
-                    )
-                    queries = []
-                    for row in rows:
-                        row = dict(row)
-                        try:
-                            queries.append(
-                                TableQuery(
-                                    query=self.format_query(row["query_text"]),
-                                    userName=row["user_name"],
-                                    startTime=str(row["start_time"]),
-                                    endTime=str(row["end_time"]),
-                                    analysisDate=row["start_time"],
-                                    aborted=self.get_aborted_status(row),
-                                    databaseName=self.get_database_name(row),
-                                    duration=row.get("duration"),
-                                    serviceName=self.config.serviceName,
-                                    databaseSchema=self.get_schema_name(row),
+                        queries = []
+                        for row in rows:
+                            row = dict(row)
+                            try:
+                                queries.append(
+                                    TableQuery(
+                                        query=self.format_query(row["query_text"]),
+                                        userName=row["user_name"],
+                                        startTime=str(row["start_time"]),
+                                        endTime=str(row["end_time"]),
+                                        analysisDate=row["start_time"],
+                                        aborted=self.get_aborted_status(row),
+                                        databaseName=self.get_database_name(row),
+                                        duration=row.get("duration"),
+                                        serviceName=self.config.serviceName,
+                                        databaseSchema=self.get_schema_name(row),
+                                    )
                                 )
-                            )
-                        except Exception as exc:
-                            logger.debug(traceback.format_exc())
-                            logger.warning(
-                                f"Unexpected exception processing row [{row}]: {exc}"
-                            )
-                yield TableQueries(queries=queries)
+                            except Exception as exc:
+                                logger.debug(traceback.format_exc())
+                                logger.warning(
+                                    f"Unexpected exception processing row [{row}]: {exc}"
+                                )
+                    yield TableQueries(queries=queries)
             except Exception as exc:
                 logger.debug(traceback.format_exc())
                 logger.error(f"Source usage processing error: {exc}")

metadata/ingestion/stage/table_usage.py

@@ -18,6 +18,7 @@ import os
 import shutil
 import traceback
 from pathlib import Path
+from typing import List, Tuple
 
 from metadata.config.common import ConfigModel
 from metadata.generated.schema.api.data.createQuery import CreateQueryRequest
@@ -25,7 +26,7 @@ from metadata.generated.schema.entity.services.connections.metadata.openMetadata
     OpenMetadataConnection,
 )
 from metadata.generated.schema.entity.teams.user import User
-from metadata.generated.schema.type.queryParserData import QueryParserData
+from metadata.generated.schema.type.queryParserData import ParsedData, QueryParserData
 from metadata.generated.schema.type.tableUsageCount import TableUsageCount
 from metadata.ingestion.api.stage import Stage
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
@@ -80,20 +81,22 @@ class TableUsageStage(Stage[QueryParserData]):
         logger.info(f"Creating the directory to store staging data in {location}")
         location.mkdir(parents=True, exist_ok=True)
 
-    def _get_user_entity(self, username: str):
+    def _get_user_entity(self, username: str) -> Tuple[List[str], List[str]]:
         if username:
             user = self.metadata.get_by_name(entity=User, fqn=username)
             if user:
-                return [user.fullyQualifiedName.__root__]
-        return []
+                return [user.fullyQualifiedName.__root__], []
+        return [], [username]
 
     def _add_sql_query(self, record, table):
+        users, used_by = self._get_user_entity(record.userName)
         if self.table_queries.get((table, record.date)):
             self.table_queries[(table, record.date)].append(
                 CreateQueryRequest(
                     query=record.sql,
-                    users=self._get_user_entity(record.userName),
+                    users=users,
                     queryDate=record.date,
+                    usedBy=used_by,
                     duration=record.duration,
                 )
             )
@@ -101,13 +104,45 @@ class TableUsageStage(Stage[QueryParserData]):
             self.table_queries[(table, record.date)] = [
                 CreateQueryRequest(
                     query=record.sql,
-                    users=self._get_user_entity(record.userName),
+                    users=users,
                     queryDate=record.date,
+                    usedBy=used_by,
                     duration=record.duration,
                 )
             ]
 
+    def _handle_table_usage(self, parsed_data: ParsedData, table: str) -> None:
+        table_joins = parsed_data.joins.get(table)
+        try:
+            self._add_sql_query(record=parsed_data, table=table)
+            table_usage_count = self.table_usage.get((table, parsed_data.date))
+            if table_usage_count is not None:
+                table_usage_count.count = table_usage_count.count + 1
+                if table_joins:
+                    table_usage_count.joins.extend(table_joins)
+            else:
+                joins = []
+                if table_joins:
+                    joins.extend(table_joins)
+
+                table_usage_count = TableUsageCount(
+                    table=table,
+                    databaseName=parsed_data.databaseName,
+                    date=parsed_data.date,
+                    joins=joins,
+                    serviceName=parsed_data.serviceName,
+                    sqlQueries=[],
+                    databaseSchema=parsed_data.databaseSchema,
+                )
+
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.warning(f"Error in staging record: {exc}")
+        self.table_usage[(table, parsed_data.date)] = table_usage_count
+        logger.info(f"Successfully record staged for {table}")
+
     def stage_record(self, record: QueryParserData) -> None:
+
         """
         Process the parsed data and store it in a file
         """
@@ -119,34 +154,7 @@ class TableUsageStage(Stage[QueryParserData]):
             if parsed_data is None:
                 continue
             for table in parsed_data.tables:
-                table_joins = parsed_data.joins.get(table)
-                try:
-                    self._add_sql_query(record=parsed_data, table=table)
-                    table_usage_count = self.table_usage.get((table, parsed_data.date))
-                    if table_usage_count is not None:
-                        table_usage_count.count = table_usage_count.count + 1
-                        if table_joins:
-                            table_usage_count.joins.extend(table_joins)
-                    else:
-                        joins = []
-                        if table_joins:
-                            joins.extend(table_joins)
-
-                        table_usage_count = TableUsageCount(
-                            table=table,
-                            databaseName=parsed_data.databaseName,
-                            date=parsed_data.date,
-                            joins=joins,
-                            serviceName=parsed_data.serviceName,
-                            sqlQueries=[],
-                            databaseSchema=parsed_data.databaseSchema,
-                        )
-
-                except Exception as exc:
-                    logger.debug(traceback.format_exc())
-                    logger.warning(f"Error in staging record: {exc}")
-                self.table_usage[(table, parsed_data.date)] = table_usage_count
-                logger.info(f"Successfully record staged for {table}")
+                self._handle_table_usage(parsed_data=parsed_data, table=table)
         self.dump_data_to_file()
 
     def dump_data_to_file(self):

metadata/profiler/interface/pandas/profiler_interface.py

@@ -8,6 +8,7 @@
 #  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
@@ -32,7 +33,6 @@ from metadata.profiler.metrics.registry import Metrics
 from metadata.profiler.processor.sampler.sampler_factory import sampler_factory_
 from metadata.readers.dataframe.models import DatalakeTableSchemaWrapper
 from metadata.utils.datalake.datalake_utils import fetch_dataframe
-from metadata.utils.dispatch import valuedispatch
 from metadata.utils.logger import profiler_interface_registry_logger
 from metadata.utils.sqa_like_column import SQALikeColumn, Type
 
@@ -109,20 +109,10 @@ class PandasProfilerInterface(ProfilerInterface, PandasInterfaceMixin):
             profile_sample_query=self.profile_query,
         )
 
-    @valuedispatch
-    def _get_metrics(self, *args, **kwargs):
-        """Generic getter method for metrics. To be used with
-        specific dispatch methods
-        """
-        logger.warning("Could not get metric. No function registered.")
-
-    # pylint: disable=unused-argument
-    @_get_metrics.register(MetricTypes.Table.value)
-    def _(
+    def _compute_table_metrics(
         self,
-        metric_type: str,
         metrics: List[Metrics],
-        dfs,
+        runner: List,
         *args,
         **kwargs,
     ):
@@ -138,7 +128,7 @@ class PandasProfilerInterface(ProfilerInterface, PandasInterfaceMixin):
 
         try:
             row_dict = {}
-            df_list = [df.where(pd.notnull(df), None) for df in dfs]
+            df_list = [df.where(pd.notnull(df), None) for df in runner]
             for metric in metrics:
                 row_dict[metric.name()] = metric().df_fn(df_list)
             return row_dict
@@ -147,13 +137,10 @@ class PandasProfilerInterface(ProfilerInterface, PandasInterfaceMixin):
             logger.warning(f"Error trying to compute profile for {exc}")
             raise RuntimeError(exc)
 
-    # pylint: disable=unused-argument
-    @_get_metrics.register(MetricTypes.Static.value)
-    def _(
+    def _compute_static_metrics(
         self,
-        metric_type: str,
         metrics: List[Metrics],
-        dfs,
+        runner: List,
         column,
         *args,
         **kwargs,
@@ -172,7 +159,7 @@ class PandasProfilerInterface(ProfilerInterface, PandasInterfaceMixin):
         try:
             row_dict = {}
             for metric in metrics:
-                metric_resp = metric(column).df_fn(dfs)
+                metric_resp = metric(column).df_fn(runner)
                 row_dict[metric.name()] = (
                     None if pd.isnull(metric_resp) else metric_resp
                 )
@@ -183,13 +170,10 @@ class PandasProfilerInterface(ProfilerInterface, PandasInterfaceMixin):
             )
             raise RuntimeError(exc)
 
-    # pylint: disable=unused-argument
-    @_get_metrics.register(MetricTypes.Query.value)
-    def _(
+    def _compute_query_metrics(
         self,
-        metric_type: str,
-        metrics: Metrics,
-        dfs,
+        metric: Metrics,
+        runner: List,
         column,
         *args,
         **kwargs,
@@ -204,18 +188,15 @@ class PandasProfilerInterface(ProfilerInterface, PandasInterfaceMixin):
             dictionnary of results
         """
         col_metric = None
-        col_metric = metrics(column).df_fn(dfs)
+        col_metric = metric(column).df_fn(runner)
         if not col_metric:
             return None
-        return {metrics.name(): col_metric}
+        return {metric.name(): col_metric}
 
-    # pylint: disable=unused-argument
-    @_get_metrics.register(MetricTypes.Window.value)
-    def _(
+    def _compute_window_metrics(
         self,
-        metric_type: str,
-        metrics: Metrics,
-        dfs,
+        metrics: List[Metrics],
+        runner: List,
         column,
         *args,
         **kwargs,
@@ -224,19 +205,21 @@ class PandasProfilerInterface(ProfilerInterface, PandasInterfaceMixin):
         Given a list of metrics, compute the given results
         and returns the values
         """
+
         try:
             metric_values = {}
             for metric in metrics:
-                metric_values[metric.name()] = metric(column).df_fn(dfs)
+                metric_values[metric.name()] = metric(column).df_fn(runner)
             return metric_values if metric_values else None
         except Exception as exc:
             logger.debug(traceback.format_exc())
             logger.warning(f"Unexpected exception computing metrics: {exc}")
             return None
 
-    @_get_metrics.register(MetricTypes.System.value)
-    def _(
+    def _compute_system_metrics(
         self,
+        metrics: Metrics,
+        runner: List,
         *args,
         **kwargs,
     ):
@@ -260,11 +243,9 @@ class PandasProfilerInterface(ProfilerInterface, PandasInterfaceMixin):
         try:
             row = None
             if self.dfs:
-                row = self._get_metrics(
-                    metric_type.value,
+                row = self._get_metric_fn[metric_type.value](
                     metrics,
                     dfs,
-                    session=self.client,
                     column=column,
                 )
         except Exception as exc:
@@ -354,7 +335,9 @@ class PandasProfilerInterface(ProfilerInterface, PandasInterfaceMixin):
                         profile_results["columns"][column].update(
                             {
                                 "name": column,
-                                "timestamp": datetime.now(tz=timezone.utc).timestamp(),
+                                "timestamp": int(
+                                    datetime.now(tz=timezone.utc).timestamp() * 1000
+                                ),
                                 **profile,
                             }
                         )

metadata/profiler/interface/profiler_interface.py

@@ -15,7 +15,7 @@ supporting sqlalchemy abstraction layer
 """
 
 from abc import ABC, abstractmethod
-from typing import Dict, Optional, Union
+from typing import Dict, List, Optional, Union
 
 from sqlalchemy import Column
 from typing_extensions import Self
@@ -36,7 +36,9 @@ from metadata.ingestion.api.processor import ProfilerProcessorStatus
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.connections import get_connection
 from metadata.profiler.api.models import ProfileSampleConfig, TableConfig
+from metadata.profiler.metrics.core import MetricTypes
 from metadata.profiler.metrics.registry import Metrics
+from metadata.profiler.processor.runner import QueryRunner
 from metadata.utils.partition import get_partition_details
 
 
@@ -78,6 +80,14 @@ class ProfilerInterface(ABC):
         )
         self.timeout_seconds = timeout_seconds
 
+        self._get_metric_fn = {
+            MetricTypes.Table.value: self._compute_table_metrics,
+            MetricTypes.Static.value: self._compute_static_metrics,
+            MetricTypes.Query.value: self._compute_query_metrics,
+            MetricTypes.Window.value: self._compute_window_metrics,
+            MetricTypes.System.value: self._compute_system_metrics,
+        }
+
     @abstractmethod
     def _get_sampler(self):
         """Get the sampler"""
@@ -222,7 +232,57 @@ class ProfilerInterface(ABC):
         raise NotImplementedError
 
     @abstractmethod
-    def _get_metrics(self, *args, **kwargs):
+    def _compute_table_metrics(
+        self,
+        metrics: List[Metrics],
+        runner,
+        *args,
+        **kwargs,
+    ):
+        """Get metrics"""
+        raise NotImplementedError
+
+    @abstractmethod
+    def _compute_static_metrics(
+        self,
+        metrics: List[Metrics],
+        runner,
+        *args,
+        **kwargs,
+    ):
+        """Get metrics"""
+        raise NotImplementedError
+
+    @abstractmethod
+    def _compute_query_metrics(
+        self,
+        metric: Metrics,
+        runner,
+        *args,
+        **kwargs,
+    ):
+        """Get metrics"""
+        raise NotImplementedError
+
+    @abstractmethod
+    def _compute_window_metrics(
+        self,
+        metrics: List[Metrics],
+        runner: QueryRunner,
+        *args,
+        **kwargs,
+    ):
+        """Get metrics"""
+        raise NotImplementedError
+
+    @abstractmethod
+    def _compute_system_metrics(
+        self,
+        metrics: Metrics,
+        runner,
+        *args,
+        **kwargs,
+    ):
         """Get metrics"""
         raise NotImplementedError
 

metadata/profiler/interface/profiler_interface_factory.py

@@ -21,6 +21,9 @@ from metadata.generated.schema.entity.services.connections.database.bigQueryConn
 from metadata.generated.schema.entity.services.connections.database.datalakeConnection import (
     DatalakeConnection,
 )
+from metadata.generated.schema.entity.services.connections.database.singleStoreConnection import (
+    SingleStoreConnection,
+)
 from metadata.generated.schema.entity.services.databaseService import DatabaseConnection
 from metadata.profiler.interface.pandas.profiler_interface import (
     PandasProfilerInterface,
@@ -32,6 +35,9 @@ from metadata.profiler.interface.sqlalchemy.bigquery.profiler_interface import (
 from metadata.profiler.interface.sqlalchemy.profiler_interface import (
     SQAProfilerInterface,
 )
+from metadata.profiler.interface.sqlalchemy.single_store.profiler_interface import (
+    SingleStoreProfilerInterface,
+)
 
 
 class ProfilerInterfaceFactory:
@@ -58,6 +64,9 @@ profiler_interface_factory.register(DatabaseConnection.__name__, SQAProfilerInte
 profiler_interface_factory.register(
     BigQueryConnection.__name__, BigQueryProfilerInterface
 )
+profiler_interface_factory.register(
+    SingleStoreConnection.__name__, SingleStoreProfilerInterface
+)
 profiler_interface_factory.register(
     DatalakeConnection.__name__, PandasProfilerInterface
 )