+352

atproto/client.go

reviewed

··· 1 1 + package atproto 2 2 + 3 3 + import ( 4 4 + "context" 5 5 + "encoding/json" 6 6 + "fmt" 7 7 + "net/http" 8 8 + "net/url" 9 9 + "sync" 10 10 + "time" 11 11 + 12 12 + "github.com/bluesky-social/indigo/atproto/atclient" 13 13 + "github.com/bluesky-social/indigo/atproto/syntax" 14 14 + ) 15 15 + 16 16 + const ( 17 17 + DefaultRateLimitPercent = 0.85 18 18 + DefaultResolverURL = "https://slingshot.microcosm.blue/xrpc/com.bad-example.identity.resolveMiniDoc" 19 19 + ) 20 20 + 21 21 + type ClientOptions struct { 22 22 + UserAgent string 23 23 + ResolverURL string 24 24 + HTTPClient *http.Client 25 25 + RateLimitPercent float32 26 26 + } 27 27 + 28 28 + type resolvedIdentity struct { 29 29 + DID string 30 30 + Handle string 31 31 + PDS string 32 32 + SigningKey string 33 33 + } 34 34 + 35 35 + type Client struct { 36 36 + client *atclient.APIClient 37 37 + resolvedIdentity resolvedIdentity 38 38 + mu sync.Mutex 39 39 + } 40 40 + 41 41 + func WithUserAgent(ua string) func(*ClientOptions) { 42 42 + return func(o *ClientOptions) { 43 43 + o.UserAgent = ua 44 44 + } 45 45 + } 46 46 + 47 47 + func WithResolverURL(url string) func(*ClientOptions) { 48 48 + return func(o *ClientOptions) { 49 49 + o.ResolverURL = url 50 50 + } 51 51 + } 52 52 + 53 53 + func WithHTTPClient(hc *http.Client) func(*ClientOptions) { 54 54 + return func(o *ClientOptions) { 55 55 + o.HTTPClient = hc 56 56 + } 57 57 + } 58 58 + 59 59 + func WithRateLimitPercent(p float32) func(*ClientOptions) { 60 60 + return func(o *ClientOptions) { 61 61 + o.RateLimitPercent = p 62 62 + } 63 63 + } 64 64 + 65 65 + func NewClientOptions() *ClientOptions { 66 66 + return &ClientOptions{ 67 67 + ResolverURL: DefaultResolverURL, 68 68 + RateLimitPercent: DefaultRateLimitPercent, 69 69 + HTTPClient: &http.Client{ 70 70 + Timeout: 10 * time.Second, 71 71 + }, 72 72 + } 73 73 + } 74 74 + 75 75 + func ResolveMiniDoc(ctx context.Context, identifier string, opts *ClientOptions) (did, pds, signingKey string, err error) { 76 76 + resolver := DefaultResolverURL 77 77 + if opts != nil && opts.ResolverURL != "" { 78 78 + resolver = opts.ResolverURL 79 79 + } 80 80 + 81 81 + parsedURL, err := url.Parse(resolver) 82 82 + if err != nil { 83 83 + return "", "", "", fmt.Errorf("invalid resolver URL: %w", err) 84 84 + } 85 85 + 86 86 + query := parsedURL.Query() 87 87 + query.Set("identifier", identifier) 88 88 + parsedURL.RawQuery = query.Encode() 89 89 + 90 90 + req, err := http.NewRequestWithContext(ctx, "GET", parsedURL.String(), nil) 91 91 + if err != nil { 92 92 + return "", "", "", fmt.Errorf("failed to create request: %w", err) 93 93 + } 94 94 + 95 95 + if opts != nil && opts.UserAgent != "" { 96 96 + req.Header.Set("User-Agent", opts.UserAgent) 97 97 + } 98 98 + 99 99 + httpClient := http.DefaultClient 100 100 + if opts != nil && opts.HTTPClient != nil { 101 101 + httpClient = opts.HTTPClient 102 102 + } 103 103 + 104 104 + resp, err := httpClient.Do(req) 105 105 + if err != nil { 106 106 + return "", "", "", fmt.Errorf("failed to resolve mini doc: %w", err) 107 107 + } 108 108 + defer resp.Body.Close() 109 109 + 110 110 + if resp.StatusCode != http.StatusOK { 111 111 + return "", "", "", fmt.Errorf("mini doc resolution failed with status: %d", resp.StatusCode) 112 112 + } 113 113 + 114 114 + var result struct { 115 115 + DID string `json:"did"` 116 116 + PDS string `json:"pds"` 117 117 + SigningKey string `json:"signing_key"` 118 118 + } 119 119 + if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { 120 120 + return "", "", "", fmt.Errorf("failed to decode mini doc response: %w", err) 121 121 + } 122 122 + 123 123 + if result.DID == "" { 124 124 + return "", "", "", fmt.Errorf("resolved mini doc missing DID") 125 125 + } 126 126 + 127 127 + return result.DID, result.PDS, result.SigningKey, nil 128 128 + } 129 129 + 130 130 + type FixedPasswordAuth struct { 131 131 + *atclient.PasswordAuth 132 132 + lk sync.RWMutex 133 133 + RefreshCallback func(ctx context.Context, session atclient.PasswordSessionData) 134 134 + } 135 135 + 136 136 + func (a *FixedPasswordAuth) DoWithAuth(c *http.Client, req *http.Request, endpoint syntax.NSID) (*http.Response, error) { 137 137 + accessToken, refreshToken := a.GetTokens() 138 138 + req.Header.Set("Authorization", "Bearer "+accessToken) 139 139 + resp, err := c.Do(req) 140 140 + if err != nil { 141 141 + return nil, err 142 142 + } 143 143 + 144 144 + if resp.StatusCode != http.StatusBadRequest { 145 145 + return resp, nil 146 146 + } 147 147 + 148 148 + if !hasJSONContent(resp.Header) { 149 149 + return resp, nil 150 150 + } 151 151 + 152 152 + defer resp.Body.Close() 153 153 + var eb atclient.ErrorBody 154 154 + if err := json.NewDecoder(resp.Body).Decode(&eb); err != nil { 155 155 + return nil, &atclient.APIError{StatusCode: resp.StatusCode} 156 156 + } 157 157 + if eb.Name != "ExpiredToken" { 158 158 + return nil, eb.APIError(resp.StatusCode) 159 159 + } 160 160 + 161 161 + if err := a.Refresh(req.Context(), c, refreshToken); err != nil { 162 162 + return nil, err 163 163 + } 164 164 + 165 165 + retry := req.Clone(req.Context()) 166 166 + if req.GetBody != nil { 167 167 + retryBody, err := req.GetBody() 168 168 + if err != nil { 169 169 + return nil, fmt.Errorf("API request retry GetBody failed: %w", err) 170 170 + } 171 171 + retry.Body = retryBody 172 172 + } 173 173 + 174 174 + accessToken, _ = a.GetTokens() 175 175 + retry.Header.Set("Authorization", "Bearer "+accessToken) 176 176 + return c.Do(retry) 177 177 + } 178 178 + 179 179 + func hasJSONContent(header http.Header) bool { 180 180 + return len(header.Get("Content-Type")) > 0 && header.Get("Content-Type")[0:19] == "application/json" 181 181 + } 182 182 + 183 183 + func (a *FixedPasswordAuth) Refresh(ctx context.Context, c *http.Client, priorRefreshToken string) error { 184 184 + a.lk.Lock() 185 185 + defer a.lk.Unlock() 186 186 + 187 187 + if priorRefreshToken != "" && priorRefreshToken != a.Session.RefreshToken { 188 188 + return nil 189 189 + } 190 190 + 191 191 + u := a.Session.Host + "/xrpc/com.atproto.server.refreshSession" 192 192 + req, err := http.NewRequestWithContext(ctx, http.MethodPost, u, nil) 193 193 + if err != nil { 194 194 + return err 195 195 + } 196 196 + req.Header.Set("User-Agent", "indigo-sdk") 197 197 + req.Header.Set("Authorization", "Bearer "+a.Session.RefreshToken) 198 198 + 199 199 + resp, err := c.Do(req) 200 200 + if err != nil { 201 201 + return err 202 202 + } 203 203 + defer resp.Body.Close() 204 204 + 205 205 + if resp.StatusCode < 200 || resp.StatusCode >= 300 { 206 206 + var eb atclient.ErrorBody 207 207 + if err := json.NewDecoder(resp.Body).Decode(&eb); err != nil { 208 208 + return &atclient.APIError{StatusCode: resp.StatusCode} 209 209 + } 210 210 + return eb.APIError(resp.StatusCode) 211 211 + } 212 212 + 213 213 + var out struct { 214 214 + AccessJwt string `json:"accessJwt"` 215 215 + RefreshJwt string `json:"refreshJwt"` 216 216 + } 217 217 + if err := json.NewDecoder(resp.Body).Decode(&out); err != nil { 218 218 + return err 219 219 + } 220 220 + 221 221 + a.Session.AccessToken = out.AccessJwt 222 222 + a.Session.RefreshToken = out.RefreshJwt 223 223 + 224 224 + if a.RefreshCallback != nil { 225 225 + snapshot := a.Session.Clone() 226 226 + a.RefreshCallback(ctx, snapshot) 227 227 + } 228 228 + 229 229 + return nil 230 230 + } 231 231 + 232 232 + func ResolveIdentity(ctx context.Context, handle string, opts *ClientOptions) (resolvedIdentity, error) { 233 233 + if handle == "" { 234 234 + return resolvedIdentity{}, fmt.Errorf("handle cannot be empty") 235 235 + } 236 236 + 237 237 + did, pds, signingKey, err := ResolveMiniDoc(ctx, handle, opts) 238 238 + if err != nil { 239 239 + return resolvedIdentity{}, fmt.Errorf("failed to resolve identity: %w", err) 240 240 + } 241 241 + 242 242 + if did == "" { 243 243 + return resolvedIdentity{}, fmt.Errorf("resolved identity missing DID") 244 244 + } 245 245 + 246 246 + if signingKey == "" { 247 247 + return resolvedIdentity{}, fmt.Errorf("resolved identity missing signing key") 248 248 + } 249 249 + 250 250 + return resolvedIdentity{ 251 251 + DID: did, 252 252 + Handle: handle, 253 253 + PDS: pds, 254 254 + SigningKey: signingKey, 255 255 + }, nil 256 256 + } 257 257 + 258 258 + func NewClient(ctx context.Context, handle, password string, opts ...func(*ClientOptions)) (*Client, error) { 259 259 + options := NewClientOptions() 260 260 + for _, opt := range opts { 261 261 + opt(options) 262 262 + } 263 263 + 264 264 + identity, err := ResolveIdentity(ctx, handle, options) 265 265 + if err != nil { 266 266 + return nil, fmt.Errorf("failed to resolve identity: %w", err) 267 267 + } 268 268 + 269 269 + pdsURL := identity.PDS 270 270 + if pdsURL == "" { 271 271 + pdsURL = "https://bsky.social" 272 272 + } 273 273 + 274 274 + client, err := atclient.LoginWithPasswordHost(ctx, pdsURL, handle, password, "", nil) 275 275 + if err != nil { 276 276 + return nil, fmt.Errorf("login failed: %w", err) 277 277 + } 278 278 + 279 279 + if pa, ok := client.Auth.(*atclient.PasswordAuth); ok { 280 280 + client.Auth = &FixedPasswordAuth{PasswordAuth: pa} 281 281 + } 282 282 + 283 283 + return &Client{ 284 284 + client: client, 285 285 + resolvedIdentity: identity, 286 286 + }, nil 287 287 + } 288 288 + 289 289 + func (c *Client) DID() string { 290 290 + c.mu.Lock() 291 291 + defer c.mu.Unlock() 292 292 + return c.resolvedIdentity.DID 293 293 + } 294 294 + 295 295 + func (c *Client) PDS() string { 296 296 + c.mu.Lock() 297 297 + defer c.mu.Unlock() 298 298 + return c.resolvedIdentity.PDS 299 299 + } 300 300 + 301 301 + func (c *Client) Handle() string { 302 302 + c.mu.Lock() 303 303 + defer c.mu.Unlock() 304 304 + return c.resolvedIdentity.Handle 305 305 + } 306 306 + 307 307 + func (c *Client) SigningKey() string { 308 308 + c.mu.Lock() 309 309 + defer c.mu.Unlock() 310 310 + return c.resolvedIdentity.SigningKey 311 311 + } 312 312 + 313 313 + func (c *Client) APIClient() *atclient.APIClient { 314 314 + c.mu.Lock() 315 315 + defer c.mu.Unlock() 316 316 + return c.client 317 317 + } 318 318 + 319 319 + func (c *Client) Close() error { 320 320 + c.mu.Lock() 321 321 + defer c.mu.Unlock() 322 322 + if c.client != nil && c.client.Auth != nil { 323 323 + if logout, ok := c.client.Auth.(*atclient.PasswordAuth); ok { 324 324 + return logout.Logout(context.Background(), c.client.Client) 325 325 + } 326 326 + } 327 327 + return nil 328 328 + } 329 329 + 330 330 + func (c *Client) HasClient() bool { 331 331 + c.mu.Lock() 332 332 + defer c.mu.Unlock() 333 333 + return c.client != nil 334 334 + } 335 335 + 336 336 + type AuthClient interface { 337 337 + APIClient() *atclient.APIClient 338 338 + DID() string 339 339 + } 340 340 + 341 341 + func BuildClient[T any](client AuthClient, customClient RepoClient[T]) (RepoClient[T], error) { 342 342 + if customClient != nil { 343 343 + return customClient, nil 344 344 + } 345 345 + 346 346 + apiClient := client.APIClient() 347 347 + if apiClient == nil { 348 348 + return nil, fmt.Errorf("API client is nil") 349 349 + } 350 350 + 351 351 + return NewRateClient[T](apiClient, client.DID(), nil), nil 352 352 + }

+706

atproto/client_test.go

reviewed

··· 1 1 + package atproto 2 2 + 3 3 + import ( 4 4 + "context" 5 5 + "encoding/json" 6 6 + "net/http" 7 7 + "net/http/httptest" 8 8 + "testing" 9 9 + "time" 10 10 + 11 11 + "github.com/bluesky-social/indigo/atproto/atclient" 12 12 + ) 13 13 + 14 14 + func TestResolveMiniDoc(t *testing.T) { 15 15 + t.Parallel() 16 16 + 17 17 + tests := []struct { 18 18 + name string 19 19 + handle string 20 20 + opts *ClientOptions 21 21 + responseStatus int 22 22 + responseBody any 23 23 + wantErr bool 24 24 + wantDID string 25 25 + wantPDS string 26 26 + }{ 27 27 + { 28 28 + name: "success with default options", 29 29 + handle: "test.user", 30 30 + opts: NewClientOptions(), 31 31 + responseStatus: http.StatusOK, 32 32 + responseBody: map[string]any{ 33 33 + "did": "did:plc:z72iitness", 34 34 + "pds": "https://pds.example.com", 35 35 + "signing_key": "-----BEGIN PUBLIC KEY-----\ntest\n-----END PUBLIC KEY-----", 36 36 + }, 37 37 + wantErr: false, 38 38 + wantDID: "did:plc:z72iitness", 39 39 + wantPDS: "https://pds.example.com", 40 40 + }, 41 41 + { 42 42 + name: "empty DID", 43 43 + handle: "test.user", 44 44 + opts: NewClientOptions(), 45 45 + responseStatus: http.StatusOK, 46 46 + responseBody: map[string]any{ 47 47 + "did": "", 48 48 + "pds": "https://pds.example.com", 49 49 + "signing_key": "key", 50 50 + }, 51 51 + wantErr: true, 52 52 + }, 53 53 + { 54 54 + name: "non-200 status", 55 55 + handle: "test.user", 56 56 + opts: NewClientOptions(), 57 57 + responseStatus: http.StatusInternalServerError, 58 58 + responseBody: nil, 59 59 + wantErr: true, 60 60 + }, 61 61 + { 62 62 + name: "invalid JSON", 63 63 + handle: "test.user", 64 64 + opts: NewClientOptions(), 65 65 + responseStatus: http.StatusOK, 66 66 + responseBody: "not json", 67 67 + wantErr: true, 68 68 + }, 69 69 + } 70 70 + 71 71 + for _, tt := range tests { 72 72 + t.Run(tt.name, func(t *testing.T) { 73 73 + t.Parallel() 74 74 + 75 75 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 76 76 + if tt.opts != nil && tt.opts.ResolverURL != "" && tt.opts.ResolverURL != DefaultResolverURL { 77 77 + if r.URL.String() == DefaultResolverURL { 78 78 + t.Errorf("expected custom resolver URL to be used") 79 79 + } 80 80 + } 81 81 + 82 82 + w.Header().Set("Content-Type", "application/json") 83 83 + 84 84 + switch v := tt.responseBody.(type) { 85 85 + case string: 86 86 + w.Write([]byte(v)) 87 87 + case map[string]any: 88 88 + json.NewEncoder(w).Encode(v) 89 89 + default: 90 90 + w.WriteHeader(tt.responseStatus) 91 91 + return 92 92 + } 93 93 + 94 94 + if tt.responseStatus != http.StatusOK { 95 95 + w.WriteHeader(tt.responseStatus) 96 96 + } 97 97 + })) 98 98 + defer server.Close() 99 99 + 100 100 + ctx := context.Background() 101 101 + 102 102 + opts := tt.opts 103 103 + if opts == nil { 104 104 + opts = NewClientOptions() 105 105 + } 106 106 + if opts.ResolverURL == DefaultResolverURL { 107 107 + opts.ResolverURL = server.URL 108 108 + } 109 109 + 110 110 + did, pds, _, err := ResolveMiniDoc(ctx, tt.handle, opts) 111 111 + 112 112 + if tt.wantErr { 113 113 + if err == nil { 114 114 + t.Error("expected error, got nil") 115 115 + } 116 116 + return 117 117 + } 118 118 + 119 119 + if err != nil { 120 120 + t.Fatalf("ResolveMiniDoc failed: %v", err) 121 121 + } 122 122 + 123 123 + if tt.wantDID != "" && did != tt.wantDID { 124 124 + t.Errorf("DID = %s, want %s", did, tt.wantDID) 125 125 + } 126 126 + if tt.wantPDS != "" && pds != tt.wantPDS { 127 127 + t.Errorf("PDS = %s, want %s", pds, tt.wantPDS) 128 128 + } 129 129 + }) 130 130 + } 131 131 + } 132 132 + 133 133 + func TestResolveMiniDoc_WithCustomResolver(t *testing.T) { 134 134 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 135 135 + w.Header().Set("Content-Type", "application/json") 136 136 + json.NewEncoder(w).Encode(map[string]any{ 137 137 + "did": "did:plc:custom", 138 138 + "pds": "https://custom.pds.example.com", 139 139 + "signing_key": "key", 140 140 + }) 141 141 + })) 142 142 + defer server.Close() 143 143 + 144 144 + opts := NewClientOptions() 145 145 + opts.ResolverURL = server.URL 146 146 + 147 147 + ctx := context.Background() 148 148 + did, pds, _, err := ResolveMiniDoc(ctx, "test.user", opts) 149 149 + if err != nil { 150 150 + t.Fatalf("ResolveMiniDoc failed: %v", err) 151 151 + } 152 152 + 153 153 + if did != "did:plc:custom" { 154 154 + t.Errorf("DID = %s, want did:plc:custom", did) 155 155 + } 156 156 + if pds != "https://custom.pds.example.com" { 157 157 + t.Errorf("PDS = %s, want https://custom.pds.example.com", pds) 158 158 + } 159 159 + } 160 160 + 161 161 + func TestResolveMiniDoc_WithUserAgent(t *testing.T) { 162 162 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 163 163 + if r.Header.Get("User-Agent") != "test-agent/1.0" { 164 164 + t.Errorf("expected User-Agent 'test-agent/1.0', got %s", r.Header.Get("User-Agent")) 165 165 + } 166 166 + w.Header().Set("Content-Type", "application/json") 167 167 + json.NewEncoder(w).Encode(map[string]any{ 168 168 + "did": "did:plc:agent", 169 169 + "pds": "https://pds.example.com", 170 170 + "signing_key": "key", 171 171 + }) 172 172 + })) 173 173 + defer server.Close() 174 174 + 175 175 + opts := NewClientOptions() 176 176 + opts.ResolverURL = server.URL 177 177 + opts.UserAgent = "test-agent/1.0" 178 178 + 179 179 + ctx := context.Background() 180 180 + did, _, _, err := ResolveMiniDoc(ctx, "test.user", opts) 181 181 + if err != nil { 182 182 + t.Fatalf("ResolveMiniDoc failed: %v", err) 183 183 + } 184 184 + 185 185 + if did != "did:plc:agent" { 186 186 + t.Errorf("DID = %s, want did:plc:agent", did) 187 187 + } 188 188 + } 189 189 + 190 190 + func TestResolveMiniDoc_ContextCancelled(t *testing.T) { 191 191 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 192 192 + select { 193 193 + case <-time.After(100 * time.Millisecond): 194 194 + w.WriteHeader(http.StatusOK) 195 195 + case <-r.Context().Done(): 196 196 + return 197 197 + } 198 198 + })) 199 199 + defer server.Close() 200 200 + 201 201 + ctx, cancel := context.WithCancel(context.Background()) 202 202 + cancel() 203 203 + 204 204 + opts := NewClientOptions() 205 205 + opts.ResolverURL = server.URL 206 206 + opts.HTTPClient = server.Client() 207 207 + 208 208 + _, _, _, err := ResolveMiniDoc(ctx, "test.user", opts) 209 209 + 210 210 + if err == nil { 211 211 + t.Error("expected error for cancelled context") 212 212 + } 213 213 + } 214 214 + 215 215 + func TestResolveMiniDoc_InvalidURL(t *testing.T) { 216 216 + ctx := context.Background() 217 217 + 218 218 + opts := NewClientOptions() 219 219 + opts.ResolverURL = "://invalid-url" 220 220 + 221 221 + _, _, _, err := ResolveMiniDoc(ctx, "test.user", opts) 222 222 + 223 223 + if err == nil { 224 224 + t.Error("expected error for invalid URL") 225 225 + } 226 226 + } 227 227 + 228 228 + func TestResolveIdentity(t *testing.T) { 229 229 + t.Parallel() 230 230 + 231 231 + tests := []struct { 232 232 + name string 233 233 + handle string 234 234 + body map[string]any 235 235 + opts *ClientOptions 236 236 + wantErr bool 237 237 + checkFunc func(t *testing.T, identity resolvedIdentity, err error) 238 238 + }{ 239 239 + { 240 240 + name: "empty handle", 241 241 + handle: "", 242 242 + opts: NewClientOptions(), 243 243 + wantErr: true, 244 244 + }, 245 245 + { 246 246 + name: "missing signing key", 247 247 + handle: "test.user", 248 248 + body: map[string]any{ 249 249 + "did": "did:plc:z72iitness", 250 250 + "pds": "https://pds.example.com", 251 251 + "signing_key": "", 252 252 + }, 253 253 + opts: NewClientOptions(), 254 254 + wantErr: true, 255 255 + }, 256 256 + { 257 257 + name: "with custom resolver", 258 258 + handle: "custom.handle", 259 259 + body: map[string]any{ 260 260 + "did": "did:plc:custom123", 261 261 + "pds": "https://custom.pds.example.com", 262 262 + "signing_key": "-----BEGIN PUBLIC KEY-----\ncustom\n-----END PUBLIC KEY-----", 263 263 + }, 264 264 + opts: NewClientOptions(), 265 265 + checkFunc: func(t *testing.T, identity resolvedIdentity, err error) { 266 266 + if identity.DID != "did:plc:custom123" { 267 267 + t.Errorf("DID = %s, want did:plc:custom123", identity.DID) 268 268 + } 269 269 + if identity.PDS != "https://custom.pds.example.com" { 270 270 + t.Errorf("PDS = %s, want https://custom.pds.example.com", identity.PDS) 271 271 + } 272 272 + if identity.Handle != "custom.handle" { 273 273 + t.Errorf("Handle = %s, want custom.handle", identity.Handle) 274 274 + } 275 275 + }, 276 276 + }, 277 277 + } 278 278 + 279 279 + for _, tt := range tests { 280 280 + t.Run(tt.name, func(t *testing.T) { 281 281 + t.Parallel() 282 282 + 283 283 + if tt.handle == "" { 284 284 + ctx := context.Background() 285 285 + _, err := ResolveIdentity(ctx, "", tt.opts) 286 286 + if err == nil { 287 287 + t.Error("expected error for empty handle") 288 288 + } 289 289 + return 290 290 + } 291 291 + 292 292 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 293 293 + w.Header().Set("Content-Type", "application/json") 294 294 + json.NewEncoder(w).Encode(tt.body) 295 295 + })) 296 296 + defer server.Close() 297 297 + 298 298 + ctx := context.Background() 299 299 + 300 300 + opts := tt.opts 301 301 + if opts == nil { 302 302 + opts = NewClientOptions() 303 303 + } 304 304 + if opts.ResolverURL == DefaultResolverURL { 305 305 + opts.ResolverURL = server.URL 306 306 + } 307 307 + 308 308 + identity, err := ResolveIdentity(ctx, tt.handle, opts) 309 309 + 310 310 + if tt.wantErr { 311 311 + if err == nil { 312 312 + t.Error("expected error, got nil") 313 313 + } 314 314 + return 315 315 + } 316 316 + 317 317 + if err != nil { 318 318 + t.Fatalf("ResolveIdentity failed: %v", err) 319 319 + } 320 320 + 321 321 + if tt.checkFunc != nil { 322 322 + tt.checkFunc(t, identity, err) 323 323 + } 324 324 + }) 325 325 + } 326 326 + } 327 327 + 328 328 + func TestClient_Getters(t *testing.T) { 329 329 + t.Parallel() 330 330 + 331 331 + tests := []struct { 332 332 + name string 333 333 + setup func() *Client 334 334 + check func(t *testing.T, c *Client) 335 335 + }{ 336 336 + { 337 337 + name: "HasClient true", 338 338 + setup: func() *Client { 339 339 + return &Client{ 340 340 + client: &atclient.APIClient{}, 341 341 + resolvedIdentity: resolvedIdentity{ 342 342 + DID: "did:plc:test", 343 343 + Handle: "test.bsky.social", 344 344 + PDS: "https://pds.example.com", 345 345 + }, 346 346 + } 347 347 + }, 348 348 + check: func(t *testing.T, c *Client) { 349 349 + if !c.HasClient() { 350 350 + t.Error("HasClient() = false, want true") 351 351 + } 352 352 + }, 353 353 + }, 354 354 + { 355 355 + name: "HasClient false", 356 356 + setup: func() *Client { 357 357 + return &Client{} 358 358 + }, 359 359 + check: func(t *testing.T, c *Client) { 360 360 + if c.HasClient() { 361 361 + t.Error("HasClient() = true, want false") 362 362 + } 363 363 + }, 364 364 + }, 365 365 + { 366 366 + name: "DID", 367 367 + setup: func() *Client { 368 368 + return &Client{ 369 369 + resolvedIdentity: resolvedIdentity{DID: "did:plc:test123"}, 370 370 + } 371 371 + }, 372 372 + check: func(t *testing.T, c *Client) { 373 373 + if got := c.DID(); got != "did:plc:test123" { 374 374 + t.Errorf("DID() = %s, want did:plc:test123", got) 375 375 + } 376 376 + }, 377 377 + }, 378 378 + { 379 379 + name: "PDS", 380 380 + setup: func() *Client { 381 381 + return &Client{ 382 382 + resolvedIdentity: resolvedIdentity{PDS: "https://pds.example.com"}, 383 383 + } 384 384 + }, 385 385 + check: func(t *testing.T, c *Client) { 386 386 + if got := c.PDS(); got != "https://pds.example.com" { 387 387 + t.Errorf("PDS() = %s, want https://pds.example.com", got) 388 388 + } 389 389 + }, 390 390 + }, 391 391 + { 392 392 + name: "Handle", 393 393 + setup: func() *Client { 394 394 + return &Client{ 395 395 + resolvedIdentity: resolvedIdentity{Handle: "test.bsky.social"}, 396 396 + } 397 397 + }, 398 398 + check: func(t *testing.T, c *Client) { 399 399 + if got := c.Handle(); got != "test.bsky.social" { 400 400 + t.Errorf("Handle() = %s, want test.bsky.social", got) 401 401 + } 402 402 + }, 403 403 + }, 404 404 + { 405 405 + name: "SigningKey", 406 406 + setup: func() *Client { 407 407 + return &Client{ 408 408 + resolvedIdentity: resolvedIdentity{SigningKey: "-----BEGIN PUBLIC KEY-----\nabc\n-----END PUBLIC KEY-----"}, 409 409 + } 410 410 + }, 411 411 + check: func(t *testing.T, c *Client) { 412 412 + if got := c.SigningKey(); got != "-----BEGIN PUBLIC KEY-----\nabc\n-----END PUBLIC KEY-----" { 413 413 + t.Errorf("SigningKey() = %s, want expected key", got) 414 414 + } 415 415 + }, 416 416 + }, 417 417 + { 418 418 + name: "APIClient", 419 419 + setup: func() *Client { 420 420 + expectedClient := &atclient.APIClient{} 421 421 + return &Client{ 422 422 + client: expectedClient, 423 423 + resolvedIdentity: resolvedIdentity{DID: "did:plc:test"}, 424 424 + } 425 425 + }, 426 426 + check: func(t *testing.T, c *Client) { 427 427 + if got := c.APIClient(); got == nil { 428 428 + t.Error("APIClient() = nil, want non-nil") 429 429 + } 430 430 + }, 431 431 + }, 432 432 + { 433 433 + name: "APIClient nil", 434 434 + setup: func() *Client { 435 435 + return &Client{} 436 436 + }, 437 437 + check: func(t *testing.T, c *Client) { 438 438 + if got := c.APIClient(); got != nil { 439 439 + t.Errorf("APIClient() = %v, want nil", got) 440 440 + } 441 441 + }, 442 442 + }, 443 443 + } 444 444 + 445 445 + for _, tt := range tests { 446 446 + t.Run(tt.name, func(t *testing.T) { 447 447 + t.Parallel() 448 448 + c := tt.setup() 449 449 + tt.check(t, c) 450 450 + }) 451 451 + } 452 452 + } 453 453 + 454 454 + func TestBuildClient(t *testing.T) { 455 455 + t.Parallel() 456 456 + 457 457 + tests := []struct { 458 458 + name string 459 459 + customClient RepoClient[map[string]any] 460 460 + apiClient *atclient.APIClient 461 461 + did string 462 462 + wantErr bool 463 463 + wantCustom bool 464 464 + }{ 465 465 + { 466 466 + name: "nil custom client", 467 467 + customClient: nil, 468 468 + apiClient: &atclient.APIClient{}, 469 469 + did: "did:plc:test", 470 470 + wantErr: false, 471 471 + wantCustom: false, 472 472 + }, 473 473 + { 474 474 + name: "with custom client", 475 475 + customClient: &mockRepoClient[map[string]any]{}, 476 476 + apiClient: &atclient.APIClient{}, 477 477 + did: "did:plc:test", 478 478 + wantErr: false, 479 479 + wantCustom: true, 480 480 + }, 481 481 + { 482 482 + name: "nil API client", 483 483 + customClient: nil, 484 484 + apiClient: nil, 485 485 + did: "did:plc:test", 486 486 + wantErr: true, 487 487 + wantCustom: false, 488 488 + }, 489 489 + } 490 490 + 491 491 + for _, tt := range tests { 492 492 + t.Run(tt.name, func(t *testing.T) { 493 493 + t.Parallel() 494 494 + 495 495 + identity := &mockAuthClient{apiClient: tt.apiClient, did: tt.did} 496 496 + result, err := BuildClient[map[string]any](identity, tt.customClient) 497 497 + 498 498 + if tt.wantErr { 499 499 + if err == nil { 500 500 + t.Error("expected error, got nil") 501 501 + } 502 502 + return 503 503 + } 504 504 + 505 505 + if err != nil { 506 506 + t.Fatalf("BuildClient failed: %v", err) 507 507 + } 508 508 + 509 509 + if result == nil { 510 510 + t.Fatal("BuildClient returned nil") 511 511 + } 512 512 + 513 513 + if tt.wantCustom && result != tt.customClient { 514 514 + t.Error("BuildClient should return custom client") 515 515 + } 516 516 + }) 517 517 + } 518 518 + } 519 519 + 520 520 + func TestClientOptions(t *testing.T) { 521 521 + t.Parallel() 522 522 + 523 523 + t.Run("WithUserAgent", func(t *testing.T) { 524 524 + t.Parallel() 525 525 + opts := NewClientOptions() 526 526 + WithUserAgent("test-agent")(opts) 527 527 + if opts.UserAgent != "test-agent" { 528 528 + t.Errorf("UserAgent = %s, want test-agent", opts.UserAgent) 529 529 + } 530 530 + }) 531 531 + 532 532 + t.Run("WithResolverURL", func(t *testing.T) { 533 533 + t.Parallel() 534 534 + opts := NewClientOptions() 535 535 + WithResolverURL("http://custom-resolver")(opts) 536 536 + if opts.ResolverURL != "http://custom-resolver" { 537 537 + t.Errorf("ResolverURL = %s, want http://custom-resolver", opts.ResolverURL) 538 538 + } 539 539 + }) 540 540 + 541 541 + t.Run("WithHTTPClient", func(t *testing.T) { 542 542 + t.Parallel() 543 543 + opts := NewClientOptions() 544 544 + customClient := &http.Client{} 545 545 + WithHTTPClient(customClient)(opts) 546 546 + if opts.HTTPClient != customClient { 547 547 + t.Errorf("HTTPClient = %v, want %v", opts.HTTPClient, customClient) 548 548 + } 549 549 + }) 550 550 + 551 551 + t.Run("WithRateLimitPercent", func(t *testing.T) { 552 552 + t.Parallel() 553 553 + opts := NewClientOptions() 554 554 + WithRateLimitPercent(0.5)(opts) 555 555 + if opts.RateLimitPercent != 0.5 { 556 556 + t.Errorf("RateLimitPercent = %f, want 0.5", opts.RateLimitPercent) 557 557 + } 558 558 + }) 559 559 + 560 560 + t.Run("defaults", func(t *testing.T) { 561 561 + t.Parallel() 562 562 + opts := NewClientOptions() 563 563 + if opts.RateLimitPercent != DefaultRateLimitPercent { 564 564 + t.Errorf("RateLimitPercent = %f, want %f", opts.RateLimitPercent, DefaultRateLimitPercent) 565 565 + } 566 566 + if opts.ResolverURL != DefaultResolverURL { 567 567 + t.Errorf("ResolverURL = %s, want %s", opts.ResolverURL, DefaultResolverURL) 568 568 + } 569 569 + if opts.HTTPClient == nil { 570 570 + t.Error("HTTPClient should not be nil") 571 571 + } 572 572 + }) 573 573 + } 574 574 + 575 575 + type mockAuthClient struct { 576 576 + apiClient *atclient.APIClient 577 577 + did string 578 578 + } 579 579 + 580 580 + func (m *mockAuthClient) APIClient() *atclient.APIClient { 581 581 + return m.apiClient 582 582 + } 583 583 + 584 584 + func (m *mockAuthClient) DID() string { 585 585 + return m.did 586 586 + } 587 587 + 588 588 + type mockRepoClient[T any] struct{} 589 589 + 590 590 + func (m *mockRepoClient[T]) ListRecords(ctx context.Context, collection string, limit int, cursor string) ([]RecordRef[T], string, error) { 591 591 + return nil, "", nil 592 592 + } 593 593 + 594 594 + func (m *mockRepoClient[T]) ApplyWrites(ctx context.Context, collection string, records []T) error { 595 595 + return nil 596 596 + } 597 597 + 598 598 + func (m *mockRepoClient[T]) DeleteRecord(ctx context.Context, collection, rkey string) error { 599 599 + return nil 600 600 + } 601 601 + 602 602 + func TestFixedPasswordAuth_RefreshUsesPOST(t *testing.T) { 603 603 + var capturedMethod string 604 604 + var capturedPath string 605 605 + 606 606 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 607 607 + capturedMethod = r.Method 608 608 + capturedPath = r.URL.Path 609 609 + 610 610 + if capturedMethod == http.MethodPost && capturedPath == "/xrpc/com.atproto.server.refreshSession" { 611 611 + w.Header().Set("Content-Type", "application/json") 612 612 + json.NewEncoder(w).Encode(map[string]any{ 613 613 + "accessJwt": "new-access-token", 614 614 + "refreshJwt": "new-refresh-token", 615 615 + }) 616 616 + return 617 617 + } 618 618 + 619 619 + if capturedPath == "/xrpc/com.atproto.server.createSession" { 620 620 + w.Header().Set("Content-Type", "application/json") 621 621 + json.NewEncoder(w).Encode(map[string]any{ 622 622 + "accessJwt": "access-token", 623 623 + "refreshJwt": "refresh-token", 624 624 + "did": "did:plc:test", 625 625 + }) 626 626 + return 627 627 + } 628 628 + 629 629 + w.WriteHeader(http.StatusBadRequest) 630 630 + })) 631 631 + defer server.Close() 632 632 + 633 633 + pdsURL := server.URL 634 634 + 635 635 + session := atclient.PasswordSessionData{ 636 636 + AccessToken: "old-access-token", 637 637 + RefreshToken: "refresh-token", 638 638 + Host: pdsURL, 639 639 + } 640 640 + 641 641 + auth := &FixedPasswordAuth{ 642 642 + PasswordAuth: &atclient.PasswordAuth{Session: session}, 643 643 + } 644 644 + 645 645 + httpClient := server.Client() 646 646 + 647 647 + ctx := context.Background() 648 648 + err := auth.Refresh(ctx, httpClient, "refresh-token") 649 649 + 650 650 + if err != nil { 651 651 + t.Fatalf("Refresh failed: %v", err) 652 652 + } 653 653 + 654 654 + if capturedMethod != http.MethodPost { 655 655 + t.Errorf("Refresh used %s, want POST", capturedMethod) 656 656 + } 657 657 + 658 658 + if capturedPath != "/xrpc/com.atproto.server.refreshSession" { 659 659 + t.Errorf("Refresh hit %s, want /xrpc/com.atproto.server.refreshSession", capturedPath) 660 660 + } 661 661 + } 662 662 + 663 663 + func TestFixedPasswordAuth_IndigoBugCheck(t *testing.T) { 664 664 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 665 665 + if r.URL.Path == "/xrpc/com.atproto.server.refreshSession" { 666 666 + if r.Method != http.MethodPost { 667 667 + t.Errorf("REFRESH BUG: indigo uses %s for refreshSession, should use POST", r.Method) 668 668 + } 669 669 + w.Header().Set("Content-Type", "application/json") 670 670 + json.NewEncoder(w).Encode(map[string]any{ 671 671 + "accessJwt": "new-access-token", 672 672 + "refreshJwt": "new-refresh-token", 673 673 + }) 674 674 + return 675 675 + } 676 676 + if r.URL.Path == "/xrpc/com.atproto.server.createSession" { 677 677 + w.Header().Set("Content-Type", "application/json") 678 678 + json.NewEncoder(w).Encode(map[string]any{ 679 679 + "accessJwt": "access-token", 680 680 + "refreshJwt": "refresh-token", 681 681 + "did": "did:plc:test", 682 682 + }) 683 683 + return 684 684 + } 685 685 + w.WriteHeader(http.StatusBadRequest) 686 686 + })) 687 687 + defer server.Close() 688 688 + 689 689 + pdsURL := server.URL 690 690 + 691 691 + pa := &atclient.PasswordAuth{ 692 692 + Session: atclient.PasswordSessionData{ 693 693 + AccessToken: "old-access-token", 694 694 + RefreshToken: "refresh-token", 695 695 + Host: pdsURL, 696 696 + }, 697 697 + } 698 698 + 699 699 + fixedAuth := &FixedPasswordAuth{PasswordAuth: pa} 700 700 + 701 701 + err := fixedAuth.Refresh(context.Background(), server.Client(), "refresh-token") 702 702 + 703 703 + if err != nil { 704 704 + t.Fatalf("FixedPasswordAuth.Refresh failed: %v", err) 705 705 + } 706 706 + }

