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

metadata/ingestion/source/database/impala/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.
+"""
+Impala lineage module
+"""
+from typing import Optional
+
+from metadata.generated.schema.entity.services.connections.database.impalaConnection import (
+    ImpalaConnection,
+)
+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 ImpalaLineageSource(LineageSource):
+    """
+    Impala 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: ImpalaConnection = config.serviceConnection.root.config
+        if not isinstance(connection, ImpalaConnection):
+            raise InvalidSourceException(
+                f"Expected ImpalaConnection, but got {connection}"
+            )
+        return cls(config, metadata)

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

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

metadata/ingestion/source/database/life_cycle_query_mixin.py

@@ -36,7 +36,7 @@ from metadata.ingestion.models.topology import TopologyContextManager
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.utils import fqn
 from metadata.utils.logger import ingestion_logger
-from metadata.utils.time_utils import convert_timestamp_to_milliseconds
+from metadata.utils.time_utils import datetime_to_timestamp
 
 logger = ingestion_logger()
 
@@ -104,10 +104,8 @@ class LifeCycleQueryMixin:
                 life_cycle = LifeCycle(
                     created=AccessDetails(
                         timestamp=Timestamp(
-                            int(
-                                convert_timestamp_to_milliseconds(
-                                    life_cycle_data.created_at.timestamp()
-                                )
+                            datetime_to_timestamp(
+                                life_cycle_data.created_at, milliseconds=True
                             )
                         )
                     )

metadata/ingestion/source/database/lineage_source.py

@@ -15,17 +15,22 @@ import csv
 import os
 import traceback
 from abc import ABC
-from typing import Iterable, Iterator, Union
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from functools import partial
+from typing import Callable, Iterable, Iterator, Union
 
 from metadata.generated.schema.api.data.createQuery import CreateQueryRequest
 from metadata.generated.schema.api.lineage.addLineage import AddLineageRequest
 from metadata.generated.schema.type.basic import FullyQualifiedEntityName, SqlQuery
 from metadata.generated.schema.type.tableQuery import TableQuery
 from metadata.ingestion.api.models import Either
-from metadata.ingestion.lineage.models import ConnectionTypeDialectMapper
+from metadata.ingestion.lineage.models import ConnectionTypeDialectMapper, Dialect
 from metadata.ingestion.lineage.sql_lineage import get_lineage_by_query
+from metadata.ingestion.models.ometa_lineage import OMetaLineageRequest
 from metadata.ingestion.source.database.query_parser_source import QueryParserSource
+from metadata.ingestion.source.models import TableView
 from metadata.utils import fqn
+from metadata.utils.db_utils import get_view_lineage
 from metadata.utils.logger import ingestion_logger
 
 logger = ingestion_logger()
@@ -43,12 +48,14 @@ class LineageSource(QueryParserSource, ABC):
     - schema
     """
 
+    dialect: Dialect
+
     def yield_table_queries_from_logs(self) -> Iterator[TableQuery]:
         """
         Method to handle the usage from query logs
         """
         try:
-            query_log_path = self.config.sourceConfig.config.queryLogFilePath
+            query_log_path = self.source_config.queryLogFilePath
             if os.path.isfile(query_log_path):
                 file_paths = [query_log_path]
             elif os.path.isdir(query_log_path):
@@ -89,6 +96,28 @@ class LineageSource(QueryParserSource, ABC):
             )
             yield from self.yield_table_query()
 
