From afabd37fa91fe2aa1bf40687cd9f0b1598d7576a Mon Sep 17 00:00:00 2001
From: Onur Tirtir <onurcantirtir@gmail.com>
Date: Thu, 3 Aug 2023 19:13:55 +0300
Subject: [PATCH] commit

---
 .../distributed/commands/utility_hook.c       |  4 +-
 .../expected/schema_based_sharding.out        | 49 ++++++++++++++++++-
 .../regress/sql/schema_based_sharding.sql     | 40 ++++++++++++++-
 3 files changed, 89 insertions(+), 4 deletions(-)

diff --git a/src/backend/distributed/commands/utility_hook.c b/src/backend/distributed/commands/utility_hook.c
index 888b3dfed..6444be234 100644
--- a/src/backend/distributed/commands/utility_hook.c
+++ b/src/backend/distributed/commands/utility_hook.c
@@ -332,10 +332,10 @@ multi_ProcessUtility(PlannedStmt *pstmt,
 			ResetConstraintDropped();
 
 			/*
-			 * We're only interested in top-level CREATE TABLE commands
+			 * We're not interested in CREATE TABLE commands in the form a subcommand
 			 * to create a tenant schema table or a Citus managed table.
 			 */
-			if (context == PROCESS_UTILITY_TOPLEVEL &&
+			if (context != PROCESS_UTILITY_SUBCOMMAND &&
 				(IsA(parsetree, CreateStmt) ||
 				 IsA(parsetree, CreateForeignTableStmt) ||
 				 IsA(parsetree, CreateTableAsStmt)))
diff --git a/src/test/regress/expected/schema_based_sharding.out b/src/test/regress/expected/schema_based_sharding.out
index 4493f9614..680116110 100644
--- a/src/test/regress/expected/schema_based_sharding.out
+++ b/src/test/regress/expected/schema_based_sharding.out
@@ -1720,8 +1720,55 @@ $$);
  {t}
 (2 rows)
 
+SET citus.enable_schema_based_sharding TO ON;
+-- verify that creating tenant tables in a DO / function block works fine
+CREATE SCHEMA do_block_test;
+DO $$
+BEGIN
+    CREATE TABLE do_block_test.tbl_1(a int, b text);
+END;
+$$;
+CREATE OR REPLACE FUNCTION do_block_test.fn_1()
+RETURNS void AS $$
+BEGIN
+    CREATE TABLE do_block_test.tbl_2(a int, b text);
+END;
+$$ LANGUAGE plpgsql;
+CREATE OR REPLACE FUNCTION do_block_test.fn_2()
+RETURNS void AS $$
+    CREATE TABLE do_block_test.tbl_3(a int, b text);
+$$ LANGUAGE sql;
+SELECT do_block_test.fn_1();
+ fn_1
+---------------------------------------------------------------------
+
+(1 row)
+
+SELECT do_block_test.fn_2();
+ fn_2
+---------------------------------------------------------------------
+
+(1 row)
+
+SELECT result FROM run_command_on_all_nodes($$
+    SELECT COUNT(*)=3 FROM citus_tables
+    WHERE
+        table_name IN (
+            'do_block_test.tbl_1'::regclass,
+            'do_block_test.tbl_2'::regclass,
+            'do_block_test.tbl_3'::regclass
+        ) AND
+        citus_table_type = 'schema';
+$$);
+ result
+---------------------------------------------------------------------
+ t
+ t
+ t
+(3 rows)
+
 SET client_min_messages TO WARNING;
-DROP SCHEMA regular_schema, tenant_3, tenant_5, tenant_7, tenant_6, type_sch, citus_sch1, citus_sch2, citus_empty_sch1, citus_empty_sch2, authschema, sc1 CASCADE;
+DROP SCHEMA regular_schema, tenant_3, tenant_5, tenant_7, tenant_6, type_sch, citus_sch1, citus_sch2, citus_empty_sch1, citus_empty_sch2, authschema, sc1, do_block_test CASCADE;
 DROP ROLE citus_schema_role, citus_schema_nonpri, authschema;
 SELECT citus_remove_node('localhost', :master_port);
  citus_remove_node
diff --git a/src/test/regress/sql/schema_based_sharding.sql b/src/test/regress/sql/schema_based_sharding.sql
index 2b9bbf516..6120afc7d 100644
--- a/src/test/regress/sql/schema_based_sharding.sql
+++ b/src/test/regress/sql/schema_based_sharding.sql
@@ -1172,8 +1172,46 @@ SELECT result FROM run_command_on_workers($$
     WHERE schemaid::regnamespace::text = 'sc1'
 $$);
 
+SET citus.enable_schema_based_sharding TO ON;
+
+-- verify that creating tenant tables in a DO / function block works fine
+
+CREATE SCHEMA do_block_test;
+
+DO $$
+BEGIN
+    CREATE TABLE do_block_test.tbl_1(a int, b text);
+END;
+$$;
+
+CREATE OR REPLACE FUNCTION do_block_test.fn_1()
+RETURNS void AS $$
+BEGIN
+    CREATE TABLE do_block_test.tbl_2(a int, b text);
+END;
+$$ LANGUAGE plpgsql;
+
+CREATE OR REPLACE FUNCTION do_block_test.fn_2()
+RETURNS void AS $$
+    CREATE TABLE do_block_test.tbl_3(a int, b text);
+$$ LANGUAGE sql;
+
+SELECT do_block_test.fn_1();
+SELECT do_block_test.fn_2();
+
+SELECT result FROM run_command_on_all_nodes($$
+    SELECT COUNT(*)=3 FROM citus_tables
+    WHERE
+        table_name IN (
+            'do_block_test.tbl_1'::regclass,
+            'do_block_test.tbl_2'::regclass,
+            'do_block_test.tbl_3'::regclass
+        ) AND
+        citus_table_type = 'schema';
+$$);
+
 SET client_min_messages TO WARNING;
-DROP SCHEMA regular_schema, tenant_3, tenant_5, tenant_7, tenant_6, type_sch, citus_sch1, citus_sch2, citus_empty_sch1, citus_empty_sch2, authschema, sc1 CASCADE;
+DROP SCHEMA regular_schema, tenant_3, tenant_5, tenant_7, tenant_6, type_sch, citus_sch1, citus_sch2, citus_empty_sch1, citus_empty_sch2, authschema, sc1, do_block_test CASCADE;
 DROP ROLE citus_schema_role, citus_schema_nonpri, authschema;
 
 SELECT citus_remove_node('localhost', :master_port);