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