+694

atproto/rate_test.go

reviewed

··· 1 1 + package atproto 2 2 + 3 3 + import ( 4 4 + "context" 5 5 + "fmt" 6 6 + "testing" 7 7 + "time" 8 8 + ) 9 9 + 10 10 + func TestNewRateLimiter(t *testing.T) { 11 11 + kv := &mockKVStore{data: make(map[string]int)} 12 12 + rl := NewRateLimiter(kv, 0.8) 13 13 + 14 14 + if rl == nil { 15 15 + t.Fatal("NewRateLimiter returned nil") 16 16 + } 17 17 + } 18 18 + 19 19 + func TestQuotaLimiter_Stats(t *testing.T) { 20 20 + kv := &mockKVStore{data: map[string]int{ 21 21 + "quota:writes:d:2024-01-15": 100, 22 22 + "quota:global:d:2024-01-15": 200, 23 23 + }} 24 24 + clock := &mockClock{now: time.Date(2024, 1, 15, 12, 0, 0, 0, time.UTC)} 25 25 + 26 26 + rl := &quotaLimiter{ 27 27 + kv: kv, 28 28 + prefix: "quota", 29 29 + clock: clock, 30 30 + rlQuota: 1.0, 31 31 + } 32 32 + 33 33 + writes, global, err := rl.Stats() 34 34 + if err != nil { 35 35 + t.Fatalf("Stats failed: %v", err) 36 36 + } 37 37 + if writes != 100 { 38 38 + t.Errorf("Stats writes = %d, want 100", writes) 39 39 + } 40 40 + if global != 200 { 41 41 + t.Errorf("Stats global = %d, want 200", global) 42 42 + } 43 43 + } 44 44 + 45 45 + func TestQuotaLimiter_AllowRead(t *testing.T) { 46 46 + kv := &mockKVStore{data: make(map[string]int)} 47 47 + clock := &mockClock{now: time.Date(2024, 1, 15, 12, 0, 0, 0, time.UTC)} 48 48 + 49 49 + rl := &quotaLimiter{ 50 50 + kv: kv, 51 51 + prefix: "quota", 52 52 + clock: clock, 53 53 + rlQuota: 1.0, 54 54 + } 55 55 + 56 56 + ctx := context.Background() 57 57 + chargedAt, err := rl.AllowRead(ctx) 58 58 + if err != nil { 59 59 + t.Fatalf("AllowRead failed: %v", err) 60 60 + } 61 61 + 62 62 + if kv.incrs == nil || len(kv.incrs) == 0 { 63 63 + t.Error("Expected kv.IncrByMulti to be called") 64 64 + } 65 65 + if !chargedAt.IsZero() { 66 66 + t.Logf("Charged at: %v", chargedAt) 67 67 + } 68 68 + } 69 69 + 70 70 + func TestQuotaLimiter_AllowBulkWrite(t *testing.T) { 71 71 + t.Parallel() 72 72 + 73 73 + tests := []struct { 74 74 + name string 75 75 + n int 76 76 + kvData map[string]int 77 77 + wantErr bool 78 78 + }{ 79 79 + { 80 80 + name: "success with 5 records", 81 81 + n: 5, 82 82 + kvData: make(map[string]int), 83 83 + wantErr: false, 84 84 + }, 85 85 + } 86 86 + 87 87 + for _, tt := range tests { 88 88 + t.Run(tt.name, func(t *testing.T) { 89 89 + t.Parallel() 90 90 + 91 91 + kv := &mockKVStore{data: tt.kvData} 92 92 + clock := &mockClock{now: time.Date(2024, 1, 15, 12, 0, 0, 0, time.UTC)} 93 93 + 94 94 + rl := &quotaLimiter{ 95 95 + kv: kv, 96 96 + prefix: "quota", 97 97 + clock: clock, 98 98 + rlQuota: 1.0, 99 99 + } 100 100 + 101 101 + ctx := context.Background() 102 102 + _, err := rl.AllowBulkWrite(ctx, tt.n) 103 103 + 104 104 + if tt.wantErr && err == nil { 105 105 + t.Error("expected error, got nil") 106 106 + } 107 107 + if !tt.wantErr && err != nil { 108 108 + t.Fatalf("AllowBulkWrite failed: %v", err) 109 109 + } 110 110 + }) 111 111 + } 112 112 + } 113 113 + 114 114 + func TestQuotaLimiter_Refund(t *testing.T) { 115 115 + t.Parallel() 116 116 + 117 117 + tests := []struct { 118 118 + name string 119 119 + setupKV func() *mockKVStore 120 120 + refundFunc func(rl *quotaLimiter, ctx context.Context) 121 121 + wantIncrCalls int 122 122 + }{ 123 123 + { 124 124 + name: "refund bulk write", 125 125 + setupKV: func() *mockKVStore { 126 126 + return &mockKVStore{data: map[string]int{ 127 127 + "quota:writes:m:2024-01-15-12-00": 10, 128 128 + "quota:writes:h:2024-01-15-12": 20, 129 129 + "quota:writes:d:2024-01-15": 30, 130 130 + "quota:global:m:2024-01-15-12-00": 5, 131 131 + "quota:global:h:2024-01-15-12": 10, 132 132 + "quota:global:d:2024-01-15": 15, 133 133 + }} 134 134 + }, 135 135 + refundFunc: func(rl *quotaLimiter, ctx context.Context) { 136 136 + rl.RefundBulkWrite(ctx, 3, time.Now()) 137 137 + }, 138 138 + wantIncrCalls: 1, 139 139 + }, 140 140 + { 141 141 + name: "refund read", 142 142 + setupKV: func() *mockKVStore { 143 143 + return &mockKVStore{data: map[string]int{ 144 144 + "quota:global:m:2024-01-15-12-00": 5, 145 145 + "quota:global:h:2024-01-15-12": 10, 146 146 + "quota:global:d:2024-01-15": 15, 147 147 + }} 148 148 + }, 149 149 + refundFunc: func(rl *quotaLimiter, ctx context.Context) { 150 150 + rl.RefundRead(ctx, time.Now()) 151 151 + }, 152 152 + wantIncrCalls: 1, 153 153 + }, 154 154 + } 155 155 + 156 156 + for _, tt := range tests { 157 157 + t.Run(tt.name, func(t *testing.T) { 158 158 + t.Parallel() 159 159 + 160 160 + kv := tt.setupKV() 161 161 + clock := &mockClock{now: time.Date(2024, 1, 15, 12, 0, 0, 0, time.UTC)} 162 162 + 163 163 + rl := &quotaLimiter{ 164 164 + kv: kv, 165 165 + prefix: "quota", 166 166 + clock: clock, 167 167 + rlQuota: 1.0, 168 168 + } 169 169 + 170 170 + ctx := context.Background() 171 171 + tt.refundFunc(rl, ctx) 172 172 + 173 173 + if len(kv.incrs) != tt.wantIncrCalls { 174 174 + t.Errorf("Expected %d IncrByMulti call, got %d", tt.wantIncrCalls, len(kv.incrs)) 175 175 + } 176 176 + }) 177 177 + } 178 178 + } 179 179 + 180 180 + func TestQuotaLimiter_EstimatedWriteTime(t *testing.T) { 181 181 + t.Parallel() 182 182 + 183 183 + tests := []struct { 184 184 + name string 185 185 + n int 186 186 + kvData map[string]int 187 187 + wantZero bool 188 188 + }{ 189 189 + { 190 190 + name: "zero records", 191 191 + n: 0, 192 192 + kvData: make(map[string]int), 193 193 + wantZero: true, 194 194 + }, 195 195 + { 196 196 + name: "negative records", 197 197 + n: -5, 198 198 + kvData: make(map[string]int), 199 199 + wantZero: true, 200 200 + }, 201 201 + { 202 202 + name: "with existing usage", 203 203 + n: 60, 204 204 + kvData: map[string]int{ 205 205 + "quota:writes:m:2024-01-15-12-00": 50, 206 206 + "quota:writes:h:2024-01-15-12": 500, 207 207 + }, 208 208 + wantZero: false, 209 209 + }, 210 210 + } 211 211 + 212 212 + for _, tt := range tests { 213 213 + t.Run(tt.name, func(t *testing.T) { 214 214 + t.Parallel() 215 215 + 216 216 + kv := &mockKVStore{data: tt.kvData} 217 217 + clock := &mockClock{now: time.Date(2024, 1, 15, 12, 0, 30, 0, time.UTC)} 218 218 + 219 219 + rl := &quotaLimiter{ 220 220 + kv: kv, 221 221 + prefix: "quota", 222 222 + clock: clock, 223 223 + rlQuota: 1.0, 224 224 + } 225 225 + 226 226 + wait := rl.EstimatedWriteTime(tt.n) 227 227 + 228 228 + if tt.wantZero && wait != 0 { 229 229 + t.Errorf("EstimatedWriteTime(%d) = %v, want 0", tt.n, wait) 230 230 + } 231 231 + if !tt.wantZero && tt.n > 0 && wait == 0 { 232 232 + t.Log("EstimatedWriteTime returned 0 with existing usage") 233 233 + } 234 234 + }) 235 235 + } 236 236 + } 237 237 + 238 238 + func TestQuotaLimiter_RemainingQuota(t *testing.T) { 239 239 + kv := &mockKVStore{data: make(map[string]int)} 240 240 + clock := &mockClock{now: time.Date(2024, 1, 15, 12, 0, 0, 0, time.UTC)} 241 241 + 242 242 + rl := &quotaLimiter{ 243 243 + kv: kv, 244 244 + prefix: "quota", 245 245 + clock: clock, 246 246 + rlQuota: 1.0, 247 247 + } 248 248 + 249 249 + writes, global, reset := rl.RemainingQuota() 250 250 + if writes != WriteLimitMinute { 251 251 + t.Errorf("RemainingQuota writes = %d, want %d (minute limit)", writes, WriteLimitMinute) 252 252 + } 253 253 + if global != GlobalLimitMinute { 254 254 + t.Errorf("RemainingQuota global = %d, want %d (minute limit)", global, GlobalLimitMinute) 255 255 + } 256 256 + if reset < 0 { 257 257 + t.Errorf("RemainingQuota reset = %v, want >= 0", reset) 258 258 + } 259 259 + } 260 260 + 261 261 + func TestUntilNextWindow(t *testing.T) { 262 262 + t.Parallel() 263 263 + 264 264 + tests := []struct { 265 265 + name string 266 266 + now time.Time 267 267 + tier int 268 268 + wantInRange bool 269 269 + minExpected time.Duration 270 270 + }{ 271 271 + { 272 272 + name: "minute tier at 30s past", 273 273 + now: time.Date(2024, 1, 15, 12, 0, 30, 0, time.UTC), 274 274 + tier: 0, 275 275 + wantInRange: true, 276 276 + minExpected: 30*time.Second + 100*time.Millisecond, 277 277 + }, 278 278 + { 279 279 + name: "hour tier at 30m past", 280 280 + now: time.Date(2024, 1, 15, 12, 30, 0, 0, time.UTC), 281 281 + tier: 1, 282 282 + wantInRange: true, 283 283 + minExpected: 30*time.Minute + 100*time.Millisecond, 284 284 + }, 285 285 + { 286 286 + name: "day tier at 6 hours past", 287 287 + now: time.Date(2024, 1, 15, 18, 0, 0, 0, time.UTC), 288 288 + tier: 2, 289 289 + wantInRange: true, 290 290 + minExpected: 6*time.Hour + 100*time.Millisecond, 291 291 + }, 292 292 + { 293 293 + name: "unknown tier", 294 294 + now: time.Now(), 295 295 + tier: 99, 296 296 + wantInRange: false, 297 297 + minExpected: time.Minute, 298 298 + }, 299 299 + } 300 300 + 301 301 + for _, tt := range tests { 302 302 + t.Run(tt.name, func(t *testing.T) { 303 303 + t.Parallel() 304 304 + 305 305 + kv := &mockKVStore{data: make(map[string]int)} 306 306 + clock := &mockClock{now: tt.now} 307 307 + 308 308 + rl := &quotaLimiter{ 309 309 + kv: kv, 310 310 + prefix: "quota", 311 311 + clock: clock, 312 312 + rlQuota: 1.0, 313 313 + } 314 314 + 315 315 + wait := rl.untilNextWindow(tt.now, tt.tier) 316 316 + 317 317 + if tt.wantInRange { 318 318 + expected := tt.minExpected 319 319 + if wait < expected-time.Second || wait > expected+time.Second { 320 320 + t.Errorf("untilNextWindow(%d) = %v, want ~%v", tt.tier, wait, expected) 321 321 + } 322 322 + } else { 323 323 + if wait != tt.minExpected { 324 324 + t.Errorf("untilNextWindow(%d) = %v, want %v", tt.tier, wait, tt.minExpected) 325 325 + } 326 326 + } 327 327 + }) 328 328 + } 329 329 + } 330 330 + 331 331 + func TestAddJitter(t *testing.T) { 332 332 + base := 100 * time.Millisecond 333 333 + 334 334 + for range 100 { 335 335 + jitter := addJitter(base) 336 336 + if jitter < base { 337 337 + t.Errorf("addJitter returned %v, want >= %v", jitter, base) 338 338 + } 339 339 + if jitter > base+base/10 { 340 340 + t.Errorf("addJitter returned %v, want <= %v + 10%%", jitter, base) 341 341 + } 342 342 + } 343 343 + } 344 344 + 345 345 + func TestQuotaLimiter_EmptyKeys(t *testing.T) { 346 346 + t.Parallel() 347 347 + 348 348 + tests := []struct { 349 349 + name string 350 350 + wKeys []string 351 351 + gKeys []string 352 352 + wCost int 353 353 + gCost int 354 354 + wantWait time.Duration 355 355 + }{ 356 356 + { 357 357 + name: "empty keys for checkQuota", 358 358 + wKeys: nil, 359 359 + gKeys: nil, 360 360 + wCost: 0, 361 361 + gCost: 0, 362 362 + wantWait: 0, 363 363 + }, 364 364 + { 365 365 + name: "empty keys for charge", 366 366 + wKeys: nil, 367 367 + gKeys: nil, 368 368 + wCost: 0, 369 369 + gCost: 0, 370 370 + wantWait: 0, 371 371 + }, 372 372 + } 373 373 + 374 374 + for _, tt := range tests { 375 375 + t.Run(tt.name, func(t *testing.T) { 376 376 + t.Parallel() 377 377 + 378 378 + kv := &mockKVStore{data: make(map[string]int)} 379 379 + clock := &mockClock{now: time.Date(2024, 1, 15, 12, 0, 0, 0, time.UTC)} 380 380 + 381 381 + rl := &quotaLimiter{ 382 382 + kv: kv, 383 383 + prefix: "quota", 384 384 + clock: clock, 385 385 + rlQuota: 1.0, 386 386 + } 387 387 + 388 388 + now := time.Now() 389 389 + 390 390 + if tt.name == "empty keys for checkQuota" { 391 391 + wait, err := rl.checkQuota(now, tt.wKeys, tt.gKeys, tt.wCost, tt.gCost) 392 392 + if err != nil { 393 393 + t.Fatalf("checkQuota failed: %v", err) 394 394 + } 395 395 + if wait != tt.wantWait { 396 396 + t.Errorf("checkQuota = %v, want %v", wait, tt.wantWait) 397 397 + } 398 398 + } else { 399 399 + err := rl.charge(tt.wKeys, tt.gKeys, tt.wCost, tt.gCost) 400 400 + if err != nil { 401 401 + t.Fatalf("charge failed: %v", err) 402 402 + } 403 403 + if len(kv.incrs) > 0 { 404 404 + t.Errorf("charge with empty keys should not call IncrByMulti") 405 405 + } 406 406 + } 407 407 + }) 408 408 + } 409 409 + } 410 410 + 411 411 + func TestRateLimiter_Interfaces(t *testing.T) { 412 412 + t.Parallel() 413 413 + 414 414 + tests := []struct { 415 415 + name string 416 416 + checkFunc func(t *testing.T) 417 417 + }{ 418 418 + { 419 419 + name: "RateLimiter interface", 420 420 + checkFunc: func(t *testing.T) { 421 421 + var _ RateLimiter = (*quotaLimiter)(nil) 422 422 + }, 423 423 + }, 424 424 + { 425 425 + name: "KVStore interface", 426 426 + checkFunc: func(t *testing.T) { 427 427 + var _ KVStore = (*mockKVStore)(nil) 428 428 + }, 429 429 + }, 430 430 + { 431 431 + name: "Clock interface - realClock", 432 432 + checkFunc: func(t *testing.T) { 433 433 + var _ Clock = (*realClock)(nil) 434 434 + }, 435 435 + }, 436 436 + { 437 437 + name: "Clock interface - mockClock", 438 438 + checkFunc: func(t *testing.T) { 439 439 + var _ Clock = (*mockClock)(nil) 440 440 + }, 441 441 + }, 442 442 + } 443 443 + 444 444 + for _, tt := range tests { 445 445 + t.Run(tt.name, func(t *testing.T) { 446 446 + t.Parallel() 447 447 + tt.checkFunc(t) 448 448 + }) 449 449 + } 450 450 + } 451 451 + 452 452 + func TestRateLimiterConstants(t *testing.T) { 453 453 + tests := []struct { 454 454 + name string 455 455 + got int 456 456 + expected int 457 457 + }{ 458 458 + {"WriteLimitMinute", WriteLimitMinute, 100}, 459 459 + {"WriteLimitHour", WriteLimitHour, 1000}, 460 460 + {"WriteLimitDay", WriteLimitDay, 10000}, 461 461 + {"GlobalLimitMinute", GlobalLimitMinute, 300}, 462 462 + {"GlobalLimitHour", GlobalLimitHour, 3000}, 463 463 + {"GlobalLimitDay", GlobalLimitDay, 35000}, 464 464 + {"ReadGlobalCost", ReadGlobalCost, 1}, 465 465 + {"WriteOnlyCost", WriteOnlyCost, 1}, 466 466 + {"WriteGlobalCost", WriteGlobalCost, 3}, 467 467 + } 468 468 + 469 469 + for _, tt := range tests { 470 470 + t.Run(tt.name, func(t *testing.T) { 471 471 + if tt.got != tt.expected { 472 472 + t.Errorf("%s = %d, want %d", tt.name, tt.got, tt.expected) 473 473 + } 474 474 + }) 475 475 + } 476 476 + } 477 477 + 478 478 + func TestQuotaLimiter_getKeys(t *testing.T) { 479 479 + kv := &mockKVStore{data: make(map[string]int)} 480 480 + clock := &mockClock{now: time.Date(2024, 1, 15, 12, 30, 45, 0, time.UTC)} 481 481 + 482 482 + rl := &quotaLimiter{ 483 483 + kv: kv, 484 484 + prefix: "test", 485 485 + clock: clock, 486 486 + rlQuota: 1.0, 487 487 + } 488 488 + 489 489 + wd, gd, wh, gh, wm, gm := rl.getKeys(time.Date(2024, 1, 15, 12, 30, 45, 0, time.UTC)) 490 490 + 491 491 + expectedKeys := []struct { 492 492 + got string 493 493 + want string 494 494 + }{ 495 495 + {wd, "test:writes:d:2024-01-15"}, 496 496 + {gd, "test:global:d:2024-01-15"}, 497 497 + {wh, "test:writes:h:2024-01-15-12"}, 498 498 + {gh, "test:global:h:2024-01-15-12"}, 499 499 + {wm, "test:writes:m:2024-01-15-12-30"}, 500 500 + {gm, "test:global:m:2024-01-15-12-30"}, 501 501 + } 502 502 + 503 503 + for _, k := range expectedKeys { 504 504 + if k.got != k.want { 505 505 + t.Errorf("key = %s, want %s", k.got, k.want) 506 506 + } 507 507 + } 508 508 + } 509 509 + 510 510 + func TestQuotaLimiter_getAllKeys(t *testing.T) { 511 511 + kv := &mockKVStore{data: make(map[string]int)} 512 512 + clock := &mockClock{now: time.Date(2024, 1, 15, 12, 30, 45, 0, time.UTC)} 513 513 + 514 514 + rl := &quotaLimiter{ 515 515 + kv: kv, 516 516 + prefix: "test", 517 517 + clock: clock, 518 518 + rlQuota: 1.0, 519 519 + } 520 520 + 521 521 + tm := time.Date(2024, 1, 15, 12, 30, 45, 0, time.UTC) 522 522 + wKeys, gKeys := rl.getAllKeys(tm) 523 523 + 524 524 + if len(wKeys) != 3 { 525 525 + t.Errorf("getAllKeys returned %d write keys, want 3", len(wKeys)) 526 526 + } 527 527 + if len(gKeys) != 3 { 528 528 + t.Errorf("getAllKeys returned %d global keys, want 3", len(gKeys)) 529 529 + } 530 530 + } 531 531 + 532 532 + func TestQuotaLimiter_KVErrors(t *testing.T) { 533 533 + t.Parallel() 534 534 + 535 535 + tests := []struct { 536 536 + name string 537 537 + testFunc func(t *testing.T, rl *quotaLimiter) 538 538 + }{ 539 539 + { 540 540 + name: "Stats error", 541 541 + testFunc: func(t *testing.T, rl *quotaLimiter) { 542 542 + _, _, err := rl.Stats() 543 543 + if err == nil { 544 544 + t.Error("expected error from KV store") 545 545 + } 546 546 + }, 547 547 + }, 548 548 + { 549 549 + name: "AllowBulkWrite error", 550 550 + testFunc: func(t *testing.T, rl *quotaLimiter) { 551 551 + ctx := context.Background() 552 552 + _, err := rl.AllowBulkWrite(ctx, 1) 553 553 + if err == nil { 554 554 + t.Error("expected error from KV store") 555 555 + } 556 556 + }, 557 557 + }, 558 558 + { 559 559 + name: "AllowRead error", 560 560 + testFunc: func(t *testing.T, rl *quotaLimiter) { 561 561 + ctx := context.Background() 562 562 + _, err := rl.AllowRead(ctx) 563 563 + if err == nil { 564 564 + t.Error("expected error from KV store") 565 565 + } 566 566 + }, 567 567 + }, 568 568 + { 569 569 + name: "charge error", 570 570 + testFunc: func(t *testing.T, rl *quotaLimiter) { 571 571 + err := rl.charge([]string{"test:writes:m:2024-01-15-12-00"}, nil, 1, 0) 572 572 + if err == nil { 573 573 + t.Error("expected error from KV store") 574 574 + } 575 575 + }, 576 576 + }, 577 577 + { 578 578 + name: "checkQuota error", 579 579 + testFunc: func(t *testing.T, rl *quotaLimiter) { 580 580 + now := time.Now() 581 581 + _, err := rl.checkQuota(now, []string{"test:writes:m:2024-01-15-12-00"}, nil, 1, 0) 582 582 + if err == nil { 583 583 + t.Error("expected error from KV store") 584 584 + } 585 585 + }, 586 586 + }, 587 587 + } 588 588 + 589 589 + for _, tt := range tests { 590 590 + t.Run(tt.name, func(t *testing.T) { 591 591 + t.Parallel() 592 592 + 593 593 + kv := &mockKVStoreWithErr{Data: make(map[string]int), Err: fmt.Errorf("KV error")} 594 594 + clock := &mockClock{now: time.Date(2024, 1, 15, 12, 0, 0, 0, time.UTC)} 595 595 + 596 596 + rl := &quotaLimiter{ 597 597 + kv: kv, 598 598 + prefix: "quota", 599 599 + clock: clock, 600 600 + rlQuota: 1.0, 601 601 + } 602 602 + 603 603 + tt.testFunc(t, rl) 604 604 + }) 605 605 + } 606 606 + } 607 607 + 608 608 + func TestQuotaLimiter_Refund_KVErrors(t *testing.T) { 609 609 + t.Parallel() 610 610 + 611 611 + tests := []struct { 612 612 + name string 613 613 + setupKV func() *mockKVStoreWithErr 614 614 + refundFn func(rl *quotaLimiter, ctx context.Context) 615 615 + }{ 616 616 + { 617 617 + name: "RefundBulkWrite error", 618 618 + setupKV: func() *mockKVStoreWithErr { 619 619 + return &mockKVStoreWithErr{Data: map[string]int{ 620 620 + "quota:writes:m:2024-01-15-12-00": 10, 621 621 + }, Err: fmt.Errorf("KV error")} 622 622 + }, 623 623 + refundFn: func(rl *quotaLimiter, ctx context.Context) { 624 624 + rl.RefundBulkWrite(ctx, 3, time.Now()) 625 625 + }, 626 626 + }, 627 627 + { 628 628 + name: "RefundRead error", 629 629 + setupKV: func() *mockKVStoreWithErr { 630 630 + return &mockKVStoreWithErr{Data: map[string]int{ 631 631 + "quota:global:m:2024-01-15-12-00": 5, 632 632 + }, Err: fmt.Errorf("KV error")} 633 633 + }, 634 634 + refundFn: func(rl *quotaLimiter, ctx context.Context) { 635 635 + rl.RefundRead(ctx, time.Now()) 636 636 + }, 637 637 + }, 638 638 + } 639 639 + 640 640 + for _, tt := range tests { 641 641 + t.Run(tt.name, func(t *testing.T) { 642 642 + t.Parallel() 643 643 + 644 644 + kv := tt.setupKV() 645 645 + clock := &mockClock{now: time.Date(2024, 1, 15, 12, 0, 0, 0, time.UTC)} 646 646 + 647 647 + rl := &quotaLimiter{ 648 648 + kv: kv, 649 649 + prefix: "quota", 650 650 + clock: clock, 651 651 + rlQuota: 1.0, 652 652 + } 653 653 + 654 654 + ctx := context.Background() 655 655 + tt.refundFn(rl, ctx) 656 656 + }) 657 657 + } 658 658 + } 659 659 + 660 660 + func TestQuotaLimiter_MutexProtection(t *testing.T) { 661 661 + kv := &mockKVStore{data: make(map[string]int)} 662 662 + clock := &mockClock{now: time.Date(2024, 1, 15, 12, 0, 0, 0, time.UTC)} 663 663 + 664 664 + rl := &quotaLimiter{ 665 665 + kv: kv, 666 666 + prefix: "quota", 667 667 + clock: clock, 668 668 + rlQuota: 1.0, 669 669 + } 670 670 + 671 671 + ctx := context.Background() 672 672 + done := make(chan bool) 673 673 + errors := make(chan error, 10) 674 674 + 675 675 + for range 10 { 676 676 + go func() { 677 677 + _, err := rl.AllowBulkWrite(ctx, 1) 678 678 + if err != nil { 679 679 + errors <- err 680 680 + } 681 681 + done <- true 682 682 + }() 683 683 + } 684 684 + 685 685 + for range 10 { 686 686 + <-done 687 687 + } 688 688 + 689 689 + select { 690 690 + case err := <-errors: 691 691 + t.Errorf("concurrent AllowBulkWrite returned error: %v", err) 692 692 + default: 693 693 + } 694 694 + }

