openmetadata-ingestion 1.6.3.2.dev1__py3-none-any.whl → 1.6.4.0__py3-none-any.whl

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

@@ -11,12 +11,10 @@
 """
 Databricks lineage module
 """
-import traceback
-from datetime import datetime
-from typing import Iterator
 
-from metadata.generated.schema.type.basic import DateTime
-from metadata.generated.schema.type.tableQuery import TableQuery
+from metadata.ingestion.source.database.databricks.queries import (
+    DATABRICKS_SQL_STATEMENT,
+)
 from metadata.ingestion.source.database.databricks.query_parser import (
     DatabricksQueryParserSource,
 )
@@ -31,23 +29,13 @@ class DatabricksLineageSource(DatabricksQueryParserSource, LineageSource):
     Databricks Lineage Legacy Source
     """
 
-    def yield_table_query(self) -> Iterator[TableQuery]:
-        data = self.client.list_query_history(
-            start_date=self.start,
-            end_date=self.end,
+    sql_stmt = DATABRICKS_SQL_STATEMENT
+
+    filters = """
+        AND (
+            lower(statement_text) LIKE '%%create%%select%%'
+            OR lower(statement_text) LIKE '%%insert%%into%%select%%'
+            OR lower(statement_text) LIKE '%%update%%'
+            OR lower(statement_text) LIKE '%%merge%%'
         )
-        for row in data or []:
-            try:
-                if self.client.is_query_valid(row):
-                    yield TableQuery(
-                        dialect=self.dialect.value,
-                        query=row.get("query_text"),
-                        userName=row.get("user_name"),
-                        startTime=str(row.get("query_start_time_ms")),
-                        endTime=str(row.get("execution_end_time_ms")),
-                        analysisDate=DateTime(datetime.now()),
-                        serviceName=self.config.serviceName,
-                    )
-            except Exception as exc:
-                logger.debug(traceback.format_exc())
-                logger.warning(f"Error processing query_dict {row}: {exc}")
+    """

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

@@ -14,6 +14,30 @@ SQL Queries used during ingestion
 
 import textwrap
 
+DATABRICKS_SQL_STATEMENT = textwrap.dedent(
+    """
+    SELECT
+      statement_type AS query_type,
+      statement_text AS query_text,
+      executed_by AS user_name,
+      start_time AS start_time,
+      null AS database_name,
+      null AS schema_name,
+      end_time AS end_time,
+      total_duration_ms/1000 AS duration
+    from {query_history}
+    WHERE statement_text NOT LIKE '/* {{"app": "OpenMetadata", %%}} */%%'
+    AND statement_text NOT LIKE '/* {{"app": "dbt", %%}} */%%'
+    AND start_time between to_timestamp('{start_time}') and to_timestamp('{end_time}')
+    {filters}
+    LIMIT {result_limit}
+    """
+)
+
+DATABRICKS_SQL_STATEMENT_TEST = """
+ SELECT statement_text from  {query_history} LIMIT 1
+"""
+
 DATABRICKS_VIEW_DEFINITIONS = textwrap.dedent(
     """
     select

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

@@ -22,7 +22,6 @@ from metadata.generated.schema.metadataIngestion.workflow import (
 )
 from metadata.ingestion.api.steps import InvalidSourceException
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
-from metadata.ingestion.source.database.databricks.client import DatabricksClient
 from metadata.ingestion.source.database.query_parser_source import QueryParserSource
 from metadata.utils.logger import ingestion_logger
 
@@ -36,18 +35,6 @@ class DatabricksQueryParserSource(QueryParserSource, ABC):
 
     filters: str
 
-    def _init_super(
-        self,
-        config: WorkflowSource,
-        metadata: OpenMetadata,
-    ):
-        super().__init__(config, metadata, False)
-
-    # pylint: disable=super-init-not-called
-    def __init__(self, config: WorkflowSource, metadata: OpenMetadata):
-        self._init_super(config=config, metadata=metadata)
-        self.client = DatabricksClient(self.service_connection)
-
     @classmethod
     def create(
         cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
@@ -61,7 +48,16 @@ class DatabricksQueryParserSource(QueryParserSource, ABC):
             )
         return cls(config, metadata)
 
-    def prepare(self):
+    def get_sql_statement(self, start_time, end_time):
         """
-        By default, there's nothing to prepare
+        returns sql statement to fetch query logs.
+
+        Override if we have specific parameters
         """
