Skip to content

Commit bae6b46

Browse files
ekuvardinekuvardin
committed
Add tests and execute script
1 parent 3b6d6e6 commit bae6b46

File tree

5 files changed

+542
-0
lines changed

5 files changed

+542
-0
lines changed

‎query/src/main/java/tech/ydb/query/QuerySession.java‎

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,11 @@
66

77
import tech.ydb.common.transaction.TxMode;
88
import tech.ydb.core.Result;
9+
import tech.ydb.proto.OperationProtos;
910
import tech.ydb.query.settings.BeginTransactionSettings;
1011
import tech.ydb.query.settings.ExecuteQuerySettings;
12+
import tech.ydb.query.settings.ExecuteScriptSettings;
13+
import tech.ydb.query.settings.FetchScriptSettings;
1114
import tech.ydb.table.query.Params;
1215

1316
/**
@@ -68,6 +71,29 @@ public interface QuerySession extends AutoCloseable {
6871
*/
6972
QueryStream createQuery(String query, TxMode tx, Params params, ExecuteQuerySettings settings);
7073

74+
/**
75+
* Create {@link QueryStream} for executing query with specified {@link TxMode}. The query can contain DML, DDL and
76+
* DCL statements. Supported mix of different statement types depends on the chosen transaction type.
77+
*
78+
* @param query text of query
79+
* @param params query parameters
80+
* @param settings additional settings of query execution
81+
* @return a ready to execute instance of {@link QueryStream}
82+
*/
83+
CompletableFuture<Result<OperationProtos.Operation>> executeScript(String query, Params params, ExecuteScriptSettings settings);
84+
85+
/**
86+
* Create {@link QueryStream} for executing query with specified {@link TxMode}. The query can contain DML, DDL and
87+
* DCL statements. Supported mix of different statement types depends on the chosen transaction type.
88+
*
89+
* @param query text of query
90+
* @param tx transaction mode
91+
* @param params query parameters
92+
* @param settings additional settings of query execution
93+
* @return a ready to execute instance of {@link QueryStream}
94+
*/
95+
QueryStream fetchScriptResults(String query, TxMode tx, Params params, FetchScriptSettings settings);
96+
7197
@Override
7298
void close();
7399

@@ -106,4 +132,27 @@ default QueryStream createQuery(String query, TxMode tx) {
106132
default CompletableFuture<Result<QueryTransaction>> beginTransaction(TxMode txMode) {
107133
return beginTransaction(txMode, BeginTransactionSettings.newBuilder().build());
108134
}
135+
136+
/**
137+
* Create {@link QueryStream} for executing query with specified {@link TxMode}. The query can contain DML, DDL and
138+
* DCL statements. Supported mix of different statement types depends on the chosen transaction type.
139+
*
140+
* @param query text of query
141+
* @return a ready to execute instance of {@link QueryStream}
142+
*/
143+
default CompletableFuture<Result<OperationProtos.Operation>> executeScript(String query) {
144+
return executeScript(query, Params.empty(), ExecuteScriptSettings.newBuilder().build());
145+
}
146+
147+
/**
148+
* Create {@link QueryStream} for executing query with specified {@link TxMode}. The query can contain DML, DDL and
149+
* DCL statements. Supported mix of different statement types depends on the chosen transaction type.
150+
*
151+
* @param query text of query
152+
* @param tx transaction mode
153+
* @return a ready to execute instance of {@link QueryStream}
154+
*/
155+
default QueryStream fetchScriptResults(String query, TxMode tx) {
156+
return fetchScriptResults(query, tx, Params.empty(), FetchScriptSettings.newBuilder().build());
157+
}
109158
}

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

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
import java.util.concurrent.atomic.AtomicBoolean;
1010
import java.util.concurrent.atomic.AtomicReference;
1111

