Annotate explain plans with routing decisions (#546)

* initial draft internal explain annotations

* clean up formatting logic in query parser

* explain plan annotation integration test

* explain annotation testing

* enable explain annotations only when expanded_explain flag set

* Enable expanded_explain in integration tests

* Tweak explain recorder

* better integration debugging

* less strict assertion :(

Author: jaggederest

integration/pgdog.toml

@@ -11,6 +11,7 @@ read_write_strategy = "aggressive"
 openmetrics_port = 9090
 openmetrics_namespace = "pgdog_"
 prepared_statements_limit = 500
+expanded_explain = true
 # dns_ttl = 15_000
 query_cache_limit = 500
 pub_sub_channel_size = 4098

integration/rust/tests/integration/explain.rs

@@ -0,0 +1,92 @@
+use rust::setup::connections_tokio;
+use tokio_postgres::SimpleQueryMessage;
+
+#[tokio::test]
+async fn explain_routing_annotations_surface() -> Result<(), Box<dyn std::error::Error>> {
+    let mut clients = connections_tokio().await;
+    let sharded = clients.swap_remove(1);
+
+    for shard in [0, 1] {
+        let drop = format!(
+            "/* pgdog_shard: {} */ DROP TABLE IF EXISTS explain_avg_test",
+            shard
+        );
+        sharded.simple_query(drop.as_str()).await.ok();
+    }
+
+    for shard in [0, 1] {
+        let create = format!(
+            "/* pgdog_shard: {} */ CREATE TABLE explain_avg_test(price DOUBLE PRECISION)",
+            shard
+        );
+        sharded.simple_query(create.as_str()).await?;
+    }
+
+    sharded
+        .simple_query(
+            "/* pgdog_shard: 0 */ INSERT INTO explain_avg_test(price) VALUES (10.0), (14.0)",
+        )
+        .await?;
+    sharded
+        .simple_query(
+            "/* pgdog_shard: 1 */ INSERT INTO explain_avg_test(price) VALUES (18.0), (22.0)",
+        )
+        .await?;
+
+    let rows = sharded
+        .simple_query("EXPLAIN SELECT AVG(price) FROM explain_avg_test")
+        .await?;
+
+    let mut plan_lines = vec![];
+    for message in rows {
+        if let SimpleQueryMessage::Row(row) = message {
+            plan_lines.push(row.get(0).unwrap_or_default().to_string());
+        }
+    }
+
+    assert!(
+        plan_lines.iter().any(|line| line.contains("Aggregate")),
+        "missing Aggregate node in EXPLAIN output: {:?}",
+        plan_lines
+    );
+    let routing_header = plan_lines
+        .iter()
+        .find(|line| line.contains("PgDog Routing:"));
+    assert!(
+        routing_header.is_some(),
+        "missing PgDog Routing header in EXPLAIN output: {:?}",
+        plan_lines
+    );
+    let summary_line = plan_lines
+        .iter()
+        .find(|line| line.contains("Summary:"))
+        .cloned();
+    assert!(
+        summary_line
+            .as_ref()
+            .map(|line| line.contains("Summary: shard=all"))
+            .unwrap_or(false),
+        "unexpected summary line: {:?} (all lines: {:?})",
+        summary_line,
+        plan_lines
+    );
+    let broadcast_line = plan_lines
+        .iter()
+        .find(|line| line.contains("no sharding key matched"))
+        .cloned();
+    assert!(
+        broadcast_line.is_some(),
+        "missing broadcast note in EXPLAIN output: {:?}",
+        plan_lines
+    );
+
+    for shard in [0, 1] {
+        let drop = format!(
+            "/* pgdog_shard: {} */ DROP TABLE IF EXISTS explain_avg_test",
+            shard
+        );
+        sharded.simple_query(drop.as_str()).await.ok();
+    }
+
+    Ok(())
+}

integration/rust/tests/integration/mod.rs

@@ -3,6 +3,7 @@ pub mod avg;
 pub mod ban;
 pub mod cross_shard_disabled;
 pub mod distinct;
+pub mod explain;
 pub mod fake_transactions;
 pub mod maintenance_mode;
 pub mod notify;

pgdog/src/config/general.rs

@@ -156,6 +156,9 @@ pub struct General {
     /// Two-phase commit automatic transactions.
     #[serde(default)]
     pub two_phase_commit_auto: Option<bool>,
+    /// Enable expanded EXPLAIN output.
+    #[serde(default = "General::expanded_explain")]
+    pub expanded_explain: bool,
 }
 
 impl Default for General {
@@ -211,6 +214,7 @@ impl Default for General {
             log_disconnections: Self::log_disconnections(),
             two_phase_commit: bool::default(),
             two_phase_commit_auto: None,
+            expanded_explain: Self::expanded_explain(),
             server_lifetime: Self::server_lifetime(),
         }
     }
@@ -479,6 +483,10 @@ impl General {
         Self::env_bool_or_default("PGDOG_LOG_DISCONNECTIONS", true)
     }
 
+    pub fn expanded_explain() -> bool {
+        Self::env_bool_or_default("PGDOG_EXPANDED_EXPLAIN", false)
+    }
+
     pub fn server_lifetime() -> u64 {
         Self::env_or_default(
             "PGDOG_SERVER_LIFETIME",

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

@@ -1,5 +1,6 @@
 use crate::{
     backend::pool::{Connection, Request},
+    config::config,
     frontend::{
         client::query_engine::hooks::QueryEngineHooks,
         router::{parser::Shard, Route},
@@ -32,6 +33,7 @@ pub mod unknown_command;
 #[cfg(test)]
 mod testing;
 
+use self::query::ExplainResponseState;
 pub use context::QueryEngineContext;
 use notify_buffer::NotifyBuffer;
 pub use two_pc::phase::TwoPcPhase;
@@ -50,6 +52,7 @@ pub struct QueryEngine {
     set_route: Option<Route>,
     two_pc: TwoPc,
     notify_buffer: NotifyBuffer,
+    pending_explain: Option<ExplainResponseState>,
 }
 
 impl QueryEngine {
@@ -130,13 +133,21 @@ impl QueryEngine {
         // to have accurate timings between queries.
         self.backend.mirror(context.client_request);
 
+        self.pending_explain = None;
+
         let command = self.router.command();
-        let route = if let Some(ref route) = self.set_route {
+        let mut route = if let Some(ref route) = self.set_route {
             route.clone()
         } else {
             command.route().clone()
         };
 
+        if let Some(trace) = route.take_explain() {
+            if config().config.general.expanded_explain {
+                self.pending_explain = Some(ExplainResponseState::new(trace));
+            }
+        }
+
         // FIXME, we should not to copy route twice.
         context.client_request.route = Some(route.clone());
 

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

@@ -1,10 +1,10 @@
 use tokio::time::timeout;
 
 use crate::{
-    frontend::client::TransactionType,
+    frontend::{client::TransactionType, router::parser::explain_trace::ExplainTrace},
     net::{
-        FromBytes, Message, Protocol, ProtocolMessage, Query, ReadyForQuery, ToBytes,
-        TransactionState,
+        DataRow, FromBytes, Message, Protocol, ProtocolMessage, Query, ReadyForQuery,
+        RowDescription, ToBytes, TransactionState,
     },
     state::State,
 };
@@ -96,9 +96,35 @@ impl QueryEngine {
         self.streaming = message.streaming();
 
         let code = message.code();
+        let payload = if code == 'T' {
+            Some(message.payload())
+        } else {
+            None
+        };
         let mut message = message.backend();
         let has_more_messages = self.backend.has_more_messages();
 
+        if let Some(bytes) = payload {
+            if let Some(state) = self.pending_explain.as_mut() {
+                if let Ok(row_description) = RowDescription::from_bytes(bytes) {
+                    state.capture_row_description(row_description);
+                } else {
+                    state.annotated = true;
+                }
+            }
+        }
+
+        if code == 'C' {
+            self.emit_explain_rows(context).await?;
+        }
+
+        if code == 'E' {
+            if let Some(state) = self.pending_explain.as_mut() {
+                state.annotated = true;
+            }
+            self.pending_explain = None;
+        }
+
         // Messages that we need to send to the client immediately.
         // ReadyForQuery (B) | CopyInResponse (B) | ErrorResponse(B) | NoticeResponse(B) | NotificationResponse (B)
         let flush = matches!(code, 'Z' | 'G' | 'E' | 'N' | 'A')
@@ -185,6 +211,38 @@ impl QueryEngine {
             context.stream.send(&message).await?;
         }
 
+        if code == 'Z' {
+            self.pending_explain = None;
+        }
+
+        Ok(())
+    }
+
+    async fn emit_explain_rows(
+        &mut self,
+        context: &mut QueryEngineContext<'_>,
+    ) -> Result<(), Error> {
+        if let Some(state) = self.pending_explain.as_mut() {
+            if !state.should_emit() {
+                return Ok(());
+            }
+
+            if state.row_description.is_none() {
+                return Ok(());
+            }
+
+            for line in state.lines.clone() {
+                let mut row = DataRow::new();
+                row.add(line);
+                let message = row.message()?;
+                let len = message.len();
+                context.stream.send(&message).await?;
+                self.stats.sent(len);
+            }
+
+            state.annotated = true;
+        }
+
         Ok(())
     }
 
@@ -310,3 +368,36 @@ impl QueryEngine {
         }
     }
 }
+
+#[derive(Debug, Default, Clone)]
+pub(super) struct ExplainResponseState {
+    lines: Vec<String>,
+    row_description: Option<RowDescription>,
+    annotated: bool,
+    supported: bool,
+}
+
+impl ExplainResponseState {
+    pub fn new(trace: ExplainTrace) -> Self {
+        Self {
+            lines: trace.render_lines(),
+            row_description: None,
+            annotated: false,
+            supported: false,
+        }
+    }
+
+    pub fn capture_row_description(&mut self, row_description: RowDescription) {
+        self.supported = row_description.fields.len() == 1
+            && matches!(row_description.field(0).map(|f| f.type_oid), Some(25));
+        if self.supported {
+            self.row_description = Some(row_description);
+        } else {
+            self.annotated = true;
+        }
+    }
+
+    pub fn should_emit(&self) -> bool {
+        self.supported && !self.annotated
+    }
+}

pgdog/src/frontend/router/parser/context.rs

@@ -44,6 +44,8 @@ pub struct QueryParserContext<'a> {
     pub(super) multi_tenant: &'a Option<MultiTenant>,
     /// Dry run enabled?
     pub(super) dry_run: bool,
+    /// Expanded EXPLAIN annotations enabled?
+    pub(super) expanded_explain: bool,
 }
 
 impl<'a> QueryParserContext<'a> {
@@ -61,6 +63,7 @@ impl<'a> QueryParserContext<'a> {
             pub_sub_enabled: config.config.general.pub_sub_enabled(),
             multi_tenant: router_context.cluster.multi_tenant(),
             dry_run: config.config.general.dry_run,
+            expanded_explain: config.config.general.expanded_explain,
             router_context,
         }
     }
@@ -136,4 +139,8 @@ impl<'a> QueryParserContext<'a> {
             params,
         }
     }
+
+    pub(super) fn expanded_explain(&self) -> bool {
+        self.expanded_explain
+    }
 }