@atproto/bsky 0.0.124 → 0.0.125

package/src/data-plane/server/db/pagination.ts

@@ -2,8 +2,8 @@ import { sql } from 'kysely'
 import { InvalidRequestError } from '@atproto/xrpc-server'
 import { AnyQb, DbRef } from './util'
 
-export type Cursor = { primary: string; secondary: string }
-export type LabeledResult = {
+type KeysetCursor = { primary: string; secondary: string }
+type KeysetLabeledResult = {
   primary: string | number
   secondary: string | number
 }
@@ -22,14 +22,14 @@ export type LabeledResult = {
  *   Result -*-> LabeledResult <-*-> Cursor <--> packed/string cursor
  *                     ↳ SQL Condition
  */
-export abstract class GenericKeyset<R, LR extends LabeledResult> {
+export abstract class GenericKeyset<R, LR extends KeysetLabeledResult> {
   constructor(
     public primary: DbRef,
     public secondary: DbRef,
   ) {}
   abstract labelResult(result: R): LR
-  abstract labeledResultToCursor(labeled: LR): Cursor
-  abstract cursorToLabeledResult(cursor: Cursor): LR
+  abstract labeledResultToCursor(labeled: LR): KeysetCursor
+  abstract cursorToLabeledResult(cursor: KeysetCursor): LR
   packFromResult(results: R | R[]): string | undefined {
     const result = Array.isArray(results) ? results.at(-1) : results
     if (!result) return
@@ -45,11 +45,11 @@ export abstract class GenericKeyset<R, LR extends LabeledResult> {
     if (!cursor) return
     return this.cursorToLabeledResult(cursor)
   }
-  packCursor(cursor?: Cursor): string | undefined {
+  packCursor(cursor?: KeysetCursor): string | undefined {
     if (!cursor) return
     return `${cursor.primary}__${cursor.secondary}`
   }
-  unpackCursor(cursorStr?: string): Cursor | undefined {
+  unpackCursor(cursorStr?: string): KeysetCursor | undefined {
     if (!cursorStr) return
     const result = cursorStr.split('__')
     const [primary, secondary, ...others] = result
@@ -79,10 +79,43 @@ export abstract class GenericKeyset<R, LR extends LabeledResult> {
       }
     }
   }
+  paginate<QB extends AnyQb>(
+    qb: QB,
+    opts: {
+      limit?: number
+      cursor?: string
+      direction?: 'asc' | 'desc'
+      tryIndex?: boolean
+      // By default, pg does nullsFirst
+      nullsLast?: boolean
+    },
+  ): QB {
+    const { limit, cursor, direction = 'desc', tryIndex, nullsLast } = opts
+    const keysetSql = this.getSql(this.unpack(cursor), direction, tryIndex)
+    return qb
+      .if(!!limit, (q) => q.limit(limit as number))
+      .if(!nullsLast, (q) =>
+        q.orderBy(this.primary, direction).orderBy(this.secondary, direction),
+      )
+      .if(!!nullsLast, (q) =>
+        q
+          .orderBy(
+            direction === 'asc'
+              ? sql`${this.primary} asc nulls last`
+              : sql`${this.primary} desc nulls last`,
+          )
+          .orderBy(
+            direction === 'asc'
+              ? sql`${this.secondary} asc nulls last`
+              : sql`${this.secondary} desc nulls last`,
+          ),
+      )
+      .if(!!keysetSql, (qb) => (keysetSql ? qb.where(keysetSql) : qb)) as QB
+  }
 }
 
 type SortAtCidResult = { sortAt: string; cid: string }
-type TimeCidLabeledResult = Cursor
+type TimeCidLabeledResult = KeysetCursor
 
 export class TimeCidKeyset<
   TimeCidResult = SortAtCidResult,
@@ -97,7 +130,7 @@ export class TimeCidKeyset<
       secondary: labeled.secondary,
     }
   }
-  cursorToLabeledResult(cursor: Cursor) {
+  cursorToLabeledResult(cursor: KeysetCursor) {
     const primaryDate = new Date(parseInt(cursor.primary, 10))
     if (isNaN(primaryDate.getTime())) {
       throw new InvalidRequestError('Malformed cursor')
@@ -127,6 +160,9 @@ export class IndexedAtDidKeyset extends TimeCidKeyset<{
   }
 }
 
+/**
+ * This is being deprecated. Use {@link GenericKeyset#paginate} instead.
+ */
 export const paginate = <
   QB extends AnyQb,
   K extends GenericKeyset<unknown, any>,
@@ -142,32 +178,119 @@ export const paginate = <
     nullsLast?: boolean
   },
 ): QB => {
-  const {
-    limit,
-    cursor,
-    keyset,
-    direction = 'desc',
-    tryIndex,
-    nullsLast,
-  } = opts
-  const keysetSql = keyset.getSql(keyset.unpack(cursor), direction, tryIndex)
-  return qb
-    .if(!!limit, (q) => q.limit(limit as number))
-    .if(!nullsLast, (q) =>
-      q.orderBy(keyset.primary, direction).orderBy(keyset.secondary, direction),
-    )
-    .if(!!nullsLast, (q) =>
-      q
-        .orderBy(
-          direction === 'asc'
-            ? sql`${keyset.primary} asc nulls last`
-            : sql`${keyset.primary} desc nulls last`,
-        )
-        .orderBy(
+  return opts.keyset.paginate(qb, opts)
+}
+
+type SingleKeyCursor = {
+  primary: string
+}
+
+type SingleKeyLabeledResult = {
+  primary: string | number
+}
+
+/**
+ * GenericSingleKey is similar to {@link GenericKeyset} but for a single key cursor.
+ */
+export abstract class GenericSingleKey<R, LR extends SingleKeyLabeledResult> {
+  constructor(public primary: DbRef) {}
+  abstract labelResult(result: R): LR
+  abstract labeledResultToCursor(labeled: LR): SingleKeyCursor
+  abstract cursorToLabeledResult(cursor: SingleKeyCursor): LR
+  packFromResult(results: R | R[]): string | undefined {
+    const result = Array.isArray(results) ? results.at(-1) : results
+    if (!result) return
+    return this.pack(this.labelResult(result))
+  }
+  pack(labeled?: LR): string | undefined {
+    if (!labeled) return
+    const cursor = this.labeledResultToCursor(labeled)
+    return this.packCursor(cursor)
+  }
+  unpack(cursorStr?: string): LR | undefined {
+    const cursor = this.unpackCursor(cursorStr)
+    if (!cursor) return
+    return this.cursorToLabeledResult(cursor)
+  }
+  packCursor(cursor?: SingleKeyCursor): string | undefined {
+    if (!cursor) return
+    return cursor.primary
+  }
+  unpackCursor(cursorStr?: string): SingleKeyCursor | undefined {
+    if (!cursorStr) return
+    const result = cursorStr.split('__')
+    const [primary, ...others] = result
+    if (!primary || others.length > 0) {
+      throw new InvalidRequestError('Malformed cursor')
+    }
+    return {
+      primary,
+    }
+  }
+  getSql(labeled?: LR, direction?: 'asc' | 'desc') {
+    if (labeled === undefined) return
+    if (direction === 'asc') {
+      return sql`${this.primary} > ${labeled.primary}`
+    }
+    return sql`${this.primary} < ${labeled.primary}`
+  }
+  paginate<QB extends AnyQb>(
+    qb: QB,
+    opts: {
+      limit?: number
+      cursor?: string
+      direction?: 'asc' | 'desc'
+      // By default, pg does nullsFirst
+      nullsLast?: boolean
+    },
+  ): QB {
+    const { limit, cursor, direction = 'desc', nullsLast } = opts
+    const keySql = this.getSql(this.unpack(cursor), direction)
+    return qb
+      .if(!!limit, (q) => q.limit(limit as number))
+      .if(!nullsLast, (q) => q.orderBy(this.primary, direction))
+      .if(!!nullsLast, (q) =>
+        q.orderBy(
           direction === 'asc'
-            ? sql`${keyset.secondary} asc nulls last`
-            : sql`${keyset.secondary} desc nulls last`,
+            ? sql`${this.primary} asc nulls last`
+            : sql`${this.primary} desc nulls last`,
         ),
-    )
-    .if(!!keysetSql, (qb) => (keysetSql ? qb.where(keysetSql) : qb)) as QB
+      )
+      .if(!!keySql, (qb) => (keySql ? qb.where(keySql) : qb)) as QB
+  }
+}
+
+type SortAtResult = { sortAt: string }
+type TimeLabeledResult = SingleKeyCursor
+
+export class IsoTimeKey<TimeResult = SortAtResult> extends GenericSingleKey<
+  TimeResult,
+  TimeLabeledResult
+> {
+  labelResult(result: TimeResult): TimeLabeledResult
+  labelResult<TimeResult extends SortAtResult>(result: TimeResult) {
+    return { primary: result.sortAt }
+  }
+  labeledResultToCursor(labeled: TimeLabeledResult) {
+    return {
+      primary: new Date(labeled.primary).toISOString(),
+    }
+  }
+  cursorToLabeledResult(cursor: SingleKeyCursor) {
+    const primaryDate = new Date(cursor.primary)
+    if (isNaN(primaryDate.getTime())) {
+      throw new InvalidRequestError('Malformed cursor')
+    }
+    return {
+      primary: primaryDate.toISOString(),
+    }
+  }
+}
+
+export class IsoSortAtKey extends IsoTimeKey<{
+  sortAt: string
+}> {
+  labelResult(result: { sortAt: string }) {
+    return { primary: result.sortAt }
+  }
 }

package/src/data-plane/server/routes/notifs.ts

@@ -3,7 +3,7 @@ import { ServiceImpl } from '@connectrpc/connect'
 import { sql } from 'kysely'
 import { Service } from '../../../proto/bsky_connect'
 import { Database } from '../db'
-import { TimeCidKeyset, paginate } from '../db/pagination'
+import { IsoSortAtKey } from '../db/pagination'
 import { countAll, notSoftDeletedClause } from '../db/util'
 
 export default (db: Database): Partial<ServiceImpl<typeof Service>> => ({
@@ -41,15 +41,10 @@ export default (db: Database): Partial<ServiceImpl<typeof Service>> => ({
       ])
       .select(priorityFollowQb.as('priority'))
 
-    const keyset = new TimeCidKeyset(
-      ref('notif.sortAt'),
-      ref('notif.recordCid'),
-    )
-    builder = paginate(builder, {
+    const key = new IsoSortAtKey(ref('notif.sortAt'))
+    builder = key.paginate(builder, {
       cursor,
       limit,
-      keyset,
-      tryIndex: true,
     })
 
     const notifsRes = await builder.execute()
@@ -63,7 +58,7 @@ export default (db: Database): Partial<ServiceImpl<typeof Service>> => ({
     }))
     return {
       notifications,
-      cursor: keyset.packFromResult(notifsRes),
+      cursor: key.packFromResult(notifsRes),
     }
   },
 

package/src/etcd.ts

@@ -0,0 +1,90 @@
+import { once } from 'node:events'
+import { Etcd3, Watcher } from 'etcd3'
+
+/**
+ * A reactive map based on the keys and values stored within etcd under a given prefix.
+ */
+export class EtcdMap {
+  inner = new Map<string, VersionedValue>()
+  watcher: Watcher
+  connecting: Promise<void> | undefined
+  handlers: ((self: EtcdMap) => void)[] = []
+
+  constructor(
+    private etcd: Etcd3,
+    private prefix = '',
+  ) {
+    this.watcher = etcd.watch().prefix(prefix).watcher()
+    this.connecting = connectWatcher(this.watcher)
+  }
+
+  async connect() {
+    this.watcher.on('put', (kv) => {
+      const key = kv.key.toString()
+      const value = kv.value.toString()
+      const rev = revToInt(kv.mod_revision)
+      this.apply(key, { value, rev })
+    })
+    this.watcher.on('delete', (kv) => {
+      const key = kv.key.toString()
+      const value = null
+      const rev = revToInt(kv.mod_revision)
+      this.apply(key, { value, rev })
+    })
+    await this.connecting?.finally(() => {
+      this.connecting = undefined
+    })
+    const { kvs } = await this.etcd.getAll().prefix(this.prefix).exec()
+    for (const kv of kvs) {
+      const key = kv.key.toString()
+      const value = kv.value.toString()
+      const rev = revToInt(kv.mod_revision)
+      this.apply(key, { value, rev })
+    }
+  }
+
+  get(key: string) {
+    return this.inner.get(key)?.value ?? null
+  }
+
+  *values() {
+    for (const { value } of this.inner.values()) {
+      if (value !== null) {
+        yield value
+      }
+    }
+  }
+
+  onUpdate(handler: (self: EtcdMap) => void) {
+    this.handlers.push(handler)
+  }
+
+  private update() {
+    for (const handler of this.handlers) {
+      handler(this)
+    }
+  }
+
+  private apply(key, vv: VersionedValue) {
+    const curr = this.inner.get(key)
+    if (curr && curr.rev > vv.rev) return
+    this.inner.set(key, vv)
+    this.update()
+  }
+}
+
+function revToInt(rev: string) {
+  return parseInt(rev, 10)
+}
+
+async function connectWatcher(watcher: Watcher) {
+  await Promise.race([
+    once(watcher, 'connected'),
+    once(watcher, 'error').then((err) => Promise.reject(err)),
+  ])
+}
+
+type VersionedValue = {
+  rev: number
+  value: string | null
+}

package/src/index.ts

@@ -3,6 +3,7 @@ import http from 'node:http'
 import { AddressInfo } from 'node:net'
 import compression from 'compression'
 import cors from 'cors'
+import { Etcd3 } from 'etcd3'
 import express from 'express'
 import { HttpTerminator, createHttpTerminator } from 'http-terminator'
 import { AtpAgent } from '@atproto/api'
@@ -16,7 +17,11 @@ import { authWithApiKey as bsyncAuth, createBsyncClient } from './bsync'
 import { ServerConfig } from './config'
 import { AppContext } from './context'
 import { authWithApiKey as courierAuth, createCourierClient } from './courier'
-import { createDataPlaneClient } from './data-plane/client'
+import {
+  BasicHostList,
+  EtcdHostList,
+  createDataPlaneClient,
+} from './data-plane/client'
 import * as error from './error'
 import { FeatureGates } from './feature-gates'
 import { Hydrator } from './hydration/hydrator'
@@ -98,7 +103,20 @@ export class BskyAppView {
       )
     }
 
-    const dataplane = createDataPlaneClient(config.dataplaneUrls, {
+    const etcd = config.etcdHosts.length
+      ? new Etcd3({ hosts: config.etcdHosts })
+      : undefined
+
+    const dataplaneHostList =
+      etcd && config.dataplaneUrlsEtcdKeyPrefix
+        ? new EtcdHostList(
+            etcd,
+            config.dataplaneUrlsEtcdKeyPrefix,
+            config.dataplaneUrls,
+          )
+        : new BasicHostList(config.dataplaneUrls)
+
+    const dataplane = createDataPlaneClient(dataplaneHostList, {
       httpVersion: config.dataplaneHttpVersion,
       rejectUnauthorized: !config.dataplaneIgnoreBadTls,
     })
@@ -147,7 +165,9 @@ export class BskyAppView {
 
     const ctx = new AppContext({
       cfg: config,
+      etcd,
       dataplane,
+      dataplaneHostList,
       searchAgent,
       suggestionsAgent,
       topicsAgent,
@@ -184,6 +204,9 @@ export class BskyAppView {
   }
 
   async start(): Promise<http.Server> {
+    if (this.ctx.dataplaneHostList instanceof EtcdHostList) {
+      await this.ctx.dataplaneHostList.connect()
+    }
     await this.ctx.featureGates.start()
     const server = this.app.listen(this.ctx.cfg.port)
     this.server = server
@@ -198,6 +221,7 @@ export class BskyAppView {
   async destroy(): Promise<void> {
     await this.terminator?.terminate()
     this.ctx.featureGates.destroy()
+    await this.ctx.etcd?.close()
   }
 }
 

package/src/lexicon/index.ts

@@ -24,8 +24,11 @@ import * as ComAtprotoAdminUpdateAccountHandle from './types/com/atproto/admin/u
 import * as ComAtprotoAdminUpdateAccountPassword from './types/com/atproto/admin/updateAccountPassword.js'
 import * as ComAtprotoAdminUpdateSubjectStatus from './types/com/atproto/admin/updateSubjectStatus.js'
 import * as ComAtprotoIdentityGetRecommendedDidCredentials from './types/com/atproto/identity/getRecommendedDidCredentials.js'
+import * as ComAtprotoIdentityRefreshIdentity from './types/com/atproto/identity/refreshIdentity.js'
 import * as ComAtprotoIdentityRequestPlcOperationSignature from './types/com/atproto/identity/requestPlcOperationSignature.js'
+import * as ComAtprotoIdentityResolveDid from './types/com/atproto/identity/resolveDid.js'
 import * as ComAtprotoIdentityResolveHandle from './types/com/atproto/identity/resolveHandle.js'
+import * as ComAtprotoIdentityResolveIdentity from './types/com/atproto/identity/resolveIdentity.js'
 import * as ComAtprotoIdentitySignPlcOperation from './types/com/atproto/identity/signPlcOperation.js'
 import * as ComAtprotoIdentitySubmitPlcOperation from './types/com/atproto/identity/submitPlcOperation.js'
 import * as ComAtprotoIdentityUpdateHandle from './types/com/atproto/identity/updateHandle.js'
@@ -436,6 +439,17 @@ export class ComAtprotoIdentityNS {
     return this._server.xrpc.method(nsid, cfg)
   }
 
+  refreshIdentity<AV extends AuthVerifier>(
+    cfg: ConfigOf<
+      AV,
+      ComAtprotoIdentityRefreshIdentity.Handler<ExtractAuth<AV>>,
+      ComAtprotoIdentityRefreshIdentity.HandlerReqCtx<ExtractAuth<AV>>
+    >,
+  ) {
+    const nsid = 'com.atproto.identity.refreshIdentity' // @ts-ignore
+    return this._server.xrpc.method(nsid, cfg)
+  }
+
   requestPlcOperationSignature<AV extends AuthVerifier>(
     cfg: ConfigOf<
       AV,
@@ -449,6 +463,17 @@ export class ComAtprotoIdentityNS {
     return this._server.xrpc.method(nsid, cfg)
   }
 
+  resolveDid<AV extends AuthVerifier>(
+    cfg: ConfigOf<
+      AV,
+      ComAtprotoIdentityResolveDid.Handler<ExtractAuth<AV>>,
+      ComAtprotoIdentityResolveDid.HandlerReqCtx<ExtractAuth<AV>>
+    >,
+  ) {
+    const nsid = 'com.atproto.identity.resolveDid' // @ts-ignore
+    return this._server.xrpc.method(nsid, cfg)
+  }
+
   resolveHandle<AV extends AuthVerifier>(
     cfg: ConfigOf<
       AV,
@@ -460,6 +485,17 @@ export class ComAtprotoIdentityNS {
     return this._server.xrpc.method(nsid, cfg)
   }
 
+  resolveIdentity<AV extends AuthVerifier>(
+    cfg: ConfigOf<
+      AV,
+      ComAtprotoIdentityResolveIdentity.Handler<ExtractAuth<AV>>,
+      ComAtprotoIdentityResolveIdentity.HandlerReqCtx<ExtractAuth<AV>>
+    >,
+  ) {
+    const nsid = 'com.atproto.identity.resolveIdentity' // @ts-ignore
+    return this._server.xrpc.method(nsid, cfg)
+  }
+
   signPlcOperation<AV extends AuthVerifier>(
     cfg: ConfigOf<
       AV,