Skip to content

Commit dacddb6

Browse files
Lexert19Lexert19
Lexert19
authored and
Lexert19
committed
Add X-TimeZone header support to REST API (#17344)
1 parent 58ec38a commit dacddb6

File tree

4 files changed

+176
-8
lines changed

4 files changed

+176
-8
lines changed

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

Lines changed: 35 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@
3939
import javax.ws.rs.ext.Provider;
4040

4141
import java.io.IOException;
42+
import java.time.DateTimeException;
4243
import java.time.ZoneId;
4344
import java.util.Base64;
4445
import java.util.UUID;
@@ -88,6 +89,12 @@ public void filter(ContainerRequestContext containerRequestContext) throws IOExc
8889
if (user == null) {
8990
return;
9091
}
92+
93+
ZoneId zoneId = parseTimeZone(containerRequestContext);
94+
if (zoneId == null) {
95+
return;
96+
}
97+
9198
String sessionid = UUID.randomUUID().toString();
9299
if (SESSION_MANAGER.getCurrSession() == null) {
93100
RestClientSession restClientSession = new RestClientSession(sessionid);
@@ -97,7 +104,7 @@ public void filter(ContainerRequestContext containerRequestContext) throws IOExc
97104
SESSION_MANAGER.getCurrSession(),
98105
user.getUserId(),
99106
user.getUsername(),
100-
ZoneId.systemDefault(),
107+
zoneId,
101108
IoTDBConstant.ClientVersion.V_1_0);
102109
}
103110
BasicSecurityContext basicSecurityContext =
@@ -147,6 +154,33 @@ private User checkLogin(
147154
return user;
148155
}
149156

