module Database.PostgreSQL.PQTypes.Internal.QueryResult
  ( QueryResult (..)
  , mkQueryResult
  , ntuples
  , nfields

    -- * Implementation
  , foldrImpl
  , foldlImpl
  ) where

import Control.Exception qualified as E
import Control.Monad
import Data.Coerce
import Data.Foldable
import Data.Functor.Identity
import Foreign.C.Types
import Foreign.ForeignPtr
import Foreign.Marshal.Alloc
import Foreign.Ptr
import GHC.Stack
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
import Database.PostgreSQL.PQTypes.Internal.Exception
import Database.PostgreSQL.PQTypes.SQL.Class

-- | Representation of a query result. Provides 'Functor'
-- and 'Foldable' instances for data transformation and
-- extraction appropriately.
data QueryResult t = forall row. FromRow row => QueryResult
  { forall t. QueryResult t -> SomeSQL
qrSQL :: !SomeSQL
  , forall t. QueryResult t -> BackendPid
qrBackendPid :: !BackendPid
  , forall t. QueryResult t -> ForeignPtr PGresult
qrResult :: !(ForeignPtr PGresult)
  , ()
qrFromRow :: !(row -> t)
  }

mkQueryResult
  :: (FromRow t, IsSQL sql)
  => sql
  -> BackendPid
  -> ForeignPtr PGresult
  -> QueryResult t
mkQueryResult :: forall t sql.
(FromRow t, IsSQL sql) =>
sql -> BackendPid -> ForeignPtr PGresult -> QueryResult t
mkQueryResult sql
sql BackendPid
pid ForeignPtr PGresult
res =
  QueryResult
    { qrSQL :: SomeSQL
qrSQL = sql -> SomeSQL
forall sql. IsSQL sql => sql -> SomeSQL
SomeSQL sql
sql
    , qrBackendPid :: BackendPid
qrBackendPid = BackendPid
pid
    , qrResult :: ForeignPtr PGresult
qrResult = ForeignPtr PGresult
res
    , qrFromRow :: t -> t
qrFromRow = t -> t
forall a. a -> a
id
    }

instance Functor QueryResult where
  a -> b
f fmap :: forall a b. (a -> b) -> QueryResult a -> QueryResult b
`fmap` QueryResult SomeSQL
ctx BackendPid
pid ForeignPtr PGresult
fres row -> a
g = SomeSQL
-> BackendPid -> ForeignPtr PGresult -> (row -> b) -> QueryResult b
forall t row.
FromRow row =>
SomeSQL
-> BackendPid -> ForeignPtr PGresult -> (row -> t) -> QueryResult t
QueryResult SomeSQL
ctx BackendPid
pid ForeignPtr PGresult
fres (a -> b
f (a -> b) -> (row -> a) -> row -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. row -> a
g)

instance Foldable QueryResult where
  foldr :: forall a b. (a -> b -> b) -> b -> QueryResult a -> b
foldr a -> b -> b
f b
acc = Identity b -> b
forall a. Identity a -> a
runIdentity (Identity b -> b)
-> (QueryResult a -> Identity b) -> QueryResult a -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bool -> (a -> b -> Identity b) -> b -> QueryResult a -> Identity b
forall (m :: * -> *) t acc.
(HasCallStack, Monad m) =>
Bool -> (t -> acc -> m acc) -> acc -> QueryResult t -> m acc
foldrImpl Bool
False ((a -> b -> b) -> a -> b -> Identity b
forall a b. Coercible a b => a -> b
coerce a -> b -> b
f) b
acc
  foldr' :: forall a b. (a -> b -> b) -> b -> QueryResult a -> b
foldr' a -> b -> b
f b
acc = Identity b -> b
forall a. Identity a -> a
runIdentity (Identity b -> b)
-> (QueryResult a -> Identity b) -> QueryResult a -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bool -> (a -> b -> Identity b) -> b -> QueryResult a -> Identity b
forall (m :: * -> *) t acc.
(HasCallStack, Monad m) =>
Bool -> (t -> acc -> m acc) -> acc -> QueryResult t -> m acc
foldrImpl Bool
True ((a -> b -> b) -> a -> b -> Identity b
forall a b. Coercible a b => a -> b
coerce a -> b -> b
f) b
acc

  foldl :: forall b a. (b -> a -> b) -> b -> QueryResult a -> b
