[RW-12277][risk=low] Update localizeRuntime and localizeApp API to su…

…pport localized all files (#8472)

* localize all files

* add test for Cromwell app

* WIP: Some UI Change for testing

* Multiple apps can be in provision state also add filename when clicking edit

* Fix unit test

* Clean up and add some comments

---------

Co-authored-by: NehaBroad <[email protected]>

api/src/main/java/org/pmiops/workbench/api/AppsController.java

@@ -97,8 +97,12 @@ public ResponseEntity<ListAppsResponse> listAppsInWorkspace(String workspaceName
     return ResponseEntity.ok(response);
   }
 
+  @Override
   public ResponseEntity<AppLocalizeResponse> localizeApp(
-      String workspaceNamespace, String appName, AppLocalizeRequest body) {
+      String workspaceNamespace,
+      String appName,
+      Boolean localizeAllFiles,
+      AppLocalizeRequest body) {
     DbUser user = userProvider.get();
     leonardoApiHelper.enforceComputeSecuritySuspension(user);
 
@@ -111,6 +115,7 @@ public ResponseEntity<AppLocalizeResponse> localizeApp(
                     body.getAppType(),
                     body.getFileNames(),
                     body.isPlaygroundMode(),
-                    false)));
+                    false,
+                    localizeAllFiles)));
   }
 }

api/src/main/java/org/pmiops/workbench/api/RuntimeController.java

@@ -274,7 +274,7 @@ public ResponseEntity<EmptyResponse> deleteRuntime(
 
   @Override
   public ResponseEntity<RuntimeLocalizeResponse> localize(
-      String workspaceNamespace, RuntimeLocalizeRequest body) {
+      String workspaceNamespace, Boolean localizeAllFiles, RuntimeLocalizeRequest body) {
     DbUser user = userProvider.get();
     leonardoApiHelper.enforceComputeSecuritySuspension(user);
     DbWorkspace dbWorkspace = workspaceService.lookupWorkspaceByNamespace(workspaceNamespace);
@@ -295,6 +295,7 @@ public ResponseEntity<RuntimeLocalizeResponse> localize(
                     appType,
                     body.getNotebookNames(),
                     body.isPlaygroundMode(),
-                    true)));
+                    true,
+                    localizeAllFiles)));
   }
 }

api/src/main/java/org/pmiops/workbench/interactiveanalysis/InteractiveAnalysisService.java

@@ -23,6 +23,8 @@
 import org.pmiops.workbench.firecloud.FireCloudService;
 import org.pmiops.workbench.leonardo.LeonardoApiClient;
 import org.pmiops.workbench.model.AppType;
+import org.pmiops.workbench.model.FileDetail;
+import org.pmiops.workbench.notebooks.NotebooksService;
 import org.pmiops.workbench.notebooks.model.StorageLink;
 import org.pmiops.workbench.rawls.model.RawlsWorkspaceDetails;
 import org.pmiops.workbench.workspaces.WorkspaceService;
