From 6acb47a87ae23c0ed784ba766e3b3b5c6c38c61b Mon Sep 17 00:00:00 2001
From: Mayur Singal <39544459+ulixius9@users.noreply.github.com>
Date: Wed, 27 Mar 2024 15:14:26 +0530
Subject: [PATCH] MINOR: Fix unity catalog test connection (#15712)

---
 .../database/unitycatalog/connection.py       | 20 +++++++++++--------
 1 file changed, 12 insertions(+), 8 deletions(-)

diff --git a/ingestion/src/metadata/ingestion/source/database/unitycatalog/connection.py b/ingestion/src/metadata/ingestion/source/database/unitycatalog/connection.py
index 1035a0fc398..2ba24d3ac32 100644
--- a/ingestion/src/metadata/ingestion/source/database/unitycatalog/connection.py
+++ b/ingestion/src/metadata/ingestion/source/database/unitycatalog/connection.py
@@ -68,16 +68,20 @@ def test_connection(
             break
 
     def get_schemas(connection: WorkspaceClient, table_obj: DatabricksTable):
-        for schema in connection.schemas.list(catalog_name=table_obj.catalog_name):
-            table_obj.schema_name = schema.name
-            break
+        for catalog in connection.catalogs.list():
+            for schema in connection.schemas.list(catalog_name=catalog.name):
+                if schema.name:
+                    table_obj.schema_name = schema.name
+                    table_obj.catalog_name = catalog.name
+                    return
 
     def get_tables(connection: WorkspaceClient, table_obj: DatabricksTable):
-        for table in connection.tables.list(
-            catalog_name=table_obj.catalog_name, schema_name=table_obj.schema_name
-        ):
-            table_obj.name = table.name
-            break
+        if table_obj.catalog_name and table_obj.schema_name:
+            for table in connection.tables.list(
+                catalog_name=table_obj.catalog_name, schema_name=table_obj.schema_name
+            ):
+                table_obj.name = table.name
+                break
 
     test_fn = {
         "CheckAccess": connection.catalogs.list,