mt_start_sql: enable enterprise features for multitenant sql servers

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)
+}