Add partitionBy_ and orderBy_ support to windows

src/Database/Esqueleto/Experimental/WindowFunctions.hs

@@ -7,28 +7,38 @@
 {-# LANGUAGE RankNTypes #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE StandaloneDeriving #-}
+{-# LANGUAGE TypeApplications #-}
 {-# LANGUAGE TypeFamilies #-}
 {-# LANGUAGE UndecidableInstances #-}
 
 module Database.Esqueleto.Experimental.WindowFunctions
     where
 
+import Control.Arrow (first)
 import Data.Coerce (coerce)
 import Data.Int (Int64)
 import Data.Semigroup (First(..))
 import qualified Data.Text.Lazy.Builder as TLB
 import Database.Esqueleto.Experimental.Aggregates
 import Database.Esqueleto.Internal.Internal
-       ( NeedParens(..)
+       ( IdentInfo
+       , NeedParens(..)
+       , OrderBy
        , SideData(..)
+       , SomeValue(..)
        , SqlExpr(..)
        , SqlQuery(..)
        , SqlSelect(..)
        , ToSomeValues(..)
        , Value(..)
+       , asc
        , noMeta
+       , parens
+       , parensM
        , select
+       , uncommas'
        , unsafeSqlFunction
+       , val
        , (?.)
        , (^.)
        )
@@ -42,23 +52,60 @@ import Database.Esqueleto.Internal.PersistentImport
        , fromPersistValue
        )
 
+--( "LAG(?) OVER (PARTITION BY ?, ? ORDER BY ? ASC ROWS BETWEEN ? PRECEEDING AND UNBOUNDED FOLLOWING)"
+--, [PersistInt64 10,PersistInt64 10,PersistBool True,PersistInt64 10,PersistInt64 1]
+--)
+example =
+    lag_ (val @Int64 10) Nothing Nothing `over_`
+        (  partitionBy_ (val @Int64 10, val True)
+        <> frame_ (rows $ between (preceeding 1) unboundedFollowing)
+        <> orderBy_ [asc (val @Int64 10)]
+        )
 
+data NeedsWindow a
+
+lag :: SqlExpr (Value a) -> WindowExpr a
+lag v = lag_ v Nothing Nothing
+
+lag_ :: SqlExpr a -> Maybe (SqlExpr Int64) -> Maybe (SqlExpr a) -> WindowExpr a
+lag_ v mOffset mDefaultVal =
+    coerce $
+    case (mOffset, mDefaultVal) of
+      (Just offset, Just defaultVal) ->
+          unsafeSqlFunction "LAG" (v, offset, defaultVal)
+      (Just offset, Nothing) ->
+          unsafeSqlFunction "LAG" (v, offset)
+      (Nothing, _) ->
+          unsafeSqlFunction "LAG" v
+
+-- Phantom helper type
+data PartitionBy
 data Window = Window
-    { windowPartitionBy :: Maybe (First (TLB.Builder, [PersistValue]))
+    { windowPartitionBy :: Maybe (First (SqlExpr PartitionBy))
-    , windowOrderBy     :: Maybe (First (TLB.Builder, [PersistValue]))
+    , windowOrderBy     :: Maybe [SqlExpr OrderBy]
     , windowFrame       :: Maybe (First Frame)
     }
 
+partitionBy_ :: ToSomeValues a => a -> Window
+partitionBy_ expr = mempty{windowPartitionBy = Just $ First $ ERaw noMeta $ \_ info ->
+    let (b, v) = uncommas' $ fmap (\(SomeValue (ERaw _ f)) -> f Never info) $ toSomeValues expr
+    in ("PARTITION BY " <> b, v)
+                          }
+
+orderBy_ :: [SqlExpr OrderBy] -> Window
+orderBy_ []    = mempty
+orderBy_ exprs = mempty{windowOrderBy=Just exprs}
+
 class RenderWindow a where
-    renderWindow :: a -> (TLB.Builder, [PersistValue])
+    renderWindow :: IdentInfo -> a -> (TLB.Builder, [PersistValue])
 instance RenderWindow () where
-    renderWindow = mempty
+    renderWindow _ = mempty
 instance RenderWindow Window where
-    renderWindow window =
+    renderWindow info window =
-        let (p, pVal) = maybe mempty getFirst $ windowPartitionBy window
+        let (partition, partitionVal) = maybe mempty ((\(ERaw _ f) -> f Never info) . getFirst) (windowPartitionBy window)
-            (o, oVal) = maybe mempty getFirst $ windowOrderBy window
+            (order, orderVal) = maybe mempty (first ((<>) " ORDER BY ") . uncommas' . fmap (\(ERaw _ f) -> f Never info)) (windowOrderBy window)
-            (f, fVal) = maybe mempty (renderWindow . getFirst) (windowFrame window)
+            (frame, frameVal) = maybe mempty (renderWindow info . getFirst) (windowFrame window)
-        in (p <> o <> f, pVal <> oVal <> fVal)
+        in (partition <> order <> frame, partitionVal <> orderVal <> frameVal)
 
 instance Semigroup Window where
     (Window a b c) <> (Window a' b' c') = Window (a <> a') (b <> b') (c <> c')
@@ -69,15 +116,15 @@ instance Monoid Window where
 
 data Frame = Frame (Maybe FrameKind) FrameBody (Maybe FrameExclusion)
 
-frame :: ToFrame frame => frame -> Window
+frame_ :: ToFrame frame => frame -> Window
-frame f = mempty{windowFrame = Just $ First $ toFrame f}
+frame_ f = mempty{windowFrame = Just $ First $ toFrame f}
 
 instance RenderWindow Frame where
-    renderWindow (Frame mKind frameBody mExclusion) =
+    renderWindow info (Frame mKind frameBody mExclusion) =
-        let (kind, kindVals) = maybe mempty renderWindow mKind
+        let (kind, kindVals) = maybe mempty (renderWindow info) mKind
-            (exclusion, exclusionVals) = maybe mempty renderWindow mExclusion
+            (exclusion, exclusionVals) = maybe mempty (renderWindow info) mExclusion
-            (body, bodyVals) = renderWindow frameBody
+            (body, bodyVals) = renderWindow info frameBody
-        in (kind <> body <> exclusion, kindVals <> bodyVals <> exclusionVals)
+        in (" " <> kind <> body <> exclusion, kindVals <> bodyVals <> exclusionVals)
 
 class ToFrame a where
     toFrame :: a -> Frame
@@ -87,7 +134,7 @@ instance ToFrame Frame where
 newtype FrameKind = FrameKind { unFrameKind :: (TLB.Builder, [PersistValue]) }
 
 instance RenderWindow FrameKind where
-    renderWindow = unFrameKind
+    renderWindow _ = unFrameKind
 
 frameKind :: ToFrame frame => TLB.Builder -> frame -> Frame
 frameKind tlb frame =
@@ -106,7 +153,7 @@ groups = frameKind "GROUPS"
 newtype FrameExclusion = FrameExclusion { unFrameExclusion :: (TLB.Builder, [PersistValue]) }
 
 instance RenderWindow FrameExclusion where
-    renderWindow = unFrameExclusion
+    renderWindow _ = unFrameExclusion
 
 frameExclusion :: ToFrame frame => TLB.Builder -> frame -> Frame
 frameExclusion tlb frame =
@@ -125,6 +172,13 @@ excludeTies = frameExclusion "TIES"
 excludeNoOthers :: ToFrame frame => frame -> Frame
 excludeNoOthers = frameExclusion "NO OTHERS"
 
+-- In order to prevent runtime errors we do some magic rewriting of queries that wouldn't be valid SQL.
+-- In the case of an implicit frame end `following 10` would become BETWEEN 10 FOLLOWING AND CURRENT ROW
+-- This is illegal so `following 10` instead becomes `BETWEEN CURRENT_ROW AND 10 FOLLOWING`
+-- Additionally `BETWEEN` requires that the frame start be before the frame end.
+-- To prevent this error the frame will be flipped automatically.
+-- i.e. `between (following 10) (preceeding 10)` becomes `BETWEEEN 10 PRECEEDING AND 10 FOLLOWING`
+-- therefore `between (following 10) (preceeding 10) === between (preceeding 10) (following 10)
 data FrameBody
     = FrameStart FrameRange
     | FrameBetween FrameRange FrameRange
@@ -133,21 +187,21 @@ instance ToFrame FrameBody where
     toFrame b = Frame Nothing b Nothing
 
 instance RenderWindow FrameBody where
-    renderWindow (FrameStart (FrameRangeFollowing b)) = renderWindow (FrameBetween FrameRangeCurrentRow (FrameRangeFollowing b))
+    renderWindow info (FrameStart (FrameRangeFollowing b)) = renderWindow info (FrameBetween FrameRangeCurrentRow (FrameRangeFollowing b))
-    renderWindow (FrameStart f) = renderWindow f
+    renderWindow info (FrameStart f) = renderWindow info f
-    renderWindow (FrameBetween startRange endRange)
+    renderWindow info (FrameBetween startRange endRange)
-          | startRange > endRange = renderWindow (FrameBetween endRange startRange)
+          | startRange > endRange = renderWindow info (FrameBetween endRange startRange)
-    renderWindow (FrameBetween r r') =
+    renderWindow  info (FrameBetween r r') =
-        let (b, v) = renderWindow r
+        let (b, v) = renderWindow info r
-            (b', v') = renderWindow r'
+            (b', v') = renderWindow info r'
         in ("BETWEEN " <> b <> " AND " <> b', v <> v')
 
 instance ToFrame FrameRange where
     toFrame r = Frame Nothing (FrameStart r) Nothing
 instance RenderWindow FrameRange where
-    renderWindow (FrameRangeCurrentRow) = ("CURRENT ROW", [])
+    renderWindow _ (FrameRangeCurrentRow) = ("CURRENT ROW", [])
-    renderWindow (FrameRangePreceeding bounds) = renderBounds bounds <> (" PRECEEDING", [])
+    renderWindow _ (FrameRangePreceeding bounds) = renderBounds bounds <> (" PRECEEDING", [])
-    renderWindow (FrameRangeFollowing bounds) = renderBounds bounds <> (" FOLLOWING", [])
+    renderWindow _ (FrameRangeFollowing bounds) = renderBounds bounds <> (" FOLLOWING", [])
 
 renderBounds :: FrameRangeBound -> (TLB.Builder, [PersistValue])
 renderBounds (FrameRangeUnbounded) = ("UNBOUNDED", [])
@@ -177,6 +231,7 @@ data FrameRangeBound
 
 instance Ord FrameRangeBound where
     FrameRangeUnbounded <= FrameRangeBounded _ = False
+    FrameRangeUnbounded <= FrameRangeUnbounded = True
     FrameRangeBounded _ <= FrameRangeUnbounded = True
     FrameRangeBounded a <= FrameRangeBounded b = a <= b
 
@@ -213,7 +268,7 @@ newtype WindowExpr a = WindowExpr { unsafeWindowExpr :: SqlExpr a }
 instance Over WindowExpr where
     (WindowExpr (ERaw _ f)) `over_` window = ERaw noMeta $ \p info ->
         let (b, v) = f Never info
-            (w, vw) = renderWindow window
+            (w, vw) = renderWindow info window
-        in (b <> " OVER (" <> w <> ")", v <> vw)
+        in (parensM p $ b <> " OVER " <> parens w , v <> vw)
 
 deriving via WindowExpr instance Over SqlAggregate

src/Database/Esqueleto/Internal/Internal.hs

@@ -2312,7 +2312,10 @@ instance UnsafeSqlFunctionArgument () where
     toArgList _ = []
 
 instance UnsafeSqlFunctionArgument (SqlExpr a) where
-    toArgList = (:[]) . veryUnsafeCoerceSqlExpr
+    toArgList = (:[]) . coerce 
+
+instance UnsafeSqlFunctionArgument (Maybe (SqlExpr a)) where
+    toArgList = maybe [] ((:[]) . coerce)
 
 instance UnsafeSqlFunctionArgument a => UnsafeSqlFunctionArgument [a] where
   toArgList = concatMap toArgList