diff --git a/api/app_analytics/analytics_db_service.py b/api/app_analytics/analytics_db_service.py index 1cb377ac5f0b..106eaed9f645 100644 --- a/api/app_analytics/analytics_db_service.py +++ b/api/app_analytics/analytics_db_service.py @@ -1,4 +1,4 @@ -from datetime import date, datetime, timedelta +from datetime import datetime, timedelta import structlog from common.core.utils import using_database_replica @@ -122,20 +122,31 @@ def get_usage_data_from_local_db( return map_annotated_api_usage_buckets_to_usage_data(qs) -def get_total_events_count(organisation) -> int: # type: ignore[no-untyped-def] +def get_total_events_count( + organisation: Organisation, + date_start: datetime | None = None, + date_stop: datetime | None = None, +) -> int: """ - Return total number of events for an organisation in the last 30 days + Return total number of events for an organisation for a range, or last 30 days """ + today = timezone.now().replace(hour=0, minute=0, second=0, microsecond=0) + date_start = date_start or (today - timedelta(days=30)) + date_stop = date_stop or today if settings.USE_POSTGRES_FOR_ANALYTICS: - count = APIUsageBucket.objects.filter( + count: int = APIUsageBucket.objects.filter( environment_id__in=_get_environment_ids_for_org(organisation), - created_at__date__lte=date.today(), - created_at__date__gt=date.today() - timedelta(days=30), + created_at__date__lte=date_stop, + created_at__date__gt=date_start, bucket_size=constants.ANALYTICS_READ_BUCKET_SIZE, ).aggregate(total_count=Sum("total_count"))["total_count"] else: - count = get_events_for_organisation(organisation.id) - return count # type: ignore[no-any-return] + count = get_events_for_organisation( + organisation.id, + date_start=date_start, + date_stop=date_stop, + ) + return count def get_feature_evaluation_data( diff --git a/api/integrations/flagsmith/data/environment.json b/api/integrations/flagsmith/data/environment.json index 811b466d4f0b..6cd4f30a6690 100644 --- a/api/integrations/flagsmith/data/environment.json +++ b/api/integrations/flagsmith/data/environment.json @@ -104,6 +104,19 @@ "feature_state_value": null, "featurestate_uuid": "f976df2f-2341-4623-8425-d6eda23a2ebc", "multivariate_feature_state_values": [] + }, + { + "django_id": 1229327, + "enabled": false, + "feature": { + "id": 195393, + "name": "get_current_api_usage_deprecated", + "type": "STANDARD" + }, + "feature_segment": null, + "feature_state_value": null, + "featurestate_uuid": "e7ed0d54-b17c-4df1-ab98-5f8dc9597127", + "multivariate_feature_state_values": [] } ], "id": 0, diff --git a/api/organisations/task_helpers.py b/api/organisations/task_helpers.py index 2e0bd1a940c6..fbfc8cb24dd9 100644 --- a/api/organisations/task_helpers.py +++ b/api/organisations/task_helpers.py @@ -7,8 +7,10 @@ from django.template.loader import render_to_string from django.utils import timezone +from app_analytics.analytics_db_service import get_total_events_count from app_analytics.influxdb_wrapper import get_current_api_usage from core.helpers import get_current_site_url +from integrations.flagsmith.client import get_client from organisations.models import ( Organisation, OrganisationAPIUsageNotification, @@ -126,7 +128,16 @@ def handle_api_usage_notification_for_organisation(organisation: Organisation) - allowed_api_calls = subscription_cache.allowed_30d_api_calls - api_usage = get_current_api_usage(organisation.id, period_starts_at) + flagsmith_client = get_client("local", local_eval=True) + flags = flagsmith_client.get_identity_flags( + organisation.flagsmith_identifier, + traits=organisation.flagsmith_on_flagsmith_api_traits, + ) + # TODO: Default to get_total_events_count — https://github.com/Flagsmith/flagsmith/issues/6985 + if flags.is_feature_enabled("get_current_api_usage_deprecated"): # pragma: no cover + api_usage = get_total_events_count(organisation, period_starts_at) + else: + api_usage = get_current_api_usage(organisation.id, period_starts_at) # For some reason the allowed API calls is set to 0 so default to the max free plan. allowed_api_calls = allowed_api_calls or MAX_API_CALLS_IN_FREE_PLAN diff --git a/api/organisations/tasks.py b/api/organisations/tasks.py index a5019a51fe23..4c689857e4d9 100644 --- a/api/organisations/tasks.py +++ b/api/organisations/tasks.py @@ -12,6 +12,7 @@ register_task_handler, ) +from app_analytics.analytics_db_service import get_total_events_count from app_analytics.influxdb_wrapper import get_current_api_usage from integrations.flagsmith.client import get_client from organisations import subscription_info_cache @@ -214,10 +215,19 @@ def charge_for_api_call_count_overages(): # type: ignore[no-untyped-def] continue subscription_cache = organisation.subscription_information_cache - api_usage = get_current_api_usage( - organisation_id=organisation.id, - date_start=subscription_cache.current_billing_term_starts_at, - ) + # TODO: Default to get_total_events_count — https://github.com/Flagsmith/flagsmith/issues/6985 + if flags.is_feature_enabled( + "get_current_api_usage_deprecated" + ): # pragma: no cover + api_usage = get_total_events_count( + organisation, + date_start=subscription_cache.current_billing_term_starts_at, + ) + else: + api_usage = get_current_api_usage( + organisation_id=organisation.id, + date_start=subscription_cache.current_billing_term_starts_at, + ) # Grace period for organisations < 200% of usage. if ( @@ -347,10 +357,19 @@ def restrict_use_due_to_api_limit_grace_period_over() -> None: OrganisationBreachedGracePeriod.objects.get_or_create(organisation=organisation) subscription_cache = organisation.subscription_information_cache - api_usage = get_current_api_usage( - organisation_id=organisation.id, - date_start=now - timedelta(days=30), - ) + # TODO: Default to get_total_events_count — https://github.com/Flagsmith/flagsmith/issues/6985 + if flags.is_feature_enabled( + "get_current_api_usage_deprecated" + ): # pragma: no cover + api_usage = get_total_events_count( + organisation, + date_start=now - timedelta(days=30), + ) + else: + api_usage = get_current_api_usage( + organisation_id=organisation.id, + date_start=now - timedelta(days=30), + ) if api_usage / subscription_cache.allowed_30d_api_calls < 1.0: logger.info( f"API use for organisation {organisation.id} has fallen to below limit, so not restricting use." diff --git a/api/tests/unit/app_analytics/test_analytics_db_service.py b/api/tests/unit/app_analytics/test_analytics_db_service.py index a3d8dda585da..84191490a917 100644 --- a/api/tests/unit/app_analytics/test_analytics_db_service.py +++ b/api/tests/unit/app_analytics/test_analytics_db_service.py @@ -45,6 +45,7 @@ def cache(organisation: Organisation) -> OrganisationSubscriptionInformationCach @pytest.mark.use_analytics_db +@pytest.mark.freeze_time("2023-01-19T09:09:47.325132+00:00") def test_get_usage_data_from_local_db__multiple_buckets__returns_aggregated_daily_data( # type: ignore[no-untyped-def] organisation, environment, settings ): @@ -197,6 +198,7 @@ def test_get_usage_data_from_local_db__environment_filter__returns_expected( @pytest.mark.use_analytics_db +@pytest.mark.freeze_time("2023-01-19T09:09:47.325132+00:00") def test_get_usage_data_from_local_db__labels_filter__returns_expected( organisation: Organisation, environment: Environment, @@ -269,6 +271,7 @@ def test_get_usage_data_from_local_db__labels_filter__returns_expected( @pytest.mark.use_analytics_db +@pytest.mark.freeze_time("2023-01-19T09:09:47.325132+00:00") def test_get_total_events_count__multiple_buckets__returns_correct_total( # type: ignore[no-untyped-def] organisation, environment, settings ): @@ -322,6 +325,7 @@ def test_get_total_events_count__multiple_buckets__returns_correct_total( # typ @pytest.mark.use_analytics_db +@pytest.mark.freeze_time("2023-01-19T09:09:47.325132+00:00") def test_get_feature_evaluation_data_from_local_db__multiple_buckets__returns_aggregated_daily_data( feature: Feature, environment: Environment, @@ -538,6 +542,7 @@ def test_get_usage_data__no_analytics_configured__no_calls_expected( mocked_get_usage_data_from_local_db.assert_not_called() +@pytest.mark.freeze_time("2023-01-19T09:09:47.325132+00:00") def test_get_total_events_count__postgres_not_configured__calls_influx( # type: ignore[no-untyped-def] mocker, settings, organisation ): @@ -553,7 +558,9 @@ def test_get_total_events_count__postgres_not_configured__calls_influx( # type: # Then assert total_events_count == mocked_get_events_for_organisation.return_value mocked_get_events_for_organisation.assert_called_once_with( - organisation_id=organisation.id + organisation.id, + date_start=datetime(2022, 12, 20, 0, 0, tzinfo=UTC), + date_stop=datetime(2023, 1, 19, 0, 0, tzinfo=UTC), ) diff --git a/api/tests/unit/organisations/test_unit_organisations_tasks.py b/api/tests/unit/organisations/test_unit_organisations_tasks.py index b06ed920dfbf..6ecba9230b49 100644 --- a/api/tests/unit/organisations/test_unit_organisations_tasks.py +++ b/api/tests/unit/organisations/test_unit_organisations_tasks.py @@ -51,6 +51,7 @@ send_org_subscription_cancelled_alert, unrestrict_after_api_limit_grace_period_is_stale, ) +from tests.types import EnableFeaturesFixture from users.models import FFAdminUser @@ -885,6 +886,7 @@ def test_charge_for_api_call_count_overages__scale_up_plan__charges_correct_addo organisation: Organisation, mocker: MockerFixture, plan: str, + enable_features: EnableFeaturesFixture, ) -> None: # Given now = timezone.now() @@ -924,10 +926,7 @@ def test_charge_for_api_call_count_overages__scale_up_plan__charges_correct_addo autospec=True, ) - get_client_mock = mocker.patch("organisations.tasks.get_client") - client_mock = MagicMock() - get_client_mock.return_value = client_mock - client_mock.get_identity_flags.return_value.is_feature_enabled.return_value = True + enable_features("api_usage_overage_charges") mock_api_usage = mocker.patch( "organisations.tasks.get_current_api_usage", @@ -1071,6 +1070,7 @@ def test_charge_for_api_call_count_overages__flagsmith_feature_disabled__skips_c def test_charge_for_api_call_count_overages__within_grace_period__creates_breached_record_without_charging( organisation: Organisation, mocker: MockerFixture, + enable_features: EnableFeaturesFixture, ) -> None: # Given now = timezone.now() @@ -1092,10 +1092,7 @@ def test_charge_for_api_call_count_overages__within_grace_period__creates_breach notified_at=now, ) - get_client_mock = mocker.patch("organisations.tasks.get_client") - client_mock = MagicMock() - get_client_mock.return_value = client_mock - client_mock.get_identity_flags.return_value.is_feature_enabled.return_value = True + enable_features("api_usage_overage_charges") mock_chargebee_update = mocker.patch( "organisations.chargebee.chargebee.chargebee_client.Subscription.update", @@ -1121,6 +1118,7 @@ def test_charge_for_api_call_count_overages__within_grace_period__creates_breach def test_charge_for_api_call_count_overages__grace_period_previously_breached__charges_overage( organisation: Organisation, mocker: MockerFixture, + enable_features: EnableFeaturesFixture, ) -> None: # Given now = timezone.now() @@ -1143,10 +1141,7 @@ def test_charge_for_api_call_count_overages__grace_period_previously_breached__c ) OrganisationBreachedGracePeriod.objects.create(organisation=organisation) - get_client_mock = mocker.patch("organisations.tasks.get_client") - client_mock = MagicMock() - get_client_mock.return_value = client_mock - client_mock.get_identity_flags.return_value.is_feature_enabled.return_value = True + enable_features("api_usage_overage_charges") mock_chargebee_update = mocker.patch( "organisations.chargebee.chargebee.chargebee_client.Subscription.update", @@ -1180,6 +1175,7 @@ def test_charge_for_api_call_count_overages__grace_period_previously_breached__c def test_charge_for_api_call_count_overages__uncovered_plan__does_not_charge( organisation: Organisation, mocker: MockerFixture, + enable_features: EnableFeaturesFixture, ) -> None: # Given now = timezone.now() @@ -1204,10 +1200,7 @@ def test_charge_for_api_call_count_overages__uncovered_plan__does_not_charge( notified_at=now, ) - get_client_mock = mocker.patch("organisations.tasks.get_client") - client_mock = MagicMock() - get_client_mock.return_value = client_mock - client_mock.get_identity_flags.return_value.is_feature_enabled.return_value = True + enable_features("api_usage_overage_charges") mocker.patch( "organisations.chargebee.chargebee.chargebee_client.Subscription.retrieve", @@ -1236,6 +1229,7 @@ def test_charge_for_api_call_count_overages__uncovered_plan__does_not_charge( def test_charge_for_api_call_count_overages__usage_under_api_limit__does_not_charge( organisation: Organisation, mocker: MockerFixture, + enable_features: EnableFeaturesFixture, ) -> None: # Given now = timezone.now() @@ -1257,10 +1251,7 @@ def test_charge_for_api_call_count_overages__usage_under_api_limit__does_not_cha notified_at=now, ) - get_client_mock = mocker.patch("organisations.tasks.get_client") - client_mock = MagicMock() - get_client_mock.return_value = client_mock - client_mock.get_identity_flags.return_value.is_feature_enabled.return_value = True + enable_features("api_usage_overage_charges") mocker.patch( "organisations.chargebee.chargebee.chargebee_client.Subscription.retrieve", @@ -1289,6 +1280,7 @@ def test_charge_for_api_call_count_overages__usage_under_api_limit__does_not_cha def test_charge_for_api_call_count_overages__startup_plan__charges_correct_addon_quantity( organisation: Organisation, mocker: MockerFixture, + enable_features: EnableFeaturesFixture, ) -> None: # Given now = timezone.now() @@ -1311,10 +1303,7 @@ def test_charge_for_api_call_count_overages__startup_plan__charges_correct_addon notified_at=now, ) - get_client_mock = mocker.patch("organisations.tasks.get_client") - client_mock = MagicMock() - get_client_mock.return_value = client_mock - client_mock.get_identity_flags.return_value.is_feature_enabled.return_value = True + enable_features("api_usage_overage_charges") mocker.patch( "organisations.chargebee.chargebee.chargebee_client.Subscription.retrieve", autospec=True, @@ -1381,6 +1370,7 @@ def test_charge_for_api_call_count_overages__non_standard_plan__logs_unknown_pla organisation: Organisation, mocker: MockerFixture, inspecting_handler: logging.Handler, + enable_features: EnableFeaturesFixture, ) -> None: # Given now = timezone.now() @@ -1407,10 +1397,7 @@ def test_charge_for_api_call_count_overages__non_standard_plan__logs_unknown_pla notified_at=now, ) - get_client_mock = mocker.patch("organisations.tasks.get_client") - client_mock = MagicMock() - get_client_mock.return_value = client_mock - client_mock.get_identity_flags.return_value.is_feature_enabled.return_value = True + enable_features("api_usage_overage_charges") mocker.patch( "organisations.chargebee.chargebee.chargebee_client.Subscription.retrieve", autospec=True, @@ -1444,6 +1431,7 @@ def test_charge_for_api_call_count_overages__billing_exception__logs_error_and_d organisation: Organisation, mocker: MockerFixture, inspecting_handler: logging.Handler, + enable_features: EnableFeaturesFixture, ) -> None: # Given now = timezone.now() @@ -1469,10 +1457,7 @@ def test_charge_for_api_call_count_overages__billing_exception__logs_error_and_d notified_at=now, ) - get_client_mock = mocker.patch("organisations.tasks.get_client") - client_mock = MagicMock() - get_client_mock.return_value = client_mock - client_mock.get_identity_flags.return_value.is_feature_enabled.return_value = True + enable_features("api_usage_overage_charges") mocker.patch( "organisations.chargebee.chargebee.chargebee_client.Subscription.retrieve", autospec=True, @@ -1506,6 +1491,7 @@ def test_charge_for_api_call_count_overages__billing_exception__logs_error_and_d def test_charge_for_api_call_count_overages__startup_plan_with_existing_billing__charges_remaining_overage( organisation: Organisation, mocker: MockerFixture, + enable_features: EnableFeaturesFixture, ) -> None: # Given now = timezone.now() @@ -1534,10 +1520,7 @@ def test_charge_for_api_call_count_overages__startup_plan_with_existing_billing_ billed_at=now, ) - get_client_mock = mocker.patch("organisations.tasks.get_client") - client_mock = MagicMock() - get_client_mock.return_value = client_mock - client_mock.get_identity_flags.return_value.is_feature_enabled.return_value = True + enable_features("api_usage_overage_charges") mocker.patch( "organisations.chargebee.chargebee.chargebee_client.Subscription.retrieve", autospec=True, @@ -1681,12 +1664,13 @@ def test_restrict_use_due_to_api_limit_grace_period_over__multiple_organisations mailoutbox: list[EmailMultiAlternatives], admin_user: FFAdminUser, staff_user: FFAdminUser, + enable_features: EnableFeaturesFixture, ) -> None: # Given - get_client_mock = mocker.patch("organisations.tasks.get_client") - client_mock = MagicMock() - get_client_mock.return_value = client_mock - client_mock.get_identity_flags.return_value.is_feature_enabled.return_value = True + enable_features( + "api_limiting_stop_serving_flags", + "api_limiting_block_access_to_admin", + ) now = timezone.now() organisation2 = Organisation.objects.create(name="Org #2") @@ -1821,33 +1805,6 @@ def test_restrict_use_due_to_api_limit_grace_period_over__multiple_organisations assert organisation6.block_access_to_admin is True assert organisation6.api_limit_access_block - client_mock.get_identity_flags.call_args_list == [ - call( - f"org.{organisation.id}", - traits={ - "organisation.id": organisation.id, - "organisation.name": organisation.name, - "subscription.plan": organisation.subscription.plan, - }, - ), - call( - f"org.{organisation2.id}", - traits={ - "organisation.id": organisation2.id, - "organisation.name": organisation2.name, - "subscription.plan": organisation2.subscription.plan, - }, - ), - call( - f"org.{organisation6.id}", - traits={ - "organisation.id": organisation6.id, - "organisation.name": organisation6.name, - "subscription.plan": organisation6.subscription.plan, - }, - ), - ] - assert len(mailoutbox) == 3 email1 = mailoutbox[0] assert email1.subject == "Flagsmith API use has been blocked due to overuse" @@ -1921,12 +1878,13 @@ def test_restrict_use_due_to_api_limit_grace_period_over__previously_breached__b mailoutbox: list[EmailMultiAlternatives], admin_user: FFAdminUser, staff_user: FFAdminUser, + enable_features: EnableFeaturesFixture, ) -> None: # Given - get_client_mock = mocker.patch("organisations.tasks.get_client") - client_mock = MagicMock() - get_client_mock.return_value = client_mock - client_mock.get_identity_flags.return_value.is_feature_enabled.return_value = True + enable_features( + "api_limiting_stop_serving_flags", + "api_limiting_block_access_to_admin", + ) now = timezone.now() @@ -1977,14 +1935,15 @@ def test_restrict_use_due_to_api_limit_grace_period_over__missing_subscription_c organisation: Organisation, freezer: FrozenDateTimeFactory, mailoutbox: list[EmailMultiAlternatives], + enable_features: EnableFeaturesFixture, ) -> None: # Given assert not organisation.has_subscription_information_cache() - get_client_mock = mocker.patch("organisations.tasks.get_client") - client_mock = MagicMock() - get_client_mock.return_value = client_mock - client_mock.get_identity_flags.return_value.is_feature_enabled.return_value = True + enable_features( + "api_limiting_stop_serving_flags", + "api_limiting_block_access_to_admin", + ) now = timezone.now() organisation.subscription.subscription_id = "fancy_sub_id23" @@ -2020,6 +1979,7 @@ def test_restrict_use_due_to_api_limit_grace_period_over__reduced_api_usage__doe freezer: FrozenDateTimeFactory, mailoutbox: list[EmailMultiAlternatives], inspecting_handler: logging.Handler, + enable_features: EnableFeaturesFixture, ) -> None: # Given assert not organisation.has_subscription_information_cache() @@ -2028,10 +1988,10 @@ def test_restrict_use_due_to_api_limit_grace_period_over__reduced_api_usage__doe logger.addHandler(inspecting_handler) - get_client_mock = mocker.patch("organisations.tasks.get_client") - client_mock = MagicMock() - get_client_mock.return_value = client_mock - client_mock.get_identity_flags.return_value.is_feature_enabled.return_value = True + enable_features( + "api_limiting_stop_serving_flags", + "api_limiting_block_access_to_admin", + ) now = timezone.now()