Skip to content

fix(postgres): concurrent UpdateItem upserts creating one item must all succeed #250

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
LeeroyHannigan merged 2 commits into from
Aug 13, 2026
+391 −125
Merged

fix(postgres): concurrent UpdateItem upserts creating one item must all succeed #250

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
24912da
fix(postgres): concurrent UpdateItem upserts creating one item must a…
LeeroyHannigan Aug 10, 2026
8dfc358
test: prove same-attribute concurrent updates serialize into a linear…
LeeroyHannigan Aug 10, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
275 changes: 150 additions & 125 deletions crates/storage-postgres/src/data/update_item.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
//! `update_item` implementation for the `PostgreSQL` backend.

use extenddb_core::expression::{self, Expr, ExpressionMaps, UpdateAction};
use extenddb_core::types::{Item, KeyType, TableKeyInfo};
use extenddb_core::types::{Item, TableKeyInfo};
use extenddb_core::validation;
use extenddb_storage::StreamCapture;
use extenddb_storage::error::StorageError;
Expand Down Expand Up @@ -54,153 +54,178 @@ impl PostgresEngine {
.load(std::sync::atomic::Ordering::Relaxed)
};

// Fetch existing item
let old_json = if let Some((sk_name, sk_type)) =
// Sort-key binding, computed once: the key is immutable across attempts.
let sk_parts = if let Some((sk_name, sk_type)) =
sk_info(&key_info.key_schema, &key_info.attribute_definitions)
{
let sk_value = key
.get(sk_name)
.ok_or_else(|| StorageError::Internal("missing sort key".to_owned()))?;
let sk = parse_sk(sk_value, sk_type)?;
let sk_col = sk_column(sk_type);
let select_sql = format!(
"SELECT item_data FROM {ddb_table} WHERE pk = $1 AND {sk_col} = $2 FOR UPDATE"
);
let row: Option<(serde_json::Value,)> =
bind_sk_fetch_optional!(&select_sql, pk_text.as_ref(), &sk, &mut *tx)?;
row.map(|(v,)| v)
Some((parse_sk(sk_value, sk_type)?, sk_column(sk_type)))
} else {
let select_sql = format!("SELECT item_data FROM {ddb_table} WHERE pk = $1 FOR UPDATE");
let row: Option<(serde_json::Value,)> = sqlx::query_as(&select_sql)
.bind(pk_text.as_ref())
.fetch_optional(&mut *tx)
.await
.map_err(|e| StorageError::Internal(e.to_string()))?;
row.map(|(v,)| v)
None
};

