Added HStoreV2 & related tests.

postgresql-simple.cabal

@@ -30,6 +30,8 @@ Library
      Database.PostgreSQL.Simple.LargeObjects
      Database.PostgreSQL.Simple.HStore
      Database.PostgreSQL.Simple.HStore.Internal
+     Database.PostgreSQL.Simple.HStoreV2
+     Database.PostgreSQL.Simple.HStoreV2.Internal
      Database.PostgreSQL.Simple.Notification
      Database.PostgreSQL.Simple.Ok
      Database.PostgreSQL.Simple.Range
@@ -50,6 +52,7 @@ Library
   Other-modules:
      Database.PostgreSQL.Simple.Compat
      Database.PostgreSQL.Simple.HStore.Implementation
+     Database.PostgreSQL.Simple.HStoreV2.Implementation
      Database.PostgreSQL.Simple.Internal.PQResultUtils
      Database.PostgreSQL.Simple.Time.Implementation
      Database.PostgreSQL.Simple.Time.Internal.Parser

src/Database/PostgreSQL/Simple/HStoreV2.hs

@@ -0,0 +1,36 @@
+------------------------------------------------------------------------------
+-- |
+-- Module:      Database.PostgreSQL.Simple.HStore
+-- Copyright:   (c) 2013 Leon P Smith
+-- License:     BSD3
+-- Maintainer:  Leon P Smith <[email protected]>
+-- Stability:   experimental
+--
+-- Parsers and printers for hstore,  a extended type bundled with
+-- PostgreSQL providing finite maps from text strings to text strings.
+-- See <https://www.postgresql.org/docs/9.5/static/hstore.html> for more
+-- information.
+--
+-- Note that in order to use this type,  a database superuser must
+-- install it by running a sql script in the share directory.  This
+-- can be done on PostgreSQL 9.1 and later with the command
+-- @CREATE EXTENSION hstore@.  See
+-- <https://www.postgresql.org/docs/9.5/static/contrib.html> for more
+-- information.
+--
+------------------------------------------------------------------------------
+
+module Database.PostgreSQL.Simple.HStoreV2
+     ( HStoreList(..)
+     , HStoreMap(..)
+     , ToHStore(..)
+     , HStoreBuilder
+     , toBuilder
+     , toLazyByteString
+     , hstore
+     , parseHStoreList
+     , ToHStoreText(..)
+     , HStoreText
+     ) where
+
+import Database.PostgreSQL.Simple.HStoreV2.Implementation

src/Database/PostgreSQL/Simple/HStoreV2/Implementation.hs

