PersistBackendHelpers.hs

{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}

-- | This helper module contains generic versions of PersistBackend functions
module Database.Groundhog.Generic.PersistBackendHelpers
  ( get,
    select,
    selectAll,
    selectStream,
    selectAllStream,
    getBy,
    project,
    projectStream,
    count,
    replace,
    replaceBy,
    update,
    delete,
    deleteBy,
    deleteAll,
    insertByAll,
    countAll,
    insertBy,
  )
where

import Data.Either (rights)
import Data.Maybe (catMaybes, fromJust, fromMaybe, mapMaybe)
import Database.Groundhog.Core (PersistBackendConn)
import Database.Groundhog.Core hiding (PersistBackendConn (..))
import qualified Database.Groundhog.Core as Core
import Database.Groundhog.Generic (firstRow, getUniqueFields, isSimple, joinStreams, mapStream, streamToList)
import Database.Groundhog.Generic.Sql

get ::
  forall conn v.
  (PersistBackendConn conn, PersistEntity v, PrimitivePersistField (Key v BackendSpecific)) =>
  RenderConfig ->
  (Utf8 -> [PersistValue] -> Action conn (RowStream [PersistValue])) ->
  Key v BackendSpecific ->
  Action conn (Maybe v)
get RenderConfig {..} queryFunc (k :: Key v BackendSpecific) = do
  let e = entityDef proxy (undefined :: v)
      proxy = undefined :: proxy conn
  if isSimple (constructors e)
    then do
      let constr = head $ constructors e
      let fields = renderFields esc (constrParams constr)
      let query = "SELECT " <> fields <> " FROM " <> tableName esc e constr <> " WHERE " <> fromJust (constrId esc constr) <> "=?"
      x <- queryFunc query [toPrimitivePersistValue k] >>= firstRow
      case x of
        Just xs -> fmap (Just . fst) $ fromEntityPersistValues $ PersistInt64 0 : xs
        Nothing -> pure Nothing
    else do
      let query = "SELECT discr FROM " <> mainTableName esc e <> " WHERE id=?"
      x <- queryFunc query [toPrimitivePersistValue k] >>= firstRow
      case x of
        Just [discr] -> do
          let constructorNum = fromPrimitivePersistValue discr
              constr = constructors e !! constructorNum
              fields = renderFields esc (constrParams constr)
              cQuery = "SELECT " <> fields <> " FROM " <> tableName esc e constr <> " WHERE " <> fromJust (constrId esc constr) <> "=?"
          x2 <- queryFunc cQuery [toPrimitivePersistValue k] >>= firstRow
          case x2 of
            Just xs -> fmap (Just . fst) $ fromEntityPersistValues $ discr : xs
            Nothing -> fail "Missing entry in constructor table"
        Just x' -> fail $ "Unexpected number of columns returned: " ++ show x'
        Nothing -> pure Nothing

select ::
  forall conn r v c opts.
  (SqlDb conn, r ~ RestrictionHolder v c, PersistBackendConn conn, PersistEntity v, EntityConstr v c, HasSelectOptions opts conn r) =>
  RenderConfig ->
  (Utf8 -> [PersistValue] -> Action conn (RowStream [PersistValue])) ->
  (opts -> RenderS conn r) ->
  Utf8 ->
  opts ->
  Action conn [v]
select conf queryFunc preColumns noLimit options = selectStream conf queryFunc preColumns noLimit options >>= streamToList

selectStream ::
  forall conn r v c opts.
  (SqlDb conn, r ~ RestrictionHolder v c, PersistBackendConn conn, PersistEntity v, EntityConstr v c, HasSelectOptions opts conn r) =>
  RenderConfig ->
  (Utf8 -> [PersistValue] -> Action conn (RowStream [PersistValue])) ->
  (opts -> RenderS conn r) ->
  Utf8 ->
  opts ->
  Action conn (RowStream v)
