use std::error::Error;
use std::fmt;
use std::sync::OnceLock;
use serde_json::Value;
use sqlx::postgres::PgRow;
use sqlx::{Postgres, Row};
use crate::{
FailureCategory, ReaderError, StateCodecError, StateSchemaId, StateSchemaVersion,
VersionedStateCodec,
};
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct KeysetColumn {
name: &'static str,
kind: KeysetColumnKind,
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
#[non_exhaustive]
pub enum KeysetColumnKind {
Text,
I64,
}
impl KeysetColumn {
#[must_use]
pub const fn text(name: &'static str) -> Self {
Self {
name,
kind: KeysetColumnKind::Text,
}
}
#[must_use]
pub const fn i64(name: &'static str) -> Self {
Self {
name,
kind: KeysetColumnKind::I64,
}
}
#[must_use]
pub const fn name(&self) -> &'static str {
self.name
}
#[must_use]
pub const fn kind(&self) -> KeysetColumnKind {
self.kind
}
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
#[non_exhaustive]
pub enum PostgresComponentConfigError {
EmptyKeyColumns,
InvalidColumnName,
InvalidFetchSize,
InvalidMaxParameters,
InvalidColumnsPerRow,
MaxParametersExceedsProtocolLimit,
}
impl fmt::Display for PostgresComponentConfigError {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str(match self {
Self::EmptyKeyColumns => "at least one ordering-key column is required",
Self::InvalidColumnName => "an ordering-key or bound column name is invalid",
Self::InvalidFetchSize => "the configured fetch/page size must be nonzero",
Self::InvalidMaxParameters => {
"the configured maximum parameters per statement is too small"
}
Self::InvalidColumnsPerRow => "the configured column count per row must be nonzero",
Self::MaxParametersExceedsProtocolLimit => {
"the configured maximum parameters per statement exceeds PostgreSQL's \
65,535-parameter protocol limit"
}
})
}
}
impl Error for PostgresComponentConfigError {}
pub(crate) fn validate_identifier(name: &str) -> Result<(), PostgresComponentConfigError> {
let mut bytes = name.bytes();
let Some(first) = bytes.next() else {
return Err(PostgresComponentConfigError::InvalidColumnName);
};
if !(first.is_ascii_alphabetic() || first == b'_') {
return Err(PostgresComponentConfigError::InvalidColumnName);
}
if !bytes.all(|byte| byte.is_ascii_alphanumeric() || byte == b'_') {
return Err(PostgresComponentConfigError::InvalidColumnName);
}
Ok(())
}
pub(crate) fn validate_key_columns(
columns: &[KeysetColumn],
) -> Result<(), PostgresComponentConfigError> {
if columns.is_empty() {
return Err(PostgresComponentConfigError::EmptyKeyColumns);
}
for column in columns {
validate_identifier(column.name())?;
}
Ok(())
}
pub(crate) fn order_by_sql(columns: &[KeysetColumn]) -> String {
columns
.iter()
.map(KeysetColumn::name)
.collect::<Vec<_>>()
.join(", ")
}
pub(crate) fn keyset_predicate_sql(columns: &[KeysetColumn], next_param: usize) -> String {
let names = columns
.iter()
.map(KeysetColumn::name)
.collect::<Vec<_>>()
.join(", ");
let placeholders = (0..columns.len())
.map(|offset| format!("${}", next_param + offset))
.collect::<Vec<_>>()
.join(", ");
format!("({names}) > ({placeholders})")
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub(crate) enum PostgresKeyValue {
Text(String),
I64(i64),
}
#[derive(Clone, Debug, Default, Eq, PartialEq)]
pub(crate) struct KeysetPosition(Vec<PostgresKeyValue>);
impl KeysetPosition {
pub(crate) const fn none() -> Self {
Self(Vec::new())
}
pub(crate) fn from_values(values: Vec<PostgresKeyValue>) -> Self {
Self(values)
}
pub(crate) fn values(&self) -> &[PostgresKeyValue] {
&self.0
}
pub(crate) fn is_empty(&self) -> bool {
self.0.is_empty()
}
}
pub(crate) fn bind_keyset<'q>(
mut query: sqlx::query::Query<'q, Postgres, sqlx::postgres::PgArguments>,
values: &'q [PostgresKeyValue],
) -> sqlx::query::Query<'q, Postgres, sqlx::postgres::PgArguments> {
for value in values {
query = match value {
PostgresKeyValue::Text(text) => query.bind(text.as_str()),
PostgresKeyValue::I64(value) => query.bind(*value),
};
}
query
}
pub(crate) fn extract_keyset(
row: &PgRow,
columns: &[KeysetColumn],
) -> Result<KeysetPosition, ReaderError> {
let values = columns
.iter()
.map(|column| {
let result = match column.kind() {
KeysetColumnKind::Text => row
.try_get::<String, _>(column.name())
.map(PostgresKeyValue::Text),
KeysetColumnKind::I64 => row
.try_get::<i64, _>(column.name())
.or_else(|_| row.try_get::<i32, _>(column.name()).map(i64::from))
.map(PostgresKeyValue::I64),
};
result.map_err(|_| ReaderError::with_category(FailureCategory::UserComponent))
})
.collect::<Result<Vec<_>, _>>()?;
Ok(KeysetPosition::from_values(values))
}
pub(crate) fn classify_pg_error(error: &sqlx::Error) -> FailureCategory {
let Some(database) = error.as_database_error() else {
return FailureCategory::TransientInfrastructure;
};
let Some(code) = database.code() else {
return FailureCategory::TransientInfrastructure;
};
if code == "57014" {
return FailureCategory::Cancelled;
}
if code.starts_with("22") || code.starts_with("23") || code.starts_with("42") {
return FailureCategory::UserComponent;
}
FailureCategory::TransientInfrastructure
}
pub struct PostgresRow<'a>(&'a PgRow);
impl<'a> PostgresRow<'a> {
pub(crate) const fn new(row: &'a PgRow) -> Self {
Self(row)
}
pub fn text(&self, column: &str) -> Result<String, ReaderError> {
self.0
.try_get(column)
.map_err(|_| ReaderError::with_category(FailureCategory::UserComponent))
}
pub fn i64(&self, column: &str) -> Result<i64, ReaderError> {
self.0
.try_get(column)
.map_err(|_| ReaderError::with_category(FailureCategory::UserComponent))
}
pub fn i32(&self, column: &str) -> Result<i32, ReaderError> {
self.0
.try_get(column)
.map_err(|_| ReaderError::with_category(FailureCategory::UserComponent))
}
pub fn bool(&self, column: &str) -> Result<bool, ReaderError> {
self.0
.try_get(column)
.map_err(|_| ReaderError::with_category(FailureCategory::UserComponent))
}
pub fn f64(&self, column: &str) -> Result<f64, ReaderError> {
self.0
.try_get::<f64, _>(column)
.or_else(|_| self.0.try_get::<f32, _>(column).map(f64::from))
.map_err(|_| ReaderError::with_category(FailureCategory::UserComponent))
}
}
fn encode_keyset(value: &KeysetPosition) -> Result<Vec<u8>, StateCodecError> {
let items: Vec<Value> = value
.values()
.iter()
.map(|value| match value {
PostgresKeyValue::Text(text) => serde_json::json!({"t": "text", "v": text}),
PostgresKeyValue::I64(value) => serde_json::json!({"t": "i64", "v": value}),
})
.collect();
serde_json::to_vec(&serde_json::json!({ "keys": items }))
.map_err(|_| StateCodecError::InvalidPayload)
}
fn decode_keyset(payload: &[u8]) -> Result<KeysetPosition, StateCodecError> {
let value: Value =
serde_json::from_slice(payload).map_err(|_| StateCodecError::InvalidPayload)?;
let array = value
.get("keys")
.and_then(Value::as_array)
.ok_or(StateCodecError::InvalidPayload)?;
let values = array
.iter()
.map(|item| {
let tag = item
.get("t")
.and_then(Value::as_str)
.ok_or(StateCodecError::InvalidPayload)?;
match tag {
"text" => item
.get("v")
.and_then(Value::as_str)
.map(|text| PostgresKeyValue::Text(text.to_owned()))
.ok_or(StateCodecError::InvalidPayload),
"i64" => item
.get("v")
.and_then(Value::as_i64)
.map(PostgresKeyValue::I64)
.ok_or(StateCodecError::InvalidPayload),
_ => Err(StateCodecError::InvalidPayload),
}
})
.collect::<Result<Vec<_>, _>>()?;
Ok(KeysetPosition::from_values(values))
}
#[derive(Clone, Copy)]
pub(crate) struct CursorKeysetSchema;
impl VersionedStateCodec<KeysetPosition> for CursorKeysetSchema {
fn schema_id(&self) -> &StateSchemaId {
static SCHEMA: OnceLock<StateSchemaId> = OnceLock::new();
#[allow(
clippy::unwrap_used,
reason = "fixed literal schema identity cannot fail validation"
)]
SCHEMA.get_or_init(|| {
StateSchemaId::new("oxide-batch.postgres-cursor-reader-position").unwrap()
})
}
fn current_version(&self) -> StateSchemaVersion {
#[allow(
clippy::unwrap_used,
reason = "fixed literal schema version cannot fail validation"
)]
StateSchemaVersion::new(1).unwrap()
}
fn encode(&self, value: &KeysetPosition) -> Result<Vec<u8>, StateCodecError> {
encode_keyset(value)
}
fn decode(&self, payload: &[u8]) -> Result<KeysetPosition, StateCodecError> {
decode_keyset(payload)
}
}
#[derive(Clone, Copy)]
pub(crate) struct PagingKeysetSchema;
impl VersionedStateCodec<KeysetPosition> for PagingKeysetSchema {
fn schema_id(&self) -> &StateSchemaId {
static SCHEMA: OnceLock<StateSchemaId> = OnceLock::new();
#[allow(
clippy::unwrap_used,
reason = "fixed literal schema identity cannot fail validation"
)]
SCHEMA.get_or_init(|| {
StateSchemaId::new("oxide-batch.postgres-paging-reader-position").unwrap()
})
}
fn current_version(&self) -> StateSchemaVersion {
#[allow(
clippy::unwrap_used,
reason = "fixed literal schema version cannot fail validation"
)]
StateSchemaVersion::new(1).unwrap()
}
fn encode(&self, value: &KeysetPosition) -> Result<Vec<u8>, StateCodecError> {
encode_keyset(value)
}
fn decode(&self, payload: &[u8]) -> Result<KeysetPosition, StateCodecError> {
decode_keyset(payload)
}
}
#[cfg(test)]
mod tests {
#![allow(clippy::unwrap_used, clippy::expect_used)]
use super::*;
#[test]
fn keyset_predicate_sql_builds_a_composite_tuple_comparison() {
let columns = [KeysetColumn::text("email"), KeysetColumn::i64("id")];
assert_eq!(keyset_predicate_sql(&columns, 1), "(email, id) > ($1, $2)");
assert_eq!(order_by_sql(&columns), "email, id");
}
#[test]
fn keyset_position_round_trips_through_json() {
let position = KeysetPosition::from_values(vec![
PostgresKeyValue::Text("a@example.com".to_owned()),
PostgresKeyValue::I64(42),
]);
let encoded = encode_keyset(&position).unwrap();
let decoded = decode_keyset(&encoded).unwrap();
assert_eq!(decoded, position);
}
#[test]
fn none_position_is_empty_and_round_trips() {
let position = KeysetPosition::none();
assert!(position.is_empty());
let encoded = encode_keyset(&position).unwrap();
let decoded = decode_keyset(&encoded).unwrap();
assert!(decoded.is_empty());
}
#[test]
fn validate_identifier_rejects_unsafe_names() {
assert!(validate_identifier("valid_name_1").is_ok());
assert!(validate_identifier("").is_err());
assert!(validate_identifier("1leading_digit").is_err());
assert!(validate_identifier("has space").is_err());
assert!(validate_identifier("has;semicolon").is_err());
assert!(validate_identifier("has'quote").is_err());
}
#[test]
fn validate_key_columns_requires_at_least_one() {
assert_eq!(
validate_key_columns(&[]),
Err(PostgresComponentConfigError::EmptyKeyColumns)
);
assert!(validate_key_columns(&[KeysetColumn::i64("id")]).is_ok());
}
#[test]
fn classify_pg_error_matches_documented_sqlstate_classes() {
fn classify_code(code: &str) -> FailureCategory {
if code == "57014" {
return FailureCategory::Cancelled;
}
if code.starts_with("22") || code.starts_with("23") || code.starts_with("42") {
return FailureCategory::UserComponent;
}
FailureCategory::TransientInfrastructure
}
assert_eq!(classify_code("57014"), FailureCategory::Cancelled);
assert_eq!(classify_code("23505"), FailureCategory::UserComponent);
assert_eq!(classify_code("22001"), FailureCategory::UserComponent);
assert_eq!(classify_code("42601"), FailureCategory::UserComponent);
assert_eq!(
classify_code("08006"),
FailureCategory::TransientInfrastructure
);
}
}