karaolidis/qrust
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add multiple asset adding route

Browse Source 
Signed-off-by: Nikolaos Karaolidis <nick@karaolidis.com>
This commit is contained in:
Nick Karaolidis
2024-03-09 20:13:36 +00:00
parent 080f91b044
commit 681d7393d7
31 changed files with 754 additions and 282 deletions
Download Patch File Download Diff File Expand all files Collapse all files

166
src/threads/data/backfill.rs
View File

@@ -2,7 +2,7 @@ use super::ThreadType;
use crate::{
config::{
Config, ALPACA_CRYPTO_DATA_API_URL, ALPACA_SOURCE, ALPACA_STOCK_DATA_API_URL,
MAX_BERT_INPUTS,
BERT_MAX_INPUTS,
},
database,
types::{
@@ -30,24 +30,14 @@ pub enum Action {
Purge,
}
impl From<super::Action> for Option<Action> {
fn from(action: super::Action) -> Self {
match action {
super::Action::Add | super::Action::Enable => Some(Action::Backfill),
super::Action::Remove => Some(Action::Purge),
super::Action::Disable => None,
}
}
}
pub struct Message {
pub action: Option<Action>,
pub action: Action,
pub symbols: Vec<String>,
pub response: oneshot::Sender<()>,
}
impl Message {
pub fn new(action: Option<Action>, symbols: Vec<String>) -> (Self, oneshot::Receiver<()>) {
pub fn new(action: Action, symbols: Vec<String>) -> (Self, oneshot::Receiver<()>) {
let (sender, receiver) = oneshot::channel::<()>();
(
Self {
@@ -62,10 +52,10 @@ impl Message {
#[async_trait]
pub trait Handler: Send + Sync {
async fn select_latest_backfill(
async fn select_latest_backfills(
&self,
symbol: String,
) -> Result<Option<Backfill>, clickhouse::error::Error>;
symbols: &[String],
) -> Result<Vec<Backfill>, clickhouse::error::Error>;
async fn delete_backfills(&self, symbol: &[String]) -> Result<(), clickhouse::error::Error>;
async fn delete_data(&self, symbol: &[String]) -> Result<(), clickhouse::error::Error>;
async fn queue_backfill(&self, symbol: &str, fetch_to: OffsetDateTime);
@@ -94,9 +84,17 @@ async fn handle_backfill_message(
let mut backfill_jobs = backfill_jobs.lock().await;
match message.action {
Some(Action::Backfill) => {
Action::Backfill => {
let log_string = handler.log_string();
let backfills = handler
.select_latest_backfills(&message.symbols)
.await
.unwrap()
.into_iter()
.map(|backfill| (backfill.symbol.clone(), backfill))
.collect::<HashMap<_, _>>();
for symbol in message.symbols {
if let Some(job) = backfill_jobs.get(&symbol) {
if !job.is_finished() {
@@ -108,33 +106,30 @@ async fn handle_backfill_message(
}
}
let fetch_from = backfills
.get(&symbol)
.map_or(OffsetDateTime::UNIX_EPOCH, |backfill| {
backfill.time + ONE_SECOND
});
let fetch_to = last_minute();
if fetch_from > fetch_to {
info!("No need to backfill {} {}.", symbol, log_string,);
return;
}
let handler = handler.clone();
backfill_jobs.insert(
symbol.clone(),
spawn(async move {
let fetch_from = match handler
.select_latest_backfill(symbol.clone())
.await
.unwrap()
{
Some(latest_backfill) => latest_backfill.time + ONE_SECOND,
None => OffsetDateTime::UNIX_EPOCH,
};
let fetch_to = last_minute();
if fetch_from > fetch_to {
info!("No need to backfill {} {}.", symbol, log_string,);
return;
}
handler.queue_backfill(&symbol, fetch_to).await;
handler.backfill(symbol, fetch_from, fetch_to).await;
}),
);
}
}
Some(Action::Purge) => {
Action::Purge => {
for symbol in &message.symbols {
if let Some(job) = backfill_jobs.remove(symbol) {
if !job.is_finished() {
@@ -150,7 +145,6 @@ async fn handle_backfill_message(
)
.unwrap();
}
None => {}
}
message.response.send(()).unwrap();
@@ -199,20 +193,34 @@ fn crypto_query_constructor(
#[async_trait]
impl Handler for BarHandler {
async fn select_latest_backfill(
async fn select_latest_backfills(
&self,
symbol: String,
) -> Result<Option<Backfill>, clickhouse::error::Error> {
database::backfills_bars::select_where_symbol(&self.config.clickhouse_client, &symbol).await
symbols: &[String],
) -> Result<Vec<Backfill>, clickhouse::error::Error> {
database::backfills_bars::select_where_symbols(
&self.config.clickhouse_client,
&self.config.clickhouse_concurrency_limiter,
symbols,
)
.await
}
async fn delete_backfills(&self, symbols: &[String]) -> Result<(), clickhouse::error::Error> {
database::backfills_bars::delete_where_symbols(&self.config.clickhouse_client, symbols)
.await
database::backfills_bars::delete_where_symbols(
&self.config.clickhouse_client,
&self.config.clickhouse_concurrency_limiter,
symbols,
)
.await
}
async fn delete_data(&self, symbols: &[String]) -> Result<(), clickhouse::error::Error> {
database::bars::delete_where_symbols(&self.config.clickhouse_client, symbols).await
database::bars::delete_where_symbols(
&self.config.clickhouse_client,
&self.config.clickhouse_concurrency_limiter,
symbols,
)
.await
}
async fn queue_backfill(&self, symbol: &str, fetch_to: OffsetDateTime) {
@@ -230,7 +238,7 @@ impl Handler for BarHandler {
let mut next_page_token = None;
loop {
let Ok(message) = alpaca::api::incoming::bar::get_historical(
let Ok(message) = alpaca::api::incoming::bar::get(
&self.config.alpaca_client,
&self.config.alpaca_rate_limiter,
self.data_url,
@@ -267,12 +275,20 @@ impl Handler for BarHandler {
let backfill = bars.last().unwrap().clone().into();
database::bars::upsert_batch(&self.config.clickhouse_client, &bars)
.await
.unwrap();
database::backfills_bars::upsert(&self.config.clickhouse_client, &backfill)
.await
.unwrap();
database::bars::upsert_batch(
&self.config.clickhouse_client,
&self.config.clickhouse_concurrency_limiter,
&bars,
)
.await
.unwrap();
database::backfills_bars::upsert(
&self.config.clickhouse_client,
&self.config.clickhouse_concurrency_limiter,
&backfill,
)
.await
.unwrap();
info!("Backfilled bars for {}.", symbol);
}
@@ -288,20 +304,34 @@ struct NewsHandler {
#[async_trait]
impl Handler for NewsHandler {
async fn select_latest_backfill(
async fn select_latest_backfills(
&self,
symbol: String,
) -> Result<Option<Backfill>, clickhouse::error::Error> {
database::backfills_news::select_where_symbol(&self.config.clickhouse_client, &symbol).await
symbols: &[String],
) -> Result<Vec<Backfill>, clickhouse::error::Error> {
database::backfills_news::select_where_symbols(
&self.config.clickhouse_client,
&self.config.clickhouse_concurrency_limiter,
symbols,
)
.await
}
async fn delete_backfills(&self, symbols: &[String]) -> Result<(), clickhouse::error::Error> {
database::backfills_news::delete_where_symbols(&self.config.clickhouse_client, symbols)
.await
database::backfills_news::delete_where_symbols(
&self.config.clickhouse_client,
&self.config.clickhouse_concurrency_limiter,
symbols,
)
.await
}
async fn delete_data(&self, symbols: &[String]) -> Result<(), clickhouse::error::Error> {
database::news::delete_where_symbols(&self.config.clickhouse_client, symbols).await
database::news::delete_where_symbols(
&self.config.clickhouse_client,
&self.config.clickhouse_concurrency_limiter,
symbols,
)
.await
}
async fn queue_backfill(&self, symbol: &str, fetch_to: OffsetDateTime) {
@@ -317,7 +347,7 @@ impl Handler for NewsHandler {
let mut next_page_token = None;
loop {
let Ok(message) = alpaca::api::incoming::news::get_historical(
let Ok(message) = alpaca::api::incoming::news::get(
&self.config.alpaca_client,
&self.config.alpaca_rate_limiter,
&alpaca::api::outgoing::news::News {
@@ -355,7 +385,7 @@ impl Handler for NewsHandler {
.map(|news| format!("{}\n\n{}", news.headline, news.content))
.collect::<Vec<_>>();
let predictions = join_all(inputs.chunks(*MAX_BERT_INPUTS).map(|inputs| async move {
let predictions = join_all(inputs.chunks(*BERT_MAX_INPUTS).map(|inputs| async move {
let sequence_classifier = self.config.sequence_classifier.lock().await;
block_in_place(|| {
sequence_classifier
@@ -381,12 +411,20 @@ impl Handler for NewsHandler {
let backfill = (news.last().unwrap().clone(), symbol.clone()).into();
database::news::upsert_batch(&self.config.clickhouse_client, &news)
.await
.unwrap();
database::backfills_news::upsert(&self.config.clickhouse_client, &backfill)
.await
.unwrap();
database::news::upsert_batch(
&self.config.clickhouse_client,
&self.config.clickhouse_concurrency_limiter,
&news,
)
.await
.unwrap();
database::backfills_news::upsert(
&self.config.clickhouse_client,
&self.config.clickhouse_concurrency_limiter,
&backfill,
)
.await
.unwrap();
info!("Backfilled news for {}.", symbol);
}

201
src/threads/data/mod.rs
View File

@@ -9,18 +9,18 @@ use crate::{
},
create_send_await, database,
types::{alpaca, Asset, Class},
utils::backoff,
};
use futures_util::{future::join_all, StreamExt};
use futures_util::StreamExt;
use itertools::{Either, Itertools};
use std::sync::Arc;
use log::error;
use std::{collections::HashMap, sync::Arc};
use tokio::{
join, select, spawn,
sync::{mpsc, oneshot},
};
use tokio_tungstenite::connect_async;
#[derive(Clone, Copy)]
#[derive(Clone, Copy, PartialEq, Eq)]
#[allow(dead_code)]
pub enum Action {
Add,
@@ -173,13 +173,6 @@ async fn handle_message(
message.action.into(),
us_equity_symbols.clone()
);
create_send_await!(
bars_us_equity_backfill_sender,
backfill::Message::new,
message.action.into(),
us_equity_symbols
);
};
let bars_crypto_future = async {
@@ -193,13 +186,6 @@ async fn handle_message(
message.action.into(),
crypto_symbols.clone()
);
create_send_await!(
bars_crypto_backfill_sender,
backfill::Message::new,
message.action.into(),
crypto_symbols
);
};
let news_future = async {
@@ -209,62 +195,127 @@ async fn handle_message(
message.action.into(),
symbols.clone()
);
create_send_await!(
news_backfill_sender,
backfill::Message::new,
message.action.into(),
symbols.clone()
);
};
join!(bars_us_equity_future, bars_crypto_future, news_future);
match message.action {
Action::Add => {
let assets = join_all(symbols.into_iter().map(|symbol| {
let config = config.clone();
async move {
let asset_future = async {
alpaca::api::incoming::asset::get_by_symbol(
&config.alpaca_client,
&config.alpaca_rate_limiter,
&symbol,
Some(backoff::infinite()),
)
.await
.unwrap()
};
let position_future = async {
alpaca::api::incoming::position::get_by_symbol(
&config.alpaca_client,
&config.alpaca_rate_limiter,
&symbol,
Some(backoff::infinite()),
)
.await
.unwrap()
};
let (asset, position) = join!(asset_future, position_future);
Asset::from((asset, position))
}
}))
.await;
database::assets::upsert_batch(&config.clickhouse_client, &assets)
let assets = async {
alpaca::api::incoming::asset::get_by_symbols(
&config.alpaca_client,
&config.alpaca_rate_limiter,
&symbols,
None,
)
.await
.unwrap();
.unwrap()
.into_iter()
.map(|asset| (asset.symbol.clone(), asset))
.collect::<HashMap<_, _>>()
};
let positions = async {
alpaca::api::incoming::position::get_by_symbols(
&config.alpaca_client,
&config.alpaca_rate_limiter,
&symbols,
None,
)
.await
.unwrap()
.into_iter()
.map(|position| (position.symbol.clone(), position))
.collect::<HashMap<_, _>>()
};
let (mut assets, mut positions) = join!(assets, positions);
let mut batch = vec![];
for symbol in &symbols {
if let Some(asset) = assets.remove(symbol) {
let position = positions.remove(symbol);
batch.push(Asset::from((asset, position)));
} else {
error!("Failed to find asset for symbol: {}", symbol);
}
}
database::assets::upsert_batch(
&config.clickhouse_client,
&config.clickhouse_concurrency_limiter,
&batch,
)
.await
.unwrap();
}
Action::Remove => {
database::assets::delete_where_symbols(&config.clickhouse_client, &symbols)
.await
.unwrap();
database::assets::delete_where_symbols(
&config.clickhouse_client,
&config.clickhouse_concurrency_limiter,
&symbols,
)
.await
.unwrap();
}
_ => {}
}
if message.action == Action::Disable {
message.response.send(()).unwrap();
return;
}
let bars_us_equity_future = async {
if us_equity_symbols.is_empty() {
return;
}
create_send_await!(
bars_us_equity_backfill_sender,
backfill::Message::new,
match message.action {
Action::Add | Action::Enable => backfill::Action::Backfill,
Action::Remove => backfill::Action::Purge,
Action::Disable => unreachable!(),
},
us_equity_symbols
);
};
let bars_crypto_future = async {
if crypto_symbols.is_empty() {
return;
}
create_send_await!(
bars_crypto_backfill_sender,
backfill::Message::new,
match message.action {
Action::Add | Action::Enable => backfill::Action::Backfill,
Action::Remove => backfill::Action::Purge,
Action::Disable => unreachable!(),
},
crypto_symbols
);
};
let news_future = async {
create_send_await!(
news_backfill_sender,
backfill::Message::new,
match message.action {
Action::Add | Action::Enable => backfill::Action::Backfill,
Action::Remove => backfill::Action::Purge,
Action::Disable => unreachable!(),
},
symbols
);
};
join!(bars_us_equity_future, bars_crypto_future, news_future);
message.response.send(()).unwrap();
}
@@ -274,13 +325,19 @@ async fn handle_clock_message(
bars_crypto_backfill_sender: mpsc::Sender<backfill::Message>,
news_backfill_sender: mpsc::Sender<backfill::Message>,
) {
database::cleanup_all(&config.clickhouse_client)
.await
.unwrap();
database::cleanup_all(
&config.clickhouse_client,
&config.clickhouse_concurrency_limiter,
)
.await
.unwrap();
let assets = database::assets::select(&config.clickhouse_client)
.await
.unwrap();
let assets = database::assets::select(
&config.clickhouse_client,
&config.clickhouse_concurrency_limiter,
)
.await
.unwrap();
let (us_equity_symbols, crypto_symbols): (Vec<_>, Vec<_>) = assets
.clone()
@@ -299,8 +356,8 @@ async fn handle_clock_message(
create_send_await!(
bars_us_equity_backfill_sender,
backfill::Message::new,
Some(backfill::Action::Backfill),
us_equity_symbols.clone()
backfill::Action::Backfill,
us_equity_symbols
);
};
@@ -308,8 +365,8 @@ async fn handle_clock_message(
create_send_await!(
bars_crypto_backfill_sender,
backfill::Message::new,
Some(backfill::Action::Backfill),
crypto_symbols.clone()
backfill::Action::Backfill,
crypto_symbols
);
};
@@ -317,7 +374,7 @@ async fn handle_clock_message(
create_send_await!(
news_backfill_sender,
backfill::Message::new,
Some(backfill::Action::Backfill),
backfill::Action::Backfill,
symbols
);
};

22
src/threads/data/websocket.rs
View File

@@ -268,9 +268,13 @@ impl Handler for BarsHandler {
let bar = Bar::from(message);
debug!("Received bar for {}: {}.", bar.symbol, bar.time);
database::bars::upsert(&self.config.clickhouse_client, &bar)
.await
.unwrap();
database::bars::upsert(
&self.config.clickhouse_client,
&self.config.clickhouse_concurrency_limiter,
&bar,
)
.await
.unwrap();
}
websocket::data::incoming::Message::Status(message) => {
debug!(
@@ -283,6 +287,7 @@ impl Handler for BarsHandler {
| websocket::data::incoming::status::Status::VolatilityTradingPause(_) => {
database::assets::update_status_where_symbol(
&self.config.clickhouse_client,
&self.config.clickhouse_concurrency_limiter,
&message.symbol,
false,
)
@@ -293,6 +298,7 @@ impl Handler for BarsHandler {
| websocket::data::incoming::status::Status::TradingResumption(_) => {
database::assets::update_status_where_symbol(
&self.config.clickhouse_client,
&self.config.clickhouse_concurrency_limiter,
&message.symbol,
true,
)
@@ -398,9 +404,13 @@ impl Handler for NewsHandler {
..news
};
database::news::upsert(&self.config.clickhouse_client, &news)
.await
.unwrap();
database::news::upsert(
&self.config.clickhouse_client,
&self.config.clickhouse_concurrency_limiter,
&news,
)
.await
.unwrap();
}
websocket::data::incoming::Message::Error(message) => {
error!("Received error message: {}.", message.message);