Add context to all calls.

This allows passing through a custom context to the underlying transport.
elastic · Sep 10, 2024 · 59be954 · 59be954
1 parent 4ceec3a
commit 59be954
Show file tree

Hide file tree

Showing 10 changed files with 40 additions and 10 deletions.
diff --git a/pkg/api/organizationapi/get.go b/pkg/api/organizationapi/get.go
@@ -18,6 +18,7 @@
 package organizationapi
 
 import (
+	"context"
 	"errors"
 	"github.com/elastic/cloud-sdk-go/pkg/api"
 	"github.com/elastic/cloud-sdk-go/pkg/api/apierror"
@@ -28,6 +29,7 @@ import (
 
 type GetParams struct {
 	*api.API
+	Context context.Context
 
 	OrganizationID string
 }
@@ -50,7 +52,8 @@ func Get(params GetParams) (*models.Organization, error) {
 
 	response, err := params.V1API.Organizations.GetOrganization(
 		organizations.NewGetOrganizationParams().
-			WithOrganizationID(params.OrganizationID),
+			WithOrganizationID(params.OrganizationID).
+			WithContext(params.Context),
 		params.AuthWriter,
 	)
 	if err != nil {

diff --git a/pkg/api/organizationapi/invitation_create.go b/pkg/api/organizationapi/invitation_create.go
@@ -18,6 +18,7 @@
 package organizationapi
 
 import (
+	"context"
 	"errors"
 	"github.com/elastic/cloud-sdk-go/pkg/api"
 	"github.com/elastic/cloud-sdk-go/pkg/api/apierror"
@@ -28,6 +29,7 @@ import (
 
 type CreateInvitationParams struct {
 	*api.API
+	Context context.Context
 
 	OrganizationID string
 
@@ -67,7 +69,8 @@ func CreateInvitation(params CreateInvitationParams) (*models.OrganizationInvita
 				Emails:          params.Emails,
 				ExpiresIn:       params.ExpiresIn,
 				RoleAssignments: params.RoleAssignments,
-			}),
+			}).
+			WithContext(params.Context),
 		params.AuthWriter,
 	)
 	if err != nil {

diff --git a/pkg/api/organizationapi/invitation_delete.go b/pkg/api/organizationapi/invitation_delete.go
@@ -18,6 +18,7 @@
 package organizationapi
 
 import (
+	"context"
 	"errors"
 	"github.com/elastic/cloud-sdk-go/pkg/api"
 	"github.com/elastic/cloud-sdk-go/pkg/api/apierror"
@@ -29,6 +30,7 @@ import (
 
 type DeleteInvitationParams struct {
 	*api.API
+	Context context.Context
 
 	OrganizationID string
 
@@ -59,7 +61,8 @@ func DeleteInvitation(params DeleteInvitationParams) (models.EmptyResponse, erro
 	response, err := params.V1API.Organizations.DeleteOrganizationInvitations(
 		organizations.NewDeleteOrganizationInvitationsParams().
 			WithOrganizationID(params.OrganizationID).
-			WithInvitationTokens(invitationTokens),
+			WithInvitationTokens(invitationTokens).
+			WithContext(params.Context),
 		params.AuthWriter,
 	)
 	if err != nil {

diff --git a/pkg/api/organizationapi/invitation_list.go b/pkg/api/organizationapi/invitation_list.go
@@ -18,6 +18,7 @@
 package organizationapi
 
 import (
+	"context"
 	"errors"
 	"github.com/elastic/cloud-sdk-go/pkg/api"
 	"github.com/elastic/cloud-sdk-go/pkg/api/apierror"
@@ -28,6 +29,7 @@ import (
 
 type ListInvitationsParams struct {
 	*api.API
+	Context context.Context
 
 	OrganizationID string
 }
@@ -50,7 +52,8 @@ func ListInvitations(params ListInvitationsParams) (*models.OrganizationInvitati
 
 	response, err := params.V1API.Organizations.ListOrganizationInvitations(
 		organizations.NewListOrganizationInvitationsParams().
-			WithOrganizationID(params.OrganizationID),
+			WithOrganizationID(params.OrganizationID).
+			WithContext(params.Context),
 		params.AuthWriter,
 	)
 	if err != nil {

diff --git a/pkg/api/organizationapi/list.go b/pkg/api/organizationapi/list.go
@@ -18,6 +18,7 @@
 package organizationapi
 
 import (
+	"context"
 	"github.com/elastic/cloud-sdk-go/pkg/api"
 	"github.com/elastic/cloud-sdk-go/pkg/api/apierror"
 	"github.com/elastic/cloud-sdk-go/pkg/client/organizations"
@@ -27,6 +28,7 @@ import (
 
 type ListParams struct {
 	*api.API
+	Context context.Context
 }
 
 func (params ListParams) Validate() error {
@@ -43,7 +45,8 @@ func List(params ListParams) ([]*models.Organization, error) {
 	}
 
 	response, err := params.V1API.Organizations.ListOrganizations(
-		organizations.NewListOrganizationsParams(),
+		organizations.NewListOrganizationsParams().
+			WithContext(params.Context),
 		params.AuthWriter,
 	)
 	if err != nil {

diff --git a/pkg/api/organizationapi/member_delete.go b/pkg/api/organizationapi/member_delete.go
@@ -18,6 +18,7 @@
 package organizationapi
 
 import (
+	"context"
 	"errors"
 	"github.com/elastic/cloud-sdk-go/pkg/api"
 	"github.com/elastic/cloud-sdk-go/pkg/api/apierror"
@@ -29,6 +30,7 @@ import (
 
 type DeleteMemberParams struct {
 	*api.API
+	Context context.Context
 
 	OrganizationID string
 
@@ -58,7 +60,8 @@ func DeleteMember(params DeleteMemberParams) (models.EmptyResponse, error) {
 	response, err := params.V1API.Organizations.DeleteOrganizationMemberships(
 		organizations.NewDeleteOrganizationMembershipsParams().
 			WithOrganizationID(params.OrganizationID).
-			WithUserIds(userIds),
+			WithUserIds(userIds).
+			WithContext(params.Context),
 		params.AuthWriter,
 	)
 	if err != nil {

diff --git a/pkg/api/organizationapi/member_list.go b/pkg/api/organizationapi/member_list.go
@@ -18,6 +18,7 @@
 package organizationapi
 
 import (
+	"context"
 	"errors"
 	"github.com/elastic/cloud-sdk-go/pkg/api"
 	"github.com/elastic/cloud-sdk-go/pkg/api/apierror"
@@ -28,6 +29,7 @@ import (
 
 type ListMembersParams struct {
 	*api.API
+	Context context.Context
 
 	OrganizationID string
 }
@@ -50,7 +52,8 @@ func ListMembers(params ListMembersParams) (*models.OrganizationMemberships, err
 
 	response, err := params.V1API.Organizations.ListOrganizationMembers(
 		organizations.NewListOrganizationMembersParams().
-			WithOrganizationID(params.OrganizationID),
+			WithOrganizationID(params.OrganizationID).
+			WithContext(params.Context),
 		params.AuthWriter,
 	)
 	if err != nil {

diff --git a/pkg/api/organizationapi/member_role_assignments_add.go b/pkg/api/organizationapi/member_role_assignments_add.go
@@ -18,6 +18,7 @@
 package organizationapi
 
 import (
+	"context"
 	"errors"
 	"github.com/elastic/cloud-sdk-go/pkg/api"
 	"github.com/elastic/cloud-sdk-go/pkg/api/apierror"
@@ -28,6 +29,7 @@ import (
 
 type AddRoleAssignmentsParams struct {
 	*api.API
+	Context context.Context
 
 	UserID          string
 	RoleAssignments models.RoleAssignments
@@ -52,7 +54,8 @@ func AddRoleAssignments(params AddRoleAssignmentsParams) (*models.EmptyResponse,
 	response, err := params.V1API.UserRoleAssignments.AddRoleAssignments(
 		user_role_assignments.NewAddRoleAssignmentsParams().
 			WithUserID(params.UserID).
-			WithBody(&params.RoleAssignments),
+			WithBody(&params.RoleAssignments).
+			WithContext(params.Context),
 		params.AuthWriter,
 	)
 	if err != nil {

diff --git a/pkg/api/organizationapi/member_role_assignments_remove.go b/pkg/api/organizationapi/member_role_assignments_remove.go
@@ -18,6 +18,7 @@
 package organizationapi
 
 import (
+	"context"
 	"errors"
 	"github.com/elastic/cloud-sdk-go/pkg/api"
 	"github.com/elastic/cloud-sdk-go/pkg/api/apierror"
@@ -28,6 +29,7 @@ import (
 
 type RemoveRoleAssignmentsParams struct {
 	*api.API
+	Context context.Context
 
 	UserID          string
 	RoleAssignments models.RoleAssignments
@@ -52,7 +54,8 @@ func RemoveRoleAssignments(params RemoveRoleAssignmentsParams) (*models.EmptyRes
 	response, err := params.V1API.UserRoleAssignments.RemoveRoleAssignments(
 		user_role_assignments.NewRemoveRoleAssignmentsParams().
 			WithUserID(params.UserID).
-			WithBody(&params.RoleAssignments),
+			WithBody(&params.RoleAssignments).
+			WithContext(params.Context),
 		params.AuthWriter,
 	)
 	if err != nil {

diff --git a/pkg/api/organizationapi/update.go b/pkg/api/organizationapi/update.go
@@ -18,6 +18,7 @@
 package organizationapi
 
 import (
+	"context"
 	"errors"
 	"github.com/elastic/cloud-sdk-go/pkg/api"
 	"github.com/elastic/cloud-sdk-go/pkg/api/apierror"
@@ -28,6 +29,7 @@ import (
 
 type UpdateParams struct {
 	*api.API
+	Context context.Context
 
 	OrganizationID                string
 	Name                          string
@@ -62,7 +64,8 @@ func Update(params UpdateParams) (*models.Organization, error) {
 				BillingContacts:                  params.BillingContacts,
 				NotificationsAllowedEmailDomains: params.AllowedEmailDomains,
 				OperationalContacts:              params.OperationalContacts,
-			}),
+			}).
+			WithContext(params.Context),
 		params.AuthWriter,
 	)
 	if err != nil {