Merge pull request #2 from riptano/python-1098

aboudreault · web-flow · commit a06f4b9c213e · 2019-08-15T10:05:07.000-04:00
PYTHON-1098: Add a new Endpoint type to support unix sockets
diff --git a/CHANGELOG.rst b/CHANGELOG.rst
@@ -5,6 +5,7 @@ Not Released
 Features
 --------
 * Add SNIEndPoint support (PYTHON-1084)
+* Add a new Endpoint type to support unix sockets (PYTHON-1098)
 
 3.19.0
 ======
diff --git a/build.yaml b/build.yaml
@@ -150,6 +150,9 @@ build:
       export PATH=$JAVA_HOME/bin:$PATH
       export PYTHONPATH=""
 
+      # Required for unix socket tests
+      sudo apt-get install socat
+
       # Install latest setuptools
       pip install --upgrade pip
       pip install -U setuptools
diff --git a/cassandra/connection.py b/cassandra/connection.py
@@ -144,6 +144,13 @@ def ssl_options(self):
         """
         return None
 
+    @property
+    def socket_family(self):
+        """
+        The socket family of the endpoint.
+        """
+        return socket.AF_UNSPEC
+
     def resolve(self):
         """
         Resolve the endpoint to an address/port. This is called
@@ -299,6 +306,47 @@ def create_from_sni(self, sni):
         return SniEndPoint(self._proxy_address, sni, self._port)
 
 
+@total_ordering
+class UnixSocketEndPoint(EndPoint):
+    """
+    Unix Socket EndPoint implementation.
+    """
+
+    def __init__(self, unix_socket_path):
+        self._unix_socket_path = unix_socket_path
+
+    @property
+    def address(self):
+        return self._unix_socket_path
+
+    @property
+    def port(self):
+        return None
+
+    @property
+    def socket_family(self):
+        return socket.AF_UNIX
+
+    def resolve(self):
+        return self.address, None
+
+    def __eq__(self, other):
+        return (isinstance(other, UnixSocketEndPoint) and
+                self._unix_socket_path == other._unix_socket_path)
+
+    def __hash__(self):
+        return hash(self._unix_socket_path)
+
+    def __lt__(self, other):
+        return self._unix_socket_path < other._unix_socket_path
+
+    def __str__(self):
+        return str("%s" % (self._unix_socket_path,))
+
+    def __repr__(self):
+        return "<%s: %s>" % (self.__class__.__name__, self._unix_socket_path)
+
+
 class _Frame(object):
     def __init__(self, version, flags, stream, opcode, body_offset, end_pos):
         self.version = version
@@ -557,13 +605,22 @@ def factory(cls, endpoint, timeout, *args, **kwargs):
         else:
             return conn
 
-    def _connect_socket(self):
-        sockerr = None
-        inet_address, port = self.endpoint.resolve()
-        addresses = socket.getaddrinfo(inet_address, port, socket.AF_UNSPEC, socket.SOCK_STREAM)
+    def _get_socket_addresses(self):
+        address, port = self.endpoint.resolve()
+
+        if self.endpoint.socket_family == socket.AF_UNIX:
+            return [(socket.AF_UNIX, socket.SOCK_STREAM, 0, None, address)]
+
+        addresses = socket.getaddrinfo(address, port, self.endpoint.socket_family, socket.SOCK_STREAM)
         if not addresses:
             raise ConnectionException("getaddrinfo returned empty list for %s" % (self.endpoint,))
-        for (af, socktype, proto, canonname, sockaddr) in addresses:
+
+        return addresses
+
+    def _connect_socket(self):
+        sockerr = None
+        addresses = self._get_socket_addresses()
+        for (af, socktype, proto, _, sockaddr) in addresses:
             try:
                 self._socket = self._socket_impl.socket(af, socktype, proto)
                 if self.ssl_context:
@@ -587,7 +644,8 @@ def _connect_socket(self):
                 sockerr = err
 
         if sockerr:
-            raise socket.error(sockerr.errno, "Tried connecting to %s. Last error: %s" % ([a[4] for a in addresses], sockerr.strerror or sockerr))
+            raise socket.error(sockerr.errno, "Tried connecting to %s. Last error: %s" %
+                               ([a[4] for a in addresses], sockerr.strerror or sockerr))
 
         if self.sockopts:
             for args in self.sockopts:
diff --git a/docs/api/cassandra/connection.rst b/docs/api/cassandra/connection.rst
@@ -17,3 +17,5 @@
 .. autoclass:: SniEndPoint
 
 .. autoclass:: SniEndPointFactory
+
+.. autoclass:: UnixSocketEndPoint
diff --git a/tests/integration/advanced/__init__.py b/tests/integration/advanced/__init__.py
@@ -0,0 +1,23 @@
+# 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
+
+try:
+    from ccmlib import common
+except ImportError as e:
+    raise unittest.SkipTest('ccm is a dependency for integration tests:', e)
diff --git a/tests/integration/advanced/test_unixsocketendpoint.py b/tests/integration/advanced/test_unixsocketendpoint.py
@@ -0,0 +1,77 @@
+# 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
+
+import time
+import subprocess
+import logging
+
+from cassandra.cluster import Cluster, ExecutionProfile, EXEC_PROFILE_DEFAULT
+from cassandra.connection import UnixSocketEndPoint
+from cassandra.policies import WhiteListRoundRobinPolicy, RoundRobinPolicy
+
+from tests import notwindows
+from tests.integration import use_single_node
+
+log = logging.getLogger()
+log.setLevel('DEBUG')
+
+UNIX_SOCKET_PATH = '/tmp/cass.sock'
+
+
+def setup_module():
+    use_single_node()
+
+
+class UnixSocketWhiteListRoundRobinPolicy(WhiteListRoundRobinPolicy):
+    def __init__(self, hosts):
+        self._allowed_hosts = self._allowed_hosts_resolved = tuple(hosts)
+        RoundRobinPolicy.__init__(self)
+
+
+@notwindows
+class UnixSocketTest(unittest.TestCase):
+
+    @classmethod
+    def setUpClass(cls):
+        log.debug("Starting socat...")
+        cls.proc = subprocess.Popen(
+            ['socat',
+             'UNIX-LISTEN:%s,fork' % UNIX_SOCKET_PATH,
+             'TCP:localhost:9042'],
+            stdout=subprocess.PIPE,
+            stderr=subprocess.STDOUT)
+
+        time.sleep(1)
+        if cls.proc.poll() is not None:
+            for line in cls.proc.stdout.readlines():
+                log.debug("socat: " + line.decode('utf-8'))
+            raise Exception("Error while starting socat. Return code: %d" % cls.proc.returncode)
+
+        lbp = UnixSocketWhiteListRoundRobinPolicy([UNIX_SOCKET_PATH])
+        ep = ExecutionProfile(load_balancing_policy=lbp)
+        endpoint = UnixSocketEndPoint(UNIX_SOCKET_PATH)
+        cls.cluster = Cluster([endpoint], execution_profiles={EXEC_PROFILE_DEFAULT: ep})
+
+    @classmethod
+    def tearDownClass(cls):
+        cls.cluster.shutdown()
+        cls.proc.terminate()
+
+    def test_unix_socket_connection(self):
+        s = self.cluster.connect()
+        s.execute('select * from system.local')
