Python: Add FT.SEARCH command

python/python/glide/async_commands/server_modules/ft.py

@@ -3,7 +3,7 @@
 module for `vector search` commands.
 """
 
-from typing import List, Optional, cast
+from typing import List, Mapping, Optional, Union, cast
 
 from glide.async_commands.server_modules.ft_options.ft_constants import (
     CommandNames,
@@ -13,6 +13,9 @@
     Field,
     FtCreateOptions,
 )
+from glide.async_commands.server_modules.ft_options.ft_search_options import (
+    FtSeachOptions,
+)
 from glide.constants import TOK, TEncodable
 from glide.glide_client import TGlideClient
 
@@ -76,3 +79,40 @@ async def dropindex(client: TGlideClient, indexName: TEncodable) -> TOK:
     """
     args: List[TEncodable] = [CommandNames.FT_DROPINDEX, indexName]
     return cast(TOK, await client.custom_command(args))
+
+
+async def search(
+    client: TGlideClient,
+    indexName: TEncodable,
+    query: TEncodable,
+    options: Optional[FtSeachOptions],
+) -> List[Union[int, Mapping[TEncodable, Mapping[TEncodable, TEncodable]]]]:
+    """
+    Uses the provided query expression to locate keys within an index.
+
+    Args:
+        client (TGlideClient): The client to execute the command.
+        indexName (TEncodable): The index name for the index to be searched.
+        query (TEncodable): The query expression to use for the search on the index.
+        options (Optional[FtSeachOptions]): Optional arguments for the FT.SEARCH command. See `FtSearchOptions`.
+
+    Returns:
+        List[Union[int, Mapping[TEncodable, Mapping[TEncodable]]]]: A list containing the search result. The first element is the total number of keys matching the query. The second element is a map of key name and field/value pair map.
+
+    Examples:
+        For the following example to work the following must already exist:
+        - An index named "idx", with fields having identifiers as "a" and "b" and prefix as "{json:}"
+        - A key named {json:}1 with value {"a": 1, "b":2}
+
+        >>> from glide.async_commands.server_modules import ft
+        >>> index = "idx"
+        >>> result = await ft.search(glide_client, index, "*", options=FtSeachOptions(return_fields=[ReturnField(field_identifier="a"),ReturnField(field_identifier="b")]))
+            [1, { b'{json:}1': {b'a': b'1', b'b' : b'2'}}]  #The first element, 1 is the number of keys returned in the search result. The second element is field/value pair map for the index.
+    """
+    args: List[TEncodable] = [CommandNames.FT_SEARCH, indexName, query]
+    if options:
+        args.extend(options.toArgs())
+    return cast(
+        List[Union[int, Mapping[TEncodable, Mapping[TEncodable, TEncodable]]]],
+        await client.custom_command(args),
+    )

python/python/glide/async_commands/server_modules/ft_options/ft_constants.py

@@ -8,6 +8,7 @@ class CommandNames:
 
     FT_CREATE = "FT.CREATE"
     FT_DROPINDEX = "FT.DROPINDEX"
+    FT_SEARCH = "FT.SEARCH"
 
 
 class FtCreateKeywords:
@@ -31,3 +32,16 @@ class FtCreateKeywords:
     M = "M"
     EF_CONSTRUCTION = "EF_CONSTRUCTION"
     EF_RUNTIME = "EF_RUNTIME"
+
+
+class FtSeachKeywords:
+    """
+    Keywords used in the FT.SEARCH command statment.
+    """
+
+    RETURN = "RETURN"
+    TIMEOUT = "TIMEOUT"
+    PARAMS = "PARAMS"
+    LIMIT = "LIMIT"
+    COUNT = "COUNT"
+    AS = "AS"

python/python/glide/async_commands/server_modules/ft_options/ft_search_options.py