+246

atproto/repo.go

reviewed

··· 1 1 + package atproto 2 2 + 3 3 + import ( 4 4 + "context" 5 5 + "encoding/json" 6 6 + "errors" 7 7 + "fmt" 8 8 + "log/slog" 9 9 + "net" 10 10 + "time" 11 11 + 12 12 + "github.com/bluesky-social/indigo/atproto/atclient" 13 13 + "github.com/bluesky-social/indigo/atproto/syntax" 14 14 + ) 15 15 + 16 16 + type RecordRef[T any] struct { 17 17 + URI string 18 18 + CID string 19 19 + Value T 20 20 + } 21 21 + 22 22 + type RepoClient[T any] interface { 23 23 + ListRecords(ctx context.Context, collection string, limit int, cursor string) ([]RecordRef[T], string, error) 24 24 + ApplyWrites(ctx context.Context, collection string, records []T) error 25 25 + DeleteRecord(ctx context.Context, collection, rkey string) error 26 26 + } 27 27 + 28 28 + type RepoClientFuncs[T any] struct { 29 29 + ListRecordsFn func(ctx context.Context, collection string, limit int, cursor string) ([]RecordRef[T], string, error) 30 30 + ApplyWritesFn func(ctx context.Context, collection string, records []T) error 31 31 + DeleteRecordFn func(ctx context.Context, collection, rkey string) error 32 32 + } 33 33 + 34 34 + func (c RepoClientFuncs[T]) ListRecords(ctx context.Context, collection string, limit int, cursor string) ([]RecordRef[T], string, error) { 35 35 + return c.ListRecordsFn(ctx, collection, limit, cursor) 36 36 + } 37 37 + 38 38 + func (c RepoClientFuncs[T]) ApplyWrites(ctx context.Context, collection string, records []T) error { 39 39 + return c.ApplyWritesFn(ctx, collection, records) 40 40 + } 41 41 + 42 42 + func (c RepoClientFuncs[T]) DeleteRecord(ctx context.Context, collection, rkey string) error { 43 43 + return c.DeleteRecordFn(ctx, collection, rkey) 44 44 + } 45 45 + 46 46 + type RateClient[T any] struct { 47 47 + client *atclient.APIClient 48 48 + did string 49 49 + limiter RateLimiter 50 50 + } 51 51 + 52 52 + func NewRateClient[T any](client *atclient.APIClient, did string, limiter RateLimiter) *RateClient[T] { 53 53 + return &RateClient[T]{ 54 54 + client: client, 55 55 + did: did, 56 56 + limiter: limiter, 57 57 + } 58 58 + } 59 59 + 60 60 + func (c *RateClient[T]) ListRecords(ctx context.Context, collection string, limit int, cursor string) ([]RecordRef[T], string, error) { 61 61 + if c.client == nil { 62 62 + return nil, "", fmt.Errorf("client cannot be nil") 63 63 + } 64 64 + 65 65 + var outResp struct { 66 66 + Records []struct { 67 67 + URI string `json:"uri"` 68 68 + CID string `json:"cid"` 69 69 + Value map[string]any `json:"value"` 70 70 + } `json:"records"` 71 71 + Cursor string `json:"cursor"` 72 72 + } 73 73 + 74 74 + var chargedAt time.Time 75 75 + if c.limiter != nil { 76 76 + slog.Debug("waiting for rate limit (read)") 77 77 + var err error 78 78 + chargedAt, err = c.limiter.AllowRead(ctx) 79 79 + if err != nil { 80 80 + slog.Error("rate limit wait cancelled/failed (read)", Error(err)) 81 81 + return nil, "", err 82 82 + } 83 83 + } 84 84 + 85 85 + err := c.client.Get(ctx, syntax.NSID("com.atproto.repo.listRecords"), map[string]any{ 86 86 + "repo": c.did, 87 87 + "collection": collection, 88 88 + "limit": limit, 89 89 + "cursor": cursor, 90 90 + }, &outResp) 91 91 + if err != nil { 92 92 + if c.limiter != nil && IsTransientError(err) { 93 93 + c.limiter.RefundRead(ctx, chargedAt) 94 94 + } 95 95 + return nil, "", err 96 96 + } 97 97 + 98 98 + out := make([]RecordRef[T], 0, len(outResp.Records)) 99 99 + for _, r := range outResp.Records { 100 100 + var value T 101 101 + if r.Value != nil { 102 102 + b, err := json.Marshal(r.Value) 103 103 + if err != nil { 104 104 + slog.Debug("failed to marshal record value", slog.String("uri", r.URI), Error(err)) 105 105 + continue 106 106 + } 107 107 + if err := json.Unmarshal(b, &value); err != nil { 108 108 + slog.Debug("failed to unmarshal record", slog.String("uri", r.URI), Error(err)) 109 109 + continue 110 110 + } 111 111 + } 112 112 + out = append(out, RecordRef[T]{ 113 113 + URI: r.URI, 114 114 + CID: r.CID, 115 115 + Value: value, 116 116 + }) 117 117 + } 118 118 + 119 119 + return out, outResp.Cursor, nil 120 120 + } 121 121 + 122 122 + func (c *RateClient[T]) ApplyWrites(ctx context.Context, collection string, records []T) error { 123 123 + if len(records) == 0 { 124 124 + return nil 125 125 + } 126 126 + 127 127 + if c.client == nil { 128 128 + return fmt.Errorf("client cannot be nil") 129 129 + } 130 130 + 131 131 + var chargedAt time.Time 132 132 + if c.limiter != nil { 133 133 + var err error 134 134 + chargedAt, err = c.limiter.AllowBulkWrite(ctx, len(records)) 135 135 + if err != nil { 136 136 + slog.Error("rate limit wait cancelled/failed (write)", 137 137 + slog.String("did", c.did), 138 138 + slog.String("collection", collection), 139 139 + Error(err), 140 140 + ) 141 141 + return err 142 142 + } 143 143 + } 144 144 + 145 145 + err := applyWrites(ctx, c.client, c.did, collection, records) 146 146 + if err != nil && IsTransientError(err) && c.limiter != nil { 147 147 + c.limiter.RefundBulkWrite(ctx, len(records), chargedAt) 148 148 + } 149 149 + return err 150 150 + } 151 151 + 152 152 + func applyWrites[T any](ctx context.Context, client *atclient.APIClient, did, collection string, records []T) error { 153 153 + if len(records) == 0 { 154 154 + return nil 155 155 + } 156 156 + 157 157 + if len(records) > 200 { 158 158 + return fmt.Errorf("too many records in one ApplyWrites call: %d (max 200)", len(records)) 159 159 + } 160 160 + 161 161 + writes, err := prepareWrites(records, collection) 162 162 + if err != nil { 163 163 + return err 164 164 + } 165 165 + 166 166 + return client.Post(ctx, syntax.NSID("com.atproto.repo.applyWrites"), map[string]any{ 167 167 + "repo": did, 168 168 + "writes": writes, 169 169 + }, nil) 170 170 + } 171 171 + 172 172 + func prepareWrites[T any](records []T, collection string) ([]map[string]any, error) { 173 173 + if len(records) == 0 { 174 174 + return nil, nil 175 175 + } 176 176 + 177 177 + writes := make([]map[string]any, len(records)) 178 178 + for i, rec := range records { 179 179 + b, err := json.Marshal(rec) 180 180 + if err != nil { 181 181 + return nil, fmt.Errorf("failed to marshal record: %w", err) 182 182 + } 183 183 + var recMap map[string]any 184 184 + if err := json.Unmarshal(b, &recMap); err != nil { 185 185 + return nil, fmt.Errorf("failed to unmarshal record: %w", err) 186 186 + } 187 187 + recMap["$type"] = "com.atproto.repo.applyWrites#create" 188 188 + recMap["collection"] = collection 189 189 + writes[i] = recMap 190 190 + } 191 191 + 192 192 + return writes, nil 193 193 + } 194 194 + 195 195 + func (c *RateClient[T]) DeleteRecord(ctx context.Context, collection, rkey string) error { 196 196 + if c.client == nil { 197 197 + return fmt.Errorf("client is nil") 198 198 + } 199 199 + 200 200 + var chargedAt time.Time 201 201 + if c.limiter != nil { 202 202 + slog.Debug("waiting for rate limit (delete)") 203 203 + var err error 204 204 + chargedAt, err = c.limiter.AllowBulkWrite(ctx, 1) 205 205 + if err != nil { 206 206 + slog.Error("rate limit wait cancelled/failed (delete)", Error(err)) 207 207 + return err 208 208 + } 209 209 + } 210 210 + 211 211 + _, err := c.client.Do(ctx, &atclient.APIRequest{ 212 212 + Method: "DELETE", 213 213 + Endpoint: syntax.NSID("com.atproto.repo.deleteRecord"), 214 214 + QueryParams: map[string][]string{ 215 215 + "repo": {c.did}, 216 216 + "collection": {collection}, 217 217 + "rkey": {rkey}, 218 218 + }, 219 219 + }) 220 220 + if err != nil && c.limiter != nil && IsTransientError(err) { 221 221 + c.limiter.RefundBulkWrite(ctx, 1, chargedAt) 222 222 + } 223 223 + return err 224 224 + } 225 225 + 226 226 + func IsTransientError(err error) bool { 227 227 + if err == nil { 228 228 + return false 229 229 + } 230 230 + 231 231 + var apiErr *atclient.APIError 232 232 + if errors.As(err, &apiErr) { 233 233 + switch apiErr.StatusCode { 234 234 + case 429, 500, 502, 503, 504: 235 235 + return true 236 236 + } 237 237 + return false 238 238 + } 239 239 + 240 240 + var netErr net.Error 241 241 + if errors.As(err, &netErr) { 242 242 + return netErr.Timeout() 243 243 + } 244 244 + 245 245 + return false 246 246 + }

+945

atproto/repo_test.go

reviewed

