omry_integration/

bridge.rs

//! Module with APIs that make it more convenient to test the Omry RPC server.
use anyhow::{Context, Result};
use omry_search::wrappers::SearchParametersWrapper;
use omry_service::archiving::{RawRecord, Record, RecordMeta, RecordParams};
use omry_service::common::RecordId;
use omry_service::search::SearchResponse;
use omry_service::service::{MetaHit, SearchRespMeta};
use omry_service::{OmryApi, PONG};
use tokio::task::JoinSet;
use tracing::instrument;

/// Trait to "bridge" the user-facing application with the flora backend
/// via [`omry_service`].
#[allow(async_fn_in_trait)]
pub trait ApiBridge {
    /// Returns a reference to the underlying tarpc client used by this bridge.
    fn api(&self) -> &OmryApi;

    /// Pings the flora server.
    ///
    /// Returns the expected [`PONG`] response on success.
    #[instrument(skip(self))]
    async fn ping(&self) -> Result<String> {
        let rx = self
            .api()
            .ping()
            .await
            .context("error sending ping request")?;
        let resp = rx.await?;
        if resp == PONG {
            Ok(resp)
        } else {
            anyhow::bail!("Unexpected ping response: {resp}")
        }
    }

    /// Adds a single record to the flora archive.
    #[instrument(skip(self, record))]
    async fn add_record(&self, record: RecordParams) -> Result<RecordId> {
        let (tx, mut rx) = self.api().add().await.context("error calling add")?;
        tx.send(record).await.context("Failed to send record")?;
        drop(tx);

        let id = rx
            .recv()
            .await?
            .ok_or_else(|| anyhow::anyhow!("got empty response"))?;
        Ok(id)
    }

    /// Adds multiple records.
    #[instrument(skip(self, records))]
    async fn add_many_records(&self, records: Vec<RecordParams>) -> Result<Vec<RecordId>> {
        let (tx, mut rx) = self.api().add().await.context("error calling add")?;
        let n_recs = records.len();

        let mut set = JoinSet::new();
        for rec in records {
            let tx = tx.clone();
            set.spawn(async move { tx.send(rec).await.context("error sending record") });
        }

        set.join_all().await;
        drop(tx);

        let mut ids = Vec::with_capacity(n_recs);
        while let Some(id) = rx.recv().await.context("error receiving response")? {
            ids.push(id);
        }

        ids.sort_unstable();
        Ok(ids)
    }

    /// Queries the backend for the number of records it currently stores.
    #[instrument(skip(self))]
    async fn count(&self) -> Result<i64> {
        let rx = self.api().count().await.context("failed calling count")?;
        Ok(rx.await?)
    }

    /// Queries the flora backend for records with the given IDs.
    #[instrument(skip(self))]
    async fn get_records(&self, record_ids: Vec<RecordId>) -> Result<Vec<Record>> {
        let (tx, mut rx) = self
            .api()
            .record()
            .await
            .context("failed calling get record")?;
        let mut set = JoinSet::new();

        for id in &record_ids {
            let tx = tx.clone();
            let id = *id;
            set.spawn(async move { tx.send(id).await });
        }

        drop(tx);

        set.join_all().await;
        let mut records = Vec::with_capacity(record_ids.len());

        // The last item from the stream is `None`, indicating we've reached the end.
        // The outer `None` would be for when the channel got closed.
        while let Some(Some(record)) = rx.recv().await? {
            records.push(record);
        }

        // the returned records may not be in the same order as our input IDs
        #[allow(clippy::expect_used)]
        records.sort_by_key(|rec| {
            let this_id = rec.id().expect("returned record should have an ID");
            record_ids
                .iter()
                .position(|id| id == &this_id)
                .expect("returned ID should be present among original IDs")
        });

        Ok(records)
    }

    /// Returns all the records in the database.
    #[instrument(skip(self))]
    async fn get_all_records(&self) -> Result<Vec<Record>> {
        let mut start_row = None;

        let mut recs = vec![];
        'paginate: loop {
            let mut rx = self
                .api()
                .all_records(start_row)
                .await
                .context("Failed to invoke 'all_records' RPC")?;

            while let Some(maybe_rec) = rx.recv().await? {
                let Some(rec) = maybe_rec else {
                    break 'paginate;
                };

                recs.push(rec);
            }

            start_row = Some(recs.len());
        }

