diff options
Diffstat (limited to 'users/Profpatsch/my-prelude/src/Postgres')
| -rw-r--r-- | users/Profpatsch/my-prelude/src/Postgres/MonadPostgres.hs | 313 |
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 :: |