@@ -0,0 +1,215 @@
+{-# LANGUAGE CPP, ViewPatterns, DeriveDataTypeable, GeneralizedNewtypeDeriving #-}
+
+------------------------------------------------------------------------------
+-- |
+-- Module:      Database.PostgreSQL.Simple.HStore.Implementation
+-- Copyright:   (c) 2013 Leon P Smith
+-- License:     BSD3
+-- Maintainer:  Leon P Smith <[email protected]>
+-- Stability:   experimental
+--
+-- This code has yet to be profiled and optimized.
+--
+------------------------------------------------------------------------------
+
+module Database.PostgreSQL.Simple.HStoreV2.Implementation where
+
+import           Control.Applicative
+import qualified Data.Attoparsec.ByteString as P
+import qualified Data.Attoparsec.ByteString.Char8 as P (isSpace_w8)
+import qualified Data.ByteString as BS
+import           Data.ByteString.Builder (Builder, byteString, char8)
+import qualified Data.ByteString.Builder as BU
+import           Data.ByteString.Internal (c2w, w2c)
+import qualified Data.ByteString.Lazy          as BL
+#if !MIN_VERSION_bytestring(0,10,0)
+import qualified Data.ByteString.Lazy.Internal as BL (foldrChunks)
+#endif
+import           Data.Map(Map)
+import qualified Data.Map as Map
+import           Data.Text(Text)
+import qualified Data.Text               as TS
+import qualified Data.Text.Encoding      as TS
+import           Data.Text.Encoding.Error(UnicodeException)
+import qualified Data.Text.Lazy          as TL
+import           Data.Typeable
+import           Data.Monoid(Monoid(..))
+import           Database.PostgreSQL.Simple.FromField
+import           Database.PostgreSQL.Simple.ToField
+
+class ToHStore a where
+   toHStore :: a -> HStoreBuilder
+
+-- | Represents valid hstore syntax.
+data HStoreBuilder
+   = Empty
+   | Comma !Builder
+     deriving (Typeable)
+
+instance ToHStore HStoreBuilder where
+   toHStore = id
+
+toBuilder :: HStoreBuilder -> Builder
+toBuilder x = case x of
+                Empty -> mempty
+                Comma x -> x
+
+toLazyByteString :: HStoreBuilder -> BL.ByteString
+toLazyByteString x = case x of
+                       Empty -> BL.empty
+                       Comma x -> BU.toLazyByteString x
+
+instance Monoid HStoreBuilder where
+    mempty = Empty
+    mappend Empty     x = x
+    mappend (Comma a) x
+        = Comma (a `mappend` case x of
+                               Empty   -> mempty
+                               Comma b -> char8 ',' `mappend` b)
+
+class ToHStoreText a where
+  toHStoreText :: a -> HStoreText
+
+-- | Represents escape text, ready to be the key or value to a hstore value
+newtype HStoreText = HStoreText Builder deriving (Typeable, Monoid)
+
+instance ToHStoreText HStoreText where
+  toHStoreText = id
+
+-- | Assumed to be UTF-8 encoded
+instance ToHStoreText BS.ByteString where
+  toHStoreText str = HStoreText (escapeAppend str mempty)
+
+-- | Assumed to be UTF-8 encoded
+instance ToHStoreText BL.ByteString where
+  toHStoreText = HStoreText . BL.foldrChunks escapeAppend mempty
+
+instance ToHStoreText TS.Text where
+  toHStoreText str = HStoreText (escapeAppend (TS.encodeUtf8 str) mempty)
+
+instance ToHStoreText TL.Text where
+  toHStoreText = HStoreText . TL.foldrChunks (escapeAppend . TS.encodeUtf8) mempty
+
+instance (ToHStoreText a) => ToHStoreText (Maybe a) where
+  toHStoreText Nothing = HStoreText $ byteString "NULL"
+  toHStoreText (Just x) = toHStoreText x
+
+escapeAppend :: BS.ByteString -> Builder -> Builder
+escapeAppend = loop
+  where
+    loop (BS.break quoteNeeded -> (a,b)) rest
+      = byteString a `mappend`
+          case BS.uncons b of
+            Nothing     ->  rest
+            Just (c,d)  ->  quoteChar c `mappend` loop d rest
+
+    quoteNeeded c = c == c2w '\"' || c == c2w '\\'
+    quoteChar c
+        | c == c2w '\"' = byteString "\\\""
+        | otherwise     = byteString "\\\\"
+
+hstore :: (ToHStoreText a, ToHStoreText b) => a -> b -> HStoreBuilder
+hstore (toHStoreText -> (HStoreText key)) (toHStoreText -> (HStoreText val)) =
+    Comma (char8 '"' `mappend` key `mappend` byteString "\"=>\""
+              `mappend` val `mappend` char8 '"')
+
+instance ToField HStoreBuilder where
+    toField  Empty    = toField (BS.empty)
+    toField (Comma x) = toField (BU.toLazyByteString x)
+
+newtype HStoreList = HStoreList {fromHStoreList :: [(Text, Maybe Text)]} deriving (Typeable, Show)
+
+-- | hstore
+instance ToHStore HStoreList where
+    toHStore (HStoreList xs) = mconcat (map (uncurry hstore) xs)
+
+instance ToField HStoreList where
+    toField xs = toField (toHStore xs)
+
+-- | hstore
+instance FromField HStoreList where
+    fromField f mdat = do
+      typ <- typename f
+      if typ /= "hstore"
+        then returnError Incompatible f ""
+        else case mdat of
+               Nothing  -> returnError UnexpectedNull f ""
+               Just dat ->
+                   case P.parseOnly (parseHStore <* P.endOfInput) dat of
+                     Left err ->
+                         returnError ConversionFailed f err
+                     Right (Left err) ->
+                         returnError ConversionFailed f "unicode exception" <|>
+                           conversionError err
+                     Right (Right val) ->
+                         return val
+
+newtype HStoreMap  = HStoreMap {fromHStoreMap :: Map Text (Maybe Text)} deriving (Eq, Ord, Typeable, Show)
+
+instance ToHStore HStoreMap where
+    toHStore (HStoreMap xs) = Map.foldrWithKey f mempty xs
+      where f k v xs = hstore k v `mappend` xs
+
+instance ToField HStoreMap where
+    toField xs = toField (toHStore xs)
+
+instance FromField HStoreMap where
+    fromField f mdat = convert <$> fromField f mdat
+      where convert (HStoreList xs) = HStoreMap (Map.fromList xs)
+
+parseHStoreList :: BS.ByteString -> Either String HStoreList
+parseHStoreList dat =
+    case P.parseOnly (parseHStore <* P.endOfInput) dat of
+      Left err          -> Left (show err)
+      Right (Left err)  -> Left (show err)
+      Right (Right val) -> Right val
+
+parseHStore :: P.Parser (Either UnicodeException HStoreList)
+parseHStore = do
+    kvs <- P.sepBy' (skipWhiteSpace *> parseHStoreKeyVal)
+                    (skipWhiteSpace *> P.word8 (c2w ','))
+    return $ HStoreList <$> sequence kvs
+
+parseHStoreKeyVal :: P.Parser (Either UnicodeException (Text, Maybe Text))
+parseHStoreKeyVal = do
+  mkey <- parseHStoreText
+  case mkey of
+    Left err -> return (Left err)
+    Right key -> do
+      skipWhiteSpace
+      _ <- P.string "=>"
+      skipWhiteSpace
+      mval <- parseHStoreText
+      case mval of
+        Left  err -> return (Left err)
+        Right "NULL" -> return (Right (key, Nothing))
+        Right val -> return (Right (key, Just val))
+
+
+skipWhiteSpace :: P.Parser ()
+skipWhiteSpace = P.skipWhile P.isSpace_w8
+
+parseHStoreText :: P.Parser (Either UnicodeException Text)
+parseHStoreText = do
+  _ <- P.word8 (c2w '"')
+  mtexts <- parseHStoreTexts id
+  case mtexts of
+    Left  err   -> return (Left err)
+    Right texts -> do
+                     _ <- P.word8 (c2w '"')
+                     return (Right (TS.concat texts))
+
+parseHStoreTexts :: ([Text] -> [Text])
+                 -> P.Parser (Either UnicodeException [Text])
+parseHStoreTexts acc = do
+  mchunk <- TS.decodeUtf8' <$> P.takeWhile (not . isSpecialChar)
+  case mchunk of
+    Left err    -> return (Left err)
+    Right chunk ->
+        (do
+          _ <- P.word8 (c2w '\\')
+          c <- TS.singleton . w2c <$> P.satisfy isSpecialChar
+          parseHStoreTexts (acc . (chunk:) . (c:))
+        ) <|> return (Right (acc [chunk]))
+ where
+   isSpecialChar c = c == c2w '\\' || c == c2w '"'

src/Database/PostgreSQL/Simple/HStoreV2/Internal.hs

@@ -0,0 +1,19 @@
+------------------------------------------------------------------------------
+-- |
+-- Module:      Database.PostgreSQL.Simple.HStore.Internal
+-- Copyright:   (c) 2013 Leon P Smith
+-- License:     BSD3
+-- Maintainer:  Leon P Smith <[email protected]>
+-- Stability:   experimental
+--
+------------------------------------------------------------------------------
+
+module Database.PostgreSQL.Simple.HStoreV2.Internal
+     ( HStoreBuilder(..)
+     , HStoreText(..)
+     , parseHStore
+     , parseHStoreKeyVal
+     , parseHStoreText
+     ) where
+
+import Database.PostgreSQL.Simple.HStoreV2.Implementation

test/Main.hs

@@ -7,6 +7,7 @@ import Common
 import Database.PostgreSQL.Simple.FromField (FromField)
 import Database.PostgreSQL.Simple.Types(Query(..),Values(..))
 import Database.PostgreSQL.Simple.HStore
+import qualified Database.PostgreSQL.Simple.HStoreV2 as V2
 import Database.PostgreSQL.Simple.Copy
 import qualified Database.PostgreSQL.Simple.Transaction as ST
 
@@ -61,6 +62,7 @@ tests env = testGroup "tests"
     , testCase "2-ary generic"      . testGeneric2
     , testCase "3-ary generic"      . testGeneric3
     , testCase "Timeout"            . testTimeout
+    , testCase "HStoreV2"           . testHStoreV2
     ]
 
 testBytea :: TestEnv -> TestTree
