--
-- GRANT_ON_FUNCTION_PROPAGATION
--
CREATE SCHEMA grant_on_function;
SET search_path TO grant_on_function, public;

-- remove one of the worker nodes to test adding a new node later
SET citus.shard_replication_factor TO 1;
SELECT 1 FROM citus_remove_node('localhost', :worker_2_port);

-- create some simple functions
CREATE OR REPLACE FUNCTION function_notice(text)
RETURNS void
LANGUAGE plpgsql AS $$
BEGIN
    RAISE NOTICE '%', $1;
END;
$$;

CREATE OR REPLACE FUNCTION function_notice()
RETURNS void
LANGUAGE plpgsql AS $$
BEGIN
    RAISE NOTICE 'Constant Notice';
END;
$$;

CREATE OR REPLACE FUNCTION function_hello()
RETURNS void
LANGUAGE plpgsql AS $$
BEGIN
    RAISE NOTICE 'Hello World';
END;
$$;

SET citus.enable_metadata_sync TO OFF;
CREATE OR REPLACE FUNCTION not_distributed_function()
RETURNS void
LANGUAGE plpgsql AS $$
BEGIN
    RAISE NOTICE 'not_distributed_function';
END;
$$;
RESET citus.enable_metadata_sync;

-- create some users and grant them permission on grant_on_function schema
CREATE USER function_user_1;
CREATE USER function_user_2;
CREATE USER function_user_3;
GRANT ALL ON SCHEMA grant_on_function TO function_user_1, function_user_2, function_user_3;

-- do some varying grants
GRANT EXECUTE ON FUNCTION function_notice() TO function_user_1;
GRANT EXECUTE ON FUNCTION function_notice() TO function_user_2 WITH GRANT OPTION;
SET ROLE function_user_2;
GRANT EXECUTE ON FUNCTION function_notice() TO function_user_3;
RESET ROLE;

SELECT create_distributed_function('function_notice()');

-- re-distributing the same function with GRANTs should be fine
SELECT create_distributed_function('function_notice()');


-- check grants propagated correctly after create_distributed_function
SELECT proname, pronargs, proacl FROM pg_proc WHERE proname = 'function_notice' ORDER BY 1, 2;
\c - - - :worker_1_port
SELECT proname, pronargs, proacl FROM pg_proc WHERE proname = 'function_notice' ORDER BY 1, 2;
\c - - - :master_port

-- do some varying revokes
REVOKE EXECUTE ON FUNCTION grant_on_function.function_notice() FROM function_user_1, function_user_3;
REVOKE GRANT OPTION FOR EXECUTE ON FUNCTION grant_on_function.function_notice() FROM function_user_2 CASCADE;

-- check revokes propagated correctly for the distributed function function_notice()
SELECT proname, pronargs, proacl FROM pg_proc WHERE proname = 'function_notice' ORDER BY 1, 2;
\c - - - :worker_1_port
SELECT proname, pronargs, proacl FROM pg_proc WHERE proname = 'function_notice' ORDER BY 1, 2;
\c - - - :master_port

REVOKE EXECUTE ON FUNCTION grant_on_function.function_notice() FROM function_user_2;

SELECT proname, pronargs, proacl FROM pg_proc WHERE proname = 'function_notice' ORDER BY 1, 2;
\c - - - :worker_1_port
SELECT proname, pronargs, proacl FROM pg_proc WHERE proname = 'function_notice' ORDER BY 1, 2;
\c - - - :master_port

-- distribute another function
SET search_path TO grant_on_function, public;
SELECT create_distributed_function('function_notice(text)');

-- GRANT .. ON ALL FUNCTIONS IN SCHEMA .. with multiple roles
GRANT EXECUTE ON ALL FUNCTIONS IN SCHEMA grant_on_function TO function_user_1, function_user_3;

SELECT proname, pronargs, proacl FROM pg_proc WHERE proname = 'function_notice' ORDER BY 1, 2;
\c - - - :worker_1_port
SELECT proname, pronargs, proacl FROM pg_proc WHERE proname = 'function_notice' ORDER BY 1, 2;
\c - - - :master_port

