Class PgConnInfo

misc/python/materialize/util.py:145–196

@dataclass
class PgConnInfo:
    user: str
    host: str
    port: int
    database: str
    password: str | None = None
    ssl: bool = False
    cluster: str | None = None
    autocommit: bool = False

    def connect(self) -> psycopg.Connection:
        conn = psycopg.connect(
            host=self.host,
            port=self.port,
            user=self.user,
            password=self.password,
            dbname=self.database,
            sslmode="require" if self.ssl else None,
        )
        # Set SO_LINGER(1, 0) so close() sends RST instead of FIN, bypassing
        # TIME_WAIT. Prevents exhausting the ~28k ephemeral port range under
        # high connection churn (e.g. benchmarks doing rapid connect/disconnect).
        self._set_linger(conn)
        if self.autocommit:
            conn.autocommit = True
        if self.cluster:
            with conn.cursor() as cur:
                cur.execute(f"SET cluster = {self.cluster}".encode())
        return conn

    @staticmethod
    def _set_linger(conn: psycopg.Connection) -> None:
        import socket
        import struct

        fd = conn.pgconn.socket
        if fd < 0:
            return
        sock = socket.socket(fileno=fd)
        try:
            sock.setsockopt(
                socket.SOL_SOCKET, socket.SO_LINGER, struct.pack("ii", 1, 0)
            )
        finally:
            sock.detach()

    def to_conn_string(self) -> str:
        return (
            f"postgres://{quote(self.user)}:{quote(self.password)}@{self.host}:{self.port}/{quote(self.database)}"
            if self.password
            else f"postgres://{quote(self.user)}@{self.host}:{self.port}/{quote(self.database)}"
        )


def parse_pg_conn_string(conn_string: str) -> PgConnInfo: