TsDB2/demo/wal_rowlevel_test/main.go

package main

import (
	"TsDB2/DB/core"
	"TsDB2/DB/core/page"
	"TsDB2/DB/core/types"
	"encoding/binary"
	"fmt"
	"os"
)

const dbPath = "wal_rowlevel_test.db"

func main() {
	// 支持命令行参数控制测试阶段
	phase := "all"
	if len(os.Args) > 1 {
		phase = os.Args[1]
	}

	switch phase {
	case "setup":
		runSetup()
	case "verify":
		runVerify()
	case "txn_setup":
		runTxnSetup()
	case "txn_verify":
		runTxnVerify()
	case "rollback_test":
		runRollbackTest()
	default:
		runSetup()
		runVerify()
		fmt.Println()
		runTxnSetup()
		runTxnVerify()
		fmt.Println()
		runRollbackTest()
	}
}

// =============================================================================
// 测试 1: 普通 CRUD 通过 WAL 保存数据 (显式 Commit)
// =============================================================================
func runSetup() {
	fmt.Println("===== 测试 1: 普通 CRUD + WAL Commit =====")
	os.Remove(dbPath)

	db, err := core.Open(dbPath)
	if err != nil {
		fmt.Println("Open DB error:", err)
		return
	}
	// 禁用防抖自动提交, 仅靠手动 Commit
	db.SetAutoCommitDelay(0)

	table, err := db.CreateTable("Users", []types.FieldInfo{
		{FieldType: types.FieldType_Uint64, FieldName: []byte("ID")},
		{FieldType: types.FieldType_String, Flag: types.WidthCodeFor(32), FieldName: []byte("Name")},
		{FieldType: types.FieldType_Uint16, FieldName: []byte("Age")},
	})
	if err != nil {
		fmt.Println("CreateTable error:", err)
		return
	}

	_, err = table.CreateIndex("pk_id", page.IndexTypePrimary, []byte{0})
	if err != nil {
		fmt.Println("CreateIndex error:", err)
		return
	}

	// 插入 10 行
	fmt.Println("\n--- 插入 10 行 ---")
	for i := uint64(1); i <= 10; i++ {
		nameBytes := make([]byte, 32)
		copy(nameBytes, fmt.Sprintf("User_%d", i))
		_, err := table.InsertRow([]types.FieldValue{
			{FieldName: []byte("ID"), Value: i},
			{FieldName: []byte("Name"), Value: nameBytes},
			{FieldName: []byte("Age"), Value: int16(20 + i)},
		})
		if err != nil {
			fmt.Printf("Insert(%d) error: %v\n", i, err)
		}
	}
	fmt.Printf("插入完成, TotalRows=%d\n", table.TotalRows())

	// 更新 2 行
	fmt.Println("\n--- 更新 ID=3, ID=5 ---")
	updateName := make([]byte, 32)
	copy(updateName, "User_3_Updated")
	key3 := make([]byte, 8)
	binary.BigEndian.PutUint64(key3, 3)
	err = table.UpdateRow("pk_id", key3, []types.FieldValue{
		{FieldName: []byte("ID"), Value: uint64(3)},
		{FieldName: []byte("Name"), Value: updateName},
		{FieldName: []byte("Age"), Value: int16(99)},
	})
	if err != nil {
		fmt.Println("Update(3) error:", err)
	} else {
		fmt.Println("Update(3) OK")
	}

	updateName5 := make([]byte, 32)
	copy(updateName5, "User_5_Updated")
	key5 := make([]byte, 8)
	binary.BigEndian.PutUint64(key5, 5)
	err = table.UpdateRow("pk_id", key5, []types.FieldValue{
		{FieldName: []byte("ID"), Value: uint64(5)},
		{FieldName: []byte("Name"), Value: updateName5},
		{FieldName: []byte("Age"), Value: int16(88)},
	})
	if err != nil {
		fmt.Println("Update(5) error:", err)
	} else {
		fmt.Println("Update(5) OK")
	}

	// 删除 2 行
	fmt.Println("\n--- 删除 ID=7, ID=9 ---")
	key7 := make([]byte, 8)
	binary.BigEndian.PutUint64(key7, 7)
	if err := table.DeleteRow("pk_id", key7); err != nil {
		fmt.Println("Delete(7) error:", err)
	} else {
		fmt.Println("Delete(7) OK")
	}
	key9 := make([]byte, 8)
	binary.BigEndian.PutUint64(key9, 9)
	if err := table.DeleteRow("pk_id", key9); err != nil {
		fmt.Println("Delete(9) error:", err)
	} else {
		fmt.Println("Delete(9) OK")
	}

	fmt.Printf("\n操作后 TotalRows=%d\n", table.TotalRows())

	// 显式 Commit (走 WAL)
	fmt.Println("\n--- 显式 Commit (WAL) ---")
	if err := db.Commit(); err != nil {
		fmt.Println("Commit error:", err)
	} else {
		fmt.Println("Commit OK")
	}

	// 验证当前可查询
	fmt.Println("\n--- 验证查询 ---")
	verifyQuery(table, "ID=3", 3)
	verifyQuery(table, "ID=5", 5)
	verifyQueryNotExist(table, "ID=7", 7)
	verifyQueryNotExist(table, "ID=9", 9)

	db.Close()
	fmt.Println("\n数据库已关闭 (不模拟崩溃, 正常关闭)")
}

