package main

import (
	"context"
	"database/sql"
	"fmt"
	"testing"
	"time"

	"github.com/pingcap/errors"
	. "github.com/pingcap/tidb-test/large_txn"
)

type LargeTxn struct{}

func (LargeTxn) LargeTxn(conn *sql.Conn) error {
	tx, err := conn.BeginTx(context.Background(), nil)
	if err != nil {
		return errors.WithStack(err)
	}

	// Make sure this is the primary key.
	_, err = tx.ExecContext(context.Background(), "select * from sbtest1 where id = 1 for update")
	if err != nil {
		return errors.WithStack(err)
	}

	_, err = tx.ExecContext(context.Background(), "update sbtest1 set k = k + 1")
	if err != nil {
		return errors.WithStack(err)
	}

	start := time.Now()
	// During the 2PC commit phase, the minCommitTS may be pushed.
	// Make sure the write operation could success eventually.
	if err := tx.Commit(); err != nil {
		return errors.WithStack(err)
	}
	dur := time.Since(start)
	fmt.Println("commit takes === ", dur)

	return nil
}

type SmallTxn struct{}

func (SmallTxn) SmallTxn(worker *Worker, conn *sql.Conn) error {
	rows, err := conn.QueryContext(context.Background(), "select * from sbtest1 limit 20")
	if err != nil {
		return errors.WithStack(err)
	}
	defer rows.Close()
	return nil
}

func TestT(t *testing.T) {
	conf := DefaultConfig()
	conf.Prepare.Task = BasePrepareTask{
		InsertCount:    1000,
		RowsEachInsert: 50,
	}
	conf.Run = RunConfig{
		Large: LargeTxnConfig{
			TxnMode: Pessimistic,
			Task:    LargeTxn{},
		},
		Small: SmallTxnConfig{
			TxnMode: Optimistic,
			// Change this value and see the commit time.
			WorkerCount: 256,
			Task:        SmallTxn{},
		},
	}
	RunTest(&conf)
}