foldl b -> a -> b
f b
acc = Identity b -> b
forall a. Identity a -> a
runIdentity (Identity b -> b)
-> (QueryResult a -> Identity b) -> QueryResult a -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bool -> (b -> a -> Identity b) -> b -> QueryResult a -> Identity b
forall (m :: * -> *) acc t.
(HasCallStack, Monad m) =>
Bool -> (acc -> t -> m acc) -> acc -> QueryResult t -> m acc
foldlImpl Bool
False ((b -> a -> b) -> b -> a -> Identity b
forall a b. Coercible a b => a -> b
coerce b -> a -> b
f) b
acc
  foldl' :: forall b a. (b -> a -> b) -> b -> QueryResult a -> b
foldl' b -> a -> b
f b
acc = Identity b -> b
forall a. Identity a -> a
runIdentity (Identity b -> b)
-> (QueryResult a -> Identity b) -> QueryResult a -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bool -> (b -> a -> Identity b) -> b -> QueryResult a -> Identity b
forall (m :: * -> *) acc t.
(HasCallStack, Monad m) =>
Bool -> (acc -> t -> m acc) -> acc -> QueryResult t -> m acc
foldlImpl Bool
True ((b -> a -> b) -> b -> a -> Identity b
forall a b. Coercible a b => a -> b
coerce b -> a -> b
f) b
acc

foldrImpl
  :: (HasCallStack, Monad m)
  => Bool
  -> (t -> acc -> m acc)
  -> acc
  -> QueryResult t
  -> m acc
