openmetadata-ingestion 1.5.13.0__py3-none-any.whl → 1.6.0.0rc2__py3-none-any.whl

metadata/ingestion/source/database/oracle/metadata.py

@@ -12,7 +12,7 @@
 # pylint: disable=protected-access
 """Oracle source module"""
 import traceback
-from typing import Dict, Iterable, List, Optional
+from typing import Iterable, Optional
 
 from sqlalchemy.dialects.oracle.base import INTERVAL, OracleDialect, ischema_names
 from sqlalchemy.engine import Inspector
@@ -49,7 +49,6 @@ from metadata.ingestion.source.database.oracle.models import (
     OracleStoredProcedure,
 )
 from metadata.ingestion.source.database.oracle.queries import (
-    ORACLE_GET_STORED_PROCEDURE_QUERIES,
     ORACLE_GET_STORED_PROCEDURES,
 )
 from metadata.ingestion.source.database.oracle.utils import (
@@ -65,12 +64,7 @@ from metadata.ingestion.source.database.oracle.utils import (
     get_view_names,
     get_view_names_dialect,
 )
-from metadata.ingestion.source.database.stored_procedures_mixin import (
-    QueryByProcedure,
-    StoredProcedureMixin,
-)
 from metadata.utils import fqn
-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,
@@ -109,7 +103,7 @@ Inspector.get_table_ddl = get_table_ddl
 OracleDialect._get_constraint_data = _get_constraint_data
 
 
-class OracleSource(StoredProcedureMixin, CommonDbSourceService):
+class OracleSource(CommonDbSourceService):
     """
     Implements the necessary methods to extract
     Database metadata from Oracle Source
@@ -250,18 +244,3 @@ class OracleSource(StoredProcedureMixin, CommonDbSourceService):
                     stackTrace=traceback.format_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 = ORACLE_GET_STORED_PROCEDURE_QUERIES.format(
-            start_date=start,
-        )
-        queries_dict = self.procedure_queries_dict(
-            query=query,
-        )
-
-        return queries_dict

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

@@ -0,0 +1,10 @@
+from metadata.ingestion.source.database.oracle.lineage import OracleLineageSource
+from metadata.ingestion.source.database.oracle.metadata import OracleSource
+from metadata.ingestion.source.database.oracle.usage import OracleUsageSource
+from metadata.utils.service_spec.default import DefaultDatabaseSpec
+
+ServiceSpec = DefaultDatabaseSpec(
+    metadata_source_class=OracleSource,
+    lineage_source_class=OracleLineageSource,
+    usage_source_class=OracleUsageSource,
+)

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

@@ -0,0 +1,46 @@
+#  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.
+"""
+Pinotdb lineage module
+"""
+from typing import Optional
+
+from metadata.generated.schema.entity.services.connections.database.pinotDBConnection import (
+    PinotDBConnection,
+)
+from metadata.generated.schema.metadataIngestion.workflow import (
+    Source as WorkflowSource,
+)
+from metadata.ingestion.api.steps import InvalidSourceException
+from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.ingestion.source.database.lineage_source import LineageSource
+from metadata.utils.logger import ingestion_logger
+
+logger = ingestion_logger()
+
+
+class PinotdbLineageSource(LineageSource):
+    """
+    Pinotdb lineage source implements view lineage
+    """
+
+    @classmethod
+    def create(
+        cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
+    ):
+        """Create class instance"""
+        config: WorkflowSource = WorkflowSource.model_validate(config_dict)
+        connection: PinotDBConnection = config.serviceConnection.root.config
+        if not isinstance(connection, PinotDBConnection):
+            raise InvalidSourceException(
+                f"Expected PinotDBConnection, but got {connection}"
+            )
+        return cls(config, metadata)

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

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

metadata/ingestion/source/database/postgres/metadata.py

@@ -13,7 +13,7 @@ Postgres source module
 """
 import traceback
 from collections import namedtuple
-from typing import Dict, Iterable, List, Optional, Tuple
+from typing import Iterable, Optional, Tuple
 
 from sqlalchemy import String as SqlAlchemyString
 from sqlalchemy import sql
@@ -67,14 +67,11 @@ from metadata.ingestion.source.database.postgres.utils import (
     get_columns,
     get_etable_owner,
     get_foreign_keys,
+    get_json_fields_and_type,
     get_table_comment,
     get_table_owner,
     get_view_definition,
 )
-from metadata.ingestion.source.database.stored_procedures_mixin import (
-    QueryByProcedure,
-    StoredProcedureMixin,
-)
 from metadata.utils import fqn
 from metadata.utils.filters import filter_by_database
 from metadata.utils.importer import import_side_effects
@@ -150,11 +147,12 @@ PGDialect.ischema_names = ischema_names
 Inspector.get_all_table_ddls = get_all_table_ddls
 Inspector.get_table_ddl = get_table_ddl
 Inspector.get_table_owner = get_etable_owner
+Inspector.get_json_fields_and_type = get_json_fields_and_type
 
 PGDialect.get_foreign_keys = get_foreign_keys
 
 
-class PostgresSource(CommonDbSourceService, MultiDBSource, StoredProcedureMixin):
+class PostgresSource(CommonDbSourceService, MultiDBSource):
     """
     Implements the necessary methods to extract
     Database metadata from Postgres Source
@@ -362,10 +360,3 @@ class PostgresSource(CommonDbSourceService, MultiDBSource, StoredProcedureMixin)
                     stackTrace=traceback.format_exc(),
                 )
             )
-
-    def get_stored_procedure_queries_dict(self) -> Dict[str, List[QueryByProcedure]]:
-        """
-        Return the dictionary associating stored procedures to the
-        queries they triggered
-        """
-        return {}

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

@@ -0,0 +1,10 @@
+from metadata.ingestion.source.database.postgres.lineage import PostgresLineageSource
+from metadata.ingestion.source.database.postgres.metadata import PostgresSource
+from metadata.ingestion.source.database.postgres.usage import PostgresUsageSource
+from metadata.utils.service_spec.default import DefaultDatabaseSpec
+
+ServiceSpec = DefaultDatabaseSpec(
+    metadata_source_class=PostgresSource,
+    lineage_source_class=PostgresLineageSource,
+    usage_source_class=PostgresUsageSource,
+)

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

@@ -13,6 +13,7 @@
 """
 Postgres SQLAlchemy util methods
 """
+import json
 import re
 import traceback
 from typing import Dict, Optional, Tuple
@@ -23,15 +24,18 @@ from sqlalchemy.dialects.postgresql.base import ENUM
 from sqlalchemy.engine import reflection
 from sqlalchemy.sql import sqltypes
 
+from metadata.generated.schema.entity.data.table import Column
 from metadata.ingestion.source.database.postgres.queries import (
     POSTGRES_COL_IDENTITY,
     POSTGRES_FETCH_FK,
+    POSTGRES_GET_JSON_FIELDS,
     POSTGRES_GET_SERVER_VERSION,
     POSTGRES_SQL_COLUMNS,
     POSTGRES_TABLE_COMMENTS,
     POSTGRES_TABLE_OWNERS,
     POSTGRES_VIEW_DEFINITIONS,
 )
+from metadata.parsers.json_schema_parser import parse_json_schema
 from metadata.utils.logger import utils_logger
 from metadata.utils.sqlalchemy_utils import (
     get_table_comment_wrapper,
@@ -186,6 +190,28 @@ def get_table_comment(
     )
 
 
+@reflection.cache
+def get_json_fields_and_type(
+    self, table_name, column_name, schema=None, **kw
+):  # pylint: disable=unused-argument
+    try:
+        query = POSTGRES_GET_JSON_FIELDS.format(
+            table_name=table_name, column_name=column_name
+        )
+        cursor = self.engine.execute(query)
+        result = cursor.fetchone()
+        if result:
+            parsed_column = parse_json_schema(json.dumps(result[0]), Column)
+            if parsed_column:
+                return parsed_column[0].children
+    except Exception as err:
+        logger.warning(
+            f"Unable to parse the json fields for {table_name}.{column_name} - {err}"
+        )
+        logger.debug(traceback.format_exc())
+    return None
+
+
 @reflection.cache
 def get_columns(  # pylint: disable=too-many-locals
     self, connection, table_name, schema=None, **kw

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

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

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

@@ -0,0 +1,9 @@
+from metadata.ingestion.source.database.query.lineage import QueryLogLineageSource
+from metadata.ingestion.source.database.query.usage import QueryLogUsageSource
+from metadata.utils.service_spec import BaseSpec
+
+ServiceSpec = BaseSpec(
+    metadata_source_class="not.implemented",
+    lineage_source_class=QueryLogLineageSource,
+    usage_source_class=QueryLogUsageSource,
+)

metadata/ingestion/source/database/query_parser_source.py

@@ -58,6 +58,7 @@ class QueryParserSource(Source, ABC):
         super().__init__()
         self.config = config
         self.metadata = metadata
+        self.service_name = self.config.serviceName
         self.service_connection = self.config.serviceConnection.root.config
         connection_type = self.service_connection.type.value
         self.dialect = ConnectionTypeDialectMapper.dialect_of(connection_type)

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

@@ -16,6 +16,7 @@ from functools import partial
 from typing import Optional
 
 from sqlalchemy.engine import Engine
+from sqlalchemy.sql import text
 
 from metadata.generated.schema.entity.automations.workflow import (
     Workflow as AutomationWorkflow,
@@ -41,6 +42,7 @@ from metadata.ingestion.connections.test_connections import (
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.connections import kill_active_connections
 from metadata.ingestion.source.database.redshift.queries import (
+    REDSHIFT_GET_ALL_RELATIONS,
     REDSHIFT_GET_DATABASE_NAMES,
     REDSHIFT_TEST_GET_QUERIES,
     REDSHIFT_TEST_PARTITION_DETAILS,
@@ -70,6 +72,11 @@ def test_connection(
     Test connection. This can be executed either as part
     of a metadata workflow or during an Automation Workflow
     """
+    table_and_view_query = text(
+        REDSHIFT_GET_ALL_RELATIONS.format(
+            schema_clause="", table_clause="", limit_clause="LIMIT 1"
+        )
+    )
 
     def test_get_queries_permissions(engine_: Engine):
         """Check if we have the right permissions to list queries"""
@@ -83,8 +90,8 @@ def test_connection(
     test_fn = {
         "CheckAccess": partial(test_connection_engine_step, engine),
         "GetSchemas": partial(execute_inspector_func, engine, "get_schema_names"),
-        "GetTables": partial(execute_inspector_func, engine, "get_table_names"),
-        "GetViews": partial(execute_inspector_func, engine, "get_view_names"),
+        "GetTables": partial(test_query, statement=table_and_view_query, engine=engine),
+        "GetViews": partial(test_query, statement=table_and_view_query, engine=engine),
         "GetQueries": partial(test_get_queries_permissions, engine),
         "GetDatabases": partial(
             test_query, statement=REDSHIFT_GET_DATABASE_NAMES, engine=engine

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