12+
import com.google.protobuf.Duration;
1213
import com.google.protobuf.TextFormat;
1314
import org.slf4j.Logger;
1415
import org.slf4j.LoggerFactory;
@@ -25,6 +26,7 @@
2526
import tech.ydb.core.settings.BaseRequestSettings;
2627
import tech.ydb.core.utils.URITools;
2728
import tech.ydb.core.utils.UpdatableOptional;
29+
import tech.ydb.proto.OperationProtos;
2830
import tech.ydb.proto.query.YdbQuery;
2931
import tech.ydb.query.QuerySession;
3032
import tech.ydb.query.QueryStream;
@@ -37,6 +39,8 @@
3739
import tech.ydb.query.settings.CreateSessionSettings;
3840
import tech.ydb.query.settings.DeleteSessionSettings;
3941
import tech.ydb.query.settings.ExecuteQuerySettings;
42+
import tech.ydb.query.settings.ExecuteScriptSettings;
43+
import tech.ydb.query.settings.FetchScriptSettings;
4044
import tech.ydb.query.settings.QueryExecMode;
4145
import tech.ydb.query.settings.QueryStatsMode;
4246
import tech.ydb.query.settings.RollbackTransactionSettings;
@@ -244,6 +248,75 @@ void handleTxMeta(String txID) {
244248
};
245249
}
246250

