[OPIK-523] batch delete for prompts (#806)

* OPIK-523 failing test * OPIK-523 failing test green [WIP] * OPIK-523 failing test green * OPIK-523 pr comments * OPIK-523 pr comments * OPIK-523 pr comments
comet-ml · Dec 4, 2024 · b837127 · b837127
1 parent 7d8f84e
commit b837127
Show file tree

Hide file tree

Showing 11 changed files with 113 additions and 58 deletions.
diff --git a/apps/opik-backend/src/main/java/com/comet/opik/api/resources/v1/priv/DatasetsResource.java b/apps/opik-backend/src/main/java/com/comet/opik/api/resources/v1/priv/DatasetsResource.java
@@ -208,13 +208,13 @@ public Response deleteDatasetByName(
             @ApiResponse(responseCode = "204", description = "No content"),
     })
     public Response deleteDatasetsBatch(
-            @RequestBody(content = @Content(schema = @Schema(implementation = BatchDelete.class))) @NotNull @Valid BatchDelete batchDelete) {
+            @NotNull @RequestBody(content = @Content(schema = @Schema(implementation = BatchDelete.class))) @NotNull @Valid BatchDelete batchDelete) {
 
         String workspaceId = requestContext.get().getWorkspaceId();
 
-        log.info("Deleting datasets by ids '{}' on workspace_id '{}'", batchDelete.ids(), workspaceId);
+        log.info("Deleting datasets by ids, count '{}' on workspace_id '{}'", batchDelete.ids().size(), workspaceId);
         service.delete(batchDelete.ids());
-        log.info("Deleted datasets by ids '{}' on workspace_id '{}'", batchDelete.ids(), workspaceId);
+        log.info("Deleted datasets by ids, count '{}' on workspace_id '{}'", batchDelete.ids().size(), workspaceId);
 
         return Response.noContent().build();
     }

diff --git a/...ackend/src/main/java/com/comet/opik/api/resources/v1/priv/FeedbackDefinitionResource.java b/...ackend/src/main/java/com/comet/opik/api/resources/v1/priv/FeedbackDefinitionResource.java
@@ -163,11 +163,13 @@ public Response deleteById(@PathParam("id") UUID id) {
             @ApiResponse(responseCode = "204", description = "No Content"),
     })
     public Response deleteFeedbackDefinitionsBatch(
-            @RequestBody(content = @Content(schema = @Schema(implementation = BatchDelete.class))) @Valid BatchDelete batchDelete) {
+            @NotNull @RequestBody(content = @Content(schema = @Schema(implementation = BatchDelete.class))) @Valid BatchDelete batchDelete) {
         String workspaceId = requestContext.get().getWorkspaceId();
-        log.info("Deleting feedback definitions by ids '{}', on workspace_id '{}'", batchDelete.ids(), workspaceId);
+        log.info("Deleting feedback definitions by ids, count '{}', on workspace_id '{}'", batchDelete.ids().size(),
+                workspaceId);
         service.delete(batchDelete.ids());
-        log.info("Deleted feedback definitions by ids '{}', on workspace_id '{}'", batchDelete.ids(), workspaceId);
+        log.info("Deleted feedback definitions by ids, count '{}', on workspace_id '{}'", batchDelete.ids().size(),
+                workspaceId);
         return Response.noContent().build();
     }
 }
diff --git a/apps/opik-backend/src/main/java/com/comet/opik/api/resources/v1/priv/ProjectsResource.java b/apps/opik-backend/src/main/java/com/comet/opik/api/resources/v1/priv/ProjectsResource.java
@@ -28,6 +28,7 @@
 import jakarta.inject.Provider;
 import jakarta.validation.Valid;
 import jakarta.validation.constraints.Min;
+import jakarta.validation.constraints.NotNull;
 import jakarta.ws.rs.BadRequestException;
 import jakarta.ws.rs.Consumes;
 import jakarta.ws.rs.DELETE;
