openmetadata-ingestion 1.5.11.0__py3-none-any.whl → 1.6.0.0rc1__py3-none-any.whl

metadata/ingestion/source/database/redshift/lineage.py

@@ -31,20 +31,30 @@ workflowConfig:
 """
 
 import traceback
-from typing import Iterator
+from typing import Dict, Iterator, List
 
 from metadata.generated.schema.type.tableQuery import TableQuery
 from metadata.ingestion.source.database.lineage_source import LineageSource
-from metadata.ingestion.source.database.redshift.queries import REDSHIFT_SQL_STATEMENT
+from metadata.ingestion.source.database.redshift.queries import (
+    REDSHIFT_GET_STORED_PROCEDURE_QUERIES,
+    REDSHIFT_SQL_STATEMENT,
+)
 from metadata.ingestion.source.database.redshift.query_parser import (
     RedshiftQueryParserSource,
 )
+from metadata.ingestion.source.database.stored_procedures_mixin import (
+    QueryByProcedure,
+    StoredProcedureLineageMixin,
+)
+from metadata.utils.helpers import get_start_and_end
 from metadata.utils.logger import ingestion_logger
 
 logger = ingestion_logger()
 
 
-class RedshiftLineageSource(RedshiftQueryParserSource, LineageSource):
+class RedshiftLineageSource(
+    RedshiftQueryParserSource, StoredProcedureLineageMixin, LineageSource
+):
     filters = """
         AND (
           querytxt ILIKE '%%create%%table%%as%%select%%'
@@ -86,3 +96,17 @@ class RedshiftLineageSource(RedshiftQueryParserSource, LineageSource):
                         logger.warning(
                             f"Error processing query_dict {query_dict}: {exc}"
                         )
+
+    def get_stored_procedure_queries_dict(self) -> Dict[str, List[QueryByProcedure]]:
+        """
+        Return the dictionary associating stored procedures to the
+        queries they triggered
+        """
+        start, _ = get_start_and_end(self.source_config.queryLogDuration)
+        query = REDSHIFT_GET_STORED_PROCEDURE_QUERIES.format(start_date=start)
+
+        queries_dict = self.procedure_queries_dict(
+            query=query,
+        )
+
+        return queries_dict

metadata/ingestion/source/database/redshift/metadata.py

@@ -14,7 +14,7 @@ Redshift source ingestion
 
 import re
 import traceback
-from typing import Dict, Iterable, List, Optional, Tuple
+from typing import Iterable, List, Optional, Tuple
 
 from sqlalchemy import sql
 from sqlalchemy.dialects.postgresql.base import PGDialect
@@ -71,7 +71,6 @@ from metadata.ingestion.source.database.redshift.models import RedshiftStoredPro
 from metadata.ingestion.source.database.redshift.queries import (
     REDSHIFT_GET_ALL_RELATION_INFO,
     REDSHIFT_GET_DATABASE_NAMES,
-    REDSHIFT_GET_STORED_PROCEDURE_QUERIES,
     REDSHIFT_GET_STORED_PROCEDURES,
     REDSHIFT_LIFE_CYCLE_QUERY,
     REDSHIFT_PARTITION_DETAILS,
@@ -85,17 +84,12 @@ from metadata.ingestion.source.database.redshift.utils import (
     get_table_comment,
     get_view_definition,
 )
-from metadata.ingestion.source.database.stored_procedures_mixin import (
-    QueryByProcedure,
-    StoredProcedureMixin,
-)
 from metadata.utils import fqn
 from metadata.utils.execution_time_tracker import (
     calculate_execution_time,
     calculate_execution_time_generator,
 )
 from metadata.utils.filters import filter_by_database
-from metadata.utils.helpers import get_start_and_end
 from metadata.utils.logger import ingestion_logger
 from metadata.utils.sqlalchemy_utils import (
     get_all_table_comments,
@@ -132,9 +126,7 @@ Inspector.get_all_table_ddls = get_all_table_ddls
 Inspector.get_table_ddl = get_table_ddl
 
 
-class RedshiftSource(
-    LifeCycleQueryMixin, StoredProcedureMixin, CommonDbSourceService, MultiDBSource
-):
+class RedshiftSource(LifeCycleQueryMixin, CommonDbSourceService, MultiDBSource):
     """
     Implements the necessary methods to extract
     Database metadata from Redshift Source
@@ -424,23 +416,6 @@ class RedshiftSource(
                 )
             )
 
-    def get_stored_procedure_queries_dict(self) -> Dict[str, List[QueryByProcedure]]:
-        """
-        Return the dictionary associating stored procedures to the
-        queries they triggered
-        """
-        start, _ = get_start_and_end(self.source_config.queryLogDuration)
-        query = REDSHIFT_GET_STORED_PROCEDURE_QUERIES.format(
-            start_date=start,
-            database_name=self.context.get().database,
-        )
-
-        queries_dict = self.procedure_queries_dict(
-            query=query,
-        )
-
-        return queries_dict
-
     def mark_tables_as_deleted(self):
         """
         Use the current inspector to mark tables as deleted

metadata/ingestion/source/database/redshift/profiler/profiler.py

@@ -0,0 +1,14 @@
+"""Redshift profiler"""
+
+from metadata.ingestion.source.database.redshift.profiler.system import (
+    RedshiftSystemMetricsComputer,
+)
+from metadata.profiler.interface.sqlalchemy.profiler_interface import (
+    SQAProfilerInterface,
+)
+from metadata.profiler.metrics.system.system import SystemMetricsComputer
+
+
+class RedshiftProfiler(SQAProfilerInterface):
+    def initialize_system_metrics_computer(self) -> SystemMetricsComputer:
+        return RedshiftSystemMetricsComputer(session=self.session)

metadata/ingestion/source/database/redshift/profiler/system.py

@@ -0,0 +1,148 @@
+"""
+Imeplemetation for the redshift system metrics source
+"""
+
+from typing import List
+
+from pydantic import TypeAdapter
+
+from metadata.generated.schema.entity.data.table import SystemProfile
+from metadata.ingestion.source.database.redshift.queries import (
+    STL_QUERY,
+    get_query_results,
+)
+from metadata.profiler.metrics.system.dml_operation import DatabaseDMLOperations
+from metadata.profiler.metrics.system.system import (
+    CacheProvider,
+    EmptySystemMetricsSource,
+    SQASessionProvider,
+    SystemMetricsComputer,
+)
+from metadata.utils.logger import profiler_logger
+from metadata.utils.profiler_utils import QueryResult
+from metadata.utils.time_utils import datetime_to_timestamp
+
+logger = profiler_logger()
+
+
+class RedshiftSystemMetricsSource(
+    SQASessionProvider, EmptySystemMetricsSource, CacheProvider
+):
+    """Redshift system metrics source class"""
+
+    def get_inserts(self, **kwargs) -> List[SystemProfile]:
+        database, schema, table = (
+            kwargs.get("database"),
+            kwargs.get("schema"),
+            kwargs.get("table"),
+        )
+        queries = self.get_or_update_cache(
+            f"{database}.{schema}",
+            self._get_insert_queries,
+            database=database,
+            schema=schema,
+        )
+        return get_metric_result(queries, table)
+
+    def get_kwargs(self, **kwargs):
+        table = kwargs.get("table")
+        return {
+            "table": table.__table__.name,
+            "database": self.get_session().get_bind().url.database,
+            "schema": table.__table__.schema,
+        }
+
+    def get_deletes(self, **kwargs) -> List[SystemProfile]:
+        database, schema, table = (
+            kwargs.get("database"),
+            kwargs.get("schema"),
+            kwargs.get("table"),
+        )
+        queries = self.get_or_update_cache(
+            f"{database}.{schema}",
+            self._get_delete_queries,
+            database=database,
+            schema=schema,
+        )
+        return get_metric_result(queries, table)
+
+    def get_updates(self, **kwargs) -> List[SystemProfile]:
+        database = kwargs.get("database")
+        schema = kwargs.get("schema")
+        table = kwargs.get("table")
+        queries = self.get_or_update_cache(
+            f"{database}.{schema}",
+            self._get_update_queries,
+            database=database,
+            schema=schema,
+        )
+        return get_metric_result(queries, table)
+
+    def _get_insert_queries(self, database: str, schema: str) -> List[QueryResult]:
+        insert_query = STL_QUERY.format(
+            alias="si",
+            join_type="LEFT",
+            condition="sd.query is null",
+            database=database,
+            schema=schema,
+        )
+        return get_query_results(
+            super().get_session(),
+            insert_query,
+            DatabaseDMLOperations.INSERT.value,
+        )
+
+    def _get_delete_queries(self, database: str, schema: str) -> List[QueryResult]:
+        delete_query = STL_QUERY.format(
+            alias="sd",
+            join_type="RIGHT",
+            condition="si.query is null",
+            database=database,
+            schema=schema,
+        )
+        return get_query_results(
+            super().get_session(),
+            delete_query,
+            DatabaseDMLOperations.DELETE.value,
+        )
+
+    def _get_update_queries(self, database: str, schema: str) -> List[QueryResult]:
+        update_query = STL_QUERY.format(
+            alias="si",
+            join_type="INNER",
+            condition="sd.query is not null",
+            database=database,
+            schema=schema,
+        )
+        return get_query_results(
+            super().get_session(),
+            update_query,
+            DatabaseDMLOperations.UPDATE.value,
+        )
+
+
+def get_metric_result(ddls: List[QueryResult], table_name: str) -> List[SystemProfile]:
+    """Given query results, retur the metric result
+
+    Args:
+        ddls (List[QueryResult]): list of query results
+        table_name (str): table name
+
+    Returns:
+        List:
+    """
+    return TypeAdapter(List[SystemProfile]).validate_python(
+        [
+            {
+                "timestamp": datetime_to_timestamp(ddl.start_time, milliseconds=True),
+                "operation": ddl.query_type,
+                "rowsAffected": ddl.rows,
+            }
+            for ddl in ddls
+            if ddl.table_name == table_name
+        ]
+    )
+
+
+class RedshiftSystemMetricsComputer(SystemMetricsComputer, RedshiftSystemMetricsSource):
+    pass

metadata/ingestion/source/database/redshift/queries.py

@@ -13,6 +13,13 @@ SQL Queries used during ingestion
 """
 
 import textwrap
+from typing import List
+
+from sqlalchemy import text
+from sqlalchemy.orm.session import Session
+
+from metadata.utils.profiler_utils import QueryResult
+from metadata.utils.time_utils import datetime_to_timestamp
 
 # Not able to use SYS_QUERY_HISTORY here. Few users not getting any results
 REDSHIFT_SQL_STATEMENT = textwrap.dedent(
@@ -236,7 +243,7 @@ REDSHIFT_TEST_PARTITION_DETAILS = "select * from SVV_TABLE_INFO limit 1"
 # hence we are appending "create view <schema>.<table> as " to select query
 # to generate the column level lineage
 REDSHIFT_GET_ALL_RELATIONS = """
-    SELECT
+    (SELECT
         c.relkind,
         n.oid as "schema_oid",
         n.nspname as "schema",
@@ -255,8 +262,9 @@ REDSHIFT_GET_ALL_RELATIONS = """
             JOIN pg_catalog.pg_user u ON u.usesysid = c.relowner
     WHERE c.relkind IN ('r', 'v', 'm', 'S', 'f')
         AND n.nspname !~ '^pg_' {schema_clause} {table_clause}
+        {limit_clause})
     UNION
-    SELECT
+    (SELECT
         'r' AS "relkind",
         s.esoid AS "schema_oid",
         s.schemaname AS "schema",
@@ -272,7 +280,8 @@ REDSHIFT_GET_ALL_RELATIONS = """
         JOIN svv_external_schemas s ON s.schemaname = t.schemaname
         JOIN pg_catalog.pg_user u ON u.usesysid = s.esowner
     where 1 {schema_clause} {table_clause}
-    ORDER BY "relkind", "schema_oid", "schema";
+    ORDER BY "relkind", "schema_oid", "schema"
+    {limit_clause});
     """
 
 
@@ -319,7 +328,7 @@ Q_HISTORY as (
         pid as query_session_id,
         starttime as query_start_time,
         endtime as query_end_time,
-        userid as query_user_name
+        b.usename as query_user_name
     from STL_QUERY q
     join pg_catalog.pg_user b
       on b.usesysid = q.userid
@@ -330,16 +339,16 @@ Q_HISTORY as (
       and userid <> 1
 )
 select
-    sp.procedure_text,
+    trim(sp.procedure_text) procedure_text,
     sp.procedure_start_time,
     sp.procedure_end_time,
-    q.query_text,
+    trim(q.query_text) query_text,
     q.query_type,
-    q.query_database_name,
+    trim(q.query_database_name) query_database_name,
     null as query_schema_name,
     q.query_start_time,
     q.query_end_time,
-    q.query_user_name
+    trim(q.query_user_name) query_user_name
 from SP_HISTORY sp
   join Q_HISTORY q
     on sp.procedure_session_id = q.query_session_id
@@ -373,3 +382,85 @@ WHERE status = 'success'
   and end_time >= '{start_date}'
 ORDER BY end_time DESC
 """
+
+
+STL_QUERY = """
+    with data as (
+        select
+            {alias}.*
+        from 
+            pg_catalog.stl_insert si
+            {join_type} join pg_catalog.stl_delete sd on si.query = sd.query
+        where 
+            {condition}
+    )
+	SELECT
+        SUM(data."rows") AS "rows",
+        sti."database",
+        sti."schema",
+        sti."table",
+        DATE_TRUNC('second', data.starttime) AS starttime
+    FROM
+        data
+        INNER JOIN  pg_catalog.svv_table_info sti ON data.tbl = sti.table_id
+    where
+        sti."database" = '{database}' AND
+       	sti."schema" = '{schema}' AND
+        "rows" != 0 AND
+        DATE(data.starttime) >= CURRENT_DATE - 1
+    GROUP BY 2,3,4,5
+    ORDER BY 5 DESC
+"""
+
+
+def get_query_results(
+    session: Session,
+    query,
+    operation,
+) -> List[QueryResult]:
+    """get query results either from cache or from the database
+
+    Args:
+        session (Session): session
+        query (_type_): query
+        operation (_type_): operation
+
+    Returns:
+        List[QueryResult]:
+    """
+    cursor = session.execute(text(query))
+    results = [
+        QueryResult(
+            database_name=row.database,
+            schema_name=row.schema,
+            table_name=row.table,
+            query_text=None,
+            query_type=operation,
+            start_time=row.starttime,
+            rows=row.rows,
+        )
+        for row in cursor
+    ]
+
+    return results
+
+
+def get_metric_result(ddls: List[QueryResult], table_name: str) -> List:
+    """Given query results, retur the metric result
+
+    Args:
+        ddls (List[QueryResult]): list of query results
+        table_name (str): table name
+
+    Returns:
+        List:
+    """
+    return [
+        {
+            "timestamp": datetime_to_timestamp(ddl.start_time, milliseconds=True),
+            "operation": ddl.query_type,
+            "rowsAffected": ddl.rows,
+        }
+        for ddl in ddls
+        if ddl.table_name == table_name
+    ]

metadata/ingestion/source/database/redshift/service_spec.py

@@ -0,0 +1,14 @@
+from metadata.ingestion.source.database.redshift.lineage import RedshiftLineageSource
+from metadata.ingestion.source.database.redshift.metadata import RedshiftSource
+from metadata.ingestion.source.database.redshift.profiler.profiler import (
+    RedshiftProfiler,
+)
+from metadata.ingestion.source.database.redshift.usage import RedshiftUsageSource
+from metadata.utils.service_spec.default import DefaultDatabaseSpec
+
+ServiceSpec = DefaultDatabaseSpec(
+    metadata_source_class=RedshiftSource,
+    lineage_source_class=RedshiftLineageSource,
+    usage_source_class=RedshiftUsageSource,
+    profiler_class=RedshiftProfiler,
+)

metadata/ingestion/source/database/redshift/utils.py

@@ -244,14 +244,14 @@ def _update_column_info(  # pylint: disable=too-many-arguments
     return column_info
 
 
+# pylint: disable=unused-argument
 def _update_coltype(coltype, args, kwargs, attype, name, is_array):
     if coltype:
         coltype = coltype(*args, **kwargs)
         if is_array:
             coltype = ischema_names["_array"](coltype)
     else:
-        util.warn(f"Did not recognize type '{attype}' of column '{name}'")
-        coltype = sqltypes.NULLTYPE
+        coltype = create_sqlalchemy_type("UNKNOWN")
     return coltype
 
 
@@ -386,7 +386,7 @@ def _get_all_relation_info(self, connection, **kw):  # pylint: disable=unused-ar
     result = connection.execute(
         sa.text(
             REDSHIFT_GET_ALL_RELATIONS.format(
-                schema_clause=schema_clause, table_clause=table_clause
+                schema_clause=schema_clause, table_clause=table_clause, limit_clause=""
             )
         )
     )

metadata/ingestion/source/database/salesforce/metadata.py

@@ -12,18 +12,16 @@
 Salesforce source ingestion
 """
 import traceback
-from typing import Any, Iterable, Optional, Tuple, Union
+from typing import Any, Iterable, Optional, Tuple
 
 from metadata.generated.schema.api.data.createDatabase import CreateDatabaseRequest
 from metadata.generated.schema.api.data.createDatabaseSchema import (
     CreateDatabaseSchemaRequest,
 )
-from metadata.generated.schema.api.data.createQuery import CreateQueryRequest
 from metadata.generated.schema.api.data.createStoredProcedure import (
     CreateStoredProcedureRequest,
 )
 from metadata.generated.schema.api.data.createTable import CreateTableRequest
-from metadata.generated.schema.api.lineage.addLineage import AddLineageRequest
 from metadata.generated.schema.entity.data.database import Database
 from metadata.generated.schema.entity.data.databaseSchema import DatabaseSchema
 from metadata.generated.schema.entity.data.table import (
@@ -324,9 +322,6 @@ class SalesforceSource(DatabaseServiceSource):
             return DataType.VARCHAR.value
         return DataType.UNKNOWN.value
 
-    def yield_view_lineage(self) -> Iterable[Either[AddLineageRequest]]:
-        yield from []
-
     def yield_tag(
         self, schema_name: str
     ) -> Iterable[Either[OMetaTagAndClassification]]:
@@ -343,12 +338,6 @@ class SalesforceSource(DatabaseServiceSource):
     def get_stored_procedure_queries(self) -> Iterable[QueryByProcedure]:
         """Not Implemented"""
 
-    def yield_procedure_lineage_and_queries(
-        self,
-    ) -> Iterable[Either[Union[AddLineageRequest, CreateQueryRequest]]]:
-        """Not Implemented"""
-        yield from []
-
     def standardize_table_name(  # pylint: disable=unused-argument
         self, schema: str, table: str
     ) -> str:

metadata/ingestion/source/database/salesforce/service_spec.py

@@ -0,0 +1,4 @@
+from metadata.ingestion.source.database.salesforce.metadata import SalesforceSource
+from metadata.utils.service_spec.default import DefaultDatabaseSpec
+
+ServiceSpec = DefaultDatabaseSpec(metadata_source_class=SalesforceSource)