1use crate::error::AwaError;
2use crate::job::{InsertOpts, InsertParams, JobRow, JobState};
3use crate::unique::compute_unique_key;
4use crate::JobArgs;
5use sqlx::postgres::PgConnection;
6use sqlx::{Connection, PgExecutor, PgPool};
7
8const COPY_NULL_SENTINEL: &str = "__AWA_NULL__";
9const ACTIVE_STORAGE_ENGINE_SQL: &str = "SELECT awa.active_storage_engine()";
10
11pub(crate) fn reject_null_bytes(value: &serde_json::Value) -> Result<(), AwaError> {
24 match value {
25 serde_json::Value::String(s) if s.contains('\0') => Err(AwaError::Validation(
26 "job args/metadata must not contain null bytes (\\u0000): Postgres JSONB does not support them".into(),
27 )),
28 serde_json::Value::Array(arr) => {
29 for v in arr {
30 reject_null_bytes(v)?;
31 }
32 Ok(())
33 }
34 serde_json::Value::Object(map) => {
35 for (k, v) in map {
36 if k.contains('\0') {
37 return Err(AwaError::Validation(
38 "job args/metadata keys must not contain null bytes (\\u0000)".into(),
39 ));
40 }
41 reject_null_bytes(v)?;
42 }
43 Ok(())
44 }
45 _ => Ok(()),
46 }
47}
48
49pub(crate) struct PreparedRow {
53 pub kind: String,
54 pub queue: String,
55 pub args: serde_json::Value,
56 pub state: JobState,
57 pub priority: i16,
58 pub max_attempts: i16,
59 pub run_at: Option<chrono::DateTime<chrono::Utc>>,
60 pub metadata: serde_json::Value,
61 pub tags: Vec<String>,
62 pub unique_key: Option<Vec<u8>>,
63 pub unique_states: Option<String>,
64}
65
66#[derive(Clone, Copy, Debug, Eq, PartialEq)]
67enum TargetTable {
68 JobsHot,
69 ScheduledJobs,
70}
71
72impl TargetTable {
73 fn as_str(self) -> &'static str {
74 match self {
75 TargetTable::JobsHot => "awa.jobs_hot",
76 TargetTable::ScheduledJobs => "awa.scheduled_jobs",
77 }
78 }
79}
80
81fn target_table_for_state(state: JobState) -> TargetTable {
82 match state {
83 JobState::Scheduled | JobState::Retryable => TargetTable::ScheduledJobs,
84 _ => TargetTable::JobsHot,
85 }
86}
87
88fn homogeneous_target_table(rows: &[PreparedRow]) -> Option<TargetTable> {
89 let first = rows.first().map(|row| target_table_for_state(row.state))?;
90 rows.iter()
91 .all(|row| target_table_for_state(row.state) == first)
92 .then_some(first)
93}
94
95fn map_sqlx_error(err: sqlx::Error) -> AwaError {
96 if let sqlx::Error::Database(ref db_err) = err {
97 if db_err.code().as_deref() == Some("23505") {
98 return AwaError::UniqueConflict {
99 constraint: db_err.constraint().map(|c| c.to_string()),
100 };
101 }
102 }
103 AwaError::Database(err)
104}
105
106async fn insert_compat_row<'e, E>(executor: E, row: &PreparedRow) -> Result<JobRow, sqlx::Error>
107where
108 E: PgExecutor<'e>,
109{
110 sqlx::query_as::<_, JobRow>(INSERT_COMPAT_SQL)
111 .bind(&row.kind)
112 .bind(&row.queue)
113 .bind(&row.args)
114 .bind(row.state)
115 .bind(row.priority)
116 .bind(row.max_attempts)
117 .bind(row.run_at)
118 .bind(&row.metadata)
119 .bind(&row.tags)
120 .bind(&row.unique_key)
121 .bind(&row.unique_states)
122 .fetch_one(executor)
123 .await
124}
125
126async fn insert_compat_rows(
127 conn: &mut PgConnection,
128 rows: &[PreparedRow],
129 ignore_unique_conflicts: bool,
130) -> Result<Vec<JobRow>, AwaError> {
131 let mut inserted = Vec::with_capacity(rows.len());
132
133 for row in rows {
134 if ignore_unique_conflicts {
135 let mut savepoint = conn.begin().await?;
136 match insert_compat_row(&mut *savepoint, row).await {
137 Ok(job) => {
138 savepoint.commit().await?;
139 inserted.push(job);
140 }
141 Err(sqlx::Error::Database(db_err)) if db_err.code().as_deref() == Some("23505") => {
142 savepoint.rollback().await?;
143 }
144 Err(err) => {
145 savepoint.rollback().await?;
146 return Err(map_sqlx_error(err));
147 }
148 }
149 } else {
150 inserted.push(
151 insert_compat_row(&mut *conn, row)
152 .await
153 .map_err(map_sqlx_error)?,
154 );
155 }
156 }
157
158 Ok(inserted)
159}
160
161const INSERT_COMPAT_SQL: &str = r#"
162 SELECT *
163 FROM awa.insert_job_compat(
164 $1,
165 $2,
166 $3::jsonb,
167 $4::awa.job_state,
168 $5::smallint,
169 $6::smallint,
170 $7::timestamptz,
171 $8::jsonb,
172 $9::text[],
173 $10::bytea,
174 $11::bit(8)
175 )
176"#;
177
178fn build_multi_insert_query(target_table: &str, count: usize) -> String {
179 let mut query = "WITH active AS (SELECT awa.active_storage_engine() AS engine), \
180 rows(ord, kind, queue, args, state, priority, max_attempts, run_at, metadata, tags, unique_key, unique_states) AS (VALUES "
181 .to_string();
182
183 let params_per_row = 11u32;
184 let mut param_index = 1u32;
185 for i in 0..count {
186 if i > 0 {
187 query.push_str(", ");
188 }
189 query.push_str(&format!(
190 "({}, ${}, ${}, ${}::jsonb, ${}::awa.job_state, ${}::smallint, ${}::smallint, ${}::timestamptz, ${}::jsonb, ${}::text[], ${}::bytea, ${}::bit(8))",
191 i + 1,
192 param_index,
193 param_index + 1,
194 param_index + 2,
195 param_index + 3,
196 param_index + 4,
197 param_index + 5,
198 param_index + 6,
199 param_index + 7,
200 param_index + 8,
201 param_index + 9,
202 param_index + 10,
203 ));
204 param_index += params_per_row;
205 }
206 query.push_str(&format!(
207 "), canonical_inserted AS ( \
208 INSERT INTO {} (kind, queue, args, state, priority, max_attempts, run_at, metadata, tags, unique_key, unique_states) \
209 SELECT rows.kind, rows.queue, rows.args, rows.state, rows.priority, rows.max_attempts, COALESCE(rows.run_at, now()), rows.metadata, rows.tags, rows.unique_key, rows.unique_states \
210 FROM rows \
211 CROSS JOIN active \
212 WHERE active.engine = 'canonical' \
213 ORDER BY rows.ord \
214 RETURNING * \
215 ), compat_inserted AS ( \
216 SELECT inserted.* \
217 FROM rows \
218 CROSS JOIN active \
219 CROSS JOIN LATERAL awa.insert_job_compat( \
220 rows.kind, \
221 rows.queue, \
222 rows.args, \
223 rows.state, \
224 rows.priority, \
225 rows.max_attempts, \
226 rows.run_at, \
227 rows.metadata, \
228 rows.tags, \
229 rows.unique_key, \
230 rows.unique_states \
231 ) AS inserted \
232 WHERE active.engine <> 'canonical' \
233 ) \
234 SELECT * FROM canonical_inserted \
235 UNION ALL \
236 SELECT * FROM compat_inserted",
237 target_table
238 ));
239 query
240}
241
242fn build_copy_insert_query(target_table: &str) -> String {
243 format!(
244 r#"
245 WITH _storage_guard AS (
246 SELECT awa.assert_writable_canonical_storage() AS ok
247 )
248 INSERT INTO {target_table} (
249 kind,
250 queue,
251 args,
252 state,
253 priority,
254 max_attempts,
255 run_at,
256 metadata,
257 tags,
258 unique_key,
259 unique_states
260 )
261 SELECT
262 s.kind,
263 s.queue,
264 s.args,
265 s.state,
266 s.priority,
267 s.max_attempts,
268 COALESCE(s.run_at, now()),
269 s.metadata,
270 s.tags,
271 s.unique_key,
272 s.unique_states
273 FROM pg_temp.awa_copy_staging AS s
274 CROSS JOIN _storage_guard
275 RETURNING *
276 "#
277 )
278}
279
280fn build_copy_unique_insert_query(target_table: &str) -> String {
281 format!(
282 r#"
283 WITH _storage_guard AS (
284 SELECT awa.assert_writable_canonical_storage() AS ok
285 ),
286 staged AS (
287 SELECT
288 s.ord,
289 s.kind,
290 s.queue,
291 s.args,
292 s.state,
293 s.priority,
294 s.max_attempts,
295 s.run_at,
296 s.metadata,
297 s.tags,
298 s.unique_key,
299 s.unique_states
300 FROM pg_temp.awa_copy_staging AS s
301 CROSS JOIN _storage_guard
302 ),
303 claimable AS (
304 SELECT DISTINCT ON (s.unique_key)
305 s.ord,
306 s.kind,
307 s.queue,
308 s.args,
309 s.state,
310 s.priority,
311 s.max_attempts,
312 s.run_at,
313 s.metadata,
314 s.tags,
315 s.unique_key,
316 s.unique_states
317 FROM staged AS s
318 WHERE s.unique_key IS NOT NULL
319 AND s.unique_states IS NOT NULL
320 AND awa.job_state_in_bitmask(s.unique_states, s.state)
321 ORDER BY s.unique_key, s.ord
322 ),
323 claimed AS (
324 INSERT INTO awa.job_unique_claims (unique_key, job_id)
325 SELECT c.unique_key, nextval('awa.jobs_id_seq')
326 FROM claimable AS c
327 ON CONFLICT (unique_key) DO NOTHING
328 RETURNING unique_key, job_id
329 ),
330 prepared AS (
331 SELECT
332 c.ord,
333 cl.job_id AS id,
334 c.kind,
335 c.queue,
336 c.args,
337 c.state,
338 c.priority,
339 c.max_attempts,
340 c.run_at,
341 c.metadata,
342 c.tags,
343 c.unique_key,
344 c.unique_states
345 FROM claimable AS c
346 JOIN claimed AS cl USING (unique_key)
347
348 UNION ALL
349
350 SELECT
351 s.ord,
352 nextval('awa.jobs_id_seq') AS id,
353 s.kind,
354 s.queue,
355 s.args,
356 s.state,
357 s.priority,
358 s.max_attempts,
359 s.run_at,
360 s.metadata,
361 s.tags,
362 s.unique_key,
363 s.unique_states
364 FROM staged AS s
365 WHERE s.unique_key IS NULL
366 OR s.unique_states IS NULL
367 OR NOT awa.job_state_in_bitmask(s.unique_states, s.state)
368 ),
369 inserted AS (
370 INSERT INTO {target_table} (
371 id,
372 kind,
373 queue,
374 args,
375 state,
376 priority,
377 max_attempts,
378 run_at,
379 metadata,
380 tags,
381 unique_key,
382 unique_states
383 )
384 SELECT
385 p.id,
386 p.kind,
387 p.queue,
388 p.args,
389 p.state,
390 p.priority,
391 p.max_attempts,
392 COALESCE(p.run_at, now()),
393 p.metadata,
394 p.tags,
395 p.unique_key,
396 p.unique_states
397 FROM prepared AS p
398 ORDER BY p.ord
399 RETURNING *
400 )
401 SELECT inserted.*
402 FROM inserted
403 JOIN prepared ON prepared.id = inserted.id
404 ORDER BY prepared.ord
405 "#
406 )
407}
408
409fn compute_unique_fields(
411 kind: &str,
412 args: &serde_json::Value,
413 opts: &InsertOpts,
414) -> (Option<Vec<u8>>, Option<String>) {
415 let unique_key = opts.unique.as_ref().map(|u| {
416 compute_unique_key(
417 kind,
418 if u.by_queue { Some(&opts.queue) } else { None },
419 if u.by_args { Some(args) } else { None },
420 u.by_period,
421 )
422 });
423
424 let unique_states = opts.unique.as_ref().map(|u| {
425 let mut bit_string = String::with_capacity(8);
429 for bit_position in 0..8 {
430 if u.states & (1 << bit_position) != 0 {
431 bit_string.push('1');
432 } else {
433 bit_string.push('0');
434 }
435 }
436 bit_string
437 });
438
439 (unique_key, unique_states)
440}
441
442pub(crate) fn prepare_row(args: &impl JobArgs, opts: InsertOpts) -> Result<PreparedRow, AwaError> {
446 let kind = args.kind_str().to_string();
447 let args_value = args.to_args()?;
448 prepare_row_raw(kind, args_value, opts)
449}
450
451pub(crate) fn prepare_row_raw(
453 kind: String,
454 args: serde_json::Value,
455 opts: InsertOpts,
456) -> Result<PreparedRow, AwaError> {
457 reject_null_bytes(&args)?;
458 reject_null_bytes(&opts.metadata)?;
459
460 let state = if opts.run_at.is_some() {
461 JobState::Scheduled
462 } else {
463 JobState::Available
464 };
465
466 let (unique_key, unique_states) = compute_unique_fields(&kind, &args, &opts);
467
468 Ok(PreparedRow {
469 kind,
470 queue: opts.queue,
471 args,
472 state,
473 priority: opts.priority,
474 max_attempts: opts.max_attempts,
475 run_at: opts.run_at,
476 metadata: opts.metadata,
477 tags: opts.tags,
478 unique_key,
479 unique_states,
480 })
481}
482
483pub async fn insert<'e, E>(executor: E, args: &impl JobArgs) -> Result<JobRow, AwaError>
487where
488 E: PgExecutor<'e>,
489{
490 insert_with(executor, args, InsertOpts::default()).await
491}
492
493#[tracing::instrument(skip(executor, args), fields(job.kind = args.kind_str(), job.queue = %opts.queue))]
495pub async fn insert_with<'e, E>(
496 executor: E,
497 args: &impl JobArgs,
498 opts: InsertOpts,
499) -> Result<JobRow, AwaError>
500where
501 E: PgExecutor<'e>,
502{
503 let row = prepare_row(args, opts)?;
504 sqlx::query_as::<_, JobRow>(INSERT_COMPAT_SQL)
505 .bind(&row.kind)
506 .bind(&row.queue)
507 .bind(&row.args)
508 .bind(row.state)
509 .bind(row.priority)
510 .bind(row.max_attempts)
511 .bind(row.run_at)
512 .bind(&row.metadata)
513 .bind(&row.tags)
514 .bind(&row.unique_key)
515 .bind(&row.unique_states)
516 .fetch_one(executor)
517 .await
518 .map_err(map_sqlx_error)
519}
520
521fn precompute_rows(jobs: &[InsertParams]) -> Result<Vec<PreparedRow>, AwaError> {
523 jobs.iter()
524 .map(|job| prepare_row_raw(job.kind.clone(), job.args.clone(), job.opts.clone()))
525 .collect()
526}
527
528#[tracing::instrument(skip(executor, jobs), fields(job.count = jobs.len()))]
533pub async fn insert_many<'e, E>(executor: E, jobs: &[InsertParams]) -> Result<Vec<JobRow>, AwaError>
534where
535 E: PgExecutor<'e>,
536{
537 if jobs.is_empty() {
538 return Ok(Vec::new());
539 }
540
541 let rows = precompute_rows(jobs)?;
542 let target_table = homogeneous_target_table(&rows)
543 .map(TargetTable::as_str)
544 .unwrap_or("awa.jobs");
545 let query = build_multi_insert_query(target_table, rows.len());
546
547 let mut sql_query = sqlx::query_as::<_, JobRow>(&query);
548
549 for row in &rows {
550 sql_query = sql_query
551 .bind(&row.kind)
552 .bind(&row.queue)
553 .bind(&row.args)
554 .bind(row.state)
555 .bind(row.priority)
556 .bind(row.max_attempts)
557 .bind(row.run_at)
558 .bind(&row.metadata)
559 .bind(&row.tags)
560 .bind(&row.unique_key)
561 .bind(&row.unique_states);
562 }
563
564 let results = sql_query
565 .fetch_all(executor)
566 .await
567 .map_err(map_sqlx_error)?;
568
569 Ok(results)
570}
571
572#[tracing::instrument(skip(conn, jobs), fields(job.count = jobs.len()))]
581pub async fn insert_many_copy(
582 conn: &mut PgConnection,
583 jobs: &[InsertParams],
584) -> Result<Vec<JobRow>, AwaError> {
585 if jobs.is_empty() {
586 return Ok(Vec::new());
587 }
588
589 let rows = precompute_rows(jobs)?;
590 let active_engine: String = sqlx::query_scalar(ACTIVE_STORAGE_ENGINE_SQL)
591 .fetch_one(&mut *conn)
592 .await?;
593 let has_unique = rows.iter().any(|r| r.unique_key.is_some());
594
595 if active_engine != "canonical" {
596 return insert_compat_rows(conn, &rows, has_unique).await;
597 }
598
599 let target_table = homogeneous_target_table(&rows)
600 .map(TargetTable::as_str)
601 .unwrap_or("awa.jobs");
602
603 sqlx::query(
609 r#"
610 CREATE TEMP TABLE IF NOT EXISTS pg_temp.awa_copy_staging (
611 ord BIGINT GENERATED ALWAYS AS IDENTITY,
612 kind TEXT NOT NULL,
613 queue TEXT NOT NULL,
614 args JSONB NOT NULL,
615 state awa.job_state NOT NULL,
616 priority SMALLINT NOT NULL,
617 max_attempts SMALLINT NOT NULL,
618 run_at TIMESTAMPTZ,
619 metadata JSONB NOT NULL,
620 tags TEXT[] NOT NULL,
621 unique_key BYTEA,
622 unique_states BIT(8)
623 ) ON COMMIT DELETE ROWS
624 "#,
625 )
626 .execute(&mut *conn)
627 .await?;
628
629 let mut csv_buf = Vec::with_capacity(rows.len() * 256);
631 for row in &rows {
632 write_csv_row(&mut csv_buf, row);
633 }
634
635 let mut copy_in = conn
636 .copy_in_raw(
637 "COPY pg_temp.awa_copy_staging (kind, queue, args, state, priority, max_attempts, run_at, metadata, tags, unique_key, unique_states) FROM STDIN WITH (FORMAT csv, NULL '__AWA_NULL__')",
638 )
639 .await?;
640 copy_in.send(csv_buf).await?;
641 copy_in.finish().await?;
642
643 let results = if has_unique {
645 let query = build_copy_unique_insert_query(target_table);
646 sqlx::query_as::<_, JobRow>(&query)
647 .fetch_all(&mut *conn)
648 .await
649 .map_err(map_sqlx_error)?
650 } else {
651 let query = build_copy_insert_query(target_table);
652 sqlx::query_as::<_, JobRow>(&query)
653 .fetch_all(&mut *conn)
654 .await
655 .map_err(map_sqlx_error)?
656 };
657
658 sqlx::query("DELETE FROM pg_temp.awa_copy_staging")
661 .execute(&mut *conn)
662 .await?;
663
664 Ok(results)
665}
666
667#[tracing::instrument(skip(pool, jobs), fields(job.count = jobs.len()))]
672pub async fn insert_many_copy_from_pool(
673 pool: &PgPool,
674 jobs: &[InsertParams],
675) -> Result<Vec<JobRow>, AwaError> {
676 if jobs.is_empty() {
677 return Ok(Vec::new());
678 }
679
680 let mut tx = pool.begin().await?;
681 let results = insert_many_copy(&mut tx, jobs).await?;
682 tx.commit().await?;
683
684 Ok(results)
685}
686
687fn write_csv_row(buf: &mut Vec<u8>, row: &PreparedRow) {
691 write_csv_field(buf, &row.kind);
693 buf.push(b',');
694 write_csv_field(buf, &row.queue);
696 buf.push(b',');
697 let args_str = serde_json::to_string(&row.args).expect("JSON serialization should not fail");
699 write_csv_field(buf, &args_str);
700 buf.push(b',');
701 write_csv_field(buf, &row.state.to_string());
703 buf.push(b',');
704 buf.extend_from_slice(row.priority.to_string().as_bytes());
706 buf.push(b',');
707 buf.extend_from_slice(row.max_attempts.to_string().as_bytes());
709 buf.push(b',');
710 match &row.run_at {
712 Some(dt) => write_csv_field(buf, &dt.to_rfc3339()),
713 None => buf.extend_from_slice(COPY_NULL_SENTINEL.as_bytes()),
714 }
715 buf.push(b',');
716 let metadata_str =
718 serde_json::to_string(&row.metadata).expect("JSON serialization should not fail");
719 write_csv_field(buf, &metadata_str);
720 buf.push(b',');
721 write_pg_text_array(buf, &row.tags);
723 buf.push(b',');
724 match &row.unique_key {
726 Some(key) => {
727 let bytea_hex = format!("\\x{}", hex::encode(key));
728 write_csv_field(buf, &bytea_hex);
729 }
730 None => buf.extend_from_slice(COPY_NULL_SENTINEL.as_bytes()),
731 }
732 buf.push(b',');
733 match &row.unique_states {
735 Some(bits) => write_csv_field(buf, bits),
736 None => buf.extend_from_slice(COPY_NULL_SENTINEL.as_bytes()),
737 }
738 buf.push(b'\n');
739}
740
741fn write_csv_field(buf: &mut Vec<u8>, value: &str) {
743 if value.contains(',')
744 || value.contains('"')
745 || value.contains('\n')
746 || value.contains('\r')
747 || value.contains('\\')
748 || value == COPY_NULL_SENTINEL
749 {
750 buf.push(b'"');
751 for byte in value.bytes() {
752 if byte == b'"' {
753 buf.push(b'"');
754 }
755 buf.push(byte);
756 }
757 buf.push(b'"');
758 } else {
759 buf.extend_from_slice(value.as_bytes());
760 }
761}
762
763fn write_pg_text_array(buf: &mut Vec<u8>, values: &[String]) {
766 buf.push(b'"');
767 buf.push(b'{');
768 for (i, val) in values.iter().enumerate() {
769 if i > 0 {
770 buf.push(b',');
771 }
772 if val.is_empty()
773 || val.contains(',')
774 || val.contains('"')
775 || val.contains('\\')
776 || val.contains('{')
777 || val.contains('}')
778 || val.contains(' ')
779 || val.eq_ignore_ascii_case("NULL")
780 {
781 buf.push(b'"');
782 buf.push(b'"');
783 for ch in val.chars() {
784 match ch {
785 '"' => buf.extend_from_slice(b"\\\"\""),
786 '\\' => buf.extend_from_slice(b"\\\\"),
787 _ => {
788 let mut utf8_buf = [0u8; 4];
789 buf.extend_from_slice(ch.encode_utf8(&mut utf8_buf).as_bytes());
790 }
791 }
792 }
793 buf.push(b'"');
794 buf.push(b'"');
795 } else {
796 buf.extend_from_slice(val.as_bytes());
797 }
798 }
799 buf.push(b'}');
800 buf.push(b'"');
801}
802
803pub fn params(args: &impl JobArgs) -> Result<InsertParams, AwaError> {
805 params_with(args, InsertOpts::default())
806}
807
808pub fn params_with(args: &impl JobArgs, opts: InsertOpts) -> Result<InsertParams, AwaError> {
810 Ok(InsertParams {
811 kind: args.kind_str().to_string(),
812 args: args.to_args()?,
813 opts,
814 })
815}