SERVER-55398 Dropping a database should also drop all its associated tags.

5 files changed, 79 insertions, 0 deletions

diff --git a/src/mongo/s/catalog/sharding_catalog_client.h b/src/mongo/s/catalog/sharding_catalog_client.h
index 1fafa130380..cc03a2baa28 100644
--- a/src/mongo/s/catalog/sharding_catalog_client.h
+++ b/src/mongo/s/catalog/sharding_catalog_client.h
@@ -213,6 +213,12 @@ public:
                                                                    const NamespaceString& nss) = 0;
 
     /**
+     * Retrieves all namespaces that have zones associated with a database.
+     */
+    virtual std::vector<NamespaceString> getAllNssThatHaveZonesForDatabase(
+        OperationContext* opCtx, const StringData& dbName) = 0;
+
+    /**
      * Retrieves all shards in this sharded cluster.
      * Returns a !OK status if an error occurs.
      */
diff --git a/src/mongo/s/catalog/sharding_catalog_client_impl.cpp b/src/mongo/s/catalog/sharding_catalog_client_impl.cpp
index d82ec1e6f40..2253e2ab55e 100644
--- a/src/mongo/s/catalog/sharding_catalog_client_impl.cpp
+++ b/src/mongo/s/catalog/sharding_catalog_client_impl.cpp
@@ -1021,6 +1021,68 @@ StatusWith<std::vector<TagsType>> ShardingCatalogClientImpl::getTagsForCollectio
     return tags;
 }
 
+std::vector<NamespaceString> ShardingCatalogClientImpl::getAllNssThatHaveZonesForDatabase(
+    OperationContext* opCtx, const StringData& dbName) {
+    auto expCtx =
+        make_intrusive<ExpressionContext>(opCtx, nullptr /*collator*/, TagsType::ConfigNS);
+    StringMap<ExpressionContext::ResolvedNamespace> resolvedNamespaces;
+    resolvedNamespaces[TagsType::ConfigNS.coll()] = {TagsType::ConfigNS,
+                                                     std::vector<BSONObj>() /* pipeline */};
+    expCtx->setResolvedNamespaces(resolvedNamespaces);
+
+    // Parse pipeline:
+    //      - First stage will find all that namespaces on 'config.tags' that are part of the
+    //      given database.
+    //      - Second stage will group namespaces to not have repetitions.
+    //
+    //      db.tags.aggregate([
+    //          {$match: {ns: {$regex : "^dbName\\..*", $options: "i"}}},
+    //          {$group: {_id : "$ns"}}
+    //      ])
+    //
+    const std::string regex = "^" + dbName + "\\..*";
+    auto matchStageBson = BSON("ns" << BSON("$regex" << regex << "$options"
+                                                     << "i"));
+    auto matchStage = DocumentSourceMatch::createFromBson(
+        Document{{"$match", std::move(matchStageBson)}}.toBson().firstElement(), expCtx);
+
+    auto groupStageBson = BSON("_id"
+                               << "$ns");
+    auto groupStage = DocumentSourceGroup::createFromBson(
+        Document{{"$group", std::move(groupStageBson)}}.toBson().firstElement(), expCtx);
+
+    // Create pipeline
+    Pipeline::SourceContainer stages;
+    stages.emplace_back(std::move(matchStage));
+    stages.emplace_back(std::move(groupStage));
+
+    const auto pipeline = Pipeline::create(stages, expCtx);
+    auto aggRequest = AggregateCommandRequest(TagsType::ConfigNS, pipeline->serializeToBson());
+
+    // Run the aggregation
+    const auto readConcern = [&]() -> repl::ReadConcernArgs {
+        if (serverGlobalParams.clusterRole == ClusterRole::ConfigServer) {
+            return {repl::ReadConcernLevel::kMajorityReadConcern};
+        } else {
+            const auto time = VectorClock::get(opCtx)->getTime();
+            return {time.configTime(), repl::ReadConcernLevel::kMajorityReadConcern};
+        }
+    }();
+
+    auto aggResult = runCatalogAggregation(opCtx,
+                                           Grid::get(opCtx)->shardRegistry()->getConfigShard(),
+                                           aggRequest,
+                                           readConcern,
+                                           Shard::kDefaultConfigCommandTimeout);
+
+    // Parse the result
+    std::vector<NamespaceString> nssList;
+    for (const auto& doc : aggResult) {
+        nssList.push_back(NamespaceString(doc.getField("_id").String()));
+    }
+    return nssList;
+}
+
 StatusWith<repl::OpTimeWith<std::vector<ShardType>>> ShardingCatalogClientImpl::getAllShards(
     OperationContext* opCtx, repl::ReadConcernLevel readConcern) {
     auto findStatus = _exhaustiveFindOnConfig(opCtx,
diff --git a/src/mongo/s/catalog/sharding_catalog_client_impl.h b/src/mongo/s/catalog/sharding_catalog_client_impl.h
index c4ab7f259f0..0a58457412b 100644
--- a/src/mongo/s/catalog/sharding_catalog_client_impl.h
+++ b/src/mongo/s/catalog/sharding_catalog_client_impl.h
@@ -115,6 +115,9 @@ public:
     StatusWith<std::vector<TagsType>> getTagsForCollection(OperationContext* opCtx,
                                                            const NamespaceString& nss) override;
 
+    std::vector<NamespaceString> getAllNssThatHaveZonesForDatabase(
+        OperationContext* opCtx, const StringData& dbName) override;
+
     StatusWith<repl::OpTimeWith<std::vector<ShardType>>> getAllShards(
         OperationContext* opCtx, repl::ReadConcernLevel readConcern) override;
 
diff --git a/src/mongo/s/catalog/sharding_catalog_client_mock.cpp b/src/mongo/s/catalog/sharding_catalog_client_mock.cpp
index a0a1b38331d..162d445aa31 100644
--- a/src/mongo/s/catalog/sharding_catalog_client_mock.cpp
+++ b/src/mongo/s/catalog/sharding_catalog_client_mock.cpp
@@ -116,6 +116,11 @@ StatusWith<std::vector<TagsType>> ShardingCatalogClientMock::getTagsForCollectio
     return {ErrorCodes::InternalError, "Method not implemented"};
 }
 
+std::vector<NamespaceString> ShardingCatalogClientMock::getAllNssThatHaveZonesForDatabase(
+    OperationContext* opCtx, const StringData& dbName) {
+    uasserted(ErrorCodes::InternalError, "Method not implemented");
+}
+
 StatusWith<repl::OpTimeWith<std::vector<ShardType>>> ShardingCatalogClientMock::getAllShards(
     OperationContext* opCtx, repl::ReadConcernLevel readConcern) {
     return {ErrorCodes::InternalError, "Method not implemented"};
diff --git a/src/mongo/s/catalog/sharding_catalog_client_mock.h b/src/mongo/s/catalog/sharding_catalog_client_mock.h
index 1954d418dd9..70ba2f74a88 100644
--- a/src/mongo/s/catalog/sharding_catalog_client_mock.h
+++ b/src/mongo/s/catalog/sharding_catalog_client_mock.h
@@ -90,6 +90,9 @@ public:
     StatusWith<std::vector<TagsType>> getTagsForCollection(OperationContext* opCtx,
                                                            const NamespaceString& nss) override;
 
+    std::vector<NamespaceString> getAllNssThatHaveZonesForDatabase(
+        OperationContext* opCtx, const StringData& dbName) override;
+
     StatusWith<repl::OpTimeWith<std::vector<ShardType>>> getAllShards(
         OperationContext* opCtx, repl::ReadConcernLevel readConcern) override;