Extending Microsoft Graph data source

docs/plugins/microsoft/data-sources/microsoft_graph.md

@@ -101,5 +101,24 @@ data microsoft_graph {
   # Optional map of string.
   # Default value:
   query_params = null
+
+  # Number of objects to be returned
+  #
+  # Optional number.
+  # Must be >= 1
+  # Default value:
+  objects_size = 50
+
+  # Return only the list of objects. If `false`, returns an object with `objects` and `totalCount` fields
+  #
+  # Optional bool.
+  # Default value:
+  only_objects = true
+
+  # If API endpoint response should be treated as a list or as an object. If set to `true`, `only_objects`, `query_params` and `objects_size` are ignored.
+  #
+  # Optional bool.
+  # Default value:
+  is_object_endpoint = false
 }
 ```

docs/plugins/plugins.json

@@ -347,6 +347,9 @@
         "arguments": [
           "api_version",
           "endpoint",
+          "is_object_endpoint",
+          "objects_size",
+          "only_objects",
           "query_params"
         ]
       },

internal/microsoft/client/graph_client.go

@@ -2,13 +2,20 @@ package client
 
 import (
 	"context"
-	"encoding/json"
 	"fmt"
+	"io"
+	"log/slog"
 	"net/http"
 	"net/url"
+	"strconv"
+
+	"github.com/blackstork-io/fabric/plugin/plugindata"
 )
 
-const graphUrl = "https://graph.microsoft.com"
+const (
+	graphUrl        = "https://graph.microsoft.com"
+	defaultPageSize = 20
+)
 
 type graphClient struct {
 	accessToken string
@@ -28,30 +35,143 @@ func (cli *graphClient) prepare(r *http.Request) {
 	r.Header.Set("Authorization", fmt.Sprintf("Bearer %s", cli.accessToken))
 }
 
-func (cli *graphClient) QueryGraph(ctx context.Context, endpoint string, queryParams url.Values) (result interface{}, err error) {
-	requestUrl, err := url.Parse(graphUrl + fmt.Sprintf("/%s%s", cli.apiVersion, endpoint))
-	if err != nil {
-		return
-	}
-	if queryParams != nil {
-		requestUrl.RawQuery = queryParams.Encode()
-	}
+func (cli *graphClient) fetchURL(ctx context.Context, requestUrl *url.URL) (result plugindata.Data, err error) {
 	r, err := http.NewRequestWithContext(ctx, http.MethodGet, requestUrl.String(), nil)
 	if err != nil {
 		return
 	}
 	cli.prepare(r)
+	slog.DebugContext(ctx, "Fetching an URL from API", "url", requestUrl.String())
 	res, err := cli.client.Do(r)
 	if err != nil {
 		return
 	}
+	raw, err := io.ReadAll(res.Body)
+	if err != nil {
+		return nil, fmt.Errorf("failed to read the results: %s", err)
+	}
+	defer res.Body.Close()
 	if res.StatusCode != http.StatusOK {
-		err = fmt.Errorf("microsoft graph client returned status code: %d", res.StatusCode)
+		slog.ErrorContext(ctx, "Error received from Microsoft Graph API", "status_code", res.StatusCode, "body", string(raw))
+		err = fmt.Errorf("Microsoft Graph client returned status code: %d", res.StatusCode)
 		return
 	}
-	defer res.Body.Close()
-	if err := json.NewDecoder(res.Body).Decode(&result); err != nil {
-		return nil, err
+	result, err = plugindata.UnmarshalJSON(raw)
+	if err != nil {
+		return nil, fmt.Errorf("failed to unmarshal results: %s", err)
 	}
 	return
 }
+
+func (cli *graphClient) QueryGraph(
+	ctx context.Context,
+	endpoint string,
+	queryParams url.Values,
+	size int,
+	onlyobjects bool,
+) (result plugindata.Data, err error) {
+	objects := make(plugindata.List, 0)
+
+	urlStr := graphUrl + fmt.Sprintf("/%s%s", cli.apiVersion, endpoint)
+	requestUrl, err := url.Parse(urlStr)
+	if err != nil {
+		return
+	}
+
+	if queryParams == nil {
+		queryParams = url.Values{}
+	}
+
+	queryParams.Set("$count", strconv.FormatBool(true))
+	queryParams.Set("$skip", strconv.Itoa(0))
+
+	limit := min(size, defaultPageSize)
+	queryParams.Set("$top", strconv.Itoa(limit))
+
+	var totalCount int = -1
+	var response plugindata.Data
+
+	for {
+
+		requestUrl.RawQuery = queryParams.Encode()
+
+		response, err = cli.fetchURL(ctx, requestUrl)
+		if err != nil {
+			slog.ErrorContext(ctx, "Error while fetching objects", "url", requestUrl.String(), "error", err)
+			return nil, err
+		}
+
+		resultMap, ok := response.(plugindata.Map)
+		if !ok {
+			return nil, fmt.Errorf("unexpected result type: %T", response)
+		}
+
+		countRaw, ok := resultMap["@odata.count"]
+		if ok {
+			totalCount = int(countRaw.(plugindata.Number))
+		}
+
+		objectsPageRaw, ok := resultMap["value"]
+		if !ok {
+			break
+		}
+
+		objectsPage, ok := objectsPageRaw.(plugindata.List)
+		if !ok {
+			return nil, fmt.Errorf("unexpected value type: %T", objectsPageRaw)
+		}
+
+		if len(objectsPage) == 0 {
+			break
+		}
+
+		slog.InfoContext(
+			ctx, "Objects fetched from Microsoft Graph API",
+			"fetched_overall", len(objects),
+			"fetched", len(objectsPage),
+			"total_available", totalCount,
+			"to_fetch_overall", size,
+		)
+
+		objects = append(objects, objectsPage...)
+		if len(objects) >= size {
+			break
+		}
+		queryParams.Set("$skip", strconv.Itoa(len(objects)))
+
+		// Number of objects to get, >= 0
+		objectsToGet := max(size-len(objects), 0)
+		pageSize := min(defaultPageSize, objectsToGet)
+		queryParams.Set("$top", strconv.Itoa(pageSize))
+	}
+
+	objectsToReturn := objects[:min(len(objects), size)]
+
+	if onlyobjects {
+		return objectsToReturn, nil
+	} else {
+		data := plugindata.Map{
+			"objects":     objectsToReturn,
+			"total_count": plugindata.Number(totalCount),
+		}
+		return data, nil
+	}
+}
+
+func (cli *graphClient) QueryGraphObject(
+	ctx context.Context,
+	endpoint string,
+) (result plugindata.Data, err error) {
+	urlStr := graphUrl + fmt.Sprintf("/%s%s", cli.apiVersion, endpoint)
+	requestUrl, err := url.Parse(urlStr)
+	if err != nil {
+		return
+	}
+	response, err := cli.fetchURL(ctx, requestUrl)
+	if err != nil {
+		slog.ErrorContext(ctx, "Error while fetching an object", "url", requestUrl.String(), "error", err)
+		return nil, err
+	}
+
+	return response, nil
+}

internal/microsoft/data_microsoft_graph.go

@@ -14,6 +14,8 @@ import (
 	"github.com/blackstork-io/fabric/plugin/plugindata"
 )
 
+const defaultSize = 50
+
 func makeMicrosoftGraphDataSource(loader MicrosoftGraphClientLoadFn) *plugin.DataSource {
 	return &plugin.DataSource{
 		Doc:      "The `microsoft_graph` data source queries Microsoft Graph.",
@@ -58,23 +60,43 @@ func makeMicrosoftGraphDataSource(loader MicrosoftGraphClientLoadFn) *plugin.Dat
 		Args: &dataspec.RootSpec{
 			Attrs: []*dataspec.AttrSpec{
 				{
-					Doc:        "The API version",
 					Name:       "api_version",
+					Doc:        "The API version",
 					Type:       cty.String,
 					DefaultVal: cty.StringVal("beta"),
 				},
 				{
-					Doc:         "The endpoint to query",
 					Name:        "endpoint",
+					Doc:         "The endpoint to query",
 					Type:        cty.String,
 					Constraints: constraint.RequiredNonNull,
 					ExampleVal:  cty.StringVal("/security/incidents"),
 				},
 				{
-					Doc:  "The query parameters",
 					Name: "query_params",
+					Doc:  "The query parameters",
 					Type: cty.Map(cty.String),
 				},
+				{
+					Name:         "objects_size",
+					Doc:          "Number of objects to be returned",
+					Type:         cty.Number,
+					Constraints:  constraint.NonNull,
+					DefaultVal:   cty.NumberIntVal(defaultSize),
+					MinInclusive: cty.NumberIntVal(1),
+				},
+				{
+					Name:       "only_objects",
+					Doc:        "Return only the list of objects. If `false`, returns an object with `objects` and `totalCount` fields",
+					Type:       cty.Bool,
+					DefaultVal: cty.BoolVal(true),
+				},
+				{
+					Name:       "is_object_endpoint",
+					Doc:        "If API endpoint response should be treated as a list or as an object. If set to `true`, `only_objects`, `query_params` and `objects_size` are ignored.",
+					Type:       cty.Bool,
+					DefaultVal: cty.BoolVal(false),
+				},
 			},
 		},
 	}
@@ -92,33 +114,39 @@ func fetchMicrosoftGraph(loader MicrosoftGraphClientLoadFn) plugin.RetrieveDataF
 			}}
 		}
 		endPoint := params.Args.GetAttrVal("endpoint").AsString()
-		queryParamsAttr := params.Args.GetAttrVal("query_params")
-		var queryParams url.Values
+		isObjectEndpoint := params.Args.GetAttrVal("is_object_endpoint")
 
-		if !queryParamsAttr.IsNull() {
-			queryParams = url.Values{}
-			queryMap := queryParamsAttr.AsValueMap()
-			for k, v := range queryMap {
-				queryParams.Add(k, v.AsString())
+		var response plugindata.Data
+
+		if isObjectEndpoint.True() {
+			response, err = cli.QueryGraphObject(ctx, endPoint)
+		} else {
+
+			queryParamsAttr := params.Args.GetAttrVal("query_params")
+			var queryParams url.Values
+
+			if !queryParamsAttr.IsNull() {
+				queryParams = url.Values{}
+				queryMap := queryParamsAttr.AsValueMap()
+				for k, v := range queryMap {
+					queryParams.Add(k, v.AsString())
+				}
 			}
-		}
 
-		response, err := cli.QueryGraph(ctx, endPoint, queryParams)
-		if err != nil {
-			return nil, diagnostics.Diag{{
-				Severity: hcl.DiagError,
-				Summary:  "Failed to query microsoft graph",
-				Detail:   err.Error(),
-			}}
+			onlyObjects := params.Args.GetAttrVal("only_objects")
+
+			size64, _ := params.Args.GetAttrVal("objects_size").AsBigFloat().Int64()
+			size := int(size64)
+
+			response, err = cli.QueryGraph(ctx, endPoint, queryParams, size, onlyObjects.True())
 		}
-		data, err := plugindata.ParseAny(response)
 		if err != nil {
 			return nil, diagnostics.Diag{{
 				Severity: hcl.DiagError,
-				Summary:  "Failed to parse response",
+				Summary:  "Failed to query microsoft graph",
 				Detail:   err.Error(),
 			}}
 		}
-		return data, nil
+		return response, nil
 	}
 }