PLAT-1243 - Add ability to pass in datalinks to be mounted to data st…

…udio by name, resourceRef or id
seqeralabs · Jan 28, 2025 · da2a534 · da2a534
1 parent 7d31056
commit da2a534
Show file tree

Hide file tree

Showing 10 changed files with 394 additions and 55 deletions.
diff --git a/src/main/java/io/seqera/tower/cli/commands/data/links/AbstractDataLinkCmd.java b/src/main/java/io/seqera/tower/cli/commands/data/links/AbstractDataLinkCmd.java
@@ -0,0 +1,72 @@
+/*
+ * Copyright 2021-2023, Seqera.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package io.seqera.tower.cli.commands.data.links;
+
+import io.seqera.tower.ApiException;
+import io.seqera.tower.cli.commands.AbstractApiCmd;
+import io.seqera.tower.cli.commands.enums.OutputType;
+import io.seqera.tower.cli.utils.ResponseHelper;
+
+public class AbstractDataLinkCmd extends AbstractApiCmd {
+
+    public boolean checkIfResultIncomplete(Long wspId, String credId, boolean wait) {
+        DataLinksFetchStatus status = checkDataLinksFetchStatus(wspId, credId);
+        if (wait && status == DataLinksFetchStatus.FETCHING) {
+            waitForDoneStatus(wspId, credId);
+        }
+
+        return !wait && status == DataLinksFetchStatus.FETCHING;
+    }
+
+    void waitForDoneStatus(Long wspId, String credId) {
+        try {
+            ResponseHelper.waitStatus(
+                    app().getOut(),
+                    app().output != OutputType.json,
+                    DataLinksFetchStatus.DONE,
+                    DataLinksFetchStatus.values(),
+                    () -> checkDataLinksFetchStatus(wspId, credId),
+                    DataLinksFetchStatus.DONE, DataLinksFetchStatus.ERROR
+            );
+        } catch (InterruptedException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    DataLinksFetchStatus checkDataLinksFetchStatus(Long wspId, String credentialsId) {
+        int status;
+        try {
+            status = api().listDataLinksWithHttpInfo(wspId, credentialsId, null, 1, 0, null).getStatusCode();
+        } catch (ApiException e) {
+            return DataLinksFetchStatus.ERROR;
+        }
+        switch (status) {
+            case 200:
+                return DataLinksFetchStatus.DONE;
+            case 202:
+                return DataLinksFetchStatus.FETCHING;
+            default:
+                return DataLinksFetchStatus.ERROR;
+        }
+    }
+
+    public enum DataLinksFetchStatus {
+        FETCHING, DONE, ERROR
+    }
+
+}
diff --git a/src/main/java/io/seqera/tower/cli/commands/data/links/ListCmd.java b/src/main/java/io/seqera/tower/cli/commands/data/links/ListCmd.java
@@ -18,14 +18,11 @@
 package io.seqera.tower.cli.commands.data.links;
 
 import io.seqera.tower.ApiException;
-import io.seqera.tower.cli.commands.AbstractApiCmd;
-import io.seqera.tower.cli.commands.enums.OutputType;
 import io.seqera.tower.cli.commands.global.PaginationOptions;
 import io.seqera.tower.cli.commands.global.WorkspaceOptionalOptions;
 import io.seqera.tower.cli.responses.Response;
 import io.seqera.tower.cli.responses.data.DataLinksList;
 import io.seqera.tower.cli.utils.PaginationInfo;
-import io.seqera.tower.cli.utils.ResponseHelper;
 import io.seqera.tower.cli.utils.data.DataLinkProvider;
 import io.seqera.tower.model.DataLinksListResponse;
 import picocli.CommandLine;
@@ -38,7 +35,7 @@
         name = "list",
         description = "List data links."
 )
-public class ListCmd extends AbstractApiCmd {
+public class ListCmd extends AbstractDataLinkCmd {
 
     @CommandLine.Mixin
     public WorkspaceOptionalOptions workspace;
@@ -70,34 +67,14 @@ protected Response exec() throws ApiException, IOException {
         String search = buildSearch(searchOption.startsWith, provider, searchOption.region, searchOption.uri);
         String visibility = visibilityOption == null ? null : visibilityOption.toString();
 
-        DataLinksFetchStatus status = checkDataLinksFetchStatus(wspId, credId);
-        if (wait && status == DataLinksFetchStatus.FETCHING) {
-            waitForDoneStatus(wspId, credId);
-        }
-
-        boolean isResultIncomplete = !wait && status == DataLinksFetchStatus.FETCHING;
+        boolean isResultIncomplete = checkIfResultIncomplete(wspId, credId, wait);
 
         DataLinksListResponse data = api().listDataLinks(wspId, credId, search, max, offset, visibility);
         return new DataLinksList(workspaceRef(wspId), data.getDataLinks(),
                 isResultIncomplete,
                 PaginationInfo.from(offset, max, data.getTotalSize()));
     }
 
-    private void waitForDoneStatus(Long wspId, String credId) {
-        try {
-            ResponseHelper.waitStatus(
-                    app().getOut(),
-                    app().output != OutputType.json,
-                    DataLinksFetchStatus.DONE,
-                    DataLinksFetchStatus.values(),
-                    () -> checkDataLinksFetchStatus(wspId, credId),
-                    DataLinksFetchStatus.DONE, DataLinksFetchStatus.ERROR
-            );
-        } catch (InterruptedException e) {
-            throw new RuntimeException(e);
-        }
-    }
-
     private static String formatProviders(List<DataLinkProvider> providers) {
         if (providers == null || providers.isEmpty()) {
             return null;
@@ -144,24 +121,4 @@ enum Visibility {
         hidden, visible, all
     }
 
-    DataLinksFetchStatus checkDataLinksFetchStatus(Long wspId, String credentialsId) {
-        int status = 0;
-        try {
-            status = api().listDataLinksWithHttpInfo(wspId, credentialsId, null, 1, 0, null).getStatusCode();
-        } catch (ApiException e) {
-            return DataLinksFetchStatus.ERROR;
-        }
-        switch (status) {
-            case 200:
-                return DataLinksFetchStatus.DONE;
-            case 202:
-                return DataLinksFetchStatus.FETCHING;
-            default:
-                return DataLinksFetchStatus.ERROR;
-        }
-    }
-
-    public enum DataLinksFetchStatus {
-        FETCHING, DONE, ERROR
-    }
 }
diff --git a/src/main/java/io/seqera/tower/cli/commands/datastudios/AbstractStudiosCmd.java b/src/main/java/io/seqera/tower/cli/commands/datastudios/AbstractStudiosCmd.java
@@ -17,20 +17,29 @@
 
 package io.seqera.tower.cli.commands.datastudios;
 
+import java.util.ArrayList;
 import java.util.List;
 import java.util.Optional;
+import java.util.function.Predicate;
 import java.util.function.Supplier;
+import java.util.stream.Collectors;
 
 import io.seqera.tower.ApiException;
-import io.seqera.tower.cli.commands.AbstractApiCmd;
+import io.seqera.tower.cli.commands.data.links.AbstractDataLinkCmd;
+import io.seqera.tower.cli.exceptions.DataLinkNotFoundException;
 import io.seqera.tower.cli.exceptions.DataStudioNotFoundException;
+import io.seqera.tower.cli.exceptions.MultipleDataLinksFoundException;
+import io.seqera.tower.cli.exceptions.TowerException;
+import io.seqera.tower.cli.exceptions.TowerRuntimeException;
+import io.seqera.tower.model.DataLinkDto;
+import io.seqera.tower.model.DataStudioConfiguration;
 import io.seqera.tower.model.DataStudioDto;
 import io.seqera.tower.model.DataStudioProgressStep;
 
 import static io.seqera.tower.model.DataStudioProgressStepStatus.ERRORED;
 import static io.seqera.tower.model.DataStudioProgressStepStatus.IN_PROGRESS;
 
-public class AbstractStudiosCmd extends AbstractApiCmd {
+public class AbstractStudiosCmd extends AbstractDataLinkCmd {
 
     protected DataStudioDto fetchDataStudio(DataStudioRefOptions dataStudioRefOptions, Long wspId) throws ApiException {
         DataStudioDto dataStudio;
@@ -57,6 +66,76 @@ private DataStudioDto getDataStudioById(Long wspId, String sessionId) throws Api
         return api().describeDataStudio(sessionId, wspId);
     }
 
+    List<String> getMountDataIds(DataStudioConfigurationOptions dataStudioConfigOptions, DataStudioConfiguration currentDataStudioConfiguration, Long wspId) throws ApiException {
+        if (dataStudioConfigOptions.dataLinkRefOptions == null || dataStudioConfigOptions.dataLinkRefOptions.dataLinkRef == null) {
+            return currentDataStudioConfiguration.getMountData();
+        }
+
+        // if DataLink IDs are supplied - use those in request directly
+        if (dataStudioConfigOptions.dataLinkRefOptions.dataLinkRef.mountDataIds != null) {
+            return dataStudioConfigOptions.dataLinkRefOptions.dataLinkRef.mountDataIds;
+        }
+
+        // Check and wait if DataLinks are still being fetched
+        boolean isResultIncomplete = checkIfResultIncomplete(wspId, null, true);
+        if (isResultIncomplete) {
+            throw new TowerException("Failed to fetch datalinks for mountData - please retry.");
+        }
+
+        List<String> dataLinkIds = new ArrayList<>();
+
+        if (dataStudioConfigOptions.dataLinkRefOptions.dataLinkRef.mountDataNames != null) {
+            dataLinkIds = dataStudioConfigOptions.dataLinkRefOptions.dataLinkRef.mountDataNames.stream()
+                    .map(name -> getDataLinkIdByName(wspId, name))
+                    .collect(Collectors.toList());
+        }
+
+        if (dataStudioConfigOptions.dataLinkRefOptions.dataLinkRef.mountDataResourceRefs != null) {
+            dataLinkIds = dataStudioConfigOptions.dataLinkRefOptions.dataLinkRef.mountDataResourceRefs.stream()
+                    .map(resourceRef -> getDataLinkIdByResourceRef(wspId, resourceRef))
+                    .collect(Collectors.toList());
+        }
+
+        return dataLinkIds;
+    }
+
+    private String getDataLinkIdByName(Long wspId, String name) {
+        return getDataLinkIdsBySearchAndFindExactMatch(wspId, name, datalink -> name.equals(datalink.getName()));
+    }
+
+    private String getDataLinkIdByResourceRef(Long wspId, String resourceRef) {
+        return getDataLinkIdsBySearchAndFindExactMatch(wspId, getResourceRefKeywordParam(resourceRef), datalink -> resourceRef.equals(datalink.getResourceRef()));
+    }
+
+    private String getDataLinkIdsBySearchAndFindExactMatch(Long wspId, String search, Predicate<DataLinkDto> filter) {
+        var datalinks = getDataLinksBySearchCriteria(wspId, search).stream()
+                .filter(filter)
+                .collect(Collectors.toList());
+
+        if (datalinks.isEmpty()) {
+            throw new DataLinkNotFoundException(search, wspId);
+        }
+
+        if (datalinks.size() > 1) {
+            var dataLinkIds = datalinks.stream().map(DataLinkDto::getId).collect(Collectors.toList());
+            throw new MultipleDataLinksFoundException(search, wspId, dataLinkIds);
+        }
+
+        return datalinks.get(0).getId();
+    }
+
+    private List<DataLinkDto> getDataLinksBySearchCriteria(Long wspId, String search) {
+        try {
+            return api().listDataLinks(wspId, null, search, null, null, null).getDataLinks();
+        } catch (ApiException e) {
+            throw new TowerRuntimeException("Encountered error while retrieving data links for " + search, e);
+        }
+    }
+
+    private String getResourceRefKeywordParam(String resourceRef) {
+        return String.format("resourceRef:%s", resourceRef);
+    }
+
     public class ProgressStepMessageSupplier implements Supplier<String> {
 
         private final String sessionId;

diff --git a/src/main/java/io/seqera/tower/cli/commands/datastudios/DataLinkRefOptions.java b/src/main/java/io/seqera/tower/cli/commands/datastudios/DataLinkRefOptions.java
@@ -0,0 +1,40 @@
+/*
+ * Copyright 2021-2023, Seqera.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package io.seqera.tower.cli.commands.datastudios;
+
+import java.util.List;
+
+import picocli.CommandLine;
+
+public class DataLinkRefOptions {
+
+    @CommandLine.ArgGroup
+    public DataLinkRef dataLinkRef;
+
+    public static class DataLinkRef {
+        @CommandLine.Option(names = {"--mount-data"}, description = "Optional configuration override for 'mountData' setting (comma separate list of data-link names)", split = ",")
+        public List<String> mountDataNames;
+
+        @CommandLine.Option(names = {"--mount-data-ids"}, description = "Optional configuration override for 'mountData' setting (comma separate list of data-link Ids)", split = ",")
+        public List<String> mountDataIds;
+
+        @CommandLine.Option(names = {"--mount-data-resource-refs"}, description = "Optional configuration override for 'mountData' setting (comma separate list of data-link resource refs)", split = ",")
+        public List<String> mountDataResourceRefs;
+    }
+
+}
diff --git a/src/main/java/io/seqera/tower/cli/commands/datastudios/DataStudioConfigurationOptions.java b/src/main/java/io/seqera/tower/cli/commands/datastudios/DataStudioConfigurationOptions.java
@@ -17,8 +17,6 @@
 
 package io.seqera.tower.cli.commands.datastudios;
 
-import java.util.List;
-
 import picocli.CommandLine;
 
 public class DataStudioConfigurationOptions {
@@ -32,7 +30,7 @@ public class DataStudioConfigurationOptions {
     @CommandLine.Option(names = {"-m", "--memory"}, description = "Optional configuration override for 'memory' setting (Integer representing memory in MBs)")
     public Integer memory;
 
-    @CommandLine.Option(names = {"--mount-data"}, description = "Optional configuration override for 'mountData' setting (comma separate list of datalinkIds)", split = ",")
-    public List<String> mountData;
+    @CommandLine.Mixin
+    public DataLinkRefOptions dataLinkRefOptions;
 
 }
diff --git a/src/main/java/io/seqera/tower/cli/commands/datastudios/StartCmd.java b/src/main/java/io/seqera/tower/cli/commands/datastudios/StartCmd.java
@@ -117,7 +117,7 @@ private DataStudioStatus checkDataStudioStatus(String sessionId, Long workspaceI
         }
     }
 
-    private DataStudioStartRequest getStartRequestWithOverridesApplied(DataStudioDto dataStudioDto) {
+    private DataStudioStartRequest getStartRequestWithOverridesApplied(DataStudioDto dataStudioDto) throws ApiException {
         DataStudioConfiguration dataStudioConfiguration = dataStudioDto.getConfiguration() == null
                 ? new DataStudioConfiguration()
                 : dataStudioDto.getConfiguration();
@@ -131,9 +131,8 @@ private DataStudioStartRequest getStartRequestWithOverridesApplied(DataStudioDto
         dataStudioConfiguration.setMemory(dataStudioConfigOptions.memory == null
                 ? dataStudioConfiguration.getMemory()
                 : dataStudioConfigOptions.memory);
-        dataStudioConfiguration.setMountData(dataStudioConfigOptions.mountData == null
-                ? dataStudioConfiguration.getMountData()
-                : dataStudioConfigOptions.mountData);
+
+        dataStudioConfiguration.setMountData(getMountDataIds(dataStudioConfigOptions, dataStudioConfiguration, dataStudioDto.getWorkspaceId()));
 
         String appliedDescription = description == null
                 ? dataStudioDto.getDescription()

diff --git a/src/main/java/io/seqera/tower/cli/exceptions/DataLinkNotFoundException.java b/src/main/java/io/seqera/tower/cli/exceptions/DataLinkNotFoundException.java
@@ -0,0 +1,26 @@
+/*
+ * Copyright 2021-2023, Seqera.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package io.seqera.tower.cli.exceptions;
+
+public class DataLinkNotFoundException extends TowerRuntimeException {
+
+    public DataLinkNotFoundException(String dataLinkIdentifier, Long workspaceId) {
+        super(String.format("DataLink '%s' not found at workspace '%s'", dataLinkIdentifier, workspaceId));
+    }
+
+}