diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,6 @@ +# hpqtypes-1.11.1.0 (2023-01-31) +* Add support for setting a custom role when establishing a connection. + # hpqtypes-1.11.0.0 (2023-01-18) * Require `resource-pool` >= 0.4 and adjust the `createPool` function to seamlessly accommodate future changes to the `resource-pool` library. diff --git a/hpqtypes.cabal b/hpqtypes.cabal --- a/hpqtypes.cabal +++ b/hpqtypes.cabal @@ -1,5 +1,5 @@ name: hpqtypes -version: 1.11.0.0 +version: 1.11.1.0 synopsis: Haskell bindings to libpqtypes description: Efficient and easy-to-use bindings to (slightly modified) @@ -21,7 +21,7 @@ category: Database build-type: Simple cabal-version: 1.24 -tested-with: GHC ==8.8.4 || ==8.10.7 || ==9.0.2 || ==9.2.4 || ==9.4.2 +tested-with: GHC ==8.8.4 || ==8.10.7 || ==9.0.2 || ==9.2.5 || ==9.4.4 extra-source-files: README.md , CHANGELOG.md @@ -86,7 +86,6 @@ , Database.PostgreSQL.PQTypes.Internal.Monad , Database.PostgreSQL.PQTypes.Internal.Notification , Database.PostgreSQL.PQTypes.Internal.QueryResult - , Database.PostgreSQL.PQTypes.Internal.Query , Database.PostgreSQL.PQTypes.Internal.State , Database.PostgreSQL.PQTypes.Internal.C.Put , Database.PostgreSQL.PQTypes.Internal.C.Types @@ -189,7 +188,6 @@ , monad-control >= 1.0.3 , mtl >= 2.1 , random >= 1.0 - , resource-pool , scientific , test-framework >= 0.8 , test-framework-hunit >= 0.3 diff --git a/src/Database/PostgreSQL/PQTypes/Class.hs b/src/Database/PostgreSQL/PQTypes/Class.hs --- a/src/Database/PostgreSQL/PQTypes/Class.hs +++ b/src/Database/PostgreSQL/PQTypes/Class.hs @@ -9,7 +9,6 @@ import Database.PostgreSQL.PQTypes.FromRow import Database.PostgreSQL.PQTypes.Internal.Connection import Database.PostgreSQL.PQTypes.Internal.Notification -import Database.PostgreSQL.PQTypes.Internal.Query import Database.PostgreSQL.PQTypes.Internal.QueryResult import Database.PostgreSQL.PQTypes.SQL.Class import Database.PostgreSQL.PQTypes.Transaction.Settings diff --git a/src/Database/PostgreSQL/PQTypes/Internal/Connection.hs b/src/Database/PostgreSQL/PQTypes/Internal/Connection.hs --- a/src/Database/PostgreSQL/PQTypes/Internal/Connection.hs +++ b/src/Database/PostgreSQL/PQTypes/Internal/Connection.hs @@ -1,4 +1,5 @@ -module Database.PostgreSQL.PQTypes.Internal.Connection ( +module Database.PostgreSQL.PQTypes.Internal.Connection + ( -- * Connection Connection(..) , ConnectionData(..) , withConnectionData @@ -11,22 +12,29 @@ , poolSource , connect , disconnect + -- * Running queries + , runQueryIO + , QueryName(..) + , runPreparedQueryIO ) where -import Control.Arrow (first) import Control.Concurrent +import Control.Concurrent.Async import Control.Monad import Control.Monad.Base import Control.Monad.Catch +import Data.Bifunctor import Data.Function import Data.IORef import Data.Kind import Data.Pool +import Data.String import Foreign.C.String +import Foreign.ForeignPtr import Foreign.Ptr import GHC.Conc (closeFdWith) import qualified Control.Exception as E -import qualified Data.ByteString as BS +import qualified Data.ByteString.Char8 as BS import qualified Data.Foldable as F import qualified Data.Set as S import qualified Data.Text as T @@ -36,13 +44,20 @@ 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.Utils +import Database.PostgreSQL.PQTypes.SQL.Class +import Database.PostgreSQL.PQTypes.SQL.Raw +import Database.PostgreSQL.PQTypes.ToSQL data ConnectionSettings = ConnectionSettings { -- | Connection info string. csConnInfo :: !T.Text -- | Client-side encoding. If set to 'Nothing', database encoding is used. , csClientEncoding :: !(Maybe T.Text) + -- | A custom role to set with "SET ROLE". + , csRole :: !(Maybe (RawSQL ())) -- | A list of composite types to register. In order to be able to -- (de)serialize specific composite types, you need to register them. , csComposites :: ![T.Text] @@ -56,6 +71,7 @@ ConnectionSettings { csConnInfo = T.empty , csClientEncoding = Just "UTF-8" + , csRole = Nothing , csComposites = [] } @@ -171,23 +187,26 @@ -- 'disconnect', otherwise there will be a resource leak. connect :: ConnectionSettings -> IO Connection connect ConnectionSettings{..} = mask $ \unmask -> do - conn <- BS.useAsCString (T.encodeUtf8 csConnInfo) (openConnection unmask) - (`onException` c_PQfinish conn) . unmask $ do - status <- c_PQstatus conn + connPtr <- BS.useAsCString (T.encodeUtf8 csConnInfo) (openConnection unmask) + (`onException` c_PQfinish connPtr) . unmask $ do + status <- c_PQstatus connPtr when (status /= c_CONNECTION_OK) $ - throwLibPQError conn fname + throwLibPQError connPtr fname F.forM_ csClientEncoding $ \enc -> do - res <- BS.useAsCString (T.encodeUtf8 enc) (c_PQsetClientEncoding conn) + res <- BS.useAsCString (T.encodeUtf8 enc) (c_PQsetClientEncoding connPtr) when (res == -1) $ - throwLibPQError conn fname - c_PQinitTypes conn - registerComposites conn csComposites - preparedQueries <- newIORef S.empty - fmap Connection . newMVar $ Just ConnectionData - { cdPtr = conn - , cdStats = initialStats - , cdPreparedQueries = preparedQueries - } + throwLibPQError connPtr fname + c_PQinitTypes connPtr + registerComposites connPtr csComposites + conn <- do + preparedQueries <- newIORef S.empty + fmap Connection . newMVar $ Just ConnectionData + { cdPtr = connPtr + , cdStats = initialStats + , cdPreparedQueries = preparedQueries + } + F.forM_ csRole $ \role -> runQueryIO conn $ "SET ROLE " <> role + pure conn where fname = "connect" @@ -236,3 +255,152 @@ Nothing -> E.throwIO (HPQTypesError "disconnect: no connection (shouldn't happen)") return Nothing + +---------------------------------------- +-- Query running + +-- | Low-level function for running an SQL query. +runQueryIO + :: IsSQL sql + => Connection + -> sql + -> IO (Int, ForeignPtr PGresult) +runQueryIO conn sql = do + runQueryImpl "runQueryIO" conn sql $ \ConnectionData{..} -> do + let allocParam = ParamAllocator $ withPGparam cdPtr + withSQL sql allocParam $ \param query -> (,) + <$> (fromIntegral <$> c_PQparamCount param) + <*> c_PQparamExec cdPtr nullPtr param query c_RESULT_BINARY + +-- | Name of a prepared query. +newtype QueryName = QueryName T.Text + deriving (Eq, Ord, Show, IsString) + +-- | Low-level function for running a prepared SQL query. +runPreparedQueryIO + :: IsSQL sql + => Connection + -> QueryName + -> sql + -> IO (Int, ForeignPtr PGresult) +runPreparedQueryIO conn (QueryName queryName) sql = do + runQueryImpl "runPreparedQueryIO" conn sql $ \ConnectionData{..} -> do + when (T.null queryName) $ do + E.throwIO DBException + { dbeQueryContext = sql + , dbeError = HPQTypesError "runPreparedQueryIO: unnamed prepared query is not supported" + } + let allocParam = ParamAllocator $ withPGparam cdPtr + withSQL sql allocParam $ \param query -> do + preparedQueries <- readIORef cdPreparedQueries + BS.useAsCString (T.encodeUtf8 queryName) $ \cname -> do + when (queryName `S.notMember` preparedQueries) . E.mask_ $ do + -- Mask asynchronous exceptions, because if preparation of the query + -- 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 + modifyIORef' cdPreparedQueries $ S.insert queryName + (,) <$> (fromIntegral <$> c_PQparamCount param) + <*> c_PQparamExecPrepared cdPtr nullPtr param cname c_RESULT_BINARY + +-- | Shared implementation of 'runQueryIO' and 'runPreparedQueryIO'. +runQueryImpl + :: IsSQL sql + => String + -> Connection + -> sql + -> (ConnectionData -> IO (Int, ForeignPtr PGresult)) + -> IO (Int, ForeignPtr PGresult) +runQueryImpl fname conn sql execSql = do + withConnDo $ \cd@ConnectionData{..} -> E.mask $ \restore -> do + -- While the query runs, the current thread will not be able to receive + -- asynchronous exceptions. This prevents clients of the library from + -- interrupting execution of the query. To remedy that we spawn a separate + -- thread for the query execution and while we wait for its completion, we + -- are able to receive asynchronous exceptions (assuming that threaded GHC + -- runtime system is used) and react appropriately. + queryRunner <- async . restore $ do + (paramCount, res) <- execSql cd + affected <- withForeignPtr res $ verifyResult sql cdPtr + stats' <- case affected of + Left _ -> return cdStats { + statsQueries = statsQueries cdStats + 1 + , statsParams = statsParams cdStats + paramCount + } + Right rows -> do + columns <- fromIntegral <$> withForeignPtr res c_PQnfields + return ConnectionStats { + statsQueries = statsQueries cdStats + 1 + , statsRows = statsRows cdStats + rows + , 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)) + -- 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 + -- exception handler uninterruptible as we can't exit from the main block + -- until the query runner thread has terminated. + E.onException (restore $ wait queryRunner) . E.uninterruptibleMask_ $ do + c_PQcancel cdPtr >>= \case + -- If query cancellation request was successfully processed, there is + -- nothing else to do apart from waiting for the runner to terminate. + Nothing -> cancel queryRunner + -- Otherwise we check what happened with the runner. If it already + -- finished we're fine, just ignore the result. If it didn't, something + -- weird is going on. Maybe the cancellation request went through when + -- the thread wasn't making a request to the server? In any case, try to + -- cancel again and wait for the thread to terminate. + Just _ -> poll queryRunner >>= \case + Just _ -> return () + Nothing -> do + void $ c_PQcancel cdPtr + cancel queryRunner + where + withConnDo = withConnectionData conn fname + +verifyResult :: IsSQL sql => sql -> Ptr PGconn -> Ptr PGresult -> IO (Either Int Int) +verifyResult sql conn res = do + -- works even if res is NULL + rst <- c_PQresultStatus res + case rst of + _ | rst == c_PGRES_COMMAND_OK -> do + sn <- c_PQcmdTuples res >>= BS.packCString + case BS.readInt sn of + Nothing + | BS.null sn -> return . Left $ 0 + | otherwise -> throwParseError sn + Just (n, rest) + | rest /= BS.empty -> throwParseError sn + | otherwise -> return . Left $ n + _ | rst == c_PGRES_TUPLES_OK -> Right . fromIntegral <$> c_PQntuples res + _ | rst == c_PGRES_FATAL_ERROR -> throwSQLError + _ | rst == c_PGRES_BAD_RESPONSE -> throwSQLError + _ | otherwise -> return . Left $ 0 + where + throwSQLError = rethrowWithContext sql =<< if res == nullPtr + then return . E.toException . QueryError + =<< safePeekCString' =<< c_PQerrorMessage conn + else E.toException <$> (DetailedQueryError + <$> field c_PG_DIAG_SEVERITY + <*> (stringToErrorCode <$> field c_PG_DIAG_SQLSTATE) + <*> field c_PG_DIAG_MESSAGE_PRIMARY + <*> mfield c_PG_DIAG_MESSAGE_DETAIL + <*> mfield c_PG_DIAG_MESSAGE_HINT + <*> ((mread =<<) <$> mfield c_PG_DIAG_STATEMENT_POSITION) + <*> ((mread =<<) <$> mfield c_PG_DIAG_INTERNAL_POSITION) + <*> mfield c_PG_DIAG_INTERNAL_QUERY + <*> mfield c_PG_DIAG_CONTEXT + <*> mfield c_PG_DIAG_SOURCE_FILE + <*> ((mread =<<) <$> mfield c_PG_DIAG_SOURCE_LINE) + <*> mfield c_PG_DIAG_SOURCE_FUNCTION) + where + field f = maybe "" id <$> mfield f + mfield f = safePeekCString =<< c_PQresultErrorField res f + + throwParseError sn = E.throwIO DBException { + dbeQueryContext = sql + , dbeError = HPQTypesError ("verifyResult: string returned by PQcmdTuples is not a valid number: " ++ show sn) + } diff --git a/src/Database/PostgreSQL/PQTypes/Internal/Monad.hs b/src/Database/PostgreSQL/PQTypes/Internal/Monad.hs --- a/src/Database/PostgreSQL/PQTypes/Internal/Monad.hs +++ b/src/Database/PostgreSQL/PQTypes/Internal/Monad.hs @@ -14,6 +14,7 @@ import Control.Monad.State.Strict import Control.Monad.Trans.Control import Control.Monad.Writer.Class +import Data.Bifunctor import qualified Control.Monad.Trans.State.Strict as S import qualified Control.Monad.Fail as MF @@ -21,7 +22,6 @@ import Database.PostgreSQL.PQTypes.Internal.Connection import Database.PostgreSQL.PQTypes.Internal.Error import Database.PostgreSQL.PQTypes.Internal.Notification -import Database.PostgreSQL.PQTypes.Internal.Query import Database.PostgreSQL.PQTypes.Internal.State import Database.PostgreSQL.PQTypes.SQL import Database.PostgreSQL.PQTypes.SQL.Class @@ -71,8 +71,10 @@ ---------------------------------------- instance (m ~ n, MonadBase IO m, MonadMask m) => MonadDB (DBT_ m n) where - runQuery sql = DBT . StateT $ liftBase . runQueryIO sql - runPreparedQuery name sql = DBT . StateT $ liftBase . runPreparedQueryIO name sql + runQuery sql = DBT . StateT $ \st -> liftBase $ do + second (updateStateWith st sql) <$> runQueryIO (dbConnection st) sql + runPreparedQuery name sql = DBT . StateT $ \st -> liftBase $ do + second (updateStateWith st sql) <$> runPreparedQueryIO (dbConnection st) name sql getLastQuery = DBT . gets $ dbLastQuery diff --git a/src/Database/PostgreSQL/PQTypes/Internal/Query.hs b/src/Database/PostgreSQL/PQTypes/Internal/Query.hs deleted file mode 100644 --- a/src/Database/PostgreSQL/PQTypes/Internal/Query.hs +++ /dev/null @@ -1,186 +0,0 @@ -module Database.PostgreSQL.PQTypes.Internal.Query - ( runQueryIO - , QueryName(..) - , runPreparedQueryIO - ) where - -import Control.Concurrent.Async -import Control.Monad -import Data.IORef -import Data.String -import Foreign.ForeignPtr -import Foreign.Ptr -import qualified Control.Exception as E -import qualified Data.ByteString.Char8 as BS -import qualified Data.Text as T -import qualified Data.Text.Encoding as T -import qualified Data.Set as S - -import Database.PostgreSQL.PQTypes.Internal.C.Interface -import Database.PostgreSQL.PQTypes.Internal.C.Types -import Database.PostgreSQL.PQTypes.Internal.Connection -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.State -import Database.PostgreSQL.PQTypes.Internal.Utils -import Database.PostgreSQL.PQTypes.SQL.Class -import Database.PostgreSQL.PQTypes.ToSQL - --- | Low-level function for running an SQL query. -runQueryIO - :: IsSQL sql - => sql - -> DBState m - -> IO (Int, DBState m) -runQueryIO sql = runQueryImpl "runQueryIO" sql $ \ConnectionData{..} -> do - let allocParam = ParamAllocator $ withPGparam cdPtr - withSQL sql allocParam $ \param query -> (,) - <$> (fromIntegral <$> c_PQparamCount param) - <*> c_PQparamExec cdPtr nullPtr param query c_RESULT_BINARY - --- | Name of a prepared query. -newtype QueryName = QueryName T.Text - deriving (Eq, Ord, Show, IsString) - --- | Low-level function for running a prepared SQL query. -runPreparedQueryIO - :: IsSQL sql - => QueryName - -> sql - -> DBState m - -> IO (Int, DBState m) -runPreparedQueryIO (QueryName queryName) sql = do - runQueryImpl "runPreparedQueryIO" sql $ \ConnectionData{..} -> do - when (T.null queryName) $ do - E.throwIO DBException - { dbeQueryContext = sql - , dbeError = HPQTypesError "runPreparedQueryIO: unnamed prepared query is not supported" - } - let allocParam = ParamAllocator $ withPGparam cdPtr - withSQL sql allocParam $ \param query -> do - preparedQueries <- readIORef cdPreparedQueries - BS.useAsCString (T.encodeUtf8 queryName) $ \cname -> do - when (queryName `S.notMember` preparedQueries) . E.mask_ $ do - -- Mask asynchronous exceptions, because if preparation of the query - -- 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 - modifyIORef' cdPreparedQueries $ S.insert queryName - (,) <$> (fromIntegral <$> c_PQparamCount param) - <*> c_PQparamExecPrepared cdPtr nullPtr param cname c_RESULT_BINARY - ----------------------------------------- --- Helpers - --- | Shared implementation of 'runQueryIO' and 'runPreparedQueryIO'. -runQueryImpl - :: IsSQL sql - => String - -> sql - -> (ConnectionData -> IO (Int, ForeignPtr PGresult)) - -> DBState m - -> IO (Int, DBState m) -runQueryImpl fname sql execSql st = do - (affected, res) <- withConnDo $ \cd@ConnectionData{..} -> E.mask $ \restore -> do - -- While the query runs, the current thread will not be able to receive - -- asynchronous exceptions. This prevents clients of the library from - -- interrupting execution of the query. To remedy that we spawn a separate - -- thread for the query execution and while we wait for its completion, we - -- are able to receive asynchronous exceptions (assuming that threaded GHC - -- runtime system is used) and react appropriately. - queryRunner <- async . restore $ do - (paramCount, res) <- execSql cd - affected <- withForeignPtr res $ verifyResult sql cdPtr - stats' <- case affected of - Left _ -> return cdStats { - statsQueries = statsQueries cdStats + 1 - , statsParams = statsParams cdStats + paramCount - } - Right rows -> do - columns <- fromIntegral <$> withForeignPtr res c_PQnfields - return ConnectionStats { - statsQueries = statsQueries cdStats + 1 - , statsRows = statsRows cdStats + rows - , 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)) - -- 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 - -- exception handler uninterruptible as we can't exit from the main block - -- until the query runner thread has terminated. - E.onException (restore $ wait queryRunner) . E.uninterruptibleMask_ $ do - c_PQcancel cdPtr >>= \case - -- If query cancellation request was successfully processed, there is - -- nothing else to do apart from waiting for the runner to terminate. - Nothing -> cancel queryRunner - -- Otherwise we check what happened with the runner. If it already - -- finished we're fine, just ignore the result. If it didn't, something - -- weird is going on. Maybe the cancellation request went through when - -- the thread wasn't making a request to the server? In any case, try to - -- cancel again and wait for the thread to terminate. - Just _ -> poll queryRunner >>= \case - Just _ -> return () - Nothing -> do - void $ c_PQcancel cdPtr - cancel queryRunner - - return (affected, st { - dbLastQuery = if dbRecordLastQuery st then SomeSQL sql else dbLastQuery st - , dbQueryResult = Just QueryResult { - qrSQL = SomeSQL sql - , qrResult = res - , qrFromRow = id - } - }) - where - withConnDo = withConnectionData (dbConnection st) fname - -verifyResult :: IsSQL sql => sql -> Ptr PGconn -> Ptr PGresult -> IO (Either Int Int) -verifyResult sql conn res = do - -- works even if res is NULL - rst <- c_PQresultStatus res - case rst of - _ | rst == c_PGRES_COMMAND_OK -> do - sn <- c_PQcmdTuples res >>= BS.packCString - case BS.readInt sn of - Nothing - | BS.null sn -> return . Left $ 0 - | otherwise -> throwParseError sn - Just (n, rest) - | rest /= BS.empty -> throwParseError sn - | otherwise -> return . Left $ n - _ | rst == c_PGRES_TUPLES_OK -> Right . fromIntegral <$> c_PQntuples res - _ | rst == c_PGRES_FATAL_ERROR -> throwSQLError - _ | rst == c_PGRES_BAD_RESPONSE -> throwSQLError - _ | otherwise -> return . Left $ 0 - where - throwSQLError = rethrowWithContext sql =<< if res == nullPtr - then return . E.toException . QueryError - =<< safePeekCString' =<< c_PQerrorMessage conn - else E.toException <$> (DetailedQueryError - <$> field c_PG_DIAG_SEVERITY - <*> (stringToErrorCode <$> field c_PG_DIAG_SQLSTATE) - <*> field c_PG_DIAG_MESSAGE_PRIMARY - <*> mfield c_PG_DIAG_MESSAGE_DETAIL - <*> mfield c_PG_DIAG_MESSAGE_HINT - <*> ((mread =<<) <$> mfield c_PG_DIAG_STATEMENT_POSITION) - <*> ((mread =<<) <$> mfield c_PG_DIAG_INTERNAL_POSITION) - <*> mfield c_PG_DIAG_INTERNAL_QUERY - <*> mfield c_PG_DIAG_CONTEXT - <*> mfield c_PG_DIAG_SOURCE_FILE - <*> ((mread =<<) <$> mfield c_PG_DIAG_SOURCE_LINE) - <*> mfield c_PG_DIAG_SOURCE_FUNCTION) - where - field f = maybe "" id <$> mfield f - mfield f = safePeekCString =<< c_PQresultErrorField res f - - throwParseError sn = E.throwIO DBException { - dbeQueryContext = sql - , dbeError = HPQTypesError ("verifyResult: string returned by PQcmdTuples is not a valid number: " ++ show sn) - } diff --git a/src/Database/PostgreSQL/PQTypes/Internal/State.hs b/src/Database/PostgreSQL/PQTypes/Internal/State.hs --- a/src/Database/PostgreSQL/PQTypes/Internal/State.hs +++ b/src/Database/PostgreSQL/PQTypes/Internal/State.hs @@ -1,9 +1,13 @@ -- | Definition of internal DBT state. -module Database.PostgreSQL.PQTypes.Internal.State ( - DBState(..) +module Database.PostgreSQL.PQTypes.Internal.State + ( DBState(..) + , updateStateWith ) where +import Foreign.ForeignPtr + import Database.PostgreSQL.PQTypes.FromRow +import Database.PostgreSQL.PQTypes.Internal.C.Types import Database.PostgreSQL.PQTypes.Internal.Connection import Database.PostgreSQL.PQTypes.Internal.QueryResult import Database.PostgreSQL.PQTypes.SQL.Class @@ -23,4 +27,14 @@ , dbRecordLastQuery :: !Bool -- | Current query result. , dbQueryResult :: !(forall row. FromRow row => Maybe (QueryResult row)) + } + +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 + } } diff --git a/test/Main.hs b/test/Main.hs --- a/test/Main.hs +++ b/test/Main.hs @@ -12,7 +12,6 @@ import Data.Char import Data.Int import Data.Maybe -import Data.Pool import Data.Time import Data.Typeable import Data.Word @@ -56,10 +55,12 @@ ---------------------------------------- -type TestData = (QCGen, ConnectionSourceM IO) +type TestData = (QCGen, ConnectionSettings) runTestEnv :: TestData -> TransactionSettings -> TestEnv a -> IO a -runTestEnv (env, cs) ts m = runDBT cs ts $ evalStateT (unTestEnv m) env +runTestEnv (env, connSettings) ts m = runDBT cs ts $ evalStateT (unTestEnv m) env + where + ConnectionSource cs = simpleSource connSettings runTimes :: Monad m => Int -> m () -> m () runTimes !n m = case n of @@ -316,6 +317,33 @@ assertEqualEq "Other connection sees autocommited data" 1 n runQuery_ $ rawSQL "DELETE FROM test1_ WHERE a = $1" sint +setRoleTest :: TestData -> Test +setRoleTest td = do + testCase "SET ROLE works" $ bracket createRole dropRole $ \case + False -> putStrLn "Cannot create role, skipping SET ROLE test" + True -> do + runDBT roledCs defaultTransactionSettings $ do + runSQL_ "SELECT CURRENT_USER::text" + role <- fetchOne (runIdentity @String) + assertEqualEq "Role set successfully" testRole role + where + testRole :: String + testRole = "hpqtypes_test_role" + + ConnectionSource roledCs = simpleSource $ (snd td) + { csRole = Just $ unsafeSQL testRole + } + + createRole = runTestEnv td defaultTransactionSettings $ do + try (runSQL_ $ "CREATE ROLE" <+> unsafeSQL testRole) >>= \case + Right () -> pure True + Left DBException{} -> pure False + + dropRole = \case + False -> pure () + True -> runTestEnv td defaultTransactionSettings $ do + runSQL_ $ "DROP ROLE" <+> unsafeSQL testRole + preparedStatementTest :: TestData -> Test preparedStatementTest td = testCase "Execution of prepared statements works" . runTestEnv td defaultTransactionSettings $ do @@ -496,6 +524,7 @@ tests :: TestData -> [Test] tests td = [ autocommitTest td + , setRoleTest td , preparedStatementTest td , xmlTest td , readOnlyTest td @@ -653,18 +682,15 @@ main :: IO () main = do (connString, args) <- getConnString - let connSettings = defaultConnectionSettings { - csConnInfo = connString + let connSettings = defaultConnectionSettings + { csConnInfo = connString , csClientEncoding = Just "latin1" } ConnectionSource connSource = simpleSource connSettings createStructures connSource - ConnectionSource connPool <- - poolSource (connSettings { csComposites = ["simple_", "nested_"] }) - (\create destroy -> defaultPoolConfig create destroy 30 16) gen <- newQCGen putStrLn $ "PRNG:" <+> show gen - finally (defaultMainWithArgs (tests (gen, connPool)) $ args) $ do + finally (defaultMainWithArgs (tests (gen, connSettings { csComposites = ["simple_", "nested_"] })) $ args) $ do dropStructures connSource