{-# LANGUAGE CPP #-}

------------------------------------------------------------------------------
-- |
-- Module:      Database.PostgreSQL.Simple.Internal.PQResultUtils
-- Copyright:   (c) 2011 MailRank, Inc.
--              (c) 2011-2012 Leon P Smith
-- License:     BSD3
-- Maintainer:  Leon P Smith <leon@melding-monads.com>
-- Stability:   experimental
--
------------------------------------------------------------------------------


module Database.PostgreSQL.Simple.Internal.PQResultUtils
    ( finishQueryWith
    , finishQueryWithV
    , finishQueryWithVU
    , getRowWith
    ) where

import           Control.Exception as E
import           Data.ByteString (ByteString)
import           Data.Foldable (for_)
import           Database.PostgreSQL.Simple.FromField (ResultError(..))
import           Database.PostgreSQL.Simple.Ok
import           Database.PostgreSQL.Simple.Types (Query(..))
import           Database.PostgreSQL.Simple.Internal as Base hiding (result, row)
import           Database.PostgreSQL.Simple.TypeInfo
import qualified Database.PostgreSQL.LibPQ as PQ
import qualified Data.ByteString.Char8 as B
import qualified Data.Vector as V
import qualified Data.Vector.Mutable as MV
import qualified Data.Vector.Unboxed as VU
import qualified Data.Vector.Unboxed.Mutable as MVU
import           Control.Monad.Trans.Reader
import           Control.Monad.Trans.State.Strict

finishQueryWith :: RowParser r -> Connection -> Query -> PQ.Result -> IO [r]
finishQueryWith :: forall r. RowParser r -> Connection -> Query -> Result -> IO [r]
finishQueryWith RowParser r
parser Connection
conn Query
q Result
result = forall a. Query -> Result -> IO a -> IO a
finishQueryWith' Query
q Result
result forall a b. (a -> b) -> a -> b
$ do
    Row
nrows <- Result -> IO Row
PQ.ntuples Result
result
    Column
ncols <- Result -> IO Column
PQ.nfields Result
result
    forall n a. (Ord n, Num n) => n -> n -> (n -> IO a) -> IO [a]
forM' Row
0 (Row
nrowsforall a. Num a => a -> a -> a
-Row
1) forall a b. (a -> b) -> a -> b
$ \Row
row ->
        forall r.
RowParser r -> Row -> Column -> Connection -> Result -> IO r
getRowWith RowParser r
parser Row
row Column
ncols Connection
conn Result
result

finishQueryWithV :: RowParser r -> Connection -> Query -> PQ.Result -> IO (V.Vector r)
finishQueryWithV :: forall r.
RowParser r -> Connection -> Query -> Result -> IO (Vector r)
finishQueryWithV RowParser r
parser Connection
conn Query
q Result
result = forall a. Query -> Result -> IO a -> IO a
finishQueryWith' Query
q Result
result forall a b. (a -> b) -> a -> b
$ do
    Row
nrows <- Result -> IO Row
PQ.ntuples Result
result
    let PQ.Row CInt
nrows' = Row
nrows
    Column
ncols <- Result -> IO Column
PQ.nfields Result
result
    MVector RealWorld r
mv <- forall (m :: * -> *) a.
PrimMonad m =>
Int -> m (MVector (PrimState m) a)
MV.unsafeNew (forall a b. (Integral a, Num b) => a -> b
fromIntegral CInt
nrows')
    forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ [ Row
0 .. Row
nrowsforall a. Num a => a -> a -> a
-Row
1 ] forall a b. (a -> b) -> a -> b
$ \Row
row -> do
        let PQ.Row CInt
row' = Row
row
        r
value <- forall r.
RowParser r -> Row -> Column -> Connection -> Result -> IO r
getRowWith RowParser r
parser Row
row Column
ncols Connection
conn Result
result
        forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> a -> m ()
MV.unsafeWrite MVector RealWorld r
mv (forall a b. (Integral a, Num b) => a -> b
fromIntegral CInt
row') r
value
    forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> m (Vector a)
V.unsafeFreeze MVector RealWorld r
mv

finishQueryWithVU :: VU.Unbox r => RowParser r -> Connection -> Query -> PQ.Result -> IO (VU.Vector r)
finishQueryWithVU :: forall r.
Unbox r =>
RowParser r -> Connection -> Query -> Result -> IO (Vector r)
finishQueryWithVU RowParser r
parser Connection
conn Query
q Result
result = forall a. Query -> Result -> IO a -> IO a
finishQueryWith' Query
q Result
result forall a b. (a -> b) -> a -> b
$ do
    Row
nrows <- Result -> IO Row
PQ.ntuples Result
result
    let PQ.Row CInt
nrows' = Row
nrows
    Column
ncols <- Result -> IO Column
PQ.nfields Result
result
    MVector RealWorld r
mv <- forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
MVU.unsafeNew (forall a b. (Integral a, Num b) => a -> b
fromIntegral CInt
nrows')
    forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ [ Row
0 .. Row
nrowsforall a. Num a => a -> a -> a
-Row
1 ] forall a b. (a -> b) -> a -> b
$ \Row
row -> do
        let PQ.Row CInt
row' = Row
row
        r
value <- forall r.
RowParser r -> Row -> Column -> Connection -> Result -> IO r
getRowWith RowParser r
parser Row
row Column
ncols Connection
conn Result
result
        forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
MVU.unsafeWrite MVector RealWorld r
mv (forall a b. (Integral a, Num b) => a -> b
fromIntegral CInt
row') r
value
    forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector RealWorld r
mv

finishQueryWith' :: Query -> PQ.Result -> IO a -> IO a
finishQueryWith' :: forall a. Query -> Result -> IO a -> IO a
finishQueryWith' Query
q Result
result IO a
k = do
  ExecStatus
status <- Result -> IO ExecStatus
PQ.resultStatus Result
result
  case ExecStatus
status of
    ExecStatus
PQ.TuplesOk -> IO a
k
    ExecStatus
PQ.EmptyQuery    -> forall {a}. String -> IO a
queryErr String
"query: Empty query"
    ExecStatus
PQ.CommandOk     -> forall {a}. String -> IO a
queryErr String
"query resulted in a command response (did you mean to use `execute` or forget a RETURNING?)"
    ExecStatus
PQ.CopyOut       -> forall {a}. String -> IO a
queryErr String
"query: COPY TO is not supported"
    ExecStatus
PQ.CopyIn        -> forall {a}. String -> IO a
queryErr String
"query: COPY FROM is not supported"
#if MIN_VERSION_postgresql_libpq(0,9,3)
    ExecStatus
PQ.CopyBoth      -> forall {a}. String -> IO a
queryErr String
"query: COPY BOTH is not supported"
#endif
#if MIN_VERSION_postgresql_libpq(0,9,2)
    ExecStatus
PQ.SingleTuple   -> forall {a}. String -> IO a
queryErr String
"query: single-row mode is not supported"
#endif
    ExecStatus
PQ.BadResponse   -> forall a. ByteString -> Result -> ExecStatus -> IO a
throwResultError ByteString
"query" Result
result ExecStatus
status
    ExecStatus
PQ.NonfatalError -> forall a. ByteString -> Result -> ExecStatus -> IO a
throwResultError ByteString
"query" Result
result ExecStatus
status
    ExecStatus
PQ.FatalError    -> forall a. ByteString -> Result -> ExecStatus -> IO a
throwResultError ByteString
"query" Result
result ExecStatus
status
  where
    queryErr :: String -> IO a
queryErr String
msg = forall e a. Exception e => e -> IO a
throwIO forall a b. (a -> b) -> a -> b
$ String -> Query -> QueryError
QueryError String
msg Query
q

getRowWith :: RowParser r -> PQ.Row -> PQ.Column -> Connection -> PQ.Result -> IO r
getRowWith :: forall r.
RowParser r -> Row -> Column -> Connection -> Result -> IO r
getRowWith RowParser r
parser Row
row Column
ncols Connection
conn Result
result = do
  let rw :: Row
rw = Row -> Result -> Row
Row Row
row Result
result
  let unCol :: Column -> Int
unCol (PQ.Col CInt
x) = forall a b. (Integral a, Num b) => a -> b
fromIntegral CInt
x :: Int
  Ok (r, Column)
okvc <- forall a. Conversion a -> Connection -> IO (Ok a)
runConversion (forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT (forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT (forall a. RowParser a -> ReaderT Row (StateT Column Conversion) a
unRP RowParser r
parser) Row
rw) Column
0) Connection
conn
  case Ok (r, Column)
okvc of
    Ok (r
val,Column
col) | Column
col forall a. Eq a => a -> a -> Bool
== Column
ncols -> forall (m :: * -> *) a. Monad m => a -> m a
return r
val
                 | Bool
otherwise -> do
                     [(TypeInfo, Maybe ByteString)]
vals <- forall n a. (Ord n, Num n) => n -> n -> (n -> IO a) -> IO [a]
forM' Column
0 (Column
ncolsforall a. Num a => a -> a -> a
-Column
1) forall a b. (a -> b) -> a -> b
$ \Column
c -> do
                         TypeInfo
tinfo <- Connection -> Oid -> IO TypeInfo
getTypeInfo Connection
conn forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Result -> Column -> IO Oid
PQ.ftype Result
result Column
c
                         Maybe ByteString
v <- Result -> Row -> Column -> IO (Maybe ByteString)
PQ.getvalue Result
result Row
row Column
c
                         forall (m :: * -> *) a. Monad m => a -> m a
return ( TypeInfo
tinfo
                                , forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ByteString -> ByteString
ellipsis Maybe ByteString
v       )
                     forall a e. Exception e => e -> a
throw (String -> Maybe Oid -> String -> String -> String -> ResultError
ConversionFailed
                      (forall a. Show a => a -> String
show (Column -> Int
unCol Column
ncols) forall a. [a] -> [a] -> [a]
++ String
" values: " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show [(TypeInfo, Maybe ByteString)]
vals)
                      forall a. Maybe a
Nothing
                      String
""
                      (forall a. Show a => a -> String
show (Column -> Int
unCol Column
col) forall a. [a] -> [a] -> [a]
++ String
" slots in target type")
                      String
"mismatch between number of columns to convert and number in target type")
    Errors []  -> forall e a. Exception e => e -> IO a
throwIO forall a b. (a -> b) -> a -> b
$ String -> Maybe Oid -> String -> String -> String -> ResultError
ConversionFailed String
"" forall a. Maybe a
Nothing String
"" String
"" String
"unknown error"
    Errors [SomeException
x] -> forall e a. Exception e => e -> IO a
throwIO SomeException
x
    Errors [SomeException]
xs  -> forall e a. Exception e => e -> IO a
throwIO forall a b. (a -> b) -> a -> b
$ [SomeException] -> ManyErrors
ManyErrors [SomeException]
xs

ellipsis :: ByteString -> ByteString
ellipsis :: ByteString -> ByteString
ellipsis ByteString
bs
    | ByteString -> Int
B.length ByteString
bs forall a. Ord a => a -> a -> Bool
> Int
15 = Int -> ByteString -> ByteString
B.take Int
10 ByteString
bs ByteString -> ByteString -> ByteString
`B.append` ByteString
"[...]"
    | Bool
otherwise        = ByteString
bs

forM' :: (Ord n, Num n) => n -> n -> (n -> IO a) -> IO [a]
forM' :: forall n a. (Ord n, Num n) => n -> n -> (n -> IO a) -> IO [a]
forM' n
lo n
hi n -> IO a
m = n -> [a] -> IO [a]
loop n
hi []
  where
    loop :: n -> [a] -> IO [a]
loop !n
n ![a]
as
      | n
n forall a. Ord a => a -> a -> Bool
< n
lo = forall (m :: * -> *) a. Monad m => a -> m a
return [a]
as
      | Bool
otherwise = do
           a
a <- n -> IO a
m n
n
           n -> [a] -> IO [a]
loop (n
nforall a. Num a => a -> a -> a
-n
1) (a
aforall a. a -> [a] -> [a]
:[a]
as)
{-# INLINE forM' #-}