diff --git a/jdbc/src/main/java/tech/ydb/jdbc/YdbTracer.java b/jdbc/src/main/java/tech/ydb/jdbc/YdbTracer.java index 16f30c1..2c78221 100644 --- a/jdbc/src/main/java/tech/ydb/jdbc/YdbTracer.java +++ b/jdbc/src/main/java/tech/ydb/jdbc/YdbTracer.java @@ -1,6 +1,5 @@ package tech.ydb.jdbc; - import tech.ydb.jdbc.impl.YdbTracerImpl; @@ -10,12 +9,20 @@ * @author Aleksandr Gorshenin */ public interface YdbTracer { - static void clear() { - YdbTracerImpl.clear(); + interface Storage { + YdbTracer get(); + + default void clear() { + + } } static YdbTracer current() { - return YdbTracerImpl.current(); + return YdbTracerImpl.ENABLED.get(); + } + + static void clear() { + YdbTracerImpl.ENABLED.clear(); } void setId(String id); diff --git a/jdbc/src/main/java/tech/ydb/jdbc/common/TypeDescription.java b/jdbc/src/main/java/tech/ydb/jdbc/common/TypeDescription.java index b890efb..9d0c914 100644 --- a/jdbc/src/main/java/tech/ydb/jdbc/common/TypeDescription.java +++ b/jdbc/src/main/java/tech/ydb/jdbc/common/TypeDescription.java @@ -42,6 +42,10 @@ private TypeDescription(Type type, this.isNull = type.getKind() == Type.Kind.NULL || type.getKind() == Type.Kind.VOID; } + public String toYqlLiteral() { + return type.toString() + (optional ? "?" : ""); + } + public boolean isNull() { return isNull; } diff --git a/jdbc/src/main/java/tech/ydb/jdbc/context/BaseYdbExecutor.java b/jdbc/src/main/java/tech/ydb/jdbc/context/BaseYdbExecutor.java index 7f9df95..1249aee 100644 --- a/jdbc/src/main/java/tech/ydb/jdbc/context/BaseYdbExecutor.java +++ b/jdbc/src/main/java/tech/ydb/jdbc/context/BaseYdbExecutor.java @@ -40,13 +40,11 @@ public abstract class BaseYdbExecutor implements YdbExecutor { private final boolean useStreamResultSet; private final AtomicReference currResult; - protected final boolean traceEnabled; protected final String prefixPragma; protected final YdbTypes types; public BaseYdbExecutor(YdbContext ctx) { this.retryCtx = ctx.getRetryCtx(); - this.traceEnabled = ctx.isTxTracerEnabled(); this.sessionTimeout = ctx.getOperationProperties().getSessionTimeout(); this.useStreamResultSet = ctx.getOperationProperties().getUseStreamResultSets(); this.tableClient = ctx.getTableClient(); diff --git a/jdbc/src/main/java/tech/ydb/jdbc/context/YdbContext.java b/jdbc/src/main/java/tech/ydb/jdbc/context/YdbContext.java index 6e15fcd..b4eb548 100644 --- a/jdbc/src/main/java/tech/ydb/jdbc/context/YdbContext.java +++ b/jdbc/src/main/java/tech/ydb/jdbc/context/YdbContext.java @@ -27,6 +27,7 @@ import tech.ydb.jdbc.YdbTracer; import tech.ydb.jdbc.common.YdbTypes; import tech.ydb.jdbc.exception.ExceptionFactory; +import tech.ydb.jdbc.impl.YdbTracerImpl; import tech.ydb.jdbc.impl.YdbTracerNone; import tech.ydb.jdbc.query.QueryType; import tech.ydb.jdbc.query.YdbPreparedQuery; @@ -90,6 +91,8 @@ public class YdbContext implements AutoCloseable { private final boolean autoResizeSessionPool; private final AtomicInteger connectionsCount = new AtomicInteger(); + private YdbTracer.Storage tracerStorage; + private YdbContext( YdbConfig config, YdbOperationProperties operationProperties, @@ -137,12 +140,18 @@ private YdbContext( prefixPath = transport.getDatabase(); prefixPragma = ""; } + + this.tracerStorage = config.isTxTracedEnabled() ? YdbTracerImpl.ENABLED : YdbTracerNone.DISABLED; } public YdbTypes getTypes() { return types; } + public void setTracerStorage(YdbTracer.Storage storage) { + this.tracerStorage = storage; + } + /** * Grpc Transport for other API YDB server clients * @@ -152,6 +161,10 @@ public GrpcTransport getGrpcTransport() { return grpcTransport; } + public YdbTracer getTracer() { + return tracerStorage.get(); + } + private String joined(String path1, String path2) { return path1.endsWith("/") || path2.startsWith("/") ? path1 + path2 : path1 + "/" + path2; } @@ -188,14 +201,6 @@ public String getUsername() { return config.getUsername(); } - public YdbTracer getTracer() { - return config.isTxTracedEnabled() ? YdbTracer.current() : YdbTracerNone.current(); - } - - public boolean isTxTracerEnabled() { - return config.isTxTracedEnabled(); - } - public YdbExecutor createExecutor() throws SQLException { if (config.isUseQueryService()) { return new QueryServiceExecutor(this, operationProps.getTransactionLevel(), operationProps.isAutoCommit()); @@ -264,6 +269,8 @@ public void register() { } public void deregister() { + tracerStorage.clear(); + int actual = connectionsCount.decrementAndGet(); int maxSize = tableClient.sessionPoolStats().getMaxSize(); if (autoResizeSessionPool && maxSize > SESSION_POOL_RESIZE_STEP) { @@ -409,9 +416,9 @@ public YdbPreparedQuery findOrPrepareParams(YdbQuery query, YdbPrepareMode mode) String tablePath = joined(getPrefixPath(), query.getYqlBatcher().getTableName()); TableDescription description = tableDescribeCache.getIfPresent(tablePath); if (description == null) { - YdbTracer tracer = getTracer(); + YdbTracer tracer = tracerStorage.get(); tracer.trace("--> describe table"); - tracer.query(tablePath); + tracer.trace(tablePath); DescribeTableSettings settings = withDefaultTimeout(new DescribeTableSettings()); Result result = retryCtx.supplyResult( @@ -449,9 +456,9 @@ public YdbPreparedQuery findOrPrepareParams(YdbQuery query, YdbPrepareMode mode) Map queryTypes = queryParamsCache.getIfPresent(query.getOriginQuery()); if (queryTypes == null) { String yql = prefixPragma + query.getPreparedYql(); - YdbTracer tracer = getTracer(); + YdbTracer tracer = tracerStorage.get(); tracer.trace("--> prepare data query"); - tracer.query(yql); + tracer.trace(yql); PrepareDataQuerySettings settings = withDefaultTimeout(new PrepareDataQuerySettings()); Result result = retryCtx.supplyResult( diff --git a/jdbc/src/main/java/tech/ydb/jdbc/impl/YdbConnectionImpl.java b/jdbc/src/main/java/tech/ydb/jdbc/impl/YdbConnectionImpl.java index bbbfe1c..02cba57 100644 --- a/jdbc/src/main/java/tech/ydb/jdbc/impl/YdbConnectionImpl.java +++ b/jdbc/src/main/java/tech/ydb/jdbc/impl/YdbConnectionImpl.java @@ -26,7 +26,6 @@ import tech.ydb.jdbc.YdbPrepareMode; import tech.ydb.jdbc.YdbPreparedStatement; import tech.ydb.jdbc.YdbStatement; -import tech.ydb.jdbc.YdbTracer; import tech.ydb.jdbc.context.YdbContext; import tech.ydb.jdbc.context.YdbExecutor; import tech.ydb.jdbc.context.YdbValidator; @@ -108,7 +107,6 @@ public void close() throws SQLException { validator.clearWarnings(); executor.close(); ctx.deregister(); - YdbTracer.clear(); } @Override @@ -218,24 +216,39 @@ public YdbStatement createStatement(int resultSetType, int resultSetConcurrency, public YdbPreparedStatement prepareStatement(String origSql, int resultSetType, int resultSetConcurrency, int resultSetHoldability) throws SQLException { checkStatementParams(resultSetType, resultSetConcurrency, resultSetHoldability); - return prepareStatement(origSql, resultSetType, YdbPrepareMode.AUTO); + return prepareStatement(origSql, resultSetType, null, YdbPrepareMode.AUTO); } @Override public YdbPreparedStatement prepareStatement(String sql, YdbPrepareMode mode) throws SQLException { - return prepareStatement(sql, ResultSet.TYPE_SCROLL_INSENSITIVE, mode); + return prepareStatement(sql, ResultSet.TYPE_SCROLL_INSENSITIVE, null, mode); } @Override public YdbPreparedStatement prepareStatement(String sql, int autoGeneratedKeys) throws SQLException { - if (autoGeneratedKeys != Statement.NO_GENERATED_KEYS) { - throw new SQLFeatureNotSupportedException(YdbConst.AUTO_GENERATED_KEYS_UNSUPPORTED); + if (autoGeneratedKeys != Statement.RETURN_GENERATED_KEYS) { + return prepareStatement(sql); } - return prepareStatement(sql, ResultSet.TYPE_FORWARD_ONLY, ResultSet.CONCUR_READ_ONLY, - ResultSet.HOLD_CURSORS_OVER_COMMIT); + return prepareStatement(sql, new String[]{"*"}); + } + + @Override + public YdbPreparedStatement prepareStatement(String sql, int[] columnIndexes) throws SQLException { + if (columnIndexes != null && columnIndexes.length == 0) { + return prepareStatement(sql); + } + throw new SQLFeatureNotSupportedException(YdbConst.AUTO_GENERATED_KEYS_UNSUPPORTED); + } + + @Override + public YdbPreparedStatement prepareStatement(String sql, String[] columnNames) throws SQLException { + if (columnNames == null || columnNames.length == 0) { + return prepareStatement(sql); + } + return prepareStatement(sql, ResultSet.TYPE_FORWARD_ONLY, columnNames, YdbPrepareMode.AUTO); } - private YdbPreparedStatement prepareStatement(String sql, int resultSetType, YdbPrepareMode mode) + private YdbPreparedStatement prepareStatement(String sql, int type, String[] columnNames, YdbPrepareMode mode) throws SQLException { validator.clearWarnings(); @@ -243,7 +256,7 @@ private YdbPreparedStatement prepareStatement(String sql, int resultSetType, Ydb YdbQuery query = ctx.findOrParseYdbQuery(sql); YdbPreparedQuery params = ctx.findOrPrepareParams(query, mode); ctx.getTracer().trace("create prepared statement"); - return new YdbPreparedStatementImpl(this, query, params, resultSetType); + return new YdbPreparedStatementImpl(this, query, params, type); } @Override @@ -331,16 +344,6 @@ public CallableStatement prepareCall(String sql, int resultSetType, int resultSe throw new SQLFeatureNotSupportedException(YdbConst.PREPARED_CALLS_UNSUPPORTED); } - @Override - public YdbPreparedStatement prepareStatement(String sql, int[] columnIndexes) throws SQLException { - throw new SQLFeatureNotSupportedException(YdbConst.AUTO_GENERATED_KEYS_UNSUPPORTED); - } - - @Override - public YdbPreparedStatement prepareStatement(String sql, String[] columnNames) throws SQLException { - throw new SQLFeatureNotSupportedException(YdbConst.AUTO_GENERATED_KEYS_UNSUPPORTED); - } - @Override public Clob createClob() throws SQLException { throw new SQLFeatureNotSupportedException(YdbConst.CLOB_UNSUPPORTED); diff --git a/jdbc/src/main/java/tech/ydb/jdbc/impl/YdbPreparedStatementImpl.java b/jdbc/src/main/java/tech/ydb/jdbc/impl/YdbPreparedStatementImpl.java index 42e3192..8f9feb3 100644 --- a/jdbc/src/main/java/tech/ydb/jdbc/impl/YdbPreparedStatementImpl.java +++ b/jdbc/src/main/java/tech/ydb/jdbc/impl/YdbPreparedStatementImpl.java @@ -96,7 +96,7 @@ public int[] executeBatch() throws SQLException { executeBulkUpsert(query, bulk.getTablePath(), bulk.getBatchedBulk()); } else { for (Params prm: prepared.getBatchParams()) { - executeDataQuery(query, prepared.getQueryText(prm), prm); + executeDataQuery(query, prepared.getBatchText(prm), prm); } } } finally { diff --git a/jdbc/src/main/java/tech/ydb/jdbc/impl/YdbTracerImpl.java b/jdbc/src/main/java/tech/ydb/jdbc/impl/YdbTracerImpl.java index cefc0e1..8b65fea 100644 --- a/jdbc/src/main/java/tech/ydb/jdbc/impl/YdbTracerImpl.java +++ b/jdbc/src/main/java/tech/ydb/jdbc/impl/YdbTracerImpl.java @@ -16,9 +16,27 @@ */ public class YdbTracerImpl implements YdbTracer { private static final Logger LOGGER = Logger.getLogger(YdbTracer.class.getName()); - private static final ThreadLocal LOCAL = new ThreadLocal<>(); + private static final ThreadLocal LOCAL = new ThreadLocal<>(); private static final AtomicLong ANONYMOUS_COUNTER = new AtomicLong(0); + public static final Storage ENABLED = new Storage() { + @Override + public YdbTracer get() { + YdbTracer tracer = LOCAL.get(); + if (tracer == null) { + tracer = new YdbTracerImpl(); + LOCAL.set(tracer); + } + + return tracer; + } + + @Override + public void clear() { + LOCAL.remove(); + } + }; + private final Date startDate = new Date(); private final long startedAt = System.currentTimeMillis(); private final List records = new ArrayList<>(); @@ -26,7 +44,6 @@ public class YdbTracerImpl implements YdbTracer { private String txID = null; private String label = null; private boolean isMarked = false; - private boolean isClosed = false; private class Record { private final long executedAt = System.currentTimeMillis(); @@ -39,20 +56,6 @@ private class Record { } } - public static void clear() { - LOCAL.remove(); - } - - public static YdbTracer current() { - YdbTracerImpl tracer = LOCAL.get(); - if (tracer == null || tracer.isClosed) { - tracer = new YdbTracerImpl(); - LOCAL.set(tracer); - } - - return tracer; - } - @Override public void trace(String message) { records.add(new Record(message, false)); @@ -82,8 +85,6 @@ public void markToPrint(String label) { @Override public void close() { - isClosed = true; - LOCAL.remove(); final Level level = isMarked ? Level.INFO : Level.FINE; diff --git a/jdbc/src/main/java/tech/ydb/jdbc/impl/YdbTracerNone.java b/jdbc/src/main/java/tech/ydb/jdbc/impl/YdbTracerNone.java index f5808de..e681ee1 100644 --- a/jdbc/src/main/java/tech/ydb/jdbc/impl/YdbTracerNone.java +++ b/jdbc/src/main/java/tech/ydb/jdbc/impl/YdbTracerNone.java @@ -7,11 +7,8 @@ * @author Aleksandr Gorshenin */ public class YdbTracerNone implements YdbTracer { - private static final YdbTracerNone INSTANCE = new YdbTracerNone(); - - public static YdbTracer current() { - return INSTANCE; - } + private static final YdbTracer NONE = new YdbTracerNone(); + public static final YdbTracer.Storage DISABLED = () -> NONE; @Override public void trace(String message) { } diff --git a/jdbc/src/main/java/tech/ydb/jdbc/query/ParamDescription.java b/jdbc/src/main/java/tech/ydb/jdbc/query/ParamDescription.java index b033970..beb84f0 100644 --- a/jdbc/src/main/java/tech/ydb/jdbc/query/ParamDescription.java +++ b/jdbc/src/main/java/tech/ydb/jdbc/query/ParamDescription.java @@ -1,5 +1,6 @@ package tech.ydb.jdbc.query; +import tech.ydb.jdbc.YdbConst; import tech.ydb.jdbc.common.TypeDescription; /** @@ -18,7 +19,7 @@ public ParamDescription(String name, String displayName, TypeDescription type) { } public ParamDescription(String name, TypeDescription type) { - this(name, name, type); + this(name, YdbConst.VARIABLE_PARAMETER_PREFIX + name, type); } public String name() { diff --git a/jdbc/src/main/java/tech/ydb/jdbc/query/YdbPreparedQuery.java b/jdbc/src/main/java/tech/ydb/jdbc/query/YdbPreparedQuery.java index 6cc135f..7779962 100644 --- a/jdbc/src/main/java/tech/ydb/jdbc/query/YdbPreparedQuery.java +++ b/jdbc/src/main/java/tech/ydb/jdbc/query/YdbPreparedQuery.java @@ -14,6 +14,7 @@ */ public interface YdbPreparedQuery { String getQueryText(Params prms) throws SQLException; + String getBatchText(Params prms) throws SQLException; void clearParameters(); diff --git a/jdbc/src/main/java/tech/ydb/jdbc/query/params/BatchedQuery.java b/jdbc/src/main/java/tech/ydb/jdbc/query/params/BatchedQuery.java index b93ad43..97d84d5 100644 --- a/jdbc/src/main/java/tech/ydb/jdbc/query/params/BatchedQuery.java +++ b/jdbc/src/main/java/tech/ydb/jdbc/query/params/BatchedQuery.java @@ -4,7 +4,6 @@ import java.sql.SQLDataException; import java.sql.SQLException; import java.util.ArrayList; -import java.util.Arrays; import java.util.Collections; import java.util.HashMap; import java.util.HashSet; @@ -36,7 +35,8 @@ * @author Aleksandr Gorshenin */ public class BatchedQuery implements YdbPreparedQuery { - private final String yql; + private final String singleQuery; + private final String batchQuery; private final String batchParamName; private final Map paramsByName; private final ParamDescription[] params; @@ -44,28 +44,26 @@ public class BatchedQuery implements YdbPreparedQuery { private final List batchList = new ArrayList<>(); private final Map> currentValues = new HashMap<>(); - protected BatchedQuery(YdbTypes types, String yql, String listName, List pnames, Map ptypes) - throws SQLException { - this.yql = yql; - this.batchParamName = listName; + protected BatchedQuery(String single, String batched, String prm, ParamDescription[] params) throws SQLException { + this.singleQuery = single; + this.batchQuery = batched; + this.batchParamName = prm; this.paramsByName = new HashMap<>(); - this.params = new ParamDescription[pnames.size()]; + this.params = params; - for (int idx = 0; idx < pnames.size(); idx += 1) { - String name = pnames.get(idx); - if (!ptypes.containsKey(name)) { - throw new SQLException(YdbConst.INVALID_BATCH_COLUMN + name); - } - TypeDescription type = types.find(ptypes.get(name)); - ParamDescription desc = new ParamDescription(name, YdbConst.VARIABLE_PARAMETER_PREFIX + name, type); - params[idx] = desc; - paramsByName.put(name, desc); + for (ParamDescription pd: params) { + paramsByName.put(pd.name(), pd); } } @Override public String getQueryText(Params prms) { - return yql; + return singleQuery != null ? singleQuery : batchQuery; + } + + @Override + public String getBatchText(Params prms) { + return batchQuery; } @Override @@ -85,7 +83,7 @@ public void clearParameters() { @Override public void addBatch() throws SQLException { - batchList.add(getCurrentValues()); + batchList.add(StructValue.of(validateValues())); currentValues.clear(); } @@ -94,13 +92,13 @@ public void clearBatch() { batchList.clear(); } - protected StructValue getCurrentValues() throws SQLException { + protected Map> validateValues() throws SQLException { for (ParamDescription prm: params) { if (!currentValues.containsKey(prm.name())) { throw new SQLDataException(YdbConst.MISSING_VALUE_FOR_PARAMETER + prm.displayName()); } } - return StructValue.of(currentValues); + return currentValues; } protected List getBatchedValues() { @@ -109,8 +107,13 @@ protected List getBatchedValues() { @Override public Params getCurrentParams() throws SQLException { - ListValue list = ListValue.of(getCurrentValues()); - return Params.of(batchParamName, list); + Map> vv = validateValues(); + if (singleQuery == null) { + return Params.of(batchParamName, ListValue.of(StructValue.of(vv))); + } + Params prms = Params.create(vv.size()); + vv.forEach((name, value) -> prms.put(YdbConst.VARIABLE_PARAMETER_PREFIX + name, value)); + return prms; } @Override @@ -216,7 +219,15 @@ public static BatchedQuery tryCreateBatched(YdbTypes types, YdbQuery query, Map< columns[idx] = param; } - return new BatchedQuery(types, query.getPreparedYql(), listName, Arrays.asList(columns), paramTypes); + ParamDescription[] descriptions = new ParamDescription[columns.length]; + for (int idx = 0; idx < columns.length; idx += 1) { + String name = columns[idx]; + if (!paramTypes.containsKey(name)) { + throw new SQLException(YdbConst.INVALID_BATCH_COLUMN + name); + } + descriptions[idx] = new ParamDescription(name, types.find(paramTypes.get(name))); + } + return new BatchedQuery(null, query.getPreparedYql(), listName, descriptions); } public static BatchedQuery createAutoBatched(YdbTypes types, YqlBatcher batcher, TableDescription description) @@ -235,33 +246,38 @@ public static BatchedQuery createAutoBatched(YdbTypes types, YqlBatcher batcher, } } - StringBuilder sb = new StringBuilder(); Map columnTypes = new HashMap<>(); - Map structTypes = new HashMap<>(); - List columns = new ArrayList<>(); - for (TableColumn column: description.getColumns()) { columnTypes.put(column.getName(), column.getType()); } - List params = new ArrayList<>(); - params.addAll(batcher.getColumns()); - params.addAll(batcher.getKeyColumns()); + List columns = new ArrayList<>(); + columns.addAll(batcher.getColumns()); + columns.addAll(batcher.getKeyColumns()); + + ParamDescription[] params = new ParamDescription[columns.size()]; - sb.append("DECLARE $batch AS List 1) { + params[idx - 1] = new ParamDescription("p" + idx, column, types.find(type)); + idx++; + } + + return new BatchedQuery(simpleQuery(batcher, params), batchQuery(batcher, params), "$batch", params); + } + + private static String batchQuery(YqlBatcher batcher, ParamDescription[] params) { + StringBuilder sb = new StringBuilder(); + sb.append("DECLARE $batch AS List 0) { sb.append(", "); } - sb.append("p").append(idx).append(":").append(type.toString()); - structTypes.put("p" + idx, type); - columns.add("p" + idx); - idx++; + sb.append(params[idx].name()).append(":").append(params[idx].type().toYqlLiteral()); } sb.append(">>;\n"); @@ -282,20 +298,97 @@ public static BatchedQuery createAutoBatched(YdbTypes types, YqlBatcher batcher, sb.append("DELETE FROM `").append(batcher.getTableName()).append("` ON SELECT "); break; default: - return null; + return "UNSUPPORTED CMD " + batcher.getCommand(); } - idx = 1; - for (String column: params) { - if (idx > 1) { + for (int idx = 0; idx < params.length; idx++) { + if (idx > 0) { sb.append(", "); } - sb.append("p").append(idx).append(" AS `").append(column).append("`"); - idx++; + sb.append(params[idx].name()).append(" AS `").append(params[idx].displayName()).append("`"); } sb.append(" FROM AS_TABLE($batch);"); + return sb.toString(); + } + + private static String simpleQuery(YqlBatcher batcher, ParamDescription[] params) { + StringBuilder sb = new StringBuilder(); + for (ParamDescription p : params) { + sb.append("DECLARE ").append(YdbConst.VARIABLE_PARAMETER_PREFIX).append(p.name()) + .append(" AS ").append(p.type().toYqlLiteral()).append(";\n"); + } - return new BatchedQuery(types, sb.toString(), "$batch", columns, structTypes); + switch (batcher.getCommand()) { + case UPSERT: + sb.append("UPSERT INTO `").append(batcher.getTableName()).append("` ("); + appendColumns(sb, params); + sb.append(") VALUES ("); + appendValues(sb, params); + sb.append(");"); + break; + case INSERT: + sb.append("INSERT INTO `").append(batcher.getTableName()).append("` ("); + appendColumns(sb, params); + sb.append(") VALUES ("); + appendValues(sb, params); + sb.append(");"); + break; + case REPLACE: + sb.append("REPLACE INTO `").append(batcher.getTableName()).append("` ("); + appendColumns(sb, params); + sb.append(") VALUES ("); + appendValues(sb, params); + sb.append(");"); + break; + case UPDATE: + sb.append("UPDATE `").append(batcher.getTableName()).append("` SET "); + for (int idx = 0; idx < batcher.getColumns().size(); idx++) { + if (idx > 0) { + sb.append(", "); + } + sb.append('`').append(params[idx].displayName()).append("` = ") + .append(YdbConst.VARIABLE_PARAMETER_PREFIX).append(params[idx].name()); + } + sb.append(" WHERE "); + appendKeys(sb, params, batcher.getColumns().size()); + break; + case DELETE: + sb.append("DELETE FROM `").append(batcher.getTableName()).append("` WHERE "); + appendKeys(sb, params, batcher.getColumns().size()); + break; + default: + break; + } + + return sb.toString(); + } + + private static void appendColumns(StringBuilder sb, ParamDescription[] params) { + for (int idx = 0; idx < params.length; idx++) { + if (idx > 0) { + sb.append(", "); + } + sb.append('`').append(params[idx].displayName()).append('`'); + } + } + + private static void appendValues(StringBuilder sb, ParamDescription[] params) { + for (int idx = 0; idx < params.length; idx++) { + if (idx > 0) { + sb.append(", "); + } + sb.append(YdbConst.VARIABLE_PARAMETER_PREFIX).append(params[idx].name()); + } + } + + private static void appendKeys(StringBuilder sb, ParamDescription[] params, int firstKeyIdx) { + for (int idx = firstKeyIdx; idx < params.length; idx++) { + if (idx > firstKeyIdx) { + sb.append(" AND "); + } + sb.append('`').append(params[idx].displayName()).append("` = ") + .append(YdbConst.VARIABLE_PARAMETER_PREFIX).append(params[idx].name()); + } } } diff --git a/jdbc/src/main/java/tech/ydb/jdbc/query/params/BulkUpsertQuery.java b/jdbc/src/main/java/tech/ydb/jdbc/query/params/BulkUpsertQuery.java index c370513..ecdbc7d 100644 --- a/jdbc/src/main/java/tech/ydb/jdbc/query/params/BulkUpsertQuery.java +++ b/jdbc/src/main/java/tech/ydb/jdbc/query/params/BulkUpsertQuery.java @@ -8,11 +8,13 @@ import java.util.stream.Collectors; import tech.ydb.jdbc.common.YdbTypes; +import tech.ydb.jdbc.query.ParamDescription; import tech.ydb.table.description.TableColumn; import tech.ydb.table.description.TableDescription; import tech.ydb.table.values.ListType; import tech.ydb.table.values.ListValue; import tech.ydb.table.values.StructType; +import tech.ydb.table.values.StructValue; import tech.ydb.table.values.Type; /** @@ -23,11 +25,10 @@ public class BulkUpsertQuery extends BatchedQuery { private final String tablePath; private final ListType bulkType; - private BulkUpsertQuery(YdbTypes types, String path, String yql, List columns, Map paramTypes) - throws SQLException { - super(types, yql, "$bulk", columns, paramTypes); - this.tablePath = path; - this.bulkType = ListType.of(StructType.of(paramTypes)); + private BulkUpsertQuery(String tablePath, String yql, ListType tp, ParamDescription[] params) throws SQLException { + super(null, yql, "$bulk", params); + this.tablePath = tablePath; + this.bulkType = tp; } public String getTablePath() { @@ -35,7 +36,7 @@ public String getTablePath() { } public ListValue getCurrentBulk() throws SQLException { - return bulkType.newValue(Collections.singletonList(getCurrentValues())); + return bulkType.newValue(Collections.singletonList(StructValue.of(validateValues()))); } public ListValue getBatchedBulk() { @@ -56,14 +57,18 @@ public static BulkUpsertQuery build(YdbTypes types, String path, List co columnTypes.put(column.getName(), column.getType()); } + ParamDescription[] params = new ParamDescription[columns.size()]; Map structTypes = new HashMap<>(); + int idx = 0; for (String column: columns) { if (!columnTypes.containsKey(column)) { throw new SQLException("Cannot parse BULK upsert: column " + column + " not found"); } - structTypes.put(column, columnTypes.get(column)); + Type type = columnTypes.get(column); + structTypes.put(column, type); + params[idx++] = new ParamDescription(column, types.find(type)); } - return new BulkUpsertQuery(types, path, yql.toString(), columns, structTypes); + return new BulkUpsertQuery(path, yql.toString(), ListType.of(StructType.of(structTypes)), params); } } diff --git a/jdbc/src/main/java/tech/ydb/jdbc/query/params/InMemoryQuery.java b/jdbc/src/main/java/tech/ydb/jdbc/query/params/InMemoryQuery.java index fca4cba..b4356df 100644 --- a/jdbc/src/main/java/tech/ydb/jdbc/query/params/InMemoryQuery.java +++ b/jdbc/src/main/java/tech/ydb/jdbc/query/params/InMemoryQuery.java @@ -61,6 +61,11 @@ public String getQueryText(Params prms) throws SQLException { return query.toString(); } + @Override + public String getBatchText(Params prms) throws SQLException { + return getQueryText(prms); + } + @Override public int parametersCount() { return parameters.size(); diff --git a/jdbc/src/main/java/tech/ydb/jdbc/query/params/PreparedQuery.java b/jdbc/src/main/java/tech/ydb/jdbc/query/params/PreparedQuery.java index 371411e..9737301 100644 --- a/jdbc/src/main/java/tech/ydb/jdbc/query/params/PreparedQuery.java +++ b/jdbc/src/main/java/tech/ydb/jdbc/query/params/PreparedQuery.java @@ -78,6 +78,11 @@ public String getQueryText(Params prms) { return yql; } + @Override + public String getBatchText(Params prms) { + return yql; + } + @Override public void setParam(int index, Object obj, int sqlType) throws SQLException { if (index <= 0 || index > paramNames.length) { diff --git a/jdbc/src/test/java/tech/ydb/jdbc/YdbDriverExampleTest.java b/jdbc/src/test/java/tech/ydb/jdbc/YdbDriverExampleTest.java index 87d9a10..2421a8b 100644 --- a/jdbc/src/test/java/tech/ydb/jdbc/YdbDriverExampleTest.java +++ b/jdbc/src/test/java/tech/ydb/jdbc/YdbDriverExampleTest.java @@ -124,11 +124,25 @@ public void testYdb() throws SQLException { psBatch.executeBatch(); } + try (PreparedStatement bulkPs = connection.prepareStatement("" + + "BULK UPSERT INTO table_sample (id, value) VALUES (?, ?)")) { + + bulkPs.setInt(1, 12); + bulkPs.setString(2, "value-12"); + bulkPs.addBatch(); + + bulkPs.setInt(1, 13); + bulkPs.setString(2, "value-13"); + bulkPs.addBatch(); + + bulkPs.executeBatch(); + } + try (PreparedStatement select = connection .prepareStatement("select count(1) as cnt from table_sample")) { ResultSet rs = select.executeQuery(); rs.next(); - Assertions.assertEquals(11, rs.getLong("cnt")); + Assertions.assertEquals(13, rs.getLong("cnt")); } } } @@ -212,11 +226,25 @@ public void testYdbNotNull() throws SQLException { psBatch.executeBatch(); } + try (PreparedStatement bulkPs = connection.prepareStatement("" + + "BULK UPSERT INTO table_sample (id, value) VALUES (?, ?)")) { + + bulkPs.setInt(1, 9); + bulkPs.setString(2, "value-9"); + bulkPs.addBatch(); + + bulkPs.setInt(1, 10); + bulkPs.setString(2, "value-10"); + bulkPs.addBatch(); + + bulkPs.executeBatch(); + } + try (PreparedStatement select = connection .prepareStatement("select count(1) as cnt from table_sample")) { ResultSet rs = select.executeQuery(); rs.next(); - Assertions.assertEquals(8, rs.getLong("cnt")); + Assertions.assertEquals(10, rs.getLong("cnt")); } } } diff --git a/jdbc/src/test/java/tech/ydb/jdbc/impl/YdbConnectionImplTest.java b/jdbc/src/test/java/tech/ydb/jdbc/impl/YdbConnectionImplTest.java index fe67b21..671e670 100644 --- a/jdbc/src/test/java/tech/ydb/jdbc/impl/YdbConnectionImplTest.java +++ b/jdbc/src/test/java/tech/ydb/jdbc/impl/YdbConnectionImplTest.java @@ -165,20 +165,22 @@ public void prepareStatement() throws SQLException { ResultSet.TYPE_FORWARD_ONLY, ResultSet.CONCUR_READ_ONLY, ResultSet.HOLD_CURSORS_OVER_COMMIT)) { TableAssert.assertSelectInt(2, statement.executeQuery()); } + try (PreparedStatement statement = jdbc.connection().prepareStatement("select 1 + 1", new int[0])) { + TableAssert.assertSelectInt(2, statement.executeQuery()); + } + try (PreparedStatement statement = jdbc.connection().prepareStatement("select 1 + 1", new String[0])) { + TableAssert.assertSelectInt(2, statement.executeQuery()); + } + try (PreparedStatement statement = jdbc.connection().prepareStatement("select 1 + 1", + Statement.RETURN_GENERATED_KEYS)) { + TableAssert.assertSelectInt(2, statement.executeQuery()); + } } @Test public void prepareStatementInvalid() throws SQLException { ExceptionAssert.sqlFeatureNotSupported("Auto-generated keys are not supported", - () -> jdbc.connection().prepareStatement(SELECT_2_2, new int[] {}) - ); - - ExceptionAssert.sqlFeatureNotSupported("Auto-generated keys are not supported", - () -> jdbc.connection().prepareStatement(SELECT_2_2, new String[] {}) - ); - - ExceptionAssert.sqlFeatureNotSupported("Auto-generated keys are not supported", - () -> jdbc.connection().prepareStatement(SELECT_2_2, Statement.RETURN_GENERATED_KEYS) + () -> jdbc.connection().prepareStatement(SELECT_2_2, new int[] {1, 2}) ); ExceptionAssert.sqlFeatureNotSupported( diff --git a/jdbc/src/test/java/tech/ydb/jdbc/impl/YdbPreparedStatementTest.java b/jdbc/src/test/java/tech/ydb/jdbc/impl/YdbPreparedStatementTest.java index ca8b219..46a0c8c 100644 --- a/jdbc/src/test/java/tech/ydb/jdbc/impl/YdbPreparedStatementTest.java +++ b/jdbc/src/test/java/tech/ydb/jdbc/impl/YdbPreparedStatementTest.java @@ -137,7 +137,7 @@ public void executeWithMissingParameter() throws SQLException { try (PreparedStatement ps = jdbc.connection().prepareStatement(upsert.apply(SqlQueries.JdbcQuery.STANDARD))) { ps.setInt(1, 1); - ExceptionAssert.sqlDataException("Missing value for parameter: $p2", ps::execute); + ExceptionAssert.sqlDataException("Missing value for parameter: c_Text", ps::execute); } try (PreparedStatement ps = jdbc.connection().prepareStatement(upsert.apply(SqlQueries.JdbcQuery.IN_MEMORY))) { diff --git a/jdbc/src/test/java/tech/ydb/jdbc/impl/YdbTableConnectionImplTest.java b/jdbc/src/test/java/tech/ydb/jdbc/impl/YdbTableConnectionImplTest.java index e3ddb5e..7a8d50e 100644 --- a/jdbc/src/test/java/tech/ydb/jdbc/impl/YdbTableConnectionImplTest.java +++ b/jdbc/src/test/java/tech/ydb/jdbc/impl/YdbTableConnectionImplTest.java @@ -166,20 +166,22 @@ public void prepareStatement() throws SQLException { ResultSet.TYPE_FORWARD_ONLY, ResultSet.CONCUR_READ_ONLY, ResultSet.HOLD_CURSORS_OVER_COMMIT)) { TableAssert.assertSelectInt(2, statement.executeQuery()); } + try (PreparedStatement statement = jdbc.connection().prepareStatement("select 1 + 1", new int[0])) { + TableAssert.assertSelectInt(2, statement.executeQuery()); + } + try (PreparedStatement statement = jdbc.connection().prepareStatement("select 1 + 1", new String[0])) { + TableAssert.assertSelectInt(2, statement.executeQuery()); + } + try (PreparedStatement statement = jdbc.connection().prepareStatement("select 1 + 1", + Statement.RETURN_GENERATED_KEYS)) { + TableAssert.assertSelectInt(2, statement.executeQuery()); + } } @Test public void prepareStatementInvalid() throws SQLException { ExceptionAssert.sqlFeatureNotSupported("Auto-generated keys are not supported", - () -> jdbc.connection().prepareStatement(SELECT_2_2, new int[] {}) - ); - - ExceptionAssert.sqlFeatureNotSupported("Auto-generated keys are not supported", - () -> jdbc.connection().prepareStatement(SELECT_2_2, new String[] {}) - ); - - ExceptionAssert.sqlFeatureNotSupported("Auto-generated keys are not supported", - () -> jdbc.connection().prepareStatement(SELECT_2_2, Statement.RETURN_GENERATED_KEYS) + () -> jdbc.connection().prepareStatement(SELECT_2_2, new int[] {1, 2}) ); ExceptionAssert.sqlFeatureNotSupported(