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