PYTHON-1129: Ability to execute Fluent Graph queries asynchronously (scylladb#24)

Ability to execute Fluent Graph queries asynchronously

Author: aboudreault

CHANGELOG.rst

@@ -7,7 +7,8 @@ Features
 * Allow passing ssl context for Twisted (PYTHON-1161)
 * Cloud Twisted support (PYTHON-1163)
 * Add additional_write_policy and read_repair to system schema parsing (PYTHON-1048)
-* Remove *read_repair_chance table options (PYTHON-1140)
+* Remove *read_repair_chance* table options (PYTHON-1140)
+* [GRAPH] Ability to execute Fluent Graph queries asynchronously (PYTHON-1129)
 
 Bug Fixes
 ---------

cassandra/datastax/graph/fluent/__init__.py

@@ -15,6 +15,8 @@
 import logging
 import copy
 
+from concurrent.futures import Future
+
 HAVE_GREMLIN = False
 try:
     import gremlin_python
@@ -117,22 +119,52 @@ def __init__(self, session, graph_name=None, execution_profile=EXEC_PROFILE_GRAP
             self.graph_name = graph_name
             self.execution_profile = execution_profile
 
-        def submit(self, bytecode):
-
+        def _prepare_query(self, bytecode):
             query = DseGraph.query_from_traversal(bytecode)
             ep = self.session.execution_profile_clone_update(self.execution_profile,
                                                              row_factory=graph_traversal_row_factory)
             graph_options = ep.graph_options.copy()
             graph_options.graph_language = DseGraph.DSE_GRAPH_QUERY_LANGUAGE
             if self.graph_name:
                 graph_options.graph_name = self.graph_name
-
             ep.graph_options = graph_options
 
+            return query, ep
+
+        @staticmethod
+        def _handle_query_results(result_set, gremlin_future):
+            try:
+                traversers = [Traverser(t) for t in result_set]
+                gremlin_future.set_result(
+                    RemoteTraversal(iter(traversers), TraversalSideEffects())
+                )
+            except Exception as e:
+                gremlin_future.set_exception(e)
+
+        @staticmethod
+        def _handle_query_error(response, gremlin_future):
+            gremlin_future.set_exception(response)
+
+        def submit(self, bytecode):
+            # the only reason I don't use submitAsync here
+            # is to avoid an unuseful future wrap
+            query, ep = self._prepare_query(bytecode)
+
             traversers = self.session.execute_graph(query, execution_profile=ep)
             traversers = [Traverser(t) for t in traversers]
             return RemoteTraversal(iter(traversers), TraversalSideEffects())
 
+        def submitAsync(self, bytecode):
+            query, ep = self._prepare_query(bytecode)
+
+            # to be compatible with gremlinpython, we need to return a concurrent.futures.Future
+            gremlin_future = Future()
+            response_future = self.session.execute_graph_async(query, execution_profile=ep)
+            response_future.add_callback(self._handle_query_results, gremlin_future)
+            response_future.add_errback(self._handle_query_error, gremlin_future)
+
+            return gremlin_future
+
         def __str__(self):
             return "<DSESessionRemoteGraphConnection: graph_name='{0}'>".format(self.graph_name)
 

docs/graph_fluent.rst

@@ -109,16 +109,45 @@ that session. We call this *implicit execution* because the `Session` is not
 explicitly involved. Everything is managed internally by TinkerPop while
 traversing the graph and the results are TinkerPop types as well.
 
-For example:
+Synchronous Example
+-------------------
 
 .. code-block:: python
 
     # Build the GraphTraversalSource
     g = DseGraph.traversal_source(session)
     # implicitly execute the query by traversing the TraversalSource
     g.addV('genre').property('genreId', 1).property('name', 'Action').next()
-    # view the results of the execution
-    pprint(g.V().toList())
+
+    # blocks until the query is completed and return the results
+    results = g.V().toList()
+    pprint(results)
+
+Asynchronous Exemple
+--------------------
+
+You can execute a graph traversal query asynchronously by using `.promise()`. It returns a
+python `Future <https://docs.python.org/3/library/concurrent.futures.html#concurrent.futures.Future>`_.
+
+.. code-block:: python
+
+    # Build the GraphTraversalSource
+    g = DseGraph.traversal_source(session)
+    # implicitly execute the query by traversing the TraversalSource
+    g.addV('genre').property('genreId', 1).property('name', 'Action').next()  # not async
+
+    # get a future and wait
+    future = g.V().promise()
+    results = list(future.result())
+    pprint(results)
+
+    # or set a callback
+    def cb(f):
+        results = list(f.result())
+        pprint(results)
+    future = g.V().promise()
+    future.add_done_callback(cb)
+    # do other stuff...
 
 Specify the Execution Profile explicitly
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

tests/integration/advanced/graph/fluent/test_graph.py

@@ -14,6 +14,9 @@
 
 
 import sys
+
+from concurrent.futures import Future
+
 from cassandra.datastax.graph.fluent import DseGraph
 from gremlin_python.process.graph_traversal import GraphTraversal, GraphTraversalSource
 from gremlin_python.process.traversal import P
@@ -35,6 +38,7 @@
 if six.PY3:
     import ipaddress
 
+
 def setup_module():
     if DSE_VERSION:
         dse_options = {'graph': {'realtime_evaluation_timeout_in_seconds': 60}}
@@ -398,13 +402,13 @@ def _check_equality(self, g, original, read_value):
 
 
 @requiredse
-class ImplicitExecutionTest(AbstractTraversalTest, BasicGraphUnitTestCase):
+class BaseImplicitExecutionTest(BasicGraphUnitTestCase):
     """
     This test class will execute all tests of the AbstractTraversalTestClass using implicit execution
     This all traversal will be run directly using toList()
     """
     def setUp(self):
-        super(ImplicitExecutionTest, self).setUp()
+        super(BaseImplicitExecutionTest, self).setUp()
         if DSE_VERSION:
             self.ep = DseGraph().create_execution_profile(self.graph_name)
             self.cluster.add_execution_profile(self.graph_name, self.ep)
@@ -487,6 +491,70 @@ def _validate_type(self, g,  vertex):
             _validate_prop(key, value, self)
 
 
+@requiredse
+class ImplicitExecutionTest(BaseImplicitExecutionTest, AbstractTraversalTest):
+    pass
+
+
+@requiredse
+class ImplicitAsyncExecutionTest(BaseImplicitExecutionTest):
+    """
+    Test to validate that the traversal async execution works properly.
+
+    @since 3.21.0
+    @jira_ticket PYTHON-1129
+
+    @test_category dse graph
+    """
+
+    def _validate_results(self, results):
+        results = list(results)
+        self.assertEqual(len(results), 2)
+        self.assertIn('vadas', results)
+        self.assertIn('josh', results)
+
+    def test_promise(self):
+        generate_classic(self.session)
+        g = self.fetch_traversal_source()
+        traversal_future = g.V().has('name', 'marko').out('knows').values('name').promise()
+        self._validate_results(traversal_future.result())
+
+    def test_promise_error_is_propagated(self):
+        generate_classic(self.session)
+        g = DseGraph().traversal_source(self.session, 'wrong_graph', execution_profile=self.ep)
+        traversal_future = g.V().has('name', 'marko').out('knows').values('name').promise()
+        with self.assertRaises(Exception):
+            traversal_future.result()
+
+    def test_promise_callback(self):
+        generate_classic(self.session)
+        g = self.fetch_traversal_source()
+        future = Future()
+
+        def cb(f):
+            future.set_result(f.result())
+
+        traversal_future = g.V().has('name', 'marko').out('knows').values('name').promise()
+        traversal_future.add_done_callback(cb)
+        self._validate_results(future.result())
+
+    def test_promise_callback_on_error(self):
+        generate_classic(self.session)
+        g = DseGraph().traversal_source(self.session, 'wrong_graph', execution_profile=self.ep)
+        future = Future()
+
+        def cb(f):
+            try:
+                f.result()
+            except Exception as e:
+                future.set_exception(e)
+
+        traversal_future = g.V().has('name', 'marko').out('knows').values('name').promise()
+        traversal_future.add_done_callback(cb)
+        with self.assertRaises(Exception):
+            future.result()
+
+
 @requiredse
 class ExplicitExecutionBase(BasicGraphUnitTestCase):
     def setUp(self):