[query/azlogs] add batch queries (Azure#23109)

* add batch

* docs

Author: gracewilcox

sdk/monitor/query/azlogs/CHANGELOG.md

@@ -1,8 +1,9 @@
 # Release History
 
-## 1.0.1 (Unreleased)
+## 1.1.0-beta.1 (Unreleased)
 
 ### Features Added
+* Added batch query functionality
 
 ### Breaking Changes
 

sdk/monitor/query/azlogs/assets.json

@@ -2,5 +2,5 @@
   "AssetsRepo": "Azure/azure-sdk-assets",
   "AssetsRepoPrefixPath": "go",
   "TagPrefix": "go/monitor/query/azlogs",
-  "Tag": "go/monitor/query/azlogs_139fabf226"
+  "Tag": "go/monitor/query/azlogs_2bd6ffba1a"
 }

sdk/monitor/query/azlogs/autorest.md

@@ -26,15 +26,10 @@ directive:
     where: $["x-ms-paths"]
     transform: >
         delete $["/{resourceId}/query?disambiguation_dummy"];
-  - from: swagger-document
-    where: $["paths"]
-    transform: >
-        delete $["/$batch"];
 
   # delete extra operations
   - remove-operation: Query_Get
   - remove-operation: Query_ResourceGet
-  - remove-operation: Query_Batch
 
   # delete metadata and batch models
   - remove-model: metadataResults
@@ -48,11 +43,6 @@ directive:
   - remove-model: metadataWorkspace
   - remove-model: metadataResource
   - remove-model: metadataPermissions
-  - remove-model: batchRequest
-  - remove-model: batchQueryRequest
-  - remove-model: batchResponse
-  - remove-model: batchQueryResponse
-  - remove-model: batchQueryResults
 
  # rename log operations to generate into a separate logs client
   - rename-operation:
@@ -61,6 +51,9 @@ directive:
   - rename-operation:
       from: Query_ResourceExecute
       to: Logs_QueryResource
+  - rename-operation:
+      from: Query_Batch
+      to: Logs_QueryBatch
 
   # rename Body.Workspaces to Body.AdditionalWorkspaces
   - from: swagger-document
@@ -69,13 +62,18 @@ directive:
 
   # rename Render to Visualization
   - from: swagger-document
-    where: $.definitions.queryResults.properties.render
+    where: $.definitions..render
     transform: $["x-ms-client-name"] = "Visualization"
 
   # rename LogsColumnType to ColumnType
   - from: swagger-document
     where: $.definitions.logsColumnType.x-ms-enum
     transform: $["name"] = "ColumnType"
+
+   # rename BatchQueryRequest.Workspace to BatchQueryRequest.WorkspaceID
+  - from: swagger-document
+    where: $.definitions.batchQueryRequest.properties.workspace
+    transform: $["x-ms-client-name"] = "WorkspaceID"
 
   # rename Prefer to Options
   - from: swagger-document

sdk/monitor/query/azlogs/client.go

@@ -205,6 +205,103 @@ func TestQueryResource_Advanced(t *testing.T) {
 	testSerde(t, &res)
 }
 
