Add new function to API for returning the result of a script/function…

… execution Signed-off-by: Ricardo Dias <ricardo.dias@percona.com>
valkey-io · Dec 18, 2024 · 3c60980 · 3c60980
1 parent 4a06209
commit 3c60980
Show file tree

Hide file tree

Showing 6 changed files with 284 additions and 67 deletions.
diff --git a/src/functions.c b/src/functions.c
@@ -524,6 +524,77 @@ int functionsUnregisterEngine(const char *engine_name) {
     return C_OK;
 }
 
+static void functionsReturnCallDeferredResult(scriptRunCtx *run_ctx,
+                                              const callResult *result) {
+    serverAssert(result->kind == VMSE_ARRAY ||
+                 result->kind == VMSE_SET ||
+                 result->kind == VMSE_MAP);
+
+    client *c = run_ctx->original_client;
+    void *replylen = addReplyDeferredLen(c);
+    long length = 0;
+    while (result->val.deferred.returnNextElem(result->val.deferred.context)) {
+        length++;
+    }
+
+    if (result->kind == VMSE_ARRAY) {
+        setDeferredArrayLen(c, replylen, length);
+    } else if (result->kind == VMSE_SET) {
+        setDeferredSetLen(c, replylen, length);
+    } else {
+        setDeferredMapLen(c, replylen, length);
+    }
+}
+
+void functionsReturnCallResult(scriptRunCtx *run_ctx, const callResult *result) {
+    client *c = run_ctx->original_client;
+    switch (result->kind) {
+    case VMSE_ERROR:
+        addReplyError(c, result->val.buffer.ptr);
+        break;
+    case VMSE_ERROR_FORMAT:
+        addReplyErrorFormatEx(c, result->val.error.flags, result->val.error.buffer.ptr);
+        break;
+    case VMSE_STATUS:
+        addReplyStatusLength(c, result->val.buffer.ptr, result->val.buffer.len);
+        break;
+    case VMSE_BULK_STRING:
+        addReplyBulkCString(c, result->val.buffer.ptr);
+        break;
+    case VMSE_VERBATIM:
+        addReplyVerbatim(c, result->val.verb.buffer.ptr,
+                         result->val.verb.buffer.len, result->val.verb.format);
+        break;
+    case VMSE_BIGNUM:
+        addReplyBigNum(c, result->val.buffer.ptr, result->val.buffer.len);
+        break;
+    case VMSE_BOOL:
+        if (run_ctx->c->resp == 2) {
+            addReply(c, result->val.i32 ? shared.cone : shared.null[c->resp]);
+        } else {
+            addReplyBool(c, result->val.i32);
+        }
+        break;
+    case VMSE_LONG_LONG:
+        addReplyLongLong(c, result->val.i64);
+        break;
+    case VMSE_DOUBLE:
+        addReplyDouble(c, result->val.d64);
+        break;
+    case VMSE_HUMAN_LONG_DOUBLE:
+        addReplyHumanLongDouble(c, result->val.d128);
+        break;
+    case VMSE_NULL:
+        addReplyNull(c);
+        break;
+    case VMSE_ARRAY:
+    case VMSE_SET:
+    case VMSE_MAP:
+        functionsReturnCallDeferredResult(run_ctx, result);
+        break;
+    }
+}
+
 /*
  * FUNCTION STATS
  */

diff --git a/src/functions.h b/src/functions.h
@@ -58,6 +58,7 @@ typedef struct functionLibInfo functionLibInfo;
 typedef ValkeyModuleScriptingEngineCtx engineCtx;
 typedef ValkeyModuleScriptingEngineFunctionCtx functionCtx;
 typedef ValkeyModuleScriptingEngineCompiledFunction compiledFunction;
+typedef ValkeyModuleScriptingEngineCallResult callResult;
 typedef ValkeyModuleScriptingEngineMethods engineMethods;
 
 typedef struct engine {
@@ -142,6 +143,7 @@ int functionsRegisterEngine(const char *engine_name,
                             void *engine_ctx,
                             engineMethods *engine_methods);
 int functionsUnregisterEngine(const char *engine_name);
+void functionsReturnCallResult(scriptRunCtx *run_ctx, const callResult *result);
 
 sds functionsCreateWithLibraryCtx(sds code, int replace, sds *err, functionsLibCtx *lib_ctx, size_t timeout);
 unsigned long functionsMemory(void);

diff --git a/src/module.c b/src/module.c
@@ -13146,6 +13146,18 @@ int VM_UnregisterScriptingEngine(ValkeyModuleCtx *ctx, const char *engine_name)
     return VALKEYMODULE_OK;
 }
 
+/* Returns the result of calling a scripting engine function back to the client.
+ *
+ * `func_ctx` is the function call runtime context.
+ *
+ * `result` is the pointer to a structure that contains the result value.
+ *
+ */
+void VM_ReturnFunctionCallResult(ValkeyModuleScriptingEngineFunctionCtx *func_ctx,
+                                 const ValkeyModuleScriptingEngineCallResult *result) {
+    functionsReturnCallResult(func_ctx, result);
+}
+
 /* MODULE command.
  *
  * MODULE LIST
@@ -14018,4 +14030,5 @@ void moduleRegisterCoreAPI(void) {
     REGISTER_API(RdbSave);
     REGISTER_API(RegisterScriptingEngine);
     REGISTER_API(UnregisterScriptingEngine);
+    REGISTER_API(ReturnFunctionCallResult);
 }