From 99ae31105a9a3cbd4525993cefcf0c346092ad9c Mon Sep 17 00:00:00 2001 From: Michal Wrobel Date: Tue, 23 Jul 2024 09:30:54 +0200 Subject: [PATCH] feat: Add support for both IPv4 and IPv6 connections - Updated the method `__get_host_address_info` to support returning both IPv4 and IPv6 address information. - The method now returns a tuple containing the address (sockaddr) and the address family (af). - This allows the caller to handle the address appropriately based on its type. - Modified the initializer to handle both IPv4 and IPv6 addresses. - The method `__get_host_address_info` is called and its response is unpacked into `hostport` and `address_family`. - Created a socket using the correct address family (`AF_INET` for IPv4 and `AF_INET6` for IPv6). - Ensured that the socket is connected using the `hostport` which now correctly handles both address types. --- redshift_connector/core.py | 28 +++++++++++++++++++++------- 1 file changed, 21 insertions(+), 7 deletions(-) diff --git a/redshift_connector/core.py b/redshift_connector/core.py index d5a3ed5..b0f8ad9 100644 --- a/redshift_connector/core.py +++ b/redshift_connector/core.py @@ -395,19 +395,28 @@ def client_os_version(self: "Connection") -> str: os_version = "unknown" return os_version + @staticmethod - def __get_host_address_info(host: str, port: int): + def __get_host_address_info(host: str, port: int) -> typing.Tuple: """ - Returns IPv4 address and port given a host name and port + Returns address and port given a host name and port, supports both IPv4 and IPv6. """ - # https://docs.python.org/3/library/socket.html#socket.getaddrinfo - response = socket.getaddrinfo(host=host, port=port, family=socket.AF_INET) + response = socket.getaddrinfo(host=host, port=port) _logger.debug("getaddrinfo response %s", response) if not response: - raise InterfaceError("Unable to determine ip for host %s port %s", host, port) + raise InterfaceError( + "Unable to determine IP for host %s port %s" % (host, port) + ) - return response[0][4] + for res in response: + af, socktype, proto, canonname, sockaddr = res + if af == socket.AF_INET or af == socket.AF_INET6: + return sockaddr, af + + raise InterfaceError( + "No suitable address found for host %s port %s" % (host, port) + ) def __init__( self: "Connection", @@ -623,8 +632,13 @@ def get_calling_module() -> str: self._usock.settimeout(timeout) if unix_sock is None and host is not None: - hostport: typing.Tuple[str, int] = Connection.__get_host_address_info(host, port) + # hostport: typing.Tuple[str, int] = Connection.__get_host_address_info(host, port) + # _logger.debug("Attempting to create connection socket with address %s", hostport) + # self._usock.connect(hostport) + hostport, address_family = Connection.__get_host_address_info(host, port) _logger.debug("Attempting to create connection socket with address %s", hostport) + + self._usock = socket.socket(address_family, socket.SOCK_STREAM) self._usock.connect(hostport) elif unix_sock is not None: _logger.debug("connecting to socket with unix socket")