From a4c227ce41f7584bf3f6589c8512eeb827ea400c Mon Sep 17 00:00:00 2001
From: Paul <paulmakles@gmail.com>
Date: Wed, 7 Jul 2021 16:21:21 +0100
Subject: [PATCH] Add "nearby" message query.

---
 set_version.sh                       |   2 +-
 src/routes/channels/message_query.rs | 127 ++++++++++++++++++++-------
 src/version.rs                       |   2 +-
 3 files changed, 95 insertions(+), 36 deletions(-)

diff --git a/set_version.sh b/set_version.sh
index 372bac6..e256b67 100755
--- a/set_version.sh
+++ b/set_version.sh
@@ -1,3 +1,3 @@
 #!/bin/bash
-export version=0.5.1-alpha.5
+export version=0.5.1-alpha.6
 echo "pub const VERSION: &str = \"${version}\";" > src/version.rs
diff --git a/src/routes/channels/message_query.rs b/src/routes/channels/message_query.rs
index 4492e83..672701b 100644
--- a/src/routes/channels/message_query.rs
+++ b/src/routes/channels/message_query.rs
@@ -3,7 +3,7 @@ use std::collections::HashSet;
 use crate::database::*;
 use crate::util::result::{Error, Result};
 
-use futures::StreamExt;
+use futures::{StreamExt, try_join};
 use mongodb::{
     bson::{doc, from_document},
     options::FindOptions,
@@ -28,6 +28,11 @@ pub struct Options {
     #[validate(length(min = 26, max = 26))]
     after: Option<String>,
     sort: Option<Sort>,
+    // Specifying 'nearby' ignores 'before', 'after' and 'sort'.
+    // It will also take half of limit rounded as the limits to each side.
+    // It also fetches the message ID specified.
+    #[validate(length(min = 26, max = 26))]
+    nearby: Option<String>,
     include_users: Option<bool>,
 }
 
@@ -48,46 +53,100 @@ pub async fn req(user: User, target: Ref, options: Form<Options>) -> Result<Json
         Err(Error::MissingPermission)?
     }
 
-    let mut query = doc! { "channel": target.id() };
-
-    if let Some(before) = &options.before {
-        query.insert("_id", doc! { "$lt": before });
-    }
-
-    if let Some(after) = &options.after {
-        query.insert("_id", doc! { "$gt": after });
-    }
+    let mut messages = vec![];
 
-    let sort = if let Sort::Latest = options.sort.as_ref().unwrap_or_else(|| &Sort::Latest) {
-        -1
-    } else {
-        1
-    };
-    let mut cursor = get_collection("messages")
-        .find(
-            query,
-            FindOptions::builder()
-                .limit(options.limit.unwrap_or(50))
-                .sort(doc! {
-                    "_id": sort
-                })
-                .build(),
+    let collection = get_collection("messages");
+    let limit = options.limit.unwrap_or(50);
+    let channel = target.id();
+    if let Some(nearby) = &options.nearby {
+        let cursors = try_join!(
+            collection.find(
+                doc! {
+                    "channel": channel,
+                    "_id": {
+                        "$gte": &nearby
+                    }
+                },
+                FindOptions::builder()
+                    .limit(limit / 2 + 1)
+                    .sort(doc! {
+                        "_id": 1
+                    })
+                    .build(),
+            ),
+            collection.find(
+                doc! {
+                    "channel": channel,
+                    "_id": {
+                        "$lt": &nearby
+                    }
+                },
+                FindOptions::builder()
+                    .limit(limit / 2)
+                    .sort(doc! {
+                        "_id": -1
+                    })
+                    .build(),
+            )
         )
-        .await
         .map_err(|_| Error::DatabaseError {
             operation: "find",
             with: "messages",
         })?;
 
-    let mut messages = vec![];
-    while let Some(result) = cursor.next().await {
-        if let Ok(doc) = result {
-            messages.push(
-                from_document::<Message>(doc).map_err(|_| Error::DatabaseError {
-                    operation: "from_document",
-                    with: "message",
-                })?,
-            );
+        for mut cursor in [ cursors.0, cursors.1 ] {
+            while let Some(result) = cursor.next().await {
+                if let Ok(doc) = result {
+                    messages.push(
+                        from_document::<Message>(doc).map_err(|_| Error::DatabaseError {
+                            operation: "from_document",
+                            with: "message",
+                        })?,
+                    );
+                }
+            }
+        }
+    } else {
+        let mut query = doc! { "channel": target.id() };
+        if let Some(before) = &options.before {
+            query.insert("_id", doc! { "$lt": before });
+        }
+
+        if let Some(after) = &options.after {
+            query.insert("_id", doc! { "$gt": after });
+        }
+
+        let sort: i32 = if let Sort::Latest = options.sort.as_ref().unwrap_or_else(|| &Sort::Latest) {
+            -1
+        } else {
+            1
+        };
+
+        let mut cursor = collection
+            .find(
+                query,
+                FindOptions::builder()
+                    .limit(limit)
+                    .sort(doc! {
+                        "_id": sort
+                    })
+                    .build(),
+            )
+            .await
+            .map_err(|_| Error::DatabaseError {
+                operation: "find",
+                with: "messages",
+            })?;
+
+        while let Some(result) = cursor.next().await {
+            if let Ok(doc) = result {
+                messages.push(
+                    from_document::<Message>(doc).map_err(|_| Error::DatabaseError {
+                        operation: "from_document",
+                        with: "message",
+                    })?,
+                );
+            }
         }
     }
 
diff --git a/src/version.rs b/src/version.rs
index ef59e77..818c98d 100644
--- a/src/version.rs
+++ b/src/version.rs
@@ -1 +1 @@
-pub const VERSION: &str = "0.5.1-alpha.5";
+pub const VERSION: &str = "0.5.1-alpha.6";
-- 
GitLab