f9ee357 · typeable/postgresql-simple

git.iko.soy/typeable/postgresql-simple

Merge pull request #16 from joeyadams/escapeByteaConn

bytea compatibility fix, tests, and error handling fixes

commitf9ee35767b7efa75a675822060c3d059c63c9c9f authorLeon P Smith <leon.p.smith+github@gmail.com> dateWed 9 May 2012, 02:08 (14 years ago) parenta7bee4b aa84100

8 files changed+209−84

compare with master

unifiedsplit

postgresql-simple.cabal+17−2 src/Database/PostgreSQL/Simple.hs+35−64 src/Database/PostgreSQL/Simple/Internal.hs+66 src/Database/PostgreSQL/Simple/ToField.hs+16−17 src/Database/PostgreSQL/Simple/Types.hs+1−1 test/Bytea.hs+36 test/Main.hs+23 test/WRITING-TESTS+15

postgresql-simple.cabal+17−2

1212Category: Database

1313Build-type: Simple

1414

15−Cabal-version: >=1.6

15+Cabal-version: >= 1.9.2

1616

1717Library

1818 hs-source-dirs: src

@@ -37,7 +37,6 @@ Library

3737 Build-depends:

3838 attoparsec >= 0.8.5.3,

3939 base < 5,

40− base16-bytestring,

4140 blaze-builder,

4241 blaze-textual,

4342 bytestring >= 0.9,

@@ -64,3 +63,19 @@ source-repository this

6463 type: git

6564 location: http://github.com/lpsmith/postgresql-simple

6665 tag: v0.1.1

66+

67+test-suite test

68+ type: exitcode-stdio-1.0

69+

70+ hs-source-dirs: test

71+ main-is: Main.hs

72+ other-modules:

73+ Bytea

74+

75+ build-depends: base

76+ , base16-bytestring

77+ , bytestring

78+ , cryptohash

79+ , HUnit

80+ , postgresql-simple

81+ , text

src/Database/PostgreSQL/Simple.hs+35−64

@@ -113,7 +113,6 @@ import Control.Exception

113113 ( Exception, onException, throw, throwIO, finally )

114114import Control.Monad (foldM)

115115import Data.ByteString (ByteString)

116−import Data.Char(ord)

117116import Data.Int (Int64)

118117import qualified Data.IntMap as IntMap

119118import Data.List (intersperse)

@@ -133,6 +132,8 @@ import Database.PostgreSQL.Simple.Internal as Base

133132import qualified Database.PostgreSQL.LibPQ as PQ

134133import Text.Regex.PCRE.Light (compile, caseless, match)

135134import qualified Data.ByteString.Char8 as B

135+import qualified Data.Text as T

136+import qualified Data.Text.Encoding as TE

136137import qualified Data.Vector as V

137138import Control.Monad.Trans.Reader

138139import Control.Monad.Trans.State.Strict

