Skip to content

Commit 632cd91

Browse files
ivanceaivancea
committed
Pass ProfileLogger as a parameter
1 parent 4fe71c9 commit 632cd91

File tree

4 files changed

+147
-75
lines changed

4 files changed

+147
-75
lines changed

‎x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/EsqlSpecTestCase.java‎

Lines changed: 5 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -39,8 +39,6 @@
3939
import org.junit.AfterClass;
4040
import org.junit.Before;
4141
import org.junit.Rule;
42-
import org.junit.rules.TestWatcher;
43-
import org.junit.runner.Description;
4442

4543
import java.io.IOException;
4644
import java.math.BigDecimal;
@@ -80,27 +78,12 @@
8078
import static org.elasticsearch.xpack.esql.action.EsqlCapabilities.Cap.RERANK;
8179
import static org.elasticsearch.xpack.esql.action.EsqlCapabilities.Cap.SEMANTIC_TEXT_FIELD_CAPS;
8280
import static org.elasticsearch.xpack.esql.action.EsqlCapabilities.Cap.SOURCE_FIELD_MAPPING;
83-
import static org.hamcrest.Matchers.anyOf;
84-
import static org.hamcrest.Matchers.is;
85-
import static org.hamcrest.Matchers.nullValue;
81+
import static org.elasticsearch.xpack.esql.qa.rest.RestEsqlTestCase.assertNotPartial;
8682

8783
// This test can run very long in serverless configurations
8884
@TimeoutSuite(millis = 30 * TimeUnits.MINUTE)
8985
public abstract class EsqlSpecTestCase extends ESRestTestCase {
9086

91-
public class ProfileLogger extends TestWatcher {
92-
private Object profile;
93-
94-
void setProfile(RestEsqlTestCase.EsqlResponse response) {
95-
profile = response.json().get("profile");
96-
}
97-
98-
@Override
99-
protected void failed(Throwable e, Description description) {
100-
LOGGER.warn("Profile: {}", profile);
101-
}
102-
}
103-
10487
@Rule(order = Integer.MIN_VALUE)
10588
public ProfileLogger profileLogger = new ProfileLogger();
10689

@@ -321,12 +304,6 @@ protected final void doTest(String query) throws Throwable {
321304
}
322305
}
323306

