diff options
| -rwxr-xr-x | runtest-moduleapi | 1 | ||||
| -rw-r--r-- | src/module.c | 64 | ||||
| -rw-r--r-- | src/redismodule.h | 2 | ||||
| -rw-r--r-- | tests/modules/Makefile | 3 | ||||
| -rw-r--r-- | tests/modules/getkeys.c | 125 | ||||
| -rw-r--r-- | tests/unit/moduleapi/getkeys.tcl | 44 |
6 files changed, 238 insertions, 1 deletions
diff --git a/runtest-moduleapi b/runtest-moduleapi index e32147d5d..f881dfd3f 100755 --- a/runtest-moduleapi +++ b/runtest-moduleapi @@ -27,4 +27,5 @@ $TCLSH tests/test_helper.tcl \ --single unit/moduleapi/auth \ --single unit/moduleapi/keyspace_events \ --single unit/moduleapi/blockedclient \ +--single unit/moduleapi/getkeys \ "${@}" diff --git a/src/module.c b/src/module.c index a48780be9..1e18f3c46 100644 --- a/src/module.c +++ b/src/module.c @@ -7887,6 +7887,69 @@ int RM_ModuleTypeReplaceValue(RedisModuleKey *key, moduleType *mt, void *new_val return REDISMODULE_OK; } +/* For a specified command, parse its arguments and return an array that + * contains the indexes of all key name arguments. This function is + * essnetially a more efficient way to do COMMAND GETKEYS. + * + * A NULL return value indicates the specified command has no keys, or + * an error condition. Error conditions are indicated by setting errno + * as folllows: + * + * ENOENT: Specified command does not exist. + * EINVAL: Invalid command arity specified. + * + * NOTE: The returned array is not a Redis Module object so it does not + * get automatically freed even when auto-memory is used. The caller + * must explicitly call RM_Free() to free it. + */ +int *RM_GetCommandKeys(RedisModuleCtx *ctx, const char *cmdname, RedisModuleString **argv, int argc, int *num_keys) { + UNUSED(ctx); + struct redisCommand *cmd; + int *res = NULL; + + /* Find command */ + if ((cmd = lookupCommandByCString(cmdname)) == NULL) { + errno = ENOENT; + return NULL; + } + + /* Bail out if command has no keys */ + if (cmd->getkeys_proc == NULL && cmd->firstkey == 0) { + errno = 0; + return NULL; + } + + if ((cmd->arity > 0 && cmd->arity != argc) || (argc < -cmd->arity)) { + errno = EINVAL; + return NULL; + } + + getKeysResult result = GETKEYS_RESULT_INIT; + getKeysFromCommand(cmd, argv, argc, &result); + + *num_keys = result.numkeys; + if (!result.numkeys) { + errno = 0; + getKeysFreeResult(&result); + return NULL; + } + + if (result.keys == result.keysbuf) { + /* If the result is using a stack based array, copy it. */ + unsigned long int size = sizeof(int) * result.numkeys; + res = zmalloc(size); + memcpy(res, result.keys, size); + } else { + /* We return the heap based array and intentionally avoid calling + * getKeysFreeResult() here, as it is the caller's responsibility + * to free this array. + */ + res = result.keys; + } + + return res; +} + /* Register all the APIs we export. Keep this function at the end of the * file so that's easy to seek it to add new entries. */ void moduleRegisterCoreAPI(void) { @@ -8122,4 +8185,5 @@ void moduleRegisterCoreAPI(void) { REGISTER_API(DeauthenticateAndCloseClient); REGISTER_API(AuthenticateClientWithACLUser); REGISTER_API(AuthenticateClientWithUser); + REGISTER_API(GetCommandKeys); } diff --git a/src/redismodule.h b/src/redismodule.h index 00ca1f578..e05ce65ee 100644 --- a/src/redismodule.h +++ b/src/redismodule.h @@ -726,6 +726,7 @@ REDISMODULE_API int (*RedisModule_SetModuleUserACL)(RedisModuleUser *user, const REDISMODULE_API int (*RedisModule_AuthenticateClientWithACLUser)(RedisModuleCtx *ctx, const char *name, size_t len, RedisModuleUserChangedFunc callback, void *privdata, uint64_t *client_id) REDISMODULE_ATTR; REDISMODULE_API int (*RedisModule_AuthenticateClientWithUser)(RedisModuleCtx *ctx, RedisModuleUser *user, RedisModuleUserChangedFunc callback, void *privdata, uint64_t *client_id) REDISMODULE_ATTR; REDISMODULE_API int (*RedisModule_DeauthenticateAndCloseClient)(RedisModuleCtx *ctx, uint64_t client_id) REDISMODULE_ATTR; +REDISMODULE_API int *(*RedisModule_GetCommandKeys)(RedisModuleCtx *ctx, const char *cmdname, RedisModuleString **argv, int argc, int *num_keys) REDISMODULE_ATTR; #endif #define RedisModule_IsAOFClient(id) ((id) == CLIENT_ID_AOF) @@ -967,6 +968,7 @@ static int RedisModule_Init(RedisModuleCtx *ctx, const char *name, int ver, int REDISMODULE_GET_API(DeauthenticateAndCloseClient); REDISMODULE_GET_API(AuthenticateClientWithACLUser); REDISMODULE_GET_API(AuthenticateClientWithUser); + REDISMODULE_GET_API(GetCommandKeys); #endif if (RedisModule_IsModuleNameBusy && RedisModule_IsModuleNameBusy(name)) return REDISMODULE_ERR; diff --git a/tests/modules/Makefile b/tests/modules/Makefile index fad6e55d8..a8b08ecf0 100644 --- a/tests/modules/Makefile +++ b/tests/modules/Makefile @@ -24,7 +24,8 @@ TEST_MODULES = \ datatype.so \ auth.so \ keyspace_events.so \ - blockedclient.so + blockedclient.so \ + getkeys.so .PHONY: all diff --git a/tests/modules/getkeys.c b/tests/modules/getkeys.c new file mode 100644 index 000000000..9cde606ab --- /dev/null +++ b/tests/modules/getkeys.c @@ -0,0 +1,125 @@ +#define REDISMODULE_EXPERIMENTAL_API + +#include "redismodule.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(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) +{ + int i; + int count = 0; + + /* Handle getkeys-api introspection */ + if (RedisModule_IsKeysPositionRequest(ctx)) { + for (i = 0; i < argc; i++) { + size_t len; + const char *str = RedisModule_StringPtrLen(argv[i], &len); + + if (len == 3 && !strncasecmp(str, "key", 3) && i + 1 < argc) + RedisModule_KeyAtPos(ctx, i + 1); + } + + return REDISMODULE_OK; + } + + /* Handle real command invocation */ + RedisModule_ReplyWithArray(ctx, REDISMODULE_POSTPONED_ARRAY_LEN); + for (i = 0; i < argc; i++) { + size_t len; + const char *str = RedisModule_StringPtrLen(argv[i], &len); + + if (len == 3 && !strncasecmp(str, "key", 3) && i + 1 < argc) { + RedisModule_ReplyWithString(ctx, argv[i+1]); + count++; + } + } + RedisModule_ReplySetArrayLength(ctx, count); + + return REDISMODULE_OK; +} + +int getkeys_fixed(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) +{ + int i; + + RedisModule_ReplyWithArray(ctx, argc - 1); + for (i = 1; i < argc; i++) { + RedisModule_ReplyWithString(ctx, argv[i]); + } + return REDISMODULE_OK; +} + +/* Introspect a command using RM_GetCommandKeys() and returns the list + * of keys. Essentially this is COMMAND GETKEYS implemented in a module. + */ +int getkeys_introspect(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) +{ + UNUSED(argv); + UNUSED(argc); + + if (argc < 3) { + RedisModule_WrongArity(ctx); + return REDISMODULE_OK; + } + + size_t cmd_len; + const char *cmd = RedisModule_StringPtrLen(argv[1], &cmd_len); + + int num_keys; + int *keyidx = RedisModule_GetCommandKeys(ctx, cmd, &argv[1], argc - 1, &num_keys); + + if (!keyidx) { + if (!errno) + RedisModule_ReplyWithEmptyArray(ctx); + else { + char err[100]; + switch (errno) { + case ENOENT: + RedisModule_ReplyWithError(ctx, "ERR ENOENT"); + break; + case EINVAL: + RedisModule_ReplyWithError(ctx, "ERR EINVAL"); + break; + default: + snprintf(err, sizeof(err) - 1, "ERR errno=%d", errno); + RedisModule_ReplyWithError(ctx, err); + break; + } + } + } else { + int i; + + RedisModule_ReplyWithArray(ctx, num_keys); + for (i = 0; i < num_keys; i++) + RedisModule_ReplyWithString(ctx, argv[1 + keyidx[i]]); + + RedisModule_Free(keyidx); + } + + return REDISMODULE_OK; +} + +int RedisModule_OnLoad(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) { + UNUSED(argv); + UNUSED(argc); + if (RedisModule_Init(ctx,"getkeys",1,REDISMODULE_APIVER_1)== REDISMODULE_ERR) + return REDISMODULE_ERR; + + if (RedisModule_CreateCommand(ctx,"getkeys.command", getkeys_command,"getkeys-api",0,0,0) == REDISMODULE_ERR) + return REDISMODULE_ERR; + + if (RedisModule_CreateCommand(ctx,"getkeys.fixed", getkeys_fixed,"",2,4,1) == REDISMODULE_ERR) + return REDISMODULE_ERR; + + if (RedisModule_CreateCommand(ctx,"getkeys.introspect", getkeys_introspect,"",0,0,0) == REDISMODULE_ERR) + return REDISMODULE_ERR; + + return REDISMODULE_OK; +} diff --git a/tests/unit/moduleapi/getkeys.tcl b/tests/unit/moduleapi/getkeys.tcl new file mode 100644 index 000000000..a24a65bae --- /dev/null +++ b/tests/unit/moduleapi/getkeys.tcl @@ -0,0 +1,44 @@ +set testmodule [file normalize tests/modules/getkeys.so] + +start_server {tags {"modules"}} { + r module load $testmodule + + test {COMMAND INFO correctly reports a movable keys module command} { + set info [lindex [r command info getkeys.command] 0] + + assert_equal {movablekeys} [lindex $info 2] + assert_equal {0} [lindex $info 3] + assert_equal {0} [lindex $info 4] + assert_equal {0} [lindex $info 5] + } + + test {COMMAND GETKEYS correctly reports a movable keys module command} { + r command getkeys getkeys.command arg1 arg2 key key1 arg3 key key2 key key3 + } {key1 key2 key3} + + test {RM_GetCommandKeys on non-existing command} { + catch {r getkeys.introspect non-command key1 key2} e + set _ $e + } {*ENOENT*} + + test {RM_GetCommandKeys on built-in fixed keys command} { + r getkeys.introspect set key1 value1 + } {key1} + + test {RM_GetCommandKeys on EVAL} { + r getkeys.introspect eval "" 4 key1 key2 key3 key4 arg1 arg2 + } {key1 key2 key3 key4} + + test {RM_GetCommandKeys on a movable keys module command} { + r getkeys.introspect getkeys.command arg1 arg2 key key1 arg3 key key2 key key3 + } {key1 key2 key3} + + test {RM_GetCommandKeys on a non-movable module command} { + r getkeys.introspect getkeys.fixed arg1 key1 key2 key3 arg2 + } {key1 key2 key3} + + test {RM_GetCommandKeys with bad arity} { + catch {r getkeys.introspect set key} e + set _ $e + } {*EINVAL*} +} |