From 98472941821fc40e8ed55498c62ed5758b7eea97 Mon Sep 17 00:00:00 2001
From: ChandanBalajiBP <104058110+ChandanBalajiBP@users.noreply.github.com>
Date: Thu, 10 Aug 2023 18:08:00 +0530
Subject: [PATCH] fix: Save and authorise flow in current editing environment
 (#26212)

Co-authored-by: manish kumar <manish@appsmith.com>
---
 app/client/src/ce/api/ApiUtils.ts                  | 11 +++++++++--
 app/client/src/entities/Datasource/index.ts        |  1 +
 .../src/pages/common/datasourceAuth/index.tsx      | 14 ++++++++++----
 app/client/src/sagas/DatasourcesSagas.ts           |  2 +-
 .../controllers/ce/DatasourceControllerCE.java     |  2 +-
 5 files changed, 22 insertions(+), 8 deletions(-)

diff --git a/app/client/src/ce/api/ApiUtils.ts b/app/client/src/ce/api/ApiUtils.ts
index ffd0dc16c5..f13d6e1c60 100644
--- a/app/client/src/ce/api/ApiUtils.ts
+++ b/app/client/src/ce/api/ApiUtils.ts
@@ -26,7 +26,10 @@ import { getAppsmithConfigs } from "@appsmith/configs";
 import * as Sentry from "@sentry/react";
 import { CONTENT_TYPE_HEADER_KEY } from "constants/ApiEditorConstants/CommonApiConstants";
 import { isAirgapped } from "@appsmith/utils/airgapHelpers";
-import { getCurrentEnvironment } from "@appsmith/utils/Environments";
+import {
+  getCurrentEnvironment,
+  getCurrentEditingEnvID,
+} from "@appsmith/utils/Environments";
 
 const executeActionRegex = /actions\/execute/;
 const timeoutErrorRegex = /timeout of (\d+)ms exceeded/;
@@ -96,7 +99,11 @@ export const apiRequestInterceptor = (config: AxiosRequestConfig) => {
   const activeEnv = getCurrentEnvironment();
 
   if (activeEnv && config.headers) {
-    config.headers.environmentId = activeEnv;
+    if (config.url?.indexOf("/code") !== -1) {
+      config.headers.environmentId = getCurrentEditingEnvID();
+    } else {
+      config.headers.environmentId = activeEnv;
+    }
   }
 
   const anonymousId = AnalyticsUtil.getAnonymousId();
diff --git a/app/client/src/entities/Datasource/index.ts b/app/client/src/entities/Datasource/index.ts
index 1ff0c638f4..36fb15a98b 100644
--- a/app/client/src/entities/Datasource/index.ts
+++ b/app/client/src/entities/Datasource/index.ts
@@ -53,6 +53,7 @@ export interface DatasourceAuthentication {
   authenticationStatus?: string;
   authenticationType?: string;
   secretExists?: Record<string, boolean>;
+  isAuthorized?: boolean;
 }
 
 export interface DatasourceColumns {
diff --git a/app/client/src/pages/common/datasourceAuth/index.tsx b/app/client/src/pages/common/datasourceAuth/index.tsx
index 3f3d3edea3..f494be3ed6 100644
--- a/app/client/src/pages/common/datasourceAuth/index.tsx
+++ b/app/client/src/pages/common/datasourceAuth/index.tsx
@@ -35,7 +35,10 @@ import { integrationEditorURL } from "RouteBuilder";
 import { getQueryParams } from "utils/URLUtils";
 import type { AppsmithLocationState } from "utils/history";
 import type { PluginType } from "entities/Action";
-import { getCurrentEnvName } from "@appsmith/utils/Environments";
+import {
+  getCurrentEnvName,
+  getCurrentEditingEnvID,
+} from "@appsmith/utils/Environments";
 
 interface Props {
   datasource: Datasource;
@@ -234,9 +237,12 @@ function DatasourceAuth({
     }
   }, [triggerSave]);
   const isAuthorized =
-    datasource?.datasourceStorages &&
-    datasource?.datasourceStorages[currentEnvironment]?.datasourceConfiguration
-      ?.authentication?.authenticationStatus === AuthenticationStatus.SUCCESS;
+    datasource?.datasourceStorages && authType === AuthType.OAUTH2
+      ? datasource?.datasourceStorages[getCurrentEditingEnvID()]
+          ?.datasourceConfiguration?.authentication?.isAuthorized
+      : datasource?.datasourceStorages[currentEnvironment]
+          ?.datasourceConfiguration?.authentication?.authenticationStatus ===
+        AuthenticationStatus.SUCCESS;
 
   // Button Operations for respective buttons.
 
diff --git a/app/client/src/sagas/DatasourcesSagas.ts b/app/client/src/sagas/DatasourcesSagas.ts
index 93cd3f93e5..26c2f33289 100644
--- a/app/client/src/sagas/DatasourcesSagas.ts
+++ b/app/client/src/sagas/DatasourcesSagas.ts
@@ -596,7 +596,7 @@ function* redirectAuthorizationCodeSaga(
   const isImport: string = yield select(getWorkspaceIdForImport);
 
   if (pluginType === PluginType.API) {
-    window.location.href = `/api/v1/datasources/${datasourceId}/pages/${pageId}/code`;
+    window.location.href = `/api/v1/datasources/${datasourceId}/pages/${pageId}/code?environmentId=${getCurrentEditingEnvID()}`;
   } else {
     try {
       // Get an "appsmith token" from the server
diff --git a/app/server/appsmith-server/src/main/java/com/appsmith/server/controllers/ce/DatasourceControllerCE.java b/app/server/appsmith-server/src/main/java/com/appsmith/server/controllers/ce/DatasourceControllerCE.java
index 532923be3c..1a90eaa475 100644
--- a/app/server/appsmith-server/src/main/java/com/appsmith/server/controllers/ce/DatasourceControllerCE.java
+++ b/app/server/appsmith-server/src/main/java/com/appsmith/server/controllers/ce/DatasourceControllerCE.java
@@ -153,7 +153,7 @@ public class DatasourceControllerCE {
             @PathVariable String datasourceId,
             @PathVariable String pageId,
             ServerWebExchange serverWebExchange,
-            @RequestHeader(name = FieldName.ENVIRONMENT_ID, required = false) String environmentId) {
+            @RequestParam String environmentId) {
         log.debug(
                 "Going to retrieve token request URL for datasource with id: {} and page id: {}", datasourceId, pageId);
         return authenticationService