0f3f55a · typeable/postgresql-simple

git.iko.soy/typeable/postgresql-simple

on master help sign in

Merge pull request #138 from s9gf4ult/master

0f3f55aLeon P Smith12 years ago

src/Database/PostgreSQL/Simple.hs+13−60

compare with master

unifiedsplit

@@ -66,9 +66,9 @@ module Database.PostgreSQL.Simple

66 -- ** Exceptions

67 , SqlError(..)

68 , PQ.ExecStatus(..)

69− , FormatError(fmtMessage, fmtQuery, fmtParams)

69+ , FormatError(..)

70− , QueryError(qeMessage, qeQuery)

70+ , QueryError(..)

71− , ResultError(errSQLType, errHaskellType, errMessage)

71+ , ResultError(..)

72 -- * Connection management

73 , Base.connectPostgreSQL

74 , Base.close

@@ -114,19 +114,18 @@ import Blaze.ByteString.Builder

114 ( Builder, fromByteString, toByteString )

115import Blaze.ByteString.Builder.Char8 (fromChar)

116import Blaze.Text ( integral )

117−import Control.Applicative ((<$>), pure)

117+import Control.Applicative ((<$>))

118import Control.Exception as E

119import Control.Monad (foldM)

120import Data.ByteString (ByteString)

121import Data.Int (Int64)

122import Data.List (intersperse)

123import Data.Monoid (mconcat)

124−import Data.Typeable (Typeable)

125import Database.PostgreSQL.Simple.Compat ( (<>) )

124import Database.PostgreSQL.Simple.Compat ( (<>) )

126import Database.PostgreSQL.Simple.FromField (ResultError(..))

125import Database.PostgreSQL.Simple.FromField (ResultError(..))

127import Database.PostgreSQL.Simple.FromRow (FromRow(..))

126import Database.PostgreSQL.Simple.FromRow (FromRow(..))

128import Database.PostgreSQL.Simple.Ok

127import Database.PostgreSQL.Simple.Ok

129−import Database.PostgreSQL.Simple.ToField (Action(..), inQuotes)

128+import Database.PostgreSQL.Simple.ToField (Action(..))

130import Database.PostgreSQL.Simple.ToRow (ToRow(..))

129import Database.PostgreSQL.Simple.ToRow (ToRow(..))

131import Database.PostgreSQL.Simple.Types

130import Database.PostgreSQL.Simple.Types

132 ( Binary(..), In(..), Only(..), Query(..), (:.)(..) )

131 ( Binary(..), In(..), Only(..), Query(..), (:.)(..) )

@@ -135,21 +134,9 @@ import Database.PostgreSQL.Simple.Transaction

135import Database.PostgreSQL.Simple.TypeInfo

134import Database.PostgreSQL.Simple.TypeInfo

136import qualified Database.PostgreSQL.LibPQ as PQ

135import qualified Database.PostgreSQL.LibPQ as PQ

137import qualified Data.ByteString.Char8 as B

136import qualified Data.ByteString.Char8 as B

138−import qualified Data.Text as T

139−import qualified Data.Text.Encoding as TE

140import Control.Monad.Trans.Reader

137import Control.Monad.Trans.Reader

141import Control.Monad.Trans.State.Strict

138import Control.Monad.Trans.State.Strict

142

139

143−-- | Exception thrown if a 'Query' could not be formatted correctly.

144−-- This may occur if the number of \'@?@\' characters in the query

145−-- string does not match the number of parameters provided.

146−data FormatError = FormatError {

147− fmtMessage :: String

148− , fmtQuery :: Query

149− , fmtParams :: [ByteString]

150− } deriving (Eq, Show, Typeable)

151−

152−instance Exception FormatError

153

140

154-- | Format a query string.

141-- | Format a query string.

155--

142--

@@ -286,40 +273,17 @@ parseTemplate template =

286

273

287 skipSpace = B.dropWhile isSpace_ascii

274 skipSpace = B.dropWhile isSpace_ascii

288

275

289−escapeStringConn :: Connection -> ByteString -> IO (Either ByteString ByteString)

290−escapeStringConn = escapeWrap PQ.escapeStringConn

291−

292−escapeIdentifier :: Connection -> ByteString -> IO (Either ByteString ByteString)

293−escapeIdentifier = escapeWrap PQ.escapeIdentifier

294−

295−escapeByteaConn :: Connection -> ByteString -> IO (Either ByteString ByteString)

296−escapeByteaConn = escapeWrap PQ.escapeByteaConn

297−

298−escapeWrap :: (PQ.Connection -> ByteString -> IO (Maybe ByteString))

