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

# Check that query_id is advertised as expected. This requires a concurrent
# session to check for the last query_id advertised by another session.
use strict;
use warnings FATAL => 'all';
use PostgreSQL::Test::Cluster;
use PostgreSQL::Test::Utils;
use Test::More;

# declare global variables
my $ret;
my $node;
my $psql_session1;
my $psql_session2;

# declare the helper functions
sub test_num_unique_qids
{
    my ($command, $expected_qids, $wait_for) = @_;

    $psql_session1->query_safe($command);

    $psql_session2->query_safe(q(INSERT INTO track_query_id 
                                 SELECT query_id FROM pg_stat_activity 
                                 WHERE application_name = 'query_id_test'
                                 AND rtrim(query) = 'SELECT * FROM test_query_id'
                                 AND query_id is NOT NULL and query_id <> 0;));

    $ret = $psql_session2->query_safe(
	    qq[
            SELECT COUNT(DISTINCT query_id) FROM track_query_id;
    ]);
    
    is($ret, $expected_qids, "# of distinct queryIds");
}

sub create_track_qid_table
{
    $node->safe_psql("postgres", qq[CREATE TABLE track_query_id (query_id BIGINT);]);
}

sub create_test_qid_table
{
    $node->safe_psql("postgres", qq[DROP TABLE IF EXISTS test_query_id;]);
    $node->safe_psql("postgres", qq[CREATE TABLE test_query_id (id INT);]);
}

# Setup the test
## create the node
$node = PostgreSQL::Test::Cluster->new('main');
$node->init;
$node->append_conf('postgresql.conf', 'compute_query_id = on');
$node->start;

## create the sessions
## session1 will run the queries and session2 will query pg_stat_activity 
## for the last query_id of session1 and write it to track_query_id
$psql_session1 = $node->background_psql('postgres');
$psql_session1->query_until(qr//, q(SET application_name = query_id_test;));
$psql_session2 = $node->background_psql('postgres');

create_track_qid_table();

# test 1 - check post-parse queryId is reported with extended query protocol
# We expected 1 queryId to be reported so far.
create_test_qid_table();
test_num_unique_qids("SELECT * FROM test_query_id \\parse n1\n", 1, qr//);

# test 2 - check bind-execute queryId is reported in extended query protocol
# Since we recreate the table of the underlysing named portal, we now expect
# 2 queryIds to be reported.
create_test_qid_table();
test_num_unique_qids("\\bind_named n1 \\g\n", 2, qr//);

# Destroy the sessions and the node
$psql_session1->quit;
$psql_session2->quit;
$node->stop;

done_testing();