From df410982f9e3add765f0787021bb67373a2af539 Mon Sep 17 00:00:00 2001
From: Lucas Vieira <vieiralucas4@gmail.com>
Date: Wed, 29 Jan 2025 11:24:12 -0400
Subject: [PATCH] fix default schema fetching

---
 apps/api/src/datasources/structure.ts   |  6 +++---
 apps/api/src/python/query/sqlalchemy.ts | 15 +++++++++++----
 2 files changed, 14 insertions(+), 7 deletions(-)

diff --git a/apps/api/src/datasources/structure.ts b/apps/api/src/datasources/structure.ts
index 3613949d..1bc7ee01 100644
--- a/apps/api/src/datasources/structure.ts
+++ b/apps/api/src/datasources/structure.ts
@@ -209,7 +209,7 @@ async function v2ToV3(
       failedAt: 'failedAt' in v2 ? new Date(v2.failedAt) : null,
       error: 'error' in v2 ? v2.error : undefined,
       defaultSchema:
-        'structure' in v2 ? (v2.structure?.defaultSchema ?? null) : null,
+        'structure' in v2 ? v2.structure?.defaultSchema ?? null : null,
     },
   })
 
@@ -544,8 +544,8 @@ async function _refreshDataSourceStructure(
   const updateQueue = new PQueue({ concurrency: 1 })
   const tables: { schema: string; table: string }[] = []
   let defaultSchema = ''
-  const onTable: OnTable = (schema, tableName, table, defaultSchema) => {
-    defaultSchema = defaultSchema || schema
+  const onTable: OnTable = (schema, tableName, table, newDefaultSchema) => {
+    defaultSchema = newDefaultSchema || schema
     tables.push({ schema, table: tableName })
     updateQueue.add(async () => {
       let embedding: number[] | null = null
diff --git a/apps/api/src/python/query/sqlalchemy.ts b/apps/api/src/python/query/sqlalchemy.ts
index 9e774a09..70c90b13 100644
--- a/apps/api/src/python/query/sqlalchemy.ts
+++ b/apps/api/src/python/query/sqlalchemy.ts
@@ -412,7 +412,7 @@ def get_columns(conn, inspector, table_name, schema_name):
 
     return inspector.get_columns(table_name, schema=schema_name)
 
-def schema_from_tables(conn, inspector, tables):
+def schema_from_tables(conn, inspector, tables, default_schema):
     made_progress = False
     exceptions = []
     for table in tables:
@@ -438,7 +438,7 @@ def schema_from_tables(conn, inspector, tables):
             "table": {
                 "columns": columns
             },
-            "defaultSchema": "public"
+            "defaultSchema": default_schema
         }
         print(json.dumps(progress, default=str))
 
@@ -460,9 +460,16 @@ def get_data_source_structure(data_source_id, credentials_info=None):
                 oracledb.init_oracle_client()
 
             inspector = inspect(engine)
+
+            default_schema = "public"
+            try:
+                default_schema = inspector.default_schema_name or default_schema
+            except:
+                pass
+
             if ${JSON.stringify(ds.type)} == "bigquery":
                 tables = inspector.get_table_names()
-                schema_from_tables(conn, inspector, tables)
+                schema_from_tables(conn, inspector, tables, default_schema)
             else:
                 tables = []
                 exceptions = []
@@ -475,7 +482,7 @@ def get_data_source_structure(data_source_id, credentials_info=None):
                         exceptions.append(e)
                         print(json.dumps({"log": f"Failed to get tables for schema {schema_name}: {str(e)}"}))
                         continue
-                schema_from_tables(conn, inspector, tables)
+                schema_from_tables(conn, inspector, tables, default_schema)
                 if len(tables) == 0 and len(exceptions) > 0:
                     raise BrieferAggregateException(exceptions)
     finally: