openmetadata-ingestion 1.2.4.0__py3-none-any.whl → 1.3.2.0rc1__py3-none-any.whl

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

@@ -11,11 +11,9 @@
 """
 Postgres Query parser module
 """
-import csv
 import traceback
 from abc import ABC
-from datetime import datetime
-from typing import Iterable, Optional
+from typing import Iterable
 
 from sqlalchemy.engine.base import Engine
 
@@ -25,7 +23,7 @@ from metadata.generated.schema.entity.services.connections.database.postgresConn
 from metadata.generated.schema.metadataIngestion.workflow import (
     Source as WorkflowSource,
 )
-from metadata.generated.schema.type.tableQuery import TableQueries, TableQuery
+from metadata.generated.schema.type.tableQuery import TableQuery
 from metadata.ingestion.api.steps import InvalidSourceException
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.connections import get_connection
@@ -75,46 +73,11 @@ class PostgresQueryParserSource(QueryParserSource, ABC):
             time_column_name=get_postgres_time_column_name(engine=self.engine),
         )
 
+    # pylint: disable=no-member
     def get_table_query(self) -> Iterable[TableQuery]:
         try:
             if self.config.sourceConfig.config.queryLogFilePath:
-                table_query_list = []
-                with open(
-                    self.config.sourceConfig.config.queryLogFilePath,
-                    "r",
-                    encoding="utf-8",
-                ) as query_log_file:
-                    for record in csv.DictReader(query_log_file):
-                        query_dict = dict(record)
-
-                        analysis_date = (
-                            datetime.utcnow()
-                            if not query_dict.get("session_start_time")
-                            else datetime.strptime(
-                                query_dict.get("session_start_time"),
-                                "%Y-%m-%d %H:%M:%S+%f",
-                            )
-                        )
-
-                        query_dict["aborted"] = query_dict["sql_state_code"] == "00000"
-                        if "statement" in query_dict["message"]:
-                            query_dict["message"] = query_dict["message"].split(":")[1]
-
-                        table_query_list.append(
-                            TableQuery(
-                                query=query_dict["message"],
-                                userName=query_dict.get("user_name", ""),
-                                startTime=query_dict.get("session_start_time", ""),
-                                endTime=query_dict.get("log_time", ""),
-                                analysisDate=analysis_date,
-                                aborted=self.get_aborted_status(query_dict),
-                                databaseName=self.get_database_name(query_dict),
-                                serviceName=self.config.serviceName,
-                                databaseSchema=self.get_schema_name(query_dict),
-                            )
-                        )
-                yield TableQueries(queries=table_query_list)
-
+                yield from super().yield_table_queries_from_logs()
             else:
                 database = self.config.serviceConnection.__root__.config.database
                 if database:
@@ -133,37 +96,6 @@ class PostgresQueryParserSource(QueryParserSource, ABC):
             logger.error(f"Source usage processing error - {err}")
             logger.debug(traceback.format_exc())
 
-    def process_table_query(self) -> Optional[Iterable[TableQuery]]:
-        """
-        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))
-            yield TableQueries(queries=queries)
-        except Exception as err:
-            logger.error(f"Source usage processing error - {err}")
-            logger.debug(traceback.format_exc())
-
     @staticmethod
     def get_database_name(data: dict) -> str:
         """

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(),
                 )
             )