ddl: fix create hash partition table makes tidb OOM (#16194) (#16218)

pingcap · Apr 14, 2020 · e729528 · e729528
1 parent dc248ea
commit e729528
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 0 deletions.
diff --git a/ddl/db_partition_test.go b/ddl/db_partition_test.go
@@ -292,6 +292,17 @@ func (s *testIntegrationSuite7) TestCreateTableWithHashPartition(c *C) {
 		store_id int
 	)
 	partition by hash( year(hired) ) partitions 4;`)
+
+	// This query makes tidb OOM without partition count check.
+	tk.MustGetErrCode(`CREATE TABLE employees1 (
+    id INT NOT NULL,
+    fname VARCHAR(30),
+    lname VARCHAR(30),
+    hired DATE NOT NULL DEFAULT '1970-01-01',
+    separated DATE NOT NULL DEFAULT '9999-12-31',
+    job_code INT,
+    store_id INT
+) PARTITION BY HASH(store_id) PARTITIONS 102400000000;`, tmysql.ErrTooManyPartitions)
 }
 
 func (s *testIntegrationSuite10) TestCreateTableWithRangeColumnPartition(c *C) {

diff --git a/ddl/partition.go b/ddl/partition.go
@@ -118,10 +118,14 @@ func buildTablePartitionInfo(ctx sessionctx.Context, d *ddl, s *ast.CreateTableS
 }
 
 func buildHashPartitionDefinitions(ctx sessionctx.Context, d *ddl, s *ast.CreateTableStmt, pi *model.PartitionInfo) error {
+	if err := checkAddPartitionTooManyPartitions(pi.Num); err != nil {
+		return err
+	}
 	genIDs, err := d.genGlobalIDs(int(pi.Num))
 	if err != nil {
 		return errors.Trace(err)
 	}
+
 	defs := make([]model.PartitionDefinition, pi.Num)
 	for i := 0; i < len(defs); i++ {
 		defs[i].ID = genIDs[i]