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

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

Generated by: LCOV version 2.1-beta