[Cosmos] some tidy up to Cosmos query engine API (Azure#24485)

* some tidy up to Cosmos query engine API

* pass completed state through PipelineResult

Author: analogrelay

sdk/data/azcosmos/cosmos_container_query_engine.go

@@ -122,24 +122,24 @@ func (c *ContainerClient) executeQueryWithEngine(queryEngine queryengine.QueryEn
 					}, nil
 				}
 				// Fetch more data from the pipeline
-				items, requests, err := queryPipeline.NextBatch()
+				result, err := queryPipeline.Run()
 				if err != nil {
 					queryPipeline.Close()
 					return QueryItemsResponse{}, err
 				}
 
 				// If we got items, we can return them, and we should do so now, to avoid making unnecessary requests.
 				// Even if there are requests in the queue, the pipeline should return the same requests again on the next call to NextBatch.
-				if len(items) > 0 {
+				if len(result.Items) > 0 {
 					return QueryItemsResponse{
 						Response: lastResponse,
-						Items:    items,
+						Items:    result.Items,
 					}, nil
 				}
 
 				// If we didn't have any items to return, we need to make requests for the items in the queue.
 				// If there are no requests, the pipeline should return true for IsComplete, so we'll stop on the next iteration.
-				for _, request := range requests {
+				for _, request := range result.Requests {
 					// Make the single-partition query request
 					qryRequest := queryRequest(request) // Cast to our type, which has toHeaders defined on it.
 					azResponse, err := c.database.client.sendQueryRequest(

sdk/data/azcosmos/internal/mock_query_engine.go

@@ -139,9 +139,9 @@ func (m *MockQueryPipeline) IsComplete() bool {
 }
 
 // NextBatch returns the next batch of items from the pipeline, as well as any requests needed to collect more data.
-func (m *MockQueryPipeline) NextBatch() ([][]byte, []queryengine.QueryRequest, error) {
+func (m *MockQueryPipeline) Run() (*queryengine.PipelineResult, error) {
 	if m.IsClosed {
-		return nil, nil, fmt.Errorf("pipeline is closed")
+		return nil, fmt.Errorf("pipeline is closed")
 	}
 
 	items := make([][]byte, 0)
@@ -154,7 +154,11 @@ func (m *MockQueryPipeline) NextBatch() ([][]byte, []queryengine.QueryRequest, e
 		for i := range m.partitionState {
 			// If any partition hasn't started yet, we can't return any items.
 			if !m.partitionState[i].started {
-				return nil, m.getRequests(), nil
+				return &queryengine.PipelineResult{
+					IsCompleted: false,
+					Items:       nil,
+					Requests:    m.getRequests(),
+				}, nil
 			}
 
 			if m.partitionState[i].IsExhausted() {
@@ -175,7 +179,7 @@ func (m *MockQueryPipeline) NextBatch() ([][]byte, []queryengine.QueryRequest, e
 			// Add the item to the result set and remove it from the queue.
 			item, err := lowestPartition.PopItem()
 			if err != nil {
-				return nil, nil, err
+				return nil, err
 			}
 			items = append(items, item)
 		}
@@ -190,7 +194,11 @@ func (m *MockQueryPipeline) NextBatch() ([][]byte, []queryengine.QueryRequest, e
 		m.completed = true
 	}
 
-	return items, requests, nil
+	return &queryengine.PipelineResult{
+		IsCompleted: m.completed,
+		Items:       items,
+		Requests:    requests,
+	}, nil
 }
 
 // getRequests returns a list of all the QueryRequests that are needed to get the next batch of items.

sdk/data/azcosmos/unstable/queryengine/cosmos_query_engine.go

@@ -56,9 +56,8 @@ type QueryPipeline interface {
 	Query() string
 	// IsComplete gets a boolean indicating if the pipeline has concluded
 	IsComplete() bool
-	// NextBatch gets the next batch of items, which will be empty if there are no more items in the buffer, and the next set of QueryRequests which must be fulfilled, which will be empty if there are no more requests.
-	// If both the items and requests are empty, the pipeline has concluded.
-	NextBatch() ([][]byte, []QueryRequest, error)
+	// Run executes a single turn of the pipeline, yielding a PipelineResult containing the items and requests for more data.
+	Run() (*PipelineResult, error)
 	// ProvideData provides more data for a given partition key range ID, using data retrieved from the server in response to making a DataRequest.
 	ProvideData(data QueryResult) error
 	// Close frees the resources associated with the pipeline.