orca/internal/store/capacity_repo.go

// Package store — capacity_repo.go implements persistence for NodeCapacity
// declarations (v0.2 P02). Capacity is declared per node via
// `orca node capacity --set` (or from `~/.orca/node.hcl` at join time).
// The dispatcher reads capacity rows to bin-pack jobs across nodes.
package store

import (
	"context"
	"database/sql"
	"errors"
	"fmt"
	"time"
)

// NodeCapacity is the per-node resource declaration consumed by the
// scheduler. Units:
//   - CPUMillicores: 1000 = 1 vCPU
//   - MemoryMiB:     mebibytes of RAM
//   - DiskMiB:       mebibytes of scratch disk
type NodeCapacity struct {
	NodeID        string
	CPUMillicores int64
	MemoryMiB     int64
	DiskMiB       int64
	UpdatedAt     time.Time
}

// CapacityRepo is the persistence layer for NodeCapacity rows.
type CapacityRepo struct {
	db *sql.DB
}

// NewCapacityRepo returns a CapacityRepo backed by the given DB.
func NewCapacityRepo(db *sql.DB) *CapacityRepo {
	return &CapacityRepo{db: db}
}

// Upsert writes the capacity row for nodeID, replacing any prior row.
// The UpdatedAt column is set to time.Now().UTC() unless the caller
// supplied a non-zero value.
func (r *CapacityRepo) Upsert(ctx context.Context, c *NodeCapacity) error {
	if c == nil {
		return errors.New("CapacityRepo.Upsert: nil capacity")
	}
	if c.NodeID == "" {
		return errors.New("CapacityRepo.Upsert: NodeID is required")
	}
	if c.UpdatedAt.IsZero() {
		c.UpdatedAt = time.Now().UTC()
	}
	_, err := r.db.ExecContext(ctx, `
		INSERT INTO node_capacity (node_id, cpu_millicores, memory_mib, disk_mib, updated_at)
		VALUES (?, ?, ?, ?, ?)
		ON CONFLICT(node_id) DO UPDATE SET
			cpu_millicores = excluded.cpu_millicores,
			memory_mib     = excluded.memory_mib,
			disk_mib       = excluded.disk_mib,
			updated_at     = excluded.updated_at
	`, c.NodeID, c.CPUMillicores, c.MemoryMiB, c.DiskMiB, c.UpdatedAt)
	if err != nil {
		return fmt.Errorf("CapacityRepo.Upsert: %w", err)
	}
	return nil
}

// Get returns the capacity for nodeID or ErrNotFound.
func (r *CapacityRepo) Get(ctx context.Context, nodeID string) (*NodeCapacity, error) {
	if nodeID == "" {
		return nil, errors.New("CapacityRepo.Get: nodeID is required")
	}
	row := r.db.QueryRowContext(ctx, `
		SELECT node_id, cpu_millicores, memory_mib, disk_mib, updated_at
		FROM node_capacity WHERE node_id = ?
	`, nodeID)
	var c NodeCapacity
	if err := row.Scan(&c.NodeID, &c.CPUMillicores, &c.MemoryMiB, &c.DiskMiB, &c.UpdatedAt); err != nil {
		if errors.Is(err, sql.ErrNoRows) {
			return nil, ErrNotFound
		}
		return nil, fmt.Errorf("CapacityRepo.Get: %w", err)
	}
	return &c, nil
}

// List returns all capacity rows ordered by node_id.
func (r *CapacityRepo) List(ctx context.Context) ([]*NodeCapacity, error) {
	rows, err := r.db.QueryContext(ctx, `
		SELECT node_id, cpu_millicores, memory_mib, disk_mib, updated_at
		FROM node_capacity ORDER BY node_id
	`)
	if err != nil {
		return nil, fmt.Errorf("CapacityRepo.List: %w", err)
	}
	defer rows.Close()
	var out []*NodeCapacity
	for rows.Next() {
		var c NodeCapacity
		if err := rows.Scan(&c.NodeID, &c.CPUMillicores, &c.MemoryMiB, &c.DiskMiB, &c.UpdatedAt); err != nil {
			return nil, fmt.Errorf("CapacityRepo.List: scan: %w", err)
		}
		out = append(out, &c)
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("CapacityRepo.List: rows: %w", err)
	}
	return out, nil
}

// Delete removes the capacity row for nodeID. Returns ErrNotFound if
// the row doesn't exist.
func (r *CapacityRepo) Delete(ctx context.Context, nodeID string) error {
	res, err := r.db.ExecContext(ctx, `DELETE FROM node_capacity WHERE node_id = ?`, nodeID)
	if err != nil {
		return fmt.Errorf("CapacityRepo.Delete: %w", err)
	}
	n, err := res.RowsAffected()
	if err != nil {
		return fmt.Errorf("CapacityRepo.Delete: rows: %w", err)
	}
	if n == 0 {
		return ErrNotFound
	}
	return nil
}