Allow limits and offsets by Fields

This work is due to @shane-circuithub

Test/Test.hs

@@ -739,6 +739,38 @@ testLimitOffset = it "" $ limitOrderShouldMatch (O.limit 2 . O.offset 2) (take 2
 testOffsetLimit :: Test
 testOffsetLimit = it "" $ limitOrderShouldMatch (O.offset 2 . O.limit 2) (drop 2 . take 2)
 
+limitFieldOrderShouldMatch
+  :: (Field O.SqlInt8 -> Select (Field O.SqlInt4, Field O.SqlInt4) -> Select (Field O.SqlInt4, Field O.SqlInt4))
+  -> (Int -> [(Int, Int)] -> [(Int, Int)])
+  -> (PGS.Connection -> Expectation)
+limitFieldOrderShouldMatch olQ ol =
+  testH
+    (nsQ >>= \n -> olQ n (orderQ table1Q))
+    ((ns >>= \n -> ol n (order table1data)) `shouldBe`)
+  where
+    orderQ = O.orderBy (O.desc snd)
+    order = L.sortBy (flip (Ord.comparing snd))
+    ns = [1, 2, 3, 4]
+    nsQ = O.values $ fromIntegral <$> ns
+
+testLimitField :: Test
+testLimitField = it "" $ limitFieldOrderShouldMatch O.limitField take
+
+testOffsetField :: Test
+testOffsetField = it "" $ limitFieldOrderShouldMatch O.offsetField drop
+
+testLimitFieldOffset :: Test
+testLimitFieldOffset = it "" $
+  limitFieldOrderShouldMatch
+    (\n -> O.limitField n . O.offsetField n)
+    (\n -> take n . drop n)
+
+testOffsetFieldLimit :: Test
+testOffsetFieldLimit = it "" $
+  limitFieldOrderShouldMatch
+    (\n -> O.offsetField n . O.limitField n)
+    (\n -> drop n . take n)
+
 testDistinctAndAggregate :: Test
 testDistinctAndAggregate = it "" $ q `selectShouldReturnSorted` expectedResult
   where q = O.distinct table1Q
@@ -1595,6 +1627,11 @@ main = do
         testOffset
         testLimitOffset
         testOffsetLimit
+      describe "limit field" $ do
+        testLimitField
+        testOffsetField
+        testLimitFieldOffset
+        testOffsetFieldLimit
       describe "double" $ do
         testDoubleDistinct
         testDoubleLeftJoin

src/Opaleye/Internal/Order.hs

@@ -60,10 +60,10 @@ orderByU os (columns, primQ) = (columns, primQ')
 orderExprs :: a -> Order a -> [HPQ.OrderExpr]
 orderExprs x (Order os) = map (uncurry HPQ.OrderExpr) (os x)
 
-limit' :: Int -> (a, PQ.PrimQuery) -> (a, PQ.PrimQuery)
+limit' :: HPQ.PrimExpr -> (a, PQ.PrimQuery) -> (a, PQ.PrimQuery)
 limit' n (x, q) = (x, PQ.Limit (PQ.LimitOp n) q)
 
-offset' :: Int -> (a, PQ.PrimQuery) -> (a, PQ.PrimQuery)
+offset' :: HPQ.PrimExpr -> (a, PQ.PrimQuery) -> (a, PQ.PrimQuery)
 offset' n (x, q) = (x, PQ.Limit (PQ.OffsetOp n) q)
 
 distinctOn :: U.Unpackspec b b -> (a -> b)
@@ -77,7 +77,7 @@ distinctOnBy ups proj ord (cols, pq) = (cols, pqOut)
             Just xs -> PQ.DistinctOnOrderBy (Just xs) oexprs pq
             Nothing -> PQ.Limit (PQ.LimitOp one) (PQ.DistinctOnOrderBy Nothing oexprs pq)
           oexprs = orderExprs cols ord
-          one = 1
+          one = HPQ.ConstExpr (HPQ.IntegerLit 1)
 
 -- | Order the results of a given query exactly, as determined by the given list
 -- of input fields. Note that this list does not have to contain an entry for

src/Opaleye/Internal/PrimQuery.hs

@@ -9,9 +9,9 @@ import qualified Opaleye.Internal.HaskellDB.Sql as HSql
 import qualified Opaleye.Internal.HaskellDB.PrimQuery as HPQ
 import           Opaleye.Internal.HaskellDB.PrimQuery (Symbol)
 
-data LimitOp = LimitOp Int
-             | OffsetOp Int
-             | LimitOffsetOp Int Int
+data LimitOp = LimitOp HPQ.PrimExpr
+             | OffsetOp HPQ.PrimExpr
+             | LimitOffsetOp HPQ.PrimExpr HPQ.PrimExpr
              deriving Show
 
 data BinOp = Except

src/Opaleye/Internal/Print.hs

@@ -192,13 +192,18 @@ ppGroupBy :: Maybe (NEL.NonEmpty HSql.SqlExpr) -> Doc
 ppGroupBy Nothing   = empty
 ppGroupBy (Just xs) = HPrint.ppGroupBy (NEL.toList xs)
 
-ppLimit :: Maybe Int -> Doc
+ppLimit :: Maybe HSql.SqlExpr -> Doc
 ppLimit Nothing = empty
-ppLimit (Just n) = text ("LIMIT " ++ show n)
+ppLimit (Just n) = text "LIMIT" <+> ppSqlExprParens n
 
-ppOffset :: Maybe Int -> Doc
+ppOffset :: Maybe HSql.SqlExpr -> Doc
 ppOffset Nothing = empty
-ppOffset (Just n) = text ("OFFSET " ++ show n)
+ppOffset (Just n) = text "OFFSET" <+> ppSqlExprParens n
+
+ppSqlExprParens :: HSql.SqlExpr -> Doc
+ppSqlExprParens = \case
+  HSql.ConstSqlExpr a -> text a
+  a -> parens (HPrint.ppSqlExpr a)
 
 ppFor :: Maybe Sql.LockStrength -> Doc
 ppFor Nothing       = empty

src/Opaleye/Internal/Sql.hs

@@ -47,8 +47,8 @@ data From = From {
   groupBy    :: Maybe (NEL.NonEmpty HSql.SqlExpr),
   orderBy    :: [(HSql.SqlExpr, HSql.SqlOrder)],
   distinctOn :: Maybe (NEL.NonEmpty HSql.SqlExpr),
-  limit      :: Maybe Int,
-  offset     :: Maybe Int,
+  limit      :: Maybe HSql.SqlExpr,
+  offset     :: Maybe HSql.SqlExpr,
   for        :: Maybe LockStrength
   }
           deriving Show
@@ -231,9 +231,9 @@ limit_ lo s = SelectFrom $ newSelect { tables = oneTable s
                                      , limit = limit'
                                      , offset = offset' }
   where (limit', offset') = case lo of
-          PQ.LimitOp n         -> (Just n, Nothing)
-          PQ.OffsetOp n        -> (Nothing, Just n)
-          PQ.LimitOffsetOp l o -> (Just l, Just o)
+          PQ.LimitOp n         -> (Just (sqlExpr n), Nothing)
+          PQ.OffsetOp n        -> (Nothing, Just (sqlExpr n))
+          PQ.LimitOffsetOp l o -> (Just (sqlExpr l), Just (sqlExpr o))
 
 join :: PQ.JoinType
      -> HPQ.PrimExpr

src/Opaleye/Order.hs

@@ -12,7 +12,9 @@ module Opaleye.Order ( -- * Order by
                      , descNullsFirst
                      -- * Limit and offset
                      , limit
+                     , limitField
                      , offset
+                     , offsetField
                      -- * Distinct on
                      , distinctOn
                      , distinctOnBy
@@ -30,6 +32,7 @@ module Opaleye.Order ( -- * Order by
 
 import qualified Data.Profunctor.Product.Default as D
 import qualified Opaleye.Field as F
+import qualified Opaleye.Internal.Column as C
 import qualified Opaleye.Internal.HaskellDB.PrimQuery as HPQ
 import qualified Opaleye.Internal.Order as O
 import qualified Opaleye.Internal.QueryArr as Q
@@ -123,7 +126,12 @@ SELECT * FROM (SELECT * FROM yourTable LIMIT 10) OFFSET 50
 @
 -}
 limit :: Int -> S.Select a -> S.Select a
-limit n a = Q.productQueryArr $ do
+limit = limitField . fromIntegral
+
+-- | A version of 'limit' that can accept a @Field@ rather than a
+-- constant @Int@.
+limitField :: F.Field T.SqlInt8 -> S.Select a -> S.Select a
+limitField (C.Column n) a = Q.productQueryArr $ do
   a_pq <- Q.runSimpleSelect a
   pure (O.limit' n a_pq)
 
@@ -135,7 +143,12 @@ that many result rows.
 'offset' with 'limit'.
 -}
 offset :: Int -> S.Select a -> S.Select a
-offset n a = Q.productQueryArr $ do
+offset = offsetField . fromIntegral
+
+-- | A version of 'offset' that can accept a @Field@ rather than a
+-- constant @Int@.
+offsetField :: F.Field T.SqlInt8 -> S.Select a -> S.Select a
+offsetField (C.Column n) a = Q.productQueryArr $ do
   a_pq <- Q.runSimpleSelect a
   pure (O.offset' n a_pq)