From 15aaac59832cb7671ddc47b79500429f7f6796b5 Mon Sep 17 00:00:00 2001
From: qwqtoday <lukas.leung@hotmail.com>
Date: Mon, 29 Jul 2024 20:01:59 +0800
Subject: [PATCH] added offset and limit

---
 setup/mscs/msc2836/msc2836.go | 14 +++++++++++---
 setup/mscs/msc2836/storage.go | 12 ++++++------
 2 files changed, 17 insertions(+), 9 deletions(-)

diff --git a/setup/mscs/msc2836/msc2836.go b/setup/mscs/msc2836/msc2836.go
index 83c7296b3..d71944a6d 100644
--- a/setup/mscs/msc2836/msc2836.go
+++ b/setup/mscs/msc2836/msc2836.go
@@ -25,6 +25,7 @@ import (
 	"net/http"
 	"slices"
 	"sort"
+	"strconv"
 	"strings"
 	"time"
 
@@ -158,8 +159,15 @@ func Enable(
 				}
 
 				roomID := vars["roomID"]
-
-				childrens, err := db.GetChildrensByRoomId(req.Context(), roomID, true)
+				limit, err := strconv.ParseUint(req.URL.Query().Get("limit"), 10, 64)
+				if err != nil {
+					limit = 50
+				}
+				offset, err := strconv.ParseUint(req.URL.Query().Get("from"), 10, 64)
+				if err != nil {
+					offset = 0
+				}
+				childrens, err := db.GetChildrensByRoomId(req.Context(), roomID, true, limit, offset)
 				if err != nil {
 					return util.ErrorResponse(err)
 				}
@@ -179,7 +187,7 @@ func Enable(
 					Code: 200,
 					JSON: map[string]any{
 						"chunk":      chunks,
-						"next_batch": "next_batch_token",
+						"next_batch": strconv.FormatUint(uint64(len(chunks))+offset, 10),
 					},
 				}
 			},
diff --git a/setup/mscs/msc2836/storage.go b/setup/mscs/msc2836/storage.go
index 8f971e90a..20cdf0862 100644
--- a/setup/mscs/msc2836/storage.go
+++ b/setup/mscs/msc2836/storage.go
@@ -44,7 +44,7 @@ type Database interface {
 	// MarkChildrenExplored sets the 'explored' flag on this event to `true`.
 	MarkChildrenExplored(ctx context.Context, eventID string) error
 
-	GetChildrensByRoomId(ctx context.Context, roomID string, recentFirst bool) ([]eventInfo, error)
+	GetChildrensByRoomId(ctx context.Context, roomID string, recentFirst bool, limit uint64, offset uint64) ([]eventInfo, error)
 }
 
 type DB struct {
@@ -150,10 +150,10 @@ func newPostgresDatabase(conMan *sqlutil.Connections, dbOpts *config.DatabaseOpt
 		LEFT JOIN msc2836_nodes ON msc2836_edges.child_event_id = msc2836_nodes.event_id
 		WHERE room_id = $1
 		ORDER BY origin_server_ts`
-	if d.selectChildrenByRoomIdForParentOldestFirstStmt, err = d.db.Prepare(selectChildrenByRoomIdQuery + "ASC"); err != nil {
+	if d.selectChildrenByRoomIdForParentOldestFirstStmt, err = d.db.Prepare(selectChildrenByRoomIdQuery + "ASC LIMIT $2 OFFSET $3"); err != nil {
 		return nil, err
 	}
-	if d.selectChildrenByRoomIdForParentRecentFirstStmt, err = d.db.Prepare(selectChildrenByRoomIdQuery + "DESC"); err != nil {
+	if d.selectChildrenByRoomIdForParentRecentFirstStmt, err = d.db.Prepare(selectChildrenByRoomIdQuery + "DESC LIMIT $2 OFFSET $3"); err != nil {
 		return nil, err
 	}
 
@@ -342,13 +342,13 @@ func (p *DB) ParentForChild(ctx context.Context, eventID, relType string) (*even
 	return &ei, nil
 }
 
-func (p *DB) GetChildrensByRoomId(ctx context.Context, roomID string, recentFirst bool) ([]eventInfo, error) {
+func (p *DB) GetChildrensByRoomId(ctx context.Context, roomID string, recentFirst bool, limit uint64, offset uint64) ([]eventInfo, error) {
 	var rows *sql.Rows
 	var err error
 	if recentFirst {
-		rows, err = p.selectChildrenForParentRecentFirstStmt.QueryContext(ctx, roomID)
+		rows, err = p.selectChildrenForParentRecentFirstStmt.QueryContext(ctx, roomID, limit, offset)
 	} else {
-		rows, err = p.selectChildrenByRoomIdForParentOldestFirstStmt.QueryContext(ctx, roomID)
+		rows, err = p.selectChildrenByRoomIdForParentOldestFirstStmt.QueryContext(ctx, roomID, limit, offset)
 	}
 	if err != nil {
 		return nil, err