From 8a52d13a2fabba1a9f084f0f11d4174596e548d9 Mon Sep 17 00:00:00 2001
From: Raveline <adrien.duclos@scrive.com>
Date: Wed, 17 Jul 2024 16:25:04 +0200
Subject: [PATCH] Add support for RECURSIVE withs

---
 .../PostgreSQL/PQTypes/SQL/Builder.hs         | 47 +++++++++++++++----
 1 file changed, 38 insertions(+), 9 deletions(-)

diff --git a/src/Database/PostgreSQL/PQTypes/SQL/Builder.hs b/src/Database/PostgreSQL/PQTypes/SQL/Builder.hs
index 3404f49..73a547c 100644
--- a/src/Database/PostgreSQL/PQTypes/SQL/Builder.hs
+++ b/src/Database/PostgreSQL/PQTypes/SQL/Builder.hs
@@ -127,10 +127,12 @@ module Database.PostgreSQL.PQTypes.SQL.Builder
   , sqlLimit
   , sqlDistinct
   , sqlWith
+  , sqlWithRecursive
   , sqlWithMaterialized
   , sqlUnion
   , sqlUnionAll
   , checkAndRememberMaterializationSupport
+  , checkAndRememberRecursiveSupport
 
   , sqlSelect
   , sqlSelect2
@@ -241,7 +243,7 @@ data SqlSelect = SqlSelect
   , sqlSelectHaving   :: [SQL]
   , sqlSelectOffset   :: Integer
   , sqlSelectLimit    :: Integer
-  , sqlSelectWith     :: [(SQL, SQL, Materialized)]
+  , sqlSelectWith     :: [(SQL, SQL, Materialized, Recursive)]
   }
 
 data SqlUpdate = SqlUpdate