// runVerify 重新打开数据库验证数据持久化
func runVerify() {
	fmt.Println("\n===== 测试 1 验证: 重新打开数据库 =====")
	db, err := core.Open(dbPath)
	if err != nil {
		fmt.Println("Open DB error:", err)
		return
	}
	defer db.Close()

	table, ok := db.FindTable("Users")
	if !ok {
		fmt.Println("Table 'Users' not found!")
		return
	}

	fmt.Printf("TotalRows=%d (期望: 8)\n", table.TotalRows())

	// 验证所有行
	fmt.Println("\n--- 验证查询 ---")
	verifyQuery(table, "ID=1", 1)
	verifyQuery(table, "ID=3", 3)         // 已更新
	verifyQuery(table, "ID=5", 5)         // 已更新
	verifyQueryNotExist(table, "ID=7", 7) // 已删除
	verifyQueryNotExist(table, "ID=9", 9) // 已删除
	verifyQuery(table, "ID=10", 10)

	fmt.Println("\n✅ 测试 1 通过: 普通 CRUD + WAL Commit 数据持久化成功")
}

// =============================================================================
// 测试 2: 事务 CRUD + WAL
// =============================================================================
func runTxnSetup() {
	fmt.Println("===== 测试 2: 事务 CRUD + WAL =====")
	os.Remove(dbPath)

	db, err := core.Open(dbPath)
	if err != nil {
		fmt.Println("Open DB error:", err)
		return
	}
	db.SetAutoCommitDelay(0)

	table, err := db.CreateTable("TxnUsers", []types.FieldInfo{
		{FieldType: types.FieldType_Uint64, FieldName: []byte("ID")},
		{FieldType: types.FieldType_String, Flag: types.WidthCodeFor(32), FieldName: []byte("Name")},
		{FieldType: types.FieldType_Uint16, FieldName: []byte("Age")},
	})
	if err != nil {
		fmt.Println("CreateTable error:", err)
		return
	}
	_, err = table.CreateIndex("pk_id", page.IndexTypePrimary, []byte{0})
	if err != nil {
		fmt.Println("CreateIndex error:", err)
		return
	}

	// 先插入 5 行 (非事务)
	fmt.Println("\n--- 非事务插入 5 行 ---")
	for i := uint64(1); i <= 5; i++ {
		nameBytes := make([]byte, 32)
		copy(nameBytes, fmt.Sprintf("Base_%d", i))
		table.InsertRow([]types.FieldValue{
			{FieldName: []byte("ID"), Value: i},
			{FieldName: []byte("Name"), Value: nameBytes},
			{FieldName: []byte("Age"), Value: int16(20 + i)},
		})
	}
	db.Commit() // 提交基础数据
	fmt.Printf("基础数据 TotalRows=%d\n", table.TotalRows())

	// 开启事务
	fmt.Println("\n--- 开启事务 ---")
	txn := db.BeginTransaction()

	// 事务内插入 5 行
	fmt.Println("--- 事务内插入 5 行 (ID=6~10) ---")
	for i := uint64(6); i <= 10; i++ {
		nameBytes := make([]byte, 32)
		copy(nameBytes, fmt.Sprintf("Txn_%d", i))
		_, err := table.InsertRow([]types.FieldValue{
			{FieldName: []byte("ID"), Value: i},
			{FieldName: []byte("Name"), Value: nameBytes},
			{FieldName: []byte("Age"), Value: int16(30 + i)},
		})
		if err != nil {
			fmt.Printf("Txn Insert(%d) error: %v\n", i, err)
		}
	}

	// 事务内更新 1 行
	fmt.Println("--- 事务内更新 ID=2 ---")
	txnName := make([]byte, 32)
	copy(txnName, "Base_2_TxnUpdated")
	key2 := make([]byte, 8)
	binary.BigEndian.PutUint64(key2, 2)
	table.UpdateRow("pk_id", key2, []types.FieldValue{
		{FieldName: []byte("ID"), Value: uint64(2)},
		{FieldName: []byte("Name"), Value: txnName},
		{FieldName: []byte("Age"), Value: int16(77)},
	})

	// 事务内删除 1 行
	fmt.Println("--- 事务内删除 ID=4 ---")
	key4 := make([]byte, 8)
	binary.BigEndian.PutUint64(key4, 4)
	table.DeleteRow("pk_id", key4)

	fmt.Printf("\n事务内操作后 TotalRows=%d (期望: 9, 5+5-1)\n", table.TotalRows())

	// 提交事务
	fmt.Println("\n--- 事务 Commit ---")
	if err := txn.Commit(); err != nil {
		fmt.Println("Txn Commit error:", err)
		return
	}
	fmt.Println("事务 Commit OK")

	db.Close()
	fmt.Println("数据库已关闭")
}

