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

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

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 {

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.
 

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

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