redict/tests/modules/scan.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 <string.h>
#include <assert.h>
#include <unistd.h>

typedef struct {
    size_t nkeys;
} scan_strings_pd;

void scan_strings_callback(RedictModuleCtx *ctx, RedictModuleString* keyname, RedictModuleKey* key, void *privdata) {
    scan_strings_pd* pd = privdata;
    int was_opened = 0;
    if (!key) {
        key = RedictModule_OpenKey(ctx, keyname, REDICTMODULE_READ);
        was_opened = 1;
    }

    if (RedictModule_KeyType(key) == REDICTMODULE_KEYTYPE_STRING) {
        size_t len;
        char * data = RedictModule_StringDMA(key, &len, REDICTMODULE_READ);
        RedictModule_ReplyWithArray(ctx, 2);
        RedictModule_ReplyWithString(ctx, keyname);
        RedictModule_ReplyWithStringBuffer(ctx, data, len);
        pd->nkeys++;
    }
    if (was_opened)
        RedictModule_CloseKey(key);
}

int scan_strings(RedictModuleCtx *ctx, RedictModuleString **argv, int argc)
{
    REDICTMODULE_NOT_USED(argv);
    REDICTMODULE_NOT_USED(argc);
    scan_strings_pd pd = {
        .nkeys = 0,
    };

    RedictModule_ReplyWithArray(ctx, REDICTMODULE_POSTPONED_LEN);

    RedictModuleScanCursor* cursor = RedictModule_ScanCursorCreate();
    while(RedictModule_Scan(ctx, cursor, scan_strings_callback, &pd));
    RedictModule_ScanCursorDestroy(cursor);

    RedictModule_ReplySetArrayLength(ctx, pd.nkeys);
    return REDICTMODULE_OK;
}

typedef struct {
    RedictModuleCtx *ctx;
    size_t nreplies;
} scan_key_pd;

void scan_key_callback(RedictModuleKey *key, RedictModuleString* field, RedictModuleString* value, void *privdata) {
    REDICTMODULE_NOT_USED(key);
    scan_key_pd* pd = privdata;
    RedictModule_ReplyWithArray(pd->ctx, 2);
    size_t fieldCStrLen;

    // The implementation of RedictModuleString is robj with lots of encodings.
    // We want to make sure the robj that passes to this callback in
    // String encoded, this is why we use RedictModule_StringPtrLen and
    // RedictModule_ReplyWithStringBuffer instead of directly use
    // RedictModule_ReplyWithString.
    const char* fieldCStr = RedictModule_StringPtrLen(field, &fieldCStrLen);
    RedictModule_ReplyWithStringBuffer(pd->ctx, fieldCStr, fieldCStrLen);
    if(value){
        size_t valueCStrLen;
        const char* valueCStr = RedictModule_StringPtrLen(value, &valueCStrLen);
        RedictModule_ReplyWithStringBuffer(pd->ctx, valueCStr, valueCStrLen);
    } else {
        RedictModule_ReplyWithNull(pd->ctx);
    }

    pd->nreplies++;
}

int scan_key(RedictModuleCtx *ctx, RedictModuleString **argv, int argc)
{
    if (argc != 2) {
        RedictModule_WrongArity(ctx);
        return REDICTMODULE_OK;
    }
    scan_key_pd pd = {
        .ctx = ctx,
        .nreplies = 0,
    };

    RedictModuleKey *key = RedictModule_OpenKey(ctx, argv[1], REDICTMODULE_READ);
    if (!key) {
        RedictModule_ReplyWithError(ctx, "not found");
        return REDICTMODULE_OK;
    }

    RedictModule_ReplyWithArray(ctx, REDICTMODULE_POSTPONED_ARRAY_LEN);

    RedictModuleScanCursor* cursor = RedictModule_ScanCursorCreate();
    while(RedictModule_ScanKey(key, cursor, scan_key_callback, &pd));
    RedictModule_ScanCursorDestroy(cursor);

    RedictModule_ReplySetArrayLength(ctx, pd.nreplies);
    RedictModule_CloseKey(key);
    return REDICTMODULE_OK;
}

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

    if (RedictModule_CreateCommand(ctx, "scan.scan_strings", scan_strings, "", 0, 0, 0) == REDICTMODULE_ERR)
        return REDICTMODULE_ERR;

    if (RedictModule_CreateCommand(ctx, "scan.scan_key", scan_key, "", 0, 0, 0) == REDICTMODULE_ERR)
        return REDICTMODULE_ERR;

    return REDICTMODULE_OK;
}