From 0a47d758708911ce02afaee1bc3bcbd6beaa87c3 Mon Sep 17 00:00:00 2001
From: Jinzhu <wosmvp@gmail.com>
Date: Wed, 30 Oct 2013 15:21:58 +0800
Subject: [PATCH] Add support for Attrs

---
 chain.go     | 10 ++++++++--
 do.go        |  7 +++++++
 gorm_test.go | 31 +++++++++++++++++++++++++++++--
 main.go      |  4 ++++
 4 files changed, 48 insertions(+), 4 deletions(-)

diff --git a/chain.go b/chain.go
index f38ba241..2c54236a 100644
--- a/chain.go
+++ b/chain.go
@@ -20,6 +20,7 @@ type Chain struct {
 
 	whereClause        []map[string]interface{}
 	orClause           []map[string]interface{}
+	initAttrs          []interface{}
 	selectStr          string
 	orderStrs          []string
 	offsetStr          string
@@ -175,9 +176,14 @@ func (s *Chain) First(out interface{}, where ...interface{}) *Chain {
 	return s
 }
 
+func (s *Chain) Attrs(attrs interface{}) *Chain {
+	s.initAttrs = append(s.initAttrs, attrs)
+	return s
+}
+
 func (s *Chain) FirstOrInit(out interface{}, where ...interface{}) *Chain {
 	if s.First(out, where...).Error != nil {
-		s.do(out).where(where...).initializeWithSearchCondition()
+		s.do(out).where(where...).where(s.initAttrs).initializeWithSearchCondition()
 		s.deleteLastError()
 	}
 	return s
@@ -185,7 +191,7 @@ func (s *Chain) FirstOrInit(out interface{}, where ...interface{}) *Chain {
 
 func (s *Chain) FirstOrCreate(out interface{}, where ...interface{}) *Chain {
 	if s.First(out, where...).Error != nil {
-		s.do(out).where(where...).initializeWithSearchCondition()
+		s.do(out).where(where...).where(s.initAttrs).initializeWithSearchCondition()
 		s.deleteLastError()
 		s.Save(out)
 	}
diff --git a/do.go b/do.go
index 641abf87..73e2801f 100644
--- a/do.go
+++ b/do.go
@@ -482,6 +482,13 @@ func (s *Do) initializeWithSearchCondition() {
 	for _, clause := range s.whereClause {
 		query := clause["query"]
 		switch query.(type) {
+		case []interface{}:
+			for _, obj := range query.([]interface{}) {
+				m := &Model{data: obj, driver: s.driver}
+				for _, field := range m.columnsHasValue("") {
+					m.setValueByColumn(field.DbName, field.Value, s.value)
+				}
+			}
 		case map[string]interface{}:
 			for key, value := range query.(map[string]interface{}) {
 				m.setValueByColumn(key, value, s.value)
diff --git a/gorm_test.go b/gorm_test.go
index f726314d..46ccd3c2 100644
--- a/gorm_test.go
+++ b/gorm_test.go
@@ -806,7 +806,7 @@ func TestSoftDelete(t *testing.T) {
 }
 
 func TestFindOrInitialize(t *testing.T) {
-	var user1, user2, user3 User
+	var user1, user2, user3, user4, user5, user6 User
 	db.Where(&User{Name: "find or init", Age: 33}).FirstOrInit(&user1)
 	if user1.Name != "find or init" || user1.Id != 0 || user1.Age != 33 {
 		t.Errorf("user should be initialized with search value")
@@ -821,10 +821,27 @@ func TestFindOrInitialize(t *testing.T) {
 	if user3.Name != "find or init 2" || user3.Id != 0 {
 		t.Errorf("user should be initialized with inline search value")
 	}
+
+	db.Where(&User{Name: "find or init"}).Attrs(User{Age: 44}).FirstOrInit(&user4)
+	if user4.Name != "find or init" || user4.Id != 0 || user4.Age != 44 {
+		t.Errorf("user should be initialized with search value and attrs")
+	}
+
+	db.Save(&User{Name: "find or init", Age: 33})
+
+	db.Where(&User{Name: "find or init"}).Attrs(User{Age: 44}).FirstOrInit(&user5)
+	if user5.Name != "find or init" || user5.Id == 0 || user5.Age != 33 {
+		t.Errorf("user should be found and not initialized by Attrs")
+	}
+
+	db.Where(&User{Name: "find or init", Age: 33}).FirstOrInit(&user6)
+	if user6.Name != "find or init" || user6.Id == 0 || user6.Age != 33 {
+		t.Errorf("user should be found with FirstOrInit")
+	}
 }
 
 func TestFindOrCreate(t *testing.T) {
-	var user1, user2, user3 User
+	var user1, user2, user3, user4, user5 User
 	db.Where(&User{Name: "find or create", Age: 33}).FirstOrCreate(&user1)
 	if user1.Name != "find or create" || user1.Id == 0 || user1.Age != 33 {
 		t.Errorf("user should be created with search value")
@@ -839,4 +856,14 @@ func TestFindOrCreate(t *testing.T) {
 	if user3.Name != "find or create 2" || user3.Id == 0 {
 		t.Errorf("user should be created with inline search value")
 	}
+
+	db.Where(&User{Name: "find or create 3"}).Attrs(User{Age: 44}).FirstOrCreate(&user4)
+	if user4.Name != "find or create 3" || user4.Id == 0 || user4.Age != 44 {
+		t.Errorf("user should be created with search value and attrs")
+	}
+
+	db.Where(&User{Name: "find or create"}).Attrs(User{Age: 44}).FirstOrInit(&user5)
+	if user5.Name != "find or create" || user5.Id == 0 || user5.Age != 33 {
+		t.Errorf("user should be found and not initialized by Attrs")
+	}
 }
diff --git a/main.go b/main.go
index bfb3d5c7..27c5f535 100644
--- a/main.go
+++ b/main.go
@@ -30,6 +30,10 @@ func (s *DB) First(out interface{}, where ...interface{}) *Chain {
 	return s.buildChain().First(out, where...)
 }
 
+func (s *DB) Attrs(attrs interface{}) *Chain {
+	return s.buildChain().Attrs(attrs)
+}
+
 func (s *DB) FirstOrInit(out interface{}, where ...interface{}) *Chain {
 	return s.buildChain().FirstOrInit(out, where...)
 }