Repository: redis/redis-py
File: redis/asyncio/sentinel.py (Python)

Code before:


class MasterNotFoundError(ConnectionError):
    pass


class SlaveNotFoundError(ConnectionError):
    pass


class SentinelManagedConnection(Connection):
    def __init__(self, **kwargs):
        self.connection_pool = kwargs.pop("connection_pool")
        super().__init__(**kwargs)

    def __repr__(self):
        s = f"<{self.__class__.__module__}.{self.__class__.__name__}"
        if self.host:
            host_info = f",host={self.host},port={self.port}"
            s += host_info
        return s + ")>"

    async def connect_to(self, address):
        self.host, self.port = address

        if self.is_connected:
            return
        try:
            await self._connect()
        except asyncio.CancelledError:
            raise  # in 3.7 and earlier, this is an Exception, not BaseException
        except (socket.timeout, asyncio.TimeoutError):
            raise TimeoutError("Timeout connecting to server")
        except OSError as e:
            raise ConnectionError(self._error_message(e))
        except Exception as exc:
            raise ConnectionError(exc) from exc

        try:
            if not self.redis_connect_func:
                # Use the default on_connect function
                await self.on_connect_check_health(
                    check_health=self.connection_pool.check_connection
                )
            else:
                # Use the passed function redis_connect_func
                (
                    await self.redis_connect_func(self)
                    if asyncio.iscoroutinefunction(self.redis_connect_func)
                    else self.redis_connect_func(self)
                )

Review comment (refactor):
That’s quite a bit of code duplication. Perhaps the approach with a flag would be preferable after all. Since `connect()` is part of the interface, I think it’s better not to modify it directly. Instead, you can use the `connect_with_health_check()` method from AbstractConnection. You could introduce an additional flag to enable or disable `retry_socket_connect`, with a default value that preserves the current behaviour.

Code after:
    pass


class SlaveNotFoundError(ConnectionError):
    pass


class SentinelManagedConnection(Connection):
    def __init__(self, **kwargs):
        self.connection_pool = kwargs.pop("connection_pool")
        super().__init__(**kwargs)

    def __repr__(self):
        s = f"<{self.__class__.__module__}.{self.__class__.__name__}"
        if self.host:
            host_info = f",host={self.host},port={self.port}"
            s += host_info
        return s + ")>"

    async def connect_to(self, address):
        self.host, self.port = address
        await self.connect_check_health(
            check_health=self.connection_pool.check_connection,
            retry_socket_connect=False,
        )

    async def _connect_retry(self):
        if self._reader:
            return  # already connected
        if self.connection_pool.is_master:
            await self.connect_to(await self.connection_pool.get_master_address())
        else:
            async for slave in self.connection_pool.rotate_slaves():
                try:
                    return await self.connect_to(slave)
                except ConnectionError:
                    continue
            raise SlaveNotFoundError  # Never be here

    async def connect(self):
        return await self.retry.call_with_retry(
            self._connect_retry,
            lambda error: asyncio.sleep(0),
        )

    async def read_response(
        self,
        disable_decoding: bool = False,
        timeout: Optional[float] = None,
        *,
        disconnect_on_error: Optional[float] = True,