Add support for async GRAPH module

redis/commands/graph/__init__.py

@@ -1,9 +1,13 @@
 from ..helpers import quote_string, random_string, stringify_param_value
-from .commands import GraphCommands
+from .commands import AsyncGraphCommands, GraphCommands
 from .edge import Edge  # noqa
 from .node import Node  # noqa
 from .path import Path  # noqa
 
+DB_LABELS = "DB.LABELS"
+DB_RAELATIONSHIPTYPES = "DB.RELATIONSHIPTYPES"
+DB_PROPERTYKEYS = "DB.PROPERTYKEYS"
+
 
 class Graph(GraphCommands):
     """
@@ -44,25 +48,19 @@ def _refresh_labels(self):
         lbls = self.labels()
 
         # Unpack data.
-        self._labels = [None] * len(lbls)
-        for i, l in enumerate(lbls):
-            self._labels[i] = l[0]
+        self._labels = [l[0] for _, l in enumerate(lbls)]
 
     def _refresh_relations(self):
         rels = self.relationship_types()
 
         # Unpack data.
-        self._relationship_types = [None] * len(rels)
-        for i, r in enumerate(rels):
-            self._relationship_types[i] = r[0]
+        self._relationship_types = [r[0] for _, r in enumerate(rels)]
 
     def _refresh_attributes(self):
         props = self.property_keys()
 
         # Unpack data.
-        self._properties = [None] * len(props)
-        for i, p in enumerate(props):
-            self._properties[i] = p[0]
+        self._properties = [p[0] for _, p in enumerate(props)]
 
     def get_label(self, idx):
         """
@@ -108,12 +106,12 @@ def get_property(self, idx):
             The index of the property
         """
         try:
-            propertie = self._properties[idx]
+            p = self._properties[idx]
         except IndexError:
             # Refresh properties.
             self._refresh_attributes()
-            propertie = self._properties[idx]
-        return propertie
+            p = self._properties[idx]
+        return p
 
     def add_node(self, node):
         """
@@ -133,6 +131,8 @@ def add_edge(self, edge):
         self.edges.append(edge)
 
     def _build_params_header(self, params):
+        if params is None:
+            return ""
         if not isinstance(params, dict):
             raise TypeError("'params' must be a dict")
         # Header starts with "CYPHER"
@@ -147,16 +147,109 @@ def call_procedure(self, procedure, *args, read_only=False, **kwagrs):
         q = f"CALL {procedure}({','.join(args)})"
 
         y = kwagrs.get("y", None)
-        if y:
-            q += f" YIELD {','.join(y)}"
+        if y is not None:
+            q += f"YIELD {','.join(y)}"
 
         return self.query(q, read_only=read_only)
 
     def labels(self):
-        return self.call_procedure("db.labels", read_only=True).result_set
+        return self.call_procedure(DB_LABELS, read_only=True).result_set
 
     def relationship_types(self):
-        return self.call_procedure("db.relationshipTypes", read_only=True).result_set
+        return self.call_procedure(DB_RAELATIONSHIPTYPES, read_only=True).result_set
 
     def property_keys(self):
-        return self.call_procedure("db.propertyKeys", read_only=True).result_set
+        return self.call_procedure(DB_PROPERTYKEYS, read_only=True).result_set
+
+
+class AsyncGraph(Graph, AsyncGraphCommands):
+    """Async version for Graph"""
+
+    async def _refresh_labels(self):
+        lbls = await self.labels()
+
+        # Unpack data.
+        self._labels = [l[0] for _, l in enumerate(lbls)]
+
+    async def _refresh_attributes(self):
+        props = await self.property_keys()
+
+        # Unpack data.
+        self._properties = [p[0] for _, p in enumerate(props)]
+
+    async def _refresh_relations(self):
+        rels = await self.relationship_types()
+
+        # Unpack data.
+        self._relationship_types = [r[0] for _, r in enumerate(rels)]
+
+    async def get_label(self, idx):
+        """
+        Returns a label by it's index
+
+        Args:
+
+        idx:
+            The index of the label
+        """
+        try:
+            label = self._labels[idx]
+        except IndexError:
+            # Refresh labels.
+            await self._refresh_labels()
+            label = self._labels[idx]
+        return label
+
+    async def get_property(self, idx):
+        """
+        Returns a property by it's index
+
+        Args:
+
+        idx:
+            The index of the property
+        """
+        try:
+            p = self._properties[idx]
+        except IndexError:
+            # Refresh properties.
+            await self._refresh_attributes()
+            p = self._properties[idx]
+        return p
+
+    async def get_relation(self, idx):
+        """
+        Returns a relationship type by it's index
+
+        Args:
+
+        idx:
+            The index of the relation
+        """
+        try:
+            relationship_type = self._relationship_types[idx]
+        except IndexError:
+            # Refresh relationship types.
+            await self._refresh_relations()
+            relationship_type = self._relationship_types[idx]
+        return relationship_type
+
+    async def call_procedure(self, procedure, *args, read_only=False, **kwagrs):
+        args = [quote_string(arg) for arg in args]
+        q = f"CALL {procedure}({','.join(args)})"
+
+        y = kwagrs.get("y", None)
+        if y is not None:
+            f"YIELD {','.join(y)}"
+        return await self.query(q, read_only=read_only)
+
+    async def labels(self):
+        return ((await self.call_procedure(DB_LABELS, read_only=True))).result_set
+
+    async def property_keys(self):
+        return (await self.call_procedure(DB_PROPERTYKEYS, read_only=True)).result_set
+
+    async def relationship_types(self):
+        return (
+            await self.call_procedure(DB_RAELATIONSHIPTYPES, read_only=True)
+        ).result_set

