100 lines
3.3 KiB
Go
100 lines
3.3 KiB
Go
package store
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
|
|
"github.com/jackc/pgx/v5"
|
|
"github.com/jackc/pgx/v5/pgconn"
|
|
"github.com/jackc/pgx/v5/pgxpool"
|
|
)
|
|
|
|
// DBTX is the interface for database query operations.
|
|
// Satisfied by *pgxpool.Pool, *pgx.Conn, and pgx.Tx.
|
|
type DBTX interface {
|
|
Exec(ctx context.Context, sql string, arguments ...any) (pgconn.CommandTag, error)
|
|
Query(ctx context.Context, sql string, args ...any) (pgx.Rows, error)
|
|
QueryRow(ctx context.Context, sql string, args ...any) pgx.Row
|
|
}
|
|
|
|
// Store provides database operations for the dev-pod API.
|
|
type Store struct {
|
|
db DBTX
|
|
}
|
|
|
|
// New creates a Store with the given database connection.
|
|
func New(db DBTX) *Store {
|
|
return &Store{db: db}
|
|
}
|
|
|
|
// NewPool creates a new pgx connection pool.
|
|
func NewPool(ctx context.Context, databaseURL string) (*pgxpool.Pool, error) {
|
|
pool, err := pgxpool.New(ctx, databaseURL)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("connect to database: %w", err)
|
|
}
|
|
if err := pool.Ping(ctx); err != nil {
|
|
pool.Close()
|
|
return nil, fmt.Errorf("ping database: %w", err)
|
|
}
|
|
return pool, nil
|
|
}
|
|
|
|
// Migrate runs database migrations to create required tables.
|
|
func (s *Store) Migrate(ctx context.Context) error {
|
|
migrations := []string{
|
|
`CREATE TABLE IF NOT EXISTS users (
|
|
id TEXT PRIMARY KEY,
|
|
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
|
max_concurrent_pods INTEGER NOT NULL DEFAULT 3,
|
|
max_cpu_per_pod INTEGER NOT NULL DEFAULT 8,
|
|
max_ram_gb_per_pod INTEGER NOT NULL DEFAULT 16,
|
|
monthly_pod_hours INTEGER NOT NULL DEFAULT 500,
|
|
monthly_ai_requests INTEGER NOT NULL DEFAULT 10000
|
|
)`,
|
|
`CREATE TABLE IF NOT EXISTS api_keys (
|
|
key_hash TEXT PRIMARY KEY,
|
|
user_id TEXT NOT NULL REFERENCES users(id) ON DELETE CASCADE,
|
|
role TEXT NOT NULL DEFAULT 'user',
|
|
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
|
last_used_at TIMESTAMPTZ
|
|
)`,
|
|
`CREATE TABLE IF NOT EXISTS usage_records (
|
|
id BIGSERIAL PRIMARY KEY,
|
|
user_id TEXT NOT NULL REFERENCES users(id) ON DELETE CASCADE,
|
|
pod_name TEXT NOT NULL,
|
|
event_type TEXT NOT NULL,
|
|
value DOUBLE PRECISION,
|
|
recorded_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
|
|
)`,
|
|
`CREATE INDEX IF NOT EXISTS idx_usage_user_month ON usage_records(user_id, recorded_at)`,
|
|
`ALTER TABLE users ADD COLUMN IF NOT EXISTS forgejo_token TEXT NOT NULL DEFAULT ''`,
|
|
`CREATE TABLE IF NOT EXISTS runners (
|
|
id TEXT PRIMARY KEY,
|
|
user_id TEXT NOT NULL REFERENCES users(id),
|
|
repo_url TEXT NOT NULL,
|
|
branch TEXT NOT NULL DEFAULT 'main',
|
|
tools TEXT NOT NULL DEFAULT '',
|
|
task TEXT NOT NULL DEFAULT '',
|
|
status TEXT NOT NULL DEFAULT 'received',
|
|
forgejo_runner_id TEXT NOT NULL DEFAULT '',
|
|
webhook_delivery_id TEXT NOT NULL DEFAULT '',
|
|
pod_name TEXT NOT NULL DEFAULT '',
|
|
cpu_req TEXT NOT NULL DEFAULT '2',
|
|
mem_req TEXT NOT NULL DEFAULT '4Gi',
|
|
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
|
claimed_at TIMESTAMPTZ,
|
|
completed_at TIMESTAMPTZ
|
|
)`,
|
|
`CREATE UNIQUE INDEX IF NOT EXISTS idx_runners_webhook_delivery
|
|
ON runners(webhook_delivery_id) WHERE webhook_delivery_id != ''`,
|
|
`CREATE INDEX IF NOT EXISTS idx_runners_status ON runners(status)`,
|
|
`ALTER TABLE users ADD COLUMN IF NOT EXISTS tailscale_key TEXT NOT NULL DEFAULT ''`,
|
|
}
|
|
for _, m := range migrations {
|
|
if _, err := s.db.Exec(ctx, m); err != nil {
|
|
return fmt.Errorf("run migration: %w", err)
|
|
}
|
|
}
|
|
return nil
|
|
}
|