From 94fe1fce992fc0541c60578f4df78e571077a7a7 Mon Sep 17 00:00:00 2001
From: Adam Holmberg <Adam.Holmberg@datastax.com>
Date: Tue, 23 Feb 2016 13:30:35 -0600
Subject: [PATCH] ResultSet.was_applied property for LWT results.

---
 cassandra/cluster.py         | 25 +++++++++++++++++++++++--
 cassandra/cqlengine/query.py | 10 ++++++----
 tests/unit/test_resultset.py | 27 +++++++++++++++++++++++++++
 3 files changed, 56 insertions(+), 6 deletions(-)

diff --git a/cassandra/cluster.py b/cassandra/cluster.py
index 1954a139..13afecd2 100644
--- a/cassandra/cluster.py
+++ b/cassandra/cluster.py
@@ -70,8 +70,8 @@ from cassandra.pool import (Host, _ReconnectionHandler, _HostReconnectionHandler
                             HostConnectionPool, HostConnection,
                             NoConnectionsAvailable)
 from cassandra.query import (SimpleStatement, PreparedStatement, BoundStatement,
-                             BatchStatement, bind_params, QueryTrace, Statement,
-                             named_tuple_factory, dict_factory, FETCH_SIZE_UNSET)
+                             BatchStatement, bind_params, QueryTrace,
+                             named_tuple_factory, dict_factory, tuple_factory, FETCH_SIZE_UNSET)
 
 
 def _is_eventlet_monkey_patched():
@@ -3409,3 +3409,24 @@ class ResultSet(object):
         See :meth:`.ResponseFuture.get_all_query_traces` for details.
         """
         return self.response_future.get_all_query_traces(max_wait_sec_per)
+
+    @property
+    def was_applied(self):
+        """
+        For LWT results, returns whether the transaction was applied.
+
+        Result is indeterminate if called on a result that was not an LWT request.
+
+        Only valid when one of tne of the internal row factories is in use.
+        """
+        if self.response_future.row_factory not in (named_tuple_factory, dict_factory, tuple_factory):
+            raise RuntimeError("Cannot determine LWT result with row factory %s" % (self.response_future.row_factsory,))
+        if len(self.current_rows) != 1:
+            raise RuntimeError("LWT result should have exactly one row. This has %d." % (len(self.current_rows)))
+
+        row = self.current_rows[0]
+        if isinstance(row, tuple):
+            return row[0]
+        else:
+            return row['[applied]']
+
diff --git a/cassandra/cqlengine/query.py b/cassandra/cqlengine/query.py
index dc64a005..d356bf5e 100644
--- a/cassandra/cqlengine/query.py
+++ b/cassandra/cqlengine/query.py
@@ -63,11 +63,13 @@ class MultipleObjectsReturned(QueryException):
 
 def check_applied(result):
     """
-    check if result contains some column '[applied]' with false value,
-    if that value is false, it means our light-weight transaction didn't
-    applied to database.
+    Raises LWTException if it looks like a failed LWT request.
     """
-    if result and '[applied]' in result[0] and not result[0]['[applied]']:
+    try:
+        applied = result.was_applied
+    except Exception:
+        applied = True  # result was not LWT form
+    if not applied:
         raise LWTException(result[0])
 
 
diff --git a/tests/unit/test_resultset.py b/tests/unit/test_resultset.py
index 252b6e05..2deeb30f 100644
--- a/tests/unit/test_resultset.py
+++ b/tests/unit/test_resultset.py
@@ -11,6 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from cassandra.query import named_tuple_factory, dict_factory, tuple_factory
 
 try:
     import unittest2 as unittest
@@ -161,3 +162,29 @@ class ResultSetTests(unittest.TestCase):
     def test_bool(self):
         self.assertFalse(ResultSet(Mock(has_more_pages=False), []))
         self.assertTrue(ResultSet(Mock(has_more_pages=False), [1]))
+
+    def test_was_applied(self):
+        # unknown row factory raises
+        with self.assertRaises(RuntimeError):
+            ResultSet(Mock(), []).was_applied
+
+        response_future = Mock(row_factory=named_tuple_factory)
+
+        # no row
+        with self.assertRaises(RuntimeError):
+            ResultSet(response_future, []).was_applied
+
+        # too many rows
+        with self.assertRaises(RuntimeError):
+            ResultSet(response_future, [tuple(), tuple()]).was_applied
+
+        # various internal row factories
+        for row_factory in (named_tuple_factory, tuple_factory):
+            for applied in (True, False):
+                rs = ResultSet(Mock(row_factory=row_factory), [(applied,)])
+                self.assertEqual(rs.was_applied, applied)
+
+        row_factory = dict_factory
+        for applied in (True, False):
+            rs = ResultSet(Mock(row_factory=row_factory), [{'[applied]': applied}])
+            self.assertEqual(rs.was_applied, applied)