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

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)

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

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

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

@@ -39,6 +39,7 @@ 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 (
@@ -135,6 +136,10 @@ def get_connection(connection: TrinoConnection) -> Engine:
     # 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
+    from trino.sqlalchemy.dialect import TrinoDialect
+
+    TrinoDialect.is_disconnect = _is_disconnect
+
     connection_copy = deepcopy(connection)
     if connection_copy.verify:
         connection_copy.connectionArguments = (
@@ -183,3 +188,11 @@ def test_connection(
         queries=queries,
         timeout_seconds=timeout_seconds,
     )
+
+
+# pylint: disable=unused-argument
+def _is_disconnect(self, e, connection, cursor):
+    """is_disconnect method for the Databricks dialect"""
+    if "JWT expired" in str(e):
+        return True
+    return False

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

@@ -11,6 +11,7 @@
 """
 Databricks Unity Catalog Lineage Source Module
 """
+import traceback
 from typing import Iterable, Optional
 
 from metadata.generated.schema.api.lineage.addLineage import AddLineageRequest
@@ -27,6 +28,7 @@ from metadata.generated.schema.type.entityLineage import (
     EntitiesEdge,
     LineageDetails,
 )
+from metadata.generated.schema.type.entityLineage import Source as LineageSource
 from metadata.generated.schema.type.entityReference import EntityReference
 from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import InvalidSourceException, Source
@@ -109,9 +111,59 @@ class UnitycatalogLineageSource(Source):
                     )
                 )
         if col_lineage:
-            return LineageDetails(columnsLineage=col_lineage)
+            return LineageDetails(
+                columnsLineage=col_lineage, source=LineageSource.QueryLineage
+            )
         return None
 
+    def _handle_upstream_table(
+        self,
+        table_streams: LineageTableStreams,
+        table: Table,
+        databricks_table_fqn: str,
+    ) -> Iterable[Either[AddLineageRequest]]:
+        for upstream_table in table_streams.upstream_tables:
+            try:
+                if not upstream_table.name:
+                    continue
+                from_entity_fqn = fqn.build(
+                    metadata=self.metadata,
+                    entity_type=Table,
+                    database_name=upstream_table.catalog_name,
+                    schema_name=upstream_table.schema_name,
+                    table_name=upstream_table.name,
+                    service_name=self.config.serviceName,
+                )
+
+                from_entity = self.metadata.get_by_name(
+                    entity=Table, fqn=from_entity_fqn
+                )
+                if from_entity:
+                    lineage_details = self._get_lineage_details(
+                        from_table=from_entity,
+                        to_table=table,
+                        databricks_table_fqn=databricks_table_fqn,
+                    )
+                    yield Either(
+                        left=None,
+                        right=AddLineageRequest(
+                            edge=EntitiesEdge(
+                                toEntity=EntityReference(id=table.id, type="table"),
+                                fromEntity=EntityReference(
+                                    id=from_entity.id, type="table"
+                                ),
+                                lineageDetails=lineage_details,
+                            )
+                        ),
+                    )
+            except Exception:
+                logger.debug(
+                    "Error while processing lineage for "
+                    f"{upstream_table.catalog_name}.{upstream_table.schema_name}.{upstream_table.name}"
+                    f" -> {databricks_table_fqn}"
+                )
+                logger.debug(traceback.format_exc())
+
     def _iter(self, *_, **__) -> Iterable[Either[AddLineageRequest]]:
         """
         Based on the query logs, prepare the lineage
@@ -128,37 +180,9 @@ class UnitycatalogLineageSource(Source):
                 table_streams: LineageTableStreams = self.client.get_table_lineage(
                     databricks_table_fqn
                 )
-                for upstream_table in table_streams.upstream_tables:
-                    from_entity_fqn = fqn.build(
-                        metadata=self.metadata,
-                        entity_type=Table,
-                        database_name=upstream_table.catalog_name,
-                        schema_name=upstream_table.schema_name,
-                        table_name=upstream_table.name,
-                        service_name=self.config.serviceName,
-                    )
-
-                    from_entity = self.metadata.get_by_name(
-                        entity=Table, fqn=from_entity_fqn
-                    )
-                    if from_entity:
-                        lineage_details = self._get_lineage_details(
-                            from_table=from_entity,
-                            to_table=table,
-                            databricks_table_fqn=databricks_table_fqn,
-                        )
-                        yield Either(
-                            left=None,
-                            right=AddLineageRequest(
-                                edge=EntitiesEdge(
-                                    toEntity=EntityReference(id=table.id, type="table"),
-                                    fromEntity=EntityReference(
-                                        id=from_entity.id, type="table"
-                                    ),
-                                    lineageDetails=lineage_details,
-                                )
-                            ),
-                        )
+                yield from self._handle_upstream_table(
+                    table_streams, table, databricks_table_fqn
+                )
 
     def test_connection(self) -> None:
         test_connection_fn = get_test_connection_fn(self.service_connection)

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

@@ -44,6 +44,13 @@ class UnityCatalogQueryParserSource(
 
     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)

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

@@ -11,6 +11,9 @@ from metadata.ingestion.source.database.unitycatalog.usage import (
 from metadata.profiler.interface.sqlalchemy.unity_catalog.profiler_interface import (
     UnityCatalogProfilerInterface,
 )
+from metadata.profiler.interface.sqlalchemy.unity_catalog.sampler_interface import (
+    UnityCatalogSamplerInterface,
+)
 from metadata.utils.service_spec.default import DefaultDatabaseSpec
 
 ServiceSpec = DefaultDatabaseSpec(
@@ -19,4 +22,5 @@ ServiceSpec = DefaultDatabaseSpec(
     usage_source_class=UnitycatalogUsageSource,
     profiler_class=UnityCatalogProfilerInterface,
     test_suite_class=UnityCatalogTestSuiteInterface,
+    sampler_class=UnityCatalogSamplerInterface,
 )

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

@@ -11,17 +11,22 @@
 """
 unity catalog usage module
 """
+import traceback
+from datetime import datetime
+from typing import Iterable
 
-from metadata.ingestion.source.database.databricks.usage import DatabricksUsageSource
+from metadata.generated.schema.type.basic import DateTime
+from metadata.generated.schema.type.tableQuery import TableQueries, TableQuery
 from metadata.ingestion.source.database.unitycatalog.query_parser import (
     UnityCatalogQueryParserSource,
 )
+from metadata.ingestion.source.database.usage_source import UsageSource
 from metadata.utils.logger import ingestion_logger
 
 logger = ingestion_logger()
 
 
-class UnitycatalogUsageSource(UnityCatalogQueryParserSource, DatabricksUsageSource):
+class UnitycatalogUsageSource(UnityCatalogQueryParserSource, UsageSource):
     """
     UnityCatalog Usage Source
 
@@ -29,3 +34,37 @@ class UnitycatalogUsageSource(UnityCatalogQueryParserSource, DatabricksUsageSour
     DatabricksUsageSource as both the sources would call
     the same API for fetching Usage Queries
     """
+
+    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}"
+                )
+
+        yield TableQueries(queries=queries)

metadata/ingestion/source/database/usage_source.py

@@ -153,7 +153,7 @@ class UsageSource(QueryParserSource, ABC):
                 if query:
                     logger.debug(
                         (
-                            f"###### USAGE QUERY #######\n{mask_query(query, self.dialect.value)}"
+                            f"###### USAGE QUERY #######\n{mask_query(query, self.dialect.value) or query}"
                             "\n##########################"
                         )
                     )

metadata/ingestion/source/pipeline/nifi/connection.py

@@ -17,8 +17,10 @@ from typing import Optional
 from metadata.generated.schema.entity.automations.workflow import (
     Workflow as AutomationWorkflow,
 )
+from metadata.generated.schema.entity.services.connections.pipeline.nifi.basicAuth import (
+    NifiBasicAuth,
+)
 from metadata.generated.schema.entity.services.connections.pipeline.nifiConnection import (
-    BasicAuthentication,
     NifiConnection,
 )
 from metadata.generated.schema.entity.services.connections.testConnectionResult import (
@@ -34,7 +36,7 @@ def get_connection(connection: NifiConnection) -> NifiClient:
     """
     Create connection
     """
-    if isinstance(connection.nifiConfig, BasicAuthentication):
+    if isinstance(connection.nifiConfig, NifiBasicAuth):
         return NifiClient(
             host_port=connection.hostPort,
             username=connection.nifiConfig.username,

metadata/profiler/interface/sqlalchemy/unity_catalog/sampler_interface.py

@@ -0,0 +1,29 @@
+#  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.
+
+"""
+Interfaces with database for all database engine
+supporting sqlalchemy abstraction layer
+"""
+from metadata.ingestion.source.database.databricks.connection import (
+    get_connection as databricks_get_connection,
+)
+from metadata.sampler.sqlalchemy.sampler import SQASampler
+
+
+class UnityCatalogSamplerInterface(SQASampler):
+    def get_client(self):
+        """client is the session for SQA"""
+        self.connection = databricks_get_connection(self.service_connection_config)
+        self.client = super().get_client()
+        self.set_catalog(self.client)
+
+        return self.client