+    def generate_lineage_in_thread(self, producer_fn: Callable, processor_fn: Callable):
+        with ThreadPoolExecutor(max_workers=self.source_config.threads) as executor:
+            futures = []
+
+            for produced_input in producer_fn():
+                futures.append(executor.submit(processor_fn, produced_input))
+
+            # Handle remaining futures after the loop
+            for future in as_completed(
+                futures, timeout=self.source_config.parsingTimeoutLimit
+            ):
+                try:
+                    results = future.result(
+                        timeout=self.source_config.parsingTimeoutLimit
+                    )
+                    yield from results
+                except Exception as exc:
+                    logger.debug(traceback.format_exc())
+                    logger.warning(
+                        f"Error processing result for {produced_input}: {exc}"
+                    )
+
     def yield_table_query(self) -> Iterator[TableQuery]:
         """
         Given an engine, iterate over the query results to
@@ -128,40 +157,103 @@ class LineageSource(QueryParserSource, ABC):
         )
         return fqn.get_query_checksum(table_query.query) in checksums or {}
 
-    def _iter(
-        self, *_, **__
+    def query_lineage_generator(
+        self, table_query: TableQuery
+    ) -> Iterable[Either[Union[AddLineageRequest, CreateQueryRequest]]]:
+        if not self._query_already_processed(table_query):
+            lineages: Iterable[Either[AddLineageRequest]] = get_lineage_by_query(
+                self.metadata,
+                query=table_query.query,
+                service_name=table_query.serviceName,
+                database_name=table_query.databaseName,
+                schema_name=table_query.databaseSchema,
+                dialect=self.dialect,
+                timeout_seconds=self.source_config.parsingTimeoutLimit,
+            )
+
+            for lineage_request in lineages or []:
+                yield lineage_request
+
+                # If we identified lineage properly, ingest the original query
+                if lineage_request.right:
+                    yield Either(
+                        right=CreateQueryRequest(
+                            query=SqlQuery(table_query.query),
+                            query_type=table_query.query_type,
+                            duration=table_query.duration,
+                            processedLineage=True,
+                            service=FullyQualifiedEntityName(self.config.serviceName),
+                        )
+                    )
+
+    def yield_query_lineage(
+        self,
     ) -> Iterable[Either[Union[AddLineageRequest, CreateQueryRequest]]]:
         """
         Based on the query logs, prepare the lineage
         and send it to the sink
         """
         connection_type = str(self.service_connection.type.value)
-        dialect = ConnectionTypeDialectMapper.dialect_of(connection_type)
-        for table_query in self.get_table_query():
-            if not self._query_already_processed(table_query):
-                lineages: Iterable[Either[AddLineageRequest]] = get_lineage_by_query(
-                    self.metadata,
-                    query=table_query.query,
-                    service_name=table_query.serviceName,
-                    database_name=table_query.databaseName,
-                    schema_name=table_query.databaseSchema,
-                    dialect=dialect,
-                    timeout_seconds=self.source_config.parsingTimeoutLimit,
-                )
+        self.dialect = ConnectionTypeDialectMapper.dialect_of(connection_type)
+        producer_fn = self.get_table_query
+        processor_fn = self.query_lineage_generator
+        yield from self.generate_lineage_in_thread(producer_fn, processor_fn)
 
-                for lineage_request in lineages or []:
-                    yield lineage_request
-
-                    # If we identified lineage properly, ingest the original query
-                    if lineage_request.right:
-                        yield Either(
-                            right=CreateQueryRequest(
-                                query=SqlQuery(table_query.query),
-                                query_type=table_query.query_type,
-                                duration=table_query.duration,
-                                processedLineage=True,
-                                service=FullyQualifiedEntityName(
-                                    self.config.serviceName
-                                ),
-                            )
+    def view_lineage_generator(
+        self, view: TableView
+    ) -> Iterable[Either[AddLineageRequest]]:
+        try:
+            for lineage in get_view_lineage(
+                view=view,
+                metadata=self.metadata,
+                service_name=self.config.serviceName,
+                connection_type=self.service_connection.type.value,
+                timeout_seconds=self.source_config.parsingTimeoutLimit,
+            ):
+                if lineage.right is not None:
+                    yield Either(
+                        right=OMetaLineageRequest(
+                            lineage_request=lineage.right,
+                            override_lineage=self.source_config.overrideViewLineage,
                         )
+                    )
+                else:
+                    yield lineage
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.warning(f"Error processing view {view}: {exc}")
+
+    def yield_view_lineage(self) -> Iterable[Either[AddLineageRequest]]:
+        logger.info("Processing View Lineage")
+        producer_fn = partial(self.metadata.yield_es_view_def, self.config.serviceName)
+        processor_fn = self.view_lineage_generator
+        yield from self.generate_lineage_in_thread(producer_fn, processor_fn)
+
+    def yield_procedure_lineage(
+        self,
+    ) -> Iterable[Either[Union[AddLineageRequest, CreateQueryRequest]]]:
+        """
+        By default stored procedure lineage is not supported.
+        """
+        logger.info(
+            f"Processing Procedure Lineage not supported for {str(self.service_connection.type.value)}"
+        )
+
+    def _iter(
+        self, *_, **__
+    ) -> Iterable[Either[Union[AddLineageRequest, CreateQueryRequest]]]:
+        """
+        Based on the query logs, prepare the lineage
+        and send it to the sink
+        """
+        if self.source_config.processViewLineage:
+            yield from self.yield_view_lineage() or []
+        if self.source_config.processStoredProcedureLineage:
+            yield from self.yield_procedure_lineage() or []
+        if self.source_config.processQueryLineage:
+            if hasattr(self.service_connection, "supportsLineageExtraction"):
+                yield from self.yield_query_lineage() or []
+            else:
+                logger.warning(
+                    f"Lineage extraction is not supported for {str(self.service_connection.type.value)} connection"
+                )

metadata/ingestion/source/database/mariadb/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.
+"""
+Mariadb lineage module
+"""
+from typing import Optional
+
+from metadata.generated.schema.entity.services.connections.database.mariaDBConnection import (
+    MariadbConnection,
+)
+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 MariadbLineageSource(LineageSource):
+    """
+    Mariadb 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: MariadbConnection = config.serviceConnection.root.config
+        if not isinstance(connection, MariadbConnection):
+            raise InvalidSourceException(
+                f"Expected MariadbConnection, but got {connection}"
+            )
+        return cls(config, metadata)

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

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

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

