Support External Postgres DB

backend/src/main/java/com/cloudera/cai/rag/configuration/JdbiConfiguration.java

@@ -66,6 +66,7 @@ private static Jdbi createJdbi() {
     if (jdbi == null) {
       synchronized (LOCK) {
         if (jdbi == null) {
+          log.info("Initializing new Jdbi instance");
           jdbi = Jdbi.create(createDataSource());
         }
       }
@@ -92,10 +93,19 @@ private static Migrator migrator(DataSource dataSource, RdbConfig dbConfig) {
   private static DatabaseConfig createDatabaseConfig() {
     String dbUrl = System.getenv().getOrDefault("DB_URL", "jdbc:h2:mem:rag");
     String rdbType = System.getenv().getOrDefault("DB_TYPE", RdbConfig.H2_DB_TYPE);
+    String password = System.getenv().get("DB_PASSWORD");
+    String username = System.getenv().get("DB_USERNAME");
     RdbConfig rdbConfiguration =
-        RdbConfig.builder().rdbUrl(dbUrl).rdbType(rdbType).rdbDatabaseName("rag").build();
+        RdbConfig.builder()
+            .rdbUrl(dbUrl)
+            .rdbType(rdbType)
+            .rdbDatabaseName("rag")
+            .rdbUsername(username)
+            .rdbPassword(password)
+            .build();
     if (rdbConfiguration.isPostgres()) {
-      rdbConfiguration = rdbConfiguration.toBuilder().rdbUsername("postgres").build();
+      rdbConfiguration =
+          rdbConfiguration.toBuilder().rdbUsername("postgres").rdbDatabaseName(null).build();
     }
     return DatabaseConfig.builder().RdbConfiguration(rdbConfiguration).build();
   }

backend/src/main/java/com/cloudera/cai/util/db/JdbiUtils.java

@@ -1,4 +1,4 @@
-/*******************************************************************************
+/*
  * CLOUDERA APPLIED MACHINE LEARNING PROTOTYPE (AMP)
  * (C) Cloudera, Inc. 2024
  * All rights reserved.
@@ -97,4 +97,46 @@ public static void createDBIfNotExists(RdbConfig rdbConfig) throws SQLException
       }
     }
   }
+
+  /**
+   * A utility class to test database connectivity using JDBC.
+   *
+   * <p>Run this with: java -cp prebuilt_artifacts/rag-api.jar
+   * -Dloader.main=com.cloudera.cai.util.db.JdbiUtils
+   * org.springframework.boot.loader.launch.PropertiesLauncher <jdbc_url> <username> <password>
+   *
+   * <p>An exit code of 0 indicates success, 1 indicates failure, and 2 indicates incorrect usage.
+   */
+  public static void main(String[] args) {
+    if (args.length != 4) {
+      System.err.println("Usage: JdbiUtils <db_url> <username> <password> <db_type>");
+      System.exit(2); // Incorrect usage
+    }
+    String dbUrl = args[0];
+    String username = args[1];
+    String password = args[2];
+    String dbType = args[3];
+    RdbConfig rdbConfiguration =
+        RdbConfig.builder()
+            .rdbUrl(dbUrl)
+            .rdbType(dbType)
+            .rdbDatabaseName("rag")
+            .rdbUsername(username)
+            .rdbPassword(password)
+            .build();
+    var connectionString = RdbConfig.buildDatabaseServerConnectionString(rdbConfiguration);
+    try (Connection connection =
+        DriverManager.getConnection(connectionString, username, password)) {
+      if (connection != null && !connection.isClosed()) {
+        System.out.println("Connection successful.");
+        System.exit(0); // Success
+      } else {
+        System.err.println("Connection failed: Connection is null or closed.");
+        System.exit(1); // Failure
+      }
+    } catch (Exception e) {
+      System.err.println("Connection failed: " + e.getMessage());
+      System.exit(1); // Failure
+    }
+  }
 }

backend/src/main/java/com/cloudera/cai/util/db/RdbConfig.java

@@ -123,7 +123,7 @@ public static String buildDatabaseConnectionString(RdbConfig rdb) {
       return adjustMsSqlRdbUrl(rdb.rdbUrl) + ";databaseName=" + rdb.getRdbDatabaseName();
     }
     if (rdb.isPostgres()) {
-      return rdb.rdbUrl + "/" + rdb.getRdbDatabaseName();
+      return rdb.rdbUrl;
     }
 
     final var url =
@@ -153,7 +153,15 @@ public static String buildDatabaseServerConnectionString(RdbConfig rdb) {
     }
 
     if (rdb.isPostgres()) {
-      return rdb.rdbUrl + "/" + rdb.getRdbDatabaseName();
+      var pattern =
+          Pattern.compile("^jdbc:postgresql:(//[^/]+/)?(\\w+)(.*)", Pattern.CASE_INSENSITIVE);
+      var matcher = pattern.matcher(rdb.rdbUrl);
+      if (!matcher.matches()) {
+        throw new IllegalStateException("URL doesn't match the expected regex");
+      }
+      var firstPart = matcher.group(1);
+      var lastPart = matcher.group(3);
+      return "jdbc:postgresql:" + firstPart + rdb.getRdbDatabaseName() + lastPart;
     }
     final var url =
         rdb.rdbUrl
@@ -181,7 +189,11 @@ private static String adjustMsSqlRdbUrl(String rdbUrl) {
   public static String buildDatabaseName(RdbConfig rdb) {
     String dbName = rdb.getRdbDatabaseName();
     if (rdb.getDbConnectionUrl() != null) {
-      dbName = getDBNameFromDBConnectionURL(rdb);
+      dbName = getDBNameFromDBConnectionURL(rdb, rdb.dbConnectionUrl);
+    }
+
+    if (dbName == null) {
+      dbName = getDBNameFromDBConnectionURL(rdb, rdb.rdbUrl);
     }
 
     if (dbName.contains("-")) {
@@ -195,25 +207,26 @@ public static String buildDatabaseName(RdbConfig rdb) {
     return dbName;
   }
 
-  private static String getDBNameFromDBConnectionURL(RdbConfig rdb) {
+  private static String getDBNameFromDBConnectionURL(RdbConfig rdb, String url) {
     String regex;
     if (rdb.isMssql()) {
       // Regex reference: https://regex101.com/r/yaU0DY/1
       regex = ";databaseName=([^;]*)";
-    } else {
+    } else if (rdb.isMysql()) {
       regex = "^jdbc:mysql:(?://[^/]+/)?(\\w+)";
+    } else if (rdb.isPostgres()) {
+      regex = "^jdbc:postgresql:(?://[^/]+/)?(\\w+)";
+    } else {
+      throw new IllegalStateException(
+          "database url parsing not supported for db type: " + rdb.rdbType);
     }
     Pattern pattern = Pattern.compile(regex, Pattern.CASE_INSENSITIVE);
     var dbName =
-        pattern
-            .matcher(rdb.dbConnectionUrl)
-            .results()
-            .map(mr -> mr.group(1))
-            .collect(Collectors.joining());
+        pattern.matcher(url).results().map(mr -> mr.group(1)).collect(Collectors.joining());
 
     if (dbName.isEmpty()) {
       throw new InvalidDbConfigException(
-          rdb.dbConnectionUrl, "Database name not found in the database connection URL");
+          url, "Database name not found in the database connection URL");
     }
     return dbName;
   }

backend/src/main/resources/migrations/postgres/25_alter_datasource_config.up.sql

@@ -38,6 +38,6 @@
 
 BEGIN;
 
-ALTER TABLE rag_data_source ALTER COLUMN chunk_overlap_percent INTEGER DEFAULT 10;
+ALTER TABLE rag_data_source ALTER COLUMN chunk_overlap_percent SET DEFAULT 10;
 
 COMMIT;

backend/src/test/java/com/cloudera/cai/util/db/RdbConfigTest.java

@@ -0,0 +1,60 @@
+/*
+ * CLOUDERA APPLIED MACHINE LEARNING PROTOTYPE (AMP)
+ * (C) Cloudera, Inc. 2025
+ * All rights reserved.
+ *
+ * Applicable Open Source License: Apache 2.0
+ *
+ * NOTE: Cloudera open source products are modular software products
+ * made up of hundreds of individual components, each of which was
+ * individually copyrighted.  Each Cloudera open source product is a
+ * collective work under U.S. Copyright Law. Your license to use the
+ * collective work is as provided in your written agreement with
+ * Cloudera.  Used apart from the collective work, this file is
+ * licensed for your use pursuant to the open source license
+ * identified above.
+ *
+ * This code is provided to you pursuant a written agreement with
+ * (i) Cloudera, Inc. or (ii) a third-party authorized to distribute
+ * this code. If you do not have a written agreement with Cloudera nor
+ * with an authorized and properly licensed third party, you do not
+ * have any rights to access nor to use this code.
+ *
+ * Absent a written agreement with Cloudera, Inc. ("Cloudera") to the
+ * contrary, A) CLOUDERA PROVIDES THIS CODE TO YOU WITHOUT WARRANTIES OF ANY
+ * KIND; (B) CLOUDERA DISCLAIMS ANY AND ALL EXPRESS AND IMPLIED
+ * WARRANTIES WITH RESPECT TO THIS CODE, INCLUDING BUT NOT LIMITED TO
+ * IMPLIED WARRANTIES OF TITLE, NON-INFRINGEMENT, MERCHANTABILITY AND
+ * FITNESS FOR A PARTICULAR PURPOSE; (C) CLOUDERA IS NOT LIABLE TO YOU,
+ * AND WILL NOT DEFEND, INDEMNIFY, NOR HOLD YOU HARMLESS FOR ANY CLAIMS
+ * ARISING FROM OR RELATED TO THE CODE; AND (D)WITH RESPECT TO YOUR EXERCISE
+ * OF ANY RIGHTS GRANTED TO YOU FOR THE CODE, CLOUDERA IS NOT LIABLE FOR ANY
+ * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, PUNITIVE OR
+ * CONSEQUENTIAL DAMAGES INCLUDING, BUT NOT LIMITED TO, DAMAGES
+ * RELATED TO LOST REVENUE, LOST PROFITS, LOSS OF INCOME, LOSS OF
+ * BUSINESS ADVANTAGE OR UNAVAILABILITY, OR LOSS OR CORRUPTION OF
+ * DATA.
+ */
+
+package com.cloudera.cai.util.db;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.junit.jupiter.api.Assertions.*;
+
+import org.junit.jupiter.api.Test;
+class RdbConfigTest {
+
+  @Test
+  void buildDatabaseConnectionString() {
+
+    var url =
+        "jdbc:postgresql://rag-dev-testing.cluster.us-west-2.rds.amazonaws.com:5432/rag?username=foo&password=bar";
+
+    var rdb =
+        RdbConfig.builder().rdbUrl(url).rdbDatabaseName("postgres").rdbType("PostgreSQL").build();
+    var result = RdbConfig.buildDatabaseServerConnectionString(rdb);
+    assertThat(result)
+        .isEqualTo(
+            "jdbc:postgresql://rag-dev-testing.cluster.us-west-2.rds.amazonaws.com:5432/postgres?username=foo&password=bar");
+  }
+}

docker-compose.yaml

@@ -10,7 +10,7 @@ services:
       - "9464:9464"
     environment:
       - API_HOST=0.0.0.0
-      - DB_URL=jdbc:postgresql://db:5432
+      - DB_URL=jdbc:postgresql://db:5432/rag
       - DB_TYPE=PostgreSQL
       - OTEL_EXPORTER_OTLP_ENDPOINT=http://tempo:4318
       - OTEL_METRICS_EXPORTER=none # we configure this by hand

llm-service/app/config.py

@@ -52,6 +52,7 @@
 SummaryStorageProviderType = Literal["Local", "S3"]
 ChatStoreProviderType = Literal["Local", "S3"]
 VectorDbProviderType = Literal["QDRANT", "OPENSEARCH"]
+MetadataDbProviderType = Literal["H2", "PostgreSQL"]
 
 
 class _Settings:

llm-service/app/routers/index/__init__.py

@@ -46,16 +46,15 @@
 from . import amp_metadata
 from . import models
 from . import metrics
-from . import chat
 
 logger = logging.getLogger(__name__)
 
 
 router = APIRouter()
-router.include_router(chat.router)
 router.include_router(summaries.router)
 router.include_router(data_source.router)
 router.include_router(sessions.router)
+router.include_router(sessions.no_id_router)
 router.include_router(amp_metadata.router)
 # include this for legacy UI calls
 router.include_router(amp_metadata.router, prefix="/index", deprecated=True)

llm-service/app/routers/index/amp_metadata/__init__.py

@@ -46,6 +46,7 @@
 from fastapi.params import Header
 
 from .... import exceptions
+from ....config import MetadataDbProviderType
 from ....services.amp_metadata import (
     ProjectConfig,
     ProjectConfigPlus,
@@ -54,6 +55,8 @@
     update_project_environment,
     get_project_environment,
     get_application_config,
+    validate_jdbc,
+    ValidationResult,
 )
 from ....services.amp_update import does_amp_need_updating
 from ....services.models.providers import CAIIModelProvider
@@ -198,6 +201,24 @@ def save_auth_token(auth_token: Annotated[str, Body(embed=True)]) -> str:
     return "Auth token saved successfully"
 
 
+@router.post(
+    "/validate-jdbc-connection",
+    summary="Validates a JDBC connection string, username, and password.",
+)
+@exceptions.propagates
+def validate_jdbc_connection(
+    db_url: Annotated[str, Body(embed=True)],
+    username: Annotated[str, Body(embed=True)],
+    password: Annotated[str, Body(embed=True)],
+    db_type: Annotated[MetadataDbProviderType, Body(embed=True)],
+) -> ValidationResult:
+    """
+    Calls the JdbiUtils main method to validate JDBC connection parameters.
+    Returns a dict with 'valid': True/False and 'message'.
+    """
+    return validate_jdbc(db_type, db_url, password, username)
+
+
 def save_cdp_token(auth_token: str) -> None:
     token_data = {"access_token": auth_token}
     with open("cdp_token", "w") as file:

llm-service/app/routers/index/models/__init__.py

@@ -64,7 +64,9 @@ def get_reranking_models() -> List[ModelResponse]:
     return models.Reranking.list_available()
 
 
-@router.get("/model_source", summary="Model source enabled - Bedrock, CAII, or Azure")
+@router.get(
+    "/model_source", summary="Model source enabled - Bedrock, CAII, OpenAI or Azure"
+)
 @exceptions.propagates
 def get_model() -> models.ModelSource:
     return models.get_model_source()