Files
pgcat/tests/ruby/auth_query_spec.rb
Mostafa Abdelraouf 0b01d70b55 Allow configuring routing decision when no shard is selected (#578)
The TL;DR for the change is that we allow QueryRouter to set the active shard to None. This signals to the Pool::get method that we have no shard selected. The get method follows a no_shard_specified_behavior config to know how to route the query.

Original PR description
Ruby-pg library makes a startup query to SET client_encoding to ... if Encoding.default_internal value is set (Code). This query is troublesome because we cannot possibly attach a routing comment to it. PgCat, by default, will route that query to the default shard.

Everything is fine until shard 0 has issues, Clients will all be attempting to send this query to shard0 which increases the connection latency significantly for all clients, even those not interested in shard0

This PR introduces no_shard_specified_behavior that defines the behavior in case we have routing-by-comment enabled but we get a query without a comment. The allowed behaviors are

random: Picks a shard at random
random_healthy: Picks a shard at random favoring shards with the least number of recent connection/checkout errors
shard_<number>: e.g. shard_0, shard_4, etc. picks a specific shard, everytime
In order to achieve this, this PR introduces an error_count on the Address Object that tracks the number of errors since the last checkout and uses that metric to sort shards by error count before making a routing decision.
I didn't want to use address stats to avoid introducing a routing dependency on internal stats (We might do that in the future but I prefer to avoid this for the time being.

I also made changes to the test environment to replace Ruby's TOML reader library, It appears to be abandoned and does not support mixed arrays (which we use in the config toml), and it also does not play nicely with single-quoted regular expressions. I opted for using yj which is a CLI tool that can convert from toml to JSON and back. So I refactor the tests to use that library.
2023-09-11 13:47:28 -05:00

216 lines
9.0 KiB
Ruby

# frozen_string_literal: true
require_relative 'spec_helper'
require_relative 'helpers/auth_query_helper'
describe "Auth Query" do
let(:configured_instances) {[5432, 10432]}
let(:config_user) { { 'username' => 'sharding_user', 'password' => 'sharding_user' } }
let(:pg_user) { { 'username' => 'sharding_user', 'password' => 'sharding_user' } }
let(:processes) { Helpers::AuthQuery.single_shard_auth_query(pool_name: "sharded_db", pg_user: pg_user, config_user: config_user, extra_conf: config, wait_until_ready: wait_until_ready ) }
let(:config) { {} }
let(:wait_until_ready) { true }
after do
unless @failing_process
processes.all_databases.map(&:reset)
processes.pgcat.shutdown
end
@failing_process = false
end
context "when auth_query is not configured" do
context 'and cleartext passwords are set' do
it "uses local passwords" do
conn = PG.connect(processes.pgcat.connection_string("sharded_db", config_user['username'], config_user['password']))
expect(conn.async_exec("SELECT 1 + 2")).not_to be_nil
end
end
context 'and cleartext passwords are not set' do
let(:config_user) { { 'username' => 'sharding_user' } }
it "does not start because it is not possible to authenticate" do
@failing_process = true
expect { processes.pgcat }.to raise_error(StandardError, /You have to specify a user password for every pool if auth_query is not specified/)
end
end
end
context 'when auth_query is configured' do
context 'with global configuration' do
around(:example) do |example|
# Set up auth query
Helpers::AuthQuery.set_up_auth_query_for_user(
user: 'md5_auth_user',
password: 'secret'
);
example.run
# Drop auth query support
Helpers::AuthQuery.tear_down_auth_query_for_user(
user: 'md5_auth_user',
password: 'secret'
);
end
context 'with correct global parameters' do
let(:config) { { 'general' => { 'auth_query' => "SELECT * FROM public.user_lookup('$1');", 'auth_query_user' => 'md5_auth_user', 'auth_query_password' => 'secret' } } }
context 'and with cleartext passwords set' do
it 'it uses local passwords' do
conn = PG.connect(processes.pgcat.connection_string("sharded_db", pg_user['username'], pg_user['password']))
expect(conn.exec("SELECT 1 + 2")).not_to be_nil
end
end
context 'and with cleartext passwords not set' do
let(:config_user) { { 'username' => 'sharding_user', 'password' => 'sharding_user' } }
it 'it uses obtained passwords' do
connection_string = processes.pgcat.connection_string("sharded_db", pg_user['username'], pg_user['password'])
conn = PG.connect(connection_string)
expect(conn.async_exec("SELECT 1 + 2")).not_to be_nil
end
it 'allows passwords to be changed without closing existing connections' do
pgconn = PG.connect(processes.pgcat.connection_string("sharded_db", pg_user['username']))
expect(pgconn.exec("SELECT 1 + 2")).not_to be_nil
Helpers::AuthQuery.exec_in_instances(query: "ALTER USER #{pg_user['username']} WITH ENCRYPTED PASSWORD 'secret2';")
expect(pgconn.exec("SELECT 1 + 4")).not_to be_nil
Helpers::AuthQuery.exec_in_instances(query: "ALTER USER #{pg_user['username']} WITH ENCRYPTED PASSWORD '#{pg_user['password']}';")
end
it 'allows passwords to be changed and that new password is needed when reconnecting' do
pgconn = PG.connect(processes.pgcat.connection_string("sharded_db", pg_user['username']))
expect(pgconn.exec("SELECT 1 + 2")).not_to be_nil
Helpers::AuthQuery.exec_in_instances(query: "ALTER USER #{pg_user['username']} WITH ENCRYPTED PASSWORD 'secret2';")
newconn = PG.connect(processes.pgcat.connection_string("sharded_db", pg_user['username'], 'secret2'))
expect(newconn.exec("SELECT 1 + 2")).not_to be_nil
Helpers::AuthQuery.exec_in_instances(query: "ALTER USER #{pg_user['username']} WITH ENCRYPTED PASSWORD '#{pg_user['password']}';")
end
end
end
context 'with wrong parameters' do
let(:config) { { 'general' => { 'auth_query' => 'SELECT 1', 'auth_query_user' => 'wrong_user', 'auth_query_password' => 'wrong' } } }
context 'and with clear text passwords set' do
it "it uses local passwords" do
conn = PG.connect(processes.pgcat.connection_string("sharded_db", pg_user['username'], pg_user['password']))
expect(conn.async_exec("SELECT 1 + 2")).not_to be_nil
end
end
context 'and with cleartext passwords not set' do
let(:config_user) { { 'username' => 'sharding_user' } }
it "it fails to start as it cannot authenticate against servers" do
@failing_process = true
expect { PG.connect(processes.pgcat.connection_string("sharded_db", pg_user['username'], pg_user['password'])) }.to raise_error(StandardError, /Error trying to obtain password from auth_query/ )
end
context 'and we fix the issue and reload' do
let(:wait_until_ready) { false }
it 'fails in the beginning but starts working after reloading config' do
connection_string = processes.pgcat.connection_string("sharded_db", pg_user['username'], pg_user['password'])
while !(processes.pgcat.logs =~ /Waiting for clients/) do
sleep 0.5
end
expect { PG.connect(connection_string)}.to raise_error(PG::ConnectionBad)
expect(processes.pgcat.logs).to match(/Error trying to obtain password from auth_query/)
current_config = processes.pgcat.current_config
config = { 'general' => { 'auth_query' => "SELECT * FROM public.user_lookup('$1');", 'auth_query_user' => 'md5_auth_user', 'auth_query_password' => 'secret' } }
processes.pgcat.update_config(current_config.deep_merge(config))
processes.pgcat.reload_config
conn = nil
expect { conn = PG.connect(connection_string)}.not_to raise_error
expect(conn.async_exec("SELECT 1 + 2")).not_to be_nil
end
end
end
end
end
context 'with per pool configuration' do
around(:example) do |example|
# Set up auth query
Helpers::AuthQuery.set_up_auth_query_for_user(
user: 'md5_auth_user',
password: 'secret'
);
Helpers::AuthQuery.set_up_auth_query_for_user(
user: 'md5_auth_user1',
password: 'secret',
database: 'shard1'
);
example.run
# Tear down auth query
Helpers::AuthQuery.tear_down_auth_query_for_user(
user: 'md5_auth_user',
password: 'secret'
);
Helpers::AuthQuery.tear_down_auth_query_for_user(
user: 'md5_auth_user1',
password: 'secret',
database: 'shard1'
);
end
context 'with correct parameters' do
let(:processes) { Helpers::AuthQuery.two_pools_auth_query(pool_names: ["sharded_db0", "sharded_db1"], pg_user: pg_user, config_user: config_user, extra_conf: config ) }
let(:config) {
{ 'pools' =>
{
'sharded_db0' => {
'auth_query' => "SELECT * FROM public.user_lookup('$1');",
'auth_query_user' => 'md5_auth_user',
'auth_query_password' => 'secret'
},
'sharded_db1' => {
'auth_query' => "SELECT * FROM public.user_lookup('$1');",
'auth_query_user' => 'md5_auth_user1',
'auth_query_password' => 'secret'
},
}
}
}
context 'and with cleartext passwords set' do
it 'it uses local passwords' do
conn = PG.connect(processes.pgcat.connection_string("sharded_db0", pg_user['username'], pg_user['password']))
expect(conn.exec("SELECT 1 + 2")).not_to be_nil
conn = PG.connect(processes.pgcat.connection_string("sharded_db1", pg_user['username'], pg_user['password']))
expect(conn.exec("SELECT 1 + 2")).not_to be_nil
end
end
context 'and with cleartext passwords not set' do
let(:config_user) { { 'username' => 'sharding_user' } }
it 'it uses obtained passwords' do
connection_string = processes.pgcat.connection_string("sharded_db0", pg_user['username'], pg_user['password'])
conn = PG.connect(connection_string)
expect(conn.async_exec("SELECT 1 + 2")).not_to be_nil
connection_string = processes.pgcat.connection_string("sharded_db1", pg_user['username'], pg_user['password'])
conn = PG.connect(connection_string)
expect(conn.async_exec("SELECT 1 + 2")).not_to be_nil
end
end
end
end
end
end