mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-07 18:07:47 +08:00
feat(data): complete portable SQL backend parity
Align MySQL and SQLite schemas, migrations, usage, stats, export, and backfill behavior with the shared data contracts. Extend gateway startup and maintenance support across all SQL drivers.
This commit is contained in:
@@ -642,23 +642,8 @@ LIMIT 1
|
||||
input: &PaymentGatewayConfigWriteInput,
|
||||
) -> Result<AdminBillingMutationOutcome<PaymentGatewayConfigRecord>, DataLayerError> {
|
||||
let provider = input.provider.trim().to_ascii_lowercase();
|
||||
let existing_secret = if input.preserve_existing_secret {
|
||||
sqlx::query_scalar::<_, String>(
|
||||
"SELECT merchant_key_encrypted FROM payment_gateway_configs WHERE provider = ?",
|
||||
)
|
||||
.bind(&provider)
|
||||
.fetch_optional(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let secret = if input.preserve_existing_secret {
|
||||
existing_secret
|
||||
} else {
|
||||
input.merchant_key_encrypted.clone()
|
||||
};
|
||||
let now = current_unix_secs_i64();
|
||||
let mut tx = self.pool.begin().await.map_sql_err()?;
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO payment_gateway_configs (
|
||||
@@ -672,7 +657,10 @@ ON DUPLICATE KEY UPDATE
|
||||
endpoint_url = VALUES(endpoint_url),
|
||||
callback_base_url = VALUES(callback_base_url),
|
||||
merchant_id = VALUES(merchant_id),
|
||||
merchant_key_encrypted = VALUES(merchant_key_encrypted),
|
||||
merchant_key_encrypted = CASE
|
||||
WHEN ? THEN merchant_key_encrypted
|
||||
ELSE VALUES(merchant_key_encrypted)
|
||||
END,
|
||||
pay_currency = VALUES(pay_currency),
|
||||
usd_exchange_rate = VALUES(usd_exchange_rate),
|
||||
min_recharge_usd = VALUES(min_recharge_usd),
|
||||
@@ -685,22 +673,35 @@ ON DUPLICATE KEY UPDATE
|
||||
.bind(&input.endpoint_url)
|
||||
.bind(input.callback_base_url.as_deref())
|
||||
.bind(&input.merchant_id)
|
||||
.bind(secret.as_deref())
|
||||
.bind(input.merchant_key_encrypted.as_deref())
|
||||
.bind(&input.pay_currency)
|
||||
.bind(input.usd_exchange_rate)
|
||||
.bind(input.min_recharge_usd)
|
||||
.bind(json_to_string(&input.channels_json)?)
|
||||
.bind(now)
|
||||
.bind(now)
|
||||
.execute(&self.pool)
|
||||
.bind(input.preserve_existing_secret)
|
||||
.execute(&mut *tx)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
match self.find_payment_gateway_config(&provider).await? {
|
||||
Some(record) => Ok(AdminBillingMutationOutcome::Applied(record)),
|
||||
None => Err(DataLayerError::UnexpectedValue(
|
||||
"upserted payment gateway config missing".to_string(),
|
||||
)),
|
||||
}
|
||||
let row = sqlx::query(
|
||||
r#"
|
||||
SELECT
|
||||
provider, enabled, endpoint_url, callback_base_url, merchant_id,
|
||||
merchant_key_encrypted, pay_currency, usd_exchange_rate, min_recharge_usd,
|
||||
channels_json, created_at AS created_at_unix_secs, updated_at AS updated_at_unix_secs
|
||||
FROM payment_gateway_configs
|
||||
WHERE provider = ?
|
||||
LIMIT 1
|
||||
"#,
|
||||
)
|
||||
.bind(&provider)
|
||||
.fetch_one(&mut *tx)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
let record = map_payment_gateway_config_mysql(&row)?;
|
||||
tx.commit().await.map_sql_err()?;
|
||||
Ok(AdminBillingMutationOutcome::Applied(record))
|
||||
}
|
||||
|
||||
async fn list_billing_plans(
|
||||
@@ -1402,7 +1403,13 @@ fn parse_required_json(raw: String) -> Result<serde_json::Value, DataLayerError>
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use serde_json::json;
|
||||
|
||||
use super::MysqlBillingReadRepository;
|
||||
use crate::run_migrations;
|
||||
use aether_data_contracts::repository::billing::{
|
||||
AdminBillingMutationOutcome, BillingReadRepository, PaymentGatewayConfigWriteInput,
|
||||
};
|
||||
|
||||
#[tokio::test]
|
||||
async fn repository_builds_from_lazy_pool() {
|
||||
@@ -1414,4 +1421,88 @@ mod tests {
|
||||
|
||||
let _repository = MysqlBillingReadRepository::new(pool);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn mysql_gateway_upsert_preserves_secret_when_url_is_set() {
|
||||
let Some(database_url) = std::env::var("AETHER_TEST_MYSQL_URL")
|
||||
.ok()
|
||||
.filter(|value| !value.trim().is_empty())
|
||||
else {
|
||||
eprintln!("skipping mysql billing parity test because AETHER_TEST_MYSQL_URL is unset");
|
||||
return;
|
||||
};
|
||||
let pool = sqlx::mysql::MySqlPoolOptions::new()
|
||||
.max_connections(1)
|
||||
.connect(&database_url)
|
||||
.await
|
||||
.expect("mysql pool should connect");
|
||||
run_migrations(&pool)
|
||||
.await
|
||||
.expect("mysql migrations should run");
|
||||
sqlx::query("DELETE FROM payment_gateway_configs WHERE provider = 'paritypay'")
|
||||
.execute(&pool)
|
||||
.await
|
||||
.expect("gateway cleanup should succeed");
|
||||
let repository = MysqlBillingReadRepository::new(pool.clone());
|
||||
|
||||
let mut input = PaymentGatewayConfigWriteInput {
|
||||
provider: " ParityPay ".to_string(),
|
||||
enabled: true,
|
||||
endpoint_url: "https://pay.example/first".to_string(),
|
||||
callback_base_url: Some("https://api.example/callback".to_string()),
|
||||
merchant_id: "merchant-1".to_string(),
|
||||
merchant_key_encrypted: Some("secret-first".to_string()),
|
||||
preserve_existing_secret: true,
|
||||
pay_currency: "USD".to_string(),
|
||||
usd_exchange_rate: 1.0,
|
||||
min_recharge_usd: 2.0,
|
||||
channels_json: json!([{"id": "card"}]),
|
||||
};
|
||||
let inserted = match repository
|
||||
.upsert_payment_gateway_config(&input)
|
||||
.await
|
||||
.expect("gateway insert should run")
|
||||
{
|
||||
AdminBillingMutationOutcome::Applied(record) => record,
|
||||
other => panic!("unexpected gateway insert outcome: {other:?}"),
|
||||
};
|
||||
assert_eq!(
|
||||
inserted.merchant_key_encrypted.as_deref(),
|
||||
Some("secret-first")
|
||||
);
|
||||
|
||||
input.merchant_key_encrypted = Some("secret-ignored".to_string());
|
||||
let preserved = match repository
|
||||
.upsert_payment_gateway_config(&input)
|
||||
.await
|
||||
.expect("gateway preserve update should run")
|
||||
{
|
||||
AdminBillingMutationOutcome::Applied(record) => record,
|
||||
other => panic!("unexpected gateway preserve outcome: {other:?}"),
|
||||
};
|
||||
assert_eq!(
|
||||
preserved.merchant_key_encrypted.as_deref(),
|
||||
Some("secret-first")
|
||||
);
|
||||
|
||||
input.preserve_existing_secret = false;
|
||||
input.merchant_key_encrypted = Some("secret-replaced".to_string());
|
||||
let replaced = match repository
|
||||
.upsert_payment_gateway_config(&input)
|
||||
.await
|
||||
.expect("gateway replacement update should run")
|
||||
{
|
||||
AdminBillingMutationOutcome::Applied(record) => record,
|
||||
other => panic!("unexpected gateway replace outcome: {other:?}"),
|
||||
};
|
||||
assert_eq!(
|
||||
replaced.merchant_key_encrypted.as_deref(),
|
||||
Some("secret-replaced")
|
||||
);
|
||||
|
||||
sqlx::query("DELETE FROM payment_gateway_configs WHERE provider = 'paritypay'")
|
||||
.execute(&pool)
|
||||
.await
|
||||
.expect("gateway cleanup should succeed");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
use std::collections::{BTreeMap, BTreeSet};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use sqlx::{mysql::MySqlRow, MySql, QueryBuilder, Row};
|
||||
use sqlx::{mysql::MySqlRow, MySql, MySqlConnection, QueryBuilder, Row};
|
||||
|
||||
use aether_data_contracts::repository::candidates::{
|
||||
request_candidate_lifecycle_would_regress, PublicHealthStatusCount, PublicHealthTimelineBucket,
|
||||
@@ -51,24 +51,6 @@ impl MysqlRequestCandidateRepository {
|
||||
pub fn new(pool: MysqlPool) -> Self {
|
||||
Self { pool }
|
||||
}
|
||||
|
||||
async fn find_by_unique(
|
||||
&self,
|
||||
request_id: &str,
|
||||
candidate_index: u32,
|
||||
retry_index: u32,
|
||||
) -> Result<Option<StoredRequestCandidate>, DataLayerError> {
|
||||
let row = sqlx::query(&format!(
|
||||
"{CANDIDATE_COLUMNS} WHERE request_id = ? AND candidate_index = ? AND retry_index = ? LIMIT 1"
|
||||
))
|
||||
.bind(request_id)
|
||||
.bind(to_i32(candidate_index)?)
|
||||
.bind(to_i32(retry_index)?)
|
||||
.fetch_optional(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
row.as_ref().map(map_candidate_row).transpose()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
@@ -237,16 +219,50 @@ impl RequestCandidateWriteRepository for MysqlRequestCandidateRepository {
|
||||
candidate: UpsertRequestCandidateRecord,
|
||||
) -> Result<StoredRequestCandidate, DataLayerError> {
|
||||
candidate.validate()?;
|
||||
let existing = self
|
||||
.find_by_unique(
|
||||
&candidate.request_id,
|
||||
candidate.candidate_index,
|
||||
candidate.retry_index,
|
||||
)
|
||||
.await?;
|
||||
let merged = merge_candidate(candidate, existing)?;
|
||||
upsert_merged_candidate(&self.pool, &merged).await?;
|
||||
Ok(merged)
|
||||
let mut tx = self.pool.begin().await.map_sql_err()?;
|
||||
match upsert_candidate_in_transaction(&mut tx, candidate).await {
|
||||
Ok(candidate) => {
|
||||
tx.commit().await.map_sql_err()?;
|
||||
Ok(candidate)
|
||||
}
|
||||
Err(err) => {
|
||||
tx.rollback().await.map_sql_err()?;
|
||||
Err(err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn upsert_many(
|
||||
&self,
|
||||
candidates: Vec<UpsertRequestCandidateRecord>,
|
||||
) -> Result<usize, DataLayerError> {
|
||||
if candidates.is_empty() {
|
||||
return Ok(0);
|
||||
}
|
||||
for candidate in &candidates {
|
||||
candidate.validate()?;
|
||||
}
|
||||
|
||||
let mut tx = self.pool.begin().await.map_sql_err()?;
|
||||
let result: Result<usize, DataLayerError> = async {
|
||||
let mut persisted = 0usize;
|
||||
for candidate in candidates {
|
||||
upsert_candidate_in_transaction(&mut tx, candidate).await?;
|
||||
persisted = persisted.saturating_add(1);
|
||||
}
|
||||
Ok(persisted)
|
||||
}
|
||||
.await;
|
||||
match result {
|
||||
Ok(persisted) => {
|
||||
tx.commit().await.map_sql_err()?;
|
||||
Ok(persisted)
|
||||
}
|
||||
Err(err) => {
|
||||
tx.rollback().await.map_sql_err()?;
|
||||
Err(err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn delete_created_before(
|
||||
@@ -282,8 +298,90 @@ WHERE id IN (
|
||||
}
|
||||
}
|
||||
|
||||
async fn upsert_candidate_in_transaction(
|
||||
tx: &mut sqlx::Transaction<'_, MySql>,
|
||||
candidate: UpsertRequestCandidateRecord,
|
||||
) -> Result<StoredRequestCandidate, DataLayerError> {
|
||||
// Write first so both existing rows and previously empty unique keys are locked
|
||||
// before the Rust merge reads their latest committed state.
|
||||
let insert_candidate = merge_candidate(candidate.clone(), None)?;
|
||||
insert_candidate_if_absent(tx, &insert_candidate).await?;
|
||||
let existing = find_by_unique_for_update(
|
||||
tx,
|
||||
&candidate.request_id,
|
||||
candidate.candidate_index,
|
||||
candidate.retry_index,
|
||||
)
|
||||
.await?
|
||||
.ok_or_else(|| {
|
||||
DataLayerError::UnexpectedValue(
|
||||
"request candidate row was not locked after insert-if-absent".to_string(),
|
||||
)
|
||||
})?;
|
||||
let merged = merge_candidate(candidate, Some(existing))?;
|
||||
upsert_merged_candidate(tx, &merged).await?;
|
||||
find_by_unique_for_update(
|
||||
tx,
|
||||
&merged.request_id,
|
||||
merged.candidate_index,
|
||||
merged.retry_index,
|
||||
)
|
||||
.await?
|
||||
.ok_or_else(|| {
|
||||
DataLayerError::UnexpectedValue(
|
||||
"request candidate row disappeared after atomic upsert".to_string(),
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
async fn insert_candidate_if_absent(
|
||||
connection: &mut MySqlConnection,
|
||||
candidate: &StoredRequestCandidate,
|
||||
) -> Result<(), DataLayerError> {
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO request_candidates (
|
||||
id, request_id, candidate_index, retry_index, status, created_at
|
||||
)
|
||||
VALUES (?, ?, ?, ?, ?, ?)
|
||||
ON DUPLICATE KEY UPDATE id = id
|
||||
"#,
|
||||
)
|
||||
.bind(&candidate.id)
|
||||
.bind(&candidate.request_id)
|
||||
.bind(to_i32(candidate.candidate_index)?)
|
||||
.bind(to_i32(candidate.retry_index)?)
|
||||
.bind(status_to_database(candidate.status))
|
||||
.bind(u64_to_i64(
|
||||
candidate.created_at_unix_ms,
|
||||
"request candidate created_at",
|
||||
)?)
|
||||
.execute(connection)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn find_by_unique_for_update(
|
||||
connection: &mut MySqlConnection,
|
||||
request_id: &str,
|
||||
candidate_index: u32,
|
||||
retry_index: u32,
|
||||
) -> Result<Option<StoredRequestCandidate>, DataLayerError> {
|
||||
let row = sqlx::query(&format!(
|
||||
"{CANDIDATE_COLUMNS} WHERE request_id = ? AND candidate_index = ? AND retry_index = ? LIMIT 1 FOR UPDATE"
|
||||
))
|
||||
.bind(request_id)
|
||||
.bind(to_i32(candidate_index)?)
|
||||
.bind(to_i32(retry_index)?)
|
||||
.fetch_optional(connection)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
row.as_ref().map(map_candidate_row).transpose()
|
||||
}
|
||||
|
||||
async fn upsert_merged_candidate(
|
||||
pool: &MysqlPool,
|
||||
connection: &mut MySqlConnection,
|
||||
candidate: &StoredRequestCandidate,
|
||||
) -> Result<(), DataLayerError> {
|
||||
sqlx::query(
|
||||
@@ -405,7 +503,7 @@ ON DUPLICATE KEY UPDATE
|
||||
candidate.finished_at_unix_ms,
|
||||
"request candidate finished_at",
|
||||
)?)
|
||||
.execute(pool)
|
||||
.execute(connection)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
Ok(())
|
||||
@@ -776,9 +874,10 @@ mod tests {
|
||||
use super::MysqlRequestCandidateRepository;
|
||||
use crate::run_migrations;
|
||||
use aether_data_contracts::repository::candidates::{
|
||||
RequestCandidateReadRepository, RequestCandidateStatus, StoredRequestCandidate,
|
||||
UpsertRequestCandidateRecord,
|
||||
RequestCandidateReadRepository, RequestCandidateStatus, RequestCandidateWriteRepository,
|
||||
StoredRequestCandidate, UpsertRequestCandidateRecord,
|
||||
};
|
||||
use serde_json::json;
|
||||
|
||||
#[tokio::test]
|
||||
async fn repository_builds_from_lazy_pool() {
|
||||
@@ -792,7 +891,7 @@ mod tests {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn mysql_atomic_conflict_keeps_candidate_lifecycle_monotonic_when_configured() {
|
||||
async fn mysql_concurrent_and_batch_upserts_are_atomic_when_configured() {
|
||||
let Some(database_url) = std::env::var("AETHER_TEST_MYSQL_URL")
|
||||
.ok()
|
||||
.filter(|value| !value.trim().is_empty())
|
||||
@@ -803,7 +902,7 @@ mod tests {
|
||||
return;
|
||||
};
|
||||
let pool = sqlx::mysql::MySqlPoolOptions::new()
|
||||
.max_connections(2)
|
||||
.max_connections(12)
|
||||
.connect(&database_url)
|
||||
.await
|
||||
.expect("mysql test pool should connect");
|
||||
@@ -811,103 +910,172 @@ mod tests {
|
||||
.await
|
||||
.expect("mysql migrations should run");
|
||||
let repository = MysqlRequestCandidateRepository::new(pool.clone());
|
||||
let request_id = format!("candidate-lifecycle-{}", uuid::Uuid::new_v4());
|
||||
|
||||
let terminal = stored_candidate(
|
||||
let request_id = format!("candidate-concurrency-{}", uuid::Uuid::new_v4());
|
||||
let mut initial = sample_upsert(
|
||||
&request_id,
|
||||
"terminal",
|
||||
0,
|
||||
RequestCandidateStatus::Success,
|
||||
Some(123),
|
||||
Some(2_000_002),
|
||||
);
|
||||
super::upsert_merged_candidate(&pool, &terminal)
|
||||
.await
|
||||
.expect("terminal candidate should insert");
|
||||
let stale_streaming = stored_candidate(
|
||||
&request_id,
|
||||
"stale-streaming",
|
||||
0,
|
||||
RequestCandidateStatus::Streaming,
|
||||
Some(9_999),
|
||||
Some(9_999_999),
|
||||
);
|
||||
super::upsert_merged_candidate(&pool, &stale_streaming)
|
||||
.await
|
||||
.expect("stale streaming conflict should execute");
|
||||
|
||||
let streaming = stored_candidate(
|
||||
&request_id,
|
||||
"streaming",
|
||||
1,
|
||||
RequestCandidateStatus::Streaming,
|
||||
None,
|
||||
None,
|
||||
);
|
||||
super::upsert_merged_candidate(&pool, &streaming)
|
||||
.await
|
||||
.expect("streaming candidate should insert");
|
||||
let stale_pending = stored_candidate(
|
||||
&request_id,
|
||||
"stale-pending",
|
||||
1,
|
||||
"initial",
|
||||
RequestCandidateStatus::Pending,
|
||||
None,
|
||||
None,
|
||||
Some(json!({"initial": true})),
|
||||
3_000_000,
|
||||
);
|
||||
super::upsert_merged_candidate(&pool, &stale_pending)
|
||||
initial.is_cached = Some(false);
|
||||
repository
|
||||
.upsert(initial)
|
||||
.await
|
||||
.expect("stale pending conflict should execute");
|
||||
.expect("initial candidate should insert");
|
||||
|
||||
let pending = stored_candidate(
|
||||
&request_id,
|
||||
"pending",
|
||||
2,
|
||||
RequestCandidateStatus::Pending,
|
||||
Some(321),
|
||||
None,
|
||||
);
|
||||
super::upsert_merged_candidate(&pool, &pending)
|
||||
.await
|
||||
.expect("pending candidate should insert");
|
||||
let stale_available = stored_candidate(
|
||||
&request_id,
|
||||
"stale-available",
|
||||
2,
|
||||
RequestCandidateStatus::Available,
|
||||
Some(9_999),
|
||||
Some(9_999_999),
|
||||
);
|
||||
super::upsert_merged_candidate(&pool, &stale_available)
|
||||
.await
|
||||
.expect("stale available conflict should execute");
|
||||
const WRITERS: usize = 8;
|
||||
let barrier = std::sync::Arc::new(tokio::sync::Barrier::new(WRITERS));
|
||||
let mut tasks = Vec::new();
|
||||
for writer in 0..WRITERS {
|
||||
let repository = repository.clone();
|
||||
let request_id = request_id.clone();
|
||||
let barrier = barrier.clone();
|
||||
tasks.push(tokio::spawn(async move {
|
||||
let status = if writer == 0 {
|
||||
RequestCandidateStatus::Success
|
||||
} else {
|
||||
RequestCandidateStatus::Streaming
|
||||
};
|
||||
let mut extra_data = serde_json::Map::new();
|
||||
extra_data.insert(format!("writer_{writer}"), json!(writer));
|
||||
let mut candidate = sample_upsert(
|
||||
&request_id,
|
||||
format!("writer-{writer}").as_str(),
|
||||
status,
|
||||
Some(serde_json::Value::Object(extra_data)),
|
||||
3_100_000 + u64::try_from(writer).expect("writer index should fit") * 10,
|
||||
);
|
||||
if writer != 0 {
|
||||
candidate.latency_ms = Some(9_000 + writer as u64);
|
||||
candidate.finished_at_unix_ms = Some(9_000_000 + writer as u64);
|
||||
}
|
||||
barrier.wait().await;
|
||||
repository.upsert(candidate).await
|
||||
}));
|
||||
}
|
||||
for task in tasks {
|
||||
task.await
|
||||
.expect("candidate writer should join")
|
||||
.expect("candidate writer should persist");
|
||||
}
|
||||
|
||||
let candidates = repository
|
||||
.list_by_request_id(&request_id)
|
||||
.await
|
||||
.expect("mysql request candidates should load");
|
||||
let terminal = candidates
|
||||
.iter()
|
||||
.find(|candidate| candidate.candidate_index == 0)
|
||||
.expect("terminal candidate should remain");
|
||||
assert_eq!(terminal.status, RequestCandidateStatus::Success);
|
||||
assert_eq!(terminal.latency_ms, Some(123));
|
||||
assert_eq!(terminal.finished_at_unix_ms, Some(2_000_002));
|
||||
let streaming = candidates
|
||||
.iter()
|
||||
.find(|candidate| candidate.candidate_index == 1)
|
||||
.expect("streaming candidate should remain");
|
||||
assert_eq!(streaming.status, RequestCandidateStatus::Streaming);
|
||||
let pending = candidates
|
||||
.iter()
|
||||
.find(|candidate| candidate.candidate_index == 2)
|
||||
.expect("pending candidate should remain");
|
||||
assert_eq!(pending.status, RequestCandidateStatus::Pending);
|
||||
assert_eq!(pending.latency_ms, Some(321));
|
||||
assert_eq!(pending.finished_at_unix_ms, None);
|
||||
assert_eq!(candidates.len(), 1);
|
||||
let candidate = &candidates[0];
|
||||
assert_eq!(candidate.id, "initial");
|
||||
assert_eq!(candidate.status, RequestCandidateStatus::Success);
|
||||
assert_eq!(candidate.latency_ms, Some(123));
|
||||
assert_eq!(candidate.finished_at_unix_ms, Some(3_100_002));
|
||||
let extra_data = candidate
|
||||
.extra_data
|
||||
.as_ref()
|
||||
.and_then(serde_json::Value::as_object)
|
||||
.expect("merged extra data should be an object");
|
||||
assert_eq!(extra_data.get("initial"), Some(&json!(true)));
|
||||
for writer in 0..WRITERS {
|
||||
assert_eq!(
|
||||
extra_data.get(format!("writer_{writer}").as_str()),
|
||||
Some(&json!(writer))
|
||||
);
|
||||
}
|
||||
|
||||
sqlx::query("DELETE FROM request_candidates WHERE request_id = ?")
|
||||
let batch_request_id = format!("candidate-batch-{}", uuid::Uuid::new_v4());
|
||||
let mut pending = sample_upsert(
|
||||
&batch_request_id,
|
||||
"batch-first",
|
||||
RequestCandidateStatus::Pending,
|
||||
Some(json!({"pending": true})),
|
||||
4_000_000,
|
||||
);
|
||||
pending.is_cached = Some(false);
|
||||
let mut streaming = sample_upsert(
|
||||
&batch_request_id,
|
||||
"batch-second",
|
||||
RequestCandidateStatus::Streaming,
|
||||
Some(json!({"streaming": true})),
|
||||
4_000_100,
|
||||
);
|
||||
streaming.is_cached = None;
|
||||
let mut success = sample_upsert(
|
||||
&batch_request_id,
|
||||
"batch-third",
|
||||
RequestCandidateStatus::Success,
|
||||
Some(json!({"success": true})),
|
||||
4_000_200,
|
||||
);
|
||||
success.is_cached = Some(true);
|
||||
let mut late_pending = sample_upsert(
|
||||
&batch_request_id,
|
||||
"batch-fourth",
|
||||
RequestCandidateStatus::Pending,
|
||||
Some(json!({"late": true})),
|
||||
4_000_300,
|
||||
);
|
||||
late_pending.is_cached = None;
|
||||
late_pending.latency_ms = Some(9_999);
|
||||
late_pending.finished_at_unix_ms = Some(9_999_999);
|
||||
assert_eq!(
|
||||
repository
|
||||
.upsert_many(vec![pending, streaming, success, late_pending])
|
||||
.await
|
||||
.expect("ordered batch should persist"),
|
||||
4
|
||||
);
|
||||
let batch_candidates = repository
|
||||
.list_by_request_id(&batch_request_id)
|
||||
.await
|
||||
.expect("batch candidate should load");
|
||||
assert_eq!(batch_candidates.len(), 1);
|
||||
assert_eq!(batch_candidates[0].id, "batch-first");
|
||||
assert_eq!(batch_candidates[0].status, RequestCandidateStatus::Success);
|
||||
assert!(batch_candidates[0].is_cached);
|
||||
assert_eq!(batch_candidates[0].latency_ms, Some(123));
|
||||
assert_eq!(batch_candidates[0].finished_at_unix_ms, Some(4_000_202));
|
||||
assert_eq!(
|
||||
batch_candidates[0].extra_data,
|
||||
Some(json!({
|
||||
"pending": true,
|
||||
"streaming": true,
|
||||
"success": true,
|
||||
"late": true
|
||||
}))
|
||||
);
|
||||
|
||||
let rollback_request_id = format!("candidate-rollback-{}", uuid::Uuid::new_v4());
|
||||
let valid = sample_upsert(
|
||||
&rollback_request_id,
|
||||
"rollback-valid",
|
||||
RequestCandidateStatus::Pending,
|
||||
None,
|
||||
5_000_000,
|
||||
);
|
||||
let mut invalid = sample_upsert(
|
||||
&rollback_request_id,
|
||||
"rollback-invalid",
|
||||
RequestCandidateStatus::Success,
|
||||
None,
|
||||
5_000_100,
|
||||
);
|
||||
invalid.candidate_index = 1;
|
||||
invalid.latency_ms = Some(u64::MAX);
|
||||
repository
|
||||
.upsert_many(vec![valid, invalid])
|
||||
.await
|
||||
.expect_err("invalid later row should roll back the batch");
|
||||
assert!(repository
|
||||
.list_by_request_id(&rollback_request_id)
|
||||
.await
|
||||
.expect("rolled-back batch should be readable")
|
||||
.is_empty());
|
||||
|
||||
sqlx::query("DELETE FROM request_candidates WHERE request_id IN (?, ?, ?)")
|
||||
.bind(&request_id)
|
||||
.bind(&batch_request_id)
|
||||
.bind(&rollback_request_id)
|
||||
.execute(&pool)
|
||||
.await
|
||||
.expect("mysql candidate test rows should clean up");
|
||||
@@ -984,40 +1152,38 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
fn stored_candidate(
|
||||
fn sample_upsert(
|
||||
request_id: &str,
|
||||
id: &str,
|
||||
candidate_index: u32,
|
||||
status: RequestCandidateStatus,
|
||||
latency_ms: Option<i32>,
|
||||
finished_at_unix_ms: Option<i64>,
|
||||
) -> StoredRequestCandidate {
|
||||
StoredRequestCandidate::new(
|
||||
id.to_string(),
|
||||
request_id.to_string(),
|
||||
Some("user-1".to_string()),
|
||||
Some("key-1".to_string()),
|
||||
None,
|
||||
None,
|
||||
i32::try_from(candidate_index).expect("candidate index should fit"),
|
||||
0,
|
||||
Some("provider-1".to_string()),
|
||||
Some("endpoint-1".to_string()),
|
||||
Some("provider-key-1".to_string()),
|
||||
extra_data: Option<serde_json::Value>,
|
||||
created_at_unix_ms: u64,
|
||||
) -> UpsertRequestCandidateRecord {
|
||||
UpsertRequestCandidateRecord {
|
||||
id: id.to_string(),
|
||||
request_id: request_id.to_string(),
|
||||
user_id: Some("user-1".to_string()),
|
||||
api_key_id: Some("key-1".to_string()),
|
||||
username: Some("user".to_string()),
|
||||
api_key_name: Some("Key".to_string()),
|
||||
candidate_index: 0,
|
||||
retry_index: 0,
|
||||
provider_id: Some("provider-1".to_string()),
|
||||
endpoint_id: Some("endpoint-1".to_string()),
|
||||
key_id: Some("provider-key-1".to_string()),
|
||||
status,
|
||||
None,
|
||||
false,
|
||||
Some(200),
|
||||
None,
|
||||
None,
|
||||
latency_ms,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
2_000_000,
|
||||
Some(2_000_001),
|
||||
finished_at_unix_ms,
|
||||
)
|
||||
.expect("stored candidate should build")
|
||||
skip_reason: None,
|
||||
is_cached: Some(false),
|
||||
status_code: Some(200),
|
||||
error_type: None,
|
||||
error_message: None,
|
||||
latency_ms: Some(123),
|
||||
concurrent_requests: Some(2),
|
||||
extra_data,
|
||||
required_capabilities: Some(json!({"streaming": true})),
|
||||
created_at_unix_ms: Some(created_at_unix_ms),
|
||||
started_at_unix_ms: Some(created_at_unix_ms + 1),
|
||||
finished_at_unix_ms: Some(created_at_unix_ms + 2),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -47,7 +47,7 @@ pub use proxy_nodes::MysqlProxyNodeReadRepository;
|
||||
pub use quota::MysqlProviderQuotaRepository;
|
||||
pub use routing_profiles::MysqlRoutingGroupRepository;
|
||||
pub use settlement::MysqlSettlementRepository;
|
||||
pub use usage::{MysqlUsageStorage, MysqlUsageWriteRepository};
|
||||
pub use usage::{MysqlUsageReadFilter, MysqlUsageStorage, MysqlUsageWriteRepository};
|
||||
pub use users::MysqlUserReadRepository;
|
||||
pub use video_tasks::MysqlVideoTaskRepository;
|
||||
pub use wallet::MysqlWalletReadRepository;
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
use sqlx::{
|
||||
migrate::{Migrate, MigrateError, Migrator},
|
||||
migrate::{AppliedMigration, Migrate, MigrateError, Migrator},
|
||||
MySqlPool,
|
||||
};
|
||||
|
||||
@@ -17,9 +17,15 @@ pub async fn pending_migrations(
|
||||
let mut conn = pool.acquire().await?;
|
||||
let applied_migrations = match conn.list_applied_migrations().await {
|
||||
Ok(applied_migrations) => applied_migrations,
|
||||
Err(err) if is_missing_sqlx_migrations_table_error(&err) => Vec::new(),
|
||||
Err(err) if is_missing_sqlx_migrations_table_error(&err) => {
|
||||
return Ok(pending_migrations_from_applied(&[]));
|
||||
}
|
||||
Err(err) => return Err(err),
|
||||
};
|
||||
if let Some(version) = conn.dirty_version().await? {
|
||||
return Err(MigrateError::Dirty(version));
|
||||
}
|
||||
validate_applied_migrations(&applied_migrations)?;
|
||||
Ok(pending_migrations_from_applied(&applied_migrations))
|
||||
}
|
||||
|
||||
@@ -56,9 +62,33 @@ fn pending_migrations_from_applied(
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn validate_applied_migrations(
|
||||
applied_migrations: &[AppliedMigration],
|
||||
) -> Result<(), MigrateError> {
|
||||
if MIGRATOR.ignore_missing {
|
||||
return Ok(());
|
||||
}
|
||||
let known_versions = MIGRATOR
|
||||
.iter()
|
||||
.map(|migration| migration.version)
|
||||
.collect::<std::collections::HashSet<_>>();
|
||||
if let Some(migration) = applied_migrations
|
||||
.iter()
|
||||
.find(|migration| !known_versions.contains(&migration.version))
|
||||
{
|
||||
return Err(MigrateError::VersionMissing(migration.version));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::MIGRATOR;
|
||||
use std::borrow::Cow;
|
||||
|
||||
use super::{
|
||||
pending_migrations, prepare_database_for_startup, validate_applied_migrations, MIGRATOR,
|
||||
};
|
||||
use sqlx::migrate::{AppliedMigration, MigrateError};
|
||||
|
||||
#[test]
|
||||
fn embeds_mysql_migration_sources() {
|
||||
@@ -69,4 +99,152 @@ mod tests {
|
||||
assert!(!versions.is_empty());
|
||||
assert!(versions.windows(2).all(|pair| pair[0] < pair[1]));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn embeds_cross_driver_schema_parity_migration() {
|
||||
let migration = MIGRATOR
|
||||
.iter()
|
||||
.find(|migration| migration.version == 20260725010000)
|
||||
.expect("cross-driver schema parity migration should be embedded");
|
||||
let sql = migration.sql.as_ref();
|
||||
|
||||
for required_fragment in [
|
||||
"CREATE TABLE IF NOT EXISTS usage_body_blobs",
|
||||
"CREATE TABLE IF NOT EXISTS usage_http_audits",
|
||||
"CREATE TABLE IF NOT EXISTS stats_summary",
|
||||
"CREATE TABLE IF NOT EXISTS user_model_usage_counts",
|
||||
"CREATE TABLE IF NOT EXISTS api_key_provider_mappings",
|
||||
"CREATE TABLE IF NOT EXISTS provider_usage_tracking",
|
||||
"ADD COLUMN `settlement_snapshot_schema_version`",
|
||||
"ADD COLUMN `billing_effective_input_tokens`",
|
||||
"ADD COLUMN `converted_request_body`",
|
||||
"ADD COLUMN `p99_first_byte_time_ms`",
|
||||
"idx_usage_stale_pending_created_request",
|
||||
] {
|
||||
assert!(
|
||||
sql.contains(required_fragment),
|
||||
"parity migration is missing {required_fragment}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn embeds_advanced_stats_parity_migration() {
|
||||
let migration = MIGRATOR
|
||||
.iter()
|
||||
.find(|migration| migration.version == 20260725020000)
|
||||
.expect("advanced stats parity migration should be embedded");
|
||||
let sql = migration.sql.as_ref();
|
||||
|
||||
for required_fragment in [
|
||||
"CREATE TABLE stats_user_summary",
|
||||
"CREATE TABLE stats_user_daily_api_format",
|
||||
"CREATE TABLE stats_user_daily_model_provider",
|
||||
"CREATE TABLE stats_daily_model_provider",
|
||||
"CREATE TABLE stats_daily_cost_savings",
|
||||
"CREATE TABLE stats_user_daily_cost_savings_model_provider",
|
||||
"ADD COLUMN completed_total_input_context",
|
||||
"ADD COLUMN settled_total_cost",
|
||||
"ADD COLUMN response_time_samples",
|
||||
"UPDATE stats_hourly SET is_complete = 0",
|
||||
"UPDATE stats_daily SET is_complete = 0",
|
||||
] {
|
||||
assert!(
|
||||
sql.contains(required_fragment),
|
||||
"advanced stats migration is missing {required_fragment}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_applied_migration_versions_unknown_to_this_binary() {
|
||||
let version = MIGRATOR
|
||||
.iter()
|
||||
.map(|migration| migration.version)
|
||||
.max()
|
||||
.expect("mysql migrations should not be empty")
|
||||
+ 1;
|
||||
let error = validate_applied_migrations(&[AppliedMigration {
|
||||
version,
|
||||
checksum: Cow::Borrowed(&[]),
|
||||
}])
|
||||
.expect_err("unknown applied migration should block startup");
|
||||
|
||||
assert!(matches!(error, MigrateError::VersionMissing(found) if found == version));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn pending_and_startup_preparation_reject_dirty_mysql_migration_state_when_url_is_set() {
|
||||
let Some(database_url) = std::env::var("AETHER_TEST_MYSQL_URL")
|
||||
.ok()
|
||||
.filter(|value| !value.trim().is_empty())
|
||||
else {
|
||||
eprintln!("skipping mysql dirty migration test because AETHER_TEST_MYSQL_URL is unset");
|
||||
return;
|
||||
};
|
||||
|
||||
let pool = sqlx::mysql::MySqlPoolOptions::new()
|
||||
.max_connections(1)
|
||||
.connect(&database_url)
|
||||
.await
|
||||
.expect("mysql test pool should connect");
|
||||
let dirty_version = MIGRATOR
|
||||
.iter()
|
||||
.next()
|
||||
.expect("mysql migrations should not be empty")
|
||||
.version;
|
||||
|
||||
let mut conn = pool
|
||||
.acquire()
|
||||
.await
|
||||
.expect("mysql connection should acquire");
|
||||
sqlx::query(
|
||||
r#"
|
||||
CREATE TEMPORARY TABLE _sqlx_migrations (
|
||||
version BIGINT PRIMARY KEY,
|
||||
description TEXT NOT NULL,
|
||||
installed_on TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
success BOOLEAN NOT NULL,
|
||||
checksum BLOB NOT NULL,
|
||||
execution_time BIGINT NOT NULL
|
||||
)
|
||||
"#,
|
||||
)
|
||||
.execute(&mut *conn)
|
||||
.await
|
||||
.expect("temporary mysql migrations table should create");
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO _sqlx_migrations (
|
||||
version,
|
||||
description,
|
||||
success,
|
||||
checksum,
|
||||
execution_time
|
||||
) VALUES (?, 'dirty test migration', FALSE, ?, 0)
|
||||
"#,
|
||||
)
|
||||
.bind(dirty_version)
|
||||
.bind(Vec::<u8>::new())
|
||||
.execute(&mut *conn)
|
||||
.await
|
||||
.expect("dirty mysql migration should insert");
|
||||
drop(conn);
|
||||
|
||||
let pending_error = pending_migrations(&pool)
|
||||
.await
|
||||
.expect_err("dirty mysql migration should fail pending inspection");
|
||||
assert!(
|
||||
matches!(&pending_error, MigrateError::Dirty(version) if *version == dirty_version),
|
||||
"unexpected pending migration error: {pending_error}"
|
||||
);
|
||||
|
||||
let preparation_error = prepare_database_for_startup(&pool)
|
||||
.await
|
||||
.expect_err("dirty mysql migration should fail startup preparation");
|
||||
assert!(
|
||||
matches!(&preparation_error, MigrateError::Dirty(version) if *version == dirty_version),
|
||||
"unexpected startup preparation error: {preparation_error}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -62,6 +62,14 @@ impl MysqlPoolFactory {
|
||||
.acquire_timeout(Duration::from_millis(acquire_timeout_ms))
|
||||
.idle_timeout(Duration::from_millis(idle_timeout_ms))
|
||||
.max_lifetime(Duration::from_millis(max_lifetime_ms))
|
||||
.after_connect(|connection, _metadata| {
|
||||
Box::pin(async move {
|
||||
sqlx::query("SET time_zone = '+00:00'")
|
||||
.execute(connection)
|
||||
.await?;
|
||||
Ok(())
|
||||
})
|
||||
})
|
||||
.connect_lazy_with(self.connect_options()?))
|
||||
}
|
||||
}
|
||||
@@ -90,4 +98,34 @@ mod tests {
|
||||
let factory = MysqlPoolFactory::new(config).expect("factory should build");
|
||||
let _pool = factory.connect_lazy().expect("lazy pool should build");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn factory_configures_utc_session_timezone_when_url_is_set() {
|
||||
let Some(database_url) = std::env::var("AETHER_TEST_MYSQL_URL")
|
||||
.ok()
|
||||
.filter(|value| !value.trim().is_empty())
|
||||
else {
|
||||
eprintln!("skipping mysql timezone test because AETHER_TEST_MYSQL_URL is unset");
|
||||
return;
|
||||
};
|
||||
let config = SqlDatabaseConfig {
|
||||
driver: DatabaseDriver::Mysql,
|
||||
url: database_url,
|
||||
pool: SqlPoolConfig {
|
||||
max_connections: 1,
|
||||
..SqlPoolConfig::default()
|
||||
},
|
||||
};
|
||||
|
||||
let pool = MysqlPoolFactory::new(config)
|
||||
.expect("factory should build")
|
||||
.connect_lazy()
|
||||
.expect("lazy pool should build");
|
||||
let timezone: String = sqlx::query_scalar("SELECT @@session.time_zone")
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.expect("mysql session timezone should load");
|
||||
|
||||
assert_eq!(timezone, "+00:00");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -9,15 +9,19 @@ use sqlx::{
|
||||
|
||||
use aether_data_contracts::repository::provider_catalog::{
|
||||
ProviderCatalogKeyAdaptiveStateUpdate, ProviderCatalogKeyHealthStateUpdate,
|
||||
ProviderCatalogKeyListQuery, ProviderCatalogKeyOAuthRuntimeStateCasUpdate,
|
||||
ProviderCatalogKeyRuntimeMetadataUpdate, ProviderCatalogKeyStatusSnapshotUpdate,
|
||||
ProviderCatalogReadRepository, ProviderCatalogSnapshot,
|
||||
ProviderCatalogKeyListOrder, ProviderCatalogKeyListQuery,
|
||||
ProviderCatalogKeyOAuthRuntimeStateCasUpdate, ProviderCatalogKeyRuntimeMetadataUpdate,
|
||||
ProviderCatalogKeyStatusSnapshotUpdate, ProviderCatalogReadRepository,
|
||||
ProviderCatalogUpstreamMetadataNamespaceUpdate, ProviderCatalogWriteRepository,
|
||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
|
||||
StoredProviderCatalogKeyMaintenanceSummary, StoredProviderCatalogKeyPage,
|
||||
StoredProviderCatalogKeyStats, StoredProviderCatalogProvider,
|
||||
};
|
||||
use aether_data_contracts::DataLayerError;
|
||||
use aether_data_query::{
|
||||
push_ci_contains_any, push_eq, push_in, push_limit_offset, push_optional_eq, SqlDialect,
|
||||
WhereClause,
|
||||
};
|
||||
|
||||
use crate::error::SqlResultExt;
|
||||
use crate::MysqlPool;
|
||||
@@ -27,6 +31,30 @@ pub struct MysqlProviderCatalogReadRepository {
|
||||
pool: MysqlPool,
|
||||
}
|
||||
|
||||
const PROVIDER_SELECT_SQL: &str = r#"
|
||||
SELECT
|
||||
id, name, description, website, provider_type, billing_type,
|
||||
monthly_quota_usd, monthly_used_usd, quota_reset_day,
|
||||
quota_last_reset_at AS quota_last_reset_at_unix_secs,
|
||||
quota_expires_at AS quota_expires_at_unix_secs,
|
||||
provider_priority, is_active, keep_priority_on_conversion,
|
||||
enable_format_conversion, concurrent_limit, max_retries, proxy,
|
||||
request_timeout, stream_first_byte_timeout, config,
|
||||
created_at AS created_at_unix_ms,
|
||||
updated_at AS updated_at_unix_secs
|
||||
FROM providers
|
||||
"#;
|
||||
|
||||
const ENDPOINT_SELECT_SQL: &str = r#"
|
||||
SELECT
|
||||
id, provider_id, api_format, api_family, endpoint_kind, is_active,
|
||||
health_score, base_url, header_rules, body_rules, max_retries,
|
||||
custom_path, config, format_acceptance_config, proxy,
|
||||
created_at AS created_at_unix_ms,
|
||||
updated_at AS updated_at_unix_secs
|
||||
FROM provider_endpoints
|
||||
"#;
|
||||
|
||||
const KEY_SELECT_SQL: &str = r#"
|
||||
SELECT
|
||||
id, provider_id, name, auth_type, capabilities, is_active, api_formats,
|
||||
@@ -54,69 +82,176 @@ SELECT
|
||||
FROM provider_api_keys
|
||||
"#;
|
||||
|
||||
const KEY_SUMMARY_SELECT_SQL: &str = r#"
|
||||
SELECT
|
||||
id,
|
||||
provider_id,
|
||||
COALESCE(NULLIF(name, ''), id) AS name,
|
||||
COALESCE(NULLIF(auth_type, ''), 'summary') AS auth_type,
|
||||
NULL AS capabilities,
|
||||
is_active,
|
||||
api_formats,
|
||||
NULL AS auth_type_by_format,
|
||||
NULL AS allow_auth_channel_mismatch_formats,
|
||||
'summary' AS api_key,
|
||||
CASE WHEN auth_config IS NULL THEN NULL ELSE '{}' END AS auth_config,
|
||||
NULL AS note,
|
||||
NULL AS internal_priority,
|
||||
NULL AS rate_multipliers,
|
||||
NULL AS global_priority_by_format,
|
||||
NULL AS allowed_models,
|
||||
NULL AS expires_at_unix_secs,
|
||||
NULL AS cache_ttl_minutes,
|
||||
NULL AS max_probe_interval_minutes,
|
||||
NULL AS proxy,
|
||||
NULL AS fingerprint,
|
||||
NULL AS rpm_limit,
|
||||
NULL AS concurrent_limit,
|
||||
NULL AS learned_rpm_limit,
|
||||
NULL AS concurrent_429_count,
|
||||
NULL AS rpm_429_count,
|
||||
NULL AS last_429_at_unix_secs,
|
||||
NULL AS last_429_type,
|
||||
NULL AS adjustment_history,
|
||||
NULL AS utilization_samples,
|
||||
NULL AS last_probe_increase_at_unix_secs,
|
||||
NULL AS last_rpm_peak,
|
||||
NULL AS request_count,
|
||||
CAST(0 AS SIGNED) AS total_tokens,
|
||||
CAST(0 AS DOUBLE) AS total_cost_usd,
|
||||
NULL AS success_count,
|
||||
NULL AS error_count,
|
||||
NULL AS total_response_time_ms,
|
||||
NULL AS last_used_at_unix_secs,
|
||||
FALSE AS auto_fetch_models,
|
||||
NULL AS last_models_fetch_at_unix_secs,
|
||||
NULL AS last_models_fetch_error,
|
||||
NULL AS locked_models,
|
||||
NULL AS model_include_patterns,
|
||||
NULL AS model_exclude_patterns,
|
||||
NULL AS upstream_metadata,
|
||||
NULL AS oauth_invalid_at_unix_secs,
|
||||
NULL AS oauth_invalid_reason,
|
||||
NULL AS status_snapshot,
|
||||
NULL AS created_at_unix_ms,
|
||||
NULL AS updated_at_unix_secs,
|
||||
health_by_format,
|
||||
NULL AS circuit_breaker_by_format
|
||||
FROM provider_api_keys
|
||||
"#;
|
||||
|
||||
const KEY_MAINTENANCE_SUMMARY_SELECT_SQL: &str = r#"
|
||||
SELECT id, provider_id, is_active, upstream_metadata
|
||||
FROM provider_api_keys
|
||||
"#;
|
||||
|
||||
const KEY_STATS_SELECT_SQL: &str = r#"
|
||||
SELECT
|
||||
provider_id,
|
||||
CAST(COUNT(*) AS SIGNED) AS total_keys,
|
||||
CAST(SUM(CASE WHEN is_active THEN 1 ELSE 0 END) AS SIGNED) AS active_keys
|
||||
FROM provider_api_keys
|
||||
"#;
|
||||
|
||||
impl MysqlProviderCatalogReadRepository {
|
||||
pub fn new(pool: MysqlPool) -> Self {
|
||||
Self { pool }
|
||||
}
|
||||
|
||||
async fn load_snapshot(&self) -> Result<ProviderCatalogSnapshot, DataLayerError> {
|
||||
Ok(ProviderCatalogSnapshot::new(
|
||||
self.load_providers().await?,
|
||||
self.load_endpoints().await?,
|
||||
self.load_keys().await?,
|
||||
))
|
||||
}
|
||||
pub async fn list_providers_by_ids(
|
||||
&self,
|
||||
provider_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogProvider>, DataLayerError> {
|
||||
if provider_ids.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
async fn load_providers(&self) -> Result<Vec<StoredProviderCatalogProvider>, DataLayerError> {
|
||||
let rows = sqlx::query(
|
||||
r#"
|
||||
SELECT
|
||||
id, name, description, website, provider_type, billing_type,
|
||||
monthly_quota_usd, monthly_used_usd, quota_reset_day,
|
||||
quota_last_reset_at AS quota_last_reset_at_unix_secs,
|
||||
quota_expires_at AS quota_expires_at_unix_secs,
|
||||
provider_priority, is_active, keep_priority_on_conversion,
|
||||
enable_format_conversion, concurrent_limit, max_retries, proxy,
|
||||
request_timeout, stream_first_byte_timeout, config,
|
||||
created_at AS created_at_unix_ms,
|
||||
updated_at AS updated_at_unix_secs
|
||||
FROM providers
|
||||
"#,
|
||||
let rows = build_in_query(
|
||||
PROVIDER_SELECT_SQL,
|
||||
"id",
|
||||
provider_ids,
|
||||
" ORDER BY name ASC",
|
||||
)
|
||||
.build()
|
||||
.fetch_all(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
rows.iter().map(map_provider_row).collect()
|
||||
}
|
||||
|
||||
async fn load_endpoints(&self) -> Result<Vec<StoredProviderCatalogEndpoint>, DataLayerError> {
|
||||
let rows = sqlx::query(
|
||||
r#"
|
||||
SELECT
|
||||
id, provider_id, api_format, api_family, endpoint_kind, is_active,
|
||||
health_score, base_url, header_rules, body_rules, max_retries,
|
||||
custom_path, config, format_acceptance_config, proxy,
|
||||
created_at AS created_at_unix_ms,
|
||||
updated_at AS updated_at_unix_secs
|
||||
FROM provider_endpoints
|
||||
WHERE api_format IS NOT NULL
|
||||
"#,
|
||||
pub async fn list_providers(
|
||||
&self,
|
||||
active_only: bool,
|
||||
) -> Result<Vec<StoredProviderCatalogProvider>, DataLayerError> {
|
||||
let mut builder = QueryBuilder::<MySql>::new(PROVIDER_SELECT_SQL);
|
||||
let mut where_clause = WhereClause::new();
|
||||
if active_only {
|
||||
push_eq(&mut builder, &mut where_clause, "is_active", true);
|
||||
}
|
||||
builder.push(" ORDER BY provider_priority ASC, name ASC");
|
||||
let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?;
|
||||
rows.iter().map(map_provider_row).collect()
|
||||
}
|
||||
|
||||
pub async fn list_endpoints_by_ids(
|
||||
&self,
|
||||
endpoint_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogEndpoint>, DataLayerError> {
|
||||
if endpoint_ids.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let rows = build_in_query(
|
||||
ENDPOINT_SELECT_SQL,
|
||||
"id",
|
||||
endpoint_ids,
|
||||
" ORDER BY api_format ASC, id ASC",
|
||||
)
|
||||
.build()
|
||||
.fetch_all(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
rows.iter().map(map_endpoint_row).collect()
|
||||
}
|
||||
|
||||
async fn load_keys(&self) -> Result<Vec<StoredProviderCatalogKey>, DataLayerError> {
|
||||
let rows = sqlx::query(KEY_SELECT_SQL)
|
||||
pub async fn list_endpoints_by_provider_ids(
|
||||
&self,
|
||||
provider_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogEndpoint>, DataLayerError> {
|
||||
if provider_ids.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let rows = build_in_query(
|
||||
ENDPOINT_SELECT_SQL,
|
||||
"provider_id",
|
||||
provider_ids,
|
||||
" ORDER BY provider_id ASC, api_format ASC, id ASC",
|
||||
)
|
||||
.build()
|
||||
.fetch_all(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
rows.iter().map(map_endpoint_row).collect()
|
||||
}
|
||||
|
||||
pub async fn list_keys_by_ids(
|
||||
&self,
|
||||
key_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogKey>, DataLayerError> {
|
||||
if key_ids.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let rows = build_in_query(KEY_SELECT_SQL, "id", key_ids, " ORDER BY name ASC, id ASC")
|
||||
.build()
|
||||
.fetch_all(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
rows.iter().map(map_key_row).collect()
|
||||
}
|
||||
|
||||
async fn list_keys_by_provider_ids_direct(
|
||||
pub async fn list_keys_by_provider_ids(
|
||||
&self,
|
||||
provider_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogKey>, DataLayerError> {
|
||||
@@ -124,12 +259,120 @@ WHERE api_format IS NOT NULL
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let rows = build_list_keys_by_provider_ids_query(provider_ids)
|
||||
let rows = build_in_query(
|
||||
KEY_SELECT_SQL,
|
||||
"provider_id",
|
||||
provider_ids,
|
||||
" ORDER BY provider_id ASC, name ASC, id ASC",
|
||||
)
|
||||
.build()
|
||||
.fetch_all(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
rows.iter().map(map_key_row).collect()
|
||||
}
|
||||
|
||||
pub async fn list_key_summaries_by_provider_ids(
|
||||
&self,
|
||||
provider_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogKey>, DataLayerError> {
|
||||
if provider_ids.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let rows = build_in_query(
|
||||
KEY_SUMMARY_SELECT_SQL,
|
||||
"provider_id",
|
||||
provider_ids,
|
||||
" ORDER BY provider_id ASC, id ASC",
|
||||
)
|
||||
.build()
|
||||
.fetch_all(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
rows.iter().map(map_key_row).collect()
|
||||
}
|
||||
|
||||
pub async fn list_key_maintenance_summaries_by_provider_ids(
|
||||
&self,
|
||||
provider_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogKeyMaintenanceSummary>, DataLayerError> {
|
||||
if provider_ids.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let rows = build_in_query(
|
||||
KEY_MAINTENANCE_SUMMARY_SELECT_SQL,
|
||||
"provider_id",
|
||||
provider_ids,
|
||||
" ORDER BY provider_id ASC, id ASC",
|
||||
)
|
||||
.build()
|
||||
.fetch_all(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
rows.iter().map(map_key_maintenance_summary_row).collect()
|
||||
}
|
||||
|
||||
pub async fn list_keys_page(
|
||||
&self,
|
||||
query: &ProviderCatalogKeyListQuery,
|
||||
) -> Result<StoredProviderCatalogKeyPage, DataLayerError> {
|
||||
validate_non_empty(&query.provider_id, "provider catalog provider_id")?;
|
||||
let offset = i64::try_from(query.offset).map_err(|_| {
|
||||
DataLayerError::InvalidInput(format!(
|
||||
"invalid provider catalog key offset: {}",
|
||||
query.offset
|
||||
))
|
||||
})?;
|
||||
let limit = i64::try_from(query.limit).map_err(|_| {
|
||||
DataLayerError::InvalidInput(format!(
|
||||
"invalid provider catalog key limit: {}",
|
||||
query.limit
|
||||
))
|
||||
})?;
|
||||
|
||||
let mut count_builder = build_key_page_count_query(query);
|
||||
let total = count_builder
|
||||
.build_query_scalar::<i64>()
|
||||
.fetch_one(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?
|
||||
.max(0) as usize;
|
||||
|
||||
let mut list_builder = build_key_page_list_query(query, limit, offset);
|
||||
let rows = list_builder
|
||||
.build()
|
||||
.fetch_all(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
rows.iter().map(map_key_row).collect()
|
||||
let items = rows
|
||||
.iter()
|
||||
.map(map_key_row)
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
|
||||
Ok(StoredProviderCatalogKeyPage { items, total })
|
||||
}
|
||||
|
||||
pub async fn list_key_stats_by_provider_ids(
|
||||
&self,
|
||||
provider_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogKeyStats>, DataLayerError> {
|
||||
if provider_ids.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let rows = build_in_query(
|
||||
KEY_STATS_SELECT_SQL,
|
||||
"provider_id",
|
||||
provider_ids,
|
||||
" GROUP BY provider_id ORDER BY provider_id ASC",
|
||||
)
|
||||
.build()
|
||||
.fetch_all(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
rows.iter().map(map_key_stats_row).collect()
|
||||
}
|
||||
|
||||
pub async fn create_provider(
|
||||
@@ -1443,84 +1686,70 @@ impl ProviderCatalogReadRepository for MysqlProviderCatalogReadRepository {
|
||||
&self,
|
||||
active_only: bool,
|
||||
) -> Result<Vec<StoredProviderCatalogProvider>, DataLayerError> {
|
||||
Ok(self.load_snapshot().await?.list_providers(active_only))
|
||||
Self::list_providers(self, active_only).await
|
||||
}
|
||||
|
||||
async fn list_providers_by_ids(
|
||||
&self,
|
||||
provider_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogProvider>, DataLayerError> {
|
||||
Ok(self
|
||||
.load_snapshot()
|
||||
.await?
|
||||
.list_providers_by_ids(provider_ids))
|
||||
Self::list_providers_by_ids(self, provider_ids).await
|
||||
}
|
||||
|
||||
async fn list_endpoints_by_ids(
|
||||
&self,
|
||||
endpoint_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogEndpoint>, DataLayerError> {
|
||||
Ok(self
|
||||
.load_snapshot()
|
||||
.await?
|
||||
.list_endpoints_by_ids(endpoint_ids))
|
||||
Self::list_endpoints_by_ids(self, endpoint_ids).await
|
||||
}
|
||||
|
||||
async fn list_endpoints_by_provider_ids(
|
||||
&self,
|
||||
provider_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogEndpoint>, DataLayerError> {
|
||||
Ok(self
|
||||
.load_snapshot()
|
||||
.await?
|
||||
.list_endpoints_by_provider_ids(provider_ids))
|
||||
Self::list_endpoints_by_provider_ids(self, provider_ids).await
|
||||
}
|
||||
|
||||
async fn list_keys_by_ids(
|
||||
&self,
|
||||
key_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogKey>, DataLayerError> {
|
||||
Ok(self.load_snapshot().await?.list_keys_by_ids(key_ids))
|
||||
Self::list_keys_by_ids(self, key_ids).await
|
||||
}
|
||||
|
||||
async fn list_keys_by_provider_ids(
|
||||
&self,
|
||||
provider_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogKey>, DataLayerError> {
|
||||
self.list_keys_by_provider_ids_direct(provider_ids).await
|
||||
Self::list_keys_by_provider_ids(self, provider_ids).await
|
||||
}
|
||||
|
||||
async fn list_key_summaries_by_provider_ids(
|
||||
&self,
|
||||
provider_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogKey>, DataLayerError> {
|
||||
self.list_keys_by_provider_ids_direct(provider_ids).await
|
||||
Self::list_key_summaries_by_provider_ids(self, provider_ids).await
|
||||
}
|
||||
|
||||
async fn list_key_maintenance_summaries_by_provider_ids(
|
||||
&self,
|
||||
provider_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogKeyMaintenanceSummary>, DataLayerError> {
|
||||
Ok(self
|
||||
.load_snapshot()
|
||||
.await?
|
||||
.list_key_maintenance_summaries_by_provider_ids(provider_ids))
|
||||
Self::list_key_maintenance_summaries_by_provider_ids(self, provider_ids).await
|
||||
}
|
||||
|
||||
async fn list_keys_page(
|
||||
&self,
|
||||
query: &ProviderCatalogKeyListQuery,
|
||||
) -> Result<StoredProviderCatalogKeyPage, DataLayerError> {
|
||||
Ok(self.load_snapshot().await?.list_keys_page(query))
|
||||
Self::list_keys_page(self, query).await
|
||||
}
|
||||
|
||||
async fn list_key_stats_by_provider_ids(
|
||||
&self,
|
||||
provider_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogKeyStats>, DataLayerError> {
|
||||
self.load_snapshot()
|
||||
.await?
|
||||
.list_key_stats_by_provider_ids(provider_ids)
|
||||
Self::list_key_stats_by_provider_ids(self, provider_ids).await
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1975,17 +2204,87 @@ fn optional_json_to_string(
|
||||
optional_json_ref_to_string(value.as_ref(), field_name)
|
||||
}
|
||||
|
||||
fn build_list_keys_by_provider_ids_query(provider_ids: &[String]) -> QueryBuilder<'_, MySql> {
|
||||
fn build_in_query<'a>(
|
||||
select_sql: &'static str,
|
||||
column: &'static str,
|
||||
values: &'a [String],
|
||||
suffix: &'static str,
|
||||
) -> QueryBuilder<'a, MySql> {
|
||||
let mut builder = QueryBuilder::<MySql>::new(select_sql);
|
||||
let mut where_clause = WhereClause::new();
|
||||
push_in(&mut builder, &mut where_clause, column, values);
|
||||
builder.push(suffix);
|
||||
builder
|
||||
}
|
||||
|
||||
fn build_key_page_count_query<'a>(
|
||||
query: &'a ProviderCatalogKeyListQuery,
|
||||
) -> QueryBuilder<'a, MySql> {
|
||||
let mut builder = QueryBuilder::<MySql>::new(
|
||||
"SELECT CAST(COUNT(*) AS SIGNED) AS total FROM provider_api_keys",
|
||||
);
|
||||
let mut where_clause = WhereClause::new();
|
||||
apply_key_page_filters(&mut builder, &mut where_clause, query);
|
||||
builder
|
||||
}
|
||||
|
||||
fn build_key_page_list_query<'a>(
|
||||
query: &'a ProviderCatalogKeyListQuery,
|
||||
limit: i64,
|
||||
offset: i64,
|
||||
) -> QueryBuilder<'a, MySql> {
|
||||
let mut builder = QueryBuilder::<MySql>::new(KEY_SELECT_SQL);
|
||||
builder.push("WHERE provider_id IN (");
|
||||
{
|
||||
let mut separated = builder.separated(", ");
|
||||
for provider_id in provider_ids {
|
||||
separated.push_bind(provider_id.clone());
|
||||
let mut where_clause = WhereClause::new();
|
||||
apply_key_page_filters(&mut builder, &mut where_clause, query);
|
||||
builder
|
||||
.push(" ORDER BY ")
|
||||
.push(key_page_order_by(&query.order));
|
||||
push_limit_offset(&mut builder, limit, offset);
|
||||
builder
|
||||
}
|
||||
|
||||
fn key_page_order_by(order: &ProviderCatalogKeyListOrder) -> &'static str {
|
||||
match order {
|
||||
ProviderCatalogKeyListOrder::Name => "internal_priority ASC, name ASC, id ASC",
|
||||
ProviderCatalogKeyListOrder::CreatedAt => {
|
||||
"internal_priority ASC, COALESCE(created_at, 0) ASC, id ASC"
|
||||
}
|
||||
ProviderCatalogKeyListOrder::CreatedAtAsc => {
|
||||
"created_at IS NULL ASC, created_at ASC, name ASC, id ASC"
|
||||
}
|
||||
ProviderCatalogKeyListOrder::CreatedAtDesc => {
|
||||
"created_at IS NULL ASC, created_at DESC, name ASC, id ASC"
|
||||
}
|
||||
ProviderCatalogKeyListOrder::LastUsedAtAsc => {
|
||||
"last_used_at IS NULL ASC, last_used_at ASC, name ASC, id ASC"
|
||||
}
|
||||
ProviderCatalogKeyListOrder::LastUsedAtDesc => {
|
||||
"last_used_at IS NULL ASC, last_used_at DESC, name ASC, id ASC"
|
||||
}
|
||||
}
|
||||
builder.push(") ORDER BY provider_id ASC, name ASC, id ASC");
|
||||
builder
|
||||
}
|
||||
|
||||
fn apply_key_page_filters<'a>(
|
||||
builder: &mut QueryBuilder<'a, MySql>,
|
||||
where_clause: &mut WhereClause,
|
||||
query: &'a ProviderCatalogKeyListQuery,
|
||||
) {
|
||||
push_eq(
|
||||
builder,
|
||||
where_clause,
|
||||
"provider_id",
|
||||
query.provider_id.clone(),
|
||||
);
|
||||
if let Some(search) = query.search.as_deref() {
|
||||
push_ci_contains_any(
|
||||
builder,
|
||||
where_clause,
|
||||
SqlDialect::MySql,
|
||||
&["name", "id"],
|
||||
search,
|
||||
);
|
||||
}
|
||||
push_optional_eq(builder, where_clause, "is_active", query.is_active);
|
||||
}
|
||||
|
||||
fn key_insert_sql() -> &'static str {
|
||||
@@ -2258,6 +2557,28 @@ fn map_endpoint_row(row: &MySqlRow) -> Result<StoredProviderCatalogEndpoint, Dat
|
||||
)
|
||||
}
|
||||
|
||||
fn map_key_stats_row(row: &MySqlRow) -> Result<StoredProviderCatalogKeyStats, DataLayerError> {
|
||||
StoredProviderCatalogKeyStats::new(
|
||||
row.try_get("provider_id").map_sql_err()?,
|
||||
row.try_get("total_keys").map_sql_err()?,
|
||||
row.try_get("active_keys").map_sql_err()?,
|
||||
)
|
||||
}
|
||||
|
||||
fn map_key_maintenance_summary_row(
|
||||
row: &MySqlRow,
|
||||
) -> Result<StoredProviderCatalogKeyMaintenanceSummary, DataLayerError> {
|
||||
Ok(StoredProviderCatalogKeyMaintenanceSummary {
|
||||
id: row.try_get("id").map_sql_err()?,
|
||||
provider_id: row.try_get("provider_id").map_sql_err()?,
|
||||
is_active: row.try_get("is_active").map_sql_err()?,
|
||||
upstream_metadata: optional_json_from_string(
|
||||
row.try_get("upstream_metadata").map_sql_err()?,
|
||||
"provider_api_keys.upstream_metadata",
|
||||
)?,
|
||||
})
|
||||
}
|
||||
|
||||
fn map_key_row(row: &MySqlRow) -> Result<StoredProviderCatalogKey, DataLayerError> {
|
||||
let total_cost_usd = row
|
||||
.try_get::<Option<f64>, _>("total_cost_usd")
|
||||
@@ -2387,9 +2708,9 @@ fn map_key_row(row: &MySqlRow) -> Result<StoredProviderCatalogKey, DataLayerErro
|
||||
.map_sql_err()?,
|
||||
"provider_api_keys.allow_auth_channel_mismatch_formats",
|
||||
)?;
|
||||
key.internal_priority = row.try_get("internal_priority").map_sql_err()?;
|
||||
key.cache_ttl_minutes = row.try_get("cache_ttl_minutes").map_sql_err()?;
|
||||
key.max_probe_interval_minutes = row.try_get("max_probe_interval_minutes").map_sql_err()?;
|
||||
key.internal_priority = row.try_get("internal_priority").unwrap_or(50);
|
||||
key.cache_ttl_minutes = row.try_get("cache_ttl_minutes").unwrap_or(5);
|
||||
key.max_probe_interval_minutes = row.try_get("max_probe_interval_minutes").unwrap_or(32);
|
||||
key.last_429_type = row.try_get("last_429_type").map_sql_err()?;
|
||||
key.utilization_samples = optional_json_from_string(
|
||||
row.try_get("utilization_samples").map_sql_err()?,
|
||||
@@ -2457,10 +2778,11 @@ mod tests {
|
||||
use super::MysqlProviderCatalogReadRepository;
|
||||
use crate::run_migrations;
|
||||
use aether_data_contracts::repository::provider_catalog::{
|
||||
ProviderCatalogReadRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
|
||||
StoredProviderCatalogProvider,
|
||||
ProviderCatalogKeyListOrder, ProviderCatalogKeyListQuery, StoredProviderCatalogEndpoint,
|
||||
StoredProviderCatalogKey, StoredProviderCatalogProvider,
|
||||
};
|
||||
use serde_json::json;
|
||||
use sqlx::Execute;
|
||||
|
||||
#[test]
|
||||
fn ordinary_key_update_does_not_own_adaptive_runtime_fields() {
|
||||
@@ -2478,23 +2800,66 @@ mod tests {
|
||||
assert!(!sql.contains(runtime_assignment));
|
||||
}
|
||||
}
|
||||
use sqlx::Execute;
|
||||
|
||||
#[tokio::test]
|
||||
async fn repository_builds_from_lazy_pool() {
|
||||
async fn empty_id_lists_do_not_connect_to_lazy_pool() {
|
||||
let pool = sqlx::mysql::MySqlPoolOptions::new().connect_lazy_with(
|
||||
"mysql://user:pass@localhost:3306/aether"
|
||||
.parse()
|
||||
.expect("mysql options should parse"),
|
||||
);
|
||||
|
||||
let _repository = MysqlProviderCatalogReadRepository::new(pool);
|
||||
let repository = MysqlProviderCatalogReadRepository::new(pool);
|
||||
assert!(repository
|
||||
.list_providers_by_ids(&[])
|
||||
.await
|
||||
.expect("empty provider ids should not connect")
|
||||
.is_empty());
|
||||
assert!(repository
|
||||
.list_endpoints_by_ids(&[])
|
||||
.await
|
||||
.expect("empty endpoint ids should not connect")
|
||||
.is_empty());
|
||||
assert!(repository
|
||||
.list_endpoints_by_provider_ids(&[])
|
||||
.await
|
||||
.expect("empty endpoint provider ids should not connect")
|
||||
.is_empty());
|
||||
assert!(repository
|
||||
.list_keys_by_ids(&[])
|
||||
.await
|
||||
.expect("empty key ids should not connect")
|
||||
.is_empty());
|
||||
assert!(repository
|
||||
.list_keys_by_provider_ids(&[])
|
||||
.await
|
||||
.expect("empty key provider ids should not connect")
|
||||
.is_empty());
|
||||
assert!(repository
|
||||
.list_key_summaries_by_provider_ids(&[])
|
||||
.await
|
||||
.expect("empty key summary provider ids should not connect")
|
||||
.is_empty());
|
||||
assert!(repository
|
||||
.list_key_maintenance_summaries_by_provider_ids(&[])
|
||||
.await
|
||||
.expect("empty maintenance provider ids should not connect")
|
||||
.is_empty());
|
||||
assert!(repository
|
||||
.list_key_stats_by_provider_ids(&[])
|
||||
.await
|
||||
.expect("empty stats provider ids should not connect")
|
||||
.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn list_keys_by_provider_ids_query_targets_index_aligned_ordering() {
|
||||
fn list_query_targets_ids_and_index_aligned_ordering() {
|
||||
let provider_ids = vec!["provider-a".to_string(), "provider-b".to_string()];
|
||||
let mut builder = super::build_list_keys_by_provider_ids_query(&provider_ids);
|
||||
let mut builder = super::build_in_query(
|
||||
super::KEY_SELECT_SQL,
|
||||
"provider_id",
|
||||
&provider_ids,
|
||||
" ORDER BY provider_id ASC, name ASC, id ASC",
|
||||
);
|
||||
let query = builder.build();
|
||||
let sql = query.sql();
|
||||
|
||||
@@ -2502,6 +2867,64 @@ mod tests {
|
||||
assert!(sql.contains("ORDER BY provider_id ASC, name ASC, id ASC"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn key_page_builders_apply_matching_filters_and_requested_order() {
|
||||
let query = ProviderCatalogKeyListQuery {
|
||||
provider_id: "provider-a".to_string(),
|
||||
search: Some("Primary".to_string()),
|
||||
is_active: Some(true),
|
||||
offset: 20,
|
||||
limit: 10,
|
||||
order: ProviderCatalogKeyListOrder::LastUsedAtDesc,
|
||||
};
|
||||
let mut count_builder = super::build_key_page_count_query(&query);
|
||||
let count_sql = count_builder.build().sql().to_string();
|
||||
let mut list_builder = super::build_key_page_list_query(&query, 10, 20);
|
||||
let list_sql = list_builder.build().sql().to_string();
|
||||
|
||||
for sql in [&count_sql, &list_sql] {
|
||||
assert!(sql.contains("provider_id = ?"));
|
||||
assert!(sql.contains("LOWER(name) LIKE ?"));
|
||||
assert!(sql.contains("LOWER(id) LIKE ?"));
|
||||
assert!(sql.contains("is_active = ?"));
|
||||
}
|
||||
assert!(list_sql
|
||||
.contains("ORDER BY last_used_at IS NULL ASC, last_used_at DESC, name ASC, id ASC"));
|
||||
assert!(list_sql.contains("LIMIT ? OFFSET ?"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn key_page_orders_match_cross_driver_contract() {
|
||||
for (order, expected) in [
|
||||
(
|
||||
ProviderCatalogKeyListOrder::Name,
|
||||
"internal_priority ASC, name ASC, id ASC",
|
||||
),
|
||||
(
|
||||
ProviderCatalogKeyListOrder::CreatedAt,
|
||||
"internal_priority ASC, COALESCE(created_at, 0) ASC, id ASC",
|
||||
),
|
||||
(
|
||||
ProviderCatalogKeyListOrder::CreatedAtAsc,
|
||||
"created_at IS NULL ASC, created_at ASC, name ASC, id ASC",
|
||||
),
|
||||
(
|
||||
ProviderCatalogKeyListOrder::CreatedAtDesc,
|
||||
"created_at IS NULL ASC, created_at DESC, name ASC, id ASC",
|
||||
),
|
||||
(
|
||||
ProviderCatalogKeyListOrder::LastUsedAtAsc,
|
||||
"last_used_at IS NULL ASC, last_used_at ASC, name ASC, id ASC",
|
||||
),
|
||||
(
|
||||
ProviderCatalogKeyListOrder::LastUsedAtDesc,
|
||||
"last_used_at IS NULL ASC, last_used_at DESC, name ASC, id ASC",
|
||||
),
|
||||
] {
|
||||
assert_eq!(super::key_page_order_by(&order), expected);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn mysql_provider_catalog_repository_round_trips_when_url_is_set() {
|
||||
let Some(database_url) = std::env::var("AETHER_TEST_MYSQL_URL")
|
||||
@@ -2634,6 +3057,34 @@ mod tests {
|
||||
.expect("key should create");
|
||||
assert_eq!(created_key.concurrent_limit, Some(3));
|
||||
|
||||
let providers = repository
|
||||
.list_providers_by_ids(std::slice::from_ref(&provider_id))
|
||||
.await
|
||||
.expect("providers should list by ids");
|
||||
assert_eq!(providers.len(), 1);
|
||||
assert_eq!(providers[0].id, provider_id);
|
||||
assert!(repository
|
||||
.list_providers(true)
|
||||
.await
|
||||
.expect("active providers should list")
|
||||
.iter()
|
||||
.any(|provider| provider.id == provider_id));
|
||||
|
||||
let endpoints = repository
|
||||
.list_endpoints_by_ids(std::slice::from_ref(&endpoint_id))
|
||||
.await
|
||||
.expect("endpoints should list by ids");
|
||||
assert_eq!(endpoints.len(), 1);
|
||||
assert_eq!(endpoints[0].id, endpoint_id);
|
||||
assert_eq!(
|
||||
repository
|
||||
.list_endpoints_by_provider_ids(std::slice::from_ref(&provider_id))
|
||||
.await
|
||||
.expect("endpoints should list by provider")
|
||||
.len(),
|
||||
1
|
||||
);
|
||||
|
||||
let keys = repository
|
||||
.list_keys_by_provider_ids(std::slice::from_ref(&provider_id))
|
||||
.await
|
||||
@@ -2641,6 +3092,36 @@ mod tests {
|
||||
assert_eq!(keys.len(), 1);
|
||||
assert_eq!(keys[0].total_tokens, 1234);
|
||||
|
||||
let summaries = repository
|
||||
.list_key_summaries_by_provider_ids(std::slice::from_ref(&provider_id))
|
||||
.await
|
||||
.expect("key summaries should list");
|
||||
assert_eq!(summaries.len(), 1);
|
||||
assert_eq!(summaries[0].encrypted_api_key.as_deref(), Some("summary"));
|
||||
assert_eq!(summaries[0].total_tokens, 0);
|
||||
|
||||
let page = repository
|
||||
.list_keys_page(&ProviderCatalogKeyListQuery {
|
||||
provider_id: provider_id.clone(),
|
||||
search: Some("default".to_string()),
|
||||
is_active: Some(true),
|
||||
offset: 0,
|
||||
limit: 10,
|
||||
order: ProviderCatalogKeyListOrder::Name,
|
||||
})
|
||||
.await
|
||||
.expect("filtered key page should list");
|
||||
assert_eq!(page.total, 1);
|
||||
assert_eq!(page.items[0].id, key_id);
|
||||
|
||||
let stats = repository
|
||||
.list_key_stats_by_provider_ids(std::slice::from_ref(&provider_id))
|
||||
.await
|
||||
.expect("key stats should list");
|
||||
assert_eq!(stats.len(), 1);
|
||||
assert_eq!(stats[0].total_keys, 1);
|
||||
assert_eq!(stats[0].active_keys, 1);
|
||||
|
||||
assert!(repository
|
||||
.update_key_upstream_metadata(
|
||||
&key_id,
|
||||
@@ -2671,6 +3152,17 @@ mod tests {
|
||||
);
|
||||
assert!(!reloaded_key.is_active);
|
||||
|
||||
let maintenance = repository
|
||||
.list_key_maintenance_summaries_by_provider_ids(std::slice::from_ref(&provider_id))
|
||||
.await
|
||||
.expect("maintenance summaries should list");
|
||||
assert_eq!(maintenance.len(), 1);
|
||||
assert!(!maintenance[0].is_active);
|
||||
assert_eq!(
|
||||
maintenance[0].upstream_metadata,
|
||||
Some(json!({"models":["gpt-4.1-mini"]}))
|
||||
);
|
||||
|
||||
assert!(repository
|
||||
.delete_key(&key_id)
|
||||
.await
|
||||
|
||||
@@ -94,6 +94,13 @@ ON DUPLICATE KEY UPDATE
|
||||
updated_at = VALUES(updated_at)
|
||||
"#;
|
||||
|
||||
const ENQUEUE_PROVIDER_MONTHLY_USAGE_DELTA_SQL: &str = r#"
|
||||
INSERT INTO usage_counter_deltas (
|
||||
id, request_id, kind, target_id, total_cost_usd_delta, created_at
|
||||
)
|
||||
VALUES (?, ?, 'provider_monthly', ?, ?, ?)
|
||||
"#;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct MysqlSettlementRepository {
|
||||
pool: MysqlPool,
|
||||
@@ -138,6 +145,36 @@ fn now_unix_secs() -> Result<i64, DataLayerError> {
|
||||
.map_err(|_| DataLayerError::InvalidInput("timestamp overflow".to_string()))
|
||||
}
|
||||
|
||||
async fn enqueue_provider_monthly_usage_delta_mysql(
|
||||
tx: &mut sqlx::Transaction<'_, sqlx::MySql>,
|
||||
request_id: &str,
|
||||
provider_id: &str,
|
||||
total_cost_usd_delta: f64,
|
||||
created_at: i64,
|
||||
) -> Result<(), DataLayerError> {
|
||||
let request_id = request_id.trim();
|
||||
let provider_id = provider_id.trim();
|
||||
if request_id.is_empty() || provider_id.is_empty() || total_cost_usd_delta == 0.0 {
|
||||
return Ok(());
|
||||
}
|
||||
if !total_cost_usd_delta.is_finite() {
|
||||
return Err(DataLayerError::UnexpectedValue(format!(
|
||||
"provider monthly usage delta is not finite for {provider_id}"
|
||||
)));
|
||||
}
|
||||
|
||||
sqlx::query(ENQUEUE_PROVIDER_MONTHLY_USAGE_DELTA_SQL)
|
||||
.bind(uuid::Uuid::new_v4().to_string())
|
||||
.bind(request_id)
|
||||
.bind(provider_id)
|
||||
.bind(total_cost_usd_delta)
|
||||
.bind(created_at)
|
||||
.execute(&mut **tx)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
struct DailyQuotaDebitResult {
|
||||
debited_usd: f64,
|
||||
@@ -622,30 +659,14 @@ WHERE id = ?
|
||||
.as_deref()
|
||||
.filter(|value| !value.is_empty())
|
||||
{
|
||||
sqlx::query(
|
||||
r#"
|
||||
UPDATE providers
|
||||
SET
|
||||
monthly_used_usd = COALESCE(monthly_used_usd, 0) + ?,
|
||||
updated_at = ?
|
||||
WHERE id = ?
|
||||
"#,
|
||||
enqueue_provider_monthly_usage_delta_mysql(
|
||||
&mut tx,
|
||||
&input.request_id,
|
||||
provider_id,
|
||||
input.actual_total_cost_usd,
|
||||
updated_at,
|
||||
)
|
||||
.bind(input.actual_total_cost_usd)
|
||||
.bind(updated_at)
|
||||
.bind(provider_id)
|
||||
.execute(&mut *tx)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
|
||||
settlement.provider_monthly_used_usd = sqlx::query_scalar::<_, Option<f64>>(
|
||||
"SELECT monthly_used_usd FROM providers WHERE id = ? LIMIT 1",
|
||||
)
|
||||
.bind(provider_id)
|
||||
.fetch_optional(&mut *tx)
|
||||
.await
|
||||
.map_sql_err()?
|
||||
.flatten();
|
||||
.await?;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -683,6 +704,10 @@ WHERE id = ?
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::MysqlSettlementRepository;
|
||||
use crate::run_migrations;
|
||||
use aether_data_contracts::repository::settlement::{
|
||||
SettlementWriteRepository, UsageSettlementInput,
|
||||
};
|
||||
|
||||
#[tokio::test]
|
||||
async fn repository_builds_from_lazy_pool() {
|
||||
@@ -694,4 +719,132 @@ mod tests {
|
||||
|
||||
let _repository = MysqlSettlementRepository::new(pool);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn mysql_repository_settles_once_and_enqueues_provider_delta_when_url_is_set() {
|
||||
let Some(database_url) = std::env::var("AETHER_TEST_MYSQL_URL")
|
||||
.ok()
|
||||
.filter(|value| !value.trim().is_empty())
|
||||
else {
|
||||
eprintln!(
|
||||
"skipping mysql settlement parity test because AETHER_TEST_MYSQL_URL is unset"
|
||||
);
|
||||
return;
|
||||
};
|
||||
let pool = sqlx::mysql::MySqlPoolOptions::new()
|
||||
.max_connections(1)
|
||||
.connect(&database_url)
|
||||
.await
|
||||
.expect("mysql pool should connect");
|
||||
run_migrations(&pool)
|
||||
.await
|
||||
.expect("mysql migrations should run");
|
||||
cleanup_settlement_rows(&pool).await;
|
||||
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO providers (id, name, provider_type, monthly_used_usd, created_at, updated_at)
|
||||
VALUES ('settlement-provider-1', 'Settlement Provider', 'openai', 5.0, 1, 1)
|
||||
"#,
|
||||
)
|
||||
.execute(&pool)
|
||||
.await
|
||||
.expect("provider should seed");
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO wallets (id, user_id, balance, gift_balance, limit_mode, created_at, updated_at)
|
||||
VALUES ('settlement-wallet-1', 'settlement-user-1', 10.0, 2.0, 'finite', 1, 1)
|
||||
"#,
|
||||
)
|
||||
.execute(&pool)
|
||||
.await
|
||||
.expect("wallet should seed");
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO `usage` (
|
||||
request_id, user_id, provider_id, status, billing_status,
|
||||
total_cost_usd, actual_total_cost_usd
|
||||
)
|
||||
VALUES (
|
||||
'settlement-request-1', 'settlement-user-1', 'settlement-provider-1',
|
||||
'completed', 'pending', 3.0, 6.0
|
||||
)
|
||||
"#,
|
||||
)
|
||||
.execute(&pool)
|
||||
.await
|
||||
.expect("usage should seed");
|
||||
|
||||
let repository = MysqlSettlementRepository::new(pool.clone());
|
||||
let input = UsageSettlementInput {
|
||||
request_id: "settlement-request-1".to_string(),
|
||||
user_id: Some("settlement-user-1".to_string()),
|
||||
api_key_id: None,
|
||||
api_key_is_standalone: false,
|
||||
provider_id: Some("settlement-provider-1".to_string()),
|
||||
status: "completed".to_string(),
|
||||
billing_status: "pending".to_string(),
|
||||
total_cost_usd: 3.0,
|
||||
actual_total_cost_usd: 6.0,
|
||||
finalized_at_unix_secs: Some(1_234),
|
||||
};
|
||||
let first = repository
|
||||
.settle_usage(input.clone())
|
||||
.await
|
||||
.expect("settlement should run")
|
||||
.expect("usage should exist");
|
||||
let second = repository
|
||||
.settle_usage(input)
|
||||
.await
|
||||
.expect("second settlement should run")
|
||||
.expect("usage should exist");
|
||||
assert_eq!(first.billing_status, "settled");
|
||||
assert_eq!(first.provider_monthly_used_usd, None);
|
||||
assert_eq!(second.finalized_at_unix_secs, Some(1_234));
|
||||
|
||||
let wallet: (f64, f64, f64) = sqlx::query_as(
|
||||
"SELECT balance, gift_balance, total_consumed FROM wallets WHERE id = 'settlement-wallet-1'",
|
||||
)
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.expect("wallet should load");
|
||||
assert_eq!(wallet, (4.0, 2.0, 6.0));
|
||||
let provider_used: f64 = sqlx::query_scalar(
|
||||
"SELECT monthly_used_usd FROM providers WHERE id = 'settlement-provider-1'",
|
||||
)
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.expect("provider should load");
|
||||
assert_eq!(provider_used, 5.0);
|
||||
let provider_delta: (i64, f64) = sqlx::query_as(
|
||||
r#"
|
||||
SELECT CAST(COUNT(*) AS SIGNED), COALESCE(SUM(total_cost_usd_delta), 0)
|
||||
FROM usage_counter_deltas
|
||||
WHERE request_id = 'settlement-request-1'
|
||||
AND kind = 'provider_monthly'
|
||||
AND target_id = 'settlement-provider-1'
|
||||
"#,
|
||||
)
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.expect("provider delta should load");
|
||||
assert_eq!(provider_delta, (1, 6.0));
|
||||
|
||||
cleanup_settlement_rows(&pool).await;
|
||||
}
|
||||
|
||||
async fn cleanup_settlement_rows(pool: &sqlx::MySqlPool) {
|
||||
for sql in [
|
||||
"DELETE FROM usage_counter_deltas WHERE request_id = 'settlement-request-1'",
|
||||
"DELETE FROM usage_settlement_snapshots WHERE request_id = 'settlement-request-1'",
|
||||
"DELETE FROM `usage` WHERE request_id = 'settlement-request-1'",
|
||||
"DELETE FROM wallets WHERE id = 'settlement-wallet-1'",
|
||||
"DELETE FROM providers WHERE id = 'settlement-provider-1'",
|
||||
] {
|
||||
sqlx::query(sql)
|
||||
.execute(pool)
|
||||
.await
|
||||
.expect("settlement cleanup should succeed");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,830 @@
|
||||
use std::io::{Read, Write};
|
||||
|
||||
use aether_data_contracts::repository::usage::{
|
||||
parse_usage_body_ref, usage_body_ref, StoredRequestUsageAudit, UpsertUsageRecord,
|
||||
UsageBodyCaptureState, UsageBodyField,
|
||||
};
|
||||
use flate2::{read::GzDecoder, write::GzEncoder, Compression};
|
||||
use serde_json::{Map, Value};
|
||||
use sqlx::{mysql::MySqlRow, Row};
|
||||
|
||||
use crate::error::SqlResultExt;
|
||||
use crate::{DataLayerError, MysqlPool};
|
||||
|
||||
#[derive(Debug)]
|
||||
pub(crate) struct PreparedUsageHttpCapture {
|
||||
request_headers: Option<String>,
|
||||
provider_request_headers: Option<String>,
|
||||
response_headers: Option<String>,
|
||||
client_response_headers: Option<String>,
|
||||
request_body: PreparedBody,
|
||||
provider_request_body: PreparedBody,
|
||||
response_body: PreparedBody,
|
||||
client_response_body: PreparedBody,
|
||||
refs: HttpAuditRefs,
|
||||
states: HttpAuditStates,
|
||||
capture_mode: &'static str,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct PreparedBody {
|
||||
field: UsageBodyField,
|
||||
payload_gzip: Option<Vec<u8>>,
|
||||
clear_existing: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
struct HttpAuditRefs {
|
||||
request_body_ref: Option<String>,
|
||||
provider_request_body_ref: Option<String>,
|
||||
response_body_ref: Option<String>,
|
||||
client_response_body_ref: Option<String>,
|
||||
}
|
||||
|
||||
impl HttpAuditRefs {
|
||||
fn any_present(&self) -> bool {
|
||||
self.request_body_ref.is_some()
|
||||
|| self.provider_request_body_ref.is_some()
|
||||
|| self.response_body_ref.is_some()
|
||||
|| self.client_response_body_ref.is_some()
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
struct HttpAuditStates {
|
||||
request_body_state: Option<UsageBodyCaptureState>,
|
||||
provider_request_body_state: Option<UsageBodyCaptureState>,
|
||||
response_body_state: Option<UsageBodyCaptureState>,
|
||||
client_response_body_state: Option<UsageBodyCaptureState>,
|
||||
}
|
||||
|
||||
impl HttpAuditStates {
|
||||
fn any_present(&self) -> bool {
|
||||
self.request_body_state.is_some()
|
||||
|| self.provider_request_body_state.is_some()
|
||||
|| self.response_body_state.is_some()
|
||||
|| self.client_response_body_state.is_some()
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn capture_update_allowed(
|
||||
previous: Option<&StoredRequestUsageAudit>,
|
||||
incoming_status: &str,
|
||||
) -> bool {
|
||||
let Some(previous) = previous else {
|
||||
return true;
|
||||
};
|
||||
if previous.billing_status != "pending" {
|
||||
return false;
|
||||
}
|
||||
if matches!(
|
||||
previous.status.as_str(),
|
||||
"completed" | "failed" | "cancelled"
|
||||
) && matches!(incoming_status, "pending" | "streaming")
|
||||
{
|
||||
return false;
|
||||
}
|
||||
!(previous.status == "streaming" && incoming_status == "pending")
|
||||
}
|
||||
|
||||
pub(crate) fn apply_previous_metadata_tombstones(
|
||||
usage: &mut UpsertUsageRecord,
|
||||
previous: Option<&StoredRequestUsageAudit>,
|
||||
) {
|
||||
if usage.request_metadata.is_some() {
|
||||
return;
|
||||
}
|
||||
let clear_request = usage.request_body_state == Some(UsageBodyCaptureState::None);
|
||||
let clear_provider_request =
|
||||
usage.provider_request_body_state == Some(UsageBodyCaptureState::None);
|
||||
if !clear_request && !clear_provider_request {
|
||||
return;
|
||||
}
|
||||
let mut metadata = previous
|
||||
.and_then(|previous| previous.request_metadata.as_ref())
|
||||
.and_then(Value::as_object)
|
||||
.cloned()
|
||||
.unwrap_or_default();
|
||||
if clear_request {
|
||||
metadata.remove("requested_reasoning_effort");
|
||||
metadata.remove("request_body_ref");
|
||||
}
|
||||
if clear_provider_request {
|
||||
metadata.remove("provider_reasoning_effort");
|
||||
metadata.remove("provider_service_tier");
|
||||
metadata.remove("provider_cache_ttl_minutes");
|
||||
metadata.remove("provider_request_body_ref");
|
||||
}
|
||||
usage.request_metadata = Some(Value::Object(metadata));
|
||||
}
|
||||
|
||||
pub(crate) fn prepare_usage_http_capture(
|
||||
usage: &mut UpsertUsageRecord,
|
||||
) -> Result<PreparedUsageHttpCapture, DataLayerError> {
|
||||
let clear_request = usage.request_body_state == Some(UsageBodyCaptureState::None);
|
||||
let clear_provider_request =
|
||||
usage.provider_request_body_state == Some(UsageBodyCaptureState::None);
|
||||
let clear_response = usage.response_body_state == Some(UsageBodyCaptureState::None);
|
||||
let clear_client_response =
|
||||
usage.client_response_body_state == Some(UsageBodyCaptureState::None);
|
||||
|
||||
let request_body_value = (!clear_request)
|
||||
.then_some(usage.request_body.as_ref())
|
||||
.flatten();
|
||||
let provider_request_body_value = (!clear_provider_request)
|
||||
.then_some(usage.provider_request_body.as_ref())
|
||||
.flatten();
|
||||
let response_body_value = (!clear_response)
|
||||
.then_some(usage.response_body.as_ref())
|
||||
.flatten();
|
||||
let client_response_body_value = (!clear_client_response)
|
||||
.then_some(usage.client_response_body.as_ref())
|
||||
.flatten();
|
||||
|
||||
let request_body = prepare_body(
|
||||
UsageBodyField::RequestBody,
|
||||
request_body_value,
|
||||
clear_request,
|
||||
)?;
|
||||
let provider_request_body = prepare_body(
|
||||
UsageBodyField::ProviderRequestBody,
|
||||
provider_request_body_value,
|
||||
clear_provider_request,
|
||||
)?;
|
||||
let response_body = prepare_body(
|
||||
UsageBodyField::ResponseBody,
|
||||
response_body_value,
|
||||
clear_response,
|
||||
)?;
|
||||
let client_response_body = prepare_body(
|
||||
UsageBodyField::ClientResponseBody,
|
||||
client_response_body_value,
|
||||
clear_client_response,
|
||||
)?;
|
||||
|
||||
let refs = HttpAuditRefs {
|
||||
request_body_ref: resolved_write_ref(
|
||||
(!clear_request)
|
||||
.then_some(usage.request_body_ref.as_deref())
|
||||
.flatten(),
|
||||
&usage.request_id,
|
||||
UsageBodyField::RequestBody,
|
||||
request_body.payload_gzip.is_some(),
|
||||
),
|
||||
provider_request_body_ref: resolved_write_ref(
|
||||
(!clear_provider_request)
|
||||
.then_some(usage.provider_request_body_ref.as_deref())
|
||||
.flatten(),
|
||||
&usage.request_id,
|
||||
UsageBodyField::ProviderRequestBody,
|
||||
provider_request_body.payload_gzip.is_some(),
|
||||
),
|
||||
response_body_ref: resolved_write_ref(
|
||||
(!clear_response)
|
||||
.then_some(usage.response_body_ref.as_deref())
|
||||
.flatten(),
|
||||
&usage.request_id,
|
||||
UsageBodyField::ResponseBody,
|
||||
response_body.payload_gzip.is_some(),
|
||||
),
|
||||
client_response_body_ref: resolved_write_ref(
|
||||
(!clear_client_response)
|
||||
.then_some(usage.client_response_body_ref.as_deref())
|
||||
.flatten(),
|
||||
&usage.request_id,
|
||||
UsageBodyField::ClientResponseBody,
|
||||
client_response_body.payload_gzip.is_some(),
|
||||
),
|
||||
};
|
||||
let states = HttpAuditStates {
|
||||
request_body_state: state_for_storage(
|
||||
usage.request_body_state,
|
||||
&request_body,
|
||||
refs.request_body_ref.as_deref(),
|
||||
),
|
||||
provider_request_body_state: state_for_storage(
|
||||
usage.provider_request_body_state,
|
||||
&provider_request_body,
|
||||
refs.provider_request_body_ref.as_deref(),
|
||||
),
|
||||
response_body_state: state_for_storage(
|
||||
usage.response_body_state,
|
||||
&response_body,
|
||||
refs.response_body_ref.as_deref(),
|
||||
),
|
||||
client_response_body_state: state_for_storage(
|
||||
usage.client_response_body_state,
|
||||
&client_response_body,
|
||||
refs.client_response_body_ref.as_deref(),
|
||||
),
|
||||
};
|
||||
|
||||
usage.request_metadata = prepare_metadata_for_body_storage(
|
||||
usage.request_metadata.take(),
|
||||
[
|
||||
(
|
||||
UsageBodyField::RequestBody,
|
||||
request_body_value.is_some(),
|
||||
usage.request_body_ref.as_deref(),
|
||||
),
|
||||
(
|
||||
UsageBodyField::ProviderRequestBody,
|
||||
provider_request_body_value.is_some(),
|
||||
usage.provider_request_body_ref.as_deref(),
|
||||
),
|
||||
(
|
||||
UsageBodyField::ResponseBody,
|
||||
response_body_value.is_some(),
|
||||
usage.response_body_ref.as_deref(),
|
||||
),
|
||||
(
|
||||
UsageBodyField::ClientResponseBody,
|
||||
client_response_body_value.is_some(),
|
||||
usage.client_response_body_ref.as_deref(),
|
||||
),
|
||||
],
|
||||
);
|
||||
|
||||
let capture_mode = if refs.any_present() {
|
||||
"ref_backed"
|
||||
} else if [
|
||||
request_body_value,
|
||||
provider_request_body_value,
|
||||
response_body_value,
|
||||
client_response_body_value,
|
||||
]
|
||||
.iter()
|
||||
.any(Option::is_some)
|
||||
{
|
||||
"inline_legacy"
|
||||
} else {
|
||||
"none"
|
||||
};
|
||||
|
||||
Ok(PreparedUsageHttpCapture {
|
||||
request_headers: json_text(usage.request_headers.as_ref())?,
|
||||
provider_request_headers: json_text(usage.provider_request_headers.as_ref())?,
|
||||
response_headers: json_text(usage.response_headers.as_ref())?,
|
||||
client_response_headers: json_text(usage.client_response_headers.as_ref())?,
|
||||
request_body,
|
||||
provider_request_body,
|
||||
response_body,
|
||||
client_response_body,
|
||||
refs,
|
||||
states,
|
||||
capture_mode,
|
||||
})
|
||||
}
|
||||
|
||||
fn prepare_body(
|
||||
field: UsageBodyField,
|
||||
value: Option<&Value>,
|
||||
clear_existing: bool,
|
||||
) -> Result<PreparedBody, DataLayerError> {
|
||||
Ok(PreparedBody {
|
||||
field,
|
||||
payload_gzip: value.map(compress_json).transpose()?,
|
||||
clear_existing,
|
||||
})
|
||||
}
|
||||
|
||||
fn compress_json(value: &Value) -> Result<Vec<u8>, DataLayerError> {
|
||||
let bytes = serde_json::to_vec(value).map_err(|err| {
|
||||
DataLayerError::UnexpectedValue(format!("failed to serialize usage body: {err}"))
|
||||
})?;
|
||||
let mut encoder = GzEncoder::new(Vec::new(), Compression::new(6));
|
||||
encoder.write_all(&bytes).map_err(|err| {
|
||||
DataLayerError::UnexpectedValue(format!("failed to gzip usage body: {err}"))
|
||||
})?;
|
||||
encoder.finish().map_err(|err| {
|
||||
DataLayerError::UnexpectedValue(format!("failed to finish usage body gzip: {err}"))
|
||||
})
|
||||
}
|
||||
|
||||
fn resolved_write_ref(
|
||||
explicit_ref: Option<&str>,
|
||||
request_id: &str,
|
||||
field: UsageBodyField,
|
||||
has_blob: bool,
|
||||
) -> Option<String> {
|
||||
explicit_ref
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned)
|
||||
.or_else(|| has_blob.then(|| usage_body_ref(request_id, field)))
|
||||
}
|
||||
|
||||
fn state_for_storage(
|
||||
incoming: Option<UsageBodyCaptureState>,
|
||||
body: &PreparedBody,
|
||||
body_ref: Option<&str>,
|
||||
) -> Option<UsageBodyCaptureState> {
|
||||
if matches!(
|
||||
incoming,
|
||||
Some(
|
||||
UsageBodyCaptureState::Disabled
|
||||
| UsageBodyCaptureState::Unavailable
|
||||
| UsageBodyCaptureState::None
|
||||
)
|
||||
) {
|
||||
return incoming;
|
||||
}
|
||||
if body.payload_gzip.is_some() || body_ref.is_some() {
|
||||
return Some(UsageBodyCaptureState::Reference);
|
||||
}
|
||||
incoming
|
||||
}
|
||||
|
||||
fn json_text(value: Option<&Value>) -> Result<Option<String>, DataLayerError> {
|
||||
value
|
||||
.map(|value| {
|
||||
serde_json::to_string(value).map_err(|err| {
|
||||
DataLayerError::UnexpectedValue(format!("failed to serialize usage JSON: {err}"))
|
||||
})
|
||||
})
|
||||
.transpose()
|
||||
}
|
||||
|
||||
fn prepare_metadata_for_body_storage(
|
||||
metadata: Option<Value>,
|
||||
body_fields: [(UsageBodyField, bool, Option<&str>); 4],
|
||||
) -> Option<Value> {
|
||||
let mut object = match metadata {
|
||||
Some(Value::Object(object)) => object,
|
||||
Some(value) => {
|
||||
let mut object = Map::new();
|
||||
object.insert("request_metadata".to_string(), value);
|
||||
object
|
||||
}
|
||||
None => Map::new(),
|
||||
};
|
||||
let should_replace = !object.is_empty()
|
||||
|| body_fields
|
||||
.iter()
|
||||
.any(|(_, has_value, explicit_ref)| *has_value || explicit_ref.is_some());
|
||||
if !should_replace {
|
||||
return None;
|
||||
}
|
||||
for (field, has_value, explicit_ref) in body_fields {
|
||||
if has_value || explicit_ref.is_some() {
|
||||
object.remove(field.as_ref_key());
|
||||
}
|
||||
}
|
||||
Some(Value::Object(object))
|
||||
}
|
||||
|
||||
pub(crate) async fn sync_usage_http_capture(
|
||||
tx: &mut sqlx::Transaction<'_, sqlx::MySql>,
|
||||
request_id: &str,
|
||||
prepared: &PreparedUsageHttpCapture,
|
||||
) -> Result<(), DataLayerError> {
|
||||
for body in [
|
||||
&prepared.request_body,
|
||||
&prepared.provider_request_body,
|
||||
&prepared.response_body,
|
||||
&prepared.client_response_body,
|
||||
] {
|
||||
sync_body(tx, request_id, body).await?;
|
||||
}
|
||||
let headers_present = prepared.request_headers.is_some()
|
||||
|| prepared.provider_request_headers.is_some()
|
||||
|| prepared.response_headers.is_some()
|
||||
|| prepared.client_response_headers.is_some();
|
||||
if !headers_present
|
||||
&& !prepared.refs.any_present()
|
||||
&& !prepared.states.any_present()
|
||||
&& prepared.capture_mode == "none"
|
||||
{
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO usage_http_audits (
|
||||
request_id,
|
||||
request_headers,
|
||||
provider_request_headers,
|
||||
response_headers,
|
||||
client_response_headers,
|
||||
request_body_ref,
|
||||
provider_request_body_ref,
|
||||
response_body_ref,
|
||||
client_response_body_ref,
|
||||
request_body_state,
|
||||
provider_request_body_state,
|
||||
response_body_state,
|
||||
client_response_body_state,
|
||||
body_capture_mode
|
||||
)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
ON DUPLICATE KEY UPDATE
|
||||
request_headers = COALESCE(VALUES(request_headers), request_headers),
|
||||
provider_request_headers = COALESCE(
|
||||
VALUES(provider_request_headers),
|
||||
provider_request_headers
|
||||
),
|
||||
response_headers = COALESCE(VALUES(response_headers), response_headers),
|
||||
client_response_headers = COALESCE(
|
||||
VALUES(client_response_headers),
|
||||
client_response_headers
|
||||
),
|
||||
request_body_ref = CASE
|
||||
WHEN VALUES(request_body_state) = 'none' THEN NULL
|
||||
ELSE COALESCE(VALUES(request_body_ref), request_body_ref)
|
||||
END,
|
||||
provider_request_body_ref = CASE
|
||||
WHEN VALUES(provider_request_body_state) = 'none' THEN NULL
|
||||
ELSE COALESCE(VALUES(provider_request_body_ref), provider_request_body_ref)
|
||||
END,
|
||||
response_body_ref = CASE
|
||||
WHEN VALUES(response_body_state) = 'none' THEN NULL
|
||||
ELSE COALESCE(VALUES(response_body_ref), response_body_ref)
|
||||
END,
|
||||
client_response_body_ref = CASE
|
||||
WHEN VALUES(client_response_body_state) = 'none' THEN NULL
|
||||
ELSE COALESCE(VALUES(client_response_body_ref), client_response_body_ref)
|
||||
END,
|
||||
request_body_state = COALESCE(VALUES(request_body_state), request_body_state),
|
||||
provider_request_body_state = COALESCE(
|
||||
VALUES(provider_request_body_state),
|
||||
provider_request_body_state
|
||||
),
|
||||
response_body_state = COALESCE(VALUES(response_body_state), response_body_state),
|
||||
client_response_body_state = COALESCE(
|
||||
VALUES(client_response_body_state),
|
||||
client_response_body_state
|
||||
),
|
||||
body_capture_mode = CASE
|
||||
WHEN VALUES(body_capture_mode) = 'none'
|
||||
AND (
|
||||
VALUES(request_body_state) = 'none'
|
||||
OR VALUES(provider_request_body_state) = 'none'
|
||||
OR VALUES(response_body_state) = 'none'
|
||||
OR VALUES(client_response_body_state) = 'none'
|
||||
)
|
||||
THEN 'none'
|
||||
ELSE COALESCE(
|
||||
NULLIF(VALUES(body_capture_mode), 'none'),
|
||||
body_capture_mode,
|
||||
'none'
|
||||
)
|
||||
END,
|
||||
updated_at = UNIX_TIMESTAMP()
|
||||
"#,
|
||||
)
|
||||
.bind(request_id)
|
||||
.bind(&prepared.request_headers)
|
||||
.bind(&prepared.provider_request_headers)
|
||||
.bind(&prepared.response_headers)
|
||||
.bind(&prepared.client_response_headers)
|
||||
.bind(prepared.refs.request_body_ref.as_deref())
|
||||
.bind(prepared.refs.provider_request_body_ref.as_deref())
|
||||
.bind(prepared.refs.response_body_ref.as_deref())
|
||||
.bind(prepared.refs.client_response_body_ref.as_deref())
|
||||
.bind(
|
||||
prepared
|
||||
.states
|
||||
.request_body_state
|
||||
.map(UsageBodyCaptureState::as_str),
|
||||
)
|
||||
.bind(
|
||||
prepared
|
||||
.states
|
||||
.provider_request_body_state
|
||||
.map(UsageBodyCaptureState::as_str),
|
||||
)
|
||||
.bind(
|
||||
prepared
|
||||
.states
|
||||
.response_body_state
|
||||
.map(UsageBodyCaptureState::as_str),
|
||||
)
|
||||
.bind(
|
||||
prepared
|
||||
.states
|
||||
.client_response_body_state
|
||||
.map(UsageBodyCaptureState::as_str),
|
||||
)
|
||||
.bind(prepared.capture_mode)
|
||||
.execute(&mut **tx)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn sync_body(
|
||||
tx: &mut sqlx::Transaction<'_, sqlx::MySql>,
|
||||
request_id: &str,
|
||||
body: &PreparedBody,
|
||||
) -> Result<(), DataLayerError> {
|
||||
let body_ref = usage_body_ref(request_id, body.field);
|
||||
if body.clear_existing || body.payload_gzip.is_some() {
|
||||
sqlx::query(clear_legacy_body_sql(body.field))
|
||||
.bind(request_id)
|
||||
.execute(&mut **tx)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
}
|
||||
if body.clear_existing {
|
||||
sqlx::query("DELETE FROM usage_body_blobs WHERE body_ref = ?")
|
||||
.bind(body_ref)
|
||||
.execute(&mut **tx)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
return Ok(());
|
||||
}
|
||||
if let Some(payload_gzip) = body.payload_gzip.as_deref() {
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO usage_body_blobs (body_ref, request_id, body_field, payload_gzip)
|
||||
VALUES (?, ?, ?, ?)
|
||||
ON DUPLICATE KEY UPDATE
|
||||
request_id = VALUES(request_id),
|
||||
body_field = VALUES(body_field),
|
||||
payload_gzip = VALUES(payload_gzip),
|
||||
updated_at = UNIX_TIMESTAMP()
|
||||
"#,
|
||||
)
|
||||
.bind(body_ref)
|
||||
.bind(request_id)
|
||||
.bind(body.field.as_storage_field())
|
||||
.bind(payload_gzip)
|
||||
.execute(&mut **tx)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn clear_legacy_body_sql(field: UsageBodyField) -> &'static str {
|
||||
match field {
|
||||
UsageBodyField::RequestBody => {
|
||||
"UPDATE `usage` SET request_body = NULL, request_body_compressed = NULL WHERE request_id = ?"
|
||||
}
|
||||
UsageBodyField::ProviderRequestBody => {
|
||||
"UPDATE `usage` SET provider_request_body = NULL, provider_request_body_compressed = NULL WHERE request_id = ?"
|
||||
}
|
||||
UsageBodyField::ResponseBody => {
|
||||
"UPDATE `usage` SET response_body = NULL, response_body_compressed = NULL WHERE request_id = ?"
|
||||
}
|
||||
UsageBodyField::ClientResponseBody => {
|
||||
"UPDATE `usage` SET client_response_body = NULL, client_response_body_compressed = NULL WHERE request_id = ?"
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn hydrate_usage_row(
|
||||
row: &MySqlRow,
|
||||
usage: &mut StoredRequestUsageAudit,
|
||||
resolve_legacy_compressed: bool,
|
||||
) -> Result<(), DataLayerError> {
|
||||
usage.request_headers = optional_json(row, "request_headers")?;
|
||||
usage.provider_request_headers = optional_json(row, "provider_request_headers")?;
|
||||
usage.response_headers = optional_json(row, "response_headers")?;
|
||||
usage.client_response_headers = optional_json(row, "client_response_headers")?;
|
||||
|
||||
let request_body = legacy_body_column(
|
||||
row,
|
||||
"request_body",
|
||||
"request_body_compressed",
|
||||
resolve_legacy_compressed,
|
||||
)?;
|
||||
let provider_request_body = legacy_body_column(
|
||||
row,
|
||||
"provider_request_body",
|
||||
"provider_request_body_compressed",
|
||||
resolve_legacy_compressed,
|
||||
)?;
|
||||
let response_body = legacy_body_column(
|
||||
row,
|
||||
"response_body",
|
||||
"response_body_compressed",
|
||||
resolve_legacy_compressed,
|
||||
)?;
|
||||
let client_response_body = legacy_body_column(
|
||||
row,
|
||||
"client_response_body",
|
||||
"client_response_body_compressed",
|
||||
resolve_legacy_compressed,
|
||||
)?;
|
||||
usage.request_body = request_body.0;
|
||||
usage.provider_request_body = provider_request_body.0;
|
||||
usage.response_body = response_body.0;
|
||||
usage.client_response_body = client_response_body.0;
|
||||
|
||||
let metadata = usage.request_metadata.as_ref().and_then(Value::as_object);
|
||||
usage.request_body_ref = resolved_read_ref(
|
||||
row.try_get("http_request_body_ref").map_sql_err()?,
|
||||
metadata,
|
||||
&usage.request_id,
|
||||
UsageBodyField::RequestBody,
|
||||
request_body.1,
|
||||
);
|
||||
usage.provider_request_body_ref = resolved_read_ref(
|
||||
row.try_get("http_provider_request_body_ref")
|
||||
.map_sql_err()?,
|
||||
metadata,
|
||||
&usage.request_id,
|
||||
UsageBodyField::ProviderRequestBody,
|
||||
provider_request_body.1,
|
||||
);
|
||||
usage.response_body_ref = resolved_read_ref(
|
||||
row.try_get("http_response_body_ref").map_sql_err()?,
|
||||
metadata,
|
||||
&usage.request_id,
|
||||
UsageBodyField::ResponseBody,
|
||||
response_body.1,
|
||||
);
|
||||
usage.client_response_body_ref = resolved_read_ref(
|
||||
row.try_get("http_client_response_body_ref").map_sql_err()?,
|
||||
metadata,
|
||||
&usage.request_id,
|
||||
UsageBodyField::ClientResponseBody,
|
||||
client_response_body.1,
|
||||
);
|
||||
usage.request_body_state = optional_state(row, "http_request_body_state")?;
|
||||
usage.provider_request_body_state = optional_state(row, "http_provider_request_body_state")?;
|
||||
usage.response_body_state = optional_state(row, "http_response_body_state")?;
|
||||
usage.client_response_body_state = optional_state(row, "http_client_response_body_state")?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn optional_json(row: &MySqlRow, column: &str) -> Result<Option<Value>, DataLayerError> {
|
||||
row.try_get::<Option<String>, _>(column)
|
||||
.map_sql_err()?
|
||||
.map(|raw| {
|
||||
serde_json::from_str(&raw).map_err(|err| {
|
||||
DataLayerError::UnexpectedValue(format!("invalid usage JSON in {column}: {err}"))
|
||||
})
|
||||
})
|
||||
.transpose()
|
||||
}
|
||||
|
||||
fn legacy_body_column(
|
||||
row: &MySqlRow,
|
||||
inline_column: &str,
|
||||
compressed_column: &str,
|
||||
resolve_compressed: bool,
|
||||
) -> Result<(Option<Value>, bool), DataLayerError> {
|
||||
let inline = optional_json(row, inline_column)?;
|
||||
if inline.is_some() {
|
||||
return Ok((inline, false));
|
||||
}
|
||||
let compressed = row
|
||||
.try_get::<Option<Vec<u8>>, _>(compressed_column)
|
||||
.map_sql_err()?;
|
||||
let has_compressed = compressed.is_some();
|
||||
let value = if resolve_compressed {
|
||||
compressed.map(|bytes| inflate_json(&bytes)).transpose()?
|
||||
} else {
|
||||
None
|
||||
};
|
||||
Ok((value, has_compressed))
|
||||
}
|
||||
|
||||
fn resolved_read_ref(
|
||||
audit_ref: Option<String>,
|
||||
metadata: Option<&Map<String, Value>>,
|
||||
request_id: &str,
|
||||
field: UsageBodyField,
|
||||
has_compressed: bool,
|
||||
) -> Option<String> {
|
||||
audit_ref
|
||||
.map(|value| value.trim().to_string())
|
||||
.filter(|value| !value.is_empty())
|
||||
.or_else(|| has_compressed.then(|| usage_body_ref(request_id, field)))
|
||||
.or_else(|| metadata_body_ref(metadata, request_id, field))
|
||||
}
|
||||
|
||||
fn metadata_body_ref(
|
||||
metadata: Option<&Map<String, Value>>,
|
||||
request_id: &str,
|
||||
field: UsageBodyField,
|
||||
) -> Option<String> {
|
||||
metadata
|
||||
.and_then(|metadata| metadata.get(field.as_ref_key()))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.and_then(parse_usage_body_ref)
|
||||
.filter(|(parsed_request_id, parsed_field)| {
|
||||
parsed_request_id == request_id && *parsed_field == field
|
||||
})
|
||||
.map(|(parsed_request_id, parsed_field)| usage_body_ref(&parsed_request_id, parsed_field))
|
||||
}
|
||||
|
||||
fn optional_state(
|
||||
row: &MySqlRow,
|
||||
column: &str,
|
||||
) -> Result<Option<UsageBodyCaptureState>, DataLayerError> {
|
||||
Ok(row
|
||||
.try_get::<Option<String>, _>(column)
|
||||
.map_sql_err()?
|
||||
.as_deref()
|
||||
.and_then(parse_state))
|
||||
}
|
||||
|
||||
fn parse_state(value: &str) -> Option<UsageBodyCaptureState> {
|
||||
match value.trim() {
|
||||
"none" => Some(UsageBodyCaptureState::None),
|
||||
"inline" => Some(UsageBodyCaptureState::Inline),
|
||||
"reference" => Some(UsageBodyCaptureState::Reference),
|
||||
"truncated" => Some(UsageBodyCaptureState::Truncated),
|
||||
"disabled" => Some(UsageBodyCaptureState::Disabled),
|
||||
"unavailable" => Some(UsageBodyCaptureState::Unavailable),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn hydrate_usage_body_refs(
|
||||
pool: &MysqlPool,
|
||||
mut usage: StoredRequestUsageAudit,
|
||||
) -> Result<StoredRequestUsageAudit, DataLayerError> {
|
||||
for field in [
|
||||
UsageBodyField::RequestBody,
|
||||
UsageBodyField::ProviderRequestBody,
|
||||
UsageBodyField::ResponseBody,
|
||||
UsageBodyField::ClientResponseBody,
|
||||
] {
|
||||
if usage.body_value(field).is_some() {
|
||||
continue;
|
||||
}
|
||||
let Some(body_ref) = usage.body_ref(field) else {
|
||||
continue;
|
||||
};
|
||||
let value = resolve_body_ref(pool, body_ref).await?;
|
||||
match field {
|
||||
UsageBodyField::RequestBody => usage.request_body = value,
|
||||
UsageBodyField::ProviderRequestBody => usage.provider_request_body = value,
|
||||
UsageBodyField::ResponseBody => usage.response_body = value,
|
||||
UsageBodyField::ClientResponseBody => usage.client_response_body = value,
|
||||
}
|
||||
}
|
||||
Ok(usage)
|
||||
}
|
||||
|
||||
pub(crate) async fn resolve_body_ref(
|
||||
pool: &MysqlPool,
|
||||
body_ref: &str,
|
||||
) -> Result<Option<Value>, DataLayerError> {
|
||||
if let Some(payload_gzip) = sqlx::query_scalar::<_, Vec<u8>>(
|
||||
"SELECT payload_gzip FROM usage_body_blobs WHERE body_ref = ? LIMIT 1",
|
||||
)
|
||||
.bind(body_ref)
|
||||
.fetch_optional(pool)
|
||||
.await
|
||||
.map_sql_err()?
|
||||
{
|
||||
return inflate_json(&payload_gzip).map(Some);
|
||||
}
|
||||
let Some((request_id, field)) = parse_usage_body_ref(body_ref) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let (inline_column, compressed_column) = usage_body_sql_columns(field);
|
||||
let row = sqlx::query(&format!(
|
||||
"SELECT CAST({inline_column} AS CHAR) AS inline_body, {compressed_column} AS compressed_body FROM `usage` WHERE request_id = ? LIMIT 1"
|
||||
))
|
||||
.bind(request_id)
|
||||
.fetch_optional(pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
let Some(row) = row.as_ref() else {
|
||||
return Ok(None);
|
||||
};
|
||||
if let Some(raw) = row
|
||||
.try_get::<Option<String>, _>("inline_body")
|
||||
.map_sql_err()?
|
||||
{
|
||||
return serde_json::from_str(&raw).map(Some).map_err(|err| {
|
||||
DataLayerError::UnexpectedValue(format!("invalid inline usage body JSON: {err}"))
|
||||
});
|
||||
}
|
||||
row.try_get::<Option<Vec<u8>>, _>("compressed_body")
|
||||
.map_sql_err()?
|
||||
.map(|bytes| inflate_json(&bytes))
|
||||
.transpose()
|
||||
}
|
||||
|
||||
fn usage_body_sql_columns(field: UsageBodyField) -> (&'static str, &'static str) {
|
||||
match field {
|
||||
UsageBodyField::RequestBody => ("request_body", "request_body_compressed"),
|
||||
UsageBodyField::ProviderRequestBody => {
|
||||
("provider_request_body", "provider_request_body_compressed")
|
||||
}
|
||||
UsageBodyField::ResponseBody => ("response_body", "response_body_compressed"),
|
||||
UsageBodyField::ClientResponseBody => {
|
||||
("client_response_body", "client_response_body_compressed")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn inflate_json(bytes: &[u8]) -> Result<Value, DataLayerError> {
|
||||
let mut decoder = GzDecoder::new(bytes);
|
||||
let mut decoded = Vec::new();
|
||||
decoder.read_to_end(&mut decoded).map_err(|err| {
|
||||
DataLayerError::UnexpectedValue(format!("failed to decompress usage body: {err}"))
|
||||
})?;
|
||||
serde_json::from_slice(&decoded).map_err(|err| {
|
||||
DataLayerError::UnexpectedValue(format!("failed to decode usage body JSON: {err}"))
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,811 @@
|
||||
use aether_data_contracts::repository::usage::{
|
||||
ProviderApiKeyWindowUsageRequest, StoredRequestUsageAudit, UsageAuditKeywordSearchQuery,
|
||||
UsageAuditListQuery, UsageMonitoringErrorCountQuery, UsageMonitoringErrorListQuery,
|
||||
};
|
||||
use aether_data_contracts::DataLayerError;
|
||||
use sqlx::{MySql, QueryBuilder};
|
||||
|
||||
use crate::error::SqlResultExt;
|
||||
|
||||
use super::{http_capture, map_usage_row, row_u64, to_i64, MysqlUsageStorage, USAGE_COLUMNS};
|
||||
|
||||
const EFFECTIVE_PROVIDER_ID_EXPR: &str = r#"CASE
|
||||
WHEN usage_routing_snapshots.request_id IS NOT NULL
|
||||
THEN usage_routing_snapshots.selected_provider_id
|
||||
ELSE `usage`.provider_id
|
||||
END"#;
|
||||
|
||||
const EFFECTIVE_PROVIDER_API_KEY_ID_EXPR: &str = r#"CASE
|
||||
WHEN usage_routing_snapshots.request_id IS NOT NULL
|
||||
THEN usage_routing_snapshots.selected_provider_api_key_id
|
||||
ELSE `usage`.provider_api_key_id
|
||||
END"#;
|
||||
|
||||
const MONITORING_ERROR_PREDICATE: &str = r#"(
|
||||
LOWER(TRIM(COALESCE(`usage`.status, ''))) IN ('failed', 'error')
|
||||
OR (`usage`.error_category IS NOT NULL AND TRIM(`usage`.error_category) <> '')
|
||||
OR (
|
||||
TRIM(COALESCE(`usage`.status, '')) = ''
|
||||
AND (
|
||||
COALESCE(`usage`.status_code, 0) >= 400
|
||||
OR (`usage`.error_message IS NOT NULL AND TRIM(`usage`.error_message) <> '')
|
||||
)
|
||||
)
|
||||
)"#;
|
||||
|
||||
/// A SQL-side superset filter used before the runtime applies complex usage analytics.
|
||||
///
|
||||
/// Every scan has explicit time bounds; an empty range returns no rows. Optional dimensions
|
||||
/// further reduce the rows sent to the in-memory analytics implementation without changing its
|
||||
/// result.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct MysqlUsageReadFilter {
|
||||
created_from_unix_secs: u64,
|
||||
created_until_unix_secs: u64,
|
||||
user_id: Option<String>,
|
||||
api_key_id: Option<String>,
|
||||
provider_name: Option<String>,
|
||||
provider_id: Option<String>,
|
||||
model: Option<String>,
|
||||
api_format: Option<String>,
|
||||
endpoint_kind: Option<String>,
|
||||
is_stream: Option<bool>,
|
||||
has_format_conversion: Option<bool>,
|
||||
finalized_only: bool,
|
||||
completed_only: bool,
|
||||
}
|
||||
|
||||
impl MysqlUsageReadFilter {
|
||||
pub fn new(created_from_unix_secs: u64, created_until_unix_secs: u64) -> Self {
|
||||
Self {
|
||||
created_from_unix_secs,
|
||||
created_until_unix_secs,
|
||||
user_id: None,
|
||||
api_key_id: None,
|
||||
provider_name: None,
|
||||
provider_id: None,
|
||||
model: None,
|
||||
api_format: None,
|
||||
endpoint_kind: None,
|
||||
is_stream: None,
|
||||
has_format_conversion: None,
|
||||
finalized_only: false,
|
||||
completed_only: false,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn with_user_id(mut self, value: Option<&str>) -> Self {
|
||||
self.user_id = value.map(ToOwned::to_owned);
|
||||
self
|
||||
}
|
||||
|
||||
pub fn with_api_key_id(mut self, value: Option<&str>) -> Self {
|
||||
self.api_key_id = value.map(ToOwned::to_owned);
|
||||
self
|
||||
}
|
||||
|
||||
pub fn with_provider_name(mut self, value: Option<&str>) -> Self {
|
||||
self.provider_name = value.map(ToOwned::to_owned);
|
||||
self
|
||||
}
|
||||
|
||||
pub fn with_provider_id(mut self, value: Option<&str>) -> Self {
|
||||
self.provider_id = value.map(ToOwned::to_owned);
|
||||
self
|
||||
}
|
||||
|
||||
pub fn with_model(mut self, value: Option<&str>) -> Self {
|
||||
self.model = value.map(ToOwned::to_owned);
|
||||
self
|
||||
}
|
||||
|
||||
pub fn with_api_format(mut self, value: Option<&str>) -> Self {
|
||||
self.api_format = value.map(ToOwned::to_owned);
|
||||
self
|
||||
}
|
||||
|
||||
pub fn with_endpoint_kind(mut self, value: Option<&str>) -> Self {
|
||||
self.endpoint_kind = value.map(ToOwned::to_owned);
|
||||
self
|
||||
}
|
||||
|
||||
pub fn with_is_stream(mut self, value: Option<bool>) -> Self {
|
||||
self.is_stream = value;
|
||||
self
|
||||
}
|
||||
|
||||
pub fn with_has_format_conversion(mut self, value: Option<bool>) -> Self {
|
||||
self.has_format_conversion = value;
|
||||
self
|
||||
}
|
||||
|
||||
pub fn finalized_only(mut self) -> Self {
|
||||
self.finalized_only = true;
|
||||
self
|
||||
}
|
||||
|
||||
pub fn completed_only(mut self) -> Self {
|
||||
self.completed_only = true;
|
||||
self
|
||||
}
|
||||
|
||||
fn is_empty(&self) -> bool {
|
||||
self.created_from_unix_secs >= self.created_until_unix_secs
|
||||
}
|
||||
}
|
||||
|
||||
impl MysqlUsageStorage {
|
||||
pub async fn find_by_id(
|
||||
&self,
|
||||
id: &str,
|
||||
) -> Result<Option<StoredRequestUsageAudit>, DataLayerError> {
|
||||
let row = sqlx::query(&format!("{USAGE_COLUMNS} WHERE `usage`.id = ? LIMIT 1"))
|
||||
.bind(id)
|
||||
.fetch_optional(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
row.as_ref()
|
||||
.map(|row| map_usage_row(row, false))
|
||||
.transpose()
|
||||
}
|
||||
|
||||
pub async fn list_by_ids(
|
||||
&self,
|
||||
ids: &[String],
|
||||
) -> Result<Vec<StoredRequestUsageAudit>, DataLayerError> {
|
||||
if ids.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let mut builder = QueryBuilder::<MySql>::new(USAGE_COLUMNS);
|
||||
builder.push(" WHERE `usage`.id IN (");
|
||||
{
|
||||
let mut separated = builder.separated(", ");
|
||||
for id in ids {
|
||||
separated.push_bind(id.clone());
|
||||
}
|
||||
}
|
||||
builder.push(") ORDER BY created_at_unix_ms DESC, `usage`.id ASC");
|
||||
self.fetch_usage_items(builder).await
|
||||
}
|
||||
|
||||
pub async fn find_by_request_id(
|
||||
&self,
|
||||
request_id: &str,
|
||||
) -> Result<Option<StoredRequestUsageAudit>, DataLayerError> {
|
||||
let row = sqlx::query(&format!(
|
||||
"{USAGE_COLUMNS} WHERE `usage`.request_id = ? LIMIT 1"
|
||||
))
|
||||
.bind(request_id)
|
||||
.fetch_optional(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
let usage = row
|
||||
.as_ref()
|
||||
.map(|row| map_usage_row(row, true))
|
||||
.transpose()?;
|
||||
match usage {
|
||||
Some(usage) => http_capture::hydrate_usage_body_refs(&self.pool, usage)
|
||||
.await
|
||||
.map(Some),
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn resolve_body_ref(
|
||||
&self,
|
||||
body_ref: &str,
|
||||
) -> Result<Option<serde_json::Value>, DataLayerError> {
|
||||
http_capture::resolve_body_ref(&self.pool, body_ref).await
|
||||
}
|
||||
|
||||
pub async fn list_usage_audits(
|
||||
&self,
|
||||
query: &UsageAuditListQuery,
|
||||
) -> Result<Vec<StoredRequestUsageAudit>, DataLayerError> {
|
||||
let mut builder = QueryBuilder::<MySql>::new(USAGE_COLUMNS);
|
||||
let mut has_where = false;
|
||||
push_list_filters(&mut builder, query, &mut has_where)?;
|
||||
push_order_limit_offset(&mut builder, query.newest_first, query.limit, query.offset)?;
|
||||
self.fetch_usage_items(builder).await
|
||||
}
|
||||
|
||||
pub async fn count_usage_audits(
|
||||
&self,
|
||||
query: &UsageAuditListQuery,
|
||||
) -> Result<u64, DataLayerError> {
|
||||
let mut builder =
|
||||
QueryBuilder::<MySql>::new("SELECT CAST(COUNT(*) AS SIGNED) AS total FROM `usage`");
|
||||
let mut has_where = false;
|
||||
push_list_filters(&mut builder, query, &mut has_where)?;
|
||||
let row = builder.build().fetch_one(&self.pool).await.map_sql_err()?;
|
||||
row_u64(&row, "total")
|
||||
}
|
||||
|
||||
pub async fn list_usage_audits_by_keyword_search(
|
||||
&self,
|
||||
query: &UsageAuditKeywordSearchQuery,
|
||||
) -> Result<Vec<StoredRequestUsageAudit>, DataLayerError> {
|
||||
let mut builder = QueryBuilder::<MySql>::new(USAGE_COLUMNS);
|
||||
let mut has_where = false;
|
||||
push_keyword_filters(&mut builder, query, &mut has_where)?;
|
||||
push_order_limit_offset(&mut builder, query.newest_first, query.limit, query.offset)?;
|
||||
self.fetch_usage_items(builder).await
|
||||
}
|
||||
|
||||
pub async fn count_usage_audits_by_keyword_search(
|
||||
&self,
|
||||
query: &UsageAuditKeywordSearchQuery,
|
||||
) -> Result<u64, DataLayerError> {
|
||||
let mut builder =
|
||||
QueryBuilder::<MySql>::new("SELECT CAST(COUNT(*) AS SIGNED) AS total FROM `usage`");
|
||||
let mut has_where = false;
|
||||
push_keyword_filters(&mut builder, query, &mut has_where)?;
|
||||
let row = builder.build().fetch_one(&self.pool).await.map_sql_err()?;
|
||||
row_u64(&row, "total")
|
||||
}
|
||||
|
||||
pub async fn load_usage_records_in_range(
|
||||
&self,
|
||||
filter: &MysqlUsageReadFilter,
|
||||
) -> Result<Vec<StoredRequestUsageAudit>, DataLayerError> {
|
||||
if filter.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let mut builder = build_range_query(filter)?;
|
||||
let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?;
|
||||
rows.iter().map(|row| map_usage_row(row, false)).collect()
|
||||
}
|
||||
|
||||
pub async fn count_monitoring_usage_errors(
|
||||
&self,
|
||||
query: &UsageMonitoringErrorCountQuery,
|
||||
) -> Result<u64, DataLayerError> {
|
||||
if query.created_from_unix_secs >= query.created_until_unix_secs {
|
||||
return Ok(0);
|
||||
}
|
||||
let row = sqlx::query(&format!(
|
||||
r#"
|
||||
SELECT CAST(COUNT(*) AS SIGNED) AS total
|
||||
FROM `usage`
|
||||
WHERE created_at_unix_ms >= ?
|
||||
AND created_at_unix_ms < ?
|
||||
AND {MONITORING_ERROR_PREDICATE}
|
||||
"#
|
||||
))
|
||||
.bind(to_i64(
|
||||
query.created_from_unix_secs,
|
||||
"usage.created_at_unix_ms",
|
||||
)?)
|
||||
.bind(to_i64(
|
||||
query.created_until_unix_secs,
|
||||
"usage.created_at_unix_ms",
|
||||
)?)
|
||||
.fetch_one(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
row_u64(&row, "total")
|
||||
}
|
||||
|
||||
pub async fn list_monitoring_usage_errors(
|
||||
&self,
|
||||
query: &UsageMonitoringErrorListQuery,
|
||||
) -> Result<Vec<StoredRequestUsageAudit>, DataLayerError> {
|
||||
if query.created_from_unix_secs >= query.created_until_unix_secs {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let mut builder = QueryBuilder::<MySql>::new(USAGE_COLUMNS);
|
||||
builder
|
||||
.push(" WHERE `usage`.created_at_unix_ms >= ")
|
||||
.push_bind(to_i64(
|
||||
query.created_from_unix_secs,
|
||||
"usage.created_at_unix_ms",
|
||||
)?)
|
||||
.push(" AND `usage`.created_at_unix_ms < ")
|
||||
.push_bind(to_i64(
|
||||
query.created_until_unix_secs,
|
||||
"usage.created_at_unix_ms",
|
||||
)?)
|
||||
.push(" AND ")
|
||||
.push(MONITORING_ERROR_PREDICATE)
|
||||
.push(" ORDER BY created_at_unix_ms DESC, `usage`.id ASC");
|
||||
if let Some(limit) = query.limit {
|
||||
builder
|
||||
.push(" LIMIT ")
|
||||
.push_bind(usize_to_i64(limit, "usage monitoring limit")?);
|
||||
}
|
||||
self.fetch_usage_items(builder).await
|
||||
}
|
||||
|
||||
pub async fn list_recent_usage_audits(
|
||||
&self,
|
||||
user_id: Option<&str>,
|
||||
limit: usize,
|
||||
) -> Result<Vec<StoredRequestUsageAudit>, DataLayerError> {
|
||||
if limit == 0 {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let mut builder = QueryBuilder::<MySql>::new(USAGE_COLUMNS);
|
||||
if let Some(user_id) = user_id {
|
||||
builder
|
||||
.push(" WHERE `usage`.user_id = ")
|
||||
.push_bind(user_id.to_string());
|
||||
}
|
||||
builder
|
||||
.push(" ORDER BY created_at_unix_ms DESC, `usage`.id ASC LIMIT ")
|
||||
.push_bind(usize_to_i64(limit, "recent usage limit")?);
|
||||
self.fetch_usage_items(builder).await
|
||||
}
|
||||
|
||||
pub async fn load_usage_records_by_api_key_ids(
|
||||
&self,
|
||||
api_key_ids: &[String],
|
||||
) -> Result<Vec<StoredRequestUsageAudit>, DataLayerError> {
|
||||
if api_key_ids.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let mut builder = QueryBuilder::<MySql>::new(USAGE_COLUMNS);
|
||||
builder.push(" WHERE `usage`.api_key_id IN (");
|
||||
push_string_list(&mut builder, api_key_ids);
|
||||
builder.push(") ORDER BY `usage`.created_at_unix_ms ASC, `usage`.request_id ASC");
|
||||
self.fetch_usage_items(builder).await
|
||||
}
|
||||
|
||||
pub async fn load_usage_records_by_provider_api_key_ids(
|
||||
&self,
|
||||
provider_api_key_ids: &[String],
|
||||
) -> Result<Vec<StoredRequestUsageAudit>, DataLayerError> {
|
||||
if provider_api_key_ids.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let mut builder = QueryBuilder::<MySql>::new(USAGE_COLUMNS);
|
||||
builder
|
||||
.push(" WHERE (")
|
||||
.push(EFFECTIVE_PROVIDER_API_KEY_ID_EXPR)
|
||||
.push(") IN (");
|
||||
push_string_list(&mut builder, provider_api_key_ids);
|
||||
builder.push(") ORDER BY `usage`.created_at_unix_ms ASC, `usage`.request_id ASC");
|
||||
self.fetch_usage_items(builder).await
|
||||
}
|
||||
|
||||
pub async fn load_usage_records_by_provider_api_key_windows(
|
||||
&self,
|
||||
requests: &[ProviderApiKeyWindowUsageRequest],
|
||||
) -> Result<Vec<StoredRequestUsageAudit>, DataLayerError> {
|
||||
if requests.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let mut builder = QueryBuilder::<MySql>::new(USAGE_COLUMNS);
|
||||
builder.push(" WHERE ");
|
||||
for request in requests {
|
||||
let provider_api_key_id = request.provider_api_key_id.trim();
|
||||
if provider_api_key_id.is_empty() {
|
||||
return Err(DataLayerError::InvalidInput(
|
||||
"provider api key window usage provider_api_key_id cannot be empty".to_string(),
|
||||
));
|
||||
}
|
||||
let window_code = request.window_code.trim();
|
||||
if window_code.is_empty() {
|
||||
return Err(DataLayerError::InvalidInput(
|
||||
"provider api key window usage window_code cannot be empty".to_string(),
|
||||
));
|
||||
}
|
||||
if request.start_unix_secs >= request.end_unix_secs {
|
||||
return Err(DataLayerError::InvalidInput(
|
||||
"provider api key window usage range must be non-empty".to_string(),
|
||||
));
|
||||
}
|
||||
}
|
||||
{
|
||||
let mut separated = builder.separated(" OR ");
|
||||
for request in requests {
|
||||
separated
|
||||
.push("((")
|
||||
.push(EFFECTIVE_PROVIDER_API_KEY_ID_EXPR)
|
||||
.push(") = ")
|
||||
.push_bind(request.provider_api_key_id.trim().to_string())
|
||||
.push(" AND `usage`.created_at_unix_ms >= ")
|
||||
.push_bind(to_i64(request.start_unix_secs, "usage.created_at_unix_ms")?)
|
||||
.push(" AND `usage`.created_at_unix_ms < ")
|
||||
.push_bind(to_i64(request.end_unix_secs, "usage.created_at_unix_ms")?)
|
||||
.push(")");
|
||||
}
|
||||
}
|
||||
builder.push(" ORDER BY `usage`.created_at_unix_ms ASC, `usage`.request_id ASC");
|
||||
self.fetch_usage_items(builder).await
|
||||
}
|
||||
|
||||
pub async fn load_usage_records_for_provider_since(
|
||||
&self,
|
||||
provider_id: &str,
|
||||
since_unix_secs: u64,
|
||||
) -> Result<Vec<StoredRequestUsageAudit>, DataLayerError> {
|
||||
let mut builder = QueryBuilder::<MySql>::new(USAGE_COLUMNS);
|
||||
builder
|
||||
.push(" WHERE (")
|
||||
.push(EFFECTIVE_PROVIDER_ID_EXPR)
|
||||
.push(") = ")
|
||||
.push_bind(provider_id.to_string())
|
||||
.push(" AND `usage`.created_at_unix_ms >= ")
|
||||
.push_bind(to_i64(since_unix_secs, "usage.created_at_unix_ms")?)
|
||||
.push(" ORDER BY `usage`.created_at_unix_ms ASC, `usage`.request_id ASC");
|
||||
self.fetch_usage_items(builder).await
|
||||
}
|
||||
|
||||
async fn fetch_usage_items(
|
||||
&self,
|
||||
mut builder: QueryBuilder<'_, MySql>,
|
||||
) -> Result<Vec<StoredRequestUsageAudit>, DataLayerError> {
|
||||
let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?;
|
||||
rows.iter().map(|row| map_usage_row(row, false)).collect()
|
||||
}
|
||||
}
|
||||
|
||||
fn build_range_query(
|
||||
filter: &MysqlUsageReadFilter,
|
||||
) -> Result<QueryBuilder<'static, MySql>, DataLayerError> {
|
||||
let mut builder = QueryBuilder::<MySql>::new(USAGE_COLUMNS);
|
||||
builder
|
||||
.push(" WHERE `usage`.created_at_unix_ms >= ")
|
||||
.push_bind(to_i64(
|
||||
filter.created_from_unix_secs,
|
||||
"usage.created_at_unix_ms",
|
||||
)?)
|
||||
.push(" AND `usage`.created_at_unix_ms < ")
|
||||
.push_bind(to_i64(
|
||||
filter.created_until_unix_secs,
|
||||
"usage.created_at_unix_ms",
|
||||
)?);
|
||||
push_optional_text_filter(&mut builder, "`usage`.user_id", filter.user_id.as_deref());
|
||||
push_optional_text_filter(
|
||||
&mut builder,
|
||||
"`usage`.api_key_id",
|
||||
filter.api_key_id.as_deref(),
|
||||
);
|
||||
push_optional_text_filter(
|
||||
&mut builder,
|
||||
"`usage`.provider_name",
|
||||
filter.provider_name.as_deref(),
|
||||
);
|
||||
if let Some(provider_id) = filter.provider_id.as_deref() {
|
||||
builder
|
||||
.push(" AND (")
|
||||
.push(EFFECTIVE_PROVIDER_ID_EXPR)
|
||||
.push(") = ")
|
||||
.push_bind(provider_id.to_string());
|
||||
}
|
||||
push_optional_text_filter(&mut builder, "`usage`.model", filter.model.as_deref());
|
||||
push_optional_text_filter(
|
||||
&mut builder,
|
||||
"`usage`.api_format",
|
||||
filter.api_format.as_deref(),
|
||||
);
|
||||
push_optional_text_filter(
|
||||
&mut builder,
|
||||
"`usage`.endpoint_kind",
|
||||
filter.endpoint_kind.as_deref(),
|
||||
);
|
||||
if let Some(is_stream) = filter.is_stream {
|
||||
builder
|
||||
.push(" AND `usage`.is_stream = ")
|
||||
.push_bind(is_stream);
|
||||
}
|
||||
if let Some(has_format_conversion) = filter.has_format_conversion {
|
||||
builder
|
||||
.push(" AND CASE WHEN usage_routing_snapshots.request_id IS NOT NULL ")
|
||||
.push("THEN COALESCE(usage_routing_snapshots.has_format_conversion, FALSE) ")
|
||||
.push("ELSE COALESCE(`usage`.has_format_conversion, FALSE) END = ")
|
||||
.push_bind(has_format_conversion);
|
||||
}
|
||||
if filter.finalized_only {
|
||||
builder.push(
|
||||
" AND `usage`.status NOT IN ('pending', 'streaming') \
|
||||
AND `usage`.provider_name NOT IN ('unknown', 'pending')",
|
||||
);
|
||||
}
|
||||
if filter.completed_only {
|
||||
builder.push(" AND `usage`.status = 'completed'");
|
||||
}
|
||||
builder.push(" ORDER BY `usage`.created_at_unix_ms ASC, `usage`.request_id ASC");
|
||||
Ok(builder)
|
||||
}
|
||||
|
||||
fn push_list_filters(
|
||||
builder: &mut QueryBuilder<'_, MySql>,
|
||||
query: &UsageAuditListQuery,
|
||||
has_where: &mut bool,
|
||||
) -> Result<(), DataLayerError> {
|
||||
if let Some(value) = query.created_from_unix_secs {
|
||||
push_where(builder, has_where);
|
||||
builder
|
||||
.push("`usage`.created_at_unix_ms >= ")
|
||||
.push_bind(to_i64(value, "usage.created_at_unix_ms")?);
|
||||
}
|
||||
if let Some(value) = query.created_until_unix_secs {
|
||||
push_where(builder, has_where);
|
||||
builder
|
||||
.push("`usage`.created_at_unix_ms < ")
|
||||
.push_bind(to_i64(value, "usage.created_at_unix_ms")?);
|
||||
}
|
||||
for (column, value) in [
|
||||
("`usage`.user_id", query.user_id.as_deref()),
|
||||
("`usage`.provider_name", query.provider_name.as_deref()),
|
||||
("`usage`.model", query.model.as_deref()),
|
||||
("`usage`.api_format", query.api_format.as_deref()),
|
||||
] {
|
||||
if let Some(value) = value {
|
||||
push_where(builder, has_where);
|
||||
builder
|
||||
.push(column)
|
||||
.push(" = ")
|
||||
.push_bind(value.to_string());
|
||||
}
|
||||
}
|
||||
if let Some(client_family) = query
|
||||
.client_family
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
{
|
||||
push_where(builder, has_where);
|
||||
builder
|
||||
.push("LOWER(COALESCE(NULLIF(TRIM(JSON_UNQUOTE(JSON_EXTRACT(`usage`.request_metadata, '$.client_session_affinity.client_family'))), ''), NULLIF(TRIM(JSON_UNQUOTE(JSON_EXTRACT(`usage`.request_metadata, '$.client_family'))), ''))) = ")
|
||||
.push_bind(client_family.to_ascii_lowercase());
|
||||
}
|
||||
if query.exclude_unknown_model_or_provider {
|
||||
push_where(builder, has_where);
|
||||
builder.push(
|
||||
"(LOWER(TRIM(COALESCE(`usage`.model, ''))) NOT IN ('unknown', 'unknow') \
|
||||
AND LOWER(TRIM(COALESCE(`usage`.provider_name, ''))) NOT IN ('unknown', 'unknow'))",
|
||||
);
|
||||
}
|
||||
if let Some(statuses) = query
|
||||
.statuses
|
||||
.as_deref()
|
||||
.filter(|values| !values.is_empty())
|
||||
{
|
||||
push_where(builder, has_where);
|
||||
builder.push("`usage`.status IN (");
|
||||
push_string_list(builder, statuses);
|
||||
builder.push(")");
|
||||
}
|
||||
if !query.exclude_status_codes.is_empty() {
|
||||
push_where(builder, has_where);
|
||||
builder.push("(`usage`.status_code IS NULL OR `usage`.status_code NOT IN (");
|
||||
{
|
||||
let mut separated = builder.separated(", ");
|
||||
for status_code in &query.exclude_status_codes {
|
||||
separated.push_bind(i64::from(*status_code));
|
||||
}
|
||||
}
|
||||
builder.push("))");
|
||||
}
|
||||
if let Some(is_stream) = query.is_stream {
|
||||
push_where(builder, has_where);
|
||||
builder.push("`usage`.is_stream = ").push_bind(is_stream);
|
||||
}
|
||||
if query.error_only {
|
||||
push_where(builder, has_where);
|
||||
builder.push(
|
||||
"(`usage`.status = 'failed' \
|
||||
OR COALESCE(`usage`.status_code, 0) >= 400 \
|
||||
OR (`usage`.error_message IS NOT NULL AND TRIM(`usage`.error_message) <> ''))",
|
||||
);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn push_keyword_filters(
|
||||
builder: &mut QueryBuilder<'_, MySql>,
|
||||
query: &UsageAuditKeywordSearchQuery,
|
||||
has_where: &mut bool,
|
||||
) -> Result<(), DataLayerError> {
|
||||
push_list_filters(
|
||||
builder,
|
||||
&UsageAuditListQuery {
|
||||
created_from_unix_secs: query.created_from_unix_secs,
|
||||
created_until_unix_secs: query.created_until_unix_secs,
|
||||
user_id: query.user_id.clone(),
|
||||
provider_name: query.provider_name.clone(),
|
||||
model: query.model.clone(),
|
||||
api_format: query.api_format.clone(),
|
||||
client_family: query.client_family.clone(),
|
||||
exclude_unknown_model_or_provider: query.exclude_unknown_model_or_provider,
|
||||
statuses: query.statuses.clone(),
|
||||
exclude_status_codes: query.exclude_status_codes.clone(),
|
||||
is_stream: query.is_stream,
|
||||
error_only: query.error_only,
|
||||
limit: None,
|
||||
offset: None,
|
||||
newest_first: query.newest_first,
|
||||
},
|
||||
has_where,
|
||||
)?;
|
||||
|
||||
for (index, keyword) in query.keywords.iter().enumerate() {
|
||||
let keyword = keyword.trim();
|
||||
if keyword.is_empty() {
|
||||
continue;
|
||||
}
|
||||
let pattern = format!("%{}%", keyword.to_ascii_lowercase());
|
||||
push_where(builder, has_where);
|
||||
builder
|
||||
.push("(LOWER(COALESCE(`usage`.model, '')) LIKE ")
|
||||
.push_bind(pattern.clone())
|
||||
.push(" OR LOWER(COALESCE(`usage`.provider_name, '')) LIKE ")
|
||||
.push_bind(pattern.clone());
|
||||
if query.auth_user_reader_available {
|
||||
if let Some(ids) = query
|
||||
.matched_user_ids_by_keyword
|
||||
.get(index)
|
||||
.filter(|ids| !ids.is_empty())
|
||||
{
|
||||
builder.push(" OR `usage`.user_id IN (");
|
||||
push_string_list(builder, ids);
|
||||
builder.push(")");
|
||||
}
|
||||
} else {
|
||||
builder
|
||||
.push(" OR LOWER(COALESCE(`usage`.username, '')) LIKE ")
|
||||
.push_bind(pattern.clone());
|
||||
}
|
||||
if query.auth_api_key_reader_available {
|
||||
if let Some(ids) = query
|
||||
.matched_api_key_ids_by_keyword
|
||||
.get(index)
|
||||
.filter(|ids| !ids.is_empty())
|
||||
{
|
||||
builder.push(" OR `usage`.api_key_id IN (");
|
||||
push_string_list(builder, ids);
|
||||
builder.push(")");
|
||||
}
|
||||
} else {
|
||||
builder
|
||||
.push(" OR LOWER(COALESCE(`usage`.api_key_name, '')) LIKE ")
|
||||
.push_bind(pattern);
|
||||
}
|
||||
builder.push(")");
|
||||
}
|
||||
|
||||
if let Some(username_keyword) = query
|
||||
.username_keyword
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
{
|
||||
push_where(builder, has_where);
|
||||
if query.auth_user_reader_available {
|
||||
if query.matched_user_ids_for_username.is_empty() {
|
||||
builder.push("FALSE");
|
||||
} else {
|
||||
builder.push("`usage`.user_id IN (");
|
||||
push_string_list(builder, &query.matched_user_ids_for_username);
|
||||
builder.push(")");
|
||||
}
|
||||
} else {
|
||||
builder
|
||||
.push("LOWER(COALESCE(`usage`.username, '')) LIKE ")
|
||||
.push_bind(format!("%{}%", username_keyword.to_ascii_lowercase()));
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn push_order_limit_offset(
|
||||
builder: &mut QueryBuilder<'_, MySql>,
|
||||
newest_first: bool,
|
||||
limit: Option<usize>,
|
||||
offset: Option<usize>,
|
||||
) -> Result<(), DataLayerError> {
|
||||
if newest_first {
|
||||
builder.push(" ORDER BY created_at_unix_ms DESC, `usage`.id ASC");
|
||||
} else {
|
||||
builder.push(" ORDER BY created_at_unix_ms ASC, `usage`.request_id ASC");
|
||||
}
|
||||
match (limit, offset) {
|
||||
(Some(limit), offset) => {
|
||||
builder
|
||||
.push(" LIMIT ")
|
||||
.push_bind(usize_to_i64(limit, "usage list limit")?);
|
||||
if let Some(offset) = offset {
|
||||
builder
|
||||
.push(" OFFSET ")
|
||||
.push_bind(usize_to_i64(offset, "usage list offset")?);
|
||||
}
|
||||
}
|
||||
(None, Some(offset)) => {
|
||||
builder
|
||||
.push(" LIMIT 18446744073709551615 OFFSET ")
|
||||
.push_bind(usize_to_i64(offset, "usage list offset")?);
|
||||
}
|
||||
(None, None) => {}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn push_where(builder: &mut QueryBuilder<'_, MySql>, has_where: &mut bool) {
|
||||
builder.push(if *has_where { " AND " } else { " WHERE " });
|
||||
*has_where = true;
|
||||
}
|
||||
|
||||
fn push_optional_text_filter(
|
||||
builder: &mut QueryBuilder<'_, MySql>,
|
||||
column: &'static str,
|
||||
value: Option<&str>,
|
||||
) {
|
||||
if let Some(value) = value {
|
||||
builder
|
||||
.push(" AND ")
|
||||
.push(column)
|
||||
.push(" = ")
|
||||
.push_bind(value.to_string());
|
||||
}
|
||||
}
|
||||
|
||||
fn push_string_list(builder: &mut QueryBuilder<'_, MySql>, values: &[String]) {
|
||||
let mut separated = builder.separated(", ");
|
||||
for value in values {
|
||||
separated.push_bind(value.clone());
|
||||
}
|
||||
}
|
||||
|
||||
fn usize_to_i64(value: usize, field: &str) -> Result<i64, DataLayerError> {
|
||||
i64::try_from(value).map_err(|_| DataLayerError::InvalidInput(format!("{field} overflow")))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn range_reads_always_have_both_time_bounds() {
|
||||
let filter = MysqlUsageReadFilter::new(100, 200)
|
||||
.with_user_id(Some("user-1"))
|
||||
.finalized_only();
|
||||
let query = build_range_query(&filter).expect("range query should build");
|
||||
let sql = query.sql();
|
||||
assert!(sql.contains("created_at_unix_ms >= ?"));
|
||||
assert!(sql.contains("created_at_unix_ms < ?"));
|
||||
assert!(sql.contains("`usage`.user_id = ?"));
|
||||
assert!(sql.contains("status NOT IN ('pending', 'streaming')"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn audit_reads_keep_pagination_in_mysql() {
|
||||
let mut query = QueryBuilder::<MySql>::new(USAGE_COLUMNS);
|
||||
push_order_limit_offset(&mut query, true, Some(25), Some(50))
|
||||
.expect("pagination should build");
|
||||
assert!(query.sql().contains("LIMIT ? OFFSET ?"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn offset_only_uses_mysql_unbounded_limit_syntax() {
|
||||
let mut query = QueryBuilder::<MySql>::new(USAGE_COLUMNS);
|
||||
push_order_limit_offset(&mut query, false, None, Some(50))
|
||||
.expect("pagination should build");
|
||||
assert!(query.sql().contains("LIMIT 18446744073709551615 OFFSET ?"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn usage_projection_and_legacy_keyword_search_keep_snapshot_names() {
|
||||
assert!(USAGE_COLUMNS.contains("`usage`.username"));
|
||||
assert!(USAGE_COLUMNS.contains("`usage`.api_key_name"));
|
||||
|
||||
let mut query = QueryBuilder::<MySql>::new(USAGE_COLUMNS);
|
||||
let mut has_where = false;
|
||||
push_keyword_filters(
|
||||
&mut query,
|
||||
&UsageAuditKeywordSearchQuery {
|
||||
keywords: vec!["legacy".to_string()],
|
||||
..UsageAuditKeywordSearchQuery::default()
|
||||
},
|
||||
&mut has_where,
|
||||
)
|
||||
.expect("keyword query should build");
|
||||
assert!(query
|
||||
.sql()
|
||||
.contains("LOWER(COALESCE(`usage`.username, '')) LIKE ?"));
|
||||
assert!(query
|
||||
.sql()
|
||||
.contains("LOWER(COALESCE(`usage`.api_key_name, '')) LIKE ?"));
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,6 +1,10 @@
|
||||
use super::{MysqlUsageStorage, MysqlUsageWriteRepository};
|
||||
use crate::run_migrations;
|
||||
use aether_data_contracts::repository::usage::{UpsertUsageRecord, UsageWriteRepository};
|
||||
use aether_data_contracts::repository::usage::{
|
||||
UpsertUsageRecord, UsageAuditListQuery, UsageBodyCaptureState, UsageCleanupExecutionMode,
|
||||
UsageCleanupTargets, UsageCleanupWindow, UsageWriteRepository,
|
||||
};
|
||||
use chrono::DateTime;
|
||||
|
||||
#[tokio::test]
|
||||
async fn repository_builds_from_lazy_pool() {
|
||||
@@ -48,6 +52,45 @@ fn mysql_dashboard_reads_imported_daily_aggregates() {
|
||||
assert!(source.contains("CAST(COALESCE(SUM(total_requests), 0) AS SIGNED) AS requests"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mysql_usage_stat_rebuilds_aggregate_in_sql() {
|
||||
let source = include_str!("../usage.rs");
|
||||
assert!(source.contains("UPDATE api_keys\nJOIN ("));
|
||||
assert!(source.contains("AND status NOT IN ('pending', 'streaming')"));
|
||||
assert!(source.contains("MYSQL_USAGE_CANONICAL_TOTAL_TOKENS_EXPR"));
|
||||
assert!(source.contains("MAX(created_at_unix_ms) AS last_used_at"));
|
||||
assert!(source.contains("UPDATE provider_api_keys\nJOIN ("));
|
||||
assert!(source.contains("GROUP BY provider_api_key_id"));
|
||||
assert!(!source.contains("struct ProviderKeyStats"));
|
||||
assert!(super::MYSQL_PROVIDER_KEY_SUCCESS_FLAG_EXPR
|
||||
.contains("status IN ('completed', 'success', 'ok', 'billed', 'settled')"));
|
||||
assert!(super::MYSQL_PROVIDER_KEY_SUCCESS_FLAG_EXPR
|
||||
.contains("error_message IS NULL OR TRIM(error_message) = ''"));
|
||||
assert!(super::MYSQL_PROVIDER_KEY_ERROR_FLAG_EXPR
|
||||
.contains("status NOT IN ('pending', 'streaming')"));
|
||||
assert!(super::MYSQL_USAGE_CANONICAL_TOTAL_TOKENS_EXPR.contains(
|
||||
"COALESCE(`usage`.input_tokens, 0) - COALESCE(`usage`.cache_read_input_tokens, 0)"
|
||||
));
|
||||
assert!(super::MYSQL_USAGE_CANONICAL_TOTAL_TOKENS_EXPR
|
||||
.contains("GREATEST(COALESCE(`usage`.output_tokens, 0), 0)"));
|
||||
assert!(super::MYSQL_USAGE_CANONICAL_TOTAL_TOKENS_EXPR
|
||||
.contains("NULLIF(GREATEST(COALESCE(`usage`.total_tokens, 0), 0), 0)"));
|
||||
let snapshot_position = super::MYSQL_USAGE_CANONICAL_TOTAL_TOKENS_EXPR
|
||||
.find("settlement.billing_effective_input_tokens")
|
||||
.expect("canonical total must read settlement snapshots");
|
||||
let raw_position = super::MYSQL_USAGE_CANONICAL_TOTAL_TOKENS_EXPR
|
||||
.find("NULLIF(GREATEST(COALESCE(`usage`.total_tokens, 0), 0), 0)")
|
||||
.expect("canonical total must preserve non-zero legacy raw totals");
|
||||
assert!(snapshot_position < raw_position);
|
||||
assert!(source.matches("canonical_total_tokens_expr =").count() >= 4);
|
||||
assert_eq!(
|
||||
source
|
||||
.matches("LEFT JOIN usage_settlement_snapshots AS settlement")
|
||||
.count(),
|
||||
4
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mysql_usage_upsert_keeps_terminal_state_when_streaming_arrives_late() {
|
||||
assert!(super::UPSERT_USAGE_SQL.contains(
|
||||
@@ -110,7 +153,7 @@ fn mysql_usage_upsert_guards_candidate_identity_metadata_and_routing_from_late_l
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn mysql_usage_write_repository_upserts_when_url_is_set() {
|
||||
async fn mysql_usage_write_repository_upserts_and_flushes_counters_when_url_is_set() {
|
||||
let Some(database_url) = std::env::var("AETHER_TEST_MYSQL_URL")
|
||||
.ok()
|
||||
.filter(|value| !value.trim().is_empty())
|
||||
@@ -133,29 +176,35 @@ async fn mysql_usage_write_repository_upserts_when_url_is_set() {
|
||||
let api_key_id = format!("api-key-{suffix}");
|
||||
let provider_id = format!("provider-{suffix}");
|
||||
let provider_key_id = format!("provider-key-{suffix}");
|
||||
let model_name = format!("model-{suffix}");
|
||||
seed_stats_targets(&pool, &user_id, &api_key_id, &provider_id, &provider_key_id).await;
|
||||
sqlx::query("INSERT INTO global_models (id, name, created_at, updated_at) VALUES (?, ?, 1, 1)")
|
||||
.bind(format!("global-model-{suffix}"))
|
||||
.bind(&model_name)
|
||||
.execute(&pool)
|
||||
.await
|
||||
.expect("global model should seed");
|
||||
|
||||
let repository = MysqlUsageWriteRepository::new(pool.clone());
|
||||
let record = repository
|
||||
.upsert(sample_usage(
|
||||
&format!("request-{suffix}"),
|
||||
&user_id,
|
||||
&api_key_id,
|
||||
&provider_id,
|
||||
&provider_key_id,
|
||||
"completed",
|
||||
"pending",
|
||||
1_000,
|
||||
))
|
||||
.await
|
||||
.expect("usage should upsert");
|
||||
let mut usage = sample_usage(
|
||||
&format!("request-{suffix}"),
|
||||
&user_id,
|
||||
&api_key_id,
|
||||
&provider_id,
|
||||
&provider_key_id,
|
||||
"completed",
|
||||
"pending",
|
||||
1_000,
|
||||
);
|
||||
usage.model.clone_from(&model_name);
|
||||
let record = repository.upsert(usage).await.expect("usage should upsert");
|
||||
|
||||
assert_eq!(record.api_key_id.as_deref(), Some(api_key_id.as_str()));
|
||||
assert_eq!(
|
||||
record.provider_api_key_id.as_deref(),
|
||||
Some(provider_key_id.as_str())
|
||||
);
|
||||
assert_eq!(record.total_tokens, 7);
|
||||
assert_eq!(record.total_tokens, 5);
|
||||
assert_eq!(
|
||||
record.request_metadata.as_ref().unwrap()["upstream_is_stream"],
|
||||
true
|
||||
@@ -168,6 +217,25 @@ async fn mysql_usage_write_repository_upserts_when_url_is_set() {
|
||||
.expect("usage stream mode should load");
|
||||
assert_eq!(upstream_is_stream, Some(true));
|
||||
|
||||
repository
|
||||
.flush_usage_counter_deltas(100)
|
||||
.await
|
||||
.expect("usage counter deltas should flush");
|
||||
assert!(
|
||||
repository
|
||||
.rebuild_api_key_usage_stats()
|
||||
.await
|
||||
.expect("api key stats should rebuild")
|
||||
>= 1
|
||||
);
|
||||
assert!(
|
||||
repository
|
||||
.rebuild_provider_api_key_usage_stats()
|
||||
.await
|
||||
.expect("provider api key stats should rebuild")
|
||||
>= 1
|
||||
);
|
||||
|
||||
let stats = sqlx::query_as::<_, (i64, i64, f64, Option<i64>)>(
|
||||
"SELECT total_requests, total_tokens, total_cost_usd, last_used_at FROM api_keys WHERE id = ?",
|
||||
)
|
||||
@@ -175,7 +243,7 @@ async fn mysql_usage_write_repository_upserts_when_url_is_set() {
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.expect("api key stats should load");
|
||||
assert_eq!(stats, (1, 7, 0.5, Some(1_000)));
|
||||
assert_eq!(stats, (1, 5, 0.5, Some(1_000)));
|
||||
|
||||
let provider_stats = sqlx::query_as::<_, (i64, i64, i64, i64, f64, i64, Option<i64>)>(
|
||||
"SELECT request_count, success_count, error_count, total_tokens, total_cost_usd, total_response_time_ms, last_used_at FROM provider_api_keys WHERE id = ?",
|
||||
@@ -184,7 +252,575 @@ async fn mysql_usage_write_repository_upserts_when_url_is_set() {
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.expect("provider key stats should load");
|
||||
assert_eq!(provider_stats, (1, 1, 0, 7, 0.5, 42, Some(1_000)));
|
||||
assert_eq!(provider_stats, (1, 1, 0, 5, 0.5, 42, Some(1_000)));
|
||||
let model_usage_count: i64 =
|
||||
sqlx::query_scalar("SELECT usage_count FROM global_models WHERE name = ?")
|
||||
.bind(&model_name)
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.expect("global model usage count should load");
|
||||
assert_eq!(model_usage_count, 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn mysql_canonical_totals_preserve_legacy_total_tokens_only_rows_when_url_is_set() {
|
||||
let Some(database_url) = std::env::var("AETHER_TEST_MYSQL_URL")
|
||||
.ok()
|
||||
.filter(|value| !value.trim().is_empty())
|
||||
else {
|
||||
eprintln!(
|
||||
"skipping mysql canonical total_tokens test because AETHER_TEST_MYSQL_URL is unset"
|
||||
);
|
||||
return;
|
||||
};
|
||||
|
||||
let pool = sqlx::mysql::MySqlPoolOptions::new()
|
||||
.max_connections(1)
|
||||
.connect(&database_url)
|
||||
.await
|
||||
.expect("mysql test pool should connect");
|
||||
run_migrations(&pool)
|
||||
.await
|
||||
.expect("mysql migrations should run");
|
||||
|
||||
let suffix = unique_suffix();
|
||||
let user_id = format!("total-only-user-{suffix}");
|
||||
let api_key_id = format!("total-only-api-key-{suffix}");
|
||||
let provider_id = format!("total-only-provider-{suffix}");
|
||||
let provider_key_id = format!("total-only-provider-key-{suffix}");
|
||||
seed_stats_targets(&pool, &user_id, &api_key_id, &provider_id, &provider_key_id).await;
|
||||
|
||||
let repository = MysqlUsageWriteRepository::new(pool.clone());
|
||||
let mut usage = sample_usage(
|
||||
&format!("total-only-request-{suffix}"),
|
||||
&user_id,
|
||||
&api_key_id,
|
||||
&provider_id,
|
||||
&provider_key_id,
|
||||
"completed",
|
||||
"pending",
|
||||
1_000,
|
||||
);
|
||||
usage.input_tokens = Some(0);
|
||||
usage.output_tokens = Some(0);
|
||||
usage.cache_creation_input_tokens = Some(0);
|
||||
usage.cache_creation_ephemeral_5m_input_tokens = Some(0);
|
||||
usage.cache_creation_ephemeral_1h_input_tokens = Some(0);
|
||||
usage.cache_read_input_tokens = Some(0);
|
||||
usage.total_tokens = Some(77);
|
||||
|
||||
let request_id = format!("total-only-request-{suffix}");
|
||||
usage.request_id.clone_from(&request_id);
|
||||
repository
|
||||
.upsert(usage)
|
||||
.await
|
||||
.expect("legacy total-only usage should upsert");
|
||||
sqlx::query("DELETE FROM usage_settlement_snapshots WHERE request_id = ?")
|
||||
.bind(&request_id)
|
||||
.execute(&pool)
|
||||
.await
|
||||
.expect("legacy fixture must not have a settlement snapshot");
|
||||
let stored = repository
|
||||
.find_by_request_id(&request_id)
|
||||
.await
|
||||
.expect("legacy total-only usage should load")
|
||||
.expect("legacy total-only usage should exist");
|
||||
assert_eq!(stored.total_tokens, 77);
|
||||
|
||||
repository
|
||||
.rebuild_api_key_usage_stats()
|
||||
.await
|
||||
.expect("api key stats should rebuild");
|
||||
repository
|
||||
.rebuild_provider_api_key_usage_stats()
|
||||
.await
|
||||
.expect("provider api key stats should rebuild");
|
||||
|
||||
let api_key_total: i64 = sqlx::query_scalar("SELECT total_tokens FROM api_keys WHERE id = ?")
|
||||
.bind(&api_key_id)
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.expect("api key total should load");
|
||||
let provider_key_total: i64 =
|
||||
sqlx::query_scalar("SELECT total_tokens FROM provider_api_keys WHERE id = ?")
|
||||
.bind(&provider_key_id)
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.expect("provider key total should load");
|
||||
assert_eq!(api_key_total, 77);
|
||||
assert_eq!(provider_key_total, 77);
|
||||
|
||||
let totals = MysqlUsageStorage::new(pool)
|
||||
.summarize_usage_totals_by_user_ids(std::slice::from_ref(&user_id))
|
||||
.await
|
||||
.expect("user usage totals should load");
|
||||
assert_eq!(totals.len(), 1);
|
||||
assert_eq!(totals[0].total_tokens, 77);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn mysql_concurrent_same_request_upserts_enqueue_counters_once_when_url_is_set() {
|
||||
let Some(database_url) = std::env::var("AETHER_TEST_MYSQL_URL")
|
||||
.ok()
|
||||
.filter(|value| !value.trim().is_empty())
|
||||
else {
|
||||
eprintln!(
|
||||
"skipping mysql usage counter concurrency test because AETHER_TEST_MYSQL_URL is unset"
|
||||
);
|
||||
return;
|
||||
};
|
||||
let pool = sqlx::mysql::MySqlPoolOptions::new()
|
||||
.max_connections(4)
|
||||
.connect(&database_url)
|
||||
.await
|
||||
.expect("mysql test pool should connect");
|
||||
run_migrations(&pool)
|
||||
.await
|
||||
.expect("mysql migrations should run");
|
||||
|
||||
let suffix = unique_suffix();
|
||||
let user_id = format!("counter-user-{suffix}");
|
||||
let api_key_id = format!("counter-api-key-{suffix}");
|
||||
let provider_id = format!("counter-provider-{suffix}");
|
||||
let provider_key_id = format!("counter-provider-key-{suffix}");
|
||||
let model_name = format!("counter-model-{suffix}");
|
||||
let request_id = format!("counter-request-{suffix}");
|
||||
seed_stats_targets(&pool, &user_id, &api_key_id, &provider_id, &provider_key_id).await;
|
||||
sqlx::query("INSERT INTO global_models (id, name, created_at, updated_at) VALUES (?, ?, 1, 1)")
|
||||
.bind(format!("counter-global-model-{suffix}"))
|
||||
.bind(&model_name)
|
||||
.execute(&pool)
|
||||
.await
|
||||
.expect("global model should seed");
|
||||
|
||||
let repository = MysqlUsageWriteRepository::new(pool.clone());
|
||||
let barrier = std::sync::Arc::new(tokio::sync::Barrier::new(2));
|
||||
let mut tasks = Vec::new();
|
||||
for _ in 0..2 {
|
||||
let repository = repository.clone();
|
||||
let barrier = barrier.clone();
|
||||
let mut usage = sample_usage(
|
||||
&request_id,
|
||||
&user_id,
|
||||
&api_key_id,
|
||||
&provider_id,
|
||||
&provider_key_id,
|
||||
"completed",
|
||||
"pending",
|
||||
1_000,
|
||||
);
|
||||
usage.model.clone_from(&model_name);
|
||||
tasks.push(tokio::spawn(async move {
|
||||
barrier.wait().await;
|
||||
repository.upsert(usage).await
|
||||
}));
|
||||
}
|
||||
for task in tasks {
|
||||
task.await
|
||||
.expect("concurrent usage writer should join")
|
||||
.expect("concurrent usage should persist");
|
||||
}
|
||||
repository
|
||||
.flush_usage_counter_deltas(1_000)
|
||||
.await
|
||||
.expect("usage counter deltas should flush");
|
||||
|
||||
let api_key_requests: i64 =
|
||||
sqlx::query_scalar("SELECT total_requests FROM api_keys WHERE id = ?")
|
||||
.bind(&api_key_id)
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.expect("api key counter should load");
|
||||
let provider_key_requests: i64 =
|
||||
sqlx::query_scalar("SELECT request_count FROM provider_api_keys WHERE id = ?")
|
||||
.bind(&provider_key_id)
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.expect("provider key counter should load");
|
||||
let model_requests: i64 =
|
||||
sqlx::query_scalar("SELECT usage_count FROM global_models WHERE name = ?")
|
||||
.bind(&model_name)
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.expect("model counter should load");
|
||||
let outbox_rows: i64 =
|
||||
sqlx::query_scalar("SELECT COUNT(*) FROM usage_counter_deltas WHERE request_id = ?")
|
||||
.bind(&request_id)
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.expect("usage counter outbox should load");
|
||||
let routing_rows: i64 =
|
||||
sqlx::query_scalar("SELECT COUNT(*) FROM usage_routing_snapshots WHERE request_id = ?")
|
||||
.bind(&request_id)
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.expect("routing snapshot should count");
|
||||
let settlement_rows: i64 =
|
||||
sqlx::query_scalar("SELECT COUNT(*) FROM usage_settlement_snapshots WHERE request_id = ?")
|
||||
.bind(&request_id)
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.expect("settlement snapshot should count");
|
||||
assert_eq!(api_key_requests, 1);
|
||||
assert_eq!(provider_key_requests, 1);
|
||||
assert_eq!(model_requests, 1);
|
||||
assert_eq!(outbox_rows, 3);
|
||||
assert_eq!(routing_rows, 1);
|
||||
assert_eq!(settlement_rows, 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn mysql_usage_http_capture_round_trips_when_url_is_set() {
|
||||
let Some(database_url) = std::env::var("AETHER_TEST_MYSQL_URL")
|
||||
.ok()
|
||||
.filter(|value| !value.trim().is_empty())
|
||||
else {
|
||||
eprintln!("skipping mysql usage HTTP capture test because AETHER_TEST_MYSQL_URL is unset");
|
||||
return;
|
||||
};
|
||||
|
||||
let pool = sqlx::mysql::MySqlPoolOptions::new()
|
||||
.max_connections(1)
|
||||
.connect(&database_url)
|
||||
.await
|
||||
.expect("mysql test pool should connect");
|
||||
run_migrations(&pool)
|
||||
.await
|
||||
.expect("mysql migrations should run");
|
||||
let suffix = unique_suffix();
|
||||
let user_id = format!("capture-user-{suffix}");
|
||||
let api_key_id = format!("capture-api-key-{suffix}");
|
||||
let provider_id = format!("capture-provider-{suffix}");
|
||||
let provider_key_id = format!("capture-provider-key-{suffix}");
|
||||
let request_id = format!("capture-request-{suffix}");
|
||||
seed_stats_targets(&pool, &user_id, &api_key_id, &provider_id, &provider_key_id).await;
|
||||
let writer = MysqlUsageWriteRepository::new(pool.clone());
|
||||
|
||||
let mut rich = sample_usage(
|
||||
&request_id,
|
||||
&user_id,
|
||||
&api_key_id,
|
||||
&provider_id,
|
||||
&provider_key_id,
|
||||
"pending",
|
||||
"pending",
|
||||
1_000,
|
||||
);
|
||||
rich.request_headers = Some(serde_json::json!({"x-client": "one"}));
|
||||
rich.provider_request_headers = Some(serde_json::json!({"x-provider": "two"}));
|
||||
rich.request_body = Some(serde_json::json!({"request": true}));
|
||||
rich.provider_request_body = Some(serde_json::json!({"provider": true}));
|
||||
rich.request_body_state = Some(UsageBodyCaptureState::Inline);
|
||||
rich.provider_request_body_state = Some(UsageBodyCaptureState::Inline);
|
||||
let stored = writer
|
||||
.upsert(rich)
|
||||
.await
|
||||
.expect("MySQL canonical capture should upsert");
|
||||
assert_eq!(
|
||||
stored.request_headers,
|
||||
Some(serde_json::json!({"x-client": "one"}))
|
||||
);
|
||||
assert_eq!(
|
||||
stored.request_body,
|
||||
Some(serde_json::json!({"request": true}))
|
||||
);
|
||||
assert_eq!(
|
||||
stored.request_body_state,
|
||||
Some(UsageBodyCaptureState::Reference)
|
||||
);
|
||||
assert_eq!(
|
||||
stored.request_body_ref.as_deref(),
|
||||
Some(format!("usage://request/{request_id}/request_body").as_str())
|
||||
);
|
||||
let blob_count: i64 =
|
||||
sqlx::query_scalar("SELECT COUNT(*) FROM usage_body_blobs WHERE request_id = ?")
|
||||
.bind(&request_id)
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.expect("MySQL canonical blobs should count");
|
||||
assert_eq!(blob_count, 2);
|
||||
let legacy_body: Option<String> =
|
||||
sqlx::query_scalar("SELECT CAST(request_body AS CHAR) FROM `usage` WHERE request_id = ?")
|
||||
.bind(&request_id)
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.expect("legacy body should load");
|
||||
assert!(legacy_body.is_none());
|
||||
|
||||
let sparse = sample_usage(
|
||||
&request_id,
|
||||
&user_id,
|
||||
&api_key_id,
|
||||
&provider_id,
|
||||
&provider_key_id,
|
||||
"streaming",
|
||||
"pending",
|
||||
1_001,
|
||||
);
|
||||
let sparse_stored = writer
|
||||
.upsert(sparse)
|
||||
.await
|
||||
.expect("MySQL sparse capture should upsert");
|
||||
assert_eq!(sparse_stored.request_headers, stored.request_headers);
|
||||
assert_eq!(sparse_stored.request_body, stored.request_body);
|
||||
|
||||
let mut clear = sample_usage(
|
||||
&request_id,
|
||||
&user_id,
|
||||
&api_key_id,
|
||||
&provider_id,
|
||||
&provider_key_id,
|
||||
"streaming",
|
||||
"pending",
|
||||
1_002,
|
||||
);
|
||||
clear.request_body = Some(serde_json::json!({"residual": true}));
|
||||
clear.request_body_state = Some(UsageBodyCaptureState::None);
|
||||
let cleared = writer
|
||||
.upsert(clear)
|
||||
.await
|
||||
.expect("MySQL explicit none should clear");
|
||||
assert!(cleared.request_body.is_none());
|
||||
assert!(cleared.request_body_ref.is_none());
|
||||
assert_eq!(
|
||||
cleared.request_body_state,
|
||||
Some(UsageBodyCaptureState::None)
|
||||
);
|
||||
assert_eq!(cleared.provider_request_body, stored.provider_request_body);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn mysql_usage_canonical_snapshots_round_trip_when_url_is_set() {
|
||||
let Some(database_url) = std::env::var("AETHER_TEST_MYSQL_URL")
|
||||
.ok()
|
||||
.filter(|value| !value.trim().is_empty())
|
||||
else {
|
||||
eprintln!("skipping mysql usage snapshot test because AETHER_TEST_MYSQL_URL is unset");
|
||||
return;
|
||||
};
|
||||
|
||||
let pool = sqlx::mysql::MySqlPoolOptions::new()
|
||||
.max_connections(1)
|
||||
.connect(&database_url)
|
||||
.await
|
||||
.expect("mysql test pool should connect");
|
||||
run_migrations(&pool)
|
||||
.await
|
||||
.expect("mysql migrations should run");
|
||||
let suffix = unique_suffix();
|
||||
let user_id = format!("snapshot-user-{suffix}");
|
||||
let api_key_id = format!("snapshot-api-key-{suffix}");
|
||||
let provider_id = format!("snapshot-provider-{suffix}");
|
||||
let provider_key_id = format!("snapshot-provider-key-{suffix}");
|
||||
let request_id = format!("snapshot-request-{suffix}");
|
||||
seed_stats_targets(&pool, &user_id, &api_key_id, &provider_id, &provider_key_id).await;
|
||||
let writer = MysqlUsageWriteRepository::new(pool.clone());
|
||||
|
||||
let mut rich = sample_usage(
|
||||
&request_id,
|
||||
&user_id,
|
||||
&api_key_id,
|
||||
&provider_id,
|
||||
&provider_key_id,
|
||||
"streaming",
|
||||
"pending",
|
||||
1_000,
|
||||
);
|
||||
rich.candidate_id = None;
|
||||
rich.candidate_index = None;
|
||||
rich.key_name = None;
|
||||
rich.planner_kind = None;
|
||||
rich.route_family = None;
|
||||
rich.route_kind = None;
|
||||
rich.execution_path = None;
|
||||
rich.input_tokens = Some(100);
|
||||
rich.output_tokens = Some(20);
|
||||
rich.cache_creation_input_tokens = Some(10);
|
||||
rich.cache_creation_ephemeral_5m_input_tokens = Some(10);
|
||||
rich.cache_read_input_tokens = Some(30);
|
||||
rich.request_metadata = Some(serde_json::json!({
|
||||
"trace_id": "rich",
|
||||
"candidate_id": "candidate-canonical",
|
||||
"candidate_index": 2,
|
||||
"key_name": "key-canonical",
|
||||
"planner_kind": "fallback",
|
||||
"route_family": "chat",
|
||||
"route_kind": "remote",
|
||||
"execution_path": "converted",
|
||||
"billing_snapshot_schema_version": "v3",
|
||||
"billing_snapshot_status": "resolved",
|
||||
"rate_multiplier": 0.5,
|
||||
"input_price_per_1m": 1.1,
|
||||
"output_price_per_1m": 2.2,
|
||||
"billing_dimensions": {
|
||||
"input_tokens": 100,
|
||||
"effective_input_tokens": 60,
|
||||
"output_tokens": 20,
|
||||
"cache_creation_tokens": 10,
|
||||
"cache_read_tokens": 30,
|
||||
"total_input_context": 100
|
||||
},
|
||||
"settlement_snapshot": {
|
||||
"schema_version": "v3",
|
||||
"total_cost": 1.25,
|
||||
"actual_total_cost": 1.0,
|
||||
"pricing_snapshot": {"pricing_source": "catalog"},
|
||||
"billing_plan_snapshot": {"rule_id": "rule-1", "rule_version": "7"}
|
||||
}
|
||||
}));
|
||||
let stored = writer
|
||||
.upsert(rich)
|
||||
.await
|
||||
.expect("MySQL canonical snapshots should upsert");
|
||||
assert_eq!(stored.routing_candidate_id(), Some("candidate-canonical"));
|
||||
assert_eq!(stored.routing_candidate_index(), Some(2));
|
||||
assert_eq!(stored.provider_id.as_deref(), Some(provider_id.as_str()));
|
||||
assert_eq!(stored.output_tokens, 20);
|
||||
assert_eq!(stored.cache_creation_input_tokens, 10);
|
||||
assert_eq!(stored.cache_read_input_tokens, 30);
|
||||
assert_eq!(stored.total_tokens, 120);
|
||||
assert_eq!(stored.total_cost_usd, 1.25);
|
||||
assert_eq!(stored.actual_total_cost_usd, 1.0);
|
||||
assert_eq!(stored.settlement_rate_multiplier(), Some(0.5));
|
||||
assert_eq!(stored.settlement_input_price_per_1m(), Some(1.1));
|
||||
|
||||
let mut sparse = sample_usage(
|
||||
&request_id,
|
||||
&user_id,
|
||||
&api_key_id,
|
||||
&provider_id,
|
||||
&provider_key_id,
|
||||
"streaming",
|
||||
"pending",
|
||||
1_001,
|
||||
);
|
||||
sparse.provider_id = None;
|
||||
sparse.provider_endpoint_id = None;
|
||||
sparse.provider_api_key_id = None;
|
||||
sparse.has_format_conversion = None;
|
||||
sparse.candidate_id = None;
|
||||
sparse.candidate_index = None;
|
||||
sparse.key_name = None;
|
||||
sparse.planner_kind = None;
|
||||
sparse.route_family = None;
|
||||
sparse.route_kind = None;
|
||||
sparse.execution_path = None;
|
||||
sparse.local_execution_runtime_miss_reason = None;
|
||||
sparse.input_tokens = None;
|
||||
sparse.output_tokens = None;
|
||||
sparse.total_tokens = None;
|
||||
sparse.cache_creation_input_tokens = None;
|
||||
sparse.cache_creation_ephemeral_5m_input_tokens = None;
|
||||
sparse.cache_creation_ephemeral_1h_input_tokens = None;
|
||||
sparse.cache_read_input_tokens = None;
|
||||
sparse.cache_creation_cost_usd = None;
|
||||
sparse.cache_read_cost_usd = None;
|
||||
sparse.output_price_per_1m = None;
|
||||
sparse.total_cost_usd = None;
|
||||
sparse.actual_total_cost_usd = None;
|
||||
sparse.request_metadata = Some(serde_json::json!({"trace_id": "sparse"}));
|
||||
let sparse_stored = writer
|
||||
.upsert(sparse)
|
||||
.await
|
||||
.expect("MySQL sparse snapshots should merge");
|
||||
assert_eq!(
|
||||
sparse_stored.routing_candidate_id(),
|
||||
Some("candidate-canonical")
|
||||
);
|
||||
assert_eq!(
|
||||
sparse_stored.provider_id.as_deref(),
|
||||
Some(provider_id.as_str())
|
||||
);
|
||||
assert_eq!(sparse_stored.output_tokens, 20);
|
||||
assert_eq!(sparse_stored.total_tokens, 120);
|
||||
assert_eq!(sparse_stored.total_cost_usd, 1.25);
|
||||
assert_eq!(sparse_stored.settlement_rate_multiplier(), Some(0.5));
|
||||
|
||||
let mut terminal = sample_usage(
|
||||
&request_id,
|
||||
&user_id,
|
||||
&api_key_id,
|
||||
&provider_id,
|
||||
&provider_key_id,
|
||||
"completed",
|
||||
"settled",
|
||||
1_002,
|
||||
);
|
||||
terminal.provider_id = None;
|
||||
terminal.provider_endpoint_id = None;
|
||||
terminal.provider_api_key_id = None;
|
||||
terminal.has_format_conversion = None;
|
||||
terminal.candidate_id = None;
|
||||
terminal.candidate_index = None;
|
||||
terminal.key_name = None;
|
||||
terminal.planner_kind = None;
|
||||
terminal.route_family = None;
|
||||
terminal.route_kind = None;
|
||||
terminal.execution_path = None;
|
||||
terminal.local_execution_runtime_miss_reason = None;
|
||||
terminal.input_tokens = None;
|
||||
terminal.output_tokens = None;
|
||||
terminal.total_tokens = None;
|
||||
terminal.cache_creation_input_tokens = None;
|
||||
terminal.cache_creation_ephemeral_5m_input_tokens = None;
|
||||
terminal.cache_creation_ephemeral_1h_input_tokens = None;
|
||||
terminal.cache_read_input_tokens = None;
|
||||
terminal.cache_creation_cost_usd = None;
|
||||
terminal.cache_read_cost_usd = None;
|
||||
terminal.output_price_per_1m = None;
|
||||
terminal.total_cost_usd = None;
|
||||
terminal.actual_total_cost_usd = None;
|
||||
terminal.request_metadata = Some(serde_json::json!({"trace_id": "terminal"}));
|
||||
let terminal_stored = writer
|
||||
.upsert(terminal)
|
||||
.await
|
||||
.expect("MySQL terminal snapshots should replace");
|
||||
assert_eq!(terminal_stored.status, "completed");
|
||||
assert_eq!(terminal_stored.billing_status, "settled");
|
||||
assert!(terminal_stored.candidate_id.is_none());
|
||||
assert!(terminal_stored.provider_id.is_none());
|
||||
assert_eq!(terminal_stored.total_tokens, 0);
|
||||
assert_eq!(terminal_stored.total_cost_usd, 0.0);
|
||||
assert_eq!(terminal_stored.settlement_rate_multiplier(), None);
|
||||
|
||||
let mut late = sample_usage(
|
||||
&request_id,
|
||||
&user_id,
|
||||
&api_key_id,
|
||||
&provider_id,
|
||||
&provider_key_id,
|
||||
"pending",
|
||||
"pending",
|
||||
1_003,
|
||||
);
|
||||
late.candidate_id = Some("late-candidate".to_string());
|
||||
late.request_metadata = Some(serde_json::json!({
|
||||
"trace_id": "late",
|
||||
"rate_multiplier": 9,
|
||||
"settlement_snapshot": {"schema_version": "late", "total_cost": 99}
|
||||
}));
|
||||
let after_late = writer
|
||||
.upsert(late)
|
||||
.await
|
||||
.expect("late MySQL pending usage should return terminal record");
|
||||
assert_eq!(after_late.status, "completed");
|
||||
assert_eq!(after_late.billing_status, "settled");
|
||||
assert!(after_late.candidate_id.is_none());
|
||||
assert_eq!(after_late.total_cost_usd, 0.0);
|
||||
|
||||
let snapshot_counts: (i64, i64) = sqlx::query_as(
|
||||
r#"
|
||||
SELECT
|
||||
(SELECT COUNT(*) FROM usage_routing_snapshots WHERE request_id = ?),
|
||||
(SELECT COUNT(*) FROM usage_settlement_snapshots WHERE request_id = ?)
|
||||
"#,
|
||||
)
|
||||
.bind(&request_id)
|
||||
.bind(&request_id)
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.expect("MySQL canonical snapshots should count");
|
||||
assert_eq!(snapshot_counts, (1, 1));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -211,12 +847,14 @@ async fn mysql_usage_read_repository_reads_usage_contract_views_when_url_is_set(
|
||||
let api_key_id = format!("api-key-read-{suffix}");
|
||||
let provider_id = format!("provider-read-{suffix}");
|
||||
let provider_key_id = format!("provider-key-read-{suffix}");
|
||||
let first_request_id = format!("request-read-1-{suffix}");
|
||||
let second_request_id = format!("request-read-2-{suffix}");
|
||||
seed_stats_targets(&pool, &user_id, &api_key_id, &provider_id, &provider_key_id).await;
|
||||
|
||||
let writer = MysqlUsageWriteRepository::new(pool.clone());
|
||||
writer
|
||||
.upsert(sample_usage(
|
||||
&format!("request-read-1-{suffix}"),
|
||||
&first_request_id,
|
||||
&user_id,
|
||||
&api_key_id,
|
||||
&provider_id,
|
||||
@@ -229,7 +867,7 @@ async fn mysql_usage_read_repository_reads_usage_contract_views_when_url_is_set(
|
||||
.expect("usage should upsert");
|
||||
writer
|
||||
.upsert(sample_usage(
|
||||
&format!("request-read-2-{suffix}"),
|
||||
&second_request_id,
|
||||
&user_id,
|
||||
&api_key_id,
|
||||
&provider_id,
|
||||
@@ -241,17 +879,33 @@ async fn mysql_usage_read_repository_reads_usage_contract_views_when_url_is_set(
|
||||
.await
|
||||
.expect("usage should upsert");
|
||||
|
||||
sqlx::query("UPDATE `usage` SET username = ?, api_key_name = ? WHERE request_id = ?")
|
||||
.bind("legacy-user")
|
||||
.bind("legacy-key")
|
||||
.bind(&first_request_id)
|
||||
.execute(&pool)
|
||||
.await
|
||||
.expect("legacy usage names should update");
|
||||
|
||||
let reader = MysqlUsageStorage::new(pool);
|
||||
let records = reader
|
||||
.load_usage_records()
|
||||
.list_usage_audits(&UsageAuditListQuery {
|
||||
created_from_unix_secs: Some(900),
|
||||
created_until_unix_secs: Some(1_100),
|
||||
user_id: Some(user_id.clone()),
|
||||
newest_first: true,
|
||||
..UsageAuditListQuery::default()
|
||||
})
|
||||
.await
|
||||
.expect("usage records should load");
|
||||
let loaded = records
|
||||
.iter()
|
||||
.find(|item| item.request_id == format!("request-read-1-{suffix}"))
|
||||
.find(|item| item.request_id == first_request_id)
|
||||
.expect("usage should exist");
|
||||
assert_eq!(loaded.total_tokens, 7);
|
||||
assert_eq!(loaded.total_tokens, 5);
|
||||
assert_eq!(loaded.billing_status, "settled");
|
||||
assert_eq!(loaded.username.as_deref(), Some("legacy-user"));
|
||||
assert_eq!(loaded.api_key_name.as_deref(), Some("legacy-key"));
|
||||
assert_eq!(
|
||||
records
|
||||
.iter()
|
||||
@@ -261,6 +915,164 @@ async fn mysql_usage_read_repository_reads_usage_contract_views_when_url_is_set(
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn mysql_usage_cleanup_executes_when_url_is_set() {
|
||||
let Some(database_url) = std::env::var("AETHER_TEST_MYSQL_URL")
|
||||
.ok()
|
||||
.filter(|value| !value.trim().is_empty())
|
||||
else {
|
||||
eprintln!("skipping mysql usage cleanup test because AETHER_TEST_MYSQL_URL is unset");
|
||||
return;
|
||||
};
|
||||
|
||||
let pool = sqlx::mysql::MySqlPoolOptions::new()
|
||||
.max_connections(1)
|
||||
.connect(&database_url)
|
||||
.await
|
||||
.expect("mysql test pool should connect");
|
||||
run_migrations(&pool)
|
||||
.await
|
||||
.expect("mysql migrations should run");
|
||||
|
||||
let suffix = unique_suffix();
|
||||
let user_id = format!("cleanup-user-{suffix}");
|
||||
let api_key_id = format!("cleanup-api-key-{suffix}");
|
||||
let provider_id = format!("cleanup-provider-{suffix}");
|
||||
let provider_key_id = format!("cleanup-provider-key-{suffix}");
|
||||
let request_id = format!("cleanup-request-{suffix}");
|
||||
seed_stats_targets(&pool, &user_id, &api_key_id, &provider_id, &provider_key_id).await;
|
||||
let repository = MysqlUsageWriteRepository::new(pool.clone());
|
||||
repository
|
||||
.upsert(sample_usage(
|
||||
&request_id,
|
||||
&user_id,
|
||||
&api_key_id,
|
||||
&provider_id,
|
||||
&provider_key_id,
|
||||
"completed",
|
||||
"settled",
|
||||
10,
|
||||
))
|
||||
.await
|
||||
.expect("cleanup usage should seed");
|
||||
sqlx::query(
|
||||
"UPDATE `usage` SET request_headers = '{\"keep\":true}', request_body = '{\"body\":true}' WHERE request_id = ?",
|
||||
)
|
||||
.bind(&request_id)
|
||||
.execute(&pool)
|
||||
.await
|
||||
.expect("cleanup fields should seed");
|
||||
|
||||
let window = UsageCleanupWindow {
|
||||
detail_cutoff: DateTime::from_timestamp(20, 0).expect("valid detail cutoff"),
|
||||
compressed_cutoff: DateTime::from_timestamp(5, 0).expect("valid compressed cutoff"),
|
||||
header_cutoff: DateTime::from_timestamp(20, 0).expect("valid header cutoff"),
|
||||
log_cutoff: DateTime::from_timestamp(5, 0).expect("valid log cutoff"),
|
||||
};
|
||||
let detail_only = UsageCleanupTargets {
|
||||
detail_body: true,
|
||||
compressed_body: false,
|
||||
headers: false,
|
||||
records: false,
|
||||
expired_keys: false,
|
||||
};
|
||||
let preview = repository
|
||||
.preview_usage_cleanup(&window, detail_only, UsageCleanupExecutionMode::Policy)
|
||||
.await
|
||||
.expect("MySQL cleanup preview should load");
|
||||
assert!(preview.detail >= 1);
|
||||
let summary = repository
|
||||
.cleanup_usage(
|
||||
&window,
|
||||
1,
|
||||
false,
|
||||
detail_only,
|
||||
UsageCleanupExecutionMode::Policy,
|
||||
)
|
||||
.await
|
||||
.expect("MySQL detail cleanup should succeed");
|
||||
assert!(summary.body_externalized >= 1);
|
||||
let body_ref: String =
|
||||
sqlx::query_scalar("SELECT request_body_ref FROM usage_http_audits WHERE request_id = ?")
|
||||
.bind(&request_id)
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.expect("externalized body ref should load");
|
||||
assert_eq!(
|
||||
body_ref,
|
||||
format!("usage://request/{request_id}/request_body")
|
||||
);
|
||||
|
||||
let headers_only = UsageCleanupTargets {
|
||||
detail_body: false,
|
||||
compressed_body: false,
|
||||
headers: true,
|
||||
records: false,
|
||||
expired_keys: false,
|
||||
};
|
||||
let summary = repository
|
||||
.cleanup_usage(
|
||||
&window,
|
||||
1,
|
||||
false,
|
||||
headers_only,
|
||||
UsageCleanupExecutionMode::Policy,
|
||||
)
|
||||
.await
|
||||
.expect("MySQL header cleanup should succeed");
|
||||
assert!(summary.header_cleaned >= 1);
|
||||
let request_headers: Option<String> =
|
||||
sqlx::query_scalar("SELECT request_headers FROM `usage` WHERE request_id = ?")
|
||||
.bind(&request_id)
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.expect("cleaned headers should load");
|
||||
assert!(request_headers.is_none());
|
||||
|
||||
let before_now_request_id = format!("cleanup-before-now-{suffix}");
|
||||
repository
|
||||
.upsert(sample_usage(
|
||||
&before_now_request_id,
|
||||
&user_id,
|
||||
&api_key_id,
|
||||
&provider_id,
|
||||
&provider_key_id,
|
||||
"completed",
|
||||
"settled",
|
||||
10,
|
||||
))
|
||||
.await
|
||||
.expect("before-now usage should seed");
|
||||
sqlx::query(
|
||||
"UPDATE `usage` SET request_body = '{\"raw\":true}', request_body_compressed = ? WHERE request_id = ?",
|
||||
)
|
||||
.bind(vec![0x1f_u8, 0x8b])
|
||||
.bind(&before_now_request_id)
|
||||
.execute(&pool)
|
||||
.await
|
||||
.expect("before-now fields should seed");
|
||||
let summary = repository
|
||||
.cleanup_usage(
|
||||
&window,
|
||||
1,
|
||||
false,
|
||||
UsageCleanupTargets::body_targets(),
|
||||
UsageCleanupExecutionMode::BeforeNowBodyFields,
|
||||
)
|
||||
.await
|
||||
.expect("MySQL before-now cleanup should succeed");
|
||||
assert!(summary.body_externalized >= 1);
|
||||
assert!(summary.body_cleaned >= 1);
|
||||
let body_fields: (Option<String>, Option<Vec<u8>>) = sqlx::query_as(
|
||||
"SELECT CAST(request_body AS CHAR), request_body_compressed FROM `usage` WHERE request_id = ?",
|
||||
)
|
||||
.bind(&before_now_request_id)
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.expect("before-now fields should load");
|
||||
assert_eq!(body_fields, (None, None));
|
||||
}
|
||||
|
||||
async fn seed_stats_targets(
|
||||
pool: &sqlx::MySqlPool,
|
||||
user_id: &str,
|
||||
|
||||
@@ -99,32 +99,6 @@ impl MysqlVideoTaskRepository {
|
||||
.map_sql_err()?;
|
||||
row.as_ref().map(map_video_task_row).transpose()
|
||||
}
|
||||
|
||||
async fn reload_ids(&self, ids: &[String]) -> Result<Vec<StoredVideoTask>, DataLayerError> {
|
||||
if ids.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let mut builder = QueryBuilder::<MySql>::new(VIDEO_TASK_COLUMNS);
|
||||
builder.push(" WHERE id IN (");
|
||||
{
|
||||
let mut separated = builder.separated(", ");
|
||||
for id in ids {
|
||||
separated.push_bind(id);
|
||||
}
|
||||
}
|
||||
builder.push(")");
|
||||
let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?;
|
||||
let mut tasks = rows
|
||||
.iter()
|
||||
.map(map_video_task_row)
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
tasks.sort_by(|left, right| {
|
||||
left.next_poll_at_unix_secs
|
||||
.cmp(&right.next_poll_at_unix_secs)
|
||||
.then_with(|| left.updated_at_unix_secs.cmp(&right.updated_at_unix_secs))
|
||||
});
|
||||
Ok(tasks)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
@@ -320,20 +294,77 @@ impl VideoTaskWriteRepository for MysqlVideoTaskRepository {
|
||||
if limit == 0 {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let due = self.list_due(now_unix_secs, limit).await?;
|
||||
let ids = due.iter().map(|task| task.id.clone()).collect::<Vec<_>>();
|
||||
for id in &ids {
|
||||
sqlx::query(
|
||||
"UPDATE video_tasks SET next_poll_at = ?, updated_at = GREATEST(updated_at, ?) WHERE id = ?",
|
||||
)
|
||||
.bind(u64_to_i64(claim_until_unix_secs, "video task claim_until")?)
|
||||
.bind(u64_to_i64(now_unix_secs, "video task now")?)
|
||||
.bind(id)
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
let now = u64_to_i64(now_unix_secs, "video task now")?;
|
||||
let claim_until = u64_to_i64(claim_until_unix_secs, "video task claim_until")?;
|
||||
let limit = limit_i64(limit, "due video task claim limit")?;
|
||||
let mut tx = self.pool.begin().await.map_sql_err()?;
|
||||
let due_rows = sqlx::query(
|
||||
r#"
|
||||
SELECT id
|
||||
FROM video_tasks
|
||||
WHERE status IN ('submitted', 'queued', 'processing')
|
||||
AND next_poll_at IS NOT NULL
|
||||
AND next_poll_at <= ?
|
||||
AND poll_count < max_poll_count
|
||||
ORDER BY next_poll_at ASC, updated_at ASC
|
||||
LIMIT ?
|
||||
FOR UPDATE SKIP LOCKED
|
||||
"#,
|
||||
)
|
||||
.bind(now)
|
||||
.bind(limit)
|
||||
.fetch_all(&mut *tx)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
let ids = due_rows
|
||||
.iter()
|
||||
.map(|row| row.try_get::<String, _>("id").map_sql_err())
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
if ids.is_empty() {
|
||||
tx.commit().await.map_sql_err()?;
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
self.reload_ids(&ids).await
|
||||
|
||||
let mut update = QueryBuilder::<MySql>::new("UPDATE video_tasks SET next_poll_at = ");
|
||||
update
|
||||
.push_bind(claim_until)
|
||||
.push(", updated_at = GREATEST(updated_at, ")
|
||||
.push_bind(now)
|
||||
.push(") WHERE id IN (");
|
||||
{
|
||||
let mut separated = update.separated(", ");
|
||||
for id in &ids {
|
||||
separated.push_bind(id);
|
||||
}
|
||||
}
|
||||
update
|
||||
.push(") AND status IN ('submitted', 'queued', 'processing')")
|
||||
.push(" AND next_poll_at IS NOT NULL AND next_poll_at <= ")
|
||||
.push_bind(now)
|
||||
.push(" AND poll_count < max_poll_count");
|
||||
update.build().execute(&mut *tx).await.map_sql_err()?;
|
||||
|
||||
let mut reload = QueryBuilder::<MySql>::new(VIDEO_TASK_COLUMNS);
|
||||
reload.push(" WHERE id IN (");
|
||||
{
|
||||
let mut separated = reload.separated(", ");
|
||||
for id in &ids {
|
||||
separated.push_bind(id);
|
||||
}
|
||||
}
|
||||
reload.push(")");
|
||||
let rows = reload.build().fetch_all(&mut *tx).await.map_sql_err()?;
|
||||
let mut tasks = rows
|
||||
.iter()
|
||||
.map(map_video_task_row)
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
tasks.sort_by(|left, right| {
|
||||
left.next_poll_at_unix_secs
|
||||
.cmp(&right.next_poll_at_unix_secs)
|
||||
.then_with(|| left.updated_at_unix_secs.cmp(&right.updated_at_unix_secs))
|
||||
});
|
||||
tx.commit().await.map_sql_err()?;
|
||||
Ok(tasks)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -676,6 +707,11 @@ fn optional_u32_to_i32(value: Option<u32>, name: &str) -> Result<Option<i32>, Da
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::MysqlVideoTaskRepository;
|
||||
use crate::run_migrations;
|
||||
use aether_data_contracts::repository::video_tasks::{
|
||||
UpsertVideoTask, VideoTaskStatus, VideoTaskWriteRepository,
|
||||
};
|
||||
use std::sync::Arc;
|
||||
|
||||
#[tokio::test]
|
||||
async fn repository_builds_from_lazy_pool() {
|
||||
@@ -687,4 +723,116 @@ mod tests {
|
||||
|
||||
let _repository = MysqlVideoTaskRepository::new(pool);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn mysql_claim_due_does_not_return_one_task_to_multiple_workers_when_url_is_set() {
|
||||
let Some(database_url) = std::env::var("AETHER_TEST_MYSQL_URL")
|
||||
.ok()
|
||||
.filter(|value| !value.trim().is_empty())
|
||||
else {
|
||||
eprintln!(
|
||||
"skipping mysql claim_due concurrency test because AETHER_TEST_MYSQL_URL is unset"
|
||||
);
|
||||
return;
|
||||
};
|
||||
|
||||
let pool = sqlx::mysql::MySqlPoolOptions::new()
|
||||
.max_connections(4)
|
||||
.connect(&database_url)
|
||||
.await
|
||||
.expect("mysql test pool should connect");
|
||||
run_migrations(&pool)
|
||||
.await
|
||||
.expect("mysql migrations should run");
|
||||
|
||||
let suffix = uuid::Uuid::new_v4().simple().to_string();
|
||||
let task_id = format!("claim-{}", &suffix[..20]);
|
||||
let repository = MysqlVideoTaskRepository::new(pool.clone());
|
||||
repository
|
||||
.upsert(claimable_task(&task_id))
|
||||
.await
|
||||
.expect("claimable mysql task should insert");
|
||||
|
||||
let barrier = Arc::new(tokio::sync::Barrier::new(3));
|
||||
let first_repository = repository.clone();
|
||||
let first_barrier = barrier.clone();
|
||||
let first = tokio::spawn(async move {
|
||||
first_barrier.wait().await;
|
||||
first_repository.claim_due(100, 130, 1).await
|
||||
});
|
||||
let second_repository = repository.clone();
|
||||
let second_barrier = barrier.clone();
|
||||
let second = tokio::spawn(async move {
|
||||
second_barrier.wait().await;
|
||||
second_repository.claim_due(100, 130, 1).await
|
||||
});
|
||||
barrier.wait().await;
|
||||
|
||||
let first_result = first.await;
|
||||
let second_result = second.await;
|
||||
let followup_result = repository.claim_due(100, 130, 1).await;
|
||||
let cleanup_result = sqlx::query("DELETE FROM video_tasks WHERE id = ?")
|
||||
.bind(&task_id)
|
||||
.execute(&pool)
|
||||
.await;
|
||||
|
||||
let first_claimed = first_result
|
||||
.expect("first mysql claim worker should join")
|
||||
.expect("first mysql claim should execute");
|
||||
let second_claimed = second_result
|
||||
.expect("second mysql claim worker should join")
|
||||
.expect("second mysql claim should execute");
|
||||
let followup_claimed = followup_result.expect("follow-up mysql claim should execute");
|
||||
cleanup_result.expect("mysql claim fixture should clean up");
|
||||
|
||||
let claimed_ids = first_claimed
|
||||
.iter()
|
||||
.chain(&second_claimed)
|
||||
.map(|task| task.id.as_str())
|
||||
.collect::<Vec<_>>();
|
||||
assert_eq!(claimed_ids, vec![task_id.as_str()]);
|
||||
assert!(followup_claimed.is_empty());
|
||||
}
|
||||
|
||||
fn claimable_task(id: &str) -> UpsertVideoTask {
|
||||
UpsertVideoTask {
|
||||
id: id.to_string(),
|
||||
short_id: None,
|
||||
request_id: format!("request-{id}"),
|
||||
user_id: None,
|
||||
api_key_id: None,
|
||||
username: None,
|
||||
api_key_name: None,
|
||||
external_task_id: Some(format!("external-{id}")),
|
||||
provider_id: None,
|
||||
endpoint_id: None,
|
||||
key_id: None,
|
||||
client_api_format: Some("openai:video".to_string()),
|
||||
provider_api_format: Some("openai:video".to_string()),
|
||||
format_converted: false,
|
||||
model: Some("sora-2".to_string()),
|
||||
prompt: Some("claim test".to_string()),
|
||||
original_request_body: None,
|
||||
duration_seconds: None,
|
||||
resolution: None,
|
||||
aspect_ratio: None,
|
||||
size: None,
|
||||
status: VideoTaskStatus::Submitted,
|
||||
progress_percent: 0,
|
||||
progress_message: None,
|
||||
retry_count: 0,
|
||||
poll_interval_seconds: 10,
|
||||
next_poll_at_unix_secs: Some(100),
|
||||
poll_count: 0,
|
||||
max_poll_count: 360,
|
||||
created_at_unix_ms: 90,
|
||||
submitted_at_unix_secs: Some(90),
|
||||
completed_at_unix_secs: None,
|
||||
updated_at_unix_secs: 100,
|
||||
error_code: None,
|
||||
error_message: None,
|
||||
video_url: None,
|
||||
request_metadata: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,10 +1,119 @@
|
||||
use super::MysqlWalletReadRepository;
|
||||
use super::{
|
||||
admin_payment_callback_list_builder, admin_payment_order_list_builder,
|
||||
admin_redeem_batch_list_builder, admin_redeem_code_list_builder,
|
||||
admin_wallet_ledger_list_builder, admin_wallet_list_builder,
|
||||
admin_wallet_refund_request_list_builder, wallets_by_owner_ids_builder,
|
||||
MysqlWalletReadRepository,
|
||||
};
|
||||
use crate::run_migrations;
|
||||
use aether_data_contracts::repository::wallet::{
|
||||
AdminPaymentOrderListQuery, AdminRedeemCodeListQuery, AdminWalletListQuery, WalletLookupKey,
|
||||
WalletReadRepository,
|
||||
AdminPaymentOrderListQuery, AdminRedeemCodeBatchListQuery, AdminRedeemCodeListQuery,
|
||||
AdminWalletLedgerQuery, AdminWalletListQuery, AdminWalletRefundRequestListQuery,
|
||||
WalletLookupKey, WalletReadRepository,
|
||||
};
|
||||
|
||||
#[test]
|
||||
fn mysql_wallet_builders_bind_filters_and_page_bounds() {
|
||||
let injected_status = "active' OR 1 = 1 --".to_string();
|
||||
let query = AdminWalletListQuery {
|
||||
status: Some(injected_status.clone()),
|
||||
owner_type: Some("user".to_string()),
|
||||
limit: 7,
|
||||
offset: 3,
|
||||
};
|
||||
let builder = admin_wallet_list_builder(&query, 7, 3);
|
||||
let sql = compact_sql(builder.sql());
|
||||
|
||||
assert!(sql.contains("AND w.status = ?"));
|
||||
assert!(sql.contains("AND w.user_id IS NOT NULL"));
|
||||
assert!(sql.contains("ORDER BY w.updated_at DESC, w.id DESC LIMIT ? OFFSET ?"));
|
||||
assert!(!sql.contains(&injected_status));
|
||||
assert_eq!(sql.matches('?').count(), 3);
|
||||
|
||||
let ids = vec!["user-1".to_string(), "user-2' OR 1 = 1 --".to_string()];
|
||||
let builder = wallets_by_owner_ids_builder("user_id", &ids);
|
||||
let sql = compact_sql(builder.sql());
|
||||
assert!(sql.contains("WHERE user_id IN (?, ?) ORDER BY id ASC"));
|
||||
assert!(!sql.contains(&ids[0]));
|
||||
assert!(!sql.contains(&ids[1]));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mysql_wallet_admin_builders_cover_filters_ordering_and_mapping_columns() {
|
||||
let ledger_query = AdminWalletLedgerQuery {
|
||||
category: Some("credit".to_string()),
|
||||
reason_code: Some("manual_adjustment".to_string()),
|
||||
owner_type: Some("api_key".to_string()),
|
||||
limit: 5,
|
||||
offset: 2,
|
||||
};
|
||||
let ledger_sql = compact_sql(admin_wallet_ledger_list_builder(&ledger_query, 5, 2).sql());
|
||||
assert!(ledger_sql.contains("tx.category = ?"));
|
||||
assert!(ledger_sql.contains("tx.reason_code = ?"));
|
||||
assert!(ledger_sql.contains("w.api_key_id IS NOT NULL"));
|
||||
assert!(ledger_sql.contains("wallet_users.username AS wallet_user_name"));
|
||||
assert!(ledger_sql.contains("ORDER BY tx.created_at DESC, tx.id DESC LIMIT ? OFFSET ?"));
|
||||
|
||||
let refund_query = AdminWalletRefundRequestListQuery {
|
||||
status: Some("pending_approval".to_string()),
|
||||
limit: 4,
|
||||
offset: 1,
|
||||
};
|
||||
let refund_sql =
|
||||
compact_sql(admin_wallet_refund_request_list_builder(&refund_query, 4, 1).sql());
|
||||
assert!(refund_sql.contains("WHERE w.user_id IS NOT NULL"));
|
||||
assert!(refund_sql.contains("rr.status = ?"));
|
||||
assert!(refund_sql.contains("rr.payout_proof"));
|
||||
assert!(refund_sql.contains("ORDER BY rr.created_at DESC, rr.id DESC LIMIT ? OFFSET ?"));
|
||||
|
||||
let order_query = AdminPaymentOrderListQuery {
|
||||
status: Some("expired".to_string()),
|
||||
payment_method: Some("card".to_string()),
|
||||
limit: 8,
|
||||
offset: 6,
|
||||
};
|
||||
let order_sql = compact_sql(admin_payment_order_list_builder(&order_query, 100, 8, 6).sql());
|
||||
assert!(order_sql.contains("payment_method = ?"));
|
||||
assert!(order_sql.contains(
|
||||
"CASE WHEN status = 'pending' AND expires_at IS NOT NULL AND expires_at < ? THEN 'expired'"
|
||||
));
|
||||
assert!(order_sql.contains("ORDER BY created_at DESC, id DESC LIMIT ? OFFSET ?"));
|
||||
|
||||
let callback_sql =
|
||||
compact_sql(admin_payment_callback_list_builder(Some("redeem_code"), 9, 4).sql());
|
||||
assert!(callback_sql.contains("payment_method = ?"));
|
||||
assert!(callback_sql.contains("ORDER BY created_at DESC, id DESC LIMIT ? OFFSET ?"));
|
||||
|
||||
let batch_query = AdminRedeemCodeBatchListQuery {
|
||||
status: Some("active".to_string()),
|
||||
limit: 2,
|
||||
offset: 1,
|
||||
};
|
||||
let batch_sql = compact_sql(admin_redeem_batch_list_builder(&batch_query, 2, 1).sql());
|
||||
assert!(batch_sql.contains("batches.status = ?"));
|
||||
assert!(batch_sql.contains("AS redeemed_count"));
|
||||
assert!(
|
||||
batch_sql.contains("ORDER BY batches.created_at DESC, batches.id DESC LIMIT ? OFFSET ?")
|
||||
);
|
||||
|
||||
let code_query = AdminRedeemCodeListQuery {
|
||||
batch_id: "batch-1".to_string(),
|
||||
status: Some("redeemed".to_string()),
|
||||
limit: 3,
|
||||
offset: 2,
|
||||
};
|
||||
let code_sql = compact_sql(admin_redeem_code_list_builder(&code_query, 3, 2).sql());
|
||||
assert!(code_sql.contains("codes.batch_id = ?"));
|
||||
assert!(code_sql.contains("codes.status = ?"));
|
||||
assert!(code_sql.contains("orders.order_no AS redeemed_order_no"));
|
||||
assert!(code_sql.contains("ORDER BY codes.created_at DESC, codes.id DESC LIMIT ? OFFSET ?"));
|
||||
assert!(!code_sql.contains("batch-1"));
|
||||
}
|
||||
|
||||
fn compact_sql(sql: &str) -> String {
|
||||
sql.split_whitespace().collect::<Vec<_>>().join(" ")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn mysql_wallet_read_repository_reads_wallet_contract_views() {
|
||||
let Some(database_url) = std::env::var("AETHER_TEST_MYSQL_URL")
|
||||
@@ -34,22 +143,93 @@ async fn mysql_wallet_read_repository_reads_wallet_contract_views() {
|
||||
.expect("wallet should exist");
|
||||
assert_eq!(wallet.total_adjusted, 3.0);
|
||||
|
||||
let user_wallets = repository
|
||||
.list_wallets_by_user_ids(&[
|
||||
"user-2".to_string(),
|
||||
"missing-user".to_string(),
|
||||
"user-1".to_string(),
|
||||
])
|
||||
.await
|
||||
.expect("user wallets should list");
|
||||
assert_eq!(
|
||||
user_wallets
|
||||
.iter()
|
||||
.map(|wallet| wallet.id.as_str())
|
||||
.collect::<Vec<_>>(),
|
||||
vec!["wallet-1", "wallet-2"]
|
||||
);
|
||||
let api_key_wallets = repository
|
||||
.list_wallets_by_api_key_ids(&["api-key-1".to_string()])
|
||||
.await
|
||||
.expect("api key wallets should list");
|
||||
assert_eq!(api_key_wallets[0].id, "wallet-api-1");
|
||||
assert!(repository
|
||||
.list_wallets_by_user_ids(&[])
|
||||
.await
|
||||
.expect("empty user ids should short circuit")
|
||||
.is_empty());
|
||||
|
||||
let page = repository
|
||||
.list_admin_wallets(&AdminWalletListQuery {
|
||||
status: Some("active".to_string()),
|
||||
owner_type: Some("user".to_string()),
|
||||
limit: 1,
|
||||
offset: 1,
|
||||
})
|
||||
.await
|
||||
.expect("admin wallets should list");
|
||||
assert_eq!(page.total, 2);
|
||||
assert_eq!(page.items.len(), 1);
|
||||
let wallet_item = &page.items[0];
|
||||
assert_eq!(wallet_item.id, "wallet-1");
|
||||
assert_eq!(wallet_item.total_adjusted, 3.0);
|
||||
|
||||
let unknown_owner = repository
|
||||
.list_admin_wallets(&AdminWalletListQuery {
|
||||
status: Some("active".to_string()),
|
||||
owner_type: Some("unknown".to_string()),
|
||||
limit: 10,
|
||||
offset: 0,
|
||||
})
|
||||
.await
|
||||
.expect("admin wallets should list");
|
||||
let wallet_item = page
|
||||
.items
|
||||
.iter()
|
||||
.find(|item| item.id == "wallet-1")
|
||||
.expect("seeded wallet should be listed");
|
||||
assert!(page.total >= 1);
|
||||
assert_eq!(wallet_item.total_adjusted, 3.0);
|
||||
.expect("unknown owner type should preserve snapshot semantics");
|
||||
assert_eq!(unknown_owner.total, 3);
|
||||
|
||||
let ledger = repository
|
||||
.list_admin_wallet_ledger(&AdminWalletLedgerQuery {
|
||||
category: Some("credit".to_string()),
|
||||
reason_code: Some("manual_adjustment".to_string()),
|
||||
owner_type: Some("user".to_string()),
|
||||
limit: 1,
|
||||
offset: 1,
|
||||
})
|
||||
.await
|
||||
.expect("admin ledger should list");
|
||||
assert_eq!(ledger.total, 2);
|
||||
assert_eq!(ledger.items[0].id, "tx-1");
|
||||
assert_eq!(ledger.items[0].wallet_user_name.as_deref(), Some("Alice"));
|
||||
|
||||
let refund_requests = repository
|
||||
.list_admin_wallet_refund_requests(&AdminWalletRefundRequestListQuery {
|
||||
status: Some("pending_approval".to_string()),
|
||||
limit: 10,
|
||||
offset: 0,
|
||||
})
|
||||
.await
|
||||
.expect("admin refund requests should list");
|
||||
assert_eq!(refund_requests.total, 1);
|
||||
assert_eq!(refund_requests.items[0].id, "refund-2");
|
||||
assert_eq!(
|
||||
refund_requests.items[0].wallet_user_name.as_deref(),
|
||||
Some("Alice")
|
||||
);
|
||||
|
||||
let transactions = repository
|
||||
.list_admin_wallet_transactions("wallet-1", 1, 1)
|
||||
.await
|
||||
.expect("wallet transactions should page");
|
||||
assert_eq!(transactions.total, 2);
|
||||
assert_eq!(transactions.items[0].id, "tx-1");
|
||||
|
||||
let orders = repository
|
||||
.list_admin_payment_orders(&AdminPaymentOrderListQuery {
|
||||
@@ -66,23 +246,99 @@ async fn mysql_wallet_read_repository_reads_wallet_contract_views() {
|
||||
true
|
||||
);
|
||||
|
||||
let expired_orders = repository
|
||||
.list_admin_payment_orders(&AdminPaymentOrderListQuery {
|
||||
status: Some("expired".to_string()),
|
||||
payment_method: Some("card".to_string()),
|
||||
limit: 10,
|
||||
offset: 0,
|
||||
})
|
||||
.await
|
||||
.expect("expired payment orders should filter by effective status");
|
||||
assert_eq!(expired_orders.total, 1);
|
||||
assert_eq!(expired_orders.items[0].id, "order-2");
|
||||
|
||||
let user_orders = repository
|
||||
.list_wallet_payment_orders_by_user_id("user-1", 2, 0)
|
||||
.await
|
||||
.expect("user payment orders should page");
|
||||
assert_eq!(user_orders.total, 3);
|
||||
assert_eq!(user_orders.items.len(), 2);
|
||||
assert_eq!(user_orders.items[0].id, "order-3");
|
||||
assert_eq!(user_orders.items[1].id, "order-2");
|
||||
assert_eq!(user_orders.items[1].status, "expired");
|
||||
assert_eq!(
|
||||
repository
|
||||
.count_pending_payment_orders_by_user_id("user-1")
|
||||
.await
|
||||
.expect("pending payment orders should count"),
|
||||
2
|
||||
);
|
||||
assert_eq!(
|
||||
repository
|
||||
.find_admin_payment_order("order-1")
|
||||
.await
|
||||
.expect("admin payment order should query")
|
||||
.expect("admin payment order should exist")
|
||||
.id,
|
||||
"order-1"
|
||||
);
|
||||
assert!(repository
|
||||
.find_wallet_payment_order_by_user_id("user-2", "order-1")
|
||||
.await
|
||||
.expect("cross-user payment order lookup should query")
|
||||
.is_none());
|
||||
|
||||
let refunds = repository
|
||||
.list_admin_wallet_refunds("wallet-1", 10, 0)
|
||||
.list_admin_wallet_refunds("wallet-1", 1, 1)
|
||||
.await
|
||||
.expect("refunds should list");
|
||||
assert_eq!(refunds.total, 1);
|
||||
assert_eq!(refunds.total, 2);
|
||||
assert_eq!(refunds.items[0].id, "refund-1");
|
||||
assert_eq!(
|
||||
refunds.items[0].payout_proof.as_ref().unwrap()["proof"],
|
||||
"ok"
|
||||
);
|
||||
assert_eq!(
|
||||
repository
|
||||
.count_pending_refunds_by_user_id("user-1")
|
||||
.await
|
||||
.expect("pending refunds should count"),
|
||||
1
|
||||
);
|
||||
assert!(repository
|
||||
.find_wallet_refund("wallet-2", "refund-1")
|
||||
.await
|
||||
.expect("cross-wallet refund lookup should query")
|
||||
.is_none());
|
||||
|
||||
let callbacks = repository
|
||||
.list_admin_payment_callbacks(Some("redeem_code"), 10, 0)
|
||||
.list_admin_payment_callbacks(Some("redeem_code"), 1, 1)
|
||||
.await
|
||||
.expect("callbacks should list");
|
||||
assert_eq!(callbacks.total, 1);
|
||||
assert_eq!(callbacks.total, 2);
|
||||
assert_eq!(callbacks.items[0].id, "callback-1");
|
||||
assert!(callbacks.items[0].signature_valid);
|
||||
|
||||
let batches = repository
|
||||
.list_admin_redeem_code_batches(&AdminRedeemCodeBatchListQuery {
|
||||
status: Some("active".to_string()),
|
||||
limit: 1,
|
||||
offset: 0,
|
||||
})
|
||||
.await
|
||||
.expect("redeem batches should filter");
|
||||
assert_eq!(batches.total, 1);
|
||||
assert_eq!(batches.items[0].id, "batch-1");
|
||||
assert_eq!(batches.items[0].redeemed_count, 1);
|
||||
assert_eq!(batches.items[0].active_count, 1);
|
||||
let batch = repository
|
||||
.find_admin_redeem_code_batch("batch-2")
|
||||
.await
|
||||
.expect("redeem batch should query")
|
||||
.expect("redeem batch should exist");
|
||||
assert_eq!(batch.active_count, 1);
|
||||
|
||||
let codes = repository
|
||||
.list_admin_redeem_codes(&AdminRedeemCodeListQuery {
|
||||
batch_id: "batch-1".to_string(),
|
||||
@@ -94,6 +350,17 @@ async fn mysql_wallet_read_repository_reads_wallet_contract_views() {
|
||||
.expect("redeem codes should list");
|
||||
assert_eq!(codes.total, 1);
|
||||
assert_eq!(codes.items[0].masked_code, "ABCD****WXYZ");
|
||||
let paged_codes = repository
|
||||
.list_admin_redeem_codes(&AdminRedeemCodeListQuery {
|
||||
batch_id: "batch-1".to_string(),
|
||||
status: None,
|
||||
limit: 1,
|
||||
offset: 1,
|
||||
})
|
||||
.await
|
||||
.expect("redeem codes should page");
|
||||
assert_eq!(paged_codes.total, 2);
|
||||
assert_eq!(paged_codes.items[0].id, "code-1");
|
||||
|
||||
let today = super::current_billing_date("UTC").expect("UTC should parse");
|
||||
sqlx::query("UPDATE wallet_daily_usage_ledgers SET billing_date = ? WHERE id = 'daily-1'")
|
||||
@@ -107,6 +374,8 @@ async fn mysql_wallet_read_repository_reads_wallet_contract_views() {
|
||||
.expect("daily usage should query")
|
||||
.expect("daily usage should exist");
|
||||
assert_eq!(daily.total_requests, 2);
|
||||
|
||||
cleanup_rows(repository.pool()).await;
|
||||
}
|
||||
|
||||
impl MysqlWalletReadRepository {
|
||||
@@ -118,14 +387,15 @@ impl MysqlWalletReadRepository {
|
||||
async fn cleanup_rows(pool: &sqlx::MySqlPool) {
|
||||
for sql in [
|
||||
"DELETE FROM wallet_daily_usage_ledgers WHERE id = 'daily-1'",
|
||||
"DELETE FROM redeem_codes WHERE id = 'code-1'",
|
||||
"DELETE FROM redeem_code_batches WHERE id = 'batch-1'",
|
||||
"DELETE FROM wallet_transactions WHERE id = 'tx-1'",
|
||||
"DELETE FROM refund_requests WHERE id = 'refund-1'",
|
||||
"DELETE FROM payment_callbacks WHERE id = 'callback-1'",
|
||||
"DELETE FROM payment_orders WHERE id = 'order-1'",
|
||||
"DELETE FROM wallets WHERE id = 'wallet-1'",
|
||||
"DELETE FROM users WHERE id = 'user-1'",
|
||||
"DELETE FROM redeem_codes WHERE id IN ('code-1', 'code-2', 'code-3')",
|
||||
"DELETE FROM redeem_code_batches WHERE id IN ('batch-1', 'batch-2')",
|
||||
"DELETE FROM wallet_transactions WHERE id IN ('tx-1', 'tx-2', 'tx-api-1')",
|
||||
"DELETE FROM refund_requests WHERE id IN ('refund-1', 'refund-2', 'refund-api-1')",
|
||||
"DELETE FROM payment_callbacks WHERE id IN ('callback-1', 'callback-2', 'callback-3')",
|
||||
"DELETE FROM payment_orders WHERE id IN ('order-1', 'order-2', 'order-3', 'order-4')",
|
||||
"DELETE FROM wallets WHERE id IN ('wallet-1', 'wallet-2', 'wallet-api-1')",
|
||||
"DELETE FROM api_keys WHERE id = 'api-key-1'",
|
||||
"DELETE FROM users WHERE id IN ('user-1', 'user-2')",
|
||||
] {
|
||||
sqlx::query(sql)
|
||||
.execute(pool)
|
||||
@@ -138,72 +408,99 @@ async fn seed_rows(pool: &sqlx::MySqlPool) {
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO users (id, username, email, auth_source, created_at, updated_at)
|
||||
VALUES ('user-1', 'Alice', 'alice@example.com', 'local', 1, 1)
|
||||
VALUES
|
||||
('user-1', 'Alice', 'alice@example.com', 'local', 1, 1),
|
||||
('user-2', 'Bob', 'bob@example.com', 'local', 1, 1)
|
||||
"#,
|
||||
)
|
||||
.execute(pool)
|
||||
.await
|
||||
.expect("user should seed");
|
||||
.expect("users should seed");
|
||||
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO api_keys (id, user_id, key_hash, name, created_at, updated_at)
|
||||
VALUES ('api-key-1', 'user-1', 'wallet-test-api-key-hash-1', 'Standalone Key', 1, 1)
|
||||
"#,
|
||||
)
|
||||
.execute(pool)
|
||||
.await
|
||||
.expect("api key should seed");
|
||||
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO wallets (
|
||||
id, user_id, balance, gift_balance, total_recharged, total_consumed,
|
||||
id, user_id, api_key_id, balance, gift_balance, total_recharged, total_consumed,
|
||||
total_refunded, total_adjusted, created_at, updated_at
|
||||
) VALUES (
|
||||
'wallet-1', 'user-1', 10.0, 2.0, 20.0, 4.0, 1.0, 3.0, 1, 2
|
||||
)
|
||||
) VALUES
|
||||
('wallet-1', 'user-1', NULL, 10.0, 2.0, 20.0, 4.0, 1.0, 3.0, 1, 2),
|
||||
('wallet-2', 'user-2', NULL, 4.0, 1.0, 5.0, 2.0, 0.0, 0.0, 1, 3),
|
||||
('wallet-api-1', NULL, 'api-key-1', 7.0, 0.0, 7.0, 0.0, 0.0, 0.0, 1, 4)
|
||||
"#,
|
||||
)
|
||||
.execute(pool)
|
||||
.await
|
||||
.expect("wallet should seed");
|
||||
.expect("wallets should seed");
|
||||
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO payment_orders (
|
||||
id, order_no, wallet_id, user_id, amount_usd, refunded_amount_usd,
|
||||
refundable_amount_usd, payment_method, gateway_response, status, created_at
|
||||
) VALUES (
|
||||
'order-1', 'order-no-1', 'wallet-1', 'user-1', 5.0, 1.0, 4.0,
|
||||
'redeem_code', '{"ok":true}', 'credited', 3
|
||||
)
|
||||
refundable_amount_usd, payment_method, gateway_response, status, created_at,
|
||||
expires_at
|
||||
) VALUES
|
||||
('order-1', 'order-no-1', 'wallet-1', 'user-1', 5.0, 1.0, 4.0,
|
||||
'redeem_code', '{"ok":true}', 'credited', 3, NULL),
|
||||
('order-2', 'order-no-2', 'wallet-1', 'user-1', 2.0, 0.0, 2.0,
|
||||
'card', NULL, 'pending', 12, 1),
|
||||
('order-3', 'order-no-3', 'wallet-1', 'user-1', 3.0, 0.0, 3.0,
|
||||
'card', NULL, 'pending', 13, 4102444800),
|
||||
('order-4', 'order-no-4', 'wallet-2', 'user-2', 4.0, 0.0, 4.0,
|
||||
'bank', NULL, 'paid', 14, 4102444800)
|
||||
"#,
|
||||
)
|
||||
.execute(pool)
|
||||
.await
|
||||
.expect("payment order should seed");
|
||||
.expect("payment orders should seed");
|
||||
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO payment_callbacks (
|
||||
id, payment_order_id, payment_method, callback_key, order_no,
|
||||
signature_valid, payload, created_at
|
||||
) VALUES (
|
||||
'callback-1', 'order-1', 'redeem_code', 'callback-key-1',
|
||||
'order-no-1', 1, '{"event":"paid"}', 4
|
||||
)
|
||||
) VALUES
|
||||
('callback-1', 'order-1', 'redeem_code', 'callback-key-1',
|
||||
'order-no-1', 1, '{"event":"paid"}', 4),
|
||||
('callback-2', 'order-1', 'redeem_code', 'callback-key-2',
|
||||
'order-no-1', 1, '{"event":"duplicate"}', 5),
|
||||
('callback-3', 'order-2', 'card', 'callback-key-3',
|
||||
'order-no-2', 0, '{"event":"created"}', 6)
|
||||
"#,
|
||||
)
|
||||
.execute(pool)
|
||||
.await
|
||||
.expect("callback should seed");
|
||||
.expect("callbacks should seed");
|
||||
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO refund_requests (
|
||||
id, refund_no, wallet_id, user_id, payment_order_id, source_type,
|
||||
refund_mode, amount_usd, status, payout_proof, created_at, updated_at
|
||||
) VALUES (
|
||||
'refund-1', 'refund-no-1', 'wallet-1', 'user-1', 'order-1',
|
||||
'payment_order', 'offline_payout', 1.0, 'completed',
|
||||
'{"proof":"ok"}', 5, 6
|
||||
)
|
||||
) VALUES
|
||||
('refund-1', 'refund-no-1', 'wallet-1', 'user-1', 'order-1',
|
||||
'payment_order', 'offline_payout', 1.0, 'completed',
|
||||
'{"proof":"ok"}', 5, 6),
|
||||
('refund-2', 'refund-no-2', 'wallet-1', 'user-1', 'order-2',
|
||||
'payment_order', 'original_channel', 0.5, 'pending_approval',
|
||||
NULL, 6, 7),
|
||||
('refund-api-1', 'refund-no-api-1', 'wallet-api-1', NULL, NULL,
|
||||
'manual', 'offline_payout', 0.5, 'pending_approval',
|
||||
NULL, 7, 8)
|
||||
"#,
|
||||
)
|
||||
.execute(pool)
|
||||
.await
|
||||
.expect("refund should seed");
|
||||
.expect("refunds should seed");
|
||||
|
||||
sqlx::query(
|
||||
r#"
|
||||
@@ -211,28 +508,31 @@ INSERT INTO wallet_transactions (
|
||||
id, wallet_id, category, reason_code, amount, balance_before,
|
||||
balance_after, recharge_balance_before, recharge_balance_after,
|
||||
gift_balance_before, gift_balance_after, created_at
|
||||
) VALUES (
|
||||
'tx-1', 'wallet-1', 'credit', 'manual_adjustment', 3.0, 7.0, 10.0,
|
||||
5.0, 8.0, 2.0, 2.0, 7
|
||||
)
|
||||
) VALUES
|
||||
('tx-1', 'wallet-1', 'credit', 'manual_adjustment', 3.0, 7.0, 10.0,
|
||||
5.0, 8.0, 2.0, 2.0, 7),
|
||||
('tx-2', 'wallet-1', 'credit', 'manual_adjustment', 1.0, 11.0, 12.0,
|
||||
9.0, 10.0, 2.0, 2.0, 8),
|
||||
('tx-api-1', 'wallet-api-1', 'credit', 'manual_adjustment', 1.0, 6.0, 7.0,
|
||||
6.0, 7.0, 0.0, 0.0, 9)
|
||||
"#,
|
||||
)
|
||||
.execute(pool)
|
||||
.await
|
||||
.expect("transaction should seed");
|
||||
.expect("transactions should seed");
|
||||
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO redeem_code_batches (
|
||||
id, name, amount_usd, total_count, created_at, updated_at
|
||||
) VALUES (
|
||||
'batch-1', 'Batch One', 5.0, 1, 8, 9
|
||||
)
|
||||
id, name, amount_usd, total_count, status, created_at, updated_at
|
||||
) VALUES
|
||||
('batch-1', 'Batch One', 5.0, 2, 'active', 8, 9),
|
||||
('batch-2', 'Batch Two', 8.0, 1, 'disabled', 9, 10)
|
||||
"#,
|
||||
)
|
||||
.execute(pool)
|
||||
.await
|
||||
.expect("redeem batch should seed");
|
||||
.expect("redeem batches should seed");
|
||||
|
||||
sqlx::query(
|
||||
r#"
|
||||
@@ -240,15 +540,18 @@ INSERT INTO redeem_codes (
|
||||
id, batch_id, code_hash, code_prefix, code_suffix, status,
|
||||
redeemed_by_user_id, redeemed_wallet_id, redeemed_payment_order_id,
|
||||
redeemed_at, created_at, updated_at
|
||||
) VALUES (
|
||||
'code-1', 'batch-1', 'hash-1', 'ABCD', 'WXYZ', 'redeemed',
|
||||
'user-1', 'wallet-1', 'order-1', 10, 8, 10
|
||||
)
|
||||
) VALUES
|
||||
('code-1', 'batch-1', 'hash-1', 'ABCD', 'WXYZ', 'redeemed',
|
||||
'user-1', 'wallet-1', 'order-1', 10, 8, 10),
|
||||
('code-2', 'batch-1', 'hash-2', 'EFGH', 'QRST', 'active',
|
||||
NULL, NULL, NULL, NULL, 9, 10),
|
||||
('code-3', 'batch-2', 'hash-3', 'IJKL', 'MNOP', 'active',
|
||||
NULL, NULL, NULL, NULL, 10, 11)
|
||||
"#,
|
||||
)
|
||||
.execute(pool)
|
||||
.await
|
||||
.expect("redeem code should seed");
|
||||
.expect("redeem codes should seed");
|
||||
|
||||
sqlx::query(
|
||||
r#"
|
||||
|
||||
Reference in New Issue
Block a user