LCOV - code coverage report
Current view: top level - libs/proxy/tokio-postgres2/src - query.rs (source / functions) Coverage Total Hit
Test: 07bee600374ccd486c69370d0972d9035964fe68.info Lines: 0.0 % 202 0
Test Date: 2025-02-20 13:11:02 Functions: 0.0 % 27 0

            Line data    Source code
       1              : use crate::client::{InnerClient, Responses};
       2              : use crate::codec::FrontendMessage;
       3              : use crate::connection::RequestMessages;
       4              : use crate::types::IsNull;
       5              : use crate::{Column, Error, ReadyForQueryStatus, Row, Statement};
       6              : use bytes::{BufMut, Bytes, BytesMut};
       7              : use fallible_iterator::FallibleIterator;
       8              : use futures_util::{ready, Stream};
       9              : use log::{debug, log_enabled, Level};
      10              : use pin_project_lite::pin_project;
      11              : use postgres_protocol2::message::backend::Message;
      12              : use postgres_protocol2::message::frontend;
      13              : use postgres_types2::{Format, ToSql, Type};
      14              : use std::fmt;
      15              : use std::marker::PhantomPinned;
      16              : use std::pin::Pin;
      17              : use std::sync::Arc;
      18              : use std::task::{Context, Poll};
      19              : 
      20              : struct BorrowToSqlParamsDebug<'a>(&'a [&'a (dyn ToSql + Sync)]);
      21              : 
      22              : impl fmt::Debug for BorrowToSqlParamsDebug<'_> {
      23            0 :     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
      24            0 :         f.debug_list().entries(self.0.iter()).finish()
      25            0 :     }
      26              : }
      27              : 
      28            0 : pub async fn query<'a, I>(
      29            0 :     client: &InnerClient,
      30            0 :     statement: Statement,
      31            0 :     params: I,
      32            0 : ) -> Result<RowStream, Error>
      33            0 : where
      34            0 :     I: IntoIterator<Item = &'a (dyn ToSql + Sync)>,
      35            0 :     I::IntoIter: ExactSizeIterator,
      36            0 : {
      37            0 :     let buf = if log_enabled!(Level::Debug) {
      38            0 :         let params = params.into_iter().collect::<Vec<_>>();
      39            0 :         debug!(
      40            0 :             "executing statement {} with parameters: {:?}",
      41            0 :             statement.name(),
      42            0 :             BorrowToSqlParamsDebug(params.as_slice()),
      43              :         );
      44            0 :         encode(client, &statement, params)?
      45              :     } else {
      46            0 :         encode(client, &statement, params)?
      47              :     };
      48            0 :     let responses = start(client, buf).await?;
      49            0 :     Ok(RowStream {
      50            0 :         statement,
      51            0 :         responses,
      52            0 :         command_tag: None,
      53            0 :         status: ReadyForQueryStatus::Unknown,
      54            0 :         output_format: Format::Binary,
      55            0 :         _p: PhantomPinned,
      56            0 :     })
      57            0 : }
      58              : 
      59            0 : pub async fn query_txt<S, I>(
      60            0 :     client: &Arc<InnerClient>,
      61            0 :     query: &str,
      62            0 :     params: I,
      63            0 : ) -> Result<RowStream, Error>
      64            0 : where
      65            0 :     S: AsRef<str>,
      66            0 :     I: IntoIterator<Item = Option<S>>,
      67            0 :     I::IntoIter: ExactSizeIterator,
      68            0 : {
      69            0 :     let params = params.into_iter();
      70              : 
      71            0 :     let buf = client.with_buf(|buf| {
      72            0 :         frontend::parse(
      73            0 :             "",                 // unnamed prepared statement
      74            0 :             query,              // query to parse
      75            0 :             std::iter::empty(), // give no type info
      76            0 :             buf,
      77            0 :         )
      78            0 :         .map_err(Error::encode)?;
      79            0 :         frontend::describe(b'S', "", buf).map_err(Error::encode)?;
      80              :         // Bind, pass params as text, retrieve as binary
      81              :         match frontend::bind(
      82            0 :             "",                 // empty string selects the unnamed portal
      83            0 :             "",                 // unnamed prepared statement
      84            0 :             std::iter::empty(), // all parameters use the default format (text)
      85            0 :             params,
      86            0 :             |param, buf| match param {
      87            0 :                 Some(param) => {
      88            0 :                     buf.put_slice(param.as_ref().as_bytes());
      89            0 :                     Ok(postgres_protocol2::IsNull::No)
      90              :                 }
      91            0 :                 None => Ok(postgres_protocol2::IsNull::Yes),
      92            0 :             },
      93            0 :             Some(0), // all text
      94            0 :             buf,
      95              :         ) {
      96            0 :             Ok(()) => Ok(()),
      97            0 :             Err(frontend::BindError::Conversion(e)) => Err(Error::to_sql(e, 0)),
      98            0 :             Err(frontend::BindError::Serialization(e)) => Err(Error::encode(e)),
      99            0 :         }?;
     100              : 
     101              :         // Execute
     102            0 :         frontend::execute("", 0, buf).map_err(Error::encode)?;
     103              :         // Sync
     104            0 :         frontend::sync(buf);
     105            0 : 
     106            0 :         Ok(buf.split().freeze())
     107            0 :     })?;
     108              : 
     109              :     // now read the responses
     110            0 :     let mut responses = client.send(RequestMessages::Single(FrontendMessage::Raw(buf)))?;
     111              : 
     112            0 :     match responses.next().await? {
     113            0 :         Message::ParseComplete => {}
     114            0 :         _ => return Err(Error::unexpected_message()),
     115              :     }
     116              : 
     117            0 :     let parameter_description = match responses.next().await? {
     118            0 :         Message::ParameterDescription(body) => body,
     119            0 :         _ => return Err(Error::unexpected_message()),
     120              :     };
     121              : 
     122            0 :     let row_description = match responses.next().await? {
     123            0 :         Message::RowDescription(body) => Some(body),
     124            0 :         Message::NoData => None,
     125            0 :         _ => return Err(Error::unexpected_message()),
     126              :     };
     127              : 
     128            0 :     match responses.next().await? {
     129            0 :         Message::BindComplete => {}
     130            0 :         _ => return Err(Error::unexpected_message()),
     131              :     }
     132              : 
     133            0 :     let mut parameters = vec![];
     134            0 :     let mut it = parameter_description.parameters();
     135            0 :     while let Some(oid) = it.next().map_err(Error::parse)? {
     136            0 :         let type_ = Type::from_oid(oid).unwrap_or(Type::UNKNOWN);
     137            0 :         parameters.push(type_);
     138            0 :     }
     139              : 
     140            0 :     let mut columns = vec![];
     141            0 :     if let Some(row_description) = row_description {
     142            0 :         let mut it = row_description.fields();
     143            0 :         while let Some(field) = it.next().map_err(Error::parse)? {
     144            0 :             let type_ = Type::from_oid(field.type_oid()).unwrap_or(Type::UNKNOWN);
     145            0 :             let column = Column::new(field.name().to_string(), type_, field);
     146            0 :             columns.push(column);
     147            0 :         }
     148            0 :     }
     149              : 
     150            0 :     Ok(RowStream {
     151            0 :         statement: Statement::new_anonymous(parameters, columns),
     152            0 :         responses,
     153            0 :         command_tag: None,
     154            0 :         status: ReadyForQueryStatus::Unknown,
     155            0 :         output_format: Format::Text,
     156            0 :         _p: PhantomPinned,
     157            0 :     })
     158            0 : }
     159              : 
     160            0 : async fn start(client: &InnerClient, buf: Bytes) -> Result<Responses, Error> {
     161            0 :     let mut responses = client.send(RequestMessages::Single(FrontendMessage::Raw(buf)))?;
     162              : 
     163            0 :     match responses.next().await? {
     164            0 :         Message::BindComplete => {}
     165            0 :         _ => return Err(Error::unexpected_message()),
     166              :     }
     167              : 
     168            0 :     Ok(responses)
     169            0 : }
     170              : 
     171            0 : pub fn encode<'a, I>(client: &InnerClient, statement: &Statement, params: I) -> Result<Bytes, Error>
     172            0 : where
     173            0 :     I: IntoIterator<Item = &'a (dyn ToSql + Sync)>,
     174            0 :     I::IntoIter: ExactSizeIterator,
     175            0 : {
     176            0 :     client.with_buf(|buf| {
     177            0 :         encode_bind(statement, params, "", buf)?;
     178            0 :         frontend::execute("", 0, buf).map_err(Error::encode)?;
     179            0 :         frontend::sync(buf);
     180            0 :         Ok(buf.split().freeze())
     181            0 :     })
     182            0 : }
     183              : 
     184            0 : pub fn encode_bind<'a, I>(
     185            0 :     statement: &Statement,
     186            0 :     params: I,
     187            0 :     portal: &str,
     188            0 :     buf: &mut BytesMut,
     189            0 : ) -> Result<(), Error>
     190            0 : where
     191            0 :     I: IntoIterator<Item = &'a (dyn ToSql + Sync)>,
     192            0 :     I::IntoIter: ExactSizeIterator,
     193            0 : {
     194            0 :     let param_types = statement.params();
     195            0 :     let params = params.into_iter();
     196            0 : 
     197            0 :     assert!(
     198            0 :         param_types.len() == params.len(),
     199            0 :         "expected {} parameters but got {}",
     200            0 :         param_types.len(),
     201            0 :         params.len()
     202              :     );
     203              : 
     204            0 :     let (param_formats, params): (Vec<_>, Vec<_>) = params
     205            0 :         .zip(param_types.iter())
     206            0 :         .map(|(p, ty)| (p.encode_format(ty) as i16, p))
     207            0 :         .unzip();
     208            0 : 
     209            0 :     let params = params.into_iter();
     210            0 : 
     211            0 :     let mut error_idx = 0;
     212            0 :     let r = frontend::bind(
     213            0 :         portal,
     214            0 :         statement.name(),
     215            0 :         param_formats,
     216            0 :         params.zip(param_types).enumerate(),
     217            0 :         |(idx, (param, ty)), buf| match param.to_sql_checked(ty, buf) {
     218            0 :             Ok(IsNull::No) => Ok(postgres_protocol2::IsNull::No),
     219            0 :             Ok(IsNull::Yes) => Ok(postgres_protocol2::IsNull::Yes),
     220            0 :             Err(e) => {
     221            0 :                 error_idx = idx;
     222            0 :                 Err(e)
     223              :             }
     224            0 :         },
     225            0 :         Some(1),
     226            0 :         buf,
     227            0 :     );
     228            0 :     match r {
     229            0 :         Ok(()) => Ok(()),
     230            0 :         Err(frontend::BindError::Conversion(e)) => Err(Error::to_sql(e, error_idx)),
     231            0 :         Err(frontend::BindError::Serialization(e)) => Err(Error::encode(e)),
     232              :     }
     233            0 : }
     234              : 
     235              : pin_project! {
     236              :     /// A stream of table rows.
     237              :     pub struct RowStream {
     238              :         statement: Statement,
     239              :         responses: Responses,
     240              :         command_tag: Option<String>,
     241              :         output_format: Format,
     242              :         status: ReadyForQueryStatus,
     243              :         #[pin]
     244              :         _p: PhantomPinned,
     245              :     }
     246              : }
     247              : 
     248              : impl Stream for RowStream {
     249              :     type Item = Result<Row, Error>;
     250              : 
     251            0 :     fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
     252            0 :         let this = self.project();
     253              :         loop {
     254            0 :             match ready!(this.responses.poll_next(cx)?) {
     255            0 :                 Message::DataRow(body) => {
     256            0 :                     return Poll::Ready(Some(Ok(Row::new(
     257            0 :                         this.statement.clone(),
     258            0 :                         body,
     259            0 :                         *this.output_format,
     260            0 :                     )?)))
     261              :                 }
     262            0 :                 Message::EmptyQueryResponse | Message::PortalSuspended => {}
     263            0 :                 Message::CommandComplete(body) => {
     264            0 :                     if let Ok(tag) = body.tag() {
     265            0 :                         *this.command_tag = Some(tag.to_string());
     266            0 :                     }
     267              :                 }
     268            0 :                 Message::ReadyForQuery(status) => {
     269            0 :                     *this.status = status.into();
     270            0 :                     return Poll::Ready(None);
     271              :                 }
     272            0 :                 _ => return Poll::Ready(Some(Err(Error::unexpected_message()))),
     273              :             }
     274              :         }
     275            0 :     }
     276              : }
     277              : 
     278              : impl RowStream {
     279              :     /// Returns information about the columns of data in the row.
     280            0 :     pub fn columns(&self) -> &[Column] {
     281            0 :         self.statement.columns()
     282            0 :     }
     283              : 
     284              :     /// Returns the command tag of this query.
     285              :     ///
     286              :     /// This is only available after the stream has been exhausted.
     287            0 :     pub fn command_tag(&self) -> Option<String> {
     288            0 :         self.command_tag.clone()
     289            0 :     }
     290              : 
     291              :     /// Returns if the connection is ready for querying, with the status of the connection.
     292              :     ///
     293              :     /// This might be available only after the stream has been exhausted.
     294            0 :     pub fn ready_status(&self) -> ReadyForQueryStatus {
     295            0 :         self.status
     296            0 :     }
     297              : }
        

Generated by: LCOV version 2.1-beta