mirror of
https://github.com/neondatabase/neon.git
synced 2026-01-08 14:02:55 +00:00
100 lines
3.3 KiB
Python
100 lines
3.3 KiB
Python
import json
|
|
from urllib.parse import urlparse
|
|
|
|
import psycopg2
|
|
import pytest
|
|
from fixtures.log_helper import log
|
|
from fixtures.neon_fixtures import PSQL, NeonProxy, VanillaPostgres
|
|
|
|
|
|
def test_proxy_select_1(static_proxy: NeonProxy):
|
|
static_proxy.safe_psql("select 1", options="project=generic-project-name")
|
|
|
|
|
|
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"
|
|
)
|
|
|
|
# Note the format of `magic`!
|
|
magic = f"project=irrelevant;{password}"
|
|
static_proxy.safe_psql("select 1", sslsni=0, user=user, password=magic)
|
|
|
|
# Must also check that invalid magic won't be accepted.
|
|
with pytest.raises(psycopg2.OperationalError):
|
|
magic = "broken"
|
|
static_proxy.safe_psql("select 1", sslsni=0, user=user, password=magic)
|
|
|
|
|
|
def get_session_id(uri_prefix, uri_line):
|
|
assert uri_prefix in uri_line
|
|
|
|
url_parts = urlparse(uri_line)
|
|
psql_session_id = url_parts.path[1:]
|
|
assert psql_session_id.isalnum(), "session_id should only contain alphanumeric chars"
|
|
|
|
return psql_session_id
|
|
|
|
|
|
async def find_auth_link(link_auth_uri_prefix, proc):
|
|
for _ in range(100):
|
|
line = (await proc.stderr.readline()).decode("utf-8").strip()
|
|
log.info(f"psql line: {line}")
|
|
if link_auth_uri_prefix in line:
|
|
log.info(f"SUCCESS, found auth url: {line}")
|
|
return line
|
|
|
|
|
|
async def activate_link_auth(local_vanilla_pg, link_proxy, psql_session_id):
|
|
pg_user = "proxy"
|
|
|
|
log.info("creating a new user for link auth test")
|
|
local_vanilla_pg.start()
|
|
local_vanilla_pg.safe_psql(f"create user {pg_user} with login superuser")
|
|
|
|
db_info = json.dumps(
|
|
{
|
|
"session_id": psql_session_id,
|
|
"result": {
|
|
"Success": {
|
|
"host": local_vanilla_pg.default_options["host"],
|
|
"port": local_vanilla_pg.default_options["port"],
|
|
"dbname": local_vanilla_pg.default_options["dbname"],
|
|
"user": pg_user,
|
|
"project": "irrelevant",
|
|
}
|
|
},
|
|
}
|
|
)
|
|
|
|
log.info("sending session activation message")
|
|
psql = await PSQL(host=link_proxy.host, port=link_proxy.mgmt_port).run(db_info)
|
|
out = (await psql.stdout.read()).decode("utf-8").strip()
|
|
assert out == "ok"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_psql_session_id(vanilla_pg: VanillaPostgres, link_proxy: NeonProxy):
|
|
psql = await PSQL(host=link_proxy.host, port=link_proxy.proxy_port).run("select 42")
|
|
|
|
uri_prefix = link_proxy.link_auth_uri_prefix
|
|
link = await find_auth_link(uri_prefix, psql)
|
|
|
|
psql_session_id = get_session_id(uri_prefix, link)
|
|
await activate_link_auth(vanilla_pg, link_proxy, psql_session_id)
|
|
|
|
assert psql.stdout is not None
|
|
out = (await psql.stdout.read()).decode("utf-8").strip()
|
|
assert out == "42"
|
|
|
|
|
|
# Pass extra options to the server.
|
|
def test_proxy_options(static_proxy: NeonProxy):
|
|
with static_proxy.connect(options="project=irrelevant -cproxytest.option=value") as conn:
|
|
with conn.cursor() as cur:
|
|
cur.execute("SHOW proxytest.option")
|
|
value = cur.fetchall()[0][0]
|
|
assert value == "value"
|