Support setting MIME type with URL query string

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

2 files changed, 40 insertions, 7 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) {