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

-- | Description: Translate insert statements.
module Squeal.QuasiQuotes.Insert (
  toSquealInsert,
) where

import Control.Monad (MonadFail(fail), mapM, when, zipWithM)
import Data.Maybe (isJust)
import Language.Haskell.TH.Syntax (Exp(AppE, ConE, LabelE, ListE, VarE), Q)
import Prelude
  ( Applicative(pure), Either(Left), Eq((/=)), Foldable(foldr, length)
  , Maybe(Just, Nothing), Semigroup((<>)), Show(show), ($), (.), error
  , otherwise
  )
import Squeal.QuasiQuotes.Common (getIdentText, renderPGTAExpr)
import Squeal.QuasiQuotes.Query (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


toSquealInsert :: PGT_AST.InsertStmt -> Q Exp
toSquealInsert :: InsertStmt -> Q Exp
toSquealInsert
  ( PGT_AST.InsertStmt
      Maybe WithClause
maybeWithClause
      InsertTarget
insertTarget
      InsertRest
insertRest
      Maybe OnConflict
maybeOnConflict
      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
$
      [Char] -> Q ()
forall a. [Char] -> Q a
forall (m :: * -> *) a. MonadFail m => [Char] -> m a
fail [Char]
"WITH clauses are not supported in INSERT statements yet."
    Bool -> Q () -> Q ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Maybe OnConflict -> Bool
forall a. Maybe a -> Bool
isJust Maybe OnConflict
maybeOnConflict) (Q () -> Q ()) -> Q () -> Q ()
forall a b. (a -> b) -> a -> b
$
      [Char] -> Q ()
forall a. [Char] -> Q a
forall (m :: * -> *) a. MonadFail m => [Char] -> m a
fail [Char]
"ON CONFLICT clauses are not supported yet."

    insertBody <-
      case InsertRest
insertRest of
        PGT_AST.SelectInsertRest Maybe InsertColumnList
maybeInsertColumnList Maybe OverrideKind
maybeOverrideKind SelectStmt
selectStmt -> do
          Bool -> Q () -> Q ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Maybe OverrideKind -> Bool
forall a. Maybe a -> Bool
isJust Maybe OverrideKind
maybeOverrideKind) (Q () -> Q ()) -> Q () -> Q ()
forall a b. (a -> b) -> a -> b
$
            [Char] -> Q ()
forall a. [Char] -> Q a
forall (m :: * -> *) a. MonadFail m => [Char] -> m a
fail [Char]
"OVERRIDING clause is not supported yet."
          queryClauseExp <-
            case SelectStmt
selectStmt of
              -- Case 1: INSERT ... VALUES ...
              Left
                (PGT_AST.SelectNoParens Maybe WithClause
_ (Left (PGT_AST.ValuesSimpleSelect ValuesClause
valuesClause)) Maybe SortClause
_ Maybe SelectLimit
_ Maybe ForLockingClause
_) ->
                  case Maybe InsertColumnList
maybeInsertColumnList of
                    Just InsertColumnList
colItems ->
                      [InsertColumnItem] -> ValuesClause -> Q Exp
renderPGTValueRows (InsertColumnList -> [InsertColumnItem]
forall a. NonEmpty a -> [a]
NE.toList InsertColumnList
colItems) ValuesClause
valuesClause
                    Maybe InsertColumnList
Nothing ->
                      [Char] -> Q Exp
forall a. [Char] -> Q a
forall (m :: * -> *) a. MonadFail m => [Char] -> m a
fail
                        [Char]
"INSERT INTO ... VALUES must specify column names for the Squeal-QQ translation."
              -- Case 2: INSERT ... SELECT ... (a general SELECT statement)
              SelectStmt
