module Database.PostgreSQL.Typed.TH
( getTPGDatabase
, withTPGTypeConnection
, 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, withMVar)
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.UTF8 as BSU
import qualified Data.Foldable as Fold
import Data.Maybe (isJust, fromMaybe)
import Data.String (fromString)
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.Protocol
import Database.PostgreSQL.Typed.TypeCache
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 PGTypeConnection)
tpgState = unsafePerformIO $ do
db <- unsafeInterleaveIO getTPGDatabase
newMVar (db, Nothing)
withTPGTypeConnection :: (PGTypeConnection -> IO a) -> IO a
withTPGTypeConnection f = do
(db, tpg') <- takeMVar tpgState
tpg <- maybe (newPGTypeConnection =<< 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 = withTPGTypeConnection (f . pgConnection)
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 (pgConnection t) db
return t{ pgConnection = c }) tpg'
else Nothing <$ Fold.mapM_ (pgDisconnect . pgConnection) tpg')
`onException` putMVar tpgState (db, Nothing)
return []
reloadTPGTypes :: TH.DecsQ
reloadTPGTypes = TH.runIO $ [] <$ withMVar tpgState (mapM_ flushPGTypeConnection . snd)
tpgType :: PGTypeConnection -> OID -> IO PGName
tpgType c o =
maybe (fail $ "Unknown PostgreSQL type: " ++ show o ++ "\nYou may need to use reloadTPGTypes or adjust search_path, or your postgresql-typed application may need to be rebuilt.") return =<< lookupPGType c o
getTPGTypeOID :: PGTypeConnection -> PGName -> IO OID
getTPGTypeOID c t =
maybe (fail $ "Unknown PostgreSQL type: " ++ show t ++ "; be sure to use the exact type name from \\dTS") return =<< findPGType c t
data TPGValueInfo = TPGValueInfo
{ tpgValueName :: BS.ByteString
, tpgValueTypeOID :: !OID
, tpgValueType :: PGName
, tpgValueNullable :: Bool
}
tpgDescribe :: BS.ByteString -> [String] -> Bool -> IO ([TPGValueInfo], [TPGValueInfo])
tpgDescribe sql types nulls = withTPGTypeConnection $ \tpg -> do
at <- mapM (getTPGTypeOID tpg . fromString) types
(pt, rt) <- pgDescribe (pgConnection tpg) (BSL.fromStrict sql) at nulls
(,)
<$> mapM (\o -> do
ot <- tpgType tpg o
return TPGValueInfo
{ tpgValueName = BS.empty
, tpgValueTypeOID = o
, tpgValueType = ot
, tpgValueNullable = True
}) pt
<*> mapM (\(c, o, n) -> do
ot <- tpgType tpg o
return TPGValueInfo
{ tpgValueName = c
, tpgValueTypeOID = o
, tpgValueType = ot
, tpgValueNullable = n && o /= 2278
}) rt
typeApply :: PGName -> 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 ''PGTypeID `TH.AppT` TH.LitT (TH.StrTyLit $ pgNameString $ 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