add.go

// Copyright 2019 The CUE Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package cmd

import (
	"bytes"
	"errors"
	"fmt"
	"html/template"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"

	"github.com/spf13/cobra"

	"cuelang.org/go/cue"
	"cuelang.org/go/cue/build"
	"cuelang.org/go/cue/format"
	"cuelang.org/go/cue/load"
	"cuelang.org/go/cue/parser"
	"cuelang.org/go/internal"
)

func newAddCmd(c *Command) *cobra.Command {
	cmd := &cobra.Command{
		// TODO: this command is still experimental, don't show it in
		// the documentation just yet.
		Hidden: true,

		Use:   "add <glob> [--list]",
		Short: "bulk append to CUE files",
		Long: `Append a common snippet of CUE to many files and commit atomically.
`,
		RunE: mkRunE(c, runAdd),
	}

	f := cmd.Flags()
	f.Bool(string(flagList), false,
		"text executed as Go template with instance info")
	f.BoolP(string(flagDryrun), "n", false,
		"only run simulation")
	f.StringP(string(flagPackage), "p", "", "package to append to")

	return cmd
}

func runAdd(cmd *Command, args []string) (err error) {
	return doAdd(cmd, args)
}

func doAdd(cmd *Command, args []string) (err error) {
	// Offsets at which to restore original files, if any, if any of the
	// appends fail.
	// Ideally this is placed below where it is used, but we want to make
	// absolutely sure that the error variable used in defer is the named
	// returned value and not some shadowed value.

	originals := []originalFile{}
	defer func() {
		if err != nil {
			restoreOriginals(cmd, originals)
		}
	}()

	// build instance cache
	builds := map[string]*build.Instance{}

	getBuild := func(path string) *build.Instance {
		if b, ok := builds[path]; ok {
			return b
		}
		b := load.Instances([]string{path}, nil)[0]
		builds[path] = b
		return b
	}

	// determine file set.

	todo := []*fileInfo{}

	done := map[string]bool{}

	for _, arg := range args {
		dir, base := filepath.Split(arg)
		dir = filepath.Clean(dir)
		matches, err := filepath.Glob(dir)
		if err != nil {
			return err
		}
		for _, m := range matches {
			if fi, err := os.Stat(m); err != nil || !fi.IsDir() {
				continue
			}
			file := filepath.Join(m, base)
			if done[file] {
				continue
			}
			if s := filepath.ToSlash(file); strings.HasPrefix(s, "pkg/") || strings.Contains(s, "/pkg/") {
				continue
			}
			done[file] = true
			fi, err := initFile(cmd, file, getBuild)
			if err != nil {
				return err
			}
			todo = append(todo, fi)
			b := fi.build
			if flagList.Bool(cmd) && (b == nil) {
				return fmt.Errorf("instance info not available for %s", fi.filename)
			}
		}
	}

	// Read text to be appended.
	text, err := ioutil.ReadAll(cmd.InOrStdin())
	if err != nil {
		return err
	}

	var tmpl *template.Template
	if flagList.Bool(cmd) {
		tmpl, err = template.New("append").Parse(string(text))
		if err != nil {
			return err
		}
	}

	for _, fi := range todo {
		if tmpl == nil {
			fi.contents.Write(text)
			continue
		}
		if err := tmpl.Execute(fi.contents, fi.build); err != nil {
			return err
		}
	}

	if flagDryrun.Bool(cmd) {
		stdout := cmd.OutOrStdout()
		for _, fi := range todo {
			fmt.Fprintln(stdout, "---", fi.filename)
			_, _ = io.Copy(stdout, fi.contents)
		}
		return nil
	}

	// All verified. Execute the todo plan
	for _, fi := range todo {
		fo, err := fi.appendAndCheck()
		if err != nil {
			return err
		}
		originals = append(originals, fo)
	}

	// Verify resulting builds
	for _, fi := range todo {
		builds = map[string]*build.Instance{}

		b := getBuild(fi.buildArg)
		if b.Err != nil {
			return b.Err
		}
		i := cue.Build([]*build.Instance{b})[0]
		if i.Err != nil {
			return i.Err
		}
		if err := i.Value().Validate(); err != nil {
			return i.Err
		}
	}

	return nil
}

type originalFile struct {
	filename string
	contents []byte
}

func restoreOriginals(cmd *Command, originals []originalFile) {
	for _, fo := range originals {
		if err := fo.restore(); err != nil {
			fmt.Fprintln(cmd.Stderr(), "Error restoring file: ", err)
		}
	}
}

func (fo *originalFile) restore() error {
	if len(fo.contents) == 0 {
		return os.Remove(fo.filename)
	}
	return ioutil.WriteFile(fo.filename, fo.contents, 0644)
}

type fileInfo struct {
	filename string
	buildArg string
	contents *bytes.Buffer
	build    *build.Instance
}

func initFile(cmd *Command, file string, getBuild func(path string) *build.Instance) (todo *fileInfo, err error) {
	defer func() {
		if err != nil {
			err = fmt.Errorf("init file: %v", err)
		}
	}()
	dir := filepath.Dir(file)
	todo = &fileInfo{file, dir, &bytes.Buffer{}, nil}

	if fi, err := os.Stat(file); err != nil {
		if !os.IsNotExist(err) {
			return nil, err
		}
		// File does not exist
		b := getBuild(dir)
		todo.build = b
		pkg := flagPackage.String(cmd)
		if pkg != "" {
			// TODO: do something more intelligent once the package name is
			// computed on a module basis, even for empty directories.
			b.PkgName = pkg
			b.Err = nil
		} else {
			pkg = b.PkgName
		}
		if pkg == "" {
			return nil, errors.New("must specify package using -p for new files")
		}
		todo.buildArg = file
		fmt.Fprintf(todo.contents, "package %s\n\n", pkg)
	} else {
		if fi.IsDir() {
			return nil, fmt.Errorf("cannot append to directory %s", file)
		}

		f, err := parser.ParseFile(file, nil)
		if err != nil {
			return nil, err
		}
		if _, pkgName, _ := internal.PackageInfo(f); pkgName != "" {
			if pkg := flagPackage.String(cmd); pkg != "" && pkgName != pkg {
				return nil, fmt.Errorf("package mismatch (%s vs %s) for file %s", pkgName, pkg, file)
			}
			todo.build = getBuild(dir)
		} else {
			if pkg := flagPackage.String(cmd); pkg != "" {
				return nil, fmt.Errorf("file %s has no package clause but package %s requested", file, pkg)
			}
			todo.build = getBuild(file)
			todo.buildArg = file
		}
	}
	return todo, nil
}

func (fi *fileInfo) appendAndCheck() (fo originalFile, err error) {
	// Read original file
	b, err := ioutil.ReadFile(fi.filename)
	if err == nil {
		fo.filename = fi.filename
		fo.contents = b
	} else if !os.IsNotExist(err) {
		return originalFile{}, err
	}

	if !bytes.HasSuffix(b, []byte("\n")) {
		b = append(b, '\n')
	}
	b = append(b, fi.contents.Bytes()...)

	b, err = format.Source(b)
	if err != nil {
		return originalFile{}, err
	}

	if err = ioutil.WriteFile(fi.filename, b, 0644); err != nil {
		// Just in case, attempt to restore original file.
		_ = fo.restore()
		return originalFile{}, err
	}

	return fo, nil
}