From b9822e2287d0fe1c671f76cb03cd1b4ad7b0ad48 Mon Sep 17 00:00:00 2001 From: Aly Raffauf Date: Fri, 3 Jul 2026 01:09:21 -0400 Subject: [PATCH] add backfill --- cmd/asterism/main.go | 27 +++++++- internal/backfill/backfill.go | 112 ++++++++++++++++++++++++++++++++++ internal/firehose/commit.go | 23 +------ internal/index/record.go | 30 +++++++++ 4 files changed, 169 insertions(+), 23 deletions(-) create mode 100644 internal/backfill/backfill.go create mode 100644 internal/index/record.go diff --git a/cmd/asterism/main.go b/cmd/asterism/main.go index 764deeb..dc10fdf 100644 --- a/cmd/asterism/main.go +++ b/cmd/asterism/main.go @@ -3,13 +3,17 @@ package main import ( "context" "flag" + "fmt" "log/slog" "net/http" "strings" + "github.com/bluesky-social/indigo/atproto/identity" + "github.com/bluesky-social/indigo/xrpc" "github.com/gorilla/websocket" "github.com/alyraffauf/asterism/internal/api" + "github.com/alyraffauf/asterism/internal/backfill" "github.com/alyraffauf/asterism/internal/firehose" "github.com/alyraffauf/asterism/internal/store" ) @@ -57,10 +61,29 @@ func main() { panic(err) } }() - defer conn.Close() + + wantedCollections := parseCollections(*collectionsFlag) + + var collections []string + for collection := range wantedCollections { + collections = append(collections, collection) + } + + if len(collections) > 0 { + bf := &backfill.Backfill{ + Client: &xrpc.Client{Host: "https://relay1.us-east.bsky.network"}, + Directory: identity.DefaultDirectory(), + Store: linkStore, + } + go func() { + if err := bf.Run(ctx, collections); err != nil { + fmt.Println("backfill error:", err) + } + }() + } consumer := &firehose.Consumer{ - WantedCollections: parseCollections(*collectionsFlag), + WantedCollections: wantedCollections, Store: linkStore, } diff --git a/internal/backfill/backfill.go b/internal/backfill/backfill.go new file mode 100644 index 0000000..b7b14d9 --- /dev/null +++ b/internal/backfill/backfill.go @@ -0,0 +1,112 @@ +package backfill + +import ( + "bytes" + "context" + "fmt" + "net/http" + "strings" + "time" + + "github.com/alyraffauf/asterism/internal/index" + "github.com/alyraffauf/asterism/internal/store" + "github.com/bluesky-social/indigo/api/atproto" + "github.com/bluesky-social/indigo/atproto/identity" + "github.com/bluesky-social/indigo/atproto/syntax" + "github.com/bluesky-social/indigo/xrpc" + "github.com/ipfs/go-cid" + + indigorepo "github.com/bluesky-social/indigo/repo" +) + +type Backfill struct { + Client *xrpc.Client + Directory identity.Directory + Store *store.Store +} + +func (b *Backfill) Run(ctx context.Context, collections []string) error { + dids := make(map[string]map[string]struct{}) + + for _, collection := range collections { + if err := b.listRepos(ctx, collection, dids); err != nil { + return fmt.Errorf("list repos for %s: %w", collection, err) + } + } + + for did, wantedCollections := range dids { + if err := b.repo(ctx, did, wantedCollections); err != nil { + fmt.Println("could not backfill repo:", did, err) + continue + } + } + + return nil +} + +func (b *Backfill) listRepos(ctx context.Context, collection string, dids map[string]map[string]struct{}) error { + cursor := "" + for { + page, err := atproto.SyncListReposByCollection(ctx, b.Client, collection, cursor, 1000) + if err != nil { + return fmt.Errorf("list repos: %w", err) + } + + for _, repo := range page.Repos { + if dids[repo.Did] == nil { + dids[repo.Did] = make(map[string]struct{}) + } + dids[repo.Did][collection] = struct{}{} + } + + if page.Cursor == nil || *page.Cursor == "" { + return nil + } + cursor = *page.Cursor + } +} + +func (b *Backfill) repo(ctx context.Context, did string, wantedCollections map[string]struct{}) error { + resolveCtx, cancel := context.WithTimeout(ctx, 10*time.Second) + identity, err := b.Directory.LookupDID(resolveCtx, syntax.DID(did)) + cancel() + if err != nil { + return fmt.Errorf("resolve did: %w", err) + } + + pdsClient := &xrpc.Client{ + Host: identity.PDSEndpoint(), + Client: &http.Client{Timeout: 5 * time.Minute}, + } + + fetchCtx, cancel := context.WithTimeout(ctx, 5*time.Minute) + defer cancel() + + carBytes, err := atproto.SyncGetRepo(fetchCtx, pdsClient, did, "") + if err != nil { + return fmt.Errorf("get repo: %w", err) + } + + repo, err := indigorepo.ReadRepoFromCar(ctx, bytes.NewReader(carBytes)) + if err != nil { + return fmt.Errorf("read repo car: %w", err) + } + + return repo.ForEach(ctx, "", func(k string, v cid.Cid) error { + collection, recordKey, ok := strings.Cut(k, "/") + if !ok { + return fmt.Errorf("bad path: %s", k) + } + + if _, wanted := wantedCollections[collection]; !wanted { + return nil + } + + recordCid, recordBytes, err := repo.GetRecordBytes(ctx, k) + if err != nil { + return fmt.Errorf("read record: %w", err) + } + + return index.Record(ctx, b.Store, did, collection, recordKey, recordCid.String(), "", *recordBytes) + }) +} diff --git a/internal/firehose/commit.go b/internal/firehose/commit.go index b2a1953..9135d04 100644 --- a/internal/firehose/commit.go +++ b/internal/firehose/commit.go @@ -6,9 +6,8 @@ import ( "fmt" "strings" - "github.com/alyraffauf/asterism/internal/backlink" + "github.com/alyraffauf/asterism/internal/index" "github.com/bluesky-social/indigo/api/atproto" - "github.com/bluesky-social/indigo/atproto/atdata" indigorepo "github.com/bluesky-social/indigo/repo" "github.com/ipfs/go-cid" ) @@ -38,7 +37,6 @@ func (c *Consumer) HandleCommit(ctx context.Context, event *atproto.SyncSubscrib func (c *Consumer) handleOperation(ctx context.Context, event *atproto.SyncSubscribeRepos_Commit, repo *indigorepo.Repo, operation *atproto.SyncSubscribeRepos_RepoOp) error { collection, recordKey, ok := strings.Cut(operation.Path, "/") - if !ok { return fmt.Errorf("bad path: %s", operation.Path) } @@ -47,8 +45,6 @@ func (c *Consumer) handleOperation(ctx context.Context, event *atproto.SyncSubsc return nil } - fmt.Println("op:", operation.Action, "collection:", collection, "rkey:", recordKey) - if operation.Action == "delete" { return c.Store.DeleteLinks(ctx, event.Repo, collection, recordKey) } @@ -67,20 +63,5 @@ func (c *Consumer) handleOperation(ctx context.Context, event *atproto.SyncSubsc return fmt.Errorf("cid mismatch: %s operation=%s record=%s", operation.Path, operationCid, recordCid) } - record, err := atdata.UnmarshalCBOR(*recordBytes) - if err != nil { - return fmt.Errorf("decode record: %w", err) - } - - base := backlink.Link{ - ActorDid: event.Repo, - Collection: collection, - RecordKey: recordKey, - RecordCid: recordCid.String(), - Rev: event.Rev, - } - - links := backlink.Extract(record, base) - - return c.Store.SaveLinks(ctx, event.Repo, collection, recordKey, links) + return index.Record(ctx, c.Store, event.Repo, collection, recordKey, recordCid.String(), event.Rev, *recordBytes) } diff --git a/internal/index/record.go b/internal/index/record.go new file mode 100644 index 0000000..adb6818 --- /dev/null +++ b/internal/index/record.go @@ -0,0 +1,30 @@ +package index + +import ( + "context" + "fmt" + + "github.com/bluesky-social/indigo/atproto/atdata" + + "github.com/alyraffauf/asterism/internal/backlink" + "github.com/alyraffauf/asterism/internal/store" +) + +func Record(ctx context.Context, s *store.Store, actorDid, collection, recordKey, recordCid, rev string, recordBytes []byte) error { + record, err := atdata.UnmarshalCBOR(recordBytes) + if err != nil { + return fmt.Errorf("decode record: %w", err) + } + + base := backlink.Link{ + ActorDid: actorDid, + Collection: collection, + RecordKey: recordKey, + RecordCid: recordCid, + Rev: rev, + } + + links := backlink.Extract(record, base) + + return s.SaveLinks(ctx, actorDid, collection, recordKey, links) +} -- 2.51.2