Merge pull request #507 from Scalingo/feat/506/count-methods

feat(count): Add count method to document package
Scalingo · Feb 3, 2023 · 3687e34 · 3687e34
2 parents 74d24eb + 8e0df0b
commit 3687e34
Show file tree

Hide file tree

Showing 3 changed files with 93 additions and 0 deletions.
diff --git a/mongo/CHANGELOG.md b/mongo/CHANGELOG.md
@@ -2,6 +2,8 @@
 
 ## To be Released
 
+* feat(count): Add count method to document package
+
 ## v1.2.2
 
 * docs: add doc for Session function [#397](https://github.com/Scalingo/go-utils/pull/397)

diff --git a/mongo/document/document.go b/mongo/document/document.go
@@ -189,6 +189,29 @@ func WhereUnscoped(ctx context.Context, collectionName string, query bson.M, dat
 	return nil
 }
 
+func Count(ctx context.Context, collectionName string, query bson.M) (int, error) {
+	if query == nil {
+		query = bson.M{}
+	}
+	if _, ok := query["deleted_at"]; !ok {
+		query["deleted_at"] = nil
+	}
+
+	return CountUnscoped(ctx, collectionName, query)
+}
+
+func CountUnscoped(ctx context.Context, collectionName string, query bson.M) (int, error) {
+	log := logger.Get(ctx)
+	// nolint: contextcheck
+	c := mongo.Session(log).Clone().DB("").C(collectionName)
+
+	if query == nil {
+		query = bson.M{}
+	}
+
+	return c.Find(query).Count()
+}
+
 func WhereIter(ctx context.Context, collectionName string, query bson.M, fun func(*mgo.Iter) error, sortFields ...SortField) error {
 	if query == nil {
 		query = bson.M{}

diff --git a/mongo/document/paranoid_test.go b/mongo/document/paranoid_test.go
@@ -25,6 +25,74 @@ func NewTestParanoidDoc(t *testing.T) (*ParanoidDoc, func()) {
 	}
 }
 
+func TestParanoid_Count(t *testing.T) {
+	examples := []struct {
+		Name         string
+		Query        bson.M
+		ParanoidDocs func(t *testing.T) ([]*ParanoidDoc, func())
+		Count        int
+	}{
+		{
+			Name: "it should count existing documents",
+			ParanoidDocs: func(t *testing.T) ([]*ParanoidDoc, func()) {
+				d1, clean1 := NewTestParanoidDoc(t)
+				d2, clean2 := NewTestParanoidDoc(t)
+				return []*ParanoidDoc{d1, d2}, func() {
+					clean1()
+					clean2()
+				}
+			},
+			Count: 2,
+		}, {
+			Name: "it should not count paranoia-deleted documents",
+			ParanoidDocs: func(t *testing.T) ([]*ParanoidDoc, func()) {
+				d1, clean1 := NewTestParanoidDoc(t)
+				err := Destroy(context.Background(), ParanoidDocsCollection, d1)
+				require.NoError(t, err)
+				d2, clean2 := NewTestParanoidDoc(t)
+				err = Destroy(context.Background(), ParanoidDocsCollection, d2)
+				require.NoError(t, err)
+				return []*ParanoidDoc{d1, d2}, func() {
+					clean1()
+					clean2()
+				}
+			},
+			Count: 0,
+		}, {
+			Name:  "it should find count document, if queried specifically",
+			Query: bson.M{"deleted_at": bson.M{"$exists": true}},
+			ParanoidDocs: func(t *testing.T) ([]*ParanoidDoc, func()) {
+				d1, clean1 := NewTestParanoidDoc(t)
+				err := Destroy(context.Background(), ParanoidDocsCollection, d1)
+				require.NoError(t, err)
+				d2, clean2 := NewTestParanoidDoc(t)
+				err = Destroy(context.Background(), ParanoidDocsCollection, d2)
+				require.NoError(t, err)
+				return []*ParanoidDoc{d1, d2}, func() {
+					clean1()
+					clean2()
+				}
+			},
+			Count: 2,
+		},
+	}
+
+	for _, example := range examples {
+		t.Run(example.Name, func(t *testing.T) {
+			_, clean := example.ParanoidDocs(t)
+			defer clean()
+
+			query := bson.M{}
+			if example.Query != nil {
+				query = example.Query
+			}
+			count, err := Count(context.Background(), ParanoidDocsCollection, query)
+			require.NoError(t, err)
+			require.Equal(t, count, example.Count)
+		})
+	}
+}
+
 func TestParanoid_Find(t *testing.T) {
 	examples := []struct {
 		Name        string