Allow MIGRATE to accept multiple keys

phpredis · Apr 23, 2017 · 9aa3dbf · 9aa3dbf
1 parent b087223
commit 9aa3dbf
Show file tree

Hide file tree

Showing 4 changed files with 93 additions and 58 deletions.
diff --git a/library.h b/library.h
@@ -4,6 +4,9 @@
 #define REDIS_CMD_APPEND_SSTR_STATIC(sstr, str) \
     redis_cmd_append_sstr(sstr, str, sizeof(str)-1);
 
+#define REDIS_CMD_APPEND_SSTR_OPT_STATIC(sstr, opt, str) \
+    if (opt) REDIS_CMD_APPEND_SSTR_STATIC(sstr, str);
+
 #define REDIS_CMD_INIT_SSTR_STATIC(sstr, argc, keyword) \
     redis_cmd_init_sstr(sstr, argc, keyword, sizeof(keyword)-1);
 

diff --git a/redis.c b/redis.c
@@ -3250,64 +3250,7 @@ PHP_METHOD(Redis, debug) {
 /* {{{ proto Redis::migrate(host port key dest-db timeout [bool copy,
  *                          bool replace]) */
 PHP_METHOD(Redis, migrate) {
-    zval *object;
-    RedisSock *redis_sock;
-    char *cmd, *host, *key;
-    int cmd_len, key_free;
-    strlen_t host_len, key_len;
-    zend_bool copy=0, replace=0;
-    zend_long port, dest_db, timeout;
-
-    // Parse arguments
-    if(zend_parse_method_parameters(ZEND_NUM_ARGS() TSRMLS_CC, getThis(),
-                                    "Oslsll|bb", &object, redis_ce, &host,
-                                    &host_len, &port, &key, &key_len, &dest_db,
-                                    &timeout, &copy, &replace) == FAILURE)
-    {
-        RETURN_FALSE;
-    }
-
-	/* Grabg our socket */
-    if ((redis_sock = redis_sock_get(object TSRMLS_CC, 0)) == NULL) {
-		RETURN_FALSE;
-	}
-
-    // Prefix our key if we need to, build our command
-    key_free = redis_key_prefix(redis_sock, &key, &key_len);
-
-    /* Construct our command */
-    if(copy && replace) {
-        cmd_len = redis_cmd_format_static(&cmd, "MIGRATE", "sdsddss",
-                                          host, host_len, port, key, key_len,
-                                          dest_db, timeout, "COPY",
-                                          sizeof("COPY")-1, "REPLACE",
-                                          sizeof("REPLACE")-1);
-    } else if(copy) {
-        cmd_len = redis_cmd_format_static(&cmd, "MIGRATE", "sdsdds",
-                                          host, host_len, port, key, key_len,
-                                          dest_db, timeout, "COPY",
-                                          sizeof("COPY")-1);
-    } else if(replace) {
-        cmd_len = redis_cmd_format_static(&cmd, "MIGRATE", "sdsdds",
-                                          host, host_len, port, key, key_len,
-                                          dest_db, timeout, "REPLACE",
-                                          sizeof("REPLACE")-1);
-    } else {
-        cmd_len = redis_cmd_format_static(&cmd, "MIGRATE", "sdsdd",
-                                          host, host_len, port, key, key_len,
-                                          dest_db, timeout);
-    }
-
-    /* Free our key if we prefixed it */
-	if(key_free) efree(key);
-
-    // Kick off our MIGRATE request
-    REDIS_PROCESS_REQUEST(redis_sock, cmd, cmd_len);
-    IF_ATOMIC() {
-        redis_boolean_response(INTERNAL_FUNCTION_PARAM_PASSTHRU, redis_sock,
-            NULL, NULL);
-    }
-    REDIS_PROCESS_RESPONSE(redis_boolean_response);
+    REDIS_PROCESS_CMD(migrate, redis_boolean_response);
 }
 
 /* {{{ proto Redis::_prefix(key) */

diff --git a/redis_commands.c b/redis_commands.c
@@ -2946,6 +2946,92 @@ int redis_georadiusbymember_cmd(INTERNAL_FUNCTION_PARAMETERS, RedisSock *redis_s
     return SUCCESS;
 }
 
+/* MIGRATE */
+int redis_migrate_cmd(INTERNAL_FUNCTION_PARAMETERS, RedisSock *redis_sock,
+                      char **cmd, int *cmd_len, short *slot, void **ctx)
+{
+    smart_string cmdstr = {0};
+    char *host, *key;
+    int argc, keyfree;
+    zval *z_keys, *z_key;
+    strlen_t hostlen, keylen;
+    zend_long destdb, port, timeout;
+    zend_bool copy = 0, replace = 0;
+    zend_string *zstr;
+
+    if (zend_parse_parameters(ZEND_NUM_ARGS() TSRMLS_CC, "slzll|bb", &host, &hostlen, &port,
+                              &z_keys, &destdb, &timeout, &copy, &replace) == FAILURE)
+    {
+        return FAILURE;
+    }
+
+    /* Protect against being passed an array with zero elements */
+    if (Z_TYPE_P(z_keys) == IS_ARRAY && zend_hash_num_elements(Z_ARRVAL_P(z_keys)) == 0) {
+        php_error_docref(NULL TSRMLS_CC, E_WARNING, "Keys array cannot be empty");
+        return FAILURE;
+    }
+
+    /* host, port, key|"", dest-db, timeout, [copy, replace] [KEYS key1..keyN] */
+    argc = 5 + copy + replace;
+    if (Z_TYPE_P(z_keys) == IS_ARRAY) {
+        /* +1 for the "KEYS" argument itself */
+        argc += 1 + zend_hash_num_elements(Z_ARRVAL_P(z_keys));
+    }
+
+    /* Initialize MIGRATE command with host and port */
+    REDIS_CMD_INIT_SSTR_STATIC(&cmdstr, argc, "MIGRATE");
+    redis_cmd_append_sstr(&cmdstr, host, hostlen);
+    redis_cmd_append_sstr_long(&cmdstr, port);
+
+    /* If passed a keys array the keys come later, otherwise pass the key to
+     * migrate here */
+    if (Z_TYPE_P(z_keys) == IS_ARRAY) {
+        REDIS_CMD_APPEND_SSTR_STATIC(&cmdstr, "");
+    } else {
+        /* Grab passed value as a string */
+        zstr = zval_get_string(z_keys);
+
+        /* We may need to prefix our string */
+        key = zstr->val;
+        keylen = zstr->len;
+        keyfree = redis_key_prefix(redis_sock, &key, &keylen);
+
+        /* Add key to migrate */
+        redis_cmd_append_sstr(&cmdstr, key, keylen);
+
+        zend_string_release(zstr);
+        if (keyfree) efree(key);
+    }
+
+    redis_cmd_append_sstr_long(&cmdstr, destdb);
+    redis_cmd_append_sstr_long(&cmdstr, timeout);
+    REDIS_CMD_APPEND_SSTR_OPT_STATIC(&cmdstr, copy, "COPY");
+    REDIS_CMD_APPEND_SSTR_OPT_STATIC(&cmdstr, replace, "REPLACE");
+
+    /* Append actual keys if we've got a keys array */
+    if (Z_TYPE_P(z_keys) == IS_ARRAY) {
+        REDIS_CMD_APPEND_SSTR_STATIC(&cmdstr, "KEYS");
+
+        ZEND_HASH_FOREACH_VAL(Z_ARRVAL_P(z_keys), z_key) {
+            zstr = zval_get_string(z_key);
+
+            key = zstr->val;
+            keylen = zstr->len;
+            keyfree = redis_key_prefix(redis_sock, &key, &keylen);
+
+            /* Append the key */
+            redis_cmd_append_sstr(&cmdstr, zstr->val, zstr->len);
+
+            zend_string_release(zstr);
+            if (keyfree) efree(key);
+        } ZEND_HASH_FOREACH_END();
+    }
+
+    *cmd = cmdstr.c;
+    *cmd_len = cmdstr.len;
+    return SUCCESS;
+}
+
 /* DEL */
 int redis_del_cmd(INTERNAL_FUNCTION_PARAMETERS, RedisSock *redis_sock,
                   char **cmd, int *cmd_len, short *slot, void **ctx)

diff --git a/redis_commands.h b/redis_commands.h
@@ -251,6 +251,9 @@ int redis_georadius_cmd(INTERNAL_FUNCTION_PARAMETERS, RedisSock *redis_sock,
 int redis_georadiusbymember_cmd(INTERNAL_FUNCTION_PARAMETERS, RedisSock *redis_sock,
     char **cmd, int *cmd_len, short *slot, void **ctx);
 
+int redis_migrate_cmd(INTERNAL_FUNCTION_PARAMETERS, RedisSock *redis_sock,
+    char **cmd, int *cmd_len, short *slot, void **ctx);
+
 /* Commands that don't communicate with Redis at all (such as getOption, 
  * setOption, _prefix, _serialize, etc).  These can be handled in one place
  * with the method of grabbing our RedisSock* object in different ways