Merge pull request datastax#939 from datastax/python-941

PYTHON-941 Improve error for batch WriteTimeouts

Author: beltran

CHANGELOG.rst

@@ -17,6 +17,7 @@ Other
 * Warn users when using the deprecated Session.default_consistency_level (PYTHON-953)
 * Add DSE smoke test to OSS driver tests (PYTHON-894)
 * Document long compilation times and workarounds (PYTHON-868)
+* Improve error for batch WriteTimeouts (PYTHON-941)
 
 3.13.0
 ======

cassandra/__init__.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import logging
+import six
 
 
 class NullHandler(logging.Handler):
@@ -213,6 +214,75 @@ def uses_keyspace_flag(cls, version):
         return version >= cls.V5
 
 
+class WriteType(object):
+    """
+    For usage with :class:`.RetryPolicy`, this describe a type
+    of write operation.
+    """
+
+    SIMPLE = 0
+    """
+    A write to a single partition key. Such writes are guaranteed to be atomic
+    and isolated.
+    """
+
+    BATCH = 1
+    """
+    A write to multiple partition keys that used the distributed batch log to
+    ensure atomicity.
+    """
+
+    UNLOGGED_BATCH = 2
+    """
+    A write to multiple partition keys that did not use the distributed batch
+    log. Atomicity for such writes is not guaranteed.
+    """
+
+    COUNTER = 3
+    """
+    A counter write (for one or multiple partition keys). Such writes should
+    not be replayed in order to avoid overcount.
+    """
+
+    BATCH_LOG = 4
+    """
+    The initial write to the distributed batch log that Cassandra performs
+    internally before a BATCH write.
+    """
+
+    CAS = 5
+    """
+    A lighweight-transaction write, such as "DELETE ... IF EXISTS".
+    """
+
+    VIEW = 6
+    """
+    This WriteType is only seen in results for requests that were unable to
+    complete MV operations.
+    """
+
+    CDC = 7
+    """
+    This WriteType is only seen in results for requests that were unable to
+    complete CDC operations.
+    """
+
+
+WriteType.name_to_value = {
+    'SIMPLE': WriteType.SIMPLE,
+    'BATCH': WriteType.BATCH,
+    'UNLOGGED_BATCH': WriteType.UNLOGGED_BATCH,
+    'COUNTER': WriteType.COUNTER,
+    'BATCH_LOG': WriteType.BATCH_LOG,
+    'CAS': WriteType.CAS,
+    'VIEW': WriteType.VIEW,
+    'CDC': WriteType.CDC
+}
+
+
+WriteType.value_to_name = {v: k for k, v in six.iteritems(WriteType.name_to_value)}
+
+
 class SchemaChangeType(object):
     DROPPED = 'DROPPED'
     CREATED = 'CREATED'
@@ -339,14 +409,21 @@ class Timeout(RequestExecutionException):
     the operation
     """
 
-    def __init__(self, summary_message, consistency=None, required_responses=None, received_responses=None):
+    def __init__(self, summary_message, consistency=None, required_responses=None,
+                 received_responses=None, **kwargs):
         self.consistency = consistency
         self.required_responses = required_responses
         self.received_responses = received_responses
-        Exception.__init__(self, summary_message + ' info=' +
-                           repr({'consistency': consistency_value_to_name(consistency),
-                                 'required_responses': required_responses,
-                                 'received_responses': received_responses}))
+
+        if "write_type" in kwargs:
+            kwargs["write_type"] = WriteType.value_to_name[kwargs["write_type"]]
+
+        info = {'consistency': consistency_value_to_name(consistency),
+                'required_responses': required_responses,
+                'received_responses': received_responses}
+        info.update(kwargs)
+
+        Exception.__init__(self, summary_message + ' info=' + repr(info))
 
 
 class ReadTimeout(Timeout):
@@ -387,6 +464,7 @@ class WriteTimeout(Timeout):
     """
 
     def __init__(self, message, write_type=None, **kwargs):
+        kwargs["write_type"] = write_type
         Timeout.__init__(self, message, **kwargs)
         self.write_type = write_type
 

cassandra/policies.py

@@ -17,6 +17,14 @@
 from random import randint, shuffle
 from threading import Lock
 import socket
+from cassandra import WriteType as WT
+
+
+# This is done this way because WriteType was originally
+# defined here and in order not to break the API.
+# It may removed in the next mayor.
+WriteType = WT
+
 
 from cassandra import ConsistencyLevel, OperationTimedOut
 
@@ -690,72 +698,6 @@ def new_schedule(self):
             i += 1
 
 
