{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TemplateHaskellQuotes #-}

-- | The @sql@ quasiquoter. It builds a 'Arbiter.Core.Sql.Query.Query' whose text, parameters, and row
-- decoder all come from one template.
--
-- Holes reference in-scope identifiers, like @NeatInterpolation@'s @${var}@:
--
--   * @${x}@ splices a fragment: 'Text' (raw clause or table name) or a
--     @Query ()@ (its parameters interleave at the splice site), via @ToFragment@.
--   * @#{ident :: CInt8}@ emits one @?@ and binds in-scope @ident@ as a
--     parameter. @Maybe CInt8@, @[CInt8]@, and @[Maybe CInt8]@ pick the
--     nullable, array, and nullable-array encoders.
--   * @\@{name :: CInt8}@ emits the identifier @name@ and adds @col \"name\"
--     CInt8@ to the decoder (@Maybe CInt8@ uses @ncol@). The quote's result type
--     is @Query@ of the tuple of these holes, or @Query ()@ when there are none.
--
-- A bare @?@ is rejected. Every placeholder comes from a hole. Use @jsonb_exists@
-- and friends for jsonb key-existence.
module Arbiter.Core.Sql.QQ
  ( sql
  ) where

import Data.Text (Text)
import Data.Text qualified as T
import Language.Haskell.TH (Exp, Name, Q)
import Language.Haskell.TH qualified as TH
import Language.Haskell.TH.Quote (QuasiQuoter (..))

import Arbiter.Core.Codec
  ( Col (..)
  , col
  , ncol
  , parr
  , pnarr
  , pnul
  , pval
  )
import Arbiter.Core.Sql.Query (param, raw, rows, toFragment)

-- | The @sql@ quasiquoter, valid in expression position.
sql :: QuasiQuoter
sql :: QuasiQuoter
sql =
  QuasiQuoter
    { quoteExp :: String -> Q Exp
quoteExp = \String
template -> (String -> Q Exp)
-> ([Piece] -> Q Exp) -> Either String [Piece] -> Q Exp
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either String -> Q Exp
forall a. HasCallStack => String -> Q a
forall (m :: * -> *) a.
(MonadFail m, HasCallStack) =>
String -> m a
fail [Piece] -> Q Exp
compile (String -> Either String [Piece]
tokenize (Text -> String
T.unpack (Text -> Text
normalizeIndent (String -> Text
T.pack String
template))))
    , quotePat :: String -> Q Pat
quotePat = String -> Q Pat
forall {m :: * -> *} {p} {a}. MonadFail m => p -> m a
badContext
    , quoteType :: String -> Q Type
quoteType = String -> Q Type
forall {m :: * -> *} {p} {a}. MonadFail m => p -> m a
badContext
    , quoteDec :: String -> Q [Dec]
quoteDec = String -> Q [Dec]
forall {m :: * -> *} {p} {a}. MonadFail m => p -> m a
badContext
    }
  where
    badContext :: p -> m a
badContext p
_ = String -> m a
forall a. HasCallStack => String -> m a
forall (m :: * -> *) a.
(MonadFail m, HasCallStack) =>
String -> m a
fail String
"sql: only valid in expression position"

-- ---------------------------------------------------------------------------
-- Parsing
-- ---------------------------------------------------------------------------

-- | How a parameter column is shaped.
data Kind = KScalar | KNullable | KArray | KNullArray

data Piece
  = Lit Text
  | -- | @${ident}@
    Splice Text
  | -- | @#{ident :: coltype}@: identifier, shape, column constructor.
    InHole Text Kind Text
  | -- | @\@{name :: coltype}@: name, nullable, column constructor.
    OutHole Text Bool Text

-- | Scan a template into pieces. A bare @?@ is an error.
tokenize :: String -> Either String [Piece]
tokenize :: String -> Either String [Piece]
tokenize = String -> String -> Either String [Piece]
go String
""
  where
    go :: String -> String -> Either String [Piece]
go String
buf String
input = case String
input of
      [] -> [Piece] -> Either String [Piece]
forall a b. b -> Either a b
Right (String -> [Piece]
flush String
buf)
      (Char
'$' : Char
'{' : String
rest) -> Char -> String -> String -> Either String [Piece]
hole Char
'$' String
buf String
rest
      (Char
'#' : Char
'{' : String
rest) -> Char -> String -> String -> Either String [Piece]
hole Char
'#' String
buf String
rest
      (Char
'@' : Char
'{' : String
rest) -> Char -> String -> String -> Either String [Piece]
hole Char
'@' String
buf String
rest
      (Char
'?' : String
_) ->
        String -> Either String [Piece]
forall a b. a -> Either a b
Left
          String
"sql: bare '?' placeholder. Use a #{} hole. \
          \For jsonb key-existence use jsonb_exists/jsonb_exists_any/jsonb_exists_all"
      (Char
ch : String
rest) -> String -> String -> Either String [Piece]
go (Char
ch Char -> String -> String
forall a. a -> [a] -> [a]
: String
buf) String
rest

    hole :: Char -> String -> String -> Either String [Piece]
hole Char
sig String
buf String
rest = do
      (inside, rest') <- String -> Either String (String, String)
takeBrace String
rest
      piece <- mkPiece sig inside
      rest'' <- go "" rest'
      Right (flush buf ++ [piece] ++ rest'')

    flush :: String -> [Piece]
flush String
buf = [Text -> Piece
Lit (String -> Text
T.pack (String -> String
forall a. [a] -> [a]
reverse String
buf)) | Bool -> Bool
not (String -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null String
buf)]

takeBrace :: String -> Either String (String, String)
takeBrace :: String -> Either String (String, String)
takeBrace String
input = case (Char -> Bool) -> String -> (String, String)
forall a. (a -> Bool) -> [a] -> ([a], [a])
break (Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
'}') String
input of
  (String
inside, Char
'}' : String
rest) -> (String, String) -> Either String (String, String)
forall a b. b -> Either a b
Right (String
inside, String
rest)
  (String, String)
_ -> String -> Either String (String, String)
forall a b. a -> Either a b
Left String
"sql: unterminated hole (missing '}')"

mkPiece :: Char -> String -> Either String Piece
mkPiece :: Char -> String -> Either String Piece
mkPiece Char
'$' String
inside =
  let name :: Text
name = Text -> Text
T.strip (String -> Text
T.pack String
inside)
   in if Text -> Bool
T.null Text
name then String -> Either String Piece
forall a b. a -> Either a b
Left String
"sql: empty ${} splice" else Piece -> Either String Piece
forall a b. b -> Either a b
Right (Text -> Piece
Splice Text
name)
mkPiece Char
'#' String
inside = do
  (ident, colType) <- String -> Either String (Text, Text)
splitAnn String
inside
  (kind, colName) <- parseColType colType
  Right (InHole ident kind colName)
mkPiece Char
'@' String
inside = do
  (name, colType) <- String -> Either String (Text, Text)
splitAnn String
inside
  (kind, colName) <- parseColType colType
  case kind of
    Kind
KScalar -> Piece -> Either String Piece
forall a b. b -> Either a b
Right (Text -> Bool -> Text -> Piece
OutHole Text
name Bool
False Text
colName)
    Kind
KNullable -> Piece -> Either String Piece
forall a b. b -> Either a b
Right (Text -> Bool -> Text -> Piece
OutHole Text
name Bool
True Text
colName)
    Kind
_ -> String -> Either String Piece
forall a b. a -> Either a b
Left String
"sql: @{} output holes cannot be array-typed"
mkPiece Char
sigil String
_ = String -> Either String Piece
forall a b. a -> Either a b
Left (String
"sql: unknown hole sigil " String -> String -> String
forall a. [a] -> [a] -> [a]
++ [Char
sigil])

-- | Split @expr :: coltype@ on the @::@.
splitAnn :: String -> Either String (Text, Text)
splitAnn :: String -> Either String (Text, Text)
splitAnn String
annotation = case HasCallStack => Text -> Text -> (Text, Text)
Text -> Text -> (Text, Text)
T.breakOn Text
"::" (String -> Text
T.pack String
annotation) of
  (Text
lhs, Text
rhs)
    | Text -> Bool
T.null Text
rhs -> String -> Either String (Text, Text)
forall a b. a -> Either a b
Left (String
"sql: hole needs a ':: ColType' annotation: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
annotation)
    | Bool
otherwise -> (Text, Text) -> Either String (Text, Text)
forall a b. b -> Either a b
Right (Text -> Text
T.strip Text
lhs, Text -> Text
T.strip (Int -> Text -> Text
T.drop Int
2 Text
rhs))

-- | Parse a column-type annotation into its shape and 'Col' constructor name.
parseColType :: Text -> Either String (Kind, Text)
parseColType :: Text -> Either String (Kind, Text)
parseColType Text
rawType =
  let trimmed :: Text
trimmed = Text -> Text
T.strip Text
rawType
   in case Text -> Maybe Text
bracketed Text
trimmed of
        Just Text
inner -> case Text -> Maybe Text
maybePrefixed Text
inner of
          Just Text
conName -> (Kind, Text) -> Either String (Kind, Text)
forall a b. b -> Either a b
Right (Kind
KNullArray, Text
conName)
          Maybe Text
Nothing -> (Kind, Text) -> Either String (Kind, Text)
forall a b. b -> Either a b
Right (Kind
KArray, Text
inner)
        Maybe Text
Nothing -> case Text -> Maybe Text
maybePrefixed Text
trimmed of
          Just Text
conName -> (Kind, Text) -> Either String (Kind, Text)
forall a b. b -> Either a b
Right (Kind
KNullable, Text
conName)
          Maybe Text
Nothing -> (Kind, Text) -> Either String (Kind, Text)
forall a b. b -> Either a b
Right (Kind
KScalar, Text
trimmed)
  where
    bracketed :: Text -> Maybe Text
bracketed Text
token = do
      afterOpen <- Text -> Text -> Maybe Text
T.stripPrefix Text
"[" (Text -> Text
T.strip Text
token)
      T.strip <$> T.stripSuffix "]" (T.strip afterOpen)
    maybePrefixed :: Text -> Maybe Text
maybePrefixed Text
token = Text -> Text
T.strip (Text -> Text) -> Maybe Text -> Maybe Text
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Text -> Text -> Maybe Text
T.stripPrefix Text
"Maybe " (Text -> Text
T.strip Text
token)

-- ---------------------------------------------------------------------------
-- Codegen
-- ---------------------------------------------------------------------------

compile :: [Piece] -> Q Exp
compile :: [Piece] -> Q Exp
compile [Piece]
pieces = do
  let textExpr :: Q Exp
textExpr = [|mconcat $([Q Exp] -> Q Exp
forall (m :: * -> *). Quote m => [m Exp] -> m Exp
TH.listE ((Piece -> Q Exp) -> [Piece] -> [Q Exp]
forall a b. (a -> b) -> [a] -> [b]
map Piece -> Q Exp
pieceExp [Piece]
pieces))|]
      outHoles :: [(Text, Bool, Text)]
outHoles = [(Text
name, Bool
nullable, Text
colType) | OutHole Text
name Bool
nullable Text
colType <- [Piece]
pieces]
  case [(Text, Bool, Text)]
outHoles of
    [] -> Q Exp
textExpr
    [(Text, Bool, Text)]
_ -> do
      dec <- [(Text, Bool, Text)] -> Q Exp
mkDecoder [(Text, Bool, Text)]
outHoles
      [|rows $(pure dec) $textExpr|]

-- | A single piece as an expression of type @Query ()@.
pieceExp :: Piece -> Q Exp
pieceExp :: Piece -> Q Exp
pieceExp (Lit Text
literal) = [|raw (T.pack $(String -> Q Exp
forall (m :: * -> *). Quote m => String -> m Exp
TH.stringE (Text -> String
T.unpack Text
literal)))|]
pieceExp (Splice Text
name) = [|toFragment $(Name -> Q Exp
forall (m :: * -> *). Quote m => Name -> m Exp
TH.varE (String -> Name
TH.mkName (Text -> String
T.unpack Text
name)))|]
pieceExp (OutHole Text
name Bool
_ Text
_) = [|raw (T.pack $(String -> Q Exp
forall (m :: * -> *). Quote m => String -> m Exp
TH.stringE (Text -> String
T.unpack Text
name)))|]
pieceExp (InHole Text
ident Kind
kind Text
colType) = do
  colN <- Text -> Q Name
colConName Text
colType
  let encoder = case Kind
kind of
        Kind
KScalar -> 'pval
        Kind
KNullable -> 'pnul
        Kind
KArray -> 'parr
        Kind
KNullArray -> 'pnarr
  [|param ($(TH.varE encoder) $(TH.conE colN) $(TH.varE (TH.mkName (T.unpack ident))))|]

-- | Applicative decoder for the output holes, as a tuple for arity >= 2.
mkDecoder :: [(Text, Bool, Text)] -> Q Exp
mkDecoder :: [(Text, Bool, Text)] -> Q Exp
mkDecoder [(Text, Bool, Text)]
holes
  | [(Text, Bool, Text)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(Text, Bool, Text)]
holes Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
8 =
      String -> Q Exp
forall a. HasCallStack => String -> Q a
forall (m :: * -> *) a.
(MonadFail m, HasCallStack) =>
String -> m a
fail String
"sql: more than 8 @{} output holes. Attach a handwritten codec with `rows` instead"
  | Bool
otherwise = case ((Text, Bool, Text) -> Q Exp) -> [(Text, Bool, Text)] -> [Q Exp]
forall a b. (a -> b) -> [a] -> [b]
map (Text, Bool, Text) -> Q Exp
colDecoder [(Text, Bool, Text)]
holes of
      [] -> String -> Q Exp
forall a. HasCallStack => String -> Q a
forall (m :: * -> *) a.
(MonadFail m, HasCallStack) =>
String -> m a
fail String
"sql: mkDecoder called with no holes"
      [Q Exp
single] -> Q Exp
single
      [Q Exp]
decoders ->
        (Q Exp -> Q Exp -> Q Exp) -> Q Exp -> [Q Exp] -> Q Exp
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\Q Exp
acc Q Exp
decoder -> [|$Q Exp
acc <*> $Q Exp
decoder|]) [|pure $(Name -> Q Exp
forall (m :: * -> *). Quote m => Name -> m Exp
TH.conE (Int -> Name
TH.tupleDataName ([Q Exp] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Q Exp]
decoders)))|] [Q Exp]
decoders
  where
    colDecoder :: (Text, Bool, Text) -> Q Exp
