Skip to content

Commit 3c0af7b

Browse files
piotrrzyskopiotrrzysko
committed
Ensure queries returned via REST API are redacted
@JsonConstructor for TrimmedBasicQueryInfo was introduced to facilitate the deserialization of server responses in tests.
1 parent 0fa189b commit 3c0af7b

File tree

5 files changed

+606
-7
lines changed

5 files changed

+606
-7
lines changed

‎core/trino-main/src/main/java/io/trino/server/ui/TrimmedBasicQueryInfo.java‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
*/
1414
package io.trino.server.ui;
1515

16+
import com.fasterxml.jackson.annotation.JsonCreator;
1617
import com.fasterxml.jackson.annotation.JsonProperty;
1718
import com.google.errorprone.annotations.Immutable;
1819
import io.trino.execution.QueryState;
@@ -56,6 +57,47 @@ public class TrimmedBasicQueryInfo
5657
private final RetryPolicy retryPolicy;
5758
private final Optional<Set<String>> clientTags;
5859

60+
@JsonCreator
61+
public TrimmedBasicQueryInfo(
62+
@JsonProperty("queryId") QueryId queryId,
63+
@JsonProperty("sessionUser") String sessionUser,
64+
@JsonProperty("sessionPrincipal") Optional<String> sessionPrincipal,
65+
@JsonProperty("sessionSource") Optional<String> sessionSource,
66+
@JsonProperty("resourceGroupId") Optional<ResourceGroupId> resourceGroupId,
67+
@JsonProperty("queryDataEncoding") Optional<String> queryDataEncoding,
68+
@JsonProperty("state") QueryState state,
69+
@JsonProperty("scheduled") boolean scheduled,
70+
@JsonProperty("self") URI self,
71+
@JsonProperty("queryTextPreview") String queryTextPreview,
72+
@JsonProperty("updateType") Optional<String> updateType,
73+
@JsonProperty("preparedQuery") Optional<String> preparedQuery,
74+
@JsonProperty("queryStats") BasicQueryStats queryStats,
75+
@JsonProperty("errorType") Optional<ErrorType> errorType,
76+
@JsonProperty("errorCode") Optional<ErrorCode> errorCode,
77+
@JsonProperty("queryType") Optional<QueryType> queryType,
78+
@JsonProperty("retryPolicy") RetryPolicy retryPolicy,
79+
@JsonProperty("clientTags") Optional<Set<String>> clientTags)
80+
{
81+
this.queryId = requireNonNull(queryId, "queryId is null");
82+
this.sessionUser = requireNonNull(sessionUser, "sessionUser is null");
83+
this.sessionPrincipal = requireNonNull(sessionPrincipal, "sessionPrincipal is null");
84+
this.sessionSource = requireNonNull(sessionSource, "sessionSource is null");
85+
this.resourceGroupId = requireNonNull(resourceGroupId, "resourceGroupId is null");
86+
this.queryDataEncoding = requireNonNull(queryDataEncoding, "queryDataEncoding is null");
87+
this.state = requireNonNull(state, "state is null");
88+
this.scheduled = scheduled;
89+
this.self = requireNonNull(self, "self is null");
90+
this.queryTextPreview = requireNonNull(queryTextPreview, "queryTextPreview is null");
91+
this.updateType = requireNonNull(updateType, "updateType is null");
92+
this.preparedQuery = requireNonNull(preparedQuery, "preparedQuery is null");
93+
this.queryStats = requireNonNull(queryStats, "queryStats is null");
94+
this.errorType = requireNonNull(errorType, "errorType is null");
95+
this.errorCode = requireNonNull(errorCode, "errorCode is null");
96+
this.queryType = requireNonNull(queryType, "queryType is null");
97+
this.retryPolicy = requireNonNull(retryPolicy, "retryPolicy is null");
98+
this.clientTags = requireNonNull(clientTags, "clientTags is null");
99+
}
100+
59101
public TrimmedBasicQueryInfo(BasicQueryInfo queryInfo)
60102
{
61103
this.queryId = requireNonNull(queryInfo.getQueryId(), "queryId is null");

‎core/trino-main/src/test/java/io/trino/server/TestQueryResource.java‎

Lines changed: 130 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
*/
1414
package io.trino.server;
1515

16+
import com.google.common.collect.ImmutableSet;
1617
import com.google.inject.Key;
1718
import io.airlift.http.client.HttpClient;
1819
import io.airlift.http.client.HttpUriBuilder;
@@ -29,6 +30,8 @@
2930
import io.trino.client.QueryDataClientJacksonModule;
3031
import io.trino.client.QueryResults;
3132
import io.trino.client.ResultRowsDecoder;
33+
import io.trino.connector.MockConnectorFactory;
34+
import io.trino.connector.MockConnectorPlugin;
3235
import io.trino.execution.QueryInfo;
3336
import io.trino.plugin.tpch.TpchPlugin;
3437
import io.trino.server.testing.TestingTrinoServer;
@@ -39,8 +42,10 @@
3942
import org.junit.jupiter.api.TestInstance;
4043

4144
import java.net.URI;
45+
import java.net.URLEncoder;
4246
import java.util.List;
4347
import java.util.Map;
48+
import java.util.Optional;
4449
import java.util.Set;
4550

4651
import static io.airlift.http.client.HttpUriBuilder.uriBuilderFrom;
@@ -65,6 +70,7 @@
6570
import static io.trino.testing.TestingAccessControlManager.TestingPrivilegeType.KILL_QUERY;
6671
import static io.trino.testing.TestingAccessControlManager.TestingPrivilegeType.VIEW_QUERY;
6772
import static io.trino.testing.TestingAccessControlManager.privilege;
73+
import static io.trino.testing.TestingNames.randomNameSuffix;
6874
import static java.nio.charset.StandardCharsets.UTF_8;
6975
import static org.assertj.core.api.Assertions.assertThat;
7076
import static org.assertj.core.api.Assertions.assertThatThrownBy;
@@ -94,6 +100,9 @@ public void setup()
94100
{
95101
client = new JettyHttpClient();
96102
server = TestingTrinoServer.create();
103+
server.installPlugin(new MockConnectorPlugin(MockConnectorFactory.builder()
104+
.withSecuritySensitivePropertyNames(ImmutableSet.of("password"))
105+
.build()));
97106
server.installPlugin(new TpchPlugin());
98107
server.createCatalog("tpch", "tpch");
99108
}
@@ -226,6 +235,114 @@ public void testGetQueryInfoExecutionFailure()
226235
assertThat(info.getFailureInfo().getErrorCode()).isEqualTo(DIVISION_BY_ZERO.toErrorCode());
227236
}
228237

