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 -> DBState m -> IO (Int, DBState m)
runQueryIO sql
sql = String
-> sql
-> (ConnectionData -> IO (Int, ForeignPtr PGresult))
-> DBState m
-> IO (Int, DBState m)
forall sql (m :: * -> *).
IsSQL sql =>
String
-> sql
-> (ConnectionData -> IO (Int, ForeignPtr PGresult))
-> DBState m
-> IO (Int, DBState m)
runQueryImpl String
"runQueryIO" sql
sql ((ConnectionData -> IO (Int, ForeignPtr PGresult))
 -> DBState m -> IO (Int, DBState m))
-> (ConnectionData -> IO (Int, ForeignPtr PGresult))
-> DBState m
-> IO (Int, DBState m)
forall a b. (a -> b) -> a -> b
$ \ConnectionData{Ptr PGconn
IORef (Set Text)
ConnectionStats
cdPreparedQueries :: ConnectionData -> IORef (Set Text)
cdStats :: ConnectionData -> ConnectionStats
cdPtr :: ConnectionData -> Ptr PGconn
cdPreparedQueries :: IORef (Set Text)
cdStats :: ConnectionStats
cdPtr :: Ptr PGconn
..} -> do
  let allocParam :: ParamAllocator
allocParam = (forall r. (Ptr PGparam -> IO r) -> IO r) -> ParamAllocator
ParamAllocator ((forall r. (Ptr PGparam -> IO r) -> IO r) -> ParamAllocator)
-> (forall r. (Ptr PGparam -> IO r) -> IO r) -> ParamAllocator
forall a b. (a -> b) -> a -> b
$ Ptr PGconn -> (Ptr PGparam -> IO r) -> IO r
forall r. Ptr PGconn -> (Ptr PGparam -> IO r) -> IO r
withPGparam Ptr PGconn
cdPtr
  sql
-> ParamAllocator
-> (Ptr PGparam -> CString -> IO (Int, ForeignPtr PGresult))
-> IO (Int, ForeignPtr PGresult)
forall sql r.
IsSQL sql =>
sql -> ParamAllocator -> (Ptr PGparam -> CString -> IO r) -> IO r
withSQL sql
sql ParamAllocator
allocParam ((Ptr PGparam -> CString -> IO (Int, ForeignPtr PGresult))
 -> IO (Int, ForeignPtr PGresult))
-> (Ptr PGparam -> CString -> IO (Int, ForeignPtr PGresult))
-> IO (Int, ForeignPtr PGresult)
forall a b. (a -> b) -> a -> b
$ \Ptr PGparam
param CString
query -> (,)
    (Int -> ForeignPtr PGresult -> (Int, ForeignPtr PGresult))
-> IO Int -> IO (ForeignPtr PGresult -> (Int, ForeignPtr PGresult))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (CInt -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (CInt -> Int) -> IO CInt -> IO Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Ptr PGparam -> IO CInt
c_PQparamCount Ptr PGparam
param)
    IO (ForeignPtr PGresult -> (Int, ForeignPtr PGresult))
-> IO (ForeignPtr PGresult) -> IO (Int, ForeignPtr PGresult)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Ptr PGconn
-> Ptr PGerror
-> Ptr PGparam
-> CString
-> ResultFormat
-> IO (ForeignPtr PGresult)
c_PQparamExec Ptr PGconn
cdPtr Ptr PGerror
forall a. Ptr a
nullPtr Ptr PGparam
param CString
query ResultFormat
c_RESULT_BINARY

-- | Name of a prepared query.
newtype QueryName = QueryName T.Text
  deriving (QueryName -> QueryName -> Bool
(QueryName -> QueryName -> Bool)
-> (QueryName -> QueryName -> Bool) -> Eq QueryName
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: QueryName -> QueryName -> Bool
$c/= :: QueryName -> QueryName -> Bool
== :: QueryName -> QueryName -> Bool
$c== :: QueryName -> QueryName -> Bool
Eq, Eq QueryName
Eq QueryName
-> (QueryName -> QueryName -> Ordering)
-> (QueryName -> QueryName -> Bool)
-> (QueryName -> QueryName -> Bool)
-> (QueryName -> QueryName -> Bool)
-> (QueryName -> QueryName -> Bool)
-> (QueryName -> QueryName -> QueryName)
-> (QueryName -> QueryName -> QueryName)
-> Ord QueryName
QueryName -> QueryName -> Bool
QueryName -> QueryName -> Ordering
QueryName -> QueryName -> QueryName
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: QueryName -> QueryName -> QueryName
$cmin :: QueryName -> QueryName -> QueryName
max :: QueryName -> QueryName -> QueryName
$cmax :: QueryName -> QueryName -> QueryName
>= :: QueryName -> QueryName -> Bool
$c>= :: QueryName -> QueryName -> Bool
> :: QueryName -> QueryName -> Bool
$c> :: QueryName -> QueryName -> Bool
<= :: QueryName -> QueryName -> Bool
$c<= :: QueryName -> QueryName -> Bool
< :: QueryName -> QueryName -> Bool
$c< :: QueryName -> QueryName -> Bool
compare :: QueryName -> QueryName -> Ordering
$ccompare :: QueryName -> QueryName -> Ordering
$cp1Ord :: Eq QueryName
Ord, Int -> QueryName -> ShowS
[QueryName] -> ShowS
QueryName -> String
(Int -> QueryName -> ShowS)
-> (QueryName -> String)
-> ([QueryName] -> ShowS)
-> Show QueryName
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [QueryName] -> ShowS
$cshowList :: [QueryName] -> ShowS
show :: QueryName -> String
$cshow :: QueryName -> String
showsPrec :: Int -> QueryName -> ShowS
$cshowsPrec :: Int -> QueryName -> ShowS
Show, String -> QueryName
(String -> QueryName) -> IsString QueryName
forall a. (String -> a) -> IsString a
fromString :: String -> QueryName
$cfromString :: String -> QueryName
IsString)

