From 6221f6abe6238312b87866da34a82dde78b29db1 Mon Sep 17 00:00:00 2001
From: Juntao Wang <juntwang@redhat.com>
Date: Tue, 24 Sep 2024 13:38:05 -0400
Subject: [PATCH] Restart all connected notebooks when the pvc size changes

---
 .../src/api/k8s/__tests__/notebooks.spec.ts   | 28 +++++++++++++++++++
 frontend/src/api/k8s/notebooks.ts             | 25 +++++++++++++++++
 .../detail/storage/ManageStorageModal.tsx     |  7 ++++-
 3 files changed, 59 insertions(+), 1 deletion(-)

diff --git a/frontend/src/api/k8s/__tests__/notebooks.spec.ts b/frontend/src/api/k8s/__tests__/notebooks.spec.ts
index 0e7d2cc6db..b989ef17db 100644
--- a/frontend/src/api/k8s/__tests__/notebooks.spec.ts
+++ b/frontend/src/api/k8s/__tests__/notebooks.spec.ts
@@ -30,6 +30,7 @@ import {
   getStopPatch,
   startPatch,
   mergePatchUpdateNotebook,
+  restartNotebook,
 } from '~/api/k8s/notebooks';
 
 import {
@@ -943,3 +944,30 @@ describe('removeNotebookSecret', () => {
     });
   });
 });
+
+describe('restartNotebook', () => {
+  it('should add restart notebook annotation', async () => {
+    const name = 'test-notebook';
+    const namespace = 'test-project';
+    const notebookMock = mockNotebookK8sResource({ uid });
+
+    k8sGetResourceMock.mockResolvedValue(notebookMock);
+    k8sPatchResourceMock.mockResolvedValue(notebookMock);
+
+    const renderResult = await restartNotebook(name, namespace);
+    expect(k8sPatchResourceMock).toHaveBeenCalledWith({
+      fetchOptions: { requestInit: {} },
+      model: NotebookModel,
+      patches: [
+        {
+          op: 'add',
+          path: '/metadata/annotations/notebooks.opendatahub.io~1notebook-restart',
+          value: 'true',
+        },
+      ],
+      queryOptions: { name, ns: namespace, queryParams: {} },
+    });
+    expect(k8sPatchResourceMock).toHaveBeenCalledTimes(1);
+    expect(renderResult).toStrictEqual(notebookMock);
+  });
+});
diff --git a/frontend/src/api/k8s/notebooks.ts b/frontend/src/api/k8s/notebooks.ts
index 6a8e29232f..0be437d25c 100644
--- a/frontend/src/api/k8s/notebooks.ts
+++ b/frontend/src/api/k8s/notebooks.ts
@@ -506,3 +506,28 @@ export const removeNotebookSecret = (
       })
       .catch(reject);
   });
+
+export const restartNotebook = (
+  notebookName: string,
+  namespace: string,
+  opts?: K8sAPIOptions,
+): Promise<NotebookKind> => {
+  const patches: Patch[] = [
+    {
+      op: 'add',
+      path: '/metadata/annotations/notebooks.opendatahub.io~1notebook-restart',
+      value: 'true',
+    },
+  ];
+
+  return k8sPatchResource<NotebookKind>(
+    applyK8sAPIOptions(
+      {
+        model: NotebookModel,
+        queryOptions: { name: notebookName, ns: namespace },
+        patches,
+      },
+      opts,
+    ),
+  );
+};
diff --git a/frontend/src/pages/projects/screens/detail/storage/ManageStorageModal.tsx b/frontend/src/pages/projects/screens/detail/storage/ManageStorageModal.tsx
index ee5714eb01..edd5474d39 100644
--- a/frontend/src/pages/projects/screens/detail/storage/ManageStorageModal.tsx
+++ b/frontend/src/pages/projects/screens/detail/storage/ManageStorageModal.tsx
@@ -1,6 +1,6 @@
 import * as React from 'react';
 import { Form, Modal, Stack, StackItem } from '@patternfly/react-core';
-import { attachNotebookPVC, createPvc, removeNotebookPVC, updatePvc } from '~/api';
+import { attachNotebookPVC, createPvc, removeNotebookPVC, restartNotebook, updatePvc } from '~/api';
 import { NotebookKind, PersistentVolumeClaimKind } from '~/k8sTypes';
 import { ProjectDetailsContext } from '~/pages/projects/ProjectDetailsContext';
 import { useCreateStorageObjectForNotebook } from '~/pages/projects/screens/spawner/storage/utils';
@@ -98,6 +98,11 @@ const ManageStorageModal: React.FC<AddStorageModalProps> = ({ existingData, isOp
       ) {
         pvcPromises.push(updatePvc(createData, existingData, namespace, { dryRun }));
       }
+      if (existingData.spec.resources.requests.storage !== createData.size) {
+        connectedNotebooks.map((connectedNotebook) =>
+          pvcPromises.push(restartNotebook(connectedNotebook.metadata.name, namespace, { dryRun })),
+        );
+      }
       if (removedNotebooks.length > 0) {
         // Remove connected pvcs
         pvcPromises.push(