PYTHON-751: Add IPV4Address/IPV6Address support for inet type (apache#828)

aboudreault · web-flow · commit 9869c2a044e5 · 2017-10-04T12:42:10.000-04:00
* Add IPV4Address/IPV6Address support for inet type
diff --git a/CHANGELOG.rst b/CHANGELOG.rst
@@ -4,6 +4,7 @@
 Features
 --------
 * Send keyspace in QUERY, PREPARE, and BATCH messages (PYTHON-678)
+* Add IPv4Address/IPv6Address support for inet types (PYTHON-751)
 
 Bug Fixes
 ---------
diff --git a/cassandra/cqltypes.py b/cassandra/cqltypes.py
@@ -43,6 +43,8 @@
 from uuid import UUID
 import warnings
 
+if six.PY3:
+    import ipaddress
 
 from cassandra.marshal import (int8_pack, int8_unpack, int16_pack, int16_unpack,
                                uint16_pack, uint16_unpack, uint32_pack, uint32_unpack,
@@ -517,12 +519,17 @@ def deserialize(byts, protocol_version):
 
     @staticmethod
     def serialize(addr, protocol_version):
-        if ':' in addr:
-            return util.inet_pton(socket.AF_INET6, addr)
-        else:
-            # util.inet_pton could also handle, but this is faster
-            # since we've already determined the AF
-            return socket.inet_aton(addr)
+        try:
+            if ':' in addr:
+                return util.inet_pton(socket.AF_INET6, addr)
+            else:
+                # util.inet_pton could also handle, but this is faster
+                # since we've already determined the AF
+                return socket.inet_aton(addr)
+        except:
+            if six.PY3 and isinstance(addr, (ipaddress.IPv4Address, ipaddress.IPv6Address)):
+                return addr.packed
+            raise ValueError("can't interpret %r as an inet address" % (addr,))
 
 
 class CounterColumnType(LongType):
diff --git a/cassandra/encoder.py b/cassandra/encoder.py
@@ -29,6 +29,9 @@
 from uuid import UUID
 import six
 
+if six.PY3:
+    import ipaddress
+
 from cassandra.util import (OrderedDict, OrderedMap, OrderedMapSerializedKey,
                             sortedset, Time, Date)
 
@@ -103,6 +106,8 @@ def __init__(self):
                 memoryview: self.cql_encode_bytes,
                 bytes: self.cql_encode_bytes,
                 type(None): self.cql_encode_none,
+                ipaddress.IPv4Address: self.cql_encode_ipaddress,
+                ipaddress.IPv6Address: self.cql_encode_ipaddress
             })
 
     def cql_encode_none(self, val):
