Skip to content

Commit daea696

Browse files
ethanndicksonethanndickson
committed
idempotent install
1 parent 802129c commit daea696

File tree

8 files changed

+278
-108
lines changed

8 files changed

+278
-108
lines changed

‎completion.go

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,8 @@
11
package serpent
22

3-
import "strings"
3+
import (
4+
"github.com/spf13/pflag"
5+
)
46

57
// CompletionModeEnv is a special environment variable that is
68
// set when the command is being run in completion mode.
@@ -12,17 +14,18 @@ func (inv *Invocation) IsCompletionMode() bool {
1214
return ok
1315
}
1416

15-
// DefaultCompletionHandler returns a handler that prints all
16-
// known flags and subcommands that haven't already been set to valid values.
17+
// DefaultCompletionHandler is a handler that prints all known flags and
18+
// subcommands that haven't been exhaustively set.
1719
func DefaultCompletionHandler(inv *Invocation) []string {
1820
var allResps []string
1921
for _, cmd := range inv.Command.Children {
2022
allResps = append(allResps, cmd.Name())
2123
}
2224
for _, opt := range inv.Command.Options {
25+
_, isSlice := opt.Value.(pflag.SliceValue)
2326
if opt.ValueSource == ValueSourceNone ||
2427
opt.ValueSource == ValueSourceDefault ||
25-
strings.Contains(opt.Value.Type(), "array") {
28+
isSlice {
2629
allResps = append(allResps, "--"+opt.Flag)
2730
}
2831
}

‎completion/all.go

Lines changed: 134 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,11 @@
11
package completion
22

33
import (
4+
"bytes"
5+
"errors"
46
"fmt"
57
"io"
8+
"io/fs"
69
"os"
710
"os/user"
811
"path/filepath"
@@ -13,11 +16,16 @@ import (
1316
"github.com/coder/serpent"
1417
)
1518

19+
const (
20+
completionStartTemplate = `# ============ BEGIN {{.Name}} COMPLETION ============`
21+
completionEndTemplate = `# ============ END {{.Name}} COMPLETION ==============`
22+
)
23+
1624
type Shell interface {
1725
Name() string
1826
InstallPath() (string, error)
19-
UsesOwnFile() bool
2027
WriteCompletion(io.Writer) error
28+
ProgramName() string
2129
}
2230

2331
const (
@@ -77,56 +85,154 @@ func DetectUserShell(programName string) (Shell, error) {
7785
return nil, fmt.Errorf("default shell not found")
7886
}
7987

80-
func generateCompletion(
81-
scriptTemplate string,
82-
) func(io.Writer, string) error {
83-
return func(w io.Writer, programName string) error {
84-
tmpl, err := template.New("script").Parse(scriptTemplate)
85-
if err != nil {
86-
return fmt.Errorf("parse template: %w", err)
87-
}
88-
89-
err = tmpl.Execute(
90-
w,
91-
map[string]string{
92-
"Name": programName,
93-
},
94-
)
95-
if err != nil {
96-
return fmt.Errorf("execute template: %w", err)
97-
}
88+
func configTemplateWriter(
89+
w io.Writer,
90+
cfgTemplate string,
91+
programName string,
92+
) error {
93+
tmpl, err := template.New("script").Parse(cfgTemplate)
94+
if err != nil {
95+
return fmt.Errorf("parse template: %w", err)
96+
}
9897

99-
return nil
98+
err = tmpl.Execute(
99+
w,
100+
map[string]string{
101+
"Name": programName,
102+
},
103+
)
104+
if err != nil {
105+
return fmt.Errorf("execute template: %w", err)
100106
}
107+
108+
return nil
101109
}
102110

103111
func InstallShellCompletion(shell Shell) error {
104112
path, err := shell.InstallPath()
105113
if err != nil {
106114
return fmt.Errorf("get install path: %w", err)
107115
}
116+
var headerBuf bytes.Buffer
117+
err = configTemplateWriter(&headerBuf, completionStartTemplate, shell.ProgramName())
118+
if err != nil {
119+
return fmt.Errorf("generate header: %w", err)
120+
}
121+
122+
var footerBytes bytes.Buffer
123+
err = configTemplateWriter(&footerBytes, completionEndTemplate, shell.ProgramName())
124+
if err != nil {
125+
return fmt.Errorf("generate footer: %w", err)
126+
}
108127

109128
err = os.MkdirAll(filepath.Dir(path), 0o755)
110129
if err != nil {
111130
return fmt.Errorf("create directories: %w", err)
112131
}
113132

114-
if shell.UsesOwnFile() {
115-
err := os.WriteFile(path, nil, 0o644)
133+
f, err := os.ReadFile(path)
134+
if err != nil && !errors.Is(err, fs.ErrNotExist) {
135+
return fmt.Errorf("read ssh config failed: %w", err)
136+
}
137+
138+
before, after, err := templateConfigSplit(headerBuf.Bytes(), footerBytes.Bytes(), f)
139+
if err != nil {
140+
return err
141+
}
142+
143+
outBuf := bytes.Buffer{}
144+
_, _ = outBuf.Write(before)
145+
if len(before) > 0 {
146+
_, _ = outBuf.Write([]byte("\n"))
147+
}
148+
_, _ = outBuf.Write(headerBuf.Bytes())
149+
err = shell.WriteCompletion(&outBuf)
150+
if err != nil {
151+
return fmt.Errorf("generate completion: %w", err)
152+
}
153+
_, _ = outBuf.Write(footerBytes.Bytes())
154+
_, _ = outBuf.Write([]byte("\n"))
155+
_, _ = outBuf.Write(after)
156+
157+
err = writeWithTempFileAndMove(path, &outBuf)
158+
if err != nil {
159+
return fmt.Errorf("write completion: %w", err)
160+
}
161+
162+
return nil
163+
}
164+
165+
func templateConfigSplit(header, footer, data []byte) (before, after []byte, err error) {
166+
startCount := bytes.Count(data, header)
167+
endCount := bytes.Count(data, footer)
168+
if startCount > 1 || endCount > 1 {
169+
return nil, nil, fmt.Errorf("Malformed config file: multiple config sections")
170+
}
171+
172+
startIndex := bytes.Index(data, header)
173+
endIndex := bytes.Index(data, footer)
174+
if startIndex == -1 && endIndex != -1 {
175+
return data, nil, fmt.Errorf("Malformed config file: missing completion header")
176+
}
177+
if startIndex != -1 && endIndex == -1 {
178+
return data, nil, fmt.Errorf("Malformed config file: missing completion footer")
179+
}
180+
if startIndex != -1 && endIndex != -1 {
181+
if startIndex > endIndex {
182+
return data, nil, fmt.Errorf("Malformed config file: completion header after footer")
183+
}
184+
// Include leading and trailing newline, if present
185+
start := startIndex
186+
if start > 0 {
187+
start--
188+
}
189+
end := endIndex + len(footer)
190+
if end < len(data) {
191+
end++
192+
}
193+
return data[:start], data[end:], nil
194+
}
195+
return data, nil, nil
196+
}
197+
198+
// writeWithTempFileAndMove writes to a temporary file in the same
199+
// directory as path and renames the temp file to the file provided in
200+
// path. This ensure we avoid trashing the file we are writing due to
201+
// unforeseen circumstance like filesystem full, command killed, etc.
202+
func writeWithTempFileAndMove(path string, r io.Reader) (err error) {
203+
dir := filepath.Dir(path)
204+
name := filepath.Base(path)
205+
206+
if err = os.MkdirAll(dir, 0o700); err != nil {
207+
return fmt.Errorf("create directory: %w", err)
208+
}
209+
210+
// Create a tempfile in the same directory for ensuring write
211+
// operation does not fail.
212+
f, err := os.CreateTemp(dir, fmt.Sprintf(".%s.", name))
213+
if err != nil {
214+
return fmt.Errorf("create temp file failed: %w", err)
215+
}
216+
defer func() {
116217
if err != nil {
117-
return fmt.Errorf("create file: %w", err)
218+
_ = os.Remove(f.Name()) // Cleanup in case a step failed.
118219
}
220+
}()
221+
222+
_, err = io.Copy(f, r)
223+
if err != nil {
224+
_ = f.Close()
225+
return fmt.Errorf("write temp file failed: %w", err)
119226
}
120227

121-
f, err := os.OpenFile(path, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0o644)
228+
err = f.Close()
122229
if err != nil {
123-
return fmt.Errorf("open file for appending: %w", err)
230+
return fmt.Errorf("close temp file failed: %w", err)
124231
}
125-
defer f.Close()
126232

127-
err = shell.WriteCompletion(f)
233+
err = os.Rename(f.Name(), path)
128234
if err != nil {
129-
return fmt.Errorf("write completion script: %w", err)
235+
return fmt.Errorf("rename temp file failed: %w", err)
130236
}
131237

132238
return nil

‎completion/bash.go

Lines changed: 6 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -23,11 +23,6 @@ func (b *bash) Name() string {
2323
return "bash"
2424
}
2525

26-
// UsesOwnFile implements Shell.
27-
func (b *bash) UsesOwnFile() bool {
28-
return false
29-
}
30-
3126
// InstallPath implements Shell.
3227
func (b *bash) InstallPath() (string, error) {
3328
homeDir, err := home.Dir()
@@ -42,12 +37,15 @@ func (b *bash) InstallPath() (string, error) {
4237

4338
// WriteCompletion implements Shell.
4439
func (b *bash) WriteCompletion(w io.Writer) error {
45-
return generateCompletion(bashCompletionTemplate)(w, b.programName)
40+
return configTemplateWriter(w, bashCompletionTemplate, b.programName)
4641
}
4742

48-
const bashCompletionTemplate = `
43+
// ProgramName implements Shell.
44+
func (b *bash) ProgramName() string {
45+
return b.programName
46+
}
4947

50-
# === BEGIN {{.Name}} COMPLETION ===
48+
const bashCompletionTemplate = `
5149
_generate_{{.Name}}_completions() {
5250
# Capture the line excluding the command, and everything after the current word
5351
local args=("${COMP_WORDS[@]:1:COMP_CWORD}")
@@ -65,6 +63,4 @@ _generate_{{.Name}}_completions() {
6563
}
6664
# Setup Bash to use the function for completions for '{{.Name}}'
6765
complete -F _generate_{{.Name}}_completions {{.Name}}
68-
# === END {{.Name}} COMPLETION ===
69-
7066
`

‎completion/fish.go

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -18,11 +18,6 @@ func Fish(goos string, programName string) Shell {
1818
return &fish{goos: goos, programName: programName}
1919
}
2020

21-
// UsesOwnFile implements Shell.
22-
func (f *fish) UsesOwnFile() bool {
23-
return true
24-
}
25-
2621
// Name implements Shell.
2722
func (f *fish) Name() string {
2823
return "fish"
@@ -39,7 +34,12 @@ func (f *fish) InstallPath() (string, error) {
3934

4035
// WriteCompletion implements Shell.
4136
func (f *fish) WriteCompletion(w io.Writer) error {
42-
return generateCompletion(fishCompletionTemplate)(w, f.programName)
37+
return configTemplateWriter(w, fishCompletionTemplate, f.programName)
38+
}
39+
40+
// ProgramName implements Shell.
41+
func (f *fish) ProgramName() string {
42+
return f.programName
4343
}
4444

4545
const fishCompletionTemplate = `

‎completion/powershell.go

Lines changed: 7 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,8 @@ type powershell struct {
1111
programName string
1212
}
1313

14+
var _ Shell = &powershell{}
15+
1416
// Name implements Shell.
1517
func (p *powershell) Name() string {
1618
return "powershell"
@@ -20,11 +22,6 @@ func Powershell(goos string, programName string) Shell {
2022
return &powershell{goos: goos, programName: programName}
2123
}
2224

23-
// UsesOwnFile implements Shell.
24-
func (p *powershell) UsesOwnFile() bool {
25-
return false
26-
}
27-
2825
// InstallPath implements Shell.
2926
func (p *powershell) InstallPath() (string, error) {
3027
var (
@@ -45,14 +42,15 @@ func (p *powershell) InstallPath() (string, error) {
4542

4643
// WriteCompletion implements Shell.
4744
func (p *powershell) WriteCompletion(w io.Writer) error {
48-
return generateCompletion(pshCompletionTemplate)(w, p.programName)
45+
return configTemplateWriter(w, pshCompletionTemplate, p.programName)
4946
}
5047

51-
var _ Shell = &powershell{}
48+
// ProgramName implements Shell.
49+
func (p *powershell) ProgramName() string {
50+
return p.programName
51+
}
5252

5353
const pshCompletionTemplate = `
54-
55-
# === BEGIN {{.Name}} COMPLETION ===
5654
# Escaping output sourced from:
5755
# https://github.com/spf13/cobra/blob/e94f6d0dd9a5e5738dca6bce03c4b1207ffbc0ec/powershell_completions.go#L47
5856
filter _{{.Name}}_escapeStringWithSpecialChars {
@@ -89,6 +87,4 @@ $_{{.Name}}_completions = {
8987
rm env:COMPLETION_MODE
9088
}
9189
Register-ArgumentCompleter -CommandName {{.Name}} -ScriptBlock $_{{.Name}}_completions
92-
# === END {{.Name}} COMPLETION ===
93-
9490
`

‎completion/zsh.go

Lines changed: 6 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -23,11 +23,6 @@ func (z *zsh) Name() string {
2323
return "zsh"
2424
}
2525

26-
// UsesOwnFile implements Shell.
27-
func (z *zsh) UsesOwnFile() bool {
28-
return false
29-
}
30-
3126
// InstallPath implements Shell.
3227
func (z *zsh) InstallPath() (string, error) {
3328
homeDir, err := home.Dir()
@@ -39,19 +34,20 @@ func (z *zsh) InstallPath() (string, error) {
3934

4035
// WriteCompletion implements Shell.
4136
func (z *zsh) WriteCompletion(w io.Writer) error {
42-
return generateCompletion(zshCompletionTemplate)(w, z.programName)
37+
return configTemplateWriter(w, zshCompletionTemplate, z.programName)
4338
}
4439

45-
const zshCompletionTemplate = `
40+
// ProgramName implements Shell.
41+
func (z *zsh) ProgramName() string {
42+
return z.programName
43+
}
4644

47-
# === BEGIN {{.Name}} COMPLETION ===
45+
const zshCompletionTemplate = `
4846
_{{.Name}}_completions() {
4947
local -a args completions
5048
args=("${words[@]:1:$#words}")
5149
completions=($(COMPLETION_MODE=1 "{{.Name}}" "${args[@]}"))
5250
compadd -a completions
5351
}
5452
compdef _{{.Name}}_completions {{.Name}}
55-
# === END {{.Name}} COMPLETION ===
56-
5753
`

0 commit comments

Comments
 (0)