demo: Get a temporary license upon demo startup

There is a wider discussion/design going on to serve licenses in
a more general way in the future, so this commit is not aiming
for a future-proof design, but instead an MVP to allow users to
demo enterprise features within `cockroach demo`.

We are not concerned about offline usage of enterprise features
as users can obtain a license and enable features manually using SET.

Fixes #40222.

Release note (cli change): cockroach demo attempts to contact a license
server to obtain a temporary license. cockroach demo now enables
telemetry for the demo cluster. This feature can be opted out of by
setting the `COCKROACH_SKIP_ENABLING_DIAGNOSTIC_REPORTING` environment variable
(https://www.cockroachlabs.com/docs/stable/diagnostics-reporting.html).

pkg/ccl/cliccl/demo.go

@@ -0,0 +1,79 @@
+// Copyright 2019 The Cockroach Authors.
+//
+// Licensed as a CockroachDB Enterprise file under the Cockroach Community
+// License (the "License"); you may not use this file except in compliance with
+// the License. You may obtain a copy of the License at
+//
+//     https://github.com/cockroachdb/cockroach/blob/master/licenses/CCL.txt
+
+package cliccl
+
+import (
+	gosql "database/sql"
+	"fmt"
+	"io/ioutil"
+	"net/http"
+	"time"
+
+	"github.com/cockroachdb/cockroach/pkg/build"
+	"github.com/cockroachdb/cockroach/pkg/cli"
+	"github.com/cockroachdb/cockroach/pkg/util/log"
+	"github.com/cockroachdb/cockroach/pkg/util/uuid"
+	"github.com/cockroachdb/errors"
+)
+
+// TODO (rohany): change this once another endpoint is setup for getting licenses.
+// This URL grants a license that is valid for 1 hour.
+const licenseURL = "https://register.cockroachdb.com/api/prodtest"
+
+func getLicense(clusterID uuid.UUID) (string, error) {
+	client := &http.Client{
+		Timeout: 5 * time.Second,
+	}
+	req, err := http.NewRequest("GET", licenseURL, nil)
+	if err != nil {
+		return "", err
+	}
+	// Send some extra information to the endpoint.
+	q := req.URL.Query()
+	q.Add("type", "demo")
+	q.Add("version", build.VersionPrefix())
+	q.Add("clusterid", clusterID.String())
+	req.URL.RawQuery = q.Encode()
+
+	resp, err := client.Do(req)
+	if err != nil {
+		return "", err
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		return "", errors.New("unable to connect to licensing endpoint")
+	}
+	bodyBytes, err := ioutil.ReadAll(resp.Body)
+	if err != nil {
+		return "", err
+	}
+	return string(bodyBytes), nil
+}
+
+func getAndApplyLicense(db *gosql.DB, clusterID uuid.UUID, org string) (bool, error) {
+	license, err := getLicense(clusterID)
+	if err != nil {
+		fmt.Fprintf(log.OrigStderr, "error when contacting licensing server: %+v\n", err)
+		return false, nil
+	}
+	if _, err := db.Exec(`SET CLUSTER SETTING cluster.organization = $1`, org); err != nil {
+		return false, err
+	}
+	if _, err := db.Exec(`SET CLUSTER SETTING enterprise.license = $1`, license); err != nil {
+		return false, err
+	}
+	return true, nil
+}
+
+func init() {
+	// Set the GetAndApplyLicense function within cockroach demo.
+	// This separation is done to avoid using enterprise features in an OSS/BSL build.
+	cli.GetAndApplyLicense = getAndApplyLicense
+}

pkg/cli/cli.go

@@ -63,8 +63,13 @@ func Main() {
 
 	defer log.RecoverAndReportPanic(context.Background(), &serverCfg.Settings.SV)
 
+	err := Run(os.Args[1:])
+	exitWithError(cmdName, err)
+}
+
+func exitWithError(cmdName string, err error) {
 	errCode := 0
-	if err := Run(os.Args[1:]); err != nil {
+	if err != nil {
 		// Display the error and its details/hints.
 		fmt.Fprintln(stderr, "Error:", err.Error())
 		maybeShowErrorDetails(stderr, err, false /* printNewline */)

pkg/cli/demo.go

@@ -22,11 +22,13 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/security"
 	"github.com/cockroachdb/cockroach/pkg/server"
+	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
 	"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/util/stop"
+	"github.com/cockroachdb/cockroach/pkg/util/uuid"
 	"github.com/cockroachdb/cockroach/pkg/workload"
 	"github.com/cockroachdb/cockroach/pkg/workload/histogram"
 	"github.com/cockroachdb/cockroach/pkg/workload/workloadsql"
@@ -45,14 +47,22 @@ 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.`,
+to avoid pre-loading a dataset.
+
+cockroach demo attempts to connect to a Cockroach Labs server to obtain a
+temporary enterprise license for demoing enterprise features and enable
+telemetry back to Cockroach Labs. In order to disable this behavior, set the
+environment variable "COCKROACH_SKIP_ENABLING_DIAGNOSTIC_REPORTING".
+`,
 	Example: `  cockroach demo`,
 	Args:    cobra.NoArgs,
 	RunE: MaybeDecorateGRPCError(func(cmd *cobra.Command, _ []string) error {
 		return runDemo(cmd, nil /* gen */)
 	}),
 }
 
+const demoOrg = "Cockroach Labs - Production Testing"
+
 const defaultGeneratorName = "movr"
 
 var defaultGenerator workload.Generator
@@ -100,6 +110,10 @@ func init() {
 	}
 }
 
+// GetAndApplyLicense is not implemented in order to keep OSS/BSL builds successful.
+// The cliccl package sets this function if enterprise features are available to demo.
+var GetAndApplyLicense func(dbConn *gosql.DB, clusterID uuid.UUID, org string) (bool, error)
+
 func setupTransientServers(
 	cmd *cobra.Command, gen workload.Generator,
 ) (connURL string, adminURL string, cleanup func(), err error) {
@@ -193,6 +207,33 @@ func setupTransientServers(
 	}
 	urlStr := url.String()
 
+	// Start up the update check loop.
+	// We don't do this in (*server.Server).Start() because we don't want it
+	// in tests.
+	if !cluster.TelemetryOptOut() {
+		s.PeriodicallyCheckForUpdates(ctx)
+		// If we allow telemetry, then also try and get an enterprise license for the demo.
+		// GetAndApplyLicense will be nil in the pure OSS/BSL build of cockroach.
+		if GetAndApplyLicense != nil {
+			db, err := gosql.Open("postgres", urlStr)
+			if err != nil {
+				return ``, ``, cleanup, err
+			}
+			// Perform license acquisition asynchronously to avoid delay in cli startup.
+			go func() {
+				defer db.Close()
+				success, err := GetAndApplyLicense(db, s.ClusterID(), demoOrg)
+				if err != nil {
+					exitWithError("demo", err)
+				}
+				if !success {
+					const msg = "Unable to acquire demo license. Enterprise features are not enabled in this session.\n"
+					fmt.Fprint(stderr, msg)
+				}
+			}()
+		}
+	}
+
 	// If there is a load generator, create its database and load its
 	// fixture.
 	if gen != nil {

pkg/cli/interactive_tests/test_demo_telemetry.tcl

@@ -0,0 +1,19 @@
+#! /usr/bin/env expect -f
+
+source [file join [file dirname $argv0] common.tcl]
+
+start_test "Check cockroach demo telemetry and license check can be disabled"
+
+# set the proper environment variable
+set env(COCKROACH_SKIP_ENABLING_DIAGNOSTIC_REPORTING) "true"
+spawn $argv demo
+# wait for the CLI to start up
+eexpect "movr>"
+# send a request for an enterprise feature
+send "alter table vehicles partition by list (city) (partition p1 values in ('nyc'));\n"
+# expect that it failed, as no license was requested.
+eexpect "use of partitions requires an enterprise license"
+# clean up after the test
+interrupt
+eexpect eof
+end_test

pkg/cli/start.go

@@ -715,7 +715,7 @@ If problems persist, please see ` + base.DocsURL("cluster-setup-troubleshooting.
 			// Start up the update check loop.
 			// We don't do this in (*server.Server).Start() because we don't want it
 			// in tests.
-			if !envutil.EnvOrDefaultBool("COCKROACH_SKIP_UPDATE_CHECK", false) {
+			if !cluster.TelemetryOptOut() {
 				s.PeriodicallyCheckForUpdates(ctx)
 			}
 

pkg/settings/cluster/settings.go

@@ -18,6 +18,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/settings"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgcode"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror"
+	"github.com/cockroachdb/cockroach/pkg/util/envutil"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/cockroach/pkg/util/protoutil"
 	"github.com/cockroachdb/cockroach/pkg/util/tracing"
@@ -74,6 +75,11 @@ type Settings struct {
 	cpuProfiling int32 // atomic
 }
 
+// TelemetryOptOut is a place for controlling whether to opt out of telemetry or not.
+func TelemetryOptOut() bool {
+	return envutil.EnvOrDefaultBool("COCKROACH_SKIP_ENABLING_DIAGNOSTIC_REPORTING", false)
+}
+
 // IsCPUProfiling returns true if a pprofui CPU profile is being recorded. This can
 // be used by moving parts across the system to add profiler labels which are
 // too expensive to be enabled at all times.

pkg/sqlmigrations/migrations.go

@@ -28,7 +28,6 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/sessiondata"
 	"github.com/cockroachdb/cockroach/pkg/sql/sqlbase"
 	"github.com/cockroachdb/cockroach/pkg/sql/types"
-	"github.com/cockroachdb/cockroach/pkg/util/envutil"
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/cockroach/pkg/util/protoutil"
@@ -555,8 +554,6 @@ func createCommentTable(ctx context.Context, r runner) error {
 	return createSystemTable(ctx, r, sqlbase.CommentsTable)
 }
 
-var reportingOptOut = envutil.EnvOrDefaultBool("COCKROACH_SKIP_ENABLING_DIAGNOSTIC_REPORTING", false)
-
 func runStmtAsRootWithRetry(
 	ctx context.Context, r runner, opName string, stmt string, qargs ...interface{},
 ) error {
@@ -585,7 +582,7 @@ var SettingsDefaultOverrides = map[string]string{
 
 func optInToDiagnosticsStatReporting(ctx context.Context, r runner) error {
 	// We're opting-out of the automatic opt-in. See discussion in updates.go.
-	if reportingOptOut {
+	if cluster.TelemetryOptOut() {
 		return nil
 	}
 	return runStmtAsRootWithRetry(