at-container-registry/pkg/hold/pds/records.go

package pds

import (
	"context"
	"database/sql"
	"fmt"
	"log/slog"
	"strings"

	"github.com/bluesky-social/indigo/repo"
	"github.com/ipfs/go-cid"
	_ "github.com/mattn/go-sqlite3"
)

// RecordsIndex provides an efficient index for listing records
// This follows the official ATProto PDS pattern of using SQL for queries
// while MST is used for sync operations.
type RecordsIndex struct {
	db *sql.DB
}

// Record represents a record in the index
type Record struct {
	Collection string
	Rkey       string
	Cid        string
}

const recordsSchema = `
CREATE TABLE IF NOT EXISTS records (
    collection TEXT NOT NULL,
    rkey TEXT NOT NULL,
    cid TEXT NOT NULL,
    PRIMARY KEY (collection, rkey)
);
CREATE INDEX IF NOT EXISTS idx_records_collection_rkey ON records(collection, rkey);
`

// NewRecordsIndex creates or opens a records index
func NewRecordsIndex(dbPath string) (*RecordsIndex, error) {
	db, err := sql.Open("sqlite3", dbPath)
	if err != nil {
		return nil, fmt.Errorf("failed to open records database: %w", err)
	}

	// Create schema
	_, err = db.Exec(recordsSchema)
	if err != nil {
		db.Close()
		return nil, fmt.Errorf("failed to create records schema: %w", err)
	}

	return &RecordsIndex{db: db}, nil
}

// Close closes the database connection
func (ri *RecordsIndex) Close() error {
	if ri.db != nil {
		return ri.db.Close()
	}
	return nil
}

// IndexRecord adds or updates a record in the index
func (ri *RecordsIndex) IndexRecord(collection, rkey, cidStr string) error {
	_, err := ri.db.Exec(`
		INSERT OR REPLACE INTO records (collection, rkey, cid)
		VALUES (?, ?, ?)
	`, collection, rkey, cidStr)
	return err
}

// DeleteRecord removes a record from the index
func (ri *RecordsIndex) DeleteRecord(collection, rkey string) error {
	_, err := ri.db.Exec(`
		DELETE FROM records WHERE collection = ? AND rkey = ?
	`, collection, rkey)
	return err
}

// ListRecords returns records for a collection with pagination support
// reverse=false (default): newest first (rkey DESC)
// reverse=true: oldest first (rkey ASC)
func (ri *RecordsIndex) ListRecords(collection string, limit int, cursor string, reverse bool) ([]Record, string, error) {
	// Build query based on sort order
	var query string
	var args []any

	if reverse {
		// Oldest first (ascending order)
		if cursor != "" {
			query = `
				SELECT collection, rkey, cid FROM records
				WHERE collection = ? AND rkey > ?
				ORDER BY rkey ASC
				LIMIT ?
			`
			args = []any{collection, cursor, limit + 1}
		} else {
			query = `
				SELECT collection, rkey, cid FROM records
				WHERE collection = ?
				ORDER BY rkey ASC
				LIMIT ?
			`
			args = []any{collection, limit + 1}
		}
	} else {
		// Newest first (descending order) - default
		if cursor != "" {
			query = `
				SELECT collection, rkey, cid FROM records
				WHERE collection = ? AND rkey < ?
				ORDER BY rkey DESC
				LIMIT ?
			`
			args = []any{collection, cursor, limit + 1}
		} else {
			query = `
				SELECT collection, rkey, cid FROM records
				WHERE collection = ?
				ORDER BY rkey DESC
				LIMIT ?
			`
			args = []any{collection, limit + 1}
		}
	}

	rows, err := ri.db.Query(query, args...)
	if err != nil {
		return nil, "", fmt.Errorf("failed to query records: %w", err)
	}
	defer rows.Close()

	var records []Record
	for rows.Next() {
		var rec Record
		if err := rows.Scan(&rec.Collection, &rec.Rkey, &rec.Cid); err != nil {
			return nil, "", fmt.Errorf("failed to scan record: %w", err)
		}
		records = append(records, rec)
	}

	if err := rows.Err(); err != nil {
		return nil, "", fmt.Errorf("error iterating records: %w", err)
	}

	// Determine next cursor
	var nextCursor string
	if len(records) > limit {
		// More records available, set cursor to the last included record
		nextCursor = records[limit-1].Rkey
		records = records[:limit]
	}

	return records, nextCursor, nil
}

// Count returns the number of records in a collection
func (ri *RecordsIndex) Count(collection string) (int, error) {
	var count int
	err := ri.db.QueryRow(`
		SELECT COUNT(*) FROM records WHERE collection = ?
	`, collection).Scan(&count)
	return count, err
}

// TotalCount returns the total number of records in the index
func (ri *RecordsIndex) TotalCount() (int, error) {
	var count int
	err := ri.db.QueryRow(`SELECT COUNT(*) FROM records`).Scan(&count)
	return count, err
}

// BackfillFromRepo populates the records index from an existing MST repo
// Compares MST count with index count - only backfills if they differ
func (ri *RecordsIndex) BackfillFromRepo(ctx context.Context, repoHandle *repo.Repo) error {
	// Count records in MST
	mstCount := 0
	err := repoHandle.ForEach(ctx, "", func(key string, c cid.Cid) error {
		mstCount++
		return nil
	})
	if err != nil {
		return fmt.Errorf("failed to count MST records: %w", err)
	}

	// Count records in index
	indexCount, err := ri.TotalCount()
	if err != nil {
		return fmt.Errorf("failed to check index count: %w", err)
	}

	// Skip if counts match
	if indexCount == mstCount {
		slog.Debug("Records index in sync with MST", "count", indexCount)
		return nil
	}

	slog.Info("Backfilling records index from MST...", "mstCount", mstCount, "indexCount", indexCount)

	// Begin transaction for bulk insert
	tx, err := ri.db.BeginTx(ctx, nil)
	if err != nil {
		return fmt.Errorf("failed to begin transaction: %w", err)
	}
	defer tx.Rollback()

	stmt, err := tx.Prepare(`
		INSERT OR REPLACE INTO records (collection, rkey, cid)
		VALUES (?, ?, ?)
	`)
	if err != nil {
		return fmt.Errorf("failed to prepare statement: %w", err)
	}
	defer stmt.Close()

	recordCount := 0
	err = repoHandle.ForEach(ctx, "", func(key string, c cid.Cid) error {
		// key format: "collection/rkey"
		parts := strings.SplitN(key, "/", 2)
		if len(parts) != 2 {
			return nil // Skip malformed keys
		}
		collection, rkey := parts[0], parts[1]

		_, err := stmt.Exec(collection, rkey, c.String())
		if err != nil {
			return fmt.Errorf("failed to index record %s: %w", key, err)
		}
		recordCount++

		// Log progress every 1000 records
		if recordCount%1000 == 0 {
			slog.Debug("Backfill progress", "count", recordCount)
		}

		return nil
	})

	if err != nil {
		return fmt.Errorf("failed to walk repo: %w", err)
	}

	if err := tx.Commit(); err != nil {
		return fmt.Errorf("failed to commit transaction: %w", err)
	}

	slog.Info("Backfill complete", "records", recordCount)
	return nil
}