-- REVOKE .. ON ALL FUNCTIONS IN SCHEMA .. with multiple roles
REVOKE EXECUTE ON ALL FUNCTIONS IN SCHEMA grant_on_function FROM function_user_1, function_user_3;

SELECT proname, pronargs, proacl FROM pg_proc WHERE proname = 'function_notice' ORDER BY 1, 2;
\c - - - :worker_1_port
SELECT proname, pronargs, proacl FROM pg_proc WHERE proname = 'function_notice' ORDER BY 1, 2;
\c - - - :master_port

-- distribute another function
SET search_path TO grant_on_function, public;
SELECT create_distributed_function('function_hello()');

-- GRANT with multiple functions and multiple roles
-- function_hello needs no arguments since no other function has that name
GRANT EXECUTE ON FUNCTION grant_on_function.function_hello, grant_on_function.function_notice(), grant_on_function.function_notice(text), grant_on_function.not_distributed_function() TO function_user_2 WITH GRANT OPTION;
SET ROLE function_user_2;
GRANT EXECUTE ON FUNCTION grant_on_function.function_hello, grant_on_function.function_notice(), grant_on_function.function_notice(text), grant_on_function.not_distributed_function() TO function_user_1, function_user_3;
RESET ROLE;

SELECT proname, pronargs, proacl FROM pg_proc WHERE proname IN ('function_notice', 'function_hello', 'not_distributed_function') ORDER BY 1, 2;
\c - - - :worker_1_port
SELECT proname, pronargs, proacl FROM pg_proc WHERE proname IN ('function_notice', 'function_hello', 'not_distributed_function') ORDER BY 1, 2;
\c - - - :master_port

-- add the previously removed node
SELECT 1 FROM citus_add_node('localhost', :worker_2_port);

-- check if the grants are propagated correctly
SELECT proname, pronargs, proacl FROM pg_proc WHERE proname IN ('function_notice', 'function_hello') ORDER BY 1, 2;
\c - - - :worker_2_port
SELECT proname, pronargs, proacl FROM pg_proc WHERE proname IN ('function_notice', 'function_hello') ORDER BY 1, 2;
\c - - - :master_port

-- check that it works correctly with a user that is not distributed
CREATE OR REPLACE FUNCTION not_propagated_function_user_test()
RETURNS void
LANGUAGE plpgsql AS $$
BEGIN
    RAISE NOTICE 'not_propagated_function_user_test';
END;
$$;
SELECT create_distributed_function('not_propagated_function_user_test()');

SET citus.enable_ddl_propagation TO off;
CREATE USER not_propagated_function_user_4;
SET citus.enable_ddl_propagation TO on;

GRANT EXECUTE ON FUNCTION not_propagated_function_user_test TO function_user_1, not_propagated_function_user_4;

-- check if the grants are propagated correctly
SELECT proname, pronargs, proacl FROM pg_proc WHERE proname IN ('not_propagated_function_user_test') ORDER BY 1, 2;
\c - - - :worker_2_port
SELECT proname, pronargs, proacl FROM pg_proc WHERE proname IN ('not_propagated_function_user_test') ORDER BY 1, 2;
\c - - - :master_port

SET search_path TO grant_on_function, public;

-- the following should fail is in plain PG
GRANT EXECUTE ON FUNCTION function_notice(), non_existent_function TO function_user_1;
GRANT EXECUTE ON FUNCTION function_notice() TO function_user_1, non_existent_user;
GRANT EXECUTE ON ALL FUNCTIONS IN SCHEMA grant_on_function, non_existent_schema TO function_user_1;

DROP FUNCTION function_notice(), function_notice(text), function_hello, not_distributed_function, not_propagated_function_user_test;

-- add similar tests for procedures

-- remove one of the worker nodes to test adding a new node later
SELECT 1 FROM citus_remove_node('localhost', :worker_2_port);

