From dbc493df5e34bf9a5b111c6fbe5e626d99a87134 Mon Sep 17 00:00:00 2001
From: Lunny Xiao <xiaolunwen@gmail.com>
Date: Mon, 24 Jul 2017 22:09:47 +0800
Subject: [PATCH] multiple Cols support (#653)

---
 session_cols_test.go | 32 ++++++++++++++++++++++++++++++++
 statement.go         | 18 +++++++++++++++++-
 2 files changed, 49 insertions(+), 1 deletion(-)

diff --git a/session_cols_test.go b/session_cols_test.go
index 331052814..f742a86a9 100644
--- a/session_cols_test.go
+++ b/session_cols_test.go
@@ -35,3 +35,35 @@ func TestSetExpr(t *testing.T) {
 	assert.NoError(t, err)
 	assert.EqualValues(t, 1, cnt)
 }
+
+func TestCols(t *testing.T) {
+	assert.NoError(t, prepareEngine())
+
+	type ColsTable struct {
+		Id   int64
+		Col1 string
+		Col2 string
+	}
+
+	assertSync(t, new(ColsTable))
+
+	_, err := testEngine.Insert(&ColsTable{
+		Col1: "1",
+		Col2: "2",
+	})
+	assert.NoError(t, err)
+
+	sess := testEngine.ID(1)
+	_, err = sess.Cols("col1").Cols("col2").Update(&ColsTable{
+		Col1: "",
+		Col2: "",
+	})
+	assert.NoError(t, err)
+
+	var tb ColsTable
+	has, err := testEngine.ID(1).Get(&tb)
+	assert.NoError(t, err)
+	assert.True(t, has)
+	assert.EqualValues(t, "", tb.Col1)
+	assert.EqualValues(t, "", tb.Col2)
+}
diff --git a/statement.go b/statement.go
index 0f90f64d7..c97dea50a 100644
--- a/statement.go
+++ b/statement.go
@@ -592,6 +592,22 @@ func (statement *Statement) col2NewColsWithQuote(columns ...string) []string {
 	return newColumns
 }
 
+func (statement *Statement) colmap2NewColsWithQuote() []string {
+	newColumns := make([]string, 0, len(statement.columnMap))
+	for col := range statement.columnMap {
+		fields := strings.Split(strings.TrimSpace(col), ".")
+		if len(fields) == 1 {
+			newColumns = append(newColumns, statement.Engine.quote(fields[0]))
+		} else if len(fields) == 2 {
+			newColumns = append(newColumns, statement.Engine.quote(fields[0])+"."+
+				statement.Engine.quote(fields[1]))
+		} else {
+			panic(errors.New("unwanted colnames"))
+		}
+	}
+	return newColumns
+}
+
 // Distinct generates "DISTINCT col1, col2 " statement
 func (statement *Statement) Distinct(columns ...string) *Statement {
 	statement.IsDistinct = true
@@ -618,7 +634,7 @@ func (statement *Statement) Cols(columns ...string) *Statement {
 		statement.columnMap[strings.ToLower(nc)] = true
 	}
 
-	newColumns := statement.col2NewColsWithQuote(columns...)
+	newColumns := statement.colmap2NewColsWithQuote()
 	statement.ColumnStr = strings.Join(newColumns, ", ")
 	statement.ColumnStr = strings.Replace(statement.ColumnStr, statement.Engine.quote("*"), "*", -1)
 	return statement