fetch_task static and AsyncQueue method (#32)

* fetch and touch a task inside of db transaction

* remove transaction from AsyncQueue

* fix `fetch_and_touch` test

* uncomment one more test
This commit is contained in:
Ayrat Badykov 2022-07-19 23:05:09 +03:00 committed by GitHub
parent abe324fc57
commit 3602097fb6
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23

View file

@ -14,6 +14,16 @@ use thiserror::Error;
use typed_builder::TypedBuilder;
use uuid::Uuid;
const INSERT_TASK_QUERY: &str = include_str!("queries/insert_task.sql");
const UPDATE_TASK_STATE_QUERY: &str = include_str!("queries/update_task_state.sql");
const FAIL_TASK_QUERY: &str = include_str!("queries/fail_task.sql");
const REMOVE_ALL_TASK_QUERY: &str = include_str!("queries/remove_all_tasks.sql");
const REMOVE_TASK_QUERY: &str = include_str!("queries/remove_task.sql");
const REMOVE_TASKS_TYPE_QUERY: &str = include_str!("queries/remove_tasks_type.sql");
const FETCH_TASK_TYPE_QUERY: &str = include_str!("queries/fetch_task_type.sql");
const DEFAULT_TASK_TYPE: &str = "common";
#[derive(Debug, Eq, PartialEq, Clone, ToSql, FromSql)]
#[postgres(name = "fang_task_state")]
pub enum FangTaskState {
@ -88,35 +98,19 @@ pub enum AsyncQueueError {
PgError(#[from] bb8_postgres::tokio_postgres::Error),
#[error("returned invalid result (expected {expected:?}, found {found:?})")]
ResultError { expected: u64, found: u64 },
#[error("Queue doesn't have a connection")]
PoolAndTransactionEmpty,
#[error("Need to create a transaction to perform this operation")]
TransactionEmpty,
}
#[derive(TypedBuilder)]
pub struct AsyncQueue<'a, Tls>
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,
{
#[builder(default, setter(into))]
pool: Option<Pool<PostgresConnectionManager<Tls>>>,
#[builder(default, setter(into))]
transaction: Option<Transaction<'a>>,
pool: Pool<PostgresConnectionManager<Tls>>,
}
const INSERT_TASK_QUERY: &str = include_str!("queries/insert_task.sql");
const UPDATE_TASK_STATE_QUERY: &str = include_str!("queries/update_task_state.sql");
const FAIL_TASK_QUERY: &str = include_str!("queries/fail_task.sql");
const REMOVE_ALL_TASK_QUERY: &str = include_str!("queries/remove_all_tasks.sql");
const REMOVE_TASK_QUERY: &str = include_str!("queries/remove_task.sql");
const REMOVE_TASKS_TYPE_QUERY: &str = include_str!("queries/remove_tasks_type.sql");
const FETCH_TASK_TYPE_QUERY: &str = include_str!("queries/fetch_task_type.sql");
impl<'a, Tls> AsyncQueue<'a, Tls>
impl<Tls> AsyncQueue<Tls>
where
Tls: MakeTlsConnect<Socket> + Clone + Send + Sync + 'static,
<Tls as MakeTlsConnect<Socket>>::Stream: Send + Sync,
@ -124,116 +118,55 @@ where
<<Tls as MakeTlsConnect<Socket>>::TlsConnect as TlsConnect<Socket>>::Future: Send,
{
pub fn new(pool: Pool<PostgresConnectionManager<Tls>>) -> Self {
AsyncQueue::builder().pool(pool).build()
AsyncQueue { pool }
}
pub fn new_with_transaction(transaction: Transaction<'a>) -> Self {
AsyncQueue::builder().transaction(transaction).build()
}
pub async fn rollback(mut self) -> Result<AsyncQueue<'a, Tls>, AsyncQueueError> {
let transaction = self.transaction;
self.transaction = None;
match transaction {
Some(tr) => {
tr.rollback().await?;
Ok(self)
}
None => Err(AsyncQueueError::TransactionEmpty),
}
}
pub async fn commit(mut self) -> Result<AsyncQueue<'a, Tls>, AsyncQueueError> {
let transaction = self.transaction;
self.transaction = None;
match transaction {
Some(tr) => {
tr.commit().await?;
Ok(self)
}
None => Err(AsyncQueueError::TransactionEmpty),
}
}
pub async fn fetch_task(
pub async fn fetch_and_touch_task(
&mut self,
task_type: &Option<String>,
) -> Result<Task, AsyncQueueError> {
let mut task = match task_type {
None => self.get_task_type("common").await?,
Some(task_type_str) => self.get_task_type(task_type_str).await?,
};
self.update_task_state(&task, FangTaskState::InProgress)
.await?;
task.state = FangTaskState::InProgress;
Ok(task)
}
pub async fn get_task_type(&mut self, task_type: &str) -> Result<Task, AsyncQueueError> {
let row: Row = self.get_row(FETCH_TASK_TYPE_QUERY, &[&task_type]).await?;
let id: Uuid = row.get("id");
let metadata: serde_json::Value = row.get("metadata");
let error_message: Option<String> = match row.try_get("error_message") {
Ok(error_message) => Some(error_message),
Err(_) => None,
};
let state: FangTaskState = FangTaskState::New;
let task_type: String = row.get("task_type");
let created_at: DateTime<Utc> = row.get("created_at");
let updated_at: DateTime<Utc> = row.get("updated_at");
let task = Task::builder()
.id(id)
.metadata(metadata)
.error_message(error_message)
.state(state)
.task_type(task_type)
.created_at(created_at)
.updated_at(updated_at)
.build();
let mut connection = self.pool.get().await?;
let mut transaction = connection.transaction().await?;
let task = Self::fetch_and_touch_task_query(&mut transaction, task_type).await?;
transaction.commit().await?;
Ok(task)
}
pub async fn get_row(
&mut self,
query: &str,
params: &[&(dyn ToSql + Sync)],
) -> Result<Row, AsyncQueueError> {
let row: Row = if let Some(pool) = &self.pool {
let connection = pool.get().await?;
let connection = self.pool.get().await?;
let row = connection.query_one(query, params).await?;
connection.query_one(query, params).await?
} else if let Some(transaction) = &self.transaction {
transaction.query_one(query, params).await?
} else {
return Err(AsyncQueueError::PoolAndTransactionEmpty);
};
Ok(row)
}
pub async fn insert_task(&mut self, task: &dyn AsyncRunnable) -> Result<u64, AsyncQueueError> {
let metadata = serde_json::to_value(task).unwrap();
let task_type = task.task_type();
self.execute(INSERT_TASK_QUERY, &[&metadata, &task_type], Some(1))
.await
}
pub async fn update_task_state(
&mut self,
task: &Task,
state: FangTaskState,
) -> Result<u64, AsyncQueueError> {
let updated_at = Utc::now();
self.execute(
UPDATE_TASK_STATE_QUERY,
&[&state, &updated_at, &task.id],
Some(1),
)
.await
pub async fn insert_task(&mut self, task: &dyn AsyncRunnable) -> Result<u64, AsyncQueueError> {
let mut connection = self.pool.get().await?;
let mut transaction = connection.transaction().await?;
Self::insert_task_query(&mut transaction, task).await
}
pub async fn remove_all_tasks(&mut self) -> Result<u64, AsyncQueueError> {
self.execute(REMOVE_ALL_TASK_QUERY, &[], None).await
}
pub async fn remove_task(&mut self, task: &Task) -> Result<u64, AsyncQueueError> {
self.execute(REMOVE_TASK_QUERY, &[&task.id], Some(1)).await
}
pub async fn remove_tasks_type(&mut self, task_type: &str) -> Result<u64, AsyncQueueError> {
self.execute(REMOVE_TASKS_TYPE_QUERY, &[&task_type], None)
.await
}
pub async fn fail_task(&mut self, task: &Task) -> Result<u64, AsyncQueueError> {
let updated_at = Utc::now();
self.execute(
@ -255,15 +188,9 @@ where
params: &[&(dyn ToSql + Sync)],
expected_result_count: Option<u64>,
) -> Result<u64, AsyncQueueError> {
let result = if let Some(pool) = &self.pool {
let connection = pool.get().await?;
let connection = self.pool.get().await?;
connection.execute(query, params).await?
} else if let Some(transaction) = &self.transaction {
transaction.execute(query, params).await?
} else {
return Err(AsyncQueueError::PoolAndTransactionEmpty);
};
let result = connection.execute(query, params).await?;
if let Some(expected_result) = expected_result_count {
if result != expected_result {
return Err(AsyncQueueError::ResultError {
@ -274,6 +201,109 @@ where
}
Ok(result)
}
pub async fn fetch_and_touch_task_query(
transaction: &mut Transaction<'_>,
task_type: &Option<String>,
) -> Result<Task, AsyncQueueError> {
let mut task = match task_type {
None => Self::get_task_type(transaction, DEFAULT_TASK_TYPE).await?,
Some(task_type_str) => Self::get_task_type(transaction, task_type_str).await?,
};
Self::update_task_state(transaction, &task, FangTaskState::InProgress).await?;
task.state = FangTaskState::InProgress;
Ok(task)
}
pub async fn get_task_type(
transaction: &mut Transaction<'_>,
task_type: &str,
) -> Result<Task, AsyncQueueError> {
let row: Row = transaction
.query_one(FETCH_TASK_TYPE_QUERY, &[&task_type])
.await?;
let task = Self::row_to_task(row);
Ok(task)
}
pub async fn update_task_state(
transaction: &mut Transaction<'_>,
task: &Task,
state: FangTaskState,
) -> Result<u64, AsyncQueueError> {
let updated_at = Utc::now();
Self::execute_query(
transaction,
UPDATE_TASK_STATE_QUERY,
&[&state, &updated_at, &task.id],
Some(1),
)
.await
}
pub async fn insert_task_query(
transaction: &mut Transaction<'_>,
task: &dyn AsyncRunnable,
) -> Result<u64, AsyncQueueError> {
let metadata = serde_json::to_value(task).unwrap();
let task_type = task.task_type();
Self::execute_query(
transaction,
INSERT_TASK_QUERY,
&[&metadata, &task_type],
Some(1),
)
.await
}
pub async fn execute_query(
transaction: &mut Transaction<'_>,
query: &str,
params: &[&(dyn ToSql + Sync)],
expected_result_count: Option<u64>,
) -> Result<u64, AsyncQueueError> {
let result = transaction.execute(query, params).await?;
if let Some(expected_result) = expected_result_count {
if result != expected_result {
return Err(AsyncQueueError::ResultError {
expected: expected_result,
found: result,
});
}
}
Ok(result)
}
fn row_to_task(row: Row) -> Task {
let id: Uuid = row.get("id");
let metadata: serde_json::Value = row.get("metadata");
let error_message: Option<String> = match row.try_get("error_message") {
Ok(error_message) => Some(error_message),
Err(_) => None,
};
let state: FangTaskState = FangTaskState::New;
let task_type: String = row.get("task_type");
let created_at: DateTime<Utc> = row.get("created_at");
let updated_at: DateTime<Utc> = row.get("updated_at");
Task::builder()
.id(id)
.metadata(metadata)
.error_message(error_message)
.state(state)
.task_type(task_type)
.created_at(created_at)
.updated_at(updated_at)
.build()
}
}
#[cfg(test)]
@ -305,71 +335,94 @@ mod async_queue_tests {
async fn insert_task_creates_new_task() {
let pool = pool().await;
let mut connection = pool.get().await.unwrap();
let transaction = connection.transaction().await.unwrap();
let mut queue = AsyncQueue::<NoTls>::new_with_transaction(transaction);
let mut transaction = connection.transaction().await.unwrap();
let result = queue.insert_task(&AsyncTask { number: 1 }).await.unwrap();
let result =
AsyncQueue::<NoTls>::insert_task_query(&mut transaction, &AsyncTask { number: 1 })
.await
.unwrap();
assert_eq!(1, result);
queue.rollback().await.unwrap();
transaction.rollback().await.unwrap();
}
#[tokio::test]
async fn remove_all_tasks_test() {
let pool = pool().await;
let mut connection = pool.get().await.unwrap();
let transaction = connection.transaction().await.unwrap();
let mut queue = AsyncQueue::<NoTls>::new_with_transaction(transaction);
// #[tokio::test]
// async fn remove_all_tasks_test() {
// let pool = pool().await;
// let mut connection = pool.get().await.unwrap();
// let transaction = connection.transaction().await.unwrap();
// let mut queue = AsyncQueue::<NoTls>::new_with_transaction(transaction);
let result = queue.insert_task(&AsyncTask { number: 1 }).await.unwrap();
assert_eq!(1, result);
let result = queue.insert_task(&AsyncTask { number: 2 }).await.unwrap();
assert_eq!(1, result);
let result = queue.remove_all_tasks().await.unwrap();
assert_eq!(2, result);
queue.rollback().await.unwrap();
}
// let result = queue.insert_task(&AsyncTask { number: 1 }).await.unwrap();
// assert_eq!(1, result);
// let result = queue.insert_task(&AsyncTask { number: 2 }).await.unwrap();
// assert_eq!(1, result);
// let result = queue.remove_all_tasks().await.unwrap();
// assert_eq!(2, result);
// queue.rollback().await.unwrap();
// }
#[tokio::test]
async fn fetch_test() {
async fn fetch_and_touch_test() {
let pool = pool().await;
let mut connection = pool.get().await.unwrap();
let transaction = connection.transaction().await.unwrap();
let mut queue = AsyncQueue::<NoTls>::new_with_transaction(transaction);
let mut transaction = connection.transaction().await.unwrap();
let result = queue.insert_task(&AsyncTask { number: 1 }).await.unwrap();
let result =
AsyncQueue::<NoTls>::insert_task_query(&mut transaction, &AsyncTask { number: 1 })
.await
.unwrap();
assert_eq!(1, result);
let result = queue.insert_task(&AsyncTask { number: 2 }).await.unwrap();
let result =
AsyncQueue::<NoTls>::insert_task_query(&mut transaction, &AsyncTask { number: 2 })
.await
.unwrap();
assert_eq!(1, result);
let task = queue.fetch_task(&None).await.unwrap();
let task = AsyncQueue::<NoTls>::fetch_and_touch_task_query(&mut transaction, &None)
.await
.unwrap();
let metadata = task.metadata.as_object().unwrap();
let number = metadata["number"].as_u64();
let type_task = metadata["type"].as_str();
assert_eq!(Some(1), number);
assert_eq!(Some("AsyncTask"), type_task);
let task = queue.fetch_task(&None).await.unwrap();
let task = AsyncQueue::<NoTls>::fetch_and_touch_task_query(&mut transaction, &None)
.await
.unwrap();
let metadata = task.metadata.as_object().unwrap();
let number = metadata["number"].as_u64();
let type_task = metadata["type"].as_str();
assert_eq!(Some(2), number);
assert_eq!(Some("AsyncTask"), type_task);
queue.rollback().await.unwrap();
}
#[tokio::test]
async fn remove_tasks_type_test() {
let pool = pool().await;
let mut connection = pool.get().await.unwrap();
let transaction = connection.transaction().await.unwrap();
let mut queue = AsyncQueue::<NoTls>::new_with_transaction(transaction);
let result = queue.insert_task(&AsyncTask { number: 1 }).await.unwrap();
assert_eq!(1, result);
let result = queue.insert_task(&AsyncTask { number: 2 }).await.unwrap();
assert_eq!(1, result);
let result = queue.remove_tasks_type("common").await.unwrap();
assert_eq!(2, result);
queue.rollback().await.unwrap();
transaction.rollback().await.unwrap();
}
// #[tokio::test]
// async fn remove_tasks_type_test() {
// let pool = pool().await;
// let mut connection = pool.get().await.unwrap();
// let transaction = connection.transaction().await.unwrap();
// let mut queue = AsyncQueue::<NoTls>::new_with_transaction(transaction);
// let result = queue.insert_task(&AsyncTask { number: 1 }).await.unwrap();
// assert_eq!(1, result);
// let result = queue.insert_task(&AsyncTask { number: 2 }).await.unwrap();
// assert_eq!(1, result);
// let result = queue.remove_tasks_type("common").await.unwrap();
// assert_eq!(2, result);
// queue.rollback().await.unwrap();
// }
async fn pool() -> Pool<PostgresConnectionManager<NoTls>> {
let pg_mgr = PostgresConnectionManager::new_from_stringlike(