Skip to content

Commit 8f8ab4c

Browse files
alex268alex268
committed
Batched query is used only for executeBatch
1 parent 1c33dcf commit 8f8ab4c

File tree

10 files changed

+199
-57
lines changed

10 files changed

+199
-57
lines changed

‎jdbc/src/main/java/tech/ydb/jdbc/common/TypeDescription.java

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,10 @@ private TypeDescription(Type type,
4242
this.isNull = type.getKind() == Type.Kind.NULL || type.getKind() == Type.Kind.VOID;
4343
}
4444

45+
public String toYqlLiteral() {
46+
return type.toString() + (optional ? "?" : "");
47+
}
48+
4549
public boolean isNull() {
4650
return isNull;
4751
}

‎jdbc/src/main/java/tech/ydb/jdbc/impl/YdbPreparedStatementImpl.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -96,7 +96,7 @@ public int[] executeBatch() throws SQLException {
9696
executeBulkUpsert(query, bulk.getTablePath(), bulk.getBatchedBulk());
9797
} else {
9898
for (Params prm: prepared.getBatchParams()) {
99-
executeDataQuery(query, prepared.getQueryText(prm), prm);
99+
executeDataQuery(query, prepared.getBatchText(prm), prm);
100100
}
101101
}
102102
} finally {

‎jdbc/src/main/java/tech/ydb/jdbc/query/ParamDescription.java

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
package tech.ydb.jdbc.query;
22

3+
import tech.ydb.jdbc.YdbConst;
34
import tech.ydb.jdbc.common.TypeDescription;
45

56
/**
@@ -18,7 +19,7 @@ public ParamDescription(String name, String displayName, TypeDescription type) {
1819
}
1920

2021
public ParamDescription(String name, TypeDescription type) {
21-
this(name, name, type);
22+
this(name, YdbConst.VARIABLE_PARAMETER_PREFIX + name, type);
2223
}
2324

2425
public String name() {

‎jdbc/src/main/java/tech/ydb/jdbc/query/YdbPreparedQuery.java

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
*/
1515
public interface YdbPreparedQuery {
1616
String getQueryText(Params prms) throws SQLException;
17+
String getBatchText(Params prms) throws SQLException;
1718

1819
void clearParameters();
1920

‎jdbc/src/main/java/tech/ydb/jdbc/query/params/BatchedQuery.java

Lines changed: 137 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,6 @@
44
import java.sql.SQLDataException;
55
import java.sql.SQLException;
66
import java.util.ArrayList;
7-
import java.util.Arrays;
87
import java.util.Collections;
98
import java.util.HashMap;
109
import java.util.HashSet;
@@ -36,36 +35,35 @@
3635
* @author Aleksandr Gorshenin
3736
*/
3837
public class BatchedQuery implements YdbPreparedQuery {
39-
private final String yql;
38+
private final String singleQuery;
39+
private final String batchQuery;
4040
private final String batchParamName;
4141
private final Map<String, ParamDescription> paramsByName;
4242
private final ParamDescription[] params;
4343

4444
private final List<StructValue> batchList = new ArrayList<>();
4545
private final Map<String, Value<?>> currentValues = new HashMap<>();
4646

47-
protected BatchedQuery(YdbTypes types, String yql, String listName, List<String> pnames, Map<String, Type> ptypes)
48-
throws SQLException {
49-
this.yql = yql;
50-
this.batchParamName = listName;
47+
protected BatchedQuery(String single, String batched, String prm, ParamDescription[] params) throws SQLException {
48+
this.singleQuery = single;
49+
this.batchQuery = batched;
50+
this.batchParamName = prm;
5151
this.paramsByName = new HashMap<>();
52-
this.params = new ParamDescription[pnames.size()];
52+
this.params = params;
5353

54-
for (int idx = 0; idx < pnames.size(); idx += 1) {
55-
String name = pnames.get(idx);
56-
if (!ptypes.containsKey(name)) {
57-
throw new SQLException(YdbConst.INVALID_BATCH_COLUMN + name);
58-
}
59-
TypeDescription type = types.find(ptypes.get(name));
60-
ParamDescription desc = new ParamDescription(name, YdbConst.VARIABLE_PARAMETER_PREFIX + name, type);
61-
params[idx] = desc;
62-
paramsByName.put(name, desc);
54+
for (ParamDescription pd: params) {
55+
paramsByName.put(pd.name(), pd);
6356
}
6457
}
6558

6659
@Override
6760
public String getQueryText(Params prms) {
68-
return yql;
61+
return singleQuery != null ? singleQuery : batchQuery;
62+
}
63+
64+
@Override
65+
public String getBatchText(Params prms) {
66+
return batchQuery;
6967
}
7068

7169
@Override
@@ -85,7 +83,7 @@ public void clearParameters() {
8583

8684
@Override
8785
public void addBatch() throws SQLException {
88-
batchList.add(getCurrentValues());
86+
batchList.add(StructValue.of(validateValues()));
8987
currentValues.clear();
9088
}
9189

@@ -94,13 +92,13 @@ public void clearBatch() {
9492
batchList.clear();
9593
}
9694

97-
protected StructValue getCurrentValues() throws SQLException {
95+
protected Map<String, Value<?>> validateValues() throws SQLException {
9896
for (ParamDescription prm: params) {
9997
if (!currentValues.containsKey(prm.name())) {
10098
throw new SQLDataException(YdbConst.MISSING_VALUE_FOR_PARAMETER + prm.displayName());
10199
}
102100
}
103-
return StructValue.of(currentValues);
101+
return currentValues;
104102
}
105103

106104
protected List<StructValue> getBatchedValues() {
@@ -109,8 +107,13 @@ protected List<StructValue> getBatchedValues() {
109107

110108
@Override
111109
public Params getCurrentParams() throws SQLException {
112-
ListValue list = ListValue.of(getCurrentValues());
113-
return Params.of(batchParamName, list);
110+
Map<String, Value<?>> vv = validateValues();
111+
if (singleQuery == null) {
112+
return Params.of(batchParamName, ListValue.of(StructValue.of(vv)));
113+
}
114+
Params prms = Params.create(vv.size());
115+
vv.forEach((name, value) -> prms.put(YdbConst.VARIABLE_PARAMETER_PREFIX + name, value));
116+
return prms;
114117
}
115118

116119
@Override
@@ -216,7 +219,15 @@ public static BatchedQuery tryCreateBatched(YdbTypes types, YdbQuery query, Map<
216219
columns[idx] = param;
217220
}
218221

219-
return new BatchedQuery(types, query.getPreparedYql(), listName, Arrays.asList(columns), paramTypes);
222+
ParamDescription[] descriptions = new ParamDescription[columns.length];
223+
for (int idx = 0; idx < columns.length; idx += 1) {
224+
String name = columns[idx];
225+
if (!paramTypes.containsKey(name)) {
226+
throw new SQLException(YdbConst.INVALID_BATCH_COLUMN + name);
227+
}
228+
descriptions[idx] = new ParamDescription(name, types.find(paramTypes.get(name)));
229+
}
230+
return new BatchedQuery(null, query.getPreparedYql(), listName, descriptions);
220231
}
221232

222233
public static BatchedQuery createAutoBatched(YdbTypes types, YqlBatcher batcher, TableDescription description)
@@ -235,33 +246,38 @@ public static BatchedQuery createAutoBatched(YdbTypes types, YqlBatcher batcher,
235246
}
236247
}
237248

238-
StringBuilder sb = new StringBuilder();
239249
Map<String, Type> columnTypes = new HashMap<>();
240-
Map<String, Type> structTypes = new HashMap<>();
241-
List<String> columns = new ArrayList<>();
242-
243250
for (TableColumn column: description.getColumns()) {
244251
columnTypes.put(column.getName(), column.getType());
245252
}
246253

247-
List<String> params = new ArrayList<>();
248-
params.addAll(batcher.getColumns());
249-
params.addAll(batcher.getKeyColumns());
254+
List<String> columns = new ArrayList<>();
255+
columns.addAll(batcher.getColumns());
256+
columns.addAll(batcher.getKeyColumns());
257+
258+
ParamDescription[] params = new ParamDescription[columns.size()];
250259

251-
sb.append("DECLARE $batch AS List<Struct<");
252260
int idx = 1;
253-
for (String column: params) {
261+
for (String column: columns) {
254262
Type type = columnTypes.get(column);
255263
if (type == null) {
256264
return null;
257265
}
258-
if (idx > 1) {
266+
params[idx - 1] = new ParamDescription("p" + idx, column, types.find(type));
267+
idx++;
268+
}
269+
270+
return new BatchedQuery(simpleQuery(batcher, params), batchQuery(batcher, params), "$batch", params);
271+
}
272+
273+
private static String batchQuery(YqlBatcher batcher, ParamDescription[] params) {
274+
StringBuilder sb = new StringBuilder();
275+
sb.append("DECLARE $batch AS List<Struct<");
276+
for (int idx = 0; idx < params.length; idx++) {
277+
if (idx > 0) {
259278
sb.append(", ");
260279
}
261-
sb.append("p").append(idx).append(":").append(type.toString());
262-
structTypes.put("p" + idx, type);
263-
columns.add("p" + idx);
264-
idx++;
280+
sb.append(params[idx].name()).append(":").append(params[idx].type().toYqlLiteral());
265281
}
266282
sb.append(">>;\n");
267283

@@ -282,20 +298,97 @@ public static BatchedQuery createAutoBatched(YdbTypes types, YqlBatcher batcher,
282298
sb.append("DELETE FROM `").append(batcher.getTableName()).append("` ON SELECT ");
283299
break;
284300
default:
285-
return null;
301+
return "UNSUPPORTED CMD " + batcher.getCommand();
286302
}
287303

288-
idx = 1;
289-
for (String column: params) {
290-
if (idx > 1) {
304+
for (int idx = 0; idx < params.length; idx++) {
305+
if (idx > 0) {
291306
sb.append(", ");
292307
}
293-
sb.append("p").append(idx).append(" AS `").append(column).append("`");
294-
idx++;
308+
sb.append(params[idx].name()).append(" AS `").append(params[idx].displayName()).append("`");
295309
}
296310

297311
sb.append(" FROM AS_TABLE($batch);");
312+
return sb.toString();
313+
}
314+
315+
private static String simpleQuery(YqlBatcher batcher, ParamDescription[] params) {
316+
StringBuilder sb = new StringBuilder();
317+
for (ParamDescription p : params) {
318+
sb.append("DECLARE ").append(YdbConst.VARIABLE_PARAMETER_PREFIX).append(p.name())
319+
.append(" AS ").append(p.type().toYqlLiteral()).append(";\n");
320+
}
298321

299-
return new BatchedQuery(types, sb.toString(), "$batch", columns, structTypes);
322+
switch (batcher.getCommand()) {
323+
case UPSERT:
324+
sb.append("UPSERT INTO `").append(batcher.getTableName()).append("` (");
325+
appendColumns(sb, params);
326+
sb.append(") VALUES (");
327+
appendValues(sb, params);
328+
sb.append(");");
329+
break;
330+
case INSERT:
331+
sb.append("INSERT INTO `").append(batcher.getTableName()).append("` (");
332+
appendColumns(sb, params);
333+
sb.append(") VALUES (");
334+
appendValues(sb, params);
335+
sb.append(");");
336+
break;
337+
case REPLACE:
338+
sb.append("REPLACE INTO `").append(batcher.getTableName()).append("` (");
339+
appendColumns(sb, params);
340+
sb.append(") VALUES (");
341+
appendValues(sb, params);
342+
sb.append(");");
343+
break;
344+
case UPDATE:
345+
sb.append("UPDATE `").append(batcher.getTableName()).append("` SET ");
346+
for (int idx = 0; idx < batcher.getColumns().size(); idx++) {
347+
if (idx > 0) {
348+
sb.append(", ");
349+
}
350+
sb.append('`').append(params[idx].displayName()).append("` = ")
351+
.append(YdbConst.VARIABLE_PARAMETER_PREFIX).append(params[idx].name());
352+
}
353+
sb.append(" WHERE ");
354+
appendKeys(sb, params, batcher.getColumns().size());
355+
break;
356+
case DELETE:
357+
sb.append("DELETE FROM `").append(batcher.getTableName()).append("` WHERE ");
358+
appendKeys(sb, params, batcher.getColumns().size());
359+
break;
360+
default:
361+
break;
362+
}
363+
364+
return sb.toString();
365+
}
366+
367+
private static void appendColumns(StringBuilder sb, ParamDescription[] params) {
368+
for (int idx = 0; idx < params.length; idx++) {
369+
if (idx > 0) {
370+
sb.append(", ");
371+
}
372+
sb.append('`').append(params[idx].displayName()).append('`');
373+
}
374+
}
375+
376+
private static void appendValues(StringBuilder sb, ParamDescription[] params) {
377+
for (int idx = 0; idx < params.length; idx++) {
378+
if (idx > 0) {
379+
sb.append(", ");
380+
}
381+
sb.append(YdbConst.VARIABLE_PARAMETER_PREFIX).append(params[idx].name());
382+
}
383+
}
384+
385+
private static void appendKeys(StringBuilder sb, ParamDescription[] params, int firstKeyIdx) {
386+
for (int idx = firstKeyIdx; idx < params.length; idx++) {
387+
if (idx > firstKeyIdx) {
388+
sb.append(" AND ");
389+
}
390+
sb.append('`').append(params[idx].displayName()).append("` = ")
391+
.append(YdbConst.VARIABLE_PARAMETER_PREFIX).append(params[idx].name());
392+
}
300393
}
301394
}

‎jdbc/src/main/java/tech/ydb/jdbc/query/params/BulkUpsertQuery.java

Lines changed: 13 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -8,11 +8,13 @@
88
import java.util.stream.Collectors;
99

1010
import tech.ydb.jdbc.common.YdbTypes;
11+
import tech.ydb.jdbc.query.ParamDescription;
1112
import tech.ydb.table.description.TableColumn;
1213
import tech.ydb.table.description.TableDescription;
1314
import tech.ydb.table.values.ListType;
1415
import tech.ydb.table.values.ListValue;
1516
import tech.ydb.table.values.StructType;
17+
import tech.ydb.table.values.StructValue;
1618
import tech.ydb.table.values.Type;
1719

1820
/**
@@ -23,19 +25,18 @@ public class BulkUpsertQuery extends BatchedQuery {
2325
private final String tablePath;
2426
private final ListType bulkType;
2527

26-
private BulkUpsertQuery(YdbTypes types, String path, String yql, List<String> columns, Map<String, Type> paramTypes)
27-
throws SQLException {
28-
super(types, yql, "$bulk", columns, paramTypes);
29-
this.tablePath = path;
30-
this.bulkType = ListType.of(StructType.of(paramTypes));
28+
private BulkUpsertQuery(String tablePath, String yql, ListType tp, ParamDescription[] params) throws SQLException {
29+
super(null, yql, "$bulk", params);
30+
this.tablePath = tablePath;
31+
this.bulkType = tp;
3132
}
3233

3334
public String getTablePath() {
3435
return tablePath;
3536
}
3637

3738
public ListValue getCurrentBulk() throws SQLException {
38-
return bulkType.newValue(Collections.singletonList(getCurrentValues()));
39+
return bulkType.newValue(Collections.singletonList(StructValue.of(validateValues())));
3940
}
4041

4142
public ListValue getBatchedBulk() {
@@ -56,14 +57,18 @@ public static BulkUpsertQuery build(YdbTypes types, String path, List<String> co
5657
columnTypes.put(column.getName(), column.getType());
5758
}
5859

60+
ParamDescription[] params = new ParamDescription[columns.size()];
5961
Map<String, Type> structTypes = new HashMap<>();
62+
int idx = 0;
6063
for (String column: columns) {
6164
if (!columnTypes.containsKey(column)) {
6265
throw new SQLException("Cannot parse BULK upsert: column " + column + " not found");
6366
}
64-
structTypes.put(column, columnTypes.get(column));
67+
Type type = columnTypes.get(column);
68+
structTypes.put(column, type);
69+
params[idx++] = new ParamDescription(column, types.find(type));
6570
}
6671

67-
return new BulkUpsertQuery(types, path, yql.toString(), columns, structTypes);
72+
return new BulkUpsertQuery(path, yql.toString(), ListType.of(StructType.of(structTypes)), params);
6873
}
6974
}

‎jdbc/src/main/java/tech/ydb/jdbc/query/params/InMemoryQuery.java

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,11 @@ public String getQueryText(Params prms) throws SQLException {
6161
return query.toString();
6262
}
6363

64+
@Override
65+
public String getBatchText(Params prms) throws SQLException {
66+
return getQueryText(prms);
67+
}
68+
6469
@Override
6570
public int parametersCount() {
6671
return parameters.size();

0 commit comments

Comments
 (0)