From a4e49963f5a7417fc34a33a98e6b5a63eb416335 Mon Sep 17 00:00:00 2001
From: "Jonathan A. Sternberg" <jonathan@influxdb.com>
Date: Wed, 10 Aug 2016 15:15:31 -0500
Subject: [PATCH] Implement text/csv content encoding for the response writer

CSV doesn't offer a way to separate different sheets from each other and
it doesn't really have a standard format. We separate sheets with a
newline so they can be imported into something like Excel or LibreOffice
more easily.

The number of columns for each sheet is inferred from the first returned
row in each statement since they should all be the same.
---
 CHANGELOG.md                      |  1 +
 services/httpd/response_writer.go | 91 +++++++++++++++++++++++++++++++
 2 files changed, 92 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index ccd97a59192..a74a6563c7e 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -5,6 +5,7 @@
 - [#7120](https://github.com/influxdata/influxdb/issues/7120): Add additional statistics to query executor.
 - [#7135](https://github.com/influxdata/influxdb/pull/7135): Support enable HTTP service over unix domain socket. Thanks @oiooj
 - [#3634](https://github.com/influxdata/influxdb/issues/3634): Support mixed duration units.
+- [#7099](https://github.com/influxdata/influxdb/pull/7099): Implement text/csv content encoding for the response writer.
 
 ### Bugfixes
 
diff --git a/services/httpd/response_writer.go b/services/httpd/response_writer.go
index d4d4ab65ba6..8bdd2926eba 100644
--- a/services/httpd/response_writer.go
+++ b/services/httpd/response_writer.go
@@ -1,9 +1,14 @@
 package httpd
 
 import (
+	"encoding/csv"
 	"encoding/json"
 	"io"
 	"net/http"
+	"strconv"
+	"time"
+
+	"github.com/influxdata/influxdb/models"
 )
 
 // ResponseWriter is an interface for writing a response.
@@ -19,6 +24,8 @@ type ResponseWriter interface {
 func NewResponseWriter(w http.ResponseWriter, r *http.Request) ResponseWriter {
 	pretty := r.URL.Query().Get("pretty") == "true"
 	switch r.Header.Get("Accept") {
+	case "application/csv", "text/csv":
+		return &csvResponseWriter{statementID: -1, ResponseWriter: w}
 	case "application/json":
 		fallthrough
 	default:
@@ -62,3 +69,87 @@ func (w *jsonResponseWriter) Flush() {
 		w.Flush()
 	}
 }
+
+type csvResponseWriter struct {
+	statementID int
+	columns     []string
+	http.ResponseWriter
+}
+
+func (w *csvResponseWriter) WriteResponse(resp Response) (n int, err error) {
+	csv := csv.NewWriter(w)
+	for _, result := range resp.Results {
+		if result.StatementID != w.statementID {
+			// If there are no series in the result, skip past this result.
+			if len(result.Series) == 0 {
+				continue
+			}
+
+			// Set the statement id and print out a newline if this is not the first statement.
+			if w.statementID >= 0 {
+				// Flush the csv writer and write a newline.
+				csv.Flush()
+				if err := csv.Error(); err != nil {
+					return n, err
+				}
+
+				if out, err := io.WriteString(w, "\n"); err != nil {
+					return n, err
+				} else {
+					n += out
+				}
+			}
+			w.statementID = result.StatementID
+
+			// Print out the column headers from the first series.
+			w.columns = make([]string, 2+len(result.Series[0].Columns))
+			w.columns[0] = "name"
+			w.columns[1] = "tags"
+			copy(w.columns[2:], result.Series[0].Columns)
+			if err := csv.Write(w.columns); err != nil {
+				return n, err
+			}
+		}
+
+		for _, row := range result.Series {
+			w.columns[0] = row.Name
+			if len(row.Tags) > 0 {
+				w.columns[1] = string(models.Tags(row.Tags).HashKey()[1:])
+			} else {
+				w.columns[1] = ""
+			}
+			for _, values := range row.Values {
+				for i, value := range values {
+					switch v := value.(type) {
+					case float64:
+						w.columns[i+2] = strconv.FormatFloat(v, 'f', -1, 64)
+					case int64:
+						w.columns[i+2] = strconv.FormatInt(v, 10)
+					case string:
+						w.columns[i+2] = v
+					case bool:
+						if v {
+							w.columns[i+2] = "true"
+						} else {
+							w.columns[i+2] = "false"
+						}
+					case time.Time:
+						w.columns[i+2] = strconv.FormatInt(v.UnixNano(), 10)
+					}
+				}
+				csv.Write(w.columns)
+			}
+		}
+	}
+	csv.Flush()
+	if err := csv.Error(); err != nil {
+		return n, err
+	}
+	return n, nil
+}
+
+func (w *csvResponseWriter) Flush() {
+	if w, ok := w.ResponseWriter.(http.Flusher); ok {
+		w.Flush()
+	}
+}