From f17f82a57e16e35ecf076466736c7c0bc47a61b8 Mon Sep 17 00:00:00 2001
From: tiancaiamao <tiancaiamao@gmail.com>
Date: Wed, 8 Apr 2020 17:42:11 +0800
Subject: [PATCH] ddl: fix create hash partition table makes tidb OOM

---
 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 00c40e6eacc24..fad386d01c71b 100644
--- a/ddl/db_partition_test.go
+++ b/ddl/db_partition_test.go
@@ -291,6 +291,17 @@ func (s *testIntegrationSuite2) 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 *testIntegrationSuite1) TestCreateTableWithRangeColumnPartition(c *C) {
diff --git a/ddl/partition.go b/ddl/partition.go
index df32ccdf988db..6e0f0a6d74402 100644
--- a/ddl/partition.go
+++ b/ddl/partition.go
@@ -121,6 +121,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 {