Files
groot 47c523dd98
All checks were successful
release-tag / release-image (push) Successful in 3m51s
RC-14
2026-08-14 06:17:30 +02:00

761 lines
31 KiB
Go

package customer
import (
"context"
"database/sql"
"errors"
"fmt"
"net/url"
"os"
"path/filepath"
"strings"
"time"
_ "modernc.org/sqlite"
)
const schema = `
PRAGMA foreign_keys=ON;
CREATE TABLE IF NOT EXISTS customers(
id TEXT PRIMARY KEY,
username TEXT NOT NULL UNIQUE COLLATE NOCASE,
password_salt TEXT NOT NULL,
password_hash TEXT NOT NULL,
reward_client_id TEXT NOT NULL DEFAULT '',
blocked INTEGER NOT NULL DEFAULT 0,
blocked_reason TEXT NOT NULL DEFAULT '',
blocked_at INTEGER,
worker_limit_bypass INTEGER NOT NULL DEFAULT 0,
created_at INTEGER NOT NULL,
updated_at INTEGER NOT NULL
);
CREATE TABLE IF NOT EXISTS customer_sessions(
id TEXT PRIMARY KEY,
customer_id TEXT NOT NULL REFERENCES customers(id) ON DELETE CASCADE,
expires_at INTEGER NOT NULL,
created_at INTEGER NOT NULL
);
CREATE INDEX IF NOT EXISTS customer_sessions_exp_idx ON customer_sessions(expires_at);
CREATE TABLE IF NOT EXISTS credit_ledger(
id INTEGER PRIMARY KEY AUTOINCREMENT,
customer_id TEXT NOT NULL REFERENCES customers(id) ON DELETE CASCADE,
delta_micros INTEGER NOT NULL,
reason TEXT NOT NULL,
reference TEXT NOT NULL UNIQUE,
created_at INTEGER NOT NULL
);
CREATE INDEX IF NOT EXISTS credit_ledger_customer_idx ON credit_ledger(customer_id,created_at DESC);
CREATE TABLE IF NOT EXISTS workers(
id TEXT PRIMARY KEY,
customer_id TEXT NOT NULL REFERENCES customers(id) ON DELETE CASCADE,
task_id TEXT NOT NULL,
beacon_path TEXT NOT NULL DEFAULT 'auto',
docker_container_id TEXT NOT NULL DEFAULT '',
controller_id TEXT NOT NULL DEFAULT '',
docker_volume TEXT NOT NULL,
status TEXT NOT NULL DEFAULT 'stopped' CHECK(status IN ('stopped','starting','running','error')),
worker_client_id TEXT NOT NULL DEFAULT '',
register_token TEXT NOT NULL,
rate_micros_per_minute INTEGER NOT NULL,
last_charge_at INTEGER,
last_lease_at INTEGER,
created_at INTEGER NOT NULL,
updated_at INTEGER NOT NULL,
last_error TEXT NOT NULL DEFAULT ''
);
CREATE INDEX IF NOT EXISTS workers_customer_idx ON workers(customer_id,created_at);
CREATE TABLE IF NOT EXISTS service_controllers(
id TEXT PRIMARY KEY,
name TEXT NOT NULL,
base_url TEXT NOT NULL,
protocol_version INTEGER NOT NULL DEFAULT 1,
max_workers INTEGER NOT NULL DEFAULT 1000,
max_running INTEGER NOT NULL DEFAULT 100,
reported_workers INTEGER NOT NULL DEFAULT 0,
reported_running INTEGER NOT NULL DEFAULT 0,
enabled INTEGER NOT NULL DEFAULT 1,
draining INTEGER NOT NULL DEFAULT 0,
last_seen INTEGER NOT NULL,
last_error TEXT NOT NULL DEFAULT '',
created_at INTEGER NOT NULL,
updated_at INTEGER NOT NULL
);
CREATE INDEX IF NOT EXISTS service_controllers_seen_idx ON service_controllers(last_seen);
CREATE TABLE IF NOT EXISTS paypal_orders(
order_id TEXT PRIMARY KEY,
customer_id TEXT NOT NULL REFERENCES customers(id) ON DELETE CASCADE,
package_id TEXT NOT NULL,
amount_cents INTEGER NOT NULL,
currency TEXT NOT NULL,
credits_micros INTEGER NOT NULL,
status TEXT NOT NULL,
capture_id TEXT NOT NULL DEFAULT '',
created_at INTEGER NOT NULL,
updated_at INTEGER NOT NULL
);
CREATE INDEX IF NOT EXISTS paypal_orders_customer_idx ON paypal_orders(customer_id,created_at DESC);
CREATE TABLE IF NOT EXISTS customer_settings(
key TEXT PRIMARY KEY,
value TEXT NOT NULL,
updated_at INTEGER NOT NULL
);
CREATE TABLE IF NOT EXISTS registration_invites(
token_hash TEXT PRIMARY KEY,
label TEXT NOT NULL DEFAULT '',
expires_at INTEGER NOT NULL,
used_by TEXT REFERENCES customers(id) ON DELETE SET NULL,
used_at INTEGER,
created_at INTEGER NOT NULL
);
CREATE INDEX IF NOT EXISTS registration_invites_exp_idx ON registration_invites(expires_at,used_at);
`
type Store struct{ DB *sql.DB }
func Open(ctx context.Context, path string) (*Store, error) {
if strings.TrimSpace(path) == "" {
path = "/customer-data/customer-service.db"
}
abs, err := filepath.Abs(path)
if err != nil {
return nil, err
}
if err := os.MkdirAll(filepath.Dir(abs), 0o750); err != nil {
return nil, err
}
u := &url.URL{Scheme: "file", Path: filepath.ToSlash(abs)}
q := u.Query()
q.Add("_pragma", "busy_timeout(10000)")
q.Add("_pragma", "foreign_keys(ON)")
q.Add("_pragma", "synchronous(NORMAL)")
q.Set("_txlock", "immediate")
u.RawQuery = q.Encode()
db, err := sql.Open("sqlite", u.String())
if err != nil {
return nil, err
}
db.SetMaxOpenConns(4)
if err := db.PingContext(ctx); err != nil {
db.Close()
return nil, err
}
if _, err := db.ExecContext(ctx, "PRAGMA journal_mode=WAL"); err != nil {
db.Close()
return nil, err
}
for i, stmt := range strings.Split(schema, ";") {
stmt = strings.TrimSpace(stmt)
if stmt == "" {
continue
}
if _, err := db.ExecContext(ctx, stmt); err != nil {
db.Close()
return nil, fmt.Errorf("customer schema %d: %w", i+1, err)
}
}
// Existing V4 customer databases need explicit column evolution because
// CREATE TABLE IF NOT EXISTS does not alter an already-created table.
for _, m := range []struct{ name, def string }{
{"blocked", "INTEGER NOT NULL DEFAULT 0"},
{"blocked_reason", "TEXT NOT NULL DEFAULT ''"},
{"blocked_at", "INTEGER"},
{"worker_limit_bypass", "INTEGER NOT NULL DEFAULT 0"},
} {
if err := ensureCustomerColumn(ctx, db, "customers", m.name, m.def); err != nil {
db.Close()
return nil, fmt.Errorf("customer add customers.%s: %w", m.name, err)
}
}
for _, m := range []struct{ name, def string }{
{"controller_id", "TEXT NOT NULL DEFAULT ''"},
{"last_lease_at", "INTEGER"},
} {
if err := ensureCustomerColumn(ctx, db, "workers", m.name, m.def); err != nil {
db.Close()
return nil, fmt.Errorf("customer add workers.%s: %w", m.name, err)
}
}
if err := ensureCustomerColumn(ctx, db, "service_controllers", "protocol_version", "INTEGER NOT NULL DEFAULT 1"); err != nil {
db.Close()
return nil, fmt.Errorf("customer add service_controllers.protocol_version: %w", err)
}
return &Store{DB: db}, nil
}
func ensureCustomerColumn(ctx context.Context, db *sql.DB, table, name, def string) error {
rows, err := db.QueryContext(ctx, `PRAGMA table_info(`+table+`)`)
if err != nil {
return err
}
defer rows.Close()
for rows.Next() {
var cid int
var col, typ string
var notnull, pk int
var dflt sql.NullString
if err := rows.Scan(&cid, &col, &typ, &notnull, &dflt, &pk); err != nil {
return err
}
if strings.EqualFold(col, name) {
return nil
}
}
_, err = db.ExecContext(ctx, `ALTER TABLE `+table+` ADD COLUMN `+name+` `+def)
return err
}
type Customer struct {
ID string `json:"id"`
Username string `json:"username"`
RewardClientID string `json:"reward_client_id"`
Blocked bool `json:"blocked"`
BlockedReason string `json:"blocked_reason,omitempty"`
BlockedAt *time.Time `json:"blocked_at,omitempty"`
WorkerLimitBypass bool `json:"worker_limit_bypass"`
CreatedAt time.Time `json:"created_at"`
}
type Worker struct {
ID string `json:"id"`
CustomerID string `json:"customer_id"`
TaskID string `json:"task_id"`
BeaconPath string `json:"beacon_path"`
ContainerID string `json:"container_id"`
ControllerID string `json:"controller_id,omitempty"`
Volume string `json:"volume"`
Status string `json:"status"`
WorkerClientID string `json:"worker_client_id"`
RegisterToken string `json:"-"`
RateMicrosPerMinute int64 `json:"rate_micros_per_minute"`
LastChargeAt *time.Time `json:"last_charge_at,omitempty"`
LastLeaseAt *time.Time `json:"last_lease_at,omitempty"`
CreatedAt time.Time `json:"created_at"`
UpdatedAt time.Time `json:"updated_at"`
LastError string `json:"last_error,omitempty"`
}
type ServiceController struct {
ID string `json:"id"`
Name string `json:"name"`
BaseURL string `json:"base_url"`
ProtocolVersion int `json:"protocol_version"`
MaxWorkers int `json:"max_workers"`
MaxRunning int `json:"max_running"`
ReportedWorkers int `json:"reported_workers"`
ReportedRunning int `json:"reported_running"`
Enabled bool `json:"enabled"`
Draining bool `json:"draining"`
LastSeen time.Time `json:"last_seen"`
LastError string `json:"last_error,omitempty"`
CreatedAt time.Time `json:"created_at"`
UpdatedAt time.Time `json:"updated_at"`
}
func scanServiceController(row interface{ Scan(...any) error }) (ServiceController, error) {
var c ServiceController
var enabled, draining int
var seen, created, updated int64
err := row.Scan(&c.ID, &c.Name, &c.BaseURL, &c.ProtocolVersion, &c.MaxWorkers, &c.MaxRunning, &c.ReportedWorkers, &c.ReportedRunning, &enabled, &draining, &seen, &c.LastError, &created, &updated)
c.Enabled = enabled != 0
c.Draining = draining != 0
c.LastSeen = time.UnixMilli(seen).UTC()
c.CreatedAt = time.UnixMilli(created).UTC()
c.UpdatedAt = time.UnixMilli(updated).UTC()
return c, err
}
const controllerCols = `id,name,base_url,protocol_version,max_workers,max_running,reported_workers,reported_running,enabled,draining,last_seen,last_error,created_at,updated_at`
func (s *Store) UpsertServiceController(ctx context.Context, c ServiceController) error {
now := time.Now().UTC().UnixMilli()
if c.MaxWorkers <= 0 {
c.MaxWorkers = 1000
}
if c.MaxRunning <= 0 {
c.MaxRunning = 100
}
if c.ProtocolVersion <= 0 {
c.ProtocolVersion = 1
}
_, err := s.DB.ExecContext(ctx, `INSERT INTO service_controllers(id,name,base_url,protocol_version,max_workers,max_running,reported_workers,reported_running,enabled,draining,last_seen,last_error,created_at,updated_at) VALUES(?,?,?,?,?,?,?,?,1,0,?,'',?,?)
ON CONFLICT(id) DO UPDATE SET name=excluded.name,base_url=excluded.base_url,protocol_version=excluded.protocol_version,max_workers=excluded.max_workers,max_running=excluded.max_running,reported_workers=excluded.reported_workers,reported_running=excluded.reported_running,last_seen=excluded.last_seen,last_error='',updated_at=excluded.updated_at`, strings.TrimSpace(c.ID), strings.TrimSpace(c.Name), strings.TrimRight(strings.TrimSpace(c.BaseURL), "/"), c.ProtocolVersion, c.MaxWorkers, c.MaxRunning, c.ReportedWorkers, c.ReportedRunning, now, now, now)
return err
}
func (s *Store) ServiceController(ctx context.Context, id string) (ServiceController, error) {
return scanServiceController(s.DB.QueryRowContext(ctx, `SELECT `+controllerCols+` FROM service_controllers WHERE id=?`, strings.TrimSpace(id)))
}
func (s *Store) ServiceControllers(ctx context.Context) ([]ServiceController, error) {
rows, err := s.DB.QueryContext(ctx, `SELECT `+controllerCols+` FROM service_controllers ORDER BY name,id`)
if err != nil {
return nil, err
}
defer rows.Close()
var out []ServiceController
for rows.Next() {
c, err := scanServiceController(rows)
if err != nil {
return nil, err
}
out = append(out, c)
}
return out, rows.Err()
}
func (s *Store) SetServiceControllerPolicy(ctx context.Context, id string, enabled, draining bool) error {
_, err := s.DB.ExecContext(ctx, `UPDATE service_controllers SET enabled=?,draining=?,updated_at=? WHERE id=?`, boolInt(enabled), boolInt(draining), time.Now().UTC().UnixMilli(), strings.TrimSpace(id))
return err
}
func (s *Store) ControllerWorkerCounts(ctx context.Context, id string) (total, running int, err error) {
err = s.DB.QueryRowContext(ctx, `SELECT count(*),COALESCE(sum(CASE WHEN status IN ('running','starting') THEN 1 ELSE 0 END),0) FROM workers WHERE controller_id=?`, strings.TrimSpace(id)).Scan(&total, &running)
return
}
func boolInt(v bool) int {
if v {
return 1
}
return 0
}
var ErrInvalidInvite = errors.New("registration invite invalid, expired or already used")
func scanCustomer(row interface{ Scan(...any) error }, withPassword bool) (Customer, string, string, error) {
var c Customer
var blocked, workerLimitBypass int
var blockedAt sql.NullInt64
var created int64
var salt, hash string
args := []any{&c.ID, &c.Username, &c.RewardClientID, &blocked, &c.BlockedReason, &blockedAt, &workerLimitBypass, &created}
if withPassword {
args = append(args, &salt, &hash)
}
err := row.Scan(args...)
c.Blocked = blocked != 0
c.WorkerLimitBypass = workerLimitBypass != 0
c.CreatedAt = time.UnixMilli(created).UTC()
if blockedAt.Valid {
v := time.UnixMilli(blockedAt.Int64).UTC()
c.BlockedAt = &v
}
return c, salt, hash, err
}
// CreateCustomer atomically creates the account, consumes an optional one-shot
// invite and grants the configured signup credits. A failed invite therefore
// cannot leave behind a partially-created account or free ledger entry.
func (s *Store) CreateCustomer(ctx context.Context, id, username, salt, hash string, signupBonusMicros int64, inviteHash string) error {
tx, err := s.DB.BeginTx(ctx, nil)
if err != nil {
return err
}
defer tx.Rollback()
now := time.Now().UTC().UnixMilli()
if _, err := tx.ExecContext(ctx, `INSERT INTO customers(id,username,password_salt,password_hash,created_at,updated_at) VALUES(?,?,?,?,?,?)`, id, strings.TrimSpace(username), salt, hash, now, now); err != nil {
return err
}
if strings.TrimSpace(inviteHash) != "" {
res, err := tx.ExecContext(ctx, `UPDATE registration_invites SET used_by=?,used_at=? WHERE token_hash=? AND used_by IS NULL AND expires_at>?`, id, now, inviteHash, now)
if err != nil {
return err
}
n, _ := res.RowsAffected()
if n != 1 {
return ErrInvalidInvite
}
}
if signupBonusMicros > 0 {
if _, err := tx.ExecContext(ctx, `INSERT INTO credit_ledger(customer_id,delta_micros,reason,reference,created_at) VALUES(?,?,?,?,?)`, id, signupBonusMicros, "signup_bonus", "signup:"+id, now); err != nil {
return err
}
}
return tx.Commit()
}
func (s *Store) CustomerByUsername(ctx context.Context, username string) (Customer, string, string, error) {
return scanCustomer(s.DB.QueryRowContext(ctx, `SELECT id,username,reward_client_id,blocked,blocked_reason,blocked_at,worker_limit_bypass,created_at,password_salt,password_hash FROM customers WHERE username=?`, strings.TrimSpace(username)), true)
}
func (s *Store) CustomerByID(ctx context.Context, id string) (Customer, error) {
c, _, _, err := scanCustomer(s.DB.QueryRowContext(ctx, `SELECT id,username,reward_client_id,blocked,blocked_reason,blocked_at,worker_limit_bypass,created_at FROM customers WHERE id=?`, id), false)
return c, err
}
func (s *Store) CustomerByRewardClientID(ctx context.Context, rewardClientID string) (Customer, error) {
c, _, _, err := scanCustomer(s.DB.QueryRowContext(ctx, `SELECT id,username,reward_client_id,blocked,blocked_reason,blocked_at,worker_limit_bypass,created_at FROM customers WHERE reward_client_id=? ORDER BY created_at LIMIT 1`, strings.TrimSpace(rewardClientID)), false)
return c, err
}
func (s *Store) SetRewardClientID(ctx context.Context, id, cid string) error {
_, err := s.DB.ExecContext(ctx, `UPDATE customers SET reward_client_id=?,updated_at=? WHERE id=?`, strings.TrimSpace(cid), time.Now().UTC().UnixMilli(), id)
return err
}
func (s *Store) SetCustomerBlocked(ctx context.Context, id string, blocked bool, reason string) error {
now := time.Now().UTC().UnixMilli()
if blocked {
_, err := s.DB.ExecContext(ctx, `UPDATE customers SET blocked=1,blocked_reason=?,blocked_at=?,updated_at=? WHERE id=?`, strings.TrimSpace(reason), now, now, id)
return err
}
_, err := s.DB.ExecContext(ctx, `UPDATE customers SET blocked=0,blocked_reason='',blocked_at=NULL,updated_at=? WHERE id=?`, now, id)
return err
}
func (s *Store) SetCustomerWorkerLimitBypass(ctx context.Context, id string, bypass bool) error {
v := 0
if bypass {
v = 1
}
res, err := s.DB.ExecContext(ctx, `UPDATE customers SET worker_limit_bypass=?,updated_at=? WHERE id=?`, v, time.Now().UTC().UnixMilli(), id)
if err != nil {
return err
}
n, _ := res.RowsAffected()
if n == 0 {
return sql.ErrNoRows
}
return nil
}
func (s *Store) DeleteCustomerSessions(ctx context.Context, cid string) error {
_, err := s.DB.ExecContext(ctx, `DELETE FROM customer_sessions WHERE customer_id=?`, cid)
return err
}
func (s *Store) CreateSession(ctx context.Context, sid, cid string, ttl time.Duration) error {
now := time.Now().UTC()
_, err := s.DB.ExecContext(ctx, `INSERT INTO customer_sessions(id,customer_id,expires_at,created_at) VALUES(?,?,?,?)`, sid, cid, now.Add(ttl).UnixMilli(), now.UnixMilli())
return err
}
func (s *Store) SessionCustomer(ctx context.Context, sid string) (string, error) {
var cid string
err := s.DB.QueryRowContext(ctx, `SELECT s.customer_id FROM customer_sessions s JOIN customers c ON c.id=s.customer_id WHERE s.id=? AND s.expires_at>? AND c.blocked=0`, sid, time.Now().UTC().UnixMilli()).Scan(&cid)
return cid, err
}
func (s *Store) DeleteSession(ctx context.Context, sid string) {
_, _ = s.DB.ExecContext(ctx, `DELETE FROM customer_sessions WHERE id=?`, sid)
}
func (s *Store) BalanceMicros(ctx context.Context, cid string) (int64, error) {
var v sql.NullInt64
err := s.DB.QueryRowContext(ctx, `SELECT sum(delta_micros) FROM credit_ledger WHERE customer_id=?`, cid).Scan(&v)
if err != nil {
return 0, err
}
return v.Int64, nil
}
func (s *Store) AddLedger(ctx context.Context, cid string, delta int64, reason, ref string) error {
if delta == 0 {
return errors.New("zero credit change")
}
_, err := s.DB.ExecContext(ctx, `INSERT INTO credit_ledger(customer_id,delta_micros,reason,reference,created_at) VALUES(?,?,?,?,?)`, cid, delta, reason, ref, time.Now().UTC().UnixMilli())
return err
}
type LedgerItem struct {
DeltaMicros int64 `json:"delta_micros"`
Reason string `json:"reason"`
Reference string `json:"reference"`
CreatedAt time.Time `json:"created_at"`
}
func (s *Store) Ledger(ctx context.Context, cid string, limit int) ([]LedgerItem, error) {
if limit < 1 || limit > 200 {
limit = 50
}
rows, err := s.DB.QueryContext(ctx, `SELECT delta_micros,reason,reference,created_at FROM credit_ledger WHERE customer_id=? ORDER BY id DESC LIMIT ?`, cid, limit)
if err != nil {
return nil, err
}
defer rows.Close()
var out []LedgerItem
for rows.Next() {
var x LedgerItem
var ms int64
if err := rows.Scan(&x.DeltaMicros, &x.Reason, &x.Reference, &ms); err != nil {
return nil, err
}
x.CreatedAt = time.UnixMilli(ms).UTC()
out = append(out, x)
}
return out, rows.Err()
}
type LedgerSummaryItem struct {
DeltaMicros int64 `json:"delta_micros"`
Reason string `json:"reason"`
Count int `json:"count"`
CreatedAt time.Time `json:"created_at"`
Day string `json:"day"`
}
// LedgerSummary intentionally collapses high-frequency worker-minute and
// positive-tip entries into daily reason buckets. The raw ledger stays intact
// in SQLite for auditing and idempotency.
func (s *Store) LedgerSummary(ctx context.Context, cid string, limit int) ([]LedgerSummaryItem, error) {
if limit < 1 || limit > 120 {
limit = 45
}
rows, err := s.DB.QueryContext(ctx, `SELECT reason,COALESCE(sum(delta_micros),0),count(*),max(created_at),strftime('%Y-%m-%d',max(created_at)/1000,'unixepoch')
FROM credit_ledger WHERE customer_id=?
GROUP BY reason,strftime('%Y-%m-%d',created_at/1000,'unixepoch')
ORDER BY max(created_at) DESC LIMIT ?`, cid, limit)
if err != nil {
return nil, err
}
defer rows.Close()
var out []LedgerSummaryItem
for rows.Next() {
var x LedgerSummaryItem
var ms int64
if err := rows.Scan(&x.Reason, &x.DeltaMicros, &x.Count, &ms, &x.Day); err != nil {
return nil, err
}
x.CreatedAt = time.UnixMilli(ms).UTC()
out = append(out, x)
}
return out, rows.Err()
}
func (s *Store) AddLedgerOnce(ctx context.Context, cid string, delta int64, reason, ref string) (bool, error) {
if delta == 0 {
return false, nil
}
res, err := s.DB.ExecContext(ctx, `INSERT OR IGNORE INTO credit_ledger(customer_id,delta_micros,reason,reference,created_at) VALUES(?,?,?,?,?)`, cid, delta, reason, ref, time.Now().UTC().UnixMilli())
if err != nil {
return false, err
}
n, err := res.RowsAffected()
return n == 1, err
}
func (s *Store) SeedSetting(ctx context.Context, key, value string) error {
_, err := s.DB.ExecContext(ctx, `INSERT OR IGNORE INTO customer_settings(key,value,updated_at) VALUES(?,?,?)`, key, value, time.Now().UTC().UnixMilli())
return err
}
func (s *Store) Setting(ctx context.Context, key, fallback string) string {
var v string
if err := s.DB.QueryRowContext(ctx, `SELECT value FROM customer_settings WHERE key=?`, key).Scan(&v); err != nil {
return fallback
}
return v
}
func (s *Store) SetSetting(ctx context.Context, key, value string) error {
_, err := s.DB.ExecContext(ctx, `INSERT INTO customer_settings(key,value,updated_at) VALUES(?,?,?) ON CONFLICT(key) DO UPDATE SET value=excluded.value,updated_at=excluded.updated_at`, key, value, time.Now().UTC().UnixMilli())
return err
}
func (s *Store) CreateRegistrationInvite(ctx context.Context, tokenHash, label string, expiresAt time.Time) error {
_, err := s.DB.ExecContext(ctx, `INSERT INTO registration_invites(token_hash,label,expires_at,created_at) VALUES(?,?,?,?)`, tokenHash, strings.TrimSpace(label), expiresAt.UTC().UnixMilli(), time.Now().UTC().UnixMilli())
return err
}
func (s *Store) ActiveInviteCount(ctx context.Context) int {
var n int
_ = s.DB.QueryRowContext(ctx, `SELECT count(*) FROM registration_invites WHERE used_by IS NULL AND expires_at>?`, time.Now().UTC().UnixMilli()).Scan(&n)
return n
}
func (s *Store) WorkerByClientID(ctx context.Context, clientID string) (Worker, error) {
return scanWorker(s.DB.QueryRowContext(ctx, `SELECT `+workerCols+` FROM workers WHERE worker_client_id=? ORDER BY created_at DESC LIMIT 1`, strings.TrimSpace(clientID)))
}
func scanWorker(row interface{ Scan(...any) error }) (Worker, error) {
var w Worker
var last, lastLease sql.NullInt64
var created, updated int64
err := row.Scan(&w.ID, &w.CustomerID, &w.TaskID, &w.BeaconPath, &w.ContainerID, &w.ControllerID, &w.Volume, &w.Status, &w.WorkerClientID, &w.RegisterToken, &w.RateMicrosPerMinute, &last, &lastLease, &created, &updated, &w.LastError)
if err != nil {
return w, err
}
if last.Valid {
v := time.UnixMilli(last.Int64).UTC()
w.LastChargeAt = &v
}
if lastLease.Valid {
v := time.UnixMilli(lastLease.Int64).UTC()
w.LastLeaseAt = &v
}
w.CreatedAt = time.UnixMilli(created).UTC()
w.UpdatedAt = time.UnixMilli(updated).UTC()
return w, nil
}
const workerCols = `id,customer_id,task_id,beacon_path,docker_container_id,controller_id,docker_volume,status,worker_client_id,register_token,rate_micros_per_minute,last_charge_at,last_lease_at,created_at,updated_at,last_error`
func (s *Store) CreateWorker(ctx context.Context, w Worker) error {
now := time.Now().UTC().UnixMilli()
_, err := s.DB.ExecContext(ctx, `INSERT INTO workers(id,customer_id,task_id,beacon_path,controller_id,docker_volume,status,register_token,rate_micros_per_minute,created_at,updated_at) VALUES(?,?,?,?,?,?,'stopped',?,?,?,?)`, w.ID, w.CustomerID, w.TaskID, w.BeaconPath, w.ControllerID, w.Volume, w.RegisterToken, w.RateMicrosPerMinute, now, now)
return err
}
func (s *Store) Worker(ctx context.Context, cid, wid string) (Worker, error) {
return scanWorker(s.DB.QueryRowContext(ctx, `SELECT `+workerCols+` FROM workers WHERE id=? AND customer_id=?`, wid, cid))
}
func (s *Store) WorkerByID(ctx context.Context, wid string) (Worker, error) {
return scanWorker(s.DB.QueryRowContext(ctx, `SELECT `+workerCols+` FROM workers WHERE id=?`, wid))
}
func (s *Store) Workers(ctx context.Context, cid string) ([]Worker, error) {
rows, err := s.DB.QueryContext(ctx, `SELECT `+workerCols+` FROM workers WHERE customer_id=? ORDER BY created_at`, cid)
if err != nil {
return nil, err
}
defer rows.Close()
var out []Worker
for rows.Next() {
w, err := scanWorker(rows)
if err != nil {
return nil, err
}
out = append(out, w)
}
return out, rows.Err()
}
func (s *Store) WorkerCount(ctx context.Context, cid string, runningOnly bool) (int, error) {
q := `SELECT count(*) FROM workers WHERE customer_id=?`
if runningOnly {
q += ` AND status='running'`
}
var n int
err := s.DB.QueryRowContext(ctx, q, cid).Scan(&n)
return n, err
}
func (s *Store) TotalWorkerCount(ctx context.Context) (int, error) {
var n int
err := s.DB.QueryRowContext(ctx, `SELECT count(*) FROM workers`).Scan(&n)
return n, err
}
func (s *Store) RunningWorkerCount(ctx context.Context) (int, error) {
var n int
err := s.DB.QueryRowContext(ctx, `SELECT count(*) FROM workers WHERE status='running'`).Scan(&n)
return n, err
}
func (s *Store) RunningWorkers(ctx context.Context) ([]Worker, error) {
rows, err := s.DB.QueryContext(ctx, `SELECT `+workerCols+` FROM workers WHERE status='running' ORDER BY created_at`)
if err != nil {
return nil, err
}
defer rows.Close()
var out []Worker
for rows.Next() {
w, err := scanWorker(rows)
if err != nil {
return nil, err
}
out = append(out, w)
}
return out, rows.Err()
}
func (s *Store) ClaimWorkerStart(ctx context.Context, cid, wid string) (bool, error) {
res, err := s.DB.ExecContext(ctx, `UPDATE workers SET status='starting',last_error='',updated_at=? WHERE id=? AND customer_id=? AND status IN ('stopped','error')`, time.Now().UTC().UnixMilli(), wid, cid)
if err != nil {
return false, err
}
n, err := res.RowsAffected()
if err != nil {
return false, err
}
return n == 1, nil
}
func (s *Store) RecoverStartingWorkers(ctx context.Context) error {
_, err := s.DB.ExecContext(ctx, `UPDATE workers SET status='stopped',last_error='recovered after Customer Service restart',updated_at=? WHERE status='starting'`, time.Now().UTC().UnixMilli())
return err
}
func (s *Store) SetWorkerRuntime(ctx context.Context, wid, status, containerID, lastErr string) error {
_, err := s.DB.ExecContext(ctx, `UPDATE workers SET status=?,docker_container_id=?,last_error=?,updated_at=? WHERE id=?`, status, containerID, lastErr, time.Now().UTC().UnixMilli(), wid)
return err
}
func (s *Store) SetWorkerController(ctx context.Context, wid, controllerID string) error {
_, err := s.DB.ExecContext(ctx, `UPDATE workers SET controller_id=?,updated_at=? WHERE id=?`, strings.TrimSpace(controllerID), time.Now().UTC().UnixMilli(), wid)
return err
}
func (s *Store) TouchWorkerLease(ctx context.Context, wid string) error {
_, err := s.DB.ExecContext(ctx, `UPDATE workers SET last_lease_at=?,updated_at=? WHERE id=?`, time.Now().UTC().UnixMilli(), time.Now().UTC().UnixMilli(), wid)
return err
}
func (s *Store) SetWorkerClient(ctx context.Context, wid, clientID string) error {
_, err := s.DB.ExecContext(ctx, `UPDATE workers SET worker_client_id=?,updated_at=? WHERE id=?`, clientID, time.Now().UTC().UnixMilli(), wid)
return err
}
func (s *Store) UpdateWorkerConfig(ctx context.Context, cid, wid, taskID, beaconPath string) error {
_, err := s.DB.ExecContext(ctx, `UPDATE workers SET task_id=?,beacon_path=?,updated_at=? WHERE id=? AND customer_id=?`, taskID, beaconPath, time.Now().UTC().UnixMilli(), wid, cid)
return err
}
func (s *Store) DeleteWorker(ctx context.Context, cid, wid string) error {
_, err := s.DB.ExecContext(ctx, `DELETE FROM workers WHERE id=? AND customer_id=?`, wid, cid)
return err
}
func (s *Store) MarkWorkerCharged(ctx context.Context, wid string, when time.Time) error {
_, err := s.DB.ExecContext(ctx, `UPDATE workers SET last_charge_at=?,updated_at=? WHERE id=?`, when.UTC().UnixMilli(), time.Now().UTC().UnixMilli(), wid)
return err
}
// ChargeWorkerMinute debits one prepaid minute atomically. It never allows a
// negative balance, so a billing loop can stop the worker as soon as funding is
// exhausted.
func (s *Store) ChargeWorkerMinute(ctx context.Context, w Worker, minute time.Time) (bool, error) {
tx, err := s.DB.BeginTx(ctx, nil)
if err != nil {
return false, err
}
defer tx.Rollback()
var bal sql.NullInt64
if err := tx.QueryRowContext(ctx, `SELECT sum(delta_micros) FROM credit_ledger WHERE customer_id=?`, w.CustomerID).Scan(&bal); err != nil {
return false, err
}
if bal.Int64 < w.RateMicrosPerMinute {
return false, nil
}
ref := fmt.Sprintf("worker:%s:%d", w.ID, minute.UTC().UnixMilli())
if _, err := tx.ExecContext(ctx, `INSERT OR IGNORE INTO credit_ledger(customer_id,delta_micros,reason,reference,created_at) VALUES(?,?,?,?,?)`, w.CustomerID, -w.RateMicrosPerMinute, "worker_minute", ref, time.Now().UTC().UnixMilli()); err != nil {
return false, err
}
if _, err := tx.ExecContext(ctx, `UPDATE workers SET last_charge_at=?,updated_at=? WHERE id=?`, minute.UTC().UnixMilli(), time.Now().UTC().UnixMilli(), w.ID); err != nil {
return false, err
}
if err := tx.Commit(); err != nil {
return false, err
}
return true, nil
}
func (s *Store) RefundWorkerStartMinute(ctx context.Context, w Worker, chargedAt time.Time, detail string) error {
ref := fmt.Sprintf("worker_start_refund:%s:%d", w.ID, chargedAt.UTC().UnixMilli())
reason := "worker_start_refund"
if strings.TrimSpace(detail) != "" {
reason += ":" + strings.TrimSpace(detail)
}
_, err := s.DB.ExecContext(ctx, `INSERT OR IGNORE INTO credit_ledger(customer_id,delta_micros,reason,reference,created_at) VALUES(?,?,?,?,?)`, w.CustomerID, w.RateMicrosPerMinute, reason, ref, time.Now().UTC().UnixMilli())
return err
}
func (s *Store) UpsertPayPalOrder(ctx context.Context, orderID, cid, pkg string, cents int64, currency string, credits int64, status string) error {
now := time.Now().UTC().UnixMilli()
_, err := s.DB.ExecContext(ctx, `INSERT INTO paypal_orders(order_id,customer_id,package_id,amount_cents,currency,credits_micros,status,created_at,updated_at) VALUES(?,?,?,?,?,?,?,?,?) ON CONFLICT(order_id) DO UPDATE SET status=excluded.status,updated_at=excluded.updated_at`, orderID, cid, pkg, cents, currency, credits, status, now, now)
return err
}
type PayPalOrder struct {
OrderID, CustomerID, PackageID, Currency, Status, CaptureID string
AmountCents, CreditsMicros int64
}
func (s *Store) PayPalOrder(ctx context.Context, orderID string) (PayPalOrder, error) {
var o PayPalOrder
err := s.DB.QueryRowContext(ctx, `SELECT order_id,customer_id,package_id,amount_cents,currency,credits_micros,status,capture_id FROM paypal_orders WHERE order_id=?`, orderID).Scan(&o.OrderID, &o.CustomerID, &o.PackageID, &o.AmountCents, &o.Currency, &o.CreditsMicros, &o.Status, &o.CaptureID)
return o, err
}
func (s *Store) CompletePayPalOrder(ctx context.Context, orderID, captureID string) error {
tx, err := s.DB.BeginTx(ctx, nil)
if err != nil {
return err
}
defer tx.Rollback()
var o PayPalOrder
if err := tx.QueryRowContext(ctx, `SELECT order_id,customer_id,package_id,amount_cents,currency,credits_micros,status,capture_id FROM paypal_orders WHERE order_id=?`, orderID).Scan(&o.OrderID, &o.CustomerID, &o.PackageID, &o.AmountCents, &o.Currency, &o.CreditsMicros, &o.Status, &o.CaptureID); err != nil {
return err
}
ref := "paypal:" + orderID
if _, err := tx.ExecContext(ctx, `INSERT OR IGNORE INTO credit_ledger(customer_id,delta_micros,reason,reference,created_at) VALUES(?,?,?,?,?)`, o.CustomerID, o.CreditsMicros, "paypal_topup", ref, time.Now().UTC().UnixMilli()); err != nil {
return err
}
if _, err := tx.ExecContext(ctx, `UPDATE paypal_orders SET status='COMPLETED',capture_id=?,updated_at=? WHERE order_id=?`, captureID, time.Now().UTC().UnixMilli(), orderID); err != nil {
return err
}
return tx.Commit()
}