Skip to content

Commit

Permalink
Add support for RECURSIVE withs
Browse files Browse the repository at this point in the history
  • Loading branch information
Raveline committed Jul 17, 2024
1 parent 3ceb4ee commit 8a52d13
Showing 1 changed file with 38 additions and 9 deletions.
47 changes: 38 additions & 9 deletions src/Database/PostgreSQL/PQTypes/SQL/Builder.hs
Original file line number Diff line number Diff line change
Expand Up @@ -127,10 +127,12 @@ module Database.PostgreSQL.PQTypes.SQL.Builder
, sqlLimit
, sqlDistinct
, sqlWith
, sqlWithRecursive
, sqlWithMaterialized
, sqlUnion
, sqlUnionAll
, checkAndRememberMaterializationSupport
, checkAndRememberRecursiveSupport

, sqlSelect
, sqlSelect2
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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)) <+>
Expand Down Expand Up @@ -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.
Expand Down

0 comments on commit 8a52d13

Please sign in to comment.