-- | Low-level function for running a prepared SQL query.
runPreparedQueryIO
  :: IsSQL sql
  => QueryName
  -> sql
  -> DBState m
  -> IO (Int, DBState m)
runPreparedQueryIO :: QueryName -> sql -> DBState m -> IO (Int, DBState m)
runPreparedQueryIO (QueryName Text
queryName) sql
sql = do
  String
-> sql
-> (ConnectionData -> IO (Int, ForeignPtr PGresult))
-> DBState m
-> IO (Int, DBState m)
forall sql (m :: * -> *).
IsSQL sql =>
String
-> sql
-> (ConnectionData -> IO (Int, ForeignPtr PGresult))
-> DBState m
-> IO (Int, DBState m)
runQueryImpl String
"runPreparedQueryIO" sql
sql ((ConnectionData -> IO (Int, ForeignPtr PGresult))
 -> DBState m -> IO (Int, DBState m))
-> (ConnectionData -> IO (Int, ForeignPtr PGresult))
-> DBState m
-> IO (Int, DBState m)
forall a b. (a -> b) -> a -> b
$ \ConnectionData{Ptr PGconn
IORef (Set Text)
ConnectionStats
cdPreparedQueries :: IORef (Set Text)
cdStats :: ConnectionStats
cdPtr :: Ptr PGconn
cdPreparedQueries :: ConnectionData -> IORef (Set Text)
cdStats :: ConnectionData -> ConnectionStats
cdPtr :: ConnectionData -> Ptr PGconn
..} -> do
    Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Text -> Bool
T.null Text
queryName) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
      DBException -> IO ()
forall e a. Exception e => e -> IO a
E.throwIO DBException :: forall e sql. (Exception e, Show sql) => sql -> e -> DBException
DBException
        { dbeQueryContext :: sql
dbeQueryContext = sql
sql
        , dbeError :: HPQTypesError
dbeError = String -> HPQTypesError
HPQTypesError String
"runPreparedQueryIO: unnamed prepared query is not supported"
        }
    let allocParam :: ParamAllocator
allocParam = (forall r. (Ptr PGparam -> IO r) -> IO r) -> ParamAllocator
ParamAllocator ((forall r. (Ptr PGparam -> IO r) -> IO r) -> ParamAllocator)
-> (forall r. (Ptr PGparam -> IO r) -> IO r) -> ParamAllocator
forall a b. (a -> b) -> a -> b
$ Ptr PGconn -> (Ptr PGparam -> IO r) -> IO r
forall r. Ptr PGconn -> (Ptr PGparam -> IO r) -> IO r
withPGparam Ptr PGconn
cdPtr
    sql
-> ParamAllocator
-> (Ptr PGparam -> CString -> IO (Int, ForeignPtr PGresult))
-> IO (Int, ForeignPtr PGresult)
forall sql r.
IsSQL sql =>
sql -> ParamAllocator -> (Ptr PGparam -> CString -> IO r) -> IO r
withSQL sql
sql ParamAllocator
allocParam ((Ptr PGparam -> CString -> IO (Int, ForeignPtr PGresult))
 -> IO (Int, ForeignPtr PGresult))
-> (Ptr PGparam -> CString -> IO (Int, ForeignPtr PGresult))
-> IO (Int, ForeignPtr PGresult)
forall a b. (a -> b) -> a -> b
$ \Ptr PGparam
param CString
query -> do
      Set Text
preparedQueries <- IORef (Set Text) -> IO (Set Text)
forall a. IORef a -> IO a
readIORef IORef (Set Text)
cdPreparedQueries
      ByteString
