{-# LANGUAGE CPP, TemplateHaskell, FlexibleInstances, MultiParamTypeClasses, DataKinds #-}
{-# LANGUAGE OverloadedStrings #-}
#if __GLASGOW_HASKELL__ >= 800
{-# LANGUAGE UndecidableSuperClasses #-}
#endif
-- |
-- Module: Database.PostgreSQL.Typed.Enum
-- Copyright: 2015 Dylan Simon
-- 
-- Support for PostgreSQL enums.

module Database.PostgreSQL.Typed.Enum
  ( PGEnum(..)
  , dataPGEnum
  ) where

import           Control.Arrow ((&&&))
import qualified Data.ByteString as BS
import qualified Data.ByteString.Char8 as BSC
import qualified Data.ByteString.Lazy as BSL
import           Data.Ix (Ix)
import           Data.Maybe (fromJust, fromMaybe)
import           Data.Tuple (swap)
import           Data.Typeable (Typeable)
import qualified Language.Haskell.TH as TH

import Database.PostgreSQL.Typed.Types
import Database.PostgreSQL.Typed.Dynamic
import Database.PostgreSQL.Typed.Protocol
import Database.PostgreSQL.Typed.TypeCache
import Database.PostgreSQL.Typed.TH

-- |A type based on a PostgreSQL enum. Automatically instantiated by 'dataPGEnum'.
class (Eq a, Ord a, Enum a, Bounded a, PGRep a) => PGEnum a where
  {-# MINIMAL pgEnumName | pgEnumValues #-}
  -- |The database name of a value.
  pgEnumName :: a -> PGName
  pgEnumName a
a = Maybe PGName -> PGName
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe PGName -> PGName) -> Maybe PGName -> PGName
forall a b. (a -> b) -> a -> b
$ a -> [(a, PGName)] -> Maybe PGName
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup a
a [(a, PGName)]
forall a. PGEnum a => [(a, PGName)]
pgEnumValues
  -- |Lookup a value matching the given database name.
  pgEnumValue :: PGName -> Maybe a
  pgEnumValue PGName
n = PGName -> [(PGName, a)] -> Maybe a
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup PGName
n ([(PGName, a)] -> Maybe a) -> [(PGName, a)] -> Maybe a
forall a b. (a -> b) -> a -> b
$ ((a, PGName) -> (PGName, a)) -> [(a, PGName)] -> [(PGName, a)]
forall a b. (a -> b) -> [a] -> [b]
map (a, PGName) -> (PGName, a)
forall a b. (a, b) -> (b, a)
swap [(a, PGName)]
forall a. PGEnum a => [(a, PGName)]
pgEnumValues
  -- |List of all the values in the enum along with their database names.
  pgEnumValues :: [(a, PGName)]
  pgEnumValues = (a -> (a, PGName)) -> [a] -> [(a, PGName)]
forall a b. (a -> b) -> [a] -> [b]
map (a -> a
forall a. a -> a
id (a -> a) -> (a -> PGName) -> a -> (a, PGName)
forall b c c'. (b -> c) -> (b -> c') -> b -> (c, c')
forall (a :: * -> * -> *) b c c'.
Arrow a =>
a b c -> a b c' -> a b (c, c')
&&& a -> PGName
forall a. PGEnum a => a -> PGName
pgEnumName) ([a] -> [(a, PGName)]) -> [a] -> [(a, PGName)]
forall a b. (a -> b) -> a -> b
$ a -> a -> [a]
forall a. Enum a => a -> a -> [a]
enumFromTo a
forall a. Bounded a => a
minBound a
forall a. Bounded a => a
maxBound

-- |Create a new enum type corresponding to the given PostgreSQL enum type.
-- For example, if you have @CREATE TYPE foo AS ENUM (\'abc\', \'DEF\')@, then
-- @dataPGEnum \"Foo\" \"foo\" (\"Foo_\"++)@ will be equivalent to:
-- 
-- > data Foo = Foo_abc | Foo_DEF deriving (Eq, Ord, Enum, Bounded, Typeable)
-- > instance PGType "foo" where PGVal "foo" = Foo
-- > instance PGParameter "foo" Foo where ...
-- > instance PGColumn "foo" Foo where ...
-- > instance PGRep Foo where PGRepType = "foo"
-- > instance PGEnum Foo where pgEnumValues = [(Foo_abc, "abc"), (Foo_DEF, "DEF")]
--
-- Requires language extensions: TemplateHaskell, FlexibleInstances, MultiParamTypeClasses, DeriveDataTypeable, DataKinds, TypeFamilies
dataPGEnum :: String -- ^ Haskell type to create
  -> PGName -- ^ PostgreSQL enum type name
  -> (String -> String) -- ^ How to generate constructor names from enum values, e.g. @(\"Type_\"++)@ (input is 'pgNameString')
  -> TH.DecsQ
dataPGEnum :: [Char] -> PGName -> ([Char] -> [Char]) -> DecsQ
dataPGEnum [Char]
typs PGName
pgenum [Char] -> [Char]
valnf = do
  (pgid, vals) <- IO (PGName, [PGName]) -> Q (PGName, [PGName])
forall a. IO a -> Q a
TH.runIO (IO (PGName, [PGName]) -> Q (PGName, [PGName]))
-> IO (PGName, [PGName]) -> Q (PGName, [PGName])
forall a b. (a -> b) -> a -> b
$ (PGTypeConnection -> IO (PGName, [PGName]))
-> IO (PGName, [PGName])
forall a. (PGTypeConnection -> IO a) -> IO a
withTPGTypeConnection ((PGTypeConnection -> IO (PGName, [PGName]))
 -> IO (PGName, [PGName]))
-> (PGTypeConnection -> IO (PGName, [PGName]))
-> IO (PGName, [PGName])
forall a b. (a -> b) -> a -> b
$ \PGTypeConnection
tpg -> do
    vals <- ([PGValue] -> (OID, PGName)) -> [[PGValue]] -> [(OID, PGName)]
forall a b. (a -> b) -> [a] -> [b]
map (\([PGValue
eo, PGValue
v]) -> (PGValue -> OID
forall a. PGRep a => PGValue -> a
pgDecodeRep PGValue
eo, PGValue -> PGName
forall a. PGRep a => PGValue -> a
pgDecodeRep PGValue
v)) ([[PGValue]] -> [(OID, PGName)])
-> ((Int, [[PGValue]]) -> [[PGValue]])
-> (Int, [[PGValue]])
-> [(OID, PGName)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int, [[PGValue]]) -> [[PGValue]]
forall a b. (a, b) -> b
snd
      ((Int, [[PGValue]]) -> [(OID, PGName)])
-> IO (Int, [[PGValue]]) -> IO [(OID, PGName)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> PGConnection -> ByteString -> IO (Int, [[PGValue]])
pgSimpleQuery (PGTypeConnection -> PGConnection
pgConnection PGTypeConnection
tpg) ([StrictByteString] -> ByteString
BSL.fromChunks
        [ StrictByteString
"SELECT enumtypid, enumlabel"
        ,  StrictByteString
" FROM pg_catalog.pg_enum"
        , StrictByteString
" WHERE enumtypid = ", PGName -> StrictByteString
forall a. PGRep a => a -> StrictByteString
pgLiteralRep PGName
pgenum, StrictByteString
"::regtype"
        , StrictByteString
" ORDER BY enumsortorder"
        ])
    case vals of
      [] -> [Char] -> IO (PGName, [PGName])
forall a. [Char] -> IO a
forall (m :: * -> *) a. MonadFail m => [Char] -> m a
fail ([Char] -> IO (PGName, [PGName]))
-> [Char] -> IO (PGName, [PGName])
forall a b. (a -> b) -> a -> b
$ [Char]
"dataPGEnum " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
typs [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
" = " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ PGName -> [Char]
forall a. Show a => a -> [Char]
show PGName
pgenum [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
": no values found"
      (OID
eo, PGName
_):[(OID, PGName)]
_ -> do
        et <- IO PGName -> (PGName -> IO PGName) -> Maybe PGName -> IO PGName
forall b a. b -> (a -> b) -> Maybe a -> b
maybe ([Char] -> IO PGName
forall a. [Char] -> IO a
forall (m :: * -> *) a. MonadFail m => [Char] -> m a
fail ([Char] -> IO PGName) -> [Char] -> IO PGName
forall a b. (a -> b) -> a -> b
$ [Char]
"dataPGEnum " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
typs [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
" = " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ PGName -> [Char]
forall a. Show a => a -> [Char]
show PGName
pgenum [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
": enum type not found (you may need to use reloadTPGTypes or adjust search_path)") PGName -> IO PGName
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return
          (Maybe PGName -> IO PGName) -> IO (Maybe PGName) -> IO PGName
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< PGTypeConnection -> OID -> IO (Maybe PGName)
lookupPGType PGTypeConnection
tpg OID
eo
        return (et, map snd vals)
  let valn = (PGName -> (Name, [Lit])) -> [PGName] -> [(Name, [Lit])]
forall a b. (a -> b) -> [a] -> [b]
map ([Char] -> Name
TH.mkName ([Char] -> Name) -> (PGName -> [Char]) -> PGName -> Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Char] -> [Char]
valnf ([Char] -> [Char]) -> (PGName -> [Char]) -> PGName -> [Char]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PGName -> [Char]
pgNameString (PGName -> Name) -> (PGName -> [Lit]) -> PGName -> (Name, [Lit])
forall b c c'. (b -> c) -> (b -> c') -> b -> (c, c')
forall (a :: * -> * -> *) b c c'.
Arrow a =>
a b c -> a b c' -> a b (c, c')
&&& (Word8 -> Lit) -> [Word8] -> [Lit]
forall a b. (a -> b) -> [a] -> [b]
map (Integer -> Lit
TH.IntegerL (Integer -> Lit) -> (Word8 -> Integer) -> Word8 -> Lit
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word8 -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral) ([Word8] -> [Lit]) -> (PGName -> [Word8]) -> PGName -> [Lit]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PGName -> [Word8]
pgNameBytes) [PGName]
vals
      typl = TyLit -> Type
TH.LitT ([Char] -> TyLit
TH.StrTyLit ([Char] -> TyLit) -> [Char] -> TyLit
forall a b. (a -> b) -> a -> b
$ PGName -> [Char]
pgNameString PGName
pgid)
  dv <- TH.newName "x"
  return $
    [ TH.DataD [] typn []
#if MIN_VERSION_template_haskell(2,11,0)
      Nothing
#endif
      (map (\(Name
n, [Lit]
_) -> Name -> [BangType] -> Con
TH.NormalC Name
n []) valn) $
#if MIN_VERSION_template_haskell(2,11,0)
#if MIN_VERSION_template_haskell(2,12,0)
      return $ TH.DerivClause Nothing $
#endif
      map TH.ConT
#endif
      [''Eq, ''Ord, ''Enum, ''Ix, ''Bounded, ''Typeable]
    , instanceD [] (TH.ConT ''PGType `TH.AppT` typl)
      [ tySynInstD ''PGVal typl typt
      ]
    , instanceD [] (TH.ConT ''PGParameter `TH.AppT` typl `TH.AppT` typt)
      [ TH.FunD 'pgEncode [TH.Clause [TH.WildP, TH.VarP dv]
        (TH.NormalB $ TH.VarE 'pgNameBS `TH.AppE` (TH.VarE 'pgEnumName `TH.AppE` TH.VarE dv))
        []]
      ]
    , instanceD [] (TH.ConT ''PGColumn `TH.AppT` typl `TH.AppT` typt)
      [ TH.FunD 'pgDecode [TH.Clause [TH.WildP, TH.VarP dv]
        (TH.NormalB $ TH.VarE 'fromMaybe
          `TH.AppE` (TH.AppE (TH.VarE 'error) $
            TH.InfixE (Just $ TH.LitE (TH.StringL ("pgEnumValue " ++ show pgid ++ ": "))) (TH.VarE '(++)) (Just $ TH.VarE 'BSC.unpack `TH.AppE` TH.VarE dv))
          `TH.AppE` (TH.VarE 'pgEnumValue `TH.AppE` (TH.ConE 'PGName
            `TH.AppE` (TH.VarE 'BS.unpack `TH.AppE` TH.VarE dv))))
        []]
      ]
    , instanceD [] (TH.ConT ''PGRep `TH.AppT` typt)
      [ tySynInstD ''PGRepType typt typl
      ]
    , instanceD [] (TH.ConT ''PGEnum `TH.AppT` typt)
      [ TH.FunD 'pgEnumName $ map (\(Name
n, [Lit]
l) -> [Pat] -> Body -> [Dec] -> Clause
TH.Clause [Name -> [Pat] -> Pat
conP Name
n []]
        (Exp -> Body
TH.NormalB (Exp -> Body) -> Exp -> Body
forall a b. (a -> b) -> a -> b
$ [Lit] -> Exp
namelit [Lit]
l)
        []) valn
      , TH.FunD 'pgEnumValue $ map (\(Name
n, [Lit]
l) ->
          [Pat] -> Body -> [Dec] -> Clause
TH.Clause [Name -> [Pat] -> Pat
conP 'PGName [[Pat] -> Pat
TH.ListP ((Lit -> Pat) -> [Lit] -> [Pat]
forall a b. (a -> b) -> [a] -> [b]
map Lit -> Pat
TH.LitP [Lit]
l)]]
            (Exp -> Body
TH.NormalB (Exp -> Body) -> Exp -> Body
forall a b. (a -> b) -> a -> b
$ Name -> Exp
TH.ConE 'Just Exp -> Exp -> Exp
`TH.AppE` Name -> Exp
TH.ConE Name
n)
            []) valn
          ++ [TH.Clause [TH.WildP] (TH.NormalB $ TH.ConE 'Nothing) []]
      , TH.FunD 'pgEnumValues [TH.Clause []
        (TH.NormalB $ TH.ListE $ map (\(Name
n, [Lit]
l) ->
          Name -> Exp
TH.ConE '(,) Exp -> Exp -> Exp
`TH.AppE` Name -> Exp
TH.ConE Name
n Exp -> Exp -> Exp
`TH.AppE` [Lit] -> Exp
namelit [Lit]
l) valn)
        []]
      ]
    , TH.PragmaD $ TH.AnnP (TH.TypeAnnotation typn) $ namelit $ map (TH.IntegerL . fromIntegral) $ pgNameBytes pgid
    ]
    ++ map (\(Name
n, [Lit]
l) ->
      Pragma -> Dec
TH.PragmaD (Pragma -> Dec) -> Pragma -> Dec
forall a b. (a -> b) -> a -> b
$ AnnTarget -> Exp -> Pragma
TH.AnnP (Name -> AnnTarget
TH.ValueAnnotation Name
n) (Exp -> Pragma) -> Exp -> Pragma
forall a b. (a -> b) -> a -> b
$ [Lit] -> Exp
namelit [Lit]
l) valn
  where
  typn :: Name
typn = [Char] -> Name
TH.mkName [Char]
typs
  typt :: Type
typt = Name -> Type
TH.ConT Name
typn
  instanceD :: Cxt -> Type -> [Dec] -> Dec
instanceD = Maybe Overlap -> Cxt -> Type -> [Dec] -> Dec
TH.InstanceD
#if MIN_VERSION_template_haskell(2,11,0)
      Maybe Overlap
forall a. Maybe a
Nothing
#endif
  tySynInstD :: Name -> Type -> Type -> Dec
tySynInstD Name
c Type
l Type
t = TySynEqn -> Dec
TH.TySynInstD
#if MIN_VERSION_template_haskell(2,15,0)
    (TySynEqn -> Dec) -> TySynEqn -> Dec
forall a b. (a -> b) -> a -> b
$ Maybe [TyVarBndr ()] -> Type -> Type -> TySynEqn
TH.TySynEqn Maybe [TyVarBndr ()]
forall a. Maybe a
Nothing (Type -> Type -> Type
TH.AppT (Name -> Type
TH.ConT Name
c) Type
l)
#else
    c $ TH.TySynEqn [l]
#endif
    Type
t
  namelit :: [Lit] -> Exp
namelit [Lit]
l = Name -> Exp
TH.ConE 'PGName Exp -> Exp -> Exp
`TH.AppE` [Exp] -> Exp
TH.ListE ((Lit -> Exp) -> [Lit] -> [Exp]
forall a b. (a -> b) -> [a] -> [b]
map Lit -> Exp
TH.LitE [Lit]
l)
  conP :: Name -> [Pat] -> Pat
conP Name
n [Pat]
p = Name -> Cxt -> [Pat] -> Pat
TH.ConP Name
n
#if MIN_VERSION_template_haskell(2,18,0)
    []
#endif
    [Pat]
p