fix: ensure PG connection is established before using it

Fixes launchbadge#1940.

Author: crepererum

sqlx-core/src/postgres/listener.rs

@@ -9,6 +9,10 @@ use either::Either;
 use futures_channel::mpsc;
 use futures_core::future::BoxFuture;
 use futures_core::stream::{BoxStream, Stream};
+use futures_util::{
+    stream::{StreamExt, TryStreamExt},
+    FutureExt,
+};
 use std::fmt::{self, Debug};
 use std::io;
 use std::str::from_utf8;
@@ -65,6 +69,7 @@ impl PgListener {
     /// The channel name is quoted here to ensure case sensitivity.
     pub async fn listen(&mut self, channel: &str) -> Result<(), Error> {
         self.connection()
+            .await?
             .execute(&*format!(r#"LISTEN "{}""#, ident(channel)))
             .await?;
 
@@ -81,11 +86,8 @@ impl PgListener {
         let beg = self.channels.len();
         self.channels.extend(channels.into_iter().map(|s| s.into()));
 
-        self.connection
-            .as_mut()
-            .unwrap()
-            .execute(&*build_listen_all_query(&self.channels[beg..]))
-            .await?;
+        let query = build_listen_all_query(&self.channels[beg..]);
+        self.connection().await?.execute(&*query).await?;
 
         Ok(())
     }
@@ -94,6 +96,7 @@ impl PgListener {
     /// The channel name is quoted here to ensure case sensitivity.
     pub async fn unlisten(&mut self, channel: &str) -> Result<(), Error> {
         self.connection()
+            .await?
             .execute(&*format!(r#"UNLISTEN "{}""#, ident(channel)))
             .await?;
 
@@ -106,7 +109,7 @@ impl PgListener {
 
     /// Stops listening for notifications on all channels.
     pub async fn unlisten_all(&mut self) -> Result<(), Error> {
-        self.connection().execute("UNLISTEN *").await?;
+        self.connection().await?.execute("UNLISTEN *").await?;
 
         self.channels.clear();
 
@@ -130,8 +133,11 @@ impl PgListener {
     }
 
     #[inline]
-    fn connection(&mut self) -> &mut PgConnection {
-        self.connection.as_mut().unwrap()
+    async fn connection(&mut self) -> Result<&mut PgConnection, Error> {
+        // Ensure we have an active connection to work with.
+        self.connect_if_needed().await?;
+
+        Ok(self.connection.as_mut().unwrap())
     }
 
     /// Receives the next notification available from any of the subscribed channels.
@@ -203,16 +209,13 @@ impl PgListener {
         }
 
         loop {
-            // Ensure we have an active connection to work with.
-            self.connect_if_needed().await?;
-
-            let message = match self.connection().stream.recv_unchecked().await {
+            let message = match self.connection().await?.stream.recv_unchecked().await {
                 Ok(message) => message,
 
                 // The connection is dead, ensure that it is dropped,
                 // update self state, and loop to try again.
                 Err(Error::Io(err)) if err.kind() == io::ErrorKind::ConnectionAborted => {
-                    self.buffer_tx = self.connection().stream.notifications.take();
+                    self.buffer_tx = self.connection().await?.stream.notifications.take();
                     self.connection = None;
 
                     // lost connection
@@ -233,7 +236,7 @@ impl PgListener {
 
                 // Mark the connection as ready for another query
                 MessageFormat::ReadyForQuery => {
-                    self.connection().pending_ready_for_query_count -= 1;
+                    self.connection().await?.pending_ready_for_query_count -= 1;
                 }
 
                 // Ignore unexpected messages
@@ -292,7 +295,13 @@ impl<'c> Executor<'c> for &'c mut PgListener {
         'c: 'e,
         E: Execute<'q, Self::Database>,
     {
-        self.connection().fetch_many(query)
+        futures_util::stream::once(async move {
+            // need some basic type annotation to help the compiler a bit
+            let res: Result<_, Error> = Ok(self.connection().await?.fetch_many(query));
+            res
+        })
+        .try_flatten()
+        .boxed()
     }
 
     fn fetch_optional<'e, 'q: 'e, E: 'q>(
@@ -303,7 +312,7 @@ impl<'c> Executor<'c> for &'c mut PgListener {
         'c: 'e,
         E: Execute<'q, Self::Database>,
     {
-        self.connection().fetch_optional(query)
+        async move { self.connection().await?.fetch_optional(query).await }.boxed()
     }
 
     fn prepare_with<'e, 'q: 'e>(
@@ -314,7 +323,13 @@ impl<'c> Executor<'c> for &'c mut PgListener {
     where
         'c: 'e,
     {
-        self.connection().prepare_with(query, parameters)
+        async move {
+            self.connection()
+                .await?
+                .prepare_with(query, parameters)
+                .await
+        }
+        .boxed()
     }
 
     #[doc(hidden)]
@@ -325,7 +340,7 @@ impl<'c> Executor<'c> for &'c mut PgListener {
     where
         'c: 'e,
     {
-        self.connection().describe(query)
+        async move { self.connection().await?.describe(query).await }.boxed()
     }
 }