// Build the working item: existing or new with key attributes only (upsert)
let mut item = if let Some(json) = old_json.clone() {
json_to_item(json)?
} else {
key.clone()
};
// Read-modify-write, retried when a concurrent writer creates the row first.
//
// The locking read below serialises writers to an item that EXISTS. It cannot
// lock a row that does not exist yet, so two writers can both decide to
// insert. The loser's `ON CONFLICT DO NOTHING` affects no rows, and at that
// moment the winner may still be uncommitted, so a plain re-read can see
// nothing at all. Re-reading `FOR UPDATE` is what makes this terminate: it
// blocks until the winner commits, returning its row, or aborts, returning
// none so the insert can be retried.
//
// Returning ConditionalCheckFailedException here instead was wrong. DynamoDB
// serialises writes to one item, and an UpdateItem carrying no condition is
// an upsert, so every writer must succeed and each update expression applies
// on top of whatever the previous writer committed. Measured against the
// service: four writers setting four different attributes on one new key
// yield an item holding all four, so the loser's expression must be re-applied
// to the winner's item rather than overwriting it.
//
// A supplied condition is re-evaluated against the winner for the same
// reason: after losing the race `attribute_exists` is now true, and failing
// it unconditionally was wrong in the opposite direction.
const MAX_CREATE_RACE_ATTEMPTS: u32 = 5;
let mut attempt: u32 = 0;
let (old_item, new_item, item, pre_mutation_item) = loop {
attempt += 1;

// Save pre-mutation item for index sync and stream capture.
let pre_mutation_item = if (!indexes.is_empty() || stream.is_some()) && old_json.is_some() {
Some(item.clone())
} else {
None
};
// Current row, locked when present. First pass: the initial read. Later
// passes: the post-conflict re-read that waits on the race winner.
let old_json: Option<serde_json::Value> = match &sk_parts {
Some((sk, sk_col)) => {
let select_sql = format!(
"SELECT item_data FROM {ddb_table} WHERE pk = $1 AND {sk_col} = $2 FOR UPDATE"
);
let row: Option<(serde_json::Value,)> =
bind_sk_fetch_optional!(&select_sql, pk_text.as_ref(), sk, &mut *tx)?;
row.map(|(v,)| v)
}
None => {
let select_sql =
format!("SELECT item_data FROM {ddb_table} WHERE pk = $1 FOR UPDATE");
let row: Option<(serde_json::Value,)> = sqlx::query_as(&select_sql)
.bind(pk_text.as_ref())
.fetch_optional(&mut *tx)
.await
.map_err(|e| StorageError::Internal(e.to_string()))?;
row.map(|(v,)| v)
}
};

let old_item = if return_old && old_json.is_some() {
Some(item.clone())
} else {
None
};
// Build the working item: existing or new with key attributes only (upsert)
let mut item = if let Some(json) = old_json.clone() {
json_to_item(json)?
} else {
key.clone()
};

// Evaluate condition against the existing item (empty if non-existent).
// DynamoDB treats a non-existent item as having no attributes at all.
let condition_item = if old_json.is_some() {
&item
} else {
&std::collections::BTreeMap::new()
};
match check_condition(condition, condition_item, maps) {
Ok(()) => {}
Err(StorageError::ConditionFailed(_)) => {
if old_json.is_some() {
return Err(StorageError::ConditionFailed(Some(item)));
// Save pre-mutation item for index sync and stream capture.
let pre_mutation_item =
if (!indexes.is_empty() || stream.is_some()) && old_json.is_some() {
Some(item.clone())
} else {
None
};

let old_item = if return_old && old_json.is_some() {
Some(item.clone())
} else {
None
};

// Evaluate condition against the existing item (empty if non-existent).
// DynamoDB treats a non-existent item as having no attributes at all.
let empty = std::collections::BTreeMap::new();
let condition_item = if old_json.is_some() { &item } else { &empty };
match check_condition(condition, condition_item, maps) {
Ok(()) => {}
Err(StorageError::ConditionFailed(_)) => {
if old_json.is_some() {
return Err(StorageError::ConditionFailed(Some(item)));
}
return Err(StorageError::ConditionFailed(None));
}
return Err(StorageError::ConditionFailed(None));
Err(e) => return Err(e),
}
Err(e) => return Err(e),
}

// Apply update actions
expression::apply_update(actions, &mut item, maps)
.map_err(|e| StorageError::Validation(e.to_string()))?;
// Apply update actions
expression::apply_update(actions, &mut item, maps)
.map_err(|e| StorageError::Validation(e.to_string()))?;

// Validate post-update item size (400 KB limit)
validation::validate_item_size(&item, self.max_item_size_bytes)
.map_err(|e| StorageError::Validation(e.to_string()))?;
// Validate post-update item size (400 KB limit)
validation::validate_item_size(&item, self.max_item_size_bytes)
.map_err(|e| StorageError::Validation(e.to_string()))?;

let new_item = if return_new { Some(item.clone()) } else { None };
let new_item = if return_new { Some(item.clone()) } else { None };

// Write the updated item back
let item_json =
serde_json::to_value(&item).map_err(|e| StorageError::Internal(e.to_string()))?;
let item_json =
serde_json::to_value(&item).map_err(|e| StorageError::Internal(e.to_string()))?;

if let Some((_, sk_type)) = sk_info(&key_info.key_schema, &key_info.attribute_definitions) {
let sk_name_ref = key_info
.key_schema
.iter()
.find(|ks| ks.key_type == KeyType::Range)
.map(|ks| ks.attribute_name.as_str())
.ok_or_else(|| StorageError::Internal("missing sort key schema".to_owned()))?;
let sk_value = key
.get(sk_name_ref)
.ok_or_else(|| StorageError::Internal("missing sort key".to_owned()))?;
let sk = parse_sk(sk_value, sk_type)?;
let sk_col = sk_column(sk_type);
if old_json.is_some() {
// Row existed — update in place.
let update_sql = format!(
"UPDATE {ddb_table} SET item_data = $3 WHERE pk = $1 AND {sk_col} = $2"
);
bind_sk_execute!(&update_sql, pk_text.as_ref(), &sk, &item_json, &mut *tx)?;
} else {
// Row didn't exist (upsert) — atomic insert, fail if someone beat us.
let insert_sql = format!(
"INSERT INTO {ddb_table} (pk, {sk_col}, item_data) VALUES ($1, $2, $3) \
ON CONFLICT (pk, {sk_col}) DO NOTHING"
);
let result =
bind_sk_execute!(&insert_sql, pk_text.as_ref(), &sk, &item_json, &mut *tx)?;
if result.rows_affected() == 0 {
// Another transaction inserted between our SELECT and INSERT.
// Fetch the winner to return with ConditionFailed.
let winner_sql = format!(
"SELECT item_data FROM {ddb_table} WHERE pk = $1 AND {sk_col} = $2"
// Row existed and is locked by the read above, so update in place.
match &sk_parts {
Some((sk, sk_col)) => {
let update_sql = format!(
"UPDATE {ddb_table} SET item_data = $3 WHERE pk = $1 AND {sk_col} = $2"
);
bind_sk_execute!(&update_sql, pk_text.as_ref(), sk, &item_json, &mut *tx)?;
}
None => {
let update_sql =
format!("UPDATE {ddb_table} SET item_data = $2 WHERE pk = $1");
sqlx::query(&update_sql)
.bind(pk_text.as_ref())
.bind(&item_json)
.execute(&mut *tx)
.await
.map_err(|e| StorageError::Internal(e.to_string()))?;
}
}
break (old_item, new_item, item, pre_mutation_item);
}

// Row absent, so insert atomically; someone may beat us to it.
let inserted = match &sk_parts {

@jcshepherd jcshepherd Aug 13, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

General comment: Rust seems to allow a great deal of mashing together assignment statements, control statements and execution statements. In this change, I'm finding it's making it very hard to understand the control structures. It's probably not idiomatic Rust, but this would really benefit from breaking it up into some sub-functions that do the actual db access, so that the flow of control is more obvious in the main routine.

Some((sk, sk_col)) => {
let insert_sql = format!(
"INSERT INTO {ddb_table} (pk, {sk_col}, item_data) VALUES ($1, $2, $3) \
ON CONFLICT (pk, {sk_col}) DO NOTHING"
);
bind_sk_execute!(&insert_sql, pk_text.as_ref(), sk, &item_json, &mut *tx)?
.rows_affected()
== 1
}
None => {
let insert_sql = format!(
"INSERT INTO {ddb_table} (pk, item_data) VALUES ($1, $2) \
ON CONFLICT (pk) DO NOTHING"
);
let winner: Option<(serde_json::Value,)> =
bind_sk_fetch_optional!(&winner_sql, pk_text.as_ref(), &sk, &mut *tx)?;
let winner_item = winner.map(|(v,)| json_to_item(v)).transpose()?;
return Err(StorageError::ConditionFailed(winner_item));
sqlx::query(&insert_sql)
.bind(pk_text.as_ref())
.bind(&item_json)
.execute(&mut *tx)
.await
.map_err(|e| StorageError::Internal(e.to_string()))?
.rows_affected()
== 1
}
};
if inserted {
break (old_item, new_item, item, pre_mutation_item);
}
} else if old_json.is_some() {
let update_sql = format!("UPDATE {ddb_table} SET item_data = $2 WHERE pk = $1");
sqlx::query(&update_sql)
.bind(pk_text.as_ref())
.bind(&item_json)
.execute(&mut *tx)
.await
.map_err(|e| StorageError::Internal(e.to_string()))?;
} else {
let insert_sql = format!(
"INSERT INTO {ddb_table} (pk, item_data) VALUES ($1, $2) \
ON CONFLICT (pk) DO NOTHING"
);
let result = sqlx::query(&insert_sql)
.bind(pk_text.as_ref())
.bind(&item_json)
.execute(&mut *tx)
.await
.map_err(|e| StorageError::Internal(e.to_string()))?;
if result.rows_affected() == 0 {
// Another transaction inserted between our SELECT and INSERT.
// Fetch the winner to return with ConditionFailed.
let winner_sql = format!("SELECT item_data FROM {ddb_table} WHERE pk = $1");
let winner: Option<(serde_json::Value,)> = sqlx::query_as(&winner_sql)
.bind(pk_text.as_ref())
.fetch_optional(&mut *tx)
.await
.map_err(|e| StorageError::Internal(e.to_string()))?;
let winner_item = winner.map(|(v,)| json_to_item(v)).transpose()?;
return Err(StorageError::ConditionFailed(winner_item));

// Lost the create race. Loop round: the locking read at the top waits for
// the winner, then this item's update expression is applied on top of it.
if attempt >= MAX_CREATE_RACE_ATTEMPTS {
return Err(StorageError::Internal(format!(
"UpdateItem could not create the item after {attempt} attempts: a \
concurrent writer repeatedly created and rolled back the row"
)));
}
}
};

// Sync GSI/LSI update within transaction (D-4).
if !indexes.is_empty() {
Expand Down
Loading
Loading