pangolin/server/internal/provision/mysqlstore.go

package provision

import (
	"context"
	"database/sql"
	"encoding/json"
	"fmt"
	"strings"
	"time"

	dbx "github.com/wangjia/pangolin/server/internal/db"
)

// MySQLStore is the production Store backed by the shared *sql.DB pool.
// (Despite the name it is dialect-aware and also works against SQLite.)
type MySQLStore struct {
	db      *sql.DB
	dialect dbx.Dialect
}

// NewMySQLStore wraps a database connection pool (MySQL or SQLite).
func NewMySQLStore(db *sql.DB) *MySQLStore {
	return &MySQLStore{db: db, dialect: dbx.DialectForDB(db)}
}

var _ Store = (*MySQLStore)(nil)

func marshalJSONList(v []string) string {
	if v == nil {
		v = []string{}
	}
	b, _ := json.Marshal(v)
	return string(b)
}

func unmarshalJSONList(s sql.NullString) []string {
	if !s.Valid || s.String == "" {
		return nil
	}
	var out []string
	_ = json.Unmarshal([]byte(s.String), &out)
	return out
}

// --- nodes ---

func (s *MySQLStore) InsertNode(ctx context.Context, n *Node) (int64, error) {
	endpoint := n.Endpoint
	if endpoint == "" {
		endpoint = pendingEndpoint
	}
	status := n.Status
	if status == "" {
		status = StatusProvisioning
	}
	weight := n.Weight
	if weight == 0 {
		weight = 100
	}
	res, err := s.db.ExecContext(ctx,
		`INSERT INTO nodes
		  (uuid, region, name_zh, name_en, role, tier, endpoint, hy2_port,
		   reality_pbk, reality_sni, provider_id, tags, status, weight, created_at)
		 VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`,
		n.UUID, n.Region, n.NameZH, n.NameEn, string(n.Role), string(n.Tier),
		endpoint, nullInt(n.HY2Port), n.RealityPBK, n.RealitySNI, n.ProviderID,
		marshalJSONList(n.Tags), string(status), weight, time.Now().UTC())
	if err != nil {
		return 0, fmt.Errorf("store.InsertNode: %w", err)
	}
	id, err := res.LastInsertId()
	if err != nil {
		return 0, fmt.Errorf("store.InsertNode last id: %w", err)
	}
	return id, nil
}

const nodeColumns = `id, uuid, region, name_zh, name_en, role, tier, endpoint,
	hy2_port, reality_pbk, reality_sni, provider_id, provider_instance_id,
	elastic_ip_id, tags, status, weight, created_at`

func scanNode(row interface{ Scan(...any) error }) (*Node, error) {
	var (
		n          Node
		hy2        sql.NullInt64
		instanceID sql.NullString
		elasticID  sql.NullString
		tags       sql.NullString
	)
	if err := row.Scan(
		&n.ID, &n.UUID, &n.Region, &n.NameZH, &n.NameEn, &n.Role, &n.Tier,
		&n.Endpoint, &hy2, &n.RealityPBK, &n.RealitySNI, &n.ProviderID,
		&instanceID, &elasticID, &tags, &n.Status, &n.Weight, &n.CreatedAt,
	); err != nil {
		return nil, err
	}
	n.HY2Port = int(hy2.Int64)
	n.ProviderInstanceID = instanceID.String
	n.ElasticIPID = elasticID.String
	n.Tags = unmarshalJSONList(tags)
	return &n, nil
}

func (s *MySQLStore) GetNode(ctx context.Context, id int64) (*Node, error) {
	row := s.db.QueryRowContext(ctx, `SELECT `+nodeColumns+` FROM nodes WHERE id=?`, id)
	n, err := scanNode(row)
	if err == sql.ErrNoRows {
		return nil, nil
	}
	if err != nil {
		return nil, fmt.Errorf("store.GetNode: %w", err)
	}
	return n, nil
}

func (s *MySQLStore) GetNodeByUUID(ctx context.Context, uuid string) (*Node, error) {
	row := s.db.QueryRowContext(ctx, `SELECT `+nodeColumns+` FROM nodes WHERE uuid=?`, uuid)
	n, err := scanNode(row)
	if err == sql.ErrNoRows {
		return nil, nil
	}
	if err != nil {
		return nil, fmt.Errorf("store.GetNodeByUUID: %w", err)
	}
	return n, nil
}

