This repository has no description
5.0 kB
173 lines
1// prefill-zoekt bulk-enqueues zoekt index tasks
2//
3// It reads a REPOS file with one repo DID per line:
4//
5// did:plc:repository
6//
7// For each DID it resolves the knot from the DID document, then resolves the
8// remote HEAD branch + commit via `git ls-remote --symref`, and POSTs an
9// enqueue request to the indexserver's /admin/enqueueIndex endpoint.
10package main
11
12import (
13 "bytes"
14 "context"
15 "encoding/json"
16 "flag"
17 "fmt"
18 "log"
19 "net/http"
20 "net/url"
21 "os"
22 "os/exec"
23 "strings"
24 "sync"
25 "sync/atomic"
26
27 "github.com/bluesky-social/indigo/atproto/identity"
28 "github.com/bluesky-social/indigo/atproto/syntax"
29 "github.com/samber/lo"
30 "github.com/sourcegraph/zoekt"
31 "tangled.org/core/repoident"
32)
33
34func main() {
35 reposPath := flag.String("repos", "REPOS", "path to repos list file (one DID per line)")
36 serverUrl := flag.String("server", "http://localhost:6060", "indexserver base url")
37 plc := flag.String("plc", "https://plc.directory", "atproto PLC directory url")
38 concurrency := flag.Int("concurrency", 5, "number of repos to process in parallel")
39 allowHttp := flag.Bool("allow-http", false, "accept repo DIDs whose knot service endpoint is plaintext http, and skip TLS verification when reading HEAD")
40 flag.Parse()
41
42 server, err := url.Parse(*serverUrl)
43 if err != nil {
44 log.Fatalf("parsing -server %q: %v", *serverUrl, err)
45 }
46 if (server.Scheme != "http" && server.Scheme != "https") || server.Host == "" {
47 log.Fatalf("-server %q must be an http or https URL with a host", *serverUrl)
48 }
49
50 data, err := os.ReadFile(*reposPath)
51 if err != nil {
52 log.Fatalf("reading %s: %v", *reposPath, err)
53 }
54
55 ctx := context.Background()
56 dir := identity.BaseDirectory{PLCURL: *plc}
57
58 var ok, fail atomic.Int64
59 var wg sync.WaitGroup
60 sem := make(chan struct{}, *concurrency)
61
62 lo.ForEach(strings.Split(string(data), "\n"), func(line string, i int) {
63 raw := strings.TrimSpace(line)
64 if raw == "" {
65 return
66 }
67
68 wg.Add(1)
69 sem <- struct{}{}
70 go func() {
71 defer wg.Done()
72 defer func() { <-sem }()
73
74 head, knot, err := prefillRepo(ctx, &dir, server, raw, *allowHttp)
75 if err != nil {
76 log.Printf("line %d: %s: %v", i+1, raw, err)
77 fail.Add(1)
78 return
79 }
80 log.Printf("line %d: %s: enqueued %s@%s (knot=%s)", i+1, raw, head.Name, head.Version, knot)
81 ok.Add(1)
82 }()
83 })
84
85 wg.Wait()
86 fmt.Printf("done: %d enqueued, %d failed\n", ok.Load(), fail.Load())
87}
88
89func prefillRepo(ctx context.Context, dir identity.Directory, server *url.URL, raw string, allowHTTP bool) (zoekt.RepositoryBranch, repoident.KnotURL, error) {
90 var knot repoident.KnotURL
91
92 repoDid, err := repoident.NewRepoDid(raw)
93 if err != nil {
94 return zoekt.RepositoryBranch{}, knot, err
95 }
96
97 ident, err := dir.LookupDID(ctx, syntax.DID(repoDid))
98 if err != nil {
99 return zoekt.RepositoryBranch{}, knot, fmt.Errorf("resolving repo DID: %w", err)
100 }
101
102 knot, err = repoident.KnotURLFromIdentity(ident, repoident.SchemeFor(allowHTTP))
103 if err != nil {
104 return zoekt.RepositoryBranch{}, knot, fmt.Errorf("resolving knot: %w", err)
105 }
106
107 head, err := resolveHead(knot, repoDid, allowHTTP)
108 if err != nil {
109 return head, knot, fmt.Errorf("resolving HEAD: %w", err)
110 }
111
112 if err := enqueue(server, repoDid, head); err != nil {
113 return head, knot, fmt.Errorf("enqueue: %w", err)
114 }
115 return head, knot, nil
116}
117
118func resolveHead(knot repoident.KnotURL, repoDid repoident.RepoDid, allowHTTP bool) (zoekt.RepositoryBranch, error) {
119 remote := knot.JoinPath(repoDid.String())
120 args := append(
121 lo.Ternary(allowHTTP, []string{"-c", "http.sslVerify=false"}, nil),
122 "ls-remote", "--symref", remote, "HEAD",
123 )
124 out, err := exec.Command("git", args...).Output()
125 if err != nil {
126 return zoekt.RepositoryBranch{}, fmt.Errorf("git ls-remote --symref %s HEAD: %w", remote, err)
127 }
128 head := lo.Reduce(
129 strings.Split(string(out), "\n"),
130 func(head zoekt.RepositoryBranch, line string, _ int) zoekt.RepositoryBranch {
131 fields := strings.Fields(line)
132 if len(fields) < 2 {
133 return head
134 }
135 switch {
136 case fields[0] == "ref:":
137 head.Name = strings.TrimPrefix(fields[1], "refs/heads/")
138 case fields[1] == "HEAD":
139 head.Version = fields[0]
140 }
141 return head
142 },
143 zoekt.RepositoryBranch{},
144 )
145 if head.Name == "" || head.Version == "" {
146 return zoekt.RepositoryBranch{}, fmt.Errorf("couldn't resolve HEAD (branch=%q sha=%q)", head.Name, head.Version)
147 }
148 return head, nil
149}
150
151func enqueue(server *url.URL, repoDid repoident.RepoDid, head zoekt.RepositoryBranch) error {
152 body, err := json.Marshal(map[string]any{
153 "repo": repoDid.String(),
154 "branches": []zoekt.RepositoryBranch{head},
155 })
156 if err != nil {
157 return err
158 }
159
160 resp, err := http.Post(server.JoinPath("admin", "enqueueIndex").String(),
161 "application/json", bytes.NewReader(body))
162 if err != nil {
163 return err
164 }
165 defer resp.Body.Close()
166
167 log.Println("status", resp.StatusCode)
168
169 if resp.StatusCode < 200 || resp.StatusCode >= 300 {
170 return fmt.Errorf("status %d", resp.StatusCode)
171 }
172 return nil
173}