From 6dfc402acec714067570a11c6a9d195501fea7ca Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?C=C3=A9dric=20Villemain?= <Cedric.Villemain@Data-Bene.io>
Date: Mon, 23 Oct 2023 19:55:12 +0200
Subject: [PATCH] Do the index check in the dedicated check function

And minor indentation fixes, removed old comment, update tests.
---
 .../distributed/metadata/metadata_utility.c   | 45 ++++++++++++-------
 .../regress/expected/multi_size_queries.out   |  6 +--
 2 files changed, 32 insertions(+), 19 deletions(-)

diff --git a/src/backend/distributed/metadata/metadata_utility.c b/src/backend/distributed/metadata/metadata_utility.c
index 13a3cf56f8b..24e6023ca3e 100644
--- a/src/backend/distributed/metadata/metadata_utility.c
+++ b/src/backend/distributed/metadata/metadata_utility.c
@@ -547,14 +547,7 @@ DistributedRelationSize(Oid relationId, SizeQueryType sizeQueryType,
 		return false;
 	}
 
-	Oid checkRelId = relationId;
-	if (relation->rd_rel->relkind == RELKIND_INDEX)
-	{
-		bool missingOk = false;
-		checkRelId = IndexGetRelation(relation->rd_id, missingOk);
-	}
-
-	ErrorIfNotSuitableToGetSize(checkRelId);
+	ErrorIfNotSuitableToGetSize(relationId);
 
 	table_close(relation, AccessShareLock);
 
@@ -584,7 +577,7 @@ DistributedRelationSize(Oid relationId, SizeQueryType sizeQueryType,
 /*
  * DistributedRelationSizeOnWorker gets the workerNode and relationId to calculate
  * size of that relation on the given workerNode by summing up the size of each
- * shard placement. If indexId is defined then the relation is an index.
+ * shard placement.
  */
 static bool
 DistributedRelationSizeOnWorker(WorkerNode *workerNode, Oid relationId,
@@ -604,12 +597,12 @@ DistributedRelationSizeOnWorker(WorkerNode *workerNode, Oid relationId,
 	PGresult *result = NULL;
 
 	/* if the relation is an index, update relationId and define indexId */
-    Oid indexId = InvalidOid;
+	Oid indexId = InvalidOid;
 	if (get_rel_relkind(relationId) == RELKIND_INDEX)
 	{
 		indexId = relationId;
 
-        bool missingOk = false;
+		bool missingOk = false;
 		relationId = IndexGetRelation(indexId, missingOk);
 	}
 
@@ -1051,11 +1044,31 @@ ErrorIfNotSuitableToGetSize(Oid relationId)
 {
 	if (!IsCitusTable(relationId))
 	{
-		char *relationName = get_rel_name(relationId);
-		char *escapedQueryString = quote_literal_cstr(relationName);
-		ereport(ERROR, (errcode(ERRCODE_INVALID_TABLE_DEFINITION),
-						errmsg("cannot calculate the size because relation %s is not "
-							   "distributed", escapedQueryString)));
+		if (get_rel_relkind(relationId) != RELKIND_INDEX)
+		{
+			char *relationName = get_rel_name(relationId);
+			char *escapedQueryString = quote_literal_cstr(relationName);
+			ereport(ERROR, (errcode(ERRCODE_INVALID_TABLE_DEFINITION),
+							errmsg(
+								"cannot calculate the size because relation %s "
+								"is not distributed",
+								escapedQueryString)));
+		}
+		bool missingOk = false;
+		Oid indexId = relationId;
+		relationId = IndexGetRelation(relationId, missingOk);
+		if (!IsCitusTable(relationId))
+		{
+			char *tableName = get_rel_name(relationId);
+			char *escapedTableName = quote_literal_cstr(tableName);
+			char *indexName = get_rel_name(indexId);
+			char *escapedIndexName = quote_literal_cstr(indexName);
+			ereport(ERROR, (errcode(ERRCODE_INVALID_OBJECT_DEFINITION),
+							errmsg(
+								"cannot calculate the size because table %s for "
+								"index %s is not distributed",
+								escapedTableName, escapedIndexName)));
+		}
 	}
 }
 
diff --git a/src/test/regress/expected/multi_size_queries.out b/src/test/regress/expected/multi_size_queries.out
index 30c508847b0..d0ee3d6d215 100644
--- a/src/test/regress/expected/multi_size_queries.out
+++ b/src/test/regress/expected/multi_size_queries.out
@@ -21,11 +21,11 @@ ERROR:  cannot calculate the size because relation 'non_distributed_table' is no
 SELECT citus_total_relation_size('non_distributed_table');
 ERROR:  cannot calculate the size because relation 'non_distributed_table' is not distributed
 SELECT citus_table_size('non_distributed_table_pkey');
-ERROR:  cannot calculate the size because relation 'non_distributed_table' is not distributed
+ERROR:  cannot calculate the size because table 'non_distributed_table' for index 'non_distributed_table_pkey' is not distributed
 SELECT citus_relation_size('non_distributed_table_pkey');
-ERROR:  cannot calculate the size because relation 'non_distributed_table' is not distributed
+ERROR:  cannot calculate the size because table 'non_distributed_table' for index 'non_distributed_table_pkey' is not distributed
 SELECT citus_total_relation_size('non_distributed_table_pkey');
-ERROR:  cannot calculate the size because relation 'non_distributed_table' is not distributed
+ERROR:  cannot calculate the size because table 'non_distributed_table' for index 'non_distributed_table_pkey' is not distributed
 DROP TABLE non_distributed_table;
 -- fix broken placements via disabling the node
 SET client_min_messages TO ERROR;