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 118921e1c..19cc9183a 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/mysql.go b/pkg/data/mysql.go index 1d362ed67..6899994ed 100644 --- a/pkg/data/mysql.go +++ b/pkg/data/mysql.go @@ -81,7 +81,9 @@ func InsertMySQLDataCMD(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/postgres.go b/pkg/data/postgres.go index 61c3af862..a44c194da 100644 --- a/pkg/data/postgres.go +++ b/pkg/data/postgres.go @@ -84,6 +84,9 @@ func InsertPostgresDataCMD(f cmdutil.Factory) *cobra.Command { log.Fatal("couldn't create tunnel, error: ", err) } 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) @@ -97,10 +100,6 @@ func InsertPostgresDataCMD(f cmdutil.Factory) *cobra.Command { } func (opts *postgresOpts) insertDataExecCmd(tunnel *portforward.Tunnel, rows int) error { - if rows <= 0 { - return fmt.Errorf("rows need to be greater than 0") - } - command := ` DO $$ BEGIN IF NOT EXISTS (SELECT 1 FROM information_schema.tables WHERE table_schema = 'public' AND table_name = 'appscode_kubedb_postgres_test_table') THEN 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 {