use sqlx::{PgPool, Row};
use tower_layer::{Identity, Layer, Stack};
use tower_service::Service;
use crate::{
error::{Error, Result},
execution::{ExecutionRequest, ExecutionResponse, ExecutionService, SharedExecutionService},
queue::Queue,
task::{TaskHandle, TaskRef, validate_task_name},
types::QueueCleanup,
};
#[derive(Clone, Debug)]
pub struct Steda {
pool: PgPool,
execution: SharedExecutionService,
}
impl Steda {
pub async fn connect(database_url: &str) -> Result<Self> {
let pool = PgPool::connect(database_url).await?;
Ok(Self::from_pool(pool))
}
pub fn from_pool(pool: PgPool) -> Self {
Self::builder(pool).build()
}
pub const fn builder(pool: PgPool) -> StedaBuilder {
StedaBuilder::new(pool)
}
pub fn queue(&self, name: impl Into<String>) -> Result<Queue> {
Queue::from_parts(self.pool.clone(), name, self.execution.clone())
}
pub fn task<Input, Output>(
&self,
task: &TaskRef<Input, Output>,
) -> Result<TaskHandle<Input, Output>> {
validate_task_name(task.task_name())?;
let queue = self.queue(task.queue_name().to_owned())?;
Ok(TaskHandle::from_ref(queue, task.clone()))
}
pub async fn queues(&self) -> Result<Vec<String>> {
let rows =
sqlx::query("SELECT name FROM steda.list_queues()").fetch_all(&self.pool).await?;
Ok(rows.into_iter().map(|row| row.get("name")).collect())
}
pub async fn cleanup(&self) -> Result<Vec<QueueCleanup>> {
let rows = sqlx::query(
r#"
SELECT queue_name, tasks_deleted
FROM steda.cleanup_all_queues(NULL::text)
"#,
)
.fetch_all(&self.pool)
.await?;
rows.into_iter()
.map(|row| {
let tasks_deleted: i32 = row.get("tasks_deleted");
let tasks_deleted = u32::try_from(tasks_deleted).map_err(|_| {
Error::Other("PostgreSQL returned a negative cleanup count".to_owned())
})?;
Ok(QueueCleanup { queue_name: row.get("queue_name"), tasks_deleted })
})
.collect()
}
pub const fn pool(&self) -> &PgPool {
&self.pool
}
}
#[derive(Debug)]
pub struct StedaBuilder<L = Identity> {
pool: PgPool,
layers: L,
}
impl StedaBuilder<Identity> {
const fn new(pool: PgPool) -> Self {
Self { pool, layers: Identity::new() }
}
}
impl<L> StedaBuilder<L> {
#[must_use]
pub fn layer<T>(self, layer: T) -> StedaBuilder<Stack<L, T>> {
StedaBuilder { pool: self.pool, layers: Stack::new(self.layers, layer) }
}
#[must_use]
pub fn build(self) -> Steda
where
L: Layer<ExecutionService>,
L::Service:
Service<ExecutionRequest, Response = ExecutionResponse, Error = Error> + Send + 'static,
<L::Service as Service<ExecutionRequest>>::Future: Send + 'static,
{
let execution = SharedExecutionService::new(self.layers.layer(ExecutionService));
Steda { pool: self.pool, execution }
}
}