picodata_driver/
client.rs

1use std::future::IntoFuture;
2use std::marker::PhantomData;
3use std::pin::Pin;
4use std::sync::Arc;
5
6use crate::{
7    config::Config,
8    error::Error,
9    manager::Manager,
10    strategy::{RoundRobin, Strategy},
11    topology::view::Topology,
12};
13
14/// Client which is used to query picodata.
15#[derive(Clone)]
16pub struct Client {
17    manager: Arc<Manager>,
18}
19
20impl Client {
21    /// Creates a new instance of picodata client.
22    ///
23    /// # Errors
24    ///
25    /// Returns [`Error`] if manager initialization fails.
26    pub async fn new(config: Config) -> Result<Self, Error> {
27        let manager = Arc::new(Manager::new(config, RoundRobin::new()).await?);
28        Ok(Self { manager })
29    }
30
31    /// Creates a new instance of picodata client with a custom strategy.
32    ///
33    /// # Errors
34    ///
35    /// Returns [`Error`] if manager initialization fails.
36    pub async fn new_with_strategy<S>(config: Config, strategy: S) -> Result<Self, Error>
37    where
38        S: Strategy,
39    {
40        let manager = Arc::new(Manager::new(config, strategy).await?);
41        Ok(Self { manager })
42    }
43
44    /// Query picodata cluster. Returns a builder that supports
45    /// `.on_master(bool)`, `.bucket_id(u64)` and `.instance(&str)` before `.await`.
46    pub fn query<'a>(
47        &'a self,
48        sql: &'a str,
49        params: &'a [&'a (dyn pg::types::ToSql + Sync)],
50    ) -> Query<'a, Vec<pg::Row>> {
51        Query::new(&self.manager, sql, params)
52    }
53
54    /// Execute on picodata cluster. Returns a builder that supports
55    /// `.on_master(bool)`, `.bucket_id(u64)` and `.instance(&str)` before `.await`.
56    pub fn execute<'a>(
57        &'a self,
58        sql: &'a str,
59        params: &'a [&'a (dyn pg::types::ToSql + Sync)],
60    ) -> Query<'a, u64> {
61        Query::new(&self.manager, sql, params)
62    }
63
64    /// Returns a snapshot of the cluster topology. [`Tier::bucket_map`](crate::Tier::bucket_map)
65    /// checks a tier's buckets.
66    ///
67    /// # Errors
68    ///
69    /// Returns [`Error::Closed`] if the client has been closed.
70    pub async fn topology(&self) -> Result<Topology, Error> {
71        self.manager.topology().await
72    }
73
74    /// Ping each instance of picodata cluster.
75    ///
76    /// # Errors
77    ///
78    /// Returns [`Error`] if any instance fails to respond.
79    pub async fn ping(&self) -> Result<(), Error> {
80        self.manager
81            .map(|pool| async move {
82                let client = pool.get().await?;
83                client.simple_query("SELECT 1").await?;
84                Ok(())
85            })
86            .await?;
87        Ok(())
88    }
89
90    /// Closes pool manager.
91    ///
92    /// # Errors
93    ///
94    /// Returns [`Error`] if the manager is already closed.
95    pub async fn close(&self) -> Result<(), Error> {
96        self.manager.close().await
97    }
98}
99
100/// Request builder returned by [`Client::query`] and [`Client::execute`].
101///
102/// Supports `.on_master(bool)`, `.bucket_id(u64)` and `.instance(&str)` before `.await`.
103pub struct Query<'a, T> {
104    manager: &'a Manager,
105    sql: &'a str,
106    params: &'a [&'a (dyn pg::types::ToSql + Sync)],
107    on_master: bool,
108    bucket_id: Option<u64>,
109    instance: Option<&'a str>,
110    marker: PhantomData<T>,
111}
112
113impl<'a, T> Query<'a, T> {
114    /// Creates a new query builder with default settings (`on_master=true`, no forced `bucket_id`).
115    fn new(
116        manager: &'a Arc<Manager>,
117        sql: &'a str,
118        params: &'a [&'a (dyn pg::types::ToSql + Sync)],
119    ) -> Self {
120        Self {
121            manager,
122            sql,
123            params,
124            on_master: true,
125            bucket_id: None,
126            instance: None,
127            marker: PhantomData,
128        }
129    }
130
131    /// Routes the query to master or replica. Defaults to `true` (master).
132    #[must_use]
133    pub fn on_master(self, on_master: bool) -> RoutedQuery<'a, T> {
134        RoutedQuery(self).on_master(on_master)
135    }
136
137    /// Forces routing to a specific `bucket_id`, bypassing automatic computation.
138    #[must_use]
139    pub fn bucket_id(self, bucket_id: u64) -> RoutedQuery<'a, T> {
140        RoutedQuery(self).bucket_id(bucket_id)
141    }
142
143    /// Pins the query to the named instance. Fails if it is unknown or offline.
144    #[must_use]
145    pub fn instance(mut self, name: &'a str) -> PinnedQuery<'a, T> {
146        self.instance = Some(name);
147        PinnedQuery(self)
148    }
149
150    /// The pool of the pinned instance, or the one routing picks.
151    async fn pool(&self) -> Result<Arc<pgpool::Pool>, Error> {
152        match self.instance {
153            Some(name) => self.manager.route_to(name).await,
154            None => {
155                self.manager
156                    .route(self.sql, self.params, self.on_master, self.bucket_id)
157                    .await
158            }
159        }
160    }
161}
162
163/// Request routed by `.on_master` or `.bucket_id`, returned by [`Query::on_master`] and
164/// [`Query::bucket_id`]. It takes no `.instance`:
165///
166/// ```no_run
167/// # async fn f(c: &picodata_driver::Client) -> Result<(), picodata_driver::Error> {
168/// c.query("SELECT 1", &[]).bucket_id(1).on_master(false).await?;
169/// c.query("SELECT 1", &[]).on_master(false).bucket_id(1).await?;
170/// c.query("SELECT 1", &[]).instance("i1").await?;
171/// # Ok(()) }
172/// ```
173///
174/// ```compile_fail,E0599
175/// # async fn f(c: &picodata_driver::Client) {
176/// let _ = c.query("SELECT 1", &[]).bucket_id(1).instance("i1").await;
177/// # }
178/// ```
179///
180/// ```compile_fail,E0599
181/// # async fn f(c: &picodata_driver::Client) {
182/// let _ = c.query("SELECT 1", &[]).on_master(false).instance("i1").await;
183/// # }
184/// ```
185pub struct RoutedQuery<'a, T>(Query<'a, T>);
186
187impl<T> RoutedQuery<'_, T> {
188    /// Routes the query to master or replica. Defaults to `true` (master).
189    #[must_use]
190    pub fn on_master(mut self, on_master: bool) -> Self {
191        self.0.on_master = on_master;
192        self
193    }
194
195    /// Forces routing to a specific `bucket_id`, bypassing automatic computation.
196    #[must_use]
197    pub fn bucket_id(mut self, bucket_id: u64) -> Self {
198        self.0.bucket_id = Some(bucket_id);
199        self
200    }
201}
202
203impl<'a, T> IntoFuture for RoutedQuery<'a, T>
204where
205    Query<'a, T>: IntoFuture,
206{
207    type Output = <Query<'a, T> as IntoFuture>::Output;
208
209    type IntoFuture = <Query<'a, T> as IntoFuture>::IntoFuture;
210
211    fn into_future(self) -> Self::IntoFuture {
212        self.0.into_future()
213    }
214}
215
216/// Request pinned to one instance, returned by [`Query::instance`].
217/// It takes neither `.on_master` nor `.bucket_id`:
218///
219/// ```compile_fail,E0599
220/// # async fn f(c: &picodata_driver::Client) {
221/// let _ = c.query("SELECT 1", &[]).instance("i1").bucket_id(1).await;
222/// # }
223/// ```
224///
225/// ```compile_fail,E0599
226/// # async fn f(c: &picodata_driver::Client) {
227/// let _ = c.query("SELECT 1", &[]).instance("i1").on_master(false).await;
228/// # }
229/// ```
230pub struct PinnedQuery<'a, T>(Query<'a, T>);
231
232impl<'a, T> IntoFuture for PinnedQuery<'a, T>
233where
234    Query<'a, T>: IntoFuture,
235{
236    type Output = <Query<'a, T> as IntoFuture>::Output;
237
238    type IntoFuture = <Query<'a, T> as IntoFuture>::IntoFuture;
239
240    fn into_future(self) -> Self::IntoFuture {
241        self.0.into_future()
242    }
243}
244
245impl<'a> IntoFuture for Query<'a, Vec<pg::Row>> {
246    type Output = Result<Vec<pg::Row>, Error>;
247
248    type IntoFuture = Pin<Box<dyn Future<Output = Self::Output> + Send + 'a>>;
249
250    fn into_future(self) -> Self::IntoFuture {
251        Box::pin(async move {
252            let pool = self.pool().await?;
253            let client = pool.get().await?;
254            match client.query(self.sql, self.params).await {
255                Ok(rows) => Ok(rows),
256                Err(e) => {
257                    self.manager.on_error(self.sql, &e).await;
258                    Err(e.into())
259                }
260            }
261        })
262    }
263}
264
265impl<'a> IntoFuture for Query<'a, u64> {
266    type Output = Result<u64, Error>;
267
268    type IntoFuture = Pin<Box<dyn Future<Output = Self::Output> + Send + 'a>>;
269
270    fn into_future(self) -> Self::IntoFuture {
271        Box::pin(async move {
272            let pool = self.pool().await?;
273            let client = pool.get().await?;
274            match client.execute(self.sql, self.params).await {
275                Ok(n) => Ok(n),
276                Err(e) => {
277                    self.manager.on_error(self.sql, &e).await;
278                    Err(e.into())
279                }
280            }
281        })
282    }
283}
picodata_driver/client.rs

picodata_driver/
client.rs
