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

450
crates/aether-data/src/repository/video_tasks/memory.rs
View File

@@ -4,7 +4,8 @@ use std::sync::RwLock;
use async_trait::async_trait;
use super::types::{
StoredVideoTask, UpsertVideoTask, VideoTaskLookupKey, VideoTaskReadRepository,
StoredVideoTask, UpsertVideoTask, VideoTaskLookupKey, VideoTaskModelCount,
VideoTaskQueryFilter, VideoTaskReadRepository, VideoTaskStatus, VideoTaskStatusCount,
VideoTaskWriteRepository,
};
use crate::DataLayerError;
@@ -47,6 +48,34 @@ impl InMemoryVideoTaskRepository {
task
}
fn matches_filter(task: &StoredVideoTask, filter: &VideoTaskQueryFilter) -> bool {
if let Some(user_id) = filter.user_id.as_deref() {
if task.user_id.as_deref() != Some(user_id) {
return false;
}
}
if let Some(status) = filter.status {
if task.status != status {
return false;
}
}
if let Some(model_substring) = filter.model_substring.as_deref() {
let needle = model_substring.trim().to_ascii_lowercase();
let Some(model) = task.model.as_deref() else {
return false;
};
if !model.to_ascii_lowercase().contains(&needle) {
return false;
}
}
if let Some(client_api_format) = filter.client_api_format.as_deref() {
if task.client_api_format.as_deref() != Some(client_api_format) {
return false;
}
}
true
}
}
#[async_trait]
@@ -92,6 +121,178 @@ impl VideoTaskReadRepository for InMemoryVideoTaskRepository {
tasks.truncate(limit);
Ok(tasks)
}
async fn list_due(
&self,
now_unix_secs: u64,
limit: usize,
) -> Result<Vec<StoredVideoTask>, DataLayerError> {
if limit == 0 {
return Ok(Vec::new());
}
let mut tasks = self
.index
.read()
.expect("video task repository lock")
.by_id
.values()
.filter(|task| {
matches!(
task.status,
super::types::VideoTaskStatus::Submitted
| super::types::VideoTaskStatus::Queued
| super::types::VideoTaskStatus::Processing
) && task.poll_count < task.max_poll_count
&& task
.next_poll_at_unix_secs
.is_some_and(|value| value <= now_unix_secs)
})
.cloned()
.collect::<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))
});
tasks.truncate(limit);
Ok(tasks)
}
async fn list_page(
&self,
filter: &VideoTaskQueryFilter,
offset: usize,
limit: usize,
) -> Result<Vec<StoredVideoTask>, DataLayerError> {
if limit == 0 {
return Ok(Vec::new());
}
let mut tasks = self
.index
.read()
.expect("video task repository lock")
.by_id
.values()
.filter(|task| Self::matches_filter(task, filter))
.cloned()
.collect::<Vec<_>>();
tasks.sort_by(|left, right| {
right
.created_at_unix_secs
.cmp(&left.created_at_unix_secs)
.then_with(|| right.updated_at_unix_secs.cmp(&left.updated_at_unix_secs))
});
Ok(tasks.into_iter().skip(offset).take(limit).collect())
}
async fn count(&self, filter: &VideoTaskQueryFilter) -> Result<u64, DataLayerError> {
Ok(self
.index
.read()
.expect("video task repository lock")
.by_id
.values()
.filter(|task| Self::matches_filter(task, filter))
.count() as u64)
}
async fn count_by_status(
&self,
filter: &VideoTaskQueryFilter,
) -> Result<Vec<VideoTaskStatusCount>, DataLayerError> {
let mut counts = BTreeMap::<VideoTaskStatus, u64>::new();
for task in self
.index
.read()
.expect("video task repository lock")
.by_id
.values()
.filter(|task| Self::matches_filter(task, filter))
{
*counts.entry(task.status).or_default() += 1;
}
Ok(counts
.into_iter()
.map(|(status, count)| VideoTaskStatusCount { status, count })
.collect())
}
async fn count_distinct_users(
&self,
filter: &VideoTaskQueryFilter,
) -> Result<u64, DataLayerError> {
let index = self.index.read().expect("video task repository lock");
let users = index
.by_id
.values()
.filter(|task| Self::matches_filter(task, filter))
.filter_map(|task| task.user_id.as_deref())
.map(str::trim)
.filter(|value| !value.is_empty())
.collect::<std::collections::BTreeSet<_>>();
Ok(users.len() as u64)
}
async fn top_models(
&self,
filter: &VideoTaskQueryFilter,
limit: usize,
) -> Result<Vec<VideoTaskModelCount>, DataLayerError> {
if limit == 0 {
return Ok(Vec::new());
}
let mut counts = BTreeMap::<String, u64>::new();
for task in self
.index
.read()
.expect("video task repository lock")
.by_id
.values()
.filter(|task| Self::matches_filter(task, filter))
{
let Some(model) = task.model.as_deref() else {
continue;
};
if model.trim().is_empty() {
continue;
}
*counts.entry(model.to_string()).or_default() += 1;
}
let mut models = counts
.into_iter()
.map(|(model, count)| VideoTaskModelCount { model, count })
.collect::<Vec<_>>();
models.sort_by(|left, right| {
right
.count
.cmp(&left.count)
.then_with(|| left.model.cmp(&right.model))
});
models.truncate(limit);
Ok(models)
}
async fn count_created_since(
&self,
filter: &VideoTaskQueryFilter,
created_since_unix_secs: u64,
) -> Result<u64, DataLayerError> {
Ok(self
.index
.read()
.expect("video task repository lock")
.by_id
.values()
.filter(|task| {
Self::matches_filter(task, filter)
&& task.created_at_unix_secs >= created_since_unix_secs
})
.count() as u64)
}
}
#[async_trait]
@@ -100,14 +301,76 @@ impl VideoTaskWriteRepository for InMemoryVideoTaskRepository {
let mut index = self.index.write().expect("video task repository lock");
Ok(Self::store_locked(&mut index, task.into_stored()))
}
async fn update_if_active(
&self,
task: UpsertVideoTask,
) -> Result<Option<StoredVideoTask>, DataLayerError> {
let mut index = self.index.write().expect("video task repository lock");
let Some(existing) = index.by_id.get(&task.id) else {
return Ok(None);
};
if !existing.status.is_active() {
return Ok(None);
}
Ok(Some(Self::store_locked(&mut index, task.into_stored())))
}
async fn claim_due(
&self,
now_unix_secs: u64,
claim_until_unix_secs: u64,
limit: usize,
) -> Result<Vec<StoredVideoTask>, DataLayerError> {
if limit == 0 {
return Ok(Vec::new());
}
let mut index = self.index.write().expect("video task repository lock");
let mut due_ids = index
.by_id
.values()
.filter(|task| {
matches!(
task.status,
VideoTaskStatus::Submitted
| VideoTaskStatus::Queued
| VideoTaskStatus::Processing
) && task.poll_count < task.max_poll_count
&& task
.next_poll_at_unix_secs
.is_some_and(|value| value <= now_unix_secs)
})
.map(|task| task.id.clone())
.collect::<Vec<_>>();
due_ids.sort_by(|left_id, right_id| {
let left = index.by_id.get(left_id).expect("task should exist");
let right = index.by_id.get(right_id).expect("task should exist");
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))
});
due_ids.truncate(limit);
let mut claimed = Vec::with_capacity(due_ids.len());
for id in due_ids {
let Some(task) = index.by_id.get_mut(&id) else {
continue;
};
task.next_poll_at_unix_secs = Some(claim_until_unix_secs);
task.updated_at_unix_secs = now_unix_secs.max(task.updated_at_unix_secs);
claimed.push(task.clone());
}
Ok(claimed)
}
}
#[cfg(test)]
mod tests {
use super::InMemoryVideoTaskRepository;
use crate::repository::video_tasks::{
UpsertVideoTask, VideoTaskLookupKey, VideoTaskReadRepository, VideoTaskStatus,
VideoTaskWriteRepository,
UpsertVideoTask, VideoTaskLookupKey, VideoTaskQueryFilter, VideoTaskReadRepository,
VideoTaskStatus, VideoTaskWriteRepository,
};
fn sample_task(
@@ -118,19 +381,41 @@ mod tests {
UpsertVideoTask {
id: id.to_string(),
short_id: Some(format!("short-{id}")),
request_id: format!("request-{id}"),
user_id: Some("user-1".to_string()),
api_key_id: Some("api-key-1".to_string()),
username: Some("user".to_string()),
api_key_name: Some("primary".to_string()),
external_task_id: Some(format!("ext-{id}")),
provider_id: Some("provider-1".to_string()),
endpoint_id: Some("endpoint-1".to_string()),
key_id: Some("provider-key-1".to_string()),
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("hello".to_string()),
original_request_body: Some(serde_json::json!({"prompt": "hello"})),
duration_seconds: Some(4),
resolution: Some("720p".to_string()),
aspect_ratio: Some("16:9".to_string()),
size: Some("1280x720".to_string()),
status,
progress_percent: 0,
progress_message: None,
retry_count: 0,
poll_interval_seconds: 10,
next_poll_at_unix_secs: Some(updated_at_unix_secs),
poll_count: 0,
max_poll_count: 360,
created_at_unix_secs: updated_at_unix_secs.saturating_sub(10),
submitted_at_unix_secs: Some(updated_at_unix_secs.saturating_sub(10)),
completed_at_unix_secs: None,
updated_at_unix_secs,
error_code: None,
error_message: None,
video_url: None,
request_metadata: None,
}
}
@@ -193,19 +478,41 @@ mod tests {
repo.upsert(UpsertVideoTask {
id: "task-1".to_string(),
short_id: Some("short-task-1b".to_string()),
request_id: "request-task-1b".to_string(),
user_id: Some("user-2".to_string()),
api_key_id: Some("api-key-2".to_string()),
username: Some("user-2".to_string()),
api_key_name: Some("secondary".to_string()),
external_task_id: Some("ext-task-1b".to_string()),
provider_id: Some("provider-2".to_string()),
endpoint_id: Some("endpoint-2".to_string()),
key_id: Some("provider-key-2".to_string()),
client_api_format: Some("gemini:video".to_string()),
provider_api_format: Some("gemini:video".to_string()),
format_converted: false,
model: Some("veo-3".to_string()),
prompt: Some("remix".to_string()),
original_request_body: Some(serde_json::json!({"prompt": "remix"})),
duration_seconds: Some(8),
resolution: Some("1080p".to_string()),
aspect_ratio: Some("16:9".to_string()),
size: Some("720p".to_string()),
status: VideoTaskStatus::Processing,
progress_percent: 50,
progress_message: Some("processing".to_string()),
retry_count: 1,
poll_interval_seconds: 10,
next_poll_at_unix_secs: Some(200),
poll_count: 2,
max_poll_count: 360,
created_at_unix_secs: 150,
submitted_at_unix_secs: Some(150),
completed_at_unix_secs: None,
updated_at_unix_secs: 200,
error_code: None,
error_message: None,
video_url: None,
request_metadata: None,
})
.await
.expect("upsert should succeed");
@@ -229,4 +536,141 @@ mod tests {
.expect("find should succeed")
.is_some());
}
#[tokio::test]
async fn list_due_returns_due_active_tasks_in_next_poll_order() {
let repo = InMemoryVideoTaskRepository::default();
repo.upsert(sample_task("task-1", VideoTaskStatus::Submitted, 300))
.await
.expect("upsert should succeed");
repo.upsert(sample_task("task-2", VideoTaskStatus::Processing, 100))
.await
.expect("upsert should succeed");
repo.upsert(UpsertVideoTask {
next_poll_at_unix_secs: Some(500),
..sample_task("task-3", VideoTaskStatus::Queued, 200)
})
.await
.expect("upsert should succeed");
let due = repo
.list_due(300, 10)
.await
.expect("list due should succeed");
assert_eq!(due.len(), 2);
assert_eq!(due[0].id, "task-2");
assert_eq!(due[1].id, "task-1");
}
#[tokio::test]
async fn update_if_active_skips_terminal_tasks() {
let repo = InMemoryVideoTaskRepository::default();
repo.upsert(sample_task("task-1", VideoTaskStatus::Completed, 100))
.await
.expect("upsert should succeed");
let updated = repo
.update_if_active(UpsertVideoTask {
progress_percent: 100,
..sample_task("task-1", VideoTaskStatus::Completed, 200)
})
.await
.expect("update should succeed");
assert!(updated.is_none());
}
#[tokio::test]
async fn list_page_and_stats_apply_filters() {
let repo = InMemoryVideoTaskRepository::default();
repo.upsert(sample_task("task-1", VideoTaskStatus::Submitted, 100))
.await
.expect("upsert should succeed");
repo.upsert(UpsertVideoTask {
model: Some("veo-3-fast".to_string()),
user_id: Some("user-2".to_string()),
client_api_format: Some("gemini:video".to_string()),
created_at_unix_secs: 250,
updated_at_unix_secs: 250,
..sample_task("task-2", VideoTaskStatus::Completed, 250)
})
.await
.expect("upsert should succeed");
repo.upsert(UpsertVideoTask {
model: Some("veo-3-fast".to_string()),
user_id: Some("user-2".to_string()),
client_api_format: Some("gemini:video".to_string()),
created_at_unix_secs: 260,
updated_at_unix_secs: 260,
..sample_task("task-3", VideoTaskStatus::Completed, 260)
})
.await
.expect("upsert should succeed");
let filter = VideoTaskQueryFilter {
user_id: Some("user-2".to_string()),
status: Some(VideoTaskStatus::Completed),
model_substring: Some("veo".to_string()),
client_api_format: Some("gemini:video".to_string()),
};
let page = repo
.list_page(&filter, 0, 10)
.await
.expect("list page should succeed");
assert_eq!(page.len(), 2);
assert_eq!(page[0].id, "task-3");
assert_eq!(page[1].id, "task-2");
let count = repo.count(&filter).await.expect("count should succeed");
assert_eq!(count, 2);
let by_status = repo
.count_by_status(&filter)
.await
.expect("status count should succeed");
assert_eq!(by_status.len(), 1);
assert_eq!(by_status[0].status, VideoTaskStatus::Completed);
assert_eq!(by_status[0].count, 2);
let top_models = repo
.top_models(&filter, 10)
.await
.expect("top models should succeed");
assert_eq!(top_models.len(), 1);
assert_eq!(top_models[0].model, "veo-3-fast");
assert_eq!(top_models[0].count, 2);
let today_count = repo
.count_created_since(&filter, 255)
.await
.expect("today count should succeed");
assert_eq!(today_count, 1);
}
#[tokio::test]
async fn claim_due_advances_claimed_tasks_until_claim_deadline() {
let repo = InMemoryVideoTaskRepository::default();
repo.upsert(sample_task("task-1", VideoTaskStatus::Submitted, 100))
.await
.expect("upsert should succeed");
repo.upsert(sample_task("task-2", VideoTaskStatus::Processing, 90))
.await
.expect("upsert should succeed");
let claimed = repo
.claim_due(100, 130, 1)
.await
.expect("claim should succeed");
assert_eq!(claimed.len(), 1);
assert_eq!(claimed[0].id, "task-2");
assert_eq!(claimed[0].next_poll_at_unix_secs, Some(130));
let remaining = repo
.list_due(100, 10)
.await
.expect("list due should succeed");
assert_eq!(remaining.len(), 1);
assert_eq!(remaining[0].id, "task-1");
}
}

7
crates/aether-data/src/repository/video_tasks/mod.rs
View File

@@ -3,8 +3,9 @@ mod sql;
mod types;
pub use memory::InMemoryVideoTaskRepository;
pub use sql::SqlxVideoTaskReadRepository;
pub use sql::{SqlxVideoTaskReadRepository, SqlxVideoTaskRepository};
pub use types::{
StoredVideoTask, UpsertVideoTask, VideoTaskLookupKey, VideoTaskReadRepository,
VideoTaskRepository, VideoTaskStatus, VideoTaskWriteRepository,
StoredVideoTask, UpsertVideoTask, VideoTaskLookupKey, VideoTaskModelCount,
VideoTaskQueryFilter, VideoTaskReadRepository, VideoTaskRepository, VideoTaskStatus,
VideoTaskStatusCount, VideoTaskWriteRepository,
};

1001
crates/aether-data/src/repository/video_tasks/sql.rs
View File

File diff suppressed because it is too large Load Diff

426
crates/aether-data/src/repository/video_tasks/types.rs
View File

@@ -1,6 +1,9 @@
use async_trait::async_trait;
use serde_json::Value;
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
#[derive(
Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, serde::Serialize, serde::Deserialize,
)]
pub enum VideoTaskStatus {
Pending,
Submitted,
@@ -43,71 +46,166 @@ impl VideoTaskStatus {
pub struct StoredVideoTask {
pub id: String,
pub short_id: Option<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 external_task_id: Option<String>,
pub provider_id: Option<String>,
pub endpoint_id: Option<String>,
pub key_id: Option<String>,
pub client_api_format: Option<String>,
pub provider_api_format: Option<String>,
pub format_converted: bool,
pub model: Option<String>,
pub prompt: Option<String>,
pub original_request_body: Option<Value>,
pub duration_seconds: Option<u32>,
pub resolution: Option<String>,
pub aspect_ratio: Option<String>,
pub size: Option<String>,
pub status: VideoTaskStatus,
pub progress_percent: u16,
pub progress_message: Option<String>,
pub retry_count: u32,
pub poll_interval_seconds: u32,
pub next_poll_at_unix_secs: Option<u64>,
pub poll_count: u32,
pub max_poll_count: u32,
pub created_at_unix_secs: u64,
pub submitted_at_unix_secs: Option<u64>,
pub completed_at_unix_secs: Option<u64>,
pub updated_at_unix_secs: u64,
pub error_code: Option<String>,
pub error_message: Option<String>,
pub video_url: Option<String>,
pub request_metadata: Option<Value>,
}
impl StoredVideoTask {
#[allow(clippy::too_many_arguments)]
pub fn new(
id: String,
short_id: Option<String>,
request_id: String,
user_id: Option<String>,
api_key_id: Option<String>,
username: Option<String>,
api_key_name: Option<String>,
external_task_id: Option<String>,
provider_id: Option<String>,
endpoint_id: Option<String>,
key_id: Option<String>,
client_api_format: Option<String>,
provider_api_format: Option<String>,
format_converted: bool,
model: Option<String>,
prompt: Option<String>,
original_request_body: Option<Value>,
duration_seconds: Option<i32>,
resolution: Option<String>,
aspect_ratio: Option<String>,
size: Option<String>,
status: VideoTaskStatus,
progress_percent: i32,
progress_message: Option<String>,
retry_count: i32,
poll_interval_seconds: i32,
next_poll_at_unix_secs: Option<i64>,
poll_count: i32,
max_poll_count: i32,
created_at_unix_secs: i64,
submitted_at_unix_secs: Option<i64>,
completed_at_unix_secs: Option<i64>,
updated_at_unix_secs: i64,
error_code: Option<String>,
error_message: Option<String>,
video_url: Option<String>,
request_metadata: Option<Value>,
) -> Result<Self, crate::DataLayerError> {
let progress_percent = u16::try_from(progress_percent).map_err(|_| {
crate::DataLayerError::UnexpectedValue(format!(
"invalid progress_percent: {progress_percent}"
))
})?;
let retry_count = u32::try_from(retry_count).map_err(|_| {
crate::DataLayerError::UnexpectedValue(format!("invalid retry_count: {retry_count}"))
})?;
let poll_interval_seconds = u32::try_from(poll_interval_seconds).map_err(|_| {
crate::DataLayerError::UnexpectedValue(format!(
"invalid poll_interval_seconds: {poll_interval_seconds}"
))
})?;
let next_poll_at_unix_secs =
coerce_optional_unix_secs(next_poll_at_unix_secs, "next_poll_at_unix_secs")?;
let poll_count = u32::try_from(poll_count).map_err(|_| {
crate::DataLayerError::UnexpectedValue(format!("invalid poll_count: {poll_count}"))
})?;
let max_poll_count = u32::try_from(max_poll_count).map_err(|_| {
crate::DataLayerError::UnexpectedValue(format!(
"invalid max_poll_count: {max_poll_count}"
))
})?;
let created_at_unix_secs = u64::try_from(created_at_unix_secs).map_err(|_| {
crate::DataLayerError::UnexpectedValue(format!(
"invalid created_at_unix_secs: {created_at_unix_secs}"
))
})?;
let submitted_at_unix_secs =
coerce_optional_unix_secs(submitted_at_unix_secs, "submitted_at_unix_secs")?;
let completed_at_unix_secs =
coerce_optional_unix_secs(completed_at_unix_secs, "completed_at_unix_secs")?;
let updated_at_unix_secs = u64::try_from(updated_at_unix_secs).map_err(|_| {
crate::DataLayerError::UnexpectedValue(format!(
"invalid updated_at_unix_secs: {updated_at_unix_secs}"
))
})?;
let duration_seconds = match duration_seconds {
Some(value) => Some(u32::try_from(value).map_err(|_| {
crate::DataLayerError::UnexpectedValue(format!("invalid duration_seconds: {value}"))
})?),
None => None,
};
Ok(Self {
id,
short_id,
request_id,
user_id,
api_key_id,
username,
api_key_name,
external_task_id,
provider_id,
endpoint_id,
key_id,
client_api_format,
provider_api_format,
format_converted,
model,
prompt,
original_request_body,
duration_seconds,
resolution,
aspect_ratio,
size,
status,
progress_percent,
progress_message,
retry_count,
poll_interval_seconds,
next_poll_at_unix_secs,
poll_count,
max_poll_count,
created_at_unix_secs,
submitted_at_unix_secs,
completed_at_unix_secs,
updated_at_unix_secs,
error_code,
error_message,
video_url,
request_metadata,
})
}
}
@@ -116,19 +214,41 @@ impl StoredVideoTask {
pub struct UpsertVideoTask {
pub id: String,
pub short_id: Option<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 external_task_id: Option<String>,
pub provider_id: Option<String>,
pub endpoint_id: Option<String>,
pub key_id: Option<String>,
pub client_api_format: Option<String>,
pub provider_api_format: Option<String>,
pub format_converted: bool,
pub model: Option<String>,
pub prompt: Option<String>,
pub original_request_body: Option<Value>,
pub duration_seconds: Option<u32>,
pub resolution: Option<String>,
pub aspect_ratio: Option<String>,
pub size: Option<String>,
pub status: VideoTaskStatus,
pub progress_percent: u16,
pub progress_message: Option<String>,
pub retry_count: u32,
pub poll_interval_seconds: u32,
pub next_poll_at_unix_secs: Option<u64>,
pub poll_count: u32,
pub max_poll_count: u32,
pub created_at_unix_secs: u64,
pub submitted_at_unix_secs: Option<u64>,
pub completed_at_unix_secs: Option<u64>,
pub updated_at_unix_secs: u64,
pub error_code: Option<String>,
pub error_message: Option<String>,
pub video_url: Option<String>,
pub request_metadata: Option<Value>,
}
impl UpsertVideoTask {
@@ -136,19 +256,85 @@ impl UpsertVideoTask {
StoredVideoTask {
id: self.id,
short_id: self.short_id,
request_id: self.request_id,
user_id: self.user_id,
api_key_id: self.api_key_id,
username: self.username,
api_key_name: self.api_key_name,
external_task_id: self.external_task_id,
provider_id: self.provider_id,
endpoint_id: self.endpoint_id,
key_id: self.key_id,
client_api_format: self.client_api_format,
provider_api_format: self.provider_api_format,
format_converted: self.format_converted,
model: self.model,
prompt: self.prompt,
original_request_body: self.original_request_body,
duration_seconds: self.duration_seconds,
resolution: self.resolution,
aspect_ratio: self.aspect_ratio,
size: self.size,
status: self.status,
progress_percent: self.progress_percent,
progress_message: self.progress_message,
retry_count: self.retry_count,
poll_interval_seconds: self.poll_interval_seconds,
next_poll_at_unix_secs: self.next_poll_at_unix_secs,
poll_count: self.poll_count,
max_poll_count: self.max_poll_count,
created_at_unix_secs: self.created_at_unix_secs,
submitted_at_unix_secs: self.submitted_at_unix_secs,
completed_at_unix_secs: self.completed_at_unix_secs,
updated_at_unix_secs: self.updated_at_unix_secs,
error_code: self.error_code,
error_message: self.error_message,
video_url: self.video_url,
request_metadata: self.request_metadata,
}
}
}
impl From<StoredVideoTask> for UpsertVideoTask {
fn from(task: StoredVideoTask) -> Self {
Self {
id: task.id,
short_id: task.short_id,
request_id: task.request_id,
user_id: task.user_id,
api_key_id: task.api_key_id,
username: task.username,
api_key_name: task.api_key_name,
external_task_id: task.external_task_id,
provider_id: task.provider_id,
endpoint_id: task.endpoint_id,
key_id: task.key_id,
client_api_format: task.client_api_format,
provider_api_format: task.provider_api_format,
format_converted: task.format_converted,
model: task.model,
prompt: task.prompt,
original_request_body: task.original_request_body,
duration_seconds: task.duration_seconds,
resolution: task.resolution,
aspect_ratio: task.aspect_ratio,
size: task.size,
status: task.status,
progress_percent: task.progress_percent,
progress_message: task.progress_message,
retry_count: task.retry_count,
poll_interval_seconds: task.poll_interval_seconds,
next_poll_at_unix_secs: task.next_poll_at_unix_secs,
poll_count: task.poll_count,
max_poll_count: task.max_poll_count,
created_at_unix_secs: task.created_at_unix_secs,
submitted_at_unix_secs: task.submitted_at_unix_secs,
completed_at_unix_secs: task.completed_at_unix_secs,
updated_at_unix_secs: task.updated_at_unix_secs,
error_code: task.error_code,
error_message: task.error_message,
video_url: task.video_url,
request_metadata: task.request_metadata,
}
}
}
@@ -163,6 +349,26 @@ pub enum VideoTaskLookupKey<'a> {
},
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct VideoTaskQueryFilter {
pub user_id: Option<String>,
pub status: Option<VideoTaskStatus>,
pub model_substring: Option<String>,
pub client_api_format: Option<String>,
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct VideoTaskStatusCount {
pub status: VideoTaskStatus,
pub count: u64,
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct VideoTaskModelCount {
pub model: String,
pub count: u64,
}
#[async_trait]
pub trait VideoTaskReadRepository: Send + Sync {
async fn find(
@@ -174,12 +380,61 @@ pub trait VideoTaskReadRepository: Send + Sync {
&self,
limit: usize,
) -> Result<Vec<StoredVideoTask>, crate::DataLayerError>;
async fn list_due(
&self,
now_unix_secs: u64,
limit: usize,
) -> Result<Vec<StoredVideoTask>, crate::DataLayerError>;
async fn list_page(
&self,
filter: &VideoTaskQueryFilter,
offset: usize,
limit: usize,
) -> Result<Vec<StoredVideoTask>, crate::DataLayerError>;
async fn count(&self, filter: &VideoTaskQueryFilter) -> Result<u64, crate::DataLayerError>;
async fn count_by_status(
&self,
filter: &VideoTaskQueryFilter,
) -> Result<Vec<VideoTaskStatusCount>, crate::DataLayerError>;
async fn count_distinct_users(
&self,
filter: &VideoTaskQueryFilter,
) -> Result<u64, crate::DataLayerError>;
async fn top_models(
&self,
filter: &VideoTaskQueryFilter,
limit: usize,
) -> Result<Vec<VideoTaskModelCount>, crate::DataLayerError>;
async fn count_created_since(
&self,
filter: &VideoTaskQueryFilter,
created_since_unix_secs: u64,
) -> Result<u64, crate::DataLayerError>;
}
#[async_trait]
pub trait VideoTaskWriteRepository: Send + Sync {
async fn upsert(&self, task: UpsertVideoTask)
-> Result<StoredVideoTask, crate::DataLayerError>;
async fn update_if_active(
&self,
task: UpsertVideoTask,
) -> Result<Option<StoredVideoTask>, crate::DataLayerError>;
async fn claim_due(
&self,
now_unix_secs: u64,
claim_until_unix_secs: u64,
limit: usize,
) -> Result<Vec<StoredVideoTask>, crate::DataLayerError>;
}
pub trait VideoTaskRepository:
@@ -192,10 +447,103 @@ impl<T> VideoTaskRepository for T where
{
}
fn coerce_optional_unix_secs(
value: Option<i64>,
field: &str,
) -> Result<Option<u64>, crate::DataLayerError> {
match value {
Some(value) => Ok(Some(u64::try_from(value).map_err(|_| {
crate::DataLayerError::UnexpectedValue(format!("invalid {field}: {value}"))
})?)),
None => Ok(None),
}
}
#[cfg(test)]
mod tests {
use super::{StoredVideoTask, VideoTaskStatus};
#[allow(clippy::type_complexity)]
fn base_new_args() -> (
String,
Option<String>,
String,
Option<String>,
Option<String>,
Option<String>,
Option<String>,
Option<String>,
Option<String>,
Option<String>,
Option<String>,
Option<String>,
Option<String>,
bool,
Option<String>,
Option<String>,
Option<serde_json::Value>,
Option<i32>,
Option<String>,
Option<String>,
Option<String>,
VideoTaskStatus,
i32,
Option<String>,
i32,
i32,
Option<i64>,
i32,
i32,
i64,
Option<i64>,
Option<i64>,
i64,
Option<String>,
Option<String>,
Option<String>,
Option<serde_json::Value>,
) {
(
"task-1".to_string(),
None,
"request-1".to_string(),
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
false,
None,
None,
None,
None,
None,
None,
None,
VideoTaskStatus::Submitted,
10,
None,
0,
10,
Some(1),
0,
360,
1,
None,
None,
1,
None,
None,
None,
None,
)
}
#[test]
fn parses_status_from_database_text() {
assert_eq!(
@@ -211,66 +559,52 @@ mod tests {
#[test]
fn rejects_invalid_numeric_fields() {
let mut args = base_new_args();
args.22 = -1;
assert!(StoredVideoTask::new(
"task-1".to_string(),
None,
None,
None,
None,
None,
None,
None,
VideoTaskStatus::Submitted,
-1,
1,
1,
None,
None,
None
args.0, args.1, args.2, args.3, args.4, args.5, args.6, args.7, args.8, args.9,
args.10, args.11, args.12, args.13, args.14, args.15, args.16, args.17, args.18,
args.19, args.20, args.21, args.22, args.23, args.24, args.25, args.26, args.27,
args.28, args.29, args.30, args.31, args.32, args.33, args.34, args.35, args.36,
)
.is_err());
}
#[test]
fn rejects_negative_updated_at_values() {
let mut args = base_new_args();
args.32 = -1;
assert!(StoredVideoTask::new(
"task-1".to_string(),
None,
None,
None,
None,
None,
None,
None,
VideoTaskStatus::Submitted,
10,
1,
-1,
None,
None,
None
args.0, args.1, args.2, args.3, args.4, args.5, args.6, args.7, args.8, args.9,
args.10, args.11, args.12, args.13, args.14, args.15, args.16, args.17, args.18,
args.19, args.20, args.21, args.22, args.23, args.24, args.25, args.26, args.27,
args.28, args.29, args.30, args.31, args.32, args.33, args.34, args.35, args.36,
)
.is_err());
}
#[test]
fn rejects_negative_created_at_values() {
let mut args = base_new_args();
args.29 = -1;
assert!(StoredVideoTask::new(
"task-1".to_string(),
None,
None,
None,
None,
None,
None,
None,
VideoTaskStatus::Submitted,
10,
-1,
1,
None,
None,
None
args.0, args.1, args.2, args.3, args.4, args.5, args.6, args.7, args.8, args.9,
args.10, args.11, args.12, args.13, args.14, args.15, args.16, args.17, args.18,
args.19, args.20, args.21, args.22, args.23, args.24, args.25, args.26, args.27,
args.28, args.29, args.30, args.31, args.32, args.33, args.34, args.35, args.36,
)
.is_err());
}
#[test]
fn rejects_negative_optional_completed_at_values() {
let mut args = base_new_args();
args.31 = Some(-1);
assert!(StoredVideoTask::new(
args.0, args.1, args.2, args.3, args.4, args.5, args.6, args.7, args.8, args.9,
args.10, args.11, args.12, args.13, args.14, args.15, args.16, args.17, args.18,
args.19, args.20, args.21, args.22, args.23, args.24, args.25, args.26, args.27,
args.28, args.29, args.30, args.31, args.32, args.33, args.34, args.35, args.36,
)
.is_err());
}