[Compatibility] Added COMMAND GETKEYS and GETKEYSANDFLAGS command

libs/resources/RespCommandsDocs.json

@@ -1485,6 +1485,50 @@
             "ArgumentFlags": "Optional, Multiple"
           }
         ]
+      },
+      {
+        "Command": "COMMAND_GETKEYS",
+        "Name": "COMMAND|GETKEYS",
+        "Summary": "Extracts the key names from an arbitrary command.",
+        "Group": "Server",
+        "Complexity": "O(N) where N is the number of arguments to the command",
+        "Arguments": [
+          {
+            "TypeDiscriminator": "RespCommandBasicArgument",
+            "Name": "COMMAND",
+            "DisplayText": "command",
+            "Type": "String"
+          },
+          {
+            "TypeDiscriminator": "RespCommandBasicArgument",
+            "Name": "ARG",
+            "DisplayText": "arg",
+            "Type": "String",
+            "ArgumentFlags": "Optional, Multiple"
+          }
+        ]
+      },
+      {
+        "Command": "COMMAND_GETKEYSANDFLAGS",
+        "Name": "COMMAND|GETKEYSANDFLAGS",
+        "Summary": "Extracts the key names and access flags for an arbitrary command.",
+        "Group": "Server",
+        "Complexity": "O(N) where N is the number of arguments to the command",
+        "Arguments": [
+          {
+            "TypeDiscriminator": "RespCommandBasicArgument",
+            "Name": "COMMAND",
+            "DisplayText": "command",
+            "Type": "String"
+          },
+          {
+            "TypeDiscriminator": "RespCommandBasicArgument",
+            "Name": "ARG",
+            "DisplayText": "arg",
+            "Type": "String",
+            "ArgumentFlags": "Optional, Multiple"
+          }
+        ]
       }
     ]
   },

libs/resources/RespCommandsInfo.json

@@ -849,6 +849,20 @@
         "Tips": [
           "nondeterministic_output_order"
         ]
+      },
+      {
+        "Command": "COMMAND_GETKEYS",
+        "Name": "COMMAND|GETKEYS",
+        "Arity": -3,
+        "Flags": "Loading, Stale",
+        "AclCategories": "Connection, Slow"
+      },
+      {
+        "Command": "COMMAND_GETKEYSANDFLAGS",
+        "Name": "COMMAND|GETKEYSANDFLAGS",
+        "Arity": -3,
+        "Flags": "Loading, Stale",
+        "AclCategories": "Connection, Slow"
       }
     ]
   },

libs/server/Resp/BasicCommands.cs

@@ -2,6 +2,7 @@
 // Licensed under the MIT license.
 
 using System;
+using System.Collections.Generic;
 using System.Diagnostics;
 using System.Text;
 using System.Threading.Tasks;
@@ -1170,6 +1171,95 @@ private bool NetworkCOMMAND_INFO()
             return true;
         }
 
