Merge pull request #6 from pupizoid/3-include-desired-fields-to-query

implemented selected columns

Author: pupizoid

ormlite.go

@@ -92,7 +92,10 @@ type Options struct {
 	OrderBy       *OrderBy `json:"order_by"`
 	RelationDepth int      `json:"relation_depth"`
 	RelatedTo     []IModel `json:"related"`
-	joins         []string
+	// Columns contains map with string keys of columns to include to the query
+	// instead of querying all model fields
+	Columns map[string]struct{} `json:"columns"`
+	joins   []string
 }
 
 // DefaultOptions returns default options for query
@@ -150,7 +153,7 @@ type columnInfo struct {
 }
 
 func isExportedField(f reflect.StructField) bool {
-	return strings.ToLower(string([]rune(f.Name)[0])) != string([]rune(f.Name)[0])
+	return f.IsExported()
 }
 
 func lookForSettingWithSep(s, setting, sep string) string {
@@ -637,6 +640,18 @@ func QueryStructContext(ctx context.Context, db *sql.DB, opts *Options, out Mode
 			continue
 		}
 
+		if opts != nil && opts.Columns != nil {
+			var colName string
+			if exp, ok := model.Field(i).Interface().(Expression); ok {
+				colName = exp.Column()
+			} else {
+				colName = getFieldColumnName(model.Type().Field(i))
+			}
+			if _, ok := opts.Columns[colName]; !ok && !strings.Contains(tag, "primary") {
+				continue
+			}
+		}
+
 		if ri := extractRelationInfo(model.Type().Field(i)); ri != nil {
 			if ri.Type == hasOne {
 				columns = append(columns, getFieldColumnName(model.Type().Field(i)))
@@ -727,6 +742,16 @@ func QuerySliceCountContext(ctx context.Context, db *sql.DB, opts *Options, out
 		return fmt.Errorf("failed to get column info for type: %v", modelType)
 	}
 
+	if opts != nil && opts.Columns != nil {
+		var selected []columnInfo
+		for _, ci := range colInfo {
+			if _, ok := opts.Columns[ci.Name]; ok || ci.Primary {
+				selected = append(selected, ci)
+			}
+		}
+		colInfo = selected
+	}
+
 	for _, ci := range colInfo {
 		if ci.RelationInfo.Type == noRelation || ci.RelationInfo.Type == hasOne {
 			if ci.Primary {
@@ -899,7 +924,7 @@ func addWhereClause(options *Options, s string, value reflect.Value) {
 	}
 }
 
-// Delete removes model object from database by it's primary key
+// Delete removes model object from database by its primary key
 func Delete(db *sql.DB, m Model) (sql.Result, error) {
 	modelValue := reflect.ValueOf(m).Elem()
 
@@ -950,6 +975,7 @@ type pkFieldInfo struct {
 	field        reflect.Value
 }
 
+// Count models in database with search options
 func Count(db *sql.DB, m Model, opts *Options) (count int64, err error) {
 	mInfo, err := getModelInfo(m)
 	if err != nil {

ormlite_test.go

@@ -1237,5 +1237,85 @@ func TestQuerySliceCount(t *testing.T) {
 	if assert.NoError(t, QuerySliceCount(db, &Options{Where: Where{"attr": 1}}, &m, &count)) {
 		assert.Len(t, m, 4)
 		assert.EqualValues(t, 4, count)
+		assert.EqualValues(t, 4, m[3].ID)
 	}
 }
+
+type SelectedColumnsSuite struct {
+	suite.Suite
+	db *sql.DB
+}
+
+type BigModel struct {
+	ID      int           `ormlite:"primary"`
+	Attr1   int           `ormlite:"col=attr1"`
+	Attr2   int           `ormlite:"col=attr2"`
+	Attr3   string        `ormlite:"col=attr3"`
+	Attr4   float64       `ormlite:"col=attr4"`
+	Related *relatedModel `ormlite:"has_one,col=rel_id"`
+}
+
+func (b BigModel) Table() string { return "big_model" }
+
+func (s *SelectedColumnsSuite) SetupSuite() {
+	db, err := sql.Open("sqlite3", ":memory:")
+	require.NoError(s.T(), err)
+
+	_, err = db.Exec("create table big_model(id integer primary key, attr1 int, attr2 int, attr3 string, attr4 float, rel_id int);" +
+		"create table related_model(id integer primary key, field text);")
+	require.NoError(s.T(), err)
+	s.db = db
+
+	require.NoError(s.T(), Insert(db, &BigModel{Attr1: 1, Attr2: 2, Attr3: "first", Attr4: 1.0}))
+	require.NoError(s.T(), Insert(db, &BigModel{Attr1: 3, Attr2: 4, Attr3: "second", Attr4: 2.0}))
+	require.NoError(s.T(), Insert(db, &BigModel{Attr1: 5, Attr2: 6, Attr3: "third", Attr4: 3.0}))
+	require.NoError(s.T(), Insert(db, &BigModel{Attr1: 7, Attr2: 8, Attr3: "forth", Attr4: 4.0}))
+	require.NoError(s.T(), Insert(db, &BigModel{Attr1: 9, Attr2: 10, Attr3: "fifth", Attr4: 5.0}))
+
+	require.NoError(s.T(), Upsert(db, &BigModel{
+		Attr1:   11,
+		Attr2:   11,
+		Attr3:   "11",
+		Attr4:   11,
+		Related: &relatedModel{Field: "Hello"},
+	}))
+}
+
+func (s *SelectedColumnsSuite) TearDownSuite() {
+	require.NoError(s.T(), s.db.Close())
+}
+
+func (s *SelectedColumnsSuite) TestQueryStruct() {
+	var m BigModel
+	require.NoError(s.T(), QueryStruct(s.db, &Options{Columns: map[string]struct{}{
+		"attr1": {},
+		"attr3": {},
+	}, Where: Where{"id": 1}}, &m))
+
+	assert.EqualValues(s.T(), 1, m.ID)
+	assert.EqualValues(s.T(), 1, m.Attr1)
+	assert.EqualValues(s.T(), 0, m.Attr2)
+	assert.EqualValues(s.T(), "first", m.Attr3)
+	assert.EqualValues(s.T(), 0.0, m.Attr4)
+}
+
+func (s *SelectedColumnsSuite) TestQuerySlice() {
+	var mm []*BigModel
+	require.NoError(s.T(), QuerySlice(s.db, &Options{Columns: map[string]struct{}{
+		"attr2": {},
+		"attr4": {},
+	}}, &mm))
+
+	if assert.NotNil(s.T(), mm) {
+		assert.Len(s.T(), mm, 6)
+		assert.EqualValues(s.T(), 6, mm[2].Attr2)
+		assert.EqualValues(s.T(), "", mm[2].Attr3)
+		assert.EqualValues(s.T(), 3.0, mm[2].Attr4)
+		assert.EqualValues(s.T(), (*relatedModel)(nil), mm[5].Related)
+	}
+
+}
+
+func TestSelectedColumns(t *testing.T) {
+	suite.Run(t, new(SelectedColumnsSuite))
+}