diff --git a/src/quartz_api/internal/service/uk_national/cache.py b/src/quartz_api/internal/service/uk_national/cache.py index b6a739bf..bd29f15e 100644 --- a/src/quartz_api/internal/service/uk_national/cache.py +++ b/src/quartz_api/internal/service/uk_national/cache.py @@ -1,13 +1,17 @@ """Cache key builder.""" +import datetime as dt import logging from collections.abc import Callable from typing import Any from fastapi import Request, Response +from .time_utils import trial_expired + log = logging.getLogger(__name__) +#TODO: verify this string with brad cache_dependent_scopes = ["read:intraday"] legacy_query_params = ["historic"] @@ -60,6 +64,10 @@ async def key_builder( .replace("True", "true") ) + #Invalidate cache if trial has expired + if trial_expired(auth, dt.datetime.now(dt.UTC)): + key += ":trial_ended=true" + log.info(f"Cache key generated: {key}") return key diff --git a/src/quartz_api/internal/service/uk_national/gsp_router.py b/src/quartz_api/internal/service/uk_national/gsp_router.py index 9d17b1fb..7c551c24 100644 --- a/src/quartz_api/internal/service/uk_national/gsp_router.py +++ b/src/quartz_api/internal/service/uk_national/gsp_router.py @@ -45,6 +45,7 @@ ) from .time_utils import ( limit_end_datetime_by_permissions, + trial_expired, ) log = logging.getLogger(__name__) @@ -73,7 +74,7 @@ async def get_forecasts_for_a_specific_gsp( request: Request, # noqa: ARG001 db: models.StorageClientDependency, - auth: AuthDependency, # noqa: ARG001 + auth: AuthDependency, gsp_id: int, start_datetime_utc: models.UTCDatetimeDefaultWindowStartShiftUK, end_datetime_utc: Annotated[ @@ -109,6 +110,11 @@ async def get_forecasts_for_a_specific_gsp( # existent GSP - so that is what is replicated here. Seems odd to me. return [] + # end_datetime_utc is already clamped to "now" for expired trials. Clamp + # start_datetime_utc too, so it can't end up later than end_datetime_utc. + if trial_expired(auth, dt.datetime.now(dt.UTC)): + start_datetime_utc = min(start_datetime_utc, end_datetime_utc) + pgvs = await db.get_predicted_generation( location_uuid=gsp_id_map[gsp_id].uuid, window_start=start_datetime_utc, @@ -302,7 +308,7 @@ async def get_all_available_forecasts( request: Request, background_tasks: BackgroundTasks, db: models.StorageClientDependency, - auth: AuthDependency, # noqa: ARG001 + auth: AuthDependency, start_datetime_utc: Annotated[ models.UTCDatetimeDefaultNowWindowStart, AfterValidator(lambda v: pd.Timestamp(v).ceil("30min").to_pydatetime()), @@ -334,6 +340,11 @@ async def get_all_available_forecasts( start_datetime_utc_set = start_datetime_utc != default_now_window_start() end_datetime_utc_set = end_datetime_utc != default_window_end() + # end_datetime_utc is already clamped to "now" for expired trials. Clamp + # start_datetime_utc too, so it can't end up later than end_datetime_utc. + if trial_expired(auth, dt.datetime.now(dt.UTC)): + start_datetime_utc = min(start_datetime_utc, end_datetime_utc) + if gsp_ids is None and start_datetime_utc != end_datetime_utc: if start_datetime_utc_set or end_datetime_utc_set: raise HTTPException( diff --git a/src/quartz_api/internal/service/uk_national/national_router.py b/src/quartz_api/internal/service/uk_national/national_router.py index 9ca7eeec..e772bd9d 100644 --- a/src/quartz_api/internal/service/uk_national/national_router.py +++ b/src/quartz_api/internal/service/uk_national/national_router.py @@ -25,7 +25,7 @@ gsp_id_map, model_names_external_to_internal, ) -from .time_utils import limit_end_datetime_by_permissions +from .time_utils import limit_end_datetime_by_permissions, trial_expired log = logging.getLogger(__name__) @@ -89,7 +89,7 @@ async def get_national_forecast( request: Request, # noqa: ARG001 db: models.StorageClientDependency, - auth: AuthDependency, # noqa: ARG001 + auth: AuthDependency, end_datetime_utc: Annotated[ models.UTCDatetimeDefaultWindowEndNonAware, Depends(limit_end_datetime_by_permissions), @@ -143,6 +143,11 @@ async def get_national_forecast( days=3, ) + # end_datetime_utc is already clamped to "now" for expired trials. Clamp + # start_datetime_utc too, so it can't end up later than end_datetime_utc. + if trial_expired(auth, dt.datetime.now(dt.UTC)): + start_datetime_utc = min(start_datetime_utc, end_datetime_utc) + windows: list[tuple[dt.datetime, dt.datetime]] = [(start_datetime_utc, end_datetime_utc)] if end_datetime_utc - start_datetime_utc > dt.timedelta(days=7): windows = [ diff --git a/src/quartz_api/internal/service/uk_national/test_time_utils.py b/src/quartz_api/internal/service/uk_national/test_time_utils.py new file mode 100644 index 00000000..ec7994a2 --- /dev/null +++ b/src/quartz_api/internal/service/uk_national/test_time_utils.py @@ -0,0 +1,18 @@ +"""Unit tests for the trial-expiry helper in uk_national time_utils.""" + +import datetime as dt + +from .time_utils import trial_expired + + +def test_trial_expired_true_when_claim_in_past() -> None: + now = dt.datetime(2026, 1, 1, tzinfo=dt.UTC) + auth = {"app_metadata": {"trial_ends_at": "2025-12-31T00:00:00.000Z"}} + assert trial_expired(auth, now) is True + + +def test_trial_expired_false_when_claim_missing_or_future() -> None: + now = dt.datetime(2026, 1, 1, tzinfo=dt.UTC) + assert trial_expired({}, now) is False + auth = {"app_metadata": {"trial_ends_at": "2026-06-01T00:00:00.000Z"}} + assert trial_expired(auth, now) is False diff --git a/src/quartz_api/internal/service/uk_national/time_utils.py b/src/quartz_api/internal/service/uk_national/time_utils.py index aeb85add..875136b2 100644 --- a/src/quartz_api/internal/service/uk_national/time_utils.py +++ b/src/quartz_api/internal/service/uk_national/time_utils.py @@ -13,12 +13,34 @@ # and the time clipping is done then INTRADAY_LIMIT_HOURS = int(os.getenv("INTRADAY_LIMIT_HOURS", 8)) +def trial_expired(auth: AuthDependency, now: dt.datetime) -> bool: + """Check whether the caller's trial has ended, per the app_metadata.trial_ends_at claim. + + No claim at all (paid/non-trial user) is treated the same as an unexpired + trial: both mean "don't clamp." Only a claim that has actually passed + counts as expired. + """ + app_metadata = auth.get("app_metadata", {}) + if not isinstance(app_metadata, dict): + return False + raw = app_metadata.get("trial_ends_at") + if raw is None: + return False + try: + ends_at = dt.datetime.fromisoformat(str(raw)) + except ValueError: + sentry_sdk.capture_message(f"Unparseable trial_ends_at claim: {raw!r}") + return False + if ends_at.tzinfo is None: + ends_at = ends_at.replace(tzinfo=dt.UTC) + return ends_at <= now + def limit_end_datetime_by_permissions( auth: AuthDependency, end_datetime_utc: models.UTCDatetimeDefaultWindowEndNonAware, ) -> dt.datetime: - """Ensures only users with required permissions can access intraday data.""" + """Ensures only users with required permissions/an active trial can see future data.""" permissions: str | list[str] = auth.get("permissions", []) if isinstance(permissions, str): permissions = [permissions] @@ -30,6 +52,8 @@ def limit_end_datetime_by_permissions( return end_datetime_utc intraday_max_allowed = dt.datetime.now(dt.UTC) + dt.timedelta(hours=INTRADAY_LIMIT_HOURS) + if trial_expired(auth, dt.datetime.now(dt.UTC)): + return min(end_datetime_utc, dt.datetime.now(dt.UTC)) if "read:uk-intraday" in permissions: return min(end_datetime_utc, intraday_max_allowed) diff --git a/src/quartz_api/internal/service/v1/cache.py b/src/quartz_api/internal/service/v1/cache.py index 1a23d4ce..7bcf0315 100644 --- a/src/quartz_api/internal/service/v1/cache.py +++ b/src/quartz_api/internal/service/v1/cache.py @@ -14,7 +14,7 @@ from .auth_scopes import ALL_COUNTRY_PERMISSIONS from .country_config import COUNTRIES -from .helpers import internal_to_api_name, timeseries_window, to_uuid +from .helpers import internal_to_api_name, timeseries_window, to_uuid, trial_expired log = logging.getLogger(__name__) @@ -70,6 +70,9 @@ async def key_builder( parts = [namespace, request.method.lower(), request.url.path, repr(sorted(params))] if tier is not None: parts.append(tier) + #Invalidte cache if trial has expired + if trial_expired(auth, dt.datetime.now(dt.UTC)): + parts.append("trial_ended=true") return ":".join(parts) diff --git a/src/quartz_api/internal/service/v1/helpers.py b/src/quartz_api/internal/service/v1/helpers.py index 16a708c5..045a32d8 100644 --- a/src/quartz_api/internal/service/v1/helpers.py +++ b/src/quartz_api/internal/service/v1/helpers.py @@ -248,6 +248,28 @@ def check_country_access(auth: AuthDependency, cfg: CountryConfig) -> bool: ) +def trial_expired(auth: AuthDependency, now: dt.datetime) -> bool: + """Check whether the caller's trial has ended, per app_metadata.trial_ends_at. + + No claim at all (paid/non-trial user) is treated the same as an unexpired + trial: both mean "don't clamp." Only a claim that has actually passed + counts as expired. + """ + app_metadata = auth.get("app_metadata", {}) + if not isinstance(app_metadata, dict): + return False + raw = app_metadata.get("trial_ends_at") + if raw is None: + return False + try: + ends_at = dt.datetime.fromisoformat(str(raw)) + except ValueError: + return False + if ends_at.tzinfo is None: + ends_at = ends_at.replace(tzinfo=dt.UTC) + return ends_at <= now + + def resolve_forecast_model( model: str | None, rt: RegionTypeConfig | None, diff --git a/src/quartz_api/internal/service/v1/routes/forecasts.py b/src/quartz_api/internal/service/v1/routes/forecasts.py index d57e7aec..f031f371 100644 --- a/src/quartz_api/internal/service/v1/routes/forecasts.py +++ b/src/quartz_api/internal/service/v1/routes/forecasts.py @@ -45,6 +45,7 @@ resolve_region_id, timeseries_window, to_uuid, + trial_expired, validate_model, validate_window, window_chunks, @@ -119,6 +120,11 @@ async def get_forecast( now = pd.Timestamp.utcnow().floor("30min").to_pydatetime() win_start = start_utc or now win_end = end_utc or now + dt.timedelta(days=2) + if trial_expired(auth, dt.datetime.now(dt.UTC)): + if win_end.tzinfo is None: + win_end = win_end.replace(tzinfo=dt.UTC) + win_end = min(win_end, dt.datetime.now(dt.UTC)) + win_start = min(win_start, win_end) validate_window(win_start, win_end) pgvs: list = [] for chunk_start, chunk_end in window_chunks(win_start, win_end): @@ -276,6 +282,8 @@ async def get_forecasts_at_time( snapshot_time = time_utc or pd.Timestamp.utcnow().floor("30min").to_pydatetime() if snapshot_time.tzinfo is None: snapshot_time = snapshot_time.replace(tzinfo=dt.UTC) + if trial_expired(auth, dt.datetime.now(dt.UTC)): + snapshot_time = min(snapshot_time, dt.datetime.now(dt.UTC)) snapshot = await db.get_predicted_generation_snapshot( location_uuids=[to_uuid(r.uuid) for r in regions], @@ -380,6 +388,9 @@ async def get_forecasts_period( ) win_start, win_end = timeseries_window(start_utc, end_utc) + if trial_expired(auth, dt.datetime.now(dt.UTC)): + win_end = min(win_end, dt.datetime.now(dt.UTC)) + win_start = min(win_start, win_end) validate_window(win_start, win_end) backend = FastAPICache.get_backend() diff --git a/src/quartz_api/internal/service/v1/test_helpers.py b/src/quartz_api/internal/service/v1/test_helpers.py new file mode 100644 index 00000000..390e821d --- /dev/null +++ b/src/quartz_api/internal/service/v1/test_helpers.py @@ -0,0 +1,18 @@ +"""Unit tests for the trial-expiry helper in v1 helpers.""" + +import datetime as dt + +from .helpers import trial_expired + + +def test_trial_expired_true_when_claim_in_past() -> None: + now = dt.datetime(2026, 1, 1, tzinfo=dt.UTC) + auth = {"app_metadata": {"trial_ends_at": "2025-12-31T00:00:00.000Z"}} + assert trial_expired(auth, now) is True + + +def test_trial_expired_false_when_claim_missing_or_future() -> None: + now = dt.datetime(2026, 1, 1, tzinfo=dt.UTC) + assert trial_expired({}, now) is False + auth = {"app_metadata": {"trial_ends_at": "2026-06-01T00:00:00.000Z"}} + assert trial_expired(auth, now) is False