+func TestQueryBatch_QuerySuccess(t *testing.T) {
+	client := startTest(t)
+	query1, query2 := query, query+" | take 2"
+	timespan := azlogs.NewTimeInterval(time.Date(2022, 3, 2, 0, 0, 0, 0, time.UTC), time.Date(2022, 3, 3, 0, 0, 0, 0, time.UTC))
+
+	batchRequest := azlogs.BatchRequest{[]azlogs.BatchQueryRequest{
+		{Body: &azlogs.QueryBody{Query: to.Ptr(query1), Timespan: to.Ptr(timespan)}, ID: to.Ptr("1"), WorkspaceID: to.Ptr(workspaceID)},
+		{Body: &azlogs.QueryBody{Query: to.Ptr(query2), Timespan: to.Ptr(timespan)}, ID: to.Ptr("2"), WorkspaceID: to.Ptr(workspaceID)},
+	}}
+	testSerde(t, &batchRequest)
+
+	res, err := client.QueryBatch(context.Background(), batchRequest, nil)
+	require.NoError(t, err)
+	require.Len(t, res.Responses, 2)
+	for _, resp := range res.Responses {
+		require.Nil(t, resp.Body.Error)
+		require.NotNil(t, resp.Body.Tables)
+		if *resp.ID == "1" && len(resp.Body.Tables[0].Rows) != 100 {
+			t.Fatal("expected 100 rows from batch request 1")
+		}
+		if *resp.ID == "2" && len(resp.Body.Tables[0].Rows) != 2 {
+			t.Fatal("expected 2 rows from batch request 2")
+		}
+	}
+	testSerde(t, &res)
+}
+
+func TestQueryBatch_BasicQueryFailure(t *testing.T) {
+	client := startTest(t)
+
+	batchRequest := azlogs.BatchRequest{[]azlogs.BatchQueryRequest{
+		{Body: &azlogs.QueryBody{Query: to.Ptr(query)}, ID: to.Ptr("1"), WorkspaceID: to.Ptr(workspaceID)},
+		{Body: &azlogs.QueryBody{Query: to.Ptr(query)}, ID: to.Ptr("1"), WorkspaceID: to.Ptr(workspaceID)},
+	}}
+	testSerde(t, &batchRequest)
+
+	res, err := client.QueryBatch(context.Background(), batchRequest, nil)
+	require.Error(t, err)
+	require.Nil(t, res.Responses)
+
+	var httpErr *azcore.ResponseError
+	require.ErrorAs(t, err, &httpErr)
+	require.Equal(t, httpErr.ErrorCode, "BadArgumentError")
+	require.Equal(t, httpErr.StatusCode, 400)
+}
+
+func TestQueryBatch_AdvancedQuerySuccess(t *testing.T) {
+	client := startTest(t)
+	timespan := azlogs.NewTimeInterval(time.Date(2022, 12, 25, 0, 0, 0, 0, time.UTC), time.Date(2022, 12, 25, 12, 0, 0, 0, time.UTC))
+	batchPrefer1 := "wait=600,include-statistics=true,include-render=true"
+	headers1 := map[string]*string{"prefer": &batchPrefer1}
+	batchPrefer2 := "wait=180,include-statistics=true,include-render=true"
+	headers2 := map[string]*string{"prefer": &batchPrefer2}
+
+	batchRequestAdvanced := azlogs.BatchRequest{[]azlogs.BatchQueryRequest{
+		{Body: &azlogs.QueryBody{Query: to.Ptr(query), Timespan: to.Ptr(timespan)}, ID: to.Ptr("1"), WorkspaceID: to.Ptr(workspaceID2), Headers: headers1},
+		{Body: &azlogs.QueryBody{Query: to.Ptr(query), Timespan: to.Ptr(timespan)}, ID: to.Ptr("2"), WorkspaceID: to.Ptr(workspaceID2), Headers: headers2},
+	}}
+	testSerde(t, &batchRequestAdvanced)
+
+	res, err := client.QueryBatch(context.Background(), batchRequestAdvanced, nil)
+	require.NoError(t, err)
+	require.Len(t, res.Responses, 2)
+	for _, resp := range res.Responses {
+		require.Nil(t, resp.Body.Error)
+		require.NotNil(t, resp.Body.Tables)
+		require.NotNil(t, resp.Body.Visualization)
+		require.NotNil(t, resp.Body.Statistics)
+		require.Len(t, resp.Body.Tables[0].Rows, 100)
+	}
+	testSerde(t, &res)
+}
+
+func TestQueryBatch_PartialError(t *testing.T) {
+	client := startTest(t)
+
+	batchRequest := azlogs.BatchRequest{[]azlogs.BatchQueryRequest{
+		{Body: &azlogs.QueryBody{Query: to.Ptr("not a valid query")}, ID: to.Ptr("1"), WorkspaceID: to.Ptr(workspaceID)},
+		{Body: &azlogs.QueryBody{Query: to.Ptr(query)}, ID: to.Ptr("2"), WorkspaceID: to.Ptr(workspaceID)},
+	}}
+
+	res, err := client.QueryBatch(context.Background(), batchRequest, nil)
+	require.NoError(t, err)
+	require.Len(t, res.Responses, 2)
+	for _, resp := range res.Responses {
+		if *resp.ID == "1" {
+			require.NotNil(t, resp.Body.Error)
+			require.Equal(t, resp.Body.Error.Code, "BadArgumentError")
+			require.Contains(t, resp.Body.Error.Error(), "BadArgumentError")
+		}
+		if *resp.ID == "2" {
+			require.Nil(t, resp.Body.Error)
+			require.Len(t, resp.Body.Tables[0].Rows, 100)
+		}
+	}
+}
+
 func TestTimeInterval(t *testing.T) {
 	timespan := azlogs.NewTimeInterval(time.Date(2022, 3, 2, 1, 2, 3, 0, time.UTC), time.Date(2022, 3, 3, 0, 0, 0, 0, time.UTC))
 	require.Equal(t, timespan, azlogs.TimeInterval("2022-03-02T01:02:03Z/2022-03-03T00:00:00Z"))

sdk/monitor/query/azlogs/client_test.go

@@ -13,7 +13,12 @@ import (
 	"github.com/Azure/azure-sdk-for-go/sdk/monitor/query/azlogs"
 )
 
-var client azlogs.Client
+var (
+	client      azlogs.Client
+	kustoQuery1 string
+	kustoQuery2 string
+	kustoQuery3 string
+)
 
 type queryResult struct {
 	Bool   bool
@@ -164,3 +169,31 @@ func ExampleClient_QueryResource() {
 		}
 	}
 }
+
+func ExampleClient_QueryBatch() {
+	// `QueryBatch` is an advanced method allowing users to execute multiple log queries in a single request.
+
+	workspaceID := "g4d1e129-fb1e-4b0a-b234-250abc987ea65" // example Azure Log Analytics Workspace ID
+	timespan := azlogs.NewTimeInterval(time.Date(2022, 12, 25, 0, 0, 0, 0, time.UTC), time.Date(2022, 12, 25, 12, 0, 0, 0, time.UTC))
+
+	batchRequest := azlogs.BatchRequest{[]azlogs.BatchQueryRequest{
+		{Body: &azlogs.QueryBody{Query: to.Ptr(kustoQuery1), Timespan: to.Ptr(timespan)}, ID: to.Ptr("1"), WorkspaceID: to.Ptr(workspaceID)},
+		{Body: &azlogs.QueryBody{Query: to.Ptr(kustoQuery2), Timespan: to.Ptr(timespan)}, ID: to.Ptr("2"), WorkspaceID: to.Ptr(workspaceID)},
+		{Body: &azlogs.QueryBody{Query: to.Ptr(kustoQuery3), Timespan: to.Ptr(timespan)}, ID: to.Ptr("3"), WorkspaceID: to.Ptr(workspaceID)},
+	}}
+
+	res, err := client.QueryBatch(context.TODO(), batchRequest, nil)
+	if err != nil {
+		//TODO: handle error
+	}
+
+	// `QueryBatch` can return results in any order, usually by time it takes each individual query to complete.
+	// Use the `ID` field to identify the correct response.
+	responses := res.Responses
+	fmt.Println("ID's of successful responses:")
+	for _, response := range responses {
+		if response.Body.Error == nil {
+			fmt.Println(*response.ID)
+		}
+	}
+}