Add X-TimeZone header support to REST API (#17344)

external-service-impl/rest/src/main/java/org/apache/iotdb/rest/protocol/filter/AuthorizationFilter.java

@@ -39,6 +39,7 @@
 import javax.ws.rs.ext.Provider;
 
 import java.io.IOException;
+import java.time.DateTimeException;
 import java.time.ZoneId;
 import java.util.Base64;
 import java.util.UUID;
@@ -88,6 +89,12 @@ public void filter(ContainerRequestContext containerRequestContext) throws IOExc
     if (user == null) {
       return;
     }
+
+    ZoneId zoneId = parseTimeZone(containerRequestContext);
+    if (zoneId == null) {
+      return;
+    }
+
     String sessionid = UUID.randomUUID().toString();
     if (SESSION_MANAGER.getCurrSession() == null) {
       RestClientSession restClientSession = new RestClientSession(sessionid);
@@ -97,7 +104,7 @@ public void filter(ContainerRequestContext containerRequestContext) throws IOExc
           SESSION_MANAGER.getCurrSession(),
           user.getUserId(),
           user.getUsername(),
-          ZoneId.systemDefault(),
+          zoneId,
           IoTDBConstant.ClientVersion.V_1_0);
     }
     BasicSecurityContext basicSecurityContext =
@@ -147,6 +154,33 @@ private User checkLogin(
     return user;
   }
 
+  /**
+   * Parses the X-TimeZone header from the request.
+   *
+   * @param requestContext the incoming HTTP request
+   * @return the parsed ZoneId, or {@code null} if the header is invalid (the request is aborted)
+   */
+  private ZoneId parseTimeZone(ContainerRequestContext requestContext) {
+    String timeZoneHeader = requestContext.getHeaderString("X-TimeZone");
+    if (timeZoneHeader == null || timeZoneHeader.isEmpty()) {
+      return ZoneId.systemDefault();
+    }
+    try {
+      return ZoneId.of(timeZoneHeader);
+    } catch (DateTimeException e) {
+      Response resp =
+          Response.status(Status.BAD_REQUEST)
+              .type(MediaType.APPLICATION_JSON)
+              .entity(
+                  new ExecutionStatus()
+                      .code(TSStatusCode.ILLEGAL_PARAMETER.getStatusCode())
+                      .message("Invalid time zone: " + timeZoneHeader))
+              .build();
+      requestContext.abortWith(resp);
+      return null;
+    }
+  }
+
   @Override
   public void filter(
       ContainerRequestContext requestContext, ContainerResponseContext responseContext)

external-service-impl/rest/src/main/java/org/apache/iotdb/rest/protocol/v2/impl/GrafanaApiServiceImpl.java

