From ac94432bab8c9abaf3af0026eb9b0709148b3881 Mon Sep 17 00:00:00 2001
From: Roman Wu <me@romanwu.com>
Date: Fri, 25 Jul 2025 19:31:56 -0400
Subject: [PATCH 1/6] Add PostgreSQL schema support to SQL connector

- Add schema configuration option to PostgreSQL connector
- Update SQL queries to filter by specified schema
- Add schema input field to frontend connection modal (PostgreSQL only)
- Default to 'public' schema when no custom schema specified
- Add getQualifiedTableName() method for schema.table format
---
 .../NewConnectionModal.jsx                     | 18 ++++++++++++++++++
 .../sql-agent/SQLConnectors/Postgresql.js      | 15 +++++++++++++--
 2 files changed, 31 insertions(+), 2 deletions(-)

diff --git a/frontend/src/pages/Admin/Agents/SQLConnectorSelection/NewConnectionModal.jsx b/frontend/src/pages/Admin/Agents/SQLConnectorSelection/NewConnectionModal.jsx
index 9c5327b99b7..0803ecd138e 100644
--- a/frontend/src/pages/Admin/Agents/SQLConnectorSelection/NewConnectionModal.jsx
+++ b/frontend/src/pages/Admin/Agents/SQLConnectorSelection/NewConnectionModal.jsx
@@ -36,6 +36,7 @@ const DEFAULT_CONFIG = {
   host: null,
   port: null,
   database: null,
+  schema: null,
   encrypt: false,
 };
 
