openmetadata-ingestion 1.5.13.2__py3-none-any.whl → 1.6.0.0__py3-none-any.whl

metadata/profiler/metrics/system/system.py

@@ -13,47 +13,18 @@
 System Metric
 """
 
-import traceback
+from abc import ABC
 from collections import defaultdict
-from typing import Dict, List, Optional
+from typing import Callable, Generic, List, TypeVar
 
-from sqlalchemy import text
-from sqlalchemy.orm import DeclarativeMeta, Session
+from sqlalchemy.orm import Session
 
 from metadata.generated.schema.configuration.profilerConfiguration import MetricType
-from metadata.generated.schema.entity.services.connections.database.bigQueryConnection import (
-    BigQueryConnection,
-)
-from metadata.generated.schema.entity.services.databaseService import DatabaseService
-from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.generated.schema.entity.data.table import SystemProfile
 from metadata.profiler.metrics.core import SystemMetric
-from metadata.profiler.metrics.system.dml_operation import (
-    DML_OPERATION_MAP,
-    DatabaseDMLOperations,
-)
-from metadata.profiler.metrics.system.queries.bigquery import (
-    DML_STAT_TO_DML_STATEMENT_MAPPING,
-    JOBS,
-    BigQueryQueryResult,
-)
-from metadata.profiler.metrics.system.queries.redshift import (
-    STL_QUERY,
-    get_metric_result,
-    get_query_results,
-)
-from metadata.profiler.metrics.system.queries.snowflake import (
-    INFORMATION_SCHEMA_QUERY,
-    get_snowflake_system_queries,
-)
-from metadata.profiler.orm.registry import Dialects
-from metadata.utils.dispatch import valuedispatch
 from metadata.utils.helpers import deep_size_of_dict
 from metadata.utils.logger import profiler_logger
-from metadata.utils.profiler_utils import (
-    SnowflakeQueryResult,
-    get_value_from_cache,
-    set_cache,
-)
+from metadata.utils.lru_cache import LRU_CACHE_SIZE, LRUCache
 
 logger = profiler_logger()
 
@@ -67,343 +38,79 @@ def recursive_dic():
 
 SYSTEM_QUERY_RESULT_CACHE = recursive_dic()
 
+T = TypeVar("T")
 
-@valuedispatch
-def get_system_metrics_for_dialect(
-    dialect: str,
-    session: Session,
-    table: DeclarativeMeta,
-    *args,
-    **kwargs,
-) -> Optional[Dict]:
-    """_summary_
-
-    Args:
-        dialect (str): database API dialect
-        session (Session): session object
-
-    Returns:
-        Optional[Dict]: For BigQuery, Snowflake, Redshift returns
-            {
-                timestamp: <timestamp>,
-                operationType: <Enum: 'INSERT', 'UPDATE', 'DELETE'>
-                rowsAffected: <int>,
-            } else returns None
-    """
-    logger.debug(f"System metrics not support for {dialect}. Skipping processing.")
-
-
-@get_system_metrics_for_dialect.register(Dialects.BigQuery)
-def _(
-    dialect: str,
-    session: Session,
-    table: DeclarativeMeta,
-    conn_config: BigQueryConnection,
-    *args,
-    **kwargs,
-) -> List[Dict]:
-    """Compute system metrics for bigquery
-
-    Args:
-        dialect (str): bigquery
-        session (Session): session Object
-        table (DeclarativeMeta): orm table
-
-    Returns:
-        List[Dict]:
-    """
-    logger.debug(f"Fetching system metrics for {dialect}")
-
-    project_id = session.get_bind().url.host
-    dataset_id = table.__table_args__["schema"]  # type: ignore
-
-    metric_results: List[Dict] = []
-
-    jobs = get_value_from_cache(
-        SYSTEM_QUERY_RESULT_CACHE, f"{Dialects.BigQuery}.{project_id}.{dataset_id}.jobs"
-    )
-
-    if not jobs:
-        cursor_jobs = session.execute(
-            text(
-                JOBS.format(
-                    usage_location=conn_config.usageLocation,
-                    dataset_id=dataset_id,
-                    project_id=project_id,
-                    insert=DatabaseDMLOperations.INSERT.value,
-                    update=DatabaseDMLOperations.UPDATE.value,
-                    delete=DatabaseDMLOperations.DELETE.value,
-                    merge=DatabaseDMLOperations.MERGE.value,
-                )
-            )
-        )
-        jobs = [
-            BigQueryQueryResult(
-                query_type=row.statement_type,
-                timestamp=row.start_time,
-                table_name=row.destination_table,
-                dml_statistics=row.dml_statistics,
-            )
-            for row in cursor_jobs
-        ]
-        set_cache(
-            SYSTEM_QUERY_RESULT_CACHE,
-            f"{Dialects.BigQuery}.{project_id}.{dataset_id}.jobs",
-            jobs,
-        )
 
-    for job in jobs:
-        if job.table_name.get("table_id") == table.__tablename__:  # type: ignore
-            rows_affected = None
-            try:
-                if job.query_type == DatabaseDMLOperations.INSERT.value:
-                    rows_affected = job.dml_statistics.get("inserted_row_count")
-                if job.query_type == DatabaseDMLOperations.DELETE.value:
-                    rows_affected = job.dml_statistics.get("deleted_row_count")
-                if job.query_type == DatabaseDMLOperations.UPDATE.value:
-                    rows_affected = job.dml_statistics.get("updated_row_count")
-            except AttributeError:
-                logger.debug(traceback.format_exc())
-                rows_affected = None
-
-            if job.query_type == DatabaseDMLOperations.MERGE.value:
-                for indx, key in enumerate(job.dml_statistics):
-                    if job.dml_statistics[key] != 0:
-                        metric_results.append(
-                            {
-                                # Merge statement can include multiple DML operations
-                                # We are padding timestamps by 0,1,2 millisesond to avoid
-                                # duplicate timestamps
-                                "timestamp": int(job.timestamp.timestamp() * 1000)
-                                + indx,
-                                "operation": DML_STAT_TO_DML_STATEMENT_MAPPING.get(key),
-                                "rowsAffected": job.dml_statistics[key],
-                            }
-                        )
-                continue
-
-            metric_results.append(
-                {
-                    "timestamp": int(job.timestamp.timestamp() * 1000),
-                    "operation": job.query_type,
-                    "rowsAffected": rows_affected,
-                }
-            )
-
-    return metric_results
-
-
-@get_system_metrics_for_dialect.register(Dialects.Redshift)
-def _(
-    dialect: str,
-    session: Session,
-    table: DeclarativeMeta,
-    *args,
-    **kwargs,
-) -> List[Dict]:
-    """List all the DML operations for reshifts tables
-
-    Args:
-        dialect (str): redshift
-        session (Session): session object
-        table (DeclarativeMeta): orm table
-
-    Returns:
-        List[Dict]:
-    """
-    logger.debug(f"Fetching system metrics for {dialect}")
-    database = session.get_bind().url.database
-    schema = table.__table_args__["schema"]  # type: ignore
-
-    metric_results: List[Dict] = []
-
-    # get inserts ddl queries
-    inserts = get_value_from_cache(
-        SYSTEM_QUERY_RESULT_CACHE, f"{Dialects.Redshift}.{database}.{schema}.inserts"
-    )
-    if not inserts:
-        insert_query = STL_QUERY.format(
-            alias="si",
-            join_type="LEFT",
-            condition="sd.query is null",
-            database=database,
-            schema=schema,
-        )
-        inserts = get_query_results(
-            session,
-            insert_query,
-            DatabaseDMLOperations.INSERT.value,
-        )
-        set_cache(
-            SYSTEM_QUERY_RESULT_CACHE,
-            f"{Dialects.Redshift}.{database}.{schema}.inserts",
-            inserts,
-        )
-    metric_results.extend(get_metric_result(inserts, table.__tablename__))  # type: ignore
-
-    # get deletes ddl queries
-    deletes = get_value_from_cache(
-        SYSTEM_QUERY_RESULT_CACHE, f"{Dialects.Redshift}.{database}.{schema}.deletes"
-    )
-    if not deletes:
-        delete_query = STL_QUERY.format(
-            alias="sd",
-            join_type="RIGHT",
-            condition="si.query is null",
-            database=database,
-            schema=schema,
-        )
-        deletes = get_query_results(
-            session,
-            delete_query,
-            DatabaseDMLOperations.DELETE.value,
-        )
-        set_cache(
-            SYSTEM_QUERY_RESULT_CACHE,
-            f"{Dialects.Redshift}.{database}.{schema}.deletes",
-            deletes,
-        )
-    metric_results.extend(get_metric_result(deletes, table.__tablename__))  # type: ignore
-
-    # get updates ddl queries
-    updates = get_value_from_cache(
-        SYSTEM_QUERY_RESULT_CACHE, f"{Dialects.Redshift}.{database}.{schema}.updates"
-    )
-    if not updates:
-        update_query = STL_QUERY.format(
-            alias="si",
-            join_type="INNER",
-            condition="sd.query is not null",
-            database=database,
-            schema=schema,
-        )
-        updates = get_query_results(
-            session,
-            update_query,
-            DatabaseDMLOperations.UPDATE.value,
-        )
-        set_cache(
-            SYSTEM_QUERY_RESULT_CACHE,
-            f"{Dialects.Redshift}.{database}.{schema}.updates",
-            updates,
-        )
-    metric_results.extend(get_metric_result(updates, table.__tablename__))  # type: ignore
-
-    return metric_results
-
-
-def _snowflake_build_query_result(
-    session: Session,
-    table: DeclarativeMeta,
-    database: str,
-    schema: str,
-    ometa_client: OpenMetadata,
-    db_service: DatabaseService,
-) -> List[SnowflakeQueryResult]:
-    """List and parse snowflake DML query results"""
-    rows = session.execute(
-        text(
-            INFORMATION_SCHEMA_QUERY.format(
-                tablename=table.__tablename__,  # type: ignore
-                insert=DatabaseDMLOperations.INSERT.value,
-                update=DatabaseDMLOperations.UPDATE.value,
-                delete=DatabaseDMLOperations.DELETE.value,
-                merge=DatabaseDMLOperations.MERGE.value,
-            )
-        )
-    )
-    query_results = []
-    for row in rows:
-        result = get_snowflake_system_queries(
-            row=row,
-            database=database,
-            schema=schema,
-            ometa_client=ometa_client,
-            db_service=db_service,
-        )
-        if result:
-            query_results.append(result)
+class CacheProvider(ABC, Generic[T]):
+    """Cache provider class to provide cache for system metrics"""
 
-    return query_results
+    def __init__(self):
+        self.cache = LRUCache[T](LRU_CACHE_SIZE)
 
+    def get_or_update_cache(
+        self,
+        cache_path: str,
+        get_queries_fn: Callable[..., List[T]],
+        *args,
+        **kwargs,
+    ):
+        if cache_path in self.cache:
+            return self.cache.get(cache_path)
+        result = get_queries_fn(*args, **kwargs)
+        self.cache.put(cache_path, result)
+        return result
 
-@get_system_metrics_for_dialect.register(Dialects.Snowflake)
-def _(
-    dialect: str,
-    session: Session,
-    table: DeclarativeMeta,
-    ometa_client: OpenMetadata,
-    db_service: DatabaseService,
-    *args,
-    **kwargs,
-) -> Optional[List[Dict]]:
-    """Fetch system metrics for Snowflake. query_history will return maximum 10K rows in one request.
-    We'll be fetching all the queries ran for the past 24 hours and filtered on specific query types
-    (INSERTS, MERGE, DELETE, UPDATE).
 
-    :waring: Unlike redshift and bigquery results are not cached as we'll be looking
-    at DDL for each table
+class EmptySystemMetricsSource:
+    """Empty system metrics source that can be used as a default. Just returns an empty list of system metrics
+    for any resource."""
 
-    To get the number of rows affected we'll use the specific query ID.
+    def get_inserts(self, **kwargs) -> List[SystemProfile]:
+        """Get insert queries"""
+        return []
 
-    Args:
-        dialect (str): dialect
-        session (Session): session object
+    def get_deletes(self, **kwargs) -> List[SystemProfile]:
+        """Get delete queries"""
+        return []
 
-    Returns:
-        Dict: system metric
-    """
-    logger.debug(f"Fetching system metrics for {dialect}")
-    database = session.get_bind().url.database
-    schema = table.__table_args__["schema"]  # type: ignore
-
-    metric_results: List[Dict] = []
-
-    query_results = _snowflake_build_query_result(
-        session=session,
-        table=table,
-        database=database,
-        schema=schema,
-        ometa_client=ometa_client,
-        db_service=db_service,
-    )
-
-    for query_result in query_results:
-        rows_affected = None
-        if query_result.query_type == DatabaseDMLOperations.INSERT.value:
-            rows_affected = query_result.rows_inserted
-        if query_result.query_type == DatabaseDMLOperations.DELETE.value:
-            rows_affected = query_result.rows_deleted
-        if query_result.query_type == DatabaseDMLOperations.UPDATE.value:
-            rows_affected = query_result.rows_updated
-        if query_result.query_type == DatabaseDMLOperations.MERGE.value:
-            if query_result.rows_inserted:
-                metric_results.append(
-                    {
-                        "timestamp": int(query_result.timestamp.timestamp() * 1000),
-                        "operation": DatabaseDMLOperations.INSERT.value,
-                        "rowsAffected": query_result.rows_inserted,
-                    }
-                )
-            if query_result.rows_updated:
-                metric_results.append(
-                    {
-                        "timestamp": int(query_result.timestamp.timestamp() * 1000),
-                        "operation": DatabaseDMLOperations.UPDATE.value,
-                        "rowsAffected": query_result.rows_updated,
-                    }
-                )
-            continue
-
-        metric_results.append(
-            {
-                "timestamp": int(query_result.timestamp.timestamp() * 1000),
-                "operation": DML_OPERATION_MAP.get(query_result.query_type),
-                "rowsAffected": rows_affected,
-            }
+    def get_updates(self, **kwargs) -> List[SystemProfile]:
+        """Get update queries"""
+        return []
+
+    def get_kwargs(self, **kwargs):
+        """Get kwargs to be used in get_inserts, get_deletes, get_updates"""
+        return {}
+
+
+class SystemMetricsComputer(EmptySystemMetricsSource):
+    def __init__(self, *args, **kwargs):
+        # collaborative constructor that initalizes upstream classes
+        super().__init__(*args, **kwargs)
+
+    def get_system_metrics(self, **kwargs) -> List[SystemProfile]:
+        """Return system metrics for a given table. Actual passed object can be a variety of types based
+        on the underlying infrastructure. For example, in the case of SQLalchemy, it can be a Table object
+        and in the case of Mongo, it can be a collection object."""
+        kwargs = super().get_kwargs(**kwargs)
+        return (
+            super().get_inserts(**kwargs)
+            + super().get_deletes(**kwargs)
+            + super().get_updates(**kwargs)
         )
 
-    return metric_results
+
+class SQASessionProvider:
+    """SQASessionProvider class to provide session to the system metrics"""
+
+    def __init__(self, *args, **kwargs):
+        self.session = kwargs.pop("session")
+        super().__init__(*args, **kwargs)
+
+    def get_session(self):
+        return self.session
+
+    def get_database(self) -> str:
+        return self.session.get_bind().url.database
 
 
 class System(SystemMetric):
@@ -454,18 +161,6 @@ class System(SystemMetric):
                 )
 
     def sql(self, session: Session, **kwargs):
-        """Implements the SQL logic to fetch system data"""
-        self._validate_attrs(["table", "ometa_client", "db_service"])
-
-        conn_config = kwargs.get("conn_config")
-
-        system_metrics = get_system_metrics_for_dialect(
-            session.get_bind().dialect.name,
-            session=session,
-            table=self.table,  # pylint: disable=no-member
-            conn_config=conn_config,
-            ometa_client=self.ometa_client,  # pylint: disable=no-member
-            db_service=self.db_service,  # pylint: disable=no-member
+        raise NotImplementedError(
+            "SQL method is not implemented for System metric. Use SystemMetricsComputer.get_system_metrics instead"
         )
-        self._manage_cache()
-        return system_metrics

metadata/profiler/orm/converter/base.py

@@ -109,10 +109,29 @@ def ometa_to_sqa_orm(
     We are building the class dynamically using
     `type` and passing SQLAlchemy `Base` class
     as the bases tuple for inheritance.
+
+    Args:
+        table (Table): OpenMetadata Table instance
+        metadata (OpenMetadata): OpenMetadata connection
+        sqa_metadata_obj (MetaData): For advanced use cases, you can pass a custom MetaData object. For most cases, this
+        can be left as None so that the global_metadata object is used.
     """