selectStream conf@RenderConfig {..} queryFunc preColumns noLimit options = doSelectQuery
  where
    SelectOptions cond limit offset ords dist _ = getSelectOptions options

    e = entityDef proxy (undefined :: v)
    proxy = undefined :: proxy conn
    orders = renderOrders conf ords
    lim = case (limit, offset) of
      (Nothing, Nothing) -> mempty
      (Nothing, o) -> RenderS (" " <> noLimit <> " OFFSET ?") (toPurePersistValues o)
      (l, Nothing) -> RenderS " LIMIT ?" (toPurePersistValues l)
      (l, o) -> RenderS " LIMIT ? OFFSET ?" (toPurePersistValues (l, o))
    cond' = renderCond conf cond
    fields = RenderS (renderFields esc (constrParams constr)) id
    distinctClause = if dist then "DISTINCT " else mempty
    RenderS query binds = "SELECT " <> distinctClause <> preColumns options <> fields <> " FROM " <> RenderS (tableName esc e constr) id <> whereClause <> orders <> lim
    whereClause = maybe "" (" WHERE " <>) cond'
    doSelectQuery = queryFunc query (binds []) >>= mapStream (\xs -> fmap fst $ fromEntityPersistValues $ toPrimitivePersistValue cNum : xs)
    cNum = entityConstrNum (undefined :: proxy v) (undefined :: c a)
    constr = constructors e !! cNum

selectAll ::
  forall conn v.
  (PersistBackendConn conn, PersistEntity v) =>
  RenderConfig ->
  (Utf8 -> [PersistValue] -> Action conn (RowStream [PersistValue])) ->
  Action conn [(AutoKey v, v)]
selectAll conf queryFunc = selectAllStream conf queryFunc >>= streamToList

-- | It may call the passed function multiple times
selectAllStream ::
  forall conn v.
  (PersistBackendConn conn, PersistEntity v) =>
  RenderConfig ->
  (Utf8 -> [PersistValue] -> Action conn (RowStream [PersistValue])) ->
  Action conn (RowStream (AutoKey v, v))
