From 54e4250df9675a03a42fd34591d9f854736690ab Mon Sep 17 00:00:00 2001
From: Brad Rydzewski <brad.rydzewski@gmail.com>
Date: Sun, 12 Apr 2015 23:08:55 -0700
Subject: [PATCH] ability to limit (or not) the log output

---
 datastore/bolt/task.go  | 16 +++++++++-------
 datastore/bolt/token.go |  1 +
 datastore/datastore.go  |  8 ++++++--
 server/tasks.go         |  6 +++++-
 4 files changed, 21 insertions(+), 10 deletions(-)

diff --git a/datastore/bolt/task.go b/datastore/bolt/task.go
index 684d2f4e0..54378d73e 100644
--- a/datastore/bolt/task.go
+++ b/datastore/bolt/task.go
@@ -1,10 +1,12 @@
 package bolt
 
 import (
+	"bytes"
+	"io"
 	"strconv"
 
-	"github.com/drone/drone/common"
 	"github.com/boltdb/bolt"
+	"github.com/drone/drone/common"
 )
 
 // GetTask gets the task at index N for the named
@@ -12,7 +14,7 @@ import (
 func (db *DB) GetTask(repo string, build int, task int) (*common.Task, error) {
 	key := []byte(repo + "/" + strconv.Itoa(build) + "/" + strconv.Itoa(task))
 	task_ := &common.Task{}
-	err := db.View(func (t *bolt.Tx) error {
+	err := db.View(func(t *bolt.Tx) error {
 		return get(t, bucketBuildTasks, key, task_)
 	})
 	return task_, err
@@ -20,17 +22,17 @@ func (db *DB) GetTask(repo string, build int, task int) (*common.Task, error) {
 
 // GetTaskLogs gets the task logs at index N for
 // the named repository and build number.
-func (db *DB) GetTaskLogs(repo string, build int, task int) ([]byte, error) {
+func (db *DB) GetTaskLogs(repo string, build int, task int) (io.Reader, error) {
 	key := []byte(repo + "/" + strconv.Itoa(build) + "/" + strconv.Itoa(task))
 
 	var log []byte
-	err := db.View(func (t *bolt.Tx) error {
+	err := db.View(func(t *bolt.Tx) error {
 		var err error
 		log, err = raw(t, bucketBuildLogs, key)
 		return err
 	})
-
-	return log, err
+	buf := bytes.NewBuffer(log)
+	return buf, err
 }
 
 // GetTaskList gets all tasks for the named repository
@@ -72,7 +74,7 @@ func (db *DB) GetTaskList(repo string, build int) ([]*common.Task, error) {
 // repository and build number.
 func (db *DB) UpsertTask(repo string, build int, task *common.Task) error {
 	key := []byte(repo + "/" + strconv.Itoa(build) + "/" + strconv.Itoa(task.Number))
-	return db.Update(func (t *bolt.Tx) error {
+	return db.Update(func(t *bolt.Tx) error {
 		return update(t, bucketBuildTasks, key, task)
 	})
 }
diff --git a/datastore/bolt/token.go b/datastore/bolt/token.go
index ca5dec049..1b28728f5 100644
--- a/datastore/bolt/token.go
+++ b/datastore/bolt/token.go
@@ -33,4 +33,5 @@ func (db *DB) DeleteToken(token *common.Token) error {
 	return db.Update(func(t *bolt.Tx) error {
 		return delete(t, bucketUser, key)
 	})
+	// TODO(bradrydzewski) remove token from users_token index
 }
diff --git a/datastore/datastore.go b/datastore/datastore.go
index 5371101ff..4c200fca8 100644
--- a/datastore/datastore.go
+++ b/datastore/datastore.go
@@ -1,6 +1,10 @@
 package datastore
 
-import "github.com/drone/drone/common"
+import (
+	"io"
+
+	"github.com/drone/drone/common"
+)
 
 type Datastore interface {
 	// GetUser gets a user by user login.
@@ -123,7 +127,7 @@ type Datastore interface {
 
 	// GetTaskLogs gets the task logs at index N for
 	// the named repository and build number.
-	GetTaskLogs(string, int, int) ([]byte, error)
+	GetTaskLogs(string, int, int) (io.Reader, error)
 
 	// GetTaskList gets all tasks for the named repository
 	// and build number.
diff --git a/server/tasks.go b/server/tasks.go
index 789904ad2..cdb2268c1 100644
--- a/server/tasks.go
+++ b/server/tasks.go
@@ -1,6 +1,7 @@
 package server
 
 import (
+	"io"
 	"strconv"
 
 	"github.com/gin-gonic/gin"
@@ -54,13 +55,16 @@ func GetTasks(c *gin.Context) {
 func GetTaskLogs(c *gin.Context) {
 	ds := ToDatastore(c)
 	repo := ToRepo(c)
+	full, _ := strconv.ParseBool(c.Params.ByName("full"))
 	build, _ := strconv.Atoi(c.Params.ByName("number"))
 	task, _ := strconv.Atoi(c.Params.ByName("task"))
 
 	logs, err := ds.GetTaskLogs(repo.FullName, build, task)
 	if err != nil {
 		c.Fail(404, err)
+	} else if full {
+		io.Copy(c.Writer, logs)
 	} else {
-		c.Writer.Write(logs)
+		io.Copy(c.Writer, io.LimitReader(logs, 2000000))
 	}
 }