@@ -148,15 +149,6 @@ data FormatError = FormatError {

148149

149150instance Exception FormatError

150151

151−-- | Exception thrown if 'query' is used to perform an @INSERT@-like

152−-- operation, or 'execute' is used to perform a @SELECT@-like operation.

153−data QueryError = QueryError {

154− qeMessage :: String

155− , qeQuery :: Query

156− } deriving (Eq, Show, Typeable)

157−

158−instance Exception QueryError

159−

160152-- | Format a query string.

161153--

162154-- This function is exposed to help with debugging and logging. Do not

@@ -202,16 +194,29 @@ formatMany conn q@(Query template) qs = do

202194 \([^?]*)$"

203195 [caseless]

204196

205−escapeStringConn :: Connection -> ByteString -> IO (Maybe ByteString)

206−escapeStringConn conn s = withConnection conn $ \c -> do

207− PQ.escapeStringConn c s

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

198+escapeStringConn conn s =

199+ withConnection conn $ \c ->

200+ PQ.escapeStringConn c s >>= checkError c

201+

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

203+escapeByteaConn conn s =

204+ withConnection conn $ \c ->

205+ PQ.escapeByteaConn c s >>= checkError c

206+

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

208+checkError c (Just x) = return $ Right x

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

208210

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

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

211− where quote = inQuotes . fromByteString . maybe undefined id

212− sub (Plain b) = pure b

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

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

213+ where quote = either (\msg -> fmtError (utf8ToString msg) q xs)

214+ (inQuotes . fromByteString)

215+ utf8ToString = T.unpack . TE.decodeUtf8

216+ sub (Plain b) = pure b

217+ sub (Escape s) = quote <$> escapeStringConn conn s

218+ sub (EscapeBytea s) = quote <$> escapeByteaConn conn s

219+ sub (Many ys) = mconcat <$> mapM sub ys

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

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

217222 zipParams (t:ts) (p:ps) = t `mappend` p `mappend` zipParams ts ps

@@ -231,12 +236,6 @@ execute conn template qs = do

231236 result <- exec conn =<< formatQuery conn template qs

232237 finishExecute conn template result

233238

234−-- | A version of 'execute' that does not perform query substitution.

235−execute_ :: Connection -> Query -> IO Int64

236−execute_ conn q@(Query stmt) = do

237− result <- exec conn stmt

238− finishExecute conn q result

239−

240239-- | Execute a multi-row @INSERT@, @UPDATE@, or other SQL query that is not

241240-- expected to return results.

242241--

@@ -249,44 +248,6 @@ executeMany conn q qs = do

249248 result <- exec conn =<< formatMany conn q qs

250249 finishExecute conn q result

251250

252−finishExecute :: Connection -> Query -> PQ.Result -> IO Int64

253−finishExecute _conn q result = do

254− status <- PQ.resultStatus result

255− case status of

256− PQ.CommandOk -> do

257− ncols <- PQ.nfields result

258− if ncols /= 0

259− then throwIO $ QueryError ("execute resulted in " ++ show ncols ++

260− "-column result") q

261− else do

262− nstr <- PQ.cmdTuples result

263− return $ case nstr of

264− Nothing -> 0 -- is this appropriate?

265− Just str -> toInteger str

266− PQ.TuplesOk -> do

267− ncols <- PQ.nfields result

268− throwIO $ QueryError ("execute resulted in " ++ show ncols ++

269− "-column result") q

270− PQ.CopyIn -> fail "FIXME: postgresql-simple does not currently support COPY IN"

271− PQ.CopyOut -> fail "FIXME: postgresql-simple does not currently support COPY OUT"

272− _ -> do

273− errormsg <- maybe "" id <$> PQ.resultErrorMessage result

274− statusmsg <- PQ.resStatus status

275− state <- maybe "" id <$> PQ.resultErrorField result PQ.DiagSqlstate

276− throwIO $ SqlError { sqlState = state

277− , sqlNativeError = fromEnum status

278− , sqlErrorMsg = B.concat [ "execute: ", statusmsg

279− , ": ", errormsg ]}

280− where

281− toInteger str = B.foldl' delta 0 str

282− where

283− delta acc c =

284− if '0' <= c && c <= '9'

285− then 10 * acc + fromIntegral (ord c - ord '0')

286− else error ("finishExecute: not an int: " ++ B.unpack str)

287−

288−

289−

290251-- | Perform a @SELECT@ or other SQL query that is expected to return

291252-- results. All results are retrieved and converted before this

292253-- function returns.

@@ -467,6 +428,8 @@ finishQuery :: (FromRow r) => Connection -> Query -> PQ.Result -> IO [r]

467428finishQuery conn q result = do

468429 status <- PQ.resultStatus result

469430 case status of

431+ PQ.EmptyQuery ->

432+ throwIO $ QueryError "query: Empty query" q

470433 PQ.CommandOk -> do

471434 throwIO $ QueryError "query resulted in a command response" q

472435 PQ.TuplesOk -> do

@@ -494,6 +457,13 @@ finishQuery conn q result = do

494457 Errors [] -> throwIO $ ConversionFailed "" "" "unknown error"

495458 Errors [x] -> throwIO x

496459 Errors xs -> throwIO $ ManyErrors xs

460+ PQ.CopyOut ->

461+ throwIO $ QueryError "query: COPY TO is not supported" q

462+ PQ.CopyIn ->

463+ throwIO $ QueryError "query: COPY FROM is not supported" q

464+ PQ.BadResponse -> throwResultError "query" result status

465+ PQ.NonfatalError -> throwResultError "query" result status

466+ PQ.FatalError -> throwResultError "query" result status

497467

498468ellipsis :: ByteString -> ByteString

499469ellipsis bs

@@ -611,9 +581,10 @@ fmtError msg q xs = throw FormatError {

611581 , fmtQuery = q

612582 , fmtParams = map twiddle xs

613583 }

614− where twiddle (Plain b) = toByteString b

615− twiddle (Escape s) = s

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

584+ where twiddle (Plain b) = toByteString b

585+ twiddle (Escape s) = s

586+ twiddle (EscapeBytea s) = s

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

617588

618589-- $use

619590--

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

@@ -25,6 +25,10 @@ import Control.Applicative

2525import Control.Exception

2626import Control.Concurrent.MVar

2727import Data.ByteString(ByteString)

28+import qualified Data.ByteString as B

29+import qualified Data.ByteString.Char8 as B8

30+import Data.Char (ord)

31+import Data.Int (Int64)

2832import qualified Data.IntMap as IntMap

2933import Data.String

3034import Data.Typeable

@@ -33,6 +37,7 @@ import Database.PostgreSQL.LibPQ(Oid(..))

3337import qualified Database.PostgreSQL.LibPQ as PQ

3438import Database.PostgreSQL.Simple.BuiltinTypes (BuiltinType)

3539import Database.PostgreSQL.Simple.Ok

40+import Database.PostgreSQL.Simple.Types (Query(..))

3641import Control.Monad.Trans.State.Strict

3742import Control.Monad.Trans.Reader

3843import qualified Data.Vector as V

@@ -86,6 +91,15 @@ data SqlError = SqlError {

8691

8792instance Exception SqlError

8893

94+-- | Exception thrown if 'query' is used to perform an @INSERT@-like

95+-- operation, or 'execute' is used to perform a @SELECT@-like operation.

96+data QueryError = QueryError {

97+ qeMessage :: String

98+ , qeQuery :: Query

99+ } deriving (Eq, Show, Typeable)

100+

101+instance Exception QueryError

102+

89103data ConnectInfo = ConnectInfo {

90104 connectHost :: String

91105 , connectPort :: Word16

@@ -138,6 +152,8 @@ connectPostgreSQL connstr = do

138152 PQ.ConnectionOk -> do

139153 connectionHandle <- newMVar conn

140154 connectionObjects <- newMVar (IntMap.empty)

155+ let wconn = Connection{..}

156+ _ <- execute_ wconn "SET standard_conforming_strings TO on"

141157 return Connection{..}

142158 _ -> do

143159 msg <- maybe "connectPostgreSQL error" id <$> PQ.errorMessage conn

@@ -198,6 +214,56 @@ exec conn sql =

198214 Just res -> do

199215 return res

200216

217+-- | A version of 'execute' that does not perform query substitution.

218+execute_ :: Connection -> Query -> IO Int64

219+execute_ conn q@(Query stmt) = do

220+ result <- exec conn stmt

221+ finishExecute conn q result

222+

223+finishExecute :: Connection -> Query -> PQ.Result -> IO Int64

224+finishExecute _conn q result = do

225+ status <- PQ.resultStatus result

226+ case status of

227+ PQ.EmptyQuery -> throwIO $ QueryError "execute: Empty query" q

228+ PQ.CommandOk -> do

229+ ncols <- PQ.nfields result

230+ if ncols /= 0

231+ then throwIO $ QueryError ("execute resulted in " ++ show ncols ++

232+ "-column result") q

233+ else do

234+ nstr <- PQ.cmdTuples result

235+ return $ case nstr of

236+ Nothing -> 0 -- is this appropriate?

237+ Just str -> toInteger str

238+ PQ.TuplesOk -> do

239+ ncols <- PQ.nfields result

240+ throwIO $ QueryError ("execute resulted in " ++ show ncols ++

241+ "-column result") q

242+ PQ.CopyOut ->

243+ throwIO $ QueryError "execute: COPY TO is not supported" q

244+ PQ.CopyIn ->

245+ throwIO $ QueryError "execute: COPY FROM is not supported" q

246+ PQ.BadResponse -> throwResultError "execute" result status

247+ PQ.NonfatalError -> throwResultError "execute" result status

248+ PQ.FatalError -> throwResultError "execute" result status

249+ where

250+ toInteger str = B8.foldl' delta 0 str

251+ where

252+ delta acc c =

253+ if '0' <= c && c <= '9'

254+ then 10 * acc + fromIntegral (ord c - ord '0')

255+ else error ("finishExecute: not an int: " ++ B8.unpack str)

256+

257+throwResultError :: ByteString -> PQ.Result -> PQ.ExecStatus -> IO a

258+throwResultError context result status = do

259+ errormsg <- maybe "" id <$> PQ.resultErrorMessage result

260+ statusmsg <- PQ.resStatus status

261+ state <- maybe "" id <$> PQ.resultErrorField result PQ.DiagSqlstate

262+ throwIO $ SqlError { sqlState = state

263+ , sqlNativeError = fromEnum status

264+ , sqlErrorMsg = B.concat [ context, ": ", statusmsg

265+ , ": ", errormsg ]}

266+

201267disconnectedError :: SqlError

202268disconnectedError = SqlError {

203269 sqlNativeError = -1,

src/Database/PostgreSQL/Simple/ToField.hs+16−17

@@ -21,13 +21,10 @@ module Database.PostgreSQL.Simple.ToField

2121 , inQuotes

2222 ) where

2323

24−import Blaze.ByteString.Builder (Builder, fromByteString, fromLazyByteString,

25− toByteString)

24+import Blaze.ByteString.Builder (Builder, fromByteString, toByteString)

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

2726import Blaze.Text (integral, double, float)

2827import Data.ByteString (ByteString)

29−import qualified Data.ByteString.Base16 as B16

30−import qualified Data.ByteString.Base16.Lazy as L16

3128import Data.Int (Int8, Int16, Int32, Int64)

3229import Data.List (intersperse)

3330import Data.Monoid (mappend)

@@ -58,14 +55,18 @@ data Action =

5855 -- ^ Escape and enclose in quotes before substituting. Use for all

5956 -- text-like types, and anything else that may contain unsafe

6057 -- characters when rendered.

58+ | EscapeBytea ByteString

59+ -- ^ Escape binary data for use as a @bytea@ literal. Include surrounding

60+ -- quotes. This is used by the 'Binary' newtype wrapper.

6161 | Many [Action]

6262 -- ^ Concatenate a series of rendering actions.

6363 deriving (Typeable)

6464

6565instance Show Action where

66− show (Plain b) = "Plain " ++ show (toByteString b)

67− show (Escape b) = "Escape " ++ show b

68− show (Many b) = "Many " ++ show b

66+ show (Plain b) = "Plain " ++ show (toByteString b)

67+ show (Escape b) = "Escape " ++ show b

68+ show (EscapeBytea b) = "EscapeBytea " ++ show b

69+ show (Many b) = "Many " ++ show b

6970

7071-- | A type that may be used as a single parameter to a SQL query.

7172class ToField a where

@@ -88,16 +89,6 @@ instance (ToField a) => ToField (In [a]) where

8889 (intersperse (Plain (fromChar ',')) . map toField $ xs) ++

8990 [Plain (fromChar ')')]

9091

91−instance ToField (Binary SB.ByteString) where

92− toField (Binary bs) = Plain $ fromByteString "'\\x" `mappend`

93− fromByteString (B16.encode bs) `mappend`

94− fromChar '\''

95−

96−instance ToField (Binary LB.ByteString) where

97− toField (Binary bs) = Plain $ fromByteString "'\\x" `mappend`

98− fromLazyByteString (L16.encode bs) `mappend`

99− fromChar '\''

100−

10192renderNull :: Action

10293renderNull = Plain (fromByteString "null")

10394

@@ -168,6 +159,14 @@ instance ToField Double where

168159 | otherwise = Plain (double v)

169160 {-# INLINE toField #-}

170161

162+instance ToField (Binary SB.ByteString) where

163+ toField (Binary bs) = EscapeBytea bs

164+ {-# INLINE toField #-}

165+

166+instance ToField (Binary LB.ByteString) where

167+ toField (Binary bs) = (EscapeBytea . SB.concat . LB.toChunks) bs

168+ {-# INLINE toField #-}

169+

171170instance ToField SB.ByteString where

172171 toField = Escape

173172 {-# INLINE toField #-}

src/Database/PostgreSQL/Simple/Types.hs+1−1

@@ -107,7 +107,7 @@ newtype Only a = Only {

107107newtype In a = In a

108108 deriving (Eq, Ord, Read, Show, Typeable, Functor)

109109

110−-- | Wrap a mostly-binary string to be escaped in hexadecimal.

110+-- | Wrap binary data for use as a @bytea@ value.

111111newtype Binary a = Binary a

112112 deriving (Eq, Ord, Read, Show, Typeable, Functor)

113113

test/Bytea.hsnew file+36

@@ -0,0 +1,36 @@

1+{-# LANGUAGE OverloadedStrings #-}

2+module Bytea where

4+import Data.ByteString (ByteString)

5+import Data.Text (Text)

6+import Database.PostgreSQL.Simple

7+import Test.HUnit

9+import qualified Crypto.Hash.MD5 as MD5

10+import qualified Data.ByteString as B

11+import qualified Data.ByteString.Base16 as Base16

12+import qualified Data.Text.Encoding as TE

13+

14+testBytea :: Connection -> Test

15+testBytea conn = TestList

16+ [ testStr "empty" []

17+ , testStr "\"hello\"" $ map (fromIntegral . fromEnum) ("hello" :: String)

18+ , testStr "ascending" [0..255]

19+ , testStr "descending" [255,254..0]

20+ , testStr "ascending, doubled up" $ doubleUp [0..255]

21+ , testStr "descending, doubled up" $ doubleUp [255,254..0]

22+ ]

23+ where

24+ testStr label bytes = TestLabel label $ TestCase $ do

25+ let bs = B.pack bytes

26+

27+ [Only h] <- query conn "SELECT md5(?::bytea)" [Binary bs]

28+ assertBool "Haskell -> SQL conversion altered the string" $ md5 bs == h

29+

30+ [Only (Binary r)] <- query conn "SELECT ?::bytea" [Binary bs]

31+ assertBool "SQL -> Haskell conversion altered the string" $ bs == r

32+

33+ doubleUp = concatMap (\x -> [x, x])

34+

35+md5 :: ByteString -> Text

36+md5 = TE.decodeUtf8 . Base16.encode . MD5.hash

test/Main.hsnew file+23

@@ -0,0 +1,23 @@

1+{-# LANGUAGE OverloadedStrings #-}

2+{-# LANGUAGE NamedFieldPuns #-}

4+import Control.Exception (bracket)

5+import Control.Monad (when)

6+import Database.PostgreSQL.Simple

7+import System.Exit (exitFailure)

8+import System.IO

9+import Test.HUnit

10+

11+import Bytea

12+

13+tests :: Connection -> [Test]

14+tests conn =

15+ [ TestLabel "Bytea" $ testBytea conn

16+ ]

17+

18+main :: IO ()

19+main = do

20+ mapM_ (`hSetBuffering` LineBuffering) [stdout, stderr]

21+ bracket (connectPostgreSQL "") close $ \conn -> do

22+ Counts{cases, tried, errors, failures} <- runTestTT $ TestList $ tests conn

23+ when (cases /= tried || errors /= 0 || failures /= 0) $ exitFailure

test/WRITING-TESTSnew file+15

@@ -0,0 +1,15 @@

1+Main.hs is a small wrapper around HUnit that opens a database connection and

2+starts the HUnit text-based test controller.

4+To add a new module to the test suite, do the following:

6+ * Create a new module that exports a function of type Connection -> Test.

7+ 'Test' is basically a tree of 'Assertion's, where 'Assertion' is just a type

8+ alias for IO (). See Bytea.hs for reference.

10+ * Add an entry to 'tests' in Main.hs (along with the corresponding

11+ module import) so the test driver will know about the new module.

12+

13+ * Add the module to postgresql-simple.cabal, under test-suite > other-modules.

14+ Otherwise, the module will be left out of the tarball generated by

15+ `cabal sdist`, and tests will fail to build when installing from Hackage.