238+
@Test
239+
public void testGetQueryInfosWithRedactedSecrets()
240+
{
241+
String catalog = "catalog_" + randomNameSuffix();
242+
runToCompletion(
243+
"""
244+
CREATE CATALOG %s USING mock
245+
WITH (
246+
"user" = 'bob',
247+
"password" = '1234'
248+
)
249+
""".formatted(catalog));
250+
251+
List<BasicQueryInfo> infos = getQueryInfos("/v1/query");
252+
assertThat(infos.size()).isEqualTo(1);
253+
assertThat(infos.getFirst().getQuery()).isEqualTo(
254+
"""
255+
CREATE CATALOG %s USING mock
256+
WITH (
257+
"user" = 'bob',
258+
"password" = '***'
259+
)\
260+
""".formatted(catalog));
261+
}
262+
263+
@Test
264+
public void testGetQueryInfosWithRedactedSecretsInPreparedStatement()
265+
{
266+
String catalog = "catalog_" + randomNameSuffix();
267+
runToCompletion(
268+
"EXECUTE create_catalog",
269+
Optional.of("create_catalog"),
270+
Optional.of(
271+
"""
272+
CREATE CATALOG %s USING mock
273+
WITH (
274+
"user" = 'bob',
275+
"password" = '1234'
276+
)
277+
""".formatted(catalog)));
278+
279+
List<BasicQueryInfo> infos = getQueryInfos("/v1/query");
280+
assertThat(infos.size()).isEqualTo(1);
281+
assertThat(infos.getFirst().getQuery()).isEqualTo("EXECUTE create_catalog");
282+
assertThat(infos.getFirst().getPreparedQuery()).isEqualTo(
283+
Optional.of(
284+
"""
285+
CREATE CATALOG %s USING mock
286+
WITH (
287+
"user" = 'bob',
288+
"password" = '***'
289+
)\
290+
""".formatted(catalog)));
291+
}
292+
293+
@Test
294+
public void testGetQueryInfoWithRedactedSecrets()
295+
{
296+
String catalog = "catalog_" + randomNameSuffix();
297+
String queryId = runToCompletion(
298+
"""
299+
CREATE CATALOG %s USING mock
300+
WITH (
301+
"user" = 'bob',
302+
"password" = '1234'
303+
)
304+
""".formatted(catalog));
305+
306+
QueryInfo queryInfo = getQueryInfo(queryId);
307+
assertThat(queryInfo.getQuery()).isEqualTo(
308+
"""
309+
CREATE CATALOG %s USING mock
310+
WITH (
311+
"user" = 'bob',
312+
"password" = '***'
313+
)\
314+
""".formatted(catalog));
315+
}
316+
317+
@Test
318+
public void testGetQueryInfoWithRedactedSecretsInPreparedStatement()
319+
{
320+
String catalog = "catalog_" + randomNameSuffix();
321+
String queryId = runToCompletion(
322+
"EXECUTE create_catalog",
323+
Optional.of("create_catalog"),
324+
Optional.of(
325+
"""
326+
CREATE CATALOG %s USING mock
327+
WITH (
328+
"user" = 'bob',
329+
"password" = '1234'
330+
)
331+
""".formatted(catalog)));
332+
333+
QueryInfo queryInfo = getQueryInfo(queryId);
334+
assertThat(queryInfo.getQuery()).isEqualTo("EXECUTE create_catalog");
335+
assertThat(queryInfo.getPreparedQuery()).isEqualTo(
336+
Optional.of(
337+
"""
338+
CREATE CATALOG %s USING mock
339+
WITH (
340+
"user" = 'bob',
341+
"password" = '***'
342+
)\
343+
""".formatted(catalog)));
344+
}
345+
229346
@Test
230347
public void testCancel()
231348
{
@@ -299,13 +416,23 @@ private void testKilled(String killType)
299416
}
300417

