From 685f946d0c3d154c36a3a3be283d0d16969e5dae Mon Sep 17 00:00:00 2001
From: Sattvik Chakravarthy <sattvik@gmail.com>
Date: Tue, 7 Jan 2025 19:14:56 +0530
Subject: [PATCH] fix: session queries

---
 .../postgresql/queries/SessionQueries.java    | 57 ++++++++++++++-----
 1 file changed, 44 insertions(+), 13 deletions(-)

diff --git a/src/main/java/io/supertokens/storage/postgresql/queries/SessionQueries.java b/src/main/java/io/supertokens/storage/postgresql/queries/SessionQueries.java
index c5108fd4..30857b08 100644
--- 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,