Skip to content

Commit 3209c4f

Browse files
ekuvardinekuvardin
committed
Add tests connecting to resource pool
1 parent 5e090f4 commit 3209c4f

File tree

3 files changed

+192
-2
lines changed

3 files changed

+192
-2
lines changed

‎query/src/main/java/tech/ydb/query/impl/SessionImpl.java‎

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -182,6 +182,16 @@ private static YdbQuery.StatsMode mapStatsMode(QueryStatsMode mode) {
182182
}
183183
}
184184

185+
private String mapPoolId(ExecuteQuerySettings settings) {
186+
String actualPoolId = settings.getResourcePool();
187+
188+
if (actualPoolId == null) {
189+
return YdbQuery.ExecuteQueryRequest.getDefaultInstance().getPoolId();
190+
}
191+
192+
return actualPoolId;
193+
}
194+
185195
GrpcReadStream<YdbQuery.ExecuteQueryResponsePart> createGrpcStream(
186196
String query, YdbQuery.TransactionControl tx, Params prms, ExecuteQuerySettings settings
187197
) {
@@ -194,7 +204,7 @@ GrpcReadStream<YdbQuery.ExecuteQueryResponsePart> createGrpcStream(
194204
.setText(query)
195205
.build()
196206
)
197-
.setPoolId(settings.getResourcePool())
207+
.setPoolId(mapPoolId(settings))
198208
.putAllParameters(prms.toPb());
199209

200210
if (tx != null) {

‎query/src/main/java/tech/ydb/query/settings/ExecuteQuerySettings.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,7 @@ public static Builder newBuilder() {
3737
public static class Builder extends BaseBuilder<Builder> {
3838
private QueryExecMode execMode = QueryExecMode.EXECUTE;
3939
private QueryStatsMode statsMode = QueryStatsMode.NONE;
40-
private String resourcePool = "default";
40+
private String resourcePool = null;
4141

4242
public Builder withExecMode(QueryExecMode mode) {
4343
this.execMode = mode;

‎query/src/test/java/tech/ydb/query/impl/QueryIntegrationResourcePoolTest.java‎

Lines changed: 180 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,180 @@
1+
package tech.ydb.query.impl;
2+
3+
import org.junit.*;
4+
import org.slf4j.Logger;
5+
import org.slf4j.LoggerFactory;
6+
7+
import tech.ydb.common.transaction.TxMode;
8+
import tech.ydb.core.Issue;
9+
import tech.ydb.core.Result;
10+
import tech.ydb.core.Status;
11+
import tech.ydb.core.StatusCode;
12+
import tech.ydb.query.QueryClient;
13+
import tech.ydb.query.QuerySession;
14+
import tech.ydb.query.QueryStream;
15+
import tech.ydb.query.QueryTransaction;
16+
import tech.ydb.query.result.QueryInfo;
17+
import tech.ydb.query.result.QueryResultPart;
18+
import tech.ydb.query.settings.ExecuteQuerySettings;
19+
import tech.ydb.query.settings.QueryExecMode;
20+
import tech.ydb.query.settings.QueryStatsMode;
21+
import tech.ydb.query.tools.QueryReader;
22+
import tech.ydb.table.SessionRetryContext;
23+
import tech.ydb.table.description.TableDescription;
24+
import tech.ydb.table.impl.SimpleTableClient;
25+
import tech.ydb.table.query.Params;
26+
import tech.ydb.table.result.ResultSetReader;
27+
import tech.ydb.table.rpc.grpc.GrpcTableRpc;
28+
import tech.ydb.table.values.PrimitiveType;
29+
import tech.ydb.table.values.PrimitiveValue;
30+
import tech.ydb.test.junit4.GrpcTransportRule;
31+
32+
import java.time.Duration;
33+
import java.util.ArrayList;
34+
import java.util.Iterator;
35+
import java.util.List;
36+
import java.util.concurrent.CompletableFuture;
37+
38+
/**
39+
* Test on resource poll.
40+
* Take an account, that resource poll with name "default" exists every time and can't be deleted
41+
* Also when we specify pool with empty string "" it's equivalent to default pool
42+
*
43+
* @author Evgeny Kuvardin
44+
*/
45+
public class QueryIntegrationResourcePoolTest {
46+
private final static Logger logger = LoggerFactory.getLogger(QueryIntegrationResourcePoolTest.class);
47+
private final static String TEST_TABLE = "query_service_test";
48+
private final static String TEST_DOUBLE_TABLE = "query_double_table";
49+
private final static String TEST_RESOURCE_POOL = "test_pool";
50+
51+
52+
@ClassRule
53+
public final static GrpcTransportRule ydbTransport = new GrpcTransportRule();
54+
55+
@BeforeClass
56+
public static void initSchema() {
57+
logger.info("Prepare database...");
58+
59+
String tablePath = ydbTransport.getDatabase() + "/" + TEST_TABLE;
60+
TableDescription tableDescription = TableDescription.newBuilder()
61+
.addNonnullColumn("id", PrimitiveType.Int32)
62+
.addNullableColumn("name", PrimitiveType.Text)
63+
.addNullableColumn("payload", PrimitiveType.Bytes)
64+
.addNullableColumn("is_valid", PrimitiveType.Bool)
65+
.setPrimaryKey("id")
66+
.build();
67+
68+
69+
String table2Path = ydbTransport.getDatabase() + "/" + TEST_DOUBLE_TABLE;
70+
TableDescription table2Description = TableDescription.newBuilder()
71+
.addNonnullColumn("id", PrimitiveType.Int32)
72+
.addNullableColumn("amount", PrimitiveType.Double)
73+
.setPrimaryKey("id")
74+
.build();
75+
76+
SimpleTableClient client = SimpleTableClient.newClient(GrpcTableRpc.useTransport(ydbTransport)).build();
77+
SessionRetryContext retryCtx = SessionRetryContext.create(client).build();
78+
retryCtx.supplyStatus(session -> session.createTable(tablePath, tableDescription)).join();
79+
retryCtx.supplyStatus(session -> session.createTable(table2Path, table2Description)).join();
80+
logger.info("Prepare database OK");
81+
82+
try (QueryClient queryClient = QueryClient.newClient(ydbTransport).build()) {
83+
try (QuerySession querySession = queryClient.createSession(Duration.ofSeconds(5)).join().getValue()) {
84+
Result<QueryInfo> result = querySession.createQuery("CREATE RESOURCE POOL " + TEST_RESOURCE_POOL + " WITH (\n" +
85+
" CONCURRENT_QUERY_LIMIT=10,\n" +
86+
" QUEUE_SIZE=1000,\n" +
87+
" DATABASE_LOAD_CPU_THRESHOLD=80,\n" +
88+
" TOTAL_CPU_LIMIT_PERCENT_PER_NODE=70);", TxMode.NONE).execute().join();
89+
90+
Assert.assertTrue(result.getValue().toString(), result.isSuccess());
91+
}
92+
}
93+
}
94+
95+
@AfterClass
96+
public static void dropAll() {
97+
logger.info("Clean database...");
98+
String tablePath = ydbTransport.getDatabase() + "/" + TEST_TABLE;
99+
String table2Path = ydbTransport.getDatabase() + "/" + TEST_DOUBLE_TABLE;
100+
101+
SimpleTableClient client = SimpleTableClient.newClient(GrpcTableRpc.useTransport(ydbTransport)).build();
102+
SessionRetryContext retryCtx = SessionRetryContext.create(client).build();
103+
retryCtx.supplyStatus(session -> session.dropTable(tablePath)).join().isSuccess();
104+
retryCtx.supplyStatus(session -> session.dropTable(table2Path)).join().isSuccess();
105+
106+
try (QueryClient queryClient = QueryClient.newClient(ydbTransport).build()) {
107+
try (QuerySession querySession = queryClient.createSession(Duration.ofSeconds(5)).join().getValue()) {
108+
Result<QueryInfo> result = querySession.createQuery("DROP RESOURCE POOL " + TEST_RESOURCE_POOL + ";", TxMode.NONE).execute().join();
109+
110+
Assert.assertTrue(result.getValue().toString(), result.isSuccess());
111+
}
112+
}
113+
114+
logger.info("Clean database OK");
115+
}
116+
117+
@Test
118+
public void selectWithResourcePoolTest() {
119+
try (QueryClient client = QueryClient.newClient(ydbTransport).build()) {
120+
try (QuerySession session = client.createSession(Duration.ofSeconds(5)).join().getValue()) {
121+
ExecuteQuerySettings settings = ExecuteQuerySettings.newBuilder()
122+
.withExecMode(QueryExecMode.EXECUTE)
123+
.withResourcePool("test_pool")
124+
.build();
125+
126+
Assert.assertTrue("Query shouldn't fall",
127+
session.createQuery("SELECT 2 + 3;", TxMode.SERIALIZABLE_RW, Params.empty(), settings).execute()
128+
.join().isSuccess());
129+
}
130+
}
131+
}
132+
133+
@Test
134+
public void selectWithDefaultResourcePoolTest() {
135+
try (QueryClient client = QueryClient.newClient(ydbTransport).build()) {
136+
try (QuerySession session = client.createSession(Duration.ofSeconds(5)).join().getValue()) {
137+
ExecuteQuerySettings settings = ExecuteQuerySettings.newBuilder()
138+
.withExecMode(QueryExecMode.EXECUTE)
139+
.withResourcePool("default")
140+
.build();
141+
142+
Assert.assertTrue("Query shouldn't fall",
143+
session.createQuery("SELECT 2 + 3;", TxMode.SERIALIZABLE_RW, Params.empty(), settings).execute()
144+
.join().isSuccess());
145+
}
146+
}
147+
}
148+
149+
@Test
150+
public void selectWithDefaultResourcePoolAndEmptyStringTest() {
151+
try (QueryClient client = QueryClient.newClient(ydbTransport).build()) {
152+
try (QuerySession session = client.createSession(Duration.ofSeconds(5)).join().getValue()) {
153+
ExecuteQuerySettings settings = ExecuteQuerySettings.newBuilder()
154+
.withExecMode(QueryExecMode.EXECUTE)
155+
.withResourcePool("")
156+
.build();
157+
158+
Assert.assertTrue("Query shouldn't fall",
159+
session.createQuery("SELECT 2 + 3;", TxMode.SERIALIZABLE_RW, Params.empty(), settings).execute()
160+
.join().isSuccess());
161+
}
162+
}
163+
}
164+
165+
@Test
166+
public void selectShouldFailWithUnknownResourcePollTest() {
167+
try (QueryClient client = QueryClient.newClient(ydbTransport).build()) {
168+
try (QuerySession session = client.createSession(Duration.ofSeconds(5)).join().getValue()) {
169+
ExecuteQuerySettings settings = ExecuteQuerySettings.newBuilder()
170+
.withExecMode(QueryExecMode.EXECUTE)
171+
.withResourcePool("some_unknown_pool")
172+
.build();
173+
174+
Assert.assertFalse("Query should fall",
175+
session.createQuery("SELECT 2 + 3;", TxMode.SERIALIZABLE_RW, Params.empty(), settings).execute()
176+
.join().isSuccess());
177+
}
178+
}
179+
}
180+
}

0 commit comments

Comments
 (0)