Line data Source code
1 : use crate::codec::{BackendMessages, FrontendMessage};
2 :
3 : use crate::config::Host;
4 : use crate::config::SslMode;
5 : use crate::connection::{Request, RequestMessages};
6 :
7 : use crate::query::RowStream;
8 : use crate::simple_query::SimpleQueryStream;
9 :
10 : use crate::types::{Oid, ToSql, Type};
11 :
12 : use crate::{
13 : prepare, query, simple_query, slice_iter, CancelToken, Error, ReadyForQueryStatus, Row,
14 : SimpleQueryMessage, Statement, ToStatement, Transaction, TransactionBuilder,
15 : };
16 : use bytes::BytesMut;
17 : use fallible_iterator::FallibleIterator;
18 : use futures_util::{future, ready, TryStreamExt};
19 : use parking_lot::Mutex;
20 : use postgres_protocol2::message::{backend::Message, frontend};
21 : use std::collections::HashMap;
22 : use std::fmt;
23 : use std::sync::Arc;
24 : use std::task::{Context, Poll};
25 : use tokio::sync::mpsc;
26 :
27 : use std::time::Duration;
28 :
29 : pub struct Responses {
30 : receiver: mpsc::Receiver<BackendMessages>,
31 : cur: BackendMessages,
32 : }
33 :
34 : impl Responses {
35 0 : pub fn poll_next(&mut self, cx: &mut Context<'_>) -> Poll<Result<Message, Error>> {
36 : loop {
37 0 : match self.cur.next().map_err(Error::parse)? {
38 0 : Some(Message::ErrorResponse(body)) => return Poll::Ready(Err(Error::db(body))),
39 0 : Some(message) => return Poll::Ready(Ok(message)),
40 0 : None => {}
41 : }
42 :
43 0 : match ready!(self.receiver.poll_recv(cx)) {
44 0 : Some(messages) => self.cur = messages,
45 0 : None => return Poll::Ready(Err(Error::closed())),
46 : }
47 : }
48 0 : }
49 :
50 0 : pub async fn next(&mut self) -> Result<Message, Error> {
51 0 : future::poll_fn(|cx| self.poll_next(cx)).await
52 0 : }
53 : }
54 :
55 : /// A cache of type info and prepared statements for fetching type info
56 : /// (corresponding to the queries in the [prepare] module).
57 : #[derive(Default)]
58 : struct CachedTypeInfo {
59 : /// A statement for basic information for a type from its
60 : /// OID. Corresponds to [TYPEINFO_QUERY](prepare::TYPEINFO_QUERY) (or its
61 : /// fallback).
62 : typeinfo: Option<Statement>,
63 : /// A statement for getting information for a composite type from its OID.
64 : /// Corresponds to [TYPEINFO_QUERY](prepare::TYPEINFO_COMPOSITE_QUERY).
65 : typeinfo_composite: Option<Statement>,
66 : /// A statement for getting information for a composite type from its OID.
67 : /// Corresponds to [TYPEINFO_QUERY](prepare::TYPEINFO_COMPOSITE_QUERY) (or
68 : /// its fallback).
69 : typeinfo_enum: Option<Statement>,
70 :
71 : /// Cache of types already looked up.
72 : types: HashMap<Oid, Type>,
73 : }
74 :
75 : pub struct InnerClient {
76 : sender: mpsc::UnboundedSender<Request>,
77 : cached_typeinfo: Mutex<CachedTypeInfo>,
78 :
79 : /// A buffer to use when writing out postgres commands.
80 : buffer: Mutex<BytesMut>,
81 : }
82 :
83 : impl InnerClient {
84 0 : pub fn send(&self, messages: RequestMessages) -> Result<Responses, Error> {
85 0 : let (sender, receiver) = mpsc::channel(1);
86 0 : let request = Request { messages, sender };
87 0 : self.sender.send(request).map_err(|_| Error::closed())?;
88 :
89 0 : Ok(Responses {
90 0 : receiver,
91 0 : cur: BackendMessages::empty(),
92 0 : })
93 0 : }
94 :
95 0 : pub fn typeinfo(&self) -> Option<Statement> {
96 0 : self.cached_typeinfo.lock().typeinfo.clone()
97 0 : }
98 :
99 0 : pub fn set_typeinfo(&self, statement: &Statement) {
100 0 : self.cached_typeinfo.lock().typeinfo = Some(statement.clone());
101 0 : }
102 :
103 0 : pub fn typeinfo_composite(&self) -> Option<Statement> {
104 0 : self.cached_typeinfo.lock().typeinfo_composite.clone()
105 0 : }
106 :
107 0 : pub fn set_typeinfo_composite(&self, statement: &Statement) {
108 0 : self.cached_typeinfo.lock().typeinfo_composite = Some(statement.clone());
109 0 : }
110 :
111 0 : pub fn typeinfo_enum(&self) -> Option<Statement> {
112 0 : self.cached_typeinfo.lock().typeinfo_enum.clone()
113 0 : }
114 :
115 0 : pub fn set_typeinfo_enum(&self, statement: &Statement) {
116 0 : self.cached_typeinfo.lock().typeinfo_enum = Some(statement.clone());
117 0 : }
118 :
119 0 : pub fn type_(&self, oid: Oid) -> Option<Type> {
120 0 : self.cached_typeinfo.lock().types.get(&oid).cloned()
121 0 : }
122 :
123 0 : pub fn set_type(&self, oid: Oid, type_: &Type) {
124 0 : self.cached_typeinfo.lock().types.insert(oid, type_.clone());
125 0 : }
126 :
127 : /// Call the given function with a buffer to be used when writing out
128 : /// postgres commands.
129 0 : pub fn with_buf<F, R>(&self, f: F) -> R
130 0 : where
131 0 : F: FnOnce(&mut BytesMut) -> R,
132 0 : {
133 0 : let mut buffer = self.buffer.lock();
134 0 : let r = f(&mut buffer);
135 0 : buffer.clear();
136 0 : r
137 0 : }
138 : }
139 :
140 : #[derive(Clone)]
141 : pub struct SocketConfig {
142 : pub host: Host,
143 : pub port: u16,
144 : pub connect_timeout: Option<Duration>,
145 : // pub keepalive: Option<KeepaliveConfig>,
146 : }
147 :
148 : /// An asynchronous PostgreSQL client.
149 : ///
150 : /// The client is one half of what is returned when a connection is established. Users interact with the database
151 : /// through this client object.
152 : pub struct Client {
153 : inner: Arc<InnerClient>,
154 :
155 : socket_config: SocketConfig,
156 : ssl_mode: SslMode,
157 : process_id: i32,
158 : secret_key: i32,
159 : }
160 :
161 : impl Client {
162 0 : pub(crate) fn new(
163 0 : sender: mpsc::UnboundedSender<Request>,
164 0 : socket_config: SocketConfig,
165 0 : ssl_mode: SslMode,
166 0 : process_id: i32,
167 0 : secret_key: i32,
168 0 : ) -> Client {
169 0 : Client {
170 0 : inner: Arc::new(InnerClient {
171 0 : sender,
172 0 : cached_typeinfo: Default::default(),
173 0 : buffer: Default::default(),
174 0 : }),
175 0 :
176 0 : socket_config,
177 0 : ssl_mode,
178 0 : process_id,
179 0 : secret_key,
180 0 : }
181 0 : }
182 :
183 : /// Returns process_id.
184 0 : pub fn get_process_id(&self) -> i32 {
185 0 : self.process_id
186 0 : }
187 :
188 0 : pub(crate) fn inner(&self) -> &Arc<InnerClient> {
189 0 : &self.inner
190 0 : }
191 :
192 : /// Creates a new prepared statement.
193 : ///
194 : /// Prepared statements can be executed repeatedly, and may contain query parameters (indicated by `$1`, `$2`, etc),
195 : /// which are set when executed. Prepared statements can only be used with the connection that created them.
196 0 : pub async fn prepare(&self, query: &str) -> Result<Statement, Error> {
197 0 : self.prepare_typed(query, &[]).await
198 0 : }
199 :
200 : /// Like `prepare`, but allows the types of query parameters to be explicitly specified.
201 : ///
202 : /// The list of types may be smaller than the number of parameters - the types of the remaining parameters will be
203 : /// inferred. For example, `client.prepare_typed(query, &[])` is equivalent to `client.prepare(query)`.
204 0 : pub async fn prepare_typed(
205 0 : &self,
206 0 : query: &str,
207 0 : parameter_types: &[Type],
208 0 : ) -> Result<Statement, Error> {
209 0 : prepare::prepare(&self.inner, query, parameter_types).await
210 0 : }
211 :
212 : /// Executes a statement, returning a vector of the resulting rows.
213 : ///
214 : /// A statement may contain parameters, specified by `$n`, where `n` is the index of the parameter of the list
215 : /// provided, 1-indexed.
216 : ///
217 : /// The `statement` argument can either be a `Statement`, or a raw query string. If the same statement will be
218 : /// repeatedly executed (perhaps with different query parameters), consider preparing the statement up front
219 : /// with the `prepare` method.
220 : ///
221 : /// # Panics
222 : ///
223 : /// Panics if the number of parameters provided does not match the number expected.
224 0 : pub async fn query<T>(
225 0 : &self,
226 0 : statement: &T,
227 0 : params: &[&(dyn ToSql + Sync)],
228 0 : ) -> Result<Vec<Row>, Error>
229 0 : where
230 0 : T: ?Sized + ToStatement,
231 0 : {
232 0 : self.query_raw(statement, slice_iter(params))
233 0 : .await?
234 0 : .try_collect()
235 0 : .await
236 0 : }
237 :
238 : /// The maximally flexible version of [`query`].
239 : ///
240 : /// A statement may contain parameters, specified by `$n`, where `n` is the index of the parameter of the list
241 : /// provided, 1-indexed.
242 : ///
243 : /// The `statement` argument can either be a `Statement`, or a raw query string. If the same statement will be
244 : /// repeatedly executed (perhaps with different query parameters), consider preparing the statement up front
245 : /// with the `prepare` method.
246 : ///
247 : /// # Panics
248 : ///
249 : /// Panics if the number of parameters provided does not match the number expected.
250 : ///
251 : /// [`query`]: #method.query
252 0 : pub async fn query_raw<'a, T, I>(&self, statement: &T, params: I) -> Result<RowStream, Error>
253 0 : where
254 0 : T: ?Sized + ToStatement,
255 0 : I: IntoIterator<Item = &'a (dyn ToSql + Sync)>,
256 0 : I::IntoIter: ExactSizeIterator,
257 0 : {
258 0 : let statement = statement.__convert().into_statement(self).await?;
259 0 : query::query(&self.inner, statement, params).await
260 0 : }
261 :
262 : /// Pass text directly to the Postgres backend to allow it to sort out typing itself and
263 : /// to save a roundtrip
264 0 : pub async fn query_raw_txt<S, I>(&self, statement: &str, params: I) -> Result<RowStream, Error>
265 0 : where
266 0 : S: AsRef<str>,
267 0 : I: IntoIterator<Item = Option<S>>,
268 0 : I::IntoIter: ExactSizeIterator,
269 0 : {
270 0 : query::query_txt(&self.inner, statement, params).await
271 0 : }
272 :
273 : /// Executes a statement, returning the number of rows modified.
274 : ///
275 : /// A statement may contain parameters, specified by `$n`, where `n` is the index of the parameter of the list
276 : /// provided, 1-indexed.
277 : ///
278 : /// The `statement` argument can either be a `Statement`, or a raw query string. If the same statement will be
279 : /// repeatedly executed (perhaps with different query parameters), consider preparing the statement up front
280 : /// with the `prepare` method.
281 : ///
282 : /// If the statement does not modify any rows (e.g. `SELECT`), 0 is returned.
283 : ///
284 : /// # Panics
285 : ///
286 : /// Panics if the number of parameters provided does not match the number expected.
287 0 : pub async fn execute<T>(
288 0 : &self,
289 0 : statement: &T,
290 0 : params: &[&(dyn ToSql + Sync)],
291 0 : ) -> Result<u64, Error>
292 0 : where
293 0 : T: ?Sized + ToStatement,
294 0 : {
295 0 : self.execute_raw(statement, slice_iter(params)).await
296 0 : }
297 :
298 : /// The maximally flexible version of [`execute`].
299 : ///
300 : /// A statement may contain parameters, specified by `$n`, where `n` is the index of the parameter of the list
301 : /// provided, 1-indexed.
302 : ///
303 : /// The `statement` argument can either be a `Statement`, or a raw query string. If the same statement will be
304 : /// repeatedly executed (perhaps with different query parameters), consider preparing the statement up front
305 : /// with the `prepare` method.
306 : ///
307 : /// # Panics
308 : ///
309 : /// Panics if the number of parameters provided does not match the number expected.
310 : ///
311 : /// [`execute`]: #method.execute
312 0 : pub async fn execute_raw<'a, T, I>(&self, statement: &T, params: I) -> Result<u64, Error>
313 0 : where
314 0 : T: ?Sized + ToStatement,
315 0 : I: IntoIterator<Item = &'a (dyn ToSql + Sync)>,
316 0 : I::IntoIter: ExactSizeIterator,
317 0 : {
318 0 : let statement = statement.__convert().into_statement(self).await?;
319 0 : query::execute(self.inner(), statement, params).await
320 0 : }
321 :
322 : /// Executes a sequence of SQL statements using the simple query protocol, returning the resulting rows.
323 : ///
324 : /// Statements should be separated by semicolons. If an error occurs, execution of the sequence will stop at that
325 : /// point. The simple query protocol returns the values in rows as strings rather than in their binary encodings,
326 : /// so the associated row type doesn't work with the `FromSql` trait. Rather than simply returning a list of the
327 : /// rows, this method returns a list of an enum which indicates either the completion of one of the commands,
328 : /// or a row of data. This preserves the framing between the separate statements in the request.
329 : ///
330 : /// # Warning
331 : ///
332 : /// Prepared statements should be use for any query which contains user-specified data, as they provided the
333 : /// functionality to safely embed that data in the request. Do not form statements via string concatenation and pass
334 : /// them to this method!
335 0 : pub async fn simple_query(&self, query: &str) -> Result<Vec<SimpleQueryMessage>, Error> {
336 0 : self.simple_query_raw(query).await?.try_collect().await
337 0 : }
338 :
339 0 : pub(crate) async fn simple_query_raw(&self, query: &str) -> Result<SimpleQueryStream, Error> {
340 0 : simple_query::simple_query(self.inner(), query).await
341 0 : }
342 :
343 : /// Executes a sequence of SQL statements using the simple query protocol.
344 : ///
345 : /// Statements should be separated by semicolons. If an error occurs, execution of the sequence will stop at that
346 : /// point. This is intended for use when, for example, initializing a database schema.
347 : ///
348 : /// # Warning
349 : ///
350 : /// Prepared statements should be use for any query which contains user-specified data, as they provided the
351 : /// functionality to safely embed that data in the request. Do not form statements via string concatenation and pass
352 : /// them to this method!
353 0 : pub async fn batch_execute(&self, query: &str) -> Result<ReadyForQueryStatus, Error> {
354 0 : simple_query::batch_execute(self.inner(), query).await
355 0 : }
356 :
357 : /// Begins a new database transaction.
358 : ///
359 : /// The transaction will roll back by default - use the `commit` method to commit it.
360 0 : pub async fn transaction(&mut self) -> Result<Transaction<'_>, Error> {
361 : struct RollbackIfNotDone<'me> {
362 : client: &'me Client,
363 : done: bool,
364 : }
365 :
366 : impl Drop for RollbackIfNotDone<'_> {
367 0 : fn drop(&mut self) {
368 0 : if self.done {
369 0 : return;
370 0 : }
371 0 :
372 0 : let buf = self.client.inner().with_buf(|buf| {
373 0 : frontend::query("ROLLBACK", buf).unwrap();
374 0 : buf.split().freeze()
375 0 : });
376 0 : let _ = self
377 0 : .client
378 0 : .inner()
379 0 : .send(RequestMessages::Single(FrontendMessage::Raw(buf)));
380 0 : }
381 : }
382 :
383 : // This is done, as `Future` created by this method can be dropped after
384 : // `RequestMessages` is synchronously send to the `Connection` by
385 : // `batch_execute()`, but before `Responses` is asynchronously polled to
386 : // completion. In that case `Transaction` won't be created and thus
387 : // won't be rolled back.
388 : {
389 0 : let mut cleaner = RollbackIfNotDone {
390 0 : client: self,
391 0 : done: false,
392 0 : };
393 0 : self.batch_execute("BEGIN").await?;
394 0 : cleaner.done = true;
395 0 : }
396 0 :
397 0 : Ok(Transaction::new(self))
398 0 : }
399 :
400 : /// Returns a builder for a transaction with custom settings.
401 : ///
402 : /// Unlike the `transaction` method, the builder can be used to control the transaction's isolation level and other
403 : /// attributes.
404 0 : pub fn build_transaction(&mut self) -> TransactionBuilder<'_> {
405 0 : TransactionBuilder::new(self)
406 0 : }
407 :
408 : /// Constructs a cancellation token that can later be used to request cancellation of a query running on the
409 : /// connection associated with this client.
410 0 : pub fn cancel_token(&self) -> CancelToken {
411 0 : CancelToken {
412 0 : socket_config: Some(self.socket_config.clone()),
413 0 : ssl_mode: self.ssl_mode,
414 0 : process_id: self.process_id,
415 0 : secret_key: self.secret_key,
416 0 : }
417 0 : }
418 :
419 : /// Query for type information
420 0 : pub async fn get_type(&self, oid: Oid) -> Result<Type, Error> {
421 0 : crate::prepare::get_type(&self.inner, oid).await
422 0 : }
423 :
424 : /// Determines if the connection to the server has already closed.
425 : ///
426 : /// In that case, all future queries will fail.
427 0 : pub fn is_closed(&self) -> bool {
428 0 : self.inner.sender.is_closed()
429 0 : }
430 : }
431 :
432 : impl fmt::Debug for Client {
433 0 : fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
434 0 : f.debug_struct("Client").finish()
435 0 : }
436 : }
|
