Extend TypeInfo to cover Composite types

lpsmith · Mar 27, 2013 · f4e2ff6 · f4e2ff6
1 parent 4e1b449
commit f4e2ff6
Show file tree

Hide file tree

Showing 3 changed files with 58 additions and 7 deletions.
diff --git a/src/Database/PostgreSQL/Simple/FromRow.hs-boot b/src/Database/PostgreSQL/Simple/FromRow.hs-boot
@@ -6,6 +6,13 @@ import                Database.PostgreSQL.Simple.Types
 class FromRow a
 
 instance (FromField a) => FromRow (Only a)
-instance (FromField a, FromField b, FromField c, FromField d) => FromRow (a,b,c,d)
-instance (FromField a, FromField b, FromField c, FromField d, FromField e) => FromRow (a,b,c,d,e)
+instance (FromField a, FromField b)
+      => FromRow (a,b)
+instance (FromField a, FromField b, FromField c, FromField d) 
+      => FromRow (a,b,c,d)
+instance (FromField a, FromField b, FromField c, FromField d, FromField e) 
+      => FromRow (a,b,c,d,e)
+instance (FromField a, FromField b, FromField c, FromField d, FromField e
+         ,FromField f) 
+      => FromRow (a,b,c,d,e,f)
 
diff --git a/src/Database/PostgreSQL/Simple/TypeInfo.hs b/src/Database/PostgreSQL/Simple/TypeInfo.hs
@@ -27,7 +27,10 @@ module Database.PostgreSQL.Simple.TypeInfo
      , TypeInfo(..)
      ) where
 
+import qualified Data.ByteString as B
 import qualified Data.IntMap as IntMap
+import qualified Data.Vector as V
+import qualified Data.Vector.Mutable as MV
 import           Control.Concurrent.MVar
 import           Control.Exception (throw)
 
@@ -57,18 +60,19 @@ getTypeInfo' conn oid oidmap =
   case IntMap.lookup (oid2int oid) oidmap of
     Just typeinfo -> return (oidmap, typeinfo)
     Nothing -> do
-      names  <- query conn "SELECT oid, typcategory, typdelim, typname, typelem\
+      names  <- query conn "SELECT oid, typcategory, typdelim, typname,\
+                         \ typelem, typrelid\
                          \ FROM pg_type WHERE oid = ?"
                            (Only oid)
       (oidmap', typeInfo) <-
           case names of
             []  -> return $ throw (fatalError "invalid type oid")
-            [(typoid, typcategory, typdelim, typname, typelem_)] -> do
+            [(typoid, typcategory, typdelim, typname, typelem_, typrelid)] -> do
                case typcategory of
                  'A' -> do
                    (oidmap', typelem) <- getTypeInfo' conn typelem_ oidmap
                    let !typeInfo = Array{..}
-                   return (oidmap', typeInfo)
+                   return $! (oidmap', typeInfo)
                  'R' -> do
                    rngsubtypeOids <- query conn "SELECT rngsubtype\
                                                \ FROM pg_range\
@@ -79,13 +83,38 @@ getTypeInfo' conn oid oidmap =
                         (oidmap', rngsubtype) <-
                             getTypeInfo' conn rngsubtype_ oidmap
                         let !typeInfo = Range{..}
-                        return (oidmap', typeInfo)
+                        return $! (oidmap', typeInfo)
                      _ -> fail "range subtype query failed to return exactly one result"
+                 'C' -> do
+                   cols <- query conn "SELECT attname, atttypid\
+                                     \ FROM pg_attribute\
+                                     \ WHERE attrelid = ?\
+                                       \ AND attnum > 0\
+                                       \ AND NOT attisdropped\
+                                     \ ORDER BY attnum"
+                                      (Only typrelid)
+                   vec <- MV.new $! length cols
+                   (oidmap', attributes) <- getAttInfos conn cols oidmap vec 0
+                   let !typeInfo = Composite{..}
+                   return $! (oidmap', typeInfo)
                  _ -> do
                    let !typeInfo = Basic{..}
-                   return (oidmap, typeInfo)
+                   return $! (oidmap, typeInfo)
             _ -> fail "typename query returned more than one result"
                    -- oid is a primary key,  so the query should
                    -- never return more than one result
       let !oidmap'' = IntMap.insert (oid2int oid) typeInfo oidmap'
       return $! (oidmap'', typeInfo)
+
+getAttInfos :: Connection -> [(B.ByteString, PQ.Oid)] -> TypeInfoCache
+            -> MV.IOVector Attribute -> Int
+            -> IO (TypeInfoCache, V.Vector Attribute)
+getAttInfos conn cols oidmap vec n =
+    case cols of
+      [] -> do
+        !attributes <- V.unsafeFreeze vec
+        return $! (oidmap, attributes)
+      ((attname, attTypeOid):xs) -> do
+        (oidmap', atttype) <- getTypeInfo' conn attTypeOid oidmap
+        MV.write vec n $! Attribute{..}
+        getAttInfos conn xs oidmap' vec (n+1)
diff --git a/src/Database/PostgreSQL/Simple/TypeInfo/Types.hs b/src/Database/PostgreSQL/Simple/TypeInfo/Types.hs
@@ -12,6 +12,7 @@ module Database.PostgreSQL.Simple.TypeInfo.Types where
 
 import Data.ByteString(ByteString)
 import Database.PostgreSQL.LibPQ(Oid)
+import Data.Vector(Vector)
 
 -- | A structure representing some of the metadata regarding a PostgreSQL
 --   type,  mostly taken from the @pg_type@ table.
@@ -38,4 +39,18 @@ data TypeInfo
           , rngsubtype  :: !TypeInfo
           }
 
+  | Composite { typoid      :: {-# UNPACK #-} !Oid
+              , typcategory :: {-# UNPACK #-} !Char
+              , typdelim    :: {-# UNPACK #-} !Char
+              , typname     :: !ByteString
+              , typrelid    :: {-# UNPACK #-} !Oid
+              , attributes  :: !(Vector Attribute)
+              }
+
     deriving (Show)
+
+data Attribute
+   = Attribute { attname :: !ByteString
+               , atttype :: !TypeInfo
+               }
+     deriving (Show)