opsxdb/apps/server/internal/db/export_test.go

package db

import (
	"bytes"
	"context"
	"encoding/csv"
	"encoding/json"
	"strings"
	"testing"
)

func newExportConn(t *testing.T) ConnReq {
	return newSQLiteFile(t,
		`CREATE TABLE t_export (id INTEGER PRIMARY KEY, name TEXT, note TEXT)`,
		`INSERT INTO t_export (id, name, note) VALUES (1, 'alice', 'hi, "q"')`,
		`INSERT INTO t_export (id, name, note) VALUES (2, 'bo''b', NULL)`,
		`INSERT INTO t_export (id, name, note) VALUES (3, 'carol', 'line1`+"\n"+`line2')`,
	)
}

// R2-B7:csv 输出形状(表头 + csv 引号转义由 encoding/csv 负责)。
func TestExportCSV(t *testing.T) {
	c := newExportConn(t)
	c.SQL = `SELECT id, name, note FROM t_export ORDER BY id`
	c.Format = "csv"
	var buf bytes.Buffer
	if err := ExportQuery(context.Background(), c, &buf); err != nil {
		t.Fatal(err)
	}
	out := buf.String()
	if !strings.HasPrefix(out, "id,name,note\n") {
		t.Fatalf("csv 表头错误: %q", out)
	}
	// 含逗号/引号的字段整体加引号、内部引号 doubling
	if !strings.Contains(out, `"hi, ""q"""`) {
		t.Fatalf("csv 转义错误: %q", out)
	}
	// 单引号 doubling 是 SQL 字面量层;csv 层原样输出 'bo''b' 存好的 bo'b
	if !strings.Contains(out, "2,bo'b,") {
		t.Fatalf("csv 内容错误: %q", out)
	}
	// 多行字面量会使单条记录跨物理行,用 csv.Reader 按记录解析
	recs, err := csv.NewReader(strings.NewReader(out)).ReadAll()
	if err != nil {
		t.Fatalf("csv 输出应可解析: %v", err)
	}
	if len(recs) != 4 { // 表头 + 3 行
		t.Fatalf("应为表头+3 行: %q", out)
	}
}

func TestExportJSON(t *testing.T) {
	c := newExportConn(t)
	c.SQL = `SELECT id, name, note FROM t_export ORDER BY id`
	c.Format = "json"
	var buf bytes.Buffer
	if err := ExportQuery(context.Background(), c, &buf); err != nil {
		t.Fatal(err)
	}
	var rows []map[string]interface{}
	if err := json.Unmarshal(buf.Bytes(), &rows); err != nil {
		t.Fatalf("json 输出应可解析: %v\n%s", err, buf.String())
	}
	if len(rows) != 3 {
		t.Fatalf("应为 3 行: %v", rows)
	}
	if rows[0]["name"] != "alice" || rows[1]["note"] != nil {
		t.Fatalf("行值错误: %v", rows)
	}
}

func TestExportInsert(t *testing.T) {
	c := newExportConn(t)
	c.SQL = `SELECT id, name, note FROM t_export ORDER BY id`
	c.Format = "insert"
	c.Table = "t_export"
	var buf bytes.Buffer
	if err := ExportQuery(context.Background(), c, &buf); err != nil {
		t.Fatal(err)
	}
	out := buf.String()
	if !strings.Contains(out, `INSERT INTO "t_export" ("id", "name", "note") VALUES (1, 'alice', 'hi, "q"');`) {
		t.Fatalf("insert 首行错误: %q", out)
	}
	// 字符串单引号 doubling;NULL 原样
	if !strings.Contains(out, `VALUES (2, 'bo''b', NULL);`) {
		t.Fatalf("insert 字面量错误: %q", out)
	}
	// 多行字面量内换行原样保留
	if !strings.Contains(out, "line1\nline2") {
		t.Fatalf("insert 换行字面量错误: %q", out)
	}
}

// limit 生效(含默认上限收敛)。
func TestExportLimit(t *testing.T) {
	c := newExportConn(t)
	c.SQL = `SELECT id FROM t_export ORDER BY id`
	c.Format = "csv"
	c.Limit = 2
	var buf bytes.Buffer
	if err := ExportQuery(context.Background(), c, &buf); err != nil {
		t.Fatal(err)
	}
	rows := strings.Split(strings.TrimSpace(buf.String()), "\n")
	if len(rows) != 3 { // 表头 + 2 行
		t.Fatalf("limit=2 应只导 2 行: %q", buf.String())
	}
}

func TestExportAllowedQuery(t *testing.T) {
	for _, ok := range []string{
		"SELECT * FROM t", "WITH x AS (SELECT 1) SELECT * FROM x",
		"SHOW TABLES", "DESCRIBE t", "EXPLAIN SELECT 1",
		"-- c\nSELECT 1",
	} {
		if !ExportAllowedQuery(ok) {
			t.Fatalf("应允许只读: %q", ok)
		}
	}
	for _, bad := range []string{
		"DELETE FROM t", "INSERT INTO t VALUES (1)", "DROP TABLE t",
		"UPDATE t SET a=1", "", "  ",
	} {
		if ExportAllowedQuery(bad) {
			t.Fatalf("不应允许: %q", bad)
		}
	}
}

func TestExportBadFormat(t *testing.T) {
	c := newExportConn(t)
	c.SQL = "SELECT 1"
	c.Format = "xml"
	var buf bytes.Buffer
	if err := ExportQuery(context.Background(), c, &buf); err == nil {
		t.Fatal("未知格式应报错")
	}
}