func runTxnVerify() {
	fmt.Println("\n===== 测试 2 验证: 重新打开数据库 =====")
	db, err := core.Open(dbPath)
	if err != nil {
		fmt.Println("Open DB error:", err)
		return
	}
	defer db.Close()

	table, ok := db.FindTable("TxnUsers")
	if !ok {
		fmt.Println("Table 'TxnUsers' not found!")
		return
	}

	fmt.Printf("TotalRows=%d (期望: 9)\n", table.TotalRows())

	// 验证基础数据
	verifyQuery(table, "ID=1", 1)
	verifyQuery(table, "ID=2", 2)         // 事务内更新
	verifyQueryNotExist(table, "ID=4", 4) // 事务内删除

	// 验证事务内插入
	verifyQuery(table, "ID=6", 6)
	verifyQuery(table, "ID=10", 10)

	fmt.Println("\n✅ 测试 2 通过: 事务 CRUD + WAL 数据持久化成功")
}

// =============================================================================
// 测试 3: 事务 Rollback
// =============================================================================
func runRollbackTest() {
	fmt.Println("\n===== 测试 3: 事务 Rollback =====")
	os.Remove(dbPath)

	db, err := core.Open(dbPath)
	if err != nil {
		fmt.Println("Open DB error:", err)
		return
	}
	db.SetAutoCommitDelay(0)

	table, err := db.CreateTable("RollbackTest", []types.FieldInfo{
		{FieldType: types.FieldType_Uint64, FieldName: []byte("ID")},
		{FieldType: types.FieldType_String, Flag: types.WidthCodeFor(32), FieldName: []byte("Name")},
	})
	if err != nil {
		fmt.Println("CreateTable error:", err)
		return
	}
	_, err = table.CreateIndex("pk_id", page.IndexTypePrimary, []byte{0})
	if err != nil {
		fmt.Println("CreateIndex error:", err)
		return
	}

	// 插入基础数据
	fmt.Println("--- 插入 5 行基础数据 ---")
	for i := uint64(1); i <= 5; i++ {
		nameBytes := make([]byte, 32)
		copy(nameBytes, fmt.Sprintf("Base_%d", i))
		table.InsertRow([]types.FieldValue{
			{FieldName: []byte("ID"), Value: i},
			{FieldName: []byte("Name"), Value: nameBytes},
		})
	}
	db.Commit()
	fmt.Printf("基础数据 TotalRows=%d\n", table.TotalRows())

	// 开启事务
	fmt.Println("\n--- 开启事务, 插入 5 行 (ID=6~10) ---")
	txn := db.BeginTransaction()
	for i := uint64(6); i <= 10; i++ {
		nameBytes := make([]byte, 32)
		copy(nameBytes, fmt.Sprintf("Txn_%d", i))
		table.InsertRow([]types.FieldValue{
			{FieldName: []byte("ID"), Value: i},
			{FieldName: []byte("Name"), Value: nameBytes},
		})
	}
	fmt.Printf("事务内 TotalRows=%d (期望: 10)\n", table.TotalRows())

	// 回滚
	fmt.Println("\n--- 事务 Rollback ---")
	if err := txn.Rollback(); err != nil {
		fmt.Println("Rollback error:", err)
		return
	}
	fmt.Println("Rollback OK")
	fmt.Printf("回滚后 TotalRows=%d (期望: 5)\n", table.TotalRows())

	// 验证回滚效果
	verifyQuery(table, "ID=1", 1)
	verifyQueryNotExist(table, "ID=6", 6)
	verifyQueryNotExist(table, "ID=10", 10)

	db.Close()

	// 重新打开验证持久化
	fmt.Println("\n--- 重新打开数据库验证 ---")
	db2, err := core.Open(dbPath)
	if err != nil {
		fmt.Println("Open DB error:", err)
		return
	}
	defer db2.Close()

	table2, ok := db2.FindTable("RollbackTest")
	if !ok {
		fmt.Println("Table not found!")
		return
	}
	fmt.Printf("重开后 TotalRows=%d (期望: 5)\n", table2.TotalRows())
	verifyQuery(table2, "ID=1", 1)
	verifyQueryNotExist(table2, "ID=6", 6)

	fmt.Println("\n✅ 测试 3 通过: 事务 Rollback 正确")
}

