Add support for WITH _ AS MATERIALIZED (_) queries

Before PostgreSQL 12, `WITH` was always materialized, but since PostgreSQL 12 this is not necessarily the case (it's up to the query planner).

However, a new syntax was added for explicitly specifying whether a `WITH` query should be `MATERIALIZED` or `NOT MATERIALIZED`. This commit adds support for this to Opaleye's core, and adds the function `withMaterialized` to the user-facing API.

src/Opaleye/Internal/Optimize.hs

@@ -72,7 +72,7 @@ removeEmpty = PQ.foldPrimQuery PQ.PrimQueryFold {
   , PQ.relExpr   = return .: PQ.RelExpr
   , PQ.rebind    = \b -> fmap . PQ.Rebind b
   , PQ.forUpdate = fmap PQ.ForUpdate
-  , PQ.with      = \recursive name cols -> liftA2 (PQ.With recursive name cols)
+  , PQ.with      = \recursive materialized name cols -> liftA2 (PQ.With recursive materialized name cols)
   }
   where -- If only the first argument is Just, do n1 on it
         -- If only the second argument is Just, do n2 on it

src/Opaleye/Internal/PrimQuery.hs

@@ -49,6 +49,9 @@ instance Monoid Lateral where
 data Recursive = NonRecursive | Recursive
   deriving Show
 
+data Materialized = Materialized | NotMaterialized
+  deriving Show
+
 aLeftJoin :: HPQ.PrimExpr -> PrimQuery -> PrimQueryArr
 aLeftJoin cond primQuery' = PrimQueryArr $ \lat primQueryL ->
   Join LeftJoin cond (NonLateral, primQueryL) (lat, primQuery')
@@ -162,7 +165,7 @@ data PrimQuery' a = Unit
                   -- ForUpdate in the future
                   --
                   -- https://www.postgresql.org/docs/current/sql-select.html#SQL-FOR-UPDATE-SHARE
-                  | With Recursive Symbol [Symbol] (PrimQuery' a) (PrimQuery' a)
+                  | With Recursive (Maybe Materialized) Symbol [Symbol] (PrimQuery' a) (PrimQuery' a)
                  deriving Show
 
 type PrimQuery = PrimQuery' ()
@@ -200,7 +203,7 @@ data PrimQueryFoldP a p p' = PrimQueryFold
     -- ^ A relation-valued expression
   , rebind            :: Bool -> Bindings HPQ.PrimExpr -> p -> p'
   , forUpdate         :: p -> p'
-  , with              :: Recursive -> Symbol -> [Symbol] -> p -> p -> p'
+  , with              :: Recursive -> Maybe Materialized -> Symbol -> [Symbol] -> p -> p -> p'
   }
 
 
@@ -248,7 +251,7 @@ dimapPrimQueryFold self g f = PrimQueryFold
   , relExpr = \pe bs -> g (relExpr f pe bs)
   , rebind = \s bs p -> g (rebind f s bs (self p))
   , forUpdate = \p -> g (forUpdate f (self p))
-  , with = \r s ss p1 p2 -> g (with f r s ss (self p1) (self p2))
+  , with = \r m s ss p1 p2 -> g (with f r m s ss (self p1) (self p2))
   }
 
 applyPrimQueryFoldF ::
@@ -271,7 +274,7 @@ applyPrimQueryFoldF f = \case
   Exists s q -> exists f s q
   Rebind star pes q -> rebind f star pes q
   ForUpdate q -> forUpdate f q
-  With recursive name cols a b -> with f recursive name cols a b
+  With recursive materialized name cols a b -> with f recursive materialized name cols a b
 
 primQueryFoldF ::
   PrimQueryFoldP a p p' -> (PrimQuery' a -> p) -> PrimQuery' a -> p'

src/Opaleye/Internal/Print.hs

