Skip to content

Commit 27a64ef

Browse files
zlwaterfieldzlwaterfield
authored
fix: reload available features upgraded (PostHog#29161)
1 parent 9a03705 commit 27a64ef

File tree

5 files changed

+48
-16
lines changed

5 files changed

+48
-16
lines changed

‎ee/api/billing.py‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -60,12 +60,11 @@ def list(self, request: Request, *args: Any, **kwargs: Any) -> Response:
6060
if org.billing.stripe_subscription_id: # type: ignore
6161
raise NotFound("Billing V1 is active for this organization")
6262

63-
plan_keys = request.query_params.get("plan_keys", None)
6463
billing_manager = self.get_billing_manager()
6564
query = {}
6665
if "include_forecasting" in request.query_params:
6766
query["include_forecasting"] = request.query_params.get("include_forecasting")
68-
response = billing_manager.get_billing(org, plan_keys, query)
67+
response = billing_manager.get_billing(org, query)
6968

7069
return Response(response)
7170

‎ee/api/test/test_billing.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -987,7 +987,7 @@ def test_deactivate_success(self, mock_get_billing, mock_deactivate_products):
987987

988988
self.assertEqual(response.status_code, status.HTTP_200_OK)
989989
mock_deactivate_products.assert_called_once_with(self.organization, "product_1")
990-
mock_get_billing.assert_called_once_with(self.organization, None, {})
990+
mock_get_billing.assert_called_once_with(self.organization, {})
991991

992992
def test_deactivate_failure(self):
993993
url = "/api/billing/deactivate"

‎ee/billing/billing_manager.py‎

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -77,7 +77,6 @@ def __init__(self, license, user: Optional[User] = None):
7777
def get_billing(
7878
self,
7979
organization: Optional[Organization],
80-
plan_keys: Optional[str],
8180
query_params: Optional[dict[str, Any]] = None,
8281
) -> dict[str, Any]:
8382
if organization and self.license and self.license.is_v2_license:
@@ -141,6 +140,20 @@ def update_billing(self, organization: Organization, data: dict[str, Any]) -> No
141140

142141
handle_billing_service_error(res)
143142

143+
def update_available_product_features(self, organization: Organization) -> list[dict[str, Any]]:
144+
res = requests.get(
145+
f"{BILLING_SERVICE_URL}/api/billing/available_product_features",
146+
headers=self.get_auth_headers(organization),
147+
)
148+
149+
handle_billing_service_error(res)
150+
151+
available_product_features = res.json()
152+
organization.available_product_features = available_product_features
153+
organization.save()
154+
155+
return available_product_features
156+
144157
def update_billing_organization_users(self, organization: Organization) -> None:
145158
try:
146159
distinct_ids = list(organization.members.values_list("distinct_id", flat=True)) # type: ignore
@@ -394,6 +407,8 @@ def activate_trial(self, organization: Organization, data: dict[str, Any]):
394407

395408
handle_billing_service_error(res)
396409

410+
self.update_available_product_features(organization)
411+
397412
return res.json()
398413

399414
def cancel_trial(self, organization: Organization, data: dict[str, Any]):
@@ -405,6 +420,8 @@ def cancel_trial(self, organization: Organization, data: dict[str, Any]):
405420

406421
handle_billing_service_error(res)
407422

423+
self.update_available_product_features(organization)
424+
408425
def authorize(self, organization: Organization):
409426
res = requests.post(
410427
f"{BILLING_SERVICE_URL}/api/activate/authorize",

‎ee/billing/test/test_billing_manager.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,7 @@ def test_get_billing_unlicensed(self, billing_patch_request_mock):
4848
organization = self.organization
4949
TEST_clear_instance_license_cache()
5050

51-
BillingManager(license=None).get_billing(organization, plan_keys=None)
51+
BillingManager(license=None).get_billing(organization)
5252
assert billing_patch_request_mock.call_count == 1
5353
billing_patch_request_mock.assert_called_with(
5454
"https://billing.posthog.com/api/products-v2", params={"plan": "standard"}, headers={}

‎frontend/src/scenes/billing/billingLogic.tsx‎

Lines changed: 27 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ import { featureFlagLogic } from 'lib/logic/featureFlagLogic'
1313
import { pluralize } from 'lib/utils'
1414
import { eventUsageLogic } from 'lib/utils/eventUsageLogic'
1515
import posthog from 'posthog-js'
16+
import { organizationLogic } from 'scenes/organizationLogic'
1617
import { preflightLogic } from 'scenes/PreflightCheck/preflightLogic'
1718
import { userLogic } from 'scenes/userLogic'
1819

@@ -105,6 +106,8 @@ export const billingLogic = kea<billingLogicType>([
105106
actions: [
106107
userLogic,
107108
['loadUser'],
109+
organizationLogic,
110+
['loadCurrentOrganization'],
108111
eventUsageLogic,
109112
['reportProductUnsubscribed'],
110113
lemonBannerLogic({ dismissKey: 'usage-limit-exceeded' }),
@@ -237,7 +240,7 @@ export const billingLogic = kea<billingLogicType>([
237240
}
238241
},
239242

240-
deactivateProduct: async (key: string) => {
243+
deactivateProduct: async (key: string, breakpoint) => {
241244
// clear upgrade params from URL
242245
// Note(@zach): This is not working properly. We need to look into this.
243246
const currentURL = new URL(window.location.href)
@@ -255,6 +258,12 @@ export const billingLogic = kea<billingLogicType>([
255258
)
256259
actions.reportProductUnsubscribed(key)
257260

261+
// Reload billing, user, and organization to get the updated available features
262+
actions.loadBilling()
263+
await breakpoint(2000) // Wait enough time for the organization to be updated
264+
actions.loadUser()
265+
actions.loadCurrentOrganization()
266+
258267
return parseBillingResponse(jsonRes)
259268
} catch (error: any) {
260269
if (error.code) {
@@ -552,20 +561,27 @@ export const billingLogic = kea<billingLogicType>([
552561
posthog.capture('credits cta hero dismissed')
553562
}
554563
},
555-
loadBillingSuccess: () => {
564+
loadBillingSuccess: async (_, breakpoint) => {
565+
actions.registerInstrumentationProps()
566+
actions.determineBillingAlert()
567+
actions.loadCreditOverview()
568+
569+
// If the activation is successful, we reload the user/organization to get the updated available features
570+
// activation can be triggered from the billing page or onboarding
556571
if (
557-
router.values.location.pathname.includes('/organization/billing') &&
558-
router.values.searchParams['success']
572+
(router.values.location.pathname.includes('/organization/billing') ||
573+
router.values.location.pathname.includes('/onboarding')) &&
574+
(router.values.searchParams['success'] || router.values.searchParams['upgraded'])
559575
) {
560-
// if the activation is successful, we reload the user to get the updated billing info on the organization
576+
// Wait enough time for the organization to be updated
577+
await breakpoint(1000)
561578
actions.loadUser()
562-
router.actions.replace('/organization/billing')
579+
actions.loadCurrentOrganization()
580+
// Clear the params from the billing page so we don't trigger the activation again
581+
if (router.values.location.pathname.includes('/organization/billing')) {
582+
router.actions.replace('/organization/billing')
583+
}
563584
}
564-
actions.registerInstrumentationProps()
565-
566-
actions.determineBillingAlert()
567-
568-
actions.loadCreditOverview()
569585
},
570586
determineBillingAlert: () => {
571587
if (values.productSpecificAlert) {

0 commit comments

Comments
 (0)