test: add test case

af/conn.go

@@ -58,8 +58,7 @@ func (conn *Connector) Close() error {
 func (conn *Connector) StmtExecute(sql string, params *param.Param) (res driver.Result, err error) {
 	stmt := NewStmt(conn.taos)
 	if stmt == nil {
-		err = &errors.TaosError{Code: 0xffff, ErrStr: "failed to init stmt"}
-		return
+		return nil, &errors.TaosError{Code: 0xffff, ErrStr: "failed to init stmt"}
 	}
 
 	defer stmt.Close()
@@ -229,6 +228,11 @@ func (conn *Connector) InsertStmt() *insertstmt.InsertStmt {
 	return insertstmt.NewInsertStmt(conn.taos)
 }
 
+// Stmt Prepare stmt
+func (conn *Connector) Stmt() *Stmt {
+	return NewStmt(conn.taos)
+}
+
 // InsertStmtWithReqID Prepare batch insert stmt with reqID
 func (conn *Connector) InsertStmtWithReqID(reqID int64) *insertstmt.InsertStmt {
 	return insertstmt.NewInsertStmtWithReqID(conn.taos, reqID)
@@ -240,8 +244,7 @@ func (conn *Connector) SelectDB(db string) error {
 	code := wrapper.TaosSelectDB(conn.taos, db)
 	locker.Unlock()
 	if code != 0 {
-		err := taosError.NewError(code, wrapper.TaosErrorStr(nil))
-		return err
+		return taosError.NewError(code, wrapper.TaosErrorStr(nil))
 	}
 	return nil
 }

af/conn_test.go

@@ -7,8 +7,10 @@ import (
 	"testing"
 	"time"
 
+	"github.com/stretchr/testify/assert"
 	"github.com/taosdata/driver-go/v3/common"
 	param2 "github.com/taosdata/driver-go/v3/common/param"
+	"github.com/taosdata/driver-go/v3/wrapper"
 )
 
 func TestMain(m *testing.M) {
@@ -925,3 +927,49 @@ func TestConnector_QueryWithReqID(t *testing.T) {
 		t.Fatal("result is error")
 	}
 }
+
+func TestNewConnector(t *testing.T) {
+	tc, err := wrapper.TaosConnect("", "root", "taosdata", "", 0)
+	assert.NoError(t, err)
+	conn, err := NewConnector(tc)
+	assert.NoError(t, err)
+	assert.Equal(t, tc, conn.taos)
+	_, err = NewConnector(nil)
+	assert.Error(t, err)
+}
+
+func TestSelectDB(t *testing.T) {
+	db, err := Open("", "", "", "", 0)
+	assert.NoError(t, err)
+	_, err = db.Exec("create database if not exists test_af precision 'us'  keep 36500")
+	assert.NoError(t, err)
+	rows, err := db.Query("select database()")
+	assert.NoError(t, err)
+	dest := make([]driver.Value, 1)
+	err = rows.Next(dest)
+	assert.NoError(t, err)
+	err = rows.Close()
+	assert.NoError(t, err)
+	assert.Nil(t, dest[0])
+	err = db.SelectDB("test_af")
+	assert.NoError(t, err)
+	rows, err = db.Query("select database()")
+	assert.NoError(t, err)
+	dest = make([]driver.Value, 1)
+	err = rows.Next(dest)
+	assert.NoError(t, err)
+	err = rows.Close()
+	assert.NoError(t, err)
+	assert.Equal(t, "test_af", dest[0])
+	err = db.Close()
+	assert.NoError(t, err)
+}
+
+func TestGetTableVGroupID(t *testing.T) {
+	db := testDatabase(t)
+	_, err := db.Exec("create table test_vg (ts timestamp,v int)")
+	assert.NoError(t, err)
+	vgID, err := db.GetTableVGroupID("test_af", "test_vg")
+	assert.NoError(t, err)
+	t.Log(vgID)
+}

af/insertstmt/stmt.go

@@ -34,13 +34,11 @@ func (stmt *InsertStmt) Prepare(sql string) error {
 	code := wrapper.TaosStmtPrepare(stmt.stmt, sql)
 	locker.Unlock()
 	if code != 0 {
-		errStr := wrapper.TaosStmtErrStr(stmt.stmt)
-		return taosError.NewError(code, errStr)
+		return stmt.stmtErr(code)
 	}
 	isInsert, code := wrapper.TaosStmtIsInsert(stmt.stmt)
 	if code != 0 {
-		errStr := wrapper.TaosStmtErrStr(stmt.stmt)
-		return taosError.NewError(code, errStr)
+		return stmt.stmtErr(code)
 	}
 	if !isInsert {
 		return errors.New("only support insert")
@@ -53,8 +51,7 @@ func (stmt *InsertStmt) SetTableName(name string) error {
 	code := wrapper.TaosStmtSetTBName(stmt.stmt, name)
 	locker.Unlock()
 	if code != 0 {
-		errStr := wrapper.TaosStmtErrStr(stmt.stmt)
-		return taosError.NewError(code, errStr)
+		return stmt.stmtErr(code)
 	}
 	return nil
 }
@@ -64,8 +61,7 @@ func (stmt *InsertStmt) SetSubTableName(name string) error {
 	code := wrapper.TaosStmtSetSubTBName(stmt.stmt, name)
 	locker.Unlock()
 	if code != 0 {
-		errStr := wrapper.TaosStmtErrStr(stmt.stmt)
-		return taosError.NewError(code, errStr)
+		return stmt.stmtErr(code)
 	}
 	return nil
 }
@@ -75,8 +71,7 @@ func (stmt *InsertStmt) SetTableNameWithTags(tableName string, tags *param.Param
 	code := wrapper.TaosStmtSetTBNameTags(stmt.stmt, tableName, tags.GetValues())
 	locker.Unlock()
 	if code != 0 {
-		errStr := wrapper.TaosStmtErrStr(stmt.stmt)
-		return taosError.NewError(code, errStr)
+		return stmt.stmtErr(code)
 	}
 	return nil
 }
@@ -95,8 +90,7 @@ func (stmt *InsertStmt) BindParam(params []*param.Param, bindType *param.ColumnT
 	code := wrapper.TaosStmtBindParamBatch(stmt.stmt, data, columnTypes)
 	locker.Unlock()
 	if code != 0 {
-		errStr := wrapper.TaosStmtErrStr(stmt.stmt)
-		return taosError.NewError(code, errStr)
+		return stmt.stmtErr(code)
 	}
 	return nil
 }
@@ -106,8 +100,7 @@ func (stmt *InsertStmt) AddBatch() error {
 	code := wrapper.TaosStmtAddBatch(stmt.stmt)
 	locker.Unlock()
 	if code != 0 {
-		errStr := wrapper.TaosStmtErrStr(stmt.stmt)
-		return taosError.NewError(code, errStr)
+		return stmt.stmtErr(code)
 	}
 	return nil
 }
@@ -117,8 +110,7 @@ func (stmt *InsertStmt) Execute() error {
 	code := wrapper.TaosStmtExecute(stmt.stmt)
 	locker.Unlock()
 	if code != 0 {
-		errStr := wrapper.TaosStmtErrStr(stmt.stmt)
-		return taosError.NewError(code, errStr)
+		return stmt.stmtErr(code)
 	}
 	return nil
 }
@@ -133,9 +125,13 @@ func (stmt *InsertStmt) Close() error {
 	locker.Unlock()
 	var err error
 	if code != 0 {
-		errStr := wrapper.TaosStmtErrStr(stmt.stmt)
-		err = taosError.NewError(code, errStr)
+		err = stmt.stmtErr(code)
 	}
 	stmt.stmt = nil
 	return err
 }
+
+func (stmt *InsertStmt) stmtErr(code int) error {
+	errStr := wrapper.TaosStmtErrStr(stmt.stmt)
+	return taosError.NewError(code, errStr)
+}

af/insertstmt/stmt_test.go

@@ -0,0 +1,113 @@
+package insertstmt
+
+import (
+	"testing"
+	"time"
+	"unsafe"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/taosdata/driver-go/v3/common"
+	"github.com/taosdata/driver-go/v3/common/param"
+	taosError "github.com/taosdata/driver-go/v3/errors"
+	"github.com/taosdata/driver-go/v3/wrapper"
+)
+
+func prepareEnv(conn unsafe.Pointer) error {
+	sqls := []string{
+		"create database if not exists insert_stmt",
+		"use insert_stmt",
+		"create table test (ts timestamp, a int, b float)",
+		"create table stb(ts timestamp, v int) tags(a binary(10))",
+	}
+	for i := 0; i < len(sqls); i++ {
+		if err := exec(conn, sqls[i]); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+func cleanEnv(conn unsafe.Pointer) error {
+	sqls := []string{
+		"drop database if exists insert_stmt",
+	}
+	for i := 0; i < len(sqls); i++ {
+		if err := exec(conn, sqls[i]); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+func TestStmt(t *testing.T) {
+	conn, err := wrapper.TaosConnect("", "root", "taosdata", "", 0)
+	assert.NoError(t, err)
+	defer wrapper.TaosClose(conn)
+	err = prepareEnv(conn)
+	assert.NoError(t, err)
+	defer cleanEnv(conn)
+	s := NewInsertStmt(conn)
+	defer s.Close()
+	err = s.Prepare("insert into ? values(?,?,?)")
+	assert.NoError(t, err)
+	err = s.SetTableName("test")
+	assert.NoError(t, err)
+	params := []*param.Param{
+		param.NewParam(1).AddTimestamp(time.Now(), common.PrecisionMilliSecond),
+		param.NewParam(1).AddInt(1),
+		param.NewParam(1).AddFloat(1.1),
+	}
+	err = s.BindParam(params, param.NewColumnType(3).AddTimestamp().AddInt().AddFloat())
+	assert.NoError(t, err)
+	err = s.AddBatch()
+	assert.NoError(t, err)
+	err = s.Execute()
+	assert.NoError(t, err)
+	affected := s.GetAffectedRows()
+	assert.Equal(t, int(1), affected)
+
+	err = s.Prepare("insert into ? using stb tags(?) values(?,?)")
+	assert.NoError(t, err)
+	err = s.SetTableNameWithTags("ctb1", param.NewParam(1).AddBinary([]byte("test")))
+	assert.NoError(t, err)
+	params = []*param.Param{
+		param.NewParam(1).AddTimestamp(time.Now(), common.PrecisionMilliSecond),
+		param.NewParam(1).AddInt(1),
+	}
+	err = s.BindParam(params, param.NewColumnType(2).AddTimestamp().AddInt())
+	assert.NoError(t, err)
+	err = s.AddBatch()
+	assert.NoError(t, err)
+	err = s.Execute()
+	assert.NoError(t, err)
+	affected = s.GetAffectedRows()
+	assert.Equal(t, int(1), affected)
+
+	err = s.Prepare("insert into ? using stb tags('ctb2') values(?,?)")
+	assert.NoError(t, err)
+	err = s.SetSubTableName("ctb2")
+	assert.NoError(t, err)
+	params = []*param.Param{
+		param.NewParam(1).AddTimestamp(time.Now(), common.PrecisionMilliSecond),
+		param.NewParam(1).AddInt(1),
+	}
+	err = s.BindParam(params, param.NewColumnType(2).AddTimestamp().AddInt())
+	assert.NoError(t, err)
+	err = s.AddBatch()
+	assert.NoError(t, err)
+	err = s.Execute()
+	assert.NoError(t, err)
+	affected = s.GetAffectedRows()
+	assert.Equal(t, int(1), affected)
+
+}
+
+func exec(conn unsafe.Pointer, sql string) error {
+	res := wrapper.TaosQuery(conn, sql)
+	defer func() {
+		wrapper.TaosFreeResult(res)
+	}()
+	if code := wrapper.TaosError(res); code != 0 {
+		errStr := wrapper.TaosErrorStr(res)
+		return taosError.NewError(code, errStr)
+	}
+	return nil
+}

af/rows.go

@@ -23,6 +23,7 @@ type rows struct {
 	blockSize   int
 	result      unsafe.Pointer
 	precision   int
+	isStmt      bool
 }
 
 func (rs *rows) Columns() []string {
@@ -55,8 +56,7 @@ func (rs *rows) Next(dest []driver.Value) error {
 	if rs.result == nil {
 		return &errors.TaosError{Code: 0xffff, ErrStr: "result is nil!"}
 	}
-
-	if rs.block == nil {
+	if rs.block == nil || rs.blockOffset >= rs.blockSize {
 		if err := rs.taosFetchBlock(); err != nil {
 			return err
 		}
@@ -67,16 +67,6 @@ func (rs *rows) Next(dest []driver.Value) error {
 		return io.EOF
 	}
 
-	if rs.blockOffset >= rs.blockSize {
-		if err := rs.taosFetchBlock(); err != nil {
-			return err
-		}
-	}
-	if rs.blockSize == 0 {
-		rs.block = nil
-		rs.freeResult()
-		return io.EOF
-	}
 	parser.ReadRow(dest, rs.block, rs.blockSize, rs.blockOffset, rs.rowsHeader.ColTypes, rs.precision)
 	rs.blockOffset++
 	return nil
@@ -111,9 +101,11 @@ func (rs *rows) asyncFetchRows() *handler.AsyncResult {
 
 func (rs *rows) freeResult() {
 	if rs.result != nil {
-		locker.Lock()
-		wrapper.TaosFreeResult(rs.result)
-		locker.Unlock()
+		if !rs.isStmt {
+			locker.Lock()
+			wrapper.TaosFreeResult(rs.result)
+			locker.Unlock()
+		}
 		rs.result = nil
 	}