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
GitHub
parent
abe324fc57
commit
3602097fb6
1 changed files with 203 additions and 150 deletions
|
|
@ -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(
|
||||
|
|
|
|||
Loading…
Reference in a new issue