13
13
*/
14
14
package io .trino .server ;
15
15
16
+ import com .google .common .collect .ImmutableSet ;
16
17
import com .google .common .io .Closer ;
17
18
import io .airlift .http .client .HttpClient ;
18
19
import io .airlift .http .client .Request ;
23
24
import io .airlift .json .ObjectMapperProvider ;
24
25
import io .airlift .units .Duration ;
25
26
import io .trino .client .QueryResults ;
27
+ import io .trino .connector .MockConnectorFactory ;
28
+ import io .trino .connector .MockConnectorPlugin ;
26
29
import io .trino .plugin .tpch .TpchPlugin ;
27
30
import io .trino .server .protocol .spooling .QueryDataJacksonModule ;
28
31
import io .trino .server .testing .TestingTrinoServer ;
47
50
import static io .airlift .json .JsonCodec .listJsonCodec ;
48
51
import static io .trino .client .ProtocolHeaders .TRINO_HEADERS ;
49
52
import static io .trino .execution .QueryState .FAILED ;
53
+ import static io .trino .execution .QueryState .FINISHING ;
50
54
import static io .trino .execution .QueryState .RUNNING ;
51
55
import static io .trino .server .TestQueryResource .BASIC_QUERY_INFO_CODEC ;
52
56
import static io .trino .testing .TestingAccessControlManager .TestingPrivilegeType .VIEW_QUERY ;
@@ -71,11 +75,15 @@ public class TestQueryStateInfoResource
71
75
private TestingTrinoServer server ;
72
76
private HttpClient client ;
73
77
private QueryResults queryResults ;
78
+ private QueryResults createCatalogResults ;
74
79
75
80
@ BeforeAll
76
81
public void setUp ()
77
82
{
78
83
server = TestingTrinoServer .create ();
84
+ server .installPlugin (new MockConnectorPlugin (MockConnectorFactory .builder ()
85
+ .withSecuritySensitivePropertyNames (ImmutableSet .of ("password" ))
86
+ .build ()));
79
87
server .installPlugin (new TpchPlugin ());
80
88
server .createCatalog ("tpch" , "tpch" );
81
89
client = new JettyHttpClient ();
@@ -96,6 +104,19 @@ public void setUp()
96
104
QueryResults queryResults2 = client .execute (request2 , createJsonResponseHandler (QUERY_RESULTS_JSON_CODEC ));
97
105
client .execute (prepareGet ().setUri (queryResults2 .getNextUri ()).build (), createJsonResponseHandler (QUERY_RESULTS_JSON_CODEC ));
98
106
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
+
99
120
// queries are started in the background, so they may not all be immediately visible
100
121
long start = System .nanoTime ();
101
122
while (Duration .nanosSince (start ).compareTo (new Duration (5 , MINUTES )) < 0 ) {
@@ -105,8 +126,8 @@ public void setUp()
105
126
.setHeader (TRINO_HEADERS .requestUser (), "unknown" )
106
127
.build (),
107
128
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 )) {
110
131
break ;
111
132
}
112
133
@@ -143,7 +164,12 @@ public void testGetAllQueryStateInfos()
143
164
.build (),
144
165
createJsonResponseHandler (listJsonCodec (QueryStateInfo .class )));
145
166
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 );
147
173
}
148
174
149
175
@ Test
@@ -185,6 +211,19 @@ public void testGetQueryStateInfo()
185
211
assertThat (info ).isNotNull ();
186
212
}
187
213
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
+
188
227
@ Test
189
228
public void testGetAllQueryStateInfosDenied ()
190
229
{
@@ -194,7 +233,7 @@ public void testGetAllQueryStateInfosDenied()
194
233
.setHeader (TRINO_HEADERS .requestUser (), "any-other-user" )
195
234
.build (),
196
235
createJsonResponseHandler (listJsonCodec (QueryStateInfo .class )));
197
- assertThat (infos ).hasSize (2 );
236
+ assertThat (infos ).hasSize (3 );
198
237
199
238
testGetAllQueryStateInfosDenied ("user1" , 1 );
200
239
testGetAllQueryStateInfosDenied ("any-other-user" , 0 );
@@ -249,4 +288,15 @@ public void testGetQueryStateInfoNo()
249
288
.isInstanceOf (UnexpectedResponseException .class )
250
289
.hasMessageMatching ("Expected response code .*, but was 404" );
251
290
}
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
+ }
252
302
}
0 commit comments