@@ -25,14 +25,15 @@ import (
25
25
"encoding/binary"
26
26
"errors"
27
27
"fmt"
28
- "github.com/apache/iotdb-client-go/common"
29
28
"log"
30
29
"net"
31
30
"reflect"
32
31
"sort"
33
32
"strings"
34
33
"time"
35
34
35
+ "github.com/apache/iotdb-client-go/common"
36
+
36
37
"github.com/apache/iotdb-client-go/rpc"
37
38
"github.com/apache/thrift/lib/go/thrift"
38
39
)
@@ -415,26 +416,34 @@ func (s *Session) SetTimeZone(timeZone string) (r *common.TSStatus, err error) {
415
416
return r , err
416
417
}
417
418
418
- func (s * Session ) ExecuteStatement ( sql string ) (* SessionDataSet , error ) {
419
+ func (s * Session ) ExecuteStatementWithContext ( ctx context. Context , sql string ) (* SessionDataSet , error ) {
419
420
request := rpc.TSExecuteStatementReq {
420
421
SessionId : s .sessionId ,
421
422
Statement : sql ,
422
423
StatementId : s .requestStatementId ,
423
424
FetchSize : & s .config .FetchSize ,
424
425
}
425
- resp , err := s .client .ExecuteStatement (context . Background () , & request )
426
+ resp , err := s .client .ExecuteStatement (ctx , & request )
426
427
427
428
if err != nil && resp == nil {
428
429
if s .reconnect () {
429
430
request .SessionId = s .sessionId
430
431
request .StatementId = s .requestStatementId
431
- resp , err = s .client .ExecuteStatement (context . Background () , & request )
432
+ resp , err = s .client .ExecuteStatement (ctx , & request )
432
433
}
433
434
}
434
435
436
+ if statusErr := VerifySuccess (resp .Status ); statusErr != nil {
437
+ return nil , statusErr
438
+ }
439
+
435
440
return s .genDataSet (sql , resp ), err
436
441
}
437
442
443
+ func (s * Session ) ExecuteStatement (sql string ) (* SessionDataSet , error ) {
444
+ return s .ExecuteStatementWithContext (context .Background (), sql )
445
+ }
446
+
438
447
func (s * Session ) ExecuteNonQueryStatement (sql string ) (r * common.TSStatus , err error ) {
439
448
request := rpc.TSExecuteStatementReq {
440
449
SessionId : s .sessionId ,
0 commit comments