added set and get tableName

apache · Oct 6, 2023 · b999b60 · b999b60
1 parent aff099f
commit b999b60
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 1 deletion.
diff --git a/...broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseBrokerRequestHandler.java b/...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
@@ -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
@@ -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();