Python: Add FT.SEARCH command

CHANGELOG.md

@@ -1,4 +1,5 @@
 #### Changes
+* Python: Add FT.SEARCH command([#2470](https://github.com/valkey-io/valkey-glide/pull/2470))
 * Python: Add commands FT.ALIASADD, FT.ALIASDEL, FT.ALIASUPDATE([#2471](https://github.com/valkey-io/valkey-glide/pull/2471))
 * Python: Python FT.DROPINDEX command ([#2437](https://github.com/valkey-io/valkey-glide/pull/2437))
 * Python: Python: Added FT.CREATE command([#2413](https://github.com/valkey-io/valkey-glide/pull/2413))

python/python/glide/__init__.py

@@ -49,6 +49,11 @@
     VectorFieldAttributesHnsw,
     VectorType,
 )
+from glide.async_commands.server_modules.ft_options.ft_search_options import (
+    FtSeachOptions,
+    FtSearchLimit,
+    ReturnField,
+)
 from glide.async_commands.sorted_set import (
     AggregationType,
     GeoSearchByBox,
@@ -265,4 +270,7 @@
     "VectorFieldAttributesFlat",
     "VectorFieldAttributesHnsw",
     "VectorType",
+    "FtSearchLimit",
+    "ReturnField",
+    "FtSeachOptions",
 ]

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
 
@@ -78,6 +81,42 @@ async def dropindex(client: TGlideClient, indexName: TEncodable) -> TOK:
     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
+        >>> result = await ft.search(glide_client, "idx", "*", options=FtSeachOptions(return_fields=[ReturnField(field_identifier="first"), ReturnField(field_identifier="second")]))
+        [1, { b'json:1': { b'first': b'42', b'second': b'33' } }]  # The first element, 1 is the number of keys returned in the search result. The second element is a map of data queried per key.
+    """
+    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),
+    )
+
+
 async def aliasadd(
     client: TGlideClient, alias: TEncodable, indexName: TEncodable
 ) -> TOK:

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"
     FT_ALIASADD = "FT.ALIASADD"
     FT_ALIASDEL = "FT.ALIASDEL"
     FT_ALIASUPDATE = "FT.ALIASUPDATE"
@@ -34,3 +35,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,131 @@
+# 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 FtSearchLimit:
+    """
+    This class represents the arguments for the LIMIT option of the FT.SEARCH command.
+    """
+
+    def __init__(self, offset: int, count: int):
+        """
+        Initialize a new FtSearchLimit 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[FtSearchLimit] = 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[FtSearchLimit]): This option provides pagination capability. Only the keys that satisfy the offset and count values are returned. See `FtSearchLimit`.
+            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)
+            args.append(str(len(self.params)))
+            for name, value in self.params.items():
+                args.append(name)
+                args.append(value)
+        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,154 @@
+# Copyright Valkey GLIDE Project Contributors - SPDX Identifier: Apache-2.0
+
+import json
+import time
+import uuid
+from typing import List, Mapping, Union, cast
+
+import pytest
+from glide.async_commands.server_modules import ft
+from glide.async_commands.server_modules import json as GlideJson
+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:
+    sleep_wait_time = 0.5  # This value is in seconds
+
+    @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 = prefix + 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 GlideJson.set(glide_client, json_key1, "$", json.dumps(json_value1))
+            == OK
+        )
+        assert (
+            await GlideJson.set(glide_client, json_key2, "$", json.dumps(json_value2))
+            == OK
+        )
+
+        # Wait for index to be updated to avoid this error - ResponseError: The index is under construction.
+        time.sleep(self.sleep_wait_time)
+
+        # 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._ft_search_deep_compare_result(
+            self,
+            result=result1,
+            json_key1=json_key1,
+            json_key2=json_key2,
+            json_value1=json_value1,
+            json_value2=json_value2,
+            fieldName1="a",
+            fieldName2="b",
+        )
+
+        # 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._ft_search_deep_compare_result(
+            self,
+            result=result2,
+            json_key1=json_key1,
+            json_key2=json_key2,
+            json_value1=json_value1,
+            json_value2=json_value2,
+            fieldName1="a",
+            fieldName2="b",
+        )
+
+        assert await ft.dropindex(glide_client, indexName=index) == OK
+
+    def _ft_search_deep_compare_result(
+        self,
+        result: List[Union[int, Mapping[TEncodable, Mapping[TEncodable, TEncodable]]]],
+        json_key1: str,
+        json_key2: str,
+        json_value1: dict,
+        json_value2: dict,
+        fieldName1: str,
+        fieldName2: str,
+    ):
+        """
+        Deep compare the keys and values in FT.SEARCH result array.
+
+        Args:
+            result (List[Union[int, Mapping[TEncodable, Mapping[TEncodable, TEncodable]]]]):
+            json_key1 (str): The first key in search result.
+            json_key2 (str): The second key in the search result.
+            json_value1 (dict): The fields map for first key in the search result.
+            json_value2 (dict): The fields map for second key in the search result.
+        """
+        assert len(result) == 2
+        assert result[0] == 2
+        searchResultMap: Mapping[TEncodable, Mapping[TEncodable, TEncodable]] = cast(
+            Mapping[TEncodable, Mapping[TEncodable, TEncodable]], result[1]
+        )
+        expectedResultMap: Mapping[TEncodable, Mapping[TEncodable, TEncodable]] = {
+            json_key1.encode(): {
+                fieldName1.encode(): str(json_value1.get(fieldName1)).encode(),
+                fieldName2.encode(): str(json_value1.get(fieldName2)).encode(),
+            },
+            json_key2.encode(): {
+                fieldName1.encode(): str(json_value2.get(fieldName1)).encode(),
+                fieldName2.encode(): str(json_value2.get(fieldName2)).encode(),
+            },
+        }
+        assert searchResultMap == expectedResultMap