f9ee357 · typeable/postgresql-simple

git.iko.soy/typeable/postgresql-simple

on master help sign in

Merge pull request #16 from joeyadams/escapeByteaConn

f9ee357Leon P Smith14 years ago

postgresql-simple.cabal+17−2

compare with master

unifiedsplit

12Category: Database

13Build-type: Simple

15−Cabal-version: >=1.6

15+Cabal-version: >= 1.9.2

17Library

18 hs-source-dirs: src

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

37 Build-depends:

38 attoparsec >= 0.8.5.3,

39 base < 5,

40− base16-bytestring,

41 blaze-builder,

40 blaze-builder,

42 blaze-textual,

41 blaze-textual,

43 bytestring >= 0.9,

42 bytestring >= 0.9,

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

64 type: git

63 type: git

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

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

66 tag: v0.1.1

65 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

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

114import Control.Monad (foldM)

115import Data.ByteString (ByteString)

116−import Data.Char(ord)

117import Data.Int (Int64)

116import Data.Int (Int64)

118import qualified Data.IntMap as IntMap

117import qualified Data.IntMap as IntMap

119import Data.List (intersperse)

118import Data.List (intersperse)

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

133import qualified Database.PostgreSQL.LibPQ as PQ

132import qualified Database.PostgreSQL.LibPQ as PQ

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

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

135import qualified Data.ByteString.Char8 as B

134import qualified Data.ByteString.Char8 as B

135+import qualified Data.Text as T

136+import qualified Data.Text.Encoding as TE

136import qualified Data.Vector as V

137import qualified Data.Vector as V

137import Control.Monad.Trans.Reader

138import Control.Monad.Trans.Reader

138import Control.Monad.Trans.State.Strict

139import Control.Monad.Trans.State.Strict

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

148

149

149instance Exception FormatError

150instance Exception FormatError

150

151

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−

160-- | Format a query string.

152-- | Format a query string.

161--

153--

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

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

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

202 \([^?]*)$"

194 \([^?]*)$"

203 [caseless]

195 [caseless]

204

196

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

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

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

198+escapeStringConn conn s =

207− PQ.escapeStringConn c 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

208

210

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

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

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

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

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

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

214+ (inQuotes . fromByteString)

215+ utf8ToString = T.unpack . TE.decodeUtf8

212− sub (Plain b) = pure b

216+ sub (Plain b) = pure b

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

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

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

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

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

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

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

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

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

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

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

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

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

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

232 finishExecute conn template result

237 finishExecute conn template result

233

238

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−

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

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

241-- expected to return results.

240-- expected to return results.

242--

241--

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

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

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

250 finishExecute conn q result

249 finishExecute conn q result

251

250

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−

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

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

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

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

292-- function returns.

253-- function returns.

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

467finishQuery conn q result = do

428finishQuery conn q result = do

468 status <- PQ.resultStatus result

429 status <- PQ.resultStatus result

469 case status of

430 case status of

431+ PQ.EmptyQuery ->

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

470 PQ.CommandOk -> do

433 PQ.CommandOk -> do

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

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

472 PQ.TuplesOk -> do

435 PQ.TuplesOk -> do

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

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

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

495 Errors [x] -> throwIO x

458 Errors [x] -> throwIO x

496 Errors xs -> throwIO $ ManyErrors xs

459 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

497

467

498ellipsis :: ByteString -> ByteString

468ellipsis :: ByteString -> ByteString

499ellipsis bs

469ellipsis bs

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

611 , fmtQuery = q

581 , fmtQuery = q

612 , fmtParams = map twiddle xs

582 , fmtParams = map twiddle xs

613 }

583 }

614− where twiddle (Plain b) = toByteString b

584+ where twiddle (Plain b) = toByteString b

615− twiddle (Escape s) = s

585+ twiddle (Escape s) = s

586+ twiddle (EscapeBytea s) = s

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

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

617

588

618-- $use

589-- $use

619--

590--

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

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

25import Control.Exception

26import Control.Concurrent.MVar

27import 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)

28import qualified Data.IntMap as IntMap

32import qualified Data.IntMap as IntMap

29import Data.String

33import Data.String

30import Data.Typeable

34import Data.Typeable

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

33import qualified Database.PostgreSQL.LibPQ as PQ

37import qualified Database.PostgreSQL.LibPQ as PQ

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

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

35import Database.PostgreSQL.Simple.Ok

39import Database.PostgreSQL.Simple.Ok

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

36import Control.Monad.Trans.State.Strict

41import Control.Monad.Trans.State.Strict

37import Control.Monad.Trans.Reader

42import Control.Monad.Trans.Reader

38import qualified Data.Vector as V

43import qualified Data.Vector as V

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

87instance Exception SqlError

92instance Exception SqlError

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+

89data ConnectInfo = ConnectInfo {

103data ConnectInfo = ConnectInfo {

90 connectHost :: String

104 connectHost :: String

91 , connectPort :: Word16

105 , connectPort :: Word16

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

138 PQ.ConnectionOk -> do

152 PQ.ConnectionOk -> do

139 connectionHandle <- newMVar conn

153 connectionHandle <- newMVar conn

140 connectionObjects <- newMVar (IntMap.empty)

154 connectionObjects <- newMVar (IntMap.empty)

155+ let wconn = Connection{..}

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

141 return Connection{..}

157 return Connection{..}

142 _ -> do

158 _ -> do

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

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

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

198 Just res -> do

214 Just res -> do

199 return res

215 return res

200

216

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+

201disconnectedError :: SqlError

267disconnectedError :: SqlError

202disconnectedError = SqlError {

268disconnectedError = SqlError {

203 sqlNativeError = -1,

269 sqlNativeError = -1,

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

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

21 , inQuotes

22 ) where

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

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

25− toByteString)

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

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

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

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

28import Data.ByteString (ByteString)

27import Data.ByteString (ByteString)

29−import qualified Data.ByteString.Base16 as B16

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

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

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

32import Data.List (intersperse)

29import Data.List (intersperse)

33import Data.Monoid (mappend)

30import Data.Monoid (mappend)

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

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

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

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

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

60 -- characters when rendered.

57 -- 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.

61 | Many [Action]

62 -- ^ Concatenate a series of rendering actions.

63 deriving (Typeable)

65instance Show Action where

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

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

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

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

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

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

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

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

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

71class ToField a where

72class ToField a where

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

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

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

89 [Plain (fromChar ')')]

90 [Plain (fromChar ')')]

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−

101renderNull :: Action

92renderNull :: Action

102renderNull = Plain (fromByteString "null")

93renderNull = Plain (fromByteString "null")

103

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

168 | otherwise = Plain (double v)

159 | otherwise = Plain (double v)

169 {-# INLINE toField #-}

160 {-# INLINE toField #-}

170

161

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+

171instance ToField SB.ByteString where

170instance ToField SB.ByteString where

172 toField = Escape

171 toField = Escape

173 {-# INLINE toField #-}

172 {-# INLINE toField #-}

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

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

107newtype In a = In a

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

109

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

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

111newtype Binary a = Binary a

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

113

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.