Skip to content

Commit a678291

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

4 files changed

Lines changed: 176 additions & 8 deletions

File tree

‎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
@@ -26,6 +26,7 @@
2626
import org.apache.iotdb.itbase.category.LocalStandaloneIT;
2727
import org.apache.iotdb.itbase.category.RemoteIT;
2828
import org.apache.iotdb.itbase.env.BaseEnv;
29+
import org.apache.iotdb.rpc.TSStatusCode;
2930

3031
import com.fasterxml.jackson.databind.ObjectMapper;
3132
import com.google.gson.JsonObject;
@@ -52,6 +53,8 @@
5253
import java.sql.ResultSet;
5354
import java.sql.SQLException;
5455
import java.sql.Statement;
56+
import java.time.ZoneId;
57+
import java.time.ZonedDateTime;
5558
import java.util.ArrayList;
5659
import java.util.Base64;
5760
import java.util.List;
@@ -2354,4 +2357,130 @@ public void queryDateAndBlobV2(CloseableHttpClient httpClient) {
23542357
}
23552358
}
23562359
}
2360+
2361+
@Test
2362+
public void testQueryWithValidTimeZoneHeaderV2() {
2363+
CloseableHttpClient httpClient = HttpClientBuilder.create().build();
2364+
try {
2365+
HttpPost insertPost = getHttpPost("http://127.0.0.1:" + port + "/rest/v2/insertTablet");
2366+
String insertJson =
2367+
"{\"timestamps\":[1774713387626],\"measurements\":[\"s3\"],\"data_types\":[\"INT32\"],\"values\":[[11]],\"is_aligned\":false,\"device\":\"root.sg25\"}";
2368+
insertPost.setEntity(new StringEntity(insertJson, Charset.defaultCharset()));
2369+
try (CloseableHttpResponse resp = executeWithRetry(insertPost, httpClient)) {
2370+
assertEquals(200, resp.getStatusLine().getStatusCode());
2371+
}
2372+
HttpPost httpPost = getHttpPost("http://127.0.0.1:" + port + "/rest/v2/query");
2373+
httpPost.setHeader("X-TimeZone", "Europe/Warsaw");
2374+
String sql =
2375+
"{\"sql\":\"SELECT count(s3) FROM root.sg25 GROUP BY ([2026-03-28T00:00:00, 2026-03-29T00:00:00), 1d)\"}";
2376+
httpPost.setEntity(new StringEntity(sql, Charset.defaultCharset()));
2377+
CloseableHttpResponse response = httpClient.execute(httpPost);
2378+
assertEquals(200, response.getStatusLine().getStatusCode());
2379+
String message = EntityUtils.toString(response.getEntity(), "utf-8");
2380+
JsonObject result = JsonParser.parseString(message).getAsJsonObject();
2381+
assertTrue(result.has("timestamps"));
2382+
assertTrue(result.getAsJsonArray("timestamps").size() > 0);
2383+
long expectedTimestamp =
2384+
ZonedDateTime.of(2026, 3, 28, 0, 0, 0, 0, ZoneId.of("Europe/Warsaw"))
2385+
.toInstant()
2386+
.toEpochMilli();
2387+
assertEquals(expectedTimestamp, result.getAsJsonArray("timestamps").get(0).getAsLong());
2388+
} catch (IOException e) {
2389+
fail(e.getMessage());
2390+
} finally {
2391+
try {
2392+
httpClient.close();
2393+
} catch (IOException e) {
2394+
}
2395+
}
2396+
}
2397+
2398+
@Test
2399+
public void testNonQueryWithValidTimeZoneHeaderV2() throws Exception {
2400+
CloseableHttpClient httpClient = HttpClientBuilder.create().build();
2401+
try {
2402+
nonQueryWithTimeZone(
2403+
httpClient,
2404+
"{\"sql\":\"CREATE TIMESERIES root.sg.d1.s1 WITH DATATYPE=INT32\"}",
2405+
"Europe/Warsaw");
2406+
nonQueryWithTimeZone(
2407+
httpClient,
2408+
"{\"sql\":\"INSERT INTO root.sg.d1(time, s1) VALUES (2026-03-28T00:00:00, 123)\"}",
2409+
"Europe/Warsaw");
2410+
2411+
HttpPost queryPost = getHttpPost("http://127.0.0.1:" + port + "/rest/v2/query");
2412+
queryPost.setEntity(
2413+
new StringEntity("{\"sql\":\"SELECT s1 FROM root.sg.d1\"}", StandardCharsets.UTF_8));
2414+
try (CloseableHttpResponse resp = httpClient.execute(queryPost)) {
2415+
String message = EntityUtils.toString(resp.getEntity(), StandardCharsets.UTF_8);
2416+
JsonObject result = JsonParser.parseString(message).getAsJsonObject();
2417+
long expected =
2418+
ZonedDateTime.of(2026, 3, 28, 0, 0, 0, 0, ZoneId.of("Europe/Warsaw"))
2419+
.toInstant()
2420+
.toEpochMilli();
2421+
assertEquals(expected, result.getAsJsonArray("timestamps").get(0).getAsLong());
2422+
}
2423+
} finally {
2424+
try {
2425+
httpClient.close();
2426+
} catch (IOException e) {
2427+
}
2428+
}
2429+
}
2430+
2431+
@Test
2432+
public void testQueryWithInvalidTimeZoneHeaderV2() {
2433+
CloseableHttpClient httpClient = HttpClientBuilder.create().build();
2434+
try {
2435+
HttpPost httpPost = getHttpPost("http://127.0.0.1:" + port + "/rest/v2/query");
2436+
httpPost.setHeader("X-TimeZone", "Invalid/Zone");
2437+
String sql = "{\"sql\":\"SELECT s3 FROM root.sg25\"}";
2438+
httpPost.setEntity(new StringEntity(sql, Charset.defaultCharset()));
2439+
CloseableHttpResponse response = executeWithRetry(httpPost, httpClient);
2440+
assertEquals(400, response.getStatusLine().getStatusCode());
2441+
String message = EntityUtils.toString(response.getEntity(), "utf-8");
2442+
JsonObject result = JsonParser.parseString(message).getAsJsonObject();
2443+
assertEquals(TSStatusCode.ILLEGAL_PARAMETER.getStatusCode(), result.get("code").getAsInt());
2444+
assertTrue(result.get("message").getAsString().contains("Invalid time zone"));
2445+
} catch (IOException e) {
2446+
fail(e.getMessage());
2447+
} finally {
2448+
try {
2449+
httpClient.close();
2450+
} catch (IOException e) {
2451+
}
2452+
}
2453+
}
2454+
2455+
private void nonQueryWithTimeZone(CloseableHttpClient httpClient, String json, String timeZone) {
2456+
HttpPost httpPost = getHttpPost("http://127.0.0.1:" + port + "/rest/v2/nonQuery");
2457+
httpPost.setHeader("X-TimeZone", timeZone);
2458+
httpPost.setEntity(new StringEntity(json, StandardCharsets.UTF_8));
2459+
try (CloseableHttpResponse response = executeWithRetry(httpPost, httpClient)) {
2460+
String message = EntityUtils.toString(response.getEntity(), StandardCharsets.UTF_8);
2461+
JsonObject result = JsonParser.parseString(message).getAsJsonObject();
2462+
assertEquals(200, result.get("code").getAsInt());
2463+
} catch (IOException e) {
2464+
fail(e.getMessage());
2465+
}
2466+
}
2467+
2468+
private CloseableHttpResponse executeWithRetry(HttpPost httpPost, CloseableHttpClient httpClient)
2469+
throws IOException {
2470+
CloseableHttpResponse response = null;
2471+
for (int i = 0; i < 30; i++) {
2472+
try {
2473+
response = httpClient.execute(httpPost);
2474+
break;
2475+
} catch (Exception e) {
2476+
if (i == 29) throw e;
2477+
try {
2478+
Thread.sleep(1000);
2479+
} catch (InterruptedException ex) {
2480+
throw new RuntimeException(ex);
2481+
}
2482+
}
2483+
}
2484+
return response;
2485+
}
23572486
}

0 commit comments

Comments
 (0)