colDecoder (Text
name, Bool
nullable, Text
colType) = do
      colN <- Text -> Q Name
colConName Text
colType
      let dec = if Bool
nullable then 'ncol else 'col
      [|$(TH.varE dec) (T.pack $(TH.stringE (T.unpack name))) $(TH.conE colN)|]

-- | The 'Col' constructor 'Name' for a column-type token.
colConName :: Text -> Q Name
colConName :: Text -> Q Name
colConName Text
token = case Text
token of
  Text
"CInt4" -> Name -> Q Name
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure 'CInt4
  Text
"CInt8" -> Name -> Q Name
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure 'CInt8
  Text
"CText" -> Name -> Q Name
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure 'CText
  Text
"CBool" -> Name -> Q Name
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure 'CBool
  Text
"CTimestamptz" -> Name -> Q Name
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure 'CTimestamptz
  Text
"CJsonb" -> Name -> Q Name
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure 'CJsonb
  Text
"CFloat8" -> Name -> Q Name
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure 'CFloat8
  Text
"CUuid" -> Name -> Q Name
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure 'CUuid
  Text
other -> String -> Q Name
forall a. HasCallStack => String -> Q a
forall (m :: * -> *) a.
(MonadFail m, HasCallStack) =>
String -> m a
fail (String
"sql: unknown column type '" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Text -> String
T.unpack Text
other String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
"'")

