From 2e1b05d84aa8214a7e43d2ca6f818f7332fab5dc Mon Sep 17 00:00:00 2001
From: Shaad7 <abdullah.alshaad@appscode.com>
Date: Tue, 22 Aug 2023 14:36:11 +0600
Subject: [PATCH] Add limit to rows for insert command

Signed-off-by: Shaad7 <abdullah.alshaad@appscode.com>
---
 pkg/data/const.go         | 4 ++++
 pkg/data/elasticsearch.go | 5 ++---
 pkg/data/mariadb.go       | 7 +++----
 pkg/data/mongodb.go       | 4 ++--
 pkg/data/redis.go         | 3 +++
 5 files changed, 14 insertions(+), 9 deletions(-)

diff --git a/pkg/data/const.go b/pkg/data/const.go
index bb168da95..eb6254d59 100644
--- a/pkg/data/const.go
+++ b/pkg/data/const.go
@@ -25,3 +25,7 @@ const (
 const (
 	actor = "kubedb-cli"
 )
+
+const (
+	maxRows = 100000
+)
diff --git a/pkg/data/elasticsearch.go b/pkg/data/elasticsearch.go
index 92d277f84..ec1f6cb62 100644
--- a/pkg/data/elasticsearch.go
+++ b/pkg/data/elasticsearch.go
@@ -90,10 +90,9 @@ func InsertElasticsearchDataCMD(f cmdutil.Factory) *cobra.Command {
 				log.Fatalln(err)
 			}
 
-			if rows <= 0 {
-				log.Fatal("rows need to be greater than 0")
+			if rows < 1 || rows > maxRows {
+				log.Fatalf("rows need to be between 1 and %d", maxRows)
 			}
-
 			err = opts.insertDataInDatabase(rows)
 			if err != nil {
 				log.Fatal(err)
diff --git a/pkg/data/mariadb.go b/pkg/data/mariadb.go
index 12ba76b31..fb0b7c5be 100644
--- a/pkg/data/mariadb.go
+++ b/pkg/data/mariadb.go
@@ -71,10 +71,6 @@ func InsertMariaDBDataCMD(f cmdutil.Factory) *cobra.Command {
 				log.Fatalln(err)
 			}
 
-			if rows <= 0 {
-				log.Fatal("Inserted rows must be greater than 0")
-			}
-
 			tunnel, err := lib.TunnelToDBService(opts.config, dbName, namespace, api.MySQLDatabasePort)
 			if err != nil {
 				log.Fatal("couldn't create tunnel, error: ", err)
@@ -82,6 +78,9 @@ func InsertMariaDBDataCMD(f cmdutil.Factory) *cobra.Command {
 
 			defer tunnel.Close()
 
+			if rows < 1 || rows > maxRows {
+				log.Fatalf("rows need to be between 1 and %d", maxRows)
+			}
 			err = opts.insertDataExecCmd(tunnel, rows)
 			if err != nil {
 				log.Fatal(err)
diff --git a/pkg/data/mongodb.go b/pkg/data/mongodb.go
index 1270dfc84..cb65b4490 100644
--- a/pkg/data/mongodb.go
+++ b/pkg/data/mongodb.go
@@ -76,8 +76,8 @@ func InsertMongoDBDataCMD(f cmdutil.Factory) *cobra.Command {
 				log.Fatalln(err)
 			}
 
-			if rows <= 0 {
-				log.Fatal("rows need to be greater than 0")
+			if rows < 1 || rows > maxRows {
+				log.Fatalf("rows need to be between 1 and %d", maxRows)
 			}
 
 			command := fmt.Sprintf("for(var i=1;i<=%d;i++){db[\"%s\"].insert({_id:\"doc\"+i,actor:\"%s\"})}", rows, mgCollectionName, actor)
diff --git a/pkg/data/redis.go b/pkg/data/redis.go
index 592aa8ed7..391e1ce08 100644
--- a/pkg/data/redis.go
+++ b/pkg/data/redis.go
@@ -115,6 +115,9 @@ func InsertRedisDataCMD(f cmdutil.Factory) *cobra.Command {
 			if err != nil {
 				log.Fatalln(err)
 			}
+			if rows < 1 || rows > maxRows {
+				log.Fatalf("rows need to be between 1 and %d", maxRows)
+			}
 
 			err = opts.insertDataInDatabase(rows)
 			if err != nil {