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

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

@@ -245,6 +245,10 @@ SNOWFLAKE_TEST_GET_TABLES = """
 SELECT TABLE_NAME FROM "{database_name}".information_schema.tables LIMIT 1
 """
 
+SNOWFLAKE_TEST_GET_VIEWS = """
+SELECT TABLE_NAME FROM "{database_name}".information_schema.views LIMIT 1
+"""
+
 SNOWFLAKE_GET_DATABASES = "SHOW DATABASES"
 
 
@@ -290,38 +294,17 @@ SELECT
   PROCEDURE_LANGUAGE AS language,
   PROCEDURE_DEFINITION AS definition,
   ARGUMENT_SIGNATURE AS signature,
-  COMMENT as comment,
-  'StoredProcedure' as procedure_type
+  COMMENT as comment
 FROM INFORMATION_SCHEMA.PROCEDURES
 WHERE PROCEDURE_CATALOG = '{database_name}'
   AND PROCEDURE_SCHEMA = '{schema_name}'
     """
 )
 
-SNOWFLAKE_GET_FUNCTIONS = textwrap.dedent(
-    """
-SELECT
-  FUNCTION_NAME AS name,
-  FUNCTION_OWNER AS owner,
-  FUNCTION_LANGUAGE AS language,
-  FUNCTION_DEFINITION AS definition,
-  ARGUMENT_SIGNATURE AS signature,
-  COMMENT as comment,
-  'UDF' as procedure_type
-FROM INFORMATION_SCHEMA.FUNCTIONS
-WHERE FUNCTION_CATALOG = '{database_name}'
-  AND FUNCTION_SCHEMA = '{schema_name}'
-    """
-)
-
 SNOWFLAKE_DESC_STORED_PROCEDURE = (
     "DESC PROCEDURE {database_name}.{schema_name}.{procedure_name}{procedure_signature}"
 )
 
-SNOWFLAKE_DESC_FUNCTION = (
-    "DESC FUNCTION {database_name}.{schema_name}.{procedure_name}{procedure_signature}"
-)
-
 SNOWFLAKE_GET_STORED_PROCEDURE_QUERIES = textwrap.dedent(
     """
 WITH SP_HISTORY AS (
@@ -376,3 +359,26 @@ ORDER BY PROCEDURE_START_TIME DESC
 SNOWFLAKE_GET_TABLE_DDL = """
 SELECT GET_DDL('TABLE','{table_name}') AS \"text\"
 """
+SNOWFLAKE_QUERY_LOG_QUERY = """
+    SELECT
+        QUERY_ID,
+        QUERY_TEXT,
+        QUERY_TYPE,
+        START_TIME,
+        DATABASE_NAME,
+        SCHEMA_NAME,
+        ROWS_INSERTED,
+        ROWS_UPDATED,
+        ROWS_DELETED
+    FROM "SNOWFLAKE"."ACCOUNT_USAGE"."QUERY_HISTORY"
+    WHERE
+    start_time>= DATEADD('DAY', -1, CURRENT_TIMESTAMP)
+    AND QUERY_TEXT ILIKE '%{tablename}%'
+    AND QUERY_TYPE IN (
+        '{insert}',
+        '{update}',
+        '{delete}',
+        '{merge}'
+    )
+    AND EXECUTION_STATUS = 'SUCCESS';
+"""

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

@@ -0,0 +1,20 @@
+from metadata.data_quality.interface.sqlalchemy.snowflake.test_suite_interface import (
+    SnowflakeTestSuiteInterface,
+)
+from metadata.ingestion.source.database.snowflake.lineage import SnowflakeLineageSource
+from metadata.ingestion.source.database.snowflake.metadata import SnowflakeSource
+from metadata.ingestion.source.database.snowflake.profiler.profiler import (
+    SnowflakeProfiler,
+)
+from metadata.ingestion.source.database.snowflake.usage import SnowflakeUsageSource
+from metadata.sampler.sqlalchemy.snowflake.sampler import SnowflakeSampler
+from metadata.utils.service_spec.default import DefaultDatabaseSpec
+
+ServiceSpec = DefaultDatabaseSpec(
+    metadata_source_class=SnowflakeSource,
+    lineage_source_class=SnowflakeLineageSource,
+    usage_source_class=SnowflakeUsageSource,
+    profiler_class=SnowflakeProfiler,
+    test_suite_class=SnowflakeTestSuiteInterface,
+    sampler_class=SnowflakeSampler,
+)

metadata/ingestion/source/database/sql_column_handler.py

@@ -196,6 +196,26 @@ class SqlColumnHandlerMixin:
             ]
         return Column(**parsed_string)
 