@@ -0,0 +1,10 @@
+from metadata.ingestion.source.database.mongodb.metadata import MongodbSource
+from metadata.profiler.interface.nosql.profiler_interface import NoSQLProfilerInterface
+from metadata.sampler.nosql.sampler import NoSQLSampler
+from metadata.utils.service_spec.default import DefaultDatabaseSpec
+
+ServiceSpec = DefaultDatabaseSpec(
+    metadata_source_class=MongodbSource,
+    profiler_class=NoSQLProfilerInterface,
+    sampler_class=NoSQLSampler,
+)

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

@@ -11,21 +11,39 @@
 """
 MSSQL lineage module
 """
+import traceback
 from datetime import datetime
+from typing import Dict, List
 
+from metadata.generated.schema.entity.services.ingestionPipelines.status import (
+    StackTraceError,
+)
 from metadata.ingestion.source.database.lineage_source import LineageSource
 from metadata.ingestion.source.database.mssql.constants import (
     DEFAULT_DATETIME_FORMAT,
     MSSQL_DATEFORMAT_DATETIME_MAP,
 )
-from metadata.ingestion.source.database.mssql.queries import MSSQL_SQL_STATEMENT
+from metadata.ingestion.source.database.mssql.queries import (
+    MSSQL_GET_STORED_PROCEDURE_QUERIES,
+    MSSQL_SQL_STATEMENT,
+)
 from metadata.ingestion.source.database.mssql.query_parser import MssqlQueryParserSource
 from metadata.ingestion.source.database.mssql.utils import (
     get_sqlalchemy_engine_dateformat,
 )
+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 MssqlLineageSource(MssqlQueryParserSource, LineageSource):
+class MssqlLineageSource(
+    MssqlQueryParserSource, StoredProcedureLineageMixin, LineageSource
+):
 
     sql_stmt = MSSQL_SQL_STATEMENT
 
@@ -55,3 +73,34 @@ class MssqlLineageSource(MssqlQueryParserSource, LineageSource):
             filters=self.get_filters(),
             result_limit=self.source_config.resultLimit,
         )
