{-# LANGUAGE OverloadedStrings #-}
module Database.PostgreSQL.Typed.TypeCache
  ( PGTypes
  , pgGetTypes
  , PGTypeConnection
  , pgConnection
  , newPGTypeConnection
  , flushPGTypeConnection
  , lookupPGType
  , findPGType
  ) where

import Data.IORef (IORef, newIORef, readIORef, writeIORef)
import qualified Data.IntMap as IntMap
import Data.List (find)

import Database.PostgreSQL.Typed.Types (PGName, OID)
import Database.PostgreSQL.Typed.Dynamic
import Database.PostgreSQL.Typed.Protocol

-- |Map keyed on fromIntegral OID.
type PGTypes = IntMap.IntMap PGName

-- |A 'PGConnection' along with cached information about types.
data PGTypeConnection = PGTypeConnection
  { PGTypeConnection -> PGConnection
pgConnection :: !PGConnection
  , PGTypeConnection -> IORef (Maybe PGTypes)
pgTypes :: IORef (Maybe PGTypes)
  }

-- |Create a 'PGTypeConnection'.
newPGTypeConnection :: PGConnection -> IO PGTypeConnection
newPGTypeConnection :: PGConnection -> IO PGTypeConnection
newPGTypeConnection PGConnection
c = do
  IORef (Maybe PGTypes)
t <- forall a. a -> IO (IORef a)
newIORef forall a. Maybe a
Nothing
  forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ PGConnection -> IORef (Maybe PGTypes) -> PGTypeConnection
PGTypeConnection PGConnection
c IORef (Maybe PGTypes)
t

-- |Flush the cached type list, forcing it to be reloaded.
flushPGTypeConnection :: PGTypeConnection -> IO ()
flushPGTypeConnection :: PGTypeConnection -> IO ()
flushPGTypeConnection PGTypeConnection
c =
  forall a. IORef a -> a -> IO ()
writeIORef (PGTypeConnection -> IORef (Maybe PGTypes)
pgTypes PGTypeConnection
c) forall a. Maybe a
Nothing

-- |Get a map of types from the database.
pgGetTypes :: PGConnection -> IO PGTypes
pgGetTypes :: PGConnection -> IO PGTypes
pgGetTypes PGConnection
c =
  forall a. [(Key, a)] -> IntMap a
IntMap.fromAscList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map (\[PGValue
to, PGValue
tn] -> (forall a b. (Integral a, Num b) => a -> b
fromIntegral (forall a. PGRep a => PGValue -> a
pgDecodeRep PGValue
to :: OID), forall a. PGRep a => PGValue -> a
pgDecodeRep PGValue
tn)) forall b c a. (b -> c) -> (a -> b) -> a -> c
.
    forall a b. (a, b) -> b
snd forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> PGConnection -> ByteString -> IO (Key, [[PGValue]])
pgSimpleQuery PGConnection
c ByteString
"SELECT oid, format_type(CASE WHEN typtype = 'd' THEN typbasetype ELSE oid END, -1) FROM pg_catalog.pg_type ORDER BY oid"

-- |Get a cached map of types.
getPGTypes :: PGTypeConnection -> IO PGTypes
getPGTypes :: PGTypeConnection -> IO PGTypes
getPGTypes (PGTypeConnection PGConnection
c IORef (Maybe PGTypes)
tr) =
  forall b a. b -> (a -> b) -> Maybe a -> b
maybe (do
      PGTypes
t <- PGConnection -> IO PGTypes
pgGetTypes PGConnection
c
      forall a. IORef a -> a -> IO ()
writeIORef IORef (Maybe PGTypes)
tr forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just PGTypes
t
      forall (m :: * -> *) a. Monad m => a -> m a
return PGTypes
t)
    forall (m :: * -> *) a. Monad m => a -> m a
return
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall a. IORef a -> IO a
readIORef IORef (Maybe PGTypes)
tr

-- |Lookup a type name by OID.
-- This is an efficient, often pure operation.
lookupPGType :: PGTypeConnection -> OID -> IO (Maybe PGName)
lookupPGType :: PGTypeConnection -> OID -> IO (Maybe PGName)
lookupPGType PGTypeConnection
c OID
o =
  forall a. Key -> IntMap a -> Maybe a
IntMap.lookup (forall a b. (Integral a, Num b) => a -> b
fromIntegral OID
o) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> PGTypeConnection -> IO PGTypes
getPGTypes PGTypeConnection
c

-- |Lookup a type OID by type name.
-- This is less common and thus less efficient than going the other way.
findPGType :: PGTypeConnection -> PGName -> IO (Maybe OID)
findPGType :: PGTypeConnection -> PGName -> IO (Maybe OID)
findPGType PGTypeConnection
c PGName
t =
  forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall a b. (Integral a, Num b) => a -> b
fromIntegral forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find (forall a. Eq a => a -> a -> Bool
(==) PGName
t forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. IntMap a -> [(Key, a)]
IntMap.toList forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> PGTypeConnection -> IO PGTypes
getPGTypes PGTypeConnection
c