bugfix: Shopify resubscribe and cancel UI + backend fixes

cdxker · cdxker · commit b11cd54de6d6 · 2025-03-31T16:55:50.000-07:00
diff --git a/frontends/dashboard/src/components/PlansTable.tsx b/frontends/dashboard/src/components/PlansTable.tsx
@@ -407,19 +407,24 @@ export const PlansTable = (props: PlansTableProps) => {
               <For each={availablePlansWithCurrent()}>
                 {(plan) => {
                   const curPlan = currentPlan();
-                  const isUpgrade = curPlan
-                    ? curPlan.amount < plan.amount
-                    : false;
 
                   const currentPeriodEnd = plan.current_period_end;
                   let actionButton = <ActiveTag text="Current Tier" />;
 
                   if (!plan.current || currentPeriodEnd) {
-                    if ((curPlan?.amount ?? 0) > 0 && !currentPeriodEnd) {
+                    if (!currentPeriodEnd) {
                       const onClickFunc = () => {
                         void updatePlan(plan);
                       };
-                      const buttonText = isUpgrade ? "Upgrade" : "Downgrade";
+
+                      const isUpgrade = curPlan
+                        ? curPlan.amount < plan.amount
+                        : true;
+
+                      const buttonText =
+                        isUpgrade || curPlan?.amount == null
+                          ? "Upgrade"
+                          : "Downgrade";
                       actionButton = (
                         <button
                           onClick={onClickFunc}
diff --git a/frontends/dashboard/src/components/UsagePlansTable.tsx b/frontends/dashboard/src/components/UsagePlansTable.tsx
@@ -51,6 +51,23 @@ export const UsagePlansTable = (props: PlansTableProps) => {
     });
   });
 
+  const refetchOrgSubPlan = async () => {
+    const selectedOrgId = props.currentOrgSubPlan?.organization.id ?? "";
+
+    const resp = await fetch(`${apiHost}/organization/${selectedOrgId}`, {
+      credentials: "include",
+      headers: {
+        "TR-Organization": selectedOrgId,
+      },
+    });
+
+    if (resp.ok) {
+      const data = (await resp.json()) as OrganizationAndSubAndPlan;
+      setCurrentSubscription(data.subscription ?? null);
+      setCurrentPlan(data.plan ?? null);
+    }
+  };
+
   const updatePlan = async (plan: StripeUsageBasedPlan) => {
     const resp = await fetch(
       `${apiHost}/stripe/subscription_plan/${
@@ -68,6 +85,8 @@ export const UsagePlansTable = (props: PlansTableProps) => {
     if (resp.ok) {
       setCurrentPlan(plan);
     }
+
+    void refetchOrgSubPlan();
   };
 
   return (
@@ -78,8 +97,10 @@ export const UsagePlansTable = (props: PlansTableProps) => {
         let actionButton = <ActiveTag text="Current Tier" />;
 
         if (plan.id !== curPlan?.id || curSub?.current_period_end != null) {
-          if (curPlan?.id !== "00000000-0000-0000-0000-000000000000") {
-            console.log("We have id", plan.id);
+          if (
+            curPlan?.id !== "00000000-0000-0000-0000-000000000000" &&
+            curSub?.current_period_end == null
+          ) {
             actionButton = (
               <button
                 onClick={() => void updatePlan(plan)}
diff --git a/server/src/data/models.rs b/server/src/data/models.rs
@@ -4128,12 +4128,24 @@ impl TrieveSubscription {
         stripe_subscription: Option<StripeSubscription>,
         stripe_usage_based_subscription: Option<StripeUsageBasedSubscription>,
     ) -> Option<Self> {
-        if let Some(usage_subscription) = stripe_usage_based_subscription {
-            Some(TrieveSubscription::UsageBased(usage_subscription))
-        } else if let Some(subscription) = stripe_subscription {
-            Some(TrieveSubscription::Flat(subscription))
-        } else {
-            None
+        match (stripe_subscription, stripe_usage_based_subscription) {
+            (Some(flat_sub), None) => Some(TrieveSubscription::Flat(flat_sub)),
+            (None, Some(usage_sub)) => Some(TrieveSubscription::UsageBased(usage_sub)),
+            (Some(flat_sub), Some(usage_sub)) => {
+                match (flat_sub.current_period_end, usage_sub.current_period_end) {
+                    (Some(flat_period_end), Some(usage_period_ned)) => {
+                        if flat_period_end > usage_period_ned {
+                            Some(TrieveSubscription::Flat(flat_sub))
+                        } else {
+                            Some(TrieveSubscription::UsageBased(usage_sub))
+                        }
+                    }
+                    (Some(_), None) => Some(TrieveSubscription::UsageBased(usage_sub)),
+                    (None, Some(_)) => Some(TrieveSubscription::Flat(flat_sub)),
+                    (None, None) => None,
+                }
+            }
+            (None, None) => None,
         }
     }
 
@@ -4159,6 +4171,13 @@ impl TrieveSubscription {
             }
         }
     }
+
+    pub fn current_period_end(&self) -> Option<chrono::NaiveDateTime> {
+        match self {
+            TrieveSubscription::Flat(subscription) => subscription.current_period_end,
+            TrieveSubscription::UsageBased(subscription) => subscription.current_period_end,
+        }
+    }
 }
 
 impl From<StripeSubscription> for TrieveSubscription {
diff --git a/server/src/handlers/stripe_handler.rs b/server/src/handlers/stripe_handler.rs
@@ -13,8 +13,8 @@ use crate::{
             create_stripe_plan_query, create_stripe_setup_checkout_session,
             create_stripe_subscription_query, create_usage_based_stripe_payment_link,
             create_usage_stripe_subscription_query, delete_subscription_by_id_query,
-            get_all_plans_query, get_all_usage_plans_query, get_invoices_for_org_query,
-            get_option_subscription_by_organization_id_query,
+            delete_usage_subscription_by_id_query, get_all_plans_query, get_all_usage_plans_query,
+            get_invoices_for_org_query, get_option_subscription_by_organization_id_query,
             get_option_usage_based_subscription_by_organization_id_query,
             get_option_usage_based_subscription_by_subscription_id_query, get_plan_by_id_query,
             get_stripe_client, get_trieve_plan_by_id_query, get_trieve_subscription_by_id_query,
@@ -153,6 +153,21 @@ pub async fn webhook(
 
                             if plan_type == "usage-based" {
                                 log::info!("Creating usage based stripe subscription");
+                                let optional_existing_subscription =
+                                    get_option_usage_based_subscription_by_organization_id_query(
+                                        organization_id,
+                                        optional_subscription_pool,
+                                    )
+                                    .await?;
+
+                                if let Some(existing_subscription) = optional_existing_subscription
+                                {
+                                    delete_usage_subscription_by_id_query(
+                                        existing_subscription.id,
+                                        pool.clone(),
+                                    )
+                                    .await?;
+                                }
                                 // record current usage
                                 let usage =
                                     get_org_usage_by_id_query(organization_id, pool.clone())
@@ -331,21 +346,20 @@ pub async fn direct_to_payment_link(
     query_params: web::Query<CreateDirectPaymentLinkQueryParams>,
     pool: web::Data<Pool>,
 ) -> Result<HttpResponse, actix_web::Error> {
-    let organization_pool = pool.clone();
-    let subscription_pool = pool.clone();
     let subscription_org_id = path_data.organization_id;
 
     let current_subscription =
-        get_option_subscription_by_organization_id_query(subscription_org_id, subscription_pool)
-            .await?;
+        get_trieve_subscription_by_id_query(subscription_org_id, pool.clone())
+            .await
+            .ok();
 
-    if current_subscription.is_some_and(|s| s.current_period_end.is_none()) {
+    if current_subscription.is_some_and(|s| s.current_period_end().is_none()) {
         return Ok(HttpResponse::Conflict().finish());
     }
 
     let organization_id = path_data.organization_id;
     let organization_id_clone = path_data.organization_id;
-    let _org_plan_sub = get_org_from_id_query(organization_id_clone, organization_pool).await?;
+    let _org_plan_sub = get_org_from_id_query(organization_id_clone, pool.clone()).await?;
     let plan_id = path_data.plan_id;
 
     let payment_link = match query_params.usage_based {
@@ -440,13 +454,13 @@ pub async fn update_subscription_plan(
     let trieve_subscription =
         get_trieve_subscription_by_id_query(current_subscription_id, pool.clone()).await?;
 
-    let trieve_plan = get_trieve_plan_by_id_query(new_plan_id, pool.clone()).await?;
+    let new_trieve_plan = get_trieve_plan_by_id_query(new_plan_id, pool.clone()).await?;
 
     if !verify_owner(&user, &trieve_subscription.organization_id()) {
         return Err(ServiceError::Forbidden.into());
     };
 
-    match trieve_plan {
+    match new_trieve_plan {
         TrievePlan::Flat(flat_plan) => {
             update_to_flat_stripe_subscription(
                 trieve_subscription.stripe_subscription_id(),
diff --git a/server/src/operators/stripe_operator.rs b/server/src/operators/stripe_operator.rs
@@ -1,9 +1,12 @@
 use std::{collections::HashMap, str::FromStr};
 
 use crate::{
-    data::models::{
-        DateRange, OrganizationUsageCount, Pool, StripeInvoice, StripePlan, StripeSubscription,
-        StripeUsageBasedPlan, StripeUsageBasedSubscription, TrievePlan, TrieveSubscription,
+    data::{
+        models::{
+            DateRange, OrganizationUsageCount, Pool, StripeInvoice, StripePlan, StripeSubscription,
+            StripeUsageBasedPlan, StripeUsageBasedSubscription, TrievePlan, TrieveSubscription,
+        },
+        schema::stripe_usage_based_subscriptions,
     },
     errors::ServiceError,
     get_env,
@@ -478,6 +481,30 @@ pub async fn get_trieve_subscription_by_id_query(
         .ok_or(ServiceError::NotFound("Subscription not found".to_string()))
 }
 
+pub async fn delete_usage_subscription_by_id_query(
+    subscription_id: uuid::Uuid,
+    pool: web::Data<Pool>,
+) -> Result<(), ServiceError> {
+    use crate::data::schema::stripe_usage_based_subscriptions::dsl as stripe_usage_based_subscriptions_columns;
+
+    let mut conn = pool
+        .get()
+        .await
+        .expect("Failed to get connection from pool");
+    diesel::delete(
+        stripe_usage_based_subscriptions_columns::stripe_usage_based_subscriptions
+            .filter(stripe_usage_based_subscriptions::id.eq(subscription_id)),
+    )
+    .execute(&mut conn)
+    .await
+    .map_err(|e| {
+        log::error!("Failed to delete stripe subscription: {}", e);
+        ServiceError::BadRequest("Failed to delete stripe subscription".to_string())
+    })?;
+
+    Ok(())
+}
+
 pub async fn delete_subscription_by_id_query(
     subscription_id: uuid::Uuid,
     pool: web::Data<Pool>,