func (s *MySQLStore) UpdateNodeStatus(ctx context.Context, id int64, status Status) error {
	_, err := s.db.ExecContext(ctx, `UPDATE nodes SET status=? WHERE id=?`, string(status), id)
	if err != nil {
		return fmt.Errorf("store.UpdateNodeStatus: %w", err)
	}
	return nil
}

func (s *MySQLStore) UpdateNodeEndpoint(ctx context.Context, id int64, endpoint string) error {
	_, err := s.db.ExecContext(ctx, `UPDATE nodes SET endpoint=? WHERE id=?`, endpoint, id)
	if err != nil {
		return fmt.Errorf("store.UpdateNodeEndpoint: %w", err)
	}
	return nil
}

func (s *MySQLStore) SetNodeInstance(ctx context.Context, id int64, instanceID, endpoint string) error {
	_, err := s.db.ExecContext(ctx,
		`UPDATE nodes SET provider_instance_id=?, endpoint=? WHERE id=?`,
		instanceID, endpoint, id)
	if err != nil {
		return fmt.Errorf("store.SetNodeInstance: %w", err)
	}
	return nil
}

func (s *MySQLStore) SetNodeWeight(ctx context.Context, id int64, weight int) error {
	_, err := s.db.ExecContext(ctx, `UPDATE nodes SET weight=? WHERE id=?`, weight, id)
	if err != nil {
		return fmt.Errorf("store.SetNodeWeight: %w", err)
	}
	return nil
}

func (s *MySQLStore) ListNodesByPool(ctx context.Context, pool Pool, status Status) ([]*Node, error) {
	q := `SELECT ` + qualify(nodeColumns, "n") + `
	      FROM nodes n JOIN providers p ON p.id = n.provider_id
	      WHERE p.pool = ?`
	args := []any{string(pool)}
	if status != "" {
		q += ` AND n.status = ?`
		args = append(args, string(status))
	}
	q += ` ORDER BY n.id`
	rows, err := s.db.QueryContext(ctx, q, args...)
	if err != nil {
		return nil, fmt.Errorf("store.ListNodesByPool: %w", err)
	}
	defer rows.Close()
	var out []*Node
	for rows.Next() {
		n, err := scanNode(rows)
		if err != nil {
			return nil, fmt.Errorf("store.ListNodesByPool scan: %w", err)
		}
		out = append(out, n)
	}
	return out, rows.Err()
}

// --- providers ---

func (s *MySQLStore) ListProviders(ctx context.Context, pool Pool) ([]*Provider, error) {
	q := `SELECT id, name, api_kind, regions, pool, enabled FROM providers WHERE enabled=TRUE`
	var args []any
	if pool != "" {
		q += ` AND pool=?`
		args = append(args, string(pool))
	}
	q += ` ORDER BY id`
	rows, err := s.db.QueryContext(ctx, q, args...)
	if err != nil {
		return nil, fmt.Errorf("store.ListProviders: %w", err)
	}
	defer rows.Close()
	var out []*Provider
	for rows.Next() {
		p, err := scanProvider(rows)
		if err != nil {
			return nil, fmt.Errorf("store.ListProviders scan: %w", err)
		}
		out = append(out, p)
	}
	return out, rows.Err()
}

func (s *MySQLStore) GetProvider(ctx context.Context, id int64) (*Provider, error) {
	row := s.db.QueryRowContext(ctx,
		`SELECT id, name, api_kind, regions, pool, enabled FROM providers WHERE id=?`, id)
	p, err := scanProvider(row)
	if err == sql.ErrNoRows {
		return nil, nil
	}
	if err != nil {
		return nil, fmt.Errorf("store.GetProvider: %w", err)
	}
	return p, nil
}

func scanProvider(row interface{ Scan(...any) error }) (*Provider, error) {
	var (
		p       Provider
		regions sql.NullString
	)
	if err := row.Scan(&p.ID, &p.Name, &p.APIKind, &regions, &p.Pool, &p.Enabled); err != nil {
		return nil, err
	}
	p.Regions = unmarshalJSONList(regions)
	return &p, nil
}

// --- events / audit / directory ---

func (s *MySQLStore) WriteNodeEvent(ctx context.Context, nodeID int64, event Event, detailJSON string) error {
	if detailJSON == "" {
		detailJSON = "null"
	}
	_, err := s.db.ExecContext(ctx,
		`INSERT INTO node_events (node_id, event, detail, at) VALUES (?, ?, ?, ?)`,
		nodeID, string(event), detailJSON, time.Now().UTC())
	if err != nil {
		return fmt.Errorf("store.WriteNodeEvent: %w", err)
	}
	return nil
}

