feat: enforce TiDB Cloud free plan limits

cmd/drive9-server/main.go

@@ -207,6 +207,10 @@ func main() {
 	if err != nil {
 		die(err)
 	}
+	tidbCloudNonFreePlanCacheTTL, err := tiDBCloudNonFreePlanCacheTTLFromEnv()
+	if err != nil {
+		die(err)
+	}
 	maxUploadBytes := server.DefaultMaxUploadBytes
 	if raw := os.Getenv("DRIVE9_MAX_UPLOAD_BYTES"); raw != "" {
 		maxUploadBytes, err = strconv.ParseInt(raw, 10, 64)
@@ -219,6 +223,15 @@ func main() {
 	}
 	backendOptions.MaxUploadBytes = maxUploadBytes
 	meta.SetDefaultMaxFileSizeBytes(maxUploadBytes)
+	maxTenantFileCount, err := maxTenantFileCountFromEnv()
+	if err != nil {
+		die(err)
+	}
+	meta.SetDefaultMaxFileCount(maxTenantFileCount)
+	tidbCloudFreePlanLimits, err := tiDBCloudFreePlanLimitsFromEnv(maxUploadBytes)
+	if err != nil {
+		die(err)
+	}
 	providerType, err = tenant.NormalizeProvider(providerType)
 	if err != nil {
 		die(err)
@@ -244,10 +257,10 @@ func main() {
 		provisioner, provisionerErr = db9.NewProvisionerFromEnv()
 	}
 	if provisionerErr == nil && providerType == tenant.ProviderTiDBCloudNativeShared {
-		validator, ok := provisioner.(interface{ ValidateSharedCredentials(context.Context) error })
+		validator, ok := provisioner.(interface{ ValidateSharedAccess(context.Context) error })
 		if !ok {
-			provisionerErr = fmt.Errorf("shared TiDB Cloud credential validation is not supported")
-		} else if err := validator.ValidateSharedCredentials(context.Background()); err != nil {
+			provisionerErr = fmt.Errorf("shared TiDB Cloud access validation is not supported")
+		} else if err := validator.ValidateSharedAccess(context.Background()); err != nil {
 			provisionerErr = err
 		}
 	}
@@ -438,14 +451,16 @@ func main() {
 	}
 
 	srv := server.NewWithConfig(server.Config{
-		Meta:                  store,
-		Pool:                  pool,
-		Provisioner:           provisioner,
-		DefaultTenantProvider: providerType,
-		SharedDBMaxTenants:    sharedDBMaxTenants,
-		SharedDBHardCapRatio:  sharedDBHardCapRatio,
-		SharedDBReopenRatio:   sharedDBReopenRatio,
-		SharedDBSpendingLimit: sharedDBDefaultSpendingLimit,
+		Meta:                         store,
+		Pool:                         pool,
+		Provisioner:                  provisioner,
+		DefaultTenantProvider:        providerType,
+		SharedDBMaxTenants:           sharedDBMaxTenants,
+		SharedDBHardCapRatio:         sharedDBHardCapRatio,
+		SharedDBReopenRatio:          sharedDBReopenRatio,
+		SharedDBSpendingLimit:        sharedDBDefaultSpendingLimit,
+		TiDBCloudNonFreePlanCacheTTL: tidbCloudNonFreePlanCacheTTL,
+		TiDBCloudFreePlanLimits:      tidbCloudFreePlanLimits,
 		ManagedSharedDBCloudBatchSize: envInt("DRIVE9_TIDBCLOUD_NATIVE_SHARED_CLOUD_BATCH_SIZE",
 			server.DefaultManagedSharedDBCloudBatchSize),
 		ManagedSharedDBMetadataWorkers: envInt("DRIVE9_TIDBCLOUD_NATIVE_SHARED_METADATA_WORKERS",
@@ -1266,6 +1281,82 @@ func sharedDBDefaultSpendingLimitFromEnv() (int64, error) {
 	return v, nil
 }
 
+func tiDBCloudNonFreePlanCacheTTLFromEnv() (time.Duration, error) {
+	const key = "DRIVE9_TIDBCLOUD_NON_FREE_PLAN_CACHE_TTL"
+	raw := strings.TrimSpace(os.Getenv(key))
+	if raw == "" {
+		return 30 * time.Minute, nil
+	}
+	ttl, err := time.ParseDuration(raw)
+	if err != nil || ttl <= 0 {
+		return 0, fmt.Errorf("invalid %s=%q: must be a positive duration", key, raw)
+	}
+	return ttl, nil
+}
+
+func tiDBCloudFreePlanLimitsFromEnv(maxUploadBytes int64) (server.TiDBCloudFreePlanLimits, error) {
+	limits := server.DefaultTiDBCloudFreePlanLimits()
+	var err error
+	if limits.TenantCount, err = positiveIntEnv("DRIVE9_TIDBCLOUD_FREE_TENANT_COUNT", limits.TenantCount); err != nil {
+		return server.TiDBCloudFreePlanLimits{}, err
+	}
+	if limits.MaxStorageBytes, err = positiveInt64Env("DRIVE9_TIDBCLOUD_FREE_MAX_STORAGE_BYTES", limits.MaxStorageBytes); err != nil {
+		return server.TiDBCloudFreePlanLimits{}, err
+	}
+	const freeMaxFileSizeKey = "DRIVE9_TIDBCLOUD_FREE_MAX_FILE_SIZE_BYTES"
+	if strings.TrimSpace(os.Getenv(freeMaxFileSizeKey)) == "" {
+		if limits.MaxFileSizeBytes > maxUploadBytes {
+			limits.MaxFileSizeBytes = maxUploadBytes
+		}
+	} else if limits.MaxFileSizeBytes, err = positiveInt64Env(freeMaxFileSizeKey, limits.MaxFileSizeBytes); err != nil {
+		return server.TiDBCloudFreePlanLimits{}, err
+	}
+	if limits.MaxFileCount, err = positiveInt64Env("DRIVE9_TIDBCLOUD_FREE_MAX_FILE_COUNT", limits.MaxFileCount); err != nil {
+		return server.TiDBCloudFreePlanLimits{}, err
+	}
+	if limits.MaxFileSizeBytes > maxUploadBytes {
+		return server.TiDBCloudFreePlanLimits{}, fmt.Errorf("DRIVE9_TIDBCLOUD_FREE_MAX_FILE_SIZE_BYTES must not exceed DRIVE9_MAX_UPLOAD_BYTES")
+	}
+	return limits, nil
+}
+
+func maxTenantFileCountFromEnv() (int64, error) {
+	const key = "DRIVE9_MAX_TENANT_FILE_COUNT"
+	raw := strings.TrimSpace(os.Getenv(key))
+	if raw == "" {
+		return 0, nil
+	}
+	value, err := strconv.ParseInt(raw, 10, 64)
+	if err != nil || value < 0 {
+		return 0, fmt.Errorf("invalid %s=%q: must be a non-negative integer", key, raw)
+	}
+	return value, nil
+}
+
+func positiveIntEnv(key string, fallback int) (int, error) {
+	raw := strings.TrimSpace(os.Getenv(key))
+	if raw == "" {
+		return fallback, nil
+	}
+	value, err := strconv.Atoi(raw)
+	if err != nil || value <= 0 {
+		return 0, fmt.Errorf("invalid %s=%q: must be a positive integer", key, raw)
+	}
+	return value, nil
+}
+
+func positiveInt64Env(key string, fallback int64) (int64, error) {
+	raw := strings.TrimSpace(os.Getenv(key))
+	if raw == "" {
+		return fallback, nil
+	}
+	value, err := strconv.ParseInt(raw, 10, 64)
+	if err != nil || value <= 0 {
+		return 0, fmt.Errorf("invalid %s=%q: must be a positive integer", key, raw)
+	}
+	return value, nil
+}
+
 func envDuration(key string, fallback time.Duration) time.Duration {
 	raw := strings.TrimSpace(os.Getenv(key))
 	if raw == "" {

cmd/drive9-server/main_test.go

@@ -331,6 +331,104 @@ func TestSharedDBCapacityConfigFromEnv(t *testing.T) {
 	}
 }
 
+func TestTiDBCloudNonFreePlanCacheTTLFromEnv(t *testing.T) {
+	const key = "DRIVE9_TIDBCLOUD_NON_FREE_PLAN_CACHE_TTL"
+	restore := snapshotEnv(t, []string{key})
+	t.Cleanup(func() { restoreEnv(t, restore) })
+
+	unsetEnv(t, []string{key})
+	got, err := tiDBCloudNonFreePlanCacheTTLFromEnv()
+	if err != nil || got != 30*time.Minute {
+		t.Fatalf("default cache TTL = %v/%v, want 30m/nil", got, err)
+	}
+	setEnv(t, key, "45m")
+	got, err = tiDBCloudNonFreePlanCacheTTLFromEnv()
+	if err != nil || got != 45*time.Minute {
+		t.Fatalf("configured cache TTL = %v/%v, want 45m/nil", got, err)
+	}
+	for _, raw := range []string{"0", "-1s", "bad"} {
+		setEnv(t, key, raw)
+		if _, err := tiDBCloudNonFreePlanCacheTTLFromEnv(); err == nil {
+			t.Fatalf("cache TTL %q error = nil, want error", raw)
+		}
+	}
+}
+
+func TestTiDBCloudFreePlanLimitsFromEnv(t *testing.T) {
+	keys := []string{
+		"DRIVE9_TIDBCLOUD_FREE_TENANT_COUNT",
+		"DRIVE9_TIDBCLOUD_FREE_MAX_STORAGE_BYTES",
+		"DRIVE9_TIDBCLOUD_FREE_MAX_FILE_SIZE_BYTES",
+		"DRIVE9_TIDBCLOUD_FREE_MAX_FILE_COUNT",
+	}
+	restore := snapshotEnv(t, keys)
+	t.Cleanup(func() { restoreEnv(t, restore) })
+
+	unsetEnv(t, keys)
+	got, err := tiDBCloudFreePlanLimitsFromEnv(server.DefaultMaxUploadBytes)
+	if err != nil {
+		t.Fatalf("default free limits: %v", err)
+	}
+	want := server.TiDBCloudFreePlanLimits{
+		TenantCount: 5, MaxStorageBytes: 5368709120,
+		MaxFileSizeBytes: 524288000, MaxFileCount: 500,
+	}
+	if got != want {
+		t.Fatalf("default free limits = %+v, want %+v", got, want)
+	}
+
+	got, err = tiDBCloudFreePlanLimitsFromEnv(1 << 20)
+	if err != nil {
+		t.Fatalf("default free max file size with smaller upload cap: %v", err)
+	}
+	if got.MaxFileSizeBytes != 1<<20 {
+		t.Fatalf("default free max file size = %d, want upload cap %d", got.MaxFileSizeBytes, 1<<20)
+	}
+
+	setEnv(t, keys[0], "5")
+	setEnv(t, keys[1], "4096")
+	setEnv(t, keys[2], "1024")
+	setEnv(t, keys[3], "2000")
+	got, err = tiDBCloudFreePlanLimitsFromEnv(server.DefaultMaxUploadBytes)
+	if err != nil || got.TenantCount != 5 || got.MaxStorageBytes != 4096 || got.MaxFileSizeBytes != 1024 || got.MaxFileCount != 2000 {
+		t.Fatalf("configured free limits = %+v/%v", got, err)
+	}
+
+	for i, key := range keys {
+		unsetEnv(t, keys)
+		setEnv(t, key, "0")
+		if _, err := tiDBCloudFreePlanLimitsFromEnv(server.DefaultMaxUploadBytes); err == nil {
+			t.Fatalf("zero %s (index %d) error = nil", key, i)
+		}
+	}
+	unsetEnv(t, keys)
+	setEnv(t, keys[2], "1048577")
+	if _, err := tiDBCloudFreePlanLimitsFromEnv(1048576); err == nil {
+		t.Fatal("free max file size above upload limit error = nil")
+	}
+}
+
+func TestGlobalMaxTenantFileCountFromEnv(t *testing.T) {
+	const key = "DRIVE9_MAX_TENANT_FILE_COUNT"
+	restore := snapshotEnv(t, []string{key})
+	t.Cleanup(func() { restoreEnv(t, restore) })
+
+	unsetEnv(t, []string{key})
+	if got, err := maxTenantFileCountFromEnv(); err != nil || got != 0 {
+		t.Fatalf("default max file count = %d/%v, want 0/nil", got, err)
+	}
+	setEnv(t, key, "2500")
+	if got, err := maxTenantFileCountFromEnv(); err != nil || got != 2500 {
+		t.Fatalf("configured max file count = %d/%v, want 2500/nil", got, err)
+	}
+	for _, raw := range []string{"-1", "bad"} {
+		setEnv(t, key, raw)
+		if _, err := maxTenantFileCountFromEnv(); err == nil {
+			t.Fatalf("max file count %q error = nil", raw)
+		}
+	}
+}
+
 func TestDBHealthProbeOptionsFromEnvDefaults(t *testing.T) {
 	keys := []string{
 		"DRIVE9_DB_HEALTH_PROBE_META_ENABLED",

pkg/meta/finalize_tenant_connection_test.go

@@ -0,0 +1,64 @@
+package meta
+
+import (
+	"context"
+	"database/sql"
+	"database/sql/driver"
+	"errors"
+	"fmt"
+	"sync/atomic"
+	"testing"
+)
+
+var finalizeRowsAffectedDriverSequence atomic.Int64
+
+type finalizeRowsAffectedErrorDriver struct {
+	err error
+}
+
+func (d finalizeRowsAffectedErrorDriver) Open(string) (driver.Conn, error) {
+	return finalizeRowsAffectedErrorConn(d), nil
+}
+
+type finalizeRowsAffectedErrorConn struct {
+	err error
+}
+
+func (c finalizeRowsAffectedErrorConn) Prepare(string) (driver.Stmt, error) {
+	return nil, errors.New("prepare is not supported")
+}
+
+func (c finalizeRowsAffectedErrorConn) Close() error { return nil }
+
+func (c finalizeRowsAffectedErrorConn) Begin() (driver.Tx, error) {
+	return nil, errors.New("transactions are not supported")
+}
+
+func (c finalizeRowsAffectedErrorConn) ExecContext(context.Context, string, []driver.NamedValue) (driver.Result, error) {
+	return finalizeRowsAffectedErrorResult(c), nil
+}
+
+type finalizeRowsAffectedErrorResult struct {
+	err error
+}
+
+func (r finalizeRowsAffectedErrorResult) LastInsertId() (int64, error) { return 0, nil }
+
+func (r finalizeRowsAffectedErrorResult) RowsAffected() (int64, error) { return 0, r.err }
+
+func TestFinalizeTenantConnectionReturnsRowsAffectedError(t *testing.T) {
+	wantErr := errors.New("rows affected unavailable")
+	driverName := fmt.Sprintf("finalize-rows-affected-error-%d", finalizeRowsAffectedDriverSequence.Add(1))
+	sql.Register(driverName, finalizeRowsAffectedErrorDriver{err: wantErr})
+	db, err := sql.Open(driverName, "")
+	if err != nil {
+		t.Fatal(err)
+	}
+	t.Cleanup(func() { _ = db.Close() })
+
+	s := &Store{db: db}
+	updated, err := s.FinalizeTenantConnection(context.Background(), "tenant-1", TenantPending, TenantProvisioning, &Tenant{})
+	if updated || !errors.Is(err, wantErr) {
+		t.Fatalf("FinalizeTenantConnection updated=%v err=%v, want false and rows-affected error", updated, err)
+	}
+}