{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE EmptyCase #-}
module Database.Beam.Postgres.CustomTypes
    ( PgType, PgTypeCheck(..)
    , PgDataTypeSchema

    , IsPgCustomDataType(..)

    , PgHasEnum(..)

    , HasSqlValueSyntax, FromBackendRow

    , pgCustomEnumSchema, pgBoundedEnumSchema

    , pgCustomEnumActionProvider
    , pgCreateEnumActionProvider
    , pgDropEnumActionProvider

    , pgChecksForTypeSchema

    , pgEnumValueSyntax, pgParseEnum

    , createEnum
    , beamTypeForCustomPg
    ) where

import           Database.Beam
import           Database.Beam.Schema.Tables
import           Database.Beam.Backend.SQL
import           Database.Beam.Migrate
import           Database.Beam.Postgres.Types
import           Database.Beam.Postgres.Syntax

import           Control.Monad
import           Control.Monad.Free.Church
import           Data.Aeson (object, (.=))
import qualified Data.ByteString.Char8 as BC
import           Data.Functor.Const
import qualified Data.HashSet as HS
import           Data.Proxy (Proxy(..))
#if !MIN_VERSION_base(4,11,0)
import           Data.Semigroup
#endif
import           Data.Text (Text)
import qualified Data.Text.Encoding as TE

import qualified Database.PostgreSQL.Simple.FromField as Pg

data PgType a
newtype PgTypeCheck = PgTypeCheck (Text -> SomeDatabasePredicate)

data PgDataTypeSchema a where
    PgDataTypeEnum :: HasSqlValueSyntax PgValueSyntax a => [a] -> PgDataTypeSchema a

class IsPgCustomDataType a where
    pgDataTypeName :: Proxy a -> Text
    pgDataTypeDescription :: PgDataTypeSchema a

pgCustomEnumSchema :: HasSqlValueSyntax PgValueSyntax a => [a] -> PgDataTypeSchema a
pgCustomEnumSchema :: forall a.
HasSqlValueSyntax PgValueSyntax a =>
[a] -> PgDataTypeSchema a
pgCustomEnumSchema = forall a.
HasSqlValueSyntax PgValueSyntax a =>
[a] -> PgDataTypeSchema a
PgDataTypeEnum

pgBoundedEnumSchema :: ( Enum a, Bounded a, HasSqlValueSyntax PgValueSyntax a )
                    => PgDataTypeSchema a
pgBoundedEnumSchema :: forall a.
(Enum a, Bounded a, HasSqlValueSyntax PgValueSyntax a) =>
PgDataTypeSchema a
pgBoundedEnumSchema = forall a.
HasSqlValueSyntax PgValueSyntax a =>
[a] -> PgDataTypeSchema a
pgCustomEnumSchema [forall a. Bounded a => a
minBound..forall a. Bounded a => a
maxBound]

pgCustomEnumActionProvider :: ActionProvider Postgres
pgCustomEnumActionProvider :: ActionProvider Postgres
pgCustomEnumActionProvider = ActionProvider Postgres
pgCreateEnumActionProvider forall a. Semigroup a => a -> a -> a
<> ActionProvider Postgres
pgDropEnumActionProvider

pgCreateEnumActionProvider :: ActionProvider Postgres
pgCreateEnumActionProvider :: ActionProvider Postgres
pgCreateEnumActionProvider =
  forall be. ActionProviderFn be -> ActionProvider be
ActionProvider forall a b. (a -> b) -> a -> b
$ \forall preCondition. Typeable preCondition => [preCondition]
findPre forall preCondition. Typeable preCondition => [preCondition]
findPost ->
  do enumP :: PgHasEnum
enumP@(PgHasEnum Text
nm [Text]
vals) <- forall preCondition. Typeable preCondition => [preCondition]
findPost
     forall (m :: * -> *) a. Alternative m => [a] -> m ()
ensuringNot_ forall a b. (a -> b) -> a -> b
$
      do (PgHasEnum Text
beforeNm [Text]
_) <- forall preCondition. Typeable preCondition => [preCondition]
findPre
         forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Text
beforeNm forall a. Eq a => a -> a -> Bool
== Text
nm)

     let cmd :: PgCommandSyntax
