Skip to content

Commit d70d968

Browse files
develop7develop7
committed
add: propagate trace state to pg_tracing via GUC
1 parent 555b7f5 commit d70d968

5 files changed

Lines changed: 47 additions & 25 deletions

File tree

‎postgrest.cabal‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -129,6 +129,7 @@ library
129129
-- ^ this is due to hs-otel-sdk is not reexporting getTracerTracerProvider
130130
-- needed to initialize OpenTelemetry.middleware
131131
, hs-opentelemetry-utils-exceptions
132+
, hs-opentelemetry-propagator-w3c
132133
, insert-ordered-containers >= 0.2.2 && < 0.3
133134
, jose-jwt >= 0.9.6 && < 0.11
134135
, lens >= 4.14 && < 5.4

‎src/PostgREST/App.hs‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -184,7 +184,9 @@ postgrestResponse appState conf@AppConfig{..} maybeSchemaCache authResult@AuthRe
184184
(parseTime, apiReq@ApiRequest{..}) <- withOTel "parse" $ withTiming $ liftEither . mapLeft Error.ApiRequestErr $ ApiRequest.userApiRequest conf prefs req body
185185
(planTime, plan) <- withOTel "plan" $ withTiming $ liftEither $ Plan.actionPlan iAction conf apiReq sCache
186186

187-
let mainQ = Query.mainQuery plan conf apiReq authResult configDbPreRequest
187+
traceContext <- lift OTel.renderTraceContext
188+
189+
let mainQ = Query.mainQuery plan conf apiReq authResult configDbPreRequest traceContext
188190
tx = MainTx.mainTx mainQ conf authResult apiReq plan sCache
189191
obsQuery s = when configLogQuery $ observer $ QueryObs mainQ s
190192

‎src/PostgREST/OpenTelemetry.hs‎

Lines changed: 35 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -13,23 +13,28 @@ In order produced spans to have correct code locations, all the functions across
1313
`inSpanM` call must have `HasCallStack` constraint, because
1414
[GHC is never inferring it](https://downloads.haskell.org/ghc/9.8.4/docs/users_guide/exts/callstack.html) for us.
1515
-}
16-
module PostgREST.OpenTelemetry (Tracer, withTracer, middleware, inSpanM) where
16+
module PostgREST.OpenTelemetry (Tracer, withTracer, middleware, inSpanM, renderTraceContext) where
1717

