Merge pull request #633 from thomasjm/limit-offset-maybe

Add limitMaybe_, offsetMaybe_ + tests

Author: LaurentRDC

beam-core/Database/Beam/Query/Combinators.hs

@@ -41,7 +41,8 @@ module Database.Beam.Query.Combinators
     , QIfCond, QIfElse
     , (<|>.)
 
-    , limit_, offset_
+    , limit_, limitMaybe_
+    , offset_, offsetMaybe_
 
     , as_
 
@@ -66,8 +67,8 @@ module Database.Beam.Query.Combinators
     , orderBy_, asc_, desc_, nullsFirst_, nullsLast_
     ) where
 
-import Database.Beam.Backend.Types
 import Database.Beam.Backend.SQL
+import Database.Beam.Backend.Types
 
 import Database.Beam.Query.Internal
 import Database.Beam.Query.Ord
@@ -83,6 +84,7 @@ import Control.Applicative
 import Data.Maybe
 import Data.Proxy
 import Data.Time (LocalTime)
+import Unsafe.Coerce (unsafeCoerce)
 
 import GHC.TypeLits (TypeError, ErrorMessage(Text))
 
@@ -335,6 +337,17 @@ limit_ :: forall s a be db
 limit_ limit' (Q q) =
   Q (liftF (QLimit limit' q (rewriteThread (Proxy @s))))
 
+-- | Conditionally limit the number of results returned by a query.
+limitMaybe_ :: forall s a be db
+        . ( Projectible be a
+          , ThreadRewritable (QNested s) a )
+        => Maybe Integer -> Q be db (QNested s) a -> Q be db s (WithRewrittenThread (QNested s) s a)
+limitMaybe_ (Just limit') (Q q) =
+  Q (liftF (QLimit limit' q (rewriteThread (Proxy @s))))
+-- This uses unsafeCoerce, but should be safe since this function is tested.
+-- See discussion on https://github.com/haskell-beam/beam/pull/633.
+limitMaybe_ Nothing (Q q) = Q (unsafeCoerce q)
+
 -- | Drop the first `offset'` results.
 offset_ :: forall s a be db
          . ( Projectible be a
@@ -343,6 +356,17 @@ offset_ :: forall s a be db
 offset_ offset' (Q q) =
   Q (liftF (QOffset offset' q (rewriteThread (Proxy @s))))
 
+-- | Conditionally drop the first `offset'` results.
+offsetMaybe_ :: forall s a be db
+         . ( Projectible be a
+           , ThreadRewritable (QNested s) a )
+        => Maybe Integer -> Q be db (QNested s) a -> Q be db s (WithRewrittenThread (QNested s) s a)
+offsetMaybe_ (Just offset') (Q q) =
+  Q (liftF (QOffset offset' q (rewriteThread (Proxy @s))))
+-- This uses unsafeCoerce, but should be safe since this function is tested.
+-- See discussion on https://github.com/haskell-beam/beam/pull/633.
+offsetMaybe_ Nothing (Q q) = Q (unsafeCoerce q)
+
 -- | Use the SQL @EXISTS@ operator to determine if the given query returns any results
 exists_ :: ( BeamSqlBackend be, HasQBuilder be, Projectible be a)
         => Q be db s a -> QExpr be s Bool

beam-core/test/Database/Beam/Test/SQL.hs

@@ -1059,7 +1059,9 @@ selectCombinators =
 limitOffset :: TestTree
 limitOffset =
   testGroup "LIMIT/OFFSET support"
-  [ limitSupport, offsetSupport, limitOffsetSupport
+  [ limitSupport, maybeLimitSupportJust, maybeLimitSupportNothing
+  , offsetSupport, maybeOffsetSupportJust, maybeOffsetSupportNothing
+  , limitOffsetSupport
 
   , limitPlacedOnUnion ]
   where
@@ -1071,6 +1073,22 @@ limitOffset =
          selectLimit @?= Just 20
          selectOffset @?= Nothing
 
+    maybeLimitSupportJust =
+      testCase "Maybe LIMIT support (Just)" $
+      do SqlSelect Select { selectLimit, selectOffset } <-
+           pure $ selectMock $ limitMaybe_ (Just 20) (all_ (_employees employeeDbSettings))
+
+         selectLimit @?= Just 20
+         selectOffset @?= Nothing
+
+    maybeLimitSupportNothing =
+      testCase "Maybe LIMIT support (Nothing)" $
+      do SqlSelect Select { selectLimit, selectOffset } <-
+           pure $ selectMock $ limitMaybe_ Nothing (all_ (_employees employeeDbSettings))
+
+         selectLimit @?= Nothing
+         selectOffset @?= Nothing
+
     offsetSupport =
       testCase "Basic OFFSET support" $
       do SqlSelect Select { selectLimit, selectOffset } <-
@@ -1079,6 +1097,22 @@ limitOffset =
          selectLimit @?= Nothing
          selectOffset @?= Just 102
 
+    maybeOffsetSupportJust =
+      testCase "Maybe OFFSET support (Just)" $
+      do SqlSelect Select { selectLimit, selectOffset } <-
+           pure $ selectMock $ offsetMaybe_ (Just 2) $ offset_ 100 (all_ (_employees employeeDbSettings))
+
+         selectLimit @?= Nothing
+         selectOffset @?= Just 102
+
+    maybeOffsetSupportNothing =
+      testCase "Maybe OFFSET support (Nothing)" $
+      do SqlSelect Select { selectLimit, selectOffset } <-
+           pure $ selectMock $ offsetMaybe_ Nothing $ offset_ 100 (all_ (_employees employeeDbSettings))
+
+         selectLimit @?= Nothing
+         selectOffset @?= Just 100
+
     limitOffsetSupport =
       testCase "Basic LIMIT .. OFFSET .. support" $
       do SqlSelect Select { selectLimit, selectOffset } <-