runifold_store_postgres/conversation/
durable.rs

1//! `PostgreSQL` checkpoint persistence and atomic durable conversation commits.
2
3use runifold_agent::{
4    ConversationStoreError, ConversationStoreFuture, ConversationVersion,
5    DurableConversationCommit, DurableConversationStore,
6};
7use runifold_core::{
8    Checkpoint, CheckpointError, CheckpointErrorKind, CheckpointId, CheckpointStore,
9};
10use serde_json::Value;
11
12use super::{
13    PostgresConversationStore,
14    support::{
15        conflict_error, conversation_uuid, decode_version, encode_error, storage_error, to_i64,
16        validate_append,
17    },
18};
19
20impl CheckpointStore for PostgresConversationStore {
21    fn load(&self, id: CheckpointId) -> Result<Checkpoint, CheckpointError> {
22        let sql = format!(
23            "SELECT record_json FROM {}_checkpoints WHERE checkpoint_id = $1",
24            self.table
25        );
26        let row = self
27            .blocking()
28            .execute(move |client| client.query_opt(&sql, &[&id.as_uuid()]))
29            .map_err(|_| checkpoint_worker())?
30            .map_err(checkpoint_storage)?
31            .ok_or_else(|| checkpoint_not_found(id))?;
32        serde_json::from_value(row.get::<_, Value>("record_json")).map_err(|_| {
33            CheckpointError::new(
34                CheckpointErrorKind::InvalidPayload,
35                "PostgreSQL checkpoint payload is invalid",
36            )
37        })
38    }
39
40    fn compare_and_swap(
41        &self,
42        checkpoint: &Checkpoint,
43        expected_revision: Option<u64>,
44    ) -> Result<(), CheckpointError> {
45        let revision = checkpoint_i64(checkpoint.revision)?;
46        let expected = expected_revision.map(checkpoint_i64).transpose()?;
47        let record = serde_json::to_value(checkpoint).map_err(|_| {
48            CheckpointError::new(
49                CheckpointErrorKind::InvalidPayload,
50                "checkpoint could not be encoded",
51            )
52        })?;
53        let table = format!("{}_checkpoints", self.table);
54        let checkpoint_id = checkpoint.id;
55        self.blocking()
56            .execute(move |client| -> Result<(), CheckpointError> {
57                let mut transaction = client.transaction().map_err(checkpoint_storage)?;
58                match expected {
59                    None if revision == 0 => {
60                        let changed = transaction
61                            .execute(
62                                &format!(
63                                    "INSERT INTO {table} (checkpoint_id, revision, record_json) \
64                                     VALUES ($1, $2, $3) \
65                                     ON CONFLICT (checkpoint_id) DO NOTHING"
66                                ),
67                                &[&checkpoint_id.as_uuid(), &revision, &record],
68                            )
69                            .map_err(checkpoint_storage)?;
70                        if changed != 1 {
71                            return Err(checkpoint_conflict(checkpoint_id));
72                        }
73                    }
74                    Some(expected)
75                        if expected.checked_add(1).is_some_and(|next| revision == next) =>
76                    {
77                        let changed = transaction
78                            .execute(
79                                &format!(
80                                    "UPDATE {table} SET revision = $1, record_json = $2 \
81                                     WHERE checkpoint_id = $3 AND revision = $4"
82                                ),
83                                &[&revision, &record, &checkpoint_id.as_uuid(), &expected],
84                            )
85                            .map_err(checkpoint_storage)?;
86                        if changed != 1 {
87                            let exists = transaction
88                                .query_opt(
89                                    &format!("SELECT 1 FROM {table} WHERE checkpoint_id = $1"),
90                                    &[&checkpoint_id.as_uuid()],
91                                )
92                                .map_err(checkpoint_storage)?
93                                .is_some();
94                            return Err(if exists {
95                                checkpoint_conflict(checkpoint_id)
96                            } else {
97                                checkpoint_not_found(checkpoint_id)
98                            });
99                        }
100                    }
101                    _ => return Err(checkpoint_conflict(checkpoint_id)),
102                }
103                transaction.commit().map_err(checkpoint_storage)
104            })
105            .map_err(|_| checkpoint_worker())?
106    }
107}
108
109impl DurableConversationStore for PostgresConversationStore {
110    fn commit_durable_turn(
111        &self,
112        command: DurableConversationCommit,
113    ) -> ConversationStoreFuture<'_, Result<ConversationVersion, ConversationStoreError>> {
114        Box::pin(async move {
115            validate_append(&command.append)?;
116            let expected_checkpoint = to_i64(command.expected_checkpoint_revision)?;
117            let checkpoint_revision = to_i64(command.checkpoint.revision)?;
118            if expected_checkpoint
119                .checked_add(1)
120                .is_none_or(|next| next != checkpoint_revision)
121            {
122                return Err(conflict_error(
123                    "durable conversation checkpoint is not the expected successor",
124                ));
125            }
126            let messages = serde_json::to_value(&command.append.messages).map_err(encode_error)?;
127            let checkpoint = serde_json::to_value(&command.checkpoint).map_err(encode_error)?;
128            let mut client = self.transaction_client.lock().await;
129            let transaction = client.transaction().await.map_err(storage_error)?;
130            let conversation_sql = format!(
131                "UPDATE {table} SET version = version + 1, updated_at = clock_timestamp() \
132                 WHERE conversation_id = $1 AND namespace = $2 AND version = $3 \
133                   AND version < 9223372036854775807 RETURNING version",
134                table = self.table
135            );
136            let Some(version_row) = transaction
137                .query_opt(
138                    &conversation_sql,
139                    &[
140                        &conversation_uuid(command.append.conversation_id),
141                        &command.namespace.as_str(),
142                        &to_i64(command.append.expected_version.get())?,
143                    ],
144                )
145                .await
146                .map_err(storage_error)?
147            else {
148                return Err(conflict_error(
149                    "durable conversation transcript version precondition failed",
150                ));
151            };
152            let transcript_sql = format!(
153                "WITH base AS (\
154                     SELECT COALESCE(MAX(sequence), 0) AS last_sequence \
155                     FROM {table}_transcript WHERE conversation_id = $1\
156                 ) \
157                 INSERT INTO {table}_transcript (conversation_id, sequence, message) \
158                 SELECT $1, base.last_sequence + payload.ordinality, payload.message \
159                 FROM base CROSS JOIN LATERAL \
160                    jsonb_array_elements($2::JSONB) WITH ORDINALITY AS payload(message, ordinality)",
161                table = self.table
162            );
163            let inserted = transaction
164                .execute(
165                    &transcript_sql,
166                    &[
167                        &conversation_uuid(command.append.conversation_id),
168                        &messages,
169                    ],
170                )
171                .await
172                .map_err(storage_error)?;
173            if inserted != u64::try_from(command.append.messages.len()).unwrap_or(u64::MAX) {
174                return Err(conflict_error(
175                    "durable conversation transcript append was incomplete",
176                ));
177            }
178            let checkpoint_sql = format!(
179                "UPDATE {table}_checkpoints SET revision = $1, record_json = $2 \
180                 WHERE checkpoint_id = $3 AND revision = $4",
181                table = self.table
182            );
183            let changed = transaction
184                .execute(
185                    &checkpoint_sql,
186                    &[
187                        &checkpoint_revision,
188                        &checkpoint,
189                        &command.checkpoint.id.as_uuid(),
190                        &expected_checkpoint,
191                    ],
192                )
193                .await
194                .map_err(storage_error)?;
195            if changed != 1 {
196                return Err(conflict_error(
197                    "durable conversation checkpoint precondition failed",
198                ));
199            }
200            transaction.commit().await.map_err(storage_error)?;
201            decode_version(version_row.get("version"))
202        })
203    }
204}
205
206fn checkpoint_i64(value: u64) -> Result<i64, CheckpointError> {
207    i64::try_from(value).map_err(|_| {
208        CheckpointError::new(
209            CheckpointErrorKind::InvalidPayload,
210            "checkpoint revision exceeds PostgreSQL BIGINT",
211        )
212    })
213}
214
215fn checkpoint_storage(_error: postgres::Error) -> CheckpointError {
216    CheckpointError::new(
217        CheckpointErrorKind::Storage,
218        "PostgreSQL checkpoint operation failed",
219    )
220}
221
222fn checkpoint_worker() -> CheckpointError {
223    CheckpointError::new(
224        CheckpointErrorKind::Storage,
225        "PostgreSQL checkpoint worker is unavailable",
226    )
227}
228
229fn checkpoint_conflict(id: CheckpointId) -> CheckpointError {
230    CheckpointError::new(
231        CheckpointErrorKind::Conflict,
232        format!("checkpoint `{id}` revision precondition failed"),
233    )
234}
235
236fn checkpoint_not_found(id: CheckpointId) -> CheckpointError {
237    CheckpointError::new(
238        CheckpointErrorKind::NotFound,
239        format!("checkpoint `{id}` does not exist"),
240    )
241}
runifold_store_postgres/conversation/durable.rs

runifold_store_postgres/conversation/
durable.rs
