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, _),
50            ):
51                yield addr
52            case (
53                AddressFamily.AF_INET6,
54                SocketKind.SOCK_STREAM,
55                _,
56                _,
57                (addr, _, _, scope),
58            ):
59                # IPv6 flow is without purpose and cannot be represented in addresses
60                if scope:
61                    yield f"{addr}%{scope}"
62                else:
63                    yield addr
64            case _:
65                continue
66
67

[docs]
68class DnsPoolMixin(SessionPool):
69    """
70    Uses DNS round robin to find pool members.
71
72    Assumes that CouchDB is not behind a vhost or uses TLS.
73    """
74
75    url: httpx.URL
76
77    def __init__(self, url: str):
78        self.url = httpx.URL(url)
79        super().__init__()
80

[docs]
81    async def iter_servers(self):
82        # The domain stack already does our shuffling
83        async for ip in _get_ips(self.url.raw_host, self.url.port):
84            if ":" in ip:
85                ip = f"[{ip}]"
86            yield str(self.url.copy_with(host=ip))