mirrors/at-container-registry
1
0
Fork 0
mirror of https://tangled.org/evan.jarrett.net/at-container-registry synced 2026-09-03 08:46:57 +00:00
Code Issues Packages Projects Releases Wiki Activity

optimize queries for admin panel

Browse Source
This commit is contained in:
Evan Jarrett
2026-02-10 22:51:51 -06:00
parent 8e45b2eee5
commit 22d5396589
11 changed files with 328 additions and 181 deletions
Download Patch File Download Diff File Expand all files Collapse all files
hold
BIN
View File
Binary file not shown.
pkg/appview/db/hold_store.go
+1 -3
View File
@@ -290,14 +290,12 @@ func GetAvailableHolds(db DBTX, userDID string) ([]AvailableHold, error) {
WHEN h.owner_did = ?1 THEN 'owner'
WHEN c.member_did IS NOT NULL THEN 'crew'
WHEN h.allow_all_crew = 1 THEN 'eligible'
WHEN h.public = 1 THEN 'public'
ELSE 'none'
END as membership,
c.permissions
FROM hold_captain_records h
LEFT JOIN hold_crew_members c ON h.hold_did = c.hold_did AND c.member_did = ?1
WHERE h.public = 1
OR h.allow_all_crew = 1
WHERE h.allow_all_crew = 1
OR h.owner_did = ?1
OR c.member_did IS NOT NULL
ORDER BY
pkg/appview/handlers/settings.go
+3 -5
View File
@@ -58,7 +58,7 @@ func (h *SettingsHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
slog.Debug("Fetched profile", "component", "settings", "did", user.DID, "default_hold", profile.DefaultHold)
// Get available holds for dropdown
var ownedHolds, crewHolds, eligibleHolds, publicHolds []HoldDisplay
var ownedHolds, crewHolds, eligibleHolds []HoldDisplay
holdDataMap := make(map[string]HoldDisplay)
if h.DB != nil {
@@ -93,8 +93,6 @@ func (h *SettingsHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
crewHolds = append(crewHolds, display)
case "eligible":
eligibleHolds = append(eligibleHolds, display)
case "public":
publicHolds = append(publicHolds, display)
}
}
}
@@ -134,12 +132,12 @@ func (h *SettingsHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
CurrentHoldDID string
CurrentHoldDisplay string
ShowCurrentHold bool
AppViewDefaultHoldDID string
AppViewDefaultHoldDisplay string
AppViewDefaultRegion string
OwnedHolds []HoldDisplay
CrewHolds []HoldDisplay
EligibleHolds []HoldDisplay
PublicHolds []HoldDisplay
HoldDataJSON template.JS
}{
PageData: NewPageData(r, &h.BaseUIHandler),
@@ -147,12 +145,12 @@ func (h *SettingsHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
CurrentHoldDID: profile.DefaultHold,
CurrentHoldDisplay: deriveDisplayName(profile.DefaultHold),
ShowCurrentHold: showCurrentHold,
AppViewDefaultHoldDID: h.DefaultHoldDID,
AppViewDefaultHoldDisplay: appViewDefaultDisplay,
AppViewDefaultRegion: appViewDefaultRegion,
OwnedHolds: ownedHolds,
CrewHolds: crewHolds,
EligibleHolds: eligibleHolds,
PublicHolds: publicHolds,
HoldDataJSON: template.JS(holdDataJSON),
}
pkg/appview/templates/pages/settings.html
+1 -10
View File
@@ -142,7 +142,7 @@
<span class="label-text">Storage Hold</span>
</label>
<select id="default-hold" name="hold_did" class="select select-bordered w-full" autocomplete="off">
<option value=""{{ if eq .CurrentHoldDID "" }} selected{{ end }}>AppView Default ({{ .AppViewDefaultHoldDisplay }}{{ if .AppViewDefaultRegion }}, {{ .AppViewDefaultRegion }}{{ end }})</option>
<option value="{{ .AppViewDefaultHoldDID }}"{{ if or (eq .CurrentHoldDID "") (eq .CurrentHoldDID .AppViewDefaultHoldDID) }} selected{{ end }}>AppView Default ({{ .AppViewDefaultHoldDisplay }}{{ if .AppViewDefaultRegion }}, {{ .AppViewDefaultRegion }}{{ end }})</option>
{{ if .ShowCurrentHold }}
<option value="{{ .CurrentHoldDID }}" selected>Current ({{ .CurrentHoldDisplay }})</option>
@@ -178,15 +178,6 @@
</optgroup>
{{ end }}
{{ if .PublicHolds }}
<optgroup label="Public Holds">
{{ range .PublicHolds }}
<option value="{{ .DID }}" {{ if eq $.CurrentHoldDID .DID }}selected{{ end }}>
{{ .DisplayName }}{{ if .Region }} ({{ .Region }}){{ end }}
</option>
{{ end }}
</optgroup>
{{ end }}
</select>
<p class="text-sm text-base-content/60 mt-1">Your images will be stored on the selected hold</p>
</fieldset>
pkg/hold/admin/handlers.go
+15 -26
View File
@@ -93,22 +93,17 @@ func (ui *AdminUI) handleStatsAPI(w http.ResponseWriter, r *http.Request) {
return
}
// Calculate total storage by summing quota for all crew members
// Calculate total storage with a single bulk query
var totalSize int64
uniqueDigests := 0
crew, err := ui.pds.ListCrewMembers(ctx)
allQuotas, err := ui.pds.GetAllUserQuotas(ctx)
if err != nil {
slog.Warn("Failed to list crew for stats", "error", err)
slog.Warn("Failed to get all user quotas", "error", err)
} else {
// Get usage for each crew member
for _, member := range crew {
quotaStats, err := ui.pds.GetQuotaForUser(ctx, member.Record.Member)
if err != nil {
continue
}
totalSize += quotaStats.TotalSize
uniqueDigests += quotaStats.UniqueBlobs
for _, q := range allQuotas {
totalSize += q.TotalSize
uniqueDigests += q.UniqueBlobs
}
}
@@ -152,28 +147,22 @@ func (ui *AdminUI) handleTopUsersAPI(w http.ResponseWriter, r *http.Request) {
}
}
// Get all crew members and their usage
crew, err := ui.pds.ListCrewMembers(ctx)
// Get all user quotas in a single bulk query
allQuotas, err := ui.pds.GetAllUserQuotas(ctx)
if err != nil {
slog.Error("Failed to list crew members for top users", "error", err)
slog.Error("Failed to get all user quotas", "error", err)
http.Error(w, "Failed to load top users", http.StatusInternalServerError)
return
}
var users []UserUsage
for _, member := range crew {
quotaStats, err := ui.pds.GetQuotaForUser(ctx, member.Record.Member)
if err != nil {
slog.Warn("Failed to get quota for user", "did", member.Record.Member, "error", err)
continue
}
for did, q := range allQuotas {
users = append(users, UserUsage{
DID: member.Record.Member,
Handle: resolveHandle(ctx, member.Record.Member),
Usage: quotaStats.TotalSize,
UsageHuman: formatHumanBytes(quotaStats.TotalSize),
BlobCount: quotaStats.UniqueBlobs,
DID: did,
Handle: resolveHandle(ctx, did),
Usage: q.TotalSize,
UsageHuman: formatHumanBytes(q.TotalSize),
BlobCount: q.UniqueBlobs,
})
}
pkg/hold/admin/handlers_crew.go
+13 -10
View File
@@ -9,6 +9,7 @@ import (
"time"
"atcr.io/pkg/atproto"
"atcr.io/pkg/hold/pds"
"github.com/go-chi/chi/v5"
)
@@ -56,14 +57,11 @@ func (ui *AdminUI) getCrewViews(ctx context.Context) ([]CrewMemberView, error) {
return nil, err
}
userUsage := make(map[string]int64)
for _, member := range crew {
quotaStats, err := ui.pds.GetQuotaForUser(ctx, member.Record.Member)
if err != nil {
slog.Warn("Failed to get quota for crew member", "did", member.Record.Member, "error", err)
continue
}
userUsage[member.Record.Member] = quotaStats.TotalSize
// Single bulk query for all user quotas
allQuotas, err := ui.pds.GetAllUserQuotas(ctx)
if err != nil {
slog.Warn("Failed to get all user quotas for crew views", "error", err)
allQuotas = make(map[string]*pds.QuotaStats)
}
defaultTier := "default"
@@ -88,11 +86,16 @@ func (ui *AdminUI) getCrewViews(ctx context.Context) ([]CrewMemberView, error) {
AddedAt: parseTime(member.Record.AddedAt),
}
usage := int64(0)
if q, ok := allQuotas[member.Record.Member]; ok {
usage = q.TotalSize
}
if ui.quotaMgr != nil && ui.quotaMgr.IsEnabled() {
if limit := ui.quotaMgr.GetTierLimit(tier); limit != nil {
view.TierLimit = formatHumanBytes(*limit)
if *limit > 0 {
view.UsagePercent = int(float64(userUsage[view.DID]) / float64(*limit) * 100)
view.UsagePercent = int(float64(usage) / float64(*limit) * 100)
}
} else {
view.TierLimit = "Unlimited"
@@ -101,7 +104,7 @@ func (ui *AdminUI) getCrewViews(ctx context.Context) ([]CrewMemberView, error) {
view.TierLimit = "Unlimited"
}
view.CurrentUsage = userUsage[view.DID]
view.CurrentUsage = usage
view.UsageHuman = formatHumanBytes(view.CurrentUsage)
crewViews = append(crewViews, view)
pkg/hold/pds/delete_test.go
+2 -2
View File
@@ -215,7 +215,7 @@ func TestDeleteAndListBlueskyPosts_WithPosts(t *testing.T) {
// Index the record
if pds.recordsIndex != nil {
err = pds.recordsIndex.IndexRecord(atproto.BskyPostCollection, rkey, "testcid", aliceDID)
err = pds.recordsIndex.IndexRecord(atproto.BskyPostCollection, rkey, "testcid", aliceDID, "", 0)
if err != nil {
t.Fatalf("Failed to index test post: %v", err)
}
@@ -304,7 +304,7 @@ func TestDeleteUserData_IncludesPosts(t *testing.T) {
// Index the record
if pds.recordsIndex != nil {
err = pds.recordsIndex.IndexRecord(atproto.BskyPostCollection, rkey, "testcid", aliceDID)
err = pds.recordsIndex.IndexRecord(atproto.BskyPostCollection, rkey, "testcid", aliceDID, "", 0)
if err != nil {
t.Fatalf("Failed to index test post: %v", err)
}
pkg/hold/pds/layer.go
+31 -91
View File
@@ -89,100 +89,48 @@ type QuotaStats struct {
Tier string `json:"tier,omitempty"` // quota tier (e.g., 'deckhand', 'bosun', 'quartermaster')
}
// GetQuotaForUser calculates storage quota for a specific user
// It iterates through all layer records, filters by userDid, deduplicates by digest,
// and sums the sizes of unique blobs.
// GetQuotaForUser calculates storage quota for a specific user.
// Uses SQL aggregation over the denormalized digest/size columns in the records index.
func (p *HoldPDS) GetQuotaForUser(ctx context.Context, userDID string) (*QuotaStats, error) {
if p.recordsIndex == nil {
return nil, fmt.Errorf("records index not available")
}
// Get session for reading record data
session, err := p.carstore.ReadOnlySession(p.uid)
uniqueBlobs, totalSize, err := p.recordsIndex.QuotaForDID(atproto.LayerCollection, userDID)
if err != nil {
return nil, fmt.Errorf("failed to create session: %w", err)
}
head, err := p.carstore.GetUserRepoHead(ctx, p.uid)
if err != nil {
return nil, fmt.Errorf("failed to get repo head: %w", err)
}
if !head.Defined() {
// Empty repo - return zero stats
return &QuotaStats{UserDID: userDID}, nil
}
repoHandle, err := repo.OpenRepo(ctx, session, head)
if err != nil {
return nil, fmt.Errorf("failed to open repo: %w", err)
}
// Track unique digests and their sizes
digestSizes := make(map[string]int64)
// Iterate all layer records via the index
cursor := ""
batchSize := 1000 // Process in batches
for {
records, nextCursor, err := p.recordsIndex.ListRecords(atproto.LayerCollection, batchSize, cursor, true)
if err != nil {
return nil, fmt.Errorf("failed to list layer records: %w", err)
}
for _, rec := range records {
// Construct record path and get the record data
recordPath := rec.Collection + "/" + rec.Rkey
_, recBytes, err := repoHandle.GetRecordBytes(ctx, recordPath)
if err != nil {
// Skip records we can't read
continue
}
// Decode the layer record
recordValue, err := lexutil.CborDecodeValue(*recBytes)
if err != nil {
continue
}
layerRecord, ok := recordValue.(*atproto.LayerRecord)
if !ok {
continue
}
// Filter by userDID
if layerRecord.UserDID != userDID {
continue
}
// Deduplicate by digest - keep the size (could be different pushes of same blob)
// Store the size - we only count each unique digest once
if _, exists := digestSizes[layerRecord.Digest]; !exists {
digestSizes[layerRecord.Digest] = layerRecord.Size
}
}
if nextCursor == "" {
break
}
cursor = nextCursor
}
// Calculate totals
var totalSize int64
for _, size := range digestSizes {
totalSize += size
return nil, fmt.Errorf("failed to query quota: %w", err)
}
return &QuotaStats{
UserDID: userDID,
UniqueBlobs: len(digestSizes),
UniqueBlobs: uniqueBlobs,
TotalSize: totalSize,
}, nil
}
// GetAllUserQuotas returns quota stats for all users in a single SQL query.
// Used by admin endpoints to avoid N+1 per-user quota lookups.
func (p *HoldPDS) GetAllUserQuotas(ctx context.Context) (map[string]*QuotaStats, error) {
if p.recordsIndex == nil {
return nil, fmt.Errorf("records index not available")
}
quotas, err := p.recordsIndex.QuotasByDID(atproto.LayerCollection)
if err != nil {
return nil, fmt.Errorf("failed to query all quotas: %w", err)
}
result := make(map[string]*QuotaStats, len(quotas))
for did, q := range quotas {
result[did] = &QuotaStats{
UserDID: did,
UniqueBlobs: q.UniqueBlobs,
TotalSize: q.TotalSize,
}
}
return result, nil
}
// GetQuotaForUserWithTier calculates quota with tier-aware limits
// It returns the base quota stats plus the tier limit and tier name.
// Captain (owner) always has unlimited quota.
@@ -262,27 +210,24 @@ func (p *HoldPDS) ListLayerRecordsForUser(ctx context.Context, userDID string) (
var records []*atproto.LayerRecord
// Iterate all layer records via the index
// Iterate layer records for this user via the index (filtered by DID in SQL)
cursor := ""
batchSize := 1000 // Process in batches
batchSize := 1000
for {
indexRecords, nextCursor, err := p.recordsIndex.ListRecords(atproto.LayerCollection, batchSize, cursor, true)
indexRecords, nextCursor, err := p.recordsIndex.ListRecordsByDID(atproto.LayerCollection, userDID, batchSize, cursor)
if err != nil {
return nil, fmt.Errorf("failed to list layer records: %w", err)
}
for _, rec := range indexRecords {
// Construct record path and get the record data
recordPath := rec.Collection + "/" + rec.Rkey
_, recBytes, err := repoHandle.GetRecordBytes(ctx, recordPath)
if err != nil {
// Skip records we can't read
continue
}
// Decode the layer record
recordValue, err := lexutil.CborDecodeValue(*recBytes)
if err != nil {
continue
@@ -293,11 +238,6 @@ func (p *HoldPDS) ListLayerRecordsForUser(ctx context.Context, userDID string) (
continue
}
// Filter by userDID
if layerRecord.UserDID != userDID {
continue
}
records = append(records, layerRecord)
}
pkg/hold/pds/records.go
+94 -13
View File
@@ -36,6 +36,8 @@ CREATE TABLE IF NOT EXISTS records (
rkey TEXT NOT NULL,
cid TEXT NOT NULL,
did TEXT,
digest TEXT,
size INTEGER,
PRIMARY KEY (collection, rkey)
);
CREATE INDEX IF NOT EXISTS idx_records_collection_rkey ON records(collection, rkey);
@@ -54,18 +56,26 @@ func NewRecordsIndex(dbPath string) (*RecordsIndex, error) {
return nil, fmt.Errorf("failed to open records database: %w", err)
}
// Check if table exists and has the did column
// Check if table exists and has required columns
needsRebuild := false
var tableName string
err = db.QueryRow(`SELECT name FROM sqlite_master WHERE type='table' AND name='records'`).Scan(&tableName)
if err == nil {
// Table exists, check for did column
var colCount int
// Check for did column
err = db.QueryRow(`SELECT COUNT(*) FROM pragma_table_info('records') WHERE name='did'`).Scan(&colCount)
if err != nil || colCount == 0 {
needsRebuild = true
slog.Info("Records index schema outdated, rebuilding with did column")
}
// Check for digest column (added for SQL-based quota queries)
if !needsRebuild {
err = db.QueryRow(`SELECT COUNT(*) FROM pragma_table_info('records') WHERE name='digest'`).Scan(&colCount)
if err != nil || colCount == 0 {
needsRebuild = true
slog.Info("Records index schema outdated, rebuilding with digest/size columns")
}
}
}
if needsRebuild {
@@ -91,17 +101,26 @@ func NewRecordsIndex(dbPath string) (*RecordsIndex, error) {
// NewRecordsIndexWithDB creates a records index using an existing *sql.DB connection.
// The caller is responsible for the DB lifecycle.
func NewRecordsIndexWithDB(db *sql.DB) (*RecordsIndex, error) {
// Check if table exists and has the did column
// Check if table exists and has required columns
needsRebuild := false
var tableName string
err := db.QueryRow(`SELECT name FROM sqlite_master WHERE type='table' AND name='records'`).Scan(&tableName)
if err == nil {
var colCount int
// Check for did column
err = db.QueryRow(`SELECT COUNT(*) FROM pragma_table_info('records') WHERE name='did'`).Scan(&colCount)
if err != nil || colCount == 0 {
needsRebuild = true
slog.Info("Records index schema outdated, rebuilding with did column")
}
// Check for digest column (added for SQL-based quota queries)
if !needsRebuild {
err = db.QueryRow(`SELECT COUNT(*) FROM pragma_table_info('records') WHERE name='digest'`).Scan(&colCount)
if err != nil || colCount == 0 {
needsRebuild = true
slog.Info("Records index schema outdated, rebuilding with digest/size columns")
}
}
}
if needsRebuild {
@@ -129,12 +148,15 @@ func (ri *RecordsIndex) Close() error {
}
// IndexRecord adds or updates a record in the index
// did parameter is optional - pass empty string if not applicable
func (ri *RecordsIndex) IndexRecord(collection, rkey, cidStr, did string) error {
// did, digest, size parameters are optional - pass empty/zero if not applicable
func (ri *RecordsIndex) IndexRecord(collection, rkey, cidStr, did, digest string, size int64) error {
_, err := ri.db.Exec(`
INSERT OR REPLACE INTO records (collection, rkey, cid, did)
VALUES (?, ?, ?, ?)
`, collection, rkey, cidStr, sql.NullString{String: did, Valid: did != ""})
INSERT OR REPLACE INTO records (collection, rkey, cid, did, digest, size)
VALUES (?, ?, ?, ?, ?, ?)
`, collection, rkey, cidStr,
sql.NullString{String: did, Valid: did != ""},
sql.NullString{String: digest, Valid: digest != ""},
sql.NullInt64{Int64: size, Valid: size > 0})
return err
}
@@ -328,8 +350,8 @@ func (ri *RecordsIndex) BackfillFromRepo(ctx context.Context, repoHandle *repo.R
defer tx.Rollback()
stmt, err := tx.Prepare(`
INSERT OR REPLACE INTO records (collection, rkey, cid, did)
VALUES (?, ?, ?, ?)
INSERT OR REPLACE INTO records (collection, rkey, cid, did, digest, size)
VALUES (?, ?, ?, ?, ?, ?)
`)
if err != nil {
return fmt.Errorf("failed to prepare statement: %w", err)
@@ -345,14 +367,19 @@ func (ri *RecordsIndex) BackfillFromRepo(ctx context.Context, repoHandle *repo.R
}
collection, rkey := parts[0], parts[1]
// Extract DID from record content based on collection type
var did string
// Extract fields from record content based on collection type
var did, digest string
var size int64
_, recBytes, err := repoHandle.GetRecordBytes(ctx, key)
if err == nil && recBytes != nil {
did = extractDIDFromRecord(collection, *recBytes)
digest, size = extractLayerFieldsFromRecord(collection, *recBytes)
}
_, err = stmt.Exec(collection, rkey, c.String(), sql.NullString{String: did, Valid: did != ""})
_, err = stmt.Exec(collection, rkey, c.String(),
sql.NullString{String: did, Valid: did != ""},
sql.NullString{String: digest, Valid: digest != ""},
sql.NullInt64{Int64: size, Valid: size > 0})
if err != nil {
return fmt.Errorf("failed to index record %s: %w", key, err)
}
@@ -392,6 +419,60 @@ func splitStatements(sql string) []string {
return out
}
// QuotaForDID returns unique blob count and total size for a single DID in a collection.
// Uses SQL aggregation over the denormalized digest/size columns.
func (ri *RecordsIndex) QuotaForDID(collection, did string) (uniqueBlobs int, totalSize int64, err error) {
err = ri.db.QueryRow(`
SELECT COUNT(*), COALESCE(SUM(size), 0)
FROM (SELECT DISTINCT digest, size FROM records WHERE collection = ? AND did = ? AND digest IS NOT NULL)
`, collection, did).Scan(&uniqueBlobs, &totalSize)
return
}
// QuotasByDID returns unique blob count and total size grouped by DID.
// Single query replaces N individual quota lookups.
func (ri *RecordsIndex) QuotasByDID(collection string) (map[string]QuotaResult, error) {
rows, err := ri.db.Query(`
SELECT did, COUNT(*), COALESCE(SUM(size), 0)
FROM (SELECT DISTINCT did, digest, size FROM records WHERE collection = ? AND did IS NOT NULL AND digest IS NOT NULL)
GROUP BY did
`, collection)
if err != nil {
return nil, fmt.Errorf("failed to query quotas by DID: %w", err)
}
defer rows.Close()
result := make(map[string]QuotaResult)
for rows.Next() {
var did string
var qr QuotaResult
if err := rows.Scan(&did, &qr.UniqueBlobs, &qr.TotalSize); err != nil {
return nil, fmt.Errorf("failed to scan quota row: %w", err)
}
result[did] = qr
}
return result, rows.Err()
}
// QuotaResult holds aggregated quota data from SQL queries
type QuotaResult struct {
UniqueBlobs int
TotalSize int64
}
// extractLayerFieldsFromRecord extracts digest and size from layer records.
// Returns empty/zero for non-layer records.
func extractLayerFieldsFromRecord(collection string, recBytes []byte) (string, int64) {
if collection != atproto.LayerCollection {
return "", 0
}
var rec atproto.LayerRecord
if err := rec.UnmarshalCBOR(bytes.NewReader(recBytes)); err != nil {
return "", 0
}
return rec.Digest, rec.Size
}
// extractDIDFromRecord extracts the associated DID from a record based on its collection type
func extractDIDFromRecord(collection string, recBytes []byte) string {
switch collection {
pkg/hold/pds/records_test.go
+154 -19
View File
@@ -50,7 +50,7 @@ func TestRecordsIndex_IndexRecord(t *testing.T) {
defer ri.Close()
// Index a record
err = ri.IndexRecord("io.atcr.hold.crew", "abc123", "bafyrei123", "")
err = ri.IndexRecord("io.atcr.hold.crew", "abc123", "bafyrei123", "", "", 0)
if err != nil {
t.Fatalf("IndexRecord() error = %v", err)
}
@@ -75,13 +75,13 @@ func TestRecordsIndex_IndexRecord_Upsert(t *testing.T) {
defer ri.Close()
// Index a record
err = ri.IndexRecord("io.atcr.hold.crew", "abc123", "bafyrei123", "")
err = ri.IndexRecord("io.atcr.hold.crew", "abc123", "bafyrei123", "", "", 0)
if err != nil {
t.Fatalf("IndexRecord() first call error = %v", err)
}
// Update the same record with new CID
err = ri.IndexRecord("io.atcr.hold.crew", "abc123", "bafyrei456", "")
err = ri.IndexRecord("io.atcr.hold.crew", "abc123", "bafyrei456", "", "", 0)
if err != nil {
t.Fatalf("IndexRecord() second call error = %v", err)
}
@@ -118,7 +118,7 @@ func TestRecordsIndex_DeleteRecord(t *testing.T) {
defer ri.Close()
// Index a record
err = ri.IndexRecord("io.atcr.hold.crew", "abc123", "bafyrei123", "")
err = ri.IndexRecord("io.atcr.hold.crew", "abc123", "bafyrei123", "", "", 0)
if err != nil {
t.Fatalf("IndexRecord() error = %v", err)
}
@@ -217,7 +217,7 @@ func TestRecordsIndex_ListRecords_Basic(t *testing.T) {
{"ccc", "cid3"},
}
for _, r := range records {
if err := ri.IndexRecord("io.atcr.hold.crew", r.rkey, r.cid, ""); err != nil {
if err := ri.IndexRecord("io.atcr.hold.crew", r.rkey, r.cid, "", "", 0); err != nil {
t.Fatalf("IndexRecord() error = %v", err)
}
}
@@ -248,7 +248,7 @@ func TestRecordsIndex_ListRecords_DefaultOrder(t *testing.T) {
// Add records with different rkeys (TIDs are lexicographically ordered by time)
rkeys := []string{"3m3aaaaaaaaa", "3m3bbbbbbbbb", "3m3ccccccccc"}
for _, rkey := range rkeys {
if err := ri.IndexRecord("io.atcr.hold.crew", rkey, "cid-"+rkey, ""); err != nil {
if err := ri.IndexRecord("io.atcr.hold.crew", rkey, "cid-"+rkey, "", "", 0); err != nil {
t.Fatalf("IndexRecord() error = %v", err)
}
}
@@ -286,7 +286,7 @@ func TestRecordsIndex_ListRecords_ReverseOrder(t *testing.T) {
// Add records
rkeys := []string{"3m3aaaaaaaaa", "3m3bbbbbbbbb", "3m3ccccccccc"}
for _, rkey := range rkeys {
if err := ri.IndexRecord("io.atcr.hold.crew", rkey, "cid-"+rkey, ""); err != nil {
if err := ri.IndexRecord("io.atcr.hold.crew", rkey, "cid-"+rkey, "", "", 0); err != nil {
t.Fatalf("IndexRecord() error = %v", err)
}
}
@@ -324,7 +324,7 @@ func TestRecordsIndex_ListRecords_Limit(t *testing.T) {
// Add 5 records
for i := range 5 {
rkey := string(rune('a' + i))
if err := ri.IndexRecord("io.atcr.hold.crew", rkey, "cid-"+rkey, ""); err != nil {
if err := ri.IndexRecord("io.atcr.hold.crew", rkey, "cid-"+rkey, "", "", 0); err != nil {
t.Fatalf("IndexRecord() error = %v", err)
}
}
@@ -355,7 +355,7 @@ func TestRecordsIndex_ListRecords_Cursor(t *testing.T) {
// Add 5 records
rkeys := []string{"a", "b", "c", "d", "e"}
for _, rkey := range rkeys {
if err := ri.IndexRecord("io.atcr.hold.crew", rkey, "cid-"+rkey, ""); err != nil {
if err := ri.IndexRecord("io.atcr.hold.crew", rkey, "cid-"+rkey, "", "", 0); err != nil {
t.Fatalf("IndexRecord() error = %v", err)
}
}
@@ -430,7 +430,7 @@ func TestRecordsIndex_ListRecords_CursorReverse(t *testing.T) {
// Add 5 records
rkeys := []string{"a", "b", "c", "d", "e"}
for _, rkey := range rkeys {
if err := ri.IndexRecord("io.atcr.hold.crew", rkey, "cid-"+rkey, ""); err != nil {
if err := ri.IndexRecord("io.atcr.hold.crew", rkey, "cid-"+rkey, "", "", 0); err != nil {
t.Fatalf("IndexRecord() error = %v", err)
}
}
@@ -474,10 +474,10 @@ func TestRecordsIndex_Count(t *testing.T) {
// Add records to two collections
for i := range 3 {
ri.IndexRecord("io.atcr.hold.crew", string(rune('a'+i)), "cid1", "")
ri.IndexRecord("io.atcr.hold.crew", string(rune('a'+i)), "cid1", "", "", 0)
}
for i := range 5 {
ri.IndexRecord("io.atcr.hold.captain", string(rune('a'+i)), "cid2", "")
ri.IndexRecord("io.atcr.hold.captain", string(rune('a'+i)), "cid2", "", "", 0)
}
// Count crew
@@ -527,10 +527,10 @@ func TestRecordsIndex_TotalCount(t *testing.T) {
defer ri.Close()
// Add records to multiple collections
ri.IndexRecord("io.atcr.hold.crew", "a", "cid1", "")
ri.IndexRecord("io.atcr.hold.crew", "b", "cid2", "")
ri.IndexRecord("io.atcr.hold.captain", "self", "cid3", "")
ri.IndexRecord("io.atcr.manifest", "abc123", "cid4", "")
ri.IndexRecord("io.atcr.hold.crew", "a", "cid1", "", "", 0)
ri.IndexRecord("io.atcr.hold.crew", "b", "cid2", "", "", 0)
ri.IndexRecord("io.atcr.hold.captain", "self", "cid3", "", "", 0)
ri.IndexRecord("io.atcr.manifest", "abc123", "cid4", "", "", 0)
count, err := ri.TotalCount()
if err != nil {
@@ -581,9 +581,9 @@ func TestRecordsIndex_MultipleCollections(t *testing.T) {
defer ri.Close()
// Add records to different collections with same rkeys
ri.IndexRecord("io.atcr.hold.crew", "abc", "cid-crew", "")
ri.IndexRecord("io.atcr.hold.captain", "abc", "cid-captain", "")
ri.IndexRecord("io.atcr.manifest", "abc", "cid-manifest", "")
ri.IndexRecord("io.atcr.hold.crew", "abc", "cid-crew", "", "", 0)
ri.IndexRecord("io.atcr.hold.captain", "abc", "cid-captain", "", "", 0)
ri.IndexRecord("io.atcr.manifest", "abc", "cid-manifest", "", "", 0)
// Listing should only return records from requested collection
records, _, err := ri.ListRecords("io.atcr.hold.crew", 10, "", false)
@@ -605,3 +605,138 @@ func TestRecordsIndex_MultipleCollections(t *testing.T) {
t.Errorf("Expected captain count 1 after deleting crew, got %d", count)
}
}
// TestRecordsIndex_QuotaForDID tests single-user quota aggregation
func TestRecordsIndex_QuotaForDID(t *testing.T) {
tmpDir := t.TempDir()
ri, err := NewRecordsIndex(filepath.Join(tmpDir, "records.db"))
if err != nil {
t.Fatalf("NewRecordsIndex() error = %v", err)
}
defer ri.Close()
// Add layer records for two users, with some duplicate digests
ri.IndexRecord("io.atcr.hold.layer", "r1", "cid1", "did:plc:alice", "sha256:aaa", 100)
ri.IndexRecord("io.atcr.hold.layer", "r2", "cid2", "did:plc:alice", "sha256:bbb", 200)
ri.IndexRecord("io.atcr.hold.layer", "r3", "cid3", "did:plc:alice", "sha256:aaa", 100) // duplicate digest
ri.IndexRecord("io.atcr.hold.layer", "r4", "cid4", "did:plc:bob", "sha256:ccc", 300)
// Alice should have 2 unique blobs, total 300 bytes
uniqueBlobs, totalSize, err := ri.QuotaForDID("io.atcr.hold.layer", "did:plc:alice")
if err != nil {
t.Fatalf("QuotaForDID() error = %v", err)
}
if uniqueBlobs != 2 {
t.Errorf("Expected 2 unique blobs for alice, got %d", uniqueBlobs)
}
if totalSize != 300 {
t.Errorf("Expected total size 300 for alice, got %d", totalSize)
}
// Bob should have 1 unique blob, total 300 bytes
uniqueBlobs, totalSize, err = ri.QuotaForDID("io.atcr.hold.layer", "did:plc:bob")
if err != nil {
t.Fatalf("QuotaForDID() error = %v", err)
}
if uniqueBlobs != 1 {
t.Errorf("Expected 1 unique blob for bob, got %d", uniqueBlobs)
}
if totalSize != 300 {
t.Errorf("Expected total size 300 for bob, got %d", totalSize)
}
// Unknown user should have 0
uniqueBlobs, totalSize, err = ri.QuotaForDID("io.atcr.hold.layer", "did:plc:unknown")
if err != nil {
t.Fatalf("QuotaForDID() error = %v", err)
}
if uniqueBlobs != 0 || totalSize != 0 {
t.Errorf("Expected 0/0 for unknown user, got %d/%d", uniqueBlobs, totalSize)
}
}
// TestRecordsIndex_QuotasByDID tests bulk quota aggregation
func TestRecordsIndex_QuotasByDID(t *testing.T) {
tmpDir := t.TempDir()
ri, err := NewRecordsIndex(filepath.Join(tmpDir, "records.db"))
if err != nil {
t.Fatalf("NewRecordsIndex() error = %v", err)
}
defer ri.Close()
// Add layer records for multiple users
ri.IndexRecord("io.atcr.hold.layer", "r1", "cid1", "did:plc:alice", "sha256:aaa", 100)
ri.IndexRecord("io.atcr.hold.layer", "r2", "cid2", "did:plc:alice", "sha256:bbb", 200)
ri.IndexRecord("io.atcr.hold.layer", "r3", "cid3", "did:plc:bob", "sha256:ccc", 500)
// Non-layer record should be excluded
ri.IndexRecord("io.atcr.hold.crew", "r4", "cid4", "did:plc:alice", "", 0)
quotas, err := ri.QuotasByDID("io.atcr.hold.layer")
if err != nil {
t.Fatalf("QuotasByDID() error = %v", err)
}
if len(quotas) != 2 {
t.Fatalf("Expected 2 users in quotas, got %d", len(quotas))
}
alice := quotas["did:plc:alice"]
if alice.UniqueBlobs != 2 {
t.Errorf("Expected 2 unique blobs for alice, got %d", alice.UniqueBlobs)
}
if alice.TotalSize != 300 {
t.Errorf("Expected total size 300 for alice, got %d", alice.TotalSize)
}
bob := quotas["did:plc:bob"]
if bob.UniqueBlobs != 1 {
t.Errorf("Expected 1 unique blob for bob, got %d", bob.UniqueBlobs)
}
if bob.TotalSize != 500 {
t.Errorf("Expected total size 500 for bob, got %d", bob.TotalSize)
}
}
// TestRecordsIndex_QuotasByDID_Empty tests bulk quota with no records
func TestRecordsIndex_QuotasByDID_Empty(t *testing.T) {
tmpDir := t.TempDir()
ri, err := NewRecordsIndex(filepath.Join(tmpDir, "records.db"))
if err != nil {
t.Fatalf("NewRecordsIndex() error = %v", err)
}
defer ri.Close()
quotas, err := ri.QuotasByDID("io.atcr.hold.layer")
if err != nil {
t.Fatalf("QuotasByDID() error = %v", err)
}
if len(quotas) != 0 {
t.Errorf("Expected empty quotas map, got %d entries", len(quotas))
}
}
// TestRecordsIndex_QuotaForDID_IgnoresNullDigest tests that records without digest are excluded
func TestRecordsIndex_QuotaForDID_IgnoresNullDigest(t *testing.T) {
tmpDir := t.TempDir()
ri, err := NewRecordsIndex(filepath.Join(tmpDir, "records.db"))
if err != nil {
t.Fatalf("NewRecordsIndex() error = %v", err)
}
defer ri.Close()
// Record with digest+size
ri.IndexRecord("io.atcr.hold.layer", "r1", "cid1", "did:plc:alice", "sha256:aaa", 100)
// Record without digest (e.g. old data before migration)
ri.IndexRecord("io.atcr.hold.layer", "r2", "cid2", "did:plc:alice", "", 0)
uniqueBlobs, totalSize, err := ri.QuotaForDID("io.atcr.hold.layer", "did:plc:alice")
if err != nil {
t.Fatalf("QuotaForDID() error = %v", err)
}
if uniqueBlobs != 1 {
t.Errorf("Expected 1 unique blob (ignoring null digest), got %d", uniqueBlobs)
}
if totalSize != 100 {
t.Errorf("Expected total size 100, got %d", totalSize)
}
}
pkg/hold/pds/server.go
+14 -2
View File
@@ -411,9 +411,10 @@ func (p *HoldPDS) CreateRecordsIndexEventHandler(broadcasterHandler func(context
if op.RecCid != nil {
cidStr = op.RecCid.String()
}
// Extract DID from record based on collection type
// Extract fields from record based on collection type
did := extractDIDFromOp(op)
if err := p.recordsIndex.IndexRecord(op.Collection, op.Rkey, cidStr, did); err != nil {
digest, size := extractLayerFieldsFromOp(op)
if err := p.recordsIndex.IndexRecord(op.Collection, op.Rkey, cidStr, did, digest, size); err != nil {
slog.Warn("Failed to index record", "collection", op.Collection, "rkey", op.Rkey, "error", err)
}
case EvtKindDeleteRecord:
@@ -454,6 +455,17 @@ func extractDIDFromOp(op RepoOp) string {
return ""
}
// extractLayerFieldsFromOp extracts digest and size from a layer record operation
func extractLayerFieldsFromOp(op RepoOp) (string, int64) {
if op.Record == nil || op.Collection != atproto.LayerCollection {
return "", 0
}
if rec, ok := op.Record.(*atproto.LayerRecord); ok {
return rec.Digest, rec.Size
}
return "", 0
}
// BackfillRecordsIndex populates the records index from existing MST data
func (p *HoldPDS) BackfillRecordsIndex(ctx context.Context) error {
if p.recordsIndex == nil {