+        /// <summary>
+        /// Processes COMMAND GETKEYS subcommand.
+        /// </summary>
+        private bool NetworkCOMMAND_GETKEYS()
+        {
+            if (parseState.Count == 0)
+            {
+                return AbortWithWrongNumberOfArguments(nameof(RespCommand.COMMAND_GETKEYS));
+            }
+
+            var cmdName = parseState.GetString(0);
+            bool cmdFound = RespCommandsInfo.TryGetRespCommandInfo(cmdName, out var cmdInfo, true, true, logger) ||
+                          storeWrapper.customCommandManager.TryGetCustomCommandInfo(cmdName, out cmdInfo);
+
+            if (!cmdFound)
+            {
+                return AbortWithErrorMessage(CmdStrings.RESP_INVALID_COMMAND_SPECIFIED);
+            }
+
+            if (cmdInfo.KeySpecifications == null || cmdInfo.KeySpecifications.Length == 0)
+            {
+                return AbortWithErrorMessage(CmdStrings.RESP_COMMAND_HAS_NO_KEY_ARGS);
+            }
+
+            parseState.TryExtractKeysFromSpecs(cmdInfo.KeySpecifications, out var keys);
+
+
+            while (!RespWriteUtils.WriteArrayLength(keys.Count, ref dcurr, dend))
+                SendAndReset();
+
+            foreach (var key in keys)
+            {
+                while (!RespWriteUtils.WriteBulkString(key.Span, ref dcurr, dend))
+                    SendAndReset();
+            }
+
+            return true;
+        }
+
+        /// <summary>
+        /// Processes COMMAND GETKEYSANDFLAGS subcommand.
+        /// </summary>
+        private bool NetworkCOMMAND_GETKEYSANDFLAGS()
+        {
+            if (parseState.Count == 0)
+            {
+                return AbortWithWrongNumberOfArguments(nameof(RespCommand.COMMAND_GETKEYSANDFLAGS));
+            }
+
+            var cmdName = parseState.GetString(0);
+            bool cmdFound = RespCommandsInfo.TryGetRespCommandInfo(cmdName, out var cmdInfo, true, true, logger) ||
+                          storeWrapper.customCommandManager.TryGetCustomCommandInfo(cmdName, out cmdInfo);
+
+            if (!cmdFound)
+            {
+                return AbortWithErrorMessage(CmdStrings.RESP_INVALID_COMMAND_SPECIFIED);
+            }
+
+            if (cmdInfo.KeySpecifications == null || cmdInfo.KeySpecifications.Length == 0)
+            {
+                return AbortWithErrorMessage(CmdStrings.RESP_COMMAND_HAS_NO_KEY_ARGS);
+            }
+
+            parseState.TryExtractKeyandFlagsFromSpecs(cmdInfo.KeySpecifications, out var keys, out var flags);
+
+            while (!RespWriteUtils.WriteArrayLength(keys.Count, ref dcurr, dend))
+                SendAndReset();
+
+            for (int i = 0; i < keys.Count; i++)
+            {
+                while (!RespWriteUtils.WriteArrayLength(2, ref dcurr, dend))
+                    SendAndReset();
+
+                while (!RespWriteUtils.WriteBulkString(keys[i].Span, ref dcurr, dend))
+                    SendAndReset();
+
+                while (!RespWriteUtils.WriteArrayLength(flags[i].Length, ref dcurr, dend))
+                    SendAndReset();
+
+                foreach (var flag in flags[i])
+                {
+                    while (!RespWriteUtils.WriteBulkString(Encoding.ASCII.GetBytes(flag), ref dcurr, dend))
+                        SendAndReset();
+                }
+            }
+
+            return true;
+        }
+
         private bool NetworkECHO()
         {
             if (parseState.Count != 1)

libs/server/Resp/CmdStrings.cs

@@ -36,6 +36,8 @@ static partial class CmdStrings
         public static ReadOnlySpan<byte> info => "info"u8;
         public static ReadOnlySpan<byte> DOCS => "DOCS"u8;
         public static ReadOnlySpan<byte> docs => "docs"u8;
+        public static ReadOnlySpan<byte> GETKEYS => "GETKEYS"u8;
+        public static ReadOnlySpan<byte> GETKEYSANDFLAGS => "GETKEYSANDFLAGS"u8;
         public static ReadOnlySpan<byte> COMMAND => "COMMAND"u8;
         public static ReadOnlySpan<byte> LATENCY => "LATENCY"u8;
         public static ReadOnlySpan<byte> CLUSTER => "CLUSTER"u8;
@@ -228,6 +230,8 @@ static partial class CmdStrings
         public static ReadOnlySpan<byte> RESP_ERR_INVALID_BITFIELD_TYPE => "ERR Invalid bitfield type. Use something like i16 u8. Note that u64 is not supported but i64 is"u8;
         public static ReadOnlySpan<byte> RESP_ERR_SCRIPT_FLUSH_OPTIONS => "ERR SCRIPT FLUSH only support SYNC|ASYNC option"u8;
         public static ReadOnlySpan<byte> RESP_ERR_LENGTH_AND_INDEXES => "If you want both the length and indexes, please just use IDX."u8;
+        public static ReadOnlySpan<byte> RESP_INVALID_COMMAND_SPECIFIED => "Invalid command specified"u8;
+        public static ReadOnlySpan<byte> RESP_COMMAND_HAS_NO_KEY_ARGS => "The command has no key arguments"u8;
 
         /// <summary>
         /// Response string templates

libs/server/Resp/Parser/RespCommand.cs

@@ -270,6 +270,8 @@ public enum RespCommand : ushort
         COMMAND_COUNT,
         COMMAND_DOCS,
         COMMAND_INFO,
+        COMMAND_GETKEYS,
+        COMMAND_GETKEYSANDFLAGS,
 
         MEMORY,
         // MEMORY_USAGE is a read-only command, so moved up
@@ -379,6 +381,8 @@ public static class RespCommandExtensions
             RespCommand.COMMAND_COUNT,
             RespCommand.COMMAND_DOCS,
             RespCommand.COMMAND_INFO,
+            RespCommand.COMMAND_GETKEYS,
+            RespCommand.COMMAND_GETKEYSANDFLAGS,
             RespCommand.MEMORY_USAGE,
             // Config
             RespCommand.CONFIG_GET,
@@ -1736,6 +1740,16 @@ private RespCommand SlowParseCommand(ref int count, ref ReadOnlySpan<byte> speci
                 {
                     return RespCommand.COMMAND_DOCS;
                 }
+
+                if (subCommand.EqualsUpperCaseSpanIgnoringCase(CmdStrings.GETKEYS))
+                {
+                    return RespCommand.COMMAND_GETKEYS;
+                }
+
+                if (subCommand.EqualsUpperCaseSpanIgnoringCase(CmdStrings.GETKEYSANDFLAGS))
+                {
+                    return RespCommand.COMMAND_GETKEYSANDFLAGS;
+                }
             }
             else if (command.SequenceEqual(CmdStrings.PING))
             {

libs/server/Resp/Parser/SessionParseState.cs

@@ -5,6 +5,7 @@
 using System.Diagnostics;
 using System.Runtime.CompilerServices;
 using System.Runtime.InteropServices;
+using System.Text;
 using Garnet.common;
 using Garnet.common.Parsing;
 using Tsavorite.core;

libs/server/Resp/Parser/SessionParseStateExtension.cs

@@ -0,0 +1,89 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+
+using System.Collections.Generic;
+
+namespace Garnet.server
+{
+    /// <summary>
+    /// Extension methods for the SessionParseState struct.
+    /// </summary>
+    internal static class SessionParseStateExtension
+    {
+        /// <summary>
+        /// Tries to extract keys from the key specifications in the given RespCommandsInfo.
+        /// </summary>
+        /// <param name="state">The SessionParseState instance.</param>
+        /// <param name="keySpecs">The RespCommandKeySpecification array contains the key specification</param>
+        /// <param name="keys">The list to store extracted keys.</param>
+        /// <returns>True if keys were successfully extracted, otherwise false.</returns>
+        internal static bool TryExtractKeysFromSpecs(this ref SessionParseState state, RespCommandKeySpecification[] keySpecs, out List<ArgSlice> keys)
+        {
+            keys = new();
+
+            foreach (var spec in keySpecs)
+            {
+                if (!ExtractKeysFromSpec(ref state, keys, spec))
+                {
+                    return false;
+                }
+            }
+
+            return true;
+        }
+
+        internal static bool TryExtractKeyandFlagsFromSpecs(this ref SessionParseState state, RespCommandKeySpecification[] keySpecs, out List<ArgSlice> keys, out List<string[]> flags)
+        {
+            keys = new();
+            flags = new();
+
+            foreach (var spec in keySpecs)
+            {
+                var prevKeyCount = keys.Count;
+                if (!ExtractKeysFromSpec(ref state, keys, spec))
+                {
+                    return false;
+                }
+
+                var keyFlags = spec.RespFormatFlags;
+                for (int i = prevKeyCount; i < keys.Count; i++)
+                {
+                    flags.Add(keyFlags);
+                }
+            }
+
+            return true;
+        }
+
+        private static bool ExtractKeysFromSpec(ref SessionParseState state, List<ArgSlice> keys, RespCommandKeySpecification spec)
+        {
+            int startIndex = 0;
+
+            if (spec.BeginSearch is BeginSearchIndex bsIndex)
+            {
+                startIndex = bsIndex.GetIndex(ref state);
+            }
+            else if (spec.BeginSearch is BeginSearchKeyword bsKeyword)
+            {
+                if (!bsKeyword.TryGetStartFrom(ref state, out startIndex))
+                {
+                    return false;
+                }
+            }
+
+            if (startIndex < 0 || startIndex >= state.Count)
+                return false;
+
+            if (spec.FindKeys is FindKeysRange range)
+            {
+                range.ExtractKeys(ref state, startIndex, keys);
+            }
+            else if (spec.FindKeys is FindKeysKeyNum keyNum)
+            {
+                keyNum.ExtractKeys(ref state, startIndex, keys);
+            }
+
+            return true;
+        }
+    }
+}