Quasi.hs

{-# LANGUAGE GADTs, QuasiQuotes, TemplateHaskell #-}
{-# OPTIONS_GHC -Wall #-}

module Database.MySQL.Simple.Quasi (
-- * Motivation
-- | So, you want to access a MySQL database from Haskell.
-- In an ideal world, your database schema would be known in its entirety to the compiler,
-- and your MySQL queries would be fully parsed at compile-time and type-checked.
--
-- In reality, constructing a full parser for a MySQL query is a huge job,
-- as is creating an EDSL for MySQL queries.  But, still, the mysql-simple package by itself is a little
-- too type-unsafe.  For example, there's nothing stopping you writing this:
--
-- > (a, b) <- query conn "select a, b, c from table where id = ? and name = ?" (True, "x", 7)
--
-- That is: there's no guarantee that the number of params or results in the query matches
-- the number you try to pass in\/receive out of the @query@ call.  Additionally, there's no
-- type checking on the inputs or outputs to the query.

-- * Description
-- | This module provides a quasi-quoter that is a half-way house between mysql-simple and
-- a full-on database library/EDSL.  You write your queries as a quasi-quote, and provide
-- type annotations /in the query itself/.  So for the previous example you would instead write:
--
-- > (a, b) <- query [qquery|select a{Int}, b{Bool}, c{Maybe String} from table where id = ?{Int} and name = ?{String}"|] (True, "x", 7)
--
-- You would then receive two type errors.  One that your tuple @(a, b)@ does not match @(Int, Bool, Maybe String)@,
-- and the other that @(True, \"x\", 7)@ does not match @(Int, String)@.
--
-- So, using the 'qquery' quasi-quoter gives the following benefits:
--
--   * It types the query with the annotated return type (and if any @?@ are used, the param types).
-- 
--   * It automatically chooses to use @QQuery@ or @QQuery_@ based on whether there are @?@ present.
-- 
--   * It automatically wraps/unwraps @Only@ types that mysql-simple uses to disambiguate its instances (but introduces no amibiguity).
-- 
--   * It prevents splicing together strings to make a query.  You must provide the entire query in one literal, and use
--   wildcards to adjust any values.  (I think this is a benefit, right?  Enforced discipline.)
--
-- One technique that I find useful to combine with this quasi-quoter is typing integer keys differently.
-- For example, let's say that you have two tables, @\"users\"@ and @\"locations\"@, each with an @\"id\"@ field.  You want
-- an inner join of these tables where the @\"id\"@s are less than a given amount.  If you use say, @Int32@, you can
-- easily get confused, even with this quasi-quoter:
--
-- > (loc, user) {- mistake -} <- [qquery|select users.id{Int32}, locations.id{Int32}
-- >   from users inner join locations
-- >   where users.location_type = locations.type and users.id < ?{Int32} and locations.id < ?{Int32}|]
-- >     (locIdThreshold, userIdThreshold) {- mistake -}
--
-- If instead you type them differently (with the appropriate instances), you will get two errors when doing this:
--
-- > newtype UserId = UserId Int
-- > newtype LocationId = LocationId Int
--
-- > (loc, user) {- mistake -} <- [qquery|select users.id{UserId}, locations.id{LocationId}
-- >   from users inner join locations
-- >   where users.location_type = locations.type and users.id < ?{UserId} and locations.id < ?{LocationId}|]
-- >     (locIdThreshold, userIdThreshold) {- mistake -}

-- * Quasi-Quoters
  qquery, qexec,
-- * Query Types
  QQuery, QQuery_, QExecute, QExecute_,
-- * Running queries
  execute, execute_, executeMany, fold, fold_, forEach, forEach_, query, query_,
-- * Internal access
  QExtractable(extractQuery)) where

import Control.Applicative ((<$>))
import Control.Monad (when)
import Data.Int (Int64)
import Data.String (fromString)

import Database.MySQL.Simple (Connection, Only(..))
import qualified Database.MySQL.Simple as Orig (Query, execute, execute_, executeMany, fold, fold_, forEach, forEach_, query, query_)
import Database.MySQL.Simple.Param (Param)
import Database.MySQL.Simple.QueryParams (QueryParams)
import Database.MySQL.Simple.QueryResults (QueryResults)

import Language.Haskell.TH.Syntax (Q, Exp(AppE, ConE, LitE, SigE, TupE, VarE), Lit(StringL), Name,
  Pred(ClassP), Type(AppT, ConT, ForallT, TupleT, VarT), TyVarBndr(PlainTV), mkName)
import Language.Haskell.TH.Quote (QuasiQuoter(..))

import Language.Haskell.Meta.Parse (parseType)

-- | A select-like query that takes `q` as its parameters and returns a list of `r` as its results.
data QQuery q r where
  QQuery :: (QueryParams q', QueryResults r') => (q -> q') -> (r' -> r) -> (String, Orig.Query) -> QQuery q r

-- | A select-like query that has no parameters, and returns a list of `r` as its results.
data QQuery_ r where
  QQuery_ :: QueryResults r' => (r' -> r) -> (String, Orig.Query) -> QQuery_ r

-- | An execute-like query that takes `q` as its parameters.
data QExecute q where
  QExecute :: QueryParams q' => (q -> q') -> (String, Orig.Query) -> QExecute q

-- | An execute-like query that has no parameters.  There's very little gain in using
-- this over using `Database.MySQL.Simple.execute_` from mysql-simple directly, but it's provided for completeness.
data QExecute_ where
  QExecute_ :: (String, Orig.Query) -> QExecute_

class QExtractable q where
  -- | Extracts the query.  This loses all the type safety
  -- of the original query and the whole point of using the library,
  -- but presumably you know what you're doing.
  extractQuery :: q -> String

instance QExtractable (QQuery q r) where extractQuery (QQuery _ _ q) = fst q

instance QExtractable (QQuery_ r) where extractQuery (QQuery_ _ q) = fst q

instance QExtractable (QExecute q) where extractQuery (QExecute _ q) = fst q

instance QExtractable QExecute_ where extractQuery (QExecute_ q) = fst q

-- | A wrapper
-- for mysql-simple's 'Database.MySQL.Simple.query' function.
--
-- Note that no instances are required for @q@ or @r@ because the 'QQuery' type
-- witnesses them at its construction.
query :: Connection -> QQuery q r -> q -> IO [r]
query conn (QQuery qf rf (_, s)) = fmap (map rf) . Orig.query conn s . qf

-- | A wrapper
-- for mysql-simple's 'Database.MySQL.Simple.query_' function.
--
-- Note that no instances are required for @r@ because the 'QQuery_' type
-- witnesses them at its construction.
query_ :: Connection -> QQuery_ r -> IO [r]
query_ conn (QQuery_ rf (_, s)) = fmap (map rf) $ Orig.query_ conn s

-- | A wrapper
-- for mysql-simple's 'Database.MySQL.Simple.fold' function.
--
-- Note that no instances are required for @q@ or @r@ because the 'QQuery' type
-- witnesses them at its construction.
fold :: Connection -> QQuery q r -> q -> a -> (a -> r -> IO a) -> IO a
fold conn (QQuery qf rf (_, s)) q x f = Orig.fold conn s (qf q) x (\y r -> f y (rf r))

-- | A wrapper
-- for mysql-simple's 'Database.MySQL.Simple.fold_' function.
--
-- Note that no instances are required for @r@ because the 'QQuery_' type
-- witnesses them at its construction.
fold_ :: Connection -> QQuery_ r -> a -> (a -> r -> IO a) -> IO a
fold_ conn (QQuery_ rf (_, s)) x f = Orig.fold_ conn s x (\y r -> f y (rf r))

-- | A wrapper
-- for mysql-simple's 'Database.MySQL.Simple.forEach' function.
--
-- Note that no instances are required for @q@ or @r@ because the 'QQuery' type
-- witnesses them at its construction.
forEach :: Connection -> QQuery q r -> q -> (r -> IO ()) -> IO ()
forEach conn (QQuery qf rf (_, s)) q f = Orig.forEach conn s (qf q) (f . rf)

-- | A wrapper
-- for mysql-simple's 'Database.MySQL.Simple.forEach_' function.
--
-- Note that no instances are required for @r@ because the 'QQuery_' type
-- witnesses them at its construction.
forEach_ :: Connection -> QQuery_ r -> (r -> IO ()) -> IO ()
forEach_ conn (QQuery_ rf (_, s)) f = Orig.forEach_ conn s (f . rf)

-- | A wrapper
-- for mysql-simple's 'Database.MySQL.Simple.execute' function.
--
-- Note that no instances are required for @q@ because the 'QExecute' type
-- witnesses them at its construction.
execute :: Connection -> QExecute q -> q -> IO Int64
execute conn (QExecute qf (_, s)) = Orig.execute conn s . qf

-- | A wrapper
-- for mysql-simple's 'Database.MySQL.Simple.execute_' function.
execute_ :: Connection -> QExecute_ -> IO Int64
execute_ conn (QExecute_ (_, s)) = Orig.execute_ conn s

-- | A wrapper
-- for mysql-simple's 'Database.MySQL.Simple.executeMany' function.
--
-- Note that no instances are required for @q@ because the 'QExecute' type
-- witnesses them at its construction.
executeMany :: Connection -> QExecute q -> [q] -> IO Int64
executeMany conn (QExecute qf (_, s)) = Orig.executeMany conn s . map qf

-- | A quasi-quoter that takes the param and result types from the query
-- string and generates a typed query.  For example:
--
-- > [qquery|select * from users|]
--
-- will turn into an expression of type @QueryResults r => QQuery_ r@.
-- This is not particularly useful.  However, this:
--
-- > [qquery|select id{Int32}, name{String} from users|]
-- 
-- becomes @QQuery_ (Int32, String)@.
--
-- Furthermore, this:
--
-- > [qquery|select id{Int32} from users where name = ?{String}|]
--
-- becomes: @QQuery String Int32@.  And this:
--
-- > [qquery| select a.*{Int, Maybe String, String}, b.value{Double}
-- >            from a inner join b on a.id = b.id
-- >            where a.name = ?{String} and b.num = ?{Int}|]
--
-- becomes: @QQuery (String, Int) (Int, Maybe String, String, Double)@.
--
-- In general:
-- 
-- * Any non-escaped question mark in the String is taken to be one
--   substitution.  It is given type @QueryParam a => a@ unless it is
--   followed immediately (no spaces) by curly brackets with a type in it,
--   in which case it uses that type.
-- 
-- * A question mark preceded by a backslash is turned into a single question mark.
-- 
-- * To insert an actual backslash, use double backslash.
-- 
-- * Any other instances of curly brackets in the String are taken to
--   be a comma-separated list of result types, which are all tupled
--   (in the order they appear in the String) into a single result type.
--   To get a literal curly bracket, put a backslash before it.
-- 
-- * If there is only a single substitution or single result, @Only@ is automatically added/removed
--   when passing it through to the mysql-simple library.
-- 
-- * If there is no @?@ substitution in the query, the resulting type is @QQuery_ r@.  If there are
--   substitutions, the resulting type is @QQuery q r@.
qquery :: QuasiQuoter
qquery = qqExp "qquery" qqueryExp

-- | Same as `qquery`, except that it produces a query of type QExecute/QExecute_
-- instead of QQuery/QQuery_, and it gives an error if there are any result annotations
-- (since executes don't return any results).
qexec :: QuasiQuoter
qexec = qqExp "qexec" qexecExp

qqExp :: String -> (String -> Q Exp) -> QuasiQuoter
qqExp name qq = QuasiQuoter {quoteExp = qq, quotePat = err, quoteType = err, quoteDec = err}
  where
    err :: String -> Q a
    err _ = fail $ name ++ " quasi-quoter only valid in expression"

emptyToNothing :: [a] -> [Maybe a]
emptyToNothing [] = [Nothing]
emptyToNothing xs = map Just xs

qqueryExp :: String -> Q Exp
qqueryExp input = do
  processed <- either fail return $ parse input
  let result = tuple QResult $ emptyToNothing $ returnTypes processed
      param = tuple QParam $ paramTypes processed
      noParams = null (paramTypes processed)
  return $
   SigE
    ((if noParams
        then ConE 'QQuery_ `AppE` wrapUnwrap result
        else (ConE 'QQuery `AppE` wrapUnwrap param) `AppE` wrapUnwrap result)
      `AppE` mkQuery (stripped processed))
    (makeContext (typeContext param ++ typeContext result) $
       if noParams
         then ConT ''QQuery_ `AppT` qtype result
         else (ConT ''QQuery `AppT` qtype param) `AppT` qtype result)
  -- Result looks like these:
  -- (QQuery_ id "select ...") :: QQuery_ (Int, String)
  -- (QQuery_ fromOnly "select ...") :: QQuery_ Int
  -- (QQuery_ id "select ...") :: (QueryResults a => QQuery_ a)
  -- (QQuery fromOnly Only "select ...") :: QQuery Int String
  -- (QQuery id Only "select ...") :: QQuery Int (String, String)
  -- (QQuery id id "select ...") :: QQuery (Int, Int) (Bool, String)
  -- (QQuery id id "select ...") :: (QueryResults a, QueryParam b) => (QQuery b a))

-- | A version of qqueryExp, above, that allows no return types:
qexecExp :: String -> Q Exp
qexecExp input = do
    processed <- either fail return $ parse input
    when (not . null $ returnTypes processed) $
      fail "Found return annotations in qexec query"
    let param = tuple QParam $ paramTypes processed
        noParams = null $ paramTypes processed
    return $
      SigE
        ((if noParams
            then ConE 'QExecute_
            else ConE 'QExecute `AppE` wrapUnwrap param)
         `AppE` mkQuery (stripped processed))
        (if noParams
           then ConT ''QExecute_
           else makeContext (typeContext param) $ ConT ''QExecute `AppT` qtype param)

mkQuery :: String -> Exp
mkQuery s = TupE [litS, VarE 'fromString `AppE` litS]
  where
    litS = LitE (StringL s)

data QResultOrParam = QResult | QParam

data QType = QT {
  -- | The expression needed to wrap/unwrap a type
  -- (Needed to add/remove the Only wrapper on single types)
  wrapUnwrap :: Exp,
  -- | Required context: class name, var name
  typeContext :: [(Name, Name)],
  -- | The actual type (may be a single type or a tuple):
  qtype :: Type }

-- | Takes a list of types and turns them 
tuple :: QResultOrParam -> [Maybe Type] -> QType
tuple rp mts = QT {wrapUnwrap = w, typeContext = ctxt, qtype = case subst (map VarT vars) mts of
  [t] -> t
  ts -> foldl AppT (TupleT (length ts)) ts}
  where
    w = case (mts, rp) of
          ([Just _], QResult) -> VarE 'fromOnly
          ([_], QParam) -> ConE 'Only
          _ -> VarE 'id
    (c, clz) = case rp of
                 QResult -> ('r', ''QueryResults)
                 QParam -> ('q', ''Param)

    ctxt = zip (repeat clz) vars
    vars = take (countNothings mts) [mkName $ c : suffix | suffix <- "" : map show [(1::Int)..]]

makeContext :: [(Name, Name)] -> (Type -> Type)
makeContext [] = id
makeContext classVars
  = ForallT (map (PlainTV . snd) classVars)
            [ClassP cls [VarT var] | (cls, var) <- classVars]

countNothings :: [Maybe a] -> Int
countNothings xs = length [() | Nothing <- xs]

subst :: [a] -> [Maybe a] -> [a]
subst _ [] = []
subst (s:ss) (Nothing:xs) = s : subst ss xs
subst [] (Nothing : _) = error "Internal error: not enough types to substitue"
subst ss (Just x : xs) = x : subst ss xs

(<:>) :: Functor m => Char -> m ProcessedQuery -> m ProcessedQuery
(<:>) c = fmap (\p -> p { stripped = c : stripped p })

toCloseCurly :: String -> Either String (String, String)
toCloseCurly src = case span (/= '}') src of
  (before, '}':after) -> return (before, after)
  _ -> Left $ "Unclosed '{' in remainder of expression: " ++ src

data ProcessedQuery = PQ { returnTypes :: [Type], paramTypes :: [Maybe Type], stripped :: String }

-- | Parses a String looking for various annotations. Gives back either an error,
-- or the query string with the annotations removed and processed into lists of types.
parse :: String -> Either String ProcessedQuery
-- Double backslash becomes single backslash:
parse ('\\':'\\':rest) = '\\' <:> parse rest
-- Backslashed question mark or curly bracket makes it literal:
parse ('\\':'{':rest) = '{' <:> parse rest
parse ('\\':'?':rest) = '?' <:> parse rest
-- Question mark then curly bracket is a param type:
parse ('?':'{':rest) = '?' <:> do
  (before, after) <- toCloseCurly rest
  types <- parseTypes before
  remainder <- parse after
  case types of
    [single] -> return $ remainder { paramTypes = Just single : paramTypes remainder }
    _ -> fail $ "Multiple types given for substitution: " ++ before
-- Any other curly bracket is a return type:
parse ('{':rest) = do
  (before, after) <- toCloseCurly rest
  types <- parseTypes before
  remainder <- parse after
  return $ remainder { returnTypes = types ++ returnTypes remainder }
parse ('?':rest) = '?' <:> do
  remainder <- parse rest
  return $ remainder { paramTypes = Nothing : paramTypes remainder }
parse (c:cs) = c <:> parse cs
parse "" = return $ PQ [] [] ""

-- | Takes a comma-separated list of types and parses them.
parseTypes :: String -> Either String [Type]
-- Slight hack; make sure a single type parses as a tuple too, by adding another type
-- so that it's at least size 2:
parseTypes s = case flattenAppT <$> parseType ("( ()," ++ s ++ ")") of
  Right (TupleT n : _ : xs) | n == 1 + length xs -> Right xs
  Left e -> Left e
  Right t -> Left $ "Bad type: " ++ show t

-- | Takes a type that is a series of nested applications and flattens it into a list.
-- For example, the type (Int, String, Bool) is really (,,) Int String Bool,
-- which is represented in an AST as:
-- ((( (,,) `AppT` Int) `AppT` String) `AppT` Bool)
-- or rather:
-- AppT (AppT (AppT (,,) Int) String) Bool
--
-- This function takes something like that and turns it into:
--
-- [(,,), Int, String, Bool]
flattenAppT :: Type -> [Type]
flattenAppT (AppT f t) = flattenAppT f ++ [t]
flattenAppT t = [t]