nom-sql: Added support for SHOW PROXIED SUPPORTED QUERIES

On a heavy-load system, SHOW PROXIED QUERIES can show a large number
of queries that might not be eligible for caching.

Added support for a variant of SHOW PROXIED QUERIES that accept an
optional [SUPPORTED] keyword. If present, only queries that are
eligible for caching are shown.

Release-Note-Core: You can now optionally display only supported proxied
  queries with the command `SHOW PROXIED SUPPORTED QUERIES`. Thanks,
  altmannmarcelo!
Change-Id: I9611ae07c3fb50e477ac198d7e86f161000de1cb
Reviewed-on: https://gerrit.readyset.name/c/readyset/+/5892
Reviewed-by: Aspen Smith <aspen@readyset.io>
Reviewed-by: Luke Osborne <luke@readyset.io>
Tested-by: Buildkite CI

nom-sql/src/show.rs

@@ -19,7 +19,7 @@ pub enum ShowStatement {
     Events,
     Tables(Tables),
     CachedQueries(Option<QueryID>),
-    ProxiedQueries(Option<QueryID>),
+    ProxiedQueries(ProxiedQueriesOptions),
     ReadySetStatus,
     ReadySetMigrationStatus(u64),
     ReadySetVersion,
@@ -40,11 +40,15 @@ impl ShowStatement {
                         write!(f, "CACHES")
                     }
                 }
