Use context-aware CRUD functions for Project Protected Branches data source

Author: timofurrer

gitlab/data_source_gitlab_project_protected_branches.go

@@ -1,17 +1,19 @@
 package gitlab
 
 import (
+	"context"
 	"fmt"
 	"log"
 
+	"github.com/hashicorp/terraform-plugin-sdk/v2/diag"
 	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
 	"github.com/mitchellh/hashstructure"
 	"github.com/xanzy/go-gitlab"
 )
 
 func dataSourceGitlabProjectProtectedBranches() *schema.Resource {
 	return &schema.Resource{
-		Read: dataSourceGitlabProjectProtectedBranchesRead,
+		ReadContext: dataSourceGitlabProjectProtectedBranchesRead,
 		Schema: map[string]*schema.Schema{
 			"project_id": {
 				Type:        schema.TypeString,
@@ -58,26 +60,26 @@ type stateProtectedBranch struct {
 	CodeOwnerApprovalRequired bool                           `json:"code_owner_approval_required,omitempty" mapstructure:"code_owner_approval_required,omitempty"`
 }
 
-func dataSourceGitlabProjectProtectedBranchesRead(d *schema.ResourceData, meta interface{}) error {
+func dataSourceGitlabProjectProtectedBranchesRead(ctx context.Context, d *schema.ResourceData, meta interface{}) diag.Diagnostics {
 	client := meta.(*gitlab.Client)
 
 	log.Printf("[INFO] Reading Gitlab protected branch")
 
 	project := d.Get("project_id")
 
-	projectObject, _, err := client.Projects.GetProject(project, &gitlab.GetProjectOptions{})
+	projectObject, _, err := client.Projects.GetProject(project, &gitlab.GetProjectOptions{}, gitlab.WithContext(ctx))
 	if err != nil {
-		return err
+		return diag.FromErr(err)
 	}
 
 	allProtectedBranches := make([]stateProtectedBranch, 0)
 	totalPages := -1
 	opts := &gitlab.ListProtectedBranchesOptions{}
 	for opts.Page = 0; opts.Page != totalPages; opts.Page++ {
 		// Get protected branch by project ID/path and branch name
-		pbs, resp, err := client.ProtectedBranches.ListProtectedBranches(project, opts)
+		pbs, resp, err := client.ProtectedBranches.ListProtectedBranches(project, opts, gitlab.WithContext(ctx))
 		if err != nil {
-			return err
+			return diag.FromErr(err)
 		}
 		totalPages = resp.TotalPages
 		for _, pb := range pbs {
@@ -94,12 +96,12 @@ func dataSourceGitlabProjectProtectedBranchesRead(d *schema.ResourceData, meta i
 
 	// lintignore:R004 // TODO: Resolve this tfproviderlint issue
 	if err := d.Set("protected_branches", allProtectedBranches); err != nil {
-		return err
+		return diag.FromErr(err)
 	}
 
 	h, err := hashstructure.Hash(*opts, nil)
 	if err != nil {
-		return err
+		return diag.FromErr(err)
 	}
 
 	d.SetId(fmt.Sprintf("%d-%d", projectObject.ID, h))