Merge pull request #3 from quackscience/scan

Add SCAN, MGET support

Author: lmangani

docs/README.md

@@ -106,6 +106,32 @@ SELECT redis_lpush('mylist', value, 'redis')
 FROM items;
 ```
 
+### Batch Operations
+```sql
+-- Get multiple keys at once
+SELECT redis_mget('key1,key2,key3', 'redis') as values;
+-- Returns comma-separated values for all keys
+
+-- Scan keys matching a pattern
+SELECT redis_scan('0', 'user:*', 10, 'redis') as result;
+-- Returns: "cursor:key1,key2,key3" where cursor is the next position for scanning
+-- Use the returned cursor for the next scan until cursor is 0
+
+-- Scan all keys matching a pattern
+WITH RECURSIVE scan(cursor, keys) AS (
+    -- Initial scan
+    SELECT split_part(redis_scan('0', 'user:*', 10, 'redis'), ':', 1),
+           split_part(redis_scan('0', 'user:*', 10, 'redis'), ':', 2)
+    UNION ALL
+    -- Continue scanning until cursor is 0
+    SELECT split_part(redis_scan(cursor, 'user:*', 10, 'redis'), ':', 1),
+           split_part(redis_scan(cursor, 'user:*', 10, 'redis'), ':', 2)
+    FROM scan
+    WHERE cursor != '0'
+)
+SELECT keys FROM scan;
+```
+
 ## Error Handling
 The extension functions will throw exceptions with descriptive error messages when:
 - Redis secret is not found or invalid

src/redis_extension.cpp

@@ -91,10 +91,14 @@ class RedisProtocol {
 
     // Key scanning
     static std::string formatScan(const std::string& cursor, const std::string& pattern = "*", int64_t count = 10) {
+        std::string cmd = "*6\r\n$4\r\nSCAN\r\n";
+        cmd += "$" + std::to_string(cursor.length()) + "\r\n" + cursor + "\r\n";
+        cmd += "$5\r\nMATCH\r\n";
+        cmd += "$" + std::to_string(pattern.length()) + "\r\n" + pattern + "\r\n";
+        cmd += "$5\r\nCOUNT\r\n";
         auto count_str = std::to_string(count);
-        return "*6\r\n$4\r\nSCAN\r\n$" + std::to_string(cursor.length()) + "\r\n" + cursor +
-               "\r\n$5\r\nMATCH\r\n$" + std::to_string(pattern.length()) + "\r\n" + pattern +
-               "\r\n$5\r\nCOUNT\r\n$" + std::to_string(count_str.length()) + "\r\n" + count_str + "\r\n";
+        cmd += "$" + std::to_string(count_str.length()) + "\r\n" + count_str + "\r\n";
+        return cmd;
     }
 
     static std::vector<std::string> parseArrayResponse(const std::string& response) {
@@ -120,6 +124,14 @@ class RedisProtocol {
         }
         return result;
     }
+
+    static std::string formatMGet(const std::vector<std::string>& keys) {
+        std::string cmd = "*" + std::to_string(keys.size() + 1) + "\r\n$4\r\nMGET\r\n";
+        for (const auto& key : keys) {
+            cmd += "$" + std::to_string(key.length()) + "\r\n" + key + "\r\n";
+        }
+        return cmd;
+    }
 };
 
 // Redis connection class
@@ -392,8 +404,96 @@ static void RedisLRangeFunction(DataChunk &args, ExpressionState &state, Vector
         });
 }
 
+static void RedisMGetFunction(DataChunk &args, ExpressionState &state, Vector &result) {
+    auto &keys_list = args.data[0];
+    auto &secret_vector = args.data[1];
+
+    UnaryExecutor::Execute<string_t, string_t>(
+        keys_list, result, args.size(),
+        [&](string_t keys_str) {
+            try {
+                // Split comma-separated keys
+                std::vector<std::string> keys;
+                std::string key_list = keys_str.GetString();
+                size_t pos = 0;
+                while ((pos = key_list.find(',')) != std::string::npos) {
+                    keys.push_back(key_list.substr(0, pos));
+                    key_list.erase(0, pos + 1);
+                }
+                if (!key_list.empty()) {
+                    keys.push_back(key_list);
+                }
+
+                string host, port, password;
+                if (!GetRedisSecret(state.GetContext(), secret_vector.GetValue(0).ToString(), 
+                                  host, port, password)) {
+                    throw InvalidInputException("Redis secret not found");
+                }
+                
+                auto conn = ConnectionPool::getInstance().getConnection(host, port, password);
+                auto response = conn->execute(RedisProtocol::formatMGet(keys));
+                auto values = RedisProtocol::parseArrayResponse(response);
+                
+                // Join results with comma
+                std::string joined;
+                for (size_t i = 0; i < values.size(); i++) {
+                    if (i > 0) joined += ",";
+                    joined += values[i];
+                }
+                return StringVector::AddString(result, joined);
+            } catch (std::exception &e) {
+                throw InvalidInputException("Redis MGET error: %s", e.what());
+            }
+        });
+}
+
+static void RedisScanFunction(DataChunk &args, ExpressionState &state, Vector &result) {
+    auto &cursor_vector = args.data[0];
+    auto &pattern_vector = args.data[1];
+    auto &count_vector = args.data[2];
+    auto &secret_vector = args.data[3];
+
+    BinaryExecutor::Execute<string_t, string_t, string_t>(
+        cursor_vector, pattern_vector, result, args.size(),
+        [&](string_t cursor, string_t pattern) {
+            try {
+                string host, port, password;
+                if (!GetRedisSecret(state.GetContext(), secret_vector.GetValue(0).ToString(), 
+                                  host, port, password)) {
+                    throw InvalidInputException("Redis secret not found");
+                }
+                
+                auto count = count_vector.GetValue(0).GetValue<int64_t>();
+                auto conn = ConnectionPool::getInstance().getConnection(host, port, password);
+                auto response = conn->execute(RedisProtocol::formatScan(
+                    cursor.GetString(), 
+                    pattern.GetString(),
+                    count
+                ));
+                auto scan_result = RedisProtocol::parseArrayResponse(response);
+                
+                if (scan_result.size() >= 2) {
+                    // First element is the new cursor, second element is array of keys
+                    std::string result_str = scan_result[0] + ":";
+                    auto keys = RedisProtocol::parseArrayResponse(scan_result[1]);
+                    for (size_t i = 0; i < keys.size(); i++) {
+                        if (i > 0) result_str += ",";
+                        result_str += keys[i];
+                    }
+                    return StringVector::AddString(result, result_str);
+                }
+                return StringVector::AddString(result, "0:");
+            } catch (std::exception &e) {
+                throw InvalidInputException("Redis SCAN error: %s", e.what());
+            }
+        });
+}
+
 static void LoadInternal(DatabaseInstance &instance) {
-    // Register Redis GET function
+    // Register the secret functions first!
+    CreateRedisSecretFunctions::Register(instance);
+
+    // Then register Redis functions
     auto redis_get_func = ScalarFunction(
         "redis_get",
         {LogicalType::VARCHAR,    // key
@@ -460,8 +560,27 @@ static void LoadInternal(DatabaseInstance &instance) {
     );
     ExtensionUtil::RegisterFunction(instance, redis_lrange_func);
 
-    // Register the secret functions
-    CreateRedisSecretFunctions::Register(instance);
+    // Register MGET
+    auto redis_mget_func = ScalarFunction(
+        "redis_mget",
+        {LogicalType::VARCHAR,    // comma-separated keys
+         LogicalType::VARCHAR},   // secret_name
+        LogicalType::VARCHAR,
+        RedisMGetFunction
+    );
+    ExtensionUtil::RegisterFunction(instance, redis_mget_func);
+
+    // Register SCAN
+    auto redis_scan_func = ScalarFunction(
+        "redis_scan",
+        {LogicalType::VARCHAR,    // cursor
+         LogicalType::VARCHAR,    // pattern
+         LogicalType::BIGINT,     // count
+         LogicalType::VARCHAR},   // secret_name
+        LogicalType::VARCHAR,
+        RedisScanFunction
+    );
+    ExtensionUtil::RegisterFunction(instance, redis_scan_func);
 }
 
 void RedisExtension::Load(DuckDB &db) {