299− -> Connection

300− -> ByteString

301− -> IO (Either ByteString ByteString)

302−escapeWrap f conn s =

303− withConnection conn $ \c ->

304− f c s >>= checkError c

305

276

306−checkError :: PQ.Connection -> Maybe a -> IO (Either ByteString a)

307−checkError _ (Just x) = return $ Right x

308−checkError c Nothing = Left . maybe "" id <$> PQ.errorMessage c

309

277

310buildQuery :: Connection -> Query -> ByteString -> [Action] -> IO Builder

278buildQuery :: Connection -> Query -> ByteString -> [Action] -> IO Builder

311−buildQuery conn q template xs = zipParams (split template) <$> mapM sub xs

279+buildQuery conn q template xs =

312− where err' msg = fmtError (utf8ToString msg) q xs

280+ zipParams (split template) <$> mapM (buildAction conn q xs) xs

313− quote = either err' (inQuotes . fromByteString)

281+ where split s =

314− utf8ToString = T.unpack . TE.decodeUtf8

282+ let (h,t) = B.break (=='?') s

315− sub (Plain b) = pure b

283+ in fromByteString h

316− sub (Escape s) = quote <$> escapeStringConn conn s

284+ : if B.null t

317− sub (EscapeByteA s) = quote <$> escapeByteaConn conn s

285+ then []

318− sub (EscapeIdentifier s) = either err' fromByteString <$>

319− escapeIdentifier conn s

320− sub (Many ys) = mconcat <$> mapM sub ys

321− split s = fromByteString h : if B.null t then [] else split (B.tail t)

286+ else split (B.tail t)

322− where (h,t) = B.break (=='?') s

323 zipParams (t:ts) (p:ps) = t <> p <> zipParams ts ps

287 zipParams (t:ts) (p:ps) = t <> p <> zipParams ts ps

324 zipParams [t] [] = t

288 zipParams [t] [] = t

