fix(users/Profpatsch/whatcd-resolver): fix postgres query log

The queries would not be interpolated anymore, because we didn’t pass
the thing down deep enough.

Also only init the `pgFormatPool` if we want to use the formatter,
this saves on a bunch of subprocesses.

Change-Id: I8d69ef5aab4d8eac1cbfb1c3991d4edaacba254f
Reviewed-on: https://cl.tvl.fyi/c/depot/+/12139
Reviewed-by: Profpatsch <mail@profpatsch.de>
Tested-by: BuildkiteCI
This commit is contained in:
Profpatsch 2024-08-06 11:42:38 +02:00
parent 2510cd6a5c
commit 13d79e04d8
3 changed files with 109 additions and 104 deletions

View file

@ -370,20 +370,19 @@ addErrorInformation msg io =
-- print the query that was run and the query parameters,
-- then rethrow inside an 'Error'.
handlePGException ::
forall a params tools m.
forall a params m.
( ToRow params,
MonadUnliftIO m,
MonadLogger m,
HasField "pgFormat" tools PgFormatPool
MonadLogger m
) =>
tools ->
PrettyPrintDatabaseQueries ->
Text ->
Query ->
-- | Depending on whether we used `format` or `formatMany`.
Either params (NonEmpty params) ->
IO a ->
Transaction m a
handlePGException tools queryType query' params io = do
handlePGException prettyQuery queryType query' params io = do
withRunInIO $ \unliftIO ->
io
`catches` [ Handler $ unliftIO . logQueryException @SqlError,
@ -397,9 +396,10 @@ handlePGException tools queryType query' params io = do
throwErr err = liftIO $ throwAsError $ prettyErrorTree $ nestedMultiError "A Postgres query failed" err
logQueryException :: (Exception e) => e -> Transaction m a
logQueryException exc = do
formattedQuery <- case params of
Left one -> pgFormatQuery' tools query' one
Right many -> pgFormatQueryMany' tools query' many
formattedQuery <-
case params of
Left one -> pgFormatQuery' prettyQuery query' one
Right many -> pgFormatQueryMany' prettyQuery query' many
throwErr
( singleError [fmt|Query Type: {queryType}|]
:| [ nestedError "Exception" (exc & showPretty & newError & singleError),
@ -533,55 +533,52 @@ runPGTransactionImpl zoom (Transaction transaction) = do
unliftIO $ runReaderT transaction conn
executeImpl ::
(ToRow params, MonadUnliftIO m, MonadLogger m, HasField "pgFormat" tools PgFormatPool, Otel.MonadTracer m) =>
m tools ->
(ToRow params, MonadUnliftIO m, MonadLogger m, Otel.MonadTracer m) =>
m (DebugLogDatabaseQueries, PrettyPrintDatabaseQueries) ->
Query ->
params ->
Transaction m (Label "numberOfRowsAffected" Natural)
{-# INLINE executeImpl #-}
executeImpl zoomTools zoomDbOptions qry params =
executeImpl zoomDbOptions qry params =
Otel.inSpan' "Postgres Query (execute)" Otel.defaultSpanArguments $ \span -> do
tools <- lift @Transaction zoomTools
(logDatabaseQueries, prettyQuery) <- lift @Transaction zoomDbOptions
traceQueryIfEnabled tools span logDatabaseQueries prettyQuery qry (HasSingleParam params)
traceQueryIfEnabled span logDatabaseQueries prettyQuery qry (HasSingleParam params)
conn <- Transaction ask
PG.execute conn qry params
& handlePGException tools "execute" qry (Left params)
& handlePGException prettyQuery "execute" qry (Left params)
>>= toNumberOfRowsAffected "executeImpl"
executeImpl_ ::
(MonadUnliftIO m, MonadLogger m, HasField "pgFormat" tools PgFormatPool, Otel.MonadTracer m) =>
m tools ->
( MonadUnliftIO m,
MonadLogger m,
Otel.MonadTracer m
) =>
m (DebugLogDatabaseQueries, PrettyPrintDatabaseQueries) ->
Query ->
Transaction m (Label "numberOfRowsAffected" Natural)
{-# INLINE executeImpl_ #-}
executeImpl_ zoomTools zoomDbOptions qry =
executeImpl_ zoomDbOptions qry =
Otel.inSpan' "Postgres Query (execute)" Otel.defaultSpanArguments $ \span -> do
tools <- lift @Transaction zoomTools
(logDatabaseQueries, prettyQuery) <- lift @Transaction zoomDbOptions
traceQueryIfEnabled @() tools span logDatabaseQueries prettyQuery qry HasNoParams
traceQueryIfEnabled @() span logDatabaseQueries prettyQuery qry HasNoParams
conn <- Transaction ask
PG.execute_ conn qry
& handlePGException tools "execute_" qry (Left ())
& handlePGException prettyQuery "execute_" qry (Left ())
>>= toNumberOfRowsAffected "executeImpl_"
executeManyImpl ::
(ToRow params, MonadUnliftIO m, MonadLogger m, HasField "pgFormat" tools PgFormatPool, Otel.MonadTracer m) =>
m tools ->
(ToRow params, MonadUnliftIO m, MonadLogger m, Otel.MonadTracer m) =>
m (DebugLogDatabaseQueries, PrettyPrintDatabaseQueries) ->
Query ->
NonEmpty params ->
Transaction m (Label "numberOfRowsAffected" Natural)
executeManyImpl zoomTools zoomDbOptions qry params =
executeManyImpl zoomDbOptions qry params =
Otel.inSpan' "Postgres Query (executeMany)" Otel.defaultSpanArguments $ \span -> do
tools <- lift @Transaction zoomTools
(logDatabaseQueries, prettyQuery) <- lift @Transaction zoomDbOptions
traceQueryIfEnabled tools span logDatabaseQueries prettyQuery qry (HasMultiParams params)
traceQueryIfEnabled span logDatabaseQueries prettyQuery qry (HasMultiParams params)
conn <- Transaction ask
PG.executeMany conn qry (params & toList)
& handlePGException tools "executeMany" qry (Right params)
& handlePGException prettyQuery "executeMany" qry (Right params)
>>= toNumberOfRowsAffected "executeManyImpl"
toNumberOfRowsAffected :: (MonadIO m) => Text -> Int64 -> m (Label "numberOfRowsAffected" Natural)
@ -595,31 +592,31 @@ toNumberOfRowsAffected functionName i64 =
<&> label @"numberOfRowsAffected"
executeManyReturningWithImpl ::
(ToRow params, MonadUnliftIO m, MonadLogger m, HasField "pgFormat" tools PgFormatPool, Otel.MonadTracer m) =>
m tools ->
( ToRow params,
MonadUnliftIO m,
MonadLogger m,
Otel.MonadTracer m
) =>
m (DebugLogDatabaseQueries, PrettyPrintDatabaseQueries) ->
Query ->
NonEmpty params ->
Decoder r ->
Transaction m [r]
{-# INLINE executeManyReturningWithImpl #-}
executeManyReturningWithImpl zoomTools zoomDbOptions qry params (Decoder fromRow) = do
executeManyReturningWithImpl zoomDbOptions qry params (Decoder fromRow) = do
Otel.inSpan' "Postgres Query (executeManyReturning)" Otel.defaultSpanArguments $ \span -> do
tools <- lift @Transaction zoomTools
(logDatabaseQueries, prettyQuery) <- lift @Transaction zoomDbOptions
traceQueryIfEnabled tools span logDatabaseQueries prettyQuery qry (HasMultiParams params)
traceQueryIfEnabled span logDatabaseQueries prettyQuery qry (HasMultiParams params)
conn <- Transaction ask
PG.returningWith fromRow conn qry (params & toList)
& handlePGException tools "executeManyReturning" qry (Right params)
& handlePGException prettyQuery "executeManyReturning" qry (Right params)
foldRowsWithAccImpl ::
( ToRow params,
MonadUnliftIO m,
MonadLogger m,
HasField "pgFormat" tools PgFormatPool,
Otel.MonadTracer m
) =>
m tools ->
m (DebugLogDatabaseQueries, PrettyPrintDatabaseQueries) ->
Query ->
params ->
@ -628,11 +625,10 @@ foldRowsWithAccImpl ::
(a -> row -> Transaction m a) ->
Transaction m a
{-# INLINE foldRowsWithAccImpl #-}
foldRowsWithAccImpl zoomTools zoomDbOptions qry params (Decoder rowParser) accumulator f = do
foldRowsWithAccImpl zoomDbOptions qry params (Decoder rowParser) accumulator f = do
Otel.inSpan' "Postgres Query (foldRowsWithAcc)" Otel.defaultSpanArguments $ \span -> do
tools <- lift @Transaction zoomTools
(logDatabaseQueries, prettyQuery) <- lift @Transaction zoomDbOptions
traceQueryIfEnabled tools span logDatabaseQueries prettyQuery qry (HasSingleParam params)
traceQueryIfEnabled span logDatabaseQueries prettyQuery qry (HasSingleParam params)
conn <- Transaction ask
withRunInIO
( \runInIO ->
@ -645,17 +641,18 @@ foldRowsWithAccImpl zoomTools zoomDbOptions qry params (Decoder rowParser) accum
params
accumulator
(\acc row -> runInIO $ f acc row)
& handlePGException tools "fold" qry (Left params)
& handlePGException prettyQuery "fold" qry (Left params)
& runInIO
)
pgFormatQueryNoParams' ::
(MonadIO m, MonadLogger m, HasField "pgFormat" tools PgFormatPool) =>
tools ->
(MonadIO m, MonadLogger m) =>
PrettyPrintDatabaseQueries ->
Query ->
Transaction m ByteString
pgFormatQueryNoParams' tools q =
lift $ pgFormatQueryByteString tools q.fromQuery
pgFormatQueryNoParams' prettyQuery q = case prettyQuery of
DontPrettyPrintDatabaseQueries -> pure q.fromQuery
PrettyPrintDatabaseQueries pool -> lift $ pgFormatQueryByteString pool q.fromQuery
pgFormatQuery ::
(ToRow params, MonadIO m) =>
@ -686,40 +683,36 @@ queryWithImpl ::
( ToRow params,
MonadUnliftIO m,
MonadLogger m,
HasField "pgFormat" tools PgFormatPool,
Otel.MonadTracer m
) =>
m tools ->
m (DebugLogDatabaseQueries, PrettyPrintDatabaseQueries) ->
Query ->
params ->
Decoder r ->
Transaction m [r]
{-# INLINE queryWithImpl #-}
queryWithImpl zoomTools zoomDbOptions qry params (Decoder fromRow) = do
queryWithImpl zoomDbOptions qry params (Decoder fromRow) = do
Otel.inSpan' "Postgres Query (queryWith)" Otel.defaultSpanArguments $ \span -> do
tools <- lift @Transaction zoomTools
(logDatabaseQueries, prettyQuery) <- lift @Transaction zoomDbOptions
traceQueryIfEnabled tools span logDatabaseQueries prettyQuery qry (HasSingleParam params)
traceQueryIfEnabled span logDatabaseQueries prettyQuery qry (HasSingleParam params)
conn <- Transaction ask
PG.queryWith fromRow conn qry params
& handlePGException tools "query" qry (Left params)
& handlePGException prettyQuery "query" qry (Left params)
queryWithImpl_ ::
( MonadUnliftIO m,
MonadLogger m,
HasField "pgFormat" tools PgFormatPool
MonadLogger m
) =>
m tools ->
m PrettyPrintDatabaseQueries ->
Query ->
Decoder r ->
Transaction m [r]
{-# INLINE queryWithImpl_ #-}
queryWithImpl_ zoomTools qry (Decoder fromRow) = do
tools <- lift @Transaction zoomTools
queryWithImpl_ zoomDbOptions qry (Decoder fromRow) = do
prettyQuery <- lift @Transaction zoomDbOptions
conn <- Transaction ask
liftIO (PG.queryWith_ fromRow conn qry)
& handlePGException tools "query" qry (Left ())
& handlePGException prettyQuery "query" qry (Left ())
data SingleRowError = SingleRowError
{ -- | How many columns were actually returned by the query
@ -733,30 +726,32 @@ instance Exception SingleRowError where
pgFormatQuery' ::
( MonadIO m,
ToRow params,
MonadLogger m,
HasField "pgFormat" tools PgFormatPool
MonadLogger m
) =>
tools ->
PrettyPrintDatabaseQueries ->
Query ->
params ->
Transaction m ByteString
pgFormatQuery' tools q p =
pgFormatQuery q p
>>= lift . pgFormatQueryByteString tools
pgFormatQuery' prettyQuery q p = case prettyQuery of
DontPrettyPrintDatabaseQueries -> pgFormatQuery q p
PrettyPrintDatabaseQueries pool ->
pgFormatQuery q p
>>= lift . pgFormatQueryByteString pool
pgFormatQueryMany' ::
( MonadIO m,
ToRow params,
MonadLogger m,
HasField "pgFormat" tools PgFormatPool
MonadLogger m
) =>
tools ->
PrettyPrintDatabaseQueries ->
Query ->
NonEmpty params ->
Transaction m ByteString
pgFormatQueryMany' tools q p =
pgFormatQueryMany q p
>>= lift . pgFormatQueryByteString tools
pgFormatQueryMany' prettyQuery q p = case prettyQuery of
DontPrettyPrintDatabaseQueries -> pgFormatQueryMany q p
PrettyPrintDatabaseQueries pool ->
pgFormatQueryMany q p
>>= lift . pgFormatQueryByteString pool
-- | Read the executable name "pg_format"
postgresToolsParser :: ToolParserT IO (Label "pgFormat" Tool)
@ -764,17 +759,16 @@ postgresToolsParser = label @"pgFormat" <$> readTool "pg_format"
pgFormatQueryByteString ::
( MonadIO m,
MonadLogger m,
HasField "pgFormat" tools PgFormatPool
MonadLogger m
) =>
tools ->
PgFormatPool ->
ByteString ->
m ByteString
pgFormatQueryByteString tools queryBytes = do
pgFormatQueryByteString pool queryBytes = do
res <-
liftIO $
runPgFormat
tools.pgFormat
pool
(queryBytes)
case res.exitCode of
ExitSuccess -> pure (res.formatted)
@ -832,8 +826,11 @@ data PrettyPrintDatabaseQueries
= -- | Do not pretty-print database querios
DontPrettyPrintDatabaseQueries
| -- | Pretty-print database queries, slow
PrettyPrintDatabaseQueries
deriving stock (Show, Enum, Bounded)
PrettyPrintDatabaseQueries PgFormatPool
instance Show PrettyPrintDatabaseQueries where
show DontPrettyPrintDatabaseQueries = "DontPrettyPrintDatabaseQueries"
show (PrettyPrintDatabaseQueries _) = "PrettyPrintDatabaseQueries"
data HasQueryParams param
= HasNoParams
@ -845,29 +842,25 @@ traceQueryIfEnabled ::
( ToRow params,
MonadUnliftIO m,
MonadLogger m,
HasField "pgFormat" tools PgFormatPool,
Otel.MonadTracer m
) =>
tools ->
Otel.Span ->
DebugLogDatabaseQueries ->
PrettyPrintDatabaseQueries ->
Query ->
HasQueryParams params ->
Transaction m ()
traceQueryIfEnabled tools span logDatabaseQueries prettyQuery qry params = do
traceQueryIfEnabled span logDatabaseQueries prettyQuery qry params = do
-- In case we have query logging enabled, we want to do that
let formattedQuery = case prettyQuery of
DontPrettyPrintDatabaseQueries -> pure qry.fromQuery
PrettyPrintDatabaseQueries -> do
withEvent
span
"Query Format start"
"Query Format end"
$ case params of
HasNoParams -> pgFormatQueryNoParams' tools qry
HasSingleParam p -> pgFormatQuery' tools qry p
HasMultiParams ps -> pgFormatQueryMany' tools qry ps
let formattedQuery =
withEvent
span
"Query Format start"
"Query Format end"
$ case params of
HasNoParams -> pgFormatQueryNoParams' prettyQuery qry
HasSingleParam p -> pgFormatQuery' prettyQuery qry p
HasMultiParams ps -> pgFormatQueryMany' prettyQuery qry ps
let doLog errs =
Otel.addAttributes
@ -884,7 +877,7 @@ traceQueryIfEnabled tools span logDatabaseQueries prettyQuery qry params = do
q <- formattedQuery
Otel.inSpan "Postgres EXPLAIN Query" Otel.defaultSpanArguments $ do
queryWithImpl_
(pure tools)
(pure prettyQuery)
( "EXPLAIN "
<> (
-- TODO: this is not nice, but the only way to get the `executeMany` form to work with this

View file

@ -25,10 +25,14 @@ import UnliftIO
import Prelude hiding (span)
data Context = Context
{ config :: T2 "logDatabaseQueries" DebugLogDatabaseQueries "prettyPrintDatabaseQueries" PrettyPrintDatabaseQueries,
{ pgConfig ::
T2
"logDatabaseQueries"
DebugLogDatabaseQueries
"prettyPrintDatabaseQueries"
PrettyPrintDatabaseQueries,
pgConnPool :: (Pool Postgres.Connection),
tracer :: Otel.Tracer,
pgFormat :: PgFormatPool,
pgConnPool :: Pool Postgres.Connection,
transmissionSessionId :: IORef (Maybe ByteString),
redactedApiKey :: ByteString
}
@ -147,17 +151,24 @@ recordException span dat = liftIO $ do
-- * Postgres
instance (MonadThrow m, MonadUnliftIO m) => MonadPostgres (AppT m) where
execute = executeImpl (AppT ask) dbConfig
executeMany = executeManyImpl (AppT ask) dbConfig
executeManyReturningWith = executeManyReturningWithImpl (AppT ask) dbConfig
queryWith = queryWithImpl (AppT ask) dbConfig
queryWith_ = queryWithImpl_ (AppT ask)
execute = executeImpl dbConfig
executeMany = executeManyImpl dbConfig
executeManyReturningWith = executeManyReturningWithImpl dbConfig
queryWith = queryWithImpl dbConfig
queryWith_ = queryWithImpl_ (dbConfig <&> snd)
foldRowsWithAcc = foldRowsWithAccImpl (AppT ask) dbConfig
foldRowsWithAcc = foldRowsWithAccImpl dbConfig
runTransaction = runPGTransaction
dbConfig :: (Monad m) => AppT m (DebugLogDatabaseQueries, PrettyPrintDatabaseQueries)
dbConfig = AppT $ asks (\c -> (c.config.logDatabaseQueries, c.config.prettyPrintDatabaseQueries))
dbConfig =
AppT $
asks
( \c ->
( c.pgConfig.logDatabaseQueries,
c.pgConfig.prettyPrintDatabaseQueries
)
)
runPGTransaction :: (MonadUnliftIO m) => Transaction (AppT m) a -> AppT m a
runPGTransaction (Transaction transaction) = do

View file

@ -776,12 +776,13 @@ httpTorrent span req =
runAppWith :: AppT IO a -> IO (Either TmpPg.StartError a)
runAppWith appT = withTracer $ \tracer -> withDb $ \db -> do
tool <- readTools (label @"toolsEnvVar" "WHATCD_RESOLVER_TOOLS") (readTool "pg_format")
pgFormat <- initPgFormatPool (label @"pgFormat" tool)
prettyPrintDatabaseQueries <-
Env.lookupEnv "WHATCD_RESOLVER_PRETTY_PRINT_DATABASE_QUERIES" <&> \case
Just _ -> PrettyPrintDatabaseQueries
Nothing -> DontPrettyPrintDatabaseQueries
let config =
Env.lookupEnv "WHATCD_RESOLVER_PRETTY_PRINT_DATABASE_QUERIES" >>= \case
Nothing -> pure DontPrettyPrintDatabaseQueries
Just _ -> do
pgFormat <- initPgFormatPool (label @"pgFormat" tool)
pure $ PrettyPrintDatabaseQueries pgFormat
let pgConfig =
T2
(label @"logDatabaseQueries" LogDatabaseQueries)
(label @"prettyPrintDatabaseQueries" prettyPrintDatabaseQueries)
@ -800,7 +801,7 @@ runAppWith appT = withTracer $ \tracer -> withDb $ \db -> do
logInfo "WHATCD_RESOLVER_REDACTED_API_KEY was not set, trying pass"
runCommandExpect0 "pass" ["internet/redacted/api-keys/whatcd-resolver"]
let newAppT = do
logInfo [fmt|Running with config: {showPretty config}|]
logInfo [fmt|Running with config: {showPretty pgConfig}|]
logInfo [fmt|Connected to database at {db & TmpPg.toDataDirectory} on socket {db & TmpPg.toConnectionString}|]
appT
runReaderT newAppT.unAppT Context {..}