# Copyright (c) 2021-2025, PostgreSQL Global Development Group

# Tests that standbys:
# - drop stats for objects when the those records are replayed
# - persist stats across graceful restarts
# - discard stats after immediate / crash restarts

use strict;
use warnings FATAL => 'all';
use PostgreSQL::Test::Cluster;
use PostgreSQL::Test::Utils;
use Test::More;

my $node_primary = PostgreSQL::Test::Cluster->new('primary');
$node_primary->init(allows_streaming => 1);
$node_primary->append_conf('postgresql.conf', "track_functions = 'all'");
$node_primary->start;

my $backup_name = 'my_backup';
$node_primary->backup($backup_name);

my $node_standby = PostgreSQL::Test::Cluster->new('standby');
$node_standby->init_from_backup($node_primary, $backup_name,
	has_streaming => 1);
$node_standby->start;


## Test that stats are cleaned up on standby after dropping table or function

my $sect = 'initial';

my ($dboid, $tableoid, $funcoid) =
  populate_standby_stats('postgres', 'public');
test_standby_func_tab_stats_status('postgres',
	$dboid, $tableoid, $funcoid, 't');

drop_table_by_oid('postgres', $tableoid);
drop_function_by_oid('postgres', $funcoid);

$sect = 'post drop';
$node_primary->wait_for_replay_catchup($node_standby);
test_standby_func_tab_stats_status('postgres',
	$dboid, $tableoid, $funcoid, 'f');


## Test that stats are cleaned up on standby after dropping indirectly

$sect = "schema creation";

$node_primary->safe_psql('postgres', "CREATE SCHEMA drop_schema_test1");
$node_primary->wait_for_replay_catchup($node_standby);

($dboid, $tableoid, $funcoid) =
  populate_standby_stats('postgres', 'drop_schema_test1');

test_standby_func_tab_stats_status('postgres',
	$dboid, $tableoid, $funcoid, 't');
$node_primary->safe_psql('postgres', "DROP SCHEMA drop_schema_test1 CASCADE");

$sect = "post schema drop";

$node_primary->wait_for_replay_catchup($node_standby);

# verify table and function stats removed from standby
test_standby_func_tab_stats_status('postgres',
	$dboid, $tableoid, $funcoid, 'f');


## Test that stats are cleaned up on standby after dropping database

$sect = "createdb";

$node_primary->safe_psql('postgres', "CREATE DATABASE test");
$node_primary->wait_for_replay_catchup($node_standby);

($dboid, $tableoid, $funcoid) = populate_standby_stats('test', 'public');

# verify stats are present
test_standby_func_tab_stats_status('test', $dboid, $tableoid, $funcoid, 't');
test_standby_db_stats_status('test', $dboid, 't');

$node_primary->safe_psql('postgres', "DROP DATABASE test");
$sect = "post dropdb";
$node_primary->wait_for_replay_catchup($node_standby);

# Test that the stats were cleaned up on standby
# Note that this connects to 'postgres' but provides the dboid of dropped db
# 'test' which we acquired previously
test_standby_func_tab_stats_status('postgres',
	$dboid, $tableoid, $funcoid, 'f');

test_standby_db_stats_status('postgres', $dboid, 'f');


## verify that stats persist across graceful restarts on a replica

# NB: Can't test database stats, they're immediately repopulated when
# reconnecting...
$sect = "pre restart";
($dboid, $tableoid, $funcoid) = populate_standby_stats('postgres', 'public');
test_standby_func_tab_stats_status('postgres',
	$dboid, $tableoid, $funcoid, 't');

$node_standby->restart();

$sect = "post non-immediate";

test_standby_func_tab_stats_status('postgres',
	$dboid, $tableoid, $funcoid, 't');

# but gone after an immediate restart
$node_standby->stop('immediate');
$node_standby->start();

$sect = "post immediate restart";

test_standby_func_tab_stats_status('postgres',
	$dboid, $tableoid, $funcoid, 'f');


done_testing();


sub populate_standby_stats
{
	my ($connect_db, $schema) = @_;

	# create objects on primary
	$node_primary->safe_psql($connect_db,
		"CREATE TABLE $schema.drop_tab_test1 AS SELECT generate_series(1,100) AS a"
	);
	$node_primary->safe_psql($connect_db,
		"CREATE FUNCTION $schema.drop_func_test1() RETURNS VOID AS 'select 2;' LANGUAGE SQL IMMUTABLE"
	);
	$node_primary->wait_for_replay_catchup($node_standby);

	# collect object oids
	my $dboid = $node_standby->safe_psql($connect_db,
		"SELECT oid FROM pg_database WHERE datname = '$connect_db'");
	my $tableoid = $node_standby->safe_psql($connect_db,
		"SELECT '$schema.drop_tab_test1'::regclass::oid");
	my $funcoid = $node_standby->safe_psql($connect_db,
		"SELECT '$schema.drop_func_test1()'::regprocedure::oid");

	# generate stats on standby
	$node_standby->safe_psql($connect_db,
		"SELECT * FROM $schema.drop_tab_test1");
	$node_standby->safe_psql($connect_db, "SELECT $schema.drop_func_test1()");

	return ($dboid, $tableoid, $funcoid);
}

sub drop_function_by_oid
{
	my ($connect_db, $funcoid) = @_;

	# Get function name from returned oid
	my $func_name = $node_primary->safe_psql($connect_db,
		"SELECT '$funcoid'::regprocedure");
	$node_primary->safe_psql($connect_db, "DROP FUNCTION $func_name");
}

sub drop_table_by_oid
{
	my ($connect_db, $tableoid) = @_;

	# Get table name from returned oid
	my $table_name =
	  $node_primary->safe_psql($connect_db, "SELECT '$tableoid'::regclass");
	$node_primary->safe_psql($connect_db, "DROP TABLE $table_name");
}

sub test_standby_func_tab_stats_status
{
	local $Test::Builder::Level = $Test::Builder::Level + 1;
	my ($connect_db, $dboid, $tableoid, $funcoid, $present) = @_;

	my %expected = (rel => $present, func => $present);
	my %stats;

	$stats{rel} = $node_standby->safe_psql($connect_db,
		"SELECT pg_stat_have_stats('relation', $dboid, $tableoid)");
	$stats{func} = $node_standby->safe_psql($connect_db,
		"SELECT pg_stat_have_stats('function', $dboid, $funcoid)");

	is_deeply(\%stats, \%expected, "$sect: standby stats as expected");

	return;
}

sub test_standby_db_stats_status
{
	local $Test::Builder::Level = $Test::Builder::Level + 1;
	my ($connect_db, $dboid, $present) = @_;

	is( $node_standby->safe_psql(
			$connect_db, "SELECT pg_stat_have_stats('database', $dboid, 0)"),
		$present,
		"$sect: standby db stats as expected");
}