use crate::dbnexus_entities::{KeyValueActiveModel, KeyValueEntity, KeyValueModel};
use crate::error::StorageError;
use crate::storage::Storage;
use async_trait::async_trait;
use chrono::Utc;
use dbnexus::{DbPool, Session};
use sea_orm::entity::prelude::*;
use std::sync::Arc;
pub struct DBNexusStorageAdapter {
pool: Arc<DbPool>,
}
impl DBNexusStorageAdapter {
pub fn new(pool: Arc<DbPool>) -> Self {
Self { pool }
}
async fn get_session(&self) -> Result<Session, StorageError> {
self.pool
.get_session("admin")
.await
.map_err(|e| StorageError::ConnectionError(e.to_string()))
}
fn get_conn(session: &Session) -> Result<&DatabaseConnection, StorageError> {
session
.connection()
.map_err(|e| StorageError::ConnectionError(e.to_string()))
}
}
#[async_trait]
impl Storage for DBNexusStorageAdapter {
async fn get(&self, key: &str) -> Result<Option<String>, StorageError> {
let session = self.get_session().await?;
let conn = Self::get_conn(&session)?;
match KeyValueEntity::find_by_id(key.to_string()).one(conn).await {
Ok(Some(model)) => {
if let Some(expires_at) = model.expires_at {
if expires_at < Utc::now() {
let _ = KeyValueEntity::delete_by_id(key.to_string())
.exec(conn)
.await;
return Ok(None);
}
}
Ok(Some(model.value))
}
Ok(None) => Ok(None),
Err(e) => Err(StorageError::QueryError(e.to_string())),
}
}
async fn set(&self, key: &str, value: &str, ttl: Option<u64>) -> Result<(), StorageError> {
let session = self.get_session().await?;
let conn = Self::get_conn(&session)?;
let now = Utc::now();
let expires_at = ttl.map(|seconds| now + chrono::Duration::seconds(seconds as i64));
let exists = KeyValueEntity::find_by_id(key.to_string())
.one(conn)
.await
.is_ok();
let model = KeyValueModel {
key: key.to_string(),
value: value.to_string(),
expires_at,
created_at: now,
updated_at: now,
};
if exists {
let active_model: KeyValueActiveModel = model.into();
active_model.save(conn).await.map_err(|e| {
StorageError::QueryError(format!("Failed to update key-value: {}", e))
})?;
} else {
let active_model: KeyValueActiveModel = model.into();
active_model.insert(conn).await.map_err(|e| {
StorageError::QueryError(format!("Failed to insert key-value: {}", e))
})?;
}
Ok(())
}
async fn delete(&self, key: &str) -> Result<(), StorageError> {
let session = self.get_session().await?;
let conn = Self::get_conn(&session)?;
KeyValueEntity::delete_by_id(key.to_string())
.exec(conn)
.await
.map_err(|e| StorageError::QueryError(format!("Failed to delete key: {}", e)))?;
Ok(())
}
}