foldrImpl :: forall (m :: * -> *) t acc.
(HasCallStack, Monad m) =>
Bool -> (t -> acc -> m acc) -> acc -> QueryResult t -> m acc
foldrImpl = (Ptr PGresult -> IO CInt)
-> (Ptr PGresult -> IO CInt)
-> (CInt -> CInt)
-> Bool
-> (t -> acc -> m acc)
-> acc
-> QueryResult t
-> m acc
forall (m :: * -> *) t acc.
(HasCallStack, Monad m) =>
(Ptr PGresult -> IO CInt)
-> (Ptr PGresult -> IO CInt)
-> (CInt -> CInt)
-> Bool
-> (t -> acc -> m acc)
-> acc
-> QueryResult t
-> m acc
foldImpl ((CInt -> CInt) -> IO CInt -> IO CInt
forall a b. (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap CInt -> CInt
forall a. Enum a => a -> a
pred (IO CInt -> IO CInt)
-> (Ptr PGresult -> IO CInt) -> Ptr PGresult -> IO CInt
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ptr PGresult -> IO CInt
c_PQntuples) (IO CInt -> Ptr PGresult -> IO CInt
forall a b. a -> b -> a
const (IO CInt -> Ptr PGresult -> IO CInt)
-> (CInt -> IO CInt) -> CInt -> Ptr PGresult -> IO CInt
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CInt -> IO CInt
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (CInt -> Ptr PGresult -> IO CInt)
-> CInt -> Ptr PGresult -> IO CInt
forall a b. (a -> b) -> a -> b
$ -CInt
1) CInt -> CInt
forall a. Enum a => a -> a
pred

foldlImpl
  :: (HasCallStack, Monad m)
  => Bool
  -> (acc -> t -> m acc)
  -> acc
  -> QueryResult t
  -> m acc
foldlImpl :: forall (m :: * -> *) acc t.
(HasCallStack, Monad m) =>
Bool -> (acc -> t -> m acc) -> acc -> QueryResult t -> m acc
foldlImpl Bool
strict = (Ptr PGresult -> IO CInt)
-> (Ptr PGresult -> IO CInt)
-> (CInt -> CInt)
-> Bool
-> (t -> acc -> m acc)
-> acc
-> QueryResult t
-> m acc
forall (m :: * -> *) t acc.
(HasCallStack, Monad m) =>
(Ptr PGresult -> IO CInt)
-> (Ptr PGresult -> IO CInt)
-> (CInt -> CInt)
-> Bool
-> (t -> acc -> m acc)
-> acc
-> QueryResult t
-> m acc
foldImpl (IO CInt -> Ptr PGresult -> IO CInt
forall a b. a -> b -> a
const (IO CInt -> Ptr PGresult -> IO CInt)
-> IO CInt -> Ptr PGresult -> IO CInt
forall a b. (a -> b) -> a -> b
$ CInt -> IO CInt
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure CInt
0) Ptr PGresult -> IO CInt
c_PQntuples CInt -> CInt
forall a. Enum a => a -> a
succ Bool
strict ((t -> acc -> m acc) -> acc -> QueryResult t -> m acc)
-> ((acc -> t -> m acc) -> t -> acc -> m acc)
-> (acc -> t -> m acc)
-> acc
-> QueryResult t
-> m acc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (acc -> t -> m acc) -> t -> acc -> m acc
forall a b c. (a -> b -> c) -> b -> a -> c
flip

foldImpl
  :: (HasCallStack, Monad m)
  => (Ptr PGresult -> IO CInt)
  -> (Ptr PGresult -> IO CInt)
  -> (CInt -> CInt)
  -> Bool
  -> (t -> acc -> m acc)
  -> acc
  -> QueryResult t
  -> m acc
foldImpl :: forall (m :: * -> *) t acc.
(HasCallStack, Monad m) =>
(Ptr PGresult -> IO CInt)
-> (Ptr PGresult -> IO CInt)
-> (CInt -> CInt)
-> Bool
-> (t -> acc -> m acc)
-> acc
-> QueryResult t
-> m acc
foldImpl Ptr PGresult -> IO CInt
initCtr Ptr PGresult -> IO CInt
termCtr CInt -> CInt
advCtr Bool
strict t -> acc -> m acc
f acc
iacc (QueryResult (SomeSQL sql
ctx) BackendPid
pid ForeignPtr PGresult
fres row -> t
g) =
  IO (m acc) -> m acc
forall a. IO a -> a
unsafePerformIO (IO (m acc) -> m acc)
-> ((Ptr PGresult -> IO (m acc)) -> IO (m acc))
-> (Ptr PGresult -> IO (m acc))
-> m acc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ForeignPtr PGresult -> (Ptr PGresult -> IO (m acc)) -> IO (m acc)
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr PGresult
fres ((Ptr PGresult -> IO (m acc)) -> m acc)
-> (Ptr PGresult -> IO (m acc)) -> m acc
forall a b. (a -> b) -> a -> b
$ \Ptr PGresult
res -> do
    -- This bit is referentially transparent iff appropriate
    -- FrowRow and FromSQL instances are (the ones provided
    -- by the library fulfil this requirement).
    Int
rowlen <- 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 PGresult -> IO CInt
c_PQnfields Ptr PGresult
res
    Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
rowlen Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Proxy row -> Int
forall t. PQFormat t => Proxy t -> Int
pqVariablesP Proxy row
rowp) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
      DBException -> IO ()
forall e a. Exception e => e -> IO a
E.throwIO
        DBException
          { dbeQueryContext :: sql
dbeQueryContext = sql
ctx
          , dbeBackendPid :: BackendPid
dbeBackendPid = BackendPid
pid
          , dbeError :: RowLengthMismatch
dbeError =
              RowLengthMismatch
                { lengthExpected :: Int
lengthExpected = Proxy row -> Int
forall t. PQFormat t => Proxy t -> Int
pqVariablesP Proxy row
rowp
                , lengthDelivered :: Int
lengthDelivered = Int
rowlen
                }
          , dbeCallStack :: CallStack
dbeCallStack = CallStack
HasCallStack => CallStack
callStack
          }
    (Ptr PGerror -> IO (m acc)) -> IO (m acc)
forall a b. Storable a => (Ptr a -> IO b) -> IO b
alloca ((Ptr PGerror -> IO (m acc)) -> IO (m acc))
-> (Ptr PGerror -> IO (m acc)) -> IO (m acc)
forall a b. (a -> b) -> a -> b
$ \Ptr PGerror
err -> do
      CInt
