samples: limit support added

BatchLabs · Sep 7, 2015 · e866a70 · e866a70
1 parent ad09b86
commit e866a70
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 0 deletions.
diff --git a/samples/csv/csv.go b/samples/csv/csv.go
@@ -42,6 +42,18 @@ func executeRequest(reader *csv.Reader, query *charlatan.Query) {
 	fmt.Println("$")
 
 	var header []string
+	var limit int64
+
+	hasLimit := query.HasLimit()
+
+	if hasLimit {
+		limit = query.Limit()
+
+		if limit <= 0 {
+			return
+		}
+	}
+
 	line := 0
 
 	for {
@@ -86,6 +98,13 @@ func executeRequest(reader *csv.Reader, query *charlatan.Query) {
 		}
 
 		fmt.Println("# ", values)
+
+		if hasLimit {
+			limit--
+			if limit <= 0 {
+				break
+			}
+		}
 	}
 }
 

diff --git a/samples/json/json.go b/samples/json/json.go
@@ -21,6 +21,8 @@ func fatalf(format string, args ...interface{}) {
 }
 
 func main() {
+	var limit int64
+
 	if len(os.Args) != 2 {
 		usage()
 	}
@@ -39,6 +41,16 @@ func main() {
 
 	decoder := json.NewDecoder(reader)
 
+	hasLimit := query.HasLimit()
+
+	if hasLimit {
+		limit = query.Limit()
+
+		if limit <= 0 {
+			return
+		}
+	}
+
 	line := 0
 
 	for {
@@ -72,5 +84,11 @@ func main() {
 		}
 
 		fmt.Println("# ", values)
+		if hasLimit {
+			limit--
+			if limit <= 0 {
+				break
+			}
+		}
 	}
 }