Only log cost if the query is not a healthcheck

MystenLabs · Jan 30, 2024 · 0133462 · 0133462
1 parent a9f9bec
commit 0133462
Show file tree

Hide file tree

Showing 5 changed files with 189 additions and 143 deletions.
diff --git a/crates/sui-graphql-rpc/src/data.rs b/crates/sui-graphql-rpc/src/data.rs
@@ -9,6 +9,7 @@ use diesel::{
     query_dsl::{methods::LimitDsl, LoadQuery},
     QueryResult,
 };
+use uuid::Uuid;
 
 use crate::error::Error;
 
@@ -69,28 +70,28 @@ pub(crate) trait DbConnection {
 
     /// Run a query that fetches a single value. `query` is a thunk that returns a query when
     /// called.
-    fn result<Q, U>(&mut self, query: impl Fn() -> Q) -> QueryResult<U>
+    fn result<Q, U>(&mut self, query: impl Fn() -> Q, query_id: &Uuid) -> QueryResult<U>
     where
         Q: diesel::query_builder::Query,
         Q: LoadQuery<'static, Self::Connection, U>,
         Q: QueryId + QueryFragment<Self::Backend>;
 
     /// Run a query that fetches multiple values. `query` is a thunk that returns a query when
     /// called.
-    fn results<Q, U>(&mut self, query: impl Fn() -> Q) -> QueryResult<Vec<U>>
+    fn results<Q, U>(&mut self, query: impl Fn() -> Q, query_id: &Uuid) -> QueryResult<Vec<U>>
     where
         Q: diesel::query_builder::Query,
         Q: LoadQuery<'static, Self::Connection, U>,
         Q: QueryId + QueryFragment<Self::Backend>;
 
     /// Helper to limit a query that fetches multiple values to return only its first value. `query`
     /// is a thunk that returns a query when called.
-    fn first<Q: LimitDsl, U>(&mut self, query: impl Fn() -> Q) -> QueryResult<U>
+    fn first<Q: LimitDsl, U>(&mut self, query: impl Fn() -> Q, query_id: &Uuid) -> QueryResult<U>
     where
         <Q as LimitDsl>::Output: diesel::query_builder::Query,
         <Q as LimitDsl>::Output: LoadQuery<'static, Self::Connection, U>,
         <Q as LimitDsl>::Output: QueryId + QueryFragment<Self::Backend>,
     {
-        self.result(move || query().limit(1i64))
+        self.result(move || query().limit(1i64), query_id)
     }
 }
diff --git a/crates/sui-graphql-rpc/src/data/pg.rs b/crates/sui-graphql-rpc/src/data/pg.rs
@@ -15,6 +15,7 @@ use diesel::{
 use sui_indexer::indexer_reader::IndexerReader;
 
 use tracing::error;
+use uuid::Uuid;
 
 pub(crate) struct PgExecutor {
     pub inner: IndexerReader,
@@ -94,23 +95,27 @@ impl<'c> super::DbConnection for PgConnection<'c> {
     type Connection = diesel::PgConnection;
     type Backend = Pg;
 
-    fn result<Q, U>(&mut self, query: impl Fn() -> Q) -> QueryResult<U>
+    fn result<Q, U>(&mut self, query: impl Fn() -> Q, query_id: &Uuid) -> QueryResult<U>
     where
         Q: diesel::query_builder::Query,
         Q: LoadQuery<'static, Self::Connection, U>,
         Q: QueryId + QueryFragment<Self::Backend>,
     {
-        query_cost::log(self.conn, self.max_cost, query());
+        if !query_id.is_nil() {
+            query_cost::log(self.conn, self.max_cost, query());
+        }
         query().get_result(self.conn)
     }
 
-    fn results<Q, U>(&mut self, query: impl Fn() -> Q) -> QueryResult<Vec<U>>
+    fn results<Q, U>(&mut self, query: impl Fn() -> Q, query_id: &Uuid) -> QueryResult<Vec<U>>
     where
         Q: diesel::query_builder::Query,
         Q: LoadQuery<'static, Self::Connection, U>,
         Q: QueryId + QueryFragment<Self::Backend>,
     {
-        query_cost::log(self.conn, self.max_cost, query());
+        if !query_id.is_nil() {
+            query_cost::log(self.conn, self.max_cost, query());
+        }
         query().get_results(self.conn)
     }
 }

diff --git a/crates/sui-graphql-rpc/src/types/cursor.rs b/crates/sui-graphql-rpc/src/types/cursor.rs
@@ -189,6 +189,7 @@ impl<C: CursorType + Eq + Clone + Send + Sync + 'static> Page<C> {
         &self,
         conn: &mut Conn<'_>,
         query: Q,
+        ctx: &Context<'_>,
     ) -> QueryResult<(bool, bool, impl Iterator<Item = T>)>
     where
         Q: Fn() -> Query<ST, T::Source, GB>,
@@ -221,7 +222,7 @@ impl<C: CursorType + Eq + Clone + Send + Sync + 'static> Page<C> {
             // Avoid the database roundtrip in the degenerate case.
             vec![]
         } else {
-            let mut results = conn.results(query)?;
+            let mut results = conn.results(query, ctx.data_unchecked())?;
             if !self.is_from_front() {
                 results.reverse();
             }