n <- Ptr PGresult -> IO CInt
termCtr Ptr PGresult
res
      let worker :: m acc -> CInt -> IO (m acc)
worker m acc
acc CInt
i =
            if CInt
i CInt -> CInt -> Bool
forall a. Eq a => a -> a -> Bool
== CInt
n
              then m acc -> IO (m acc)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure m acc
acc
              else do
                -- mask asynchronous exceptions so they won't be wrapped in DBException
                t
obj <- IO t -> IO t
forall a. IO a -> IO a
E.mask_ (row -> t
g (row -> t) -> IO row -> IO t
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Ptr PGresult -> Ptr PGerror -> CInt -> CInt -> IO row
forall row.
FromRow row =>
Ptr PGresult -> Ptr PGerror -> CInt -> CInt -> IO row
fromRow Ptr PGresult
res Ptr PGerror
err CInt
0 CInt
i IO row -> (SomeException -> IO row) -> IO row
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
`E.catch` sql -> BackendPid -> SomeException -> IO row
forall sql a.
(HasCallStack, IsSQL sql) =>
sql -> BackendPid -> SomeException -> IO a
rethrowWithContext sql
ctx BackendPid
pid)
                m acc -> CInt -> IO (m acc)
worker (m acc -> CInt -> IO (m acc)) -> m acc -> CInt -> IO (m acc)
`apply` (t -> acc -> m acc
f t
obj (acc -> m acc) -> m acc -> m acc
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< m acc
acc) (CInt -> IO (m acc)) -> CInt -> IO (m acc)
forall a b. (a -> b) -> a -> b
$ CInt -> CInt
advCtr CInt
i
      m acc -> CInt -> IO (m acc)
worker (acc -> m acc
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure acc
iacc) (CInt -> IO (m acc)) -> IO CInt -> IO (m acc)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Ptr PGresult -> IO CInt
initCtr Ptr PGresult
res
  where
    -- ⊥ of existential type hidden in QueryResult
    row :: row
row = let t
_ = row -> t
g row
row in row
row
    rowp :: Proxy row
rowp = row -> Proxy row
forall a. a -> Proxy a
forall (f :: * -> *) a. Applicative f => a -> f a
pure row
row

    apply :: (m acc -> CInt -> IO (m acc)) -> m acc -> CInt -> IO (m acc)
apply = if Bool
strict then (m acc -> CInt -> IO (m acc)) -> m acc -> CInt -> IO (m acc)
forall a b. (a -> b) -> a -> b
($!) else (m acc -> CInt -> IO (m acc)) -> m acc -> CInt -> IO (m acc)
forall a b. (a -> b) -> a -> b
($)

-- Note: c_PQntuples/c_PQnfields are pure on a C level and QueryResult
-- constructor is not exported to the end user (so it's not possible
-- to enforce premature finalization via finalizeForeignPtr), which
-- makes usage of unsafeDupablePerformIO fine here.

-- | Extract number of returned tuples (rows) from query result.
ntuples :: QueryResult t -> Int
ntuples :: forall a. QueryResult a -> Int
ntuples QueryResult t
qr = IO Int -> Int
forall a. IO a -> a
unsafeDupablePerformIO (IO Int -> Int) -> IO Int -> Int
forall a b. (a -> b) -> a -> b
$ do
  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 (QueryResult t -> ForeignPtr PGresult
forall t. QueryResult t -> ForeignPtr PGresult
qrResult QueryResult t
qr) Ptr PGresult -> IO CInt
c_PQntuples

-- | Extract number of returned fields (columns) from query result.
nfields :: QueryResult t -> Int
nfields :: forall a. QueryResult a -> Int
nfields QueryResult t
qr = IO Int -> Int
forall a. IO a -> a
unsafeDupablePerformIO (IO Int -> Int) -> IO Int -> Int
forall a b. (a -> b) -> a -> b
$ do
  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 (QueryResult t -> ForeignPtr PGresult
forall t. QueryResult t -> ForeignPtr PGresult
qrResult QueryResult t
qr) Ptr PGresult -> IO CInt
c_PQnfields