Allow table type to be used in target list

src/backend/distributed/deparser/ruleutils_11.c

@@ -429,6 +429,7 @@ static void printSubscripts(ArrayRef *aref, deparse_context *context);
 static char *get_relation_name(Oid relid);
 static char *generate_relation_or_shard_name(Oid relid, Oid distrelid,
 				int64 shardid, List *namespaces);
+static char *generate_rte_shard_name(RangeTblEntry *rangeTableEntry);
 static char *generate_fragment_name(char *schemaName, char *tableName);
 static char *generate_function_name(Oid funcid, int nargs,
 					   List *argnames, Oid *argtypes,
@@ -3750,10 +3751,22 @@ get_variable(Var *var, int levelsup, bool istoplevel, deparse_context *context)
 	else
 	{
 		appendStringInfoChar(buf, '*');
+
 		if (istoplevel)
-			appendStringInfo(buf, "::%s",
+		{
-							 format_type_with_typemod(var->vartype,
+			if (GetRangeTblKind(rte) == CITUS_RTE_SHARD)
-													  var->vartypmod));
+			{
+				/* use rel.*::shard_name instead of rel.*::table_name */
+				appendStringInfo(buf, "::%s",
+								 generate_rte_shard_name(rte));
+			}
+			else
+			{
+				appendStringInfo(buf, "::%s",
+								 format_type_with_typemod(var->vartype,
+														  var->vartypmod));
+			}
+		}
 	}
 
 	return attname;
@@ -7736,6 +7749,26 @@ generate_relation_name(Oid relid, List *namespaces)
 	return result;
 }
 
+
+/*
+ * generate_rte_shard_name returns the qualified name of the shard given a
+ * CITUS_RTE_SHARD range table entry.
+ */
+static char *
+generate_rte_shard_name(RangeTblEntry *rangeTableEntry)
+{
+	char *shardSchemaName = NULL;
+	char *shardTableName = NULL;
+
+	Assert(GetRangeTblKind(rangeTableEntry) == CITUS_RTE_SHARD);
+
+	ExtractRangeTblExtraData(rangeTableEntry, NULL, &shardSchemaName, &shardTableName,
+							 NULL);
+
+	return generate_fragment_name(shardSchemaName, shardTableName);
+}
+
+
 /*
  * generate_fragment_name
  *		Compute the name to display for a shard or merged table

src/backend/distributed/deparser/ruleutils_12.c

@@ -430,6 +430,7 @@ static void printSubscripts(SubscriptingRef *aref, deparse_context *context);
 static char *get_relation_name(Oid relid);
 static char *generate_relation_or_shard_name(Oid relid, Oid distrelid,
 				int64 shardid, List *namespaces);
+static char *generate_rte_shard_name(RangeTblEntry *rangeTableEntry);
 static char *generate_fragment_name(char *schemaName, char *tableName);
 static char *generate_function_name(Oid funcid, int nargs,
 					   List *argnames, Oid *argtypes,
@@ -3763,10 +3764,22 @@ get_variable(Var *var, int levelsup, bool istoplevel, deparse_context *context)
 	else
 	{
 		appendStringInfoChar(buf, '*');
+
 		if (istoplevel)
-			appendStringInfo(buf, "::%s",
+		{
-							 format_type_with_typemod(var->vartype,
+			if (GetRangeTblKind(rte) == CITUS_RTE_SHARD)
-													  var->vartypmod));
+			{
+				/* use rel.*::shard_name instead of rel.*::table_name */
+				appendStringInfo(buf, "::%s",
+								 generate_rte_shard_name(rte));
+			}
+			else
+			{
+				appendStringInfo(buf, "::%s",
+								 format_type_with_typemod(var->vartype,
+														  var->vartypmod));
+			}
+		}
 	}
 
 	return attname;
@@ -7737,6 +7750,26 @@ generate_relation_name(Oid relid, List *namespaces)
 	return result;
 }
 