+
+    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)
+        server_date_format = get_sqlalchemy_engine_dateformat(self.engine)
+        current_datetime_format = MSSQL_DATEFORMAT_DATETIME_MAP.get(
+            server_date_format, DEFAULT_DATETIME_FORMAT
+        )
+        start = start.strftime(current_datetime_format)
+        query = MSSQL_GET_STORED_PROCEDURE_QUERIES.format(
+            start_date=start,
+        )
+        try:
+            queries_dict = self.procedure_queries_dict(
+                query=query,
+            )
+        except Exception as ex:  # pylint: disable=broad-except
+            logger.debug(f"Error runnning query:\n{query}")
+            self.status.failed(
+                StackTraceError(
+                    name="Stored Procedure",
+                    error=f"Error trying to get stored procedure queries: {ex}",
+                    stackTrace=traceback.format_exc(),
+                )
+            )
+            return {}
+
+        return queries_dict

metadata/ingestion/source/database/mssql/metadata.py

@@ -10,7 +10,7 @@
 #  limitations under the License.
 """MSSQL source module"""
 import traceback
-from typing import Dict, Iterable, List, Optional
+from typing import Iterable, Optional
 
 from sqlalchemy.dialects.mssql.base import MSDialect, ischema_names
 from sqlalchemy.engine.reflection import Inspector