@@ -91,8 +91,8 @@ public Response variables(SQL sql, SecurityContext securityContext) {
     try {
       RequestValidationHandler.validateSQL(sql);
 
-      Statement statement =
-          StatementGenerator.createStatement(sql.getSql(), ZoneId.systemDefault());
+      ZoneId zoneId = SESSION_MANAGER.getCurrSession().getZoneId();
+      Statement statement = StatementGenerator.createStatement(sql.getSql(), zoneId);
       if (!(statement instanceof ShowStatement) && !(statement instanceof QueryStatement)) {
         return Response.ok()
             .entity(
@@ -168,7 +168,8 @@ public Response expression(ExpressionRequest expressionRequest, SecurityContext
         sql += " " + expressionRequest.getControl();
       }
 
-      Statement statement = StatementGenerator.createStatement(sql, ZoneId.systemDefault());
+      ZoneId zoneId = SESSION_MANAGER.getCurrSession().getZoneId();
+      Statement statement = StatementGenerator.createStatement(sql, zoneId);
 
       Response response = authorizationHandler.checkAuthority(securityContext, statement);
       if (response != null) {
@@ -232,7 +233,8 @@ public Response node(List<String> requestBody, SecurityContext securityContext)
         // TODO: necessary to create a PartialPath
         PartialPath path = new PartialPath(Joiner.on(".").join(requestBody));
         String sql = "show child paths " + path;
-        Statement statement = StatementGenerator.createStatement(sql, ZoneId.systemDefault());
+        ZoneId zoneId = SESSION_MANAGER.getCurrSession().getZoneId();
+        Statement statement = StatementGenerator.createStatement(sql, zoneId);
 
         Response response = authorizationHandler.checkAuthority(securityContext, statement);
         if (response != null) {

external-service-impl/rest/src/main/java/org/apache/iotdb/rest/protocol/v2/impl/RestApiServiceImpl.java

@@ -270,7 +270,8 @@ public Response executeNonQueryStatement(SQL sql, SecurityContext securityContex
     boolean finish = false;
     try {
       RequestValidationHandler.validateSQL(sql);
-      statement = StatementGenerator.createStatement(sql.getSql(), ZoneId.systemDefault());
+      ZoneId zoneId = SESSION_MANAGER.getCurrSession().getZoneId();
+      statement = StatementGenerator.createStatement(sql.getSql(), zoneId);
       if (statement == null) {
         return Response.ok()
             .entity(
@@ -310,9 +311,10 @@ public Response executeNonQueryStatement(SQL sql, SecurityContext securityContex
       return Response.ok().entity(ExceptionHandler.tryCatchException(e)).build();
     } finally {
       long costTime = System.nanoTime() - startTime;
-      if (statement != null)
+      if (statement != null) {
         CommonUtils.addStatementExecutionLatency(
             OperationType.EXECUTE_NON_QUERY_PLAN, statement.getType().name(), costTime);
+      }
       if (queryId != null) {
         if (finish) {
           long executionTime = COORDINATOR.getTotalExecutionTime(queryId);
@@ -332,7 +334,8 @@ public Response executeQueryStatement(SQL sql, SecurityContext securityContext)
     boolean finish = false;
     try {
       RequestValidationHandler.validateSQL(sql);
-      statement = StatementGenerator.createStatement(sql.getSql(), ZoneId.systemDefault());
+      ZoneId zoneId = SESSION_MANAGER.getCurrSession().getZoneId();
+      statement = StatementGenerator.createStatement(sql.getSql(), zoneId);
 
       if (statement == null) {
         return Response.ok()

integration-test/src/test/java/org/apache/iotdb/db/it/IoTDBRestServiceIT.java

@@ -27,6 +27,7 @@
 import org.apache.iotdb.itbase.category.LocalStandaloneIT;
 import org.apache.iotdb.itbase.category.RemoteIT;
 import org.apache.iotdb.itbase.env.BaseEnv;
+import org.apache.iotdb.rpc.TSStatusCode;
 
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.gson.JsonObject;
@@ -55,6 +56,8 @@
 import java.sql.ResultSetMetaData;
 import java.sql.SQLException;
 import java.sql.Statement;
+import java.time.ZoneId;
+import java.time.ZonedDateTime;
 import java.util.ArrayList;
 import java.util.Base64;
 import java.util.List;
@@ -2391,4 +2394,130 @@ public void queryDateAndBlobV2(CloseableHttpClient httpClient) {
       }
     }
   }
+
+  @Test
+  public void testQueryWithValidTimeZoneHeaderV2() {
+    CloseableHttpClient httpClient = HttpClientBuilder.create().build();
+    try {
+      HttpPost insertPost = getHttpPost("http://127.0.0.1:" + port + "/rest/v2/insertTablet");
+      String insertJson =
+          "{\"timestamps\":[1774713387626],\"measurements\":[\"s3\"],\"data_types\":[\"INT32\"],\"values\":[[11]],\"is_aligned\":false,\"device\":\"root.sg25\"}";
+      insertPost.setEntity(new StringEntity(insertJson, Charset.defaultCharset()));
+      try (CloseableHttpResponse resp = executeWithRetry(insertPost, httpClient)) {
+        assertEquals(200, resp.getStatusLine().getStatusCode());
+      }
+      HttpPost httpPost = getHttpPost("http://127.0.0.1:" + port + "/rest/v2/query");
+      httpPost.setHeader("X-TimeZone", "Europe/Warsaw");
+      String sql =
+          "{\"sql\":\"SELECT count(s3) FROM root.sg25 GROUP BY ([2026-03-28T00:00:00, 2026-03-29T00:00:00), 1d)\"}";
+      httpPost.setEntity(new StringEntity(sql, Charset.defaultCharset()));
+      CloseableHttpResponse response = httpClient.execute(httpPost);
+      assertEquals(200, response.getStatusLine().getStatusCode());
+      String message = EntityUtils.toString(response.getEntity(), "utf-8");
+      JsonObject result = JsonParser.parseString(message).getAsJsonObject();
+      assertTrue(result.has("timestamps"));
+      assertTrue(result.getAsJsonArray("timestamps").size() > 0);
+      long expectedTimestamp =
+          ZonedDateTime.of(2026, 3, 28, 0, 0, 0, 0, ZoneId.of("Europe/Warsaw"))
+              .toInstant()
+              .toEpochMilli();
+      assertEquals(expectedTimestamp, result.getAsJsonArray("timestamps").get(0).getAsLong());
+    } catch (IOException e) {
+      fail(e.getMessage());
+    } finally {
+      try {
+        httpClient.close();
+      } catch (IOException e) {
+      }
+    }
+  }
+
+  @Test
+  public void testNonQueryWithValidTimeZoneHeaderV2() throws Exception {
+    CloseableHttpClient httpClient = HttpClientBuilder.create().build();
+    try {
+      nonQueryWithTimeZone(
+          httpClient,
+          "{\"sql\":\"CREATE TIMESERIES root.sg.d1.s1 WITH DATATYPE=INT32\"}",
+          "Europe/Warsaw");
+      nonQueryWithTimeZone(
+          httpClient,
+          "{\"sql\":\"INSERT INTO root.sg.d1(time, s1) VALUES (2026-03-28T00:00:00, 123)\"}",
+          "Europe/Warsaw");
+
+      HttpPost queryPost = getHttpPost("http://127.0.0.1:" + port + "/rest/v2/query");
+      queryPost.setEntity(
+          new StringEntity("{\"sql\":\"SELECT s1 FROM root.sg.d1\"}", StandardCharsets.UTF_8));
+      try (CloseableHttpResponse resp = httpClient.execute(queryPost)) {
+        String message = EntityUtils.toString(resp.getEntity(), StandardCharsets.UTF_8);
+        JsonObject result = JsonParser.parseString(message).getAsJsonObject();
+        long expected =
+            ZonedDateTime.of(2026, 3, 28, 0, 0, 0, 0, ZoneId.of("Europe/Warsaw"))
+                .toInstant()
+                .toEpochMilli();
+        assertEquals(expected, result.getAsJsonArray("timestamps").get(0).getAsLong());
+      }
+    } finally {
+      try {
+        httpClient.close();
+      } catch (IOException e) {
+      }
+    }
+  }
+
+  @Test
+  public void testQueryWithInvalidTimeZoneHeaderV2() {
+    CloseableHttpClient httpClient = HttpClientBuilder.create().build();
+    try {
+      HttpPost httpPost = getHttpPost("http://127.0.0.1:" + port + "/rest/v2/query");
+      httpPost.setHeader("X-TimeZone", "Invalid/Zone");
+      String sql = "{\"sql\":\"SELECT s3 FROM root.sg25\"}";
+      httpPost.setEntity(new StringEntity(sql, Charset.defaultCharset()));
+      CloseableHttpResponse response = executeWithRetry(httpPost, httpClient);
+      assertEquals(400, response.getStatusLine().getStatusCode());
+      String message = EntityUtils.toString(response.getEntity(), "utf-8");
+      JsonObject result = JsonParser.parseString(message).getAsJsonObject();
+      assertEquals(TSStatusCode.ILLEGAL_PARAMETER.getStatusCode(), result.get("code").getAsInt());
+      assertTrue(result.get("message").getAsString().contains("Invalid time zone"));
+    } catch (IOException e) {
+      fail(e.getMessage());
+    } finally {
+      try {
+        httpClient.close();
+      } catch (IOException e) {
+      }
+    }
+  }
+
+  private void nonQueryWithTimeZone(CloseableHttpClient httpClient, String json, String timeZone) {
+    HttpPost httpPost = getHttpPost("http://127.0.0.1:" + port + "/rest/v2/nonQuery");
+    httpPost.setHeader("X-TimeZone", timeZone);
+    httpPost.setEntity(new StringEntity(json, StandardCharsets.UTF_8));
+    try (CloseableHttpResponse response = executeWithRetry(httpPost, httpClient)) {
+      String message = EntityUtils.toString(response.getEntity(), StandardCharsets.UTF_8);
+      JsonObject result = JsonParser.parseString(message).getAsJsonObject();
+      assertEquals(200, result.get("code").getAsInt());
+    } catch (IOException e) {
+      fail(e.getMessage());
+    }
+  }
+
+  private CloseableHttpResponse executeWithRetry(HttpPost httpPost, CloseableHttpClient httpClient)
+      throws IOException {
+    CloseableHttpResponse response = null;
+    for (int i = 0; i < 30; i++) {
+      try {
+        response = httpClient.execute(httpPost);
+        break;
+      } catch (Exception e) {
+        if (i == 29) throw e;
+        try {
+          Thread.sleep(1000);
+        } catch (InterruptedException ex) {
+          throw new RuntimeException(ex);
+        }
+      }
+    }
+    return response;
+  }
 }