1 {-# LANGUAGE FunctionalDependencies, TemplateHaskell #-}
    2 module Databrary.Model.SQL.Select
    3   ( SelectOutput(..)
    4   , Selector(..)
    5   , selector
    6   , selectColumn
    7   , selectColumns
    8   , addSelects
    9   , fromMap
   10   , fromAlias
   11   , crossJoin
   12   , joinOn
   13   , joinUsing
   14   , maybeJoinOn
   15   , maybeJoinUsing
   16   , selectJoin
   17   , selectMap
   18   , makeQuery
   19   , selectDistinctQuery
   20   , nameRef
   21   ) where
   22 
   23 import Control.Arrow (second)
   24 import Control.Monad.State (StateT(..))
   25 import Control.Monad.Trans.Class (lift)
   26 import Data.Char (isLetter, toLower)
   27 import Data.List (intercalate, unfoldr)
   28 import Database.PostgreSQL.Typed.Query (QueryFlags, parseQueryFlags, makePGQuery)
   29 import qualified Language.Haskell.TH as TH
   30 
   31 import Databrary.Service.DB (useTDB)
   32 
   33 data SelectOutput
   34   = SelectColumn { _selectTable, _selectColumn :: String }
   35   | SelectExpr String
   36   | OutputJoin { outputNullable :: !Bool, outputJoiner :: TH.Name, outputJoin :: [SelectOutput] }
   37   | OutputMap { outputNullable :: !Bool, outputMapper :: TH.Exp -> TH.Exp, outputMap :: SelectOutput }
   38 
   39 _outputTuple :: [SelectOutput] -> SelectOutput
   40 _outputTuple l = OutputJoin False (TH.tupleDataName $ length l) l
   41 
   42 outputMaybe :: SelectOutput -> SelectOutput
   43 outputMaybe (OutputJoin False f l) = OutputJoin True f l
   44 outputMaybe (OutputMap False f l) = OutputMap True f l
   45 outputMaybe s = s
   46 
   47 outputColumns :: SelectOutput -> [String]
   48 outputColumns (SelectColumn t c) = [t ++ '.' : c]
   49 outputColumns (SelectExpr s) = [s]
   50 outputColumns (OutputJoin _ _ o) = concatMap outputColumns o
   51 outputColumns (OutputMap _ _ o) = outputColumns o
   52 
   53 outputParser :: SelectOutput -> StateT [TH.Name] TH.Q TH.Exp
   54 outputParser (OutputJoin mb f ol) = do
   55   fi <- lift $ TH.reify f
   56   (fe, ft) <- case fi of
   57     TH.ClassOpI _ t _ -> return (TH.VarE f, t)
   58     TH.DataConI _ t _ -> return (TH.ConE f, t)
   59     TH.VarI _ t _ -> return (TH.VarE f, t)
   60     _ -> die "wrong kind"
   61   if mb
   62     then do
   63       let am = unfoldr argMaybe ft
   64       (bl, ae) <- bindArgs am ol
   65       -- when (null bl) $ die "function with at least one non-Maybe argument required"
   66       return $ TH.DoE $ bl ++ [TH.NoBindS $ TH.AppE (TH.ConE 'Just) $ foldl TH.AppE fe ae]
   67     else foldl TH.AppE fe <$> mapM outputParser ol
   68   where
   69   bindArgs (False:m) (o:l) = do
   70     n <- lift $ TH.newName "cm"
   71     a <- outputParser (outputMaybe o)
   72     (bl, al) <- bindArgs m l
   73     return $ (TH.BindS (TH.VarP n) a : bl, TH.VarE n : al)
   74   bindArgs (True:m) (o:l) = do
   75     a <- outputParser o
   76     second (a:) <$> bindArgs m l
   77   bindArgs _ o = (,) [] <$> mapM outputParser o
   78   argMaybe (TH.ArrowT `TH.AppT` a `TH.AppT` r) = Just (isMaybeT a, r)
   79   argMaybe _ = Nothing
   80   isMaybeT (TH.AppT (TH.ConT m) _) = m == ''Maybe
   81   isMaybeT _ = False
   82   die s = fail $ "outputParser " ++ show f ++ ": " ++ s
   83 outputParser (OutputMap False f o) =
   84   f <$> outputParser o
   85 outputParser (OutputMap True f o) = do
   86   x <- lift $ TH.newName "x"
   87   ((TH.VarE 'fmap `TH.AppE` (TH.LamE [TH.VarP x] $ f $ TH.VarE x)) `TH.AppE`)
   88     <$> outputParser (outputMaybe o)
   89 outputParser _ = StateT st where
   90   st (i:l) = return (TH.VarE i, l)
   91   st [] = fail "outputParser: insufficient values"
   92 
   93 data Selector = Selector
   94   { selectOutput :: SelectOutput
   95   , selectSource :: String
   96   , selectJoined :: String
   97   }
   98 
   99 selector :: String -> SelectOutput -> Selector
  100 selector t o = Selector o t (',':t)
  101 
  102 selectColumn :: String -> String -> Selector
  103 selectColumn t c = selector t $ SelectColumn t c
  104 
  105 selectColumns :: TH.Name -> String -> [String] -> Selector
  106 selectColumns f t c =
  107   selector t $ OutputJoin False f $ map (SelectColumn t) c
  108 
  109 addSelects :: TH.Name -> Selector -> [SelectOutput] -> Selector
  110 addSelects f s c = s
  111   { selectOutput = OutputJoin False f (selectOutput s : c) }
  112 
  113 fromMap :: (String -> String) -> Selector -> Selector
  114 fromMap f sel = sel
  115   { selectSource = f $ selectSource sel
  116   , selectJoined = f $ selectJoined sel
  117   }
  118 
  119 outputFromAlias :: String -> SelectOutput -> SelectOutput
  120 outputFromAlias t (SelectColumn _ c) = SelectColumn t c
  121 outputFromAlias _ (SelectExpr e) = error $ "fromAlias (SelectExpr " ++ show e ++ ")"
  122 outputFromAlias t o@OutputJoin{ outputJoin = l } = o{ outputJoin = map (outputFromAlias t) l }
  123 outputFromAlias t o@OutputMap{ outputMap = l } = o{ outputMap = outputFromAlias t l }
  124 
  125 fromAlias :: Selector -> String -> Selector
  126 fromAlias sel as = fromMap (++ " AS " ++ as) sel
  127   { selectOutput = outputFromAlias as $ selectOutput sel }
  128 
  129 joinWith :: (String -> String) -> Selector -> Selector
  130 joinWith j sel = sel{ selectJoined = j (selectSource sel) }
  131 
  132 maybeJoinWith :: (String -> String) -> Selector -> Selector
  133 maybeJoinWith j sel = sel
  134   { selectJoined = j (selectSource sel)
  135   , selectOutput = outputMaybe (selectOutput sel) }
  136 
  137 crossJoin :: Selector -> Selector
  138 crossJoin = joinWith (" CROSS JOIN " ++)
  139 
  140 joinOn :: String -> Selector -> Selector
  141 joinOn on = joinWith (\s -> " JOIN " ++ s ++ " ON " ++ on)
  142 
  143 joinUsing :: [String] -> Selector -> Selector
  144 joinUsing using = joinWith (\s -> " JOIN " ++ s ++ " USING (" ++ intercalate "," using ++ ")")
  145 
  146 maybeJoinOn :: String -> Selector -> Selector
  147 maybeJoinOn on = maybeJoinWith (\s -> " LEFT JOIN " ++ s ++ " ON " ++ on)
  148 
  149 maybeJoinUsing :: [String] -> Selector -> Selector
  150 maybeJoinUsing using = maybeJoinWith (\s -> " LEFT JOIN " ++ s ++ " USING (" ++ intercalate "," using ++ ")")
  151 
  152 selectJoin :: TH.Name -> [Selector] -> Selector
  153 selectJoin f l@(h:t) = Selector
  154   { selectOutput = OutputJoin False f $ map selectOutput l
  155   , selectSource = selectSource h ++ joins
  156   , selectJoined = selectJoined h ++ joins
  157   } where joins = concatMap selectJoined t
  158 selectJoin _ [] = error "selectJoin: empty list"
  159 
  160 selectMap :: (TH.Exp -> TH.Exp) -> Selector -> Selector
  161 selectMap f s = s{ selectOutput = OutputMap False f (selectOutput s) }
  162 
  163 
  164 takeWhileEnd :: (a -> Bool) -> [a] -> [a]
  165 takeWhileEnd p = fst . foldr go ([], False) where
  166   go x (rest, done)
  167     | not done && p x = (x:rest, False)
  168     | otherwise = (rest, True)
  169 
  170 makeQuery :: QueryFlags -> (String -> String) -> SelectOutput -> TH.ExpQ
  171 makeQuery flags sql output = do
  172   _ <- useTDB
  173   nl <- mapM (TH.newName . ('v':) . colVar) cols
  174   (parse, []) <- runStateT (outputParser output) nl
  175   TH.AppE (TH.VarE 'fmap `TH.AppE` TH.LamE [TH.TupP $ map TH.VarP nl] parse)
  176     <$> makePGQuery flags (sql $ intercalate "," cols)
  177   where
  178   colVar s = case takeWhileEnd isLetter s of
  179     [] -> "c"
  180     (h:l) -> toLower h : l
  181   cols = outputColumns output
  182 
  183 selectDistinctQuery :: Maybe [String] -> Selector -> String -> TH.ExpQ
  184 selectDistinctQuery dist (Selector{ selectOutput = o, selectSource = s }) sqlf =
  185   makeQuery flags (\c -> select dist ++ c ++ " FROM " ++ s ++ ' ':sql) o
  186   where
  187   (flags, sql) = parseQueryFlags sqlf
  188   select Nothing = "SELECT " -- ALL
  189   select (Just []) = "SELECT DISTINCT "
  190   select (Just l) = "SELECT DISTINCT ON (" ++ intercalate "," l ++ ") "
  191 
  192 nameRef :: TH.Name -> String
  193 nameRef n = maybe b (++ '.' : b) $ TH.nameModule n where b = TH.nameBase n