// =============================================================================
// 辅助函数
// =============================================================================
func verifyQuery(table *core.Table, label string, id uint64) {
	key := make([]byte, 8)
	binary.BigEndian.PutUint64(key, id)
	row, err := table.SearchByIndex("pk_id", key)
	if err != nil {
		fmt.Printf("  ❌ %s: 查询失败: %v\n", label, err)
		return
	}
	fmt.Printf("  ✅ %s: 找到 -> ", label)
	for _, fv := range row {
		fmt.Printf("%s=%v ", string(fv.FieldName), formatValue(fv.Value))
	}
	fmt.Println()
}

func verifyQueryNotExist(table *core.Table, label string, id uint64) {
	key := make([]byte, 8)
	binary.BigEndian.PutUint64(key, id)
	_, err := table.SearchByIndex("pk_id", key)
	if err != nil {
		fmt.Printf("  ✅ %s: 正确不存在\n", label)
		return
	}
	fmt.Printf("  ❌ %s: 应该不存在但找到了!\n", label)
}

func formatValue(v interface{}) string {
	switch val := v.(type) {
	case []byte:
		// 去除尾部零字节
		end := len(val)
		for end > 0 && val[end-1] == 0 {
			end--
		}
		return string(val[:end])
	default:
		return fmt.Sprintf("%v", val)
	}
}