TSPS-405 Add quota units to pipeline quotas (#205)

DataBiosphere · Feb 26, 2025 · 0a482f0 · 0a482f0
1 parent cb2c82d
commit 0a482f0
Show file tree

Hide file tree

Showing 12 changed files with 117 additions and 6 deletions.
diff --git a/common/openapi.yml b/common/openapi.yml
@@ -932,18 +932,25 @@ components:
         Maximum allowable quota for the pipeline
       type: integer
 
+    QuotaUnits:
+      description: |
+        Units for pipeline quota
+      type: string
+
     QuotaWithDetails:
       description: |
         Object containing the pipeline identifier, quota limit, and quota usage of a Pipeline.
       type: object
-      required: [ pipelineName, quotaLimit, quotaConsumed ]
+      required: [ pipelineName, quotaLimit, quotaConsumed, quotaUnits ]
       properties:
         pipelineName:
           $ref: "#/components/schemas/PipelineName"
         quotaLimit:
           $ref: "#/components/schemas/QuotaLimit"
         quotaConsumed:
           $ref: "#/components/schemas/QuotaConsumed"
+        quotaUnits:
+          $ref: "#/components/schemas/QuotaUnits"
 
     StartPipelineRunRequestBody:
       description: |

diff --git a/service/src/main/java/bio/terra/pipelines/app/controller/QuotasController.java b/service/src/main/java/bio/terra/pipelines/app/controller/QuotasController.java
@@ -4,6 +4,7 @@
 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.common.utils.QuotaUnitsEnum;
 import bio.terra.pipelines.db.entities.UserQuota;
 import bio.terra.pipelines.generated.api.QuotasApi;
 import bio.terra.pipelines.generated.model.ApiQuotaWithDetails;
@@ -51,15 +52,18 @@ public ResponseEntity<ApiQuotaWithDetails> getQuotaForPipeline(String pipelineNa
     UserQuota userQuota =
         quotasService.getOrCreateQuotaForUserAndPipeline(
             user.getSubjectId(), validatedPipelineName);
+    QuotaUnitsEnum quotaUnits = quotasService.getQuotaUnitsForPipeline(validatedPipelineName);
 
-    return new ResponseEntity<>(quotasToApiQuotaWithDetails(userQuota), HttpStatus.OK);
+    return new ResponseEntity<>(quotasToApiQuotaWithDetails(userQuota, quotaUnits), HttpStatus.OK);
   }
 
-  static ApiQuotaWithDetails quotasToApiQuotaWithDetails(UserQuota userQuota) {
+  static ApiQuotaWithDetails quotasToApiQuotaWithDetails(
+      UserQuota userQuota, QuotaUnitsEnum quotaUnits) {
 
     return new ApiQuotaWithDetails()
         .pipelineName(userQuota.getPipelineName().getValue())
         .quotaConsumed(userQuota.getQuotaConsumed())
-        .quotaLimit(userQuota.getQuota());
+        .quotaLimit(userQuota.getQuota())
+        .quotaUnits(quotaUnits.getValue());
   }
 }
diff --git a/service/src/main/java/bio/terra/pipelines/common/utils/QuotaUnitsEnum.java b/service/src/main/java/bio/terra/pipelines/common/utils/QuotaUnitsEnum.java
@@ -0,0 +1,15 @@
+package bio.terra.pipelines.common.utils;
+
+public enum QuotaUnitsEnum {
+  SAMPLES("samples");
+
+  private final String value;
+
+  QuotaUnitsEnum(String value) {
+    this.value = value;
+  }
+
+  public String getValue() {
+    return value;
+  }
+}
diff --git a/service/src/main/java/bio/terra/pipelines/db/entities/PipelineQuota.java b/service/src/main/java/bio/terra/pipelines/db/entities/PipelineQuota.java
@@ -1,6 +1,7 @@
 package bio.terra.pipelines.db.entities;
 
 import bio.terra.pipelines.common.utils.PipelinesEnum;
+import bio.terra.pipelines.common.utils.QuotaUnitsEnum;
 import jakarta.persistence.*;
 import lombok.Getter;
 import lombok.NoArgsConstructor;
@@ -28,9 +29,17 @@ public class PipelineQuota {
   @Column(name = "min_quota_consumed")
   private int minQuotaConsumed;
 
-  public PipelineQuota(PipelinesEnum pipelineName, int defaultQuota, int minQuotaConsumed) {
+  @Column(name = "quota_units")
+  private QuotaUnitsEnum quotaUnits;
+
+  public PipelineQuota(
+      PipelinesEnum pipelineName,
+      int defaultQuota,
+      int minQuotaConsumed,
+      QuotaUnitsEnum quotaUnits) {
     this.pipelineName = pipelineName;
     this.defaultQuota = defaultQuota;
     this.minQuotaConsumed = minQuotaConsumed;
+    this.quotaUnits = quotaUnits;
   }
 }
diff --git a/service/src/main/java/bio/terra/pipelines/db/entities/converters/QuotaUnitsConverter.java b/service/src/main/java/bio/terra/pipelines/db/entities/converters/QuotaUnitsConverter.java
@@ -0,0 +1,19 @@
+package bio.terra.pipelines.db.entities.converters;
+
+import bio.terra.pipelines.common.utils.QuotaUnitsEnum;
+import jakarta.persistence.AttributeConverter;
+import jakarta.persistence.Converter;
+
+// inspired by https://www.baeldung.com/jpa-persisting-enums-in-jpa
+@Converter(autoApply = true)
+public class QuotaUnitsConverter implements AttributeConverter<QuotaUnitsEnum, String> {
+  @Override
+  public String convertToDatabaseColumn(QuotaUnitsEnum quotaUnitsEnum) {
+    return quotaUnitsEnum.getValue();
+  }
+
+  @Override
+  public QuotaUnitsEnum convertToEntityAttribute(String quotaUnitsString) {
+    return QuotaUnitsEnum.valueOf(quotaUnitsString.toUpperCase());
+  }
+}
diff --git a/service/src/main/java/bio/terra/pipelines/db/repositories/PipelineQuotasRepository.java b/service/src/main/java/bio/terra/pipelines/db/repositories/PipelineQuotasRepository.java
@@ -1,9 +1,14 @@
 package bio.terra.pipelines.db.repositories;
 
 import bio.terra.pipelines.common.utils.PipelinesEnum;
+import bio.terra.pipelines.common.utils.QuotaUnitsEnum;
 import bio.terra.pipelines.db.entities.PipelineQuota;
+import org.springframework.data.jpa.repository.Query;
 import org.springframework.data.repository.CrudRepository;
 
 public interface PipelineQuotasRepository extends CrudRepository<PipelineQuota, Long> {
   PipelineQuota findByPipelineName(PipelinesEnum pipelineName);
+
+  @Query("SELECT p.quotaUnits FROM PipelineQuota p WHERE p.pipelineName = ?1")
+  QuotaUnitsEnum findQuotaUnitsByPipelineName(PipelinesEnum pipelineName);
 }
diff --git a/service/src/main/java/bio/terra/pipelines/service/QuotasService.java b/service/src/main/java/bio/terra/pipelines/service/QuotasService.java
@@ -3,6 +3,7 @@
 import bio.terra.common.db.WriteTransaction;
 import bio.terra.common.exception.InternalServerErrorException;
 import bio.terra.pipelines.common.utils.PipelinesEnum;
+import bio.terra.pipelines.common.utils.QuotaUnitsEnum;
 import bio.terra.pipelines.db.entities.PipelineQuota;
 import bio.terra.pipelines.db.entities.UserQuota;
 import bio.terra.pipelines.db.repositories.PipelineQuotasRepository;
@@ -33,6 +34,11 @@ public PipelineQuota getPipelineQuota(PipelinesEnum pipelineName) {
     return pipelineQuotasRepository.findByPipelineName(pipelineName);
   }
 
+  /** This method gets the quota units value for a given pipeline. */
+  public QuotaUnitsEnum getQuotaUnitsForPipeline(PipelinesEnum pipelineName) {
+    return pipelineQuotasRepository.findQuotaUnitsByPipelineName(pipelineName);
+  }
+
   /**
    * This method gets the quota for a given user and pipeline. If the user quota does not exist, it
    * will create a new row in the user quotas table with the default quota for the pipeline.

diff --git a/service/src/main/resources/db/changelog.xml b/service/src/main/resources/db/changelog.xml
@@ -11,6 +11,7 @@
   <include file="changesets/20241212_add_quota_consumed_update_pipeline_display_name.yaml" relativeToChangelogFile="true"/>
   <include file="changesets/20250127_add_expects_custom_value_field_to_input_defs.yaml" relativeToChangelogFile="true"/>
   <include file="changesets/20250220_rename_wdl_method_version_and_wdl_method_name.yaml" relativeToChangelogFile="true"/>
+  <include file="changesets/20250225_add_quota_units.yaml" relativeToChangelogFile="true"/>
 
   <include file="changesets/testdata.yaml" relativeToChangelogFile="true"/>
 

diff --git a/service/src/main/resources/db/changesets/20250225_add_quota_units.yaml b/service/src/main/resources/db/changesets/20250225_add_quota_units.yaml
@@ -0,0 +1,28 @@
+# add quota_units field to pipeline_quotas table and set value for array_imputation pipeline
+
+databaseChangeLog:
+  -  changeSet:
+       id:  add and populate quota_units column in pipeline_quotas table
+       author:  mma
+       changes:
+         # add quota_units column to pipeline_quotas table
+         - addColumn:
+             tableName: pipeline_quotas
+             columns:
+               - column:
+                   name: quota_units
+                   type: text
+                   constraints:
+                     nullable: true
+         # update quota_units value for array_imputation pipeline
+         - update:
+             tableName: pipeline_quotas
+             columns:
+               - column:
+                  name: quota_units
+                  value: 'samples'
+             where: pipeline_name='array_imputation'
+        # make quota_units non-nullable
+         -  addNotNullConstraint:
+                tableName: pipeline_quotas
+                columnName: quota_units
diff --git a/service/src/test/java/bio/terra/pipelines/controller/QuotasControllerTest.java b/service/src/test/java/bio/terra/pipelines/controller/QuotasControllerTest.java
@@ -15,6 +15,7 @@
 import bio.terra.pipelines.app.controller.GlobalExceptionHandler;
 import bio.terra.pipelines.app.controller.QuotasController;
 import bio.terra.pipelines.common.utils.PipelinesEnum;
+import bio.terra.pipelines.common.utils.QuotaUnitsEnum;
 import bio.terra.pipelines.db.entities.UserQuota;
 import bio.terra.pipelines.db.exception.InvalidPipelineException;
 import bio.terra.pipelines.generated.model.ApiQuotaWithDetails;
@@ -52,6 +53,8 @@ void beforeEach() {
     when(quotasServiceMock.getOrCreateQuotaForUserAndPipeline(
             testUser.getSubjectId(), PipelinesEnum.ARRAY_IMPUTATION))
         .thenReturn(testUserQuota);
+    when(quotasServiceMock.getQuotaUnitsForPipeline(PipelinesEnum.ARRAY_IMPUTATION))
+        .thenReturn(QuotaUnitsEnum.SAMPLES);
   }
 
   @Test
@@ -71,6 +74,7 @@ void getQuotaOk() throws Exception {
     assertEquals(testUserQuota.getQuota(), response.getQuotaLimit());
     assertEquals(testUserQuota.getQuotaConsumed(), response.getQuotaConsumed());
     assertEquals(testUserQuota.getPipelineName().getValue(), response.getPipelineName());
+    assertEquals(QuotaUnitsEnum.SAMPLES.getValue(), response.getQuotaUnits());
   }
 
   @Test

diff --git a/service/src/test/java/bio/terra/pipelines/db/entities/PipelineQuotaTest.java b/service/src/test/java/bio/terra/pipelines/db/entities/PipelineQuotaTest.java
@@ -3,16 +3,19 @@
 import static org.junit.jupiter.api.Assertions.assertEquals;
 
 import bio.terra.pipelines.common.utils.PipelinesEnum;
+import bio.terra.pipelines.common.utils.QuotaUnitsEnum;
 import bio.terra.pipelines.testutils.BaseTest;
 import org.junit.jupiter.api.Test;
 
 class PipelineQuotaTest extends BaseTest {
 
   @Test
   void testPipelineQuotaConstructor() {
-    PipelineQuota pipelineQuota = new PipelineQuota(PipelinesEnum.ARRAY_IMPUTATION, 10000, 500);
+    PipelineQuota pipelineQuota =
+        new PipelineQuota(PipelinesEnum.ARRAY_IMPUTATION, 10000, 500, QuotaUnitsEnum.SAMPLES);
     assertEquals(PipelinesEnum.ARRAY_IMPUTATION, pipelineQuota.getPipelineName());
     assertEquals(10000, pipelineQuota.getDefaultQuota());
     assertEquals(500, pipelineQuota.getMinQuotaConsumed());
+    assertEquals(QuotaUnitsEnum.SAMPLES, pipelineQuota.getQuotaUnits());
   }
 }
diff --git a/service/src/test/java/bio/terra/pipelines/service/QuotasServiceTest.java b/service/src/test/java/bio/terra/pipelines/service/QuotasServiceTest.java
@@ -4,6 +4,7 @@
 
 import bio.terra.common.exception.InternalServerErrorException;
 import bio.terra.pipelines.common.utils.PipelinesEnum;
+import bio.terra.pipelines.common.utils.QuotaUnitsEnum;
 import bio.terra.pipelines.db.entities.PipelineQuota;
 import bio.terra.pipelines.db.entities.UserQuota;
 import bio.terra.pipelines.db.repositories.PipelineQuotasRepository;
@@ -25,6 +26,15 @@ void getPipelineQuota() {
     assertEquals(PipelinesEnum.ARRAY_IMPUTATION, pipelineQuota.getPipelineName());
     assertEquals(10000, pipelineQuota.getDefaultQuota());
     assertEquals(500, pipelineQuota.getMinQuotaConsumed());
+    assertEquals(QuotaUnitsEnum.SAMPLES, pipelineQuota.getQuotaUnits());
+  }
+
+  @Test
+  void getQuotaUnitsForPipeline() {
+    QuotaUnitsEnum quotaUnits =
+        quotasService.getQuotaUnitsForPipeline(PipelinesEnum.ARRAY_IMPUTATION);
+
+    assertEquals(QuotaUnitsEnum.SAMPLES, quotaUnits);
   }
 
   @Test