cli/demo: fix enterprise features for multitenancy

pkg/ccl/cliccl/BUILD.bazel

@@ -14,12 +14,12 @@ go_library(
     deps = [
         "//pkg/base",
         "//pkg/blobs",
-        "//pkg/build",
         "//pkg/ccl/backupccl",
         "//pkg/ccl/baseccl",
         "//pkg/ccl/cliccl/cliflagsccl",
         "//pkg/ccl/storageccl",
         "//pkg/ccl/storageccl/engineccl/enginepbccl",
+        "//pkg/ccl/utilccl",
         "//pkg/ccl/workloadccl/cliccl",
         "//pkg/cli",
         "//pkg/cli/clierrorplus",
@@ -43,10 +43,8 @@ go_library(
         "//pkg/sql/sem/tree",
         "//pkg/storage",
         "//pkg/storage/enginepb",
-        "//pkg/util/envutil",
         "//pkg/util/hlc",
         "//pkg/util/humanizeutil",
-        "//pkg/util/log",
         "//pkg/util/protoutil",
         "//pkg/util/stop",
         "//pkg/util/timeutil",

pkg/ccl/cliccl/demo.go

@@ -10,74 +10,23 @@ package cliccl
 
 import (
 	gosql "database/sql"
-	"fmt"
-	"io/ioutil"
-	"net/http"
-	"time"
 
-	"github.com/cockroachdb/cockroach/pkg/build"
+	"github.com/cockroachdb/cockroach/pkg/ccl/utilccl"
 	"github.com/cockroachdb/cockroach/pkg/cli/democluster"
-	"github.com/cockroachdb/cockroach/pkg/util/envutil"
-	"github.com/cockroachdb/cockroach/pkg/util/log"
-	"github.com/cockroachdb/cockroach/pkg/util/uuid"
-	"github.com/cockroachdb/errors"
 )
 
-// This URL grants a license that is valid for 24 hours.
-const licenseDefaultURL = "https://register.cockroachdb.com/api/license"
-
-// We make licenseURL configurable for use in tests.
-var licenseURL = envutil.EnvOrDefaultString("COCKROACH_DEMO_LICENSE_URL", licenseDefaultURL)
-
-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()
-	// Let the endpoint know we are requesting a demo license.
-	q.Add("kind", "demo")
-	q.Add("version", build.BinaryVersionPrefix())
-	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)
+// enableEnterpriseForDemo enables enterprise features for 'cockroach demo'.
+// It is not intended for use for non-demo servers.
+func enableEnterpriseForDemo(db *gosql.DB, org string) (func(), error) {
+	_, err := db.Exec(`SET CLUSTER SETTING cluster.organization = $1`, org)
 	if err != nil {
-		fmt.Fprintf(log.OrigStderr, "\nerror while contacting licensing server:\n%+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 nil, err
 	}
-	return true, nil
+	return utilccl.TestingEnableEnterprise(), nil
 }
 
 func init() {
-	// Set the GetAndApplyLicense function within cockroach demo.
+	// Set the EnableEnterprise function within cockroach demo.
 	// This separation is done to avoid using enterprise features in an OSS/BSL build.
-	democluster.GetAndApplyLicense = getAndApplyLicense
+	democluster.EnableEnterprise = enableEnterpriseForDemo
 }

pkg/cli/cliflags/flags.go

@@ -1284,7 +1284,7 @@ and the system tenant using the \connect command.`,
 	DemoNoLicense = FlagInfo{
 		Name: "disable-demo-license",
 		Description: `
-If set, disable cockroach demo from attempting to obtain a temporary license.`,
+If set, disable enterprise features.`,
 	}
 
 	UseEmptyDatabase = FlagInfo{

pkg/cli/context.go

@@ -578,8 +578,13 @@ func setConvContextDefaults() {
 
 // demoCtx captures the command-line parameters of the `demo` command.
 // See below for defaults.
-var demoCtx = democluster.Context{
-	CliCtx: &cliCtx.Context,
+var demoCtx = struct {
+	democluster.Context
+	disableEnterpriseFeatures bool
+}{
+	Context: democluster.Context{
+		CliCtx: &cliCtx.Context,
+	},
 }
 
 // setDemoContextDefaults set the default values in demoCtx.  This
@@ -595,7 +600,6 @@ func setDemoContextDefaults() {
 	demoCtx.Localities = nil
 	demoCtx.GeoPartitionedReplicas = false
 	demoCtx.DisableTelemetry = false
-	demoCtx.DisableLicenseAcquisition = false
 	demoCtx.DefaultKeySize = defaultKeySize
 	demoCtx.DefaultCALifetime = defaultCALifetime
 	demoCtx.DefaultCertLifetime = defaultCertLifetime
@@ -604,6 +608,8 @@ func setDemoContextDefaults() {
 	demoCtx.HTTPPort, _ = strconv.Atoi(base.DefaultHTTPPort)
 	demoCtx.WorkloadMaxQPS = 25
 	demoCtx.Multitenant = true
+
+	demoCtx.disableEnterpriseFeatures = false
 }
 
 // stmtDiagCtx captures the command-line parameters of the 'statement-diag'

pkg/cli/demo.go

@@ -41,8 +41,7 @@ subcommands: e.g. "cockroach demo startrek". See --help for a full list.
 By default, the 'movr' dataset is pre-loaded. You can also use --no-example-database
 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
+cockroach demo attempts to connect to a Cockroach Labs server to send
 telemetry back to Cockroach Labs. In order to disable this behavior, set the
 environment variable "COCKROACH_SKIP_ENABLING_DIAGNOSTIC_REPORTING" to true.
 `,
@@ -143,15 +142,19 @@ func checkDemoConfiguration(
 	}
 
 	demoCtx.DisableTelemetry = cluster.TelemetryOptOut()
-	// disableLicenseAcquisition can also be set by the user as an
-	// input flag, so make sure it include it when considering the final
-	// value of disableLicenseAcquisition.
-	demoCtx.DisableLicenseAcquisition =
-		demoCtx.DisableTelemetry || (democluster.GetAndApplyLicense == nil) || demoCtx.DisableLicenseAcquisition
+
+	// Whether or not we enable enterprise feature is a combination of:
+	//
+	// - whether the user wants them (they can disable enterprise
+	//   features explicitly with --no-license, e.g. for testing what
+	//   errors come up if no license is available)
+	// - whether enterprise features are enabled in this build, depending
+	//   on whether CCL code is included (and sets democluster.EnableEnterprise).
+	demoCtx.disableEnterpriseFeatures = democluster.EnableEnterprise == nil || demoCtx.disableEnterpriseFeatures
 
 	if demoCtx.GeoPartitionedReplicas {
 		geoFlag := "--" + cliflags.DemoGeoPartitionedReplicas.Name
-		if demoCtx.DisableLicenseAcquisition {
+		if demoCtx.disableEnterpriseFeatures {
 			return nil, errors.Newf("enterprise features are needed for this demo (%s)", geoFlag)
 		}
 
@@ -217,7 +220,7 @@ func runDemo(cmd *cobra.Command, gen workload.Generator) (resErr error) {
 
 	demoCtx.WorkloadGenerator = gen
 
-	c, err := democluster.NewDemoCluster(ctx, &demoCtx,
+	c, err := democluster.NewDemoCluster(ctx, &demoCtx.Context,
 		log.Infof,
 		log.Warningf,
 		log.Ops.Shoutf,
@@ -274,25 +277,29 @@ func runDemo(cmd *cobra.Command, gen workload.Generator) (resErr error) {
 		// Only print details about the telemetry configuration if the
 		// user has control over it.
 		if demoCtx.DisableTelemetry {
-			cliCtx.PrintlnUnlessEmbedded("#\n# Telemetry and automatic license acquisition disabled by configuration.")
-		} else if demoCtx.DisableLicenseAcquisition {
-			cliCtx.PrintlnUnlessEmbedded("#\n# Enterprise features disabled by OSS-only build.")
+			cliCtx.PrintlnUnlessEmbedded("#\n# Telemetry disabled by configuration.")
 		} else {
-			cliCtx.PrintlnUnlessEmbedded("#\n# This demo session will attempt to enable enterprise features\n" +
-				"# by acquiring a temporary license from Cockroach Labs in the background.\n" +
+			cliCtx.PrintlnUnlessEmbedded("#\n" +
+				"# This demo session will send telemetry to Cockroach Labs in the background.\n" +
 				"# To disable this behavior, set the environment variable\n" +
 				"# COCKROACH_SKIP_ENABLING_DIAGNOSTIC_REPORTING=true.")
 		}
+
+		if demoCtx.disableEnterpriseFeatures {
+			cliCtx.PrintlnUnlessEmbedded("#\n# Enterprise features are disabled by configuration.\n")
+		}
 	}
 
-	// Start license acquisition in the background.
-	licenseDone, err := c.AcquireDemoLicense(ctx)
-	if err != nil {
-		return clierrorplus.CheckAndMaybeShout(err)
+	if !demoCtx.disableEnterpriseFeatures {
+		fn, err := c.EnableEnterprise(ctx)
+		if err != nil {
+			return clierrorplus.CheckAndMaybeShout(err)
+		}
+		defer fn()
 	}
 
 	// Initialize the workload, if requested.
-	if err := c.SetupWorkload(ctx, licenseDone); err != nil {
+	if err := c.SetupWorkload(ctx); err != nil {
 		return clierrorplus.CheckAndMaybeShout(err)
 	}
 
@@ -331,21 +338,6 @@ func runDemo(cmd *cobra.Command, gen workload.Generator) (resErr error) {
 				certsDir,
 			)
 		}
-
-		// It's ok to do this twice (if workload setup already waited) because
-		// then the error return is guaranteed to be nil.
-		go func() {
-			if err := waitForLicense(licenseDone); err != nil {
-				_ = clierrorplus.CheckAndMaybeShout(err)
-			}
-		}()
-	} else {
-		// If we are not running an interactive shell, we need to wait to ensure
-		// that license acquisition is successful. If license acquisition is
-		// disabled, then a read on this channel will return immediately.
-		if err := waitForLicense(licenseDone); err != nil {
-			return clierrorplus.CheckAndMaybeShout(err)
-		}
 	}
 
 	conn, err := sqlCtx.MakeConn(c.GetConnURL())
@@ -357,8 +349,3 @@ func runDemo(cmd *cobra.Command, gen workload.Generator) (resErr error) {
 	sqlCtx.ShellCtx.ParseURL = makeURLParser(cmd)
 	return sqlCtx.Run(conn)
 }
-
-func waitForLicense(licenseDone <-chan error) error {
-	err := <-licenseDone
-	return err
-}

pkg/cli/democluster/BUILD.bazel

@@ -30,6 +30,7 @@ go_library(
         "//pkg/server/pgurl",
         "//pkg/server/serverpb",
         "//pkg/server/status",
+        "//pkg/sql",
         "//pkg/sql/catalog/catalogkeys",
         "//pkg/sql/distsql",
         "//pkg/sql/sem/catconstants",
@@ -41,7 +42,6 @@ go_library(
         "//pkg/util/netutil/addr",
         "//pkg/util/retry",
         "//pkg/util/stop",
-        "//pkg/util/uuid",
         "//pkg/workload",
         "//pkg/workload/histogram",
         "//pkg/workload/workloadsql",

pkg/cli/democluster/api.go

@@ -17,7 +17,6 @@ import (
 	democlusterapi "github.com/cockroachdb/cockroach/pkg/cli/democluster/api"
 	"github.com/cockroachdb/cockroach/pkg/security/username"
 	"github.com/cockroachdb/cockroach/pkg/server"
-	"github.com/cockroachdb/cockroach/pkg/util/uuid"
 )
 
 // DemoCluster represents a demo cluster.
@@ -45,16 +44,15 @@ type DemoCluster interface {
 	// Close shuts down the demo cluster.
 	Close(ctx context.Context)
 
-	// AcquireDemoLicense acquires the demo license if configured,
-	// otherwise does nothing. In any case, if there is no error, it
-	// returns a channel that will either produce an error or a nil
-	// value.
-	AcquireDemoLicense(ctx context.Context) (chan error, error)
+	// EnableEnterprise enables enterprise features for this demo,
+	// if available in this build. The returned callback should be called
+	// before terminating the demo.
+	EnableEnterprise(ctx context.Context) (func(), error)
 
 	// SetupWorkload initializes the workload generator if defined.
-	SetupWorkload(ctx context.Context, licenseDone <-chan error) error
+	SetupWorkload(ctx context.Context) error
 }
 
-// GetAndApplyLicense is not implemented in order to keep OSS/BSL builds successful.
+// EnableEnterprise is not implemented here 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)
+var EnableEnterprise func(db *gosql.DB, org string) (func(), error)

pkg/cli/democluster/context.go

@@ -38,10 +38,6 @@ type Context struct {
 	// DisableTelemetry requests that telemetry be disabled.
 	DisableTelemetry bool
 
-	// DisableLicenseAcquisition requests that no evaluation license be
-	// automatically acquired to enable enterprise features.
-	DisableLicenseAcquisition bool
-
 	// NoExampleDatabase prevents the auto-creation of a demo database
 	// from a workload.
 	NoExampleDatabase bool