From a266ffc60e12f77109558c37abb9cb62a41116d2 Mon Sep 17 00:00:00 2001 From: Bianca Yang <21014901+xrdt@users.noreply.github.com> Date: Tue, 7 May 2024 16:25:22 -0700 Subject: [PATCH] feat: Update tiers usage information with local org usage information (#21957) * update tiers usage with local org usage information * type finagling * update addon tiers as well * test * Update ee/billing/billing_manager.py * update tests --------- Co-authored-by: Raquel Smith --- .vscode/launch.json | 2 +- ee/api/test/test_billing.py | 236 ++++++++++++++++++++++++++++++++-- ee/billing/billing_manager.py | 90 ++++++++++++- ee/billing/billing_types.py | 41 +++++- posthog/hogql/property.py | 16 +-- 5 files changed, 355 insertions(+), 30 deletions(-) diff --git a/.vscode/launch.json b/.vscode/launch.json index f2e15ed84c0d8..99b5d08fcb750 100644 --- a/.vscode/launch.json +++ b/.vscode/launch.json @@ -71,7 +71,7 @@ "DATABASE_URL": "postgres://posthog:posthog@localhost:5432/posthog", "SKIP_SERVICE_VERSION_REQUIREMENTS": "1", "PRINT_SQL": "1", - "BILLING_SERVICE_URL": "https://billing.dev.posthog.dev" + "BILLING_SERVICE_URL": "http://localhost:8100" //"https://billing.dev.posthog.dev" }, "console": "integratedTerminal", "python": "${workspaceFolder}/env/bin/python", diff --git a/ee/api/test/test_billing.py b/ee/api/test/test_billing.py index 94eed34d29d79..22d7a49c6dec0 100644 --- a/ee/api/test/test_billing.py +++ b/ee/api/test/test_billing.py @@ -11,7 +11,12 @@ from rest_framework import status from ee.api.test.base import APILicensedTest -from ee.billing.billing_types import BillingPeriod, CustomerInfo, CustomerProduct +from ee.billing.billing_types import ( + BillingPeriod, + CustomerInfo, + CustomerProduct, + CustomerProductAddon, +) from ee.models.license import License from posthog.cloud_utils import ( TEST_clear_instance_license_cache, @@ -76,7 +81,7 @@ def create_billing_customer(**kwargs) -> CustomerInfo: { "unit_amount_usd": "0.00045", "up_to": 2000000, - "current_amount_usd": None, + "current_amount_usd": "0.00", }, ], tiered=True, @@ -89,6 +94,39 @@ def create_billing_customer(**kwargs) -> CustomerInfo: projected_usage=0, projected_amount_usd="0.00", usage_key="events", + addons=[ + CustomerProductAddon( + name="Addon", + description="Test Addon", + price_description=None, + type="events", + image_url="https://posthog.com/static/images/product-os.png", + free_allocation=10000, + tiers=[ + { + "unit_amount_usd": "0.00", + "up_to": 1000000, + "current_amount_usd": "0.00", + }, + { + "unit_amount_usd": "0.0000135", + "up_to": 2000000, + "current_amount_usd": "0.00", + }, + ], + tiered=True, + unit_amount_usd="0.00", + current_amount_usd="0.00", + current_usage=0, + usage_limit=None, + has_exceeded_limit=False, + percentage_usage=0, + projected_usage=0, + projected_amount_usd="0.00", + usage_key="events", + subscribed=True, + ) + ], ) ], billing_period=BillingPeriod( @@ -121,22 +159,72 @@ def create_billing_products_response(**kwargs) -> dict[str, list[CustomerProduct "unit_amount_usd": "0.00", "up_to": 1000000, "current_amount_usd": "0.00", + "current_usage": 0, + "flat_amount_usd": "0", + "projected_amount_usd": "None", + "projected_usage": None, }, { "unit_amount_usd": "0.00045", "up_to": 2000000, - "current_amount_usd": None, + "current_amount_usd": "0.00", + "current_usage": 0, + "flat_amount_usd": "0", + "projected_amount_usd": "None", + "projected_usage": None, + }, + ], + addons=[ + { + "current_amount_usd": 0.0, + "current_usage": 0, + "description": "Test Addon", + "free_allocation": 10000, + "has_exceeded_limit": False, + "image_url": "https://posthog.com/static/images/product-os.png", + "name": "Addon", + "percentage_usage": 0, + "price_description": None, + "projected_amount_usd": "0.00", + "projected_usage": 0, + "subscribed": True, + "tiered": True, + "tiers": [ + { + "current_amount_usd": "0.00", + "current_usage": 0, + "flat_amount_usd": "0", + "projected_amount_usd": "None", + "projected_usage": None, + "unit_amount_usd": "0.00", + "up_to": 1000000, + }, + { + "current_amount_usd": "0.00", + "current_usage": 0, + "flat_amount_usd": "0", + "projected_amount_usd": "None", + "projected_usage": None, + "unit_amount_usd": "0.0000135", + "up_to": 2000000, + }, + ], + "type": "events", + "unit_amount_usd": "0.00", + "usage_key": "events", + "usage_limit": None, }, ], tiered=True, unit_amount_usd="0.00", - current_amount_usd="0.00", + current_amount_usd=0.0, current_usage=0, usage_limit=None, has_exceeded_limit=False, percentage_usage=0, projected_usage=0, - projected_amount_usd="0.00", + projected_amount=0, + projected_amount_usd=0.00, usage_key="events", ) ] @@ -271,15 +359,23 @@ def mock_implementation(url: str, headers: Any = None, params: Any = None) -> Ma "unit_amount_usd": "0.00", "up_to": 1000000, "current_amount_usd": "0.00", + "current_usage": 0, + "flat_amount_usd": "0", + "projected_amount_usd": "None", + "projected_usage": None, }, { "unit_amount_usd": "0.00045", "up_to": 2000000, - "current_amount_usd": None, + "current_amount_usd": "0.00", + "current_usage": 0, + "flat_amount_usd": "0", + "projected_amount_usd": "None", + "projected_usage": None, }, ], "tiered": True, - "current_amount_usd": "0.00", + "current_amount_usd": 0.00, "current_usage": 0, "usage_limit": None, "percentage_usage": 0, @@ -288,7 +384,48 @@ def mock_implementation(url: str, headers: Any = None, params: Any = None) -> Ma "projected_amount_usd": "0.00", "projected_usage": 0, "usage_key": "events", - } + "addons": [ + { + "current_amount_usd": 0.00, + "current_usage": 0, + "description": "Test Addon", + "free_allocation": 10000, + "has_exceeded_limit": False, + "image_url": "https://posthog.com/static/images/product-os.png", + "name": "Addon", + "percentage_usage": 0, + "price_description": None, + "projected_amount_usd": "0.00", + "projected_usage": 0, + "subscribed": True, + "tiered": True, + "tiers": [ + { + "current_amount_usd": "0.00", + "current_usage": 0, + "flat_amount_usd": "0", + "projected_amount_usd": "None", + "projected_usage": None, + "unit_amount_usd": "0.00", + "up_to": 1000000, + }, + { + "current_amount_usd": "0.00", + "current_usage": 0, + "flat_amount_usd": "0", + "projected_amount_usd": "None", + "projected_usage": None, + "unit_amount_usd": "0.0000135", + "up_to": 2000000, + }, + ], + "type": "events", + "unit_amount_usd": "0.00", + "usage_key": "events", + "usage_limit": None, + }, + ], + }, ], "billing_period": { "current_period_start": "2022-10-07T11:12:48", @@ -342,25 +479,75 @@ def mock_implementation(url: str, headers: Any = None, params: Any = None) -> Ma "unit_amount_usd": "0.00", "up_to": 1000000, "current_amount_usd": "0.00", + "current_usage": 0, + "flat_amount_usd": "0", + "projected_amount_usd": "None", + "projected_usage": None, }, { "unit_amount_usd": "0.00045", "up_to": 2000000, - "current_amount_usd": None, + "current_amount_usd": "0.00", + "current_usage": 0, + "flat_amount_usd": "0", + "projected_amount_usd": "None", + "projected_usage": None, }, ], "current_usage": 0, - "percentage_usage": 0.0, - "current_amount_usd": "0.00", + "percentage_usage": 0, + "current_amount_usd": 0.0, "has_exceeded_limit": False, - "projected_amount_usd": "0.00", + "projected_amount_usd": 0.0, + "projected_amount": 0, "projected_usage": 0, "tiered": True, "unit_amount_usd": "0.00", "usage_limit": None, "image_url": "https://posthog.com/static/images/product-os.png", - "percentage_usage": 0.0, + "percentage_usage": 0, "usage_key": "events", + "addons": [ + { + "current_amount_usd": 0.0, + "current_usage": 0, + "description": "Test Addon", + "free_allocation": 10000, + "has_exceeded_limit": False, + "image_url": "https://posthog.com/static/images/product-os.png", + "name": "Addon", + "percentage_usage": 0, + "price_description": None, + "projected_amount_usd": "0.00", + "projected_usage": 0, + "subscribed": True, + "tiered": True, + "tiers": [ + { + "current_amount_usd": "0.00", + "current_usage": 0, + "flat_amount_usd": "0", + "projected_amount_usd": "None", + "projected_usage": None, + "unit_amount_usd": "0.00", + "up_to": 1000000, + }, + { + "current_amount_usd": "0.00", + "current_usage": 0, + "flat_amount_usd": "0", + "projected_amount_usd": "None", + "projected_usage": None, + "unit_amount_usd": "0.0000135", + "up_to": 2000000, + }, + ], + "type": "events", + "unit_amount_usd": "0.00", + "usage_key": "events", + "usage_limit": None, + }, + ], } ], "billing_period": { @@ -502,7 +689,7 @@ def mock_implementation(url: str, headers: Any = None, params: Any = None) -> Ma elif "api/billing" in url: mock.status_code = 200 mock.json.return_value = create_billing_response( - customer=create_billing_customer(has_active_subscription=True) + customer=create_billing_customer(has_active_subscription=True), ) mock.json.return_value["customer"]["usage_summary"]["events"]["usage"] = 1000 @@ -537,6 +724,27 @@ def mock_implementation(url: str, headers: Any = None, params: Any = None) -> Ma "period": ["2022-10-07T11:12:48", "2022-11-07T11:12:48"], } + self.organization.usage = {"events": {"limit": None, "usage": 1000, "todays_usage": 1100000}} + self.organization.save() + + res = self.client.get("/api/billing-v2") + assert res.status_code == 200 + res_json = res.json() + # Should update product usage to reflect today's usage + assert res_json["products"][0]["current_usage"] == 1101000 + assert res_json["products"][0]["current_amount_usd"] == 45.45 + assert res_json["products"][0]["tiers"][0]["current_usage"] == 1000000 + assert res_json["products"][0]["tiers"][0]["current_amount_usd"] == "0.00" + assert res_json["products"][0]["tiers"][1]["current_usage"] == 101000 + assert res_json["products"][0]["tiers"][1]["current_amount_usd"] == "45.45" + + assert res_json["products"][0]["addons"][0]["current_usage"] == 1101000 + assert res_json["products"][0]["addons"][0]["current_amount_usd"] == 1.36 + assert res_json["products"][0]["addons"][0]["tiers"][0]["current_usage"] == 1000000 + assert res_json["products"][0]["addons"][0]["tiers"][0]["current_amount_usd"] == "0.00" + assert res_json["products"][0]["addons"][0]["tiers"][1]["current_usage"] == 101000 + assert res_json["products"][0]["addons"][0]["tiers"][1]["current_amount_usd"] == "1.36" + def mock_implementation_missing_customer(url: str, headers: Any = None, params: Any = None) -> MagicMock: mock = MagicMock() mock.status_code = 404 diff --git a/ee/billing/billing_manager.py b/ee/billing/billing_manager.py index dee55589a051e..ac3ace24ddb8d 100644 --- a/ee/billing/billing_manager.py +++ b/ee/billing/billing_manager.py @@ -1,5 +1,6 @@ from datetime import datetime, timedelta -from typing import Any, Optional, cast +from decimal import Decimal +from typing import Any, Optional, Union, cast import jwt import requests @@ -9,7 +10,7 @@ from rest_framework.exceptions import NotAuthenticated from sentry_sdk import capture_exception -from ee.billing.billing_types import BillingStatus +from ee.billing.billing_types import BillingStatus, Tier from ee.billing.quota_limiting import set_org_usage_summary, sync_org_quota_limits from ee.models import License from ee.settings import BILLING_SERVICE_URL @@ -52,6 +53,61 @@ def handle_billing_service_error(res: requests.Response, valid_codes=(200, 404, raise Exception(f"Billing service returned bad status code: {res.status_code}", f"body:", res.text) +def compute_usage_per_tier(current_usage: int, projected_usage: int, tiers): + remaining_usage = current_usage + remaining_projected_usage = projected_usage or 0 + previous_tier: Optional[dict[str, Any]] = None + tier_max_usage: Union[int, float] = 0 + + result: list[Tier] = [] + for tier in tiers: + if previous_tier and previous_tier.get("up_to"): + previous_tier_up_to = previous_tier["up_to"] + else: + previous_tier_up_to = 0 + + if tier.get("up_to"): + tier_max_usage = tier["up_to"] - previous_tier_up_to + else: + tier_max_usage = float("inf") + + flat_amount_usd = Decimal(tier.get("flat_amount_usd") or 0) + unit_amount_usd = Decimal(tier.get("unit_amount_usd") or 0) + usage_this_tier = int(min(remaining_usage, tier_max_usage)) + remaining_usage -= usage_this_tier + current_amount_usd = Decimal(unit_amount_usd * usage_this_tier + flat_amount_usd).quantize(Decimal("0.01")) + previous_tier = tier + if projected_usage: + projected_usage_this_tier = int(min(remaining_projected_usage, tier_max_usage)) + remaining_projected_usage -= projected_usage_this_tier + projected_amount_usd = Decimal(unit_amount_usd * projected_usage_this_tier + flat_amount_usd).quantize( + Decimal("0.01") + ) + else: + projected_usage_this_tier = None + projected_amount_usd = None + + result.append( + Tier( + flat_amount_usd=str(flat_amount_usd), + unit_amount_usd=str(unit_amount_usd), + up_to=tier.get("up_to", None), + current_amount_usd=str(current_amount_usd), + current_usage=usage_this_tier, + projected_usage=projected_usage_this_tier, + projected_amount_usd=str(projected_amount_usd), + ) + ) + return result + + +def sum_total_across_tiers(tiers): + total = Decimal(0) + for tier in tiers: + total += Decimal(tier["current_amount_usd"]) + return total + + class BillingManager: license: Optional[License] @@ -89,7 +145,6 @@ def get_billing(self, organization: Optional[Organization], plan_keys: Optional[ } # Extend the products with accurate usage_limit info - for product in response["products"]: usage_key = product.get("usage_key", None) if not usage_key: @@ -105,6 +160,34 @@ def get_billing(self, organization: Optional[Organization], plan_keys: Optional[ product["current_usage"] = current_usage product["percentage_usage"] = current_usage / usage_limit if usage_limit else 0 + # Also update the tiers + if product["tiers"]: + product["tiers"] = compute_usage_per_tier(current_usage, product["projected_usage"], product["tiers"]) + product["current_amount_usd"] = sum_total_across_tiers(product["tiers"]) + + # Update the add on tiers + # TODO: enhanced_persons: make sure this updates properly for addons with different usage keys + for addon in product["addons"]: + if not addon["subscribed"]: + continue + addon_usage_key = addon.get("usage_key") + if not usage_key: + continue + if addon_usage_key != usage_key: + usage = response.get("usage_summary", {}).get(addon_usage_key, {}) + usage_limit = usage.get("limit") + current_usage = usage.get("usage") or 0 + if ( + organization + and organization.usage + and organization.usage.get(usage_key, {}).get("todays_usage", None) + ): + todays_usage = organization.usage[usage_key]["todays_usage"] + current_usage = current_usage + todays_usage + addon["current_usage"] = current_usage + addon["tiers"] = compute_usage_per_tier(current_usage, addon["projected_usage"], addon["tiers"]) + addon["current_amount_usd"] = sum_total_across_tiers(addon["tiers"]) + return response def update_billing(self, organization: Organization, data: dict[str, Any]) -> None: @@ -180,7 +263,6 @@ def _get_billing(self, organization: Organization) -> BillingStatus: f"{BILLING_SERVICE_URL}/api/billing", headers=self.get_auth_headers(organization), ) - handle_billing_service_error(res) data = res.json() diff --git a/ee/billing/billing_types.py b/ee/billing/billing_types.py index 0761e02e807ef..fd25af67fdbba 100644 --- a/ee/billing/billing_types.py +++ b/ee/billing/billing_types.py @@ -5,10 +5,43 @@ class Tier(TypedDict): - flat_amount_usd: Decimal - unit_amount_usd: Decimal - current_amount_usd: Decimal + flat_amount_usd: str + unit_amount_usd: str + current_amount_usd: str up_to: Optional[int] + current_usage: int + projected_usage: Optional[int] + projected_amount_usd: Optional[str] + + +class CustomerProductAddon(TypedDict): + name: str + description: str + price_description: Optional[str] + image_url: Optional[str] + icon_key: str + docs_url: Optional[str] + type: str + tiers: Optional[Tier] + tiered: bool + included_with_main_product: ( + bool # if the addon is included in the main product subscription, not paid for separately + ) + inclusion_only: ( + # if the addon subscription state is dependent on the main product subscription state. + # Ie. addon is automatically sub'd when parent is sub'd. + bool + ) + subscribed: bool + unit: Optional[str] + unit_amount_usd: Optional[Decimal] + current_amount_usd: Optional[Decimal] + current_usage: int + projected_usage: Optional[int] + projected_amount_usd: Optional[Decimal] + contact_support: bool + usage_key: Optional[str] + usage_limit: Optional[int] class CustomerProduct(TypedDict): @@ -28,7 +61,9 @@ class CustomerProduct(TypedDict): percentage_usage: float projected_usage: int projected_amount: Decimal + projected_amount_usd: Decimal usage_key: str + addons: list[CustomerProductAddon] class LicenseInfo(TypedDict): diff --git a/posthog/hogql/property.py b/posthog/hogql/property.py index 84a73c6a53314..413d1ddefc384 100644 --- a/posthog/hogql/property.py +++ b/posthog/hogql/property.py @@ -1,39 +1,39 @@ import re -from typing import Optional, Union, cast, Literal +from typing import Literal, Optional, Union, cast from pydantic import BaseModel from posthog.constants import ( AUTOCAPTURE_EVENT, - PropertyOperatorType, + PAGEVIEW_EVENT, TREND_FILTER_TYPE_ACTIONS, TREND_FILTER_TYPE_EVENTS, - PAGEVIEW_EVENT, + PropertyOperatorType, ) from posthog.hogql import ast from posthog.hogql.base import AST -from posthog.hogql.functions import find_hogql_aggregation from posthog.hogql.errors import NotImplementedError +from posthog.hogql.functions import find_hogql_aggregation from posthog.hogql.parser import parse_expr from posthog.hogql.visitor import TraversingVisitor, clone_expr from posthog.models import ( Action, Cohort, Property, - Team, PropertyDefinition, + Team, ) from posthog.models.event import Selector from posthog.models.property import PropertyGroup from posthog.models.property.util import build_selector_regex from posthog.models.property_definition import PropertyType from posthog.schema import ( - PropertyOperator, + EmptyPropertyFilter, + FilterLogicalOperator, PropertyGroupFilter, PropertyGroupFilterValue, - FilterLogicalOperator, + PropertyOperator, RetentionEntity, - EmptyPropertyFilter, ) from posthog.warehouse.models import DataWarehouseJoin, DataWarehouseSavedQuery, DataWarehouseTable from posthog.utils import get_from_dict_or_attr