Paginate listRecords based on cursor (bluesky-social#751)

paginate listRecords based on cursor
qwexvf · Apr 4, 2023 · ba52791 · ba52791
1 parent 25c23b6
commit ba52791
Show file tree

Hide file tree

Showing 8 changed files with 66 additions and 41 deletions.
diff --git a/lexicons/com/atproto/repo/listRecords.json b/lexicons/com/atproto/repo/listRecords.json
@@ -12,8 +12,9 @@
           "repo": {"type": "string", "format": "at-identifier", "description": "The handle or DID of the repo."},
           "collection": {"type": "string", "format": "nsid", "description": "The NSID of the record type."},
           "limit": {"type": "integer", "minimum": 1, "maximum": 100, "default": 50, "description": "The number of records to return."},
-          "rkeyStart": {"type": "string", "description": "The lowest sort-ordered rkey to start from (exclusive)"},
-          "rkeyEnd": {"type": "string", "description": "The highest sort-ordered rkey to stop at (exclusive)"},
+          "cursor": {"type": "string"},
+          "rkeyStart": {"type": "string", "description": "DEPRECATED: The lowest sort-ordered rkey to start from (exclusive)"},
+          "rkeyEnd": {"type": "string", "description": "DEPRECATED: The highest sort-ordered rkey to stop at (exclusive)"},
           "reverse": {"type": "boolean", "description": "Reverse the order of the returned records?"}
         }
       },

diff --git a/packages/api/src/client/lexicons.ts b/packages/api/src/client/lexicons.ts
@@ -1428,6 +1428,9 @@ export const schemaDict = {
               default: 50,
               description: 'The number of records to return.',
             },
