From bb2ab9e93c76db57b0353825b246077e0d23ac59 Mon Sep 17 00:00:00 2001
From: Gokhan Gulbiz <ggulbiz@gmail.com>
Date: Wed, 22 Mar 2023 16:39:10 +0300
Subject: [PATCH] Move adding tenant statistic attributes to local_executor to
 keep track of both local shard queries and local cached queries

---
 src/backend/distributed/executor/local_executor.c     | 3 +++
 src/backend/distributed/planner/deparse_shard_query.c | 1 -
 src/backend/distributed/planner/distributed_planner.c | 7 +++++--
 src/backend/distributed/utils/attribute.c             | 3 ---
 4 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/src/backend/distributed/executor/local_executor.c b/src/backend/distributed/executor/local_executor.c
index ffd063ca0..899b50774 100644
--- a/src/backend/distributed/executor/local_executor.c
+++ b/src/backend/distributed/executor/local_executor.c
@@ -103,6 +103,7 @@
 #include "optimizer/optimizer.h"
 #include "nodes/params.h"
 #include "utils/snapmgr.h"
+#include "distributed/utils/attribute.h"
 
 /* controlled via a GUC */
 bool EnableLocalExecution = true;
@@ -646,6 +647,8 @@ LocallyExecuteTaskPlan(PlannedStmt *taskPlan, char *queryString,
 		LocalExecutorShardId = task->anchorShardId;
 	}
 
+	AttributeTask(task->partitionColumn, task->colocationId, taskPlan->commandType);
+
 	PG_TRY();
 	{
 		processedRows = ExecuteTaskPlan(taskPlan, queryString, tupleDest, task,
diff --git a/src/backend/distributed/planner/deparse_shard_query.c b/src/backend/distributed/planner/deparse_shard_query.c
index 4be978c78..afb4597b8 100644
--- a/src/backend/distributed/planner/deparse_shard_query.c
+++ b/src/backend/distributed/planner/deparse_shard_query.c
@@ -400,7 +400,6 @@ SetTaskQueryIfShouldLazyDeparse(Task *task, Query *query)
 		task->taskQuery.queryType = TASK_QUERY_OBJECT;
 		task->taskQuery.data.jobQueryReferenceForLazyDeparsing = query;
 		task->queryCount = 1;
-		AttributeTask(task->partitionColumn, task->colocationId, query->commandType);
 
 		return;
 	}
diff --git a/src/backend/distributed/planner/distributed_planner.c b/src/backend/distributed/planner/distributed_planner.c
index e069ec417..4e2370130 100644
--- a/src/backend/distributed/planner/distributed_planner.c
+++ b/src/backend/distributed/planner/distributed_planner.c
@@ -146,8 +146,11 @@ distributed_planner(Query *parse,
 	bool fastPathRouterQuery = false;
 	Node *distributionKeyValue = NULL;
 
-	AttributeQueryIfAnnotated(query_string, parse->commandType);
-
+	if (query_string != NULL)
+	{
+		AttributeQueryIfAnnotated(query_string, parse->commandType);
+	}
+	
 	List *rangeTableList = ExtractRangeTableEntryList(parse);
 
 	if (cursorOptions & CURSOR_OPT_FORCE_DISTRIBUTED)
diff --git a/src/backend/distributed/utils/attribute.c b/src/backend/distributed/utils/attribute.c
index db97f9c8a..29862dce5 100644
--- a/src/backend/distributed/utils/attribute.c
+++ b/src/backend/distributed/utils/attribute.c
@@ -205,9 +205,6 @@ AttributeQueryIfAnnotated(const char *query_string, CmdType commandType)
  */
 void AttributeTask(char *tenantId, int colocationId, CmdType commandType)
 {
-	ereport(NOTICE, (errmsg("MyProcPid: %d", MyProcPid)));
-	sleep(10);
-
 	colocationGroupId = colocationId;
 	strcpy_s(attributeToTenant, sizeof(attributeToTenant), tenantId);
 	attributeCommandType = commandType;