LCOV - code coverage report
Current view: top level - libs/proxy/tokio-postgres2/src - prepare.rs (source / functions) Coverage Total Hit
Test: 5fe7fa8d483b39476409aee736d6d5e32728bfac.info Lines: 0.0 % 153 0
Test Date: 2025-03-12 16:10:49 Functions: 0.0 % 20 0

            Line data    Source code
       1              : use std::future::Future;
       2              : use std::pin::Pin;
       3              : use std::sync::Arc;
       4              : 
       5              : use bytes::Bytes;
       6              : use fallible_iterator::FallibleIterator;
       7              : use futures_util::{TryStreamExt, pin_mut};
       8              : use log::debug;
       9              : use postgres_protocol2::message::backend::Message;
      10              : use postgres_protocol2::message::frontend;
      11              : 
      12              : use crate::client::InnerClient;
      13              : use crate::codec::FrontendMessage;
      14              : use crate::connection::RequestMessages;
      15              : use crate::types::{Field, Kind, Oid, Type};
      16              : use crate::{Column, Error, Statement, query, slice_iter};
      17              : 
      18              : pub(crate) const TYPEINFO_QUERY: &str = "\
      19              : SELECT t.typname, t.typtype, t.typelem, r.rngsubtype, t.typbasetype, n.nspname, t.typrelid
      20              : FROM pg_catalog.pg_type t
      21              : LEFT OUTER JOIN pg_catalog.pg_range r ON r.rngtypid = t.oid
      22              : INNER JOIN pg_catalog.pg_namespace n ON t.typnamespace = n.oid
      23              : WHERE t.oid = $1
      24              : ";
      25              : 
      26              : const TYPEINFO_ENUM_QUERY: &str = "\
      27              : SELECT enumlabel
      28              : FROM pg_catalog.pg_enum
      29              : WHERE enumtypid = $1
      30              : ORDER BY enumsortorder
      31              : ";
      32              : 
      33              : pub(crate) const TYPEINFO_COMPOSITE_QUERY: &str = "\
      34              : SELECT attname, atttypid
      35              : FROM pg_catalog.pg_attribute
      36              : WHERE attrelid = $1
      37              : AND NOT attisdropped
      38              : AND attnum > 0
      39              : ORDER BY attnum
      40              : ";
      41              : 
      42            0 : pub async fn prepare(
      43            0 :     client: &Arc<InnerClient>,
      44            0 :     name: &'static str,
      45            0 :     query: &str,
      46            0 :     types: &[Type],
      47            0 : ) -> Result<Statement, Error> {
      48            0 :     let buf = encode(client, name, query, types)?;
      49            0 :     let mut responses = client.send(RequestMessages::Single(FrontendMessage::Raw(buf)))?;
      50              : 
      51            0 :     match responses.next().await? {
      52            0 :         Message::ParseComplete => {}
      53            0 :         _ => return Err(Error::unexpected_message()),
      54              :     }
      55              : 
      56            0 :     let parameter_description = match responses.next().await? {
      57            0 :         Message::ParameterDescription(body) => body,
      58            0 :         _ => return Err(Error::unexpected_message()),
      59              :     };
      60              : 
      61            0 :     let row_description = match responses.next().await? {
      62            0 :         Message::RowDescription(body) => Some(body),
      63            0 :         Message::NoData => None,
      64            0 :         _ => return Err(Error::unexpected_message()),
      65              :     };
      66              : 
      67            0 :     let mut parameters = vec![];
      68            0 :     let mut it = parameter_description.parameters();
      69            0 :     while let Some(oid) = it.next().map_err(Error::parse)? {
      70            0 :         let type_ = get_type(client, oid).await?;
      71            0 :         parameters.push(type_);
      72              :     }
      73              : 
      74            0 :     let mut columns = vec![];
      75            0 :     if let Some(row_description) = row_description {
      76            0 :         let mut it = row_description.fields();
      77            0 :         while let Some(field) = it.next().map_err(Error::parse)? {
      78            0 :             let type_ = get_type(client, field.type_oid()).await?;
      79            0 :             let column = Column::new(field.name().to_string(), type_, field);
      80            0 :             columns.push(column);
      81              :         }
      82            0 :     }
      83              : 
      84            0 :     Ok(Statement::new(client, name, parameters, columns))
      85            0 : }
      86              : 
      87            0 : fn prepare_rec<'a>(
      88            0 :     client: &'a Arc<InnerClient>,
      89            0 :     name: &'static str,
      90            0 :     query: &'a str,
      91            0 :     types: &'a [Type],
      92            0 : ) -> Pin<Box<dyn Future<Output = Result<Statement, Error>> + 'a + Send>> {
      93            0 :     Box::pin(prepare(client, name, query, types))
      94            0 : }
      95              : 
      96            0 : fn encode(client: &InnerClient, name: &str, query: &str, types: &[Type]) -> Result<Bytes, Error> {
      97            0 :     if types.is_empty() {
      98            0 :         debug!("preparing query {}: {}", name, query);
      99              :     } else {
     100            0 :         debug!("preparing query {} with types {:?}: {}", name, types, query);
     101              :     }
     102              : 
     103            0 :     client.with_buf(|buf| {
     104            0 :         frontend::parse(name, query, types.iter().map(Type::oid), buf).map_err(Error::encode)?;
     105            0 :         frontend::describe(b'S', name, buf).map_err(Error::encode)?;
     106            0 :         frontend::sync(buf);
     107            0 :         Ok(buf.split().freeze())
     108            0 :     })
     109            0 : }
     110              : 
     111            0 : pub async fn get_type(client: &Arc<InnerClient>, oid: Oid) -> Result<Type, Error> {
     112            0 :     if let Some(type_) = Type::from_oid(oid) {
     113            0 :         return Ok(type_);
     114            0 :     }
     115              : 
     116            0 :     if let Some(type_) = client.type_(oid) {
     117            0 :         return Ok(type_);
     118            0 :     }
     119              : 
     120            0 :     let stmt = typeinfo_statement(client).await?;
     121              : 
     122            0 :     let rows = query::query(client, stmt, slice_iter(&[&oid])).await?;
     123            0 :     pin_mut!(rows);
     124              : 
     125            0 :     let row = match rows.try_next().await? {
     126            0 :         Some(row) => row,
     127            0 :         None => return Err(Error::unexpected_message()),
     128              :     };
     129              : 
     130            0 :     let name: String = row.try_get(0)?;
     131            0 :     let type_: i8 = row.try_get(1)?;
     132            0 :     let elem_oid: Oid = row.try_get(2)?;
     133            0 :     let rngsubtype: Option<Oid> = row.try_get(3)?;
     134            0 :     let basetype: Oid = row.try_get(4)?;
     135            0 :     let schema: String = row.try_get(5)?;
     136            0 :     let relid: Oid = row.try_get(6)?;
     137              : 
     138            0 :     let kind = if type_ == b'e' as i8 {
     139            0 :         let variants = get_enum_variants(client, oid).await?;
     140            0 :         Kind::Enum(variants)
     141            0 :     } else if type_ == b'p' as i8 {
     142            0 :         Kind::Pseudo
     143            0 :     } else if basetype != 0 {
     144            0 :         let type_ = get_type_rec(client, basetype).await?;
     145            0 :         Kind::Domain(type_)
     146            0 :     } else if elem_oid != 0 {
     147            0 :         let type_ = get_type_rec(client, elem_oid).await?;
     148            0 :         Kind::Array(type_)
     149            0 :     } else if relid != 0 {
     150            0 :         let fields = get_composite_fields(client, relid).await?;
     151            0 :         Kind::Composite(fields)
     152            0 :     } else if let Some(rngsubtype) = rngsubtype {
     153            0 :         let type_ = get_type_rec(client, rngsubtype).await?;
     154            0 :         Kind::Range(type_)
     155              :     } else {
     156            0 :         Kind::Simple
     157              :     };
     158              : 
     159            0 :     let type_ = Type::new(name, oid, kind, schema);
     160            0 :     client.set_type(oid, &type_);
     161            0 : 
     162            0 :     Ok(type_)
     163            0 : }
     164              : 
     165            0 : fn get_type_rec<'a>(
     166            0 :     client: &'a Arc<InnerClient>,
     167            0 :     oid: Oid,
     168            0 : ) -> Pin<Box<dyn Future<Output = Result<Type, Error>> + Send + 'a>> {
     169            0 :     Box::pin(get_type(client, oid))
     170            0 : }
     171              : 
     172            0 : async fn typeinfo_statement(client: &Arc<InnerClient>) -> Result<Statement, Error> {
     173            0 :     if let Some(stmt) = client.typeinfo() {
     174            0 :         return Ok(stmt);
     175            0 :     }
     176            0 : 
     177            0 :     let typeinfo = "neon_proxy_typeinfo";
     178            0 :     let stmt = prepare_rec(client, typeinfo, TYPEINFO_QUERY, &[]).await?;
     179              : 
     180            0 :     client.set_typeinfo(&stmt);
     181            0 :     Ok(stmt)
     182            0 : }
     183              : 
     184            0 : async fn get_enum_variants(client: &Arc<InnerClient>, oid: Oid) -> Result<Vec<String>, Error> {
     185            0 :     let stmt = typeinfo_enum_statement(client).await?;
     186              : 
     187            0 :     query::query(client, stmt, slice_iter(&[&oid]))
     188            0 :         .await?
     189            0 :         .and_then(|row| async move { row.try_get(0) })
     190            0 :         .try_collect()
     191            0 :         .await
     192            0 : }
     193              : 
     194            0 : async fn typeinfo_enum_statement(client: &Arc<InnerClient>) -> Result<Statement, Error> {
     195            0 :     if let Some(stmt) = client.typeinfo_enum() {
     196            0 :         return Ok(stmt);
     197            0 :     }
     198            0 : 
     199            0 :     let typeinfo = "neon_proxy_typeinfo_enum";
     200            0 :     let stmt = prepare_rec(client, typeinfo, TYPEINFO_ENUM_QUERY, &[]).await?;
     201              : 
     202            0 :     client.set_typeinfo_enum(&stmt);
     203            0 :     Ok(stmt)
     204            0 : }
     205              : 
     206            0 : async fn get_composite_fields(client: &Arc<InnerClient>, oid: Oid) -> Result<Vec<Field>, Error> {
     207            0 :     let stmt = typeinfo_composite_statement(client).await?;
     208              : 
     209            0 :     let rows = query::query(client, stmt, slice_iter(&[&oid]))
     210            0 :         .await?
     211            0 :         .try_collect::<Vec<_>>()
     212            0 :         .await?;
     213              : 
     214            0 :     let mut fields = vec![];
     215            0 :     for row in rows {
     216            0 :         let name = row.try_get(0)?;
     217            0 :         let oid = row.try_get(1)?;
     218            0 :         let type_ = get_type_rec(client, oid).await?;
     219            0 :         fields.push(Field::new(name, type_));
     220              :     }
     221              : 
     222            0 :     Ok(fields)
     223            0 : }
     224              : 
     225            0 : async fn typeinfo_composite_statement(client: &Arc<InnerClient>) -> Result<Statement, Error> {
     226            0 :     if let Some(stmt) = client.typeinfo_composite() {
     227            0 :         return Ok(stmt);
     228            0 :     }
     229            0 : 
     230            0 :     let typeinfo = "neon_proxy_typeinfo_composite";
     231            0 :     let stmt = prepare_rec(client, typeinfo, TYPEINFO_COMPOSITE_QUERY, &[]).await?;
     232              : 
     233            0 :     client.set_typeinfo_composite(&stmt);
     234            0 :     Ok(stmt)
     235            0 : }
        

Generated by: LCOV version 2.1-beta