{-# LANGUAGE DuplicateRecordFields    #-}
{-# LANGUAGE NamedFieldPuns           #-}
{-# LANGUAGE TemplateHaskell          #-}
module Preql.QuasiQuoter.Syntax.TH where

import Preql.Imports
import Preql.QuasiQuoter.Common
import Preql.QuasiQuoter.Syntax.Params
import Preql.QuasiQuoter.Syntax.Parser (parseStatement, parseSelect)
import Preql.QuasiQuoter.Syntax.Printer (formatAsByteString)
import Preql.QuasiQuoter.Syntax.Syntax as Syntax hiding (select)
import Preql.Wire.Internal as Wire (Query(..))

import Language.Haskell.TH
import Language.Haskell.TH.Quote

import qualified Data.Text as T

tupleType :: [Name] -> Type
tupleType :: [Name] -> Type
tupleType [Name
v] = Name -> Type
VarT Name
v
tupleType [Name]
names = (Type -> Name -> Type) -> Type -> [Name] -> Type
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (\Type
expr Name
v -> Type -> Type -> Type
AppT Type
expr (Name -> Type
VarT Name
v)) (Int -> Type
TupleT Int
n) [Name]
names
    where n :: Int
n = [Name] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Name]
names

-- | Synthesize a Query tagged with the number of returned columns.
makeArityQuery :: Statement -> Q Exp
makeArityQuery :: Statement -> Q Exp
makeArityQuery Statement
parsed = do
  let
    width :: Q Type
width = case Statement -> Maybe Int
countColumnsReturned Statement
parsed of
              Just Int
n -> Type -> Q Type
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TyLit -> Type
LitT (Integer -> TyLit
NumTyLit (Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n)))
              Maybe Int
Nothing -> Name -> Type
VarT (Name -> Type) -> Q Name -> Q Type
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> Q Name
newName String
"r" -- SELECT *
    formatted :: ByteString
formatted = Statement -> ByteString
forall a. FormatSql a => a -> ByteString
formatAsByteString Statement
parsed
  [e| Wire.Query formatted :: Wire.Query $(width) |]

-- | This quasiquoter will accept most syntactically valid SELECT
-- queries.  Language features not yet implemented include type casts,
-- lateral joins, EXTRACT, INTO, string & XML operators, and
-- user-defined operators.  For now, please fall back to
-- 'Preql.QuasiQuoter.Raw.TH.sql' for these less-frequently used SQL
-- features, or file a bug report if a commonly used feature is not
-- parsed correctly.
--
-- @select@ accepts antiquotes with the same syntax as 'sql'.
select :: QuasiQuoter
select :: QuasiQuoter
select = String -> (String -> Q Exp) -> QuasiQuoter
expressionOnly String
"select" ((String -> String -> Either String SelectStmt)
-> (SelectStmt -> Statement) -> String -> Q Exp
forall a.
(String -> String -> Either String a)
-> (a -> Statement) -> String -> Q Exp
aritySql String -> String -> Either String SelectStmt
parseSelect SelectStmt -> Statement
QS)

-- | This quasiquoter will accept all queries accepted by 'select',
-- and limited INSERT, UPDATE, and DELETE queries.  For details of
-- what can be parsed, consult Parser.y
validSql :: QuasiQuoter
validSql :: QuasiQuoter
validSql = String -> (String -> Q Exp) -> QuasiQuoter
expressionOnly String
"validSql" ((String -> String -> Either String Statement)
-> (Statement -> Statement) -> String -> Q Exp
forall a.
(String -> String -> Either String a)
-> (a -> Statement) -> String -> Q Exp
aritySql String -> String -> Either String Statement
parseStatement Statement -> Statement
forall a. a -> a
id)

aritySql  :: (String -> String -> Either String a) -> (a -> Statement) -> String -> Q Exp
aritySql :: (String -> String -> Either String a)
-> (a -> Statement) -> String -> Q Exp
aritySql String -> String -> Either String a
parse a -> Statement
mkStatement String
raw = do
    Loc
loc <- Q Loc
location
    let e_ast :: Either String Statement
e_ast = a -> Statement
mkStatement (a -> Statement) -> Either String a -> Either String Statement
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> String -> Either String a
parse (Loc -> String
forall a. Show a => a -> String
show Loc
loc) String
raw
    case Either String Statement
e_ast of
        Right Statement
parsed -> do
            let
                positionalCount :: Word
positionalCount = Statement -> Word
maxParam Statement
parsed
                (Statement
rewritten, AntiquoteState
aqs) = Word -> Statement -> (Statement, AntiquoteState)
numberAntiquotes Word
positionalCount Statement
parsed
                antiNames :: [Name]
antiNames = (Text -> Name) -> [Text] -> [Name]
forall a b. (a -> b) -> [a] -> [b]
map (String -> Name
mkName (String -> Name) -> (Text -> String) -> Text -> Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> String
T.unpack) (AntiquoteState -> [Text]
haskellExpressions AntiquoteState
aqs)
            Exp
typedQuery <- Statement -> Q Exp
makeArityQuery Statement
rewritten
            case Word
positionalCount of
                Word
0 -> -- only antiquotes (or no params)
                    Exp -> Q Exp
forall (m :: * -> *) a. Monad m => a -> m a
return (Exp -> Q Exp) -> Exp -> Q Exp
forall a b. (a -> b) -> a -> b
$ [Exp] -> Exp
tupleE [Exp
typedQuery, [Name] -> Exp
tupleOrSingle [Name]
antiNames]
                Word
1 -> do -- one positional param, doesn't take a tuple
                    Name
patternName <- String -> Q Name
newName String
"c"
                    Exp -> Q Exp
forall (m :: * -> *) a. Monad m => a -> m a
return (Exp -> Q Exp) -> Exp -> Q Exp
forall a b. (a -> b) -> a -> b
$ [Pat] -> Exp -> Exp
LamE [Name -> Pat
VarP Name
patternName]
                        ([Exp] -> Exp
tupleE [Exp
typedQuery, [Name] -> Exp
tupleOrSingle (Name
patternName Name -> [Name] -> [Name]
forall a. a -> [a] -> [a]
: [Name]
antiNames)])
                Word
_ -> do -- at least two positional parameters
                    [Name]
patternNames <- Char -> Int -> Q [Name]
cNames Char
'q' (Word -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word
positionalCount)
                    Exp -> Q Exp
forall (m :: * -> *) a. Monad m => a -> m a
return (Exp -> Q Exp) -> Exp -> Q Exp
forall a b. (a -> b) -> a -> b
$ [Pat] -> Exp -> Exp
LamE
                        [[Pat] -> Pat
TupP ((Name -> Pat) -> [Name] -> [Pat]
forall a b. (a -> b) -> [a] -> [b]
map Name -> Pat
VarP [Name]
patternNames)]
                        ([Exp] -> Exp
tupleE [Exp
typedQuery, [Name] -> Exp
tupleOrSingle ([Name]
patternNames [Name] -> [Name] -> [Name]
forall a. [a] -> [a] -> [a]
++ [Name]
antiNames)])
        Left String
err -> String -> Q Exp
forall a. HasCallStack => String -> a
error String
err

countColumnsReturned :: Statement -> Maybe Int
countColumnsReturned :: Statement -> Maybe Int
countColumnsReturned (QS SelectStmt
selectQ) = SelectStmt -> Maybe Int
go SelectStmt
selectQ where
  go :: SelectStmt -> Maybe Int
go SelectStmt
s = case SelectStmt
s of
      SelectValues NonEmpty (NonEmpty Expr)
rows -> Int -> Maybe Int
forall a. a -> Maybe a
Just ((Int -> Int -> Int) -> Int -> NonEmpty Int -> Int
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
0 ((NonEmpty Expr -> Int) -> NonEmpty (NonEmpty Expr) -> NonEmpty Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap NonEmpty Expr -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length NonEmpty (NonEmpty Expr)
rows))
      Simple Select {[ResTarget]
$sel:targetList:Select :: Select -> [ResTarget]
targetList :: [ResTarget]
targetList} -> if ResTarget
Star ResTarget -> [ResTarget] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [ResTarget]
targetList
          then Maybe Int
forall a. Maybe a
Nothing
          else Int -> Maybe Int
forall a. a -> Maybe a
Just ([ResTarget] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [ResTarget]
targetList)
      S SelectStmt
ss SelectOptions
_ -> SelectStmt -> Maybe Int
go SelectStmt
ss
      Set SetOp
_ AllOrDistinct
_ SelectStmt
a SelectStmt
b -> case (SelectStmt -> Maybe Int
go SelectStmt
a, SelectStmt -> Maybe Int
go SelectStmt
b) of
        (Just Int
m, Just Int
n) | Int
m Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
n -> Int -> Maybe Int
forall a. a -> Maybe a
Just Int
n
        (Maybe Int, Maybe Int)
_ -> Maybe Int
forall a. Maybe a
Nothing
countColumnsReturned Statement
_                       = Int -> Maybe Int
forall a. a -> Maybe a
Just Int
0
-- TODO INSERT ... RETURNING &c