From c13ee114d6169ed3e91349eea859a11c4d814e17 Mon Sep 17 00:00:00 2001
From: zhangxiao <140503120+PeachThinking@users.noreply.github.com>
Date: Fri, 12 Jul 2024 11:37:36 +0800
Subject: [PATCH] fix(session): support auto recreate session when upgrade
 (#2959)

* fix failed to auto recreate session when upgrade

* The impact of code changes is limited to multiCloud

* response to reivew

* response to review

* response to review
---
 .../service/session/ConnectSessionService.java    | 15 ++++++++++++---
 1 file changed, 12 insertions(+), 3 deletions(-)

diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/ConnectSessionService.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/ConnectSessionService.java
index 90ed9fcda9..d3010025ee 100644
--- a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/ConnectSessionService.java
+++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/ConnectSessionService.java
@@ -276,7 +276,11 @@ public ConnectionSession create(@NotNull CreateSessionReq req) {
         DefaultConnectSessionIdGenerator idGenerator = new DefaultConnectSessionIdGenerator();
         idGenerator.setDatabaseId(req.getDbId());
         idGenerator.setFixRealId(StringUtils.isBlank(req.getRealId()) ? null : req.getRealId());
-        idGenerator.setHost(stateHostGenerator.getHost());
+        if (Objects.nonNull(req.getFrom()) && isOBCloudEnvironment()) {
+            idGenerator.setHost(req.getFrom());
+        } else {
+            idGenerator.setHost(stateHostGenerator.getHost());
+        }
         sessionFactory.setIdGenerator(idGenerator);
         long timeoutMillis = TimeUnit.MILLISECONDS.convert(sessionProperties.getTimeoutMins(), TimeUnit.MINUTES);
         timeoutMillis = timeoutMillis + this.connectionSessionManager.getScanIntervalMillis();
@@ -330,8 +334,8 @@ public ConnectionSession nullSafeGet(@NotNull String sessionId, boolean autoCrea
         ConnectionSession session = connectionSessionManager.getSession(sessionId);
         if (session == null) {
             CreateSessionReq req = new DefaultConnectSessionIdGenerator().getKeyFromId(sessionId);
-            boolean autoRecreate = cloudMetadataClient.supportsCloudMetadata();
-            if (!autoCreate || (!StringUtils.equals(req.getFrom(), stateHostGenerator.getHost()) && !autoRecreate)) {
+            if (!autoCreate
+                    || (!StringUtils.equals(req.getFrom(), stateHostGenerator.getHost()) && !isOBCloudEnvironment())) {
                 throw new NotFoundException(ResourceType.ODC_SESSION, "ID", sessionId);
             }
             Lock lock = this.sessionId2Lock.computeIfAbsent(sessionId, s -> new ReentrantLock());
@@ -466,4 +470,9 @@ private void preCheckSessionLimit() {
         }
     }
 
+    private boolean isOBCloudEnvironment() {
+        return cloudMetadataClient.supportsCloudMetadata()
+                && Boolean.FALSE.equals(cloudMetadataClient.supportsCloudParentUid());
+    }
+
 }