Auto migration using flyway

build.gradle

@@ -19,6 +19,12 @@ dependencies {
     // https://mvnrepository.com/artifact/org.postgresql/postgresql
     implementation group: 'org.postgresql', name: 'postgresql', version: '42.2.10'
 
+    // https://mvnrepository.com/artifact/org.flywaydb/flyway-core
+    implementation group: 'org.flywaydb', name: 'flyway-core', version: '7.15.0'
+
+    // https://mvnrepository.com/artifact/com.googlecode.libphonenumber/libphonenumber
+    compileOnly group: 'com.googlecode.libphonenumber', name: 'libphonenumber', version: '8.13.25'
+
     // https://mvnrepository.com/artifact/com.fasterxml.jackson.dataformat/jackson-dataformat-yaml
     compileOnly group: 'com.fasterxml.jackson.dataformat', name: 'jackson-dataformat-yaml', version: '2.14.0'
 

implementationDependencies.json

@@ -15,6 +15,11 @@
       "jar": "https://repo1.maven.org/maven2/org/slf4j/slf4j-api/1.7.25/slf4j-api-1.7.25.jar",
       "name": "SLF4j API 1.7.25",
       "src": "https://repo1.maven.org/maven2/org/slf4j/slf4j-api/1.7.25/slf4j-api-1.7.25-sources.jar"
+    },
+    {
+      "jar": "https://repo1.maven.org/maven2/org/flywaydb/flyway-core/7.15.0/flyway-core-7.15.0.jar",
+      "name": "Flyway Core 7.15.0",
+      "src": "https://repo1.maven.org/maven2/org/flywaydb/flyway-core/7.15.0/flyway-core-7.15.0-sources.jar"
     }
   ]
 }

src/main/java/io/supertokens/storage/postgresql/ConnectionPool.java

@@ -198,6 +198,16 @@ public static Connection getConnection(Start start) throws SQLException {
         return getInstance(start).hikariDataSource.getConnection();
     }
 