··· 1 1 + package atproto 2 2 + 3 3 + import ( 4 4 + "context" 5 5 + "encoding/json" 6 6 + "errors" 7 7 + "net/http" 8 8 + "net/http/httptest" 9 9 + "testing" 10 10 + "time" 11 11 + 12 12 + "github.com/bluesky-social/indigo/atproto/atclient" 13 13 + ) 14 14 + 15 15 + func TestRateClient_ListRecords(t *testing.T) { 16 16 + t.Parallel() 17 17 + 18 18 + tests := []struct { 19 19 + name string 20 20 + setupServer func() *httptest.Server 21 21 + client *atclient.APIClient 22 22 + limiter RateLimiter 23 23 + collection string 24 24 + limit int 25 25 + cursor string 26 26 + wantRecords int 27 27 + wantCursor string 28 28 + wantErr bool 29 29 + }{ 30 30 + { 31 31 + name: "success", 32 32 + setupServer: func() *httptest.Server { 33 33 + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 34 34 + if r.URL.Path != "/xrpc/com.atproto.repo.listRecords" { 35 35 + t.Errorf("unexpected path: %s", r.URL.Path) 36 36 + } 37 37 + if r.URL.Query().Get("repo") != "did:plc:test" { 38 38 + t.Errorf("unexpected repo: %s", r.URL.Query().Get("repo")) 39 39 + } 40 40 + if r.URL.Query().Get("collection") != "app.bsky.feed.post" { 41 41 + t.Errorf("unexpected collection: %s", r.URL.Query().Get("collection")) 42 42 + } 43 43 + 44 44 + response := map[string]any{ 45 45 + "records": []any{ 46 46 + map[string]any{ 47 47 + "uri": "at://did:plc:test/app.bsky.feed.post/3k5x3x2x1", 48 48 + "cid": "bafyre...", 49 49 + "value": map[string]any{ 50 50 + "$type": "app.bsky.feed.post", 51 51 + "text": "Hello world", 52 52 + }, 53 53 + }, 54 54 + }, 55 55 + "cursor": "next-cursor", 56 56 + } 57 57 + w.Header().Set("Content-Type", "application/json") 58 58 + json.NewEncoder(w).Encode(response) 59 59 + })) 60 60 + }, 61 61 + client: func() *atclient.APIClient { 62 62 + return &atclient.APIClient{} 63 63 + }(), 64 64 + collection: "app.bsky.feed.post", 65 65 + limit: 10, 66 66 + cursor: "", 67 67 + wantRecords: 1, 68 68 + wantCursor: "next-cursor", 69 69 + wantErr: false, 70 70 + }, 71 71 + { 72 72 + name: "nil client", 73 73 + setupServer: func() *httptest.Server { 74 74 + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})) 75 75 + }, 76 76 + client: nil, 77 77 + collection: "app.bsky.feed.post", 78 78 + limit: 10, 79 79 + cursor: "", 80 80 + wantRecords: 0, 81 81 + wantErr: true, 82 82 + }, 83 83 + { 84 84 + name: "with limiter", 85 85 + setupServer: func() *httptest.Server { 86 86 + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 87 87 + response := map[string]any{ 88 88 + "records": []any{}, 89 89 + } 90 90 + w.Header().Set("Content-Type", "application/json") 91 91 + json.NewEncoder(w).Encode(response) 92 92 + })) 93 93 + }, 94 94 + client: func() *atclient.APIClient { 95 95 + return &atclient.APIClient{} 96 96 + }(), 97 97 + limiter: NewRateLimiter(&mockKVStore{data: make(map[string]int)}, 1.0), 98 98 + collection: "app.bsky.feed.post", 99 99 + limit: 10, 100 100 + cursor: "", 101 101 + wantRecords: 0, 102 102 + wantErr: false, 103 103 + }, 104 104 + { 105 105 + name: "empty records", 106 106 + setupServer: func() *httptest.Server { 107 107 + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 108 108 + w.Header().Set("Content-Type", "application/json") 109 109 + json.NewEncoder(w).Encode(map[string]any{ 110 110 + "records": []any{}, 111 111 + "cursor": "", 112 112 + }) 113 113 + })) 114 114 + }, 115 115 + client: func() *atclient.APIClient { 116 116 + return &atclient.APIClient{} 117 117 + }(), 118 118 + collection: "app.bsky.feed.post", 119 119 + limit: 10, 120 120 + cursor: "", 121 121 + wantRecords: 0, 122 122 + wantCursor: "", 123 123 + wantErr: false, 124 124 + }, 125 125 + { 126 126 + name: "4xx error", 127 127 + setupServer: func() *httptest.Server { 128 128 + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 129 129 + w.WriteHeader(http.StatusBadRequest) 130 130 + })) 131 131 + }, 132 132 + client: func() *atclient.APIClient { 133 133 + return &atclient.APIClient{} 134 134 + }(), 135 135 + collection: "app.bsky.feed.post", 136 136 + limit: 10, 137 137 + cursor: "", 138 138 + wantRecords: 0, 139 139 + wantErr: true, 140 140 + }, 141 141 + { 142 142 + name: "5xx error", 143 143 + setupServer: func() *httptest.Server { 144 144 + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 145 145 + w.WriteHeader(http.StatusInternalServerError) 146 146 + })) 147 147 + }, 148 148 + client: func() *atclient.APIClient { 149 149 + return &atclient.APIClient{} 150 150 + }(), 151 151 + collection: "app.bsky.feed.post", 152 152 + limit: 10, 153 153 + cursor: "", 154 154 + wantRecords: 0, 155 155 + wantErr: true, 156 156 + }, 157 157 + } 158 158 + 159 159 + for _, tt := range tests { 160 160 + t.Run(tt.name, func(t *testing.T) { 161 161 + t.Parallel() 162 162 + 163 163 + server := tt.setupServer() 164 164 + defer server.Close() 165 165 + 166 166 + if tt.client != nil && tt.client.Client == nil { 167 167 + tt.client.Client = server.Client() 168 168 + tt.client.Host = server.URL 169 169 + } 170 170 + 171 171 + rateClient := NewRateClient[map[string]any](tt.client, "did:plc:test", tt.limiter) 172 172 + ctx := context.Background() 173 173 + 174 174 + records, cursor, err := rateClient.ListRecords(ctx, tt.collection, tt.limit, tt.cursor) 175 175 + 176 176 + if tt.wantErr { 177 177 + if err == nil { 178 178 + t.Error("expected error, got nil") 179 179 + } 180 180 + return 181 181 + } 182 182 + 183 183 + if err != nil { 184 184 + t.Fatalf("ListRecords failed: %v", err) 185 185 + } 186 186 + 187 187 + if len(records) != tt.wantRecords { 188 188 + t.Errorf("ListRecords returned %d records, want %d", len(records), tt.wantRecords) 189 189 + } 190 190 + if tt.wantCursor != "" && cursor != tt.wantCursor { 191 191 + t.Errorf("ListRecords cursor = %s, want %s", cursor, tt.wantCursor) 192 192 + } 193 193 + }) 194 194 + } 195 195 + } 196 196 + 197 197 + func TestRateClient_ListRecords_NetworkTimeout(t *testing.T) { 198 198 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 199 199 + time.Sleep(100 * time.Millisecond) 200 200 + w.WriteHeader(http.StatusOK) 201 201 + })) 202 202 + defer server.Close() 203 203 + 204 204 + client := atclient.APIClient{ 205 205 + Client: &http.Client{Timeout: 10 * time.Millisecond}, 206 206 + Host: server.URL, 207 207 + } 208 208 + 209 209 + rateClient := NewRateClient[map[string]any](&client, "did:plc:test", nil) 210 210 + ctx := context.Background() 211 211 + 212 212 + _, _, err := rateClient.ListRecords(ctx, "app.bsky.feed.post", 10, "") 213 213 + 214 214 + if err == nil { 215 215 + t.Error("expected error for network timeout") 216 216 + } 217 217 + } 218 218 + 219 219 + func TestRateClient_ListRecords_ContextCancelled(t *testing.T) { 220 220 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 221 221 + select { 222 222 + case <-time.After(100 * time.Millisecond): 223 223 + w.WriteHeader(http.StatusOK) 224 224 + case <-r.Context().Done(): 225 225 + return 226 226 + } 227 227 + })) 228 228 + defer server.Close() 229 229 + 230 230 + client := atclient.APIClient{ 231 231 + Client: server.Client(), 232 232 + Host: server.URL, 233 233 + } 234 234 + 235 235 + rateClient := NewRateClient[map[string]any](&client, "did:plc:test", nil) 236 236 + ctx, cancel := context.WithCancel(context.Background()) 237 237 + cancel() 238 238 + 239 239 + _, _, err := rateClient.ListRecords(ctx, "app.bsky.feed.post", 10, "") 240 240 + 241 241 + if err == nil { 242 242 + t.Error("expected error for cancelled context") 243 243 + } 244 244 + } 245 245 + 246 246 + func TestRateClient_ListRecords_WithPagination(t *testing.T) { 247 247 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 248 248 + cursor := r.URL.Query().Get("cursor") 249 249 + if cursor == "page1" { 250 250 + w.Header().Set("Content-Type", "application/json") 251 251 + json.NewEncoder(w).Encode(map[string]any{ 252 252 + "records": []any{}, 253 253 + "cursor": "page2", 254 254 + }) 255 255 + } else { 256 256 + w.Header().Set("Content-Type", "application/json") 257 257 + json.NewEncoder(w).Encode(map[string]any{ 258 258 + "records": []any{ 259 259 + map[string]any{ 260 260 + "uri": "at://did:plc:test/app.bsky.feed.post/1", 261 261 + "cid": "bafyre...", 262 262 + "value": map[string]any{ 263 263 + "$type": "app.bsky.feed.post", 264 264 + "text": "test", 265 265 + }, 266 266 + }, 267 267 + }, 268 268 + "cursor": "page1", 269 269 + }) 270 270 + } 271 271 + })) 272 272 + defer server.Close() 273 273 + 274 274 + client := atclient.APIClient{ 275 275 + Client: server.Client(), 276 276 + Host: server.URL, 277 277 + } 278 278 + 279 279 + rateClient := NewRateClient[map[string]any](&client, "did:plc:test", nil) 280 280 + ctx := context.Background() 281 281 + 282 282 + records, cursor, err := rateClient.ListRecords(ctx, "app.bsky.feed.post", 10, "") 283 283 + if err != nil { 284 284 + t.Fatalf("ListRecords failed: %v", err) 285 285 + } 286 286 + if len(records) != 1 { 287 287 + t.Errorf("ListRecords returned %d records, want 1", len(records)) 288 288 + } 289 289 + if cursor != "page1" { 290 290 + t.Errorf("ListRecords cursor = %s, want page1", cursor) 291 291 + } 292 292 + } 293 293 + 294 294 + func TestRateClient_ApplyWrites(t *testing.T) { 295 295 + t.Parallel() 296 296 + 297 297 + tests := []struct { 298 298 + name string 299 299 + setupServer func() *httptest.Server 300 300 + client *atclient.APIClient 301 301 + records []map[string]any 302 302 + wantErr bool 303 303 + }{ 304 304 + { 305 305 + name: "success", 306 306 + setupServer: func() *httptest.Server { 307 307 + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 308 308 + if r.URL.Path != "/xrpc/com.atproto.repo.applyWrites" { 309 309 + t.Errorf("unexpected path: %s", r.URL.Path) 310 310 + } 311 311 + w.WriteHeader(http.StatusOK) 312 312 + })) 313 313 + }, 314 314 + client: func() *atclient.APIClient { 315 315 + return &atclient.APIClient{} 316 316 + }(), 317 317 + records: []map[string]any{ 318 318 + {"$type": "app.bsky.feed.post", "text": "Hello"}, 319 319 + {"$type": "app.bsky.feed.post", "text": "World"}, 320 320 + }, 321 321 + wantErr: false, 322 322 + }, 323 323 + { 324 324 + name: "empty records", 325 325 + setupServer: func() *httptest.Server { 326 326 + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})) 327 327 + }, 328 328 + client: &atclient.APIClient{}, 329 329 + records: []map[string]any{}, 330 330 + wantErr: false, 331 331 + }, 332 332 + { 333 333 + name: "nil client", 334 334 + setupServer: func() *httptest.Server { 335 335 + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})) 336 336 + }, 337 337 + client: nil, 338 338 + records: []map[string]any{ 339 339 + {"$type": "app.bsky.feed.post", "text": "Hello"}, 340 340 + }, 341 341 + wantErr: true, 342 342 + }, 343 343 + { 344 344 + name: "too many records", 345 345 + setupServer: func() *httptest.Server { 346 346 + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})) 347 347 + }, 348 348 + client: &atclient.APIClient{}, 349 349 + records: make([]map[string]any, 201), 350 350 + wantErr: true, 351 351 + }, 352 352 + { 353 353 + name: "4xx error", 354 354 + setupServer: func() *httptest.Server { 355 355 + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 356 356 + w.WriteHeader(http.StatusBadRequest) 357 357 + })) 358 358 + }, 359 359 + client: func() *atclient.APIClient { 360 360 + return &atclient.APIClient{} 361 361 + }(), 362 362 + records: []map[string]any{ 363 363 + {"$type": "app.bsky.feed.post", "text": "Hello"}, 364 364 + }, 365 365 + wantErr: true, 366 366 + }, 367 367 + { 368 368 + name: "5xx error", 369 369 + setupServer: func() *httptest.Server { 370 370 + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 371 371 + w.WriteHeader(http.StatusInternalServerError) 372 372 + })) 373 373 + }, 374 374 + client: func() *atclient.APIClient { 375 375 + return &atclient.APIClient{} 376 376 + }(), 377 377 + records: []map[string]any{ 378 378 + {"$type": "app.bsky.feed.post", "text": "Hello"}, 379 379 + }, 380 380 + wantErr: true, 381 381 + }, 382 382 + } 383 383 + 384 384 + for _, tt := range tests { 385 385 + t.Run(tt.name, func(t *testing.T) { 386 386 + t.Parallel() 387 387 + 388 388 + server := tt.setupServer() 389 389 + defer server.Close() 390 390 + 391 391 + if tt.client != nil && tt.client.Client == nil { 392 392 + tt.client.Client = server.Client() 393 393 + tt.client.Host = server.URL 394 394 + } 395 395 + 396 396 + rateClient := NewRateClient[map[string]any](tt.client, "did:plc:test", nil) 397 397 + ctx := context.Background() 398 398 + 399 399 + err := rateClient.ApplyWrites(ctx, "app.bsky.feed.post", tt.records) 400 400 + 401 401 + if tt.wantErr && err == nil { 402 402 + t.Error("expected error, got nil") 403 403 + } 404 404 + if !tt.wantErr && err != nil { 405 405 + t.Fatalf("ApplyWrites failed: %v", err) 406 406 + } 407 407 + }) 408 408 + } 409 409 + } 410 410 + 411 411 + func TestRateClient_ApplyWrites_NetworkTimeout(t *testing.T) { 412 412 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 413 413 + time.Sleep(100 * time.Millisecond) 414 414 + w.WriteHeader(http.StatusOK) 415 415 + })) 416 416 + defer server.Close() 417 417 + 418 418 + client := atclient.APIClient{ 419 419 + Client: &http.Client{Timeout: 10 * time.Millisecond}, 420 420 + Host: server.URL, 421 421 + } 422 422 + 423 423 + rateClient := NewRateClient[map[string]any](&client, "did:plc:test", nil) 424 424 + ctx := context.Background() 425 425 + 426 426 + records := []map[string]any{ 427 427 + {"$type": "app.bsky.feed.post", "text": "Hello"}, 428 428 + } 429 429 + 430 430 + err := rateClient.ApplyWrites(ctx, "app.bsky.feed.post", records) 431 431 + 432 432 + if err == nil { 433 433 + t.Error("expected error for network timeout") 434 434 + } 435 435 + } 436 436 + 437 437 + func TestRateClient_ApplyWrites_WithLimiter(t *testing.T) { 438 438 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 439 439 + w.WriteHeader(http.StatusOK) 440 440 + })) 441 441 + defer server.Close() 442 442 + 443 443 + client := atclient.APIClient{ 444 444 + Client: server.Client(), 445 445 + Host: server.URL, 446 446 + } 447 447 + 448 448 + mockKV := &mockKVStore{data: make(map[string]int)} 449 449 + limiter := NewRateLimiter(mockKV, 1.0) 450 450 + 451 451 + rateClient := NewRateClient[map[string]any](&client, "did:plc:test", limiter) 452 452 + ctx := context.Background() 453 453 + 454 454 + records := []map[string]any{ 455 455 + {"$type": "app.bsky.feed.post", "text": "Hello"}, 456 456 + } 457 457 + 458 458 + err := rateClient.ApplyWrites(ctx, "app.bsky.feed.post", records) 459 459 + if err != nil { 460 460 + t.Fatalf("ApplyWrites failed: %v", err) 461 461 + } 462 462 + } 463 463 + 464 464 + func TestRateClient_DeleteRecord(t *testing.T) { 465 465 + t.Parallel() 466 466 + 467 467 + tests := []struct { 468 468 + name string 469 469 + setupServer func() *httptest.Server 470 470 + client *atclient.APIClient 471 471 + wantErr bool 472 472 + }{ 473 473 + { 474 474 + name: "success", 475 475 + setupServer: func() *httptest.Server { 476 476 + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 477 477 + if r.Method != "DELETE" { 478 478 + t.Errorf("unexpected method: %s", r.Method) 479 479 + } 480 480 + if r.URL.Path != "/xrpc/com.atproto.repo.deleteRecord" { 481 481 + t.Errorf("unexpected path: %s", r.URL.Path) 482 482 + } 483 483 + if r.URL.Query().Get("repo") != "did:plc:test" { 484 484 + t.Errorf("unexpected repo: %s", r.URL.Query().Get("repo")) 485 485 + } 486 486 + if r.URL.Query().Get("collection") != "app.bsky.feed.post" { 487 487 + t.Errorf("unexpected collection: %s", r.URL.Query().Get("collection")) 488 488 + } 489 489 + if r.URL.Query().Get("rkey") != "3k5x3x2x1" { 490 490 + t.Errorf("unexpected rkey: %s", r.URL.Query().Get("rkey")) 491 491 + } 492 492 + w.WriteHeader(http.StatusOK) 493 493 + })) 494 494 + }, 495 495 + client: func() *atclient.APIClient { 496 496 + return &atclient.APIClient{} 497 497 + }(), 498 498 + wantErr: false, 499 499 + }, 500 500 + { 501 501 + name: "nil client", 502 502 + setupServer: func() *httptest.Server { 503 503 + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})) 504 504 + }, 505 505 + client: nil, 506 506 + wantErr: true, 507 507 + }, 508 508 + { 509 509 + name: "4xx error (skipped)", 510 510 + setupServer: func() *httptest.Server { 511 511 + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 512 512 + w.WriteHeader(http.StatusBadRequest) 513 513 + })) 514 514 + }, 515 515 + client: func() *atclient.APIClient { 516 516 + return &atclient.APIClient{} 517 517 + }(), 518 518 + wantErr: true, 519 519 + }, 520 520 + { 521 521 + name: "5xx error (skipped)", 522 522 + setupServer: func() *httptest.Server { 523 523 + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 524 524 + w.WriteHeader(http.StatusInternalServerError) 525 525 + })) 526 526 + }, 527 527 + client: func() *atclient.APIClient { 528 528 + return &atclient.APIClient{} 529 529 + }(), 530 530 + wantErr: true, 531 531 + }, 532 532 + } 533 533 + 534 534 + for _, tt := range tests { 535 535 + t.Run(tt.name, func(t *testing.T) { 536 536 + t.Parallel() 537 537 + 538 538 + server := tt.setupServer() 539 539 + defer server.Close() 540 540 + 541 541 + if tt.client != nil && tt.client.Client == nil { 542 542 + tt.client.Client = server.Client() 543 543 + tt.client.Host = server.URL 544 544 + } 545 545 + 546 546 + rateClient := NewRateClient[map[string]any](tt.client, "did:plc:test", nil) 547 547 + ctx := context.Background() 548 548 + 549 549 + err := rateClient.DeleteRecord(ctx, "app.bsky.feed.post", "3k5x3x2x1") 550 550 + 551 551 + if tt.wantErr && err == nil { 552 552 + t.Skip("atclient.Do may not return errors for HTTP status codes") 553 553 + } 554 554 + if !tt.wantErr && err != nil { 555 555 + t.Fatalf("DeleteRecord failed: %v", err) 556 556 + } 557 557 + }) 558 558 + } 559 559 + } 560 560 + 561 561 + func TestRateClient_DeleteRecord_NetworkTimeout(t *testing.T) { 562 562 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 563 563 + time.Sleep(100 * time.Millisecond) 564 564 + w.WriteHeader(http.StatusOK) 565 565 + })) 566 566 + defer server.Close() 567 567 + 568 568 + client := atclient.APIClient{ 569 569 + Client: &http.Client{Timeout: 10 * time.Millisecond}, 570 570 + Host: server.URL, 571 571 + } 572 572 + 573 573 + rateClient := NewRateClient[map[string]any](&client, "did:plc:test", nil) 574 574 + ctx := context.Background() 575 575 + 576 576 + err := rateClient.DeleteRecord(ctx, "app.bsky.feed.post", "3k5x3x2x1") 577 577 + 578 578 + if err == nil { 579 579 + t.Error("expected error for network timeout") 580 580 + } 581 581 + } 582 582 + 583 583 + func TestRepoClientFuncs(t *testing.T) { 584 584 + t.Parallel() 585 585 + 586 586 + tests := []struct { 587 587 + name string 588 588 + client *RepoClientFuncs[map[string]any] 589 589 + checkFunc func(t *testing.T, err error) 590 590 + }{ 591 591 + { 592 592 + name: "ListRecords", 593 593 + client: &RepoClientFuncs[map[string]any]{ 594 594 + ListRecordsFn: func(ctx context.Context, collection string, limit int, cursor string) ([]RecordRef[map[string]any], string, error) { 595 595 + return []RecordRef[map[string]any]{ 596 596 + {URI: "at://did:plc:test/app.bsky.feed.post/1", Value: map[string]any{"text": "test"}}, 597 597 + }, "cursor123", nil 598 598 + }, 599 599 + }, 600 600 + checkFunc: func(t *testing.T, err error) { 601 601 + if err != nil { 602 602 + t.Fatalf("ListRecords failed: %v", err) 603 603 + } 604 604 + }, 605 605 + }, 606 606 + { 607 607 + name: "ApplyWrites", 608 608 + client: &RepoClientFuncs[map[string]any]{ 609 609 + ApplyWritesFn: func(ctx context.Context, collection string, records []map[string]any) error { 610 610 + return errors.New("write failed") 611 611 + }, 612 612 + }, 613 613 + checkFunc: func(t *testing.T, err error) { 614 614 + if err == nil { 615 615 + t.Error("expected error, got nil") 616 616 + } 617 617 + }, 618 618 + }, 619 619 + { 620 620 + name: "DeleteRecord", 621 621 + client: &RepoClientFuncs[map[string]any]{ 622 622 + DeleteRecordFn: func(ctx context.Context, collection, rkey string) error { 623 623 + return nil 624 624 + }, 625 625 + }, 626 626 + checkFunc: func(t *testing.T, err error) { 627 627 + if err != nil { 628 628 + t.Errorf("DeleteRecord failed: %v", err) 629 629 + } 630 630 + }, 631 631 + }, 632 632 + } 633 633 + 634 634 + for _, tt := range tests { 635 635 + t.Run(tt.name, func(t *testing.T) { 636 636 + t.Parallel() 637 637 + ctx := context.Background() 638 638 + 639 639 + switch tt.name { 640 640 + case "ListRecords": 641 641 + _, _, err := tt.client.ListRecords(ctx, "app.bsky.feed.post", 10, "") 642 642 + tt.checkFunc(t, err) 643 643 + case "ApplyWrites": 644 644 + err := tt.client.ApplyWrites(ctx, "app.bsky.feed.post", []map[string]any{}) 645 645 + tt.checkFunc(t, err) 646 646 + case "DeleteRecord": 647 647 + err := tt.client.DeleteRecord(ctx, "app.bsky.feed.post", "3k5x3x2x1") 648 648 + tt.checkFunc(t, err) 649 649 + } 650 650 + }) 651 651 + } 652 652 + } 653 653 + 654 654 + func TestPrepareWrites(t *testing.T) { 655 655 + t.Parallel() 656 656 + 657 657 + tests := []struct { 658 658 + name string 659 659 + records []map[string]any 660 660 + collection string 661 661 + wantCount int 662 662 + wantErr bool 663 663 + checkFunc func(t *testing.T, writes []map[string]any) 664 664 + }{ 665 665 + { 666 666 + name: "empty records", 667 667 + records: []map[string]any{}, 668 668 + collection: "app.bsky.feed.post", 669 669 + wantCount: 0, 670 670 + wantErr: false, 671 671 + checkFunc: func(t *testing.T, writes []map[string]any) { 672 672 + if writes != nil { 673 673 + t.Errorf("prepareWrites returned %v, want nil", writes) 674 674 + } 675 675 + }, 676 676 + }, 677 677 + { 678 678 + name: "success", 679 679 + records: []map[string]any{ 680 680 + {"$type": "app.bsky.feed.post", "text": "Hello"}, 681 681 + {"$type": "app.bsky.feed.post", "text": "World"}, 682 682 + }, 683 683 + collection: "app.bsky.feed.post", 684 684 + wantCount: 2, 685 685 + wantErr: false, 686 686 + checkFunc: func(t *testing.T, writes []map[string]any) { 687 687 + for i, w := range writes { 688 688 + if w["$type"] != "com.atproto.repo.applyWrites#create" { 689 689 + t.Errorf("write[%d] $type = %s, want com.atproto.repo.applyWrites#create", i, w["$type"]) 690 690 + } 691 691 + if w["collection"] != "app.bsky.feed.post" { 692 692 + t.Errorf("write[%d] collection = %s, want app.bsky.feed.post", i, w["collection"]) 693 693 + } 694 694 + } 695 695 + }, 696 696 + }, 697 697 + { 698 698 + name: "invalid JSON", 699 699 + records: []map[string]any{ 700 700 + {"text": "test", "channel": make(chan int)}, 701 701 + }, 702 702 + collection: "app.bsky.feed.post", 703 703 + wantCount: 0, 704 704 + wantErr: true, 705 705 + checkFunc: func(t *testing.T, writes []map[string]any) {}, 706 706 + }, 707 707 + { 708 708 + name: "with JSON-encoded data", 709 709 + records: []map[string]any{ 710 710 + {"raw": `{"$type":"app.bsky.feed.post","text":"test"}`}, 711 711 + }, 712 712 + collection: "app.bsky.feed.post", 713 713 + wantCount: 1, 714 714 + wantErr: false, 715 715 + checkFunc: func(t *testing.T, writes []map[string]any) {}, 716 716 + }, 717 717 + } 718 718 + 719 719 + for _, tt := range tests { 720 720 + t.Run(tt.name, func(t *testing.T) { 721 721 + t.Parallel() 722 722 + 723 723 + writes, err := prepareWrites(tt.records, tt.collection) 724 724 + 725 725 + if tt.wantErr { 726 726 + if err == nil { 727 727 + t.Error("expected error, got nil") 728 728 + } 729 729 + return 730 730 + } 731 731 + 732 732 + if err != nil { 733 733 + t.Fatalf("prepareWrites failed: %v", err) 734 734 + } 735 735 + 736 736 + if len(writes) != tt.wantCount { 737 737 + t.Errorf("prepareWrites returned %d writes, want %d", len(writes), tt.wantCount) 738 738 + } 739 739 + 740 740 + if tt.checkFunc != nil { 741 741 + tt.checkFunc(t, writes) 742 742 + } 743 743 + }) 744 744 + } 745 745 + } 746 746 + 747 747 + func TestApplyWrites_Empty(t *testing.T) { 748 748 + err := applyWrites[map[string]any](context.Background(), nil, "did:plc:test", "app.bsky.feed.post", []map[string]any{}) 749 749 + if err != nil { 750 750 + t.Errorf("applyWrites with empty records should not fail: %v", err) 751 751 + } 752 752 + } 753 753 + 754 754 + func TestIsTransientError(t *testing.T) { 755 755 + t.Parallel() 756 756 + 757 757 + tests := []struct { 758 758 + name string 759 759 + statusCode int 760 760 + expectError bool 761 761 + }{ 762 762 + {"TooManyRequests", 429, true}, 763 763 + {"InternalServerError", 500, true}, 764 764 + {"BadGateway", 502, true}, 765 765 + {"ServiceUnavailable", 503, true}, 766 766 + {"GatewayTimeout", 504, true}, 767 767 + {"OK", 200, false}, 768 768 + {"Created", 201, false}, 769 769 + {"BadRequest", 400, false}, 770 770 + {"Unauthorized", 401, false}, 771 771 + {"NotFound", 404, false}, 772 772 + } 773 773 + 774 774 + for _, tt := range tests { 775 775 + t.Run(tt.name, func(t *testing.T) { 776 776 + t.Parallel() 777 777 + 778 778 + apiErr := &atclient.APIError{StatusCode: tt.statusCode} 779 779 + result := IsTransientError(apiErr) 780 780 + 781 781 + if result != tt.expectError { 782 782 + t.Errorf("IsTransientError(status=%d) = %v, want %v", tt.statusCode, result, tt.expectError) 783 783 + } 784 784 + }) 785 785 + } 786 786 + } 787 787 + 788 788 + func TestIsTransientError_NetError(t *testing.T) { 789 789 + t.Parallel() 790 790 + 791 791 + tests := []struct { 792 792 + name string 793 793 + timeout bool 794 794 + wantError bool 795 795 + }{ 796 796 + {"timeout error", true, true}, 797 797 + {"non-timeout error", false, false}, 798 798 + } 799 799 + 800 800 + for _, tt := range tests { 801 801 + t.Run(tt.name, func(t *testing.T) { 802 802 + t.Parallel() 803 803 + 804 804 + err := &netError{timeout: tt.timeout} 805 805 + result := IsTransientError(err) 806 806 + 807 807 + if result != tt.wantError { 808 808 + t.Errorf("IsTransientError(timeout=%v) = %v, want %v", tt.timeout, result, tt.wantError) 809 809 + } 810 810 + }) 811 811 + } 812 812 + } 813 813 + 814 814 + func TestIsTransientError_Nil(t *testing.T) { 815 815 + result := IsTransientError(nil) 816 816 + 817 817 + if result { 818 818 + t.Error("IsTransientError(nil) should return false") 819 819 + } 820 820 + } 821 821 + 822 822 + func TestRecordRef(t *testing.T) { 823 823 + ref := RecordRef[string]{ 824 824 + URI: "at://did:plc:test/app.bsky.feed.post/1", 825 825 + CID: "bafyre...", 826 826 + Value: "test value", 827 827 + } 828 828 + 829 829 + if ref.URI != "at://did:plc:test/app.bsky.feed.post/1" { 830 830 + t.Errorf("RecordRef.URI = %s, want at://did:plc:test/app.bsky.feed.post/1", ref.URI) 831 831 + } 832 832 + if ref.CID != "bafyre..." { 833 833 + t.Errorf("RecordRef.CID = %s, want bafyre...", ref.CID) 834 834 + } 835 835 + if ref.Value != "test value" { 836 836 + t.Errorf("RecordRef.Value = %s, want test value", ref.Value) 837 837 + } 838 838 + } 839 839 + 840 840 + func TestNewRateClient(t *testing.T) { 841 841 + t.Parallel() 842 842 + 843 843 + tests := []struct { 844 844 + name string 845 845 + client *atclient.APIClient 846 846 + did string 847 847 + limiter RateLimiter 848 848 + check func(t *testing.T, rateClient *RateClient[map[string]any]) 849 849 + }{ 850 850 + { 851 851 + name: "nil limiter", 852 852 + client: &atclient.APIClient{}, 853 853 + did: "did:plc:test", 854 854 + limiter: nil, 855 855 + check: func(t *testing.T, rc *RateClient[map[string]any]) { 856 856 + if rc.limiter != nil { 857 857 + t.Error("NewRateClient with nil limiter should have nil limiter") 858 858 + } 859 859 + }, 860 860 + }, 861 861 + { 862 862 + name: "with limiter", 863 863 + client: &atclient.APIClient{}, 864 864 + did: "did:plc:test", 865 865 + limiter: NewRateLimiter(&mockKVStore{data: make(map[string]int)}, 1.0), 866 866 + check: func(t *testing.T, rc *RateClient[map[string]any]) { 867 867 + if rc.limiter == nil { 868 868 + t.Error("NewRateClient should store limiter") 869 869 + } 870 870 + }, 871 871 + }, 872 872 + } 873 873 + 874 874 + for _, tt := range tests { 875 875 + t.Run(tt.name, func(t *testing.T) { 876 876 + t.Parallel() 877 877 + 878 878 + rateClient := NewRateClient[map[string]any](tt.client, tt.did, tt.limiter) 879 879 + 880 880 + if rateClient.client != tt.client { 881 881 + t.Error("NewRateClient should store client") 882 882 + } 883 883 + if rateClient.did != tt.did { 884 884 + t.Error("NewRateClient should store did") 885 885 + } 886 886 + if tt.check != nil { 887 887 + tt.check(t, rateClient) 888 888 + } 889 889 + }) 890 890 + } 891 891 + } 892 892 + 893 893 + func TestRepoClient_Interfaces(t *testing.T) { 894 894 + var _ RepoClient[map[string]any] = (*RateClient[map[string]any])(nil) 895 895 + var _ RepoClient[map[string]any] = (*RepoClientFuncs[map[string]any])(nil) 896 896 + } 897 897 + 898 898 + func TestRateClient_ListRecords_WithTypedRecords(t *testing.T) { 899 899 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 900 900 + type testRecord struct { 901 901 + Text string `json:"text"` 902 902 + } 903 903 + records := []struct { 904 904 + URI string `json:"uri"` 905 905 + CID string `json:"cid"` 906 906 + Value testRecord `json:"value"` 907 907 + }{ 908 908 + { 909 909 + URI: "at://did:plc:test/app.bsky.feed.post/1", 910 910 + CID: "bafyre...", 911 911 + Value: testRecord{ 912 912 + Text: "Hello world", 913 913 + }, 914 914 + }, 915 915 + } 916 916 + w.Header().Set("Content-Type", "application/json") 917 917 + json.NewEncoder(w).Encode(map[string]any{ 918 918 + "records": records, 919 919 + }) 920 920 + })) 921 921 + defer server.Close() 922 922 + 923 923 + client := atclient.APIClient{ 924 924 + Client: server.Client(), 925 925 + Host: server.URL, 926 926 + } 927 927 + 928 928 + type testRecord struct { 929 929 + Text string `json:"text"` 930 930 + } 931 931 + 932 932 + rateClient := NewRateClient[testRecord](&client, "did:plc:test", nil) 933 933 + ctx := context.Background() 934 934 + 935 935 + records, _, err := rateClient.ListRecords(ctx, "app.bsky.feed.post", 10, "") 936 936 + if err != nil { 937 937 + t.Fatalf("ListRecords failed: %v", err) 938 938 + } 939 939 + if len(records) != 1 { 940 940 + t.Fatalf("ListRecords returned %d records, want 1", len(records)) 941 941 + } 942 942 + if records[0].Value.Text != "Hello world" { 943 943 + t.Errorf("Record value text = %s, want 'Hello world'", records[0].Value.Text) 944 944 + } 945 945 + }

+70

atproto/testmock.go

reviewed

