From 710a0d4d2f5085e6532d93d130ad63096713868f Mon Sep 17 00:00:00 2001
From: SuZhou-Joe <suzhou@amazon.com>
Date: Mon, 1 Apr 2024 09:43:53 +0800
Subject: [PATCH] feat: add APIs to support plugin state in request (#312)

* feat: add APIs to support plugin state in request

Signed-off-by: SuZhou-Joe <suzhou@amazon.com>

* feat: add APIs to support plugin state in request

Signed-off-by: SuZhou-Joe <suzhou@amazon.com>

---------

Signed-off-by: SuZhou-Joe <suzhou@amazon.com>
---
 src/core/server/utils/index.ts          |  1 +
 src/core/server/utils/workspace.test.ts | 19 +++++++++++
 src/core/server/utils/workspace.ts      | 43 +++++++++++++++++++++++++
 src/legacy/server/osd_server.d.ts       |  8 +++++
 4 files changed, 71 insertions(+)
 create mode 100644 src/core/server/utils/workspace.test.ts
 create mode 100644 src/core/server/utils/workspace.ts

diff --git a/src/core/server/utils/index.ts b/src/core/server/utils/index.ts
index 42b01e72b0d1..a20b8c4c4e5b 100644
--- a/src/core/server/utils/index.ts
+++ b/src/core/server/utils/index.ts
@@ -33,3 +33,4 @@ export * from './from_root';
 export * from './package_json';
 export * from './streams';
 export { getWorkspaceIdFromUrl, cleanWorkspaceId } from '../../utils';
+export { updateWorkspaceState, getWorkspaceState } from './workspace';
diff --git a/src/core/server/utils/workspace.test.ts b/src/core/server/utils/workspace.test.ts
new file mode 100644
index 000000000000..49382cfac38f
--- /dev/null
+++ b/src/core/server/utils/workspace.test.ts
@@ -0,0 +1,19 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { httpServerMock } from '../mocks';
+import { getWorkspaceState, updateWorkspaceState } from './workspace';
+
+describe('updateWorkspaceState', () => {
+  it('update with payload', () => {
+    const requestMock = httpServerMock.createOpenSearchDashboardsRequest();
+    updateWorkspaceState(requestMock, {
+      id: 'foo',
+    });
+    expect(getWorkspaceState(requestMock)).toEqual({
+      id: 'foo',
+    });
+  });
+});
diff --git a/src/core/server/utils/workspace.ts b/src/core/server/utils/workspace.ts
new file mode 100644
index 000000000000..c5dcf84b92d9
--- /dev/null
+++ b/src/core/server/utils/workspace.ts
@@ -0,0 +1,43 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+/**
+ * This file is using {@link PluginsStates} to store workspace info into request.
+ * The best practice would be using {@link Server.register} to register plugins into the hapi server
+ * but OSD is wrappering the hapi server and the hapi server instance is hidden as internal implementation.
+ */
+import { PluginsStates } from '@hapi/hapi';
+import { OpenSearchDashboardsRequest, ensureRawRequest } from '../http/router';
+
+/**
+ * This function will be used as a proxy
+ * because `ensureRequest` is only importable from core module.
+ *
+ * @param workspaceId string
+ * @returns void
+ */
+export const updateWorkspaceState = (
+  request: OpenSearchDashboardsRequest,
+  payload: Partial<PluginsStates['workspace']>
+) => {
+  const rawRequest = ensureRawRequest(request);
+
+  if (!rawRequest.plugins) {
+    rawRequest.plugins = {};
+  }
+
+  if (!rawRequest.plugins.workspace) {
+    rawRequest.plugins.workspace = {};
+  }
+
+  rawRequest.plugins.workspace = {
+    ...rawRequest.plugins.workspace,
+    ...payload,
+  };
+};
+
+export const getWorkspaceState = (request: OpenSearchDashboardsRequest) => {
+  return ensureRawRequest(request).plugins?.workspace;
+};
diff --git a/src/legacy/server/osd_server.d.ts b/src/legacy/server/osd_server.d.ts
index 1f8535b59e87..9d94349cf1c0 100644
--- a/src/legacy/server/osd_server.d.ts
+++ b/src/legacy/server/osd_server.d.ts
@@ -60,6 +60,14 @@ declare module 'hapi' {
   }
 }
 
+declare module '@hapi/hapi' {
+  interface PluginsStates {
+    workspace?: {
+      id?: string;
+    };
+  }
+}
+
 type OsdMixinFunc = (osdServer: OsdServer, server: Server, config: any) => Promise<any> | void;
 
 export interface PluginsSetup {