use deadpool_postgres::Pool;
use super::PostgresBackend;
use crate::backend::{clamped, Backend, ReadCursor, StagingCursor};
use crate::pg_common::{err, is_undefined_table, size_u64};
use crate::Error;
#[allow(dead_code)]
impl PostgresBackend {
pub(crate) async fn delete_batch(&self, keys: &[Vec<u8>]) -> Result<u64, Error> {
if keys.is_empty() {
return Ok(0);
}
let client = self.pool.get().await.map_err(err)?;
let sql = format!("DELETE FROM {} WHERE key = ANY($1)", self.table);
let stmt = match client.prepare_cached(&sql).await {
Ok(stmt) => stmt,
Err(e) if is_undefined_table(&e) => return Ok(0),
Err(e) => return Err(err(e)),
};
client.execute(&stmt, &[&keys]).await.map_err(err)
}
}
pub struct PostgresCursor {
value: Vec<u8>,
len: u64,
}
impl ReadCursor for PostgresCursor {
fn len(&self) -> u64 {
self.len
}
async fn read_at(&self, offset: u64, len: u64) -> Result<Vec<u8>, Error> {
let (start, end) = clamped(offset, len, self.value.len());
Ok(self.value[start..end].to_vec())
}
}
pub struct PostgresStaging {
staged: Vec<u8>,
pool: Pool,
table: String,
#[cfg(test)]
fail_commit: bool,
}
impl ReadCursor for PostgresStaging {
fn len(&self) -> u64 {
self.staged.len() as u64
}
async fn read_at(&self, offset: u64, len: u64) -> Result<Vec<u8>, Error> {
let (start, end) = clamped(offset, len, self.staged.len());
Ok(self.staged[start..end].to_vec())
}
}
impl StagingCursor for PostgresStaging {
async fn append(&mut self, chunk: &[u8]) -> Result<(), Error> {
self.staged.extend_from_slice(chunk);
Ok(())
}
async fn commit(mut self, key: &[u8]) -> Result<(), Error> {
let value = std::mem::take(&mut self.staged);
#[cfg(test)]
if self.fail_commit {
return Err(Error::Io(
"engine commit failure (postgres, injected)".to_owned(),
));
}
let size = value.len() as i64;
let client = self.pool.get().await.map_err(err)?;
let sql = format!(
"INSERT INTO {} (key, value, size) VALUES ($1, $2, $3) ON CONFLICT DO NOTHING",
self.table
);
let stmt = client.prepare_cached(&sql).await.map_err(err)?;
client
.execute(&stmt, &[&key, &value, &size])
.await
.map_err(err)?;
Ok(())
}
async fn discard(mut self) {
self.staged.clear();
}
}
impl Backend for PostgresBackend {
type Cursor = PostgresCursor;
type Staging = PostgresStaging;
fn name(&self) -> &'static str {
"postgres"
}
async fn has(&self, key: &[u8]) -> Result<bool, Error> {
let client = self.pool.get().await.map_err(err)?;
let sql = format!("SELECT EXISTS(SELECT 1 FROM {} WHERE key = $1)", self.table);
let stmt = match client.prepare_cached(&sql).await {
Ok(stmt) => stmt,
Err(e) if is_undefined_table(&e) => return Ok(false),
Err(e) => return Err(err(e)),
};
let present = client
.query_one(&stmt, &[&key])
.await
.map_err(err)?
.try_get::<_, bool>(0)
.map_err(err)?;
Ok(present)
}
async fn get(&self, key: &[u8]) -> Result<Option<Self::Cursor>, Error> {
let client = self.pool.get().await.map_err(err)?;
let sql = format!("SELECT value, size FROM {} WHERE key = $1", self.table);
let stmt = match client.prepare_cached(&sql).await {
Ok(stmt) => stmt,
Err(e) if is_undefined_table(&e) => return Ok(None),
Err(e) => return Err(err(e)),
};
match client.query_opt(&stmt, &[&key]).await.map_err(err)? {
Some(row) => {
let value = row.try_get::<_, Vec<u8>>(0).map_err(err)?;
let len = size_u64(row.try_get::<_, i64>(1).map_err(err)?);
Ok(Some(PostgresCursor { value, len }))
}
None => Ok(None),
}
}
async fn put(&self, key: &[u8], value: &[u8]) -> Result<(), Error> {
let client = self.pool.get().await.map_err(err)?;
let sql = format!(
"INSERT INTO {} (key, value, size) VALUES ($1, $2, $3) ON CONFLICT DO NOTHING",
self.table
);
let stmt = client.prepare_cached(&sql).await.map_err(err)?;
let size = value.len() as i64;
client
.execute(&stmt, &[&key, &value, &size])
.await
.map_err(err)?;
Ok(())
}
async fn put_staged(&self) -> Result<Self::Staging, Error> {
Ok(PostgresStaging {
staged: Vec::new(),
pool: self.pool.clone(),
table: self.table.clone(),
#[cfg(test)]
fail_commit: self.fail_commit,
})
}
async fn delete(&self, key: &[u8]) -> Result<(), Error> {
let client = self.pool.get().await.map_err(err)?;
let sql = format!("DELETE FROM {} WHERE key = $1", self.table);
let stmt = match client.prepare_cached(&sql).await {
Ok(stmt) => stmt,
Err(e) if is_undefined_table(&e) => return Err(Error::Missing),
Err(e) => return Err(err(e)),
};
match client.execute(&stmt, &[&key]).await.map_err(err)? {
0 => Err(Error::Missing),
_ => Ok(()),
}
}
async fn list(&self) -> Result<Vec<Vec<u8>>, Error> {
let client = self.pool.get().await.map_err(err)?;
let sql = format!("SELECT key FROM {} ORDER BY key", self.table);
let stmt = match client.prepare_cached(&sql).await {
Ok(stmt) => stmt,
Err(e) if is_undefined_table(&e) => return Ok(Vec::new()),
Err(e) => return Err(err(e)),
};
let rows = client.query(&stmt, &[]).await.map_err(err)?;
rows.iter()
.map(|row| row.try_get::<_, Vec<u8>>(0).map_err(err))
.collect::<Result<Vec<Vec<u8>>, _>>()
.map_err(err)
}
async fn size(&self, key: &[u8]) -> Result<Option<u64>, Error> {
let client = self.pool.get().await.map_err(err)?;
let sql = format!("SELECT size FROM {} WHERE key = $1", self.table);
let stmt = match client.prepare_cached(&sql).await {
Ok(stmt) => stmt,
Err(e) if is_undefined_table(&e) => return Ok(None),
Err(e) => return Err(err(e)),
};
match client.query_opt(&stmt, &[&key]).await.map_err(err)? {
Some(row) => Ok(Some(size_u64(row.try_get::<_, i64>(0).map_err(err)?))),
None => Ok(None),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::Arc;
use parking_lot::Mutex;
use crate::postgres::testing::{harness, staged, stored_value, CONTENT};
#[tokio::test]
async fn backend_method_contract() {
let Some(backend) = harness("engine").await else {
return;
};
let key = b"\x01opaque-key-bytes".to_vec();
assert_eq!(backend.name(), "postgres", "reporting identity");
assert!(
!backend.has(&key).await.unwrap(),
"virgin store: has is a no-op read"
);
assert!(
backend.get(&key).await.unwrap().is_none(),
"virgin store: get is a no-op read"
);
assert_eq!(
backend.size(&key).await.unwrap(),
None,
"virgin store: size is a no-op read"
);
assert!(
backend.list().await.unwrap().is_empty(),
"virgin store: list is empty"
);
assert_eq!(
backend.delete(&key).await.unwrap_err(),
Error::Missing,
"virgin store: delete reports Missing"
);
backend.put(&key, CONTENT).await.unwrap();
assert!(backend.has(&key).await.unwrap(), "put lands");
assert_eq!(
backend.size(&key).await.unwrap(),
Some(CONTENT.len() as u64),
"size is metadata-only and exact"
);
backend.put(&key, b"overwritten").await.unwrap();
assert_eq!(
stored_value(&backend, &key).await.as_deref(),
Some(CONTENT),
"CAS first-writer wins: re-put is a no-op"
);
assert_eq!(
backend.list().await.unwrap(),
vec![key.clone()],
"list is complete, keys not values"
);
backend.delete(&key).await.unwrap();
assert!(!backend.has(&key).await.unwrap(), "delete lands");
assert_eq!(
backend.delete(&key).await.unwrap_err(),
Error::Missing,
"delete reports Missing when absent"
);
}
#[tokio::test]
async fn list_is_complete_and_well_formed() {
let Some(backend) = harness("list").await else {
return;
};
let keys = [
b"\x02b-key".to_vec(),
b"\x02a-key".to_vec(),
b"\x02c-key".to_vec(),
];
for key in &keys {
backend.put(key, b"value").await.unwrap();
}
let mut listed = backend.list().await.unwrap();
let mut expected = keys.to_vec();
listed.sort();
expected.sort();
assert_eq!(listed, expected, "list is complete and well-formed");
for key in &expected {
assert!(
backend.has(key).await.unwrap(),
"no spurious or mangled keys"
);
}
assert_eq!(
backend.list().await.unwrap(),
expected,
"the single-statement posture lists keys ordered"
);
}
#[tokio::test]
async fn cursor_contract_materialized_body() {
let Some(backend) = harness("cursor").await else {
return;
};
let key = b"\x01cursor-key".to_vec();
backend.put(&key, CONTENT).await.unwrap();
let cursor = backend
.get(&key)
.await
.expect("get succeeds")
.expect("entry present");
let len = CONTENT.len() as u64;
assert_eq!(cursor.len(), len, "len is the entry length");
assert_eq!(cursor.read_at(0, len).await.unwrap(), CONTENT, "full read");
assert_eq!(
cursor.read_at(0, len + 17).await.unwrap(),
CONTENT,
"over-read clamps at EOF, pread-style"
);
let mid = (CONTENT.len() / 2) as u64;
assert_eq!(
cursor.read_at(mid, len).await.unwrap(),
&CONTENT[mid as usize..],
"tail read clamps"
);
assert_eq!(
cursor.read_at(len, 5).await.unwrap(),
Vec::<u8>::new(),
"read at EOF is empty, never an error"
);
assert_eq!(
cursor.read_at(0, 0).await.unwrap(),
Vec::<u8>::new(),
"zero-length read"
);
let (slice, digest) = cursor.read_range(mid, 7).await.unwrap();
assert_eq!(slice, &CONTENT[mid as usize..(mid as usize + 7)]);
assert_eq!(digest, crate::backend::slice_digest(&slice));
}
#[tokio::test]
async fn empty_entry_round_trips() {
let Some(backend) = harness("empty").await else {
return;
};
let key = b"\x01empty-key".to_vec();
backend.put(&key, b"").await.unwrap();
assert!(backend.has(&key).await.unwrap(), "empty value still lands");
assert_eq!(backend.size(&key).await.unwrap(), Some(0), "exact zero");
let cursor = backend.get(&key).await.unwrap().expect("entry present");
assert_eq!(cursor.len(), 0);
assert_eq!(cursor.read_at(0, 8).await.unwrap(), Vec::<u8>::new());
let staged_key = b"\x01empty-staged-key".to_vec();
let mut staging = backend.put_staged().await.unwrap();
staging.append(b"").await.unwrap();
assert_eq!(staging.len(), 0, "empty append stays empty");
staging
.commit(&staged_key)
.await
.expect("empty commit publishes");
assert!(
backend.has(&staged_key).await.unwrap(),
"empty entry landed"
);
assert_eq!(backend.size(&staged_key).await.unwrap(), Some(0));
let committed = backend
.get(&staged_key)
.await
.expect("get succeeds")
.expect("entry present");
assert_eq!(
committed.read_at(0, 8).await.unwrap(),
Vec::<u8>::new(),
"committed empty body reads empty"
);
}
#[tokio::test]
async fn staged_put_round_trips_through_the_read_seam() {
let Some(backend) = harness("staged").await else {
return;
};
let key = b"\x01staged-key".to_vec();
let mut staging = backend.put_staged().await.unwrap();
staging.append(&CONTENT[..8]).await.unwrap();
staging.append(&CONTENT[8..16]).await.unwrap();
staging.append(&CONTENT[16..]).await.unwrap();
assert_eq!(
staging.len() as usize,
CONTENT.len(),
"staged length is the appended length"
);
let (bytes, digest) = staging
.read_range(0, staging.len())
.await
.expect("staged bytes readable");
assert_eq!(bytes, CONTENT, "staged artifact readable as a whole");
assert_eq!(digest, crate::backend::slice_digest(CONTENT));
staging.commit(&key).await.expect("commit publishes");
assert!(backend.has(&key).await.unwrap(), "entry published");
let cursor = backend
.get(&key)
.await
.expect("get succeeds")
.expect("entry present");
let (bytes, digest) = cursor
.read_range(0, cursor.len())
.await
.expect("committed read");
assert_eq!(bytes, CONTENT);
assert_eq!(digest, crate::backend::slice_digest(CONTENT));
backend.put(&key, b"overwritten").await.unwrap();
assert_eq!(
stored_value(&backend, &key).await.as_deref(),
Some(CONTENT),
"whole-value re-put does not overwrite"
);
staged(&backend, &key, b"overwritten-by-stage")
.await
.expect("staged re-commit is dedup, not failure");
assert_eq!(
stored_value(&backend, &key).await.as_deref(),
Some(CONTENT),
"staged re-commit does not overwrite"
);
}
#[tokio::test]
async fn staged_discard_leaves_zero_residue() {
let Some(backend) = harness("discard").await else {
return;
};
let key = b"\x01discard-key".to_vec();
let mut staging = backend.put_staged().await.unwrap();
staging.append(b"residue").await.unwrap();
staging.discard().await;
assert!(
!backend.has(&key).await.unwrap(),
"discarded stage published nothing"
);
assert_eq!(backend.size(&key).await.unwrap(), None, "no entry landed");
assert!(
backend.list().await.unwrap().is_empty(),
"the table holds exactly zero entries after discard"
);
}
#[tokio::test]
async fn commit_failure_self_discards() {
let Some(backend) = harness("fail").await else {
return;
};
let failing = PostgresBackend {
pool: backend.pool.clone(),
table: backend.table.clone(),
schema: backend.schema.clone(),
pins_table: backend.pins_table.clone(),
window_table: backend.window_table.clone(),
pin_ttl: backend.pin_ttl,
lease: Mutex::new(None),
liveness_sources: Arc::new(Mutex::new(Vec::new())),
#[cfg(test)]
fail_commit: true,
#[cfg(test)]
fail_joint_commit: false,
#[cfg(test)]
fail_window_stage: false,
};
let key = b"\x01failure-key".to_vec();
let mut staging = failing.put_staged().await.unwrap();
staging.append(CONTENT).await.unwrap();
let err = staging.commit(&key).await.unwrap_err();
assert!(matches!(err, Error::Io(_)), "engine media failure");
assert!(
!backend.has(&key).await.unwrap(),
"failed commit published nothing"
);
assert_eq!(backend.size(&key).await.unwrap(), None, "zero residue");
assert!(
backend.list().await.unwrap().is_empty(),
"zero residue in the table itself"
);
}
#[tokio::test]
async fn size_probe_is_exact_and_rides_the_pk_index() {
let Some(backend) = harness("size").await else {
return;
};
let key = b"\x01size-probe-key".to_vec();
let large_value = vec![0xAB_u8; 256 * 1024];
backend.put(&key, &large_value).await.unwrap();
assert_eq!(
backend.size(&key).await.unwrap(),
Some(large_value.len() as u64),
"size is metadata-only and exact"
);
assert_eq!(backend.size(b"\x01absent").await.unwrap(), None);
let client = backend.pool.get().await.unwrap();
let explain = format!(
"EXPLAIN SELECT value FROM {} WHERE key = '\\x{}'",
backend.table,
hex::encode(&key)
);
let rows = client.query(explain.as_str(), &[]).await.unwrap();
let plan: String = rows
.iter()
.map(|row| row.get::<_, String>(0))
.collect::<Vec<_>>()
.join(" ");
assert!(
plan.contains("Index Scan") && plan.contains("blobs_pkey"),
"the size probe searches the PK index: {plan}"
);
}
#[tokio::test]
async fn delete_batch_helper_contract() {
let Some(backend) = harness("batch").await else {
return;
};
let keys = [
b"\x01batch-a".to_vec(),
b"\x01batch-b".to_vec(),
b"\x01batch-c".to_vec(),
b"\x01batch-d".to_vec(),
b"\x01batch-e".to_vec(),
];
for key in &keys {
backend.put(key, b"value").await.unwrap();
}
let removed = backend
.delete_batch(&keys[1..4])
.await
.expect("batch delete");
assert_eq!(removed, 3, "rows-affected count returned");
assert!(backend.has(&keys[0]).await.unwrap());
assert!(!backend.has(&keys[1]).await.unwrap());
assert!(!backend.has(&keys[2]).await.unwrap());
assert!(!backend.has(&keys[3]).await.unwrap());
assert!(backend.has(&keys[4]).await.unwrap());
let absent = [b"\x01batch-absent".to_vec()];
assert_eq!(
backend.delete_batch(&absent).await.unwrap(),
0,
"absent keys are simply not deleted"
);
assert_eq!(
backend.delete_batch(&[]).await.unwrap(),
0,
"an empty request removes nothing"
);
}
}