··· 1 1 + package atproto 2 2 + 3 3 + import ( 4 4 + "sync" 5 5 + "time" 6 6 + ) 7 7 + 8 8 + type mockKVStore struct { 9 9 + mu sync.Mutex 10 10 + data map[string]int 11 11 + gets [][]string 12 12 + incrs []map[string]int 13 13 + } 14 14 + 15 15 + func (m *mockKVStore) GetMulti(keys []string) (map[string]int, error) { 16 16 + m.mu.Lock() 17 17 + defer m.mu.Unlock() 18 18 + m.gets = append(m.gets, keys) 19 19 + result := make(map[string]int) 20 20 + for _, k := range keys { 21 21 + result[k] = m.data[k] 22 22 + } 23 23 + return result, nil 24 24 + } 25 25 + 26 26 + func (m *mockKVStore) IncrByMulti(counts map[string]int) error { 27 27 + m.mu.Lock() 28 28 + defer m.mu.Unlock() 29 29 + m.incrs = append(m.incrs, counts) 30 30 + for k, v := range counts { 31 31 + m.data[k] += v 32 32 + } 33 33 + return nil 34 34 + } 35 35 + 36 36 + type mockClock struct { 37 37 + mu sync.Mutex 38 38 + now time.Time 39 39 + nows []time.Time 40 40 + } 41 41 + 42 42 + func (m *mockClock) Now() time.Time { 43 43 + m.mu.Lock() 44 44 + defer m.mu.Unlock() 45 45 + m.nows = append(m.nows, m.now) 46 46 + return m.now 47 47 + } 48 48 + 49 49 + type mockKVStoreWithErr struct { 50 50 + mockKVStore 51 51 + Data map[string]int 52 52 + Err error 53 53 + } 54 54 + 55 55 + func (m *mockKVStoreWithErr) GetMulti(keys []string) (map[string]int, error) { 56 56 + return nil, m.Err 57 57 + } 58 58 + 59 59 + func (m *mockKVStoreWithErr) IncrByMulti(counts map[string]int) error { 60 60 + return m.Err 61 61 + } 62 62 + 63 63 + type netError struct { 64 64 + timeout bool 65 65 + temporary bool 66 66 + } 67 67 + 68 68 + func (e *netError) Error() string { return "net error" } 69 69 + func (e *netError) Timeout() bool { return e.timeout } 70 70 + func (e *netError) Temporary() bool { return e.temporary }

-164

flags.go

reviewed

··· 1 1 - package main 2 2 - 3 3 - import ( 4 4 - "github.com/urfave/cli/v3" 5 5 - 6 6 - "tangled.org/karitham.dev/lazuli/sync" 7 7 - ) 8 8 - 9 9 - const ( 10 10 - DefaultBatchSize = 20 11 11 - ) 12 12 - 13 13 - const ( 14 14 - EnvHandle = "LAZULI_HANDLE" 15 15 - EnvPassword = "LAZULI_PASSWORD" 16 16 - EnvVerbose = "LAZULI_VERBOSE" 17 17 - EnvQuiet = "LAZULI_QUIET" 18 18 - EnvReverse = "LAZULI_REVERSE" 19 19 - EnvDryRun = "LAZULI_DRY_RUN" 20 20 - EnvFresh = "LAZULI_FRESH" 21 21 - EnvClearCache = "LAZULI_CLEAR_CACHE" 22 22 - EnvYes = "LAZULI_YES" 23 23 - ) 24 24 - 25 25 - var lastfmFlag = &cli.StringFlag{ 26 26 - Name: "lastfm", 27 27 - Usage: "Path to Last.fm CSV file or directory", 28 28 - Sources: cli.EnvVars("LAZULI_LASTFM"), 29 29 - } 30 30 - 31 31 - var spotifyFlag = &cli.StringFlag{ 32 32 - Name: "spotify", 33 33 - Usage: "Path to Spotify JSON/directory/zip", 34 34 - Sources: cli.EnvVars("LAZULI_SPOTIFY"), 35 35 - } 36 36 - 37 37 - var ( 38 38 - verboseCount int 39 39 - quietCount int 40 40 - ) 41 41 - 42 42 - var commonFlags = []cli.Flag{ 43 43 - &cli.StringFlag{ 44 44 - Name: "handle", 45 45 - Usage: "Bluesky handle", 46 46 - Sources: cli.EnvVars(EnvHandle), 47 47 - }, 48 48 - &cli.StringFlag{ 49 49 - Name: "password", 50 50 - Usage: "App password", 51 51 - Sources: cli.EnvVars(EnvPassword), 52 52 - }, 53 53 - &cli.BoolFlag{ 54 54 - Name: "verbose", 55 55 - Usage: "Enable verbose logging (-v for debug, -vv for trace)", 56 56 - Aliases: []string{"v"}, 57 57 - Sources: cli.EnvVars(EnvVerbose), 58 58 - Config: cli.BoolConfig{Count: &verboseCount}, 59 59 - }, 60 60 - &cli.BoolFlag{ 61 61 - Name: "quiet", 62 62 - Usage: "Suppress non-essential output (-q for warn, -qq for errors, -qqq for silent)", 63 63 - Aliases: []string{"q"}, 64 64 - Sources: cli.EnvVars(EnvQuiet), 65 65 - Config: cli.BoolConfig{Count: &quietCount}, 66 66 - }, 67 67 - &cli.StringFlag{ 68 68 - Name: "output-format", 69 69 - Usage: "Output format: text or json", 70 70 - Value: "text", 71 71 - Sources: cli.EnvVars("LAZULI_OUTPUT_FORMAT"), 72 72 - }, 73 73 - } 74 74 - 75 75 - var exportFlags = []cli.Flag{ 76 76 - lastfmFlag, 77 77 - spotifyFlag, 78 78 - &cli.StringFlag{ 79 79 - Name: "output", 80 80 - Usage: "Output file (stdout if not set)", 81 81 - Sources: cli.EnvVars("LAZULI_OUTPUT"), 82 82 - }, 83 83 - &cli.BoolFlag{ 84 84 - Name: "reverse", 85 85 - Usage: "Sort records reverse chronologically", 86 86 - Sources: cli.EnvVars(EnvReverse), 87 87 - }, 88 88 - &cli.DurationFlag{ 89 89 - Name: "tolerance", 90 90 - Usage: "Time tolerance for cross-source deduplication (e.g., 5m, 10m)", 91 91 - Value: sync.DefaultCrossSourceTolerance, 92 92 - Sources: cli.EnvVars("LAZULI_TOLERANCE"), 93 93 - }, 94 94 - } 95 95 - 96 96 - var importFlags = []cli.Flag{ 97 97 - lastfmFlag, 98 98 - spotifyFlag, 99 99 - &cli.StringFlag{ 100 100 - Name: "mode", 101 101 - Usage: "Import mode: lastfm, spotify, combined (default: combined)", 102 102 - Value: "combined", 103 103 - Sources: cli.EnvVars("LAZULI_MODE"), 104 104 - }, 105 105 - &cli.BoolFlag{ 106 106 - Name: "dry-run", 107 107 - Usage: "Preview without publishing", 108 108 - Sources: cli.EnvVars(EnvDryRun), 109 109 - }, 110 110 - &cli.BoolFlag{ 111 111 - Name: "reverse", 112 112 - Usage: "Import in reverse order", 113 113 - Sources: cli.EnvVars(EnvReverse), 114 114 - }, 115 115 - &cli.BoolFlag{ 116 116 - Name: "fresh", 117 117 - Usage: "Don't use cached Bluesky records", 118 118 - Sources: cli.EnvVars(EnvFresh), 119 119 - }, 120 120 - &cli.BoolFlag{ 121 121 - Name: "clear-cache", 122 122 - Usage: "Clear cache before running", 123 123 - Sources: cli.EnvVars(EnvClearCache), 124 124 - }, 125 125 - &cli.IntFlag{ 126 126 - Name: "batch-size", 127 127 - Usage: "Records per batch (default: 20)", 128 128 - Value: DefaultBatchSize, 129 129 - Sources: cli.EnvVars("LAZULI_BATCH_SIZE"), 130 130 - }, 131 131 - &cli.DurationFlag{ 132 132 - Name: "tolerance", 133 133 - Usage: "Time tolerance for cross-source deduplication (e.g., 5m, 10m)", 134 134 - Value: sync.DefaultCrossSourceTolerance, 135 135 - Sources: cli.EnvVars("LAZULI_TOLERANCE"), 136 136 - }, 137 137 - } 138 138 - 139 139 - var syncFlags = []cli.Flag{ 140 140 - &cli.BoolFlag{ 141 141 - Name: "fresh", 142 142 - Usage: "Force refresh cache", 143 143 - Sources: cli.EnvVars(EnvFresh), 144 144 - }, 145 145 - } 146 146 - 147 147 - var dedupeFlags = []cli.Flag{ 148 148 - &cli.BoolFlag{ 149 149 - Name: "dry-run", 150 150 - Usage: "Preview without deleting", 151 151 - Sources: cli.EnvVars(EnvDryRun), 152 152 - }, 153 153 - &cli.BoolFlag{ 154 154 - Name: "fresh", 155 155 - Usage: "Force refresh cache", 156 156 - Sources: cli.EnvVars(EnvFresh), 157 157 - }, 158 158 - &cli.BoolFlag{ 159 159 - Name: "yes", 160 160 - Usage: "Skip confirmation prompt", 161 161 - Aliases: []string{"y"}, 162 162 - Sources: cli.EnvVars(EnvYes), 163 163 - }, 164 164 - }

+225 -80

main.go

reviewed

