Add async supoort for SEARCH commands (#2096)

* Add async supoort for SEARCH commands

* linters

* linters

* linters

* linters

* linters

Author: dvora-h

redis/asyncio/client.py

@@ -41,8 +41,8 @@
 )
 from redis.commands import (
     AsyncCoreCommands,
+    AsyncRedisModuleCommands,
     AsyncSentinelCommands,
-    RedisModuleCommands,
     list_or_args,
 )
 from redis.compat import Protocol, TypedDict
@@ -81,7 +81,7 @@ async def __call__(self, response: Any, **kwargs):
 
 
 class Redis(
-    AbstractRedis, RedisModuleCommands, AsyncCoreCommands, AsyncSentinelCommands
+    AbstractRedis, AsyncRedisModuleCommands, AsyncCoreCommands, AsyncSentinelCommands
 ):
     """
     Implementation of the Redis protocol.

redis/commands/__init__.py

@@ -2,7 +2,7 @@
 from .core import AsyncCoreCommands, CoreCommands
 from .helpers import list_or_args
 from .parser import CommandsParser
-from .redismodules import RedisModuleCommands
+from .redismodules import AsyncRedisModuleCommands, RedisModuleCommands
 from .sentinel import AsyncSentinelCommands, SentinelCommands
 
 __all__ = [
@@ -11,6 +11,7 @@
     "AsyncCoreCommands",
     "CoreCommands",
     "list_or_args",
+    "AsyncRedisModuleCommands",
     "RedisModuleCommands",
     "AsyncSentinelCommands",
     "SentinelCommands",

redis/commands/redismodules.py

@@ -81,3 +81,13 @@ def graph(self, index_name="idx"):
 
         g = Graph(client=self, name=index_name)
         return g
+
+
+class AsyncRedisModuleCommands(RedisModuleCommands):
+    def ft(self, index_name="idx"):
+        """Access the search namespace, providing support for redis search."""
+
+        from .search import AsyncSearch
+
+        s = AsyncSearch(client=self, index_name=index_name)
+        return s

redis/commands/search/__init__.py

@@ -1,6 +1,6 @@
 import redis
 
-from .commands import SearchCommands
+from .commands import AsyncSearchCommands, SearchCommands
 
 
 class Search(SearchCommands):
@@ -112,5 +112,67 @@ def pipeline(self, transaction=True, shard_hint=None):
         return p
 
 
+class AsyncSearch(Search, AsyncSearchCommands):
+    class BatchIndexer(Search.BatchIndexer):
+        """
+        A batch indexer allows you to automatically batch
+        document indexing in pipelines, flushing it every N documents.
+        """
+
+        async def add_document(
+            self,
+            doc_id,
+            nosave=False,
+            score=1.0,
+            payload=None,
+            replace=False,
+            partial=False,
+            no_create=False,
+            **fields,
+        ):
+            """
+            Add a document to the batch query
+            """
+            self.client._add_document(
+                doc_id,
+                conn=self._pipeline,
+                nosave=nosave,
+                score=score,
+                payload=payload,
+                replace=replace,
+                partial=partial,
+                no_create=no_create,
+                **fields,
+            )
+            self.current_chunk += 1
+            self.total += 1
+            if self.current_chunk >= self.chunk_size:
+                await self.commit()
+
+        async def commit(self):
+            """
+            Manually commit and flush the batch indexing query
+            """
+            await self._pipeline.execute()
+            self.current_chunk = 0
+
+    def pipeline(self, transaction=True, shard_hint=None):
+        """Creates a pipeline for the SEARCH module, that can be used for executing
+        SEARCH commands, as well as classic core commands.
+        """
+        p = AsyncPipeline(
+            connection_pool=self.client.connection_pool,
+            response_callbacks=self.MODULE_CALLBACKS,
+            transaction=transaction,
+            shard_hint=shard_hint,
+        )
+        p.index_name = self.index_name
+        return p
+
+
 class Pipeline(SearchCommands, redis.client.Pipeline):
     """Pipeline for the module."""
+
+
+class AsyncPipeline(AsyncSearchCommands, redis.asyncio.client.Pipeline):
+    """AsyncPipeline for the module."""

redis/commands/search/commands.py

@@ -857,3 +857,245 @@ def syndump(self):
         """  # noqa
         raw = self.execute_command(SYNDUMP_CMD, self.index_name)
         return {raw[i]: raw[i + 1] for i in range(0, len(raw), 2)}
