Propagate context into datasource read functions (#1282)

Signed-off-by: Dave Henderson <dhenderson@gmail.com>

23 files changed, 99 insertions, 69 deletions

diff --git a/data/datasource.go b/data/datasource.go
index 77dc7990..5f684a4c 100644
--- a/data/datasource.go
+++ b/data/datasource.go
@@ -37,7 +37,7 @@ func init() {
 
 // registerReaders registers the source-reader functions
 func (d *Data) registerReaders() {
-	d.sourceReaders = make(map[string]func(*Source, ...string) ([]byte, error))
+	d.sourceReaders = make(map[string]func(context.Context, *Source, ...string) ([]byte, error))
 
 	d.sourceReaders["aws+smp"] = readAWSSMP
 	d.sourceReaders["aws+sm"] = readAWSSecretsManager
@@ -64,7 +64,7 @@ func (d *Data) registerReaders() {
 }
 
 // lookupReader - return the reader function for the given scheme
-func (d *Data) lookupReader(scheme string) (func(*Source, ...string) ([]byte, error), error) {
+func (d *Data) lookupReader(scheme string) (func(context.Context, *Source, ...string) ([]byte, error), error) {
 	if d.sourceReaders == nil {
 		d.registerReaders()
 	}
@@ -81,7 +81,7 @@ type Data struct {
 
 	Sources map[string]*Source
 
-	sourceReaders map[string]func(*Source, ...string) ([]byte, error)
+	sourceReaders map[string]func(context.Context, *Source, ...string) ([]byte, error)
 	cache         map[string][]byte
 
 	// headers from the --datasource-header/-H option that don't reference datasources from the commandline
@@ -279,12 +279,12 @@ func (d *Data) lookupSource(alias string) (*Source, error) {
 	return source, nil
 }
 
-func (d *Data) readDataSource(alias string, args ...string) (data, mimeType string, err error) {
+func (d *Data) readDataSource(ctx context.Context, alias string, args ...string) (data, mimeType string, err error) {
 	source, err := d.lookupSource(alias)
 	if err != nil {
 		return "", "", err
 	}
-	b, err := d.readSource(source, args...)
+	b, err := d.readSource(ctx, source, args...)
 	if err != nil {
 		return "", "", errors.Wrapf(err, "Couldn't read datasource '%s'", alias)
 	}
@@ -302,13 +302,13 @@ func (d *Data) readDataSource(alias string, args ...string) (data, mimeType stri
 
 // Include -
 func (d *Data) Include(alias string, args ...string) (string, error) {
-	data, _, err := d.readDataSource(alias, args...)
+	data, _, err := d.readDataSource(d.ctx, alias, args...)
 	return data, err
 }
 
 // Datasource -
 func (d *Data) Datasource(alias string, args ...string) (interface{}, error) {
-	data, mimeType, err := d.readDataSource(alias, args...)
+	data, mimeType, err := d.readDataSource(d.ctx, alias, args...)
 	if err != nil {
 		return nil, err
 	}
@@ -353,13 +353,13 @@ func (d *Data) DatasourceReachable(alias string, args ...string) bool {
 	if !ok {
 		return false
 	}
-	_, err := d.readSource(source, args...)
+	_, err := d.readSource(d.ctx, source, args...)
 	return err == nil
 }
 
 // readSource returns the (possibly cached) data from the given source,
 // as referenced by the given args
-func (d *Data) readSource(source *Source, args ...string) ([]byte, error) {
+func (d *Data) readSource(ctx context.Context, source *Source, args ...string) ([]byte, error) {
 	if d.cache == nil {
 		d.cache = make(map[string][]byte)
 	}
@@ -375,7 +375,7 @@ func (d *Data) readSource(source *Source, args ...string) ([]byte, error) {
 	if err != nil {
 		return nil, errors.Wrap(err, "Datasource not yet supported")
 	}
-	data, err := r(source, args...)
+	data, err := r(ctx, source, args...)
 	if err != nil {
 		return nil, err
 	}
diff --git a/data/datasource_aws_sm.go b/data/datasource_aws_sm.go
index e6c100ed..f64bf5fa 100644
--- a/data/datasource_aws_sm.go
+++ b/data/datasource_aws_sm.go
@@ -1,12 +1,14 @@
 package data
 
 import (
+	"context"
 	"fmt"
 	"net/url"
 	"path"
 	"strings"
 
 	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/aws/request"
 	"github.com/aws/aws-sdk-go/service/secretsmanager"
 	"github.com/pkg/errors"
 
@@ -15,7 +17,7 @@ import (
 
 // awsSecretsManagerGetter - A subset of Secrets Manager API for use in unit testing
 type awsSecretsManagerGetter interface {
-	GetSecretValue(input *secretsmanager.GetSecretValueInput) (*secretsmanager.GetSecretValueOutput, error)
+	GetSecretValueWithContext(ctx context.Context, input *secretsmanager.GetSecretValueInput, opts ...request.Option) (*secretsmanager.GetSecretValueOutput, error)
 }
 
 func parseDatasourceURLArgs(sourceURL *url.URL, args ...string) (params map[string]interface{}, p string, err error) {
@@ -55,7 +57,7 @@ func parseDatasourceURLArgs(sourceURL *url.URL, args ...string) (params map[stri
 	return params, p, nil
 }
 
-func readAWSSecretsManager(source *Source, args ...string) (output []byte, err error) {
+func readAWSSecretsManager(ctx context.Context, source *Source, args ...string) (output []byte, err error) {
 	if source.awsSecretsManager == nil {
 		source.awsSecretsManager = secretsmanager.New(gaws.SDKSession())
 	}
@@ -65,15 +67,15 @@ func readAWSSecretsManager(source *Source, args ...string) (output []byte, err e
 		return nil, err
 	}
 
-	return readAWSSecretsManagerParam(source, paramPath)
+	return readAWSSecretsManagerParam(ctx, source, paramPath)
 }
 
-func readAWSSecretsManagerParam(source *Source, paramPath string) ([]byte, error) {
+func readAWSSecretsManagerParam(ctx context.Context, source *Source, paramPath string) ([]byte, error) {
 	input := &secretsmanager.GetSecretValueInput{
 		SecretId: aws.String(paramPath),
 	}
 
-	response, err := source.awsSecretsManager.GetSecretValue(input)
+	response, err := source.awsSecretsManager.GetSecretValueWithContext(ctx, input)
 	if err != nil {
 		return nil, errors.Wrapf(err, "Error reading aws+sm from AWS using GetSecretValue with input %v", input)
 	}
diff --git a/data/datasource_aws_sm_test.go b/data/datasource_aws_sm_test.go
index c29a9b0b..892d5429 100644
--- a/data/datasource_aws_sm_test.go
+++ b/data/datasource_aws_sm_test.go
@@ -1,11 +1,13 @@
 package data
 
 import (
+	"context"
 	"net/url"
 	"testing"
 
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/aws/awserr"
+	"github.com/aws/aws-sdk-go/aws/request"
 	"github.com/aws/aws-sdk-go/service/secretsmanager"
 	"github.com/stretchr/testify/assert"
 )
@@ -18,7 +20,7 @@ type DummyAWSSecretsManagerSecretGetter struct {
 	mockGetSecretValue func(input *secretsmanager.GetSecretValueInput) (*secretsmanager.GetSecretValueOutput, error)
 }
 
-func (d DummyAWSSecretsManagerSecretGetter) GetSecretValue(input *secretsmanager.GetSecretValueInput) (*secretsmanager.GetSecretValueOutput, error) {
+func (d DummyAWSSecretsManagerSecretGetter) GetSecretValueWithContext(ctx context.Context, input *secretsmanager.GetSecretValueInput, opts ...request.Option) (*secretsmanager.GetSecretValueOutput, error) {
 	if d.mockGetSecretValue != nil {
 		output, err := d.mockGetSecretValue(input)
 		return output, err
@@ -107,7 +109,7 @@ func TestAWSSecretsManager_GetParameterSetup(t *testing.T) {
 		},
 	})
 
-	_, err := readAWSSecretsManager(s, "/bar")
+	_, err := readAWSSecretsManager(context.Background(), s, "/bar")
 	assert.True(t, calledOk)
 	assert.Nil(t, err)
 }
@@ -123,7 +125,7 @@ func TestAWSSecretsManager_GetParameterSetupWrongArgs(t *testing.T) {
 		},
 	})
 
-	_, err := readAWSSecretsManager(s, "/bar", "/foo", "/bla")
+	_, err := readAWSSecretsManager(context.Background(), s, "/bar", "/foo", "/bla")
 	assert.False(t, calledOk)
 	assert.Error(t, err)
 }
@@ -135,7 +137,7 @@ func TestAWSSecretsManager_GetParameterMissing(t *testing.T) {
 		err: expectedErr,
 	})
 
-	_, err := readAWSSecretsManager(s, "")
+	_, err := readAWSSecretsManager(context.Background(), s, "")
 	assert.Error(t, err, "Test of error message")
 }
 
@@ -150,7 +152,7 @@ func TestAWSSecretsManager_ReadSecret(t *testing.T) {
 		},
 	})
 
-	output, err := readAWSSecretsManagerParam(s, "/foo/bar")
+	output, err := readAWSSecretsManagerParam(context.Background(), s, "/foo/bar")
 	assert.True(t, calledOk)
 	assert.NoError(t, err)
 	assert.Equal(t, []byte("blub"), output)
diff --git a/data/datasource_awssmp.go b/data/datasource_awssmp.go
index 26d93855..d61affe8 100644
--- a/data/datasource_awssmp.go
+++ b/data/datasource_awssmp.go
@@ -18,8 +18,7 @@ type awssmpGetter interface {
 	GetParametersByPathWithContext(ctx context.Context, input *ssm.GetParametersByPathInput, opts ...request.Option) (*ssm.GetParametersByPathOutput, error)
 }
 
-func readAWSSMP(source *Source, args ...string) (data []byte, err error) {
-	ctx := context.TODO()
+func readAWSSMP(ctx context.Context, source *Source, args ...string) (data []byte, err error) {
 	if source.asmpg == nil {
 		source.asmpg = ssm.New(gaws.SDKSession())
 	}
diff --git a/data/datasource_awssmp_test.go b/data/datasource_awssmp_test.go
index 3c14e325..7fac1084 100644
--- a/data/datasource_awssmp_test.go
+++ b/data/datasource_awssmp_test.go
@@ -71,7 +71,7 @@ func TestAWSSMP_GetParameterSetup(t *testing.T) {
 		},
 	})
 
-	_, err := readAWSSMP(s, "/bar")
+	_, err := readAWSSMP(context.Background(), s, "/bar")
 	assert.True(t, calledOk)
 	assert.Nil(t, err)
 }
@@ -88,7 +88,7 @@ func TestAWSSMP_GetParameterValidOutput(t *testing.T) {
 		param: expected,
 	})
 
-	output, err := readAWSSMP(s, "")
+	output, err := readAWSSMP(context.Background(), s, "")
 	assert.Nil(t, err)
 	actual := &ssm.Parameter{}
 	err = json.Unmarshal(output, &actual)
@@ -104,7 +104,7 @@ func TestAWSSMP_GetParameterMissing(t *testing.T) {
 		err: expectedErr,
 	})
 
-	_, err := readAWSSMP(s, "")
+	_, err := readAWSSMP(context.Background(), s, "")
 	assert.Error(t, err, "Test of error message")
 }
 
diff --git a/data/datasource_blob.go b/data/datasource_blob.go
index b40ae198..bb316e7f 100644
--- a/data/datasource_blob.go
+++ b/data/datasource_blob.go
@@ -20,13 +20,11 @@ import (
 	"gocloud.dev/gcp"
 )
 
-func readBlob(source *Source, args ...string) (output []byte, err error) {
+func readBlob(ctx context.Context, source *Source, args ...string) (output []byte, err error) {
 	if len(args) >= 2 {
 		return nil, errors.New("maximum two arguments to blob datasource: alias, extraPath")
 	}
 
-	ctx := context.TODO()
-
 	key := source.URL.Path
 	if len(args) == 1 {
 		key = path.Join(key, args[0])
diff --git a/data/datasource_blob_test.go b/data/datasource_blob_test.go
index 75d466f1..7763d8e9 100644
--- a/data/datasource_blob_test.go
+++ b/data/datasource_blob_test.go
@@ -2,6 +2,7 @@ package data
 
 import (
 	"bytes"
+	"context"
 	"net/http/httptest"
 	"net/url"
 	"os"
@@ -56,7 +57,7 @@ func putFile(backend gofakes3.Backend, bucket, file, mime, content string) error
 }
 
 func TestReadBlob(t *testing.T) {
-	_, err := readBlob(nil, "foo", "bar")
+	_, err := readBlob(context.Background(), nil, "foo", "bar")
 	assert.Error(t, err)
 
 	ts, u := setupTestBucket(t)
diff --git a/data/datasource_boltdb.go b/data/datasource_boltdb.go
index 81682e04..f7198c74 100644
--- a/data/datasource_boltdb.go
+++ b/data/datasource_boltdb.go
@@ -1,11 +1,13 @@
 package data
 
 import (
+	"context"
+
 	"github.com/hairyhenderson/gomplate/v3/libkv"
 	"github.com/pkg/errors"
 )
 
-func readBoltDB(source *Source, args ...string) (data []byte, err error) {
+func readBoltDB(ctx context.Context, source *Source, args ...string) (data []byte, err error) {
 	if source.kv == nil {
 		source.kv, err = libkv.NewBoltDB(source.URL)
 		if err != nil {
diff --git a/data/datasource_consul.go b/data/datasource_consul.go
index e0ec623c..da804395 100644
--- a/data/datasource_consul.go
+++ b/data/datasource_consul.go
@@ -1,12 +1,13 @@
 package data
 
 import (
+	"context"
 	"strings"
 
 	"github.com/hairyhenderson/gomplate/v3/libkv"
 )
 
-func readConsul(source *Source, args ...string) (data []byte, err error) {
+func readConsul(ctx context.Context, source *Source, args ...string) (data []byte, err error) {
 	if source.kv == nil {
 		source.kv, err = libkv.NewConsul(source.URL)
 		if err != nil {
diff --git a/data/datasource_env.go b/data/datasource_env.go
index 15d56a1d..f1e2e5aa 100644
--- a/data/datasource_env.go
+++ b/data/datasource_env.go
@@ -1,12 +1,13 @@
 package data
 
 import (
+	"context"
 	"strings"
 
 	"github.com/hairyhenderson/gomplate/v3/env"
 )
 
-func readEnv(source *Source, args ...string) (b []byte, err error) {
+func readEnv(ctx context.Context, source *Source, args ...string) (b []byte, err error) {
 	n := source.URL.Path
 	n = strings.TrimPrefix(n, "/")
 	if n == "" {
diff --git a/data/datasource_env_test.go b/data/datasource_env_test.go
index e368397a..c7723b3b 100644
--- a/data/datasource_env_test.go
+++ b/data/datasource_env_test.go
@@ -1,6 +1,7 @@
 package data
 
 import (
+	"context"
 	"net/url"
 	"os"
 	"testing"
@@ -14,6 +15,8 @@ func mustParseURL(in string) *url.URL {
 }
 
 func TestReadEnv(t *testing.T) {
+	ctx := context.Background()
+
 	content := []byte(`hello world`)
 	os.Setenv("HELLO_WORLD", "hello world")
 	defer os.Unsetenv("HELLO_WORLD")
@@ -22,31 +25,31 @@ func TestReadEnv(t *testing.T) {
 
 	source := &Source{Alias: "foo", URL: mustParseURL("env:HELLO_WORLD")}
 
-	actual, err := readEnv(source)
+	actual, err := readEnv(ctx, source)
 	assert.NoError(t, err)
 	assert.Equal(t, content, actual)
 
 	source = &Source{Alias: "foo", URL: mustParseURL("env:/HELLO_WORLD")}
 
-	actual, err = readEnv(source)
+	actual, err = readEnv(ctx, source)
 	assert.NoError(t, err)
 	assert.Equal(t, content, actual)
 
 	source = &Source{Alias: "foo", URL: mustParseURL("env:///HELLO_WORLD")}
 
-	actual, err = readEnv(source)
+	actual, err = readEnv(ctx, source)
 	assert.NoError(t, err)
 	assert.Equal(t, content, actual)
 
 	source = &Source{Alias: "foo", URL: mustParseURL("env:HELLO_WORLD?foo=bar")}
 
-	actual, err = readEnv(source)
+	actual, err = readEnv(ctx, source)
 	assert.NoError(t, err)
 	assert.Equal(t, content, actual)
 
 	source = &Source{Alias: "foo", URL: mustParseURL("env:///HELLO_WORLD?foo=bar")}
 
-	actual, err = readEnv(source)
+	actual, err = readEnv(ctx, source)
 	assert.NoError(t, err)
 	assert.Equal(t, content, actual)
 }
diff --git a/data/datasource_file.go b/data/datasource_file.go
index 175f92f0..ce86c0a2 100644
--- a/data/datasource_file.go
+++ b/data/datasource_file.go
@@ -2,6 +2,7 @@ package data
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"io/ioutil"
 	"net/url"
@@ -14,7 +15,7 @@ import (
 	"github.com/pkg/errors"
 )
 
-func readFile(source *Source, args ...string) ([]byte, error) {
+func readFile(ctx context.Context, source *Source, args ...string) ([]byte, error) {
 	if source.fs == nil {
 		source.fs = afero.NewOsFs()
 	}
diff --git a/data/datasource_file_test.go b/data/datasource_file_test.go
index aa51dbca..8f2a7f0a 100644
--- a/data/datasource_file_test.go
+++ b/data/datasource_file_test.go
@@ -1,6 +1,7 @@
 package data
 
 import (
+	"context"
 	"testing"
 
 	"github.com/spf13/afero"
@@ -9,6 +10,8 @@ import (
 )
 
 func TestReadFile(t *testing.T) {
+	ctx := context.Background()
+
 	content := []byte(`hello world`)
 	fs := afero.NewMemMapFs()
 
@@ -26,30 +29,30 @@ func TestReadFile(t *testing.T) {
 	source := &Source{Alias: "foo", URL: mustParseURL("file:///tmp/foo")}
 	source.fs = fs
 
-	actual, err := readFile(source)
+	actual, err := readFile(ctx, source)
 	assert.NoError(t, err)
 	assert.Equal(t, content, actual)
 
 	source = &Source{Alias: "bogus", URL: mustParseURL("file:///bogus")}
 	source.fs = fs
-	_, err = readFile(source)
+	_, err = readFile(ctx, source)
 	assert.Error(t, err)
 
 	source = &Source{Alias: "partial", URL: mustParseURL("file:///tmp/partial")}
 	source.fs = fs
-	actual, err = readFile(source, "foo.txt")
+	actual, err = readFile(ctx, source, "foo.txt")
 	assert.NoError(t, err)
 	assert.Equal(t, content, actual)
 
 	source = &Source{Alias: "dir", URL: mustParseURL("file:///tmp/partial/")}
 	source.fs = fs
-	actual, err = readFile(source)
+	actual, err = readFile(ctx, source)
 	assert.NoError(t, err)
 	assert.Equal(t, []byte(`["bar.txt","baz.txt","foo.txt"]`), actual)
 
 	source = &Source{Alias: "dir", URL: mustParseURL("file:///tmp/partial/?type=application/json")}
 	source.fs = fs
-	actual, err = readFile(source)
+	actual, err = readFile(ctx, source)
 	assert.NoError(t, err)
 	assert.Equal(t, []byte(`["bar.txt","baz.txt","foo.txt"]`), actual)
 	mime, err := source.mimeType("")
@@ -58,7 +61,7 @@ func TestReadFile(t *testing.T) {
 
 	source = &Source{Alias: "dir", URL: mustParseURL("file:///tmp/partial/?type=application/json")}
 	source.fs = fs
-	actual, err = readFile(source, "foo.txt")
+	actual, err = readFile(ctx, source, "foo.txt")
 	assert.NoError(t, err)
 	assert.Equal(t, content, actual)
 	mime, err = source.mimeType("")
diff --git a/data/datasource_git.go b/data/datasource_git.go
index bd1811b4..c2673dd3 100644
--- a/data/datasource_git.go
+++ b/data/datasource_git.go
@@ -27,8 +27,7 @@ import (
 	"github.com/go-git/go-git/v5/storage/memory"
 )
 
-func readGit(source *Source, args ...string) ([]byte, error) {
-	ctx := context.TODO()
+func readGit(ctx context.Context, source *Source, args ...string) ([]byte, error) {
 	g := gitsource{}
 
 	u := source.URL
diff --git a/data/datasource_git_test.go b/data/datasource_git_test.go
index 33771410..0c20c558 100644
--- a/data/datasource_git_test.go
+++ b/data/datasource_git_test.go
@@ -380,6 +380,7 @@ func TestOpenBareFileRepo(t *testing.T) {
 }
 
 func TestReadGit(t *testing.T) {
+	ctx := context.Background()
 	repoFS := setupGitRepo(t)
 
 	overrideFSLoader(repoFS)
@@ -389,7 +390,7 @@ func TestReadGit(t *testing.T) {
 		Alias: "hi",
 		URL:   mustParseURL("git+file:///bare.git//hello.txt"),
 	}
-	b, err := readGit(s)
+	b, err := readGit(ctx, s)
 	assert.NilError(t, err)
 	assert.Equal(t, "hello world", string(b))
 
@@ -397,7 +398,7 @@ func TestReadGit(t *testing.T) {
 		Alias: "hi",
 		URL:   mustParseURL("git+file:///bare.git"),
 	}
-	b, err = readGit(s)
+	b, err = readGit(ctx, s)
 	assert.NilError(t, err)
 	assert.Equal(t, "application/array+json", s.mediaType)
 	assert.Equal(t, `["hello.txt"]`, string(b))
diff --git a/data/datasource_http.go b/data/datasource_http.go
index e809768e..ed1c97b7 100644
--- a/data/datasource_http.go
+++ b/data/datasource_http.go
@@ -1,6 +1,7 @@
 package data
 
 import (
+	"context"
 	"io/ioutil"
 	"mime"
 	"net/http"
@@ -21,7 +22,7 @@ func buildURL(base *url.URL, args ...string) (*url.URL, error) {
 	return base.ResolveReference(p), nil
 }
 
-func readHTTP(source *Source, args ...string) ([]byte, error) {
+func readHTTP(ctx context.Context, source *Source, args ...string) ([]byte, error) {
 	if source.hc == nil {
 		source.hc = &http.Client{Timeout: time.Second * 5}
 	}
@@ -29,7 +30,7 @@ func readHTTP(source *Source, args ...string) ([]byte, error) {
 	if err != nil {
 		return nil, err
 	}
-	req, err := http.NewRequest("GET", u.String(), nil)
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, u.String(), nil)
 	if err != nil {
 		return nil, err
 	}
diff --git a/data/datasource_http_test.go b/data/datasource_http_test.go
index 7f8515cc..e2c13a8f 100644
--- a/data/datasource_http_test.go
+++ b/data/datasource_http_test.go
@@ -1,6 +1,7 @@
 package data
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"net/http"
@@ -57,6 +58,7 @@ func TestHTTPFile(t *testing.T) {
 		hc: client,
 	}
 	data := &Data{
+		ctx:     context.Background(),
 		Sources: sources,
 	}
 
@@ -94,6 +96,7 @@ func TestHTTPFileWithHeaders(t *testing.T) {
 		},
 	}
 	data := &Data{
+		ctx:     context.Background(),
 		Sources: sources,
 	}
 	expected := http.Header{
@@ -110,6 +113,7 @@ func TestHTTPFileWithHeaders(t *testing.T) {
 		"User-Agent":      {"Go-http-client/1.1"},
 	}
 	data = &Data{
+		ctx:          context.Background(),
 		Sources:      sources,
 		extraHeaders: map[string]http.Header{server.URL: expected},
 	}
diff --git a/data/datasource_merge.go b/data/datasource_merge.go
index 044302a2..136a3779 100644
--- a/data/datasource_merge.go
+++ b/data/datasource_merge.go
@@ -1,6 +1,7 @@
 package data
 
 import (
+	"context"
 	"strings"
 
 	"github.com/hairyhenderson/gomplate/v3/coll"
@@ -19,7 +20,7 @@ import (
 // the source data. To merge datasources with query strings or fragments, define
 // separate sources first and specify the alias names. HTTP headers are also not
 // supported directly.
-func (d *Data) readMerge(source *Source, args ...string) ([]byte, error) {
+func (d *Data) readMerge(ctx context.Context, source *Source, args ...string) ([]byte, error) {
 	opaque := source.URL.Opaque
 	parts := strings.Split(opaque, "|")
 	if len(parts) < 2 {
@@ -42,7 +43,7 @@ func (d *Data) readMerge(source *Source, args ...string) ([]byte, error) {
 		}
 		subSource.inherit(source)
 
-		b, err := d.readSource(subSource)
+		b, err := d.readSource(ctx, subSource)
 		if err != nil {
 			return nil, errors.Wrapf(err, "Couldn't read datasource '%s'", part)
 		}
diff --git a/data/datasource_merge_test.go b/data/datasource_merge_test.go
index a7427154..2d81c98a 100644
--- a/data/datasource_merge_test.go
+++ b/data/datasource_merge_test.go
@@ -1,6 +1,7 @@
 package data
 
 import (
+	"context"
 	"net/url"
 	"os"
 	"path/filepath"
@@ -12,6 +13,8 @@ import (
 )
 
 func TestReadMerge(t *testing.T) {
+	ctx := context.Background()
+
 	jsonContent := `{"hello": "world"}`
 	yamlContent := "hello: earth\ngoodnight: moon\n"
 	arrayContent := `["hello", "world"]`
@@ -55,38 +58,38 @@ func TestReadMerge(t *testing.T) {
 		},
 	}
 
-	actual, err := d.readMerge(source)
+	actual, err := d.readMerge(ctx, source)
 	assert.NoError(t, err)
 	assert.Equal(t, mergedContent, string(actual))
 
 	source.URL = mustParseURL("merge:bar|baz")
-	actual, err = d.readMerge(source)
+	actual, err = d.readMerge(ctx, source)
 	assert.NoError(t, err)
 	assert.Equal(t, mergedContent, string(actual))
 
 	source.URL = mustParseURL("merge:./jsonfile.json|baz")
-	actual, err = d.readMerge(source)
+	actual, err = d.readMerge(ctx, source)
 	assert.NoError(t, err)
 	assert.Equal(t, mergedContent, string(actual))
 
 	source.URL = mustParseURL("merge:file:///tmp/jsonfile.json")
-	_, err = d.readMerge(source)
+	_, err = d.readMerge(ctx, source)
 	assert.Error(t, err)
 
 	source.URL = mustParseURL("merge:bogusalias|file:///tmp/jsonfile.json")
-	_, err = d.readMerge(source)
+	_, err = d.readMerge(ctx, source)
 	assert.Error(t, err)
 
 	source.URL = mustParseURL("merge:file:///tmp/jsonfile.json|badscheme")
-	_, err = d.readMerge(source)
+	_, err = d.readMerge(ctx, source)
 	assert.Error(t, err)
 
 	source.URL = mustParseURL("merge:file:///tmp/jsonfile.json|badtype")
-	_, err = d.readMerge(source)
+	_, err = d.readMerge(ctx, source)
 	assert.Error(t, err)
 
 	source.URL = mustParseURL("merge:file:///tmp/jsonfile.json|array")
-	_, err = d.readMerge(source)
+	_, err = d.readMerge(ctx, source)
 	assert.Error(t, err)
 }
 
diff --git a/data/datasource_stdin.go b/data/datasource_stdin.go
index 8216573b..007d935b 100644
--- a/data/datasource_stdin.go
+++ b/data/datasource_stdin.go
@@ -1,6 +1,7 @@
 package data
 
 import (
+	"context"
 	"io"
 	"io/ioutil"
 
@@ -10,7 +11,7 @@ import (
 // stdin - for overriding in tests
 var stdin io.Reader
 
-func readStdin(source *Source, args ...string) ([]byte, error) {
+func readStdin(ctx context.Context, source *Source, args ...string) ([]byte, error) {
 	b, err := ioutil.ReadAll(stdin)
 	if err != nil {
 		return nil, errors.Wrapf(err, "Can't read %s", stdin)
diff --git a/data/datasource_stdin_test.go b/data/datasource_stdin_test.go
index 7fc02318..1b92914e 100644
--- a/data/datasource_stdin_test.go
+++ b/data/datasource_stdin_test.go
@@ -1,6 +1,7 @@
 package data
 
 import (
+	"context"
 	"strings"
 	"testing"
 
@@ -8,15 +9,17 @@ import (
 )
 
 func TestReadStdin(t *testing.T) {
+	ctx := context.Background()
+
 	defer func() {
 		stdin = nil
 	}()
 	stdin = strings.NewReader("foo")
-	out, err := readStdin(nil)
+	out, err := readStdin(ctx, nil)
 	assert.NoError(t, err)
 	assert.Equal(t, []byte("foo"), out)
 
 	stdin = errorReader{}
-	_, err = readStdin(nil)
+	_, err = readStdin(ctx, nil)
 	assert.Error(t, err)
 }
diff --git a/data/datasource_vault.go b/data/datasource_vault.go
index 6f3fe142..09f6dcb1 100644
--- a/data/datasource_vault.go
+++ b/data/datasource_vault.go
@@ -1,6 +1,7 @@
 package data
 
 import (
+	"context"
 	"strings"
 
 	"github.com/pkg/errors"
@@ -8,7 +9,7 @@ import (
 	"github.com/hairyhenderson/gomplate/v3/vault"
 )
 
-func readVault(source *Source, args ...string) (data []byte, err error) {
+func readVault(ctx context.Context, source *Source, args ...string) (data []byte, err error) {
 	if source.vc == nil {
 		source.vc, err = vault.New(source.URL)
 		if err != nil {
diff --git a/data/datasource_vault_test.go b/data/datasource_vault_test.go
index 47aff14d..8b470abc 100644
--- a/data/datasource_vault_test.go
+++ b/data/datasource_vault_test.go
@@ -1,6 +1,7 @@
 package data
 
 import (
+	"context"
 	"net/url"
 	"testing"
 
@@ -9,6 +10,8 @@ import (
 )
 
 func TestReadVault(t *testing.T) {
+	ctx := context.Background()
+
 	expected := "{\"value\":\"foo\"}\n"
 	server, v := vault.MockServer(200, `{"data":`+expected+`}`)
 	defer server.Close()
@@ -20,20 +23,20 @@ func TestReadVault(t *testing.T) {
 		vc:        v,
 	}
 
-	r, err := readVault(source)
+	r, err := readVault(ctx, source)
 	assert.NoError(t, err)
 	assert.Equal(t, []byte(expected), r)
 
-	r, err = readVault(source, "bar")
+	r, err = readVault(ctx, source, "bar")
 	assert.NoError(t, err)
 	assert.Equal(t, []byte(expected), r)
 
-	r, err = readVault(source, "?param=value")
+	r, err = readVault(ctx, source, "?param=value")
 	assert.NoError(t, err)
 	assert.Equal(t, []byte(expected), r)
 
 	source.URL, _ = url.Parse("vault:///secret/foo?param1=value1&param2=value2")
-	r, err = readVault(source)
+	r, err = readVault(ctx, source)
 	assert.NoError(t, err)
 	assert.Equal(t, []byte(expected), r)
 
@@ -41,7 +44,7 @@ func TestReadVault(t *testing.T) {
 	server, source.vc = vault.MockServer(200, `{"data":{"keys":`+expected+`}}`)
 	defer server.Close()
 	source.URL, _ = url.Parse("vault:///secret/foo/")
-	r, err = readVault(source)
+	r, err = readVault(ctx, source)
 	assert.NoError(t, err)
 	assert.Equal(t, []byte(expected), r)
 }