From a972c6bb008763703ac83cbfaf9022c5fa4acf0d Mon Sep 17 00:00:00 2001
From: Clayton Coleman <ccoleman@redhat.com>
Date: Wed, 15 Nov 2017 21:00:32 -0500
Subject: [PATCH] Allow resource.Builder to modify requests per client

Gives the builder a hook point to add settings to each request. These
settings are applied before the request is created and so are unable to
view the request. Intended to set controls on a per request basis.
---
 pkg/kubectl/resource/BUILD           |  1 +
 pkg/kubectl/resource/builder.go      | 10 ++++++++
 pkg/kubectl/resource/builder_test.go | 21 +++++++++++++++
 pkg/kubectl/resource/interfaces.go   | 38 ++++++++++++++++++++++++++++
 4 files changed, 70 insertions(+)

diff --git a/pkg/kubectl/resource/BUILD b/pkg/kubectl/resource/BUILD
index 87c405fefa3..e6dd3e8b680 100644
--- a/pkg/kubectl/resource/BUILD
+++ b/pkg/kubectl/resource/BUILD
@@ -73,6 +73,7 @@ go_test(
         "//vendor/k8s.io/apimachinery/pkg/util/errors:go_default_library",
         "//vendor/k8s.io/apimachinery/pkg/watch:go_default_library",
         "//vendor/k8s.io/client-go/kubernetes/scheme:go_default_library",
+        "//vendor/k8s.io/client-go/rest:go_default_library",
         "//vendor/k8s.io/client-go/rest/fake:go_default_library",
         "//vendor/k8s.io/client-go/rest/watch:go_default_library",
         "//vendor/k8s.io/client-go/util/testing:go_default_library",
diff --git a/pkg/kubectl/resource/builder.go b/pkg/kubectl/resource/builder.go
index a45c342c520..205356dd046 100644
--- a/pkg/kubectl/resource/builder.go
+++ b/pkg/kubectl/resource/builder.go
@@ -58,6 +58,7 @@ type Builder struct {
 	selectAll            bool
 	includeUninitialized bool
 	limitChunks          int64
+	requestTransforms    []RequestTransform
 
 	resources []string
 
@@ -354,6 +355,13 @@ func (b *Builder) RequestChunksOf(chunkSize int64) *Builder {
 	return b
 }
 
+// TransformRequests alters API calls made by clients requested from this builder. Pass
+// an empty list to clear modifiers.
+func (b *Builder) TransformRequests(opts ...RequestTransform) *Builder {
+	b.requestTransforms = opts
+	return b
+}
+
 // SelectEverythingParam
 func (b *Builder) SelectAllParam(selectAll bool) *Builder {
 	if selectAll && (b.labelSelector != nil || b.fieldSelector != nil) {
@@ -656,6 +664,7 @@ func (b *Builder) visitBySelector() *Result {
 			result.err = err
 			return result
 		}
+		client = NewClientWithOptions(client, b.requestTransforms...)
 		selectorNamespace := b.namespace
 		if mapping.Scope.Name() != meta.RESTScopeNameNamespace {
 			selectorNamespace = ""
@@ -705,6 +714,7 @@ func (b *Builder) visitByResource() *Result {
 			result.err = err
 			return result
 		}
+		client = NewClientWithOptions(client, b.requestTransforms...)
 		clients[s] = client
 	}
 
diff --git a/pkg/kubectl/resource/builder_test.go b/pkg/kubectl/resource/builder_test.go
index b9ab9012b4c..fc03370a3ef 100644
--- a/pkg/kubectl/resource/builder_test.go
+++ b/pkg/kubectl/resource/builder_test.go
@@ -42,6 +42,7 @@ import (
 	"k8s.io/apimachinery/pkg/runtime/serializer/streaming"
 	utilerrors "k8s.io/apimachinery/pkg/util/errors"
 	"k8s.io/apimachinery/pkg/watch"
+	"k8s.io/client-go/rest"
 	"k8s.io/client-go/rest/fake"
 	restclientwatch "k8s.io/client-go/rest/watch"
 	utiltesting "k8s.io/client-go/util/testing"
@@ -609,6 +610,26 @@ func TestMultipleResourceByTheSameName(t *testing.T) {
 	}
 }
 
+func TestRequestModifier(t *testing.T) {
+	var got *rest.Request
+	b := NewBuilder(restmapper, categories.LegacyCategoryExpander, scheme.Scheme, fakeClientWith("test", t, nil), corev1Codec).
+		NamespaceParam("foo").
+		TransformRequests(func(req *rest.Request) {
+			got = req
+		}).
+		ResourceNames("", "services/baz").
+		RequireObject(false)
+
+	i, err := b.Do().Infos()
+	if err != nil {
+		t.Fatal(err)
+	}
+	req := i[0].Client.Get()
+	if got != req {
+		t.Fatalf("request was not received by modifier: %#v", req)
+	}
+}
+
 func TestResourceNames(t *testing.T) {
 	pods, svc := testData()
 	b := NewBuilder(restmapper, categories.LegacyCategoryExpander, scheme.Scheme, fakeClientWith("", t, map[string]string{
diff --git a/pkg/kubectl/resource/interfaces.go b/pkg/kubectl/resource/interfaces.go
index bb7a956cde2..8b75d41dbf9 100644
--- a/pkg/kubectl/resource/interfaces.go
+++ b/pkg/kubectl/resource/interfaces.go
@@ -44,3 +44,41 @@ type ClientMapperFunc func(mapping *meta.RESTMapping) (RESTClient, error)
 func (f ClientMapperFunc) ClientForMapping(mapping *meta.RESTMapping) (RESTClient, error) {
 	return f(mapping)
 }
+
+// RequestTransform is a function that is given a chance to modify the outgoing request.
+type RequestTransform func(*client.Request)
+
+// NewClientWithOptions wraps the provided RESTClient and invokes each transform on each
+// newly created request.
+func NewClientWithOptions(c RESTClient, transforms ...RequestTransform) RESTClient {
+	return &clientOptions{c: c, transforms: transforms}
+}
+
+type clientOptions struct {
+	c          RESTClient
+	transforms []RequestTransform
+}
+
+func (c *clientOptions) modify(req *client.Request) *client.Request {
+	for _, transform := range c.transforms {
+		transform(req)
+	}
+	return req
+}
+
+func (c *clientOptions) Get() *client.Request {
+	return c.modify(c.c.Get())
+}
+
+func (c *clientOptions) Post() *client.Request {
+	return c.modify(c.c.Post())
+}
+func (c *clientOptions) Patch(t types.PatchType) *client.Request {
+	return c.modify(c.c.Patch(t))
+}
+func (c *clientOptions) Delete() *client.Request {
+	return c.modify(c.c.Delete())
+}
+func (c *clientOptions) Put() *client.Request {
+	return c.modify(c.c.Put())
+}