mirror of
https://github.com/bitwarden/server
synced 2026-01-19 00:43:47 +00:00
[PM-30460] update storage job to also update database max storage (#6803)
* [PM-30460] update storage job to also update database max storage * dry run logs * more logging fixes and pr feedback, forgot sql scripts * claude feedback * pr feedback, redesign of entity id reverse lookup * claude feedback
This commit is contained in:
Kyle Denney
@@ -4,6 +4,7 @@ using Bit.Billing.Services;
|
||||
using Bit.Core;
|
||||
using Bit.Core.Billing.Constants;
|
||||
using Bit.Core.Jobs;
|
||||
using Bit.Core.Repositories;
|
||||
using Bit.Core.Services;
|
||||
using Quartz;
|
||||
using Stripe;
|
||||
@@ -13,12 +14,23 @@ namespace Bit.Billing.Jobs;
|
||||
public class ReconcileAdditionalStorageJob(
|
||||
IStripeFacade stripeFacade,
|
||||
ILogger<ReconcileAdditionalStorageJob> logger,
|
||||
IFeatureService featureService) : BaseJob(logger)
|
||||
IFeatureService featureService,
|
||||
IUserRepository userRepository,
|
||||
IOrganizationRepository organizationRepository,
|
||||
IStripeEventUtilityService stripeEventUtilityService) : BaseJob(logger)
|
||||
{
|
||||
private const string _storageGbMonthlyPriceId = "storage-gb-monthly";
|
||||
private const string _storageGbAnnuallyPriceId = "storage-gb-annually";
|
||||
private const string _personalStorageGbAnnuallyPriceId = "personal-storage-gb-annually";
|
||||
private const int _storageGbToRemove = 4;
|
||||
private const short _includedStorageGb = 5;
|
||||
|
||||
public enum SubscriptionPlanTier
|
||||
{
|
||||
Personal,
|
||||
Organization,
|
||||
Unknown
|
||||
}
|
||||
|
||||
protected override async Task ExecuteJobAsync(IJobExecutionContext context)
|
||||
{
|
||||
@@ -34,6 +46,7 @@ public class ReconcileAdditionalStorageJob(
|
||||
var subscriptionsFound = 0;
|
||||
var subscriptionsUpdated = 0;
|
||||
var subscriptionsWithErrors = 0;
|
||||
var databaseUpdatesFailed = 0;
|
||||
var failures = new List<string>();
|
||||
|
||||
logger.LogInformation("Starting ReconcileAdditionalStorageJob (live mode: {LiveMode})", liveMode);
|
||||
@@ -51,11 +64,13 @@ public class ReconcileAdditionalStorageJob(
|
||||
{
|
||||
logger.LogWarning(
|
||||
"Job cancelled!! Exiting. Progress at time of cancellation: Subscriptions found: {SubscriptionsFound}, " +
|
||||
"Updated: {SubscriptionsUpdated}, Errors: {SubscriptionsWithErrors}{Failures}",
|
||||
"Stripe updates: {StripeUpdates}, Database updates: {DatabaseFailed} failed, " +
|
||||
"Errors: {SubscriptionsWithErrors}{Failures}",
|
||||
subscriptionsFound,
|
||||
liveMode
|
||||
? subscriptionsUpdated
|
||||
: $"(In live mode, would have updated) {subscriptionsUpdated}",
|
||||
databaseUpdatesFailed,
|
||||
subscriptionsWithErrors,
|
||||
failures.Count > 0
|
||||
? $", Failures: {Environment.NewLine}{string.Join(Environment.NewLine, failures)}"
|
||||
@@ -99,20 +114,68 @@ public class ReconcileAdditionalStorageJob(
|
||||
|
||||
subscriptionsUpdated++;
|
||||
|
||||
if (!liveMode)
|
||||
// Now, prepare the database update so we can log details out if not in live mode
|
||||
var (organizationId, userId, _) = stripeEventUtilityService.GetIdsFromMetadata(subscription.Metadata ?? new Dictionary<string, string>());
|
||||
var subscriptionPlanTier = DetermineSubscriptionPlanTier(userId, organizationId);
|
||||
|
||||
if (subscriptionPlanTier == SubscriptionPlanTier.Unknown)
|
||||
{
|
||||
logger.LogInformation(
|
||||
"Not live mode (dry-run): Would have updated subscription {SubscriptionId} with item changes: {NewLine}{UpdateOptions}",
|
||||
subscription.Id,
|
||||
Environment.NewLine,
|
||||
JsonSerializer.Serialize(updateOptions));
|
||||
logger.LogError(
|
||||
"Cannot determine subscription plan tier for {SubscriptionId}. Skipping subscription. ",
|
||||
subscription.Id);
|
||||
subscriptionsWithErrors++;
|
||||
continue;
|
||||
}
|
||||
|
||||
var entityId =
|
||||
subscriptionPlanTier switch
|
||||
{
|
||||
SubscriptionPlanTier.Personal => userId!.Value,
|
||||
SubscriptionPlanTier.Organization => organizationId!.Value,
|
||||
_ => throw new ArgumentOutOfRangeException(nameof(subscriptionPlanTier), subscriptionPlanTier, null)
|
||||
};
|
||||
|
||||
// Calculate new MaxStorageGb
|
||||
var currentStorageQuantity = GetCurrentStorageQuantityFromSubscription(subscription, priceId);
|
||||
var newMaxStorageGb = CalculateNewMaxStorageGb(currentStorageQuantity, updateOptions);
|
||||
|
||||
if (!liveMode)
|
||||
{
|
||||
logger.LogInformation(
|
||||
"Not live mode (dry-run): Would have updated subscription {SubscriptionId} with item changes: {NewLine}{UpdateOptions}" +
|
||||
"{NewLine2}And would have updated database record tier: {Tier} to new MaxStorageGb: {MaxStorageGb}",
|
||||
subscription.Id,
|
||||
Environment.NewLine,
|
||||
JsonSerializer.Serialize(updateOptions),
|
||||
Environment.NewLine,
|
||||
subscriptionPlanTier,
|
||||
newMaxStorageGb);
|
||||
continue;
|
||||
}
|
||||
|
||||
// Live mode enabled - continue with updates to stripe and database
|
||||
try
|
||||
{
|
||||
await stripeFacade.UpdateSubscription(subscription.Id, updateOptions);
|
||||
logger.LogInformation("Successfully updated subscription: {SubscriptionId}", subscription.Id);
|
||||
logger.LogInformation("Successfully updated Stripe subscription: {SubscriptionId}", subscription.Id);
|
||||
|
||||
logger.LogInformation(
|
||||
"Updating MaxStorageGb in database for subscription {SubscriptionId} ({Type}): New MaxStorageGb: {MaxStorage}",
|
||||
subscription.Id,
|
||||
subscriptionPlanTier,
|
||||
newMaxStorageGb);
|
||||
|
||||
var dbUpdateSuccess = await UpdateDatabaseMaxStorageAsync(
|
||||
subscriptionPlanTier,
|
||||
entityId,
|
||||
newMaxStorageGb,
|
||||
subscription.Id);
|
||||
|
||||
if (!dbUpdateSuccess)
|
||||
{
|
||||
databaseUpdatesFailed++;
|
||||
failures.Add($"Subscription {subscription.Id}: Database update failed");
|
||||
}
|
||||
}
|
||||
catch (Exception ex)
|
||||
{
|
||||
@@ -125,12 +188,14 @@ public class ReconcileAdditionalStorageJob(
|
||||
}
|
||||
|
||||
logger.LogInformation(
|
||||
"ReconcileAdditionalStorageJob completed. Subscriptions found: {SubscriptionsFound}, " +
|
||||
"Updated: {SubscriptionsUpdated}, Errors: {SubscriptionsWithErrors}{Failures}",
|
||||
"ReconcileAdditionalStorageJob FINISHED. Subscriptions found: {SubscriptionsFound}, " +
|
||||
"Subscriptions updated: {SubscriptionsUpdated}, Database failures: {DatabaseFailed}, " +
|
||||
"Total Subscriptions With Errors: {SubscriptionsWithErrors}{Failures}",
|
||||
subscriptionsFound,
|
||||
liveMode
|
||||
? subscriptionsUpdated
|
||||
: $"(In live mode, would have updated) {subscriptionsUpdated}",
|
||||
databaseUpdatesFailed,
|
||||
subscriptionsWithErrors,
|
||||
failures.Count > 0
|
||||
? $", Failures: {Environment.NewLine}{string.Join(Environment.NewLine, failures)}"
|
||||
@@ -182,6 +247,117 @@ public class ReconcileAdditionalStorageJob(
|
||||
return hasUpdates ? updateOptions : null;
|
||||
}
|
||||
|
||||
public SubscriptionPlanTier DetermineSubscriptionPlanTier(
|
||||
Guid? userId,
|
||||
Guid? organizationId)
|
||||
{
|
||||
return userId.HasValue
|
||||
? SubscriptionPlanTier.Personal
|
||||
: organizationId.HasValue
|
||||
? SubscriptionPlanTier.Organization
|
||||
: SubscriptionPlanTier.Unknown;
|
||||
}
|
||||
|
||||
public long GetCurrentStorageQuantityFromSubscription(
|
||||
Subscription subscription,
|
||||
string storagePriceId)
|
||||
{
|
||||
return subscription.Items?.Data?.FirstOrDefault(item => item?.Price?.Id == storagePriceId)?.Quantity ?? 0;
|
||||
}
|
||||
|
||||
public short CalculateNewMaxStorageGb(
|
||||
long currentQuantity,
|
||||
SubscriptionUpdateOptions? updateOptions)
|
||||
{
|
||||
if (updateOptions?.Items == null)
|
||||
{
|
||||
return (short)(_includedStorageGb + currentQuantity);
|
||||
}
|
||||
|
||||
// If the update marks item as deleted, new quantity is whatever the base storage gb
|
||||
if (updateOptions.Items.Any(i => i.Deleted == true))
|
||||
{
|
||||
return _includedStorageGb;
|
||||
}
|
||||
|
||||
// If the update has a new quantity, use it to calculate the new max
|
||||
var updatedItem = updateOptions.Items.FirstOrDefault(i => i.Quantity.HasValue);
|
||||
if (updatedItem?.Quantity != null)
|
||||
{
|
||||
return (short)(_includedStorageGb + updatedItem.Quantity.Value);
|
||||
}
|
||||
|
||||
// Otherwise, no change
|
||||
return (short)(_includedStorageGb + currentQuantity);
|
||||
}
|
||||
|
||||
public async Task<bool> UpdateDatabaseMaxStorageAsync(
|
||||
SubscriptionPlanTier subscriptionPlanTier,
|
||||
Guid entityId,
|
||||
short newMaxStorageGb,
|
||||
string subscriptionId)
|
||||
{
|
||||
try
|
||||
{
|
||||
switch (subscriptionPlanTier)
|
||||
{
|
||||
case SubscriptionPlanTier.Personal:
|
||||
{
|
||||
var user = await userRepository.GetByIdAsync(entityId);
|
||||
if (user == null)
|
||||
{
|
||||
logger.LogError(
|
||||
"User not found for subscription {SubscriptionId}. Database not updated.",
|
||||
subscriptionId);
|
||||
return false;
|
||||
}
|
||||
|
||||
user.MaxStorageGb = newMaxStorageGb;
|
||||
await userRepository.ReplaceAsync(user);
|
||||
|
||||
logger.LogInformation(
|
||||
"Successfully updated User {UserId} MaxStorageGb to {MaxStorageGb} for subscription {SubscriptionId}",
|
||||
user.Id,
|
||||
newMaxStorageGb,
|
||||
subscriptionId);
|
||||
return true;
|
||||
}
|
||||
case SubscriptionPlanTier.Organization:
|
||||
{
|
||||
var organization = await organizationRepository.GetByIdAsync(entityId);
|
||||
if (organization == null)
|
||||
{
|
||||
logger.LogError(
|
||||
"Organization not found for subscription {SubscriptionId}. Database not updated.",
|
||||
subscriptionId);
|
||||
return false;
|
||||
}
|
||||
|
||||
organization.MaxStorageGb = newMaxStorageGb;
|
||||
await organizationRepository.ReplaceAsync(organization);
|
||||
|
||||
logger.LogInformation(
|
||||
"Successfully updated Organization {OrganizationId} MaxStorageGb to {MaxStorageGb} for subscription {SubscriptionId}",
|
||||
organization.Id,
|
||||
newMaxStorageGb,
|
||||
subscriptionId);
|
||||
return true;
|
||||
}
|
||||
case SubscriptionPlanTier.Unknown:
|
||||
default:
|
||||
return false;
|
||||
}
|
||||
}
|
||||
catch (Exception ex)
|
||||
{
|
||||
logger.LogError(ex,
|
||||
"Failed to update database MaxStorageGb for subscription {SubscriptionId} (Plan Tier: {SubscriptionType})",
|
||||
subscriptionId,
|
||||
subscriptionPlanTier);
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
public static ITrigger GetTrigger()
|
||||
{
|
||||
return TriggerBuilder.Create()
|
||||
|
||||
Reference in New Issue
Block a user