Quasi-quoter for postgresql-simple

HereQ.hs

{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE FlexibleContexts #-}
{-# OPTIONS_GHC -fno-warn-missing-fields #-}

-- | Interpolate any type that inhabits both IsString and Semigroup
-- This code is based on the "here" package. It would be nice to strip
-- leading whitespace, as the "neat-interpolation" package does.
-- Additionally, "neat-interpolation" has a simpler interpolation
-- placeholder style which avoids treating "\" specially
module HereQ
  ( hereQ
  ) where

import Control.Monad.State
import Data.Char
import Language.Haskell.Meta
import Language.Haskell.TH
import Language.Haskell.TH.Quote
import Text.Parsec
import Text.Parsec.String

data StringPart
  = Lit String
  | Esc Char
  | Anti (Q Exp)

data HsChompState = HsChompState
  { quoteState :: QuoteState
  , braceCt :: Int
  , consumed :: String
  , prevCharWasIdentChar :: Bool
  }

data QuoteState
  = None
  | Single EscapeState
  | Double EscapeState

data EscapeState
  = Escaped
  | Unescaped

-- | Quote a here doc with embedded antiquoted expressions.
--
-- The result type must have 'IsString' and 'Semigroup' instances, and
-- any expression of the same type occurring between @${@ and @}@
-- instances will be interpolated into the quoted result using
-- semigroup concatenation.
--
-- Characters preceded by a backslash are treated literally. This enables the
-- inclusion of the literal substring @${@ within your quoted text by writing
-- it as @\\${@. The literal sequence @\\${@ may be written as @\\\\${@.
hereQ :: QuasiQuoter
hereQ = QuasiQuoter {quoteExp = quoteInterp}

quoteInterp :: String -> Q Exp
quoteInterp s = either (handleError s) combineParts (parseInterp s)

handleError :: String -> ParseError -> Q Exp
handleError expStr parseError =
  error $
  "Failed to parse interpolated expression in string: " ++
  expStr ++ "\n" ++ show parseError

combineParts :: [StringPart] -> Q Exp
combineParts = combine . map toExpQ
  where
    toExpQ (Lit s) = stringE s
    toExpQ (Esc c) = stringE [c]
    toExpQ (Anti expq) = expq
    combine [] = stringE ""
    combine parts = foldr1 (\subExpr acc -> [|$subExpr <> $acc|]) parts

parseInterp :: String -> Either ParseError [StringPart]
parseInterp = parse pInterp ""

pInterp :: Parser [StringPart]
pInterp = manyTill pStringPart eof

pStringPart :: Parser StringPart
pStringPart = pAnti <|> pEsc <|> pLit

pAnti :: Parser StringPart
pAnti = Anti <$> between (try pAntiOpen) pAntiClose pAntiExpr

pAntiOpen :: Parser String
pAntiOpen = string "${"

pAntiClose :: Parser String
pAntiClose = string "}"

pAntiExpr :: Parser (Q Exp)
pAntiExpr =
  pUntilUnbalancedCloseBrace >>= either fail (return . return) . parseExp

pUntilUnbalancedCloseBrace :: Parser String
pUntilUnbalancedCloseBrace = evalStateT go $ HsChompState None 0 "" False
  where
    go = do
      c <- lift anyChar
      modify $ \st@HsChompState {consumed} -> st {consumed = c : consumed}
      HsChompState {..} <- get
      let next = setIdentifierCharState c >> go
      case quoteState of
        None ->
          case c of
            '{' -> incBraceCt 1 >> next
            '}'
              | braceCt > 0 -> incBraceCt (-1) >> next
              | otherwise -> stepBack >> return (reverse $ tail consumed)
            '\'' ->
              unless prevCharWasIdentChar (setQuoteState $ Single Unescaped) >>
              next
            '"' -> setQuoteState (Double Unescaped) >> next
            _ -> next
        Single Unescaped -> do
          case c of
            '\\' -> setQuoteState (Single Escaped)
            '\'' -> setQuoteState None
            _ -> return ()
          next
        Single Escaped -> setQuoteState (Single Unescaped) >> next
        Double Unescaped -> do
          case c of
            '\\' -> setQuoteState (Double Escaped)
            '"' -> setQuoteState None
            _ -> return ()
          next
        Double Escaped -> setQuoteState (Double Unescaped) >> next
    stepBack =
      lift $
      updateParserState (\s -> s {statePos = incSourceColumn (statePos s) (-1)}) >>
      getInput >>=
      setInput . ('}' :)
    incBraceCt n =
      modify $ \st@HsChompState {braceCt} -> st {braceCt = braceCt + n}
    setQuoteState qs = modify $ \st -> st {quoteState = qs}
    setIdentifierCharState c =
      modify $ \st ->
        st
          { prevCharWasIdentChar =
              or [isLetter c, isDigit c, c == '_', c == '\'']
          }

pEsc :: Parser StringPart
pEsc = Esc <$> (char '\\' *> anyChar)

pLit :: Parser StringPart
pLit =
  fmap Lit $
  try (litCharTil $ try $ lookAhead pAntiOpen <|> lookAhead (string "\\")) <|>
  litCharTil eof
  where
    litCharTil = manyTill $ noneOf ['\\']

QFrag.hs

{-# LANGUAGE OverloadedStrings #-}

-- | You can use this to build SQL strings piece by piece, where each
-- piece may embed one or more "?" placeholders and therefore be
-- associated with a matching number of values to be interpolated at
-- those positions later. Use the Monoid instance to concatenate
-- fragments.
module QFrag
  ( QFrag
  , frag
  , frag1
  , frag2
  , frag3
  , frag4
  , fragN
  , param
  , unpack
  , applyQ
  ) where

import Data.Semigroup (Semigroup, (<>))
import Data.String (IsString, fromString)

import Data.Text (Text)
import qualified Data.Text as Text
import Database.PostgreSQL.Simple (Only(..), Query)
import Database.PostgreSQL.Simple.ToField
import Database.PostgreSQL.Simple.ToRow

data QFrag =
  QFrag Text
        [Action]

instance Semigroup QFrag where
  (QFrag s ps) <> (QFrag s' ps') =
    QFrag (s <> s') (ps <> ps')

instance Monoid QFrag where
  mappend = (<>)
  mempty = QFrag mempty mempty

instance IsString QFrag where
  fromString = frag . Text.pack

param :: ToField a => a -> QFrag
param = frag1 "?"

frag :: Text -> QFrag
frag s = QFrag s mempty

frag1 :: ToField a => Text -> a -> QFrag
frag1 s a = fragN s (Only a)

frag2 :: (ToField a, ToField b) => Text -> a -> b -> QFrag
frag2 s a b = fragN s (a, b)

frag3 :: (ToField a, ToField b, ToField c) => Text -> a -> b -> c -> QFrag
frag3 s a b c = fragN s (a, b, c)

frag4 ::
     (ToField a, ToField b, ToField c, ToField d)
  => Text
  -> a
  -> b
  -> c
  -> d
  -> QFrag
frag4 s a b c d = fragN s (a, b, c, d)

-- | If you pass a bunch of params at the same time, they have to be
-- of the same type, so for heterogenous groups of params, you might
-- want to use toAction directly in order to coerce them to a matching
-- type.
fragN :: ToRow a => Text -> a -> QFrag
fragN s params = QFrag s (toRow params)

unpack :: QFrag -> ([Action], Query)
unpack (QFrag s params) = (params, (fromString . Text.unpack) s)

-- | Convert a Database.Postgresql.Transaction function which takes
-- params and a query into one which takes a QFrag.
applyQ :: ([Action] -> Query -> a) -> QFrag -> a
applyQ f = uncurry f . unpack