Support partial indexes

domodwyer · Mar 29, 2017 · 5efe8ec · 5efe8ec
1 parent 0cf7028
commit 5efe8ec
Show file tree

Hide file tree

Showing 2 changed files with 63 additions and 35 deletions.
diff --git a/session.go b/session.go
@@ -1058,30 +1058,32 @@ func (db *Database) RemoveUser(user string) error {
 }
 
 type indexSpec struct {
-	Name, NS         string
-	Key              bson.D
-	Unique           bool    ",omitempty"
-	DropDups         bool    "dropDups,omitempty"
-	Background       bool    ",omitempty"
-	Sparse           bool    ",omitempty"
-	Bits             int     ",omitempty"
-	Min, Max         float64 ",omitempty"
-	BucketSize       float64 "bucketSize,omitempty"
-	ExpireAfter      int     "expireAfterSeconds,omitempty"
-	Weights          bson.D  ",omitempty"
-	DefaultLanguage  string  "default_language,omitempty"
-	LanguageOverride string  "language_override,omitempty"
-	TextIndexVersion int     "textIndexVersion,omitempty"
+	Name, NS                string
+	Key                     bson.D
+	Unique                  bool    ",omitempty"
+	DropDups                bool    "dropDups,omitempty"
+	Background              bool    ",omitempty"
+	Sparse                  bool    ",omitempty"
+	Bits                    int     ",omitempty"
+	Min, Max                float64 ",omitempty"
+	BucketSize              float64 "bucketSize,omitempty"
+	ExpireAfter             int     "expireAfterSeconds,omitempty"
+	Weights                 bson.D  ",omitempty"
+	DefaultLanguage         string  "default_language,omitempty"
+	LanguageOverride        string  "language_override,omitempty"
+	TextIndexVersion        int     "textIndexVersion,omitempty"
+	PartialFilterExpression bson.M  "partialFilterExpression,omitempty"
 
 	Collation *Collation "collation,omitempty"
 }
 
 type Index struct {
-	Key        []string // Index key fields; prefix name with dash (-) for descending order
-	Unique     bool     // Prevent two documents from having the same index key
-	DropDups   bool     // Drop documents with the same index key as a previously indexed one
-	Background bool     // Build index in background and return immediately
-	Sparse     bool     // Only index documents containing the Key fields
+	Key           []string // Index key fields; prefix name with dash (-) for descending order
+	Unique        bool     // Prevent two documents from having the same index key
+	DropDups      bool     // Drop documents with the same index key as a previously indexed one
+	Background    bool     // Build index in background and return immediately
+	Sparse        bool     // Only index documents containing the Key fields
+	PartialFilter bson.M   // Partial index filter expression
 
 	// If ExpireAfter is defined the server will periodically delete
 	// documents with indexed time.Time older than the provided delta.
@@ -1334,6 +1336,10 @@ func (c *Collection) EnsureIndexKey(key ...string) error {
 //     http://www.mongodb.org/display/DOCS/Multikeys
 //
 func (c *Collection) EnsureIndex(index Index) error {
+	if index.Sparse && index.PartialFilter != nil {
+		return errors.New("cannot mix sparse and partial indexes")
+	}
+
 	keyInfo, err := parseIndexKey(index.Key)
 	if err != nil {
 		return err
@@ -1346,22 +1352,23 @@ func (c *Collection) EnsureIndex(index Index) error {
 	}
 
 	spec := indexSpec{
-		Name:             keyInfo.name,
-		NS:               c.FullName,
-		Key:              keyInfo.key,
-		Unique:           index.Unique,
-		DropDups:         index.DropDups,
-		Background:       index.Background,
-		Sparse:           index.Sparse,
-		Bits:             index.Bits,
-		Min:              index.Minf,
-		Max:              index.Maxf,
-		BucketSize:       index.BucketSize,
-		ExpireAfter:      int(index.ExpireAfter / time.Second),
-		Weights:          keyInfo.weights,
-		DefaultLanguage:  index.DefaultLanguage,
-		LanguageOverride: index.LanguageOverride,
-		Collation:        index.Collation,
+		Name:                    keyInfo.name,
+		NS:                      c.FullName,
+		Key:                     keyInfo.key,
+		Unique:                  index.Unique,
+		DropDups:                index.DropDups,
+		Background:              index.Background,
+		Sparse:                  index.Sparse,
+		Bits:                    index.Bits,
+		Min:                     index.Minf,
+		Max:                     index.Maxf,
+		BucketSize:              index.BucketSize,
+		ExpireAfter:             int(index.ExpireAfter / time.Second),
+		Weights:                 keyInfo.weights,
+		DefaultLanguage:         index.DefaultLanguage,
+		LanguageOverride:        index.LanguageOverride,
+		Collation:               index.Collation,
+		PartialFilterExpression: index.PartialFilter,
 	}
 
 	if spec.Min == 0 && spec.Max == 0 {
@@ -1577,6 +1584,7 @@ func indexFromSpec(spec indexSpec) Index {
 		LanguageOverride: spec.LanguageOverride,
 		ExpireAfter:      time.Duration(spec.ExpireAfter) * time.Second,
 		Collation:        spec.Collation,
+		PartialFilter:    spec.PartialFilterExpression,
 	}
 	if float64(int(spec.Min)) == spec.Min && float64(int(spec.Max)) == spec.Max {
 		index.Min = int(spec.Min)

diff --git a/session_test.go b/session_test.go
@@ -3104,6 +3104,21 @@ var indexTests = []struct {
 		"key":  M{"cn": 1},
 		"ns":   "mydb.mycoll",
 	},
+}, {
+	mgo.Index{
+		Key: []string{"partial"},
+		PartialFilter: bson.M{
+			"b": bson.M{"$gt": 42},
+		},
+	},
+	M{
+		"name": "partial_1",
+		"ns":   "mydb.mycoll",
+		"partialFilterExpression": M{
+			"b": M{"$gt": 42},
+		},
+		"key": M{"partial": 1},
+	},
 }}
 
 func (s *S) TestEnsureIndex(c *C) {
@@ -3119,6 +3134,11 @@ func (s *S) TestEnsureIndex(c *C) {
 			continue
 		}
 
+		// Only test partial indexes on
+		if !s.versionAtLeast(3, 2) && test.expected["name"] == "partial_1" {
+			continue
+		}
+
 		err = coll.EnsureIndex(test.index)
 		msg := "text search not enabled"
 		if err != nil && strings.Contains(err.Error(), msg) {