{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TemplateHaskellQuotes #-}
{-# LANGUAGE ViewPatterns #-}

-- | Description: Translate delete statements.
module Squeal.QuasiQuotes.Delete (
  toSquealDelete,
) where

import Control.Monad (when)
import Data.Foldable (Foldable(foldr), foldlM)
import Data.Maybe (isJust)
import Language.Haskell.TH.Syntax (Exp(AppE, ConE, LabelE, VarE), Q)
import Prelude
  ( Applicative(pure), Maybe(Just, Nothing), MonadFail(fail), Semigroup((<>))
  , ($), (<$>)
  )
import Squeal.QuasiQuotes.Query
  ( getIdentText, renderPGTAExpr, renderPGTTableRef, renderPGTTargetList
  , toSquealQuery
  )
import qualified Data.List.NonEmpty as NE
import qualified Data.Text as Text
import qualified PostgresqlSyntax.Ast as PGT_AST
import qualified Squeal.PostgreSQL as S


renderPGTWithClause :: PGT_AST.WithClause -> Q ([Text.Text], Exp)
renderPGTWithClause :: WithClause -> Q ([Text], Exp)
renderPGTWithClause (PGT_AST.WithClause Bool
recursive NonEmpty CommonTableExpr
ctes) = do
    Bool -> Q () -> Q ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
recursive (Q () -> Q ()) -> Q () -> Q ()
forall a b. (a -> b) -> a -> b
$ String -> Q ()
forall a. String -> Q a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"Recursive WITH clauses are not supported yet."
    let
      cteList :: [CommonTableExpr]
cteList = NonEmpty CommonTableExpr -> [CommonTableExpr]
forall a. NonEmpty a -> [a]
NE.toList NonEmpty CommonTableExpr
ctes
    (finalCteNames, renderedCtes) <-
      (([Text], [Exp]) -> CommonTableExpr -> Q ([Text], [Exp]))
-> ([Text], [Exp]) -> [CommonTableExpr] -> Q ([Text], [Exp])
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldlM
        ( \([Text]
names, [Exp]
exps) CommonTableExpr
cte -> do
            (name, exp) <- [Text] -> CommonTableExpr -> Q (Text, Exp)
renderCte [Text]
names CommonTableExpr
cte
            pure (names <> [name], exps <> [exp])
        )
        ([], [])
        [CommonTableExpr]
cteList

    let
      withExp =
        (Exp -> Exp -> Exp) -> Exp -> [Exp] -> Exp
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr
          (\Exp
cte Exp
acc -> Name -> Exp
ConE '(S.:>>) Exp -> Exp -> Exp
`AppE` Exp
cte Exp -> Exp -> Exp
`AppE` Exp
acc)
          (Name -> Exp
ConE 'S.Done)
          [Exp]
renderedCtes
    pure (finalCteNames, withExp)
  where
    renderCte :: [Text.Text] -> PGT_AST.CommonTableExpr -> Q (Text.Text, Exp)
    renderCte :: [Text] -> CommonTableExpr -> Q (Text, Exp)
renderCte [Text]
existingCteNames (PGT_AST.CommonTableExpr Ident
ident Maybe (NonEmpty Ident)
maybeColNames Maybe Bool
maybeMaterialized PreparableStmt
stmt) = do
      Bool -> Q () -> Q ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Maybe Bool -> Bool
forall a. Maybe a -> Bool
isJust Maybe Bool
maybeMaterialized) (Q () -> Q ()) -> Q () -> Q ()
forall a b. (a -> b) -> a -> b
$
        String -> Q ()
forall a. String -> Q a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"MATERIALIZED/NOT MATERIALIZED for CTEs is not supported yet."
      cteStmtExp <-
        case PreparableStmt
stmt of
          PGT_AST.SelectPreparableStmt SelectStmt
selectStmt -> do
            queryExp <- [Text] -> Maybe (NonEmpty Ident) -> SelectStmt -> Q Exp
toSquealQuery [Text]
existingCteNames Maybe (NonEmpty Ident)
maybeColNames SelectStmt
selectStmt
            pure $ VarE 'S.queryStatement `AppE` queryExp
          PreparableStmt
_ -> String -> Q Exp
forall a. String -> Q a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"Only SELECT statements are supported in CTEs for now."
      let
        cteName = Ident -> Text
getIdentText Ident
ident
      pure
        (cteName, VarE 'S.as `AppE` cteStmtExp `AppE` LabelE (Text.unpack cteName))


toSquealDelete :: PGT_AST.DeleteStmt -> Q Exp
toSquealDelete :: DeleteStmt -> Q Exp
toSquealDelete
  ( PGT_AST.DeleteStmt
      Maybe WithClause
maybeWithClause
      RelationExprOptAlias
relationExprOptAlias
      Maybe UsingClause
maybeUsingClause
      Maybe WhereOrCurrentClause
maybeWhereClause
      Maybe ReturningClause
maybeReturningClause
    ) = do
    (cteNames, renderedWithClause) <-
      case Maybe WithClause
maybeWithClause of
        Maybe WithClause
Nothing -> ([Text], Maybe Exp) -> Q ([Text], Maybe Exp)
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([], Maybe Exp
forall a. Maybe a
Nothing)
        Just WithClause
withClause -> do
          (names, exp) <- WithClause -> Q ([Text], Exp)
renderPGTWithClause WithClause
withClause
          pure (names, Just exp)
    -- whereOrCurrentClause with cursor is not supported yet

    targetTableExp <- renderPGTRelationExprOptAlias' relationExprOptAlias

    usingClauseExp <-
      case maybeUsingClause of
        Maybe UsingClause
Nothing -> Exp -> Q Exp
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> Q Exp) -> Exp -> Q Exp
forall a b. (a -> b) -> a -> b
$ Name -> Exp
ConE 'S.NoUsing
        Just UsingClause
usingClause -> Exp -> Exp -> Exp
AppE (Name -> Exp
ConE 'S.Using) (Exp -> Exp) -> Q Exp -> Q Exp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Text] -> UsingClause -> Q Exp
renderPGTTableRef [Text]
cteNames UsingClause
usingClause

    whereConditionExp <-
      case maybeWhereClause of
        Maybe WhereOrCurrentClause
Nothing ->
          String -> Q Exp
forall a. String -> Q a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail
            String
"DELETE statements must have a WHERE clause for safety. Use WHERE TRUE to delete all rows."
        Just (PGT_AST.ExprWhereOrCurrentClause AExpr
whereAExpr) -> [Text] -> AExpr -> Q Exp
renderPGTAExpr [Text]
cteNames AExpr
whereAExpr
        Just (PGT_AST.CursorWhereOrCurrentClause Ident
_) -> String -> Q Exp
forall a. String -> Q a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"WHERE CURRENT OF is not supported."

    returningClauseExp <-
      case maybeReturningClause of
        Maybe ReturningClause
Nothing -> Exp -> Q Exp
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> Q Exp) -> Exp -> Q Exp
forall a b. (a -> b) -> a -> b
$ Name -> Exp
ConE 'S.Returning_ Exp -> Exp -> Exp
`AppE` Name -> Exp
ConE 'S.Nil
        Just ReturningClause
returningClause -> Exp -> Exp -> Exp
AppE (Name -> Exp
ConE 'S.Returning) (Exp -> Exp) -> Q Exp -> Q Exp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Text] -> ReturningClause -> Q Exp
renderPGTTargetList [Text]
cteNames ReturningClause
returningClause

    let
      deleteBody =
        ( Name -> Exp
VarE 'S.deleteFrom
            Exp -> Exp -> Exp
`AppE` Exp
targetTableExp
            Exp -> Exp -> Exp
`AppE` Exp
usingClauseExp
            Exp -> Exp -> Exp
`AppE` Exp
whereConditionExp
            Exp -> Exp -> Exp
`AppE` Exp
returningClauseExp
        )
    let
      finalExp = case Maybe Exp
renderedWithClause of
        Maybe Exp
Nothing -> Exp
deleteBody
        Just Exp
withExp -> Name -> Exp
VarE 'S.with Exp -> Exp -> Exp
`AppE` Exp
withExp Exp -> Exp -> Exp
`AppE` Exp
deleteBody
    pure finalExp


renderPGTRelationExprOptAlias' :: PGT_AST.RelationExprOptAlias -> Q Exp
renderPGTRelationExprOptAlias' :: RelationExprOptAlias -> Q Exp
renderPGTRelationExprOptAlias' (PGT_AST.RelationExprOptAlias RelationExpr
relationExpr Maybe (Bool, Ident)
maybeAlias) = do
  (tableName, schemaName) <-
    case RelationExpr
relationExpr of
      PGT_AST.SimpleRelationExpr (PGT_AST.SimpleQualifiedName Ident
ident) Bool
_ ->
        (Text, Maybe Text) -> Q (Text, Maybe Text)
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Ident -> Text
getIdentText Ident
ident, Maybe Text
forall a. Maybe a
Nothing)
      PGT_AST.SimpleRelationExpr
        ( PGT_AST.IndirectedQualifiedName
            Ident
schemaIdent
            (PGT_AST.AttrNameIndirectionEl Ident
tableIdent NE.:| [])
          )
        Bool
_ ->
          (Text, Maybe Text) -> Q (Text, Maybe Text)
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Ident -> Text
getIdentText Ident
tableIdent, Text -> Maybe Text
forall a. a -> Maybe a
Just (Ident -> Text
getIdentText Ident
schemaIdent))
      RelationExpr
_ -> String -> Q (Text, Maybe Text)
forall a. String -> Q a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"Unsupported relation expression in DELETE statement."

  let
    aliasName = case Maybe (Bool, Ident)
maybeAlias of
      Just (Bool
_, Ident
colId) -> Ident -> Text
getIdentText Ident
colId
      Maybe (Bool, Ident)
Nothing -> Text
tableName

  let
    qualifiedAlias = case Maybe Text
schemaName of
      Maybe Text
Nothing -> String -> Exp
LabelE (Text -> String
Text.unpack Text
tableName)
      Just Text
schema ->
        Name -> Exp
VarE '(S.!)
          Exp -> Exp -> Exp
`AppE` String -> Exp
LabelE (Text -> String
Text.unpack Text
schema)
          Exp -> Exp -> Exp
`AppE` String -> Exp
LabelE (Text -> String
Text.unpack Text
tableName)

  pure $ VarE 'S.as `AppE` qualifiedAlias `AppE` LabelE (Text.unpack aliasName)