func (s *MySQLStore) WriteAuditLog(ctx context.Context, actor, action, target, metaJSON string) error {
	if metaJSON == "" {
		metaJSON = "null"
	}
	_, err := s.db.ExecContext(ctx,
		`INSERT INTO audit_log (actor, action, target, meta, at) VALUES (?, ?, ?, ?, ?)`,
		actor, action, target, metaJSON, time.Now().UTC())
	if err != nil {
		return fmt.Errorf("store.WriteAuditLog: %w", err)
	}
	return nil
}

func (s *MySQLStore) BumpDirectoryVersion(ctx context.Context) (int64, error) {
	if _, err := s.db.ExecContext(ctx,
		`UPDATE directory_version SET version = version + 1 WHERE id = 1`); err != nil {
		return 0, fmt.Errorf("store.BumpDirectoryVersion: %w", err)
	}
	var v int64
	if err := s.db.QueryRowContext(ctx,
		`SELECT version FROM directory_version WHERE id = 1`).Scan(&v); err != nil {
		return 0, fmt.Errorf("store.BumpDirectoryVersion read: %w", err)
	}
	return v, nil
}

// --- idempotency ---

func (s *MySQLStore) LookupIdempotency(ctx context.Context, key string) (string, bool, error) {
	var uuid string
	err := s.db.QueryRowContext(ctx,
		`SELECT node_uuid FROM provision_idempotency WHERE idempotency_key=?`, key).Scan(&uuid)
	if err == sql.ErrNoRows {
		return "", false, nil
	}
	if err != nil {
		return "", false, fmt.Errorf("store.LookupIdempotency: %w", err)
	}
	return uuid, true, nil
}

func (s *MySQLStore) SaveIdempotency(ctx context.Context, key, nodeUUID string) error {
	_, err := s.db.ExecContext(ctx,
		`INSERT INTO provision_idempotency (idempotency_key, node_uuid, created_at)
		 VALUES (?, ?, ?) `+
			s.dialect.Upsert([]string{"idempotency_key"}),
		key, nodeUUID, time.Now().UTC())
	if err != nil {
		return fmt.Errorf("store.SaveIdempotency: %w", err)
	}
	return nil
}

// --- replacements ---

func (s *MySQLStore) CreateReplacement(ctx context.Context, r *Replacement) error {
	_, err := s.db.ExecContext(ctx,
		`INSERT INTO replacements (uuid, old_node_id, new_node_id, pool, status, step)
		 VALUES (?, ?, ?, ?, ?, ?)`,
		r.UUID, r.OldNodeID, nullInt64(r.NewNodeID), string(r.Pool),
		string(r.Status), string(r.Step))
	if err != nil {
		return fmt.Errorf("store.CreateReplacement: %w", err)
	}
	return nil
}

func (s *MySQLStore) GetReplacement(ctx context.Context, uuid string) (*Replacement, error) {
	var (
		r     Replacement
		newID sql.NullInt64
	)
	err := s.db.QueryRowContext(ctx,
		`SELECT uuid, old_node_id, new_node_id, pool, status, step, created_at, updated_at
		 FROM replacements WHERE uuid=?`, uuid).
		Scan(&r.UUID, &r.OldNodeID, &newID, &r.Pool, &r.Status, &r.Step, &r.CreatedAt, &r.UpdatedAt)
	if err == sql.ErrNoRows {
		return nil, nil
	}
	if err != nil {
		return nil, fmt.Errorf("store.GetReplacement: %w", err)
	}
	r.NewNodeID = newID.Int64
	return &r, nil
}

func (s *MySQLStore) UpdateReplacement(ctx context.Context, r *Replacement) error {
	_, err := s.db.ExecContext(ctx,
		`UPDATE replacements SET new_node_id=?, status=?, step=? WHERE uuid=?`,
		nullInt64(r.NewNodeID), string(r.Status), string(r.Step), r.UUID)
	if err != nil {
		return fmt.Errorf("store.UpdateReplacement: %w", err)
	}
	return nil
}

// --- helpers ---

func nullInt(v int) interface{} {
	if v == 0 {
		return nil
	}
	return v
}

func nullInt64(v int64) interface{} {
	if v == 0 {
		return nil
	}
	return v
}

// qualify prefixes every comma-separated column in cols with the alias.
func qualify(cols, alias string) string {
	parts := strings.Split(cols, ",")
	for i, p := range parts {
		parts[i] = alias + "." + strings.TrimSpace(p)
	}
	return strings.Join(parts, ", ")
}