From 277dcd79b2ace0ab74efe7c997bbb0b8d8349c89 Mon Sep 17 00:00:00 2001
From: Phil Freeman <paf31@cantab.net>
Date: Fri, 7 Aug 2020 14:50:49 -0700
Subject: [PATCH 1/2] server: set hasura.tracecontext in RQL mutations [#5542]

---
 server/src-lib/Hasura/App.hs                         | 4 ++--
 server/src-lib/Hasura/GraphQL/Transport/WebSocket.hs | 3 ++-
 server/src-lib/Hasura/RQL/Types/Run.hs               | 6 ++++--
 server/src-lib/Hasura/Server/API/Query.hs            | 3 ++-
 server/src-lib/Hasura/Server/SchemaUpdate.hs         | 2 +-
 5 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/server/src-lib/Hasura/App.hs b/server/src-lib/Hasura/App.hs
index 50cbc0775d535..1504d85f458e3 100644
--- a/server/src-lib/Hasura/App.hs
+++ b/server/src-lib/Hasura/App.hs
@@ -235,7 +235,7 @@ migrateCatalogSchema env logger pool httpManager sqlGenCtx = do
   let pgExecCtx = mkPGExecCtx Q.Serializable pool
       adminRunCtx = RunCtx adminUserInfo httpManager sqlGenCtx
   currentTime <- liftIO Clock.getCurrentTime
-  initialiseResult <- runExceptT $ peelRun adminRunCtx pgExecCtx Q.ReadWrite $
+  initialiseResult <- runExceptT $ peelRun adminRunCtx pgExecCtx Q.ReadWrite Nothing $
     (,) <$> migrateCatalog env currentTime <*> liftTx fetchLastUpdate
 
   ((migrationResult, schemaCache), lastUpdateEvent) <-
@@ -586,7 +586,7 @@ runAsAdmin
 runAsAdmin pool sqlGenCtx httpManager m = do
   let runCtx = RunCtx adminUserInfo httpManager sqlGenCtx
       pgCtx = mkPGExecCtx Q.Serializable pool
-  runExceptT $ peelRun runCtx pgCtx Q.ReadWrite m
+  runExceptT $ peelRun runCtx pgCtx Q.ReadWrite Nothing m
 
 execQuery
   :: ( HasVersion
diff --git a/server/src-lib/Hasura/GraphQL/Transport/WebSocket.hs b/server/src-lib/Hasura/GraphQL/Transport/WebSocket.hs
index d56e5523a0ee1..6587b6649e130 100644
--- a/server/src-lib/Hasura/GraphQL/Transport/WebSocket.hs
+++ b/server/src-lib/Hasura/GraphQL/Transport/WebSocket.hs
@@ -368,8 +368,9 @@ onStart env serverEnv wsConn (StartMsg opId q) = catchAndIgnore $ do
           Tracing.interpTraceT id $ executeQuery queryParsed asts genSql pgExecCtx Q.ReadOnly opTx
       -- Response headers discarded over websockets
       E.ExOpMutation _ opTx -> Tracing.trace "pg" do
+        ctx <- Tracing.currentContext
         execQueryOrMut Telem.Mutation Nothing $
-          Tracing.interpTraceT (runLazyTx pgExecCtx Q.ReadWrite . withUserInfo userInfo) opTx
+          Tracing.interpTraceT (runLazyTx pgExecCtx Q.ReadWrite . withTraceContext ctx . withUserInfo userInfo) opTx
       E.ExOpSubs lqOp -> do
         -- log the graphql query
         logQueryLog logger query Nothing reqId
diff --git a/server/src-lib/Hasura/RQL/Types/Run.hs b/server/src-lib/Hasura/RQL/Types/Run.hs
index bbbc35e38282e..b0b285e40c52c 100644
--- a/server/src-lib/Hasura/RQL/Types/Run.hs
+++ b/server/src-lib/Hasura/RQL/Types/Run.hs
@@ -16,6 +16,7 @@ import           Control.Monad.Trans.Control (MonadBaseControl)
 import           Control.Monad.Unique
 
 import           Hasura.RQL.Types
+import qualified Hasura.Tracing              as Tracing
 
 data RunCtx
   = RunCtx
@@ -50,7 +51,8 @@ peelRun
   => RunCtx
   -> PGExecCtx
   -> Q.TxAccess
+  -> Maybe Tracing.TraceContext
   -> Run a
   -> ExceptT QErr m a
-peelRun runCtx@(RunCtx userInfo _ _) pgExecCtx txAccess (Run m) =
-  runLazyTx pgExecCtx txAccess $ withUserInfo userInfo $ runReaderT m runCtx
+peelRun runCtx@(RunCtx userInfo _ _) pgExecCtx txAccess ctx (Run m) =
+  runLazyTx pgExecCtx txAccess $ maybe id withTraceContext ctx $ withUserInfo userInfo $ runReaderT m runCtx
diff --git a/server/src-lib/Hasura/Server/API/Query.hs b/server/src-lib/Hasura/Server/API/Query.hs
index 357f00901bdef..fa4dc8f03d162 100644
--- a/server/src-lib/Hasura/Server/API/Query.hs
+++ b/server/src-lib/Hasura/Server/API/Query.hs
@@ -198,10 +198,11 @@ runQuery
   -> SQLGenCtx -> SystemDefined -> RQLQuery -> m (EncJSON, RebuildableSchemaCache Run)
 runQuery env pgExecCtx instanceId userInfo sc hMgr sqlGenCtx systemDefined query = do
   accessMode <- getQueryAccessMode query
+  traceCtx <- Tracing.currentContext
   resE <- runQueryM env query & Tracing.interpTraceT \x -> do
     a <- x & runHasSystemDefinedT systemDefined
            & runCacheRWT sc
-           & peelRun runCtx pgExecCtx accessMode
+           & peelRun runCtx pgExecCtx accessMode (Just traceCtx)
            & runExceptT
            & liftIO
     pure (either 
diff --git a/server/src-lib/Hasura/Server/SchemaUpdate.hs b/server/src-lib/Hasura/Server/SchemaUpdate.hs
index a7a7fc2e73c17..7c2d200e35203 100644
--- a/server/src-lib/Hasura/Server/SchemaUpdate.hs
+++ b/server/src-lib/Hasura/Server/SchemaUpdate.hs
@@ -221,7 +221,7 @@ refreshSchemaCache sqlGenCtx pool logger httpManager cacheRef invalidations thre
     rebuildableCache <- fst <$> liftIO (readIORef $ _scrCache cacheRef)
     ((), cache, _) <- buildSchemaCacheWithOptions CatalogSync invalidations
       & runCacheRWT rebuildableCache
-      & peelRun runCtx pgCtx PG.ReadWrite
+      & peelRun runCtx pgCtx PG.ReadWrite Nothing
     pure ((), cache)
   case resE of
     Left e   -> logError logger threadType $ TEQueryError e

From ac5e6105a9598b7e0d7ac42ce8f9d5d6e6f28825 Mon Sep 17 00:00:00 2001
From: Phil Freeman <paf31@cantab.net>
Date: Mon, 10 Aug 2020 09:09:17 -0700
Subject: [PATCH 2/2] Update test suite

---
 server/src-test/Main.hs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/server/src-test/Main.hs b/server/src-test/Main.hs
index f9e95964ab307..b37db2ce08886 100644
--- a/server/src-test/Main.hs
+++ b/server/src-test/Main.hs
@@ -83,7 +83,7 @@ buildPostgresSpecs pgConnOptions = do
 
             runAsAdmin :: Run a -> IO a
             runAsAdmin =
-                  peelRun runContext pgContext Q.ReadWrite
+                  peelRun runContext pgContext Q.ReadWrite Nothing
               >>> runExceptT
               >=> flip onLeft printErrJExit