mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-03 01:40:21 +08:00
fix provider pool quota status handling
This commit is contained in:
fawney19
@@ -1333,15 +1333,20 @@ const REBUILD_API_KEY_USAGE_STATS_SQL: &str =
|
||||
const APPLY_PROVIDER_API_KEY_USAGE_DELTA_SQL: &str =
|
||||
include_str!("queries/apply_provider_api_key_usage_delta_sql.sql");
|
||||
|
||||
const APPLY_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_DELTA_SQL: &str =
|
||||
include_str!("queries/apply_provider_api_key_codex_window_usage_delta_sql.sql");
|
||||
|
||||
const RESET_PROVIDER_API_KEY_USAGE_STATS_SQL: &str =
|
||||
include_str!("queries/reset_provider_api_key_usage_stats_sql.sql");
|
||||
|
||||
const REBUILD_PROVIDER_API_KEY_USAGE_STATS_SQL: &str =
|
||||
include_str!("queries/rebuild_provider_api_key_usage_stats_sql.sql");
|
||||
|
||||
const REBUILD_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_STATS_SQL: &str =
|
||||
include_str!("queries/rebuild_provider_api_key_codex_window_usage_stats_sql.sql");
|
||||
|
||||
const LIST_USAGE_AUDITS_PREFIX: &str = include_str!("queries/list_usage_audits_prefix.sql");
|
||||
const USAGE_RESERVED_PROVIDER_LABELS_FILTER_SQL: &str =
|
||||
" AND BTRIM(COALESCE(\"usage\".provider_name, '')) <> '' AND lower(BTRIM(COALESCE(\"usage\".provider_name, ''))) NOT IN ('unknown', 'unknow', 'pending')";
|
||||
const USAGE_RESERVED_PROVIDER_LABELS_FILTER_SQL: &str = " AND BTRIM(COALESCE(\"usage\".provider_name, '')) <> '' AND lower(BTRIM(COALESCE(\"usage\".provider_name, ''))) NOT IN ('unknown', 'unknow', 'pending')";
|
||||
|
||||
struct UsageAuditAggregationSqlFragments {
|
||||
filtered_extra_where: &'static str,
|
||||
@@ -6395,33 +6400,38 @@ WHERE stats_daily_api_key.date >=
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let (table_name, group_column, display_name_expr, avg_response_time_expr, success_count_expr) =
|
||||
match group_by {
|
||||
UsageAuditAggregationGroupBy::Model => (
|
||||
"stats_user_daily_model",
|
||||
"model",
|
||||
"NULL::varchar",
|
||||
"NULL::DOUBLE PRECISION",
|
||||
"NULL::BIGINT",
|
||||
),
|
||||
UsageAuditAggregationGroupBy::Provider => (
|
||||
"stats_user_daily_provider",
|
||||
"provider_name",
|
||||
"provider_name",
|
||||
"CASE WHEN COALESCE(SUM(response_time_samples), 0) > 0 THEN COALESCE(SUM(response_time_sum_ms), 0) / COALESCE(SUM(response_time_samples), 0) ELSE NULL END",
|
||||
"COALESCE(SUM(success_requests), 0)::BIGINT",
|
||||
),
|
||||
UsageAuditAggregationGroupBy::ApiFormat => (
|
||||
"stats_user_daily_api_format",
|
||||
"api_format",
|
||||
"NULL::varchar",
|
||||
"CASE WHEN COALESCE(SUM(response_time_samples), 0) > 0 THEN COALESCE(SUM(response_time_sum_ms), 0) / COALESCE(SUM(response_time_samples), 0) ELSE NULL END",
|
||||
"NULL::BIGINT",
|
||||
),
|
||||
UsageAuditAggregationGroupBy::User => {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
};
|
||||
let (
|
||||
table_name,
|
||||
group_column,
|
||||
display_name_expr,
|
||||
avg_response_time_expr,
|
||||
success_count_expr,
|
||||
) = match group_by {
|
||||
UsageAuditAggregationGroupBy::Model => (
|
||||
"stats_user_daily_model",
|
||||
"model",
|
||||
"NULL::varchar",
|
||||
"NULL::DOUBLE PRECISION",
|
||||
"NULL::BIGINT",
|
||||
),
|
||||
UsageAuditAggregationGroupBy::Provider => (
|
||||
"stats_user_daily_provider",
|
||||
"provider_name",
|
||||
"provider_name",
|
||||
"CASE WHEN COALESCE(SUM(response_time_samples), 0) > 0 THEN COALESCE(SUM(response_time_sum_ms), 0) / COALESCE(SUM(response_time_samples), 0) ELSE NULL END",
|
||||
"COALESCE(SUM(success_requests), 0)::BIGINT",
|
||||
),
|
||||
UsageAuditAggregationGroupBy::ApiFormat => (
|
||||
"stats_user_daily_api_format",
|
||||
"api_format",
|
||||
"NULL::varchar",
|
||||
"CASE WHEN COALESCE(SUM(response_time_samples), 0) > 0 THEN COALESCE(SUM(response_time_sum_ms), 0) / COALESCE(SUM(response_time_samples), 0) ELSE NULL END",
|
||||
"NULL::BIGINT",
|
||||
),
|
||||
UsageAuditAggregationGroupBy::User => {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
};
|
||||
|
||||
let provider_extra_where = if matches!(group_by, UsageAuditAggregationGroupBy::Provider) {
|
||||
" AND BTRIM(COALESCE(provider_name, '')) <> '' AND lower(BTRIM(COALESCE(provider_name, ''))) NOT IN ('unknown', 'unknow', 'pending')"
|
||||
@@ -7925,6 +7935,10 @@ ORDER BY "usage".user_id ASC
|
||||
.await
|
||||
.map_postgres_err()?
|
||||
.rows_affected();
|
||||
sqlx::query(REBUILD_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_STATS_SQL)
|
||||
.execute(&mut **tx)
|
||||
.await
|
||||
.map_postgres_err()?;
|
||||
Ok(rows_affected)
|
||||
}) as BoxFuture<'_, Result<u64, DataLayerError>>
|
||||
})
|
||||
@@ -8371,6 +8385,40 @@ async fn apply_provider_api_key_usage_delta_in_tx(
|
||||
.execute(&mut **tx)
|
||||
.await
|
||||
.map_postgres_err()?;
|
||||
apply_provider_api_key_codex_window_usage_delta_in_tx(tx, key_id, delta).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn apply_provider_api_key_codex_window_usage_delta_in_tx(
|
||||
tx: &mut sqlx::Transaction<'_, Postgres>,
|
||||
key_id: &str,
|
||||
delta: &ProviderApiKeyUsageDelta,
|
||||
) -> Result<(), DataLayerError> {
|
||||
let Some(usage_created_at_unix_secs) = delta.usage_created_at_unix_secs else {
|
||||
return Ok(());
|
||||
};
|
||||
if delta.request_count == 0 && delta.total_tokens == 0 && delta.total_cost_usd == 0.0 {
|
||||
return Ok(());
|
||||
}
|
||||
let total_cost_usd_delta = if delta.total_cost_usd.is_finite() {
|
||||
delta.total_cost_usd
|
||||
} else {
|
||||
0.0
|
||||
};
|
||||
|
||||
sqlx::query(APPLY_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_DELTA_SQL)
|
||||
.bind(key_id)
|
||||
.bind(i64::try_from(usage_created_at_unix_secs).map_err(|_| {
|
||||
DataLayerError::UnexpectedValue(format!(
|
||||
"provider api key window usage timestamp exceeds i64: {usage_created_at_unix_secs}"
|
||||
))
|
||||
})?)
|
||||
.bind(delta.request_count)
|
||||
.bind(delta.total_tokens)
|
||||
.bind(total_cost_usd_delta)
|
||||
.execute(&mut **tx)
|
||||
.await
|
||||
.map_postgres_err()?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,166 @@
|
||||
WITH target_key AS (
|
||||
SELECT
|
||||
id,
|
||||
COALESCE(status_snapshot::jsonb, '{}'::jsonb) AS snapshot
|
||||
FROM provider_api_keys
|
||||
WHERE id = $1
|
||||
AND jsonb_typeof((status_snapshot::jsonb) -> 'quota' -> 'windows') = 'array'
|
||||
AND lower(BTRIM(COALESCE((status_snapshot::jsonb) -> 'quota' ->> 'provider_type', ''))) = 'codex'
|
||||
FOR UPDATE
|
||||
),
|
||||
window_items AS (
|
||||
SELECT
|
||||
target_key.id,
|
||||
window_rows.window_item,
|
||||
window_rows.window_ordinality
|
||||
FROM target_key
|
||||
CROSS JOIN LATERAL jsonb_array_elements(target_key.snapshot -> 'quota' -> 'windows')
|
||||
WITH ORDINALITY AS window_rows(window_item, window_ordinality)
|
||||
),
|
||||
parsed_windows AS (
|
||||
SELECT
|
||||
window_items.id,
|
||||
window_items.window_item,
|
||||
window_items.window_ordinality,
|
||||
lower(BTRIM(COALESCE(window_items.window_item ->> 'code', ''))) AS window_code,
|
||||
CASE
|
||||
WHEN text_values.reset_at_text ~ '^[0-9]+$'
|
||||
AND (
|
||||
length(text_values.reset_at_text) < 19
|
||||
OR (
|
||||
length(text_values.reset_at_text) = 19
|
||||
AND text_values.reset_at_text <= '9223372036854775807'
|
||||
)
|
||||
)
|
||||
THEN text_values.reset_at_text::BIGINT
|
||||
ELSE NULL
|
||||
END AS reset_at,
|
||||
CASE
|
||||
WHEN text_values.window_minutes_text ~ '^[0-9]+$'
|
||||
AND (
|
||||
length(text_values.window_minutes_text) < 19
|
||||
OR (
|
||||
length(text_values.window_minutes_text) = 19
|
||||
AND text_values.window_minutes_text <= '9223372036854775807'
|
||||
)
|
||||
)
|
||||
THEN text_values.window_minutes_text::BIGINT
|
||||
ELSE NULL
|
||||
END AS window_minutes,
|
||||
CASE
|
||||
WHEN text_values.usage_reset_at_text ~ '^[0-9]+$'
|
||||
AND (
|
||||
length(text_values.usage_reset_at_text) < 19
|
||||
OR (
|
||||
length(text_values.usage_reset_at_text) = 19
|
||||
AND text_values.usage_reset_at_text <= '9223372036854775807'
|
||||
)
|
||||
)
|
||||
THEN text_values.usage_reset_at_text::BIGINT
|
||||
ELSE NULL
|
||||
END AS usage_reset_at,
|
||||
CASE
|
||||
WHEN text_values.request_count_text ~ '^[0-9]+$'
|
||||
AND (
|
||||
length(text_values.request_count_text) < 19
|
||||
OR (
|
||||
length(text_values.request_count_text) = 19
|
||||
AND text_values.request_count_text <= '9223372036854775807'
|
||||
)
|
||||
)
|
||||
THEN text_values.request_count_text::BIGINT
|
||||
ELSE 0
|
||||
END AS current_request_count,
|
||||
CASE
|
||||
WHEN text_values.total_tokens_text ~ '^[0-9]+$'
|
||||
AND (
|
||||
length(text_values.total_tokens_text) < 19
|
||||
OR (
|
||||
length(text_values.total_tokens_text) = 19
|
||||
AND text_values.total_tokens_text <= '9223372036854775807'
|
||||
)
|
||||
)
|
||||
THEN text_values.total_tokens_text::BIGINT
|
||||
ELSE 0
|
||||
END AS current_total_tokens,
|
||||
CASE
|
||||
WHEN text_values.total_cost_usd_text ~ '^[-+]?[0-9]+([.][0-9]+)?$'
|
||||
THEN text_values.total_cost_usd_text::DOUBLE PRECISION
|
||||
ELSE 0
|
||||
END AS current_total_cost_usd
|
||||
FROM window_items
|
||||
CROSS JOIN LATERAL (
|
||||
SELECT
|
||||
BTRIM(COALESCE(window_items.window_item ->> 'reset_at', '')) AS reset_at_text,
|
||||
BTRIM(COALESCE(window_items.window_item ->> 'window_minutes', '')) AS window_minutes_text,
|
||||
BTRIM(COALESCE(window_items.window_item ->> 'usage_reset_at', '')) AS usage_reset_at_text,
|
||||
BTRIM(COALESCE(window_items.window_item -> 'usage' ->> 'request_count', '')) AS request_count_text,
|
||||
BTRIM(COALESCE(window_items.window_item -> 'usage' ->> 'total_tokens', '')) AS total_tokens_text,
|
||||
BTRIM(COALESCE(window_items.window_item -> 'usage' ->> 'total_cost_usd', '')) AS total_cost_usd_text
|
||||
) AS text_values
|
||||
),
|
||||
window_usage AS (
|
||||
SELECT
|
||||
parsed_windows.*,
|
||||
CASE
|
||||
WHEN parsed_windows.window_minutes BETWEEN 0 AND 153722867280912930
|
||||
THEN parsed_windows.window_minutes * 60
|
||||
ELSE NULL
|
||||
END AS window_seconds
|
||||
FROM parsed_windows
|
||||
),
|
||||
updated_windows AS (
|
||||
SELECT
|
||||
window_usage.id,
|
||||
jsonb_agg(
|
||||
CASE
|
||||
WHEN window_usage.window_code IN ('5h', 'weekly')
|
||||
AND window_usage.reset_at IS NOT NULL
|
||||
AND window_usage.window_seconds IS NOT NULL
|
||||
AND window_usage.reset_at >= window_usage.window_seconds
|
||||
AND window_usage.reset_at > $2
|
||||
AND GREATEST(
|
||||
window_usage.reset_at - window_usage.window_seconds,
|
||||
COALESCE(window_usage.usage_reset_at, 0)
|
||||
) <= $2
|
||||
THEN jsonb_set(
|
||||
window_usage.window_item,
|
||||
'{usage}',
|
||||
jsonb_build_object(
|
||||
'request_count',
|
||||
LEAST(
|
||||
GREATEST(window_usage.current_request_count::NUMERIC + $3::NUMERIC, 0),
|
||||
9223372036854775807
|
||||
)::BIGINT,
|
||||
'total_tokens',
|
||||
LEAST(
|
||||
GREATEST(window_usage.current_total_tokens::NUMERIC + $4::NUMERIC, 0),
|
||||
9223372036854775807
|
||||
)::BIGINT,
|
||||
'total_cost_usd',
|
||||
to_char(
|
||||
GREATEST(COALESCE(window_usage.current_total_cost_usd, 0) + $5, 0),
|
||||
'FM999999999999999990.00000000'
|
||||
)
|
||||
),
|
||||
true
|
||||
)
|
||||
ELSE window_usage.window_item
|
||||
END
|
||||
ORDER BY window_usage.window_ordinality
|
||||
) AS windows
|
||||
FROM window_usage
|
||||
GROUP BY window_usage.id
|
||||
)
|
||||
UPDATE provider_api_keys AS keys
|
||||
SET
|
||||
status_snapshot = jsonb_set(
|
||||
target_key.snapshot,
|
||||
'{quota,windows}',
|
||||
updated_windows.windows,
|
||||
true
|
||||
)::json,
|
||||
updated_at = NOW()
|
||||
FROM target_key
|
||||
JOIN updated_windows ON updated_windows.id = target_key.id
|
||||
WHERE keys.id = target_key.id
|
||||
@@ -0,0 +1,164 @@
|
||||
WITH target_keys AS (
|
||||
SELECT
|
||||
id,
|
||||
COALESCE(status_snapshot::jsonb, '{}'::jsonb) AS snapshot
|
||||
FROM provider_api_keys
|
||||
WHERE jsonb_typeof((status_snapshot::jsonb) -> 'quota' -> 'windows') = 'array'
|
||||
AND lower(BTRIM(COALESCE((status_snapshot::jsonb) -> 'quota' ->> 'provider_type', ''))) = 'codex'
|
||||
),
|
||||
window_items AS (
|
||||
SELECT
|
||||
target_keys.id,
|
||||
window_rows.window_item,
|
||||
window_rows.window_ordinality
|
||||
FROM target_keys
|
||||
CROSS JOIN LATERAL jsonb_array_elements(target_keys.snapshot -> 'quota' -> 'windows')
|
||||
WITH ORDINALITY AS window_rows(window_item, window_ordinality)
|
||||
),
|
||||
parsed_windows AS (
|
||||
SELECT
|
||||
window_items.id,
|
||||
window_items.window_item,
|
||||
window_items.window_ordinality,
|
||||
lower(BTRIM(COALESCE(window_items.window_item ->> 'code', ''))) AS window_code,
|
||||
CASE
|
||||
WHEN text_values.reset_at_text ~ '^[0-9]+$'
|
||||
AND (
|
||||
length(text_values.reset_at_text) < 19
|
||||
OR (
|
||||
length(text_values.reset_at_text) = 19
|
||||
AND text_values.reset_at_text <= '9223372036854775807'
|
||||
)
|
||||
)
|
||||
THEN text_values.reset_at_text::BIGINT
|
||||
ELSE NULL
|
||||
END AS reset_at,
|
||||
CASE
|
||||
WHEN text_values.window_minutes_text ~ '^[0-9]+$'
|
||||
AND (
|
||||
length(text_values.window_minutes_text) < 19
|
||||
OR (
|
||||
length(text_values.window_minutes_text) = 19
|
||||
AND text_values.window_minutes_text <= '9223372036854775807'
|
||||
)
|
||||
)
|
||||
THEN text_values.window_minutes_text::BIGINT
|
||||
ELSE NULL
|
||||
END AS window_minutes,
|
||||
CASE
|
||||
WHEN text_values.usage_reset_at_text ~ '^[0-9]+$'
|
||||
AND (
|
||||
length(text_values.usage_reset_at_text) < 19
|
||||
OR (
|
||||
length(text_values.usage_reset_at_text) = 19
|
||||
AND text_values.usage_reset_at_text <= '9223372036854775807'
|
||||
)
|
||||
)
|
||||
THEN text_values.usage_reset_at_text::BIGINT
|
||||
ELSE NULL
|
||||
END AS usage_reset_at
|
||||
FROM window_items
|
||||
CROSS JOIN LATERAL (
|
||||
SELECT
|
||||
BTRIM(COALESCE(window_items.window_item ->> 'reset_at', '')) AS reset_at_text,
|
||||
BTRIM(COALESCE(window_items.window_item ->> 'window_minutes', '')) AS window_minutes_text,
|
||||
BTRIM(COALESCE(window_items.window_item ->> 'usage_reset_at', '')) AS usage_reset_at_text
|
||||
) AS text_values
|
||||
),
|
||||
window_usage AS (
|
||||
SELECT
|
||||
parsed_windows.*,
|
||||
CASE
|
||||
WHEN parsed_windows.window_minutes BETWEEN 0 AND 153722867280912930
|
||||
THEN parsed_windows.window_minutes * 60
|
||||
ELSE NULL
|
||||
END AS window_seconds
|
||||
FROM parsed_windows
|
||||
),
|
||||
window_bounds AS (
|
||||
SELECT
|
||||
window_usage.*,
|
||||
CASE
|
||||
WHEN window_usage.window_code IN ('5h', 'weekly')
|
||||
AND window_usage.reset_at IS NOT NULL
|
||||
AND window_usage.window_seconds IS NOT NULL
|
||||
AND window_usage.reset_at >= window_usage.window_seconds
|
||||
THEN GREATEST(
|
||||
window_usage.reset_at - window_usage.window_seconds,
|
||||
COALESCE(window_usage.usage_reset_at, 0)
|
||||
)
|
||||
ELSE NULL
|
||||
END AS window_start,
|
||||
CASE
|
||||
WHEN window_usage.window_code IN ('5h', 'weekly')
|
||||
AND window_usage.reset_at IS NOT NULL
|
||||
AND window_usage.window_seconds IS NOT NULL
|
||||
AND window_usage.reset_at >= window_usage.window_seconds
|
||||
THEN window_usage.reset_at
|
||||
ELSE NULL
|
||||
END AS window_end
|
||||
FROM window_usage
|
||||
),
|
||||
aggregated AS (
|
||||
SELECT
|
||||
window_bounds.id,
|
||||
window_bounds.window_ordinality,
|
||||
COUNT("usage".id)::BIGINT AS request_count,
|
||||
COALESCE(SUM(GREATEST(COALESCE("usage".total_tokens, 0), 0)::BIGINT), 0)::BIGINT AS total_tokens,
|
||||
CAST(COALESCE(SUM(COALESCE("usage".total_cost_usd, 0)), 0) AS DOUBLE PRECISION) AS total_cost_usd
|
||||
FROM window_bounds
|
||||
LEFT JOIN usage_billing_facts AS "usage"
|
||||
ON window_bounds.window_start IS NOT NULL
|
||||
AND window_bounds.window_end IS NOT NULL
|
||||
AND "usage".provider_api_key_id = window_bounds.id
|
||||
AND "usage".created_at >= to_timestamp(window_bounds.window_start::DOUBLE PRECISION)
|
||||
AND "usage".created_at < to_timestamp(window_bounds.window_end::DOUBLE PRECISION)
|
||||
GROUP BY
|
||||
window_bounds.id,
|
||||
window_bounds.window_ordinality
|
||||
),
|
||||
updated_windows AS (
|
||||
SELECT
|
||||
window_bounds.id,
|
||||
jsonb_agg(
|
||||
CASE
|
||||
WHEN window_bounds.window_start IS NOT NULL
|
||||
AND window_bounds.window_end IS NOT NULL
|
||||
THEN jsonb_set(
|
||||
window_bounds.window_item,
|
||||
'{usage}',
|
||||
jsonb_build_object(
|
||||
'request_count',
|
||||
COALESCE(aggregated.request_count, 0),
|
||||
'total_tokens',
|
||||
COALESCE(aggregated.total_tokens, 0),
|
||||
'total_cost_usd',
|
||||
to_char(
|
||||
GREATEST(COALESCE(aggregated.total_cost_usd, 0), 0),
|
||||
'FM999999999999999990.00000000'
|
||||
)
|
||||
),
|
||||
true
|
||||
)
|
||||
ELSE window_bounds.window_item
|
||||
END
|
||||
ORDER BY window_bounds.window_ordinality
|
||||
) AS windows
|
||||
FROM window_bounds
|
||||
LEFT JOIN aggregated
|
||||
ON aggregated.id = window_bounds.id
|
||||
AND aggregated.window_ordinality = window_bounds.window_ordinality
|
||||
GROUP BY window_bounds.id
|
||||
)
|
||||
UPDATE provider_api_keys AS keys
|
||||
SET
|
||||
status_snapshot = jsonb_set(
|
||||
target_keys.snapshot,
|
||||
'{quota,windows}',
|
||||
updated_windows.windows,
|
||||
true
|
||||
)::json,
|
||||
updated_at = NOW()
|
||||
FROM target_keys
|
||||
JOIN updated_windows ON updated_windows.id = target_keys.id
|
||||
WHERE keys.id = target_keys.id
|
||||
@@ -235,18 +235,31 @@ fn usage_sql_summarizes_usage_by_provider_api_key_ids_in_database() {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn usage_sql_summarizes_provider_key_window_usage_from_billing_facts() {
|
||||
assert!(super::SUMMARIZE_PROVIDER_API_KEY_WINDOW_USAGE_SQL.contains("UNNEST"));
|
||||
assert!(super::SUMMARIZE_PROVIDER_API_KEY_WINDOW_USAGE_SQL
|
||||
.contains("LEFT JOIN usage_billing_facts AS \"usage\""));
|
||||
fn usage_sql_materializes_provider_key_window_usage_in_status_snapshot() {
|
||||
assert!(super::APPLY_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_DELTA_SQL
|
||||
.contains("UPDATE provider_api_keys AS keys"));
|
||||
assert!(super::APPLY_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_DELTA_SQL.contains("jsonb_set"));
|
||||
assert!(
|
||||
super::SUMMARIZE_PROVIDER_API_KEY_WINDOW_USAGE_SQL.contains("created_at >= to_timestamp")
|
||||
super::APPLY_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_DELTA_SQL.contains("'{quota,windows}'")
|
||||
);
|
||||
assert!(super::APPLY_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_DELTA_SQL.contains("'usage'"));
|
||||
assert!(
|
||||
super::SUMMARIZE_PROVIDER_API_KEY_WINDOW_USAGE_SQL.contains("created_at < to_timestamp")
|
||||
!super::APPLY_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_DELTA_SQL.contains("usage_billing_facts")
|
||||
);
|
||||
assert!(super::SUMMARIZE_PROVIDER_API_KEY_WINDOW_USAGE_SQL
|
||||
.contains("COUNT(\"usage\".id)::BIGINT AS request_count"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn usage_sql_rebuilds_provider_key_window_usage_into_status_snapshot() {
|
||||
assert!(super::REBUILD_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_STATS_SQL
|
||||
.contains("UPDATE provider_api_keys AS keys"));
|
||||
assert!(super::REBUILD_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_STATS_SQL
|
||||
.contains("usage_billing_facts"));
|
||||
assert!(super::REBUILD_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_STATS_SQL
|
||||
.contains("provider_type', ''))) = 'codex'"));
|
||||
assert!(
|
||||
super::REBUILD_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_STATS_SQL.contains("'{quota,windows}'")
|
||||
);
|
||||
assert!(super::REBUILD_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_STATS_SQL.contains("'{usage}'"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -470,7 +483,7 @@ fn usage_sql_raw_aggregates_use_canonical_billing_facts() {
|
||||
.contains("FROM usage_billing_facts AS \"usage\""));
|
||||
assert!(super::SUMMARIZE_USAGE_TOTALS_BY_USER_IDS_SQL
|
||||
.contains("FROM usage_billing_facts AS \"usage\""));
|
||||
assert!(super::SUMMARIZE_PROVIDER_API_KEY_WINDOW_USAGE_SQL
|
||||
assert!(!super::APPLY_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_DELTA_SQL
|
||||
.contains("usage_billing_facts AS \"usage\""));
|
||||
}
|
||||
|
||||
@@ -498,7 +511,10 @@ fn usage_billing_facts_projects_upstream_stream_mode() {
|
||||
assert!(migration.contains("COALESCE(usage_rows.upstream_is_stream"));
|
||||
assert!(migration.contains("COALESCE(usage_rows.is_stream, FALSE)"));
|
||||
assert!(migration.contains("ADD COLUMN IF NOT EXISTS upstream_is_stream boolean"));
|
||||
assert!(migration.contains("request_metadata->>'upstream_is_stream'"));
|
||||
assert!(
|
||||
!migration.contains("request_metadata->>'upstream_is_stream'"),
|
||||
"migration should avoid backfilling historical usage rows from request metadata"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
Reference in New Issue
Block a user