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

metadata/ingestion/source/database/bigquery/metadata.py

@@ -34,9 +34,11 @@ from metadata.generated.schema.entity.data.database import Database
 from metadata.generated.schema.entity.data.databaseSchema import DatabaseSchema
 from metadata.generated.schema.entity.data.storedProcedure import StoredProcedureCode
 from metadata.generated.schema.entity.data.table import (
+    ConstraintType,
     PartitionColumnDetails,
     PartitionIntervalTypes,
     Table,
+    TableConstraint,
     TablePartition,
     TableType,
 )
@@ -96,6 +98,7 @@ from metadata.ingestion.source.database.life_cycle_query_mixin import (
 from metadata.ingestion.source.database.multi_db_source import MultiDBSource
 from metadata.utils import fqn
 from metadata.utils.credentials import GOOGLE_CREDENTIALS
+from metadata.utils.execution_time_tracker import calculate_execution_time
 from metadata.utils.filters import filter_by_database, filter_by_schema
 from metadata.utils.logger import ingestion_logger
 from metadata.utils.sqlalchemy_utils import (
@@ -657,6 +660,42 @@ class BigquerySource(LifeCycleQueryMixin, CommonDbSourceService, MultiDBSource):
             )
         return None
 
+    @calculate_execution_time()
+    def update_table_constraints(
+        self,
+        table_name,
+        schema_name,
+        db_name,
+        table_constraints,
+        foreign_columns,
+        columns,
+    ) -> List[TableConstraint]:
+        """
+        From topology.
+        process the table constraints of all tables
+        """
+        table_constraints = super().update_table_constraints(
+            table_name,
+            schema_name,
+            db_name,
+            table_constraints,
+            foreign_columns,
+            columns,
+        )
+        try:
+            table = self.client.get_table(fqn._build(db_name, schema_name, table_name))
+            if hasattr(table, "clustering_fields") and table.clustering_fields:
+                table_constraints.append(
+                    TableConstraint(
+                        constraintType=ConstraintType.CLUSTER_KEY,
+                        columns=table.clustering_fields,
+                    )
+                )
+        except Exception as exc:
+            logger.warning(f"Error getting clustering fields for {table_name}: {exc}")
+            logger.debug(traceback.format_exc())
+        return table_constraints
+
     def get_table_partition_details(
         self, table_name: str, schema_name: str, inspector: Inspector
     ) -> Tuple[bool, Optional[TablePartition]]:
@@ -667,8 +706,10 @@ class BigquerySource(LifeCycleQueryMixin, CommonDbSourceService, MultiDBSource):
             database = self.context.get().database
             table = self.client.get_table(fqn._build(database, schema_name, table_name))
             columns = inspector.get_columns(table_name, schema_name, db_name=database)
-            if hasattr(table, "external_data_configuration") and hasattr(
-                table.external_data_configuration, "hive_partitioning"
+            if (
+                hasattr(table, "external_data_configuration")
+                and hasattr(table.external_data_configuration, "hive_partitioning")
+                and table.external_data_configuration.hive_partitioning
             ):
                 # Ingesting External Hive Partitioned Tables
                 from google.cloud.bigquery.external_config import (  # pylint: disable=import-outside-toplevel
@@ -733,6 +774,30 @@ class BigquerySource(LifeCycleQueryMixin, CommonDbSourceService, MultiDBSource):
                     table_partition.interval = table.range_partitioning.range_.interval
                 table_partition.columnName = table.range_partitioning.field
                 return True, TablePartition(columns=[table_partition])
+            if (
+                hasattr(table, "_properties")
+                and table._properties.get("partitionDefinition")
+                and table._properties.get("partitionDefinition").get(
+                    "partitionedColumn"
+                )
+            ):
+
+                return True, TablePartition(
+                    columns=[
+                        PartitionColumnDetails(
+                            columnName=self._get_partition_column_name(
+                                columns=columns,
+                                partition_field_name=field.get("field"),
+                            ),
+                            intervalType=PartitionIntervalTypes.OTHER,
+                        )
+                        for field in table._properties.get("partitionDefinition").get(
+                            "partitionedColumn"
+                        )
+                        if field and field.get("field")
+                    ]
+                )
+
         except Exception as exc:
             logger.debug(traceback.format_exc())
             logger.warning(

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

@@ -38,9 +38,9 @@ from metadata.ingestion.connections.test_connections import (
     test_connection_steps,
 )
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
-from metadata.ingestion.source.database.databricks.client import DatabricksClient
 from metadata.ingestion.source.database.databricks.queries import (
     DATABRICKS_GET_CATALOGS,
+    DATABRICKS_SQL_STATEMENT_TEST,
 )
 from metadata.utils.constants import THREE_MIN
 from metadata.utils.logger import ingestion_logger
@@ -81,7 +81,6 @@ def test_connection(
     Test connection. This can be executed either as part
     of a metadata workflow or during an Automation Workflow
     """
-    client = DatabricksClient(service_connection)
 
     def test_database_query(engine: Engine, statement: str):
         """
@@ -106,7 +105,13 @@ def test_connection(
             engine=connection,
             statement=DATABRICKS_GET_CATALOGS,
         ),
-        "GetQueries": client.test_query_api_access,
+        "GetQueries": partial(
+            test_database_query,
+            engine=connection,
+            statement=DATABRICKS_SQL_STATEMENT_TEST.format(
+                query_history=service_connection.queryHistoryTable
+            ),
+        ),
     }
 
     return test_connection_steps(

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