-> (CString -> IO (Int, ForeignPtr PGresult))
-> IO (Int, ForeignPtr PGresult)
forall a. ByteString -> (CString -> IO a) -> IO a
BS.useAsCString (Text -> ByteString
T.encodeUtf8 Text
queryName) ((CString -> IO (Int, ForeignPtr PGresult))
 -> IO (Int, ForeignPtr PGresult))
-> (CString -> IO (Int, ForeignPtr PGresult))
-> IO (Int, ForeignPtr PGresult)
forall a b. (a -> b) -> a -> b
$ \CString
cname -> do
        Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Text
queryName Text -> Set Text -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.notMember` Set Text
preparedQueries) (IO () -> IO ()) -> (IO () -> IO ()) -> IO () -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IO () -> IO ()
forall a. IO a -> IO a
E.mask_ (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ 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.
          ForeignPtr PGresult
res <- Ptr PGconn
-> Ptr PGerror
-> Ptr PGparam
-> CString
-> CString
-> IO (ForeignPtr PGresult)
c_PQparamPrepare Ptr PGconn
cdPtr Ptr PGerror
forall a. Ptr a
nullPtr Ptr PGparam
param CString
cname CString
query
          IO (Either Int Int) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO (Either Int Int) -> IO ())
-> ((Ptr PGresult -> IO (Either Int Int)) -> IO (Either Int Int))
-> (Ptr PGresult -> IO (Either Int Int))
-> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ForeignPtr PGresult
-> (Ptr PGresult -> IO (Either Int Int)) -> IO (Either Int Int)
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr PGresult
res ((Ptr PGresult -> IO (Either Int Int)) -> IO ())
-> (Ptr PGresult -> IO (Either Int Int)) -> IO ()
forall a b. (a -> b) -> a -> b
$ sql -> Ptr PGconn -> Ptr PGresult -> IO (Either Int Int)
forall sql.
IsSQL sql =>
sql -> Ptr PGconn -> Ptr PGresult -> IO (Either Int Int)
verifyResult sql
sql Ptr PGconn
cdPtr
          IORef (Set Text) -> (Set Text -> Set Text) -> IO ()
forall a. IORef a -> (a -> a) -> IO ()
modifyIORef' IORef (Set Text)
cdPreparedQueries ((Set Text -> Set Text) -> IO ())
-> (Set Text -> Set Text) -> IO ()
forall a b. (a -> b) -> a -> b
$ Text -> Set Text -> Set Text
forall a. Ord a => a -> Set a -> Set a
S.insert Text
queryName
        (,) (Int -> ForeignPtr PGresult -> (Int, ForeignPtr PGresult))
-> IO Int -> IO (ForeignPtr PGresult -> (Int, ForeignPtr PGresult))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (CInt -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (CInt -> Int) -> IO CInt -> IO Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Ptr PGparam -> IO CInt
c_PQparamCount Ptr PGparam
param)
            IO (ForeignPtr PGresult -> (Int, ForeignPtr PGresult))
-> IO (ForeignPtr PGresult) -> IO (Int, ForeignPtr PGresult)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Ptr PGconn
-> Ptr PGerror
-> Ptr PGparam
-> CString
-> ResultFormat
-> IO (ForeignPtr PGresult)
c_PQparamExecPrepared Ptr PGconn
cdPtr Ptr PGerror
forall a. Ptr a
nullPtr Ptr PGparam
param CString
cname ResultFormat
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 :: String
-> sql
-> (ConnectionData -> IO (Int, ForeignPtr PGresult))
-> DBState m
-> IO (Int, DBState m)
runQueryImpl String
fname sql
sql ConnectionData -> IO (Int, ForeignPtr PGresult)
execSql DBState m
st = do
  (Int
affected, ForeignPtr PGresult
res) <- (ConnectionData -> IO (ConnectionData, (Int, ForeignPtr PGresult)))
-> IO (Int, ForeignPtr PGresult)
withConnDo ((ConnectionData
  -> IO (ConnectionData, (Int, ForeignPtr PGresult)))
 -> IO (Int, ForeignPtr PGresult))
-> (ConnectionData
    -> IO (ConnectionData, (Int, ForeignPtr PGresult)))
-> IO (Int, ForeignPtr PGresult)
forall a b. (a -> b) -> a -> b
$ \cd :: ConnectionData
cd@ConnectionData{Ptr PGconn
IORef (Set Text)
ConnectionStats
cdPreparedQueries :: IORef (Set Text)
cdStats :: ConnectionStats
cdPtr :: Ptr PGconn
cdPreparedQueries :: ConnectionData -> IORef (Set Text)
cdStats :: ConnectionData -> ConnectionStats
cdPtr :: ConnectionData -> Ptr PGconn
..} -> ((forall a. IO a -> IO a)
 -> IO (ConnectionData, (Int, ForeignPtr PGresult)))
-> IO (ConnectionData, (Int, ForeignPtr PGresult))
forall b. ((forall a. IO a -> IO a) -> IO b) -> IO b
E.mask (((forall a. IO a -> IO a)
  -> IO (ConnectionData, (Int, ForeignPtr PGresult)))
 -> IO (ConnectionData, (Int, ForeignPtr PGresult)))
-> ((forall a. IO a -> IO a)
    -> IO (ConnectionData, (Int, ForeignPtr PGresult)))
-> IO (ConnectionData, (Int, ForeignPtr PGresult))
forall a b. (a -> b) -> a -> b
$ \forall a. IO a -> IO a
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.
    Async (ConnectionData, (Int, ForeignPtr PGresult))
queryRunner <- IO (ConnectionData, (Int, ForeignPtr PGresult))
-> IO (Async (ConnectionData, (Int, ForeignPtr PGresult)))
forall a. IO a -> IO (Async a)
async (IO (ConnectionData, (Int, ForeignPtr PGresult))
 -> IO (Async (ConnectionData, (Int, ForeignPtr PGresult))))
-> (IO (ConnectionData, (Int, ForeignPtr PGresult))
    -> IO (ConnectionData, (Int, ForeignPtr PGresult)))
-> IO (ConnectionData, (Int, ForeignPtr PGresult))
-> IO (Async (ConnectionData, (Int, ForeignPtr PGresult)))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IO (ConnectionData, (Int, ForeignPtr PGresult))
-> IO (ConnectionData, (Int, ForeignPtr PGresult))
forall a. IO a -> IO a
restore (IO (ConnectionData, (Int, ForeignPtr PGresult))
 -> IO (Async (ConnectionData, (Int, ForeignPtr PGresult))))
-> IO (ConnectionData, (Int, ForeignPtr PGresult))
-> IO (Async (ConnectionData, (Int, ForeignPtr PGresult)))
forall a b. (a -> b) -> a -> b
$ do
      (Int
paramCount, ForeignPtr PGresult
res) <- ConnectionData -> IO (Int, ForeignPtr PGresult)
execSql ConnectionData
cd
      Either Int Int
affected <- ForeignPtr PGresult
-> (Ptr PGresult -> IO (Either Int Int)) -> IO (Either Int Int)
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr PGresult
res ((Ptr PGresult -> IO (Either Int Int)) -> IO (Either Int Int))
-> (Ptr PGresult -> IO (Either Int Int)) -> IO (Either Int Int)
forall a b. (a -> b) -> a -> b
$ sql -> Ptr PGconn -> Ptr PGresult -> IO (Either Int Int)
forall sql.
IsSQL sql =>
sql -> Ptr PGconn -> Ptr PGresult -> IO (Either Int Int)
verifyResult sql
sql Ptr PGconn
cdPtr
      ConnectionStats
stats' <- case Either Int Int
affected of
        Left Int
_ -> ConnectionStats -> IO ConnectionStats
forall (m :: * -> *) a. Monad m => a -> m a
return ConnectionStats
cdStats {
          statsQueries :: Int
statsQueries = ConnectionStats -> Int
statsQueries ConnectionStats
cdStats Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
        , statsParams :: Int
statsParams  = ConnectionStats -> Int
statsParams ConnectionStats
cdStats Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
paramCount
        }
        Right Int
rows -> do
          Int
columns <- CInt -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (CInt -> Int) -> IO CInt -> IO Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ForeignPtr PGresult -> (Ptr PGresult -> IO CInt) -> IO CInt
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr PGresult
res Ptr PGresult -> IO CInt
c_PQnfields
          ConnectionStats -> IO ConnectionStats
forall (m :: * -> *) a. Monad m => a -> m a
return ConnectionStats :: Int -> Int -> Int -> Int -> ConnectionStats
ConnectionStats {
            statsQueries :: Int
statsQueries = ConnectionStats -> Int
statsQueries ConnectionStats
cdStats Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
          , statsRows :: Int
statsRows    = ConnectionStats -> Int
statsRows ConnectionStats
cdStats Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
rows
          , statsValues :: Int
statsValues  = ConnectionStats -> Int
statsValues ConnectionStats
cdStats Int -> Int -> Int
forall a. Num a => a -> a -> a
+ (Int
rows Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
columns)
          , statsParams :: Int
statsParams  = ConnectionStats -> Int
statsParams ConnectionStats
cdStats Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
paramCount
          }
      -- Force evaluation of modified stats to squash a space leak.
      ConnectionStats
stats' ConnectionStats
-> IO (ConnectionData, (Int, ForeignPtr PGresult))
-> IO (ConnectionData, (Int, ForeignPtr PGresult))
`seq` (ConnectionData, (Int, ForeignPtr PGresult))
-> IO (ConnectionData, (Int, ForeignPtr PGresult))
forall (m :: * -> *) a. Monad m => a -> m a
return (ConnectionData
cd { cdStats :: ConnectionStats
cdStats = ConnectionStats
stats' }, ((Int -> Int) -> (Int -> Int) -> Either Int Int -> Int
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either Int -> Int
forall a. a -> a
id Int -> Int
forall a. a -> a
id Either Int Int
affected, ForeignPtr PGresult
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.
    IO (ConnectionData, (Int, ForeignPtr PGresult))
-> IO () -> IO (ConnectionData, (Int, ForeignPtr PGresult))
forall a b. IO a -> IO b -> IO a
E.onException (IO (ConnectionData, (Int, ForeignPtr PGresult))
-> IO (ConnectionData, (Int, ForeignPtr PGresult))
forall a. IO a -> IO a
restore (IO (ConnectionData, (Int, ForeignPtr PGresult))
 -> IO (ConnectionData, (Int, ForeignPtr PGresult)))
-> IO (ConnectionData, (Int, ForeignPtr PGresult))
-> IO (ConnectionData, (Int, ForeignPtr PGresult))
forall a b. (a -> b) -> a -> b
$ Async (ConnectionData, (Int, ForeignPtr PGresult))
-> IO (ConnectionData, (Int, ForeignPtr PGresult))
forall a. Async a -> IO a
wait Async (ConnectionData, (Int, ForeignPtr PGresult))
queryRunner) (IO () -> IO (ConnectionData, (Int, ForeignPtr PGresult)))
-> (IO () -> IO ())
-> IO ()
-> IO (ConnectionData, (Int, ForeignPtr PGresult))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IO () -> IO ()
forall a. IO a -> IO a
E.uninterruptibleMask_ (IO () -> IO (ConnectionData, (Int, ForeignPtr PGresult)))
-> IO () -> IO (ConnectionData, (Int, ForeignPtr PGresult))
forall a b. (a -> b) -> a -> b
$ do
      Ptr PGconn -> IO (Maybe String)
c_PQcancel Ptr PGconn
cdPtr IO (Maybe String) -> (Maybe String -> IO ()) -> IO ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
        -- If query cancellation request was successfully processed, there is
        -- nothing else to do apart from waiting for the runner to terminate.
        Maybe String
Nothing -> Async (ConnectionData, (Int, ForeignPtr PGresult)) -> IO ()
forall a. Async a -> IO ()
cancel Async (ConnectionData, (Int, ForeignPtr PGresult))
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 String
_ -> Async (ConnectionData, (Int, ForeignPtr PGresult))
-> IO
     (Maybe
        (Either
           SomeException (ConnectionData, (Int, ForeignPtr PGresult))))
forall a. Async a -> IO (Maybe (Either SomeException a))
poll Async (ConnectionData, (Int, ForeignPtr PGresult))
queryRunner IO
  (Maybe
     (Either
        SomeException (ConnectionData, (Int, ForeignPtr PGresult))))
-> (Maybe
      (Either SomeException (ConnectionData, (Int, ForeignPtr PGresult)))
    -> IO ())
-> IO ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
          Just Either SomeException (ConnectionData, (Int, ForeignPtr PGresult))
_  -> () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
          Maybe
  (Either SomeException (ConnectionData, (Int, ForeignPtr PGresult)))
Nothing -> do
            IO (Maybe String) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO (Maybe String) -> IO ()) -> IO (Maybe String) -> IO ()
forall a b. (a -> b) -> a -> b
$ Ptr PGconn -> IO (Maybe String)
c_PQcancel Ptr PGconn
cdPtr
            Async (ConnectionData, (Int, ForeignPtr PGresult)) -> IO ()
forall a. Async a -> IO ()
cancel Async (ConnectionData, (Int, ForeignPtr PGresult))
queryRunner

  (Int, DBState m) -> IO (Int, DBState m)
forall (m :: * -> *) a. Monad m => a -> m a
return (Int
affected, DBState m
st {
    dbLastQuery :: SomeSQL
dbLastQuery = if DBState m -> Bool
forall (m :: * -> *). DBState m -> Bool
dbRecordLastQuery DBState m
st then sql -> SomeSQL
forall sql. IsSQL sql => sql -> SomeSQL
SomeSQL sql
sql else DBState m -> SomeSQL
forall (m :: * -> *). DBState m -> SomeSQL
dbLastQuery DBState m
st
  , dbQueryResult :: forall row. FromRow row => Maybe (QueryResult row)
dbQueryResult = QueryResult row -> Maybe (QueryResult row)
forall a. a -> Maybe a
Just QueryResult :: forall t row.
FromRow row =>
SomeSQL -> ForeignPtr PGresult -> (row -> t) -> QueryResult t
QueryResult {
      qrSQL :: SomeSQL
qrSQL = sql -> SomeSQL
forall sql. IsSQL sql => sql -> SomeSQL
SomeSQL sql
sql
    , qrResult :: ForeignPtr PGresult
qrResult = ForeignPtr PGresult
res
    , qrFromRow :: row -> row
qrFromRow = row -> row
forall a. a -> a
id
    }
  })
  where
    withConnDo :: (ConnectionData -> IO (ConnectionData, (Int, ForeignPtr PGresult)))
-> IO (Int, ForeignPtr PGresult)
withConnDo = Connection
-> String
-> (ConnectionData
    -> IO (ConnectionData, (Int, ForeignPtr PGresult)))
-> IO (Int, ForeignPtr PGresult)
forall r.
Connection
-> String -> (ConnectionData -> IO (ConnectionData, r)) -> IO r
withConnectionData (DBState m -> Connection
forall (m :: * -> *). DBState m -> Connection
dbConnection DBState m
st) String
fname

verifyResult :: IsSQL sql => sql -> Ptr PGconn -> Ptr PGresult -> IO (Either Int Int)
verifyResult :: sql -> Ptr PGconn -> Ptr PGresult -> IO (Either Int Int)
verifyResult sql
sql Ptr PGconn
conn Ptr PGresult
res = do
  -- works even if res is NULL
  ExecStatusType
rst <- Ptr PGresult -> IO ExecStatusType
c_PQresultStatus Ptr PGresult
res
  case ExecStatusType
rst of
    ExecStatusType
_ | ExecStatusType
rst ExecStatusType -> ExecStatusType -> Bool
forall a. Eq a => a -> a -> Bool
== ExecStatusType
c_PGRES_COMMAND_OK -> do
      ByteString
sn <- Ptr PGresult -> IO CString
c_PQcmdTuples Ptr PGresult
res IO CString -> (CString -> IO ByteString) -> IO ByteString
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= CString -> IO ByteString
BS.packCString
      case ByteString -> Maybe (Int, ByteString)
BS.readInt ByteString
sn of
        Maybe (Int, ByteString)
Nothing
          | ByteString -> Bool
BS.null ByteString
sn -> Either Int Int -> IO (Either Int Int)
forall (m :: * -> *) a. Monad m => a -> m a
return (Either Int Int -> IO (Either Int Int))
-> (Int -> Either Int Int) -> Int -> IO (Either Int Int)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Either Int Int
forall a b. a -> Either a b
Left (Int -> IO (Either Int Int)) -> Int -> IO (Either Int Int)
forall a b. (a -> b) -> a -> b
$ Int
0
          | Bool
otherwise  -> ByteString -> IO (Either Int Int)
throwParseError ByteString
sn
        Just (Int
n, ByteString
rest)
          | ByteString
rest ByteString -> ByteString -> Bool
forall a. Eq a => a -> a -> Bool
/= ByteString
BS.empty -> ByteString -> IO (Either Int Int)
throwParseError ByteString
sn
          | Bool
otherwise        -> Either Int Int -> IO (Either Int Int)
forall (m :: * -> *) a. Monad m => a -> m a
return (Either Int Int -> IO (Either Int Int))
-> (Int -> Either Int Int) -> Int -> IO (Either Int Int)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Either Int Int
forall a b. a -> Either a b
Left (Int -> IO (Either Int Int)) -> Int -> IO (Either Int Int)
forall a b. (a -> b) -> a -> b
$ Int
n
    ExecStatusType
_ | ExecStatusType
rst ExecStatusType -> ExecStatusType -> Bool
forall a. Eq a => a -> a -> Bool
== ExecStatusType
c_PGRES_TUPLES_OK    -> Int -> Either Int Int
forall a b. b -> Either a b
Right (Int -> Either Int Int) -> (CInt -> Int) -> CInt -> Either Int Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CInt -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (CInt -> Either Int Int) -> IO CInt -> IO (Either Int Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Ptr PGresult -> IO CInt
c_PQntuples Ptr PGresult
res
    ExecStatusType
_ | ExecStatusType
rst ExecStatusType -> ExecStatusType -> Bool
forall a. Eq a => a -> a -> Bool
== ExecStatusType
c_PGRES_FATAL_ERROR  -> IO (Either Int Int)
throwSQLError
    ExecStatusType
_ | ExecStatusType
rst ExecStatusType -> ExecStatusType -> Bool
forall a. Eq a => a -> a -> Bool
== ExecStatusType
c_PGRES_BAD_RESPONSE -> IO (Either Int Int)
throwSQLError
    ExecStatusType
_ | Bool
otherwise                  -> Either Int Int -> IO (Either Int Int)
forall (m :: * -> *) a. Monad m => a -> m a
return (Either Int Int -> IO (Either Int Int))
-> (Int -> Either Int Int) -> Int -> IO (Either Int Int)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Either Int Int
forall a b. a -> Either a b
Left (Int -> IO (Either Int Int)) -> Int -> IO (Either Int Int)
forall a b. (a -> b) -> a -> b
$ Int
0
    where
      throwSQLError :: IO (Either Int Int)
throwSQLError = sql -> SomeException -> IO (Either Int Int)
forall sql a. IsSQL sql => sql -> SomeException -> IO a
rethrowWithContext sql
sql (SomeException -> IO (Either Int Int))
-> IO SomeException -> IO (Either Int Int)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< if Ptr PGresult
res Ptr PGresult -> Ptr PGresult -> Bool
forall a. Eq a => a -> a -> Bool
== Ptr PGresult
forall a. Ptr a
nullPtr
        then SomeException -> IO SomeException
forall (m :: * -> *) a. Monad m => a -> m a
return (SomeException -> IO SomeException)
-> (String -> SomeException) -> String -> IO SomeException
forall b c a. (b -> c) -> (a -> b) -> a -> c
. QueryError -> SomeException
forall e. Exception e => e -> SomeException
E.toException (QueryError -> SomeException)
-> (String -> QueryError) -> String -> SomeException
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> QueryError
QueryError
          (String -> IO SomeException) -> IO String -> IO SomeException
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< CString -> IO String
safePeekCString' (CString -> IO String) -> IO CString -> IO String
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Ptr PGconn -> IO CString
c_PQerrorMessage Ptr PGconn
conn
        else DetailedQueryError -> SomeException
forall e. Exception e => e -> SomeException
E.toException (DetailedQueryError -> SomeException)
-> IO DetailedQueryError -> IO SomeException
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (String
-> ErrorCode
-> String
-> Maybe String
-> Maybe String
-> Maybe Int
-> Maybe Int
-> Maybe String
-> Maybe String
-> Maybe String
-> Maybe Int
-> Maybe String
-> DetailedQueryError
DetailedQueryError
          (String
 -> ErrorCode
 -> String
 -> Maybe String
 -> Maybe String
 -> Maybe Int
 -> Maybe Int
 -> Maybe String
 -> Maybe String
 -> Maybe String
 -> Maybe Int
 -> Maybe String
 -> DetailedQueryError)
-> IO String
-> IO
     (ErrorCode
      -> String
      -> Maybe String
      -> Maybe String
      -> Maybe Int
      -> Maybe Int
      -> Maybe String
      -> Maybe String
      -> Maybe String
      -> Maybe Int
      -> Maybe String
      -> DetailedQueryError)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ErrorField -> IO String
field ErrorField
c_PG_DIAG_SEVERITY
          IO
  (ErrorCode
   -> String
   -> Maybe String
   -> Maybe String
   -> Maybe Int
   -> Maybe Int
   -> Maybe String
   -> Maybe String
   -> Maybe String
   -> Maybe Int
   -> Maybe String
   -> DetailedQueryError)
-> IO ErrorCode
-> IO
     (String
      -> Maybe String
      -> Maybe String
      -> Maybe Int
      -> Maybe Int
      -> Maybe String
      -> Maybe String
      -> Maybe String
      -> Maybe Int
      -> Maybe String
      -> DetailedQueryError)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (String -> ErrorCode
stringToErrorCode (String -> ErrorCode) -> IO String -> IO ErrorCode
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ErrorField -> IO String
field ErrorField
c_PG_DIAG_SQLSTATE)
          IO
  (String
   -> Maybe String
   -> Maybe String
   -> Maybe Int
   -> Maybe Int
   -> Maybe String
   -> Maybe String
   -> Maybe String
   -> Maybe Int
   -> Maybe String
   -> DetailedQueryError)
-> IO String
-> IO
     (Maybe String
      -> Maybe String
      -> Maybe Int
      -> Maybe Int
      -> Maybe String
      -> Maybe String
      -> Maybe String
      -> Maybe Int
      -> Maybe String
      -> DetailedQueryError)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ErrorField -> IO String
field ErrorField
c_PG_DIAG_MESSAGE_PRIMARY
          IO
  (Maybe String
   -> Maybe String
   -> Maybe Int
   -> Maybe Int
   -> Maybe String
   -> Maybe String
   -> Maybe String
   -> Maybe Int
   -> Maybe String
   -> DetailedQueryError)
-> IO (Maybe String)
-> IO
     (Maybe String
      -> Maybe Int
      -> Maybe Int
      -> Maybe String
      -> Maybe String
      -> Maybe String
      -> Maybe Int
      -> Maybe String
      -> DetailedQueryError)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ErrorField -> IO (Maybe String)
mfield ErrorField
c_PG_DIAG_MESSAGE_DETAIL
          IO
  (Maybe String
   -> Maybe Int
   -> Maybe Int
   -> Maybe String
   -> Maybe String
   -> Maybe String
   -> Maybe Int
   -> Maybe String
   -> DetailedQueryError)
-> IO (Maybe String)
-> IO
     (Maybe Int
      -> Maybe Int
      -> Maybe String
      -> Maybe String
      -> Maybe String
      -> Maybe Int
      -> Maybe String
      -> DetailedQueryError)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ErrorField -> IO (Maybe String)
mfield ErrorField
c_PG_DIAG_MESSAGE_HINT
          IO
  (Maybe Int
   -> Maybe Int
   -> Maybe String
   -> Maybe String
   -> Maybe String
   -> Maybe Int
   -> Maybe String
   -> DetailedQueryError)
-> IO (Maybe Int)
-> IO
     (Maybe Int
      -> Maybe String
      -> Maybe String
      -> Maybe String
      -> Maybe Int
      -> Maybe String
      -> DetailedQueryError)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ((String -> Maybe Int
forall a. Read a => String -> Maybe a
mread (String -> Maybe Int) -> Maybe String -> Maybe Int
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<<) (Maybe String -> Maybe Int) -> IO (Maybe String) -> IO (Maybe Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ErrorField -> IO (Maybe String)
mfield ErrorField
c_PG_DIAG_STATEMENT_POSITION)
          IO
  (Maybe Int
   -> Maybe String
   -> Maybe String
   -> Maybe String
   -> Maybe Int
   -> Maybe String
   -> DetailedQueryError)
-> IO (Maybe Int)
-> IO
     (Maybe String
      -> Maybe String
      -> Maybe String
      -> Maybe Int
      -> Maybe String
      -> DetailedQueryError)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ((String -> Maybe Int
forall a. Read a => String -> Maybe a
mread (String -> Maybe Int) -> Maybe String -> Maybe Int
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<<) (Maybe String -> Maybe Int) -> IO (Maybe String) -> IO (Maybe Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ErrorField -> IO (Maybe String)
mfield ErrorField
c_PG_DIAG_INTERNAL_POSITION)
          IO
  (Maybe String
   -> Maybe String
   -> Maybe String
   -> Maybe Int
   -> Maybe String
   -> DetailedQueryError)
-> IO (Maybe String)
-> IO
     (Maybe String
      -> Maybe String -> Maybe Int -> Maybe String -> DetailedQueryError)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ErrorField -> IO (Maybe String)
mfield ErrorField
c_PG_DIAG_INTERNAL_QUERY
          IO
  (Maybe String
   -> Maybe String -> Maybe Int -> Maybe String -> DetailedQueryError)
-> IO (Maybe String)
-> IO
     (Maybe String -> Maybe Int -> Maybe String -> DetailedQueryError)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ErrorField -> IO (Maybe String)
mfield ErrorField
c_PG_DIAG_CONTEXT
          IO
  (Maybe String -> Maybe Int -> Maybe String -> DetailedQueryError)
-> IO (Maybe String)
-> IO (Maybe Int -> Maybe String -> DetailedQueryError)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ErrorField -> IO (Maybe String)
mfield ErrorField
c_PG_DIAG_SOURCE_FILE
          IO (Maybe Int -> Maybe String -> DetailedQueryError)
-> IO (Maybe Int) -> IO (Maybe String -> DetailedQueryError)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ((String -> Maybe Int
forall a. Read a => String -> Maybe a
mread (String -> Maybe Int) -> Maybe String -> Maybe Int
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<<) (Maybe String -> Maybe Int) -> IO (Maybe String) -> IO (Maybe Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ErrorField -> IO (Maybe String)
mfield ErrorField
c_PG_DIAG_SOURCE_LINE)
          IO (Maybe String -> DetailedQueryError)
-> IO (Maybe String) -> IO DetailedQueryError
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ErrorField -> IO (Maybe String)
mfield ErrorField
c_PG_DIAG_SOURCE_FUNCTION)
        where
          field :: ErrorField -> IO String
field ErrorField
f = String -> ShowS -> Maybe String -> String
forall b a. b -> (a -> b) -> Maybe a -> b
maybe String
"" ShowS
forall a. a -> a
id (Maybe String -> String) -> IO (Maybe String) -> IO String
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ErrorField -> IO (Maybe String)
mfield ErrorField
f
          mfield :: ErrorField -> IO (Maybe String)
mfield ErrorField
f = CString -> IO (Maybe String)
safePeekCString (CString -> IO (Maybe String)) -> IO CString -> IO (Maybe String)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Ptr PGresult -> ErrorField -> IO CString
c_PQresultErrorField Ptr PGresult
res ErrorField
f

      throwParseError :: ByteString -> IO (Either Int Int)
throwParseError ByteString
sn = DBException -> IO (Either Int Int)
forall e a. Exception e => e -> IO a
E.throwIO DBException :: forall e sql. (Exception e, Show sql) => sql -> e -> DBException
DBException {
        dbeQueryContext :: sql
dbeQueryContext = sql
sql
      , dbeError :: HPQTypesError
dbeError = String -> HPQTypesError
HPQTypesError (String
"verifyResult: string returned by PQcmdTuples is not a valid number: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ ByteString -> String
forall a. Show a => a -> String
show ByteString
sn)
      }