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

crates/collab/src/api/billing.rs +4 −2
crates/collab/src/stripe_billing.rs +21 −42
crates/collab/src/stripe_client.rs +31 −1
crates/collab/src/stripe_client/fake_stripe_client.rs +20 −1
crates/collab/src/stripe_client/real_stripe_client.rs +68 −2
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	499	let flow = match body.intent {
500	500	ManageSubscriptionIntent::ManageSubscription => None,
501	501	ManageSubscriptionIntent::UpgradeToPro => {
502		− let zed_pro_price_id = stripe_billing.zed_pro_price_id().await?;
503		− let zed_free_price_id = stripe_billing.zed_free_price_id().await?;
	502	+ let zed_pro_price_id: stripe::PriceId =
	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	507	let stripe_subscription =
506	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
4	4	use chrono::Utc;
5	5	use collections::HashMap;
6	6	use serde::{Deserialize, Serialize};
7		−use stripe::{PriceId, SubscriptionStatus};
	7	+use stripe::SubscriptionStatus;
8	8	use tokio::sync::RwLock;
9	9	use uuid::Uuid;
10	10
11	11	use crate::Result;
12	12	use crate::db::billing_subscription::SubscriptionKind;
13	13	use crate::llm::AGENT_EXTENDED_TRIAL_FEATURE_FLAG;
14		−use crate::stripe_client::{RealStripeClient, StripeClient, StripeCustomerId};
	14	+use crate::stripe_client::{
	15	+ RealStripeClient, StripeClient, StripeCustomerId, StripeMeter, StripePrice, StripePriceId,
	16	+};
15	17
16	18	pub struct StripeBilling {
17	19	state: RwLock<StripeBillingState>,
@@ -22,8 +24,8 @@ pub struct StripeBilling { · 4 unmodified lines
22	24	#[derive(Default)]
23	25	struct StripeBillingState {
24	26	meters_by_event_name: HashMap<String, StripeMeter>,
25		− price_ids_by_meter_id: HashMap<String, stripe::PriceId>,
26		− prices_by_lookup_key: HashMap<String, stripe::Price>,
	27	+ price_ids_by_meter_id: HashMap<String, StripePriceId>,
	28	+ prices_by_lookup_key: HashMap<String, StripePrice>,
27	29	}
28	30
29	31	impl StripeBilling {
@@ -50,24 +52,16 @@ impl StripeBilling { · 20 unmodified lines
50	52
51	53	let mut state = self.state.write().await;
52	54
53		− let (meters, prices) = futures::try_join!(
54		− StripeMeter::list(&self.real_client),
55		− stripe::Price::list(
56		− &self.real_client,
57		− &stripe::ListPrices {
58		− limit: Some(100),
59		− ..Default::default()
60		− }
61		− )
62		− )?;
	55	+ let (meters, prices) =
	56	+ futures::try_join!(self.client.list_meters(), self.client.list_prices())?;
63	57
64		− for meter in meters.data {
	58	+ for meter in meters {
65	59	state
66	60	.meters_by_event_name
67	61	.insert(meter.event_name.clone(), meter);
68	62	}
69	63
70		− for price in prices.data {
	64	+ for price in prices {
71	65	if let Some(lookup_key) = price.lookup_key.clone() {
72	66	state.prices_by_lookup_key.insert(lookup_key, price.clone());
73	67	}
@@ -84,15 +78,15 @@ impl StripeBilling { · 10 unmodified lines
84	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	82	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	86	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	90	self.state
97	91	.read()
98	92	.await
@@ -102,7 +96,7 @@ impl StripeBilling { · 3 unmodified lines
102	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	100	self.state
107	101	.read()
108	102	.await
@@ -116,8 +110,10 @@ impl StripeBilling { · 7 unmodified lines
116	110	&self,
117	111	subscription: &stripe::Subscription,
118	112	) -> Option<SubscriptionKind> {
119		− let zed_pro_price_id = self.zed_pro_price_id().await.ok()?;
120		− let zed_free_price_id = self.zed_free_price_id().await.ok()?;
	113	+ let zed_pro_price_id: stripe::PriceId =
	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	118	subscription.items.data.iter().find_map(\|item\| {
123	119	let price = item.price.as_ref()?;
@@ -171,12 +167,13 @@ impl StripeBilling { · 47 unmodified lines
171	167	pub async fn subscribe_to_price(
172	168	&self,
173	169	subscription_id: &stripe::SubscriptionId,
174		− price: &stripe::Price,
	170	+ price: &StripePrice,
175	171	) -> Result<()> {
176	172	let subscription =
177	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	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)]
379		−struct StripeMeter {
380		− id: String,
381		− event_name: String,
382		−}
383		−
384		−impl 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	375	#[derive(Deserialize)]
397	376	struct StripeMeterEvent {
398	377	identifier: String,

crates/collab/src/stripe_client.rs+31−1

@@ -10,8 +10,9 @@ use async_trait::async_trait; · 9 unmodified lines
10	10	#[cfg(test)]
11	11	pub use fake_stripe_client::*;
12	12	pub use real_stripe_client::*;
	13	+use serde::Deserialize;
13	14
14		−#[derive(Debug, PartialEq, Eq, Hash, Clone)]
	15	+#[derive(Debug, PartialEq, Eq, Hash, Clone, derive_more::Display)]
15	16	pub struct StripeCustomerId(pub Arc<str>);
16	17
17	18	#[derive(Debug, Clone)]
@@ -25,9 +26,38 @@ pub struct CreateCustomerParams<'a> { · 7 unmodified lines
25	26	pub email: Option<&'a str>,
26	27	}
27	28
	29	+#[derive(Debug, PartialEq, Eq, Hash, Clone, derive_more::Display)]
	30	+pub struct StripePriceId(pub Arc<str>);
	31	+
	32	+#[derive(Debug, Clone)]
	33	+pub 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)]
	41	+pub struct StripePriceRecurring {
	42	+ pub meter: Option<String>,
	43	+}
	44	+
	45	+#[derive(Debug, PartialEq, Eq, Hash, Clone, derive_more::Display, Deserialize)]
	46	+pub struct StripeMeterId(pub Arc<str>);
	47	+
	48	+#[derive(Debug, Clone, Deserialize)]
	49	+pub struct StripeMeter {
	50	+ pub id: StripeMeterId,
	51	+ pub event_name: String,
	52	+}
	53	+
28	54	#[async_trait]
29	55	pub trait StripeClient: Send + Sync {
30	56	async fn list_customers_by_email(&self, email: &str) -> Result<Vec<StripeCustomer>>;
31	57
32	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
6	6	use parking_lot::Mutex;
7	7	use uuid::Uuid;
8	8
9		−use crate::stripe_client::{CreateCustomerParams, StripeClient, StripeCustomer, StripeCustomerId};
	9	+use crate::stripe_client::{
	10	+ CreateCustomerParams, StripeClient, StripeCustomer, StripeCustomerId, StripeMeter,
	11	+ StripeMeterId, StripePrice, StripePriceId,
	12	+};
10	13
11	14	pub struct FakeStripeClient {
12	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
15	20	impl FakeStripeClient {
16	21	pub fn new() -> Self {
17	22	Self {
18	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	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
4	4	use anyhow::{Context as _, Result};
5	5	use async_trait::async_trait;
6		−use stripe::{CreateCustomer, Customer, CustomerId, ListCustomers};
	6	+use serde::Serialize;
	7	+use stripe::{CreateCustomer, Customer, CustomerId, ListCustomers, Price, PriceId, Recurring};
7	8
8		−use crate::stripe_client::{CreateCustomerParams, StripeClient, StripeCustomer, StripeCustomerId};
	9	+use crate::stripe_client::{
	10	+ CreateCustomerParams, StripeClient, StripeCustomer, StripeCustomerId, StripeMeter, StripePrice,
	11	+ StripePriceId, StripePriceRecurring,
	12	+};
9	13
10	14	pub struct RealStripeClient {
11	15	client: Arc<stripe::Client>,
@@ -48,6 +52,37 @@ impl StripeClient for RealStripeClient { · 36 unmodified lines
48	52
49	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
53	88	impl From<CustomerId> for StripeCustomerId {
@@ -72,3 +107,34 @@ impl From<Customer> for StripeCustomer { · 18 unmodified lines
72	107	}
73	108	}
74	109	}
	110	+
	111	+impl From<PriceId> for StripePriceId {
	112	+ fn from(value: PriceId) -> Self {
	113	+ Self(value.as_str().into())
	114	+ }
	115	+}
	116	+
	117	+impl 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	+
	125	+impl 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	+
	136	+impl 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
3	3	use pretty_assertions::assert_eq;
4	4
5	5	use crate::stripe_billing::StripeBilling;
6		−use crate::stripe_client::FakeStripeClient;
	6	+use crate::stripe_client::{
	7	+ FakeStripeClient, StripeMeter, StripeMeterId, StripePrice, StripePriceId, StripePriceRecurring,
	8	+};
7	9
8	10	fn make_stripe_billing() -> (StripeBilling, Arc<FakeStripeClient>) {
9	11	let stripe_client = Arc::new(FakeStripeClient::new());
@@ -12,6 +14,87 @@ fn make_stripe_billing() -> (StripeBilling, Arc<FakeStripeClient>) { · 2 unmodified lines
12	14	(stripe_billing, stripe_client)
13	15	}
14	16
	17	+#[gpui::test]
	18	+async 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	98	#[gpui::test]
16	99	async fn test_find_or_create_customer_by_email() {
17	100	let (stripe_billing, stripe_client) = make_stripe_billing();