fix: core db crash (#215)

* fix: core db crash

* fix: exception

* fix: PR comments

* fix: update test

* fix: update test

* fix: update test

* fix: version

* fix: rename

* fix: pr comments

* fix: pr comments

* fix: pr comments

* fix: update tests

Author: sattvikc

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()

pluginInterfaceSupported.json

@@ -1,6 +1,6 @@
 {
   "_comment": "contains a list of plugin interfaces branch names that this core supports",
   "versions": [
-    "6.0"
+    "6.1"
   ]
 }

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

@@ -20,6 +20,7 @@
 import com.zaxxer.hikari.HikariConfig;
 import com.zaxxer.hikari.HikariDataSource;
 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 +36,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 +102,19 @@ private synchronized void initialiseHikariDataSource() throws SQLException {
         } catch (Exception e) {
             throw new SQLException(e);
         }
+
+        try {
+            try (Connection con = hikariDataSource.getConnection()) {
+                this.postConnectCallback.apply(con);
+            }
+        } 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 +149,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 +162,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 +205,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 +232,9 @@ static void close(Start start) {
             }
         }
     }
+
+    @FunctionalInterface
+    public static interface PostConnectCallback {
+        void apply(Connection connection) 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;
 
@@ -218,7 +215,7 @@ public void stopLogging() {
     }
 
     @Override
-    public void initStorage(boolean shouldWait) throws DbInitException {
+    public void initStorage(boolean shouldWait, List<TenantIdentifier> tenantIdentifiers) throws DbInitException {
         if (ConnectionPool.isAlreadyInitialised(this)) {
             return;
         }
@@ -228,8 +225,20 @@ public void initStorage(boolean shouldWait) throws DbInitException {
             mainThread = Thread.currentThread();
         }
         try {
-            ConnectionPool.initPool(this, shouldWait);
-            GeneralQueries.createTablesIfNotExists(this);
+            ConnectionPool.initPool(this, shouldWait, (con) -> {
+                try {
+                    GeneralQueries.createTablesIfNotExists(this, con);
+                } catch (SQLException e) {
+                    throw new StorageQueryException(e);
+                }
+                for (TenantIdentifier tenantIdentifier : tenantIdentifiers) {
+                    try {
+                        this.addTenantIdInTargetStorage_Transaction(con, tenantIdentifier);
+                    } catch (DuplicateTenantException e) {
+                        // ignore
+                    }
+                }
+            });
         } catch (Exception e) {
             throw new DbInitException(e);
         }
@@ -465,7 +474,7 @@ public void deleteAllInformation() throws StorageQueryException {
         }
         ProcessState.getInstance(this).clear();
         try {
-            initStorage(false);
+            initStorage(false, new ArrayList<>());
             enabled = true; // Allow get connection to work, to delete the data
             GeneralQueries.deleteAllTables(this);
 
@@ -2283,6 +2292,22 @@ public void addTenantIdInTargetStorage(TenantIdentifier tenantIdentifier)
         }
     }
 
+    public void addTenantIdInTargetStorage_Transaction(Connection con, TenantIdentifier tenantIdentifier)
+            throws DuplicateTenantException, StorageQueryException {
+        try {
+            MultitenancyQueries.addTenantIdInTargetStorage_Transaction(this, con, tenantIdentifier);
+        } catch (SQLException e) {
+            if (e instanceof PSQLException) {
+                PostgreSQLConfig config = Config.getConfig(this);
+                if (isPrimaryKeyError(((PSQLException) e).getServerErrorMessage(),
+                        config.getTenantsTable())) {
+                    throw new DuplicateTenantException();
+                }
+            }
+            throw new StorageQueryException(e);
+        }
+    }
+
     @Override
     public void overwriteTenantConfig(TenantConfig tenantConfig)
             throws TenantOrAppNotFoundException, StorageQueryException, DuplicateThirdPartyIdException,