{-# language DuplicateRecordFields #-}
{-# language FlexibleContexts #-}
{-# language GADTs #-}
{-# language LambdaCase #-}
{-# language NamedFieldPuns #-}
{-# language OverloadedStrings #-}
{-# language RecordWildCards #-}
{-# language StandaloneKindSignatures #-}
{-# language StrictData #-}
{-# language TypeOperators #-}

module Rel8.Internal.Statement.OnConflict
  ( OnConflict(..)
  , Conflict (..)
  , Index (..)
  , Upsert(..)
  , ppOnConflict
  )
where

-- base
import Data.Kind ( Type )
import Prelude

-- opaleye
import qualified Opaleye.Internal.HaskellDB.Sql.Print as Opaleye
import qualified Opaleye.Internal.Sql as Opaleye

-- pretty
import Text.PrettyPrint ( Doc, (<+>), ($$), parens, text )

-- rel8
import Rel8.Internal.Expr ( Expr )
import Rel8.Internal.Expr.Opaleye (toPrimExpr)
import Rel8.Internal.Schema.Escape (escape)
import Rel8.Internal.Schema.Name ( Selects )
import Rel8.Internal.Schema.HTable (hfoldMap)
import Rel8.Internal.Schema.Table ( TableSchema(..) )
import Rel8.Internal.Statement.Set ( ppSet )
import Rel8.Internal.Statement.Where ( ppWhere )
import Rel8.Internal.Table ( Table, toColumns )
import Rel8.Internal.Table.Opaleye (attributes, view)


-- | 'OnConflict' represents the @ON CONFLICT@ clause of an @INSERT@
-- statement. This specifies what ought to happen when one or more of the
-- rows proposed for insertion conflict with an existing row in the table.
type OnConflict :: Type -> Type
data OnConflict exprs
  = Abort
    -- ^ Abort the transaction if there are conflicting rows (Postgres' default)
  | DoNothing (Maybe (Conflict exprs))
    -- ^ @ON CONFLICT DO NOTHING@, or @ON CONFLICT (...) DO NOTHING@ if an
    -- explicit conflict target  is supplied. Specifying a conflict target is
    -- essential when your table has has deferrable constraints — @ON
    -- CONFLICT@ can't work on deferrable constraints, so it's necessary
    -- to explicitly name one of its non-deferrable constraints in order to
    -- use @ON CONFLICT@.
  | DoUpdate (Upsert exprs)
    -- ^ @ON CONFLICT (...) DO UPDATE ...@


-- | The @ON CONFLICT (...) DO UPDATE@ clause of an @INSERT@ statement, also
-- known as \"upsert\".
--
-- When an existing row conflicts with a row proposed for insertion,
-- @ON CONFLICT DO UPDATE@ allows you to instead update this existing row. The
-- conflicting row proposed for insertion is then \"excluded\", but its values
-- can still be referenced from the @SET@ and @WHERE@ clauses of the @UPDATE@
-- statement.
--
-- Upsert in Postgres a \"conflict target\" to be specified — this is the
-- @UNIQUE@ index from conflicts with which we would like to recover. Indexes
-- are specified by listing the columns that comprise them along with an
-- optional predicate in the case of partial indexes.
type Upsert :: Type -> Type
data Upsert exprs where
  Upsert :: excluded ~ exprs =>
    { forall exprs. Upsert exprs -> Conflict exprs
conflict :: Conflict exprs
      -- ^ The conflict target to supply to @DO UPDATE@.
    , ()
set :: excluded -> exprs -> exprs
      -- ^ How to update each selected row.
    , ()
updateWhere :: excluded -> exprs -> Expr Bool
      -- ^ Which rows to select for update.
    }
    -> Upsert exprs


-- | Represents what PostgreSQL calls a
-- [@conflict_target@](https://www.postgresql.org/docs/current/sql-insert.html#SQL-ON-CONFLICT)
-- in an @ON CONFLICT@ clause of an @INSERT@ statement.
type Conflict :: Type -> Type
data Conflict exprs
  = OnConstraint String
  -- ^ Use a specific named constraint for the conflict target. This
  -- corresponds the the syntax @ON CONFLICT constraint@ in PostgreSQL.
  | OnIndex (Index exprs)
  -- ^ Have PostgreSQL perform what it calls _unique index inference_ by
  -- giving it a description of the target index.


-- | A description of the target unique index — its columns (and/or
-- expressions) and, in the case of partial indexes, a predicate.
type Index :: Type -> Type
data Index exprs where
  Index :: Table Expr index =>
    { ()
columns :: exprs -> index
      -- ^ The set of columns and/or expressions comprising the @UNIQUE@ index
    , forall exprs. Index exprs -> Maybe (exprs -> Expr Bool)
predicate :: Maybe (exprs -> Expr Bool)
      -- ^ An optional predicate used to specify a
      -- [partial index](https://www.postgresql.org/docs/current/indexes-partial.html).
    }
    -> Index exprs


ppOnConflict :: Selects names exprs => TableSchema names -> OnConflict exprs -> Doc
ppOnConflict :: forall names exprs.
Selects names exprs =>
TableSchema names -> OnConflict exprs -> Doc
ppOnConflict schema :: TableSchema names
schema@TableSchema {names
columns :: names
columns :: forall names. TableSchema names -> names
columns} = \case
  OnConflict exprs
Abort -> Doc
forall a. Monoid a => a
mempty
  DoNothing Maybe (Conflict exprs)
conflict -> String -> Doc
text String
"ON CONFLICT" Doc -> Doc -> Doc
<+> (Conflict exprs -> Doc) -> Maybe (Conflict exprs) -> Doc
forall m a. Monoid m => (a -> m) -> Maybe a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (exprs -> Conflict exprs -> Doc
forall exprs. exprs -> Conflict exprs -> Doc
ppConflict exprs
row) Maybe (Conflict exprs)
conflict Doc -> Doc -> Doc
<+> String -> Doc
text String
"DO NOTHING"
  DoUpdate Upsert exprs
upsert -> TableSchema names -> exprs -> Upsert exprs -> Doc
forall names exprs.
Selects names exprs =>
TableSchema names -> exprs -> Upsert exprs -> Doc
ppUpsert TableSchema names
schema exprs
row Upsert exprs
upsert
  where
    row :: exprs
row = names -> exprs
forall names exprs. Selects names exprs => names -> exprs
view names
columns


ppConflict :: exprs -> Conflict exprs -> Doc
ppConflict :: forall exprs. exprs -> Conflict exprs -> Doc
ppConflict exprs
row = \case
  OnConstraint String
name -> Doc
"ON CONSTRAINT" Doc -> Doc -> Doc
<+> String -> Doc
escape String
name
  OnIndex Index exprs
index -> exprs -> Index exprs -> Doc
forall exprs. exprs -> Index exprs -> Doc
ppIndex exprs
row Index exprs
index


ppIndex :: exprs -> Index exprs -> Doc
ppIndex :: forall exprs. exprs -> Index exprs -> Doc
ppIndex exprs
row Index {exprs -> index
columns :: ()
columns :: exprs -> index
columns, Maybe (exprs -> Expr Bool)
predicate :: forall exprs. Index exprs -> Maybe (exprs -> Expr Bool)
predicate :: Maybe (exprs -> Expr Bool)
predicate} =
  Doc -> Doc
parens ((Doc -> Doc) -> [Doc] -> Doc
forall a. (a -> Doc) -> [a] -> Doc
Opaleye.commaH Doc -> Doc
forall a. a -> a
id [Doc]
exprs) Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<>
  ((exprs -> Expr Bool) -> Doc) -> Maybe (exprs -> Expr Bool) -> Doc
forall m a. Monoid m => (a -> m) -> Maybe a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (Expr Bool -> Doc
ppPredicate (Expr Bool -> Doc)
-> ((exprs -> Expr Bool) -> Expr Bool)
-> (exprs -> Expr Bool)
-> Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((exprs -> Expr Bool) -> exprs -> Expr Bool
forall a b. (a -> b) -> a -> b
$ exprs
row)) Maybe (exprs -> Expr Bool)
predicate
  where
    exprs :: [Doc]
exprs = (forall a. Expr a -> [Doc]) -> Columns index Expr -> [Doc]
forall (t :: HTable) s (context :: * -> *).
(HTable t, Semigroup s) =>
(forall a. context a -> s) -> t context -> s
hfoldMap (Doc -> [Doc]
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Doc -> [Doc]) -> (Expr a -> Doc) -> Expr a -> [Doc]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Doc -> Doc
parens (Doc -> Doc) -> (Expr a -> Doc) -> Expr a -> Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Expr a -> Doc
forall a. Expr a -> Doc
ppExpr) (Columns index Expr -> [Doc]) -> Columns index Expr -> [Doc]
forall a b. (a -> b) -> a -> b
$ index -> Columns index Expr
forall (context :: * -> *) a.
Table context a =>
a -> Columns a context
toColumns (index -> Columns index Expr) -> index -> Columns index Expr
forall a b. (a -> b) -> a -> b
$ exprs -> index
columns exprs
row


ppPredicate :: Expr Bool -> Doc
ppPredicate :: Expr Bool -> Doc
ppPredicate Expr Bool
condition = String -> Doc
text String
"WHERE" Doc -> Doc -> Doc
<+> Expr Bool -> Doc
forall a. Expr a -> Doc
ppExpr Expr Bool
condition


ppUpsert :: Selects names exprs => TableSchema names -> exprs -> Upsert exprs -> Doc
ppUpsert :: forall names exprs.
Selects names exprs =>
TableSchema names -> exprs -> Upsert exprs -> Doc
ppUpsert schema :: TableSchema names
schema@TableSchema {names
columns :: forall names. TableSchema names -> names
columns :: names
columns} exprs
row Upsert {Conflict exprs
excluded -> exprs -> exprs
excluded -> exprs -> Expr Bool
conflict :: forall exprs. Upsert exprs -> Conflict exprs
set :: ()
updateWhere :: ()
conflict :: Conflict exprs
set :: excluded -> exprs -> exprs
updateWhere :: excluded -> exprs -> Expr Bool
..} =
  String -> Doc
text String
"ON CONFLICT" Doc -> Doc -> Doc
<+> exprs -> Conflict exprs -> Doc
forall exprs. exprs -> Conflict exprs -> Doc
ppConflict exprs
row Conflict exprs
conflict Doc -> Doc -> Doc
<+> Doc
"DO UPDATE" Doc -> Doc -> Doc
$$
  TableSchema names -> (exprs -> exprs) -> Doc
forall names exprs.
Selects names exprs =>
TableSchema names -> (exprs -> exprs) -> Doc
ppSet TableSchema names
schema (excluded -> exprs -> exprs
set excluded
excluded) Doc -> Doc -> Doc
$$
  TableSchema names -> (exprs -> Expr Bool) -> Doc
forall names exprs.
Selects names exprs =>
TableSchema names -> (exprs -> Expr Bool) -> Doc
ppWhere TableSchema names
schema (excluded -> exprs -> Expr Bool
updateWhere excluded
excluded)
  where
    excluded :: excluded
excluded = TableSchema names -> excluded
forall names exprs.
Selects names exprs =>
TableSchema names -> exprs
attributes TableSchema
      { name :: QualifiedName
name = QualifiedName
"excluded"
      , names
columns :: names
columns :: names
columns
      }


ppExpr :: Expr a -> Doc
ppExpr :: forall a. Expr a -> Doc
ppExpr = SqlExpr -> Doc
Opaleye.ppSqlExpr (SqlExpr -> Doc) -> (Expr a -> SqlExpr) -> Expr a -> Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PrimExpr -> SqlExpr
Opaleye.sqlExpr (PrimExpr -> SqlExpr) -> (Expr a -> PrimExpr) -> Expr a -> SqlExpr
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Expr a -> PrimExpr
forall a. Expr a -> PrimExpr
toPrimExpr