cmd = Text -> [PgValueSyntax] -> PgCommandSyntax
pgCreateEnumSyntax Text
nm (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall expr ty. HasSqlValueSyntax expr ty => ty -> expr
sqlValueSyntax [Text]
vals)
     forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall be.
HashSet SomeDatabasePredicate
-> HashSet SomeDatabasePredicate
-> Seq (MigrationCommand be)
-> Text
-> Int
-> PotentialAction be
PotentialAction forall a. Monoid a => a
mempty (forall a. (Eq a, Hashable a) => [a] -> HashSet a
HS.fromList [forall p. DatabasePredicate p => p -> SomeDatabasePredicate
p PgHasEnum
enumP])
                           (forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall be.
BeamSqlBackendSyntax be -> MigrationDataLoss -> MigrationCommand be
MigrationCommand PgCommandSyntax
cmd MigrationDataLoss
MigrationKeepsData))
                           (Text
"Create the enumeration " forall a. Semigroup a => a -> a -> a
<> Text
nm) Int
1)

pgDropEnumActionProvider :: ActionProvider Postgres
pgDropEnumActionProvider :: ActionProvider Postgres
pgDropEnumActionProvider =
  forall be. ActionProviderFn be -> ActionProvider be
ActionProvider forall a b. (a -> b) -> a -> b
$ \forall preCondition. Typeable preCondition => [preCondition]
findPre forall preCondition. Typeable preCondition => [preCondition]
findPost ->
  do enumP :: PgHasEnum
enumP@(PgHasEnum Text
nm [Text]
_) <- forall preCondition. Typeable preCondition => [preCondition]
findPre
     forall (m :: * -> *) a. Alternative m => [a] -> m ()
ensuringNot_ forall a b. (a -> b) -> a -> b
$
      do (PgHasEnum Text
afterNm [Text]
_) <- forall preCondition. Typeable preCondition => [preCondition]
findPost
         forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Text
afterNm forall a. Eq a => a -> a -> Bool
== Text
nm)

     let cmd :: PgCommandSyntax
cmd = Text -> PgCommandSyntax
pgDropTypeSyntax Text
nm
     forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall be.
HashSet SomeDatabasePredicate
-> HashSet SomeDatabasePredicate
-> Seq (MigrationCommand be)
-> Text
-> Int
-> PotentialAction be
PotentialAction (forall a. (Eq a, Hashable a) => [a] -> HashSet a
HS.fromList [forall p. DatabasePredicate p => p -> SomeDatabasePredicate
p PgHasEnum
enumP]) forall a. Monoid a => a
mempty
                           (forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall be.
BeamSqlBackendSyntax be -> MigrationDataLoss -> MigrationCommand be
MigrationCommand PgCommandSyntax
cmd MigrationDataLoss
MigrationKeepsData))
                           (Text
"Drop the enumeration type " forall a. Semigroup a => a -> a -> a
<> Text
nm) Int
1)

pgChecksForTypeSchema :: PgDataTypeSchema a -> [ PgTypeCheck ]
pgChecksForTypeSchema :: forall a. PgDataTypeSchema a -> [PgTypeCheck]
pgChecksForTypeSchema (PgDataTypeEnum [a]
vals) =
  let valTxts :: [Text]
valTxts = forall a b. (a -> b) -> [a] -> [b]
map forall {ty}. HasSqlValueSyntax PgValueSyntax ty => ty -> Text
encodeToString [a]
vals

      -- TODO better reporting
      encodeToString :: ty -> Text
encodeToString ty
val =
        let PgValueSyntax (PgSyntax PgSyntaxM ()
syntax) = forall expr ty. HasSqlValueSyntax expr ty => ty -> expr
sqlValueSyntax ty
val
        in forall (f :: * -> *) a.
F f a -> forall r. (a -> r) -> (f r -> r) -> r
runF PgSyntaxM ()
syntax (\()
_ -> forall a. HasCallStack => [Char] -> a
error [Char]
"Expecting a simple text encoding for enumeration type")
                       (\case
                           EmitByteString ByteString
"'" Text
next -> Text
next
                           EscapeString ByteString
s Text
_ -> ByteString -> Text
TE.decodeUtf8 ByteString
s -- TODO Make this more robust
                           PgSyntaxF Text
_ -> forall a. HasCallStack => [Char] -> a
error [Char]
"Expecting a simple text encoding for enumeration type")
  in [ (Text -> SomeDatabasePredicate) -> PgTypeCheck
PgTypeCheck (\Text
nm -> forall p. DatabasePredicate p => p -> SomeDatabasePredicate
p (Text -> [Text] -> PgHasEnum
PgHasEnum Text
nm [Text]
valTxts)) ]