@@ -269,6 +270,23 @@ export default function NewSQLConnection({
                   />
                 </div>
 
+                {engine === "postgresql" && (
+                  <div className="flex flex-col">
+                    <label className="block mb-2 text-sm font-medium text-white">
+                      Schema (optional)
+                    </label>
+                    <input
+                      type="text"
+                      name="schema"
+                      className="border-none bg-theme-settings-input-bg w-full text-white placeholder:text-theme-settings-input-placeholder text-sm rounded-lg focus:outline-primary-button active:outline-primary-button outline-none block w-full p-2.5"
+                      placeholder="public (default schema if not specified)"
+                      required={false}
+                      autoComplete="off"
+                      spellCheck={false}
+                    />
+                  </div>
+                )}
+
                 {engine === "sql-server" && (
                   <div className="flex items-center justify-between">
                     <label className="relative inline-flex items-center cursor-pointer">
diff --git a/server/utils/agents/aibitat/plugins/sql-agent/SQLConnectors/Postgresql.js b/server/utils/agents/aibitat/plugins/sql-agent/SQLConnectors/Postgresql.js
index 93eb3362f46..8c6ee54c2cc 100644
--- a/server/utils/agents/aibitat/plugins/sql-agent/SQLConnectors/Postgresql.js
+++ b/server/utils/agents/aibitat/plugins/sql-agent/SQLConnectors/Postgresql.js
@@ -5,9 +5,11 @@ class PostgresSQLConnector {
   constructor(
     config = {
       connectionString: null,
+      schema: null,
     }
   ) {
     this.connectionString = config.connectionString;
+    this.schema = config.schema || 'public';
     this._client = new pgSql.Client({
       connectionString: this.connectionString,
     });
@@ -54,10 +56,19 @@ class PostgresSQLConnector {
   }
 
   getTablesSql() {
-    return `SELECT * FROM pg_catalog.pg_tables WHERE schemaname = 'public'`;
+    return `SELECT * FROM pg_catalog.pg_tables WHERE schemaname = '${this.schema}'`;
   }
   getTableSchemaSql(table_name) {
-    return ` select column_name, data_type, character_maximum_length, column_default, is_nullable from INFORMATION_SCHEMA.COLUMNS where table_name = '${table_name}'`;
+    return ` select column_name, data_type, character_maximum_length, column_default, is_nullable from INFORMATION_SCHEMA.COLUMNS where table_name = '${table_name}' AND table_schema = '${this.schema}'`;
+  }
+
+  /**
+   * Returns the qualified table name with schema prefix
+   * @param {string} table_name - The table name
+   * @returns {string} - The qualified table name (schema.table)
+   */
+  getQualifiedTableName(table_name) {
+    return `${this.schema}.${table_name}`;
   }
 }
 

From 994d2eedae1c01bd23ad8ea82c9e8d7288e2a173 Mon Sep 17 00:00:00 2001
From: Roman Wu <me@romanwu.com>
Date: Fri, 25 Jul 2025 19:36:45 -0400
Subject: [PATCH 2/6] Fix code formatting after linting

---
 .../aibitat/plugins/sql-agent/SQLConnectors/Postgresql.js       | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/server/utils/agents/aibitat/plugins/sql-agent/SQLConnectors/Postgresql.js b/server/utils/agents/aibitat/plugins/sql-agent/SQLConnectors/Postgresql.js
index 8c6ee54c2cc..f995f01f883 100644
--- a/server/utils/agents/aibitat/plugins/sql-agent/SQLConnectors/Postgresql.js
+++ b/server/utils/agents/aibitat/plugins/sql-agent/SQLConnectors/Postgresql.js
@@ -9,7 +9,7 @@ class PostgresSQLConnector {
     }
   ) {
     this.connectionString = config.connectionString;
-    this.schema = config.schema || 'public';
+    this.schema = config.schema || "public";
     this._client = new pgSql.Client({
       connectionString: this.connectionString,
     });

From 0dc0dc08e2fd5f12c38e32580a19a2ffe6afca1b Mon Sep 17 00:00:00 2001
From: Roman Wu <me@romanwu.com>
Date: Mon, 28 Jul 2025 10:51:28 -0400
Subject: [PATCH 3/6] Update Postgresql.js

---
 .../plugins/sql-agent/SQLConnectors/Postgresql.js      | 10 ----------
 1 file changed, 10 deletions(-)

diff --git a/server/utils/agents/aibitat/plugins/sql-agent/SQLConnectors/Postgresql.js b/server/utils/agents/aibitat/plugins/sql-agent/SQLConnectors/Postgresql.js
index f995f01f883..a056132c57c 100644
--- a/server/utils/agents/aibitat/plugins/sql-agent/SQLConnectors/Postgresql.js
+++ b/server/utils/agents/aibitat/plugins/sql-agent/SQLConnectors/Postgresql.js
@@ -62,14 +62,4 @@ class PostgresSQLConnector {
     return ` select column_name, data_type, character_maximum_length, column_default, is_nullable from INFORMATION_SCHEMA.COLUMNS where table_name = '${table_name}' AND table_schema = '${this.schema}'`;
   }
 
-  /**
-   * Returns the qualified table name with schema prefix
-   * @param {string} table_name - The table name
-   * @returns {string} - The qualified table name (schema.table)
-   */
-  getQualifiedTableName(table_name) {
-    return `${this.schema}.${table_name}`;
-  }
-}
-
 module.exports.PostgresSQLConnector = PostgresSQLConnector;

From 6c957cf756635bd69b35e0f93482d015f15da7bd Mon Sep 17 00:00:00 2001
From: timothycarambat <rambat1010@gmail.com>
Date: Mon, 28 Jul 2025 08:28:34 -0700
Subject: [PATCH 4/6] add back removal of ending curly brace

---
 .../agents/aibitat/plugins/sql-agent/SQLConnectors/Postgresql.js | 1 +
 1 file changed, 1 insertion(+)

diff --git a/server/utils/agents/aibitat/plugins/sql-agent/SQLConnectors/Postgresql.js b/server/utils/agents/aibitat/plugins/sql-agent/SQLConnectors/Postgresql.js
index a056132c57c..15ca6dde428 100644
--- a/server/utils/agents/aibitat/plugins/sql-agent/SQLConnectors/Postgresql.js
+++ b/server/utils/agents/aibitat/plugins/sql-agent/SQLConnectors/Postgresql.js
@@ -61,5 +61,6 @@ class PostgresSQLConnector {
   getTableSchemaSql(table_name) {
     return ` select column_name, data_type, character_maximum_length, column_default, is_nullable from INFORMATION_SCHEMA.COLUMNS where table_name = '${table_name}' AND table_schema = '${this.schema}'`;
   }
+}
 
 module.exports.PostgresSQLConnector = PostgresSQLConnector;

From a2febbb81da3b23dda34765e8573799f7dfd8fbb Mon Sep 17 00:00:00 2001
From: Roman Wu <me@romanwu.com>
Date: Mon, 28 Jul 2025 11:32:06 -0400
Subject: [PATCH 5/6] Update Postgresql.js

---
 .../agents/aibitat/plugins/sql-agent/SQLConnectors/Postgresql.js | 1 -
 1 file changed, 1 deletion(-)

diff --git a/server/utils/agents/aibitat/plugins/sql-agent/SQLConnectors/Postgresql.js b/server/utils/agents/aibitat/plugins/sql-agent/SQLConnectors/Postgresql.js
index 15ca6dde428..a056132c57c 100644
--- a/server/utils/agents/aibitat/plugins/sql-agent/SQLConnectors/Postgresql.js
+++ b/server/utils/agents/aibitat/plugins/sql-agent/SQLConnectors/Postgresql.js
@@ -61,6 +61,5 @@ class PostgresSQLConnector {
   getTableSchemaSql(table_name) {
     return ` select column_name, data_type, character_maximum_length, column_default, is_nullable from INFORMATION_SCHEMA.COLUMNS where table_name = '${table_name}' AND table_schema = '${this.schema}'`;
   }
-}
 
 module.exports.PostgresSQLConnector = PostgresSQLConnector;

From f45f0d1742c418bb6f8f3a7a3ddd52658f25bc2e Mon Sep 17 00:00:00 2001
From: timothycarambat <rambat1010@gmail.com>
Date: Mon, 28 Jul 2025 08:33:41 -0700
Subject: [PATCH 6/6] add back removal of ending curly brace (again?)

---
 .../agents/aibitat/plugins/sql-agent/SQLConnectors/Postgresql.js | 1 +
 1 file changed, 1 insertion(+)

diff --git a/server/utils/agents/aibitat/plugins/sql-agent/SQLConnectors/Postgresql.js b/server/utils/agents/aibitat/plugins/sql-agent/SQLConnectors/Postgresql.js
index a056132c57c..15ca6dde428 100644
--- a/server/utils/agents/aibitat/plugins/sql-agent/SQLConnectors/Postgresql.js
+++ b/server/utils/agents/aibitat/plugins/sql-agent/SQLConnectors/Postgresql.js
@@ -61,5 +61,6 @@ class PostgresSQLConnector {
   getTableSchemaSql(table_name) {
     return ` select column_name, data_type, character_maximum_length, column_default, is_nullable from INFORMATION_SCHEMA.COLUMNS where table_name = '${table_name}' AND table_schema = '${this.schema}'`;
   }
+}
 
 module.exports.PostgresSQLConnector = PostgresSQLConnector;