feat: add trial check and don't allow forecast access

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

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(

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 = [

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

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)
 

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)
 
 

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,

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()

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