Add Support for Numbered Databases (#259)

Adding support for numbered databases 

---------

Signed-off-by: Nilanshu Sharma <[email protected]>
Co-authored-by: Nilanshu Sharma <[email protected]>

Author: nilanshu-sharma

Sources/Valkey/Connection/ValkeyChannelHandler.swift

@@ -58,6 +58,7 @@ final class ValkeyChannelHandler: ChannelInboundHandler {
         let blockingCommandTimeout: TimeAmount
         let clientName: String?
         let readOnly: Bool
+        let databaseNumber: Int
     }
     @usableFromInline
     struct PendingCommand {
@@ -292,6 +293,13 @@ final class ValkeyChannelHandler: ChannelInboundHandler {
         helloCommand.encode(into: &self.encoder)
         clientInfoLibName.encode(into: &self.encoder)
         clientInfoLibVersion.encode(into: &self.encoder)
+
+        // Select DB if needed
+        if self.configuration.databaseNumber > 0 {
+            numberOfPendingCommands += 1
+            SELECT(index: self.configuration.databaseNumber).encode(into: &self.encoder)
+        }
+
         if self.configuration.readOnly {
             numberOfPendingCommands += 1
             READONLY().encode(into: &self.encoder)
@@ -539,7 +547,8 @@ extension ValkeyChannelHandler.Configuration {
             commandTimeout: .init(other.commandTimeout),
             blockingCommandTimeout: .init(other.blockingCommandTimeout),
             clientName: other.clientName,
-            readOnly: other.readOnly
+            readOnly: other.readOnly,
+            databaseNumber: other.databaseNumber
         )
     }
 }

Sources/Valkey/Connection/ValkeyConnectionConfiguration.swift

@@ -121,6 +121,9 @@ public struct ValkeyConnectionConfiguration: Sendable {
     /// Readonly connections can run readonly commands on replica nodes
     public var readOnly: Bool
 
+    /// The number of Valkey Database
+    public var databaseNumber: Int = 0
+
     #if DistributedTracingSupport
     /// The distributed tracing configuration to use for this connection.
     /// Defaults to using the globally bootstrapped tracer with OpenTelemetry semantic conventions.
@@ -139,20 +142,23 @@ public struct ValkeyConnectionConfiguration: Sendable {
     ///   - tls: TLS configuration for secure connections. Defaults to `.disable` for unencrypted connections.
     ///   - clientName: Optional name to identify this client connection on the server. Defaults to `nil`.
     ///   - readOnly: Is the connection a readonly connection
+    ///   - databaseNumber: Database Number to use for the connection
     public init(
         authentication: Authentication? = nil,
         commandTimeout: Duration = .seconds(30),
         blockingCommandTimeout: Duration = .seconds(120),
         tls: TLS = .disable,
         clientName: String? = nil,
-        readOnly: Bool = false
+        readOnly: Bool = false,
+        databaseNumber: Int = 0
     ) {
         self.authentication = authentication
         self.commandTimeout = commandTimeout
         self.blockingCommandTimeout = blockingCommandTimeout
         self.tls = tls
         self.clientName = clientName
         self.readOnly = readOnly
+        self.databaseNumber = databaseNumber
     }
 }
 

Sources/Valkey/Connection/ValkeyConnectionFactory.swift

@@ -112,7 +112,8 @@ package final class ValkeyConnectionFactory: Sendable {
             blockingCommandTimeout: self.configuration.blockingCommandTimeout,
             tls: tls,
             clientName: nil,
-            readOnly: readOnly
+            readOnly: readOnly,
+            databaseNumber: self.configuration.databaseNumber
         )
 
         #if DistributedTracingSupport

Sources/Valkey/ValkeyClientConfiguration.swift

@@ -171,6 +171,9 @@ public struct ValkeyClientConfiguration: Sendable {
     /// The TLS to use for the Valkey connection.
     public var tls: TLS
 
+    /// Database Number to use for the Valkey Connection
+    public var databaseNumber: Int = 0
+
     #if DistributedTracingSupport
     /// The distributed tracing configuration to use for the Valkey connection.
     /// Defaults to using the globally bootstrapped tracer with OpenTelemetry semantic conventions.
@@ -187,14 +190,16 @@ public struct ValkeyClientConfiguration: Sendable {
     ///   - commandTimeout: The timeout for a connection response.
     ///   - blockingCommandTimeout: The timeout for a blocking command response.
     ///   - tls: The TLS configuration.
+    ///   - databaseNumber: The Valkey Database number.
     public init(
         authentication: Authentication? = nil,
         connectionPool: ConnectionPool = .init(),
         keepAliveBehavior: KeepAliveBehavior = .init(),
         retryParameters: RetryParameters = .init(),
         commandTimeout: Duration = .seconds(30),
         blockingCommandTimeout: Duration = .seconds(120),
-        tls: TLS = .disable
+        tls: TLS = .disable,
+        databaseNumber: Int = 0
     ) {
         self.authentication = authentication
         self.connectionPool = connectionPool
@@ -203,5 +208,6 @@ public struct ValkeyClientConfiguration: Sendable {
         self.commandTimeout = commandTimeout
         self.blockingCommandTimeout = blockingCommandTimeout
         self.tls = tls
+        self.databaseNumber = databaseNumber
     }
 }

Tests/IntegrationTests/ClientIntegrationTests.swift

@@ -592,4 +592,39 @@ struct ClientIntegratedTests {
             #expect(clients.firstRange(of: "lib-ver=\(valkeySwiftLibraryVersion)") != nil)
         }
     }
+
+    @Test
+    @available(valkeySwift 1.0, *)
+    func testMultipleDB() async throws {
+        var logger = Logger(label: "Valkey")
+        logger.logLevel = .debug
+        // Test all default enabled databases in range {0,15}
+        for dbNum in 0...15 {
+            let clientConfig: ValkeyClientConfiguration = .init(databaseNumber: dbNum)
+            try await withValkeyConnection(.hostname(valkeyHostname, port: 6379), configuration: clientConfig, logger: logger) { connection in
+                // Verify ClientInfo contains dbNum
+                let clientInfo = String(buffer: try await connection.clientInfo())
+                #expect(clientInfo.contains("db=\(dbNum)"))
+
+                // Verify via setting and getting keys on all the DBs
+                let key = "key-\(dbNum)"
+                let value = "value-\(dbNum)"
+                try await connection.set(ValkeyKey(key), value: value)
+                let response = try await connection.get(ValkeyKey(key)).map { String(buffer: $0) }
+                #expect(response == value)
+
+                // Verify key belonging to other DBs don't exist in this DB
+                for otherDbNum in 0...15 {
+                    let otherKey = "key-\(otherDbNum)"
+                    if otherDbNum == dbNum { continue }
+                    let otherResponse = try await connection.get(ValkeyKey(otherKey)).map { String(buffer: $0) }
+                    #expect(otherResponse == nil)
+                }
+
+                let delCount = try await connection.del(keys: [ValkeyKey(key)])
+                #expect(delCount == 1)
+            }
+        }
+    }
+
 }