From 971a4869a9f09d456e976cfaba8adab9f1432b7f Mon Sep 17 00:00:00 2001
From: gurkanindibay <gindibay@microsoft.com>
Date: Tue, 20 Feb 2024 13:30:01 +0300
Subject: [PATCH] Fixes review notes

---
 .../distributed/commands/utility_hook.c       | 23 ++++++++++++++++++-
 ...n_database_propagation_from_non_maindb.out |  1 +
 ...n_database_propagation_from_non_maindb.sql |  1 -
 .../sql/metadata_sync_from_non_maindb.sql     |  9 --------
 4 files changed, 23 insertions(+), 11 deletions(-)

diff --git a/src/backend/distributed/commands/utility_hook.c b/src/backend/distributed/commands/utility_hook.c
index 86709edd6..244b82dd0 100644
--- a/src/backend/distributed/commands/utility_hook.c
+++ b/src/backend/distributed/commands/utility_hook.c
@@ -1703,7 +1703,28 @@ IsStatementSupportedFromNonMainDb(Node *parsetree)
 	{
 		if (type == NonMainDbSupportedStatements[i].statementType)
 		{
-			return true;
+			if (NonMainDbSupportedStatements[i].supportedObjectTypes == NULL)
+			{
+				return true;
+			}
+			else
+			{
+				if (type == T_GrantStmt)
+				{
+					GrantStmt *stmt = castNode(GrantStmt, parsetree);
+					/* check if stmt->objtype is in supportedObjectTypes */
+					for (int j = 0; j <
+						 NonMainDbSupportedStatements[i].supportedObjectTypesSize; j++)
+					{
+						if (stmt->objtype ==
+							NonMainDbSupportedStatements[i].supportedObjectTypes[j])
+						{
+							return true;
+						}
+					}
+					return false;
+				}
+			}
 		}
 	}
 
diff --git a/src/test/regress/expected/grant_on_database_propagation_from_non_maindb.out b/src/test/regress/expected/grant_on_database_propagation_from_non_maindb.out
index 92bb3f5e5..cc732dd84 100644
--- a/src/test/regress/expected/grant_on_database_propagation_from_non_maindb.out
+++ b/src/test/regress/expected/grant_on_database_propagation_from_non_maindb.out
@@ -467,4 +467,5 @@ grant connect,temp,temporary  on database test_2pc_db to public;
 drop database test_2pc_db;
 set citus.enable_create_database_propagation to off;
 DROP SCHEMA grant_on_database_propagation_non_maindb CASCADE;
+reset citus.enable_create_database_propagation;
 ---------------------------------------------------------------------
diff --git a/src/test/regress/sql/grant_on_database_propagation_from_non_maindb.sql b/src/test/regress/sql/grant_on_database_propagation_from_non_maindb.sql
index 21f6b0cc9..ef5e10d32 100644
--- a/src/test/regress/sql/grant_on_database_propagation_from_non_maindb.sql
+++ b/src/test/regress/sql/grant_on_database_propagation_from_non_maindb.sql
@@ -4,7 +4,6 @@
 set citus.enable_create_database_propagation to on;
 create database test_2pc_db;
 show citus.main_db;
-
 revoke connect,temp,temporary  on database test_2pc_db from public;
 
 CREATE SCHEMA grant_on_database_propagation_non_maindb;
diff --git a/src/test/regress/sql/metadata_sync_from_non_maindb.sql b/src/test/regress/sql/metadata_sync_from_non_maindb.sql
index ed9d09c88..157e5e4e3 100644
--- a/src/test/regress/sql/metadata_sync_from_non_maindb.sql
+++ b/src/test/regress/sql/metadata_sync_from_non_maindb.sql
@@ -1,16 +1,11 @@
 
 CREATE SCHEMA metadata_sync_2pc_schema;
-
 SET search_path TO metadata_sync_2pc_schema;
-
 set citus.enable_create_database_propagation to on;
-
-
 CREATE DATABASE metadata_sync_2pc_db;
 
 revoke connect,temp,temporary  on database metadata_sync_2pc_db from public;
 
-
 \c metadata_sync_2pc_db
 SHOW citus.main_db;
 
@@ -44,7 +39,6 @@ select check_database_privileges('grant_role2pc''_user3','metadata_sync_2pc_db',
 set citus.enable_create_database_propagation to on;
 select 1 from citus_add_node('localhost', :worker_2_port);
 
-
 select result FROM run_command_on_all_nodes($$
 SELECT array_to_json(array_agg(row_to_json(t)))
 FROM (
@@ -60,7 +54,6 @@ select check_database_privileges('grant_role2pc''_user1','metadata_sync_2pc_db',
 select check_database_privileges('grant_role2pc''_user2','metadata_sync_2pc_db',ARRAY['CONNECT']);
 select check_database_privileges('grant_role2pc''_user3','metadata_sync_2pc_db',ARRAY['CREATE','CONNECT','TEMP','TEMPORARY']);
 
-
 \c metadata_sync_2pc_db
 revoke "grant_role2pc'_user1","grant_role2pc'_user2" from grant_role2pc_user4,grant_role2pc_user5 granted by "grant_role2pc'_user3";
 
@@ -76,9 +69,7 @@ revoke CREATE on database metadata_sync_2pc_db from "grant_role2pc'_user1";
 drop user "grant_role2pc'_user1","grant_role2pc'_user2","grant_role2pc'_user3",grant_role2pc_user4,grant_role2pc_user5;
 set citus.enable_create_database_propagation to on;
 drop database metadata_sync_2pc_db;
-
 drop schema metadata_sync_2pc_schema;
 
 reset citus.enable_create_database_propagation;
 reset search_path;
-