@@ -198,11 +199,11 @@ public Response retrieveProject(
             @ApiResponse(responseCode = "204", description = "No Content"),
     })
     public Response deleteProjectsBatch(
-            @RequestBody(content = @Content(schema = @Schema(implementation = BatchDelete.class))) @Valid BatchDelete batchDelete) {
+            @NotNull @RequestBody(content = @Content(schema = @Schema(implementation = BatchDelete.class))) @Valid BatchDelete batchDelete) {
         String workspaceId = requestContext.get().getWorkspaceId();
-        log.info("Deleting projects by ids '{}', on workspace_id '{}'", batchDelete.ids(), workspaceId);
+        log.info("Deleting projects by ids, count '{}', on workspace_id '{}'", batchDelete.ids().size(), workspaceId);
         projectService.delete(batchDelete.ids());
-        log.info("Deleted projects by ids '{}', on workspace_id '{}'", batchDelete.ids(), workspaceId);
+        log.info("Deleted projects by ids, count '{}', on workspace_id '{}'", batchDelete.ids().size(), workspaceId);
         return Response.noContent().build();
     }
 

diff --git a/apps/opik-backend/src/main/java/com/comet/opik/api/resources/v1/priv/PromptResource.java b/apps/opik-backend/src/main/java/com/comet/opik/api/resources/v1/priv/PromptResource.java
@@ -1,6 +1,7 @@
 package com.comet.opik.api.resources.v1.priv;
 
 import com.codahale.metrics.annotation.Timed;
+import com.comet.opik.api.BatchDelete;
 import com.comet.opik.api.CreatePromptVersion;
 import com.comet.opik.api.Prompt;
 import com.comet.opik.api.Prompt.PromptPage;
@@ -23,6 +24,7 @@
 import jakarta.inject.Provider;
 import jakarta.validation.Valid;
 import jakarta.validation.constraints.Min;
+import jakarta.validation.constraints.NotNull;
 import jakarta.ws.rs.Consumes;
 import jakarta.ws.rs.DELETE;
 import jakarta.ws.rs.DefaultValue;
@@ -163,6 +165,20 @@ public Response deletePrompt(@PathParam("id") UUID id) {
         return Response.noContent().build();
     }
 