325 zipParams _ _ = fmtError (show (B.count '?' template) ++

289 zipParams _ _ = fmtError (show (B.count '?' template) ++

@@ -659,17 +623,6 @@ ellipsis bs

659 | B.length bs > 15 = B.take 10 bs `B.append` "[...]"

623 | B.length bs > 15 = B.take 10 bs `B.append` "[...]"

660 | otherwise = bs

624 | otherwise = bs

661

625

662−fmtError :: String -> Query -> [Action] -> a

663−fmtError msg q xs = throw FormatError {

664− fmtMessage = msg

665− , fmtQuery = q

666− , fmtParams = map twiddle xs

667− }

668− where twiddle (Plain b) = toByteString b

669− twiddle (Escape s) = s

670− twiddle (EscapeByteA s) = s

671− twiddle (EscapeIdentifier s) = s

672− twiddle (Many ys) = B.concat (map twiddle ys)

673

626

674-- $use

627-- $use

675--

628--

src/Database/PostgreSQL/Simple/Internal.hs+72

@@ -20,24 +20,31 @@

21module Database.PostgreSQL.Simple.Internal where

23+import Blaze.ByteString.Builder

24+ ( Builder, fromByteString, toByteString )

23import Control.Applicative

25import Control.Applicative

24import Control.Exception

26import Control.Exception

25import Control.Concurrent.MVar

27import Control.Concurrent.MVar

26import Control.Monad(MonadPlus(..))

28import Control.Monad(MonadPlus(..))

27import Data.ByteString(ByteString)

29import Data.ByteString(ByteString)

30+import qualified Data.ByteString as B

28import qualified Data.ByteString.Char8 as B8

31import qualified Data.ByteString.Char8 as B8

29import Data.Char (ord)

32import Data.Char (ord)

30import Data.Int (Int64)

33import Data.Int (Int64)

31import qualified Data.IntMap as IntMap

34import qualified Data.IntMap as IntMap

32import Data.IORef

35import Data.IORef

33import Data.Maybe(fromMaybe)

36import Data.Maybe(fromMaybe)

37+import Data.Monoid

34import Data.String

38import Data.String

39+import qualified Data.Text as T

40+import qualified Data.Text.Encoding as TE

35import Data.Typeable

41import Data.Typeable

36import Data.Word

42import Data.Word

37import Database.PostgreSQL.LibPQ(Oid(..))

43import Database.PostgreSQL.LibPQ(Oid(..))

38import qualified Database.PostgreSQL.LibPQ as PQ

44import qualified Database.PostgreSQL.LibPQ as PQ

39import Database.PostgreSQL.LibPQ(ExecStatus(..))

45import Database.PostgreSQL.LibPQ(ExecStatus(..))

40import Database.PostgreSQL.Simple.Ok

46import Database.PostgreSQL.Simple.Ok

47+import Database.PostgreSQL.Simple.ToField (Action(..), inQuotes)

41import Database.PostgreSQL.Simple.Types (Query(..))

48import Database.PostgreSQL.Simple.Types (Query(..))

42import Database.PostgreSQL.Simple.TypeInfo.Types(TypeInfo)

49import Database.PostgreSQL.Simple.TypeInfo.Types(TypeInfo)

43import Control.Monad.Trans.State.Strict

50import Control.Monad.Trans.State.Strict

@@ -92,6 +99,17 @@ data QueryError = QueryError {

93instance Exception QueryError

100instance Exception QueryError

101

102+-- | Exception thrown if a 'Query' could not be formatted correctly.

103+-- This may occur if the number of \'@?@\' characters in the query

104+-- string does not match the number of parameters provided.

105+data FormatError = FormatError {

106+ fmtMessage :: String

107+ , fmtQuery :: Query

108+ , fmtParams :: [ByteString]

109+ } deriving (Eq, Show, Typeable)

110+

111+instance Exception FormatError

112+

95data ConnectInfo = ConnectInfo {

113data ConnectInfo = ConnectInfo {

96 connectHost :: String

114 connectHost :: String

97 , connectPort :: Word16

115 , connectPort :: Word16

@@ -501,3 +519,57 @@ throwLibPQError :: PQ.Connection -> ByteString -> IO a

501throwLibPQError conn default_desc = do

519throwLibPQError conn default_desc = do

502 msg <- maybe default_desc id <$> PQ.errorMessage conn

520 msg <- maybe default_desc id <$> PQ.errorMessage conn

503 throwIO $! libPQError msg

521 throwIO $! libPQError msg

522+

523+

524+fmtError :: String -> Query -> [Action] -> a

525+fmtError msg q xs = throw FormatError {

526+ fmtMessage = msg

527+ , fmtQuery = q

528+ , fmtParams = map twiddle xs

529+ }

530+ where twiddle (Plain b) = toByteString b

531+ twiddle (Escape s) = s

532+ twiddle (EscapeByteA s) = s

533+ twiddle (EscapeIdentifier s) = s

534+ twiddle (Many ys) = B.concat (map twiddle ys)

535+

536+fmtErrorBs :: Query -> [Action] -> ByteString -> a

537+fmtErrorBs q xs msg = fmtError (T.unpack $ TE.decodeUtf8 msg) q xs

538+

539+-- | Quote bytestring or throw 'FormatError'

540+quote :: Query -> [Action] -> Either ByteString ByteString -> Builder

541+quote q xs = either (fmtErrorBs q xs) (inQuotes . fromByteString)

542+

543+buildAction :: Connection -- ^ Connection for string escaping

544+ -> Query -- ^ Query for message error

545+ -> [Action] -- ^ List of parameters for message error

546+ -> Action -- ^ Action to build

547+ -> IO Builder

548+buildAction _ _ _ (Plain b) = pure b

549+buildAction conn q xs (Escape s) = quote q xs <$> escapeStringConn conn s

550+buildAction conn q xs (EscapeByteA s) = quote q xs <$> escapeByteaConn conn s

551+buildAction conn q xs (EscapeIdentifier s) =

552+ either (fmtErrorBs q xs) fromByteString <$> escapeIdentifier conn s

553+buildAction conn q xs (Many ys) =

554+ mconcat <$> mapM (buildAction conn q xs) ys

555+

556+checkError :: PQ.Connection -> Maybe a -> IO (Either ByteString a)

557+checkError _ (Just x) = return $ Right x

558+checkError c Nothing = Left . maybe "" id <$> PQ.errorMessage c

559+

560+escapeWrap :: (PQ.Connection -> ByteString -> IO (Maybe ByteString))

561+ -> Connection

562+ -> ByteString

563+ -> IO (Either ByteString ByteString)

564+escapeWrap f conn s =

565+ withConnection conn $ \c ->

566+ f c s >>= checkError c

567+

568+escapeStringConn :: Connection -> ByteString -> IO (Either ByteString ByteString)

569+escapeStringConn = escapeWrap PQ.escapeStringConn

570+

571+escapeIdentifier :: Connection -> ByteString -> IO (Either ByteString ByteString)

572+escapeIdentifier = escapeWrap PQ.escapeIdentifier

573+

574+escapeByteaConn :: Connection -> ByteString -> IO (Either ByteString ByteString)

575+escapeByteaConn = escapeWrap PQ.escapeByteaConn