+    _metadata = sqa_metadata_obj or Base.metadata
     table.serviceType = cast(
         databaseService.DatabaseServiceType, table.serviceType
     )  # satisfy mypy
+
+    orm_database_name = get_orm_database(table, metadata)
+    # SQLite does not support schemas
+    orm_schema_name = (
+        get_orm_schema(table, metadata)
+        if table.serviceType != databaseService.DatabaseServiceType.SQLite
+        else None
+    )
+    orm_name = f"{orm_database_name}_{orm_schema_name}_{table.name.root}".replace(
+        ".", "_"
+    )
+
     cols = {
         (
             col.name.root + "_"
@@ -122,12 +141,6 @@ def ometa_to_sqa_orm(
         for idx, col in enumerate(table.columns)
     }
 
-    orm_database_name = get_orm_database(table, metadata)
-    orm_schema_name = get_orm_schema(table, metadata)
-    orm_name = f"{orm_database_name}_{orm_schema_name}_{table.name.root}".replace(
-        ".", "_"
-    )
-
     # Type takes positional arguments in the form of (name, bases, dict)
     orm = type(
         orm_name,  # Output class name
@@ -135,23 +148,20 @@ def ometa_to_sqa_orm(
         {
             "__tablename__": str(table.name.root),
             "__table_args__": {
-                # SQLite does not support schemas
-                "schema": orm_schema_name
-                if table.serviceType != databaseService.DatabaseServiceType.SQLite
-                else None,
+                "schema": orm_schema_name,
                 "extend_existing": True,  # Recreates the table ORM object if it already exists. Useful for testing
                 "quote": check_snowflake_case_sensitive(
                     table.serviceType, table.name.root
-                ),
+                )
+                or None,
             },
             **cols,
-            "metadata": sqa_metadata_obj or Base.metadata,
+            "metadata": _metadata,
         },
     )
 
     if not isinstance(orm, DeclarativeMeta):
         raise ValueError("OMeta to ORM did not create a DeclarativeMeta")
-
     return orm
 
 

metadata/profiler/orm/converter/bigquery/converter.py

@@ -22,7 +22,7 @@ from sqlalchemy.sql.sqltypes import TypeEngine
 from metadata.generated.schema.entity.data.database import databaseService
 from metadata.generated.schema.entity.data.table import Column, DataType
 from metadata.profiler.orm.converter.common import CommonMapTypes
-from metadata.profiler.source.bigquery.type_mapper import bigquery_type_mapper
+from metadata.profiler.source.database.bigquery.type_mapper import bigquery_type_mapper
 
 
 class BigqueryMapTypes(CommonMapTypes):

metadata/profiler/orm/converter/converter_registry.py

@@ -19,6 +19,7 @@ from metadata.generated.schema.entity.services.databaseService import (
 )
 from metadata.profiler.orm.converter.bigquery.converter import BigqueryMapTypes
 from metadata.profiler.orm.converter.common import CommonMapTypes
+from metadata.profiler.orm.converter.mssql.converter import MssqlMapTypes
 from metadata.profiler.orm.converter.redshift.converter import RedshiftMapTypes
 from metadata.profiler.orm.converter.snowflake.converter import SnowflakeMapTypes
 
@@ -26,3 +27,4 @@ converter_registry = defaultdict(lambda: CommonMapTypes)
 converter_registry[DatabaseServiceType.BigQuery] = BigqueryMapTypes
 converter_registry[DatabaseServiceType.Snowflake] = SnowflakeMapTypes
 converter_registry[DatabaseServiceType.Redshift] = RedshiftMapTypes
+converter_registry[DatabaseServiceType.Mssql] = MssqlMapTypes

metadata/profiler/orm/converter/mssql/converter.py

@@ -16,10 +16,24 @@ Map Types to convert/cast mssql related data types to relevant data types
 
 from sqlalchemy import NVARCHAR, TEXT
 
-from metadata.profiler.orm.registry import CustomImage
+from metadata.profiler.orm.converter.common import CommonMapTypes
+from metadata.profiler.orm.registry import CustomImage, CustomTypes, DataType
 
 cast_dict = {
     CustomImage: "VARBINARY(max)",
     TEXT: "VARCHAR(max)",
     NVARCHAR: "NVARCHAR(max)",
 }
+
+
+class MssqlMapTypes(CommonMapTypes):
+    """
+    Mssql type mapper
+    """
+
+    def __init__(self) -> None:
+        self._TYPE_MAP.update(
+            {
+                DataType.TIMESTAMP: CustomTypes.TIMESTAMP.value,
+            }
+        )

metadata/profiler/orm/functions/count.py

@@ -18,6 +18,7 @@ from sqlalchemy.sql.sqltypes import NVARCHAR, TEXT
 
 from metadata.profiler.metrics.core import CACHE
 from metadata.profiler.orm.registry import Dialects
+from metadata.profiler.orm.types.custom_hex_byte_string import HexByteString
 
 # Keep SQA docs style defining custom constructs
 # pylint: disable=consider-using-f-string,duplicate-code
@@ -36,6 +37,14 @@ def _(element, compiler, **kw):
     return compiler.process(element.clauses, **kw)
 
 
+@compiles(CountFn, Dialects.Oracle)
+def _(element, compiler, **kw):
+    col_type = element.clauses.clauses[0].type
+    if isinstance(col_type, HexByteString):
+        return f"DBMS_LOB.GETLENGTH({compiler.process(element.clauses, **kw)})"
+    return compiler.process(element.clauses, **kw)
+
+
 @compiles(CountFn, Dialects.MSSQL)
 def _(element, compiler, **kw):
     col_type = element.clauses.clauses[0].type

metadata/profiler/orm/functions/md5.py

@@ -0,0 +1,42 @@
+#  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.
+
+"""
+Define MD5 hashing function
+"""
+from sqlalchemy.ext.compiler import compiles
+from sqlalchemy.sql.functions import FunctionElement
+
+from metadata.profiler.metrics.core import CACHE
+from metadata.profiler.orm.registry import PythonDialects
+from metadata.utils.logger import profiler_logger
+
+logger = profiler_logger()
+
+
+class MD5(FunctionElement):
+    inherit_cache = CACHE
+
+
+@compiles(MD5)
+def _(element, compiler, **kw):
+    return f"MD5({compiler.process(element.clauses, **kw)})"
+
+
+@compiles(MD5, PythonDialects.MSSQL.value)
+def _(element, compiler, **kw):
+    # TODO requires separate where clauses for each table
+    return f"CONVERT(VARCHAR(8), HashBytes('MD5', {compiler.process(element.clauses, **kw)}), 2)"
+
+
+@compiles(MD5, PythonDialects.BigQuery.value)
+def _(element, compiler, **kw):
+    return f"TO_HEX(MD5(CAST({compiler.process(element.clauses, **kw)} AS STRING)))"