helm client: add methods that accept contexts so callers can cancel long-running operations

Signed-off-by: Benjamen Keroack <benjamen@dollarshaveclub.com>

pkg/helm/client.go

@@ -85,8 +85,29 @@ func (h *Client) InstallRelease(chstr, ns string, opts ...InstallOption) (*rls.I
 	return h.InstallReleaseFromChart(chart, ns, opts...)
 }
 
+// InstallReleaseWithContext loads a chart from chstr, installs it, and returns the release response while accepting a context.
+func (h *Client) InstallReleaseWithContext(ctx context.Context, chstr, ns string, opts ...InstallOption) (*rls.InstallReleaseResponse, error) {
+	// load the chart to install
+	chart, err := chartutil.Load(chstr)
+	if err != nil {
+		return nil, err
+	}
+
+	return h.installReleaseFromChartWithContext(ctx, chart, ns, opts...)
+}
+
 // InstallReleaseFromChart installs a new chart and returns the release response.
 func (h *Client) InstallReleaseFromChart(chart *chart.Chart, ns string, opts ...InstallOption) (*rls.InstallReleaseResponse, error) {
+	return h.installReleaseFromChartWithContext(NewContext(), chart, ns, opts...)
+}
+
+// InstallReleaseFromChartWithContext installs a new chart and returns the release response while accepting a context.
+func (h *Client) InstallReleaseFromChartWithContext(ctx context.Context, chart *chart.Chart, ns string, opts ...InstallOption) (*rls.InstallReleaseResponse, error) {
+	return h.installReleaseFromChartWithContext(ctx, chart, ns, opts...)
+}
+
+// InstallReleaseFromChartWithContext installs a new chart and returns the release response while accepting a context.
+func (h *Client) installReleaseFromChartWithContext(ctx context.Context, chart *chart.Chart, ns string, opts ...InstallOption) (*rls.InstallReleaseResponse, error) {
 	// apply the install options
 	reqOpts := h.opts
 	for _, opt := range opts {
@@ -99,7 +120,7 @@ func (h *Client) InstallReleaseFromChart(chart *chart.Chart, ns string, opts ...
 	req.DisableHooks = reqOpts.disableHooks
 	req.DisableCrdHook = reqOpts.disableCRDHook
 	req.ReuseName = reqOpts.reuseName
-	ctx := NewContext()
+	ctx = FromContext(ctx)
 
 	if reqOpts.before != nil {
 		if err := reqOpts.before(ctx, req); err != nil {
@@ -159,8 +180,29 @@ func (h *Client) UpdateRelease(rlsName string, chstr string, opts ...UpdateOptio
 	return h.UpdateReleaseFromChart(rlsName, chart, opts...)
 }
 
+// UpdateReleaseWithContext loads a chart from chstr and updates a release to a new/different chart while accepting a context.
+func (h *Client) UpdateReleaseWithContext(ctx context.Context, rlsName string, chstr string, opts ...UpdateOption) (*rls.UpdateReleaseResponse, error) {
+	// load the chart to update
+	chart, err := chartutil.Load(chstr)
+	if err != nil {
+		return nil, err
+	}
+
+	return h.updateReleaseFromChartWithContext(ctx, rlsName, chart, opts...)
+}
+
 // UpdateReleaseFromChart updates a release to a new/different chart.
 func (h *Client) UpdateReleaseFromChart(rlsName string, chart *chart.Chart, opts ...UpdateOption) (*rls.UpdateReleaseResponse, error) {
+	return h.updateReleaseFromChartWithContext(NewContext(), rlsName, chart, opts...)
+}
+
+// UpdateReleaseFromChartWithContext updates a release to a new/different chart while accepting a context.
+func (h *Client) UpdateReleaseFromChartWithContext(ctx context.Context, rlsName string, chart *chart.Chart, opts ...UpdateOption) (*rls.UpdateReleaseResponse, error) {
+	return h.updateReleaseFromChartWithContext(ctx, rlsName, chart, opts...)
+}
+
+// updateReleaseFromChartWithContext updates a release to a new/different chart and accepts a context.
+func (h *Client) updateReleaseFromChartWithContext(ctx context.Context, rlsName string, chart *chart.Chart, opts ...UpdateOption) (*rls.UpdateReleaseResponse, error) {
 	// apply the update options
 	reqOpts := h.opts
 	for _, opt := range opts {
@@ -175,7 +217,7 @@ func (h *Client) UpdateReleaseFromChart(rlsName string, chart *chart.Chart, opts
 	req.Force = reqOpts.force
 	req.ResetValues = reqOpts.resetValues
 	req.ReuseValues = reqOpts.reuseValues
-	ctx := NewContext()
+	ctx = FromContext(ctx)
 
 	if reqOpts.before != nil {
 		if err := reqOpts.before(ctx, req); err != nil {

pkg/helm/fake.go

@@ -24,6 +24,7 @@ import (
 	"sync"
 
 	"github.com/golang/protobuf/ptypes/timestamp"
+	"golang.org/x/net/context"
 	"k8s.io/helm/pkg/chartutil"
 	"k8s.io/helm/pkg/manifest"
 	"k8s.io/helm/pkg/proto/hapi/chart"
@@ -87,6 +88,16 @@ func (c *FakeClient) InstallRelease(chStr, ns string, opts ...InstallOption) (*r
 	return c.InstallReleaseFromChart(chart, ns, opts...)
 }
 
+// InstallReleaseWithContext creates a new release and returns a InstallReleaseResponse containing that release and accepts a context
+func (c *FakeClient) InstallReleaseWithContext(ctx context.Context, chStr, ns string, opts ...InstallOption) (*rls.InstallReleaseResponse, error) {
+	return c.InstallRelease(chStr, ns, opts...)
+}
+
+// InstallReleaseFromChartWithContext adds a new MockRelease to the fake client and returns a InstallReleaseResponse containing that release and accepts a context
+func (c *FakeClient) InstallReleaseFromChartWithContext(ctx context.Context, chart *chart.Chart, ns string, opts ...InstallOption) (*rls.InstallReleaseResponse, error) {
+	return c.InstallReleaseFromChart(chart, ns, opts...)
+}
+
 // InstallReleaseFromChart adds a new MockRelease to the fake client and returns a InstallReleaseResponse containing that release
 func (c *FakeClient) InstallReleaseFromChart(chart *chart.Chart, ns string, opts ...InstallOption) (*rls.InstallReleaseResponse, error) {
 	for _, opt := range opts {
@@ -155,6 +166,16 @@ func (c *FakeClient) UpdateRelease(rlsName string, chStr string, opts ...UpdateO
 	return c.UpdateReleaseFromChart(rlsName, &chart.Chart{}, opts...)
 }
 
+// UpdateReleaseWithContext returns an UpdateReleaseResponse containing the updated release, if it exists and accepts a context
+func (c *FakeClient) UpdateReleaseWithContext(ctx context.Context, rlsName string, chStr string, opts ...UpdateOption) (*rls.UpdateReleaseResponse, error) {
+	return c.UpdateRelease(rlsName, chStr, opts...)
+}
+
+// UpdateReleaseFromChartWithContext returns an UpdateReleaseResponse containing the updated release, if it exists and accepts a context
+func (c *FakeClient) UpdateReleaseFromChartWithContext(ctx context.Context, rlsName string, newChart *chart.Chart, opts ...UpdateOption) (*rls.UpdateReleaseResponse, error) {
+	return c.UpdateReleaseFromChart(rlsName, newChart, opts...)
+}
+
 // UpdateReleaseFromChart returns an UpdateReleaseResponse containing the updated release, if it exists
 func (c *FakeClient) UpdateReleaseFromChart(rlsName string, newChart *chart.Chart, opts ...UpdateOption) (*rls.UpdateReleaseResponse, error) {
 	for _, opt := range opts {

pkg/helm/interface.go

@@ -17,6 +17,7 @@ limitations under the License.
 package helm
 
 import (
+	"golang.org/x/net/context"
 	"k8s.io/helm/pkg/proto/hapi/chart"
 	rls "k8s.io/helm/pkg/proto/hapi/services"
 )
@@ -25,11 +26,15 @@ import (
 type Interface interface {
 	ListReleases(opts ...ReleaseListOption) (*rls.ListReleasesResponse, error)
 	InstallRelease(chStr, namespace string, opts ...InstallOption) (*rls.InstallReleaseResponse, error)
+	InstallReleaseWithContext(ctx context.Context, chStr, namespace string, opts ...InstallOption) (*rls.InstallReleaseResponse, error)
 	InstallReleaseFromChart(chart *chart.Chart, namespace string, opts ...InstallOption) (*rls.InstallReleaseResponse, error)
+	InstallReleaseFromChartWithContext(ctx context.Context, chart *chart.Chart, namespace string, opts ...InstallOption) (*rls.InstallReleaseResponse, error)
 	DeleteRelease(rlsName string, opts ...DeleteOption) (*rls.UninstallReleaseResponse, error)
 	ReleaseStatus(rlsName string, opts ...StatusOption) (*rls.GetReleaseStatusResponse, error)
 	UpdateRelease(rlsName, chStr string, opts ...UpdateOption) (*rls.UpdateReleaseResponse, error)
+	UpdateReleaseWithContext(ctx context.Context, rlsName, chStr string, opts ...UpdateOption) (*rls.UpdateReleaseResponse, error)
 	UpdateReleaseFromChart(rlsName string, chart *chart.Chart, opts ...UpdateOption) (*rls.UpdateReleaseResponse, error)
+	UpdateReleaseFromChartWithContext(ctx context.Context, rlsName string, chart *chart.Chart, opts ...UpdateOption) (*rls.UpdateReleaseResponse, error)
 	RollbackRelease(rlsName string, opts ...RollbackOption) (*rls.RollbackReleaseResponse, error)
 	ReleaseContent(rlsName string, opts ...ContentOption) (*rls.GetReleaseContentResponse, error)
 	ReleaseHistory(rlsName string, opts ...HistoryOption) (*rls.GetHistoryResponse, error)

pkg/helm/option.go

@@ -507,8 +507,13 @@ func WithMaxHistory(max int32) HistoryOption {
 
 // NewContext creates a versioned context.
 func NewContext() context.Context {
+	return FromContext(context.TODO())
+}
+
+// FromContext returns a versioned context from a parent context
+func FromContext(ctx context.Context) context.Context {
 	md := metadata.Pairs("x-helm-api-client", version.GetVersion())
-	return metadata.NewOutgoingContext(context.TODO(), md)
+	return metadata.NewOutgoingContext(ctx, md)
 }
 
 // ReleaseTestOption allows configuring optional request data for

pkg/releasetesting/test_suite_test.go

@@ -415,7 +415,9 @@ func (rs mockStream) SendHeader(m metadata.MD) error { return nil }
 func (rs mockStream) SetTrailer(m metadata.MD)       {}
 func (rs mockStream) SendMsg(v interface{}) error    { return nil }
 func (rs mockStream) RecvMsg(v interface{}) error    { return nil }
-func (rs mockStream) Context() context.Context       { return helm.NewContext() }
+func (rs mockStream) Context() context.Context {
+	return helm.NewContext()
+}
 
 type podSucceededKubeClient struct {
 	tillerEnv.PrintingKubeClient

pkg/tiller/release_server_test.go

@@ -584,7 +584,9 @@ func (rs mockRunReleaseTestServer) SendHeader(m metadata.MD) error { return nil
 func (rs mockRunReleaseTestServer) SetTrailer(m metadata.MD)       {}
 func (rs mockRunReleaseTestServer) SendMsg(v interface{}) error    { return nil }
 func (rs mockRunReleaseTestServer) RecvMsg(v interface{}) error    { return nil }
-func (rs mockRunReleaseTestServer) Context() context.Context       { return helm.NewContext() }
+func (rs mockRunReleaseTestServer) Context() context.Context {
+	return helm.NewContext()
+}
 
 type mockHooksManifest struct {
 	Metadata struct {

pkg/tiller/release_update_test.go

@@ -18,6 +18,7 @@ package tiller
 
 import (
 	"fmt"
+	"reflect"
 	"strings"
 	"testing"
 
@@ -28,7 +29,6 @@ import (
 	"k8s.io/helm/pkg/proto/hapi/chart"
 	"k8s.io/helm/pkg/proto/hapi/release"
 	"k8s.io/helm/pkg/proto/hapi/services"
-	"reflect"
 )
 
 func TestUpdateRelease(t *testing.T) {