@@ -225,3 +230,11 @@ def cql_encode_all_types(self, val):
         if :attr:`~Encoder.mapping` does not contain an entry for the type.
         """
         return self.mapping.get(type(val), self.cql_encode_object)(val)
+
+    if six.PY3:
+        def cql_encode_ipaddress(self, val):
+            """
+            Converts an ipaddress (IPV4Address, IPV6Address) to a CQL string. This
+            is suitable for ``inet`` type columns.
+            """
+            return "'%s'" % val.compressed
diff --git a/tests/__init__.py b/tests/__init__.py
@@ -59,6 +59,33 @@ def is_monkey_patched():
     return is_gevent_monkey_patched() or is_eventlet_monkey_patched()
 
 
+EVENT_LOOP_MANAGER = os.getenv('EVENT_LOOP_MANAGER', "libev")
+if "gevent" in EVENT_LOOP_MANAGER:
+    import gevent.monkey
+    gevent.monkey.patch_all()
+    from cassandra.io.geventreactor import GeventConnection
+    connection_class = GeventConnection
+elif "eventlet" in EVENT_LOOP_MANAGER:
+    from eventlet import monkey_patch
+    monkey_patch()
+
+    from cassandra.io.eventletreactor import EventletConnection
+    connection_class = EventletConnection
+elif "async" in EVENT_LOOP_MANAGER:
+    from cassandra.io.asyncorereactor import AsyncoreConnection
+    connection_class = AsyncoreConnection
+elif "twisted" in EVENT_LOOP_MANAGER:
+    from cassandra.io.twistedreactor import TwistedConnection
+    connection_class = TwistedConnection
+
+else:
+    try:
+        from cassandra.io.libevreactor import LibevConnection
+        connection_class = LibevConnection
+    except ImportError:
+        connection_class = None
+
+
 MONKEY_PATCH_LOOP = bool(os.getenv('MONKEY_PATCH_LOOP', False))
 
 notwindows = unittest.skipUnless(not "Windows" in platform.system(), "This test is not adequate for windows")
diff --git a/tests/integration/__init__.py b/tests/integration/__init__.py
@@ -12,31 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import os
-
-EVENT_LOOP_MANAGER = os.getenv('EVENT_LOOP_MANAGER', "libev")
-if "gevent" in EVENT_LOOP_MANAGER:
-    import gevent.monkey
-    gevent.monkey.patch_all()
-    from cassandra.io.geventreactor import GeventConnection
-    connection_class = GeventConnection
-elif "eventlet" in EVENT_LOOP_MANAGER:
-    from eventlet import monkey_patch
-    monkey_patch()
-
-    from cassandra.io.eventletreactor import EventletConnection
-    connection_class = EventletConnection
-elif "async" in EVENT_LOOP_MANAGER:
-    from cassandra.io.asyncorereactor import AsyncoreConnection
-    connection_class = AsyncoreConnection
-elif "twisted" in EVENT_LOOP_MANAGER:
-    from cassandra.io.twistedreactor import TwistedConnection
-    connection_class = TwistedConnection
-
-else:
-    from cassandra.io.libevreactor import LibevConnection
-    connection_class = LibevConnection
-
 from cassandra.cluster import Cluster
+
+from tests import connection_class, EVENT_LOOP_MANAGER
 Cluster.connection_class = connection_class
 
 try:
diff --git a/tests/integration/cqlengine/columns/test_validation.py b/tests/integration/cqlengine/columns/test_validation.py
@@ -529,7 +529,6 @@ def test_default_zero_fields_validate(self):
 
 
 class TestAscii(BaseCassEngTestCase):
-
     def test_min_length(self):
         """ Test arbitrary minimal lengths requirements. """
 
diff --git a/tests/integration/datatype_utils.py b/tests/integration/datatype_utils.py
@@ -15,6 +15,7 @@
 from decimal import Decimal
 from datetime import datetime, date, time
 from uuid import uuid1, uuid4
+import six
 
 from cassandra.util import OrderedMap, Date, Time, sortedset, Duration
 
@@ -90,7 +91,11 @@ def get_sample_data():
             sample_data[datatype] = 3.4028234663852886e+38
 
         elif datatype == 'inet':
-            sample_data[datatype] = '123.123.123.123'
+            sample_data[datatype] = ('123.123.123.123', '2001:db8:85a3:8d3:1319:8a2e:370:7348')
+            if six.PY3:
+                import ipaddress
+                sample_data[datatype] += (ipaddress.IPv4Address("123.123.123.123"),
+                                          ipaddress.IPv6Address('2001:db8:85a3:8d3:1319:8a2e:370:7348'))
 
         elif datatype == 'int':
             sample_data[datatype] = 2147483647
@@ -140,10 +145,20 @@ def get_sample(datatype):
     """
     Helper method to access created sample data for primitive types
     """
-
+    if isinstance(SAMPLE_DATA[datatype], tuple):
+        return SAMPLE_DATA[datatype][0]
     return SAMPLE_DATA[datatype]
 
 
+def get_all_samples(datatype):
+    """
+    Helper method to access created sample data for primitive types
+    """
+    if isinstance(SAMPLE_DATA[datatype], tuple):
+        return SAMPLE_DATA[datatype]
+    return SAMPLE_DATA[datatype],
+
+
 def get_collection_sample(collection_type, datatype):
     """
     Helper method to access created sample data for collection types
diff --git a/tests/integration/standard/test_cluster.py b/tests/integration/standard/test_cluster.py
@@ -34,7 +34,7 @@
 from cassandra.pool import Host
 from cassandra.query import SimpleStatement, TraceUnavailable, tuple_factory
 
-
+from tests import notwindows
 from tests.integration import use_singledc, PROTOCOL_VERSION, get_server_versions, CASSANDRA_VERSION, \
     execute_until_pass, execute_with_long_wait_retry, get_node, MockLoggingHandler, get_unsupported_lower_protocol, \
     get_unsupported_upper_protocol, protocolv5, local, CASSANDRA_IP
