From 4ba02d7a202d20bed54c35471842bd255fe46e25 Mon Sep 17 00:00:00 2001
From: Shunsuke Suzuki <suzuki-shunsuke@users.noreply.github.com>
Date: Sun, 19 May 2024 07:07:15 +0900
Subject: [PATCH] Fix a bug that `aqua g -i` fails if aqua.yaml doesn't have
 the field `packages` (#2904)

* fix(generate): fix a bug that aqua g -i fails if aqua.yaml doesn't have the field packages

* fix: remove useless codes

* refactor: refactoring
---
 pkg/controller/generate/output/insert.go | 48 ++++++++++++++++++++----
 1 file changed, 40 insertions(+), 8 deletions(-)

diff --git a/pkg/controller/generate/output/insert.go b/pkg/controller/generate/output/insert.go
index 45028d6dd..893a28902 100644
--- a/pkg/controller/generate/output/insert.go
+++ b/pkg/controller/generate/output/insert.go
@@ -3,6 +3,7 @@ package output
 import (
 	"errors"
 	"fmt"
+	"strings"
 
 	wast "github.com/aquaproj/aqua/v2/pkg/ast"
 	"github.com/aquaproj/aqua/v2/pkg/config/aqua"
@@ -29,8 +30,8 @@ func (o *Outputter) generateInsert(cfgFilePath string, pkgs []*aqua.Package) err
 			"num_of_docs": len(file.Docs),
 		})
 	}
-
-	if err := updateASTFile(file.Docs[0].Body, pkgs); err != nil {
+	s, err := o.getAppendedTxt(cfgFilePath, file, pkgs)
+	if err != nil {
 		return err
 	}
 
@@ -38,12 +39,28 @@ func (o *Outputter) generateInsert(cfgFilePath string, pkgs []*aqua.Package) err
 	if err != nil {
 		return fmt.Errorf("get configuration file stat: %w", err)
 	}
-	if err := afero.WriteFile(o.fs, cfgFilePath, []byte(file.String()), stat.Mode()); err != nil {
+	if err := afero.WriteFile(o.fs, cfgFilePath, []byte(s), stat.Mode()); err != nil {
 		return fmt.Errorf("write the configuration file: %w", err)
 	}
 	return nil
 }
 
+func (o *Outputter) getAppendedTxt(cfgFilePath string, file *ast.File, pkgs []*aqua.Package) (string, error) {
+	body := file.Docs[0].Body
+	values, err := wast.FindMappingValueFromNode(body, "packages")
+	if err != nil {
+		return "", fmt.Errorf(`find a mapping value node "packages": %w`, err)
+	}
+	if values == nil {
+		return o.appendPkgsTxt(cfgFilePath, pkgs)
+	}
+
+	if err := updateASTFile(values, pkgs); err != nil {
+		return "", err
+	}
+	return file.String(), nil
+}
+
 func appendPkgsNode(mapValue *ast.MappingValueNode, node ast.Node) error {
 	switch mapValue.Value.Type() {
 	case ast.NullType:
@@ -59,16 +76,31 @@ func appendPkgsNode(mapValue *ast.MappingValueNode, node ast.Node) error {
 	}
 }
 
-func updateASTFile(body ast.Node, pkgs []*aqua.Package) error {
+func updateASTFile(values *ast.MappingValueNode, pkgs []*aqua.Package) error {
 	node, err := yaml.ValueToNode(pkgs)
 	if err != nil {
 		return fmt.Errorf("convert packages to node: %w", err)
 	}
 
-	values, err := wast.FindMappingValueFromNode(body, "packages")
+	return appendPkgsNode(values, node)
+}
+
+func (o *Outputter) appendPkgsTxt(cfgFilePath string, pkgs []*aqua.Package) (string, error) {
+	a, err := yaml.Marshal(struct {
+		Packages []*aqua.Package `yaml:"packages"`
+	}{
+		Packages: pkgs,
+	})
 	if err != nil {
-		return fmt.Errorf(`find a mapping value node "packages": %w`, err)
+		return "", fmt.Errorf("marshal packages: %w", err)
 	}
-
-	return appendPkgsNode(values, node)
+	b, err := afero.ReadFile(o.fs, cfgFilePath)
+	if err != nil {
+		return "", fmt.Errorf("read a configuration file: %w", err)
+	}
+	sb := string(b)
+	if !strings.HasSuffix(sb, "\n") {
+		sb += "\n"
+	}
+	return sb + string(a), nil
 }