+
+/*
+ * generate_rte_shard_name returns the qualified name of the shard given a
+ * CITUS_RTE_SHARD range table entry.
+ */
+static char *
+generate_rte_shard_name(RangeTblEntry *rangeTableEntry)
+{
+	char *shardSchemaName = NULL;
+	char *shardTableName = NULL;
+
+	Assert(GetRangeTblKind(rangeTableEntry) == CITUS_RTE_SHARD);
+
+	ExtractRangeTblExtraData(rangeTableEntry, NULL, &shardSchemaName, &shardTableName,
+							 NULL);
+
+	return generate_fragment_name(shardSchemaName, shardTableName);
+}
+
+
 /*
  * generate_fragment_name
  *		Compute the name to display for a shard or merged table

src/test/regress/expected/aggregate_support.out

@@ -430,8 +430,17 @@ RESET citus.task_executor_type;
 -- This fails due to table types not being managed properly
 select key, count(distinct aggdata)
 from aggdata group by key order by 1, 2;
-ERROR:  type "aggregate_support.aggdata" does not exist
+ key | count
-CONTEXT:  while executing command on localhost:xxxxx
+---------------------------------------------------------------------
+   1 |     2
+   2 |     3
+   3 |     1
+   5 |     1
+   6 |     2
+   7 |     1
+   9 |     1
+(7 rows)
+
 -- GROUPING parses to GroupingFunc, distinct from Aggref
 -- These three queries represent edge cases implementation would have to consider
 -- For now we error out of all three

src/test/regress/expected/subquery_complex_target_list.out

@@ -1,5 +1,5 @@
 -- ===================================================================
--- test recursive planning functionality with complex target entries
+-- test subquery functionality with complex target entries
 -- and some utilities
 -- ===================================================================
 CREATE SCHEMA subquery_complex;
@@ -499,7 +499,99 @@ GROUP BY key ORDER BY 3, 2, 1;
  key-2 | value-1 |   2
 (2 rows)
 
+-- whole rows in the output should work
+SELECT a.key, a, count(b)
+FROM items a LEFT JOIN other_items b ON (a.key = b.key)
+GROUP BY a.key ORDER BY 3, 2, 1;
+  key  |                     a                      | count
+---------------------------------------------------------------------
+ key-2 | (key-2,value-1,"Sun Feb 02 00:00:00 2020") |     0
+ key-1 | (key-1,value-2,"Wed Jan 01 00:00:00 2020") |     1
+(2 rows)
+
+SELECT a FROM items a ORDER BY key;
+                     a
+---------------------------------------------------------------------
+ (key-1,value-2,"Wed Jan 01 00:00:00 2020")
+ (key-2,value-1,"Sun Feb 02 00:00:00 2020")
+(2 rows)
+
+SELECT a FROM items a WHERE key = 'key-1';
+                     a
+---------------------------------------------------------------------
+ (key-1,value-2,"Wed Jan 01 00:00:00 2020")
+(1 row)
+
+SELECT a FROM (SELECT a, random() FROM items a) b ORDER BY a;
+                     a
+---------------------------------------------------------------------
+ (key-1,value-2,"Wed Jan 01 00:00:00 2020")
+ (key-2,value-1,"Sun Feb 02 00:00:00 2020")
+(2 rows)
+
+-- whole rows when table name needs escaping
+create table "another table" (key text primary key, value text not null, t timestamp);
+SELECT create_distributed_table('"another table"','key');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+INSERT INTO "another table" TABLE items;
+SELECT a.key, a, count(b)
+FROM "another table" a LEFT JOIN other_items b ON (a.key = b.key)
+GROUP BY a.key ORDER BY 3, 2, 1;
+  key  |                     a                      | count
+---------------------------------------------------------------------
+ key-2 | (key-2,value-1,"Sun Feb 02 00:00:00 2020") |     0
+ key-1 | (key-1,value-2,"Wed Jan 01 00:00:00 2020") |     1
+(2 rows)
+
+-- subquery output is not recognised yet
+SELECT b FROM (SELECT a FROM items a GROUP BY key) b ORDER BY b;
+ERROR:  input of anonymous composite types is not implemented
+-- table output in recursive planning
+WITH a AS (SELECT a FROM items a OFFSET 0)
+SELECT count(a) FROM a;
+ count
+---------------------------------------------------------------------
+     2
+(1 row)
+
+WITH a AS (SELECT a AS b FROM items a OFFSET 0)
+SELECT b FROM a ORDER BY b;
+                     b
+---------------------------------------------------------------------
+ (key-1,value-2,"Wed Jan 01 00:00:00 2020")
+ (key-2,value-1,"Sun Feb 02 00:00:00 2020")
+(2 rows)
+
+BEGIN;
+WITH a AS (DELETE FROM items RETURNING items)
+SELECT count(a) FROM a;
+ count
+---------------------------------------------------------------------
+     2
+(1 row)
+
+ROLLBACK;
+-- mix custom types with table type
+CREATE TYPE full_item AS (key text, value text);
+ALTER TABLE items ADD COLUMN fi subquery_complex.full_item;
+UPDATE items SET fi = (items.key, items.value);
+SELECT a, a.fi, (a.fi).* FROM items a ORDER BY 1,2,3;
+                              a                               |       fi        |  key  |  value
+---------------------------------------------------------------------
+ (key-1,value-2,"Wed Jan 01 00:00:00 2020","(key-1,value-2)") | (key-1,value-2) | key-1 | value-2
+ (key-2,value-1,"Sun Feb 02 00:00:00 2020","(key-2,value-1)") | (key-2,value-1) | key-2 | value-1
+(2 rows)
+
+SELECT a, b.fi, (b.fi).*, (a).fi, ((a).fi).* FROM (SELECT a, fi FROM items a GROUP BY key) b ORDER BY 1,2,3;
+                              a                               |       fi        |  key  |  value  |       fi        |  key  |  value
+---------------------------------------------------------------------
+ (key-1,value-2,"Wed Jan 01 00:00:00 2020","(key-1,value-2)") | (key-1,value-2) | key-1 | value-2 | (key-1,value-2) | key-1 | value-2
+ (key-2,value-1,"Sun Feb 02 00:00:00 2020","(key-2,value-1)") | (key-2,value-1) | key-2 | value-1 | (key-2,value-1) | key-2 | value-1
+(2 rows)
+
+SET client_min_messages TO 'WARNING';
 DROP SCHEMA subquery_complex CASCADE;
-NOTICE:  drop cascades to 2 other objects
-DETAIL:  drop cascades to table items
-drop cascades to table other_items

src/test/regress/sql/subquery_complex_target_list.sql

@@ -1,5 +1,5 @@
 -- ===================================================================
--- test recursive planning functionality with complex target entries
+-- test subquery functionality with complex target entries
 -- and some utilities
 -- ===================================================================
 CREATE SCHEMA subquery_complex;
@@ -346,4 +346,45 @@ SELECT key, a.value, sum(b)
 FROM items a JOIN generate_series(1,10) b ON (a.key = 'key-'||b)
 GROUP BY key ORDER BY 3, 2, 1;
 
+-- whole rows in the output should work
+SELECT a.key, a, count(b)
+FROM items a LEFT JOIN other_items b ON (a.key = b.key)
+GROUP BY a.key ORDER BY 3, 2, 1;
+
+SELECT a FROM items a ORDER BY key;
+SELECT a FROM items a WHERE key = 'key-1';
+SELECT a FROM (SELECT a, random() FROM items a) b ORDER BY a;
+
+-- whole rows when table name needs escaping
+create table "another table" (key text primary key, value text not null, t timestamp);
+SELECT create_distributed_table('"another table"','key');
+INSERT INTO "another table" TABLE items;
+SELECT a.key, a, count(b)
+FROM "another table" a LEFT JOIN other_items b ON (a.key = b.key)
+GROUP BY a.key ORDER BY 3, 2, 1;
+
+-- subquery output is not recognised yet
+SELECT b FROM (SELECT a FROM items a GROUP BY key) b ORDER BY b;
+
+-- table output in recursive planning
+WITH a AS (SELECT a FROM items a OFFSET 0)
+SELECT count(a) FROM a;
+
+WITH a AS (SELECT a AS b FROM items a OFFSET 0)
+SELECT b FROM a ORDER BY b;
+
+BEGIN;
+WITH a AS (DELETE FROM items RETURNING items)
+SELECT count(a) FROM a;
+ROLLBACK;
+
+-- mix custom types with table type
+CREATE TYPE full_item AS (key text, value text);
+ALTER TABLE items ADD COLUMN fi subquery_complex.full_item;
+UPDATE items SET fi = (items.key, items.value);
+
+SELECT a, a.fi, (a.fi).* FROM items a ORDER BY 1,2,3;
+SELECT a, b.fi, (b.fi).*, (a).fi, ((a).fi).* FROM (SELECT a, fi FROM items a GROUP BY key) b ORDER BY 1,2,3;
+
+SET client_min_messages TO 'WARNING';
 DROP SCHEMA subquery_complex CASCADE;