async insert task query (#23)

* async insert task query

* return the number of rows
This commit is contained in:
Ayrat Badykov 2022-07-17 19:22:00 +03:00 committed by GitHub
parent 5b61458b5f
commit c2c83d1031
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
6 changed files with 154 additions and 9 deletions

View file

@ -22,5 +22,7 @@ typetag = "0.2"
log = "0.4" log = "0.4"
serde = { version = "1", features = ["derive"] } serde = { version = "1", features = ["derive"] }
thiserror = "1.0" thiserror = "1.0"
bb8-postgres = "0.8" bb8-postgres = {version = "0.8", features = ["with-serde_json-1"]}
tokio = { version = "1.20", features = ["full"] } tokio = { version = "1.20", features = ["full"] }
async-trait = "0.1"
typed-builder = "0.10"

View file

@ -1,11 +1,131 @@
use crate::asynk::AsyncRunnable;
use bb8_postgres::bb8::Pool;
use bb8_postgres::bb8::RunError;
use bb8_postgres::tokio_postgres::tls::MakeTlsConnect;
use bb8_postgres::tokio_postgres::tls::TlsConnect;
use bb8_postgres::tokio_postgres::types::ToSql;
use bb8_postgres::tokio_postgres::Socket;
use bb8_postgres::PostgresConnectionManager;
use thiserror::Error;
use typed_builder::TypedBuilder;
#[derive(Error, Debug)]
pub enum AsyncQueueError {
#[error(transparent)]
PoolError(#[from] RunError<bb8_postgres::tokio_postgres::Error>),
#[error(transparent)]
PgError(#[from] bb8_postgres::tokio_postgres::Error),
#[error("returned invalid result (expected {expected:?}, found {found:?})")]
ResultError { expected: u64, found: u64 },
}
#[derive(Debug, TypedBuilder)]
pub struct AsyncQueue<Tls>
where
Tls: MakeTlsConnect<Socket> + Clone + Send + Sync + 'static,
<Tls as MakeTlsConnect<Socket>>::Stream: Send + Sync,
<Tls as MakeTlsConnect<Socket>>::TlsConnect: Send,
<<Tls as MakeTlsConnect<Socket>>::TlsConnect as TlsConnect<Socket>>::Future: Send,
{
pool: Pool<PostgresConnectionManager<Tls>>,
#[builder(default = false)]
test: bool,
}
const INSERT_TASK_QUERY: &str = include_str!("queries/insert_task.sql");
impl<Tls> AsyncQueue<Tls>
where
Tls: MakeTlsConnect<Socket> + Clone + Send + Sync + 'static,
<Tls as MakeTlsConnect<Socket>>::Stream: Send + Sync,
<Tls as MakeTlsConnect<Socket>>::TlsConnect: Send,
<<Tls as MakeTlsConnect<Socket>>::TlsConnect as TlsConnect<Socket>>::Future: Send,
{
pub fn new(pool: Pool<PostgresConnectionManager<Tls>>) -> Self {
AsyncQueue::builder().pool(pool).build()
}
pub async fn insert_task(&mut self, task: &dyn AsyncRunnable) -> Result<u64, AsyncQueueError> {
let json_task = serde_json::to_value(task).unwrap();
let task_type = task.task_type();
self.execute_one(INSERT_TASK_QUERY, &[&json_task, &task_type])
.await
}
async fn execute_one(
&mut self,
query: &str,
params: &[&(dyn ToSql + Sync)],
) -> Result<u64, AsyncQueueError> {
let mut connection = self.pool.get().await?;
let result = if self.test {
let transaction = connection.transaction().await?;
let result = transaction.execute(query, params).await?;
transaction.rollback().await?;
result
} else {
connection.execute(query, params).await?
};
if result != 1 {
return Err(AsyncQueueError::ResultError {
expected: 1,
found: result,
});
}
Ok(result)
}
}
#[cfg(test)]
mod async_queue_tests {
use super::AsyncQueue;
use crate::asynk::AsyncRunnable;
use crate::asynk::Error;
use async_trait::async_trait;
use bb8_postgres::bb8::Pool;
use bb8_postgres::tokio_postgres::Client; use bb8_postgres::tokio_postgres::Client;
use bb8_postgres::tokio_postgres::NoTls;
use bb8_postgres::PostgresConnectionManager;
use serde::{Deserialize, Serialize};
pub struct AsyncQueue { #[derive(Serialize, Deserialize)]
pg_client: Client, struct Job {
pub number: u16,
} }
impl AsyncQueue { #[typetag::serde]
pub fn new(pg_client: Client) -> Self { #[async_trait]
Self { pg_client } impl AsyncRunnable for Job {
async fn run(&self, _connection: &Client) -> Result<(), Error> {
Ok(())
}
}
#[tokio::test]
async fn insert_task_creates_new_task() {
let mut queue = queue().await;
let result = queue.insert_task(&Job { number: 1 }).await.unwrap();
assert_eq!(1, result);
}
async fn queue() -> AsyncQueue<NoTls> {
let pg_mgr = PostgresConnectionManager::new_from_stringlike(
"postgres://postgres:postgres@localhost/fang",
NoTls,
)
.unwrap();
let pool = Pool::builder().build(pg_mgr).await.unwrap();
AsyncQueue::builder().pool(pool).test(true).build()
} }
} }

View file

@ -0,0 +1,19 @@
use async_trait::async_trait;
use bb8_postgres::tokio_postgres::Client;
const COMMON_TYPE: &str = "common";
#[derive(Debug)]
pub struct Error {
pub description: String,
}
#[typetag::serde(tag = "type")]
#[async_trait]
pub trait AsyncRunnable {
async fn run(&self, client: &Client) -> Result<(), Error>;
fn task_type(&self) -> String {
COMMON_TYPE.to_string()
}
}

View file

@ -1 +1,5 @@
pub mod async_queue; pub mod async_queue;
pub mod async_runnable;
pub use async_runnable::AsyncRunnable;
pub use async_runnable::Error;

View file

@ -1 +1 @@
INSERT INTO "fang_tasks" ("metadata", "created_at") VALUES ($1, $2) INSERT INTO "fang_tasks" ("metadata", "task_type") VALUES ($1, $2)

View file

@ -1 +1 @@
INSERT INTO "fang_tasks" ("metadata", "error_message", "state", "task_type") VALUES ($1, $2, $3, $4) INSERT INTO "fang_tasks" ("metadata", "task_type") VALUES ($1, $2)