+        return self.sql_stmt.format(
+            start_time=start_time,
+            end_time=end_time,
+            filters=self.get_filters(),
+            result_limit=self.source_config.resultLimit,
+            query_history=self.service_connection.queryHistoryTable,
+        )

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

@@ -11,12 +11,10 @@
 """
 Databricks usage module
 """
-import traceback
-from datetime import datetime
-from typing import Iterable
 
-from metadata.generated.schema.type.basic import DateTime
-from metadata.generated.schema.type.tableQuery import TableQueries, TableQuery
+from metadata.ingestion.source.database.databricks.queries import (
+    DATABRICKS_SQL_STATEMENT,
+)
 from metadata.ingestion.source.database.databricks.query_parser import (
     DatabricksQueryParserSource,
 )
@@ -31,36 +29,8 @@ class DatabricksUsageSource(DatabricksQueryParserSource, UsageSource):
     Databricks Usage Source
     """
 
-    def yield_table_queries(self) -> Iterable[TableQuery]:
-        """
-        Method to yield TableQueries
-        """
-        queries = []
-        data = self.client.list_query_history(
-            start_date=self.start,
-            end_date=self.end,
-        )
-        for row in data or []:
-            try:
-                if self.client.is_query_valid(row):
-                    queries.append(
-                        TableQuery(
-                            dialect=self.dialect.value,
-                            query=row.get("query_text"),
-                            userName=row.get("user_name"),
-                            startTime=str(row.get("query_start_time_ms")),
-                            endTime=str(row.get("execution_end_time_ms")),
-                            analysisDate=DateTime(datetime.now()),
-                            serviceName=self.config.serviceName,
-                            duration=row.get("duration")
-                            if row.get("duration")
-                            else None,
-                        )
-                    )
-            except Exception as err:
-                logger.debug(traceback.format_exc())
-                logger.warning(
-                    f"Failed to process query {row.get('query_text')} due to: {err}"
-                )
+    sql_stmt = DATABRICKS_SQL_STATEMENT
 
-        yield TableQueries(queries=queries)
+    filters = """
+        AND statement_type NOT IN ('SHOW', 'DESCRIBE', 'USE')
+    """

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

@@ -50,7 +50,7 @@ def get_connection(connection: Db2Connection) -> Engine:
             "w",
             encoding=UTF_8,
         ) as file:
-            file.write(connection.license)
+            file.write(connection.license.encode(UTF_8).decode("unicode-escape"))
 
     return create_generic_db_connection(
         connection=connection,

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

@@ -1,4 +1,3 @@
-from metadata.ingestion.source.database.db2.lineage import Db2LineageSource
 from metadata.ingestion.source.database.db2.metadata import Db2Source
 from metadata.profiler.interface.sqlalchemy.db2.profiler_interface import (
     DB2ProfilerInterface,
@@ -6,7 +5,5 @@ from metadata.profiler.interface.sqlalchemy.db2.profiler_interface import (
 from metadata.utils.service_spec.default import DefaultDatabaseSpec
 
 ServiceSpec = DefaultDatabaseSpec(
-    metadata_source_class=Db2Source,
-    profiler_class=DB2ProfilerInterface,
-    lineage_source_class=Db2LineageSource,
+    metadata_source_class=Db2Source, profiler_class=DB2ProfilerInterface
 )

metadata/ingestion/source/database/dbt/metadata.py

@@ -1093,7 +1093,9 @@ class DbtSource(DbtServiceSource):
 
                 # Create the test case result object
                 test_case_result = TestCaseResult(
-                    timestamp=Timestamp(datetime_to_timestamp(dbt_timestamp)),
+                    timestamp=Timestamp(
+                        datetime_to_timestamp(dbt_timestamp, milliseconds=True)
+                    ),
                     testCaseStatus=test_case_status,
                     testResultValue=[
                         TestResultValue(

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

@@ -1,8 +1,4 @@
-from metadata.ingestion.source.database.druid.lineage import DruidLineageSource
 from metadata.ingestion.source.database.druid.metadata import DruidSource
 from metadata.utils.service_spec.default import DefaultDatabaseSpec
 
-ServiceSpec = DefaultDatabaseSpec(
-    metadata_source_class=DruidSource,
-    lineage_source_class=DruidLineageSource,
-)
+ServiceSpec = DefaultDatabaseSpec(metadata_source_class=DruidSource)

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

@@ -1,8 +1,4 @@
-from metadata.ingestion.source.database.greenplum.lineage import GreenplumLineageSource
 from metadata.ingestion.source.database.greenplum.metadata import GreenplumSource
 from metadata.utils.service_spec.default import DefaultDatabaseSpec
 
-ServiceSpec = DefaultDatabaseSpec(
-    metadata_source_class=GreenplumSource,
-    lineage_source_class=GreenplumLineageSource,
-)
+ServiceSpec = DefaultDatabaseSpec(metadata_source_class=GreenplumSource)

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

@@ -1,8 +1,4 @@
-from metadata.ingestion.source.database.hive.lineage import HiveLineageSource
 from metadata.ingestion.source.database.hive.metadata import HiveSource
 from metadata.utils.service_spec.default import DefaultDatabaseSpec
 
-ServiceSpec = DefaultDatabaseSpec(
-    metadata_source_class=HiveSource,
-    lineage_source_class=HiveLineageSource,
-)
+ServiceSpec = DefaultDatabaseSpec(metadata_source_class=HiveSource)

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

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

metadata/ingestion/source/database/lineage_source.py

@@ -26,7 +26,10 @@ from metadata.generated.schema.type.basic import FullyQualifiedEntityName, SqlQu
 from metadata.generated.schema.type.tableQuery import TableQuery
 from metadata.ingestion.api.models import Either
 from metadata.ingestion.lineage.models import ConnectionTypeDialectMapper, Dialect
-from metadata.ingestion.lineage.sql_lineage import get_lineage_by_query
+from metadata.ingestion.lineage.sql_lineage import (
+    get_lineage_by_graph,
+    get_lineage_by_query,
+)
 from metadata.ingestion.models.ometa_lineage import OMetaLineageRequest
 from metadata.ingestion.models.topology import Queue
 from metadata.ingestion.source.database.query_parser_source import QueryParserSource
@@ -195,6 +198,12 @@ class LineageSource(QueryParserSource, ABC):
     def query_lineage_generator(
         self, table_queries: List[TableQuery], queue: Queue
     ) -> Iterable[Either[Union[AddLineageRequest, CreateQueryRequest]]]:
+        if self.graph is None and self.source_config.enableTempTableLineage:
+            import networkx as nx
+
+            # Create a directed graph
+            self.graph = nx.DiGraph()
+
         for table_query in table_queries or []:
             if not self._query_already_processed(table_query):
                 lineages: Iterable[Either[AddLineageRequest]] = get_lineage_by_query(
@@ -205,6 +214,7 @@ class LineageSource(QueryParserSource, ABC):
                     schema_name=table_query.databaseSchema,
                     dialect=self.dialect,
                     timeout_seconds=self.source_config.parsingTimeoutLimit,
+                    graph=self.graph,
                 )
 
                 for lineage_request in lineages or []:
@@ -298,6 +308,7 @@ class LineageSource(QueryParserSource, ABC):
         if self.source_config.processQueryLineage:
             if hasattr(self.service_connection, "supportsLineageExtraction"):
                 yield from self.yield_query_lineage() or []
+                yield from get_lineage_by_graph(graph=self.graph)
             else:
                 logger.warning(
                     f"Lineage extraction is not supported for {str(self.service_connection.type.value)} connection"

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

@@ -1,12 +1,4 @@
-from metadata.ingestion.source.database.mariadb.lineage import MariadbLineageSource
 from metadata.ingestion.source.database.mariadb.metadata import MariadbSource
-from metadata.profiler.interface.sqlalchemy.mariadb.profiler_interface import (
-    MariaDBProfilerInterface,
-)
 from metadata.utils.service_spec.default import DefaultDatabaseSpec
 
-ServiceSpec = DefaultDatabaseSpec(
-    metadata_source_class=MariadbSource,
-    lineage_source_class=MariadbLineageSource,
-    profiler_class=MariaDBProfilerInterface,
-)
+ServiceSpec = DefaultDatabaseSpec(metadata_source_class=MariadbSource)

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

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

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

@@ -15,6 +15,11 @@ import traceback
 from datetime import datetime
 from typing import Iterable
 
+from sqlalchemy.exc import OperationalError
+
+from metadata.generated.schema.entity.services.ingestionPipelines.status import (
+    StackTraceError,
+)
 from metadata.generated.schema.type.basic import DateTime
 from metadata.generated.schema.type.tableQuery import TableQueries, TableQuery
 from metadata.ingestion.source.connections import get_connection
@@ -67,6 +72,16 @@ class PostgresUsageSource(PostgresQueryParserSource, UsageSource):
                         logger.error(str(err))
             if queries:
                 yield TableQueries(queries=queries)
+
+        except OperationalError as err:
+            self.status.failed(
+                StackTraceError(
+                    name="Usage",
+                    error=f"Source Usage failed due to - {err}",
+                    stackTrace=traceback.format_exc(),
+                )
+            )
+
         except Exception as err:
             if query:
                 logger.debug(

metadata/ingestion/source/database/query_parser_source.py

@@ -67,6 +67,7 @@ class QueryParserSource(Source, ABC):
         self.engine = (
             get_ssl_connection(self.service_connection) if get_engine else None
         )
+        self.graph = None
 
     @property
     def name(self) -> str:

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

@@ -11,7 +11,6 @@
 """
 Redshift source ingestion
 """
-
 import re
 import traceback
 from typing import Iterable, List, Optional, Tuple
@@ -57,6 +56,9 @@ from metadata.ingestion.source.database.common_db_source import (
     CommonDbSourceService,
     TableNameAndType,
 )
+from metadata.ingestion.source.database.external_table_lineage_mixin import (
+    ExternalTableLineageMixin,
+)
 from metadata.ingestion.source.database.incremental_metadata_extraction import (
     IncrementalConfig,
 )
@@ -69,6 +71,7 @@ from metadata.ingestion.source.database.redshift.incremental_table_processor imp
 )
 from metadata.ingestion.source.database.redshift.models import RedshiftStoredProcedure
 from metadata.ingestion.source.database.redshift.queries import (
+    REDSHIFT_EXTERNAL_TABLE_LOCATION,
     REDSHIFT_GET_ALL_RELATION_INFO,
     REDSHIFT_GET_DATABASE_NAMES,
     REDSHIFT_GET_STORED_PROCEDURES,
@@ -121,12 +124,13 @@ RedshiftDialect.get_view_definition = get_view_definition
 RedshiftDialect._get_all_relation_info = (  # pylint: disable=protected-access
     _get_all_relation_info
 )
-
 Inspector.get_all_table_ddls = get_all_table_ddls
 Inspector.get_table_ddl = get_table_ddl
 
 
-class RedshiftSource(LifeCycleQueryMixin, CommonDbSourceService, MultiDBSource):
+class RedshiftSource(
+    ExternalTableLineageMixin, LifeCycleQueryMixin, CommonDbSourceService, MultiDBSource
+):
     """
     Implements the necessary methods to extract
     Database metadata from Redshift Source
