Adding directory support for file datasources

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

7 files changed, 174 insertions, 59 deletions

diff --git a/data/datasource.go b/data/datasource.go
index 8f9857b4..5d95895e 100644
--- a/data/datasource.go
+++ b/data/datasource.go
@@ -23,11 +23,6 @@ import (
 	"github.com/hairyhenderson/gomplate/vault"
 )
 
-const (
-	plaintext    = "text/plain"
-	jsonMimetype = "application/json"
-)
-
 // stdin - for overriding in tests
 var stdin io.Reader
 
@@ -40,11 +35,11 @@ func regExtension(ext, typ string) {
 
 func init() {
 	// Add some types we want to be able to handle which can be missing by default
-	regExtension(".json", "application/json")
-	regExtension(".yml", "application/yaml")
-	regExtension(".yaml", "application/yaml")
-	regExtension(".csv", "text/csv")
-	regExtension(".toml", "application/toml")
+	regExtension(".json", jsonMimetype)
+	regExtension(".yml", yamlMimetype)
+	regExtension(".yaml", yamlMimetype)
+	regExtension(".csv", csvMimetype)
+	regExtension(".toml", tomlMimetype)
 
 	sourceReaders = make(map[string]func(*Source, ...string) ([]byte, error))
 
@@ -156,7 +151,7 @@ func NewSource(alias string, URL *url.URL) (*Source, error) {
 		s.Params = params
 	}
 	if s.Type == "" {
-		s.Type = plaintext
+		s.Type = textMimetype
 	}
 	return s, nil
 }
@@ -244,15 +239,15 @@ func (d *Data) Datasource(alias string, args ...string) (interface{}, error) {
 	switch source.Type {
 	case jsonMimetype:
 		out = JSON(s)
-	case "application/array+json":
+	case jsonArrayMimetype:
 		out = JSONArray(s)
-	case "application/yaml":
+	case yamlMimetype:
 		out = YAML(s)
-	case "text/csv":
+	case csvMimetype:
 		out = CSV(s)
-	case "application/toml":
+	case tomlMimetype:
 		out = TOML(s)
-	case plaintext:
+	case textMimetype:
 		out = s
 	default:
 		return nil, errors.Errorf("Datasources of type %s not yet supported", source.Type)
@@ -309,31 +304,6 @@ func (d *Data) ReadSource(source *Source, args ...string) ([]byte, error) {
 	return nil, errors.Errorf("Datasources with scheme %s not yet supported", source.URL.Scheme)
 }
 
-func readFile(source *Source, args ...string) ([]byte, error) {
-	if source.FS == nil {
-		source.FS = vfs.OS()
-	}
-
-	p := filepath.FromSlash(source.URL.Path)
-
-	// make sure we can access the file
-	_, err := source.FS.Stat(p)
-	if err != nil {
-		return nil, errors.Wrapf(err, "Can't stat %s", p)
-	}
-
-	f, err := source.FS.OpenFile(p, os.O_RDONLY, 0)
-	if err != nil {
-		return nil, errors.Wrapf(err, "Can't open %s", p)
-	}
-
-	b, err := ioutil.ReadAll(f)
-	if err != nil {
-		return nil, errors.Wrapf(err, "Can't read %s", p)
-	}
-	return b, nil
-}
-
 func readStdin(source *Source, args ...string) ([]byte, error) {
 	if stdin == nil {
 		stdin = os.Stdin
diff --git a/data/datasource_file.go b/data/datasource_file.go
new file mode 100644
index 00000000..e66615c0
--- /dev/null
+++ b/data/datasource_file.go
@@ -0,0 +1,79 @@
+package data
+
+import (
+	"bytes"
+	"encoding/json"
+	"io/ioutil"
+	"net/url"
+	"os"
+	"path/filepath"
+	"strings"
+
+	"github.com/pkg/errors"
+
+	"github.com/blang/vfs"
+)
+
+func readFile(source *Source, args ...string) ([]byte, error) {
+	if source.FS == nil {
+		source.FS = vfs.OS()
+	}
+
+	p := filepath.FromSlash(source.URL.Path)
+
+	if len(args) == 1 {
+		parsed, err := url.Parse(args[0])
+		if err != nil {
+			return nil, err
+		}
+
+		if parsed.Path != "" {
+			p = p + "/" + parsed.Path
+		}
+	}
+
+	// make sure we can access the file
+	i, err := source.FS.Stat(p)
+	if err != nil {
+		return nil, errors.Wrapf(err, "Can't stat %s", p)
+	}
+
+	if strings.HasSuffix(p, "/") {
+		source.Type = jsonArrayMimetype
+		if i.IsDir() {
+			return readFileDir(source, p)
+		}
+		return nil, errors.Errorf("%s is not a directory", p)
+	}
+
+	f, err := source.FS.OpenFile(p, os.O_RDONLY, 0)
+	if err != nil {
+		return nil, errors.Wrapf(err, "Can't open %s", p)
+	}
+
+	b, err := ioutil.ReadAll(f)
+	if err != nil {
+		return nil, errors.Wrapf(err, "Can't read %s", p)
+	}
+	return b, nil
+}
+
+func readFileDir(source *Source, p string) ([]byte, error) {
+	names, err := source.FS.ReadDir(p)
+	if err != nil {
+		return nil, err
+	}
+	files := make([]string, len(names))
+	for i, v := range names {
+		files[i] = v.Name()
+	}
+
+	var buf bytes.Buffer
+	enc := json.NewEncoder(&buf)
+	if err := enc.Encode(files); err != nil {
+		return nil, err
+	}
+	b := buf.Bytes()
+	// chop off the newline added by the json encoder
+	return b[:len(b)-1], nil
+}
diff --git a/data/datasource_file_test.go b/data/datasource_file_test.go
new file mode 100644
index 00000000..6fc92b83
--- /dev/null
+++ b/data/datasource_file_test.go
@@ -0,0 +1,56 @@
+// +build !windows
+
+package data
+
+import (
+	"net/url"
+	"testing"
+
+	"github.com/blang/vfs"
+	"github.com/blang/vfs/memfs"
+	"github.com/stretchr/testify/assert"
+)
+
+func mustParseURL(in string) *url.URL {
+	u, _ := url.Parse(in)
+	return u
+}
+
+func TestReadFile(t *testing.T) {
+	content := []byte(`hello world`)
+	fs := memfs.Create()
+
+	_ = fs.Mkdir("/tmp", 0777)
+	f, _ := vfs.Create(fs, "/tmp/foo")
+	_, _ = f.Write(content)
+
+	_ = fs.Mkdir("/tmp/partial", 0777)
+	f, _ = vfs.Create(fs, "/tmp/partial/foo.txt")
+	_, _ = f.Write(content)
+	_, _ = vfs.Create(fs, "/tmp/partial/bar.txt")
+	_, _ = vfs.Create(fs, "/tmp/partial/baz.txt")
+
+	source, _ := NewSource("foo", mustParseURL("file:///tmp/foo"))
+	source.FS = fs
+
+	actual, err := readFile(source)
+	assert.NoError(t, err)
+	assert.Equal(t, content, actual)
+
+	source, _ = NewSource("bogus", mustParseURL("file:///bogus"))
+	source.FS = fs
+	_, err = readFile(source)
+	assert.Error(t, err)
+
+	source, _ = NewSource("partial", mustParseURL("file:///tmp/partial"))
+	source.FS = fs
+	actual, err = readFile(source, "foo.txt")
+	assert.NoError(t, err)
+	assert.Equal(t, content, actual)
+
+	source, _ = NewSource("dir", mustParseURL("file:///tmp/partial/"))
+	source.FS = fs
+	actual, err = readFile(source)
+	assert.NoError(t, err)
+	assert.Equal(t, []byte(`["bar.txt","baz.txt","foo.txt"]`), actual)
+}
diff --git a/data/datasource_test.go b/data/datasource_test.go
index ab71a060..4d3b37ac 100644
--- a/data/datasource_test.go
+++ b/data/datasource_test.go
@@ -22,7 +22,7 @@ func TestNewSource(t *testing.T) {
 		Path:   "/foo.json",
 	})
 	assert.NoError(t, err)
-	assert.Equal(t, "application/json", s.Type)
+	assert.Equal(t, jsonMimetype, s.Type)
 	assert.Equal(t, ".json", s.Ext)
 
 	s, err = NewSource("foo", &url.URL{
@@ -30,7 +30,7 @@ func TestNewSource(t *testing.T) {
 		Path:   "/foo",
 	})
 	assert.NoError(t, err)
-	assert.Equal(t, "text/plain", s.Type)
+	assert.Equal(t, textMimetype, s.Type)
 	assert.Equal(t, "", s.Ext)
 
 	s, err = NewSource("foo", &url.URL{
@@ -39,7 +39,7 @@ func TestNewSource(t *testing.T) {
 		Path:   "/foo.json",
 	})
 	assert.NoError(t, err)
-	assert.Equal(t, "application/json", s.Type)
+	assert.Equal(t, jsonMimetype, s.Type)
 	assert.Equal(t, ".json", s.Ext)
 
 	s, err = NewSource("foo", &url.URL{
@@ -48,7 +48,7 @@ func TestNewSource(t *testing.T) {
 		Path:   "/foo.json",
 	})
 	assert.NoError(t, err)
-	assert.Equal(t, "application/json", s.Type)
+	assert.Equal(t, jsonMimetype, s.Type)
 	assert.Equal(t, ".json", s.Ext)
 
 	s, err = NewSource("foo", &url.URL{
@@ -58,7 +58,7 @@ func TestNewSource(t *testing.T) {
 		RawQuery: "type=application/json%3Bcharset=utf-8",
 	})
 	assert.NoError(t, err)
-	assert.Equal(t, "application/json", s.Type)
+	assert.Equal(t, jsonMimetype, s.Type)
 	assert.Equal(t, ".blarb", s.Ext)
 	assert.Equal(t, map[string]string{"charset": "utf-8"}, s.Params)
 
@@ -69,7 +69,7 @@ func TestNewSource(t *testing.T) {
 		RawQuery: "type=application/json",
 	})
 	assert.NoError(t, err)
-	assert.Equal(t, "application/json", s.Type)
+	assert.Equal(t, jsonMimetype, s.Type)
 	assert.Equal(t, "", s.Ext)
 	assert.Equal(t, map[string]string{}, s.Params)
 }
@@ -116,7 +116,7 @@ func TestParseSourceWithAlias(t *testing.T) {
 	assert.NoError(t, err)
 	assert.Equal(t, "data", s.Alias)
 	assert.Equal(t, "file", s.URL.Scheme)
-	assert.Equal(t, "application/json", s.Type)
+	assert.Equal(t, jsonMimetype, s.Type)
 	assert.True(t, s.URL.IsAbs())
 
 	s, err = ParseSource("data=/otherdir/foo.json")
@@ -161,10 +161,10 @@ func TestDatasource(t *testing.T) {
 		assert.Equal(t, expected, actual)
 	}
 
-	test("json", "application/json", []byte(`{"hello":{"cruel":"world"}}`))
-	test("yml", "application/yaml", []byte("hello:\n  cruel: world\n"))
+	test("json", jsonMimetype, []byte(`{"hello":{"cruel":"world"}}`))
+	test("yml", yamlMimetype, []byte("hello:\n  cruel: world\n"))
 
-	d := setup("", "text/plain", nil)
+	d := setup("", textMimetype, nil)
 	actual, err := d.Datasource("foo")
 	assert.NoError(t, err)
 	assert.Equal(t, "", actual)
@@ -182,7 +182,7 @@ func TestDatasourceReachable(t *testing.T) {
 			Alias: "foo",
 			URL:   &url.URL{Scheme: "file", Path: "/tmp/" + fname},
 			Ext:   "json",
-			Type:  "application/json",
+			Type:  jsonMimetype,
 			FS:    fs,
 		},
 		"bar": {
@@ -256,7 +256,7 @@ func TestHTTPFile(t *testing.T) {
 }
 
 func TestHTTPFileWithHeaders(t *testing.T) {
-	server, client := setupHTTP(200, "application/json", "")
+	server, client := setupHTTP(200, jsonMimetype, "")
 	defer server.Close()
 
 	sources := make(map[string]*Source)
@@ -294,7 +294,7 @@ func TestParseHeaderArgs(t *testing.T) {
 	}
 	expected := map[string]http.Header{
 		"foo": {
-			"Accept": {"application/json"},
+			"Accept": {jsonMimetype},
 		},
 		"bar": {
 			"Authorization": {"Bearer supersecret"},
@@ -319,7 +319,7 @@ func TestParseHeaderArgs(t *testing.T) {
 	}
 	expected = map[string]http.Header{
 		"foo": {
-			"Accept": {"application/json"},
+			"Accept": {jsonMimetype},
 			"Foo":    {"bar", "baz", "qux"},
 		},
 		"bar": {
@@ -345,7 +345,7 @@ func TestInclude(t *testing.T) {
 			Alias: "foo",
 			URL:   &url.URL{Scheme: "file", Path: "/tmp/" + fname},
 			Ext:   ext,
-			Type:  "text/plain",
+			Type:  textMimetype,
 			FS:    fs,
 		},
 	}
diff --git a/data/datasource_vault.go b/data/datasource_vault.go
index 9033a804..46cd43c0 100644
--- a/data/datasource_vault.go
+++ b/data/datasource_vault.go
@@ -46,11 +46,11 @@ func readVault(source *Source, args ...string) ([]byte, error) {
 
 	var data []byte
 
-	source.Type = "application/json"
+	source.Type = jsonMimetype
 	if len(params) > 0 {
 		data, err = source.VC.Write(p, params)
 	} else if strings.HasSuffix(p, "/") {
-		source.Type = "application/array+json"
+		source.Type = jsonArrayMimetype
 		data, err = source.VC.List(p)
 	} else {
 		data, err = source.VC.Read(p)
diff --git a/data/datasource_vault_test.go b/data/datasource_vault_test.go
index fc7acec1..368d0ef3 100644
--- a/data/datasource_vault_test.go
+++ b/data/datasource_vault_test.go
@@ -17,7 +17,7 @@ func TestReadVault(t *testing.T) {
 		Alias: "foo",
 		URL:   &url.URL{Scheme: "vault", Path: "/secret/foo"},
 		Ext:   "",
-		Type:  "text/plain",
+		Type:  textMimetype,
 		VC:    v,
 	}
 
diff --git a/data/mimetypes.go b/data/mimetypes.go
new file mode 100644
index 00000000..99e952ce
--- /dev/null
+++ b/data/mimetypes.go
@@ -0,0 +1,10 @@
+package data
+
+const (
+	textMimetype      = "text/plain"
+	csvMimetype       = "text/csv"
+	jsonMimetype      = "application/json"
+	jsonArrayMimetype = "application/array+json"
+	tomlMimetype      = "application/toml"
+	yamlMimetype      = "application/yaml"
+)