diff --git a/CHANGELOG.md b/CHANGELOG.md index fc89b888bd..19303101ad 100644 --- a/CHANGELOG.md +++ b/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)) diff --git a/python/python/glide/__init__.py b/python/python/glide/__init__.py index 05910eb480..5bb31c75fb 100644 --- a/python/python/glide/__init__.py +++ b/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", ] diff --git a/python/python/glide/async_commands/server_modules/ft.py b/python/python/glide/async_commands/server_modules/ft.py index ccd1fd8735..82118e9070 100644 --- a/python/python/glide/async_commands/server_modules/ft.py +++ b/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. Once located, the count and/or the content of indexed fields within those keys can be returned. + + Args: + client (TGlideClient): The client to execute the command. + indexName (TEncodable): The index name to search into. + query (TEncodable): The text query to search. + options (Optional[FtSeachOptions]): The search options. See `FtSearchOptions`. + + Returns: + List[Union[int, Mapping[TEncodable, Mapping[TEncodable, TEncodable]]]]: A two element array, where first element is count of documents in result set, and the second element, which has the format Mapping[TEncodable, Mapping[TEncodable, TEncodable]] is a mapping between document names and map of their attributes. + If count(option in `FtSearchOptions`) is set to true or limit(option in `FtSearchOptions`) is set to FtSearchLimit(0, 0), the command returns array with only one element - the count of the documents. + 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: diff --git a/python/python/glide/async_commands/server_modules/ft_options/ft_constants.py b/python/python/glide/async_commands/server_modules/ft_options/ft_constants.py index 14fef2a681..541b286d83 100644 --- a/python/python/glide/async_commands/server_modules/ft_options/ft_constants.py +++ b/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" @@ -15,7 +16,7 @@ class CommandNames: class FtCreateKeywords: """ - Keywords used in the FT.CREATE command statment. + Keywords used in the FT.CREATE command. """ SCHEMA = "SCHEMA" @@ -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. + """ + + RETURN = "RETURN" + TIMEOUT = "TIMEOUT" + PARAMS = "PARAMS" + LIMIT = "LIMIT" + COUNT = "COUNT" + AS = "AS" diff --git a/python/python/glide/async_commands/server_modules/ft_options/ft_search_options.py b/python/python/glide/async_commands/server_modules/ft_options/ft_search_options.py new file mode 100644 index 0000000000..79f5422edc --- /dev/null +++ b/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 diff --git a/python/python/tests/tests_server_modules/search/test_ft_search.py b/python/python/tests/tests_server_modules/search/test_ft_search.py new file mode 100644 index 0000000000..80d8319676 --- /dev/null +++ b/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