Skip to content

Commit d4fe2fd

Browse files
alex268alex268
committed
Add auto declare of types
1 parent 71fae45 commit d4fe2fd

File tree

5 files changed

+62
-44
lines changed

5 files changed

+62
-44
lines changed

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

Lines changed: 45 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -2,12 +2,14 @@
22

33

44
import java.util.ArrayList;
5-
import java.util.Collections;
65
import java.util.List;
6+
import java.util.Map;
77
import java.util.function.Supplier;
88

99
import tech.ydb.jdbc.YdbConst;
1010
import tech.ydb.jdbc.settings.YdbOperationProperties;
11+
import tech.ydb.table.query.Params;
12+
import tech.ydb.table.values.Value;
1113

1214
/**
1315
*
@@ -31,7 +33,7 @@ private YdbQuery(Builder builder) {
3133
String sql = updateAlternativePrefix(originSQL);
3234
if (builder.props.isDisableJdbcParametersSupport()) {
3335
this.originYQL = sql;
34-
this.extraParams = Collections.emptyList();
36+
this.extraParams = null;
3537
} else {
3638
ArgNameGenerator generator = new ArgNameGenerator(sql);
3739
this.originYQL = JdbcLexer.jdbc2yql(sql, generator);
@@ -43,17 +45,22 @@ public String originSQL() {
4345
return originSQL;
4446
}
4547

46-
private static String updateAlternativePrefix(String sql) {
47-
String updated = sql;
48-
for (QueryType type : QueryType.values()) {
49-
if (updated.contains(type.getAlternativePrefix())) {
50-
updated = updated.replace(type.getAlternativePrefix(), type.getPrefix()); // Support alternative mode
51-
}
48+
public boolean hasFreeParameters() {
49+
return extraParams != null && !extraParams.isEmpty();
50+
}
51+
52+
public String getParameterName(int parameterIndex) {
53+
if (extraParams == null) {
54+
return YdbConst.INDEXED_PARAMETER_PREFIX + parameterIndex;
5255
}
53-
return updated;
56+
57+
if (parameterIndex <= 0 || parameterIndex > extraParams.size()) {
58+
throw new IllegalArgumentException("Wrong argument index " + parameterIndex);
59+
}
60+
return extraParams.get(parameterIndex - 1);
5461
}
5562

56-
public String nativeSql(List<String> argTypes) {
63+
public String getYqlQuery(Params params) {
5764
StringBuilder yql = new StringBuilder();
5865

5966
if (enforceV1) {
@@ -63,21 +70,26 @@ public String nativeSql(List<String> argTypes) {
6370
}
6471
}
6572

66-
if (argTypes != null) {
67-
if (argTypes.size() != extraParams.size()) {
68-
throw new IllegalArgumentException("Wrong arguments count, expected "
69-
+ extraParams.size() + ", but got " + argTypes.size());
70-
}
73+
if (extraParams != null) {
74+
if (params != null) {
75+
Map<String, Value<?>> values = params.values();
76+
for (int idx = 0; idx < extraParams.size(); idx += 1) {
77+
String prm = extraParams.get(idx);
78+
if (!values.containsKey(prm)) {
79+
throw new IllegalArgumentException("Empty JDBC argument " + idx);
80+
}
81+
82+
String prmType = values.get(prm).getType().toString();
83+
yql.append("DECLARE ")
84+
.append(prm)
85+
.append(" AS ")
86+
.append(prmType)
87+
.append(";\n");
7188

72-
for (int idx = 0; idx < extraParams.size(); idx += 1) {
73-
yql.append("DECLARE ")
74-
.append(extraParams.get(idx))
75-
.append(" AS ")
76-
.append(argTypes.get(idx))
77-
.append(";\n");
89+
}
90+
} else if (!extraParams.isEmpty()) {
91+
yql.append("-- DECLARE ").append(extraParams.size()).append(" PARAMETERS").append("\n");
7892
}
79-
} else if (!extraParams.isEmpty()) {
80-
yql.append("-- DECLARE ").append(extraParams.size()).append(" PARAMETERS").append("\n");
8193
}
8294

8395
yql.append(originYQL);
@@ -92,6 +104,16 @@ public QueryType type() {
92104
return type;
93105
}
94106

107+
private static String updateAlternativePrefix(String sql) {
108+
String updated = sql;
109+
for (QueryType type : QueryType.values()) {
110+
if (updated.contains(type.getAlternativePrefix())) {
111+
updated = updated.replace(type.getAlternativePrefix(), type.getPrefix()); // Support alternative mode
112+
}
113+
}
114+
return updated;
115+
}
116+
95117
private static QueryType decodeQueryType(String sql, boolean autoDetect) {
96118
/*
97119
Need some logic to figure out - if this is a scheme, data, scan or explain plan query.

‎src/main/java/tech/ydb/jdbc/connection/YdbConnectionImpl.java‎

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -91,7 +91,7 @@ public YdbPreparedStatement prepareStatement(String sql) throws SQLException {
9191

9292
@Override
9393
public String nativeSQL(String sql) {
94-
return YdbQuery.from(ctx.getOperationProperties(), sql).build().nativeSql(null);
94+
return YdbQuery.from(ctx.getOperationProperties(), sql).build().getYqlQuery(null);
9595
}
9696

9797
private void updateState(YdbTxState newState) {
@@ -246,7 +246,7 @@ public void executeSchemeQuery(YdbQuery query, YdbExecutor executor) throws SQLE
246246

247247
// Scheme query does not affect transactions or result sets
248248
ExecuteSchemeQuerySettings settings = executor.withTimeouts(new ExecuteSchemeQuerySettings());
249-
final String yql = query.nativeSql(null);
249+
final String yql = query.getYqlQuery(null);
250250

251251
try (Session session = executor.createSession(ctx)) {
252252
executor.execute(QueryType.SCHEME_QUERY + " >>\n" + yql, () -> session.executeSchemeQuery(yql, settings));
@@ -261,7 +261,7 @@ public DataQueryResult executeDataQuery(YdbQuery query, Params params, YdbExecut
261261
if (!query.keepInCache()) {
262262
settings.disableQueryCache();
263263
}
264-
final String yql = query.nativeSql(null);
264+
final String yql = query.getYqlQuery(params);
265265

266266
Session session = state.getSession(ctx, executor);
267267
try {
@@ -284,7 +284,7 @@ public ResultSetReader executeScanQuery(YdbQuery query, Params params, YdbExecut
284284
Duration scanQueryTimeout = ctx.getOperationProperties().getScanQueryTimeout();
285285
ExecuteScanQuerySettings settings = ExecuteScanQuerySettings.newBuilder()
286286
.withRequestTimeout(scanQueryTimeout).build();
287-
String yql = query.nativeSql(null);
287+
String yql = query.getYqlQuery(params);
288288

289289
Collection<ResultSetReader> resultSets = new LinkedBlockingQueue<>();
290290
try (Session session = executor.createSession(ctx)) {
@@ -300,7 +300,7 @@ public ExplainDataQueryResult executeExplainQuery(YdbQuery query, YdbExecutor ex
300300
ensureOpened();
301301

302302
ExplainDataQuerySettings settings = executor.withTimeouts(new ExplainDataQuerySettings());
303-
String yql = query.nativeSql(null);
303+
String yql = query.getYqlQuery(null);
304304

305305
try (Session session = executor.createSession(ctx)) {
306306
return executor.call(
@@ -311,7 +311,7 @@ public ExplainDataQueryResult executeExplainQuery(YdbQuery query, YdbExecutor ex
311311

312312
private DataQuery prepareDataQuery(YdbQuery query) throws SQLException {
313313
PrepareDataQuerySettings cfg = executor.withTimeouts(new PrepareDataQuerySettings());
314-
String yql = query.nativeSql(null);
314+
String yql = query.getYqlQuery(null);
315315

316316
try (Session session = executor.createSession(ctx)) {
317317
return executor.call("Preparing Query >>\n" + yql, () -> session.prepareDataQuery(yql, cfg));
@@ -395,7 +395,8 @@ private YdbPreparedStatement prepareStatementImpl(String sql, int resultSetType,
395395
YdbQuery query = YdbQuery.from(props, sql).build();
396396

397397
if (mode == PreparedStatementMode.IN_MEMORY
398-
|| (mode == PreparedStatementMode.DEFAULT && !props.isAlwaysPrepareDataQuery())) {
398+
|| (mode == PreparedStatementMode.DEFAULT && !props.isAlwaysPrepareDataQuery())
399+
|| query.hasFreeParameters()) {
399400
return new YdbPreparedStatementImpl(this, query, resultSetType);
400401
}
401402

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

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,6 @@
2929

3030
import static tech.ydb.jdbc.YdbConst.BATCH_INVALID;
3131
import static tech.ydb.jdbc.YdbConst.DEFAULT_BATCH_PARAMETER;
32-
import static tech.ydb.jdbc.YdbConst.INDEXED_PARAMETER_PREFIX;
3332
import static tech.ydb.jdbc.YdbConst.PARAMETER_TYPE_UNKNOWN;
3433
import static tech.ydb.jdbc.YdbConst.TRY_EXECUTE_ON_BATCH_STATEMENT;
3534
import static tech.ydb.jdbc.YdbConst.UNKNOWN_PARAMETER_IN_BATCH;
@@ -47,12 +46,12 @@ public YdbPreparedStatementImpl(YdbConnection connection, YdbQuery query, int re
4746

4847
YdbOperationProperties properties = connection.getCtx().getOperationProperties();
4948
this.enforceVariablePrefix = properties.isEnforceVariablePrefix();
50-
this.clearParameters();
49+
this.state.params = new LinkedHashMap<>(this.state.descriptions.size());
5150
}
5251

5352
@Override
5453
public void clearParameters() {
55-
this.state.params = new LinkedHashMap<>(this.state.descriptions.size());
54+
this.state.params.clear();
5655
}
5756

5857
@Override
@@ -127,7 +126,7 @@ protected void setImpl(String parameterName, @Nullable Object value,
127126
protected void setImpl(int parameterIndex, @Nullable Object value,
128127
int sqlType, @Nullable String typeName, @Nullable Type type)
129128
throws SQLException {
130-
setImpl(INDEXED_PARAMETER_PREFIX + parameterIndex, value, sqlType, typeName, type);
129+
setImpl(query.getParameterName(parameterIndex), value, sqlType, typeName, type);
131130
}
132131

133132
@Override

‎src/main/java/tech/ydb/jdbc/statement/YdbPreparedStatementWithDataQueryImpl.java‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,6 @@
1818
import tech.ydb.table.values.Value;
1919

2020
import static tech.ydb.jdbc.YdbConst.BATCH_UNSUPPORTED;
21-
import static tech.ydb.jdbc.YdbConst.INDEXED_PARAMETER_PREFIX;
2221
import static tech.ydb.jdbc.YdbConst.PARAMETER_NOT_FOUND;
2322
import static tech.ydb.jdbc.YdbConst.VARIABLE_PARAMETER_PREFIX;
2423

@@ -37,7 +36,7 @@ public YdbPreparedStatementWithDataQueryImpl(YdbConnectionImpl connection,
3736
this.dataQuery = dataQuery;
3837
this.cfg = asPreparedConfiguration(dataQuery.types());
3938
this.enforceVariablePrefix = connection.getCtx().getOperationProperties().isEnforceVariablePrefix();
40-
this.clearParameters();
39+
this.state.params = dataQuery.newParams();
4140
}
4241

4342
@Override
@@ -86,7 +85,7 @@ protected void setImpl(String origParameterName, @Nullable Object x,
8685
protected void setImpl(int parameterIndex, @Nullable Object x,
8786
int sqlType, @Nullable String typeName, @Nullable Type type)
8887
throws SQLException {
89-
setImpl(INDEXED_PARAMETER_PREFIX + parameterIndex, x, sqlType, typeName, null);
88+
setImpl(query.getParameterName(parameterIndex), x, sqlType, typeName, null);
9089
}
9190

9291
@Override

‎src/test/java/tech/ydb/jdbc/YdbDriverIntegrationTest.java‎

Lines changed: 4 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -136,17 +136,14 @@ public void testYdb() throws SQLException {
136136

137137
PreparedStatement ps = connection
138138
.prepareStatement("" +
139-
"declare $p1 as Int32;\n" +
140-
"declare $p2 as Text;\n" +
141-
"upsert into jdbc_table_sample (id, value) values ($p1, $p2)");
139+
"upsert into jdbc_table_sample (id, value) values (?, ?)");
142140
ps.setInt(1, 1);
143141
ps.setString(2, "value-1");
144142
ps.executeUpdate();
145143

146-
YdbPreparedStatement yps = ps.unwrap(YdbPreparedStatement.class);
147-
yps.setInt("p1", 2);
148-
yps.setString("p2", "value-2");
149-
yps.executeUpdate();
144+
ps.setInt(1, 2);
145+
ps.setString(2, "value-2");
146+
ps.executeUpdate();
150147

151148
connection.commit();
152149

0 commit comments

Comments
 (0)