diff options
Diffstat (limited to 'users/Profpatsch/my-prelude/src/Postgres/MonadPostgres.hs')
| -rw-r--r-- | users/Profpatsch/my-prelude/src/Postgres/MonadPostgres.hs | 276 |
1 files changed, 175 insertions, 101 deletions
diff --git a/users/Profpatsch/my-prelude/src/Postgres/MonadPostgres.hs b/users/Profpatsch/my-prelude/src/Postgres/MonadPostgres.hs index 2c9a48d134ef..87928678a052 100644 --- a/users/Profpatsch/my-prelude/src/Postgres/MonadPostgres.hs +++ b/users/Profpatsch/my-prelude/src/Postgres/MonadPostgres.hs @@ -34,10 +34,11 @@ import Database.PostgreSQL.Simple qualified as Postgres import Database.PostgreSQL.Simple.FromRow qualified as PG import Database.PostgreSQL.Simple.ToField (ToField) import Database.PostgreSQL.Simple.ToRow (ToRow (toRow)) -import Database.PostgreSQL.Simple.Types (Query (..)) +import Database.PostgreSQL.Simple.Types (PGArray (PGArray), Query (..)) import GHC.IO.Handle (Handle) import GHC.Records (getField) import Label +import Language.Haskell.TH.Quote (QuasiQuoter) import OpenTelemetry.Trace.Core (NewEvent (newEventName)) import OpenTelemetry.Trace.Core qualified as Otel hiding (inSpan, inSpan') import OpenTelemetry.Trace.Monad qualified as Otel @@ -45,6 +46,7 @@ import PossehlAnalyticsPrelude import Postgres.Decoder import Postgres.Decoder qualified as Dec import Pretty (showPretty) +import PyF qualified import Seconds import System.Exit (ExitCode (..)) import Tool @@ -140,6 +142,10 @@ class (Monad m) => MonadPostgres (m :: Type -> Type) where -- Only handlers should run transactions. runTransaction :: Transaction m a -> m a +-- | Quasi-Quoter for multi-line SQL literals. Trims leading whitespace up to the least-indented line. +sql :: QuasiQuoter +sql = PyF.fmtTrim + -- | Run a query, passing parameters. Prefer 'queryWith' if possible. query :: forall m params r. @@ -364,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, @@ -391,13 +396,14 @@ 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), - nestedError "Query" (formattedQuery & newError & singleError) + nestedError "Query" (formattedQuery & bytesToTextUtf8Lenient & newError & singleError) ] ) logFormatException :: FormatError -> Transaction m a @@ -527,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 -> - m DebugLogDatabaseQueries -> + (ToRow params, MonadUnliftIO m, MonadLogger m, Otel.MonadTracer m) => + m (DebugLogDatabaseQueries, PrettyPrintDatabaseQueries) -> Query -> params -> Transaction m (Label "numberOfRowsAffected" Natural) {-# INLINE executeImpl #-} -executeImpl zoomTools zoomDebugLogDatabaseQueries qry params = +executeImpl zoomDbOptions qry params = Otel.inSpan' "Postgres Query (execute)" Otel.defaultSpanArguments $ \span -> do - tools <- lift @Transaction zoomTools - logDatabaseQueries <- lift @Transaction zoomDebugLogDatabaseQueries - traceQueryIfEnabled tools span logDatabaseQueries qry (HasSingleParam params) + (logDatabaseQueries, prettyQuery) <- lift @Transaction zoomDbOptions + 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 -> - m DebugLogDatabaseQueries -> + ( MonadUnliftIO m, + MonadLogger m, + Otel.MonadTracer m + ) => + m (DebugLogDatabaseQueries, PrettyPrintDatabaseQueries) -> Query -> Transaction m (Label "numberOfRowsAffected" Natural) {-# INLINE executeImpl_ #-} -executeImpl_ zoomTools zoomDebugLogDatabaseQueries qry = +executeImpl_ zoomDbOptions qry = Otel.inSpan' "Postgres Query (execute)" Otel.defaultSpanArguments $ \span -> do - tools <- lift @Transaction zoomTools - logDatabaseQueries <- lift @Transaction zoomDebugLogDatabaseQueries - traceQueryIfEnabled @() tools span logDatabaseQueries qry HasNoParams + (logDatabaseQueries, prettyQuery) <- lift @Transaction zoomDbOptions + 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 -> - m DebugLogDatabaseQueries -> + (ToRow params, MonadUnliftIO m, MonadLogger m, Otel.MonadTracer m) => + m (DebugLogDatabaseQueries, PrettyPrintDatabaseQueries) -> Query -> NonEmpty params -> Transaction m (Label "numberOfRowsAffected" Natural) -executeManyImpl zoomTools zoomDebugLogDatabaseQueries qry params = +executeManyImpl zoomDbOptions qry params = Otel.inSpan' "Postgres Query (executeMany)" Otel.defaultSpanArguments $ \span -> do - tools <- lift @Transaction zoomTools - logDatabaseQueries <- lift @Transaction zoomDebugLogDatabaseQueries - traceQueryIfEnabled tools span logDatabaseQueries qry (HasMultiParams params) + (logDatabaseQueries, prettyQuery) <- lift @Transaction zoomDbOptions + 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) @@ -589,32 +592,32 @@ toNumberOfRowsAffected functionName i64 = <&> label @"numberOfRowsAffected" executeManyReturningWithImpl :: - (ToRow params, MonadUnliftIO m, MonadLogger m, HasField "pgFormat" tools PgFormatPool, Otel.MonadTracer m) => - m tools -> - m DebugLogDatabaseQueries -> + ( ToRow params, + MonadUnliftIO m, + MonadLogger m, + Otel.MonadTracer m + ) => + m (DebugLogDatabaseQueries, PrettyPrintDatabaseQueries) -> Query -> NonEmpty params -> Decoder r -> Transaction m [r] {-# INLINE executeManyReturningWithImpl #-} -executeManyReturningWithImpl zoomTools zoomDebugLogDatabaseQueries 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 <- lift @Transaction zoomDebugLogDatabaseQueries - traceQueryIfEnabled tools span logDatabaseQueries qry (HasMultiParams params) + (logDatabaseQueries, prettyQuery) <- lift @Transaction zoomDbOptions + 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 -> + m (DebugLogDatabaseQueries, PrettyPrintDatabaseQueries) -> Query -> params -> Decoder row -> @@ -622,11 +625,10 @@ foldRowsWithAccImpl :: (a -> row -> Transaction m a) -> Transaction m a {-# INLINE foldRowsWithAccImpl #-} -foldRowsWithAccImpl zoomTools zoomDebugLogDatabaseQueries 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 <- lift @Transaction zoomDebugLogDatabaseQueries - traceQueryIfEnabled tools span logDatabaseQueries qry (HasSingleParam params) + (logDatabaseQueries, prettyQuery) <- lift @Transaction zoomDbOptions + traceQueryIfEnabled span logDatabaseQueries prettyQuery qry (HasSingleParam params) conn <- Transaction ask withRunInIO ( \runInIO -> @@ -639,17 +641,18 @@ foldRowsWithAccImpl zoomTools zoomDebugLogDatabaseQueries qry params (Decoder ro 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 Text -pgFormatQueryNoParams' tools q = - lift $ pgFormatQueryByteString tools q.fromQuery + Transaction m ByteString +pgFormatQueryNoParams' prettyQuery q = case prettyQuery of + DontPrettyPrintDatabaseQueries -> pure q.fromQuery + PrettyPrintDatabaseQueries pool -> lift $ pgFormatQueryByteString pool q.fromQuery pgFormatQuery :: (ToRow params, MonadIO m) => @@ -680,40 +683,36 @@ queryWithImpl :: ( ToRow params, MonadUnliftIO m, MonadLogger m, - HasField "pgFormat" tools PgFormatPool, Otel.MonadTracer m ) => - m tools -> - m DebugLogDatabaseQueries -> + m (DebugLogDatabaseQueries, PrettyPrintDatabaseQueries) -> Query -> params -> Decoder r -> Transaction m [r] {-# INLINE queryWithImpl #-} -queryWithImpl zoomTools zoomDebugLogDatabaseQueries 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 <- lift @Transaction zoomDebugLogDatabaseQueries - traceQueryIfEnabled tools span logDatabaseQueries qry (HasSingleParam params) + (logDatabaseQueries, prettyQuery) <- lift @Transaction zoomDbOptions + 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 @@ -727,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 Text -pgFormatQuery' tools q p = - pgFormatQuery q p - >>= lift . pgFormatQueryByteString tools + Transaction m ByteString +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 Text -pgFormatQueryMany' tools q p = - pgFormatQueryMany q p - >>= lift . pgFormatQueryByteString tools + Transaction m ByteString +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) @@ -758,20 +759,19 @@ postgresToolsParser = label @"pgFormat" <$> readTool "pg_format" pgFormatQueryByteString :: ( MonadIO m, - MonadLogger m, - HasField "pgFormat" tools PgFormatPool + MonadLogger m ) => - tools -> + PgFormatPool -> ByteString -> - m Text -pgFormatQueryByteString tools queryBytes = do + m ByteString +pgFormatQueryByteString pool queryBytes = do res <- liftIO $ runPgFormat - tools.pgFormat + pool (queryBytes) case res.exitCode of - ExitSuccess -> pure (res.formatted & bytesToTextUtf8Lenient) + ExitSuccess -> pure (res.formatted) ExitFailure status -> do logWarn [fmt|pg_format failed with status {status} while formatting the query, using original query string. Is there a syntax error?|] logDebug @@ -784,7 +784,7 @@ pgFormatQueryByteString tools queryBytes = do ) ) logDebug [fmt|pg_format stdout: stderr|] - pure (queryBytes & bytesToTextUtf8Lenient) + pure (queryBytes) pgFormatStartCommandWaitForInput :: ( MonadIO m, @@ -821,6 +821,17 @@ data DebugLogDatabaseQueries LogDatabaseQueriesAndExplain deriving stock (Show, Enum, Bounded) +-- | Whether to pipe database queries thru `pg_format` before logging them. This takes a long (long! 200ms+) time per query, so should only be used in debugging environments where speed is not an issue. +data PrettyPrintDatabaseQueries + = -- | Do not pretty-print database querios + DontPrettyPrintDatabaseQueries + | -- | Pretty-print database queries, slow + PrettyPrintDatabaseQueries PgFormatPool + +instance Show PrettyPrintDatabaseQueries where + show DontPrettyPrintDatabaseQueries = "DontPrettyPrintDatabaseQueries" + show (PrettyPrintDatabaseQueries _) = "PrettyPrintDatabaseQueries" + data HasQueryParams param = HasNoParams | HasSingleParam param @@ -831,32 +842,31 @@ 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 qry params = do +traceQueryIfEnabled span logDatabaseQueries prettyQuery qry params = do -- In case we have query logging enabled, we want to do that - let formattedQuery = do + let formattedQuery = 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 + HasNoParams -> pgFormatQueryNoParams' prettyQuery qry + HasSingleParam p -> pgFormatQuery' prettyQuery qry p + HasMultiParams ps -> pgFormatQueryMany' prettyQuery qry ps let doLog errs = Otel.addAttributes span $ HashMap.fromList - $ ( ("_.postgres.query", Otel.toAttribute @Text errs.query) + $ ( ("_.postgres.query", Otel.toAttribute @Text (errs.query & bytesToTextUtf8Lenient)) : ( errs.explain & \case Nothing -> [] @@ -867,12 +877,12 @@ traceQueryIfEnabled tools span logDatabaseQueries 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 -- because we need the query with all elements already interpolated. - Query (q & textToBytesUtf8) + Query q ) ) (Dec.fromField @Text) @@ -920,6 +930,70 @@ withEvent span start end act = do ) pure res +unzipPGArray :: + forall l1 t1 l2 t2 r. + ( HasField l1 r t1, + HasField l2 r t2 + ) => + [r] -> + (PGArray t1, PGArray t2) +{-# INLINEABLE unzipPGArray #-} +unzipPGArray xs = + ( PGArray $ getField @l1 <$> xs, + PGArray $ getField @l2 <$> xs + ) + +unzip3PGArray :: + forall l1 t1 l2 t2 l3 t3 r. + ( HasField l1 r t1, + HasField l2 r t2, + HasField l3 r t3 + ) => + [r] -> + (PGArray t1, PGArray t2, PGArray t3) +{-# INLINEABLE unzip3PGArray #-} +unzip3PGArray xs = + ( PGArray $ getField @l1 <$> xs, + PGArray $ getField @l2 <$> xs, + PGArray $ getField @l3 <$> xs + ) + +unzip4PGArray :: + forall l1 t1 l2 t2 l3 t3 l4 t4 r. + ( HasField l1 r t1, + HasField l2 r t2, + HasField l3 r t3, + HasField l4 r t4 + ) => + [r] -> + (PGArray t1, PGArray t2, PGArray t3, PGArray t4) +{-# INLINEABLE unzip4PGArray #-} +unzip4PGArray xs = + ( PGArray $ getField @l1 <$> xs, + PGArray $ getField @l2 <$> xs, + PGArray $ getField @l3 <$> xs, + PGArray $ getField @l4 <$> xs + ) + +unzip5PGArray :: + forall l1 t1 l2 t2 l3 t3 l4 t4 l5 t5 r. + ( HasField l1 r t1, + HasField l2 r t2, + HasField l3 r t3, + HasField l4 r t4, + HasField l5 r t5 + ) => + [r] -> + (PGArray t1, PGArray t2, PGArray t3, PGArray t4, PGArray t5) +{-# INLINEABLE unzip5PGArray #-} +unzip5PGArray xs = + ( PGArray $ getField @l1 <$> xs, + PGArray $ getField @l2 <$> xs, + PGArray $ getField @l3 <$> xs, + PGArray $ getField @l4 <$> xs, + PGArray $ getField @l5 <$> xs + ) + instance (ToField t1) => ToRow (Label l1 t1) where toRow t2 = toRow $ PG.Only $ getField @l1 t2 |