Track more errors in query engine (#582)

Author: levkk

pgdog/src/frontend/client/query_engine/connect.rs

@@ -61,10 +61,16 @@ impl QueryEngine {
 
                 if err.no_server() {
                     error!("{} [{:?}]", err, context.stream.peer_addr());
+
+                    let error = ErrorResponse::from_err(&err);
+
+                    self.hooks.on_engine_error(context, &error)?;
+
                     let bytes_sent = context
                         .stream
-                        .error(ErrorResponse::from_err(&err), context.in_transaction())
+                        .error(error, context.in_transaction())
                         .await?;
+
                     self.stats.sent(bytes_sent);
                     self.backend.disconnect();
                     self.router.reset();

pgdog/src/frontend/client/query_engine/hooks/mod.rs

@@ -45,4 +45,12 @@ impl QueryEngineHooks {
     ) -> Result<(), Error> {
         Ok(())
     }
+
+    pub(super) fn on_engine_error(
+        &mut self,
+        context: &mut QueryEngineContext<'_>,
+        error: &ErrorResponse,
+    ) -> Result<(), Error> {
+        Ok(())
+    }
 }

pgdog/src/frontend/client/query_engine/query.rs

@@ -365,12 +365,13 @@ impl QueryEngine {
             && context.client_request.executable()
             && !route.rollback_savepoint()
         {
+            let error = ErrorResponse::in_failed_transaction();
+
+            self.hooks.on_engine_error(context, &error)?;
+
             let bytes_sent = context
                 .stream
-                .error(
-                    ErrorResponse::in_failed_transaction(),
-                    context.in_transaction(),
-                )
+                .error(error, context.in_transaction())
                 .await?;
             self.stats.sent(bytes_sent);
 

pgdog/src/frontend/client/query_engine/route_query.rs

@@ -31,12 +31,14 @@ impl QueryEngine {
             }
             Err(err) => {
                 error!("{:?} [{:?}]", err, context.stream.peer_addr());
+
+                let error = ErrorResponse::syntax(err.to_string().as_str());
+
+                self.hooks.on_engine_error(context, &error)?;
+
                 let bytes_sent = context
                     .stream
-                    .error(
-                        ErrorResponse::syntax(err.to_string().as_str()),
-                        context.in_transaction(),
-                    )
+                    .error(error, context.in_transaction())
                     .await?;
                 self.stats.sent(bytes_sent);
                 return Ok(false);