-
Notifications
You must be signed in to change notification settings - Fork 3.8k
/
demo.go
216 lines (186 loc) · 5.97 KB
/
demo.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
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
// Copyright 2018 The Cockroach Authors.
//
// Use of this software is governed by the Business Source License
// included in the file licenses/BSL.txt.
//
// As of the Change Date specified in that file, in accordance with
// the Business Source License, use of this software will be governed
// by the Apache License, Version 2.0, included in the file
// licenses/APL.txt.
package cli
import (
"context"
gosql "database/sql"
"fmt"
"net/url"
"github.com/cockroachdb/cockroach/pkg/base"
"github.com/cockroachdb/cockroach/pkg/cli/cliflags"
"github.com/cockroachdb/cockroach/pkg/security"
"github.com/cockroachdb/cockroach/pkg/server"
"github.com/cockroachdb/cockroach/pkg/sql/sqlbase"
"github.com/cockroachdb/cockroach/pkg/util"
"github.com/cockroachdb/cockroach/pkg/util/log"
"github.com/cockroachdb/cockroach/pkg/util/log/logflags"
"github.com/cockroachdb/cockroach/pkg/workload"
"github.com/cockroachdb/cockroach/pkg/workload/workloadsql"
"github.com/spf13/cobra"
"github.com/spf13/pflag"
)
var demoCmd = &cobra.Command{
Use: "demo",
Short: "open a demo sql shell",
Long: `
Start an in-memory, standalone, single-node CockroachDB instance, and open an
interactive SQL prompt to it. Various datasets are available to be preloaded as
subcommands: e.g. "cockroach demo startrek". See --help for a full list.
By default, the 'movr' dataset is pre-loaded. You can also use --empty
to avoid pre-loading a dataset.`,
Example: ` cockroach demo`,
Args: cobra.NoArgs,
RunE: MaybeDecorateGRPCError(func(cmd *cobra.Command, _ []string) error {
return runDemo(cmd, nil /* gen */)
}),
}
const defaultGeneratorName = "movr"
var defaultGenerator workload.Generator
func init() {
for _, meta := range workload.Registered() {
gen := meta.New()
if meta.Name == defaultGeneratorName {
// Save the default for use in the top-level 'demo' command
// without argument.
defaultGenerator = gen
}
var genFlags *pflag.FlagSet
if f, ok := gen.(workload.Flagser); ok {
genFlags = f.Flags().FlagSet
}
genDemoCmd := &cobra.Command{
Use: meta.Name,
Short: meta.Description,
Args: cobra.ArbitraryArgs,
RunE: MaybeDecorateGRPCError(func(cmd *cobra.Command, _ []string) error {
return runDemo(cmd, gen)
}),
}
demoCmd.AddCommand(genDemoCmd)
genDemoCmd.Flags().AddFlagSet(genFlags)
}
}
func setupTransientServers(
cmd *cobra.Command, gen workload.Generator,
) (connURL string, adminURL string, cleanup func(), err error) {
cleanup = func() {}
ctx := context.Background()
// Set up logging. For demo/transient server we use non-standard
// behavior where we avoid file creation if possible.
df := cmd.Flags().Lookup(cliflags.LogDir.Name)
sf := cmd.Flags().Lookup(logflags.LogToStderrName)
if !df.Changed && !sf.Changed {
// User did not request logging flags; shut down logging under
// errors and make logs appear on stderr.
// Otherwise, the demo command would cause a cockroach-data
// directory to appear in the current directory just for logs.
_ = df.Value.Set("")
df.Changed = true
_ = sf.Value.Set(log.Severity_ERROR.String())
sf.Changed = true
}
stopper, err := setupAndInitializeLoggingAndProfiling(ctx, cmd)
if err != nil {
return connURL, adminURL, cleanup, err
}
cleanup = func() { stopper.Stop(ctx) }
// Create the first transient server. The others will join this one.
args := base.TestServerArgs{
PartOfCluster: true,
Insecure: true,
Stopper: stopper,
}
serverFactory := server.TestServerFactory
s := serverFactory.New(args).(*server.TestServer)
if err := s.Start(args); err != nil {
return connURL, adminURL, cleanup, err
}
args.JoinAddr = s.ServingRPCAddr()
for i := 0; i < demoCtx.nodes-1; i++ {
s := serverFactory.New(args).(*server.TestServer)
if err := s.Start(args); err != nil {
return connURL, adminURL, cleanup, err
}
}
// Prepare the URL for use by the SQL shell.
options := url.Values{}
options.Add("sslmode", "disable")
options.Add("application_name", sqlbase.ReportableAppNamePrefix+"cockroach demo")
url := url.URL{
Scheme: "postgres",
User: url.User(security.RootUser),
Host: s.ServingSQLAddr(),
RawQuery: options.Encode(),
}
if demoCtx.nodes < 3 {
// Set up the default zone configuration. We are using an in-memory store
// so we really want to disable replication.
// Note: we use the raw URL without a database (Path)
// as it is not creaed yet.
if err := cliDisableReplication(ctx, url.String()); err != nil {
return ``, ``, cleanup, err
}
}
if gen != nil {
url.Path = gen.Meta().Name
}
urlStr := url.String()
// If there is a load generator, create its database and load its
// fixture.
if gen != nil {
db, err := gosql.Open("postgres", urlStr)
if err != nil {
return ``, ``, cleanup, err
}
defer db.Close()
if _, err := db.Exec(`CREATE DATABASE ` + gen.Meta().Name); err != nil {
return ``, ``, cleanup, err
}
ctx := context.TODO()
var l workloadsql.InsertsDataLoader
if _, err := workloadsql.Setup(ctx, db, gen, l); err != nil {
return ``, ``, cleanup, err
}
}
return urlStr, s.AdminURL(), cleanup, nil
}
func runDemo(cmd *cobra.Command, gen workload.Generator) error {
if gen == nil && !demoCtx.useEmptyDatabase {
// Use a default dataset unless prevented by --empty.
gen = defaultGenerator
}
connURL, adminURL, cleanup, err := setupTransientServers(cmd, gen)
defer cleanup()
if err != nil {
return checkAndMaybeShout(err)
}
checkInteractive()
if cliCtx.isInteractive {
fmt.Printf(`#
# Welcome to the CockroachDB demo database!
#
# You are connected to a temporary, in-memory CockroachDB cluster of %d node%s.
`, demoCtx.nodes, util.Pluralize(int64(demoCtx.nodes)))
if gen != nil {
fmt.Printf("# The cluster has been preloaded with the %q dataset\n# (%s).\n",
gen.Meta().Name, gen.Meta().Description)
}
fmt.Printf(`#
# Your changes will not be saved!
#
# Web UI: %s
#
`, adminURL)
}
checkTzDatabaseAvailability(context.Background())
conn := makeSQLConn(connURL)
defer conn.Close()
return runClient(cmd, conn)
}