-- create some simple procedures
CREATE OR REPLACE PROCEDURE procedure_notice(text)
LANGUAGE PLPGSQL AS $proc$
BEGIN
  RAISE NOTICE '%', $1;
END;
$proc$;

CREATE OR REPLACE PROCEDURE procedure_notice()
LANGUAGE PLPGSQL AS $proc$
BEGIN
  RAISE NOTICE 'Constant Notice';
END;
$proc$;

CREATE OR REPLACE PROCEDURE procedure_hello()
LANGUAGE PLPGSQL AS $proc$
BEGIN
  RAISE NOTICE 'Hello World';
END;
$proc$;

SET citus.enable_metadata_sync TO OFF;
CREATE OR REPLACE PROCEDURE not_distributed_procedure()
LANGUAGE PLPGSQL AS $proc$
BEGIN
  RAISE NOTICE 'not_distributed_procedure';
END;
$proc$;
RESET citus.enable_metadata_sync;

-- create some users and grant them permission on grant_on_function schema
CREATE USER procedure_user_1;
CREATE USER procedure_user_2;
CREATE USER procedure_user_3;
GRANT ALL ON SCHEMA grant_on_function TO procedure_user_1, procedure_user_2, procedure_user_3;

-- do some varying grants
GRANT EXECUTE ON PROCEDURE procedure_notice() TO procedure_user_1;
GRANT EXECUTE ON PROCEDURE procedure_notice() TO procedure_user_2 WITH GRANT OPTION;
SET ROLE procedure_user_2;
GRANT EXECUTE ON PROCEDURE procedure_notice() TO procedure_user_3;
RESET ROLE;

SELECT create_distributed_function('procedure_notice()');

-- check grants propagated correctly after create_distributed_function
SELECT proname, pronargs, proacl FROM pg_proc WHERE proname = 'procedure_notice' ORDER BY 1, 2;
\c - - - :worker_1_port
SELECT proname, pronargs, proacl FROM pg_proc WHERE proname = 'procedure_notice' ORDER BY 1, 2;
\c - - - :master_port

-- do some varying revokes
REVOKE EXECUTE ON PROCEDURE grant_on_function.procedure_notice() FROM procedure_user_1, procedure_user_3;
REVOKE GRANT OPTION FOR EXECUTE ON PROCEDURE grant_on_function.procedure_notice() FROM procedure_user_2 CASCADE;

-- check revokes propagated correctly for the distributed procedure procedure_notice()
SELECT proname, pronargs, proacl FROM pg_proc WHERE proname = 'procedure_notice' ORDER BY 1, 2;
\c - - - :worker_1_port
SELECT proname, pronargs, proacl FROM pg_proc WHERE proname = 'procedure_notice' ORDER BY 1, 2;
\c - - - :master_port

REVOKE EXECUTE ON PROCEDURE grant_on_function.procedure_notice() FROM procedure_user_2;

SELECT proname, pronargs, proacl FROM pg_proc WHERE proname = 'procedure_notice' ORDER BY 1, 2;
\c - - - :worker_1_port
SELECT proname, pronargs, proacl FROM pg_proc WHERE proname = 'procedure_notice' ORDER BY 1, 2;
\c - - - :master_port

-- distribute another procedure
SET search_path TO grant_on_function, public;
SELECT create_distributed_function('procedure_notice(text)');

-- GRANT .. ON ALL PROCEDURES IN SCHEMA .. with multiple roles
GRANT EXECUTE ON ALL PROCEDURES IN SCHEMA grant_on_function TO procedure_user_1, procedure_user_3;

SELECT proname, pronargs, proacl FROM pg_proc WHERE proname = 'procedure_notice' ORDER BY 1, 2;
\c - - - :worker_1_port
SELECT proname, pronargs, proacl FROM pg_proc WHERE proname = 'procedure_notice' ORDER BY 1, 2;
\c - - - :master_port

-- REVOKE .. ON ALL PROCEDURES IN SCHEMA .. with multiple roles
REVOKE EXECUTE ON ALL PROCEDURES IN SCHEMA grant_on_function FROM procedure_user_1, procedure_user_3;

