implement Rows.HasNextResultSet and Rows.NextResultSet

Fixes alexbrainman#81
Fixes alexbrainman#127

Author: alexbrainman

api/api.go

@@ -58,6 +58,7 @@ type (
 //sys	SQLGetData(statementHandle SQLHSTMT, colOrParamNum SQLUSMALLINT, targetType SQLSMALLINT, targetValuePtr SQLPOINTER, bufferLength SQLLEN, vallen *SQLLEN) (ret SQLRETURN) = odbc32.SQLGetData
 //sys	SQLGetDiagRec(handleType SQLSMALLINT, handle SQLHANDLE, recNumber SQLSMALLINT, sqlState *SQLWCHAR, nativeErrorPtr *SQLINTEGER, messageText *SQLWCHAR, bufferLength SQLSMALLINT, textLengthPtr *SQLSMALLINT) (ret SQLRETURN) = odbc32.SQLGetDiagRecW
 //sys	SQLNumParams(statementHandle SQLHSTMT, parameterCountPtr *SQLSMALLINT) (ret SQLRETURN) = odbc32.SQLNumParams
+//sys	SQLMoreResults(statementHandle SQLHSTMT) (ret SQLRETURN) = odbc32.SQLMoreResults
 //sys	SQLNumResultCols(statementHandle SQLHSTMT, columnCountPtr *SQLSMALLINT)  (ret SQLRETURN) = odbc32.SQLNumResultCols
 //sys	SQLPrepare(statementHandle SQLHSTMT, statementText *SQLWCHAR, textLength SQLINTEGER) (ret SQLRETURN) = odbc32.SQLPrepareW
 //sys	SQLRowCount(statementHandle SQLHSTMT, rowCountPtr *SQLLEN) (ret SQLRETURN) = odbc32.SQLRowCount

api/zapi_unix.go

@@ -95,6 +95,11 @@ func SQLNumParams(statementHandle SQLHSTMT, parameterCountPtr *SQLSMALLINT) (ret
 	return SQLRETURN(r)
 }
 
+func SQLMoreResults(statementHandle SQLHSTMT) (ret SQLRETURN) {
+	r := C.SQLMoreResults(C.SQLHSTMT(statementHandle))
+	return SQLRETURN(r)
+}
+
 func SQLNumResultCols(statementHandle SQLHSTMT, columnCountPtr *SQLSMALLINT) (ret SQLRETURN) {
 	r := C.SQLNumResultCols(C.SQLHSTMT(statementHandle), (*C.SQLSMALLINT)(columnCountPtr))
 	return SQLRETURN(r)

api/zapi_windows.go

@@ -54,6 +54,7 @@ var (
 	procSQLGetData         = mododbc32.NewProc("SQLGetData")
 	procSQLGetDiagRecW     = mododbc32.NewProc("SQLGetDiagRecW")
 	procSQLNumParams       = mododbc32.NewProc("SQLNumParams")
+	procSQLMoreResults     = mododbc32.NewProc("SQLMoreResults")
 	procSQLNumResultCols   = mododbc32.NewProc("SQLNumResultCols")
 	procSQLPrepareW        = mododbc32.NewProc("SQLPrepareW")
 	procSQLRowCount        = mododbc32.NewProc("SQLRowCount")
@@ -151,6 +152,12 @@ func SQLNumParams(statementHandle SQLHSTMT, parameterCountPtr *SQLSMALLINT) (ret
 	return
 }
 
+func SQLMoreResults(statementHandle SQLHSTMT) (ret SQLRETURN) {
+	r0, _, _ := syscall.Syscall(procSQLMoreResults.Addr(), 1, uintptr(statementHandle), 0, 0)
+	ret = SQLRETURN(r0)
+	return
+}
+
 func SQLNumResultCols(statementHandle SQLHSTMT, columnCountPtr *SQLSMALLINT) (ret SQLRETURN) {
 	r0, _, _ := syscall.Syscall(procSQLNumResultCols.Addr(), 2, uintptr(statementHandle), uintptr(unsafe.Pointer(columnCountPtr)), 0)
 	ret = SQLRETURN(r0)

mssql_test.go

@@ -1733,3 +1733,146 @@ func TestMSSQLMarkBeginBadConn(t *testing.T) {
 		testFn(next.label, next.fn)
 	}
 }
+
+func testMSSQLNextResultSet(t *testing.T, verifyBatch func(rows *sql.Rows)) {
+	db, sc, err := mssqlConnect()
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer closeDB(t, db, sc, sc)
+
+	db.Exec("drop table dbo.temp")
+	exec(t, db, `create table dbo.temp (name varchar(50))`)
+	exec(t, db, `insert into dbo.temp (name) values ('russ')`)
+	exec(t, db, `insert into dbo.temp (name) values ('brad')`)
+
+	rows, err := db.Query(`
+select name from dbo.temp where name = 'russ';
+select name from dbo.temp where name = 'brad';
+`)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer rows.Close()
+
+	verifyBatch(rows)
+
+	exec(t, db, "drop table dbo.temp")
+}
+
+func TestMSSQLNextResultSet(t *testing.T) {
+	checkName := func(rows *sql.Rows, name string) {
+		if !rows.Next() {
+			if err := rows.Err(); err != nil {
+				t.Fatalf("executing Next for %q failed: %v", name, err)
+			}
+			t.Fatalf("checking %q: at least one row expected", name)
+		}
+		var have string
+		err := rows.Scan(&have)
+		if err != nil {
+			t.Fatalf("executing Scan for %q failed: %v", name, err)
+		}
+		if name != have {
+			t.Fatalf("want %q, but %q found", name, have)
+		}
+	}
+	testMSSQLNextResultSet(t,
+		func(rows *sql.Rows) {
+			checkName(rows, "russ")
+			if !rows.NextResultSet() {
+				if err := rows.Err(); err != nil {
+					t.Fatal(err)
+				}
+				t.Fatal("more result sets expected")
+			}
+			checkName(rows, "brad")
+			if rows.NextResultSet() {
+				if !isFreeTDS() { // not sure why it does not work on FreeTDS
+					t.Fatal("unexpected result set found")
+				}
+			} else {
+				if err := rows.Err(); err != nil {
+					t.Fatal(err)
+				}
+			}
+		})
+}
+
+func TestMSSQLHasNextResultSet(t *testing.T) {
+	checkName := func(rows *sql.Rows, name string) {
+		var reccount int
+		for rows.Next() { // reading till the end of data set to trigger call into HasNextResultSet
+			var have string
+			err := rows.Scan(&have)
+			if err != nil {
+				t.Fatalf("executing Scan for %q failed: %v", name, err)
+			}
+			if name != have {
+				t.Fatalf("want %q, but %q found", name, have)
+			}
+			reccount++
+		}
+		if err := rows.Err(); err != nil {
+			t.Fatalf("executing Next for %q failed: %v", name, err)
+		}
+		if reccount != 1 {
+			t.Fatalf("checking %q: expected 1 row returned, but %v found", name, reccount)
+		}
+	}
+	testMSSQLNextResultSet(t,
+		func(rows *sql.Rows) {
+			checkName(rows, "russ")
+			if !rows.NextResultSet() {
+				if err := rows.Err(); err != nil {
+					t.Fatal(err)
+				}
+				t.Fatal("more result sets expected")
+			}
+			checkName(rows, "brad")
+			if rows.NextResultSet() {
+				t.Fatal("unexpected result set found")
+			} else {
+				if err := rows.Err(); err != nil {
+					t.Fatal(err)
+				}
+			}
+		})
+}
+
+func TestMSSQLIssue127(t *testing.T) {
+	db, sc, err := mssqlConnect()
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer closeDB(t, db, sc, sc)
+
+	db.Exec("drop table dbo.temp")
+	exec(t, db, "create table dbo.temp (id int, a varchar(255))")
+
+	tx, err := db.Begin()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	stmt, err := tx.Prepare(`
+DECLARE @id INT, @a VARCHAR(255)
+SELECT @id = ?, @a = ?
+UPDATE dbo.temp SET a = @a WHERE id = @id
+IF @@ROWCOUNT = 0
+  INSERT INTO dbo.temp (id, a) VALUES (@id, @a)
+`)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if _, err = stmt.Exec(1, "test"); err != nil {
+		t.Errorf("Failed to insert record with ID 1: %s", err)
+	}
+	if _, err = stmt.Exec(1, "test2"); err != nil {
+		t.Errorf("Failed to update record with ID 1: %s", err)
+	}
+
+	if err = tx.Commit(); err != nil {
+		t.Fatal(err)
+	}
+}

rows.go

@@ -44,3 +44,23 @@ func (r *Rows) Next(dest []driver.Value) error {
 func (r *Rows) Close() error {
 	return r.os.closeByRows()
 }
+
+func (r *Rows) HasNextResultSet() bool {
+	return true
+}
+
+func (r *Rows) NextResultSet() error {
+	ret := api.SQLMoreResults(r.os.h)
+	if ret == api.SQL_NO_DATA {
+		return io.EOF
+	}
+	if IsError(ret) {
+		return NewError("SQLMoreResults", r.os.h)
+	}
+	// TODO: !!!!!!!!!!!!!!! maybe need to call r.os.BindColumns here !!!!!
+	//err := r.os.BindColumns()
+	//if err != nil {
+	//	return err
+	//}
+	return nil
+}