|
| 1 | +{-# LANGUAGE CApiFFI #-} |
| 2 | + |
| 3 | +-- | |
| 4 | +-- Programmatic capture of server notices (RAISE WARNING et al). |
| 5 | +-- |
| 6 | +-- Registers a libpq notice receiver that copies the structured diagnostic |
| 7 | +-- fields off each notice's PGresult before libpq frees it. Notices accumulate |
| 8 | +-- in a bounded per-connection buffer; sessions drain it after each command. |
| 9 | +-- |
| 10 | +-- This bypasses postgresql-libpq's NoticeBuffer machinery entirely |
| 11 | +-- ("enableNoticeReporting" / "getNotice"), which flattens every notice to its |
| 12 | +-- rendered text and loses severity, SQLSTATE, detail and hint. |
| 13 | +module Hasql.LibPq14.Notices |
| 14 | + ( Notice (..), |
| 15 | + NoticeChannel, |
| 16 | + newNoticeChannel, |
| 17 | + registerNoticeReceiver, |
| 18 | + destroyNoticeChannel, |
| 19 | + drainNotices, |
| 20 | + noticeChannelCapacity, |
| 21 | + ) |
| 22 | +where |
| 23 | + |
| 24 | +import qualified Data.ByteString as BS |
| 25 | +import Data.IORef |
| 26 | +import qualified Database.PostgreSQL.LibPQ as LibPQ |
| 27 | +import Database.PostgreSQL.LibPQ.Internal (PGconn, withConn) |
| 28 | +import Foreign.C.String (CString) |
| 29 | +import Foreign.C.Types (CInt (..)) |
| 30 | +import Foreign.Ptr (FunPtr, Ptr, |
| 31 | + freeHaskellFunPtr, nullPtr) |
| 32 | +import Hasql.Prelude |
| 33 | + |
| 34 | +-- | |
| 35 | +-- A single non-fatal message received from the server. |
| 36 | +-- |
| 37 | +-- Fields mirror the libpq @PG_DIAG_*@ diagnostics of the notice's PGresult. |
| 38 | +data Notice = Notice |
| 39 | + { -- | e.g. @WARNING@ or @NOTICE@ |
| 40 | + noticeSeverity :: BS.ByteString, |
| 41 | + -- | SQLSTATE code |
| 42 | + noticeCode :: BS.ByteString, |
| 43 | + -- | Primary human-readable message |
| 44 | + noticeMessage :: BS.ByteString, |
| 45 | + noticeDetail :: Maybe BS.ByteString, |
| 46 | + noticeHint :: Maybe BS.ByteString |
| 47 | + } |
| 48 | + deriving (Show, Eq) |
| 49 | + |
| 50 | +-- | Maximum notices buffered per connection. Overflow drops the oldest, |
| 51 | +-- bounding memory on connections that receive notices outside of sessions |
| 52 | +-- (e.g. dedicated LISTEN connections that never drain). |
| 53 | +noticeChannelCapacity :: Int |
| 54 | +noticeChannelCapacity = 100 |
| 55 | + |
| 56 | +-- | Callback signature libpq expects: @void (*)(void *arg, const PGresult *res)@. |
| 57 | +-- The result pointer is typed as @Ptr ()@ because postgresql-libpq's Internal |
| 58 | +-- module does not export @PGresult@; the same approach is used by |
| 59 | +-- "Hasql.LibPq14.Ffi" for its @PQresultStatus@ import. |
| 60 | +type NoticeReceiverCb = Ptr () -> Ptr () -> IO () |
| 61 | + |
| 62 | +foreign import ccall "wrapper" |
| 63 | + mkNoticeReceiver :: NoticeReceiverCb -> IO (FunPtr NoticeReceiverCb) |
| 64 | + |
| 65 | +foreign import capi "libpq-fe.h PQsetNoticeReceiver" |
| 66 | + pqSetNoticeReceiver :: Ptr PGconn -> FunPtr NoticeReceiverCb -> Ptr () -> IO (FunPtr NoticeReceiverCb) |
| 67 | + |
| 68 | +foreign import capi "libpq-fe.h PQresultErrorField" |
| 69 | + pqResultErrorField :: Ptr () -> CInt -> IO CString |
| 70 | + |
| 71 | +foreign import capi "postgres_ext.h value PG_DIAG_SEVERITY" diagSeverityField :: CInt |
| 72 | +foreign import capi "postgres_ext.h value PG_DIAG_SQLSTATE" diagSqlstateField :: CInt |
| 73 | +foreign import capi "postgres_ext.h value PG_DIAG_MESSAGE_PRIMARY" diagMessagePrimaryField :: CInt |
| 74 | +foreign import capi "postgres_ext.h value PG_DIAG_MESSAGE_DETAIL" diagMessageDetailField :: CInt |
| 75 | +foreign import capi "postgres_ext.h value PG_DIAG_MESSAGE_HINT" diagMessageHintField :: CInt |
| 76 | + |
| 77 | +-- | Per-connection channel: the accumulating buffer plus the registered |
| 78 | +-- receiver closure, so the 'FunPtr' can be freed on connection release. |
| 79 | +data NoticeChannel = NoticeChannel |
| 80 | + { noticeChannelRef :: !(IORef [Notice]), |
| 81 | + noticeChannelFunPtr :: !(FunPtr NoticeReceiverCb) |
| 82 | + } |
| 83 | + |
| 84 | +-- | Allocate an empty channel with its receiver closure already wired to it. |
| 85 | +newNoticeChannel :: IO NoticeChannel |
| 86 | +newNoticeChannel = do |
| 87 | + ref <- newIORef [] |
| 88 | + funPtr <- mkNoticeReceiver (\_ result -> receiveNotice ref result) |
| 89 | + pure (NoticeChannel ref funPtr) |
| 90 | + |
| 91 | +-- | Install the channel's receiver on the given connection. The previously |
| 92 | +-- installed receiver returned by libpq is dropped without freeing: if it was |
| 93 | +-- libpq's default handler it is a static address, and freeing foreign static |
| 94 | +-- function pointers is undefined behavior. |
| 95 | +registerNoticeReceiver :: NoticeChannel -> LibPQ.Connection -> IO () |
| 96 | +registerNoticeReceiver channel connection = |
| 97 | + withConn connection $ \conn -> do |
| 98 | + _ <- pqSetNoticeReceiver conn (noticeChannelFunPtr channel) nullPtr |
| 99 | + pure () |
| 100 | + |
| 101 | +-- | Free the receiver closure. Must be called exactly once per channel, |
| 102 | +-- after the connection using it is finished. |
| 103 | +destroyNoticeChannel :: NoticeChannel -> IO () |
| 104 | +destroyNoticeChannel = |
| 105 | + freeHaskellFunPtr . noticeChannelFunPtr |
| 106 | + |
| 107 | +-- | Remove and return all buffered notices, oldest first. |
| 108 | +drainNotices :: NoticeChannel -> IO [Notice] |
| 109 | +drainNotices channel = |
| 110 | + atomicModifyIORef' (noticeChannelRef channel) (\old -> ([], old)) |
| 111 | + |
| 112 | +-- | Receiver entry point: decode the PGresult's diagnostics and buffer them. |
| 113 | +-- Runs inside libpq's input processing, while the calling session holds the |
| 114 | +-- connection lock, so 'atomicModifyIORef'' suffices. |
| 115 | +receiveNotice :: IORef [Notice] -> Ptr () -> IO () |
| 116 | +receiveNotice buffer result = do |
| 117 | + mNotice <- decodeNotice result |
| 118 | + traverse_ (appendNotice buffer) mNotice |
| 119 | + |
| 120 | +decodeNotice :: Ptr () -> IO (Maybe Notice) |
| 121 | +decodeNotice result = do |
| 122 | + mSeverity <- field diagSeverityField |
| 123 | + mCode <- field diagSqlstateField |
| 124 | + mMessage <- field diagMessagePrimaryField |
| 125 | + mDetail <- field diagMessageDetailField |
| 126 | + mHint <- field diagMessageHintField |
| 127 | + case (mSeverity, mCode, mMessage) of |
| 128 | + (Just severity, Just code, Just message) -> |
| 129 | + pure (Just (Notice severity code message mDetail mHint)) |
| 130 | + _ -> |
| 131 | + pure Nothing |
| 132 | + where |
| 133 | + field :: |
| 134 | + CInt -> |
| 135 | + IO (Maybe BS.ByteString) |
| 136 | + field code = do |
| 137 | + cstr <- pqResultErrorField result code |
| 138 | + if cstr == nullPtr |
| 139 | + then pure Nothing |
| 140 | + else Just <$> BS.packCString cstr |
| 141 | + |
| 142 | +appendNotice :: IORef [Notice] -> Notice -> IO () |
| 143 | +appendNotice buffer notice = |
| 144 | + atomicModifyIORef' buffer $ \old -> |
| 145 | + let grown = old ++ [notice] |
| 146 | + excess = max 0 (length grown - noticeChannelCapacity) |
| 147 | + in (drop excess grown, ()) |
0 commit comments