AT Protocol IPLD-CAR Repository toolkit (CLI)
atscan.net
1package engine
2
3import (
4 "bytes"
5 "context"
6 "errors"
7 "io"
8 "log"
9 "os"
10 "strings"
11 "time"
12
13 "github.com/atscan/atr/cli"
14 "github.com/atscan/atr/repo"
15 "github.com/klauspost/compress/zstd"
16 "github.com/schollz/progressbar/v3"
17)
18
19func Load(ctx *cli.Context, fn string) (repo.RepoSnapshot, error) {
20 var ss repo.RepoSnapshot
21 var err error
22 if strings.HasSuffix(fn, ".car.zst") {
23 ss, err = LoadCompressed(ctx, fn)
24 } else if strings.HasSuffix(fn, ".car") {
25 ss, err = LoadRaw(ctx, fn)
26 }
27 if err != nil {
28 log.Fatal("Cannot load: ", fn)
29 return ss, errors.New("Cannot load")
30 }
31 ss.File = fn
32 return ss, nil
33}
34
35func LoadRaw(ctx *cli.Context, fn string) (repo.RepoSnapshot, error) {
36 stat, err := os.Stat(fn)
37 if err != nil {
38 return repo.RepoSnapshot{}, err
39 }
40 dat, err := os.Open(fn)
41 if err != nil {
42 return repo.RepoSnapshot{}, err
43 }
44 defer dat.Close()
45
46 if err != nil {
47 return repo.RepoSnapshot{}, err
48 }
49 return LoadFromStream(ctx, dat, stat.Size())
50}
51
52func LoadCompressed(ctx *cli.Context, fn string) (repo.RepoSnapshot, error) {
53 dat, err := os.Open(fn)
54 if err != nil {
55 return repo.RepoSnapshot{}, err
56 }
57 defer dat.Close()
58
59 if err != nil {
60 return repo.RepoSnapshot{}, err
61 }
62 decoder, _ := zstd.NewReader(dat)
63 return LoadFromStream(ctx, decoder, -1)
64}
65
66func LoadFromStream(ctx *cli.Context, input io.Reader, sz int64) (repo.RepoSnapshot, error) {
67 ss, err := LoadRepoFromStream(ctx, input, sz)
68 if err != nil {
69 return ss, err
70 }
71
72 ss.Root = ss.Repo.Head()
73
74 return ss, nil
75}
76
77func LoadRepoFromStream(ctx *cli.Context, input io.Reader, sz int64) (repo.RepoSnapshot, error) {
78 rctx := context.TODO()
79 ss := repo.RepoSnapshot{}
80
81 var bar *progressbar.ProgressBar
82 if sz > 150_000_000 && !ctx.Raw {
83 bar = progressbar.NewOptions(int(sz),
84 progressbar.OptionSetDescription("loading"),
85 progressbar.OptionSetWriter(os.Stderr),
86 progressbar.OptionShowBytes(true),
87 //progressbar.OptionSetWidth(10),
88 progressbar.OptionThrottle(35*time.Millisecond),
89 progressbar.OptionShowCount(),
90 //progressbar.OptionOnCompletion(func() {
91 //}),
92 //progressbar.OptionSpinnerType(14),
93 //progressbar.OptionFullWidth(),
94 //progressbar.OptionUseANSICodes(true),
95 progressbar.OptionSetRenderBlankState(false),
96 /*progressbar.OptionSetTheme(progressbar.Theme{
97 Saucer: "=",
98 SaucerHead: ">",
99 SaucerPadding: " ",
100 BarStart: "[",
101 BarEnd: "]",
102 }),*/
103 )
104 ctx.ProgressBar = bar
105 }
106
107 buf := new(bytes.Buffer)
108 var err error
109 var size int64
110 if bar != nil {
111 size, err = io.Copy(io.MultiWriter(buf, bar), input)
112 } else {
113 size, err = io.Copy(buf, input)
114 }
115 if err != nil {
116 log.Fatal(err)
117 }
118 ss.Size = int(size)
119
120 if bar != nil {
121 defer bar.Finish()
122 defer bar.Clear()
123 }
124
125 r, err := repo.ReadRepoFromCar(rctx, buf)
126 if err != nil {
127 return ss, err
128 }
129 ss.Repo = *r
130 return ss, nil
131}