@@ -203,6 +205,22 @@ testHStore TestEnv{..} = do
       m' <- query conn "SELECT ?::hstore" m
       [m] @?= m'
 
+-- Should be able to handle NULL values assigned to hstore keys.
+testHStoreV2 :: TestEnv -> Assertion
+testHStoreV2 TestEnv{..} = do
+    execute_ conn "CREATE EXTENSION IF NOT EXISTS hstore"
+    roundTrip []
+    roundTrip [("foo", Just "bar"),("bar", Just "baz"),("baz", Just "hello")]
+    roundTrip [("fo\"o", Just "bar"),("b\\ar", Just "baz"),("baz", Just "\"value\\with\"escapes")]
+    roundTrip [("fo\"o", Just "bar"),("b\\ar", Just "baz"),("baz_with_null", Nothing)]
+  where
+    roundTrip :: [(Text, Maybe Text)] -> Assertion
+    roundTrip xs = do
+      let m = Only (V2.HStoreMap (Map.fromList xs))
+      m' <- query conn "SELECT ?::hstore" m
+      [m] @?= m'
+
+
 testJSON :: TestEnv -> Assertion
 testJSON TestEnv{..} = do
     roundTrip (Map.fromList [] :: Map Text Text)
@@ -469,7 +487,7 @@ isSyntaxError SqlError{..} = sqlState == "42601"
 -- Note that some tests, such as Notify, use multiple connections, and assume
 -- that 'testConnect' connects to the same database every time it is called.
 testConnect :: IO Connection
-testConnect = connectPostgreSQL ""
+testConnect = connectPostgreSQL "host='localhost' port=5432 user='pgsimple' password='pgsimple'"
 
 withTestEnv :: (TestEnv -> IO a) -> IO a
 withTestEnv cb =