From a0477f94dd97ef33a442aadf7c710ac03d4a0590 Mon Sep 17 00:00:00 2001
From: Jinzhu <wosmvp@gmail.com>
Date: Sun, 19 Jul 2020 21:48:58 +0800
Subject: [PATCH] Allow Omit with Query, close #3165

---
 callbacks/query.go  |  8 ++++++++
 tests/query_test.go | 15 +++++++++++++++
 2 files changed, 23 insertions(+)

diff --git a/callbacks/query.go b/callbacks/query.go
index 9601f9bd..5c322a05 100644
--- a/callbacks/query.go
+++ b/callbacks/query.go
@@ -64,6 +64,14 @@ func BuildQuerySQL(db *gorm.DB) {
 				clauseSelect.Columns[idx] = clause.Column{Name: name, Raw: true}
 			}
 		}
+	} else if db.Statement.Schema != nil && len(db.Statement.Omits) > 0 {
+		selectColumns, _ := db.Statement.SelectAndOmitColumns(false, false)
+		clauseSelect.Columns = make([]clause.Column, 0, len(db.Statement.Schema.DBNames))
+		for _, dbName := range db.Statement.Schema.DBNames {
+			if v, ok := selectColumns[dbName]; (ok && v) || !ok {
+				clauseSelect.Columns = append(clauseSelect.Columns, clause.Column{Name: dbName})
+			}
+		}
 	} else if db.Statement.Schema != nil && db.Statement.ReflectValue.IsValid() {
 		smallerStruct := false
 		switch db.Statement.ReflectValue.Kind() {
diff --git a/tests/query_test.go b/tests/query_test.go
index 22807377..59f1130b 100644
--- a/tests/query_test.go
+++ b/tests/query_test.go
@@ -330,6 +330,21 @@ func TestSelect(t *testing.T) {
 	}
 }
 
+func TestOmit(t *testing.T) {
+	user := User{Name: "OmitUser1", Age: 20}
+	DB.Save(&user)
+
+	var result User
+	DB.Where("name = ?", user.Name).Omit("name").Find(&result)
+	if result.ID == 0 {
+		t.Errorf("Should not have ID because only selected name, %+v", result.ID)
+	}
+
+	if result.Name != "" || result.Age != 20 {
+		t.Errorf("User Name should be omitted, got %v, Age should be ok, got %v", result.Name, result.Age)
+	}
+}
+
 func TestPluckWithSelect(t *testing.T) {
 	users := []User{
 		{Name: "pluck_with_select_1", Age: 25},