301418
private String runToCompletion(String sql)
419+
{
420+
return runToCompletion(sql, Optional.empty(), Optional.empty());
421+
}
422+
423+
private String runToCompletion(String sql, Optional<String> preparedStatementName, Optional<String> preparedStatement)
302424
{
303425
URI uri = uriBuilderFrom(server.getBaseUrl().resolve("/v1/statement")).build();
304-
Request request = preparePost()
426+
Request.Builder requestBuilder = preparePost()
305427
.setHeader(TRINO_HEADERS.requestUser(), "user")
306428
.setUri(uri)
307-
.setBodyGenerator(createStaticBodyGenerator(sql, UTF_8))
308-
.build();
429+
.setBodyGenerator(createStaticBodyGenerator(sql, UTF_8));
430+
431+
if (preparedStatementName.isPresent() && preparedStatement.isPresent()) {
432+
requestBuilder.setHeader(TRINO_HEADERS.requestPreparedStatement(), preparedStatementName.get() + "=" + URLEncoder.encode(preparedStatement.get(), UTF_8));
433+
}
434+
435+
Request request = requestBuilder.build();
309436
QueryResults queryResults = client.execute(request, createJsonResponseHandler(QUERY_RESULTS_JSON_CODEC));
310437
while (queryResults.getNextUri() != null) {
311438
request = prepareGet()

‎core/trino-main/src/test/java/io/trino/server/TestQueryStateInfoResource.java‎

Lines changed: 54 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
*/
1414
package io.trino.server;
1515

16+
import com.google.common.collect.ImmutableSet;
1617
import com.google.common.io.Closer;
1718
import io.airlift.http.client.HttpClient;
1819
import io.airlift.http.client.Request;
@@ -23,6 +24,8 @@
2324
import io.airlift.json.ObjectMapperProvider;
2425
import io.airlift.units.Duration;
2526
import io.trino.client.QueryResults;
27+
import io.trino.connector.MockConnectorFactory;
28+
import io.trino.connector.MockConnectorPlugin;
2629
import io.trino.plugin.tpch.TpchPlugin;
2730
import io.trino.server.protocol.spooling.QueryDataJacksonModule;
2831
import io.trino.server.testing.TestingTrinoServer;
@@ -47,6 +50,7 @@
4750
import static io.airlift.json.JsonCodec.listJsonCodec;
4851
import static io.trino.client.ProtocolHeaders.TRINO_HEADERS;
4952
import static io.trino.execution.QueryState.FAILED;
53+
import static io.trino.execution.QueryState.FINISHING;
5054
import static io.trino.execution.QueryState.RUNNING;
5155
import static io.trino.server.TestQueryResource.BASIC_QUERY_INFO_CODEC;
5256
import static io.trino.testing.TestingAccessControlManager.TestingPrivilegeType.VIEW_QUERY;
@@ -71,11 +75,15 @@ public class TestQueryStateInfoResource
7175
private TestingTrinoServer server;
7276
private HttpClient client;
7377
private QueryResults queryResults;
78+
private QueryResults createCatalogResults;
7479

7580
@BeforeAll
7681
public void setUp()
7782
{
7883
server = TestingTrinoServer.create();
84+
server.installPlugin(new MockConnectorPlugin(MockConnectorFactory.builder()
85+
.withSecuritySensitivePropertyNames(ImmutableSet.of("password"))
86+
.build()));
7987
server.installPlugin(new TpchPlugin());
8088
server.createCatalog("tpch", "tpch");
8189
client = new JettyHttpClient();
@@ -96,6 +104,19 @@ public void setUp()
96104
QueryResults queryResults2 = client.execute(request2, createJsonResponseHandler(QUERY_RESULTS_JSON_CODEC));
97105
client.execute(prepareGet().setUri(queryResults2.getNextUri()).build(), createJsonResponseHandler(QUERY_RESULTS_JSON_CODEC));
98106

107+
Request createCatalogRequest = preparePost()
108+
.setUri(uriBuilderFrom(server.getBaseUrl()).replacePath("/v1/statement").build())
109+
.setBodyGenerator(createStaticBodyGenerator("""
110+
CREATE CATALOG test_catalog USING mock
111+
WITH (
112+
"user" = 'bob',
113+
"password" = '1234'
114+
)""", UTF_8))
115+
.setHeader(TRINO_HEADERS.requestUser(), "catalogCreator")
116+
.build();
117+
createCatalogResults = client.execute(createCatalogRequest, createJsonResponseHandler(jsonCodec(QueryResults.class)));
118+
client.execute(prepareGet().setUri(createCatalogResults.getNextUri()).build(), createJsonResponseHandler(QUERY_RESULTS_JSON_CODEC));
119+
99120
// queries are started in the background, so they may not all be immediately visible
100121
long start = System.nanoTime();
101122
while (Duration.nanosSince(start).compareTo(new Duration(5, MINUTES)) < 0) {
@@ -105,8 +126,8 @@ public void setUp()
105126
.setHeader(TRINO_HEADERS.requestUser(), "unknown")
106127
.build(),
107128
createJsonResponseHandler(BASIC_QUERY_INFO_CODEC));
108-
if (queryInfos.size() == 2) {
109-
if (queryInfos.stream().allMatch(info -> info.getState() == RUNNING)) {
129+
if (queryInfos.size() == 3) {
130+
if (queryInfos.stream().allMatch(info -> info.getState() == RUNNING || info.getState() == FINISHING)) {
110131
break;
111132
}
112133

@@ -143,7 +164,12 @@ public void testGetAllQueryStateInfos()
143164
.build(),
144165
createJsonResponseHandler(listJsonCodec(QueryStateInfo.class)));
145166

146-
assertThat(infos).hasSize(2);
167+
assertThat(infos.size()).isEqualTo(3);
168+
QueryStateInfo createCatalogInfo = infos.stream()
169+
.filter(info -> info.getQueryId().getId().equals(createCatalogResults.getId()))
170+
.findFirst()
171+
.orElse(null);
172+
assertCreateCatalogQueryIsRedacted(createCatalogInfo);
147173
}
148174

149175
@Test
@@ -185,6 +211,19 @@ public void testGetQueryStateInfo()
185211
assertThat(info).isNotNull();
186212
}
187213

214+
@Test
215+
public void testGetQueryStateInfoWithRedactedSecrets()
216+
{
217+
QueryStateInfo info = client.execute(
218+
prepareGet()
219+
.setUri(server.resolve("/v1/queryState/" + createCatalogResults.getId()))
220+
.setHeader(TRINO_HEADERS.requestUser(), "unknown")
221+
.build(),
222+
createJsonResponseHandler(jsonCodec(QueryStateInfo.class)));
223+
224+
assertCreateCatalogQueryIsRedacted(info);
225+
}
226+
188227
@Test
189228
public void testGetAllQueryStateInfosDenied()
190229
{
@@ -194,7 +233,7 @@ public void testGetAllQueryStateInfosDenied()
194233
.setHeader(TRINO_HEADERS.requestUser(), "any-other-user")
195234
.build(),
196235
createJsonResponseHandler(listJsonCodec(QueryStateInfo.class)));
197-
assertThat(infos).hasSize(2);
236+
assertThat(infos).hasSize(3);
198237

199238
testGetAllQueryStateInfosDenied("user1", 1);
200239
testGetAllQueryStateInfosDenied("any-other-user", 0);
@@ -249,4 +288,15 @@ public void testGetQueryStateInfoNo()
249288
.isInstanceOf(UnexpectedResponseException.class)
250289
.hasMessageMatching("Expected response code .*, but was 404");
251290
}
291+
292+
private static void assertCreateCatalogQueryIsRedacted(QueryStateInfo info)
293+
{
294+
assertThat(info).isNotNull();
295+
assertThat(info.getQuery()).isEqualTo("""
296+
CREATE CATALOG test_catalog USING mock
297+
WITH (
298+
"user" = 'bob',
299+
"password" = '***'
300+
)""");
301+
}
252302
}

0 commit comments

Comments
 (0)