-- ---------------------------------------------------------------------------
-- Indentation
-- ---------------------------------------------------------------------------

-- | Strip the common leading indentation and surrounding blank lines, matching
-- how @NeatInterpolation@ normalizes a multi-line quote.
normalizeIndent :: Text -> Text
normalizeIndent :: Text -> Text
normalizeIndent Text
template =
  let templateLines :: [Text]
templateLines = [Text] -> [Text]
dropTrailingBlank ((Text -> Bool) -> [Text] -> [Text]
forall a. (a -> Bool) -> [a] -> [a]
dropWhile Text -> Bool
isBlank (Text -> [Text]
T.lines Text
template))
      indent :: Int
indent = [Int] -> Int
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
minimum (Int
forall a. Bounded a => a
maxBound Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: [Text -> Int
leading Text
line | Text
line <- [Text]
templateLines, Bool -> Bool
not (Text -> Bool
isBlank Text
line)])
   in Text -> [Text] -> Text
T.intercalate Text
"\n" ((Text -> Text) -> [Text] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> Text -> Text
T.drop Int
indent) [Text]
templateLines)
  where
    isBlank :: Text -> Bool
isBlank = Text -> Bool
T.null (Text -> Bool) -> (Text -> Text) -> Text -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> Text
T.strip
    leading :: Text -> Int
leading = Text -> Int
T.length (Text -> Int) -> (Text -> Text) -> Text -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Char -> Bool) -> Text -> Text
T.takeWhile (Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
' ')
    dropTrailingBlank :: [Text] -> [Text]
dropTrailingBlank = [Text] -> [Text]
forall a. [a] -> [a]
reverse ([Text] -> [Text]) -> ([Text] -> [Text]) -> [Text] -> [Text]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Text -> Bool) -> [Text] -> [Text]
forall a. (a -> Bool) -> [a] -> [a]
dropWhile Text -> Bool
isBlank ([Text] -> [Text]) -> ([Text] -> [Text]) -> [Text] -> [Text]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Text] -> [Text]
forall a. [a] -> [a]
reverse