@@ -43,6 +45,7 @@ public class InteractiveAnalysisService {
   private final UserRecentResourceService userRecentResourceService;
 
   private final LeonardoApiClient leonardoNotebooksClient;
+  private final NotebooksService notebooksService;
 
   private static final String AOU_CONFIG_FILENAME = ".all_of_us_config.json";
   private static final String WORKSPACE_NAMESPACE_KEY = "WORKSPACE_NAMESPACE";
@@ -69,13 +72,15 @@ public InteractiveAnalysisService(
       Provider<WorkbenchConfig> workbenchConfigProvider,
       Provider<DbUser> userProvider,
       UserRecentResourceService userRecentResourceService,
-      LeonardoApiClient leonardoNotebooksClient) {
+      LeonardoApiClient leonardoNotebooksClient,
+      NotebooksService notebooksService) {
     this.workspaceService = workspaceService;
     this.fireCloudService = fireCloudService;
     this.workbenchConfigProvider = workbenchConfigProvider;
     this.userProvider = userProvider;
     this.userRecentResourceService = userRecentResourceService;
     this.leonardoNotebooksClient = leonardoNotebooksClient;
+    this.notebooksService = notebooksService;
   }
 
   public String localize(
@@ -84,7 +89,8 @@ public String localize(
       @Nullable AppType appType,
       List<String> fileNames,
       boolean isPlayground,
-      boolean isGceRuntime) {
+      boolean isGceRuntime,
+      boolean localizeAllFiles) {
     DbWorkspace dbWorkspace = workspaceService.lookupWorkspaceByNamespace(workspaceNamespace);
     final RawlsWorkspaceDetails firecloudWorkspace;
     try {
@@ -161,12 +167,35 @@ public String localize(
 
     localizeMap.put(aouConfigEditDir, aouConfigUri);
 
-    // Localize the requested notebooks, if any.
-    localizeMap.putAll(
-        fileNames.stream()
-            .collect(
-                Collectors.toMap(
-                    name -> localizeTargetDir + name, name -> gcsNotebooksDir + "/" + name)));
+    // Localize all files if localizeAllFiles is true, otherwise, localize the requested notebooks.
+    if (localizeAllFiles) {
+      List<FileDetail> notebooks;
+      if (isGceRuntime) {
+        notebooks =
+            notebooksService.getAllJupyterNotebooks(
+                firecloudWorkspace.getBucketName(),
+                workspaceNamespace,
+                firecloudWorkspace.getName());
+      } else {
+        notebooks =
+            notebooksService.getAllNotebooksByAppType(
+                firecloudWorkspace.getBucketName(),
+                workspaceNamespace,
+                firecloudWorkspace.getName(),
+                appType);
+      }
+      localizeMap.putAll(
+          notebooks.stream()
+              .collect(
+                  Collectors.toMap(
+                      notebook -> localizeTargetDir + notebook.getName(), FileDetail::getPath)));
+    } else {
+      localizeMap.putAll(
+          fileNames.stream()
+              .collect(
+                  Collectors.toMap(
+                      name -> localizeTargetDir + name, name -> gcsNotebooksDir + "/" + name)));
+    }
 
     if (isGceRuntime) {
       localizeMap.put(playgroundDir + "/" + AOU_CONFIG_FILENAME, aouConfigUri);

api/src/main/java/org/pmiops/workbench/notebooks/NotebooksService.java

@@ -3,6 +3,7 @@
 import com.google.cloud.storage.Blob;
 import java.util.List;
 import org.json.JSONObject;
+import org.pmiops.workbench.model.AppType;
 import org.pmiops.workbench.model.FileDetail;
 import org.pmiops.workbench.model.KernelTypeEnum;
 
@@ -18,6 +19,12 @@ public interface NotebooksService {
   List<FileDetail> getNotebooksAsService(
       String bucketName, String workspaceNamespace, String workspaceName);
 
+  List<FileDetail> getAllNotebooksByAppType(
+      String bucketName, String workspaceNamespace, String workspaceName, AppType appType);
+
+  List<FileDetail> getAllJupyterNotebooks(
+      String bucketName, String workspaceNamespace, String workspaceName);
+
   /**
    * Is this a notebook file which is managed (localized and delocalized) by the Workbench?
    *

api/src/main/java/org/pmiops/workbench/notebooks/NotebooksServiceImpl.java

@@ -5,6 +5,7 @@
 import java.nio.charset.StandardCharsets;
 import java.sql.Timestamp;
 import java.time.Clock;
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
 import java.util.Optional;
@@ -28,6 +29,7 @@
 import org.pmiops.workbench.firecloud.FireCloudService;
 import org.pmiops.workbench.google.CloudStorageClient;
 import org.pmiops.workbench.google.GoogleCloudLocators;
+import org.pmiops.workbench.model.AppType;
 import org.pmiops.workbench.model.FileDetail;
 import org.pmiops.workbench.model.KernelTypeEnum;
 import org.pmiops.workbench.model.WorkspaceAccessLevel;
@@ -124,6 +126,32 @@ public List<FileDetail> getNotebooksAsService(
         .collect(Collectors.toList());
   }
 
+  @Override
+  public List<FileDetail> getAllNotebooksByAppType(
+      String bucketName, String workspaceNamespace, String workspaceName, AppType appType) {
+    List<FileDetail> allNotebooks =
+        getNotebooksAsService(bucketName, workspaceNamespace, workspaceName);
+    if (appType.equals(AppType.RSTUDIO)) {
+      return allNotebooks.stream()
+          .filter(fileDetail -> NotebookUtils.isRStudioFile(fileDetail.getName()))
+          .collect(Collectors.toList());
+    } else if (appType.equals(AppType.SAS)) {
+      return allNotebooks.stream()
+          .filter(fileDetail -> NotebookUtils.isSasFile(fileDetail.getName()))
+          .collect(Collectors.toList());
+    } else {
+      return new ArrayList<>();
+    }
+  }
+
+  @Override
+  public List<FileDetail> getAllJupyterNotebooks(
+      String bucketName, String workspaceNamespace, String workspaceName) {
+    return getNotebooksAsService(bucketName, workspaceNamespace, workspaceName).stream()
+        .filter(fileDetail -> NotebookUtils.isJupyterNotebook(fileDetail.getName()))
+        .collect(Collectors.toList());
+  }
+
   @Override
   public boolean isManagedNotebookBlob(Blob blob) {
     // Blobs have notebooks/ directory

api/src/main/resources/workbench-api.yaml

@@ -832,6 +832,15 @@ paths:
         required: true
         schema:
           type: string
+      - in: query
+        name: localizeAllFiles
+        description: |
+          Optional query that will localize all files associate with this app. 
+          If true, it will ignore fileNames in body and localize all files instead.
+        required: false
+        schema:
+          type: boolean
+          default: false
       requestBody:
         description: Localization request.
         content:
@@ -1102,6 +1111,15 @@ paths:
         required: true
         schema:
           type: string
+      - in: query
+        name: localizeAllFiles
+        description: |
+          Optional query that will localize all files associate with this app. 
+          If true, it will ignore fileNames in body and localize all files instead.
+        required: false
+        schema:
+          type: boolean
+          default: false
       requestBody:
         description: Localization request.
         content:

api/src/test/java/org/pmiops/workbench/api/RuntimeControllerTest.java

@@ -1321,7 +1321,7 @@ public void testLocalize_securitySuspended() {
 
     assertThrows(
         FailedPreconditionException.class,
-        () -> runtimeController.localize(WORKSPACE_NS, new RuntimeLocalizeRequest()));
+        () -> runtimeController.localize(WORKSPACE_NS, false, new RuntimeLocalizeRequest()));
   }
 
   @Test
@@ -1331,7 +1331,8 @@ public void localize_validateActiveBilling() {
         .validateActiveBilling(WORKSPACE_NS, WORKSPACE_ID);
 
     RuntimeLocalizeRequest req = new RuntimeLocalizeRequest();
-    assertThrows(ForbiddenException.class, () -> runtimeController.localize(WORKSPACE_NS, req));
+    assertThrows(
+        ForbiddenException.class, () -> runtimeController.localize(WORKSPACE_NS, false, req));
   }
 
   @Test
@@ -1342,7 +1343,8 @@ public void localize_validateActiveBilling_checkAccessFirst() {
 
     RuntimeLocalizeRequest req = new RuntimeLocalizeRequest();
 
-    assertThrows(ForbiddenException.class, () -> runtimeController.localize(WORKSPACE_NS, req));
+    assertThrows(
+        ForbiddenException.class, () -> runtimeController.localize(WORKSPACE_NS, false, req));
     verify(mockWorkspaceAuthService, never()).validateActiveBilling(anyString(), anyString());
   }
 

api/src/test/java/org/pmiops/workbench/interactiveanalysis/InteractiveAnalysisServiceTest.java

@@ -8,6 +8,7 @@
 import static org.pmiops.workbench.interactiveanalysis.InteractiveAnalysisService.SAS_DELOC_PATTERN;
 import static org.pmiops.workbench.interactiveanalysis.InteractiveAnalysisService.aouConfigDataUri;
 
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
@@ -22,6 +23,7 @@
 import org.pmiops.workbench.firecloud.FireCloudService;
 import org.pmiops.workbench.leonardo.LeonardoApiClient;
 import org.pmiops.workbench.model.AppType;
+import org.pmiops.workbench.notebooks.NotebooksService;
 import org.pmiops.workbench.notebooks.model.StorageLink;
 import org.pmiops.workbench.rawls.model.RawlsWorkspaceDetails;
 import org.pmiops.workbench.rawls.model.RawlsWorkspaceResponse;
@@ -74,6 +76,7 @@ DbUser user() {
   @MockBean WorkspaceService mockWorkspaceService;
   @MockBean LeonardoApiClient mockLeonardoApiClient;
   @MockBean FireCloudService mockFireCloudService;
+  @MockBean NotebooksService mockNotebooksService;
 
   @Autowired InteractiveAnalysisService interactiveAnalysisService;
 
@@ -96,7 +99,7 @@ public void setUp() {
             .setFirecloudName(FIRECLOUD_WS_NAME);
     when(mockWorkspaceService.lookupWorkspaceByNamespace(WORKSPACE_NS)).thenReturn(dbWorkspace);
     RawlsWorkspaceDetails rawlsWorkspaceDetails =
-        new RawlsWorkspaceDetails().bucketName(BUCKET_NAME);
+        new RawlsWorkspaceDetails().bucketName(BUCKET_NAME).name(FIRECLOUD_WS_NAME);
     when(mockFireCloudService.getWorkspace(WORKSPACE_NS, FIRECLOUD_WS_NAME))
         .thenReturn(new RawlsWorkspaceResponse().workspace(rawlsWorkspaceDetails));
 
@@ -122,7 +125,7 @@ public void testLocalize_jupyter_editMode() {
 
     AppType appType = null; // Jupyter uses GCE, so it doesn't have a GKE App Type
     interactiveAnalysisService.localize(
-        WORKSPACE_NS, APP_NAME, appType, notebookLists, false, true);
+        WORKSPACE_NS, APP_NAME, appType, notebookLists, false, true, false);
     verify(mockLeonardoApiClient)
         .createStorageLinkForRuntime(GOOGLE_PROJECT_ID, APP_NAME, expectedStorageLink);
     verify(mockLeonardoApiClient)
@@ -146,7 +149,8 @@ public void testLocalize_jupyter_playground() {
     expectedLocalizeMap.put(playgroundDir + "/foo.ipynb", NOTEBOOK_DIR + "/foo.ipynb");
 
     AppType appType = null; // Jupyter uses GCE, so it doesn't have a GKE App Type
-    interactiveAnalysisService.localize(WORKSPACE_NS, APP_NAME, appType, notebookLists, true, true);
+    interactiveAnalysisService.localize(
+        WORKSPACE_NS, APP_NAME, appType, notebookLists, true, true, false);
     verify(mockLeonardoApiClient)
         .createStorageLinkForRuntime(GOOGLE_PROJECT_ID, APP_NAME, expectedStorageLink);
     verify(mockLeonardoApiClient)
@@ -167,12 +171,36 @@ public void testLocalize_RStudio_editMode() {
     expectedLocalizeMap.put("foo.Rmd", NOTEBOOK_DIR + "/foo.Rmd");
 
     interactiveAnalysisService.localize(
-        WORKSPACE_NS, APP_NAME, AppType.RSTUDIO, notebookLists, false, false);
+        WORKSPACE_NS, APP_NAME, AppType.RSTUDIO, notebookLists, false, false, false);
     verify(mockLeonardoApiClient)
         .createStorageLinkForApp(GOOGLE_PROJECT_ID, APP_NAME, expectedStorageLink);
     verify(mockLeonardoApiClient).localizeForApp(GOOGLE_PROJECT_ID, APP_NAME, expectedLocalizeMap);
   }
 
+  @Test
+  public void testLocalize_allFiles_rstudio() {
+    interactiveAnalysisService.localize(
+        WORKSPACE_NS, APP_NAME, AppType.RSTUDIO, new ArrayList<>(), false, false, true);
+    verify(mockNotebooksService)
+        .getAllNotebooksByAppType(BUCKET_NAME, WORKSPACE_NS, FIRECLOUD_WS_NAME, AppType.RSTUDIO);
+  }
+
+  @Test
+  public void testLocalize_allFiles_sas() {
+    interactiveAnalysisService.localize(
+        WORKSPACE_NS, APP_NAME, AppType.SAS, new ArrayList<>(), false, false, true);
+    verify(mockNotebooksService)
+        .getAllNotebooksByAppType(BUCKET_NAME, WORKSPACE_NS, FIRECLOUD_WS_NAME, AppType.SAS);
+  }
+
+  @Test
+  public void testLocalize_allFiles_jupyter() {
+    interactiveAnalysisService.localize(
+        WORKSPACE_NS, APP_NAME, null, new ArrayList<>(), false, true, true);
+    verify(mockNotebooksService)
+        .getAllJupyterNotebooks(BUCKET_NAME, WORKSPACE_NS, FIRECLOUD_WS_NAME);
+  }
+
   @Test
   public void testLocalize_SAS_editMode() {
     String editDir = "";
@@ -187,7 +215,7 @@ public void testLocalize_SAS_editMode() {
     expectedLocalizeMap.put("foo.sas", NOTEBOOK_DIR + "/foo.sas");
 
     interactiveAnalysisService.localize(
-        WORKSPACE_NS, APP_NAME, AppType.SAS, notebookLists, false, false);
+        WORKSPACE_NS, APP_NAME, AppType.SAS, notebookLists, false, false, false);
     verify(mockLeonardoApiClient)
         .createStorageLinkForApp(GOOGLE_PROJECT_ID, APP_NAME, expectedStorageLink);
     verify(mockLeonardoApiClient).localizeForApp(GOOGLE_PROJECT_ID, APP_NAME, expectedLocalizeMap);
@@ -200,6 +228,12 @@ public void testLocalize_appType_not_supported() {
         NotImplementedException.class,
         () ->
             interactiveAnalysisService.localize(
-                WORKSPACE_NS, APP_NAME, unsupportedAppType, Collections.emptyList(), false, false));
+                WORKSPACE_NS,
+                APP_NAME,
+                unsupportedAppType,
+                Collections.emptyList(),
+                false,
+                false,
+                false));
   }
 }