Merge branch 'test-cicd/upgrade-postgresql-to-13' of github.com:super…

…tokens/supertokens-postgresql-plugin into test-cicd/upgrade-postgresql-to-13
supertokens · Jan 9, 2025 · 904e547 · 904e547
2 parents 1c78d8f + f1359a4
commit 904e547
Show file tree

Hide file tree

Showing 5 changed files with 59 additions and 18 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -10,6 +10,7 @@ to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 ## [7.3.0] 
 
 - Adds tables and queries for Bulk Import
+- Optimize getUserIdMappingWithEitherSuperTokensUserIdOrExternalUserId query
 
 ### Migration
 
@@ -34,6 +35,7 @@ CREATE INDEX IF NOT EXISTS bulk_import_users_pagination_index1 ON bulk_import_us
  
 CREATE INDEX IF NOT EXISTS bulk_import_users_pagination_index2 ON bulk_import_users (app_id, created_at DESC, id DESC);
 ```
+
 ## [7.2.0] - 2024-10-03
 
 - Compatible with plugin interface version 6.3

diff --git a/jar/postgresql-plugin-7.3.0.jar b/jar/postgresql-plugin-7.3.0.jar
diff --git a/src/main/java/io/supertokens/storage/postgresql/queries/SessionQueries.java b/src/main/java/io/supertokens/storage/postgresql/queries/SessionQueries.java
@@ -126,9 +126,6 @@ public static void createNewSession(Start start, TenantIdentifier tenantIdentifi
     public static SessionInfo getSessionInfo_Transaction(Start start, Connection con, TenantIdentifier tenantIdentifier,
                                                          String sessionHandle)
             throws SQLException, StorageQueryException {
-        // we do this as two separate queries and not one query with left join cause psql does not
-        // support left join with for update if the right table returns null.
-
         String QUERY =
                 "SELECT session_handle, user_id, refresh_token_hash_2, session_data, " +
                         "expires_at, created_at_time, jwt_user_payload, use_static_key FROM " +
@@ -149,21 +146,55 @@ public static SessionInfo getSessionInfo_Transaction(Start start, Connection con
             return null;
         }
 
-        QUERY = "SELECT primary_or_recipe_user_id FROM " + getConfig(start).getUsersTable()
-                + " WHERE app_id = ? AND user_id = ?";
-
-        return execute(con, QUERY, pst -> {
+        QUERY = "SELECT external_user_id " +
+                "FROM " + getConfig(start).getUserIdMappingTable() + " um2 " +
+                "WHERE um2.app_id = ? AND um2.supertokens_user_id IN (" +
+                    "SELECT primary_or_recipe_user_id " + 
+                    "FROM " + getConfig(start).getUsersTable() + " " +
+                    "WHERE app_id = ? AND user_id IN (" +
+                        "SELECT um1.supertokens_user_id as user_id " +
+                        "FROM " + getConfig(start).getUserIdMappingTable() + " um1 " +
+                        "WHERE um1.app_id = ? AND um1.external_user_id = ? " +
+                        "UNION ALL " +
+                        "SELECT ? " +
+                        "LIMIT 1" +
+                    ")" +
+                ") " +
+                "UNION ALL " +
+                "SELECT primary_or_recipe_user_id " +
+                "FROM " + getConfig(start).getUsersTable() + " " +
+                "WHERE app_id = ? AND user_id IN (" +
+                    "SELECT um1.supertokens_user_id as user_id " +
+                    "FROM " + getConfig(start).getUserIdMappingTable() + " um1 " +
+                    "WHERE um1.app_id = ? AND um1.external_user_id = ? " +
+                    "UNION ALL " +
+                    "SELECT ? " +
+                    "LIMIT 1" +
+                ") " +
+                "LIMIT 1";
+
+        String finalUserId = execute(con, QUERY, pst -> {
             pst.setString(1, tenantIdentifier.getAppId());
-            pst.setString(2, sessionInfo.recipeUserId);
+            pst.setString(2, tenantIdentifier.getAppId());
+            pst.setString(3, tenantIdentifier.getAppId());
+            pst.setString(4, sessionInfo.recipeUserId);
+            pst.setString(5, sessionInfo.recipeUserId);
+            pst.setString(6, tenantIdentifier.getAppId());
+            pst.setString(7, tenantIdentifier.getAppId());
+            pst.setString(8, sessionInfo.recipeUserId);
+            pst.setString(9, sessionInfo.recipeUserId);
         }, result -> {
             if (result.next()) {
-                String primaryUserId = result.getString("primary_or_recipe_user_id");
-                if (primaryUserId != null) {
-                    sessionInfo.userId = primaryUserId;
-                }
+                return result.getString(1);
             }
-            return sessionInfo;
+            return sessionInfo.recipeUserId;
         });
+
+        if (finalUserId != null) {
+            sessionInfo.userId = finalUserId;
+        }
+
+        return sessionInfo;
     }
 
     public static void updateSessionInfo_Transaction(Start start, Connection con, TenantIdentifier tenantIdentifier,

diff --git a/src/main/java/io/supertokens/storage/postgresql/queries/UserIdMappingQueries.java b/src/main/java/io/supertokens/storage/postgresql/queries/UserIdMappingQueries.java
@@ -139,12 +139,16 @@ public static UserIdMapping[] getUserIdMappingWithEitherSuperTokensUserIdOrExter
                                                                                               String userId)
             throws SQLException, StorageQueryException {
         String QUERY = "SELECT * FROM " + Config.getConfig(start).getUserIdMappingTable()
-                + " WHERE app_id = ? AND (supertokens_user_id = ? OR external_user_id = ?)";
+                + " WHERE app_id = ? AND supertokens_user_id = ?"
+                + " UNION ALL "
+                + "SELECT * FROM " + Config.getConfig(start).getUserIdMappingTable()
+                + " WHERE app_id = ? AND external_user_id = ?";
 
         return execute(start, QUERY, pst -> {
             pst.setString(1, appIdentifier.getAppId());
             pst.setString(2, userId);
-            pst.setString(3, userId);
+            pst.setString(3, appIdentifier.getAppId());
+            pst.setString(4, userId);
         }, result -> {
             ArrayList<UserIdMapping> userIdMappingArray = new ArrayList<>();
             while (result.next()) {
@@ -375,12 +379,16 @@ public static UserIdMapping[] getUserIdMappingWithEitherSuperTokensUserIdOrExter
                                                                                                           String userId)
             throws SQLException, StorageQueryException {
         String QUERY = "SELECT * FROM " + Config.getConfig(start).getUserIdMappingTable()
-                + " WHERE app_id = ? AND (supertokens_user_id = ? OR external_user_id = ?)";
+                + " WHERE app_id = ? AND supertokens_user_id = ?"
+                + " UNION ALL "
+                + "SELECT * FROM " + Config.getConfig(start).getUserIdMappingTable()
+                + " WHERE app_id = ? AND external_user_id = ?";
 
         return execute(sqlCon, QUERY, pst -> {
             pst.setString(1, appIdentifier.getAppId());
             pst.setString(2, userId);
-            pst.setString(3, userId);
+            pst.setString(3, appIdentifier.getAppId());
+            pst.setString(4, userId);
         }, result -> {
             ArrayList<UserIdMapping> userIdMappingArray = new ArrayList<>();
             while (result.next()) {

diff --git a/src/test/java/io/supertokens/storage/postgresql/test/OneMillionUsersTest.java b/src/test/java/io/supertokens/storage/postgresql/test/OneMillionUsersTest.java
@@ -990,7 +990,7 @@ public void testWithOneMillionUsers() throws Exception {
 
         Main main = startCronProcess(String.valueOf(4)); // ci uses instances with 4 cores..
 
-        int NUMBER_OF_USERS_TO_UPLOAD = 300000;
+        int NUMBER_OF_USERS_TO_UPLOAD = 300000; // 300k
 
         if (StorageLayer.getBaseStorage(main).getType() != STORAGE_TYPE.SQL || StorageLayer.isInMemDb(main)) {
             return;