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