Add endpoint to connection string

This commit is contained in:
Vadim Kharitonov
2023-05-09 15:46:34 +02:00
committed by Vadim Kharitonov
parent e7514cc15e
commit 4f64be4a98
7 changed files with 153 additions and 39 deletions

View File

@@ -5,16 +5,18 @@ import pytest
from fixtures.neon_fixtures import PSQL, NeonProxy, VanillaPostgres
def test_proxy_select_1(static_proxy: NeonProxy):
@pytest.mark.parametrize("option_name", ["project", "endpoint"])
def test_proxy_select_1(static_proxy: NeonProxy, option_name: str):
"""
A simplest smoke test: check proxy against a local postgres instance.
"""
out = static_proxy.safe_psql("select 1", options="project=generic-project-name")
out = static_proxy.safe_psql("select 1", options=f"{option_name}=generic-project-name")
assert out[0][0] == 1
def test_password_hack(static_proxy: NeonProxy):
@pytest.mark.parametrize("option_name", ["project", "endpoint"])
def test_password_hack(static_proxy: NeonProxy, option_name: str):
"""
Check the PasswordHack auth flow: an alternative to SCRAM auth for
clients which can't provide the project/endpoint name via SNI or `options`.
@@ -23,11 +25,12 @@ def test_password_hack(static_proxy: NeonProxy):
user = "borat"
password = "password"
static_proxy.safe_psql(
f"create role {user} with login password '{password}'", options="project=irrelevant"
f"create role {user} with login password '{password}'",
options=f"{option_name}=irrelevant",
)
# Note the format of `magic`!
magic = f"project=irrelevant;{password}"
magic = f"{option_name}=irrelevant;{password}"
static_proxy.safe_psql("select 1", sslsni=0, user=user, password=magic)
# Must also check that invalid magic won't be accepted.
@@ -56,55 +59,62 @@ async def test_link_auth(vanilla_pg: VanillaPostgres, link_proxy: NeonProxy):
assert out == "42"
def test_proxy_options(static_proxy: NeonProxy):
@pytest.mark.parametrize("option_name", ["project", "endpoint"])
def test_proxy_options(static_proxy: NeonProxy, option_name: str):
"""
Check that we pass extra `options` to the PostgreSQL server:
* `project=...` shouldn't be passed at all (otherwise postgres will raise an error).
* `project=...` and `endpoint=...` shouldn't be passed at all
* (otherwise postgres will raise an error).
* everything else should be passed as-is.
"""
options = "project=irrelevant -cproxytest.option=value"
options = f"{option_name}=irrelevant -cproxytest.option=value"
out = static_proxy.safe_psql("show proxytest.option", options=options)
assert out[0][0] == "value"
options = "-c proxytest.foo=\\ str project=irrelevant"
options = f"-c proxytest.foo=\\ str {option_name}=irrelevant"
out = static_proxy.safe_psql("show proxytest.foo", options=options)
assert out[0][0] == " str"
def test_auth_errors(static_proxy: NeonProxy):
@pytest.mark.parametrize("option_name", ["project", "endpoint"])
def test_auth_errors(static_proxy: NeonProxy, option_name: str):
"""
Check that we throw very specific errors in some unsuccessful auth scenarios.
"""
# User does not exist
with pytest.raises(psycopg2.Error) as exprinfo:
static_proxy.connect(user="pinocchio", options="project=irrelevant")
static_proxy.connect(user="pinocchio", options=f"{option_name}=irrelevant")
text = str(exprinfo.value).strip()
assert text.endswith("password authentication failed for user 'pinocchio'")
static_proxy.safe_psql(
"create role pinocchio with login password 'magic'", options="project=irrelevant"
"create role pinocchio with login password 'magic'",
options=f"{option_name}=irrelevant",
)
# User exists, but password is missing
with pytest.raises(psycopg2.Error) as exprinfo:
static_proxy.connect(user="pinocchio", password=None, options="project=irrelevant")
static_proxy.connect(user="pinocchio", password=None, options=f"{option_name}=irrelevant")
text = str(exprinfo.value).strip()
assert text.endswith("password authentication failed for user 'pinocchio'")
# User exists, but password is wrong
with pytest.raises(psycopg2.Error) as exprinfo:
static_proxy.connect(user="pinocchio", password="bad", options="project=irrelevant")
static_proxy.connect(user="pinocchio", password="bad", options=f"{option_name}=irrelevant")
text = str(exprinfo.value).strip()
assert text.endswith("password authentication failed for user 'pinocchio'")
# Finally, check that the user can connect
with static_proxy.connect(user="pinocchio", password="magic", options="project=irrelevant"):
with static_proxy.connect(
user="pinocchio", password="magic", options=f"{option_name}=irrelevant"
):
pass
def test_forward_params_to_client(static_proxy: NeonProxy):
@pytest.mark.parametrize("option_name", ["project", "endpoint"])
def test_forward_params_to_client(static_proxy: NeonProxy, option_name: str):
"""
Check that we forward all necessary PostgreSQL server params to client.
"""
@@ -130,7 +140,7 @@ def test_forward_params_to_client(static_proxy: NeonProxy):
where name = any(%s)
"""
with static_proxy.connect(options="project=irrelevant") as conn:
with static_proxy.connect(options=f"{option_name}=irrelevant") as conn:
with conn.cursor() as cur:
cur.execute(query, (reported_params_subset,))
for name, value in cur.fetchall():
@@ -138,17 +148,18 @@ def test_forward_params_to_client(static_proxy: NeonProxy):
assert conn.get_parameter_status(name) == value
@pytest.mark.parametrize("option_name", ["project", "endpoint"])
@pytest.mark.timeout(5)
def test_close_on_connections_exit(static_proxy: NeonProxy):
def test_close_on_connections_exit(static_proxy: NeonProxy, option_name: str):
# Open two connections, send SIGTERM, then ensure that proxy doesn't exit
# until after connections close.
with static_proxy.connect(options="project=irrelevant"), static_proxy.connect(
options="project=irrelevant"
with static_proxy.connect(options=f"{option_name}=irrelevant"), static_proxy.connect(
options=f"{option_name}=irrelevant"
):
static_proxy.terminate()
with pytest.raises(subprocess.TimeoutExpired):
static_proxy.wait_for_exit(timeout=2)
# Ensure we don't accept any more connections
with pytest.raises(psycopg2.OperationalError):
static_proxy.connect(options="project=irrelevant")
static_proxy.connect(options=f"{option_name}=irrelevant")
static_proxy.wait_for_exit()