{-|
Module: Squeal.PostgreSQL.Query.With
Description: with statements
Copyright: (c) Eitan Chatav, 2019
Maintainer: eitan@morphism.tech
Stability: experimental

with statements
-}

{-# LANGUAGE
    ConstraintKinds
  , DeriveGeneric
  , DerivingStrategies
  , FlexibleContexts
  , FlexibleInstances
  , GADTs
  , GeneralizedNewtypeDeriving
  , LambdaCase
  , MultiParamTypeClasses
  , OverloadedLabels
  , OverloadedStrings
  , QuantifiedConstraints
  , ScopedTypeVariables
  , StandaloneDeriving
  , TypeApplications
  , TypeFamilies
  , TypeInType
  , TypeOperators
  , RankNTypes
  , UndecidableInstances
  #-}

module Squeal.PostgreSQL.Query.With
  ( -- ** With
    With (..)
  , CommonTableExpression (..)
  , withRecursive
  , Materialization (..)
  , materialized
  , notMaterialized
  ) where

import Data.Quiver.Functor
import GHC.TypeLits

import qualified GHC.Generics as GHC
import qualified Generics.SOP as SOP

import Squeal.PostgreSQL.Type.Alias
import Squeal.PostgreSQL.Query
import Squeal.PostgreSQL.Type.List
import Squeal.PostgreSQL.Render
import Squeal.PostgreSQL.Type.Schema

-- $setup
-- >>> import Squeal.PostgreSQL

{- | `with` provides a way to write auxiliary statements for use in a larger query.
These statements, referred to as `CommonTableExpression`s, can be thought of as
defining temporary tables that exist just for one query.

`with` can be used for a `Query`. Multiple `CommonTableExpression`s can be
chained together with the `Path` constructor `:>>`, and each `CommonTableExpression`
is constructed via overloaded `as`.

>>> type Columns = '["col1" ::: 'NoDef :=> 'NotNull 'PGint4, "col2" ::: 'NoDef :=> 'NotNull 'PGint4]
>>> type Schema = '["tab" ::: 'Table ('[] :=> Columns)]
>>> :{
let
  qry :: Query lat with (Public Schema) params '["col1" ::: 'NotNull 'PGint4, "col2" ::: 'NotNull 'PGint4]
  qry = with (
    select Star (from (table #tab)) `as` #cte1 :>>
    select Star (from (common #cte1)) `as` #cte2
    ) (select Star (from (common #cte2)))
in printSQL qry
:}
WITH "cte1" AS (SELECT * FROM "tab" AS "tab"), "cte2" AS (SELECT * FROM "cte1" AS "cte1") SELECT * FROM "cte2" AS "cte2"

You can use data-modifying statements in `with`. This allows you to perform several
different operations in the same query. An example is:

>>> type ProductsColumns = '["product" ::: 'NoDef :=> 'NotNull 'PGtext, "date" ::: 'Def :=> 'NotNull 'PGdate]
>>> type ProductsSchema = '["products" ::: 'Table ('[] :=> ProductsColumns), "products_deleted" ::: 'Table ('[] :=> ProductsColumns)]
>>> :{
let
  manp :: Manipulation with (Public ProductsSchema) '[ 'NotNull 'PGdate] '[]
  manp = with
    (deleteFrom #products NoUsing (#date .< param @1) (Returning Star) `as` #del)
    (insertInto_ #products_deleted (Subquery (select Star (from (common #del)))))
in printSQL manp
:}
WITH "del" AS (DELETE FROM "products" AS "products" WHERE ("date" < ($1 :: date)) RETURNING *) INSERT INTO "products_deleted" AS "products_deleted" SELECT * FROM "del" AS "del"
-}
class With statement where
  with
    :: Path (CommonTableExpression statement db params) with0 with1
    -- ^ common table expressions
    -> statement with1 db params row
    -- ^ larger query
    -> statement with0 db params row
instance With (Query lat) where
  with :: Path (CommonTableExpression (Query lat) db params) with0 with1
-> Query lat with1 db params row -> Query lat with0 db params row
with Path (CommonTableExpression (Query lat) db params) with0 with1
Done Query lat with1 db params row
query = Query lat with0 db params row
Query lat with1 db params row
query
  with Path (CommonTableExpression (Query lat) db params) with0 with1
ctes Query lat with1 db params row
query = ByteString -> Query lat with0 db params row
forall (lat :: FromType) (with :: FromType) (db :: SchemasType)
       (params :: [NullType]) (row :: RowType).
ByteString -> Query lat with db params row
UnsafeQuery (ByteString -> Query lat with0 db params row)
-> ByteString -> Query lat with0 db params row
forall a b. (a -> b) -> a -> b
$
    ByteString
"WITH" ByteString -> ByteString -> ByteString
<+> [ByteString] -> ByteString
commaSeparated ((forall (x1 :: FromType) (y1 :: FromType).
 CommonTableExpression (Query lat) db params x1 y1 -> ByteString)
-> Path (CommonTableExpression (Query lat) db params) with0 with1
-> [ByteString]
forall k (c :: (k -> k -> *) -> k -> k -> *) (p :: k -> k -> *) a
       (x :: k) (y :: k).
QFoldable c =>
(forall (x1 :: k) (y1 :: k). p x1 y1 -> a) -> c p x y -> [a]
qtoList forall (x1 :: FromType) (y1 :: FromType).
CommonTableExpression (Query lat) db params x1 y1 -> ByteString
forall sql. RenderSQL sql => sql -> ByteString
renderSQL Path (CommonTableExpression (Query lat) db params) with0 with1
ctes) ByteString -> ByteString -> ByteString
<+> Query lat with1 db params row -> ByteString
forall sql. RenderSQL sql => sql -> ByteString
renderSQL Query lat with1 db params row
query

{- | A `withRecursive` `Query` can refer to its own output.
A very simple example is this query to sum the integers from 1 through 100:

>>> import Data.Monoid (Sum (..))
>>> import Data.Int (Int64)
>>> :{
  let
    sum100 :: Statement db () (Sum Int64)
    sum100 = query $
      withRecursive
        ( values_ ((1 & astype int) `as` #n)
          `unionAll`
          select_ ((#n + 1) `as` #n)
            (from (common #t) & where_ (#n .< 100)) `as` #t )
        ( select_
            (fromNull 0 (sum_ (All #n)) `as` #getSum)
            (from (common #t) & groupBy Nil) )
  in printSQL sum100
:}
WITH RECURSIVE "t" AS ((SELECT * FROM (VALUES (((1 :: int4) :: int))) AS t ("n")) UNION ALL (SELECT ("n" + (1 :: int4)) AS "n" FROM "t" AS "t" WHERE ("n" < (100 :: int4)))) SELECT COALESCE(sum(ALL "n"), (0 :: int8)) AS "getSum" FROM "t" AS "t"

The general form of a recursive WITH query is always a non-recursive term,
then `union` (or `unionAll`), then a recursive term, where
only the recursive term can contain a reference to the query's own output.
-}
withRecursive
  :: Aliased (Query lat (recursive ': with) db params) recursive
  -- ^ recursive query
  -> Query lat (recursive ': with) db params row
  -- ^ larger query
  -> Query lat with db params row
withRecursive :: Aliased (Query lat (recursive : with) db params) recursive
-> Query lat (recursive : with) db params row
-> Query lat with db params row
withRecursive (Query lat (recursive : with) db params ty
recursive `As` Alias alias
cte) Query lat (recursive : with) db params row
query = ByteString -> Query lat with db params row
forall (lat :: FromType) (with :: FromType) (db :: SchemasType)
       (params :: [NullType]) (row :: RowType).
ByteString -> Query lat with db params row
UnsafeQuery (ByteString -> Query lat with db params row)
-> ByteString -> Query lat with db params row
forall a b. (a -> b) -> a -> b
$
  ByteString
"WITH RECURSIVE" ByteString -> ByteString -> ByteString
<+> Alias alias -> ByteString
forall sql. RenderSQL sql => sql -> ByteString
renderSQL Alias alias
cte
    ByteString -> ByteString -> ByteString
<+> ByteString
"AS" ByteString -> ByteString -> ByteString
<+> ByteString -> ByteString
parenthesized (Query lat (recursive : with) db params ty -> ByteString
forall sql. RenderSQL sql => sql -> ByteString
renderSQL Query lat (recursive : with) db params ty
recursive)
    ByteString -> ByteString -> ByteString
<+> Query lat (recursive : with) db params row -> ByteString
forall sql. RenderSQL sql => sql -> ByteString
renderSQL Query lat (recursive : with) db params row
query

-- | Whether the contents of the WITH clause are materialized.
-- If a WITH query is non-recursive and side-effect-free (that is, it is a SELECT containing no volatile functions) then it can be folded into the parent query, allowing joint optimization of the two query levels.
--
-- Note: Use of `Materialized` or `NotMaterialized` requires PostgreSQL version 12 or higher. For earlier versions, use `DefaultMaterialization` which in those earlier versions of PostgreSQL behaves as `Materialized`. PostgreSQL 12 both changes the default behavior as well as adds options for customizing the materialization behavior.
data Materialization =
  DefaultMaterialization -- ^ By default, folding happens if the parent query references the WITH query just once, but not if it references the WITH query more than once. Note: this is the behavior in PostgreSQL 12+. In PostgreSQL 11 and earlier, all CTEs are materialized.
  | Materialized -- ^ You can override that decision by specifying MATERIALIZED to force separate calculation of the WITH query. Requires PostgreSQL 12+.
  | NotMaterialized -- ^ or by specifying NOT MATERIALIZED to force it to be merged into the parent query. Requires PostgreSQL 12+.
  deriving (Materialization -> Materialization -> Bool
(Materialization -> Materialization -> Bool)
-> (Materialization -> Materialization -> Bool)
-> Eq Materialization
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Materialization -> Materialization -> Bool
$c/= :: Materialization -> Materialization -> Bool
== :: Materialization -> Materialization -> Bool
$c== :: Materialization -> Materialization -> Bool
Eq, Eq Materialization
Eq Materialization
-> (Materialization -> Materialization -> Ordering)
-> (Materialization -> Materialization -> Bool)
-> (Materialization -> Materialization -> Bool)
-> (Materialization -> Materialization -> Bool)
-> (Materialization -> Materialization -> Bool)
-> (Materialization -> Materialization -> Materialization)
-> (Materialization -> Materialization -> Materialization)
-> Ord Materialization
Materialization -> Materialization -> Bool
Materialization -> Materialization -> Ordering
Materialization -> Materialization -> Materialization
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: Materialization -> Materialization -> Materialization
$cmin :: Materialization -> Materialization -> Materialization
max :: Materialization -> Materialization -> Materialization
$cmax :: Materialization -> Materialization -> Materialization
>= :: Materialization -> Materialization -> Bool
$c>= :: Materialization -> Materialization -> Bool
> :: Materialization -> Materialization -> Bool
$c> :: Materialization -> Materialization -> Bool
<= :: Materialization -> Materialization -> Bool
$c<= :: Materialization -> Materialization -> Bool
< :: Materialization -> Materialization -> Bool
$c< :: Materialization -> Materialization -> Bool
compare :: Materialization -> Materialization -> Ordering
$ccompare :: Materialization -> Materialization -> Ordering
$cp1Ord :: Eq Materialization
Ord, Int -> Materialization -> ShowS
[Materialization] -> ShowS
Materialization -> String
(Int -> Materialization -> ShowS)
-> (Materialization -> String)
-> ([Materialization] -> ShowS)
-> Show Materialization
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Materialization] -> ShowS
$cshowList :: [Materialization] -> ShowS
show :: Materialization -> String
$cshow :: Materialization -> String
showsPrec :: Int -> Materialization -> ShowS
$cshowsPrec :: Int -> Materialization -> ShowS
Show, ReadPrec [Materialization]
ReadPrec Materialization
Int -> ReadS Materialization
ReadS [Materialization]
(Int -> ReadS Materialization)
-> ReadS [Materialization]
-> ReadPrec Materialization
-> ReadPrec [Materialization]
-> Read Materialization
forall a.
(Int -> ReadS a)
-> ReadS [a] -> ReadPrec a -> ReadPrec [a] -> Read a
readListPrec :: ReadPrec [Materialization]
$creadListPrec :: ReadPrec [Materialization]
readPrec :: ReadPrec Materialization
$creadPrec :: ReadPrec Materialization
readList :: ReadS [Materialization]
$creadList :: ReadS [Materialization]
readsPrec :: Int -> ReadS Materialization
$creadsPrec :: Int -> ReadS Materialization
Read, Int -> Materialization
Materialization -> Int
Materialization -> [Materialization]
Materialization -> Materialization
Materialization -> Materialization -> [Materialization]
Materialization
-> Materialization -> Materialization -> [Materialization]
(Materialization -> Materialization)
-> (Materialization -> Materialization)
-> (Int -> Materialization)
-> (Materialization -> Int)
-> (Materialization -> [Materialization])
-> (Materialization -> Materialization -> [Materialization])
-> (Materialization -> Materialization -> [Materialization])
-> (Materialization
    -> Materialization -> Materialization -> [Materialization])
-> Enum Materialization
forall a.
(a -> a)
-> (a -> a)
-> (Int -> a)
-> (a -> Int)
-> (a -> [a])
-> (a -> a -> [a])
-> (a -> a -> [a])
-> (a -> a -> a -> [a])
-> Enum a
enumFromThenTo :: Materialization
-> Materialization -> Materialization -> [Materialization]
$cenumFromThenTo :: Materialization
-> Materialization -> Materialization -> [Materialization]
enumFromTo :: Materialization -> Materialization -> [Materialization]
$cenumFromTo :: Materialization -> Materialization -> [Materialization]
enumFromThen :: Materialization -> Materialization -> [Materialization]
$cenumFromThen :: Materialization -> Materialization -> [Materialization]
enumFrom :: Materialization -> [Materialization]
$cenumFrom :: Materialization -> [Materialization]
fromEnum :: Materialization -> Int
$cfromEnum :: Materialization -> Int
toEnum :: Int -> Materialization
$ctoEnum :: Int -> Materialization
pred :: Materialization -> Materialization
$cpred :: Materialization -> Materialization
succ :: Materialization -> Materialization
$csucc :: Materialization -> Materialization
Enum, (forall x. Materialization -> Rep Materialization x)
-> (forall x. Rep Materialization x -> Materialization)
-> Generic Materialization
forall x. Rep Materialization x -> Materialization
forall x. Materialization -> Rep Materialization x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep Materialization x -> Materialization
$cfrom :: forall x. Materialization -> Rep Materialization x
GHC.Generic)
instance SOP.Generic Materialization
instance SOP.HasDatatypeInfo Materialization
instance RenderSQL Materialization where
  renderSQL :: Materialization -> ByteString
renderSQL = \case
    Materialization
DefaultMaterialization -> ByteString
""
    Materialization
Materialized -> ByteString
"MATERIALIZED"
    Materialization
NotMaterialized -> ByteString
"NOT MATERIALIZED"

-- | A `CommonTableExpression` is an auxiliary statement in a `with` clause.
data CommonTableExpression statement
  (db :: SchemasType)
  (params :: [NullType])
  (with0 :: FromType)
  (with1 :: FromType) where
  CommonTableExpression
    :: Aliased (statement with db params) (cte ::: common)
    -- ^ aliased statement
    -> Materialization
    -- ^ materialization of the CTE output
    -> CommonTableExpression statement db params with (cte ::: common ': with)
instance
  ( KnownSymbol cte
  , with1 ~ (cte ::: common ': with)
  ) => Aliasable cte
    (statement with db params common)
    (CommonTableExpression statement db params with with1) where
      statement with db params common
statement `as` Alias cte
cte = Aliased (statement with db params) (cte ::: common)
-> Materialization
-> CommonTableExpression
     statement db params with ((cte ::: common) : with)
forall (statement :: FromType
                     -> SchemasType -> [NullType] -> RowType -> *)
       (with :: FromType) (db :: SchemasType) (params :: [NullType])
       (cte :: Symbol) (common :: RowType).
Aliased (statement with db params) (cte ::: common)
-> Materialization
-> CommonTableExpression
     statement db params with ((cte ::: common) : with)
CommonTableExpression (statement with db params common
statement `as` Alias cte
cte) Materialization
DefaultMaterialization
instance
  ( KnownSymbol cte
  , with1 ~ (cte ::: common ': with)
  ) => Aliasable cte
    (statement with db params common)
    (Path (CommonTableExpression statement db params) with with1) where
      statement with db params common
statement `as` Alias cte
cte = CommonTableExpression statement db params with with1
-> Path (CommonTableExpression statement db params) with with1
forall k k1 (c :: (k -> k1 -> *) -> k -> k1 -> *)
       (p :: k -> k1 -> *) (x :: k) (y :: k1).
QPointed c =>
p x y -> c p x y
qsingle (statement with db params common
statement `as` Alias cte
cte)

instance (forall c s p r. RenderSQL (statement c s p r)) => RenderSQL
  (CommonTableExpression statement db params with0 with1) where
    renderSQL :: CommonTableExpression statement db params with0 with1 -> ByteString
renderSQL (CommonTableExpression (statement with0 db params ty
statement `As` Alias alias
cte) Materialization
materialization) =
      Alias alias -> ByteString
forall sql. RenderSQL sql => sql -> ByteString
renderSQL Alias alias
cte
        ByteString -> ByteString -> ByteString
<+> ByteString
"AS"
        ByteString -> ByteString -> ByteString
<+> Materialization -> ByteString
forall sql. RenderSQL sql => sql -> ByteString
renderSQL Materialization
materialization
        ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> case Materialization
materialization of
              Materialization
DefaultMaterialization -> ByteString
""
              Materialization
_ -> ByteString
" "
        ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString -> ByteString
parenthesized (statement with0 db params ty -> ByteString
forall sql. RenderSQL sql => sql -> ByteString
renderSQL statement with0 db params ty
statement)

{- | Force separate calculation of the WITH query.

>>> type Columns = '["col1" ::: 'NoDef :=> 'NotNull 'PGint4, "col2" ::: 'NoDef :=> 'NotNull 'PGint4]
>>> type Schema = '["tab" ::: 'Table ('[] :=> Columns)]
>>> :{
let
  qry :: Query lat with (Public Schema) params '["col1" ::: 'NotNull 'PGint4, "col2" ::: 'NotNull 'PGint4]
  qry = with (
    materialized (select Star (from (table #tab)) `as` #cte1) :>>
    select Star (from (common #cte1)) `as` #cte2
    ) (select Star (from (common #cte2)))
in printSQL qry
:}
WITH "cte1" AS MATERIALIZED (SELECT * FROM "tab" AS "tab"), "cte2" AS (SELECT * FROM "cte1" AS "cte1") SELECT * FROM "cte2" AS "cte2"

Note: if the last CTE has `materialized` or `notMaterialized` you must add `:>> Done`.

Requires PostgreSQL 12 or higher.
-}
materialized
  :: Aliased (statement with db params) (cte ::: common) -- ^ CTE
  -> CommonTableExpression statement db params with (cte ::: common ': with)
materialized :: Aliased (statement with db params) (cte ::: common)
-> CommonTableExpression
     statement db params with ((cte ::: common) : with)
materialized Aliased (statement with db params) (cte ::: common)
stmt = Aliased (statement with db params) (cte ::: common)
-> Materialization
-> CommonTableExpression
     statement db params with ((cte ::: common) : with)
forall (statement :: FromType
                     -> SchemasType -> [NullType] -> RowType -> *)
       (with :: FromType) (db :: SchemasType) (params :: [NullType])
       (cte :: Symbol) (common :: RowType).
Aliased (statement with db params) (cte ::: common)
-> Materialization
-> CommonTableExpression
     statement db params with ((cte ::: common) : with)
CommonTableExpression Aliased (statement with db params) (cte ::: common)
stmt Materialization
Materialized

{- | Force the WITH query to be merged into the parent query.

>>> type Columns = '["col1" ::: 'NoDef :=> 'NotNull 'PGint4, "col2" ::: 'NoDef :=> 'NotNull 'PGint4]
>>> type Schema = '["tab" ::: 'Table ('[] :=> Columns)]
>>> :{
let
  qry :: Query lat with (Public Schema) params '["col1" ::: 'NotNull 'PGint4, "col2" ::: 'NotNull 'PGint4]
  qry = with (
    select Star (from (table #tab)) `as` #cte1 :>>
    notMaterialized (select Star (from (common #cte1)) `as` #cte2) :>>
    Done
    ) (select Star (from (common #cte2)))
in printSQL qry
:}
WITH "cte1" AS (SELECT * FROM "tab" AS "tab"), "cte2" AS NOT MATERIALIZED (SELECT * FROM "cte1" AS "cte1") SELECT * FROM "cte2" AS "cte2"

Note: if the last CTE has `materialized` or `notMaterialized` you must add `:>> Done` to finish the `Path`.

Requires PostgreSQL 12 or higher.
-}
notMaterialized
  :: Aliased (statement with db params) (cte ::: common) -- ^ CTE
  -> CommonTableExpression statement db params with (cte ::: common ': with)
notMaterialized :: Aliased (statement with db params) (cte ::: common)
-> CommonTableExpression
     statement db params with ((cte ::: common) : with)
notMaterialized Aliased (statement with db params) (cte ::: common)
stmt = Aliased (statement with db params) (cte ::: common)
-> Materialization
-> CommonTableExpression
     statement db params with ((cte ::: common) : with)
forall (statement :: FromType
                     -> SchemasType -> [NullType] -> RowType -> *)
       (with :: FromType) (db :: SchemasType) (params :: [NullType])
       (cte :: Symbol) (common :: RowType).
Aliased (statement with db params) (cte ::: common)
-> Materialization
-> CommonTableExpression
     statement db params with ((cte ::: common) : with)
CommonTableExpression Aliased (statement with db params) (cte ::: common)
stmt Materialization
NotMaterialized