TsDB2/migrate/abort_test.go

package migrate

import (
	"git.ycl.pub/yl32940/TsDB2/DB/core/page"
	"git.ycl.pub/yl32940/TsDB2/DB/core/types"
	"database/sql"
	"errors"
	"path/filepath"
	"testing"
)

// failSource 一个"扫描到一半就报错"的源。
//
// 前 failAt 行正常吐出来, 第 failAt+1 行直接返回错误 —— 用来把"扫描失败时目标端
// 到底提交了什么"钉死。failAt < 0 表示永不失败 (当对照组: 正常路径必须照常提交)。
type failSource struct {
	table  string
	sc     *TableSchema
	rows   int
	failAt int
}

func (s *failSource) Kind() Kind       { return KindSQLite } // 借 SQL 的字面量即可, 只用于日志/报告
func (s *failSource) Identity() string { return "fail-source" }
func (s *failSource) Close() error     { return nil }

func (s *failSource) Tables() ([]string, error) { return []string{s.table}, nil }

func (s *failSource) Schema(string) (*TableSchema, error) { return s.sc.Clone(), nil }

func (s *failSource) Scan(_ string, _ *TableSchema, fn func(Row) error) error {
	for i := 0; i < s.rows; i++ {
		if i == s.failAt {
			return errors.New("模拟扫描中途失败")
		}
		if err := fn(Row{
			PK:     uint64(i + 1),
			HasPK:  true,
			Values: []any{int64(i + 1), int64(i * 10)},
		}); err != nil {
			return err
		}
	}
	return nil
}

// TestScanFailureDoesNotCommitPartialRows 扫描失败时, 缓冲区里那批行必须被丢弃。
//
// 修复前 runTable 在出错路径上调的是 writer.Close() —— 而 SQL 写端的 Close 是
// 【成功语义】: flush 剩余缓冲 + Commit。于是攒在 RowWriter 缓冲里、还没写出去的行被
// 当成成功提交上去, 调用方却拿到一个 error、以为这张表没进去 —— 库里躺了半张表,
// 而且这份"半截数据"下一次续迁还会撞主键。
//
// 修复后出错路径走 Discarder.Discard (Rollback), 只有真正提交过的批次留在库里。
func TestScanFailureDoesNotCommitPartialRows(t *testing.T) {
	cases := []struct {
		name     string
		rows     int
		failAt   int
		wantRows int
	}{
		{"扫到第 3 行失败", 5, 2, 0},
		{"扫到第 5 行失败", 5, 4, 0},
		{"对照组: 正常迁完", 5, -1, 5},
	}

	for _, c := range cases {
		t.Run(c.name, func(t *testing.T) {
			dir := t.TempDir()
			dstPath := filepath.Join(dir, "dst.sqlite")

			dst, err := OpenSQLite(dstPath, quiet())
			if err != nil {
				t.Fatal(err)
			}

			// 只留两列整数: 避开字符串宽度档位这些与本次验证无关的东西。
			sc := &TableSchema{
				Name: "t1",
				Columns: []Column{
					{Name: "id", Type: types.FieldType_Int64},
					{Name: "n", Type: types.FieldType_Int64},
				},
				PKColumn:    "id",
				PKIndexType: page.IndexTypePrimary,
			}
			src := &failSource{table: "t1", sc: sc, rows: c.rows, failAt: c.failAt}

			// 批大小比行数大得多, 于是所有行都只会留在写入器的缓冲里 ——
			// 这正是"出错时 Close 会把它们提交上去"的那一批。
			o := quiet()
			_, runErr := Run(src, dst, o)
			if err := dst.Close(); err != nil {
				t.Fatal(err)
			}

			if c.failAt >= 0 && runErr == nil {
				t.Fatal("扫描失败必须把错误抛给调用方")
			}
			if c.failAt < 0 && runErr != nil {
				t.Fatalf("对照组不该失败: %v", runErr)
			}

			// 关键断言: 库里到底留了几行。
			db, err := sql.Open("sqlite", dstPath)
			if err != nil {
				t.Fatal(err)
			}
			defer db.Close()

			var n int
			if err := db.QueryRow(`SELECT COUNT(*) FROM t1`).Scan(&n); err != nil {
				t.Fatalf("查行数: %v", err)
			}
			if n != c.wantRows {
				t.Fatalf("目标表应当留下 %d 行, 实际 %d 行; Run 错误: %v", c.wantRows, n, runErr)
			}
		})
	}
}

// TestAbortWriterPrefersDiscard 盯住 runTable 出错路径的选择: 实现了 Discarder 就
// 走 Discard, 没实现才退回 Close。VRTX 写端没实现它 (它的 Close 不提交), 所以两条
// 路都必须留着 —— 少了哪一支, 要么 SQL 侧提交半截数据, 要么 VRTX 侧收尾漏掉。
func TestAbortWriterPrefersDiscard(t *testing.T) {
	d := &discardProbe{}
	abortWriter(d)
	if !d.discarded {
		t.Fatal("实现了 Discarder 的写入器必须走 Discard")
	}
	if d.closed {
		t.Fatal("走了 Discard 就不该再走 Close")
	}

	c := &closeOnlyProbe{}
	abortWriter(c)
	if !c.closed {
		t.Fatal("没实现 Discarder 的写入器必须退回 Close")
	}
}

type discardProbe struct {
	discarded bool
	closed    bool
}

func (p *discardProbe) Write([]Row) error { return nil }
func (p *discardProbe) Close() error      { p.closed = true; return nil }
func (p *discardProbe) Discard() error    { p.discarded = true; return nil }

type closeOnlyProbe struct {
	closed bool
}

func (p *closeOnlyProbe) Write([]Row) error { return nil }
func (p *closeOnlyProbe) Close() error      { p.closed = true; return nil }