diff --git a/Sources/Valkey/Connection/ValkeyChannelHandler.swift b/Sources/Valkey/Connection/ValkeyChannelHandler.swift index 8d209373..0da29275 100644 --- a/Sources/Valkey/Connection/ValkeyChannelHandler.swift +++ b/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 ) } } diff --git a/Sources/Valkey/Connection/ValkeyConnectionConfiguration.swift b/Sources/Valkey/Connection/ValkeyConnectionConfiguration.swift index bf0f4f57..1e09968c 100644 --- a/Sources/Valkey/Connection/ValkeyConnectionConfiguration.swift +++ b/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,13 +142,15 @@ 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 @@ -153,6 +158,7 @@ public struct ValkeyConnectionConfiguration: Sendable { self.tls = tls self.clientName = clientName self.readOnly = readOnly + self.databaseNumber = databaseNumber } } diff --git a/Sources/Valkey/Connection/ValkeyConnectionFactory.swift b/Sources/Valkey/Connection/ValkeyConnectionFactory.swift index a63c2607..146dc7c3 100644 --- a/Sources/Valkey/Connection/ValkeyConnectionFactory.swift +++ b/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 diff --git a/Sources/Valkey/ValkeyClientConfiguration.swift b/Sources/Valkey/ValkeyClientConfiguration.swift index c22d0008..f1a80920 100644 --- a/Sources/Valkey/ValkeyClientConfiguration.swift +++ b/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,6 +190,7 @@ 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(), @@ -194,7 +198,8 @@ public struct ValkeyClientConfiguration: Sendable { 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 } } diff --git a/Tests/IntegrationTests/ClientIntegrationTests.swift b/Tests/IntegrationTests/ClientIntegrationTests.swift index 4c2b8606..9d3182c5 100644 --- a/Tests/IntegrationTests/ClientIntegrationTests.swift +++ b/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) + } + } + } + }