func (c *multiBankClient) DumpState(ctx context.Context) (interface{}, error) {
    txn, err := c.db.Begin()

    if err != nil {
        return nil, err
    }
    defer txn.Rollback()

    balances := make([]int64, 0, c.accountNum)
    for i := 0; i < c.accountNum; i++ {
        var balance int64
        sql := fmt.Sprintf("select balance from accounts_%d", i)
        if err = txn.QueryRowContext(ctx, sql).Scan(&balance); err != nil {
            return nil, err
        }
        balances = append(balances, balance)
    }
    return balances, nil
}

func (c *multiBankClient) Invoke(ctx context.Context, node string, r interface{}) interface{} {
    arg := r.(bankRequest)
    if arg.Op == 0 {
        return c.invokeRead(ctx, arg)
    }

    txn, err := c.db.Begin()

    if err != nil {
        return bankResponse{Ok: false}
    }
    defer txn.Rollback()

    var (
        fromBalance int64
        toBalance   int64
        tso         uint64
    )

    if err = txn.QueryRow("select @@tidb_current_ts").Scan(&tso); err != nil {
        return bankResponse{Ok: false}
    }

    if err = txn.QueryRowContext(ctx, fmt.Sprintf("select balance from accounts_%d where id = ? for update", arg.From), arg.From).Scan(&fromBalance); err != nil {
        return bankResponse{Ok: false}
    }

    if err = txn.QueryRowContext(ctx, fmt.Sprintf("select balance from accounts_%d where id = ? for update", arg.To), arg.To).Scan(&toBalance); err != nil {
        return bankResponse{Ok: false}
    }

    if fromBalance < arg.Amount {
        return bankResponse{Ok: false}
    }

    if _, err = txn.ExecContext(ctx, fmt.Sprintf("update accounts_%d set balance = balance - ? where id = ?", arg.From), arg.Amount, arg.From); err != nil {
        return bankResponse{Ok: false}
    }

    if _, err = txn.ExecContext(ctx, fmt.Sprintf("update accounts_%d set balance = balance + ? where id = ?", arg.To), arg.Amount, arg.To); err != nil {
        return bankResponse{Ok: false}
    }

    if err = txn.Commit(); err != nil {
        return bankResponse{Unknown: true, Tso: tso, FromBalance: fromBalance, ToBalance: toBalance}
    }

    return bankResponse{Ok: true, Tso: tso, FromBalance: fromBalance, ToBalance: toBalance}
}

func (c *multiBankClient) SetUp(ctx context.Context, nodes []string, node string) error {
    c.r = rand.New(rand.NewSource(time.Now().UnixNano()))
    db, err := sql.Open("mysql", fmt.Sprintf("root@tcp(%s:4000)/test", node))
    if err != nil {
        return err
    }
    c.db = db

    db.SetMaxIdleConns(1 + c.accountNum)

    // Do SetUp in the first node
    if node != nodes[0] {
        return nil
    }

    log.Printf("begin to create table accounts on node %s", node)

    for i := 0; i < c.accountNum; i++ {
        sql := fmt.Sprintf(`create table if not exists accounts_%d
            (id     int not null primary key,
            balance bigint not null)`, i)

        if _, err = db.ExecContext(ctx, sql); err != nil {
            return err
        }

        sql = fmt.Sprintf("insert into accounts_%d values (?, ?)", i)
        if _, err = db.ExecContext(ctx, sql, i, initBalance); err != nil {
            return err
        }

    }

    return nil
}

func (c *multiBankClient) TearDown(ctx context.Context, nodes []string, node string) error {
    return c.db.Close()
}