From 9a048f2f85599337c8a9b726eba18003257cfde5 Mon Sep 17 00:00:00 2001
From: Andrea Condoluci <andrea.condoluci@tweag.io>
Date: Tue, 15 Dec 2020 14:36:10 +0000
Subject: [PATCH] Add allowDiskUse option for aggregates

Aggregation pipeline stages have a limit of 100Mb of RAM.
In case of large datasets, one can cross that limit by setting
allowDiskUse = True and making stages write data
to temporary files.

See also https://docs.mongodb.com/manual/reference/method/db.collection.aggregate .
---
 Database/MongoDB/Query.hs | 22 +++++++++++++++++-----
 1 file changed, 17 insertions(+), 5 deletions(-)

diff --git a/Database/MongoDB/Query.hs b/Database/MongoDB/Query.hs
index 1d76de0..c814317 100644
--- a/Database/MongoDB/Query.hs
+++ b/Database/MongoDB/Query.hs
@@ -1338,16 +1338,28 @@ aggregate :: (MonadIO m, MonadFail m) => Collection -> Pipeline -> Action m [Doc
 aggregate aColl agg = do
     aggregateCursor aColl agg def >>= rest
 
-data AggregateConfig = AggregateConfig {}
-                       deriving Show
+data AggregateConfig = AggregateConfig
+  { allowDiskUse :: Bool -- ^ Enable writing to temporary files (aggregations have a 100Mb RAM limit)
+  }
+  deriving Show
 
 instance Default AggregateConfig where
-  def = AggregateConfig {}
+  def = AggregateConfig
+    { allowDiskUse = False
+    }
+
+aggregateCommand :: Collection -> Pipeline -> AggregateConfig -> Document
+aggregateCommand aColl agg AggregateConfig {..} =
+  [ "aggregate" =: aColl
+  , "pipeline" =: agg
+  , "cursor" =: ([] :: Document)
+  , "allowDiskUse" =: allowDiskUse
+  ]
 
 aggregateCursor :: (MonadIO m, MonadFail m) => Collection -> Pipeline -> AggregateConfig -> Action m Cursor
 -- ^ Runs an aggregate and unpacks the result. See <http://docs.mongodb.org/manual/core/aggregation/> for details.
-aggregateCursor aColl agg _ = do
-    response <- runCommand ["aggregate" =: aColl, "pipeline" =: agg, "cursor" =: ([] :: Document)]
+aggregateCursor aColl agg cfg = do
+    response <- runCommand (aggregateCommand aColl agg cfg)
     getCursorFromResponse aColl response
       >>= either (liftIO . throwIO . AggregateFailure) return