openmetadata-ingestion 1.2.5.1__py3-none-any.whl → 1.3.0.0__py3-none-any.whl

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

@@ -11,11 +11,20 @@
 """
 Postgres usage module
 """
+import traceback
+from datetime import datetime
+from typing import Iterable
+
+from metadata.generated.schema.type.tableQuery import TableQueries, TableQuery
+from metadata.ingestion.source.connections import get_connection
 from metadata.ingestion.source.database.postgres.queries import POSTGRES_SQL_STATEMENT
 from metadata.ingestion.source.database.postgres.query_parser import (
     PostgresQueryParserSource,
 )
 from metadata.ingestion.source.database.usage_source import UsageSource
+from metadata.utils.logger import ingestion_logger
+
+logger = ingestion_logger()
 
 
 class PostgresUsageSource(PostgresQueryParserSource, UsageSource):
@@ -25,3 +34,35 @@ class PostgresUsageSource(PostgresQueryParserSource, UsageSource):
 
     sql_stmt = POSTGRES_SQL_STATEMENT
     filters = ""
+
+    def process_table_query(self) -> Iterable[TableQueries]:
+        """
+        Process Query
+        """
+        try:
+            with get_connection(self.service_connection).connect() as conn:
+                rows = conn.execute(self.get_sql_statement())
+                queries = []
+                for row in rows:
+                    row = dict(row)
+                    try:
+                        queries.append(
+                            TableQuery(
+                                query=row["query_text"],
+                                userName=row["usename"],
+                                analysisDate=datetime.now(),
+                                aborted=self.get_aborted_status(row),
+                                databaseName=self.get_database_name(row),
+                                serviceName=self.config.serviceName,
+                                databaseSchema=self.get_schema_name(row),
+                                duration=row.get("duration"),
+                            )
+                        )
+                    except Exception as err:
+                        logger.debug(traceback.format_exc())
+                        logger.error(str(err))
+            if queries:
+                yield TableQueries(queries=queries)
+        except Exception as err:
+            logger.error(f"Source usage processing error - {err}")
+            logger.debug(traceback.format_exc())

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

@@ -28,11 +28,13 @@ from metadata.ingestion.source.database.postgres.queries import (
     POSTGRES_GET_SERVER_VERSION,
     POSTGRES_SQL_COLUMNS,
     POSTGRES_TABLE_COMMENTS,
+    POSTGRES_TABLE_OWNERS,
     POSTGRES_VIEW_DEFINITIONS,
 )
 from metadata.utils.logger import utils_logger
 from metadata.utils.sqlalchemy_utils import (
     get_table_comment_wrapper,
+    get_table_owner_wrapper,
     get_view_definition_wrapper,
 )
 
@@ -41,6 +43,38 @@ logger = utils_logger()
 OLD_POSTGRES_VERSION = "13.0"
 
 
