redict/tests/modules/getkeys.c

// SPDX-FileCopyrightText: 2024 Redict Contributors
// SPDX-FileCopyrightText: 2024 Salvatore Sanfilippo <antirez at gmail dot com>
//
// SPDX-License-Identifier: BSD-3-Clause
// SPDX-License-Identifier: LGPL-3.0-only

#include "redictmodule.h"
#include <strings.h>
#include <assert.h>
#include <unistd.h>
#include <errno.h>

#define UNUSED(V) ((void) V)

/* A sample movable keys command that returns a list of all
 * arguments that follow a KEY argument, i.e.
 */
int getkeys_command(RedictModuleCtx *ctx, RedictModuleString **argv, int argc)
{
    int i;
    int count = 0;

    /* Handle getkeys-api introspection */
    if (RedictModule_IsKeysPositionRequest(ctx)) {
        for (i = 0; i < argc; i++) {
            size_t len;
            const char *str = RedictModule_StringPtrLen(argv[i], &len);

            if (len == 3 && !strncasecmp(str, "key", 3) && i + 1 < argc)
                RedictModule_KeyAtPos(ctx, i + 1);
        }

        return REDICTMODULE_OK;
    }

    /* Handle real command invocation */
    RedictModule_ReplyWithArray(ctx, REDICTMODULE_POSTPONED_LEN);
    for (i = 0; i < argc; i++) {
        size_t len;
        const char *str = RedictModule_StringPtrLen(argv[i], &len);

        if (len == 3 && !strncasecmp(str, "key", 3) && i + 1 < argc) {
            RedictModule_ReplyWithString(ctx, argv[i+1]);
            count++;
        }
    }
    RedictModule_ReplySetArrayLength(ctx, count);

    return REDICTMODULE_OK;
}

int getkeys_command_with_flags(RedictModuleCtx *ctx, RedictModuleString **argv, int argc)
{
    int i;
    int count = 0;

    /* Handle getkeys-api introspection */
    if (RedictModule_IsKeysPositionRequest(ctx)) {
        for (i = 0; i < argc; i++) {
            size_t len;
            const char *str = RedictModule_StringPtrLen(argv[i], &len);

            if (len == 3 && !strncasecmp(str, "key", 3) && i + 1 < argc)
                RedictModule_KeyAtPosWithFlags(ctx, i + 1, REDICTMODULE_CMD_KEY_RO | REDICTMODULE_CMD_KEY_ACCESS);
        }

        return REDICTMODULE_OK;
    }

    /* Handle real command invocation */
    RedictModule_ReplyWithArray(ctx, REDICTMODULE_POSTPONED_LEN);
    for (i = 0; i < argc; i++) {
        size_t len;
        const char *str = RedictModule_StringPtrLen(argv[i], &len);

        if (len == 3 && !strncasecmp(str, "key", 3) && i + 1 < argc) {
            RedictModule_ReplyWithString(ctx, argv[i+1]);
            count++;
        }
    }
    RedictModule_ReplySetArrayLength(ctx, count);

    return REDICTMODULE_OK;
}

int getkeys_fixed(RedictModuleCtx *ctx, RedictModuleString **argv, int argc)
{
    int i;

    RedictModule_ReplyWithArray(ctx, argc - 1);
    for (i = 1; i < argc; i++) {
        RedictModule_ReplyWithString(ctx, argv[i]);
    }
    return REDICTMODULE_OK;
}

/* Introspect a command using RM_GetCommandKeys() and returns the list
 * of keys. Essentially this is COMMAND GETKEYS implemented in a module.
 * INTROSPECT <with-flags> <cmd> <args>
 */
int getkeys_introspect(RedictModuleCtx *ctx, RedictModuleString **argv, int argc)
{
    long long with_flags = 0;

    if (argc < 4) {
        RedictModule_WrongArity(ctx);
        return REDICTMODULE_OK;
    }

    if (RedictModule_StringToLongLong(argv[1],&with_flags) != REDICTMODULE_OK)
        return RedictModule_ReplyWithError(ctx,"ERR invalid integer");

    int num_keys, *keyflags = NULL;
    int *keyidx = RedictModule_GetCommandKeysWithFlags(ctx, &argv[2], argc - 2, &num_keys, with_flags ? &keyflags : NULL);

    if (!keyidx) {
        if (!errno)
            RedictModule_ReplyWithEmptyArray(ctx);
        else {
            char err[100];
            switch (errno) {
                case ENOENT:
                    RedictModule_ReplyWithError(ctx, "ERR ENOENT");
                    break;
                case EINVAL:
                    RedictModule_ReplyWithError(ctx, "ERR EINVAL");
                    break;
                default:
                    snprintf(err, sizeof(err) - 1, "ERR errno=%d", errno);
                    RedictModule_ReplyWithError(ctx, err);
                    break;
            }
        }
    } else {
        int i;

        RedictModule_ReplyWithArray(ctx, num_keys);
        for (i = 0; i < num_keys; i++) {
            if (!with_flags) {
                RedictModule_ReplyWithString(ctx, argv[2 + keyidx[i]]);
                continue;
            }
            RedictModule_ReplyWithArray(ctx, 2);
            RedictModule_ReplyWithString(ctx, argv[2 + keyidx[i]]);
            char* sflags = "";
            if (keyflags[i] & REDICTMODULE_CMD_KEY_RO)
                sflags = "RO";
            else if (keyflags[i] & REDICTMODULE_CMD_KEY_RW)
                sflags = "RW";
            else if (keyflags[i] & REDICTMODULE_CMD_KEY_OW)
                sflags = "OW";
            else if (keyflags[i] & REDICTMODULE_CMD_KEY_RM)
                sflags = "RM";
            RedictModule_ReplyWithCString(ctx, sflags);
        }

        RedictModule_Free(keyidx);
        RedictModule_Free(keyflags);
    }

    return REDICTMODULE_OK;
}

int RedictModule_OnLoad(RedictModuleCtx *ctx, RedictModuleString **argv, int argc) {
    UNUSED(argv);
    UNUSED(argc);
    if (RedictModule_Init(ctx,"getkeys",1,REDICTMODULE_APIVER_1)== REDICTMODULE_ERR)
        return REDICTMODULE_ERR;

    if (RedictModule_CreateCommand(ctx,"getkeys.command", getkeys_command,"getkeys-api",0,0,0) == REDICTMODULE_ERR)
        return REDICTMODULE_ERR;

    if (RedictModule_CreateCommand(ctx,"getkeys.command_with_flags", getkeys_command_with_flags,"getkeys-api",0,0,0) == REDICTMODULE_ERR)
        return REDICTMODULE_ERR;

    if (RedictModule_CreateCommand(ctx,"getkeys.fixed", getkeys_fixed,"",2,4,1) == REDICTMODULE_ERR)
        return REDICTMODULE_ERR;

    if (RedictModule_CreateCommand(ctx,"getkeys.introspect", getkeys_introspect,"",0,0,0) == REDICTMODULE_ERR)
        return REDICTMODULE_ERR;

    return REDICTMODULE_OK;
}