90CNServer/server/internal/platform/onlinepush/persistent_backends.go

package onlinepush

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

	cachekit "longheng.io/server/internal/platform/cache"
)

var (
	// ErrCacheRequired 表示缓存状态后端缺失。
	ErrCacheRequired = errors.New("online push cache is required")
	// ErrSQLRequired 表示 SQL 状态后端缺失。
	ErrSQLRequired = errors.New("online push sql db is required")
	// ErrInvalidKey 表示状态 key 为空或非法。
	ErrInvalidKey = errors.New("online push state key is invalid")
	// ErrInvalidTable 表示 SQL 表名为空或非法。
	ErrInvalidTable = errors.New("online push table is invalid")
)

// CacheStateStore 使用框架 cache.Store 保存在线推送状态快照。
type CacheStateStore struct {
	Cache cachekit.Store
	Key   string
	TTL   time.Duration
}

// NewCacheStateStore 创建缓存状态后端。
func NewCacheStateStore(cache cachekit.Store, key string, ttl time.Duration) *CacheStateStore {
	return &CacheStateStore{Cache: cache, Key: strings.TrimSpace(key), TTL: ttl}
}

// Load 从缓存读取在线推送状态。
func (s *CacheStateStore) Load(ctx context.Context) (State, bool, error) {
	if err := ctxErr(ctx); err != nil {
		return State{}, false, err
	}
	if s == nil || s.Cache == nil {
		return State{}, false, ErrCacheRequired
	}
	key := strings.TrimSpace(s.Key)
	if key == "" {
		return State{}, false, ErrInvalidKey
	}
	data, ok, err := s.Cache.Get(ctx, key)
	if err != nil || !ok {
		return State{}, ok, err
	}
	var state State
	if err := json.Unmarshal(data, &state); err != nil {
		return State{}, false, err
	}
	return state, true, nil
}

// Save 将在线推送状态保存到缓存。
func (s *CacheStateStore) Save(ctx context.Context, state State) error {
	if err := ctxErr(ctx); err != nil {
		return err
	}
	if s == nil || s.Cache == nil {
		return ErrCacheRequired
	}
	key := strings.TrimSpace(s.Key)
	if key == "" {
		return ErrInvalidKey
	}
	data, err := json.Marshal(state)
	if err != nil {
		return err
	}
	return s.Cache.Set(ctx, key, data, s.TTL)
}

// Close 关闭缓存状态后端。
func (s *CacheStateStore) Close() error {
	return nil
}

// SQLStateStore 使用 SQL 单行 JSON 保存在线推送状态快照。
type SQLStateStore struct {
	DB     *sql.DB
	Table  string
	Key    string
	Now    func() time.Time
	stmts  SQLStatements
	initMu sync.Mutex
}

// SQLStatements 保存在线推送 SQL 状态后端的语句模板。
type SQLStatements struct {
	Table        string
	EnsureSchema string
	Select       string
	Upsert       string
}

var tableNamePattern = regexp.MustCompile(`^[A-Za-z0-9_]+$`)

// NewSQLStateStore 创建 SQL 状态后端。
func NewSQLStateStore(db *sql.DB, table, key string, now func() time.Time) (*SQLStateStore, error) {
	if db == nil {
		return nil, ErrSQLRequired
	}
	key = strings.TrimSpace(key)
	if key == "" {
		return nil, ErrInvalidKey
	}
	stmts, err := NewSQLStatements(table)
	if err != nil {
		return nil, err
	}
	if now == nil {
		now = func() time.Time { return time.Now().UTC() }
	}
	return &SQLStateStore{DB: db, Table: stmts.Table, Key: key, Now: now, stmts: stmts}, nil
}

// NewSQLStatements 为指定白名单表名生成 SQL 语句。
func NewSQLStatements(table string) (SQLStatements, error) {
	table = strings.TrimSpace(table)
	if table == "" {
		table = "online_push_state"
	}
	if !tableNamePattern.MatchString(table) {
		return SQLStatements{}, ErrInvalidTable
	}
	quoted := "`" + table + "`"
	return SQLStatements{
		Table: table,
		EnsureSchema: fmt.Sprintf(`CREATE TABLE IF NOT EXISTS %s (
  state_key VARCHAR(191) NOT NULL,
  state_json JSON NOT NULL,
  updated_at TIMESTAMP(6) NOT NULL,
  PRIMARY KEY (state_key)
)`, quoted),
		Select: fmt.Sprintf(`SELECT state_json FROM %s WHERE state_key=?`, quoted),
		Upsert: fmt.Sprintf(`INSERT INTO %s (state_key, state_json, updated_at) VALUES (?, ?, ?) ON DUPLICATE KEY UPDATE state_json=VALUES(state_json), updated_at=VALUES(updated_at)`, quoted),
	}, nil
}

func (s *SQLStateStore) ensureReady(requireKey bool) error {
	if s == nil || s.DB == nil {
		return ErrSQLRequired
	}
	s.initMu.Lock()
	defer s.initMu.Unlock()
	table := strings.TrimSpace(s.Table)
	if table == "" {
		table = strings.TrimSpace(s.stmts.Table)
	}
	stmts, err := NewSQLStatements(table)
	if err != nil {
		return err
	}
	if s.stmts.Table != stmts.Table || s.stmts.EnsureSchema == "" || s.stmts.Select == "" || s.stmts.Upsert == "" {
		s.stmts = stmts
	}
	s.Table = s.stmts.Table
	s.Key = strings.TrimSpace(s.Key)
	if requireKey && s.Key == "" {
		return ErrInvalidKey
	}
	if s.Now == nil {
		s.Now = func() time.Time { return time.Now().UTC() }
	}
	return nil
}

// EnsureSchema 确保 SQL 状态表存在。
func (s *SQLStateStore) EnsureSchema(ctx context.Context) error {
	if err := ctxErr(ctx); err != nil {
		return err
	}
	if err := s.ensureReady(false); err != nil {
		return err
	}
	_, err := s.DB.ExecContext(ctx, s.stmts.EnsureSchema)
	return err
}

// Load 从 SQL 状态表读取在线推送状态。
func (s *SQLStateStore) Load(ctx context.Context) (State, bool, error) {
	if err := ctxErr(ctx); err != nil {
		return State{}, false, err
	}
	if err := s.ensureReady(true); err != nil {
		return State{}, false, err
	}
	var data []byte
	err := s.DB.QueryRowContext(ctx, s.stmts.Select, s.Key).Scan(&data)
	if errors.Is(err, sql.ErrNoRows) {
		return State{}, false, nil
	}
	if err != nil {
		return State{}, false, err
	}
	var state State
	if err := json.Unmarshal(data, &state); err != nil {
		return State{}, false, err
	}
	return state, true, nil
}

// Save 将在线推送状态保存到 SQL 状态表。
func (s *SQLStateStore) Save(ctx context.Context, state State) error {
	if err := ctxErr(ctx); err != nil {
		return err
	}
	if err := s.ensureReady(true); err != nil {
		return err
	}
	data, err := json.Marshal(state)
	if err != nil {
		return err
	}
	_, err = s.DB.ExecContext(ctx, s.stmts.Upsert, s.Key, data, s.Now().UTC()) // #nosec G701 -- Upsert 由 NewSQLStatements 使用白名单表名生成，业务值通过参数绑定。
	return err
}

// Close 关闭 SQL 状态后端。
func (s *SQLStateStore) Close() error {
	return nil
}