+def get_etable_owner(
+    self, connection, table_name=None, schema=None
+):  # pylint: disable=unused-argument
+    """Return all owners.
+
+    :param schema: Optional, retrieve names from a non-default schema.
+        For special quoting, use :class:`.quoted_name`.
+
+    """
+
+    with self._operation_context() as conn:
+        return self.dialect.get_table_owner(
+            connection=conn,
+            query=POSTGRES_TABLE_OWNERS,
+            table_name=table_name,
+            schema=schema,
+        )
+
+
+@reflection.cache
+def get_table_owner(
+    self, connection, table_name, schema=None, **kw
+):  # pylint: disable=unused-argument
+    return get_table_owner_wrapper(
+        self,
+        connection=connection,
+        query=POSTGRES_TABLE_OWNERS,
+        table_name=table_name,
+        schema=schema,
+    )
+
+
 @reflection.cache
 def get_table_comment(
     self, connection, table_name, schema=None, **kw

metadata/ingestion/source/database/query_parser_source.py

@@ -13,7 +13,7 @@ Query Parser Source module. Parent class for Lineage & Usage workflows
 """
 from abc import ABC, abstractmethod
 from datetime import datetime
-from typing import Iterator
+from typing import Iterator, Optional
 
 from metadata.generated.schema.metadataIngestion.workflow import (
     Source as WorkflowSource,
@@ -57,6 +57,10 @@ class QueryParserSource(Source, ABC):
         self.start, self.end = get_start_and_end(self.source_config.queryLogDuration)
         self.engine = get_connection(self.service_connection) if get_engine else None
 
+    @property
+    def name(self) -> str:
+        return self.service_connection.type.name
+
     def prepare(self):
         """By default, there's nothing to prepare"""
 
@@ -90,7 +94,9 @@ class QueryParserSource(Source, ABC):
         )
 
     def check_life_cycle_query(
-        self, query_type: str  # pylint: disable=unused-argument
+        self,
+        query_type: Optional[str],  # pylint: disable=unused-argument
+        query_text: Optional[str],  # pylint: disable=unused-argument
     ) -> bool:
         """
         returns true if query is to be used for life cycle processing.

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

@@ -32,7 +32,6 @@ from metadata.generated.schema.entity.data.storedProcedure import (
 )
 from metadata.generated.schema.entity.data.table import (
     ConstraintType,
-    EntityName,
     IntervalType,
     TableConstraint,
     TablePartition,
@@ -41,16 +40,23 @@ from metadata.generated.schema.entity.data.table import (
 from metadata.generated.schema.entity.services.connections.database.redshiftConnection import (
     RedshiftConnection,
 )
+from metadata.generated.schema.entity.services.ingestionPipelines.status import (
+    StackTraceError,
+)
 from metadata.generated.schema.metadataIngestion.workflow import (
     Source as WorkflowSource,
 )
-from metadata.ingestion.api.models import Either, StackTraceError
+from metadata.generated.schema.type.basic import EntityName
+from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import InvalidSourceException
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.database.common_db_source import (
     CommonDbSourceService,
     TableNameAndType,
 )
+from metadata.ingestion.source.database.life_cycle_query_mixin import (
+    LifeCycleQueryMixin,
+)
 from metadata.ingestion.source.database.multi_db_source import MultiDBSource
 from metadata.ingestion.source.database.redshift.models import RedshiftStoredProcedure
 from metadata.ingestion.source.database.redshift.queries import (
@@ -58,6 +64,7 @@ from metadata.ingestion.source.database.redshift.queries import (
     REDSHIFT_GET_DATABASE_NAMES,
     REDSHIFT_GET_STORED_PROCEDURE_QUERIES,
     REDSHIFT_GET_STORED_PROCEDURES,
+    REDSHIFT_LIFE_CYCLE_QUERY,
     REDSHIFT_PARTITION_DETAILS,
 )
 from metadata.ingestion.source.database.redshift.utils import (
@@ -103,7 +110,9 @@ RedshiftDialect._get_all_relation_info = (  # pylint: disable=protected-access
 )
 
 
-class RedshiftSource(StoredProcedureMixin, CommonDbSourceService, MultiDBSource):
+class RedshiftSource(
+    LifeCycleQueryMixin, StoredProcedureMixin, CommonDbSourceService, MultiDBSource
+):
     """
     Implements the necessary methods to extract
     Database metadata from Redshift Source
@@ -112,6 +121,7 @@ class RedshiftSource(StoredProcedureMixin, CommonDbSourceService, MultiDBSource)
     def __init__(self, config, metadata):
         super().__init__(config, metadata)
         self.partition_details = {}
+        self.life_cycle_query = REDSHIFT_LIFE_CYCLE_QUERY
 
     @classmethod
     def create(cls, config_dict, metadata: OpenMetadata):