selectAllStream RenderConfig {..} queryFunc = start
  where
    start = joinStreams $ zipWith selectConstr [0 ..] $ constructors e
    selectConstr cNum constr = queryFunc query [] >>= mapStream (mkEntity cNum)
      where
        fields = maybe id (\key cont -> key <> fromChar ',' <> cont) (constrId esc constr) $ renderFields esc (constrParams constr)
        query = "SELECT " <> fields <> " FROM " <> tableName esc e constr
    e = entityDef proxy (undefined :: v)
    proxy = undefined :: proxy conn
    mkEntity cNum xs = do
      let (k, xs') = fromPurePersistValues xs
      (v, _) <- fromEntityPersistValues (toPrimitivePersistValue (cNum :: Int) : xs')
      pure (k, v)

getBy ::
  forall conn v u.
  (PersistBackendConn conn, PersistEntity v, IsUniqueKey (Key v (Unique u))) =>
  RenderConfig ->
  -- | function to run query
  (Utf8 -> [PersistValue] -> Action conn (RowStream [PersistValue])) ->
  Key v (Unique u) ->
  Action conn (Maybe v)
getBy conf@RenderConfig {..} queryFunc (k :: Key v (Unique u)) = do
  uniques <- toPersistValues k
  let e = entityDef proxy (undefined :: v)
      proxy = undefined :: proxy conn
      u = (undefined :: Key v (Unique u) -> u (UniqueMarker v)) k
      uFields = renderChain conf (fieldChain proxy u) []
      RenderS cond vals = intercalateS " AND " $ mkUniqueCond uFields uniques
      constr = head $ constructors e
      fields = renderFields esc (constrParams constr)
      query = "SELECT " <> fields <> " FROM " <> tableName esc e constr <> " WHERE " <> cond
  x <- queryFunc query (vals []) >>= firstRow
  case x of
    Just xs -> fmap (Just . fst) $ fromEntityPersistValues $ PersistInt64 0 : xs
    Nothing -> pure Nothing

project ::
  forall conn r v c p opts a'.
  (SqlDb conn, r ~ RestrictionHolder v c, PersistBackendConn conn, PersistEntity v, EntityConstr v c, Projection p a', ProjectionDb p conn, ProjectionRestriction p r, HasSelectOptions opts conn r) =>
  RenderConfig ->
  (Utf8 -> [PersistValue] -> Action conn (RowStream [PersistValue])) ->
  (opts -> RenderS conn r) ->
  Utf8 ->
  p ->
  opts ->
  Action conn [a']
project conf queryFunc preColumns noLimit p options = projectStream conf queryFunc preColumns noLimit p options >>= streamToList

projectStream ::
  forall conn r v c p opts a'.
  (SqlDb conn, r ~ RestrictionHolder v c, PersistBackendConn conn, PersistEntity v, EntityConstr v c, Projection p a', ProjectionDb p conn, ProjectionRestriction p r, HasSelectOptions opts conn r) =>
  RenderConfig ->
  (Utf8 -> [PersistValue] -> Action conn (RowStream [PersistValue])) ->
  (opts -> RenderS conn r) ->
  Utf8 ->
  p ->
  opts ->
  Action conn (RowStream a')
projectStream conf@RenderConfig {..} queryFunc preColumns noLimit p options = doSelectQuery
  where
    SelectOptions cond limit offset ords dist _ = getSelectOptions options
    e = entityDef proxy (undefined :: v)
    proxy = undefined :: proxy conn
    orders = renderOrders conf ords
    lim = case (limit, offset) of
      (Nothing, Nothing) -> mempty
      (Nothing, o) -> RenderS (" " <> noLimit <> " OFFSET ?") (toPurePersistValues o)
      (l, Nothing) -> RenderS " LIMIT ?" (toPurePersistValues l)
      (l, o) -> RenderS " LIMIT ? OFFSET ?" (toPurePersistValues (l, o))
    cond' = renderCond conf cond
    chains = projectionExprs p []
    fields = commasJoin $ concatMap (renderExprExtended conf 0) chains
    distinctClause = if dist then "DISTINCT " else mempty
    RenderS query binds = "SELECT " <> distinctClause <> preColumns options <> fields <> " FROM " <> RenderS (tableName esc e constr) id <> whereClause <> orders <> lim
    whereClause = maybe "" (" WHERE " <>) cond'
    doSelectQuery = queryFunc query (binds []) >>= mapStream (fmap fst . projectionResult p)
    constr = constructors e !! entityConstrNum (undefined :: proxy v) (undefined :: c a)

count ::
  forall conn r v c.
  (SqlDb conn, r ~ RestrictionHolder v c, PersistBackendConn conn, PersistEntity v, EntityConstr v c) =>
  RenderConfig ->
  (Utf8 -> [PersistValue] -> Action conn (RowStream [PersistValue])) ->
  Cond conn r ->
  Action conn Int
count conf@RenderConfig {..} queryFunc cond = do
  let e = entityDef proxy (undefined :: v)
      proxy = undefined :: proxy conn
      cond' = renderCond conf cond
      constr = constructors e !! entityConstrNum (undefined :: proxy v) (undefined :: c a)
      query = "SELECT COUNT(*) FROM " <> tableName esc e constr <> whereClause
        where
          whereClause = maybe "" (\c -> " WHERE " <> getQuery c) cond'
  x <- queryFunc query (maybe [] (`getValues` []) cond') >>= firstRow
  case x of
    Just [num] -> pure $ fromPrimitivePersistValue num
    Just xs -> fail $ "requested 1 column, returned " ++ show (length xs)
    Nothing -> fail "COUNT returned no rows"

replace ::
  forall conn r v.
  (PersistBackendConn conn, PersistEntity v, PrimitivePersistField (Key v BackendSpecific)) =>
  RenderConfig ->
  (Utf8 -> [PersistValue] -> Action conn (RowStream [PersistValue])) ->
  (Utf8 -> [PersistValue] -> Action conn ()) ->
  (Bool -> Utf8 -> ConstructorDef -> [PersistValue] -> RenderS conn r) ->
  Key v BackendSpecific ->
  v ->
  Action conn ()
replace RenderConfig {..} queryFunc execFunc insertIntoConstructorTable k v = do
  vals <- toEntityPersistValues' v
  let e = entityDef proxy v
      proxy = undefined :: proxy conn
      constructorNum = fromPrimitivePersistValue (head vals)
      constr = constructors e !! constructorNum
      k' = toPrimitivePersistValue k
      RenderS upds updsVals = commasJoin $ zipWith f fields $ tail vals
        where
          fields = foldr (flatten esc) [] $ constrParams constr
          f f1 f2 = RenderS f1 id <> fromChar '=' <> renderPersistValue f2
      updateQuery = "UPDATE " <> tableName esc e constr <> " SET " <> upds <> " WHERE " <> fromString (fromJust $ constrAutoKeyName constr) <> "=?"

  if isSimple (constructors e)
    then execFunc updateQuery (updsVals [k'])
    else do
      let query = "SELECT discr FROM " <> mainTableName esc e <> " WHERE id=?"
      x <- queryFunc query [k'] >>= fmap (fmap $ fromPrimitivePersistValue . head) . firstRow
      case x of
        Just discr -> do
          let cName = tableName esc e constr

          if discr == constructorNum
            then execFunc updateQuery (updsVals [k'])
            else do
              let RenderS insQuery vals' = insertIntoConstructorTable True cName constr (k' : tail vals)
              execFunc insQuery (vals' [])

              let oldConstr = constructors e !! discr
              let delQuery = "DELETE FROM " <> tableName esc e oldConstr <> " WHERE " <> fromJust (constrId esc oldConstr) <> "=?"
              execFunc delQuery [k']

              let updateDiscrQuery = "UPDATE " <> mainTableName esc e <> " SET discr=? WHERE id=?"
              execFunc updateDiscrQuery [head vals, k']
        Nothing -> pure ()

replaceBy ::
  forall conn v u.
  (PersistBackendConn conn, PersistEntity v, IsUniqueKey (Key v (Unique u))) =>
  RenderConfig ->
  -- | function to execute query
  (Utf8 -> [PersistValue] -> Action conn ()) ->
  u (UniqueMarker v) ->
  v ->
  Action conn ()
replaceBy conf@RenderConfig {..} execFunc u v = do
  uniques <- toPersistValues (extractUnique v `asTypeOf` (undefined :: u (UniqueMarker v) -> Key v (Unique u)) u)
  vals <- toEntityPersistValues' v
  let e = entityDef proxy (undefined :: v)
      proxy = undefined :: proxy conn
      uFields = renderChain conf (fieldChain proxy u) []
      RenderS cond condVals = intercalateS " AND " $ mkUniqueCond uFields uniques
      constr = head $ constructors e
      RenderS upds updsVals = commasJoin $ zipWith f fields $ tail vals
        where
          fields = foldr (flatten esc) [] $ constrParams constr
          f f1 f2 = RenderS f1 id <> fromChar '=' <> renderPersistValue f2
      updateQuery = "UPDATE " <> tableName esc e constr <> " SET " <> upds <> " WHERE " <> cond
  execFunc updateQuery (updsVals . condVals $ [])

update ::
  forall conn r v c.
  (SqlDb conn, r ~ RestrictionHolder v c, PersistBackendConn conn, PersistEntity v, EntityConstr v c) =>
  RenderConfig ->
  (Utf8 -> [PersistValue] -> Action conn ()) ->
  [Update conn r] ->
  Cond conn r ->
  Action conn ()
update conf@RenderConfig {..} execFunc upds cond = do
  let e = entityDef proxy (undefined :: v)
      proxy = undefined :: proxy conn
  case renderUpdates conf upds of
    Just upds' -> do
      let cond' = renderCond conf cond
          constr = constructors e !! entityConstrNum (undefined :: proxy v) (undefined :: c a)
          query = "UPDATE " <> tableName esc e constr <> " SET " <> whereClause
            where
              whereClause = maybe (getQuery upds') (\c -> getQuery upds' <> " WHERE " <> getQuery c) cond'
      execFunc query (getValues upds' <> maybe mempty getValues cond' $ [])
    Nothing -> pure ()

delete ::
  forall conn r v c.
  (SqlDb conn, r ~ RestrictionHolder v c, PersistBackendConn conn, PersistEntity v, EntityConstr v c) =>
  RenderConfig ->
  (Utf8 -> [PersistValue] -> Action conn ()) ->
  Cond conn r ->
  Action conn ()
delete conf@RenderConfig {..} execFunc cond = execFunc query (maybe [] (($ []) . getValues) cond')
  where
    e = entityDef proxy (undefined :: v)
    proxy = undefined :: proxy conn
    constr = constructors e !! entityConstrNum (undefined :: proxy v) (undefined :: c a)
    cond' = renderCond conf cond
    whereClause = maybe "" (\c -> " WHERE " <> getQuery c) cond'
    query =
      if isSimple (constructors e)
        then "DELETE FROM " <> tableName esc e constr <> whereClause
        else -- the entries in the constructor table are deleted because of the reference on delete cascade
          "DELETE FROM " <> mainTableName esc e <> " WHERE id IN(SELECT " <> fromJust (constrId esc constr) <> " FROM " <> tableName esc e constr <> whereClause <> ")"

insertByAll ::
  forall conn v.
  (PersistBackendConn conn, PersistEntity v) =>
  RenderConfig ->
  -- | function to run query
  (Utf8 -> [PersistValue] -> Action conn (RowStream [PersistValue])) ->
  -- | allow multiple duplication of uniques with nulls
  Bool ->
  v ->
  Action conn (Either (AutoKey v) (AutoKey v))
insertByAll RenderConfig {..} queryFunc manyNulls v = do
  let e = entityDef proxy v
      proxy = undefined :: proxy conn
      (constructorNum, uniques) = getUniques v
      constr = constructors e !! constructorNum
      uniqueDefs = constrUniques constr

      query = "SELECT " <> fromMaybe "1" (constrId esc constr) <> " FROM " <> tableName esc e constr <> " WHERE " <> cond
      conds = catMaybes $ zipWith (\uFields (_, uVals) -> checkNulls uVals $ intercalateS " AND " $ mkUniqueCond uFields uVals) (mapMaybe f uniqueDefs) uniques
        where
          f u@(UniqueDef _ _ uFields) =
            if null $ rights uFields
              then Just $ foldr (flatten esc) [] $ getUniqueFields u
              else Nothing
      -- skip condition if any value is NULL. It allows to insert many values with duplicate unique key
      checkNulls uVals x = if manyNulls && elem PersistNull (uVals []) then Nothing else Just x
      RenderS cond vals = intercalateS " OR " conds
  if null conds
    then Right <$> Core.insert v
    else do
      x <- queryFunc query (vals []) >>= firstRow
      case x of
        Nothing -> Right <$> Core.insert v
        Just xs -> pure $ Left $ fst $ fromPurePersistValues xs

deleteBy ::
  forall conn v.
  (PersistBackendConn conn, PersistEntity v, PrimitivePersistField (Key v BackendSpecific)) =>
  RenderConfig ->
  (Utf8 -> [PersistValue] -> Action conn ()) ->
  Key v BackendSpecific ->
  Action conn ()
deleteBy RenderConfig {..} execFunc k = execFunc query [toPrimitivePersistValue k]
  where
    e = entityDef proxy ((undefined :: Key v u -> v) k)
    proxy = undefined :: proxy conn
    constr = head $ constructors e
    idName =
      if isSimple (constructors e)
        then fromJust $ constrId esc constr
        else "id"
    -- the entries in the constructor table are deleted because of the reference on delete cascade
    query = "DELETE FROM " <> mainTableName esc e <> " WHERE " <> idName <> "=?"

deleteAll ::
  forall conn v.
  (PersistBackendConn conn, PersistEntity v) =>
  RenderConfig ->
  (Utf8 -> [PersistValue] -> Action conn ()) ->
  v ->
  Action conn ()
deleteAll RenderConfig {..} execFunc (_ :: v) = execFunc query []
  where
    e = entityDef proxy (undefined :: v)
    proxy = undefined :: proxy conn
    query = "DELETE FROM " <> mainTableName esc e

countAll ::
  forall conn v.
  (PersistBackendConn conn, PersistEntity v) =>
  RenderConfig ->
  (Utf8 -> [PersistValue] -> Action conn (RowStream [PersistValue])) ->
  v ->
  Action conn Int
countAll RenderConfig {..} queryFunc (_ :: v) = do
  let e = entityDef proxy (undefined :: v)
      proxy = undefined :: proxy conn
      query = "SELECT COUNT(*) FROM " <> mainTableName esc e
  x <- queryFunc query [] >>= firstRow
  case x of
    Just [num] -> pure $ fromPrimitivePersistValue num
    Just xs -> fail $ "requested 1 column, returned " ++ show (length xs)
    Nothing -> fail "COUNT returned no rows"

insertBy ::
  forall conn v u.
  (PersistBackendConn conn, PersistEntity v, IsUniqueKey (Key v (Unique u))) =>
  RenderConfig ->
  (Utf8 -> [PersistValue] -> Action conn (RowStream [PersistValue])) ->
  Bool ->
  u (UniqueMarker v) ->
  v ->
  Action conn (Either (AutoKey v) (AutoKey v))
insertBy conf@RenderConfig {..} queryFunc manyNulls u v = do
  uniques <- toPersistValues (extractUnique v `asTypeOf` (undefined :: u (UniqueMarker v) -> Key v (Unique u)) u)
  let e = entityDef proxy v
      proxy = undefined :: proxy conn
      uFields = renderChain conf (fieldChain proxy u) []
      RenderS cond vals = intercalateS " AND " $ mkUniqueCond uFields uniques
      -- skip condition if any value is NULL. It allows to insert many values with duplicate unique key
      checkNulls uVals = manyNulls && elem PersistNull (uVals [])
      -- this is safe because unique keys exist only for entities with one constructor
      constr = head $ constructors e
      query = "SELECT " <> fromMaybe "1" (constrId esc constr) <> " FROM " <> tableName esc e constr <> " WHERE " <> cond
  if checkNulls uniques
    then Right <$> Core.insert v
    else do
      x <- queryFunc query (vals []) >>= firstRow
      case x of
        Nothing -> Right <$> Core.insert v
        Just [k] -> pure $ Left $ fst $ fromPurePersistValues [k]
        Just xs -> fail $ "unexpected query result: " ++ show xs

constrId :: (Utf8 -> Utf8) -> ConstructorDef -> Maybe Utf8
constrId escape = fmap (escape . fromString) . constrAutoKeyName

toEntityPersistValues' :: (PersistBackendConn conn, PersistEntity v) => v -> Action conn [PersistValue]
toEntityPersistValues' = fmap ($ []) . toEntityPersistValues

mkUniqueCond :: [Utf8] -> ([PersistValue] -> [PersistValue]) -> [RenderS conn r]
mkUniqueCond u vals = zipWith f u (vals [])
  where
    f a PersistNull = RenderS (a <> " IS NULL") id
    f a x = RenderS (a <> "=?") (x :)