@@ -126,12 +126,17 @@ ppRecursive :: Sql.Recursive -> Doc
 ppRecursive Sql.Recursive = text "RECURSIVE"
 ppRecursive Sql.NonRecursive = mempty
 
+ppMaterialized :: Sql.Materialized -> Doc
+ppMaterialized Sql.Materialized = text "MATERIALIZED"
+ppMaterialized Sql.NotMaterialized = text "NOT MATERIALIZED"
+
 ppWith :: With -> Doc
 ppWith w
   =  text "WITH" <+> ppRecursive (Sql.wRecursive w)
   <+> HPrint.ppTable (Sql.wTable w)
   <+> parens (HPrint.commaV unColumn (Sql.wCols w))
   <+> text "AS"
+  <+> foldMap ppMaterialized (Sql.wMaterialized w)
   $$ parens (ppSql (Sql.wWith w))
   $$ ppSql (Sql.wSelect w)
   where unColumn (HSql.SqlColumn col) = text col

src/Opaleye/Internal/Sql.hs

@@ -83,12 +83,14 @@ data BinOp = Except | ExceptAll | Union | UnionAll | Intersect | IntersectAll de
 data Lateral = Lateral | NonLateral deriving Show
 data LockStrength = Update deriving Show
 data Recursive = NonRecursive | Recursive deriving Show
+data Materialized = Materialized | NotMaterialized deriving Show
 data With = With {
-  wTable     :: HSql.SqlTable, -- The name of the result, i.e. WITH <name> AS
-  wCols      :: [HSql.SqlColumn],
-  wRecursive :: Recursive,
-  wWith      :: Select,
-  wSelect    :: Select
+  wTable        :: HSql.SqlTable, -- The name of the result, i.e. WITH <name> AS
+  wCols         :: [HSql.SqlColumn],
+  wRecursive    :: Recursive,
+  wMaterialized :: Maybe Materialized,
+  wWith         :: Select,
+  wSelect       :: Select
 } deriving Show
 
 
@@ -263,8 +265,8 @@ binary op (select1, select2) = SelectBinary Binary {
   bSelect2 = select2
   }
 
