Materialized CTEs for Postgres

bitemyapp · Jan 19, 2024 · 07dd001 · 07dd001
1 parent f689e22
commit 07dd001
Show file tree

Hide file tree

Showing 8 changed files with 213 additions and 13 deletions.
diff --git a/changelog.md b/changelog.md
@@ -1,3 +1,9 @@
+3.5.12.0
+========
+- @JoelMcCracken
+    - [#354](https://github.com/bitemyapp/esqueleto/pull/354)
+      - Add `withMaterialized`, `withNotMaterialized` to the PostgreSQL module
+
 3.5.11.0
 ========
 - @9999years, @halogenandtoast

diff --git a/esqueleto.cabal b/esqueleto.cabal
@@ -2,7 +2,7 @@ cabal-version: 1.12
 
 name:           esqueleto
 
-version:        3.5.11.1
+version:        3.5.12.0
 synopsis:       Type-safe EDSL for SQL queries on persistent backends.
 description:    @esqueleto@ is a bare bones, type-safe EDSL for SQL queries that works with unmodified @persistent@ SQL backends.  Its language closely resembles SQL, so you don't have to learn new concepts, just new syntax, and it's fairly easy to predict the generated SQL and optimize it for your backend. Most kinds of errors committed when writing SQL are caught as compile-time errors---although it is possible to write type-checked @esqueleto@ queries that fail at runtime.
                 .

diff --git a/src/Database/Esqueleto/Experimental/From/CommonTableExpression.hs b/src/Database/Esqueleto/Experimental/From/CommonTableExpression.hs
@@ -38,7 +38,8 @@ import Database.Esqueleto.Internal.Internal hiding (From(..), from, on)
 -- PostgreSQL 12, non-recursive and side-effect-free queries may be inlined and
 -- optimized accordingly if not declared @MATERIALIZED@ to get the previous
 -- behaviour. See [the PostgreSQL CTE documentation](https://www.postgresql.org/docs/current/queries-with.html#id-1.5.6.12.7),
--- section Materialization, for more information.
+-- section Materialization, for more information. To use a @MATERIALIZED@ query
+-- in Esquelto, see functions 'withMaterialized' and 'withRecursiveMaterialized'.
 --
 -- /Since: 3.4.0.0/
 with :: ( ToAlias a
@@ -50,7 +51,7 @@ with query = do
     aliasedValue <- toAlias ret
     let aliasedQuery = Q $ W.WriterT $ pure (aliasedValue, sideData)
     ident <- newIdentFor (DBName "cte")
-    let clause = CommonTableExpressionClause NormalCommonTableExpression ident (\info -> toRawSql SELECT info aliasedQuery)
+    let clause = CommonTableExpressionClause NormalCommonTableExpression (\_ _ -> "")  ident (\info -> toRawSql SELECT info aliasedQuery)
     Q $ W.tell mempty{sdCteClause = [clause]}
     ref <- toAliasReference ident aliasedValue
     pure $ From $ pure (ref, (\_ info -> (useIdent info ident, mempty)))
@@ -103,7 +104,8 @@ withRecursive baseCase unionKind recursiveCase = do
     ref <- toAliasReference ident aliasedValue
     let refFrom = From (pure (ref, (\_ info -> (useIdent info ident, mempty))))
     let recursiveQuery = recursiveCase refFrom
-    let clause = CommonTableExpressionClause RecursiveCommonTableExpression ident
+    let noModifier _ _ = ""
+    let clause = CommonTableExpressionClause RecursiveCommonTableExpression noModifier ident
                  (\info -> (toRawSql SELECT info aliasedQuery)
                         <> ("\n" <> (unUnionKind unionKind)  <> "\n", mempty)
                         <> (toRawSql SELECT info recursiveQuery)

diff --git a/src/Database/Esqueleto/Internal/Internal.hs b/src/Database/Esqueleto/Internal/Internal.hs
@@ -1876,8 +1876,10 @@ data CommonTableExpressionKind
     | NormalCommonTableExpression
     deriving Eq
 
-data CommonTableExpressionClause =
-    CommonTableExpressionClause CommonTableExpressionKind Ident (IdentInfo -> (TLB.Builder, [PersistValue]))
+type CommonTableExpressionModifierAfterAs = CommonTableExpressionClause -> IdentInfo -> TLB.Builder
+
+data CommonTableExpressionClause
+  = CommonTableExpressionClause CommonTableExpressionKind CommonTableExpressionModifierAfterAs Ident (IdentInfo -> (TLB.Builder, [PersistValue]))
 
 data SubQueryType
     = NormalSubQuery
@@ -3090,14 +3092,15 @@ makeCte info cteClauses =
         | hasRecursive = "WITH RECURSIVE "
         | otherwise = "WITH "
       where
+
         hasRecursive =
             elem RecursiveCommonTableExpression
-            $ fmap (\(CommonTableExpressionClause cteKind _ _) -> cteKind)
+            $ fmap (\(CommonTableExpressionClause cteKind _ _ _) -> cteKind)
             $ cteClauses
 
-    cteClauseToText (CommonTableExpressionClause _ cteIdent cteFn) =
+    cteClauseToText clause@(CommonTableExpressionClause _ cteModifier cteIdent cteFn) =
         first
-            (\tlb -> useIdent info cteIdent <> " AS " <> parens tlb)
+            (\tlb -> useIdent info cteIdent <> " AS " <> cteModifier clause info <> parens tlb)
             (cteFn info)
 
     cteBody =

diff --git a/src/Database/Esqueleto/PostgreSQL.hs b/src/Database/Esqueleto/PostgreSQL.hs
@@ -34,6 +34,8 @@ module Database.Esqueleto.PostgreSQL
     , forShareOf
     , filterWhere
     , values
+    , withMaterialized
+    , withNotMaterialized
     -- * Internal
     , unsafeSqlAggregateFunction
     ) where
@@ -46,15 +48,22 @@ import Control.Exception (throw)
 import Control.Monad (void)
 import Control.Monad.IO.Class (MonadIO(..))
 import qualified Control.Monad.Trans.Reader as R
+import qualified Control.Monad.Trans.Writer as W
 import Data.Int (Int64)
 import qualified Data.List.NonEmpty as NE
 import Data.Maybe
 import Data.Proxy (Proxy(..))
 import qualified Data.Text.Internal.Builder as TLB
 import qualified Data.Text.Lazy as TL
+import qualified Data.Text.Lazy.Builder as TLB
 import Data.Time.Clock (UTCTime)
 import qualified Database.Esqueleto.Experimental as Ex
-import Database.Esqueleto.Internal.Internal hiding (random_)
+import qualified Database.Esqueleto.Experimental.From as Ex
+import Database.Esqueleto.Experimental.From.CommonTableExpression
+import Database.Esqueleto.Experimental.From.SqlSetOperation
+import Database.Esqueleto.Experimental.ToAlias
+import Database.Esqueleto.Experimental.ToAliasReference
+import Database.Esqueleto.Internal.Internal hiding (From(..), from, on, random_)
 import Database.Esqueleto.Internal.PersistentImport hiding
        (uniqueFields, upsert, upsertBy)
 import Database.Persist.SqlBackend
@@ -477,3 +486,79 @@ forUpdateOf lockableEntities onLockedBehavior =
 forShareOf :: LockableEntity a => a -> OnLockedBehavior -> SqlQuery ()
 forShareOf lockableEntities onLockedBehavior =
   putLocking $ PostgresLockingClauses [PostgresLockingKind PostgresForShare (Just $ LockingOfClause lockableEntities) onLockedBehavior]
+
+-- | @WITH@ @MATERIALIZED@ clause is used to introduce a
+-- [Common Table Expression (CTE)](https://en.wikipedia.org/wiki/Hierarchical_and_recursive_queries_in_SQL#Common_table_expression)
+-- with the MATERIALIZED keyword. The MATERIALIZED keyword is only supported in PostgreSQL >= version 12.
+-- In Esqueleto, CTEs should be used as a subquery memoization tactic. PostgreSQL treats a materialized CTE as an optimization fence.
+-- A materialized CTE is always fully calculated, and is not "inlined" with other table joins.
+-- Without the MATERIALIZED keyword, PostgreSQL >= 12 may "inline" the CTE as though it was any other join.
+-- You should always verify that using a materialized CTE will in fact improve your performance
+-- over a regular subquery.
+--
+-- @
+-- select $ do
+-- cte <- withMaterialized subQuery
+-- cteResult <- from cte
+-- where_ $ cteResult ...
+-- pure cteResult
+-- @
+--
+--
+-- For more information on materialized CTEs, see the PostgreSQL manual documentation on
+-- [Common Table Expression Materialization](https://www.postgresql.org/docs/14/queries-with.html#id-1.5.6.12.7).
+--
+-- /Since: 3.5.12.0/
+withMaterialized :: ( ToAlias a
+                    , ToAliasReference a
+                    , SqlSelect a r
+                    ) => SqlQuery a -> SqlQuery (Ex.From a)
+withMaterialized query = do
+    (ret, sideData) <- Q $ W.censor (\_ -> mempty) $ W.listen $ unQ query
+    aliasedValue <- toAlias ret
+    let aliasedQuery = Q $ W.WriterT $ pure (aliasedValue, sideData)
+    ident <- newIdentFor (DBName "cte")
+    let clause = CommonTableExpressionClause NormalCommonTableExpression (\_ _ -> "MATERIALIZED ") ident (\info -> toRawSql SELECT info aliasedQuery)
+    Q $ W.tell mempty{sdCteClause = [clause]}
+    ref <- toAliasReference ident aliasedValue
+    pure $ Ex.From $ pure (ref, (\_ info -> (useIdent info ident, mempty)))
+
+-- | @WITH@ @NOT@ @MATERIALIZED@ clause is used to introduce a
+-- [Common Table Expression (CTE)](https://en.wikipedia.org/wiki/Hierarchical_and_recursive_queries_in_SQL#Common_table_expression)
+-- with the NOT MATERIALIZED keywords. These are only supported in PostgreSQL >=
+-- version 12. In Esqueleto, CTEs should be used as a subquery memoization
+-- tactic. PostgreSQL treats a materialized CTE as an optimization fence. A
+-- MATERIALIZED CTE is always fully calculated, and is not "inlined" with other
+-- table joins. Sometimes, this is undesirable, so postgres provides the NOT
+-- MATERIALIZED modifier to prevent this behavior, thus enabling it to possibly
+-- decide to treat the CTE as any other join.
+--
+-- Given the above, it is unlikely that this function will be useful, as a
+-- normal join should be used instead, but is provided for completeness.
+--
+-- @
+-- select $ do
+-- cte <- withNotMaterialized subQuery
+-- cteResult <- from cte
+-- where_ $ cteResult ...
+-- pure cteResult
+-- @
+--
+--
+-- For more information on materialized CTEs, see the PostgreSQL manual documentation on
+-- [Common Table Expression Materialization](https://www.postgresql.org/docs/14/queries-with.html#id-1.5.6.12.7).
+--
+-- /Since: 3.5.12.0/
+withNotMaterialized :: ( ToAlias a
+                    , ToAliasReference a
+                    , SqlSelect a r
+                    ) => SqlQuery a -> SqlQuery (Ex.From a)
+withNotMaterialized query = do
+    (ret, sideData) <- Q $ W.censor (\_ -> mempty) $ W.listen $ unQ query
+    aliasedValue <- toAlias ret
+    let aliasedQuery = Q $ W.WriterT $ pure (aliasedValue, sideData)
+    ident <- newIdentFor (DBName "cte")
+    let clause = CommonTableExpressionClause NormalCommonTableExpression (\_ _ -> "NOT MATERIALIZED ") ident (\info -> toRawSql SELECT info aliasedQuery)
+    Q $ W.tell mempty{sdCteClause = [clause]}
+    ref <- toAliasReference ident aliasedValue
+    pure $ Ex.From $ pure (ref, (\_ info -> (useIdent info ident, mempty)))
diff --git a/test/MySQL/Test.hs b/test/MySQL/Test.hs
@@ -23,14 +23,16 @@ import Database.Persist.MySQL
        , connectPassword
        , connectPort
        , connectUser
+       , createMySQLPool
        , defaultConnectInfo
        , withMySQLConn
-       , createMySQLPool
        )
 
 import Test.Hspec
 
 import Common.Test
+import Data.Maybe (fromMaybe)
+import System.Environment (lookupEnv)
 
 testMysqlSum :: SpecDb
 testMysqlSum = do
@@ -189,6 +191,7 @@ migrateIt = do
 mkConnectionPool :: IO ConnectionPool
 mkConnectionPool = do
     ci <- isCI
+    mysqlHost <- (fromMaybe "localhost" <$> lookupEnv "MYSQL_HOST")
     let connInfo
             | ci =
                 defaultConnectInfo
@@ -200,7 +203,7 @@ mkConnectionPool = do
                     }
             | otherwise =
                 defaultConnectInfo
-                    { connectHost     = "localhost"
+                    { connectHost     = mysqlHost
                     , connectUser     = "travis"
                     , connectPassword = "esqutest"
                     , connectDatabase = "esqutest"

diff --git a/test/PostgreSQL/Test.hs b/test/PostgreSQL/Test.hs
@@ -40,7 +40,8 @@ import Database.Esqueleto hiding (random_)
 import Database.Esqueleto.Experimental hiding (from, on, random_)
 import qualified Database.Esqueleto.Experimental as Experimental
 import qualified Database.Esqueleto.Internal.Internal as ES
-import Database.Esqueleto.PostgreSQL (random_)
+import Database.Esqueleto.PostgreSQL
+       (random_, withMaterialized, withNotMaterialized)
 import qualified Database.Esqueleto.PostgreSQL as EP
 import Database.Esqueleto.PostgreSQL.JSON hiding ((-.), (?.), (||.))
 import qualified Database.Esqueleto.PostgreSQL.JSON as JSON
@@ -1232,6 +1233,78 @@ testCommonTableExpressions = do
             pure res
         asserting $ vals `shouldBe` fmap Value [2..11]
 
+    describe "MATERIALIZED CTEs" $ do
+      describe "withNotMaterialized" $ do
+        itDb "successfully executes query" $ do
+          void $ select $ do
+            limitedLordsCte <-
+                withNotMaterialized $ do
+                    lords <- Experimental.from $ Experimental.table @Lord
+                    limit 10
+                    pure lords
+            lords <- Experimental.from limitedLordsCte
+            orderBy [asc $ lords ^. LordId]
+            pure lords
+
+          asserting noExceptions
+
+        itDb "generates the expected SQL" $ do
+          (sql, _) <- showQuery ES.SELECT $ do
+                  limitedLordsCte <-
+                      withNotMaterialized $ do
+                          lords <- Experimental.from $ Experimental.table @Lord
+                          limit 10
+                          pure lords
+                  lords <- Experimental.from limitedLordsCte
+                  orderBy [asc $ lords ^. LordId]
+                  pure lords
+
+          asserting $ sql `shouldBe` T.unlines
+            [ "WITH \"cte\" AS NOT MATERIALIZED (SELECT \"Lord\".\"county\" AS \"v_county\", \"Lord\".\"dogs\" AS \"v_dogs\""
+            , "FROM \"Lord\""
+            , " LIMIT 10)"
+            , "SELECT \"cte\".\"v_county\", \"cte\".\"v_dogs\""
+            , "FROM \"cte\""
+            , "ORDER BY \"cte\".\"v_county\" ASC"
+            ]
+          asserting noExceptions
+
+
+      describe "withMaterialized" $ do
+        itDb "generates the expected SQL" $ do
+          (sql, _) <- showQuery ES.SELECT $ do
+                  limitedLordsCte <-
+                      withMaterialized $ do
+                          lords <- Experimental.from $ Experimental.table @Lord
+                          limit 10
+                          pure lords
+                  lords <- Experimental.from limitedLordsCte
+                  orderBy [asc $ lords ^. LordId]
+                  pure lords
+
+          asserting $ sql `shouldBe` T.unlines
+            [ "WITH \"cte\" AS MATERIALIZED (SELECT \"Lord\".\"county\" AS \"v_county\", \"Lord\".\"dogs\" AS \"v_dogs\""
+            , "FROM \"Lord\""
+            , " LIMIT 10)"
+            , "SELECT \"cte\".\"v_county\", \"cte\".\"v_dogs\""
+            , "FROM \"cte\""
+            , "ORDER BY \"cte\".\"v_county\" ASC"
+            ]
+          asserting noExceptions
+
+        itDb "successfully executes query" $ do
+            void $ select $ do
+                  limitedLordsCte <-
+                      withMaterialized $ do
+                          lords <- Experimental.from $ Experimental.table @Lord
+                          limit 10
+                          pure lords
+                  lords <- Experimental.from limitedLordsCte
+                  orderBy [asc $ lords ^. LordId]
+                  pure lords
+
+            asserting noExceptions
+
 testPostgresqlLocking :: SpecDb
 testPostgresqlLocking = do
     describe "Monoid instance" $ do

diff --git a/test/docker-compose.yml b/test/docker-compose.yml
@@ -0,0 +1,28 @@
+# docker-compose file for running postgres and mysql DBMS
+
+# If using this to run the tests,
+# while these containers are running (i.e. after something like)
+#     (cd test; docker-compose up -d)
+# the tests must be told to use the hostname via MYSQL_HOST environment variable
+# e.g. something like:
+#     MYSQL_HOST=127.0.0.1 stack test
+
+version: '3'
+services:
+  postgres:
+    image: 'postgres:15.2-alpine'
+    environment:
+      POSTGRES_USER: esqutest
+      POSTGRES_PASSWORD: esqutest
+      POSTGRES_DB: esqutest
+    ports:
+    - 5432:5432
+  mysql:
+    image: 'mysql:8.0.32'
+    environment:
+      MYSQL_USER: travis
+      MYSQL_PASSWORD: esqutest
+      MYSQL_ROOT_PASSWORD: esqutest
+      MYSQL_DATABASE: esqutest
+    ports:
+    - 3306:3306