Add mongo/collectionopt

GODRIVER-272

Change-Id: I801d503bdb6452de4b4537cbca32b7fb79f88a80

Author: rfblue2

mongo/collection.go

@@ -20,6 +20,7 @@ import (
 	"github.com/mongodb/mongo-go-driver/core/readpref"
 	"github.com/mongodb/mongo-go-driver/core/writeconcern"
 	"github.com/mongodb/mongo-go-driver/mongo/aggregateopt"
+	"github.com/mongodb/mongo-go-driver/mongo/collectionopt"
 	"github.com/mongodb/mongo-go-driver/mongo/countopt"
 	"github.com/mongodb/mongo-go-driver/mongo/deleteopt"
 	"github.com/mongodb/mongo-go-driver/mongo/distinctopt"
@@ -41,14 +42,34 @@ type Collection struct {
 	writeSelector  description.ServerSelector
 }
 
-func newCollection(db *Database, name string) *Collection {
+func newCollection(db *Database, name string, opts ...collectionopt.Option) *Collection {
+	collOpt, err := collectionopt.BundleCollection(opts...).Unbundle()
+	if err != nil {
+		return nil
+	}
+
+	rc := db.readConcern
+	if collOpt.ReadConcern != nil {
+		rc = collOpt.ReadConcern
+	}
+
+	wc := db.writeConcern
+	if collOpt.WriteConcern != nil {
+		wc = collOpt.WriteConcern
+	}
+
+	rp := db.readPreference
+	if collOpt.ReadPreference != nil {
+		rp = collOpt.ReadPreference
+	}
+
 	coll := &Collection{
 		client:         db.client,
 		db:             db,
 		name:           name,
-		readPreference: db.readPreference,
-		readConcern:    db.readConcern,
-		writeConcern:   db.writeConcern,
+		readPreference: rp,
+		readConcern:    rc,
+		writeConcern:   wc,
 		readSelector:   db.readSelector,
 		writeSelector:  db.writeSelector,
 	}

mongo/collection_internal_test.go

@@ -15,9 +15,12 @@ import (
 	"github.com/mongodb/mongo-go-driver/bson"
 	"github.com/mongodb/mongo-go-driver/bson/objectid"
 	"github.com/mongodb/mongo-go-driver/core/option"
+	"github.com/mongodb/mongo-go-driver/core/readconcern"
+	"github.com/mongodb/mongo-go-driver/core/readpref"
 	"github.com/mongodb/mongo-go-driver/core/writeconcern"
 	"github.com/mongodb/mongo-go-driver/internal/testutil"
 	"github.com/mongodb/mongo-go-driver/mongo/aggregateopt"
+	"github.com/mongodb/mongo-go-driver/mongo/collectionopt"
 	"github.com/mongodb/mongo-go-driver/mongo/countopt"
 	"github.com/mongodb/mongo-go-driver/mongo/deleteopt"
 	"github.com/mongodb/mongo-go-driver/mongo/distinctopt"
@@ -29,15 +32,15 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-func createTestCollection(t *testing.T, dbName *string, collName *string) *Collection {
+func createTestCollection(t *testing.T, dbName *string, collName *string, opts ...collectionopt.Option) *Collection {
 	if collName == nil {
 		coll := testutil.ColName(t)
 		collName = &coll
 	}
 
 	db := createTestDatabase(t, dbName)
 
-	return db.Collection(*collName)
+	return db.Collection(*collName, opts...)
 }
 
 func initCollection(t *testing.T, coll *Collection) {
@@ -76,6 +79,73 @@ func TestCollection_initialize(t *testing.T) {
 	require.NotNil(t, coll.db)
 }
 
+func compareColls(t *testing.T, expected *Collection, got *Collection) {
+	switch {
+	case expected.readPreference != got.readPreference:
+		t.Errorf("expected read preference %#v. got %#v", expected.readPreference, got.readPreference)
+	case expected.readConcern != got.readConcern:
+		t.Errorf("expected read concern %#v. got %#v", expected.readConcern, got.readConcern)
+	case expected.writeConcern != got.writeConcern:
+		t.Errorf("expected write concern %#v. got %#v", expected.writeConcern, got.writeConcern)
+	}
+}
+
+func TestCollection_Options(t *testing.T) {
+	name := "testDb_options"
+	rpPrimary := readpref.Primary()
+	rpSecondary := readpref.Secondary()
+	wc1 := writeconcern.New(writeconcern.W(5))
+	wc2 := writeconcern.New(writeconcern.W(10))
+	rcLocal := readconcern.Local()
+	rcMajority := readconcern.Majority()
+
+	opts := []collectionopt.Option{collectionopt.ReadPreference(rpPrimary), collectionopt.ReadConcern(rcLocal), collectionopt.WriteConcern(wc1),
+		collectionopt.ReadPreference(rpSecondary), collectionopt.ReadConcern(rcMajority), collectionopt.WriteConcern(wc2)}
+
+	dbName := "collection_internal_test_db1"
+
+	expectedColl := &Collection{
+		readConcern:    rcMajority,
+		readPreference: rpSecondary,
+		writeConcern:   wc2,
+	}
+
+	t.Run("IndividualOptions", func(t *testing.T) {
+		// if options specified multiple times, last instance should take precedence
+		coll := createTestCollection(t, &dbName, &name, opts...)
+		compareColls(t, expectedColl, coll)
+	})
+
+	t.Run("Bundle", func(t *testing.T) {
+		coll := createTestCollection(t, &dbName, &name, collectionopt.BundleCollection(opts...))
+		compareColls(t, expectedColl, coll)
+	})
+}
+
+func TestCollection_InheritOptions(t *testing.T) {
+	name := "testDb_options_inherit"
+	client := createTestClient(t)
+
+	rpPrimary := readpref.Primary()
+	rcLocal := readconcern.Local()
+	wc1 := writeconcern.New(writeconcern.W(10))
+
+	db := client.Database("collection_internal_test_db2")
+	db.readPreference = rpPrimary
+	db.readConcern = rcLocal
+	coll := db.Collection(name, collectionopt.WriteConcern(wc1))
+
+	// coll should inherit read preference and read concern from client
+	switch {
+	case coll.readPreference != rpPrimary:
+		t.Errorf("expected read preference primary. got %#v", coll.readPreference)
+	case coll.readConcern != rcLocal:
+		t.Errorf("expected read concern local. got %#v", coll.readConcern)
+	case coll.writeConcern != wc1:
+		t.Errorf("expected write concern %#v. got %#v", wc1, coll.writeConcern)
+	}
+}
+
 func TestCollection_namespace(t *testing.T) {
 	t.Parallel()
 

mongo/collectionopt/collectionopt.go

@@ -0,0 +1,168 @@
+package collectionopt
+
+import (
+	"reflect"
+
+	"github.com/mongodb/mongo-go-driver/core/readconcern"
+	"github.com/mongodb/mongo-go-driver/core/readpref"
+	"github.com/mongodb/mongo-go-driver/core/writeconcern"
+)
+
+var collectionBundle = new(CollectionBundle)
+
+// Option represents a collection option.
+type Option interface {
+	collectionOption()
+}
+
+// optionFunc adds the option to the client.
+type optionFunc func(*Collection) error
+
+// Collection represents a collection.
+type Collection struct {
+	ReadConcern    *readconcern.ReadConcern
+	WriteConcern   *writeconcern.WriteConcern
+	ReadPreference *readpref.ReadPref
+}
+
+// CollectionBundle is a bundle of collection options.
+type CollectionBundle struct {
+	option Option
+	next   *CollectionBundle
+}
+
+// CollectionBundle implements Option.
+func (*CollectionBundle) collectionOption() {}
+
+// OptionFunc implements Option.
+func (optionFunc) collectionOption() {}
+
+// BundleCollection bundles collection options.
+func BundleCollection(opts ...Option) *CollectionBundle {
+	head := collectionBundle
+
+	for _, opt := range opts {
+		newBundle := CollectionBundle{
+			option: opt,
+			next:   head,
+		}
+
+		head = &newBundle
+	}
+
+	return head
+}
+
+// ReadConcern sets the read concern.
+func (cb *CollectionBundle) ReadConcern(rc *readconcern.ReadConcern) *CollectionBundle {
+	return &CollectionBundle{
+		option: ReadConcern(rc),
+		next:   cb,
+	}
+}
+
+// WriteConcern sets the write concern.
+func (cb *CollectionBundle) WriteConcern(wc *writeconcern.WriteConcern) *CollectionBundle {
+	return &CollectionBundle{
+		option: WriteConcern(wc),
+		next:   cb,
+	}
+}
+
+// ReadPreference sets the read preference.
+func (cb *CollectionBundle) ReadPreference(rp *readpref.ReadPref) *CollectionBundle {
+	return &CollectionBundle{
+		option: ReadPreference(rp),
+		next:   cb,
+	}
+}
+
+// String prints a string representation of the bundle for debug purposes
+func (cb *CollectionBundle) String() string {
+	if cb == nil {
+		return ""
+	}
+
+	debugStr := ""
+	for head := cb; head != nil && head.option != nil; head = head.next {
+		switch opt := head.option.(type) {
+		case *CollectionBundle:
+			debugStr += opt.String()
+		case optionFunc:
+			debugStr += reflect.TypeOf(opt).String() + "\n"
+		default:
+			return debugStr + "(error: CollectionOption can only be *CollectionBundle or optionFunc)"
+		}
+	}
+
+	return debugStr
+}
+
+// Unbundle unbundles the options, returning a collection.
+func (cb *CollectionBundle) Unbundle() (*Collection, error) {
+	client := &Collection{}
+	err := cb.unbundle(client)
+	if err != nil {
+		return nil, err
+	}
+
+	return client, nil
+}
+
+// Helper that recursively unwraps the bundle.
+func (cb *CollectionBundle) unbundle(client *Collection) error {
+	if cb == nil {
+		return nil
+	}
+
+	for head := cb; head != nil && head.option != nil; head = head.next {
+		var err error
+		switch opt := head.option.(type) {
+		case *CollectionBundle:
+			err = opt.unbundle(client) // add all bundle's options to client
+		case optionFunc:
+			err = opt(client) // add option to client
+		default:
+			return nil
+		}
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+
+}
+
+// ReadConcern sets the read concern.
+func ReadConcern(rc *readconcern.ReadConcern) Option {
+	return optionFunc(
+		func(c *Collection) error {
+			if c.ReadConcern == nil {
+				c.ReadConcern = rc
+			}
+			return nil
+		})
+}
+
+// WriteConcern sets the write concern.
+func WriteConcern(wc *writeconcern.WriteConcern) Option {
+	return optionFunc(
+		func(c *Collection) error {
+			if c.WriteConcern == nil {
+				c.WriteConcern = wc
+			}
+			return nil
+		})
+}
+
+// ReadPreference sets the read preference.
+func ReadPreference(rp *readpref.ReadPref) Option {
+	return optionFunc(
+		func(c *Collection) error {
+			if c.ReadPreference == nil {
+				c.ReadPreference = rp
+			}
+			return nil
+		})
+}