opsxdb/server/internal/db/redis_exec_test.go

package db

import (
	"context"
	"net"
	"strconv"
	"testing"
	"time"

	"github.com/alicebob/miniredis/v2"
)

func redisTestReq(t *testing.T, mr *miniredis.Miniredis) ConnReq {
	t.Helper()
	host, portStr, err := net.SplitHostPort(mr.Addr())
	if err != nil {
		t.Fatal(err)
	}
	port, _ := strconv.Atoi(portStr)
	return ConnReq{Driver: "redis", Host: host, Port: port, Database: "0"}
}

// 写命令白名单:SET EX / HSET / EXPIRE / MSET / INCR 等应直接可用。
func TestExecRedisWriteWhitelist(t *testing.T) {
	mr, err := miniredis.Run()
	if err != nil {
		t.Fatal(err)
	}
	defer mr.Close()
	base := redisTestReq(t, mr)
	ctx := context.Background()

	run := func(sql string) ExecResult {
		c := base
		c.SQL = sql
		return execRedis(ctx, c)
	}

	if r := run(`SET k1 "v with space" EX 100`); r.Error != "" {
		t.Fatalf("SET EX: %s", r.Error)
	}
	if v, _ := mr.Get("k1"); v != "v with space" {
		t.Fatalf("SET 引号分组值被破坏: %q", v)
	}
	if ttl := mr.TTL("k1"); ttl != 100*time.Second {
		t.Fatalf("SET EX TTL: %v", ttl)
	}

	if r := run("HSET h1 f1 v1 f2 v2"); r.Error != "" {
		t.Fatalf("HSET: %s", r.Error)
	}
	if mr.HGet("h1", "f2") != "v2" {
		t.Fatal("HSET 未生效")
	}

	if r := run("EXPIRE h1 60"); r.Error != "" {
		t.Fatalf("EXPIRE: %s", r.Error)
	}
	if ttl := mr.TTL("h1"); ttl != 60*time.Second {
		t.Fatalf("EXPIRE 未生效: %v", ttl)
	}
	if r := run("PERSIST h1"); r.Error != "" {
		t.Fatalf("PERSIST: %s", r.Error)
	}
	if ttl := mr.TTL("h1"); ttl != 0 {
		t.Fatalf("PERSIST 未生效: %v", ttl)
	}

	if r := run("MSET a 1 b 2"); r.Error != "" {
		t.Fatalf("MSET: %s", r.Error)
	}
	if r := run("MGET a b"); r.Error != "" || r.RowCount != 2 {
		t.Fatalf("MGET: %#v", r)
	}
	if r := run("INCR n1"); r.Error != "" {
		t.Fatalf("INCR: %s", r.Error)
	}
	if r := run("SADD s1 m1 m2"); r.Error != "" {
		t.Fatalf("SADD: %s", r.Error)
	}
	if r := run("ZADD z1 1.5 zm"); r.Error != "" {
		t.Fatalf("ZADD: %s", r.Error)
	}
	if r := run("LPUSH l1 x y"); r.Error != "" {
		t.Fatalf("LPUSH: %s", r.Error)
	}
	if r := run("RENAME a a2"); r.Error != "" {
		t.Fatalf("RENAME: %s", r.Error)
	}
	if !mr.Exists("a2") {
		t.Fatal("RENAME 未生效")
	}
}

// 危险命令二段确认:FLUSHALL/KEYS 未确认不执行,确认后放行。
func TestExecRedisDangerConfirm(t *testing.T) {
	mr, err := miniredis.Run()
	if err != nil {
		t.Fatal(err)
	}
	defer mr.Close()
	base := redisTestReq(t, mr)
	ctx := context.Background()
	_ = mr.Set("k", "v")

	run := func(sql string, confirm bool) ExecResult {
		c := base
		c.SQL = sql
		c.ConfirmDanger = confirm
		return execRedis(ctx, c)
	}

	if r := run("FLUSHALL", false); !r.NeedConfirm {
		t.Fatalf("FLUSHALL 未确认应 needConfirm: %#v", r)
	}
	if !mr.Exists("k") {
		t.Fatal("未确认的 FLUSHALL 不应执行")
	}
	if r := run("KEYS *", false); !r.NeedConfirm {
		t.Fatalf("KEYS 未确认应 needConfirm: %#v", r)
	}
	if r := run("KEYS *", true); r.Error != "" || r.RowCount != 1 {
		t.Fatalf("KEYS 确认后应放行: %#v", r)
	}
	if r := run("FLUSHDB", true); r.Error != "" {
		t.Fatalf("FLUSHDB 确认后应放行: %#v", r)
	}
	if mr.Exists("k") {
		t.Fatal("确认后的 FLUSHDB 应已执行")
	}
}

func TestExecRedisGenericDoAndScan(t *testing.T) {
	mr, err := miniredis.Run()
	if err != nil {
		t.Fatal(err)
	}
	defer mr.Close()
	base := redisTestReq(t, mr)
	ctx := context.Background()
	_ = mr.Set("user:1", "a")
	_ = mr.Set("user:2", "b")
	_ = mr.Set("other", "c")

	run := func(sql string) ExecResult {
		c := base
		c.SQL = sql
		return execRedis(ctx, c)
	}

	if r := run("EXISTS user:1"); r.Error != "" || r.RowCount != 1 {
		t.Fatalf("EXISTS via Do: %#v", r)
	}
	if r := run("SCAN 0 MATCH user:* COUNT 50"); r.Error != "" {
		t.Fatalf("SCAN MATCH: %#v", r)
	} else {
		if r.RowCount != 2 {
			t.Fatalf("SCAN MATCH want 2 got %d %#v", r.RowCount, r.Rows)
		}
	}
	if r := run("ECHO hello"); r.Error != "" {
		t.Fatalf("ECHO via Do: %#v", r)
	}
}