Map mobc pool errors to type that would invalidate connection

wolf4ood · Aug 27, 2024 · 802a20c · 802a20c
1 parent 093c04c
commit 802a20c
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 5 deletions.
diff --git a/gremlin-client/src/aio/connection.rs b/gremlin-client/src/aio/connection.rs
@@ -196,7 +196,9 @@ impl Conn {
             .map_err(|e| {
                 //If there's been an websocket layer error, mark the connection as invalid
                 match e {
-                    GremlinError::WebSocket(_) | GremlinError::WebSocketAsync(_) => {
+                    GremlinError::WebSocket(_)
+                    | GremlinError::WebSocketAsync(_)
+                    | GremlinError::WebSocketPoolAsync(_) => {
                         self.valid = false;
                     }
                     _ => {}

diff --git a/gremlin-client/src/error.rs b/gremlin-client/src/error.rs
@@ -39,6 +39,9 @@ pub enum GremlinError {
     WebSocketAsync(#[from] Arc<async_tungstenite::tungstenite::Error>),
     #[cfg(feature = "async_gremlin")]
     #[error(transparent)]
+    WebSocketPoolAsync(#[from] Arc<mobc::Error<GremlinError>>),
+    #[cfg(feature = "async_gremlin")]
+    #[error(transparent)]
     ChannelSend(#[from] futures::channel::mpsc::SendError),
     #[error(transparent)]
     Uuid(#[from] uuid::Error),
@@ -49,10 +52,7 @@ impl From<mobc::Error<GremlinError>> for GremlinError {
     fn from(e: mobc::Error<GremlinError>) -> GremlinError {
         match e {
             mobc::Error::Inner(e) => e,
-            mobc::Error::BadConn => {
-                GremlinError::Generic(String::from("Async pool bad connection"))
-            }
-            mobc::Error::Timeout => GremlinError::Generic(String::from("Async pool timeout")),
+            other => GremlinError::WebSocketPoolAsync(Arc::new(other)),
         }
     }
 }