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

# Logical replication tests for
# ALTER COLUMN c ADD GENERATED ALWAYS STORED USING CONSTRAINT name
#
# This tries to exercise the replication code paths when the column is
# converted to be a stored generated column and the publication is set up to
# *not* publish them. The scenario is what a DBA would do to add a stored
# generated column without taking the table offline.
#
use strict;
use warnings FATAL => 'all';
use PostgreSQL::Test::Cluster;
use PostgreSQL::Test::Utils;
use Test::More;

# Initialize publisher node
my $node_publisher = PostgreSQL::Test::Cluster->new('publisher');
$node_publisher->init(allows_streaming => 'logical');
$node_publisher->start;

# Create subscriber node
my $node_subscriber = PostgreSQL::Test::Cluster->new('subscriber');
$node_subscriber->init;
$node_subscriber->start;

# We will use the same user throughout the test, so let's just fix it here.
sub sql
{
	local $Carp::CarpLevel = $Carp::CarpLevel + 1;
	my ($node, $sql_code) = @_;
	$node->safe_psql('postgres', $sql_code);
}

my $schema_ddl = qq[
	CREATE SCHEMA sch1;
	CREATE TABLE sch1.tab1 (a INT PRIMARY KEY);
];
sql($node_publisher, $schema_ddl);
sql($node_subscriber, $schema_ddl);

sql($node_publisher, qq[
	INSERT INTO sch1.tab1 (a) values (1);
]);

# Set up replication
my $publisher_connstr = $node_publisher->connstr . ' dbname=postgres';
sql($node_publisher, qq[
	CREATE PUBLICATION tap_pub_schema FOR TABLES IN SCHEMA sch1
	WITH (publish_generated_columns = none)
]);
sql($node_subscriber, qq[
	CREATE SUBSCRIPTION tap_sub_schema CONNECTION '$publisher_connstr'
	    PUBLICATION tap_pub_schema
]);
$node_subscriber->wait_for_subscription_sync($node_publisher,
	'tap_sub_schema');

my $result = sql($node_subscriber, "SELECT a FROM sch1.tab1");
is($result, "1",
	'sanity check: initial data has been synced');

# Now we want to add a stored generated column `b`. Since we are not set up
# to publish stored generated column, we need to set up the subscriber first.
sql($node_subscriber, qq[
	ALTER TABLE sch1.tab1 ADD COLUMN b INT;
]);
sql($node_publisher, qq[
	ALTER TABLE sch1.tab1 ADD COLUMN b INT;
]);

sql($node_publisher, qq[
	INSERT INTO sch1.tab1 (a) values (2);
]);

# Take care of new/updated rows
sql($node_publisher, qq[
	CREATE FUNCTION sch1.generate_b () RETURNS TRIGGER LANGUAGE plpgsql AS \$\$
	BEGIN
	  NEW.b = NEW.a * 2; RETURN NEW;
	END
	\$\$;
	CREATE TRIGGER trig_gen BEFORE INSERT OR UPDATE ON sch1.tab1
		FOR EACH ROW EXECUTE FUNCTION sch1.generate_b();

	ALTER TABLE sch1.tab1 ADD CONSTRAINT check_gen CHECK (b IS NOT DISTINCT FROM a * 2) NOT VALID;
]);

sql($node_publisher, qq[
	INSERT INTO sch1.tab1 (a) values (3);
]);

# Backfill
sql($node_publisher, qq[
	UPDATE sch1.tab1 SET b = a * 2 WHERE b IS NULL;
	ALTER TABLE sch1.tab1 VALIDATE CONSTRAINT check_gen;
]);
$node_publisher->wait_for_catchup('tap_sub_schema');

# When we switch "b" to a stored gen column on the publisher, it will not be
# synced anymore. Let's set up the replica first, in order to not lose data.
sql($node_subscriber, qq[
	CREATE FUNCTION sch1.generate_b () RETURNS TRIGGER LANGUAGE plpgsql AS \$\$
	BEGIN
	  NEW.b = NEW.a * 2; RETURN NEW;
	END
	\$\$;
	CREATE TRIGGER trig_gen BEFORE INSERT OR UPDATE ON sch1.tab1
		FOR EACH ROW EXECUTE FUNCTION sch1.generate_b();
    ALTER TABLE sch1.tab1 ENABLE REPLICA TRIGGER trig_gen;

	ALTER TABLE sch1.tab1 ADD CONSTRAINT check_gen CHECK (b IS NOT DISTINCT FROM a * 2) NOT VALID;
	ALTER TABLE sch1.tab1 VALIDATE CONSTRAINT check_gen;
]);

sql($node_publisher, qq[
	INSERT INTO sch1.tab1 (a) values (4);
]);

# Convert b to a stored generated column on the publisher first: b will not
# be synced anymore, but the trigger and constraint on the subscriber guarantee
# that it writes the correct values.
sql($node_publisher, qq[
	ALTER TABLE sch1.tab1 ALTER COLUMN b ADD GENERATED ALWAYS STORED USING CONSTRAINT check_gen;
]);
sql($node_publisher, qq[
	INSERT INTO sch1.tab1 (a) values (5);
]);
$node_publisher->wait_for_catchup('tap_sub_schema');
sql($node_subscriber, qq[
	ALTER TABLE sch1.tab1 ALTER COLUMN b ADD GENERATED ALWAYS STORED USING CONSTRAINT check_gen;
]);

sql($node_publisher, qq[
	INSERT INTO sch1.tab1 (a) values (6);
]);

$node_publisher->wait_for_catchup('tap_sub_schema');

# Check: all the rows have the correct values in b
$result = sql($node_subscriber, "SELECT * FROM sch1.tab1 ORDER BY a");
is($result, qq[1|2
2|4
3|6
4|8
5|10
6|12], 'check: all of "b" has been replicated');

$node_subscriber->stop('fast');
$node_publisher->stop('fast');

done_testing();