324-
static void assertNotPartial(Map<String, Object> answer) {
325-
var clusters = answer.get("_clusters");
326-
var reason = "unexpected partial results" + (clusters != null ? ": _clusters=" + clusters : "");
327-
assertThat(reason, answer.get("is_partial"), anyOf(nullValue(), is(false)));
328-
}
329-
330307
private Map<?, ?> tooks() throws IOException {
331308
Request request = new Request("GET", "/_xpack/usage");
332309
HttpEntity entity = client().performRequest(request).getEntity();
@@ -354,20 +331,11 @@ protected boolean deduplicateExactWarnings() {
354331
}
355332

356333
private Map<String, Object> runEsql(RequestObjectBuilder requestObject, AssertWarnings assertWarnings) throws IOException {
357-
requestObject.profile(true);
358-
359-
RestEsqlTestCase.EsqlResponse response = mode == Mode.ASYNC
360-
? RestEsqlTestCase.runEsqlAsyncNoWarningsChecks(requestObject, randomBoolean())
361-
: RestEsqlTestCase.runEsqlSyncNoWarningsChecks(requestObject);
362-
363-
profileLogger.setProfile(response);
364-
365-
RestEsqlTestCase.assertWarnings(response.response(), assertWarnings);
366-
if (response.asyncInitialResponse() != response.response()) {
367-
RestEsqlTestCase.assertWarnings(response.response(), assertWarnings);
334+
if (mode == Mode.ASYNC) {
335+
return RestEsqlTestCase.runEsqlAsync(requestObject, assertWarnings, profileLogger);
336+
} else {
337+
return RestEsqlTestCase.runEsqlSync(requestObject, assertWarnings, profileLogger);
368338
}
369-
370-
return response.json();
371339
}
372340

373341
protected void assertResults(

‎x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/ProfileLogger.java‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the Elastic License
4+
* 2.0; you may not use this file except in compliance with the Elastic License
5+
* 2.0.
6+
*/
7+
8+
package org.elasticsearch.xpack.esql.qa.rest;
9+
10+
import org.elasticsearch.logging.LogManager;
11+
import org.elasticsearch.logging.Logger;
12+
import org.junit.rules.TestWatcher;
13+
import org.junit.runner.Description;
14+
15+
import java.util.Map;
16+
17+
/**
18+
* A JUnit rule that logs the profile information from an ESQL response if the test fails.
19+
* <p>
20+
* Usage:
21+
* </p>
22+
* <pre>
23+
* {@code
24+
* @Rule(order = Integer.MIN_VALUE)
25+
* public ProfileLogger profileLogger = new ProfileLogger();
26+
*
27+
* public void test() {
28+
* var response = RestEsqlTestCase.runEsqlSync(..., profileLogger);
29+
* }
30+
* }
31+
* </pre>
32+
*/
33+
public class ProfileLogger extends TestWatcher {
34+
private static final Logger LOGGER = LogManager.getLogger(ProfileLogger.class);
35+
36+
private Object profile;
37+
38+
void extractProfile(Map<String, Object> jsonResponse, Boolean originalProfileParameter) {
39+
if (jsonResponse.containsKey("profile") == false) {
40+
return;
41+
}
42+
43+
profile = jsonResponse.get("profile");
44+
45+
if (Boolean.TRUE.equals(originalProfileParameter) == false) {
46+
jsonResponse.remove("profile");
47+
}
48+
}
49+
50+
@Override
51+
protected void failed(Throwable e, Description description) {
52+
LOGGER.info("Profile: {}", profile);
53+
}
54+
}

‎x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/RestEsqlTestCase.java‎

Lines changed: 79 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -60,11 +60,11 @@
6060
import static org.elasticsearch.test.MapMatcher.assertMap;
6161
import static org.elasticsearch.test.MapMatcher.matchesMap;
6262
import static org.elasticsearch.xpack.esql.EsqlTestUtils.as;
63-
import static org.elasticsearch.xpack.esql.qa.rest.EsqlSpecTestCase.assertNotPartial;
6463
import static org.elasticsearch.xpack.esql.qa.rest.RestEsqlTestCase.Mode.ASYNC;
6564
import static org.elasticsearch.xpack.esql.qa.rest.RestEsqlTestCase.Mode.SYNC;
6665
import static org.elasticsearch.xpack.esql.type.EsqlDataTypeConverter.dateTimeToString;
6766
import static org.hamcrest.Matchers.any;
67+
import static org.hamcrest.Matchers.anyOf;
6868
import static org.hamcrest.Matchers.containsString;
6969
import static org.hamcrest.Matchers.either;
7070
import static org.hamcrest.Matchers.emptyOrNullString;
@@ -1256,10 +1256,20 @@ public static Map<String, Object> runEsql(
12561256
AssertWarnings assertWarnings,
12571257
Mode mode,
12581258
boolean checkPartialResults
1259+
) throws IOException {
1260+
return runEsql(requestObject, assertWarnings, null, mode, checkPartialResults);
1261+
}
1262+
1263+
public static Map<String, Object> runEsql(
1264+
RequestObjectBuilder requestObject,
1265+
AssertWarnings assertWarnings,
1266+
@Nullable ProfileLogger profileLogger,
1267+
Mode mode,
1268+
boolean checkPartialResults
12591269
) throws IOException {
12601270
var results = mode == ASYNC
1261-
? runEsqlAsync(requestObject, randomBoolean(), assertWarnings)
1262-
: runEsqlSync(requestObject, assertWarnings);
1271+
? runEsqlAsync(requestObject, randomBoolean(), assertWarnings, profileLogger)
1272+
: runEsqlSync(requestObject, assertWarnings, profileLogger);
12631273
if (checkPartialResults) {
12641274
assertNotPartial(results);
12651275
}
@@ -1271,53 +1281,70 @@ public static Map<String, Object> runEsql(RequestObjectBuilder requestObject, As
12711281
return runEsql(requestObject, assertWarnings, mode, true);
12721282
}
12731283

1274-
/**
1275-
* A response from an ESQL query.
1276-
* @param response the HTTP response that contains the JSON response body
1277-
* @param json the parsed JSON response body
1278-
* @param asyncInitialResponse for async only. The response that was returned from the initial async request.
1279-
* May be the same as {@code response} if the async request completed immediately.
1280-
*/
1281-
public record EsqlResponse(Response response, Map<String, Object> json, Response asyncInitialResponse) {
1282-
public EsqlResponse(Response response, Map<String, Object> json) {
1283-
this(response, json, null);
1284-
}
1284+
public static Map<String, Object> runEsql(RequestObjectBuilder requestObject, AssertWarnings assertWarnings,
1285+
@Nullable ProfileLogger profileLogger, Mode mode)
1286+
throws IOException {
1287+
return runEsql(requestObject, assertWarnings, profileLogger, mode, true);
12851288
}
12861289

12871290
public static Map<String, Object> runEsqlSync(RequestObjectBuilder requestObject, AssertWarnings assertWarnings) throws IOException {
1288-
EsqlResponse response = runEsqlSyncNoWarningsChecks(requestObject);
1289-
assertWarnings(response.response, assertWarnings);
1290-
return response.json;
1291+
return runEsqlSync(requestObject, assertWarnings, null);
12911292
}
12921293

1293-
public static EsqlResponse runEsqlSyncNoWarningsChecks(RequestObjectBuilder requestObject) throws IOException {
1294+
public static Map<String, Object> runEsqlSync(
1295+
RequestObjectBuilder requestObject,
1296+
AssertWarnings assertWarnings,
1297+
@Nullable ProfileLogger profileLogger
1298+
) throws IOException {
1299+
Boolean profileEnabled = requestObject.profile;
1300+
if (profileLogger != null) {
1301+
requestObject.profile(true);
1302+
}
12941303
Request request = prepareRequestWithOptions(requestObject, SYNC);
1304+
12951305
Response response = performRequest(request);
1296-
Map<String, Object> json = entityToMap(response.getEntity(), requestObject.contentType());
1297-
return new EsqlResponse(response, json);
1306+
HttpEntity entity = response.getEntity();
1307+
Map<String, Object> json = entityToMap(entity, requestObject.contentType());
1308+
1309+
if (profileLogger != null) {
1310+
profileLogger.extractProfile(json, profileEnabled);
1311+
}
1312+
1313+
assertWarnings(response, assertWarnings);
1314+
1315+
return json;
12981316
}
12991317

13001318
public static Map<String, Object> runEsqlAsync(RequestObjectBuilder requestObject, AssertWarnings assertWarnings) throws IOException {
13011319
return runEsqlAsync(requestObject, randomBoolean(), assertWarnings);
13021320
}
13031321

1322+
public static Map<String, Object> runEsqlAsync(
1323+
RequestObjectBuilder requestObject,
1324+
AssertWarnings assertWarnings,
1325+
@Nullable ProfileLogger profileLogger
1326+
) throws IOException {
1327+
return runEsqlAsync(requestObject, randomBoolean(), assertWarnings, profileLogger);
1328+
}
1329+
13041330
public static Map<String, Object> runEsqlAsync(
13051331
RequestObjectBuilder requestObject,
13061332
boolean keepOnCompletion,
13071333
AssertWarnings assertWarnings
13081334
) throws IOException {
1309-
EsqlResponse response = runEsqlAsyncNoWarningsChecks(requestObject, keepOnCompletion);
1310-
1311-
assertWarnings(response.response(), assertWarnings);
1312-
if (response.asyncInitialResponse() != response.response()) {
1313-
assertWarnings(response.asyncInitialResponse(), assertWarnings);
1314-
}
1315-
1316-
return response.json;
1335+
return runEsqlAsync(requestObject, keepOnCompletion, assertWarnings, null);
13171336
}
13181337

1319-
public static EsqlResponse runEsqlAsyncNoWarningsChecks(RequestObjectBuilder requestObject, boolean keepOnCompletion)
1320-
throws IOException {
1338+
public static Map<String, Object> runEsqlAsync(
1339+
RequestObjectBuilder requestObject,
1340+
boolean keepOnCompletion,
1341+
AssertWarnings assertWarnings,
1342+
@Nullable ProfileLogger profileLogger
1343+
) throws IOException {
1344+
Boolean profileEnabled = requestObject.profile;
1345+
if (profileLogger != null) {
1346+
requestObject.profile(true);
1347+
}
13211348
addAsyncParameters(requestObject, keepOnCompletion);
13221349
Request request = prepareRequestWithOptions(requestObject, ASYNC);
13231350

@@ -1326,7 +1353,6 @@ public static EsqlResponse runEsqlAsyncNoWarningsChecks(RequestObjectBuilder req
13261353
}
13271354

13281355
Response response = performRequest(request);
1329-
Response initialResponse = response;
13301356
HttpEntity entity = response.getEntity();
13311357

13321358
Object initialColumns = null;
@@ -1346,15 +1372,23 @@ public static EsqlResponse runEsqlAsyncNoWarningsChecks(RequestObjectBuilder req
13461372
assertThat(response.getHeader("X-Elasticsearch-Async-Id"), nullValue());
13471373
assertThat(response.getHeader("X-Elasticsearch-Async-Is-Running"), is("?0"));
13481374
}
1375+
if (profileLogger != null) {
1376+
profileLogger.extractProfile(json, profileEnabled);
1377+
}
1378+
assertWarnings(response, assertWarnings);
13491379
json.remove("is_running"); // remove this to not mess up later map assertions
1350-
return new EsqlResponse(response, Collections.unmodifiableMap(json), initialResponse);
1380+
return Collections.unmodifiableMap(json);
13511381
} else {
13521382
// async may not return results immediately, so may need an async get
13531383
assertThat(id, is(not(emptyOrNullString())));
13541384
boolean isRunning = (boolean) json.get("is_running");
13551385
if (isRunning == false) {
13561386
// must have completed immediately so keep_on_completion must be true
13571387
assertThat(requestObject.keepOnCompletion(), is(true));
1388+
if (profileLogger != null) {
1389+
profileLogger.extractProfile(json, profileEnabled);
1390+
}
1391+
assertWarnings(response, assertWarnings);
13581392
// we already have the results, but let's remember them so that we can compare to async get
13591393
initialColumns = json.get("columns");
13601394
initialValues = json.get("values");
@@ -1391,8 +1425,12 @@ public static EsqlResponse runEsqlAsyncNoWarningsChecks(RequestObjectBuilder req
13911425
assertEquals(initialValues, result.get("values"));
13921426
}
13931427

1428+
if (profileLogger != null) {
1429+
profileLogger.extractProfile(result, profileEnabled);
1430+
}
1431+
assertWarnings(response, assertWarnings);
13941432
assertDeletable(id);
1395-
return new EsqlResponse(response, removeAsyncProperties(result), initialResponse);
1433+
return removeAsyncProperties(result);
13961434
}
13971435

13981436
private static Object removeOriginalTypesAndSuggestedCast(Object response) {
@@ -1624,8 +1662,8 @@ static String runEsqlAsTextWithFormat(RequestObjectBuilder builder, String forma
16241662

16251663
Response response = performRequest(request);
16261664
assertWarnings(response, new AssertWarnings.NoWarnings());
1627-
16281665
HttpEntity entity = response.getEntity();
1666+
16291667
// get the content, it could be empty because the request might have not completed
16301668
String initialValue = Streams.copyToString(new InputStreamReader(entity.getContent(), StandardCharsets.UTF_8));
16311669
String id = response.getHeader("X-Elasticsearch-Async-Id");
@@ -1727,7 +1765,13 @@ protected static Response performRequest(Request request) throws IOException {
17271765
return response;
17281766
}
17291767

1730-
public static void assertWarnings(Response response, AssertWarnings assertWarnings) {
1768+
static void assertNotPartial(Map<String, Object> answer) {
1769+
var clusters = answer.get("_clusters");
1770+
var reason = "unexpected partial results" + (clusters != null ? ": _clusters=" + clusters : "");
1771+
assertThat(reason, answer.get("is_partial"), anyOf(nullValue(), is(false)));
1772+
}
1773+
1774+
private static void assertWarnings(Response response, AssertWarnings assertWarnings) {
17311775
List<String> warnings = new ArrayList<>(response.getWarnings());
17321776
warnings.removeAll(mutedWarnings());
17331777
if (shouldLog()) {

‎x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/GenerativeRestTest.java‎

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -12,10 +12,12 @@
1212
import org.elasticsearch.test.rest.ESRestTestCase;
1313
import org.elasticsearch.xpack.esql.AssertWarnings;
1414
import org.elasticsearch.xpack.esql.CsvTestsDataLoader;
15+
import org.elasticsearch.xpack.esql.qa.rest.ProfileLogger;
1516
import org.elasticsearch.xpack.esql.qa.rest.RestEsqlTestCase;
1617
import org.elasticsearch.xpack.esql.qa.rest.generative.command.CommandGenerator;
1718
import org.junit.AfterClass;
1819
import org.junit.Before;
20+
import org.junit.Rule;
1921

2022
import java.io.IOException;
2123
import java.util.ArrayList;
@@ -32,6 +34,9 @@
3234

3335
public abstract class GenerativeRestTest extends ESRestTestCase {
3436

37+
@Rule(order = Integer.MIN_VALUE)
38+
public ProfileLogger profileLogger = new ProfileLogger();
39+
3540
public static final int ITERATIONS = 100;
3641
public static final int MAX_DEPTH = 20;
3742

@@ -165,13 +170,14 @@ private void checkException(EsqlQueryGenerator.QueryExecuted query) {
165170
@SuppressWarnings("unchecked")
166171
private EsqlQueryGenerator.QueryExecuted execute(String command, int depth) {
167172
try {
168-
Map<String, Object> a = RestEsqlTestCase.runEsql(
173+
Map<String, Object> json = RestEsqlTestCase.runEsql(
169174
new RestEsqlTestCase.RequestObjectBuilder().query(command).build(),
170175
new AssertWarnings.AllowedRegexes(List.of(Pattern.compile(".*"))),// we don't care about warnings
176+
profileLogger,
171177
RestEsqlTestCase.Mode.SYNC
172178
);
173-
List<EsqlQueryGenerator.Column> outputSchema = outputSchema(a);
174-
List<List<Object>> values = (List<List<Object>>) a.get("values");
179+
List<EsqlQueryGenerator.Column> outputSchema = outputSchema(json);
180+
List<List<Object>> values = (List<List<Object>>) json.get("values");
175181
return new EsqlQueryGenerator.QueryExecuted(command, depth, outputSchema, values, null);
176182
} catch (Exception e) {
177183
return new EsqlQueryGenerator.QueryExecuted(command, depth, null, null, e);

0 commit comments

Comments
 (0)