From 3874f512008671b15e506e1e9095d3ccb0e0bc78 Mon Sep 17 00:00:00 2001
From: xumin <xumin.zhou@konghq.com>
Date: Thu, 21 Nov 2024 17:14:43 +0800
Subject: [PATCH] fix(module): allowing page size 1

---
 lib/resty/lmdb/prefix.lua | 20 ++++++++++++++---
 t/10-prefix.t             | 46 ++++++++++++++++++++++++++++++++-------
 2 files changed, 55 insertions(+), 11 deletions(-)

diff --git a/lib/resty/lmdb/prefix.lua b/lib/resty/lmdb/prefix.lua
index 989530b..a8acdab 100644
--- a/lib/resty/lmdb/prefix.lua
+++ b/lib/resty/lmdb/prefix.lua
@@ -31,10 +31,18 @@ function _M.page(start, prefix, db, page_size)
         page_size = DEFAULT_OPS_SIZE
     end
 
-    assert(page_size >= 2, "'page_size' can not be less than 2")
+    assert(page_size >= 1, "page_size must be at least 1")
+
+    local query_size = page_size
+
+    -- the function ngx_lua_resty_lmdb_ffi_prefix requires at least 2 operations
+    -- special handling for the case when page_size is less than 2
+    if query_size < 2 then
+        query_size = 2
+    end
 
     local value_buf_size = get_string_buf_size()
-    local ops = ffi_new("ngx_lua_resty_lmdb_operation_t[?]", page_size)
+    local ops = ffi_new("ngx_lua_resty_lmdb_operation_t[?]", query_size)
 
     ops[0].opcode = C.NGX_LMDB_OP_PREFIX
     ops[0].key.data = start
@@ -56,7 +64,7 @@ function _M.page(start, prefix, db, page_size)
 
 ::again::
     local buf = get_string_buf(value_buf_size, false)
-    local ret = C.ngx_lua_resty_lmdb_ffi_prefix(ops, page_size,
+    local ret = C.ngx_lua_resty_lmdb_ffi_prefix(ops, query_size,
                     buf, value_buf_size, err_ptr)
     if ret == NGX_ERROR then
         return nil, ffi_string(err_ptr[0])
@@ -74,6 +82,12 @@ function _M.page(start, prefix, db, page_size)
 
     assert(ret > 0)
 
+    -- special handling for the case when page_size is less than 2
+    if ret > page_size then
+        -- we then blindly use ret == page_size to indicate there are more keys
+        ret = page_size
+    end
+
     local res = table_new(ret, 0)
 
     for i = 1, ret do
diff --git a/t/10-prefix.t b/t/10-prefix.t
index 8ca2352..30f71ca 100644
--- a/t/10-prefix.t
+++ b/t/10-prefix.t
@@ -328,7 +328,7 @@ key: test3 value: value3
 
 
 
-=== TEST 8: prefix.page() operation with invalid page size
+=== TEST 8: prefix.page() operation with page size 1
 --- http_config eval: $::HttpConfig
 --- main_config eval: $::MainConfig
 --- config
@@ -336,19 +336,49 @@ key: test3 value: value3
         content_by_lua_block {
             local l = require("resty.lmdb")
 
-            ngx.say(l.db_drop(true))
-            local p = require("resty.lmdb.prefix")
+            l.db_drop(true)
+            l.set("test", "value")
+            l.set("test1", "value1")
+            l.set("test2", "value2")
+            l.set("test3", "value3")
+            l.set("u", "value4")
+            l.set("u1", "value5")
 
-            ngx.say(l.set("test", "value"))
-            ngx.say(pcall(p.page, "test", "test", nil, 1))
+            local p = require("resty.lmdb.prefix")    
+
+            ngx.say("PAGE SIZE 1")
+            local res, more = assert(p.page("test", "test", nil, 1))
+
+            assert(#res == 1)
+            ngx.say(res[1].key, " ", res[1].value)
+            assert(more)
+
+            ngx.say("PAGE SIZE 2")
+            res, more = p.page("test", "test", nil, 2)
+            assert(more)
+
+            assert(#res == 2)
+            for _, pair in ipairs(res) do
+                ngx.say(pair.key, " ", pair.value)
+            end
+
+            ngx.say("PAGE SIZE 0")
+            local ok, err= pcall(p.page, "test", "test", nil, 0)
+            if not ok then
+                ngx.say(err)
+            end
         }
     }
 --- request
 GET /t
 --- response_body
-true
-true
-false.../lua-resty-lmdb/lua-resty-lmdb/lib/resty/lmdb/prefix.lua:34: 'page_size' can not be less than 2
+PAGE SIZE 1
+test value
+PAGE SIZE 2
+test value
+test1 value1
+PAGE SIZE 0
+page_size must be at least 1
 --- no_error_log
 [error]
 [warn]