-
Notifications
You must be signed in to change notification settings - Fork 3
/
cockroach_test_server.go
148 lines (133 loc) · 3.62 KB
/
cockroach_test_server.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
package db
import (
"bytes"
"context"
"io"
"io/ioutil"
"os"
"os/exec"
"path/filepath"
"regexp"
"pixur.org/pixur/be/status"
)
var cockroachReadyMatcher = regexp.MustCompile(`sql:[\s]*postgresql://([^@]+)@([^:]+):([\d]+)\?`)
type scanUntilCockroachReady struct {
io.WriteCloser
buf bytes.Buffer
done chan []string
}
func (scan *scanUntilCockroachReady) Write(data []byte) (int, error) {
if scan.done != nil {
scan.buf.Write(data)
match := cockroachReadyMatcher.FindStringSubmatch(scan.buf.String())
if len(match) > 0 {
select {
case scan.done <- match[1:]:
close(scan.done)
scan.done = nil
default:
panic("already done")
}
}
}
return scan.WriteCloser.Write(data)
}
func (scan *scanUntilCockroachReady) Close() error {
if scan.done != nil {
close(scan.done)
scan.done = nil
}
return scan.WriteCloser.Close()
}
type testCockroachPostgresServer struct {
cmd *exec.Cmd
testdir string
user, host, port string
}
func (s *testCockroachPostgresServer) start(ctx context.Context) (stscap status.S) {
var defers []func()
defer func() {
for i := len(defers) - 1; i >= 0; i-- {
defers[i]()
}
}()
testdir, err := ioutil.TempDir("", "postgrespixurtest")
if err != nil {
return status.Unknown(err, "can't create temp dir")
}
defers = append(defers, func() {
if err := os.RemoveAll(testdir); err != nil {
status.ReplaceOrSuppress(&stscap, status.Unknown(err, "failed to remove testdir while cleaning up"))
}
})
// Don't use context command as it would kill the cmd even after starting successfully
cmd := exec.Command(
"cockroach",
"start",
"--insecure",
"--host", "localhost",
"--port", "0",
"--store", "path="+testdir,
"--http-port", "0",
"--http-host", "0.0.0.0",
)
stderr, err := os.Create(filepath.Join(testdir, "STDERR"))
if err != nil {
return status.Unknown(err, "can't create stderr")
}
defers = append(defers, func() {
if err := stderr.Close(); err != nil {
status.ReplaceOrSuppress(&stscap, status.Unknown(err, "failed to close stderr while cleaning up"))
}
})
stdout, err := os.Create(filepath.Join(testdir, "STDOUT"))
if err != nil {
return status.Unknown(err, "can't create stdout")
}
defers = append(defers, func() {
if err := stdout.Close(); err != nil {
status.ReplaceOrSuppress(&stscap, status.Unknown(err, "failed to close stdout while cleaning up"))
}
})
doneparts := make(chan []string, 1)
scan := &scanUntilCockroachReady{
WriteCloser: stdout,
done: doneparts,
}
cmd.Stderr = stderr
cmd.Stdout = scan
if err := cmd.Start(); err != nil {
return status.Unknown(err, "can't start")
}
defers = append(defers, func() {
if err := cmd.Process.Kill(); err != nil {
status.ReplaceOrSuppress(&stscap, status.Unknown(err, "failed to kill process while cleaning up"))
}
})
select {
case parts := <-doneparts:
s.user, s.host, s.port = parts[0], parts[1], parts[2]
case <-ctx.Done():
return status.From(ctx.Err())
}
s.cmd = cmd
s.testdir = testdir
defers = nil
return nil
}
func (s *testCockroachPostgresServer) stop() status.S {
var sts status.S
if err := s.cmd.Process.Kill(); err != nil {
status.ReplaceOrSuppress(&sts, status.Unknown(err, "failed to kill process"))
}
if err := s.cmd.Stderr.(io.Closer).Close(); err != nil {
status.ReplaceOrSuppress(&sts, status.Unknown(err, "failed to close stderr"))
}
if err := s.cmd.Stdout.(io.Closer).Close(); err != nil {
status.ReplaceOrSuppress(&sts, status.Unknown(err, "failed to close stdout"))
}
if err := os.RemoveAll(s.testdir); err != nil {
status.ReplaceOrSuppress(&sts, status.Unknown(err, "failed to remove testdir"))
}
return sts
}