-with :: PQ.Recursive -> Symbol -> [Symbol] -> Select -> Select -> Select
-with recursive name cols wWith wSelect =
+with :: PQ.Recursive -> Maybe PQ.Materialized -> Symbol -> [Symbol] -> Select -> Select -> Select
+with recursive materialized name cols wWith wSelect =
   SelectFrom
     newSelect
       { attrs = Star
@@ -275,6 +277,10 @@ with recursive name cols wWith wSelect =
    wRecursive = case recursive of
      PQ.NonRecursive -> NonRecursive
      PQ.Recursive -> Recursive
+   wMaterialized = case materialized of
+     Nothing -> Nothing
+     Just PQ.Materialized -> Just Materialized
+     Just PQ.NotMaterialized -> Just NotMaterialized
    wCols = map (HSql.SqlColumn . sqlSymbol) cols
 
 

src/Opaleye/With.hs

@@ -2,11 +2,13 @@
 
 module Opaleye.With
   ( with,
+    withMaterialized,
     withRecursive,
     withRecursiveDistinct,
 
     -- * Explicit versions
     withExplicit,
+    withMaterializedExplicit,
     withRecursiveExplicit,
     withRecursiveDistinctExplicit,
   )
@@ -30,6 +32,9 @@ import Opaleye.Internal.Unpackspec (Unpackspec (..), runUnpackspec)
 with :: Default Unpackspec a a => Select a -> (Select a -> Select b) -> Select b
 with = withExplicit def
 
+withMaterialized :: Default Unpackspec a a => Select a -> (Select a -> Select b) -> Select b
+withMaterialized = withMaterializedExplicit def
+
 -- | Denotionally, @withRecursive s f@ is the smallest set of rows @r@ such
 -- that
 --
@@ -65,7 +70,13 @@ withRecursiveDistinct = withRecursiveDistinctExplicit def
 
 withExplicit :: Unpackspec a a -> Select a -> (Select a -> Select b) -> Select b
 withExplicit unpackspec rhsSelect bodySelect = productQueryArr $ do
-  withG unpackspec PQ.NonRecursive (\_ -> rebind rhsSelect) bodySelect
+  withG unpackspec PQ.NonRecursive Nothing (\_ -> rebind rhsSelect) bodySelect
+  where
+    rebind = (>>> rebindExplicitPrefixNoStar "rebind" unpackspec)
+
+withMaterializedExplicit :: Unpackspec a a -> Select a -> (Select a -> Select b) -> Select b
+withMaterializedExplicit unpackspec rhsSelect bodySelect = productQueryArr $ do
+  withG unpackspec PQ.NonRecursive (Just PQ.Materialized) (\_ -> rebind rhsSelect) bodySelect
   where
     rebind = (>>> rebindExplicitPrefixNoStar "rebind" unpackspec)
 
@@ -74,7 +85,7 @@ withRecursiveExplicit binaryspec base recursive = productQueryArr $ do
   let bodySelect selectCte = selectCte
   let rhsSelect selectCte = unionAllExplicit binaryspec base (selectCte >>= recursive)
 
-  withG unpackspec PQ.Recursive rhsSelect bodySelect
+  withG unpackspec PQ.Recursive Nothing rhsSelect bodySelect
   where
     unpackspec = binaryspecToUnpackspec binaryspec
 
@@ -83,17 +94,18 @@ withRecursiveDistinctExplicit binaryspec base recursive = productQueryArr $ do
   let bodySelect selectCte = selectCte
   let rhsSelect selectCte = unionExplicit binaryspec base (selectCte >>= recursive)
 
-  withG unpackspec PQ.Recursive rhsSelect bodySelect
+  withG unpackspec PQ.Recursive Nothing rhsSelect bodySelect
   where
     unpackspec = binaryspecToUnpackspec binaryspec
 
 withG ::
   Unpackspec a a ->
   PQ.Recursive ->
+  Maybe PQ.Materialized ->
   (Select a -> Select a) ->
   (Select a -> Select b) ->
   State Tag.Tag (b, PQ.PrimQuery)
-withG unpackspec recursive rhsSelect bodySelect = do
+withG unpackspec recursive materialized rhsSelect bodySelect = do
   (selectCte, withCte) <- freshCte unpackspec
 
   let rhsSelect' = rhsSelect selectCte
@@ -102,14 +114,14 @@ withG unpackspec recursive rhsSelect bodySelect = do
   (_, rhsQ) <- runSimpleSelect rhsSelect'
   bodyQ <- runSimpleSelect bodySelect'
 
-  pure (withCte recursive rhsQ bodyQ)
+  pure (withCte recursive materialized rhsQ bodyQ)
 
 freshCte ::
   Unpackspec a a ->
   State
     Tag.Tag
     ( Select a,
-      PQ.Recursive -> PQ.PrimQuery -> (b, PQ.PrimQuery) -> (b, PQ.PrimQuery)
+      PQ.Recursive -> Maybe PQ.Materialized -> PQ.PrimQuery -> (b, PQ.PrimQuery) -> (b, PQ.PrimQuery)
     )
 freshCte unpackspec = do
   cteName <- HPQ.Symbol "cte" <$> Tag.fresh
@@ -131,8 +143,8 @@ freshCte unpackspec = do
 
   pure
     ( selectCte,
-      \recursive withQ (withedCols, withedQ) ->
-        (withedCols, PQ.With recursive cteName (map fst cteBindings) withQ withedQ)
+      \recursive materialized withQ (withedCols, withedQ) ->
+        (withedCols, PQ.With recursive materialized cteName (map fst cteBindings) withQ withedQ)
     )
 
 binaryspecToUnpackspec :: Binaryspec a a -> Unpackspec a a