Merge #81737 #81797 #81800

81737: cli/demo: fix enterprise features for multitenancy r=rafiss a=knz

First commit from #81762.

Fixes #80270.
Fixes #40434.

This patch achieves two things:

- it ensures that enterprise features work properly in multi-tenant `cockroach
  demo` when `--disable-demo-license` is not passed as argument.
- it ensures `cockroach demo` does not require a working licensing
  endpoint to run (by not retrieving a license over the network on startup).

The net result will be less spurious failures in CI, as well as
avoiding startup delays/errors during interactive uses, which provide
poor UX.

Release note: None

81797: sql: assert that table writers use root txns r=yuzefovich a=yuzefovich

We assume that all mutation statements use the Root txn, but this was
never enforced in practice. This commit adds an explicit check when
initializing the table writer that the root txn is provided.

Release note: None

81800: stats: add logging to TestDefaultColumns to help debug flake r=rytaft a=rytaft

This commit adds some logging to the TestDefaultColumns test in order
to help debug a test failure.

Informs #81513

Release note: None

Co-authored-by: Raphael 'kena' Poss <knz@thaumogen.net>
Co-authored-by: Yahor Yuzefovich <yahor@cockroachlabs.com>
Co-authored-by: Rebecca Taft <becca@cockroachlabs.com>

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

@@ -42,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