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

-- | Description: Translate update statements.
module Squeal.QuasiQuotes.Update (
  toSquealUpdate,
) where

import Data.Text (Text)
import Control.Monad (when)
import Data.Maybe (isJust)
import Language.Haskell.TH.Syntax (Exp(AppE, ConE, LabelE, VarE), Q)
import Prelude
  ( Applicative(pure), Foldable(foldr), Maybe(Just, Nothing), MonadFail(fail)
  , Traversable(mapM), ($), (<$>)
  )
import Squeal.QuasiQuotes.Common
  ( getIdentText, renderPGTAExpr, renderPGTTableRef, renderPGTTargetList
  )
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


toSquealUpdate :: PGT_AST.UpdateStmt -> Q Exp
toSquealUpdate :: UpdateStmt -> Q Exp
toSquealUpdate
  ( PGT_AST.UpdateStmt
      Maybe WithClause
maybeWithClause
      RelationExprOptAlias
relationExprOptAlias
      SetClauseList
setClauseList
      Maybe FromClause
maybeFromClause
      Maybe WhereOrCurrentClause
maybeWhereClause
      Maybe ReturningClause
maybeReturningClause
    ) = do
    Bool -> Q () -> Q ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Maybe WithClause -> Bool
forall a. Maybe a -> Bool
isJust Maybe WithClause
maybeWithClause) (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
"WITH clauses are not supported in UPDATE statements yet."
    targetTableExp <- RelationExprOptAlias -> Q Exp
renderPGTRelationExprOptAlias' RelationExprOptAlias
relationExprOptAlias

    setClauseExp <- renderPGTSetClauseList setClauseList

    usingClauseExp <-
      case maybeFromClause of
        Maybe FromClause
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 FromClause
fromClause -> 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
<$> FromClause -> Q Exp
renderPGTTableRef FromClause
fromClause

    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
"UPDATE statements must have a WHERE clause for safety. Use WHERE TRUE to update all rows."
        Just (PGT_AST.ExprWhereOrCurrentClause AExpr
whereAExpr) -> AExpr -> Q Exp
renderPGTAExpr AExpr
whereAExpr
        Just (PGT_AST.CursorWhereOrCurrentClause CursorName
_) -> 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
<$> ReturningClause -> Q Exp
renderPGTTargetList ReturningClause
returningClause

    pure $
      ( VarE 'S.update
          `AppE` targetTableExp
          `AppE` setClauseExp
          `AppE` usingClauseExp
          `AppE` whereConditionExp
          `AppE` returningClauseExp
      )


renderPGTRelationExprOptAlias' :: PGT_AST.RelationExprOptAlias -> Q Exp
renderPGTRelationExprOptAlias' :: RelationExprOptAlias -> Q Exp
renderPGTRelationExprOptAlias' (PGT_AST.RelationExprOptAlias RelationExpr
relationExpr Maybe (Bool, CursorName)
maybeAlias) = do
  (tableName, schemaName) <-
    case RelationExpr
relationExpr of
      PGT_AST.SimpleRelationExpr (PGT_AST.SimpleQualifiedName CursorName
ident) Bool
_ ->
        (Text, Maybe Text) -> Q (Text, Maybe Text)
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (CursorName -> Text
getIdentText CursorName
ident, Maybe Text
forall a. Maybe a
Nothing)
      PGT_AST.SimpleRelationExpr
        ( PGT_AST.IndirectedQualifiedName
            CursorName
schemaIdent
            (PGT_AST.AttrNameIndirectionEl CursorName
tableIdent NE.:| [])
          )
        Bool
_ ->
          (Text, Maybe Text) -> Q (Text, Maybe Text)
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (CursorName -> Text
getIdentText CursorName
tableIdent, Text -> Maybe Text
forall a. a -> Maybe a
Just (CursorName -> Text
getIdentText CursorName
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 UPDATE statement."

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

    qualifiedAlias :: Exp
    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)


renderPGTSetClauseList :: PGT_AST.SetClauseList -> Q Exp
renderPGTSetClauseList :: SetClauseList -> Q Exp
renderPGTSetClauseList SetClauseList
setClauses = do
  renderedItems <- (SetClause -> Q Exp) -> [SetClause] -> Q [Exp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM SetClause -> Q Exp
renderPGTSetClause (SetClauseList -> [SetClause]
forall a. NonEmpty a -> [a]
NE.toList SetClauseList
setClauses)
  pure $
    foldr
      (\Exp
item Exp
acc -> Name -> Exp
ConE '(S.:*) Exp -> Exp -> Exp
`AppE` Exp
item Exp -> Exp -> Exp
`AppE` Exp
acc)
      (ConE 'S.Nil)
      renderedItems


renderPGTSetClause :: PGT_AST.SetClause -> Q Exp
renderPGTSetClause :: SetClause -> Q Exp
renderPGTSetClause = \case
  PGT_AST.TargetSetClause (PGT_AST.SetTarget CursorName
colId Maybe (NonEmpty IndirectionEl)
maybeIndirection) AExpr
aExpr -> do
    Bool -> Q () -> Q ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Maybe (NonEmpty IndirectionEl) -> Bool
forall a. Maybe a -> Bool
isJust Maybe (NonEmpty IndirectionEl)
maybeIndirection) (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
"UPDATE SET with indirection (e.g., array access) is not supported."
    let colNameStr :: String
colNameStr = Text -> String
Text.unpack (CursorName -> Text
getIdentText CursorName
colId)
    renderedExpr <- AExpr -> Q Exp
renderPGTAExpr AExpr
aExpr
    pure $
      VarE 'S.as
        `AppE` (ConE 'S.Set `AppE` renderedExpr)
        `AppE` LabelE colNameStr
  PGT_AST.TargetListSetClause SetTargetList
_ AExpr
_ ->
    String -> Q Exp
forall a. String -> Q a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail
      String
"UPDATE with multiple SET targets (e.g. (col1, col2) = (val1, val2)) is not yet supported."