SQL preflight connection validation (#4150)

* wip sql connection string validation

* handle failed sql connections in frontend

* sql preflight connection validation on modal save

* revert unneeded be/fe changes

* linting, form updates

---------

Co-authored-by: timothycarambat <rambat1010@gmail.com>

frontend/src/models/system.js

@@ -773,6 +773,25 @@ const System = {
     return newVersion;
   },
 
+  /**
+   * Validates a SQL connection string.
+   * @param {'postgresql'|'mysql'|'sql-server'} engine - the database engine identifier
+   * @param {string} connectionString - the connection string to validate
+   * @returns {Promise<{success: boolean, error: string | null}>}
+   */
+  validateSQLConnection: async function (engine, connectionString) {
+    return fetch(`${API_BASE}/system/validate-sql-connection`, {
+      method: "POST",
+      headers: baseHeaders(),
+      body: JSON.stringify({ engine, connectionString }),
+    })
+      .then((res) => res.json())
+      .catch((e) => {
+        console.error("Failed to validate SQL connection:", e);
+        return { success: false, error: e.message };
+      });
+  },
+
   experimentalFeatures: {
     liveSync: LiveDocumentSync,
     agentPlugins: AgentPlugins,

frontend/src/pages/Admin/Agents/SQLConnectorSelection/DBConnection.jsx

@@ -9,18 +9,16 @@ export const DB_LOGOS = {
   "sql-server": MSSQLLogo,
 };
 
-export default function DBConnection({ connection, onRemove, setHasChanges }) {
+export default function DBConnection({ connection, onRemove }) {
   const { database_id, engine } = connection;
   function removeConfirmation() {
     if (
       !window.confirm(
         `Delete ${database_id} from the list of available SQL connections? This cannot be undone.`
       )
-    ) {
+    )
       return false;
-    }
     onRemove(database_id);
-    setHasChanges(true);
   }
 
   return (

frontend/src/pages/Admin/Agents/SQLConnectorSelection/NewConnectionModal.jsx

@@ -3,6 +3,8 @@ import { createPortal } from "react-dom";
 import ModalWrapper from "@/components/ModalWrapper";
 import { WarningOctagon, X } from "@phosphor-icons/react";
 import { DB_LOGOS } from "./DBConnection";
+import System from "@/models/system";
+import showToast from "@/utils/toast";
 
 function assembleConnectionString({
   engine,
@@ -37,9 +39,15 @@ const DEFAULT_CONFIG = {
   encrypt: false,
 };
 
-export default function NewSQLConnection({ isOpen, closeModal, onSubmit }) {
+export default function NewSQLConnection({
+  isOpen,
+  closeModal,
+  onSubmit,
+  setHasChanges,
+}) {
   const [engine, setEngine] = useState(DEFAULT_ENGINE);
   const [config, setConfig] = useState(DEFAULT_CONFIG);
+  const [isValidating, setIsValidating] = useState(false);
   if (!isOpen) return null;
 
   function handleClose() {
@@ -48,8 +56,8 @@ export default function NewSQLConnection({ isOpen, closeModal, onSubmit }) {
     closeModal();
   }
 
-  function onFormChange() {
-    const form = new FormData(document.getElementById("sql-connection-form"));
+  function onFormChange(e) {
+    const form = new FormData(e.target.form);
     setConfig({
       username: form.get("username").trim(),
       password: form.get("password"),
@@ -64,12 +72,41 @@ export default function NewSQLConnection({ isOpen, closeModal, onSubmit }) {
     e.preventDefault();
     e.stopPropagation();
     const form = new FormData(e.target);
-    onSubmit({
-      engine,
-      database_id: form.get("name"),
-      connectionString: assembleConnectionString({ engine, ...config }),
-    });
-    handleClose();
+    const connectionString = assembleConnectionString({ engine, ...config });
+
+    setIsValidating(true);
+    try {
+      const { success, error } = await System.validateSQLConnection(
+        engine,
+        connectionString
+      );
+      if (!success) {
+        showToast(
+          error ||
+            "Failed to establish database connection. Please check your connection details.",
+          "error"
+        );
+        setIsValidating(false);
+        return;
+      }
+
+      onSubmit({
+        engine,
+        database_id: form.get("name"),
+        connectionString,
+      });
+      setHasChanges(true);
+      handleClose();
+    } catch (error) {
+      console.error("Error validating connection:", error);
+      showToast(
+        error?.message ||
+          "Failed to validate connection. Please check your connection details.",
+        "error"
+      );
+    } finally {
+      setIsValidating(false);
+    }
     return false;
   }
 
@@ -95,8 +132,8 @@ export default function NewSQLConnection({ isOpen, closeModal, onSubmit }) {
           </div>
           <form
             id="sql-connection-form"
-            onSubmit={handleUpdate}
             onChange={onFormChange}
+            onSubmit={handleUpdate}
           >
             <div className="px-7 py-6">
               <div className="space-y-6 max-h-[60vh] overflow-y-auto pr-2">
@@ -238,7 +275,6 @@ export default function NewSQLConnection({ isOpen, closeModal, onSubmit }) {
                         name="encrypt"
                         value="true"
                         className="sr-only peer"
-                        onChange={onFormChange}
                         checked={config.encrypt}
                       />
                       <div className="w-11 h-6 bg-theme-settings-input-bg peer-focus:outline-none peer-focus:ring-4 peer-focus:ring-blue-800 rounded-full peer peer-checked:after:translate-x-full peer-checked:after:border-white after:content-[''] after:absolute after:top-[2px] after:left-[2px] after:bg-white after:border-gray-300 after:border after:rounded-full after:h-5 after:w-5 after:transition-all peer-checked:bg-blue-600"></div>
@@ -265,9 +301,10 @@ export default function NewSQLConnection({ isOpen, closeModal, onSubmit }) {
               <button
                 type="submit"
                 form="sql-connection-form"
-                className="transition-all duration-300 bg-white text-black hover:opacity-60 px-4 py-2 rounded-lg text-sm"
+                disabled={isValidating}
+                className="transition-all duration-300 bg-white text-black hover:opacity-60 px-4 py-2 rounded-lg text-sm disabled:opacity-50"
               >
-                Save connection
+                {isValidating ? "Validating..." : "Save connection"}
               </button>
             </div>
           </form>

frontend/src/pages/Admin/Agents/SQLConnectorSelection/index.jsx

@@ -21,6 +21,17 @@ export default function AgentSQLConnectorSelection({
       .catch(() => setConnections([]));
   }, []);
 
+  function handleRemoveConnection(databaseId) {
+    setHasChanges(true);
+    setConnections((prev) =>
+      prev.map((conn) => {
+        if (conn.database_id === databaseId)
+          return { ...conn, action: "remove" };
+        return conn;
+      })
+    );
+  }
+
   return (
     <>
       <div className="p-2">
@@ -81,16 +92,7 @@ export default function AgentSQLConnectorSelection({
                       <DBConnection
                         key={connection.database_id}
                         connection={connection}
-                        onRemove={(databaseId) => {
-                          setHasChanges(true);
-                          setConnections((prev) =>
-                            prev.map((conn) => {
-                              if (conn.database_id === databaseId)
-                                return { ...conn, action: "remove" };
-                              return conn;
-                            })
-                          );
-                        }}
+                        onRemove={handleRemoveConnection}
                       />
                     ))}
                   <button
@@ -120,6 +122,7 @@ export default function AgentSQLConnectorSelection({
       <NewSQLConnection
         isOpen={isOpen}
         closeModal={closeModal}
+        setHasChanges={setHasChanges}
         onSubmit={(newDb) =>
           setConnections((prev) => [...prev, { action: "add", ...newDb }])
         }

server/endpoints/system.js

@@ -1383,6 +1383,42 @@ function systemEndpoints(app) {
       }
     }
   );
+
+  app.post(
+    "/system/validate-sql-connection",
+    [validatedRequest, flexUserRoleValid([ROLES.admin])],
+    async (request, response) => {
+      try {
+        const { engine, connectionString } = reqBody(request);
+        if (!engine || !connectionString) {
+          return response.status(400).json({
+            success: false,
+            error: "Both engine and connection details are required.",
+          });
+        }
+
+        const {
+          validateConnection,
+        } = require("../utils/agents/aibitat/plugins/sql-agent/SQLConnectors");
+        const result = await validateConnection(engine, { connectionString });
+
+        if (!result.success) {
+          return response.status(200).json({
+            success: false,
+            error: `Unable to connect to ${engine}. Please verify your connection details.`,
+          });
+        }
+
+        response.status(200).json(result);
+      } catch (error) {
+        console.error("SQL validation error:", error);
+        response.status(500).json({
+          success: false,
+          error: `Unable to connect to ${engine}. Please verify your connection details.`,
+        });
+      }
+    }
+  );
 }
 
 module.exports = { systemEndpoints };

server/utils/agents/aibitat/plugins/sql-agent/SQLConnectors/MSSQL.js

@@ -61,7 +61,7 @@ class MSSQLConnector {
   /**
    *
    * @param {string} queryString the SQL query to be run
-   * @returns {import(".").QueryResult}
+   * @returns {Promise<import(".").QueryResult>}
    */
   async runQuery(queryString = "") {
     const result = { rows: [], count: 0, error: null };
@@ -75,12 +75,24 @@ class MSSQLConnector {
       console.log(this.constructor.name, err);
       result.error = err.message;
     } finally {
-      await this._client.close();
-      this.#connected = false;
+      // Check client is connected before closing since we use this for validation
+      if (this._client) {
+        await this._client.close();
+        this.#connected = false;
+      }
     }
     return result;
   }
 
+  async validateConnection() {
+    try {
+      const result = await this.runQuery("SELECT 1");
+      return { success: !result.error, error: result.error };
+    } catch (error) {
+      return { success: false, error: error.message };
+    }
+  }
+
   getTablesSql() {
     return `SELECT name FROM sysobjects WHERE xtype='U';`;
   }

server/utils/agents/aibitat/plugins/sql-agent/SQLConnectors/MySQL.js

@@ -29,7 +29,7 @@ class MySQLConnector {
   /**
    *
    * @param {string} queryString the SQL query to be run
-   * @returns {import(".").QueryResult}
+   * @returns {Promise<import(".").QueryResult>}
    */
   async runQuery(queryString = "") {
     const result = { rows: [], count: 0, error: null };
@@ -42,12 +42,24 @@ class MySQLConnector {
       console.log(this.constructor.name, err);
       result.error = err.message;
     } finally {
-      await this._client.end();
-      this.#connected = false;
+      // Check client is connected before closing since we use this for validation
+      if (this._client) {
+        await this._client.end();
+        this.#connected = false;
+      }
     }
     return result;
   }
 
+  async validateConnection() {
+    try {
+      const result = await this.runQuery("SELECT 1");
+      return { success: !result.error, error: result.error };
+    } catch (error) {
+      return { success: false, error: error.message };
+    }
+  }
+
   getTablesSql() {
     return `SELECT table_name FROM information_schema.tables WHERE table_schema = '${this.database_id}'`;
   }

server/utils/agents/aibitat/plugins/sql-agent/SQLConnectors/Postgresql.js

@@ -22,7 +22,7 @@ class PostgresSQLConnector {
   /**
    *
    * @param {string} queryString the SQL query to be run
-   * @returns {import(".").QueryResult}
+   * @returns {Promise<import(".").QueryResult>}
    */
   async runQuery(queryString = "") {
     const result = { rows: [], count: 0, error: null };
@@ -35,12 +35,24 @@ class PostgresSQLConnector {
       console.log(this.constructor.name, err);
       result.error = err.message;
     } finally {
-      await this._client.end();
-      this.#connected = false;
+      // Check client is connected before closing since we use this for validation
+      if (this._client) {
+        await this._client.end();
+        this.#connected = false;
+      }
     }
     return result;
   }
 
+  async validateConnection() {
+    try {
+      const result = await this.runQuery("SELECT 1");
+      return { success: !result.error, error: result.error };
+    } catch (error) {
+      return { success: false, error: error.message };
+    }
+  }
+
   getTablesSql() {
     return `SELECT * FROM pg_catalog.pg_tables WHERE schemaname = 'public'`;
   }

server/utils/agents/aibitat/plugins/sql-agent/SQLConnectors/index.js

@@ -54,7 +54,27 @@ async function listSQLConnections() {
   );
 }
 
+/**
+ * Validates a SQL connection by attempting to connect and run a simple query
+ * @param {SQLEngine} identifier - The SQL engine type
+ * @param {object} connectionConfig - The connection configuration
+ * @returns {Promise<{success: boolean, error: string|null}>}
+ */
+async function validateConnection(identifier = "", connectionConfig = {}) {
+  try {
+    const client = getDBClient(identifier, connectionConfig);
+    return await client.validateConnection();
+  } catch (error) {
+    console.log(`Failed to connect to ${identifier} database.`);
+    return {
+      success: false,
+      error: `Unable to connect to ${identifier}. Please verify your connection details.`,
+    };
+  }
+}
+
 module.exports = {
   getDBClient,
   listSQLConnections,
+  validateConnection,
 };