+            cursor: {
+              type: 'string',
+            },
             rkeyStart: {
               type: 'string',
               description:

diff --git a/packages/api/src/client/types/com/atproto/repo/listRecords.ts b/packages/api/src/client/types/com/atproto/repo/listRecords.ts
@@ -14,6 +14,7 @@ export interface QueryParams {
   collection: string
   /** The number of records to return. */
   limit?: number
+  cursor?: string
   /** The lowest sort-ordered rkey to start from (exclusive) */
   rkeyStart?: string
   /** The highest sort-ordered rkey to stop at (exclusive) */

diff --git a/packages/pds/src/api/com/atproto/repo/listRecords.ts b/packages/pds/src/api/com/atproto/repo/listRecords.ts
@@ -5,23 +5,30 @@ import AppContext from '../../../../context'
 
 export default function (server: Server, ctx: AppContext) {
   server.com.atproto.repo.listRecords(async ({ params }) => {
-    const { repo, collection, limit, rkeyStart, rkeyEnd, reverse } = params
+    const {
+      repo,
+      collection,
+      limit = 50,
+      cursor,
+      rkeyStart,
+      rkeyEnd,
+      reverse = false,
+    } = params
 
     const did = await ctx.services.account(ctx.db).getDidForActor(repo)
     if (!did) {
       throw new InvalidRequestError(`Could not find repo: ${repo}`)
     }
 
-    const records = await ctx.services
-      .record(ctx.db)
-      .listRecordsForCollection(
-        did,
-        collection,
-        limit || 50,
-        reverse || false,
-        rkeyStart,
-        rkeyEnd,
-      )
+    const records = await ctx.services.record(ctx.db).listRecordsForCollection({
+      did,
+      collection,
+      limit,
+      reverse,
+      cursor,
+      rkeyStart,
+      rkeyEnd,
+    })
 
     const lastRecord = records.at(-1)
     const lastUri = lastRecord && new AtUri(lastRecord?.uri)

diff --git a/packages/pds/src/lexicon/lexicons.ts b/packages/pds/src/lexicon/lexicons.ts
@@ -1428,6 +1428,9 @@ export const schemaDict = {
               default: 50,
               description: 'The number of records to return.',
             },
+            cursor: {
+              type: 'string',
+            },
             rkeyStart: {
               type: 'string',
               description:

diff --git a/packages/pds/src/lexicon/types/com/atproto/repo/listRecords.ts b/packages/pds/src/lexicon/types/com/atproto/repo/listRecords.ts
@@ -15,6 +15,7 @@ export interface QueryParams {
   collection: string
   /** The number of records to return. */
   limit: number
+  cursor?: string
   /** The lowest sort-ordered rkey to start from (exclusive) */
   rkeyStart?: string
   /** The highest sort-ordered rkey to stop at (exclusive) */

diff --git a/packages/pds/src/services/record/index.ts b/packages/pds/src/services/record/index.ts
@@ -105,15 +105,27 @@ export class RecordService {
     return collections.map((row) => row.collection)
   }
 
-  async listRecordsForCollection(
-    did: string,
-    collection: string,
-    limit: number,
-    reverse: boolean,
-    rkeyStart?: string,
-    rkeyEnd?: string,
-    includeSoftDeleted = false,
-  ): Promise<{ uri: string; cid: string; value: object }[]> {
+  async listRecordsForCollection(opts: {
+    did: string
+    collection: string
+    limit: number
+    reverse: boolean
+    cursor?: string
+    rkeyStart?: string
+    rkeyEnd?: string
+    includeSoftDeleted?: boolean
+  }): Promise<{ uri: string; cid: string; value: object }[]> {
+    const {
+      did,
+      collection,
+      limit,
+      reverse,
+      cursor,
+      rkeyStart,
+      rkeyEnd,
+      includeSoftDeleted = false,
+    } = opts
+
     const { ref } = this.db.db.dynamic
     let builder = this.db.db
       .selectFrom('record')
@@ -131,11 +143,20 @@ export class RecordService {
       .limit(limit)
       .selectAll()
 
-    if (rkeyStart !== undefined) {
-      builder = builder.where('record.rkey', '>', rkeyStart)
-    }
-    if (rkeyEnd !== undefined) {
-      builder = builder.where('record.rkey', '<', rkeyEnd)
+    // prioritize cursor but fall back to soon-to-be-depcreated rkey start/end
+    if (cursor !== undefined) {
+      if (reverse) {
+        builder = builder.where('record.rkey', '>', cursor)
+      } else {
+        builder = builder.where('record.rkey', '<', cursor)
+      }
+    } else {
+      if (rkeyStart !== undefined) {
+        builder = builder.where('record.rkey', '>', rkeyStart)
+      }
+      if (rkeyEnd !== undefined) {
+        builder = builder.where('record.rkey', '<', rkeyEnd)
+      }
     }
     const res = await builder.execute()
     return res.map((row) => {

diff --git a/packages/pds/tests/crud.test.ts b/packages/pds/tests/crud.test.ts
@@ -324,7 +324,7 @@ describe('crud operations', () => {
       const paginator = async (cursor?: string) => {
         const res = await agent.api.app.bsky.feed.post.list({
           repo: alice.did,
-          rkeyEnd: cursor,
+          cursor,
           limit: 2,
         })
         return res
@@ -349,7 +349,7 @@ describe('crud operations', () => {
         const res = await agent.api.app.bsky.feed.post.list({
           repo: alice.did,
           reverse: true,
-          rkeyStart: cursor,
+          cursor,
           limit: 2,
         })
         return res
@@ -369,18 +369,6 @@ describe('crud operations', () => {
       expect(results(paginatedAll)).toEqual(results([full]))
     })
 
-    it('between two records', async () => {
-      const list = await agent.api.app.bsky.feed.post.list({
-        repo: alice.did,
-        rkeyStart: uri1.rkey,
-        rkeyEnd: uri5.rkey,
-      })
-      expect(list.records.length).toBe(3)
-      expect(list.records[0].uri).toBe(uri4.toString())
-      expect(list.records[1].uri).toBe(uri3.toString())
-      expect(list.records[2].uri).toBe(uri2.toString())
-    })
-
     it('reverses', async () => {
       const forwards = await agent.api.app.bsky.feed.post.list({
         repo: alice.did,