From ea1be666743f0bf987495523f8d5a83e3ce3a0d3 Mon Sep 17 00:00:00 2001
From: Francisco Diaz <francisco.d@47deg.com>
Date: Mon, 13 Nov 2023 15:29:00 +0100
Subject: [PATCH] Replace host, port, and database properties by url in
 VectorStore configuration (#535)

* Replace host, port, and database properties by url

* Add invoke method
---
 .../store/migrations/PsqlVectorStoreConfig.kt |  4 +--
 .../server/db/psql/PSQLVectorStoreConfig.kt   | 36 +++++++++----------
 .../services/PostgresVectorStoreService.kt    |  8 +----
 3 files changed, 21 insertions(+), 27 deletions(-)

diff --git a/integrations/postgresql/src/main/kotlin/com/xebia/functional/xef/store/migrations/PsqlVectorStoreConfig.kt b/integrations/postgresql/src/main/kotlin/com/xebia/functional/xef/store/migrations/PsqlVectorStoreConfig.kt
index 984fed951..0f93526f9 100644
--- a/integrations/postgresql/src/main/kotlin/com/xebia/functional/xef/store/migrations/PsqlVectorStoreConfig.kt
+++ b/integrations/postgresql/src/main/kotlin/com/xebia/functional/xef/store/migrations/PsqlVectorStoreConfig.kt
@@ -7,7 +7,7 @@ import org.flywaydb.core.api.configuration.FluentConfiguration
 import org.flywaydb.core.api.output.MigrateResult
 
 class PsqlVectorStoreConfig(
-  val uri: String,
+  val url: String,
   val driver: String,
   val user: String,
   val password: String,
@@ -18,7 +18,7 @@ class PsqlVectorStoreConfig(
     withContext(Dispatchers.IO) {
       val migration: FluentConfiguration = Flyway.configure()
         .dataSource(
-          uri,
+          url,
           user,
           password
         )
diff --git a/server/src/main/kotlin/com/xebia/functional/xef/server/db/psql/PSQLVectorStoreConfig.kt b/server/src/main/kotlin/com/xebia/functional/xef/server/db/psql/PSQLVectorStoreConfig.kt
index 6147d4a60..fd216bb5f 100644
--- a/server/src/main/kotlin/com/xebia/functional/xef/server/db/psql/PSQLVectorStoreConfig.kt
+++ b/server/src/main/kotlin/com/xebia/functional/xef/server/db/psql/PSQLVectorStoreConfig.kt
@@ -16,9 +16,7 @@ import org.slf4j.Logger
 
 @Serializable
 class PSQLVectorStoreConfig(
-  val host: String,
-  val port: Int,
-  val database: String,
+  val url: String,
   val driver: String,
   val user: String,
   val password: String,
@@ -26,29 +24,33 @@ class PSQLVectorStoreConfig(
   val vectorSize: Int
 ) : VectorStoreConfig {
 
-  fun getUrl(): String = "jdbc:postgresql://$host:$port/$database"
-
   override suspend fun getVectorStoreService(logger: Logger): PostgresVectorStoreService {
-    val vectorStoreHikariDataSource =
-      RepositoryService.getHikariDataSource(getUrl(), user, password)
+    val vectorStoreHikariDataSource = RepositoryService.getHikariDataSource(url, user, password)
     return PostgresVectorStoreService(toPGVectorStoreConfig(), logger, vectorStoreHikariDataSource)
   }
 
   private fun toPGVectorStoreConfig() =
     PostgreSQLXef.PGVectorStoreConfig(
-      dbConfig =
-        PostgreSQLXef.DBConfig(
-          host = host,
-          port = port,
-          database = database,
-          user = user,
-          password = password
-        ),
+      dbConfig = PostgreSQLXef.DBConfig(url = url, user = user, password = password),
       collectionName = collectionName,
       vectorSize = vectorSize
     )
 
   companion object {
+    operator fun invoke(
+      host: String,
+      port: Int,
+      database: String,
+      driver: String,
+      user: String,
+      password: String,
+      collectionName: String,
+      vectorSize: Int
+    ): PSQLVectorStoreConfig {
+      val url = "jdbc:postgresql://${host}:${port}/${database}"
+      return PSQLVectorStoreConfig(url, driver, user, password, collectionName, vectorSize)
+    }
+
     @OptIn(ExperimentalSerializationApi::class)
     suspend fun load(configNamespace: String, config: Config?): PSQLVectorStoreConfig =
       withContext(Dispatchers.IO) {
@@ -61,9 +63,7 @@ class PSQLVectorStoreConfig(
 
     private fun PSQLVectorStoreConfig.toPSQLConfig(): PsqlVectorStoreConfig =
       PsqlVectorStoreConfig(
-        host = this.host,
-        port = this.port,
-        database = this.database,
+        url = this.url,
         driver = this.driver,
         user = this.user,
         password = this.password,
diff --git a/server/src/main/kotlin/com/xebia/functional/xef/server/services/PostgresVectorStoreService.kt b/server/src/main/kotlin/com/xebia/functional/xef/server/services/PostgresVectorStoreService.kt
index 63a96fe53..b661af17f 100644
--- a/server/src/main/kotlin/com/xebia/functional/xef/server/services/PostgresVectorStoreService.kt
+++ b/server/src/main/kotlin/com/xebia/functional/xef/server/services/PostgresVectorStoreService.kt
@@ -14,13 +14,7 @@ import kotlinx.uuid.generateUUID
 import org.slf4j.Logger
 
 object PostgreSQLXef {
-  data class DBConfig(
-    val host: String,
-    val port: Int,
-    val database: String,
-    val user: String,
-    val password: String
-  )
+  data class DBConfig(val url: String, val user: String, val password: String)
 
   data class PGVectorStoreConfig(
     val dbConfig: DBConfig,