-class WriteType(object):
-    """
-    For usage with :class:`.RetryPolicy`, this describe a type
-    of write operation.
-    """
-
-    SIMPLE = 0
-    """
-    A write to a single partition key. Such writes are guaranteed to be atomic
-    and isolated.
-    """
-
-    BATCH = 1
-    """
-    A write to multiple partition keys that used the distributed batch log to
-    ensure atomicity.
-    """
-
-    UNLOGGED_BATCH = 2
-    """
-    A write to multiple partition keys that did not use the distributed batch
-    log. Atomicity for such writes is not guaranteed.
-    """
-
-    COUNTER = 3
-    """
-    A counter write (for one or multiple partition keys). Such writes should
-    not be replayed in order to avoid overcount.
-    """
-
-    BATCH_LOG = 4
-    """
-    The initial write to the distributed batch log that Cassandra performs
-    internally before a BATCH write.
-    """
-
-    CAS = 5
-    """
-    A lighweight-transaction write, such as "DELETE ... IF EXISTS".
-    """
-
-    VIEW = 6
-    """
-    This WriteType is only seen in results for requests that were unable to
-    complete MV operations.
-    """
-
-    CDC = 7
-    """
-    This WriteType is only seen in results for requests that were unable to
-    complete CDC operations.
-    """
-
-
-WriteType.name_to_value = {
-    'SIMPLE': WriteType.SIMPLE,
-    'BATCH': WriteType.BATCH,
-    'UNLOGGED_BATCH': WriteType.UNLOGGED_BATCH,
-    'COUNTER': WriteType.COUNTER,
-    'BATCH_LOG': WriteType.BATCH_LOG,
-    'CAS': WriteType.CAS,
-    'VIEW': WriteType.VIEW,
-    'CDC': WriteType.CDC
-}
-
-
 class RetryPolicy(object):
     """
     A policy that describes whether to retry, rethrow, or ignore coordinator

cassandra/protocol.py

@@ -40,7 +40,7 @@
                                 UTF8Type, VarcharType, UUIDType, UserType,
                                 TupleType, lookup_casstype, SimpleDateType,
                                 TimeType, ByteType, ShortType, DurationType)
-from cassandra.policies import WriteType
+from cassandra import WriteType
 from cassandra.cython_deps import HAVE_CYTHON, HAVE_NUMPY
 from cassandra import util
 

tests/integration/simulacron/__init__.py

@@ -17,7 +17,12 @@
     import unittest  # noqa
 
 from tests.integration.simulacron.utils import stop_simulacron, clear_queries
+from tests.integration import PROTOCOL_VERSION, SIMULACRON_JAR,  CASSANDRA_VERSION
+from tests.integration.simulacron.utils import start_and_prime_singledc
 
+from cassandra.cluster import Cluster
+
+from packaging.version import Version
 
 def teardown_package():
     stop_simulacron()
@@ -27,3 +32,23 @@ class SimulacronBase(unittest.TestCase):
     def tearDown(self):
         clear_queries()
         stop_simulacron()
+
+
+class SimulacronCluster(SimulacronBase):
+    @classmethod
+    def setUpClass(cls):
+        if SIMULACRON_JAR is None or CASSANDRA_VERSION < Version("2.1"):
+            return
+
+        start_and_prime_singledc()
+        cls.cluster = Cluster(protocol_version=PROTOCOL_VERSION, compression=False)
+        cls.session = cls.cluster.connect(wait_for_all_pools=True)
+
+    @classmethod
+    def tearDownClass(cls):
+        if SIMULACRON_JAR is None or CASSANDRA_VERSION < Version("2.1"):
+            return
+
+        cls.cluster.shutdown()
+        stop_simulacron()
+

tests/integration/simulacron/test_cluster.py

@@ -0,0 +1,55 @@
+# Copyright DataStax, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+try:
+    import unittest2 as unittest
+except ImportError:
+    import unittest  # noqa
+
+from tests.integration.simulacron import SimulacronCluster
+from tests.integration import requiressimulacron
+from tests.integration.simulacron.utils import prime_query
+
+from cassandra import WriteTimeout, WriteType, ConsistencyLevel
+
+@requiressimulacron
+class ClusterTests(SimulacronCluster):
+    def test_writetimeout(self):
+        write_type = "UNLOGGED_BATCH"
+        consistency = "LOCAL_QUORUM"
+        received_responses = 1
+        required_responses = 4
+
+        query_to_prime_simple = "SELECT * from simulacron_keyspace.simple"
+        then = {
+            "result": "write_timeout",
+            "delay_in_ms": 0,
+            "consistency_level": consistency,
+            "received": received_responses,
+            "block_for": required_responses,
+            "write_type": write_type,
+            "ignore_on_prepare": True
+        }
+        prime_query(query_to_prime_simple, then=then, rows=None, column_types=None)
+
+        with self.assertRaises(WriteTimeout) as assert_raised_context:
+            self.session.execute(query_to_prime_simple)
+        wt = assert_raised_context.exception
+        self.assertEqual(wt.write_type, WriteType.name_to_value[write_type])
+        self.assertEqual(wt.consistency, ConsistencyLevel.name_to_value[consistency])
+        self.assertEqual(wt.received_responses, received_responses)
+        self.assertEqual(wt.required_responses, required_responses)
+        self.assertIn(write_type, str(wt))
+        self.assertIn(consistency, str(wt))
+        self.assertIn(str(received_responses), str(wt))
+        self.assertIn(str(required_responses), str(wt))

tests/integration/simulacron/test_policies.py

@@ -29,8 +29,8 @@
 from itertools import count
 from packaging.version import Version
 
-class BadRoundRobinPolicy(RoundRobinPolicy):
 
+class BadRoundRobinPolicy(RoundRobinPolicy):
     def make_query_plan(self, working_keyspace=None, query=None):
         pos = self._position
         self._position += 1
@@ -293,6 +293,7 @@ def test_retry_policy_ignores_and_rethrows(self):
 
         with self.assertRaises(WriteTimeout):
             self.session.execute(query_to_prime_simple)
+
         #CDC should be ignored
         self.session.execute(query_to_prime_cdc)