fix(proxy): rework the budgets list onto the merged list contract

PR #35308 landed a different shape than this branch was written against: `where`
is a tuple of frozen predicates rather than a Prisma-shaped mapping, `ListSpec`
carries both the row and the wire type, and `where_sql` / `order_by_sql` render
for a raw-SQL executor. The budgets executor now queries through `query_raw` the
way the spend logs facet does, selecting only the columns it serves.

Also casts datetime binds in `where_sql`. They cross into the query engine as
JSON, so an uncast placeholder arrives as text and Postgres refuses
`timestamp >= text` outright; every `filter[created_at][gte|lte]` was answering
500. The cast reads the bind as an instant and drops it to naive UTC to match
Prisma's TIMESTAMP(3) column, the same one /spend/logs/ui applies.
This commit is contained in:
Yuneng Jiang 2026-07-31 10:19:24 -07:00
parent f0d1bc3c5a
commit 78c756dff9
No known key found for this signature in database
4 changed files with 87 additions and 29 deletions

View file

@ -1,8 +1,9 @@
"""`GET /management/v1/budgets`."""
from collections.abc import Sequence
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from datetime import datetime
from types import MappingProxyType
from typing import Annotated
from fastapi import APIRouter, Depends, Request
@ -112,15 +113,19 @@ def _scope(caller: UserAPIKeyAuth) -> Scope:
# budget_duration is deliberately absent from `sortable`: the column holds strings
# like "7d" and "30d", so a lexicographic ORDER BY puts "30d" ahead of "7d".
BUDGET_FILTERS: Mapping[str, FilterSpec] = MappingProxyType(
{ # mutable-ok: an immutable mapping has no literal form; MappingProxyType freezes this one and it never escapes
"budget_duration": FilterSpec(type=str, ops=frozenset(("in", "is_null"))),
"max_budget": FilterSpec(type=float, ops=frozenset(("gte", "lte", "is_null"))),
"created_at": FilterSpec(type=datetime, ops=frozenset(("gte", "lte"))),
}
)
BUDGETS_LIST_SPEC: ListSpec[BudgetListItem, BudgetListItem] = ListSpec(
resource="budgets",
sortable=frozenset({"budget_id", "max_budget", "tpm_limit", "rpm_limit", "created_at"}),
searchable=frozenset({"budget_id"}),
filters={
"budget_duration": FilterSpec(type=str, ops=frozenset({"in", "is_null"})),
"max_budget": FilterSpec(type=float, ops=frozenset({"gte", "lte", "is_null"})),
"created_at": FilterSpec(type=datetime, ops=frozenset({"gte", "lte"})),
},
sortable=frozenset(("budget_id", "max_budget", "tpm_limit", "rpm_limit", "created_at")),
searchable=frozenset(("budget_id",)),
filters=BUDGET_FILTERS,
default_sort=(SortKey(field="created_at", descending=True),),
default_page_size=50,
max_page_size=100,
@ -132,8 +137,8 @@ BUDGETS_LIST_SPEC: ListSpec[BudgetListItem, BudgetListItem] = ListSpec(
@router.get(
"/budgets",
tags=["budget management"],
dependencies=[Depends(user_api_key_auth)],
tags=("budget management",),
dependencies=(Depends(user_api_key_auth),),
response_model=ListResponse[BudgetListItem],
)
async def list_budgets(

View file

@ -213,12 +213,23 @@ def _sql_operator(op: ComparisonOp) -> str:
assert_never(op)
def _placeholder(index: int, value: FilterValue) -> str:
"""`$n`, cast when the bind is a datetime.
Binds cross into the query engine as JSON, so a datetime arrives as text and
Postgres refuses `timestamp >= text` outright. Prisma stores DateTime as a naive
`TIMESTAMP(3)` holding UTC, so the bind is read as an instant and then dropped to
naive UTC to match the column, the same cast `/spend/logs/ui` applies.
"""
return f"${index}::timestamptz AT TIME ZONE 'UTC'" if isinstance(value, datetime) else f"${index}"
def _render(predicate: Predicate, index: int) -> tuple[str, tuple[object, ...]]:
match predicate:
case IsNull(field=field, negated=negated):
return f'"{field}" IS {"NOT NULL" if negated else "NULL"}', ()
case Within(field=field, values=values):
placeholders = ", ".join(f"${index + offset}" for offset in range(len(values)))
placeholders = ", ".join(_placeholder(index + offset, value) for offset, value in enumerate(values))
return f'"{field}" IN ({placeholders})', values
case AnyOf(clauses=clauses):
rendered, params = _render_all(clauses, index)
@ -226,7 +237,7 @@ def _render(predicate: Predicate, index: int) -> tuple[str, tuple[object, ...]]:
case Compare(field=field, op="contains", value=value):
return f"\"{field}\" ILIKE ${index} ESCAPE '\\'", (f"%{escape_like(str(value))}%",)
case Compare(field=field, op=op, value=value):
return f'"{field}" {_sql_operator(op)} ${index}', (value,)
return f'"{field}" {_sql_operator(op)} {_placeholder(index, value)}', (value,)
case _:
assert_never(predicate)

View file

@ -371,15 +371,28 @@ def test_in_filter_binds_each_requested_duration(query_raw, as_proxy_admin):
def test_created_at_range_is_bound_as_a_timestamp(query_raw, as_proxy_admin):
"""The bind crosses into the query engine as JSON, so an uncast placeholder reaches
Postgres as text and `timestamp >= text` is a hard error, not a wrong answer."""
_serve(query_raw, [])
_get("filter[created_at][gte]=2026-07-01T00:00:00Z")
sql, *params = _select_call(query_raw)
assert '"created_at" >= $1' in sql
assert "\"created_at\" >= $1::timestamptz AT TIME ZONE 'UTC'" in sql
assert params[0] == datetime(2026, 7, 1, tzinfo=timezone.utc)
def test_a_non_datetime_bind_is_not_cast(query_raw, as_proxy_admin):
"""Guards the cast above from being applied to every placeholder."""
_serve(query_raw, [])
_get("filter[max_budget][gte]=5")
sql = _select_call(query_raw)[0]
assert '"max_budget" >= $1' in sql
assert "timestamptz" not in sql
def test_an_offsetless_created_at_bound_is_read_as_utc(query_raw, as_proxy_admin):
"""The dashboard sends 'YYYY-MM-DDTHH:MM:SS' with no offset. Left naive, Postgres
would compare it in the session timezone and shift the window off the rows shown."""

View file

@ -7187,10 +7187,11 @@ export interface paths {
*
* `sort` takes a comma-separated list of `budget_id`, `max_budget`, `tpm_limit`,
* `rpm_limit` or `created_at`, each optionally prefixed with `-` for descending,
* and defaults to `-created_at,budget_id`. `q` is a case-insensitive substring
* match on `budget_id`. `page_size` defaults to 50 and is capped at 100.
* Filters are `filter[budget_duration][in|is_null]`,
* `filter[max_budget][gte|lte|is_null]` and `filter[created_at][gte|lte]`.
* and defaults to `-created_at`. `budget_id` is appended to every sort as the
* tiebreaker. `q` is a case-insensitive substring match on `budget_id`.
* `page_size` defaults to 50 and is capped at 100. Filters are
* `filter[budget_duration][in|is_null]`, `filter[max_budget][gte|lte|is_null]`
* and `filter[created_at][gte|lte]`.
*
* Example curl:
* ```
@ -21570,6 +21571,40 @@ export interface components {
/** Reset At */
reset_at?: string | null;
};
/**
* BudgetListItem
* @description One budget as the Budgets page reads it, and as it comes back off the table.
*
* Validating the raw row through here is what makes `tpm_limit` / `rpm_limit`
* numbers: they are `BigInt?` in the schema, which the query engine hands back as
* decimal strings, and a quoted "60000" breaks arithmetic in the dashboard.
*/
BudgetListItem: {
/** Budget Duration */
budget_duration?: string | null;
/** Budget Id */
budget_id: string;
/** Budget Reset At */
budget_reset_at?: string | null;
/**
* Created At
* Format: date-time
*/
created_at: string;
/** Max Budget */
max_budget?: number | null;
/** Rpm Limit */
rpm_limit?: number | null;
/** Soft Budget */
soft_budget?: number | null;
/** Tpm Limit */
tpm_limit?: number | null;
/**
* Updated At
* Format: date-time
*/
updated_at: string;
};
/** BudgetNewRequest */
BudgetNewRequest: {
/**
@ -24814,7 +24849,6 @@ export interface components {
/** Updated By */
updated_by?: string | null;
};
JsonValue: unknown;
/** KeyHealthResponse */
KeyHealthResponse: {
/**
@ -24968,7 +25002,7 @@ export interface components {
};
/**
* ListLinks
* @description Hypermedia for an entity list. `first`/`last` exist here because `total_pages` is known.
* @description Page-mode counterpart to `PageLinks`. `first`/`last` are knowable here because the total count is.
*/
ListLinks: {
/** First */
@ -24984,7 +25018,7 @@ export interface components {
};
/**
* ListMeta
* @description An entity list can afford the COUNT(*) a facet cannot, so it reports a real total.
* @description Page-mode counterpart to `PageMeta`: an entity list pays for the COUNT(*) so the table can show a page count.
*/
ListMeta: {
/** Page */
@ -25011,15 +25045,10 @@ export interface components {
/** Prompts */
prompts: components["schemas"]["PromptSpec"][];
};
/**
* ListResponse
* @description One page of an entity collection. Rows are flat: no `{type, id, attributes}` wrapper.
*/
ListResponse: {
/** ListResponse[BudgetListItem] */
ListResponse_BudgetListItem_: {
/** Data */
data: {
[key: string]: components["schemas"]["JsonValue"];
}[];
data: components["schemas"]["BudgetListItem"][];
links: components["schemas"]["ListLinks"];
meta: components["schemas"]["ListMeta"];
};
@ -43574,7 +43603,7 @@ export interface operations {
[name: string]: unknown;
};
content: {
"application/json": components["schemas"]["ListResponse"];
"application/json": components["schemas"]["ListResponse_BudgetListItem_"];
};
};
};