diff --git a/src/bin/main.rs b/src/bin/main.rs index 50f6ebe..9dc3396 100644 --- a/src/bin/main.rs +++ b/src/bin/main.rs @@ -1,10 +1,8 @@ use clap::Parser; -use serde::Deserialize; use std::time::Duration; -use tokio_postgres::NoTls; use url::Url; -use allegedly::{Dt, ExportPage, bin_init, poll_upstream, week_to_pages}; +use allegedly::{Db, Dt, ExportPage, Op, bin_init, poll_upstream, week_to_pages}; const EXPORT_PAGE_QUEUE_SIZE: usize = 0; // rendezvous for now const WEEK_IN_SECONDS: u64 = 7 * 86400; @@ -42,17 +40,6 @@ struct Args { postgres: String, } -#[derive(Deserialize)] -#[serde(rename_all = "camelCase")] -struct Op<'a> { - pub did: &'a str, - pub cid: &'a str, - pub created_at: Dt, - pub nullified: bool, - #[serde(borrow)] - pub operation: &'a serde_json::value::RawValue, -} - async fn bulk_backfill((upstream, epoch): (Url, u64), tx: flume::Sender) { let immutable_cutoff = std::time::SystemTime::now() - Duration::from_secs((7 + 4) * 86400); let immutable_ts = (immutable_cutoff.duration_since(std::time::SystemTime::UNIX_EPOCH)) @@ -181,52 +168,31 @@ async fn get_latest(pg_client: &tokio_postgres::Client) -> Option
{ } #[tokio::main] -async fn main() { +async fn main() -> anyhow::Result<()> { bin_init("main"); let args = Args::parse(); - - log::trace!("connecting postgres..."); - let (pg_client, connection) = tokio_postgres::connect(&args.postgres, NoTls) - .await - .unwrap(); - - // send the connection away to do the actual communication work - // TODO: error and shutdown handling - let conn_task = tokio::task::spawn(async move { - if let Err(e) = connection.await { - eprintln!("connection error: {e}"); - } - }); - - log::trace!("connecting postgres 2..."); - let (pg_client2, connection2) = tokio_postgres::connect(&args.postgres, NoTls) - .await - .unwrap(); - - // send the connection away to do the actual communication work - // TODO: error and shutdown handling - let conn_task2 = tokio::task::spawn(async move { - if let Err(e) = connection2.await { - eprintln!("connection error: {e}"); - } - }); - + let db = Db::new(&args.postgres); let (tx, rx) = flume::bounded(EXPORT_PAGE_QUEUE_SIZE); + log::trace!("connecting postgres for export task..."); + let pg_client = db.connect().await?; let export_task = tokio::task::spawn(export_upstream( args.upstream, (args.upstream_bulk, args.bulk_epoch), tx, - pg_client2, + pg_client, )); + + log::trace!("connecting postgres for writer task..."); + let pg_client = db.connect().await?; let writer_task = tokio::task::spawn(write_pages(rx, pg_client)); tokio::select! { - z = conn_task => log::warn!("connection task ended: {z:?}"), - z = conn_task2 => log::warn!("connection task ended: {z:?}"), z = export_task => log::warn!("export task ended: {z:?}"), z = writer_task => log::warn!("writer task ended: {z:?}"), }; log::error!("todo: shutdown"); + + Ok(()) } diff --git a/src/lib.rs b/src/lib.rs index ea0efe0..3053fc7 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -2,10 +2,12 @@ use serde::Deserialize; mod backfill; mod client; +mod plc_pg; mod poll; pub use backfill::week_to_pages; pub use client::CLIENT; +pub use plc_pg::Db; pub use poll::poll_upstream; pub type Dt = chrono::DateTime; diff --git a/src/plc_pg.rs b/src/plc_pg.rs new file mode 100644 index 0000000..3050b69 --- /dev/null +++ b/src/plc_pg.rs @@ -0,0 +1,31 @@ +use tokio_postgres::{Client, Error as PgError, NoTls, connect}; + +/// a little tokio-postgres helper +#[derive(Debug, Clone)] +pub struct Db { + pg_uri: String, +} + +impl Db { + pub fn new(pg_uri: &str) -> Self { + Self { + pg_uri: pg_uri.to_string(), + } + } + + pub async fn connect(&self) -> Result { + log::trace!("connecting postgres..."); + let (client, connection) = connect(&self.pg_uri, NoTls).await?; + + // send the connection away to do the actual communication work + // apparently the connection will complete when the client drops + tokio::task::spawn(async move { + connection + .await + .inspect_err(|e| log::error!("connection ended with error: {e}")) + .unwrap(); + }); + + Ok(client) + } +}