limit # collections to 2022 when error marshaling

keyhole.go

@@ -224,7 +224,7 @@ func Run(fullVersion string) {
 		}
 		stats.Print()
 		if ofile, data, err = stats.OutputBSON(); err != nil {
-			return
+			log.Fatalf("failed to output bson file\n%v", err)
 		}
 		GenerateMaobiReport(*maobiURL, data, ofile)
 		return

mdb/cluster_stats.go

@@ -265,7 +265,21 @@ func (p *ClusterStats) OutputBSON() (string, []byte, error) {
 		return ofile, data, errors.New(result)
 	}
 	if data, err = bson.Marshal(p); err != nil {
-		return ofile, data, err
+		maxCollections := 2022
+		left := maxCollections
+		gox.GetLogger("").Errorf("error marshaling bson: %v, retry and limit # of collections to %v", err, maxCollections)
+		for i, db := range p.Databases {
+			if len(db.Collections) < left {
+				left -= len(db.Collections)
+			} else {
+				p.Databases[i].Collections = db.Collections[:left]
+				p.Databases = p.Databases[:i+1]
+				break
+			}
+		}
+		if data, err = bson.Marshal(p); err != nil {
+			return ofile, data, err
+		}
 	}
 
 	os.Mkdir(outdir, 0755)

mdb/cluster_stats_test.go

@@ -0,0 +1,43 @@
+// Copyright 2019 Kuei-chun Chen. All rights reserved.
+
+package mdb
+
+import (
+	"log"
+	"os"
+	"testing"
+
+	"go.mongodb.org/mongo-driver/x/mongo/driver/connstring"
+)
+
+func TestOutputBSON(t *testing.T) {
+	uri := "mongodb://admin:secret@localhost:30309/"
+	if os.Getenv("DATABASE_URL") != "" {
+		uri = os.Getenv("DATABASE_URL")
+	}
+	var connString connstring.ConnString
+	connString, err := ParseURI(uri)
+	if err != nil {
+		t.Fatal(err)
+	}
+	client, err := NewMongoClient(uri)
+	if err != nil {
+		t.Fatal(err)
+	}
+	stats := NewClusterStats("simagix/keyhole test")
+	if err = stats.GetClusterStats(client, connString); err != nil {
+		log.Fatalf("error GetClusterStats: %v", err)
+	}
+
+	// test bson.Marshal() overflow
+	// for index := range stats.Databases {
+	// 	for i := 0; i < 12; i++ {
+	// 		colls := stats.Databases[index].Collections
+	// 		stats.Databases[index].Collections = append(stats.Databases[index].Collections, colls...)
+	// 	}
+	// }
+
+	if _, _, err = stats.OutputBSON(); err != nil {
+		log.Fatalf("failed to output bson file\n%v", err)
+	}
+}

version

@@ -1 +1 @@
-1.2.6
+1.2.7