From b2105646ec1ef3c4f52f6623b261221fd3374034 Mon Sep 17 00:00:00 2001
From: pingcap-github-bot <sre-bot@pingcap.com>
Date: Thu, 9 Apr 2020 11:38:39 +0800
Subject: [PATCH] ddl: fix create hash partition table makes tidb OOM (#16194)
 (#16219)

---
 ddl/db_partition_test.go | 11 +++++++++++
 ddl/partition.go         |  4 ++++
 2 files changed, 15 insertions(+)

diff --git a/ddl/db_partition_test.go b/ddl/db_partition_test.go
index 85aed311aedb4..86aa06f368718 100644
--- 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 employees (
+    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
index cb58884f3f79a..0ee90343a188c 100644
--- a/ddl/partition.go
+++ b/ddl/partition.go
@@ -122,6 +122,10 @@ func buildTablePartitionInfo(ctx sessionctx.Context, s *ast.CreateTableStmt) (*m
 }
 
 func buildHashPartitionDefinitions(ctx sessionctx.Context, s *ast.CreateTableStmt, pi *model.PartitionInfo) error {
+	if err := checkAddPartitionTooManyPartitions(pi.Num); err != nil {
+		return err
+	}
+
 	defs := make([]model.PartitionDefinition, pi.Num)
 	for i := 0; i < len(defs); i++ {
 		if len(s.Partition.Definitions) == 0 {