{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TemplateHaskellQuotes #-}
{-# OPTIONS_HADDOCK not-home #-}

-- | Internal to the arbiter packages. Not covered by the PVP.
--
-- The @sql@ quasiquoter. It builds a 'Arbiter.Core.Sql.Query.Query' whose text, parameters, and row
-- decoder all come from one template. The @stmt@ quasiquoter renders the text once and makes each distinct
-- @#{...}@ identifier a function argument.
--
-- 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 t'Arbiter.Core.Sql.Query.ToFragment'.
--   * @#{ident :: CInt8}@ emits one parameter hole and binds in-scope @ident@ as
--     its value. @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 ()@ with no @\@{...}@ holes, @Query@ of the bare value with
--     one, and @Query@ of a tuple with 2 to 'maxOutputHoles'.
--
-- A @?@ in the template is literal SQL. Do not use it in a template with holes.
-- The @?@-placeholder form cannot tell it from a hole.
module Arbiter.Core.Sql.QQ
  ( sql
  , stmt
  , maxOutputHoles
  ) where

import Data.Containers.ListUtils (nubOrd)
import Data.List (dropWhileEnd)
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 (Piece (..), Query (..), mkQuery, toFragment)

-- | The @sql@ quasiquoter, valid in expression position.
sql :: QuasiQuoter
sql :: QuasiQuoter
sql = ([Token] -> Q Exp) -> QuasiQuoter
quoter [Token] -> Q Exp
compile

-- | Like 'sql', but each distinct @#{ident :: coltype}@ hole becomes a parameter of
-- the result, in order of first appearance. The text is rendered once.
stmt :: QuasiQuoter
stmt :: QuasiQuoter
stmt = ([Token] -> Q Exp) -> QuasiQuoter
quoter [Token] -> Q Exp
compileStmt

quoter :: ([Token] -> Q Exp) -> QuasiQuoter
quoter :: ([Token] -> Q Exp) -> QuasiQuoter
quoter [Token] -> Q Exp
codegen =
  QuasiQuoter
    { quoteExp :: String -> Q Exp
quoteExp = \String
template -> (String -> Q Exp)
-> ([Token] -> Q Exp) -> Either String [Token] -> 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 [Token] -> Q Exp
codegen (String -> Either String [Token]
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 Token
  = Raw 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 tokens.
tokenize :: String -> Either String [Token]
tokenize :: String -> Either String [Token]
tokenize = String -> String -> Either String [Token]
go String
""
  where
    go :: String -> String -> Either String [Token]
go String
buf String
input = case String
input of
      [] -> [Token] -> Either String [Token]
forall a b. b -> Either a b
Right (String -> [Token]
flush String
buf)
      (Char
'$' : Char
'{' : String
rest) -> Char -> String -> String -> Either String [Token]
hole Char
'$' String
buf String
rest
      (Char
'#' : Char
'{' : String
rest) -> Char -> String -> String -> Either String [Token]
hole Char
'#' String
buf String
rest
      (Char
'@' : Char
'{' : String
rest) -> Char -> String -> String -> Either String [Token]
hole Char
'@' String
buf String
rest
      (Char
ch : String
rest) -> String -> String -> Either String [Token]
go (Char
ch Char -> String -> String
forall a. a -> [a] -> [a]
: String
buf) String
rest

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

    flush :: String -> [Token]
flush String
buf = [Text -> Token
Raw (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 '}')"

mkToken :: Char -> String -> Either String Token
mkToken :: Char -> String -> Either String Token
mkToken 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 Token
forall a b. a -> Either a b
Left String
"sql: empty ${} splice" else Token -> Either String Token
forall a b. b -> Either a b
Right (Text -> Token
Splice Text
name)
mkToken Char
'#' String
inside = do
  (ident, colType) <- String -> Either String (Text, Text)
splitAnn String
inside
  (kind, colName) <- parseColType colType
  Right (InHole ident kind colName)
mkToken Char
'@' String
inside = do
  (name, colType) <- String -> Either String (Text, Text)
splitAnn String
inside
  (kind, colName) <- parseColType colType
  case kind of
    Kind
KScalar -> Token -> Either String Token
forall a b. b -> Either a b
Right (Text -> Bool -> Text -> Token
OutHole Text
name Bool
False Text
colName)
    Kind
KNullable -> Token -> Either String Token
forall a b. b -> Either a b
Right (Text -> Bool -> Text -> Token
OutHole Text
name Bool
True Text
colName)
    Kind
_ -> String -> Either String Token
forall a b. a -> Either a b
Left String
"sql: @{} output holes cannot be array-typed"
mkToken Char
sigil String
_ = String -> Either String Token
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
token
      T.strip <$> T.stripSuffix "]" 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
token

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

-- | The 'sql' form.
compile :: [Token] -> Q Exp
compile :: [Token] -> Q Exp
compile [Token]
tokens = [|mkQuery $([Token] -> Q Exp
piecesE [Token]
tokens) $([Token] -> Q Exp
paramsE [Token]
tokens) $([Token] -> Q Exp
decoderE [Token]
tokens)|]

-- | The 'stmt' form. The staged query is shared by every application.
compileStmt :: [Token] -> Q Exp
compileStmt :: [Token] -> Q Exp
compileStmt [Token]
tokens = do
  stagedN <- String -> Q Name
forall (m :: * -> *). Quote m => String -> m Name
TH.newName String
"staged"
  let binders = (Text -> Pat) -> [Text] -> [Pat]
forall a b. (a -> b) -> [a] -> [b]
map (Name -> Pat
TH.VarP (Name -> Pat) -> (Text -> Name) -> Text -> Pat
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Name
TH.mkName (String -> Name) -> (Text -> String) -> Text -> Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> String
T.unpack) ([Text] -> [Text]
forall a. Ord a => [a] -> [a]
nubOrd [Text
ident | InHole Text
ident Kind
_ Text
_ <- [Token]
tokens])
  stagedE <- [|mkQuery $(piecesE tokens) [] $(decoderE tokens)|]
  body <- [|$(TH.varE stagedN) {qParams = $(paramsE tokens)}|]
  pure (TH.LetE [TH.ValD (TH.VarP stagedN) (TH.NormalB stagedE) []] (TH.LamE binders body))

piecesE :: [Token] -> Q Exp
piecesE :: [Token] -> Q Exp
piecesE [Token]
tokens = [|concat $([Q Exp] -> Q Exp
forall (m :: * -> *). Quote m => [m Exp] -> m Exp
TH.listE ((Token -> Q Exp) -> [Token] -> [Q Exp]
forall a b. (a -> b) -> [a] -> [b]
map Token -> Q Exp
piecesExp [Token]
tokens))|]

paramsE :: [Token] -> Q Exp
paramsE :: [Token] -> Q Exp
paramsE [Token]
tokens = [|concat $([Q Exp] -> Q Exp
forall (m :: * -> *). Quote m => [m Exp] -> m Exp
TH.listE ((Token -> Q Exp) -> [Token] -> [Q Exp]
forall a b. (a -> b) -> [a] -> [b]
map Token -> Q Exp
paramsExp [Token]
tokens))|]

decoderE :: [Token] -> Q Exp
decoderE :: [Token] -> Q Exp
decoderE [Token]
tokens = [(Text, Bool, Text)] -> Q Exp
mkDecoder [(Text
name, Bool
nullable, Text
colType) | OutHole Text
name Bool
nullable Text
colType <- [Token]
tokens]

-- | A token's contribution to the query's pieces.
piecesExp :: Token -> Q Exp
piecesExp :: Token -> Q Exp
piecesExp (Raw Text
literal) = [|[Lit (T.pack $(String -> Q Exp
forall (m :: * -> *). Quote m => String -> m Exp
TH.stringE (Text -> String
T.unpack Text
literal)))]|]
piecesExp (Splice Text
name) = [|qPieces (toFragment $(Name -> Q Exp
forall (m :: * -> *). Quote m => Name -> m Exp
TH.varE (String -> Name
TH.mkName (Text -> String
T.unpack Text
name))))|]
piecesExp (OutHole Text
name Bool
_ Text
_) = [|[Lit (T.pack $(String -> Q Exp
forall (m :: * -> *). Quote m => String -> m Exp
TH.stringE (Text -> String
T.unpack Text
name)))]|]
piecesExp (InHole Text
_ Kind
_ Text
_) = [|[Hole]|]

-- | A token's contribution to the parameters.
paramsExp :: Token -> Q Exp
paramsExp :: Token -> Q Exp
paramsExp (Splice Text
name) = [|qParams (toFragment $(Name -> Q Exp
forall (m :: * -> *). Quote m => Name -> m Exp
TH.varE (String -> Name
TH.mkName (Text -> String
T.unpack Text
name))))|]
paramsExp (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
  [|[$(TH.varE encoder) $(TH.conE colN) $(TH.varE (TH.mkName (T.unpack ident)))]|]
paramsExp Token
_ = [|[]|]

-- | 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
maxOutputHoles =
      String -> Q Exp
forall a. HasCallStack => String -> Q a
forall (m :: * -> *) a.
(MonadFail m, HasCallStack) =>
String -> m a
fail (String
"sql: more than " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show Int
maxOutputHoles String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
" @{} 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
      [] -> [|pure ()|]
      [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 most @\@{...}@ holes one quote decodes as a tuple.
maxOutputHoles :: Int
maxOutputHoles :: Int
maxOutputHoles = Int
8

-- | 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 -> Bool) -> [Text] -> [Text]
forall a. (a -> Bool) -> [a] -> [a]
dropWhileEnd Text -> Bool
isBlank