Skip to content

Commit 9c9d982

Browse files
develop7develop7
committed
add: propagate trace state to pg_tracing via GUC
1 parent 4cba0e4 commit 9c9d982

5 files changed

Lines changed: 45 additions & 25 deletions

File tree

‎postgrest.cabal‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -134,6 +134,7 @@ library
134134
-- ^ this is due to hs-otel-sdk is not reexporting getTracerTracerProvider
135135
-- needed to initialize OpenTelemetry.middleware
136136
, hs-opentelemetry-utils-exceptions
137+
, hs-opentelemetry-propagator-w3c
137138
, insert-ordered-containers >= 0.2.2 && < 0.3
138139
, jose-jwt >= 0.9.6 && < 0.11
139140
, lens >= 4.14 && < 5.4

‎src/PostgREST/App.hs‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -192,7 +192,9 @@ postgrestResponse appState conf@AppConfig{..} maybeSchemaCache jwtTime authResul
192192
(parseTime, apiReq@ApiRequest{..}) <- withRequestPhase "parse" $ liftEither . mapLeft Error.ApiRequestErr $ ApiRequest.userApiRequest conf prefs req body
193193
(planTime, plan) <- withRequestPhase "plan" $ liftEither $ Plan.actionPlan iAction conf apiReq sCache
194194

195-
let mainQ = Query.mainQuery plan conf apiReq authResult configDbPreRequest
195+
traceContext <- lift OTel.renderTraceContext
196+
197+
let mainQ = Query.mainQuery plan conf apiReq authResult configDbPreRequest traceContext
196198
tx = MainTx.mainTx mainQ conf authResult apiReq plan sCache
197199
obsQuery s = when configLogQuery $ observer $ QueryObs mainQ s
198200

‎src/PostgREST/OpenTelemetry.hs‎

Lines changed: 33 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -13,23 +13,27 @@ In order for produced spans to have correct code locations, all the functions ac
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, inSpanM, inSpanMDefault, oTelMiddleware) where
16+
module PostgREST.OpenTelemetry (Tracer, withTracer, inSpanM, inSpanMDefault, oTelMiddleware, 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-
defaultSpanArguments,
25-
initializeGlobalTracerProvider,
26-
makeTracer,
27-
shutdownTracerProvider,
28-
tracerOptions)
29-
import OpenTelemetry.Trace.Core (getTracerTracerProvider)
30-
import OpenTelemetry.Utils.Exceptions (inSpanM'')
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+
defaultSpanArguments,
29+
initializeGlobalTracerProvider,
30+
makeTracer,
31+
shutdownTracerProvider,
32+
tracerOptions)
33+
import OpenTelemetry.Trace.Core (getTracerTracerProvider)
34+
import OpenTelemetry.Utils.Exceptions (inSpanM'')
35+
import PostgREST.Version (prettyVersion)
36+
import Protolude
3337

3438
-- | Wrap user's code with OpenTelemetry Tracer, initializing it with sensible defaults
3539
withTracer :: (Tracer -> IO c) -> IO c
@@ -69,3 +73,16 @@ oTelMiddleware Nothing = identity
6973

7074
inSpanMDefault :: (MonadIO m, MonadMask m, HasCallStack) => Maybe Tracer -> Text -> m a -> m a
7175
inSpanMDefault t name = inSpanM t name defaultSpanArguments
76+
renderTraceContext :: IO (Maybe ByteString)
77+
renderTraceContext = do
78+
-- Render TraceContext in W3C Trace Context format
79+
-- https://www.w3.org/TR/trace-context/#traceparent-header
80+
-- https://www.w3.org/TR/trace-context/#tracestate-header
81+
--
82+
ctx <- getContext
83+
case Context.lookupSpan ctx of
84+
Nothing -> pure Nothing
85+
Just span -> do
86+
(tParent, _tState) <- encodeSpanContext span
87+
let traceContext = "traceparent='" <> tParent <> "'"
88+
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
@@ -16,7 +16,6 @@ import qualified Data.HashMap.Strict as HM
1616
import qualified Hasql.DynamicStatements.Snippet as SQL hiding (sql)
1717

1818

19-
2019
import PostgREST.ApiRequest (ApiRequest (..))
2120
import PostgREST.ApiRequest.Preferences (PreferTimezone (..),
2221
Preferences (..))
@@ -35,12 +34,12 @@ import PostgREST.SchemaCache.Routine (Routine (..))
3534
import Protolude hiding (Handler)
3635

3736
-- sets transaction variables
38-
txVarQuery :: DbActionPlan -> AppConfig -> AuthResult -> ApiRequest -> SQL.Snippet
39-
txVarQuery dbActPlan AppConfig{..} AuthResult{..} ApiRequest{..} =
37+
txVarQuery :: DbActionPlan -> AppConfig -> AuthResult -> ApiRequest -> Maybe ByteString -> SQL.Snippet
38+
txVarQuery dbActPlan AppConfig{..} AuthResult{..} ApiRequest{..} traceContext =
4039
-- To ensure `GRANT SET ON PARAMETER <superuser_setting> TO authenticator` works, the role settings must be set before the impersonated role.
4140
-- Otherwise the GRANT SET would have to be applied to the impersonated role. See https://github.com/PostgREST/postgrest/issues/3045
4241
"select " <> intercalateSnippet ", " (
43-
searchPathSql : roleSettingsSql ++ roleSql ++ claimsSql ++ [methodSql, pathSql] ++ headersSql ++ cookiesSql ++ timezoneSql ++ funcSettingsSql ++ appSettingsSql
42+
searchPathSql : roleSettingsSql ++ roleSql ++ claimsSql ++ [methodSql, pathSql] ++ headersSql ++ cookiesSql ++ timezoneSql ++ funcSettingsSql ++ appSettingsSql ++ traceContextSql
4443
)
4544
where
4645
methodSql = setConfigWithConstantName ("request.method", iMethod)
@@ -52,6 +51,7 @@ txVarQuery dbActPlan AppConfig{..} AuthResult{..} ApiRequest{..} =
5251
claims = authClaims & KM.insert "role" (JSON.String $ decodeUtf8 authRole) -- insert "role" to claims as well
5352

5453
roleSql = [setConfigWithConstantName ("role", authRole)]
54+
traceContextSql = maybe mempty (\tc -> [setConfigWithConstantName ("pg_tracing.trace_context", tc)]) traceContext
5555
roleSettingsSql = setConfigWithDynamicName <$> HM.toList (fromMaybe mempty $ HM.lookup authRole configRoleSettings)
5656
appSettingsSql = setConfigWithDynamicName . join bimap toUtf8 <$> configAppSettings
5757
timezoneSql = maybe mempty (\(PreferTimezone tz) -> [setConfigWithConstantName ("timezone", tz)]) $ preferTimezone iPreferences

0 commit comments

Comments
 (0)