Merge pull request #234 from hairyhenderson/support-type-setting-with-url-query

Support setting MIME type with URL query string

4 files changed, 64 insertions, 8 deletions

diff --git a/data/datasource.go b/data/datasource.go
index ef693761..054e422c 100644
--- a/data/datasource.go
+++ b/data/datasource.go
@@ -118,17 +118,20 @@ func (s *Source) cleanup() {
 }
 
 // NewSource - builds a &Source
-func NewSource(alias string, URL *url.URL) (s *Source) {
+func NewSource(alias string, URL *url.URL) *Source {
 	ext := filepath.Ext(URL.Path)
 
-	s = &Source{
+	s := &Source{
 		Alias: alias,
 		URL:   URL,
 		Ext:   ext,
 	}
 
-	if ext != "" && URL.Scheme != "boltdb" {
-		mediatype := mime.TypeByExtension(ext)
+	mediatype := s.URL.Query().Get("type")
+	if mediatype == "" {
+		mediatype = mime.TypeByExtension(ext)
+	}
+	if mediatype != "" {
 		t, params, err := mime.ParseMediaType(mediatype)
 		if err != nil {
 			log.Fatal(err)
@@ -136,7 +139,10 @@ func NewSource(alias string, URL *url.URL) (s *Source) {
 		s.Type = t
 		s.Params = params
 	}
-	return
+	if s.Type == "" {
+		s.Type = plaintext
+	}
+	return s
 }
 
 // String is the method to format the flag's value, part of the flag.Value interface.
@@ -416,7 +422,6 @@ func readConsul(source *Source, args ...string) ([]byte, error) {
 	if err != nil {
 		return nil, err
 	}
-	source.Type = plaintext
 
 	return data, nil
 }
@@ -435,7 +440,6 @@ func readBoltDB(source *Source, args ...string) ([]byte, error) {
 	if err != nil {
 		return nil, err
 	}
-	source.Type = plaintext
 
 	return data, nil
 }
diff --git a/data/datasource_test.go b/data/datasource_test.go
index ac7b884a..44649ed8 100644
--- a/data/datasource_test.go
+++ b/data/datasource_test.go
@@ -42,6 +42,13 @@ func TestNewSource(t *testing.T) {
 	assert.Equal(t, ".json", s.Ext)
 
 	s = NewSource("foo", &url.URL{
+		Scheme: "file",
+		Path:   "/foo",
+	})
+	assert.Equal(t, "text/plain", s.Type)
+	assert.Equal(t, "", s.Ext)
+
+	s = NewSource("foo", &url.URL{
 		Scheme: "http",
 		Host:   "example.com",
 		Path:   "/foo.json",
@@ -56,6 +63,28 @@ func TestNewSource(t *testing.T) {
 	})
 	assert.Equal(t, "application/json", s.Type)
 	assert.Equal(t, ".json", s.Ext)
+
+	s = NewSource("foo", &url.URL{
+		Scheme:   "ftp",
+		Host:     "example.com",
+		Path:     "/foo.blarb",
+		RawQuery: "type=application/json%3Bcharset=utf-8",
+	})
+
+	assert.Equal(t, "application/json", s.Type)
+	assert.Equal(t, ".blarb", s.Ext)
+	assert.Equal(t, map[string]string{"charset": "utf-8"}, s.Params)
+
+	s = NewSource("foo", &url.URL{
+		Scheme:   "stdin",
+		Host:     "",
+		Path:     "",
+		RawQuery: "type=application/json",
+	})
+
+	assert.Equal(t, "application/json", s.Type)
+	assert.Equal(t, "", s.Ext)
+	assert.Equal(t, map[string]string{}, s.Params)
 }
 
 func TestNewData(t *testing.T) {
diff --git a/docs/content/functions/data.md b/docs/content/functions/data.md
index 0243b25b..dbbf45f2 100644
--- a/docs/content/functions/data.md
+++ b/docs/content/functions/data.md
@@ -11,7 +11,7 @@ A collection of functions that retrieve, parse, and convert structured data.
 
 Parses a given datasource (provided by the [`--datasource/-d`](#--datasource-d) argument).
 
-Currently, `file://`, `stdin://`, `http://`, `https://`, and `vault://` URLs are supported.
+Currently, `file://`, `stdin://`, `http://`, `https://`, `vault://`, and `boltdb://` URLs are supported.
 
 Currently-supported formats are JSON, YAML, TOML, and CSV.
 
@@ -53,6 +53,22 @@ $ echo 'foo' | gomplate -i '{{ include "data" }}' -d data=stdin:
 foo
 ```
 
+### Overriding the MIME type
+
+On occasion it's necessary to override the MIME type a datasource is parsed with.
+To accomplish this, gomplate supports a `type` query string parameter on
+datasource URLs. This can contain the same value as a standard
+[HTTP Content-Type](https://tools.ietf.org/html/rfc7231#section-3.1.1.1)
+header.
+
+For example, to force a file named `data.txt` to be parsed as a JSON document:
+
+```console
+$ echo '{"foo": "bar"}' > /tmp/data.txt
+$ gomplate -d data=file:///tmp/data.txt?type=application/json -i '{{ (ds "data").foo }}'
+bar
+```
+
 ### Usage with HTTP data
 
 ```console
diff --git a/test/integration/datasources_consul.bats b/test/integration/datasources_consul.bats
index 70ab14ca..8eb5da77 100644
--- a/test/integration/datasources_consul.bats
+++ b/test/integration/datasources_consul.bats
@@ -21,6 +21,13 @@ function teardown () {
   [[ "${output}" == "$BATS_TEST_DESCRIPTION" ]]
 }
 
+@test "Consul datasource works with MIME override" {
+  consul kv put foo "{\"desc\":$BATS_TEST_DESCRIPTION}"
+  gomplate -d consul=consul://?type=application/json -i '{{(datasource "consul" "foo").desc}}'
+  [ "$status" -eq 0 ]
+  [[ "${output}" == "$BATS_TEST_DESCRIPTION" ]]
+}
+
 @test "Consul datasource works with hostname in URL" {
   consul kv put foo "$BATS_TEST_DESCRIPTION"
   unset CONSUL_HTTP_ADDR