TSPS-357 Add admin endpoint to get a user's quota and update their quota

common/openapi.yml

@@ -4,7 +4,7 @@ info:
   version: 1.0.0
 paths:
   # Admin APIs
-  /api/admin/v1/pipeline/{pipelineName}:
+  /api/admin/v1/pipelines/{pipelineName}:
     parameters:
       - $ref: '#/components/parameters/PipelineName'
     get:
@@ -48,6 +48,51 @@ paths:
         500:
           $ref: '#/components/responses/ServerError'
 
+  /api/admin/v1/quotas/{pipelineName}/{userId}:
+    parameters:
+      - $ref: '#/components/parameters/PipelineName'
+      - $ref: '#/components/parameters/UserId'
+    get:
+      summary: Get quota for a given pipeline and user.
+      tags: [ admin ]
+      operationId: getQuotaForPipelineAndUser
+      responses:
+        200:
+          description: Success
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/AdminQuota'
+        400:
+          $ref: '#/components/responses/BadRequest'
+        403:
+          $ref: '#/components/responses/PermissionDenied'
+        500:
+          $ref: '#/components/responses/ServerError'
+    patch:
+      summary: Update quota limit for a given pipeline and user.
+      tags: [ admin ]
+      operationId: updateQuotaLimitForPipelineAndUser
+      requestBody:
+        required: true
+        content:
+          application/json:
+            schema:
+              $ref: '#/components/schemas/UpdateQuotaLimitRequestBody'
+      responses:
+        200:
+          description: Success
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/AdminQuota'
+        400:
+          $ref: '#/components/responses/BadRequest'
+        403:
+          $ref: '#/components/responses/PermissionDenied'
+        500:
+          $ref: '#/components/responses/ServerError'
+
   # Pipeline related APIs
   /api/pipelines/v1:
     get:
@@ -355,6 +400,14 @@ components:
       schema:
         type: string
 
+    UserId:
+      name: userId
+      in: path
+      description: A string identifier to used to identify a Terra user
+      required: true
+      schema:
+        type: string
+
     WorkspaceId:
       name: workspaceId
       in: path
@@ -528,6 +581,21 @@ components:
         workspaceGoogleProject:
           $ref: "#/components/schemas/PipelineWorkspaceGoogleProject"
 
+    AdminQuota:
+      description: |
+        Object containing the use id, pipeline identifier, quota limit, and quota usage of a Pipeline for a user.
+      type: object
+      required: [ userId, pipelineName, quotaLimit, quotaConsumed ]
+      properties:
+        userId:
+          $ref: "#/components/schemas/UserId"
+        pipelineName:
+          $ref: "#/components/schemas/PipelineName"
+        quotaLimit:
+          $ref: "#/components/schemas/QuotaLimit"
+        quotaConsumed:
+          $ref: "#/components/schemas/QuotaConsumed"
+
     AsyncPipelineRunResponse:
       description: Result of an asynchronous pipeline run request.
       type: object
@@ -859,6 +927,15 @@ components:
         wdlMethodVersion:
           $ref: "#/components/schemas/PipelineWdlMethodVersion"
 
+    UpdateQuotaLimitRequestBody:
+      description: |
+        json object containing the admin provided information used to update a user's quota limit
+      type: object
+      required: [ quotaLimit ]
+      properties:
+        quotaLimit:
+          $ref: "#/components/schemas/QuotaLimit"
+
     UserId:
       description: |
         The identifier string for the user who submitted a job request.

service/src/main/java/bio/terra/pipelines/app/controller/AdminApiController.java

@@ -1,16 +1,20 @@
 package bio.terra.pipelines.app.controller;
 
+import bio.terra.common.exception.BadRequestException;
 import bio.terra.common.iam.SamUser;
 import bio.terra.common.iam.SamUserFactory;
 import bio.terra.pipelines.app.configuration.external.SamConfiguration;
 import bio.terra.pipelines.common.utils.PipelinesEnum;
 import bio.terra.pipelines.db.entities.Pipeline;
+import bio.terra.pipelines.db.entities.UserQuota;
 import bio.terra.pipelines.dependencies.sam.SamService;
 import bio.terra.pipelines.generated.api.AdminApi;
 import bio.terra.pipelines.generated.model.*;
 import bio.terra.pipelines.service.PipelinesService;
+import bio.terra.pipelines.service.QuotasService;
 import io.swagger.annotations.Api;
 import jakarta.servlet.http.HttpServletRequest;
