1use crate::error::{map_sqlx_error, AwaError};
2use crate::job::{JobRow, JobState};
3use crate::queue_storage::QueueStorage;
4use chrono::{DateTime, Utc};
5use serde::{Deserialize, Serialize};
6use sqlx::postgres::PgRow;
7use sqlx::types::Json;
8use sqlx::{FromRow, Row};
9use sqlx::{PgPool, Postgres, Transaction};
10use uuid::Uuid;
11
12const DEFAULT_CHUNK_SIZE: i64 = 100;
13const PREVIEW_SAMPLE_LIMIT: i64 = 10;
14const SCAN_PAGE_SIZE: i64 = 10_000;
15const BATCH_RUNNER_LOCK_KEY: i64 = 0x4157_415f_4241_5443;
16
17fn default_retention_days() -> i64 {
18 std::env::var("AWA_BATCH_OP_RETENTION_DAYS")
19 .ok()
20 .and_then(|value| value.parse::<i64>().ok())
21 .filter(|days| *days > 0)
22 .unwrap_or(90)
23}
24
25#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
26#[serde(rename_all = "snake_case")]
27pub enum BatchOperationKind {
28 SetPriority,
29 MoveQueue,
30}
31
32impl BatchOperationKind {
33 fn as_str(self) -> &'static str {
34 match self {
35 Self::SetPriority => "set_priority",
36 Self::MoveQueue => "move_queue",
37 }
38 }
39}
40
41impl TryFrom<&str> for BatchOperationKind {
42 type Error = AwaError;
43
44 fn try_from(value: &str) -> Result<Self, Self::Error> {
45 match value {
46 "set_priority" => Ok(Self::SetPriority),
47 "move_queue" => Ok(Self::MoveQueue),
48 _ => Err(AwaError::Validation(format!(
49 "unknown batch operation kind: {value}"
50 ))),
51 }
52 }
53}
54
55#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
56#[serde(rename_all = "snake_case")]
57pub enum BatchOperationState {
58 Pending,
59 Scanning,
60 Running,
61 Cancelling,
62 Completed,
63 Cancelled,
64 Failed,
65}
66
67impl TryFrom<&str> for BatchOperationState {
68 type Error = AwaError;
69
70 fn try_from(value: &str) -> Result<Self, Self::Error> {
71 match value {
72 "pending" => Ok(Self::Pending),
73 "scanning" => Ok(Self::Scanning),
74 "running" => Ok(Self::Running),
75 "cancelling" => Ok(Self::Cancelling),
76 "completed" => Ok(Self::Completed),
77 "cancelled" => Ok(Self::Cancelled),
78 "failed" => Ok(Self::Failed),
79 _ => Err(AwaError::Validation(format!(
80 "unknown batch operation state: {value}"
81 ))),
82 }
83 }
84}
85
86impl BatchOperationState {
87 fn as_str(self) -> &'static str {
88 match self {
89 Self::Pending => "pending",
90 Self::Scanning => "scanning",
91 Self::Running => "running",
92 Self::Cancelling => "cancelling",
93 Self::Completed => "completed",
94 Self::Cancelled => "cancelled",
95 Self::Failed => "failed",
96 }
97 }
98}
99
100#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
101pub struct BatchOperationFilter {
102 pub kind: Option<String>,
103 pub queue: Option<String>,
104 pub ids: Option<Vec<i64>>,
105 pub tag: Option<String>,
106 pub state: Option<JobState>,
107 pub created_at_gte: Option<DateTime<Utc>>,
108 pub created_at_lt: Option<DateTime<Utc>>,
109}
110
111impl BatchOperationFilter {
112 pub fn is_empty(&self) -> bool {
113 self.kind.is_none()
114 && self.queue.is_none()
115 && self.ids.as_ref().is_none_or(Vec::is_empty)
116 && self.tag.is_none()
117 && self.state.is_none()
118 && self.created_at_gte.is_none()
119 && self.created_at_lt.is_none()
120 }
121}
122
123#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
124#[serde(rename_all = "snake_case", tag = "op_kind")]
125pub enum BatchOperationSpec {
126 SetPriority {
127 priority: i16,
128 },
129 MoveQueue {
130 queue: String,
131 priority: Option<i16>,
132 },
133}
134
135impl BatchOperationSpec {
136 pub fn kind(&self) -> BatchOperationKind {
137 match self {
138 Self::SetPriority { .. } => BatchOperationKind::SetPriority,
139 Self::MoveQueue { .. } => BatchOperationKind::MoveQueue,
140 }
141 }
142
143 fn validate(&self) -> Result<(), AwaError> {
144 match self {
145 Self::SetPriority { priority } => validate_priority(*priority),
146 Self::MoveQueue { queue, priority } => {
147 if queue.is_empty() || queue.len() > 200 {
148 return Err(AwaError::Validation(
149 "destination queue must be 1..=200 characters".to_string(),
150 ));
151 }
152 if let Some(priority) = priority {
153 validate_priority(*priority)?;
154 }
155 Ok(())
156 }
157 }
158 }
159}
160
161fn validate_priority(priority: i16) -> Result<(), AwaError> {
162 if (1..=4).contains(&priority) {
163 Ok(())
164 } else {
165 Err(AwaError::Validation(
166 "priority must be between 1 and 4".to_string(),
167 ))
168 }
169}
170
171#[derive(Debug, Clone, Serialize, Deserialize)]
172pub struct SubmitBatchOperation {
173 pub spec: BatchOperationSpec,
174 pub filter: BatchOperationFilter,
175 pub submitted_by: Option<String>,
176 pub allow_all: bool,
177}
178
179#[derive(Debug, Clone, Serialize)]
180pub struct BatchOperation {
181 pub id: Uuid,
182 pub op_kind: BatchOperationKind,
183 pub filter: Json<BatchOperationFilter>,
184 pub spec: Json<BatchOperationSpec>,
185 pub state: BatchOperationState,
186 pub submitted_by: Option<String>,
187 pub submitted_at: DateTime<Utc>,
188 pub started_at: Option<DateTime<Utc>>,
189 pub finalized_at: Option<DateTime<Utc>>,
190 pub cursor: Option<Json<BatchOperationCursor>>,
191 pub total_matched: Option<i64>,
192 pub processed: i64,
193 pub skipped: i64,
194 pub errored: i64,
195 pub last_error: Option<String>,
196 pub runner_instance: Option<Uuid>,
197 pub retention_until: Option<DateTime<Utc>>,
198 pub updated_at: DateTime<Utc>,
199}
200
201impl<'r> FromRow<'r, PgRow> for BatchOperation {
202 fn from_row(row: &'r PgRow) -> Result<Self, sqlx::Error> {
203 let op_kind: String = row.try_get("op_kind")?;
204 let state: String = row.try_get("state")?;
205 Ok(Self {
206 id: row.try_get("id")?,
207 op_kind: BatchOperationKind::try_from(op_kind.as_str()).map_err(|err| {
208 sqlx::Error::ColumnDecode {
209 index: "op_kind".to_string(),
210 source: Box::new(err),
211 }
212 })?,
213 filter: row.try_get("filter")?,
214 spec: row.try_get("spec")?,
215 state: BatchOperationState::try_from(state.as_str()).map_err(|err| {
216 sqlx::Error::ColumnDecode {
217 index: "state".to_string(),
218 source: Box::new(err),
219 }
220 })?,
221 submitted_by: row.try_get("submitted_by")?,
222 submitted_at: row.try_get("submitted_at")?,
223 started_at: row.try_get("started_at")?,
224 finalized_at: row.try_get("finalized_at")?,
225 cursor: row.try_get("cursor")?,
226 total_matched: row.try_get("total_matched")?,
227 processed: row.try_get("processed")?,
228 skipped: row.try_get("skipped")?,
229 errored: row.try_get("errored")?,
230 last_error: row.try_get("last_error")?,
231 runner_instance: row.try_get("runner_instance")?,
232 retention_until: row.try_get("retention_until")?,
233 updated_at: row.try_get("updated_at")?,
234 })
235 }
236}
237
238#[derive(Debug, Clone, Serialize, Deserialize)]
239pub struct BatchOperationCursor {
240 pub after_job_id: i64,
241 pub max_job_id: i64,
242}
243
244#[derive(Debug, Clone, Serialize)]
245pub struct BatchOperationPreview {
246 pub total_matched: i64,
247 pub sample: Vec<JobRow>,
248}
249
250#[derive(Debug, Clone, Copy, Default, Serialize)]
251pub struct BatchOperationRunOutcome {
252 pub claimed: bool,
253 pub finalized: bool,
254 pub processed: i64,
255 pub skipped: i64,
256 pub errored: i64,
257}
258
259#[derive(Debug, Clone, Default)]
260pub struct ListBatchOperationsFilter {
261 pub state: Option<BatchOperationState>,
262 pub limit: Option<i64>,
263}
264
265pub async fn submit_batch_operation(
266 pool: &PgPool,
267 request: SubmitBatchOperation,
268) -> Result<BatchOperation, AwaError> {
269 request.spec.validate()?;
270 validate_filter_for_spec(&request.spec, &request.filter)?;
271 if request.filter.is_empty() && !request.allow_all {
272 return Err(AwaError::Validation(
273 "batch operation requires a filter or allow_all=true".to_string(),
274 ));
275 }
276 let op_kind = request.spec.kind();
277 let id = Uuid::new_v4();
278 sqlx::query_as::<_, BatchOperation>(
279 r#"
280 INSERT INTO awa.batch_operations (id, op_kind, filter, spec, state, submitted_by)
281 VALUES ($1, $2, $3, $4, 'pending', $5)
282 RETURNING *
283 "#,
284 )
285 .bind(id)
286 .bind(op_kind.as_str())
287 .bind(Json(request.filter))
288 .bind(Json(request.spec))
289 .bind(request.submitted_by)
290 .fetch_one(pool)
291 .await
292 .map_err(map_sqlx_error)
293}
294
295pub async fn list_batch_operations(
296 pool: &PgPool,
297 filter: &ListBatchOperationsFilter,
298) -> Result<Vec<BatchOperation>, AwaError> {
299 let limit = filter.limit.unwrap_or(100).clamp(1, 1000);
300 sqlx::query_as::<_, BatchOperation>(
301 r#"
302 SELECT *
303 FROM awa.batch_operations
304 WHERE ($1::text IS NULL OR state = $1)
305 ORDER BY submitted_at DESC
306 LIMIT $2
307 "#,
308 )
309 .bind(filter.state.map(BatchOperationState::as_str))
310 .bind(limit)
311 .fetch_all(pool)
312 .await
313 .map_err(map_sqlx_error)
314}
315
316pub async fn get_batch_operation(pool: &PgPool, id: Uuid) -> Result<BatchOperation, AwaError> {
317 sqlx::query_as::<_, BatchOperation>("SELECT * FROM awa.batch_operations WHERE id = $1")
318 .bind(id)
319 .fetch_optional(pool)
320 .await
321 .map_err(map_sqlx_error)?
322 .ok_or_else(|| AwaError::Validation(format!("batch operation not found: {id}")))
323}
324
325pub async fn request_batch_operation_cancellation(
326 pool: &PgPool,
327 id: Uuid,
328) -> Result<BatchOperation, AwaError> {
329 sqlx::query_as::<_, BatchOperation>(
330 r#"
331 UPDATE awa.batch_operations
332 SET state = 'cancelling', updated_at = clock_timestamp()
333 WHERE id = $1 AND state IN ('pending', 'scanning', 'running')
334 RETURNING *
335 "#,
336 )
337 .bind(id)
338 .fetch_optional(pool)
339 .await
340 .map_err(map_sqlx_error)?
341 .ok_or_else(|| AwaError::Validation(format!("batch operation cannot be cancelled: {id}")))
342}
343
344pub async fn preview_batch_operation(
345 pool: &PgPool,
346 spec: BatchOperationSpec,
347 filter: BatchOperationFilter,
348) -> Result<BatchOperationPreview, AwaError> {
349 spec.validate()?;
350 validate_filter_for_spec(&spec, &filter)?;
351 let effective_filter = filter.clone();
352 let total_matched = count_matching_jobs(pool, &effective_filter, None, None).await?;
353 let sample =
354 load_matching_jobs(pool, &effective_filter, 0, i64::MAX, PREVIEW_SAMPLE_LIMIT).await?;
355 Ok(BatchOperationPreview {
356 total_matched,
357 sample,
358 })
359}
360
361pub async fn run_one_batch_operation_chunk(
362 pool: &PgPool,
363 runner_instance: Uuid,
364 chunk_size: i64,
365) -> Result<BatchOperationRunOutcome, AwaError> {
366 let mut lock_conn = pool.acquire().await.map_err(map_sqlx_error)?;
367 let locked: bool = sqlx::query_scalar("SELECT pg_try_advisory_lock($1)")
368 .bind(BATCH_RUNNER_LOCK_KEY)
369 .fetch_one(&mut *lock_conn)
370 .await
371 .map_err(map_sqlx_error)?;
372 if !locked {
373 return Ok(BatchOperationRunOutcome::default());
374 }
375
376 let result = run_one_batch_operation_chunk_locked(pool, runner_instance, chunk_size).await;
377 let _ = sqlx::query("SELECT pg_advisory_unlock($1)")
378 .bind(BATCH_RUNNER_LOCK_KEY)
379 .execute(&mut *lock_conn)
380 .await;
381 result
382}
383
384async fn run_one_batch_operation_chunk_locked(
385 pool: &PgPool,
386 runner_instance: Uuid,
387 chunk_size: i64,
388) -> Result<BatchOperationRunOutcome, AwaError> {
389 let chunk_size = chunk_size.max(1);
390 let Some(mut operation) = claim_next_operation(pool, runner_instance).await? else {
391 return Ok(BatchOperationRunOutcome::default());
392 };
393
394 if operation.state == BatchOperationState::Cancelling {
395 skip_pending_items(pool, operation.id).await?;
396 finalize_operation(pool, operation.id, BatchOperationState::Cancelled, None).await?;
397 return Ok(BatchOperationRunOutcome {
398 claimed: true,
399 finalized: true,
400 ..Default::default()
401 });
402 }
403
404 if operation.state == BatchOperationState::Scanning {
405 operation = scan_operation(pool, operation).await?;
406 }
407
408 if operation.state != BatchOperationState::Running {
409 return Ok(BatchOperationRunOutcome {
410 claimed: true,
411 finalized: matches!(
412 operation.state,
413 BatchOperationState::Completed
414 | BatchOperationState::Cancelled
415 | BatchOperationState::Failed
416 ),
417 ..Default::default()
418 });
419 }
420
421 let job_ids = claim_pending_items(pool, operation.id, chunk_size).await?;
422
423 if job_ids.is_empty() {
424 finalize_operation(pool, operation.id, BatchOperationState::Completed, None).await?;
425 return Ok(BatchOperationRunOutcome {
426 claimed: true,
427 finalized: true,
428 ..Default::default()
429 });
430 }
431
432 let mut processed = 0;
433 let mut skipped = 0;
434 let mut errored = 0;
435
436 for job_id in job_ids {
437 match apply_item(pool, operation.id, &operation.spec.0, job_id).await {
438 Ok(ItemOutcome::Processed) => processed += 1,
439 Ok(ItemOutcome::Skipped) => skipped += 1,
440 Ok(ItemOutcome::Errored(_)) => {
441 errored += 1;
442 }
443 Err(_) => {
444 errored += 1;
445 }
446 }
447 }
448
449 let pending_count: i64 = sqlx::query_scalar(
450 r#"
451 SELECT count(*)::bigint
452 FROM awa.batch_operation_items
453 WHERE operation_id = $1 AND state = 'pending'
454 "#,
455 )
456 .bind(operation.id)
457 .fetch_one(pool)
458 .await
459 .map_err(map_sqlx_error)?;
460
461 let state_after_chunk: String =
462 sqlx::query_scalar("SELECT state FROM awa.batch_operations WHERE id = $1")
463 .bind(operation.id)
464 .fetch_one(pool)
465 .await
466 .map_err(map_sqlx_error)?;
467
468 let finalized = if pending_count == 0 {
469 let final_state = if state_after_chunk == BatchOperationState::Cancelling.as_str() {
470 skip_pending_items(pool, operation.id).await?;
471 BatchOperationState::Cancelled
472 } else {
473 BatchOperationState::Completed
474 };
475 finalize_operation(pool, operation.id, final_state, None).await?;
476 true
477 } else {
478 false
479 };
480
481 Ok(BatchOperationRunOutcome {
482 claimed: true,
483 finalized,
484 processed,
485 skipped,
486 errored,
487 })
488}
489
490pub async fn cleanup_expired_batch_operations(pool: &PgPool, limit: i64) -> Result<u64, AwaError> {
491 let result = sqlx::query(
492 r#"
493 DELETE FROM awa.batch_operations
494 WHERE id IN (
495 SELECT id
496 FROM awa.batch_operations
497 WHERE retention_until IS NOT NULL
498 AND retention_until < clock_timestamp()
499 AND state IN ('completed', 'cancelled', 'failed')
500 ORDER BY retention_until ASC
501 LIMIT $1
502 )
503 "#,
504 )
505 .bind(limit.max(1))
506 .execute(pool)
507 .await
508 .map_err(map_sqlx_error)?;
509 Ok(result.rows_affected())
510}
511
512pub async fn purge_batch_operations_before(
513 pool: &PgPool,
514 before: DateTime<Utc>,
515 limit: i64,
516) -> Result<u64, AwaError> {
517 let result = sqlx::query(
518 r#"
519 DELETE FROM awa.batch_operations
520 WHERE id IN (
521 SELECT id
522 FROM awa.batch_operations
523 WHERE finalized_at IS NOT NULL
524 AND finalized_at < $1
525 AND state IN ('completed', 'cancelled', 'failed')
526 ORDER BY finalized_at ASC
527 LIMIT $2
528 )
529 "#,
530 )
531 .bind(before)
532 .bind(limit.max(1))
533 .execute(pool)
534 .await
535 .map_err(map_sqlx_error)?;
536 Ok(result.rows_affected())
537}
538
539pub async fn finalize_expired_batch_operation_retention(
540 pool: &PgPool,
541 id: Uuid,
542 state: BatchOperationState,
543 retention_days: i64,
544 last_error: Option<String>,
545) -> Result<(), AwaError> {
546 sqlx::query(
547 r#"
548 UPDATE awa.batch_operations
549 SET state = $2,
550 finalized_at = COALESCE(finalized_at, clock_timestamp()),
551 retention_until = COALESCE(retention_until, clock_timestamp() + ($3 * interval '1 day')),
552 last_error = COALESCE($4, last_error),
553 updated_at = clock_timestamp()
554 WHERE id = $1
555 "#,
556 )
557 .bind(id)
558 .bind(state.as_str())
559 .bind(retention_days.max(1))
560 .bind(last_error)
561 .execute(pool)
562 .await
563 .map_err(map_sqlx_error)?;
564 Ok(())
565}
566
567async fn claim_next_operation(
568 pool: &PgPool,
569 runner_instance: Uuid,
570) -> Result<Option<BatchOperation>, AwaError> {
571 sqlx::query_as::<_, BatchOperation>(
572 r#"
573 WITH target AS (
574 SELECT id, state
575 FROM awa.batch_operations
576 WHERE state IN ('pending', 'scanning', 'running', 'cancelling')
577 ORDER BY submitted_at ASC
578 LIMIT 1
579 FOR UPDATE SKIP LOCKED
580 )
581 UPDATE awa.batch_operations AS op
582 SET state = CASE WHEN target.state = 'pending' THEN 'scanning' ELSE op.state END,
583 started_at = COALESCE(op.started_at, clock_timestamp()),
584 runner_instance = $1,
585 updated_at = clock_timestamp()
586 FROM target
587 WHERE op.id = target.id
588 RETURNING op.*
589 "#,
590 )
591 .bind(runner_instance)
592 .fetch_optional(pool)
593 .await
594 .map_err(map_sqlx_error)
595}
596
597async fn scan_operation(
598 pool: &PgPool,
599 operation: BatchOperation,
600) -> Result<BatchOperation, AwaError> {
601 validate_filter_for_spec(&operation.spec.0, &operation.filter.0)?;
602 let effective_filter = operation.filter.0.clone();
603 let max_job_id = max_matching_job_id(pool, &effective_filter)
604 .await?
605 .unwrap_or(0);
606 let total_matched = if max_job_id == 0 {
607 0
608 } else {
609 count_matching_jobs(pool, &effective_filter, Some(0), Some(max_job_id)).await?
610 };
611 let operation_id = operation.id;
612 let mut after_job_id = 0;
613 while after_job_id < max_job_id {
614 let ids = load_matching_job_ids(
615 pool,
616 &effective_filter,
617 after_job_id,
618 max_job_id,
619 SCAN_PAGE_SIZE,
620 )
621 .await?;
622 let Some(last_job_id) = ids.last().copied() else {
623 break;
624 };
625 after_job_id = last_job_id;
626
627 let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
628 insert_operation_items(&mut tx, operation_id, &ids).await?;
629 tx.commit().await.map_err(map_sqlx_error)?;
630
631 let state: String =
632 sqlx::query_scalar("SELECT state FROM awa.batch_operations WHERE id = $1")
633 .bind(operation_id)
634 .fetch_one(pool)
635 .await
636 .map_err(map_sqlx_error)?;
637 if state == BatchOperationState::Cancelling.as_str() {
638 return get_batch_operation(pool, operation_id).await;
639 }
640 }
641
642 let operation = sqlx::query_as::<_, BatchOperation>(
643 r#"
644 UPDATE awa.batch_operations
645 SET state = 'running',
646 total_matched = $2,
647 cursor = $3,
648 updated_at = clock_timestamp()
649 WHERE id = $1 AND state = 'scanning'
650 RETURNING *
651 "#,
652 )
653 .bind(operation_id)
654 .bind(total_matched)
655 .bind(Json(BatchOperationCursor {
656 after_job_id: 0,
657 max_job_id,
658 }))
659 .fetch_optional(pool)
660 .await
661 .map_err(map_sqlx_error)?;
662
663 match operation {
664 Some(operation) => Ok(operation),
665 None => get_batch_operation(pool, operation_id).await,
666 }
667}
668
669async fn finalize_operation(
670 pool: &PgPool,
671 id: Uuid,
672 state: BatchOperationState,
673 last_error: Option<String>,
674) -> Result<(), AwaError> {
675 sqlx::query(
676 r#"
677 UPDATE awa.batch_operations
678 SET state = $2,
679 finalized_at = COALESCE(finalized_at, clock_timestamp()),
680 retention_until = COALESCE(retention_until, clock_timestamp() + ($3 * interval '1 day')),
681 last_error = COALESCE($4, last_error),
682 updated_at = clock_timestamp()
683 WHERE id = $1
684 "#,
685 )
686 .bind(id)
687 .bind(state.as_str())
688 .bind(default_retention_days())
689 .bind(last_error)
690 .execute(pool)
691 .await
692 .map_err(map_sqlx_error)?;
693 Ok(())
694}
695
696async fn insert_operation_items(
697 tx: &mut Transaction<'_, Postgres>,
698 operation_id: Uuid,
699 ids: &[i64],
700) -> Result<(), AwaError> {
701 if ids.is_empty() {
702 return Ok(());
703 }
704 sqlx::query(
705 r#"
706 INSERT INTO awa.batch_operation_items (operation_id, job_id)
707 SELECT $1, job_id
708 FROM unnest($2::bigint[]) AS job_id
709 ON CONFLICT (operation_id, job_id) DO NOTHING
710 "#,
711 )
712 .bind(operation_id)
713 .bind(ids)
714 .execute(tx.as_mut())
715 .await
716 .map_err(map_sqlx_error)?;
717 Ok(())
718}
719
720async fn claim_pending_items(
721 pool: &PgPool,
722 operation_id: Uuid,
723 limit: i64,
724) -> Result<Vec<i64>, AwaError> {
725 sqlx::query_scalar(
726 r#"
727 SELECT job_id
728 FROM awa.batch_operation_items
729 WHERE operation_id = $1 AND state = 'pending'
730 ORDER BY job_id ASC
731 LIMIT $2
732 FOR UPDATE SKIP LOCKED
733 "#,
734 )
735 .bind(operation_id)
736 .bind(limit.max(1))
737 .fetch_all(pool)
738 .await
739 .map_err(map_sqlx_error)
740}
741
742async fn skip_pending_items(pool: &PgPool, operation_id: Uuid) -> Result<i64, AwaError> {
743 let skipped: i64 = sqlx::query_scalar(
744 r#"
745 WITH skipped AS (
746 UPDATE awa.batch_operation_items
747 SET state = 'skipped', processed_at = clock_timestamp()
748 WHERE operation_id = $1 AND state = 'pending'
749 RETURNING job_id
750 ), count_skipped AS (
751 SELECT count(*)::bigint AS count FROM skipped
752 )
753 UPDATE awa.batch_operations
754 SET skipped = skipped + count_skipped.count,
755 updated_at = clock_timestamp()
756 FROM count_skipped
757 WHERE id = $1
758 RETURNING count_skipped.count
759 "#,
760 )
761 .bind(operation_id)
762 .fetch_one(pool)
763 .await
764 .map_err(map_sqlx_error)?;
765 Ok(skipped)
766}
767
768fn validate_filter_for_spec(
769 spec: &BatchOperationSpec,
770 filter: &BatchOperationFilter,
771) -> Result<(), AwaError> {
772 match spec {
773 BatchOperationSpec::SetPriority { .. } | BatchOperationSpec::MoveQueue { .. } => {
774 if let Some(state) = filter.state {
775 if !matches!(state, JobState::Scheduled | JobState::Available) {
776 return Err(AwaError::Validation(format!(
777 "{} only supports available or scheduled state filters",
778 spec.kind().as_str()
779 )));
780 }
781 }
782 }
783 }
784 Ok(())
785}
786
787#[derive(Debug, Clone)]
788enum ItemOutcome {
789 Processed,
790 Skipped,
791 Errored(String),
792}
793
794async fn apply_item(
795 pool: &PgPool,
796 operation_id: Uuid,
797 spec: &BatchOperationSpec,
798 job_id: i64,
799) -> Result<ItemOutcome, AwaError> {
800 let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
801 let outcome = match apply_to_job_tx(&mut tx, spec, job_id).await {
802 Ok(true) => ItemOutcome::Processed,
803 Ok(false) => ItemOutcome::Skipped,
804 Err(err) => {
805 tx.rollback().await.map_err(map_sqlx_error)?;
806 let outcome = ItemOutcome::Errored(format!("job_id={job_id}: {err}"));
807 let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
808 record_item_outcome_tx(&mut tx, operation_id, job_id, &outcome).await?;
809 tx.commit().await.map_err(map_sqlx_error)?;
810 return Ok(outcome);
811 }
812 };
813 record_item_outcome_tx(&mut tx, operation_id, job_id, &outcome).await?;
814 tx.commit().await.map_err(map_sqlx_error)?;
815 Ok(outcome)
816}
817
818async fn record_item_outcome_tx(
819 tx: &mut Transaction<'_, Postgres>,
820 operation_id: Uuid,
821 job_id: i64,
822 outcome: &ItemOutcome,
823) -> Result<(), AwaError> {
824 let (state, error, processed_delta, skipped_delta, errored_delta) = match &outcome {
825 ItemOutcome::Processed => ("processed", None, 1_i64, 0_i64, 0_i64),
826 ItemOutcome::Skipped => ("skipped", None, 0_i64, 1_i64, 0_i64),
827 ItemOutcome::Errored(error) => ("errored", Some(error.as_str()), 0_i64, 0_i64, 1_i64),
828 };
829 sqlx::query(
830 r#"
831 UPDATE awa.batch_operation_items
832 SET state = $3,
833 error = $4,
834 processed_at = clock_timestamp()
835 WHERE operation_id = $1 AND job_id = $2 AND state = 'pending'
836 "#,
837 )
838 .bind(operation_id)
839 .bind(job_id)
840 .bind(state)
841 .bind(error)
842 .execute(tx.as_mut())
843 .await
844 .map_err(map_sqlx_error)?;
845 sqlx::query(
846 r#"
847 UPDATE awa.batch_operations
848 SET processed = processed + $2,
849 skipped = skipped + $3,
850 errored = errored + $4,
851 last_error = COALESCE($5, last_error),
852 updated_at = clock_timestamp()
853 WHERE id = $1
854 "#,
855 )
856 .bind(operation_id)
857 .bind(processed_delta)
858 .bind(skipped_delta)
859 .bind(errored_delta)
860 .bind(error)
861 .execute(tx.as_mut())
862 .await
863 .map_err(map_sqlx_error)?;
864 Ok(())
865}
866
867async fn apply_to_job_tx(
868 tx: &mut Transaction<'_, Postgres>,
869 spec: &BatchOperationSpec,
870 job_id: i64,
871) -> Result<bool, AwaError> {
872 if let Some(store) = QueueStorage::active_schema_in_tx(tx)
873 .await?
874 .map(QueueStorage::from_existing_schema)
875 .transpose()?
876 {
877 return match spec {
878 BatchOperationSpec::SetPriority { priority } => {
879 store.set_priority_tx(tx, job_id, *priority).await
880 }
881 BatchOperationSpec::MoveQueue { queue, priority } => {
882 store.move_queue_tx(tx, job_id, queue, *priority).await
883 }
884 };
885 }
886
887 match spec {
888 BatchOperationSpec::SetPriority { priority } => {
889 canonical_set_priority_tx(tx, job_id, *priority).await
890 }
891 BatchOperationSpec::MoveQueue { queue, priority } => {
892 canonical_move_queue_tx(tx, job_id, queue, *priority).await
893 }
894 }
895}
896
897async fn canonical_set_priority_tx(
898 tx: &mut Transaction<'_, Postgres>,
899 job_id: i64,
900 priority: i16,
901) -> Result<bool, AwaError> {
902 let result = sqlx::query(
903 r#"
904 UPDATE awa.jobs
905 SET priority = $2,
906 metadata = CASE
907 WHEN metadata ? '_awa_original_priority' THEN metadata
908 ELSE jsonb_set(metadata, '{_awa_original_priority}', to_jsonb(priority), true)
909 END
910 WHERE id = $1 AND state IN ('available', 'scheduled') AND priority <> $2
911 "#,
912 )
913 .bind(job_id)
914 .bind(priority)
915 .execute(tx.as_mut())
916 .await
917 .map_err(map_sqlx_error)?;
918 Ok(result.rows_affected() > 0)
919}
920
921async fn canonical_move_queue_tx(
922 tx: &mut Transaction<'_, Postgres>,
923 job_id: i64,
924 queue: &str,
925 priority: Option<i16>,
926) -> Result<bool, AwaError> {
927 let result = sqlx::query(
928 r#"
929 WITH target AS (
930 SELECT id, queue, priority, metadata
931 FROM awa.jobs
932 WHERE id = $1 AND state IN ('available', 'scheduled')
933 AND (queue <> $2 OR ($3::smallint IS NOT NULL AND priority <> $3))
934 FOR UPDATE
935 ), stamped AS (
936 SELECT
937 id,
938 CASE
939 WHEN $3::smallint IS NULL OR metadata ? '_awa_original_priority' THEN metadata
940 ELSE jsonb_set(metadata, '{_awa_original_priority}', to_jsonb(priority), true)
941 END AS with_priority,
942 queue AS original_queue
943 FROM target
944 )
945 UPDATE awa.jobs AS jobs
946 SET queue = $2,
947 priority = COALESCE($3, jobs.priority),
948 metadata = CASE
949 WHEN stamped.with_priority ? '_awa_original_queue' THEN stamped.with_priority
950 ELSE jsonb_set(stamped.with_priority, '{_awa_original_queue}', to_jsonb(stamped.original_queue), true)
951 END
952 FROM stamped
953 WHERE jobs.id = stamped.id
954 "#,
955 )
956 .bind(job_id)
957 .bind(queue)
958 .bind(priority)
959 .execute(tx.as_mut())
960 .await
961 .map_err(map_sqlx_error)?;
962 Ok(result.rows_affected() > 0)
963}
964
965async fn count_matching_jobs(
966 pool: &PgPool,
967 filter: &BatchOperationFilter,
968 after_job_id: Option<i64>,
969 max_job_id: Option<i64>,
970) -> Result<i64, AwaError> {
971 if QueueStorage::active_schema(pool).await?.is_some() {
972 count_matching_queue_storage_jobs(pool, filter, after_job_id, max_job_id).await
973 } else {
974 count_matching_canonical_jobs(pool, filter, after_job_id, max_job_id).await
975 }
976}
977
978async fn max_matching_job_id(
979 pool: &PgPool,
980 filter: &BatchOperationFilter,
981) -> Result<Option<i64>, AwaError> {
982 if QueueStorage::active_schema(pool).await?.is_some() {
983 max_matching_queue_storage_job_id(pool, filter).await
984 } else {
985 max_matching_canonical_job_id(pool, filter).await
986 }
987}
988
989async fn max_matching_canonical_job_id(
990 pool: &PgPool,
991 filter: &BatchOperationFilter,
992) -> Result<Option<i64>, AwaError> {
993 sqlx::query_scalar(
994 r#"
995 SELECT max(id)
996 FROM awa.jobs
997 WHERE state IN ('available', 'scheduled')
998 AND ($1::awa.job_state IS NULL OR state = $1)
999 AND ($2::text IS NULL OR kind = $2)
1000 AND ($3::text IS NULL OR queue = $3)
1001 AND ($4::text IS NULL OR tags @> ARRAY[$4]::text[])
1002 AND ($5::bigint[] IS NULL OR id = ANY($5))
1003 AND ($6::timestamptz IS NULL OR created_at >= $6)
1004 AND ($7::timestamptz IS NULL OR created_at < $7)
1005 "#,
1006 )
1007 .bind(filter.state)
1008 .bind(&filter.kind)
1009 .bind(&filter.queue)
1010 .bind(&filter.tag)
1011 .bind(filter.ids.as_deref())
1012 .bind(filter.created_at_gte)
1013 .bind(filter.created_at_lt)
1014 .fetch_one(pool)
1015 .await
1016 .map_err(map_sqlx_error)
1017}
1018
1019async fn max_matching_queue_storage_job_id(
1020 pool: &PgPool,
1021 filter: &BatchOperationFilter,
1022) -> Result<Option<i64>, AwaError> {
1023 let schema = QueueStorage::active_schema(pool)
1024 .await?
1025 .ok_or_else(|| AwaError::Validation("queue storage is not active".to_string()))?;
1026 sqlx::query_scalar(&format!(
1027 r#"
1028 WITH candidates AS (
1029 SELECT
1030 ready.job_id AS id,
1031 ready.kind,
1032 ready.queue,
1033 'available'::awa.job_state AS state,
1034 ready.created_at,
1035 COALESCE(ready.payload, '{{}}'::jsonb) AS payload
1036 FROM {schema}.ready_entries AS ready
1037 JOIN {schema}.queue_claim_heads AS claims
1038 ON claims.queue = ready.queue
1039 AND claims.priority = ready.priority
1040 AND claims.enqueue_shard = ready.enqueue_shard
1041 WHERE ready.lane_seq >= {schema}.sequence_next_value(claims.seq_name)
1042 AND NOT EXISTS (
1043 SELECT 1 FROM {schema}.ready_tombstones AS tomb
1044 WHERE tomb.queue = ready.queue
1045 AND tomb.priority = ready.priority
1046 AND tomb.enqueue_shard = ready.enqueue_shard
1047 AND tomb.lane_seq = ready.lane_seq
1048 AND tomb.ready_slot = ready.ready_slot
1049 AND tomb.ready_generation = ready.ready_generation
1050 )
1051 UNION ALL
1052 SELECT
1053 deferred.job_id AS id,
1054 deferred.kind,
1055 deferred.queue,
1056 deferred.state,
1057 deferred.created_at,
1058 COALESCE(deferred.payload, '{{}}'::jsonb) AS payload
1059 FROM {schema}.deferred_jobs AS deferred
1060 WHERE deferred.state = 'scheduled'
1061 )
1062 SELECT max(id)
1063 FROM candidates
1064 WHERE ($1::awa.job_state IS NULL OR state = $1)
1065 AND ($2::text IS NULL OR kind = $2)
1066 AND ($3::text IS NULL OR queue = $3)
1067 AND ($4::text IS NULL OR COALESCE(ARRAY(SELECT jsonb_array_elements_text(payload->'tags')), ARRAY[]::text[]) @> ARRAY[$4]::text[])
1068 AND ($5::bigint[] IS NULL OR id = ANY($5))
1069 AND ($6::timestamptz IS NULL OR created_at >= $6)
1070 AND ($7::timestamptz IS NULL OR created_at < $7)
1071 "#
1072 ))
1073 .bind(filter.state)
1074 .bind(&filter.kind)
1075 .bind(&filter.queue)
1076 .bind(&filter.tag)
1077 .bind(filter.ids.as_deref())
1078 .bind(filter.created_at_gte)
1079 .bind(filter.created_at_lt)
1080 .fetch_one(pool)
1081 .await
1082 .map_err(map_sqlx_error)
1083}
1084
1085async fn load_matching_jobs(
1086 pool: &PgPool,
1087 filter: &BatchOperationFilter,
1088 after_job_id: i64,
1089 max_job_id: i64,
1090 limit: i64,
1091) -> Result<Vec<JobRow>, AwaError> {
1092 if let Some(store) = QueueStorage::active_schema(pool)
1093 .await?
1094 .map(QueueStorage::from_existing_schema)
1095 .transpose()?
1096 {
1097 load_matching_queue_storage_jobs(&store, pool, filter, after_job_id, max_job_id, limit)
1098 .await
1099 } else {
1100 load_matching_canonical_jobs(pool, filter, after_job_id, max_job_id, limit).await
1101 }
1102}
1103
1104async fn load_matching_job_ids(
1105 pool: &PgPool,
1106 filter: &BatchOperationFilter,
1107 after_job_id: i64,
1108 max_job_id: i64,
1109 limit: i64,
1110) -> Result<Vec<i64>, AwaError> {
1111 if QueueStorage::active_schema(pool).await?.is_some() {
1112 load_matching_queue_storage_job_ids(pool, filter, after_job_id, max_job_id, limit).await
1113 } else {
1114 load_matching_canonical_job_ids(pool, filter, after_job_id, max_job_id, limit).await
1115 }
1116}
1117
1118async fn load_matching_canonical_job_ids(
1119 pool: &PgPool,
1120 filter: &BatchOperationFilter,
1121 after_job_id: i64,
1122 max_job_id: i64,
1123 limit: i64,
1124) -> Result<Vec<i64>, AwaError> {
1125 sqlx::query_scalar(
1126 r#"
1127 SELECT id
1128 FROM awa.jobs
1129 WHERE state IN ('available', 'scheduled')
1130 AND ($1::awa.job_state IS NULL OR state = $1)
1131 AND ($2::text IS NULL OR kind = $2)
1132 AND ($3::text IS NULL OR queue = $3)
1133 AND ($4::text IS NULL OR tags @> ARRAY[$4]::text[])
1134 AND ($5::bigint[] IS NULL OR id = ANY($5))
1135 AND ($6::timestamptz IS NULL OR created_at >= $6)
1136 AND ($7::timestamptz IS NULL OR created_at < $7)
1137 AND id > $8
1138 AND id <= $9
1139 ORDER BY id ASC
1140 LIMIT $10
1141 "#,
1142 )
1143 .bind(filter.state)
1144 .bind(&filter.kind)
1145 .bind(&filter.queue)
1146 .bind(&filter.tag)
1147 .bind(filter.ids.as_deref())
1148 .bind(filter.created_at_gte)
1149 .bind(filter.created_at_lt)
1150 .bind(after_job_id)
1151 .bind(max_job_id)
1152 .bind(limit.max(1))
1153 .fetch_all(pool)
1154 .await
1155 .map_err(map_sqlx_error)
1156}
1157
1158async fn load_matching_queue_storage_job_ids(
1159 pool: &PgPool,
1160 filter: &BatchOperationFilter,
1161 after_job_id: i64,
1162 max_job_id: i64,
1163 limit: i64,
1164) -> Result<Vec<i64>, AwaError> {
1165 let schema = QueueStorage::active_schema(pool)
1166 .await?
1167 .ok_or_else(|| AwaError::Validation("queue storage is not active".to_string()))?;
1168 sqlx::query_scalar(&format!(
1169 r#"
1170 WITH candidates AS (
1171 SELECT
1172 ready.job_id AS id,
1173 ready.kind,
1174 ready.queue,
1175 'available'::awa.job_state AS state,
1176 ready.created_at,
1177 COALESCE(ready.payload, '{{}}'::jsonb) AS payload
1178 FROM {schema}.ready_entries AS ready
1179 JOIN {schema}.queue_claim_heads AS claims
1180 ON claims.queue = ready.queue
1181 AND claims.priority = ready.priority
1182 AND claims.enqueue_shard = ready.enqueue_shard
1183 WHERE ready.lane_seq >= {schema}.sequence_next_value(claims.seq_name)
1184 AND NOT EXISTS (
1185 SELECT 1 FROM {schema}.ready_tombstones AS tomb
1186 WHERE tomb.queue = ready.queue
1187 AND tomb.priority = ready.priority
1188 AND tomb.enqueue_shard = ready.enqueue_shard
1189 AND tomb.lane_seq = ready.lane_seq
1190 AND tomb.ready_slot = ready.ready_slot
1191 AND tomb.ready_generation = ready.ready_generation
1192 )
1193 UNION ALL
1194 SELECT
1195 deferred.job_id AS id,
1196 deferred.kind,
1197 deferred.queue,
1198 deferred.state,
1199 deferred.created_at,
1200 COALESCE(deferred.payload, '{{}}'::jsonb) AS payload
1201 FROM {schema}.deferred_jobs AS deferred
1202 WHERE deferred.state = 'scheduled'
1203 )
1204 SELECT id
1205 FROM candidates
1206 WHERE ($1::awa.job_state IS NULL OR state = $1)
1207 AND ($2::text IS NULL OR kind = $2)
1208 AND ($3::text IS NULL OR queue = $3)
1209 AND ($4::text IS NULL OR COALESCE(ARRAY(SELECT jsonb_array_elements_text(payload->'tags')), ARRAY[]::text[]) @> ARRAY[$4]::text[])
1210 AND ($5::bigint[] IS NULL OR id = ANY($5))
1211 AND ($6::timestamptz IS NULL OR created_at >= $6)
1212 AND ($7::timestamptz IS NULL OR created_at < $7)
1213 AND id > $8
1214 AND id <= $9
1215 ORDER BY id ASC
1216 LIMIT $10
1217 "#
1218 ))
1219 .bind(filter.state)
1220 .bind(&filter.kind)
1221 .bind(&filter.queue)
1222 .bind(&filter.tag)
1223 .bind(filter.ids.as_deref())
1224 .bind(filter.created_at_gte)
1225 .bind(filter.created_at_lt)
1226 .bind(after_job_id)
1227 .bind(max_job_id)
1228 .bind(limit.max(1))
1229 .fetch_all(pool)
1230 .await
1231 .map_err(map_sqlx_error)
1232}
1233
1234async fn count_matching_canonical_jobs(
1235 pool: &PgPool,
1236 filter: &BatchOperationFilter,
1237 after_job_id: Option<i64>,
1238 max_job_id: Option<i64>,
1239) -> Result<i64, AwaError> {
1240 sqlx::query_scalar(
1241 r#"
1242 SELECT count(*)::bigint
1243 FROM awa.jobs
1244 WHERE state IN ('available', 'scheduled')
1245 AND ($1::awa.job_state IS NULL OR state = $1)
1246 AND ($2::text IS NULL OR kind = $2)
1247 AND ($3::text IS NULL OR queue = $3)
1248 AND ($4::text IS NULL OR tags @> ARRAY[$4]::text[])
1249 AND ($5::bigint[] IS NULL OR id = ANY($5))
1250 AND ($6::timestamptz IS NULL OR created_at >= $6)
1251 AND ($7::timestamptz IS NULL OR created_at < $7)
1252 AND ($8::bigint IS NULL OR id > $8)
1253 AND ($9::bigint IS NULL OR id <= $9)
1254 "#,
1255 )
1256 .bind(filter.state)
1257 .bind(&filter.kind)
1258 .bind(&filter.queue)
1259 .bind(&filter.tag)
1260 .bind(filter.ids.as_deref())
1261 .bind(filter.created_at_gte)
1262 .bind(filter.created_at_lt)
1263 .bind(after_job_id)
1264 .bind(max_job_id)
1265 .fetch_one(pool)
1266 .await
1267 .map_err(map_sqlx_error)
1268}
1269
1270async fn load_matching_canonical_jobs(
1271 pool: &PgPool,
1272 filter: &BatchOperationFilter,
1273 after_job_id: i64,
1274 max_job_id: i64,
1275 limit: i64,
1276) -> Result<Vec<JobRow>, AwaError> {
1277 sqlx::query_as::<_, JobRow>(
1278 r#"
1279 SELECT *
1280 FROM awa.jobs
1281 WHERE state IN ('available', 'scheduled')
1282 AND ($1::awa.job_state IS NULL OR state = $1)
1283 AND ($2::text IS NULL OR kind = $2)
1284 AND ($3::text IS NULL OR queue = $3)
1285 AND ($4::text IS NULL OR tags @> ARRAY[$4]::text[])
1286 AND ($5::bigint[] IS NULL OR id = ANY($5))
1287 AND ($6::timestamptz IS NULL OR created_at >= $6)
1288 AND ($7::timestamptz IS NULL OR created_at < $7)
1289 AND id > $8
1290 AND id <= $9
1291 ORDER BY id ASC
1292 LIMIT $10
1293 "#,
1294 )
1295 .bind(filter.state)
1296 .bind(&filter.kind)
1297 .bind(&filter.queue)
1298 .bind(&filter.tag)
1299 .bind(filter.ids.as_deref())
1300 .bind(filter.created_at_gte)
1301 .bind(filter.created_at_lt)
1302 .bind(after_job_id)
1303 .bind(max_job_id)
1304 .bind(limit.max(1))
1305 .fetch_all(pool)
1306 .await
1307 .map_err(map_sqlx_error)
1308}
1309
1310async fn count_matching_queue_storage_jobs(
1311 pool: &PgPool,
1312 filter: &BatchOperationFilter,
1313 after_job_id: Option<i64>,
1314 max_job_id: Option<i64>,
1315) -> Result<i64, AwaError> {
1316 let schema = QueueStorage::active_schema(pool)
1317 .await?
1318 .ok_or_else(|| AwaError::Validation("queue storage is not active".to_string()))?;
1319 sqlx::query_scalar(&format!(
1320 r#"
1321 WITH candidates AS (
1322 SELECT
1323 ready.job_id AS id,
1324 ready.kind,
1325 ready.queue,
1326 'available'::awa.job_state AS state,
1327 ready.created_at,
1328 COALESCE(ready.payload, '{{}}'::jsonb) AS payload
1329 FROM {schema}.ready_entries AS ready
1330 JOIN {schema}.queue_claim_heads AS claims
1331 ON claims.queue = ready.queue
1332 AND claims.priority = ready.priority
1333 AND claims.enqueue_shard = ready.enqueue_shard
1334 WHERE ready.lane_seq >= {schema}.sequence_next_value(claims.seq_name)
1335 AND NOT EXISTS (
1336 SELECT 1 FROM {schema}.ready_tombstones AS tomb
1337 WHERE tomb.queue = ready.queue
1338 AND tomb.priority = ready.priority
1339 AND tomb.enqueue_shard = ready.enqueue_shard
1340 AND tomb.lane_seq = ready.lane_seq
1341 AND tomb.ready_slot = ready.ready_slot
1342 AND tomb.ready_generation = ready.ready_generation
1343 )
1344 UNION ALL
1345 SELECT
1346 deferred.job_id AS id,
1347 deferred.kind,
1348 deferred.queue,
1349 deferred.state,
1350 deferred.created_at,
1351 COALESCE(deferred.payload, '{{}}'::jsonb) AS payload
1352 FROM {schema}.deferred_jobs AS deferred
1353 WHERE deferred.state = 'scheduled'
1354 )
1355 SELECT count(*)::bigint
1356 FROM candidates
1357 WHERE ($1::awa.job_state IS NULL OR state = $1)
1358 AND ($2::text IS NULL OR kind = $2)
1359 AND ($3::text IS NULL OR queue = $3)
1360 AND ($4::text IS NULL OR COALESCE(ARRAY(SELECT jsonb_array_elements_text(payload->'tags')), ARRAY[]::text[]) @> ARRAY[$4]::text[])
1361 AND ($5::bigint[] IS NULL OR id = ANY($5))
1362 AND ($6::timestamptz IS NULL OR created_at >= $6)
1363 AND ($7::timestamptz IS NULL OR created_at < $7)
1364 AND ($8::bigint IS NULL OR id > $8)
1365 AND ($9::bigint IS NULL OR id <= $9)
1366 "#
1367 ))
1368 .bind(filter.state)
1369 .bind(&filter.kind)
1370 .bind(&filter.queue)
1371 .bind(&filter.tag)
1372 .bind(filter.ids.as_deref())
1373 .bind(filter.created_at_gte)
1374 .bind(filter.created_at_lt)
1375 .bind(after_job_id)
1376 .bind(max_job_id)
1377 .fetch_one(pool)
1378 .await
1379 .map_err(map_sqlx_error)
1380}
1381
1382async fn load_matching_queue_storage_jobs(
1383 _store: &QueueStorage,
1384 pool: &PgPool,
1385 filter: &BatchOperationFilter,
1386 after_job_id: i64,
1387 max_job_id: i64,
1388 limit: i64,
1389) -> Result<Vec<JobRow>, AwaError> {
1390 let schema = _store.schema();
1391 let ids: Vec<i64> = sqlx::query_scalar(&format!(
1392 r#"
1393 WITH candidates AS (
1394 SELECT
1395 ready.job_id AS id,
1396 ready.kind,
1397 ready.queue,
1398 'available'::awa.job_state AS state,
1399 ready.created_at,
1400 COALESCE(ready.payload, '{{}}'::jsonb) AS payload
1401 FROM {schema}.ready_entries AS ready
1402 JOIN {schema}.queue_claim_heads AS claims
1403 ON claims.queue = ready.queue
1404 AND claims.priority = ready.priority
1405 AND claims.enqueue_shard = ready.enqueue_shard
1406 WHERE ready.lane_seq >= {schema}.sequence_next_value(claims.seq_name)
1407 AND NOT EXISTS (
1408 SELECT 1 FROM {schema}.ready_tombstones AS tomb
1409 WHERE tomb.queue = ready.queue
1410 AND tomb.priority = ready.priority
1411 AND tomb.enqueue_shard = ready.enqueue_shard
1412 AND tomb.lane_seq = ready.lane_seq
1413 AND tomb.ready_slot = ready.ready_slot
1414 AND tomb.ready_generation = ready.ready_generation
1415 )
1416 UNION ALL
1417 SELECT
1418 deferred.job_id AS id,
1419 deferred.kind,
1420 deferred.queue,
1421 deferred.state,
1422 deferred.created_at,
1423 COALESCE(deferred.payload, '{{}}'::jsonb) AS payload
1424 FROM {schema}.deferred_jobs AS deferred
1425 WHERE deferred.state = 'scheduled'
1426 )
1427 SELECT id
1428 FROM candidates
1429 WHERE ($1::awa.job_state IS NULL OR state = $1)
1430 AND ($2::text IS NULL OR kind = $2)
1431 AND ($3::text IS NULL OR queue = $3)
1432 AND ($4::text IS NULL OR COALESCE(ARRAY(SELECT jsonb_array_elements_text(payload->'tags')), ARRAY[]::text[]) @> ARRAY[$4]::text[])
1433 AND ($5::bigint[] IS NULL OR id = ANY($5))
1434 AND ($6::timestamptz IS NULL OR created_at >= $6)
1435 AND ($7::timestamptz IS NULL OR created_at < $7)
1436 AND id > $8
1437 AND id <= $9
1438 ORDER BY id ASC
1439 LIMIT $10
1440 "#
1441 ))
1442 .bind(filter.state)
1443 .bind(&filter.kind)
1444 .bind(&filter.queue)
1445 .bind(&filter.tag)
1446 .bind(filter.ids.as_deref())
1447 .bind(filter.created_at_gte)
1448 .bind(filter.created_at_lt)
1449 .bind(after_job_id)
1450 .bind(max_job_id)
1451 .bind(limit.max(1))
1452 .fetch_all(pool)
1453 .await
1454 .map_err(map_sqlx_error)?;
1455
1456 let mut jobs = Vec::with_capacity(ids.len());
1457 for id in ids {
1458 if let Some(job) = _store.load_job(pool, id).await? {
1459 if matches_batch_filter(&job, filter, after_job_id, max_job_id) {
1460 jobs.push(job);
1461 }
1462 }
1463 }
1464 jobs.sort_by_key(|job| job.id);
1465 Ok(jobs)
1466}
1467
1468fn matches_batch_filter(
1469 job: &JobRow,
1470 filter: &BatchOperationFilter,
1471 after_job_id: i64,
1472 max_job_id: i64,
1473) -> bool {
1474 if job.id <= after_job_id || job.id > max_job_id {
1475 return false;
1476 }
1477 if !matches!(job.state, JobState::Available | JobState::Scheduled) {
1478 return false;
1479 }
1480 if filter.state.is_some_and(|state| job.state != state) {
1481 return false;
1482 }
1483 if filter.kind.as_ref().is_some_and(|kind| &job.kind != kind) {
1484 return false;
1485 }
1486 if filter
1487 .queue
1488 .as_ref()
1489 .is_some_and(|queue| &job.queue != queue)
1490 {
1491 return false;
1492 }
1493 if filter
1494 .tag
1495 .as_ref()
1496 .is_some_and(|tag| !job.tags.iter().any(|job_tag| job_tag == tag))
1497 {
1498 return false;
1499 }
1500 if let Some(ids) = &filter.ids {
1501 if !ids.contains(&job.id) {
1502 return false;
1503 }
1504 }
1505 if filter
1506 .created_at_gte
1507 .is_some_and(|created_at_gte| job.created_at < created_at_gte)
1508 {
1509 return false;
1510 }
1511 if filter
1512 .created_at_lt
1513 .is_some_and(|created_at_lt| job.created_at >= created_at_lt)
1514 {
1515 return false;
1516 }
1517 true
1518}
1519
1520pub async fn run_one_default_chunk(
1521 pool: &PgPool,
1522 runner_instance: Uuid,
1523) -> Result<BatchOperationRunOutcome, AwaError> {
1524 run_one_batch_operation_chunk(pool, runner_instance, DEFAULT_CHUNK_SIZE).await
1525}