boos/BoosAPI
2
0
Fork 2
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: sync upstream analytics schema support

Browse Source
This commit is contained in:
boos
2026-06-06 16:07:13 +08:00
parent 0d231ec467
commit 339f91d014
7 changed files with 1233 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/services/billing/mod.rs
+468 -1
View File
@@ -2,8 +2,9 @@ use crate::{
error::{AppError, AppResult},
services::shared::{enforce_key_request_limit, increment_key_usage},
};
use chrono::Timelike;
use serde::{Deserialize, Serialize};
use sqlx::{FromRow, SqlitePool};
use sqlx::{FromRow, Row, SqlitePool};
use std::sync::Arc;
#[derive(Clone)]
@@ -152,6 +153,215 @@ impl UsageTracker {
Ok(())
}
pub async fn project_proxy_usage_for_db(db: &SqlitePool) -> AppResult<i64> {
Self::ensure_schema(db).await?;
let last_id: i64 = sqlx::query_scalar(
"SELECT COALESCE(last_proxy_log_id, 0)
FROM analytics_projection_checkpoints
WHERE projector_key = 'usage-aggregates-v1'",
)
.fetch_optional(db)
.await?
.unwrap_or(0);
let rows = sqlx::query(
r#"
SELECT
l.id,
l.created_at,
l.status,
l.latency_ms,
l.total_tokens,
l.estimated_cost,
COALESCE(NULLIF(l.model_actual, ''), NULLIF(l.model_requested, ''), 'unknown') AS model,
a.site_id,
s.platform
FROM proxy_logs l
LEFT JOIN accounts a ON a.id = l.account_id
LEFT JOIN sites s ON s.id = a.site_id
WHERE l.id > ?
ORDER BY l.id ASC
LIMIT 5000
"#,
)
.bind(last_id)
.fetch_all(db)
.await?;
if rows.is_empty() {
touch_projection_checkpoint(db, last_id, None).await?;
return Ok(0);
}
let mut max_id = last_id;
let mut tx = db.begin().await?;
for row in rows {
let id: i64 = row.try_get("id").unwrap_or(0);
max_id = max_id.max(id);
let site_id: Option<i64> = row.try_get("site_id").ok().flatten();
let Some(site_id) = site_id.filter(|value| *value > 0) else {
continue;
};
let created_at: Option<String> = row.try_get("created_at").ok().flatten();
let Some(created_at) = created_at else {
continue;
};
let local_day = sqlite_day_key(&created_at);
let bucket_start_utc = sqlite_hour_key(&created_at);
let status: Option<String> = row.try_get("status").ok().flatten();
let is_success = status
.as_deref()
.unwrap_or_default()
.eq_ignore_ascii_case("success");
let latency_ms = non_negative_i64(row.try_get("latency_ms").ok().flatten());
let latency_count = i64::from(latency_ms > 0);
let total_tokens = non_negative_i64(row.try_get("total_tokens").ok().flatten());
let estimated_cost = non_negative_f64(row.try_get("estimated_cost").ok().flatten());
let platform: Option<String> = row.try_get("platform").ok().flatten();
let model: String = row
.try_get::<Option<String>, _>("model")
.ok()
.flatten()
.filter(|value| !value.trim().is_empty())
.unwrap_or_else(|| "unknown".to_string());
let summary_spend = if estimated_cost > 0.0 {
estimated_cost
} else if platform
.as_deref()
.unwrap_or_default()
.trim()
.eq_ignore_ascii_case("veloera")
{
total_tokens as f64 / 1_000_000.0
} else {
total_tokens as f64 / 500_000.0
};
let site_spend = if estimated_cost > 0.0 {
estimated_cost
} else {
fallback_token_cost(total_tokens, platform.as_deref())
};
let model_spend = if estimated_cost > 0.0 {
estimated_cost
} else {
total_tokens as f64 / 500_000.0
};
sqlx::query(
r#"
INSERT INTO site_day_usage
(local_day, site_id, total_calls, success_calls, failed_calls,
total_tokens, total_summary_spend, total_site_spend,
total_latency_ms, latency_count, updated_at)
VALUES (?, ?, 1, ?, ?, ?, ?, ?, ?, ?, datetime('now'))
ON CONFLICT(local_day, site_id) DO UPDATE SET
total_calls = total_calls + 1,
success_calls = success_calls + excluded.success_calls,
failed_calls = failed_calls + excluded.failed_calls,
total_tokens = total_tokens + excluded.total_tokens,
total_summary_spend = total_summary_spend + excluded.total_summary_spend,
total_site_spend = total_site_spend + excluded.total_site_spend,
total_latency_ms = total_latency_ms + excluded.total_latency_ms,
latency_count = latency_count + excluded.latency_count,
updated_at = datetime('now')
"#,
)
.bind(&local_day)
.bind(site_id)
.bind(i64::from(is_success))
.bind(i64::from(!is_success))
.bind(total_tokens)
.bind(summary_spend)
.bind(site_spend)
.bind(latency_ms)
.bind(latency_count)
.execute(&mut *tx)
.await?;
sqlx::query(
r#"
INSERT INTO site_hour_usage
(bucket_start_utc, site_id, total_calls, success_calls, failed_calls,
total_tokens, total_summary_spend, total_site_spend,
total_latency_ms, latency_count, updated_at)
VALUES (?, ?, 1, ?, ?, ?, ?, ?, ?, ?, datetime('now'))
ON CONFLICT(bucket_start_utc, site_id) DO UPDATE SET
total_calls = total_calls + 1,
success_calls = success_calls + excluded.success_calls,
failed_calls = failed_calls + excluded.failed_calls,
total_tokens = total_tokens + excluded.total_tokens,
total_summary_spend = total_summary_spend + excluded.total_summary_spend,
total_site_spend = total_site_spend + excluded.total_site_spend,
total_latency_ms = total_latency_ms + excluded.total_latency_ms,
latency_count = latency_count + excluded.latency_count,
updated_at = datetime('now')
"#,
)
.bind(&bucket_start_utc)
.bind(site_id)
.bind(i64::from(is_success))
.bind(i64::from(!is_success))
.bind(total_tokens)
.bind(summary_spend)
.bind(site_spend)
.bind(latency_ms)
.bind(latency_count)
.execute(&mut *tx)
.await?;
sqlx::query(
r#"
INSERT INTO model_day_usage
(local_day, site_id, model, total_calls, success_calls, failed_calls,
total_tokens, total_spend, total_latency_ms, latency_count, updated_at)
VALUES (?, ?, ?, 1, ?, ?, ?, ?, ?, ?, datetime('now'))
ON CONFLICT(local_day, site_id, model) DO UPDATE SET
total_calls = total_calls + 1,
success_calls = success_calls + excluded.success_calls,
failed_calls = failed_calls + excluded.failed_calls,
total_tokens = total_tokens + excluded.total_tokens,
total_spend = total_spend + excluded.total_spend,
total_latency_ms = total_latency_ms + excluded.total_latency_ms,
latency_count = latency_count + excluded.latency_count,
updated_at = datetime('now')
"#,
)
.bind(&local_day)
.bind(site_id)
.bind(&model)
.bind(i64::from(is_success))
.bind(i64::from(!is_success))
.bind(total_tokens)
.bind(model_spend)
.bind(latency_ms)
.bind(latency_count)
.execute(&mut *tx)
.await?;
}
sqlx::query(
r#"
INSERT INTO analytics_projection_checkpoints
(projector_key, time_zone, last_proxy_log_id, watermark_created_at,
last_projected_at, last_successful_at, last_error, updated_at)
VALUES ('usage-aggregates-v1', 'Local', ?, datetime('now'), datetime('now'), datetime('now'), NULL, datetime('now'))
ON CONFLICT(projector_key) DO UPDATE SET
last_proxy_log_id = excluded.last_proxy_log_id,
watermark_created_at = excluded.watermark_created_at,
last_projected_at = excluded.last_projected_at,
last_successful_at = excluded.last_successful_at,
last_error = NULL,
updated_at = datetime('now')
"#,
)
.bind(max_id)
.execute(&mut *tx)
.await?;
tx.commit().await?;
Ok(max_id - last_id)
}
async fn ensure_schema(db: &SqlitePool) -> AppResult<()> {
sqlx::query(
"CREATE TABLE IF NOT EXISTS usage_logs (
@@ -190,10 +400,179 @@ impl UsageTracker {
)
.execute(db)
.await?;
sqlx::query(
"CREATE TABLE IF NOT EXISTS analytics_projection_checkpoints (
projector_key TEXT PRIMARY KEY NOT NULL,
time_zone TEXT NOT NULL DEFAULT 'Local',
last_proxy_log_id INTEGER NOT NULL DEFAULT 0,
watermark_created_at TEXT,
recompute_from_id INTEGER,
recompute_requested_at TEXT,
recompute_reason TEXT,
recompute_started_at TEXT,
recompute_completed_at TEXT,
last_projected_at TEXT,
last_successful_at TEXT,
last_error TEXT,
lease_owner TEXT,
lease_token TEXT,
lease_expires_at TEXT,
created_at TEXT DEFAULT (datetime('now')),
updated_at TEXT DEFAULT (datetime('now'))
)",
)
.execute(db)
.await?;
sqlx::query(
"CREATE TABLE IF NOT EXISTS model_day_usage (
id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
local_day TEXT NOT NULL,
site_id INTEGER NOT NULL,
model TEXT NOT NULL,
total_calls INTEGER NOT NULL DEFAULT 0,
success_calls INTEGER NOT NULL DEFAULT 0,
failed_calls INTEGER NOT NULL DEFAULT 0,
total_tokens INTEGER NOT NULL DEFAULT 0,
total_spend REAL NOT NULL DEFAULT 0,
total_latency_ms INTEGER NOT NULL DEFAULT 0,
latency_count INTEGER NOT NULL DEFAULT 0,
created_at TEXT DEFAULT (datetime('now')),
updated_at TEXT DEFAULT (datetime('now')),
UNIQUE(local_day, site_id, model)
)",
)
.execute(db)
.await?;
sqlx::query(
"CREATE TABLE IF NOT EXISTS site_day_usage (
id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
local_day TEXT NOT NULL,
site_id INTEGER NOT NULL,
total_calls INTEGER NOT NULL DEFAULT 0,
success_calls INTEGER NOT NULL DEFAULT 0,
failed_calls INTEGER NOT NULL DEFAULT 0,
total_tokens INTEGER NOT NULL DEFAULT 0,
total_summary_spend REAL NOT NULL DEFAULT 0,
total_site_spend REAL NOT NULL DEFAULT 0,
total_latency_ms INTEGER NOT NULL DEFAULT 0,
latency_count INTEGER NOT NULL DEFAULT 0,
created_at TEXT DEFAULT (datetime('now')),
updated_at TEXT DEFAULT (datetime('now')),
UNIQUE(local_day, site_id)
)",
)
.execute(db)
.await?;
sqlx::query(
"CREATE TABLE IF NOT EXISTS site_hour_usage (
id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
bucket_start_utc TEXT NOT NULL,
site_id INTEGER NOT NULL,
total_calls INTEGER NOT NULL DEFAULT 0,
success_calls INTEGER NOT NULL DEFAULT 0,
failed_calls INTEGER NOT NULL DEFAULT 0,
total_tokens INTEGER NOT NULL DEFAULT 0,
total_summary_spend REAL NOT NULL DEFAULT 0,
total_site_spend REAL NOT NULL DEFAULT 0,
total_latency_ms INTEGER NOT NULL DEFAULT 0,
latency_count INTEGER NOT NULL DEFAULT 0,
created_at TEXT DEFAULT (datetime('now')),
updated_at TEXT DEFAULT (datetime('now')),
UNIQUE(bucket_start_utc, site_id)
)",
)
.execute(db)
.await?;
Ok(())
}
}
async fn touch_projection_checkpoint(
db: &SqlitePool,
last_proxy_log_id: i64,
error: Option<&str>,
) -> AppResult<()> {
sqlx::query(
r#"
INSERT INTO analytics_projection_checkpoints
(projector_key, time_zone, last_proxy_log_id, last_projected_at,
last_successful_at, last_error, updated_at)
VALUES ('usage-aggregates-v1', 'Local', ?, datetime('now'), datetime('now'), ?, datetime('now'))
ON CONFLICT(projector_key) DO UPDATE SET
last_projected_at = datetime('now'),
last_successful_at = CASE WHEN ? IS NULL THEN datetime('now') ELSE last_successful_at END,
last_error = ?,
updated_at = datetime('now')
"#,
)
.bind(last_proxy_log_id)
.bind(error)
.bind(error)
.bind(error)
.execute(db)
.await?;
Ok(())
}
fn non_negative_i64(value: Option<i64>) -> i64 {
value.unwrap_or(0).max(0)
}
fn non_negative_f64(value: Option<f64>) -> f64 {
value.unwrap_or(0.0).max(0.0)
}
fn sqlite_day_key(created_at: &str) -> String {
parse_sqlite_datetime(created_at)
.map(|dt| dt.format("%Y-%m-%d").to_string())
.unwrap_or_else(|| created_at.chars().take(10).collect())
}
fn sqlite_hour_key(created_at: &str) -> String {
parse_sqlite_datetime(created_at)
.map(|dt| {
dt.with_minute(0)
.and_then(|dt| dt.with_second(0))
.and_then(|dt| dt.with_nanosecond(0))
.unwrap_or(dt)
.format("%Y-%m-%d %H:%M:%S")
.to_string()
})
.unwrap_or_else(|| {
let prefix: String = created_at.chars().take(13).collect();
if prefix.len() == 13 {
format!("{prefix}:00:00")
} else {
created_at.to_string()
}
})
}
fn parse_sqlite_datetime(value: &str) -> Option<chrono::NaiveDateTime> {
chrono::NaiveDateTime::parse_from_str(value, "%Y-%m-%d %H:%M:%S")
.ok()
.or_else(|| {
chrono::DateTime::parse_from_rfc3339(value)
.ok()
.map(|dt| dt.naive_utc())
})
}
fn fallback_token_cost(tokens: i64, platform: Option<&str>) -> f64 {
if tokens <= 0 {
return 0.0;
}
match platform
.unwrap_or_default()
.trim()
.to_ascii_lowercase()
.as_str()
{
"veloera" => tokens as f64 / 1_000_000.0,
_ => tokens as f64 / 500_000.0,
}
}
#[cfg(test)]
mod tests {
use super::*;
@@ -282,4 +661,92 @@ mod tests {
.unwrap();
assert_eq!(count, 1);
}
#[tokio::test]
async fn projects_proxy_logs_to_site_and_model_usage() {
let pool = pool().await;
sqlx::query(
"CREATE TABLE sites (
id INTEGER PRIMARY KEY,
name TEXT NOT NULL,
platform TEXT
)",
)
.execute(&pool)
.await
.unwrap();
sqlx::query(
"CREATE TABLE accounts (
id INTEGER PRIMARY KEY,
site_id INTEGER NOT NULL
)",
)
.execute(&pool)
.await
.unwrap();
sqlx::query(
"CREATE TABLE proxy_logs (
id INTEGER PRIMARY KEY,
account_id INTEGER,
model_requested TEXT,
model_actual TEXT,
status TEXT,
latency_ms INTEGER,
total_tokens INTEGER,
estimated_cost REAL,
created_at TEXT
)",
)
.execute(&pool)
.await
.unwrap();
sqlx::query("INSERT INTO sites (id, name, platform) VALUES (10, 'Site', 'new-api')")
.execute(&pool)
.await
.unwrap();
sqlx::query("INSERT INTO accounts (id, site_id) VALUES (20, 10)")
.execute(&pool)
.await
.unwrap();
sqlx::query(
"INSERT INTO proxy_logs
(id, account_id, model_requested, model_actual, status, latency_ms, total_tokens, estimated_cost, created_at)
VALUES
(1, 20, 'gpt-test', NULL, 'success', 120, 1000, 0.12, '2026-06-06 08:12:30')",
)
.execute(&pool)
.await
.unwrap();
let projected = UsageTracker::project_proxy_usage_for_db(&pool).await.unwrap();
assert_eq!(projected, 1);
let model_usage: (i64, i64, f64, i64) = sqlx::query_as(
"SELECT total_calls, total_tokens, total_spend, total_latency_ms
FROM model_day_usage
WHERE site_id = 10 AND model = 'gpt-test'",
)
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(model_usage, (1, 1000, 0.12, 120));
let site_usage: (i64, i64, f64) = sqlx::query_as(
"SELECT total_calls, success_calls, total_site_spend
FROM site_day_usage
WHERE site_id = 10 AND local_day = '2026-06-06'",
)
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(site_usage, (1, 1, 0.12));
let hour_bucket: String = sqlx::query_scalar(
"SELECT bucket_start_utc FROM site_hour_usage WHERE site_id = 10",
)
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(hour_bucket, "2026-06-06 08:00:00");
}
}
src/services/cron/mod.rs
+4
View File
@@ -79,6 +79,10 @@ impl CronScheduler {
async fn aggregate_usage(db: Arc<SqlitePool>) -> AppResult<()> {
UsageTracker::create_daily_usage_snapshot_for_db(&db).await?;
let projected = UsageTracker::project_proxy_usage_for_db(&db).await?;
if projected > 0 {
tracing::debug!(projected, "usage aggregates projected");
}
sqlx::query(
"UPDATE downstream_api_keys
SET updated_at = datetime('now')