From c62dcae4cc29b7a9e35dbd3e7aded99f21dbc9d4 Mon Sep 17 00:00:00 2001 From: Jelte Fennema-Nio Date: Fri, 26 Dec 2025 12:31:43 +0100 Subject: [PATCH v6 5/7] Convert load balance tests from perl to python --- src/interfaces/libpq/Makefile | 1 + src/interfaces/libpq/meson.build | 7 +- src/interfaces/libpq/pyt/test_load_balance.py | 170 ++++++++++++++++++ .../libpq/t/003_load_balance_host_list.pl | 94 ---------- .../libpq/t/004_load_balance_dns.pl | 144 --------------- 5 files changed, 176 insertions(+), 240 deletions(-) create mode 100644 src/interfaces/libpq/pyt/test_load_balance.py delete mode 100644 src/interfaces/libpq/t/003_load_balance_host_list.pl delete mode 100644 src/interfaces/libpq/t/004_load_balance_dns.pl diff --git a/src/interfaces/libpq/Makefile b/src/interfaces/libpq/Makefile index 9fe321147fc..41ea88c7388 100644 --- a/src/interfaces/libpq/Makefile +++ b/src/interfaces/libpq/Makefile @@ -167,6 +167,7 @@ check installcheck: export PATH := $(CURDIR)/test:$(PATH) check: test-build all $(prove_check) + $(pytest_check) installcheck: test-build all $(prove_installcheck) diff --git a/src/interfaces/libpq/meson.build b/src/interfaces/libpq/meson.build index b259c998fa2..6d62ac17edb 100644 --- a/src/interfaces/libpq/meson.build +++ b/src/interfaces/libpq/meson.build @@ -150,8 +150,6 @@ tests += { 'tests': [ 't/001_uri.pl', 't/002_api.pl', - 't/003_load_balance_host_list.pl', - 't/004_load_balance_dns.pl', 't/005_negotiate_encryption.pl', 't/006_service.pl', ], @@ -162,6 +160,11 @@ tests += { }, 'deps': libpq_test_deps, }, + 'pytest': { + 'tests': [ + 'pyt/test_load_balance.py', + ], + }, } subdir('po', if_found: libintl) diff --git a/src/interfaces/libpq/pyt/test_load_balance.py b/src/interfaces/libpq/pyt/test_load_balance.py new file mode 100644 index 00000000000..0af46d8f37d --- /dev/null +++ b/src/interfaces/libpq/pyt/test_load_balance.py @@ -0,0 +1,170 @@ +# Copyright (c) 2025, PostgreSQL Global Development Group + +""" +Tests for load_balance_hosts connection parameter. + +These tests verify that libpq correctly handles load balancing across multiple +PostgreSQL servers specified in the connection string. +""" + +import platform +import re + +import pytest + +from libpq import LibpqError +import pypg + + +@pytest.fixture(scope="module") +def load_balance_nodes_hostlist(create_pg_module): + """ + Create 3 PostgreSQL nodes with different socket directories. + + Each node has its own Unix socket directory for isolation. + Returns a tuple of (nodes, connect). + """ + nodes = [create_pg_module() for _ in range(3)] + + hostlist = ",".join(node.host for node in nodes) + portlist = ",".join(str(node.port) for node in nodes) + + def connect(**kwargs): + return nodes[0].connect(host=hostlist, port=portlist, **kwargs) + + return nodes, connect + + +@pytest.fixture(scope="module") +def load_balance_nodes_dns(create_pg_module): + """ + Create 3 PostgreSQL nodes on the same port but different IP addresses. + + Uses 127.0.0.1, 127.0.0.2, 127.0.0.3 with a shared port, so that + connections to 'pg-loadbalancetest' can be load balanced via DNS. + + Since setting up a DNS server is more effort than we consider reasonable to + run this test, this situation is instead imitated by using a hosts file + where a single hostname maps to multiple different IP addresses. This test + requires the administrator to add the following lines to the hosts file (if + we detect that this hasn't happened we skip the test): + + 127.0.0.1 pg-loadbalancetest + 127.0.0.2 pg-loadbalancetest + 127.0.0.3 pg-loadbalancetest + + Windows or Linux are required to run this test because these OSes allow + binding to 127.0.0.2 and 127.0.0.3 addresses by default, but other OSes + don't. We need to bind to different IP addresses, so that we can use these + different IP addresses in the hosts file. + + The hosts file needs to be prepared before running this test. We don't do + it on the fly, because it requires root permissions to change the hosts + file. In CI we set up the previously mentioned rules in the hosts file, so + that this load balancing method is tested. + + Requires PG_TEST_EXTRA=load_balance because it requires this manual hosts + file configuration and also uses TCP with trust auth, which is potentially + unsafe on multiuser systems. + """ + pypg.skip_unless_test_extras("load_balance") + + if platform.system() not in ("Linux", "Windows"): + pytest.skip("DNS load balance test only supported on Linux and Windows") + + if platform.system() == "Windows": + hosts_path = r"c:\Windows\System32\Drivers\etc\hosts" + else: + hosts_path = "/etc/hosts" + + try: + with open(hosts_path) as f: + hosts_content = f.read() + except (OSError, IOError): + pytest.skip(f"Could not read hosts file: {hosts_path}") + + count = len(re.findall(r"127\.0\.0\.[1-3]\s+pg-loadbalancetest", hosts_content)) + if count != 3: + pytest.skip("hosts file not prepared for DNS load balance test") + + first_node = create_pg_module(hostaddr="127.0.0.1") + nodes = [ + first_node, + create_pg_module(hostaddr="127.0.0.2", port=first_node.port), + create_pg_module(hostaddr="127.0.0.3", port=first_node.port), + ] + + # Allow trust authentication for TCP connections from loopback + for node in nodes: + hba_path = node.datadir / "pg_hba.conf" + with open(hba_path, "r") as f: + original_content = f.read() + with open(hba_path, "w") as f: + f.write("host all all 127.0.0.0/8 trust\n") + f.write(original_content) + node.pg_ctl("reload") + + def connect(**kwargs): + return nodes[0].connect(host="pg-loadbalancetest", **kwargs) + + return nodes, connect + + +@pytest.fixture(scope="module", params=["hostlist", "dns"]) +def load_balance_nodes(request): + """ + Parametrized fixture providing both load balancing test environments. + """ + return request.getfixturevalue(f"load_balance_nodes_{request.param}") + + +def test_load_balance_hosts_invalid_value(load_balance_nodes): + """load_balance_hosts doesn't accept unknown values.""" + _, connect = load_balance_nodes + + with pytest.raises( + LibpqError, match='invalid load_balance_hosts value: "doesnotexist"' + ): + connect(load_balance_hosts="doesnotexist") + + +def test_load_balance_hosts_disable(load_balance_nodes): + """load_balance_hosts=disable always connects to the first node.""" + nodes, connect = load_balance_nodes + + with nodes[0].log_contains("connection received"): + connect(load_balance_hosts="disable") + + +def test_load_balance_hosts_random_distribution(load_balance_nodes): + """load_balance_hosts=random distributes connections across all nodes.""" + nodes, connect = load_balance_nodes + + for _ in range(50): + connect(load_balance_hosts="random") + + occurrences = [ + len(re.findall("connection received", node.log_content())) for node in nodes + ] + + # Statistically, each node should receive at least one connection. + # The probability of any node receiving 0 connections is (2/3)^50 ≈ 1.57e-9 + assert occurrences[0] > 0, "node1 should receive at least one connection" + assert occurrences[1] > 0, "node2 should receive at least one connection" + assert occurrences[2] > 0, "node3 should receive at least one connection" + assert sum(occurrences) == 50, "total connections should be 50" + + +def test_load_balance_hosts_failover(load_balance_nodes): + """load_balance_hosts continues trying hosts until it finds a working one.""" + nodes, connect = load_balance_nodes + + nodes[0].stop() + nodes[1].stop() + + with nodes[2].log_contains("connection received"): + connect(load_balance_hosts="disable") + + with nodes[2].log_contains("connection received", times=5): + for _ in range(5): + connect(load_balance_hosts="random") diff --git a/src/interfaces/libpq/t/003_load_balance_host_list.pl b/src/interfaces/libpq/t/003_load_balance_host_list.pl deleted file mode 100644 index 7a4c14ada98..00000000000 --- a/src/interfaces/libpq/t/003_load_balance_host_list.pl +++ /dev/null @@ -1,94 +0,0 @@ -# Copyright (c) 2023-2025, PostgreSQL Global Development Group -use strict; -use warnings FATAL => 'all'; -use Config; -use PostgreSQL::Test::Utils; -use PostgreSQL::Test::Cluster; -use Test::More; - -# This tests load balancing across the list of different hosts in the host -# parameter of the connection string. - -# Cluster setup which is shared for testing both load balancing methods -my $node1 = PostgreSQL::Test::Cluster->new('node1'); -my $node2 = PostgreSQL::Test::Cluster->new('node2', own_host => 1); -my $node3 = PostgreSQL::Test::Cluster->new('node3', own_host => 1); - -# Create a data directory with initdb -$node1->init(); -$node2->init(); -$node3->init(); - -# Start the PostgreSQL server -$node1->start(); -$node2->start(); -$node3->start(); - -# Start the tests for load balancing method 1 -my $hostlist = $node1->host . ',' . $node2->host . ',' . $node3->host; -my $portlist = $node1->port . ',' . $node2->port . ',' . $node3->port; - -$node1->connect_fails( - "host=$hostlist port=$portlist load_balance_hosts=doesnotexist", - "load_balance_hosts doesn't accept unknown values", - expected_stderr => qr/invalid load_balance_hosts value: "doesnotexist"/); - -# load_balance_hosts=disable should always choose the first one. -$node1->connect_ok( - "host=$hostlist port=$portlist load_balance_hosts=disable", - "load_balance_hosts=disable connects to the first node", - sql => "SELECT 'connect1'", - log_like => [qr/statement: SELECT 'connect1'/]); - -# Statistically the following loop with load_balance_hosts=random will almost -# certainly connect at least once to each of the nodes. The chance of that not -# happening is so small that it's negligible: (2/3)^50 = 1.56832855e-9 -foreach my $i (1 .. 50) -{ - $node1->connect_ok( - "host=$hostlist port=$portlist load_balance_hosts=random", - "repeated connections with random load balancing", - sql => "SELECT 'connect2'"); -} - -my $node1_occurrences = () = - $node1->log_content() =~ /statement: SELECT 'connect2'/g; -my $node2_occurrences = () = - $node2->log_content() =~ /statement: SELECT 'connect2'/g; -my $node3_occurrences = () = - $node3->log_content() =~ /statement: SELECT 'connect2'/g; - -my $total_occurrences = - $node1_occurrences + $node2_occurrences + $node3_occurrences; - -cmp_ok($node1_occurrences, '>', 1, - "received at least one connection on node1"); -cmp_ok($node2_occurrences, '>', 1, - "received at least one connection on node2"); -cmp_ok($node3_occurrences, '>', 1, - "received at least one connection on node3"); -is($total_occurrences, 50, "received 50 connections across all nodes"); - -$node1->stop(); -$node2->stop(); - -# load_balance_hosts=disable should continue trying hosts until it finds a -# working one. -$node3->connect_ok( - "host=$hostlist port=$portlist load_balance_hosts=disable", - "load_balance_hosts=disable continues until it connects to the a working node", - sql => "SELECT 'connect3'", - log_like => [qr/statement: SELECT 'connect3'/]); - -# Also with load_balance_hosts=random we continue to the next nodes if previous -# ones are down. Connect a few times to make sure it's not just lucky. -foreach my $i (1 .. 5) -{ - $node3->connect_ok( - "host=$hostlist port=$portlist load_balance_hosts=random", - "load_balance_hosts=random continues until it connects to the a working node", - sql => "SELECT 'connect4'", - log_like => [qr/statement: SELECT 'connect4'/]); -} - -done_testing(); diff --git a/src/interfaces/libpq/t/004_load_balance_dns.pl b/src/interfaces/libpq/t/004_load_balance_dns.pl deleted file mode 100644 index 2b4bd261c3d..00000000000 --- a/src/interfaces/libpq/t/004_load_balance_dns.pl +++ /dev/null @@ -1,144 +0,0 @@ -# Copyright (c) 2023-2025, PostgreSQL Global Development Group -use strict; -use warnings FATAL => 'all'; -use Config; -use PostgreSQL::Test::Utils; -use PostgreSQL::Test::Cluster; -use Test::More; - -if (!$ENV{PG_TEST_EXTRA} || $ENV{PG_TEST_EXTRA} !~ /\bload_balance\b/) -{ - plan skip_all => - 'Potentially unsafe test load_balance not enabled in PG_TEST_EXTRA'; -} - -# This tests loadbalancing based on a DNS entry that contains multiple records -# for different IPs. Since setting up a DNS server is more effort than we -# consider reasonable to run this test, this situation is instead imitated by -# using a hosts file where a single hostname maps to multiple different IP -# addresses. This test requires the administrator to add the following lines to -# the hosts file (if we detect that this hasn't happened we skip the test): -# -# 127.0.0.1 pg-loadbalancetest -# 127.0.0.2 pg-loadbalancetest -# 127.0.0.3 pg-loadbalancetest -# -# Windows or Linux are required to run this test because these OSes allow -# binding to 127.0.0.2 and 127.0.0.3 addresses by default, but other OSes -# don't. We need to bind to different IP addresses, so that we can use these -# different IP addresses in the hosts file. -# -# The hosts file needs to be prepared before running this test. We don't do it -# on the fly, because it requires root permissions to change the hosts file. In -# CI we set up the previously mentioned rules in the hosts file, so that this -# load balancing method is tested. - -# Cluster setup which is shared for testing both load balancing methods -my $can_bind_to_127_0_0_2 = - $Config{osname} eq 'linux' || $PostgreSQL::Test::Utils::windows_os; - -# Checks for the requirements for testing load balancing method 2 -if (!$can_bind_to_127_0_0_2) -{ - plan skip_all => 'load_balance test only supported on Linux and Windows'; -} - -my $hosts_path; -if ($windows_os) -{ - $hosts_path = 'c:\Windows\System32\Drivers\etc\hosts'; -} -else -{ - $hosts_path = '/etc/hosts'; -} - -my $hosts_content = PostgreSQL::Test::Utils::slurp_file($hosts_path); - -my $hosts_count = () = - $hosts_content =~ /127\.0\.0\.[1-3] pg-loadbalancetest/g; -if ($hosts_count != 3) -{ - # Host file is not prepared for this test - plan skip_all => "hosts file was not prepared for DNS load balance test"; -} - -$PostgreSQL::Test::Cluster::use_tcp = 1; -$PostgreSQL::Test::Cluster::test_pghost = '127.0.0.1'; -my $port = PostgreSQL::Test::Cluster::get_free_port(); -my $node1 = PostgreSQL::Test::Cluster->new('node1', port => $port); -my $node2 = - PostgreSQL::Test::Cluster->new('node2', port => $port, own_host => 1); -my $node3 = - PostgreSQL::Test::Cluster->new('node3', port => $port, own_host => 1); - -# Create a data directory with initdb -$node1->init(); -$node2->init(); -$node3->init(); - -# Start the PostgreSQL server -$node1->start(); -$node2->start(); -$node3->start(); - -# load_balance_hosts=disable should always choose the first one. -$node1->connect_ok( - "host=pg-loadbalancetest port=$port load_balance_hosts=disable", - "load_balance_hosts=disable connects to the first node", - sql => "SELECT 'connect1'", - log_like => [qr/statement: SELECT 'connect1'/]); - - -# Statistically the following loop with load_balance_hosts=random will almost -# certainly connect at least once to each of the nodes. The chance of that not -# happening is so small that it's negligible: (2/3)^50 = 1.56832855e-9 -foreach my $i (1 .. 50) -{ - $node1->connect_ok( - "host=pg-loadbalancetest port=$port load_balance_hosts=random", - "repeated connections with random load balancing", - sql => "SELECT 'connect2'"); -} - -my $node1_occurrences = () = - $node1->log_content() =~ /statement: SELECT 'connect2'/g; -my $node2_occurrences = () = - $node2->log_content() =~ /statement: SELECT 'connect2'/g; -my $node3_occurrences = () = - $node3->log_content() =~ /statement: SELECT 'connect2'/g; - -my $total_occurrences = - $node1_occurrences + $node2_occurrences + $node3_occurrences; - -cmp_ok($node1_occurrences, '>', 1, - "received at least one connection on node1"); -cmp_ok($node2_occurrences, '>', 1, - "received at least one connection on node2"); -cmp_ok($node3_occurrences, '>', 1, - "received at least one connection on node3"); -is($total_occurrences, 50, "received 50 connections across all nodes"); - -$node1->stop(); -$node2->stop(); - -# load_balance_hosts=disable should continue trying hosts until it finds a -# working one. -$node3->connect_ok( - "host=pg-loadbalancetest port=$port load_balance_hosts=disable", - "load_balance_hosts=disable continues until it connects to the a working node", - sql => "SELECT 'connect3'", - log_like => [qr/statement: SELECT 'connect3'/]); - -# Also with load_balance_hosts=random we continue to the next nodes if previous -# ones are down. Connect a few times to make sure it's not just lucky. -foreach my $i (1 .. 5) -{ - $node3->connect_ok( - "host=pg-loadbalancetest port=$port load_balance_hosts=random", - "load_balance_hosts=random continues until it connects to the a working node", - sql => "SELECT 'connect4'", - log_like => [qr/statement: SELECT 'connect4'/]); -} - -done_testing(); -- 2.52.0