i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-02 17:30:23 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: 扩展 Rust gateway 全功能模块，新增 billing/crypto/wallet crate 及完整数据层

Browse Source 
- 新增 aether-billing、aether-crypto、aether-wallet 独立 crate
- aether-data 扩展 repository 层：announcements、auth_modules、billing、
  candidate_selection、gemini_file_mappings、global_models、management_tokens、
  oauth_providers、proxy_nodes、quota、users、wallet 等模块
- aether-gateway 新增 api/auth/billing/control/middleware/scheduler/usage/
  video_tasks/hooks/maintenance/model_fetch/provider_transport 等功能模块
- 重构 executor decision 和 gateway state 为模块目录结构
- 新增 gateway router、frontdoor 路由层及对应测试
- Python 侧 API 路由重构，新增 compat/support 模块
- 前端 Logo 组件更新及 Provider 管理页面调整
This commit is contained in:
fawney19
2026-03-31 19:19:04 +08:00
parent b5a0070023
commit ddf18fed9a
690 changed files with 235087 additions and 16301 deletions
Download Patch File Download Diff File Expand all files Collapse all files

466
crates/aether-data/src/repository/candidates/memory.rs
View File

@@ -1,9 +1,13 @@
use std::collections::BTreeMap;
use std::collections::{BTreeMap, BTreeSet};
use std::sync::RwLock;
use async_trait::async_trait;
use super::types::{RequestCandidateReadRepository, StoredRequestCandidate};
use super::types::{
PublicHealthStatusCount, PublicHealthTimelineBucket, RequestCandidateReadRepository,
RequestCandidateStatus, RequestCandidateWriteRepository, StoredRequestCandidate,
UpsertRequestCandidateRecord,
};
use crate::DataLayerError;
#[derive(Debug, Default)]
@@ -68,13 +72,340 @@ impl RequestCandidateReadRepository for InMemoryRequestCandidateRepository {
rows.truncate(limit);
Ok(rows)
}
async fn list_by_provider_id(
&self,
provider_id: &str,
limit: usize,
) -> Result<Vec<StoredRequestCandidate>, DataLayerError> {
if limit == 0 {
return Ok(Vec::new());
}
let mut rows = self
.by_id
.read()
.expect("request candidate repository lock")
.values()
.filter(|row| row.provider_id.as_deref() == Some(provider_id))
.cloned()
.collect::<Vec<_>>();
rows.sort_by(|left, right| right.created_at_unix_secs.cmp(&left.created_at_unix_secs));
rows.truncate(limit);
Ok(rows)
}
async fn list_finalized_by_endpoint_ids_since(
&self,
endpoint_ids: &[String],
since_unix_secs: u64,
limit: usize,
) -> Result<Vec<StoredRequestCandidate>, DataLayerError> {
if endpoint_ids.is_empty() || limit == 0 {
return Ok(Vec::new());
}
let endpoint_ids = endpoint_ids.iter().cloned().collect::<BTreeSet<_>>();
let mut rows = self
.by_id
.read()
.expect("request candidate repository lock")
.values()
.filter(|row| {
row.endpoint_id
.as_ref()
.is_some_and(|endpoint_id| endpoint_ids.contains(endpoint_id))
&& row.created_at_unix_secs >= since_unix_secs
&& matches!(
row.status,
RequestCandidateStatus::Success
| RequestCandidateStatus::Failed
| RequestCandidateStatus::Skipped
)
})
.cloned()
.collect::<Vec<_>>();
rows.sort_by(|left, right| right.created_at_unix_secs.cmp(&left.created_at_unix_secs));
rows.truncate(limit);
Ok(rows)
}
async fn count_finalized_statuses_by_endpoint_ids_since(
&self,
endpoint_ids: &[String],
since_unix_secs: u64,
) -> Result<Vec<PublicHealthStatusCount>, DataLayerError> {
if endpoint_ids.is_empty() {
return Ok(Vec::new());
}
let endpoint_ids = endpoint_ids.iter().cloned().collect::<BTreeSet<_>>();
let mut counts = BTreeMap::<(String, &'static str), u64>::new();
for row in self
.by_id
.read()
.expect("request candidate repository lock")
.values()
{
let Some(endpoint_id) = row.endpoint_id.as_ref() else {
continue;
};
if !endpoint_ids.contains(endpoint_id) || row.created_at_unix_secs < since_unix_secs {
continue;
}
if !matches!(
row.status,
RequestCandidateStatus::Success
| RequestCandidateStatus::Failed
| RequestCandidateStatus::Skipped
) {
continue;
}
let status_key = match row.status {
RequestCandidateStatus::Success => "success",
RequestCandidateStatus::Failed => "failed",
RequestCandidateStatus::Skipped => "skipped",
_ => continue,
};
*counts.entry((endpoint_id.clone(), status_key)).or_insert(0) += 1;
}
Ok(counts
.into_iter()
.map(|((endpoint_id, status_key), count)| {
let status = match status_key {
"success" => RequestCandidateStatus::Success,
"failed" => RequestCandidateStatus::Failed,
"skipped" => RequestCandidateStatus::Skipped,
_ => unreachable!("filtered status should stay finalized"),
};
PublicHealthStatusCount {
endpoint_id,
status,
count,
}
})
.collect())
}
async fn aggregate_finalized_timeline_by_endpoint_ids_since(
&self,
endpoint_ids: &[String],
since_unix_secs: u64,
until_unix_secs: u64,
segments: u32,
) -> Result<Vec<PublicHealthTimelineBucket>, DataLayerError> {
if endpoint_ids.is_empty() || segments == 0 || until_unix_secs < since_unix_secs {
return Ok(Vec::new());
}
let endpoint_ids = endpoint_ids.iter().cloned().collect::<BTreeSet<_>>();
let span = until_unix_secs.saturating_sub(since_unix_secs);
let mut buckets = BTreeMap::<(String, u32), PublicHealthTimelineBucket>::new();
for row in self
.by_id
.read()
.expect("request candidate repository lock")
.values()
{
let Some(endpoint_id) = row.endpoint_id.as_ref() else {
continue;
};
if !endpoint_ids.contains(endpoint_id)
|| row.created_at_unix_secs < since_unix_secs
|| row.created_at_unix_secs > until_unix_secs
{
continue;
}
if !matches!(
row.status,
RequestCandidateStatus::Success
| RequestCandidateStatus::Failed
| RequestCandidateStatus::Skipped
) {
continue;
}
let segment_idx = if span == 0 {
0
} else {
let offset = row.created_at_unix_secs.saturating_sub(since_unix_secs);
let idx = ((offset as u128) * (segments as u128) / (span as u128)) as u32;
idx.min(segments.saturating_sub(1))
};
let bucket = buckets
.entry((endpoint_id.clone(), segment_idx))
.or_insert_with(|| PublicHealthTimelineBucket {
endpoint_id: endpoint_id.clone(),
segment_idx,
total_count: 0,
success_count: 0,
failed_count: 0,
min_created_at_unix_secs: None,
max_created_at_unix_secs: None,
});
bucket.total_count += 1;
if row.status == RequestCandidateStatus::Success {
bucket.success_count += 1;
} else if row.status == RequestCandidateStatus::Failed {
bucket.failed_count += 1;
}
bucket.min_created_at_unix_secs = Some(
bucket
.min_created_at_unix_secs
.map(|value| value.min(row.created_at_unix_secs))
.unwrap_or(row.created_at_unix_secs),
);
bucket.max_created_at_unix_secs = Some(
bucket
.max_created_at_unix_secs
.map(|value| value.max(row.created_at_unix_secs))
.unwrap_or(row.created_at_unix_secs),
);
}
Ok(buckets.into_values().collect())
}
}
#[async_trait]
impl RequestCandidateWriteRepository for InMemoryRequestCandidateRepository {
async fn upsert(
&self,
candidate: UpsertRequestCandidateRecord,
) -> Result<StoredRequestCandidate, DataLayerError> {
candidate.validate()?;
let mut by_id = self
.by_id
.write()
.expect("request candidate repository lock");
let existing = by_id
.values()
.find(|row| {
row.request_id == candidate.request_id
&& row.candidate_index == candidate.candidate_index
&& row.retry_index == candidate.retry_index
})
.cloned();
let created_at_unix_secs = existing
.as_ref()
.map(|row| row.created_at_unix_secs)
.or(candidate.created_at_unix_secs)
.or(candidate.started_at_unix_secs)
.or(candidate.finished_at_unix_secs)
.unwrap_or_default();
let stored = StoredRequestCandidate {
id: existing
.as_ref()
.map(|row| row.id.clone())
.unwrap_or_else(|| candidate.id.clone()),
request_id: candidate.request_id.clone(),
user_id: candidate
.user_id
.or_else(|| existing.as_ref().and_then(|row| row.user_id.clone())),
api_key_id: candidate
.api_key_id
.or_else(|| existing.as_ref().and_then(|row| row.api_key_id.clone())),
username: candidate
.username
.or_else(|| existing.as_ref().and_then(|row| row.username.clone())),
api_key_name: candidate
.api_key_name
.or_else(|| existing.as_ref().and_then(|row| row.api_key_name.clone())),
candidate_index: candidate.candidate_index,
retry_index: candidate.retry_index,
provider_id: candidate
.provider_id
.or_else(|| existing.as_ref().and_then(|row| row.provider_id.clone())),
endpoint_id: candidate
.endpoint_id
.or_else(|| existing.as_ref().and_then(|row| row.endpoint_id.clone())),
key_id: candidate
.key_id
.or_else(|| existing.as_ref().and_then(|row| row.key_id.clone())),
status: candidate.status,
skip_reason: candidate
.skip_reason
.or_else(|| existing.as_ref().and_then(|row| row.skip_reason.clone())),
is_cached: candidate
.is_cached
.unwrap_or_else(|| existing.as_ref().map(|row| row.is_cached).unwrap_or(false)),
status_code: candidate
.status_code
.or_else(|| existing.as_ref().and_then(|row| row.status_code)),
error_type: candidate
.error_type
.or_else(|| existing.as_ref().and_then(|row| row.error_type.clone())),
error_message: candidate
.error_message
.or_else(|| existing.as_ref().and_then(|row| row.error_message.clone())),
latency_ms: candidate
.latency_ms
.or_else(|| existing.as_ref().and_then(|row| row.latency_ms)),
concurrent_requests: candidate
.concurrent_requests
.or_else(|| existing.as_ref().and_then(|row| row.concurrent_requests)),
extra_data: candidate
.extra_data
.or_else(|| existing.as_ref().and_then(|row| row.extra_data.clone())),
required_capabilities: candidate.required_capabilities.or_else(|| {
existing
.as_ref()
.and_then(|row| row.required_capabilities.clone())
}),
created_at_unix_secs,
started_at_unix_secs: candidate
.started_at_unix_secs
.or_else(|| existing.as_ref().and_then(|row| row.started_at_unix_secs)),
finished_at_unix_secs: candidate
.finished_at_unix_secs
.or_else(|| existing.as_ref().and_then(|row| row.finished_at_unix_secs)),
};
by_id.insert(stored.id.clone(), stored.clone());
Ok(stored)
}
async fn delete_created_before(
&self,
created_before_unix_secs: u64,
limit: usize,
) -> Result<usize, DataLayerError> {
if limit == 0 {
return Ok(0);
}
let mut by_id = self
.by_id
.write()
.expect("request candidate repository lock");
let mut ids = by_id
.values()
.filter(|row| row.created_at_unix_secs < created_before_unix_secs)
.map(|row| (row.created_at_unix_secs, row.id.clone()))
.collect::<Vec<_>>();
ids.sort_by(|left, right| left.cmp(right));
let mut deleted = 0usize;
for (_, id) in ids.into_iter().take(limit) {
if by_id.remove(&id).is_some() {
deleted += 1;
}
}
Ok(deleted)
}
}
#[cfg(test)]
mod tests {
use super::InMemoryRequestCandidateRepository;
use crate::repository::candidates::{
RequestCandidateReadRepository, RequestCandidateStatus, StoredRequestCandidate,
RequestCandidateReadRepository, RequestCandidateStatus, RequestCandidateWriteRepository,
StoredRequestCandidate, UpsertRequestCandidateRecord,
};
fn sample_candidate(
@@ -145,4 +476,133 @@ mod tests {
assert_eq!(rows[0].id, "cand-2");
assert_eq!(rows[1].id, "cand-1");
}
#[tokio::test]
async fn aggregates_finalized_health_data_by_endpoint_ids() {
let repository = InMemoryRequestCandidateRepository::seed(vec![
sample_candidate("cand-1", "req-1", 100),
sample_candidate("cand-2", "req-2", 200),
]);
let counts = repository
.count_finalized_statuses_by_endpoint_ids_since(&["endpoint-1".to_string()], 0)
.await
.expect("count should succeed");
assert_eq!(counts.len(), 1);
assert_eq!(counts[0].endpoint_id, "endpoint-1");
assert_eq!(counts[0].status, RequestCandidateStatus::Success);
assert_eq!(counts[0].count, 2);
let timeline = repository
.aggregate_finalized_timeline_by_endpoint_ids_since(
&["endpoint-1".to_string()],
0,
300,
3,
)
.await
.expect("timeline should succeed");
assert_eq!(timeline.len(), 2);
let attempts = repository
.list_finalized_by_endpoint_ids_since(&["endpoint-1".to_string()], 0, 1)
.await
.expect("attempt list should succeed");
assert_eq!(attempts.len(), 1);
assert_eq!(attempts[0].id, "cand-2");
}
#[tokio::test]
async fn upsert_writes_and_updates_request_candidate() {
let repository = InMemoryRequestCandidateRepository::default();
let created = repository
.upsert(UpsertRequestCandidateRecord {
id: "cand-1".to_string(),
request_id: "req-1".to_string(),
user_id: Some("user-1".to_string()),
api_key_id: Some("api-key-1".to_string()),
username: Some("alice".to_string()),
api_key_name: Some("default".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("key-1".to_string()),
status: RequestCandidateStatus::Available,
skip_reason: None,
is_cached: Some(false),
status_code: None,
error_type: None,
error_message: None,
latency_ms: None,
concurrent_requests: None,
extra_data: None,
required_capabilities: None,
created_at_unix_secs: Some(100),
started_at_unix_secs: None,
finished_at_unix_secs: None,
})
.await
.expect("create should succeed");
assert_eq!(created.id, "cand-1");
assert_eq!(created.status, RequestCandidateStatus::Available);
let updated = repository
.upsert(UpsertRequestCandidateRecord {
id: "cand-1-replacement".to_string(),
request_id: "req-1".to_string(),
user_id: None,
api_key_id: None,
username: None,
api_key_name: None,
candidate_index: 0,
retry_index: 0,
provider_id: None,
endpoint_id: None,
key_id: None,
status: RequestCandidateStatus::Success,
skip_reason: None,
is_cached: None,
status_code: Some(200),
error_type: None,
error_message: None,
latency_ms: Some(25),
concurrent_requests: Some(2),
extra_data: None,
required_capabilities: None,
created_at_unix_secs: None,
started_at_unix_secs: Some(101),
finished_at_unix_secs: Some(102),
})
.await
.expect("update should succeed");
assert_eq!(updated.id, "cand-1");
assert_eq!(updated.status, RequestCandidateStatus::Success);
assert_eq!(updated.status_code, Some(200));
assert_eq!(updated.latency_ms, Some(25));
assert_eq!(updated.started_at_unix_secs, Some(101));
}
#[tokio::test]
async fn delete_created_before_removes_oldest_matching_rows_up_to_limit() {
let repository = InMemoryRequestCandidateRepository::seed(vec![
sample_candidate("cand-1", "req-1", 100),
sample_candidate("cand-2", "req-2", 200),
sample_candidate("cand-3", "req-3", 400),
]);
let deleted = repository
.delete_created_before(350, 1)
.await
.expect("delete should succeed");
assert_eq!(deleted, 1);
let rows = repository
.list_recent(10)
.await
.expect("list recent should succeed");
assert_eq!(rows.len(), 2);
assert_eq!(rows[0].id, "cand-3");
assert_eq!(rows[1].id, "cand-2");
}
}

5
crates/aether-data/src/repository/candidates/mod.rs
View File

@@ -5,6 +5,7 @@ mod types;
pub use memory::InMemoryRequestCandidateRepository;
pub use sql::SqlxRequestCandidateReadRepository;
pub use types::{
RequestCandidateReadRepository, RequestCandidateRepository, RequestCandidateStatus,
StoredRequestCandidate,
PublicHealthStatusCount, PublicHealthTimelineBucket, RequestCandidateReadRepository,
RequestCandidateRepository, RequestCandidateStatus, RequestCandidateWriteRepository,
StoredRequestCandidate, UpsertRequestCandidateRecord,
};

545
crates/aether-data/src/repository/candidates/sql.rs
View File

@@ -1,9 +1,14 @@
use async_trait::async_trait;
use futures_util::future::BoxFuture;
use sqlx::{PgPool, Row};
use uuid::Uuid;
use super::types::{
RequestCandidateReadRepository, RequestCandidateStatus, StoredRequestCandidate,
PublicHealthStatusCount, PublicHealthTimelineBucket, RequestCandidateReadRepository,
RequestCandidateStatus, RequestCandidateWriteRepository, StoredRequestCandidate,
UpsertRequestCandidateRecord,
};
use crate::postgres::PostgresTransactionRunner;
use crate::DataLayerError;
const LIST_BY_REQUEST_ID_SQL: &str = r#"
@@ -68,20 +73,235 @@ ORDER BY created_at DESC
LIMIT $1
"#;
const LIST_BY_PROVIDER_ID_SQL: &str = r#"
SELECT
id,
request_id,
user_id,
api_key_id,
username,
api_key_name,
candidate_index,
retry_index,
provider_id,
endpoint_id,
key_id,
status,
skip_reason,
is_cached,
status_code,
error_type,
error_message,
latency_ms,
concurrent_requests,
extra_data,
required_capabilities,
CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_secs,
CAST(EXTRACT(EPOCH FROM started_at) AS BIGINT) AS started_at_unix_secs,
CAST(EXTRACT(EPOCH FROM finished_at) AS BIGINT) AS finished_at_unix_secs
FROM request_candidates
WHERE provider_id = $1
ORDER BY created_at DESC
LIMIT $2
"#;
const LIST_FINALIZED_BY_ENDPOINT_IDS_SINCE_SQL: &str = r#"
SELECT
id,
request_id,
user_id,
api_key_id,
username,
api_key_name,
candidate_index,
retry_index,
provider_id,
endpoint_id,
key_id,
status,
skip_reason,
is_cached,
status_code,
error_type,
error_message,
latency_ms,
concurrent_requests,
extra_data,
required_capabilities,
CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_secs,
CAST(EXTRACT(EPOCH FROM started_at) AS BIGINT) AS started_at_unix_secs,
CAST(EXTRACT(EPOCH FROM finished_at) AS BIGINT) AS finished_at_unix_secs
FROM request_candidates
WHERE endpoint_id = ANY($1)
AND created_at >= TO_TIMESTAMP($2)
AND status IN ('success', 'failed', 'skipped')
ORDER BY created_at DESC
LIMIT $3
"#;
const COUNT_FINALIZED_STATUSES_BY_ENDPOINT_IDS_SINCE_SQL: &str = r#"
SELECT
endpoint_id,
status,
COUNT(id) AS count
FROM request_candidates
WHERE endpoint_id = ANY($1)
AND created_at >= TO_TIMESTAMP($2)
AND status IN ('success', 'failed', 'skipped')
GROUP BY endpoint_id, status
"#;
const AGGREGATE_FINALIZED_TIMELINE_BY_ENDPOINT_IDS_SINCE_SQL: &str = r#"
SELECT
endpoint_id,
FLOOR(EXTRACT(EPOCH FROM (created_at - TO_TIMESTAMP($2))) / $4)::BIGINT AS segment_idx,
COUNT(id) AS total_count,
SUM(CASE WHEN status = 'success' THEN 1 ELSE 0 END) AS success_count,
SUM(CASE WHEN status = 'failed' THEN 1 ELSE 0 END) AS failed_count,
CAST(EXTRACT(EPOCH FROM MIN(created_at)) AS BIGINT) AS min_created_at_unix_secs,
CAST(EXTRACT(EPOCH FROM MAX(created_at)) AS BIGINT) AS max_created_at_unix_secs
FROM request_candidates
WHERE endpoint_id = ANY($1)
AND created_at >= TO_TIMESTAMP($2)
AND created_at <= TO_TIMESTAMP($3)
AND status IN ('success', 'failed', 'skipped')
GROUP BY
endpoint_id,
FLOOR(EXTRACT(EPOCH FROM (created_at - TO_TIMESTAMP($2))) / $4)::BIGINT
"#;
const UPSERT_SQL: &str = r#"
INSERT INTO request_candidates (
id,
request_id,
user_id,
api_key_id,
username,
api_key_name,
candidate_index,
retry_index,
provider_id,
endpoint_id,
key_id,
status,
skip_reason,
is_cached,
status_code,
error_type,
error_message,
latency_ms,
concurrent_requests,
extra_data,
required_capabilities,
created_at,
started_at,
finished_at
)
VALUES (
$1,
$2,
$3,
$4,
$5,
$6,
$7,
$8,
$9,
$10,
$11,
$12,
$13,
COALESCE($14, false),
$15,
$16,
$17,
$18,
$19,
$20,
$21,
TO_TIMESTAMP(COALESCE($22, 0)),
TO_TIMESTAMP($23),
TO_TIMESTAMP($24)
)
ON CONFLICT (request_id, candidate_index, retry_index)
DO UPDATE SET
user_id = COALESCE(EXCLUDED.user_id, request_candidates.user_id),
api_key_id = COALESCE(EXCLUDED.api_key_id, request_candidates.api_key_id),
username = COALESCE(EXCLUDED.username, request_candidates.username),
api_key_name = COALESCE(EXCLUDED.api_key_name, request_candidates.api_key_name),
provider_id = COALESCE(EXCLUDED.provider_id, request_candidates.provider_id),
endpoint_id = COALESCE(EXCLUDED.endpoint_id, request_candidates.endpoint_id),
key_id = COALESCE(EXCLUDED.key_id, request_candidates.key_id),
status = EXCLUDED.status,
skip_reason = COALESCE(EXCLUDED.skip_reason, request_candidates.skip_reason),
is_cached = COALESCE($14, request_candidates.is_cached),
status_code = COALESCE(EXCLUDED.status_code, request_candidates.status_code),
error_type = COALESCE(EXCLUDED.error_type, request_candidates.error_type),
error_message = COALESCE(EXCLUDED.error_message, request_candidates.error_message),
latency_ms = COALESCE(EXCLUDED.latency_ms, request_candidates.latency_ms),
concurrent_requests = COALESCE(EXCLUDED.concurrent_requests, request_candidates.concurrent_requests),
extra_data = COALESCE(EXCLUDED.extra_data, request_candidates.extra_data),
required_capabilities = COALESCE(EXCLUDED.required_capabilities, request_candidates.required_capabilities),
started_at = COALESCE(EXCLUDED.started_at, request_candidates.started_at),
finished_at = COALESCE(EXCLUDED.finished_at, request_candidates.finished_at)
RETURNING
id,
request_id,
user_id,
api_key_id,
username,
api_key_name,
candidate_index,
retry_index,
provider_id,
endpoint_id,
key_id,
status,
skip_reason,
is_cached,
status_code,
error_type,
error_message,
latency_ms,
concurrent_requests,
extra_data,
required_capabilities,
CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_secs,
CAST(EXTRACT(EPOCH FROM started_at) AS BIGINT) AS started_at_unix_secs,
CAST(EXTRACT(EPOCH FROM finished_at) AS BIGINT) AS finished_at_unix_secs
"#;
const DELETE_CREATED_BEFORE_SQL: &str = r#"
DELETE FROM request_candidates
WHERE id IN (
SELECT id
FROM request_candidates
WHERE created_at < TO_TIMESTAMP($1)
ORDER BY created_at ASC, id ASC
LIMIT $2
)
"#;
#[derive(Debug, Clone)]
pub struct SqlxRequestCandidateReadRepository {
pool: PgPool,
tx_runner: PostgresTransactionRunner,
}
impl SqlxRequestCandidateReadRepository {
pub fn new(pool: PgPool) -> Self {
Self { pool }
let tx_runner = PostgresTransactionRunner::new(pool.clone());
Self { pool, tx_runner }
}
pub fn pool(&self) -> &PgPool {
&self.pool
}
pub fn transaction_runner(&self) -> &PostgresTransactionRunner {
&self.tx_runner
}
pub async fn list_by_request_id(
&self,
request_id: &str,
@@ -111,6 +331,240 @@ impl SqlxRequestCandidateReadRepository {
.await?;
rows.iter().map(map_request_candidate_row).collect()
}
pub async fn list_by_provider_id(
&self,
provider_id: &str,
limit: usize,
) -> Result<Vec<StoredRequestCandidate>, DataLayerError> {
if limit == 0 {
return Ok(Vec::new());
}
let limit_value = i64::try_from(limit).map_err(|_| {
DataLayerError::UnexpectedValue(format!(
"invalid provider request candidate limit: {limit}"
))
})?;
let rows = sqlx::query(LIST_BY_PROVIDER_ID_SQL)
.bind(provider_id)
.bind(limit_value)
.fetch_all(&self.pool)
.await?;
rows.iter().map(map_request_candidate_row).collect()
}
pub async fn list_finalized_by_endpoint_ids_since(
&self,
endpoint_ids: &[String],
since_unix_secs: u64,
limit: usize,
) -> Result<Vec<StoredRequestCandidate>, DataLayerError> {
if endpoint_ids.is_empty() || limit == 0 {
return Ok(Vec::new());
}
let rows = sqlx::query(LIST_FINALIZED_BY_ENDPOINT_IDS_SINCE_SQL)
.bind(endpoint_ids)
.bind(since_unix_secs as f64)
.bind(i64::try_from(limit).map_err(|_| {
DataLayerError::UnexpectedValue(format!(
"invalid finalized request candidate limit: {limit}"
))
})?)
.fetch_all(&self.pool)
.await?;
rows.iter().map(map_request_candidate_row).collect()
}
pub async fn count_finalized_statuses_by_endpoint_ids_since(
&self,
endpoint_ids: &[String],
since_unix_secs: u64,
) -> Result<Vec<PublicHealthStatusCount>, DataLayerError> {
if endpoint_ids.is_empty() {
return Ok(Vec::new());
}
let rows = sqlx::query(COUNT_FINALIZED_STATUSES_BY_ENDPOINT_IDS_SINCE_SQL)
.bind(endpoint_ids)
.bind(since_unix_secs as f64)
.fetch_all(&self.pool)
.await?;
rows.iter()
.map(|row| {
let status = RequestCandidateStatus::from_database(
row.try_get::<String, _>("status")?.as_str(),
)?;
Ok(PublicHealthStatusCount {
endpoint_id: row.try_get("endpoint_id")?,
status,
count: u64::try_from(row.try_get::<i64, _>("count")?).map_err(|_| {
DataLayerError::UnexpectedValue(
"public health status count out of range".to_string(),
)
})?,
})
})
.collect()
}
pub async fn aggregate_finalized_timeline_by_endpoint_ids_since(
&self,
endpoint_ids: &[String],
since_unix_secs: u64,
until_unix_secs: u64,
segments: u32,
) -> Result<Vec<PublicHealthTimelineBucket>, DataLayerError> {
if endpoint_ids.is_empty() || segments == 0 || until_unix_secs < since_unix_secs {
return Ok(Vec::new());
}
let span_seconds = until_unix_secs.saturating_sub(since_unix_secs);
let segment_seconds = if span_seconds == 0 {
1.0
} else {
(span_seconds as f64) / (segments as f64)
};
let rows = sqlx::query(AGGREGATE_FINALIZED_TIMELINE_BY_ENDPOINT_IDS_SINCE_SQL)
.bind(endpoint_ids)
.bind(since_unix_secs as f64)
.bind(until_unix_secs as f64)
.bind(segment_seconds)
.fetch_all(&self.pool)
.await?;
rows.iter()
.map(|row| {
let raw_segment_idx = row.try_get::<i64, _>("segment_idx")?;
let segment_idx = if raw_segment_idx < 0 {
0
} else {
u32::try_from(raw_segment_idx).map_err(|_| {
DataLayerError::UnexpectedValue(format!(
"public health segment idx out of range: {raw_segment_idx}"
))
})?
}
.min(segments.saturating_sub(1));
Ok(PublicHealthTimelineBucket {
endpoint_id: row.try_get("endpoint_id")?,
segment_idx,
total_count: u64::try_from(row.try_get::<i64, _>("total_count")?).map_err(
|_| {
DataLayerError::UnexpectedValue(
"public health total_count out of range".to_string(),
)
},
)?,
success_count: u64::try_from(row.try_get::<i64, _>("success_count")?).map_err(
|_| {
DataLayerError::UnexpectedValue(
"public health success_count out of range".to_string(),
)
},
)?,
failed_count: u64::try_from(row.try_get::<i64, _>("failed_count")?).map_err(
|_| {
DataLayerError::UnexpectedValue(
"public health failed_count out of range".to_string(),
)
},
)?,
min_created_at_unix_secs: row
.try_get::<Option<i64>, _>("min_created_at_unix_secs")?
.map(|value| {
u64::try_from(value).map_err(|_| {
DataLayerError::UnexpectedValue(format!(
"public health min_created_at_unix_secs out of range: {value}"
))
})
})
.transpose()?,
max_created_at_unix_secs: row
.try_get::<Option<i64>, _>("max_created_at_unix_secs")?
.map(|value| {
u64::try_from(value).map_err(|_| {
DataLayerError::UnexpectedValue(format!(
"public health max_created_at_unix_secs out of range: {value}"
))
})
})
.transpose()?,
})
})
.collect()
}
pub async fn upsert(
&self,
candidate: UpsertRequestCandidateRecord,
) -> Result<StoredRequestCandidate, DataLayerError> {
candidate.validate()?;
self.tx_runner
.run_read_write(|tx| {
Box::pin(async move {
let row = sqlx::query(UPSERT_SQL)
.bind(if candidate.id.trim().is_empty() {
Uuid::new_v4().to_string()
} else {
candidate.id.clone()
})
.bind(&candidate.request_id)
.bind(&candidate.user_id)
.bind(&candidate.api_key_id)
.bind(&candidate.username)
.bind(&candidate.api_key_name)
.bind(to_i32(candidate.candidate_index)?)
.bind(to_i32(candidate.retry_index)?)
.bind(&candidate.provider_id)
.bind(&candidate.endpoint_id)
.bind(&candidate.key_id)
.bind(status_to_database(candidate.status))
.bind(&candidate.skip_reason)
.bind(candidate.is_cached)
.bind(candidate.status_code.map(i32::from))
.bind(&candidate.error_type)
.bind(&candidate.error_message)
.bind(candidate.latency_ms.map(to_i32_u64).transpose()?)
.bind(candidate.concurrent_requests.map(to_i32).transpose()?)
.bind(&candidate.extra_data)
.bind(&candidate.required_capabilities)
.bind(candidate.created_at_unix_secs.map(|value| value as f64))
.bind(candidate.started_at_unix_secs.map(|value| value as f64))
.bind(candidate.finished_at_unix_secs.map(|value| value as f64))
.fetch_one(&mut **tx)
.await?;
map_request_candidate_row(&row)
}) as BoxFuture<'_, Result<StoredRequestCandidate, DataLayerError>>
})
.await
}
pub async fn delete_created_before(
&self,
created_before_unix_secs: u64,
limit: usize,
) -> Result<usize, DataLayerError> {
if limit == 0 {
return Ok(0);
}
let result = sqlx::query(DELETE_CREATED_BEFORE_SQL)
.bind(created_before_unix_secs as f64)
.bind(i64::try_from(limit).map_err(|_| {
DataLayerError::UnexpectedValue(format!(
"invalid request candidate delete limit: {limit}"
))
})?)
.execute(&self.pool)
.await?;
Ok(result.rows_affected() as usize)
}
}
#[async_trait]
@@ -128,6 +582,67 @@ impl RequestCandidateReadRepository for SqlxRequestCandidateReadRepository {
) -> Result<Vec<StoredRequestCandidate>, DataLayerError> {
Self::list_recent(self, limit).await
}
async fn list_finalized_by_endpoint_ids_since(
&self,
endpoint_ids: &[String],
since_unix_secs: u64,
limit: usize,
) -> Result<Vec<StoredRequestCandidate>, DataLayerError> {
Self::list_finalized_by_endpoint_ids_since(self, endpoint_ids, since_unix_secs, limit).await
}
async fn list_by_provider_id(
&self,
provider_id: &str,
limit: usize,
) -> Result<Vec<StoredRequestCandidate>, DataLayerError> {
Self::list_by_provider_id(self, provider_id, limit).await
}
async fn count_finalized_statuses_by_endpoint_ids_since(
&self,
endpoint_ids: &[String],
since_unix_secs: u64,
) -> Result<Vec<PublicHealthStatusCount>, DataLayerError> {
Self::count_finalized_statuses_by_endpoint_ids_since(self, endpoint_ids, since_unix_secs)
.await
}
async fn aggregate_finalized_timeline_by_endpoint_ids_since(
&self,
endpoint_ids: &[String],
since_unix_secs: u64,
until_unix_secs: u64,
segments: u32,
) -> Result<Vec<PublicHealthTimelineBucket>, DataLayerError> {
Self::aggregate_finalized_timeline_by_endpoint_ids_since(
self,
endpoint_ids,
since_unix_secs,
until_unix_secs,
segments,
)
.await
}
}
#[async_trait]
impl RequestCandidateWriteRepository for SqlxRequestCandidateReadRepository {
async fn upsert(
&self,
candidate: UpsertRequestCandidateRecord,
) -> Result<StoredRequestCandidate, DataLayerError> {
Self::upsert(self, candidate).await
}
async fn delete_created_before(
&self,
created_before_unix_secs: u64,
limit: usize,
) -> Result<usize, DataLayerError> {
Self::delete_created_before(self, created_before_unix_secs, limit).await
}
}
fn map_request_candidate_row(
@@ -163,6 +678,31 @@ fn map_request_candidate_row(
)
}
fn status_to_database(status: RequestCandidateStatus) -> &'static str {
match status {
RequestCandidateStatus::Available => "available",
RequestCandidateStatus::Unused => "unused",
RequestCandidateStatus::Pending => "pending",
RequestCandidateStatus::Streaming => "streaming",
RequestCandidateStatus::Success => "success",
RequestCandidateStatus::Failed => "failed",
RequestCandidateStatus::Cancelled => "cancelled",
RequestCandidateStatus::Skipped => "skipped",
}
}
fn to_i32(value: u32) -> Result<i32, DataLayerError> {
i32::try_from(value).map_err(|_| {
DataLayerError::UnexpectedValue(format!("request candidate value out of range: {value}"))
})
}
fn to_i32_u64(value: u64) -> Result<i32, DataLayerError> {
i32::try_from(value).map_err(|_| {
DataLayerError::UnexpectedValue(format!("request candidate value out of range: {value}"))
})
}
#[cfg(test)]
mod tests {
use super::SqlxRequestCandidateReadRepository;
@@ -185,5 +725,6 @@ mod tests {
let pool = factory.connect_lazy().expect("pool should build");
let repository = SqlxRequestCandidateReadRepository::new(pool);
let _ = repository.pool();
let _ = repository.transaction_runner();
}
}

147
crates/aether-data/src/repository/candidates/types.rs
View File

@@ -185,6 +185,24 @@ impl StoredRequestCandidate {
}
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct PublicHealthStatusCount {
pub endpoint_id: String,
pub status: RequestCandidateStatus,
pub count: u64,
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct PublicHealthTimelineBucket {
pub endpoint_id: String,
pub segment_idx: u32,
pub total_count: u64,
pub success_count: u64,
pub failed_count: u64,
pub min_created_at_unix_secs: Option<u64>,
pub max_created_at_unix_secs: Option<u64>,
}
#[async_trait]
pub trait RequestCandidateReadRepository: Send + Sync {
async fn list_by_request_id(
@@ -196,15 +214,106 @@ pub trait RequestCandidateReadRepository: Send + Sync {
&self,
limit: usize,
) -> Result<Vec<StoredRequestCandidate>, crate::DataLayerError>;
async fn list_by_provider_id(
&self,
provider_id: &str,
limit: usize,
) -> Result<Vec<StoredRequestCandidate>, crate::DataLayerError>;
async fn list_finalized_by_endpoint_ids_since(
&self,
endpoint_ids: &[String],
since_unix_secs: u64,
limit: usize,
) -> Result<Vec<StoredRequestCandidate>, crate::DataLayerError>;
async fn count_finalized_statuses_by_endpoint_ids_since(
&self,
endpoint_ids: &[String],
since_unix_secs: u64,
) -> Result<Vec<PublicHealthStatusCount>, crate::DataLayerError>;
async fn aggregate_finalized_timeline_by_endpoint_ids_since(
&self,
endpoint_ids: &[String],
since_unix_secs: u64,
until_unix_secs: u64,
segments: u32,
) -> Result<Vec<PublicHealthTimelineBucket>, crate::DataLayerError>;
}
pub trait RequestCandidateRepository: RequestCandidateReadRepository + Send + Sync {}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct UpsertRequestCandidateRecord {
pub id: String,
pub request_id: String,
pub user_id: Option<String>,
pub api_key_id: Option<String>,
pub username: Option<String>,
pub api_key_name: Option<String>,
pub candidate_index: u32,
pub retry_index: u32,
pub provider_id: Option<String>,
pub endpoint_id: Option<String>,
pub key_id: Option<String>,
pub status: RequestCandidateStatus,
pub skip_reason: Option<String>,
pub is_cached: Option<bool>,
pub status_code: Option<u16>,
pub error_type: Option<String>,
pub error_message: Option<String>,
pub latency_ms: Option<u64>,
pub concurrent_requests: Option<u32>,
pub extra_data: Option<serde_json::Value>,
pub required_capabilities: Option<serde_json::Value>,
pub created_at_unix_secs: Option<u64>,
pub started_at_unix_secs: Option<u64>,
pub finished_at_unix_secs: Option<u64>,
}
impl<T> RequestCandidateRepository for T where T: RequestCandidateReadRepository + Send + Sync {}
impl UpsertRequestCandidateRecord {
pub fn validate(&self) -> Result<(), crate::DataLayerError> {
if self.id.trim().is_empty() {
return Err(crate::DataLayerError::InvalidInput(
"request candidate upsert id cannot be empty".to_string(),
));
}
if self.request_id.trim().is_empty() {
return Err(crate::DataLayerError::InvalidInput(
"request candidate upsert request_id cannot be empty".to_string(),
));
}
Ok(())
}
}
#[async_trait]
pub trait RequestCandidateWriteRepository: Send + Sync {
async fn upsert(
&self,
candidate: UpsertRequestCandidateRecord,
) -> Result<StoredRequestCandidate, crate::DataLayerError>;
async fn delete_created_before(
&self,
created_before_unix_secs: u64,
limit: usize,
) -> Result<usize, crate::DataLayerError>;
}
pub trait RequestCandidateRepository:
RequestCandidateReadRepository + RequestCandidateWriteRepository + Send + Sync
{
}
impl<T> RequestCandidateRepository for T where
T: RequestCandidateReadRepository + RequestCandidateWriteRepository + Send + Sync
{
}
#[cfg(test)]
mod tests {
use super::{RequestCandidateStatus, StoredRequestCandidate};
use super::{RequestCandidateStatus, StoredRequestCandidate, UpsertRequestCandidateRecord};
#[test]
fn parses_status_from_database_text() {
@@ -286,4 +395,36 @@ mod tests {
assert!(!RequestCandidateStatus::Pending.is_attempted(None));
assert!(RequestCandidateStatus::Pending.is_attempted(Some(1)));
}
#[test]
fn rejects_invalid_upsert_payload() {
assert!(UpsertRequestCandidateRecord {
id: "".to_string(),
request_id: "".to_string(),
user_id: None,
api_key_id: None,
username: None,
api_key_name: None,
candidate_index: 0,
retry_index: 0,
provider_id: None,
endpoint_id: None,
key_id: None,
status: RequestCandidateStatus::Available,
skip_reason: None,
is_cached: None,
status_code: None,
error_type: None,
error_message: None,
latency_ms: None,
concurrent_requests: None,
extra_data: None,
required_capabilities: None,
created_at_unix_secs: None,
started_at_unix_secs: None,
finished_at_unix_secs: None,
}
.validate()
.is_err());
}
}