Repository: Azure/azure-sdk-for-go
Pull request #19769 (2023-01-06 23:06:48)

Cosmos DB - Add support for querying databases and containers
This PR adds support to list/query databases within an account and containers within a database, including support for parametrized querying:

```
queryPager := client.NewQueryDatabasesPager("select * from dbs d", nil)

opt := &azcosmos.QueryDatabasesOptions{
		QueryParameters: []azcosmos.QueryParameter{
			{"@id", "the_id"},
		},
	}

queryPager := client.NewQueryDatabasesPager("select * from dbs d where d.id = @id", opt)
```

```
queryPager := database.NewQueryContainersPager("select * from colls c", nil)

opt := &azcosmos.QueryContainersOptions{
		QueryParameters: []azcosmos.QueryParameter{
			{"@id", "the_id"},
		},
	}

queryPager := database.NewQueryContainersPager("select * from colls c where c.id = @id", opt)
```

## APIView

https://apiview.dev/Assemblies/Review/b128254cd7ac437390ea6c0c6ef4a5de/7294484ab73e4677bc962c0137b42104?diffRevisionId=1f3ba5313df24d4c9800d33af3254c9c&doc=False&diffOnly=True


- [x] The purpose of this PR is explained in this or a referenced issue.
- [x] The PR does not update generated files.
- [x] Tests are included and/or updated for code changes.
- [x] Updates to [CHANGELOG.md][] are included.
- [x] MIT license headers are included in each file.

Closes https://github.com/Azure/azure-sdk-for-go/issues/19749




Patch:
diff --git a/eng/common/pipelines/templates/stages/archetype-sdk-tool-pwsh.yml b/eng/common/pipelines/templates/stages/archetype-sdk-tool-pwsh.yml
new file mode 100644
index 0000000000..0b46240cc7
--- /dev/null
+++ b/eng/common/pipelines/templates/stages/archetype-sdk-tool-pwsh.yml
@@ -0,0 +1,43 @@
+parameters:
+  - name: TargetDirectory
+    type: string
+  - name: CustomTestSteps
+    type: object
+    default: []
+  - name: EnvVars
+    type: object
+    default: {}
+  - name: TargetTags
+    type: string
+    default: ''
+
+variables:
+  - template: /eng/pipelines/templates/variables/globals.yml
+
+stages:
+  - stage: 'eng_script_tests'
+    jobs:
+      - job: 'Test'
+        strategy:
+          matrix:
+            Windows:
+              Pool: 'azsdk-pool-mms-win-2022-general'
+              Image: 'MMS2022'
+            Linux:
+              Pool: azsdk-pool-mms-ubuntu-2204-general
+              Image: MMSUbuntu22.04
+            Mac:
+              Pool: 'Azure Pipelines'
+              Image: 'macos-11'
+
+        pool:
+          name: $(Pool)
+          vmImage: $(Image)
+
+        steps:
+          - template: /eng/common/pipelines/templates/steps/run-pester-tests.yml
+            parameters:
+              TargetDirectory: ${{ parameters.TargetDirectory }}
+              CustomTestSteps: ${{ parameters.CustomTestSteps }}
+              EnvVars: ${{ parameters.EnvVars }}
+              TargetTags: ${{ parameters.TargetTags }}
diff --git a/eng/common/scripts/job-matrix/README.md b/eng/common/scripts/job-matrix/README.md
index 3d08881523..5ca842602c 100644
--- a/eng/common/scripts/job-matrix/README.md
+++ b/eng/common/scripts/job-matrix/README.md
@@ -2,4 +2,4 @@
 
 The documentation for Azure Pipelines Matrix Generator can be found at
 
