    {-# LANGUAGE RecordWildCards #-}

    ------------------------------------------------------------------------------
    -- |
    -- Module:      Database.PostgreSQL.Simple.TypeInfo
    -- Copyright:   (c) 2013 Leon P Smith
    -- License:     BSD3
    -- Maintainer:  Leon P Smith <leon@melding-monads.com>
    -- Stability:   experimental
    --
    -- This module provides convenient and efficient access to parts of the
    -- @pg_type@ metatable.  At the moment, this requires PostgreSQL 8.4 if
    -- you need to work with types that do not appear in
    -- 'Database.PostgreSQL.Simple.TypeInfo.Static'.
    --
    -- The current scheme could be more efficient, especially for some use
    -- cases.  In particular,  connection pools that use many user-added
    -- types and connect to a set of servers with identical (or at least
    -- compatible) @pg_type@ and associated tables could share a common
    -- typeinfo cache,  thus saving memory and communication between the
    -- client and server.
    --
    ------------------------------------------------------------------------------

    module Database.PostgreSQL.Simple.TypeInfo
         ( getTypeInfo
         , TypeInfo(..)
         , Attribute(..)
         ) 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)

    import qualified Database.PostgreSQL.LibPQ as PQ
    import {-# SOURCE #-} Database.PostgreSQL.Simple
    import                Database.PostgreSQL.Simple.Internal
    import                Database.PostgreSQL.Simple.Types
    import                Database.PostgreSQL.Simple.TypeInfo.Types
    import                Database.PostgreSQL.Simple.TypeInfo.Static

    -- | Returns the metadata of the type with a particular oid.  To find
    --   this data, 'getTypeInfo' first consults postgresql-simple's
    --   built-in 'staticTypeInfo' table,  then checks  the connection's
    --   typeinfo cache.   Finally,  the database's 'pg_type' table will
    --   be queried only if necessary,  and the result will be stored
    --   in the connections's cache.

    getTypeInfo :: Connection -> PQ.Oid -> IO TypeInfo
    getTypeInfo :: Connection -> Oid -> IO TypeInfo
    getTypeInfo conn :: Connection
    conn@Connection{MVar TypeInfoCache
    MVar Connection
    IORef Int64
    connectionHandle :: MVar Connection
    connectionObjects :: MVar TypeInfoCache
    connectionTempNameCounter :: IORef Int64
    connectionTempNameCounter :: Connection -> IORef Int64
    connectionObjects :: Connection -> MVar TypeInfoCache
    connectionHandle :: Connection -> MVar Connection
    ..} Oid
    oid' =
      case Oid -> Maybe TypeInfo
    staticTypeInfo Oid
    oid' of
        Just TypeInfo
    name' -> TypeInfo -> IO TypeInfo
    forall a. a -> IO a
    forall (m :: * -> *) a. Monad m => a -> m a
    return TypeInfo
    name'
        Maybe TypeInfo
    Nothing -> MVar TypeInfoCache
    -> (TypeInfoCache -> IO (TypeInfoCache, TypeInfo)) -> IO TypeInfo
    forall a b. MVar a -> (a -> IO (a, b)) -> IO b
    modifyMVar MVar TypeInfoCache
    connectionObjects ((TypeInfoCache -> IO (TypeInfoCache, TypeInfo)) -> IO TypeInfo)
    -> (TypeInfoCache -> IO (TypeInfoCache, TypeInfo)) -> IO TypeInfo
    forall a b. (a -> b) -> a -> b
    $ Connection -> Oid -> TypeInfoCache -> IO (TypeInfoCache, TypeInfo)
    getTypeInfo' Connection
    conn Oid
    oid'

    getTypeInfo' :: Connection -> PQ.Oid -> TypeInfoCache
                 -> IO (TypeInfoCache, TypeInfo)
    getTypeInfo' :: Connection -> Oid -> TypeInfoCache -> IO (TypeInfoCache, TypeInfo)
    getTypeInfo' Connection
    conn Oid
    oid' TypeInfoCache
    oidmap =
      case Int -> TypeInfoCache -> Maybe TypeInfo
    forall a. Int -> IntMap a -> Maybe a
    IntMap.lookup (Oid -> Int
    oid2int Oid
    oid') TypeInfoCache
    oidmap of
        Just TypeInfo
    typeinfo -> (TypeInfoCache, TypeInfo) -> IO (TypeInfoCache, TypeInfo)
    forall a. a -> IO a
    forall (m :: * -> *) a. Monad m => a -> m a
    return (TypeInfoCache
    oidmap, TypeInfo
    typeinfo)
        Maybe TypeInfo
    Nothing -> do
          names  <- Connection
    -> Query
    -> Only Oid
    -> IO [(Oid, Char, Char, ByteString, Oid, Oid)]
    forall q r.
    (ToRow q, FromRow r) =>
    Connection -> Query -> q -> IO [r]
    query Connection
    conn Query
    "SELECT oid, typcategory, typdelim, typname,\
                             \ typelem, typrelid\
                             \ FROM pg_type WHERE oid = ?"
                               (Oid -> Only Oid
    forall a. a -> Only a
    Only Oid
    oid')
          (oidmap', typeInfo) <-
              case names of
                []  -> (TypeInfoCache, TypeInfo) -> IO (TypeInfoCache, TypeInfo)
    forall a. a -> IO a
    forall (m :: * -> *) a. Monad m => a -> m a
    return ((TypeInfoCache, TypeInfo) -> IO (TypeInfoCache, TypeInfo))
    -> (TypeInfoCache, TypeInfo) -> IO (TypeInfoCache, TypeInfo)
    forall a b. (a -> b) -> a -> b
    $ SqlError -> (TypeInfoCache, TypeInfo)
    forall a e. (HasCallStack, Exception e) => e -> a
    throw (ByteString -> SqlError
    fatalError ByteString
    "invalid type oid")
                [(Oid
    typoid, Char
    typcategory, Char
    typdelim, ByteString
    typname, Oid
    typelem_, Oid
    typrelid)] -> do
                   case Char
    typcategory of
                     Char
    'A' -> do
                       (oidmap', typelem) <- Connection -> Oid -> TypeInfoCache -> IO (TypeInfoCache, TypeInfo)
    getTypeInfo' Connection
    conn Oid
    typelem_ TypeInfoCache
    oidmap
                       let !typeInfo = Array{Char
    ByteString
    Oid
    TypeInfo
    typoid :: Oid
    typcategory :: Char
    typdelim :: Char
    typname :: ByteString
    typelem :: TypeInfo
    typname :: ByteString
    typdelim :: Char
    typcategory :: Char
    typoid :: Oid
    typelem :: TypeInfo
    ..}
                       return $! (oidmap', typeInfo)
                     Char
    'R' -> do
                       rngsubtypeOids <- Connection -> Query -> Only Oid -> IO [Only Oid]
    forall q r.
    (ToRow q, FromRow r) =>
    Connection -> Query -> q -> IO [r]
    query Connection
    conn Query
    "SELECT rngsubtype\
                                                   \ FROM pg_range\
                                                   \ WHERE rngtypid = ?"
                                                    (Oid -> Only Oid
    forall a. a -> Only a
    Only Oid
    oid')
                       case rngsubtypeOids of
                         [Only Oid
    rngsubtype_] -> do
                            (oidmap', rngsubtype) <-
                                Connection -> Oid -> TypeInfoCache -> IO (TypeInfoCache, TypeInfo)
    getTypeInfo' Connection
    conn Oid
    rngsubtype_ TypeInfoCache
    oidmap
                            let !typeInfo = Range{Char
    ByteString
    Oid
    TypeInfo
    typoid :: Oid
    typcategory :: Char
    typdelim :: Char
    typname :: ByteString
    typname :: ByteString
    typdelim :: Char
    typcategory :: Char
    typoid :: Oid
    rngsubtype :: TypeInfo
    rngsubtype :: TypeInfo
    ..}
                            return $! (oidmap', typeInfo)
                         [Only Oid]
    _ -> String -> IO (TypeInfoCache, TypeInfo)
    forall a. String -> IO a
    forall (m :: * -> *) a. MonadFail m => String -> m a
    fail String
    "range subtype query failed to return exactly one result"
                     Char
    'C' -> do
                       cols <- Connection -> Query -> Only Oid -> IO [(ByteString, Oid)]
    forall q r.
    (ToRow q, FromRow r) =>
    Connection -> Query -> q -> IO [r]
    query Connection
    conn Query
    "SELECT attname, atttypid\
                                         \ FROM pg_attribute\
                                         \ WHERE attrelid = ?\
                                           \ AND attnum > 0\
                                           \ AND NOT attisdropped\
                                         \ ORDER BY attnum"
                                          (Oid -> Only Oid
    forall a. a -> Only a
    Only Oid
    typrelid)
                       vec <- MV.new $! length cols
                       (oidmap', attributes) <- getAttInfos conn cols oidmap vec 0
                       let !typeInfo = Composite{Char
    Vector Attribute
    ByteString
    Oid
    typoid :: Oid
    typcategory :: Char
    typdelim :: Char
    typname :: ByteString
    typrelid :: Oid
    typname :: ByteString
    typdelim :: Char
    typcategory :: Char
    typoid :: Oid
    attributes :: Vector Attribute
    attributes :: Vector Attribute
    typrelid :: Oid
    ..}
                       return $! (oidmap', typeInfo)
                     Char
    _ -> do
                       let !typeInfo :: TypeInfo
    typeInfo = Basic{Char
    ByteString
    Oid
    typoid :: Oid
    typcategory :: Char
    typdelim :: Char
    typname :: ByteString
    typname :: ByteString
    typdelim :: Char
    typcategory :: Char
    typoid :: Oid
    ..}
                       (TypeInfoCache, TypeInfo) -> IO (TypeInfoCache, TypeInfo)
    forall a. a -> IO a
    forall (m :: * -> *) a. Monad m => a -> m a
    return ((TypeInfoCache, TypeInfo) -> IO (TypeInfoCache, TypeInfo))
    -> (TypeInfoCache, TypeInfo) -> IO (TypeInfoCache, TypeInfo)
    forall a b. (a -> b) -> a -> b
    $! (TypeInfoCache
    oidmap, TypeInfo
    typeInfo)
                [(Oid, Char, Char, ByteString, Oid, Oid)]
    _ -> String -> IO (TypeInfoCache, TypeInfo)
    forall a. String -> IO a
    forall (m :: * -> *) a. MonadFail m => String -> m a
    fail String
    "typename query returned more than one result"
                       -- oid is a primary key,  so the query should
                       -- never return more than one result
          let !oidmap'' = Int -> TypeInfo -> TypeInfoCache -> TypeInfoCache
    forall a. Int -> a -> IntMap a -> IntMap a
    IntMap.insert (Oid -> Int
    oid2int Oid
    oid') TypeInfo
    typeInfo TypeInfoCache
    oidmap'
          return $! (oidmap'', typeInfo)

    getAttInfos :: Connection -> [(B.ByteString, PQ.Oid)] -> TypeInfoCache
                -> MV.IOVector Attribute -> Int
                -> IO (TypeInfoCache, V.Vector Attribute)
    getAttInfos :: Connection
    -> [(ByteString, Oid)]
    -> TypeInfoCache
    -> IOVector Attribute
    -> Int
    -> IO (TypeInfoCache, Vector Attribute)
    getAttInfos Connection
    conn [(ByteString, Oid)]
    cols TypeInfoCache
    oidmap IOVector Attribute
    vec Int
    n =
        case [(ByteString, Oid)]
    cols of
          [] -> do
            !attributes <- MVector (PrimState IO) Attribute -> IO (Vector Attribute)
    forall (m :: * -> *) a.
    PrimMonad m =>
    MVector (PrimState m) a -> m (Vector a)
    V.unsafeFreeze IOVector Attribute
    MVector (PrimState IO) Attribute
    vec
            return $! (oidmap, attributes)
          ((ByteString
    attname, Oid
    attTypeOid):[(ByteString, Oid)]
    xs) -> do
            (oidmap', atttype) <- Connection -> Oid -> TypeInfoCache -> IO (TypeInfoCache, TypeInfo)
    getTypeInfo' Connection
    conn Oid
    attTypeOid TypeInfoCache
    oidmap
            MV.write vec n $! Attribute{..}
            getAttInfos conn xs oidmap' vec (n+1)