feat: implement copy shared access operation (#283)(#271)

epam · Mar 14, 2024 · 52ba2a7 · 52ba2a7
1 parent 6a434e2
commit 52ba2a7
Show file tree

Hide file tree

Showing 9 changed files with 438 additions and 11 deletions.
diff --git a/src/main/java/com/epam/aidial/core/AiDial.java b/src/main/java/com/epam/aidial/core/AiDial.java
@@ -113,7 +113,7 @@ void start() throws Exception {
             LockService lockService = new LockService(redis, storage.getPrefix());
             resourceService = new ResourceService(vertx, redis, storage, lockService, settings("resources"), storage.getPrefix());
             InvitationService invitationService = new InvitationService(resourceService, encryptionService, settings("invitations"));
-            ShareService shareService = new ShareService(resourceService, invitationService, encryptionService);
+            ShareService shareService = new ShareService(resourceService, invitationService, encryptionService, storage);
             PublicationService publicationService = new PublicationService(encryptionService, resourceService, storage, generator, clock);
             ResourceOperationService resourceOperationService = new ResourceOperationService(resourceService, storage, invitationService, shareService);
 

diff --git a/src/main/java/com/epam/aidial/core/controller/ControllerSelector.java b/src/main/java/com/epam/aidial/core/controller/ControllerSelector.java
@@ -45,7 +45,7 @@ public class ControllerSelector {
     private static final Pattern PATTERN_TOKENIZE = Pattern.compile("^/+v1/deployments/([^/]+)/tokenize$");
     private static final Pattern PATTERN_TRUNCATE_PROMPT = Pattern.compile("^/+v1/deployments/([^/]+)/truncate_prompt$");
 
-    private static final Pattern SHARE_RESOURCE_OPERATIONS = Pattern.compile("^/v1/ops/resource/share/(create|list|discard|revoke)$");
+    private static final Pattern SHARE_RESOURCE_OPERATIONS = Pattern.compile("^/v1/ops/resource/share/(create|list|discard|revoke|copy)$");
     private static final Pattern INVITATIONS = Pattern.compile("^/v1/invitations$");
     private static final Pattern INVITATION = Pattern.compile("^/v1/invitations/([a-zA-Z0-9]+)$");
     private static final Pattern PUBLICATIONS = Pattern.compile("^/v1/ops/publications/(list|get|create|delete|approve|reject)$");

diff --git a/src/main/java/com/epam/aidial/core/controller/ShareController.java b/src/main/java/com/epam/aidial/core/controller/ShareController.java
@@ -2,17 +2,22 @@
 
 import com.epam.aidial.core.Proxy;
 import com.epam.aidial.core.ProxyContext;
+import com.epam.aidial.core.data.CopySharedAccessRequest;
 import com.epam.aidial.core.data.ListSharedResourcesRequest;
 import com.epam.aidial.core.data.ResourceLinkCollection;
 import com.epam.aidial.core.data.ShareResourcesRequest;
 import com.epam.aidial.core.security.EncryptionService;
 import com.epam.aidial.core.service.InvitationService;
 import com.epam.aidial.core.service.LockService;
+import com.epam.aidial.core.service.ResourceService;
 import com.epam.aidial.core.service.ShareService;
+import com.epam.aidial.core.storage.BlobStorage;
 import com.epam.aidial.core.storage.BlobStorageUtil;
+import com.epam.aidial.core.storage.ResourceDescription;
 import com.epam.aidial.core.util.HttpException;
 import com.epam.aidial.core.util.HttpStatus;
 import com.epam.aidial.core.util.ProxyUtil;
+import com.epam.aidial.core.util.ResourceUtil;
 import io.vertx.core.Future;
 import io.vertx.core.buffer.Buffer;
 import lombok.extern.slf4j.Slf4j;
@@ -30,6 +35,8 @@ public class ShareController {
     private final EncryptionService encryptionService;
     private final LockService lockService;
     private final InvitationService invitationService;
+    private final ResourceService resourceService;
+    private final BlobStorage storage;
 
     public ShareController(Proxy proxy, ProxyContext context) {
         this.proxy = proxy;
@@ -38,6 +45,8 @@ public ShareController(Proxy proxy, ProxyContext context) {
         this.encryptionService = proxy.getEncryptionService();
         this.lockService = proxy.getLockService();
         this.invitationService = proxy.getInvitationService();
+        this.resourceService = proxy.getResourceService();
+        this.storage = proxy.getStorage();
     }
 
     public Future<?> handle(Operation operation) {
@@ -46,6 +55,7 @@ public Future<?> handle(Operation operation) {
             case CREATE -> createSharedResources();
             case REVOKE -> revokeSharedResources();
             case DISCARD -> discardSharedResources();
+            case COPY -> copySharedAccess();
             default ->
                     context.respond(HttpStatus.INTERNAL_SERVER_ERROR, "Operation %s is not supported".formatted(operation));
         }
@@ -137,6 +147,49 @@ public Future<?> revokeSharedResources() {
                 .onFailure(this::handleServiceError);
     }
 
+    public Future<?> copySharedAccess() {
+        return context.getRequest()
+                .body()
+                .compose(buffer -> {
+                    CopySharedAccessRequest request;
+                    try {
+                        request = ProxyUtil.convertToObject(buffer, CopySharedAccessRequest.class);
+                    } catch (Exception e) {
+                        log.error("Invalid request body provided", e);
+                        throw new IllegalArgumentException("Can't initiate copy shared access request. Incorrect body provided");
+                    }
+
+                    String sourceUrl = request.sourceUrl();
+                    if (sourceUrl == null) {
+                        throw new IllegalArgumentException("sourceUrl must be provided");
+                    }
+                    String destinationUrl = request.destinationUrl();
+                    if (destinationUrl == null) {
+                        throw new IllegalArgumentException("destinationUrl must be provided");
+                    }
+
+                    String bucketLocation = BlobStorageUtil.buildInitiatorBucket(context);
+                    String bucket = encryptionService.encrypt(bucketLocation);
+
+                    ResourceDescription source = ResourceDescription.fromPrivateUrl(sourceUrl, encryptionService);
+                    if (!bucket.equals(source.getBucketName())) {
+                        throw new IllegalArgumentException("sourceUrl does not belong to the user");
+                    }
+                    ResourceDescription destination = ResourceDescription.fromPrivateUrl(destinationUrl, encryptionService);
+                    if (!bucket.equals(destination.getBucketName())) {
+                        throw new IllegalArgumentException("destinationUrl does not belong to the user");
+                    }
+
+                    return proxy.getVertx().executeBlocking(() ->
+                            lockService.underBucketLock(bucketLocation, () -> {
+                                shareService.copySharedAccess(bucket, bucketLocation, source, destination);
+                                return null;
+                            }));
+                })
+                .onSuccess(ignore -> context.respond(HttpStatus.OK))
+                .onFailure(this::handleServiceError);
+    }
+
     private void handleServiceError(Throwable error) {
         if (error instanceof IllegalArgumentException) {
             context.respond(HttpStatus.BAD_REQUEST, error.getMessage());
@@ -157,7 +210,11 @@ private ResourceLinkCollection getResourceLinkCollection(Buffer buffer, Operatio
         }
     }
 
+    private boolean hasResource(ResourceDescription resource) {
+        return ResourceUtil.hasResource(resource, resourceService, storage);
+    }
+
     public enum Operation {
-        CREATE, LIST, DISCARD, REVOKE
+        CREATE, LIST, DISCARD, REVOKE, COPY
     }
 }
diff --git a/src/main/java/com/epam/aidial/core/data/CopySharedAccessRequest.java b/src/main/java/com/epam/aidial/core/data/CopySharedAccessRequest.java
@@ -0,0 +1,4 @@
+package com.epam.aidial.core.data;
+
+public record CopySharedAccessRequest(String sourceUrl, String destinationUrl) {
+}
diff --git a/src/main/java/com/epam/aidial/core/service/ResourceOperationService.java b/src/main/java/com/epam/aidial/core/service/ResourceOperationService.java
@@ -3,6 +3,7 @@
 import com.epam.aidial.core.data.ResourceType;
 import com.epam.aidial.core.storage.BlobStorage;
 import com.epam.aidial.core.storage.ResourceDescription;
+import com.epam.aidial.core.util.ResourceUtil;
 import lombok.AllArgsConstructor;
 
 @AllArgsConstructor
@@ -35,30 +36,34 @@ public void moveResource(String bucket, String location, ResourceDescription sou
                             .formatted(sourceResourceUrl, destinationResourceUrl));
                 }
                 storage.copy(sourceResourcePath, destinationResourcePath);
-                storage.delete(sourceResourcePath);
             }
             case CONVERSATION, PROMPT -> {
                 boolean copied = resourceService.copyResource(source, destination, overwriteIfExists);
                 if (!copied) {
                     throw new IllegalArgumentException("Can't move resource %s to %s, because destination resource already exists"
                             .formatted(sourceResourceUrl, destinationResourceUrl));
                 }
-                resourceService.deleteResource(source);
             }
             default -> throw new IllegalArgumentException("Unsupported resource type " + resourceType);
         }
         // move source links to destination if any
         invitationService.moveResource(bucket, location, source, destination);
         // move shared access if any
         shareService.moveSharedAccess(bucket, location, source, destination);
+
+        deleteResource(source);
     }
 
     private boolean hasResource(ResourceDescription resource) {
-        return switch (resource.getType()) {
-            case FILE -> storage.exists(resource.getAbsoluteFilePath());
-            case CONVERSATION, PROMPT -> resourceService.hasResource(resource);
+        return ResourceUtil.hasResource(resource, resourceService, storage);
+    }
+
+    private void deleteResource(ResourceDescription resource) {
+        switch (resource.getType()) {
+            case FILE -> storage.delete(resource.getAbsoluteFilePath());
+            case CONVERSATION, PROMPT -> resourceService.deleteResource(resource);
             default -> throw new IllegalArgumentException("Unsupported resource type " + resource.getType());
-        };
+        }
     }
 
 }
diff --git a/src/main/java/com/epam/aidial/core/service/ShareService.java b/src/main/java/com/epam/aidial/core/service/ShareService.java
@@ -13,9 +13,11 @@
 import com.epam.aidial.core.data.SharedByMeDto;
 import com.epam.aidial.core.data.SharedResourcesResponse;
 import com.epam.aidial.core.security.EncryptionService;
+import com.epam.aidial.core.storage.BlobStorage;
 import com.epam.aidial.core.storage.BlobStorageUtil;
 import com.epam.aidial.core.storage.ResourceDescription;
 import com.epam.aidial.core.util.ProxyUtil;
+import com.epam.aidial.core.util.ResourceUtil;
 import lombok.AllArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
 
@@ -36,6 +38,7 @@ public class ShareService {
     private final ResourceService resourceService;
     private final InvitationService invitationService;
     private final EncryptionService encryptionService;
+    private final BlobStorage storage;
 
     /**
      * Returns a list of resources shared with user.
@@ -304,6 +307,14 @@ public void discardSharedAccess(String bucket, String location, ResourceLinkColl
     }
 
     public void copySharedAccess(String bucket, String location, ResourceDescription source, ResourceDescription destination) {
+        if (!hasResource(source)) {
+            throw new IllegalArgumentException("source resource %s does not exists".formatted(source.getUrl()));
+        }
+
+        if (!hasResource(destination)) {
+            throw new IllegalArgumentException("destination resource %s dos not exists".formatted(destination.getUrl()));
+        }
+
         ResourceType sourceResourceType = source.getType();
         ResourceDescription sharedByMeResource = getShareResource(ResourceType.SHARED_BY_ME, sourceResourceType, bucket, location);
         SharedByMeDto sharedByMeDto = ProxyUtil.convertToObject(resourceService.getResource(sharedByMeResource), SharedByMeDto.class);
@@ -361,10 +372,12 @@ private void addSharedResource(String bucket, String location, String link, Reso
         ResourceDescription sharedByMeResource = getShareResource(ResourceType.SHARED_WITH_ME, resourceType, bucket, location);
         resourceService.computeResource(sharedByMeResource, state -> {
             ResourceLinkCollection sharedWithMe = ProxyUtil.convertToObject(state, ResourceLinkCollection.class);
-            if (sharedWithMe != null) {
-                sharedWithMe.getResources().add(new ResourceLink(link));
+            if (sharedWithMe == null) {
+                sharedWithMe = new ResourceLinkCollection(new HashSet<>());
             }
 
+            sharedWithMe.getResources().add(new ResourceLink(link));
+
             return ProxyUtil.convertToString(sharedWithMe);
         });
     }
@@ -389,6 +402,10 @@ private ResourceDescription getResourceFromLink(String url) {
         }
     }
 
+    private boolean hasResource(ResourceDescription resource) {
+        return ResourceUtil.hasResource(resource, resourceService, storage);
+    }
+
     private ResourceDescription getShareResource(ResourceType shareResourceType, ResourceType requestedResourceType, String bucket, String location) {
         return ResourceDescription.fromDecoded(shareResourceType, bucket, location,
                 requestedResourceType.getGroup() + BlobStorageUtil.PATH_SEPARATOR + SHARE_RESOURCE_FILENAME);

diff --git a/src/main/java/com/epam/aidial/core/util/ResourceUtil.java b/src/main/java/com/epam/aidial/core/util/ResourceUtil.java
@@ -0,0 +1,18 @@
+package com.epam.aidial.core.util;
+
+import com.epam.aidial.core.service.ResourceService;
+import com.epam.aidial.core.storage.BlobStorage;
+import com.epam.aidial.core.storage.ResourceDescription;
+import lombok.experimental.UtilityClass;
+
+@UtilityClass
+public class ResourceUtil {
+
+    public static boolean hasResource(ResourceDescription resource, ResourceService resourceService, BlobStorage storage) {
+        return switch (resource.getType()) {
+            case FILE -> storage.exists(resource.getAbsoluteFilePath());
+            case CONVERSATION, PROMPT -> resourceService.hasResource(resource);
+            default -> throw new IllegalArgumentException("Unsupported resource type " + resource.getType());
+        };
+    }
+}
diff --git a/src/test/java/com/epam/aidial/core/ResourceBaseTest.java b/src/test/java/com/epam/aidial/core/ResourceBaseTest.java
@@ -67,6 +67,15 @@ public class ResourceBaseTest {
             }
             """;
 
+    public static final String PROMPT_BODY = """
+                        {
+            "id": "prompt_id",
+            "name": "prompt",
+            "folderId": "folder",
+            "content": "content"
+            }
+            """;
+
     RedisServer redis;
     AiDial dial;
     Path testDir;