From a1929a130edb3142398df7300bef047a53dafc6d Mon Sep 17 00:00:00 2001
From: Alda Liu <aliu@venacorp.com>
Date: Wed, 19 Jul 2023 09:55:13 -0400
Subject: [PATCH] Add support for entire database namespaces

This commit changes namespace parsing to support databases.
First we check if any namespaces contain only the database name,
then we fetch all collections from that database and append it to
the namespace slice.
---
 internal/verifier/check.go           | 13 +++++++++++++
 internal/verifier/list_namespaces.go |  8 +++++++-
 internal/verifier/util.go            | 24 ++++++++++++++++++++++++
 3 files changed, 44 insertions(+), 1 deletion(-)

diff --git a/internal/verifier/check.go b/internal/verifier/check.go
index 3a4bcc4a..ef50bb03 100644
--- a/internal/verifier/check.go
+++ b/internal/verifier/check.go
@@ -285,6 +285,19 @@ func (verifier *Verifier) CreateInitialTasks() error {
 			verifier.logger.Error().Msgf("%s", err)
 			return err
 		}
+		// Check for database names and append their collections
+		if NamespacesContainsDBName(verifier.srcNamespaces) {
+			namespaces, dbs := ParseDBNamespaces(verifier.srcNamespaces)
+
+			dbNamespaces, err := ListUserCollectionsForDBs(context.Background(), verifier.logger, verifier.srcClient, true /* include views */, dbs)
+			if err != nil {
+				verifier.logger.Error().Msgf("Failed to parse database namespaces: %s", err)
+				return err
+			}
+
+			verifier.srcNamespaces = append(namespaces, dbNamespaces...)
+			verifier.SetNamespaceMap()
+		}
 	}
 	isPrimary, err := verifier.CheckIsPrimary()
 	if err != nil {
diff --git a/internal/verifier/list_namespaces.go b/internal/verifier/list_namespaces.go
index 2ab0603a..89d92c97 100644
--- a/internal/verifier/list_namespaces.go
+++ b/internal/verifier/list_namespaces.go
@@ -37,8 +37,14 @@ func ListAllUserCollections(ctx context.Context, logger *logger.Logger, client *
 	}
 	logger.Debug().Msgf("All user databases: %+v", dbNames)
 
+	return ListUserCollectionsForDBs(ctx, logger, client, includeViews, dbNames)
+}
+
+func ListUserCollectionsForDBs(ctx context.Context, logger *logger.Logger, client *mongo.Client, includeViews bool,
+	databases []string) ([]string, error) {
+
 	collectionNamespaces := []string{}
-	for _, dbName := range dbNames {
+	for _, dbName := range databases {
 		db := client.Database(dbName)
 		filter := bson.D{{"name", bson.D{{"$nin", bson.A{ExcludedSystemCollRegex}}}}}
 		if !includeViews {
diff --git a/internal/verifier/util.go b/internal/verifier/util.go
index 6d200f4e..bf4ff53c 100644
--- a/internal/verifier/util.go
+++ b/internal/verifier/util.go
@@ -154,3 +154,27 @@ func GetLastOpTimeAndSyncShardClusterTime(
 	t, i := rawOperationTime.Timestamp()
 	return &primitive.Timestamp{T: t, I: i}, nil
 }
+
+func NamespacesContainsDBName(namespaces []string) bool {
+	for _, namespace := range namespaces {
+		if strings.Index(namespace, ".") < 0 {
+			return true
+		}
+	}
+	return false
+}
+
+func ParseDBNamespaces(namespaces []string) ([]string, []string) {
+	databases := []string{}
+	parsedNamespaces := []string{}
+	// strip database names from namespaces - db collections will be appended to the namespaces
+	for _, name := range namespaces {
+		db, coll := SplitNamespace(name)
+		if coll == "" {
+			databases = append(databases, db)
+		} else {
+			parsedNamespaces = append(parsedNamespaces, name)
+		}
+	}
+	return parsedNamespaces, databases
+}