use std::{fmt, future::Future, sync::Arc};
use futures::future::Either;
use crate::state_stores::{self, config::StateStoreConfig};
use crate::state::record::{
StateKey, StateListRequest, StateListResponse, StateRecord, StateRecordSnapshot,
};
pub trait StateStoreBackend: Send + Sync + 'static {
type Transaction: StateTransactionBackend;
fn load(
&self,
key: StateKey,
) -> impl Future<Output = Result<Option<StateRecordSnapshot>, StateStoreError>> + Send + '_;
fn list(
&self,
request: StateListRequest,
) -> impl Future<Output = Result<StateListResponse, StateStoreError>> + Send + '_;
fn begin_transaction(
&self,
) -> impl Future<Output = Result<Self::Transaction, StateStoreError>> + Send + '_;
}
pub trait StateTransactionBackend: Send + 'static {
fn load(
&mut self,
key: StateKey,
) -> impl Future<Output = Result<Option<StateRecordSnapshot>, StateStoreError>> + Send + '_;
fn put(
&mut self,
record: StateRecord,
) -> impl Future<Output = Result<(), StateStoreError>> + Send + '_;
fn delete(
&mut self,
key: StateKey,
) -> impl Future<Output = Result<(), StateStoreError>> + Send + '_;
fn commit(self) -> impl Future<Output = Result<(), StateStoreError>> + Send
where
Self: Sized;
fn rollback(self) -> impl Future<Output = Result<(), StateStoreError>> + Send
where
Self: Sized;
}
#[derive(Clone)]
pub struct StateStore {
inner: Arc<StateStoreInner>,
}
enum StateStoreInner {
Lance(state_stores::lance::LanceStateStore),
#[cfg(feature = "redb")]
Redb(state_stores::redb::RedbStateStore),
#[cfg(feature = "postgres")]
Postgres(state_stores::postgres::PostgresStateStore),
#[cfg(feature = "sqlite")]
Sqlite(state_stores::sqlite::SqliteStateStore),
}
impl StateStore {
pub async fn connect(config: StateStoreConfig) -> Result<Self, StateStoreError> {
config.validate()?;
match config {
StateStoreConfig::Lance(config) => {
state_stores::lance::LanceStateStore::connect(config)
.await
.map(Self::from_lance)
}
StateStoreConfig::Redb(config) => {
if cfg!(feature = "redb") {
#[cfg(feature = "redb")]
return state_stores::redb::RedbStateStore::connect(config)
.await
.map(Self::from_redb);
}
drop(config);
Err(StateStoreError::InvalidConfig(
"Redb requires the `redb` feature".into(),
))
}
StateStoreConfig::Postgres(config) => {
if cfg!(feature = "postgres") {
#[cfg(feature = "postgres")]
return state_stores::postgres::PostgresStateStore::connect(config)
.await
.map(Self::from_postgres);
}
drop(config);
Err(StateStoreError::InvalidConfig(
"PostgreSQL requires the `postgres` feature".into(),
))
}
#[cfg(feature = "sqlite")]
StateStoreConfig::Sqlite(config) => {
state_stores::sqlite::SqliteStateStore::connect(config)
.await
.map(Self::from_sqlite)
}
}
}
#[cfg(feature = "redb")]
fn from_redb(store: state_stores::redb::RedbStateStore) -> Self {
Self {
inner: Arc::new(StateStoreInner::Redb(store)),
}
}
fn from_lance(store: state_stores::lance::LanceStateStore) -> Self {
Self {
inner: Arc::new(StateStoreInner::Lance(store)),
}
}
#[cfg(feature = "postgres")]
fn from_postgres(store: state_stores::postgres::PostgresStateStore) -> Self {
Self {
inner: Arc::new(StateStoreInner::Postgres(store)),
}
}
#[cfg(feature = "sqlite")]
fn from_sqlite(store: state_stores::sqlite::SqliteStateStore) -> Self {
Self {
inner: Arc::new(StateStoreInner::Sqlite(store)),
}
}
}
impl fmt::Debug for StateStore {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str("StateStore { .. }")
}
}
impl StateStoreBackend for StateStore {
type Transaction = StateTransaction;
async fn load(&self, key: StateKey) -> Result<Option<StateRecordSnapshot>, StateStoreError> {
self.inner.load(key).await
}
async fn list(&self, request: StateListRequest) -> Result<StateListResponse, StateStoreError> {
self.inner.list(request).await
}
async fn begin_transaction(&self) -> Result<Self::Transaction, StateStoreError> {
self.inner.begin_transaction().await
}
}
impl StateStoreBackend for StateStoreInner {
type Transaction = StateTransaction;
async fn load(&self, key: StateKey) -> Result<Option<StateRecordSnapshot>, StateStoreError> {
match self {
Self::Lance(store) => store.load(key).await,
#[cfg(feature = "redb")]
Self::Redb(store) => store.load(key).await,
#[cfg(feature = "postgres")]
Self::Postgres(store) => store.load(key).await,
#[cfg(feature = "sqlite")]
Self::Sqlite(store) => store.load(key).await,
}
}
async fn list(&self, request: StateListRequest) -> Result<StateListResponse, StateStoreError> {
match self {
Self::Lance(store) => store.list(request).await,
#[cfg(feature = "redb")]
Self::Redb(store) => store.list(request).await,
#[cfg(feature = "postgres")]
Self::Postgres(store) => store.list(request).await,
#[cfg(feature = "sqlite")]
Self::Sqlite(store) => store.list(request).await,
}
}
async fn begin_transaction(&self) -> Result<Self::Transaction, StateStoreError> {
match self {
Self::Lance(store) => store
.begin_transaction()
.await
.map(StateTransaction::from_lance),
#[cfg(feature = "redb")]
Self::Redb(store) => store
.begin_transaction()
.await
.map(StateTransaction::from_redb),
#[cfg(feature = "postgres")]
Self::Postgres(store) => store
.begin_transaction()
.await
.map(StateTransaction::from_postgres),
#[cfg(feature = "sqlite")]
Self::Sqlite(store) => store
.begin_transaction()
.await
.map(StateTransaction::from_sqlite),
}
}
}
pub struct StateTransaction {
inner: StateTransactionInner,
}
enum StateTransactionInner {
Lance(state_stores::lance::LanceStateTransaction),
#[cfg(feature = "redb")]
Redb(state_stores::redb::RedbStateTransaction),
#[cfg(feature = "postgres")]
Postgres(state_stores::postgres::PostgresStateTransaction),
#[cfg(feature = "sqlite")]
Sqlite(state_stores::sqlite::SqliteStateTransaction),
}
impl StateTransaction {
fn from_lance(transaction: state_stores::lance::LanceStateTransaction) -> Self {
Self {
inner: StateTransactionInner::Lance(transaction),
}
}
#[cfg(feature = "redb")]
fn from_redb(transaction: state_stores::redb::RedbStateTransaction) -> Self {
Self {
inner: StateTransactionInner::Redb(transaction),
}
}
#[cfg(feature = "postgres")]
fn from_postgres(transaction: state_stores::postgres::PostgresStateTransaction) -> Self {
Self {
inner: StateTransactionInner::Postgres(transaction),
}
}
#[cfg(feature = "sqlite")]
fn from_sqlite(transaction: state_stores::sqlite::SqliteStateTransaction) -> Self {
Self {
inner: StateTransactionInner::Sqlite(transaction),
}
}
}
impl fmt::Debug for StateTransaction {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str("StateTransaction { .. }")
}
}
impl StateTransactionBackend for StateTransaction {
async fn load(
&mut self,
key: StateKey,
) -> Result<Option<StateRecordSnapshot>, StateStoreError> {
self.inner.load(key).await
}
fn put(
&mut self,
record: StateRecord,
) -> impl Future<Output = Result<(), StateStoreError>> + Send {
self.inner.put(record)
}
async fn delete(&mut self, key: StateKey) -> Result<(), StateStoreError> {
self.inner.delete(key).await
}
async fn commit(self) -> Result<(), StateStoreError> {
self.inner.commit().await
}
async fn rollback(self) -> Result<(), StateStoreError> {
self.inner.rollback().await
}
}
impl StateTransactionBackend for StateTransactionInner {
async fn load(
&mut self,
key: StateKey,
) -> Result<Option<StateRecordSnapshot>, StateStoreError> {
match self {
Self::Lance(transaction) => transaction.load(key).await,
#[cfg(feature = "redb")]
Self::Redb(transaction) => transaction.load(key).await,
#[cfg(feature = "postgres")]
Self::Postgres(transaction) => transaction.load(key).await,
#[cfg(feature = "sqlite")]
Self::Sqlite(transaction) => transaction.load(key).await,
}
}
fn put(
&mut self,
record: StateRecord,
) -> impl Future<Output = Result<(), StateStoreError>> + Send {
let prepared: Either<_, (&mut Self, StateRecord)> = match self {
Self::Lance(transaction) => Either::Left(transaction.put(record)),
#[cfg(feature = "redb")]
transaction @ Self::Redb(_) => Either::Right((transaction, record)),
#[cfg(feature = "postgres")]
transaction @ Self::Postgres(_) => Either::Right((transaction, record)),
#[cfg(feature = "sqlite")]
transaction @ Self::Sqlite(_) => Either::Right((transaction, record)),
};
async move {
match prepared {
Either::Left(operation) => operation.await,
Either::Right((transaction, record)) => match transaction {
Self::Lance(transaction) => transaction.put(record).await,
#[cfg(feature = "redb")]
Self::Redb(transaction) => transaction.put(record).await,
#[cfg(feature = "postgres")]
Self::Postgres(transaction) => transaction.put(record).await,
#[cfg(feature = "sqlite")]
Self::Sqlite(transaction) => transaction.put(record).await,
},
}
}
}
async fn delete(&mut self, key: StateKey) -> Result<(), StateStoreError> {
match self {
Self::Lance(transaction) => transaction.delete(key).await,
#[cfg(feature = "redb")]
Self::Redb(transaction) => transaction.delete(key).await,
#[cfg(feature = "postgres")]
Self::Postgres(transaction) => transaction.delete(key).await,
#[cfg(feature = "sqlite")]
Self::Sqlite(transaction) => transaction.delete(key).await,
}
}
async fn commit(self) -> Result<(), StateStoreError> {
match self {
Self::Lance(transaction) => transaction.commit().await,
#[cfg(feature = "redb")]
Self::Redb(transaction) => transaction.commit().await,
#[cfg(feature = "postgres")]
Self::Postgres(transaction) => transaction.commit().await,
#[cfg(feature = "sqlite")]
Self::Sqlite(transaction) => transaction.commit().await,
}
}
async fn rollback(self) -> Result<(), StateStoreError> {
match self {
Self::Lance(transaction) => transaction.rollback().await,
#[cfg(feature = "redb")]
Self::Redb(transaction) => transaction.rollback().await,
#[cfg(feature = "postgres")]
Self::Postgres(transaction) => transaction.rollback().await,
#[cfg(feature = "sqlite")]
Self::Sqlite(transaction) => transaction.rollback().await,
}
}
}
#[derive(Clone, Debug, Eq, PartialEq, thiserror::Error)]
pub enum StateStoreError {
#[error("invalid state store configuration: {0}")]
InvalidConfig(String),
#[error("invalid state store request: {0}")]
InvalidRequest(String),
#[error("state store is unavailable during `{operation}`: {message}")]
Unavailable {
operation: String,
message: String,
},
#[error("state store operation `{operation}` failed: {message}")]
BackendOperation {
operation: String,
message: String,
},
#[error("state store transaction `{operation}` conflicted: {message}")]
TransactionConflict {
operation: String,
message: String,
},
#[error("state store operation `{operation}` timed out after {timeout_ms} ms")]
TimedOut {
operation: String,
timeout_ms: u64,
},
#[error("state store serialization failed: {0}")]
Serialization(String),
}
impl StateStoreError {
pub fn invalid_config(message: impl Into<String>) -> Self {
Self::InvalidConfig(message.into())
}
pub fn invalid_request(message: impl Into<String>) -> Self {
Self::InvalidRequest(message.into())
}
pub fn unavailable(operation: impl Into<String>, message: impl Into<String>) -> Self {
Self::Unavailable {
operation: operation.into(),
message: message.into(),
}
}
pub fn backend_operation(operation: impl Into<String>, message: impl Into<String>) -> Self {
Self::BackendOperation {
operation: operation.into(),
message: message.into(),
}
}
pub fn transaction_conflict(operation: impl Into<String>, message: impl Into<String>) -> Self {
Self::TransactionConflict {
operation: operation.into(),
message: message.into(),
}
}
pub fn timed_out(operation: impl Into<String>, timeout_ms: u64) -> Self {
Self::TimedOut {
operation: operation.into(),
timeout_ms,
}
}
pub fn serialization(message: impl Into<String>) -> Self {
Self::Serialization(message.into())
}
}
impl From<serde_json::Error> for StateStoreError {
fn from(error: serde_json::Error) -> Self {
Self::serialization(error.to_string())
}
}