Source code for chaise.helpers

 1"""
 2Extra bits to make your job easier.
 3"""
 4
 5import random
 6from socket import AddressFamily, SocketKind
 7
 8
 9import anyio
10import httpx
11
12from . import SessionPool
13
14

[docs]
15class ConstantPoolMixin(SessionPool):
16    """
17    A trivial pool mixin that uses the passed in URL(s).
18    """
19
20    urls: list[str]
21
22    def __init__(self, urls: str | list[str]):
23        if isinstance(urls, str):
24            self.urls = [urls]
25        else:
26            self.urls = list(urls)
27        super().__init__()
28

[docs]
29    async def iter_servers(self):
30        u = self.urls[:]
31        random.shuffle(u)
32        for url in u:
33            yield url


34
35
36async def _get_ips(hostname: str | bytes, port: int | None = None):
37    """
38    Look up the IPs for a given hostname
39    """
40    for gai in await anyio.getaddrinfo(hostname, port):
41        match gai:
42            case (AddressFamily.AF_INET, SocketKind.SOCK_STREAM, _, _, (addr, _)):
43                yield addr
44            case (
45                AddressFamily.AF_INET6,
46                SocketKind.SOCK_STREAM,
47                _,
48                _,
49                (addr, _, _, scope),
50            ):
51                # IPv6 flow is without purpose and cannot be represented in addresses
52                if scope:
53                    yield f"{addr}%{scope}"
54                else:
55                    yield addr
56            case _:
57                continue
58
59

[docs]
60class DnsPoolMixin(SessionPool):
61    """
62    Uses DNS round robin to find pool members.
63
64    Assumes that CouchDB is not behind a vhost or uses TLS.
65    """
66
67    url: httpx.URL
68
69    def __init__(self, url: str):
70        self.url = httpx.URL(url)
71        super().__init__()
72

[docs]
73    async def iter_servers(self):
74        # The domain stack already does our shuffling
75        async for ip in _get_ips(self.url.raw_host, self.url.port):
76            if ":" in ip:
77                ip = f"[{ip}]"
78            yield str(self.url.copy_with(host=ip))