Merge pull request #169 from winebarrel/add_WaitQueryJSON

Add WaitQueryJSON

README.md

@@ -64,27 +64,10 @@ func main() {
 		panic(err)
 	}
 
-	if job != nil {
-		for {
-			job, err := client.GetJob(ctx, job.Job.ID)
+  err = client.WaitQueryJSON(ctx, query.ID, job, nil, &buf)
 
-			if err != nil {
-				panic(err)
-			}
-
-			if job.Job.Status != redash.JobStatusPending && job.Job.Status != redash.JobStatusStarted {
-				buf = bytes.Buffer{}
-				err := client.GetQueryResultsJSON(ctx, query.ID, &buf)
-
-				if err != nil {
-					panic(err)
-				}
-
-				break
-			}
-
-			time.Sleep(1 * time.Second)
-		}
+	if err != nil {
+		panic(err)
 	}
 
 	fmt.Println(buf.String())
@@ -98,25 +81,19 @@ input := &redash.ExecQueryJSONInput{
   WithoutOmittingMaxAge: true,
 }
 
-// If `max_age=0`, no result is returned.
-// Results should be obtained with the GetQueryResultsXXX method.
 job, err := client.ExecQueryJSON(ctx, query.ID, input, nil)
 
 if err != nil {
   panic(err)
 }
 
-if job != nil {
-  // Waiting...
-}
-
-out, err := client.GetQueryResultsStruct(context.Background(), query.ID)
+err = client.WaitQueryJSON(ctx, query.ID, job, nil, &buf)
 
 if err != nil {
   panic(err)
 }
 
-fmt.Println(out)
+fmt.Println(buf.String())
 ```
 
 ### Set debug mode

codecov.yml

@@ -1,3 +1,11 @@
 ignore:
   - ".*_without_ctx\\.go"
   - "tools"
+coverage:
+  status:
+    project:
+      default:
+        informational: true
+    patch:
+      default:
+        informational: true

query.go

@@ -8,6 +8,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"slices"
 	"time"
 
 	"github.com/winebarrel/redash-go/v2/internal/util"
@@ -344,6 +345,63 @@ func (client *Client) ExecQueryJSON(ctx context.Context, id int, input *ExecQuer
 	return nil, err
 }
 
+var (
+	defaultWaitQueryJSONOptionWaitStatuses = []int{
+		JobStatusPending,
+		JobStatusStarted,
+	}
+)
+
+const (
+	defaultWaitQueryJSONOptionInterval = 1 * time.Second
+)
+
+type WaitQueryJSONOption struct {
+	WaitStatuses []int
+	Interval     time.Duration
+}
+
+func (client *Client) WaitQueryJSON(ctx context.Context, queryId int, job *JobResponse, option *WaitQueryJSONOption, out io.Writer) error {
+	if job == nil || job.Job.ID == "" {
+		return nil
+	}
+
+	waitStatus := defaultWaitQueryJSONOptionWaitStatuses
+	interval := defaultWaitQueryJSONOptionInterval
+
+	if option != nil {
+		if len(option.WaitStatuses) > 0 {
+			waitStatus = option.WaitStatuses
+		}
+
+		if option.Interval > 0 {
+			interval = option.Interval
+		}
+	}
+
+	for {
+		job, err := client.GetJob(ctx, job.Job.ID)
+
+		if err != nil {
+			return err
+		}
+
+		if !slices.Contains(waitStatus, job.Job.Status) {
+			err := client.GetQueryResultsJSON(ctx, queryId, out)
+
+			if err != nil {
+				return err
+			}
+
+			break
+		}
+
+		time.Sleep(interval)
+	}
+
+	return nil
+}
+
 type QueryTags struct {
 	Tags []QueryTagsTag `json:"tags"`
 }

query_test.go

@@ -1025,6 +1025,76 @@ func Test_ExecQueryJSON_ReturnJob(t *testing.T) {
 	}, job)
 }
 
+func Test_WaitQueryJSON(t *testing.T) {
+	assert := assert.New(t)
+	require := require.New(t)
+	httpmock.Activate()
+	defer httpmock.DeactivateAndReset()
+
+	httpmock.RegisterResponder(http.MethodPost, "https://redash.example.com/api/queries/1/results", func(req *http.Request) (*http.Response, error) {
+		assert.Equal(
+			http.Header(
+				http.Header{
+					"Authorization": []string{"Key " + testRedashAPIKey},
+					"Content-Type":  []string{"application/json"},
+					"User-Agent":    []string{"redash-go"},
+				},
+			),
+			req.Header,
+		)
+		require.NotNil(req.Body)
+		body, _ := io.ReadAll(req.Body)
+		assert.Equal(`{}`, string(body))
+		return httpmock.NewStringResponse(http.StatusOK, `{"job": {"status": 1, "error": "", "id": "623b290a-7fd9-4ea6-a2a6-96f9c9101f51", "query_result_id": null,	"status": 1, "updated_at": 0}}`), nil
+	})
+
+	httpmock.RegisterResponder(http.MethodGet, "https://redash.example.com/api/jobs/623b290a-7fd9-4ea6-a2a6-96f9c9101f51", func(req *http.Request) (*http.Response, error) {
+		assert.Equal(
+			http.Header(
+				http.Header{
+					"Authorization": []string{"Key " + testRedashAPIKey},
+					"Content-Type":  []string{"application/json"},
+					"User-Agent":    []string{"redash-go"},
+				},
+			),
+			req.Header,
+		)
+		return httpmock.NewStringResponse(http.StatusOK, `
+			{
+				"job": {
+					"error": "",
+					"id": "623b290a-7fd9-4ea6-a2a6-96f9c9101f51",
+					"query_result_id": 1,
+					"status": 3,
+					"updated_at": 0
+				}
+			}
+		`), nil
+	})
+
+	httpmock.RegisterResponder(http.MethodGet, "https://redash.example.com/api/queries/1/results.json", func(req *http.Request) (*http.Response, error) {
+		assert.Equal(
+			http.Header(
+				http.Header{
+					"Authorization": []string{"Key " + testRedashAPIKey},
+					"Content-Type":  []string{"application/json"},
+					"User-Agent":    []string{"redash-go"},
+				},
+			),
+			req.Header,
+		)
+		return httpmock.NewStringResponse(http.StatusOK, `{"foo":"bar"}`), nil
+	})
+
+	client, _ := redash.NewClient("https://redash.example.com", testRedashAPIKey)
+	var buf bytes.Buffer
+	job, err := client.ExecQueryJSON(context.Background(), 1, &redash.ExecQueryJSONInput{}, &buf)
+	assert.NoError(err)
+	err = client.WaitQueryJSON(context.Background(), 1, job, nil, &buf)
+	assert.NoError(err)
+	assert.Equal(`{"foo":"bar"}`, buf.String())
+}
+
 func Test_GetQueryTags_OK(t *testing.T) {
 	assert := assert.New(t)
 	httpmock.Activate()
@@ -1980,3 +2050,56 @@ func Test_Query_IgnoreCache_Acc(t *testing.T) {
 		assert.NotEqual(cachedNow, now)
 	}
 }
+
+func Test_WaitQueryJSON_Acc(t *testing.T) {
+	if !testAcc {
+		t.Skip()
+	}
+
+	assert := assert.New(t)
+	require := require.New(t)
+	client, _ := redash.NewClient(testRedashEndpoint, testRedashAPIKey)
+	ds, err := client.CreateDataSource(context.Background(), &redash.CreateDataSourceInput{
+		Name: "test-postgres-1",
+		Type: "pg",
+		Options: map[string]any{
+			"dbname": "postgres",
+			"host":   "postgres",
+			"port":   5432,
+			"user":   "postgres",
+		},
+	})
+	require.NoError(err)
+
+	defer func() {
+		client.DeleteDataSource(context.Background(), ds.ID) //nolint:errcheck
+	}()
+
+	_, err = client.ListQueries(context.Background(), nil)
+	require.NoError(err)
+
+	query, err := client.CreateQuery(context.Background(), &redash.CreateQueryInput{
+		DataSourceID: ds.ID,
+		Name:         "test-query-1",
+		Query:        "select 1",
+	})
+	require.NoError(err)
+	assert.Equal("test-query-1", query.Name)
+
+	var buf bytes.Buffer
+	job, err := client.ExecQueryJSON(context.Background(), query.ID, nil, &buf)
+	require.NoError(err)
+	err = client.WaitQueryJSON(context.Background(), query.ID, job, nil, &buf)
+	require.NoError(err)
+	assert.True(strings.HasPrefix(buf.String(), `{"query_result"`))
+
+	buf.Reset()
+	job, err = client.ExecQueryJSON(context.Background(), query.ID, nil, &buf)
+	require.NoError(err)
+	err = client.WaitQueryJSON(context.Background(), query.ID, job, &redash.WaitQueryJSONOption{
+		WaitStatuses: []int{redash.JobStatusPending, redash.JobStatusStarted},
+		Interval:     500 * time.Microsecond,
+	}, &buf)
+	require.NoError(err)
+	assert.True(strings.HasPrefix(buf.String(), `{"query_result"`))
+}