picodata-driver 3.0.0

use std::future::IntoFuture;
use std::marker::PhantomData;
use std::pin::Pin;
use std::sync::Arc;

use crate::{
    config::Config,
    error::Error,
    manager::Manager,
    strategy::{RoundRobin, Strategy},
    topology::view::Topology,
};

/// Client which is used to query picodata.
#[derive(Clone)]
pub struct Client {
    manager: Arc<Manager>,
}

impl Client {
    /// Creates a new instance of picodata client.
    ///
    /// # Errors
    ///
    /// Returns [`Error`] if manager initialization fails.
    pub async fn new(config: Config) -> Result<Self, Error> {
        let manager = Arc::new(Manager::new(config, RoundRobin::new()).await?);
        Ok(Self { manager })
    }

    /// Creates a new instance of picodata client with a custom strategy.
    ///
    /// # Errors
    ///
    /// Returns [`Error`] if manager initialization fails.
    pub async fn new_with_strategy<S>(config: Config, strategy: S) -> Result<Self, Error>
    where
        S: Strategy,
    {
        let manager = Arc::new(Manager::new(config, strategy).await?);
        Ok(Self { manager })
    }

    /// Query picodata cluster. Returns a builder that supports
    /// `.on_master(bool)`, `.bucket_id(u64)` and `.instance(&str)` before `.await`.
    pub fn query<'a>(
        &'a self,
        sql: &'a str,
        params: &'a [&'a (dyn pg::types::ToSql + Sync)],
    ) -> Query<'a, Vec<pg::Row>> {
        Query::new(&self.manager, sql, params)
    }

    /// Execute on picodata cluster. Returns a builder that supports
    /// `.on_master(bool)`, `.bucket_id(u64)` and `.instance(&str)` before `.await`.
    pub fn execute<'a>(
        &'a self,
        sql: &'a str,
        params: &'a [&'a (dyn pg::types::ToSql + Sync)],
    ) -> Query<'a, u64> {
        Query::new(&self.manager, sql, params)
    }

    /// Returns a snapshot of the cluster topology. [`Tier::bucket_map`](crate::Tier::bucket_map)
    /// checks a tier's buckets.
    ///
    /// # Errors
    ///
    /// Returns [`Error::Closed`] if the client has been closed.
    pub async fn topology(&self) -> Result<Topology, Error> {
        self.manager.topology().await
    }

    /// Ping each instance of picodata cluster.
    ///
    /// # Errors
    ///
    /// Returns [`Error`] if any instance fails to respond.
    pub async fn ping(&self) -> Result<(), Error> {
        self.manager
            .map(|pool| async move {
                let client = pool.get().await?;
                client.simple_query("SELECT 1").await?;
                Ok(())
            })
            .await?;
        Ok(())
    }

    /// Closes pool manager.
    ///
    /// # Errors
    ///
    /// Returns [`Error`] if the manager is already closed.
    pub async fn close(&self) -> Result<(), Error> {
        self.manager.close().await
    }
}

/// Request builder returned by [`Client::query`] and [`Client::execute`].
///
/// Supports `.on_master(bool)`, `.bucket_id(u64)` and `.instance(&str)` before `.await`.
pub struct Query<'a, T> {
    manager: &'a Manager,
    sql: &'a str,
    params: &'a [&'a (dyn pg::types::ToSql + Sync)],
    on_master: bool,
    bucket_id: Option<u64>,
    instance: Option<&'a str>,
    marker: PhantomData<T>,
}

impl<'a, T> Query<'a, T> {
    /// Creates a new query builder with default settings (`on_master=true`, no forced `bucket_id`).
    fn new(
        manager: &'a Arc<Manager>,
        sql: &'a str,
        params: &'a [&'a (dyn pg::types::ToSql + Sync)],
    ) -> Self {
        Self {
            manager,
            sql,
            params,
            on_master: true,
            bucket_id: None,
            instance: None,
            marker: PhantomData,
        }
    }

    /// Routes the query to master or replica. Defaults to `true` (master).
    #[must_use]
    pub fn on_master(self, on_master: bool) -> RoutedQuery<'a, T> {
        RoutedQuery(self).on_master(on_master)
    }

    /// Forces routing to a specific `bucket_id`, bypassing automatic computation.
    #[must_use]
    pub fn bucket_id(self, bucket_id: u64) -> RoutedQuery<'a, T> {
        RoutedQuery(self).bucket_id(bucket_id)
    }

    /// Pins the query to the named instance. Fails if it is unknown or offline.
    #[must_use]
    pub fn instance(mut self, name: &'a str) -> PinnedQuery<'a, T> {
        self.instance = Some(name);
        PinnedQuery(self)
    }

    /// The pool of the pinned instance, or the one routing picks.
    async fn pool(&self) -> Result<Arc<pgpool::Pool>, Error> {
        match self.instance {
            Some(name) => self.manager.route_to(name).await,
            None => {
                self.manager
                    .route(self.sql, self.params, self.on_master, self.bucket_id)
                    .await
            }
        }
    }
}