@@ -35,24 +35,18 @@ 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
-from metadata.ingestion.source.database.mssql.constants import (
-    DEFAULT_DATETIME_FORMAT,
-    MSSQL_DATEFORMAT_DATETIME_MAP,
-)
 from metadata.ingestion.source.database.mssql.models import (
     STORED_PROC_LANGUAGE_MAP,
     MssqlStoredProcedure,
 )
 from metadata.ingestion.source.database.mssql.queries import (
     MSSQL_GET_DATABASE,
-    MSSQL_GET_STORED_PROCEDURE_QUERIES,
     MSSQL_GET_STORED_PROCEDURES,
 )
 from metadata.ingestion.source.database.mssql.utils import (
     get_columns,
     get_foreign_keys,
     get_pk_constraint,
-    get_sqlalchemy_engine_dateformat,
     get_table_comment,
     get_table_names,
     get_unique_constraints,
@@ -60,13 +54,8 @@ from metadata.ingestion.source.database.mssql.utils import (
     get_view_names,
 )
 from metadata.ingestion.source.database.multi_db_source import MultiDBSource
-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.helpers import get_start_and_end
 from metadata.utils.logger import ingestion_logger
 from metadata.utils.sqa_utils import update_mssql_ischema_names
 from metadata.utils.sqlalchemy_utils import (
@@ -99,7 +88,7 @@ Inspector.get_all_table_ddls = get_all_table_ddls
 Inspector.get_table_ddl = get_table_ddl
 
 
-class MssqlSource(StoredProcedureMixin, CommonDbSourceService, MultiDBSource):
+class MssqlSource(CommonDbSourceService, MultiDBSource):
     """
     Implements the necessary methods to extract
     Database metadata from MSSQL Source
@@ -213,34 +202,3 @@ class MssqlSource(StoredProcedureMixin, CommonDbSourceService, MultiDBSource):
                     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)
-        server_date_format = get_sqlalchemy_engine_dateformat(self.engine)
-        current_datetime_format = MSSQL_DATEFORMAT_DATETIME_MAP.get(
-            server_date_format, DEFAULT_DATETIME_FORMAT
-        )
-        start = start.strftime(current_datetime_format)
-        query = MSSQL_GET_STORED_PROCEDURE_QUERIES.format(
-            start_date=start,
-        )
-        try:
-            queries_dict = self.procedure_queries_dict(
-                query=query,
-            )
-        except Exception as ex:  # pylint: disable=broad-except
-            logger.debug(f"Error runnning query:\n{query}")
-            self.status.failed(
-                StackTraceError(
-                    name="Stored Procedure",
-                    error=f"Error trying to get stored procedure queries: {ex}",
-                    stackTrace=traceback.format_exc(),
-                )
-            )
-            return {}
-
-        return queries_dict

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

@@ -0,0 +1,10 @@
+from metadata.ingestion.source.database.mssql.lineage import MssqlLineageSource
+from metadata.ingestion.source.database.mssql.metadata import MssqlSource
+from metadata.ingestion.source.database.mssql.usage import MssqlUsageSource
+from metadata.utils.service_spec.default import DefaultDatabaseSpec
+
+ServiceSpec = DefaultDatabaseSpec(
+    metadata_source_class=MssqlSource,
+    lineage_source_class=MssqlLineageSource,
+    usage_source_class=MssqlUsageSource,
+)

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

@@ -17,9 +17,7 @@ from metadata.ingestion.source.database.mysql.query_parser import MysqlQueryPars
 
 
 class MysqlLineageSource(MysqlQueryParserSource, LineageSource):
-
     sql_stmt = MYSQL_SQL_STATEMENT
-
     filters = """
         AND (
             lower(argument) LIKE '%%create%%table%%select%%'

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

@@ -11,7 +11,6 @@
 """
 SQL Queries used during ingestion
 """
-
 import textwrap
 
 MYSQL_SQL_STATEMENT = textwrap.dedent(
@@ -37,7 +36,6 @@ LIMIT {result_limit};
 """
 )
 
-
 MYSQL_TEST_GET_QUERIES = textwrap.dedent(
     """
 SELECT `argument` from mysql.general_log limit 1;

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

@@ -0,0 +1,10 @@
+from metadata.ingestion.source.database.mysql.lineage import MysqlLineageSource
+from metadata.ingestion.source.database.mysql.metadata import MysqlSource
+from metadata.ingestion.source.database.mysql.usage import MysqlUsageSource
+from metadata.utils.service_spec.default import DefaultDatabaseSpec
+
+ServiceSpec = DefaultDatabaseSpec(
+    metadata_source_class=MysqlSource,
+    lineage_source_class=MysqlLineageSource,
+    usage_source_class=MysqlUsageSource,
+)

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

@@ -17,9 +17,7 @@ from metadata.ingestion.source.database.usage_source import UsageSource
 
 
 class MysqlUsageSource(MysqlQueryParserSource, UsageSource):
-
     sql_stmt = MYSQL_SQL_STATEMENT
-
     filters = ""
 
     def format_query(self, query: bytes) -> str:

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

@@ -30,16 +30,26 @@ workflowConfig:
       jwtToken: "token"
 """
 
+from typing import Dict, List
+
 from metadata.ingestion.source.database.lineage_source import LineageSource
 from metadata.ingestion.source.database.oracle.queries import (
+    ORACLE_GET_STORED_PROCEDURE_QUERIES,
     ORACLE_QUERY_HISTORY_STATEMENT,
 )
 from metadata.ingestion.source.database.oracle.query_parser import (
     OracleQueryParserSource,
 )
+from metadata.ingestion.source.database.stored_procedures_mixin import (
+    QueryByProcedure,
+    StoredProcedureLineageMixin,
+)
+from metadata.utils.helpers import get_start_and_end
 
 
-class OracleLineageSource(OracleQueryParserSource, LineageSource):
+class OracleLineageSource(
+    OracleQueryParserSource, StoredProcedureLineageMixin, LineageSource
+):
     filters = """
         AND COMMAND_TYPE IN (1, 2) AND (
             lower(SQL_FULLTEXT) LIKE '%%create%%table%%as%%select%%'
@@ -50,3 +60,20 @@ class OracleLineageSource(OracleQueryParserSource, LineageSource):
         """
 
     sql_stmt = ORACLE_QUERY_HISTORY_STATEMENT
+
+    stored_procedure_query = ORACLE_GET_STORED_PROCEDURE_QUERIES
+
+    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 = self.stored_procedure_query.format(
+            start_date=start,
+        )
+        queries_dict = self.procedure_queries_dict(
+            query=query,
+        )
+
+        return queries_dict