@@ -146,6 +150,7 @@ class RedshiftSource(LifeCycleQueryMixin, CommonDbSourceService, MultiDBSource):
         self.incremental_table_processor: Optional[
             RedshiftIncrementalTableProcessor
         ] = None
+        self.external_location_map = {}
 
         if self.incremental.enabled:
             logger.info(
@@ -168,6 +173,14 @@ class RedshiftSource(LifeCycleQueryMixin, CommonDbSourceService, MultiDBSource):
         )
         return cls(config, metadata, incremental_config)
 
+    def get_location_path(self, table_name: str, schema_name: str) -> Optional[str]:
+        """
+        Method to fetch the location path of the table
+        """
+        return self.external_location_map.get(
+            (self.context.get().database, schema_name, table_name)
+        )
+
     def get_partition_details(self) -> None:
         """
         Populate partition details
@@ -275,15 +288,23 @@ class RedshiftSource(LifeCycleQueryMixin, CommonDbSourceService, MultiDBSource):
                 for schema_name, table_name in self.incremental_table_processor.get_deleted()
             )
 
+    def set_external_location_map(self, database_name: str) -> None:
+        self.external_location_map.clear()
+        results = self.engine.execute(
+            REDSHIFT_EXTERNAL_TABLE_LOCATION.format(database_name=database_name)
+        ).all()
+        self.external_location_map = {
+            (database_name, row.schemaname, row.tablename): row.location
+            for row in results
+        }
+
     def get_database_names(self) -> Iterable[str]:
         if not self.config.serviceConnection.root.config.ingestAllDatabases:
+            configured_db = self.config.serviceConnection.root.config.database
             self.get_partition_details()
-
-            self._set_incremental_table_processor(
-                self.config.serviceConnection.root.config.database
-            )
-
-            yield self.config.serviceConnection.root.config.database
+            self._set_incremental_table_processor(configured_db)
+            self.set_external_location_map(configured_db)
+            yield configured_db
         else:
             for new_database in self.get_database_names_raw():
                 database_fqn = fqn.build(
@@ -307,9 +328,8 @@ class RedshiftSource(LifeCycleQueryMixin, CommonDbSourceService, MultiDBSource):
                 try:
                     self.set_inspector(database_name=new_database)
                     self.get_partition_details()
-
                     self._set_incremental_table_processor(new_database)
-
+                    self.set_external_location_map(new_database)
                     yield new_database
                 except Exception as exc:
                     logger.debug(traceback.format_exc())

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

@@ -203,6 +203,11 @@ REDSHIFT_GET_SCHEMA_COLUMN_INFO = textwrap.dedent(
             """
 )
 
+REDSHIFT_EXTERNAL_TABLE_LOCATION = """
+  SELECT schemaname, tablename, location
+    FROM svv_external_tables
+    where redshift_database_name='{database_name}'
+"""
 
 REDSHIFT_PARTITION_DETAILS = """
   select "schema", "table", diststyle

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

@@ -1,6 +1,3 @@
-from metadata.ingestion.source.database.singlestore.lineage import (
-    SinglestoreLineageSource,
-)
 from metadata.ingestion.source.database.singlestore.metadata import SinglestoreSource
 from metadata.profiler.interface.sqlalchemy.single_store.profiler_interface import (
     SingleStoreProfilerInterface,
@@ -8,7 +5,5 @@ from metadata.profiler.interface.sqlalchemy.single_store.profiler_interface impo
 from metadata.utils.service_spec.default import DefaultDatabaseSpec
 
 ServiceSpec = DefaultDatabaseSpec(
-    metadata_source_class=SinglestoreSource,
-    profiler_class=SingleStoreProfilerInterface,
-    lineage_source_class=SinglestoreLineageSource,
+    metadata_source_class=SinglestoreSource, profiler_class=SingleStoreProfilerInterface
 )

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

@@ -193,10 +193,18 @@ def test_connection(
             engine_wrapper=engine_wrapper,
         ),
         "GetQueries": partial(
-            test_query, statement=SNOWFLAKE_TEST_GET_QUERIES, engine=engine
+            test_query,
+            statement=SNOWFLAKE_TEST_GET_QUERIES.format(
+                account_usage=service_connection.accountUsageSchema
+            ),
+            engine=engine,
         ),
         "GetTags": partial(
-            test_query, statement=SNOWFLAKE_TEST_FETCH_TAG, engine=engine
+            test_query,
+            statement=SNOWFLAKE_TEST_FETCH_TAG.format(
+                account_usage=service_connection.accountUsageSchema
+            ),
+            engine=engine,
         ),
     }
 

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

@@ -47,6 +47,8 @@ class SnowflakeLineageSource(
         AND (
             QUERY_TYPE IN ('MERGE', 'UPDATE','CREATE_TABLE_AS_SELECT')
             OR (QUERY_TYPE = 'INSERT' and query_text ILIKE '%%insert%%into%%select%%')
+            OR (QUERY_TYPE = 'ALTER' and query_text ILIKE '%%alter%%table%%swap%%')
+            OR (QUERY_TYPE = 'CREATE_TABLE' and query_text ILIKE '%%clone%%')
         )
     """
 
@@ -60,6 +62,7 @@ class SnowflakeLineageSource(
         start, _ = get_start_and_end(self.source_config.queryLogDuration)
         query = self.stored_procedure_query.format(
             start_date=start,
+            account_usage=self.service_connection.accountUsageSchema,
         )
         queries_dict = self.procedure_queries_dict(
             query=query,

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

@@ -418,6 +418,7 @@ class SnowflakeSource(
                     SNOWFLAKE_FETCH_ALL_TAGS.format(
                         database_name=self.context.get().database,
                         schema_name=schema_name,
+                        account_usage=self.service_connection.accountUsageSchema,
                     )
                 )
 
@@ -431,6 +432,7 @@ class SnowflakeSource(
                         SNOWFLAKE_FETCH_ALL_TAGS.format(
                             database_name=f'"{self.context.get().database}"',
                             schema_name=f'"{self.context.get().database_schema}"',
+                            account_usage=self.service_connection.accountUsageSchema,
                         )
                     )
                 except Exception as inner_exc:
@@ -635,6 +637,7 @@ class SnowflakeSource(
             query.format(
                 database_name=self.context.get().database,
                 schema_name=self.context.get().database_schema,
+                account_usage=self.service_connection.accountUsageSchema,
             )
         ).all()
         for row in results:

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

@@ -25,7 +25,7 @@ SNOWFLAKE_SQL_STATEMENT = textwrap.dedent(
       start_time "start_time",
       end_time "end_time",
       total_elapsed_time "duration"
-    from snowflake.account_usage.query_history
+    from {account_usage}.query_history
     WHERE query_text NOT LIKE '/* {{"app": "OpenMetadata", %%}} */%%'
     AND query_text NOT LIKE '/* {{"app": "dbt", %%}} */%%'
     AND start_time between to_timestamp_ltz('{start_time}') and to_timestamp_ltz('{end_time}')
@@ -39,7 +39,7 @@ SNOWFLAKE_SESSION_TAG_QUERY = 'ALTER SESSION SET QUERY_TAG="{query_tag}"'
 SNOWFLAKE_FETCH_ALL_TAGS = textwrap.dedent(
     """
     select TAG_NAME, TAG_VALUE, OBJECT_DATABASE, OBJECT_SCHEMA, OBJECT_NAME, COLUMN_NAME
-    from snowflake.account_usage.tag_references
+    from {account_usage}.tag_references
     where OBJECT_DATABASE = '{database_name}'
       and OBJECT_SCHEMA = '{schema_name}'
 """
@@ -234,11 +234,11 @@ SHOW EXTERNAL TABLES IN DATABASE "{database_name}"
 """
 
 SNOWFLAKE_TEST_FETCH_TAG = """
-select TAG_NAME from snowflake.account_usage.tag_references limit 1
+select TAG_NAME from {account_usage}.tag_references limit 1
 """
 
 SNOWFLAKE_TEST_GET_QUERIES = """
-SELECT query_text from snowflake.account_usage.query_history limit 1
+SELECT query_text from {account_usage}.query_history limit 1
 """
 
 SNOWFLAKE_TEST_GET_TABLES = """
@@ -296,7 +296,7 @@ SELECT
   ARGUMENT_SIGNATURE AS signature,
   COMMENT as comment,
   'StoredProcedure' as procedure_type
-FROM SNOWFLAKE.ACCOUNT_USAGE.PROCEDURES
+FROM {account_usage}.PROCEDURES
 WHERE PROCEDURE_CATALOG = '{database_name}'
   AND PROCEDURE_SCHEMA = '{schema_name}'
   AND DELETED IS NULL
@@ -313,7 +313,7 @@ SELECT
   ARGUMENT_SIGNATURE AS signature,
   COMMENT as comment,
   'UDF' as procedure_type
-FROM SNOWFLAKE.ACCOUNT_USAGE.FUNCTIONS
+FROM {account_usage}.FUNCTIONS
 WHERE FUNCTION_CATALOG = '{database_name}'
   AND FUNCTION_SCHEMA = '{schema_name}'
   AND DELETED IS NULL
@@ -336,7 +336,7 @@ WITH SP_HISTORY AS (
       SESSION_ID,
       START_TIME,
       END_TIME
-    FROM SNOWFLAKE.ACCOUNT_USAGE.QUERY_HISTORY SP
+    FROM {account_usage}.QUERY_HISTORY SP
     WHERE QUERY_TYPE = 'CALL'
       AND START_TIME >= '{start_date}'
       AND QUERY_TEXT <> ''
@@ -353,7 +353,7 @@ Q_HISTORY AS (
       USER_NAME,
       SCHEMA_NAME,
       DATABASE_NAME
-    FROM SNOWFLAKE.ACCOUNT_USAGE.QUERY_HISTORY SP
+    FROM {account_usage}.QUERY_HISTORY SP
     WHERE QUERY_TYPE <> 'CALL'
       AND QUERY_TEXT NOT LIKE '/* {{"app": "OpenMetadata", %%}} */%%'
       AND QUERY_TEXT NOT LIKE '/* {{"app": "dbt", %%}} */%%'

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

@@ -60,6 +60,7 @@ class SnowflakeQueryParserSource(QueryParserSource, ABC):
             end_time=end_time,
             result_limit=self.config.sourceConfig.config.resultLimit,
             filters=self.get_filters(),
+            account_usage=self.service_connection.accountUsageSchema,
         )
 
     def check_life_cycle_query(

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

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