From 78cbb06c573d85411c07950fafc91093763f97c2 Mon Sep 17 00:00:00 2001
From: Roman Wu <me@romanwu.com>
Date: Mon, 28 Jul 2025 11:43:15 -0400
Subject: [PATCH] Modify the PostgreSQL SQL connector to support querying
 tables from schemas other than the default public schema. (#4202)

* 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

* Fix code formatting after linting

* Update Postgresql.js

* add back removal of ending curly brace

* Update Postgresql.js

* add back removal of ending curly brace (again?)

---------

Co-authored-by: timothycarambat <rambat1010@gmail.com>
---
 .../NewConnectionModal.jsx                     | 18 ++++++++++++++++++
 .../sql-agent/SQLConnectors/Postgresql.js      |  6 ++++--
 2 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/frontend/src/pages/Admin/Agents/SQLConnectorSelection/NewConnectionModal.jsx b/frontend/src/pages/Admin/Agents/SQLConnectorSelection/NewConnectionModal.jsx
index 9c5327b9..0803ecd1 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 93eb3362..15ca6dde 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,10 @@ 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}'`;
   }
 }