-- PostgreSQL 17 introduces MAINTAIN privilege, this change the default output
-- of some test queries.
--
-- Alternative "grant_on_table_propagation_0.out" and this comment can be
-- removed when support for PostgreSQL versions < 17 is dropped.

SHOW server_version \gset
SELECT substring(:'server_version', '\d+')::int >= 17 AS server_version_ge_17;

--
-- GRANT_ON_TABLE_PROPAGATION
--
SET citus.shard_replication_factor TO 1;
CREATE SCHEMA grant_on_table;
SET search_path TO grant_on_table;

-- create some simple tables: 1 local on all nodes and 2 managed by citus
-- null_privs ACL must not be updated in anyway.
CREATE TABLE dist_table (id bigint GENERATED BY DEFAULT AS IDENTITY primary key
    , test_a int   -- test for INSERT
    , test_r text  -- test for SELECT
    , test_w text  -- test for UPDATE
    , test_mix int -- test for ALL/mixed
    , null_privs text
    );
SELECT create_distributed_table('dist_table', 'id');

CREATE TABLE ref_table (id bigint GENERATED BY DEFAULT AS IDENTITY primary key
    , test_a int   -- test for INSERT
    , test_r text  -- test for SELECT
    , test_w text  -- test for UPDATE
    , test_mix int -- test for ALL
    , null_privs text
    );
SELECT create_reference_table('ref_table');

SELECT result FROM run_command_on_all_nodes('CREATE TABLE grant_on_table.local_table (id int GENERATED BY DEFAULT AS IDENTITY primary key    , test_a int, test_r text, test_w text, test_mix int, null_privs text)');

-- queries used to check privileges:
\set verify_grant_table 'SELECT nodeport, unnest(result::text[]) FROM run_command_on_all_nodes($$SELECT array_agg((relname, relacl) order by 1) FROM pg_class WHERE relname IN (''dist_table'', ''ref_table'', ''local_table'') $$) join pg_dist_node using (nodeid) ORDER BY 2, 1'

\set verify_grant_attributes 'SELECT nodeport, unnest(coalesce(nullif(result, ''''), ''{}'')::text[]) FROM run_command_on_all_nodes($$SELECT array_agg((attrelid::regclass, attname, attacl) order by 1, 2) FROM pg_attribute WHERE attrelid  IN (''grant_on_table.dist_table''::regclass, ''grant_on_table.ref_table''::regclass, ''grant_on_table.local_table''::regclass) AND attacl IS NOT NULL $$) join pg_dist_node using (nodeid) ORDER BY 2, 1'

-- create some users
CREATE USER grant_user_0;
GRANT USAGE ON SCHEMA grant_on_table TO grant_user_0;
CREATE USER grant_user_1;
GRANT USAGE ON SCHEMA grant_on_table TO grant_user_1;
-- this one should not be granted anything:
CREATE USER nogrant_user;
GRANT USAGE ON SCHEMA grant_on_table TO nogrant_user;

--
-- tests related to columns ACL
--

--
-- when executing a table level grant, "postgres" is add/listed in pg_class.relacl
-- but nothing happens with revoke, as a result pg_class.relacl is not stable.
-- in order to have immutable cleanup results in those tests, init the
-- "postgres" special case
GRANT SELECT ON ref_table TO grant_user_0;
REVOKE SELECT ON ref_table FROM grant_user_0;

--
-- check we are able to propagate a single attribute privilege
-- we use only SELECT
--
SET ROLE grant_user_0;
-- not granted yet:
SELECT test_r FROM ref_table;
RESET ROLE;
GRANT SELECT (test_r) ON ref_table TO grant_user_0;
SET ROLE grant_user_0;
-- granted:
SELECT test_r FROM ref_table;
-- not granted:
SELECT test_a FROM ref_table;
SET ROLE grant_user_1;
-- not granted:
SELECT test_r FROM ref_table;
RESET ROLE;

-- check on coordinator and workers
:verify_grant_table ;
:verify_grant_attributes ;

-- cleanup
REVOKE SELECT (test_r) ON ref_table FROM grant_user_0;

-- check on coordinator and workers
:verify_grant_table ;
:verify_grant_attributes ;

--
-- check we are able to propagate a privilege to multiple attributes, users and tables at once
-- we use only INSERT/UPDATE
--
SET ROLE grant_user_0;
-- not granted yet:
INSERT INTO ref_table (test_a) VALUES (1);
RESET ROLE;
-- we would prefer not have to grant INSERT (id) as expected in the standard but
-- Citus rewrite queries with such attributes, which prevent standard to be applied.
GRANT INSERT (id, test_a), UPDATE (test_w, test_mix) ON ref_table, dist_table TO grant_user_0, grant_user_1;
SET ROLE grant_user_0;
-- granted:
INSERT INTO ref_table (test_a) VALUES (1);
-- granted:
UPDATE ref_table SET test_w = 2, test_mix = 2;
-- not granted:
INSERT INTO ref_table (test_w) VALUES (1);
SET ROLE grant_user_1;
-- granted:
INSERT INTO dist_table (test_a) VALUES (1);
-- granted:
UPDATE dist_table SET test_w = 3, test_mix = 3;
RESET ROLE;

-- check on coordinator and workers
:verify_grant_table ;
:verify_grant_attributes ;

-- cleanup
REVOKE INSERT (id, test_a), UPDATE (test_w, test_mix) ON ref_table, dist_table FROM grant_user_0, grant_user_1;

-- check on coordinator and workers
:verify_grant_table ;
:verify_grant_attributes ;

--
-- check we are able to propagate a table privilege associated with an attribute level privilege
-- we use only SELECT/DELETE
--
SET ROLE grant_user_0;
-- not granted yet:
SELECT test_r, test_mix FROM ref_table;
RESET ROLE;
GRANT SELECT (test_r, test_mix), DELETE ON ref_table TO grant_user_0, grant_user_1;
SET ROLE grant_user_0;
-- granted:
SELECT test_r, test_mix FROM ref_table;
-- granted:
DELETE FROM ref_table;
-- not granted:
SELECT test_a FROM ref_table;
-- not granted:
UPDATE ref_table SET null_privs = 3;
RESET ROLE;

-- check on coordinator and workers
:verify_grant_table ;
:verify_grant_attributes ;

-- cleanup
REVOKE SELECT (test_r, test_mix), DELETE ON ref_table FROM grant_user_0, grant_user_1;

-- check on coordinator and workers
:verify_grant_table ;
:verify_grant_attributes ;

--
-- check we also propagate system columns
-- we use only SELECT
--
SET ROLE grant_user_0;
-- not granted yet:
SELECT ctid, xmin FROM ref_table;
RESET ROLE;
GRANT SELECT (ctid, xmin) ON ref_table TO grant_user_0;
SET ROLE grant_user_0;
-- granted:
SELECT ctid, xmin FROM ref_table;
-- not granted:
SELECT ctid, test_a FROM ref_table;
RESET ROLE;

-- check on coordinator and workers
:verify_grant_table ;
:verify_grant_attributes ;

-- cleanup
REVOKE SELECT (ctid, xmin) ON ref_table FROM grant_user_0;

-- check on coordinator and workers
:verify_grant_table ;
:verify_grant_attributes ;

--
-- check we correctly propagate ALL, which has a few special cases
-- we use only ALL/SELECT
--
SET ROLE grant_user_0;
-- not granted yet:
INSERT INTO ref_table (id) VALUES (13);
SET ROLE grant_user_1;
-- not granted yet:
INSERT INTO ref_table (id, test_mix) VALUES (9, 3);
RESET ROLE;
GRANT ALL (id) ON ref_table TO grant_user_0;
GRANT ALL (id, test_mix) ON ref_table TO grant_user_1;
-- ALL cannot be mixed with other privs
-- should error:
GRANT SELECT (null_privs), ALL (null_privs) ON ref_table TO nogrant_user;
-- should error:
GRANT ALL (null_privs), SELECT (null_privs) ON ref_table TO nogrant_user;
SET ROLE grant_user_0;
-- granted:
INSERT INTO ref_table (id) VALUES (13);
SET ROLE grant_user_1;
-- granted:
INSERT INTO ref_table (id, test_mix) VALUES (9, 3);
-- not granted:
INSERT INTO ref_table (null_privs) VALUES (3);
-- granted:
SELECT id, test_mix FROM ref_table;
-- not granted:
SELECT null_privs FROM ref_table;
-- not granted:
DELETE FROM ref_table;
RESET ROLE;

-- check on coordinator and workers
:verify_grant_table ;
:verify_grant_attributes ;

-- cleanup
REVOKE ALL (id) ON ref_table FROM grant_user_0;
REVOKE ALL (id, test_mix) ON ref_table FROM grant_user_1;
TRUNCATE ref_table;

-- check on coordinator and workers
:verify_grant_table ;
:verify_grant_attributes ;

--
-- check we correctly propagate when mixed with local table, but only on the
-- current local table, not others
-- we use only INSERT/SELECT
--
SET ROLE grant_user_0;
-- not granted yet:
SELECT id FROM ref_table
UNION ALL SELECT id FROM dist_table
UNION ALL SELECT id FROM local_table;
RESET ROLE;
-- NOTE:
-- test special case: ALL TABLES IN SCHEMA is not supposed to be correct
-- but is accepted by PostgreSQL - non documented feature
GRANT SELECT (id) ON ALL TABLES IN SCHEMA grant_on_table TO grant_user_0;
-- check non propagation for local table (we'll just check ACL later, no INSERT testing)
GRANT INSERT (test_mix) ON local_table TO grant_user_0;
-- check we can propagate also when mixed with distributed table:
GRANT SELECT (test_r, test_mix) ON local_table, dist_table TO grant_user_1;
SET ROLE grant_user_0;
-- granted:
SELECT id FROM ref_table
UNION ALL SELECT id FROM dist_table
UNION ALL SELECT id FROM local_table;
SET ROLE grant_user_1;
-- granted:
SELECT test_r, test_mix FROM dist_table
UNION ALL SELECT test_r, test_mix FROM local_table;
RESET ROLE;

-- check on coordinator and workers
-- we pay special attention to local_table privileges here:
:verify_grant_table ;
:verify_grant_attributes ;

-- cleanup
REVOKE SELECT (id) ON ALL TABLES IN SCHEMA grant_on_table FROM grant_user_0;
-- check non propagation for local table (we'll just check ACL later, no INSERT testing)
REVOKE INSERT (test_mix) ON local_table FROM grant_user_0;
-- check we can propagate also when mixed with distributed table:
REVOKE SELECT (test_r, test_mix) ON local_table, dist_table FROM grant_user_1;

-- check on coordinator and workers
:verify_grant_table ;
:verify_grant_attributes ;

--
-- check TRUNCATE is not propagated (inccorect grammar)
-- also ensure no privs are propagated at all with "partially" incorrect grammar
-- we use only TRUNCATE/SELECT
--
SET ROLE grant_user_0;
-- not granted yet:
SELECT  test_r FROM ref_table;
RESET ROLE;
GRANT TRUNCATE (null_privs), SELECT (null_privs) ON ref_table TO nogrant_user;
SET ROLE grant_user_0;
-- still not granted:
SELECT test_r FROM ref_table;
-- still not granted:
TRUNCATE ref_table;
RESET ROLE;

-- check on coordinator and workers
:verify_grant_table ;
:verify_grant_attributes ;

-- no cleanup required

--
-- check we do not propage from a worker
-- we use only SELECT
--
\c - - - :worker_1_port
SET search_path TO grant_on_table;
GRANT SELECT (test_r, test_mix) ON ref_table TO grant_user_0;
\c - - - :master_port
SET search_path TO grant_on_table;

-- check on coordinator and workers
:verify_grant_table ;
:verify_grant_attributes ;

-- no cleanup required

--
-- check we do propagate WITH GRANT OPTION
-- we use only SELECT
--
SET ROLE grant_user_0;
-- not granted yet:
SELECT test_r, test_mix FROM ref_table;
SET ROLE grant_user_1;
-- not granted yet:
SELECT test_r, test_mix FROM ref_table;
RESET ROLE;
-- grant with grant option
GRANT SELECT (test_r, test_mix) ON ref_table TO grant_user_0 WITH GRANT OPTION;
SET ROLE grant_user_0;
-- grant using a role with grant option
GRANT SELECT (test_r, test_mix) ON ref_table TO grant_user_1;
SET ROLE grant_user_1;
-- granted:
SELECT test_r, test_mix FROM ref_table;
RESET ROLE;

-- check on coordinator and workers
:verify_grant_table ;
:verify_grant_attributes ;

-- cleanup and further checks:
SET ROLE grant_user_0;
-- revoke as grant_user_0:
REVOKE SELECT (test_r, test_mix) ON ref_table FROM grant_user_1;
RESET ROLE;
-- check on coordinator and workers
:verify_grant_table ;
:verify_grant_attributes ;

-- revoke only grant options from grant_user_0:
REVOKE GRANT OPTION FOR SELECT (test_r, test_mix) ON ref_table FROM grant_user_0;
-- check on coordinator and workers
:verify_grant_table ;
:verify_grant_attributes ;

-- revoke select from grant_user_0:
REVOKE SELECT (test_r, test_mix) ON ref_table FROM grant_user_0;
-- check on coordinator and workers
:verify_grant_table ;
:verify_grant_attributes ;

-- prevent useless messages on DROP objects.
SET client_min_messages TO ERROR;
DROP TABLE dist_table, ref_table;
SELECT result FROM run_command_on_all_nodes('DROP TABLE grant_on_table.local_table');
RESET client_min_messages;

--
-- check we propagate privileges when GRANTed before distributed and when adding a node
-- we use only SELECT
--
-- test propagation on columns when distributing the table after GRANT has been executed
CREATE TABLE grant_table_propagated (id int primary key);
GRANT SELECT (id) ON grant_table_propagated TO grant_user_0;
SELECT create_distributed_table('grant_table_propagated', 'id');
SET ROLE grant_user_0;
-- granted:
SELECT id FROM grant_table_propagated;
RESET ROLE;

-- check on coordinator and workers
\set verify_grant_table 'SELECT nodeport, unnest(result::text[]) FROM run_command_on_all_nodes($$SELECT array_agg((relname, relacl) order by 1) FROM pg_class WHERE relname IN (''grant_table_propagated'') $$) join pg_dist_node using (nodeid) ORDER BY 2, 1'

\set verify_grant_attributes 'SELECT nodeport, unnest(coalesce(nullif(result, ''''), ''{}'')::text[]) FROM run_command_on_all_nodes($$SELECT array_agg((attrelid::regclass, attname, attacl) order by 1, 2) FROM pg_attribute WHERE attrelid  IN (''grant_on_table.grant_table_propagated''::regclass) AND attacl IS NOT NULL $$) join pg_dist_node using (nodeid) ORDER BY 2, 1'

:verify_grant_table ;
:verify_grant_attributes ;

-- cleanup
REVOKE SELECT (id) ON grant_table_propagated FROM grant_user_0;

:verify_grant_table ;
:verify_grant_attributes ;

-- prevent useless messages on DROP objects.
SET client_min_messages TO ERROR;
DROP TABLE grant_table_propagated;
RESET client_min_messages;

-- similar test but adding a node after the fact
-- remove one of the worker nodes:
SELECT citus_remove_node('localhost', :worker_2_port);
CREATE TABLE grant_table_propagated_after (id int primary key);
SET citus.shard_replication_factor TO 1;
SELECT create_distributed_table('grant_table_propagated_after', 'id');
SET citus.shard_replication_factor TO 2;
GRANT SELECT (id) ON grant_table_propagated_after TO grant_user_0;
-- add back the worker node
SELECT FROM citus_add_node('localhost', :worker_2_port);

-- granted:
SELECT id FROM grant_table_propagated_after;

-- check on coordinator and workers
\set verify_grant_table 'SELECT nodeport, unnest(result::text[]) FROM run_command_on_all_nodes($$SELECT array_agg((relname, relacl) order by 1) FROM pg_class WHERE relname IN (''grant_table_propagated_after'') $$) join pg_dist_node using (nodeid) ORDER BY 2, 1'

\set verify_grant_attributes 'SELECT nodeport, unnest(coalesce(nullif(result, ''''), ''{}'')::text[]) FROM run_command_on_all_nodes($$SELECT array_agg((attrelid::regclass, attname, attacl) order by 1, 2) FROM pg_attribute WHERE attrelid  IN (''grant_on_table.grant_table_propagated_after''::regclass) AND attacl IS NOT NULL $$) join pg_dist_node using (nodeid) ORDER BY 2, 1'

:verify_grant_table ;
:verify_grant_attributes ;

-- cleanup
REVOKE SELECT (id) ON grant_table_propagated_after FROM grant_user_0;

:verify_grant_table ;
:verify_grant_attributes ;

-- prevent useless messages on DROP objects.
SET client_min_messages TO ERROR;
DROP TABLE grant_table_propagated_after;
RESET client_min_messages;

-- global cleanup
-- prevent useless messages on DROP objects.
SET client_min_messages TO ERROR;
DROP SCHEMA grant_on_table CASCADE;
DROP ROLE grant_user_0, grant_user_1;
RESET client_min_messages;
RESET search_path;