@@ -476,6 +476,7 @@ def test_refresh_schema_type(self):
         cluster.shutdown()
 
     @local
+    @notwindows
     def test_refresh_schema_no_wait(self):
         contact_points = [CASSANDRA_IP]
         cluster = Cluster(protocol_version=PROTOCOL_VERSION, max_schema_agreement_wait=10,
diff --git a/tests/integration/standard/test_types.py b/tests/integration/standard/test_types.py
@@ -33,7 +33,7 @@
 from tests.integration import use_singledc, PROTOCOL_VERSION, execute_until_pass, notprotocolv1, \
     BasicSharedKeyspaceUnitTestCase, greaterthancass21, lessthancass30, greaterthanorequalcass3_10
 from tests.integration.datatype_utils import update_datatypes, PRIMITIVE_DATATYPES, COLLECTION_TYPES, PRIMITIVE_DATATYPES_KEYS, \
-    get_sample, get_collection_sample
+    get_sample, get_all_samples, get_collection_sample
 
 
 def setup_module():
@@ -161,8 +161,30 @@ def test_can_insert_primitive_datatypes(self):
         for expected, actual in zip(params, results):
             self.assertEqual(actual, expected)
 
+        # try the same thing sending one insert at the time
+        s.execute("TRUNCATE alltypes;")
+        for i, datatype in enumerate(PRIMITIVE_DATATYPES):
+            single_col_name = chr(start_index + i)
+            single_col_names = ["zz", single_col_name]
+            placeholders = ','.join(["%s"] * len(single_col_names))
+            single_columns_string = ', '.join(single_col_names)
+            for j, data_sample in enumerate(get_all_samples(datatype)):
+                key = i + 1000 * j
+                single_params = (key, data_sample)
+                s.execute("INSERT INTO alltypes ({0}) VALUES ({1})".format(single_columns_string, placeholders),
+                          single_params)
+                # verify data
+                result = s.execute("SELECT {0} FROM alltypes WHERE zz=%s".format(single_columns_string), (key,))[0][1]
+                compare_value = data_sample
+                if six.PY3:
+                    import ipaddress
+                    if isinstance(data_sample, ipaddress.IPv4Address) or isinstance(data_sample, ipaddress.IPv6Address):
+                        compare_value = str(data_sample)
+                self.assertEqual(result, compare_value)
+
         # try the same thing with a prepared statement
         placeholders = ','.join(["?"] * len(col_names))
+        s.execute("TRUNCATE alltypes;")
         insert = s.prepare("INSERT INTO alltypes ({0}) VALUES ({1})".format(columns_string, placeholders))
         s.execute(insert.bind(params))
 
diff --git a/tests/unit/test_cluster.py b/tests/unit/test_cluster.py
@@ -27,11 +27,7 @@
 from cassandra.query import SimpleStatement, named_tuple_factory, tuple_factory
 from cassandra.pool import Host
 from tests.unit.utils import mock_session_pools
-
-try:
-    from cassandra.io.libevreactor import LibevConnection
-except ImportError:
-    LibevConnection = None  # noqa
+from tests import connection_class
 
 
 class ExceptionTypeTest(unittest.TestCase):
@@ -129,9 +125,9 @@ def test_event_delay_timing(self, *_):
 
 class SessionTest(unittest.TestCase):
     def setUp(self):
-        if LibevConnection is None:
+        if connection_class is None:
             raise unittest.SkipTest('libev does not appear to be installed correctly')
-        LibevConnection.initialize_reactor()
+        connection_class.initialize_reactor()
 
     # TODO: this suite could be expanded; for now just adding a test covering a PR
     @mock_session_pools
@@ -164,9 +160,9 @@ def test_default_serial_consistency_level(self, *_):
 
 class ExecutionProfileTest(unittest.TestCase):
     def setUp(self):
-        if LibevConnection is None:
+        if connection_class is None:
             raise unittest.SkipTest('libev does not appear to be installed correctly')
-        LibevConnection.initialize_reactor()
+        connection_class.initialize_reactor()
 
     def _verify_response_future_profile(self, rf, prof):
         self.assertEqual(rf._load_balancer, prof.load_balancing_policy)
