Fix pg schema (#266)

Fix : schema now works for Postgres
infiniticio · Oct 14, 2024 · 1112147 · 1112147
1 parent 8b63e2f
commit 1112147
Show file tree

Hide file tree

Showing 11 changed files with 97 additions and 40 deletions.
diff --git a/.github/workflows/engine-ci.yml b/.github/workflows/engine-ci.yml
@@ -61,4 +61,4 @@ jobs:
 
       - name: Test with Gradle
         run: ./gradlew test --no-daemon --parallel --max-workers=8  # <= -d used to debug if needed
-        timeout-minutes: 12 # max time allocated (useful if some tests hang)
+        timeout-minutes: 10 # max time allocated (useful if some tests hang)
diff --git a/infinitic-storage/src/main/kotlin/io/infinitic/storage/config/PostgresConfig.kt b/infinitic-storage/src/main/kotlin/io/infinitic/storage/config/PostgresConfig.kt
@@ -64,6 +64,7 @@ interface PostgresConfigInterface {
   val username: String
   val password: String?
   val database: String
+  val schema: String
   val keySetTable: String
   val keyValueTable: String
   val maximumPoolSize: Int?
@@ -80,6 +81,7 @@ data class PostgresConfig(
   override val username: String,
   override val password: String? = null,
   override val database: String = DEFAULT_DATABASE,
+  override val schema: String = DEFAULT_SCHEMA,
   override val keySetTable: String = DEFAULT_KEY_SET_TABLE,
   override val keyValueTable: String = DEFAULT_KEY_VALUE_TABLE,
   override val maximumPoolSize: Int? = null,
@@ -89,8 +91,8 @@ data class PostgresConfig(
   override val maxLifetime: Long? = null // milli seconds
 ) : PostgresConfigInterface {
 
-  private val jdbcUrl = "jdbc:postgresql://$host:$port/$database"
-  private val jdbcUrlDefault = "jdbc:postgresql://$host:$port/postgres"
+  private val jdbcUrlBase = "jdbc:postgresql://$host:$port"
+  private val jdbcUrl = "$jdbcUrlBase/$database"
   private val driverClassName = "org.postgresql.Driver"
 
   init {
@@ -111,13 +113,16 @@ data class PostgresConfig(
     }
 
     require(database.isValidDatabaseName()) {
-      "Invalid value for '${::database.name}': '$database' is not a valid MySQL database name"
+      "Invalid value for '${::database.name}': '$database' is not a valid Postgres database name"
+    }
+    require(database.isValidSchemaName()) {
+      "Invalid value for '${::schema.name}': '$schema' is not a valid Postgres database name"
     }
     require(keySetTable.isValidTableName()) {
-      "Invalid value for '${::keySetTable.name}': '$keySetTable' is not a valid MySQL table name"
+      "Invalid value for '${::keySetTable.name}': '$keySetTable' is not a valid Postgres table name"
     }
     require(keyValueTable.isValidTableName()) {
-      "Invalid value for '${::keyValueTable.name}': '$keyValueTable' is not a valid MySQL table name"
+      "Invalid value for '${::keyValueTable.name}': '$keyValueTable' is not a valid Postgres table name"
     }
   }
 
@@ -127,7 +132,7 @@ data class PostgresConfig(
    */
   override fun toString() =
       "${this::class.java.simpleName}(host='$host', port=$port, username='$username', password='******', " +
-          "database=$database, keySetTable=$keySetTable, keyValueTable=$keyValueTable" +
+          "database=$database, schema=$schema, keySetTable=$keySetTable, keyValueTable=$keyValueTable" +
           (maximumPoolSize?.let { ", maximumPoolSize=$it" } ?: "") +
           (minimumIdle?.let { ", minimumIdle=$it" } ?: "") +
           (idleTimeout?.let { ", idleTimeout=$it" } ?: "") +
@@ -140,7 +145,8 @@ data class PostgresConfig(
 
     internal const val DEFAULT_KEY_VALUE_TABLE = "key_value_storage"
     internal const val DEFAULT_KEY_SET_TABLE = "key_set_storage"
-    internal const val DEFAULT_DATABASE = "infinitic"
+    internal const val DEFAULT_DATABASE = "postgres"
+    internal const val DEFAULT_SCHEMA = "infinitic"
   }
 
   fun close() {
@@ -151,6 +157,8 @@ data class PostgresConfig(
   fun getPool(): HikariDataSource = pools.getOrPut(this) {
     // Create the Database if needed
     initDatabase()
+    // Create the Schema if needed
+    initSchema()
     // create pool
     HikariDataSource(hikariConfig)
   }
@@ -162,6 +170,7 @@ data class PostgresConfig(
       driverClassName = config.driverClassName
       username = config.username
       password = config.password
+      schema = config.schema
       config.maximumPoolSize?.let { maximumPoolSize = it }
       config.minimumIdle?.let { minimumIdle = it }
       config.idleTimeout?.let { idleTimeout = it }
@@ -185,7 +194,7 @@ data class PostgresConfig(
       }
 
   private fun initDatabase() {
-    getDefaultPool().use { pool ->
+    getDefaultPool(DEFAULT_DATABASE).use { pool ->
       if (!pool.databaseExists(database)) {
         pool.connection.use { connection ->
           connection.createStatement().use {
@@ -196,10 +205,31 @@ data class PostgresConfig(
     }
   }
 
-  private fun getDefaultPool() = HikariDataSource(
+  private fun HikariDataSource.schemaExists(schemaName: String): Boolean =
+      connection.use { conn ->
+        conn.metaData.schemas.use { resultSet ->
+          generateSequence {
+            if (resultSet.next()) resultSet.getString("TABLE_SCHEM") else null
+          }.any { it.lowercase() == schemaName.lowercase() }
+        }
+      }
+
+  private fun initSchema() {
+    getDefaultPool(database).use { pool ->
+      if (!pool.schemaExists(schema)) {
+        pool.connection.use { connection ->
+          connection.createStatement().use { statement ->
+            statement.executeUpdate("CREATE SCHEMA $schema")
+          }
+        }
+      }
+    }
+  }
+
+  private fun getDefaultPool(database: String) = HikariDataSource(
       HikariConfig().apply {
         // use a default source
-        jdbcUrl = this@PostgresConfig.jdbcUrlDefault
+        jdbcUrl = this@PostgresConfig.jdbcUrlBase + "/$database"
         driverClassName = this@PostgresConfig.driverClassName
         username = this@PostgresConfig.username
         password = this@PostgresConfig.password
@@ -211,6 +241,11 @@ data class PostgresConfig(
     return isNotEmpty() && matches(regex)
   }
 
+  private fun String.isValidSchemaName(): Boolean {
+    val regex = "^[a-zA-Z_][a-zA-Z0-9_]{0,62}$".toRegex()
+    return isNotEmpty() && matches(regex)
+  }
+
   private fun String.isValidTableName(): Boolean {
     // Check length
     // Note that since Postgres uses bytes and Kotlin uses UTF-16 characters,

diff --git a/infinitic-storage/src/main/kotlin/io/infinitic/storage/config/PostgresStorageConfig.kt b/infinitic-storage/src/main/kotlin/io/infinitic/storage/config/PostgresStorageConfig.kt
@@ -27,6 +27,7 @@ import io.infinitic.storage.compression.CompressionConfig
 import io.infinitic.storage.config.PostgresConfig.Companion.DEFAULT_DATABASE
 import io.infinitic.storage.config.PostgresConfig.Companion.DEFAULT_KEY_SET_TABLE
 import io.infinitic.storage.config.PostgresConfig.Companion.DEFAULT_KEY_VALUE_TABLE
+import io.infinitic.storage.config.PostgresConfig.Companion.DEFAULT_SCHEMA
 import io.infinitic.storage.databases.postgres.PostgresKeySetStorage
 import io.infinitic.storage.databases.postgres.PostgresKeyValueStorage
 import io.infinitic.storage.keySet.KeySetStorage
@@ -64,6 +65,7 @@ data class PostgresStorageConfig(
     private var username: String? = null
     private var password: String? = null
     private var database: String = DEFAULT_DATABASE
+    private var schema: String = DEFAULT_SCHEMA
     private var keySetTable: String = DEFAULT_KEY_SET_TABLE
     private var keyValueTable: String = DEFAULT_KEY_VALUE_TABLE
     private var maximumPoolSize: Int? = null
@@ -79,6 +81,7 @@ data class PostgresStorageConfig(
     fun setUsername(user: String) = apply { this.username = user }
     fun setPassword(password: String) = apply { this.password = password }
     fun setDatabase(database: String) = apply { this.database = database }
+    fun setSchema(schema: String) = apply { this.schema = schema }
     fun setKeySetTable(keySetTable: String) = apply { this.keySetTable = keySetTable }
     fun setKeyValueTable(keyValueTable: String) = apply { this.keyValueTable = keyValueTable }
     fun setMaximumPoolSize(maximumPoolSize: Int) = apply { this.maximumPoolSize = maximumPoolSize }
@@ -98,6 +101,7 @@ data class PostgresStorageConfig(
           username = username!!,
           password = password,
           database = database,
+          schema = schema,
           keySetTable = keySetTable,
           keyValueTable = keyValueTable,
           maximumPoolSize = maximumPoolSize,

diff --git a/...-storage/src/main/kotlin/io/infinitic/storage/databases/postgres/PostgresKeySetStorage.kt b/...-storage/src/main/kotlin/io/infinitic/storage/databases/postgres/PostgresKeySetStorage.kt
@@ -29,11 +29,16 @@ import org.jetbrains.annotations.TestOnly
 
 class PostgresKeySetStorage(
   internal val pool: HikariDataSource,
-  private val tableName: String
+  private val tableName: String,
+  private val schema: String
 ) : KeySetStorage {
 
   companion object {
-    fun from(config: PostgresConfig) = PostgresKeySetStorage(config.getPool(), config.keySetTable)
+    fun from(config: PostgresConfig) = PostgresKeySetStorage(
+        config.getPool(),
+        config.keySetTable,
+        config.schema,
+    )
   }
 
   init {
@@ -47,7 +52,9 @@ class PostgresKeySetStorage(
 
   override suspend fun get(key: String): Set<ByteArray> =
       pool.connection.use { connection ->
-        connection.prepareStatement("SELECT value FROM $tableName WHERE key = ?")
+        connection.prepareStatement(
+            "SELECT value FROM $schema.$tableName WHERE key = ?",
+        )
             .use { statement ->
               statement.setString(1, key)
               statement.executeQuery().use {
@@ -62,7 +69,7 @@ class PostgresKeySetStorage(
 
   override suspend fun add(key: String, value: ByteArray) {
     pool.connection.use { connection ->
-      connection.prepareStatement("INSERT INTO $tableName (key, value) VALUES (?, ?)").use {
+      connection.prepareStatement("INSERT INTO $schema.$tableName (key, value) VALUES (?, ?)").use {
         it.setString(1, key)
         it.setBytes(2, value)
         it.executeUpdate()
@@ -72,18 +79,19 @@ class PostgresKeySetStorage(
 
   override suspend fun remove(key: String, value: ByteArray) {
     pool.connection.use { connection ->
-      connection.prepareStatement("DELETE FROM $tableName WHERE key = ? AND value = ?").use {
-        it.setString(1, key)
-        it.setBytes(2, value)
-        it.executeUpdate()
-      }
+      connection.prepareStatement("DELETE FROM $schema.$tableName WHERE key = ? AND value = ?")
+          .use {
+            it.setString(1, key)
+            it.setBytes(2, value)
+            it.executeUpdate()
+          }
     }
   }
 
   @TestOnly
   override fun flush() {
     pool.connection.use { connection ->
-      connection.prepareStatement("TRUNCATE $tableName").use { it.executeUpdate() }
+      connection.prepareStatement("TRUNCATE $schema.$tableName").use { it.executeUpdate() }
     }
   }
 
@@ -92,19 +100,19 @@ class PostgresKeySetStorage(
     // And value is typically a workflowId
     pool.connection.use { connection ->
       connection.prepareStatement(
-          "CREATE TABLE IF NOT EXISTS $tableName (" +
+          "CREATE TABLE IF NOT EXISTS $schema.$tableName (" +
               "id SERIAL PRIMARY KEY," +
               "key VARCHAR(255) NOT NULL," +
               "value BYTEA NOT NULL" +
               ");",
       ).use { it.executeUpdate() }
 
       connection.prepareStatement(
-          "CREATE INDEX IF NOT EXISTS index_key ON $tableName (key);",
+          "CREATE INDEX IF NOT EXISTS index_key ON $schema.$tableName (key);",
       ).use { it.executeUpdate() }
 
       connection.prepareStatement(
-          "CREATE INDEX IF NOT EXISTS index_key_value ON $tableName (key, value);",
+          "CREATE INDEX IF NOT EXISTS index_key_value ON $schema.$tableName (key, value);",
       ).use { it.executeUpdate() }
     }
   }

diff --git a/...torage/src/main/kotlin/io/infinitic/storage/databases/postgres/PostgresKeyValueStorage.kt b/...torage/src/main/kotlin/io/infinitic/storage/databases/postgres/PostgresKeyValueStorage.kt
@@ -30,12 +30,16 @@ import kotlin.math.ceil
 
 class PostgresKeyValueStorage(
   internal val pool: HikariDataSource,
-  private val tableName: String
+  private val tableName: String,
+  private val schema: String
 ) : KeyValueStorage {
 
   companion object {
-    fun from(config: PostgresConfig) =
-        PostgresKeyValueStorage(config.getPool(), config.keyValueTable)
+    fun from(config: PostgresConfig) = PostgresKeyValueStorage(
+        config.getPool(),
+        config.keyValueTable,
+        config.schema,
+    )
   }
 
   init {
@@ -49,7 +53,7 @@ class PostgresKeyValueStorage(
 
   override suspend fun get(key: String): ByteArray? =
       pool.connection.use { connection ->
-        connection.prepareStatement("SELECT value FROM $tableName WHERE key=?")
+        connection.prepareStatement("SELECT value FROM $schema.$tableName WHERE key=?")
             .use {
               it.setString(1, key)
               it.executeQuery().use { resultSet ->
@@ -64,7 +68,7 @@ class PostgresKeyValueStorage(
     pool.connection.use { connection ->
       connection
           .prepareStatement(
-              "INSERT INTO $tableName (key, value, value_size_in_KiB) VALUES (?, ?, ?) " +
+              "INSERT INTO $schema.$tableName (key, value, value_size_in_KiB) VALUES (?, ?, ?) " +
                   "ON CONFLICT (key) DO UPDATE SET value = EXCLUDED.value",
           )
           .use {
@@ -78,7 +82,7 @@ class PostgresKeyValueStorage(
 
   override suspend fun del(key: String) {
     pool.connection.use { connection ->
-      connection.prepareStatement("DELETE FROM $tableName WHERE key=?").use {
+      connection.prepareStatement("DELETE FROM $schema.$tableName WHERE key=?").use {
         it.setString(1, key)
         it.executeUpdate()
       }
@@ -88,14 +92,14 @@ class PostgresKeyValueStorage(
   @TestOnly
   override fun flush() {
     pool.connection.use { connection ->
-      connection.prepareStatement("TRUNCATE $tableName").use { it.executeUpdate() }
+      connection.prepareStatement("TRUNCATE $schema.$tableName").use { it.executeUpdate() }
     }
   }
 
   private fun initKeyValueTable() {
     pool.connection.use { connection ->
       connection.prepareStatement(
-          "CREATE TABLE IF NOT EXISTS $tableName (" +
+          "CREATE TABLE IF NOT EXISTS $schema.$tableName (" +
               "id BIGSERIAL PRIMARY KEY," +
               "key VARCHAR(255) NOT NULL UNIQUE," +
               "value BYTEA NOT NULL," +
@@ -105,7 +109,7 @@ class PostgresKeyValueStorage(
       ).use { it.executeUpdate() }
 
       connection.prepareStatement(
-          "CREATE INDEX IF NOT EXISTS value_size_index ON $tableName(value_size_in_KiB);",
+          "CREATE INDEX IF NOT EXISTS value_size_index ON $schema.$tableName(value_size_in_KiB);",
       ).use { it.executeUpdate() }
     }
   }

diff --git a/infinitic-storage/src/test/java/io/infinitic/storage/config/PostgresConfigTest.java b/infinitic-storage/src/test/java/io/infinitic/storage/config/PostgresConfigTest.java
@@ -48,7 +48,8 @@ void testDefaultParameters() {
         assertEquals(5432, config.getPort());
         assertEquals("postgres", config.getUsername());
         assertEquals("password", config.getPassword());
-        assertEquals("infinitic", config.getDatabase());
+        assertEquals("postgres", config.getDatabase());
+        assertEquals("infinitic", config.getSchema());
         assertEquals("key_set_storage", config.getKeySetTable());
         assertEquals("key_value_storage", config.getKeyValueTable());
     }

diff --git a/infinitic-storage/src/test/kotlin/io/infinitic/storage/config/PostgresConfigTests.kt b/infinitic-storage/src/test/kotlin/io/infinitic/storage/config/PostgresConfigTests.kt
@@ -38,7 +38,8 @@ class PostgresConfigTests : StringSpec(
       }
 
       "Check PostgresConfig default values do not change to ensure backward compatibility" {
-        config.database shouldBe "infinitic"
+        config.database shouldBe "postgres"
+        config.schema shouldBe "infinitic"
         config.keySetTable shouldBe "key_set_storage"
         config.keyValueTable shouldBe "key_value_storage"
       }
@@ -101,7 +102,7 @@ class PostgresConfigTests : StringSpec(
 
       "toString() should obfuscate password" {
         config.toString() shouldBe "PostgresConfig(host='${config.host}', port=${config.port}, username='${config.username}', password='******', " +
-            "database=${config.database}, keySetTable=${config.keySetTable}, keyValueTable=${config.keyValueTable})"
+            "database=${config.database}, schema=${config.schema}, keySetTable=${config.keySetTable}, keyValueTable=${config.keyValueTable})"
       }
 
       "Can not load from yaml with no host" {
@@ -165,6 +166,7 @@ storage:
             yaml + """
     password: pass
     database: azerty
+    schema: infinitic
     keySetTable: keySet
     keyValueTable: keyVal
     maximumPoolSize: 1
@@ -182,6 +184,7 @@ storage:
                 "root",
                 "pass",
                 "azerty",
+                "infinitic",
                 "keySet",
                 "keyVal",
                 1,

diff --git a/...age/src/test/kotlin/io/infinitic/storage/databases/postgres/PostgresKeySetStorageTests.kt b/...age/src/test/kotlin/io/infinitic/storage/databases/postgres/PostgresKeySetStorageTests.kt
@@ -39,7 +39,7 @@ class PostgresKeySetStorageTests :
               startupAttempts = 1
               withUsername("test")
               withPassword("password")
-              withDatabaseName("infinitic")
+              withDatabaseName("storageTest")
             }
             .also { it.start() }
 
@@ -49,6 +49,7 @@ class PostgresKeySetStorageTests :
             username = postgresServer.username,
             password = postgresServer.password,
             database = postgresServer.databaseName,
+            schema = "infiniticTest",
         )
 
         val storage = PostgresKeySetStorage.from(config)