redis/commands/graph/commands.py

@@ -3,7 +3,16 @@
 
 from .exceptions import VersionMismatchException
 from .execution_plan import ExecutionPlan
-from .query_result import QueryResult
+from .query_result import AsyncQueryResult, QueryResult
+
+PROFILE_CMD = "GRAPH.PROFILE"
+RO_QUERY_CMD = "GRAPH.RO_QUERY"
+QUERY_CMD = "GRAPH.QUERY"
+DELETE_CMD = "GRAPH.DELETE"
+SLOWLOG_CMD = "GRAPH.SLOWLOG"
+CONFIG_CMD = "GRAPH.CONFIG"
+LIST_CMD = "GRAPH.LIST"
+EXPLAIN_CMD = "GRAPH.EXPLAIN"
 
 
 class GraphCommands:
@@ -52,33 +61,28 @@ def query(self, q, params=None, timeout=None, read_only=False, profile=False):
         query = q
 
         # handle query parameters
-        if params is not None:
-            query = self._build_params_header(params) + query
+        query = self._build_params_header(params) + query
 
         # construct query command
         # ask for compact result-set format
         # specify known graph version
         if profile:
-            cmd = "GRAPH.PROFILE"
+            cmd = PROFILE_CMD
         else:
-            cmd = "GRAPH.RO_QUERY" if read_only else "GRAPH.QUERY"
+            cmd = RO_QUERY_CMD if read_only else QUERY_CMD
         command = [cmd, self.name, query, "--compact"]
 
         # include timeout is specified
-        if timeout:
-            if not isinstance(timeout, int):
-                raise Exception("Timeout argument must be a positive integer")
-            command += ["timeout", timeout]
+        if isinstance(timeout, int):
+            command.extend(["timeout", timeout])
+        elif timeout is not None:
+            raise Exception("Timeout argument must be a positive integer")
 
         # issue query
         try:
             response = self.execute_command(*command)
             return QueryResult(self, response, profile)
         except ResponseError as e:
-            if "wrong number of arguments" in str(e):
-                print(
-                    "Note: RedisGraph Python requires server version 2.2.8 or above"
-                )  # noqa
             if "unknown command" in str(e) and read_only:
                 # `GRAPH.RO_QUERY` is unavailable in older versions.
                 return self.query(q, params, timeout, read_only=False)
@@ -106,7 +110,7 @@ def delete(self):
         For more information see `DELETE <https://redis.io/commands/graph.delete>`_. # noqa
         """
         self._clear_schema()
-        return self.execute_command("GRAPH.DELETE", self.name)
+        return self.execute_command(DELETE_CMD, self.name)
 
     # declared here, to override the built in redis.db.flush()
     def flush(self):
@@ -146,7 +150,7 @@ def slowlog(self):
         3. The issued query.
         4. The amount of time needed for its execution, in milliseconds.
         """
