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

# Logical replication tests for
# ALTER COLUMN c ADD GENERATED ALWAYS STORED USING CONSTRAINT name
#
# This tests the configuration where the publication is set up to publish
# stored generated columns.
#
# When publishing stored generated columns, it is not supported for the same
# column to be generated on both the publisher and the subscriber. The only
# valid configuration is for the column to be a regular column on the side of
# the subscriber.
#
use strict;
use warnings FATAL => 'all';
use PostgreSQL::Test::Cluster;
use PostgreSQL::Test::Utils;
use Test::More;

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

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);
}

# Schema and replication setup
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);

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 = stored);
]);

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');

# Initial data
sql($node_publisher, qq[INSERT INTO sch1.tab1 (a) VALUES (1)]);
$node_publisher->wait_for_catchup('tap_sub_schema');

# Non-locking migration to add a stored generated column
# b = (a * 2)
sql($node_subscriber, qq[
	ALTER TABLE sch1.tab1 ADD COLUMN b INT;
]);
sql($node_publisher, qq[
	INSERT INTO sch1.tab1 (a) VALUES (2);
	ALTER TABLE sch1.tab1 ADD COLUMN b INT;
	INSERT INTO sch1.tab1 (a) VALUES (3);
]);

sql($node_publisher, qq[
	-- Take care of new and updated rows, first.
	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;

	INSERT INTO sch1.tab1 (a) VALUES (4);

	-- Now, backfill the table. In production, this might be done in batches.
	UPDATE sch1.tab1 SET b = a * 2 WHERE b IS NULL;

	ALTER TABLE sch1.tab1 VALIDATE CONSTRAINT check_gen;

	INSERT INTO sch1.tab1 (a) VALUES (5);

	-- Now, we can convert the column without a rewrite while holding an AccessExclusiveLock.
	ALTER TABLE sch1.tab1 ALTER COLUMN b ADD GENERATED ALWAYS STORED USING CONSTRAINT check_gen;

	INSERT INTO sch1.tab1 (a) VALUES (6);
]);

$node_publisher->wait_for_catchup('tap_sub_schema');

sql($node_subscriber, qq[
	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 (7)]);

$node_publisher->wait_for_catchup('tap_sub_schema');

my $result = sql($node_subscriber, qq[SELECT a, b FROM sch1.tab1 ORDER BY a]);
is($result, qq[1|2
2|4
3|6
4|8
5|10
6|12
7|14], 'check: fully replicated');

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

done_testing();