poloniex_dump/src/markets/poloniex.rs

use std::pin::Pin;

use chrono::NaiveDateTime;
use futures_util::{SinkExt, Stream, TryStreamExt as _};
use reqwest::Url;
use reqwest_websocket::Message;

use crate::{
    error::{AppError, AppResult},
    models::{
        Candle, CandleInterval, SubscriptionRequest, SubscriptionResponse,
        SubscriptionResponseData, Trade,
    },
};

use super::Market;

pub struct PoloniexClient {
    rest_base_url: Url,
    ws_base_url: Url,
}

impl PoloniexClient {
    pub fn new(rest_base_url: &Url, ws_base_url: &Url) -> Self {
        Self {
            rest_base_url: rest_base_url.to_owned(),
            ws_base_url: ws_base_url.to_owned(),
        }
    }
}

#[async_trait]
impl Market for PoloniexClient {
    async fn get_historical_candles(
        &self,
        pair: &str,
        interval: CandleInterval,
        start_date: NaiveDateTime,
        end_date: NaiveDateTime,
        limit: usize,
    ) -> AppResult<Vec<Candle>> {
        let mut req = self
            .rest_base_url
            .join(&format!("/markets/{}/candles", pair))?;

        req.query_pairs_mut()
            .append_pair(
                "startTime",
                &start_date.and_utc().timestamp_millis().to_string(),
            )
            .append_pair(
                "endTime",
                &end_date.and_utc().timestamp_millis().to_string(),
            )
            .append_pair("interval", interval.as_ref())
            .append_pair("limit", &limit.to_string());

        let result = reqwest::get(req).await?;

        if let Err(e) = result.error_for_status_ref() {
            println!("{:?}", result.text().await);
            return Err(e.into());
        };
        result.json().await.map_err(AppError::from)
    }

    async fn recent_trades_stream(
        &self,
        pairs: &[String],
    ) -> AppResult<Pin<Box<dyn Stream<Item = AppResult<Trade>> + Send>>> {
        let mut ws = reqwest_websocket::websocket(self.ws_base_url.clone()).await?;

        let req = SubscriptionRequest {
            event: "subscribe".to_string(),
            channel: vec!["trades".to_string()],
            symbols: pairs.to_vec(),
        };

        ws.send(Message::Text(serde_json::to_string(&req)?)).await?;

        Ok(Box::pin(async_stream::stream! {
            while let Some(message) = ws.try_next().await? {
                match message {
                    Message::Text(text) => {
                        if let Ok(sub) = serde_json::from_str::<SubscriptionResponse>(&text) {
                            println!("{sub:?}");
                            continue
                        }

                        let trades = serde_json::from_str::<SubscriptionResponseData<Trade>>(&text)?;
                        for trade in trades.data {
                            yield Ok(trade)
                        }
                    },
                    Message::Close {..} => {
                        eprintln!("trades stream closed: {message:?}");
                        break
                    },
                    m => eprintln!("unknown message {m:?}")
                }
            }
        }))
    }
}