module Database.PostgreSQL.Typed.TH
( getTPGDatabase
, withTPGConnection
, useTPGDatabase
, reloadTPGTypes
, TPGValueInfo(..)
, tpgDescribe
, tpgTypeEncoder
, tpgTypeDecoder
, tpgTypeBinary
) where
#if !MIN_VERSION_base(4,8,0)
import Control.Applicative ((<$>), (<$))
#endif
import Control.Applicative ((<|>))
import Control.Concurrent.MVar (MVar, newMVar, takeMVar, putMVar, modifyMVar_)
import Control.Exception (onException, finally)
import Control.Monad (liftM2)
import qualified Data.ByteString as BS
import qualified Data.ByteString.Lazy as BSL
import qualified Data.ByteString.Lazy.Char8 as BSLC
import qualified Data.ByteString.UTF8 as BSU
import qualified Data.Foldable as Fold
import qualified Data.IntMap.Lazy as IntMap
import Data.List (find)
import Data.Maybe (isJust, fromMaybe)
import qualified Data.Traversable as Tv
import qualified Language.Haskell.TH as TH
import Network (PortID(UnixSocket, PortNumber), PortNumber)
import System.Environment (lookupEnv)
import System.IO.Unsafe (unsafePerformIO, unsafeInterleaveIO)
import Database.PostgreSQL.Typed.Types
import Database.PostgreSQL.Typed.Dynamic
import Database.PostgreSQL.Typed.Protocol
type TPGType = String
getTPGDatabase :: IO PGDatabase
getTPGDatabase = do
user <- fromMaybe "postgres" <$> liftM2 (<|>) (lookupEnv "TPG_USER") (lookupEnv "USER")
db <- fromMaybe user <$> lookupEnv "TPG_DB"
host <- fromMaybe "localhost" <$> lookupEnv "TPG_HOST"
pnum <- maybe (5432 :: PortNumber) ((fromIntegral :: Int -> PortNumber) . read) <$> lookupEnv "TPG_PORT"
port <- maybe (PortNumber pnum) UnixSocket <$> lookupEnv "TPG_SOCK"
pass <- fromMaybe "" <$> lookupEnv "TPG_PASS"
debug <- isJust <$> lookupEnv "TPG_DEBUG"
return $ defaultPGDatabase
{ pgDBHost = host
, pgDBPort = port
, pgDBName = BSU.fromString db
, pgDBUser = BSU.fromString user
, pgDBPass = BSU.fromString pass
, pgDBDebug = debug
}
tpgState :: MVar (PGDatabase, Maybe TPGState)
tpgState = unsafePerformIO $ do
db <- unsafeInterleaveIO getTPGDatabase
newMVar (db, Nothing)
data TPGState = TPGState
{ tpgConnection :: PGConnection
, tpgTypes :: IntMap.IntMap TPGType
}
tpgLoadTypes :: TPGState -> IO TPGState
tpgLoadTypes tpg = do
t <- IntMap.fromAscList . map (\[to, tn] -> (fromIntegral (pgDecodeRep to :: OID), pgDecodeRep tn)) .
snd <$> pgSimpleQuery (tpgConnection tpg) (BSLC.pack "SELECT typ.oid, format_type(CASE WHEN typtype = 'd' THEN typbasetype ELSE typ.oid END, -1) FROM pg_catalog.pg_type typ JOIN pg_catalog.pg_namespace nsp ON typnamespace = nsp.oid WHERE nspname <> 'pg_toast' AND nspname <> 'information_schema' ORDER BY typ.oid")
return tpg{ tpgTypes = t }
tpgInit :: PGConnection -> IO TPGState
tpgInit c = tpgLoadTypes TPGState{ tpgConnection = c, tpgTypes = undefined }
withTPGState :: (TPGState -> IO a) -> IO a
withTPGState f = do
(db, tpg') <- takeMVar tpgState
tpg <- maybe (tpgInit =<< pgConnect db) return tpg'
`onException` putMVar tpgState (db, Nothing)
f tpg `finally` putMVar tpgState (db, Just tpg)
withTPGConnection :: (PGConnection -> IO a) -> IO a
withTPGConnection f = withTPGState (f . tpgConnection)
useTPGDatabase :: PGDatabase -> TH.DecsQ
useTPGDatabase db = TH.runIO $ do
(db', tpg') <- takeMVar tpgState
putMVar tpgState . (,) db =<<
(if db == db'
then Tv.mapM (\t -> do
c <- pgReconnect (tpgConnection t) db
return t{ tpgConnection = c }) tpg'
else Nothing <$ Fold.mapM_ (pgDisconnect . tpgConnection) tpg')
`onException` putMVar tpgState (db, Nothing)
return []
reloadTPGTypes :: TH.DecsQ
reloadTPGTypes = TH.runIO $ [] <$ modifyMVar_ tpgState (\(d, c) -> (,) d <$> Tv.mapM tpgLoadTypes c)
tpgType :: TPGState -> OID -> TPGType
tpgType TPGState{ tpgTypes = types } t =
IntMap.findWithDefault (error $ "Unknown PostgreSQL type: " ++ show t) (fromIntegral t) types
getTPGTypeOID :: Monad m => TPGState -> String -> m OID
getTPGTypeOID TPGState{ tpgTypes = types } t =
maybe (fail $ "Unknown PostgreSQL type: " ++ t ++ "; be sure to use the exact type name from \\dTS") (return . fromIntegral . fst)
$ find ((==) t . snd) $ IntMap.toList types
data TPGValueInfo = TPGValueInfo
{ tpgValueName :: BS.ByteString
, tpgValueTypeOID :: !OID
, tpgValueType :: TPGType
, tpgValueNullable :: Bool
}
tpgDescribe :: BS.ByteString -> [String] -> Bool -> IO ([TPGValueInfo], [TPGValueInfo])
tpgDescribe sql types nulls = withTPGState $ \tpg -> do
at <- mapM (getTPGTypeOID tpg) types
(pt, rt) <- pgDescribe (tpgConnection tpg) (BSL.fromStrict sql) at nulls
return
( map (\o -> TPGValueInfo
{ tpgValueName = BS.empty
, tpgValueTypeOID = o
, tpgValueType = tpgType tpg o
, tpgValueNullable = True
}) pt
, map (\(c, o, n) -> TPGValueInfo
{ tpgValueName = c
, tpgValueTypeOID = o
, tpgValueType = tpgType tpg o
, tpgValueNullable = n && o /= 2278
}) rt
)
typeApply :: TPGType -> TH.Name -> TH.Name -> TH.Exp
typeApply t f e =
TH.VarE f `TH.AppE` TH.VarE e
`TH.AppE` (TH.ConE 'PGTypeProxy `TH.SigE` (TH.ConT ''PGTypeName `TH.AppT` TH.LitT (TH.StrTyLit t)))
tpgTypeEncoder :: Bool -> TPGValueInfo -> TH.Name -> TH.Exp
tpgTypeEncoder lit v = typeApply (tpgValueType v) $
if lit
then 'pgEscapeParameter
else 'pgEncodeParameter
tpgTypeDecoder :: Bool -> TPGValueInfo -> TH.Name -> TH.Exp
tpgTypeDecoder nulls v = typeApply (tpgValueType v) $
if nulls && tpgValueNullable v
then 'pgDecodeColumn
else 'pgDecodeColumnNotNull
tpgTypeBinary :: TPGValueInfo -> TH.Name -> TH.Exp
tpgTypeBinary v = typeApply (tpgValueType v) 'pgBinaryColumn