Skip to content

Commit 654d3e2

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 f28846e commit 654d3e2

File tree

5 files changed

+441
-4
lines changed

5 files changed

+441
-4
lines changed

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

+40
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;
@@ -54,6 +55,45 @@ public class TrimmedBasicQueryInfo
5455
private final Optional<QueryType> queryType;
5556
private final RetryPolicy retryPolicy;
5657

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

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

+48
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;
@@ -65,6 +68,7 @@
6568
import static io.trino.testing.TestingAccessControlManager.TestingPrivilegeType.KILL_QUERY;
6669
import static io.trino.testing.TestingAccessControlManager.TestingPrivilegeType.VIEW_QUERY;
6770
import static io.trino.testing.TestingAccessControlManager.privilege;
71+
import static io.trino.testing.TestingNames.randomNameSuffix;
6872
import static java.nio.charset.StandardCharsets.UTF_8;
6973
import static org.assertj.core.api.Assertions.assertThat;
7074
import static org.assertj.core.api.Assertions.assertThatThrownBy;
@@ -94,6 +98,9 @@ public void setup()
9498
{
9599
client = new JettyHttpClient();
96100
server = TestingTrinoServer.create();
101+
server.installPlugin(new MockConnectorPlugin(MockConnectorFactory.builder()
102+
.withSecuritySensitivePropertyNames(ImmutableSet.of("password"))
103+
.build()));
97104
server.installPlugin(new TpchPlugin());
98105
server.createCatalog("tpch", "tpch");
99106
}
@@ -226,6 +233,47 @@ public void testGetQueryInfoExecutionFailure()
226233
assertThat(info.getFailureInfo().getErrorCode()).isEqualTo(DIVISION_BY_ZERO.toErrorCode());
227234
}
228235

236+
@Test
237+
public void testGetQueryInfosWithRedactedSecrets()
238+
{
239+
String catalog = "catalog_" + randomNameSuffix();
240+
runToCompletion("""
241+
CREATE CATALOG %s USING mock
242+
WITH (
243+
"user" = 'bob',
244+
"password" = '1234'
245+
)""".formatted(catalog));
246+
247+
List<BasicQueryInfo> infos = getQueryInfos("/v1/query");
248+
assertThat(infos.size()).isEqualTo(1);
249+
assertThat(infos.getFirst().getQuery()).isEqualTo("""
250+
CREATE CATALOG %s USING mock
251+
WITH (
252+
"user" = 'bob',
253+
"password" = '***'
254+
)""".formatted(catalog));
255+
}
256+
257+
@Test
258+
public void testGetQueryInfoWithRedactedSecrets()
259+
{
260+
String catalog = "catalog_" + randomNameSuffix();
261+
String queryId = runToCompletion("""
262+
CREATE CATALOG %s USING mock
263+
WITH (
264+
"user" = 'bob',
265+
"password" = '1234'
266+
)""".formatted(catalog));
267+
268+
QueryInfo queryInfo = getQueryInfo(queryId);
269+
assertThat(queryInfo.getQuery()).isEqualTo("""
270+
CREATE CATALOG %s USING mock
271+
WITH (
272+
"user" = 'bob',
273+
"password" = '***'
274+
)""".formatted(catalog));
275+
}
276+
229277
@Test
230278
public void testCancel()
231279
{

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

+54-4
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)