Skip to content

Commit 7691fd0

Browse files
hantmachantmac
authored
Merge pull request #319 from databendlabs/fix/executeUpdate-count
fix: the return count of executeUpdate
2 parents c78ca12 + 74daec1 commit 7691fd0

File tree

5 files changed

+191
-84
lines changed

5 files changed

+191
-84
lines changed

‎databend-client/src/main/java/com/databend/client/QueryStats.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -68,7 +68,7 @@ public QueryProgress getResultProgress() {
6868
public String toString() {
6969
return toStringHelper(this)
7070
.add("runningTimeMS", runningTimeMS)
71-
.add("scamProgress", scanProgress)
71+
.add("scanProgress", scanProgress)
7272
.add("writeProgress", writeProgress)
7373
.add("readProgress", resultProgress)
7474
.toString();

‎databend-jdbc/src/main/java/com/databend/jdbc/DatabendPreparedStatement.java‎

Lines changed: 44 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -73,24 +73,23 @@ public class DatabendPreparedStatement extends DatabendStatement implements Prep
7373
static final DateTimeFormatter TIME_FORMATTER = DateTimeFormat.forPattern("HH:mm:ss.SSS");
7474
static final DateTimeFormatter TIMESTAMP_FORMATTER = DateTimeFormat.forPattern("yyyy-MM-dd HH:mm:ss.SSS");
7575
private final DatabendParameterMetaData paramMetaData;
76-
private static final java.time.format.DateTimeFormatter LOCAL_DATE_TIME_FORMATTER =
77-
new DateTimeFormatterBuilder()
78-
.append(ISO_LOCAL_DATE)
79-
.appendLiteral(' ')
80-
.append(ISO_LOCAL_TIME)
81-
.toFormatter();
82-
private static final java.time.format.DateTimeFormatter OFFSET_TIME_FORMATTER =
83-
new DateTimeFormatterBuilder()
84-
.append(ISO_LOCAL_TIME)
85-
.appendOffset("+HH:mm", "+00:00")
86-
.toFormatter();
76+
private static final java.time.format.DateTimeFormatter LOCAL_DATE_TIME_FORMATTER = new DateTimeFormatterBuilder()
77+
.append(ISO_LOCAL_DATE)
78+
.appendLiteral(' ')
79+
.append(ISO_LOCAL_TIME)
80+
.toFormatter();
81+
private static final java.time.format.DateTimeFormatter OFFSET_TIME_FORMATTER = new DateTimeFormatterBuilder()
82+
.append(ISO_LOCAL_TIME)
83+
.appendOffset("+HH:mm", "+00:00")
84+
.toFormatter();
8785
private final String originalSql;
8886
private final List<String[]> batchValues;
8987
private final Optional<BatchInsertUtils> batchInsertUtils;
9088
private final String statementName;
9189
private int batchSize = 0;
9290

93-
DatabendPreparedStatement(DatabendConnection connection, Consumer<DatabendStatement> onClose, String statementName, String sql) {
91+
DatabendPreparedStatement(DatabendConnection connection, Consumer<DatabendStatement> onClose, String statementName,
92+
String sql) {
9493
super(connection, onClose);
9594
this.statementName = requireNonNull(statementName, "statementName is null");
9695
this.originalSql = requireNonNull(sql, "sql is null");
@@ -142,13 +141,13 @@ private static String formatBigDecimalLiteral(BigDecimal x) {
142141
return x.toString();
143142
}
144143

145-
146144
private static String formatBytesLiteral(byte[] x) {
147145
return new String(x, StandardCharsets.UTF_8);
148146
}
149147

150148
static IllegalArgumentException invalidConversion(Object x, String toType) {
151-
return new IllegalArgumentException(format("Cannot convert instance of %s to %s", x.getClass().getName(), toType));
149+
return new IllegalArgumentException(
150+
format("Cannot convert instance of %s to %s", x.getClass().getName(), toType));
152151
}
153152

154153
@Override
@@ -183,7 +182,9 @@ private DatabendCopyParams uploadBatchesForCopyInto() throws SQLException {
183182
DatabendStage databendStage = DatabendStage.builder().stageName(stageName).path(stagePrefix).build();
184183
List<String> files = new ArrayList<>();
185184
files.add(fileName);
186-
DatabendCopyParams databendCopyParams = DatabendCopyParams.builder().setFiles(files).setCopyOptions(copyOptions).setDatabaseTableName(batchInsertUtils.get().getDatabaseTableName()).setDatabendStage(databendStage).build();
185+
DatabendCopyParams databendCopyParams = DatabendCopyParams.builder().setFiles(files)
186+
.setCopyOptions(copyOptions).setDatabaseTableName(batchInsertUtils.get().getDatabaseTableName())
187+
.setDatabendStage(databendStage).build();
187188
return databendCopyParams;
188189
} catch (Exception e) {
189190
throw new SQLException(e);
@@ -235,11 +236,14 @@ private StageAttachment uploadBatches() throws SQLException {
235236
}
236237

237238
/**
238-
* This method is used to build a StageAttachment object which represents a stage in Databend.
239-
* A stage in Databend is a temporary storage area where data files are stored before being loaded into the Databend database.
239+
* This method is used to build a StageAttachment object which represents a
240+
* stage in Databend.
241+
* A stage in Databend is a temporary storage area where data files are stored
242+
* before being loaded into the Databend database.
240243
*
241-
* @param connection The DatabendConnection object which contains the connection details to the Databend database.
242-
* @param stagePath The path of the stage in the Databend database.
244+
* @param connection The DatabendConnection object which contains the connection
245+
* details to the Databend database.
246+
* @param stagePath The path of the stage in the Databend database.
243247
* @return A StageAttachment object which contains the details of the stage.
244248
*/
245249
public static StageAttachment buildStateAttachment(DatabendConnection connection, String stagePath) {
@@ -291,18 +295,19 @@ private boolean dropStageAttachment(StageAttachment attachment) {
291295
public int[] executeBatchByAttachment() throws SQLException {
292296
int[] batchUpdateCounts = new int[batchValues.size()];
293297
if (!batchInsertUtils.isPresent() || batchValues == null || batchValues.isEmpty()) {
294-
// super.execute(this.originalSql);
298+
// super.execute(this.originalSql);
295299
return batchUpdateCounts;
296300
}
297301
StageAttachment attachment = uploadBatches();
298302
ResultSet r = null;
299303
if (attachment == null) {
300-
// logger.fine("use normal execute instead of batch insert");
301-
// super.execute(batchInsertUtils.get().getSql());
304+
// logger.fine("use normal execute instead of batch insert");
305+
// super.execute(batchInsertUtils.get().getSql());
302306
return batchUpdateCounts;
303307
}
304308
try {
305-
logger.fine(String.format("use batch insert instead of normal insert, attachment: %s, sql: %s", attachment, batchInsertUtils.get().getSql()));
309+
logger.fine(String.format("use batch insert instead of normal insert, attachment: %s, sql: %s", attachment,
310+
batchInsertUtils.get().getSql()));
306311
super.internalExecute(batchInsertUtils.get().getSql(), attachment);
307312
r = getResultSet();
308313
while (r.next()) {
@@ -347,7 +352,7 @@ public int[] executeBatchByCopyInto() throws SQLException {
347352

348353
public int[] executeBatchDelete() throws SQLException {
349354
if (!batchInsertUtils.isPresent() || batchValues == null || batchValues.isEmpty()) {
350-
return new int[]{};
355+
return new int[] {};
351356
}
352357
int[] batchUpdateCounts = new int[batchValues.size()];
353358
try {
@@ -382,7 +387,6 @@ public static String convertSQLWithBatchValues(String baseSql, List<String[]> ba
382387
return convertedSqlBuilder.toString();
383388
}
384389

385-
386390
@Override
387391
public int[] executeBatch() throws SQLException {
388392
if (originalSql.toLowerCase().contains("delete from")) {
@@ -449,7 +453,7 @@ protected void handleBatchInsert() throws SQLException {
449453
@Override
450454
public int executeUpdate() throws SQLException {
451455
this.execute(prepareSQL(batchInsertUtils.get().getProvideParams()));
452-
return batchInsertUtils.get().getProvideParams().size();
456+
return getUpdateCount();
453457
}
454458

455459
@Override
@@ -519,7 +523,8 @@ public void setDouble(int i, double v)
519523
public void setBigDecimal(int i, BigDecimal bigDecimal)
520524
throws SQLException {
521525
checkOpen();
522-
batchInsertUtils.ifPresent(insertUtils -> insertUtils.setPlaceHolderValue(i, formatBigDecimalLiteral(bigDecimal)));
526+
batchInsertUtils
527+
.ifPresent(insertUtils -> insertUtils.setPlaceHolderValue(i, formatBigDecimalLiteral(bigDecimal)));
523528
}
524529

525530
@Override
@@ -535,7 +540,8 @@ public void setString(int i, String s)
535540
s = s.replace("'", "\\\'");
536541
}
537542
String finalS = s;
538-
batchInsertUtils.ifPresent(insertUtils -> insertUtils.setPlaceHolderValue(i, String.format("%s%s%s", "'", finalS, "'")));
543+
batchInsertUtils.ifPresent(
544+
insertUtils -> insertUtils.setPlaceHolderValue(i, String.format("%s%s%s", "'", finalS, "'")));
539545
}
540546
}
541547

@@ -554,7 +560,8 @@ public void setDate(int i, Date date)
554560
batchInsertUtils.ifPresent(insertUtils -> insertUtils.setPlaceHolderValue(i, null));
555561
} else {
556562
if (originalSql.toLowerCase().startsWith("select")) {
557-
batchInsertUtils.ifPresent(insertUtils -> insertUtils.setPlaceHolderValue(i, String.format("%s%s%s", "'", date, "'")));
563+
batchInsertUtils.ifPresent(
564+
insertUtils -> insertUtils.setPlaceHolderValue(i, String.format("%s%s%s", "'", date, "'")));
558565
} else {
559566
batchInsertUtils.ifPresent(insertUtils -> insertUtils.setPlaceHolderValue(i, toDateLiteral(date)));
560567
}
@@ -569,7 +576,8 @@ public void setTime(int i, Time time)
569576
batchInsertUtils.ifPresent(insertUtils -> insertUtils.setPlaceHolderValue(i, null));
570577
} else {
571578
if (originalSql.toLowerCase().startsWith("select")) {
572-
batchInsertUtils.ifPresent(insertUtils -> insertUtils.setPlaceHolderValue(i, String.format("%s%s%s", "'", time, "'")));
579+
batchInsertUtils.ifPresent(
580+
insertUtils -> insertUtils.setPlaceHolderValue(i, String.format("%s%s%s", "'", time, "'")));
573581
} else {
574582
batchInsertUtils.ifPresent(insertUtils -> insertUtils.setPlaceHolderValue(i, toTimeLiteral(time)));
575583
}
@@ -584,9 +592,11 @@ public void setTimestamp(int i, Timestamp timestamp)
584592
batchInsertUtils.ifPresent(insertUtils -> insertUtils.setPlaceHolderValue(i, null));
585593
} else {
586594
if (originalSql.toLowerCase().startsWith("select")) {
587-
batchInsertUtils.ifPresent(insertUtils -> insertUtils.setPlaceHolderValue(i, String.format("%s%s%s", "'", timestamp, "'")));
595+
batchInsertUtils.ifPresent(insertUtils -> insertUtils.setPlaceHolderValue(i,
596+
String.format("%s%s%s", "'", timestamp, "'")));
588597
} else {
589-
batchInsertUtils.ifPresent(insertUtils -> insertUtils.setPlaceHolderValue(i, toTimestampLiteral(timestamp)));
598+
batchInsertUtils
599+
.ifPresent(insertUtils -> insertUtils.setPlaceHolderValue(i, toTimestampLiteral(timestamp)));
590600
}
591601
}
592602
}
@@ -759,7 +769,6 @@ public static String convertArrayListToString(ArrayList<?> arrayList) {
759769
return builder.toString();
760770
}
761771

762-
763772
@Override
764773
public void addBatch()
765774
throws SQLException {
@@ -854,7 +863,8 @@ public void setURL(int i, URL url)
854863
throw new SQLFeatureNotSupportedException("PreparedStatement", "setURL");
855864
}
856865

857-
// If you want to use ps.getParameterMetaData().* methods, you need to use a valid sql such as
866+
// If you want to use ps.getParameterMetaData().* methods, you need to use a
867+
// valid sql such as
858868
// insert into table_name (col1 type1, col2 typ2, col3 type3) values (?, ?, ?)
859869
@Override
860870
public ParameterMetaData getParameterMetaData() throws SQLException {
@@ -1006,7 +1016,6 @@ public void setNClob(int i, Reader reader)
10061016
throw new SQLFeatureNotSupportedException("PreparedStatement", "setNClob");
10071017
}
10081018

1009-
10101019
private String toDateLiteral(Object value) throws IllegalArgumentException {
10111020
requireNonNull(value, "value is null");
10121021
if (value instanceof java.util.Date) {

‎databend-jdbc/src/main/java/com/databend/jdbc/DatabendStatement.java‎

Lines changed: 27 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -185,16 +185,10 @@ final boolean internalExecute(String sql, StageAttachment attachment) throws SQL
185185
throw resultsException(client.getResults(), sql);
186186
}
187187
}
188-
if (isQueryStatement(sql)) {
189-
currentUpdateCount = -1;// Always -1 when returning a ResultSet with query statement
190-
} else {
191-
currentUpdateCount = client.getResults().getStats().getScanProgress().getRows().intValue();
192-
}
193188
executingClient.set(client);
194189
while (client.hasNext()) {
195190
QueryResults results = client.getResults();
196191
List<List<Object>> data = results.getData();
197-
// List<QueryRowField> schema = results.getSchema();
198192
if (data == null || data.isEmpty()) {
199193
client.advance();
200194
} else {
@@ -203,9 +197,34 @@ final boolean internalExecute(String sql, StageAttachment attachment) throws SQL
203197
}
204198
resultSet = DatabendResultSet.create(this, client, maxRows.get());
205199
currentResult.set(resultSet);
200+
if (isQueryStatement(sql)) {
201+
currentUpdateCount = -1;// Always -1 when returning a ResultSet with query statement
202+
} else {
203+
QueryResults results = client.getResults();
204+
if (sql.toLowerCase().startsWith("update") || sql.toLowerCase().startsWith("delete")) {
205+
List<List<Object>> data = results.getData();
206+
if (data != null && !data.isEmpty() && data.get(0) != null && !data.get(0).isEmpty()) {
207+
Object updateCount = data.get(0).get(0);
208+
if (updateCount instanceof Number) {
209+
currentUpdateCount = ((Number) updateCount).intValue();
210+
} else {
211+
// if can't find, use writeProgress.rows
212+
currentUpdateCount = results.getStats().getWriteProgress().getRows().intValue();
213+
}
214+
} else {
215+
// if data is empty, use writeProgress.rows
216+
currentUpdateCount = results.getStats().getWriteProgress().getRows().intValue();
217+
}
218+
} else {
219+
System.out.println("sql is : " + sql);
220+
System.out.println("[DEBUG] Query Write Progress: " + results.getStats().getWriteProgress());
221+
currentUpdateCount = results.getStats().getWriteProgress().getRows().intValue();
222+
}
223+
}
206224
return true;
207225
} catch (RuntimeException e) {
208-
throw new SQLException("Error executing query: " + "SQL: " + sql + " " + e.getMessage() + " cause: " + e.getCause(), e);
226+
throw new SQLException(
227+
"Error executing query: " + "SQL: " + sql + " " + e.getMessage() + " cause: " + e.getCause(), e);
209228
} finally {
210229
executingClient.set(null);
211230
if (currentResult.get() == null) {
@@ -436,7 +455,7 @@ protected final DatabendConnection connection()
436455
}
437456

438457
public QueryLiveness queryLiveness() {
439-
DatabendResultSet r = currentResult.get();
458+
DatabendResultSet r = currentResult.get();
440459

441460
if (r != null) {
442461
return r.getLiveness();
@@ -448,5 +467,3 @@ protected final Optional<DatabendConnection> optionalConnection() {
448467
return Optional.ofNullable(connection.get());
449468
}
450469
}
451-
452-

‎databend-jdbc/src/test/java/com/databend/jdbc/TestBasicDriver.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -158,7 +158,7 @@ public void TestMergeinto() throws SQLException {
158158
" INSERT *;\n");
159159
ResultSet r = statement.getResultSet();
160160
r.next();
161-
Assert.assertEquals(6, statement.getUpdateCount());
161+
Assert.assertEquals(3, statement.getUpdateCount());
162162
System.out.println(statement.getUpdateCount());
163163
} catch (SQLException throwables) {
164164
throwables.printStackTrace();

0 commit comments

Comments
 (0)