forked from boos/BoosAPI
feat: sync upstream analytics schema support
This commit is contained in:
boos
+468
-1
@@ -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");
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user