··· 15 15 "tangled.org/karitham.dev/lazuli/sources/lastfm" 16 16 "tangled.org/karitham.dev/lazuli/sources/spotify" 17 17 "tangled.org/karitham.dev/lazuli/sync" 18 18 - "tangled.org/karitham.dev/lazuli/sync/logutil" 19 18 20 19 "github.com/failsafe-go/failsafe-go" 21 20 "github.com/failsafe-go/failsafe-go/retrypolicy" ··· 24 23 25 24 var Version = "dev" 26 25 26 26 + const ( 27 27 + DefaultBatchSize = 20 28 28 + ) 29 29 + 30 30 + const ( 31 31 + EnvHandle = "LAZULI_HANDLE" 32 32 + EnvPassword = "LAZULI_PASSWORD" 33 33 + EnvVerbose = "LAZULI_VERBOSE" 34 34 + EnvQuiet = "LAZULI_QUIET" 35 35 + EnvReverse = "LAZULI_REVERSE" 36 36 + EnvDryRun = "LAZULI_DRY_RUN" 37 37 + EnvFresh = "LAZULI_FRESH" 38 38 + EnvClearCache = "LAZULI_CLEAR_CACHE" 39 39 + EnvYes = "LAZULI_YES" 40 40 + ) 41 41 + 42 42 + var ( 43 43 + verboseCount int 44 44 + quietCount int 45 45 + ) 46 46 + 27 47 type App struct { 28 48 log *slog.Logger 29 49 outputFormat string ··· 40 60 } 41 61 42 62 func run() error { 43 43 - sync.UserAgent = "lazuli/" + Version 44 63 sync.ClientAgent = "lazuli/" + Version 45 64 46 65 storage, err := cache.NewBoltStorage() ··· 154 173 } 155 174 } 156 175 176 176 + func (a *App) failedCommand() *cli.Command { 177 177 + flags := make([]cli.Flag, 0, len(commonFlags)) 178 178 + flags = append(flags, commonFlags...) 179 179 + return &cli.Command{ 180 180 + Name: "failed", 181 181 + Usage: "List records that failed to publish", 182 182 + UsageText: " lazuli failed --handle=user.bsky.social", 183 183 + Flags: flags, 184 184 + Action: a.runFailed, 185 185 + Before: a.initLoggerBefore, 186 186 + } 187 187 + } 188 188 + 189 189 + func (a *App) retryCommand() *cli.Command { 190 190 + flags := make([]cli.Flag, 0, len(commonFlags)+1) 191 191 + flags = append(flags, commonFlags...) 192 192 + flags = append(flags, &cli.BoolFlag{ 193 193 + Name: "dry-run", 194 194 + Usage: "Preview what will be retried", 195 195 + Sources: cli.EnvVars(EnvDryRun), 196 196 + }) 197 197 + return &cli.Command{ 198 198 + Name: "retry", 199 199 + Usage: "Retry failed records one by one", 200 200 + UsageText: " lazuli retry --handle=user.bsky.social", 201 201 + Flags: flags, 202 202 + Action: a.runRetry, 203 203 + Before: a.initLoggerBefore, 204 204 + } 205 205 + } 206 206 + 207 207 + func (a *App) versionCommand() *cli.Command { 208 208 + return &cli.Command{ 209 209 + Name: "version", 210 210 + Usage: "Print the version number", 211 211 + Action: func(ctx context.Context, cmd *cli.Command) error { 212 212 + fmt.Println(Version) 213 213 + return nil 214 214 + }, 215 215 + } 216 216 + } 217 217 + 157 218 func (a *App) runStats(ctx context.Context, cmd *cli.Command) error { 158 219 stats, err := a.storage.Stats() 159 220 if err != nil { ··· 208 269 return nil 209 270 } 210 271 211 211 - func (a *App) failedCommand() *cli.Command { 212 212 - flags := make([]cli.Flag, 0, len(commonFlags)) 213 213 - flags = append(flags, commonFlags...) 214 214 - return &cli.Command{ 215 215 - Name: "failed", 216 216 - Usage: "List records that failed to publish", 217 217 - UsageText: " lazuli failed --handle=user.bsky.social", 218 218 - Flags: flags, 219 219 - Action: a.runFailed, 220 220 - Before: a.initLoggerBefore, 221 221 - } 222 222 - } 223 223 - 224 224 - func (a *App) retryCommand() *cli.Command { 225 225 - flags := make([]cli.Flag, 0, len(commonFlags)+1) 226 226 - flags = append(flags, commonFlags...) 227 227 - flags = append(flags, &cli.BoolFlag{ 228 228 - Name: "dry-run", 229 229 - Usage: "Preview what will be retried", 230 230 - Sources: cli.EnvVars(EnvDryRun), 231 231 - }) 232 232 - return &cli.Command{ 233 233 - Name: "retry", 234 234 - Usage: "Retry failed records one by one", 235 235 - UsageText: " lazuli retry --handle=user.bsky.social", 236 236 - Flags: flags, 237 237 - Action: a.runRetry, 238 238 - Before: a.initLoggerBefore, 239 239 - } 240 240 - } 241 241 - 242 272 func (a *App) runRetry(ctx context.Context, cmd *cli.Command) error { 243 273 authClient, err := a.prepareAuth(ctx, cmd) 244 274 if err != nil { 245 275 return err 246 276 } 247 247 - did := authClient.GetDID() 277 277 + did := authClient.DID() 248 278 dryRun := cmd.Bool("dry-run") 249 279 250 280 limiter := sync.NewRateLimiter(a.storage, 0.9) 251 251 - repoClient := sync.NewRateClient(authClient.GetAPIClient(), did, limiter) 281 281 + repoClient := sync.NewRateClient(authClient.APIClient(), did, limiter) 252 282 253 283 var failedRecords []struct { 254 284 key string ··· 287 317 continue 288 318 } 289 319 290 290 - // Check rate limit for 1 write 291 291 - 292 320 res := sync.PublishBatch(ctx, repoClient, did, []sync.PlayRecord{fr.rec}, a.storage) 293 321 294 322 if res == nil { 295 323 fmt.Printf("Successfully retried: %s - %s\n", fr.rec.ArtistName(), fr.rec.TrackName) 296 296 - // Mark as published (updates processedBucket to 1) 297 324 if err := a.storage.MarkPublished(did, fr.key); err != nil { 298 298 - a.log.Error("Failed to mark record as published", logutil.Error(err), slog.String("key", fr.key)) 325 325 + a.log.Error("Failed to mark record as published", sync.Error(err), slog.String("key", fr.key)) 299 326 } 300 300 - // Remove from failedBucket 301 327 if err := a.storage.RemoveFailed(did, fr.key); err != nil { 302 302 - a.log.Error("Failed to remove record from failed list", logutil.Error(err), slog.String("key", fr.key)) 328 328 + a.log.Error("Failed to remove record from failed list", sync.Error(err), slog.String("key", fr.key)) 303 329 } 304 330 successCount++ 305 331 } else { 306 332 fmt.Printf("Failed again: %s - %s: %v\n", fr.rec.ArtistName(), fr.rec.TrackName, res) 307 333 errorCount++ 308 334 } 309 309 - 310 310 - // Optional: small delay between retries? 311 311 - // The rate limiter already handles the delay. 312 335 } 313 336 314 337 fmt.Printf("\nRetry complete: %d succeeded, %d failed.\n", successCount, errorCount) ··· 316 339 } 317 340 318 341 func (a *App) runFailed(ctx context.Context, cmd *cli.Command) error { 319 319 - // We need the DID to look up the user's failed records. 320 320 - // If it's already in the handle, use it, otherwise we might need to resolve it. 321 321 - // For simplicity, we'll try to get it from auth or resolve it. 322 342 authClient, err := a.prepareAuth(ctx, cmd) 323 343 if err != nil { 324 344 return err 325 345 } 326 326 - did := authClient.GetDID() 346 346 + did := authClient.DID() 327 347 328 348 type FailedRecord struct { 329 349 Key string `json:"key"` ··· 369 389 return nil 370 390 } 371 391 372 372 - func (a *App) versionCommand() *cli.Command { 373 373 - return &cli.Command{ 374 374 - Name: "version", 375 375 - Usage: "Print the version number", 376 376 - Action: func(ctx context.Context, cmd *cli.Command) error { 377 377 - fmt.Println(Version) 378 378 - return nil 379 379 - }, 380 380 - } 381 381 - } 382 382 - 383 392 func (a *App) runDebugFetch(ctx context.Context, cmd *cli.Command) error { 384 393 authClient, err := a.prepareAuth(ctx, cmd) 385 394 if err != nil { 386 395 return fmt.Errorf("authentication failed: %w", err) 387 396 } 388 397 389 389 - repoClient := sync.NewRateClient(authClient.GetAPIClient(), authClient.GetDID(), nil) 398 398 + repoClient := sync.NewRateClient(authClient.APIClient(), authClient.DID(), nil) 390 399 391 400 records, _, err := repoClient.ListRecords(ctx, sync.RecordType, 10, "") 392 401 if err != nil { ··· 495 504 return err 496 505 } 497 506 498 498 - a.log.Info("Starting import operation", logutil.DID(handle)) 507 507 + a.log.Info("Starting import operation", sync.DID(handle)) 499 508 500 509 lastfmPath := cmd.String("lastfm") 501 510 spotifyPath := cmd.String("spotify") ··· 509 518 510 519 if clearCache { 511 520 if err := a.storage.ClearAll(); err != nil { 512 512 - a.log.Error("Failed to clear cache", logutil.Error(err)) 521 521 + a.log.Error("Failed to clear cache", sync.Error(err)) 513 522 } else { 514 523 a.log.Info("Cache cleared") 515 524 } ··· 549 558 if err != nil { 550 559 return fmt.Errorf("create auth client: %w", err) 551 560 } 552 552 - a.log.Info("Authenticated", logutil.DID(authClient.GetDID()), slog.String("pds", authClient.GetPDS())) 561 561 + a.log.Info("Authenticated", sync.DID(authClient.DID()), slog.String("pds", authClient.PDS())) 553 562 554 563 limiter := sync.NewRateLimiter(a.storage, 0.9) 555 555 - repoClient := sync.NewRateClient(authClient.GetAPIClient(), authClient.GetDID(), limiter) 564 564 + repoClient := sync.NewRateClient(authClient.APIClient(), authClient.DID(), limiter) 556 565 557 557 - existingRecords, err := sync.FetchExisting(ctx, repoClient, authClient.GetDID(), a.storage, fresh) 566 566 + existingRecords, err := sync.FetchExisting(ctx, repoClient, authClient.DID(), a.storage, fresh) 558 567 if err != nil { 559 568 return fmt.Errorf("fetch existing records: %w", err) 560 569 } 561 570 a.log.Info("Fetched existing records", slog.Int("count", len(existingRecords))) 562 571 563 563 - published, _ := a.storage.GetPublished(authClient.GetDID()) 572 572 + published, _ := a.storage.GetPublished(authClient.DID()) 564 573 newRecords := sync.FilterNew(records, existingRecords, published) 565 574 skippedCount := len(records) - len(newRecords) 566 575 a.log.Info("Filtered to new records", ··· 584 593 value, _ := json.Marshal(rec) 585 594 newEntries[key] = value 586 595 } 587 587 - if err := a.storage.SaveRecords(authClient.GetDID(), newEntries); err != nil { 596 596 + if err := a.storage.SaveRecords(authClient.DID(), newEntries); err != nil { 588 597 return fmt.Errorf("save new records to storage: %w", err) 589 598 } 590 599 } ··· 663 672 } 664 673 665 674 fresh := cmd.Bool("fresh") 666 666 - a.log.Info("Starting sync operation", logutil.DID(authClient.GetDID()), slog.Bool("fresh", fresh)) 675 675 + a.log.Info("Starting sync operation", sync.DID(authClient.DID()), slog.Bool("fresh", fresh)) 667 676 668 677 limiter := sync.NewRateLimiter(a.storage, 0.85) 669 669 - repoClient := sync.NewRateClient(authClient.GetAPIClient(), authClient.GetDID(), limiter) 678 678 + repoClient := sync.NewRateClient(authClient.APIClient(), authClient.DID(), limiter) 670 679 671 680 if fresh { 672 672 - if err := a.storage.Clear(authClient.GetDID()); err != nil { 673 673 - a.log.Error("Failed to clear cache", logutil.Error(err)) 681 681 + if err := a.storage.Clear(authClient.DID()); err != nil { 682 682 + a.log.Error("Failed to clear cache", sync.Error(err)) 674 683 } else { 675 684 a.log.Info("Cache cleared") 676 685 } 677 686 } 678 687 679 679 - existingRecords, err := sync.FetchExisting(ctx, repoClient, authClient.GetDID(), a.storage, fresh) 688 688 + existingRecords, err := sync.FetchExisting(ctx, repoClient, authClient.DID(), a.storage, fresh) 680 689 if err != nil { 681 690 return fmt.Errorf("fetch existing records: %w", err) 682 691 } ··· 696 705 fresh := cmd.Bool("fresh") 697 706 yes := cmd.Bool("yes") 698 707 a.log.Info("Starting dedupe operation", 699 699 - logutil.DID(authClient.GetDID()), 708 708 + sync.DID(authClient.DID()), 700 709 slog.Bool("dry_run", dryRun), 701 710 slog.Bool("fresh", fresh)) 702 711 703 712 limiter := sync.NewRateLimiter(a.storage, 0.9) 704 704 - repoClient := sync.NewRateClient(authClient.GetAPIClient(), authClient.GetDID(), limiter) 713 713 + repoClient := sync.NewRateClient(authClient.APIClient(), authClient.DID(), limiter) 705 714 706 715 if fresh { 707 707 - if err := a.storage.Clear(authClient.GetDID()); err != nil { 708 708 - a.log.Error("Failed to clear cache", logutil.Error(err)) 716 716 + if err := a.storage.Clear(authClient.DID()); err != nil { 717 717 + a.log.Error("Failed to clear cache", sync.Error(err)) 709 718 } else { 710 719 a.log.Info("Cache cleared") 711 720 } 712 721 } 713 722 714 714 - existingRecords, err := sync.FetchExisting(ctx, repoClient, authClient.GetDID(), a.storage, fresh) 723 723 + existingRecords, err := sync.FetchExisting(ctx, repoClient, authClient.DID(), a.storage, fresh) 715 724 if err != nil { 716 725 return fmt.Errorf("failed to fetch existing records: %w", err) 717 726 } ··· 773 782 OnRetryScheduled(func(e failsafe.ExecutionScheduledEvent[any]) { 774 783 a.log.Warn("Delete failed with transient error, retrying", 775 784 slog.Duration("retryDelay", e.Delay), 776 776 - logutil.Error(e.LastError()), 785 785 + sync.Error(e.LastError()), 777 786 slog.Int("attempt", e.Attempts()), 778 787 slog.String("uri", uri)) 779 788 }). ··· 784 793 }) 785 794 786 795 if err != nil { 787 787 - a.log.Error("Failed to delete record", logutil.Error(err), slog.String("uri", uri)) 796 796 + a.log.Error("Failed to delete record", sync.Error(err), slog.String("uri", uri)) 788 797 } else { 789 798 a.log.Info("Deleted duplicate", slog.String("uri", uri), slog.String("track", rec.Value.TrackName)) 790 799 } 791 800 } 792 801 } 793 802 794 794 - if err := a.storage.Clear(authClient.GetDID()); err != nil { 803 803 + if err := a.storage.Clear(authClient.DID()); err != nil { 795 804 a.log.Error("Failed to clear cache", "err", err) 796 805 } 797 806 ··· 818 827 819 828 return nil 820 829 } 830 830 + 831 831 + var commonFlags = []cli.Flag{ 832 832 + &cli.StringFlag{ 833 833 + Name: "handle", 834 834 + Usage: "Bluesky handle", 835 835 + Sources: cli.EnvVars(EnvHandle), 836 836 + }, 837 837 + &cli.StringFlag{ 838 838 + Name: "password", 839 839 + Usage: "App password", 840 840 + Sources: cli.EnvVars(EnvPassword), 841 841 + }, 842 842 + &cli.BoolFlag{ 843 843 + Name: "verbose", 844 844 + Usage: "Enable verbose logging (-v for debug, -vv for trace)", 845 845 + Aliases: []string{"v"}, 846 846 + Sources: cli.EnvVars(EnvVerbose), 847 847 + Config: cli.BoolConfig{Count: &verboseCount}, 848 848 + }, 849 849 + &cli.BoolFlag{ 850 850 + Name: "quiet", 851 851 + Usage: "Suppress non-essential output (-q for warn, -qq for errors, -qqq for silent)", 852 852 + Aliases: []string{"q"}, 853 853 + Sources: cli.EnvVars(EnvQuiet), 854 854 + Config: cli.BoolConfig{Count: &quietCount}, 855 855 + }, 856 856 + &cli.StringFlag{ 857 857 + Name: "output-format", 858 858 + Usage: "Output format: text or json", 859 859 + Value: "text", 860 860 + Sources: cli.EnvVars("LAZULI_OUTPUT_FORMAT"), 861 861 + }, 862 862 + } 863 863 + 864 864 + var lastfmFlag = &cli.StringFlag{ 865 865 + Name: "lastfm", 866 866 + Usage: "Path to Last.fm CSV file or directory", 867 867 + Sources: cli.EnvVars("LAZULI_LASTFM"), 868 868 + } 869 869 + 870 870 + var spotifyFlag = &cli.StringFlag{ 871 871 + Name: "spotify", 872 872 + Usage: "Path to Spotify JSON/directory/zip", 873 873 + Sources: cli.EnvVars("LAZULI_SPOTIFY"), 874 874 + } 875 875 + 876 876 + var exportFlags = []cli.Flag{ 877 877 + lastfmFlag, 878 878 + spotifyFlag, 879 879 + &cli.StringFlag{ 880 880 + Name: "output", 881 881 + Usage: "Output file (stdout if not set)", 882 882 + Sources: cli.EnvVars("LAZULI_OUTPUT"), 883 883 + }, 884 884 + &cli.BoolFlag{ 885 885 + Name: "reverse", 886 886 + Usage: "Sort records reverse chronologically", 887 887 + Sources: cli.EnvVars(EnvReverse), 888 888 + }, 889 889 + &cli.DurationFlag{ 890 890 + Name: "tolerance", 891 891 + Usage: "Time tolerance for cross-source deduplication (e.g., 5m, 10m)", 892 892 + Value: sync.DefaultCrossSourceTolerance, 893 893 + Sources: cli.EnvVars("LAZULI_TOLERANCE"), 894 894 + }, 895 895 + } 896 896 + 897 897 + var importFlags = []cli.Flag{ 898 898 + lastfmFlag, 899 899 + spotifyFlag, 900 900 + &cli.StringFlag{ 901 901 + Name: "mode", 902 902 + Usage: "Import mode: lastfm, spotify, combined (default: combined)", 903 903 + Value: "combined", 904 904 + Sources: cli.EnvVars("LAZULI_MODE"), 905 905 + }, 906 906 + &cli.BoolFlag{ 907 907 + Name: "dry-run", 908 908 + Usage: "Preview without publishing", 909 909 + Sources: cli.EnvVars(EnvDryRun), 910 910 + }, 911 911 + &cli.BoolFlag{ 912 912 + Name: "reverse", 913 913 + Usage: "Import in reverse order", 914 914 + Sources: cli.EnvVars(EnvReverse), 915 915 + }, 916 916 + &cli.BoolFlag{ 917 917 + Name: "fresh", 918 918 + Usage: "Don't use cached Bluesky records", 919 919 + Sources: cli.EnvVars(EnvFresh), 920 920 + }, 921 921 + &cli.BoolFlag{ 922 922 + Name: "clear-cache", 923 923 + Usage: "Clear cache before running", 924 924 + Sources: cli.EnvVars(EnvClearCache), 925 925 + }, 926 926 + &cli.IntFlag{ 927 927 + Name: "batch-size", 928 928 + Usage: "Records per batch (default: 20)", 929 929 + Value: DefaultBatchSize, 930 930 + Sources: cli.EnvVars("LAZULI_BATCH_SIZE"), 931 931 + }, 932 932 + &cli.DurationFlag{ 933 933 + Name: "tolerance", 934 934 + Usage: "Time tolerance for cross-source deduplication (e.g., 5m, 10m)", 935 935 + Value: sync.DefaultCrossSourceTolerance, 936 936 + Sources: cli.EnvVars("LAZULI_TOLERANCE"), 937 937 + }, 938 938 + } 939 939 + 940 940 + var syncFlags = []cli.Flag{ 941 941 + &cli.BoolFlag{ 942 942 + Name: "fresh", 943 943 + Usage: "Force refresh cache", 944 944 + Sources: cli.EnvVars(EnvFresh), 945 945 + }, 946 946 + } 947 947 + 948 948 + var dedupeFlags = []cli.Flag{ 949 949 + &cli.BoolFlag{ 950 950 + Name: "dry-run", 951 951 + Usage: "Preview without deleting", 952 952 + Sources: cli.EnvVars(EnvDryRun), 953 953 + }, 954 954 + &cli.BoolFlag{ 955 955 + Name: "fresh", 956 956 + Usage: "Force refresh cache", 957 957 + Sources: cli.EnvVars(EnvFresh), 958 958 + }, 959 959 + &cli.BoolFlag{ 960 960 + Name: "yes", 961 961 + Usage: "Skip confirmation prompt", 962 962 + Aliases: []string{"y"}, 963 963 + Sources: cli.EnvVars(EnvYes), 964 964 + }, 965 965 + }

-486

sync/adapter.go

reviewed

··· 1 1 - package sync 2 2 - 3 3 - import ( 4 4 - "context" 5 5 - "encoding/json" 6 6 - "errors" 7 7 - "fmt" 8 8 - "log/slog" 9 9 - "net" 10 10 - "strings" 11 11 - "time" 12 12 - 13 13 - "github.com/bluesky-social/indigo/atproto/atclient" 14 14 - "github.com/bluesky-social/indigo/atproto/syntax" 15 15 - "github.com/failsafe-go/failsafe-go" 16 16 - "github.com/failsafe-go/failsafe-go/retrypolicy" 17 17 - 18 18 - "tangled.org/karitham.dev/lazuli/cache" 19 19 - "tangled.org/karitham.dev/lazuli/sync/logutil" 20 20 - ) 21 21 - 22 22 - type RepoClient interface { 23 23 - ListRecords(ctx context.Context, collection string, limit int, cursor string) ([]RecordRef, string, error) 24 24 - ApplyWrites(ctx context.Context, collection string, records []PlayRecord) error 25 25 - DeleteRecord(ctx context.Context, collection, rkey string) error 26 26 - } 27 27 - 28 28 - type RecordRef struct { 29 29 - URI string 30 30 - CID string 31 31 - Value PlayRecord 32 32 - } 33 33 - 34 34 - type RateClient struct { 35 35 - client *atclient.APIClient 36 36 - did string 37 37 - limiter RateLimiter 38 38 - } 39 39 - 40 40 - func NewRateClient(client *atclient.APIClient, did string, limiter RateLimiter) *RateClient { 41 41 - return &RateClient{ 42 42 - client: client, 43 43 - did: did, 44 44 - limiter: limiter, 45 45 - } 46 46 - } 47 47 - 48 48 - func (c *RateClient) ListRecords(ctx context.Context, collection string, limit int, cursor string) ([]RecordRef, string, error) { 49 49 - if c.client == nil { 50 50 - return nil, "", fmt.Errorf("client cannot be nil") 51 51 - } 52 52 - 53 53 - var outResp struct { 54 54 - Records []struct { 55 55 - URI string `json:"uri"` 56 56 - CID string `json:"cid"` 57 57 - Value map[string]any `json:"value"` 58 58 - } `json:"records"` 59 59 - Cursor string `json:"cursor"` 60 60 - } 61 61 - 62 62 - var chargedAt time.Time 63 63 - if c.limiter != nil { 64 64 - slog.Debug("waiting for rate limit (read)") 65 65 - var err error 66 66 - chargedAt, err = c.limiter.AllowRead(ctx) 67 67 - if err != nil { 68 68 - slog.Error("rate limit wait cancelled/failed (read)", logutil.Error(err)) 69 69 - return nil, "", err 70 70 - } 71 71 - } 72 72 - 73 73 - err := c.client.Get(ctx, syntax.NSID("com.atproto.repo.listRecords"), map[string]any{ 74 74 - "repo": c.did, 75 75 - "collection": collection, 76 76 - "limit": limit, 77 77 - "cursor": cursor, 78 78 - }, &outResp) 79 79 - if err != nil { 80 80 - if c.limiter != nil && isTransientError(err) { 81 81 - c.limiter.RefundRead(ctx, chargedAt) 82 82 - } 83 83 - return nil, "", err 84 84 - } 85 85 - 86 86 - out := make([]RecordRef, 0, len(outResp.Records)) 87 87 - for _, r := range outResp.Records { 88 88 - var playRecord PlayRecord 89 89 - if r.Value != nil { 90 90 - b, err := json.Marshal(r.Value) 91 91 - if err != nil { 92 92 - slog.Debug("failed to marshal record value", slog.String("uri", r.URI), logutil.Error(err)) 93 93 - continue 94 94 - } 95 95 - if err := json.Unmarshal(b, &playRecord); err != nil { 96 96 - slog.Debug("failed to unmarshal record", slog.String("uri", r.URI), logutil.Error(err)) 97 97 - continue 98 98 - } 99 99 - slog.Debug("parsed record", slog.String("uri", r.URI), logutil.Track(playRecord.TrackName, playRecord.ArtistName(), playRecord.PlayedTime.Time)) 100 100 - } 101 101 - out = append(out, RecordRef{ 102 102 - URI: r.URI, 103 103 - CID: r.CID, 104 104 - Value: playRecord, 105 105 - }) 106 106 - } 107 107 - 108 108 - return out, outResp.Cursor, nil 109 109 - } 110 110 - 111 111 - func (c *RateClient) ApplyWrites(ctx context.Context, collection string, records []PlayRecord) error { 112 112 - if len(records) == 0 { 113 113 - return nil 114 114 - } 115 115 - 116 116 - if c.client == nil { 117 117 - return fmt.Errorf("client cannot be nil") 118 118 - } 119 119 - 120 120 - var chargedAt time.Time 121 121 - if c.limiter != nil { 122 122 - var err error 123 123 - chargedAt, err = c.limiter.AllowBulkWrite(ctx, len(records)) 124 124 - if err != nil { 125 125 - slog.Error("rate limit wait cancelled/failed (write)", 126 126 - logutil.DID(c.did), 127 127 - slog.String("collection", collection), 128 128 - logutil.Error(err), 129 129 - ) 130 130 - return err 131 131 - } 132 132 - } 133 133 - 134 134 - err := applyWrites(ctx, c.client, c.did, collection, records) 135 135 - if err != nil && isTransientError(err) && c.limiter != nil { 136 136 - c.limiter.RefundBulkWrite(ctx, len(records), chargedAt) 137 137 - } 138 138 - return err 139 139 - } 140 140 - 141 141 - func applyWrites(ctx context.Context, client *atclient.APIClient, did, collection string, records []PlayRecord) error { 142 142 - if len(records) == 0 { 143 143 - return nil 144 144 - } 145 145 - 146 146 - if len(records) > 200 { 147 147 - return fmt.Errorf("too many records in one ApplyWrites call: %d (max 200)", len(records)) 148 148 - } 149 149 - 150 150 - writes, err := prepareWrites(records, collection) 151 151 - if err != nil { 152 152 - return err 153 153 - } 154 154 - 155 155 - return client.Post(ctx, syntax.NSID("com.atproto.repo.applyWrites"), map[string]any{ 156 156 - "repo": did, 157 157 - "writes": writes, 158 158 - }, nil) 159 159 - } 160 160 - 161 161 - type atprotoClientAdapter struct { 162 162 - client *atclient.APIClient 163 163 - did string 164 164 - } 165 165 - 166 166 - func (a *atprotoClientAdapter) ApplyWrites(ctx context.Context, collection string, records []PlayRecord) error { 167 167 - return applyWrites(ctx, a.client, a.did, collection, records) 168 168 - } 169 169 - 170 170 - func prepareRecords(batch []PlayRecord) []PlayRecord { 171 171 - atprotoRecords := make([]PlayRecord, 0, len(batch)) 172 172 - for _, record := range batch { 173 173 - record.Type = RecordType 174 174 - record.SubmissionClientAgent = ClientAgent 175 175 - atprotoRecords = append(atprotoRecords, record) 176 176 - } 177 177 - return atprotoRecords 178 178 - } 179 179 - 180 180 - func defaultProgressLog(f func(ProgressReport)) func(ProgressReport) { 181 181 - if f != nil { 182 182 - return f 183 183 - } 184 184 - return func(pr ProgressReport) { 185 185 - slog.Info("sync progress", 186 186 - slog.Int("completed", pr.Completed), 187 187 - slog.Int("total", pr.Total), 188 188 - slog.Float64("percent", pr.Percent), 189 189 - slog.String("elapsed", pr.Elapsed), 190 190 - slog.String("eta", pr.ETA), 191 191 - slog.String("rate", pr.Rate), 192 192 - slog.Int("errors", pr.Errors), 193 193 - ) 194 194 - } 195 195 - } 196 196 - 197 197 - func defaultBatchSize(size int) int { 198 198 - if size > 0 { 199 199 - return size 200 200 - } 201 201 - return DefaultBatchSize 202 202 - } 203 203 - 204 204 - func buildClient(client AuthClient, customClient ATProtoClient) (ATProtoClient, error) { 205 205 - if customClient != nil { 206 206 - return customClient, nil 207 207 - } 208 208 - 209 209 - apiClient := client.GetAPIClient() 210 210 - if apiClient == nil { 211 211 - slog.Error("failed to get API client", logutil.Error(fmt.Errorf("client is nil"))) 212 212 - return nil, fmt.Errorf("API client is nil") 213 213 - } 214 214 - 215 215 - return &atprotoClientAdapter{client: apiClient, did: client.GetDID()}, nil 216 216 - } 217 217 - 218 218 - func newPublishResult(success, errors, total int, start time.Time, cancelled bool) PublishResult { 219 219 - return PublishResult{ 220 220 - SuccessCount: success, 221 221 - ErrorCount: errors, 222 222 - Cancelled: cancelled, 223 223 - Duration: time.Since(start), 224 224 - TotalRecords: total, 225 225 - RecordsPerMinute: ratePerMinute(success, time.Since(start)), 226 226 - } 227 227 - } 228 228 - 229 229 - func logResult(success, errors int, startTime time.Time) { 230 230 - if errors > 0 { 231 231 - slog.Warn("import completed with errors", 232 232 - slog.Int("success", success), 233 233 - slog.Int("errors", errors)) 234 234 - } 235 235 - slog.Info("import completed", 236 236 - slog.Int("success", success), 237 237 - slog.Int("errors", errors), 238 238 - slog.Duration("duration", time.Since(startTime)), 239 239 - slog.String("rate", formatRate(ratePerMinute(success, time.Since(startTime))))) 240 240 - } 241 241 - 242 242 - func PublishBatch(ctx context.Context, client ATProtoClient, did string, batch []PlayRecord, storage cache.Storage) error { 243 243 - if len(batch) == 0 { 244 244 - return nil 245 245 - } 246 246 - 247 247 - atprotoRecords := prepareRecords(batch) 248 248 - 249 249 - err := client.ApplyWrites(ctx, RecordType, atprotoRecords) 250 250 - if err != nil { 251 251 - slog.Error("batch publish failed", logutil.Error(err)) 252 252 - return err 253 253 - } 254 254 - 255 255 - if storage != nil && did != "" { 256 256 - keys := CreateRecordKeys(atprotoRecords) 257 257 - cacheEntries := make(map[string][]byte) 258 258 - for i, rec := range atprotoRecords { 259 259 - key := keys[i] 260 260 - value, _ := json.Marshal(rec) 261 261 - cacheEntries[key] = value 262 262 - } 263 263 - 264 264 - if err := storage.SaveRecords(did, cacheEntries); err != nil { 265 265 - return fmt.Errorf("failed to save records to storage: %w", err) 266 266 - } 267 267 - 268 268 - if err := storage.MarkPublished(did, keys...); err != nil { 269 269 - return fmt.Errorf("failed to mark records as published: %w", err) 270 270 - } 271 271 - } 272 272 - 273 273 - return nil 274 274 - } 275 275 - 276 276 - type ATProtoClient interface { 277 277 - ApplyWrites(ctx context.Context, collection string, records []PlayRecord) error 278 278 - } 279 279 - 280 280 - func ratePerMinute(count int, duration time.Duration) float64 { 281 281 - if duration == 0 { 282 282 - return 0 283 283 - } 284 284 - return float64(count) / duration.Minutes() 285 285 - } 286 286 - 287 287 - type AuthClient interface { 288 288 - GetAPIClient() *atclient.APIClient 289 289 - GetDID() string 290 290 - } 291 291 - 292 292 - func IsTransientError(err error) bool { 293 293 - return isTransientError(err) 294 294 - } 295 295 - 296 296 - func isTransientError(err error) bool { 297 297 - if err == nil { 298 298 - return false 299 299 - } 300 300 - 301 301 - var apiErr *atclient.APIError 302 302 - if errors.As(err, &apiErr) { 303 303 - switch apiErr.StatusCode { 304 304 - case 429, 500, 502, 503, 504: 305 305 - return true 306 306 - } 307 307 - return false 308 308 - } 309 309 - 310 310 - var netErr net.Error 311 311 - if errors.As(err, &netErr) { 312 312 - return netErr.Timeout() 313 313 - } 314 314 - 315 315 - return false 316 316 - } 317 317 - 318 318 - func (c *RateClient) DeleteRecord(ctx context.Context, collection, rkey string) error { 319 319 - if c.client == nil { 320 320 - return fmt.Errorf("client is nil") 321 321 - } 322 322 - 323 323 - var chargedAt time.Time 324 324 - if c.limiter != nil { 325 325 - slog.Debug("waiting for rate limit (delete)") 326 326 - var err error 327 327 - chargedAt, err = c.limiter.AllowBulkWrite(ctx, 1) 328 328 - if err != nil { 329 329 - slog.Error("rate limit wait cancelled/failed (delete)", logutil.Error(err)) 330 330 - return err 331 331 - } 332 332 - } 333 333 - 334 334 - _, err := c.client.Do(ctx, &atclient.APIRequest{ 335 335 - Method: "DELETE", 336 336 - Endpoint: syntax.NSID("com.atproto.repo.deleteRecord"), 337 337 - QueryParams: map[string][]string{ 338 338 - "repo": {c.did}, 339 339 - "collection": {collection}, 340 340 - "rkey": {rkey}, 341 341 - }, 342 342 - }) 343 343 - if err != nil && c.limiter != nil && isTransientError(err) { 344 344 - c.limiter.RefundBulkWrite(ctx, 1, chargedAt) 345 345 - } 346 346 - return err 347 347 - } 348 348 - 349 349 - func FetchExisting(ctx context.Context, client RepoClient, did string, storage cache.Storage, forceRefresh bool) ([]ExistingRecord, error) { 350 350 - if !forceRefresh && storage != nil { 351 351 - published, err := storage.GetPublished(did) 352 352 - if err == nil && len(published) > 0 && storage.IsValid(did) { 353 353 - records := make([]ExistingRecord, 0, len(published)) 354 354 - err := storage.IteratePublished(did, func(key string, data []byte) error { 355 355 - var value PlayRecord 356 356 - if err := json.Unmarshal(data, &value); err != nil { 357 357 - return nil 358 358 - } 359 359 - records = append(records, ExistingRecord{ 360 360 - URI: generateRecordURI(did, value), 361 361 - Value: value, 362 362 - }) 363 363 - return nil 364 364 - }) 365 365 - if err == nil { 366 366 - slog.Debug("loaded from cache", slog.Int("count", len(records))) 367 367 - return records, nil 368 368 - } 369 369 - } 370 370 - } 371 371 - 372 372 - select { 373 373 - case <-ctx.Done(): 374 374 - return nil, ctx.Err() 375 375 - default: 376 376 - } 377 377 - 378 378 - allRecords := make([]ExistingRecord, 0, 1024) 379 379 - return fetchExistingLoop(ctx, client, did, storage, allRecords) 380 380 - } 381 381 - 382 382 - func fetchExistingLoop(ctx context.Context, client RepoClient, did string, storage cache.Storage, allRecords []ExistingRecord) ([]ExistingRecord, error) { 383 383 - const batchSize = 100 384 384 - var cursor string 385 385 - 386 386 - type fetchResult struct { 387 387 - records []RecordRef 388 388 - cursor string 389 389 - } 390 390 - 391 391 - retryPolicy := retrypolicy.NewBuilder[fetchResult](). 392 392 - WithMaxRetries(10). 393 393 - WithBackoff(BaseRetryDelay, 5*time.Minute). 394 394 - HandleIf(func(_ fetchResult, err error) bool { 395 395 - return isTransientError(err) 396 396 - }). 397 397 - OnRetryScheduled(func(e failsafe.ExecutionScheduledEvent[fetchResult]) { 398 398 - slog.Warn("fetch failed with transient error, retrying", 399 399 - slog.Duration("retryDelay", e.Delay), 400 400 - logutil.Error(e.LastError()), 401 401 - slog.Int("attempt", e.Attempts())) 402 402 - }). 403 403 - Build() 404 404 - 405 405 - for { 406 406 - select { 407 407 - case <-ctx.Done(): 408 408 - return nil, ctx.Err() 409 409 - default: 410 410 - } 411 411 - 412 412 - result, err := failsafe.With(retryPolicy). 413 413 - WithContext(ctx). 414 414 - Get(func() (fetchResult, error) { 415 415 - recs, next, err := client.ListRecords(ctx, RecordType, batchSize, cursor) 416 416 - if err != nil { 417 417 - return fetchResult{}, err 418 418 - } 419 419 - 420 420 - return fetchResult{records: recs, cursor: next}, nil 421 421 - }) 422 422 - if err != nil { 423 423 - return nil, err 424 424 - } 425 425 - 426 426 - for _, rec := range result.records { 427 427 - allRecords = append(allRecords, ExistingRecord(rec)) 428 428 - } 429 429 - 430 430 - if result.cursor == "" || len(result.records) < batchSize { 431 431 - break 432 432 - } 433 433 - cursor = result.cursor 434 434 - } 435 435 - 436 436 - if storage != nil { 437 437 - cacheEntries := make(map[string][]byte) 438 438 - keys := make([]string, 0, len(allRecords)) 439 439 - for _, rec := range allRecords { 440 440 - parts := strings.Split(rec.URI, "/") 441 441 - key := parts[len(parts)-1] 442 442 - if key == "" { 443 443 - key = CreateRecordKey(rec.Value) 444 444 - } 445 445 - value, _ := json.Marshal(rec.Value) 446 446 - cacheEntries[key] = value 447 447 - keys = append(keys, key) 448 448 - } 449 449 - 450 450 - if err := storage.SaveRecords(did, cacheEntries); err != nil { 451 451 - return nil, err 452 452 - } 453 453 - 454 454 - if err := storage.MarkPublished(did, keys...); err != nil { 455 455 - return nil, err 456 456 - } 457 457 - 458 458 - slog.Debug("saved to cache and marked as published", slog.Int("count", len(allRecords))) 459 459 - } 460 460 - 461 461 - return allRecords, nil 462 462 - } 463 463 - 464 464 - func prepareWrites(records []PlayRecord, collection string) ([]map[string]any, error) { 465 465 - if len(records) == 0 { 466 466 - return nil, nil 467 467 - } 468 468 - 469 469 - writes := make([]map[string]any, len(records)) 470 470 - keys := CreateRecordKeys(records) 471 471 - 472 472 - for i, rec := range records { 473 473 - writes[i] = map[string]any{ 474 474 - "$type": "com.atproto.repo.applyWrites#create", 475 475 - "collection": collection, 476 476 - "rkey": keys[i], 477 477 - "value": rec, 478 478 - } 479 479 - } 480 480 - 481 481 - return writes, nil 482 482 - } 483 483 - 484 484 - func generateRecordURI(did string, record PlayRecord) string { 485 485 - return fmt.Sprintf("at://%s/%s/%s", did, RecordType, CreateRecordKey(record)) 486 486 - }

-60

sync/atproto.go

reviewed

··· 1 1 - package sync 2 2 - 3 3 - import ( 4 4 - "context" 5 5 - "encoding/json" 6 6 - "fmt" 7 7 - "net/http" 8 8 - "net/url" 9 9 - "time" 10 10 - ) 11 11 - 12 12 - var UserAgent string 13 13 - 14 14 - var httpClient = &http.Client{ 15 15 - Timeout: 10 * time.Second, 16 16 - } 17 17 - 18 18 - func ResolveMiniDoc(ctx context.Context, identifier string) (did, pds, signingKey string, err error) { 19 19 - parsedURL, err := url.Parse(SlingshotResolverURL) 20 20 - if err != nil { 21 21 - return "", "", "", fmt.Errorf("invalid resolver URL: %w", err) 22 22 - } 23 23 - 24 24 - query := parsedURL.Query() 25 25 - query.Set("identifier", identifier) 26 26 - parsedURL.RawQuery = query.Encode() 27 27 - 28 28 - req, err := http.NewRequestWithContext(ctx, "GET", parsedURL.String(), nil) 29 29 - if err != nil { 30 30 - return "", "", "", fmt.Errorf("failed to create request: %w", err) 31 31 - } 32 32 - if UserAgent != "" { 33 33 - req.Header.Set("User-Agent", UserAgent) 34 34 - } 35 35 - 36 36 - resp, err := httpClient.Do(req) 37 37 - if err != nil { 38 38 - return "", "", "", fmt.Errorf("failed to resolve mini doc: %w", err) 39 39 - } 40 40 - defer resp.Body.Close() 41 41 - 42 42 - if resp.StatusCode != http.StatusOK { 43 43 - return "", "", "", fmt.Errorf("mini doc resolution failed with status: %d", resp.StatusCode) 44 44 - } 45 45 - 46 46 - var result struct { 47 47 - DID string `json:"did"` 48 48 - PDS string `json:"pds"` 49 49 - SigningKey string `json:"signing_key"` 50 50 - } 51 51 - if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { 52 52 - return "", "", "", fmt.Errorf("failed to decode mini doc response: %w", err) 53 53 - } 54 54 - 55 55 - if result.DID == "" { 56 56 - return "", "", "", fmt.Errorf("resolved mini doc missing DID") 57 57 - } 58 58 - 59 59 - return result.DID, result.PDS, result.SigningKey, nil 60 60 - }

-105

sync/atproto_auth.go

reviewed

··· 1 1 - package sync 2 2 - 3 3 - import ( 4 4 - "context" 5 5 - "encoding/json" 6 6 - "fmt" 7 7 - "net/http" 8 8 - "strings" 9 9 - "sync" 10 10 - 11 11 - "github.com/bluesky-social/indigo/atproto/atclient" 12 12 - "github.com/bluesky-social/indigo/atproto/syntax" 13 13 - ) 14 14 - 15 15 - type FixedPasswordAuth struct { 16 16 - *atclient.PasswordAuth 17 17 - lk sync.RWMutex 18 18 - } 19 19 - 20 20 - func (a *FixedPasswordAuth) DoWithAuth(c *http.Client, req *http.Request, endpoint syntax.NSID) (*http.Response, error) { 21 21 - accessToken, refreshToken := a.GetTokens() 22 22 - req.Header.Set("Authorization", "Bearer "+accessToken) 23 23 - resp, err := c.Do(req) 24 24 - if err != nil { 25 25 - return nil, err 26 26 - } 27 27 - 28 28 - if resp.StatusCode != http.StatusBadRequest || !strings.HasPrefix(resp.Header.Get("Content-Type"), "application/json") { 29 29 - return resp, nil 30 30 - } 31 31 - 32 32 - defer resp.Body.Close() 33 33 - var eb atclient.ErrorBody 34 34 - if err := json.NewDecoder(resp.Body).Decode(&eb); err != nil { 35 35 - return nil, &atclient.APIError{StatusCode: resp.StatusCode} 36 36 - } 37 37 - if eb.Name != "ExpiredToken" { 38 38 - return nil, eb.APIError(resp.StatusCode) 39 39 - } 40 40 - 41 41 - if err := a.Refresh(req.Context(), c, refreshToken); err != nil { 42 42 - return nil, err 43 43 - } 44 44 - 45 45 - retry := req.Clone(req.Context()) 46 46 - if req.GetBody != nil { 47 47 - retry.Body, err = req.GetBody() 48 48 - if err != nil { 49 49 - return nil, fmt.Errorf("API request retry GetBody failed: %w", err) 50 50 - } 51 51 - } 52 52 - 53 53 - accessToken, _ = a.GetTokens() 54 54 - retry.Header.Set("Authorization", "Bearer "+accessToken) 55 55 - return c.Do(retry) 56 56 - } 57 57 - 58 58 - func (a *FixedPasswordAuth) Refresh(ctx context.Context, c *http.Client, priorRefreshToken string) error { 59 59 - a.lk.Lock() 60 60 - defer a.lk.Unlock() 61 61 - 62 62 - if priorRefreshToken != "" && priorRefreshToken != a.Session.RefreshToken { 63 63 - return nil 64 64 - } 65 65 - 66 66 - u := a.Session.Host + "/xrpc/com.atproto.server.refreshSession" 67 67 - req, err := http.NewRequestWithContext(ctx, http.MethodPost, u, nil) 68 68 - if err != nil { 69 69 - return err 70 70 - } 71 71 - req.Header.Set("User-Agent", "indigo-sdk") 72 72 - req.Header.Set("Authorization", "Bearer "+a.Session.RefreshToken) 73 73 - 74 74 - resp, err := c.Do(req) 75 75 - if err != nil { 76 76 - return err 77 77 - } 78 78 - defer resp.Body.Close() 79 79 - 80 80 - if resp.StatusCode < 200 || resp.StatusCode >= 300 { 81 81 - var eb atclient.ErrorBody 82 82 - if err := json.NewDecoder(resp.Body).Decode(&eb); err != nil { 83 83 - return &atclient.APIError{StatusCode: resp.StatusCode} 84 84 - } 85 85 - return eb.APIError(resp.StatusCode) 86 86 - } 87 87 - 88 88 - var out struct { 89 89 - AccessJwt string `json:"accessJwt"` 90 90 - RefreshJwt string `json:"refreshJwt"` 91 91 - } 92 92 - if err := json.NewDecoder(resp.Body).Decode(&out); err != nil { 93 93 - return err 94 94 - } 95 95 - 96 96 - a.Session.AccessToken = out.AccessJwt 97 97 - a.Session.RefreshToken = out.RefreshJwt 98 98 - 99 99 - if a.RefreshCallback != nil { 100 100 - snapshot := a.Session.Clone() 101 101 - a.RefreshCallback(ctx, snapshot) 102 102 - } 103 103 - 104 104 - return nil 105 105 - }

-88

sync/atproto_auth_test.go

reviewed

··· 1 1 - package sync 2 2 - 3 3 - import ( 4 4 - "context" 5 5 - "encoding/json" 6 6 - "net/http" 7 7 - "net/http/httptest" 8 8 - "strings" 9 9 - "testing" 10 10 - 11 11 - "github.com/bluesky-social/indigo/atproto/atclient" 12 12 - ) 13 13 - 14 14 - func TestFixedPasswordAuth_Refresh(t *testing.T) { 15 15 - t.Run("Refreshes with POST", func(t *testing.T) { 16 16 - methodUsed := "" 17 17 - server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 18 18 - if strings.HasSuffix(r.URL.Path, "refreshSession") { 19 19 - methodUsed = r.Method 20 20 - w.Header().Set("Content-Type", "application/json") 21 21 - json.NewEncoder(w).Encode(map[string]string{ 22 22 - "accessJwt": "new-access", 23 23 - "refreshJwt": "new-refresh", 24 24 - }) 25 25 - return 26 26 - } 27 27 - w.WriteHeader(http.StatusNotFound) 28 28 - })) 29 29 - defer server.Close() 30 30 - 31 31 - pa := &atclient.PasswordAuth{ 32 32 - Session: atclient.PasswordSessionData{ 33 33 - Host: server.URL, 34 34 - AccessToken: "old-access", 35 35 - RefreshToken: "old-refresh", 36 36 - }, 37 37 - } 38 38 - fixed := &FixedPasswordAuth{PasswordAuth: pa} 39 39 - 40 40 - err := fixed.Refresh(context.Background(), http.DefaultClient, "old-refresh") 41 41 - if err != nil { 42 42 - t.Fatalf("Refresh failed: %v", err) 43 43 - } 44 44 - 45 45 - if methodUsed != http.MethodPost { 46 46 - t.Errorf("Expected method POST, got %s", methodUsed) 47 47 - } 48 48 - 49 49 - if pa.Session.AccessToken != "new-access" { 50 50 - t.Errorf("Expected new access token, got %s", pa.Session.AccessToken) 51 51 - } 52 52 - }) 53 53 - } 54 54 - 55 55 - func TestLibraryPasswordAuth_Refresh_Method(t *testing.T) { 56 56 - // This test documents the library's buggy behavior (GET instead of POST) 57 57 - methodUsed := "" 58 58 - server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 59 59 - if strings.HasSuffix(r.URL.Path, "refreshSession") { 60 60 - methodUsed = r.Method 61 61 - w.Header().Set("Content-Type", "application/json") 62 62 - json.NewEncoder(w).Encode(map[string]string{ 63 63 - "accessJwt": "new-access", 64 64 - "refreshJwt": "new-refresh", 65 65 - }) 66 66 - return 67 67 - } 68 68 - })) 69 69 - defer server.Close() 70 70 - 71 71 - pa := &atclient.PasswordAuth{ 72 72 - Session: atclient.PasswordSessionData{ 73 73 - Host: server.URL, 74 74 - AccessToken: "old-access", 75 75 - RefreshToken: "old-refresh", 76 76 - }, 77 77 - } 78 78 - 79 79 - // We call the library's method directly 80 80 - _ = pa.Refresh(context.Background(), http.DefaultClient, "old-refresh") 81 81 - 82 82 - switch methodUsed { 83 83 - case http.MethodGet: 84 84 - t.Log("Confirmed: Library uses GET for refreshSession (Buggy)") 85 85 - case http.MethodPost: 86 86 - t.Log("Library uses POST for refreshSession") 87 87 - } 88 88 - }

-121

sync/auth.go

reviewed

··· 1 1 - package sync 2 2 - 3 3 - import ( 4 4 - "context" 5 5 - "fmt" 6 6 - "sync" 7 7 - 8 8 - "github.com/bluesky-social/indigo/atproto/atclient" 9 9 - ) 10 10 - 11 11 - type ResolvedIdentity struct { 12 12 - DID string `json:"did"` 13 13 - Handle string `json:"handle"` 14 14 - PDS string `json:"pds"` 15 15 - SigningKey string `json:"signingKey"` 16 16 - } 17 17 - 18 18 - type Client struct { 19 19 - client *atclient.APIClient 20 20 - resolvedIdentity ResolvedIdentity 21 21 - mu sync.Mutex 22 22 - } 23 23 - 24 24 - func NewClient(ctx context.Context, handle, password string) (*Client, error) { 25 25 - identity, err := ResolveIdentity(ctx, handle) 26 26 - if err != nil { 27 27 - return nil, fmt.Errorf("failed to resolve identity: %w", err) 28 28 - } 29 29 - 30 30 - pdsURL := identity.PDS 31 31 - if pdsURL == "" { 32 32 - pdsURL = "https://bsky.social" 33 33 - } 34 34 - 35 35 - client, err := atclient.LoginWithPasswordHost(ctx, pdsURL, handle, password, "", nil) 36 36 - if err != nil { 37 37 - return nil, fmt.Errorf("login failed: %w", err) 38 38 - } 39 39 - 40 40 - if pa, ok := client.Auth.(*atclient.PasswordAuth); ok { 41 41 - client.Auth = &FixedPasswordAuth{PasswordAuth: pa} 42 42 - } 43 43 - 44 44 - return &Client{ 45 45 - client: client, 46 46 - resolvedIdentity: identity, 47 47 - }, nil 48 48 - } 49 49 - 50 50 - func ResolveIdentity(ctx context.Context, handle string) (ResolvedIdentity, error) { 51 51 - if handle == "" { 52 52 - return ResolvedIdentity{}, fmt.Errorf("handle cannot be empty") 53 53 - } 54 54 - 55 55 - did, pds, signingKey, err := ResolveMiniDoc(ctx, handle) 56 56 - if err != nil { 57 57 - return ResolvedIdentity{}, fmt.Errorf("failed to resolve identity: %w", err) 58 58 - } 59 59 - 60 60 - if did == "" { 61 61 - return ResolvedIdentity{}, fmt.Errorf("resolved identity missing DID") 62 62 - } 63 63 - 64 64 - if signingKey == "" { 65 65 - return ResolvedIdentity{}, fmt.Errorf("resolved identity missing signing key") 66 66 - } 67 67 - 68 68 - return ResolvedIdentity{ 69 69 - DID: did, 70 70 - Handle: handle, 71 71 - PDS: pds, 72 72 - SigningKey: signingKey, 73 73 - }, nil 74 74 - } 75 75 - 76 76 - func (c *Client) GetDID() string { 77 77 - c.mu.Lock() 78 78 - defer c.mu.Unlock() 79 79 - return c.resolvedIdentity.DID 80 80 - } 81 81 - 82 82 - func (c *Client) GetPDS() string { 83 83 - c.mu.Lock() 84 84 - defer c.mu.Unlock() 85 85 - return c.resolvedIdentity.PDS 86 86 - } 87 87 - 88 88 - func (c *Client) GetHandle() string { 89 89 - c.mu.Lock() 90 90 - defer c.mu.Unlock() 91 91 - return c.resolvedIdentity.Handle 92 92 - } 93 93 - 94 94 - func (c *Client) GetSigningKey() string { 95 95 - c.mu.Lock() 96 96 - defer c.mu.Unlock() 97 97 - return c.resolvedIdentity.SigningKey 98 98 - } 99 99 - 100 100 - func (c *Client) GetAPIClient() *atclient.APIClient { 101 101 - c.mu.Lock() 102 102 - defer c.mu.Unlock() 103 103 - return c.client 104 104 - } 105 105 - 106 106 - func (c *Client) Close() error { 107 107 - c.mu.Lock() 108 108 - defer c.mu.Unlock() 109 109 - if c.client != nil && c.client.Auth != nil { 110 110 - if logout, ok := c.client.Auth.(*atclient.PasswordAuth); ok { 111 111 - return logout.Logout(context.Background(), c.client.Client) 112 112 - } 113 113 - } 114 114 - return nil 115 115 - } 116 116 - 117 117 - func (c *Client) HasClient() bool { 118 118 - c.mu.Lock() 119 119 - defer c.mu.Unlock() 120 120 - return c.client != nil 121 121 - }

+27 -10

sync/batch_test.go

reviewed

··· 13 13 "time" 14 14 15 15 "github.com/bluesky-social/indigo/atproto/atclient" 16 16 + "tangled.org/karitham.dev/lazuli/atproto" 16 17 "tangled.org/karitham.dev/lazuli/cache" 17 18 ) 18 19 ··· 138 139 139 140 // Mock ATProtoClient 140 141 type mockATProtoClient struct { 141 141 - applyWritesFunc func(ctx context.Context, collection string, records []PlayRecord) error 142 142 + applyWritesFunc func(ctx context.Context, collection string, records []PlayRecord) error 143 143 + listRecordsFunc func(ctx context.Context, collection string, limit int, cursor string) ([]atproto.RecordRef[PlayRecord], string, error) 144 144 + deleteRecordFunc func(ctx context.Context, collection, rkey string) error 142 145 } 143 146 144 147 func (m *mockATProtoClient) ApplyWrites(ctx context.Context, collection string, records []PlayRecord) error { 145 148 if m.applyWritesFunc != nil { 146 149 return m.applyWritesFunc(ctx, collection, records) 150 150 + } 151 151 + return nil 152 152 + } 153 153 + 154 154 + func (m *mockATProtoClient) ListRecords(ctx context.Context, collection string, limit int, cursor string) ([]atproto.RecordRef[PlayRecord], string, error) { 155 155 + if m.listRecordsFunc != nil { 156 156 + return m.listRecordsFunc(ctx, collection, limit, cursor) 157 157 + } 158 158 + return nil, "", nil 159 159 + } 160 160 + 161 161 + func (m *mockATProtoClient) DeleteRecord(ctx context.Context, collection, rkey string) error { 162 162 + if m.deleteRecordFunc != nil { 163 163 + return m.deleteRecordFunc(ctx, collection, rkey) 147 164 } 148 165 return nil 149 166 } ··· 153 170 did string 154 171 } 155 172 156 156 - func (m *mockAuthClient) GetAPIClient() *atclient.APIClient { return nil } 157 157 - func (m *mockAuthClient) GetDID() string { return m.did } 173 173 + func (m *mockAuthClient) APIClient() *atclient.APIClient { return nil } 174 174 + func (m *mockAuthClient) DID() string { return m.did } 158 175 159 176 type timeoutError struct{} 160 177 ··· 190 207 func TestApplyWrites_RateClient(t *testing.T) { 191 208 ctx := context.Background() 192 209 193 193 - t.Run("Empty records", func(t *testing.T) { 210 210 + t.Run("Empty anyRecords(records)", func(t *testing.T) { 194 211 limiter := &mockLimiter{} 195 195 - client := NewRateClient(nil, "did:example:123", limiter) 212 212 + client := atproto.NewRateClient[any](nil, "did:example:123", limiter) 196 213 err := client.ApplyWrites(ctx, "test", nil) 197 214 if err != nil { 198 215 t.Errorf("ApplyWrites(nil) error = %v", err) ··· 201 218 202 219 t.Run("Too many records", func(t *testing.T) { 203 220 limiter := &mockLimiter{} 204 204 - client := NewRateClient(&atclient.APIClient{}, "did:example:123", limiter) 205 205 - records := make([]PlayRecord, 201) 206 206 - err := client.ApplyWrites(ctx, "test", records) 221 221 + client := atproto.NewRateClient[any](&atclient.APIClient{}, "did:example:123", limiter) 222 222 + recs := make([]any, 201) 223 223 + err := client.ApplyWrites(ctx, "test", recs) 207 224 if err == nil { 208 225 t.Fatal("expected error for > 200 records") 209 226 } ··· 222 239 Body: http.NoBody, 223 240 }, nil 224 241 }) 225 225 - client := NewRateClient(apiClient, "did:example:123", limiter) 242 242 + client := atproto.NewRateClient[PlayRecord](apiClient, "did:example:123", limiter) 226 243 227 244 err := client.ApplyWrites(ctx, "test", []PlayRecord{{TrackName: "Song 1"}}) 228 245 if err == nil { ··· 242 259 Body: http.NoBody, 243 260 }, nil 244 261 }) 245 245 - client := NewRateClient(apiClient, "did:example:123", limiter) 262 262 + client := atproto.NewRateClient[PlayRecord](apiClient, "did:example:123", limiter) 246 263 247 264 err := client.ApplyWrites(ctx, "test", []PlayRecord{{TrackName: "Song 1"}}) 248 265 if err == nil {

