mt_start_sql: enable enterprise features for multitenant sql servers

Enterprise features are controlled by the enterprise.license setting.
Currently this setting applies only to the host tenant cluster. This
change allows tenants to read the license from an environment variable.
This should be replaced once it is possible to read the license directly
from the host cluster.

Release note: None

pkg/ccl/serverccl/BUILD.bazel

@@ -23,6 +23,7 @@ go_test(
         "//pkg/ccl",
         "//pkg/ccl/kvccl",
         "//pkg/ccl/utilccl",
+        "//pkg/ccl/utilccl/licenseccl",
         "//pkg/roachpb:with-mocks",
         "//pkg/security",
         "//pkg/security/securitytest",
@@ -35,6 +36,7 @@ go_test(
         "//pkg/testutils/sqlutils",
         "//pkg/testutils/testcluster",
         "//pkg/util",
+        "//pkg/util/envutil",
         "//pkg/util/httputil",
         "//pkg/util/leaktest",
         "//pkg/util/log",

pkg/ccl/serverccl/server_sql_test.go

@@ -16,12 +16,15 @@ import (
 	"testing"
 
 	"github.com/cockroachdb/cockroach/pkg/base"
+	"github.com/cockroachdb/cockroach/pkg/ccl/utilccl"
+	"github.com/cockroachdb/cockroach/pkg/ccl/utilccl/licenseccl"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/security"
 	"github.com/cockroachdb/cockroach/pkg/sql"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgcode"
 	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
 	"github.com/cockroachdb/cockroach/pkg/testutils/sqlutils"
+	"github.com/cockroachdb/cockroach/pkg/util/envutil"
 	"github.com/cockroachdb/cockroach/pkg/util/httputil"
 	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
@@ -84,6 +87,29 @@ func TestTenantCannotSetClusterSetting(t *testing.T) {
 	require.Equal(t, pq.ErrorCode(pgcode.InsufficientPrivilege.String()), pqErr.Code, "err %v has unexpected code", err)
 }
 
+func TestTenantCanUseEnterpriseFeatures(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	license, _ := (&licenseccl.License{
+		Type: licenseccl.License_Enterprise,
+	}).Encode()
+
+	defer utilccl.TestingDisableEnterprise()()
+	defer envutil.TestSetEnv(t, "COCKROACH_TENANT_LICENSE", license)()
+
+	tc := serverutils.StartNewTestCluster(t, 1, base.TestClusterArgs{})
+	defer tc.Stopper().Stop(context.Background())
+
+	_, db := serverutils.StartTenant(t, tc.Server(0), base.TestTenantArgs{TenantID: serverutils.TestTenantID(), AllowSettingClusterSettings: false})
+	defer db.Close()
+
+	_, err := db.Exec(`BACKUP INTO 'userfile:///backup'`)
+	require.NoError(t, err)
+	_, err = db.Exec(`BACKUP INTO LATEST IN 'userfile:///backup'`)
+	require.NoError(t, err)
+}
+
 func TestTenantUnauthenticatedAccess(t *testing.T) {
 	defer leaktest.AfterTest(t)()
 	defer log.Scope(t).Close(t)

pkg/ccl/utilccl/BUILD.bazel

@@ -13,13 +13,15 @@ go_library(
         "//pkg/base",
         "//pkg/ccl/utilccl/licenseccl",
         "//pkg/kv/kvclient/kvcoord:with-mocks",
+        "//pkg/server",
         "//pkg/settings",
         "//pkg/settings/cluster",
         "//pkg/sql/catalog/colinfo",
         "//pkg/sql/flowinfra",
         "//pkg/sql/pgwire/pgcode",
         "//pkg/sql/pgwire/pgerror",
         "//pkg/sql/types",
+        "//pkg/util/envutil",
         "//pkg/util/grpcutil",
         "//pkg/util/timeutil",
         "//pkg/util/uuid",
@@ -40,6 +42,7 @@ go_test(
         "//pkg/ccl/utilccl/licenseccl",
         "//pkg/settings/cluster",
         "//pkg/testutils",
+        "//pkg/util/envutil",
         "//pkg/util/timeutil",
         "//pkg/util/uuid",
         "@com_github_stretchr_testify//require",

pkg/ccl/utilccl/license_check.go

@@ -17,10 +17,12 @@ import (
 
 	"github.com/cockroachdb/cockroach/pkg/base"
 	"github.com/cockroachdb/cockroach/pkg/ccl/utilccl/licenseccl"
+	"github.com/cockroachdb/cockroach/pkg/server"
 	"github.com/cockroachdb/cockroach/pkg/settings"
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
 	"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/timeutil"
 	"github.com/cockroachdb/cockroach/pkg/util/uuid"
 	"github.com/cockroachdb/errors"
@@ -45,16 +47,13 @@ var enterpriseLicense = func() *settings.StringSetting {
 	return s
 }()
 
-// testingEnterprise determines whether the cluster is enabled
-// or disabled for the purposes of testing.
-// It should be loaded and stored using atomic as it can race with an
-// in progress kv reader during TestingDisableEnterprise /
-// TestingEnableEnterprise.
-var testingEnterprise int32
+// enterpriseStatus determines whether the cluster is enabled
+// for enterprise features or if enterprise status depends on the license.
+var enterpriseStatus int32 = deferToLicense
 
 const (
-	testingEnterpriseDisabled = 0
-	testingEnterpriseEnabled  = 1
+	deferToLicense    = 0
+	enterpriseEnabled = 1
 )
 
 // errEnterpriseRequired is returned by check() when the caller does
@@ -68,22 +67,38 @@ type licenseCacheKey string
 
 // TestingEnableEnterprise allows overriding the license check in tests.
 func TestingEnableEnterprise() func() {
-	before := atomic.LoadInt32(&testingEnterprise)
-	atomic.StoreInt32(&testingEnterprise, testingEnterpriseEnabled)
+	before := atomic.LoadInt32(&enterpriseStatus)
+	atomic.StoreInt32(&enterpriseStatus, enterpriseEnabled)
 	return func() {
-		atomic.StoreInt32(&testingEnterprise, before)
+		atomic.StoreInt32(&enterpriseStatus, before)
 	}
 }
 
 // TestingDisableEnterprise allows re-enabling the license check in tests.
 func TestingDisableEnterprise() func() {
-	before := atomic.LoadInt32(&testingEnterprise)
-	atomic.StoreInt32(&testingEnterprise, testingEnterpriseDisabled)
+	before := atomic.LoadInt32(&enterpriseStatus)
+	atomic.StoreInt32(&enterpriseStatus, deferToLicense)
 	return func() {
-		atomic.StoreInt32(&testingEnterprise, before)
+		atomic.StoreInt32(&enterpriseStatus, before)
 	}
 }
 
+// ApplyTenantLicense verifies the COCKROACH_TENANT_LICENSE environment variable
+// and enables enterprise features for the process. This is a bit of a hack and
+// should be replaced once it is possible to read the host cluster's
+// enterprise.license setting.
+func ApplyTenantLicense() error {
+	license, ok := envutil.EnvString("COCKROACH_TENANT_LICENSE", 0)
+	if !ok {
+		return nil
+	}
+	if _, err := decode(license); err != nil {
+		return errors.Wrap(err, "COCKROACH_TENANT_LICENSE encoding is invalid")
+	}
+	atomic.StoreInt32(&enterpriseStatus, enterpriseEnabled)
+	return nil
+}
+
 // CheckEnterpriseEnabled returns a non-nil error if the requested enterprise
 // feature is not enabled, including information or a link explaining how to
 // enable it.
@@ -108,6 +123,7 @@ func init() {
 	base.CheckEnterpriseEnabled = CheckEnterpriseEnabled
 	base.LicenseType = getLicenseType
 	base.TimeToEnterpriseLicenseExpiry = TimeToEnterpriseLicenseExpiry
+	server.ApplyTenantLicense = ApplyTenantLicense
 }
 
 // TimeToEnterpriseLicenseExpiry returns a Duration from `asOf` until the current
@@ -128,7 +144,7 @@ func TimeToEnterpriseLicenseExpiry(
 func checkEnterpriseEnabledAt(
 	st *cluster.Settings, at time.Time, cluster uuid.UUID, org, feature string, withDetails bool,
 ) error {
-	if atomic.LoadInt32(&testingEnterprise) == testingEnterpriseEnabled {
+	if atomic.LoadInt32(&enterpriseStatus) == enterpriseEnabled {
 		return nil
 	}
 	license, err := getLicense(st)

pkg/ccl/utilccl/license_check_test.go

@@ -16,6 +16,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/ccl/utilccl/licenseccl"
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
 	"github.com/cockroachdb/cockroach/pkg/testutils"
+	"github.com/cockroachdb/cockroach/pkg/util/envutil"
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
 	"github.com/cockroachdb/cockroach/pkg/util/uuid"
 	"github.com/stretchr/testify/require"
@@ -188,3 +189,40 @@ func TestTimeToEnterpriseLicenseExpiry(t *testing.T) {
 		})
 	}
 }
+
+func TestApplyTenantLicenseWithLicense(t *testing.T) {
+	license, _ := (&licenseccl.License{
+		Type: licenseccl.License_Enterprise,
+	}).Encode()
+
+	defer TestingDisableEnterprise()()
+	defer envutil.TestSetEnv(t, "COCKROACH_TENANT_LICENSE", license)()
+
+	settings := cluster.MakeClusterSettings()
+
+	require.Error(t, CheckEnterpriseEnabled(settings, uuid.MakeV4(), "", ""))
+	require.False(t, IsEnterpriseEnabled(settings, uuid.MakeV4(), "", ""))
+	require.NoError(t, ApplyTenantLicense())
+	require.NoError(t, CheckEnterpriseEnabled(settings, uuid.MakeV4(), "", ""))
+	require.True(t, IsEnterpriseEnabled(settings, uuid.MakeV4(), "", ""))
+}
+
+func TestApplyTenantLicenseWithoutLicense(t *testing.T) {
+	defer TestingDisableEnterprise()()
+
+	settings := cluster.MakeClusterSettings()
+	_, ok := envutil.EnvString("COCKROACH_TENANT_LICENSE", 0)
+	envutil.ClearEnvCache()
+	require.False(t, ok)
+
+	require.Error(t, CheckEnterpriseEnabled(settings, uuid.MakeV4(), "", ""))
+	require.False(t, IsEnterpriseEnabled(settings, uuid.MakeV4(), "", ""))
+	require.NoError(t, ApplyTenantLicense())
+	require.Error(t, CheckEnterpriseEnabled(settings, uuid.MakeV4(), "", ""))
+	require.False(t, IsEnterpriseEnabled(settings, uuid.MakeV4(), "", ""))
+}
+
+func TestApplyTenantLicenseWithInvalidLicense(t *testing.T) {
+	defer envutil.TestSetEnv(t, "COCKROACH_TENANT_LICENSE", "THIS IS NOT A VALID LICENSE")()
+	require.Error(t, ApplyTenantLicense())
+}

pkg/server/tenant.go

@@ -57,6 +57,11 @@ func StartTenant(
 	baseCfg BaseConfig,
 	sqlCfg SQLConfig,
 ) (sqlServer *SQLServer, pgAddr string, httpAddr string, _ error) {
+	err := ApplyTenantLicense()
+	if err != nil {
+		return nil, "", "", err
+	}
+
 	args, err := makeTenantSQLServerArgs(stopper, kvClusterName, baseCfg, sqlCfg)
 	if err != nil {
 		return nil, "", "", err
@@ -405,6 +410,11 @@ var NewTenantSideCostController = func(
 	return noopTenantSideCostController{}, nil
 }
 
+// ApplyTenantLicense is a hook for CCL code which enables enterprise features
+// for the tenant process if the COCKROACH_TENANT_LICENSE environment variable
+// is set.
+var ApplyTenantLicense = func() error { return nil /* no-op */ }
+
 // noopTenantSideCostController is a no-op implementation of
 // TenantSideCostController.
 type noopTenantSideCostController struct{}

pkg/util/envutil/BUILD.bazel

@@ -17,4 +17,5 @@ go_test(
     size = "small",
     srcs = ["env_test.go"],
     embed = [":envutil"],
+    deps = ["@com_github_stretchr_testify//require"],
 )

pkg/util/envutil/env.go

@@ -19,6 +19,7 @@ import (
 	"sort"
 	"strconv"
 	"strings"
+	"testing"
 	"time"
 
 	"github.com/cockroachdb/cockroach/pkg/util/humanizeutil"
@@ -361,3 +362,26 @@ func EnvOrDefaultDuration(name string, value time.Duration) time.Duration {
 	}
 	return value
 }
+
+// TestSetEnv sets an environment variable and the cleanup function
+// resets it to the original value.
+func TestSetEnv(t *testing.T, name string, value string) func() {
+	ClearEnvCache()
+	before, exists := os.LookupEnv(name)
+
+	if err := os.Setenv(name, value); err != nil {
+		t.Fatal(err)
+	}
+	return func() {
+		if exists {
+			if err := os.Setenv(name, before); err != nil {
+				t.Fatal(err)
+			}
+		} else {
+			if err := os.Unsetenv(name); err != nil {
+				t.Fatal(err)
+			}
+		}
+		ClearEnvCache()
+	}
+}

pkg/util/envutil/env_test.go

@@ -13,6 +13,8 @@ package envutil
 import (
 	"os"
 	"testing"
+
+	"github.com/stretchr/testify/require"
 )
 
 func TestEnvOrDefault(t *testing.T) {
@@ -27,3 +29,43 @@ func TestEnvOrDefault(t *testing.T) {
 		t.Errorf("expected %d, got %d", def, act)
 	}
 }
+
+func TestTestSetEnvExists(t *testing.T) {
+	key := "COCKROACH_ENVUTIL_TESTSETTING"
+	require.NoError(t, os.Setenv(key, "before"))
+
+	ClearEnvCache()
+	value, ok := EnvString(key, 0)
+	require.True(t, ok)
+	require.Equal(t, value, "before")
+
+	cleanup := TestSetEnv(t, key, "testvalue")
+	value, ok = EnvString(key, 0)
+	require.True(t, ok)
+	require.Equal(t, value, "testvalue")
+
+	cleanup()
+
+	value, ok = EnvString(key, 0)
+	require.True(t, ok)
+	require.Equal(t, value, "before")
+}
+
+func TestTestSetEnvDoesNotExist(t *testing.T) {
+	key := "COCKROACH_ENVUTIL_TESTSETTING"
+	require.NoError(t, os.Unsetenv(key))
+
+	ClearEnvCache()
+	_, ok := EnvString(key, 0)
+	require.False(t, ok)
+
+	cleanup := TestSetEnv(t, key, "foo")
+	value, ok := EnvString(key, 0)
+	require.True(t, ok)
+	require.Equal(t, value, "foo")
+
+	cleanup()
+
+	_, ok = EnvString(key, 0)
+	require.False(t, ok)
+}