Box Pgconnection fields (#3529)

* Update PgConnection code

* rustfmt

Author: joeydewaal

sqlx-postgres/src/connection/describe.rs

@@ -163,7 +163,7 @@ impl PgConnection {
         }
 
         // next we check a local cache for user-defined type names <-> object id
-        if let Some(info) = self.cache_type_info.get(&oid) {
+        if let Some(info) = self.inner.cache_type_info.get(&oid) {
             return Ok(info.clone());
         }
 
@@ -173,8 +173,9 @@ impl PgConnection {
 
             // cache the type name <-> oid relationship in a paired hashmap
             // so we don't come down this road again
-            self.cache_type_info.insert(oid, info.clone());
-            self.cache_type_oid
+            self.inner.cache_type_info.insert(oid, info.clone());
+            self.inner
+                .cache_type_oid
                 .insert(info.0.name().to_string().into(), oid);
 
             Ok(info)
@@ -374,7 +375,7 @@ WHERE rngtypid = $1
     }
 
     pub(crate) async fn fetch_type_id_by_name(&mut self, name: &str) -> Result<Oid, Error> {
-        if let Some(oid) = self.cache_type_oid.get(name) {
+        if let Some(oid) = self.inner.cache_type_oid.get(name) {
             return Ok(*oid);
         }
 
@@ -387,15 +388,18 @@ WHERE rngtypid = $1
                 type_name: name.into(),
             })?;
 
-        self.cache_type_oid.insert(name.to_string().into(), oid);
+        self.inner
+            .cache_type_oid
+            .insert(name.to_string().into(), oid);
         Ok(oid)
     }
 
     pub(crate) async fn fetch_array_type_id(&mut self, array: &PgArrayOf) -> Result<Oid, Error> {
         if let Some(oid) = self
+            .inner
             .cache_type_oid
             .get(&array.elem_name)
-            .and_then(|elem_oid| self.cache_elem_type_to_array.get(elem_oid))
+            .and_then(|elem_oid| self.inner.cache_elem_type_to_array.get(elem_oid))
         {
             return Ok(*oid);
         }
@@ -411,10 +415,13 @@ WHERE rngtypid = $1
                 })?;
 
         // Avoids copying `elem_name` until necessary
-        self.cache_type_oid
+        self.inner
+            .cache_type_oid
             .entry_ref(&array.elem_name)
             .insert(elem_oid);
-        self.cache_elem_type_to_array.insert(elem_oid, array_oid);
+        self.inner
+            .cache_elem_type_to_array
+            .insert(elem_oid, array_oid);
 
         Ok(array_oid)
     }
@@ -475,8 +482,16 @@ WHERE rngtypid = $1
             })?;
 
         // If the server is CockroachDB or Materialize, skip this step (#1248).
-        if !self.stream.parameter_statuses.contains_key("crdb_version")
-            && !self.stream.parameter_statuses.contains_key("mz_version")
+        if !self
+            .inner
+            .stream
+            .parameter_statuses
+            .contains_key("crdb_version")
+            && !self
+                .inner
+                .stream
+                .parameter_statuses
+                .contains_key("mz_version")
         {
             // patch up our null inference with data from EXPLAIN
             let nullable_patch = self

sqlx-postgres/src/connection/establish.rs

@@ -9,6 +9,8 @@ use crate::message::{
 };
 use crate::{PgConnectOptions, PgConnection};
 
+use super::PgConnectionInner;
+
 // https://www.postgresql.org/docs/current/protocol-flow.html#id-1.10.5.7.3
 // https://www.postgresql.org/docs/current/protocol-flow.html#id-1.10.5.7.11
 
@@ -134,18 +136,20 @@ impl PgConnection {
         }
 
         Ok(PgConnection {
-            stream,
-            process_id,
-            secret_key,
-            transaction_status,
-            transaction_depth: 0,
-            pending_ready_for_query_count: 0,
-            next_statement_id: StatementId::NAMED_START,
-            cache_statement: StatementCache::new(options.statement_cache_capacity),
-            cache_type_oid: HashMap::new(),
-            cache_type_info: HashMap::new(),
-            cache_elem_type_to_array: HashMap::new(),
-            log_settings: options.log_settings.clone(),
+            inner: Box::new(PgConnectionInner {
+                stream,
+                process_id,
+                secret_key,
+                transaction_status,
+                transaction_depth: 0,
+                pending_ready_for_query_count: 0,
+                next_statement_id: StatementId::NAMED_START,
+                cache_statement: StatementCache::new(options.statement_cache_capacity),
+                cache_type_oid: HashMap::new(),
+                cache_type_info: HashMap::new(),
+                cache_elem_type_to_array: HashMap::new(),
+                log_settings: options.log_settings.clone(),
+            }),
         })
     }
 }