+    @POST
+    @Path("/delete")
+    @Operation(operationId = "deletePromptsBatch", summary = "Delete prompts", description = "Delete prompts batch", responses = {
+            @ApiResponse(responseCode = "204", description = "No Content"),
+    })
+    public Response deletePromptsBatch(
+            @NotNull @RequestBody(content = @Content(schema = @Schema(implementation = BatchDelete.class))) @Valid BatchDelete batchDelete) {
+        String workspaceId = requestContext.get().getWorkspaceId();
+        log.info("Deleting prompts by ids, count '{}', on workspace_id '{}'", batchDelete.ids().size(), workspaceId);
+        promptService.delete(batchDelete.ids());
+        log.info("Deleted prompts by ids, count '{}', on workspace_id '{}'", batchDelete.ids().size(), workspaceId);
+        return Response.noContent().build();
+    }
+
     @POST
     @Path("/versions")
     @Operation(operationId = "createPromptVersion", summary = "Create prompt version", description = "Create prompt version", responses = {

diff --git a/apps/opik-backend/src/main/java/com/comet/opik/domain/BatchDeleteUtils.java b/apps/opik-backend/src/main/java/com/comet/opik/domain/BatchDeleteUtils.java
diff --git a/apps/opik-backend/src/main/java/com/comet/opik/domain/DatasetService.java b/apps/opik-backend/src/main/java/com/comet/opik/domain/DatasetService.java
@@ -287,16 +287,16 @@ public void delete(@NonNull UUID id) {
     @Override
     public void delete(Set<UUID> ids) {
         if (ids.isEmpty()) {
+            log.info("ids list is empty, returning");
             return;
         }
 
         String workspaceId = requestContext.get().getWorkspaceId();
 
-        template.inTransaction(WRITE, BatchDeleteUtils.getHandler(
-                DatasetDAO.class,
-                repository -> repository.findByIds(ids, workspaceId),
-                Dataset::id,
-                (repository, idsToDelete) -> repository.delete(idsToDelete, workspaceId)));
+        template.inTransaction(WRITE, handle -> {
+            handle.attach(DatasetDAO.class).delete(ids, workspaceId);
+            return null;
+        });
     }
 
     @Override

diff --git a/apps/opik-backend/src/main/java/com/comet/opik/domain/FeedbackDefinitionService.java b/apps/opik-backend/src/main/java/com/comet/opik/domain/FeedbackDefinitionService.java
@@ -201,16 +201,16 @@ public void delete(@NonNull UUID id) {
     @Override
     public void delete(Set<UUID> ids) {
         if (ids.isEmpty()) {
+            log.info("ids list is empty, returning");
             return;
         }
 
         String workspaceId = requestContext.get().getWorkspaceId();
 
-        template.inTransaction(WRITE, BatchDeleteUtils.getHandler(
-                FeedbackDefinitionDAO.class,
-                repository -> repository.findByIds(ids, workspaceId),
-                FeedbackDefinitionModel::id,
-                (repository, idsToDelete) -> repository.delete(idsToDelete, workspaceId)));
+        template.inTransaction(WRITE, handle -> {
+            handle.attach(FeedbackDefinitionDAO.class).delete(ids, workspaceId);
+            return null;
+        });
     }
 
     private NotFoundException createNotFoundError() {

diff --git a/apps/opik-backend/src/main/java/com/comet/opik/domain/ProjectService.java b/apps/opik-backend/src/main/java/com/comet/opik/domain/ProjectService.java
@@ -223,16 +223,16 @@ public void delete(@NonNull UUID id) {
     @Override
     public void delete(Set<UUID> ids) {
         if (ids.isEmpty()) {
+            log.info("ids list is empty, returning");
             return;
         }
 
         String workspaceId = requestContext.get().getWorkspaceId();
 
-        template.inTransaction(WRITE, BatchDeleteUtils.getHandler(
-                ProjectDAO.class,
-                repository -> repository.findByIds(ids, workspaceId),
-                Project::id,
-                (repository, idsToDelete) -> repository.delete(idsToDelete, workspaceId)));
+        template.inTransaction(WRITE, handle -> {
+            handle.attach(ProjectDAO.class).delete(ids, workspaceId);
+            return null;
+        });
     }
 
     @Override

diff --git a/apps/opik-backend/src/main/java/com/comet/opik/domain/PromptDAO.java b/apps/opik-backend/src/main/java/com/comet/opik/domain/PromptDAO.java
@@ -8,13 +8,15 @@
 import org.jdbi.v3.sqlobject.config.RegisterConstructorMapper;
 import org.jdbi.v3.sqlobject.customizer.AllowUnusedBindings;
 import org.jdbi.v3.sqlobject.customizer.Bind;
+import org.jdbi.v3.sqlobject.customizer.BindList;
 import org.jdbi.v3.sqlobject.customizer.BindMethods;
 import org.jdbi.v3.sqlobject.customizer.Define;
 import org.jdbi.v3.sqlobject.statement.SqlQuery;
 import org.jdbi.v3.sqlobject.statement.SqlUpdate;
 import org.jdbi.v3.stringtemplate4.UseStringTemplateEngine;
 
 import java.util.List;
+import java.util.Set;
 import java.util.UUID;
 
 @RegisterColumnMapper(PromptVersionColumnMapper.class)
@@ -91,4 +93,6 @@ List<Prompt> find(@Define("name") @Bind("name") String name, @Bind("workspace_id
     @SqlUpdate("DELETE FROM prompts WHERE id = :id AND workspace_id = :workspace_id")
     int delete(@Bind("id") UUID id, @Bind("workspace_id") String workspaceId);
 
+    @SqlUpdate("DELETE FROM prompts WHERE id IN (<ids>) AND workspace_id = :workspaceId")
+    void delete(@BindList("ids") Set<UUID> ids, @Bind("workspaceId") String workspaceId);
 }
diff --git a/apps/opik-backend/src/main/java/com/comet/opik/domain/PromptService.java b/apps/opik-backend/src/main/java/com/comet/opik/domain/PromptService.java
@@ -45,6 +45,8 @@ public interface PromptService {
 
     void delete(UUID id);
 
+    void delete(Set<UUID> ids);
+
     Prompt getById(UUID id);
 
     PromptVersionPage getVersionsByPromptId(UUID promptId, int page, int size);
@@ -284,6 +286,21 @@ public void delete(@NonNull UUID id) {
         });
     }
 
+    @Override
+    public void delete(Set<UUID> ids) {
+        if (ids.isEmpty()) {
+            log.info("ids list is empty, returning");
+            return;
+        }
+
+        String workspaceId = requestContext.get().getWorkspaceId();
+
+        transactionTemplate.inTransaction(WRITE, handle -> {
+            handle.attach(PromptDAO.class).delete(ids, workspaceId);
+            return null;
+        });
+    }
+
     private PromptVersion retryableCreateVersion(String workspaceId, CreatePromptVersion request, Prompt prompt,
             String userName) {
         return EntityConstraintHandler.handle(() -> {

diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/PromptResourceTest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/PromptResourceTest.java
@@ -1,5 +1,6 @@
 package com.comet.opik.api.resources.v1.priv;
 
+import com.comet.opik.api.BatchDelete;
 import com.comet.opik.api.CreatePromptVersion;
 import com.comet.opik.api.Prompt;
 import com.comet.opik.api.PromptVersion;
@@ -50,6 +51,7 @@
 
 import java.sql.SQLException;
 import java.time.Instant;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Set;
 import java.util.UUID;
@@ -1117,6 +1119,53 @@ void when__promptDoesNotExist__thenReturnNotFound() {
 
             getPromptAndAssertNotFound(promptId, API_KEY, TEST_WORKSPACE);
         }
+
+        @Test
+        @DisplayName("delete batch of prompts")
+        void deleteBatch() {
+            var apiKey = UUID.randomUUID().toString();
+            var workspaceName = UUID.randomUUID().toString();
+            var workspaceId = UUID.randomUUID().toString();
+            mockTargetWorkspace(apiKey, workspaceName, workspaceId);
+
+            var ids = PodamFactoryUtils.manufacturePojoList(factory,
+                    Prompt.class).stream()
+                    .map(prompt -> createPrompt(prompt.toBuilder()
+                            .lastUpdatedBy(USER)
+                            .createdBy(USER)
+                            .build(), apiKey, workspaceName))
+                    .toList();
+            var idsToDelete = ids.subList(0, 3);
+            var notDeletedIds = ids.subList(3, ids.size());
+
+            try (var actualResponse = client.target(RESOURCE_PATH.formatted(baseURI))
+                    .path("delete")
+                    .request()
+                    .header(HttpHeaders.AUTHORIZATION, apiKey)
+                    .header(WORKSPACE_HEADER, workspaceName)
+                    .post(Entity.json(new BatchDelete(new HashSet<>(idsToDelete))))) {
+
+                assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(HttpStatus.SC_NO_CONTENT);
+                assertThat(actualResponse.hasEntity()).isFalse();
+            }
+
+            var actualResponse = client.target(RESOURCE_PATH.formatted(baseURI))
+                    .queryParam("size", ids.size())
+                    .queryParam("page", 1)
+                    .request()
+                    .header(HttpHeaders.AUTHORIZATION, apiKey)
+                    .header(WORKSPACE_HEADER, workspaceName)
+                    .get();
+
+            var actualEntity = actualResponse.readEntity(Prompt.PromptPage.class);
+
+            assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(HttpStatus.SC_OK);
+            assertThat(actualEntity.size()).isEqualTo(notDeletedIds.size());
+            assertThat(actualEntity.content().stream().map(Prompt::id).toList())
+                    .usingRecursiveComparison()
+                    .ignoringCollectionOrder()
+                    .isEqualTo(notDeletedIds);
+        }
     }
 
     @Nested