Jnm/add cte materialized to postgres (#417)

* Materialized CTEs for Postgres * changelog update * fix tests --------- Co-authored-by: Joel McCracken <mccracken.joel@gmail.com>
bitemyapp · Dec 26, 2024 · e739a91 · e739a91
1 parent 324de6a
commit e739a91
Show file tree

Hide file tree

Showing 7 changed files with 212 additions and 13 deletions.
diff --git a/changelog.md b/changelog.md
@@ -11,6 +11,9 @@
     - [#341](https://github.com/bitemyapp/esqueleto/pull/341/)
         - Add functions for `NULLS FIRST` and `NULLS LAST` in the Postgresql
           module
+- @JoelMcCracken
+    - [#354](https://github.com/bitemyapp/esqueleto/pull/354), [#417](https://github.com/bitemyapp/esqueleto/pull/417)
+      - Add `withMaterialized`, `withNotMaterialized` to the PostgreSQL module
 
 3.5.13.2
 ========

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 $ do
@@ -107,7 +108,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
@@ -1998,8 +1998,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
@@ -3212,14 +3214,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
@@ -36,6 +36,8 @@ module Database.Esqueleto.PostgreSQL
     , forKeyShareOf
     , filterWhere
     , values
+    , withMaterialized
+    , withNotMaterialized
     , ascNullsFirst
     , ascNullsLast
     , descNullsFirst
@@ -52,15 +54,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
@@ -490,7 +499,7 @@ forNoKeyUpdateOf lockableEntities onLockedBehavior =
 forShareOf :: LockableEntity a => a -> OnLockedBehavior -> SqlQuery ()
 forShareOf lockableEntities onLockedBehavior =
   putLocking $ PostgresLockingClauses [PostgresLockingKind PostgresForShare (Just $ LockingOfClause lockableEntities) onLockedBehavior]
-  
+
 -- | `FOR KEY SHARE OF` syntax for postgres locking
 -- allows locking of specific tables with a key share lock in a view or join
 --
@@ -499,6 +508,82 @@ forKeyShareOf :: LockableEntity a => a -> OnLockedBehavior -> SqlQuery ()
 forKeyShareOf lockableEntities onLockedBehavior =
   putLocking $ PostgresLockingClauses [PostgresLockingKind PostgresForKeyShare (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.14.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.14.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)))
+
 -- | Ascending order of this field or SqlExpression with nulls coming first.
 --
 -- @since 3.5.14.0

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,80 @@ 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