Add after_connected hook (#558)

Author: levkk

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

@@ -2,8 +2,15 @@
 #![allow(unused_variables, dead_code)]
 use super::*;
 
+#[derive(Debug)]
 pub struct QueryEngineHooks;
 
+impl Default for QueryEngineHooks {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
 impl QueryEngineHooks {
     pub(super) fn new() -> Self {
         Self {}
@@ -16,6 +23,14 @@ impl QueryEngineHooks {
         Ok(())
     }
 
+    pub(super) fn after_connected(
+        &mut self,
+        context: &mut QueryEngineContext<'_>,
+        backend: &Connection,
+    ) -> Result<(), Error> {
+        Ok(())
+    }
+
     pub(super) fn after_execution(
         &mut self,
         context: &mut QueryEngineContext<'_>,

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

@@ -53,6 +53,7 @@ pub struct QueryEngine {
     two_pc: TwoPc,
     notify_buffer: NotifyBuffer,
     pending_explain: Option<ExplainResponseState>,
+    hooks: QueryEngineHooks,
 }
 
 impl QueryEngine {
@@ -72,6 +73,7 @@ impl QueryEngine {
             backend,
             client_id: comms.client_id(),
             comms: comms.clone(),
+            hooks: QueryEngineHooks::new(),
             #[cfg(test)]
             test_mode: true,
             #[cfg(not(test))]
@@ -125,8 +127,7 @@ impl QueryEngine {
             return Ok(());
         }
 
-        let mut hooks = QueryEngineHooks::new();
-        hooks.before_execution(context)?;
+        self.hooks.before_execution(context)?;
 
         // Queue up request to mirrors, if any.
         // Do this before sending query to actual server
@@ -223,7 +224,7 @@ impl QueryEngine {
             command => self.unknown_command(context, command.clone()).await?,
         }
 
-        hooks.after_execution(context)?;
+        self.hooks.after_execution(context)?;
 
         if context.in_error() {
             self.backend.mirror_clear();

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

@@ -68,6 +68,8 @@ impl QueryEngine {
             }
         }
 
+        self.hooks.after_connected(context, &self.backend)?;
+
         self.backend
             .handle_client_request(context.client_request, &mut self.router, self.streaming)
             .await?;