diff --git a/CHANGELOG.md b/CHANGELOG.md index e91de63..c9b7b37 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,7 @@ # hpqtypes-1.12.0.0 (????-??-??) * Drop support for GHC 8.8. * Attach `CallStack` to `DBException`. +* Store ID of the server process attached to the current session. # hpqtypes-1.11.1.2 (2023-11-08) * Support multihost setups and the `connect_timeout` parameter in the connection diff --git a/hpqtypes.cabal b/hpqtypes.cabal index 2dea4ac..8ba6cce 100644 --- a/hpqtypes.cabal +++ b/hpqtypes.cabal @@ -77,6 +77,7 @@ library , Database.PostgreSQL.PQTypes.SQL.Class , Database.PostgreSQL.PQTypes.Transaction.Settings , Database.PostgreSQL.PQTypes.XML + , Database.PostgreSQL.PQTypes.Internal.BackendPid , Database.PostgreSQL.PQTypes.Internal.Error , Database.PostgreSQL.PQTypes.Internal.Error.Code , Database.PostgreSQL.PQTypes.Internal.Composite @@ -146,6 +147,7 @@ library , ConstraintKinds , DataKinds , DeriveFunctor + , DerivingStrategies , ExistentialQuantification , FlexibleContexts , FlexibleInstances @@ -213,6 +215,7 @@ test-suite hpqtypes-tests , ConstraintKinds , DataKinds , DeriveFunctor + , DerivingStrategies , ExistentialQuantification , FlexibleContexts , FlexibleInstances diff --git a/src/Database/PostgreSQL/PQTypes/Class.hs b/src/Database/PostgreSQL/PQTypes/Class.hs index 008170d..3bb4c56 100644 --- a/src/Database/PostgreSQL/PQTypes/Class.hs +++ b/src/Database/PostgreSQL/PQTypes/Class.hs @@ -1,6 +1,10 @@ module Database.PostgreSQL.PQTypes.Class - ( QueryName (..) - , MonadDB (..) + ( -- * Class + MonadDB (..) + + -- * Misc + , BackendPid (..) + , QueryName (..) ) where import Control.Monad.Trans @@ -8,6 +12,7 @@ import Control.Monad.Trans.Control import GHC.Stack import Database.PostgreSQL.PQTypes.FromRow +import Database.PostgreSQL.PQTypes.Internal.BackendPid import Database.PostgreSQL.PQTypes.Internal.Connection import Database.PostgreSQL.PQTypes.Internal.Notification import Database.PostgreSQL.PQTypes.Internal.QueryResult @@ -32,6 +37,9 @@ class (Applicative m, Monad m) => MonadDB m where -- 'getLastQuery'. withFrozenLastQuery :: m a -> m a + -- | Get ID of the server process attached to the current session. + getBackendPid :: m BackendPid + -- | Get current connection statistics. getConnectionStats :: HasCallStack => m ConnectionStats @@ -89,6 +97,7 @@ instance runPreparedQuery name = withFrozenCallStack $ lift . runPreparedQuery name getLastQuery = lift getLastQuery withFrozenLastQuery m = controlT $ \run -> withFrozenLastQuery (run m) + getBackendPid = lift getBackendPid getConnectionStats = withFrozenCallStack $ lift getConnectionStats getQueryResult = lift getQueryResult clearQueryResult = lift clearQueryResult diff --git a/src/Database/PostgreSQL/PQTypes/Internal/BackendPid.hs b/src/Database/PostgreSQL/PQTypes/Internal/BackendPid.hs new file mode 100644 index 0000000..5a789fc --- /dev/null +++ b/src/Database/PostgreSQL/PQTypes/Internal/BackendPid.hs @@ -0,0 +1,7 @@ +module Database.PostgreSQL.PQTypes.Internal.BackendPid + ( BackendPid (..) + ) where + +-- | Process ID of the server process attached to the current session. +newtype BackendPid = BackendPid Int + deriving newtype (Eq, Ord, Show) diff --git a/src/Database/PostgreSQL/PQTypes/Internal/Connection.hs b/src/Database/PostgreSQL/PQTypes/Internal/Connection.hs index de26d92..5ee30b8 100644 --- a/src/Database/PostgreSQL/PQTypes/Internal/Connection.hs +++ b/src/Database/PostgreSQL/PQTypes/Internal/Connection.hs @@ -1,6 +1,9 @@ +{-# LANGUAGE TypeApplications #-} + module Database.PostgreSQL.PQTypes.Internal.Connection ( -- * Connection Connection (..) + , getBackendPidIO , ConnectionData (..) , withConnectionData , ConnectionStats (..) @@ -26,10 +29,11 @@ import Control.Exception qualified as E import Control.Monad import Control.Monad.Base import Control.Monad.Catch -import Data.Bifunctor import Data.ByteString.Char8 qualified as BS import Data.Foldable qualified as F +import Data.Functor.Identity import Data.IORef +import Data.Int import Data.Kind import Data.Pool import Data.Set qualified as S @@ -42,12 +46,14 @@ import Foreign.Ptr import GHC.Conc (closeFdWith) import GHC.Stack +import Database.PostgreSQL.PQTypes.Internal.BackendPid import Database.PostgreSQL.PQTypes.Internal.C.Interface import Database.PostgreSQL.PQTypes.Internal.C.Types import Database.PostgreSQL.PQTypes.Internal.Composite import Database.PostgreSQL.PQTypes.Internal.Error import Database.PostgreSQL.PQTypes.Internal.Error.Code import Database.PostgreSQL.PQTypes.Internal.Exception +import Database.PostgreSQL.PQTypes.Internal.QueryResult import Database.PostgreSQL.PQTypes.Internal.Utils import Database.PostgreSQL.PQTypes.SQL.Class import Database.PostgreSQL.PQTypes.SQL.Raw @@ -114,6 +120,8 @@ initialStats = data ConnectionData = ConnectionData { cdPtr :: !(Ptr PGconn) -- ^ Pointer to connection object. + , cdBackendPid :: !BackendPid + -- ^ Process ID of the server process attached to the current session. , cdStats :: !ConnectionStats -- ^ Statistics associated with the connection. , cdPreparedQueries :: !(IORef (S.Set T.Text)) @@ -125,6 +133,11 @@ newtype Connection = Connection { unConnection :: MVar (Maybe ConnectionData) } +getBackendPidIO :: Connection -> IO BackendPid +getBackendPidIO conn = do + withConnectionData conn "getBackendPidIO" $ \cd -> do + pure (cd, cdBackendPid cd) + withConnectionData :: Connection -> String @@ -133,7 +146,9 @@ withConnectionData withConnectionData (Connection mvc) fname f = modifyMVar mvc $ \mc -> case mc of Nothing -> hpqTypesError $ fname ++ ": no connection" - Just cd -> first Just <$> f cd + Just cd -> do + (cd', r) <- f cd + cd' `seq` pure (Just cd', r) -- | Database connection supplier. newtype ConnectionSourceM m = ConnectionSourceM @@ -215,12 +230,25 @@ connect ConnectionSettings {..} = mask $ \unmask -> do Just ConnectionData { cdPtr = connPtr + , cdBackendPid = noBackendPid , cdStats = initialStats , cdPreparedQueries = preparedQueries } F.forM_ csRole $ \role -> runQueryIO conn $ "SET ROLE " <> role + + let selectPid = "SELECT pg_backend_pid()" :: RawSQL () + (_, res) <- runQueryIO conn selectPid + case F.toList $ mkQueryResult @(Identity Int32) selectPid noBackendPid res of + [pid] -> withConnectionData conn fname $ \cd -> do + pure (cd {cdBackendPid = BackendPid $ fromIntegral pid}, ()) + pids -> do + let err = HPQTypesError $ "unexpected backend pid: " ++ show pids + rethrowWithContext selectPid noBackendPid $ toException err + pure conn where + noBackendPid = BackendPid 0 + fname = "connect" openConnection :: (forall r. IO r -> IO r) -> CString -> IO (Ptr PGconn) @@ -317,6 +345,7 @@ runPreparedQueryIO conn (QueryName queryName) sql = do E.throwIO DBException { dbeQueryContext = sql + , dbeBackendPid = cdBackendPid , dbeError = HPQTypesError "runPreparedQueryIO: unnamed prepared query is not supported" , dbeCallStack = callStack } @@ -329,7 +358,7 @@ runPreparedQueryIO conn (QueryName queryName) sql = do -- succeeds, we need to reflect that fact in cdPreparedQueries since -- you can't prepare a query with the same name more than once. res <- c_PQparamPrepare cdPtr nullPtr param cname query - void . withForeignPtr res $ verifyResult sql cdPtr + void . withForeignPtr res $ verifyResult sql cdBackendPid cdPtr modifyIORef' cdPreparedQueries $ S.insert queryName (,) <$> (fromIntegral <$> c_PQparamCount param) @@ -353,7 +382,7 @@ runQueryImpl fname conn sql execSql = do -- runtime system is used) and react appropriately. queryRunner <- async . restore $ do (paramCount, res) <- execSql cd - affected <- withForeignPtr res $ verifyResult sql cdPtr + affected <- withForeignPtr res $ verifyResult sql cdBackendPid cdPtr stats' <- case affected of Left _ -> return @@ -370,8 +399,7 @@ runQueryImpl fname conn sql execSql = do , statsValues = statsValues cdStats + (rows * columns) , statsParams = statsParams cdStats + paramCount } - -- Force evaluation of modified stats to squash a space leak. - stats' `seq` return (cd {cdStats = stats'}, (either id id affected, res)) + return (cd {cdStats = stats'}, (either id id affected, res)) -- If we receive an exception while waiting for the execution to complete, -- we need to send a request to PostgreSQL for query cancellation and wait -- for the query runner thread to terminate. It is paramount we make the @@ -399,10 +427,11 @@ runQueryImpl fname conn sql execSql = do verifyResult :: (HasCallStack, IsSQL sql) => sql + -> BackendPid -> Ptr PGconn -> Ptr PGresult -> IO (Either Int Int) -verifyResult sql conn res = do +verifyResult sql pid conn res = do -- works even if res is NULL rst <- c_PQresultStatus res case rst of @@ -421,7 +450,7 @@ verifyResult sql conn res = do _ | otherwise -> return . Left $ 0 where throwSQLError = - rethrowWithContext sql + rethrowWithContext sql pid =<< if res == nullPtr then return . E.toException . QueryError @@ -451,6 +480,7 @@ verifyResult sql conn res = do E.throwIO DBException { dbeQueryContext = sql + , dbeBackendPid = pid , dbeError = HPQTypesError ("verifyResult: string returned by PQcmdTuples is not a valid number: " ++ show sn) , dbeCallStack = callStack } diff --git a/src/Database/PostgreSQL/PQTypes/Internal/Exception.hs b/src/Database/PostgreSQL/PQTypes/Internal/Exception.hs index d0433c0..427f63a 100644 --- a/src/Database/PostgreSQL/PQTypes/Internal/Exception.hs +++ b/src/Database/PostgreSQL/PQTypes/Internal/Exception.hs @@ -7,6 +7,7 @@ module Database.PostgreSQL.PQTypes.Internal.Exception import Control.Exception qualified as E import GHC.Stack +import Database.PostgreSQL.PQTypes.Internal.BackendPid import Database.PostgreSQL.PQTypes.SQL.Class -- | Main exception type. All exceptions thrown by @@ -14,6 +15,8 @@ import Database.PostgreSQL.PQTypes.SQL.Class data DBException = forall e sql. (E.Exception e, Show sql) => DBException { dbeQueryContext :: !sql -- ^ Last SQL query that was executed. + , dbeBackendPid :: !BackendPid + -- ^ Process ID of the server process attached to the current session. , dbeError :: !e -- ^ Specific error. , dbeCallStack :: CallStack @@ -24,11 +27,17 @@ deriving instance Show DBException instance E.Exception DBException -- | Rethrow supplied exception enriched with given SQL. -rethrowWithContext :: (HasCallStack, IsSQL sql) => sql -> E.SomeException -> IO a -rethrowWithContext sql (E.SomeException e) = +rethrowWithContext + :: (HasCallStack, IsSQL sql) + => sql + -> BackendPid + -> E.SomeException + -> IO a +rethrowWithContext sql pid (E.SomeException e) = E.throwIO DBException { dbeQueryContext = sql + , dbeBackendPid = pid , dbeError = e , dbeCallStack = callStack } diff --git a/src/Database/PostgreSQL/PQTypes/Internal/Monad.hs b/src/Database/PostgreSQL/PQTypes/Internal/Monad.hs index 43fedda..01b5c97 100644 --- a/src/Database/PostgreSQL/PQTypes/Internal/Monad.hs +++ b/src/Database/PostgreSQL/PQTypes/Internal/Monad.hs @@ -17,7 +17,6 @@ import Control.Monad.State.Strict import Control.Monad.Trans.Control import Control.Monad.Trans.State.Strict qualified as S import Control.Monad.Writer.Class -import Data.Bifunctor import GHC.Stack import Database.PostgreSQL.PQTypes.Class @@ -77,9 +76,9 @@ mapDBT f g m = DBT . StateT $ g . runStateT (unDBT m) . f instance (m ~ n, MonadBase IO m, MonadMask m) => MonadDB (DBT_ m n) where runQuery sql = withFrozenCallStack $ DBT . StateT $ \st -> liftBase $ do - second (updateStateWith st sql) <$> runQueryIO (dbConnection st) sql + updateStateWith st sql =<< runQueryIO (dbConnection st) sql runPreparedQuery name sql = withFrozenCallStack $ DBT . StateT $ \st -> liftBase $ do - second (updateStateWith st sql) <$> runPreparedQueryIO (dbConnection st) name sql + updateStateWith st sql =<< runPreparedQueryIO (dbConnection st) name sql getLastQuery = DBT . gets $ dbLastQuery @@ -88,6 +87,9 @@ instance (m ~ n, MonadBase IO m, MonadMask m) => MonadDB (DBT_ m n) where (x, st'') <- runStateT (unDBT callback) st' pure (x, st'' {dbRecordLastQuery = dbRecordLastQuery st}) + getBackendPid = DBT . StateT $ \st -> do + (,st) <$> liftBase (getBackendPidIO $ dbConnection st) + getConnectionStats = withFrozenCallStack $ do mconn <- DBT $ liftBase . readMVar =<< gets (unConnection . dbConnection) case mconn of @@ -100,9 +102,8 @@ instance (m ~ n, MonadBase IO m, MonadMask m) => MonadDB (DBT_ m n) where getTransactionSettings = DBT . gets $ dbTransactionSettings setTransactionSettings ts = DBT . modify $ \st -> st {dbTransactionSettings = ts} - getNotification time = DBT . StateT $ \st -> - (,st) - <$> liftBase (getNotificationIO st time) + getNotification time = DBT . StateT $ \st -> do + (,st) <$> liftBase (getNotificationIO st time) withNewConnection m = DBT . StateT $ \st -> do let cs = dbConnectionSource st diff --git a/src/Database/PostgreSQL/PQTypes/Internal/QueryResult.hs b/src/Database/PostgreSQL/PQTypes/Internal/QueryResult.hs index f6906b9..60ebf8b 100644 --- a/src/Database/PostgreSQL/PQTypes/Internal/QueryResult.hs +++ b/src/Database/PostgreSQL/PQTypes/Internal/QueryResult.hs @@ -2,6 +2,7 @@ module Database.PostgreSQL.PQTypes.Internal.QueryResult ( QueryResult (..) + , mkQueryResult , ntuples , nfields @@ -24,6 +25,7 @@ import System.IO.Unsafe import Database.PostgreSQL.PQTypes.Format import Database.PostgreSQL.PQTypes.FromRow +import Database.PostgreSQL.PQTypes.Internal.BackendPid import Database.PostgreSQL.PQTypes.Internal.C.Interface import Database.PostgreSQL.PQTypes.Internal.C.Types import Database.PostgreSQL.PQTypes.Internal.Error @@ -35,12 +37,27 @@ import Database.PostgreSQL.PQTypes.SQL.Class -- extraction appropriately. data QueryResult t = forall row. FromRow row => QueryResult { qrSQL :: !SomeSQL + , qrBackendPid :: !BackendPid , qrResult :: !(ForeignPtr PGresult) , qrFromRow :: !(row -> t) } +mkQueryResult + :: (FromRow t, IsSQL sql) + => sql + -> BackendPid + -> ForeignPtr PGresult + -> QueryResult t +mkQueryResult sql pid res = + QueryResult + { qrSQL = SomeSQL sql + , qrBackendPid = pid + , qrResult = res + , qrFromRow = id + } + instance Functor QueryResult where - f `fmap` QueryResult ctx fres g = QueryResult ctx fres (f . g) + f `fmap` QueryResult ctx pid fres g = QueryResult ctx pid fres (f . g) instance Foldable QueryResult where foldr f acc = runIdentity . foldrImpl False (coerce f) acc @@ -77,7 +94,7 @@ foldImpl -> acc -> QueryResult t -> m acc -foldImpl initCtr termCtr advCtr strict f iacc (QueryResult (SomeSQL ctx) fres g) = +foldImpl initCtr termCtr advCtr strict f iacc (QueryResult (SomeSQL ctx) pid fres g) = unsafePerformIO $ withForeignPtr fres $ \res -> do -- This bit is referentially transparent iff appropriate -- FrowRow and FromSQL instances are (the ones provided @@ -87,6 +104,7 @@ foldImpl initCtr termCtr advCtr strict f iacc (QueryResult (SomeSQL ctx) fres g) E.throwIO DBException { dbeQueryContext = ctx + , dbeBackendPid = pid , dbeError = RowLengthMismatch { lengthExpected = pqVariablesP rowp @@ -101,7 +119,7 @@ foldImpl initCtr termCtr advCtr strict f iacc (QueryResult (SomeSQL ctx) fres g) then return acc else do -- mask asynchronous exceptions so they won't be wrapped in DBException - obj <- E.mask_ (g <$> fromRow res err 0 i `E.catch` rethrowWithContext ctx) + obj <- E.mask_ (g <$> fromRow res err 0 i `E.catch` rethrowWithContext ctx pid) worker `apply` (f obj =<< acc) $ advCtr i worker (pure iacc) =<< initCtr res where diff --git a/src/Database/PostgreSQL/PQTypes/Internal/State.hs b/src/Database/PostgreSQL/PQTypes/Internal/State.hs index 1d23b3b..9f8712f 100644 --- a/src/Database/PostgreSQL/PQTypes/Internal/State.hs +++ b/src/Database/PostgreSQL/PQTypes/Internal/State.hs @@ -29,15 +29,21 @@ data DBState m = DBState -- ^ Current query result. } -updateStateWith :: IsSQL sql => DBState m -> sql -> ForeignPtr PGresult -> DBState m -updateStateWith st sql res = - st - { dbLastQuery = if dbRecordLastQuery st then SomeSQL sql else dbLastQuery st - , dbQueryResult = - Just - QueryResult - { qrSQL = SomeSQL sql - , qrResult = res - , qrFromRow = id - } - } +updateStateWith + :: IsSQL sql + => DBState m + -> sql + -> (r, ForeignPtr PGresult) + -> IO (r, DBState m) +updateStateWith st sql (r, res) = do + pid <- getBackendPidIO $ dbConnection st + pure + ( r + , st + { dbLastQuery = + if dbRecordLastQuery st + then SomeSQL sql + else dbLastQuery st + , dbQueryResult = Just $ mkQueryResult sql pid res + } + ) diff --git a/src/Database/PostgreSQL/PQTypes/Utils.hs b/src/Database/PostgreSQL/PQTypes/Utils.hs index bfa291c..63c75c6 100644 --- a/src/Database/PostgreSQL/PQTypes/Utils.hs +++ b/src/Database/PostgreSQL/PQTypes/Utils.hs @@ -38,9 +38,11 @@ throwDB e = case fromException $ toException e of Just (dbe :: DBException) -> throwM dbe Nothing -> do SomeSQL sql <- getLastQuery + pid <- getBackendPid throwM DBException { dbeQueryContext = sql + , dbeBackendPid = pid , dbeError = e , dbeCallStack = callStack }