Add option to select a logical database (valkey-io#244)

Set the option `select_db` in `valkeyClusterOptions` to select a logical
database after a successful connect.
The option is supported in both the Synchronous and Asynchronous API,
and the `SELECT` command is only sent when configured to a non-zero
value.

Additionally adds new flags to the clusterclients `[--select-db NUM]`
and update tests.

Signed-off-by: Björn Svensson <bjorn.a.svensson@est.tech>

Author: bjosv

include/valkey/cluster.h

@@ -93,6 +93,7 @@ typedef struct valkeyClusterContext {
     int max_retry_count;             /* Allowed retry attempts */
     char *username;                  /* Authenticate using user */
     char *password;                  /* Authentication password */
+    int select_db;
 
     struct dict *nodes;        /* Known valkeyClusterNode's */
     uint64_t route_version;    /* Increased when the node lookup table changes */
@@ -163,6 +164,10 @@ typedef struct {
     const char *password;                  /* Authentication password. */
     int max_retry;                         /* Allowed retry attempts. */
 
+    /* Select a logical database after a successful connect.
+     * Default 0, i.e. the SELECT command is not sent. */
+    int select_db;
+
     /* Common callbacks. */
 
     /* A hook to get notified when certain events occur. The `event` is set to

src/cluster.c

@@ -358,6 +358,25 @@ static int authenticate(valkeyClusterContext *cc, valkeyContext *c) {
     return VALKEY_ERR;
 }
 
+/* Select a logical database by sending the SELECT command. */
+static int select_db(valkeyClusterContext *cc, valkeyContext *c) {
+    if (cc->select_db == 0)
+        return VALKEY_OK;
+
+    valkeyReply *reply = valkeyCommand(c, "SELECT %d", cc->select_db);
+    if (reply == NULL) {
+        valkeyClusterSetError(cc, VALKEY_ERR_OTHER, "Failed to select logical database");
+        return VALKEY_ERR;
+    }
+    if (reply->type == VALKEY_REPLY_ERROR) {
+        valkeyClusterSetError(cc, VALKEY_ERR_OTHER, reply->str);
+        freeReplyObject(reply);
+        return VALKEY_ERR;
+    }
+    freeReplyObject(reply);
+    return VALKEY_OK;
+}
+
 /**
  * Return a new node with the "cluster slots" command reply.
  */
@@ -1189,6 +1208,9 @@ static int valkeyClusterContextInit(valkeyClusterContext *cc,
     } else {
         cc->max_retry_count = CLUSTER_DEFAULT_MAX_RETRY_COUNT;
     }
+    if (options->select_db > 0) {
+        cc->select_db = options->select_db;
+    }
     if (options->initial_nodes != NULL &&
         valkeyClusterSetOptionAddNodes(cc, options->initial_nodes) != VALKEY_OK) {
         return VALKEY_ERR; /* err and errstr already set. */
@@ -1563,8 +1585,10 @@ valkeyContext *valkeyClusterGetValkeyContext(valkeyClusterContext *cc,
             if (cc->tls && cc->tls_init_fn(c, cc->tls) != VALKEY_OK) {
                 valkeyClusterSetError(cc, c->err, c->errstr);
             }
-
-            authenticate(cc, c); // err and errstr handled in function
+            /* Authenticate and select a logical database when configured.
+             * cc->err and cc->errstr are set when failing. */
+            authenticate(cc, c);
+            select_db(cc, c);
         }
 
         return c;
@@ -1606,6 +1630,10 @@ valkeyContext *valkeyClusterGetValkeyContext(valkeyClusterContext *cc,
         valkeyFree(c);
         return NULL;
     }
+    if (select_db(cc, c) != VALKEY_OK) {
+        valkeyFree(c);
+        return NULL;
+    }
 
     node->con = c;
 
@@ -2582,6 +2610,18 @@ static void unlinkAsyncContextAndNode(void *data) {
     }
 }
 
+/* Reply callback function for SELECT */
+void selectReplyCallback(valkeyAsyncContext *ac, void *r, void *privdata) {
+    valkeyReply *reply = (valkeyReply *)r;
+    valkeyClusterAsyncContext *acc = (valkeyClusterAsyncContext *)privdata;
+
+    if (reply == NULL || reply->type == VALKEY_REPLY_ERROR) {
+        valkeyClusterAsyncSetError(acc, VALKEY_ERR_OTHER,
+                                   "Failed to select logical database");
+        valkeyAsyncDisconnect(ac);
+    }
+}
+
 valkeyAsyncContext *
 valkeyClusterGetValkeyAsyncContext(valkeyClusterAsyncContext *acc,
                                    valkeyClusterNode *node) {
@@ -2658,6 +2698,15 @@ valkeyClusterGetValkeyAsyncContext(valkeyClusterAsyncContext *acc,
             return NULL;
         }
     }
+    // Select logical database when needed
+    if (acc->cc.select_db > 0) {
+        ret = valkeyAsyncCommand(ac, selectReplyCallback, acc, "SELECT %d", acc->cc.select_db);
+        if (ret != VALKEY_OK) {
+            valkeyClusterAsyncSetError(acc, ac->c.err, ac->c.errstr);
+            valkeyAsyncFree(ac);
+            return NULL;
+        }
+    }
 
     if (acc->attach_fn) {
         ret = acc->attach_fn(ac, acc->attach_data);

tests/clusterclient.c

@@ -61,27 +61,45 @@ void eventCallback(const valkeyClusterContext *cc, int event, void *privdata) {
     printf("Event: %s\n", e);
 }
 
+void connectCallback(const valkeyContext *c, int status) {
+    const char *s = "";
+    if (status != VALKEY_OK)
+        s = "failed to ";
+    printf("Event: %sconnect to %s:%d\n", s, c->tcp.host, c->tcp.port);
+}
+
 int main(int argc, char **argv) {
     int show_events = 0;
     int use_cluster_nodes = 0;
     int send_to_all = 0;
+    int show_connection_events = 0;
+    int select_db = 0;
 
     int argindex;
     for (argindex = 1; argindex < argc && argv[argindex][0] == '-';
          argindex++) {
         if (strcmp(argv[argindex], "--events") == 0) {
             show_events = 1;
+        } else if (strcmp(argv[argindex], "--connection-events") == 0) {
+            show_connection_events = 1;
         } else if (strcmp(argv[argindex], "--use-cluster-nodes") == 0) {
             use_cluster_nodes = 1;
+        } else if (strcmp(argv[argindex], "--select-db") == 0) {
+            if (++argindex < argc) /* Need an additional argument */
+                select_db = atoi(argv[argindex]);
+            if (select_db == 0) {
+                fprintf(stderr, "Missing or faulty argument for --select-db\n");
+                exit(1);
+            }
         } else {
             fprintf(stderr, "Unknown argument: '%s'\n", argv[argindex]);
             exit(1);
         }
     }
 
     if (argindex >= argc) {
-        fprintf(stderr, "Usage: clusterclient [--events] [--use-cluster-nodes] "
-                        "HOST:PORT\n");
+        fprintf(stderr, "Usage: clusterclient [--events] [--connection-events] "
+                        "[--use-cluster-nodes] [--select-db NUM] HOST:PORT\n");
         exit(1);
     }
     const char *initnode = argv[argindex];
@@ -97,6 +115,12 @@ int main(int argc, char **argv) {
     if (show_events) {
         options.event_callback = eventCallback;
     }
+    if (show_connection_events) {
+        options.connect_callback = connectCallback;
+    }
+    if (select_db > 0) {
+        options.select_db = select_db;
+    }
 
     valkeyClusterContext *cc = valkeyClusterConnectWithOptions(&options);
     if (cc == NULL || cc->err) {

tests/clusterclient_async.c

@@ -242,6 +242,7 @@ void disconnectCallback(const valkeyAsyncContext *ac, int status) {
 int main(int argc, char **argv) {
     int use_cluster_nodes = 0;
     int show_connection_events = 0;
+    int select_db = 0;
 
     int optind;
     for (optind = 1; optind < argc && argv[optind][0] == '-'; optind++) {
@@ -253,14 +254,22 @@ int main(int argc, char **argv) {
             show_connection_events = 1;
         } else if (strcmp(argv[optind], "--blocking-initial-update") == 0) {
             blocking_initial_update = 1;
+        } else if (strcmp(argv[optind], "--select-db") == 0) {
+            if (++optind < argc) /* Need an additional argument */
+                select_db = atoi(argv[optind]);
+            if (select_db == 0) {
+                fprintf(stderr, "Missing or faulty argument for --select-db\n");
+                exit(1);
+            }
         } else {
             fprintf(stderr, "Unknown argument: '%s'\n", argv[optind]);
         }
     }
 
     if (optind >= argc) {
         fprintf(stderr,
-                "Usage: clusterclient_async [--use-cluster-nodes] HOST:PORT\n");
+                "Usage: clusterclient_async [--events] [--connection-events] "
+                "[--use-cluster-nodes] [--select-db NUM] HOST:PORT\n");
         exit(1);
     }
     const char *initnode = argv[optind];
@@ -283,6 +292,9 @@ int main(int argc, char **argv) {
         options.async_connect_callback = connectCallback;
         options.async_disconnect_callback = disconnectCallback;
     }
+    if (select_db > 0) {
+        options.select_db = select_db;
+    }
     valkeyClusterOptionsUseLibevent(&options, base);
 
     valkeyClusterAsyncContext *acc = valkeyClusterAsyncConnectWithOptions(&options);

tests/scripts/client-disconnect-test.sh

@@ -16,11 +16,15 @@ syncpid1=$!;
 # Start simulated valkey node
 timeout 5s ./simulated-valkey.pl -p 7401 -d --sigcont $syncpid1 <<'EOF' &
 EXPECT CONNECT
+EXPECT ["SELECT", "2"]
+SEND +OK
 EXPECT ["CLUSTER", "SLOTS"]
 SEND [[0, 16383, ["127.0.0.1", 7401, "nodeid1"]]]
 EXPECT CLOSE
 
 EXPECT CONNECT
+EXPECT ["SELECT", "2"]
+SEND +OK
 EXPECT ["SET", "foo", "initial"]
 SEND +OK
 
@@ -35,7 +39,7 @@ server1=$!
 wait $syncpid1;
 
 # Run client
-timeout 4s "$clientprog" --blocking-initial-update --connection-events 127.0.0.1:7401 > "$testname.out" <<'EOF'
+timeout 4s "$clientprog" --blocking-initial-update --connection-events --select-db 2 127.0.0.1:7401 > "$testname.out" <<'EOF'
 SET foo initial
 
 # Send a command that is expected to be redirected just before

tests/scripts/slots-not-served-test.sh

@@ -13,6 +13,8 @@ syncpid1=$!;
 timeout 5s ./simulated-valkey.pl -p 7401 -d --sigcont $syncpid1 <<'EOF' &
 # The initial slotmap is not covering all slots.
 EXPECT CONNECT
+EXPECT ["SELECT", "5"]
+SEND +OK
 EXPECT ["CLUSTER", "SLOTS"]
 SEND [[0, 1, ["127.0.0.1", 7401, "nodeid7401"]]]
 
@@ -36,7 +38,7 @@ server1=$!
 wait $syncpid1;
 
 # Run client
-timeout 3s "$clientprog" --events 127.0.0.1:7401 > "$testname.out" <<'EOF'
+timeout 3s "$clientprog" --events --connection-events  --select-db 5 127.0.0.1:7401 > "$testname.out" <<'EOF'
 GET foo1
 GET foo2
 EOF
@@ -56,7 +58,8 @@ if [ $clientexit -ne 0 ]; then
 fi
 
 # Check the output from clusterclient
-expected="Event: slotmap-updated
+expected="Event: connect to 127.0.0.1:7401
+Event: slotmap-updated
 Event: ready
 Event: slotmap-updated
 error: slot not served by any node