SWIFT-1447 Add filtering by tag sets to finding suitable servers (#727)

Author: nspektor

Sources/MongoSwift/MongoConnectionString.swift

@@ -20,7 +20,7 @@ public struct MongoConnectionString: Codable, LosslessStringConvertible {
     /// Characters that do not need to be percent-encoded when reconstructing URI options.
     fileprivate static let allowedForOptionEncoding = CharacterSet(charactersIn: "=&,:").inverted
 
-    fileprivate enum OptionName: String {
+    internal enum OptionName: String {
         case appName = "appname"
         case authSource = "authsource"
         case authMechanism = "authmechanism"

Sources/MongoSwift/ReadPreference.swift

@@ -202,6 +202,16 @@ public struct ReadPreference: Equatable {
                 )
             }
         }
+        if let tagSets = tagSets {
+            if mode == .primary {
+                guard tagSets == [BSONDocument()] else {
+                    throw MongoError.InvalidArgumentError(
+                        message: "Invalid \(MongoConnectionString.OptionName.readPreferenceTags) \(tagSets): " +
+                            "when mode is primary, tag_sets must be empty"
+                    )
+                }
+            }
+        }
         self.mode = mode
         self.tagSets = tagSets
         self.maxStalenessSeconds = maxStalenessSeconds

Sources/MongoSwift/SDAM.swift

@@ -225,7 +225,7 @@ public struct ServerDescription {
     }
 
     // For testing purposes
-    internal init(type: ServerType) {
+    internal init(type: ServerType, tags: [String: String]? = nil) {
         self.type = type
         self.address = ServerAddress(host: "fake", port: 80)
         self.serverId = 0
@@ -243,7 +243,7 @@ public struct ServerDescription {
         self.hosts = []
         self.passives = []
         self.arbiters = []
-        self.tags = [:]
+        self.tags = tags ?? [:]
     }
 }
 
@@ -417,7 +417,7 @@ extension TopologyDescription {
                 }
                 let secondaries = self.servers.filter { $0.type == .rsSecondary }
                 return self.filterReplicaSetServers(readPreference: readPreference, servers: secondaries)
-            default:
+            default: // or .primary
                 // the default mode is 'primary'.
                 return self.servers.filter { $0.type == .rsPrimary }
             }
@@ -427,12 +427,22 @@ extension TopologyDescription {
     }
 
     internal func filterReplicaSetServers(
-        readPreference _: ReadPreference?,
+        readPreference: ReadPreference?,
         servers: [ServerDescription]
     ) -> [ServerDescription] {
         // TODO: Filter out servers staler than maxStalenessSeconds
-        // TODO: Select servers matching the tag_sets
-        // While waiting for the above to be implemented, this helper just returns the servers it was passed
-        servers
+
+        // Filter by tag_sets
+        let tagSets = readPreference?.tagSets ?? []
+        if tagSets.isEmpty {
+            return servers
+        }
+        for tagSet in tagSets {
+            let matches = servers.filter { server in tagSet.allSatisfy { server.tags[$0.key] == $0.value.stringValue } }
+            if !matches.isEmpty {
+                return matches
+            }
+        }
+        return []
     }
 }

Tests/MongoSwiftTests/ServerSelectionTests.swift

@@ -6,56 +6,70 @@ import TestsCommon
 import XCTest
 
 final class ServerSelectionTests: MongoSwiftTestCase {
-    func testServerSelection() throws {
-        let standaloneServer = ServerDescription(type: .standalone)
-        let rsPrimaryServer = ServerDescription(type: .rsPrimary)
-        let rsSecondaryServer1 = ServerDescription(type: .rsSecondary)
-        let rsSecondaryServer2 = ServerDescription(type: .rsSecondary)
-        let mongosServer = ServerDescription(type: .mongos)
-
-        // unknown
+    // Servers
+    let standaloneServer = ServerDescription(type: .standalone)
+    let rsPrimaryServer = ServerDescription(type: .rsPrimary)
+    let rsSecondaryServer1 = ServerDescription(type: .rsSecondary, tags: ["dc": "ny", "rack": "2", "size": "large"])
+    let rsSecondaryServer2 = ServerDescription(type: .rsSecondary)
+    let rsSecondaryServer3 = ServerDescription(type: .rsSecondary, tags: ["dc": "ny", "rack": "3", "size": "small"])
+    let mongosServer = ServerDescription(type: .mongos)
+
+    // Read Preferences
+    let primaryReadPreference = ReadPreference(.primary)
+    let primaryPrefReadPreferemce = ReadPreference(.primaryPreferred)
+
+    // Tag Sets
+    let tagSet: BSONDocument = ["dc": "ny", "rack": "2"]
+    let tagSet2: BSONDocument = ["dc": "ny"]
+    let tagSet3: BSONDocument = ["size": "small"]
+
+    func testUnknownTopology() {
         let unkownTopology = TopologyDescription(type: .unknown, servers: [standaloneServer])
         expect(unkownTopology.findSuitableServers()).to(haveCount(0))
+    }
 
-        // single
+    func testSingleTopology() {
         let singleTopology = TopologyDescription(type: .single, servers: [standaloneServer])
         expect(singleTopology.findSuitableServers()[0].type).to(equal(.standalone))
+    }
 
-        // replica set with primary
+    func testReplicaSetWithPrimaryTopology() {
         let replicaSetTopology = TopologyDescription(type: .replicaSetWithPrimary, servers: [
             rsPrimaryServer,
             rsSecondaryServer1,
             rsSecondaryServer2
         ])
-        let primaryReadPreference = ReadPreference(.primary)
-        let replicaSetSuitableServers = replicaSetTopology.findSuitableServers(readPreference: primaryReadPreference)
+        let replicaSetSuitableServers = replicaSetTopology
+            .findSuitableServers(readPreference: self.primaryReadPreference)
         expect(replicaSetSuitableServers[0].type).to(equal(.rsPrimary))
         expect(replicaSetSuitableServers).to(haveCount(1))
 
-        let primaryPrefReadPreferemce = ReadPreference(.primaryPreferred)
         let replicaSetSuitableServers2 = replicaSetTopology
-            .findSuitableServers(readPreference: primaryPrefReadPreferemce)
+            .findSuitableServers(readPreference: self.primaryPrefReadPreferemce)
         expect(replicaSetSuitableServers2[0].type).to(equal(.rsPrimary))
         expect(replicaSetSuitableServers2).to(haveCount(1))
+    }
 
-        // replica set without primary
+    func testReplicaSetNoPrimaryTopology() {
         let replicaSetNoPrimaryTopology = TopologyDescription(type: .replicaSetNoPrimary, servers: [
             rsSecondaryServer1,
             rsSecondaryServer2
         ])
-        let replicaSetNoPrimarySuitableServers = replicaSetNoPrimaryTopology
-            .findSuitableServers(readPreference: primaryReadPreference)
-        expect(replicaSetNoPrimarySuitableServers).to(haveCount(0))
+        let suitable1 = replicaSetNoPrimaryTopology
+            .findSuitableServers(readPreference: self.primaryReadPreference)
+        expect(suitable1).to(haveCount(0))
 
-        let replicaSetNoPrimarySuitableServer2 = replicaSetNoPrimaryTopology.findSuitableServers(readPreference: nil)
-        expect(replicaSetNoPrimarySuitableServer2).to(haveCount(0))
+        let suitable2 = replicaSetNoPrimaryTopology
+            .findSuitableServers(readPreference: nil)
+        expect(suitable2).to(haveCount(0))
 
-        let replicaSetNoPrimarySuitableServer3 = replicaSetNoPrimaryTopology
-            .findSuitableServers(readPreference: primaryPrefReadPreferemce)
-        expect(replicaSetNoPrimarySuitableServer3[0].type).to(equal(.rsSecondary))
-        expect(replicaSetNoPrimarySuitableServer3).to(haveCount(2))
+        let suitable3 = replicaSetNoPrimaryTopology
+            .findSuitableServers(readPreference: self.primaryPrefReadPreferemce)
+        expect(suitable3[0].type).to(equal(.rsSecondary))
+        expect(suitable3).to(haveCount(2))
+    }
 
-        // sharded
+    func testShardedTopology() {
         let shardedTopology = TopologyDescription(type: .sharded, servers: [
             mongosServer
         ])
@@ -64,4 +78,59 @@ final class ServerSelectionTests: MongoSwiftTestCase {
             .to(equal(.mongos))
         expect(shardedSuitableServers).to(haveCount(1))
     }
+
+    func testTagSets() throws {
+        // tag set 1
+        let topology = TopologyDescription(type: .replicaSetNoPrimary, servers: [
+            rsSecondaryServer1,
+            rsSecondaryServer2,
+            rsSecondaryServer3
+        ])
+        let secondaryReadPreferenceWithTagSet = try ReadPreference(
+            .secondaryPreferred,
+            tagSets: [tagSet, tagSet3], // tagSet3 should be ignored, because tagSet matches some servers
+            maxStalenessSeconds: nil
+        )
+
+        let suitable = topology.findSuitableServers(readPreference: secondaryReadPreferenceWithTagSet)
+        expect(suitable[0].type).to(equal(.rsSecondary))
+        expect(suitable).to(haveCount(1))
+
+        // tag set 2
+        let secondaryReadPreferenceWithTagSet2 = try ReadPreference(
+            .secondaryPreferred,
+            tagSets: [tagSet2],
+            maxStalenessSeconds: nil
+        )
+
+        let suitable2 = topology.findSuitableServers(readPreference: secondaryReadPreferenceWithTagSet2)
+        expect(suitable2[0].type).to(equal(.rsSecondary))
+        expect(suitable2).to(haveCount(2))
+
+        // invalid tag set passing
+        expect(try ReadPreference(
+            .primary,
+            tagSets: [self.tagSet],
+            maxStalenessSeconds: nil
+        )).to(throwError(errorType: MongoError.InvalidArgumentError.self))
+
+        // valid tag set passing
+        let replicaSetTopology = TopologyDescription(type: .replicaSetWithPrimary, servers: [
+            rsPrimaryServer,
+            rsSecondaryServer1,
+            rsSecondaryServer2
+        ])
+
+        let emptyTagSet: BSONDocument = [:]
+
+        let primaryReadPreferenceWithEmptyTagSet = try ReadPreference(
+            .primary,
+            tagSets: [emptyTagSet],
+            maxStalenessSeconds: nil
+        )
+        let replicaSetSuitableServers = replicaSetTopology
+            .findSuitableServers(readPreference: primaryReadPreferenceWithEmptyTagSet)
+        expect(replicaSetSuitableServers[0].type).to(equal(.rsPrimary))
+        expect(replicaSetSuitableServers).to(haveCount(1))
+    }
 }