+    @calculate_execution_time()
+    def process_json_type_column_fields(  # pylint: disable=too-many-locals
+        self, schema_name: str, table_name: str, column_name: str, inspector: Inspector
+    ) -> Optional[List[Column]]:
+        """
+        Parse fields column with json data types
+        """
+        try:
+            if hasattr(inspector, "get_json_fields_and_type"):
+                result = inspector.get_json_fields_and_type(
+                    table_name, column_name, schema_name
+                )
+                return result
+
+        except NotImplementedError:
+            logger.debug(
+                "Cannot parse json fields for table column [{schema_name}.{table_name}.{col_name}]: NotImplementedError"
+            )
+        return None
+
     @calculate_execution_time()
     def get_columns_and_constraints(  # pylint: disable=too-many-locals
         self, schema_name: str, table_name: str, db_name: str, inspector: Inspector
@@ -271,6 +291,12 @@ class SqlColumnHandlerMixin:
                         precision,
                     )
                     col_data_length = 1 if col_data_length is None else col_data_length
+
+                    if col_type == "JSON":
+                        children = self.process_json_type_column_fields(
+                            schema_name, table_name, column.get("name"), inspector
+                        )
+
                     om_column = Column(
                         name=ColumnName(
                             root=column["name"]

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

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

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

metadata/ingestion/source/database/stored_procedures_mixin.py

@@ -11,6 +11,7 @@
 """
 Mixin class with common Stored Procedures logic aimed at lineage.
 """
+import json
 import re
 import traceback
 from abc import ABC, abstractmethod
@@ -27,8 +28,8 @@ from metadata.generated.schema.entity.data.storedProcedure import StoredProcedur
 from metadata.generated.schema.entity.services.ingestionPipelines.status import (
     StackTraceError,
 )
-from metadata.generated.schema.metadataIngestion.databaseServiceMetadataPipeline import (
-    DatabaseServiceMetadataPipeline,
+from metadata.generated.schema.metadataIngestion.databaseServiceQueryLineagePipeline import (
+    DatabaseServiceQueryLineagePipeline,
 )
 from metadata.generated.schema.type.basic import SqlQuery, Timestamp
 from metadata.generated.schema.type.entityLineage import Source as LineageSource
@@ -37,7 +38,6 @@ from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.status import Status
 from metadata.ingestion.lineage.models import ConnectionTypeDialectMapper
 from metadata.ingestion.lineage.sql_lineage import get_lineage_by_query
-from metadata.ingestion.models.topology import TopologyContextManager
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.utils.logger import ingestion_logger
 from metadata.utils.stored_procedures import get_procedure_name_from_call
@@ -66,7 +66,18 @@ class QueryByProcedure(BaseModel):
     model_config = ConfigDict(populate_by_name=True)
 
 
-class StoredProcedureMixin(ABC):
+class ProcedureAndQuery(BaseModel):
+    """
+    Model to hold the procedure and its queries
+    """
+
+    procedure: StoredProcedure
+    query_by_procedure: QueryByProcedure
+
+    model_config = ConfigDict(populate_by_name=True)
+
+
+class StoredProcedureLineageMixin(ABC):
     """
     The full flow is:
     1. List Stored Procedures
@@ -80,10 +91,10 @@ class StoredProcedureMixin(ABC):
     It should be inherited in those Sources that implement Stored Procedure ingestion.
     """
 
-    context: TopologyContextManager
     status: Status
-    source_config: DatabaseServiceMetadataPipeline
+    source_config: DatabaseServiceQueryLineagePipeline
     engine: Engine
+    stored_procedure_query_lineage: bool
     metadata: OpenMetadata
 
     @abstractmethod
@@ -143,29 +154,34 @@ class StoredProcedureMixin(ABC):
 
         return False
 
-    def yield_procedure_lineage(
+    def _yield_procedure_lineage(
         self, query_by_procedure: QueryByProcedure, procedure: StoredProcedure
     ) -> Iterable[Either[AddLineageRequest]]:
         """Add procedure lineage from its query"""
-        self.context.get().stored_procedure_query_lineage = False
+        self.stored_procedure_query_lineage = False
         if self.is_lineage_query(
             query_type=query_by_procedure.query_type,
             query_text=query_by_procedure.query_text,
         ):
-            self.context.get().stored_procedure_query_lineage = True
+            self.stored_procedure_query_lineage = True
             for either_lineage in get_lineage_by_query(
                 self.metadata,
                 query=query_by_procedure.query_text,
-                service_name=self.context.get().database_service,
+                service_name=self.service_name,
                 database_name=query_by_procedure.query_database_name,
                 schema_name=query_by_procedure.query_schema_name,
                 dialect=ConnectionTypeDialectMapper.dialect_of(
                     self.service_connection.type.value
                 ),
-                timeout_seconds=self.source_config.queryParsingTimeoutLimit,
+                timeout_seconds=self.source_config.parsingTimeoutLimit,
                 lineage_source=LineageSource.QueryLineage,
             ):
-                if either_lineage.right.edge.lineageDetails:
+                print("&& " * 100)
+                print(either_lineage)
+                if (
+                    either_lineage.left is None
+                    and either_lineage.right.edge.lineageDetails
+                ):
                     either_lineage.right.edge.lineageDetails.pipeline = EntityReference(
                         id=procedure.id,
                         type="storedProcedure",
@@ -192,49 +208,59 @@ class StoredProcedureMixin(ABC):
                     id=procedure.id,
                     type="storedProcedure",
                 ),
-                processedLineage=bool(
-                    self.context.get().stored_procedure_query_lineage
-                ),
-                service=self.context.get().database_service,
+                processedLineage=bool(self.stored_procedure_query_lineage),
+                service=self.service_name,
             )
         )
 
-    def yield_procedure_lineage_and_queries(
+    def procedure_lineage_processor(
+        self, procedure_and_query: ProcedureAndQuery
+    ) -> Iterable[Either[Union[AddLineageRequest, CreateQueryRequest]]]:
+
+        yield from self._yield_procedure_lineage(
+            query_by_procedure=procedure_and_query.query_by_procedure,
+            procedure=procedure_and_query.procedure,
+        )
+        yield from self.yield_procedure_query(
+            query_by_procedure=procedure_and_query.query_by_procedure,
+            procedure=procedure_and_query.procedure,
+        )
+
+    def procedure_lineage_generator(self) -> Iterable[ProcedureAndQuery]:
+        query = {
+            "query": {
+                "bool": {
+                    "must": [
+                        {"term": {"service.name.keyword": self.service_name}},
+                        {"term": {"deleted": False}},
+                    ]
+                }
+            }
+        }
+        query_filter = json.dumps(query)
+
+        logger.info("Processing Lineage for Stored Procedures")
+        # First, get all the query history
+        queries_dict = self.get_stored_procedure_queries_dict()
+        # Then for each procedure, iterate over all its queries
+        for procedure in (
+            self.metadata.paginate_es(entity=StoredProcedure, query_filter=query_filter)
+            or []
+        ):
+            if procedure:
+                logger.debug(f"Processing Lineage for [{procedure.name}]")
+                for query_by_procedure in (
+                    queries_dict.get(procedure.name.root.lower()) or []
+                ):
+                    yield ProcedureAndQuery(
+                        procedure=procedure, query_by_procedure=query_by_procedure
+                    )
+
+    def yield_procedure_lineage(
         self,
     ) -> Iterable[Either[Union[AddLineageRequest, CreateQueryRequest]]]:
         """Get all the queries and procedures list and yield them"""
-        if self.context.get().stored_procedures:
-            logger.info("Processing Lineage for Stored Procedures")
-            # First, get all the query history
-            queries_dict = self.get_stored_procedure_queries_dict()
-            # Then for each procedure, iterate over all its queries
-            for procedure_fqn in self.context.get().stored_procedures:
-                procedure: StoredProcedure = self.metadata.get_by_name(
-                    entity=StoredProcedure, fqn=procedure_fqn
-                )
-                if procedure:
-                    logger.debug(f"Processing Lineage for [{procedure.name}]")
-                    for query_by_procedure in (
-                        queries_dict.get(procedure.name.root.lower()) or []
-                    ):
-                        try:
-                            yield from self.yield_procedure_lineage(
-                                query_by_procedure=query_by_procedure,
-                                procedure=procedure,
-                            )
-                        except Exception as exc:
-                            logger.debug(traceback.format_exc())
-                            logger.warning(
-                                f"Could not get lineage for store procedure '{procedure_fqn}' due to [{exc}]."
-                            )
-
-                        try:
-                            yield from self.yield_procedure_query(
-                                query_by_procedure=query_by_procedure,
-                                procedure=procedure,
-                            )
-                        except Exception as exc:
-                            logger.debug(traceback.format_exc())
-                            logger.warning(
-                                f"Could not get query for store procedure '{procedure_fqn}' due to [{exc}]."
-                            )
+        logger.info("Processing Lineage for Stored Procedures")
+        producer_fn = self.procedure_lineage_generator
+        processor_fn = self.procedure_lineage_processor
+        yield from self.generate_lineage_in_thread(producer_fn, processor_fn)

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

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

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

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

@@ -12,11 +12,13 @@
 """
 Source connection handler
 """
+from copy import deepcopy
 from typing import Optional
 from urllib.parse import quote_plus
 
 from requests import Session
 from sqlalchemy.engine import Engine
+from trino.auth import BasicAuthentication, JWTAuthentication, OAuth2Authentication
 
 from metadata.clients.azure_client import AzureClient
 from metadata.generated.schema.entity.automations.workflow import (
@@ -25,6 +27,7 @@ from metadata.generated.schema.entity.automations.workflow import (
 from metadata.generated.schema.entity.services.connections.database.common import (
     basicAuth,
     jwtAuth,
+    noConfigAuthenticationTypes,
 )
 from metadata.generated.schema.entity.services.connections.database.trinoConnection import (
     TrinoConnection,
@@ -36,7 +39,6 @@ from metadata.ingestion.connections.builders import (
     create_generic_db_connection,
     get_connection_args_common,
     init_empty_connection_arguments,
-    init_empty_connection_options,
 )
 from metadata.ingestion.connections.secrets import connection_with_options_secrets
 from metadata.ingestion.connections.test_connections import (
@@ -52,26 +54,17 @@ def get_connection_url(connection: TrinoConnection) -> str:
     Prepare the connection url for trino
     """
     url = f"{connection.scheme.value}://"
+
+    # leaving username here as, even though with basic auth is used directly
+    # in BasicAuthentication class, it's often also required as a part of url.
+    # For example - it will be used by OAuth2Authentication to persist token in
+    # cache more efficiently (per user instead of per host)
     if connection.username:
-        # we need to encode twice because trino dialect internally
-        # url decodes the username and if there is an special char in username
-        # it will fail to authenticate
-        url += f"{quote_plus(quote_plus(connection.username))}"
-        if (
-            isinstance(connection.authType, basicAuth.BasicAuth)
-            and connection.authType.password
-        ):
-            url += f":{quote_plus(connection.authType.password.get_secret_value())}"
-        url += "@"
+        url += f"{quote_plus(connection.username)}@"
+
     url += f"{connection.hostPort}"
     if connection.catalog:
         url += f"/{connection.catalog}"
-    if isinstance(connection.authType, jwtAuth.JwtAuth):
-        if not connection.connectionOptions:
-            connection.connectionOptions = init_empty_connection_options()
-        connection.connectionOptions.root[
-            "access_token"
-        ] = connection.authType.jwt.get_secret_value()
     if connection.connectionOptions is not None:
         params = "&".join(
             f"{key}={quote_plus(value)}"
@@ -84,14 +77,54 @@ def get_connection_url(connection: TrinoConnection) -> str:
 
 @connection_with_options_secrets
 def get_connection_args(connection: TrinoConnection):
+    if not connection.connectionArguments:
+        connection.connectionArguments = init_empty_connection_arguments()
+
     if connection.proxies:
         session = Session()
         session.proxies = connection.proxies
-        if not connection.connectionArguments:
-            connection.connectionArguments = init_empty_connection_arguments()
 
         connection.connectionArguments.root["http_session"] = session
 
+    if isinstance(connection.authType, basicAuth.BasicAuth):
+        connection.connectionArguments.root["auth"] = BasicAuthentication(
+            connection.username,
+            connection.authType.password.get_secret_value()
+            if connection.authType.password
+            else None,
+        )
+        connection.connectionArguments.root["http_scheme"] = "https"
+
+    elif isinstance(connection.authType, jwtAuth.JwtAuth):
+        connection.connectionArguments.root["auth"] = JWTAuthentication(
+            connection.authType.jwt.get_secret_value()
+        )
+        connection.connectionArguments.root["http_scheme"] = "https"
+
+    elif hasattr(connection.authType, "azureConfig"):
+        if not connection.authType.azureConfig.scopes:
+            raise ValueError(
+                "Azure Scopes are missing, please refer https://learn.microsoft.com/en-gb/azure/mysql/flexible-server/how-to-azure-ad#2---retrieve-microsoft-entra-access-token and fetch the resource associated with it, for e.g. https://ossrdbms-aad.database.windows.net/.default"
+            )
+
+        azure_client = AzureClient(connection.authType.azureConfig).create_client()
+
+        access_token_obj = azure_client.get_token(
+            *connection.authType.azureConfig.scopes.split(",")
+        )
+
+        connection.connectionArguments.root["auth"] = JWTAuthentication(
+            access_token_obj.token
+        )
+        connection.connectionArguments.root["http_scheme"] = "https"
+
+    elif (
+        connection.authType
+        == noConfigAuthenticationTypes.NoConfigAuthenticationTypes.OAuth2
+    ):
+        connection.connectionArguments.root["auth"] = OAuth2Authentication()
+        connection.connectionArguments.root["http_scheme"] = "https"
+
     return get_connection_args_common(connection)
 
 
@@ -99,9 +132,13 @@ def get_connection(connection: TrinoConnection) -> Engine:
     """
     Create connection
     """
-    if connection.verify:
-        connection.connectionArguments = (
-            connection.connectionArguments or init_empty_connection_arguments()
+    # here we are creating a copy of connection, because we need to dynamically
+    # add auth params to connectionArguments, which we do no intend to store
+    # in original connection object and in OpenMetadata database
+    connection_copy = deepcopy(connection)
+    if connection_copy.verify:
+        connection_copy.connectionArguments = (
+            connection_copy.connectionArguments or init_empty_connection_arguments()
         )
         connection.connectionArguments.root["verify"] = {"verify": connection.verify}
     if hasattr(connection.authType, "azureConfig"):
@@ -117,7 +154,7 @@ def get_connection(connection: TrinoConnection) -> Engine:
             connection.connectionOptions = init_empty_connection_options()
         connection.connectionOptions.root["access_token"] = access_token_obj.token
     return create_generic_db_connection(
-        connection=connection,
+        connection=connection_copy,
         get_connection_url_fn=get_connection_url,
         get_connection_args_fn=get_connection_args,
     )