-`azure-sdk-tools/doc/common/matrix_generator.md`
+[`azure-sdk-tools/doc/common/matrix_generator.md`](https://github.com/Azure/azure-sdk-tools/blob/main/doc/common/matrix_generator.md)
diff --git a/sdk/data/azcosmos/CHANGELOG.md b/sdk/data/azcosmos/CHANGELOG.md
index b325af2d3b..79a4a1600b 100644
--- a/sdk/data/azcosmos/CHANGELOG.md
+++ b/sdk/data/azcosmos/CHANGELOG.md
@@ -1,15 +1,10 @@
 # Release History
 
-## 0.3.3 (Unreleased)
+## 0.3.3 (2023-01-10)
 
 ### Features Added
 * Added `PatchItem` function to patch documents
-
-### Breaking Changes
-
-### Bugs Fixed
-
-### Other Changes
+* Added support for querying databases and containers
 
 ## 0.3.2 (2022-08-09)
 
diff --git a/sdk/data/azcosmos/cosmos_client.go b/sdk/data/azcosmos/cosmos_client.go
index 45b71e2830..94c455aed1 100644
--- a/sdk/data/azcosmos/cosmos_client.go
+++ b/sdk/data/azcosmos/cosmos_client.go
@@ -15,6 +15,7 @@ import (
 
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/runtime"
 	azruntime "github.com/Azure/azure-sdk-for-go/sdk/azcore/runtime"
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/streaming"
 )
@@ -179,6 +180,53 @@ func (c *Client) CreateDatabase(
 	return newDatabaseResponse(azResponse)
 }
 
+// NewQueryDatabasesPager executes query for databases.
+// query - The SQL query to execute.
+// o - Options for the operation.
+func (c *Client) NewQueryDatabasesPager(query string, o *QueryDatabasesOptions) *runtime.Pager[QueryDatabasesResponse] {
+	queryOptions := &QueryDatabasesOptions{}
+	if o != nil {
+		originalOptions := *o
+		queryOptions = &originalOptions
+	}
+
+	operationContext := pipelineRequestOptions{
+		resourceType:    resourceTypeDatabase,
+		resourceAddress: "",
+	}
+
+	path, _ := generatePathForNameBased(resourceTypeDatabase, operationContext.resourceAddress, true)
+
+	return runtime.NewPager(runtime.PagingHandler[QueryDatabasesResponse]{
+		More: func(page QueryDatabasesResponse) bool {
+			return page.ContinuationToken != ""
+		},
+		Fetcher: func(ctx context.Context, page *QueryDatabasesResponse) (QueryDatabasesResponse, error) {
+			if page != nil {
+				if page.ContinuationToken != "" {
+					// Use the previous page continuation if available
+					queryOptions.ContinuationToken = page.ContinuationToken
+				}
+			}
+
+			azResponse, err := c.sendQueryRequest(
+				path,
+				ctx,
+				query,
+				queryOptions.QueryParameters,
+				operationContext,
+				queryOptions,
+				nil)
+
+			if err != nil {
+				return QueryDatabasesResponse{}, err
+			}
+
+			return newDatabasesQueryResponse(azResponse)
+		},
+	})
+}
+
 func (c *Client) sendPostRequest(
 	path string,
 	ctx context.Context,
diff --git a/sdk/data/azcosmos/cosmos_container.go b/sdk/data/azcosmos/cosmos_container.go
index 418fa13f2a..c9f4a7fac3 100644
--- a/sdk/data/azcosmos/cosmos_container.go
+++ b/sdk/data/azcosmos/cosmos_container.go
@@ -403,7 +403,6 @@ func (c *ContainerClient) DeleteItem(
 }
 
 // NewQueryItemsPager executes a single partition query in a Cosmos container.
-// ctx - The context for the request.
 // query - The SQL query to execute.
 // partitionKey - The partition key to scope the query on.
 // o - Options for the operation.
diff --git a/sdk/data/azcosmos/cosmos_container_request_options.go b/sdk/data/azcosmos/cosmos_container_request_options.go
index 2424e65200..5dcbc9ab91 100644
--- a/sdk/data/azcosmos/cosmos_container_request_options.go
+++ b/sdk/data/azcosmos/cosmos_container_request_options.go
@@ -40,3 +40,24 @@ type DeleteContainerOptions struct{}
 func (options *DeleteContainerOptions) toHeaders() *map[string]string {
 	return nil
 }
+
+// QueryContainersOptions are options to query containers
+type QueryContainersOptions struct {
+	// ContinuationToken to be used to continue a previous query execution.
+	// Obtained from QueryItemsResponse.ContinuationToken.
+	ContinuationToken string
+
+	// QueryParameters allows execution of parametrized queries.
+	// See https://docs.microsoft.com/azure/cosmos-db/sql/sql-query-parameterized-queries
+	QueryParameters []QueryParameter
+}
+
+func (options *QueryContainersOptions) toHeaders() *map[string]string {
+	headers := make(map[string]string)
+
+	if options.ContinuationToken != "" {
+		headers[cosmosHeaderContinuationToken] = options.ContinuationToken
+	}
+
+	return &headers
+}
diff --git a/sdk/data/azcosmos/cosmos_database.go b/sdk/data/azcosmos/cosmos_database.go
index 7d0f6b50dc..957aff3afd 100644
--- a/sdk/data/azcosmos/cosmos_database.go
+++ b/sdk/data/azcosmos/cosmos_database.go
@@ -6,6 +6,8 @@ package azcosmos
 import (
 	"context"
 	"errors"
+
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/runtime"
 )
 
 // DatabaseClient lets you perform read, update, change throughput, and delete database operations.
@@ -76,6 +78,53 @@ func (db *DatabaseClient) CreateContainer(
 	return newContainerResponse(azResponse)
 }
 
+// NewQueryContainersPager executes query for containers within a database.
+// query - The SQL query to execute.
+// o - Options for the operation.
+func (c *DatabaseClient) NewQueryContainersPager(query string, o *QueryContainersOptions) *runtime.Pager[QueryContainersResponse] {
+	queryOptions := &QueryContainersOptions{}
+	if o != nil {
+		originalOptions := *o
+		queryOptions = &originalOptions
+	}
+
+	operationContext := pipelineRequestOptions{
+		resourceType:    resourceTypeCollection,
+		resourceAddress: c.link,
+	}
+
+	path, _ := generatePathForNameBased(resourceTypeCollection, operationContext.resourceAddress, true)
+
+	return runtime.NewPager(runtime.PagingHandler[QueryContainersResponse]{
+		More: func(page QueryContainersResponse) bool {
+			return page.ContinuationToken != ""
+		},
+		Fetcher: func(ctx context.Context, page *QueryContainersResponse) (QueryContainersResponse, error) {
+			if page != nil {
+				if page.ContinuationToken != "" {
+					// Use the previous page continuation if available
+					queryOptions.ContinuationToken = page.ContinuationToken
+				}
+			}
+
+			azResponse, err := c.client.sendQueryRequest(
+				path,
+				ctx,
+				query,
+				queryOptions.QueryParameters,
+				operationContext,
+				queryOptions,
+				nil)
+
+			if err != nil {
+				return QueryContainersResponse{}, err
+			}
+
+			return newContainersQueryResponse(azResponse)
+		},
+	})
+}
+
 // Read obtains the information for a Cosmos database.
 // ctx - The context for the request.
 // o - Options for Read operation.
diff --git a/sdk/data/azcosmos/cosmos_database_request_options.go b/sdk/data/azcosmos/cosmos_database_request_options.go
index d5dd5c3d6e..7c2491a125 100644
--- a/sdk/data/azcosmos/cosmos_database_request_options.go
+++ b/sdk/data/azcosmos/cosmos_database_request_options.go
@@ -3,7 +3,9 @@
 
 package azcosmos
 
-import "github.com/Azure/azure-sdk-for-go/sdk/azcore"
+import (
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
+)
 
 // ReadDatabaseOptions includes options ReadDatabase operation.
 type ReadDatabaseOptions struct {
@@ -52,3 +54,24 @@ type CreateDatabaseOptions struct {
 	// ThroughputProperties: Optional throughput configuration of the database
 	ThroughputProperties *ThroughputProperties
 }
+
+// QueryDatabasesOptions are options to query databases
+type QueryDatabasesOptions struct {
+	// ContinuationToken to be used to continue a previous query execution.
+	// Obtained from QueryItemsResponse.ContinuationToken.
+	ContinuationToken string
+
+	// QueryParameters allows execution of parametrized queries.
+	// See https://docs.microsoft.com/azure/cosmos-db/sql/sql-query-parameterized-queries
+	QueryParameters []QueryParameter
+}
+
+func (options *QueryDatabasesOptions) toHeaders() *map[string]string {
+	headers := make(map[string]string)
+
+	if options.ContinuationToken != "" {
+		headers[cosmosHeaderContinuationToken] = options.ContinuationToken
+	}
+
+	return &headers
+}
diff --git a/sdk/data/azcosmos/cosmos_query_response.go b/sdk/data/azcosmos/cosmos_query_response.go
index 529fe8c6d1..2cd4e16362 100644
--- a/sdk/data/azcosmos/cosmos_query_response.go
+++ b/sdk/data/azcosmos/cosmos_query_response.go
@@ -10,7 +10,7 @@ import (
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/runtime"
 )
 
-// QueryItemsResponse contains response from the query operation.
+// QueryItemsResponse contains response from the item query operation.
 type QueryItemsResponse struct {
 	Response
 	// ContinuationToken contains the value of the x-ms-continuation header in the response.
@@ -60,3 +60,65 @@ func newQueryResponse(resp *http.Response) (QueryItemsResponse, error) {
 type queryServiceResponse struct {
 	Documents []map[string]interface{} `json:"Documents,omitempty"`
 }
+
+// QueryContainersResponse contains response from the container query operation.
+type QueryContainersResponse struct {
+	Response
+	// ContinuationToken contains the value of the x-ms-continuation header in the response.
+	// It can be used to stop a query and resume it later.
+	ContinuationToken string
+	// List of containers.
+	Containers []ContainerProperties
+}
+
+func newContainersQueryResponse(resp *http.Response) (QueryContainersResponse, error) {
+	response := QueryContainersResponse{
+		Response: newResponse(resp),
+	}
+
+	response.ContinuationToken = resp.Header.Get(cosmosHeaderContinuationToken)
+
+	result := queryContainersServiceResponse{}
+	if err := runtime.UnmarshalAsJSON(resp, &result); err != nil {
+		return QueryContainersResponse{}, err
+	}
+
+	response.Containers = result.Containers
+
+	return response, nil
+}
+
+type queryContainersServiceResponse struct {
+	Containers []ContainerProperties `json:"DocumentCollections,omitempty"`
+}
+
+// QueryDatabasesResponse contains response from the database query operation.
+type QueryDatabasesResponse struct {
+	Response
+	// ContinuationToken contains the value of the x-ms-continuation header in the response.
+	// It can be used to stop a query and resume it later.
+	ContinuationToken string
+	// List of databases.
+	Databases []DatabaseProperties
+}
+
+func newDatabasesQueryResponse(resp *http.Response) (QueryDatabasesResponse, error) {
+	response := QueryDatabasesResponse{
+		Response: newResponse(resp),
+	}
+
+	response.ContinuationToken = resp.Header.Get(cosmosHeaderContinuationToken)
+
+	result := queryDatabasesServiceResponse{}
+	if err := runtime.UnmarshalAsJSON(resp, &result); err != nil {
+		return QueryDatabasesResponse{}, err
+	}
+
+	response.Databases = result.Databases
+
+	return response, nil
+}
+
+type queryDatabasesServiceResponse struct {
+	Databases []DatabaseProperties `json:"Databases,omitempty"`
+}


Test patch:
diff --git a/eng/common/pipelines/templates/steps/run-pester-tests.yml b/eng/common/pipelines/templates/steps/run-pester-tests.yml
new file mode 100644
index 0000000000..02082a6470
--- /dev/null
+++ b/eng/common/pipelines/templates/steps/run-pester-tests.yml
@@ -0,0 +1,55 @@
+parameters:
+  - name: TargetDirectory
+    type: string
+  - name: CustomTestSteps
+    type: object
+    default: []
+  - name: EnvVars
+    type: object
+    default: {}
+  - name: TargetTags
+    type: string
+    default: ''
+
+steps:
+  - pwsh: |
+     Install-Module -Name Pester -Force
+    displayName: Install Pester
+
+   # default test steps
+  - ${{ if eq(length(parameters.CustomTestSteps), 0) }}:
+    - pwsh: |
+        $tags = "${{ parameters.TargetTags }}" -Split "," | ForEach-Object { return $_.Trim() }
+
+        $config = New-PesterConfiguration
+        $config.CodeCoverage.Enabled = $true
+        $config.TestResult.Enabled = $true
+        $config.Run.PassThru = $true
+
+        if ($tags) {
+          $config.Filter.Tag = $tags
+        }
+
+        Invoke-Pester -Configuration $config
+      displayName: Run Tests
+      env: ${{ parameters.EnvVars }}
+      workingDirectory: $(Build.SourcesDirectory)/${{ parameters.TargetDirectory }}
+
+  - ${{ if not(eq(length(parameters.CustomTestSteps), 0)) }}:
+    - ${{ parameters.CustomTestSteps }}
+
+  - task: PublishTestResults@2
+    displayName: 'Publish Test Results'
+    condition: succeededOrFailed()
+    inputs:
+      testResultsFormat: 'NUnit'
+      testResultsFiles: $(Build.SourcesDirectory)/${{ parameters.TargetDirectory }}/testResults.xml
+      testRunTitle: '$(System.StageName)_$(Agent.JobName)_Tests'
+
+  - task: PublishCodeCoverageResults@1
+    displayName: 'Publish Code Coverage to Azure DevOps'
+    condition: succeededOrFailed()
+    inputs:
+      codeCoverageTool: 'JaCoCo'
+      summaryFileLocation: '$(Build.SourcesDirectory)/${{ parameters.TargetDirectory }}/coverage.xml'
+      pathToSources: '$(Build.SourcesDirectory)/${{ parameters.TargetDirectory }}'
\ No newline at end of file
diff --git a/sdk/data/azcosmos/cosmos_client_test.go b/sdk/data/azcosmos/cosmos_client_test.go
index 786f38b9a3..f33e730722 100644
--- a/sdk/data/azcosmos/cosmos_client_test.go
+++ b/sdk/data/azcosmos/cosmos_client_test.go
@@ -9,6 +9,7 @@ import (
 	"io"
 	"net/http"
 	"net/url"
+	"strconv"
 	"testing"
 
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
@@ -556,6 +557,97 @@ func TestCreateScopeFromEndpoint(t *testing.T) {
 	}
 }
 
+func TestQueryDatabases(t *testing.T) {
+	jsonStringpage1 := []byte(`{"Databases":[{"id":"doc1"},{"id":"doc2"}]}`)
+	jsonStringpage2 := []byte(`{"Databases":[{"id":"doc3"},{"id":"doc4"},{"id":"doc5"}]}`)
+
+	srv, close := mock.NewTLSServer()
+	defer close()
+	srv.AppendResponse(
+		mock.WithBody(jsonStringpage1),
+		mock.WithHeader(cosmosHeaderEtag, "someEtag"),
+		mock.WithHeader(cosmosHeaderActivityId, "someActivityId"),
+		mock.WithHeader(cosmosHeaderRequestCharge, "13.42"),
+		mock.WithHeader(cosmosHeaderContinuationToken, "someContinuationToken"),
+		mock.WithStatusCode(200))
+	srv.AppendResponse(
+		mock.WithBody(jsonStringpage2),
+		mock.WithHeader(cosmosHeaderEtag, "someEtag"),
+		mock.WithHeader(cosmosHeaderActivityId, "someActivityId"),
+		mock.WithHeader(cosmosHeaderRequestCharge, "13.42"),
+		mock.WithStatusCode(200))
+
+	verifier := pipelineVerifier{}
+
+	pl := azruntime.NewPipeline("azcosmostest", "v1.0.0", azruntime.PipelineOptions{PerCall: []policy.Policy{&verifier}}, &policy.ClientOptions{Transport: srv})
+	client := &Client{endpoint: srv.URL(), pipeline: pl}
+
+	receivedIds := []string{}
+	queryPager := client.NewQueryDatabasesPager("select * from c", nil)
+	for queryPager.More() {
+		queryResponse, err := queryPager.NextPage(context.TODO())
+		if err != nil {
+			t.Fatalf("Failed to query items: %v", err)
+		}
+
+		for _, dbs := range queryResponse.Databases {
+			receivedIds = append(receivedIds, dbs.ID)
+		}
+
+		if queryPager.More() && queryResponse.ContinuationToken != "someContinuationToken" {
+			t.Errorf("Expected ContinuationToken to be %s, but got %s", "someContinuationToken", queryResponse.ContinuationToken)
+		}
+
+		if queryResponse.ActivityID == "" {
+			t.Fatal("Activity id was not returned")
+		}
+
+		if queryResponse.RequestCharge == 0 {
+			t.Fatal("Request charge was not returned")
+		}
+	}
+
+	for i := 0; i < 5; i++ {
+		if receivedIds[i] != "doc"+strconv.Itoa(i+1) {
+			t.Fatalf("Expected id %d, got %s", i, receivedIds[i])
+		}
+	}
+
+	if len(verifier.requests) != 2 {
+		t.Fatalf("Expected 2 requests, got %d", len(verifier.requests))
+	}
+
+	for index, request := range verifier.requests {
+		if request.method != http.MethodPost {
+			t.Errorf("Expected method to be %s, but got %s", http.MethodPost, request.method)
+		}
+
+		if request.url.RequestURI() != "/dbs" {
+			t.Errorf("Expected url to be %s, but got %s", "/dbs", request.url.RequestURI())
+		}
+
+		if !request.isQuery {
+			t.Errorf("Expected request to be a query, but it was not")
+		}
+
+		if request.body != "{\"query\":\"select * from c\"}" {
+			t.Errorf("Expected %v, but got %v", "{\"query\":\"select * from c\"}", request.body)
+		}
+
+		if request.contentType != cosmosHeaderValuesQuery {
+			t.Errorf("Expected %v, but got %v", cosmosHeaderValuesQuery, request.contentType)
+		}
+
+		if index == 0 && request.headers.Get(cosmosHeaderContinuationToken) != "" {
+			t.Errorf("Expected ContinuationToken to be %s, but got %s", "", request.headers.Get(cosmosHeaderContinuationToken))
+		}
+
+		if index == 1 && request.headers.Get(cosmosHeaderContinuationToken) != "someContinuationToken" {
+			t.Errorf("Expected ContinuationToken to be %s, but got %s", "someContinuationToken", request.headers.Get(cosmosHeaderContinuationToken))
+		}
+	}
+}
+
 type pipelineVerifier struct {
 	requests []pipelineVerifierRequest
 }
diff --git a/sdk/data/azcosmos/cosmos_container_request_options_test.go b/sdk/data/azcosmos/cosmos_container_request_options_test.go
index df667d0ad5..2575c40150 100644
--- a/sdk/data/azcosmos/cosmos_container_request_options_test.go
+++ b/sdk/data/azcosmos/cosmos_container_request_options_test.go
@@ -39,3 +39,17 @@ func TestContainerRequestOptionsToHeaders(t *testing.T) {
 		t.Error("toHeaders should return nil")
 	}
 }
+
+func TestQueryContainersRequestOptionsToHeaders(t *testing.T) {
+	options := &QueryContainersOptions{}
+	options.ContinuationToken = "continuationToken"
+	header := options.toHeaders()
+	if header == nil {
+		t.Fatal("toHeaders should return non-nil")
+	}
+
+	headers := *header
+	if headers[cosmosHeaderContinuationToken] != "continuationToken" {
+		t.Errorf("ContinuationToken should be continuationToken but got %v", headers[cosmosHeaderContinuationToken])
+	}
+}
diff --git a/sdk/data/azcosmos/cosmos_database_request_options_test.go b/sdk/data/azcosmos/cosmos_database_request_options_test.go
index ad6b9868f0..fd030b6473 100644
--- a/sdk/data/azcosmos/cosmos_database_request_options_test.go
+++ b/sdk/data/azcosmos/cosmos_database_request_options_test.go
@@ -56,3 +56,17 @@ func TestDeleteDatabaseOptionsToHeaders(t *testing.T) {
 		t.Errorf("IfNoneMatchEtag not set matching expected %v got %v", options.IfNoneMatchEtag, headers[headerIfNoneMatch])
 	}
 }
+
+func TestQueryDatabasesRequestOptionsToHeaders(t *testing.T) {
+	options := &QueryDatabasesOptions{}
+	options.ContinuationToken = "continuationToken"
+	header := options.toHeaders()
+	if header == nil {
+		t.Fatal("toHeaders should return non-nil")
+	}
+
+	headers := *header
+	if headers[cosmosHeaderContinuationToken] != "continuationToken" {
+		t.Errorf("ContinuationToken should be continuationToken but got %v", headers[cosmosHeaderContinuationToken])
+	}
+}
diff --git a/sdk/data/azcosmos/cosmos_database_test.go b/sdk/data/azcosmos/cosmos_database_test.go
new file mode 100644
index 0000000000..d4a199306e
--- /dev/null
+++ b/sdk/data/azcosmos/cosmos_database_test.go
@@ -0,0 +1,108 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package azcosmos
+
+import (
+	"context"
+	"net/http"
+	"strconv"
+	"testing"
+
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
+	azruntime "github.com/Azure/azure-sdk-for-go/sdk/azcore/runtime"
+	"github.com/Azure/azure-sdk-for-go/sdk/internal/mock"
+)
+
+func TestDatabaseQueryContainers(t *testing.T) {
+	jsonStringpage1 := []byte(`{"DocumentCollections":[{"id":"doc1"},{"id":"doc2"}]}`)
+	jsonStringpage2 := []byte(`{"DocumentCollections":[{"id":"doc3"},{"id":"doc4"},{"id":"doc5"}]}`)
+
+	srv, close := mock.NewTLSServer()
+	defer close()
+	srv.AppendResponse(
+		mock.WithBody(jsonStringpage1),
+		mock.WithHeader(cosmosHeaderEtag, "someEtag"),
+		mock.WithHeader(cosmosHeaderActivityId, "someActivityId"),
+		mock.WithHeader(cosmosHeaderRequestCharge, "13.42"),
+		mock.WithHeader(cosmosHeaderContinuationToken, "someContinuationToken"),
+		mock.WithStatusCode(200))
+	srv.AppendResponse(
+		mock.WithBody(jsonStringpage2),
+		mock.WithHeader(cosmosHeaderEtag, "someEtag"),
+		mock.WithHeader(cosmosHeaderActivityId, "someActivityId"),
+		mock.WithHeader(cosmosHeaderRequestCharge, "13.42"),
+		mock.WithStatusCode(200))
+
+	verifier := pipelineVerifier{}
+
+	pl := azruntime.NewPipeline("azcosmostest", "v1.0.0", azruntime.PipelineOptions{PerCall: []policy.Policy{&verifier}}, &policy.ClientOptions{Transport: srv})
+	client := &Client{endpoint: srv.URL(), pipeline: pl}
+
+	database, _ := newDatabase("databaseId", client)
+
+	receivedIds := []string{}
+	queryPager := database.NewQueryContainersPager("select * from c", nil)
+	for queryPager.More() {
+		queryResponse, err := queryPager.NextPage(context.TODO())
+		if err != nil {
+			t.Fatalf("Failed to query items: %v", err)
+		}
+
+		for _, container := range queryResponse.Containers {
+			receivedIds = append(receivedIds, container.ID)
+		}
+
+		if queryPager.More() && queryResponse.ContinuationToken != "someContinuationToken" {
+			t.Errorf("Expected ContinuationToken to be %s, but got %s", "someContinuationToken", queryResponse.ContinuationToken)
+		}
+
+		if queryResponse.ActivityID == "" {
+			t.Fatal("Activity id was not returned")
+		}
+
+		if queryResponse.RequestCharge == 0 {
+			t.Fatal("Request charge was not returned")
+		}
+	}
+
+	for i := 0; i < 5; i++ {
+		if receivedIds[i] != "doc"+strconv.Itoa(i+1) {
+			t.Fatalf("Expected id %d, got %s", i, receivedIds[i])
+		}
+	}
+
+	if len(verifier.requests) != 2 {
+		t.Fatalf("Expected 2 requests, got %d", len(verifier.requests))
+	}
+
+	for index, request := range verifier.requests {
+		if request.method != http.MethodPost {
+			t.Errorf("Expected method to be %s, but got %s", http.MethodPost, request.method)
+		}
+
+		if request.url.RequestURI() != "/dbs/databaseId/colls" {
+			t.Errorf("Expected url to be %s, but got %s", "/dbs/databaseId/colls", request.url.RequestURI())
+		}
+
+		if !request.isQuery {
+			t.Errorf("Expected request to be a query, but it was not")
+		}
+
+		if request.body != "{\"query\":\"select * from c\"}" {
+			t.Errorf("Expected %v, but got %v", "{\"query\":\"select * from c\"}", request.body)
+		}
+
+		if request.contentType != cosmosHeaderValuesQuery {
+			t.Errorf("Expected %v, but got %v", cosmosHeaderValuesQuery, request.contentType)
+		}
+
+		if index == 0 && request.headers.Get(cosmosHeaderContinuationToken) != "" {
+			t.Errorf("Expected ContinuationToken to be %s, but got %s", "", request.headers.Get(cosmosHeaderContinuationToken))
+		}
+
+		if index == 1 && request.headers.Get(cosmosHeaderContinuationToken) != "someContinuationToken" {
+			t.Errorf("Expected ContinuationToken to be %s, but got %s", "someContinuationToken", request.headers.Get(cosmosHeaderContinuationToken))
+		}
+	}
+}
diff --git a/sdk/data/azcosmos/cosmos_query_response_test.go b/sdk/data/azcosmos/cosmos_query_response_test.go
index a2424a17ea..1bb0e95900 100644
--- a/sdk/data/azcosmos/cosmos_query_response_test.go
+++ b/sdk/data/azcosmos/cosmos_query_response_test.go
@@ -94,3 +94,131 @@ func TestQueryResponseParsing(t *testing.T) {
 		}
 	}
 }
+
+func TestQueryContainersResponseParsing(t *testing.T) {
+	queryResponseRaw := map[string][]map[string]string{
+		"DocumentCollections": {
+			{"id": "id1"},
+			{"id": "id2"},
+		},
+	}
+
+	jsonString, err := json.Marshal(queryResponseRaw)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	srv, close := mock.NewTLSServer()
+	defer close()
+	srv.SetResponse(
+		mock.WithBody(jsonString),
+		mock.WithHeader(cosmosHeaderEtag, "someEtag"),
+		mock.WithHeader(cosmosHeaderActivityId, "someActivityId"),
+		mock.WithHeader(cosmosHeaderRequestCharge, "13.42"))
+
+	req, err := azruntime.NewRequest(context.Background(), http.MethodGet, srv.URL())
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	pl := azruntime.NewPipeline("azcosmostest", "v1.0.0", azruntime.PipelineOptions{}, &policy.ClientOptions{Transport: srv})
+	resp, _ := pl.Do(req)
+	parsedResponse, err := newContainersQueryResponse(resp)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if parsedResponse.RawResponse == nil {
+		t.Fatal("parsedResponse.RawResponse is nil")
+	}
+
+	if parsedResponse.ActivityID != "someActivityId" {
+		t.Errorf("Expected ActivityId to be %s, but got %s", "someActivityId", parsedResponse.ActivityID)
+	}
+
+	if parsedResponse.RequestCharge != 13.42 {
+		t.Errorf("Expected RequestCharge to be %f, but got %f", 13.42, parsedResponse.RequestCharge)
+	}
+
+	if parsedResponse.ETag != "someEtag" {
+		t.Errorf("Expected ETag to be %s, but got %s", "someEtag", parsedResponse.ETag)
+	}
+
+	if len(parsedResponse.Containers) != 2 {
+		t.Errorf("Expected 2 containers, but got %d", len(parsedResponse.Containers))
+	}
+
+	for index, container := range parsedResponse.Containers {
+		if err != nil {
+			t.Fatalf("Failed to unmarshal item response: %v", err)
+		}
+
+		if container.ID != ("id" + strconv.Itoa(index+1)) {
+			t.Errorf("Expected id to be %s, but got %s", "id"+strconv.Itoa(index+1), container.ID)
+		}
+	}
+}
+
+func TestQueryDatabasesResponseParsing(t *testing.T) {
+	queryResponseRaw := map[string][]map[string]string{
+		"Databases": {
+			{"id": "id1"},
+			{"id": "id2"},
+		},
+	}
+
+	jsonString, err := json.Marshal(queryResponseRaw)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	srv, close := mock.NewTLSServer()
+	defer close()
+	srv.SetResponse(
+		mock.WithBody(jsonString),
+		mock.WithHeader(cosmosHeaderEtag, "someEtag"),
+		mock.WithHeader(cosmosHeaderActivityId, "someActivityId"),
+		mock.WithHeader(cosmosHeaderRequestCharge, "13.42"))
+
+	req, err := azruntime.NewRequest(context.Background(), http.MethodGet, srv.URL())
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	pl := azruntime.NewPipeline("azcosmostest", "v1.0.0", azruntime.PipelineOptions{}, &policy.ClientOptions{Transport: srv})
+	resp, _ := pl.Do(req)
+	parsedResponse, err := newDatabasesQueryResponse(resp)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if parsedResponse.RawResponse == nil {
+		t.Fatal("parsedResponse.RawResponse is nil")
+	}
+
+	if parsedResponse.ActivityID != "someActivityId" {
+		t.Errorf("Expected ActivityId to be %s, but got %s", "someActivityId", parsedResponse.ActivityID)
+	}
+
+	if parsedResponse.RequestCharge != 13.42 {
+		t.Errorf("Expected RequestCharge to be %f, but got %f", 13.42, parsedResponse.RequestCharge)
+	}
+
+	if parsedResponse.ETag != "someEtag" {
+		t.Errorf("Expected ETag to be %s, but got %s", "someEtag", parsedResponse.ETag)
+	}
+
+	if len(parsedResponse.Databases) != 2 {
+		t.Errorf("Expected 2 containers, but got %d", len(parsedResponse.Databases))
+	}
+
+	for index, db := range parsedResponse.Databases {
+		if err != nil {
+			t.Fatalf("Failed to unmarshal item response: %v", err)
+		}
+
+		if db.ID != ("id" + strconv.Itoa(index+1)) {
+			t.Errorf("Expected id to be %s, but got %s", "id"+strconv.Itoa(index+1), db.ID)
+		}
+	}
+}
diff --git a/sdk/data/azcosmos/emulator_cosmos_container_test.go b/sdk/data/azcosmos/emulator_cosmos_container_test.go
index bf314f3b66..10924b85c9 100644
--- a/sdk/data/azcosmos/emulator_cosmos_container_test.go
+++ b/sdk/data/azcosmos/emulator_cosmos_container_test.go
@@ -52,6 +52,28 @@ func TestContainerCRUD(t *testing.T) {
 		t.Fatalf("Failed to read container: %v", err)
 	}
 
+	receivedIds := []string{}
+	opt := QueryContainersOptions{
+		QueryParameters: []QueryParameter{
+			{"@id", "aContainer"},
+		},
+	}
+	queryPager := database.NewQueryContainersPager("SELECT * FROM root r WHERE r.id = @id", &opt)
+	for queryPager.More() {
+		queryResponse, err := queryPager.NextPage(context.TODO())
+		if err != nil {
+			t.Fatalf("Failed to query databases: %v", err)
+		}
+
+		for _, db := range queryResponse.Containers {
+			receivedIds = append(receivedIds, db.ID)
+		}
+	}
+
+	if len(receivedIds) != 1 {
+		t.Fatalf("Expected 1 container, got %d", len(receivedIds))
+	}
+
 	updatedProperties := ContainerProperties{
 		ID: "aContainer",
 		PartitionKeyDefinition: PartitionKeyDefinition{
diff --git a/sdk/data/azcosmos/emulator_cosmos_database_test.go b/sdk/data/azcosmos/emulator_cosmos_database_test.go
index a517109cbe..6559713003 100644
--- a/sdk/data/azcosmos/emulator_cosmos_database_test.go
+++ b/sdk/data/azcosmos/emulator_cosmos_database_test.go
@@ -33,6 +33,28 @@ func TestDatabaseCRUD(t *testing.T) {
 		t.Errorf("Unexpected id match: %v", resp.DatabaseProperties)
 	}
 
+	receivedIds := []string{}
+	opt := QueryDatabasesOptions{
+		QueryParameters: []QueryParameter{
+			{"@id", "baseDbTest"},
+		},
+	}
+	queryPager := client.NewQueryDatabasesPager("SELECT * FROM root r WHERE r.id = @id", &opt)
+	for queryPager.More() {
+		queryResponse, err := queryPager.NextPage(context.TODO())
+		if err != nil {
+			t.Fatalf("Failed to query databases: %v", err)
+		}
+
+		for _, db := range queryResponse.Databases {
+			receivedIds = append(receivedIds, db.ID)
+		}
+	}
+
+	if len(receivedIds) != 1 {
+		t.Fatalf("Expected 1 database, got %d", len(receivedIds))
+	}
+
 	throughputResponse, err := db.ReadThroughput(context.TODO(), nil)
 	if err == nil {
 		t.Fatalf("Expected not finding throughput but instead got : %v", throughputResponse)
diff --git a/sdk/data/azcosmos/example_test.go b/sdk/data/azcosmos/example_test.go
index d8c4484fc6..c91c88d293 100644
--- a/sdk/data/azcosmos/example_test.go
+++ b/sdk/data/azcosmos/example_test.go
@@ -110,6 +110,44 @@ func ExampleClient_CreateDatabase() {
 	fmt.Printf("Database created. ActivityId %s", databaseResponse.ActivityID)
 }
 
+func ExampleClient_NewQueryDatabasesPager() {
+	endpoint, ok := os.LookupEnv("AZURE_COSMOS_ENDPOINT")
+	if !ok {
+		panic("AZURE_COSMOS_ENDPOINT could not be found")
+	}
+
+	key, ok := os.LookupEnv("AZURE_COSMOS_KEY")
+	if !ok {
+		panic("AZURE_COSMOS_KEY could not be found")
+	}
+
+	cred, err := azcosmos.NewKeyCredential(key)
+	if err != nil {
+		panic(err)
+	}
+
+	client, err := azcosmos.NewClientWithKey(endpoint, cred, nil)
+	if err != nil {
+		panic(err)
+	}
+
+	queryPager := client.NewQueryDatabasesPager("select * from dbs d", nil)
+	for queryPager.More() {
+		queryResponse, err := queryPager.NextPage(context.Background())
+		if err != nil {
+			var responseErr *azcore.ResponseError
+			errors.As(err, &responseErr)
+			panic(responseErr)
+		}
+
+		for _, container := range queryResponse.Databases {
+			fmt.Printf("Received database %s", container.ID)
+		}
+
+		fmt.Printf("Query page received with %v databases. ActivityId %s consuming %v RU", len(queryResponse.Databases), queryResponse.ActivityID, queryResponse.RequestCharge)
+	}
+}
+
 func ExampleDatabaseClient_CreateContainer() {
 	endpoint, ok := os.LookupEnv("AZURE_COSMOS_ENDPOINT")
 	if !ok {
@@ -154,6 +192,48 @@ func ExampleDatabaseClient_CreateContainer() {
 
 	fmt.Printf("Container created. ActivityId %s", resp.ActivityID)
 }
+func ExampleDatabaseClient_NewQueryContainersPager() {
+	endpoint, ok := os.LookupEnv("AZURE_COSMOS_ENDPOINT")
+	if !ok {
+		panic("AZURE_COSMOS_ENDPOINT could not be found")
+	}
+
+	key, ok := os.LookupEnv("AZURE_COSMOS_KEY")
+	if !ok {
+		panic("AZURE_COSMOS_KEY could not be found")
+	}
+
+	cred, err := azcosmos.NewKeyCredential(key)
+	if err != nil {
+		panic(err)
+	}
+
+	client, err := azcosmos.NewClientWithKey(endpoint, cred, nil)
+	if err != nil {
+		panic(err)
+	}
+
+	database, err := client.NewDatabase("databaseName")
+	if err != nil {
+		panic(err)
+	}
+
+	queryPager := database.NewQueryContainersPager("select * from containers c", nil)
+	for queryPager.More() {
+		queryResponse, err := queryPager.NextPage(context.Background())
+		if err != nil {
+			var responseErr *azcore.ResponseError
+			errors.As(err, &responseErr)
+			panic(responseErr)
+		}
+
+		for _, container := range queryResponse.Containers {
+			fmt.Printf("Received container %s", container.ID)
+		}
+
+		fmt.Printf("Query page received with %v containers. ActivityId %s consuming %v RU", len(queryResponse.Containers), queryResponse.ActivityID, queryResponse.RequestCharge)
+	}
+}
 
 func ExampleContainerClient_ReplaceThroughput() {
 	endpoint, ok := os.LookupEnv("AZURE_COSMOS_ENDPOINT")