personal/monzo-ingestion
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

rustfmt
All checks were successful
Build and Publish Docker Container / build (push) Successful in 6m33s
Details

Browse Source
This commit is contained in:
Joshua Coles 2024-05-28 17:38:45 +01:00
parent fbf473b3b4
commit 8ce60ce278
1 changed files with 49 additions and 36 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

85
src/ingestion/db.rs
View File

@ -1,10 +1,13 @@
use crate::error::AppError;
use anyhow::anyhow; use anyhow::anyhow;
use entity::{expenditure, transaction}; use entity::{expenditure, transaction};
use sea_orm::sea_query::OnConflict;
use sea_orm::{ConnectionTrait, DatabaseBackend, DatabaseTransaction, DbErr, QueryFilter, QueryTrait, Statement};
use sea_orm::{ColumnTrait, DatabaseConnection, EntityTrait, Iterable, TransactionTrait};
use migration::PostgresQueryBuilder; use migration::PostgresQueryBuilder;
use crate::error::AppError; use sea_orm::sea_query::OnConflict;
use sea_orm::{ColumnTrait, DatabaseConnection, EntityTrait, Iterable, TransactionTrait};
use sea_orm::{
ConnectionTrait, DatabaseBackend, DatabaseTransaction, DbErr, QueryFilter, QueryTrait,
Statement,
};
pub struct Insertion { pub struct Insertion {
pub transaction: transaction::ActiveModel, pub transaction: transaction::ActiveModel,
@ -31,16 +34,16 @@ async fn update_expenditures(
.flat_map(|i| &i.contained_expenditures) .flat_map(|i| &i.contained_expenditures)
.cloned(), .cloned(),
) )
.on_conflict( .on_conflict(
OnConflict::columns(vec![ OnConflict::columns(vec![
expenditure::Column::TransactionId, expenditure::Column::TransactionId,
expenditure::Column::Category, expenditure::Column::Category,
]) ])
.update_columns(expenditure::Column::iter()) .update_columns(expenditure::Column::iter())
.to_owned(), .to_owned(),
) )
.exec(tx) .exec(tx)
.await?; .await?;
Ok(()) Ok(())
} }
@ -48,7 +51,10 @@ async fn update_expenditures(
// Note while this is more efficient in db calls, it does bind together the entire group. // Note while this is more efficient in db calls, it does bind together the entire group.
// We employ a batching process for now to try balance speed and failure rate, but it is worth // We employ a batching process for now to try balance speed and failure rate, but it is worth
// trying to move failures earlier and improve reporting. // trying to move failures earlier and improve reporting.
pub async fn insert(db: &DatabaseConnection, insertions: Vec<Insertion>) -> Result<Vec<String>, AppError> { pub async fn insert(
db: &DatabaseConnection,
insertions: Vec<Insertion>,
) -> Result<Vec<String>, AppError> {
let mut new_transaction_ids = Vec::new(); let mut new_transaction_ids = Vec::new();
for insertions in insertions.chunks(400) { for insertions in insertions.chunks(400) {
@ -68,35 +74,42 @@ pub async fn insert(db: &DatabaseConnection, insertions: Vec<Insertion>) -> Resu
Ok(new_transaction_ids) Ok(new_transaction_ids)
} }
async fn update_transactions(insertions: &[Insertion], tx: &DatabaseTransaction) -> Result<Vec<String>, AppError> { async fn update_transactions(
let insert = transaction::Entity::insert_many(insertions.iter().map(|i| &i.transaction).cloned()) insertions: &[Insertion],
.on_conflict( tx: &DatabaseTransaction,
OnConflict::column(transaction::Column::Id) ) -> Result<Vec<String>, AppError> {
.update_columns(transaction::Column::iter()) let insert =
.to_owned(), transaction::Entity::insert_many(insertions.iter().map(|i| &i.transaction).cloned())
) .on_conflict(
.into_query() OnConflict::column(transaction::Column::Id)
.returning_col(transaction::Column::Id) .update_columns(transaction::Column::iter())
.build(PostgresQueryBuilder); .to_owned(),
)
.into_query()
.returning_col(transaction::Column::Id)
.build(PostgresQueryBuilder);
let inserted_transaction_ids = tx.query_all(Statement::from_sql_and_values( let inserted_transaction_ids = tx
DatabaseBackend::Postgres, .query_all(Statement::from_sql_and_values(
insert.0, DatabaseBackend::Postgres,
insert.1, insert.0,
)).await? insert.1,
))
.await?
.iter() .iter()
.map(|r| r.try_get_by("id")) .map(|r| r.try_get_by("id"))
.collect::<Result<Vec<String>, _>>()?; .collect::<Result<Vec<String>, _>>()?;
Ok(inserted_transaction_ids) Ok(inserted_transaction_ids)
} }
async fn notify_new_transactions(db: &DatabaseConnection, new_transaction_ids: &[String]) -> Result<(), AppError> { async fn notify_new_transactions(
let payload = serde_json::to_string(&new_transaction_ids) db: &DatabaseConnection,
.map_err(|e| anyhow!(e))?; new_transaction_ids: &[String],
) -> Result<(), AppError> {
let payload = serde_json::to_string(&new_transaction_ids).map_err(|e| anyhow!(e))?;
db.execute_unprepared( db.execute_unprepared(&format!(r#"NOTIFY monzo_new_transactions, {payload}"#,))
&format!(r#"NOTIFY monzo_new_transactions, {payload}"#, ), .await?;
).await?;
Ok(()) Ok(())
} }