From 8c0e4118f21c30455d4911fabc4bf50bfaeca712 Mon Sep 17 00:00:00 2001
From: Timo Sirainen <timo.sirainen@dovecot.fi>
Date: Thu, 10 Nov 2016 23:00:59 +0200
Subject: [PATCH] dict-client: Forward dict_iterate_set_limit() to dict server

---
 src/dict/dict-commands.c   | 12 ++++++++----
 src/lib-dict/dict-client.c | 33 +++++++++++++++++++++++++--------
 2 files changed, 33 insertions(+), 12 deletions(-)

diff --git a/src/dict/dict-commands.c b/src/dict/dict-commands.c
index 9c0cf30d9f..75c1d769f7 100644
--- a/src/dict/dict-commands.c
+++ b/src/dict/dict-commands.c
@@ -194,18 +194,22 @@ static int cmd_iterate(struct dict_connection_cmd *cmd, const char *line)
 {
 	const char *const *args;
 	unsigned int flags;
+	uint64_t max_rows;
 
 	args = t_strsplit_tabescaped(line);
-	if (str_array_length(args) < 2 ||
-	    str_to_uint(args[0], &flags) < 0) {
+	if (str_array_length(args) < 3 ||
+	    str_to_uint(args[0], &flags) < 0 ||
+	    str_to_uint64(args[1], &max_rows) < 0) {
 		i_error("dict client: ITERATE: broken input");
 		return -1;
 	}
 
-	/* <flags> <path> */
+	/* <flags> <max_rows> <path> */
 	flags |= DICT_ITERATE_FLAG_ASYNC;
-	cmd->iter = dict_iterate_init_multiple(cmd->conn->dict, args+1, flags);
+	cmd->iter = dict_iterate_init_multiple(cmd->conn->dict, args+2, flags);
 	cmd->iter_flags = flags;
+	if (max_rows > 0)
+		dict_iterate_set_limit(cmd->iter, max_rows);
 	dict_iterate_set_async_callback(cmd->iter, cmd_iterate_callback, cmd);
 	dict_connection_cmd_output_more(cmd);
 	return 1;
diff --git a/src/lib-dict/dict-client.c b/src/lib-dict/dict-client.c
index 3c75047fb5..ab26553725 100644
--- a/src/lib-dict/dict-client.c
+++ b/src/lib-dict/dict-client.c
@@ -89,12 +89,14 @@ struct client_dict_iter_result {
 struct client_dict_iterate_context {
 	struct dict_iterate_context ctx;
 	char *error;
+	const char **paths;
 	enum dict_iterate_flags flags;
 
 	pool_t results_pool;
 	ARRAY(struct client_dict_iter_result) results;
 	unsigned int result_idx;
 
+	bool cmd_sent;
 	bool seen_results;
 	bool finished;
 	bool deinit;
@@ -956,22 +958,32 @@ static struct dict_iterate_context *
 client_dict_iterate_init(struct dict *_dict, const char *const *paths,
 			 enum dict_iterate_flags flags)
 {
-	struct client_dict *dict = (struct client_dict *)_dict;
         struct client_dict_iterate_context *ctx;
-	struct client_dict_cmd *cmd;
-	string_t *query = t_str_new(256);
-	unsigned int i;
 
 	ctx = i_new(struct client_dict_iterate_context, 1);
 	ctx->ctx.dict = _dict;
 	ctx->results_pool = pool_alloconly_create("client dict iteration", 512);
 	ctx->flags = flags;
+	ctx->paths = p_strarray_dup(system_pool, paths);
 	i_array_init(&ctx->results, 64);
+	return &ctx->ctx;
+}
+
+static void
+client_dict_iterate_cmd_send(struct client_dict_iterate_context *ctx)
+{
+	struct client_dict *dict = (struct client_dict *)ctx->ctx.dict;
+	struct client_dict_cmd *cmd;
+	unsigned int i;
+	string_t *query = t_str_new(256);
 
-	str_printfa(query, "%c%d", DICT_PROTOCOL_CMD_ITERATE, flags);
-	for (i = 0; paths[i] != NULL; i++) {
+	/* we can't do this query in _iterate_init(), because
+	   _set_limit() hasn't been called yet at that point. */
+	str_printfa(query, "%c%d\t%llu", DICT_PROTOCOL_CMD_ITERATE, ctx->flags,
+		    (unsigned long long)ctx->ctx.max_rows);
+	for (i = 0; ctx->paths[i] != NULL; i++) {
 		str_append_c(query, '\t');
-			str_append(query, str_tabescape(paths[i]));
+		str_append(query, str_tabescape(ctx->paths[i]));
 	}
 
 	cmd = client_dict_cmd_init(dict, str_c(query));
@@ -980,7 +992,6 @@ client_dict_iterate_init(struct dict *_dict, const char *const *paths,
 	cmd->retry_errors = TRUE;
 
 	client_dict_cmd_send(dict, &cmd, NULL);
-	return &ctx->ctx;
 }
 
 static bool client_dict_iterate(struct dict_iterate_context *_ctx,
@@ -1005,6 +1016,11 @@ static bool client_dict_iterate(struct dict_iterate_context *_ctx,
 		ctx->seen_results = TRUE;
 		return TRUE;
 	}
+	if (!ctx->cmd_sent) {
+		ctx->cmd_sent = TRUE;
+		client_dict_iterate_cmd_send(ctx);
+		return client_dict_iterate(_ctx, key_r, value_r);
+	}
 	ctx->ctx.has_more = !ctx->finished;
 	ctx->result_idx = 0;
 	array_clear(&ctx->results);
@@ -1030,6 +1046,7 @@ static int client_dict_iterate_deinit(struct dict_iterate_context *_ctx,
 	*error_r = t_strdup(ctx->error);
 	array_free(&ctx->results);
 	pool_unref(&ctx->results_pool);
+	i_free(ctx->paths);
 	client_dict_iterate_free(ctx);
 
 	client_dict_add_timeout(dict);
-- 
2.47.3