        Ok(recs)
    }

    /// Returns the metadata for all the records in the database.
    #[instrument(skip(self))]
    async fn get_all_records_meta(&self) -> Result<Vec<RecordMeta>> {
        let mut rec_metas = vec![];
        let mut start_row = None;

        'paginate: loop {
            let mut rx = self
                .api()
                .all_meta(start_row)
                .await
                .context("Failed to invoke 'all_meta' RPC")?;

            while let Some(maybe_rec) = rx.recv().await? {
                let Some(rec) = maybe_rec else {
                    break 'paginate;
                };

                rec_metas.push(rec);
            }

            start_row = Some(rec_metas.len());
        }

        Ok(rec_metas)
    }

    /// Returns raw data for all the records in the DB.
    #[instrument(skip(self))]
    async fn get_all_raw(&self) -> Result<Vec<RawRecord>> {
        let mut raw_recs = vec![];
        let mut start_row = None;

        'paginate: loop {
            let mut rx = self
                .api()
                .all_raw(start_row)
                .await
                .context("error invoking 'all_raw'")?;

            while let Some(maybe_rec) = rx.recv().await? {
                let Some(rec) = maybe_rec else {
                    break 'paginate;
                };

                raw_recs.push(rec);
            }

            start_row = Some(raw_recs.len());
        }

        Ok(raw_recs)
    }

    /// Returns record metadatas, most recent first.
    #[instrument(skip(self))]
    async fn most_recent_meta(&self) -> Result<Vec<RecordMeta>> {
        let mut rec_metas = vec![];
        let mut start_idx = None;

        'paginate: loop {
            let mut rx = self
                .api()
                .most_recent_meta(start_idx)
                .await
                .context("Error invoking 'most_recent_meta'")?;

            while let Some(maybe_rec) = rx.recv().await? {
                let Some(rec) = maybe_rec else {
                    break 'paginate;
                };

                rec_metas.push(rec);
            }

            start_idx = Some(rec_metas.len());
        }

        Ok(rec_metas)
    }

    /// Sends a search request to the flora backend.
    #[instrument(skip(self))]
    async fn search(&self, search_parameters: SearchParametersWrapper) -> Result<SearchResponse> {
        let mut rx = self
            .api()
            .search(search_parameters)
            .await
            .context("failed while calling search")?;

        let mut metahits = vec![];
        while let Some(item) = rx.recv().await? {
            metahits.push(item);
        }

        metahits.sort_by_key(|hit| hit.index);
        let last_mh = metahits.pop().context("No search response")?;
        let MetaHit {
            meta,
            index: _,
            hit: last_hit,
        } = last_mh;
        let SearchRespMeta {
            found,
            out_of,
            page,
            request_params,
            search_time,
        } = meta;

        let mut hits = Vec::with_capacity(found);
        for mh in metahits {
            let Some(hit) = mh.hit else {
                // any hit being None means there aren't any hits,
                // just one item with metadata.
                break;
            };
            hits.push(hit);
        }

        if let Some(hit) = last_hit {
            hits.push(hit);
        }

        let resp = SearchResponse {
            found,
            out_of,
            page,
            request_params,
            search_time,
            hits,
        };
        Ok(resp)
    }

    /// Sends a requests for the Omry API to delete record(s) with
    /// the given IDs.
    #[instrument(skip(self))]
    async fn delete(&self, ids: Vec<RecordId>) -> Result<usize> {
        tracing::debug!(
            "Will make {} RPCs to delete the given IDs (one per each)",
            ids.len()
        );

        let (tx, mut rx) = self.api().delete_record().await?;
        let mut set = JoinSet::new();

        for id in ids {
            let tx = tx.clone();
            set.spawn(async move { tx.send(id).await });
        }

        drop(tx);

        set.join_all().await;

        let mut n_deleted = 0usize;
        let mut n_errors = 0usize;

        loop {
            match rx.recv().await {
                Ok(Some(num_affected)) => n_deleted += num_affected,
                Ok(None) => break,
                Err(e) => {
                    tracing::error!("{e}");
                    n_errors += 1;
                }
            }
        }

        if n_errors != 0 {
            anyhow::bail!("{n_errors} records couldn't be scheduled for deletion due to errors.");
        }

        Ok(n_deleted)
    }

    /// Sends a request to update records in Omry,
    /// and returns a list of IDs for records that were successfully updated.
    #[instrument(skip_all)]
    async fn update(&self, records: Vec<RecordParams>) -> Result<Vec<RecordId>> {
        let (tx, mut rx) = self
            .api()
            .update_record()
            .await
            .context("error invoking 'update_record' RPC")?;

        let mut set = JoinSet::new();
        for rec in records {
            let tx = tx.clone();
            set.spawn(async move { tx.send(rec).await });
        }

        drop(tx);
        set.join_all().await;

        let mut ids = vec![];
        while let Some(id) = rx.recv().await? {
            ids.push(id);
        }

        Ok(ids)
    }

    /// Returns the version information for the server application.
    #[instrument(skip(self))]
    async fn version(&self) -> Result<String> {
        let rx = self
            .api()
            .version()
            .await
            .context("Failed to invoke 'version' RPC")?;

        Ok(rx.await?)
    }
}