1use std::collections::BTreeSet;
2use std::error::Error;
3use std::fmt;
4use std::panic::{AssertUnwindSafe, catch_unwind};
5use std::str::FromStr;
6use std::sync::{Arc, Mutex};
7use std::time::{Duration, SystemTime, UNIX_EPOCH};
8
9use serde_json::{Map, Value, json};
10use sha2::{Digest, Sha256};
11use sqlx::migrate::Migrator;
12use sqlx::pool::PoolConnection;
13use sqlx::postgres::{PgArguments, PgConnectOptions, PgPoolOptions, PgRow, PgSslMode};
14use sqlx::types::Json;
15use sqlx::{AssertSqlSafe, Connection, PgConnection, PgPool, Postgres, Row};
16
17use oxide_batch_repository::{
18 PartitionMutationError, aggregate_partition_parent, map_partition_aggregation,
19 recovered_execution,
20};
21
22use crate::{
23 ActorRef, BatchStatus, BoxFuture, BusinessStatement, BusinessTransaction,
24 BusinessTransactionError, BusinessValueKind, BusinessWriteResult, Checkpoint,
25 ChunkCommitReceipt, ChunkCounts, ChunkFaultProgress, ChunkTransaction, ChunkTransactionContext,
26 ChunkTransactionError, ChunkTransactionManager, ClassifierRevision, Clock, CursorError,
27 CursorKey, DefinitionDescriptor, DefinitionIdentity, DefinitionRevision, DefinitionUpgrade,
28 DurableStateKind, ExecutionContext, ExecutionCounts, ExecutionMetadata, ExecutionTimestamps,
29 ExecutionVersion, ExitCode, ExitStatus, ExplorerError, ExplorerQuery, ExplorerRepository,
30 FailureCategory, FailureId, FailureSummary, FaultPhase, FaultPolicy, FaultProgress,
31 FaultStateEntry, FaultStateEnvelope, FaultStateError, FaultStateFormatError, FaultStateStore,
32 FlowDecision, FlowDecisionId, FlowDecisionRequest, FlowDecisionSequence, FlowStepState,
33 FlowTarget, FlowTransitionKind, IdentifierKind, InheritedStepProgress, JobExecution,
34 JobExecutionId, JobExecutionProjection, JobInstance, JobInstanceId, JobInstanceKey,
35 JobInstanceProjection, JobInstanceSelection, JobName, JobParameter, JobParameters,
36 JobRepository, LifecycleError, LifecycleTransition, MAX_PARTITION_CONTEXT_BYTES,
37 MAX_PARTITIONS, NodeId, OperationId, OperatorAction, OperatorOutcomeClass, OperatorRecord,
38 OperatorRecordDraft, OperatorRejection, OperatorRequestId, ParameterDescriptor, ParameterName,
39 ParameterRole, ParameterValue, ParameterValueKind, PartitionKey, PartitionPlanEntry,
40 PartitionResult, PurgeBatchBound, PurgeCandidate, PurgeCounts, PurgePlan, PurgePlanRequest,
41 PurgeSurvey, QueryWindow, ReasonCode, RecoveryDecision, RecoveryDecisionId, RecoveryRequest,
42 RecoveryResult, RepositoryCapability, RepositoryDescriptor, RepositoryError,
43 RepositoryUnitOfWork, RequestDigest, RetentionAction, RetentionActionId, RetentionHold,
44 RetentionOutcome, RetentionRecord, RetentionRecordDraft, RetryCounts, RetryKey, RetryLimit,
45 RetryOrdinal, RetryReservation, RetryStateLimit, SkipCounts, StartLimit,
46 StateEnvelopeDescriptor, StateLimits, StateSchemaId, StateSchemaVersion, StepExecution,
47 StepExecutionId, StepExecutionProjection, StepName, StepPartition, StepPartitionId,
48 StepPartitionProjection, TerminalKind,
49};
50
51const SUPPORTED_SCHEMA_VERSION: u32 = 3;
52const MAX_INSTANCE_KEY_INPUT: usize = 1024 * 1024;
53const MAX_POOL_SIZE: u32 = 1024;
54const MAX_SHORT_TIMEOUT: Duration = Duration::from_mins(5);
55const MAX_STATEMENT_TIMEOUT: Duration = Duration::from_hours(24);
56const MAX_CONNECTION_LIFETIME: Duration = Duration::from_hours(7 * 24);
57const MAX_CA_CERTIFICATE_BYTES: usize = 1024 * 1024;
58const DEFAULT_CONTEXT_SCHEMA: &str = "oxide_batch.empty.v1";
59
60static MIGRATOR: Migrator = sqlx::migrate!("./migrations");
61
62#[derive(Clone, Eq, PartialEq)]
64#[non_exhaustive]
65pub enum TlsMode {
66 VerifyFull {
68 ca_certificate: Option<CaCertificate>,
70 },
71 Plaintext,
73}
74
75#[derive(Clone, Eq, PartialEq)]
77pub struct CaCertificate(Vec<u8>);
78
79impl CaCertificate {
80 pub fn new(pem: impl Into<Vec<u8>>) -> Result<Self, PostgresConfigError> {
87 let pem = pem.into();
88 if pem.is_empty() {
89 return Err(PostgresConfigError::EmptyCaCertificate);
90 }
91 if pem.len() > MAX_CA_CERTIFICATE_BYTES {
92 return Err(PostgresConfigError::CaCertificateTooLarge {
93 max_bytes: MAX_CA_CERTIFICATE_BYTES,
94 });
95 }
96 Ok(Self(pem))
97 }
98
99 fn as_bytes(&self) -> &[u8] {
100 &self.0
101 }
102}
103
104impl fmt::Debug for CaCertificate {
105 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
106 formatter
107 .debug_struct("CaCertificate")
108 .field("byte_length", &self.0.len())
109 .field("contents", &"<redacted>")
110 .finish()
111 }
112}
113
114impl Default for TlsMode {
115 fn default() -> Self {
116 Self::VerifyFull {
117 ca_certificate: None,
118 }
119 }
120}
121
122impl fmt::Debug for TlsMode {
123 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
124 match self {
125 Self::VerifyFull { ca_certificate } => formatter
126 .debug_struct("VerifyFull")
127 .field(
128 "ca_certificate",
129 &ca_certificate.as_ref().map(|_| "<redacted>"),
130 )
131 .finish(),
132 Self::Plaintext => formatter.write_str("Plaintext"),
133 }
134 }
135}
136
137#[derive(Clone)]
141pub struct PostgresConfig {
142 connection_string: String,
143 tls_mode: TlsMode,
144 pool_size: u32,
145 acquire_timeout: Duration,
146 connect_timeout: Duration,
147 statement_timeout: Duration,
148 lock_timeout: Duration,
149 idle_transaction_timeout: Duration,
150 connection_idle_timeout: Duration,
151 connection_max_lifetime: Duration,
152 pool_close_timeout: Duration,
153}
154
155impl PostgresConfig {
156 pub fn new(connection_string: impl Into<String>) -> Result<Self, PostgresConfigError> {
165 let config = Self {
166 connection_string: connection_string.into(),
167 tls_mode: TlsMode::default(),
168 pool_size: 10,
169 acquire_timeout: Duration::from_secs(30),
170 connect_timeout: Duration::from_secs(10),
171 statement_timeout: Duration::from_secs(30),
172 lock_timeout: Duration::from_secs(5),
173 idle_transaction_timeout: Duration::from_mins(1),
174 connection_idle_timeout: Duration::from_mins(10),
175 connection_max_lifetime: Duration::from_mins(30),
176 pool_close_timeout: Duration::from_secs(30),
177 };
178 config.validate()?;
179 Ok(config)
180 }
181
182 #[must_use]
184 pub fn with_tls_mode(mut self, tls_mode: TlsMode) -> Self {
185 self.tls_mode = tls_mode;
186 self
187 }
188
189 pub fn with_pool_size(mut self, value: u32) -> Result<Self, PostgresConfigError> {
195 self.pool_size = value;
196 self.validate()?;
197 Ok(self)
198 }
199
200 pub fn with_acquire_timeout(mut self, value: Duration) -> Result<Self, PostgresConfigError> {
206 self.acquire_timeout = value;
207 self.validate()?;
208 Ok(self)
209 }
210
211 pub fn with_connect_timeout(mut self, value: Duration) -> Result<Self, PostgresConfigError> {
217 self.connect_timeout = value;
218 self.validate()?;
219 Ok(self)
220 }
221
222 pub fn with_statement_timeout(mut self, value: Duration) -> Result<Self, PostgresConfigError> {
228 self.statement_timeout = value;
229 self.validate()?;
230 Ok(self)
231 }
232
233 pub fn with_lock_timeout(mut self, value: Duration) -> Result<Self, PostgresConfigError> {
240 self.lock_timeout = value;
241 self.validate()?;
242 Ok(self)
243 }
244
245 pub fn with_idle_transaction_timeout(
251 mut self,
252 value: Duration,
253 ) -> Result<Self, PostgresConfigError> {
254 self.idle_transaction_timeout = value;
255 self.validate()?;
256 Ok(self)
257 }
258
259 pub fn with_connection_idle_timeout(
265 mut self,
266 value: Duration,
267 ) -> Result<Self, PostgresConfigError> {
268 self.connection_idle_timeout = value;
269 self.validate()?;
270 Ok(self)
271 }
272
273 pub fn with_connection_max_lifetime(
279 mut self,
280 value: Duration,
281 ) -> Result<Self, PostgresConfigError> {
282 self.connection_max_lifetime = value;
283 self.validate()?;
284 Ok(self)
285 }
286
287 pub fn with_pool_close_timeout(mut self, value: Duration) -> Result<Self, PostgresConfigError> {
293 self.pool_close_timeout = value;
294 self.validate()?;
295 Ok(self)
296 }
297
298 fn validate(&self) -> Result<(), PostgresConfigError> {
299 if self.connection_string.is_empty() {
300 return Err(PostgresConfigError::EmptyConnectionString);
301 }
302 validate_connection_query(&self.connection_string)?;
303 if !(1..=MAX_POOL_SIZE).contains(&self.pool_size) {
304 return Err(PostgresConfigError::PoolSize);
305 }
306 validate_duration(
307 self.acquire_timeout,
308 Duration::from_millis(1),
309 MAX_SHORT_TIMEOUT,
310 "acquire",
311 )?;
312 validate_duration(
313 self.connect_timeout,
314 Duration::from_millis(1),
315 MAX_SHORT_TIMEOUT,
316 "connect",
317 )?;
318 validate_duration(
319 self.statement_timeout,
320 Duration::from_millis(1),
321 MAX_STATEMENT_TIMEOUT,
322 "statement",
323 )?;
324 validate_duration(
325 self.lock_timeout,
326 Duration::from_millis(1),
327 MAX_SHORT_TIMEOUT,
328 "lock",
329 )?;
330 validate_duration(
331 self.idle_transaction_timeout,
332 Duration::from_secs(1),
333 MAX_STATEMENT_TIMEOUT,
334 "idle transaction",
335 )?;
336 validate_duration(
337 self.connection_idle_timeout,
338 Duration::from_secs(1),
339 MAX_STATEMENT_TIMEOUT,
340 "connection idle",
341 )?;
342 validate_duration(
343 self.connection_max_lifetime,
344 Duration::from_mins(1),
345 MAX_CONNECTION_LIFETIME,
346 "connection lifetime",
347 )?;
348 validate_duration(
349 self.pool_close_timeout,
350 Duration::from_millis(1),
351 MAX_SHORT_TIMEOUT,
352 "pool close",
353 )?;
354 if self.lock_timeout > self.statement_timeout {
355 return Err(PostgresConfigError::LockExceedsStatement);
356 }
357 if self.acquire_timeout > self.pool_close_timeout {
358 return Err(PostgresConfigError::AcquireExceedsClose);
359 }
360 Ok(())
361 }
362
363 fn connect_options(&self) -> Result<PgConnectOptions, PostgresConfigError> {
364 let mut options = PgConnectOptions::from_str(&self.connection_string)
365 .map_err(|_| PostgresConfigError::InvalidConnectionString)?;
366 options = options.application_name("oxide-batch");
367 options = match &self.tls_mode {
368 TlsMode::VerifyFull { ca_certificate } => {
369 let options = options.ssl_mode(PgSslMode::VerifyFull);
370 if let Some(certificate) = ca_certificate {
371 options.ssl_root_cert_from_pem(certificate.as_bytes().to_vec())
372 } else {
373 options
374 }
375 }
376 TlsMode::Plaintext => options.ssl_mode(PgSslMode::Disable),
377 };
378 Ok(options)
379 }
380}
381
382impl fmt::Debug for PostgresConfig {
383 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
384 formatter
385 .debug_struct("PostgresConfig")
386 .field("connection", &"<redacted>")
387 .field(
388 "tls_mode",
389 &match self.tls_mode {
390 TlsMode::VerifyFull { .. } => "verify-full",
391 TlsMode::Plaintext => "plaintext",
392 },
393 )
394 .field("pool_size", &self.pool_size)
395 .field("acquire_timeout", &self.acquire_timeout)
396 .field("connect_timeout", &self.connect_timeout)
397 .field("statement_timeout", &self.statement_timeout)
398 .field("lock_timeout", &self.lock_timeout)
399 .field("idle_transaction_timeout", &self.idle_transaction_timeout)
400 .field("connection_idle_timeout", &self.connection_idle_timeout)
401 .field("connection_max_lifetime", &self.connection_max_lifetime)
402 .field("pool_close_timeout", &self.pool_close_timeout)
403 .finish_non_exhaustive()
404 }
405}
406
407#[derive(Clone, Debug, Eq, PartialEq)]
409#[non_exhaustive]
410pub enum PostgresConfigError {
411 EmptyConnectionString,
413 InvalidConnectionString,
415 PoolSize,
417 TlsOptionInConnectionString,
419 EmptyCaCertificate,
421 CaCertificateTooLarge {
423 max_bytes: usize,
425 },
426 Timeout {
428 class: &'static str,
430 },
431 LockExceedsStatement,
433 AcquireExceedsClose,
435}
436
437impl fmt::Display for PostgresConfigError {
438 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
439 match self {
440 Self::EmptyConnectionString => {
441 formatter.write_str("PostgreSQL connection string is empty")
442 }
443 Self::InvalidConnectionString => {
444 formatter.write_str("PostgreSQL connection string is invalid")
445 }
446 Self::PoolSize => formatter.write_str("PostgreSQL pool size must be from 1 to 1024"),
447 Self::TlsOptionInConnectionString => formatter
448 .write_str("PostgreSQL TLS options must use facade-owned TLS configuration"),
449 Self::EmptyCaCertificate => {
450 formatter.write_str("PostgreSQL CA certificate bundle is empty")
451 }
452 Self::CaCertificateTooLarge { max_bytes } => write!(
453 formatter,
454 "PostgreSQL CA certificate bundle exceeds {max_bytes} bytes"
455 ),
456 Self::Timeout { class } => {
457 write!(
458 formatter,
459 "PostgreSQL {class} timeout is outside its bounds"
460 )
461 }
462 Self::LockExceedsStatement => {
463 formatter.write_str("PostgreSQL lock timeout exceeds statement timeout")
464 }
465 Self::AcquireExceedsClose => {
466 formatter.write_str("PostgreSQL acquire timeout exceeds pool close timeout")
467 }
468 }
469 }
470}
471
472impl Error for PostgresConfigError {}
473
474fn validate_duration(
475 value: Duration,
476 minimum: Duration,
477 maximum: Duration,
478 class: &'static str,
479) -> Result<(), PostgresConfigError> {
480 if value < minimum || value > maximum {
481 return Err(PostgresConfigError::Timeout { class });
482 }
483 Ok(())
484}
485
486fn validate_connection_query(connection_string: &str) -> Result<(), PostgresConfigError> {
487 let Some((_, query)) = connection_string.split_once('?') else {
488 return Ok(());
489 };
490 for pair in query.split('&') {
491 let key = pair
492 .split_once('=')
493 .map_or(pair, |(key, _)| key)
494 .to_ascii_lowercase();
495 if matches!(
496 key.as_str(),
497 "sslmode"
498 | "ssl-mode"
499 | "sslrootcert"
500 | "ssl-root-cert"
501 | "ssl-ca"
502 | "sslcert"
503 | "ssl-cert"
504 | "sslkey"
505 | "ssl-key"
506 ) {
507 return Err(PostgresConfigError::TlsOptionInConnectionString);
508 }
509 let recognized = matches!(
510 key.as_str(),
511 "statement-cache-capacity"
512 | "host"
513 | "hostaddr"
514 | "port"
515 | "dbname"
516 | "user"
517 | "password"
518 | "application_name"
519 | "options"
520 ) || (key.starts_with("options[") && key.ends_with(']'));
521 if !recognized {
522 return Err(PostgresConfigError::InvalidConnectionString);
523 }
524 }
525 Ok(())
526}
527
528#[derive(Clone, Copy, Debug, Default)]
530pub struct PostgresMigrator;
531
532impl PostgresMigrator {
533 #[must_use]
535 pub const fn supported_schema_version() -> u32 {
536 SUPPORTED_SCHEMA_VERSION
537 }
538
539 pub async fn installed_schema_version(
552 config: &PostgresConfig,
553 ) -> Result<Option<u32>, RepositoryError> {
554 config
555 .validate()
556 .map_err(|_| RepositoryError::Unavailable)?;
557 let options = config
558 .connect_options()
559 .map_err(|_| RepositoryError::Unavailable)?;
560 let mut connection =
561 tokio::time::timeout(config.connect_timeout, PgConnection::connect_with(&options))
562 .await
563 .map_err(|_| RepositoryError::Unavailable)?
564 .map_err(|_| RepositoryError::Unavailable)?;
565 match read_schema_version(&mut connection).await {
566 Ok(version) => Ok(Some(version)),
567 Err(RepositoryError::SchemaUninitialized) => Ok(None),
568 Err(error) => Err(error),
569 }
570 }
571
572 pub async fn migrate(config: &PostgresConfig) -> Result<(), RepositoryError> {
582 config
583 .validate()
584 .map_err(|_| RepositoryError::Unavailable)?;
585 let options = config
586 .connect_options()
587 .map_err(|_| RepositoryError::Unavailable)?;
588 let mut connection =
589 tokio::time::timeout(config.connect_timeout, PgConnection::connect_with(&options))
590 .await
591 .map_err(|_| RepositoryError::Unavailable)?
592 .map_err(|_| RepositoryError::Unavailable)?;
593 let lock_timeout = duration_millis(config.lock_timeout)?;
594 sqlx::query("SELECT set_config('lock_timeout', $1, false)")
595 .bind(lock_timeout.to_string())
596 .execute(&mut connection)
597 .await
598 .map_err(|_| RepositoryError::Unavailable)?;
599 tokio::time::timeout(
600 config.lock_timeout,
601 sqlx::query(
602 "SELECT pg_advisory_lock(hashtextextended('oxide_batch.schema.migrations', 0))",
603 )
604 .execute(&mut connection),
605 )
606 .await
607 .map_err(|_| RepositoryError::Unavailable)?
608 .map_err(|_| RepositoryError::Unavailable)?;
609
610 let result = async {
611 match read_schema_version(&mut connection).await {
612 Ok(current) if current > SUPPORTED_SCHEMA_VERSION => {
613 return Err(RepositoryError::NewerSchema {
614 current,
615 supported: SUPPORTED_SCHEMA_VERSION,
616 });
617 }
618 Ok(_) => {}
619 Err(RepositoryError::SchemaUninitialized) => {
620 let schema_exists: bool = sqlx::query_scalar(
621 "SELECT EXISTS(\
622 SELECT 1 FROM pg_catalog.pg_namespace WHERE nspname = 'oxide_batch')",
623 )
624 .fetch_one(&mut connection)
625 .await
626 .map_err(|_| RepositoryError::Unavailable)?;
627 if !schema_exists {
628 sqlx::query("CREATE SCHEMA oxide_batch")
629 .execute(&mut connection)
630 .await
631 .map_err(|_| RepositoryError::Unavailable)?;
632 }
633 }
634 Err(error) => return Err(error),
635 }
636 sqlx::query("SET search_path TO oxide_batch, pg_catalog")
637 .execute(&mut connection)
638 .await
639 .map_err(|_| RepositoryError::Unavailable)?;
640 MIGRATOR
641 .run(&mut connection)
642 .await
643 .map_err(|_| RepositoryError::Unavailable)?;
644 let installed = read_schema_version(&mut connection).await?;
645 verify_schema_version(installed)
646 }
647 .await;
648
649 let _unlock = sqlx::query(
650 "SELECT pg_advisory_unlock(hashtextextended('oxide_batch.schema.migrations', 0))",
651 )
652 .execute(&mut connection)
653 .await;
654 result
655 }
656}
657
658#[derive(Clone)]
663pub struct PostgresJobRepository {
664 pool: PgPool,
665 clock: Arc<dyn Clock>,
666 config: PostgresConfig,
667}
668
669impl PostgresJobRepository {
670 pub async fn connect(
677 config: PostgresConfig,
678 clock: Arc<dyn Clock>,
679 ) -> Result<Self, RepositoryError> {
680 config
681 .validate()
682 .map_err(|_| RepositoryError::Unavailable)?;
683 let options = config
684 .connect_options()
685 .map_err(|_| RepositoryError::Unavailable)?;
686 let pool = tokio::time::timeout(
687 config.connect_timeout,
688 PgPoolOptions::new()
689 .max_connections(config.pool_size)
690 .acquire_timeout(config.acquire_timeout)
691 .idle_timeout(Some(config.connection_idle_timeout))
692 .max_lifetime(Some(config.connection_max_lifetime))
693 .connect_with(options),
694 )
695 .await
696 .map_err(|_| RepositoryError::Unavailable)?
697 .map_err(|_| RepositoryError::Unavailable)?;
698 let current = read_schema_version(&pool).await?;
699 verify_schema_version(current)?;
700 Ok(Self {
701 pool,
702 clock,
703 config,
704 })
705 }
706
707 pub async fn close(&self) -> Result<(), RepositoryError> {
714 tokio::time::timeout(self.config.pool_close_timeout, self.pool.close())
715 .await
716 .map_err(|_| RepositoryError::Unavailable)
717 }
718
719 async fn begin_connection(&self) -> Result<PoolConnection<Postgres>, RepositoryError> {
720 let mut connection = self
721 .pool
722 .acquire()
723 .await
724 .map_err(|_| RepositoryError::Unavailable)?;
725 if sqlx::query("BEGIN")
726 .execute(&mut *connection)
727 .await
728 .is_err()
729 {
730 connection.close_on_drop();
731 return Err(RepositoryError::Unavailable);
732 }
733 if let Err(error) = configure_transaction(&mut connection, &self.config).await {
734 connection.close_on_drop();
735 return Err(error);
736 }
737 Ok(connection)
738 }
739}
740
741impl fmt::Debug for PostgresJobRepository {
742 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
743 formatter
744 .debug_struct("PostgresJobRepository")
745 .field("pool_size", &self.pool.size())
746 .field("pool_idle", &self.pool.num_idle())
747 .finish_non_exhaustive()
748 }
749}
750
751impl JobRepository for PostgresJobRepository {
752 fn connection_capacity(&self) -> u32 {
753 self.config.pool_size
754 }
755
756 fn descriptor(&self) -> RepositoryDescriptor {
760 RepositoryDescriptor::new(
761 SUPPORTED_SCHEMA_VERSION,
762 [
763 RepositoryCapability::ExecutionOwnership,
764 RepositoryCapability::InstanceHolds,
765 RepositoryCapability::OperatorRequests,
766 RepositoryCapability::RetentionPurge,
767 RepositoryCapability::StepPartitions,
768 RepositoryCapability::StopRequests,
769 ],
770 )
771 }
772
773 fn begin<'a>(
774 &'a self,
775 ) -> BoxFuture<'a, Result<Box<dyn RepositoryUnitOfWork + 'a>, RepositoryError>> {
776 Box::pin(async move {
777 let connection = self.begin_connection().await?;
778 Ok(Box::new(PostgresUnitOfWork {
779 repository: self,
780 connection: Some(connection),
781 definition_override: None,
782 created_partition_plans: BTreeSet::new(),
783 }) as Box<dyn RepositoryUnitOfWork + 'a>)
784 })
785 }
786}
787
788#[derive(Clone, Debug, Eq, PartialEq)]
790pub struct PostgresDurableStepState {
791 step_execution: StepExecution,
792 checkpoint: Checkpoint,
793 execution_context: ExecutionContext,
794 fault_progress: FaultProgress,
795 fault_state: FaultStateEnvelope,
796}
797
798impl PostgresDurableStepState {
799 #[must_use]
801 pub const fn step_execution(&self) -> &StepExecution {
802 &self.step_execution
803 }
804
805 #[must_use]
807 pub const fn checkpoint(&self) -> &Checkpoint {
808 &self.checkpoint
809 }
810
811 #[must_use]
813 pub const fn execution_context(&self) -> &ExecutionContext {
814 &self.execution_context
815 }
816
817 #[must_use]
819 pub const fn fault_progress(&self) -> FaultProgress {
820 self.fault_progress
821 }
822
823 #[must_use]
825 pub const fn fault_state(&self) -> &FaultStateEnvelope {
826 &self.fault_state
827 }
828}
829
830pub struct PostgresFaultState {
842 repository: PostgresJobRepository,
843 revision: ClassifierRevision,
844 retry_limit: RetryLimit,
845 state_limit: RetryStateLimit,
846 bound: Mutex<Option<ChunkTransactionContext>>,
847}
848
849impl PostgresFaultState {
850 #[must_use]
854 pub fn new(repository: PostgresJobRepository, policy: &FaultPolicy) -> Self {
855 Self {
856 repository,
857 revision: policy.classifier().revision().clone(),
858 retry_limit: policy.retry_limit(),
859 state_limit: policy.retry_state_limit(),
860 bound: Mutex::new(None),
861 }
862 }
863
864 fn context(&self) -> Result<ChunkTransactionContext, FaultStateError> {
865 (*self
866 .bound
867 .lock()
868 .unwrap_or_else(std::sync::PoisonError::into_inner))
869 .ok_or(FaultStateError::Unbound)
870 }
871
872 async fn load<'executor, E>(
873 &self,
874 executor: E,
875 lock: bool,
876 ) -> Result<PostgresFaultRow, FaultStateError>
877 where
878 E: sqlx::Executor<'executor, Database = Postgres>,
879 {
880 let context = self.context()?;
881 let query = format!(
882 "SELECT execution.version, execution.status, \
883 execution.checkpoint_format, execution.checkpoint_schema, \
884 execution.checkpoint_schema_version, execution.checkpoint_payload, \
885 execution.fault_state_format, execution.fault_state_schema, \
886 execution.fault_state_schema_version, execution.fault_state_payload, \
887 execution.fault_state_checksum \
888 FROM oxide_batch.ob_step_execution execution \
889 WHERE execution.id = $1 AND execution.job_execution_id = $2{}",
890 if lock { " FOR UPDATE" } else { "" }
891 );
892 let row = sqlx::query(AssertSqlSafe(query))
893 .bind(
894 database_id(
895 context.step_execution_id().get(),
896 IdentifierKind::StepExecution,
897 )
898 .map_err(|_| FaultStateError::Unavailable)?,
899 )
900 .bind(
901 database_id(
902 context.job_execution_id().get(),
903 IdentifierKind::JobExecution,
904 )
905 .map_err(|_| FaultStateError::Unavailable)?,
906 )
907 .fetch_optional(executor)
908 .await
909 .map_err(|_| FaultStateError::Unavailable)?
910 .ok_or(FaultStateError::Unavailable)?;
911 let checkpoint: Checkpoint = decode_durable_state(
912 &row,
913 "checkpoint_format",
914 "checkpoint_schema",
915 "checkpoint_schema_version",
916 "checkpoint_payload",
917 "oxide-batch.checkpoint",
918 Checkpoint::from_json,
919 )
920 .map_err(|_| FaultStateError::Unavailable)?;
921 let checkpoint_digest = checkpoint.generation_digest();
922 let envelope = decode_fault_state(&row).map_err(FaultStateError::Corrupt)?;
923 envelope.validate_for(self.retry_limit, self.state_limit, &checkpoint_digest)?;
924 Ok(PostgresFaultRow {
925 version: ExecutionVersion::new(
926 read_u64(&row, "version").map_err(|_| FaultStateError::Unavailable)?,
927 ),
928 started: row
929 .try_get::<String, _>("status")
930 .map_err(|_| FaultStateError::Unavailable)?
931 == BatchStatus::Started.to_string(),
932 checkpoint_digest,
933 envelope,
934 })
935 }
936}
937
938struct PostgresFaultRow {
939 version: ExecutionVersion,
940 started: bool,
941 checkpoint_digest: [u8; 32],
942 envelope: FaultStateEnvelope,
943}
944
945impl fmt::Debug for PostgresFaultState {
946 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
947 formatter
948 .debug_struct("PostgresFaultState")
949 .field("retry_limit", &self.retry_limit)
950 .field("retry_state_limit", &self.state_limit)
951 .finish_non_exhaustive()
952 }
953}
954
955impl FaultStateStore for PostgresFaultState {
956 fn bind(&self, context: ChunkTransactionContext) -> BoxFuture<'_, Result<(), FaultStateError>> {
957 Box::pin(async move {
958 *self
959 .bound
960 .lock()
961 .unwrap_or_else(std::sync::PoisonError::into_inner) = Some(context);
962 self.load(&self.repository.pool, false).await.map(|_| ())
963 })
964 }
965
966 fn reserved_ordinal(
967 &self,
968 key: RetryKey,
969 ) -> BoxFuture<'_, Result<Option<RetryOrdinal>, FaultStateError>> {
970 Box::pin(async move {
971 let row = self.load(&self.repository.pool, false).await?;
972 Ok(row.envelope.reserved_ordinal(key))
973 })
974 }
975
976 fn reserve(&self, reservation: RetryReservation) -> BoxFuture<'_, Result<(), FaultStateError>> {
977 Box::pin(async move {
978 let context = self.context()?;
979 let mut connection = self
980 .repository
981 .begin_connection()
982 .await
983 .map_err(|_| FaultStateError::Unavailable)?;
984 let result = self
985 .reserve_locked(&mut connection, context, reservation)
986 .await;
987 match result {
988 Ok(()) => commit_postgres_connection(connection)
989 .await
990 .map_err(|()| FaultStateError::Unavailable),
991 Err(error) => {
992 rollback_chunk_connection(&mut connection).await;
993 Err(error)
994 }
995 }
996 })
997 }
998
999 fn resolve(&self, _key: RetryKey) -> BoxFuture<'_, Result<(), FaultStateError>> {
1000 Box::pin(std::future::ready(Ok(())))
1001 }
1002
1003 fn clear_resolved(&self) -> BoxFuture<'_, Result<(), FaultStateError>> {
1004 Box::pin(std::future::ready(Ok(())))
1005 }
1006
1007 fn unresolved(&self) -> BoxFuture<'_, Result<u32, FaultStateError>> {
1008 Box::pin(async move {
1009 let row = self.load(&self.repository.pool, false).await?;
1010 u32::try_from(row.envelope.len()).map_err(|_| FaultStateError::Unavailable)
1011 })
1012 }
1013}
1014
1015impl PostgresFaultState {
1016 async fn reserve_locked(
1017 &self,
1018 connection: &mut PoolConnection<Postgres>,
1019 context: ChunkTransactionContext,
1020 reservation: RetryReservation,
1021 ) -> Result<(), FaultStateError> {
1022 let row = self.load(&mut **connection, true).await?;
1023 if !row.started {
1024 return Err(FaultStateError::StaleReservation);
1025 }
1026 let entry = FaultStateEntry::new(
1027 reservation.key(),
1028 reservation.phase(),
1029 reservation.category(),
1030 reservation.ordinal(),
1031 self.revision.clone(),
1032 );
1033 let next = row
1034 .envelope
1035 .reserved(entry, row.checkpoint_digest, self.state_limit)?;
1036 let payload: Value = serde_json::from_slice(&next.to_canonical_json()?)
1037 .map_err(|_| FaultStateError::Unavailable)?;
1038 let checksum = next.checksum()?;
1039 let retry_column = match reservation.phase() {
1040 FaultPhase::Read => "read_retry_count",
1041 FaultPhase::Process => "process_retry_count",
1042 FaultPhase::Write => "write_retry_count",
1043 _ => return Err(FaultStateError::StaleReservation),
1044 };
1045 let next_version = row
1046 .version
1047 .next()
1048 .map_err(|_| FaultStateError::Unavailable)?;
1049 let updated = sqlx::query(AssertSqlSafe(format!(
1050 "UPDATE oxide_batch.ob_step_execution SET \
1051 {retry_column} = {retry_column} + 1, rollback_count = rollback_count + 1, \
1052 fault_state_format = $1, fault_state_schema = $2, \
1053 fault_state_schema_version = $3, fault_state_payload = $4, \
1054 fault_state_checksum = $5, \
1055 updated_at = to_timestamp($6::double precision / 1000.0), version = $7 \
1056 WHERE id = $8 AND job_execution_id = $9 AND version = $10 AND status = 'STARTED'"
1057 )))
1058 .bind(
1059 i16::try_from(FaultStateEnvelope::FORMAT_VERSION)
1060 .map_err(|_| FaultStateError::Unavailable)?,
1061 )
1062 .bind(FaultStateEnvelope::FORMAT)
1063 .bind(
1064 i32::try_from(FaultStateEnvelope::SCHEMA_VERSION)
1065 .map_err(|_| FaultStateError::Unavailable)?,
1066 )
1067 .bind(Json(payload))
1068 .bind(checksum.as_slice())
1069 .bind(
1070 system_time_millis(self.repository.clock.now())
1071 .map_err(|_| FaultStateError::Unavailable)?,
1072 )
1073 .bind(database_version(next_version).map_err(|_| FaultStateError::Unavailable)?)
1074 .bind(
1075 database_id(
1076 context.step_execution_id().get(),
1077 IdentifierKind::StepExecution,
1078 )
1079 .map_err(|_| FaultStateError::Unavailable)?,
1080 )
1081 .bind(
1082 database_id(
1083 context.job_execution_id().get(),
1084 IdentifierKind::JobExecution,
1085 )
1086 .map_err(|_| FaultStateError::Unavailable)?,
1087 )
1088 .bind(database_version(row.version).map_err(|_| FaultStateError::Unavailable)?)
1089 .execute(&mut **connection)
1090 .await
1091 .map_err(|_| FaultStateError::Unavailable)?;
1092 if updated.rows_affected() == 1 {
1093 Ok(())
1094 } else {
1095 Err(FaultStateError::StaleReservation)
1096 }
1097 }
1098}
1099
1100#[derive(Clone)]
1106pub struct PostgresChunkTransactionManager {
1107 repository: PostgresJobRepository,
1108 state_provider: Arc<dyn PostgresChunkStateProvider>,
1109}
1110
1111impl PostgresChunkTransactionManager {
1112 #[must_use]
1114 pub const fn new(
1115 repository: PostgresJobRepository,
1116 state_provider: Arc<dyn PostgresChunkStateProvider>,
1117 ) -> Self {
1118 Self {
1119 repository,
1120 state_provider,
1121 }
1122 }
1123
1124 pub async fn load_committed_state(
1131 &self,
1132 context: ChunkTransactionContext,
1133 ) -> Result<PostgresDurableStepState, RepositoryError> {
1134 let row = sqlx::query(AssertSqlSafe(durable_step_select(
1135 "WHERE execution.id = $1 AND execution.job_execution_id = $2",
1136 )))
1137 .bind(database_id(
1138 context.step_execution_id().get(),
1139 IdentifierKind::StepExecution,
1140 )?)
1141 .bind(database_id(
1142 context.job_execution_id().get(),
1143 IdentifierKind::JobExecution,
1144 )?)
1145 .fetch_optional(&self.repository.pool)
1146 .await
1147 .map_err(|_| RepositoryError::Unavailable)?
1148 .ok_or(RepositoryError::StepExecutionNotFound {
1149 id: context.step_execution_id(),
1150 })?;
1151 decode_durable_step_state(&row)
1152 }
1153}
1154
1155pub trait PostgresChunkStateProvider: Send + Sync {
1161 fn state_for_commit(
1168 &self,
1169 committed_counts: ExecutionCounts,
1170 chunk_counts: ChunkCounts,
1171 ) -> Result<ChunkCommitReceipt, PostgresChunkStateError>;
1172}
1173
1174impl<F> PostgresChunkStateProvider for F
1175where
1176 F: Fn(ExecutionCounts, ChunkCounts) -> Result<ChunkCommitReceipt, PostgresChunkStateError>
1177 + Send
1178 + Sync,
1179{
1180 fn state_for_commit(
1181 &self,
1182 committed_counts: ExecutionCounts,
1183 chunk_counts: ChunkCounts,
1184 ) -> Result<ChunkCommitReceipt, PostgresChunkStateError> {
1185 self(committed_counts, chunk_counts)
1186 }
1187}
1188
1189#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
1191pub struct PostgresChunkStateError;
1192
1193impl PostgresChunkStateError {
1194 #[must_use]
1196 pub const fn new() -> Self {
1197 Self
1198 }
1199}
1200
1201impl fmt::Display for PostgresChunkStateError {
1202 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
1203 formatter.write_str("PostgreSQL chunk state preparation failed")
1204 }
1205}
1206
1207impl Error for PostgresChunkStateError {}
1208
1209impl fmt::Debug for PostgresChunkTransactionManager {
1210 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
1211 formatter
1212 .debug_struct("PostgresChunkTransactionManager")
1213 .field("repository", &self.repository)
1214 .field("durable_state", &"<redacted>")
1215 .finish_non_exhaustive()
1216 }
1217}
1218
1219impl ChunkTransactionManager for PostgresChunkTransactionManager {
1220 fn begin(
1221 &self,
1222 ) -> BoxFuture<'_, Result<Box<dyn ChunkTransaction + '_>, ChunkTransactionError>> {
1223 Box::pin(async { Err(ChunkTransactionError::NotCommitted) })
1224 }
1225
1226 fn begin_for(
1227 &self,
1228 context: ChunkTransactionContext,
1229 ) -> BoxFuture<'_, Result<Box<dyn ChunkTransaction + '_>, ChunkTransactionError>> {
1230 Box::pin(async move {
1231 let step_id = database_id(
1232 context.step_execution_id().get(),
1233 IdentifierKind::StepExecution,
1234 )
1235 .map_err(|_| ChunkTransactionError::NotCommitted)?;
1236 let job_id = database_id(
1237 context.job_execution_id().get(),
1238 IdentifierKind::JobExecution,
1239 )
1240 .map_err(|_| ChunkTransactionError::NotCommitted)?;
1241 let mut connection = self
1242 .repository
1243 .begin_connection()
1244 .await
1245 .map_err(|_| ChunkTransactionError::NotCommitted)?;
1246 let row = sqlx::query(AssertSqlSafe(durable_step_select(
1247 "WHERE execution.id = $1 AND execution.job_execution_id = $2",
1248 )))
1249 .bind(step_id)
1250 .bind(job_id)
1251 .fetch_optional(&mut *connection)
1252 .await;
1253 let Ok(row) = row else {
1254 rollback_chunk_connection(&mut connection).await;
1255 return Err(ChunkTransactionError::NotCommitted);
1256 };
1257 let Some(row) = row else {
1258 rollback_chunk_connection(&mut connection).await;
1259 return Err(ChunkTransactionError::NotCommitted);
1260 };
1261 let Ok(durable) = decode_durable_step_state(&row) else {
1262 rollback_chunk_connection(&mut connection).await;
1263 return Err(ChunkTransactionError::NotCommitted);
1264 };
1265 if durable.step_execution.metadata().status() != BatchStatus::Started {
1266 rollback_chunk_connection(&mut connection).await;
1267 return Err(ChunkTransactionError::NotCommitted);
1268 }
1269 Ok(Box::new(PostgresChunkTransaction {
1270 connection: Some(connection),
1271 context,
1272 expected_version: durable.step_execution.version(),
1273 committed_counts: durable.step_execution.metadata().counts(),
1274 clock: Arc::clone(&self.repository.clock),
1275 state_provider: Arc::clone(&self.state_provider),
1276 }) as Box<dyn ChunkTransaction>)
1277 })
1278 }
1279
1280 fn inherited_progress(
1281 &self,
1282 context: ChunkTransactionContext,
1283 ) -> BoxFuture<'_, Result<InheritedStepProgress, ChunkTransactionError>> {
1284 Box::pin(async move {
1285 let durable = self
1286 .load_committed_state(context)
1287 .await
1288 .map_err(|_| ChunkTransactionError::NotCommitted)?;
1289 let digest = durable.checkpoint.generation_digest();
1290 if !durable.fault_state.is_empty() && durable.fault_state.checkpoint_digest() != &digest
1291 {
1292 return Err(ChunkTransactionError::NotCommitted);
1293 }
1294 Ok(InheritedStepProgress::new(
1295 durable.step_execution.metadata().counts().read(),
1296 digest,
1297 durable.fault_progress,
1298 ))
1299 })
1300 }
1301}
1302
1303struct PostgresUnitOfWork<'repository> {
1304 repository: &'repository PostgresJobRepository,
1305 connection: Option<PoolConnection<Postgres>>,
1306 definition_override: Option<DefinitionIdentity>,
1307 created_partition_plans: BTreeSet<StepExecutionId>,
1308}
1309
1310impl PostgresUnitOfWork<'_> {
1311 fn transaction(&mut self) -> Result<&mut PoolConnection<Postgres>, RepositoryError> {
1312 self.connection.as_mut().ok_or(RepositoryError::Unavailable)
1313 }
1314
1315 async fn job_execution(
1316 &mut self,
1317 id: JobExecutionId,
1318 ) -> Result<Option<JobExecution>, RepositoryError> {
1319 let id = database_id(id.get(), IdentifierKind::JobExecution)?;
1320 let row = sqlx::query(AssertSqlSafe(job_execution_select(
1321 "WHERE execution.id = $1",
1322 )))
1323 .bind(id)
1324 .fetch_optional(&mut **self.transaction()?)
1325 .await
1326 .map_err(|_| RepositoryError::Unavailable)?;
1327 row.map(|row| decode_job_execution(&row)).transpose()
1328 }
1329
1330 async fn step_execution(
1331 &mut self,
1332 id: StepExecutionId,
1333 ) -> Result<Option<StepExecution>, RepositoryError> {
1334 let id = database_id(id.get(), IdentifierKind::StepExecution)?;
1335 let row = sqlx::query(AssertSqlSafe(step_execution_select(
1336 "WHERE execution.id = $1",
1337 )))
1338 .bind(id)
1339 .fetch_optional(&mut **self.transaction()?)
1340 .await
1341 .map_err(|_| RepositoryError::Unavailable)?;
1342 row.map(|row| decode_step_execution(&row)).transpose()
1343 }
1344
1345 async fn classify_job_cas(
1346 &mut self,
1347 id: JobExecutionId,
1348 expected: ExecutionVersion,
1349 ) -> RepositoryError {
1350 match self.job_execution(id).await {
1351 Ok(Some(actual)) => RepositoryError::Lifecycle(LifecycleError::StaleVersion {
1352 expected,
1353 actual: actual.version(),
1354 }),
1355 Ok(None) => RepositoryError::JobExecutionNotFound { id },
1356 Err(error) => error,
1357 }
1358 }
1359
1360 async fn classify_step_cas(
1361 &mut self,
1362 id: StepExecutionId,
1363 expected: ExecutionVersion,
1364 ) -> RepositoryError {
1365 match self.step_execution(id).await {
1366 Ok(Some(actual)) => RepositoryError::Lifecycle(LifecycleError::StaleVersion {
1367 expected,
1368 actual: actual.version(),
1369 }),
1370 Ok(None) => RepositoryError::StepExecutionNotFound { id },
1371 Err(error) => error,
1372 }
1373 }
1374
1375 async fn purge_delete(
1376 &mut self,
1377 statement: &'static str,
1378 ids: &[i64],
1379 ) -> Result<u64, RepositoryError> {
1380 if ids.is_empty() {
1381 return Ok(0);
1382 }
1383 Ok(sqlx::query(statement)
1384 .bind(ids)
1385 .execute(&mut **self.transaction()?)
1386 .await
1387 .map_err(|_| RepositoryError::Unavailable)?
1388 .rows_affected())
1389 }
1390
1391 async fn purge_count(
1392 &mut self,
1393 statement: &'static str,
1394 ids: &[i64],
1395 ) -> Result<u64, RepositoryError> {
1396 if ids.is_empty() {
1397 return Ok(0);
1398 }
1399 let row = sqlx::query(statement)
1400 .bind(ids)
1401 .fetch_one(&mut **self.transaction()?)
1402 .await
1403 .map_err(|_| RepositoryError::Unavailable)?;
1404 read_u64(&row, "matched")
1405 }
1406
1407 async fn purge_counts(
1408 &mut self,
1409 candidates: &[PurgeCandidate],
1410 ) -> Result<PurgeCounts, RepositoryError> {
1411 let executions = candidate_execution_ids(candidates)?;
1412 let instances = candidate_instance_ids(candidates)?;
1413 let flow_decisions = self
1414 .purge_count(
1415 "SELECT count(*) AS matched FROM oxide_batch.ob_flow_decision \
1416 WHERE job_execution_id = ANY($1)",
1417 &executions,
1418 )
1419 .await?;
1420 let recovery_decisions = self
1421 .purge_count(
1422 "SELECT count(*) AS matched FROM oxide_batch.ob_recovery_decision \
1423 WHERE job_execution_id = ANY($1)",
1424 &executions,
1425 )
1426 .await?;
1427 let operator_requests = self
1428 .purge_count(
1429 "SELECT count(*) AS matched FROM oxide_batch.ob_operator_request \
1430 WHERE job_execution_id = ANY($1)",
1431 &executions,
1432 )
1433 .await?;
1434 let step_partitions = self
1435 .purge_count(
1436 "SELECT count(*) AS matched FROM oxide_batch.ob_step_partition \
1437 WHERE step_execution_id IN ( \
1438 SELECT id FROM oxide_batch.ob_step_execution \
1439 WHERE job_execution_id = ANY($1))",
1440 &executions,
1441 )
1442 .await?;
1443 let step_executions = self
1444 .purge_count(
1445 "SELECT count(*) AS matched FROM oxide_batch.ob_step_execution \
1446 WHERE job_execution_id = ANY($1)",
1447 &executions,
1448 )
1449 .await?;
1450 let job_instances = if instances.is_empty() {
1451 0
1452 } else {
1453 let row = sqlx::query(
1454 "SELECT count(*) AS matched FROM oxide_batch.ob_job_instance instance \
1455 WHERE instance.id = ANY($2) AND NOT EXISTS ( \
1456 SELECT 1 FROM oxide_batch.ob_job_execution execution \
1457 WHERE execution.job_instance_id = instance.id \
1458 AND execution.id <> ALL($1))",
1459 )
1460 .bind(&executions)
1461 .bind(&instances)
1462 .fetch_one(&mut **self.transaction()?)
1463 .await
1464 .map_err(|_| RepositoryError::Unavailable)?;
1465 read_u64(&row, "matched")?
1466 };
1467 Ok(PurgeCounts::new(
1468 flow_decisions,
1469 recovery_decisions,
1470 operator_requests,
1471 step_partitions,
1472 step_executions,
1473 u64::try_from(candidates.len()).unwrap_or(u64::MAX),
1474 job_instances,
1475 ))
1476 }
1477}
1478
1479impl RepositoryUnitOfWork for PostgresUnitOfWork<'_> {
1480 fn register_definition_upgrade<'a>(
1481 &'a mut self,
1482 job_name: &'a JobName,
1483 upgrade: &'a DefinitionUpgrade,
1484 ) -> BoxFuture<'a, Result<(), RepositoryError>> {
1485 Box::pin(async move {
1486 let registered_at = self.repository.clock.now();
1487 let from_id = ensure_definition(
1488 &mut **self.transaction()?,
1489 job_name.as_str(),
1490 upgrade.from(),
1491 registered_at,
1492 )
1493 .await?;
1494 let to_id = ensure_definition(
1495 &mut **self.transaction()?,
1496 job_name.as_str(),
1497 upgrade.to(),
1498 registered_at,
1499 )
1500 .await?;
1501 let mapping = upgrade
1502 .step_mapping()
1503 .iter()
1504 .map(|(source, target)| {
1505 (
1506 source.as_str().to_owned(),
1507 Value::String(target.as_str().to_owned()),
1508 )
1509 })
1510 .collect::<Map<String, Value>>();
1511 let registered_ms = system_time_millis(registered_at)?;
1512 sqlx::query(
1513 "INSERT INTO oxide_batch.ob_definition_upgrade \
1514 (from_definition_id, to_definition_id, upgrade_key, step_mapping, registered_at) \
1515 VALUES ($1, $2, $3, $4, to_timestamp($5::double precision / 1000.0)) \
1516 ON CONFLICT (from_definition_id, to_definition_id) DO NOTHING",
1517 )
1518 .bind(from_id)
1519 .bind(to_id)
1520 .bind(upgrade.key().as_str())
1521 .bind(Json(Value::Object(mapping.clone())))
1522 .bind(registered_ms)
1523 .execute(&mut **self.transaction()?)
1524 .await
1525 .map_err(|_| RepositoryError::Unavailable)?;
1526 let matches: bool = sqlx::query_scalar(
1527 "SELECT EXISTS(SELECT 1 FROM oxide_batch.ob_definition_upgrade \
1528 WHERE from_definition_id = $1 AND to_definition_id = $2 \
1529 AND upgrade_key = $3 AND step_mapping = $4)",
1530 )
1531 .bind(from_id)
1532 .bind(to_id)
1533 .bind(upgrade.key().as_str())
1534 .bind(Json(Value::Object(mapping)))
1535 .fetch_one(&mut **self.transaction()?)
1536 .await
1537 .map_err(|_| RepositoryError::Unavailable)?;
1538 if !matches {
1539 return Err(RepositoryError::DefinitionUpgradeConflict {
1540 job_name: job_name.clone(),
1541 });
1542 }
1543 Ok(())
1544 })
1545 }
1546
1547 fn select_or_create_job_instance<'a>(
1548 &'a mut self,
1549 key: &'a JobInstanceKey,
1550 ) -> BoxFuture<'a, Result<JobInstanceSelection, RepositoryError>> {
1551 Box::pin(async move {
1552 let encoded = encode_identifying_parameters(key)?;
1553 let instance_key = key.digest();
1554 let created_ms = system_time_millis(self.repository.clock.now())?;
1555 let inserted = sqlx::query(
1556 "INSERT INTO oxide_batch.ob_job_instance \
1557 (job_name, instance_key, identifying_parameters, created_at) \
1558 VALUES ($1, $2, $3, to_timestamp($4::double precision / 1000.0)) \
1559 ON CONFLICT (job_name, instance_key) DO NOTHING \
1560 RETURNING id, job_name, identifying_parameters",
1561 )
1562 .bind(key.job_name().as_str())
1563 .bind(&instance_key[..])
1564 .bind(Json(encoded))
1565 .bind(created_ms)
1566 .fetch_optional(&mut **self.transaction()?)
1567 .await
1568 .map_err(|_| RepositoryError::Unavailable)?;
1569
1570 if let Some(row) = inserted {
1571 return Ok(JobInstanceSelection::Created(decode_job_instance(&row)?));
1572 }
1573
1574 let row = sqlx::query(
1575 "SELECT id, job_name, identifying_parameters \
1576 FROM oxide_batch.ob_job_instance \
1577 WHERE job_name = $1 AND instance_key = $2",
1578 )
1579 .bind(key.job_name().as_str())
1580 .bind(&instance_key[..])
1581 .fetch_one(&mut **self.transaction()?)
1582 .await
1583 .map_err(|_| RepositoryError::Unavailable)?;
1584 Ok(JobInstanceSelection::Existing(decode_job_instance(&row)?))
1585 })
1586 }
1587
1588 #[allow(clippy::too_many_lines)]
1589 fn create_job_execution(
1590 &mut self,
1591 job_instance_id: JobInstanceId,
1592 ) -> BoxFuture<'_, Result<JobExecution, RepositoryError>> {
1593 Box::pin(async move {
1594 let definition = self
1595 .definition_override
1596 .take()
1597 .unwrap_or_else(DefinitionIdentity::legacy);
1598 let instance_database_id =
1599 database_id(job_instance_id.get(), IdentifierKind::JobInstance)?;
1600 let instance = sqlx::query(
1601 "SELECT job_name, identifying_parameters \
1602 FROM oxide_batch.ob_job_instance WHERE id = $1 FOR UPDATE",
1603 )
1604 .bind(instance_database_id)
1605 .fetch_optional(&mut **self.transaction()?)
1606 .await
1607 .map_err(|_| RepositoryError::Unavailable)?
1608 .ok_or(RepositoryError::JobInstanceNotFound {
1609 id: job_instance_id,
1610 })?;
1611
1612 let latest = sqlx::query(AssertSqlSafe(job_execution_select(
1613 "WHERE execution.job_instance_id = $1 ORDER BY execution.attempt DESC LIMIT 1",
1614 )))
1615 .bind(instance_database_id)
1616 .fetch_optional(&mut **self.transaction()?)
1617 .await
1618 .map_err(|_| RepositoryError::Unavailable)?
1619 .map(|row| decode_job_execution(&row))
1620 .transpose()?;
1621
1622 if let Some(execution) = &latest {
1623 match execution.metadata().status() {
1624 BatchStatus::Stopped | BatchStatus::Failed => {}
1625 BatchStatus::Completed => {
1626 return Err(RepositoryError::CompletedInstance {
1627 id: job_instance_id,
1628 });
1629 }
1630 BatchStatus::Abandoned => {
1631 return Err(RepositoryError::AbandonedInstance {
1632 id: job_instance_id,
1633 });
1634 }
1635 status => {
1636 return Err(RepositoryError::ExecutionAlreadyActive {
1637 instance_id: job_instance_id,
1638 execution_id: execution.id(),
1639 status,
1640 });
1641 }
1642 }
1643 }
1644
1645 let job_name: String = instance
1646 .try_get("job_name")
1647 .map_err(|_| RepositoryError::Unavailable)?;
1648 let parameters: Json<Value> = instance
1649 .try_get("identifying_parameters")
1650 .map_err(|_| RepositoryError::Unavailable)?;
1651 let registered_at = self.repository.clock.now();
1652 let definition_id = ensure_definition(
1653 &mut **self.transaction()?,
1654 &job_name,
1655 &definition,
1656 registered_at,
1657 )
1658 .await?;
1659 let mut upgrade_from = None;
1660 if let Some(previous) = &latest {
1661 let previous_definition_id: i64 = sqlx::query_scalar(
1662 "SELECT definition_id FROM oxide_batch.ob_job_execution WHERE id = $1",
1663 )
1664 .bind(database_id(
1665 previous.id().get(),
1666 IdentifierKind::JobExecution,
1667 )?)
1668 .fetch_one(&mut **self.transaction()?)
1669 .await
1670 .map_err(|_| RepositoryError::Unavailable)?;
1671 if previous_definition_id != definition_id {
1672 let compatible: bool = sqlx::query_scalar(
1673 "SELECT EXISTS(SELECT 1 FROM oxide_batch.ob_definition_upgrade \
1674 WHERE from_definition_id = $1 AND to_definition_id = $2)",
1675 )
1676 .bind(previous_definition_id)
1677 .bind(definition_id)
1678 .fetch_one(&mut **self.transaction()?)
1679 .await
1680 .map_err(|_| RepositoryError::Unavailable)?;
1681 if !compatible {
1682 return Err(RepositoryError::IncompatibleDefinition {
1683 instance_id: job_instance_id,
1684 });
1685 }
1686 upgrade_from = Some(previous_definition_id);
1687 }
1688 }
1689 let attempt: i32 = sqlx::query_scalar(
1690 "SELECT COALESCE(MAX(attempt), 0) + 1 \
1691 FROM oxide_batch.ob_job_execution WHERE job_instance_id = $1",
1692 )
1693 .bind(instance_database_id)
1694 .fetch_one(&mut **self.transaction()?)
1695 .await
1696 .map_err(|_| RepositoryError::Unavailable)?;
1697 let restart_of = latest
1698 .as_ref()
1699 .map(|execution| database_id(execution.id().get(), IdentifierKind::JobExecution))
1700 .transpose()?;
1701 let created_at = self.repository.clock.now();
1702 let created_ms = system_time_millis(created_at)?;
1703 let context = Json(json!({}));
1704 let id: i64 = sqlx::query_scalar(
1705 "INSERT INTO oxide_batch.ob_job_execution \
1706 (job_instance_id, definition_id, upgrade_from_definition_id, \
1707 restart_of_execution_id, attempt, \
1708 status, exit_code, parameters, context_format, context_schema, \
1709 context_schema_version, context_payload, created_at, updated_at, version) \
1710 VALUES ($1, $2, $3, $4, $5, 'STARTING', 'UNKNOWN', $6, 1, $7, 1, $8, \
1711 to_timestamp($9::double precision / 1000.0), \
1712 to_timestamp($9::double precision / 1000.0), 0) \
1713 RETURNING id",
1714 )
1715 .bind(instance_database_id)
1716 .bind(definition_id)
1717 .bind(upgrade_from)
1718 .bind(restart_of)
1719 .bind(attempt)
1720 .bind(parameters)
1721 .bind(DEFAULT_CONTEXT_SCHEMA)
1722 .bind(context)
1723 .bind(created_ms)
1724 .fetch_one(&mut **self.transaction()?)
1725 .await
1726 .map_err(|_| RepositoryError::Unavailable)?;
1727 let id_value =
1728 JobExecutionId::new(u64::try_from(id).map_err(|_| RepositoryError::Unavailable)?)?;
1729 Ok(JobExecution::new(
1730 id_value,
1731 job_instance_id,
1732 starting_metadata(created_at)?,
1733 ))
1734 })
1735 }
1736
1737 fn create_job_execution_with_definition<'a>(
1738 &'a mut self,
1739 job_instance_id: JobInstanceId,
1740 definition: &'a DefinitionIdentity,
1741 ) -> BoxFuture<'a, Result<JobExecution, RepositoryError>> {
1742 Box::pin(async move {
1743 self.definition_override = Some(definition.clone());
1744 self.create_job_execution(job_instance_id).await
1745 })
1746 }
1747
1748 #[allow(
1749 clippy::too_many_lines,
1750 reason = "exact and mapped restart-state insertion remain visible in one transaction"
1751 )]
1752 fn create_step_execution<'a>(
1753 &'a mut self,
1754 job_execution_id: JobExecutionId,
1755 step_name: &'a StepName,
1756 ) -> BoxFuture<'a, Result<StepExecution, RepositoryError>> {
1757 Box::pin(async move {
1758 let job_id = database_id(job_execution_id.get(), IdentifierKind::JobExecution)?;
1759 let execution_row = sqlx::query(
1760 "SELECT restart_of_execution_id, upgrade_from_definition_id \
1761 FROM oxide_batch.ob_job_execution WHERE id = $1",
1762 )
1763 .bind(job_id)
1764 .fetch_optional(&mut **self.transaction()?)
1765 .await
1766 .map_err(|_| RepositoryError::Unavailable)?
1767 .ok_or(RepositoryError::JobExecutionNotFound {
1768 id: job_execution_id,
1769 })?;
1770 let restart_source = execution_row
1771 .try_get::<Option<i64>, _>("restart_of_execution_id")
1772 .map_err(|_| RepositoryError::Unavailable)?;
1773 let upgraded = execution_row
1774 .try_get::<Option<i64>, _>("upgrade_from_definition_id")
1775 .map_err(|_| RepositoryError::Unavailable)?
1776 .is_some();
1777 let created_at = self.repository.clock.now();
1778 let created_ms = system_time_millis(created_at)?;
1779 let restarted_id = if let Some(source_job_id) = restart_source {
1780 let source_step_name = if upgraded {
1781 sqlx::query_scalar(
1782 "SELECT mapping.key \
1783 FROM oxide_batch.ob_job_execution execution \
1784 JOIN oxide_batch.ob_definition_upgrade upgrade \
1785 ON upgrade.from_definition_id = execution.upgrade_from_definition_id \
1786 AND upgrade.to_definition_id = execution.definition_id \
1787 CROSS JOIN LATERAL jsonb_each_text(upgrade.step_mapping) mapping \
1788 WHERE execution.id = $1 AND mapping.value = $2",
1789 )
1790 .bind(job_id)
1791 .bind(step_name.as_str())
1792 .fetch_optional(&mut **self.transaction()?)
1793 .await
1794 .map_err(|_| RepositoryError::Unavailable)?
1795 .ok_or(RepositoryError::InvalidDefinitionUpgrade {
1796 execution_id: job_execution_id,
1797 })?
1798 } else {
1799 step_name.as_str().to_owned()
1800 };
1801 sqlx::query_scalar(
1802 "INSERT INTO oxide_batch.ob_step_execution \
1803 (job_execution_id, step_name, step_logical_id, status, exit_code, \
1804 read_count, processed_count, write_count, filter_count, commit_count, \
1805 rollback_count, checkpoint_format, checkpoint_schema, \
1806 checkpoint_schema_version, checkpoint_payload, context_format, \
1807 context_schema, context_schema_version, context_payload, \
1808 read_retry_count, process_retry_count, write_retry_count, \
1809 read_skip_count, process_skip_count, write_skip_count, \
1810 no_rollback_count, fault_state_format, fault_state_schema, \
1811 fault_state_schema_version, fault_state_payload, fault_state_checksum, \
1812 created_at, updated_at, version) \
1813 SELECT $1, $2, $2, 'STARTING', 'UNKNOWN', source.read_count, \
1814 source.processed_count, source.write_count, source.filter_count, \
1815 source.commit_count, source.rollback_count, source.checkpoint_format, \
1816 source.checkpoint_schema, source.checkpoint_schema_version, \
1817 source.checkpoint_payload, source.context_format, source.context_schema, \
1818 source.context_schema_version, source.context_payload, \
1819 source.read_retry_count, source.process_retry_count, \
1820 source.write_retry_count, source.read_skip_count, \
1821 source.process_skip_count, source.write_skip_count, \
1822 source.no_rollback_count, source.fault_state_format, \
1823 source.fault_state_schema, source.fault_state_schema_version, \
1824 source.fault_state_payload, source.fault_state_checksum, \
1825 to_timestamp($3::double precision / 1000.0), \
1826 to_timestamp($3::double precision / 1000.0), 0 \
1827 FROM oxide_batch.ob_step_execution source \
1828 WHERE source.job_execution_id = $4 AND source.step_name = $5 \
1829 ORDER BY source.id DESC LIMIT 1 \
1830 RETURNING id",
1831 )
1832 .bind(job_id)
1833 .bind(step_name.as_str())
1834 .bind(created_ms)
1835 .bind(source_job_id)
1836 .bind(source_step_name)
1837 .fetch_optional(&mut **self.transaction()?)
1838 .await
1839 .map_err(|_| RepositoryError::Unavailable)?
1840 } else {
1841 None
1842 };
1843 let id: i64 = match (restart_source, restarted_id) {
1844 (Some(_), Some(id)) => id,
1845 (Some(_), None) => {
1846 return Err(RepositoryError::RestartStateNotFound {
1847 execution_id: job_execution_id,
1848 step_name: step_name.clone(),
1849 });
1850 }
1851 (None, _) => sqlx::query_scalar(
1852 "INSERT INTO oxide_batch.ob_step_execution \
1853 (job_execution_id, step_name, step_logical_id, status, exit_code, \
1854 checkpoint_format, checkpoint_schema, checkpoint_schema_version, \
1855 checkpoint_payload, context_format, context_schema, \
1856 context_schema_version, context_payload, created_at, updated_at, version) \
1857 VALUES ($1, $2, $2, 'STARTING', 'UNKNOWN', 1, $3, 1, $4, 1, $3, 1, $4, \
1858 to_timestamp($5::double precision / 1000.0), \
1859 to_timestamp($5::double precision / 1000.0), 0) \
1860 RETURNING id",
1861 )
1862 .bind(job_id)
1863 .bind(step_name.as_str())
1864 .bind(DEFAULT_CONTEXT_SCHEMA)
1865 .bind(Json(json!({})))
1866 .bind(created_ms)
1867 .fetch_one(&mut **self.transaction()?)
1868 .await
1869 .map_err(|_| RepositoryError::Unavailable)?,
1870 };
1871 let id_value =
1872 StepExecutionId::new(u64::try_from(id).map_err(|_| RepositoryError::Unavailable)?)?;
1873 let execution = self
1874 .step_execution(id_value)
1875 .await?
1876 .ok_or(RepositoryError::StepExecutionNotFound { id: id_value })?;
1877 Ok(execution)
1878 })
1879 }
1880
1881 #[allow(
1882 clippy::too_many_lines,
1883 reason = "the instance lock, start-limit check, and state-copy insert form one visible atomic rule"
1884 )]
1885 fn create_flow_step_execution<'a>(
1886 &'a mut self,
1887 job_execution_id: JobExecutionId,
1888 step_name: &'a StepName,
1889 node_id: &'a NodeId,
1890 start_limit: StartLimit,
1891 ) -> BoxFuture<'a, Result<StepExecution, RepositoryError>> {
1892 Box::pin(async move {
1893 let job_id = database_id(job_execution_id.get(), IdentifierKind::JobExecution)?;
1894 let instance_id: i64 = sqlx::query_scalar(
1895 "SELECT job_instance_id FROM oxide_batch.ob_job_execution WHERE id = $1",
1896 )
1897 .bind(job_id)
1898 .fetch_optional(&mut **self.transaction()?)
1899 .await
1900 .map_err(|_| RepositoryError::Unavailable)?
1901 .ok_or(RepositoryError::JobExecutionNotFound {
1902 id: job_execution_id,
1903 })?;
1904 sqlx::query("SELECT id FROM oxide_batch.ob_job_instance WHERE id = $1 FOR UPDATE")
1905 .bind(instance_id)
1906 .fetch_one(&mut **self.transaction()?)
1907 .await
1908 .map_err(|_| RepositoryError::Unavailable)?;
1909 let starts: i64 = sqlx::query_scalar(
1910 "SELECT count(*) FROM oxide_batch.ob_step_execution step \
1911 JOIN oxide_batch.ob_job_execution job ON job.id = step.job_execution_id \
1912 WHERE job.job_instance_id = $1 AND step.step_logical_id = $2",
1913 )
1914 .bind(instance_id)
1915 .bind(node_id.as_str())
1916 .fetch_one(&mut **self.transaction()?)
1917 .await
1918 .map_err(|_| RepositoryError::Unavailable)?;
1919 if u64::try_from(starts).map_err(|_| RepositoryError::FlowStateCorrupt)?
1920 >= u64::from(start_limit.get())
1921 {
1922 return Err(RepositoryError::StartLimitExceeded {
1923 instance_id: JobInstanceId::new(
1924 u64::try_from(instance_id)
1925 .map_err(|_| RepositoryError::FlowStateCorrupt)?,
1926 )?,
1927 node_id: node_id.clone(),
1928 limit: start_limit,
1929 });
1930 }
1931
1932 let created_ms = system_time_millis(self.repository.clock.now())?;
1933 let source_id: Option<i64> = sqlx::query_scalar(
1934 "SELECT step.id FROM oxide_batch.ob_step_execution step \
1935 JOIN oxide_batch.ob_job_execution job ON job.id = step.job_execution_id \
1936 WHERE job.job_instance_id = $1 AND step.step_logical_id = $2 \
1937 ORDER BY job.attempt DESC, step.id DESC LIMIT 1",
1938 )
1939 .bind(instance_id)
1940 .bind(node_id.as_str())
1941 .fetch_optional(&mut **self.transaction()?)
1942 .await
1943 .map_err(|_| RepositoryError::Unavailable)?;
1944
1945 let id: i64 = if let Some(source_id) = source_id {
1946 sqlx::query_scalar(
1947 "INSERT INTO oxide_batch.ob_step_execution \
1948 (job_execution_id, step_name, step_logical_id, status, exit_code, \
1949 read_count, processed_count, write_count, filter_count, commit_count, \
1950 rollback_count, checkpoint_format, checkpoint_schema, \
1951 checkpoint_schema_version, checkpoint_payload, context_format, \
1952 context_schema, context_schema_version, context_payload, \
1953 read_retry_count, process_retry_count, write_retry_count, \
1954 read_skip_count, process_skip_count, write_skip_count, \
1955 no_rollback_count, fault_state_format, fault_state_schema, \
1956 fault_state_schema_version, fault_state_payload, fault_state_checksum, \
1957 created_at, updated_at, version) \
1958 SELECT $1, $2, $3, 'STARTING', 'UNKNOWN', source.read_count, \
1959 source.processed_count, source.write_count, source.filter_count, \
1960 source.commit_count, source.rollback_count, source.checkpoint_format, \
1961 source.checkpoint_schema, source.checkpoint_schema_version, \
1962 source.checkpoint_payload, source.context_format, source.context_schema, \
1963 source.context_schema_version, source.context_payload, \
1964 source.read_retry_count, source.process_retry_count, \
1965 source.write_retry_count, source.read_skip_count, \
1966 source.process_skip_count, source.write_skip_count, \
1967 source.no_rollback_count, source.fault_state_format, \
1968 source.fault_state_schema, source.fault_state_schema_version, \
1969 source.fault_state_payload, source.fault_state_checksum, \
1970 to_timestamp($4::double precision / 1000.0), \
1971 to_timestamp($4::double precision / 1000.0), 0 \
1972 FROM oxide_batch.ob_step_execution source WHERE source.id = $5 \
1973 RETURNING id",
1974 )
1975 .bind(job_id)
1976 .bind(step_name.as_str())
1977 .bind(node_id.as_str())
1978 .bind(created_ms)
1979 .bind(source_id)
1980 .fetch_one(&mut **self.transaction()?)
1981 .await
1982 .map_err(|_| RepositoryError::ConcurrentModification)?
1983 } else {
1984 sqlx::query_scalar(
1985 "INSERT INTO oxide_batch.ob_step_execution \
1986 (job_execution_id, step_name, step_logical_id, status, exit_code, \
1987 checkpoint_format, checkpoint_schema, checkpoint_schema_version, \
1988 checkpoint_payload, context_format, context_schema, \
1989 context_schema_version, context_payload, created_at, updated_at, version) \
1990 VALUES ($1, $2, $3, 'STARTING', 'UNKNOWN', 1, $4, 1, $5, 1, $4, 1, $5, \
1991 to_timestamp($6::double precision / 1000.0), \
1992 to_timestamp($6::double precision / 1000.0), 0) RETURNING id",
1993 )
1994 .bind(job_id)
1995 .bind(step_name.as_str())
1996 .bind(node_id.as_str())
1997 .bind(DEFAULT_CONTEXT_SCHEMA)
1998 .bind(Json(json!({})))
1999 .bind(created_ms)
2000 .fetch_one(&mut **self.transaction()?)
2001 .await
2002 .map_err(|_| RepositoryError::ConcurrentModification)?
2003 };
2004 let id = StepExecutionId::new(
2005 u64::try_from(id).map_err(|_| RepositoryError::FlowStateCorrupt)?,
2006 )?;
2007 self.step_execution(id)
2008 .await?
2009 .ok_or(RepositoryError::StepExecutionNotFound { id })
2010 })
2011 }
2012
2013 fn transition_job_execution(
2014 &mut self,
2015 id: JobExecutionId,
2016 expected_version: ExecutionVersion,
2017 transition: LifecycleTransition,
2018 ) -> BoxFuture<'_, Result<JobExecution, RepositoryError>> {
2019 Box::pin(async move {
2020 let mut execution = self
2021 .job_execution(id)
2022 .await?
2023 .ok_or(RepositoryError::JobExecutionNotFound { id })?;
2024 execution.transition(expected_version, transition)?;
2025 let affected = update_job_execution(
2026 &mut **self.transaction()?,
2027 &execution,
2028 transition.transitioned_at(),
2029 expected_version,
2030 )
2031 .await?;
2032 if affected != 1 {
2033 return Err(self.classify_job_cas(id, expected_version).await);
2034 }
2035 Ok(execution)
2036 })
2037 }
2038
2039 fn enrich_job_exit_status<'a>(
2040 &'a mut self,
2041 id: JobExecutionId,
2042 expected_version: ExecutionVersion,
2043 exit_status: &'a ExitStatus,
2044 ) -> BoxFuture<'a, Result<JobExecution, RepositoryError>> {
2045 Box::pin(async move {
2046 let mut execution = self
2047 .job_execution(id)
2048 .await?
2049 .ok_or(RepositoryError::JobExecutionNotFound { id })?;
2050 execution.enrich_exit_status(expected_version, exit_status.clone())?;
2051 let updated_at = self.repository.clock.now();
2052 let affected = update_job_execution(
2053 &mut **self.transaction()?,
2054 &execution,
2055 updated_at,
2056 expected_version,
2057 )
2058 .await?;
2059 if affected != 1 {
2060 return Err(self.classify_job_cas(id, expected_version).await);
2061 }
2062 Ok(execution)
2063 })
2064 }
2065
2066 fn transition_step_execution(
2067 &mut self,
2068 id: StepExecutionId,
2069 expected_version: ExecutionVersion,
2070 transition: LifecycleTransition,
2071 ) -> BoxFuture<'_, Result<StepExecution, RepositoryError>> {
2072 Box::pin(async move {
2073 let mut execution = self
2074 .step_execution(id)
2075 .await?
2076 .ok_or(RepositoryError::StepExecutionNotFound { id })?;
2077 execution.transition(expected_version, transition)?;
2078 let affected = update_step_execution(
2079 &mut **self.transaction()?,
2080 &execution,
2081 transition.transitioned_at(),
2082 expected_version,
2083 )
2084 .await?;
2085 if affected != 1 {
2086 return Err(self.classify_step_cas(id, expected_version).await);
2087 }
2088 Ok(execution)
2089 })
2090 }
2091
2092 fn enrich_step_exit_status<'a>(
2093 &'a mut self,
2094 id: StepExecutionId,
2095 expected_version: ExecutionVersion,
2096 exit_status: &'a ExitStatus,
2097 ) -> BoxFuture<'a, Result<StepExecution, RepositoryError>> {
2098 Box::pin(async move {
2099 let mut execution = self
2100 .step_execution(id)
2101 .await?
2102 .ok_or(RepositoryError::StepExecutionNotFound { id })?;
2103 execution.enrich_exit_status(expected_version, exit_status.clone())?;
2104 let updated_at = self.repository.clock.now();
2105 let affected = update_step_execution(
2106 &mut **self.transaction()?,
2107 &execution,
2108 updated_at,
2109 expected_version,
2110 )
2111 .await?;
2112 if affected != 1 {
2113 return Err(self.classify_step_cas(id, expected_version).await);
2114 }
2115 Ok(execution)
2116 })
2117 }
2118
2119 fn find_job_instance<'a>(
2120 &'a mut self,
2121 key: &'a JobInstanceKey,
2122 ) -> BoxFuture<'a, Result<Option<JobInstance>, RepositoryError>> {
2123 Box::pin(async move {
2124 let digest = key.digest();
2125 let row = sqlx::query(
2126 "SELECT id, job_name, identifying_parameters \
2127 FROM oxide_batch.ob_job_instance \
2128 WHERE job_name = $1 AND instance_key = $2",
2129 )
2130 .bind(key.job_name().as_str())
2131 .bind(&digest[..])
2132 .fetch_optional(&mut **self.transaction()?)
2133 .await
2134 .map_err(|_| RepositoryError::Unavailable)?;
2135 row.map(|row| decode_job_instance(&row)).transpose()
2136 })
2137 }
2138
2139 fn get_job_instance(
2140 &mut self,
2141 id: JobInstanceId,
2142 ) -> BoxFuture<'_, Result<Option<JobInstance>, RepositoryError>> {
2143 Box::pin(async move {
2144 let id = database_id(id.get(), IdentifierKind::JobInstance)?;
2145 let row = sqlx::query(
2146 "SELECT id, job_name, identifying_parameters \
2147 FROM oxide_batch.ob_job_instance WHERE id = $1",
2148 )
2149 .bind(id)
2150 .fetch_optional(&mut **self.transaction()?)
2151 .await
2152 .map_err(|_| RepositoryError::Unavailable)?;
2153 row.map(|row| decode_job_instance(&row)).transpose()
2154 })
2155 }
2156
2157 fn get_job_execution(
2158 &mut self,
2159 id: JobExecutionId,
2160 ) -> BoxFuture<'_, Result<Option<JobExecution>, RepositoryError>> {
2161 Box::pin(async move { self.job_execution(id).await })
2162 }
2163
2164 fn job_executions(
2165 &mut self,
2166 job_instance_id: JobInstanceId,
2167 ) -> BoxFuture<'_, Result<Vec<JobExecution>, RepositoryError>> {
2168 Box::pin(async move {
2169 let instance_id = database_id(job_instance_id.get(), IdentifierKind::JobInstance)?;
2170 let exists: bool = sqlx::query_scalar(
2171 "SELECT EXISTS(SELECT 1 FROM oxide_batch.ob_job_instance WHERE id = $1)",
2172 )
2173 .bind(instance_id)
2174 .fetch_one(&mut **self.transaction()?)
2175 .await
2176 .map_err(|_| RepositoryError::Unavailable)?;
2177 if !exists {
2178 return Err(RepositoryError::JobInstanceNotFound {
2179 id: job_instance_id,
2180 });
2181 }
2182 let rows = sqlx::query(AssertSqlSafe(job_execution_select(
2183 "WHERE execution.job_instance_id = $1 ORDER BY execution.attempt",
2184 )))
2185 .bind(instance_id)
2186 .fetch_all(&mut **self.transaction()?)
2187 .await
2188 .map_err(|_| RepositoryError::Unavailable)?;
2189 rows.iter().map(decode_job_execution).collect()
2190 })
2191 }
2192
2193 fn get_step_execution(
2194 &mut self,
2195 id: StepExecutionId,
2196 ) -> BoxFuture<'_, Result<Option<StepExecution>, RepositoryError>> {
2197 Box::pin(async move { self.step_execution(id).await })
2198 }
2199
2200 fn step_executions(
2201 &mut self,
2202 job_execution_id: JobExecutionId,
2203 ) -> BoxFuture<'_, Result<Vec<StepExecution>, RepositoryError>> {
2204 Box::pin(async move {
2205 let job_id = database_id(job_execution_id.get(), IdentifierKind::JobExecution)?;
2206 let exists: bool = sqlx::query_scalar(
2207 "SELECT EXISTS(SELECT 1 FROM oxide_batch.ob_job_execution WHERE id = $1)",
2208 )
2209 .bind(job_id)
2210 .fetch_one(&mut **self.transaction()?)
2211 .await
2212 .map_err(|_| RepositoryError::Unavailable)?;
2213 if !exists {
2214 return Err(RepositoryError::JobExecutionNotFound {
2215 id: job_execution_id,
2216 });
2217 }
2218 let rows = sqlx::query(AssertSqlSafe(step_execution_select(
2219 "WHERE execution.job_execution_id = $1 ORDER BY execution.id",
2220 )))
2221 .bind(job_id)
2222 .fetch_all(&mut **self.transaction()?)
2223 .await
2224 .map_err(|_| RepositoryError::Unavailable)?;
2225 rows.iter().map(decode_step_execution).collect()
2226 })
2227 }
2228
2229 fn latest_flow_step<'a>(
2230 &'a mut self,
2231 job_instance_id: JobInstanceId,
2232 node_id: &'a NodeId,
2233 ) -> BoxFuture<'a, Result<Option<FlowStepState>, RepositoryError>> {
2234 Box::pin(async move {
2235 let instance_id = database_id(job_instance_id.get(), IdentifierKind::JobInstance)?;
2236 let exists: bool = sqlx::query_scalar(
2237 "SELECT EXISTS(SELECT 1 FROM oxide_batch.ob_job_instance WHERE id = $1)",
2238 )
2239 .bind(instance_id)
2240 .fetch_one(&mut **self.transaction()?)
2241 .await
2242 .map_err(|_| RepositoryError::Unavailable)?;
2243 if !exists {
2244 return Err(RepositoryError::JobInstanceNotFound {
2245 id: job_instance_id,
2246 });
2247 }
2248 let row = sqlx::query(AssertSqlSafe(durable_step_select(
2249 "JOIN oxide_batch.ob_job_execution flow_job \
2250 ON flow_job.id = execution.job_execution_id \
2251 WHERE flow_job.job_instance_id = $1 AND execution.step_logical_id = $2 \
2252 ORDER BY flow_job.attempt DESC, execution.id DESC LIMIT 1",
2253 )))
2254 .bind(instance_id)
2255 .bind(node_id.as_str())
2256 .fetch_optional(&mut **self.transaction()?)
2257 .await
2258 .map_err(|_| RepositoryError::Unavailable)?;
2259 row.map(|row| {
2260 let durable = decode_durable_step_state(&row)?;
2261 Ok(FlowStepState::new(
2262 node_id.clone(),
2263 durable.step_execution,
2264 Some(durable.execution_context),
2265 ))
2266 })
2267 .transpose()
2268 })
2269 }
2270
2271 #[allow(clippy::too_many_lines)]
2272 fn append_flow_decision<'a>(
2273 &'a mut self,
2274 request: &'a FlowDecisionRequest,
2275 ) -> BoxFuture<'a, Result<FlowDecision, RepositoryError>> {
2276 Box::pin(async move {
2277 let job_id = database_id(
2278 request.job_execution_id().get(),
2279 IdentifierKind::JobExecution,
2280 )?;
2281 let (fingerprint, Json(manifest)): (Vec<u8>, Json<Value>) = sqlx::query_as(
2282 "SELECT definition.manifest_digest, definition.manifest \
2283 FROM oxide_batch.ob_job_execution execution \
2284 JOIN oxide_batch.ob_job_definition definition \
2285 ON definition.id = execution.definition_id WHERE execution.id = $1",
2286 )
2287 .bind(job_id)
2288 .fetch_optional(&mut **self.transaction()?)
2289 .await
2290 .map_err(|_| RepositoryError::Unavailable)?
2291 .ok_or(RepositoryError::JobExecutionNotFound {
2292 id: request.job_execution_id(),
2293 })?;
2294 if fingerprint.as_slice() != request.plan_fingerprint() {
2295 return Err(RepositoryError::FlowStateCorrupt);
2296 }
2297 if !crate::flow::decision_matches_manifest(&manifest, request) {
2298 return Err(RepositoryError::FlowStateCorrupt);
2299 }
2300 if let Some(step_id) = request.source_step_execution_id() {
2301 let valid: bool = sqlx::query_scalar(
2302 "SELECT EXISTS( \
2303 SELECT 1 FROM oxide_batch.ob_step_execution source \
2304 JOIN oxide_batch.ob_job_execution source_job \
2305 ON source_job.id = source.job_execution_id \
2306 JOIN oxide_batch.ob_job_execution target_job ON target_job.id = $1 \
2307 WHERE source.id = $2 \
2308 AND source_job.job_instance_id = target_job.job_instance_id \
2309 AND source.step_logical_id = $3)",
2310 )
2311 .bind(job_id)
2312 .bind(database_id(step_id.get(), IdentifierKind::StepExecution)?)
2313 .bind(request.source_node_id().as_str())
2314 .fetch_one(&mut **self.transaction()?)
2315 .await
2316 .map_err(|_| RepositoryError::Unavailable)?;
2317 if !valid {
2318 return Err(RepositoryError::FlowStateCorrupt);
2319 }
2320 } else if !matches!(
2321 request.kind(),
2322 FlowTransitionKind::Decider | FlowTransitionKind::SplitAggregate
2323 ) {
2324 return Err(RepositoryError::FlowStateCorrupt);
2325 }
2326 if let Some(reused_id) = request.reused_decision_id() {
2327 let valid: bool = sqlx::query_scalar(
2328 "SELECT EXISTS( \
2329 SELECT 1 FROM oxide_batch.ob_flow_decision prior \
2330 JOIN oxide_batch.ob_job_execution prior_job \
2331 ON prior_job.id = prior.job_execution_id \
2332 JOIN oxide_batch.ob_job_execution target_job ON target_job.id = $1 \
2333 WHERE prior.id = $2 \
2334 AND prior_job.job_instance_id = target_job.job_instance_id \
2335 AND prior.source_node_id = $3 \
2336 AND prior.plan_fingerprint = $4 \
2337 AND prior.input_digest = $5 \
2338 AND prior.observed_outcome = $6 \
2339 AND prior.target_node_id IS NOT DISTINCT FROM $7 \
2340 AND prior.terminal_kind IS NOT DISTINCT FROM $8)",
2341 )
2342 .bind(job_id)
2343 .bind(database_id(reused_id.get(), IdentifierKind::FlowDecision)?)
2344 .bind(request.source_node_id().as_str())
2345 .bind(request.plan_fingerprint().as_slice())
2346 .bind(request.input_digest().as_slice())
2347 .bind(request.observed_outcome().as_str())
2348 .bind(flow_target_node(request.target()))
2349 .bind(flow_terminal_code(request.target())?)
2350 .fetch_one(&mut **self.transaction()?)
2351 .await
2352 .map_err(|_| RepositoryError::Unavailable)?;
2353 if !valid {
2354 return Err(RepositoryError::FlowStateCorrupt);
2355 }
2356 }
2357 let decided_ms = system_time_millis(request.decided_at())?;
2358 let id: i64 = sqlx::query_scalar(
2359 "INSERT INTO oxide_batch.ob_flow_decision \
2360 (job_execution_id, source_step_execution_id, reused_decision_id, sequence, \
2361 source_node_id, observed_outcome, target_node_id, transition_kind, \
2362 terminal_kind, plan_fingerprint, input_digest, decided_at) \
2363 VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, \
2364 to_timestamp($12::double precision / 1000.0)) RETURNING id",
2365 )
2366 .bind(job_id)
2367 .bind(
2368 request
2369 .source_step_execution_id()
2370 .map(|id| database_id(id.get(), IdentifierKind::StepExecution))
2371 .transpose()?,
2372 )
2373 .bind(
2374 request
2375 .reused_decision_id()
2376 .map(|id| database_id(id.get(), IdentifierKind::FlowDecision))
2377 .transpose()?,
2378 )
2379 .bind(database_id(
2380 request.sequence().get(),
2381 IdentifierKind::FlowDecision,
2382 )?)
2383 .bind(request.source_node_id().as_str())
2384 .bind(request.observed_outcome().as_str())
2385 .bind(flow_target_node(request.target()))
2386 .bind(request.kind().durable_code())
2387 .bind(flow_terminal_code(request.target())?)
2388 .bind(request.plan_fingerprint().as_slice())
2389 .bind(request.input_digest().as_slice())
2390 .bind(decided_ms)
2391 .fetch_one(&mut **self.transaction()?)
2392 .await
2393 .map_err(|_| RepositoryError::ConcurrentModification)?;
2394 Ok(FlowDecision::new(
2395 FlowDecisionId::new(
2396 u64::try_from(id).map_err(|_| RepositoryError::FlowStateCorrupt)?,
2397 )?,
2398 request.job_execution_id(),
2399 request.sequence(),
2400 request.source_node_id().clone(),
2401 request.source_step_execution_id(),
2402 request.kind(),
2403 request.observed_outcome().clone(),
2404 request.target().clone(),
2405 *request.plan_fingerprint(),
2406 *request.input_digest(),
2407 request.reused_decision_id(),
2408 request.decided_at(),
2409 ))
2410 })
2411 }
2412
2413 fn find_reusable_flow_decision<'a>(
2414 &'a mut self,
2415 job_instance_id: JobInstanceId,
2416 node_id: &'a NodeId,
2417 plan_fingerprint: &'a [u8; 32],
2418 input_digest: &'a [u8; 32],
2419 kind: FlowTransitionKind,
2420 ) -> BoxFuture<'a, Result<Option<FlowDecision>, RepositoryError>> {
2421 Box::pin(async move {
2422 let instance_id = database_id(job_instance_id.get(), IdentifierKind::JobInstance)?;
2423 let row = sqlx::query(AssertSqlSafe(flow_decision_select(
2424 "JOIN oxide_batch.ob_job_execution flow_job \
2425 ON flow_job.id = decision.job_execution_id \
2426 WHERE flow_job.job_instance_id = $1 AND decision.source_node_id = $2 \
2427 AND decision.plan_fingerprint = $3 AND decision.input_digest = $4 \
2428 AND decision.transition_kind = $5 \
2429 ORDER BY flow_job.attempt DESC, decision.sequence DESC LIMIT 1",
2430 )))
2431 .bind(instance_id)
2432 .bind(node_id.as_str())
2433 .bind(plan_fingerprint.as_slice())
2434 .bind(input_digest.as_slice())
2435 .bind(kind.durable_code())
2436 .fetch_optional(&mut **self.transaction()?)
2437 .await
2438 .map_err(|_| RepositoryError::Unavailable)?;
2439 row.map(|row| decode_flow_decision(&row)).transpose()
2440 })
2441 }
2442
2443 fn flow_decisions(
2444 &mut self,
2445 job_execution_id: JobExecutionId,
2446 ) -> BoxFuture<'_, Result<Vec<FlowDecision>, RepositoryError>> {
2447 Box::pin(async move {
2448 let job_id = database_id(job_execution_id.get(), IdentifierKind::JobExecution)?;
2449 let exists: bool = sqlx::query_scalar(
2450 "SELECT EXISTS(SELECT 1 FROM oxide_batch.ob_job_execution WHERE id = $1)",
2451 )
2452 .bind(job_id)
2453 .fetch_one(&mut **self.transaction()?)
2454 .await
2455 .map_err(|_| RepositoryError::Unavailable)?;
2456 if !exists {
2457 return Err(RepositoryError::JobExecutionNotFound {
2458 id: job_execution_id,
2459 });
2460 }
2461 let rows = sqlx::query(AssertSqlSafe(flow_decision_select(
2462 "WHERE decision.job_execution_id = $1 ORDER BY decision.sequence",
2463 )))
2464 .bind(job_id)
2465 .fetch_all(&mut **self.transaction()?)
2466 .await
2467 .map_err(|_| RepositoryError::Unavailable)?;
2468 rows.iter().map(decode_flow_decision).collect()
2469 })
2470 }
2471
2472 fn create_step_partition_plan<'a>(
2473 &'a mut self,
2474 step_execution_id: StepExecutionId,
2475 entries: &'a [PartitionPlanEntry],
2476 ) -> BoxFuture<'a, Result<Vec<StepPartition>, RepositoryError>> {
2477 Box::pin(async move {
2478 if entries.is_empty() {
2479 return Err(RepositoryError::EmptyPartitionPlan);
2480 }
2481 if entries.len() > usize::from(MAX_PARTITIONS) {
2482 return Err(RepositoryError::PartitionPlanTooLarge {
2483 max: usize::from(MAX_PARTITIONS),
2484 });
2485 }
2486 let mut keys = BTreeSet::new();
2487 for entry in entries {
2488 if !keys.insert(entry.key()) {
2489 return Err(RepositoryError::DuplicatePartitionKey);
2490 }
2491 }
2492
2493 let parent_id = database_id(step_execution_id.get(), IdentifierKind::StepExecution)?;
2494 let parent_status = sqlx::query_scalar::<_, String>(
2495 "SELECT status FROM oxide_batch.ob_step_execution \
2496 WHERE id = $1 FOR UPDATE",
2497 )
2498 .bind(parent_id)
2499 .fetch_optional(&mut **self.transaction()?)
2500 .await
2501 .map_err(|_| RepositoryError::Unavailable)?;
2502 let parent_status = parent_status
2503 .ok_or(RepositoryError::StepExecutionNotFound {
2504 id: step_execution_id,
2505 })
2506 .and_then(|status| decode_status(&status))?;
2507 if !matches!(parent_status, BatchStatus::Starting | BatchStatus::Started) {
2508 return Err(RepositoryError::PartitionParentNotActive {
2509 step_execution_id,
2510 status: parent_status,
2511 });
2512 }
2513 let plan_exists: bool = sqlx::query_scalar(
2514 "SELECT EXISTS(SELECT 1 FROM oxide_batch.ob_step_partition \
2515 WHERE step_execution_id = $1)",
2516 )
2517 .bind(parent_id)
2518 .fetch_one(&mut **self.transaction()?)
2519 .await
2520 .map_err(|_| RepositoryError::Unavailable)?;
2521 if plan_exists {
2522 return Err(RepositoryError::PartitionPlanExists { step_execution_id });
2523 }
2524
2525 let mut partitions = Vec::with_capacity(entries.len());
2526 for (index, entry) in entries.iter().enumerate() {
2527 let ordinal = u32::try_from(index.saturating_add(1))
2528 .map_err(|_| RepositoryError::PartitionStateCorrupt)?;
2529 let payload = entry
2530 .context()
2531 .payload_json()
2532 .map_err(|_| RepositoryError::PartitionStateCorrupt)?;
2533 let payload = serde_json::from_slice::<Value>(&payload)
2534 .map_err(|_| RepositoryError::PartitionStateCorrupt)?;
2535 let envelope = entry
2536 .context()
2537 .to_json()
2538 .map_err(|_| RepositoryError::PartitionStateCorrupt)?;
2539 let checksum: [u8; 32] = Sha256::digest(&envelope).into();
2540 let database_ordinal =
2541 i32::try_from(ordinal).map_err(|_| RepositoryError::PartitionStateCorrupt)?;
2542 let format = i16::try_from(entry.context().format_version())
2543 .map_err(|_| RepositoryError::PartitionStateCorrupt)?;
2544 let schema_version = i32::try_from(entry.context().schema_version().get())
2545 .map_err(|_| RepositoryError::PartitionStateCorrupt)?;
2546 let database_partition_id: i64 = sqlx::query_scalar(
2547 "INSERT INTO oxide_batch.ob_step_partition \
2548 (step_execution_id, partition_key, partition_ordinal, status, \
2549 context_format, context_schema, context_schema_version, \
2550 context_payload, context_checksum, version) \
2551 VALUES ($1, $2, $3, 'STARTING', $4, $5, $6, $7, $8, 0) \
2552 RETURNING id",
2553 )
2554 .bind(parent_id)
2555 .bind(entry.key().as_str())
2556 .bind(database_ordinal)
2557 .bind(format)
2558 .bind(entry.context().schema_id().as_str())
2559 .bind(schema_version)
2560 .bind(Json(payload))
2561 .bind(&checksum[..])
2562 .fetch_one(&mut **self.transaction()?)
2563 .await
2564 .map_err(|_| RepositoryError::Unavailable)?;
2565 let id = StepPartitionId::new(
2566 u64::try_from(database_partition_id)
2567 .map_err(|_| RepositoryError::PartitionStateCorrupt)?,
2568 )?;
2569 partitions.push(StepPartition::starting(
2570 id,
2571 step_execution_id,
2572 ordinal,
2573 entry.clone(),
2574 ));
2575 }
2576 self.created_partition_plans.insert(step_execution_id);
2577 Ok(partitions)
2578 })
2579 }
2580
2581 fn step_partition_plan(
2582 &mut self,
2583 step_execution_id: StepExecutionId,
2584 ) -> BoxFuture<'_, Result<Vec<StepPartition>, RepositoryError>> {
2585 Box::pin(async move {
2586 let parent_id = database_id(step_execution_id.get(), IdentifierKind::StepExecution)?;
2587 let parent_exists: bool = sqlx::query_scalar(
2588 "SELECT EXISTS(SELECT 1 FROM oxide_batch.ob_step_execution WHERE id = $1)",
2589 )
2590 .bind(parent_id)
2591 .fetch_one(&mut **self.transaction()?)
2592 .await
2593 .map_err(|_| RepositoryError::Unavailable)?;
2594 if !parent_exists {
2595 return Err(RepositoryError::StepExecutionNotFound {
2596 id: step_execution_id,
2597 });
2598 }
2599 let rows = sqlx::query(AssertSqlSafe(partition_select(
2600 "WHERE partition.step_execution_id = $1 ORDER BY partition.partition_key",
2601 )))
2602 .bind(parent_id)
2603 .fetch_all(&mut **self.transaction()?)
2604 .await
2605 .map_err(|_| RepositoryError::Unavailable)?;
2606 rows.iter().map(decode_step_partition).collect()
2607 })
2608 }
2609
2610 #[allow(
2611 clippy::too_many_lines,
2612 reason = "source validation and complete bounded carry-forward form one transaction rule"
2613 )]
2614 fn restart_step_partition_plan(
2615 &mut self,
2616 source_step_execution_id: StepExecutionId,
2617 target_step_execution_id: StepExecutionId,
2618 ) -> BoxFuture<'_, Result<Vec<StepPartition>, RepositoryError>> {
2619 Box::pin(async move {
2620 let source_id = database_id(
2621 source_step_execution_id.get(),
2622 IdentifierKind::StepExecution,
2623 )?;
2624 let target_id = database_id(
2625 target_step_execution_id.get(),
2626 IdentifierKind::StepExecution,
2627 )?;
2628 let rows = sqlx::query(
2629 "SELECT source.step_logical_id = target.step_logical_id AS same_logical_id, \
2630 source_job.status AS source_job_status, target.status AS target_status \
2631 FROM oxide_batch.ob_step_execution source \
2632 JOIN oxide_batch.ob_job_execution source_job \
2633 ON source_job.id = source.job_execution_id \
2634 CROSS JOIN oxide_batch.ob_step_execution target \
2635 WHERE source.id = $1 AND target.id = $2 \
2636 FOR UPDATE OF source_job, source, target",
2637 )
2638 .bind(source_id)
2639 .bind(target_id)
2640 .fetch_optional(&mut **self.transaction()?)
2641 .await
2642 .map_err(|_| RepositoryError::Unavailable)?
2643 .ok_or(RepositoryError::PartitionStateCorrupt)?;
2644 let same_logical_id = rows
2645 .try_get::<bool, _>("same_logical_id")
2646 .map_err(|_| RepositoryError::PartitionStateCorrupt)?;
2647 let source_job_status = decode_status(
2648 rows.try_get::<String, _>("source_job_status")
2649 .map_err(|_| RepositoryError::PartitionStateCorrupt)?
2650 .as_str(),
2651 )?;
2652 let target_status = decode_status(
2653 rows.try_get::<String, _>("target_status")
2654 .map_err(|_| RepositoryError::PartitionStateCorrupt)?
2655 .as_str(),
2656 )?;
2657 if !same_logical_id
2658 || !matches!(
2659 source_job_status,
2660 BatchStatus::Failed | BatchStatus::Stopped
2661 )
2662 {
2663 return Err(RepositoryError::PartitionStateCorrupt);
2664 }
2665 if !matches!(target_status, BatchStatus::Starting | BatchStatus::Started) {
2666 return Err(RepositoryError::PartitionParentNotActive {
2667 step_execution_id: target_step_execution_id,
2668 status: target_status,
2669 });
2670 }
2671 let target_exists: bool = sqlx::query_scalar(
2672 "SELECT EXISTS(SELECT 1 FROM oxide_batch.ob_step_partition \
2673 WHERE step_execution_id = $1)",
2674 )
2675 .bind(target_id)
2676 .fetch_one(&mut **self.transaction()?)
2677 .await
2678 .map_err(|_| RepositoryError::Unavailable)?;
2679 if target_exists {
2680 return Err(RepositoryError::PartitionPlanExists {
2681 step_execution_id: target_step_execution_id,
2682 });
2683 }
2684 let source_rows = sqlx::query(AssertSqlSafe(partition_select(
2685 "WHERE partition.step_execution_id = $1 \
2686 ORDER BY partition.partition_ordinal FOR UPDATE",
2687 )))
2688 .bind(source_id)
2689 .fetch_all(&mut **self.transaction()?)
2690 .await
2691 .map_err(|_| RepositoryError::Unavailable)?;
2692 if source_rows.is_empty() {
2693 return Err(RepositoryError::PartitionStateCorrupt);
2694 }
2695 let sources = source_rows
2696 .iter()
2697 .map(decode_step_partition)
2698 .collect::<Result<Vec<_>, _>>()?;
2699 let mut copied = Vec::with_capacity(sources.len());
2700 for source in sources {
2701 let context_payload = source
2702 .context()
2703 .payload_json()
2704 .map_err(|_| RepositoryError::PartitionStateCorrupt)?;
2705 let payload = serde_json::from_slice::<Value>(&context_payload)
2706 .map_err(|_| RepositoryError::PartitionStateCorrupt)?;
2707 let envelope = source
2708 .context()
2709 .to_json()
2710 .map_err(|_| RepositoryError::PartitionStateCorrupt)?;
2711 let checksum: [u8; 32] = Sha256::digest(&envelope).into();
2712 let completed = source.status() == BatchStatus::Completed;
2713 if completed && source.worker_step_execution_id().is_none() {
2714 return Err(RepositoryError::PartitionStateCorrupt);
2715 }
2716 let worker_id = source
2717 .worker_step_execution_id()
2718 .map(|id| database_id(id.get(), IdentifierKind::StepExecution))
2719 .transpose()?;
2720 let counts = if completed {
2721 source.counts()
2722 } else {
2723 crate::ExecutionCounts::default()
2724 };
2725 let status = if completed {
2726 BatchStatus::Completed
2727 } else {
2728 BatchStatus::Starting
2729 };
2730 let exit = if completed {
2731 Some(source.exit_status().code().as_str())
2732 } else {
2733 None
2734 };
2735 let database_partition_id: i64 = sqlx::query_scalar(
2736 "INSERT INTO oxide_batch.ob_step_partition \
2737 (step_execution_id, worker_step_execution_id, partition_key, \
2738 partition_ordinal, status, exit_code, read_count, processed_count, \
2739 write_count, filter_count, commit_count, rollback_count, \
2740 context_format, context_schema, context_schema_version, \
2741 context_payload, context_checksum, version) \
2742 VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, \
2743 $13, $14, $15, $16, $17, 0) RETURNING id",
2744 )
2745 .bind(target_id)
2746 .bind(if completed { worker_id } else { None })
2747 .bind(source.key().as_str())
2748 .bind(
2749 i32::try_from(source.ordinal())
2750 .map_err(|_| RepositoryError::PartitionStateCorrupt)?,
2751 )
2752 .bind(status.as_str())
2753 .bind(exit)
2754 .bind(partition_count(counts.read())?)
2755 .bind(partition_count(counts.processed())?)
2756 .bind(partition_count(counts.written())?)
2757 .bind(partition_count(counts.filtered())?)
2758 .bind(partition_count(counts.committed())?)
2759 .bind(partition_count(counts.rolled_back())?)
2760 .bind(
2761 i16::try_from(source.context().format_version())
2762 .map_err(|_| RepositoryError::PartitionStateCorrupt)?,
2763 )
2764 .bind(source.context().schema_id().as_str())
2765 .bind(
2766 i32::try_from(source.context().schema_version().get())
2767 .map_err(|_| RepositoryError::PartitionStateCorrupt)?,
2768 )
2769 .bind(Json(payload))
2770 .bind(&checksum[..])
2771 .fetch_one(&mut **self.transaction()?)
2772 .await
2773 .map_err(|_| RepositoryError::Unavailable)?;
2774 copied.push(StepPartition::from_snapshot(
2775 StepPartitionId::new(
2776 u64::try_from(database_partition_id)
2777 .map_err(|_| RepositoryError::PartitionStateCorrupt)?,
2778 )?,
2779 target_step_execution_id,
2780 if completed {
2781 source.worker_step_execution_id()
2782 } else {
2783 None
2784 },
2785 source.key().clone(),
2786 source.ordinal(),
2787 status,
2788 if completed {
2789 source.exit_status().clone()
2790 } else {
2791 crate::ExitStatus::unknown()
2792 },
2793 counts,
2794 source.context().clone(),
2795 ExecutionVersion::INITIAL,
2796 ));
2797 }
2798 self.created_partition_plans
2799 .insert(target_step_execution_id);
2800 Ok(copied)
2801 })
2802 }
2803
2804 #[allow(
2805 clippy::too_many_lines,
2806 reason = "parent, partition, and worker validation form one lock-ordered assignment rule"
2807 )]
2808 fn assign_step_partition(
2809 &mut self,
2810 id: StepPartitionId,
2811 expected_version: ExecutionVersion,
2812 worker_step_execution_id: StepExecutionId,
2813 ) -> BoxFuture<'_, Result<StepPartition, RepositoryError>> {
2814 Box::pin(async move {
2815 let database_partition_id = database_id(id.get(), IdentifierKind::StepPartition)?;
2816 let parent_id: i64 = sqlx::query_scalar(
2817 "SELECT step_execution_id FROM oxide_batch.ob_step_partition WHERE id = $1",
2818 )
2819 .bind(database_partition_id)
2820 .fetch_optional(&mut **self.transaction()?)
2821 .await
2822 .map_err(|_| RepositoryError::Unavailable)?
2823 .ok_or(RepositoryError::StepPartitionNotFound { id })?;
2824 let parent_status: String = sqlx::query_scalar(
2825 "SELECT status FROM oxide_batch.ob_step_execution WHERE id = $1 FOR UPDATE",
2826 )
2827 .bind(parent_id)
2828 .fetch_one(&mut **self.transaction()?)
2829 .await
2830 .map_err(|_| RepositoryError::Unavailable)?;
2831 let parent_status = decode_status(&parent_status)?;
2832 if parent_status != BatchStatus::Started {
2833 return Err(RepositoryError::PartitionParentNotActive {
2834 step_execution_id: StepExecutionId::new(
2835 u64::try_from(parent_id)
2836 .map_err(|_| RepositoryError::PartitionStateCorrupt)?,
2837 )?,
2838 status: parent_status,
2839 });
2840 }
2841 let row = sqlx::query(AssertSqlSafe(partition_select(
2842 "WHERE partition.id = $1 FOR UPDATE",
2843 )))
2844 .bind(database_partition_id)
2845 .fetch_optional(&mut **self.transaction()?)
2846 .await
2847 .map_err(|_| RepositoryError::Unavailable)?
2848 .ok_or(RepositoryError::StepPartitionNotFound { id })?;
2849 let mut partition = decode_step_partition(&row)?;
2850 if self
2851 .created_partition_plans
2852 .contains(&partition.step_execution_id())
2853 {
2854 return Err(RepositoryError::PartitionPlanNotCommitted {
2855 step_execution_id: partition.step_execution_id(),
2856 });
2857 }
2858 partition
2859 .assign(expected_version, worker_step_execution_id)
2860 .map_err(|error| map_partition_mutation(id, error))?;
2861 let worker_id = database_id(
2862 worker_step_execution_id.get(),
2863 IdentifierKind::StepExecution,
2864 )?;
2865 let same_job: Option<bool> = sqlx::query_scalar(
2866 "SELECT parent.job_execution_id = worker.job_execution_id \
2867 FROM oxide_batch.ob_step_execution parent \
2868 CROSS JOIN oxide_batch.ob_step_execution worker \
2869 WHERE parent.id = $1 AND worker.id = $2 FOR UPDATE OF worker",
2870 )
2871 .bind(parent_id)
2872 .bind(worker_id)
2873 .fetch_optional(&mut **self.transaction()?)
2874 .await
2875 .map_err(|_| RepositoryError::Unavailable)?;
2876 match same_job {
2877 None => {
2878 return Err(RepositoryError::StepExecutionNotFound {
2879 id: worker_step_execution_id,
2880 });
2881 }
2882 Some(same_job)
2883 if !same_job || partition.step_execution_id() == worker_step_execution_id =>
2884 {
2885 return Err(RepositoryError::PartitionWorkerMismatch {
2886 partition_id: id,
2887 worker_step_execution_id,
2888 });
2889 }
2890 Some(true) => {}
2891 Some(false) => return Err(RepositoryError::PartitionStateCorrupt),
2892 }
2893 let worker_assigned: bool = sqlx::query_scalar(
2894 "SELECT EXISTS(SELECT 1 FROM oxide_batch.ob_step_partition \
2895 WHERE worker_step_execution_id = $1)",
2896 )
2897 .bind(worker_id)
2898 .fetch_one(&mut **self.transaction()?)
2899 .await
2900 .map_err(|_| RepositoryError::Unavailable)?;
2901 if worker_assigned {
2902 return Err(RepositoryError::PartitionWorkerAlreadyAssigned {
2903 worker_step_execution_id,
2904 });
2905 }
2906 let affected = sqlx::query(
2907 "UPDATE oxide_batch.ob_step_partition \
2908 SET worker_step_execution_id = $1, status = 'STARTED', exit_code = NULL, \
2909 read_count = 0, processed_count = 0, write_count = 0, filter_count = 0, \
2910 commit_count = 0, rollback_count = 0, version = version + 1 \
2911 WHERE id = $2 AND version = $3",
2912 )
2913 .bind(worker_id)
2914 .bind(database_partition_id)
2915 .bind(database_version(expected_version)?)
2916 .execute(&mut **self.transaction()?)
2917 .await
2918 .map_err(|_| RepositoryError::Unavailable)?
2919 .rows_affected();
2920 if affected != 1 {
2921 return Err(RepositoryError::ConcurrentModification);
2922 }
2923 Ok(partition)
2924 })
2925 }
2926
2927 fn complete_step_partition(
2928 &mut self,
2929 id: StepPartitionId,
2930 expected_version: ExecutionVersion,
2931 worker_step_execution_id: StepExecutionId,
2932 ) -> BoxFuture<'_, Result<StepPartition, RepositoryError>> {
2933 Box::pin(async move {
2934 let database_partition_id = database_id(id.get(), IdentifierKind::StepPartition)?;
2935 let parent_id: i64 = sqlx::query_scalar(
2936 "SELECT step_execution_id FROM oxide_batch.ob_step_partition WHERE id = $1",
2937 )
2938 .bind(database_partition_id)
2939 .fetch_optional(&mut **self.transaction()?)
2940 .await
2941 .map_err(|_| RepositoryError::Unavailable)?
2942 .ok_or(RepositoryError::StepPartitionNotFound { id })?;
2943 let parent_row = sqlx::query(AssertSqlSafe(step_execution_select(
2944 "WHERE execution.id = $1 FOR UPDATE",
2945 )))
2946 .bind(parent_id)
2947 .fetch_one(&mut **self.transaction()?)
2948 .await
2949 .map_err(|_| RepositoryError::Unavailable)?;
2950 let parent = decode_step_execution(&parent_row)?;
2951 if !matches!(
2952 parent.metadata().status(),
2953 BatchStatus::Started | BatchStatus::Stopping
2954 ) {
2955 return Err(RepositoryError::PartitionParentNotActive {
2956 step_execution_id: parent.id(),
2957 status: parent.metadata().status(),
2958 });
2959 }
2960 let row = sqlx::query(AssertSqlSafe(partition_select(
2961 "WHERE partition.id = $1 FOR UPDATE",
2962 )))
2963 .bind(database_partition_id)
2964 .fetch_optional(&mut **self.transaction()?)
2965 .await
2966 .map_err(|_| RepositoryError::Unavailable)?
2967 .ok_or(RepositoryError::StepPartitionNotFound { id })?;
2968 let mut partition = decode_step_partition(&row)?;
2969 if partition.worker_step_execution_id() != Some(worker_step_execution_id) {
2970 return Err(RepositoryError::PartitionWorkerStale {
2971 partition_id: id,
2972 worker_step_execution_id,
2973 });
2974 }
2975 let worker_id = database_id(
2976 worker_step_execution_id.get(),
2977 IdentifierKind::StepExecution,
2978 )?;
2979 let worker_row = sqlx::query(AssertSqlSafe(step_execution_select(
2980 "WHERE execution.id = $1 FOR UPDATE",
2981 )))
2982 .bind(worker_id)
2983 .fetch_optional(&mut **self.transaction()?)
2984 .await
2985 .map_err(|_| RepositoryError::Unavailable)?
2986 .ok_or(RepositoryError::StepExecutionNotFound {
2987 id: worker_step_execution_id,
2988 })?;
2989 let worker = decode_step_execution(&worker_row)?;
2990 if worker.job_execution_id() != parent.job_execution_id() {
2991 return Err(RepositoryError::PartitionWorkerMismatch {
2992 partition_id: id,
2993 worker_step_execution_id,
2994 });
2995 }
2996 let result = PartitionResult::from_worker(&worker).map_err(|_| {
2997 RepositoryError::PartitionAggregationIncomplete {
2998 step_execution_id: parent.id(),
2999 status: worker.metadata().status(),
3000 }
3001 })?;
3002 partition
3003 .complete(expected_version, &result)
3004 .map_err(|error| map_partition_mutation(id, error))?;
3005 let counts = result.counts();
3006 let affected = sqlx::query(
3007 "UPDATE oxide_batch.ob_step_partition \
3008 SET status = $1, exit_code = $2, read_count = $3, processed_count = $4, \
3009 write_count = $5, filter_count = $6, commit_count = $7, \
3010 rollback_count = $8, version = version + 1 \
3011 WHERE id = $9 AND version = $10",
3012 )
3013 .bind(result.status().as_str())
3014 .bind(result.exit_status().code().as_str())
3015 .bind(partition_count(counts.read())?)
3016 .bind(partition_count(counts.processed())?)
3017 .bind(partition_count(counts.written())?)
3018 .bind(partition_count(counts.filtered())?)
3019 .bind(partition_count(counts.committed())?)
3020 .bind(partition_count(counts.rolled_back())?)
3021 .bind(database_partition_id)
3022 .bind(database_version(expected_version)?)
3023 .execute(&mut **self.transaction()?)
3024 .await
3025 .map_err(|_| RepositoryError::Unavailable)?
3026 .rows_affected();
3027 if affected != 1 {
3028 return Err(RepositoryError::ConcurrentModification);
3029 }
3030 Ok(partition)
3031 })
3032 }
3033
3034 fn aggregate_step_partitions(
3035 &mut self,
3036 step_execution_id: StepExecutionId,
3037 expected_version: ExecutionVersion,
3038 transitioned_at: SystemTime,
3039 ) -> BoxFuture<'_, Result<StepExecution, RepositoryError>> {
3040 Box::pin(async move {
3041 let parent_id = database_id(step_execution_id.get(), IdentifierKind::StepExecution)?;
3042 let parent_row = sqlx::query(AssertSqlSafe(step_execution_select(
3043 "WHERE execution.id = $1 FOR UPDATE",
3044 )))
3045 .bind(parent_id)
3046 .fetch_optional(&mut **self.transaction()?)
3047 .await
3048 .map_err(|_| RepositoryError::Unavailable)?
3049 .ok_or(RepositoryError::StepExecutionNotFound {
3050 id: step_execution_id,
3051 })?;
3052 let parent = decode_step_execution(&parent_row)?;
3053 let rows = sqlx::query(AssertSqlSafe(partition_select(
3054 "WHERE partition.step_execution_id = $1 \
3055 ORDER BY partition.partition_key FOR UPDATE",
3056 )))
3057 .bind(parent_id)
3058 .fetch_all(&mut **self.transaction()?)
3059 .await
3060 .map_err(|_| RepositoryError::Unavailable)?;
3061 let partitions = rows
3062 .iter()
3063 .map(decode_step_partition)
3064 .collect::<Result<Vec<_>, _>>()?;
3065 let aggregate = crate::aggregate_step_partitions(&partitions)
3066 .map_err(|error| map_partition_aggregation(step_execution_id, error))?;
3067 for partition in &partitions {
3068 let worker_id = partition.worker_step_execution_id().ok_or(
3069 RepositoryError::PartitionAggregationIncomplete {
3070 step_execution_id,
3071 status: partition.status(),
3072 },
3073 )?;
3074 let worker_row = sqlx::query(AssertSqlSafe(step_execution_select(
3075 "WHERE execution.id = $1 FOR UPDATE",
3076 )))
3077 .bind(database_id(worker_id.get(), IdentifierKind::StepExecution)?)
3078 .fetch_optional(&mut **self.transaction()?)
3079 .await
3080 .map_err(|_| RepositoryError::Unavailable)?
3081 .ok_or(RepositoryError::PartitionStateCorrupt)?;
3082 let worker = decode_step_execution(&worker_row)?;
3083 if worker.metadata().status() != partition.status()
3084 || worker.metadata().exit_status() != partition.exit_status()
3085 || worker.metadata().counts() != partition.counts()
3086 {
3087 return Err(RepositoryError::PartitionStateCorrupt);
3088 }
3089 }
3090 let selected_worker = self
3091 .step_execution(aggregate.selected_worker_step_execution_id())
3092 .await?
3093 .ok_or(RepositoryError::PartitionStateCorrupt)?;
3094 let failure = selected_worker.metadata().failure();
3095 if let Some(next) = expected_version.get().checked_add(1)
3096 && parent.version().get() == next
3097 && parent.metadata().status() == aggregate.status()
3098 && parent.metadata().exit_status() == aggregate.exit_status()
3099 && parent.metadata().counts() == aggregate.counts()
3100 && parent.metadata().failure() == failure
3101 {
3102 return Ok(parent);
3103 }
3104 if !matches!(
3105 parent.metadata().status(),
3106 BatchStatus::Started | BatchStatus::Stopping
3107 ) {
3108 return Err(RepositoryError::PartitionParentNotActive {
3109 step_execution_id,
3110 status: parent.metadata().status(),
3111 });
3112 }
3113 let aggregated = aggregate_partition_parent(
3114 &parent,
3115 expected_version,
3116 &aggregate,
3117 transitioned_at,
3118 failure,
3119 )?;
3120 let affected = update_step_execution(
3121 &mut **self.transaction()?,
3122 &aggregated,
3123 transitioned_at,
3124 expected_version,
3125 )
3126 .await?;
3127 if affected != 1 {
3128 return Err(self
3129 .classify_step_cas(step_execution_id, expected_version)
3130 .await);
3131 }
3132 Ok(aggregated)
3133 })
3134 }
3135
3136 fn recover_job_execution<'a>(
3137 &'a mut self,
3138 id: JobExecutionId,
3139 request: &'a RecoveryRequest,
3140 ) -> BoxFuture<'a, Result<RecoveryResult, RepositoryError>> {
3141 Box::pin(async move {
3142 let database_execution_id = database_id(id.get(), IdentifierKind::JobExecution)?;
3143 let row = sqlx::query(AssertSqlSafe(job_execution_select(
3144 "WHERE execution.id = $1 FOR UPDATE",
3145 )))
3146 .bind(database_execution_id)
3147 .fetch_optional(&mut **self.transaction()?)
3148 .await
3149 .map_err(|_| RepositoryError::Unavailable)?
3150 .ok_or(RepositoryError::JobExecutionNotFound { id })?;
3151 let prior = decode_job_execution(&row)?;
3152 let decided_at = self.repository.clock.now();
3153 let recovered = recovered_execution(&prior, request, decided_at)?;
3154 let decided_ms = system_time_millis(decided_at)?;
3155 let prior_version = database_version(request.expected_version())?;
3156 let resulting_status = recovered.metadata().status().to_string();
3157 sqlx::query("SAVEPOINT ob_recovery_decision")
3158 .execute(&mut **self.transaction()?)
3159 .await
3160 .map_err(|_| RepositoryError::Unavailable)?;
3161 let affected = update_job_execution(
3162 &mut **self.transaction()?,
3163 &recovered,
3164 decided_at,
3165 request.expected_version(),
3166 )
3167 .await?;
3168 if affected != 1 {
3169 rollback_recovery_savepoint(&mut **self.transaction()?).await;
3170 return Err(self.classify_job_cas(id, request.expected_version()).await);
3171 }
3172 let insert = sqlx::query(
3173 "INSERT INTO oxide_batch.ob_recovery_decision \
3174 (job_execution_id, execution_version, prior_status, resulting_status, \
3175 reason_code, operator_reference, evidence_digest, decided_at) \
3176 VALUES ($1, $2, $3, $4, $5, $6, $7, \
3177 to_timestamp($8::double precision / 1000.0)) \
3178 RETURNING id",
3179 )
3180 .bind(database_execution_id)
3181 .bind(prior_version)
3182 .bind(prior.metadata().status().to_string())
3183 .bind(&resulting_status)
3184 .bind(request.reason_code())
3185 .bind(request.operator_reference())
3186 .bind(&request.evidence_digest()[..])
3187 .bind(decided_ms)
3188 .fetch_one(&mut **self.transaction()?)
3189 .await;
3190 if insert.is_err() {
3191 rollback_recovery_savepoint(&mut **self.transaction()?).await;
3192 return Err(RepositoryError::ConcurrentModification);
3193 }
3194 sqlx::query("RELEASE SAVEPOINT ob_recovery_decision")
3195 .execute(&mut **self.transaction()?)
3196 .await
3197 .map_err(|_| RepositoryError::Unavailable)?;
3198 let inserted = insert.map_err(|_| RepositoryError::Unavailable)?;
3199 let decision = RecoveryDecision::new(
3200 RecoveryDecisionId::new(read_u64(&inserted, "id")?)?,
3201 id,
3202 request.expected_version(),
3203 prior.metadata().status(),
3204 recovered.metadata().status(),
3205 request.reason_code().to_owned(),
3206 request.operator_reference().to_owned(),
3207 *request.evidence_digest(),
3208 decided_at,
3209 );
3210 Ok(RecoveryResult::new(recovered, decision))
3211 })
3212 }
3213
3214 fn recovery_decision(
3215 &mut self,
3216 id: JobExecutionId,
3217 ) -> BoxFuture<'_, Result<Option<RecoveryDecision>, RepositoryError>> {
3218 Box::pin(async move {
3219 let database_execution_id = database_id(id.get(), IdentifierKind::JobExecution)?;
3220 let exists: bool = sqlx::query_scalar(
3221 "SELECT EXISTS(SELECT 1 FROM oxide_batch.ob_job_execution WHERE id = $1)",
3222 )
3223 .bind(database_execution_id)
3224 .fetch_one(&mut **self.transaction()?)
3225 .await
3226 .map_err(|_| RepositoryError::Unavailable)?;
3227 if !exists {
3228 return Err(RepositoryError::JobExecutionNotFound { id });
3229 }
3230 let row = sqlx::query(
3231 "SELECT id, execution_version, prior_status, resulting_status, reason_code, \
3232 operator_reference, evidence_digest, \
3233 (extract(epoch FROM decided_at) * 1000)::bigint AS decided_ms \
3234 FROM oxide_batch.ob_recovery_decision \
3235 WHERE job_execution_id = $1 ORDER BY decided_at DESC, id DESC LIMIT 1",
3236 )
3237 .bind(database_execution_id)
3238 .fetch_optional(&mut **self.transaction()?)
3239 .await
3240 .map_err(|_| RepositoryError::Unavailable)?;
3241 row.as_ref()
3242 .map(|row| decode_recovery_decision(id, row))
3243 .transpose()
3244 })
3245 }
3246
3247 fn find_operator_request<'a>(
3248 &'a mut self,
3249 action: OperatorAction,
3250 operation_id: &'a OperationId,
3251 ) -> BoxFuture<'a, Result<Option<OperatorRecord>, RepositoryError>> {
3252 Box::pin(async move {
3253 let row = sqlx::query(AssertSqlSafe(operator_request_select(
3254 "WHERE request.action = $1 AND request.operation_id = $2",
3255 )))
3256 .bind(action.as_str())
3257 .bind(operation_id.as_str())
3258 .fetch_optional(&mut **self.transaction()?)
3259 .await
3260 .map_err(|_| RepositoryError::Unavailable)?;
3261 row.as_ref().map(decode_operator_record).transpose()
3262 })
3263 }
3264
3265 fn append_operator_request<'a>(
3266 &'a mut self,
3267 draft: &'a OperatorRecordDraft,
3268 ) -> BoxFuture<'a, Result<OperatorRecord, RepositoryError>> {
3269 Box::pin(async move {
3270 let instance_id = draft
3271 .job_instance_id()
3272 .map(|id| database_id(id.get(), IdentifierKind::JobInstance))
3273 .transpose()?;
3274 let execution_id = draft
3275 .job_execution_id()
3276 .map(|id| database_id(id.get(), IdentifierKind::JobExecution))
3277 .transpose()?;
3278 let requested_ms = system_time_millis(draft.requested_at())?;
3279 let row = sqlx::query(
3280 "INSERT INTO oxide_batch.ob_operator_request \
3281 (job_instance_id, job_execution_id, action, authorization_class, \
3282 operation_id, actor_ref, reason_code, request_digest, observed_version, \
3283 prior_status, result_status, outcome_class, rejection_code, requested_at) \
3284 VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, \
3285 to_timestamp($14::double precision / 1000.0)) \
3286 RETURNING id",
3287 )
3288 .bind(instance_id)
3289 .bind(execution_id)
3290 .bind(draft.action().as_str())
3291 .bind(draft.action().authorization_class().as_str())
3292 .bind(draft.operation_id().as_str())
3293 .bind(draft.actor().as_str())
3294 .bind(draft.reason().map(ReasonCode::as_str))
3295 .bind(&draft.digest().as_bytes()[..])
3296 .bind(draft.observed_version().map(database_version).transpose()?)
3297 .bind(draft.prior_status().map(BatchStatus::as_str))
3298 .bind(draft.result_status().map(BatchStatus::as_str))
3299 .bind(draft.outcome().as_str())
3300 .bind(draft.rejection().map(OperatorRejection::as_str))
3301 .bind(requested_ms)
3302 .fetch_one(&mut **self.transaction()?)
3303 .await
3304 .map_err(|_| RepositoryError::ConcurrentModification)?;
3305 let id = OperatorRequestId::new(read_u64(&row, "id")?)?;
3306 Ok(OperatorRecord::from_parts(id, draft.clone()))
3307 })
3308 }
3309
3310 fn request_execution_stop<'a>(
3311 &'a mut self,
3312 id: JobExecutionId,
3313 expected_version: ExecutionVersion,
3314 actor: &'a ActorRef,
3315 requested_at: SystemTime,
3316 ) -> BoxFuture<'a, Result<JobExecution, RepositoryError>> {
3317 Box::pin(async move {
3318 let database_execution_id = database_id(id.get(), IdentifierKind::JobExecution)?;
3319 let requested_ms = system_time_millis(requested_at)?;
3320 let affected = sqlx::query(
3321 "UPDATE oxide_batch.ob_job_execution \
3322 SET stop_requested_at = to_timestamp($1::double precision / 1000.0), \
3323 stop_requested_by = $2, \
3324 updated_at = greatest(updated_at, \
3325 to_timestamp($1::double precision / 1000.0)) \
3326 WHERE id = $3 AND version = $4",
3327 )
3328 .bind(requested_ms)
3329 .bind(actor.as_str())
3330 .bind(database_execution_id)
3331 .bind(database_version(expected_version)?)
3332 .execute(&mut **self.transaction()?)
3333 .await
3334 .map_err(|_| RepositoryError::Unavailable)?
3335 .rows_affected();
3336 if affected != 1 {
3337 return Err(self.classify_job_cas(id, expected_version).await);
3338 }
3339 self.job_execution(id)
3340 .await?
3341 .ok_or(RepositoryError::JobExecutionNotFound { id })
3342 })
3343 }
3344
3345 fn claim_execution_owner<'a>(
3346 &'a mut self,
3347 id: JobExecutionId,
3348 expected_version: ExecutionVersion,
3349 owner: &'a crate::OwnerToken,
3350 claimed_at: SystemTime,
3351 ) -> BoxFuture<'a, Result<JobExecution, RepositoryError>> {
3352 Box::pin(async move {
3353 let database_execution_id = database_id(id.get(), IdentifierKind::JobExecution)?;
3354 let claimed_ms = system_time_millis(claimed_at)?;
3355 let affected = sqlx::query(
3356 "UPDATE oxide_batch.ob_job_execution \
3357 SET owner_token = $1, updated_at = greatest(updated_at, \
3358 to_timestamp($2::double precision / 1000.0)) \
3359 WHERE id = $3 AND version = $4 AND status = 'STARTING' \
3360 AND (owner_token IS NULL OR owner_token = $1)",
3361 )
3362 .bind(&owner.as_bytes()[..])
3363 .bind(claimed_ms)
3364 .bind(database_execution_id)
3365 .bind(database_version(expected_version)?)
3366 .execute(&mut **self.transaction()?)
3367 .await
3368 .map_err(|_| RepositoryError::Unavailable)?
3369 .rows_affected();
3370 if affected != 1 {
3371 let row = sqlx::query(
3372 "SELECT version, owner_token, status \
3373 FROM oxide_batch.ob_job_execution WHERE id = $1",
3374 )
3375 .bind(database_execution_id)
3376 .fetch_optional(&mut **self.transaction()?)
3377 .await
3378 .map_err(|_| RepositoryError::Unavailable)?
3379 .ok_or(RepositoryError::JobExecutionNotFound { id })?;
3380 let actual = ExecutionVersion::new(read_u64(&row, "version")?);
3381 if actual != expected_version {
3382 return Err(RepositoryError::Lifecycle(LifecycleError::StaleVersion {
3383 expected: expected_version,
3384 actual,
3385 }));
3386 }
3387 let status = decode_status(&read_text(&row, "status")?)?;
3388 if status != BatchStatus::Starting {
3389 return Err(RepositoryError::ExecutionOwnershipNotAllowed { id, status });
3390 }
3391 return Err(RepositoryError::ExecutionOwned { id });
3392 }
3393 self.job_execution(id)
3394 .await?
3395 .ok_or(RepositoryError::JobExecutionNotFound { id })
3396 })
3397 }
3398
3399 fn observe_execution_control<'a>(
3400 &'a mut self,
3401 id: JobExecutionId,
3402 owner: &'a crate::OwnerToken,
3403 observed_at: SystemTime,
3404 ) -> BoxFuture<'a, Result<crate::ExecutionControl, RepositoryError>> {
3405 Box::pin(async move {
3406 let database_execution_id = database_id(id.get(), IdentifierKind::JobExecution)?;
3407 let row = sqlx::query(
3408 "SELECT owner_token = $2 AS owner_matches, \
3409 stop_requested_at IS NOT NULL AS stop_requested \
3410 FROM oxide_batch.ob_job_execution WHERE id = $1 FOR UPDATE",
3411 )
3412 .bind(database_execution_id)
3413 .bind(&owner.as_bytes()[..])
3414 .fetch_optional(&mut **self.transaction()?)
3415 .await
3416 .map_err(|_| RepositoryError::Unavailable)?
3417 .ok_or(RepositoryError::JobExecutionNotFound { id })?;
3418 let owner_matches = row
3419 .try_get::<Option<bool>, _>("owner_matches")
3420 .map_err(|_| RepositoryError::Unavailable)?
3421 .unwrap_or(false);
3422 let stop_requested = row
3423 .try_get::<bool, _>("stop_requested")
3424 .map_err(|_| RepositoryError::Unavailable)?;
3425 let mut execution = self
3426 .job_execution(id)
3427 .await?
3428 .ok_or(RepositoryError::JobExecutionNotFound { id })?;
3429 if owner_matches
3430 && stop_requested
3431 && matches!(
3432 execution.metadata().status(),
3433 BatchStatus::Starting | BatchStatus::Started
3434 )
3435 {
3436 execution = self
3437 .transition_job_execution(
3438 id,
3439 execution.version(),
3440 LifecycleTransition::new(BatchStatus::Stopping, observed_at),
3441 )
3442 .await?;
3443 }
3444 Ok(crate::ExecutionControl::new(
3445 execution,
3446 owner_matches,
3447 stop_requested,
3448 ))
3449 })
3450 }
3451
3452 fn job_instance_hold(
3453 &mut self,
3454 id: JobInstanceId,
3455 ) -> BoxFuture<'_, Result<Option<RetentionHold>, RepositoryError>> {
3456 Box::pin(async move {
3457 let database_instance_id = database_id(id.get(), IdentifierKind::JobInstance)?;
3458 let row = sqlx::query(
3459 "SELECT hold_actor, hold_reason, \
3460 (extract(epoch FROM hold_placed_at) * 1000)::bigint AS placed_ms \
3461 FROM oxide_batch.ob_job_instance WHERE id = $1",
3462 )
3463 .bind(database_instance_id)
3464 .fetch_optional(&mut **self.transaction()?)
3465 .await
3466 .map_err(|_| RepositoryError::Unavailable)?
3467 .ok_or(RepositoryError::JobInstanceNotFound { id })?;
3468 decode_retention_hold(id, &row)
3469 })
3470 }
3471
3472 fn place_instance_hold<'a>(
3473 &'a mut self,
3474 id: JobInstanceId,
3475 actor: &'a ActorRef,
3476 reason: &'a ReasonCode,
3477 placed_at: SystemTime,
3478 ) -> BoxFuture<'a, Result<RetentionHold, RepositoryError>> {
3479 Box::pin(async move {
3480 let database_instance_id = database_id(id.get(), IdentifierKind::JobInstance)?;
3481 let placed_ms = system_time_millis(placed_at)?;
3482 let affected = sqlx::query(
3483 "UPDATE oxide_batch.ob_job_instance \
3484 SET hold_actor = $1, hold_reason = $2, \
3485 hold_placed_at = to_timestamp($3::double precision / 1000.0) \
3486 WHERE id = $4",
3487 )
3488 .bind(actor.as_str())
3489 .bind(reason.as_str())
3490 .bind(placed_ms)
3491 .bind(database_instance_id)
3492 .execute(&mut **self.transaction()?)
3493 .await
3494 .map_err(|_| RepositoryError::Unavailable)?
3495 .rows_affected();
3496 if affected != 1 {
3497 return Err(RepositoryError::JobInstanceNotFound { id });
3498 }
3499 Ok(RetentionHold::new(
3500 id,
3501 actor.clone(),
3502 reason.clone(),
3503 placed_at,
3504 ))
3505 })
3506 }
3507
3508 fn release_instance_hold(
3509 &mut self,
3510 id: JobInstanceId,
3511 ) -> BoxFuture<'_, Result<Option<RetentionHold>, RepositoryError>> {
3512 Box::pin(async move {
3513 let existing = self.job_instance_hold(id).await?;
3514 let database_instance_id = database_id(id.get(), IdentifierKind::JobInstance)?;
3515 sqlx::query(
3516 "UPDATE oxide_batch.ob_job_instance \
3517 SET hold_actor = NULL, hold_reason = NULL, hold_placed_at = NULL \
3518 WHERE id = $1",
3519 )
3520 .bind(database_instance_id)
3521 .execute(&mut **self.transaction()?)
3522 .await
3523 .map_err(|_| RepositoryError::Unavailable)?;
3524 Ok(existing)
3525 })
3526 }
3527
3528 fn find_retention_action<'a>(
3529 &'a mut self,
3530 action: RetentionAction,
3531 operation_id: &'a OperationId,
3532 ) -> BoxFuture<'a, Result<Option<RetentionRecord>, RepositoryError>> {
3533 Box::pin(async move {
3534 let row = sqlx::query(AssertSqlSafe(retention_action_select(
3535 "WHERE retention.action = $1 AND retention.operation_id = $2",
3536 )))
3537 .bind(action.as_str())
3538 .bind(operation_id.as_str())
3539 .fetch_optional(&mut **self.transaction()?)
3540 .await
3541 .map_err(|_| RepositoryError::Unavailable)?;
3542 row.as_ref().map(decode_retention_record).transpose()
3543 })
3544 }
3545
3546 fn append_retention_action<'a>(
3547 &'a mut self,
3548 draft: &'a RetentionRecordDraft,
3549 ) -> BoxFuture<'a, Result<RetentionRecord, RepositoryError>> {
3550 Box::pin(async move {
3551 let instance_id = draft
3552 .job_instance_id()
3553 .map(|id| database_id(id.get(), IdentifierKind::JobInstance))
3554 .transpose()?;
3555 let applied_ms = system_time_millis(draft.applied_at())?;
3556 let counts = draft.counts();
3557 let row = sqlx::query(
3558 "INSERT INTO oxide_batch.ob_retention_action \
3559 (job_instance_id, action, operation_id, actor_ref, reason_code, \
3560 plan_digest, batch_bound, deleted_flow_decisions, \
3561 deleted_recovery_decisions, deleted_operator_requests, \
3562 deleted_step_partitions, deleted_step_executions, \
3563 deleted_job_executions, deleted_job_instances, outcome_class, \
3564 applied_at) \
3565 VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, \
3566 $15, to_timestamp($16::double precision / 1000.0)) \
3567 RETURNING id",
3568 )
3569 .bind(instance_id)
3570 .bind(draft.action().as_str())
3571 .bind(draft.operation_id().as_str())
3572 .bind(draft.actor().as_str())
3573 .bind(draft.reason().as_str())
3574 .bind(draft.plan_digest().map(|digest| digest.to_vec()))
3575 .bind(
3576 draft
3577 .batch_bound()
3578 .map(|bound| i32::try_from(bound.get()))
3579 .transpose()
3580 .map_err(|_| RepositoryError::Unavailable)?,
3581 )
3582 .bind(retention_count(counts.flow_decisions())?)
3583 .bind(retention_count(counts.recovery_decisions())?)
3584 .bind(retention_count(counts.operator_requests())?)
3585 .bind(retention_count(counts.step_partitions())?)
3586 .bind(retention_count(counts.step_executions())?)
3587 .bind(retention_count(counts.job_executions())?)
3588 .bind(retention_count(counts.job_instances())?)
3589 .bind(draft.outcome().as_str())
3590 .bind(applied_ms)
3591 .fetch_one(&mut **self.transaction()?)
3592 .await
3593 .map_err(|_| RepositoryError::ConcurrentModification)?;
3594 let id = RetentionActionId::new(read_u64(&row, "id")?)?;
3595 Ok(RetentionRecord::from_parts(id, draft.clone()))
3596 })
3597 }
3598
3599 fn purge_survey<'a>(
3600 &'a mut self,
3601 request: &'a PurgePlanRequest,
3602 ) -> BoxFuture<'a, Result<PurgeSurvey, RepositoryError>> {
3603 Box::pin(async move {
3604 let statuses = request
3605 .statuses()
3606 .iter()
3607 .map(|status| status.as_str().to_owned())
3608 .collect::<Vec<_>>();
3609 let now = self.repository.clock.now();
3610 let threshold = system_time_millis(
3611 now.checked_sub(request.minimum_age())
3612 .ok_or(RepositoryError::Unavailable)?,
3613 )?;
3614 let limit = i64::from(request.batch().get());
3615 let rows = sqlx::query(
3616 "SELECT execution.job_instance_id, execution.id, execution.version \
3617 FROM oxide_batch.ob_job_execution execution \
3618 JOIN oxide_batch.ob_job_instance instance \
3619 ON instance.id = execution.job_instance_id \
3620 WHERE instance.job_name = $1 \
3621 AND instance.hold_actor IS NULL \
3622 AND execution.status = ANY($2) \
3623 AND execution.updated_at < to_timestamp($3::double precision / 1000.0) \
3624 AND NOT EXISTS ( \
3625 SELECT 1 FROM oxide_batch.ob_job_execution sibling \
3626 WHERE sibling.job_instance_id = execution.job_instance_id \
3627 AND sibling.status IN ('STARTING', 'STARTED', 'STOPPING', 'UNKNOWN')) \
3628 ORDER BY execution.job_instance_id, execution.id \
3629 LIMIT $4",
3630 )
3631 .bind(request.job_name().as_str())
3632 .bind(&statuses)
3633 .bind(threshold)
3634 .bind(limit)
3635 .fetch_all(&mut **self.transaction()?)
3636 .await
3637 .map_err(|_| RepositoryError::Unavailable)?;
3638 let mut candidates = Vec::with_capacity(rows.len());
3639 for row in &rows {
3640 candidates.push(PurgeCandidate::new(
3641 JobInstanceId::new(read_u64(row, "job_instance_id")?)?,
3642 JobExecutionId::new(read_u64(row, "id")?)?,
3643 ExecutionVersion::new(read_u64(row, "version")?),
3644 ));
3645 }
3646 let counts = self.purge_counts(&candidates).await?;
3647 Ok(PurgeSurvey::new(candidates, counts))
3648 })
3649 }
3650
3651 #[allow(clippy::too_many_lines)]
3652 fn apply_purge<'a>(
3653 &'a mut self,
3654 plan: &'a PurgePlan,
3655 ) -> BoxFuture<'a, Result<PurgeCounts, RepositoryError>> {
3656 Box::pin(async move {
3657 let executions = candidate_execution_ids(plan.candidates())?;
3658 let instances = candidate_instance_ids(plan.candidates())?;
3659 let versions = plan
3660 .candidates()
3661 .iter()
3662 .map(|candidate| database_version(candidate.version()))
3663 .collect::<Result<Vec<_>, _>>()?;
3664 let statuses = plan
3665 .request()
3666 .statuses()
3667 .iter()
3668 .map(|status| status.as_str().to_owned())
3669 .collect::<Vec<_>>();
3670 let confirmed = sqlx::query(
3671 "SELECT count(*) AS matched \
3672 FROM unnest($1::bigint[], $2::bigint[]) AS candidate(id, version) \
3673 JOIN oxide_batch.ob_job_execution execution \
3674 ON execution.id = candidate.id AND execution.version = candidate.version \
3675 JOIN oxide_batch.ob_job_instance instance \
3676 ON instance.id = execution.job_instance_id \
3677 WHERE instance.hold_actor IS NULL \
3678 AND execution.status = ANY($3) \
3679 AND NOT EXISTS ( \
3680 SELECT 1 FROM oxide_batch.ob_job_execution sibling \
3681 WHERE sibling.job_instance_id = execution.job_instance_id \
3682 AND sibling.status IN ('STARTING', 'STARTED', 'STOPPING', 'UNKNOWN'))",
3683 )
3684 .bind(&executions)
3685 .bind(&versions)
3686 .bind(&statuses)
3687 .fetch_one(&mut **self.transaction()?)
3688 .await
3689 .map_err(|_| RepositoryError::Unavailable)?;
3690 if read_u64(&confirmed, "matched")?
3691 != u64::try_from(executions.len()).unwrap_or(u64::MAX)
3692 {
3693 return Err(RepositoryError::RetentionPlanStale);
3694 }
3695 sqlx::query(
3699 "UPDATE oxide_batch.ob_flow_decision SET reused_decision_id = NULL \
3700 WHERE reused_decision_id IN ( \
3701 SELECT id FROM oxide_batch.ob_flow_decision \
3702 WHERE job_execution_id = ANY($1)) \
3703 AND job_execution_id <> ALL($1)",
3704 )
3705 .bind(&executions)
3706 .execute(&mut **self.transaction()?)
3707 .await
3708 .map_err(|_| RepositoryError::Unavailable)?;
3709 let flow_decisions = self
3710 .purge_delete(
3711 "DELETE FROM oxide_batch.ob_flow_decision WHERE job_execution_id = ANY($1)",
3712 &executions,
3713 )
3714 .await?;
3715 let recovery_decisions = self
3716 .purge_delete(
3717 "DELETE FROM oxide_batch.ob_recovery_decision WHERE job_execution_id = ANY($1)",
3718 &executions,
3719 )
3720 .await?;
3721 let operator_requests = self
3722 .purge_delete(
3723 "DELETE FROM oxide_batch.ob_operator_request WHERE job_execution_id = ANY($1)",
3724 &executions,
3725 )
3726 .await?;
3727 let step_partitions = self
3728 .purge_delete(
3729 "DELETE FROM oxide_batch.ob_step_partition WHERE step_execution_id IN ( \
3730 SELECT id FROM oxide_batch.ob_step_execution \
3731 WHERE job_execution_id = ANY($1))",
3732 &executions,
3733 )
3734 .await?;
3735 let step_executions = self
3736 .purge_delete(
3737 "DELETE FROM oxide_batch.ob_step_execution WHERE job_execution_id = ANY($1)",
3738 &executions,
3739 )
3740 .await?;
3741 let job_executions = self
3742 .purge_delete(
3743 "DELETE FROM oxide_batch.ob_job_execution WHERE id = ANY($1)",
3744 &executions,
3745 )
3746 .await?;
3747 let orphaned_requests = self
3748 .purge_delete(
3749 "DELETE FROM oxide_batch.ob_operator_request \
3750 WHERE job_execution_id IS NULL AND job_instance_id = ANY($1) \
3751 AND NOT EXISTS ( \
3752 SELECT 1 FROM oxide_batch.ob_job_execution execution \
3753 WHERE execution.job_instance_id \
3754 = oxide_batch.ob_operator_request.job_instance_id)",
3755 &instances,
3756 )
3757 .await?;
3758 sqlx::query(
3761 "UPDATE oxide_batch.ob_retention_action SET job_instance_id = NULL \
3762 WHERE job_instance_id = ANY($1) AND NOT EXISTS ( \
3763 SELECT 1 FROM oxide_batch.ob_job_execution execution \
3764 WHERE execution.job_instance_id \
3765 = oxide_batch.ob_retention_action.job_instance_id)",
3766 )
3767 .bind(&instances)
3768 .execute(&mut **self.transaction()?)
3769 .await
3770 .map_err(|_| RepositoryError::Unavailable)?;
3771 let job_instances = self
3772 .purge_delete(
3773 "DELETE FROM oxide_batch.ob_job_instance WHERE id = ANY($1) \
3774 AND NOT EXISTS ( \
3775 SELECT 1 FROM oxide_batch.ob_job_execution execution \
3776 WHERE execution.job_instance_id = oxide_batch.ob_job_instance.id)",
3777 &instances,
3778 )
3779 .await?;
3780 Ok(PurgeCounts::new(
3781 flow_decisions,
3782 recovery_decisions,
3783 operator_requests.saturating_add(orphaned_requests),
3784 step_partitions,
3785 step_executions,
3786 job_executions,
3787 job_instances,
3788 ))
3789 })
3790 }
3791
3792 fn commit<'a>(mut self: Box<Self>) -> BoxFuture<'a, Result<(), RepositoryError>>
3793 where
3794 Self: 'a,
3795 {
3796 Box::pin(async move {
3797 let connection = self.connection.take().ok_or(RepositoryError::Unavailable)?;
3798 commit_postgres_connection(connection)
3799 .await
3800 .map_err(|()| RepositoryError::CommitOutcomeUnknown)
3801 })
3802 }
3803
3804 fn rollback<'a>(mut self: Box<Self>) -> BoxFuture<'a, Result<(), RepositoryError>>
3805 where
3806 Self: 'a,
3807 {
3808 Box::pin(async move {
3809 let mut connection = self.connection.take().ok_or(RepositoryError::Unavailable)?;
3810 if sqlx::query("ROLLBACK")
3811 .execute(&mut *connection)
3812 .await
3813 .is_err()
3814 {
3815 connection.close_on_drop();
3816 return Err(RepositoryError::Unavailable);
3817 }
3818 Ok(())
3819 })
3820 }
3821}
3822
3823impl Drop for PostgresUnitOfWork<'_> {
3824 fn drop(&mut self) {
3825 if let Some(connection) = &mut self.connection {
3826 connection.close_on_drop();
3827 }
3828 }
3829}
3830
3831struct PostgresChunkTransaction {
3832 connection: Option<PoolConnection<Postgres>>,
3833 context: ChunkTransactionContext,
3834 expected_version: ExecutionVersion,
3835 committed_counts: ExecutionCounts,
3836 clock: Arc<dyn Clock>,
3837 state_provider: Arc<dyn PostgresChunkStateProvider>,
3838}
3839
3840impl PostgresChunkTransaction {
3841 fn connection(&mut self) -> Result<&mut PoolConnection<Postgres>, ChunkTransactionError> {
3842 self.connection
3843 .as_mut()
3844 .ok_or(ChunkTransactionError::NotCommitted)
3845 }
3846
3847 fn discard_connection(&mut self) {
3848 if let Some(connection) = &mut self.connection {
3849 connection.close_on_drop();
3850 }
3851 }
3852}
3853
3854impl BusinessTransaction for PostgresChunkTransaction {
3855 fn execute<'a>(
3856 &'a mut self,
3857 statement: BusinessStatement<'a>,
3858 ) -> BoxFuture<'a, Result<BusinessWriteResult, BusinessTransactionError>> {
3859 Box::pin(async move {
3860 let statement_text = String::from(statement.text());
3861 let mut query = sqlx::query(AssertSqlSafe(statement_text));
3862 for value in statement.values() {
3863 query = match value.kind() {
3864 BusinessValueKind::Text => {
3865 query.bind(value.as_text().ok_or(BusinessTransactionError::Rejected)?)
3866 }
3867 BusinessValueKind::Bytes => {
3868 query.bind(value.as_bytes().ok_or(BusinessTransactionError::Rejected)?)
3869 }
3870 BusinessValueKind::I64 => {
3871 query.bind(value.as_i64().ok_or(BusinessTransactionError::Rejected)?)
3872 }
3873 BusinessValueKind::Bool => {
3874 query.bind(value.as_bool().ok_or(BusinessTransactionError::Rejected)?)
3875 }
3876 BusinessValueKind::Null => query.bind(Option::<String>::None),
3877 };
3878 }
3879 match query
3880 .execute(
3881 &mut **self
3882 .connection()
3883 .map_err(|_| BusinessTransactionError::Infrastructure)?,
3884 )
3885 .await
3886 {
3887 Ok(result) => Ok(BusinessWriteResult::new(result.rows_affected())),
3888 Err(error) => {
3889 let classified = classify_business_error(&error);
3890 if classified != BusinessTransactionError::Rejected {
3891 self.discard_connection();
3892 }
3893 Err(classified)
3894 }
3895 }
3896 })
3897 }
3898}
3899
3900impl ChunkTransaction for PostgresChunkTransaction {
3901 fn business_transaction(&mut self) -> Option<&mut dyn BusinessTransaction> {
3902 Some(self)
3903 }
3904
3905 #[allow(
3906 clippy::too_many_lines,
3907 reason = "the atomic business/progress bind and commit boundary remains visible"
3908 )]
3909 fn commit(
3910 &mut self,
3911 counts: ChunkCounts,
3912 fault: ChunkFaultProgress,
3913 ) -> BoxFuture<'_, Result<ChunkCommitReceipt, ChunkTransactionError>> {
3914 Box::pin(async move {
3915 let next_counts = add_chunk_counts(self.committed_counts, counts)?;
3916 let empty_state = FaultStateEnvelope::empty();
3917 let empty_payload = durable_fault_payload(&empty_state)?;
3918 let empty_checksum = empty_state
3919 .checksum()
3920 .map_err(|_| ChunkTransactionError::NotCommitted)?;
3921 let receipt = catch_unwind(AssertUnwindSafe(|| {
3922 self.state_provider
3923 .state_for_commit(self.committed_counts, counts)
3924 }))
3925 .ok()
3926 .and_then(Result::ok)
3927 .ok_or(ChunkTransactionError::NotCommitted)?;
3928 let checkpoint_payload = durable_payload(receipt.checkpoint())?;
3929 let context_payload = durable_payload(receipt.execution_context())?;
3930 let next_version = self
3931 .expected_version
3932 .get()
3933 .checked_add(1)
3934 .map(ExecutionVersion::new)
3935 .ok_or(ChunkTransactionError::NotCommitted)?;
3936 let result = sqlx::query(
3937 "UPDATE oxide_batch.ob_step_execution SET \
3938 read_count = $1, processed_count = $2, write_count = $3, \
3939 filter_count = $4, commit_count = $5, rollback_count = $6, \
3940 checkpoint_format = $7, checkpoint_schema = $8, \
3941 checkpoint_schema_version = $9, checkpoint_payload = $10, \
3942 context_format = $11, context_schema = $12, \
3943 context_schema_version = $13, context_payload = $14, \
3944 read_skip_count = read_skip_count + $20, \
3945 process_skip_count = process_skip_count + $21, \
3946 write_skip_count = write_skip_count + $22, \
3947 no_rollback_count = no_rollback_count + $23, \
3948 fault_state_format = $24, fault_state_schema = $25, \
3949 fault_state_schema_version = $26, fault_state_payload = $27, \
3950 fault_state_checksum = $28, \
3951 updated_at = to_timestamp($15::double precision / 1000.0), version = $16 \
3952 WHERE id = $17 AND job_execution_id = $18 \
3953 AND version = $19 AND status = 'STARTED'",
3954 )
3955 .bind(chunk_database_count(next_counts.read())?)
3956 .bind(chunk_database_count(next_counts.processed())?)
3957 .bind(chunk_database_count(next_counts.written())?)
3958 .bind(chunk_database_count(next_counts.filtered())?)
3959 .bind(chunk_database_count(next_counts.committed())?)
3960 .bind(chunk_database_count(next_counts.rolled_back())?)
3961 .bind(
3962 i16::try_from(receipt.checkpoint().format_version())
3963 .map_err(|_| ChunkTransactionError::NotCommitted)?,
3964 )
3965 .bind(receipt.checkpoint().schema_id().as_str())
3966 .bind(
3967 i32::try_from(receipt.checkpoint().schema_version().get())
3968 .map_err(|_| ChunkTransactionError::NotCommitted)?,
3969 )
3970 .bind(Json(checkpoint_payload))
3971 .bind(
3972 i16::try_from(receipt.execution_context().format_version())
3973 .map_err(|_| ChunkTransactionError::NotCommitted)?,
3974 )
3975 .bind(receipt.execution_context().schema_id().as_str())
3976 .bind(
3977 i32::try_from(receipt.execution_context().schema_version().get())
3978 .map_err(|_| ChunkTransactionError::NotCommitted)?,
3979 )
3980 .bind(Json(context_payload))
3981 .bind(
3982 system_time_millis(self.clock.now())
3983 .map_err(|_| ChunkTransactionError::NotCommitted)?,
3984 )
3985 .bind(database_version(next_version).map_err(|_| ChunkTransactionError::NotCommitted)?)
3986 .bind(
3987 database_id(
3988 self.context.step_execution_id().get(),
3989 IdentifierKind::StepExecution,
3990 )
3991 .map_err(|_| ChunkTransactionError::NotCommitted)?,
3992 )
3993 .bind(
3994 database_id(
3995 self.context.job_execution_id().get(),
3996 IdentifierKind::JobExecution,
3997 )
3998 .map_err(|_| ChunkTransactionError::NotCommitted)?,
3999 )
4000 .bind(
4001 database_version(self.expected_version)
4002 .map_err(|_| ChunkTransactionError::NotCommitted)?,
4003 )
4004 .bind(chunk_database_count(fault.skips().read())?)
4005 .bind(chunk_database_count(fault.skips().process())?)
4006 .bind(chunk_database_count(fault.skips().write())?)
4007 .bind(chunk_database_count(fault.no_rollbacks())?)
4008 .bind(
4009 i16::try_from(FaultStateEnvelope::FORMAT_VERSION)
4010 .map_err(|_| ChunkTransactionError::NotCommitted)?,
4011 )
4012 .bind(FaultStateEnvelope::FORMAT)
4013 .bind(
4014 i32::try_from(FaultStateEnvelope::SCHEMA_VERSION)
4015 .map_err(|_| ChunkTransactionError::NotCommitted)?,
4016 )
4017 .bind(Json(empty_payload))
4018 .bind(empty_checksum.as_slice())
4019 .execute(&mut **self.connection()?)
4020 .await;
4021
4022 let Ok(result) = result else {
4023 rollback_chunk_transaction(&mut self.connection).await;
4024 return Err(ChunkTransactionError::NotCommitted);
4025 };
4026 let affected = result.rows_affected();
4027 if affected != 1 {
4028 rollback_chunk_transaction(&mut self.connection).await;
4029 return Err(ChunkTransactionError::NotCommitted);
4030 }
4031
4032 let connection = self
4033 .connection
4034 .take()
4035 .ok_or(ChunkTransactionError::NotCommitted)?;
4036 commit_postgres_connection(connection)
4037 .await
4038 .map_err(|()| ChunkTransactionError::CommitOutcomeUnknown)?;
4039 self.expected_version = next_version;
4040 self.committed_counts = next_counts;
4041 Ok(receipt)
4042 })
4043 }
4044
4045 fn rollback(&mut self) -> BoxFuture<'_, Result<(), ChunkTransactionError>> {
4046 Box::pin(async move {
4047 let Some(mut connection) = self.connection.take() else {
4048 return Ok(());
4049 };
4050 if sqlx::query("ROLLBACK")
4051 .execute(&mut *connection)
4052 .await
4053 .is_err()
4054 {
4055 connection.close_on_drop();
4056 }
4057 Ok(())
4058 })
4059 }
4060}
4061
4062impl Drop for PostgresChunkTransaction {
4063 fn drop(&mut self) {
4064 if let Some(connection) = &mut self.connection {
4065 connection.close_on_drop();
4066 }
4067 }
4068}
4069
4070fn classify_business_error(error: &sqlx::Error) -> BusinessTransactionError {
4071 let Some(database) = error.as_database_error() else {
4072 return BusinessTransactionError::Infrastructure;
4073 };
4074 let Some(code) = database.code() else {
4075 return BusinessTransactionError::Infrastructure;
4076 };
4077 if code == "57014" {
4078 return BusinessTransactionError::Cancelled;
4079 }
4080 if code.starts_with("22") || code.starts_with("23") || code.starts_with("42") {
4081 return BusinessTransactionError::Rejected;
4082 }
4083 BusinessTransactionError::Infrastructure
4084}
4085
4086fn add_chunk_counts(
4087 current: ExecutionCounts,
4088 chunk: ChunkCounts,
4089) -> Result<ExecutionCounts, ChunkTransactionError> {
4090 Ok(ExecutionCounts::new(
4091 current
4092 .read()
4093 .checked_add(chunk.read().get())
4094 .ok_or(ChunkTransactionError::NotCommitted)?,
4095 current
4096 .processed()
4097 .checked_add(chunk.processed().get())
4098 .ok_or(ChunkTransactionError::NotCommitted)?,
4099 current
4100 .written()
4101 .checked_add(chunk.written().get())
4102 .ok_or(ChunkTransactionError::NotCommitted)?,
4103 current
4104 .filtered()
4105 .checked_add(chunk.filtered().get())
4106 .ok_or(ChunkTransactionError::NotCommitted)?,
4107 current
4108 .committed()
4109 .checked_add(1)
4110 .ok_or(ChunkTransactionError::NotCommitted)?,
4111 current.rolled_back(),
4112 ))
4113}
4114
4115fn chunk_database_count(value: u64) -> Result<i64, ChunkTransactionError> {
4116 i64::try_from(value).map_err(|_| ChunkTransactionError::NotCommitted)
4117}
4118
4119fn durable_fault_payload(state: &FaultStateEnvelope) -> Result<Value, ChunkTransactionError> {
4120 let bytes = state
4121 .to_canonical_json()
4122 .map_err(|_| ChunkTransactionError::NotCommitted)?;
4123 serde_json::from_slice(&bytes).map_err(|_| ChunkTransactionError::NotCommitted)
4124}
4125
4126fn durable_payload(state: &impl DurablePayload) -> Result<Value, ChunkTransactionError> {
4127 let bytes = state
4128 .payload_json()
4129 .map_err(|_| ChunkTransactionError::NotCommitted)?;
4130 serde_json::from_slice(&bytes).map_err(|_| ChunkTransactionError::NotCommitted)
4131}
4132
4133trait DurablePayload {
4134 fn payload_json(&self) -> Result<Vec<u8>, crate::StateError>;
4135}
4136
4137impl DurablePayload for Checkpoint {
4138 fn payload_json(&self) -> Result<Vec<u8>, crate::StateError> {
4139 Checkpoint::payload_json(self)
4140 }
4141}
4142
4143impl DurablePayload for ExecutionContext {
4144 fn payload_json(&self) -> Result<Vec<u8>, crate::StateError> {
4145 ExecutionContext::payload_json(self)
4146 }
4147}
4148
4149async fn rollback_chunk_connection(connection: &mut PoolConnection<Postgres>) {
4150 if sqlx::query("ROLLBACK")
4151 .execute(&mut **connection)
4152 .await
4153 .is_err()
4154 {
4155 connection.close_on_drop();
4156 }
4157}
4158
4159async fn rollback_recovery_savepoint(connection: &mut PgConnection) {
4160 let _ = sqlx::query("ROLLBACK TO SAVEPOINT ob_recovery_decision")
4161 .execute(connection)
4162 .await;
4163}
4164
4165async fn rollback_chunk_transaction(connection: &mut Option<PoolConnection<Postgres>>) {
4166 let Some(mut connection) = connection.take() else {
4167 return;
4168 };
4169 rollback_chunk_connection(&mut connection).await;
4170}
4171
4172async fn commit_postgres_connection(mut connection: PoolConnection<Postgres>) -> Result<(), ()> {
4173 if sqlx::query("COMMIT")
4174 .execute(&mut *connection)
4175 .await
4176 .is_err()
4177 {
4178 connection.close_on_drop();
4179 return Err(());
4180 }
4181 Ok(())
4182}
4183
4184async fn configure_transaction(
4185 connection: &mut PoolConnection<Postgres>,
4186 config: &PostgresConfig,
4187) -> Result<(), RepositoryError> {
4188 for (name, value) in [
4189 (
4190 "statement_timeout",
4191 duration_millis(config.statement_timeout)?,
4192 ),
4193 ("lock_timeout", duration_millis(config.lock_timeout)?),
4194 (
4195 "idle_in_transaction_session_timeout",
4196 duration_millis(config.idle_transaction_timeout)?,
4197 ),
4198 ] {
4199 sqlx::query("SELECT set_config($1, $2, true)")
4200 .bind(name)
4201 .bind(value.to_string())
4202 .execute(&mut **connection)
4203 .await
4204 .map_err(|_| RepositoryError::Unavailable)?;
4205 }
4206 Ok(())
4207}
4208
4209async fn read_schema_version<'executor, E>(executor: E) -> Result<u32, RepositoryError>
4210where
4211 E: sqlx::Executor<'executor, Database = Postgres>,
4212{
4213 let version = sqlx::query_scalar::<_, i32>(
4214 "SELECT version FROM oxide_batch.ob_schema_version WHERE singleton = TRUE",
4215 )
4216 .fetch_optional(executor)
4217 .await
4218 .map_err(|error| {
4219 if error
4220 .as_database_error()
4221 .and_then(sqlx::error::DatabaseError::code)
4222 .as_deref()
4223 .is_some_and(|code| code == "42P01" || code == "3F000")
4224 {
4225 RepositoryError::SchemaUninitialized
4226 } else {
4227 RepositoryError::Unavailable
4228 }
4229 })?
4230 .ok_or(RepositoryError::SchemaUninitialized)?;
4231 u32::try_from(version).map_err(|_| RepositoryError::Unavailable)
4232}
4233
4234fn verify_schema_version(current: u32) -> Result<(), RepositoryError> {
4235 match current.cmp(&SUPPORTED_SCHEMA_VERSION) {
4236 std::cmp::Ordering::Less => Err(RepositoryError::MigrationRequired {
4237 current,
4238 supported: SUPPORTED_SCHEMA_VERSION,
4239 }),
4240 std::cmp::Ordering::Equal => Ok(()),
4241 std::cmp::Ordering::Greater => Err(RepositoryError::NewerSchema {
4242 current,
4243 supported: SUPPORTED_SCHEMA_VERSION,
4244 }),
4245 }
4246}
4247
4248fn duration_millis(duration: Duration) -> Result<i64, RepositoryError> {
4249 i64::try_from(duration.as_millis()).map_err(|_| RepositoryError::Unavailable)
4250}
4251
4252fn database_id(value: u64, kind: IdentifierKind) -> Result<i64, RepositoryError> {
4253 i64::try_from(value).map_err(|_| RepositoryError::IdentifierOutOfRange { kind, value })
4254}
4255
4256fn system_time_millis(value: SystemTime) -> Result<i64, RepositoryError> {
4257 let duration = value
4258 .duration_since(UNIX_EPOCH)
4259 .map_err(|_| RepositoryError::Unavailable)?;
4260 i64::try_from(duration.as_millis()).map_err(|_| RepositoryError::Unavailable)
4261}
4262
4263fn millis_system_time(value: i64) -> Result<SystemTime, RepositoryError> {
4264 let value = u64::try_from(value).map_err(|_| RepositoryError::Unavailable)?;
4265 Ok(UNIX_EPOCH + Duration::from_millis(value))
4266}
4267
4268fn starting_metadata(created_at: SystemTime) -> Result<ExecutionMetadata, RepositoryError> {
4269 Ok(ExecutionMetadata::new(
4270 BatchStatus::Starting,
4271 ExitStatus::unknown(),
4272 ExecutionTimestamps::new(created_at, None, None)?,
4273 ExecutionCounts::default(),
4274 None,
4275 )?)
4276}
4277
4278fn encode_identifying_parameters(key: &JobInstanceKey) -> Result<Value, RepositoryError> {
4279 let mut object = Map::new();
4280 for (name, kind) in key.identifying_fields() {
4281 let parameter = key
4282 .identifying_value(name)
4283 .ok_or(RepositoryError::Unavailable)?;
4284 let (kind_name, value) = match kind {
4285 ParameterValueKind::String => (
4286 "string",
4287 Value::String(
4288 parameter
4289 .as_str()
4290 .ok_or(RepositoryError::Unavailable)?
4291 .to_owned(),
4292 ),
4293 ),
4294 ParameterValueKind::I64 => (
4295 "i64",
4296 Value::Number(
4297 parameter
4298 .as_i64()
4299 .ok_or(RepositoryError::Unavailable)?
4300 .into(),
4301 ),
4302 ),
4303 ParameterValueKind::U64 => (
4304 "u64",
4305 Value::Number(
4306 parameter
4307 .as_u64()
4308 .ok_or(RepositoryError::Unavailable)?
4309 .into(),
4310 ),
4311 ),
4312 ParameterValueKind::Bool => (
4313 "bool",
4314 Value::Bool(parameter.as_bool().ok_or(RepositoryError::Unavailable)?),
4315 ),
4316 _ => return Err(RepositoryError::Unavailable),
4320 };
4321 object.insert(
4322 name.as_str().to_owned(),
4323 json!({"type": kind_name, "identifying": true, "value": value}),
4324 );
4325 }
4326 let value = Value::Object(object);
4327 let size = serde_json::to_vec(&value)
4328 .map_err(|_| RepositoryError::Unavailable)?
4329 .len();
4330 if size > MAX_INSTANCE_KEY_INPUT {
4331 return Err(RepositoryError::Unavailable);
4332 }
4333 Ok(value)
4334}
4335
4336fn decode_job_instance(row: &PgRow) -> Result<JobInstance, RepositoryError> {
4337 let id = row
4338 .try_get::<i64, _>("id")
4339 .map_err(|_| RepositoryError::Unavailable)?;
4340 let job_name = row
4341 .try_get::<String, _>("job_name")
4342 .map_err(|_| RepositoryError::Unavailable)?;
4343 let parameters = row
4344 .try_get::<Json<Value>, _>("identifying_parameters")
4345 .map_err(|_| RepositoryError::Unavailable)?;
4346 let parameters = decode_identifying_parameters(¶meters.0)?;
4347 let id = JobInstanceId::new(u64::try_from(id).map_err(|_| RepositoryError::Unavailable)?)?;
4348 let key = JobInstanceKey::new(JobName::new(job_name)?, ¶meters);
4349 Ok(JobInstance::new(id, key))
4350}
4351
4352fn decode_identifying_parameters(value: &Value) -> Result<JobParameters, RepositoryError> {
4353 let object = value.as_object().ok_or(RepositoryError::Unavailable)?;
4354 let mut parameters = JobParameters::new();
4355 for (raw_name, envelope) in object {
4356 let envelope = envelope.as_object().ok_or(RepositoryError::Unavailable)?;
4357 if envelope.get("identifying") != Some(&Value::Bool(true)) {
4358 return Err(RepositoryError::Unavailable);
4359 }
4360 let kind = envelope
4361 .get("type")
4362 .and_then(Value::as_str)
4363 .ok_or(RepositoryError::Unavailable)?;
4364 let raw_value = envelope.get("value").ok_or(RepositoryError::Unavailable)?;
4365 let value = match kind {
4366 "string" => {
4367 ParameterValue::string(raw_value.as_str().ok_or(RepositoryError::Unavailable)?)?
4368 }
4369 "i64" => ParameterValue::from(raw_value.as_i64().ok_or(RepositoryError::Unavailable)?),
4370 "u64" => ParameterValue::from(raw_value.as_u64().ok_or(RepositoryError::Unavailable)?),
4371 "bool" => {
4372 ParameterValue::from(raw_value.as_bool().ok_or(RepositoryError::Unavailable)?)
4373 }
4374 _ => return Err(RepositoryError::Unavailable),
4375 };
4376 parameters.insert(
4377 ParameterName::new(raw_name.clone())?,
4378 JobParameter::new(value, ParameterRole::Identifying),
4379 )?;
4380 }
4381 Ok(parameters)
4382}
4383
4384async fn ensure_definition(
4385 transaction: &mut PgConnection,
4386 job_name: &str,
4387 definition: &DefinitionIdentity,
4388 registered_at: SystemTime,
4389) -> Result<i64, RepositoryError> {
4390 let expected_job_name = JobName::new(job_name.to_owned())?;
4391 if let Some(actual) = definition.job_name()
4392 && actual != &expected_job_name
4393 {
4394 return Err(RepositoryError::DefinitionJobMismatch {
4395 expected: expected_job_name,
4396 actual: actual.clone(),
4397 });
4398 }
4399 oxide_batch_core::check_manifest_format(definition.manifest_format()).map_err(|_| {
4400 RepositoryError::UnsupportedManifestVersion {
4401 format: definition.manifest_format(),
4402 }
4403 })?;
4404 let manifest: Value = serde_json::from_slice(definition.canonical_manifest())
4405 .map_err(|_| RepositoryError::Unavailable)?;
4406 let registered_ms = system_time_millis(registered_at)?;
4407 sqlx::query(
4408 "INSERT INTO oxide_batch.ob_job_definition \
4409 (job_name, definition_revision, manifest_format, manifest_digest, manifest, registered_at) \
4410 VALUES ($1, $2, $3, $4, $5, to_timestamp($6::double precision / 1000.0)) \
4411 ON CONFLICT DO NOTHING",
4412 )
4413 .bind(job_name)
4414 .bind(definition.revision().as_str())
4415 .bind(i16::try_from(definition.manifest_format()).map_err(|_| {
4416 RepositoryError::UnsupportedManifestVersion {
4417 format: definition.manifest_format(),
4418 }
4419 })?)
4420 .bind(&definition.manifest_digest()[..])
4421 .bind(Json(manifest))
4422 .bind(registered_ms)
4423 .execute(&mut *transaction)
4424 .await
4425 .map_err(|_| RepositoryError::Unavailable)?;
4426 let id = sqlx::query_scalar(
4427 "SELECT id FROM oxide_batch.ob_job_definition \
4428 WHERE job_name = $1 AND manifest_digest = $2",
4429 )
4430 .bind(job_name)
4431 .bind(&definition.manifest_digest()[..])
4432 .fetch_optional(&mut *transaction)
4433 .await
4434 .map_err(|_| RepositoryError::Unavailable)?;
4435 match id {
4436 Some(id) => Ok(id),
4437 None => Err(RepositoryError::DefinitionDrift {
4438 job_name: expected_job_name,
4439 revision: definition.revision().clone(),
4440 }),
4441 }
4442}
4443
4444fn job_execution_select(suffix: &str) -> String {
4445 format!(
4446 "SELECT execution.id, execution.job_instance_id, execution.status, \
4447 execution.exit_code, 0::bigint AS read_count, 0::bigint AS processed_count, \
4448 0::bigint AS write_count, 0::bigint AS filter_count, \
4449 0::bigint AS commit_count, 0::bigint AS rollback_count, \
4450 execution.failure_category, execution.failure_id, \
4451 (extract(epoch FROM execution.created_at) * 1000)::bigint AS created_ms, \
4452 (extract(epoch FROM execution.started_at) * 1000)::bigint AS started_ms, \
4453 (extract(epoch FROM execution.ended_at) * 1000)::bigint AS ended_ms, \
4454 execution.version FROM oxide_batch.ob_job_execution execution {suffix}"
4455 )
4456}
4457
4458fn step_execution_select(suffix: &str) -> String {
4459 format!(
4460 "SELECT execution.id, execution.job_execution_id, execution.step_name, \
4461 execution.status, execution.exit_code, execution.read_count, \
4462 execution.processed_count, execution.write_count, execution.filter_count, \
4463 execution.commit_count, execution.rollback_count, execution.failure_category, \
4464 execution.failure_id, \
4465 (extract(epoch FROM execution.created_at) * 1000)::bigint AS created_ms, \
4466 (extract(epoch FROM execution.started_at) * 1000)::bigint AS started_ms, \
4467 (extract(epoch FROM execution.ended_at) * 1000)::bigint AS ended_ms, \
4468 execution.version FROM oxide_batch.ob_step_execution execution {suffix}"
4469 )
4470}
4471
4472fn durable_step_select(suffix: &str) -> String {
4473 format!(
4474 "SELECT execution.id, execution.job_execution_id, execution.step_name, \
4475 execution.status, execution.exit_code, execution.read_count, \
4476 execution.processed_count, execution.write_count, execution.filter_count, \
4477 execution.commit_count, execution.rollback_count, execution.failure_category, \
4478 execution.failure_id, execution.checkpoint_format, execution.checkpoint_schema, \
4479 execution.checkpoint_schema_version, execution.checkpoint_payload, \
4480 execution.context_format, execution.context_schema, \
4481 execution.context_schema_version, execution.context_payload, \
4482 execution.read_retry_count, execution.process_retry_count, \
4483 execution.write_retry_count, execution.read_skip_count, \
4484 execution.process_skip_count, execution.write_skip_count, \
4485 execution.no_rollback_count, execution.fault_state_format, \
4486 execution.fault_state_schema, execution.fault_state_schema_version, \
4487 execution.fault_state_payload, execution.fault_state_checksum, \
4488 (extract(epoch FROM execution.created_at) * 1000)::bigint AS created_ms, \
4489 (extract(epoch FROM execution.started_at) * 1000)::bigint AS started_ms, \
4490 (extract(epoch FROM execution.ended_at) * 1000)::bigint AS ended_ms, \
4491 execution.version FROM oxide_batch.ob_step_execution execution {suffix}"
4492 )
4493}
4494
4495fn flow_decision_select(suffix: &str) -> String {
4496 format!(
4497 "SELECT decision.id, decision.job_execution_id, \
4498 decision.source_step_execution_id, decision.reused_decision_id, \
4499 decision.sequence, decision.source_node_id, decision.observed_outcome, \
4500 decision.target_node_id, decision.transition_kind, decision.terminal_kind, \
4501 decision.plan_fingerprint, decision.input_digest, \
4502 (extract(epoch FROM decision.decided_at) * 1000)::bigint AS decided_ms \
4503 FROM oxide_batch.ob_flow_decision decision {suffix}"
4504 )
4505}
4506
4507fn partition_select(suffix: &str) -> String {
4508 format!(
4509 "SELECT partition.id, partition.step_execution_id, \
4510 partition.worker_step_execution_id, partition.partition_key, \
4511 partition.partition_ordinal, partition.status, partition.exit_code, \
4512 partition.read_count, partition.processed_count, partition.write_count, \
4513 partition.filter_count, partition.commit_count, partition.rollback_count, \
4514 partition.context_format, partition.context_schema, \
4515 partition.context_schema_version, partition.context_payload, \
4516 partition.context_checksum, partition.version \
4517 FROM oxide_batch.ob_step_partition partition {suffix}"
4518 )
4519}
4520
4521fn decode_step_partition(row: &PgRow) -> Result<StepPartition, RepositoryError> {
4522 let id = StepPartitionId::new(read_u64(row, "id")?)?;
4523 let step_execution_id = StepExecutionId::new(read_u64(row, "step_execution_id")?)?;
4524 let worker_step_execution_id = read_optional_u64(row, "worker_step_execution_id")?
4525 .map(StepExecutionId::new)
4526 .transpose()?;
4527 let key = PartitionKey::new(read_text(row, "partition_key")?)
4528 .map_err(|_| RepositoryError::PartitionStateCorrupt)?;
4529 let ordinal = row
4530 .try_get::<i32, _>("partition_ordinal")
4531 .map_err(|_| RepositoryError::PartitionStateCorrupt)
4532 .and_then(|value| {
4533 u32::try_from(value).map_err(|_| RepositoryError::PartitionStateCorrupt)
4534 })?;
4535 if ordinal == 0 || ordinal > u32::from(MAX_PARTITIONS) {
4536 return Err(RepositoryError::PartitionStateCorrupt);
4537 }
4538 let status = decode_status(&read_text(row, "status")?)?;
4539 let exit_status = ExitStatus::new(
4540 ExitCode::new(
4541 read_optional_text(row, "exit_code")?.unwrap_or_else(|| String::from("UNKNOWN")),
4542 )
4543 .map_err(|_| RepositoryError::PartitionStateCorrupt)?,
4544 );
4545 let counts = ExecutionCounts::new(
4546 read_u64(row, "read_count")?,
4547 read_u64(row, "processed_count")?,
4548 read_u64(row, "write_count")?,
4549 read_u64(row, "filter_count")?,
4550 read_u64(row, "commit_count")?,
4551 read_u64(row, "rollback_count")?,
4552 );
4553 let context = decode_partition_context(row)?;
4554 let version = ExecutionVersion::new(read_u64(row, "version")?);
4555 if matches!(status, BatchStatus::Starting) && worker_step_execution_id.is_some()
4556 || !matches!(status, BatchStatus::Starting) && worker_step_execution_id.is_none()
4557 {
4558 return Err(RepositoryError::PartitionStateCorrupt);
4559 }
4560 Ok(StepPartition::from_snapshot(
4561 id,
4562 step_execution_id,
4563 worker_step_execution_id,
4564 key,
4565 ordinal,
4566 status,
4567 exit_status,
4568 counts,
4569 context,
4570 version,
4571 ))
4572}
4573
4574fn decode_partition_context(row: &PgRow) -> Result<ExecutionContext, RepositoryError> {
4575 let format_version = u16::try_from(
4576 row.try_get::<i16, _>("context_format")
4577 .map_err(|_| RepositoryError::PartitionStateCorrupt)?,
4578 )
4579 .map_err(|_| RepositoryError::PartitionStateCorrupt)?;
4580 let schema = read_text(row, "context_schema")?;
4581 let schema_version = u32::try_from(
4582 row.try_get::<i32, _>("context_schema_version")
4583 .map_err(|_| RepositoryError::PartitionStateCorrupt)?,
4584 )
4585 .map_err(|_| RepositoryError::PartitionStateCorrupt)?;
4586 let Json(payload): Json<Value> = row
4587 .try_get("context_payload")
4588 .map_err(|_| RepositoryError::PartitionStateCorrupt)?;
4589 let envelope = json!({
4590 "format": "oxide-batch.execution-context",
4591 "format_version": format_version,
4592 "schema": schema,
4593 "schema_version": schema_version,
4594 "payload": payload,
4595 });
4596 let bytes =
4597 serde_json::to_vec(&envelope).map_err(|_| RepositoryError::PartitionStateCorrupt)?;
4598 let limits = StateLimits::new(MAX_PARTITION_CONTEXT_BYTES, 16)
4599 .map_err(|_| RepositoryError::PartitionStateCorrupt)?;
4600 let context = ExecutionContext::from_json(&bytes, limits)
4601 .map_err(|_| RepositoryError::PartitionStateCorrupt)?;
4602 let stored_checksum: [u8; 32] = row
4603 .try_get::<Vec<u8>, _>("context_checksum")
4604 .map_err(|_| RepositoryError::PartitionStateCorrupt)?
4605 .try_into()
4606 .map_err(|_| RepositoryError::PartitionStateCorrupt)?;
4607 let actual_checksum: [u8; 32] = Sha256::digest(
4608 context
4609 .to_json()
4610 .map_err(|_| RepositoryError::PartitionStateCorrupt)?,
4611 )
4612 .into();
4613 if stored_checksum != actual_checksum {
4614 return Err(RepositoryError::PartitionStateCorrupt);
4615 }
4616 Ok(context)
4617}
4618
4619fn partition_count(value: u64) -> Result<i64, RepositoryError> {
4620 i64::try_from(value).map_err(|_| RepositoryError::PartitionStateCorrupt)
4621}
4622
4623fn execution_count(value: u64) -> Result<i64, RepositoryError> {
4624 i64::try_from(value).map_err(|_| RepositoryError::Unavailable)
4625}
4626
4627fn map_partition_mutation(id: StepPartitionId, error: PartitionMutationError) -> RepositoryError {
4628 match error {
4629 PartitionMutationError::StaleVersion { expected, actual } => {
4630 RepositoryError::Lifecycle(LifecycleError::StaleVersion { expected, actual })
4631 }
4632 PartitionMutationError::InvalidState { status } => {
4633 RepositoryError::PartitionUpdateNotAllowed { id, status }
4634 }
4635 PartitionMutationError::VersionExhausted => RepositoryError::PartitionStateCorrupt,
4636 }
4637}
4638
4639fn flow_target_node(target: &FlowTarget) -> Option<&str> {
4640 match target {
4641 FlowTarget::Node(node) => Some(node.as_str()),
4642 FlowTarget::Terminal(_) => None,
4643 }
4644}
4645
4646fn flow_terminal_code(target: &FlowTarget) -> Result<Option<&'static str>, RepositoryError> {
4647 Ok(match target {
4648 FlowTarget::Node(_) => None,
4649 FlowTarget::Terminal(TerminalKind::Complete) => Some("COMPLETE"),
4650 FlowTarget::Terminal(TerminalKind::Fail) => Some("FAIL"),
4651 FlowTarget::Terminal(TerminalKind::Stop) => Some("STOP"),
4652 FlowTarget::Terminal(_) => return Err(RepositoryError::Unavailable),
4655 })
4656}
4657
4658fn decode_flow_decision(row: &PgRow) -> Result<FlowDecision, RepositoryError> {
4659 let id = FlowDecisionId::new(read_u64(row, "id")?)?;
4660 let job_execution_id = JobExecutionId::new(read_u64(row, "job_execution_id")?)?;
4661 let sequence = FlowDecisionSequence::new(read_u64(row, "sequence")?)
4662 .map_err(|_| RepositoryError::FlowStateCorrupt)?;
4663 let source_node_id = NodeId::new(
4664 row.try_get::<String, _>("source_node_id")
4665 .map_err(|_| RepositoryError::FlowStateCorrupt)?,
4666 )
4667 .map_err(|_| RepositoryError::FlowStateCorrupt)?;
4668 let source_step_execution_id = row
4669 .try_get::<Option<i64>, _>("source_step_execution_id")
4670 .map_err(|_| RepositoryError::FlowStateCorrupt)?
4671 .map(|value| {
4672 StepExecutionId::new(
4673 u64::try_from(value).map_err(|_| RepositoryError::FlowStateCorrupt)?,
4674 )
4675 .map_err(RepositoryError::from)
4676 })
4677 .transpose()?;
4678 let reused_decision_id = row
4679 .try_get::<Option<i64>, _>("reused_decision_id")
4680 .map_err(|_| RepositoryError::FlowStateCorrupt)?
4681 .map(|value| {
4682 FlowDecisionId::new(
4683 u64::try_from(value).map_err(|_| RepositoryError::FlowStateCorrupt)?,
4684 )
4685 .map_err(RepositoryError::from)
4686 })
4687 .transpose()?;
4688 let kind = FlowTransitionKind::from_durable_code(
4689 &row.try_get::<String, _>("transition_kind")
4690 .map_err(|_| RepositoryError::FlowStateCorrupt)?,
4691 )
4692 .ok_or(RepositoryError::FlowStateCorrupt)?;
4693 let observed_outcome = ExitCode::new(
4694 row.try_get::<String, _>("observed_outcome")
4695 .map_err(|_| RepositoryError::FlowStateCorrupt)?,
4696 )?;
4697 let target_node = row
4698 .try_get::<Option<String>, _>("target_node_id")
4699 .map_err(|_| RepositoryError::FlowStateCorrupt)?;
4700 let terminal = row
4701 .try_get::<Option<String>, _>("terminal_kind")
4702 .map_err(|_| RepositoryError::FlowStateCorrupt)?;
4703 let target = match (target_node, terminal.as_deref()) {
4704 (Some(node), None) => {
4705 FlowTarget::Node(NodeId::new(node).map_err(|_| RepositoryError::FlowStateCorrupt)?)
4706 }
4707 (None, Some("COMPLETE")) => FlowTarget::Terminal(TerminalKind::Complete),
4708 (None, Some("FAIL")) => FlowTarget::Terminal(TerminalKind::Fail),
4709 (None, Some("STOP")) => FlowTarget::Terminal(TerminalKind::Stop),
4710 _ => return Err(RepositoryError::FlowStateCorrupt),
4711 };
4712 let fingerprint: [u8; 32] = row
4713 .try_get::<Vec<u8>, _>("plan_fingerprint")
4714 .map_err(|_| RepositoryError::FlowStateCorrupt)?
4715 .try_into()
4716 .map_err(|_| RepositoryError::FlowStateCorrupt)?;
4717 let input_digest: [u8; 32] = row
4718 .try_get::<Vec<u8>, _>("input_digest")
4719 .map_err(|_| RepositoryError::FlowStateCorrupt)?
4720 .try_into()
4721 .map_err(|_| RepositoryError::FlowStateCorrupt)?;
4722 Ok(FlowDecision::new(
4723 id,
4724 job_execution_id,
4725 sequence,
4726 source_node_id,
4727 source_step_execution_id,
4728 kind,
4729 observed_outcome,
4730 target,
4731 fingerprint,
4732 input_digest,
4733 reused_decision_id,
4734 millis_system_time(read_i64(row, "decided_ms")?)?,
4735 ))
4736}
4737
4738fn decode_durable_step_state(row: &PgRow) -> Result<PostgresDurableStepState, RepositoryError> {
4739 let checkpoint = decode_durable_state(
4740 row,
4741 "checkpoint_format",
4742 "checkpoint_schema",
4743 "checkpoint_schema_version",
4744 "checkpoint_payload",
4745 "oxide-batch.checkpoint",
4746 Checkpoint::from_json,
4747 )?;
4748 let execution_context = decode_durable_state(
4749 row,
4750 "context_format",
4751 "context_schema",
4752 "context_schema_version",
4753 "context_payload",
4754 "oxide-batch.execution-context",
4755 ExecutionContext::from_json,
4756 )?;
4757 Ok(PostgresDurableStepState {
4758 step_execution: decode_step_execution(row)?,
4759 checkpoint,
4760 execution_context,
4761 fault_progress: decode_fault_progress(row)?,
4762 fault_state: decode_fault_state(row).map_err(|_| RepositoryError::FaultStateCorrupt)?,
4763 })
4764}
4765
4766fn decode_fault_progress(row: &PgRow) -> Result<FaultProgress, RepositoryError> {
4767 Ok(FaultProgress::new(
4768 RetryCounts::new(
4769 read_u64(row, "read_retry_count")?,
4770 read_u64(row, "process_retry_count")?,
4771 read_u64(row, "write_retry_count")?,
4772 ),
4773 SkipCounts::new(
4774 read_u64(row, "read_skip_count")?,
4775 read_u64(row, "process_skip_count")?,
4776 read_u64(row, "write_skip_count")?,
4777 ),
4778 read_u64(row, "rollback_count")?,
4779 read_u64(row, "no_rollback_count")?,
4780 ))
4781}
4782
4783fn decode_fault_state(row: &PgRow) -> Result<FaultStateEnvelope, FaultStateFormatError> {
4784 let format_version = u16::try_from(
4785 row.try_get::<i16, _>("fault_state_format")
4786 .map_err(|_| FaultStateFormatError::Malformed)?,
4787 )
4788 .map_err(|_| FaultStateFormatError::UnsupportedFormat)?;
4789 let schema: String = row
4790 .try_get("fault_state_schema")
4791 .map_err(|_| FaultStateFormatError::Malformed)?;
4792 let schema_version = u32::try_from(
4793 row.try_get::<i32, _>("fault_state_schema_version")
4794 .map_err(|_| FaultStateFormatError::Malformed)?,
4795 )
4796 .map_err(|_| FaultStateFormatError::UnsupportedSchemaVersion)?;
4797 let Json(payload): Json<Value> = row
4798 .try_get("fault_state_payload")
4799 .map_err(|_| FaultStateFormatError::Malformed)?;
4800 let checksum: Vec<u8> = row
4801 .try_get("fault_state_checksum")
4802 .map_err(|_| FaultStateFormatError::Malformed)?;
4803 let checksum: [u8; 32] = checksum
4804 .try_into()
4805 .map_err(|_| FaultStateFormatError::ChecksumMismatch)?;
4806 let bytes = canonical_fault_bytes(&payload)?;
4807 FaultStateEnvelope::from_canonical_json(
4808 format_version,
4809 &schema,
4810 schema_version,
4811 &bytes,
4812 &checksum,
4813 )
4814}
4815
4816fn canonical_fault_bytes(payload: &Value) -> Result<Vec<u8>, FaultStateFormatError> {
4821 let object = payload
4822 .as_object()
4823 .ok_or(FaultStateFormatError::Malformed)?;
4824 let mut canonical = serde_json::Map::new();
4825 for member in ["checkpoint", "entries"] {
4826 canonical.insert(
4827 String::from(member),
4828 object
4829 .get(member)
4830 .cloned()
4831 .ok_or(FaultStateFormatError::Malformed)?,
4832 );
4833 }
4834 serde_json::to_vec(&Value::Object(canonical)).map_err(|_| FaultStateFormatError::Malformed)
4835}
4836
4837#[allow(clippy::too_many_arguments)]
4838fn decode_durable_state<T>(
4839 row: &PgRow,
4840 format_column: &str,
4841 schema_column: &str,
4842 schema_version_column: &str,
4843 payload_column: &str,
4844 format: &str,
4845 decode: impl FnOnce(&[u8], StateLimits) -> Result<T, crate::StateError>,
4846) -> Result<T, RepositoryError> {
4847 let format_version = u16::try_from(
4848 row.try_get::<i16, _>(format_column)
4849 .map_err(|_| RepositoryError::Unavailable)?,
4850 )
4851 .map_err(|_| RepositoryError::Unavailable)?;
4852 let schema: String = row
4853 .try_get(schema_column)
4854 .map_err(|_| RepositoryError::Unavailable)?;
4855 let schema_version = u32::try_from(
4856 row.try_get::<i32, _>(schema_version_column)
4857 .map_err(|_| RepositoryError::Unavailable)?,
4858 )
4859 .map_err(|_| RepositoryError::Unavailable)?;
4860 let Json(payload): Json<Value> = row
4861 .try_get(payload_column)
4862 .map_err(|_| RepositoryError::Unavailable)?;
4863 let envelope = json!({
4864 "format": format,
4865 "format_version": format_version,
4866 "schema": schema,
4867 "schema_version": schema_version,
4868 "payload": payload,
4869 });
4870 let bytes = serde_json::to_vec(&envelope).map_err(|_| RepositoryError::Unavailable)?;
4871 let limits = StateLimits::new(1024 * 1024, 64).map_err(|_| RepositoryError::Unavailable)?;
4872 decode(&bytes, limits).map_err(|_| RepositoryError::Unavailable)
4873}
4874
4875fn decode_job_execution(row: &PgRow) -> Result<JobExecution, RepositoryError> {
4876 let id = JobExecutionId::new(read_u64(row, "id")?)?;
4877 let instance_id = JobInstanceId::new(read_u64(row, "job_instance_id")?)?;
4878 let metadata = decode_execution_metadata(row)?;
4879 let version = ExecutionVersion::new(read_u64(row, "version")?);
4880 Ok(JobExecution::from_snapshot(
4881 id,
4882 instance_id,
4883 metadata,
4884 version,
4885 ))
4886}
4887
4888fn decode_step_execution(row: &PgRow) -> Result<StepExecution, RepositoryError> {
4889 let id = StepExecutionId::new(read_u64(row, "id")?)?;
4890 let job_id = JobExecutionId::new(read_u64(row, "job_execution_id")?)?;
4891 let name = StepName::new(
4892 row.try_get::<String, _>("step_name")
4893 .map_err(|_| RepositoryError::Unavailable)?,
4894 )?;
4895 let metadata = decode_execution_metadata(row)?;
4896 let version = ExecutionVersion::new(read_u64(row, "version")?);
4897 Ok(StepExecution::from_snapshot(
4898 id, job_id, name, metadata, version,
4899 ))
4900}
4901
4902fn decode_recovery_decision(
4903 id: JobExecutionId,
4904 row: &PgRow,
4905) -> Result<RecoveryDecision, RepositoryError> {
4906 let digest = row
4907 .try_get::<Vec<u8>, _>("evidence_digest")
4908 .map_err(|_| RepositoryError::Unavailable)?;
4909 let evidence_digest: [u8; 32] = digest
4910 .try_into()
4911 .map_err(|_| RepositoryError::Unavailable)?;
4912 Ok(RecoveryDecision::new(
4913 RecoveryDecisionId::new(read_u64(row, "id")?)?,
4914 id,
4915 ExecutionVersion::new(read_u64(row, "execution_version")?),
4916 decode_status(
4917 &row.try_get::<String, _>("prior_status")
4918 .map_err(|_| RepositoryError::Unavailable)?,
4919 )?,
4920 decode_status(
4921 &row.try_get::<String, _>("resulting_status")
4922 .map_err(|_| RepositoryError::Unavailable)?,
4923 )?,
4924 row.try_get("reason_code")
4925 .map_err(|_| RepositoryError::Unavailable)?,
4926 row.try_get("operator_reference")
4927 .map_err(|_| RepositoryError::Unavailable)?,
4928 evidence_digest,
4929 millis_system_time(read_i64(row, "decided_ms")?)?,
4930 ))
4931}
4932
4933fn decode_execution_metadata(row: &PgRow) -> Result<ExecutionMetadata, RepositoryError> {
4934 let status = decode_status(
4935 &row.try_get::<String, _>("status")
4936 .map_err(|_| RepositoryError::Unavailable)?,
4937 )?;
4938 let exit_status = ExitStatus::new(ExitCode::new(
4939 row.try_get::<String, _>("exit_code")
4940 .map_err(|_| RepositoryError::Unavailable)?,
4941 )?);
4942 let timestamps = ExecutionTimestamps::new(
4943 millis_system_time(read_i64(row, "created_ms")?)?,
4944 read_optional_i64(row, "started_ms")?
4945 .map(millis_system_time)
4946 .transpose()?,
4947 read_optional_i64(row, "ended_ms")?
4948 .map(millis_system_time)
4949 .transpose()?,
4950 )?;
4951 let counts = ExecutionCounts::new(
4952 read_u64(row, "read_count")?,
4953 read_u64(row, "processed_count")?,
4954 read_u64(row, "write_count")?,
4955 read_u64(row, "filter_count")?,
4956 read_u64(row, "commit_count")?,
4957 read_u64(row, "rollback_count")?,
4958 );
4959 let category = row
4960 .try_get::<Option<String>, _>("failure_category")
4961 .map_err(|_| RepositoryError::Unavailable)?;
4962 let failure_id = read_optional_i64(row, "failure_id")?;
4963 let failure = match (category, failure_id) {
4964 (None, None) => None,
4965 (Some(category), Some(id)) => Some(FailureSummary::new(
4966 decode_failure_category(&category)?,
4967 FailureId::new(u64::try_from(id).map_err(|_| RepositoryError::Unavailable)?)?,
4968 )),
4969 _ => return Err(RepositoryError::Unavailable),
4970 };
4971 ExecutionMetadata::new(status, exit_status, timestamps, counts, failure)
4972 .map_err(RepositoryError::from)
4973}
4974
4975async fn update_job_execution(
4976 transaction: &mut PgConnection,
4977 execution: &JobExecution,
4978 updated_at: SystemTime,
4979 expected: ExecutionVersion,
4980) -> Result<u64, RepositoryError> {
4981 update_execution(
4982 transaction,
4983 "oxide_batch.ob_job_execution",
4984 execution.id().get(),
4985 IdentifierKind::JobExecution,
4986 execution.metadata(),
4987 execution.version(),
4988 updated_at,
4989 expected,
4990 )
4991 .await
4992}
4993
4994async fn update_step_execution(
4995 transaction: &mut PgConnection,
4996 execution: &StepExecution,
4997 updated_at: SystemTime,
4998 expected: ExecutionVersion,
4999) -> Result<u64, RepositoryError> {
5000 let metadata = execution.metadata();
5001 let timestamps = metadata.timestamps();
5002 let failure = metadata.failure();
5003 let failure_category = failure.map(|value| encode_failure_category(value.category()));
5004 let failure_id = failure
5005 .map(|value| database_id(value.failure_id().get(), IdentifierKind::Failure))
5006 .transpose()?;
5007 let counts = metadata.counts();
5008 let result = sqlx::query(
5009 "UPDATE oxide_batch.ob_step_execution \
5010 SET status = $1, exit_code = $2, failure_category = $3, failure_id = $4, \
5011 started_at = CASE WHEN $5::bigint IS NULL THEN NULL \
5012 ELSE to_timestamp($5::double precision / 1000.0) END, \
5013 ended_at = CASE WHEN $6::bigint IS NULL THEN NULL \
5014 ELSE to_timestamp($6::double precision / 1000.0) END, \
5015 read_count = $7, processed_count = $8, write_count = $9, \
5016 filter_count = $10, commit_count = $11, rollback_count = $12, \
5017 updated_at = to_timestamp($13::double precision / 1000.0), version = $14 \
5018 WHERE id = $15 AND version = $16",
5019 )
5020 .bind(metadata.status().to_string())
5021 .bind(metadata.exit_status().code().as_str())
5022 .bind(failure_category)
5023 .bind(failure_id)
5024 .bind(
5025 timestamps
5026 .started_at()
5027 .map(system_time_millis)
5028 .transpose()?,
5029 )
5030 .bind(timestamps.ended_at().map(system_time_millis).transpose()?)
5031 .bind(execution_count(counts.read())?)
5032 .bind(execution_count(counts.processed())?)
5033 .bind(execution_count(counts.written())?)
5034 .bind(execution_count(counts.filtered())?)
5035 .bind(execution_count(counts.committed())?)
5036 .bind(execution_count(counts.rolled_back())?)
5037 .bind(system_time_millis(updated_at)?)
5038 .bind(database_version(execution.version())?)
5039 .bind(database_id(
5040 execution.id().get(),
5041 IdentifierKind::StepExecution,
5042 )?)
5043 .bind(database_version(expected)?)
5044 .execute(transaction)
5045 .await
5046 .map_err(|_| RepositoryError::Unavailable)?;
5047 Ok(result.rows_affected())
5048}
5049
5050#[allow(clippy::too_many_arguments)]
5051async fn update_execution(
5052 transaction: &mut PgConnection,
5053 table: &'static str,
5054 id: u64,
5055 kind: IdentifierKind,
5056 metadata: &ExecutionMetadata,
5057 new_version: ExecutionVersion,
5058 updated_at: SystemTime,
5059 expected: ExecutionVersion,
5060) -> Result<u64, RepositoryError> {
5061 let query = format!(
5062 "UPDATE {table} SET status = $1, exit_code = $2, failure_category = $3, \
5063 failure_id = $4, started_at = CASE WHEN $5::bigint IS NULL THEN NULL \
5064 ELSE to_timestamp($5::double precision / 1000.0) END, \
5065 ended_at = CASE WHEN $6::bigint IS NULL THEN NULL \
5066 ELSE to_timestamp($6::double precision / 1000.0) END, \
5067 updated_at = to_timestamp($7::double precision / 1000.0), version = $8 \
5068 WHERE id = $9 AND version = $10"
5069 );
5070 let timestamps = metadata.timestamps();
5071 let failure = metadata.failure();
5072 let failure_category = failure.map(|value| encode_failure_category(value.category()));
5073 let failure_id = failure
5074 .map(|value| database_id(value.failure_id().get(), IdentifierKind::Failure))
5075 .transpose()?;
5076 let result = sqlx::query(AssertSqlSafe(query))
5077 .bind(metadata.status().to_string())
5078 .bind(metadata.exit_status().code().as_str())
5079 .bind(failure_category)
5080 .bind(failure_id)
5081 .bind(
5082 timestamps
5083 .started_at()
5084 .map(system_time_millis)
5085 .transpose()?,
5086 )
5087 .bind(timestamps.ended_at().map(system_time_millis).transpose()?)
5088 .bind(system_time_millis(updated_at)?)
5089 .bind(database_version(new_version)?)
5090 .bind(database_id(id, kind)?)
5091 .bind(database_version(expected)?)
5092 .execute(transaction)
5093 .await
5094 .map_err(|_| RepositoryError::Unavailable)?;
5095 Ok(result.rows_affected())
5096}
5097
5098fn database_version(version: ExecutionVersion) -> Result<i64, RepositoryError> {
5099 i64::try_from(version.get()).map_err(|_| RepositoryError::Unavailable)
5100}
5101
5102fn read_i64(row: &PgRow, name: &str) -> Result<i64, RepositoryError> {
5103 row.try_get(name).map_err(|_| RepositoryError::Unavailable)
5104}
5105
5106fn read_optional_i64(row: &PgRow, name: &str) -> Result<Option<i64>, RepositoryError> {
5107 row.try_get(name).map_err(|_| RepositoryError::Unavailable)
5108}
5109
5110fn read_optional_u64(row: &PgRow, name: &str) -> Result<Option<u64>, RepositoryError> {
5111 read_optional_i64(row, name)?
5112 .map(|value| u64::try_from(value).map_err(|_| RepositoryError::Unavailable))
5113 .transpose()
5114}
5115
5116fn read_u64(row: &PgRow, name: &str) -> Result<u64, RepositoryError> {
5117 u64::try_from(read_i64(row, name)?).map_err(|_| RepositoryError::Unavailable)
5118}
5119
5120fn decode_status(value: &str) -> Result<BatchStatus, RepositoryError> {
5121 match value {
5122 "STARTING" => Ok(BatchStatus::Starting),
5123 "STARTED" => Ok(BatchStatus::Started),
5124 "STOPPING" => Ok(BatchStatus::Stopping),
5125 "STOPPED" => Ok(BatchStatus::Stopped),
5126 "FAILED" => Ok(BatchStatus::Failed),
5127 "COMPLETED" => Ok(BatchStatus::Completed),
5128 "ABANDONED" => Ok(BatchStatus::Abandoned),
5129 "UNKNOWN" => Ok(BatchStatus::Unknown),
5130 _ => Err(RepositoryError::Unavailable),
5131 }
5132}
5133
5134const fn encode_failure_category(value: FailureCategory) -> &'static str {
5140 value.durable_code()
5141}
5142
5143fn decode_failure_category(value: &str) -> Result<FailureCategory, RepositoryError> {
5144 FailureCategory::from_durable_code(value).ok_or(RepositoryError::Unavailable)
5145}
5146
5147#[cfg(test)]
5148mod tests {
5149 use super::*;
5150
5151 #[test]
5152 fn canonical_instance_key_matches_version_one_golden_vector() -> Result<(), Box<dyn Error>> {
5153 let parameters = JobParameters::try_from_iter([
5154 (
5155 ParameterName::new("region")?,
5156 JobParameter::new(ParameterValue::string("서울")?, ParameterRole::Identifying),
5157 ),
5158 (
5159 ParameterName::new("limit")?,
5160 JobParameter::new(
5161 ParameterValue::from(1_u64 << 63),
5162 ParameterRole::Identifying,
5163 ),
5164 ),
5165 (
5166 ParameterName::new("count")?,
5167 JobParameter::new(ParameterValue::from(-2_i64), ParameterRole::Identifying),
5168 ),
5169 (
5170 ParameterName::new("active")?,
5171 JobParameter::new(ParameterValue::from(true), ParameterRole::Identifying),
5172 ),
5173 ])?;
5174 let key = JobInstanceKey::new(JobName::new("golden_job")?, ¶meters);
5175 assert_eq!(
5176 key.digest(),
5177 [
5178 0x71, 0xf1, 0x2d, 0xb9, 0xe3, 0x88, 0x7d, 0xe2, 0xcf, 0x92, 0xe9, 0x3b, 0xb6, 0x3f,
5179 0xd4, 0xe9, 0xe7, 0xc5, 0x36, 0xdf, 0x8f, 0xa2, 0x02, 0x21, 0x24, 0x45, 0xd1, 0x8b,
5180 0xf2, 0xe4, 0x36, 0x04,
5181 ]
5182 );
5183 Ok(())
5184 }
5185}
5186
5187fn candidate_execution_ids(candidates: &[PurgeCandidate]) -> Result<Vec<i64>, RepositoryError> {
5188 candidates
5189 .iter()
5190 .map(|candidate| {
5191 database_id(
5192 candidate.job_execution_id().get(),
5193 IdentifierKind::JobExecution,
5194 )
5195 })
5196 .collect()
5197}
5198
5199fn candidate_instance_ids(candidates: &[PurgeCandidate]) -> Result<Vec<i64>, RepositoryError> {
5200 let mut ids = candidates
5201 .iter()
5202 .map(|candidate| {
5203 database_id(
5204 candidate.job_instance_id().get(),
5205 IdentifierKind::JobInstance,
5206 )
5207 })
5208 .collect::<Result<Vec<_>, _>>()?;
5209 ids.sort_unstable();
5210 ids.dedup();
5211 Ok(ids)
5212}
5213
5214fn retention_count(value: u64) -> Result<i64, RepositoryError> {
5215 i64::try_from(value).map_err(|_| RepositoryError::Unavailable)
5216}
5217
5218fn operator_request_select(suffix: &str) -> String {
5219 format!(
5220 "SELECT request.id, request.job_instance_id, request.job_execution_id, \
5221 request.action, request.operation_id, request.actor_ref, request.reason_code, \
5222 request.request_digest, request.observed_version, request.prior_status, \
5223 request.result_status, request.outcome_class, request.rejection_code, \
5224 (extract(epoch FROM request.requested_at) * 1000)::bigint AS requested_ms \
5225 FROM oxide_batch.ob_operator_request request {suffix}"
5226 )
5227}
5228
5229fn retention_action_select(suffix: &str) -> String {
5230 format!(
5231 "SELECT retention.id, retention.job_instance_id, retention.action, \
5232 retention.operation_id, retention.actor_ref, retention.reason_code, \
5233 retention.plan_digest, retention.batch_bound, \
5234 retention.deleted_flow_decisions, retention.deleted_recovery_decisions, \
5235 retention.deleted_operator_requests, retention.deleted_step_partitions, \
5236 retention.deleted_step_executions, retention.deleted_job_executions, \
5237 retention.deleted_job_instances, retention.outcome_class, \
5238 (extract(epoch FROM retention.applied_at) * 1000)::bigint AS applied_ms \
5239 FROM oxide_batch.ob_retention_action retention {suffix}"
5240 )
5241}
5242
5243fn read_text(row: &PgRow, name: &str) -> Result<String, RepositoryError> {
5244 row.try_get::<String, _>(name)
5245 .map_err(|_| RepositoryError::Unavailable)
5246}
5247
5248fn read_optional_text(row: &PgRow, name: &str) -> Result<Option<String>, RepositoryError> {
5249 row.try_get::<Option<String>, _>(name)
5250 .map_err(|_| RepositoryError::Unavailable)
5251}
5252
5253fn read_digest(row: &PgRow, name: &str) -> Result<[u8; 32], RepositoryError> {
5254 row.try_get::<Vec<u8>, _>(name)
5255 .map_err(|_| RepositoryError::Unavailable)?
5256 .try_into()
5257 .map_err(|_| RepositoryError::Unavailable)
5258}
5259
5260fn decode_operator_action(value: &str) -> Result<OperatorAction, RepositoryError> {
5261 Ok(match value {
5262 "LAUNCH" => OperatorAction::Launch,
5263 "RESTART" => OperatorAction::Restart,
5264 "STOP" => OperatorAction::Stop,
5265 "ABANDON" => OperatorAction::Abandon,
5266 "RECOVER" => OperatorAction::Recover,
5267 _ => return Err(RepositoryError::Unavailable),
5268 })
5269}
5270
5271fn decode_retention_action(value: &str) -> Result<RetentionAction, RepositoryError> {
5272 Ok(match value {
5273 "HOLD" => RetentionAction::Hold,
5274 "RELEASE_HOLD" => RetentionAction::ReleaseHold,
5275 "APPLY_PURGE" => RetentionAction::ApplyPurge,
5276 _ => return Err(RepositoryError::Unavailable),
5277 })
5278}
5279
5280fn decode_operator_outcome(value: &str) -> Result<OperatorOutcomeClass, RepositoryError> {
5281 Ok(match value {
5282 "APPLIED" => OperatorOutcomeClass::Applied,
5283 "REJECTED" => OperatorOutcomeClass::Rejected,
5284 _ => return Err(RepositoryError::Unavailable),
5285 })
5286}
5287
5288fn decode_retention_outcome(value: &str) -> Result<RetentionOutcome, RepositoryError> {
5289 Ok(match value {
5290 "APPLIED" => RetentionOutcome::Applied,
5291 "REJECTED" => RetentionOutcome::Rejected,
5292 _ => return Err(RepositoryError::Unavailable),
5293 })
5294}
5295
5296fn decode_operator_rejection(
5297 value: &str,
5298 row: &PgRow,
5299) -> Result<OperatorRejection, RepositoryError> {
5300 Ok(match value {
5301 "OPTIMISTIC_CONFLICT" => OperatorRejection::OptimisticConflict {
5302 current: ExecutionVersion::new(
5303 read_optional_u64(row, "observed_version")?.unwrap_or(0),
5304 ),
5305 },
5306 "INVALID_STATE" => OperatorRejection::InvalidState {
5307 status: read_optional_text(row, "prior_status")?
5308 .as_deref()
5309 .map(decode_status)
5310 .transpose()?
5311 .unwrap_or(BatchStatus::Unknown),
5312 },
5313 "INSTANCE_COMPLETED" => OperatorRejection::InstanceCompleted,
5314 "INSTANCE_ABANDONED" => OperatorRejection::InstanceAbandoned,
5315 "EXECUTION_ALREADY_ACTIVE" => OperatorRejection::ExecutionAlreadyActive {
5316 execution_id: JobExecutionId::new(read_u64(row, "job_execution_id")?)?,
5317 status: read_optional_text(row, "prior_status")?
5318 .as_deref()
5319 .map(decode_status)
5320 .transpose()?
5321 .unwrap_or(BatchStatus::Unknown),
5322 },
5323 "INCOMPATIBLE_DEFINITION" => OperatorRejection::IncompatibleDefinition,
5324 "RESTART_WITHOUT_PRIOR_ATTEMPT" => OperatorRejection::RestartWithoutPriorAttempt,
5325 "START_LIMIT_EXCEEDED" => OperatorRejection::StartLimitExceeded,
5326 "UNRESOLVED_RECOVERY_REQUIRED" => OperatorRejection::UnresolvedRecoveryRequired,
5327 "EXECUTION_NOT_FOUND" => OperatorRejection::ExecutionNotFound,
5328 "INSTANCE_NOT_FOUND" => OperatorRejection::InstanceNotFound,
5329 "UNSUPPORTED_ACTION" => OperatorRejection::UnsupportedAction,
5330 _ => return Err(RepositoryError::Unavailable),
5331 })
5332}
5333
5334fn decode_operator_record(row: &PgRow) -> Result<OperatorRecord, RepositoryError> {
5335 let id = OperatorRequestId::new(read_u64(row, "id")?)?;
5336 let action = decode_operator_action(&read_text(row, "action")?)?;
5337 let outcome = decode_operator_outcome(&read_text(row, "outcome_class")?)?;
5338 let rejection = read_optional_text(row, "rejection_code")?
5339 .map(|code| decode_operator_rejection(&code, row))
5340 .transpose()?;
5341 let draft = OperatorRecordDraft::from_durable(
5342 action,
5343 OperationId::new(read_text(row, "operation_id")?)
5344 .map_err(|_| RepositoryError::Unavailable)?,
5345 ActorRef::new(read_text(row, "actor_ref")?).map_err(|_| RepositoryError::Unavailable)?,
5346 read_optional_text(row, "reason_code")?
5347 .map(ReasonCode::new)
5348 .transpose()
5349 .map_err(|_| RepositoryError::Unavailable)?,
5350 RequestDigest::from_bytes(read_digest(row, "request_digest")?),
5351 read_optional_u64(row, "job_instance_id")?
5352 .map(JobInstanceId::new)
5353 .transpose()?,
5354 read_optional_u64(row, "job_execution_id")?
5355 .map(JobExecutionId::new)
5356 .transpose()?,
5357 read_optional_u64(row, "observed_version")?.map(ExecutionVersion::new),
5358 read_optional_text(row, "prior_status")?
5359 .as_deref()
5360 .map(decode_status)
5361 .transpose()?,
5362 read_optional_text(row, "result_status")?
5363 .as_deref()
5364 .map(decode_status)
5365 .transpose()?,
5366 outcome,
5367 rejection,
5368 millis_system_time(read_i64(row, "requested_ms")?)?,
5369 );
5370 Ok(OperatorRecord::from_parts(id, draft))
5371}
5372
5373fn decode_retention_record(row: &PgRow) -> Result<RetentionRecord, RepositoryError> {
5374 let id = RetentionActionId::new(read_u64(row, "id")?)?;
5375 let counts = PurgeCounts::new(
5376 read_u64(row, "deleted_flow_decisions")?,
5377 read_u64(row, "deleted_recovery_decisions")?,
5378 read_u64(row, "deleted_operator_requests")?,
5379 read_u64(row, "deleted_step_partitions")?,
5380 read_u64(row, "deleted_step_executions")?,
5381 read_u64(row, "deleted_job_executions")?,
5382 read_u64(row, "deleted_job_instances")?,
5383 );
5384 let batch_bound = row
5385 .try_get::<Option<i32>, _>("batch_bound")
5386 .map_err(|_| RepositoryError::Unavailable)?
5387 .map(|bound| {
5388 u32::try_from(bound)
5389 .map_err(|_| RepositoryError::Unavailable)
5390 .and_then(|bound| {
5391 PurgeBatchBound::new(bound).map_err(|_| RepositoryError::Unavailable)
5392 })
5393 })
5394 .transpose()?;
5395 let plan_digest = row
5396 .try_get::<Option<Vec<u8>>, _>("plan_digest")
5397 .map_err(|_| RepositoryError::Unavailable)?
5398 .map(|digest| <[u8; 32]>::try_from(digest).map_err(|_| RepositoryError::Unavailable))
5399 .transpose()?;
5400 let draft = RetentionRecordDraft::from_durable(
5401 decode_retention_action(&read_text(row, "action")?)?,
5402 OperationId::new(read_text(row, "operation_id")?)
5403 .map_err(|_| RepositoryError::Unavailable)?,
5404 ActorRef::new(read_text(row, "actor_ref")?).map_err(|_| RepositoryError::Unavailable)?,
5405 ReasonCode::new(read_text(row, "reason_code")?)
5406 .map_err(|_| RepositoryError::Unavailable)?,
5407 read_optional_u64(row, "job_instance_id")?
5408 .map(JobInstanceId::new)
5409 .transpose()?,
5410 plan_digest,
5411 counts,
5412 batch_bound,
5413 decode_retention_outcome(&read_text(row, "outcome_class")?)?,
5414 millis_system_time(read_i64(row, "applied_ms")?)?,
5415 );
5416 Ok(RetentionRecord::from_parts(id, draft))
5417}
5418
5419fn decode_retention_hold(
5420 id: JobInstanceId,
5421 row: &PgRow,
5422) -> Result<Option<RetentionHold>, RepositoryError> {
5423 let Some(actor) = read_optional_text(row, "hold_actor")? else {
5424 return Ok(None);
5425 };
5426 let reason = read_optional_text(row, "hold_reason")?.ok_or(RepositoryError::Unavailable)?;
5427 let placed_ms = read_optional_i64(row, "placed_ms")?.ok_or(RepositoryError::Unavailable)?;
5428 Ok(Some(RetentionHold::new(
5429 id,
5430 ActorRef::new(actor).map_err(|_| RepositoryError::Unavailable)?,
5431 ReasonCode::new(reason).map_err(|_| RepositoryError::Unavailable)?,
5432 millis_system_time(placed_ms)?,
5433 )))
5434}
5435
5436#[derive(Clone)]
5447pub struct PostgresExplorer {
5448 repository: PostgresJobRepository,
5449}
5450
5451impl PostgresExplorer {
5452 #[must_use]
5454 pub const fn new(repository: PostgresJobRepository) -> Self {
5455 Self { repository }
5456 }
5457
5458 async fn fetch_all(
5459 &self,
5460 query: sqlx::query::Query<'_, Postgres, PgArguments>,
5461 ) -> Result<Vec<PgRow>, ExplorerError> {
5462 let mut connection = self
5463 .repository
5464 .begin_connection()
5465 .await
5466 .map_err(ExplorerError::Repository)?;
5467 let result = query.fetch_all(&mut *connection).await;
5468 let _ = sqlx::query("ROLLBACK").execute(&mut *connection).await;
5469 result.map_err(|error| classify_explorer_error(&error))
5470 }
5471
5472 async fn fetch_optional(
5473 &self,
5474 query: sqlx::query::Query<'_, Postgres, PgArguments>,
5475 ) -> Result<Option<PgRow>, ExplorerError> {
5476 let mut connection = self
5477 .repository
5478 .begin_connection()
5479 .await
5480 .map_err(ExplorerError::Repository)?;
5481 let result = query.fetch_optional(&mut *connection).await;
5482 let _ = sqlx::query("ROLLBACK").execute(&mut *connection).await;
5483 result.map_err(|error| classify_explorer_error(&error))
5484 }
5485}
5486
5487impl fmt::Debug for PostgresExplorer {
5488 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
5489 formatter
5490 .debug_struct("PostgresExplorer")
5491 .finish_non_exhaustive()
5492 }
5493}
5494
5495fn classify_explorer_error(error: &sqlx::Error) -> ExplorerError {
5496 if let sqlx::Error::Database(database) = error
5497 && database.code().as_deref() == Some("57014")
5498 {
5499 return ExplorerError::Timeout;
5500 }
5501 ExplorerError::Repository(RepositoryError::Unavailable)
5502}
5503
5504const fn ceiling_source(query: &ExplorerQuery) -> Option<&'static str> {
5505 Some(match query {
5506 ExplorerQuery::JobNames => "SELECT max(id) AS ceiling FROM oxide_batch.ob_job_definition",
5507 ExplorerQuery::Instances { .. } => {
5508 "SELECT max(id) AS ceiling FROM oxide_batch.ob_job_instance WHERE job_name = $1"
5509 }
5510 ExplorerQuery::Executions { .. } | ExplorerQuery::UnresolvedExecutions { .. } => {
5511 "SELECT max(id) AS ceiling FROM oxide_batch.ob_job_execution"
5512 }
5513 ExplorerQuery::StepExecutions { .. } => {
5514 "SELECT max(id) AS ceiling FROM oxide_batch.ob_step_execution"
5515 }
5516 ExplorerQuery::RecoveryDecisions { .. } => {
5517 "SELECT max(id) AS ceiling FROM oxide_batch.ob_recovery_decision"
5518 }
5519 ExplorerQuery::FlowDecisions { .. } => {
5520 "SELECT max(id) AS ceiling FROM oxide_batch.ob_flow_decision"
5521 }
5522 ExplorerQuery::StepPartitions { .. } => {
5523 "SELECT max(id) AS ceiling FROM oxide_batch.ob_step_partition"
5524 }
5525 ExplorerQuery::OperatorRequests { .. } => {
5526 "SELECT max(id) AS ceiling FROM oxide_batch.ob_operator_request"
5527 }
5528 _ => return None,
5532 })
5533}
5534
5535fn window_limit(window: &QueryWindow) -> i64 {
5536 i64::from(window.limit())
5537}
5538
5539fn window_ceiling(window: &QueryWindow) -> Result<i64, ExplorerError> {
5540 i64::try_from(window.ceiling())
5541 .map_err(|_| ExplorerError::Repository(RepositoryError::Unavailable))
5542}
5543
5544fn window_identity(window: &QueryWindow) -> Result<Option<i64>, ExplorerError> {
5545 match window.after() {
5546 Some(CursorKey::Identity(value)) => i64::try_from(*value)
5547 .map(Some)
5548 .map_err(|_| ExplorerError::Cursor(CursorError::CursorInvalid)),
5549 Some(_) => Err(ExplorerError::Cursor(CursorError::CursorInvalid)),
5550 None => Ok(None),
5551 }
5552}
5553
5554fn window_ordered(window: &QueryWindow) -> Result<Option<(i64, i64)>, ExplorerError> {
5555 match window.after() {
5556 Some(CursorKey::Ordered { primary, identity }) => {
5557 let primary = i64::try_from(*primary)
5558 .map_err(|_| ExplorerError::Cursor(CursorError::CursorInvalid))?;
5559 let identity = i64::try_from(*identity)
5560 .map_err(|_| ExplorerError::Cursor(CursorError::CursorInvalid))?;
5561 Ok(Some((primary, identity)))
5562 }
5563 Some(_) => Err(ExplorerError::Cursor(CursorError::CursorInvalid)),
5564 None => Ok(None),
5565 }
5566}
5567
5568fn window_name(window: &QueryWindow) -> Result<Option<String>, ExplorerError> {
5569 match window.after() {
5570 Some(CursorKey::Name(value)) => Ok(Some(value.clone())),
5571 Some(_) => Err(ExplorerError::Cursor(CursorError::CursorInvalid)),
5572 None => Ok(None),
5573 }
5574}
5575
5576fn explorer_id(value: u64, kind: IdentifierKind) -> Result<i64, ExplorerError> {
5577 database_id(value, kind).map_err(ExplorerError::Repository)
5578}
5579
5580fn job_execution_projection_select(suffix: &str) -> String {
5581 format!(
5582 "SELECT execution.id, execution.job_instance_id, instance.job_name, \
5583 execution.attempt, execution.status, execution.exit_code, \
5584 execution.failure_category, execution.failure_id, \
5585 (extract(epoch FROM execution.created_at) * 1000)::bigint AS created_ms, \
5586 (extract(epoch FROM execution.started_at) * 1000)::bigint AS started_ms, \
5587 (extract(epoch FROM execution.ended_at) * 1000)::bigint AS ended_ms, \
5588 (extract(epoch FROM execution.updated_at) * 1000)::bigint AS updated_ms, \
5589 (extract(epoch FROM execution.stop_requested_at) * 1000)::bigint AS stop_ms, \
5590 (execution.owner_token IS NOT NULL) AS owner_recorded, execution.version, \
5591 definition.definition_revision, definition.manifest_format, \
5592 definition.manifest_digest, execution.context_format, \
5593 execution.context_schema, execution.context_schema_version, \
5594 pg_column_size(execution.context_payload)::bigint AS context_bytes \
5595 FROM oxide_batch.ob_job_execution execution \
5596 JOIN oxide_batch.ob_job_instance instance \
5597 ON instance.id = execution.job_instance_id \
5598 JOIN oxide_batch.ob_job_definition definition \
5599 ON definition.id = execution.definition_id {suffix}"
5600 )
5601}
5602
5603fn step_execution_projection_select(suffix: &str) -> String {
5604 format!(
5605 "SELECT execution.id, execution.job_execution_id, execution.step_name, \
5606 execution.step_logical_id, execution.status, execution.exit_code, \
5607 execution.read_count, execution.processed_count, execution.write_count, \
5608 execution.filter_count, execution.commit_count, execution.rollback_count, \
5609 execution.failure_category, execution.failure_id, execution.version, \
5610 (extract(epoch FROM execution.created_at) * 1000)::bigint AS created_ms, \
5611 (extract(epoch FROM execution.started_at) * 1000)::bigint AS started_ms, \
5612 (extract(epoch FROM execution.ended_at) * 1000)::bigint AS ended_ms, \
5613 execution.checkpoint_format, execution.checkpoint_schema, \
5614 execution.checkpoint_schema_version, \
5615 pg_column_size(execution.checkpoint_payload)::bigint AS checkpoint_bytes, \
5616 execution.context_format, execution.context_schema, \
5617 execution.context_schema_version, \
5618 pg_column_size(execution.context_payload)::bigint AS context_bytes \
5619 FROM oxide_batch.ob_step_execution execution {suffix}"
5620 )
5621}
5622
5623fn step_partition_projection_select(suffix: &str) -> String {
5624 format!(
5625 "SELECT partition.id, partition.step_execution_id, \
5626 partition.worker_step_execution_id, partition.partition_key, \
5627 partition.partition_ordinal, partition.status, partition.exit_code, \
5628 partition.read_count, partition.processed_count, partition.write_count, \
5629 partition.filter_count, partition.commit_count, partition.rollback_count, \
5630 partition.version, partition.context_format, partition.context_schema, \
5631 partition.context_schema_version, \
5632 pg_column_size(partition.context_payload)::bigint AS context_bytes \
5633 FROM oxide_batch.ob_step_partition partition {suffix}"
5634 )
5635}
5636
5637fn decode_state_descriptor(
5638 row: &PgRow,
5639 kind: DurableStateKind,
5640 format: &str,
5641 schema: &str,
5642 schema_version: &str,
5643 bytes: &str,
5644) -> Result<StateEnvelopeDescriptor, ExplorerError> {
5645 let unavailable = || ExplorerError::Repository(RepositoryError::Unavailable);
5646 let format_version = row
5647 .try_get::<i16, _>(format)
5648 .map_err(|_| unavailable())
5649 .and_then(|value| u16::try_from(value).map_err(|_| unavailable()))?;
5650 let schema_id = StateSchemaId::new(read_text(row, schema).map_err(ExplorerError::Repository)?)
5651 .map_err(|_| unavailable())?;
5652 let schema_version = row
5653 .try_get::<i32, _>(schema_version)
5654 .map_err(|_| unavailable())
5655 .and_then(|value| u32::try_from(value).map_err(|_| unavailable()))
5656 .and_then(|value| StateSchemaVersion::new(value).map_err(|_| unavailable()))?;
5657 let encoded_len = read_optional_i64(row, bytes)
5658 .map_err(ExplorerError::Repository)?
5659 .and_then(|value| usize::try_from(value).ok())
5660 .unwrap_or(0);
5661 Ok(StateEnvelopeDescriptor::new(
5662 kind,
5663 format_version,
5664 schema_id,
5665 schema_version,
5666 encoded_len,
5667 ))
5668}
5669
5670fn decode_projection_counts(row: &PgRow) -> Result<ExecutionCounts, ExplorerError> {
5671 let read = |name: &str| read_u64(row, name).map_err(ExplorerError::Repository);
5672 Ok(ExecutionCounts::new(
5673 read("read_count")?,
5674 read("processed_count")?,
5675 read("write_count")?,
5676 read("filter_count")?,
5677 read("commit_count")?,
5678 read("rollback_count")?,
5679 ))
5680}
5681
5682fn decode_projection_failure(row: &PgRow) -> Result<Option<FailureSummary>, ExplorerError> {
5683 let unavailable = || ExplorerError::Repository(RepositoryError::Unavailable);
5684 let category =
5685 read_optional_text(row, "failure_category").map_err(ExplorerError::Repository)?;
5686 let failure_id = read_optional_u64(row, "failure_id").map_err(ExplorerError::Repository)?;
5687 match (category, failure_id) {
5688 (Some(category), Some(failure_id)) => Ok(Some(FailureSummary::new(
5689 decode_failure_category(&category).map_err(|_| unavailable())?,
5690 FailureId::new(failure_id).map_err(|_| unavailable())?,
5691 ))),
5692 (None, None) => Ok(None),
5693 _ => Err(unavailable()),
5694 }
5695}
5696
5697fn decode_projection_timestamps(row: &PgRow) -> Result<ExecutionTimestamps, ExplorerError> {
5698 let unavailable = || ExplorerError::Repository(RepositoryError::Unavailable);
5699 let created =
5700 millis_system_time(read_i64(row, "created_ms").map_err(ExplorerError::Repository)?)
5701 .map_err(|_| unavailable())?;
5702 let started = read_optional_i64(row, "started_ms")
5703 .map_err(ExplorerError::Repository)?
5704 .map(millis_system_time)
5705 .transpose()
5706 .map_err(|_| unavailable())?;
5707 let ended = read_optional_i64(row, "ended_ms")
5708 .map_err(ExplorerError::Repository)?
5709 .map(millis_system_time)
5710 .transpose()
5711 .map_err(|_| unavailable())?;
5712 ExecutionTimestamps::new(created, started, ended).map_err(|_| unavailable())
5713}
5714
5715fn decode_job_execution_projection(row: &PgRow) -> Result<JobExecutionProjection, ExplorerError> {
5716 let unavailable = || ExplorerError::Repository(RepositoryError::Unavailable);
5717 let id = JobExecutionId::new(read_u64(row, "id").map_err(ExplorerError::Repository)?)
5718 .map_err(|_| unavailable())?;
5719 let instance_id =
5720 JobInstanceId::new(read_u64(row, "job_instance_id").map_err(ExplorerError::Repository)?)
5721 .map_err(|_| unavailable())?;
5722 let job_name = JobName::new(read_text(row, "job_name").map_err(ExplorerError::Repository)?)
5723 .map_err(|_| unavailable())?;
5724 let attempt = row
5725 .try_get::<i32, _>("attempt")
5726 .map_err(|_| unavailable())
5727 .and_then(|value| u32::try_from(value).map_err(|_| unavailable()))?;
5728 let status = decode_status(&read_text(row, "status").map_err(ExplorerError::Repository)?)
5729 .map_err(|_| unavailable())?;
5730 let exit_status = ExitStatus::new(
5731 ExitCode::new(read_text(row, "exit_code").map_err(ExplorerError::Repository)?)
5732 .map_err(|_| unavailable())?,
5733 );
5734 let definition = DefinitionDescriptor::new(
5735 DefinitionRevision::new(
5736 read_text(row, "definition_revision").map_err(ExplorerError::Repository)?,
5737 )
5738 .map_err(|_| unavailable())?,
5739 row.try_get::<i16, _>("manifest_format")
5740 .map_err(|_| unavailable())
5741 .and_then(|value| u16::try_from(value).map_err(|_| unavailable()))?,
5742 read_digest(row, "manifest_digest").map_err(ExplorerError::Repository)?,
5743 );
5744 let context = decode_state_descriptor(
5745 row,
5746 DurableStateKind::ExecutionContext,
5747 "context_format",
5748 "context_schema",
5749 "context_schema_version",
5750 "context_bytes",
5751 )?;
5752 let stop_requested_at = read_optional_i64(row, "stop_ms")
5753 .map_err(ExplorerError::Repository)?
5754 .map(millis_system_time)
5755 .transpose()
5756 .map_err(|_| unavailable())?;
5757 let owner_recorded = row
5758 .try_get::<bool, _>("owner_recorded")
5759 .map_err(|_| unavailable())?;
5760 Ok(JobExecutionProjection::new(
5761 id,
5762 instance_id,
5763 job_name,
5764 attempt,
5765 status,
5766 exit_status,
5767 ExecutionCounts::default(),
5768 ExecutionVersion::new(read_u64(row, "version").map_err(ExplorerError::Repository)?),
5769 decode_projection_timestamps(row)?,
5770 millis_system_time(read_i64(row, "updated_ms").map_err(ExplorerError::Repository)?)
5771 .map_err(|_| unavailable())?,
5772 decode_projection_failure(row)?,
5773 Some(definition),
5774 Some(context),
5775 stop_requested_at,
5776 owner_recorded,
5777 ))
5778}
5779
5780fn decode_job_instance_projection(row: &PgRow) -> Result<JobInstanceProjection, ExplorerError> {
5781 let unavailable = || ExplorerError::Repository(RepositoryError::Unavailable);
5782 let id = JobInstanceId::new(read_u64(row, "id").map_err(ExplorerError::Repository)?)
5783 .map_err(|_| unavailable())?;
5784 let job_name = JobName::new(read_text(row, "job_name").map_err(ExplorerError::Repository)?)
5785 .map_err(|_| unavailable())?;
5786 let parameters = row
5787 .try_get::<Json<Value>, _>("identifying_parameters")
5788 .map_err(|_| unavailable())?;
5789 let parameters = decode_identifying_parameters(¶meters.0).map_err(|_| unavailable())?;
5790 let descriptors = parameters
5791 .iter()
5792 .map(|(name, parameter)| {
5793 ParameterDescriptor::new(
5794 name.clone(),
5795 parameter.value().kind(),
5796 parameter.is_identifying(),
5797 )
5798 })
5799 .collect();
5800 let created_at = read_optional_i64(row, "created_ms")
5801 .map_err(ExplorerError::Repository)?
5802 .map(millis_system_time)
5803 .transpose()
5804 .map_err(|_| unavailable())?;
5805 let hold = decode_retention_hold(id, row).map_err(ExplorerError::Repository)?;
5806 Ok(JobInstanceProjection::new(
5807 id,
5808 job_name,
5809 read_digest(row, "instance_key").map_err(ExplorerError::Repository)?,
5810 descriptors,
5811 created_at,
5812 hold,
5813 ))
5814}
5815
5816fn decode_step_execution_projection(row: &PgRow) -> Result<StepExecutionProjection, ExplorerError> {
5817 let unavailable = || ExplorerError::Repository(RepositoryError::Unavailable);
5818 let id = StepExecutionId::new(read_u64(row, "id").map_err(ExplorerError::Repository)?)
5819 .map_err(|_| unavailable())?;
5820 let job_execution_id =
5821 JobExecutionId::new(read_u64(row, "job_execution_id").map_err(ExplorerError::Repository)?)
5822 .map_err(|_| unavailable())?;
5823 let step_name = StepName::new(read_text(row, "step_name").map_err(ExplorerError::Repository)?)
5824 .map_err(|_| unavailable())?;
5825 let node_id = read_optional_text(row, "step_logical_id")
5826 .map_err(ExplorerError::Repository)?
5827 .map(NodeId::new)
5828 .transpose()
5829 .map_err(|_| unavailable())?;
5830 let status = decode_status(&read_text(row, "status").map_err(ExplorerError::Repository)?)
5831 .map_err(|_| unavailable())?;
5832 let exit_status = ExitStatus::new(
5833 ExitCode::new(read_text(row, "exit_code").map_err(ExplorerError::Repository)?)
5834 .map_err(|_| unavailable())?,
5835 );
5836 Ok(StepExecutionProjection::new(
5837 id,
5838 job_execution_id,
5839 step_name,
5840 node_id,
5841 status,
5842 exit_status,
5843 decode_projection_counts(row)?,
5844 ExecutionVersion::new(read_u64(row, "version").map_err(ExplorerError::Repository)?),
5845 decode_projection_timestamps(row)?,
5846 decode_projection_failure(row)?,
5847 Some(decode_state_descriptor(
5848 row,
5849 DurableStateKind::Checkpoint,
5850 "checkpoint_format",
5851 "checkpoint_schema",
5852 "checkpoint_schema_version",
5853 "checkpoint_bytes",
5854 )?),
5855 Some(decode_state_descriptor(
5856 row,
5857 DurableStateKind::ExecutionContext,
5858 "context_format",
5859 "context_schema",
5860 "context_schema_version",
5861 "context_bytes",
5862 )?),
5863 ))
5864}
5865
5866fn decode_step_partition_projection(row: &PgRow) -> Result<StepPartitionProjection, ExplorerError> {
5867 let unavailable = || ExplorerError::Repository(RepositoryError::Unavailable);
5868 let id = StepPartitionId::new(read_u64(row, "id").map_err(ExplorerError::Repository)?)
5869 .map_err(|_| unavailable())?;
5870 let step_execution_id = StepExecutionId::new(
5871 read_u64(row, "step_execution_id").map_err(ExplorerError::Repository)?,
5872 )
5873 .map_err(|_| unavailable())?;
5874 let worker = read_optional_u64(row, "worker_step_execution_id")
5875 .map_err(ExplorerError::Repository)?
5876 .map(StepExecutionId::new)
5877 .transpose()
5878 .map_err(|_| unavailable())?;
5879 let status = decode_status(&read_text(row, "status").map_err(ExplorerError::Repository)?)
5880 .map_err(|_| unavailable())?;
5881 let exit_status = ExitStatus::new(
5882 ExitCode::new(
5883 read_optional_text(row, "exit_code")
5884 .map_err(ExplorerError::Repository)?
5885 .unwrap_or_else(|| String::from("UNKNOWN")),
5886 )
5887 .map_err(|_| unavailable())?,
5888 );
5889 let ordinal = row
5890 .try_get::<i32, _>("partition_ordinal")
5891 .map_err(|_| unavailable())
5892 .and_then(|value| u32::try_from(value).map_err(|_| unavailable()))?;
5893 Ok(StepPartitionProjection::new(
5894 id,
5895 step_execution_id,
5896 read_text(row, "partition_key").map_err(ExplorerError::Repository)?,
5897 ordinal,
5898 status,
5899 exit_status,
5900 decode_projection_counts(row)?,
5901 ExecutionVersion::new(read_u64(row, "version").map_err(ExplorerError::Repository)?),
5902 worker,
5903 Some(decode_state_descriptor(
5904 row,
5905 DurableStateKind::ExecutionContext,
5906 "context_format",
5907 "context_schema",
5908 "context_schema_version",
5909 "context_bytes",
5910 )?),
5911 ))
5912}
5913
5914impl ExplorerRepository for PostgresExplorer {
5915 fn identity_ceiling<'a>(
5916 &'a self,
5917 query: &'a ExplorerQuery,
5918 ) -> BoxFuture<'a, Result<u64, ExplorerError>> {
5919 Box::pin(async move {
5920 let Some(source) = ceiling_source(query) else {
5921 return Err(ExplorerError::UnsupportedCapability);
5922 };
5923 let statement = sqlx::query(source);
5924 let statement = match query {
5925 ExplorerQuery::Instances { job_name } => statement.bind(job_name.as_str()),
5926 _ => statement,
5927 };
5928 let row = self.fetch_optional(statement).await?;
5929 let ceiling = row
5930 .as_ref()
5931 .map(|row| read_optional_i64(row, "ceiling"))
5932 .transpose()
5933 .map_err(ExplorerError::Repository)?
5934 .flatten()
5935 .unwrap_or(0);
5936 u64::try_from(ceiling)
5937 .map_err(|_| ExplorerError::Repository(RepositoryError::Unavailable))
5938 })
5939 }
5940
5941 fn job_names<'a>(
5942 &'a self,
5943 window: &'a QueryWindow,
5944 ) -> BoxFuture<'a, Result<Vec<JobName>, ExplorerError>> {
5945 Box::pin(async move {
5946 let rows = self
5947 .fetch_all(
5948 sqlx::query(
5949 "SELECT DISTINCT job_name FROM oxide_batch.ob_job_definition \
5950 WHERE id <= $1 AND ($2::text IS NULL OR job_name > $2) \
5951 ORDER BY job_name LIMIT $3",
5952 )
5953 .bind(window_ceiling(window)?)
5954 .bind(window_name(window)?)
5955 .bind(window_limit(window)),
5956 )
5957 .await?;
5958 rows.iter()
5959 .map(|row| {
5960 JobName::new(read_text(row, "job_name").map_err(ExplorerError::Repository)?)
5961 .map_err(|_| ExplorerError::Repository(RepositoryError::Unavailable))
5962 })
5963 .collect()
5964 })
5965 }
5966
5967 fn instances<'a>(
5968 &'a self,
5969 job_name: &'a JobName,
5970 window: &'a QueryWindow,
5971 ) -> BoxFuture<'a, Result<Vec<JobInstanceProjection>, ExplorerError>> {
5972 Box::pin(async move {
5973 let rows = self
5974 .fetch_all(
5975 sqlx::query(
5976 "SELECT id, job_name, instance_key, identifying_parameters, \
5977 (extract(epoch FROM created_at) * 1000)::bigint AS created_ms, \
5978 hold_actor, hold_reason, \
5979 (extract(epoch FROM hold_placed_at) * 1000)::bigint AS placed_ms \
5980 FROM oxide_batch.ob_job_instance \
5981 WHERE job_name = $1 AND id <= $2 \
5982 AND ($3::bigint IS NULL OR id < $3) \
5983 ORDER BY id DESC LIMIT $4",
5984 )
5985 .bind(job_name.as_str())
5986 .bind(window_ceiling(window)?)
5987 .bind(window_identity(window)?)
5988 .bind(window_limit(window)),
5989 )
5990 .await?;
5991 rows.iter().map(decode_job_instance_projection).collect()
5992 })
5993 }
5994
5995 fn executions<'a>(
5996 &'a self,
5997 job_instance_id: JobInstanceId,
5998 window: &'a QueryWindow,
5999 ) -> BoxFuture<'a, Result<Vec<JobExecutionProjection>, ExplorerError>> {
6000 Box::pin(async move {
6001 let after = window_ordered(window)?;
6002 let rows = self
6003 .fetch_all(
6004 sqlx::query(AssertSqlSafe(job_execution_projection_select(
6005 "WHERE execution.job_instance_id = $1 AND execution.id <= $2 \
6006 AND ($3::bigint IS NULL \
6007 OR (execution.attempt, execution.id) < ($3, $4)) \
6008 ORDER BY execution.attempt DESC, execution.id DESC LIMIT $5",
6009 )))
6010 .bind(explorer_id(
6011 job_instance_id.get(),
6012 IdentifierKind::JobInstance,
6013 )?)
6014 .bind(window_ceiling(window)?)
6015 .bind(after.map(|(primary, _)| primary))
6016 .bind(after.map_or(0, |(_, identity)| identity))
6017 .bind(window_limit(window)),
6018 )
6019 .await?;
6020 rows.iter().map(decode_job_execution_projection).collect()
6021 })
6022 }
6023
6024 fn execution(
6025 &self,
6026 job_execution_id: JobExecutionId,
6027 ) -> BoxFuture<'_, Result<Option<JobExecutionProjection>, ExplorerError>> {
6028 Box::pin(async move {
6029 let row = self
6030 .fetch_optional(
6031 sqlx::query(AssertSqlSafe(job_execution_projection_select(
6032 "WHERE execution.id = $1",
6033 )))
6034 .bind(explorer_id(
6035 job_execution_id.get(),
6036 IdentifierKind::JobExecution,
6037 )?),
6038 )
6039 .await?;
6040 row.as_ref()
6041 .map(decode_job_execution_projection)
6042 .transpose()
6043 })
6044 }
6045
6046 fn step_executions<'a>(
6047 &'a self,
6048 job_execution_id: JobExecutionId,
6049 window: &'a QueryWindow,
6050 ) -> BoxFuture<'a, Result<Vec<StepExecutionProjection>, ExplorerError>> {
6051 Box::pin(async move {
6052 let rows = self
6053 .fetch_all(
6054 sqlx::query(AssertSqlSafe(step_execution_projection_select(
6055 "WHERE execution.job_execution_id = $1 AND execution.id <= $2 \
6056 AND ($3::bigint IS NULL OR execution.id > $3) \
6057 ORDER BY execution.id LIMIT $4",
6058 )))
6059 .bind(explorer_id(
6060 job_execution_id.get(),
6061 IdentifierKind::JobExecution,
6062 )?)
6063 .bind(window_ceiling(window)?)
6064 .bind(window_identity(window)?)
6065 .bind(window_limit(window)),
6066 )
6067 .await?;
6068 rows.iter().map(decode_step_execution_projection).collect()
6069 })
6070 }
6071
6072 fn unresolved_executions<'a>(
6073 &'a self,
6074 minimum_age: Duration,
6075 window: &'a QueryWindow,
6076 ) -> BoxFuture<'a, Result<Vec<JobExecutionProjection>, ExplorerError>> {
6077 Box::pin(async move {
6078 let seconds = i64::try_from(minimum_age.as_secs())
6079 .map_err(|_| ExplorerError::Repository(RepositoryError::Unavailable))?;
6080 let rows = self
6081 .fetch_all(
6082 sqlx::query(AssertSqlSafe(job_execution_projection_select(
6083 "WHERE execution.status IN \
6084 ('STARTING', 'STARTED', 'STOPPING', 'UNKNOWN') \
6085 AND execution.updated_at \
6086 < CURRENT_TIMESTAMP - make_interval(secs => $1::double precision) \
6087 AND execution.id <= $2 AND ($3::bigint IS NULL OR execution.id > $3) \
6088 ORDER BY execution.id LIMIT $4",
6089 )))
6090 .bind(seconds)
6091 .bind(window_ceiling(window)?)
6092 .bind(window_identity(window)?)
6093 .bind(window_limit(window)),
6094 )
6095 .await?;
6096 rows.iter().map(decode_job_execution_projection).collect()
6097 })
6098 }
6099
6100 fn recovery_decisions<'a>(
6101 &'a self,
6102 job_execution_id: JobExecutionId,
6103 window: &'a QueryWindow,
6104 ) -> BoxFuture<'a, Result<Vec<RecoveryDecision>, ExplorerError>> {
6105 Box::pin(async move {
6106 let rows = self
6107 .fetch_all(
6108 sqlx::query(
6109 "SELECT id, execution_version, prior_status, resulting_status, \
6110 reason_code, operator_reference, evidence_digest, \
6111 (extract(epoch FROM decided_at) * 1000)::bigint AS decided_ms \
6112 FROM oxide_batch.ob_recovery_decision \
6113 WHERE job_execution_id = $1 AND id <= $2 \
6114 AND ($3::bigint IS NULL OR id > $3) \
6115 ORDER BY id LIMIT $4",
6116 )
6117 .bind(explorer_id(
6118 job_execution_id.get(),
6119 IdentifierKind::JobExecution,
6120 )?)
6121 .bind(window_ceiling(window)?)
6122 .bind(window_identity(window)?)
6123 .bind(window_limit(window)),
6124 )
6125 .await?;
6126 rows.iter()
6127 .map(|row| {
6128 decode_recovery_decision(job_execution_id, row)
6129 .map_err(ExplorerError::Repository)
6130 })
6131 .collect()
6132 })
6133 }
6134
6135 fn flow_decisions<'a>(
6136 &'a self,
6137 job_execution_id: JobExecutionId,
6138 window: &'a QueryWindow,
6139 ) -> BoxFuture<'a, Result<Vec<FlowDecision>, ExplorerError>> {
6140 Box::pin(async move {
6141 let after = window_ordered(window)?;
6142 let rows = self
6143 .fetch_all(
6144 sqlx::query(AssertSqlSafe(flow_decision_select(
6145 "WHERE decision.job_execution_id = $1 AND decision.id <= $2 \
6146 AND ($3::bigint IS NULL \
6147 OR (decision.sequence, decision.id) > ($3, $4)) \
6148 ORDER BY decision.sequence, decision.id LIMIT $5",
6149 )))
6150 .bind(explorer_id(
6151 job_execution_id.get(),
6152 IdentifierKind::JobExecution,
6153 )?)
6154 .bind(window_ceiling(window)?)
6155 .bind(after.map(|(primary, _)| primary))
6156 .bind(after.map_or(0, |(_, identity)| identity))
6157 .bind(window_limit(window)),
6158 )
6159 .await?;
6160 rows.iter()
6161 .map(|row| decode_flow_decision(row).map_err(ExplorerError::Repository))
6162 .collect()
6163 })
6164 }
6165
6166 fn step_partitions<'a>(
6167 &'a self,
6168 step_execution_id: StepExecutionId,
6169 window: &'a QueryWindow,
6170 ) -> BoxFuture<'a, Result<Vec<StepPartitionProjection>, ExplorerError>> {
6171 Box::pin(async move {
6172 let rows = self
6173 .fetch_all(
6174 sqlx::query(AssertSqlSafe(step_partition_projection_select(
6175 "WHERE partition.step_execution_id = $1 AND partition.id <= $2 \
6176 AND ($3::bigint IS NULL OR partition.id > $3) \
6177 ORDER BY partition.id LIMIT $4",
6178 )))
6179 .bind(explorer_id(
6180 step_execution_id.get(),
6181 IdentifierKind::StepExecution,
6182 )?)
6183 .bind(window_ceiling(window)?)
6184 .bind(window_identity(window)?)
6185 .bind(window_limit(window)),
6186 )
6187 .await?;
6188 rows.iter().map(decode_step_partition_projection).collect()
6189 })
6190 }
6191
6192 fn operator_requests<'a>(
6193 &'a self,
6194 job_execution_id: JobExecutionId,
6195 window: &'a QueryWindow,
6196 ) -> BoxFuture<'a, Result<Vec<OperatorRecord>, ExplorerError>> {
6197 Box::pin(async move {
6198 let rows = self
6199 .fetch_all(
6200 sqlx::query(AssertSqlSafe(operator_request_select(
6201 "WHERE request.job_execution_id = $1 AND request.id <= $2 \
6202 AND ($3::bigint IS NULL OR request.id > $3) \
6203 ORDER BY request.id LIMIT $4",
6204 )))
6205 .bind(explorer_id(
6206 job_execution_id.get(),
6207 IdentifierKind::JobExecution,
6208 )?)
6209 .bind(window_ceiling(window)?)
6210 .bind(window_identity(window)?)
6211 .bind(window_limit(window)),
6212 )
6213 .await?;
6214 rows.iter()
6215 .map(|row| decode_operator_record(row).map_err(ExplorerError::Repository))
6216 .collect()
6217 })
6218 }
6219}
6220
6221impl crate::RecoveryRepository for PostgresExplorer {
6222 #[allow(
6223 clippy::too_many_lines,
6224 reason = "one bounded snapshot query keeps its redacted decode mapping adjacent"
6225 )]
6226 fn recovery_snapshot<'a>(
6227 &'a self,
6228 execution_id: JobExecutionId,
6229 current_owner: &'a crate::OwnerToken,
6230 ) -> BoxFuture<'a, Result<crate::RecoverySnapshot, RepositoryError>> {
6231 Box::pin(async move {
6232 let row = self
6233 .fetch_optional(
6234 sqlx::query(
6235 "SELECT execution.status, execution.attempt, execution.version, \
6236 (extract(epoch FROM execution.updated_at) * 1000)::bigint AS updated_ms, \
6237 (extract(epoch FROM clock_timestamp()) * 1000)::bigint AS server_ms, \
6238 CASE WHEN execution.owner_token IS NULL THEN 'ABSENT' \
6239 WHEN execution.owner_token = $2 THEN 'CURRENT' ELSE 'OTHER' END \
6240 AS owner_observation, \
6241 latest_step.id AS step_id, latest_step.status AS step_status, \
6242 latest_step.checkpoint_format, latest_step.checkpoint_schema, \
6243 latest_step.checkpoint_schema_version, \
6244 pg_column_size(latest_step.checkpoint_payload)::bigint AS checkpoint_bytes, \
6245 (execution.status = 'UNKNOWN' \
6246 OR COALESCE(execution.failure_category = 'UNKNOWN_COMMIT', false) \
6247 OR COALESCE(latest_step.status = 'UNKNOWN', false)) \
6248 AS unknown_commit, \
6249 EXISTS (SELECT 1 FROM oxide_batch.ob_step_partition partition \
6250 JOIN oxide_batch.ob_step_execution parent \
6251 ON parent.id = partition.step_execution_id \
6252 WHERE parent.job_execution_id = execution.id \
6253 AND partition.status = 'COMPLETED') AS completed_partition, \
6254 EXISTS (SELECT 1 FROM oxide_batch.ob_flow_decision decision \
6255 WHERE decision.job_execution_id = execution.id) \
6256 AS committed_flow_decision, \
6257 COALESCE(definition.manifest ->> 'delivery_mode', 'ambiguous') \
6258 <> 'atomic_same_resource' AS ambiguous_external_effect \
6259 FROM oxide_batch.ob_job_execution execution \
6260 JOIN oxide_batch.ob_job_definition definition \
6261 ON definition.id = execution.definition_id \
6262 LEFT JOIN LATERAL ( \
6263 SELECT step.id, step.status, step.checkpoint_format, \
6264 step.checkpoint_schema, step.checkpoint_schema_version, \
6265 step.checkpoint_payload \
6266 FROM oxide_batch.ob_step_execution step \
6267 WHERE step.job_execution_id = execution.id \
6268 ORDER BY step.id DESC LIMIT 1 \
6269 ) latest_step ON true \
6270 WHERE execution.id = $1",
6271 )
6272 .bind(database_id(
6273 execution_id.get(),
6274 IdentifierKind::JobExecution,
6275 )?)
6276 .bind(¤t_owner.as_bytes()[..]),
6277 )
6278 .await
6279 .map_err(|error| match error {
6280 ExplorerError::Repository(error) => error,
6281 _ => RepositoryError::Unavailable,
6282 })?
6283 .ok_or(RepositoryError::JobExecutionNotFound { id: execution_id })?;
6284 let owner = match read_text(&row, "owner_observation")?.as_str() {
6285 "ABSENT" => crate::OwnerObservation::Absent,
6286 "CURRENT" => crate::OwnerObservation::CurrentProcess,
6287 "OTHER" => crate::OwnerObservation::OtherProcess,
6288 _ => return Err(RepositoryError::Unavailable),
6289 };
6290 let latest_step = match read_optional_u64(&row, "step_id")? {
6291 Some(id) => {
6292 let descriptor = decode_state_descriptor(
6293 &row,
6294 DurableStateKind::Checkpoint,
6295 "checkpoint_format",
6296 "checkpoint_schema",
6297 "checkpoint_schema_version",
6298 "checkpoint_bytes",
6299 )
6300 .map_err(|_| RepositoryError::Unavailable)?;
6301 Some(crate::RecoveryStepEvidence::new(
6302 StepExecutionId::new(id)?,
6303 decode_status(&read_text(&row, "step_status")?)?,
6304 Some(descriptor),
6305 ))
6306 }
6307 None => None,
6308 };
6309 Ok(crate::RecoverySnapshot::new(
6310 execution_id,
6311 decode_status(&read_text(&row, "status")?)?,
6312 u32::try_from(
6317 row.try_get::<i32, _>("attempt")
6318 .map_err(|_| RepositoryError::Unavailable)?,
6319 )
6320 .map_err(|_| RepositoryError::Unavailable)?,
6321 ExecutionVersion::new(read_u64(&row, "version")?),
6322 owner,
6323 millis_system_time(read_i64(&row, "updated_ms")?)?,
6324 millis_system_time(read_i64(&row, "server_ms")?)?,
6325 latest_step,
6326 crate::RecoveryMarkers::new()
6327 .with_unknown_commit(
6328 row.try_get::<bool, _>("unknown_commit")
6329 .map_err(|_| RepositoryError::Unavailable)?,
6330 )
6331 .with_completed_partition(
6332 row.try_get::<bool, _>("completed_partition")
6333 .map_err(|_| RepositoryError::Unavailable)?,
6334 )
6335 .with_committed_flow_decision(
6336 row.try_get::<bool, _>("committed_flow_decision")
6337 .map_err(|_| RepositoryError::Unavailable)?,
6338 )
6339 .with_ambiguous_external_effect(
6340 row.try_get::<bool, _>("ambiguous_external_effect")
6341 .map_err(|_| RepositoryError::Unavailable)?,
6342 ),
6343 ))
6344 })
6345 }
6346}