Skip to content
Browse files

Add preliminary support for converting SQL arrays

  • Loading branch information...
1 parent b58cbb1 commit eb04ca39c5c22e3f4d083ba4986ab9e8339ed7d1 @basvandijk committed Jul 30, 2012
View
56 src/Database/PostgreSQL/Simple.hs
@@ -4,6 +4,7 @@
{-# LANGUAGE PatternGuards #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ViewPatterns #-}
+{-# LANGUAGE QuasiQuotes #-}
------------------------------------------------------------------------------
-- |
@@ -136,6 +137,7 @@ import Database.PostgreSQL.Simple.ToRow (ToRow(..))
import Database.PostgreSQL.Simple.Types
( Binary(..), In(..), Only(..), Query(..), (:.)(..) )
import Database.PostgreSQL.Simple.Internal as Base
+import Database.PostgreSQL.Simple.SqlQQ (sql)
import qualified Database.PostgreSQL.LibPQ as PQ
import qualified Data.ByteString.Char8 as B
import qualified Data.Text as T
@@ -531,19 +533,19 @@ finishQuery conn q result = do
PQ.TuplesOk -> do
ncols <- PQ.nfields result
let unCol (PQ.Col x) = fromIntegral x :: Int
- typenames <- V.generateM (unCol ncols)
+ typeinfos <- V.generateM (unCol ncols)
(\(PQ.Col . fromIntegral -> col) -> do
- getTypename conn =<< PQ.ftype result col)
+ getTypeInfo conn =<< PQ.ftype result col)
nrows <- PQ.ntuples result
ncols <- PQ.nfields result
forM' 0 (nrows-1) $ \row -> do
- let rw = Row row typenames result
+ let rw = Row row typeinfos result
case runStateT (runReaderT (unRP fromRow) rw) 0 of
Ok (val,col) | col == ncols -> return val
| otherwise -> do
vals <- forM' 0 (ncols-1) $ \c -> do
v <- PQ.getvalue result row c
- return ( typenames V.! unCol c
+ return ( typeinfos V.! unCol c
, fmap ellipsis v )
throw (ConversionFailed
(show (unCol ncols) ++ " values: " ++ show vals)
@@ -931,24 +933,36 @@ fmtError msg q xs = throw FormatError {
-- wrong results. In such cases, write a @newtype@ wrapper and a
-- custom 'Result' instance to handle your encoding.
-getTypename :: Connection -> PQ.Oid -> IO ByteString
-getTypename conn@Connection{..} oid =
+getTypeInfo :: Connection -> PQ.Oid -> IO TypeInfo
+getTypeInfo conn@Connection{..} oid =
case oid2typname oid of
- Just name -> return name
+ Just name -> return $! TypeInfo { typ = NamedOid oid name
+ , typelem = Nothing
+ }
Nothing -> modifyMVar connectionObjects $ \oidmap -> do
case IntMap.lookup (oid2int oid) oidmap of
- Just name -> return (oidmap, name)
+ Just typeinfo -> return (oidmap, typeinfo)
Nothing -> do
- names <- query conn "SELECT typname FROM pg_type WHERE oid=?"
- (Only oid)
- name <- case names of
- [] -> return $ throw SqlError {
- sqlNativeError = -1,
- sqlErrorMsg = "invalid type oid",
- sqlState = ""
- }
- [Only x] -> return x
- _ -> fail "typename query returned more than one result"
- -- oid is a primary key, so the query should
- -- never return more than one result
- return (IntMap.insert (oid2int oid) name oidmap, name)
+ names <- query conn
+ [sql| SELECT p.oid, p.typname, c.oid, c.typname
+ FROM pg_type AS p LEFT OUTER JOIN pg_type AS c
+ ON c.oid = p.typelem
+ WHERE p.oid = ?
+ |] (Only oid)
+ typinf <- case names of
+ [] -> return $ throw SqlError {
+ sqlNativeError = -1,
+ sqlErrorMsg = "invalid type oid",
+ sqlState = ""
+ }
+ [(pOid, pTypName, mbCOid, mbCTypName)] ->
+ return $! TypeInfo { typ = NamedOid pOid pTypName
+ , typelem = do
+ cOid <- mbCOid
+ cTypName <- mbCTypName
+ return $ NamedOid cOid cTypName
+ }
+ _ -> fail "typename query returned more than one result"
+ -- oid is a primary key, so the query should
+ -- never return more than one result
+ return (IntMap.insert (oid2int oid) typinf oidmap, typinf)
View
50 src/Database/PostgreSQL/Simple/FromField.hs
@@ -1,6 +1,7 @@
{-# LANGUAGE CPP, DeriveDataTypeable, DeriveFunctor #-}
{-# LANGUAGE FlexibleInstances, TypeSynonymInstances #-}
{-# LANGUAGE PatternGuards, ScopedTypeVariables #-}
+{-# LANGUAGE RecordWildCards, OverloadedStrings #-}
------------------------------------------------------------------------------
-- |
@@ -53,6 +54,7 @@ import Data.ByteString (ByteString)
import qualified Data.ByteString.Char8 as B
import Data.Int (Int16, Int32, Int64)
import Data.List (foldl')
+import Data.Maybe ( fromMaybe )
import Data.Ratio (Ratio)
import Data.Time ( UTCTime, ZonedTime, LocalTime, Day, TimeOfDay )
import Data.Typeable (Typeable, typeOf)
@@ -120,6 +122,52 @@ instance (FromField a) => FromField (Maybe a) where
fromField _ Nothing = pure Nothing
fromField f bs = Just <$> fromField f bs
+instance (Typeable a, FromField a) => FromField [a] where
+ fromField f Nothing = returnError UnexpectedNull f ""
+ fromField f (Just bs)
+ | nesting r == 0 = mapM (fromField f' . Just) $ reverse $ parts r
+ | otherwise = returnError ConversionFailed f ""
+ where
+ r = B8.foldl' go ArrayState { nesting = 0
+ , currentIx = 0
+ , currentPartBeginIx = 1
+ , parts = []
+ } bs
+
+ go s '{' = s {nesting = nesting s + 1, currentIx = currentIx s + 1}
+ go s '}' | nesting s == 1 = s { nesting = nesting s - 1
+ , parts = B8.take (currentIx s - currentPartBeginIx s)
+ (B8.drop (currentPartBeginIx s) bs)
+ : parts s
+ , currentIx = nextIx
+ , currentPartBeginIx = nextIx
+ }
+ | otherwise = s {nesting = nesting s - 1, currentIx = currentIx s + 1}
+ where
+ nextIx = currentIx s + 1
+ go s ',' | nesting s == 1 = s { parts = B8.take (currentIx s - currentPartBeginIx s)
+ (B8.drop (currentPartBeginIx s) bs)
+ : parts s
+ , currentIx = nextIx
+ , currentPartBeginIx = nextIx
+ }
+ where
+ nextIx = currentIx s + 1
+ go s _ = s {currentIx = currentIx s + 1}
+
+ TypeInfo{..} = typeinfo f
+
+ f' = f {typeinfo = TypeInfo { typ = fromMaybe typ typelem
+ , typelem = Nothing
+ }
+ }
+
+data ArrayState = ArrayState { nesting :: !Int
+ , currentIx :: !Int
+ , currentPartBeginIx :: !Int
+ , parts :: ![ByteString]
+ } deriving Show
+
instance FromField Null where
fromField _ Nothing = pure Null
fromField f (Just _) = returnError ConversionFailed f "data is not null"
@@ -270,7 +318,7 @@ doFromField :: forall a . (Typeable a)
=> Field -> Compat -> (ByteString -> Ok a)
-> Maybe ByteString -> Ok a
doFromField f types cvt (Just bs)
- | Just typ <- oid2builtin (typeOid f)
+ | Just typ <- oid2builtin (typoid $ typ $ typeinfo f)
, mkCompat typ `compat` types = cvt bs
| otherwise = returnError Incompatible f "types incompatible"
doFromField f _ _ _ = returnError UnexpectedNull f ""
View
6 src/Database/PostgreSQL/Simple/FromRow.hs
@@ -68,13 +68,13 @@ class FromRow a where
fieldWith :: FieldParser a -> RowParser a
fieldWith fieldP = RP $ do
let unCol (PQ.Col x) = fromIntegral x :: Int
- Row{..} <- ask
+ r@Row{..} <- ask
column <- lift get
lift (put (column + 1))
let ncols = nfields rowresult
if (column >= ncols)
then do
- let vals = map (\c -> ( typenames ! (unCol c)
+ let vals = map (\c -> ( typenames r ! (unCol c)
, fmap ellipsis (getvalue rowresult row c) ))
[0..ncols-1]
convertError = ConversionFailed
@@ -85,7 +85,7 @@ fieldWith fieldP = RP $ do
\convert and number in target type"
lift (lift (Errors [SomeException convertError]))
else do
- let typename = typenames ! unCol column
+ let typeinfo = typeinfos ! unCol column
result = rowresult
field = Field{..}
lift (lift (fieldP field (getvalue result row column)))
View
20 src/Database/PostgreSQL/Simple/Internal.hs
@@ -52,9 +52,20 @@ import System.IO.Unsafe (unsafePerformIO)
data Field = Field {
result :: !PQ.Result
, column :: {-# UNPACK #-} !PQ.Column
- , typename :: !ByteString
+ , typeinfo :: !TypeInfo
}
+data NamedOid = NamedOid { typoid :: !PQ.Oid
+ , typname :: !ByteString
+ } deriving Show
+
+data TypeInfo = TypeInfo { typ :: !NamedOid
+ , typelem :: !(Maybe NamedOid)
+ } deriving Show
+
+typename :: Field -> ByteString
+typename = typname . typ . typeinfo
+
name :: Field -> Maybe ByteString
name Field{..} = unsafePerformIO (PQ.fname result column)
@@ -76,7 +87,7 @@ typeOid Field{..} = unsafePerformIO (PQ.ftype result column)
data Connection = Connection {
connectionHandle :: {-# UNPACK #-} !(MVar PQ.Connection)
- , connectionObjects :: {-# UNPACK #-} !(MVar (IntMap.IntMap ByteString))
+ , connectionObjects :: {-# UNPACK #-} !(MVar (IntMap.IntMap TypeInfo))
}
data SqlType
@@ -301,10 +312,13 @@ newNullConnection = do
data Row = Row {
row :: {-# UNPACK #-} !PQ.Row
- , typenames :: !(V.Vector ByteString)
+ , typeinfos :: !(V.Vector TypeInfo)
, rowresult :: !PQ.Result
}
+typenames :: Row -> V.Vector ByteString
+typenames = V.map (typname . typ) . typeinfos
+
newtype RowParser a = RP { unRP :: ReaderT Row (StateT PQ.Column Ok) a }
deriving ( Functor, Applicative, Alternative, Monad )
View
6 src/Database/PostgreSQL/Simple/ToField.hs
@@ -80,6 +80,12 @@ instance (ToField a) => ToField (Maybe a) where
toField (Just a) = toField a
{-# INLINE toField #-}
+instance (ToField a) => ToField [a] where
+ toField xs = Many $
+ Plain (fromByteString "ARRAY[") :
+ (intersperse (Plain (fromChar ',')) . map toField $ xs) ++
+ [Plain (fromChar ']')]
+
instance (ToField a) => ToField (In [a]) where
toField (In []) = Plain $ fromByteString "(null)"
toField (In xs) = Many $

1 comment on commit eb04ca3

@lpsmith

Yes, all in all this is fairly close to what I had in mind. There are a couple of issues I'll comment on though.

Please sign in to comment.
Something went wrong with that request. Please try again.