instance IsDatabaseEntity Postgres (PgType a) where

  data DatabaseEntityDescriptor Postgres (PgType a) where
      PgTypeDescriptor :: Maybe Text -> Text -> PgDataTypeSyntax
                       -> DatabaseEntityDescriptor Postgres (PgType a)

  type DatabaseEntityDefaultRequirements Postgres (PgType a) =
      ( HasSqlValueSyntax PgValueSyntax a
      , FromBackendRow Postgres a
      , IsPgCustomDataType a)

  type DatabaseEntityRegularRequirements Postgres (PgType a) =
      ( HasSqlValueSyntax PgValueSyntax a
      , FromBackendRow Postgres a )

  dbEntityName :: Lens' (DatabaseEntityDescriptor Postgres (PgType a)) Text
dbEntityName Text -> f Text
f (PgTypeDescriptor Maybe Text
sch Text
nm PgDataTypeSyntax
ty) = (\Text
nm' -> forall a.
Maybe Text
-> Text
-> PgDataTypeSyntax
-> DatabaseEntityDescriptor Postgres (PgType a)
PgTypeDescriptor Maybe Text
sch Text
nm' PgDataTypeSyntax
ty) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Text -> f Text
f Text
nm
  dbEntitySchema :: Traversal'
  (DatabaseEntityDescriptor Postgres (PgType a)) (Maybe Text)
dbEntitySchema Maybe Text -> f (Maybe Text)
f (PgTypeDescriptor Maybe Text
sch Text
nm PgDataTypeSyntax
ty) = forall a.
Maybe Text
-> Text
-> PgDataTypeSyntax
-> DatabaseEntityDescriptor Postgres (PgType a)
PgTypeDescriptor forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe Text -> f (Maybe Text)
f Maybe Text
sch forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) a. Applicative f => a -> f a
pure Text
nm forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) a. Applicative f => a -> f a
pure PgDataTypeSyntax
ty
  dbEntityAuto :: DatabaseEntityDefaultRequirements Postgres (PgType a) =>
Text -> DatabaseEntityDescriptor Postgres (PgType a)
dbEntityAuto Text
_ = forall a.
Maybe Text
-> Text
-> PgDataTypeSyntax
-> DatabaseEntityDescriptor Postgres (PgType a)
PgTypeDescriptor forall a. Maybe a
Nothing Text
typeName
                                    (PgDataTypeDescr
-> PgSyntax -> BeamSerializedDataType -> PgDataTypeSyntax
PgDataTypeSyntax (Text -> PgDataTypeDescr
PgDataTypeDescrDomain Text
typeName)
                                                      (Text -> PgSyntax
pgQuotedIdentifier Text
typeName)
                                                      (Value -> BeamSerializedDataType
pgDataTypeJSON ([Pair] -> Value
object [ Key
"customType" forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
.= Text
typeName])))
      where
        typeName :: Text
typeName = forall a. IsPgCustomDataType a => Proxy a -> Text
pgDataTypeName (forall {k} (t :: k). Proxy t
Proxy @a)

instance IsCheckedDatabaseEntity Postgres (PgType a) where
    data CheckedDatabaseEntityDescriptor Postgres (PgType a) where
        CheckedPgTypeDescriptor :: DatabaseEntityDescriptor Postgres (PgType a)
                                -> [ PgTypeCheck ]
                                -> CheckedDatabaseEntityDescriptor Postgres (PgType a)
    type CheckedDatabaseEntityDefaultRequirements Postgres (PgType a) =
        DatabaseEntityDefaultRequirements Postgres (PgType a)

    unChecked :: Lens'
  (CheckedDatabaseEntityDescriptor Postgres (PgType a))
  (DatabaseEntityDescriptor Postgres (PgType a))
unChecked DatabaseEntityDescriptor Postgres (PgType a)
-> f (DatabaseEntityDescriptor Postgres (PgType a))
f (CheckedPgTypeDescriptor DatabaseEntityDescriptor Postgres (PgType a)
ty [PgTypeCheck]
d) = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\DatabaseEntityDescriptor Postgres (PgType a)
ty' -> forall a.
DatabaseEntityDescriptor Postgres (PgType a)
-> [PgTypeCheck]
-> CheckedDatabaseEntityDescriptor Postgres (PgType a)
CheckedPgTypeDescriptor DatabaseEntityDescriptor Postgres (PgType a)
ty' [PgTypeCheck]
d) (DatabaseEntityDescriptor Postgres (PgType a)
-> f (DatabaseEntityDescriptor Postgres (PgType a))
f DatabaseEntityDescriptor Postgres (PgType a)
ty)
    collectEntityChecks :: CheckedDatabaseEntityDescriptor Postgres (PgType a)
-> [SomeDatabasePredicate]
collectEntityChecks (CheckedPgTypeDescriptor DatabaseEntityDescriptor Postgres (PgType a)
e [PgTypeCheck]
chks) =
        forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\(PgTypeCheck Text -> SomeDatabasePredicate
mkCheck) -> Text -> SomeDatabasePredicate
mkCheck (forall {k} a (b :: k). Const a b -> a
getConst (forall be entityType.
IsDatabaseEntity be entityType =>
Lens' (DatabaseEntityDescriptor be entityType) Text
dbEntityName forall {k} a (b :: k). a -> Const a b
Const DatabaseEntityDescriptor Postgres (PgType a)
e))) [PgTypeCheck]
chks
    checkedDbEntityAuto :: CheckedDatabaseEntityDefaultRequirements Postgres (PgType a) =>
Text -> CheckedDatabaseEntityDescriptor Postgres (PgType a)
checkedDbEntityAuto Text
nm = forall a.
DatabaseEntityDescriptor Postgres (PgType a)
-> [PgTypeCheck]
-> CheckedDatabaseEntityDescriptor Postgres (PgType a)
CheckedPgTypeDescriptor (forall be entityType.
(IsDatabaseEntity be entityType,
 DatabaseEntityDefaultRequirements be entityType) =>
Text -> DatabaseEntityDescriptor be entityType
dbEntityAuto Text
nm)
                                                     (forall a. PgDataTypeSchema a -> [PgTypeCheck]
pgChecksForTypeSchema (forall a. IsPgCustomDataType a => PgDataTypeSchema a
pgDataTypeDescription @a))

instance RenamableWithRule (FieldRenamer (DatabaseEntityDescriptor Postgres (PgType a))) where
    renamingFields :: (NonEmpty Text -> Text)
-> FieldRenamer (DatabaseEntityDescriptor Postgres (PgType a))
renamingFields NonEmpty Text -> Text
_ = forall entity. (entity -> entity) -> FieldRenamer entity
FieldRenamer forall a. a -> a
id

createEnum :: forall a db
            . ( HasSqlValueSyntax PgValueSyntax a
              , Enum a, Bounded a )
           => Text -> Migration Postgres (CheckedDatabaseEntity Postgres db (PgType a))
createEnum :: forall a (db :: (* -> *) -> *).
(HasSqlValueSyntax PgValueSyntax a, Enum a, Bounded a) =>
Text
-> Migration
     Postgres (CheckedDatabaseEntity Postgres db (PgType a))
createEnum Text
nm = do
  forall be.
BeamSqlBackendSyntax be
-> Maybe (BeamSqlBackendSyntax be) -> Migration be ()
upDown (Text -> [PgValueSyntax] -> PgCommandSyntax
pgCreateEnumSyntax Text
nm (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall expr ty. HasSqlValueSyntax expr ty => ty -> expr
sqlValueSyntax [forall a. Bounded a => a
minBound..(forall a. Bounded a => a
maxBound::a)]))
         (forall a. a -> Maybe a
Just (Text -> PgCommandSyntax
pgDropTypeSyntax Text
nm))

  let tyDesc :: DatabaseEntityDescriptor Postgres (PgType a)
tyDesc = forall a.
Maybe Text
-> Text
-> PgDataTypeSyntax
-> DatabaseEntityDescriptor Postgres (PgType a)
PgTypeDescriptor forall a. Maybe a
Nothing Text
nm forall a b. (a -> b) -> a -> b
$
               PgDataTypeDescr