@@ -0,0 +1,127 @@
+# Copyright Valkey GLIDE Project Contributors - SPDX Identifier: Apache-2.0
+
+from typing import List, Mapping, Optional
+
+from glide.async_commands.server_modules.ft_options.ft_constants import FtSeachKeywords
+from glide.constants import TEncodable
+
+
+class Limit:
+    """
+    This class represents the arguments for the LIMIT option of the FT.SEARCH command.
+    """
+
+    def __init__(self, offset: int, count: int):
+        """
+        Initialize a new Limit instance.
+
+        Args:
+            offset (int): The number of keys to skip before returning the result for the FT.SEARCH command.
+            count (int): The total number of keys to be returned by FT.SEARCH command.
+        """
+        self.offset = offset
+        self.count = count
+
+    def toArgs(self) -> List[TEncodable]:
+        """
+        Get the arguments for the LIMIT option of FT.SEARCH.
+
+        Returns:
+            List[TEncodable]: A list of LIMIT option arguments.
+        """
+        args: List[TEncodable] = [
+            FtSeachKeywords.LIMIT,
+            str(self.offset),
+            str(self.count),
+        ]
+        return args
+
+
+class ReturnField:
+    """
+    This class represents the arguments for the RETURN option of the FT.SEARCH command.
+    """
+
+    def __init__(
+        self, field_identifier: TEncodable, alias: Optional[TEncodable] = None
+    ):
+        """
+        Initialize a new ReturnField instance.
+
+        Args:
+            field_identifier (TEncodable): The identifier for the field of the key that has to returned as a result of FT.SEARCH command.
+            alias (Optional[TEncodable]): The alias to override the name of the field in the FT.SEARCH result.
+        """
+        self.field_identifier = field_identifier
+        self.alias = alias
+
+    def toArgs(self) -> List[TEncodable]:
+        """
+        Get the arguments for the RETURN option of FT.SEARCH.
+
+        Returns:
+            List[TEncodable]: A list of RETURN option arguments.
+        """
+        args: List[TEncodable] = [self.field_identifier]
+        if self.alias:
+            args.append(FtSeachKeywords.AS)
+            args.append(self.alias)
+        return args
+
+
+class FtSeachOptions:
+    """
+    This class represents the input options to be used in the FT.SEARCH command.
+    All fields in this class are optional inputs for FT.SEARCH.
+    """
+
+    def __init__(
+        self,
+        return_fields: Optional[List[ReturnField]] = None,
+        timeout: Optional[int] = None,
+        params: Optional[Mapping[TEncodable, TEncodable]] = None,
+        limit: Optional[Limit] = None,
+        count: Optional[bool] = False,
+    ):
+        """
+        Initialize the FT.SEARCH optional fields.
+
+        Args:
+            return_fields (Optional[List[ReturnField]]): The fields of a key that are returned by FT.SEARCH command. See `ReturnField`.
+            timeout (Optional[int]): This value overrides the timeout parameter of the module. The unit for the timout is in milliseconds.
+            params (Optional[Mapping[TEncodable, TEncodable]]): Param key/value pairs that can be referenced from within the query expression.
+            limit (Optional[Limit]): This option provides pagination capability. Only the keys that satisfy the offset and count values are returned. See `Limit`.
+            count (Optional[bool]): This flag option suppresses returning the contents of keys. Only the number of keys is returned.
+        """
+        self.return_fields = return_fields
+        self.timeout = timeout
+        self.params = params
+        self.limit = limit
+        self.count = count
+
+    def toArgs(self) -> List[TEncodable]:
+        """
+        Get the optional arguments for the FT.SEARCH command.
+
+        Returns:
+            List[TEncodable]:
+                List of FT.SEARCH optional agruments.
+        """
+        args: List[TEncodable] = []
+        if self.return_fields:
+            args.append(FtSeachKeywords.RETURN)
+            return_field_args: List[TEncodable] = []
+            for return_field in self.return_fields:
+                return_field_args.extend(return_field.toArgs())
+            args.append(str(len(return_field_args)))
+            args.extend(return_field_args)
+        if self.timeout:
+            args.append(FtSeachKeywords.TIMEOUT)
+            args.append(str(self.timeout))
+        if self.params:
+            args.append(FtSeachKeywords.PARAMS)
+        if self.limit:
+            args.extend(self.limit.toArgs())
+        if self.count:
+            args.append(FtSeachKeywords.COUNT)
+        return args

python/python/tests/tests_server_modules/search/test_ft_search.py