+    public static HikariDataSource getHikariDataSource(Start start) throws SQLException {
+        if (getInstance(start) == null || getInstance(start).hikariDataSource == null ) {
+            throw new IllegalStateException("Please call initPool before getHikariDataSource");
+        }
+        if (!start.enabled) {
+            throw new SQLException("Storage layer disabled");
+        }
+        return getInstance(start).hikariDataSource;
+    }
+
     static void close(Start start) {
         if (getInstance(start) == null) {
             return;

src/main/java/io/supertokens/storage/postgresql/FlywayMigration.java

@@ -0,0 +1,66 @@
+/*
+ *    Copyright (c) 2023, VRAI Labs and/or its affiliates. All rights reserved.
+ *
+ *    This software is licensed under the Apache License, Version 2.0 (the
+ *    "License") as published by the Apache Software Foundation.
+ *
+ *    You may not use this file except in compliance with the License. You may
+ *    obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *    Unless required by applicable law or agreed to in writing, software
+ *    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ *    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ *    License for the specific language governing permissions and limitations
+ *    under the License.
+ */
+
+package io.supertokens.storage.postgresql;
+
+import io.supertokens.pluginInterface.exceptions.StorageQueryException;
+import io.supertokens.storage.postgresql.config.Config;
+import io.supertokens.storage.postgresql.output.Logging;
+import io.supertokens.storage.postgresql.queries.BaselineMigrationQueries;
+import org.flywaydb.core.Flyway;
+import java.sql.SQLException;
+import java.util.HashMap;
+import java.util.Map;
+
+public final class FlywayMigration {
+
+    private static final String LOCATION = "classpath:/io/supertokens/storage/postgresql/migrations";
+    private FlywayMigration() {}
+
+    public static void startMigration(Start start) throws SQLException, StorageQueryException {
+        String baseline = BaselineMigrationQueries.getBaselineMigrationVersion(start);
+        if (Integer.parseInt(baseline) >= BaselineMigrationQueries.LAST_MIGRATION) {
+            return;
+        }
+
+        Logging.info(start, "Starting migration.", true);
+        MigrationContextManager.putContext(start.getProcessId(), start);
+        int maxRetries = 5;
+
+        try {
+            Flyway flyway = Flyway.configure()
+                    .dataSource(ConnectionPool.getHikariDataSource(start))
+                    .baselineOnMigrate(true)
+                    .baselineVersion(baseline)
+                    .table(Config.getConfig(start).getFlywaySchemaHistory())
+                    .connectRetries(maxRetries)
+                    .lockRetryCount(maxRetries)
+                    .locations(LOCATION)
+                    .placeholders(getPlaceholders(start))
+                    .load();
+            flyway.migrate();
+        } finally {
+            MigrationContextManager.removeContext(start.getProcessId());
+        }
+    }
+
+    private static Map<String, String> getPlaceholders(Start start) {
+        Map<String, String> ph = new HashMap<>();
+        ph.put("process_id", start.getProcessId());
+        ph.put("access_token_signing_key_dynamic", String.valueOf( Config.getConfig(start).getAccessTokenSigningKeyDynamic()));
+        return ph;
+    }
+}

src/main/java/io/supertokens/storage/postgresql/MigrationContextManager.java

@@ -0,0 +1,36 @@
+/*
+ *    Copyright (c) 2024, VRAI Labs and/or its affiliates. All rights reserved.
+ *
+ *    This software is licensed under the Apache License, Version 2.0 (the
+ *    "License") as published by the Apache Software Foundation.
+ *
+ *    You may not use this file except in compliance with the License. You may
+ *    obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *    Unless required by applicable law or agreed to in writing, software
+ *    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ *    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ *    License for the specific language governing permissions and limitations
+ *    under the License.
+ */
+
+package io.supertokens.storage.postgresql;
+
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+
+public class MigrationContextManager {
+    private static Map<String, Start> contextMap = new ConcurrentHashMap<>();
+
+    public static void putContext(String key, Start start) {
+        contextMap.put(key, start);
+    }
+
+    public static Start getContext(String key) {
+        return contextMap.get(key);
+    }
+
+    public static void removeContext(String key) {
+        contextMap.remove(key);
+    }
+}

src/main/java/io/supertokens/storage/postgresql/ProcessState.java

@@ -61,7 +61,7 @@ public synchronized void clear() {
      * DEADLOCK_FOUND: In case of a deadlock situation, we put this event
      */
     public enum PROCESS_STATE {
-        CREATING_NEW_TABLE, DEADLOCK_FOUND, DEADLOCK_NOT_RESOLVED
+        CREATING_NEW_TABLE, STARTING_MIGRATION, DEADLOCK_FOUND, DEADLOCK_NOT_RESOLVED
     }
 
     public static class EventAndException {

src/main/java/io/supertokens/storage/postgresql/QueryExecutorTemplate.java

@@ -58,5 +58,4 @@ static int update(Connection con, String QUERY, PreparedStatementValueSetter set
             return pst.executeUpdate();
         }
     }
-
 }

src/main/java/io/supertokens/storage/postgresql/Start.java

@@ -193,11 +193,16 @@ public void initFileLogging(String infoLogPath, String errorLogPath) {
              * nothing will be handling logging and hikari's logs would not be outputed
              * anywhere.
              */
-            final Logger infoLog = (Logger) LoggerFactory.getLogger("com.zaxxer.hikari");
+            final Logger hikariLog = (Logger) LoggerFactory.getLogger("com.zaxxer.hikari");
+            final Logger flywayLog = (Logger) LoggerFactory.getLogger("org.flywaydb.core");
             appender = new HikariLoggingAppender(this);
-            if (infoLog.getAppender(HikariLoggingAppender.NAME) == null) {
-                infoLog.setAdditive(false);
-                infoLog.addAppender(appender);
+            if (hikariLog.getAppender(HikariLoggingAppender.NAME) == null) {
+                hikariLog.setAdditive(false);
+                hikariLog.addAppender(appender);
+            }
+            if (flywayLog.getAppender(HikariLoggingAppender.NAME) == null) {
+                flywayLog.setAdditive(false);
+                flywayLog.addAppender(appender);
             }
         }
     }
@@ -208,9 +213,13 @@ public void stopLogging() {
             synchronized (appenderLock) {
                 Logging.stopLogging(this);
 
-                final Logger infoLog = (Logger) LoggerFactory.getLogger("com.zaxxer.hikari");
-                if (infoLog.getAppender(HikariLoggingAppender.NAME) != null) {
-                    infoLog.detachAppender(HikariLoggingAppender.NAME);
+                final Logger hikariLog = (Logger) LoggerFactory.getLogger("com.zaxxer.hikari");
+                final Logger flywayLog = (Logger) LoggerFactory.getLogger("org.flywaydb.core");
+                if (hikariLog.getAppender(HikariLoggingAppender.NAME) != null) {
+                    hikariLog.detachAppender(HikariLoggingAppender.NAME);
+                }
+                if (flywayLog.getAppender(HikariLoggingAppender.NAME) != null) {
+                    flywayLog.detachAppender(HikariLoggingAppender.NAME);
                 }
             }
         }
@@ -228,7 +237,7 @@ public void initStorage(boolean shouldWait) throws DbInitException {
         }
         try {
             ConnectionPool.initPool(this, shouldWait);
-            GeneralQueries.createTablesIfNotExists(this);
+            FlywayMigration.startMigration(this);
         } catch (Exception e) {
             throw new DbInitException(e);
         }

src/main/java/io/supertokens/storage/postgresql/config/PostgreSQLConfig.java

@@ -115,6 +115,10 @@ public class PostgreSQLConfig {
     @IgnoreForAnnotationCheck
     boolean isValidAndNormalised = false;
 
+    @JsonProperty
+    @IgnoreForAnnotationCheck
+    private boolean access_token_signing_key_dynamic = true;
+
     public static Set<String> getValidFields() {
         PostgreSQLConfig config = new PostgreSQLConfig();
         JsonObject configObj = new GsonBuilder().serializeNulls().create().toJsonTree(config).getAsJsonObject();
@@ -234,6 +238,10 @@ public String getThirdPartyUsersTable() {
         return postgresql_thirdparty_users_table_name;
     }
 
+    public boolean getAccessTokenSigningKeyDynamic() {
+        return access_token_signing_key_dynamic;
+    }
+
     public String getThirdPartyUserToTenantTable() {
         return addSchemaAndPrefixToTableName("thirdparty_user_to_tenant");
     }
@@ -302,6 +310,10 @@ public String getTotpUsedCodesTable() {
         return addSchemaAndPrefixToTableName("totp_used_codes");
     }
 
+    public String getFlywaySchemaHistory() {
+        return addSchemaAndPrefixToTableName("flyway_schema_history");
+    }
+
     private String addSchemaAndPrefixToTableName(String tableName) {
         return addSchemaToTableName(postgresql_table_names_prefix + tableName);
     }

src/main/java/io/supertokens/storage/postgresql/migrations/V1__init_database.java

@@ -0,0 +1,37 @@
+/*
+ *    Copyright (c) 2024, VRAI Labs and/or its affiliates. All rights reserved.
+ *
+ *    This software is licensed under the Apache License, Version 2.0 (the
+ *    "License") as published by the Apache Software Foundation.
+ *
+ *    You may not use this file except in compliance with the License. You may
+ *    obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *    Unless required by applicable law or agreed to in writing, software
+ *    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ *    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ *    License for the specific language governing permissions and limitations
+ *    under the License.
+ */
+
+package io.supertokens.storage.postgresql.migrations;
+
+import io.supertokens.storage.postgresql.MigrationContextManager;
+import io.supertokens.storage.postgresql.Start;
+import io.supertokens.storage.postgresql.queries.GeneralQueries;
+import org.flywaydb.core.api.migration.BaseJavaMigration;
+import org.flywaydb.core.api.migration.Context;
+
+import java.util.Map;
+
+
+public class V1__init_database extends BaseJavaMigration {
+
+    @Override
+    public void migrate(Context context) throws Exception {
+        Map<String, String> ph = context.getConfiguration().getPlaceholders();
+        Start start = MigrationContextManager.getContext(ph.get("process_id"));
+        GeneralQueries.createTablesIfNotExists(start);
+    }
+
+}

src/main/java/io/supertokens/storage/postgresql/migrations/V2__plugin_version_3_0_0.java

@@ -0,0 +1,52 @@
+/*
+ *    Copyright (c) 2023, VRAI Labs and/or its affiliates. All rights reserved.
+ *
+ *    This software is licensed under the Apache License, Version 2.0 (the
+ *    "License") as published by the Apache Software Foundation.
+ *
+ *    You may not use this file except in compliance with the License. You may
+ *    obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *    Unless required by applicable law or agreed to in writing, software
+ *    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ *    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ *    License for the specific language governing permissions and limitations
+ *    under the License.
+ */
+
+package io.supertokens.storage.postgresql.migrations;
+
+import io.supertokens.storage.postgresql.MigrationContextManager;
+import io.supertokens.storage.postgresql.Start;
+import io.supertokens.storage.postgresql.config.Config;
+import org.flywaydb.core.api.migration.BaseJavaMigration;
+import org.flywaydb.core.api.migration.Context;
+import java.sql.Statement;
+import java.util.Map;
+
+public class V2__plugin_version_3_0_0 extends BaseJavaMigration {
+
+    @Override
+    public void migrate(Context context) throws Exception {
+        Map<String, String> ph = context.getConfiguration().getPlaceholders();
+        Start start = MigrationContextManager.getContext(ph.get("process_id"));
+        String sessionInfoTable = Config.getConfig(start).getSessionInfoTable();
+        String JWTSigningKeysTable = Config.getConfig(start).getJWTSigningKeysTable();
+        String accessTokenSigningKeysTable = Config.getConfig(start).getAccessTokenSigningKeysTable();
+
+        try (Statement statement = context.getConnection().createStatement()) {
+            // Add a new column with a default value
+            statement.execute("ALTER TABLE " + sessionInfoTable + " ADD COLUMN IF NOT EXISTS use_static_key BOOLEAN NOT NULL DEFAULT" +
+                    "(" + !Boolean.parseBoolean(ph.get("access_token_signing_key_dynamic")) + ")");
+            // Alter the column to drop the default value
+            statement.execute("ALTER TABLE " + sessionInfoTable + " ALTER COLUMN " +
+                    "use_static_key DROP DEFAULT");
+
+            // Insert data into jwt_signing_keys from session_access_token_signing_keys
+            statement.execute("INSERT INTO " + JWTSigningKeysTable + " (key_id, key_string, algorithm, " +
+                    "created_at) " +
+                    "SELECT CONCAT('s-', created_at_time) as key_id, value as key_string, 'RS256' as algorithm, created_at_time as created_at " +
+                    "FROM " + accessTokenSigningKeysTable);
+        }
+    }
+}