SELECT proname, pronargs, proacl FROM pg_proc WHERE proname = 'procedure_notice' ORDER BY 1, 2;
\c - - - :worker_1_port
SELECT proname, pronargs, proacl FROM pg_proc WHERE proname = 'procedure_notice' ORDER BY 1, 2;
\c - - - :master_port

-- distribute another procedure
SET search_path TO grant_on_function, public;
SELECT create_distributed_function('procedure_hello()');

-- GRANT with multiple procedures and multiple roles
-- procedure_hello needs no arguments since no other procedure has that name
GRANT EXECUTE ON PROCEDURE grant_on_function.procedure_hello, grant_on_function.procedure_notice(), grant_on_function.procedure_notice(text), grant_on_function.not_distributed_procedure() TO procedure_user_2 WITH GRANT OPTION;
SET ROLE procedure_user_2;
GRANT EXECUTE ON PROCEDURE grant_on_function.procedure_hello, grant_on_function.procedure_notice(), grant_on_function.procedure_notice(text), grant_on_function.not_distributed_procedure() TO procedure_user_1, procedure_user_3;
RESET ROLE;

SELECT proname, pronargs, proacl FROM pg_proc WHERE proname IN ('procedure_notice', 'procedure_hello', 'not_distributed_procedure') ORDER BY 1, 2;
\c - - - :worker_1_port
SELECT proname, pronargs, proacl FROM pg_proc WHERE proname IN ('procedure_notice', 'procedure_hello', 'not_distributed_procedure') ORDER BY 1, 2;
\c - - - :master_port

-- add the previously removed node
SELECT 1 FROM citus_add_node('localhost', :worker_2_port);

-- check if the grants are propagated correctly
SELECT proname, pronargs, proacl FROM pg_proc WHERE proname IN ('procedure_notice', 'procedure_hello') ORDER BY 1, 2;
\c - - - :worker_2_port
SELECT proname, pronargs, proacl FROM pg_proc WHERE proname IN ('procedure_notice', 'procedure_hello') ORDER BY 1, 2;
\c - - - :master_port

-- check that it works correctly with a user that is not distributed
CREATE OR REPLACE PROCEDURE not_propagated_procedure_user_test()
LANGUAGE PLPGSQL AS $proc$
BEGIN
  RAISE NOTICE 'not_propagated_procedure_user_test';
END;
$proc$;
SELECT create_distributed_function('not_propagated_procedure_user_test()');

SET citus.enable_ddl_propagation TO off;
CREATE USER not_propagated_procedure_user_4;
SET citus.enable_ddl_propagation TO on;

GRANT EXECUTE ON PROCEDURE not_propagated_procedure_user_test TO procedure_user_1, not_propagated_procedure_user_4;

-- check if the grants are propagated correctly
SELECT proname, pronargs, proacl FROM pg_proc WHERE proname IN ('not_propagated_procedure_user_test') ORDER BY 1, 2;
\c - - - :worker_2_port
SELECT proname, pronargs, proacl FROM pg_proc WHERE proname IN ('not_propagated_procedure_user_test') ORDER BY 1, 2;
\c - - - :master_port

SET search_path TO grant_on_function, public;

-- the following should fail is in plain PG
GRANT EXECUTE ON PROCEDURE procedure_notice(), non_existent_procedure TO procedure_user_1;
GRANT EXECUTE ON PROCEDURE procedure_notice() TO procedure_user_1, non_existent_user;
GRANT EXECUTE ON ALL PROCEDURES IN SCHEMA grant_on_function, non_existent_schema TO procedure_user_1;

DROP PROCEDURE procedure_notice(), procedure_notice(text), procedure_hello, not_distributed_procedure, not_propagated_procedure_user_test;

SELECT run_command_on_coordinator_and_workers('DROP SCHEMA grant_on_function CASCADE');
DROP USER function_user_1, function_user_2, function_user_3, not_propagated_function_user_4;
DROP USER procedure_user_1, procedure_user_2, procedure_user_3, not_propagated_procedure_user_4;