@@ -0,0 +1,161 @@
+# Copyright Valkey GLIDE Project Contributors - SPDX Identifier: Apache-2.0
+
+import json as OuterJson
+import time
+import uuid
+from typing import List, Mapping, Union, cast
+
+import pytest
+from glide.async_commands.server_modules import ft, json
+from glide.async_commands.server_modules.ft_options.ft_create_options import (
+    DataType,
+    FtCreateOptions,
+    NumericField,
+)
+from glide.async_commands.server_modules.ft_options.ft_search_options import (
+    FtSeachOptions,
+    ReturnField,
+)
+from glide.config import ProtocolVersion
+from glide.constants import OK, TEncodable
+from glide.glide_client import GlideClusterClient
+
+
+@pytest.mark.asyncio
+class TestFtSearch:
+    @pytest.mark.parametrize("cluster_mode", [True])
+    @pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3])
+    async def test_ft_search(self, glide_client: GlideClusterClient):
+        prefix = "{json-search-" + str(uuid.uuid4()) + "}:"
+        json_key1 = prefix + str(uuid.uuid4())
+        json_key2 = prefix + str(uuid.uuid4())
+        json_value1 = {"a": 11111, "b": 2, "c": 3}
+        json_value2 = {"a": 22222, "b": 2, "c": 3}
+        prefixes: List[TEncodable] = []
+        prefixes.append(prefix)
+        index = "{json-search}:" + str(uuid.uuid4())
+
+        # Create an index
+        assert (
+            await ft.create(
+                glide_client,
+                index,
+                schema=[
+                    NumericField("$.a", "a"),
+                    NumericField("$.b", "b"),
+                ],
+                options=FtCreateOptions(DataType.JSON),
+            )
+            == OK
+        )
+
+        # Create a json key
+        assert (
+            await json.set(glide_client, json_key1, "$", OuterJson.dumps(json_value1))
+            == OK
+        )
+        assert (
+            await json.set(glide_client, json_key2, "$", OuterJson.dumps(json_value2))
+            == OK
+        )
+
+        # Wait for index to be updated to avoid this error - ResponseError: The index is under construction.
+        time.sleep(0.5)
+
+        # Search the index for string inputs
+        result1 = await ft.search(
+            glide_client,
+            index,
+            "*",
+            options=FtSeachOptions(
+                return_fields=[
+                    ReturnField(field_identifier="a", alias="a_new"),
+                    ReturnField(field_identifier="b", alias="b_new"),
+                ]
+            ),
+        )
+        # Check if we get the expected result from ft.search for string inputs
+        TestFtSearch.f(
+            self,
+            result=result1,
+            json_key1=json_key1,
+            json_key2=json_key2,
+            json_value1=json_value1,
+            json_value2=json_value2,
+        )
+
+        # Search the index for byte inputs
+        result2 = await ft.search(
+            glide_client,
+            bytes(index, "utf-8"),
+            b"*",
+            options=FtSeachOptions(
+                return_fields=[
+                    ReturnField(field_identifier=b"a", alias=b"a_new"),
+                    ReturnField(field_identifier=b"b", alias=b"b_new"),
+                ]
+            ),
+        )
+
+        # Check if we get the expected result from ft.search from byte inputs
+        TestFtSearch.f(
+            self,
+            result=result2,
+            json_key1=json_key1,
+            json_key2=json_key2,
+            json_value1=json_value1,
+            json_value2=json_value2,
+        )
+
+    def f(
+        self,
+        result: List[Union[int, Mapping[TEncodable, Mapping[TEncodable, TEncodable]]]],
+        json_key1: str,
+        json_key2: str,
+        json_value1: dict,
+        json_value2: dict,
+    ):
+        type_name_bytes = "bytes"
+        assert len(result) == 2
+        assert result[0] == 2
+        searchResultMap: Mapping[TEncodable, Mapping[TEncodable, TEncodable]] = cast(
+            Mapping[TEncodable, Mapping[TEncodable, TEncodable]], result[1]
+        )
+        for key, fieldsMap in searchResultMap.items():
+            keyString = key
+            if type(key).__name__ == type_name_bytes:
+                print(type(key).__name__)
+                keyString = cast(bytes, key).decode(encoding="utf-8")
+            assert keyString == json_key1 or keyString == json_key2
+            if keyString == json_key1:
+                for fieldName, fieldValue in fieldsMap.items():
+                    fieldNameString = fieldName
+                    if type(fieldName).__name__ == type_name_bytes:
+                        fieldNameString = cast(bytes, fieldName).decode(
+                            encoding="utf-8"
+                        )
+                    fieldValueInt = int(fieldValue)
+                    if type(fieldValue).__name__ == type_name_bytes:
+                        fieldValueInt = int(
+                            cast(bytes, fieldValue).decode(encoding="utf-8")
+                        )
+                    assert fieldNameString == "a" or fieldNameString == "b"
+                    assert fieldValueInt == json_value1.get(
+                        "a"
+                    ) or fieldValueInt == json_value1.get("b")
+            if keyString == json_key2:
+                for fieldName, fieldValue in fieldsMap.items():
+                    fieldNameString = fieldName
+                    if type(fieldName).__name__ == type_name_bytes:
+                        fieldNameString = cast(bytes, fieldName).decode(
+                            encoding="utf-8"
+                        )
+                    fieldValueInt = int(fieldValue)
+                    if type(fieldValue).__name__ == type_name_bytes:
+                        fieldValueInt = int(
+                            cast(bytes, fieldValue).decode(encoding="utf-8")
+                        )
+                    assert fieldNameString == "a" or fieldNameString == "b"
+                    assert fieldValueInt == json_value2.get(
+                        "a"
+                    ) or fieldValueInt == json_value2.get("b")