fix(bigquery): wait on job status instead of getQueryResults (BRU-5103)

pkg/bigquery/db.go

@@ -273,6 +273,43 @@ func (d *Client) IsValid(ctx context.Context, query *query.Query) (bool, error)
 	return true, nil
 }
 
+// waitForJobCompletion polls job.Status until the job is terminal and returns its
+// error, if any. Unlike job.Read (jobs.getQueryResults), job.Status reports a
+// failed job's state immediately instead of retrying it as a transient error, so a
+// rate-limited job surfaces at once rather than hanging (BRU-5103).
+func waitForJobCompletion(ctx context.Context, job *bigquery.Job) error {
+	const (
+		initialPollInterval = time.Second
+		maxPollInterval     = 10 * time.Second
+	)
+	wait := initialPollInterval
+	for {
+		// Use an already-terminal status (from submit or a prior poll) to skip an RPC.
+		status := job.LastStatus()
+		if status == nil || !status.Done() {
+			fresh, err := job.Status(ctx)
+			if err != nil {
+				return err
+			}
+			status = fresh
+		}
+		if status.Done() {
+			return status.Err()
+		}
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		case <-time.After(wait):
+		}
+		if wait < maxPollInterval {
+			wait *= 2
+			if wait > maxPollInterval {
+				wait = maxPollInterval
+			}
+		}
+	}
+}
+
 func (d *Client) RunQueryWithoutResult(ctx context.Context, q *query.Query) error {
 	if err := d.ensureClientInitialized(ctx); err != nil {
 		return err
@@ -286,8 +323,7 @@ func (d *Client) RunQueryWithoutResult(ctx context.Context, q *query.Query) erro
 		return formatError(err)
 	}
 	query.LogOrSinkQueryID(ctx, "BigQuery", job.ID())
-	_, err = job.Read(ctx)
-	if err != nil {
+	if err := waitForJobCompletion(ctx, job); err != nil {
 		return formatError(err)
 	}
 
@@ -307,6 +343,10 @@ func (d *Client) Select(ctx context.Context, q *query.Query) ([][]interface{}, e
 		return nil, formatError(err)
 	}
 	query.LogOrSinkQueryID(ctx, "BigQuery", job.ID())
+	// Surface a failed job immediately; job.Read below then returns without retrying (BRU-5103).
+	if err := waitForJobCompletion(ctx, job); err != nil {
+		return nil, formatError(err)
+	}
 	rows, err := job.Read(ctx)
 	if err != nil {
 		return nil, formatError(err)
@@ -347,6 +387,10 @@ func (d *Client) SelectWithSchema(ctx context.Context, queryObj *query.Query) (*
 		return nil, fmt.Errorf("failed to run query: %w", formatError(err))
 	}
 	query.LogOrSinkQueryID(ctx, "BigQuery", job.ID())
+	// Surface a failed job immediately; job.Read below then returns without retrying (BRU-5103).
+	if err := waitForJobCompletion(ctx, job); err != nil {
+		return nil, fmt.Errorf("failed to wait for query completion: %w", formatError(err))
+	}
 	rows, err := job.Read(ctx)
 	if err != nil {
 		return nil, fmt.Errorf("failed to read query results: %w", formatError(err))
@@ -718,17 +762,31 @@ func (d *Client) UpdateTableMetadataIfNotExist(ctx context.Context, asset *pipel
 
 func formatError(err error) error {
 	var googleError *googleapi.Error
-	if !errors.As(err, &googleError) {
-		return err
+	if errors.As(err, &googleError) {
+		if googleError.Code == 404 || googleError.Code == 400 {
+			return fmt.Errorf("%s", googleError.Message)
+		}
+		return googleError
 	}
 
-	if googleError.Code == 404 || googleError.Code == 400 {
-		return fmt.Errorf("%s", googleError.Message)
+	// Terminal job failures surface as *bigquery.Error (via job.Status), whose default
+	// string is a verbose struct dump. Present its message cleanly while preserving the
+	// structured error so callers can still errors.As it for Reason/Location.
+	var bqError *bigquery.Error
+	if errors.As(err, &bqError) && bqError.Message != "" {
+		return &jobError{err: bqError}
 	}
 
-	return googleError
+	return err
 }
 
+// jobError presents a *bigquery.Error's message (its default string is a verbose
+// struct dump) while preserving the underlying error for errors.As.
+type jobError struct{ err *bigquery.Error }
+
+func (e *jobError) Error() string { return e.err.Message }
+func (e *jobError) Unwrap() error { return e.err }
+
 // Test runs a simple query (SELECT 1) to validate the connection.
 func (d *Client) Ping(ctx context.Context) error {
 	// Define the test query

pkg/bigquery/db_test.go

@@ -267,6 +267,72 @@ func TestDB_RunQueryWithoutResult(t *testing.T) {
 	}
 }
 
+// Regression test for BRU-5103: a job that submits as RUNNING then completes with a
+// retryable failure reason (rateLimitExceeded) must surface the error promptly
+// rather than being retried for hours via job.Read/getQueryResults.
+func TestDB_RunQueryWithoutResultSurfacesTerminalJobError(t *testing.T) {
+	t.Parallel()
+
+	projectID := testProjectID
+	jobID := testJobID
+
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		switch {
+		case r.Method == http.MethodPost && strings.HasPrefix(r.RequestURI, fmt.Sprintf("/projects/%s/jobs", projectID)):
+			// jobs.insert: job accepted, still running.
+			writeBqTestJSON(t, w, &bigquery2.Job{
+				JobReference: &bigquery2.JobReference{JobId: jobID, ProjectId: projectID, Location: "US"},
+				Status:       &bigquery2.JobStatus{State: "RUNNING"},
+			})
+		case r.Method == http.MethodGet && strings.HasPrefix(r.RequestURI, fmt.Sprintf("/projects/%s/jobs/%s", projectID, jobID)):
+			// jobs.get: job finished with a rate-limit error.
+			writeBqTestJSON(t, w, &bigquery2.Job{
+				JobReference: &bigquery2.JobReference{JobId: jobID, ProjectId: projectID, Location: "US"},
+				Status: &bigquery2.JobStatus{
+					State: "DONE",
+					ErrorResult: &bigquery2.ErrorProto{
+						Reason:  "rateLimitExceeded",
+						Message: "Exceeded rate limits: too many table update operations for this table",
+					},
+				},
+			})
+		default:
+			w.WriteHeader(http.StatusInternalServerError)
+			_, err := w.Write([]byte("unexpected request: " + r.Method + " " + r.RequestURI))
+			if err != nil {
+				t.Fatal(err)
+			}
+		}
+	}))
+	defer server.Close()
+
+	client, err := bigquery.NewClient(
+		t.Context(),
+		projectID,
+		option.WithEndpoint(server.URL),
+		option.WithCredentials(&google.Credentials{
+			ProjectID: projectID,
+			TokenSource: oauth2.StaticTokenSource(&oauth2.Token{
+				AccessToken: "some-token",
+			}),
+		}),
+	)
+	require.NoError(t, err)
+	client.Location = "US"
+
+	d := Client{client: client}
+
+	err = d.RunQueryWithoutResult(t.Context(), &query.Query{Query: "CREATE TABLE IF NOT EXISTS dataset.t (id INT64)"})
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "Exceeded rate limits")
+	assert.NotContains(t, err.Error(), "Location:") // clean message, not the *bigquery.Error struct dump
+
+	// The structured error is preserved for classification.
+	var bqErr *bigquery.Error
+	require.ErrorAs(t, err, &bqErr)
+	assert.Equal(t, "rateLimitExceeded", bqErr.Reason)
+}
+
 func TestClientValidateQueryLimits(t *testing.T) {
 	t.Parallel()
 
@@ -462,41 +528,34 @@ type queryResultResponse struct {
 	statusCode int
 }
 
+// mockBqHandler serves the same jsr for jobs.insert and jobs.get, so jobs.insert
+// must report a terminal (DONE) state — a RUNNING jsr would make job.Status polling
+// loop forever. For a submit-RUNNING-then-complete flow, use a dedicated handler.
 func mockBqHandler(t *testing.T, projectID, jobID string, jsr jobSubmitResponse, qrr queryResultResponse) http.Handler {
-	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		if r.Method == http.MethodGet && strings.HasPrefix(r.RequestURI, fmt.Sprintf("/projects/%s/queries/%s?", projectID, jobID)) {
-			w.WriteHeader(qrr.statusCode)
-
-			response, err := json.Marshal(qrr.response)
-			if err != nil {
-				t.Fatal(err)
-			}
-
-			_, err = w.Write(response)
-			if err != nil {
-				t.Fatal(err)
-			}
-			return
-		} else if r.Method == http.MethodPost && strings.HasPrefix(r.RequestURI, fmt.Sprintf("/projects/%s/jobs", projectID)) {
-			// Handle jobs.insert (used by q.Run)
-			w.WriteHeader(jsr.statusCode)
-
-			response, err := json.Marshal(jsr.response)
-			if err != nil {
-				t.Fatal(err)
-			}
+	write := func(w http.ResponseWriter, statusCode int, body interface{}) {
+		w.WriteHeader(statusCode)
+		response, err := json.Marshal(body)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if _, err := w.Write(response); err != nil {
+			t.Fatal(err)
+		}
+	}
 
-			_, err = w.Write(response)
-			if err != nil {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		switch {
+		case r.Method == http.MethodGet && strings.HasPrefix(r.RequestURI, fmt.Sprintf("/projects/%s/queries/%s?", projectID, jobID)):
+			write(w, qrr.statusCode, qrr.response) // jobs.getQueryResults (job.Read)
+		case r.Method == http.MethodGet && strings.HasPrefix(r.RequestURI, fmt.Sprintf("/projects/%s/jobs/%s", projectID, jobID)):
+			write(w, jsr.statusCode, jsr.response) // jobs.get (job.Status)
+		case r.Method == http.MethodPost && strings.HasPrefix(r.RequestURI, fmt.Sprintf("/projects/%s/jobs", projectID)):
+			write(w, jsr.statusCode, jsr.response) // jobs.insert (q.Run)
+		default:
+			w.WriteHeader(http.StatusInternalServerError)
+			if _, err := w.Write([]byte("there is no test definition found for the given request: " + r.Method + " " + r.RequestURI)); err != nil {
 				t.Fatal(err)
 			}
-			return
-		}
-
-		w.WriteHeader(http.StatusInternalServerError)
-		_, err := w.Write([]byte("there is no test definition found for the given request: " + r.Method + " " + r.RequestURI))
-		if err != nil {
-			t.Fatal(err)
 		}
 	})
 }