From 6d6007e680e9a324611be5dff643ac9b2cf2ed62 Mon Sep 17 00:00:00 2001
From: Lunny Xiao <xiaolunwen@gmail.com>
Date: Mon, 5 Jun 2017 12:26:08 +0800
Subject: [PATCH] add parse for composite keys

---
 dialect_sqlite3.go | 14 ++++++++++++++
 session_pk_test.go |  7 +++++++
 2 files changed, 21 insertions(+)

diff --git a/dialect_sqlite3.go b/dialect_sqlite3.go
index c32524f4..a55b1615 100644
--- a/dialect_sqlite3.go
+++ b/dialect_sqlite3.go
@@ -306,11 +306,25 @@ func (db *sqlite3) GetColumns(tableName string) ([]string, map[string]*core.Colu
 	for _, colStr := range colCreates {
 		reg = regexp.MustCompile(`,\s`)
 		colStr = reg.ReplaceAllString(colStr, ",")
+		if strings.HasPrefix(strings.TrimSpace(colStr), "PRIMARY KEY") {
+			parts := strings.Split(strings.TrimSpace(colStr), "(")
+			if len(parts) == 2 {
+				pkCols := strings.Split(strings.TrimRight(strings.TrimSpace(parts[1]), ")"), ",")
+				for _, pk := range pkCols {
+					if col, ok := cols[strings.Trim(strings.TrimSpace(pk), "`")]; ok {
+						col.IsPrimaryKey = true
+					}
+				}
+			}
+			continue
+		}
+
 		fields := strings.Fields(strings.TrimSpace(colStr))
 		col := new(core.Column)
 		col.Indexes = make(map[string]int)
 		col.Nullable = true
 		col.DefaultIsEmpty = true
+
 		for idx, field := range fields {
 			if idx == 0 {
 				col.Name = strings.Trim(strings.Trim(field, "`[] "), `"`)
diff --git a/session_pk_test.go b/session_pk_test.go
index 3030fca4..0dff9df7 100644
--- a/session_pk_test.go
+++ b/session_pk_test.go
@@ -1131,4 +1131,11 @@ func TestCompositePK(t *testing.T) {
 	assertSync(t, new(TaskSolution))
 
 	assert.NoError(t, testEngine.Sync2(new(TaskSolution)))
+	tables, err := testEngine.DBMetas()
+	assert.NoError(t, err)
+	assert.EqualValues(t, 1, len(tables))
+	pkCols := tables[0].PKColumns()
+	assert.EqualValues(t, 2, len(pkCols))
+	assert.EqualValues(t, "uid", pkCols[0].Name)
+	assert.EqualValues(t, "tid", pkCols[1].Name)
 }