+2 -2

sync/import_test.go

reviewed

··· 54 54 did string 55 55 } 56 56 57 57 - func (m *mockAuthClient) GetAPIClient() *atclient.APIClient { return nil } 58 58 - func (m *mockAuthClient) GetDID() string { return m.did } 57 57 + func (m *mockAuthClient) APIClient() *atclient.APIClient { return nil } 58 58 + func (m *mockAuthClient) DID() string { return m.did } 59 59 60 60 type mockKV struct { 61 61 data map[string]int

+38

sync/logutil.go

reviewed

··· 1 1 + package sync 2 2 + 3 3 + import ( 4 4 + "log/slog" 5 5 + "time" 6 6 + 7 7 + "tangled.org/karitham.dev/lazuli/atproto" 8 8 + ) 9 9 + 10 10 + type TrackInfo struct { 11 11 + Name string 12 12 + Artist string 13 13 + PlayedAt time.Time 14 14 + } 15 15 + 16 16 + func Track(name, artist string, playedAt time.Time) slog.Attr { 17 17 + return slog.Group("track", 18 18 + slog.String("name", name), 19 19 + slog.String("artist", artist), 20 20 + slog.Time("played_at", playedAt), 21 21 + ) 22 22 + } 23 23 + 24 24 + func DID(did string) slog.Attr { 25 25 + return slog.String("did", did) 26 26 + } 27 27 + 28 28 + func Error(err error) slog.Attr { 29 29 + if err == nil { 30 30 + return slog.Attr{} 31 31 + } 32 32 + 33 33 + if atproto.IsTransientError(err) { 34 34 + return slog.String("error", err.Error()) 35 35 + } 36 36 + 37 37 + return slog.String("error", err.Error()) 38 38 + }

-47

sync/logutil/logutil.go

reviewed

··· 1 1 - package logutil 2 2 - 3 3 - import ( 4 4 - "errors" 5 5 - "log/slog" 6 6 - "time" 7 7 - 8 8 - "github.com/bluesky-social/indigo/atproto/atclient" 9 9 - ) 10 10 - 11 11 - type TrackInfo struct { 12 12 - Name string 13 13 - Artist string 14 14 - PlayedAt time.Time 15 15 - } 16 16 - 17 17 - // Track returns slog attributes for a play record. 18 18 - func Track(name, artist string, playedAt time.Time) slog.Attr { 19 19 - return slog.Group("track", 20 20 - slog.String("name", name), 21 21 - slog.String("artist", artist), 22 22 - slog.Time("played_at", playedAt), 23 23 - ) 24 24 - } 25 25 - 26 26 - // DID returns a slog attribute for a DID. 27 27 - func DID(did string) slog.Attr { 28 28 - return slog.String("did", did) 29 29 - } 30 30 - 31 31 - // Error returns a slog attribute for an error. 32 32 - func Error(err error) slog.Attr { 33 33 - if err == nil { 34 34 - return slog.Attr{} 35 35 - } 36 36 - 37 37 - var apiErr *atclient.APIError 38 38 - if errors.As(err, &apiErr) { 39 39 - return slog.Group("error", 40 40 - slog.Int("status", apiErr.StatusCode), 41 41 - slog.String("name", apiErr.Name), 42 42 - slog.String("message", apiErr.Message), 43 43 - ) 44 44 - } 45 45 - 46 46 - return slog.String("error", err.Error()) 47 47 - }

+291 -16

sync/publish.go

reviewed

··· 5 5 "encoding/json" 6 6 "fmt" 7 7 "log/slog" 8 8 + "strings" 8 9 "time" 9 10 11 11 + "github.com/bluesky-social/indigo/atproto/atclient" 10 12 "github.com/bluesky-social/indigo/atproto/syntax" 11 13 "github.com/failsafe-go/failsafe-go" 12 14 "github.com/failsafe-go/failsafe-go/retrypolicy" 13 15 16 16 + "tangled.org/karitham.dev/lazuli/atproto" 14 17 "tangled.org/karitham.dev/lazuli/cache" 15 15 - "tangled.org/karitham.dev/lazuli/sync/logutil" 18 18 + ) 19 19 + 20 20 + type ( 21 21 + ATProtoClient = atproto.RepoClient[PlayRecord] 22 22 + AuthClient = atproto.AuthClient 23 23 + RateLimiter = atproto.RateLimiter 16 24 ) 17 25 26 26 + var ( 27 27 + WriteLimitDay = atproto.WriteLimitDay 28 28 + GlobalLimitDay = atproto.GlobalLimitDay 29 29 + ) 30 30 + 31 31 + func NewRateLimiter(kv atproto.KVStore, maxPercent float32) RateLimiter { 32 32 + return atproto.NewRateLimiter(kv, maxPercent) 33 33 + } 34 34 + 35 35 + func NewRateClient(client *atclient.APIClient, did string, limiter RateLimiter) *atproto.RateClient[PlayRecord] { 36 36 + return atproto.NewRateClient[PlayRecord](client, did, limiter) 37 37 + } 38 38 + 39 39 + func IsTransientError(err error) bool { 40 40 + return atproto.IsTransientError(err) 41 41 + } 42 42 + 43 43 + type Client = atproto.Client 44 44 + 45 45 + func NewClient(ctx context.Context, handle, password string, opts ...func(*atproto.ClientOptions)) (*Client, error) { 46 46 + return atproto.NewClient(ctx, handle, password, opts...) 47 47 + } 48 48 + 49 49 + type RepoClient[T any] = atproto.RepoClient[T] 50 50 + 51 51 + type RecordRef = atproto.RecordRef[PlayRecord] 52 52 + 18 53 type PublishOptions struct { 19 54 BatchSize int 20 55 DryRun bool ··· 29 64 30 65 batchSize := defaultBatchSize(opts.BatchSize) 31 66 32 32 - atprotoClient, err := buildClient(client, opts.ATProtoClient) 67 67 + atprotoClient, err := atproto.BuildClient(client, opts.ATProtoClient) 33 68 if err != nil { 34 69 return PublishResult{ 35 70 SuccessCount: 0, ··· 42 77 43 78 totalRecords := 0 44 79 if opts.Storage != nil { 45 45 - _ = opts.Storage.IterateUnpublished(client.GetDID(), func(key string, rec []byte) error { 80 80 + _ = opts.Storage.IterateUnpublished(client.DID(), func(key string, rec []byte) error { 46 81 totalRecords++ 47 82 return nil 48 83 }) ··· 55 90 slog.Info("starting iterative import", 56 91 slog.Int("total_records", totalRecords), 57 92 slog.Int("batch_size", batchSize), 58 58 - slog.Int("daily_write_limit", WriteLimitDay), 59 59 - slog.Int("daily_token_limit", GlobalLimitDay), 60 60 - slog.String("rate_limit", fmt.Sprintf("1 write per %.1fs", 86400.0/WriteLimitDay))) 93 93 + slog.Int("daily_write_limit", atproto.WriteLimitDay), 94 94 + slog.Int("daily_token_limit", atproto.GlobalLimitDay), 95 95 + slog.String("rate_limit", fmt.Sprintf("1 write per %.1fs", 86400.0/atproto.WriteLimitDay))) 61 96 62 97 tracker := NewProgressTracker(totalRecords, opts.Limiter) 63 98 progressLog := defaultProgressLog(opts.ProgressLog) ··· 76 111 for _, r := range batch { 77 112 tid := syntax.NewTIDFromTime(r.PlayedTime.Time, 0) 78 113 slog.Info("would publish record (dry run)", 79 79 - logutil.Track(r.TrackName, r.ArtistName(), r.PlayedTime.Time), 114 114 + Track(r.TrackName, r.ArtistName(), r.PlayedTime.Time), 80 115 slog.String("rkey", string(tid))) 81 116 } 82 117 totalSuccess += len(batch) ··· 86 121 return nil 87 122 } 88 123 89 89 - did := client.GetDID() 124 124 + did := client.DID() 90 125 retryPolicy := retrypolicy.NewBuilder[any](). 91 126 WithMaxRetries(10). 92 127 WithBackoff(BaseRetryDelay, 5*time.Minute). ··· 97 132 slog.Warn("batch failed with transient error, retrying", 98 133 slog.Int("count", len(batch)), 99 134 slog.Duration("retryDelay", e.Delay), 100 100 - logutil.Error(e.LastError()), 135 135 + Error(e.LastError()), 101 136 slog.Int("attempt", e.Attempts())) 102 137 }). 103 138 Build() ··· 107 142 }) 108 143 if err != nil { 109 144 slog.Error("batch failed after retries", 110 110 - logutil.Error(err), 145 145 + Error(err), 111 146 slog.Int("count", len(batch))) 112 147 113 148 if opts.Storage != nil { 114 149 if markErr := opts.Storage.MarkFailed(did, batchKeys, err.Error()); markErr != nil { 115 115 - slog.Error("failed to mark records as failed", logutil.Error(markErr)) 150 150 + slog.Error("failed to mark records as failed", Error(markErr)) 116 151 } 117 152 } 118 153 ··· 121 156 122 157 batch = batch[:0] 123 158 batchKeys = batchKeys[:0] 124 124 - return nil // Return nil so we continue with the next batch 159 159 + return nil 125 160 } 126 161 127 162 totalSuccess += len(batch) ··· 136 171 return nil 137 172 } 138 173 139 139 - err = opts.Storage.IterateUnpublished(client.GetDID(), func(key string, rec []byte) error { 174 174 + err = opts.Storage.IterateUnpublished(client.DID(), func(key string, rec []byte) error { 140 175 select { 141 176 case <-ctx.Done(): 142 177 return ctx.Err() ··· 145 180 146 181 var record PlayRecord 147 182 if err := json.Unmarshal(rec, &record); err != nil { 148 148 - slog.Error("malformed record in storage", slog.String("key", key), logutil.Error(err)) 183 183 + slog.Error("malformed record in storage", slog.String("key", key), Error(err)) 149 184 if opts.Storage != nil { 150 150 - _ = opts.Storage.MarkFailed(client.GetDID(), []string{key}, "malformed record") 185 185 + _ = opts.Storage.MarkFailed(client.DID(), []string{key}, "malformed record") 151 186 } 152 187 totalErrors++ 153 188 tracker.IncrementErrors(1) ··· 171 206 172 207 cancelled := false 173 208 if err != nil { 174 174 - slog.Error("import interrupted", logutil.Error(err)) 209 209 + slog.Error("import interrupted", Error(err)) 175 210 cancelled = true 176 211 } 177 212 178 213 logResult(totalSuccess, totalErrors, startTime) 179 214 return newPublishResult(totalSuccess, totalErrors, totalRecords, startTime, cancelled) 180 215 } 216 216 + 217 217 + func isTransientError(err error) bool { 218 218 + return atproto.IsTransientError(err) 219 219 + } 220 220 + 221 221 + func defaultProgressLog(f func(ProgressReport)) func(ProgressReport) { 222 222 + if f != nil { 223 223 + return f 224 224 + } 225 225 + return func(pr ProgressReport) { 226 226 + slog.Info("sync progress", 227 227 + slog.Int("completed", pr.Completed), 228 228 + slog.Int("total", pr.Total), 229 229 + slog.Float64("percent", pr.Percent), 230 230 + slog.String("elapsed", pr.Elapsed), 231 231 + slog.String("eta", pr.ETA), 232 232 + slog.String("rate", pr.Rate), 233 233 + slog.Int("errors", pr.Errors), 234 234 + ) 235 235 + } 236 236 + } 237 237 + 238 238 + func defaultBatchSize(size int) int { 239 239 + if size > 0 { 240 240 + return size 241 241 + } 242 242 + return DefaultBatchSize 243 243 + } 244 244 + 245 245 + func newPublishResult(success, errors, total int, start time.Time, cancelled bool) PublishResult { 246 246 + return PublishResult{ 247 247 + SuccessCount: success, 248 248 + ErrorCount: errors, 249 249 + Cancelled: cancelled, 250 250 + Duration: time.Since(start), 251 251 + TotalRecords: total, 252 252 + RecordsPerMinute: ratePerMinute(success, time.Since(start)), 253 253 + } 254 254 + } 255 255 + 256 256 + func logResult(success, errors int, startTime time.Time) { 257 257 + if errors > 0 { 258 258 + slog.Warn("import completed with errors", 259 259 + slog.Int("success", success), 260 260 + slog.Int("errors", errors)) 261 261 + } 262 262 + slog.Info("import completed", 263 263 + slog.Int("success", success), 264 264 + slog.Int("errors", errors), 265 265 + slog.Duration("duration", time.Since(startTime)), 266 266 + slog.String("rate", formatRate(ratePerMinute(success, time.Since(startTime))))) 267 267 + } 268 268 + 269 269 + func PublishBatch(ctx context.Context, client ATProtoClient, did string, batch []PlayRecord, storage cache.Storage) error { 270 270 + if len(batch) == 0 { 271 271 + return nil 272 272 + } 273 273 + 274 274 + err := client.ApplyWrites(ctx, RecordType, batch) 275 275 + if err != nil { 276 276 + slog.Error("batch publish failed", Error(err)) 277 277 + return err 278 278 + } 279 279 + 280 280 + if storage != nil && did != "" { 281 281 + keys := CreateRecordKeys(batch) 282 282 + cacheEntries := make(map[string][]byte) 283 283 + for i, rec := range batch { 284 284 + key := keys[i] 285 285 + value, _ := json.Marshal(rec) 286 286 + cacheEntries[key] = value 287 287 + } 288 288 + 289 289 + if err := storage.SaveRecords(did, cacheEntries); err != nil { 290 290 + return fmt.Errorf("failed to save records to storage: %w", err) 291 291 + } 292 292 + 293 293 + if err := storage.MarkPublished(did, keys...); err != nil { 294 294 + return fmt.Errorf("failed to mark records as published: %w", err) 295 295 + } 296 296 + } 297 297 + 298 298 + return nil 299 299 + } 300 300 + 301 301 + func prepareRecords(batch []PlayRecord) []PlayRecord { 302 302 + atprotoRecords := make([]PlayRecord, 0, len(batch)) 303 303 + for _, record := range batch { 304 304 + record.Type = RecordType 305 305 + record.SubmissionClientAgent = ClientAgent 306 306 + atprotoRecords = append(atprotoRecords, record) 307 307 + } 308 308 + return atprotoRecords 309 309 + } 310 310 + 311 311 + func ratePerMinute(count int, duration time.Duration) float64 { 312 312 + if duration == 0 { 313 313 + return 0 314 314 + } 315 315 + return float64(count) / duration.Minutes() 316 316 + } 317 317 + 318 318 + func FetchExisting(ctx context.Context, client RepoClient[PlayRecord], did string, storage cache.Storage, forceRefresh bool) ([]ExistingRecord, error) { 319 319 + if !forceRefresh && storage != nil { 320 320 + published, err := storage.GetPublished(did) 321 321 + if err == nil && len(published) > 0 && storage.IsValid(did) { 322 322 + records := make([]ExistingRecord, 0, len(published)) 323 323 + err := storage.IteratePublished(did, func(key string, data []byte) error { 324 324 + var value PlayRecord 325 325 + if err := json.Unmarshal(data, &value); err != nil { 326 326 + return nil 327 327 + } 328 328 + records = append(records, ExistingRecord{ 329 329 + URI: generateRecordURI(did, value), 330 330 + Value: value, 331 331 + }) 332 332 + return nil 333 333 + }) 334 334 + if err == nil { 335 335 + slog.Debug("loaded from cache", slog.Int("count", len(records))) 336 336 + return records, nil 337 337 + } 338 338 + } 339 339 + } 340 340 + 341 341 + select { 342 342 + case <-ctx.Done(): 343 343 + return nil, ctx.Err() 344 344 + default: 345 345 + } 346 346 + 347 347 + allRecords := make([]ExistingRecord, 0, 1024) 348 348 + return fetchExistingLoop(ctx, client, did, storage, allRecords) 349 349 + } 350 350 + 351 351 + func fetchExistingLoop(ctx context.Context, client RepoClient[PlayRecord], did string, storage cache.Storage, allRecords []ExistingRecord) ([]ExistingRecord, error) { 352 352 + const batchSize = 100 353 353 + var cursor string 354 354 + 355 355 + type fetchResult struct { 356 356 + records []atproto.RecordRef[PlayRecord] 357 357 + cursor string 358 358 + } 359 359 + 360 360 + retryPolicy := retrypolicy.NewBuilder[fetchResult](). 361 361 + WithMaxRetries(10). 362 362 + WithBackoff(BaseRetryDelay, 5*time.Minute). 363 363 + HandleIf(func(_ fetchResult, err error) bool { 364 364 + return IsTransientError(err) 365 365 + }). 366 366 + OnRetryScheduled(func(e failsafe.ExecutionScheduledEvent[fetchResult]) { 367 367 + slog.Warn("fetch failed with transient error, retrying", 368 368 + slog.Duration("retryDelay", e.Delay), 369 369 + Error(e.LastError()), 370 370 + slog.Int("attempt", e.Attempts())) 371 371 + }). 372 372 + Build() 373 373 + 374 374 + for { 375 375 + select { 376 376 + case <-ctx.Done(): 377 377 + return nil, ctx.Err() 378 378 + default: 379 379 + } 380 380 + 381 381 + result, err := failsafe.With(retryPolicy). 382 382 + WithContext(ctx). 383 383 + Get(func() (fetchResult, error) { 384 384 + recs, next, err := client.ListRecords(ctx, RecordType, batchSize, cursor) 385 385 + if err != nil { 386 386 + return fetchResult{}, err 387 387 + } 388 388 + 389 389 + return fetchResult{records: recs, cursor: next}, nil 390 390 + }) 391 391 + if err != nil { 392 392 + return nil, err 393 393 + } 394 394 + 395 395 + for _, rec := range result.records { 396 396 + allRecords = append(allRecords, ExistingRecord(rec)) 397 397 + } 398 398 + 399 399 + if result.cursor == "" || len(result.records) < batchSize { 400 400 + break 401 401 + } 402 402 + cursor = result.cursor 403 403 + } 404 404 + 405 405 + if storage != nil { 406 406 + cacheEntries := make(map[string][]byte) 407 407 + keys := make([]string, 0, len(allRecords)) 408 408 + for _, rec := range allRecords { 409 409 + parts := strings.Split(rec.URI, "/") 410 410 + key := parts[len(parts)-1] 411 411 + if key == "" { 412 412 + key = CreateRecordKey(rec.Value) 413 413 + } 414 414 + value, _ := json.Marshal(rec.Value) 415 415 + cacheEntries[key] = value 416 416 + keys = append(keys, key) 417 417 + } 418 418 + 419 419 + if err := storage.SaveRecords(did, cacheEntries); err != nil { 420 420 + return nil, err 421 421 + } 422 422 + 423 423 + if err := storage.MarkPublished(did, keys...); err != nil { 424 424 + return nil, err 425 425 + } 426 426 + 427 427 + slog.Debug("saved to cache and marked as published", slog.Int("count", len(allRecords))) 428 428 + } 429 429 + 430 430 + return allRecords, nil 431 431 + } 432 432 + 433 433 + func generateRecordURI(did string, record PlayRecord) string { 434 434 + return fmt.Sprintf("at://%s/%s/%s", did, RecordType, CreateRecordKey(record)) 435 435 + } 436 436 + 437 437 + func prepareWrites(records []PlayRecord, collection string) ([]map[string]any, error) { 438 438 + if len(records) == 0 { 439 439 + return nil, nil 440 440 + } 441 441 + 442 442 + writes := make([]map[string]any, len(records)) 443 443 + keys := CreateRecordKeys(records) 444 444 + 445 445 + for i, rec := range records { 446 446 + writes[i] = map[string]any{ 447 447 + "$type": "com.atproto.repo.applyWrites#create", 448 448 + "collection": collection, 449 449 + "rkey": keys[i], 450 450 + "value": rec, 451 451 + } 452 452 + } 453 453 + 454 454 + return writes, nil 455 455 + }

+16 -35

sync/rate.go atproto/rate.go

reviewed

··· 1 1 - package sync 1 1 + package atproto 2 2 3 3 import ( 4 4 "context" ··· 9 9 "math" 10 10 "sync" 11 11 "time" 12 12 - 13 13 - "tangled.org/karitham.dev/lazuli/sync/logutil" 14 12 ) 15 13 16 14 const ( 17 17 - // Limits 18 15 WriteLimitMinute = 100 19 16 WriteLimitHour = 1000 20 17 WriteLimitDay = 10000 ··· 23 20 GlobalLimitHour = 3000 24 21 GlobalLimitDay = 35000 25 22 26 26 - // Costs 27 23 ReadGlobalCost = 1 28 24 WriteOnlyCost = 1 29 25 WriteGlobalCost = 3 30 26 ) 31 27 32 28 type RateLimiter interface { 33 33 - // AllowBulkWrite blocks or returns error until N writes are permissible. 34 34 - // Returns the timestamp of when the quota was charged for bucket-accurate refunds. 35 29 AllowBulkWrite(ctx context.Context, n int) (time.Time, error) 36 36 - // AllowRead blocks or returns error until a read is permissible. 37 30 AllowRead(ctx context.Context) (time.Time, error) 38 38 - // RefundBulkWrite restores N writes to the quota using the original charge time. 39 31 RefundBulkWrite(ctx context.Context, n int, chargedAt time.Time) 40 40 - // RefundRead restores a read to the quota using the original charge time. 41 32 RefundRead(ctx context.Context, chargedAt time.Time) 42 42 - // Stats returns current consumption (writes, global) 43 33 Stats() (int, int, error) 44 44 - // EstimatedWriteTime returns how long it would take to write n records 45 45 - // given current consumption and rate limits. 46 34 EstimatedWriteTime(n int) time.Duration 47 47 - // RemainingQuota returns remaining write and global quota at each tier, 48 48 - // plus the time until the next tier reset (the maximum wait). 49 35 RemainingQuota() (writesRemaining, globalRemaining int, timeUntilReset time.Duration) 50 36 } 51 37 ··· 115 101 if maxWait > 0 { 116 102 l.mu.Unlock() 117 103 slog.Debug("Rate limit reached, sleeping until next window", slog.Duration("wait", maxWait.Round(time.Second))) 118 118 - // Add a tiny bit of buffer + jitter to ensure we are definitely in the next window 119 104 wait := maxWait + 100*time.Millisecond + addJitter(100*time.Millisecond) 120 105 if err := l.sleep(ctx, wait); err != nil { 121 106 return now, err ··· 123 108 continue 124 109 } 125 110 126 126 - // Charge quota while holding the lock 127 111 err = l.charge(wKeys, gKeys, wCost, gCost) 128 112 l.mu.Unlock() 129 113 ··· 158 142 continue 159 143 } 160 144 161 161 - // Charge quota while holding the lock 162 145 err = l.charge(nil, gKeys, 0, gCost) 163 146 l.mu.Unlock() 164 147 ··· 191 174 } 192 175 193 176 if err := l.kv.IncrByMulti(updates); err != nil { 194 194 - slog.Error("failed to refund write quota", logutil.Error(err)) 177 177 + slog.Error("failed to refund write quota", Error(err)) 195 178 } 196 179 } 197 180 ··· 208 191 } 209 192 210 193 if err := l.kv.IncrByMulti(updates); err != nil { 211 211 - slog.Error("failed to refund global quota (read)", logutil.Error(err)) 194 194 + slog.Error("failed to refund global quota (read)", Error(err)) 212 195 } 213 196 } 214 197 215 215 - // checkQuota checks if the proposed cost fits within the limits. 216 216 - // Returns the wait duration if over limit (0 if OK), or error. 217 217 - // Must be called with lock held. 218 198 func (l *quotaLimiter) checkQuota(now time.Time, wKeys, gKeys []string, wCost, gCost int) (time.Duration, error) { 219 199 wLimits := []int{WriteLimitMinute, WriteLimitHour, WriteLimitDay} 220 200 gLimits := []int{GlobalLimitMinute, GlobalLimitHour, GlobalLimitDay} ··· 251 231 return maxWait, nil 252 232 } 253 233 254 254 - // charge applies the cost to the keys. 255 255 - // Must be called with lock held. 256 234 func (l *quotaLimiter) charge(wKeys, gKeys []string, wCost, gCost int) error { 257 235 updates := make(map[string]int, len(wKeys)+len(gKeys)) 258 236 for _, k := range wKeys { ··· 274 252 275 253 func (l *quotaLimiter) untilNextWindow(now time.Time, tier int) time.Duration { 276 254 switch tier { 277 277 - case 0: // Minute 255 255 + case 0: 278 256 return now.Truncate(time.Minute).Add(time.Minute).Sub(now) 279 279 - case 1: // Hour 257 257 + case 1: 280 258 return now.Truncate(time.Hour).Add(time.Hour).Sub(now) 281 281 - case 2: // Day 259 259 + case 2: 282 260 return now.Truncate(24 * time.Hour).Add(24 * time.Hour).Sub(now) 283 261 default: 284 284 - // Safety fallback: if unknown tier, wait a reasonable amount (e.g. 1 minute) 285 285 - // to prevent busy loops or bypassing limits. 286 262 slog.Warn("unknown rate limit tier encountered", slog.Int("tier", tier)) 287 263 return time.Minute 288 264 } ··· 292 268 var b [8]byte 293 269 _, _ = rand.Read(b[:]) 294 270 n := binary.LittleEndian.Uint64(b[:]) 295 295 - // Add 0-10% jitter 296 271 jitter := float64(d) * 0.1 * (float64(n) / math.MaxUint64) 297 272 return d + time.Duration(jitter) 298 273 } ··· 336 311 for i, k := range wKeys { 337 312 curr := values[k] 338 313 limit := int(float32(WriteLimitMinute) * l.rlQuota) 339 339 - if i == 1 { 314 314 + switch i { 315 315 + case 1: 340 316 limit = int(float32(WriteLimitHour) * l.rlQuota) 341 341 - } else if i == 2 { 317 317 + case 2: 342 318 limit = int(float32(WriteLimitDay) * l.rlQuota) 343 319 } 344 320 if curr+wCost > limit { ··· 349 325 for i, k := range gKeys { 350 326 curr := values[k] 351 327 limit := int(float32(GlobalLimitMinute) * l.rlQuota) 352 352 - if i == 1 { 328 328 + switch i { 329 329 + case 1: 353 330 limit = int(float32(GlobalLimitHour) * l.rlQuota) 354 354 - } else if i == 2 { 331 331 + case 2: 355 332 limit = int(float32(GlobalLimitDay) * l.rlQuota) 356 333 } 357 334 if curr+gCost > limit { ··· 426 403 427 404 return minWritesRemaining, minGlobalRemaining, maxResetTime 428 405 } 406 406 + 407 407 + func Error(err error) slog.Attr { 408 408 + return slog.String("error", err.Error()) 409 409 + }

+182 -121

sync/rate_test.go

reviewed

··· 3 3 import ( 4 4 "context" 5 5 "encoding/json" 6 6 - "strings" 6 6 + "fmt" 7 7 "testing" 8 8 "time" 9 9 10 10 "github.com/bluesky-social/indigo/atproto/atclient" 11 11 + "tangled.org/karitham.dev/lazuli/atproto" 11 12 ) 12 13 13 14 type mockKV struct { ··· 29 30 return nil 30 31 } 31 32 32 32 - // Helper for tests that inspect internal state directly 33 33 - func (m *mockKV) Get(key string) (int, error) { 34 34 - return m.data[key], nil 33 33 + type testClock struct { 34 34 + now time.Time 35 35 } 36 36 + 37 37 + func (m *testClock) Now() time.Time { return m.now } 36 38 37 39 func TestRateLimiter_Refunds(t *testing.T) { 38 40 kv := &mockKV{data: make(map[string]int)} 39 39 - clock := &mockClock{now: time.Date(2026, 1, 22, 12, 0, 0, 0, time.UTC)} 40 40 - limiter := &quotaLimiter{ 41 41 + clock := &testClock{now: time.Date(2026, 1, 22, 12, 0, 0, 0, time.UTC)} 42 42 + limiter := &testRateLimiter{ 41 43 kv: kv, 42 44 prefix: "quota", 43 45 clock: clock, ··· 45 47 } 46 48 ctx := context.Background() 47 49 48 48 - // Test Bulk Write Refund 49 50 chargedAt, _ := limiter.AllowBulkWrite(ctx, 10) 50 51 limiter.RefundBulkWrite(ctx, 10, chargedAt) 51 52 ··· 54 55 t.Errorf("BulkWrite refund failed: w=%d, g=%d", w, g) 55 56 } 56 57 57 57 - // Test Read Refund 58 58 chargedAt, _ = limiter.AllowRead(ctx) 59 59 limiter.RefundRead(ctx, chargedAt) 60 60 ··· 64 64 } 65 65 } 66 66 67 67 + type testRateLimiter struct { 68 68 + kv *mockKV 69 69 + prefix string 70 70 + clock *testClock 71 71 + rlQuota float32 72 72 + } 73 73 + 74 74 + func (l *testRateLimiter) Stats() (int, int, error) { 75 75 + wd, gd, _, _, _, _ := l.getKeys(l.clock.now) 76 76 + vals, err := l.kv.GetMulti([]string{wd, gd}) 77 77 + if err != nil { 78 78 + return 0, 0, err 79 79 + } 80 80 + return vals[wd], vals[gd], nil 81 81 + } 82 82 + 83 83 + func (l *testRateLimiter) AllowBulkWrite(ctx context.Context, n int) (time.Time, error) { 84 84 + wCost := n * atproto.WriteOnlyCost 85 85 + gCost := n * atproto.WriteGlobalCost 86 86 + 87 87 + for { 88 88 + now := l.clock.now 89 89 + wKeys, gKeys := l.getAllKeys(now) 90 90 + 91 91 + maxWait, err := l.checkQuota(now, wKeys, gKeys, wCost, gCost) 92 92 + if err != nil { 93 93 + return now, err 94 94 + } 95 95 + 96 96 + if maxWait > 0 { 97 97 + return now, context.DeadlineExceeded 98 98 + } 99 99 + 100 100 + err = l.charge(wKeys, gKeys, wCost, gCost) 101 101 + if err != nil { 102 102 + return now, err 103 103 + } 104 104 + return now, nil 105 105 + } 106 106 + } 107 107 + 108 108 + func (l *testRateLimiter) AllowRead(ctx context.Context) (time.Time, error) { 109 109 + gCost := atproto.ReadGlobalCost 110 110 + 111 111 + for { 112 112 + now := l.clock.now 113 113 + _, gKeys := l.getAllKeys(now) 114 114 + 115 115 + maxWait, err := l.checkQuota(now, nil, gKeys, 0, gCost) 116 116 + if err != nil { 117 117 + return now, err 118 118 + } 119 119 + 120 120 + if maxWait > 0 { 121 121 + return now, context.DeadlineExceeded 122 122 + } 123 123 + 124 124 + err = l.charge(nil, gKeys, 0, gCost) 125 125 + if err != nil { 126 126 + return now, err 127 127 + } 128 128 + return now, nil 129 129 + } 130 130 + } 131 131 + 132 132 + func (l *testRateLimiter) RefundBulkWrite(ctx context.Context, n int, chargedAt time.Time) { 133 133 + wKeys, gKeys := l.getAllKeys(chargedAt) 134 134 + wCost := n * atproto.WriteOnlyCost 135 135 + gCost := n * atproto.WriteGlobalCost 136 136 + 137 137 + updates := make(map[string]int, len(wKeys)+len(gKeys)) 138 138 + for _, k := range wKeys { 139 139 + updates[k] = -wCost 140 140 + } 141 141 + for _, k := range gKeys { 142 142 + updates[k] = -gCost 143 143 + } 144 144 + 145 145 + l.kv.IncrByMulti(updates) 146 146 + } 147 147 + 148 148 + func (l *testRateLimiter) RefundRead(ctx context.Context, chargedAt time.Time) { 149 149 + _, gKeys := l.getAllKeys(chargedAt) 150 150 + gCost := atproto.ReadGlobalCost 151 151 + 152 152 + updates := make(map[string]int, len(gKeys)) 153 153 + for _, k := range gKeys { 154 154 + updates[k] = -gCost 155 155 + } 156 156 + 157 157 + l.kv.IncrByMulti(updates) 158 158 + } 159 159 + 160 160 + func (l *testRateLimiter) getKeys(t time.Time) (string, string, string, string, string, string) { 161 161 + day := t.Format("2006-01-02") 162 162 + hour := t.Format("2006-01-02-15") 163 163 + minute := t.Format("2006-01-02-15-04") 164 164 + return fmt.Sprintf("%s:writes:d:%s", l.prefix, day), fmt.Sprintf("%s:global:d:%s", l.prefix, day), 165 165 + fmt.Sprintf("%s:writes:h:%s", l.prefix, hour), fmt.Sprintf("%s:global:h:%s", l.prefix, hour), 166 166 + fmt.Sprintf("%s:writes:m:%s", l.prefix, minute), fmt.Sprintf("%s:global:m:%s", l.prefix, minute) 167 167 + } 168 168 + 169 169 + func (l *testRateLimiter) getAllKeys(t time.Time) ([]string, []string) { 170 170 + wd, gd, wh, gh, wm, gm := l.getKeys(t) 171 171 + return []string{wm, wh, wd}, []string{gm, gh, gd} 172 172 + } 173 173 + 174 174 + func (l *testRateLimiter) checkQuota(now time.Time, wKeys, gKeys []string, wCost, gCost int) (time.Duration, error) { 175 175 + wLimits := []int{atproto.WriteLimitMinute, atproto.WriteLimitHour, atproto.WriteLimitDay} 176 176 + gLimits := []int{atproto.GlobalLimitMinute, atproto.GlobalLimitHour, atproto.GlobalLimitDay} 177 177 + 178 178 + allKeys := make([]string, 0, len(wKeys)+len(gKeys)) 179 179 + allKeys = append(allKeys, wKeys...) 180 180 + allKeys = append(allKeys, gKeys...) 181 181 + 182 182 + if len(allKeys) == 0 { 183 183 + return 0, nil 184 184 + } 185 185 + 186 186 + values, err := l.kv.GetMulti(allKeys) 187 187 + if err != nil { 188 188 + return 0, err 189 189 + } 190 190 + 191 191 + maxWait := time.Duration(0) 192 192 + 193 193 + for i, k := range wKeys { 194 194 + curr := values[k] 195 195 + if curr+wCost > int(float32(wLimits[i])*l.rlQuota) { 196 196 + maxWait = max(l.untilNextWindow(now, i), maxWait) 197 197 + } 198 198 + } 199 199 + 200 200 + for i, k := range gKeys { 201 201 + curr := values[k] 202 202 + if curr+gCost > int(float32(gLimits[i])*l.rlQuota) { 203 203 + maxWait = max(l.untilNextWindow(now, i), maxWait) 204 204 + } 205 205 + } 206 206 + 207 207 + return maxWait, nil 208 208 + } 209 209 + 210 210 + func (l *testRateLimiter) charge(wKeys, gKeys []string, wCost, gCost int) error { 211 211 + updates := make(map[string]int, len(wKeys)+len(gKeys)) 212 212 + for _, k := range wKeys { 213 213 + updates[k] = wCost 214 214 + } 215 215 + for _, k := range gKeys { 216 216 + updates[k] = gCost 217 217 + } 218 218 + 219 219 + if len(updates) == 0 { 220 220 + return nil 221 221 + } 222 222 + 223 223 + return l.kv.IncrByMulti(updates) 224 224 + } 225 225 + 226 226 + func (l *testRateLimiter) untilNextWindow(now time.Time, tier int) time.Duration { 227 227 + switch tier { 228 228 + case 0: 229 229 + return now.Truncate(time.Minute).Add(time.Minute).Sub(now) 230 230 + case 1: 231 231 + return now.Truncate(time.Hour).Add(time.Hour).Sub(now) 232 232 + case 2: 233 233 + return now.Truncate(24 * time.Hour).Add(24 * time.Hour).Sub(now) 234 234 + default: 235 235 + return time.Minute 236 236 + } 237 237 + } 238 238 + 67 239 func TestRateLimiter_Weighting(t *testing.T) { 68 240 kv := &mockKV{data: make(map[string]int)} 69 69 - limiter := NewRateLimiter(kv, 1) 241 241 + limiter := atproto.NewRateLimiter(kv, 1) 70 242 ctx := context.Background() 71 243 72 72 - // 1 Read = 1 Global 73 244 _, err := limiter.AllowRead(ctx) 74 245 if err != nil { 75 246 t.Fatal(err) ··· 82 253 t.Errorf("expected 1 global unit, got %d", g) 83 254 } 84 255 85 85 - // 1 Write = 1 Write-Only + 3 Global 86 256 _, err = limiter.AllowBulkWrite(ctx, 1) 87 257 if err != nil { 88 258 t.Fatal(err) ··· 94 264 if w != 1 { 95 265 t.Errorf("expected 1 write unit, got %d", w) 96 266 } 97 97 - if g != 4 { // 1 from read + 3 from write 267 267 + if g != 4 { 98 268 t.Errorf("expected 4 global units, got %d", g) 99 269 } 100 270 101 101 - // Bulk Write (10 elements) = 10 Write-Only + 30 Global 102 271 _, err = limiter.AllowBulkWrite(ctx, 10) 103 272 if err != nil { 104 273 t.Fatal(err) ··· 112 281 } 113 282 if g != 34 { 114 283 t.Errorf("expected 34 global units, got %d", g) 115 115 - } 116 116 - } 117 117 - 118 118 - func TestRateLimiter_Smoothing(t *testing.T) { 119 119 - kv := &mockKV{data: make(map[string]int)} 120 120 - // Ensure we are at the very beginning of the minute to avoid window edge issues in test 121 121 - clock := &mockClock{now: time.Date(2026, 1, 22, 1, 0, 0, 0, time.UTC)} 122 122 - limiter := &quotaLimiter{ 123 123 - kv: kv, 124 124 - prefix: "quota", 125 125 - clock: clock, 126 126 - rlQuota: 1, 127 127 - } 128 128 - 129 129 - wd, gd, wh, gh, wm, gm := limiter.getKeys(clock.now) 130 130 - kv.data[wm] = WriteLimitMinute 131 131 - kv.data[wh] = 0 132 132 - kv.data[wd] = 0 133 133 - kv.data[gm] = 0 134 134 - kv.data[gh] = 0 135 135 - kv.data[gd] = 0 136 136 - 137 137 - // Use a context that is already cancelled to simulate what happens 138 138 - // when we hit the rate limit and can't proceed. 139 139 - // Ensure the store thinks we are over the limit 140 140 - kv.data[wm] = WriteLimitMinute + 1 141 141 - 142 142 - ctx, cancel := context.WithCancel(context.Background()) 143 143 - cancel() 144 144 - 145 145 - _, err := limiter.AllowBulkWrite(ctx, 1) 146 146 - if err == nil { 147 147 - t.Errorf("expected error, got nil") 148 284 } 149 285 } 150 286 ··· 176 312 } 177 313 if res.ErrorCount != 1 { 178 314 t.Errorf("expected 1 error, got %d", res.ErrorCount) 179 179 - } 180 180 - 181 181 - storage.mu.Lock() 182 182 - errStr, ok := storage.failed["k1"] 183 183 - storage.mu.Unlock() 184 184 - 185 185 - if !ok { 186 186 - t.Error("expected record k1 to be marked as failed") 187 187 - } 188 188 - if !strings.Contains(errStr, "503") { 189 189 - t.Errorf("expected error string to contain 503, got %q", errStr) 190 190 - } 191 191 - } 192 192 - 193 193 - type mockClock struct { 194 194 - now time.Time 195 195 - } 196 196 - 197 197 - func (m *mockClock) Now() time.Time { return m.now } 198 198 - 199 199 - func TestRateLimiter_MidnightRollover(t *testing.T) { 200 200 - kv := &mockKV{data: make(map[string]int)} 201 201 - clock := &mockClock{now: time.Date(2026, 1, 22, 23, 59, 59, 0, time.UTC)} 202 202 - limiter := &quotaLimiter{ 203 203 - kv: kv, 204 204 - prefix: "quota", 205 205 - clock: clock, 206 206 - rlQuota: 1, 207 207 - } 208 208 - ctx := context.Background() 209 209 - 210 210 - // 1. Consume some quota on day 1 211 211 - _, err := limiter.AllowBulkWrite(ctx, 10) 212 212 - if err != nil { 213 213 - t.Fatal(err) 214 214 - } 215 215 - 216 216 - w1, g1, err := limiter.Stats() 217 217 - if err != nil { 218 218 - t.Fatal(err) 219 219 - } 220 220 - if w1 != 10 || g1 != 30 { 221 221 - t.Errorf("expected w=10, g=30 on day 1, got w=%d, g=%d", w1, g1) 222 222 - } 223 223 - 224 224 - // 2. Advance time to next day 225 225 - clock.now = clock.now.Add(2 * time.Second) // 00:00:01 on 2026-01-23 226 226 - 227 227 - // 3. Stats should now reflect day 2 (0) 228 228 - w2, g2, err := limiter.Stats() 229 229 - if err != nil { 230 230 - t.Fatal(err) 231 231 - } 232 232 - if w2 != 0 || g2 != 0 { 233 233 - t.Errorf("expected w=0, g=0 on day 2, got w=%d, g=%d", w2, g2) 234 234 - } 235 235 - 236 236 - // 4. Consumption on day 2 should not affect day 1 237 237 - _, err = limiter.AllowBulkWrite(ctx, 5) 238 238 - if err != nil { 239 239 - t.Fatal(err) 240 240 - } 241 241 - 242 242 - w2, g2, err = limiter.Stats() 243 243 - if err != nil { 244 244 - t.Fatal(err) 245 245 - } 246 246 - if w2 != 5 || g2 != 15 { 247 247 - t.Errorf("expected w=5, g=15 on day 2, got w=%d, g=%d", w2, g2) 248 248 - } 249 249 - 250 250 - // Verify day 1 keys are still there but not accessed by Stats() 251 251 - day1WKey := "quota:writes:d:2026-01-22" 252 252 - if val, _ := kv.Get(day1WKey); val != 10 { 253 253 - t.Errorf("expected day 1 write key to still be 10, got %d", val) 254 315 } 255 316 }

+29

sync/record.go

reviewed

··· 287 287 } 288 288 return duplicates 289 289 } 290 290 + 291 291 + type Timestamp struct { 292 292 + time.Time 293 293 + } 294 294 + 295 295 + func (t Timestamp) MarshalJSON() ([]byte, error) { 296 296 + return []byte(`"` + t.Format(time.RFC3339Nano) + `"`), nil 297 297 + } 298 298 + 299 299 + func (t *Timestamp) UnmarshalJSON(data []byte) error { 300 300 + if string(data) == "null" { 301 301 + *t = Timestamp{} 302 302 + return nil 303 303 + } 304 304 + 305 305 + s := string(data) 306 306 + s = s[1 : len(s)-1] 307 307 + 308 308 + tm, err := time.Parse(time.RFC3339Nano, s) 309 309 + if err != nil { 310 310 + tm, err = time.Parse(time.RFC3339, s) 311 311 + if err != nil { 312 312 + return fmt.Errorf("failed to parse timestamp %q: %w", s, err) 313 313 + } 314 314 + } 315 315 + 316 316 + *t = Timestamp{Time: tm} 317 317 + return nil 318 318 + }

-35

sync/timestamp.go

reviewed

··· 1 1 - package sync 2 2 - 3 3 - import ( 4 4 - "fmt" 5 5 - "time" 6 6 - ) 7 7 - 8 8 - type Timestamp struct { 9 9 - time.Time 10 10 - } 11 11 - 12 12 - func (t Timestamp) MarshalJSON() ([]byte, error) { 13 13 - return []byte(`"` + t.Format(time.RFC3339Nano) + `"`), nil 14 14 - } 15 15 - 16 16 - func (t *Timestamp) UnmarshalJSON(data []byte) error { 17 17 - if string(data) == "null" { 18 18 - *t = Timestamp{} 19 19 - return nil 20 20 - } 21 21 - 22 22 - s := string(data) 23 23 - s = s[1 : len(s)-1] 24 24 - 25 25 - tm, err := time.Parse(time.RFC3339Nano, s) 26 26 - if err != nil { 27 27 - tm, err = time.Parse(time.RFC3339, s) 28 28 - if err != nil { 29 29 - return fmt.Errorf("failed to parse timestamp %q: %w", s, err) 30 30 - } 31 31 - } 32 32 - 33 33 - *t = Timestamp{Time: tm} 34 34 - return nil 35 35 - }