mirror of
https://github.com/neondatabase/neon.git
synced 2026-01-08 05:52:55 +00:00
1497a42296
## Problem neon_fixtures.py has grown to unmanageable size. It attracts conflicts. When adding specific utils under for example `fixtures/pageserver` things sometimes need to import stuff from `neon_fixtures.py` which creates circular import. This is usually only needed for type annotations, so `typing.TYPE_CHECKING` flag can mask the issue. Nevertheless I believe that splitting neon_fixtures.py into smaller parts is a better approach. Currently the PR contains small things, but I plan to continue and move NeonEnv to its own `fixtures.env` module. To keep the diff small I think this PR can already be merged to cause less conflicts. UPD: it looks like currently its not really possible to fully avoid usage of `typing.TYPE_CHECKING`, because some components directly depend on each other. I e Env -> Cli -> Env cycle. But its still worth it to avoid it in as many places as possible. And decreasing neon_fixture's size still makes sense.
143 lines
4.0 KiB
Python
143 lines
4.0 KiB
Python
import socket
|
|
import subprocess
|
|
from pathlib import Path
|
|
from types import TracebackType
|
|
from typing import Optional, Type
|
|
|
|
import backoff
|
|
from fixtures.log_helper import log
|
|
from fixtures.neon_fixtures import PgProtocol, VanillaPostgres
|
|
from fixtures.port_distributor import PortDistributor
|
|
|
|
|
|
def generate_tls_cert(cn, certout, keyout):
|
|
subprocess.run(
|
|
[
|
|
"openssl",
|
|
"req",
|
|
"-new",
|
|
"-x509",
|
|
"-days",
|
|
"365",
|
|
"-nodes",
|
|
"-out",
|
|
certout,
|
|
"-keyout",
|
|
keyout,
|
|
"-subj",
|
|
f"/CN={cn}",
|
|
]
|
|
)
|
|
|
|
|
|
class PgSniRouter(PgProtocol):
|
|
def __init__(
|
|
self,
|
|
neon_binpath: Path,
|
|
port: int,
|
|
destination: str,
|
|
tls_cert: Path,
|
|
tls_key: Path,
|
|
test_output_dir: Path,
|
|
):
|
|
# Must use a hostname rather than IP here, for SNI to work
|
|
host = "localhost"
|
|
super().__init__(host=host, port=port)
|
|
|
|
self.host = host
|
|
self.neon_binpath = neon_binpath
|
|
self.port = port
|
|
self.destination = destination
|
|
self.tls_cert = tls_cert
|
|
self.tls_key = tls_key
|
|
self._popen: Optional[subprocess.Popen[bytes]] = None
|
|
self.test_output_dir = test_output_dir
|
|
|
|
def start(self) -> "PgSniRouter":
|
|
assert self._popen is None
|
|
args = [
|
|
str(self.neon_binpath / "pg_sni_router"),
|
|
*["--listen", f"127.0.0.1:{self.port}"],
|
|
*["--tls-cert", str(self.tls_cert)],
|
|
*["--tls-key", str(self.tls_key)],
|
|
*["--destination", self.destination],
|
|
]
|
|
|
|
router_log_path = self.test_output_dir / "pg_sni_router.log"
|
|
router_log = open(router_log_path, "w")
|
|
|
|
self._popen = subprocess.Popen(args, stderr=router_log)
|
|
self._wait_until_ready()
|
|
log.info(f"pg_sni_router started, log file: {router_log_path}")
|
|
return self
|
|
|
|
@backoff.on_exception(backoff.expo, OSError, max_time=10)
|
|
def _wait_until_ready(self):
|
|
socket.create_connection((self.host, self.port))
|
|
|
|
# Sends SIGTERM to the proxy if it has been started
|
|
def terminate(self):
|
|
if self._popen:
|
|
self._popen.terminate()
|
|
|
|
# Waits for proxy to exit if it has been opened with a default timeout of
|
|
# two seconds. Raises subprocess.TimeoutExpired if the proxy does not exit in time.
|
|
def wait_for_exit(self, timeout=2):
|
|
if self._popen:
|
|
self._popen.wait(timeout=2)
|
|
|
|
def __enter__(self) -> "PgSniRouter":
|
|
return self
|
|
|
|
def __exit__(
|
|
self,
|
|
exc_type: Optional[Type[BaseException]],
|
|
exc: Optional[BaseException],
|
|
tb: Optional[TracebackType],
|
|
):
|
|
if self._popen is not None:
|
|
self._popen.terminate()
|
|
try:
|
|
self._popen.wait(timeout=5)
|
|
except subprocess.TimeoutExpired:
|
|
log.warning("failed to gracefully terminate pg_sni_router; killing")
|
|
self._popen.kill()
|
|
|
|
|
|
def test_pg_sni_router(
|
|
vanilla_pg: VanillaPostgres,
|
|
port_distributor: PortDistributor,
|
|
neon_binpath: Path,
|
|
test_output_dir: Path,
|
|
):
|
|
generate_tls_cert(
|
|
"endpoint.namespace.localtest.me",
|
|
test_output_dir / "router.crt",
|
|
test_output_dir / "router.key",
|
|
)
|
|
|
|
# Start a stand-alone Postgres to test with
|
|
vanilla_pg.start()
|
|
pg_port = vanilla_pg.default_options["port"]
|
|
|
|
router_port = port_distributor.get_port()
|
|
|
|
with PgSniRouter(
|
|
neon_binpath=neon_binpath,
|
|
port=router_port,
|
|
destination="localtest.me",
|
|
tls_cert=test_output_dir / "router.crt",
|
|
tls_key=test_output_dir / "router.key",
|
|
test_output_dir=test_output_dir,
|
|
) as router:
|
|
router.start()
|
|
|
|
out = router.safe_psql(
|
|
"select 1",
|
|
dbname="postgres",
|
|
sslmode="require",
|
|
host=f"endpoint--namespace--{pg_port}.localtest.me",
|
|
hostaddr="127.0.0.1",
|
|
)
|
|
assert out[0][0] == 1
|