TH.hs

{-# LANGUAGE CPP #-}
module Composite.Opaleye.TH where

import Control.Lens ((<&>))
import qualified Data.ByteString.Char8 as BSC8
import Data.List.Split (splitOn)
import Data.Maybe (fromMaybe)
import Data.Profunctor.Product.Default (Default, def)
import Data.Traversable (for)
import Database.PostgreSQL.Simple (ResultError(ConversionFailed, Incompatible, UnexpectedNull))
import Database.PostgreSQL.Simple.FromField (FromField, fromField, typename, returnError)
import Language.Haskell.TH
  ( Q, Name, mkName, nameBase, newName, pprint, reify
  , Info(TyConI), Dec(DataD), Con(NormalC)
  , conT
  , dataD, instanceD
  , lamE, varE, caseE, conE
  , conP, varP, wildP, litP, stringL
  , caseE, match
  , funD, clause
  , normalB, normalGE, guardedB
  , cxt
  )
import Language.Haskell.TH.Syntax (lift)
import Opaleye
  ( Column, DefaultFromField, ToFields, fromPGSFromField, defaultFromField
  )
import Opaleye.Internal.PGTypes (IsSqlType, showSqlType, literalColumn)
import Opaleye.Internal.HaskellDB.PrimQuery (Literal(StringLit))

getLastComponent :: String -> String
getLastComponent str = case reverse (splitOn "." str) of
  x:_ -> x
  [] -> str

-- |Derive the various instances required to make a Haskell enumeration map to a PostgreSQL @enum@ type.
--
-- In @deriveOpaleyeEnum ''HaskellType "schema.sqltype" hsConToSqlValue@, @''HaskellType@ is the sum type (data declaration) to make instances for, 
-- @"schema.sqltype"@ is the PostgreSQL type name, and @hsConToSqlValue@ is a function to map names of constructors to SQL values.
--
-- The function @hsConToSqlValue@ is of the type @String -> Maybe String@ in order to make using 'stripPrefix' convenient. The function is applied to each
-- constructor name and for @Just value@ that value is used, otherwise for @Nothing@ the constructor name is used.
--
-- For example, given the Haskell type:
--
-- @
--     data MyEnum = MyFoo | MyBar
-- @
--
-- And PostgreSQL type:
--
-- @
--     CREATE TYPE myenum AS ENUM('foo', 'bar');
-- @
--
-- The splice:
--
-- @
--     deriveOpaleyeEnum ''MyEnum "myschema.myenum" ('stripPrefix' "my" . 'map' 'toLower')
-- @
--
-- Will create @PGMyEnum@ and instances required to use @MyEnum@ / @Column MyEnum@ in Opaleye.
--
-- The Haskell generated by this splice for the example is something like:
--
-- @
--     data PGMyEnum
--
--     instance 'IsSqlType' PGMyEnum where
--       'showSqlType' _ = "myschema.myenum"
--
--     instance 'FromField' MyEnum where
--       'fromField' f mbs = do
--         tname <- 'typename' f
--         case mbs of
--           _ | 'getLastComponent' ('BSC8.unpack' tname) /= "myenum" -> 'returnError' 'Incompatible' f ""
--           Just "foo" -> pure MyFoo
--           Just "bar" -> pure MyBar
--           Just other -> 'returnError' 'ConversionFailed' f ("Unexpected myschema.myenum value: " <> 'BSC8.unpack' other)
--           Nothing    -> 'returnError' 'UnexpectedNull' f ""
--
--     instance 'DefaultFromField' PGMyEnum MyEnum where
--       defaultFromField = 'fromPGSFromField'
--
--     instance 'Default' 'ToFields' MyEnum ('Column' PGMyEnum) where
--       def = 'ToFields' $ \ a ->
--         'literalColumn' . 'stringLit' $ case a of
--           MyFoo -> "foo"
--           MyBar -> "bar"
-- @
deriveOpaleyeEnum :: Name -> String -> (String -> Maybe String) -> Q [Dec]
deriveOpaleyeEnum hsName sqlName hsConToSqlValue = do
  let sqlTypeName = mkName $ "PG" ++ nameBase hsName
      sqlType = conT sqlTypeName
      hsType = conT hsName
      unqualSqlName = getLastComponent sqlName

  rawCons <- reify hsName >>= \ case
    TyConI (DataD _cxt _name _tvVarBndrs _maybeKind cons _derivingCxt) ->
      pure cons
    other ->
      fail $ "expected " <> show hsName <> " to name a data declaration, not:\n" <> pprint other

  nullaryCons <- for rawCons $ \ case
    NormalC conName [] ->
      pure conName
    other ->
      fail $ "expected every constructor of " <> show hsName <> " to be a regular nullary constructor, not:\n" <> pprint other

  let conPairs = nullaryCons <&> \ conName ->
        (conName, fromMaybe (nameBase conName) (hsConToSqlValue (nameBase conName)))

  sqlTypeDecl <-
    dataD
      (cxt [])
      sqlTypeName
      []
      Nothing
      []
#if MIN_VERSION_template_haskell(2,12,0)
      []
#else
      (cxt [])
#endif

  isSqlTypeInst <- instanceD (cxt []) [t| IsSqlType $sqlType |] . (:[]) $ do
    funD 'showSqlType
      [ clause
          [wildP]
          (normalB (lift sqlName))
          []
      ]

  fromFieldInst <- instanceD (cxt []) [t| FromField $hsType |] . (:[]) $ do
    field <- newName "field"
    mbs   <- newName "mbs"
    tname <- newName "tname"
    other <- newName "other"

    let bodyCase = caseE (varE mbs) $
          [ match
              wildP
              (guardedB [ normalGE [| getLastComponent (BSC8.unpack $(varE tname)) /= $(lift unqualSqlName) |]
                                   [| returnError Incompatible $(varE field) "" |] ])
              []
          ] ++
          (
            conPairs <&> \ (conName, value) ->
              match
                [p| Just $(litP $ stringL value) |]
                (normalB [| pure $(conE conName) |])
                []
          ) ++
          [ match 
              [p| Just $(varP other) |]
              (normalB [| returnError ConversionFailed $(varE field) ("Unexpected " <> $(lift sqlName) <> " value: " <> BSC8.unpack $(varE other)) |])
              []
          , match
              [p| Nothing |]
              (normalB [| returnError UnexpectedNull $(varE field) "" |])
              []
          ]

    funD 'fromField
      [ clause
          [varP field, varP mbs]
          (normalB [|
            do
              $(varP tname) <- typename $(varE field)
              $bodyCase
            |])
          []
      ]

  defaultFromFieldInst <- instanceD (cxt []) [t| DefaultFromField $sqlType $hsType |] . (:[]) $
    funD 'defaultFromField
      [ clause
          []
          (normalB [| fromPGSFromField |])
          []
      ]

  defaultInst <- instanceD (cxt []) [t| Default ToFields $hsType (Column $sqlType) |] . (:[]) $ do
    s <- newName "s"
    let body = lamE [varP s] $
          caseE (varE s) $
            conPairs <&> \ (conName, value) ->
              match
                (conP conName [])
                (normalB $ lift value)
                []

    funD 'def
      [ clause
          []
          (normalB [| ToFields (literalColumn . StringLit . $body) |])
          []
      ]

  pure [sqlTypeDecl, isSqlTypeInst, fromFieldInst, defaultFromFieldInst, defaultInst]