Merge pull request #528 from RedisAI/Not_allowing_blocking_executions_via_multi_and_lua

Return an error when trying to run MODELRUN/SCRIPTRUN/DAGRUN through multi or lua

Author: alonre24

src/dag.c

@@ -1312,6 +1312,11 @@ void DAG_ReplyAndUnblock(RedisAI_OnFinishCtx ctx, void *private_data) {
 int RedisAI_ProcessDagRunCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc,
                                  int dagMode) {
 
+    int flags = RedisModule_GetContextFlags(ctx);
+    bool blocking_not_allowed = (flags & (REDISMODULE_CTX_FLAGS_MULTI | REDISMODULE_CTX_FLAGS_LUA));
+    if (blocking_not_allowed)
+        return RedisModule_ReplyWithError(
+            ctx, "ERR Cannot run RedisAI command within a transaction or a LUA script");
     RedisAI_RunInfo *rinfo = NULL;
     if (RAI_InitRunInfo(&rinfo) == REDISMODULE_ERR) {
         RedisModule_ReplyWithError(

tests/flow/tests_common.py

@@ -301,4 +301,24 @@ def test_info_modules(env):
     env.assertEqual( 'ai_self_used_cpu_user' in ret, True )
     env.assertEqual( 'ai_children_used_cpu_sys' in ret, True )
     env.assertEqual( 'ai_children_used_cpu_user' in ret, True )
-    env.assertEqual( 'ai_queue_CPU_bthread_#1_used_cpu_total' in ret, True )
+    env.assertEqual( 'ai_queue_CPU_bthread_#1_used_cpu_total' in ret, True )
+
+def test_lua_multi(env):
+    con = env.getConnection()
+    ret = con.execute_command('MULTI')
+    env.assertEqual(ret, b'OK')
+    ret = con.execute_command('AI.MODELRUN', "no_model{1}", "INPUTS", "no_input{1}", "OUTPUTS", "no_output{1}")
+    env.assertEqual(ret, b'QUEUED')
+    try:
+        ret = con.execute_command('EXEC')
+    except Exception as e:
+        exception = e
+        env.assertEqual(type(exception), redis.exceptions.ResponseError)
+        env.assertEqual("ERR Cannot run RedisAI command within a transaction or a LUA script", exception.__str__())
+    try:
+        ret = con.execute_command('EVAL', "return redis.pcall('AI.MODELRUN', 'no_model{1}', 'INPUTS', 'NO_INPUT{1}',"
+                                          " 'OUTPUTS', 'NO_OUTPUT{1}')", 0)
+    except Exception as e:
+        exception = e
+        env.assertEqual(type(exception), redis.exceptions.ResponseError)
+        env.assertEqual("Cannot run RedisAI command within a transaction or a LUA script", exception.__str__())