-        return self.execute_command("GRAPH.SLOWLOG", self.name)
+        return self.execute_command(SLOWLOG_CMD, self.name)
 
     def config(self, name, value=None, set=False):
         """
@@ -170,14 +174,14 @@ def config(self, name, value=None, set=False):
                 raise DataError(
                     "``value`` can be provided only when ``set`` is True"
                 )  # noqa
-        return self.execute_command("GRAPH.CONFIG", *params)
+        return self.execute_command(CONFIG_CMD, *params)
 
     def list_keys(self):
         """
         Lists all graph keys in the keyspace.
         For more information see `GRAPH.LIST <https://redis.io/commands/graph.list>`_. # noqa
         """
-        return self.execute_command("GRAPH.LIST")
+        return self.execute_command(LIST_CMD)
 
     def execution_plan(self, query, params=None):
         """
@@ -188,10 +192,9 @@ def execution_plan(self, query, params=None):
             query: the query that will be executed
             params: query parameters
         """
-        if params is not None:
-            query = self._build_params_header(params) + query
+        query = self._build_params_header(params) + query
 
-        plan = self.execute_command("GRAPH.EXPLAIN", self.name, query)
+        plan = self.execute_command(EXPLAIN_CMD, self.name, query)
         if isinstance(plan[0], bytes):
             plan = [b.decode() for b in plan]
         return "\n".join(plan)
@@ -206,8 +209,105 @@ def explain(self, query, params=None):
             query: the query that will be executed
             params: query parameters
         """
-        if params is not None:
-            query = self._build_params_header(params) + query
+        query = self._build_params_header(params) + query
+
+        plan = self.execute_command(EXPLAIN_CMD, self.name, query)
+        return ExecutionPlan(plan)
+
+
+class AsyncGraphCommands(GraphCommands):
+    async def query(self, q, params=None, timeout=None, read_only=False, profile=False):
+        """
+        Executes a query against the graph.
+        For more information see `GRAPH.QUERY <https://oss.redis.com/redisgraph/master/commands/#graphquery>`_. # noqa
+
+        Args:
+
+        q : str
+            The query.
+        params : dict
+            Query parameters.
+        timeout : int
+            Maximum runtime for read queries in milliseconds.
+        read_only : bool
+            Executes a readonly query if set to True.
+        profile : bool
+            Return details on results produced by and time
+            spent in each operation.
+        """
+
+        # maintain original 'q'
+        query = q
+
+        # handle query parameters
+        query = self._build_params_header(params) + query
+
+        # construct query command
+        # ask for compact result-set format
+        # specify known graph version
+        if profile:
+            cmd = PROFILE_CMD
+        else:
+            cmd = RO_QUERY_CMD if read_only else QUERY_CMD
+        command = [cmd, self.name, query, "--compact"]
+
+        # include timeout is specified
+        if isinstance(timeout, int):
+            command.extend(["timeout", timeout])
+        elif timeout is not None:
+            raise Exception("Timeout argument must be a positive integer")
+
+        # issue query
+        try:
+            response = await self.execute_command(*command)
+            return await AsyncQueryResult().initialize(self, response, profile)
+        except ResponseError as e:
+            if "unknown command" in str(e) and read_only:
+                # `GRAPH.RO_QUERY` is unavailable in older versions.
+                return await self.query(q, params, timeout, read_only=False)
+            raise e
+        except VersionMismatchException as e:
+            # client view over the graph schema is out of sync
+            # set client version and refresh local schema
+            self.version = e.version
+            self._refresh_schema()
+            # re-issue query
+            return await self.query(q, params, timeout, read_only)
+
+    async def execution_plan(self, query, params=None):
+        """
+        Get the execution plan for given query,
+        GRAPH.EXPLAIN returns an array of operations.
+
+        Args:
+            query: the query that will be executed
+            params: query parameters
+        """
+        query = self._build_params_header(params) + query
 
-        plan = self.execute_command("GRAPH.EXPLAIN", self.name, query)
+        plan = await self.execute_command(EXPLAIN_CMD, self.name, query)
+        if isinstance(plan[0], bytes):
+            plan = [b.decode() for b in plan]
+        return "\n".join(plan)
+
+    async def explain(self, query, params=None):
+        """
+        Get the execution plan for given query,
+        GRAPH.EXPLAIN returns ExecutionPlan object.
+
+        Args:
+            query: the query that will be executed
+            params: query parameters
+        """
+        query = self._build_params_header(params) + query
+
+        plan = await self.execute_command(EXPLAIN_CMD, self.name, query)
         return ExecutionPlan(plan)
+
+    async def flush(self):
+        """
+        Commit the graph and reset the edges and the nodes to zero length.
+        """
+        await self.commit()
+        self.nodes = {}
+        self.edges = []