[RW-13428][risk=no] Added cron job for checking credit expiration (#8735

)
all-of-us · Aug 30, 2024 · eea11b1 · eea11b1
1 parent 7a951e0
commit eea11b1
Show file tree

Hide file tree

Showing 27 changed files with 580 additions and 42 deletions.
diff --git a/api/config/config_local.json b/api/config/config_local.json
@@ -150,7 +150,8 @@
   "offlineBatch": {
     "unsafeCloudTasksForwardingHost": "http:\/\/localhost:8081",
     "usersPerAuditTask": 20,
-    "usersPerSynchronizeAccessTask": 50
+    "usersPerSynchronizeAccessTask": 50,
+    "usersPerCheckInitialCreditsExpirationTask": 50
   },
   "egressAlertRemediationPolicy": {
     "enableJiraTicketing": false,

diff --git a/api/config/config_preprod.json b/api/config/config_preprod.json
@@ -145,7 +145,8 @@
   },
   "offlineBatch": {
     "usersPerAuditTask": 20,
-    "usersPerSynchronizeAccessTask": 50
+    "usersPerSynchronizeAccessTask": 50,
+    "usersPerCheckInitialCreditsExpirationTask": 50
   },
   "termsOfService": {
     "latestAouVersion": 1

diff --git a/api/config/config_prod.json b/api/config/config_prod.json
@@ -151,7 +151,8 @@
   },
   "offlineBatch": {
     "usersPerAuditTask": 20,
-    "usersPerSynchronizeAccessTask": 50
+    "usersPerSynchronizeAccessTask": 50,
+    "usersPerCheckInitialCreditsExpirationTask": 50
   },
   "egressAlertRemediationPolicy": {
     "enableJiraTicketing": true,

diff --git a/api/config/config_stable.json b/api/config/config_stable.json
@@ -149,7 +149,8 @@
   },
   "offlineBatch": {
     "usersPerAuditTask": 20,
-    "usersPerSynchronizeAccessTask": 50
+    "usersPerSynchronizeAccessTask": 50,
+    "usersPerCheckInitialCreditsExpirationTask": 50
   },
   "egressAlertRemediationPolicy": {
     "enableJiraTicketing": false,

diff --git a/api/config/config_staging.json b/api/config/config_staging.json
@@ -146,7 +146,8 @@
   },
   "offlineBatch": {
     "usersPerAuditTask": 20,
-    "usersPerSynchronizeAccessTask": 50
+    "usersPerSynchronizeAccessTask": 50,
+    "usersPerCheckInitialCreditsExpirationTask": 50
   },
   "egressAlertRemediationPolicy": {
     "enableJiraTicketing": false,

diff --git a/api/config/config_test.json b/api/config/config_test.json
@@ -151,7 +151,8 @@
   },
   "offlineBatch": {
     "usersPerAuditTask": 20,
-    "usersPerSynchronizeAccessTask": 50
+    "usersPerSynchronizeAccessTask": 50,
+    "usersPerCheckInitialCreditsExpirationTask": 50
   },
   "egressAlertRemediationPolicy": {
     "enableJiraTicketing": false,

diff --git a/api/db/changelog/db.changelog-233-add-user-initial-credit-expiration-notification-column.xml b/api/db/changelog/db.changelog-233-add-user-initial-credit-expiration-notification-column.xml
@@ -0,0 +1,15 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<databaseChangeLog
+  xmlns="http://www.liquibase.org/xml/ns/dbchangelog"
+  xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+  xsi:schemaLocation="http://www.liquibase.org/xml/ns/dbchangelog
+                      http://www.liquibase.org/xml/ns/dbchangelog/dbchangelog-3.1.xsd">
+  <changeSet author="evrii" id="changelog-232-add-user_initial_credits_expiration">
+    <addColumn tableName="user_initial_credits_expiration">
+      <column name="notification_status" type="ENUM('NO_NOTIFICATION_SENT','EXPIRATION_NOTIFICATION_SENT')" defaultValue="NO_NOTIFICATION_SENT">
+        <constraints nullable="false"/>
+      </column>
+    </addColumn>
+
+  </changeSet>
+</databaseChangeLog>
diff --git a/api/db/changelog/db.changelog-master.xml b/api/db/changelog/db.changelog-master.xml
@@ -240,6 +240,7 @@
   <include file="changelog/db.changelog-230-create-featured-workspace-table.xml"/>
   <include file="changelog/db.changelog-231-drop-featured-workspace-description-col.xml"/>
   <include file="changelog/db.changelog-232-add-user_initial_credits_expiration.xml"/>
+  <include file="changelog/db.changelog-233-add-user-initial-credit-expiration-notification-column.xml"/>
   <!--
    Note: to update the DB locally, do the following:
    - Migrate schema changes: `./project.rb run-local-all-migrations`

diff --git a/api/src/main/java/org/pmiops/workbench/api/CloudTaskUserController.java b/api/src/main/java/org/pmiops/workbench/api/CloudTaskUserController.java
@@ -17,6 +17,7 @@
 import org.pmiops.workbench.db.model.DbUser;
 import org.pmiops.workbench.exceptions.WorkbenchException;
 import org.pmiops.workbench.google.CloudResourceManagerService;
+import org.pmiops.workbench.initialcredits.InitialCreditsExpirationService;
 import org.pmiops.workbench.model.AccessModuleStatus;
 import org.pmiops.workbench.model.AuditProjectAccessRequest;
 import org.pmiops.workbench.model.SynchronizeUserAccessRequest;
@@ -55,18 +56,21 @@ public class CloudTaskUserController implements CloudTaskUserApiDelegate {
   private final UserService userService;
   private final AccessModuleService accessModuleService;
   private final FreeTierBillingBatchUpdateService freeTierBillingUpdateService;
+  private final InitialCreditsExpirationService initialCreditsExpirationService;
 
   CloudTaskUserController(
       UserDao userDao,
       CloudResourceManagerService cloudResourceManagerService,
       UserService userService,
       AccessModuleService accessModuleService,
-      FreeTierBillingBatchUpdateService freeTierBillingUpdateService) {
+      FreeTierBillingBatchUpdateService freeTierBillingUpdateService,
+      InitialCreditsExpirationService initialCreditsExpirationService) {
     this.userDao = userDao;
     this.cloudResourceManagerService = cloudResourceManagerService;
     this.userService = userService;
     this.accessModuleService = accessModuleService;
     this.freeTierBillingUpdateService = freeTierBillingUpdateService;
+    this.initialCreditsExpirationService = initialCreditsExpirationService;
   }
 
   @Override
@@ -173,6 +177,18 @@ public ResponseEntity<Void> synchronizeUserAccess(SynchronizeUserAccessRequest r
     return ResponseEntity.noContent().build();
   }
 
+  /**
+   * Takes in batch of user Ids check whether users have initial credits that have expired
+   *
+   * @param body : Batch of user IDs from cloud task queue: checkCreditsExpirationForUserIDsQueue
+   * @return
+   */
+  @Override
+  public ResponseEntity<Void> checkCreditsExpirationForUserIDs(List<Long> userIdsList) {
+    initialCreditsExpirationService.checkCreditsExpirationForUserIDs(userIdsList);
+    return ResponseEntity.noContent().build();
+  }
+
   // v1 cloudresourcemanager project.getParent() returned a ResourceId with type and id fields
   // v3 returns a string instead, in the format type/id
 

diff --git a/api/src/main/java/org/pmiops/workbench/api/OfflineUserController.java b/api/src/main/java/org/pmiops/workbench/api/OfflineUserController.java
@@ -40,4 +40,9 @@ public ResponseEntity<Void> sendAccessExpirationEmails() {
     userService.getAllUsers().forEach(userService::maybeSendAccessTierExpirationEmails);
     return ResponseEntity.noContent().build();
   }
+
+  public ResponseEntity<Void> checkInitialCreditsExpiration() {
+    taskQueueService.groupAndPushCheckInitialCreditExpirationTasks(userService.getAllUserIds());
+    return ResponseEntity.noContent().build();
+  }
 }
diff --git a/api/src/main/java/org/pmiops/workbench/billing/FreeTierBillingService.java b/api/src/main/java/org/pmiops/workbench/billing/FreeTierBillingService.java
@@ -8,8 +8,13 @@
 import java.sql.Timestamp;
 import java.time.Duration;
 import java.time.LocalDateTime;
-import java.util.*;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.List;
+import java.util.Map;
 import java.util.Map.Entry;
+import java.util.Optional;
+import java.util.Set;
 import java.util.stream.Collectors;
 import org.pmiops.workbench.actionaudit.auditors.UserServiceAuditor;
 import org.pmiops.workbench.cloudtasks.TaskQueueService;
@@ -20,6 +25,7 @@
 import org.pmiops.workbench.db.model.DbUser;
 import org.pmiops.workbench.db.model.DbWorkspace;
 import org.pmiops.workbench.db.model.DbWorkspaceFreeTierUsage;
+import org.pmiops.workbench.initialcredits.InitialCreditsExpirationService;
 import org.pmiops.workbench.model.BillingStatus;
 import org.pmiops.workbench.utils.CostComparisonUtils;
 import org.pmiops.workbench.workspaces.WorkspaceUtils;
@@ -39,6 +45,7 @@ public class FreeTierBillingService {
   private final WorkspaceDao workspaceDao;
   private final WorkspaceFreeTierUsageDao workspaceFreeTierUsageDao;
   private final WorkspaceFreeTierUsageService workspaceFreeTierUsageService;
+  private final InitialCreditsExpirationService initialCreditsExpirationService;
 
   private static final Logger logger = LoggerFactory.getLogger(FreeTierBillingService.class);
 
@@ -50,14 +57,16 @@ public FreeTierBillingService(
       UserServiceAuditor userServiceAuditor,
       WorkspaceDao workspaceDao,
       WorkspaceFreeTierUsageDao workspaceFreeTierUsageDao,
-      WorkspaceFreeTierUsageService workspaceFreeTierUsageService) {
+      WorkspaceFreeTierUsageService workspaceFreeTierUsageService,
+      InitialCreditsExpirationService initialCreditsExpirationService) {
     this.taskQueueService = taskQueueService;
     this.userDao = userDao;
     this.workbenchConfigProvider = workbenchConfigProvider;
     this.userServiceAuditor = userServiceAuditor;
     this.workspaceDao = workspaceDao;
     this.workspaceFreeTierUsageDao = workspaceFreeTierUsageDao;
     this.workspaceFreeTierUsageService = workspaceFreeTierUsageService;
+    this.initialCreditsExpirationService = initialCreditsExpirationService;
   }
 
   public double getWorkspaceFreeTierBillingUsage(DbWorkspace dbWorkspace) {
@@ -193,9 +202,11 @@ public double getUserFreeTierDollarLimit(DbUser user) {
   public boolean maybeSetDollarLimitOverride(DbUser user, double newDollarLimit) {
     final Double previousLimitMaybe = user.getFreeTierCreditsLimitDollarsOverride();
 
-    if (previousLimitMaybe != null
-        || CostComparisonUtils.costsDiffer(
-            newDollarLimit, workbenchConfigProvider.get().billing.defaultFreeCreditsDollarLimit)) {
+    if (!initialCreditsExpirationService.haveCreditsExpired(user)
+        && (previousLimitMaybe != null
+            || CostComparisonUtils.costsDiffer(
+                newDollarLimit,
+                workbenchConfigProvider.get().billing.defaultFreeCreditsDollarLimit))) {
 
       // TODO: prevent setting this limit directly except in this method?
       user.setFreeTierCreditsLimitDollarsOverride(newDollarLimit);

diff --git a/api/src/main/java/org/pmiops/workbench/cloudtasks/TaskQueueService.java b/api/src/main/java/org/pmiops/workbench/cloudtasks/TaskQueueService.java
@@ -44,7 +44,8 @@ public class TaskQueueService {
   private static final String DELETE_TEST_WORKSPACES_PATH = BASE_PATH + "/deleteTestUserWorkspaces";
   private static final String DELETE_RAWLS_TEST_WORKSPACES_PATH =
       BASE_PATH + "/deleteTestUserWorkspacesInRawls";
-
+  private static final String CHECK_CREDITS_EXPIRATION_FOR_USER_IDS_PATH =
+      BASE_PATH + "/checkCreditsExpirationForUserIDs";
   private static final String CHECK_AND_ALERT_FREE_TIER_USAGE =
       BASE_PATH + "/checkAndAlertFreeTierBillingUsage";
 
@@ -58,16 +59,17 @@ public class TaskQueueService {
   private static final String DELETE_TEST_WORKSPACES_QUEUE_NAME = "deleteTestUserWorkspacesQueue";
   private static final String DELETE_RAWLS_TEST_WORKSPACES_QUEUE_NAME =
       "deleteTestUserRawlsWorkspacesQueue";
-
   private static final String FREE_TIER_BILLING_QUEUE = "freeTierBillingQueue";
-
   private static final String EXPIRED_FREE_CREDITS_QUEUE_NAME = "expiredFreeCreditsQueue";
+  private static final String CHECK_CREDITS_EXPIRATION_FOR_USER_IDS_QUEUE_NAME =
+      "checkCreditsExpirationForUserIDsQueue";
+
   private static final Logger LOGGER = Logger.getLogger(TaskQueueService.class.getName());
 
-  private WorkbenchLocationConfigService locationConfigService;
-  private Provider<CloudTasksClient> cloudTasksClientProvider;
-  private Provider<WorkbenchConfig> workbenchConfigProvider;
-  private Provider<UserAuthentication> userAuthenticationProvider;
+  private final WorkbenchLocationConfigService locationConfigService;
+  private final Provider<CloudTasksClient> cloudTasksClientProvider;
+  private final Provider<WorkbenchConfig> workbenchConfigProvider;
+  private final Provider<UserAuthentication> userAuthenticationProvider;
 
   public TaskQueueService(
       WorkbenchLocationConfigService locationConfigService,
@@ -219,6 +221,18 @@ public void pushInitialCreditsExpiryTask(
             .liveCostByCreator(liveCostByCreator));
   }
 
+  public void groupAndPushCheckInitialCreditExpirationTasks(List<Long> userIds) {
+    WorkbenchConfig workbenchConfig = workbenchConfigProvider.get();
+    CloudTasksUtils.partitionList(
+            userIds, workbenchConfig.offlineBatch.usersPerCheckInitialCreditsExpirationTask)
+        .forEach(
+            batch ->
+                createAndPushTask(
+                    CHECK_CREDITS_EXPIRATION_FOR_USER_IDS_QUEUE_NAME,
+                    CHECK_CREDITS_EXPIRATION_FOR_USER_IDS_PATH,
+                    batch));
+  }
+
   private String createAndPushTask(String queueName, String taskUri, Object jsonBody) {
     return createAndPushTask(queueName, taskUri, jsonBody, ImmutableMap.of());
   }

diff --git a/api/src/main/java/org/pmiops/workbench/config/WorkbenchConfig.java b/api/src/main/java/org/pmiops/workbench/config/WorkbenchConfig.java
@@ -400,6 +400,8 @@ public static class OfflineBatchConfig {
     public Integer usersPerAuditTask;
     // Number of users to process within a single access synchronization task.
     public Integer usersPerSynchronizeAccessTask;
+    // Number of users to process within a single check initial credits expiration task.
+    public Integer usersPerCheckInitialCreditsExpirationTask;
   }
 
   /**

diff --git a/api/src/main/java/org/pmiops/workbench/db/model/DbUserInitialCreditsExpiration.java b/api/src/main/java/org/pmiops/workbench/db/model/DbUserInitialCreditsExpiration.java
@@ -20,6 +20,7 @@ public class DbUserInitialCreditsExpiration {
   private Timestamp expirationTime;
   private boolean bypassed;
   private int extensionCount;
+  private NotificationStatus notificationStatus;
 
   @Id
   @GeneratedValue(strategy = GenerationType.IDENTITY)
@@ -83,4 +84,20 @@ public DbUserInitialCreditsExpiration setExtensionCount(int extensionCount) {
     this.extensionCount = extensionCount;
     return this;
   }
+
+  @Column(name = "notification_status")
+  public NotificationStatus getNotificationStatus() {
+    return notificationStatus;
+  }
+
+  public DbUserInitialCreditsExpiration setNotificationStatus(
+      NotificationStatus notificationStatus) {
+    this.notificationStatus = notificationStatus;
+    return this;
+  }
+
+  public enum NotificationStatus {
+    NO_NOTIFICATION_SENT,
+    EXPIRATION_NOTIFICATION_SENT
+  }
 }
diff --git a/api/src/main/java/org/pmiops/workbench/initialcredits/InitialCreditsExpirationService.java b/api/src/main/java/org/pmiops/workbench/initialcredits/InitialCreditsExpirationService.java
@@ -1,9 +1,14 @@
 package org.pmiops.workbench.initialcredits;
 
 import java.sql.Timestamp;
+import java.util.List;
 import java.util.Optional;
 import org.pmiops.workbench.db.model.DbUser;
 
 public interface InitialCreditsExpirationService {
   Optional<Timestamp> getCreditsExpiration(DbUser user);
+
+  void checkCreditsExpirationForUserIDs(List<Long> userIdsList);
+
+  boolean haveCreditsExpired(DbUser user);
 }