{-# LANGUAGE RecordWildCards #-}

------------------------------------------------------------------------------
-- |
-- Module:      Database.PostgreSQL.Simple.TypeInfo
-- Copyright:   (c) 2013 Leon P Smith
-- License:     BSD3
-- Maintainer:  Leon P Smith <[email protected]>
-- 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
      [(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 Query
"SELECT oid, typcategory, typdelim, typname,\
                         \ typelem, typrelid\
                         \ FROM pg_type WHERE oid = ?"
                           (Oid -> Only Oid
forall a. a -> Only a
Only Oid
oid')
      (TypeInfoCache
oidmap', TypeInfo
typeInfo) <-
          case [(Oid, Char, Char, ByteString, Oid, Oid)]
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. 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
                   (TypeInfoCache
oidmap', TypeInfo
typelem) <- Connection -> Oid -> TypeInfoCache -> IO (TypeInfoCache, TypeInfo)
getTypeInfo' Connection
conn Oid
typelem_ TypeInfoCache
oidmap
                   let !typeInfo :: TypeInfo
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
..}
                   (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)
                 Char
'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 Query
"SELECT rngsubtype\
                                               \ FROM pg_range\
                                               \ WHERE rngtypid = ?"
                                                (Oid -> Only Oid
forall a. a -> Only a
Only Oid
oid')
                   case [Only Oid]
rngsubtypeOids of
                     [Only Oid
rngsubtype_] -> do
                        (TypeInfoCache
oidmap', TypeInfo
rngsubtype) <-
                            Connection -> Oid -> TypeInfoCache -> IO (TypeInfoCache, TypeInfo)
getTypeInfo' Connection
conn Oid
rngsubtype_ TypeInfoCache
oidmap
                        let !typeInfo :: TypeInfo
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
..}
                        (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)
                     [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
                   [(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 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)
                   IOVector Attribute
vec <- Int -> IO (IOVector Attribute)
Int -> IO (MVector (PrimState IO) Attribute)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> m (MVector (PrimState m) a)
MV.new (Int -> IO (IOVector Attribute)) -> Int -> IO (IOVector Attribute)
forall a b. (a -> b) -> a -> b
$! [(ByteString, Oid)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(ByteString, Oid)]
cols
                   (TypeInfoCache
oidmap', Vector Attribute
attributes) <- Connection
-> [(ByteString, Oid)]
-> TypeInfoCache
-> IOVector Attribute
-> Int
-> IO (TypeInfoCache, Vector Attribute)
getAttInfos Connection
conn [(ByteString, Oid)]
cols TypeInfoCache
oidmap IOVector Attribute
vec Int
0
                   let !typeInfo :: TypeInfo
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
..}
                   (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)
                 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'' :: TypeInfoCache
oidmap'' = Int -> TypeInfo -> TypeInfoCache -> TypeInfoCache
forall a. Int -> a -> IntMap a -> IntMap a
IntMap.insert (Oid -> Int
oid2int Oid
oid') TypeInfo
typeInfo TypeInfoCache
oidmap'
      (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)

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