diff --git a/libs/proxy/tokio-postgres2/src/client.rs b/libs/proxy/tokio-postgres2/src/client.rs index a7cd53afc3..90523121b9 100644 --- a/libs/proxy/tokio-postgres2/src/client.rs +++ b/libs/proxy/tokio-postgres2/src/client.rs @@ -5,22 +5,19 @@ use crate::config::SslMode; use crate::connection::{Request, RequestMessages}; use crate::query::RowStream; -use crate::simple_query::SimpleQueryStream; -use crate::types::{Oid, ToSql, Type}; +use crate::types::{Oid, Type}; use crate::{ - prepare, query, simple_query, slice_iter, CancelToken, Error, ReadyForQueryStatus, Row, - SimpleQueryMessage, Statement, ToStatement, Transaction, TransactionBuilder, + query, simple_query, CancelToken, Error, ReadyForQueryStatus, Statement, Transaction, + TransactionBuilder, }; use bytes::BytesMut; use fallible_iterator::FallibleIterator; -use futures_util::{future, ready, TryStreamExt}; -use parking_lot::Mutex; +use futures_util::{future, ready}; use postgres_protocol2::message::{backend::Message, frontend}; use std::collections::HashMap; use std::fmt; -use std::sync::Arc; use std::task::{Context, Poll}; use tokio::sync::mpsc; @@ -74,10 +71,10 @@ struct CachedTypeInfo { pub struct InnerClient { sender: mpsc::UnboundedSender, - cached_typeinfo: Mutex, + cached_typeinfo: CachedTypeInfo, /// A buffer to use when writing out postgres commands. - buffer: Mutex, + buffer: BytesMut, } impl InnerClient { @@ -92,47 +89,46 @@ impl InnerClient { }) } - pub fn typeinfo(&self) -> Option { - self.cached_typeinfo.lock().typeinfo.clone() + pub fn typeinfo(&mut self) -> Option { + self.cached_typeinfo.typeinfo.clone() } - pub fn set_typeinfo(&self, statement: &Statement) { - self.cached_typeinfo.lock().typeinfo = Some(statement.clone()); + pub fn set_typeinfo(&mut self, statement: &Statement) { + self.cached_typeinfo.typeinfo = Some(statement.clone()); } - pub fn typeinfo_composite(&self) -> Option { - self.cached_typeinfo.lock().typeinfo_composite.clone() + pub fn typeinfo_composite(&mut self) -> Option { + self.cached_typeinfo.typeinfo_composite.clone() } - pub fn set_typeinfo_composite(&self, statement: &Statement) { - self.cached_typeinfo.lock().typeinfo_composite = Some(statement.clone()); + pub fn set_typeinfo_composite(&mut self, statement: &Statement) { + self.cached_typeinfo.typeinfo_composite = Some(statement.clone()); } - pub fn typeinfo_enum(&self) -> Option { - self.cached_typeinfo.lock().typeinfo_enum.clone() + pub fn typeinfo_enum(&mut self) -> Option { + self.cached_typeinfo.typeinfo_enum.clone() } - pub fn set_typeinfo_enum(&self, statement: &Statement) { - self.cached_typeinfo.lock().typeinfo_enum = Some(statement.clone()); + pub fn set_typeinfo_enum(&mut self, statement: &Statement) { + self.cached_typeinfo.typeinfo_enum = Some(statement.clone()); } - pub fn type_(&self, oid: Oid) -> Option { - self.cached_typeinfo.lock().types.get(&oid).cloned() + pub fn type_(&mut self, oid: Oid) -> Option { + self.cached_typeinfo.types.get(&oid).cloned() } - pub fn set_type(&self, oid: Oid, type_: &Type) { - self.cached_typeinfo.lock().types.insert(oid, type_.clone()); + pub fn set_type(&mut self, oid: Oid, type_: &Type) { + self.cached_typeinfo.types.insert(oid, type_.clone()); } /// Call the given function with a buffer to be used when writing out /// postgres commands. - pub fn with_buf(&self, f: F) -> R + pub fn with_buf(&mut self, f: F) -> R where F: FnOnce(&mut BytesMut) -> R, { - let mut buffer = self.buffer.lock(); - let r = f(&mut buffer); - buffer.clear(); + let r = f(&mut self.buffer); + self.buffer.clear(); r } } @@ -150,7 +146,7 @@ pub struct SocketConfig { /// The client is one half of what is returned when a connection is established. Users interact with the database /// through this client object. pub struct Client { - inner: Arc, + inner: InnerClient, socket_config: SocketConfig, ssl_mode: SslMode, @@ -167,11 +163,11 @@ impl Client { secret_key: i32, ) -> Client { Client { - inner: Arc::new(InnerClient { + inner: InnerClient { sender, cached_typeinfo: Default::default(), buffer: Default::default(), - }), + }, socket_config, ssl_mode, @@ -185,159 +181,23 @@ impl Client { self.process_id } - pub(crate) fn inner(&self) -> &Arc { - &self.inner - } - - /// Creates a new prepared statement. - /// - /// Prepared statements can be executed repeatedly, and may contain query parameters (indicated by `$1`, `$2`, etc), - /// which are set when executed. Prepared statements can only be used with the connection that created them. - pub async fn prepare(&self, query: &str) -> Result { - self.prepare_typed(query, &[]).await - } - - /// Like `prepare`, but allows the types of query parameters to be explicitly specified. - /// - /// The list of types may be smaller than the number of parameters - the types of the remaining parameters will be - /// inferred. For example, `client.prepare_typed(query, &[])` is equivalent to `client.prepare(query)`. - pub async fn prepare_typed( - &self, - query: &str, - parameter_types: &[Type], - ) -> Result { - prepare::prepare(&self.inner, query, parameter_types).await - } - - /// Executes a statement, returning a vector of the resulting rows. - /// - /// A statement may contain parameters, specified by `$n`, where `n` is the index of the parameter of the list - /// provided, 1-indexed. - /// - /// The `statement` argument can either be a `Statement`, or a raw query string. If the same statement will be - /// repeatedly executed (perhaps with different query parameters), consider preparing the statement up front - /// with the `prepare` method. - /// - /// # Panics - /// - /// Panics if the number of parameters provided does not match the number expected. - pub async fn query( - &self, - statement: &T, - params: &[&(dyn ToSql + Sync)], - ) -> Result, Error> - where - T: ?Sized + ToStatement, - { - self.query_raw(statement, slice_iter(params)) - .await? - .try_collect() - .await - } - - /// The maximally flexible version of [`query`]. - /// - /// A statement may contain parameters, specified by `$n`, where `n` is the index of the parameter of the list - /// provided, 1-indexed. - /// - /// The `statement` argument can either be a `Statement`, or a raw query string. If the same statement will be - /// repeatedly executed (perhaps with different query parameters), consider preparing the statement up front - /// with the `prepare` method. - /// - /// # Panics - /// - /// Panics if the number of parameters provided does not match the number expected. - /// - /// [`query`]: #method.query - pub async fn query_raw<'a, T, I>(&self, statement: &T, params: I) -> Result - where - T: ?Sized + ToStatement, - I: IntoIterator, - I::IntoIter: ExactSizeIterator, - { - let statement = statement.__convert().into_statement(self).await?; - query::query(&self.inner, statement, params).await + pub(crate) fn inner(&mut self) -> &mut InnerClient { + &mut self.inner } /// Pass text directly to the Postgres backend to allow it to sort out typing itself and /// to save a roundtrip - pub async fn query_raw_txt(&self, statement: &str, params: I) -> Result + pub async fn query_raw_txt( + &mut self, + statement: &str, + params: I, + ) -> Result where S: AsRef, I: IntoIterator>, I::IntoIter: ExactSizeIterator, { - query::query_txt(&self.inner, statement, params).await - } - - /// Executes a statement, returning the number of rows modified. - /// - /// A statement may contain parameters, specified by `$n`, where `n` is the index of the parameter of the list - /// provided, 1-indexed. - /// - /// The `statement` argument can either be a `Statement`, or a raw query string. If the same statement will be - /// repeatedly executed (perhaps with different query parameters), consider preparing the statement up front - /// with the `prepare` method. - /// - /// If the statement does not modify any rows (e.g. `SELECT`), 0 is returned. - /// - /// # Panics - /// - /// Panics if the number of parameters provided does not match the number expected. - pub async fn execute( - &self, - statement: &T, - params: &[&(dyn ToSql + Sync)], - ) -> Result - where - T: ?Sized + ToStatement, - { - self.execute_raw(statement, slice_iter(params)).await - } - - /// The maximally flexible version of [`execute`]. - /// - /// A statement may contain parameters, specified by `$n`, where `n` is the index of the parameter of the list - /// provided, 1-indexed. - /// - /// The `statement` argument can either be a `Statement`, or a raw query string. If the same statement will be - /// repeatedly executed (perhaps with different query parameters), consider preparing the statement up front - /// with the `prepare` method. - /// - /// # Panics - /// - /// Panics if the number of parameters provided does not match the number expected. - /// - /// [`execute`]: #method.execute - pub async fn execute_raw<'a, T, I>(&self, statement: &T, params: I) -> Result - where - T: ?Sized + ToStatement, - I: IntoIterator, - I::IntoIter: ExactSizeIterator, - { - let statement = statement.__convert().into_statement(self).await?; - query::execute(self.inner(), statement, params).await - } - - /// Executes a sequence of SQL statements using the simple query protocol, returning the resulting rows. - /// - /// Statements should be separated by semicolons. If an error occurs, execution of the sequence will stop at that - /// point. The simple query protocol returns the values in rows as strings rather than in their binary encodings, - /// so the associated row type doesn't work with the `FromSql` trait. Rather than simply returning a list of the - /// rows, this method returns a list of an enum which indicates either the completion of one of the commands, - /// or a row of data. This preserves the framing between the separate statements in the request. - /// - /// # Warning - /// - /// Prepared statements should be use for any query which contains user-specified data, as they provided the - /// functionality to safely embed that data in the request. Do not form statements via string concatenation and pass - /// them to this method! - pub async fn simple_query(&self, query: &str) -> Result, Error> { - self.simple_query_raw(query).await?.try_collect().await - } - - pub(crate) async fn simple_query_raw(&self, query: &str) -> Result { - simple_query::simple_query(self.inner(), query).await + query::query_txt(&mut self.inner, statement, params).await } /// Executes a sequence of SQL statements using the simple query protocol. @@ -350,7 +210,7 @@ impl Client { /// Prepared statements should be use for any query which contains user-specified data, as they provided the /// functionality to safely embed that data in the request. Do not form statements via string concatenation and pass /// them to this method! - pub async fn batch_execute(&self, query: &str) -> Result { + pub async fn batch_execute(&mut self, query: &str) -> Result { simple_query::batch_execute(self.inner(), query).await } @@ -359,7 +219,7 @@ impl Client { /// The transaction will roll back by default - use the `commit` method to commit it. pub async fn transaction(&mut self) -> Result, Error> { struct RollbackIfNotDone<'me> { - client: &'me Client, + client: &'me mut Client, done: bool, } @@ -390,7 +250,7 @@ impl Client { client: self, done: false, }; - self.batch_execute("BEGIN").await?; + cleaner.client.batch_execute("BEGIN").await?; cleaner.done = true; } @@ -417,8 +277,8 @@ impl Client { } /// Query for type information - pub async fn get_type(&self, oid: Oid) -> Result { - crate::prepare::get_type(&self.inner, oid).await + pub async fn get_type(&mut self, oid: Oid) -> Result { + crate::prepare::get_type(&mut self.inner, oid).await } /// Determines if the connection to the server has already closed. diff --git a/libs/proxy/tokio-postgres2/src/generic_client.rs b/libs/proxy/tokio-postgres2/src/generic_client.rs index 768213f8ed..8f331fab26 100644 --- a/libs/proxy/tokio-postgres2/src/generic_client.rs +++ b/libs/proxy/tokio-postgres2/src/generic_client.rs @@ -21,7 +21,7 @@ pub trait GenericClient: private::Sealed { I::IntoIter: ExactSizeIterator + Sync + Send; /// Query for type information - async fn get_type(&self, oid: Oid) -> Result; + async fn get_type(&mut self, oid: Oid) -> Result; } impl private::Sealed for Client {} @@ -38,8 +38,8 @@ impl GenericClient for Client { } /// Query for type information - async fn get_type(&self, oid: Oid) -> Result { - self.get_type(oid).await + async fn get_type(&mut self, oid: Oid) -> Result { + crate::prepare::get_type(self.inner(), oid).await } } @@ -58,7 +58,7 @@ impl GenericClient for Transaction<'_> { } /// Query for type information - async fn get_type(&self, oid: Oid) -> Result { - self.client().get_type(oid).await + async fn get_type(&mut self, oid: Oid) -> Result { + crate::prepare::get_type(self.client().inner(), oid).await } } diff --git a/libs/proxy/tokio-postgres2/src/lib.rs b/libs/proxy/tokio-postgres2/src/lib.rs index 901ed0c96c..e2a57a23c5 100644 --- a/libs/proxy/tokio-postgres2/src/lib.rs +++ b/libs/proxy/tokio-postgres2/src/lib.rs @@ -14,7 +14,7 @@ pub use crate::row::{Row, SimpleQueryRow}; pub use crate::simple_query::SimpleQueryStream; pub use crate::statement::{Column, Statement}; pub use crate::tls::NoTls; -pub use crate::to_statement::ToStatement; +// pub use crate::to_statement::ToStatement; pub use crate::transaction::Transaction; pub use crate::transaction_builder::{IsolationLevel, TransactionBuilder}; use crate::types::ToSql; @@ -65,7 +65,7 @@ pub mod row; mod simple_query; mod statement; pub mod tls; -mod to_statement; +// mod to_statement; mod transaction; mod transaction_builder; pub mod types; diff --git a/libs/proxy/tokio-postgres2/src/prepare.rs b/libs/proxy/tokio-postgres2/src/prepare.rs index da0c755c5b..d5bb45fc45 100644 --- a/libs/proxy/tokio-postgres2/src/prepare.rs +++ b/libs/proxy/tokio-postgres2/src/prepare.rs @@ -14,7 +14,6 @@ use postgres_protocol2::message::frontend; use std::future::Future; use std::pin::Pin; use std::sync::atomic::{AtomicUsize, Ordering}; -use std::sync::Arc; pub(crate) const TYPEINFO_QUERY: &str = "\ SELECT t.typname, t.typtype, t.typelem, r.rngsubtype, t.typbasetype, n.nspname, t.typrelid @@ -59,7 +58,7 @@ ORDER BY attnum static NEXT_ID: AtomicUsize = AtomicUsize::new(0); pub async fn prepare( - client: &Arc, + client: &mut InnerClient, query: &str, types: &[Type], ) -> Result { @@ -100,18 +99,18 @@ pub async fn prepare( } } - Ok(Statement::new(client, name, parameters, columns)) + Ok(Statement::new(name, parameters, columns)) } fn prepare_rec<'a>( - client: &'a Arc, + client: &'a mut InnerClient, query: &'a str, types: &'a [Type], ) -> Pin> + 'a + Send>> { Box::pin(prepare(client, query, types)) } -fn encode(client: &InnerClient, name: &str, query: &str, types: &[Type]) -> Result { +fn encode(client: &mut InnerClient, name: &str, query: &str, types: &[Type]) -> Result { if types.is_empty() { debug!("preparing query {}: {}", name, query); } else { @@ -126,7 +125,7 @@ fn encode(client: &InnerClient, name: &str, query: &str, types: &[Type]) -> Resu }) } -pub async fn get_type(client: &Arc, oid: Oid) -> Result { +pub async fn get_type(client: &mut InnerClient, oid: Oid) -> Result { if let Some(type_) = Type::from_oid(oid) { return Ok(type_); } @@ -181,13 +180,13 @@ pub async fn get_type(client: &Arc, oid: Oid) -> Result( - client: &'a Arc, + client: &'a mut InnerClient, oid: Oid, ) -> Pin> + Send + 'a>> { Box::pin(get_type(client, oid)) } -async fn typeinfo_statement(client: &Arc) -> Result { +async fn typeinfo_statement(client: &mut InnerClient) -> Result { if let Some(stmt) = client.typeinfo() { return Ok(stmt); } @@ -204,7 +203,7 @@ async fn typeinfo_statement(client: &Arc) -> Result, oid: Oid) -> Result, Error> { +async fn get_enum_variants(client: &mut InnerClient, oid: Oid) -> Result, Error> { let stmt = typeinfo_enum_statement(client).await?; query::query(client, stmt, slice_iter(&[&oid])) @@ -214,7 +213,7 @@ async fn get_enum_variants(client: &Arc, oid: Oid) -> Result) -> Result { +async fn typeinfo_enum_statement(client: &mut InnerClient) -> Result { if let Some(stmt) = client.typeinfo_enum() { return Ok(stmt); } @@ -231,7 +230,7 @@ async fn typeinfo_enum_statement(client: &Arc) -> Result, oid: Oid) -> Result, Error> { +async fn get_composite_fields(client: &mut InnerClient, oid: Oid) -> Result, Error> { let stmt = typeinfo_composite_statement(client).await?; let rows = query::query(client, stmt, slice_iter(&[&oid])) @@ -250,7 +249,7 @@ async fn get_composite_fields(client: &Arc, oid: Oid) -> Result) -> Result { +async fn typeinfo_composite_statement(client: &mut InnerClient) -> Result { if let Some(stmt) = client.typeinfo_composite() { return Ok(stmt); } diff --git a/libs/proxy/tokio-postgres2/src/query.rs b/libs/proxy/tokio-postgres2/src/query.rs index 534195a707..b29e441a2b 100644 --- a/libs/proxy/tokio-postgres2/src/query.rs +++ b/libs/proxy/tokio-postgres2/src/query.rs @@ -14,7 +14,6 @@ use postgres_types2::{Format, ToSql, Type}; use std::fmt; use std::marker::PhantomPinned; use std::pin::Pin; -use std::sync::Arc; use std::task::{Context, Poll}; struct BorrowToSqlParamsDebug<'a>(&'a [&'a (dyn ToSql + Sync)]); @@ -26,7 +25,7 @@ impl fmt::Debug for BorrowToSqlParamsDebug<'_> { } pub async fn query<'a, I>( - client: &InnerClient, + client: &mut InnerClient, statement: Statement, params: I, ) -> Result @@ -57,7 +56,7 @@ where } pub async fn query_txt( - client: &Arc, + client: &mut InnerClient, query: &str, params: I, ) -> Result @@ -157,49 +156,6 @@ where }) } -pub async fn execute<'a, I>( - client: &InnerClient, - statement: Statement, - params: I, -) -> Result -where - I: IntoIterator, - I::IntoIter: ExactSizeIterator, -{ - let buf = if log_enabled!(Level::Debug) { - let params = params.into_iter().collect::>(); - debug!( - "executing statement {} with parameters: {:?}", - statement.name(), - BorrowToSqlParamsDebug(params.as_slice()), - ); - encode(client, &statement, params)? - } else { - encode(client, &statement, params)? - }; - let mut responses = start(client, buf).await?; - - let mut rows = 0; - loop { - match responses.next().await? { - Message::DataRow(_) => {} - Message::CommandComplete(body) => { - rows = body - .tag() - .map_err(Error::parse)? - .rsplit(' ') - .next() - .unwrap() - .parse() - .unwrap_or(0); - } - Message::EmptyQueryResponse => rows = 0, - Message::ReadyForQuery(_) => return Ok(rows), - _ => return Err(Error::unexpected_message()), - } - } -} - async fn start(client: &InnerClient, buf: Bytes) -> Result { let mut responses = client.send(RequestMessages::Single(FrontendMessage::Raw(buf)))?; @@ -211,7 +167,7 @@ async fn start(client: &InnerClient, buf: Bytes) -> Result { Ok(responses) } -pub fn encode<'a, I>(client: &InnerClient, statement: &Statement, params: I) -> Result +pub fn encode<'a, I>(client: &mut InnerClient, statement: &Statement, params: I) -> Result where I: IntoIterator, I::IntoIter: ExactSizeIterator, diff --git a/libs/proxy/tokio-postgres2/src/simple_query.rs b/libs/proxy/tokio-postgres2/src/simple_query.rs index fb2550377b..d8d86cf902 100644 --- a/libs/proxy/tokio-postgres2/src/simple_query.rs +++ b/libs/proxy/tokio-postgres2/src/simple_query.rs @@ -31,22 +31,8 @@ impl SimpleColumn { } } -pub async fn simple_query(client: &InnerClient, query: &str) -> Result { - debug!("executing simple query: {}", query); - - let buf = encode(client, query)?; - let responses = client.send(RequestMessages::Single(FrontendMessage::Raw(buf)))?; - - Ok(SimpleQueryStream { - responses, - columns: None, - status: ReadyForQueryStatus::Unknown, - _p: PhantomPinned, - }) -} - pub async fn batch_execute( - client: &InnerClient, + client: &mut InnerClient, query: &str, ) -> Result { debug!("executing statement batch: {}", query); @@ -66,7 +52,7 @@ pub async fn batch_execute( } } -pub(crate) fn encode(client: &InnerClient, query: &str) -> Result { +pub(crate) fn encode(client: &mut InnerClient, query: &str) -> Result { client.with_buf(|buf| { frontend::query(query, buf).map_err(Error::encode)?; Ok(buf.split().freeze()) diff --git a/libs/proxy/tokio-postgres2/src/statement.rs b/libs/proxy/tokio-postgres2/src/statement.rs index 22e160fc05..6e5cdfe02e 100644 --- a/libs/proxy/tokio-postgres2/src/statement.rs +++ b/libs/proxy/tokio-postgres2/src/statement.rs @@ -1,36 +1,13 @@ -use crate::client::InnerClient; -use crate::codec::FrontendMessage; -use crate::connection::RequestMessages; use crate::types::Type; -use postgres_protocol2::{ - message::{backend::Field, frontend}, - Oid, -}; -use std::{ - fmt, - sync::{Arc, Weak}, -}; +use postgres_protocol2::{message::backend::Field, Oid}; +use std::{fmt, sync::Arc}; struct StatementInner { - client: Weak, name: String, params: Vec, columns: Vec, } -impl Drop for StatementInner { - fn drop(&mut self) { - if let Some(client) = self.client.upgrade() { - let buf = client.with_buf(|buf| { - frontend::close(b'S', &self.name, buf).unwrap(); - frontend::sync(buf); - buf.split().freeze() - }); - let _ = client.send(RequestMessages::Single(FrontendMessage::Raw(buf))); - } - } -} - /// A prepared statement. /// /// Prepared statements can only be used with the connection that created them. @@ -38,14 +15,8 @@ impl Drop for StatementInner { pub struct Statement(Arc); impl Statement { - pub(crate) fn new( - inner: &Arc, - name: String, - params: Vec, - columns: Vec, - ) -> Statement { + pub(crate) fn new(name: String, params: Vec, columns: Vec) -> Statement { Statement(Arc::new(StatementInner { - client: Arc::downgrade(inner), name, params, columns, @@ -54,7 +25,6 @@ impl Statement { pub(crate) fn new_anonymous(params: Vec, columns: Vec) -> Statement { Statement(Arc::new(StatementInner { - client: Weak::new(), name: String::new(), params, columns, diff --git a/libs/proxy/tokio-postgres2/src/to_statement.rs b/libs/proxy/tokio-postgres2/src/to_statement.rs deleted file mode 100644 index 427f77dd79..0000000000 --- a/libs/proxy/tokio-postgres2/src/to_statement.rs +++ /dev/null @@ -1,57 +0,0 @@ -use crate::to_statement::private::{Sealed, ToStatementType}; -use crate::Statement; - -mod private { - use crate::{Client, Error, Statement}; - - pub trait Sealed {} - - pub enum ToStatementType<'a> { - Statement(&'a Statement), - Query(&'a str), - } - - impl<'a> ToStatementType<'a> { - pub async fn into_statement(self, client: &Client) -> Result { - match self { - ToStatementType::Statement(s) => Ok(s.clone()), - ToStatementType::Query(s) => client.prepare(s).await, - } - } - } -} - -/// A trait abstracting over prepared and unprepared statements. -/// -/// Many methods are generic over this bound, so that they support both a raw query string as well as a statement which -/// was prepared previously. -/// -/// This trait is "sealed" and cannot be implemented by anything outside this crate. -pub trait ToStatement: Sealed { - #[doc(hidden)] - fn __convert(&self) -> ToStatementType<'_>; -} - -impl ToStatement for Statement { - fn __convert(&self) -> ToStatementType<'_> { - ToStatementType::Statement(self) - } -} - -impl Sealed for Statement {} - -impl ToStatement for str { - fn __convert(&self) -> ToStatementType<'_> { - ToStatementType::Query(self) - } -} - -impl Sealed for str {} - -impl ToStatement for String { - fn __convert(&self) -> ToStatementType<'_> { - ToStatementType::Query(self) - } -} - -impl Sealed for String {} diff --git a/libs/proxy/tokio-postgres2/src/transaction.rs b/libs/proxy/tokio-postgres2/src/transaction.rs index 03a57e4947..b150f6a371 100644 --- a/libs/proxy/tokio-postgres2/src/transaction.rs +++ b/libs/proxy/tokio-postgres2/src/transaction.rs @@ -53,7 +53,7 @@ impl<'a> Transaction<'a> { } /// Like `Client::query_raw_txt`. - pub async fn query_raw_txt(&self, statement: &str, params: I) -> Result + pub async fn query_raw_txt(&mut self, statement: &str, params: I) -> Result where S: AsRef, I: IntoIterator>, @@ -68,7 +68,7 @@ impl<'a> Transaction<'a> { } /// Returns a reference to the underlying `Client`. - pub fn client(&self) -> &Client { + pub fn client(&mut self) -> &mut Client { self.client } } diff --git a/proxy/src/serverless/backend.rs b/proxy/src/serverless/backend.rs index 251aa47084..5b9067b76d 100644 --- a/proxy/src/serverless/backend.rs +++ b/proxy/src/serverless/backend.rs @@ -340,7 +340,7 @@ impl PoolingBackend { debug!("setting up backend session state"); // initiates the auth session - if let Err(e) = client.execute("select auth.init()", &[]).await { + if let Err(e) = client.batch_execute("select auth.init();").await { discard.discard(); return Err(e.into()); } diff --git a/proxy/src/serverless/local_conn_pool.rs b/proxy/src/serverless/local_conn_pool.rs index b84cde9e25..ac7178131b 100644 --- a/proxy/src/serverless/local_conn_pool.rs +++ b/proxy/src/serverless/local_conn_pool.rs @@ -23,7 +23,6 @@ use jose_jwk::jose_b64::base64ct::{Base64UrlUnpadded, Encoding}; use p256::ecdsa::{Signature, SigningKey}; use parking_lot::RwLock; use postgres_client::tls::NoTlsStream; -use postgres_client::types::ToSql; use postgres_client::AsyncMessage; use serde_json::value::RawValue; use signature::Signer; @@ -287,12 +286,11 @@ impl ClientInnerCommon { let token = resign_jwt(&local_data.key, payload, local_data.jti)?; // initiates the auth session - self.inner.batch_execute("discard all").await?; + // the token contains only `[a-zA-Z1-9_-\.]+` so it cannot escape the string literal formatting. self.inner - .execute( - "select auth.jwt_session_init($1)", - &[&&*token as &(dyn ToSql + Sync)], - ) + .batch_execute(&format!( + "discard all; select auth.jwt_session_init('{token}');" + )) .await?; let pid = self.inner.get_process_id(); diff --git a/proxy/src/serverless/sql_over_http.rs b/proxy/src/serverless/sql_over_http.rs index 5e85f5ec40..7bc650ba58 100644 --- a/proxy/src/serverless/sql_over_http.rs +++ b/proxy/src/serverless/sql_over_http.rs @@ -797,7 +797,7 @@ impl QueryData { let cancel_token = inner.cancel_token(); let res = match select( - pin!(query_to_json(config, &*inner, self, &mut 0, parsed_headers)), + pin!(query_to_json(config, &mut *inner, self, &mut 0, parsed_headers)), pin!(cancel.cancelled()), ) .await @@ -881,7 +881,7 @@ impl BatchQueryData { builder = builder.deferrable(true); } - let transaction = builder.start().await.inspect_err(|_| { + let mut transaction = builder.start().await.inspect_err(|_| { // if we cannot start a transaction, we should return immediately // and not return to the pool. connection is clearly broken discard.discard(); @@ -890,7 +890,7 @@ impl BatchQueryData { let json_output = match query_batch( config, cancel.child_token(), - &transaction, + &mut transaction, self, parsed_headers, ) @@ -934,7 +934,7 @@ impl BatchQueryData { async fn query_batch( config: &'static HttpConfig, cancel: CancellationToken, - transaction: &Transaction<'_>, + transaction: &mut Transaction<'_>, queries: BatchQueryData, parsed_headers: HttpHeaders, ) -> Result { @@ -972,7 +972,7 @@ async fn query_batch( async fn query_to_json( config: &'static HttpConfig, - client: &T, + client: &mut T, data: QueryData, current_size: &mut usize, parsed_headers: HttpHeaders,