sqlx-postgres/src/connection/executor.rs

@@ -26,8 +26,8 @@ async fn prepare(
     parameters: &[PgTypeInfo],
     metadata: Option<Arc<PgStatementMetadata>>,
 ) -> Result<(StatementId, Arc<PgStatementMetadata>), Error> {
-    let id = conn.next_statement_id;
-    conn.next_statement_id = id.next();
+    let id = conn.inner.next_statement_id;
+    conn.inner.next_statement_id = id.next();
 
     // build a list of type OIDs to send to the database in the PARSE command
     // we have not yet started the query sequence, so we are *safe* to cleanly make
@@ -43,23 +43,25 @@ async fn prepare(
     conn.wait_until_ready().await?;
 
     // next we send the PARSE command to the server
-    conn.stream.write_msg(Parse {
+    conn.inner.stream.write_msg(Parse {
         param_types: &param_types,
         query: sql,
         statement: id,
     })?;
 
     if metadata.is_none() {
         // get the statement columns and parameters
-        conn.stream.write_msg(message::Describe::Statement(id))?;
+        conn.inner
+            .stream
+            .write_msg(message::Describe::Statement(id))?;
     }
 
     // we ask for the server to immediately send us the result of the PARSE command
     conn.write_sync();
-    conn.stream.flush().await?;
+    conn.inner.stream.flush().await?;
 
     // indicates that the SQL query string is now successfully parsed and has semantic validity
-    conn.stream.recv_expect::<ParseComplete>().await?;
+    conn.inner.stream.recv_expect::<ParseComplete>().await?;
 
     let metadata = if let Some(metadata) = metadata {
         // each SYNC produces one READY FOR QUERY
@@ -94,11 +96,11 @@ async fn prepare(
 }
 
 async fn recv_desc_params(conn: &mut PgConnection) -> Result<ParameterDescription, Error> {
-    conn.stream.recv_expect().await
+    conn.inner.stream.recv_expect().await
 }
 
 async fn recv_desc_rows(conn: &mut PgConnection) -> Result<Option<RowDescription>, Error> {
-    let rows: Option<RowDescription> = match conn.stream.recv().await? {
+    let rows: Option<RowDescription> = match conn.inner.stream.recv().await? {
         // describes the rows that will be returned when the statement is eventually executed
         message if message.format == BackendMessageFormat::RowDescription => {
             Some(message.decode()?)
@@ -123,7 +125,7 @@ impl PgConnection {
     pub(super) async fn wait_for_close_complete(&mut self, mut count: usize) -> Result<(), Error> {
         // we need to wait for the [CloseComplete] to be returned from the server
         while count > 0 {
-            match self.stream.recv().await? {
+            match self.inner.stream.recv().await? {
                 message if message.format == BackendMessageFormat::PortalSuspended => {
                     // there was an open portal
                     // this can happen if the last time a statement was used it was not fully executed
@@ -148,12 +150,13 @@ impl PgConnection {
 
     #[inline(always)]
     pub(crate) fn write_sync(&mut self) {
-        self.stream
+        self.inner
+            .stream
             .write_msg(message::Sync)
             .expect("BUG: Sync should not be too big for protocol");
 
         // all SYNC messages will return a ReadyForQuery
-        self.pending_ready_for_query_count += 1;
+        self.inner.pending_ready_for_query_count += 1;
     }
 
     async fn get_or_prepare<'a>(
@@ -166,18 +169,18 @@ impl PgConnection {
         // a statement object
         metadata: Option<Arc<PgStatementMetadata>>,
     ) -> Result<(StatementId, Arc<PgStatementMetadata>), Error> {
-        if let Some(statement) = self.cache_statement.get_mut(sql) {
+        if let Some(statement) = self.inner.cache_statement.get_mut(sql) {
             return Ok((*statement).clone());
         }
 
         let statement = prepare(self, sql, parameters, metadata).await?;
 
-        if store_to_cache && self.cache_statement.is_enabled() {
-            if let Some((id, _)) = self.cache_statement.insert(sql, statement.clone()) {
-                self.stream.write_msg(Close::Statement(id))?;
+        if store_to_cache && self.inner.cache_statement.is_enabled() {
+            if let Some((id, _)) = self.inner.cache_statement.insert(sql, statement.clone()) {
+                self.inner.stream.write_msg(Close::Statement(id))?;
                 self.write_sync();
 
-                self.stream.flush().await?;
+                self.inner.stream.flush().await?;
 
                 self.wait_for_close_complete(1).await?;
                 self.recv_ready_for_query().await?;
@@ -195,7 +198,7 @@ impl PgConnection {
         persistent: bool,
         metadata_opt: Option<Arc<PgStatementMetadata>>,
     ) -> Result<impl Stream<Item = Result<Either<PgQueryResult, PgRow>, Error>> + 'e, Error> {
-        let mut logger = QueryLogger::new(query, self.log_settings.clone());
+        let mut logger = QueryLogger::new(query, self.inner.log_settings.clone());
 
         // before we continue, wait until we are "ready" to accept more queries
         self.wait_until_ready().await?;
@@ -231,7 +234,7 @@ impl PgConnection {
             self.wait_until_ready().await?;
 
             // bind to attach the arguments to the statement and create a portal
-            self.stream.write_msg(Bind {
+            self.inner.stream.write_msg(Bind {
                 portal: PortalId::UNNAMED,
                 statement,
                 formats: &[PgValueFormat::Binary],
@@ -242,7 +245,7 @@ impl PgConnection {
 
             // executes the portal up to the passed limit
             // the protocol-level limit acts nearly identically to the `LIMIT` in SQL
-            self.stream.write_msg(message::Execute {
+            self.inner.stream.write_msg(message::Execute {
                 portal: PortalId::UNNAMED,
                 limit: limit.into(),
             })?;
@@ -255,7 +258,9 @@ impl PgConnection {
 
             // we ask the database server to close the unnamed portal and free the associated resources
             // earlier - after the execution of the current query.
-            self.stream.write_msg(Close::Portal(PortalId::UNNAMED))?;
+            self.inner
+                .stream
+                .write_msg(Close::Portal(PortalId::UNNAMED))?;
 
             // finally, [Sync] asks postgres to process the messages that we sent and respond with
             // a [ReadyForQuery] message when it's completely done. Theoretically, we could send
@@ -268,8 +273,8 @@ impl PgConnection {
             PgValueFormat::Binary
         } else {
             // Query will trigger a ReadyForQuery
-            self.stream.write_msg(Query(query))?;
-            self.pending_ready_for_query_count += 1;
+            self.inner.stream.write_msg(Query(query))?;
+            self.inner.pending_ready_for_query_count += 1;
 
             // metadata starts out as "nothing"
             metadata = Arc::new(PgStatementMetadata::default());
@@ -278,11 +283,11 @@ impl PgConnection {
             PgValueFormat::Text
         };
 
-        self.stream.flush().await?;
+        self.inner.stream.flush().await?;
 
         Ok(try_stream! {
             loop {
-                let message = self.stream.recv().await?;
+                let message = self.inner.stream.recv().await?;
 
                 match message.format {
                     BackendMessageFormat::BindComplete