From b999b60c18c63f224bcfcc0ce3a2d8b7cc0be8ee Mon Sep 17 00:00:00 2001
From: tibrewalpratik <tibrewalpratik@uber.com>
Date: Fri, 6 Oct 2023 21:02:03 +0530
Subject: [PATCH] added set and get tableName

---
 .../requesthandler/BaseBrokerRequestHandler.java    |  2 +-
 .../pinot/spi/trace/DefaultRequestContext.java      | 13 +++++++++++++
 .../org/apache/pinot/spi/trace/RequestContext.java  |  4 ++++
 3 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseBrokerRequestHandler.java b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseBrokerRequestHandler.java
index fe6cced4bfa..49ae0d9d941 100644
--- a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseBrokerRequestHandler.java
+++ b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseBrokerRequestHandler.java
@@ -366,7 +366,7 @@ protected BrokerResponse handleRequest(long requestId, String query, @Nullable S
       String tableName = getActualTableName(dataSource.getTableName(), _tableCache);
       dataSource.setTableName(tableName);
       String rawTableName = TableNameBuilder.extractRawTableName(tableName);
-      requestContext.setTableNames(Collections.singletonList(rawTableName));
+      requestContext.setTableName(rawTableName);
 
       try {
         boolean isCaseInsensitive = _tableCache.isIgnoreCase();
diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/trace/DefaultRequestContext.java b/pinot-spi/src/main/java/org/apache/pinot/spi/trace/DefaultRequestContext.java
index fcb8ffb5fb6..6ce063d2536 100644
--- a/pinot-spi/src/main/java/org/apache/pinot/spi/trace/DefaultRequestContext.java
+++ b/pinot-spi/src/main/java/org/apache/pinot/spi/trace/DefaultRequestContext.java
@@ -184,6 +184,11 @@ public void setQuery(String query) {
     _query = query;
   }
 
+  @Override
+  public void setTableName(String tableName) {
+    _tableNames.add(tableName);
+  }
+
   @Override
   public void setTableNames(List<String> tableNames) {
     _tableNames.addAll(tableNames);
@@ -254,6 +259,14 @@ public String getQuery() {
     return _query;
   }
 
+  @Override
+  public String getTableName() {
+    if (_tableNames.size() == 0) {
+      return DEFAULT_TABLE_NAME;
+    }
+    return _tableNames.get(0);
+  }
+
   @Override
   public List<String> getTableNames() {
     return _tableNames;
diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/trace/RequestContext.java b/pinot-spi/src/main/java/org/apache/pinot/spi/trace/RequestContext.java
index ffce2968d06..f8ec35a921d 100644
--- a/pinot-spi/src/main/java/org/apache/pinot/spi/trace/RequestContext.java
+++ b/pinot-spi/src/main/java/org/apache/pinot/spi/trace/RequestContext.java
@@ -67,6 +67,8 @@ default boolean isSampledRequest() {
 
   void setQuery(String pql);
 
+  void setTableName(String tableName);
+
   void setTableNames(List<String> tableNames);
 
   void setQueryProcessingTime(long processingTimeMillis);
@@ -95,6 +97,8 @@ default boolean isSampledRequest() {
 
   String getQuery();
 
+  String getTableName();
+
   List<String> getTableNames();
 
   long getProcessingTimeMillis();