*: set default shard_row_id_bits and pre_split_regions for all tables (#55717)

close #55716

Author: jackysp

pkg/ddl/create_table.go

@@ -674,6 +674,12 @@ func BuildTableInfoWithStmt(ctx sessionctx.Context, s *ast.CreateTableStmt, dbCh
 		return nil, errors.Trace(err)
 	}
 
+	// set default shard row id bits and pre-split regions for table.
+	if !tbInfo.HasClusteredIndex() && tbInfo.TempTableType == model.TempTableNone {
+		tbInfo.ShardRowIDBits = ctx.GetSessionVars().ShardRowIDBits
+		tbInfo.PreSplitRegions = ctx.GetSessionVars().PreSplitRegions
+	}
+
 	if err = handleTableOptions(s.Options, tbInfo); err != nil {
 		return nil, errors.Trace(err)
 	}
@@ -799,8 +805,8 @@ func handleTableOptions(options []*ast.TableOption, tbInfo *model.TableInfo) err
 				return dbterror.ErrUnsupportedShardRowIDBits
 			}
 			tbInfo.ShardRowIDBits = op.UintValue
-			if tbInfo.ShardRowIDBits > shardRowIDBitsMax {
-				tbInfo.ShardRowIDBits = shardRowIDBitsMax
+			if tbInfo.ShardRowIDBits > variable.MaxShardRowIDBits {
+				tbInfo.ShardRowIDBits = variable.MaxShardRowIDBits
 			}
 			tbInfo.MaxShardRowIDBits = tbInfo.ShardRowIDBits
 		case ast.TableOptionPreSplitRegion:

pkg/ddl/ddl.go

@@ -77,8 +77,6 @@ const (
 	addingDDLJobPrefix = "/tidb/ddl/add_ddl_job_"
 	ddlPrompt          = "ddl"
 
-	shardRowIDBitsMax = 15
-
 	batchAddingJobs = 100
 
 	reorgWorkerCnt   = 10

pkg/ddl/executor.go

@@ -1848,8 +1848,8 @@ func (e *executor) AlterTable(ctx context.Context, sctx sessionctx.Context, stmt
 			for i, opt := range spec.Options {
 				switch opt.Tp {
 				case ast.TableOptionShardRowID:
-					if opt.UintValue > shardRowIDBitsMax {
-						opt.UintValue = shardRowIDBitsMax
+					if opt.UintValue > variable.MaxShardRowIDBits {
+						opt.UintValue = variable.MaxShardRowIDBits
 					}
 					err = e.ShardRowID(sctx, ident, opt.UintValue)
 				case ast.TableOptionAutoIncrement:

pkg/ddl/schematracker/checker.go

@@ -18,6 +18,7 @@ import (
 	"bytes"
 	"context"
 	"fmt"
+	"regexp"
 	"strings"
 	"sync/atomic"
 
@@ -164,8 +165,24 @@ func (d *Checker) checkTableInfo(ctx sessionctx.Context, dbName, tableName pmode
 	s1 := removeClusteredIndexComment(result.String())
 	s2 := removeClusteredIndexComment(result2.String())
 
+	// Remove shard_row_id_bits and pre_split_regions comments.
+	if ctx.GetSessionVars().ShardRowIDBits != 0 || ctx.GetSessionVars().PreSplitRegions != 0 {
+		removeShardPreSplitComment := func(s string) string {
+			pattern := ` \/\*T! SHARD_ROW_ID_BITS=.*?\*\/`
+			re := regexp.MustCompile(pattern)
+			ret := re.ReplaceAllString(s, "")
+			pattern = ` \/\*T! PRE_SPLIT_REGIONS=.*?\*\/`
+			re = regexp.MustCompile(pattern)
+			ret = re.ReplaceAllString(ret, "")
+			return ret
+		}
+
+		s1 = removeShardPreSplitComment(s1)
+		s2 = removeShardPreSplitComment(s2)
+	}
+
 	if s1 != s2 {
-		errStr := fmt.Sprintf("%s != %s", s1, s2)
+		errStr := fmt.Sprintf("%s\n!=\n%s", s1, s2)
 		panic(errStr)
 	}
 }

pkg/executor/test/ddl/BUILD.bazel

@@ -8,7 +8,7 @@ go_test(
         "main_test.go",
     ],
     flaky = True,
-    shard_count = 19,
+    shard_count = 20,
     deps = [
         "//pkg/config",
         "//pkg/ddl/schematracker",

pkg/executor/test/ddl/ddl_test.go

@@ -1106,3 +1106,41 @@ func TestRenameMultiTables(t *testing.T) {
 	tk.MustExec("drop database rename2")
 	tk.MustExec("drop database rename3")
 }
+
+func TestDefShardTables(t *testing.T) {
+	store := testkit.CreateMockStore(t, mockstore.WithDDLChecker())
+
+	tk := testkit.NewTestKit(t, store)
+
+	tk.MustExec("set @@session.tidb_enable_clustered_index = off")
+	tk.MustExec("set @@session.tidb_shard_row_id_bits = 4")
+	tk.MustExec("set @@session.tidb_pre_split_regions = 4")
+	tk.MustExec("use test")
+	tk.MustExec("create table t (i int primary key)")
+	result := tk.MustQuery("show create table t")
+	createSQL := result.Rows()[0][1]
+	expected := "CREATE TABLE `t` (\n  `i` int(11) NOT NULL,\n  PRIMARY KEY (`i`) /*T![clustered_index] NONCLUSTERED */\n) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin /*T! SHARD_ROW_ID_BITS=4 PRE_SPLIT_REGIONS=4 */"
+	require.Equal(t, expected, createSQL)
+
+	// test for manual setup shard_row_id_bits and pre_split_regions
+	tk.MustExec("create table t0 (i int primary key) /*T! SHARD_ROW_ID_BITS=2 PRE_SPLIT_REGIONS=2 */")
+	result = tk.MustQuery("show create table t0")
+	createSQL = result.Rows()[0][1]
+	expected = "CREATE TABLE `t0` (\n  `i` int(11) NOT NULL,\n  PRIMARY KEY (`i`) /*T![clustered_index] NONCLUSTERED */\n) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin /*T! SHARD_ROW_ID_BITS=2 PRE_SPLIT_REGIONS=2 */"
+	require.Equal(t, expected, createSQL)
+
+	// test for clustered index table
+	tk.MustExec("set @@session.tidb_enable_clustered_index = on")
+	tk.MustExec("create table t1 (i int primary key)")
+	result = tk.MustQuery("show create table t1")
+	createSQL = result.Rows()[0][1]
+	expected = "CREATE TABLE `t1` (\n  `i` int(11) NOT NULL,\n  PRIMARY KEY (`i`) /*T![clustered_index] CLUSTERED */\n) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin"
+	require.Equal(t, expected, createSQL)
+
+	// test for global temporary table
+	tk.MustExec("create global temporary table tengine (id int) engine = 'innodb' on commit delete rows")
+	result = tk.MustQuery("show create table tengine")
+	createSQL = result.Rows()[0][1]
+	expected = "CREATE GLOBAL TEMPORARY TABLE `tengine` (\n  `id` int(11) DEFAULT NULL\n) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin ON COMMIT DELETE ROWS"
+	require.Equal(t, expected, createSQL)
+}

pkg/sessionctx/variable/session.go

@@ -837,6 +837,12 @@ type SessionVars struct {
 	// status stands for the session status. e.g. in transaction or not, auto commit is on or off, and so on.
 	status atomic.Uint32
 
+	// ShardRowIDBits is the number of shard bits for user table row ID.
+	ShardRowIDBits uint64
+
+	// PreSplitRegions is the number of regions that should be pre-split for the table.
+	PreSplitRegions uint64
+
 	// ClientCapability is client's capability.
 	ClientCapability uint32
 

pkg/sessionctx/variable/sysvar.go

@@ -2060,6 +2060,14 @@ var defaultSysVars = []*SysVar{
 		s.AllowBatchCop = int(TidbOptInt64(val, DefTiDBAllowBatchCop))
 		return nil
 	}},
+	{Scope: ScopeGlobal | ScopeSession, Name: TiDBShardRowIDBits, Value: strconv.Itoa(DefShardRowIDBits), Type: TypeInt, MinValue: 0, MaxValue: MaxShardRowIDBits, SetSession: func(s *SessionVars, val string) error {
+		s.ShardRowIDBits = TidbOptUint64(val, DefShardRowIDBits)
+		return nil
+	}},
+	{Scope: ScopeGlobal | ScopeSession, Name: TiDBPreSplitRegions, Value: strconv.Itoa(DefPreSplitRegions), Type: TypeInt, MinValue: 0, MaxValue: MaxPreSplitRegions, SetSession: func(s *SessionVars, val string) error {
+		s.PreSplitRegions = TidbOptUint64(val, DefPreSplitRegions)
+		return nil
+	}},
 	{Scope: ScopeGlobal | ScopeSession, Name: TiDBInitChunkSize, Value: strconv.Itoa(DefInitChunkSize), Type: TypeUnsigned, MinValue: 1, MaxValue: initChunkSizeUpperBound, SetSession: func(s *SessionVars, val string) error {
 		s.InitChunkSize = tidbOptPositiveInt32(val, DefInitChunkSize)
 		return nil

pkg/sessionctx/variable/tidb_vars.go

@@ -387,6 +387,14 @@ const (
 	// The default value is 0
 	TiDBAllowBatchCop = "tidb_allow_batch_cop"
 
+	// TiDBShardRowIDBits means all the tables created in the current session will be sharded.
+	// The default value is 0
+	TiDBShardRowIDBits = "tidb_shard_row_id_bits"
+
+	// TiDBPreSplitRegions means all the tables created in the current session will be pre-splited.
+	// The default value is 0
+	TiDBPreSplitRegions = "tidb_pre_split_regions"
+
 	// TiDBAllowMPPExecution means if we should use mpp way to execute query or not.
 	// Default value is `true`, means to be determined by the optimizer.
 	// Value set to `false` means never use mpp.
@@ -1207,6 +1215,12 @@ const (
 	// MaxConfigurableConcurrency is the maximum number of "threads" (goroutines) that can be specified
 	// for any type of configuration item that has concurrent workers.
 	MaxConfigurableConcurrency = 256
+
+	// MaxShardRowIDBits is the maximum number of bits that can be used for row-id sharding.
+	MaxShardRowIDBits = 15
+
+	// MaxPreSplitRegions is the maximum number of regions that can be pre-split.
+	MaxPreSplitRegions = 15
 )
 
 // Default TiDB system variable values.
@@ -1279,6 +1293,8 @@ const (
 	DefTiDBEnableOuterJoinReorder           = true
 	DefTiDBEnableNAAJ                       = true
 	DefTiDBAllowBatchCop                    = 1
+	DefShardRowIDBits                       = 0
+	DefPreSplitRegions                      = 0
 	DefBlockEncryptionMode                  = "aes-128-ecb"
 	DefTiDBAllowMPPExecution                = true
 	DefTiDBAllowTiFlashCop                  = false