+
+
+class AsyncSearchCommands(SearchCommands):
+    async def info(self):
+        """
+        Get info an stats about the the current index, including the number of
+        documents, memory consumption, etc
+
+        For more information https://oss.redis.com/redisearch/Commands/#ftinfo
+        """
+
+        res = await self.execute_command(INFO_CMD, self.index_name)
+        it = map(to_string, res)
+        return dict(zip(it, it))
+
+    async def search(
+        self,
+        query: Union[str, Query],
+        query_params: Dict[str, Union[str, int, float]] = None,
+    ):
+        """
+        Search the index for a given query, and return a result of documents
+
+        ### Parameters
+
+        - **query**: the search query. Either a text for simple queries with
+                     default parameters, or a Query object for complex queries.
+                     See RediSearch's documentation on query format
+
+        For more information: https://oss.redis.com/redisearch/Commands/#ftsearch
+        """  # noqa
+        args, query = self._mk_query_args(query, query_params=query_params)
+        st = time.time()
+        res = await self.execute_command(SEARCH_CMD, *args)
+
+        if isinstance(res, Pipeline):
+            return res
+
+        return Result(
+            res,
+            not query._no_content,
+            duration=(time.time() - st) * 1000.0,
+            has_payload=query._with_payloads,
+            with_scores=query._with_scores,
+        )
+
+    async def aggregate(
+        self,
+        query: Union[str, Query],
+        query_params: Dict[str, Union[str, int, float]] = None,
+    ):
+        """
+        Issue an aggregation query.
+
+        ### Parameters
+
+        **query**: This can be either an `AggregateRequest`, or a `Cursor`
+
+        An `AggregateResult` object is returned. You can access the rows from
+        its `rows` property, which will always yield the rows of the result.
+
+        For more information: https://oss.redis.com/redisearch/Commands/#ftaggregate
+        """  # noqa
+        if isinstance(query, AggregateRequest):
+            has_cursor = bool(query._cursor)
+            cmd = [AGGREGATE_CMD, self.index_name] + query.build_args()
+        elif isinstance(query, Cursor):
+            has_cursor = True
+            cmd = [CURSOR_CMD, "READ", self.index_name] + query.build_args()
+        else:
+            raise ValueError("Bad query", query)
+        if query_params is not None:
+            cmd += self.get_params_args(query_params)
+
+        raw = await self.execute_command(*cmd)
+        return self._get_aggregate_result(raw, query, has_cursor)
+
+    async def spellcheck(self, query, distance=None, include=None, exclude=None):
+        """
+        Issue a spellcheck query
+
+        ### Parameters
+
+        **query**: search query.
+        **distance***: the maximal Levenshtein distance for spelling
+                       suggestions (default: 1, max: 4).
+        **include**: specifies an inclusion custom dictionary.
+        **exclude**: specifies an exclusion custom dictionary.
+
+        For more information: https://oss.redis.com/redisearch/Commands/#ftspellcheck
+        """  # noqa
+        cmd = [SPELLCHECK_CMD, self.index_name, query]
+        if distance:
+            cmd.extend(["DISTANCE", distance])
+
+        if include:
+            cmd.extend(["TERMS", "INCLUDE", include])
+
+        if exclude:
+            cmd.extend(["TERMS", "EXCLUDE", exclude])
+
+        raw = await self.execute_command(*cmd)
+
+        corrections = {}
+        if raw == 0:
+            return corrections
+
+        for _correction in raw:
+            if isinstance(_correction, int) and _correction == 0:
+                continue
+
+            if len(_correction) != 3:
+                continue
+            if not _correction[2]:
+                continue
+            if not _correction[2][0]:
+                continue
+
+            corrections[_correction[1]] = [
+                {"score": _item[0], "suggestion": _item[1]} for _item in _correction[2]
+            ]
+
+        return corrections
+
+    async def config_set(self, option, value):
+        """Set runtime configuration option.
+
+        ### Parameters
+
+        - **option**: the name of the configuration option.
+        - **value**: a value for the configuration option.
+
+        For more information: https://oss.redis.com/redisearch/Commands/#ftconfig
+        """  # noqa
+        cmd = [CONFIG_CMD, "SET", option, value]
+        raw = await self.execute_command(*cmd)
+        return raw == "OK"
+
+    async def config_get(self, option):
+        """Get runtime configuration option value.
+
+        ### Parameters
+
+        - **option**: the name of the configuration option.
+
+        For more information: https://oss.redis.com/redisearch/Commands/#ftconfig
+        """  # noqa
+        cmd = [CONFIG_CMD, "GET", option]
+        res = {}
+        raw = await self.execute_command(*cmd)
+        if raw:
+            for kvs in raw:
+                res[kvs[0]] = kvs[1]
+        return res
+
+    async def load_document(self, id):
+        """
+        Load a single document by id
+        """
+        fields = await self.client.hgetall(id)
+        f2 = {to_string(k): to_string(v) for k, v in fields.items()}
+        fields = f2
+
+        try:
+            del fields["id"]
+        except KeyError:
+            pass
+
+        return Document(id=id, **fields)
+
+    async def sugadd(self, key, *suggestions, **kwargs):
+        """
+        Add suggestion terms to the AutoCompleter engine. Each suggestion has
+        a score and string.
+        If kwargs["increment"] is true and the terms are already in the
+        server's dictionary, we increment their scores.
+
+        For more information: https://oss.redis.com/redisearch/master/Commands/#ftsugadd
+        """  # noqa
+        # If Transaction is not False it will MULTI/EXEC which will error
+        pipe = self.pipeline(transaction=False)
+        for sug in suggestions:
+            args = [SUGADD_COMMAND, key, sug.string, sug.score]
+            if kwargs.get("increment"):
+                args.append("INCR")
+            if sug.payload:
+                args.append("PAYLOAD")
+                args.append(sug.payload)
+
+            pipe.execute_command(*args)
+
+        return (await pipe.execute())[-1]
+
+    async def sugget(
+        self, key, prefix, fuzzy=False, num=10, with_scores=False, with_payloads=False
+    ):
+        """
+        Get a list of suggestions from the AutoCompleter, for a given prefix.
+
+        Parameters:
+
+        prefix : str
+            The prefix we are searching. **Must be valid ascii or utf-8**
+        fuzzy : bool
+            If set to true, the prefix search is done in fuzzy mode.
+            **NOTE**: Running fuzzy searches on short (<3 letters) prefixes
+            can be very
+            slow, and even scan the entire index.
+        with_scores : bool
+            If set to true, we also return the (refactored) score of
+            each suggestion.
+            This is normally not needed, and is NOT the original score
+            inserted into the index.
+        with_payloads : bool
+            Return suggestion payloads
+        num : int
+            The maximum number of results we return. Note that we might
+            return less. The algorithm trims irrelevant suggestions.
+
+        Returns:
+
+        list:
+             A list of Suggestion objects. If with_scores was False, the
+             score of all suggestions is 1.
+
+        For more information: https://oss.redis.com/redisearch/master/Commands/#ftsugget
+        """  # noqa
+        args = [SUGGET_COMMAND, key, prefix, "MAX", num]
+        if fuzzy:
+            args.append(FUZZY)
+        if with_scores:
+            args.append(WITHSCORES)
+        if with_payloads:
+            args.append(WITHPAYLOADS)
+
+        ret = await self.execute_command(*args)
+        results = []
+        if not ret:
+            return results
+
+        parser = SuggestionParser(with_scores, with_payloads, ret)
+        return [s for s in parser]

tests/test_asyncio/conftest.py

@@ -202,7 +202,7 @@ async def wait_for_command(
         # generate key
         redis_version = REDIS_INFO["version"]
         if Version(redis_version) >= Version("5.0.0"):
-            id_str = str(client.client_id())
+            id_str = str(await client.client_id())
         else:
             id_str = f"{random.randrange(2 ** 32):08x}"
         key = f"__REDIS-PY-{id_str}__"