fix: poll query bug (#120)

* fix: poll query bug

* add test

restful.go

@@ -372,9 +372,10 @@ func (c *APIClient) applySessionState(response *QueryResponse) {
 }
 
 func (c *APIClient) PollUntilQueryEnd(ctx context.Context, resp *QueryResponse) (*QueryResponse, error) {
+	var err error
 	for !resp.ReadFinished() {
 		data := resp.Data
-		resp, err := c.PollQuery(ctx, resp.NextURI)
+		resp, err = c.PollQuery(ctx, resp.NextURI)
 		if err != nil {
 			return nil, err
 		}

tests/main_test.go

@@ -1,6 +1,7 @@
 package tests
 
 import (
+	"context"
 	"database/sql"
 	"encoding/json"
 	"fmt"
@@ -9,6 +10,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/pkg/errors"
 	"github.com/stretchr/testify/require"
 	"github.com/test-go/testify/suite"
 
@@ -327,6 +329,20 @@ func (s *DatabendTestSuite) TestTransactionRollback() {
 	s.r.NoError(rows.Close())
 }
 
+func (s *DatabendTestSuite) TestLongExec() {
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
+	_, err := s.db.ExecContext(ctx, "SELECT number from numbers(100000) order by number")
+	if err != nil {
+		if errors.Is(err, context.DeadlineExceeded) {
+			s.T().Errorf("Query execution exceeded the 10s timeout")
+		} else {
+			s.r.Nil(err)
+		}
+	}
+}
+
 func scanValues(rows *sql.Rows) (interface{}, error) {
 	var err error
 	var result [][]interface{}