251+
@Override
252+
public CompletableFuture<Result<OperationProtos.Operation>> executeScript(String query, Params params, ExecuteScriptSettings settings) {
253+
YdbQuery.ExecuteScriptRequest.Builder request = YdbQuery.ExecuteScriptRequest.newBuilder()
254+
.setExecMode(mapExecMode(settings.getExecMode()))
255+
.setStatsMode(mapStatsMode(settings.getStatsMode()))
256+
.setScriptContent(YdbQuery.QueryContent.newBuilder()
257+
.setSyntax(YdbQuery.Syntax.SYNTAX_YQL_V1)
258+
.setText(query)
259+
.build());
260+
261+
java.time.Duration ttl = settings.getTtl();
262+
if(ttl != null) {
263+
request.setResultsTtl(Duration.newBuilder().setNanos(settings.getTtl().getNano()));
264+
}
265+
266+
String resourcePool = settings.getResourcePool();
267+
if (resourcePool != null && !resourcePool.isEmpty()) {
268+
request.setPoolId(resourcePool);
269+
}
270+
271+
request.putAllParameters(params.toPb());
272+
273+
GrpcRequestSettings.Builder options = makeOptions(settings);
274+
275+
return rpc.executeScript(request.build(), options.build());
276+
}
277+
278+
@Override
279+
public QueryStream fetchScriptResults(String query, TxMode tx, Params params, FetchScriptSettings settings) {
280+
YdbQuery.FetchScriptResultsRequest.Builder request = YdbQuery.FetchScriptResultsRequest.newBuilder()
281+
.setOperationId(settings.getOperationId())
282+
.setFetchToken(settings.getFetchToken())
283+
.setRowsLimit(10)
284+
.setResultSetIndex(0L);
285+
286+
GrpcRequestSettings.Builder options = makeOptions(settings);
287+
CompletableFuture<Result<YdbQuery.FetchScriptResultsResponse>> response
288+
= rpc.fetchScriptResults(request.build(), options.build());
289+
290+
return null;
291+
/* response.thenApply(result -> result.getValue() )
292+
293+
Result<YdbQuery.FetchScriptResultsResponse> resp = response.get();
294+
YdbQuery.FetchScriptResultsResponse rs = resp.getValue();
295+
296+
rs.getResultSet();
297+
rs.get*/
298+
/*
299+
rs := response.GetResultSet()
300+
columns := rs.GetColumns()
301+
columnNames := make([]string, len(columns))
302+
columnTypes := make([]types.Type, len(columns))
303+
for i := range columns {
304+
columnNames[i] = columns[i].GetName()
305+
columnTypes[i] = types.TypeFromYDB(columns[i].GetType())
306+
}
307+
rows := make([]query.Row, len(rs.GetRows()))
308+
for i, r := range rs.GetRows() {
309+
rows[i] = NewRow(columns, r)
310+
}
311+
312+
return &options.FetchScriptResult{
313+
ResultSetIndex: response.GetResultSetIndex(),
314+
ResultSet: MaterializedResultSet(int(response.GetResultSetIndex()), columnNames, columnTypes, rows),
315+
NextToken: response.GetNextFetchToken(),
316+
}, nil
317+
*/
318+
}
319+
247320
public CompletableFuture<Result<YdbQuery.DeleteSessionResponse>> delete(DeleteSessionSettings settings) {
248321
YdbQuery.DeleteSessionRequest request = YdbQuery.DeleteSessionRequest.newBuilder()
249322
.setSessionId(sessionId)

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

Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,90 @@
1+
package tech.ydb.query.settings;
2+
3+
import tech.ydb.core.grpc.GrpcFlowControl;
4+
import tech.ydb.core.settings.BaseRequestSettings;
5+
6+
import java.time.Duration;
7+
8+
/**
9+
10+
*/
11+
public class ExecuteScriptSettings extends BaseRequestSettings {
12+
private final QueryExecMode execMode;
13+
private final QueryStatsMode statsMode;
14+
private final String resourcePool;
15+
private final Duration ttl;
16+
17+
private ExecuteScriptSettings(Builder builder) {
18+
super(builder);
19+
this.execMode = builder.execMode;
20+
this.statsMode = builder.statsMode;
21+
this.ttl = builder.ttl;
22+
this.resourcePool = builder.resourcePool;
23+
}
24+
25+
public QueryExecMode getExecMode() {
26+
return this.execMode;
27+
}
28+
29+
public Duration getTtl() {
30+
return ttl;
31+
}
32+
33+
public QueryStatsMode getStatsMode() {
34+
return this.statsMode;
35+
}
36+
37+
/**
38+
* Get resource pool for query execution
39+
* @return resource pool name
40+
*/
41+
public String getResourcePool() {
42+
return this.resourcePool;
43+
}
44+
45+
46+
public static Builder newBuilder() {
47+
return new Builder();
48+
}
49+
50+
public static class Builder extends BaseBuilder<Builder> {
51+
private QueryExecMode execMode = QueryExecMode.EXECUTE;
52+
private QueryStatsMode statsMode = QueryStatsMode.NONE;
53+
private String resourcePool = null;
54+
private Duration ttl =null;
55+
56+
public Builder withExecMode(QueryExecMode mode) {
57+
this.execMode = mode;
58+
return this;
59+
}
60+
61+
public Builder withStatsMode(QueryStatsMode mode) {
62+
this.statsMode = mode;
63+
return this;
64+
}
65+
66+
public Builder withTtl(Duration value) {
67+
this.ttl = value;
68+
return this;
69+
}
70+
71+
/**
72+
* Set resource pool which query try to use.
73+
* If no pool specify or poolId is empty or poolId equals "default"
74+
* the unremovable resource pool "default" will be used
75+
*
76+
* @param poolId resource pool identifier
77+
*
78+
* @return builder
79+
*/
80+
public Builder withResourcePool(String poolId) {
81+
this.resourcePool = poolId;
82+
return this;
83+
}
84+
85+
@Override
86+
public ExecuteScriptSettings build() {
87+
return new ExecuteScriptSettings(this);
88+
}
89+
}
90+
}

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

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,52 @@
1+
package tech.ydb.query.settings;
2+
3+
import java.time.Duration;
4+
5+
import tech.ydb.core.settings.BaseRequestSettings;
6+
7+
/**
8+
* operationId_ = "";
9+
* fetchToken_ = "";
10+
*/
11+
public class FetchScriptSettings extends BaseRequestSettings {
12+
private final String operationId;
13+
private final String fetchToken;
14+
15+
private FetchScriptSettings(Builder builder) {
16+
super(builder);
17+
this.operationId = builder.operationId;
18+
this.fetchToken = builder.fetchToken;
19+
}
20+
21+
public String getOperationId() {
22+
return operationId;
23+
}
24+
25+
public String getFetchToken() {
26+
return fetchToken;
27+
}
28+
29+
public static Builder newBuilder() {
30+
return new Builder();
31+
}
32+
33+
public static class Builder extends BaseBuilder<Builder> {
34+
private String operationId = "";
35+
private String fetchToken = "";
36+
37+
public Builder withEOperationId(String operationId ) {
38+
this.operationId = operationId;
39+
return this;
40+
}
41+
42+
public Builder withFetchToken(String fetchToken ) {
43+
this.fetchToken = fetchToken;
44+
return this;
45+
}
46+
47+
@Override
48+
public FetchScriptSettings build() {
49+
return new FetchScriptSettings(this);
50+
}
51+
}
52+
}

0 commit comments

Comments
 (0)