@@ -0,0 +1,187 @@
1+{-# LANGUAGE ViewPatterns, DeriveDataTypeable, GeneralizedNewtypeDeriving #-}
2+
3+------------------------------------------------------------------------------
4+-- |
5+-- Module: Database.PostgreSQL.Simple.HStore.Implementation
6+-- Copyright: (c) 2013 Leon P Smith
7+-- License: BSD3
8+-- Maintainer: Leon P Smith <leon@melding-monads.com>
9+-- Stability: experimental
10+-- Portability: portable
11+--
12+-- This code has yet to be profiled and optimized.
13+--
14+------------------------------------------------------------------------------
15+
16+module Database.PostgreSQL.Simple.HStore.Implementation where
17+
18+import Control.Applicative
19+import Blaze.ByteString.Builder
20+ ( Builder, toLazyByteString, copyByteString )
21+import Blaze.ByteString.Builder.Char8 (fromChar)
22+import qualified Data.Attoparsec.ByteString as P
23+import qualified Data.Attoparsec.ByteString.Char8 as P (isSpace_w8)
24+import qualified Data.ByteString as BS
25+import Data.ByteString.Internal (c2w, w2c)
26+import qualified Data.ByteString.Lazy as BL
27+import qualified Data.ByteString.Lazy.Internal as BL
28+import Data.Map(Map)
29+import qualified Data.Map as Map
30+import Data.Text(Text)
31+import qualified Data.Text as TS
32+import qualified Data.Text.Encoding as TS
33+import Data.Text.Encoding.Error(UnicodeException)
34+import qualified Data.Text.Lazy as TL
35+import Data.Typeable
36+import Data.Monoid(Monoid(..))
37+import Database.PostgreSQL.Simple.FromField
38+import Database.PostgreSQL.Simple.ToField
39+
40+data HStoreBuilder
41+ = Empty
42+ | Comma !Builder
43+ deriving (Typeable)
44+
45+instance Monoid HStoreBuilder where
46+ mempty = Empty
47+ mappend Empty x = x
48+ mappend (Comma a) x
49+ = Comma (a `mappend` case x of
50+ Empty -> mempty
51+ Comma b -> fromChar ',' `mappend` b)
52+
53+class ToHStoreText a where
54+ toHStoreText :: a -> HStoreText
55+
56+-- | Represents escape text, ready to be the key or value to a hstore value
57+newtype HStoreText = HStoreText Builder deriving (Typeable, Monoid)
58+
59+instance ToHStoreText HStoreText where
60+ toHStoreText = id
61+
62+-- | Assumed to be UTF-8 encoded
63+instance ToHStoreText BS.ByteString where
64+ toHStoreText str = HStoreText (escapeAppend str mempty)
65+
66+-- | Assumed to be UTF-8 encoded
67+instance ToHStoreText BL.ByteString where
68+ toHStoreText = HStoreText . BL.foldrChunks escapeAppend mempty
69+
70+instance ToHStoreText TS.Text where
71+ toHStoreText str = HStoreText (escapeAppend (TS.encodeUtf8 str) mempty)
72+
73+instance ToHStoreText TL.Text where
74+ toHStoreText = HStoreText . TL.foldrChunks (escapeAppend . TS.encodeUtf8) mempty
75+
76+escapeAppend :: BS.ByteString -> Builder -> Builder
77+escapeAppend = loop
78+ where
79+ loop (BS.break quoteNeeded -> (a,b)) rest
80+ = copyByteString a `mappend`
81+ case BS.uncons b of
82+ Nothing -> rest
83+ Just (c,d) -> quoteChar c `mappend` loop d rest
84+
85+ quoteNeeded c = c == c2w '\"' || c == c2w '\\'
86+ quoteChar c
87+ | c == c2w '\"' = copyByteString "\\\""
88+ | otherwise = copyByteString "\\\\"
89+
90+hstore :: (ToHStoreText a, ToHStoreText b) => a -> b -> HStoreBuilder
91+hstore (toHStoreText -> (HStoreText key)) (toHStoreText -> (HStoreText val)) =
92+ Comma (fromChar '"' `mappend` key `mappend` copyByteString "\"=>\""
93+ `mappend` val `mappend` fromChar '"')
94+
95+instance ToField HStoreBuilder where
96+ toField Empty = toField (BS.empty)
97+ toField (Comma x) = toField (toLazyByteString x)
98+
99+newtype HStoreList = HStoreList [(Text,Text)] deriving (Typeable, Show)
100+
101+instance ToField HStoreList where
102+ toField (HStoreList xs) = toField (mconcat (map (uncurry hstore) xs))
103+
104+instance FromField HStoreList where
105+ fromField f mdat = do
106+ typ <- typename f
107+ if typ /= "hstore"
108+ then returnError Incompatible f ""
109+ else case mdat of
110+ Nothing -> returnError UnexpectedNull f ""
111+ Just dat ->
112+ case P.parseOnly parseHStore dat of
113+ Left err ->
114+ returnError ConversionFailed f err
115+ Right (Left err) ->
116+ returnError ConversionFailed f "unicode exception" <|>
117+ conversionError err
118+ Right (Right val) ->
119+ return val
120+
121+newtype HStoreMap = HStoreMap (Map Text Text) deriving (Eq, Ord, Typeable, Show)
122+
123+instance ToField HStoreMap where
124+ toField (HStoreMap xs) = toField (Map.foldWithKey f mempty xs)
125+ where f k v xs = hstore k v `mappend` xs
126+
127+instance FromField HStoreMap where
128+ fromField f mdat = convert <$> fromField f mdat
129+ where convert (HStoreList xs) = HStoreMap (Map.fromList xs)
130+
131+parseHStore :: P.Parser (Either UnicodeException HStoreList)
132+parseHStore = skipWhiteSpace >> loop id
133+ where
134+ loop acc = do
135+ mkv <- parseHStoreKeyVal
136+ case mkv of
137+ Left err -> return (Left err)
138+ Right kv -> do
139+ skipWhiteSpace
140+ (do
141+ _ <- P.word8 (c2w ',')
142+ skipWhiteSpace
143+ loop (acc . (kv:))
144+ ) <|> return (Right (HStoreList (acc [kv])))
145+
146+parseHStoreKeyVal :: P.Parser (Either UnicodeException (Text,Text))
147+parseHStoreKeyVal = do
148+ mkey <- parseHStoreText
149+ case mkey of
150+ Left err -> return (Left err)
151+ Right key -> do
152+ skipWhiteSpace
153+ _ <- P.string "=>"
154+ skipWhiteSpace
155+ mval <- parseHStoreText
156+ case mval of
157+ Left err -> return (Left err)
158+ Right val -> return (Right (key,val))
159+
160+
161+skipWhiteSpace :: P.Parser ()
162+skipWhiteSpace = P.skipWhile P.isSpace_w8
163+
164+parseHStoreText :: P.Parser (Either UnicodeException Text)
165+parseHStoreText = do
166+ _ <- P.word8 (c2w '"')
167+ mtexts <- parseHStoreTexts id
168+ case mtexts of
169+ Left err -> return (Left err)
170+ Right texts -> do
171+ _ <- P.word8 (c2w '"')
172+ return (Right (TS.concat texts))
173+
174+parseHStoreTexts :: ([Text] -> [Text])
175+ -> P.Parser (Either UnicodeException [Text])
176+parseHStoreTexts acc = do
177+ mchunk <- TS.decodeUtf8' <$> P.takeWhile (not . isSpecialChar)
178+ case mchunk of
179+ Left err -> return (Left err)
180+ Right chunk ->
181+ (do
182+ _ <- P.word8 (c2w '\\')
183+ c <- TS.singleton . w2c <$> P.satisfy isSpecialChar
184+ parseHStoreTexts (acc . (chunk:) . (c:))
185+ ) <|> return (Right (acc [chunk]))
186+ where
187+ isSpecialChar c = c == c2w '\\' || c == c2w '"'