Identify client in query engine (#570)

Author: levkk

pgdog/src/backend/pool/connection/mirror/mod.rs

@@ -15,7 +15,7 @@ use crate::frontend::client::timeouts::Timeouts;
 use crate::frontend::client::TransactionType;
 use crate::frontend::comms::comms;
 use crate::frontend::PreparedStatements;
-use crate::net::{Parameter, Parameters, Stream};
+use crate::net::{BackendKeyData, Parameter, Parameters, Stream};
 
 use crate::frontend::ClientRequest;
 
@@ -33,6 +33,8 @@ pub use request::*;
 /// to PgDog.
 #[derive(Debug)]
 pub struct Mirror {
+    /// Random identifier for this mirror connection.
+    pub id: BackendKeyData,
     /// Mirror's prepared statements. Should be similar
     /// to client's statements, if exposure is high.
     pub prepared_statements: PreparedStatements,
@@ -51,6 +53,7 @@ pub struct Mirror {
 impl Mirror {
     fn new(params: &Parameters, config: &ConfigAndUsers) -> Self {
         Self {
+            id: BackendKeyData::new(),
             prepared_statements: PreparedStatements::new(),
             params: params.clone(),
             timeouts: Timeouts::from_config(&config.config.general),

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

@@ -4,12 +4,15 @@ use crate::{
         client::{timeouts::Timeouts, TransactionType},
         Client, ClientRequest, PreparedStatements,
     },
-    net::{Parameters, Stream},
+    net::{BackendKeyData, Parameters, Stream},
     stats::memory::MemoryUsage,
 };
 
+#[allow(dead_code)]
 /// Context passed to the query engine to execute a query.
 pub struct QueryEngineContext<'a> {
+    /// Client ID running the query.
+    pub(super) id: &'a BackendKeyData,
     /// Prepared statements cache.
     pub(super) prepared_statements: &'a mut PreparedStatements,
     /// Client session parameters.
@@ -39,6 +42,7 @@ impl<'a> QueryEngineContext<'a> {
         let memory_usage = client.memory_usage();
 
         Self {
+            id: &client.id,
             prepared_statements: &mut client.prepared_statements,
             params: &mut client.params,
             client_request: &mut client.client_request,
@@ -62,6 +66,7 @@ impl<'a> QueryEngineContext<'a> {
     /// Create context from mirror.
     pub fn new_mirror(mirror: &'a mut Mirror, buffer: &'a mut ClientRequest) -> Self {
         Self {
+            id: &mirror.id,
             prepared_statements: &mut mirror.prepared_statements,
             params: &mut mirror.params,
             client_request: buffer,