-                Self::ProxiedQueries(maybe_query_id) => {
-                    if let Some(query_id) = maybe_query_id {
-                        write!(f, "PROXIED QUERIES WHERE query_id = {}", query_id)
+                Self::ProxiedQueries(options) => {
+                    write!(f, "PROXIED ")?;
+                    if options.only_supported {
+                        write!(f, "SUPPORTED ")?;
+                    }
+                    if let Some(query_id) = &options.query_id {
+                        write!(f, "QUERIES WHERE query_id = {}", query_id)
                     } else {
-                        write!(f, "PROXIED QUERIES")
+                        write!(f, "QUERIES")
                     }
                 }
                 Self::ReadySetStatus => write!(f, "READYSET STATUS"),
@@ -79,16 +83,31 @@ fn cached_queries(
     }
 }
 
+#[derive(Clone, Debug, Eq, Hash, PartialEq, Serialize, Deserialize)]
+pub struct ProxiedQueriesOptions {
+    pub query_id: Option<String>,
+    pub only_supported: bool,
+}
+
 fn proxied_queries(
     dialect: Dialect,
 ) -> impl Fn(LocatedSpan<&[u8]>) -> NomSqlResult<&[u8], ShowStatement> {
     move |i| {
         let (i, _) = tag_no_case("proxied")(i)?;
+        let (i, only_supported) = map(
+            opt(tuple((whitespace1, tag_no_case("supported")))),
+            |only_supported| only_supported.is_some(),
+        )(i)?;
         let (i, _) = whitespace1(i)?;
         let (i, _) = tag_no_case("queries")(i)?;
-        let (i, q_id) = opt(preceded(whitespace1, where_query_id(dialect)))(i)?;
-
-        Ok((i, ShowStatement::ProxiedQueries(q_id)))
+        let (i, query_id) = opt(preceded(whitespace1, where_query_id(dialect)))(i)?;
+        Ok((
+            i,
+            ShowStatement::ProxiedQueries(ProxiedQueriesOptions {
+                query_id,
+                only_supported,
+            }),
+        ))
     }
 }
 
@@ -337,8 +356,42 @@ mod tests {
         let res2 = show(Dialect::MySQL)(LocatedSpan::new(qstring2.as_bytes()))
             .unwrap()
             .1;
-        assert_eq!(res1, ShowStatement::ProxiedQueries(None));
-        assert_eq!(res2, ShowStatement::ProxiedQueries(None));
+        let qstring3 = "SHOW PROXIED SUPPORTED QUERIES";
+        let res3 = show(Dialect::MySQL)(LocatedSpan::new(qstring3.as_bytes()))
+            .unwrap()
+            .1;
+        let qstring4 = "SHOW\tPROXIED\tSUPPORTED\tQUERIES";
+        let res4 = show(Dialect::MySQL)(LocatedSpan::new(qstring4.as_bytes()))
+            .unwrap()
+            .1;
+        assert_eq!(
+            res1,
+            ShowStatement::ProxiedQueries(ProxiedQueriesOptions {
+                query_id: None,
+                only_supported: false
+            })
+        );
+        assert_eq!(
+            res2,
+            ShowStatement::ProxiedQueries(ProxiedQueriesOptions {
+                query_id: None,
+                only_supported: false
+            })
+        );
+        assert_eq!(
+            res3,
+            ShowStatement::ProxiedQueries(ProxiedQueriesOptions {
+                query_id: None,
+                only_supported: true
+            })
+        );
+        assert_eq!(
+            res4,
+            ShowStatement::ProxiedQueries(ProxiedQueriesOptions {
+                query_id: None,
+                only_supported: true
+            })
+        );
     }
 
     #[test]
@@ -347,9 +400,23 @@ mod tests {
         let res1 = show(Dialect::MySQL)(LocatedSpan::new(qstring1.as_bytes()))
             .unwrap()
             .1;
+        let qstring2 = "SHOW PROXIED SUPPORTED QUERIES where query_id = 'test'";
+        let res2 = show(Dialect::MySQL)(LocatedSpan::new(qstring2.as_bytes()))
+            .unwrap()
+            .1;
         assert_eq!(
             res1,
-            ShowStatement::ProxiedQueries(Some("test".to_string()))
+            ShowStatement::ProxiedQueries(ProxiedQueriesOptions {
+                query_id: Some("test".to_string()),
+                only_supported: false
+            })
+        );
+        assert_eq!(
+            res2,
+            ShowStatement::ProxiedQueries(ProxiedQueriesOptions {
+                query_id: Some("test".to_string()),
+                only_supported: true
+            })
         );
     }
 

readyset-adapter/src/backend.rs

@@ -1715,11 +1715,17 @@ where
     async fn show_proxied_queries(
         &mut self,
         query_id: &Option<String>,
+        only_supported: bool,
     ) -> ReadySetResult<noria_connector::QueryResult<'static>> {
         let mut queries = self.state.query_status_cache.deny_list();
         if let Some(q_id) = query_id {
             queries.retain(|q| &q.id.to_string() == q_id);
         }
+
+        if only_supported {
+            queries.retain(|q| q.status.migration_state.is_supported());
+        }
+
         let select_schema = SelectSchema {
             schema: Cow::Owned(vec![
                 create_dummy_column("query id"),
@@ -1918,7 +1924,7 @@ where
             }
             SqlQuery::Show(ShowStatement::ReadySetVersion) => readyset_version(),
             SqlQuery::Show(ShowStatement::ReadySetTables) => self.noria.table_statuses().await,
-            SqlQuery::Show(ShowStatement::ProxiedQueries(q_id)) => {
+            SqlQuery::Show(ShowStatement::ProxiedQueries(proxied_queries_options)) => {
                 // Log a telemetry event
                 if let Some(ref telemetry_sender) = self.telemetry_sender {
                     if let Err(e) = telemetry_sender.send_event(TelemetryEvent::ShowProxiedQueries)
@@ -1929,7 +1935,11 @@ where
                     trace!("No telemetry sender. not sending metric for SHOW PROXIED QUERIES");
                 }
 
-                self.show_proxied_queries(q_id).await
+                self.show_proxied_queries(
+                    &proxied_queries_options.query_id,
+                    proxied_queries_options.only_supported,
+                )
+                .await
             }
             _ => {
                 drop(_t);

readyset-client/src/query.rs

@@ -366,6 +366,14 @@ impl MigrationState {
             MigrationState::Pending | MigrationState::DryRunSucceeded
         )
     }
+
+    /// Returns true if the query should be considered "supported"
+    pub fn is_supported(&self) -> bool {
+        matches!(
+            self,
+            MigrationState::Dropped | MigrationState::DryRunSucceeded | MigrationState::Successful
+        )
+    }
 }
 
 impl Display for MigrationState {