157+
/**
158+
* Parses the X-TimeZone header from the request.
159+
*
160+
* @param requestContext the incoming HTTP request
161+
* @return the parsed ZoneId, or {@code null} if the header is invalid (the request is aborted)
162+
*/
163+
private ZoneId parseTimeZone(ContainerRequestContext requestContext) {
164+
String timeZoneHeader = requestContext.getHeaderString("X-TimeZone");
165+
if (timeZoneHeader == null || timeZoneHeader.isEmpty()) {
166+
return ZoneId.systemDefault();
167+
}
168+
try {
169+
return ZoneId.of(timeZoneHeader);
170+
} catch (DateTimeException e) {
171+
Response resp =
172+
Response.status(Status.BAD_REQUEST)
173+
.type(MediaType.APPLICATION_JSON)
174+
.entity(
175+
new ExecutionStatus()
176+
.code(TSStatusCode.ILLEGAL_PARAMETER.getStatusCode())
177+
.message("Invalid time zone: " + timeZoneHeader))
178+
.build();
179+
requestContext.abortWith(resp);
180+
return null;
181+
}
182+
}
183+
150184
@Override
151185
public void filter(
152186
ContainerRequestContext requestContext, ContainerResponseContext responseContext)

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

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -91,8 +91,8 @@ public Response variables(SQL sql, SecurityContext securityContext) {
9191
try {
9292
RequestValidationHandler.validateSQL(sql);
9393

94-
Statement statement =
95-
StatementGenerator.createStatement(sql.getSql(), ZoneId.systemDefault());
94+
ZoneId zoneId = SESSION_MANAGER.getCurrSession().getZoneId();
95+
Statement statement = StatementGenerator.createStatement(sql.getSql(), zoneId);
9696
if (!(statement instanceof ShowStatement) && !(statement instanceof QueryStatement)) {
9797
return Response.ok()
9898
.entity(
@@ -168,7 +168,8 @@ public Response expression(ExpressionRequest expressionRequest, SecurityContext
168168
sql += " " + expressionRequest.getControl();
169169
}
170170

171-
Statement statement = StatementGenerator.createStatement(sql, ZoneId.systemDefault());
171+
ZoneId zoneId = SESSION_MANAGER.getCurrSession().getZoneId();
172+
Statement statement = StatementGenerator.createStatement(sql, zoneId);
172173

173174
Response response = authorizationHandler.checkAuthority(securityContext, statement);
174175
if (response != null) {
@@ -232,7 +233,8 @@ public Response node(List<String> requestBody, SecurityContext securityContext)
232233
// TODO: necessary to create a PartialPath
233234
PartialPath path = new PartialPath(Joiner.on(".").join(requestBody));
234235
String sql = "show child paths " + path;
235-
Statement statement = StatementGenerator.createStatement(sql, ZoneId.systemDefault());
236+
ZoneId zoneId = SESSION_MANAGER.getCurrSession().getZoneId();
237+
Statement statement = StatementGenerator.createStatement(sql, zoneId);
236238

237239
Response response = authorizationHandler.checkAuthority(securityContext, statement);
238240
if (response != null) {

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

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -270,7 +270,8 @@ public Response executeNonQueryStatement(SQL sql, SecurityContext securityContex
270270
boolean finish = false;
271271
try {
272272
RequestValidationHandler.validateSQL(sql);
273-
statement = StatementGenerator.createStatement(sql.getSql(), ZoneId.systemDefault());
273+
ZoneId zoneId = SESSION_MANAGER.getCurrSession().getZoneId();
274+
statement = StatementGenerator.createStatement(sql.getSql(), zoneId);
274275
if (statement == null) {
275276
return Response.ok()
276277
.entity(
@@ -310,9 +311,10 @@ public Response executeNonQueryStatement(SQL sql, SecurityContext securityContex
310311
return Response.ok().entity(ExceptionHandler.tryCatchException(e)).build();
311312
} finally {
312313
long costTime = System.nanoTime() - startTime;
313-
if (statement != null)
314+
if (statement != null) {
314315
CommonUtils.addStatementExecutionLatency(
315316
OperationType.EXECUTE_NON_QUERY_PLAN, statement.getType().name(), costTime);
317+
}
316318
if (queryId != null) {
317319
if (finish) {
318320
long executionTime = COORDINATOR.getTotalExecutionTime(queryId);
@@ -332,7 +334,8 @@ public Response executeQueryStatement(SQL sql, SecurityContext securityContext)
332334
boolean finish = false;
333335
try {
334336
RequestValidationHandler.validateSQL(sql);
335-
statement = StatementGenerator.createStatement(sql.getSql(), ZoneId.systemDefault());
337+
ZoneId zoneId = SESSION_MANAGER.getCurrSession().getZoneId();
338+
statement = StatementGenerator.createStatement(sql.getSql(), zoneId);
336339

337340
if (statement == null) {
338341
return Response.ok()

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

Lines changed: 129 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@
2727
import org.apache.iotdb.itbase.category.LocalStandaloneIT;
2828
import org.apache.iotdb.itbase.category.RemoteIT;
2929
import org.apache.iotdb.itbase.env.BaseEnv;
30+
import org.apache.iotdb.rpc.TSStatusCode;
3031

3132
import com.fasterxml.jackson.databind.ObjectMapper;
3233
import com.google.gson.JsonObject;
@@ -55,6 +56,8 @@
5556
import java.sql.ResultSetMetaData;
5657
import java.sql.SQLException;
5758
import java.sql.Statement;
59+
import java.time.ZoneId;
60+
import java.time.ZonedDateTime;
5861
import java.util.ArrayList;
5962
import java.util.Base64;
6063
import java.util.List;
@@ -2391,4 +2394,130 @@ public void queryDateAndBlobV2(CloseableHttpClient httpClient) {
23912394
}
23922395
}
23932396
}
2397+
2398+
@Test
2399+
public void testQueryWithValidTimeZoneHeaderV2() {
2400+
CloseableHttpClient httpClient = HttpClientBuilder.create().build();
2401+
try {
2402+
HttpPost insertPost = getHttpPost("http://127.0.0.1:" + port + "/rest/v2/insertTablet");
2403+
String insertJson =
2404+
"{\"timestamps\":[1774713387626],\"measurements\":[\"s3\"],\"data_types\":[\"INT32\"],\"values\":[[11]],\"is_aligned\":false,\"device\":\"root.sg25\"}";
2405+
insertPost.setEntity(new StringEntity(insertJson, Charset.defaultCharset()));
2406+
try (CloseableHttpResponse resp = executeWithRetry(insertPost, httpClient)) {
2407+
assertEquals(200, resp.getStatusLine().getStatusCode());
2408+
}
2409+
HttpPost httpPost = getHttpPost("http://127.0.0.1:" + port + "/rest/v2/query");
2410+
httpPost.setHeader("X-TimeZone", "Europe/Warsaw");
2411+
String sql =
2412+
"{\"sql\":\"SELECT count(s3) FROM root.sg25 GROUP BY ([2026-03-28T00:00:00, 2026-03-29T00:00:00), 1d)\"}";
2413+
httpPost.setEntity(new StringEntity(sql, Charset.defaultCharset()));
2414+
CloseableHttpResponse response = httpClient.execute(httpPost);
2415+
assertEquals(200, response.getStatusLine().getStatusCode());
2416+
String message = EntityUtils.toString(response.getEntity(), "utf-8");
2417+
JsonObject result = JsonParser.parseString(message).getAsJsonObject();
2418+
assertTrue(result.has("timestamps"));
2419+
assertTrue(result.getAsJsonArray("timestamps").size() > 0);
2420+
long expectedTimestamp =
2421+
ZonedDateTime.of(2026, 3, 28, 0, 0, 0, 0, ZoneId.of("Europe/Warsaw"))
2422+
.toInstant()
2423+
.toEpochMilli();
2424+
assertEquals(expectedTimestamp, result.getAsJsonArray("timestamps").get(0).getAsLong());
2425+
} catch (IOException e) {
2426+
fail(e.getMessage());
2427+
} finally {
2428+
try {
2429+
httpClient.close();
2430+
} catch (IOException e) {
2431+
}
2432+
}
2433+
}
2434+
2435+
@Test
2436+
public void testNonQueryWithValidTimeZoneHeaderV2() throws Exception {
2437+
CloseableHttpClient httpClient = HttpClientBuilder.create().build();
2438+
try {
2439+
nonQueryWithTimeZone(
2440+
httpClient,
2441+
"{\"sql\":\"CREATE TIMESERIES root.sg.d1.s1 WITH DATATYPE=INT32\"}",
2442+
"Europe/Warsaw");
2443+
nonQueryWithTimeZone(
2444+
httpClient,
2445+
"{\"sql\":\"INSERT INTO root.sg.d1(time, s1) VALUES (2026-03-28T00:00:00, 123)\"}",
2446+
"Europe/Warsaw");
2447+
2448+
HttpPost queryPost = getHttpPost("http://127.0.0.1:" + port + "/rest/v2/query");
2449+
queryPost.setEntity(
2450+
new StringEntity("{\"sql\":\"SELECT s1 FROM root.sg.d1\"}", StandardCharsets.UTF_8));
2451+
try (CloseableHttpResponse resp = httpClient.execute(queryPost)) {
2452+
String message = EntityUtils.toString(resp.getEntity(), StandardCharsets.UTF_8);
2453+
JsonObject result = JsonParser.parseString(message).getAsJsonObject();
2454+
long expected =
2455+
ZonedDateTime.of(2026, 3, 28, 0, 0, 0, 0, ZoneId.of("Europe/Warsaw"))
2456+
.toInstant()
2457+
.toEpochMilli();
2458+
assertEquals(expected, result.getAsJsonArray("timestamps").get(0).getAsLong());
2459+
}
2460+
} finally {
2461+
try {
2462+
httpClient.close();
2463+
} catch (IOException e) {
2464+
}
2465+
}
2466+
}
2467+
2468+
@Test
2469+
public void testQueryWithInvalidTimeZoneHeaderV2() {
2470+
CloseableHttpClient httpClient = HttpClientBuilder.create().build();
2471+
try {
2472+
HttpPost httpPost = getHttpPost("http://127.0.0.1:" + port + "/rest/v2/query");
2473+
httpPost.setHeader("X-TimeZone", "Invalid/Zone");
2474+
String sql = "{\"sql\":\"SELECT s3 FROM root.sg25\"}";
2475+
httpPost.setEntity(new StringEntity(sql, Charset.defaultCharset()));
2476+
CloseableHttpResponse response = executeWithRetry(httpPost, httpClient);
2477+
assertEquals(400, response.getStatusLine().getStatusCode());
2478+
String message = EntityUtils.toString(response.getEntity(), "utf-8");
2479+
JsonObject result = JsonParser.parseString(message).getAsJsonObject();
2480+
assertEquals(TSStatusCode.ILLEGAL_PARAMETER.getStatusCode(), result.get("code").getAsInt());
2481+
assertTrue(result.get("message").getAsString().contains("Invalid time zone"));
2482+
} catch (IOException e) {
2483+
fail(e.getMessage());
2484+
} finally {
2485+
try {
2486+
httpClient.close();
2487+
} catch (IOException e) {
2488+
}
2489+
}
2490+
}
2491+
2492+
private void nonQueryWithTimeZone(CloseableHttpClient httpClient, String json, String timeZone) {
2493+
HttpPost httpPost = getHttpPost("http://127.0.0.1:" + port + "/rest/v2/nonQuery");
2494+
httpPost.setHeader("X-TimeZone", timeZone);
2495+
httpPost.setEntity(new StringEntity(json, StandardCharsets.UTF_8));
2496+
try (CloseableHttpResponse response = executeWithRetry(httpPost, httpClient)) {
2497+
String message = EntityUtils.toString(response.getEntity(), StandardCharsets.UTF_8);
2498+
JsonObject result = JsonParser.parseString(message).getAsJsonObject();
2499+
assertEquals(200, result.get("code").getAsInt());
2500+
} catch (IOException e) {
2501+
fail(e.getMessage());
2502+
}
2503+
}
2504+
2505+
private CloseableHttpResponse executeWithRetry(HttpPost httpPost, CloseableHttpClient httpClient)
2506+
throws IOException {
2507+
CloseableHttpResponse response = null;
2508+
for (int i = 0; i < 30; i++) {
2509+
try {
2510+
response = httpClient.execute(httpPost);
2511+
break;
2512+
} catch (Exception e) {
2513+
if (i == 29) throw e;
2514+
try {
2515+
Thread.sleep(1000);
2516+
} catch (InterruptedException ex) {
2517+
throw new RuntimeException(ex);
2518+
}
2519+
}
2520+
}
2521+
return response;
2522+
}
23942523
}

0 commit comments

Comments
 (0)