fix: core db crash

CHANGELOG.md

@@ -7,6 +7,10 @@ to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## [Unreleased]
 
+## [7.0.1] - 2024-04-17
+
+- Fixes issues with partial failures during tenant creation
+
 ## [7.0.0] - 2024-03-13
 
 - Replace `TotpNotEnabledError` with `UnknownUserIdTotpError`.

build.gradle

@@ -2,7 +2,7 @@ plugins {
     id 'java-library'
 }
 
-version = "7.0.0"
+version = "7.0.1"
 
 repositories {
     mavenCentral()

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

@@ -19,7 +19,9 @@
 
 import com.zaxxer.hikari.HikariConfig;
 import com.zaxxer.hikari.HikariDataSource;
+import io.supertokens.pluginInterface.Storage;
 import io.supertokens.pluginInterface.exceptions.DbInitException;
+import io.supertokens.pluginInterface.exceptions.StorageQueryException;
 import io.supertokens.storage.postgresql.config.Config;
 import io.supertokens.storage.postgresql.config.PostgreSQLConfig;
 import io.supertokens.storage.postgresql.output.Logging;
@@ -35,12 +37,14 @@ public class ConnectionPool extends ResourceDistributor.SingletonResource {
     private static final String RESOURCE_KEY = "io.supertokens.storage.postgresql.ConnectionPool";
     private HikariDataSource hikariDataSource;
     private final Start start;
+    private PostConnectCallback postConnectCallback;
 
-    private ConnectionPool(Start start) {
+    private ConnectionPool(Start start, PostConnectCallback postConnectCallback) {
         this.start = start;
+        this.postConnectCallback = postConnectCallback;
     }
 
-    private synchronized void initialiseHikariDataSource() throws SQLException {
+    private synchronized void initialiseHikariDataSource() throws SQLException, StorageQueryException {
         if (this.hikariDataSource != null) {
             return;
         }
@@ -99,6 +103,17 @@ private synchronized void initialiseHikariDataSource() throws SQLException {
         } catch (Exception e) {
             throw new SQLException(e);
         }
+
+        try {
+            this.postConnectCallback.apply();
+        } catch (StorageQueryException e) {
+            // if an exception happens here, we want to set the hikariDataSource to null once again so that
+            // whenever the getConnection is called again, we want to re-attempt creation of tables and tenant
+            // entries for this storage
+            hikariDataSource.close();
+            hikariDataSource = null;
+            throw e;
+        }
     }
 
     private static int getTimeToWaitToInit(Start start) {
@@ -133,7 +148,7 @@ static boolean isAlreadyInitialised(Start start) {
         return getInstance(start) != null && getInstance(start).hikariDataSource != null;
     }
 
-    static void initPool(Start start, boolean shouldWait) throws DbInitException {
+    static void initPool(Start start, boolean shouldWait, PostConnectCallback postConnectCallback) throws DbInitException {
         if (isAlreadyInitialised(start)) {
             return;
         }
@@ -146,7 +161,7 @@ static void initPool(Start start, boolean shouldWait) throws DbInitException {
                         " specified the correct values for ('postgresql_host' and 'postgresql_port') or for "
                         + "'postgresql_connection_uri'";
         try {
-            ConnectionPool con = new ConnectionPool(start);
+            ConnectionPool con = new ConnectionPool(start, postConnectCallback);
             start.getResourceDistributor().setResource(RESOURCE_KEY, con);
             while (true) {
                 try {
@@ -189,7 +204,7 @@ static void initPool(Start start, boolean shouldWait) throws DbInitException {
         }
     }
 
-    public static Connection getConnection(Start start) throws SQLException {
+    public static Connection getConnection(Start start) throws SQLException, StorageQueryException {
         if (getInstance(start) == null) {
             throw new IllegalStateException("Please call initPool before getConnection");
         }
@@ -216,4 +231,9 @@ static void close(Start start) {
             }
         }
     }
+
+    @FunctionalInterface
+    public static interface PostConnectCallback {
+        void apply() throws StorageQueryException;
+    }
 }

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

@@ -98,10 +98,7 @@
 import java.sql.Connection;
 import java.sql.SQLException;
 import java.sql.SQLTransactionRollbackException;
-import java.util.ArrayList;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Set;
+import java.util.*;
 
 import static io.supertokens.storage.postgresql.QueryExecutorTemplate.execute;
 
@@ -129,6 +126,7 @@ public class Start
     boolean enabled = true;
     static Thread mainThread = Thread.currentThread();
     private Thread shutdownHook;
+    private Set<TenantIdentifier> tenantIdentifiers = new HashSet<>();
 
     private boolean isBaseTenant = false;
 
@@ -228,13 +226,30 @@ public void initStorage(boolean shouldWait) throws DbInitException {
             mainThread = Thread.currentThread();
         }
         try {
-            ConnectionPool.initPool(this, shouldWait);
-            GeneralQueries.createTablesIfNotExists(this);
+            ConnectionPool.initPool(this, shouldWait, () -> {
+                try {
+                    GeneralQueries.createTablesIfNotExists(this);
+                } catch (SQLException e) {
+                    throw new StorageQueryException(e);
+                }
+                for (TenantIdentifier tenantIdentifier : this.tenantIdentifiers) {
+                    try {
+                        this.addTenantIdInTargetStorage(tenantIdentifier);
+                    } catch (DuplicateTenantException e) {
+                        // ignore
+                    }
+                }
+            });
         } catch (Exception e) {
             throw new DbInitException(e);
         }
     }
 
+    @Override
+    public synchronized void addTenantIdentifier(TenantIdentifier tenantIdentifier) {
+        this.tenantIdentifiers.add(tenantIdentifier);
+    }
+
     @Override
     public <T> T startTransaction(TransactionLogic<T> logic)
             throws StorageTransactionLogicException, StorageQueryException {

src/main/java/io/supertokens/storage/postgresql/queries/MultitenancyQueries.java

@@ -37,6 +37,8 @@
 import static io.supertokens.storage.postgresql.config.Config.getConfig;
 
 public class MultitenancyQueries {
+    public static boolean simulateErrorInAddingTenantIdInTargetStorage_forTesting = false;
+
     static String getQueryToCreateTenantConfigsTable(Start start) {
         String schema = Config.getConfig(start).getTableSchema();
         String tenantConfigsTable = Config.getConfig(start).getTenantConfigsTable();
@@ -274,6 +276,10 @@ public static TenantConfig[] getAllTenants(Start start) throws StorageQueryExcep
     public static void addTenantIdInTargetStorage(Start start, TenantIdentifier tenantIdentifier) throws
             StorageTransactionLogicException, StorageQueryException {
         {
+            if (Start.isTesting && simulateErrorInAddingTenantIdInTargetStorage_forTesting) {
+                throw new StorageTransactionLogicException(new SQLException("Simulated error in addTenantIdInTargetStorage"));
+            }
+
             start.startTransaction(con -> {
                 Connection sqlCon = (Connection) con.getConnection();
                 long currentTime = System.currentTimeMillis();

src/test/java/io/supertokens/storage/postgresql/test/Utils.java

@@ -20,6 +20,7 @@
 import io.supertokens.Main;
 import io.supertokens.pluginInterface.PluginInterfaceTesting;
 import io.supertokens.storage.postgresql.Start;
+import io.supertokens.storage.postgresql.queries.MultitenancyQueries;
 import io.supertokens.storageLayer.StorageLayer;
 import org.apache.tomcat.util.http.fileupload.FileUtils;
 import org.junit.rules.TestRule;
@@ -78,6 +79,8 @@ public static void reset() {
         PluginInterfaceTesting.isTesting = true;
         Start.isTesting = true;
         Main.makeConsolePrintSilent = true;
+        MultitenancyQueries.simulateErrorInAddingTenantIdInTargetStorage_forTesting = false;
+
         String installDir = "../";
         try {
             // if the default config is not the same as the current config, we must reset