@@ -282,7 +292,7 @@ class RedshiftSource(StoredProcedureMixin, CommonDbSourceService, MultiDBSource)
                 left=StackTraceError(
                     name=stored_procedure.name,
                     error=f"Error yielding Stored Procedure [{stored_procedure.name}] due to [{exc}]",
-                    stack_trace=traceback.format_exc(),
+                    stackTrace=traceback.format_exc(),
                 )
             )
 

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

@@ -300,7 +300,6 @@ REDSHIFT_GET_STORED_PROCEDURE_QUERIES = textwrap.dedent(
     """
 with SP_HISTORY as (
     select
-        query as procedure_id,
         querytxt as procedure_text,
         starttime as procedure_start_time,
         endtime as procedure_end_time,
@@ -311,7 +310,6 @@ with SP_HISTORY as (
 ),
 Q_HISTORY as (
     select
-        query as query_id,
         querytxt as query_text,
         case
             when querytxt ilike '%%MERGE%%' then 'MERGE'
@@ -334,11 +332,9 @@ Q_HISTORY as (
       and userid <> 1
 )
 select
-    sp.procedure_id,
     sp.procedure_text,
     sp.procedure_start_time,
     sp.procedure_end_time,
-    q.query_id,
     q.query_text,
     q.query_type,
     q.query_database_name,
@@ -354,3 +350,13 @@ from SP_HISTORY sp
 order by procedure_start_time DESC
     """
 )
+
+REDSHIFT_LIFE_CYCLE_QUERY = textwrap.dedent(
+    """
+select "table" as table_name,
+create_time as created_at
+from pg_catalog.svv_table_info o
+where o.schema = '{schema_name}'
+and o.database = '{database_name}'
+"""
+)

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

@@ -11,8 +11,10 @@
 """
 Redshift usage module
 """
+import re
 from abc import ABC
 from datetime import datetime
+from typing import Optional
 
 from metadata.generated.schema.entity.services.connections.database.redshiftConnection import (
     RedshiftConnection,
@@ -55,3 +57,17 @@ class RedshiftQueryParserSource(QueryParserSource, ABC):
             filters=self.get_filters(),
             result_limit=self.source_config.resultLimit,
         )
+
+    def check_life_cycle_query(
+        self, query_type: Optional[str], query_text: Optional[str]
+    ) -> bool:
+        """
+        returns true if query is to be used for life cycle processing.
+
+        Override if we have specific parameters
+        """
+        create_pattern = re.compile(r".*\s*CREATE", re.IGNORECASE)
+        insert_pattern = re.compile(r".*\s*INSERT", re.IGNORECASE)
+        if re.match(create_pattern, query_text) or re.match(insert_pattern, query_text):
+            return True
+        return False

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

@@ -23,8 +23,6 @@ class RedshiftUsageSource(RedshiftQueryParserSource, UsageSource):
         AND querytxt NOT ILIKE 'fetch%%'
         AND querytxt NOT ILIKE 'padb_fetch_sample:%%'
         AND querytxt NOT ILIKE 'Undoing%%transactions%%on%%table%%with%%current%%xid%%'
-        AND querytxt NOT ILIKE '%%create%%table%%as%%select%%'
-        AND querytxt NOT ILIKE '%%insert%%'
     """
 
     sql_stmt = REDSHIFT_SQL_STATEMENT

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

@@ -36,13 +36,16 @@ from metadata.generated.schema.entity.data.table import (
 from metadata.generated.schema.entity.services.connections.database.salesforceConnection import (
     SalesforceConnection,
 )
+from metadata.generated.schema.entity.services.ingestionPipelines.status import (
+    StackTraceError,
+)
 from metadata.generated.schema.metadataIngestion.databaseServiceMetadataPipeline import (
     DatabaseServiceMetadataPipeline,
 )
 from metadata.generated.schema.metadataIngestion.workflow import (
     Source as WorkflowSource,
 )
-from metadata.ingestion.api.models import Either, StackTraceError
+from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import InvalidSourceException
 from metadata.ingestion.models.ometa_classification import OMetaTagAndClassification
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
@@ -148,6 +151,7 @@ class SalesforceSource(DatabaseServiceSource):
         :return: tables or views, depending on config
         """
         schema_name = self.context.database_schema
