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

metadata/ingestion/ometa/routes.py

@@ -92,6 +92,7 @@ from metadata.generated.schema.api.services.createStorageService import (
 from metadata.generated.schema.api.services.ingestionPipelines.createIngestionPipeline import (
     CreateIngestionPipelineRequest,
 )
+from metadata.generated.schema.api.teams.createPersona import CreatePersonaRequest
 from metadata.generated.schema.api.teams.createRole import CreateRoleRequest
 from metadata.generated.schema.api.teams.createTeam import CreateTeamRequest
 from metadata.generated.schema.api.teams.createUser import CreateUserRequest
@@ -157,6 +158,7 @@ from metadata.generated.schema.entity.services.mlmodelService import MlModelServ
 from metadata.generated.schema.entity.services.pipelineService import PipelineService
 from metadata.generated.schema.entity.services.searchService import SearchService
 from metadata.generated.schema.entity.services.storageService import StorageService
+from metadata.generated.schema.entity.teams.persona import Persona
 from metadata.generated.schema.entity.teams.role import Role
 from metadata.generated.schema.entity.teams.team import Team
 from metadata.generated.schema.entity.teams.user import AuthenticationMechanism, User
@@ -217,6 +219,8 @@ ROUTES = {
     CreateTeamRequest.__name__: "/teams",
     User.__name__: "/users",
     CreateUserRequest.__name__: "/users",
+    Persona.__name__: "/personas",
+    CreatePersonaRequest.__name__: "/personas",
     AuthenticationMechanism.__name__: "/users/auth-mechanism",
     Bot.__name__: "/bots",
     CreateBot.__name__: "/bots",

metadata/ingestion/source/api/rest/metadata.py

@@ -171,7 +171,7 @@ class RestSource(ApiServiceSource):
                     break
             return filtered_paths
         except Exception as err:
-            logger.info(
+            logger.warning(
                 f"Error while filtering endpoints for collection {collection.name.root}"
             )
             return None
@@ -184,7 +184,7 @@ class RestSource(ApiServiceSource):
                 endpoint.name = f"{path} - {method_type}"
             return endpoint
         except Exception as err:
-            logger.info(f"Error while parsing endpoint data: {err}")
+            logger.warning(f"Error while parsing endpoint data: {err}")
         return None
 
     def _generate_collection_url(self, collection_name: str) -> Optional[AnyUrl]:
@@ -195,7 +195,7 @@ class RestSource(ApiServiceSource):
                     f"{self.config.serviceConnection.root.config.openAPISchemaURL}#tag/{collection_name}"
                 )
         except Exception as err:
-            logger.info(f"Error while generating collection url: {err}")
+            logger.warning(f"Error while generating collection url: {err}")
         return None
 
     def _generate_endpoint_url(self, endpoint_name: str) -> AnyUrl:
@@ -203,15 +203,14 @@ class RestSource(ApiServiceSource):
         base_url = self.config.serviceConnection.root.config.openAPISchemaURL
         if endpoint_name:
             return AnyUrl(f"{base_url}#operation/{endpoint_name}")
-        else:
-            return AnyUrl(base_url)
+        return AnyUrl(base_url)
 
     def _get_api_request_method(self, method_type: str) -> Optional[str]:
         """fetch endpoint request method"""
         try:
             return ApiRequestMethod[method_type.upper()]
         except KeyError as err:
-            logger.info(f"Keyerror while fetching request method: {err}")
+            logger.warning(f"Keyerror while fetching request method: {err}")
         return None
 
     def _get_request_schema(self, info: dict) -> Optional[APISchema]:
@@ -227,9 +226,9 @@ class RestSource(ApiServiceSource):
             if not schema_ref:
                 logger.debug("No request schema found for the endpoint")
                 return None
-            return self._process_schema(schema_ref)
+            return APISchema(schemaFields=self.process_schema_fields(schema_ref))
         except Exception as err:
-            logger.info(f"Error while parsing request schema: {err}")
+            logger.warning(f"Error while parsing request schema: {err}")
         return None
 
     def _get_response_schema(self, info: dict) -> Optional[APISchema]:
@@ -246,33 +245,43 @@ class RestSource(ApiServiceSource):
             if not schema_ref:
                 logger.debug("No response schema found for the endpoint")
                 return None
-            return self._process_schema(schema_ref)
+            return APISchema(schemaFields=self.process_schema_fields(schema_ref))
         except Exception as err:
-            logger.info(f"Error while parsing response schema: {err}")
+            logger.warning(f"Error while parsing response schema: {err}")
         return None
 
-    def _process_schema(self, schema_ref: str) -> Optional[List[APISchema]]:
-        """process schema"""
+    def process_schema_fields(self, schema_ref: str) -> Optional[List[FieldModel]]:
         try:
-            schema_ref = schema_ref.split("/")[-1]
+            schema_name = schema_ref.split("/")[-1]
             schema_fields = (
-                self.json_response.get("components").get("schemas").get(schema_ref)
+                self.json_response.get("components").get("schemas").get(schema_name)
             )
 
             fetched_fields = []
             for key, val in schema_fields.get("properties", {}).items():
                 dtype = val.get("type")
-                if not dtype:
-                    continue
-                fetched_fields.append(
-                    FieldModel(
-                        name=key,
-                        dataType=DataTypeTopic[dtype.upper()]
+                if dtype:
+                    parsed_dtype = (
+                        DataTypeTopic[dtype.upper()]
                         if dtype.upper() in DataTypeTopic.__members__
-                        else DataTypeTopic.UNKNOWN,
+                        else DataTypeTopic.UNKNOWN
                     )
-                )
-            return APISchema(schemaFields=fetched_fields)
+                    fetched_fields.append(FieldModel(name=key, dataType=parsed_dtype))
+                else:
+                    # If type of field is not defined then check for sub-schema
+                    # Check if it's `object` type field
+                    # check infinite recrusrion by comparing with parent(schema_ref)
+                    object_children = None
+                    if val.get("$ref") and val.get("$ref") != schema_ref:
+                        object_children = self.process_schema_fields(val.get("$ref"))
+                    fetched_fields.append(
+                        FieldModel(
+                            name=key,
+                            dataType=DataTypeTopic.UNKNOWN,
+                            children=object_children,
+                        )
+                    )
+            return fetched_fields
         except Exception as err:
-            logger.info(f"Error while processing request schema: {err}")
+            logger.warning(f"Error while processing schema fields: {err}")
         return None

metadata/ingestion/source/dashboard/looker/metadata.py

@@ -21,6 +21,7 @@ Notes:
 """
 import copy
 import os
+import re
 import traceback
 from datetime import datetime
 from pathlib import Path
@@ -39,6 +40,7 @@ from typing import (
 
 import giturlparse
 import lkml
+import networkx as nx
 from liquid import Template
 from looker_sdk.sdk.api40.methods import Looker40SDK
 from looker_sdk.sdk.api40.models import Dashboard as LookerDashboard
@@ -134,6 +136,10 @@ logger = ingestion_logger()
 LIST_DASHBOARD_FIELDS = ["id", "title"]
 IMPORTED_PROJECTS_DIR = "imported_projects"
 
+# we need to find the derived references in the SQL query using regex
+# https://cloud.google.com/looker/docs/derived-tables#referencing_derived_tables_in_other_derived_tables
+DERIVED_REFERENCES = r"\${([\w\s\d_.]+)\.SQL_TABLE_NAME}"
+
 # Here we can update the fields to get further information, such as:
 # created_at, updated_at, last_updater_id, deleted_at, deleter_id, favorite_count, last_viewed_at
 GET_DASHBOARD_FIELDS = [
@@ -165,6 +171,13 @@ def build_datamodel_name(model_name: str, explore_name: str) -> str:
     return clean_dashboard_name(model_name + "_" + explore_name)
 
 
+def find_derived_references(sql_query: str) -> List[str]:
+    if sql_query is None:
+        return []
+    matches = re.findall(DERIVED_REFERENCES, sql_query)
+    return matches
+
+
 class LookerSource(DashboardServiceSource):
     """
     Looker Source Class.
@@ -172,6 +185,8 @@ class LookerSource(DashboardServiceSource):
     Its client uses Looker 40 from the SDK: client = looker_sdk.init40()
     """
 
+    # pylint: disable=too-many-instance-attributes
+
     config: WorkflowSource
     metadata: OpenMetadata
     client: Looker40SDK
@@ -192,6 +207,10 @@ class LookerSource(DashboardServiceSource):
         self._main__lookml_manifest: Optional[LookMLManifest] = None
         self._view_data_model: Optional[DashboardDataModel] = None
 
+        self._parsed_views: Optional[Dict[str, str]] = {}
+        self._unparsed_views: Optional[Dict[str, str]] = {}
+        self._derived_dependencies = nx.DiGraph()
+
         self._added_lineage: Optional[Dict] = {}
 
     @classmethod
@@ -557,6 +576,68 @@ class LookerSource(DashboardServiceSource):
                     )
                 )
 
+    def replace_derived_references(self, sql_query):
+        """
+        Replace all derived references with the parsed views sql query
+        will replace the derived references in the SQL query using regex
+        for e.g. It will replace ${view_name.SQL_TABLE_NAME} with the parsed view query for view_name
+        https://cloud.google.com/looker/docs/derived-tables#referencing_derived_tables_in_other_derived_tables
+        """
+        try:
+            sql_query = re.sub(
+                DERIVED_REFERENCES,
+                # from `${view_name.SQL_TABLE_NAME}` we want the `view_name`.
+                # match.group(1) will give us the `view_name`
+                lambda match: f"({self._parsed_views.get(match.group(1), match.group(0))})",
+                sql_query,
+            )
+        except Exception as e:
+            logger.warning(
+                f"Something went wrong while replacing derived view references: {e}"
+            )
+        return sql_query
+
+    def build_lineage_for_unparsed_views(self) -> Iterable[Either[AddLineageRequest]]:
+        """
+        build lineage by parsing the unparsed views containing derived references
+        """
+        try:
+            # Doing a reversed topological sort to process the views in the right order
+            for view_name in reversed(
+                list(nx.topological_sort(self._derived_dependencies))
+            ):
+                if view_name in self._parsed_views:
+                    # Skip if already processed
+                    continue
+                sql_query = self.replace_derived_references(
+                    self._unparsed_views[view_name]
+                )
+                if view_references := find_derived_references(sql_query):
+                    # There are still derived references in the view query
+                    logger.debug(
+                        f"Views {view_references} not found for {view_name}. Skipping."
+                    )
+                    continue
+                self._parsed_views[view_name] = sql_query
+                del self._unparsed_views[view_name]
+                yield from self._build_lineage_for_view(view_name, sql_query)
+
+        except Exception as err:
+            yield Either(
+                left=StackTraceError(
+                    name="parse_unparsed_views",
+                    error=f"Error parsing unparsed views: {err}",
+                    stackTrace=traceback.format_exc(),
+                )
+            )
+
+    def _add_dependency_edge(self, view_name: str, view_references: List[str]):
+        """
+        Add a dependency edge between the view and the derived reference
+        """
+        for dependent_view_name in view_references:
+            self._derived_dependencies.add_edge(view_name, dependent_view_name)
+
     def add_view_lineage(
         self, view: LookMlView, explore: LookmlModelExplore
     ) -> Iterable[Either[AddLineageRequest]]:
@@ -589,6 +670,7 @@ class LookerSource(DashboardServiceSource):
                 for db_service_name in db_service_names or []:
                     dialect = self._get_db_dialect(db_service_name)
                     source_table_name = self._clean_table_name(sql_table_name, dialect)
+                    self._parsed_views[view.name] = source_table_name
 
                     # View to the source is only there if we are informing the dbServiceNames
                     yield self.build_lineage_request(
@@ -601,20 +683,19 @@ class LookerSource(DashboardServiceSource):
                 sql_query = view.derived_table.sql
                 if not sql_query:
                     return
+                if find_derived_references(sql_query):
+                    sql_query = self.replace_derived_references(sql_query)
+                    # If we still have derived references, we cannot process the view
+                    if view_references := find_derived_references(sql_query):
+                        self._add_dependency_edge(view.name, view_references)
+                        logger.warning(
+                            f"Not all references are replaced for view [{view.name}]. Parsing it later."
+                        )
+                        return
                 logger.debug(f"Processing view [{view.name}] with SQL: \n[{sql_query}]")
-                for db_service_name in db_service_names or []:
-                    lineage_parser = LineageParser(
-                        sql_query,
-                        self._get_db_dialect(db_service_name),
-                        timeout_seconds=30,
-                    )
-                    if lineage_parser.source_tables:
-                        for from_table_name in lineage_parser.source_tables:
-                            yield self.build_lineage_request(
-                                source=str(from_table_name),
-                                db_service_name=db_service_name,
-                                to_entity=self._view_data_model,
-                            )
+                yield from self._build_lineage_for_view(view.name, sql_query)
+                if self._unparsed_views:
+                    self.build_lineage_for_unparsed_views()
 
         except Exception as err:
             yield Either(
@@ -625,6 +706,27 @@ class LookerSource(DashboardServiceSource):
                 )
             )
 
+    def _build_lineage_for_view(
+        self, view_name: str, sql_query: str
+    ) -> Iterable[Either[AddLineageRequest]]:
+        """
+        Parse the SQL query and build lineage for the view.
+        """
+        for db_service_name in self.get_db_service_names() or []:
+            lineage_parser = LineageParser(
+                sql_query,
+                self._get_db_dialect(db_service_name),
+                timeout_seconds=30,
+            )
+            if lineage_parser.source_tables:
+                self._parsed_views[view_name] = sql_query
+                for from_table_name in lineage_parser.source_tables:
+                    yield self.build_lineage_request(
+                        source=str(from_table_name),
+                        db_service_name=db_service_name,
+                        to_entity=self._view_data_model,
+                    )
+
     def _get_db_dialect(self, db_service_name) -> Dialect:
         db_service = self.metadata.get_by_name(DatabaseService, db_service_name)
         return ConnectionTypeDialectMapper.dialect_of(

metadata/ingestion/source/dashboard/powerbi/client.py

@@ -47,7 +47,10 @@ from metadata.utils.logger import utils_logger
 
 logger = utils_logger()
 
-
+GETGROUPS_DEFAULT_PARAMS = {"$top": "1", "$skip": "0"}
+API_RESPONSE_MESSAGE_KEY = "message"
+AUTH_TOKEN_MAX_RETRIES = 5
+AUTH_TOKEN_RETRY_WAIT = 120
 # Similar inner methods with mode client. That's fine.
 # pylint: disable=duplicate-code
 class PowerBiApiClient:
@@ -59,6 +62,9 @@ class PowerBiApiClient:
 
     def __init__(self, config: PowerBIConnection):
         self.config = config
+        self.pagination_entity_per_page = min(
+            100, self.config.pagination_entity_per_page
+        )
         self.msal_client = msal.ConfidentialClientApplication(
             client_id=self.config.clientId,
             client_credential=self.config.clientSecret.get_secret_value(),
@@ -82,42 +88,84 @@ class PowerBiApiClient:
         """
         logger.info("Generating PowerBi access token")
 
-        response_data = self.msal_client.acquire_token_silent(
-            scopes=self.config.scope, account=None
-        )
-
+        response_data = self.get_auth_token_from_cache()
         if not response_data:
             logger.info("Token does not exist in the cache. Getting a new token.")
-            response_data = self.msal_client.acquire_token_for_client(
-                scopes=self.config.scope
-            )
+            response_data = self.generate_new_auth_token()
+        response_data = response_data or {}
         auth_response = PowerBiToken(**response_data)
         if not auth_response.access_token:
             raise InvalidSourceException(
-                "Failed to generate the PowerBi access token. Please check provided config"
+                f"Failed to generate the PowerBi access token. Please check provided config {response_data}"
             )
 
         logger.info("PowerBi Access Token generated successfully")
         return auth_response.access_token, auth_response.expires_in
 
+    def generate_new_auth_token(self) -> Optional[dict]:
+        """generate new auth token"""
+        retry = AUTH_TOKEN_MAX_RETRIES
+        while retry:
+            try:
+                response_data = self.msal_client.acquire_token_for_client(
+                    scopes=self.config.scope
+                )
+                return response_data
+            except Exception as exc:
+                logger.debug(traceback.format_exc())
+                logger.warning(f"Error generating new auth token: {exc}")
+                # wait for time and retry
+                retry -= 1
+                if retry:
+                    logger.warning(
+                        f"Error generating new token: {exc}, "
+                        f"sleep {AUTH_TOKEN_RETRY_WAIT} seconds retrying {retry} more times.."
+                    )
+                    sleep(AUTH_TOKEN_RETRY_WAIT)
+                else:
+                    logger.warning(
+                        "Could not generate new token after maximum retries, "
+                        "Please check provided configs"
+                    )
+        return None
+
+    def get_auth_token_from_cache(self) -> Optional[dict]:
+        """fetch auth token from cache"""
+        retry = AUTH_TOKEN_MAX_RETRIES
+        while retry:
+            try:
+                response_data = self.msal_client.acquire_token_silent(
+                    scopes=self.config.scope, account=None
+                )
+                return response_data
+            except Exception as exc:
+                logger.debug(traceback.format_exc())
+                logger.warning(f"Error getting token from cache: {exc}")
+                retry -= 1
+                if retry:
+                    logger.warning(
+                        f"Error getting token from cache: {exc}, "
+                        f"sleep {AUTH_TOKEN_RETRY_WAIT} seconds retrying {retry} more times.."
+                    )
+                    sleep(AUTH_TOKEN_RETRY_WAIT)
+                else:
+                    logger.warning(
+                        "Could not get token from cache after maximum retries, "
+                        "Please check provided configs"
+                    )
+        return None
+
     def fetch_dashboards(self) -> Optional[List[PowerBIDashboard]]:
         """Get dashboards method
         Returns:
             List[PowerBIDashboard]
         """
-        try:
-            if self.config.useAdminApis:
-                response_data = self.client.get("/myorg/admin/dashboards")
-                response = DashboardsResponse(**response_data)
-                return response.value
-            group = self.fetch_all_workspaces()[0]
-            return self.fetch_all_org_dashboards(group_id=group.id)
-
-        except Exception as exc:  # pylint: disable=broad-except
-            logger.debug(traceback.format_exc())
-            logger.warning(f"Error fetching dashboards: {exc}")
-
-        return None
+        if self.config.useAdminApis:
+            response_data = self.client.get("/myorg/admin/dashboards")
+            response = DashboardsResponse(**response_data)
+            return response.value
+        group = self.fetch_all_workspaces()[0]
+        return self.fetch_all_org_dashboards(group_id=group.id)
 
     def fetch_all_org_dashboards(
         self, group_id: str
@@ -205,6 +253,7 @@ class PowerBiApiClient:
 
         return None
 
+    # pylint: disable=too-many-branches,too-many-statements
     def fetch_all_workspaces(self) -> Optional[List[Group]]:
         """Method to fetch all powerbi workspace details
         Returns:
@@ -213,28 +262,94 @@ class PowerBiApiClient:
         try:
             admin = "admin/" if self.config.useAdminApis else ""
             api_url = f"/myorg/{admin}groups"
-            entities_per_page = self.config.pagination_entity_per_page
-            params_data = {"$top": "1"}
-            response_data = self.client.get(api_url, data=params_data)
-            response = GroupsResponse(**response_data)
-            count = response.odata_count
+            entities_per_page = self.pagination_entity_per_page
+            failed_indexes = []
+            params_data = GETGROUPS_DEFAULT_PARAMS
+            response = self.client.get(api_url, data=params_data)
+            if (
+                not response
+                or API_RESPONSE_MESSAGE_KEY in response
+                or len(response) != len(GroupsResponse.__annotations__)
+            ):
+                logger.warning("Error fetching workspaces between results: (0, 1)")
+                if response and response.get(API_RESPONSE_MESSAGE_KEY):
+                    logger.warning(
+                        "Error message from API response: "
+                        f"{str(response.get(API_RESPONSE_MESSAGE_KEY))}"
+                    )
+                failed_indexes.append(params_data)
+                count = 0
+            else:
+                try:
+                    response = GroupsResponse(**response)
+                    count = response.odata_count
+                except Exception as exc:
+                    logger.warning(f"Error processing GetGroups response: {exc}")
+                    count = 0
             indexes = math.ceil(count / entities_per_page)
-
             workspaces = []
             for index in range(indexes):
                 params_data = {
                     "$top": str(entities_per_page),
                     "$skip": str(index * entities_per_page),
                 }
-                response_data = self.client.get(api_url, data=params_data)
-                if not response_data:
-                    logger.error(
-                        "Error fetching workspaces between results: "
-                        f"{str(index * entities_per_page)} - {str(entities_per_page)}"
+                response = self.client.get(api_url, data=params_data)
+                if (
+                    not response
+                    or API_RESPONSE_MESSAGE_KEY in response
+                    or len(response) != len(GroupsResponse.__annotations__)
+                ):
+                    index_range = (
+                        int(params_data.get("$skip")),
+                        int(params_data.get("$skip")) + int(params_data.get("$top")),
+                    )
+                    logger.warning(
+                        f"Error fetching workspaces between results: {str(index_range)}"
                     )
+                    if response and response.get(API_RESPONSE_MESSAGE_KEY):
+                        logger.warning(
+                            "Error message from API response: "
+                            f"{str(response.get(API_RESPONSE_MESSAGE_KEY))}"
+                        )
+                    failed_indexes.append(params_data)
                     continue
-                response = GroupsResponse(**response_data)
-                workspaces.extend(response.value)
+                try:
+                    response = GroupsResponse(**response)
+                    workspaces.extend(response.value)
+                except Exception as exc:
+                    logger.warning(f"Error processing GetGroups response: {exc}")
+
+            if failed_indexes:
+                logger.info(
+                    "Retrying one more time on failed indexes to get workspaces"
+                )
+                for params_data in failed_indexes:
+                    response = self.client.get(api_url, data=params_data)
+                    if (
+                        not response
+                        or API_RESPONSE_MESSAGE_KEY in response
+                        or len(response) != len(GroupsResponse.__annotations__)
+                    ):
+                        index_range = (
+                            int(params_data.get("$skip")),
+                            int(params_data.get("$skip"))
+                            + int(params_data.get("$top")),
+                        )
+                        logger.warning(
+                            f"Workspaces between results {str(index_range)} "
+                            "could not be fetched on multiple attempts"
+                        )
+                        if response and response.get(API_RESPONSE_MESSAGE_KEY):
+                            logger.warning(
+                                "Error message from API response: "
+                                f"{str(response.get(API_RESPONSE_MESSAGE_KEY))}"
+                            )
+                        continue
+                    try:
+                        response = GroupsResponse(**response)
+                        workspaces.extend(response.value)
+                    except Exception as exc:
+                        logger.warning(f"Error processing GetGroups response: {exc}")
             return workspaces
         except Exception as exc:  # pylint: disable=broad-except
             logger.debug(traceback.format_exc())

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]]:

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(