18-
import Control.Monad.Catch (MonadMask)
19-
import Network.Wai (Middleware)
20-
import OpenTelemetry.Attributes (emptyAttributes)
21-
import OpenTelemetry.Instrumentation.Wai (newOpenTelemetryWaiMiddleware')
22-
import OpenTelemetry.Trace (InstrumentationLibrary (..),
23-
SpanArguments, Tracer,
24-
initializeGlobalTracerProvider,
25-
makeTracer,
26-
shutdownTracerProvider,
27-
tracerOptions)
28-
import OpenTelemetry.Trace.Core (getTracerTracerProvider)
29-
import OpenTelemetry.Utils.Exceptions (inSpanM'')
30-
import PostgREST.AppState (AppState, getOTelTracer)
31-
import PostgREST.Version (prettyVersion)
32-
import Protolude
18+
import Control.Monad.Catch (MonadMask)
19+
import Network.Wai (Middleware)
20+
import OpenTelemetry.Attributes (emptyAttributes)
21+
import qualified OpenTelemetry.Context as Context
22+
import OpenTelemetry.Context.ThreadLocal (getContext)
23+
import OpenTelemetry.Instrumentation.Wai (newOpenTelemetryWaiMiddleware')
24+
import OpenTelemetry.Propagator.W3CTraceContext (encodeSpanContext)
25+
import OpenTelemetry.Trace (InstrumentationLibrary (..),
26+
SpanArguments,
27+
Tracer,
28+
initializeGlobalTracerProvider,
29+
makeTracer,
30+
shutdownTracerProvider,
31+
tracerOptions)
32+
import OpenTelemetry.Trace.Core (getTracerTracerProvider)
33+
import OpenTelemetry.Utils.Exceptions (inSpanM'')
34+
import PostgREST.AppState (AppState,
35+
getOTelTracer)
36+
import PostgREST.Version (prettyVersion)
37+
import Protolude
3338

3439
{- | Wrap user's code with OpenTelemetry Tracer, initializing it with sensible defaults -}
3540
withTracer :: (Tracer -> IO c) -> IO c
@@ -69,3 +74,17 @@ inSpanM
6974
-> m a
7075
inSpanM (Just t) n args m = inSpanM'' t callStack n args (const m)
7176
inSpanM Nothing _ _ m = m
77+
78+
renderTraceContext :: IO (Maybe ByteString)
79+
renderTraceContext = do
80+
-- Render TraceContext in W3C Trace Context format
81+
-- https://www.w3.org/TR/trace-context/#traceparent-header
82+
-- https://www.w3.org/TR/trace-context/#tracestate-header
83+
--
84+
ctx <- getContext
85+
case Context.lookupSpan ctx of
86+
Nothing -> pure Nothing
87+
Just span -> do
88+
(tParent, _tState) <- encodeSpanContext span
89+
let traceContext = "traceparent='" <> tParent <> "'"
90+
pure $ Just traceContext

‎src/PostgREST/Query.hs‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -41,10 +41,10 @@ data MainQuery = MainQuery
4141
, mqExplain :: Maybe SQL.Snippet -- ^ the explain query that gets generated for the "Prefer: count=estimated" case
4242
}
4343

44-
mainQuery :: ActionPlan -> AppConfig -> ApiRequest -> AuthResult -> Maybe QualifiedIdentifier -> MainQuery
45-
mainQuery (NoDb _) _ _ _ _ = MainQuery mempty Nothing mempty (mempty, mempty, mempty) mempty
46-
mainQuery (Db plan) conf@AppConfig{..} apiReq@ApiRequest{iTopLevelRange=range, iPreferences=Preferences{..}} authRes preReq =
47-
let genQ = MainQuery (PreQuery.txVarQuery plan conf authRes apiReq) (PreQuery.preReqQuery <$> preReq) in
44+
mainQuery :: ActionPlan -> AppConfig -> ApiRequest -> AuthResult -> Maybe QualifiedIdentifier -> Maybe ByteString -> MainQuery
45+
mainQuery (NoDb _) _ _ _ _ _ = MainQuery mempty Nothing mempty (mempty, mempty, mempty) mempty
46+
mainQuery (Db plan) conf@AppConfig{..} apiReq@ApiRequest{iTopLevelRange=range, iPreferences=Preferences{..}} authRes preReq traceContext =
47+
let genQ = MainQuery (PreQuery.txVarQuery plan conf authRes apiReq traceContext ) (PreQuery.preReqQuery <$> preReq) in
4848
case plan of
4949
DbCrud _ WrappedReadPlan{..} ->
5050
let countQuery = QueryBuilder.readPlanToCountQuery wrReadPlan in

‎src/PostgREST/Query/PreQuery.hs‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,6 @@ import qualified Data.HashMap.Strict as HM
1515
import qualified Hasql.DynamicStatements.Snippet as SQL hiding (sql)
1616

1717

18-
1918
import PostgREST.ApiRequest (ApiRequest (..))
2019
import PostgREST.ApiRequest.Preferences (PreferTimezone (..),
2120
Preferences (..))
@@ -34,12 +33,12 @@ import PostgREST.SchemaCache.Routine (Routine (..))
3433
import Protolude hiding (Handler)
3534

3635
-- sets transaction variables
37-
txVarQuery :: DbActionPlan -> AppConfig -> AuthResult -> ApiRequest -> SQL.Snippet
38-
txVarQuery dbActPlan AppConfig{..} AuthResult{..} ApiRequest{..} =
36+
txVarQuery :: DbActionPlan -> AppConfig -> AuthResult -> ApiRequest -> Maybe ByteString -> SQL.Snippet
37+
txVarQuery dbActPlan AppConfig{..} AuthResult{..} ApiRequest{..} traceContext =
3938
-- To ensure `GRANT SET ON PARAMETER <superuser_setting> TO authenticator` works, the role settings must be set before the impersonated role.
4039
-- Otherwise the GRANT SET would have to be applied to the impersonated role. See https://github.com/PostgREST/postgrest/issues/3045
4140
"select " <> intercalateSnippet ", " (
42-
searchPathSql : roleSettingsSql ++ roleSql ++ claimsSql ++ [methodSql, pathSql] ++ headersSql ++ cookiesSql ++ timezoneSql ++ funcSettingsSql ++ appSettingsSql
41+
searchPathSql : roleSettingsSql ++ roleSql ++ claimsSql ++ [methodSql, pathSql] ++ headersSql ++ cookiesSql ++ timezoneSql ++ funcSettingsSql ++ appSettingsSql ++ traceContextSql
4342
)
4443
where
4544
methodSql = setConfigWithConstantName ("request.method", iMethod)
@@ -48,6 +47,7 @@ txVarQuery dbActPlan AppConfig{..} AuthResult{..} ApiRequest{..} =
4847
cookiesSql = setConfigWithConstantNameJSON "request.cookies" iCookies
4948
claimsSql = [setConfigWithConstantName ("request.jwt.claims", LBS.toStrict $ JSON.encode authClaims)]
5049
roleSql = [setConfigWithConstantName ("role", authRole)]
50+
traceContextSql = maybe mempty (\tc -> [setConfigWithConstantName ("pg_tracing.trace_context", tc)]) traceContext
5151
roleSettingsSql = setConfigWithDynamicName <$> HM.toList (fromMaybe mempty $ HM.lookup authRole configRoleSettings)
5252
appSettingsSql = setConfigWithDynamicName . join bimap toUtf8 <$> configAppSettings
5353
timezoneSql = maybe mempty (\(PreferTimezone tz) -> [setConfigWithConstantName ("timezone", tz)]) $ preferTimezone iPreferences

0 commit comments

Comments
 (0)