/// Request routed by `.on_master` or `.bucket_id`, returned by [`Query::on_master`] and
/// [`Query::bucket_id`]. It takes no `.instance`:
///
/// ```no_run
/// # async fn f(c: &picodata_driver::Client) -> Result<(), picodata_driver::Error> {
/// c.query("SELECT 1", &[]).bucket_id(1).on_master(false).await?;
/// c.query("SELECT 1", &[]).on_master(false).bucket_id(1).await?;
/// c.query("SELECT 1", &[]).instance("i1").await?;
/// # Ok(()) }
/// ```
///
/// ```compile_fail,E0599
/// # async fn f(c: &picodata_driver::Client) {
/// let _ = c.query("SELECT 1", &[]).bucket_id(1).instance("i1").await;
/// # }
/// ```
///
/// ```compile_fail,E0599
/// # async fn f(c: &picodata_driver::Client) {
/// let _ = c.query("SELECT 1", &[]).on_master(false).instance("i1").await;
/// # }
/// ```
pub struct RoutedQuery<'a, T>(Query<'a, T>);

impl<T> RoutedQuery<'_, T> {
    /// Routes the query to master or replica. Defaults to `true` (master).
    #[must_use]
    pub fn on_master(mut self, on_master: bool) -> Self {
        self.0.on_master = on_master;
        self
    }

    /// Forces routing to a specific `bucket_id`, bypassing automatic computation.
    #[must_use]
    pub fn bucket_id(mut self, bucket_id: u64) -> Self {
        self.0.bucket_id = Some(bucket_id);
        self
    }
}

impl<'a, T> IntoFuture for RoutedQuery<'a, T>
where
    Query<'a, T>: IntoFuture,
{
    type Output = <Query<'a, T> as IntoFuture>::Output;

    type IntoFuture = <Query<'a, T> as IntoFuture>::IntoFuture;

    fn into_future(self) -> Self::IntoFuture {
        self.0.into_future()
    }
}

/// Request pinned to one instance, returned by [`Query::instance`].
/// It takes neither `.on_master` nor `.bucket_id`:
///
/// ```compile_fail,E0599
/// # async fn f(c: &picodata_driver::Client) {
/// let _ = c.query("SELECT 1", &[]).instance("i1").bucket_id(1).await;
/// # }
/// ```
///
/// ```compile_fail,E0599
/// # async fn f(c: &picodata_driver::Client) {
/// let _ = c.query("SELECT 1", &[]).instance("i1").on_master(false).await;
/// # }
/// ```
pub struct PinnedQuery<'a, T>(Query<'a, T>);

impl<'a, T> IntoFuture for PinnedQuery<'a, T>
where
    Query<'a, T>: IntoFuture,
{
    type Output = <Query<'a, T> as IntoFuture>::Output;

    type IntoFuture = <Query<'a, T> as IntoFuture>::IntoFuture;

    fn into_future(self) -> Self::IntoFuture {
        self.0.into_future()
    }
}

impl<'a> IntoFuture for Query<'a, Vec<pg::Row>> {
    type Output = Result<Vec<pg::Row>, Error>;

    type IntoFuture = Pin<Box<dyn Future<Output = Self::Output> + Send + 'a>>;

    fn into_future(self) -> Self::IntoFuture {
        Box::pin(async move {
            let pool = self.pool().await?;
            let client = pool.get().await?;
            match client.query(self.sql, self.params).await {
                Ok(rows) => Ok(rows),
                Err(e) => {
                    self.manager.on_error(self.sql, &e).await;
                    Err(e.into())
                }
            }
        })
    }
}

impl<'a> IntoFuture for Query<'a, u64> {
    type Output = Result<u64, Error>;

    type IntoFuture = Pin<Box<dyn Future<Output = Self::Output> + Send + 'a>>;

    fn into_future(self) -> Self::IntoFuture {
        Box::pin(async move {
            let pool = self.pool().await?;
            let client = pool.get().await?;
            match client.execute(self.sql, self.params).await {
                Ok(n) => Ok(n),
                Err(e) => {
                    self.manager.on_error(self.sql, &e).await;
                    Err(e.into())
                }
            }
        })
    }
}