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