+
         try:
             if self.service_connection.sobjectName:
                 table_name = self.standardize_table_name(
@@ -184,10 +188,34 @@ class SalesforceSource(DatabaseServiceSource):
                 StackTraceError(
                     name=schema_name,
                     error=f"Unexpected exception for schema name [{schema_name}]: {exc}",
-                    stack_trace=traceback.format_exc(),
+                    stackTrace=traceback.format_exc(),
                 )
             )
 
+    def get_table_description(self, table_name: str) -> Optional[str]:
+        """
+        Method to get the table description for salesforce with Tooling API
+        """
+        try:
+            result = self.client.toolingexecute(
+                f"query/?q=SELECT+Description+FROM+EntityDefinition+WHERE+QualifiedApiName='{table_name}'"
+            )
+            return result["records"][0]["Description"]
+        except KeyError as err:
+            logger.warning(
+                f"Unable to get required key from Tooling API response for table [{table_name}]: {err}"
+            )
+        except IndexError as err:
+            logger.warning(
+                f"Unable to get row for table [{table_name}] from EntityDefinition: {err}"
+            )
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.warning(
+                f"Unable to get description with Tooling API for table [{table_name}]: {exc}"
+            )
+        return None
+
     def yield_table(
         self, table_name_and_type: Tuple[str, str]
     ) -> Iterable[Either[CreateTableRequest]]:
@@ -206,6 +234,7 @@ class SalesforceSource(DatabaseServiceSource):
             table_request = CreateTableRequest(
                 name=table_name,
                 tableType=table_type,
+                description=self.get_table_description(table_name),
                 columns=columns,
                 tableConstraints=table_constraints,
                 databaseSchema=fqn.build(
@@ -226,7 +255,7 @@ class SalesforceSource(DatabaseServiceSource):
                 left=StackTraceError(
                     name=table_name,
                     error=f"Unexpected exception for table [{table_name}]: {exc}",
-                    stack_trace=traceback.format_exc(),
+                    stackTrace=traceback.format_exc(),
                 )
             )
 

metadata/ingestion/source/database/sample_data.py

@@ -51,7 +51,13 @@ from metadata.generated.schema.api.lineage.addLineage import AddLineageRequest
 from metadata.generated.schema.api.teams.createRole import CreateRoleRequest
 from metadata.generated.schema.api.teams.createTeam import CreateTeamRequest
 from metadata.generated.schema.api.teams.createUser import CreateUserRequest
+from metadata.generated.schema.api.tests.createCustomMetric import (
+    CreateCustomMetricRequest,
+)
 from metadata.generated.schema.api.tests.createTestCase import CreateTestCaseRequest
+from metadata.generated.schema.api.tests.createTestCaseResolutionStatus import (
+    CreateTestCaseResolutionStatus,
+)
 from metadata.generated.schema.api.tests.createTestSuite import CreateTestSuiteRequest
 from metadata.generated.schema.entity.data.container import Container
 from metadata.generated.schema.entity.data.dashboard import Dashboard
@@ -65,7 +71,10 @@ from metadata.generated.schema.entity.data.mlmodel import (
     MlStore,
 )
 from metadata.generated.schema.entity.data.pipeline import Pipeline, PipelineStatus
-from metadata.generated.schema.entity.data.storedProcedure import StoredProcedureCode
+from metadata.generated.schema.entity.data.storedProcedure import (
+    StoredProcedure,
+    StoredProcedureCode,
+)
 from metadata.generated.schema.entity.data.table import (
     ColumnProfile,
     SystemProfile,
@@ -90,7 +99,9 @@ from metadata.generated.schema.entity.teams.user import User
 from metadata.generated.schema.metadataIngestion.workflow import (
     Source as WorkflowSource,
 )
+from metadata.generated.schema.tests.assigned import Assigned
 from metadata.generated.schema.tests.basic import TestCaseResult, TestResultValue
+from metadata.generated.schema.tests.resolved import Resolved, TestCaseFailureReasonType
 from metadata.generated.schema.tests.testCase import TestCase, TestCaseParameterValue
 from metadata.generated.schema.tests.testSuite import TestSuite
 from metadata.generated.schema.type.basic import Timestamp
@@ -107,6 +118,7 @@ from metadata.ingestion.models.pipeline_status import OMetaPipelineStatus
 from metadata.ingestion.models.profile_data import OMetaTableProfileSampleData
 from metadata.ingestion.models.tests_data import (
     OMetaLogicalTestSuiteSample,
+    OMetaTestCaseResolutionStatus,
     OMetaTestCaseResultsSample,
     OMetaTestCaseSample,
     OMetaTestSuiteSample,
@@ -117,7 +129,7 @@ from metadata.parsers.schema_parsers import (
     InvalidSchemaTypeException,
     schema_parser_config_registry,
 )
-from metadata.utils import fqn
+from metadata.utils import entity_link, fqn
 from metadata.utils.constants import UTF_8
 from metadata.utils.fqn import FQN_SEPARATOR
 from metadata.utils.helpers import get_standard_chart_type
@@ -555,6 +567,7 @@ class SampleDataSource(
         yield from self.ingest_test_suite()
         yield from self.ingest_test_case()
         yield from self.ingest_test_case_results()
+        yield from self.ingest_incidents()
         yield from self.ingest_logical_test_suite()
         yield from self.ingest_data_insights()
         yield from self.ingest_life_cycle()
@@ -714,6 +727,21 @@ class SampleDataSource(
                     ),
                 )
 
+            if table.get("customMetrics"):
+                for custom_metric in table["customMetrics"]:
+                    self.metadata.create_or_update_custom_metric(
+                        CreateCustomMetricRequest(**custom_metric),
+                        table_entity.id.__root__,
+                    )
+
+            for column in table.get("columns"):
+                if column.get("customMetrics"):
+                    for custom_metric in column["customMetrics"]:
+                        self.metadata.create_or_update_custom_metric(
+                            CreateCustomMetricRequest(**custom_metric),
+                            table_entity.id.__root__,
+                        )
+
     def ingest_stored_procedures(self) -> Iterable[Either[Entity]]:
         """Ingest Sample Stored Procedures"""
 
@@ -770,6 +798,34 @@ class SampleDataSource(
 
             yield Either(right=stored_procedure)
 
+        # Create table and stored procedure lineage
+        for lineage_entities in self.stored_procedures["lineage"]:
+            from_table = self.metadata.get_by_name(
+                entity=Table, fqn=lineage_entities["from_table_fqn"]
+            )
+            stored_procedure_entity = self.metadata.get_by_name(
+                entity=StoredProcedure, fqn=lineage_entities["stored_procedure_fqn"]
+            )
+            to_table = self.metadata.get_by_name(
+                entity=Table, fqn=lineage_entities["to_table_fqn"]
+            )
+            yield Either(
+                right=AddLineageRequest(
+                    edge=EntitiesEdge(
+                        fromEntity=EntityReference(
+                            id=from_table.id.__root__, type="table"
+                        ),
+                        toEntity=EntityReference(id=to_table.id.__root__, type="table"),
+                        lineageDetails=LineageDetails(
+                            pipeline=EntityReference(
+                                id=stored_procedure_entity.id.__root__,
+                                type="storedProcedure",
+                            )
+                        ),
+                    )
+                )
+            )
+
     def ingest_topics(self) -> Iterable[CreateTopicRequest]:
         """
         Ingest Sample Topics
@@ -999,9 +1055,9 @@ class SampleDataSource(
         for pipeline in self.pipelines["pipelines"]:
             owner = None
             if pipeline.get("owner"):
-                user = self.metadata.get_user_by_email(email=pipeline.get("owner"))
-                if user:
-                    owner = EntityReference(id=user.id.__root__, type="user")
+                owner = self.metadata.get_reference_by_email(
+                    email=pipeline.get("owner")
+                )
             pipeline_ev = CreatePipelineRequest(
                 name=pipeline["name"],
                 displayName=pipeline["displayName"],
@@ -1251,6 +1307,7 @@ class SampleDataSource(
                             rowCount=profile["rowCount"],
                             createDateTime=profile.get("createDateTime"),
                             sizeInByte=profile.get("sizeInByte"),
+                            customMetrics=profile.get("customMetrics"),
                             timestamp=int(
                                 (
                                     datetime.now(tz=timezone.utc) - timedelta(days=days)
@@ -1331,6 +1388,7 @@ class SampleDataSource(
             )
 
     def ingest_test_case(self) -> Iterable[Either[OMetaTestCaseSample]]:
+        """Ingest test cases"""
         for test_suite in self.tests_suites["tests"]:
             suite = self.metadata.get_by_name(
                 fqn=test_suite["testSuiteName"], entity=TestSuite
@@ -1347,10 +1405,66 @@ class SampleDataSource(
                             TestCaseParameterValue(**param_values)
                             for param_values in test_case["parameterValues"]
                         ],
-                    )
+                    )  # type: ignore
                 )
                 yield Either(right=test_case_req)
 
+    def ingest_incidents(self) -> Iterable[Either[OMetaTestCaseResolutionStatus]]:
+        """
+        Ingest incidents after the first test failures have been added.
+
+        The test failure already creates the incident with NEW, so we
+        start always from ACK in the sample flows.
+        """
+        for test_suite in self.tests_suites["tests"]:
+            for test_case in test_suite["testCases"]:
+                test_case_fqn = f"{entity_link.get_table_or_column_fqn(test_case['entityLink'])}.{test_case['name']}"
+
+                for _, resolutions in test_case["resolutions"].items():
+                    for resolution in resolutions:
+                        create_test_case_resolution = CreateTestCaseResolutionStatus(
+                            testCaseResolutionStatusType=resolution[
+                                "testCaseResolutionStatusType"
+                            ],
+                            testCaseReference=test_case_fqn,
+                            severity=resolution["severity"],
+                        )
+
+                        if resolution["testCaseResolutionStatusType"] == "Assigned":
+                            user: User = self.metadata.get_by_name(
+                                User, fqn=resolution["assignee"]
+                            )
+                            create_test_case_resolution.testCaseResolutionStatusDetails = Assigned(
+                                assignee=EntityReference(
+                                    id=user.id.__root__,
+                                    type="user",
+                                    name=user.name.__root__,
+                                    fullyQualifiedName=user.fullyQualifiedName.__root__,
+                                )
+                            )
+                        if resolution["testCaseResolutionStatusType"] == "Resolved":
+                            user: User = self.metadata.get_by_name(
+                                User, fqn=resolution["resolver"]
+                            )
+                            create_test_case_resolution.testCaseResolutionStatusDetails = Resolved(
+                                resolvedBy=EntityReference(
+                                    id=user.id.__root__,
+                                    type="user",
+                                    name=user.name.__root__,
+                                    fullyQualifiedName=user.fullyQualifiedName.__root__,
+                                ),
+                                testCaseFailureReason=random.choice(
+                                    list(TestCaseFailureReasonType)
+                                ),
+                                testCaseFailureComment="Resolution comment",
+                            )
+
+                        yield Either(
+                            right=OMetaTestCaseResolutionStatus(
+                                test_case_resolution=create_test_case_resolution
+                            )
+                        )
+
     def ingest_test_case_results(self) -> Iterable[Either[OMetaTestCaseResultsSample]]:
         """Iterate over all the testSuite and testCase and ingest them"""
         for test_case_results in self.tests_case_results["testCaseResults"]: