fix: PR comments

supertokens · Apr 16, 2024 · cf5b3b2 · cf5b3b2
1 parent 1ff1abc
commit cf5b3b2
Show file tree

Hide file tree

Showing 5 changed files with 329 additions and 10 deletions.
diff --git a/src/main/java/io/supertokens/storage/postgresql/ConnectionPool.java b/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,14 +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 Runnable postConnectCallback;
+    private PostConnectCallback postConnectCallback;
 
-    private ConnectionPool(Start start, Runnable postConnectCallback) {
+    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;
         }
@@ -98,10 +100,20 @@ private synchronized void initialiseHikariDataSource() throws SQLException {
         config.setPoolName(start.getUserPoolId() + "~" + start.getConnectionPoolId());
         try {
             hikariDataSource = new HikariDataSource(config);
-            this.postConnectCallback.run();
         } 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) {
@@ -136,7 +148,7 @@ static boolean isAlreadyInitialised(Start start) {
         return getInstance(start) != null && getInstance(start).hikariDataSource != null;
     }
 
-    static void initPool(Start start, boolean shouldWait, Runnable postConnectCallback) throws DbInitException {
+    static void initPool(Start start, boolean shouldWait, PostConnectCallback postConnectCallback) throws DbInitException {
         if (isAlreadyInitialised(start)) {
             return;
         }
@@ -192,7 +204,7 @@ static void initPool(Start start, boolean shouldWait, Runnable postConnectCallba
         }
     }
 
-    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");
         }
@@ -219,4 +231,9 @@ static void close(Start start) {
             }
         }
     }
+
+    @FunctionalInterface
+    public static interface PostConnectCallback {
+        void apply() throws StorageQueryException;
+    }
 }
diff --git a/src/main/java/io/supertokens/storage/postgresql/Start.java b/src/main/java/io/supertokens/storage/postgresql/Start.java
@@ -229,16 +229,14 @@ public void initStorage(boolean shouldWait) throws DbInitException {
             ConnectionPool.initPool(this, shouldWait, () -> {
                 try {
                     GeneralQueries.createTablesIfNotExists(this);
-                } catch (Exception e2) {
-                    throw new RuntimeException(e2);
+                } catch (SQLException e) {
+                    throw new StorageQueryException(e);
                 }
                 for (TenantIdentifier tenantIdentifier : this.tenantIdentifiers) {
                     try {
                         this.addTenantIdInTargetStorage(tenantIdentifier);
                     } catch (DuplicateTenantException e) {
                         // ignore
-                    } catch (StorageQueryException e) {
-                        throw new RuntimeException(e);
                     }
                 }
             });

diff --git a/src/main/java/io/supertokens/storage/postgresql/queries/MultitenancyQueries.java b/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 = 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) {
+                throw new StorageTransactionLogicException(new SQLException("Simulated error in addTenantIdInTargetStorage"));
+            }
+
             start.startTransaction(con -> {
                 Connection sqlCon = (Connection) con.getConnection();
                 long currentTime = System.currentTimeMillis();

diff --git a/src/test/java/io/supertokens/storage/postgresql/test/Utils.java b/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 = false;
+
         String installDir = "../";
         try {
             // if the default config is not the same as the current config, we must reset