_ ->
                -- selectStmt is not a ValuesSimpleSelect (i.e., it's a general query)
                case Maybe InsertColumnList
maybeInsertColumnList of
                  Just InsertColumnList
_ ->
                    [Char] -> Q Exp
forall a. [Char] -> Q a
forall (m :: * -> *) a. MonadFail m => [Char] -> m a
fail
                      [Char]
"INSERT INTO table (columns) SELECT ... is not yet supported by Squeal-QQ. Please use INSERT INTO table SELECT ... and ensure your SELECT statement provides all columns for the table, matching the table's column order and types."
                  Maybe InsertColumnList
Nothing -> do
                    squealQueryExp <- SelectStmt -> Q Exp
toSquealQuery SelectStmt
selectStmt -- from Squeal.QuasiQuotes.Query
                    pure (ConE 'S.Subquery `AppE` squealQueryExp)
          let
            table = InsertTarget -> Exp
renderPGTInsertTarget InsertTarget
insertTarget
          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
VarE 'S.insertInto_ Exp -> Exp -> Exp
`AppE` Exp
table Exp -> Exp -> Exp
`AppE` Exp
queryClauseExp
            Just (ReturningClause -> [TargetEl]
forall a. NonEmpty a -> [a]
NE.toList -> [TargetEl
PGT_AST.AsteriskTargetEl]) -> do
              let
                returning :: Exp
returning = Name -> Exp
ConE 'S.Returning Exp -> Exp -> Exp
`AppE` Name -> Exp
ConE 'S.Star
              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
VarE 'S.insertInto
                  Exp -> Exp -> Exp
`AppE` Exp
table
                  Exp -> Exp -> Exp
`AppE` Exp
queryClauseExp
                  Exp -> Exp -> Exp
`AppE` Name -> Exp
ConE 'S.OnConflictDoRaise
                  Exp -> Exp -> Exp
`AppE` Exp
returning
            Just ReturningClause
targetList -> do
              returningProj <- [TargetEl] -> Q Exp
renderTargetList (ReturningClause -> [TargetEl]
forall a. NonEmpty a -> [a]
NE.toList ReturningClause
targetList)
              let
                returning = Name -> Exp
ConE 'S.Returning Exp -> Exp -> Exp
`AppE` (Name -> Exp
ConE 'S.List Exp -> Exp -> Exp
`AppE` Exp
returningProj)
              pure $
                VarE 'S.insertInto
                  `AppE` table
                  `AppE` queryClauseExp
                  `AppE` ConE 'S.OnConflictDoRaise
                  `AppE` returning
        InsertRest
PGT_AST.DefaultValuesInsertRest ->
          [Char] -> Q Exp
forall a. [Char] -> Q a
forall (m :: * -> *) a. MonadFail m => [Char] -> m a
fail [Char]
"INSERT INTO ... DEFAULT VALUES is not yet supported by Squeal-QQ."

    pure insertBody


renderTargetList :: [PGT_AST.TargetEl] -> Q Exp
renderTargetList :: [TargetEl] -> Q Exp
renderTargetList [TargetEl]
targetEls = do
  exps <- (TargetEl -> Q Exp) -> [TargetEl] -> 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 TargetEl -> Q Exp
renderTargetEl [TargetEl]
targetEls
  pure $
    foldr
      (\Exp
h Exp
t -> Name -> Exp
ConE '(S.:*) Exp -> Exp -> Exp
`AppE` Exp
h Exp -> Exp -> Exp
`AppE` Exp
t)
      (ConE 'S.Nil)
      exps


renderTargetEl :: PGT_AST.TargetEl -> Q Exp
renderTargetEl :: TargetEl -> Q Exp
renderTargetEl = \case
  PGT_AST.ExprTargetEl AExpr
expr -> do
    exprExp <- AExpr -> Q Exp
renderPGTAExpr AExpr
expr
    case expr of
      PGT_AST.CExprAExpr (PGT_AST.ColumnrefCExpr (PGT_AST.Columnref ColId
ident Maybe Indirection
Nothing)) ->
        let
          colName :: Text
colName = ColId -> Text
getIdentText ColId
ident
        in
          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
VarE 'S.as Exp -> Exp -> Exp
`AppE` Exp
exprExp Exp -> Exp -> Exp
`AppE` [Char] -> Exp
LabelE (Text -> [Char]
Text.unpack Text
colName)
      AExpr
_ ->
        [Char] -> Q Exp
forall a. [Char] -> Q a
forall (m :: * -> *) a. MonadFail m => [Char] -> m a
fail
          [Char]
"Returning expression without an alias is not supported for this expression type. Please add an alias."
  PGT_AST.AliasedExprTargetEl AExpr
expr ColId
alias -> do
    exprExp <- AExpr -> Q Exp
renderPGTAExpr AExpr
expr
    pure $
      VarE 'S.as `AppE` exprExp `AppE` LabelE (Text.unpack (getIdentText alias))
  TargetEl
PGT_AST.AsteriskTargetEl -> [Char] -> Q Exp
forall a. [Char] -> Q a
forall (m :: * -> *) a. MonadFail m => [Char] -> m a
fail [Char]
"should be handled by toSquealInsert"
  PGT_AST.ImplicitlyAliasedExprTargetEl{} ->
    [Char] -> Q Exp
forall a. [Char] -> Q a
forall (m :: * -> *) a. MonadFail m => [Char] -> m a
fail [Char]
"Implicitly aliased expressions in RETURNING are not supported"


getUnqualifiedNameFromAst :: PGT_AST.QualifiedName -> Text.Text
getUnqualifiedNameFromAst :: QualifiedName -> Text
getUnqualifiedNameFromAst (PGT_AST.SimpleQualifiedName ColId
ident) = ColId -> Text
getIdentText ColId
ident
getUnqualifiedNameFromAst (PGT_AST.IndirectedQualifiedName ColId
_ (PGT_AST.AttrNameIndirectionEl ColId
ident NE.:| [])) = ColId -> Text
getIdentText ColId
ident
getUnqualifiedNameFromAst QualifiedName
unsupported =
  [Char] -> Text
forall a. HasCallStack => [Char] -> a
error ([Char] -> Text) -> [Char] -> Text
forall a b. (a -> b) -> a -> b
$
    [Char]
"Unsupported qualified name structure for extracting unqualified name: "
      [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> QualifiedName -> [Char]
forall a. Show a => a -> [Char]
show QualifiedName
unsupported


renderPGTInsertTarget :: PGT_AST.InsertTarget -> Exp
renderPGTInsertTarget :: InsertTarget -> Exp
renderPGTInsertTarget (PGT_AST.InsertTarget QualifiedName
qualifiedName Maybe ColId
maybeAsAlias) =
  let
    tableIdentifierExp :: Exp
tableIdentifierExp = QualifiedName -> Exp
renderPGTQualifiedName QualifiedName
qualifiedName
    targetAliasText :: Text
targetAliasText = case Maybe ColId
maybeAsAlias of
      Maybe ColId
Nothing -> QualifiedName -> Text
getUnqualifiedNameFromAst QualifiedName
qualifiedName
      Just ColId
asAliasColId -> ColId -> Text
getIdentText ColId
asAliasColId
    targetAliasExp :: Exp
targetAliasExp = [Char] -> Exp
LabelE (Text -> [Char]
Text.unpack Text
targetAliasText)
  in
    Name -> Exp
VarE 'S.as Exp -> Exp -> Exp
`AppE` Exp
tableIdentifierExp Exp -> Exp -> Exp
`AppE` Exp
targetAliasExp


renderPGTQualifiedName :: PGT_AST.QualifiedName -> Exp
renderPGTQualifiedName :: QualifiedName -> Exp
renderPGTQualifiedName = \case
  PGT_AST.SimpleQualifiedName ColId
ident -> [Char] -> Exp
LabelE (Text -> [Char]
Text.unpack (ColId -> Text
getIdentText ColId
ident)) -- Defaults to public schema
  PGT_AST.IndirectedQualifiedName
    ColId
schemaIdent
    (PGT_AST.AttrNameIndirectionEl ColId
colIdent NE.:| []) ->
      -- Assuming simple schema.table form
      Name -> Exp
VarE '(S.!)
        Exp -> Exp -> Exp
`AppE` [Char] -> Exp
LabelE (Text -> [Char]
Text.unpack (ColId -> Text
getIdentText ColId
schemaIdent))
        Exp -> Exp -> Exp
`AppE` [Char] -> Exp
LabelE (Text -> [Char]
Text.unpack (ColId -> Text
getIdentText ColId
colIdent))
  QualifiedName
unsupported -> [Char] -> Exp
forall a. HasCallStack => [Char] -> a
error ([Char] -> Exp) -> [Char] -> Exp
forall a b. (a -> b) -> a -> b
$ [Char]
"Unsupported qualified name structure: " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> QualifiedName -> [Char]
forall a. Show a => a -> [Char]
show QualifiedName
unsupported


renderPGTValueRows
  :: [PGT_AST.InsertColumnItem] -> PGT_AST.ValuesClause -> Q Exp
renderPGTValueRows :: [InsertColumnItem] -> ValuesClause -> Q Exp
renderPGTValueRows [InsertColumnItem]
colItems (ValuesClause
valuesClauseRows) =
  -- valuesClauseRows is NonEmpty (NonEmpty AExpr)
  case ValuesClause -> [ExprList]
forall a. NonEmpty a -> [a]
NE.toList ValuesClause
valuesClauseRows of
    [] -> [Char] -> Q Exp
forall a. [Char] -> Q a
forall (m :: * -> *) a. MonadFail m => [Char] -> m a
fail [Char]
"Insert statement has no value rows."
    ExprList
row : [ExprList]
moreRows -> do
      firstRowExp <- [InsertColumnItem] -> [AExpr] -> Q Exp
renderPGTValueRow [InsertColumnItem]
colItems (ExprList -> [AExpr]
forall a. NonEmpty a -> [a]
NE.toList ExprList
row)
      moreRowsExp <- mapM (renderPGTValueRow colItems . NE.toList) moreRows
      pure $
        ConE 'S.Values
          `AppE` firstRowExp
          `AppE` ListE moreRowsExp


renderPGTValueRow :: [PGT_AST.InsertColumnItem] -> [PGT_AST.AExpr] -> Q Exp
renderPGTValueRow :: [InsertColumnItem] -> [AExpr] -> Q Exp
renderPGTValueRow [InsertColumnItem]
colItems [AExpr]
exprs
    | [InsertColumnItem] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [InsertColumnItem]
colItems Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= [AExpr] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [AExpr]
exprs =
        [Char] -> Q Exp
forall a. [Char] -> Q a
forall (m :: * -> *) a. MonadFail m => [Char] -> m a
fail [Char]
"Mismatched number of column names and values in INSERT statement."
    | Bool
otherwise = do
        processedItems <- (InsertColumnItem -> AExpr -> Q Exp)
-> [InsertColumnItem] -> [AExpr] -> Q [Exp]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM InsertColumnItem -> AExpr -> Q Exp
processItem [InsertColumnItem]
colItems [AExpr]
exprs
        pure $ foldr nvpToCons (ConE 'S.Nil) processedItems
  where
    nvpToCons :: Exp -> Exp -> Exp
    nvpToCons :: Exp -> Exp -> Exp
nvpToCons Exp
item Exp
acc = Name -> Exp
ConE '(S.:*) Exp -> Exp -> Exp
`AppE` Exp
item Exp -> Exp -> Exp
`AppE` Exp
acc

    processItem :: PGT_AST.InsertColumnItem -> PGT_AST.AExpr -> Q Exp
    processItem :: InsertColumnItem -> AExpr -> Q Exp
processItem (PGT_AST.InsertColumnItem ColId
colId Maybe Indirection
maybeIndirection) AExpr
expr = do
      Bool -> Q () -> Q ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Maybe Indirection -> Bool
forall a. Maybe a -> Bool
isJust Maybe Indirection
maybeIndirection) (Q () -> Q ()) -> Q () -> Q ()
forall a b. (a -> b) -> a -> b
$
        [Char] -> Q ()
forall a. [Char] -> Q a
forall (m :: * -> *) a. MonadFail m => [Char] -> m a
fail [Char]
"INSERT with indirection (e.g., array access) is not supported."
      let
        colNameStr :: [Char]
colNameStr = Text -> [Char]
Text.unpack (ColId -> Text
getIdentText ColId
colId)
      case AExpr
expr of
        AExpr
PGT_AST.DefaultAExpr ->
          -- Check for DEFAULT keyword
          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
VarE 'S.as
              Exp -> Exp -> Exp
`AppE` Name -> Exp
ConE 'S.Default
              Exp -> Exp -> Exp
`AppE` [Char] -> Exp
LabelE [Char]
colNameStr
        AExpr
_ -> do
          renderedExpr <- AExpr -> Q Exp
renderPGTAExpr AExpr
expr
          pure $
            VarE 'S.as
              `AppE` (ConE 'S.Set `AppE` renderedExpr)
              `AppE` LabelE colNameStr