support directory walking and multiple documents per file

This extends the directory support to include recursive directories for further organization. Additionally, this adds support for multiple yaml documents in a single file which allows for simple organization without cluttering folders. Signed-off-by: Mike Mason <[email protected]>
infratographer · Apr 15, 2024 · 6c27a43 · 6c27a43
1 parent 9417171
commit 6c27a43
Show file tree

Hide file tree

Showing 3 changed files with 101 additions and 57 deletions.
diff --git a/cmd/schema.go b/cmd/schema.go
@@ -3,7 +3,7 @@ package cmd
 import (
 	"context"
 	"fmt"
-	"os"
+	"io/fs"
 	"path/filepath"
 	"strings"
 
@@ -74,17 +74,28 @@ func writeSchema(_ context.Context, dryRun bool, cfg *config.AppConfig) {
 
 	if viper.GetBool("mermaid") || viper.GetBool("mermaid-markdown") {
 		if cfg.SpiceDB.PolicyDir != "" {
-			files, err := os.ReadDir(cfg.SpiceDB.PolicyDir)
-			if err != nil {
-				logger.Fatalw("failed to read policy files from directory", "error", err)
-			}
+			var filePaths []string
 
-			filePaths := make([]string, 0, len(files))
+			err := filepath.WalkDir(cfg.SpiceDB.PolicyDir, func(path string, entry fs.DirEntry, err error) error {
+				if err != nil {
+					return err
+				}
 
-			for _, file := range files {
-				if !file.IsDir() && (strings.EqualFold(filepath.Ext(file.Name()), ".yml") || strings.EqualFold(filepath.Ext(file.Name()), ".yaml")) {
-					filePaths = append(filePaths, cfg.SpiceDB.PolicyDir+"/"+file.Name())
+				if entry.IsDir() {
+					return nil
 				}
+
+				ext := filepath.Ext(entry.Name())
+
+				if strings.EqualFold(ext, ".yml") || strings.EqualFold(ext, ".yaml") {
+					filePaths = append(filePaths, path)
+				}
+
+				return nil
+			})
+
+			if err != nil {
+				logger.Fatalw("failed to read policy files from directory", "error", err)
 			}
 
 			outputPolicyMermaid(filePaths, viper.GetBool("mermaid-markdown"))

diff --git a/cmd/schema_mermaid.go b/cmd/schema_mermaid.go
@@ -3,11 +3,8 @@ package cmd
 import (
 	"bytes"
 	"fmt"
-	"os"
 	"text/template"
 
-	"gopkg.in/yaml.v3"
-
 	"go.infratographer.com/permissions-api/internal/iapl"
 )
 
@@ -70,23 +67,15 @@ type mermaidContext struct {
 }
 
 func outputPolicyMermaid(filePaths []string, markdown bool) {
-	policy := iapl.PolicyDocument{}
+	var (
+		policy iapl.PolicyDocument
+		err    error
+	)
 
 	if len(filePaths) > 0 {
-		for _, filePath := range filePaths {
-			file, err := os.Open(filePath)
-			if err != nil {
-				logger.Fatalw("failed to open policy document file", "error", err)
-			}
-			defer file.Close()
-
-			var filePolicy iapl.PolicyDocument
-
-			if err := yaml.NewDecoder(file).Decode(&filePolicy); err != nil {
-				logger.Fatalw("failed to open policy document file", "error", err)
-			}
-
-			policy = policy.MergeWithPolicyDocument(filePolicy)
+		policy, err = iapl.LoadPolicyDocumentFromFiles(filePaths...)
+		if err != nil {
+			logger.Fatalw("failed to load policy documents", "error", err)
 		}
 	} else {
 		policy = iapl.DefaultPolicyDocument()

diff --git a/internal/iapl/policy.go b/internal/iapl/policy.go
@@ -1,7 +1,10 @@
 package iapl
 
 import (
+	"errors"
 	"fmt"
+	"io"
+	"io/fs"
 	"os"
 	"path/filepath"
 	"strings"
@@ -151,62 +154,103 @@ func (p PolicyDocument) MergeWithPolicyDocument(other PolicyDocument) PolicyDocu
 	return p
 }
 
-// NewPolicyFromFile reads the provided file path and returns a new Policy.
-func NewPolicyFromFile(filePath string) (Policy, error) {
+func loadPolicyDocumentFromFile(filePath string) (PolicyDocument, error) {
 	file, err := os.Open(filePath)
 	if err != nil {
-		return nil, err
+		return PolicyDocument{}, fmt.Errorf("%s: %w", filePath, err)
 	}
 
-	var policy PolicyDocument
+	defer file.Close()
 
-	if err := yaml.NewDecoder(file).Decode(&policy); err != nil {
-		return nil, err
+	var (
+		finalPolicyDocument = PolicyDocument{}
+		decoder             = yaml.NewDecoder(file)
+		documentIndex       int
+	)
+
+	for {
+		var policyDocument PolicyDocument
+
+		if err = decoder.Decode(&policyDocument); err != nil {
+			if !errors.Is(err, io.EOF) {
+				return PolicyDocument{}, fmt.Errorf("%s document %d: %w", filePath, documentIndex, err)
+			}
+
+			break
+		}
+
+		if finalPolicyDocument.RBAC != nil && policyDocument.RBAC != nil {
+			return PolicyDocument{}, fmt.Errorf("%s document %d: %w", filePath, documentIndex, ErrorDuplicateRBACDefinition)
+		}
+
+		finalPolicyDocument = finalPolicyDocument.MergeWithPolicyDocument(policyDocument)
+
+		documentIndex++
 	}
 
-	return NewPolicy(policy), nil
+	return finalPolicyDocument, nil
 }
 
-// NewPolicyFromFiles reads the provided file paths, merges them, and returns a new Policy.
-func NewPolicyFromFiles(filePaths []string) (Policy, error) {
-	mergedPolicy := PolicyDocument{}
+// LoadPolicyDocumentFromFiles loads all policy documents in the order provided and returns a merged PolicyDocument.
+func LoadPolicyDocumentFromFiles(filePaths ...string) (PolicyDocument, error) {
+	var policyDocument PolicyDocument
 
 	for _, filePath := range filePaths {
-		file, err := os.Open(filePath)
+		filePolicyDocument, err := loadPolicyDocumentFromFile(filePath)
 		if err != nil {
-			return nil, err
+			return PolicyDocument{}, err
 		}
-		defer file.Close()
 
-		var filePolicy PolicyDocument
-
-		if err := yaml.NewDecoder(file).Decode(&filePolicy); err != nil {
-			return nil, err
-		}
+		policyDocument = policyDocument.MergeWithPolicyDocument(filePolicyDocument)
+	}
 
-		if mergedPolicy.RBAC != nil && filePolicy.RBAC != nil {
-			return nil, ErrorDuplicateRBACDefinition
-		}
+	return policyDocument, nil
+}
 
-		mergedPolicy = mergedPolicy.MergeWithPolicyDocument(filePolicy)
+// NewPolicyFromFile reads the provided file path and returns a new Policy.
+func NewPolicyFromFile(filePath string) (Policy, error) {
+	policyDocument, err := LoadPolicyDocumentFromFiles(filePath)
+	if err != nil {
+		return nil, err
 	}
 
-	return NewPolicy(mergedPolicy), nil
+	return NewPolicy(policyDocument), nil
 }
 
-// NewPolicyFromDirectory reads the provided directory path, reads all files in the directory, merges them, and returns a new Policy.
-func NewPolicyFromDirectory(directoryPath string) (Policy, error) {
-	files, err := os.ReadDir(directoryPath)
+// NewPolicyFromFiles reads the provided file paths, merges them, and returns a new Policy.
+func NewPolicyFromFiles(filePaths []string) (Policy, error) {
+	policyDocument, err := LoadPolicyDocumentFromFiles(filePaths...)
 	if err != nil {
 		return nil, err
 	}
 
-	filePaths := make([]string, 0, len(files))
+	return NewPolicy(policyDocument), nil
+}
+
+// NewPolicyFromDirectory reads the provided directory path, reads all files in the directory, merges them, and returns a new Policy.
+func NewPolicyFromDirectory(directoryPath string) (Policy, error) {
+	var filePaths []string
 
-	for _, file := range files {
-		if !file.IsDir() && (strings.EqualFold(filepath.Ext(file.Name()), ".yml") || strings.EqualFold(filepath.Ext(file.Name()), ".yaml")) {
-			filePaths = append(filePaths, directoryPath+"/"+file.Name())
+	err := filepath.WalkDir(directoryPath, func(path string, entry fs.DirEntry, err error) error {
+		if err != nil {
+			return err
+		}
+
+		if entry.IsDir() {
+			return nil
+		}
+
+		ext := filepath.Ext(entry.Name())
+
+		if strings.EqualFold(ext, ".yml") || strings.EqualFold(ext, ".yaml") {
+			filePaths = append(filePaths, path)
 		}
+
+		return nil
+	})
+
+	if err != nil {
+		return nil, err
 	}
 
 	return NewPolicyFromFiles(filePaths)