@@ -341,7 +343,7 @@ instance IsSQL SqlDelete where
 instance Sqlable SqlSelect where
   toSQLCommand cmd = smconcat
     [ emitClausesSepComma "WITH" $
-      map (\(name,command,mat) -> name <+> "AS" <+> materializedClause mat <+> parenthesize command) (sqlSelectWith cmd)
+      map (\(name,command,mat,recurse) -> recursiveClause recurse <+> name <+> "AS" <+> materializedClause mat <+> parenthesize command) (sqlSelectWith cmd)
     , if hasUnion || hasUnionAll
       then emitClausesSep "" unionKeyword (mainSelectClause : unionCmd)
       else mainSelectClause
@@ -443,6 +445,16 @@ checkAndRememberMaterializationSupport = do
     fetchOne runIdentity
   liftIO $ writeIORef withMaterializedSupported (isRight res)
 
+-- This function has to be called as one of first things in your program
+-- for the library to make sure that it is aware if the "WITH RECURSIVE"
+-- clause is supported by your PostgreSQL version.
+checkAndRememberRecursiveSupport :: (MonadDB m, MonadIO m, MonadMask m) => m ()
+checkAndRememberRecursiveSupport = do
+  res :: Either DBException Int64 <- try . withNewConnection $ do
+    runSQL01_ "WITH RECURSIVE t(n) AS (SELECT (1 :: bigint)) SELECT n FROM t LIMIT 1"
+    fetchOne runIdentity
+  liftIO $ writeIORef withRecursiveSupported (isRight res)
+
 withMaterializedSupported :: IORef Bool
 {-# NOINLINE withMaterializedSupported #-}
 withMaterializedSupported = unsafePerformIO $ newIORef False
@@ -451,10 +463,22 @@ isWithMaterializedSupported :: Bool
 {-# NOINLINE isWithMaterializedSupported #-}
 isWithMaterializedSupported = unsafePerformIO $ readIORef withMaterializedSupported
 
+withRecursiveSupported :: IORef Bool
+{-# NOINLINE withRecursiveSupported #-}
+withRecursiveSupported = unsafePerformIO $ newIORef False
+
+isWithRecursiveSupported :: Bool
+{-# NOINLINE isWithRecursiveSupported #-}
+isWithRecursiveSupported = unsafePerformIO $ readIORef withRecursiveSupported
+
 materializedClause :: Materialized -> SQL
 materializedClause Materialized = if isWithMaterializedSupported then "MATERIALIZED" else ""
 materializedClause NonMaterialized = if isWithMaterializedSupported then "NOT MATERIALIZED" else ""
 
+recursiveClause :: Recursive -> SQL
+recursiveClause Recursive = if isWithRecursiveSupported then "RECURSIVE" else ""
+recursiveClause NonRecursive = ""
+
 instance Sqlable SqlUpdate where
   toSQLCommand cmd =
     emitClausesSepComma "WITH" (map (\(name,command,mat) -> name <+> "AS" <+> materializedClause mat <+> parenthesize command) (sqlUpdateWith cmd)) <+>
@@ -522,28 +546,33 @@ sqlDelete table refine =
 
 
 data Materialized = Materialized | NonMaterialized
+data Recursive = Recursive | NonRecursive
 
 class SqlWith a where
-  sqlWith1 :: a -> SQL -> SQL -> Materialized -> a
+  sqlWith1 :: a -> SQL -> SQL -> Materialized -> Recursive -> a
 
 
 instance SqlWith SqlSelect where
-  sqlWith1 cmd name sql mat = cmd { sqlSelectWith = sqlSelectWith cmd ++ [(name,sql, mat)] }
+  sqlWith1 cmd name sql mat recurse = cmd { sqlSelectWith = sqlSelectWith cmd ++ [(name,sql,mat,recurse)] }
 
 instance SqlWith SqlInsertSelect where
-  sqlWith1 cmd name sql mat = cmd { sqlInsertSelectWith = sqlInsertSelectWith cmd ++ [(name,sql,mat)] }
+  sqlWith1 cmd name sql mat _ = cmd { sqlInsertSelectWith = sqlInsertSelectWith cmd ++ [(name,sql,mat)] }
 
 instance SqlWith SqlUpdate where
-  sqlWith1 cmd name sql mat = cmd { sqlUpdateWith = sqlUpdateWith cmd ++ [(name,sql,mat)] }
+  sqlWith1 cmd name sql mat _ = cmd { sqlUpdateWith = sqlUpdateWith cmd ++ [(name,sql,mat)] }
 
 instance SqlWith SqlDelete where
-  sqlWith1 cmd name sql mat = cmd { sqlDeleteWith = sqlDeleteWith cmd ++ [(name,sql,mat)] }
+  sqlWith1 cmd name sql mat _ = cmd { sqlDeleteWith = sqlDeleteWith cmd ++ [(name,sql,mat)] }
 
 sqlWith :: (MonadState v m, SqlWith v, Sqlable s) => SQL -> s -> m ()
-sqlWith name sql = modify (\cmd -> sqlWith1 cmd name (toSQLCommand sql) NonMaterialized)
+sqlWith name sql = modify (\cmd -> sqlWith1 cmd name (toSQLCommand sql) NonMaterialized NonRecursive)
 
 sqlWithMaterialized :: (MonadState v m, SqlWith v, Sqlable s) => SQL -> s -> m ()
-sqlWithMaterialized name sql = modify (\cmd -> sqlWith1 cmd name (toSQLCommand sql) Materialized)
+sqlWithMaterialized name sql = modify (\cmd -> sqlWith1 cmd name (toSQLCommand sql) Materialized NonRecursive)
+
+-- | Note: RECURSIVE with are only available with a SELECT. Using any other type of query inside will fail.
+sqlWithRecursive :: (MonadState v m, SqlWith v, Sqlable s) => SQL -> s -> m ()
+sqlWithRecursive name sql = modify (\cmd -> sqlWith1 cmd name (toSQLCommand sql) NonMaterialized Recursive)
 
 -- | Note: WHERE clause of the main SELECT is treated specially, i.e. it only
 -- applies to the main SELECT, not the whole union.