History

collab: Use `StripeClient` to retrieve prices and meters from Stripe (#31624)

This PR updates `StripeBilling` to use the `StripeClient` trait to retrieve prices and meters from Stripe instead of using the `stripe::Client` directly. Release Notes: - N/A

Marshall Bowers committed 1 year ago

Commit 75e69a51 parent 05afe956 files · +228 −49Browse files

Unified · Split

6 files
  1. crates/collab/src/api/billing.rs +4 −2
  2. crates/collab/src/stripe_billing.rs +21 −42
  3. crates/collab/src/stripe_client.rs +31 −1
  4. crates/collab/src/stripe_client/fake_stripe_client.rs +20 −1
  5. crates/collab/src/stripe_client/real_stripe_client.rs +68 −2
  6. crates/collab/src/tests/stripe_billing_tests.rs +84 −1
crates/collab/src/api/billing.rs+4−2
@@ -499,8 +499,10 @@ async fn manage_billing_subscription( · 498 unmodified lines
499 let flow = match body.intent {499 let flow = match body.intent {
500 ManageSubscriptionIntent::ManageSubscription => None,500 ManageSubscriptionIntent::ManageSubscription => None,
501 ManageSubscriptionIntent::UpgradeToPro => {501 ManageSubscriptionIntent::UpgradeToPro => {
502 let zed_pro_price_id = stripe_billing.zed_pro_price_id().await?;502 let zed_pro_price_id: stripe::PriceId =
503 let zed_free_price_id = stripe_billing.zed_free_price_id().await?;503 stripe_billing.zed_pro_price_id().await?.try_into()?;
504 let zed_free_price_id: stripe::PriceId =
505 stripe_billing.zed_free_price_id().await?.try_into()?;
504 506
505 let stripe_subscription =507 let stripe_subscription =
506 Subscription::retrieve(&stripe_client, &subscription_id, &[]).await?;508 Subscription::retrieve(&stripe_client, &subscription_id, &[]).await?;
crates/collab/src/stripe_billing.rs+21−42
@@ -4,14 +4,16 @@ use anyhow::{Context as _, anyhow}; · 3 unmodified lines
4use chrono::Utc;4use chrono::Utc;
5use collections::HashMap;5use collections::HashMap;
6use serde::{Deserialize, Serialize};6use serde::{Deserialize, Serialize};
7use stripe::{PriceId, SubscriptionStatus};7use stripe::SubscriptionStatus;
8use tokio::sync::RwLock;8use tokio::sync::RwLock;
9use uuid::Uuid;9use uuid::Uuid;
10 10
11use crate::Result;11use crate::Result;
12use crate::db::billing_subscription::SubscriptionKind;12use crate::db::billing_subscription::SubscriptionKind;
13use crate::llm::AGENT_EXTENDED_TRIAL_FEATURE_FLAG;13use crate::llm::AGENT_EXTENDED_TRIAL_FEATURE_FLAG;
14use crate::stripe_client::{RealStripeClient, StripeClient, StripeCustomerId};14use crate::stripe_client::{
15 RealStripeClient, StripeClient, StripeCustomerId, StripeMeter, StripePrice, StripePriceId,
16};
15 17
16pub struct StripeBilling {18pub struct StripeBilling {
17 state: RwLock<StripeBillingState>,19 state: RwLock<StripeBillingState>,
@@ -22,8 +24,8 @@ pub struct StripeBilling { · 4 unmodified lines
22#[derive(Default)]24#[derive(Default)]
23struct StripeBillingState {25struct StripeBillingState {
24 meters_by_event_name: HashMap<String, StripeMeter>,26 meters_by_event_name: HashMap<String, StripeMeter>,
25 price_ids_by_meter_id: HashMap<String, stripe::PriceId>,27 price_ids_by_meter_id: HashMap<String, StripePriceId>,
26 prices_by_lookup_key: HashMap<String, stripe::Price>,28 prices_by_lookup_key: HashMap<String, StripePrice>,
27}29}
28 30
29impl StripeBilling {31impl StripeBilling {
@@ -50,24 +52,16 @@ impl StripeBilling { · 20 unmodified lines
50 52
51 let mut state = self.state.write().await;53 let mut state = self.state.write().await;
52 54
53 let (meters, prices) = futures::try_join!(55 let (meters, prices) =
54 StripeMeter::list(&self.real_client),56 futures::try_join!(self.client.list_meters(), self.client.list_prices())?;
55 stripe::Price::list(
56 &self.real_client,
57 &stripe::ListPrices {
58 limit: Some(100),
59 ..Default::default()
60 }
61 )
62 )?;
63 57
64 for meter in meters.data {58 for meter in meters {
65 state59 state
66 .meters_by_event_name60 .meters_by_event_name
67 .insert(meter.event_name.clone(), meter);61 .insert(meter.event_name.clone(), meter);
68 }62 }
69 63
70 for price in prices.data {64 for price in prices {
71 if let Some(lookup_key) = price.lookup_key.clone() {65 if let Some(lookup_key) = price.lookup_key.clone() {
72 state.prices_by_lookup_key.insert(lookup_key, price.clone());66 state.prices_by_lookup_key.insert(lookup_key, price.clone());
73 }67 }
@@ -84,15 +78,15 @@ impl StripeBilling { · 10 unmodified lines
84 Ok(())78 Ok(())
85 }79 }
86 80
87 pub async fn zed_pro_price_id(&self) -> Result<PriceId> {81 pub async fn zed_pro_price_id(&self) -> Result<StripePriceId> {
88 self.find_price_id_by_lookup_key("zed-pro").await82 self.find_price_id_by_lookup_key("zed-pro").await
89 }83 }
90 84
91 pub async fn zed_free_price_id(&self) -> Result<PriceId> {85 pub async fn zed_free_price_id(&self) -> Result<StripePriceId> {
92 self.find_price_id_by_lookup_key("zed-free").await86 self.find_price_id_by_lookup_key("zed-free").await
93 }87 }
94 88
95 pub async fn find_price_id_by_lookup_key(&self, lookup_key: &str) -> Result<PriceId> {89 pub async fn find_price_id_by_lookup_key(&self, lookup_key: &str) -> Result<StripePriceId> {
96 self.state90 self.state
97 .read()91 .read()
98 .await92 .await
@@ -102,7 +96,7 @@ impl StripeBilling { · 3 unmodified lines
102 .ok_or_else(|| crate::Error::Internal(anyhow!("no price ID found for {lookup_key:?}")))96 .ok_or_else(|| crate::Error::Internal(anyhow!("no price ID found for {lookup_key:?}")))
103 }97 }
104 98
105 pub async fn find_price_by_lookup_key(&self, lookup_key: &str) -> Result<stripe::Price> {99 pub async fn find_price_by_lookup_key(&self, lookup_key: &str) -> Result<StripePrice> {
106 self.state100 self.state
107 .read()101 .read()
108 .await102 .await
@@ -116,8 +110,10 @@ impl StripeBilling { · 7 unmodified lines
116 &self,110 &self,
117 subscription: &stripe::Subscription,111 subscription: &stripe::Subscription,
118 ) -> Option<SubscriptionKind> {112 ) -> Option<SubscriptionKind> {
119 let zed_pro_price_id = self.zed_pro_price_id().await.ok()?;113 let zed_pro_price_id: stripe::PriceId =
120 let zed_free_price_id = self.zed_free_price_id().await.ok()?;114 self.zed_pro_price_id().await.ok()?.try_into().ok()?;
115 let zed_free_price_id: stripe::PriceId =
116 self.zed_free_price_id().await.ok()?.try_into().ok()?;
121 117
122 subscription.items.data.iter().find_map(|item| {118 subscription.items.data.iter().find_map(|item| {
123 let price = item.price.as_ref()?;119 let price = item.price.as_ref()?;
@@ -171,12 +167,13 @@ impl StripeBilling { · 47 unmodified lines
171 pub async fn subscribe_to_price(167 pub async fn subscribe_to_price(
172 &self,168 &self,
173 subscription_id: &stripe::SubscriptionId,169 subscription_id: &stripe::SubscriptionId,
174 price: &stripe::Price,170 price: &StripePrice,
175 ) -> Result<()> {171 ) -> Result<()> {
176 let subscription =172 let subscription =
177 stripe::Subscription::retrieve(&self.real_client, &subscription_id, &[]).await?;173 stripe::Subscription::retrieve(&self.real_client, &subscription_id, &[]).await?;
178 174
179 if subscription_contains_price(&subscription, &price.id) {175 let price_id = price.id.clone().try_into()?;
176 if subscription_contains_price(&subscription, &price_id) {
180 return Ok(());177 return Ok(());
181 }178 }
182 179
@@ -375,24 +372,6 @@ impl StripeBilling { · 192 unmodified lines
375 }372 }
376}373}
377 374
378#[derive(Clone, Deserialize)]
379struct StripeMeter {
380 id: String,
381 event_name: String,
382}
383
384impl StripeMeter {
385 pub fn list(client: &stripe::Client) -> stripe::Response<stripe::List<Self>> {
386 #[derive(Serialize)]
387 struct Params {
388 #[serde(skip_serializing_if = "Option::is_none")]
389 limit: Option<u64>,
390 }
391
392 client.get_query("/billing/meters", Params { limit: Some(100) })
393 }
394}
395
396#[derive(Deserialize)]375#[derive(Deserialize)]
397struct StripeMeterEvent {376struct StripeMeterEvent {
398 identifier: String,377 identifier: String,
crates/collab/src/stripe_client.rs+31−1
@@ -10,8 +10,9 @@ use async_trait::async_trait; · 9 unmodified lines
10#[cfg(test)]10#[cfg(test)]
11pub use fake_stripe_client::*;11pub use fake_stripe_client::*;
12pub use real_stripe_client::*;12pub use real_stripe_client::*;
13use serde::Deserialize;
13 14
14#[derive(Debug, PartialEq, Eq, Hash, Clone)]15#[derive(Debug, PartialEq, Eq, Hash, Clone, derive_more::Display)]
15pub struct StripeCustomerId(pub Arc<str>);16pub struct StripeCustomerId(pub Arc<str>);
16 17
17#[derive(Debug, Clone)]18#[derive(Debug, Clone)]
@@ -25,9 +26,38 @@ pub struct CreateCustomerParams<'a> { · 7 unmodified lines
25 pub email: Option<&'a str>,26 pub email: Option<&'a str>,
26}27}
27 28
29#[derive(Debug, PartialEq, Eq, Hash, Clone, derive_more::Display)]
30pub struct StripePriceId(pub Arc<str>);
31
32#[derive(Debug, Clone)]
33pub struct StripePrice {
34 pub id: StripePriceId,
35 pub unit_amount: Option<i64>,
36 pub lookup_key: Option<String>,
37 pub recurring: Option<StripePriceRecurring>,
38}
39
40#[derive(Debug, Clone)]
41pub struct StripePriceRecurring {
42 pub meter: Option<String>,
43}
44
45#[derive(Debug, PartialEq, Eq, Hash, Clone, derive_more::Display, Deserialize)]
46pub struct StripeMeterId(pub Arc<str>);
47
48#[derive(Debug, Clone, Deserialize)]
49pub struct StripeMeter {
50 pub id: StripeMeterId,
51 pub event_name: String,
52}
53
28#[async_trait]54#[async_trait]
29pub trait StripeClient: Send + Sync {55pub trait StripeClient: Send + Sync {
30 async fn list_customers_by_email(&self, email: &str) -> Result<Vec<StripeCustomer>>;56 async fn list_customers_by_email(&self, email: &str) -> Result<Vec<StripeCustomer>>;
31 57
32 async fn create_customer(&self, params: CreateCustomerParams<'_>) -> Result<StripeCustomer>;58 async fn create_customer(&self, params: CreateCustomerParams<'_>) -> Result<StripeCustomer>;
59
60 async fn list_prices(&self) -> Result<Vec<StripePrice>>;
61
62 async fn list_meters(&self) -> Result<Vec<StripeMeter>>;
33}63}
crates/collab/src/stripe_client/fake_stripe_client.rs+20−1
@@ -6,16 +6,23 @@ use collections::HashMap; · 5 unmodified lines
6use parking_lot::Mutex;6use parking_lot::Mutex;
7use uuid::Uuid;7use uuid::Uuid;
8 8
9use crate::stripe_client::{CreateCustomerParams, StripeClient, StripeCustomer, StripeCustomerId};9use crate::stripe_client::{
10 CreateCustomerParams, StripeClient, StripeCustomer, StripeCustomerId, StripeMeter,
11 StripeMeterId, StripePrice, StripePriceId,
12};
10 13
11pub struct FakeStripeClient {14pub struct FakeStripeClient {
12 pub customers: Arc<Mutex<HashMap<StripeCustomerId, StripeCustomer>>>,15 pub customers: Arc<Mutex<HashMap<StripeCustomerId, StripeCustomer>>>,
16 pub prices: Arc<Mutex<HashMap<StripePriceId, StripePrice>>>,
17 pub meters: Arc<Mutex<HashMap<StripeMeterId, StripeMeter>>>,
13}18}
14 19
15impl FakeStripeClient {20impl FakeStripeClient {
16 pub fn new() -> Self {21 pub fn new() -> Self {
17 Self {22 Self {
18 customers: Arc::new(Mutex::new(HashMap::default())),23 customers: Arc::new(Mutex::new(HashMap::default())),
24 prices: Arc::new(Mutex::new(HashMap::default())),
25 meters: Arc::new(Mutex::new(HashMap::default())),
19 }26 }
20 }27 }
21}28}
@@ -44,4 +51,16 @@ impl StripeClient for FakeStripeClient { · 22 unmodified lines
44 51
45 Ok(customer)52 Ok(customer)
46 }53 }
54
55 async fn list_prices(&self) -> Result<Vec<StripePrice>> {
56 let prices = self.prices.lock().values().cloned().collect();
57
58 Ok(prices)
59 }
60
61 async fn list_meters(&self) -> Result<Vec<StripeMeter>> {
62 let meters = self.meters.lock().values().cloned().collect();
63
64 Ok(meters)
65 }
47}66}
crates/collab/src/stripe_client/real_stripe_client.rs+68−2
@@ -3,9 +3,13 @@ use std::sync::Arc; · 2 unmodified lines
3 3
4use anyhow::{Context as _, Result};4use anyhow::{Context as _, Result};
5use async_trait::async_trait;5use async_trait::async_trait;
6use stripe::{CreateCustomer, Customer, CustomerId, ListCustomers};6use serde::Serialize;
7use stripe::{CreateCustomer, Customer, CustomerId, ListCustomers, Price, PriceId, Recurring};
7 8
8use crate::stripe_client::{CreateCustomerParams, StripeClient, StripeCustomer, StripeCustomerId};9use crate::stripe_client::{
10 CreateCustomerParams, StripeClient, StripeCustomer, StripeCustomerId, StripeMeter, StripePrice,
11 StripePriceId, StripePriceRecurring,
12};
9 13
10pub struct RealStripeClient {14pub struct RealStripeClient {
11 client: Arc<stripe::Client>,15 client: Arc<stripe::Client>,
@@ -48,6 +52,37 @@ impl StripeClient for RealStripeClient { · 36 unmodified lines
48 52
49 Ok(StripeCustomer::from(customer))53 Ok(StripeCustomer::from(customer))
50 }54 }
55
56 async fn list_prices(&self) -> Result<Vec<StripePrice>> {
57 let response = stripe::Price::list(
58 &self.client,
59 &stripe::ListPrices {
60 limit: Some(100),
61 ..Default::default()
62 },
63 )
64 .await?;
65
66 Ok(response.data.into_iter().map(StripePrice::from).collect())
67 }
68
69 async fn list_meters(&self) -> Result<Vec<StripeMeter>> {
70 #[derive(Serialize)]
71 struct Params {
72 #[serde(skip_serializing_if = "Option::is_none")]
73 limit: Option<u64>,
74 }
75
76 let response = self
77 .client
78 .get_query::<stripe::List<StripeMeter>, _>(
79 "/billing/meters",
80 Params { limit: Some(100) },
81 )
82 .await?;
83
84 Ok(response.data)
85 }
51}86}
52 87
53impl From<CustomerId> for StripeCustomerId {88impl From<CustomerId> for StripeCustomerId {
@@ -72,3 +107,34 @@ impl From<Customer> for StripeCustomer { · 18 unmodified lines
72 }107 }
73 }108 }
74}109}
110
111impl From<PriceId> for StripePriceId {
112 fn from(value: PriceId) -> Self {
113 Self(value.as_str().into())
114 }
115}
116
117impl TryFrom<StripePriceId> for PriceId {
118 type Error = anyhow::Error;
119
120 fn try_from(value: StripePriceId) -> Result<Self, Self::Error> {
121 Self::from_str(value.0.as_ref()).context("failed to parse Stripe price ID")
122 }
123}
124
125impl From<Price> for StripePrice {
126 fn from(value: Price) -> Self {
127 Self {
128 id: value.id.into(),
129 unit_amount: value.unit_amount,
130 lookup_key: value.lookup_key,
131 recurring: value.recurring.map(StripePriceRecurring::from),
132 }
133 }
134}
135
136impl From<Recurring> for StripePriceRecurring {
137 fn from(value: Recurring) -> Self {
138 Self { meter: value.meter }
139 }
140}
crates/collab/src/tests/stripe_billing_tests.rs+84−1
@@ -3,7 +3,9 @@ use std::sync::Arc; · 2 unmodified lines
3use pretty_assertions::assert_eq;3use pretty_assertions::assert_eq;
4 4
5use crate::stripe_billing::StripeBilling;5use crate::stripe_billing::StripeBilling;
6use crate::stripe_client::FakeStripeClient;6use crate::stripe_client::{
7 FakeStripeClient, StripeMeter, StripeMeterId, StripePrice, StripePriceId, StripePriceRecurring,
8};
7 9
8fn make_stripe_billing() -> (StripeBilling, Arc<FakeStripeClient>) {10fn make_stripe_billing() -> (StripeBilling, Arc<FakeStripeClient>) {
9 let stripe_client = Arc::new(FakeStripeClient::new());11 let stripe_client = Arc::new(FakeStripeClient::new());
@@ -12,6 +14,87 @@ fn make_stripe_billing() -> (StripeBilling, Arc<FakeStripeClient>) { · 2 unmodified lines
12 (stripe_billing, stripe_client)14 (stripe_billing, stripe_client)
13}15}
14 16
17#[gpui::test]
18async fn test_initialize() {
19 let (stripe_billing, stripe_client) = make_stripe_billing();
20
21 // Add test meters
22 let meter1 = StripeMeter {
23 id: StripeMeterId("meter_1".into()),
24 event_name: "event_1".to_string(),
25 };
26 let meter2 = StripeMeter {
27 id: StripeMeterId("meter_2".into()),
28 event_name: "event_2".to_string(),
29 };
30 stripe_client
31 .meters
32 .lock()
33 .insert(meter1.id.clone(), meter1);
34 stripe_client
35 .meters
36 .lock()
37 .insert(meter2.id.clone(), meter2);
38
39 // Add test prices
40 let price1 = StripePrice {
41 id: StripePriceId("price_1".into()),
42 unit_amount: Some(1_000),
43 lookup_key: Some("zed-pro".to_string()),
44 recurring: None,
45 };
46 let price2 = StripePrice {
47 id: StripePriceId("price_2".into()),
48 unit_amount: Some(0),
49 lookup_key: Some("zed-free".to_string()),
50 recurring: None,
51 };
52 let price3 = StripePrice {
53 id: StripePriceId("price_3".into()),
54 unit_amount: Some(500),
55 lookup_key: None,
56 recurring: Some(StripePriceRecurring {
57 meter: Some("meter_1".to_string()),
58 }),
59 };
60 stripe_client
61 .prices
62 .lock()
63 .insert(price1.id.clone(), price1);
64 stripe_client
65 .prices
66 .lock()
67 .insert(price2.id.clone(), price2);
68 stripe_client
69 .prices
70 .lock()
71 .insert(price3.id.clone(), price3);
72
73 // Initialize the billing system
74 stripe_billing.initialize().await.unwrap();
75
76 // Verify that prices can be found by lookup key
77 let zed_pro_price_id = stripe_billing.zed_pro_price_id().await.unwrap();
78 assert_eq!(zed_pro_price_id.to_string(), "price_1");
79
80 let zed_free_price_id = stripe_billing.zed_free_price_id().await.unwrap();
81 assert_eq!(zed_free_price_id.to_string(), "price_2");
82
83 // Verify that a price can be found by lookup key
84 let zed_pro_price = stripe_billing
85 .find_price_by_lookup_key("zed-pro")
86 .await
87 .unwrap();
88 assert_eq!(zed_pro_price.id.to_string(), "price_1");
89 assert_eq!(zed_pro_price.unit_amount, Some(1_000));
90
91 // Verify that finding a non-existent lookup key returns an error
92 let result = stripe_billing
93 .find_price_by_lookup_key("non-existent")
94 .await;
95 assert!(result.is_err());
96}
97
15#[gpui::test]98#[gpui::test]
16async fn test_find_or_create_customer_by_email() {99async fn test_find_or_create_customer_by_email() {
17 let (stripe_billing, stripe_client) = make_stripe_billing();100 let (stripe_billing, stripe_client) = make_stripe_billing();