-> PgSyntax -> BeamSerializedDataType -> PgDataTypeSyntax
PgDataTypeSyntax (Text -> PgDataTypeDescr
PgDataTypeDescrDomain Text
nm)
                                (Text -> PgSyntax
pgQuotedIdentifier Text
nm)
                                (Value -> BeamSerializedDataType
pgDataTypeJSON ([Pair] -> Value
object [ Key
"customType" forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
.= Text
nm ]))

  forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall be entityType (db :: (* -> *) -> *).
IsCheckedDatabaseEntity be entityType =>
CheckedDatabaseEntityDescriptor be entityType
-> [SomeDatabasePredicate]
-> CheckedDatabaseEntity be db entityType
CheckedDatabaseEntity
          (forall a.
DatabaseEntityDescriptor Postgres (PgType a)
-> [PgTypeCheck]
-> CheckedDatabaseEntityDescriptor Postgres (PgType a)
CheckedPgTypeDescriptor DatabaseEntityDescriptor Postgres (PgType a)
tyDesc
             (forall a. PgDataTypeSchema a -> [PgTypeCheck]
pgChecksForTypeSchema (forall a.
HasSqlValueSyntax PgValueSyntax a =>
[a] -> PgDataTypeSchema a
PgDataTypeEnum [forall a. Bounded a => a
minBound..forall a. Bounded a => a
maxBound::a])))
          [])


pgEnumValueSyntax :: (a -> String) -> a -> PgValueSyntax
pgEnumValueSyntax :: forall a. (a -> [Char]) -> a -> PgValueSyntax
pgEnumValueSyntax a -> [Char]
namer = forall expr ty. HasSqlValueSyntax expr ty => ty -> expr
sqlValueSyntax forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> [Char]
namer

newtype PgRawString = PgRawString String
instance FromBackendRow Postgres PgRawString
instance Pg.FromField PgRawString where
    fromField :: FieldParser PgRawString
fromField Field
f Maybe ByteString
Nothing = forall a err.
(Typeable a, Exception err) =>
([Char] -> Maybe Oid -> [Char] -> [Char] -> [Char] -> err)
-> Field -> [Char] -> Conversion a
Pg.returnError [Char] -> Maybe Oid -> [Char] -> [Char] -> [Char] -> ResultError
Pg.UnexpectedNull Field
f [Char]
"When parsing enumeration string"
    fromField Field
_ (Just ByteString
d) = forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Char] -> PgRawString
PgRawString (ByteString -> [Char]
BC.unpack ByteString
d))

pgParseEnum :: (Enum a, Bounded a) => (a -> String)
            -> FromBackendRowM Postgres a
pgParseEnum :: forall a.
(Enum a, Bounded a) =>
(a -> [Char]) -> FromBackendRowM Postgres a
pgParseEnum a -> [Char]
namer =
  let allNames :: [([Char], a)]
allNames = forall a b. (a -> b) -> [a] -> [b]
map (\a
x -> (a -> [Char]
namer a
x, a
x)) [forall a. Bounded a => a
minBound..forall a. Bounded a => a
maxBound]
  in do
    PgRawString [Char]
name <- forall be a. FromBackendRow be a => FromBackendRowM be a
fromBackendRow
    case forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup [Char]
name [([Char], a)]
allNames of
      Maybe a
Nothing -> forall (m :: * -> *) a. MonadFail m => [Char] -> m a
fail ([Char]
"Invalid postgres enumeration value: " forall a. [a] -> [a] -> [a]
++ [Char]
name)
      Just  a
v -> forall (f :: * -> *) a. Applicative f => a -> f a
pure a
v

beamTypeForCustomPg :: CheckedDatabaseEntity Postgres db (PgType a) -> DataType Postgres a
beamTypeForCustomPg :: forall (db :: (* -> *) -> *) a.
CheckedDatabaseEntity Postgres db (PgType a) -> DataType Postgres a
beamTypeForCustomPg (CheckedDatabaseEntity (CheckedPgTypeDescriptor (PgTypeDescriptor Maybe Text
_ Text
_ PgDataTypeSyntax
dt) [PgTypeCheck]
_) [SomeDatabasePredicate]
_)
    = forall be a. BeamSqlBackendCastTargetSyntax be -> DataType be a
DataType PgDataTypeSyntax
dt