1 /* This module current tests a small subset but should be extended in the future
2 * for general ModuleDataType coverage.
3 */
4
5 #include "redismodule.h"
6
7 static RedisModuleType *datatype = NULL;
8
9 typedef struct {
10 long long intval;
11 RedisModuleString *strval;
12 } DataType;
13
datatype_load(RedisModuleIO * io,int encver)14 static void *datatype_load(RedisModuleIO *io, int encver) {
15 (void) encver;
16
17 int intval = RedisModule_LoadSigned(io);
18 if (RedisModule_IsIOError(io)) return NULL;
19
20 RedisModuleString *strval = RedisModule_LoadString(io);
21 if (RedisModule_IsIOError(io)) return NULL;
22
23 DataType *dt = (DataType *) RedisModule_Alloc(sizeof(DataType));
24 dt->intval = intval;
25 dt->strval = strval;
26 return dt;
27 }
28
datatype_save(RedisModuleIO * io,void * value)29 static void datatype_save(RedisModuleIO *io, void *value) {
30 DataType *dt = (DataType *) value;
31 RedisModule_SaveSigned(io, dt->intval);
32 RedisModule_SaveString(io, dt->strval);
33 }
34
datatype_free(void * value)35 static void datatype_free(void *value) {
36 if (value) {
37 DataType *dt = (DataType *) value;
38
39 if (dt->strval) RedisModule_FreeString(NULL, dt->strval);
40 RedisModule_Free(dt);
41 }
42 }
43
datatype_set(RedisModuleCtx * ctx,RedisModuleString ** argv,int argc)44 static int datatype_set(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
45 if (argc != 4) {
46 RedisModule_WrongArity(ctx);
47 return REDISMODULE_OK;
48 }
49
50 long long intval;
51
52 if (RedisModule_StringToLongLong(argv[2], &intval) != REDISMODULE_OK) {
53 RedisModule_ReplyWithError(ctx, "Invalid integr value");
54 return REDISMODULE_OK;
55 }
56
57 RedisModuleKey *key = RedisModule_OpenKey(ctx, argv[1], REDISMODULE_WRITE);
58 DataType *dt = RedisModule_Calloc(sizeof(DataType), 1);
59 dt->intval = intval;
60 dt->strval = argv[3];
61 RedisModule_RetainString(ctx, dt->strval);
62
63 RedisModule_ModuleTypeSetValue(key, datatype, dt);
64 RedisModule_CloseKey(key);
65 RedisModule_ReplyWithSimpleString(ctx, "OK");
66
67 return REDISMODULE_OK;
68 }
69
datatype_restore(RedisModuleCtx * ctx,RedisModuleString ** argv,int argc)70 static int datatype_restore(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
71 if (argc != 3) {
72 RedisModule_WrongArity(ctx);
73 return REDISMODULE_OK;
74 }
75
76 DataType *dt = RedisModule_LoadDataTypeFromString(argv[2], datatype);
77 if (!dt) {
78 RedisModule_ReplyWithError(ctx, "Invalid data");
79 return REDISMODULE_OK;
80 }
81
82 RedisModuleKey *key = RedisModule_OpenKey(ctx, argv[1], REDISMODULE_WRITE);
83 RedisModule_ModuleTypeSetValue(key, datatype, dt);
84 RedisModule_CloseKey(key);
85 RedisModule_ReplyWithSimpleString(ctx, "OK");
86
87 return REDISMODULE_OK;
88 }
89
datatype_get(RedisModuleCtx * ctx,RedisModuleString ** argv,int argc)90 static int datatype_get(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
91 if (argc != 2) {
92 RedisModule_WrongArity(ctx);
93 return REDISMODULE_OK;
94 }
95
96 RedisModuleKey *key = RedisModule_OpenKey(ctx, argv[1], REDISMODULE_READ);
97 DataType *dt = RedisModule_ModuleTypeGetValue(key);
98 RedisModule_CloseKey(key);
99
100 RedisModule_ReplyWithArray(ctx, 2);
101 RedisModule_ReplyWithLongLong(ctx, dt->intval);
102 RedisModule_ReplyWithString(ctx, dt->strval);
103 return REDISMODULE_OK;
104 }
105
datatype_dump(RedisModuleCtx * ctx,RedisModuleString ** argv,int argc)106 static int datatype_dump(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
107 if (argc != 2) {
108 RedisModule_WrongArity(ctx);
109 return REDISMODULE_OK;
110 }
111
112 RedisModuleKey *key = RedisModule_OpenKey(ctx, argv[1], REDISMODULE_READ);
113 DataType *dt = RedisModule_ModuleTypeGetValue(key);
114 RedisModule_CloseKey(key);
115
116 RedisModuleString *reply = RedisModule_SaveDataTypeToString(ctx, dt, datatype);
117 if (!reply) {
118 RedisModule_ReplyWithError(ctx, "Failed to save");
119 return REDISMODULE_OK;
120 }
121
122 RedisModule_ReplyWithString(ctx, reply);
123 RedisModule_FreeString(ctx, reply);
124 return REDISMODULE_OK;
125 }
126
datatype_swap(RedisModuleCtx * ctx,RedisModuleString ** argv,int argc)127 static int datatype_swap(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
128 if (argc != 3) {
129 RedisModule_WrongArity(ctx);
130 return REDISMODULE_OK;
131 }
132
133 RedisModuleKey *a = RedisModule_OpenKey(ctx, argv[1], REDISMODULE_WRITE);
134 RedisModuleKey *b = RedisModule_OpenKey(ctx, argv[2], REDISMODULE_WRITE);
135 void *val = RedisModule_ModuleTypeGetValue(a);
136
137 int error = (RedisModule_ModuleTypeReplaceValue(b, datatype, val, &val) == REDISMODULE_ERR ||
138 RedisModule_ModuleTypeReplaceValue(a, datatype, val, NULL) == REDISMODULE_ERR);
139 if (!error)
140 RedisModule_ReplyWithSimpleString(ctx, "OK");
141 else
142 RedisModule_ReplyWithError(ctx, "ERR failed");
143
144 RedisModule_CloseKey(a);
145 RedisModule_CloseKey(b);
146
147 return REDISMODULE_OK;
148 }
149
RedisModule_OnLoad(RedisModuleCtx * ctx,RedisModuleString ** argv,int argc)150 int RedisModule_OnLoad(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
151 REDISMODULE_NOT_USED(argv);
152 REDISMODULE_NOT_USED(argc);
153
154 if (RedisModule_Init(ctx,"datatype",1,REDISMODULE_APIVER_1) == REDISMODULE_ERR)
155 return REDISMODULE_ERR;
156
157 RedisModule_SetModuleOptions(ctx, REDISMODULE_OPTIONS_HANDLE_IO_ERRORS);
158
159 RedisModuleTypeMethods datatype_methods = {
160 .version = REDISMODULE_TYPE_METHOD_VERSION,
161 .rdb_load = datatype_load,
162 .rdb_save = datatype_save,
163 .free = datatype_free,
164 };
165
166 datatype = RedisModule_CreateDataType(ctx, "test___dt", 1, &datatype_methods);
167 if (datatype == NULL)
168 return REDISMODULE_ERR;
169
170 if (RedisModule_CreateCommand(ctx,"datatype.set", datatype_set,"deny-oom",1,1,1) == REDISMODULE_ERR)
171 return REDISMODULE_ERR;
172
173 if (RedisModule_CreateCommand(ctx,"datatype.get", datatype_get,"",1,1,1) == REDISMODULE_ERR)
174 return REDISMODULE_ERR;
175
176 if (RedisModule_CreateCommand(ctx,"datatype.restore", datatype_restore,"deny-oom",1,1,1) == REDISMODULE_ERR)
177 return REDISMODULE_ERR;
178
179 if (RedisModule_CreateCommand(ctx,"datatype.dump", datatype_dump,"",1,1,1) == REDISMODULE_ERR)
180 return REDISMODULE_ERR;
181
182 if (RedisModule_CreateCommand(ctx,"datatype.swap", datatype_swap,"",1,1,1) == REDISMODULE_ERR)
183 return REDISMODULE_ERR;
184
185 return REDISMODULE_OK;
186 }
187