+import java.util.Optional;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.springframework.beans.factory.annotation.Autowired;
@@ -27,19 +31,22 @@ public class AdminApiController implements AdminApi {
   private final HttpServletRequest request;
   private final PipelinesService pipelinesService;
   private final SamService samService;
+  private final QuotasService quotasService;
 
   @Autowired
   public AdminApiController(
       SamConfiguration samConfiguration,
       SamUserFactory samUserFactory,
       HttpServletRequest request,
       PipelinesService pipelinesService,
-      SamService samService) {
+      SamService samService,
+      QuotasService quotasService) {
     this.samConfiguration = samConfiguration;
     this.samUserFactory = samUserFactory;
     this.request = request;
     this.pipelinesService = pipelinesService;
     this.samService = samService;
+    this.quotasService = quotasService;
   }
 
   private static final Logger logger = LoggerFactory.getLogger(AdminApiController.class);
@@ -58,6 +65,28 @@ public ResponseEntity<ApiAdminPipeline> getPipeline(String pipelineName) {
     return new ResponseEntity<>(pipelineToApiAdminPipeline(updatedPipeline), HttpStatus.OK);
   }
 
+  @Override
+  public ResponseEntity<ApiAdminQuota> getQuotaForPipelineAndUser(
+      String pipelineName, String userId) {
+    final SamUser authedUser = getAuthenticatedInfo();
+    samService.checkAdminAuthz(authedUser);
+    PipelinesEnum validatedPipelineName =
+        PipelineApiUtils.validatePipelineName(pipelineName, logger);
+
+    // Check if a row exists for this user + pipeline. Throw an error if it doesn't
+    // A row should exist if a user has run into a quota issue before.
+    Optional<UserQuota> userQuota =
+        quotasService.getQuotaForUserAndPipeline(userId, validatedPipelineName);
+    if (userQuota.isEmpty()) {
+      throw new BadRequestException(
+          String.format(
+              "User quota not found for user %s and pipeline %s",
+              userId, validatedPipelineName.getValue()));
+    }
+
+    return new ResponseEntity<>(userQuotaToApiAdminQuota(userQuota.get()), HttpStatus.OK);
+  }
+
   @Override
   public ResponseEntity<ApiAdminPipeline> updatePipeline(
       String pipelineName, ApiUpdatePipelineRequestBody body) {
@@ -74,6 +103,30 @@ public ResponseEntity<ApiAdminPipeline> updatePipeline(
     return new ResponseEntity<>(pipelineToApiAdminPipeline(updatedPipeline), HttpStatus.OK);
   }
 
+  @Override
+  public ResponseEntity<ApiAdminQuota> updateQuotaLimitForPipelineAndUser(
+      String pipelineName, String userId, ApiUpdateQuotaLimitRequestBody body) {
+    final SamUser authedUser = getAuthenticatedInfo();
+    samService.checkAdminAuthz(authedUser);
+    PipelinesEnum validatedPipelineName =
+        PipelineApiUtils.validatePipelineName(pipelineName, logger);
+
+    // Check if a row exists for this user + pipeline. Throw an error if it doesn't
+    // A row should exist if a user has run into a quota issue before.
+    Optional<UserQuota> userQuota =
+        quotasService.getQuotaForUserAndPipeline(userId, validatedPipelineName);
+    if (userQuota.isEmpty()) {
+      throw new BadRequestException(
+          String.format(
+              "User quota not found for user %s and pipeline %s",
+              userId, validatedPipelineName.getValue()));
+    }
+    int newQuotaLimit = body.getQuotaLimit();
+    UserQuota updatedUserQuota =
+        quotasService.adminUpdateQuotaLimit(userQuota.get(), newQuotaLimit);
+    return new ResponseEntity<>(userQuotaToApiAdminQuota(updatedUserQuota), HttpStatus.OK);
+  }
+
   public ApiAdminPipeline pipelineToApiAdminPipeline(Pipeline pipeline) {
     return new ApiAdminPipeline()
         .pipelineName(pipeline.getName().getValue())
@@ -85,4 +138,12 @@ public ApiAdminPipeline pipelineToApiAdminPipeline(Pipeline pipeline) {
         .workspaceGoogleProject(pipeline.getWorkspaceGoogleProject())
         .wdlMethodVersion(pipeline.getWdlMethodVersion());
   }
+
+  public ApiAdminQuota userQuotaToApiAdminQuota(UserQuota userQuota) {
+    return new ApiAdminQuota()
+        .userId(userQuota.getUserId())
+        .pipelineName(userQuota.getPipelineName().getValue())
+        .quotaLimit(userQuota.getQuota())
+        .quotaConsumed(userQuota.getQuotaConsumed());
+  }
 }

service/src/main/java/bio/terra/pipelines/app/controller/QuotasController.java

@@ -49,7 +49,8 @@ public ResponseEntity<ApiQuotaWithDetails> getQuotaForPipeline(String pipelineNa
     PipelinesEnum validatedPipelineName =
         PipelineApiUtils.validatePipelineName(pipelineName, logger);
     UserQuota userQuota =
-        quotasService.getQuotaForUserAndPipeline(user.getSubjectId(), validatedPipelineName);
+        quotasService.getOrCreateQuotaForUserAndPipeline(
+            user.getSubjectId(), validatedPipelineName);
 
     return new ResponseEntity<>(quotasToApiQuotaWithDetails(userQuota), HttpStatus.OK);
   }

service/src/main/java/bio/terra/pipelines/service/QuotasService.java

@@ -33,10 +33,9 @@ public class QuotasService {
    * will create a new row in the user quotas table with the default quota for the pipeline.
    */
   @WriteTransaction
-  public UserQuota getQuotaForUserAndPipeline(String userId, PipelinesEnum pipelineName) {
+  public UserQuota getOrCreateQuotaForUserAndPipeline(String userId, PipelinesEnum pipelineName) {
     // try to get the user quota
-    Optional<UserQuota> userQuota =
-        userQuotasRepository.findByUserIdAndPipelineName(userId, pipelineName);
+    Optional<UserQuota> userQuota = getQuotaForUserAndPipeline(userId, pipelineName);
     // if the user quota is not found, grab the default pipeline quota and make a new row in user
     // quotas table
     if (userQuota.isEmpty()) {
@@ -55,6 +54,14 @@ public UserQuota getQuotaForUserAndPipeline(String userId, PipelinesEnum pipelin
     return userQuota.get();
   }
 
+  /**
+   * This method gets the quota for a given user and pipeline. If the user quota does not exist, it
+   * will return an empty optional.
+   */
+  public Optional<UserQuota> getQuotaForUserAndPipeline(String userId, PipelinesEnum pipelineName) {
+    return userQuotasRepository.findByUserIdAndPipelineName(userId, pipelineName);
+  }
+
   /**
    * This method updates the quota consumed for a given user and pipeline. It will return the
    * updated user quota object.
@@ -75,4 +82,23 @@ public UserQuota updateQuotaConsumed(UserQuota userQuota, int newQuotaConsumed)
     userQuota.setQuotaConsumed(newQuotaConsumed);
     return userQuotasRepository.save(userQuota);
   }
+
+  /**
+   * This method updates the quota limit for a given user quota. This should only be called from the
+   * Admin Controller
+   *
+   * @param userQuota - the user quota to update
+   * @param newQuotaLimit - the new quota limit
+   * @return - the updated user quota
+   */
+  public UserQuota adminUpdateQuotaLimit(UserQuota userQuota, int newQuotaLimit) {
+    if (newQuotaLimit < userQuota.getQuotaConsumed()) {
+      throw new InternalServerErrorException(
+          String.format(
+              "New quota limit: %d, is less than the quota consumed: %d",
+              newQuotaLimit, userQuota.getQuotaConsumed()));
+    }
+    userQuota.setQuota(newQuotaLimit);
+    return userQuotasRepository.save(userQuota);
+  }
 }

service/src/main/java/bio/terra/pipelines/stairway/QuotaConsumedValidationStep.java

@@ -43,7 +43,7 @@ public StepResult doStep(FlightContext flightContext) {
     // check if user quota used plus quota consumed is less than or equal to user quota
     Integer quotaUsedForThisRun =
         workingMap.get(ImputationJobMapKeys.QUOTA_CONSUMED, Integer.class);
-    UserQuota userQuota = quotasService.getQuotaForUserAndPipeline(userId, pipelineName);
+    UserQuota userQuota = quotasService.getOrCreateQuotaForUserAndPipeline(userId, pipelineName);
 
     // user quota has been exceeded, fail the flight
     int totalQuotaConsumed = userQuota.getQuotaConsumed() + quotaUsedForThisRun;
@@ -79,7 +79,7 @@ public StepResult undoStep(FlightContext flightContext) {
     // update the user quota to be what it was before this run's quota was added
     int quotaUsedForThisRun = workingMap.get(ImputationJobMapKeys.QUOTA_CONSUMED, Integer.class);
 
-    UserQuota userQuota = quotasService.getQuotaForUserAndPipeline(userId, pipelineName);
+    UserQuota userQuota = quotasService.getOrCreateQuotaForUserAndPipeline(userId, pipelineName);
 
     quotasService.updateQuotaConsumed(
         userQuota, userQuota.getQuotaConsumed() - quotaUsedForThisRun);