willdot.net / at-container-registry
forked from evan.jarrett.net/at-container-registry
A container registry that uses the AT Protocol for manifest storage and S3 for blob storage.
0
fork

Configure Feed

Select the types of activity you want to include in your feed.

optimize queries for admin panel

Evan Jarrett 22d53965 8e45b2ee

+324 -177
hold

This is a binary file and will not be displayed.

+1 -3
pkg/appview/db/hold_store.go
··· 290 290 WHEN h.owner_did = ?1 THEN 'owner' 291 291 WHEN c.member_did IS NOT NULL THEN 'crew' 292 292 WHEN h.allow_all_crew = 1 THEN 'eligible' 293 - WHEN h.public = 1 THEN 'public' 294 293 ELSE 'none' 295 294 END as membership, 296 295 c.permissions 297 296 FROM hold_captain_records h 298 297 LEFT JOIN hold_crew_members c ON h.hold_did = c.hold_did AND c.member_did = ?1 299 - WHERE h.public = 1 300 - OR h.allow_all_crew = 1 298 + WHERE h.allow_all_crew = 1 301 299 OR h.owner_did = ?1 302 300 OR c.member_did IS NOT NULL 303 301 ORDER BY
+3 -5
pkg/appview/handlers/settings.go
··· 58 58 slog.Debug("Fetched profile", "component", "settings", "did", user.DID, "default_hold", profile.DefaultHold) 59 59 60 60 // Get available holds for dropdown 61 - var ownedHolds, crewHolds, eligibleHolds, publicHolds []HoldDisplay 61 + var ownedHolds, crewHolds, eligibleHolds []HoldDisplay 62 62 holdDataMap := make(map[string]HoldDisplay) 63 63 64 64 if h.DB != nil { ··· 93 93 crewHolds = append(crewHolds, display) 94 94 case "eligible": 95 95 eligibleHolds = append(eligibleHolds, display) 96 - case "public": 97 - publicHolds = append(publicHolds, display) 98 96 } 99 97 } 100 98 } ··· 134 132 CurrentHoldDID string 135 133 CurrentHoldDisplay string 136 134 ShowCurrentHold bool 135 + AppViewDefaultHoldDID string 137 136 AppViewDefaultHoldDisplay string 138 137 AppViewDefaultRegion string 139 138 OwnedHolds []HoldDisplay 140 139 CrewHolds []HoldDisplay 141 140 EligibleHolds []HoldDisplay 142 - PublicHolds []HoldDisplay 143 141 HoldDataJSON template.JS 144 142 }{ 145 143 PageData: NewPageData(r, &h.BaseUIHandler), ··· 147 145 CurrentHoldDID: profile.DefaultHold, 148 146 CurrentHoldDisplay: deriveDisplayName(profile.DefaultHold), 149 147 ShowCurrentHold: showCurrentHold, 148 + AppViewDefaultHoldDID: h.DefaultHoldDID, 150 149 AppViewDefaultHoldDisplay: appViewDefaultDisplay, 151 150 AppViewDefaultRegion: appViewDefaultRegion, 152 151 OwnedHolds: ownedHolds, 153 152 CrewHolds: crewHolds, 154 153 EligibleHolds: eligibleHolds, 155 - PublicHolds: publicHolds, 156 154 HoldDataJSON: template.JS(holdDataJSON), 157 155 } 158 156
+1 -10
pkg/appview/templates/pages/settings.html
··· 142 142 <span class="label-text">Storage Hold</span> 143 143 </label> 144 144 <select id="default-hold" name="hold_did" class="select select-bordered w-full" autocomplete="off"> 145 - <option value=""{{ if eq .CurrentHoldDID "" }} selected{{ end }}>AppView Default ({{ .AppViewDefaultHoldDisplay }}{{ if .AppViewDefaultRegion }}, {{ .AppViewDefaultRegion }}{{ end }})</option> 145 + <option value="{{ .AppViewDefaultHoldDID }}"{{ if or (eq .CurrentHoldDID "") (eq .CurrentHoldDID .AppViewDefaultHoldDID) }} selected{{ end }}>AppView Default ({{ .AppViewDefaultHoldDisplay }}{{ if .AppViewDefaultRegion }}, {{ .AppViewDefaultRegion }}{{ end }})</option> 146 146 147 147 {{ if .ShowCurrentHold }} 148 148 <option value="{{ .CurrentHoldDID }}" selected>Current ({{ .CurrentHoldDisplay }})</option> ··· 178 178 </optgroup> 179 179 {{ end }} 180 180 181 - {{ if .PublicHolds }} 182 - <optgroup label="Public Holds"> 183 - {{ range .PublicHolds }} 184 - <option value="{{ .DID }}" {{ if eq $.CurrentHoldDID .DID }}selected{{ end }}> 185 - {{ .DisplayName }}{{ if .Region }} ({{ .Region }}){{ end }} 186 - </option> 187 - {{ end }} 188 - </optgroup> 189 - {{ end }} 190 181 </select> 191 182 <p class="text-sm text-base-content/60 mt-1">Your images will be stored on the selected hold</p> 192 183 </fieldset>
+15 -26
pkg/hold/admin/handlers.go
··· 93 93 return 94 94 } 95 95 96 - // Calculate total storage by summing quota for all crew members 96 + // Calculate total storage with a single bulk query 97 97 var totalSize int64 98 98 uniqueDigests := 0 99 99 100 - crew, err := ui.pds.ListCrewMembers(ctx) 100 + allQuotas, err := ui.pds.GetAllUserQuotas(ctx) 101 101 if err != nil { 102 - slog.Warn("Failed to list crew for stats", "error", err) 102 + slog.Warn("Failed to get all user quotas", "error", err) 103 103 } else { 104 - // Get usage for each crew member 105 - for _, member := range crew { 106 - quotaStats, err := ui.pds.GetQuotaForUser(ctx, member.Record.Member) 107 - if err != nil { 108 - continue 109 - } 110 - totalSize += quotaStats.TotalSize 111 - uniqueDigests += quotaStats.UniqueBlobs 104 + for _, q := range allQuotas { 105 + totalSize += q.TotalSize 106 + uniqueDigests += q.UniqueBlobs 112 107 } 113 108 } 114 109 ··· 152 147 } 153 148 } 154 149 155 - // Get all crew members and their usage 156 - crew, err := ui.pds.ListCrewMembers(ctx) 150 + // Get all user quotas in a single bulk query 151 + allQuotas, err := ui.pds.GetAllUserQuotas(ctx) 157 152 if err != nil { 158 - slog.Error("Failed to list crew members for top users", "error", err) 153 + slog.Error("Failed to get all user quotas", "error", err) 159 154 http.Error(w, "Failed to load top users", http.StatusInternalServerError) 160 155 return 161 156 } 162 157 163 158 var users []UserUsage 164 - for _, member := range crew { 165 - quotaStats, err := ui.pds.GetQuotaForUser(ctx, member.Record.Member) 166 - if err != nil { 167 - slog.Warn("Failed to get quota for user", "did", member.Record.Member, "error", err) 168 - continue 169 - } 170 - 159 + for did, q := range allQuotas { 171 160 users = append(users, UserUsage{ 172 - DID: member.Record.Member, 173 - Handle: resolveHandle(ctx, member.Record.Member), 174 - Usage: quotaStats.TotalSize, 175 - UsageHuman: formatHumanBytes(quotaStats.TotalSize), 176 - BlobCount: quotaStats.UniqueBlobs, 161 + DID: did, 162 + Handle: resolveHandle(ctx, did), 163 + Usage: q.TotalSize, 164 + UsageHuman: formatHumanBytes(q.TotalSize), 165 + BlobCount: q.UniqueBlobs, 177 166 }) 178 167 } 179 168
+13 -10
pkg/hold/admin/handlers_crew.go
··· 9 9 "time" 10 10 11 11 "atcr.io/pkg/atproto" 12 + "atcr.io/pkg/hold/pds" 12 13 "github.com/go-chi/chi/v5" 13 14 ) 14 15 ··· 56 57 return nil, err 57 58 } 58 59 59 - userUsage := make(map[string]int64) 60 - for _, member := range crew { 61 - quotaStats, err := ui.pds.GetQuotaForUser(ctx, member.Record.Member) 62 - if err != nil { 63 - slog.Warn("Failed to get quota for crew member", "did", member.Record.Member, "error", err) 64 - continue 65 - } 66 - userUsage[member.Record.Member] = quotaStats.TotalSize 60 + // Single bulk query for all user quotas 61 + allQuotas, err := ui.pds.GetAllUserQuotas(ctx) 62 + if err != nil { 63 + slog.Warn("Failed to get all user quotas for crew views", "error", err) 64 + allQuotas = make(map[string]*pds.QuotaStats) 67 65 } 68 66 69 67 defaultTier := "default" ··· 88 86 AddedAt: parseTime(member.Record.AddedAt), 89 87 } 90 88 89 + usage := int64(0) 90 + if q, ok := allQuotas[member.Record.Member]; ok { 91 + usage = q.TotalSize 92 + } 93 + 91 94 if ui.quotaMgr != nil && ui.quotaMgr.IsEnabled() { 92 95 if limit := ui.quotaMgr.GetTierLimit(tier); limit != nil { 93 96 view.TierLimit = formatHumanBytes(*limit) 94 97 if *limit > 0 { 95 - view.UsagePercent = int(float64(userUsage[view.DID]) / float64(*limit) * 100) 98 + view.UsagePercent = int(float64(usage) / float64(*limit) * 100) 96 99 } 97 100 } else { 98 101 view.TierLimit = "Unlimited" ··· 101 104 view.TierLimit = "Unlimited" 102 105 } 103 106 104 - view.CurrentUsage = userUsage[view.DID] 107 + view.CurrentUsage = usage 105 108 view.UsageHuman = formatHumanBytes(view.CurrentUsage) 106 109 107 110 crewViews = append(crewViews, view)
+2 -2
pkg/hold/pds/delete_test.go
··· 215 215 216 216 // Index the record 217 217 if pds.recordsIndex != nil { 218 - err = pds.recordsIndex.IndexRecord(atproto.BskyPostCollection, rkey, "testcid", aliceDID) 218 + err = pds.recordsIndex.IndexRecord(atproto.BskyPostCollection, rkey, "testcid", aliceDID, "", 0) 219 219 if err != nil { 220 220 t.Fatalf("Failed to index test post: %v", err) 221 221 } ··· 304 304 305 305 // Index the record 306 306 if pds.recordsIndex != nil { 307 - err = pds.recordsIndex.IndexRecord(atproto.BskyPostCollection, rkey, "testcid", aliceDID) 307 + err = pds.recordsIndex.IndexRecord(atproto.BskyPostCollection, rkey, "testcid", aliceDID, "", 0) 308 308 if err != nil { 309 309 t.Fatalf("Failed to index test post: %v", err) 310 310 }
+27 -87
pkg/hold/pds/layer.go
··· 89 89 Tier string `json:"tier,omitempty"` // quota tier (e.g., 'deckhand', 'bosun', 'quartermaster') 90 90 } 91 91 92 - // GetQuotaForUser calculates storage quota for a specific user 93 - // It iterates through all layer records, filters by userDid, deduplicates by digest, 94 - // and sums the sizes of unique blobs. 92 + // GetQuotaForUser calculates storage quota for a specific user. 93 + // Uses SQL aggregation over the denormalized digest/size columns in the records index. 95 94 func (p *HoldPDS) GetQuotaForUser(ctx context.Context, userDID string) (*QuotaStats, error) { 96 95 if p.recordsIndex == nil { 97 96 return nil, fmt.Errorf("records index not available") 98 97 } 99 98 100 - // Get session for reading record data 101 - session, err := p.carstore.ReadOnlySession(p.uid) 99 + uniqueBlobs, totalSize, err := p.recordsIndex.QuotaForDID(atproto.LayerCollection, userDID) 102 100 if err != nil { 103 - return nil, fmt.Errorf("failed to create session: %w", err) 101 + return nil, fmt.Errorf("failed to query quota: %w", err) 104 102 } 105 103 106 - head, err := p.carstore.GetUserRepoHead(ctx, p.uid) 107 - if err != nil { 108 - return nil, fmt.Errorf("failed to get repo head: %w", err) 109 - } 104 + return &QuotaStats{ 105 + UserDID: userDID, 106 + UniqueBlobs: uniqueBlobs, 107 + TotalSize: totalSize, 108 + }, nil 109 + } 110 110 111 - if !head.Defined() { 112 - // Empty repo - return zero stats 113 - return &QuotaStats{UserDID: userDID}, nil 111 + // GetAllUserQuotas returns quota stats for all users in a single SQL query. 112 + // Used by admin endpoints to avoid N+1 per-user quota lookups. 113 + func (p *HoldPDS) GetAllUserQuotas(ctx context.Context) (map[string]*QuotaStats, error) { 114 + if p.recordsIndex == nil { 115 + return nil, fmt.Errorf("records index not available") 114 116 } 115 117 116 - repoHandle, err := repo.OpenRepo(ctx, session, head) 118 + quotas, err := p.recordsIndex.QuotasByDID(atproto.LayerCollection) 117 119 if err != nil { 118 - return nil, fmt.Errorf("failed to open repo: %w", err) 120 + return nil, fmt.Errorf("failed to query all quotas: %w", err) 119 121 } 120 122 121 - // Track unique digests and their sizes 122 - digestSizes := make(map[string]int64) 123 - 124 - // Iterate all layer records via the index 125 - cursor := "" 126 - batchSize := 1000 // Process in batches 127 - 128 - for { 129 - records, nextCursor, err := p.recordsIndex.ListRecords(atproto.LayerCollection, batchSize, cursor, true) 130 - if err != nil { 131 - return nil, fmt.Errorf("failed to list layer records: %w", err) 132 - } 133 - 134 - for _, rec := range records { 135 - // Construct record path and get the record data 136 - recordPath := rec.Collection + "/" + rec.Rkey 137 - 138 - _, recBytes, err := repoHandle.GetRecordBytes(ctx, recordPath) 139 - if err != nil { 140 - // Skip records we can't read 141 - continue 142 - } 143 - 144 - // Decode the layer record 145 - recordValue, err := lexutil.CborDecodeValue(*recBytes) 146 - if err != nil { 147 - continue 148 - } 149 - 150 - layerRecord, ok := recordValue.(*atproto.LayerRecord) 151 - if !ok { 152 - continue 153 - } 154 - 155 - // Filter by userDID 156 - if layerRecord.UserDID != userDID { 157 - continue 158 - } 159 - 160 - // Deduplicate by digest - keep the size (could be different pushes of same blob) 161 - // Store the size - we only count each unique digest once 162 - if _, exists := digestSizes[layerRecord.Digest]; !exists { 163 - digestSizes[layerRecord.Digest] = layerRecord.Size 164 - } 123 + result := make(map[string]*QuotaStats, len(quotas)) 124 + for did, q := range quotas { 125 + result[did] = &QuotaStats{ 126 + UserDID: did, 127 + UniqueBlobs: q.UniqueBlobs, 128 + TotalSize: q.TotalSize, 165 129 } 166 - 167 - if nextCursor == "" { 168 - break 169 - } 170 - cursor = nextCursor 171 130 } 172 - 173 - // Calculate totals 174 - var totalSize int64 175 - for _, size := range digestSizes { 176 - totalSize += size 177 - } 178 - 179 - return &QuotaStats{ 180 - UserDID: userDID, 181 - UniqueBlobs: len(digestSizes), 182 - TotalSize: totalSize, 183 - }, nil 131 + return result, nil 184 132 } 185 133 186 134 // GetQuotaForUserWithTier calculates quota with tier-aware limits ··· 262 210 263 211 var records []*atproto.LayerRecord 264 212 265 - // Iterate all layer records via the index 213 + // Iterate layer records for this user via the index (filtered by DID in SQL) 266 214 cursor := "" 267 - batchSize := 1000 // Process in batches 215 + batchSize := 1000 268 216 269 217 for { 270 - indexRecords, nextCursor, err := p.recordsIndex.ListRecords(atproto.LayerCollection, batchSize, cursor, true) 218 + indexRecords, nextCursor, err := p.recordsIndex.ListRecordsByDID(atproto.LayerCollection, userDID, batchSize, cursor) 271 219 if err != nil { 272 220 return nil, fmt.Errorf("failed to list layer records: %w", err) 273 221 } 274 222 275 223 for _, rec := range indexRecords { 276 - // Construct record path and get the record data 277 224 recordPath := rec.Collection + "/" + rec.Rkey 278 225 279 226 _, recBytes, err := repoHandle.GetRecordBytes(ctx, recordPath) 280 227 if err != nil { 281 - // Skip records we can't read 282 228 continue 283 229 } 284 230 285 - // Decode the layer record 286 231 recordValue, err := lexutil.CborDecodeValue(*recBytes) 287 232 if err != nil { 288 233 continue ··· 290 235 291 236 layerRecord, ok := recordValue.(*atproto.LayerRecord) 292 237 if !ok { 293 - continue 294 - } 295 - 296 - // Filter by userDID 297 - if layerRecord.UserDID != userDID { 298 238 continue 299 239 } 300 240
+94 -13
pkg/hold/pds/records.go
··· 36 36 rkey TEXT NOT NULL, 37 37 cid TEXT NOT NULL, 38 38 did TEXT, 39 + digest TEXT, 40 + size INTEGER, 39 41 PRIMARY KEY (collection, rkey) 40 42 ); 41 43 CREATE INDEX IF NOT EXISTS idx_records_collection_rkey ON records(collection, rkey); ··· 54 56 return nil, fmt.Errorf("failed to open records database: %w", err) 55 57 } 56 58 57 - // Check if table exists and has the did column 59 + // Check if table exists and has required columns 58 60 needsRebuild := false 59 61 var tableName string 60 62 err = db.QueryRow(`SELECT name FROM sqlite_master WHERE type='table' AND name='records'`).Scan(&tableName) 61 63 if err == nil { 62 - // Table exists, check for did column 63 64 var colCount int 65 + // Check for did column 64 66 err = db.QueryRow(`SELECT COUNT(*) FROM pragma_table_info('records') WHERE name='did'`).Scan(&colCount) 65 67 if err != nil || colCount == 0 { 66 68 needsRebuild = true 67 69 slog.Info("Records index schema outdated, rebuilding with did column") 68 70 } 71 + // Check for digest column (added for SQL-based quota queries) 72 + if !needsRebuild { 73 + err = db.QueryRow(`SELECT COUNT(*) FROM pragma_table_info('records') WHERE name='digest'`).Scan(&colCount) 74 + if err != nil || colCount == 0 { 75 + needsRebuild = true 76 + slog.Info("Records index schema outdated, rebuilding with digest/size columns") 77 + } 78 + } 69 79 } 70 80 71 81 if needsRebuild { ··· 91 101 // NewRecordsIndexWithDB creates a records index using an existing *sql.DB connection. 92 102 // The caller is responsible for the DB lifecycle. 93 103 func NewRecordsIndexWithDB(db *sql.DB) (*RecordsIndex, error) { 94 - // Check if table exists and has the did column 104 + // Check if table exists and has required columns 95 105 needsRebuild := false 96 106 var tableName string 97 107 err := db.QueryRow(`SELECT name FROM sqlite_master WHERE type='table' AND name='records'`).Scan(&tableName) 98 108 if err == nil { 99 109 var colCount int 110 + // Check for did column 100 111 err = db.QueryRow(`SELECT COUNT(*) FROM pragma_table_info('records') WHERE name='did'`).Scan(&colCount) 101 112 if err != nil || colCount == 0 { 102 113 needsRebuild = true 103 114 slog.Info("Records index schema outdated, rebuilding with did column") 104 115 } 116 + // Check for digest column (added for SQL-based quota queries) 117 + if !needsRebuild { 118 + err = db.QueryRow(`SELECT COUNT(*) FROM pragma_table_info('records') WHERE name='digest'`).Scan(&colCount) 119 + if err != nil || colCount == 0 { 120 + needsRebuild = true 121 + slog.Info("Records index schema outdated, rebuilding with digest/size columns") 122 + } 123 + } 105 124 } 106 125 107 126 if needsRebuild { ··· 129 148 } 130 149 131 150 // IndexRecord adds or updates a record in the index 132 - // did parameter is optional - pass empty string if not applicable 133 - func (ri *RecordsIndex) IndexRecord(collection, rkey, cidStr, did string) error { 151 + // did, digest, size parameters are optional - pass empty/zero if not applicable 152 + func (ri *RecordsIndex) IndexRecord(collection, rkey, cidStr, did, digest string, size int64) error { 134 153 _, err := ri.db.Exec(` 135 - INSERT OR REPLACE INTO records (collection, rkey, cid, did) 136 - VALUES (?, ?, ?, ?) 137 - `, collection, rkey, cidStr, sql.NullString{String: did, Valid: did != ""}) 154 + INSERT OR REPLACE INTO records (collection, rkey, cid, did, digest, size) 155 + VALUES (?, ?, ?, ?, ?, ?) 156 + `, collection, rkey, cidStr, 157 + sql.NullString{String: did, Valid: did != ""}, 158 + sql.NullString{String: digest, Valid: digest != ""}, 159 + sql.NullInt64{Int64: size, Valid: size > 0}) 138 160 return err 139 161 } 140 162 ··· 328 350 defer tx.Rollback() 329 351 330 352 stmt, err := tx.Prepare(` 331 - INSERT OR REPLACE INTO records (collection, rkey, cid, did) 332 - VALUES (?, ?, ?, ?) 353 + INSERT OR REPLACE INTO records (collection, rkey, cid, did, digest, size) 354 + VALUES (?, ?, ?, ?, ?, ?) 333 355 `) 334 356 if err != nil { 335 357 return fmt.Errorf("failed to prepare statement: %w", err) ··· 345 367 } 346 368 collection, rkey := parts[0], parts[1] 347 369 348 - // Extract DID from record content based on collection type 349 - var did string 370 + // Extract fields from record content based on collection type 371 + var did, digest string 372 + var size int64 350 373 _, recBytes, err := repoHandle.GetRecordBytes(ctx, key) 351 374 if err == nil && recBytes != nil { 352 375 did = extractDIDFromRecord(collection, *recBytes) 376 + digest, size = extractLayerFieldsFromRecord(collection, *recBytes) 353 377 } 354 378 355 - _, err = stmt.Exec(collection, rkey, c.String(), sql.NullString{String: did, Valid: did != ""}) 379 + _, err = stmt.Exec(collection, rkey, c.String(), 380 + sql.NullString{String: did, Valid: did != ""}, 381 + sql.NullString{String: digest, Valid: digest != ""}, 382 + sql.NullInt64{Int64: size, Valid: size > 0}) 356 383 if err != nil { 357 384 return fmt.Errorf("failed to index record %s: %w", key, err) 358 385 } ··· 390 417 } 391 418 } 392 419 return out 420 + } 421 + 422 + // QuotaForDID returns unique blob count and total size for a single DID in a collection. 423 + // Uses SQL aggregation over the denormalized digest/size columns. 424 + func (ri *RecordsIndex) QuotaForDID(collection, did string) (uniqueBlobs int, totalSize int64, err error) { 425 + err = ri.db.QueryRow(` 426 + SELECT COUNT(*), COALESCE(SUM(size), 0) 427 + FROM (SELECT DISTINCT digest, size FROM records WHERE collection = ? AND did = ? AND digest IS NOT NULL) 428 + `, collection, did).Scan(&uniqueBlobs, &totalSize) 429 + return 430 + } 431 + 432 + // QuotasByDID returns unique blob count and total size grouped by DID. 433 + // Single query replaces N individual quota lookups. 434 + func (ri *RecordsIndex) QuotasByDID(collection string) (map[string]QuotaResult, error) { 435 + rows, err := ri.db.Query(` 436 + SELECT did, COUNT(*), COALESCE(SUM(size), 0) 437 + FROM (SELECT DISTINCT did, digest, size FROM records WHERE collection = ? AND did IS NOT NULL AND digest IS NOT NULL) 438 + GROUP BY did 439 + `, collection) 440 + if err != nil { 441 + return nil, fmt.Errorf("failed to query quotas by DID: %w", err) 442 + } 443 + defer rows.Close() 444 + 445 + result := make(map[string]QuotaResult) 446 + for rows.Next() { 447 + var did string 448 + var qr QuotaResult 449 + if err := rows.Scan(&did, &qr.UniqueBlobs, &qr.TotalSize); err != nil { 450 + return nil, fmt.Errorf("failed to scan quota row: %w", err) 451 + } 452 + result[did] = qr 453 + } 454 + return result, rows.Err() 455 + } 456 + 457 + // QuotaResult holds aggregated quota data from SQL queries 458 + type QuotaResult struct { 459 + UniqueBlobs int 460 + TotalSize int64 461 + } 462 + 463 + // extractLayerFieldsFromRecord extracts digest and size from layer records. 464 + // Returns empty/zero for non-layer records. 465 + func extractLayerFieldsFromRecord(collection string, recBytes []byte) (string, int64) { 466 + if collection != atproto.LayerCollection { 467 + return "", 0 468 + } 469 + var rec atproto.LayerRecord 470 + if err := rec.UnmarshalCBOR(bytes.NewReader(recBytes)); err != nil { 471 + return "", 0 472 + } 473 + return rec.Digest, rec.Size 393 474 } 394 475 395 476 // extractDIDFromRecord extracts the associated DID from a record based on its collection type
+154 -19
pkg/hold/pds/records_test.go
··· 50 50 defer ri.Close() 51 51 52 52 // Index a record 53 - err = ri.IndexRecord("io.atcr.hold.crew", "abc123", "bafyrei123", "") 53 + err = ri.IndexRecord("io.atcr.hold.crew", "abc123", "bafyrei123", "", "", 0) 54 54 if err != nil { 55 55 t.Fatalf("IndexRecord() error = %v", err) 56 56 } ··· 75 75 defer ri.Close() 76 76 77 77 // Index a record 78 - err = ri.IndexRecord("io.atcr.hold.crew", "abc123", "bafyrei123", "") 78 + err = ri.IndexRecord("io.atcr.hold.crew", "abc123", "bafyrei123", "", "", 0) 79 79 if err != nil { 80 80 t.Fatalf("IndexRecord() first call error = %v", err) 81 81 } 82 82 83 83 // Update the same record with new CID 84 - err = ri.IndexRecord("io.atcr.hold.crew", "abc123", "bafyrei456", "") 84 + err = ri.IndexRecord("io.atcr.hold.crew", "abc123", "bafyrei456", "", "", 0) 85 85 if err != nil { 86 86 t.Fatalf("IndexRecord() second call error = %v", err) 87 87 } ··· 118 118 defer ri.Close() 119 119 120 120 // Index a record 121 - err = ri.IndexRecord("io.atcr.hold.crew", "abc123", "bafyrei123", "") 121 + err = ri.IndexRecord("io.atcr.hold.crew", "abc123", "bafyrei123", "", "", 0) 122 122 if err != nil { 123 123 t.Fatalf("IndexRecord() error = %v", err) 124 124 } ··· 217 217 {"ccc", "cid3"}, 218 218 } 219 219 for _, r := range records { 220 - if err := ri.IndexRecord("io.atcr.hold.crew", r.rkey, r.cid, ""); err != nil { 220 + if err := ri.IndexRecord("io.atcr.hold.crew", r.rkey, r.cid, "", "", 0); err != nil { 221 221 t.Fatalf("IndexRecord() error = %v", err) 222 222 } 223 223 } ··· 248 248 // Add records with different rkeys (TIDs are lexicographically ordered by time) 249 249 rkeys := []string{"3m3aaaaaaaaa", "3m3bbbbbbbbb", "3m3ccccccccc"} 250 250 for _, rkey := range rkeys { 251 - if err := ri.IndexRecord("io.atcr.hold.crew", rkey, "cid-"+rkey, ""); err != nil { 251 + if err := ri.IndexRecord("io.atcr.hold.crew", rkey, "cid-"+rkey, "", "", 0); err != nil { 252 252 t.Fatalf("IndexRecord() error = %v", err) 253 253 } 254 254 } ··· 286 286 // Add records 287 287 rkeys := []string{"3m3aaaaaaaaa", "3m3bbbbbbbbb", "3m3ccccccccc"} 288 288 for _, rkey := range rkeys { 289 - if err := ri.IndexRecord("io.atcr.hold.crew", rkey, "cid-"+rkey, ""); err != nil { 289 + if err := ri.IndexRecord("io.atcr.hold.crew", rkey, "cid-"+rkey, "", "", 0); err != nil { 290 290 t.Fatalf("IndexRecord() error = %v", err) 291 291 } 292 292 } ··· 324 324 // Add 5 records 325 325 for i := range 5 { 326 326 rkey := string(rune('a' + i)) 327 - if err := ri.IndexRecord("io.atcr.hold.crew", rkey, "cid-"+rkey, ""); err != nil { 327 + if err := ri.IndexRecord("io.atcr.hold.crew", rkey, "cid-"+rkey, "", "", 0); err != nil { 328 328 t.Fatalf("IndexRecord() error = %v", err) 329 329 } 330 330 } ··· 355 355 // Add 5 records 356 356 rkeys := []string{"a", "b", "c", "d", "e"} 357 357 for _, rkey := range rkeys { 358 - if err := ri.IndexRecord("io.atcr.hold.crew", rkey, "cid-"+rkey, ""); err != nil { 358 + if err := ri.IndexRecord("io.atcr.hold.crew", rkey, "cid-"+rkey, "", "", 0); err != nil { 359 359 t.Fatalf("IndexRecord() error = %v", err) 360 360 } 361 361 } ··· 430 430 // Add 5 records 431 431 rkeys := []string{"a", "b", "c", "d", "e"} 432 432 for _, rkey := range rkeys { 433 - if err := ri.IndexRecord("io.atcr.hold.crew", rkey, "cid-"+rkey, ""); err != nil { 433 + if err := ri.IndexRecord("io.atcr.hold.crew", rkey, "cid-"+rkey, "", "", 0); err != nil { 434 434 t.Fatalf("IndexRecord() error = %v", err) 435 435 } 436 436 } ··· 474 474 475 475 // Add records to two collections 476 476 for i := range 3 { 477 - ri.IndexRecord("io.atcr.hold.crew", string(rune('a'+i)), "cid1", "") 477 + ri.IndexRecord("io.atcr.hold.crew", string(rune('a'+i)), "cid1", "", "", 0) 478 478 } 479 479 for i := range 5 { 480 - ri.IndexRecord("io.atcr.hold.captain", string(rune('a'+i)), "cid2", "") 480 + ri.IndexRecord("io.atcr.hold.captain", string(rune('a'+i)), "cid2", "", "", 0) 481 481 } 482 482 483 483 // Count crew ··· 527 527 defer ri.Close() 528 528 529 529 // Add records to multiple collections 530 - ri.IndexRecord("io.atcr.hold.crew", "a", "cid1", "") 531 - ri.IndexRecord("io.atcr.hold.crew", "b", "cid2", "") 532 - ri.IndexRecord("io.atcr.hold.captain", "self", "cid3", "") 533 - ri.IndexRecord("io.atcr.manifest", "abc123", "cid4", "") 530 + ri.IndexRecord("io.atcr.hold.crew", "a", "cid1", "", "", 0) 531 + ri.IndexRecord("io.atcr.hold.crew", "b", "cid2", "", "", 0) 532 + ri.IndexRecord("io.atcr.hold.captain", "self", "cid3", "", "", 0) 533 + ri.IndexRecord("io.atcr.manifest", "abc123", "cid4", "", "", 0) 534 534 535 535 count, err := ri.TotalCount() 536 536 if err != nil { ··· 581 581 defer ri.Close() 582 582 583 583 // Add records to different collections with same rkeys 584 - ri.IndexRecord("io.atcr.hold.crew", "abc", "cid-crew", "") 585 - ri.IndexRecord("io.atcr.hold.captain", "abc", "cid-captain", "") 586 - ri.IndexRecord("io.atcr.manifest", "abc", "cid-manifest", "") 584 + ri.IndexRecord("io.atcr.hold.crew", "abc", "cid-crew", "", "", 0) 585 + ri.IndexRecord("io.atcr.hold.captain", "abc", "cid-captain", "", "", 0) 586 + ri.IndexRecord("io.atcr.manifest", "abc", "cid-manifest", "", "", 0) 587 587 588 588 // Listing should only return records from requested collection 589 589 records, _, err := ri.ListRecords("io.atcr.hold.crew", 10, "", false) ··· 605 605 t.Errorf("Expected captain count 1 after deleting crew, got %d", count) 606 606 } 607 607 } 608 + 609 + // TestRecordsIndex_QuotaForDID tests single-user quota aggregation 610 + func TestRecordsIndex_QuotaForDID(t *testing.T) { 611 + tmpDir := t.TempDir() 612 + ri, err := NewRecordsIndex(filepath.Join(tmpDir, "records.db")) 613 + if err != nil { 614 + t.Fatalf("NewRecordsIndex() error = %v", err) 615 + } 616 + defer ri.Close() 617 + 618 + // Add layer records for two users, with some duplicate digests 619 + ri.IndexRecord("io.atcr.hold.layer", "r1", "cid1", "did:plc:alice", "sha256:aaa", 100) 620 + ri.IndexRecord("io.atcr.hold.layer", "r2", "cid2", "did:plc:alice", "sha256:bbb", 200) 621 + ri.IndexRecord("io.atcr.hold.layer", "r3", "cid3", "did:plc:alice", "sha256:aaa", 100) // duplicate digest 622 + ri.IndexRecord("io.atcr.hold.layer", "r4", "cid4", "did:plc:bob", "sha256:ccc", 300) 623 + 624 + // Alice should have 2 unique blobs, total 300 bytes 625 + uniqueBlobs, totalSize, err := ri.QuotaForDID("io.atcr.hold.layer", "did:plc:alice") 626 + if err != nil { 627 + t.Fatalf("QuotaForDID() error = %v", err) 628 + } 629 + if uniqueBlobs != 2 { 630 + t.Errorf("Expected 2 unique blobs for alice, got %d", uniqueBlobs) 631 + } 632 + if totalSize != 300 { 633 + t.Errorf("Expected total size 300 for alice, got %d", totalSize) 634 + } 635 + 636 + // Bob should have 1 unique blob, total 300 bytes 637 + uniqueBlobs, totalSize, err = ri.QuotaForDID("io.atcr.hold.layer", "did:plc:bob") 638 + if err != nil { 639 + t.Fatalf("QuotaForDID() error = %v", err) 640 + } 641 + if uniqueBlobs != 1 { 642 + t.Errorf("Expected 1 unique blob for bob, got %d", uniqueBlobs) 643 + } 644 + if totalSize != 300 { 645 + t.Errorf("Expected total size 300 for bob, got %d", totalSize) 646 + } 647 + 648 + // Unknown user should have 0 649 + uniqueBlobs, totalSize, err = ri.QuotaForDID("io.atcr.hold.layer", "did:plc:unknown") 650 + if err != nil { 651 + t.Fatalf("QuotaForDID() error = %v", err) 652 + } 653 + if uniqueBlobs != 0 || totalSize != 0 { 654 + t.Errorf("Expected 0/0 for unknown user, got %d/%d", uniqueBlobs, totalSize) 655 + } 656 + } 657 + 658 + // TestRecordsIndex_QuotasByDID tests bulk quota aggregation 659 + func TestRecordsIndex_QuotasByDID(t *testing.T) { 660 + tmpDir := t.TempDir() 661 + ri, err := NewRecordsIndex(filepath.Join(tmpDir, "records.db")) 662 + if err != nil { 663 + t.Fatalf("NewRecordsIndex() error = %v", err) 664 + } 665 + defer ri.Close() 666 + 667 + // Add layer records for multiple users 668 + ri.IndexRecord("io.atcr.hold.layer", "r1", "cid1", "did:plc:alice", "sha256:aaa", 100) 669 + ri.IndexRecord("io.atcr.hold.layer", "r2", "cid2", "did:plc:alice", "sha256:bbb", 200) 670 + ri.IndexRecord("io.atcr.hold.layer", "r3", "cid3", "did:plc:bob", "sha256:ccc", 500) 671 + // Non-layer record should be excluded 672 + ri.IndexRecord("io.atcr.hold.crew", "r4", "cid4", "did:plc:alice", "", 0) 673 + 674 + quotas, err := ri.QuotasByDID("io.atcr.hold.layer") 675 + if err != nil { 676 + t.Fatalf("QuotasByDID() error = %v", err) 677 + } 678 + 679 + if len(quotas) != 2 { 680 + t.Fatalf("Expected 2 users in quotas, got %d", len(quotas)) 681 + } 682 + 683 + alice := quotas["did:plc:alice"] 684 + if alice.UniqueBlobs != 2 { 685 + t.Errorf("Expected 2 unique blobs for alice, got %d", alice.UniqueBlobs) 686 + } 687 + if alice.TotalSize != 300 { 688 + t.Errorf("Expected total size 300 for alice, got %d", alice.TotalSize) 689 + } 690 + 691 + bob := quotas["did:plc:bob"] 692 + if bob.UniqueBlobs != 1 { 693 + t.Errorf("Expected 1 unique blob for bob, got %d", bob.UniqueBlobs) 694 + } 695 + if bob.TotalSize != 500 { 696 + t.Errorf("Expected total size 500 for bob, got %d", bob.TotalSize) 697 + } 698 + } 699 + 700 + // TestRecordsIndex_QuotasByDID_Empty tests bulk quota with no records 701 + func TestRecordsIndex_QuotasByDID_Empty(t *testing.T) { 702 + tmpDir := t.TempDir() 703 + ri, err := NewRecordsIndex(filepath.Join(tmpDir, "records.db")) 704 + if err != nil { 705 + t.Fatalf("NewRecordsIndex() error = %v", err) 706 + } 707 + defer ri.Close() 708 + 709 + quotas, err := ri.QuotasByDID("io.atcr.hold.layer") 710 + if err != nil { 711 + t.Fatalf("QuotasByDID() error = %v", err) 712 + } 713 + if len(quotas) != 0 { 714 + t.Errorf("Expected empty quotas map, got %d entries", len(quotas)) 715 + } 716 + } 717 + 718 + // TestRecordsIndex_QuotaForDID_IgnoresNullDigest tests that records without digest are excluded 719 + func TestRecordsIndex_QuotaForDID_IgnoresNullDigest(t *testing.T) { 720 + tmpDir := t.TempDir() 721 + ri, err := NewRecordsIndex(filepath.Join(tmpDir, "records.db")) 722 + if err != nil { 723 + t.Fatalf("NewRecordsIndex() error = %v", err) 724 + } 725 + defer ri.Close() 726 + 727 + // Record with digest+size 728 + ri.IndexRecord("io.atcr.hold.layer", "r1", "cid1", "did:plc:alice", "sha256:aaa", 100) 729 + // Record without digest (e.g. old data before migration) 730 + ri.IndexRecord("io.atcr.hold.layer", "r2", "cid2", "did:plc:alice", "", 0) 731 + 732 + uniqueBlobs, totalSize, err := ri.QuotaForDID("io.atcr.hold.layer", "did:plc:alice") 733 + if err != nil { 734 + t.Fatalf("QuotaForDID() error = %v", err) 735 + } 736 + if uniqueBlobs != 1 { 737 + t.Errorf("Expected 1 unique blob (ignoring null digest), got %d", uniqueBlobs) 738 + } 739 + if totalSize != 100 { 740 + t.Errorf("Expected total size 100, got %d", totalSize) 741 + } 742 + }
+14 -2
pkg/hold/pds/server.go
··· 411 411 if op.RecCid != nil { 412 412 cidStr = op.RecCid.String() 413 413 } 414 - // Extract DID from record based on collection type 414 + // Extract fields from record based on collection type 415 415 did := extractDIDFromOp(op) 416 - if err := p.recordsIndex.IndexRecord(op.Collection, op.Rkey, cidStr, did); err != nil { 416 + digest, size := extractLayerFieldsFromOp(op) 417 + if err := p.recordsIndex.IndexRecord(op.Collection, op.Rkey, cidStr, did, digest, size); err != nil { 417 418 slog.Warn("Failed to index record", "collection", op.Collection, "rkey", op.Rkey, "error", err) 418 419 } 419 420 case EvtKindDeleteRecord: ··· 452 453 } 453 454 } 454 455 return "" 456 + } 457 + 458 + // extractLayerFieldsFromOp extracts digest and size from a layer record operation 459 + func extractLayerFieldsFromOp(op RepoOp) (string, int64) { 460 + if op.Record == nil || op.Collection != atproto.LayerCollection { 461 + return "", 0 462 + } 463 + if rec, ok := op.Record.(*atproto.LayerRecord); ok { 464 + return rec.Digest, rec.Size 465 + } 466 + return "", 0 455 467 } 456 468 457 469 // BackfillRecordsIndex populates the records index from existing MST data