Add Streaming Bulk API 2.0 Query Results Via Callback

lib/bulk2.go

@@ -525,10 +525,10 @@ func (f *Force) GetBulk2QueryJobInfoWithContext(ctx context.Context, jobId strin
 	}
 }
 
-// GetBulk2QueryResults retrieves results for a Bulk API 2.0 query job
-// locator is used for pagination, pass empty string for the first page
-// maxRecords limits the number of records returned (0 for default)
-func (f *Force) GetBulk2QueryResults(jobId string, locator string, maxRecords int) (Bulk2QueryResults, error) {
+// bulk2QueryResultsUrl builds the results URL for a Bulk API 2.0 query job.
+// locator paginates results (empty for the first page); maxRecords caps the
+// page size (0 leaves the page size up to Salesforce).
+func (f *Force) bulk2QueryResultsUrl(jobId string, locator string, maxRecords int) string {
 	url := fmt.Sprintf("%s/%s/results", f.bulk2QueryUrl(), jobId)
 	if locator != "" {
 		url = fmt.Sprintf("%s?locator=%s", url, locator)
@@ -538,9 +538,18 @@ func (f *Force) GetBulk2QueryResults(jobId string, locator string, maxRecords in
 	} else if maxRecords > 0 {
 		url = fmt.Sprintf("%s?maxRecords=%d", url, maxRecords)
 	}
+	return url
+}
 
+// GetBulk2QueryResults retrieves results for a Bulk API 2.0 query job
+// locator is used for pagination, pass empty string for the first page
+// maxRecords limits the number of records returned (0 for default)
+//
+// The entire page is buffered in memory. For large result sets, prefer
+// GetBulk2QueryResultsWithCallback to stream the page without buffering it.
+func (f *Force) GetBulk2QueryResults(jobId string, locator string, maxRecords int) (Bulk2QueryResults, error) {
 	req := NewRequest("GET").
-		AbsoluteUrl(url).
+		AbsoluteUrl(f.bulk2QueryResultsUrl(jobId, locator, maxRecords)).
 		WithHeader("Accept", "text/csv").
 		ReadResponseBody()
 
@@ -576,6 +585,40 @@ func (f *Force) GetBulk2QueryResultsWithContext(ctx context.Context, jobId strin
 	}
 }
 
+// GetBulk2QueryResultsWithCallback retrieves results for a Bulk API 2.0 query
+// job, invoking callback with the raw HTTP response so the caller can stream the
+// CSV body (for example, to a file) instead of buffering the whole page in
+// memory. The callback is responsible for reading and closing the response
+// body. The next-page locator is available from the "Sforce-Locator" response
+// header; an empty or "null" value indicates the final page.
+//
+// The callback is only invoked for a successful response; HTTP errors are
+// returned directly and the body is not passed to the callback.
+func (f *Force) GetBulk2QueryResultsWithCallback(jobId string, locator string, maxRecords int, callback HttpCallback) error {
+	req := NewRequest("GET").
+		AbsoluteUrl(f.bulk2QueryResultsUrl(jobId, locator, maxRecords)).
+		WithHeader("Accept", "text/csv").
+		WithResponseCallback(callback)
+	_, err := f.ExecuteRequest(req)
+	return err
+}
+
+// GetBulk2QueryResultsWithCallbackWithContext retrieves results for a Bulk API 2.0 query job with context
+func (f *Force) GetBulk2QueryResultsWithCallbackWithContext(ctx context.Context, jobId string, locator string, maxRecords int, callback HttpCallback) error {
+	done := make(chan struct{})
+	var err error
+	go func() {
+		defer close(done)
+		err = f.GetBulk2QueryResultsWithCallback(jobId, locator, maxRecords, callback)
+	}()
+	select {
+	case <-ctx.Done():
+		return fmt.Errorf("GetBulk2QueryResultsWithCallback canceled: %w", ctx.Err())
+	case <-done:
+		return err
+	}
+}
+
 // AbortBulk2QueryJob aborts a Bulk API 2.0 query job
 func (f *Force) AbortBulk2QueryJob(jobId string) (Bulk2QueryJobInfo, error) {
 	url := fmt.Sprintf("%s/%s", f.bulk2QueryUrl(), jobId)

lib/bulk2_test.go

@@ -1181,6 +1181,74 @@ func TestGetBulk2QueryResults_WithMaxRecords(t *testing.T) {
 	}
 }
 
+func TestGetBulk2QueryResultsWithCallback(t *testing.T) {
+	expectedCSV := "Id,Name\n001000000000001AAA,Test Account\n"
+
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if !strings.Contains(r.URL.RawQuery, "maxRecords=100") {
+			t.Errorf("Expected maxRecords=100, got: %s", r.URL.RawQuery)
+		}
+		w.Header().Set("Content-Type", "text/csv")
+		w.Header().Set("Sforce-Locator", "ABC123")
+		w.WriteHeader(http.StatusOK)
+		w.Write([]byte(expectedCSV))
+	}))
+	defer server.Close()
+
+	force := &Force{
+		Credentials: &ForceSession{
+			InstanceUrl: server.URL,
+			AccessToken: "test-token",
+		},
+	}
+
+	var buf bytes.Buffer
+	var locator string
+	err := force.GetBulk2QueryResultsWithCallback("7501234567890QUERY", "", 100, func(resp *http.Response) error {
+		defer resp.Body.Close()
+		locator = resp.Header.Get("Sforce-Locator")
+		_, err := io.Copy(&buf, resp.Body)
+		return err
+	})
+	if err != nil {
+		t.Fatalf("GetBulk2QueryResultsWithCallback failed: %v", err)
+	}
+	if buf.String() != expectedCSV {
+		t.Errorf("Expected data '%s', got '%s'", expectedCSV, buf.String())
+	}
+	if locator != "ABC123" {
+		t.Errorf("Expected locator 'ABC123', got '%s'", locator)
+	}
+}
+
+func TestGetBulk2QueryResultsWithCallback_HttpError(t *testing.T) {
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "application/json")
+		w.WriteHeader(http.StatusBadRequest)
+		w.Write([]byte(`[{"errorCode":"INVALIDJOB","message":"job not found"}]`))
+	}))
+	defer server.Close()
+
+	force := &Force{
+		Credentials: &ForceSession{
+			InstanceUrl: server.URL,
+			AccessToken: "test-token",
+		},
+	}
+
+	called := false
+	err := force.GetBulk2QueryResultsWithCallback("7501234567890QUERY", "", 0, func(resp *http.Response) error {
+		called = true
+		return nil
+	})
+	if err == nil {
+		t.Fatal("Expected an error for a non-2xx response, got nil")
+	}
+	if called {
+		t.Error("Callback should not be invoked for an HTTP error response")
+	}
+}
+
 func TestAbortBulk2QueryJob(t *testing.T) {
 	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		if r.Method != "PATCH" {