server: set tracecontext and userInfo for DML actions on PG sources

https://github.com/hasura/graphql-engine-mono/pull/2174

GitOrigin-RevId: e8cfb4e330938e7dfb7232e58d2c1fc07bf97896

CHANGELOG.md

@@ -4,6 +4,7 @@
 
 (Add entries below in the order of server, console, cli, docs, others)
 
+- server: set `tracecontext` and `userInfo` for DML actions on Postgres sources
 - cli: add progress bar for `migrate apply` command (#4795)
 
 ## v2.0.8

server/src-lib/Hasura/Backends/Postgres/Connection.hs

@@ -9,6 +9,7 @@ module Hasura.Backends.Postgres.Connection
 
   , runLazyTx
   , runQueryTx
+  , runQueryTxWithCtx
   , withUserInfo
   , withTraceContext
   , setHeadersTx
@@ -136,6 +137,18 @@ runQueryTx
 runQueryTx pgExecCtx ltx =
   liftEither =<< liftIO (runExceptT $ _pecRunReadNoTx pgExecCtx (unLazyTxT ltx))
 
+-- NOTE: Same warning as 'runQueryTx' applies here.
+-- This variant of 'runQueryTx' allows passing the `userInfo` context and `tracecontext`.
+runQueryTxWithCtx
+  :: (MonadIO m, MonadError QErr m)
+  => UserInfo
+  -> Tracing.TraceContext
+  -> PGExecCtx
+  -> LazyTxT QErr IO a
+  -> m a
+runQueryTxWithCtx userInfo traceCtx pgExecCtx =
+  runQueryTx pgExecCtx . withUserInfo userInfo . withTraceContext traceCtx
+
 setHeadersTx :: (MonadIO m) => SessionVariables -> Q.TxET QErr m ()
 setHeadersTx session = do
   Q.unitQE defaultTxErrorHandler setSess () False

server/src-lib/Hasura/Backends/Postgres/DDL/RunSQL.hs

@@ -16,6 +16,7 @@ import           Data.Aeson
 import           Data.Text.Extended
 
 import qualified Hasura.SQL.AnyBackend               as AB
+import qualified Hasura.Tracing                      as Tracing
 
 import           Hasura.Backends.Postgres.DDL.Source (ToMetadataFetchQuery, fetchFunctionMetadata,
                                                       fetchTableMetadata)
@@ -29,8 +30,9 @@ import           Hasura.RQL.DDL.Schema.Common
 import           Hasura.RQL.DDL.Schema.Diff
 import           Hasura.RQL.Types                    hiding (ConstraintName, fmFunction,
                                                       tmComputedFields, tmTable)
+import           Hasura.RQL.Types.Run
 import           Hasura.Server.Utils                 (quoteRegex)
-
+import           Hasura.Session
 
 data RunSQL
   = RunSQL
@@ -140,17 +142,25 @@ runRunSQL
      , MonadBaseControl IO m
      , MonadError QErr m
      , MonadIO m
+     , Tracing.MonadTrace m
+     , UserInfoM m
      )
   => RunSQL
   -> m EncJSON
-runRunSQL q@RunSQL {..}
-  -- see Note [Checking metadata consistency in run_sql]
-  | isSchemaCacheBuildRequiredRunSQL q
-  = withMetadataCheck @pgKind rSource rCascade rTxAccessMode $ execRawSQL rSql
-  | otherwise
-  = askSourceConfig @('Postgres pgKind) rSource >>= \sourceConfig ->
-      liftEitherM $ runExceptT $
-      runLazyTx (_pscExecCtx sourceConfig) rTxAccessMode $ execRawSQL rSql
+runRunSQL q@RunSQL{..} = do
+  sourceConfig <- askSourceConfig @('Postgres pgKind) rSource
+  traceCtx <- Tracing.currentContext
+  userInfo <- askUserInfo
+  let pgExecCtx = _pscExecCtx sourceConfig
+  if (isSchemaCacheBuildRequiredRunSQL q)
+    then do
+      -- see Note [Checking metadata consistency in run_sql]
+      withMetadataCheck @pgKind rSource rCascade rTxAccessMode
+        $ withTraceContext traceCtx
+        $ withUserInfo userInfo
+        $ execRawSQL rSql
+    else do
+      runQueryLazyTx pgExecCtx rTxAccessMode $ execRawSQL rSql
   where
     execRawSQL :: (MonadTx n) => Text -> n EncJSON
     execRawSQL =

server/src-lib/Hasura/Backends/Postgres/Instances/Transport.hs

@@ -62,11 +62,13 @@ runPGQuery
   -> Maybe EQ.PreparedSql
   -> m (DiffTime, EncJSON)
   -- ^ Also return the time spent in the PG query; for telemetry.
-runPGQuery reqId query fieldName _userInfo logger sourceConfig tx genSql = do
+runPGQuery reqId query fieldName userInfo logger sourceConfig tx genSql = do
   -- log the generated SQL and the graphql query
+  traceCtx <- currentContext
   logQueryLog logger $ mkQueryLog query fieldName genSql reqId
   withElapsedTime $ trace ("Postgres Query for root field " <>> fieldName) $
-    Tracing.interpTraceT id $ hoist (runQueryTx $ _pscExecCtx sourceConfig) tx
+    Tracing.interpTraceT id $
+    hoist (runQueryTxWithCtx userInfo traceCtx $ _pscExecCtx sourceConfig) tx
 
 runPGMutation
   :: ( MonadIO m
@@ -98,13 +100,13 @@ runPGMutation reqId query fieldName userInfo logger sourceConfig tx _genSql =  d
       ) tx
 
 runPGSubscription
-  :: ( MonadIO m
-     )
+  :: MonadIO m
   => SourceConfig ('Postgres pgKind)
   -> MultiplexedQuery ('Postgres pgKind)
   -> [(CohortId, CohortVariables)]
   -> m (DiffTime, Either QErr [(CohortId, B.ByteString)])
-runPGSubscription sourceConfig query variables = withElapsedTime
+runPGSubscription sourceConfig query variables =
+  withElapsedTime
   $ runExceptT
   $ runQueryTx (_pscExecCtx sourceConfig)
   $ PGL.executeMultiplexedQuery query variables

server/src-lib/Hasura/GraphQL/Transport/Backend.hs

@@ -57,8 +57,7 @@ class BackendExecute b => BackendTransport (b :: BackendType) where
     -> m (DiffTime, EncJSON)
   runDBSubscription
     :: forall m
-     . ( MonadIO m
-       )
+     . MonadIO m
     => SourceConfig b
     -> MultiplexedQuery b
     -> [(CohortId, CohortVariables)]