about summary refs log tree commit diff
path: root/users/Profpatsch/my-prelude/src/Postgres/MonadPostgres.hs
diff options
context:
space:
mode:
authorProfpatsch <mail@profpatsch.de>2024-03-16T22·26+0100
committerclbot <clbot@tvl.fyi>2024-03-16T22·36+0000
commit11a2098e0b3a4f202d35a61da06a0bf1c724b8c9 (patch)
tree96eacf02a41d5ef755c29d95408df8b86c48e148 /users/Profpatsch/my-prelude/src/Postgres/MonadPostgres.hs
parent8335076173d2fd83a9bc13134d554255a527a8aa (diff)
feat(users/Profpatsch/my-prelude): update libraries r/7712
The latest and greatest!

Change-Id: I34c0e9f41b3b3cc727d9ea89c7ce6a43271b3170
Reviewed-on: https://cl.tvl.fyi/c/depot/+/11169
Autosubmit: Profpatsch <mail@profpatsch.de>
Tested-by: BuildkiteCI
Reviewed-by: Profpatsch <mail@profpatsch.de>
Diffstat (limited to 'users/Profpatsch/my-prelude/src/Postgres/MonadPostgres.hs')
-rw-r--r--users/Profpatsch/my-prelude/src/Postgres/MonadPostgres.hs313
1 files changed, 236 insertions, 77 deletions
diff --git a/users/Profpatsch/my-prelude/src/Postgres/MonadPostgres.hs b/users/Profpatsch/my-prelude/src/Postgres/MonadPostgres.hs
index 78e3897ef5f3..bd8ddd12f775 100644
--- a/users/Profpatsch/my-prelude/src/Postgres/MonadPostgres.hs
+++ b/users/Profpatsch/my-prelude/src/Postgres/MonadPostgres.hs
@@ -1,14 +1,15 @@
+{-# LANGUAGE AllowAmbiguousTypes #-}
 {-# LANGUAGE DeriveAnyClass #-}
 {-# LANGUAGE QuasiQuotes #-}
-{-# LANGUAGE TemplateHaskell #-}
 {-# OPTIONS_GHC -Wno-orphans #-}
 
 module Postgres.MonadPostgres where
 
 import AtLeast (AtLeast)
 import Control.Exception
+import Control.Foldl qualified as Fold
 import Control.Monad.Except
-import Control.Monad.Logger (MonadLogger, logDebug, logWarn)
+import Control.Monad.Logger.CallStack (MonadLogger, logDebug, logWarn)
 import Control.Monad.Reader (MonadReader (ask), ReaderT (..))
 import Control.Monad.Trans.Resource
 import Data.Aeson (FromJSON)
@@ -28,7 +29,7 @@ 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 GHC.Records (HasField (..))
+import GHC.Records (getField)
 import Label
 import OpenTelemetry.Trace.Core qualified as Otel hiding (inSpan, inSpan')
 import OpenTelemetry.Trace.Monad qualified as Otel
@@ -42,7 +43,7 @@ import Tool
 import UnliftIO (MonadUnliftIO (withRunInIO))
 import UnliftIO.Process qualified as Process
 import UnliftIO.Resource qualified as Resource
-import Prelude hiding (span)
+import Prelude hiding (init, span)
 
 -- | Postgres queries/commands that can be executed within a running transaction.
 --
@@ -52,28 +53,46 @@ class (Monad m) => MonadPostgres (m :: Type -> Type) where
   -- | Execute an INSERT, UPDATE, or other SQL query that is not expected to return results.
 
   -- Returns the number of rows affected.
-  execute :: (ToRow params, Typeable params) => Query -> params -> Transaction m (Label "numberOfRowsAffected" Natural)
-
-  -- | Execute an INSERT, UPDATE, or other SQL query that is not expected to return results. Does not take parameters.
-
-  -- Returns the number of rows affected.
-  execute_ :: Query -> Transaction m (Label "numberOfRowsAffected" Natural)
+  execute ::
+    (ToRow params, Typeable params) =>
+    Query ->
+    params ->
+    Transaction m (Label "numberOfRowsAffected" Natural)
 
   -- | Execute a multi-row INSERT, UPDATE, or other SQL query that is not expected to return results.
   --
-  -- Returns the number of rows affected. If the list of parameters is empty, this function will simply return 0 without issuing the query to the backend. If this is not desired, consider using the 'PG.Values' constructor instead.
-  executeMany :: (ToRow params, Typeable params) => Query -> [params] -> Transaction m (Label "numberOfRowsAffected" Natural)
-
-  -- | Execute INSERT ... RETURNING, UPDATE ... RETURNING, or other SQL query that accepts multi-row input and is expected to return results. Note that it is possible to write query conn "INSERT ... RETURNING ..." ... in cases where you are only inserting a single row, and do not need functionality analogous to 'executeMany'.
+  -- Returns the number of rows affected. If the list of parameters is empty,
+  -- this function will simply return 0 without issuing the query to the backend.
+  -- If this is not desired, consider using the 'PG.Values' constructor instead.
+  executeMany ::
+    (ToRow params, Typeable params) =>
+    Query ->
+    NonEmpty params ->
+    Transaction m (Label "numberOfRowsAffected" Natural)
+
+  -- | Execute INSERT ... RETURNING, UPDATE ... RETURNING,
+  -- or other SQL query that accepts multi-row input and is expected to return results.
+  -- Note that it is possible to write query conn "INSERT ... RETURNING ..." ...
+  -- in cases where you are only inserting a single row,
+  -- and do not need functionality analogous to 'executeMany'.
   --
   -- If the list of parameters is empty, this function will simply return [] without issuing the query to the backend. If this is not desired, consider using the 'PG.Values' constructor instead.
-  executeManyReturningWith :: (ToRow q) => Query -> [q] -> Decoder r -> Transaction m [r]
+  executeManyReturningWith :: (ToRow q) => Query -> NonEmpty q -> Decoder r -> Transaction m [r]
 
   -- | Run a query, passing parameters and result row parser.
-  queryWith :: (PG.ToRow params, Typeable params, Typeable r) => PG.Query -> params -> Decoder r -> Transaction m [r]
+  queryWith ::
+    (PG.ToRow params, Typeable params, Typeable r) =>
+    PG.Query ->
+    params ->
+    Decoder r ->
+    Transaction m [r]
 
   -- | Run a query without any parameters and result row parser.
-  queryWith_ :: (Typeable r) => PG.Query -> Decoder r -> Transaction m [r]
+  queryWith_ ::
+    (Typeable r) =>
+    PG.Query ->
+    Decoder r ->
+    Transaction m [r]
 
   -- | Run a query, passing parameters, and fold over the resulting rows.
   --
@@ -82,13 +101,15 @@ class (Monad m) => MonadPostgres (m :: Type -> Type) where
   --
   -- When dealing with small results, it may be simpler (and perhaps faster) to use query instead.
   --
-  -- This fold is _not_ strict. The stream consumer is responsible for forcing the evaluation of its result to avoid space leaks.
+  -- This fold is _not_ strict. The stream consumer is responsible
+  -- for forcing the evaluation of its result to avoid space leaks.
   --
   -- If you can, prefer aggregating in the database itself.
-  foldRows ::
-    (FromRow row, ToRow params, Typeable row, Typeable params) =>
+  foldRowsWithAcc ::
+    (ToRow params, Typeable row, Typeable params) =>
     Query ->
     params ->
+    Decoder row ->
     a ->
     (a -> row -> Transaction m a) ->
     Transaction m a
@@ -109,12 +130,23 @@ class (Monad m) => MonadPostgres (m :: Type -> Type) where
   -- Only handlers should run transactions.
   runTransaction :: Transaction m a -> m a
 
--- | Run a query, passing parameters.
-query :: forall m params r. (PG.ToRow params, PG.FromRow r, Typeable params, Typeable r, MonadPostgres m) => PG.Query -> params -> Transaction m [r]
+-- | Run a query, passing parameters. Prefer 'queryWith' if possible.
+query ::
+  forall m params r.
+  (PG.ToRow params, PG.FromRow r, Typeable params, Typeable r, MonadPostgres m) =>
+  PG.Query ->
+  params ->
+  Transaction m [r]
 query qry params = queryWith qry params (Decoder PG.fromRow)
 
--- | Run a query without any parameters.
-query_ :: forall m r. (Typeable r, PG.FromRow r, MonadPostgres m) => PG.Query -> Transaction m [r]
+-- | Run a query without any parameters. Prefer 'queryWith' if possible.
+--
+-- TODO: I think(?) this can always be replaced by passing @()@ to 'query', remove?
+query_ ::
+  forall m r.
+  (Typeable r, PG.FromRow r, MonadPostgres m) =>
+  PG.Query ->
+  Transaction m [r]
 query_ qry = queryWith_ qry (Decoder PG.fromRow)
 
 -- TODO: implement via fold, so that the result doesn’t have to be realized in memory
@@ -153,7 +185,10 @@ querySingleRowMaybe qry params = do
     -- that a database function can error out, should probably handled by the instances.
     more -> throwM $ SingleRowError {numberOfRowsReturned = (List.length more)}
 
-ensureSingleRow :: (MonadThrow m) => [a] -> m a
+ensureSingleRow ::
+  (MonadThrow m) =>
+  [a] ->
+  m a
 ensureSingleRow = \case
   -- TODO: Should we MonadThrow this here? It’s really an implementation detail of MonadPostgres
   -- that a database function can error out, should probably handled by the instances.
@@ -167,6 +202,52 @@ ensureSingleRow = \case
             List.length more
         }
 
+ensureNoneOrSingleRow ::
+  (MonadThrow m) =>
+  [a] ->
+  m (Maybe a)
+ensureNoneOrSingleRow = \case
+  -- TODO: Should we MonadThrow this here? It’s really an implementation detail of MonadPostgres
+  -- that a database function can error out, should probably handled by the instances.
+  [] -> pure Nothing
+  [one] -> pure $ Just one
+  more ->
+    throwM $
+      SingleRowError
+        { numberOfRowsReturned =
+            -- TODO: this is VERY bad, because it requires to parse the full database output, even if there’s 10000000000 elements
+            List.length more
+        }
+
+-- | Run a query, passing parameters, and fold over the resulting rows.
+--
+-- This doesn’t have to realize the full list of results in memory,
+-- rather results are streamed incrementally from the database.
+--
+-- When dealing with small results, it may be simpler (and perhaps faster) to use query instead.
+--
+-- The results are folded strictly by the 'Fold.Fold' that is passed.
+--
+-- If you can, prefer aggregating in the database itself.
+foldRowsWith ::
+  forall row params m b.
+  ( MonadPostgres m,
+    PG.ToRow params,
+    Typeable row,
+    Typeable params
+  ) =>
+  PG.Query ->
+  params ->
+  Decoder row ->
+  Fold.Fold row b ->
+  Transaction m b
+foldRowsWith qry params decoder = Fold.purely f
+  where
+    f :: forall x. (x -> row -> x) -> x -> (x -> b) -> Transaction m b
+    f acc init extract = do
+      x <- foldRowsWithAcc qry params decoder init (\a r -> pure $ acc a r)
+      pure $ extract x
+
 newtype Transaction m a = Transaction {unTransaction :: (ReaderT Connection m a)}
   deriving newtype
     ( Functor,
@@ -180,9 +261,6 @@ newtype Transaction m a = Transaction {unTransaction :: (ReaderT Connection m a)
       Otel.MonadTracer
     )
 
-runTransaction' :: Connection -> Transaction m a -> m a
-runTransaction' conn transaction = runReaderT transaction.unTransaction conn
-
 -- | [Resource Pool](http://hackage.haskell.org/package/resource-pool-0.2.3.2/docs/Data-Pool.html) configuration.
 data PoolingInfo = PoolingInfo
   { -- | Minimal amount of resources that are
@@ -237,17 +315,41 @@ initMonadPostgres logInfoFn connectInfo poolingInfo = do
       IO ()
     destroyPGConnPool p = Pool.destroyAllResources p
 
+-- | Improve a possible error message, by adding some context to it.
+--
+-- The given Exception type is caught, 'show'n and pretty-printed.
+--
+-- In case we get an `IOError`, we display it in a reasonable fashion.
+addErrorInformation ::
+  forall exc a.
+  (Exception exc) =>
+  Text.Text ->
+  IO a ->
+  IO a
+addErrorInformation msg io =
+  io
+    & try @exc
+    <&> first (showPretty >>> newError >>> errorContext msg)
+    & try @IOError
+    <&> first (showToError >>> errorContext "IOError" >>> errorContext msg)
+    <&> join @(Either Error)
+    >>= unwrapIOError
+
 -- | Catch any Postgres exception that gets thrown,
 -- print the query that was run and the query parameters,
 -- then rethrow inside an 'Error'.
 handlePGException ::
   forall a params tools m.
-  (ToRow params, MonadUnliftIO m, MonadLogger m, HasField "pgFormat" tools Tool) =>
+  ( ToRow params,
+    MonadUnliftIO m,
+    MonadLogger m,
+    HasField "pgFormat" tools Tool
+  ) =>
   tools ->
   Text ->
   Query ->
   -- | Depending on whether we used `format` or `formatMany`.
-  Either params [params] ->
+  Either params (NonEmpty params) ->
   IO a ->
   Transaction m a
 handlePGException tools queryType query' params io = do
@@ -289,7 +391,11 @@ withPGTransaction connPool f =
     connPool
     (\conn -> Postgres.withTransaction conn (f conn))
 
-runPGTransactionImpl :: (MonadUnliftIO m) => m (Pool Postgres.Connection) -> Transaction m a -> m a
+runPGTransactionImpl ::
+  (MonadUnliftIO m) =>
+  m (Pool Postgres.Connection) ->
+  Transaction m a ->
+  m a
 {-# INLINE runPGTransactionImpl #-}
 runPGTransactionImpl zoom (Transaction transaction) = do
   pool <- zoom
@@ -337,7 +443,7 @@ executeManyImpl ::
   m tools ->
   m DebugLogDatabaseQueries ->
   Query ->
-  [params] ->
+  NonEmpty params ->
   Transaction m (Label "numberOfRowsAffected" Natural)
 executeManyImpl zoomTools zoomDebugLogDatabaseQueries qry params =
   Otel.inSpan' "Postgres Query (execute)" Otel.defaultSpanArguments $ \span -> do
@@ -345,7 +451,7 @@ executeManyImpl zoomTools zoomDebugLogDatabaseQueries qry params =
     logDatabaseQueries <- lift @Transaction zoomDebugLogDatabaseQueries
     traceQueryIfEnabled tools span logDatabaseQueries qry (HasMultiParams params)
     conn <- Transaction ask
-    PG.executeMany conn qry params
+    PG.executeMany conn qry (params & toList)
       & handlePGException tools "executeMany" qry (Right params)
       >>= toNumberOfRowsAffected "executeManyImpl"
 
@@ -364,7 +470,7 @@ executeManyReturningWithImpl ::
   m tools ->
   m DebugLogDatabaseQueries ->
   Query ->
-  [params] ->
+  NonEmpty params ->
   Decoder r ->
   Transaction m [r]
 {-# INLINE executeManyReturningWithImpl #-}
@@ -374,33 +480,45 @@ executeManyReturningWithImpl zoomTools zoomDebugLogDatabaseQueries qry params (D
     logDatabaseQueries <- lift @Transaction zoomDebugLogDatabaseQueries
     traceQueryIfEnabled tools span logDatabaseQueries qry (HasMultiParams params)
     conn <- Transaction ask
-    PG.returningWith fromRow conn qry params
+    PG.returningWith fromRow conn qry (params & toList)
       & handlePGException tools "executeManyReturning" qry (Right params)
 
-foldRowsImpl ::
-  (FromRow row, ToRow params, MonadUnliftIO m, MonadLogger m, HasField "pgFormat" tools Tool) =>
+foldRowsWithAccImpl ::
+  ( ToRow params,
+    MonadUnliftIO m,
+    MonadLogger m,
+    HasField "pgFormat" tools Tool,
+    Otel.MonadTracer m
+  ) =>
   m tools ->
+  m DebugLogDatabaseQueries ->
   Query ->
   params ->
+  Decoder row ->
   a ->
   (a -> row -> Transaction m a) ->
   Transaction m a
-{-# INLINE foldRowsImpl #-}
-foldRowsImpl zoomTools qry params accumulator f = do
-  conn <- Transaction ask
-  tools <- lift @Transaction zoomTools
-  withRunInIO
-    ( \runInIO ->
-        do
-          PG.fold
-            conn
-            qry
-            params
-            accumulator
-            (\acc row -> runInIO $ f acc row)
-            & handlePGException tools "fold" qry (Left params)
-            & runInIO
-    )
+{-# INLINE foldRowsWithAccImpl #-}
+foldRowsWithAccImpl zoomTools zoomDebugLogDatabaseQueries 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)
+    conn <- Transaction ask
+    withRunInIO
+      ( \runInIO ->
+          do
+            PG.foldWithOptionsAndParser
+              PG.defaultFoldOptions
+              rowParser
+              conn
+              qry
+              params
+              accumulator
+              (\acc row -> runInIO $ f acc row)
+              & handlePGException tools "fold" qry (Left params)
+              & runInIO
+      )
 
 pgFormatQueryNoParams' ::
   (MonadIO m, MonadLogger m, HasField "pgFormat" tools Tool) =>
@@ -410,18 +528,38 @@ pgFormatQueryNoParams' ::
 pgFormatQueryNoParams' tools q =
   lift $ pgFormatQueryByteString tools q.fromQuery
 
-pgFormatQuery :: (ToRow params, MonadIO m) => Query -> params -> Transaction m ByteString
+pgFormatQuery ::
+  (ToRow params, MonadIO m) =>
+  Query ->
+  params ->
+  Transaction m ByteString
 pgFormatQuery qry params = Transaction $ do
   conn <- ask
   liftIO $ PG.formatQuery conn qry params
 
-pgFormatQueryMany :: (MonadIO m, ToRow params) => Query -> [params] -> Transaction m ByteString
+pgFormatQueryMany ::
+  (MonadIO m, ToRow params) =>
+  Query ->
+  NonEmpty params ->
+  Transaction m ByteString
 pgFormatQueryMany qry params = Transaction $ do
   conn <- ask
-  liftIO $ PG.formatMany conn qry params
+  liftIO $
+    PG.formatMany
+      conn
+      qry
+      ( params
+          -- upstream is partial on empty list, see https://github.com/haskellari/postgresql-simple/issues/129
+          & toList
+      )
 
 queryWithImpl ::
-  (ToRow params, MonadUnliftIO m, MonadLogger m, HasField "pgFormat" tools Tool, Otel.MonadTracer m) =>
+  ( ToRow params,
+    MonadUnliftIO m,
+    MonadLogger m,
+    HasField "pgFormat" tools Tool,
+    Otel.MonadTracer m
+  ) =>
   m tools ->
   m DebugLogDatabaseQueries ->
   Query ->
@@ -438,7 +576,15 @@ queryWithImpl zoomTools zoomDebugLogDatabaseQueries qry params (Decoder fromRow)
     PG.queryWith fromRow conn qry params
       & handlePGException tools "query" qry (Left params)
 
-queryWithImpl_ :: (MonadUnliftIO m, MonadLogger m, HasField "pgFormat" tools Tool) => m tools -> Query -> Decoder r -> Transaction m [r]
+queryWithImpl_ ::
+  ( MonadUnliftIO m,
+    MonadLogger m,
+    HasField "pgFormat" tools Tool
+  ) =>
+  m tools ->
+  Query ->
+  Decoder r ->
+  Transaction m [r]
 {-# INLINE queryWithImpl_ #-}
 queryWithImpl_ zoomTools qry (Decoder fromRow) = do
   tools <- lift @Transaction zoomTools
@@ -446,18 +592,6 @@ queryWithImpl_ zoomTools qry (Decoder fromRow) = do
   liftIO (PG.queryWith_ fromRow conn qry)
     & handlePGException tools "query" qry (Left ())
 
-pgQuery :: (ToRow params, FromRow r, MonadUnliftIO m, MonadLogger m, HasField "pgFormat" tools Tool) => tools -> Query -> params -> Transaction m [r]
-pgQuery tools qry params = do
-  conn <- Transaction ask
-  PG.query conn qry params
-    & handlePGException tools "query" qry (Left params)
-
-pgQuery_ :: (FromRow r, MonadUnliftIO m, MonadLogger m, HasField "pgFormat" tools Tool) => tools -> Query -> Transaction m [r]
-pgQuery_ tools qry = do
-  conn <- Transaction ask
-  PG.query_ conn qry
-    & handlePGException tools "query_" qry (Left ())
-
 data SingleRowError = SingleRowError
   { -- | How many columns were actually returned by the query
     numberOfRowsReturned :: Int
@@ -467,12 +601,30 @@ data SingleRowError = SingleRowError
 instance Exception SingleRowError where
   displayException (SingleRowError {..}) = [fmt|Single row expected from SQL query result, {numberOfRowsReturned} rows were returned instead."|]
 
-pgFormatQuery' :: (MonadIO m, ToRow params, MonadLogger m, HasField "pgFormat" tools Tool) => tools -> Query -> params -> Transaction m Text
+pgFormatQuery' ::
+  ( MonadIO m,
+    ToRow params,
+    MonadLogger m,
+    HasField "pgFormat" tools Tool
+  ) =>
+  tools ->
+  Query ->
+  params ->
+  Transaction m Text
 pgFormatQuery' tools q p =
   pgFormatQuery q p
     >>= lift . pgFormatQueryByteString tools
 
-pgFormatQueryMany' :: (MonadIO m, ToRow params, MonadLogger m, HasField "pgFormat" tools Tool) => tools -> Query -> [params] -> Transaction m Text
+pgFormatQueryMany' ::
+  ( MonadIO m,
+    ToRow params,
+    MonadLogger m,
+    HasField "pgFormat" tools Tool
+  ) =>
+  tools ->
+  Query ->
+  NonEmpty params ->
+  Transaction m Text
 pgFormatQueryMany' tools q p =
   pgFormatQueryMany q p
     >>= lift . pgFormatQueryByteString tools
@@ -481,7 +633,14 @@ pgFormatQueryMany' tools q p =
 postgresToolsParser :: ToolParserT IO (Label "pgFormat" Tool)
 postgresToolsParser = label @"pgFormat" <$> readTool "pg_format"
 
-pgFormatQueryByteString :: (MonadIO m, MonadLogger m, HasField "pgFormat" tools Tool) => tools -> ByteString -> m Text
+pgFormatQueryByteString ::
+  ( MonadIO m,
+    MonadLogger m,
+    HasField "pgFormat" tools Tool
+  ) =>
+  tools ->
+  ByteString ->
+  m Text
 pgFormatQueryByteString tools queryBytes = do
   do
     (exitCode, stdout, stderr) <-
@@ -492,8 +651,8 @@ pgFormatQueryByteString tools queryBytes = do
     case exitCode of
       ExitSuccess -> pure (stdout & stringToText)
       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
+        logWarn [fmt|pg_format failed with status {status} while formatting the query, using original query string. Is there a syntax error?|]
+        logDebug
           ( prettyErrorTree
               ( nestedMultiError
                   "pg_format output"
@@ -502,7 +661,7 @@ pgFormatQueryByteString tools queryBytes = do
                   )
               )
           )
-        $logDebug [fmt|pg_format stdout: stderr|]
+        logDebug [fmt|pg_format stdout: stderr|]
         pure (queryBytes & bytesToTextUtf8Lenient)
 
 data DebugLogDatabaseQueries
@@ -517,7 +676,7 @@ data DebugLogDatabaseQueries
 data HasQueryParams param
   = HasNoParams
   | HasSingleParam param
-  | HasMultiParams [param]
+  | HasMultiParams (NonEmpty param)
 
 -- | Log the postgres query depending on the given setting
 traceQueryIfEnabled ::