use async_trait::async_trait; use futures_util::{stream::TryStream, TryStreamExt}; use sqlx::{postgres::PgRow, PgPool, Postgres, QueryBuilder, Row}; use aether_data_contracts::repository::video_tasks::{ StoredVideoTask, UpsertVideoTask, VideoTaskLookupKey, VideoTaskModelCount, VideoTaskQueryFilter, VideoTaskReadRepository, VideoTaskStatus, VideoTaskStatusCount, VideoTaskWriteRepository, }; use aether_data_contracts::DataLayerError; use crate::error::SqlxResultExt; const SELECT_VIDEO_TASK_COLUMNS_PREFIX: &str = r#" 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, "#; const SELECT_VIDEO_TASK_COLUMNS_SUFFIX: &str = r#" status, progress_percent, progress_message, retry_count, poll_interval_seconds, CAST(EXTRACT(EPOCH FROM next_poll_at) AS BIGINT) AS next_poll_at_unix_secs, poll_count, max_poll_count, CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_ms, CAST(EXTRACT(EPOCH FROM submitted_at) AS BIGINT) AS submitted_at_unix_secs, CAST(EXTRACT(EPOCH FROM completed_at) AS BIGINT) AS completed_at_unix_secs, CAST(EXTRACT(EPOCH FROM updated_at) AS BIGINT) AS updated_at_unix_secs, error_code, error_message, video_url, request_metadata "#; fn select_video_task_columns( prompt_sql: &str, original_request_body_sql: &str, duration_seconds_sql: &str, resolution_sql: &str, aspect_ratio_sql: &str, size_sql: &str, ) -> String { format!( "{SELECT_VIDEO_TASK_COLUMNS_PREFIX} {prompt_sql} AS prompt, {original_request_body_sql} AS original_request_body, {duration_seconds_sql} AS duration_seconds, {resolution_sql} AS resolution, {aspect_ratio_sql} AS aspect_ratio, {size_sql} AS size, {SELECT_VIDEO_TASK_COLUMNS_SUFFIX}" ) } fn select_video_task_full_columns() -> String { select_video_task_columns( "prompt", "original_request_body", "duration_seconds", "resolution", "aspect_ratio", "size", ) } fn select_video_task_claim_columns() -> String { select_video_task_columns( "prompt", "NULL::jsonb", "duration_seconds", "resolution", "aspect_ratio", "size", ) } fn select_video_task_sql(where_clause: &str) -> String { format!( "SELECT\n{}\nFROM video_tasks\n{where_clause}\n", select_video_task_full_columns() ) } fn find_by_id_sql() -> String { select_video_task_sql("WHERE id = $1\nLIMIT 1") } fn find_by_id_for_user_sql() -> String { select_video_task_sql("WHERE id = $1 AND user_id = $2\nLIMIT 1") } fn find_by_short_id_sql() -> String { select_video_task_sql("WHERE short_id = $1\nLIMIT 1") } fn find_by_short_id_for_user_sql() -> String { select_video_task_sql("WHERE short_id = $1 AND user_id = $2\nLIMIT 1") } fn find_by_user_external_sql() -> String { select_video_task_sql("WHERE user_id = $1 AND external_task_id = $2\nLIMIT 1") } fn list_active_sql() -> String { select_video_task_sql("WHERE status = ANY($1)\nORDER BY updated_at DESC\nLIMIT $2") } fn list_due_sql() -> String { select_video_task_sql( "WHERE status = ANY($1)\n AND next_poll_at IS NOT NULL\n AND next_poll_at <= TO_TIMESTAMP($2)\n AND poll_count < max_poll_count\nORDER BY next_poll_at ASC\nLIMIT $3", ) } fn select_video_task_page_summary_columns() -> &'static str { r#" id, NULL::TEXT AS short_id, request_id, user_id, NULL::TEXT AS api_key_id, username, NULL::TEXT AS api_key_name, external_task_id, provider_id, NULL::TEXT AS endpoint_id, NULL::TEXT AS key_id, NULL::TEXT AS client_api_format, NULL::TEXT AS provider_api_format, FALSE AS format_converted, model, CASE WHEN prompt IS NULL THEN NULL WHEN char_length(prompt) <= 100 THEN prompt ELSE LEFT(prompt, 100) || '...' END AS prompt, NULL::jsonb AS original_request_body, duration_seconds, resolution, aspect_ratio, NULL::TEXT AS size, status, progress_percent, progress_message, 0::INTEGER AS retry_count, 1::INTEGER AS poll_interval_seconds, NULL::BIGINT AS next_poll_at_unix_secs, poll_count, max_poll_count, CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_ms, CAST(EXTRACT(EPOCH FROM submitted_at) AS BIGINT) AS submitted_at_unix_secs, CAST(EXTRACT(EPOCH FROM completed_at) AS BIGINT) AS completed_at_unix_secs, CAST(EXTRACT(EPOCH FROM updated_at) AS BIGINT) AS updated_at_unix_secs, error_code, error_message, video_url, NULL::jsonb AS request_metadata "# } fn claim_due_sql() -> String { let columns = select_video_task_claim_columns(); format!( "WITH due AS ( SELECT id FROM video_tasks WHERE status = ANY($1) AND next_poll_at IS NOT NULL AND next_poll_at <= TO_TIMESTAMP($2) AND poll_count < max_poll_count ORDER BY next_poll_at ASC, updated_at ASC FOR UPDATE SKIP LOCKED LIMIT $3 ) UPDATE video_tasks SET next_poll_at = TO_TIMESTAMP($4), updated_at = TO_TIMESTAMP($5) WHERE id IN (SELECT id FROM due) RETURNING {columns} " ) } fn upsert_sql() -> String { let columns = select_video_task_full_columns(); format!( "INSERT INTO video_tasks ( 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, poll_count, max_poll_count, video_url, error_code, error_message, request_metadata, created_at, submitted_at, completed_at, updated_at ) VALUES ( $1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, $18, $19, $20, $21, $22, $23, $24, $25, $26, TO_TIMESTAMP($27), $28, $29, $30, $31, $32, $33, TO_TIMESTAMP($34), TO_TIMESTAMP($35), TO_TIMESTAMP($36), TO_TIMESTAMP($37) ) ON CONFLICT (id) DO UPDATE SET short_id = EXCLUDED.short_id, request_id = EXCLUDED.request_id, user_id = EXCLUDED.user_id, api_key_id = EXCLUDED.api_key_id, username = EXCLUDED.username, api_key_name = EXCLUDED.api_key_name, external_task_id = EXCLUDED.external_task_id, provider_id = EXCLUDED.provider_id, endpoint_id = EXCLUDED.endpoint_id, key_id = EXCLUDED.key_id, client_api_format = EXCLUDED.client_api_format, provider_api_format = EXCLUDED.provider_api_format, format_converted = EXCLUDED.format_converted, model = EXCLUDED.model, prompt = EXCLUDED.prompt, original_request_body = EXCLUDED.original_request_body, duration_seconds = EXCLUDED.duration_seconds, resolution = EXCLUDED.resolution, aspect_ratio = EXCLUDED.aspect_ratio, size = EXCLUDED.size, status = EXCLUDED.status, progress_percent = EXCLUDED.progress_percent, progress_message = EXCLUDED.progress_message, retry_count = EXCLUDED.retry_count, poll_interval_seconds = EXCLUDED.poll_interval_seconds, next_poll_at = EXCLUDED.next_poll_at, poll_count = EXCLUDED.poll_count, max_poll_count = EXCLUDED.max_poll_count, video_url = EXCLUDED.video_url, error_code = EXCLUDED.error_code, error_message = EXCLUDED.error_message, request_metadata = EXCLUDED.request_metadata, created_at = COALESCE(video_tasks.created_at, EXCLUDED.created_at), submitted_at = EXCLUDED.submitted_at, completed_at = EXCLUDED.completed_at, updated_at = EXCLUDED.updated_at WHERE video_tasks.short_id IS NOT DISTINCT FROM EXCLUDED.short_id AND video_tasks.request_id IS NOT DISTINCT FROM EXCLUDED.request_id AND video_tasks.user_id IS NOT DISTINCT FROM EXCLUDED.user_id AND video_tasks.api_key_id IS NOT DISTINCT FROM EXCLUDED.api_key_id AND video_tasks.external_task_id IS NOT DISTINCT FROM EXCLUDED.external_task_id AND video_tasks.provider_id IS NOT DISTINCT FROM EXCLUDED.provider_id AND video_tasks.endpoint_id IS NOT DISTINCT FROM EXCLUDED.endpoint_id AND video_tasks.key_id IS NOT DISTINCT FROM EXCLUDED.key_id AND video_tasks.client_api_format IS NOT DISTINCT FROM EXCLUDED.client_api_format AND video_tasks.provider_api_format IS NOT DISTINCT FROM EXCLUDED.provider_api_format AND video_tasks.format_converted IS NOT DISTINCT FROM EXCLUDED.format_converted AND video_tasks.model IS NOT DISTINCT FROM EXCLUDED.model AND video_tasks.duration_seconds IS NOT DISTINCT FROM EXCLUDED.duration_seconds AND video_tasks.resolution IS NOT DISTINCT FROM EXCLUDED.resolution AND video_tasks.aspect_ratio IS NOT DISTINCT FROM EXCLUDED.aspect_ratio AND video_tasks.size IS NOT DISTINCT FROM EXCLUDED.size RETURNING {columns} " ) } fn update_if_active_sql() -> String { let columns = select_video_task_full_columns(); format!( "UPDATE video_tasks SET short_id = $2, request_id = $3, user_id = $4, api_key_id = $5, username = $6, api_key_name = $7, external_task_id = $8, provider_id = $9, endpoint_id = $10, key_id = $11, client_api_format = $12, provider_api_format = $13, format_converted = $14, model = $15, prompt = $16, original_request_body = $17, duration_seconds = $18, resolution = $19, aspect_ratio = $20, size = $21, status = $22, progress_percent = $23, progress_message = $24, retry_count = $25, poll_interval_seconds = $26, next_poll_at = TO_TIMESTAMP($27), poll_count = $28, max_poll_count = $29, video_url = $30, error_code = $31, error_message = $32, request_metadata = $33, created_at = COALESCE(created_at, TO_TIMESTAMP($34)), submitted_at = TO_TIMESTAMP($35), completed_at = TO_TIMESTAMP($36), updated_at = TO_TIMESTAMP($37) WHERE id = $1 AND status = ANY($38) AND short_id IS NOT DISTINCT FROM $2 AND request_id IS NOT DISTINCT FROM $3 AND user_id IS NOT DISTINCT FROM $4 AND api_key_id IS NOT DISTINCT FROM $5 AND external_task_id IS NOT DISTINCT FROM $8 AND provider_id IS NOT DISTINCT FROM $9 AND endpoint_id IS NOT DISTINCT FROM $10 AND key_id IS NOT DISTINCT FROM $11 AND client_api_format IS NOT DISTINCT FROM $12 AND provider_api_format IS NOT DISTINCT FROM $13 AND format_converted IS NOT DISTINCT FROM $14 AND model IS NOT DISTINCT FROM $15 AND duration_seconds IS NOT DISTINCT FROM $18 AND resolution IS NOT DISTINCT FROM $19 AND aspect_ratio IS NOT DISTINCT FROM $20 AND size IS NOT DISTINCT FROM $21 RETURNING {columns} " ) } #[derive(Debug, Clone)] pub struct SqlxVideoTaskRepository { pool: PgPool, } pub type SqlxVideoTaskReadRepository = SqlxVideoTaskRepository; impl SqlxVideoTaskRepository { pub fn new(pool: PgPool) -> Self { Self { pool } } pub fn pool(&self) -> &PgPool { &self.pool } pub async fn find( &self, key: VideoTaskLookupKey<'_>, ) -> Result, DataLayerError> { match key { VideoTaskLookupKey::Id(id) => self.find_by_id(id).await, VideoTaskLookupKey::ShortId(short_id) => self.find_by_short_id(short_id).await, VideoTaskLookupKey::UserExternal { user_id, external_task_id, } => self.find_by_user_external(user_id, external_task_id).await, } } pub async fn find_for_user( &self, key: VideoTaskLookupKey<'_>, user_id: &str, ) -> Result, DataLayerError> { match key { VideoTaskLookupKey::Id(id) => self.find_by_id_for_user(id, user_id).await, VideoTaskLookupKey::ShortId(short_id) => { self.find_by_short_id_for_user(short_id, user_id).await } VideoTaskLookupKey::UserExternal { user_id: lookup_user_id, external_task_id, } if lookup_user_id == user_id => { self.find_by_user_external(user_id, external_task_id).await } VideoTaskLookupKey::UserExternal { .. } => Ok(None), } } pub async fn find_by_id(&self, id: &str) -> Result, DataLayerError> { let sql = find_by_id_sql(); let row = sqlx::query(&sql) .bind(id) .fetch_optional(&self.pool) .await .map_postgres_err()?; row.as_ref().map(map_video_task_row).transpose() } pub async fn find_by_id_for_user( &self, id: &str, user_id: &str, ) -> Result, DataLayerError> { let sql = find_by_id_for_user_sql(); let row = sqlx::query(&sql) .bind(id) .bind(user_id) .fetch_optional(&self.pool) .await .map_postgres_err()?; row.as_ref().map(map_video_task_row).transpose() } pub async fn find_by_short_id( &self, short_id: &str, ) -> Result, DataLayerError> { let sql = find_by_short_id_sql(); let row = sqlx::query(&sql) .bind(short_id) .fetch_optional(&self.pool) .await .map_postgres_err()?; row.as_ref().map(map_video_task_row).transpose() } pub async fn find_by_short_id_for_user( &self, short_id: &str, user_id: &str, ) -> Result, DataLayerError> { let sql = find_by_short_id_for_user_sql(); let row = sqlx::query(&sql) .bind(short_id) .bind(user_id) .fetch_optional(&self.pool) .await .map_postgres_err()?; row.as_ref().map(map_video_task_row).transpose() } pub async fn find_by_user_external( &self, user_id: &str, external_task_id: &str, ) -> Result, DataLayerError> { let sql = find_by_user_external_sql(); let row = sqlx::query(&sql) .bind(user_id) .bind(external_task_id) .fetch_optional(&self.pool) .await .map_postgres_err()?; row.as_ref().map(map_video_task_row).transpose() } pub async fn list_active(&self, limit: usize) -> Result, DataLayerError> { if limit == 0 { return Ok(Vec::new()); } let active_statuses = vec!["pending", "submitted", "queued", "processing"]; let sql = list_active_sql(); collect_query_rows( sqlx::query(&sql) .bind(active_statuses) .bind(i64::try_from(limit).map_err(|_| { DataLayerError::UnexpectedValue(format!("invalid active task limit: {limit}")) })?) .fetch(&self.pool), map_video_task_row, ) .await } pub async fn list_due( &self, now_unix_secs: u64, limit: usize, ) -> Result, DataLayerError> { if limit == 0 { return Ok(Vec::new()); } let active_statuses = vec!["submitted", "queued", "processing"]; let sql = list_due_sql(); collect_query_rows( sqlx::query(&sql) .bind(active_statuses) .bind(now_unix_secs as f64) .bind(i64::try_from(limit).map_err(|_| { DataLayerError::UnexpectedValue(format!("invalid due task limit: {limit}")) })?) .fetch(&self.pool), map_video_task_row, ) .await } pub async fn list_page( &self, filter: &VideoTaskQueryFilter, offset: usize, limit: usize, ) -> Result, DataLayerError> { if limit == 0 { return Ok(Vec::new()); } let offset = i64::try_from(offset) .map_err(|_| DataLayerError::UnexpectedValue(format!("invalid offset: {offset}")))?; let limit = i64::try_from(limit) .map_err(|_| DataLayerError::UnexpectedValue(format!("invalid limit: {limit}")))?; let mut builder = QueryBuilder::::new("SELECT\n"); builder.push(select_video_task_full_columns()); builder.push("\nFROM video_tasks"); push_video_task_filter(&mut builder, filter, None); builder.push("\nORDER BY created_at DESC, updated_at DESC"); builder.push("\nOFFSET "); builder.push_bind(offset); builder.push("\nLIMIT "); builder.push_bind(limit); let query = builder.build(); collect_query_rows(query.fetch(&self.pool), map_video_task_row).await } pub async fn list_page_summary( &self, filter: &VideoTaskQueryFilter, offset: usize, limit: usize, ) -> Result, DataLayerError> { if limit == 0 { return Ok(Vec::new()); } let offset = i64::try_from(offset) .map_err(|_| DataLayerError::UnexpectedValue(format!("invalid offset: {offset}")))?; let limit = i64::try_from(limit) .map_err(|_| DataLayerError::UnexpectedValue(format!("invalid limit: {limit}")))?; let mut builder = QueryBuilder::::new("SELECT\n"); builder.push(select_video_task_page_summary_columns()); builder.push("\nFROM video_tasks"); push_video_task_filter(&mut builder, filter, None); builder.push("\nORDER BY created_at DESC, updated_at DESC"); builder.push("\nOFFSET "); builder.push_bind(offset); builder.push("\nLIMIT "); builder.push_bind(limit); let query = builder.build(); collect_query_rows(query.fetch(&self.pool), map_video_task_row).await } pub async fn count(&self, filter: &VideoTaskQueryFilter) -> Result { let mut builder = QueryBuilder::::new("SELECT COUNT(id) AS total FROM video_tasks"); push_video_task_filter(&mut builder, filter, None); let row = builder .build() .fetch_one(&self.pool) .await .map_postgres_err()?; let total = row.try_get::("total").map_postgres_err()?; u64::try_from(total) .map_err(|_| DataLayerError::UnexpectedValue(format!("invalid count result: {total}"))) } pub async fn count_by_status( &self, filter: &VideoTaskQueryFilter, ) -> Result, DataLayerError> { let mut builder = QueryBuilder::::new("SELECT status, COUNT(id) AS total FROM video_tasks"); push_video_task_filter(&mut builder, filter, None); builder.push("\nGROUP BY status\nORDER BY status ASC"); let query = builder.build(); let mut rows = query.fetch(&self.pool); let mut items = Vec::new(); while let Some(row) = rows.try_next().await.map_postgres_err()? { let entry = { let status = VideoTaskStatus::from_database( row.try_get::("status") .map_postgres_err()? .as_str(), )?; let total = row.try_get::("total").map_postgres_err()?; VideoTaskStatusCount { status, count: u64::try_from(total).map_err(|_| { DataLayerError::UnexpectedValue(format!( "invalid status count result: {total}" )) })?, } }; items.push(entry); } Ok(items) } pub async fn count_distinct_users( &self, filter: &VideoTaskQueryFilter, ) -> Result { let mut builder = QueryBuilder::::new( "SELECT COUNT(DISTINCT user_id) AS total FROM video_tasks", ); push_video_task_filter(&mut builder, filter, None); let mut has_where = has_video_task_filter(filter, None); push_sql_clause(&mut builder, &mut has_where, "user_id IS NOT NULL"); push_sql_clause(&mut builder, &mut has_where, "user_id <> ''"); let row = builder .build() .fetch_one(&self.pool) .await .map_postgres_err()?; let total = row.try_get::("total").map_postgres_err()?; u64::try_from(total) .map_err(|_| DataLayerError::UnexpectedValue(format!("invalid count result: {total}"))) } pub async fn top_models( &self, filter: &VideoTaskQueryFilter, limit: usize, ) -> Result, DataLayerError> { if limit == 0 { return Ok(Vec::new()); } let limit = i64::try_from(limit) .map_err(|_| DataLayerError::UnexpectedValue(format!("invalid limit: {limit}")))?; let mut builder = QueryBuilder::::new("SELECT model, COUNT(id) AS total FROM video_tasks"); push_video_task_filter(&mut builder, filter, None); let mut has_where = has_video_task_filter(filter, None); push_sql_clause(&mut builder, &mut has_where, "model IS NOT NULL"); push_sql_clause(&mut builder, &mut has_where, "model <> ''"); builder.push("\nGROUP BY model\nORDER BY total DESC, model ASC\nLIMIT "); builder.push_bind(limit); let query = builder.build(); let mut rows = query.fetch(&self.pool); let mut items = Vec::new(); while let Some(row) = rows.try_next().await.map_postgres_err()? { let entry = { let model = row.try_get::("model").map_postgres_err()?; let total = row.try_get::("total").map_postgres_err()?; VideoTaskModelCount { model, count: u64::try_from(total).map_err(|_| { DataLayerError::UnexpectedValue(format!( "invalid model count result: {total}" )) })?, } }; items.push(entry); } Ok(items) } pub async fn count_created_since( &self, filter: &VideoTaskQueryFilter, created_since_unix_secs: u64, ) -> Result { let mut builder = QueryBuilder::::new("SELECT COUNT(id) AS total FROM video_tasks"); push_video_task_filter(&mut builder, filter, Some(created_since_unix_secs)); let row = builder .build() .fetch_one(&self.pool) .await .map_postgres_err()?; let total = row.try_get::("total").map_postgres_err()?; u64::try_from(total) .map_err(|_| DataLayerError::UnexpectedValue(format!("invalid count result: {total}"))) } pub async fn upsert( &self, mut task: UpsertVideoTask, ) -> Result { task.sanitize_for_persistence(); let task_id = task.id.clone(); let sql = upsert_sql(); let row = sqlx::query(&sql) .bind(task.id) .bind(task.short_id) .bind(task.request_id) .bind(task.user_id) .bind(task.api_key_id) .bind(task.username) .bind(task.api_key_name) .bind(task.external_task_id) .bind(task.provider_id) .bind(task.endpoint_id) .bind(task.key_id) .bind(task.client_api_format) .bind(task.provider_api_format) .bind(task.format_converted) .bind(task.model) .bind(task.prompt) .bind(task.original_request_body) .bind( task.duration_seconds .map(i32::try_from) .transpose() .map_err(|_| { DataLayerError::UnexpectedValue( "invalid video task duration_seconds".to_string(), ) })?, ) .bind(task.resolution) .bind(task.aspect_ratio) .bind(task.size) .bind(map_status_for_database(task.status)) .bind(i32::from(task.progress_percent)) .bind(task.progress_message) .bind(i32::try_from(task.retry_count).map_err(|_| { DataLayerError::UnexpectedValue("invalid video task retry_count".to_string()) })?) .bind(i32::try_from(task.poll_interval_seconds).map_err(|_| { DataLayerError::UnexpectedValue( "invalid video task poll_interval_seconds".to_string(), ) })?) .bind(task.next_poll_at_unix_secs.map(|value| value as f64)) .bind(i32::try_from(task.poll_count).map_err(|_| { DataLayerError::UnexpectedValue("invalid video task poll_count".to_string()) })?) .bind(i32::try_from(task.max_poll_count).map_err(|_| { DataLayerError::UnexpectedValue("invalid video task max_poll_count".to_string()) })?) .bind(task.video_url) .bind(task.error_code) .bind(task.error_message) .bind(task.request_metadata) .bind(task.created_at_unix_ms as f64) .bind(task.submitted_at_unix_secs.map(|value| value as f64)) .bind(task.completed_at_unix_secs.map(|value| value as f64)) .bind(task.updated_at_unix_secs as f64) .fetch_optional(&self.pool) .await .map_postgres_err()? .ok_or_else(|| { DataLayerError::InvalidInput(format!( "video task {task_id} conflicts with persisted immutable identity" )) })?; map_video_task_row(&row) } pub async fn update_if_active( &self, mut task: UpsertVideoTask, ) -> Result, DataLayerError> { task.sanitize_for_persistence(); let sql = update_if_active_sql(); let row = sqlx::query(&sql) .bind(task.id) .bind(task.short_id) .bind(task.request_id) .bind(task.user_id) .bind(task.api_key_id) .bind(task.username) .bind(task.api_key_name) .bind(task.external_task_id) .bind(task.provider_id) .bind(task.endpoint_id) .bind(task.key_id) .bind(task.client_api_format) .bind(task.provider_api_format) .bind(task.format_converted) .bind(task.model) .bind(task.prompt) .bind(task.original_request_body) .bind( task.duration_seconds .map(i32::try_from) .transpose() .map_err(|_| { DataLayerError::UnexpectedValue( "invalid video task duration_seconds".to_string(), ) })?, ) .bind(task.resolution) .bind(task.aspect_ratio) .bind(task.size) .bind(map_status_for_database(task.status)) .bind(i32::from(task.progress_percent)) .bind(task.progress_message) .bind(i32::try_from(task.retry_count).map_err(|_| { DataLayerError::UnexpectedValue("invalid video task retry_count".to_string()) })?) .bind(i32::try_from(task.poll_interval_seconds).map_err(|_| { DataLayerError::UnexpectedValue( "invalid video task poll_interval_seconds".to_string(), ) })?) .bind(task.next_poll_at_unix_secs.map(|value| value as f64)) .bind(i32::try_from(task.poll_count).map_err(|_| { DataLayerError::UnexpectedValue("invalid video task poll_count".to_string()) })?) .bind(i32::try_from(task.max_poll_count).map_err(|_| { DataLayerError::UnexpectedValue("invalid video task max_poll_count".to_string()) })?) .bind(task.video_url) .bind(task.error_code) .bind(task.error_message) .bind(task.request_metadata) .bind(task.created_at_unix_ms as f64) .bind(task.submitted_at_unix_secs.map(|value| value as f64)) .bind(task.completed_at_unix_secs.map(|value| value as f64)) .bind(task.updated_at_unix_secs as f64) .bind(vec!["pending", "submitted", "queued", "processing"]) .fetch_optional(&self.pool) .await .map_postgres_err()?; row.as_ref().map(map_video_task_row).transpose() } pub async fn claim_due( &self, now_unix_secs: u64, claim_until_unix_secs: u64, limit: usize, ) -> Result, DataLayerError> { if limit == 0 { return Ok(Vec::new()); } let limit = i64::try_from(limit) .map_err(|_| DataLayerError::UnexpectedValue(format!("invalid limit: {limit}")))?; let sql = claim_due_sql(); let mut tasks = collect_query_rows( sqlx::query(&sql) .bind(vec!["submitted", "queued", "processing"]) .bind(now_unix_secs as f64) .bind(limit) .bind(claim_until_unix_secs as f64) .bind(now_unix_secs as f64) .fetch(&self.pool), map_video_task_row, ) .await?; 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] impl VideoTaskReadRepository for SqlxVideoTaskRepository { async fn find( &self, key: VideoTaskLookupKey<'_>, ) -> Result, DataLayerError> { Self::find(self, key).await } async fn find_for_user( &self, key: VideoTaskLookupKey<'_>, user_id: &str, ) -> Result, DataLayerError> { Self::find_for_user(self, key, user_id).await } async fn list_active(&self, limit: usize) -> Result, DataLayerError> { Self::list_active(self, limit).await } async fn list_due( &self, now_unix_secs: u64, limit: usize, ) -> Result, DataLayerError> { Self::list_due(self, now_unix_secs, limit).await } async fn list_page( &self, filter: &VideoTaskQueryFilter, offset: usize, limit: usize, ) -> Result, DataLayerError> { Self::list_page(self, filter, offset, limit).await } async fn list_page_summary( &self, filter: &VideoTaskQueryFilter, offset: usize, limit: usize, ) -> Result, DataLayerError> { Self::list_page_summary(self, filter, offset, limit).await } async fn count(&self, filter: &VideoTaskQueryFilter) -> Result { Self::count(self, filter).await } async fn count_by_status( &self, filter: &VideoTaskQueryFilter, ) -> Result, DataLayerError> { Self::count_by_status(self, filter).await } async fn count_distinct_users( &self, filter: &VideoTaskQueryFilter, ) -> Result { Self::count_distinct_users(self, filter).await } async fn top_models( &self, filter: &VideoTaskQueryFilter, limit: usize, ) -> Result, DataLayerError> { Self::top_models(self, filter, limit).await } async fn count_created_since( &self, filter: &VideoTaskQueryFilter, created_since_unix_secs: u64, ) -> Result { Self::count_created_since(self, filter, created_since_unix_secs).await } } #[async_trait] impl VideoTaskWriteRepository for SqlxVideoTaskRepository { async fn upsert(&self, task: UpsertVideoTask) -> Result { Self::upsert(self, task).await } async fn update_if_active( &self, task: UpsertVideoTask, ) -> Result, DataLayerError> { Self::update_if_active(self, task).await } async fn claim_due( &self, now_unix_secs: u64, claim_until_unix_secs: u64, limit: usize, ) -> Result, DataLayerError> { Self::claim_due(self, now_unix_secs, claim_until_unix_secs, limit).await } } fn has_video_task_filter( filter: &VideoTaskQueryFilter, created_since_unix_secs: Option, ) -> bool { filter.user_id.is_some() || filter.status.is_some() || filter.model_substring.is_some() || filter.client_api_format.is_some() || created_since_unix_secs.is_some() } fn push_sql_clause<'args>( builder: &mut QueryBuilder<'args, Postgres>, has_where: &mut bool, clause: &str, ) { if *has_where { builder.push("\n AND "); } else { builder.push("\nWHERE "); *has_where = true; } builder.push(clause); } fn push_video_task_filter<'args>( builder: &mut QueryBuilder<'args, Postgres>, filter: &'args VideoTaskQueryFilter, created_since_unix_secs: Option, ) { let mut has_where = false; if let Some(user_id) = filter.user_id.as_deref() { push_sql_clause(builder, &mut has_where, "user_id = "); builder.push_bind(user_id); } if let Some(status) = filter.status { push_sql_clause(builder, &mut has_where, "status = "); builder.push_bind(map_status_for_database(status)); } if let Some(model_substring) = filter.model_substring.as_deref() { push_sql_clause(builder, &mut has_where, "model ILIKE "); builder.push_bind(format!("%{}%", escape_like_pattern(model_substring.trim()))); builder.push(" ESCAPE '\\'"); } if let Some(client_api_format) = filter.client_api_format.as_deref() { push_sql_clause(builder, &mut has_where, "client_api_format = "); builder.push_bind(client_api_format); } if let Some(created_since_unix_secs) = created_since_unix_secs { push_sql_clause(builder, &mut has_where, "created_at >= TO_TIMESTAMP("); builder.push_bind(created_since_unix_secs as f64); builder.push(")"); } } fn escape_like_pattern(value: &str) -> String { value .replace('\\', "\\\\") .replace('%', "\\%") .replace('_', "\\_") } fn map_status_for_database(status: VideoTaskStatus) -> &'static str { match status { VideoTaskStatus::Pending => "pending", VideoTaskStatus::Submitted => "submitted", VideoTaskStatus::Queued => "queued", VideoTaskStatus::Processing => "processing", VideoTaskStatus::Completed => "completed", VideoTaskStatus::Failed => "failed", VideoTaskStatus::Cancelled => "cancelled", VideoTaskStatus::Expired => "expired", VideoTaskStatus::Deleted => "deleted", } } async fn collect_query_rows( mut rows: S, map_row: fn(&PgRow) -> Result, ) -> Result, DataLayerError> where S: TryStream + Unpin, { let mut items = Vec::new(); while let Some(row) = rows.try_next().await.map_postgres_err()? { items.push(map_row(&row)?); } Ok(items) } fn map_video_task_row(row: &PgRow) -> Result { let status = VideoTaskStatus::from_database( row.try_get::("status") .map_postgres_err()? .as_str(), )?; StoredVideoTask::new( row.try_get("id").map_postgres_err()?, row.try_get("short_id").map_postgres_err()?, row.try_get("request_id").map_postgres_err()?, row.try_get("user_id").map_postgres_err()?, row.try_get("api_key_id").map_postgres_err()?, row.try_get("username").map_postgres_err()?, row.try_get("api_key_name").map_postgres_err()?, row.try_get("external_task_id").map_postgres_err()?, row.try_get("provider_id").map_postgres_err()?, row.try_get("endpoint_id").map_postgres_err()?, row.try_get("key_id").map_postgres_err()?, row.try_get("client_api_format").map_postgres_err()?, row.try_get("provider_api_format").map_postgres_err()?, row.try_get("format_converted").map_postgres_err()?, row.try_get("model").map_postgres_err()?, row.try_get("prompt").map_postgres_err()?, row.try_get("original_request_body").map_postgres_err()?, row.try_get("duration_seconds").map_postgres_err()?, row.try_get("resolution").map_postgres_err()?, row.try_get("aspect_ratio").map_postgres_err()?, row.try_get("size").map_postgres_err()?, status, row.try_get("progress_percent").map_postgres_err()?, row.try_get("progress_message").map_postgres_err()?, row.try_get("retry_count").map_postgres_err()?, row.try_get("poll_interval_seconds").map_postgres_err()?, row.try_get("next_poll_at_unix_secs").map_postgres_err()?, row.try_get("poll_count").map_postgres_err()?, row.try_get("max_poll_count").map_postgres_err()?, row.try_get("created_at_unix_ms").map_postgres_err()?, row.try_get("submitted_at_unix_secs").map_postgres_err()?, row.try_get("completed_at_unix_secs").map_postgres_err()?, row.try_get("updated_at_unix_secs").map_postgres_err()?, row.try_get("error_code").map_postgres_err()?, row.try_get("error_message").map_postgres_err()?, row.try_get("video_url").map_postgres_err()?, row.try_get("request_metadata").map_postgres_err()?, ) } #[cfg(test)] mod tests { use super::{claim_due_sql, update_if_active_sql, upsert_sql, SqlxVideoTaskRepository}; use crate::{PostgresPoolConfig, PostgresPoolFactory}; use aether_data_contracts::repository::video_tasks::{ UpsertVideoTask, VideoTaskLookupKey, VideoTaskQueryFilter, VideoTaskReadRepository, VideoTaskStatus, VideoTaskWriteRepository, }; fn build_pool() -> sqlx::PgPool { let factory = PostgresPoolFactory::new(PostgresPoolConfig { database_url: "postgres://localhost/aether".to_string(), min_connections: 1, max_connections: 4, acquire_timeout_ms: 1_000, idle_timeout_ms: 5_000, max_lifetime_ms: 30_000, statement_cache_capacity: 64, require_ssl: false, }) .expect("factory should build"); factory.connect_lazy().expect("pool should build") } #[test] fn write_sql_guards_every_immutable_identity_field() { let upsert = upsert_sql(); let update = update_if_active_sql(); let immutable_columns = [ ("short_id", "$2"), ("request_id", "$3"), ("user_id", "$4"), ("api_key_id", "$5"), ("external_task_id", "$8"), ("provider_id", "$9"), ("endpoint_id", "$10"), ("key_id", "$11"), ("client_api_format", "$12"), ("provider_api_format", "$13"), ("format_converted", "$14"), ("model", "$15"), ("duration_seconds", "$18"), ("resolution", "$19"), ("aspect_ratio", "$20"), ("size", "$21"), ]; for (column, parameter) in immutable_columns { assert!( upsert.contains(&format!( "video_tasks.{column} IS NOT DISTINCT FROM EXCLUDED.{column}" )), "upsert should guard {column}" ); assert!( update.contains(&format!("{column} IS NOT DISTINCT FROM {parameter}")), "active update should guard {column}" ); } assert!( upsert.contains("created_at = COALESCE(video_tasks.created_at, EXCLUDED.created_at)") ); assert!(update.contains("created_at = COALESCE(created_at, TO_TIMESTAMP($34))")); } #[test] fn poll_claim_returns_business_fields_required_by_identity_guards() { let sql = claim_due_sql(); for field in [ "prompt", "duration_seconds", "resolution", "aspect_ratio", "size", ] { assert!( sql.contains(&format!("{field} AS {field}")), "claim must retain {field}" ); } assert!(sql.contains("NULL::jsonb AS original_request_body")); assert!(sql.contains("FOR UPDATE SKIP LOCKED")); } #[tokio::test] #[ignore = "requires AETHER_TEST_DATABASE_URL and PostgreSQL migrations"] async fn live_video_task_capture_claim_and_completion_preserve_business_fields() { let database_url = std::env::var("AETHER_TEST_DATABASE_URL") .expect("AETHER_TEST_DATABASE_URL must point at the test database"); let pool = sqlx::postgres::PgPoolOptions::new() .max_connections(1) .connect(&database_url) .await .expect("test database should connect"); crate::run_migrations(&pool) .await .expect("test database should migrate"); sqlx::query("CREATE TEMP TABLE video_tasks (LIKE public.video_tasks INCLUDING ALL)") .execute(&pool) .await .expect("isolated task table should be created"); let repository = SqlxVideoTaskRepository::new(pool); for api_format in ["openai:video", "gemini:video"] { let task_id = uuid::Uuid::new_v4().to_string(); let original = UpsertVideoTask { id: task_id.clone(), short_id: Some(uuid::Uuid::new_v4().simple().to_string()[..16].to_string()), request_id: format!("request-{task_id}"), user_id: None, api_key_id: None, username: Some("alice".to_string()), api_key_name: Some("video-client".to_string()), external_task_id: Some("upstream-task-1".to_string()), provider_id: None, endpoint_id: None, key_id: None, client_api_format: Some(api_format.to_string()), provider_api_format: Some(api_format.to_string()), format_converted: false, model: Some("video-model".to_string()), prompt: Some("business prompt".to_string()), original_request_body: Some(serde_json::json!({"token": "private"})), duration_seconds: Some(8), resolution: Some("1080p".to_string()), aspect_ratio: Some("16:9".to_string()), size: Some("1920x1080".to_string()), status: VideoTaskStatus::Submitted, progress_percent: 0, progress_message: None, retry_count: 0, poll_interval_seconds: 10, next_poll_at_unix_secs: Some(10), poll_count: 0, max_poll_count: 360, created_at_unix_ms: 1, submitted_at_unix_secs: Some(1), completed_at_unix_secs: None, updated_at_unix_secs: 1, error_code: None, error_message: None, video_url: None, request_metadata: Some(serde_json::json!({"authorization": "private"})), }; let stored = repository .upsert(original.clone()) .await .expect("task should persist"); assert_eq!(stored.prompt, original.prompt); assert_eq!(stored.username, original.username); assert_eq!(stored.api_key_name, original.api_key_name); assert!(stored.original_request_body.is_none()); assert!(stored.request_metadata.is_none()); let mut claimed = repository .claim_due(20, 50, 10) .await .expect("task should be claimed"); assert_eq!(claimed.len(), 1); let mut completion: UpsertVideoTask = claimed.pop().expect("claimed task").into(); stored .ensure_immutable_identity_matches(&completion) .expect("claim must preserve task identity"); assert_eq!(completion.prompt, original.prompt); let mut mismatched = completion.clone(); mismatched.duration_seconds = Some(99); assert!(repository .update_if_active(mismatched) .await .expect("guarded update should execute") .is_none()); completion.status = VideoTaskStatus::Completed; completion.progress_percent = 100; completion.next_poll_at_unix_secs = None; completion.completed_at_unix_secs = Some(21); completion.updated_at_unix_secs = 21; completion.video_url = Some( "https://cdn.example.test/video.mp4?alt=media&signature=a%2Fb%2Bc%3D&part=2&part=1" .to_string(), ); let completed = repository .update_if_active(completion.clone()) .await .expect("completion should execute") .expect("matching active task should complete"); assert_eq!(completed.video_url, completion.video_url); let reloaded = repository .find(VideoTaskLookupKey::Id(&task_id)) .await .expect("task should reload") .expect("task should exist"); assert_eq!(reloaded.status, VideoTaskStatus::Completed); assert_eq!(reloaded.prompt, original.prompt); assert_eq!(reloaded.video_url, completion.video_url); assert_eq!(reloaded.duration_seconds, original.duration_seconds); assert_eq!(reloaded.size, original.size); assert_eq!(reloaded.username, original.username); assert!(reloaded.request_metadata.is_none()); } repository.pool().close().await; } #[tokio::test] async fn repository_constructs_from_lazy_pool() { let repository = SqlxVideoTaskRepository::new(build_pool()); let _ = repository.pool(); } #[tokio::test] async fn read_trait_delegates_to_sqlx_repository() { let repository = SqlxVideoTaskRepository::new(build_pool()); let _ = VideoTaskReadRepository::find(&repository, VideoTaskLookupKey::Id("task-1")).await; } #[tokio::test] async fn read_trait_delegates_due_listing_to_sqlx_repository() { let repository = SqlxVideoTaskRepository::new(build_pool()); let _ = VideoTaskReadRepository::list_due(&repository, 1, 10).await; } #[tokio::test] async fn read_trait_delegates_query_methods_to_sqlx_repository() { let repository = SqlxVideoTaskRepository::new(build_pool()); let filter = VideoTaskQueryFilter::default(); let _ = VideoTaskReadRepository::list_page(&repository, &filter, 0, 10).await; let _ = VideoTaskReadRepository::count(&repository, &filter).await; let _ = VideoTaskReadRepository::count_by_status(&repository, &filter).await; let _ = VideoTaskReadRepository::top_models(&repository, &filter, 10).await; let _ = VideoTaskReadRepository::count_created_since(&repository, &filter, 0).await; } #[tokio::test] async fn write_trait_delegates_to_sqlx_repository() { let repository = SqlxVideoTaskRepository::new(build_pool()); let _ = VideoTaskWriteRepository::upsert( &repository, UpsertVideoTask { id: "task-1".to_string(), short_id: Some("short-task-1".to_string()), request_id: "request-1".to_string(), user_id: None, api_key_id: None, username: None, api_key_name: None, external_task_id: Some("ext-1".to_string()), provider_id: Some("provider-1".to_string()), endpoint_id: Some("endpoint-1".to_string()), key_id: Some("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: VideoTaskStatus::Submitted, progress_percent: 0, progress_message: None, retry_count: 0, poll_interval_seconds: 10, next_poll_at_unix_secs: Some(10), poll_count: 0, max_poll_count: 360, created_at_unix_ms: 1, submitted_at_unix_secs: Some(1), completed_at_unix_secs: None, updated_at_unix_secs: 1, error_code: None, error_message: None, video_url: None, request_metadata: None, }, ) .await; } #[tokio::test] async fn write_trait_delegates_active_update_to_sqlx_repository() { let repository = SqlxVideoTaskRepository::new(build_pool()); let _ = VideoTaskWriteRepository::update_if_active( &repository, UpsertVideoTask { id: "task-1".to_string(), short_id: Some("short-task-1".to_string()), request_id: "request-1".to_string(), user_id: None, api_key_id: None, username: None, api_key_name: None, external_task_id: Some("ext-1".to_string()), provider_id: Some("provider-1".to_string()), endpoint_id: Some("endpoint-1".to_string()), key_id: Some("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: VideoTaskStatus::Processing, progress_percent: 50, progress_message: Some("polling".to_string()), retry_count: 1, poll_interval_seconds: 10, next_poll_at_unix_secs: Some(20), poll_count: 1, max_poll_count: 360, created_at_unix_ms: 1, submitted_at_unix_secs: Some(1), completed_at_unix_secs: None, updated_at_unix_secs: 20, error_code: None, error_message: None, video_url: None, request_metadata: None, }, ) .await; } #[tokio::test] async fn write_trait_delegates_due_claim_to_sqlx_repository() { let repository = SqlxVideoTaskRepository::new(build_pool()); let _ = VideoTaskWriteRepository::claim_due(&repository, 1, 30, 10).await; } }