{-# 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{..} oid' :: Oid
oid' =
  case Oid -> Maybe TypeInfo
staticTypeInfo Oid
oid' of
    Just name' :: TypeInfo
name' -> TypeInfo -> IO TypeInfo
forall (m :: * -> *) a. Monad m => a -> m a
return TypeInfo
name'
    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' conn :: Connection
conn oid' :: Oid
oid' oidmap :: TypeInfoCache
oidmap =
  case Key -> TypeInfoCache -> Maybe TypeInfo
forall a. Key -> IntMap a -> Maybe a
IntMap.lookup (Oid -> Key
oid2int Oid
oid') TypeInfoCache
oidmap of
    Just typeinfo :: TypeInfo
typeinfo -> (TypeInfoCache, TypeInfo) -> IO (TypeInfoCache, TypeInfo)
forall (m :: * -> *) a. Monad m => a -> m a
return (TypeInfoCache
oidmap, TypeInfo
typeinfo)
    Nothing -> do
      [(Oid, Char, Char, ByteString, Oid, Oid)]
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 "SELECT oid, typcategory, typdelim, typname,\
                         \ typelem, typrelid\
                         \ FROM pg_type WHERE oid = ?"
                           (Oid -> Only Oid
forall a. a -> Only a
Only Oid
oid')
      (oidmap' :: TypeInfoCache
oidmap', typeInfo :: TypeInfo
typeInfo) <-
          case [(Oid, Char, Char, ByteString, Oid, Oid)]
names of
            []  -> (TypeInfoCache, TypeInfo) -> IO (TypeInfoCache, TypeInfo)
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. Exception e => e -> a
throw (ByteString -> SqlError
fatalError "invalid type oid")
            [(typoid :: Oid
typoid, typcategory :: Char
typcategory, typdelim :: Char
typdelim, typname :: ByteString
typname, typelem_ :: Oid
typelem_, typrelid :: Oid
typrelid)] -> do
               case Char
typcategory of
                 'A' -> do
                   (oidmap' :: TypeInfoCache
oidmap', typelem :: TypeInfo
typelem) <- Connection -> Oid -> TypeInfoCache -> IO (TypeInfoCache, TypeInfo)
getTypeInfo' Connection
conn Oid
typelem_ TypeInfoCache
oidmap
                   let !typeInfo :: TypeInfo
typeInfo = $WArray :: Oid -> Char -> Char -> ByteString -> TypeInfo -> TypeInfo
Array{..}
                   (TypeInfoCache, TypeInfo) -> IO (TypeInfoCache, TypeInfo)
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)
                 'R' -> do
                   [Only Oid]
rngsubtypeOids <- Connection -> Query -> Only Oid -> IO [Only Oid]
forall q r.
(ToRow q, FromRow r) =>
Connection -> Query -> q -> IO [r]
query Connection
conn "SELECT rngsubtype\
                                               \ FROM pg_range\
                                               \ WHERE rngtypid = ?"
                                                (Oid -> Only Oid
forall a. a -> Only a
Only Oid
oid')
                   case [Only Oid]
rngsubtypeOids of
                     [Only rngsubtype_ :: Oid
rngsubtype_] -> do
                        (oidmap' :: TypeInfoCache
oidmap', rngsubtype :: TypeInfo
rngsubtype) <-
                            Connection -> Oid -> TypeInfoCache -> IO (TypeInfoCache, TypeInfo)
getTypeInfo' Connection
conn Oid
rngsubtype_ TypeInfoCache
oidmap
                        let !typeInfo :: TypeInfo
typeInfo = $WRange :: Oid -> Char -> Char -> ByteString -> TypeInfo -> TypeInfo
Range{..}
                        (TypeInfoCache, TypeInfo) -> IO (TypeInfoCache, TypeInfo)
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)
                     _ -> String -> IO (TypeInfoCache, TypeInfo)
forall (m :: * -> *) a. MonadFail m => String -> m a
fail "range subtype query failed to return exactly one result"
                 'C' -> do
                   [(ByteString, Oid)]
cols <- Connection -> Query -> Only Oid -> IO [(ByteString, Oid)]
forall q r.
(ToRow q, FromRow r) =>
Connection -> Query -> q -> IO [r]
query Connection
conn "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)
                   IOVector Attribute
vec <- Key -> IO (IOVector Attribute)
forall (m :: * -> *) a.
PrimMonad m =>
Key -> m (MVector (PrimState m) a)
MV.new (Key -> IO (IOVector Attribute)) -> Key -> IO (IOVector Attribute)
forall a b. (a -> b) -> a -> b
$! [(ByteString, Oid)] -> Key
forall (t :: * -> *) a. Foldable t => t a -> Key
length [(ByteString, Oid)]
cols
                   (oidmap' :: TypeInfoCache
oidmap', attributes :: Vector Attribute
attributes) <- Connection
-> [(ByteString, Oid)]
-> TypeInfoCache
-> IOVector Attribute
-> Key
-> IO (TypeInfoCache, Vector Attribute)
getAttInfos Connection
conn [(ByteString, Oid)]
cols TypeInfoCache
oidmap IOVector Attribute
vec 0
                   let !typeInfo :: TypeInfo
typeInfo = $WComposite :: Oid
-> Char
-> Char
-> ByteString
-> Oid
-> Vector Attribute
-> TypeInfo
Composite{..}
                   (TypeInfoCache, TypeInfo) -> IO (TypeInfoCache, TypeInfo)
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)
                 _ -> do
                   let !typeInfo :: TypeInfo
typeInfo = $WBasic :: Oid -> Char -> Char -> ByteString -> TypeInfo
Basic{..}
                   (TypeInfoCache, TypeInfo) -> IO (TypeInfoCache, TypeInfo)
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)
            _ -> String -> IO (TypeInfoCache, TypeInfo)
forall (m :: * -> *) a. MonadFail m => String -> m a
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'' :: TypeInfoCache
oidmap'' = Key -> TypeInfo -> TypeInfoCache -> TypeInfoCache
forall a. Key -> a -> IntMap a -> IntMap a
IntMap.insert (Oid -> Key
oid2int Oid
oid') TypeInfo
typeInfo TypeInfoCache
oidmap'
      (TypeInfoCache, TypeInfo) -> IO (TypeInfoCache, TypeInfo)
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)

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