From 7155c587ac31cff6b61d5612a74c80cdaf570cec Mon Sep 17 00:00:00 2001 From: Yang Xiufeng Date: Thu, 11 Sep 2025 00:26:53 +0800 Subject: [PATCH 1/5] fix: PrepareStatement 1. execute() do not affect added batches. 2. fix delete and insert updateCount. 3. fix executeBatch, use loop. 4. forbid multi statement. --- .gitignore | 1 + .../com/databend/jdbc/DatabendConnection.java | 13 +- .../jdbc/DatabendPreparedStatement.java | 354 +++++------------- .../com/databend/jdbc/DatabendStatement.java | 10 +- .../java/com/databend/jdbc/StatementUtil.java | 24 +- .../jdbc/constant/DatabendConstant.java | 5 +- .../jdbc/parser/BatchInsertUtils.java | 42 +-- .../databend/jdbc/TestPrepareStatement.java | 183 +++++---- .../test/java/com/databend/jdbc/Utils.java | 12 +- .../jdbc/parser/TestBatchInsertUtils.java | 12 +- 10 files changed, 249 insertions(+), 407 deletions(-) diff --git a/.gitignore b/.gitignore index 676fbe0b..bee2073d 100644 --- a/.gitignore +++ b/.gitignore @@ -9,3 +9,4 @@ target/ tests/data tests/compatibility/*.jar test-output +.DS_Store diff --git a/databend-jdbc/src/main/java/com/databend/jdbc/DatabendConnection.java b/databend-jdbc/src/main/java/com/databend/jdbc/DatabendConnection.java index f4ad77a7..1e60854c 100644 --- a/databend-jdbc/src/main/java/com/databend/jdbc/DatabendConnection.java +++ b/databend-jdbc/src/main/java/com/databend/jdbc/DatabendConnection.java @@ -198,7 +198,6 @@ public static URI parseRouteHint(String routeHint) { if (routeHint == null || routeHint.isEmpty()) { return null; } - URI target; try { if (routeHint.charAt(routeHint.length() - 1) != SPECIAL_CHAR) { return null; @@ -228,14 +227,6 @@ private static void checkResultSet(int resultSetType, int resultSetConcurrency) } } - // Databend DOES NOT support transaction now - private static void checkHoldability(int resultSetHoldability) - throws SQLFeatureNotSupportedException { - if (resultSetHoldability != ResultSet.HOLD_CURSORS_OVER_COMMIT) { - throw new SQLFeatureNotSupportedException("Result set holdability must be HOLD_CURSORS_OVER_COMMIT"); - } - } - public static String getCopyIntoSql(String database, DatabendCopyParams params) { StringBuilder sb = new StringBuilder(); sb.append("COPY INTO "); @@ -297,7 +288,7 @@ synchronized private void unregisterStatement(DatabendStatement statement) { public PreparedStatement prepareStatement(String s) throws SQLException { - return this.prepareStatement(s, 0, 0); + return this.prepareStatement(s, ResultSet.TYPE_FORWARD_ONLY, ResultSet.CONCUR_READ_ONLY); } @Override @@ -435,7 +426,7 @@ public Statement createStatement(int resultSetType, int resultSetConcurrency) @Override public PreparedStatement prepareStatement(String s, int i, int i1) throws SQLException { - DatabendPreparedStatement statement = new DatabendPreparedStatement(this, this::unregisterStatement, "test", s); + DatabendPreparedStatement statement = new DatabendPreparedStatement(this, this::unregisterStatement, s); registerStatement(statement); return statement; } diff --git a/databend-jdbc/src/main/java/com/databend/jdbc/DatabendPreparedStatement.java b/databend-jdbc/src/main/java/com/databend/jdbc/DatabendPreparedStatement.java index 2de43722..e093eaf5 100644 --- a/databend-jdbc/src/main/java/com/databend/jdbc/DatabendPreparedStatement.java +++ b/databend-jdbc/src/main/java/com/databend/jdbc/DatabendPreparedStatement.java @@ -2,12 +2,9 @@ import com.databend.client.StageAttachment; import com.databend.client.data.DatabendRawType; -import com.databend.jdbc.cloud.DatabendCopyParams; -import com.databend.jdbc.cloud.DatabendStage; import com.databend.jdbc.parser.BatchInsertUtils; import com.fasterxml.jackson.core.JsonProcessingException; import com.fasterxml.jackson.databind.ObjectMapper; -import lombok.NonNull; import org.joda.time.format.DateTimeFormat; import org.joda.time.format.DateTimeFormatter; import org.joda.time.format.ISODateTimeFormat; @@ -54,10 +51,10 @@ import java.util.List; import java.util.Map; import java.util.Objects; -import java.util.Optional; import java.util.UUID; import java.util.function.Consumer; import java.util.logging.Logger; +import java.util.regex.Matcher; import java.util.stream.Collectors; import static com.databend.jdbc.ObjectCasts.*; @@ -84,22 +81,21 @@ public class DatabendPreparedStatement extends DatabendStatement implements Prep .append(ISO_LOCAL_TIME) .appendOffset("+HH:mm", "+00:00") .toFormatter(); - private final String originalSql; private final List batchValues; - private final Optional batchInsertUtils; - private final String statementName; - private int batchSize = 0; + private final List batchValuesCSV; + private final BatchInsertUtils batchInsertUtils; private static final ObjectMapper objectMapper = new ObjectMapper(); - DatabendPreparedStatement(DatabendConnection connection, Consumer onClose, String statementName, - String sql) { + DatabendPreparedStatement(DatabendConnection connection, Consumer onClose, String sql) throws SQLException { super(connection, onClose); - this.statementName = requireNonNull(statementName, "statementName is null"); - this.originalSql = requireNonNull(sql, "sql is null"); this.batchValues = new ArrayList<>(); - this.batchInsertUtils = BatchInsertUtils.tryParseInsertSql(sql); + this.batchValuesCSV = new ArrayList<>(); + this.batchInsertUtils = new BatchInsertUtils(sql); this.rawStatement = StatementUtil.parseToRawStatementWrapper(sql); + if(this.rawStatement.getSubStatements().size() > 1) { + throw new SQLException("Databend do not support multi statement for now"); + } Map params = StatementUtil.extractColumnTypes(sql); List list = params.entrySet().stream().map(entry -> { String type = entry.getValue(); @@ -160,55 +156,12 @@ public void close() super.close(); } - private DatabendCopyParams uploadBatchesForCopyInto() throws SQLException { - if (this.batchValues == null || this.batchValues.size() == 0) { - return null; - } - File saved = batchInsertUtils.get().saveBatchToCSV(batchValues); - try (FileInputStream fis = new FileInputStream(saved);) { - DatabendConnection c = (DatabendConnection) getConnection(); - String uuid = UUID.randomUUID().toString().replace("-", ""); - // format %Y/%m/%d/%H/%M/%S/fileName.csv - String stagePrefix = String.format("%s/%s/%s/%s/%s/%s/%s/", - LocalDateTime.now().getYear(), - LocalDateTime.now().getMonthValue(), - LocalDateTime.now().getDayOfMonth(), - LocalDateTime.now().getHour(), - LocalDateTime.now().getMinute(), - LocalDateTime.now().getSecond(), - uuid); - String fileName = saved.getName(); - c.uploadStream(null, stagePrefix, fis, fileName, saved.length(), false); - String stageName = "~"; - Map copyOptions = new HashMap<>(); - copyOptions.put("PURGE", String.valueOf(c.copyPurge())); - copyOptions.put("NULL_DISPLAY", String.valueOf(c.nullDisplay())); - DatabendStage databendStage = DatabendStage.builder().stageName(stageName).path(stagePrefix).build(); - List files = new ArrayList<>(); - files.add(fileName); - DatabendCopyParams databendCopyParams = DatabendCopyParams.builder().setFiles(files) - .setCopyOptions(copyOptions).setDatabaseTableName(batchInsertUtils.get().getDatabaseTableName()) - .setDatabendStage(databendStage).build(); - return databendCopyParams; - } catch (Exception e) { - throw new SQLException(e); - } finally { - try { - if (saved != null) { - saved.delete(); - } - } catch (Exception e) { - // ignore - } - } - } - private StageAttachment uploadBatches() throws SQLException { - if (this.batchValues == null || this.batchValues.size() == 0) { + if (this.batchValuesCSV == null || this.batchValuesCSV.size() == 0) { return null; } - File saved = batchInsertUtils.get().saveBatchToCSV(batchValues); - try (FileInputStream fis = new FileInputStream(saved);) { + File saved = batchInsertUtils.saveBatchToCSV(batchValuesCSV); + try (FileInputStream fis = new FileInputStream(saved)) { DatabendConnection c = (DatabendConnection) getConnection(); String uuid = UUID.randomUUID().toString().replace("-", ""); // format %Y/%m/%d/%H/%M/%S/fileName.csv @@ -224,8 +177,7 @@ private StageAttachment uploadBatches() throws SQLException { // upload to stage c.uploadStream(null, stagePrefix, fis, fileName, saved.length(), false); String stagePath = "@~/" + stagePrefix + fileName; - StageAttachment attachment = buildStateAttachment(c, stagePath); - return attachment; + return buildStateAttachment(c, stagePath); } catch (Exception e) { throw new SQLException(e); } finally { @@ -277,7 +229,6 @@ public static StageAttachment buildStateAttachment(DatabendConnection connection /** * delete stage file on stage attachment * - * @param attachment * @return true if delete success or resource not found */ private boolean dropStageAttachment(StageAttachment attachment) { @@ -289,274 +240,171 @@ private boolean dropStageAttachment(StageAttachment attachment) { execute(sql); return true; } catch (SQLException e) { - if (e.getErrorCode() == 1003) { - return true; - } - return false; + return e.getErrorCode() == 1003; } } public int[] executeBatchByAttachment() throws SQLException { int[] batchUpdateCounts = new int[batchValues.size()]; - if (!batchInsertUtils.isPresent() || batchValues == null || batchValues.isEmpty()) { - // super.execute(this.originalSql); + if (batchValues.isEmpty()) { return batchUpdateCounts; } StageAttachment attachment = uploadBatches(); - ResultSet r = null; if (attachment == null) { - // logger.fine("use normal execute instead of batch insert"); - // super.execute(batchInsertUtils.get().getSql()); return batchUpdateCounts; } try { logger.fine(String.format("use batch insert instead of normal insert, attachment: %s, sql: %s", attachment, - batchInsertUtils.get().getSql())); - super.internalExecute(batchInsertUtils.get().getSql(), attachment); - r = getResultSet(); - while (r.next()) { + batchInsertUtils.getSql())); + super.internalExecute(batchInsertUtils.getSql(), attachment); + try (ResultSet r = getResultSet()) { + while (r.next()) { + } } Arrays.fill(batchUpdateCounts, 1); return batchUpdateCounts; - } catch (RuntimeException e) { - throw new SQLException(e); } finally { - clearBatch(); dropStageAttachment(attachment); + clearBatch(); } } - public int[] executeBatchByCopyInto() throws SQLException { - int[] batchUpdateCounts = new int[batchValues.size()]; - if (!batchInsertUtils.isPresent() || batchValues == null || batchValues.isEmpty()) { - super.execute(this.originalSql); - return batchUpdateCounts; - } - DatabendCopyParams databendCopyParams = uploadBatchesForCopyInto(); - ResultSet r = null; - if (databendCopyParams == null) { - logger.fine("use normal execute instead of batch insert"); - super.execute(batchInsertUtils.get().getSql()); - return batchUpdateCounts; - } - try { - String sql = DatabendConnection.getCopyIntoSql(null, databendCopyParams); - logger.fine(String.format("use copy into instead of normal insert, copy into SQL: %s", sql)); - super.internalExecute(sql, null); - r = getResultSet(); - while (r.next()) { - - } - Arrays.fill(batchUpdateCounts, 1); - return batchUpdateCounts; - } catch (RuntimeException e) { - throw new SQLException(e); - } - } - - public int[] executeBatchDelete() throws SQLException { - if (!batchInsertUtils.isPresent() || batchValues == null || batchValues.isEmpty()) { - return new int[] {}; - } - int[] batchUpdateCounts = new int[batchValues.size()]; - try { - String sql = convertSQLWithBatchValues(this.originalSql, this.batchValues); - logger.fine(String.format("use copy into instead of normal insert, copy into SQL: %s", sql)); - super.internalExecute(sql, null); - ResultSet r = getResultSet(); - while (r.next()) { - - } - return batchUpdateCounts; - } catch (RuntimeException e) { - throw new SQLException(e); - } - } - - public static String convertSQLWithBatchValues(String baseSql, List batchValues) { - StringBuilder convertedSqlBuilder = new StringBuilder(); - - if (batchValues != null && !batchValues.isEmpty()) { - for (String[] values : batchValues) { - if (values != null && values.length > 0) { - String convertedSql = baseSql; - for (int i = 0; i < values.length; i++) { - convertedSql = convertedSql.replaceFirst("\\?", values[i]); - } - convertedSqlBuilder.append(convertedSql).append(";\n"); - } - } - } - - return convertedSqlBuilder.toString(); - } - - @Override - public int[] executeBatch() throws SQLException { - if (originalSql.toLowerCase().contains("delete from")) { - return executeBatchDelete(); - } - return executeBatchByAttachment(); - } - @Override public ResultSet executeQuery() throws SQLException { - String sql = replaceParameterMarksWithValues(batchInsertUtils.get().getProvideParams(), this.originalSql) - .get(0) - .getSql(); - internalExecute(sql, null); + execute(); return getResultSet(); } - private List prepareSQL(@NonNull Map params) { - return replaceParameterMarksWithValues(params, this.rawStatement); + + @Override + public int executeUpdate() throws SQLException { + execute(); + return getUpdateCount(); } @Override public boolean execute() throws SQLException { - boolean r; - try { - r = this.execute(prepareSQL(batchInsertUtils.get().getProvideParams())); - } catch (Exception e) { - throw new SQLException(e); - } finally { - clearBatch(); - } - return r; + String sql = replaceParameterMarksWithValues(batchInsertUtils.getProvideParams(), this.rawStatement).get(0).getSql(); + return execute(sql); } - protected boolean execute(List statements) throws SQLException { - try { - for (int i = 0; i < statements.size(); i++) { - String sql = statements.get(i).getSql(); - if (isBatchInsert(sql)) { - handleBatchInsert(); - } else { - execute(sql); + @Override + public int[] executeBatch() throws SQLException { + if (isBatchInsert(batchInsertUtils.getSql())) { + return executeBatchByAttachment(); + } else { + int[] batchUpdateCounts = new int[batchValues.size()]; + for (int i = 0; i < batchValues.size(); i++) { + String [] values = batchValues.get(i); + Map m = new HashMap<>(); + for (int j = 0; j< values.length; j++){ + m.put(j + 1, values[j]); } - return true; + String sql = replaceParameterMarksWithValues(m, this.rawStatement).get(0).getSql(); + this.execute(sql); + batchUpdateCounts[i]= getUpdateCount(); } - } catch (Exception e) { - throw new SQLException(e); - } finally { + return batchUpdateCounts; } - return true; } - private boolean isBatchInsert(String sql) { - return sql.toLowerCase().contains(DATABEND_KEYWORDS_INSERT_INTO) && !sql.toLowerCase().contains(DATABEND_KEYWORDS_SELECT); + private static boolean isBatchInsert(String sql) { + sql = sql.toLowerCase(); + Matcher matcher = INSERT_INTO_PATTERN.matcher(sql); + return matcher.find() && !sql.contains(DATABEND_KEYWORDS_SELECT); } - protected void handleBatchInsert() throws SQLException { - try { - addBatch(); - executeBatch(); - } catch (Exception e) { - throw new SQLException(e); - } + + private void setValueSimple(int index, String value) { + batchInsertUtils.setPlaceHolderValue(index, value, value); } - @Override - public int executeUpdate() throws SQLException { - this.execute(prepareSQL(batchInsertUtils.get().getProvideParams())); - return getUpdateCount(); + private void setValue(int index, String value, String csvValue) { + batchInsertUtils.setPlaceHolderValue(index, value, csvValue); } @Override public void setNull(int i, int i1) throws SQLException { checkOpen(); - if (this.originalSql.toLowerCase().contains("insert") || - this.originalSql.toLowerCase().contains("replace")) { - // Databend uses \N as default null representation for csv and tsv format - // https://github.com/datafuselabs/databend/pull/6453 - batchInsertUtils.ifPresent(insertUtils -> insertUtils.setPlaceHolderValue(i, "\\N")); - } else { - batchInsertUtils.ifPresent(insertUtils -> insertUtils.setPlaceHolderValue(i, "null")); - } + setValue(i, "null", "\\N"); } @Override public void setBoolean(int i, boolean b) throws SQLException { checkOpen(); - batchInsertUtils.ifPresent(insertUtils -> insertUtils.setPlaceHolderValue(i, formatBooleanLiteral(b))); + setValueSimple(i, formatBooleanLiteral(b)); } @Override public void setByte(int i, byte b) throws SQLException { checkOpen(); - batchInsertUtils.ifPresent(insertUtils -> insertUtils.setPlaceHolderValue(i, formatByteLiteral(b))); + setValueSimple(i, formatByteLiteral(b)); } @Override public void setShort(int i, short i1) throws SQLException { checkOpen(); - batchInsertUtils.ifPresent(insertUtils -> insertUtils.setPlaceHolderValue(i, formatShortLiteral(i1))); + setValueSimple(i, formatShortLiteral(i1)); } @Override public void setInt(int i, int i1) throws SQLException { checkOpen(); - batchInsertUtils.ifPresent(insertUtils -> insertUtils.setPlaceHolderValue(i, formatIntLiteral(i1))); + setValueSimple(i, formatIntLiteral(i1)); } @Override public void setLong(int i, long l) throws SQLException { checkOpen(); - batchInsertUtils.ifPresent(insertUtils -> insertUtils.setPlaceHolderValue(i, formatLongLiteral(l))); + setValueSimple(i, formatLongLiteral(l)); } @Override public void setFloat(int i, float v) throws SQLException { checkOpen(); - batchInsertUtils.ifPresent(insertUtils -> insertUtils.setPlaceHolderValue(i, formatFloatLiteral(v))); + setValueSimple(i, formatFloatLiteral(v)); } @Override public void setDouble(int i, double v) throws SQLException { checkOpen(); - batchInsertUtils.ifPresent(insertUtils -> insertUtils.setPlaceHolderValue(i, formatDoubleLiteral(v))); + setValueSimple(i, formatDoubleLiteral(v)); } @Override - public void setBigDecimal(int i, BigDecimal bigDecimal) + public void setBigDecimal(int i, BigDecimal v) throws SQLException { checkOpen(); - batchInsertUtils - .ifPresent(insertUtils -> insertUtils.setPlaceHolderValue(i, formatBigDecimalLiteral(bigDecimal))); + setValueSimple(i, formatBigDecimalLiteral(v)); } @Override public void setString(int i, String s) throws SQLException { checkOpen(); - if ((originalSql.toLowerCase().startsWith("insert") || - originalSql.toLowerCase().startsWith("replace")) && !originalSql.toLowerCase().contains("select")) { - String finalS1 = s; - batchInsertUtils.ifPresent(insertUtils -> insertUtils.setPlaceHolderValue(i, finalS1)); - } else { - if (s.contains("'")) { - s = s.replace("'", "\\\'"); - } - String finalS = s; - batchInsertUtils.ifPresent( - insertUtils -> insertUtils.setPlaceHolderValue(i, String.format("%s%s%s", "'", finalS, "'"))); + String quoted = s; + if (s.contains("'")) { + quoted = s.replace("'", "\\'"); } + quoted = String.format("'%s'", quoted); + + setValue(i, quoted, s); } @Override - public void setBytes(int i, byte[] bytes) + public void setBytes(int i, byte[] v) throws SQLException { checkOpen(); - batchInsertUtils.ifPresent(insertUtils -> insertUtils.setPlaceHolderValue(i, formatBytesLiteral(bytes))); + setValueSimple(i, formatBytesLiteral(v)); } @Override @@ -564,47 +412,31 @@ public void setDate(int i, Date date) throws SQLException { checkOpen(); if (date == null) { - batchInsertUtils.ifPresent(insertUtils -> insertUtils.setPlaceHolderValue(i, null)); + setValueSimple(i, null); } else { - if (originalSql.toLowerCase().startsWith("select")) { - batchInsertUtils.ifPresent( - insertUtils -> insertUtils.setPlaceHolderValue(i, String.format("%s%s%s", "'", date, "'"))); - } else { - batchInsertUtils.ifPresent(insertUtils -> insertUtils.setPlaceHolderValue(i, toDateLiteral(date))); - } + setValue(i, String.format("'%s'", date), toDateLiteral(date)); } } @Override - public void setTime(int i, Time time) + public void setTime(int i, Time v) throws SQLException { checkOpen(); - if (time == null) { - batchInsertUtils.ifPresent(insertUtils -> insertUtils.setPlaceHolderValue(i, null)); + if (v == null) { + setValueSimple(i, null); } else { - if (originalSql.toLowerCase().startsWith("select")) { - batchInsertUtils.ifPresent( - insertUtils -> insertUtils.setPlaceHolderValue(i, String.format("%s%s%s", "'", time, "'"))); - } else { - batchInsertUtils.ifPresent(insertUtils -> insertUtils.setPlaceHolderValue(i, toTimeLiteral(time))); - } + setValue(i, String.format("'%s'", v), toTimeLiteral(v)); } } @Override - public void setTimestamp(int i, Timestamp timestamp) + public void setTimestamp(int i, Timestamp v) throws SQLException { checkOpen(); - if (timestamp == null) { - batchInsertUtils.ifPresent(insertUtils -> insertUtils.setPlaceHolderValue(i, null)); + if (v == null) { + setValueSimple(i, null); } else { - if (originalSql.toLowerCase().startsWith("select")) { - batchInsertUtils.ifPresent(insertUtils -> insertUtils.setPlaceHolderValue(i, - String.format("%s%s%s", "'", timestamp, "'"))); - } else { - batchInsertUtils - .ifPresent(insertUtils -> insertUtils.setPlaceHolderValue(i, toTimestampLiteral(timestamp))); - } + setValue(i, String.format("'%s'", v), toTimestampLiteral(v)); } } @@ -630,7 +462,7 @@ public void setBinaryStream(int i, InputStream inputStream, int i1) public void clearParameters() throws SQLException { checkOpen(); - batchInsertUtils.ifPresent(BatchInsertUtils::clean); + batchInsertUtils.clean(); } @Override @@ -783,20 +615,22 @@ public static String convertArrayListToString(ArrayList arrayList) { public void addBatch() throws SQLException { checkOpen(); - if (batchInsertUtils.isPresent()) { - String[] val = batchInsertUtils.get().getValues(); - batchValues.add(val); - batchInsertUtils.get().clean(); - batchSize++; - } + + String[] val = batchInsertUtils.getValues(); + batchValues.add(val); + + val = batchInsertUtils.getValuesCSV(); + batchValuesCSV.add(val); + + batchInsertUtils.clean(); } @Override public void clearBatch() throws SQLException { checkOpen(); batchValues.clear(); - batchSize = 0; - batchInsertUtils.ifPresent(BatchInsertUtils::clean); + batchValuesCSV.clear(); + batchInsertUtils.clean(); } @Override @@ -974,10 +808,10 @@ public void setBinaryStream(int i, InputStream inputStream) byte[] bytes = buffer.toByteArray(); if (BASE64_STR.equalsIgnoreCase(connection().binaryFormat())) { String base64String = bytesToBase64(bytes); - batchInsertUtils.ifPresent(insertUtils -> insertUtils.setPlaceHolderValue(i, base64String)); + setValueSimple(i, base64String); } else { String hexString = bytesToHex(bytes); - batchInsertUtils.ifPresent(insertUtils -> insertUtils.setPlaceHolderValue(i, hexString)); + setValueSimple(i, hexString); } } catch (IOException e) { throw new SQLException("Error reading InputStream", e); diff --git a/databend-jdbc/src/main/java/com/databend/jdbc/DatabendStatement.java b/databend-jdbc/src/main/java/com/databend/jdbc/DatabendStatement.java index 681defa9..d7fbbe44 100644 --- a/databend-jdbc/src/main/java/com/databend/jdbc/DatabendStatement.java +++ b/databend-jdbc/src/main/java/com/databend/jdbc/DatabendStatement.java @@ -12,7 +12,6 @@ import java.sql.SQLWarning; import java.sql.Statement; import java.util.List; -import java.util.Optional; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicLong; import java.util.concurrent.atomic.AtomicReference; @@ -110,7 +109,6 @@ public int getQueryTimeout() @Override public void setQueryTimeout(int i) throws SQLException { - } @Override @@ -208,7 +206,7 @@ final boolean internalExecute(String sql, StageAttachment attachment) throws SQL if (updateCount instanceof Number) { currentUpdateCount = ((Number) updateCount).intValue(); } else { - // if can't find, use writeProgress.rows + // if not found, use writeProgress.rows currentUpdateCount = results.getStats().getWriteProgress().getRows().intValue(); } } else { @@ -332,7 +330,7 @@ public boolean getMoreResults(int i) return false; } - if (i != KEEP_CURRENT_RESULT && i != CLOSE_CURRENT_RESULT) { + if (i != KEEP_CURRENT_RESULT) { throw new SQLException("Invalid value for getMoreResults: " + i); } throw new SQLFeatureNotSupportedException("Multiple results not supported"); @@ -460,8 +458,4 @@ public QueryLiveness queryLiveness() { } return null; } - - protected final Optional optionalConnection() { - return Optional.ofNullable(connection.get()); - } } diff --git a/databend-jdbc/src/main/java/com/databend/jdbc/StatementUtil.java b/databend-jdbc/src/main/java/com/databend/jdbc/StatementUtil.java index 90f84dd8..d7d9e3f6 100644 --- a/databend-jdbc/src/main/java/com/databend/jdbc/StatementUtil.java +++ b/databend-jdbc/src/main/java/com/databend/jdbc/StatementUtil.java @@ -81,16 +81,6 @@ public static Map extractColumnTypes(String sql) { return columnTypes; } - /** - * Parse the sql statement to a list of {@link StatementInfoWrapper} - * - * @param sql the sql statement - * @return a list of {@link StatementInfoWrapper} - */ - public List parseToStatementInfoWrappers(String sql) { - return parseToRawStatementWrapper(sql).getSubStatements().stream().map(StatementInfoWrapper::of) - .collect(Collectors.toList()); - } /** * Parse sql statement to a {@link RawStatementWrapper}. The method construct @@ -250,6 +240,12 @@ public static List replaceParameterMarksWithValues(@NonNul */ public List replaceParameterMarksWithValues(@NonNull Map params, @NonNull RawStatementWrapper rawStatement) { + if (params.size() != rawStatement.getTotalParams()) { + throw new IllegalArgumentException(String.format( + "The number of parameters passed does not equal the number of parameter markers in the SQL query. Provided: %d, Parameter markers in the SQL query: %d", + params.size(), rawStatement.getTotalParams())); + } + List subQueries = new ArrayList<>(); for (int subqueryIndex = 0; subqueryIndex < rawStatement.getSubStatements().size(); subqueryIndex++) { int currentPos; @@ -260,12 +256,6 @@ public List replaceParameterMarksWithValues(@NonNull Map> extractPropertyPair(String cleanStatement "Cannot parse the additional properties provided in the statement: " + sql); } } + + } diff --git a/databend-jdbc/src/main/java/com/databend/jdbc/constant/DatabendConstant.java b/databend-jdbc/src/main/java/com/databend/jdbc/constant/DatabendConstant.java index e8f146c8..09c89cca 100644 --- a/databend-jdbc/src/main/java/com/databend/jdbc/constant/DatabendConstant.java +++ b/databend-jdbc/src/main/java/com/databend/jdbc/constant/DatabendConstant.java @@ -1,5 +1,7 @@ package com.databend.jdbc.constant; +import java.util.regex.Pattern; + /** * databend constant * @@ -8,6 +10,7 @@ public class DatabendConstant { public static final String ENABLE_STR = "enable"; public static final String BASE64_STR = "base64"; - public static final String DATABEND_KEYWORDS_INSERT_INTO = "insert into"; + public static final Pattern INSERT_INTO_PATTERN = Pattern.compile("(insert|replace)\\s+into"); + public static final String DATABEND_KEYWORDS_SELECT = "select"; } diff --git a/databend-jdbc/src/main/java/com/databend/jdbc/parser/BatchInsertUtils.java b/databend-jdbc/src/main/java/com/databend/jdbc/parser/BatchInsertUtils.java index cbbb551a..95611d14 100644 --- a/databend-jdbc/src/main/java/com/databend/jdbc/parser/BatchInsertUtils.java +++ b/databend-jdbc/src/main/java/com/databend/jdbc/parser/BatchInsertUtils.java @@ -8,41 +8,28 @@ import java.io.IOException; import java.util.List; import java.util.Map; -import java.util.Optional; import java.util.TreeMap; import java.util.UUID; -import java.util.logging.Logger; import java.util.regex.Matcher; import java.util.regex.Pattern; -import java.util.stream.Collectors; -import java.util.stream.Stream; public class BatchInsertUtils { - private static final Logger logger = Logger.getLogger(BatchInsertUtils.class.getPackage().getName()); - private final String sql; private String databaseTableName; // prepareValues[i] is null if the i-th value is a placeholder - private TreeMap placeHolderEntries; + private final TreeMap placeHolderEntries; + private final TreeMap placeHolderEntriesCSV; - private BatchInsertUtils(String sql) { + public BatchInsertUtils(String sql) { this.sql = sql; // sort key in ascending order this.placeHolderEntries = new TreeMap<>(); + this.placeHolderEntriesCSV = new TreeMap<>(); // this.databaseTableName = getDatabaseTableName(); } - /** - * Parse the sql to get insert AST - * - * @param sql candidate sql - * @return BatchInertUtils if the sql is a batch insert sql - */ - public static Optional tryParseInsertSql(String sql) { - return Optional.of(new BatchInsertUtils(sql)); - } public String getSql() { return sql; @@ -67,10 +54,11 @@ public String getDatabaseTableName() { return databaseTableName; } - public void setPlaceHolderValue(int index, String value) throws IllegalArgumentException { + public void setPlaceHolderValue(int index, String value, String valueCSV) throws IllegalArgumentException { int i = index - 1; placeHolderEntries.put(i, value); + placeHolderEntriesCSV.put(i, valueCSV); } // get the sql with placeholder replaced by value @@ -84,6 +72,17 @@ public String[] getValues() { } return values; } + public String[] getValuesCSV() { + if (placeHolderEntriesCSV.isEmpty()) { + return null; + } + String[] values = new String[placeHolderEntriesCSV.lastKey() + 1]; + for (Map.Entry elem : placeHolderEntriesCSV.entrySet()) { + values[elem.getKey()] = elem.getValue(); + } + return values; + } + public File saveBatchToCSV(List values) { // get a temporary directory @@ -93,16 +92,10 @@ public File saveBatchToCSV(List values) { return saveBatchToCSV(values, tempFile); } - private String convertToCSV(String[] data) { - return Stream.of(data) - .collect(Collectors.joining(",")); - } public File saveBatchToCSV(List values, File file) { - int rowSize = 0; for (String[] row : values) { if (row != null) { - rowSize = row.length; break; } throw new RuntimeException("batch values is empty"); @@ -122,5 +115,6 @@ public File saveBatchToCSV(List values, File file) { public void clean() { placeHolderEntries.clear(); + placeHolderEntriesCSV.clear(); } } diff --git a/databend-jdbc/src/test/java/com/databend/jdbc/TestPrepareStatement.java b/databend-jdbc/src/test/java/com/databend/jdbc/TestPrepareStatement.java index 3849c94e..150a62b7 100644 --- a/databend-jdbc/src/test/java/com/databend/jdbc/TestPrepareStatement.java +++ b/databend-jdbc/src/test/java/com/databend/jdbc/TestPrepareStatement.java @@ -15,10 +15,12 @@ import java.sql.Statement; import java.sql.Timestamp; import java.sql.Types; -import java.util.ArrayList; -import java.util.List; +import java.util.Arrays; import java.util.Properties; +import static com.databend.jdbc.Utils.countTable; +import static org.testng.Assert.*; + public class TestPrepareStatement { @@ -110,34 +112,12 @@ public void TestBatchInsertWithNULL() throws SQLException { } } - @Test(groups = "UNIT") - public void TestConvertSQLWithBatchValues() { - List batchValues = new ArrayList<>(); - // Add string arrays to batchValues - String[] values1 = { "1" }; - String[] values2 = { "2" }; - batchValues.add(values1); - batchValues.add(values2); - - String originalSql = "delete from table where id = ?"; - String expectedSql = "delete from table where id = 1;\ndelete from table where id = 2;\n"; - Assert.assertEquals(DatabendPreparedStatement.convertSQLWithBatchValues(originalSql, batchValues), expectedSql); - - List batchValues1 = new ArrayList<>(); - // Add string arrays to batchValues - String[] values3 = { "1", "2" }; - String[] values4 = { "3", "4" }; - batchValues1.add(values3); - batchValues1.add(values4); - - String originalSql1 = "delete from table where id = ? and uuid = ?"; - String expectedSql1 = "delete from table where id = 1 and uuid = 2;\ndelete from table where id = 3 and uuid = 4;\n"; - Assert.assertEquals(DatabendPreparedStatement.convertSQLWithBatchValues(originalSql1, batchValues1), - expectedSql1); - } @Test(groups = "IT") public void TestBatchDelete() throws SQLException { + if (Compatibility.skipDriverBugLowerThen("0.4.1")) { + return; + } try (Connection c = getConn(); Statement s = c.createStatement()) { c.createStatement().execute("create or replace table t1(a int, b string)"); @@ -165,10 +145,10 @@ public void TestBatchDelete() throws SQLException { PreparedStatement deletePs = c.prepareStatement("delete from t1 where a = ?"); deletePs.setInt(1, 1); deletePs.addBatch(); - int[] ansDel = deletePs.executeBatch(); - Assert.assertEquals(ansDel.length, 1); - // todo: fix this, currently == 0 - // Assert.assertEquals(ansDel[0], 1); + deletePs.setInt(1, 2); + deletePs.addBatch(); + int[] counts = deletePs.executeBatch(); + Assert.assertEquals(counts, new int[] {1, 0}, Arrays.toString(counts)); s.execute("SELECT * from t1"); ResultSet r1 = s.getResultSet(); @@ -296,7 +276,7 @@ public void TestBatchReplaceInto() throws SQLException { @Test(groups = "IT") public void testPrepareStatementExecute() throws SQLException { try (Connection c = getConn(); - Statement s = c.createStatement()) { + Statement s = c.createStatement()) { s.execute("create or replace table t1 (a int, b string)"); String insertSql = "insert into t1 values (?,?)"; try (PreparedStatement ps = c.prepareStatement(insertSql)) { @@ -342,8 +322,7 @@ public void testUpdateSetNull() throws SQLException { statement.setString(2, "b"); statement.addBatch(); int[] result = statement.executeBatch(); - System.out.println(result); - Assert.assertEquals(1, result.length); + Assert.assertEquals(result, new int[]{1}); } String updateSQL = "update t1 set b = ? where a = ?"; try (PreparedStatement statement = conn.prepareStatement(updateSQL)) { @@ -380,16 +359,14 @@ public void testUpdateStatement() throws SQLException { statement.setString(2, "b"); statement.addBatch(); int[] result = statement.executeBatch(); - System.out.println(result); - Assert.assertEquals(1, result.length); + Assert.assertEquals(result, new int[]{1}); } String updateSQL = "update t1 set b = ? where a = ?"; try (PreparedStatement statement = conn.prepareStatement(updateSQL)) { statement.setInt(2, 1); statement.setObject(1, "c'c"); int result = statement.executeUpdate(); - System.out.println(result); - Assert.assertEquals(1, result); + Assert.assertEquals(result, 1); } try (PreparedStatement statement = conn .prepareStatement("select a, regexp_replace(b, '\\d', '*') from t1 where a = ?")) { @@ -411,75 +388,77 @@ public void testAllPreparedStatement() throws SQLException { String sql = "insert into t1 values (?,?)"; try (PreparedStatement statement = conn.prepareStatement(sql)) { statement.setInt(1, 1); - statement.setString(2, "b"); + statement.setString(2, "r1"); statement.addBatch(); statement.setInt(1, 2); - statement.setString(2, "z"); + statement.setString(2, "r2"); statement.addBatch(); statement.setInt(1, 3); - statement.setString(2, "x"); + statement.setString(2, "r3"); statement.addBatch(); statement.setInt(1, 4); - statement.setString(2, "dd"); + statement.setString(2, "r3"); statement.addBatch(); statement.setInt(1, 5); - statement.setString(2, "ddd"); + statement.setString(2, "r5"); statement.addBatch(); int[] result = statement.executeBatch(); - System.out.println(result); - Assert.assertEquals(5, result.length); + Assert.assertEquals(result, new int[] {1, 1, 1, 1, 1}); } String updateSQL = "update t1 set b = ? where b = ?"; try (PreparedStatement statement = conn.prepareStatement(updateSQL)) { - statement.setString(1, "c"); - statement.setString(2, "b"); + statement.setString(1, "r1_new"); + statement.setString(2, "r1"); int result = statement.executeUpdate(); Assert.assertEquals(1, result); } try (PreparedStatement statement = conn - .prepareStatement("select a, regexp_replace(b, '\\d', '*') from t1 where b = ?")) { - statement.setString(1, "c"); + .prepareStatement("select a, b from t1 where b = ?")) { + statement.setString(1, "r1_new"); ResultSet r = statement.executeQuery(); while (r.next()) { Assert.assertEquals(1, r.getInt(1)); - Assert.assertEquals("c", r.getString(2)); + Assert.assertEquals("r1_new", r.getString(2)); } } String replaceIntoSQL = "replace into t1 on(a) values (?,?)"; try (PreparedStatement statement = conn.prepareStatement(replaceIntoSQL)) { statement.setInt(1, 1); - statement.setString(2, "d"); + statement.setString(2, "r1_new2"); statement.addBatch(); Assert.assertEquals(statement.executeBatch(), new int[]{1}); } ResultSet r2 = conn.createStatement().executeQuery("select * from t1"); + int n = 0; while (r2.next()) { - System.out.println(r2.getInt(1)); - System.out.println(r2.getString(2)); + n +=1; } + Assert.assertEquals(n, 5); String deleteSQL = "delete from t1 where a = ?"; try (PreparedStatement statement = conn.prepareStatement(deleteSQL)) { statement.setInt(1, 1); boolean result = statement.execute(); + // TODO: fix this + // Assert.assertFalse(result); System.out.println(result); + Assert.assertEquals(statement.getUpdateCount(), 1); } String deleteSQLVarchar = "delete from t1 where b = ?"; try (PreparedStatement statement = conn.prepareStatement(deleteSQLVarchar)) { - statement.setString(1, "1"); + statement.setString(1, "not exists"); int result = statement.executeUpdate(); - System.out.println(result); + Assert.assertEquals(result, 0); } ResultSet r3 = conn.createStatement().executeQuery("select * from t1"); - Assert.assertEquals(0, r3.getRow()); + n = 0; while (r3.next()) { - // noting print - System.out.println(r3.getInt(1)); - System.out.println(r3.getString(2)); + n +=1; } + Assert.assertEquals(n, 4); } } @@ -510,8 +489,7 @@ public void testSelectWithClusterKey() throws SQLException { statement.setString(2, "c"); statement.addBatch(); int[] result = statement.executeBatch(); - System.out.println(result); - Assert.assertEquals(2, result.length); + Assert.assertEquals(result, new int[] {1, 1}); } conn.createStatement().execute("alter table t1 cluster by (a)"); String selectSQL = String.format("select * from clustering_information('%s','t1')", DB_NAME.get()); @@ -529,15 +507,15 @@ public void testSelectWithClusterKey() throws SQLException { @Test(groups = "IT") public void testEncodePass() throws SQLException { - try (Connection conn = Utils.createConnection(); - Connection conn2 = Utils.createConnection()) { + try (Connection conn = Utils.createConnection()) { conn.createStatement().execute("create user if not exists 'u01' identified by 'mS%aFRZW*GW';"); conn.createStatement().execute("GRANT ALL PRIVILEGES ON default.* TO 'u01'@'%'"); Properties p = new Properties(); p.setProperty("user", "u01"); p.setProperty("password", "mS%aFRZW*GW"); - - conn2.createStatement().execute("select 1"); + try(Connection conn2 = Utils.createConnection("default", p)) { + conn2.createStatement().execute("select 1"); + } conn.createStatement().execute("drop user if exists 'u01'"); } } @@ -634,17 +612,20 @@ public void testInsertWithSelect() throws SQLException { ps.setInt(1, 1); ps.setString(2, "a"); int insertedRows = ps.executeUpdate(); - // TODO: fix this - // Assert.assertEquals(1, insertedRows, "should insert 1 rows"); - System.out.println(insertedRows); + + // else 0 + if (!Compatibility.skipDriverBugLowerThen("0.4.1")) { + Assert.assertEquals(1, insertedRows, "should insert 1 rows"); + } ps.setInt(1, 2); ps.setString(2, "b"); insertedRows = ps.executeUpdate(); - // TODO: fix this - // Assert.assertEquals(1, insertedRows, "should insert 1 rows"); - System.out.println(insertedRows); + if (!Compatibility.skipDriverBugLowerThen("0.4.1")) { + Assert.assertEquals(1, insertedRows, "should insert 1 rows"); + } } + Assert.assertEquals(countTable(s, "t1"), 2); // Now try to insert again with select try (PreparedStatement ps = conn.prepareStatement(insertSql)) { @@ -652,13 +633,63 @@ public void testInsertWithSelect() throws SQLException { int insertedRows = ps.executeUpdate(); Assert.assertEquals(1, insertedRows, "should insert 1 row from the select"); } + Assert.assertEquals(countTable(s, "t1"), 3); + } + } - ResultSet rs = conn.createStatement().executeQuery("select * from t1 order by a"); - int count = 0; - while (rs.next()) { - count++; + @Test(groups = "IT") + public void testMultiStatement() throws SQLException { + if (Compatibility.skipDriverBugLowerThen("0.4.1")) { + return; + } + try (Connection conn = getConn()) { + assertThrows(SQLException.class, () -> conn.prepareStatement("select 1; select 1")); + } + } + + @Test(groups = "IT") + public void testBatchAndNoBatch() throws SQLException { + if (Compatibility.skipDriverBugLowerThen("0.4.1")) { + return; + } + try (Connection conn = getConn(); + Statement s = conn.createStatement()) { + s.execute("create or replace table t1(a int, b string)"); + String insertSql = "insert into t1 values (?, ?)"; + try (PreparedStatement ps = conn.prepareStatement(insertSql)) { + ps.setInt(1, 1); + ps.setString(2, "v1"); + ps.addBatch(); + + ps.setInt(1, 2); + ps.setString(2, "v2"); + int insertedRows = ps.executeUpdate(); + Assert.assertEquals(1, insertedRows); + + insertedRows = ps.executeUpdate(); + Assert.assertEquals(1, insertedRows); + + try(ResultSet rs = s.executeQuery("select * from t1")) { + for (int i = 0; i < 2; i++) { + assertTrue(rs.next()); + assertEquals(rs.getInt(1) ,2); + assertEquals(rs.getString(2) ,"v2"); + } + assertFalse(rs.next()); + } + s.execute("truncate table t1"); + assertEquals(countTable(s, "t1"), 0); + + int[] counts = ps.executeBatch(); + Assert.assertEquals(counts, new int[] {1}); + + try(ResultSet rs = s.executeQuery("select * from t1")) { + assertTrue(rs.next()); + assertEquals(rs.getInt(1) ,1); + assertEquals(rs.getString(2) ,"v1"); + assertFalse(rs.next()); + } } - Assert.assertEquals(3, count, "should have 3 rows in the table after insert with select"); } } } diff --git a/databend-jdbc/src/test/java/com/databend/jdbc/Utils.java b/databend-jdbc/src/test/java/com/databend/jdbc/Utils.java index b15f322e..f59dc3f1 100644 --- a/databend-jdbc/src/test/java/com/databend/jdbc/Utils.java +++ b/databend-jdbc/src/test/java/com/databend/jdbc/Utils.java @@ -1,10 +1,6 @@ package com.databend.jdbc; -import com.vdurmont.semver4j.Semver; - -import java.sql.Connection; -import java.sql.DriverManager; -import java.sql.SQLException; +import java.sql.*; import java.util.Properties; public class Utils { @@ -46,5 +42,11 @@ public static Connection createConnectionWithPresignedUrlDisable() throws SQLExc String url = baseURL() + "?presigned_url_disabled=true"; return DriverManager.getConnection(url, "databend", "databend"); } + + public static int countTable(Statement statement, String table) throws SQLException { + ResultSet r = statement.executeQuery(String.format("select count(*) from %s", table)); + r.next(); + return r.getInt(1); + } } diff --git a/databend-jdbc/src/test/java/com/databend/jdbc/parser/TestBatchInsertUtils.java b/databend-jdbc/src/test/java/com/databend/jdbc/parser/TestBatchInsertUtils.java index 52ebb43b..607797e3 100644 --- a/databend-jdbc/src/test/java/com/databend/jdbc/parser/TestBatchInsertUtils.java +++ b/databend-jdbc/src/test/java/com/databend/jdbc/parser/TestBatchInsertUtils.java @@ -14,7 +14,7 @@ public class TestBatchInsertUtils { public void testFiles() throws IOException { List data = new ArrayList<>(); data.add(new String[]{"1", "2", "{\"a\": 1, \"b\": \"2\"}", "hello, world 321"}); - BatchInsertUtils b = BatchInsertUtils.tryParseInsertSql("sq").get(); + BatchInsertUtils b = new BatchInsertUtils("sq"); File f = b.saveBatchToCSV(data); System.out.println(f.getAbsolutePath()); try (FileReader fr = new FileReader(f)) { @@ -28,15 +28,15 @@ public void testFiles() throws IOException { @Test(groups = "UNIT") public void testGetDatabaseTableName() { - BatchInsertUtils b = BatchInsertUtils.tryParseInsertSql("INSERT INTO tb01(id,d,x,x,x,x,xt,col1) VALUES").get(); + BatchInsertUtils b = new BatchInsertUtils("INSERT INTO tb01(id,d,x,x,x,x,xt,col1) VALUES"); Assert.assertEquals("tb01", b.getDatabaseTableName()); - BatchInsertUtils b1 = BatchInsertUtils.tryParseInsertSql("INSERT INTO db.tb_test VALUES").get(); + BatchInsertUtils b1 = new BatchInsertUtils("INSERT INTO db.tb_test VALUES"); Assert.assertEquals("db.tb_test", b1.getDatabaseTableName()); - BatchInsertUtils b2 = BatchInsertUtils.tryParseInsertSql("INSERT INTO tb01 (id,d,x,x,x,x,xt,col1) VALUES").get(); + BatchInsertUtils b2 = new BatchInsertUtils("INSERT INTO tb01 (id,d,x,x,x,x,xt,col1) VALUES"); Assert.assertEquals("tb01", b2.getDatabaseTableName()); - BatchInsertUtils b3 = BatchInsertUtils.tryParseInsertSql("insert into tb01 values").get(); + BatchInsertUtils b3 = new BatchInsertUtils("insert into tb01 values"); Assert.assertEquals("tb01", b3.getDatabaseTableName()); - BatchInsertUtils b4 = BatchInsertUtils.tryParseInsertSql("INSERT INTO `test`(`x`, `y`) VALUES (?, ?)").get(); + BatchInsertUtils b4 = new BatchInsertUtils("INSERT INTO `test`(`x`, `y`) VALUES (?, ?)"); Assert.assertEquals("test", b4.getDatabaseTableName()); } } From 2c62fa7bced4160c892ff82a8a5dbff24a810f79 Mon Sep 17 00:00:00 2001 From: Yang Xiufeng Date: Sun, 14 Sep 2025 16:29:26 +0800 Subject: [PATCH 2/5] ci: enhance test for temp table. --- .../java/com/databend/jdbc/TestTempTable.java | 48 ++++++++++++++----- 1 file changed, 36 insertions(+), 12 deletions(-) diff --git a/databend-jdbc/src/test/java/com/databend/jdbc/TestTempTable.java b/databend-jdbc/src/test/java/com/databend/jdbc/TestTempTable.java index e30d4e1c..7cf0a879 100644 --- a/databend-jdbc/src/test/java/com/databend/jdbc/TestTempTable.java +++ b/databend-jdbc/src/test/java/com/databend/jdbc/TestTempTable.java @@ -10,19 +10,43 @@ public class TestTempTable { - @Test( groups = {"IT"}) + void checkNoTempTable() throws SQLException { + try (Connection c1 = Utils.createConnection(); + Statement statement = c1.createStatement()) { + for (int i = 0; i < 3; i++) { + ResultSet rs = statement.executeQuery("SELECT name FROM system.temporary_tables"); + Assert.assertFalse(rs.next()); + } + } + } + + @Test(groups = {"IT"}) public void testTempTable() throws SQLException { - try(Connection c1 = Utils.createConnection()) { - Statement statement= c1.createStatement(); - statement.execute("create or replace temp table table1(i int)"); - statement.execute("insert into table1 values (1), (2)"); - statement.executeQuery("select * from table1"); - ResultSet rs = statement.getResultSet(); - Assert.assertEquals(true, rs.next()); - Assert.assertEquals(1, rs.getInt(1)); - Assert.assertEquals(true, rs.next()); - Assert.assertEquals(2, rs.getInt(1)); - Assert.assertEquals(false, rs.next()); + try (Connection c1 = Utils.createConnection()) { + + // test drop table + try (Statement statement = c1.createStatement()) { + for (int i = 0; i < 10; i++) { + String tableName = "test_temp_table_" + i; + statement.execute(String.format("create or replace temp table %s(i int)", tableName)); + statement.execute(String.format("insert into %s values (1), (2)", tableName)); + statement.executeQuery("select * from " + tableName); + ResultSet rs = statement.getResultSet(); + Assert.assertTrue(rs.next()); + Assert.assertEquals(1, rs.getInt(1)); + Assert.assertTrue(rs.next()); + Assert.assertEquals(2, rs.getInt(1)); + Assert.assertFalse(rs.next()); + statement.execute("drop table " + tableName); + } + ResultSet rs = statement.executeQuery("SELECT name FROM system.temporary_tables where is_current_session = true"); + Assert.assertFalse(rs.next()); + checkNoTempTable(); + + // closed when close connection + statement.execute("create or replace temp table test_temp_table(i int)"); + } } + checkNoTempTable(); } } From c6b157a1194f8672108a4542eaccb70482d8cf61 Mon Sep 17 00:00:00 2001 From: Yang Xiufeng Date: Sun, 14 Sep 2025 22:41:16 +0800 Subject: [PATCH 3/5] refactor: add interface DatabendConnection. --- README.md | 2 +- .../java/com/databend/jdbc/Capability.java | 2 +- .../databend/jdbc/ConnectionProperties.java | 2 +- .../DatabendClientLoadBalancingPolicy.java | 2 +- .../com/databend/jdbc/DatabendColumnInfo.java | 2 +- .../com/databend/jdbc/DatabendConnection.java | 1389 +---------------- .../databend/jdbc/DatabendConnectionImpl.java | 1358 ++++++++++++++++ .../jdbc/{constant => }/DatabendConstant.java | 2 +- .../jdbc/DatabendDatabaseMetaData.java | 6 +- .../com/databend/jdbc/DatabendDriverUri.java | 4 +- .../com/databend/jdbc/DatabendNodeRouter.java | 2 +- .../java/com/databend/jdbc/DatabendNodes.java | 4 +- .../jdbc/DatabendParameterMetaData.java | 2 +- .../jdbc/DatabendPreparedStatement.java | 14 +- .../com/databend/jdbc/DatabendResultSet.java | 4 +- .../jdbc/DatabendResultSetMetaData.java | 2 +- .../com/databend/jdbc/DatabendStatement.java | 14 +- .../jdbc/DatabendUnboundQueryResultSet.java | 2 +- .../com/databend/jdbc/FileTransferAPI.java | 25 +- .../com/databend/jdbc/JdbcTypeMapping.java | 2 +- .../java/com/databend/jdbc/JdbcWrapper.java | 2 +- .../java/com/databend/jdbc/LoggerUtil.java | 2 +- .../java/com/databend/jdbc/LoginRequest.java | 2 +- .../databend/jdbc/NonQueryRawStatement.java | 2 +- .../jdbc/NonRegisteringDatabendDriver.java | 6 +- .../java/com/databend/jdbc/ParamMarker.java | 2 +- .../com/databend/jdbc/PresignContext.java | 9 +- .../java/com/databend/jdbc/QueryLiveness.java | 2 +- .../com/databend/jdbc/QueryRawStatement.java | 2 +- .../java/com/databend/jdbc/RawStatement.java | 2 +- .../databend/jdbc/RawStatementWrapper.java | 2 +- .../databend/jdbc/SetParamRawStatement.java | 2 +- .../databend/jdbc/StatementInfoWrapper.java | 2 +- .../java/com/databend/jdbc/StatementType.java | 2 +- .../java/com/databend/jdbc/StatementUtil.java | 2 +- .../jdbc/cloud/DatabendPresignClient.java | 8 +- .../examples/DatabendConnectionFactory.java | 16 +- .../jdbc/examples/DatabendConnectionPool.java | 10 +- .../com/databend/jdbc/TestBasicDriver.java | 8 +- .../java/com/databend/jdbc/TestCopyInto.java | 6 +- .../jdbc/TestDatabendDatabaseMetaData.java | 2 +- .../databend/jdbc/TestDatabendDriverUri.java | 13 +- .../TestDatabendPresignClient.java | 8 +- .../com/databend/jdbc/TestFileTransfer.java | 18 +- .../java/com/databend/jdbc/TestHeartbeat.java | 2 +- .../java/com/databend/jdbc/TestMultiHost.java | 2 +- .../databend/jdbc/TestPrepareStatement.java | 4 +- .../com/databend/jdbc/TestPresignContext.java | 13 +- 48 files changed, 1537 insertions(+), 1454 deletions(-) create mode 100644 databend-jdbc/src/main/java/com/databend/jdbc/DatabendConnectionImpl.java rename databend-jdbc/src/main/java/com/databend/jdbc/{constant => }/DatabendConstant.java (91%) rename databend-jdbc/src/test/java/com/databend/jdbc/{cloud => }/TestDatabendPresignClient.java (85%) diff --git a/README.md b/README.md index 3584bd8d..ee5e9d73 100644 --- a/README.md +++ b/README.md @@ -159,7 +159,7 @@ The following example demonstrates how to upload data and load it into a table: ```java // 1. Upload a file to the internal stage Connection conn = DriverManager.getConnection("jdbc:databend://localhost:8000"); -FileTransferAPI api = conn.unwrap(DatabendConnection.class); +DatabendConnection api = conn.unwrap(DatabendConnection.class); FileInputStream fileStream = new FileInputStream("data.csv"); api.uploadStream( diff --git a/databend-jdbc/src/main/java/com/databend/jdbc/Capability.java b/databend-jdbc/src/main/java/com/databend/jdbc/Capability.java index cced8eef..37ad0608 100644 --- a/databend-jdbc/src/main/java/com/databend/jdbc/Capability.java +++ b/databend-jdbc/src/main/java/com/databend/jdbc/Capability.java @@ -2,7 +2,7 @@ import com.vdurmont.semver4j.Semver; -public class Capability { +class Capability { private final boolean streamingLoad; private final boolean heartbeat; public Capability(Semver ver) { diff --git a/databend-jdbc/src/main/java/com/databend/jdbc/ConnectionProperties.java b/databend-jdbc/src/main/java/com/databend/jdbc/ConnectionProperties.java index ef45f260..10abeea8 100644 --- a/databend-jdbc/src/main/java/com/databend/jdbc/ConnectionProperties.java +++ b/databend-jdbc/src/main/java/com/databend/jdbc/ConnectionProperties.java @@ -9,7 +9,7 @@ import java.util.Set; // all possible JDBC properties options currently supported by databend driver -public final class ConnectionProperties { +final class ConnectionProperties { public static final ConnectionProperty USER = new User(); public static final ConnectionProperty PASSWORD = new Password(); public static final ConnectionProperty SSL = new Ssl(); diff --git a/databend-jdbc/src/main/java/com/databend/jdbc/DatabendClientLoadBalancingPolicy.java b/databend-jdbc/src/main/java/com/databend/jdbc/DatabendClientLoadBalancingPolicy.java index 0fa59b17..6f799585 100644 --- a/databend-jdbc/src/main/java/com/databend/jdbc/DatabendClientLoadBalancingPolicy.java +++ b/databend-jdbc/src/main/java/com/databend/jdbc/DatabendClientLoadBalancingPolicy.java @@ -3,7 +3,7 @@ import java.net.URI; import java.util.List; -public class DatabendClientLoadBalancingPolicy { +class DatabendClientLoadBalancingPolicy { static class DisabledPolicy extends DatabendClientLoadBalancingPolicy { @Override public String toString() { diff --git a/databend-jdbc/src/main/java/com/databend/jdbc/DatabendColumnInfo.java b/databend-jdbc/src/main/java/com/databend/jdbc/DatabendColumnInfo.java index eda2d51b..2d04b790 100644 --- a/databend-jdbc/src/main/java/com/databend/jdbc/DatabendColumnInfo.java +++ b/databend-jdbc/src/main/java/com/databend/jdbc/DatabendColumnInfo.java @@ -10,7 +10,7 @@ import static java.util.Objects.requireNonNull; -public class DatabendColumnInfo { +class DatabendColumnInfo { private static final int VARBINARY_MAX = 1024 * 1024 * 1024; // current longest time zone is 32 private static final int TIME_ZONE_MAX = 40; diff --git a/databend-jdbc/src/main/java/com/databend/jdbc/DatabendConnection.java b/databend-jdbc/src/main/java/com/databend/jdbc/DatabendConnection.java index 1e60854c..25148d3a 100644 --- a/databend-jdbc/src/main/java/com/databend/jdbc/DatabendConnection.java +++ b/databend-jdbc/src/main/java/com/databend/jdbc/DatabendConnection.java @@ -1,1356 +1,75 @@ package com.databend.jdbc; -import com.databend.client.*; - -import static com.databend.client.JsonCodec.jsonCodec; -import com.databend.jdbc.annotation.NotImplemented; -import com.databend.jdbc.cloud.DatabendCopyParams; -import com.databend.jdbc.cloud.DatabendPresignClient; -import com.databend.jdbc.cloud.DatabendPresignClientV1; -import com.databend.jdbc.exception.DatabendFailedToPingException; -import com.fasterxml.jackson.core.JsonProcessingException; -import com.fasterxml.jackson.databind.JsonNode; -import com.fasterxml.jackson.databind.ObjectMapper; -import com.vdurmont.semver4j.Semver; -import okhttp3.*; -import okio.BufferedSink; -import okio.Okio; -import okio.Source; - -import java.io.ByteArrayInputStream; -import java.io.ByteArrayOutputStream; -import java.io.File; -import java.io.IOException; import java.io.InputStream; -import java.net.ConnectException; -import java.net.URI; -import java.nio.charset.StandardCharsets; -import java.sql.Array; -import java.sql.Blob; -import java.sql.CallableStatement; -import java.sql.Clob; -import java.sql.Connection; -import java.sql.DatabaseMetaData; -import java.sql.NClob; -import java.sql.PreparedStatement; -import java.sql.ResultSet; -import java.sql.SQLClientInfoException; import java.sql.SQLException; -import java.sql.SQLFeatureNotSupportedException; -import java.sql.SQLWarning; -import java.sql.SQLXML; -import java.sql.Savepoint; -import java.sql.Statement; -import java.sql.Struct; -import java.time.Instant; -import java.util.*; -import java.util.concurrent.ConcurrentHashMap; -import java.util.concurrent.Executor; -import java.util.concurrent.Executors; -import java.util.concurrent.ExecutorService; -import java.util.concurrent.ScheduledFuture; -import java.util.concurrent.ScheduledExecutorService; -import java.util.concurrent.atomic.AtomicBoolean; -import java.util.concurrent.atomic.AtomicReference; -import java.util.function.Consumer; -import java.util.logging.FileHandler; -import java.util.logging.Level; -import java.util.logging.Logger; -import java.util.logging.SimpleFormatter; -import java.util.zip.GZIPOutputStream; - -import static com.databend.client.ClientSettings.*; -import static com.databend.client.DatabendClientV1.MEDIA_TYPE_JSON; -import static com.databend.client.DatabendClientV1.USER_AGENT_VALUE; -import static com.google.common.base.Preconditions.checkNotNull; -import static com.google.common.base.Preconditions.checkState; -import static java.net.URI.create; -import static java.util.Objects.requireNonNull; -import static java.util.concurrent.TimeUnit.MILLISECONDS; - - -public class DatabendConnection implements Connection, FileTransferAPI, Consumer { - private static final Logger logger = Logger.getLogger(DatabendConnection.class.getPackage().getName()); - public static final String STREAMING_LOAD_PATH = "/v1/streaming_load"; - public static final String LOGIN_PATH = "/v1/session/login"; - public static final String LOGOUT_PATH = "/v1/session/logout"; - public static final String HEARTBEAT_PATH = "/v1/session/heartbeat"; - private static final ObjectMapper objectMapper = new ObjectMapper(); - private static final JsonCodec SESSION_JSON_CODEC = jsonCodec(DatabendSession.class); - - private final AtomicBoolean closed = new AtomicBoolean(); - private final AtomicBoolean autoCommit = new AtomicBoolean(true); - private final URI httpUri; - private final AtomicReference schema = new AtomicReference<>(); - private final OkHttpClient httpClient; - private final ConcurrentHashMap statements = new ConcurrentHashMap<>(); - private final DatabendDriverUri driverUri; - private boolean autoDiscovery; - private final AtomicReference session = new AtomicReference<>(); - - private String routeHint = ""; - private final AtomicReference lastNodeID = new AtomicReference<>(); - private Semver serverVersion = null; - private Capability serverCapability = null; - - static volatile ExecutorService heartbeatScheduler = null; - private final HeartbeatManager heartbeatManager = new HeartbeatManager(); - - private void initializeFileHandler() { - if (this.debug()) { - File file = new File("databend-jdbc-debug.log"); - if (!file.canWrite()) { - logger.warning("No write access to file: " + file.getAbsolutePath()); - return; - } - try { - // 2GB，Integer.MAX_VALUE - System.setProperty("java.util.logging.FileHandler.limit", "2147483647"); - System.setProperty("java.util.logging.FileHandler.count", "200"); - // Enable log file reuse - System.setProperty("java.util.logging.FileHandler.append", "true"); - FileHandler fileHandler= new FileHandler(file.getAbsolutePath(), Integer.parseInt(System.getProperty("java.util.logging.FileHandler.limit")), - Integer.parseInt(System.getProperty("java.util.logging.FileHandler.count")), true); - fileHandler.setLevel(Level.ALL); - fileHandler.setFormatter(new SimpleFormatter()); - logger.addHandler(fileHandler); - } catch (Exception e) { - throw new RuntimeException("Failed to create FileHandler", e); - } - } - } - - - DatabendConnection(DatabendDriverUri uri, OkHttpClient httpClient) throws SQLException { - requireNonNull(uri, "uri is null"); - // only used for presign url on non-object storage, which mainly served for demo pupose. - // TODO: may also add query id and load balancing on the part. - this.httpUri = uri.getUri(); - this.httpClient = httpClient; - this.driverUri = uri; - this.schema.set(uri.getDatabase()); - this.routeHint = randRouteHint(); - // it maybe closed due to unsupported server versioning. - this.autoDiscovery = uri.autoDiscovery(); - DatabendSession session = new DatabendSession.Builder().setDatabase(this.getSchema()).setSettings(uri.getSessionSettings()).build(); - this.setSession(session); - - initializeFileHandler(); - this.login(); - } - - public Semver getServerVersion() { - return this.serverVersion; - } - - public Capability getServerCapability() { - return this.serverCapability; - } - - private void login() throws SQLException { - RetryPolicy retryPolicy = new RetryPolicy(true, true); - - HashMap headers = new HashMap<>(); - headers.put("Accept", "application/json"); - headers.put("Content-Type", "application/json"); - try { - LoginRequest req = new LoginRequest(); - req.database = this.getSchema(); - req.settings = this.driverUri.getSessionSettings(); - String bodyString = objectMapper.writeValueAsString(req); - RequestBody requestBody= RequestBody.create(MEDIA_TYPE_JSON, bodyString); - - ResponseWithBody response = requestHelper(LOGIN_PATH, "post", requestBody, headers, retryPolicy); - // old server do not support this API - if (response.response.code() != 400) { - String version = objectMapper.readTree(response.body).get("version").asText(); - if (version != null) { - this.serverVersion = new Semver(version); - this.serverCapability = new Capability(this.serverVersion); - } - } - } catch(JsonProcessingException e){ - throw new RuntimeException(e); - } - } - - public static String randRouteHint() { - String charset = "abcdef0123456789"; - Random rand = new Random(); - StringBuilder sb = new StringBuilder(16); - for (int i = 0; i < 16; i++) { - sb.append(charset.charAt(rand.nextInt(charset.length()))); - } - return sb.toString(); - } - - private static final char SPECIAL_CHAR = '#'; - - public static String uriRouteHint(String URI) { - // Encode the URI using Base64 - String encodedUri = Base64.getEncoder().encodeToString(URI.getBytes()); - - // Append the special character - return encodedUri + SPECIAL_CHAR; - } - - public static URI parseRouteHint(String routeHint) { - if (routeHint == null || routeHint.isEmpty()) { - return null; - } - try { - if (routeHint.charAt(routeHint.length() - 1) != SPECIAL_CHAR) { - return null; - } - // Remove the special character - String encodedUri = routeHint.substring(0, routeHint.length() - 1); - - // Decode the Base64 string - byte[] decodedBytes = Base64.getDecoder().decode(encodedUri); - String decodedUri = new String(decodedBytes); - - return create(decodedUri); - } catch (Exception e) { - logger.log(Level.FINE, "Failed to parse route hint: " + routeHint, e); - return null; - } - } - - - private static void checkResultSet(int resultSetType, int resultSetConcurrency) - throws SQLFeatureNotSupportedException { - if (resultSetType != ResultSet.TYPE_FORWARD_ONLY) { - throw new SQLFeatureNotSupportedException("Result set type must be TYPE_FORWARD_ONLY"); - } - if (resultSetConcurrency != ResultSet.CONCUR_READ_ONLY) { - throw new SQLFeatureNotSupportedException("Result set concurrency must be CONCUR_READ_ONLY"); - } - } - - public static String getCopyIntoSql(String database, DatabendCopyParams params) { - StringBuilder sb = new StringBuilder(); - sb.append("COPY INTO "); - if (database != null) { - sb.append(database).append("."); - } - sb.append(params.getDatabaseTableName()).append(" "); - sb.append("FROM "); - sb.append(params.getDatabendStage().toString()); - sb.append(" "); - sb.append(params); - return sb.toString(); - } - - public DatabendSession getSession() { - return this.session.get(); - } - - public boolean inActiveTransaction() { - if (this.session.get() == null) { - return false; - } - return this.session.get().inActiveTransaction(); - } - - public void setSession(DatabendSession session) { - if (session == null) { - return; - } - this.session.set(session); - } - - public OkHttpClient getHttpClient() { - return httpClient; - } - - @Override - public Statement createStatement() - throws SQLException { - return doCreateStatement(); - } - - private DatabendStatement doCreateStatement() throws SQLException { - checkOpen(); - DatabendStatement statement = new DatabendStatement(this, this::unregisterStatement); - registerStatement(statement); - return statement; - } - - synchronized private void registerStatement(DatabendStatement statement) { - checkState(statements.put(statement, true) == null, "Statement is already registered"); - } - - synchronized private void unregisterStatement(DatabendStatement statement) { - checkNotNull(statements.remove(statement), "Statement is not registered"); - } - - @Override - public PreparedStatement prepareStatement(String s) - throws SQLException { - - return this.prepareStatement(s, ResultSet.TYPE_FORWARD_ONLY, ResultSet.CONCUR_READ_ONLY); - } - - @Override - public CallableStatement prepareCall(String s) - throws SQLException { - throw new SQLFeatureNotSupportedException("prepareCall"); - } - - @Override - public String nativeSQL(String sql) - throws SQLException { - checkOpen(); - return sql; - } - - private void checkOpen() - throws SQLException { - if (isClosed()) { - throw new SQLException("Connection is closed"); - } - } - - @Override - public void commit() - throws SQLException { - checkOpen(); - try { - this.startQuery("commit"); - } catch (SQLException e) { - throw new SQLException("Failed to commit", e); - } - } - - @Override - public boolean getAutoCommit() - throws SQLException { - checkOpen(); - return autoCommit.get(); - } - - @Override - public void setAutoCommit(boolean b) - throws SQLException { - this.session.get().setAutoCommit(b); - autoCommit.set(b); - } - - @Override - public void rollback() - throws SQLException { - checkOpen(); - try { - this.startQuery("rollback"); - } catch (SQLException e) { - throw new SQLException("Failed to rollback", e); - } - } - - @Override - public void close() - throws SQLException { - for (Statement stmt : statements.keySet()) { - stmt.close(); - } - logout(); - } - - @Override - public boolean isClosed() - throws SQLException { - return closed.get(); - } - - @Override - public DatabaseMetaData getMetaData() - throws SQLException { - return new DatabendDatabaseMetaData(this); - } - - @Override - public boolean isReadOnly() - throws SQLException { - return false; - } - - @Override - public void setReadOnly(boolean b) - throws SQLException { - - } - - @Override - public String getCatalog() - throws SQLException { - return null; - } - - @Override - public void setCatalog(String s) - throws SQLException { - - } - - @Override - public int getTransactionIsolation() - throws SQLException { - return Connection.TRANSACTION_NONE; - } - - @Override - public void setTransactionIsolation(int i) - throws SQLException { - - } - - @Override - public SQLWarning getWarnings() - throws SQLException { - return null; - } - - @Override - public void clearWarnings() - throws SQLException { - - } - - @Override - public Statement createStatement(int resultSetType, int resultSetConcurrency) - throws SQLException { - checkResultSet(resultSetType, resultSetConcurrency); - return createStatement(); - } - - @Override - public PreparedStatement prepareStatement(String s, int i, int i1) - throws SQLException { - DatabendPreparedStatement statement = new DatabendPreparedStatement(this, this::unregisterStatement, s); - registerStatement(statement); - return statement; - } - - @Override - public CallableStatement prepareCall(String s, int i, int i1) - throws SQLException { - throw new SQLFeatureNotSupportedException("prepareCall"); - } - - @Override - public Map> getTypeMap() - throws SQLException { - throw new SQLFeatureNotSupportedException("getTypeMap"); - } - - @Override - public void setTypeMap(Map> map) - throws SQLException { - throw new SQLFeatureNotSupportedException("setTypeMap"); - } - @Override - public int getHoldability() throws SQLException { - return 0; - } - public int getMaxFailoverRetries() { - return this.driverUri.getMaxFailoverRetry(); - } - - @Override - @NotImplemented - public void setHoldability(int holdability) throws SQLException { - // No support for transaction - } - - @Override - public Savepoint setSavepoint() - throws SQLException { - throw new SQLFeatureNotSupportedException("setSavepoint"); - } - - @Override - public Savepoint setSavepoint(String s) - throws SQLException { - throw new SQLFeatureNotSupportedException("setSavepoint"); - } - - @Override - public void rollback(Savepoint savepoint) - throws SQLException { - throw new SQLFeatureNotSupportedException("rollback"); - - } - - @Override - public void releaseSavepoint(Savepoint savepoint) - throws SQLException { - throw new SQLFeatureNotSupportedException("releaseSavepoint"); - - } - - @Override - public Statement createStatement(int resultSetType, int resultSetConcurrency, int resultSetHoldability) - throws SQLException { -// checkHoldability(resultSetHoldability); - return createStatement(resultSetType, resultSetConcurrency); - } - - @Override - public PreparedStatement prepareStatement(String sql, int resultSetType, int resultSetConcurrency, int resultSetHoldability) - throws SQLException { -// checkHoldability(resultSetHoldability); - return prepareStatement(sql, resultSetType, resultSetConcurrency); - } - - @Override - public CallableStatement prepareCall(String sql, int resultSetType, int resultSetConcurrency, int resultSetHoldability) - throws SQLException { - return prepareCall(sql, resultSetType, resultSetConcurrency); - } - - @Override - public PreparedStatement prepareStatement(String s, int autoGeneratedKeys) - throws SQLException { - return prepareStatement(s); - } - - @Override - public PreparedStatement prepareStatement(String s, int[] ints) - throws SQLException { - throw new SQLFeatureNotSupportedException("prepareStatement"); - } - - @Override - public PreparedStatement prepareStatement(String s, String[] strings) - throws SQLException { - throw new SQLFeatureNotSupportedException("prepareStatement"); - } - - @Override - public Clob createClob() - throws SQLException { - throw new SQLFeatureNotSupportedException("createClob"); - } - - @Override - public Blob createBlob() - throws SQLException { - throw new SQLFeatureNotSupportedException("createBlob"); - } - - @Override - public NClob createNClob() - throws SQLException { - throw new SQLFeatureNotSupportedException("createNClob"); - } - - @Override - public SQLXML createSQLXML() - throws SQLException { - throw new SQLFeatureNotSupportedException("createSQLXML"); - } - - @Override - public boolean isValid(int i) - throws SQLException { - return !isClosed(); - } +/** + * The SnowflakeConnection interface contains Snowflake-specific methods. + * providing specialized methods for interacting with Databend stages and loading data efficiently. + *

+ * This interface extends standard JDBC connection capabilities to support streaming + * operations for uploading/downloading files to/from internal stages, as well as direct + * streaming data loading into target tables. Ideal for handling large files or continuous + * data streams in Databend. + *

+ */ +public interface DatabendConnection { - @Override - public void setClientInfo(String s, String s1) - throws SQLClientInfoException { - - } - - @Override - public String getClientInfo(String s) - throws SQLException { - return null; - } - - @Override - public Properties getClientInfo() - throws SQLException { - return null; - } - - @Override - public void setClientInfo(Properties properties) - throws SQLClientInfoException { - - } - - @Override - public Array createArrayOf(String s, Object[] objects) - throws SQLException { - throw new SQLFeatureNotSupportedException("createArrayOf"); - } - - @Override - public Struct createStruct(String s, Object[] objects) - throws SQLException { - throw new SQLFeatureNotSupportedException("createStruct"); - } - - @Override - public String getSchema() - throws SQLException { - checkOpen(); - return schema.get(); - } - - @Override - public void setSchema(String schema) - throws SQLException { - checkOpen(); - this.schema.set(schema); - this.startQuery("use " + schema); - } - - @Override - public void abort(Executor executor) - throws SQLException { - close(); - } - - @Override - public void setNetworkTimeout(Executor executor, int i) - throws SQLException { - - } - - @Override - public int getNetworkTimeout() - throws SQLException { - return 0; - } - - @Override - public T unwrap(Class aClass) - throws SQLException { - if (isWrapperFor(aClass)) { - return (T) this; - } - throw new SQLException("No wrapper for " + aClass); - } - - @Override - public boolean isWrapperFor(Class aClass) - throws SQLException { - return aClass.isInstance(this); - } - - public boolean presignedUrlDisabled() { - return this.driverUri.presignedUrlDisabled(); - } - - public boolean copyPurge() { - return this.driverUri.copyPurge(); - } - - public boolean isAutoDiscovery() { - return this.autoDiscovery; - } - - public String warehouse() { - return this.driverUri.getWarehouse(); - } - - public Boolean strNullAsNull() { - return this.driverUri.getStrNullAsNull(); - } - - public Boolean useVerify() { - return this.driverUri.getUseVerify(); - } - - public Boolean debug() { - return this.driverUri.getDebug(); - } - - public String tenant() { - return this.driverUri.getTenant(); - } - - public String nullDisplay() { - return this.driverUri.nullDisplay(); - } - - public String binaryFormat() { - return this.driverUri.binaryFormat(); - } - - public PaginationOptions getPaginationOptions() { - PaginationOptions.Builder builder = PaginationOptions.builder(); - builder.setWaitTimeSecs(this.driverUri.getWaitTimeSecs()); - builder.setMaxRowsInBuffer(this.driverUri.getMaxRowsInBuffer()); - builder.setMaxRowsPerPage(this.driverUri.getMaxRowsPerPage()); - return builder.build(); - } - - public URI getURI() { - return this.httpUri; - } - - private String buildUrlWithQueryRequest(ClientSettings settings, String querySql) { - QueryRequest req = QueryRequest.builder() - .setSession(settings.getSession()) - .setStageAttachment(settings.getStageAttachment()) - .setPaginationOptions(settings.getPaginationOptions()) - .setSql(querySql) - .build(); - String reqString = req.toString(); - if (reqString == null || reqString.isEmpty()) { - throw new IllegalArgumentException("Invalid request: " + req); - } - return reqString; - } + /** + * Enumeration of available loading strategies for streaming data into tables. + */ + enum LoadMethod { + /** + * Load strategy that first uploads the stream to a Databend internal stage, + * then loads the data from the stage into the target table. + * Suitable for large files or scenarios requiring temporary storage. + */ + STAGE, - public void PingDatabendClientV1() throws SQLException { - try (Statement statement = this.createStatement()) { - statement.execute("select 1"); - ResultSet r = statement.getResultSet(); - while (r.next()) { - } - } catch (SQLException e) { - throw new DatabendFailedToPingException(String.format("failed to ping databend server: %s", e.getMessage())); - } - } - @Override - public void accept(DatabendSession session) { - setSession(session); + /** + * Direct streaming strategy that loads data directly into the target table + * without intermediate stage storage. + * Optimized for real-time or small-to-medium size data streams. + */ + STREAMING } /** - * Retry executing a query in case of connection errors. fail over mechanism is used to retry the query when connect error occur - * It will find next target host based on configured Load balancing Policy. + * Upload inputStream to the databend internal stage, the data would be uploaded as one file with no split. + * Caller should close the input stream after the upload is done. * - * @param sql The SQL statement to execute. - * @param attach The stage attachment to use for the query. - * @return A DatabendClient instance representing the successful query execution. - * @throws SQLException If the query fails after retrying the specified number of times. - * @see DatabendClientLoadBalancingPolicy + * @param stageName the stage which receive uploaded file + * @param destPrefix the prefix of the file name in the stage + * @param inputStream the input stream of the file + * @param destFileName the destination file name in the stage + * @param fileSize the file size in the stage + * @param compressData whether to compress the data + * @throws SQLException failed to upload input stream */ - DatabendClient startQueryWithFailover(String sql, StageAttachment attach) throws SQLException { - int maxRetries = getMaxFailoverRetries(); - SQLException lastException = null; - - for (int attempt = 0; attempt <= maxRetries; attempt++) { - try { - String queryId = UUID.randomUUID().toString().replace("-", ""); - String candidateHost = selectHostForQuery(queryId); - - // configure the client settings - ClientSettings.Builder sb = this.makeClientSettings(queryId, candidateHost); - if (attach != null) { - sb.setStageAttachment(attach); - } - ClientSettings settings = sb.build(); - - logger.log(Level.FINE, "execute query #{0}: SQL: {1} host: {2}", - new Object[]{attempt + 1, sql, settings.getHost()}); - - // need to retry the auto discovery in case of connection error - if (this.autoDiscovery) { - tryAutoDiscovery(httpClient, settings); - } - - return new DatabendClientV1(httpClient, sql, settings, this, lastNodeID); - } catch (Exception e) { - // handle the exception and retry the query - if (shouldRetryException(e) && attempt < maxRetries) { - lastException = wrapException("query failed", sql, e); - try { - // back off retry - Thread.sleep(Math.min(100 * (1 << attempt), 5000)); - } catch (InterruptedException ie) { - Thread.currentThread().interrupt(); - throw wrapException("query interrupt", sql, ie); - } - } else { - // throw the exception - if (e instanceof SQLException) { - throw (SQLException) e; - } else { - throw wrapException("Query failed，no need to retry", sql, e); - } - } - } - } - - throw new SQLException("after" + maxRetries + "times retry and failed: SQL: " + sql, lastException); - } - - private boolean shouldRetryException(Exception e) { - Throwable cause = e.getCause(); - // connection error - if (cause instanceof ConnectException) { - return true; - } - - if (e instanceof RuntimeException) { - String message = e.getMessage(); - return message != null && ( - message.contains("520") || - message.contains("timeout") || - message.contains("retry") - ); - } - - return false; - } - - private String selectHostForQuery(String queryId) { - String candidateHost = this.driverUri.getUri(queryId).toString(); + void uploadStream(String stageName, String destPrefix, InputStream inputStream, String destFileName, long fileSize, boolean compressData) throws SQLException; - if (!inActiveTransaction()) { - this.routeHint = uriRouteHint(candidateHost); - } - - if (this.routeHint != null && !this.routeHint.isEmpty()) { - URI uri = parseRouteHint(this.routeHint); - if (uri != null) { - candidateHost = uri.toString(); - } - } - - return candidateHost; - } - - private SQLException wrapException(String prefix, String sql, Exception e) { - String message = prefix + ": SQL: " + sql; - if (e.getMessage() != null) { - message += " - " + e.getMessage(); - } - if (e.getCause() != null) { - message += " (Reason: " + e.getCause().getMessage() + ")"; - } - return new SQLException(message, e); - } /** - * Try to auto discovery the databend nodes it will log exceptions when auto discovery failed and not affect real query execution + * Download a file from the databend internal stage, the data would be downloaded as one file with no split. * - * @param client the http client to query on - * @param settings the client settings to use + * @param stageName the stage which contains the file + * @param sourceFileName the file name in the stage + * @param decompress whether to decompress the data + * @return the input stream of the file + * @throws SQLException failed to download input stream */ - void tryAutoDiscovery(OkHttpClient client, ClientSettings settings) { - if (this.autoDiscovery) { - if (this.driverUri.enableMock()) { - settings.getAdditionalHeaders().put("~mock.unsupported.discovery", "true"); - } - DatabendNodes nodes = this.driverUri.getNodes(); - if (nodes != null && nodes.needDiscovery()) { - try { - nodes.discoverUris(client, settings); - } catch (UnsupportedOperationException e) { - logger.log(Level.WARNING, "Current Query Node do not support auto discovery, close the functionality: " + e.getMessage()); - this.autoDiscovery = false; - } catch (Exception e) { - logger.log(Level.FINE, "Error auto discovery: " + " cause: " + e.getCause() + " message: " + e.getMessage()); - } - } - } - - } - - DatabendClient startQuery(String sql) throws SQLException { - return startQuery(sql, null); - } - - DatabendClient startQuery(String sql, StageAttachment attach) throws SQLException { - DatabendClient client = startQueryWithFailover(sql, attach); - Long timeout = client.getResults().getResultTimeoutSecs(); - if (timeout != null && timeout != 0) { - heartbeatManager.onStartQuery(timeout); - } - return client; - } - - private ClientSettings.Builder makeClientSettings(String queryID, String host) { - PaginationOptions options = getPaginationOptions(); - Map additionalHeaders = setAdditionalHeaders(); - additionalHeaders.put(X_Databend_Query_ID, queryID); - return new Builder(). - setSession(this.session.get()). - setHost(host). - setQueryTimeoutSecs(this.driverUri.getQueryTimeout()). - setConnectionTimeout(this.driverUri.getConnectionTimeout()). - setSocketTimeout(this.driverUri.getSocketTimeout()). - setPaginationOptions(options). - setAdditionalHeaders(additionalHeaders); - } - - private Map setAdditionalHeaders() { - Map additionalHeaders = new HashMap<>(); - - DatabendSession session = this.getSession(); - String warehouse = null; - if (session != null ) { - Map settings = session.getSettings(); - if (settings != null) { - warehouse = settings.get("warehouse"); - } - } - if (warehouse == null && !this.driverUri.getWarehouse().isEmpty()) { - warehouse = this.driverUri.getWarehouse(); - } - if (warehouse!=null) { - additionalHeaders.put(DatabendWarehouseHeader, warehouse); - } - - if (!this.driverUri.getTenant().isEmpty()) { - additionalHeaders.put(DatabendTenantHeader, this.driverUri.getTenant()); - } - if (!this.routeHint.isEmpty()) { - additionalHeaders.put(X_DATABEND_ROUTE_HINT, this.routeHint); - } - additionalHeaders.put("User-Agent", USER_AGENT_VALUE); - return additionalHeaders; - } - + InputStream downloadStream(String stageName, String sourceFileName) throws SQLException; /** - * Method to put data from a stream at a stage location. The data will be uploaded as one file. No - * splitting is done in this method. + * Loads data from an input stream directly into a target Databend table using the specified SQL command. + * Supports two loading strategies via {@link LoadMethod}. * - *

Stream size must match the total size of data in the input stream unless compressData - * parameter is set to true. - * - *

caller is responsible for passing the correct size for the data in the stream and releasing - * the inputStream after the method is called. - * - *

Note this method is deprecated since streamSize is not required now. Keep the function - * signature for backward compatibility - * - * @param stageName stage name: e.g. ~ or table name or stage name - * @param destPrefix path prefix under which the data should be uploaded on the stage - * @param inputStream input stream from which the data will be uploaded - * @param destFileName destination file name to use - * @param fileSize data size in the stream - * @throws SQLException failed to put data from a stream at stage + * @param sql SQL command with Databend's load syntax: + * {@code INSERT INTO [()] FROM @_databend_load [file_format=(...)]} + * @param inputStream Input stream containing the data to load into the table + * @param fileSize Size of the data (in bytes) to be loaded + * @param loadMethod Loading strategy ({@link LoadMethod#STAGE} or {@link LoadMethod#STREAMING}) + * @return Number of rows successfully loaded into the target table + * @throws SQLException If the load operation fails (e.g., invalid SQL, stream errors, or data format issues) */ - @Override - public void uploadStream(String stageName, String destPrefix, InputStream inputStream, String destFileName, long fileSize, boolean compressData) - throws SQLException { - /* - remove / in the end of stage name - remove / in the beginning of destPrefix and end of destPrefix - */ - String s; - if (stageName == null) { - s = "~"; - } else { - s = stageName.replaceAll("/$", ""); - } - String p = destPrefix.replaceAll("^/", "").replaceAll("/$", ""); - String dest = p + "/" + destFileName; - try { - InputStream dataStream = inputStream; - if (compressData) { - // Wrap the input stream with a GZIPOutputStream for compression - ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream(); - try (GZIPOutputStream gzipOutputStream = new GZIPOutputStream(byteArrayOutputStream)) { - byte[] buffer = new byte[1024]; - int len; - while ((len = inputStream.read(buffer)) != -1) { - gzipOutputStream.write(buffer, 0, len); - } - } - dataStream = new ByteArrayInputStream(byteArrayOutputStream.toByteArray()); - // Update the file size to the compressed size - fileSize = byteArrayOutputStream.size(); - } - if (this.driverUri.presignedUrlDisabled()) { - DatabendPresignClient cli = new DatabendPresignClientV1(httpClient, this.httpUri.toString()); - cli.presignUpload(null, dataStream, s, p + "/", destFileName, fileSize, true); - } else { -// logger.log(Level.FINE, "presign to @" + s + "/" + dest); - long presignStartTime = System.nanoTime(); - PresignContext ctx = PresignContext.getPresignContext(this, PresignContext.PresignMethod.UPLOAD, s, dest); - long presignEndTime = System.nanoTime(); - if (this.debug()) { - logger.info("presign cost time: " + (presignEndTime - presignStartTime) / 1000000.0 + "ms"); - } - Headers h = ctx.getHeaders(); - String presignUrl = ctx.getUrl(); - DatabendPresignClient cli = new DatabendPresignClientV1(new OkHttpClient(), this.httpUri.toString()); - long uploadStartTime = System.nanoTime(); - cli.presignUpload(null, dataStream, h, presignUrl, fileSize, true); - long uploadEndTime = System.nanoTime(); - if (this.debug()) { - logger.info("upload cost time: " + (uploadEndTime - uploadStartTime) / 1000000.0 + "ms"); - } - } - } catch (RuntimeException | IOException e) { - logger.warning("failed to upload input stream, file size is:" + fileSize / 1024.0 + e.getMessage()); - throw new SQLException(e); - } - } - - @Override - public InputStream downloadStream(String stageName, String sourceFileName, boolean decompress) - throws SQLException { - String s = stageName.replaceAll("/$", ""); - DatabendPresignClient cli = new DatabendPresignClientV1(httpClient, this.httpUri.toString()); - try { - PresignContext ctx = PresignContext.getPresignContext(this, PresignContext.PresignMethod.DOWNLOAD, s, sourceFileName); - Headers h = ctx.getHeaders(); - String presignUrl = ctx.getUrl(); - return cli.presignDownloadStream(h, presignUrl); - } catch (RuntimeException e) { - throw new SQLException(e); - } - } - - @Override - public void copyIntoTable(String database, String tableName, DatabendCopyParams params) - throws SQLException { - DatabendCopyParams p = params == null ? DatabendCopyParams.builder().build() : params; - requireNonNull(p.getDatabaseTableName(), "tableName is null"); - requireNonNull(p.getDatabendStage(), "stage is null"); - String sql = getCopyIntoSql(database, p); - Statement statement = this.createStatement(); - statement.execute(sql); - ResultSet rs = statement.getResultSet(); - while (rs.next()) { - } - } - @Override - public int loadStreamToTable(String sql, InputStream inputStream, long fileSize, String loadMethod) throws SQLException { - loadMethod = loadMethod.toLowerCase(); - if (!"stage".equals(loadMethod) && !"streaming".equals(loadMethod)) { - throw new SQLException("invalid value for loadMethod(" + loadMethod + ") only accept \"stage\" or \" streaming\""); - } - - if (!this.serverCapability.streamingLoad()) { - throw new SQLException("please upgrade databend-query to >1.2.781 to use loadStreamToTable, current version=" + this.serverVersion); - } - - if (!sql.contains("@_databend_load")) { - throw new SQLException("invalid sql: must contain @_databend_load when used in loadStreamToTable "); - } - - if ("streaming".equals(loadMethod)) { - return streamingLoad(sql, inputStream, fileSize); - } else { - Instant now = Instant.now(); - long nanoTimestamp = now.getEpochSecond() * 1_000_000_000 + now.getNano(); - String fileName = String.valueOf(nanoTimestamp); - String location = "~/_databend_load/" + fileName; - sql = sql.replace("_databend_load", location); - uploadStream("~", "_databend_load", inputStream, fileName, fileSize, false); - Statement statement = this.createStatement(); - statement.execute(sql); - ResultSet rs = statement.getResultSet(); - while (rs.next()) { - } - return statement.getUpdateCount(); - } - } - - MultipartBody buildMultiPart(InputStream inputStream, long fileSize) { - RequestBody requestBody = new RequestBody() { - @Override - public MediaType contentType() { - return MediaType.parse("application/octet-stream"); - } - - @Override - public long contentLength() { - return fileSize; - } - - @Override - public void writeTo(BufferedSink sink) throws IOException { - try (Source source = Okio.source(inputStream)) { - sink.writeAll(source); - } - } - }; - return new MultipartBody.Builder() - .setType(MultipartBody.FORM) - .addFormDataPart( - "upload", - "java.io.InputStream", - requestBody - ).build(); - } - - int streamingLoad(String sql, InputStream inputStream, long fileSize) throws SQLException { - RetryPolicy retryPolicy = new RetryPolicy(true, true); - - try { - HashMap headers = new HashMap<>(); - DatabendSession session = this.session.get(); - if (session != null) { - String sessionString = objectMapper.writeValueAsString(session); - headers.put(DatabendQueryContextHeader, sessionString); - } - headers.put(DatabendSQLHeader, sql); - headers.put("Accept", "application/json"); - RequestBody requestBody = buildMultiPart(inputStream, fileSize); - ResponseWithBody response = requestHelper(STREAMING_LOAD_PATH, "put", requestBody, headers, retryPolicy); - JsonNode json = objectMapper.readTree(response.body); - JsonNode error = json.get("error"); - if (error != null) { - throw new SQLException("streaming load fail: code = " + error.get("code").asText() + ", message=" + error.get("message").asText()); - } - String base64 = response.response.headers().get(DatabendQueryContextHeader); - if (base64 != null) { - byte[] bytes = Base64.getUrlDecoder().decode(base64); - String str = new String(bytes, StandardCharsets.UTF_8); - try { - session = SESSION_JSON_CODEC.fromJson(str); - } catch(Exception e) { - throw new RuntimeException(e); - } - if (session != null) { - this.session.set(session); - } - } - JsonNode stats = json.get("stats"); - if (stats != null) { - int rows = stats.get("rows").asInt(-1); - if (rows != -1) { - return rows; - } - } - throw new SQLException("invalid response for " + STREAMING_LOAD_PATH + ": " + response.body); - } catch(JsonProcessingException e) { - throw new RuntimeException(e); - } - } - - void logout() throws SQLException { - DatabendSession session = this.session.get(); - if (session == null || !session.getNeedKeepAlive()) { - return; - } - RetryPolicy retryPolicy = new RetryPolicy(false, false); - RequestBody body = RequestBody.create(MEDIA_TYPE_JSON, "{}"); - requestHelper(LOGOUT_PATH, "post", body, new HashMap<>(), retryPolicy); - } - - - HttpUrl getUrl(String path) { - String host = this.driverUri.getUri().toString(); - HttpUrl url = HttpUrl.get(host); - return url.newBuilder().encodedPath(path).build(); - } - - ResponseWithBody sendRequestWithRetry(Request request, RetryPolicy retryPolicy, String path) throws SQLException { - String failReason = null; - - for (int j = 1; j <= 3; j++) { - try (Response response = httpClient.newCall(request).execute()) { - int code = response.code(); - if (code != 200) { - if (retryPolicy.shouldIgnore(code)) { - return new ResponseWithBody(response, ""); - } else { - failReason = "status code =" + response.code() + ", body = " + response.body().string(); - if (!retryPolicy.shouldRetry(code)) - break; - } - } else { - String body = response.body().string(); - return new ResponseWithBody(response, body); - } - } catch (IOException e) { - if (retryPolicy.shouldRetry(e)) { - if (failReason == null) { - failReason = e.getMessage(); - } - } else { - break; - } - } - if (j < 3) { - try { - MILLISECONDS.sleep(j * 100); - } catch (InterruptedException e2) { - Thread.currentThread().interrupt(); - return null; - } - } - } - throw new SQLException("Error accessing " + path + ": " + failReason); - } - - ResponseWithBody requestHelper(String path, String method, RequestBody body, Map headers, RetryPolicy retryPolicy) throws SQLException { - DatabendSession session = this.session.get(); - HttpUrl url = getUrl(path); - - Request.Builder builder = new Request.Builder().url(url); - this.setAdditionalHeaders().forEach(builder::addHeader); - if (headers != null) { - headers.forEach(builder::addHeader); - } - if (session.getNeedSticky()) { - builder.addHeader(ClientSettings.X_DATABEND_ROUTE_HINT, url.host()); - String lastNodeID = this.lastNodeID.get(); - if (lastNodeID != null) { - builder.addHeader(ClientSettings.X_DATABEND_STICKY_NODE, lastNodeID); - } - } - if ("post".equals(method)) { - builder = builder.post(body); - } else if ("put".equals(method)) { - builder = builder.put(body); - } else { - builder = builder.get(); - } - Request request = builder.build(); - return sendRequestWithRetry(request, retryPolicy, path); - } - - class HeartbeatManager implements Runnable { - private ScheduledFuture heartbeatFuture; - private long heartbeatIntervalMillis = 30000; - private long lastHeartbeatStartTimeMillis = 0; - private ScheduledExecutorService getScheduler() { - if (heartbeatScheduler == null) { - synchronized (HeartbeatManager.class) { - if (heartbeatScheduler == null) { - // create daemon thread so that it will not block JVM from exiting. - heartbeatScheduler = - Executors.newScheduledThreadPool( - 1, - runnable -> { - Thread thread = Executors.defaultThreadFactory().newThread(runnable); - thread.setName("heartbeat (" + thread.getId() + ")"); - thread.setDaemon(true); - return thread; - }); - } - } - } - return (ScheduledExecutorService) heartbeatScheduler; - } - - private void scheduleHeartbeat() { - long delay = Math.max(heartbeatIntervalMillis - (System.currentTimeMillis() - lastHeartbeatStartTimeMillis), 0); - heartbeatFuture = getScheduler().schedule(this, delay, MILLISECONDS); - } - - private ArrayList queryLiveness() { - ArrayList arr = new ArrayList<>(); - for (DatabendStatement stmt : statements.keySet()) { - QueryLiveness ql = stmt.queryLiveness(); - if (ql != null && !ql.stopped && ql.serverSupportHeartBeat) { - arr.add(ql); - } - } - return arr; - } - - private void doHeartbeat(ArrayList queryLivenesses ) { - long now = System.currentTimeMillis(); - lastHeartbeatStartTimeMillis = now; - Map> nodeToQueryID = new HashMap<>(); - Map queries = new HashMap<>(); - - for (QueryLiveness ql: queryLivenesses) { - if (now - ql.lastRequestTime.get() >= ql.resultTimeoutSecs * 1000 / 2) { - nodeToQueryID.computeIfAbsent(ql.nodeID, k -> new ArrayList<>()).add(ql.queryID); - queries.put(ql.queryID, ql); - } - } - if (nodeToQueryID.isEmpty()) { - return; - } - - Map map = new HashMap<>(); - map.put("node_to_queries", nodeToQueryID); - - try { - String body = objectMapper.writeValueAsString(map); - RequestBody requestBody = RequestBody.create(MEDIA_TYPE_JSON, body); - RetryPolicy retryPolicy = new RetryPolicy(true, false); - body = requestHelper(HEARTBEAT_PATH, "post", requestBody, null, retryPolicy).body; - JsonNode toRemove = objectMapper.readTree(body).get("queries_to_remove"); - if (toRemove.isArray()) { - for (JsonNode element : toRemove) { - String queryId = element.asText(); - queries.get(queryId).stopped = true; - } - } - } catch (JsonProcessingException e) { - logger.warning("fail to encode heartbeat body: " + e); - } catch (SQLException e) { - logger.warning("fail to send heartbeat: " + e); - } catch (Exception e) { - logger.warning("fail to send heartbeat: " + e); - throw new RuntimeException(e); - } - } - - public void onStartQuery(Long timeoutSecs) { - synchronized (DatabendConnection.this) { - if (timeoutSecs * 1000 / 4 < heartbeatIntervalMillis) { - heartbeatIntervalMillis = timeoutSecs * 1000 / 4; - if (heartbeatFuture != null) { - heartbeatFuture.cancel(false); - heartbeatFuture = null; - } - } - if (heartbeatFuture == null) { - scheduleHeartbeat(); - } - } - } - - @Override - public void run() { - ArrayList arr = queryLiveness(); - doHeartbeat(arr); - - heartbeatFuture = null; - synchronized (DatabendConnection.this) { - if (arr.size() > 0) { - if (heartbeatFuture == null) { - scheduleHeartbeat(); - } - } else { - heartbeatFuture = null; - } - } - } - } - - boolean isHeartbeatStopped() { - return heartbeatManager.heartbeatFuture == null; - } - - static class RetryPolicy { - boolean ignore404; - boolean retry503; - RetryPolicy(boolean ignore404, boolean retry503) { - this.ignore404 = ignore404; - this.retry503 = retry503; - } - - boolean shouldIgnore(int code) { - return ignore404 && code == 404; - } - - boolean shouldRetry(int code) { - return retry503 && (code == 502 || code == 503); - } - - boolean shouldRetry(IOException e) { - return e.getCause() instanceof ConnectException; - } - } - - static class ResponseWithBody { - public Response response; - public String body; - - ResponseWithBody(Response response, String body) { - this.response = response; - this.body = body; - } - } + int loadStreamToTable(String sql, InputStream inputStream, long fileSize, LoadMethod loadMethod) throws SQLException; } diff --git a/databend-jdbc/src/main/java/com/databend/jdbc/DatabendConnectionImpl.java b/databend-jdbc/src/main/java/com/databend/jdbc/DatabendConnectionImpl.java new file mode 100644 index 00000000..0c5db399 --- /dev/null +++ b/databend-jdbc/src/main/java/com/databend/jdbc/DatabendConnectionImpl.java @@ -0,0 +1,1358 @@ +package com.databend.jdbc; + +import com.databend.client.*; + +import static com.databend.client.JsonCodec.jsonCodec; +import com.databend.jdbc.annotation.NotImplemented; +import com.databend.jdbc.cloud.DatabendCopyParams; +import com.databend.jdbc.cloud.DatabendPresignClient; +import com.databend.jdbc.cloud.DatabendPresignClientV1; +import com.databend.jdbc.exception.DatabendFailedToPingException; +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.vdurmont.semver4j.Semver; +import okhttp3.*; +import okio.BufferedSink; +import okio.Okio; +import okio.Source; + +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.File; +import java.io.IOException; +import java.io.InputStream; +import java.net.ConnectException; +import java.net.URI; +import java.nio.charset.StandardCharsets; +import java.sql.Array; +import java.sql.Blob; +import java.sql.CallableStatement; +import java.sql.Clob; +import java.sql.Connection; +import java.sql.DatabaseMetaData; +import java.sql.NClob; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.sql.SQLClientInfoException; +import java.sql.SQLException; +import java.sql.SQLFeatureNotSupportedException; +import java.sql.SQLWarning; +import java.sql.SQLXML; +import java.sql.Savepoint; +import java.sql.Statement; +import java.sql.Struct; +import java.time.Instant; +import java.util.*; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.Executor; +import java.util.concurrent.Executors; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.ScheduledFuture; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Consumer; +import java.util.logging.FileHandler; +import java.util.logging.Level; +import java.util.logging.Logger; +import java.util.logging.SimpleFormatter; +import java.util.zip.GZIPOutputStream; + +import static com.databend.client.ClientSettings.*; +import static com.databend.client.DatabendClientV1.MEDIA_TYPE_JSON; +import static com.databend.client.DatabendClientV1.USER_AGENT_VALUE; +import static com.google.common.base.Preconditions.checkNotNull; +import static com.google.common.base.Preconditions.checkState; +import static java.net.URI.create; +import static java.util.Objects.requireNonNull; +import static java.util.concurrent.TimeUnit.MILLISECONDS; + + +public class DatabendConnectionImpl implements Connection, DatabendConnection, FileTransferAPI, Consumer { + private static final Logger logger = Logger.getLogger(DatabendConnectionImpl.class.getPackage().getName()); + private static final String STREAMING_LOAD_PATH = "/v1/streaming_load"; + private static final String LOGIN_PATH = "/v1/session/login"; + private static final String LOGOUT_PATH = "/v1/session/logout"; + private static final String HEARTBEAT_PATH = "/v1/session/heartbeat"; + private static final ObjectMapper objectMapper = new ObjectMapper(); + private static final JsonCodec SESSION_JSON_CODEC = jsonCodec(DatabendSession.class); + + private final AtomicBoolean closed = new AtomicBoolean(); + private final AtomicBoolean autoCommit = new AtomicBoolean(true); + private final URI httpUri; + private final AtomicReference schema = new AtomicReference<>(); + private final OkHttpClient httpClient; + private final ConcurrentHashMap statements = new ConcurrentHashMap<>(); + private final DatabendDriverUri driverUri; + private boolean autoDiscovery; + private final AtomicReference session = new AtomicReference<>(); + + private String routeHint = ""; + private final AtomicReference lastNodeID = new AtomicReference<>(); + private Semver serverVersion = null; + private Capability serverCapability = null; + + static volatile ExecutorService heartbeatScheduler = null; + private final HeartbeatManager heartbeatManager = new HeartbeatManager(); + + private void initializeFileHandler() { + if (this.debug()) { + File file = new File("databend-jdbc-debug.log"); + if (!file.canWrite()) { + logger.warning("No write access to file: " + file.getAbsolutePath()); + return; + } + try { + // 2GB，Integer.MAX_VALUE + System.setProperty("java.util.logging.FileHandler.limit", "2147483647"); + System.setProperty("java.util.logging.FileHandler.count", "200"); + // Enable log file reuse + System.setProperty("java.util.logging.FileHandler.append", "true"); + FileHandler fileHandler= new FileHandler(file.getAbsolutePath(), Integer.parseInt(System.getProperty("java.util.logging.FileHandler.limit")), + Integer.parseInt(System.getProperty("java.util.logging.FileHandler.count")), true); + fileHandler.setLevel(Level.ALL); + fileHandler.setFormatter(new SimpleFormatter()); + logger.addHandler(fileHandler); + } catch (Exception e) { + throw new RuntimeException("Failed to create FileHandler", e); + } + } + } + + + DatabendConnectionImpl(DatabendDriverUri uri, OkHttpClient httpClient) throws SQLException { + requireNonNull(uri, "uri is null"); + // only used for presign url on non-object storage, which mainly served for demo pupose. + // TODO: may also add query id and load balancing on the part. + this.httpUri = uri.getUri(); + this.httpClient = httpClient; + this.driverUri = uri; + this.schema.set(uri.getDatabase()); + this.routeHint = randRouteHint(); + // it maybe closed due to unsupported server versioning. + this.autoDiscovery = uri.autoDiscovery(); + DatabendSession session = new DatabendSession.Builder().setDatabase(this.getSchema()).setSettings(uri.getSessionSettings()).build(); + this.setSession(session); + + initializeFileHandler(); + this.login(); + } + + Semver getServerVersion() { + return this.serverVersion; + } + + Capability getServerCapability() { + return this.serverCapability; + } + + private void login() throws SQLException { + RetryPolicy retryPolicy = new RetryPolicy(true, true); + + HashMap headers = new HashMap<>(); + headers.put("Accept", "application/json"); + headers.put("Content-Type", "application/json"); + try { + LoginRequest req = new LoginRequest(); + req.database = this.getSchema(); + req.settings = this.driverUri.getSessionSettings(); + String bodyString = objectMapper.writeValueAsString(req); + RequestBody requestBody= RequestBody.create(MEDIA_TYPE_JSON, bodyString); + + ResponseWithBody response = requestHelper(LOGIN_PATH, "post", requestBody, headers, retryPolicy); + // old server do not support this API + if (response.response.code() != 400) { + String version = objectMapper.readTree(response.body).get("version").asText(); + if (version != null) { + this.serverVersion = new Semver(version); + this.serverCapability = new Capability(this.serverVersion); + } + } + } catch(JsonProcessingException e){ + throw new RuntimeException(e); + } + } + + private static String randRouteHint() { + String charset = "abcdef0123456789"; + Random rand = new Random(); + StringBuilder sb = new StringBuilder(16); + for (int i = 0; i < 16; i++) { + sb.append(charset.charAt(rand.nextInt(charset.length()))); + } + return sb.toString(); + } + + private static final char SPECIAL_CHAR = '#'; + + private static String uriRouteHint(String URI) { + // Encode the URI using Base64 + String encodedUri = Base64.getEncoder().encodeToString(URI.getBytes()); + + // Append the special character + return encodedUri + SPECIAL_CHAR; + } + + private static URI parseRouteHint(String routeHint) { + if (routeHint == null || routeHint.isEmpty()) { + return null; + } + try { + if (routeHint.charAt(routeHint.length() - 1) != SPECIAL_CHAR) { + return null; + } + // Remove the special character + String encodedUri = routeHint.substring(0, routeHint.length() - 1); + + // Decode the Base64 string + byte[] decodedBytes = Base64.getDecoder().decode(encodedUri); + String decodedUri = new String(decodedBytes); + + return create(decodedUri); + } catch (Exception e) { + logger.log(Level.FINE, "Failed to parse route hint: " + routeHint, e); + return null; + } + } + + + private static void checkResultSet(int resultSetType, int resultSetConcurrency) + throws SQLFeatureNotSupportedException { + if (resultSetType != ResultSet.TYPE_FORWARD_ONLY) { + throw new SQLFeatureNotSupportedException("Result set type must be TYPE_FORWARD_ONLY"); + } + if (resultSetConcurrency != ResultSet.CONCUR_READ_ONLY) { + throw new SQLFeatureNotSupportedException("Result set concurrency must be CONCUR_READ_ONLY"); + } + } + + static String getCopyIntoSql(String database, DatabendCopyParams params) { + StringBuilder sb = new StringBuilder(); + sb.append("COPY INTO "); + if (database != null) { + sb.append(database).append("."); + } + sb.append(params.getDatabaseTableName()).append(" "); + sb.append("FROM "); + sb.append(params.getDatabendStage().toString()); + sb.append(" "); + sb.append(params); + return sb.toString(); + } + + DatabendSession getSession() { + return this.session.get(); + } + + private boolean inActiveTransaction() { + if (this.session.get() == null) { + return false; + } + return this.session.get().inActiveTransaction(); + } + + private void setSession(DatabendSession session) { + if (session == null) { + return; + } + this.session.set(session); + } + + public OkHttpClient getHttpClient() { + return httpClient; + } + + @Override + public Statement createStatement() + throws SQLException { + return doCreateStatement(); + } + + private DatabendStatement doCreateStatement() throws SQLException { + checkOpen(); + DatabendStatement statement = new DatabendStatement(this, this::unregisterStatement); + registerStatement(statement); + return statement; + } + + synchronized private void registerStatement(DatabendStatement statement) { + checkState(statements.put(statement, true) == null, "Statement is already registered"); + } + + synchronized private void unregisterStatement(DatabendStatement statement) { + checkNotNull(statements.remove(statement), "Statement is not registered"); + } + + @Override + public PreparedStatement prepareStatement(String s) + throws SQLException { + + return this.prepareStatement(s, ResultSet.TYPE_FORWARD_ONLY, ResultSet.CONCUR_READ_ONLY); + } + + @Override + public CallableStatement prepareCall(String s) + throws SQLException { + throw new SQLFeatureNotSupportedException("prepareCall"); + } + + @Override + public String nativeSQL(String sql) + throws SQLException { + checkOpen(); + return sql; + } + + private void checkOpen() + throws SQLException { + if (isClosed()) { + throw new SQLException("Connection is closed"); + } + } + + @Override + public void commit() + throws SQLException { + checkOpen(); + try { + this.startQuery("commit"); + } catch (SQLException e) { + throw new SQLException("Failed to commit", e); + } + } + + @Override + public boolean getAutoCommit() + throws SQLException { + checkOpen(); + return autoCommit.get(); + } + + @Override + public void setAutoCommit(boolean b) + throws SQLException { + this.session.get().setAutoCommit(b); + autoCommit.set(b); + } + + @Override + public void rollback() + throws SQLException { + checkOpen(); + try { + this.startQuery("rollback"); + } catch (SQLException e) { + throw new SQLException("Failed to rollback", e); + } + } + + @Override + public void close() + throws SQLException { + for (Statement stmt : statements.keySet()) { + stmt.close(); + } + logout(); + } + + @Override + public boolean isClosed() + throws SQLException { + return closed.get(); + } + + @Override + public DatabaseMetaData getMetaData() + throws SQLException { + return new DatabendDatabaseMetaData(this); + } + + @Override + public boolean isReadOnly() + throws SQLException { + return false; + } + + @Override + public void setReadOnly(boolean b) + throws SQLException { + + } + + @Override + public String getCatalog() + throws SQLException { + return null; + } + + @Override + public void setCatalog(String s) + throws SQLException { + + } + + @Override + public int getTransactionIsolation() + throws SQLException { + return Connection.TRANSACTION_NONE; + } + + @Override + public void setTransactionIsolation(int i) + throws SQLException { + + } + + @Override + public SQLWarning getWarnings() + throws SQLException { + return null; + } + + @Override + public void clearWarnings() + throws SQLException { + + } + + @Override + public Statement createStatement(int resultSetType, int resultSetConcurrency) + throws SQLException { + checkResultSet(resultSetType, resultSetConcurrency); + return createStatement(); + } + + @Override + public PreparedStatement prepareStatement(String s, int i, int i1) + throws SQLException { + DatabendPreparedStatement statement = new DatabendPreparedStatement(this, this::unregisterStatement, s); + registerStatement(statement); + return statement; + } + + @Override + public CallableStatement prepareCall(String s, int i, int i1) + throws SQLException { + throw new SQLFeatureNotSupportedException("prepareCall"); + } + + @Override + public Map> getTypeMap() + throws SQLException { + throw new SQLFeatureNotSupportedException("getTypeMap"); + } + + @Override + public void setTypeMap(Map> map) + throws SQLException { + throw new SQLFeatureNotSupportedException("setTypeMap"); + } + @Override + public int getHoldability() throws SQLException { + return 0; + } + + public int getMaxFailoverRetries() { + return this.driverUri.getMaxFailoverRetry(); + } + + @Override + @NotImplemented + public void setHoldability(int holdability) throws SQLException { + // No support for transaction + } + + @Override + public Savepoint setSavepoint() + throws SQLException { + throw new SQLFeatureNotSupportedException("setSavepoint"); + } + + @Override + public Savepoint setSavepoint(String s) + throws SQLException { + throw new SQLFeatureNotSupportedException("setSavepoint"); + } + + @Override + public void rollback(Savepoint savepoint) + throws SQLException { + throw new SQLFeatureNotSupportedException("rollback"); + + } + + @Override + public void releaseSavepoint(Savepoint savepoint) + throws SQLException { + throw new SQLFeatureNotSupportedException("releaseSavepoint"); + + } + + @Override + public Statement createStatement(int resultSetType, int resultSetConcurrency, int resultSetHoldability) + throws SQLException { +// checkHoldability(resultSetHoldability); + return createStatement(resultSetType, resultSetConcurrency); + } + + @Override + public PreparedStatement prepareStatement(String sql, int resultSetType, int resultSetConcurrency, int resultSetHoldability) + throws SQLException { +// checkHoldability(resultSetHoldability); + return prepareStatement(sql, resultSetType, resultSetConcurrency); + } + + @Override + public CallableStatement prepareCall(String sql, int resultSetType, int resultSetConcurrency, int resultSetHoldability) + throws SQLException { + return prepareCall(sql, resultSetType, resultSetConcurrency); + } + + @Override + public PreparedStatement prepareStatement(String s, int autoGeneratedKeys) + throws SQLException { + return prepareStatement(s); + } + + @Override + public PreparedStatement prepareStatement(String s, int[] ints) + throws SQLException { + throw new SQLFeatureNotSupportedException("prepareStatement"); + } + + @Override + public PreparedStatement prepareStatement(String s, String[] strings) + throws SQLException { + throw new SQLFeatureNotSupportedException("prepareStatement"); + } + + @Override + public Clob createClob() + throws SQLException { + throw new SQLFeatureNotSupportedException("createClob"); + } + + @Override + public Blob createBlob() + throws SQLException { + throw new SQLFeatureNotSupportedException("createBlob"); + } + + @Override + public NClob createNClob() + throws SQLException { + throw new SQLFeatureNotSupportedException("createNClob"); + } + + @Override + public SQLXML createSQLXML() + throws SQLException { + throw new SQLFeatureNotSupportedException("createSQLXML"); + } + + @Override + public boolean isValid(int i) + throws SQLException { + return !isClosed(); + } + + @Override + public void setClientInfo(String s, String s1) + throws SQLClientInfoException { + + } + + @Override + public String getClientInfo(String s) + throws SQLException { + return null; + } + + @Override + public Properties getClientInfo() + throws SQLException { + return null; + } + + @Override + public void setClientInfo(Properties properties) + throws SQLClientInfoException { + + } + + @Override + public Array createArrayOf(String s, Object[] objects) + throws SQLException { + throw new SQLFeatureNotSupportedException("createArrayOf"); + } + + @Override + public Struct createStruct(String s, Object[] objects) + throws SQLException { + throw new SQLFeatureNotSupportedException("createStruct"); + } + + @Override + public String getSchema() + throws SQLException { + checkOpen(); + return schema.get(); + } + + @Override + public void setSchema(String schema) + throws SQLException { + checkOpen(); + this.schema.set(schema); + this.startQuery("use " + schema); + } + + @Override + public void abort(Executor executor) + throws SQLException { + close(); + } + + @Override + public void setNetworkTimeout(Executor executor, int i) + throws SQLException { + + } + + @Override + public int getNetworkTimeout() + throws SQLException { + return 0; + } + + @Override + public T unwrap(Class aClass) + throws SQLException { + if (isWrapperFor(aClass)) { + return (T) this; + } + throw new SQLException("No wrapper for " + aClass); + } + + @Override + public boolean isWrapperFor(Class aClass) + throws SQLException { + return aClass.isInstance(this); + } + + boolean presignedUrlDisabled() { + return this.driverUri.presignedUrlDisabled(); + } + + boolean copyPurge() { + return this.driverUri.copyPurge(); + } + + boolean isAutoDiscovery() { + return this.autoDiscovery; + } + + String warehouse() { + return this.driverUri.getWarehouse(); + } + + Boolean strNullAsNull() { + return this.driverUri.getStrNullAsNull(); + } + + Boolean useVerify() { + return this.driverUri.getUseVerify(); + } + + Boolean debug() { + return this.driverUri.getDebug(); + } + + String tenant() { + return this.driverUri.getTenant(); + } + + String nullDisplay() { + return this.driverUri.nullDisplay(); + } + + String binaryFormat() { + return this.driverUri.binaryFormat(); + } + + PaginationOptions getPaginationOptions() { + PaginationOptions.Builder builder = PaginationOptions.builder(); + builder.setWaitTimeSecs(this.driverUri.getWaitTimeSecs()); + builder.setMaxRowsInBuffer(this.driverUri.getMaxRowsInBuffer()); + builder.setMaxRowsPerPage(this.driverUri.getMaxRowsPerPage()); + return builder.build(); + } + + public URI getURI() { + return this.httpUri; + } + + private String buildUrlWithQueryRequest(ClientSettings settings, String querySql) { + QueryRequest req = QueryRequest.builder() + .setSession(settings.getSession()) + .setStageAttachment(settings.getStageAttachment()) + .setPaginationOptions(settings.getPaginationOptions()) + .setSql(querySql) + .build(); + String reqString = req.toString(); + if (reqString == null || reqString.isEmpty()) { + throw new IllegalArgumentException("Invalid request: " + req); + } + return reqString; + } + + void pingDatabendClientV1() throws SQLException { + try (Statement statement = this.createStatement()) { + statement.execute("select 1"); + ResultSet r = statement.getResultSet(); + while (r.next()) { + } + } catch (SQLException e) { + throw new DatabendFailedToPingException(String.format("failed to ping databend server: %s", e.getMessage())); + } + } + @Override + public void accept(DatabendSession session) { + setSession(session); + } + + /** + * Retry executing a query in case of connection errors. fail over mechanism is used to retry the query when connect error occur + * It will find next target host based on configured Load balancing Policy. + * + * @param sql The SQL statement to execute. + * @param attach The stage attachment to use for the query. + * @return A DatabendClient instance representing the successful query execution. + * @throws SQLException If the query fails after retrying the specified number of times. + * @see DatabendClientLoadBalancingPolicy + */ + DatabendClient startQueryWithFailover(String sql, StageAttachment attach) throws SQLException { + int maxRetries = getMaxFailoverRetries(); + SQLException lastException = null; + + for (int attempt = 0; attempt <= maxRetries; attempt++) { + try { + String queryId = UUID.randomUUID().toString().replace("-", ""); + String candidateHost = selectHostForQuery(queryId); + + // configure the client settings + ClientSettings.Builder sb = this.makeClientSettings(queryId, candidateHost); + if (attach != null) { + sb.setStageAttachment(attach); + } + ClientSettings settings = sb.build(); + + logger.log(Level.FINE, "execute query #{0}: SQL: {1} host: {2}", + new Object[]{attempt + 1, sql, settings.getHost()}); + + // need to retry the auto discovery in case of connection error + if (this.autoDiscovery) { + tryAutoDiscovery(httpClient, settings); + } + + return new DatabendClientV1(httpClient, sql, settings, this, lastNodeID); + } catch (Exception e) { + // handle the exception and retry the query + if (shouldRetryException(e) && attempt < maxRetries) { + lastException = wrapException("query failed", sql, e); + try { + // back off retry + Thread.sleep(Math.min(100 * (1 << attempt), 5000)); + } catch (InterruptedException ie) { + Thread.currentThread().interrupt(); + throw wrapException("query interrupt", sql, ie); + } + } else { + // throw the exception + if (e instanceof SQLException) { + throw (SQLException) e; + } else { + throw wrapException("Query failed，no need to retry", sql, e); + } + } + } + } + + throw new SQLException("after" + maxRetries + "times retry and failed: SQL: " + sql, lastException); + } + + private boolean shouldRetryException(Exception e) { + Throwable cause = e.getCause(); + // connection error + if (cause instanceof ConnectException) { + return true; + } + + if (e instanceof RuntimeException) { + String message = e.getMessage(); + return message != null && ( + message.contains("520") || + message.contains("timeout") || + message.contains("retry") + ); + } + + return false; + } + + private String selectHostForQuery(String queryId) { + String candidateHost = this.driverUri.getUri(queryId).toString(); + + if (!inActiveTransaction()) { + this.routeHint = uriRouteHint(candidateHost); + } + + if (this.routeHint != null && !this.routeHint.isEmpty()) { + URI uri = parseRouteHint(this.routeHint); + if (uri != null) { + candidateHost = uri.toString(); + } + } + + return candidateHost; + } + + private SQLException wrapException(String prefix, String sql, Exception e) { + String message = prefix + ": SQL: " + sql; + if (e.getMessage() != null) { + message += " - " + e.getMessage(); + } + if (e.getCause() != null) { + message += " (Reason: " + e.getCause().getMessage() + ")"; + } + return new SQLException(message, e); + } + /** + * Try to auto discovery the databend nodes it will log exceptions when auto discovery failed and not affect real query execution + * + * @param client the http client to query on + * @param settings the client settings to use + */ + void tryAutoDiscovery(OkHttpClient client, ClientSettings settings) { + if (this.autoDiscovery) { + if (this.driverUri.enableMock()) { + settings.getAdditionalHeaders().put("~mock.unsupported.discovery", "true"); + } + DatabendNodes nodes = this.driverUri.getNodes(); + if (nodes != null && nodes.needDiscovery()) { + try { + nodes.discoverUris(client, settings); + } catch (UnsupportedOperationException e) { + logger.log(Level.WARNING, "Current Query Node do not support auto discovery, close the functionality: " + e.getMessage()); + this.autoDiscovery = false; + } catch (Exception e) { + logger.log(Level.FINE, "Error auto discovery: " + " cause: " + e.getCause() + " message: " + e.getMessage()); + } + } + } + + } + + DatabendClient startQuery(String sql) throws SQLException { + return startQuery(sql, null); + } + + DatabendClient startQuery(String sql, StageAttachment attach) throws SQLException { + DatabendClient client = startQueryWithFailover(sql, attach); + Long timeout = client.getResults().getResultTimeoutSecs(); + if (timeout != null && timeout != 0) { + heartbeatManager.onStartQuery(timeout); + } + return client; + } + + private ClientSettings.Builder makeClientSettings(String queryID, String host) { + PaginationOptions options = getPaginationOptions(); + Map additionalHeaders = setAdditionalHeaders(); + additionalHeaders.put(X_Databend_Query_ID, queryID); + return new Builder(). + setSession(this.session.get()). + setHost(host). + setQueryTimeoutSecs(this.driverUri.getQueryTimeout()). + setConnectionTimeout(this.driverUri.getConnectionTimeout()). + setSocketTimeout(this.driverUri.getSocketTimeout()). + setPaginationOptions(options). + setAdditionalHeaders(additionalHeaders); + } + + private Map setAdditionalHeaders() { + Map additionalHeaders = new HashMap<>(); + + DatabendSession session = this.getSession(); + String warehouse = null; + if (session != null ) { + Map settings = session.getSettings(); + if (settings != null) { + warehouse = settings.get("warehouse"); + } + } + if (warehouse == null && !this.driverUri.getWarehouse().isEmpty()) { + warehouse = this.driverUri.getWarehouse(); + } + if (warehouse!=null) { + additionalHeaders.put(DatabendWarehouseHeader, warehouse); + } + + if (!this.driverUri.getTenant().isEmpty()) { + additionalHeaders.put(DatabendTenantHeader, this.driverUri.getTenant()); + } + if (!this.routeHint.isEmpty()) { + additionalHeaders.put(X_DATABEND_ROUTE_HINT, this.routeHint); + } + additionalHeaders.put("User-Agent", USER_AGENT_VALUE); + return additionalHeaders; + } + + + /** + * Method to put data from a stream at a stage location. The data will be uploaded as one file. No + * splitting is done in this method. + * + *

Stream size must match the total size of data in the input stream unless compressData + * parameter is set to true. + * + *

caller is responsible for passing the correct size for the data in the stream and releasing + * the inputStream after the method is called. + * + *

Note this method is deprecated since streamSize is not required now. Keep the function + * signature for backward compatibility + * + * @param stageName stage name: e.g. ~ or table name or stage name + * @param destPrefix path prefix under which the data should be uploaded on the stage + * @param inputStream input stream from which the data will be uploaded + * @param destFileName destination file name to use + * @param fileSize data size in the stream + * @throws SQLException failed to put data from a stream at stage + */ + @Override + public void uploadStream(String stageName, String destPrefix, InputStream inputStream, String destFileName, long fileSize, boolean compressData) + throws SQLException { + /* + remove / in the end of stage name + remove / in the beginning of destPrefix and end of destPrefix + */ + String s; + if (stageName == null) { + s = "~"; + } else { + s = stageName.replaceAll("/$", ""); + } + String p = destPrefix.replaceAll("^/", "").replaceAll("/$", ""); + String dest = p + "/" + destFileName; + try { + InputStream dataStream = inputStream; + if (compressData) { + // Wrap the input stream with a GZIPOutputStream for compression + ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream(); + try (GZIPOutputStream gzipOutputStream = new GZIPOutputStream(byteArrayOutputStream)) { + byte[] buffer = new byte[1024]; + int len; + while ((len = inputStream.read(buffer)) != -1) { + gzipOutputStream.write(buffer, 0, len); + } + } + dataStream = new ByteArrayInputStream(byteArrayOutputStream.toByteArray()); + // Update the file size to the compressed size + fileSize = byteArrayOutputStream.size(); + } + if (this.driverUri.presignedUrlDisabled()) { + DatabendPresignClient cli = new DatabendPresignClientV1(httpClient, this.httpUri.toString()); + cli.presignUpload(null, dataStream, s, p + "/", destFileName, fileSize, true); + } else { +// logger.log(Level.FINE, "presign to @" + s + "/" + dest); + long presignStartTime = System.nanoTime(); + PresignContext ctx = PresignContext.getPresignContext(this, PresignContext.PresignMethod.UPLOAD, s, dest); + long presignEndTime = System.nanoTime(); + if (this.debug()) { + logger.info("presign cost time: " + (presignEndTime - presignStartTime) / 1000000.0 + "ms"); + } + Headers h = ctx.getHeaders(); + String presignUrl = ctx.getUrl(); + DatabendPresignClient cli = new DatabendPresignClientV1(new OkHttpClient(), this.httpUri.toString()); + long uploadStartTime = System.nanoTime(); + cli.presignUpload(null, dataStream, h, presignUrl, fileSize, true); + long uploadEndTime = System.nanoTime(); + if (this.debug()) { + logger.info("upload cost time: " + (uploadEndTime - uploadStartTime) / 1000000.0 + "ms"); + } + } + } catch (RuntimeException | IOException e) { + logger.warning("failed to upload input stream, file size is:" + fileSize / 1024.0 + e.getMessage()); + throw new SQLException(e); + } + } + + @Override + public InputStream downloadStream(String stageName, String path) + throws SQLException { + String s = stageName.replaceAll("/$", ""); + DatabendPresignClient cli = new DatabendPresignClientV1(httpClient, this.httpUri.toString()); + try { + PresignContext ctx = PresignContext.getPresignContext(this, PresignContext.PresignMethod.DOWNLOAD, s, path); + Headers h = ctx.getHeaders(); + String presignUrl = ctx.getUrl(); + return cli.presignDownloadStream(h, presignUrl); + } catch (RuntimeException e) { + throw new SQLException(e); + } + } + + @Override + public InputStream downloadStream(String stageName, String path, boolean decompress) + throws SQLException { + return downloadStream(stageName, path); + } + + @Override + public void copyIntoTable(String database, String tableName, DatabendCopyParams params) + throws SQLException { + DatabendCopyParams p = params == null ? DatabendCopyParams.builder().build() : params; + requireNonNull(p.getDatabaseTableName(), "tableName is null"); + requireNonNull(p.getDatabendStage(), "stage is null"); + String sql = getCopyIntoSql(database, p); + Statement statement = this.createStatement(); + statement.execute(sql); + ResultSet rs = statement.getResultSet(); + while (rs.next()) { + } + } + + @Override + public int loadStreamToTable(String sql, InputStream inputStream, long fileSize, LoadMethod loadMethod) throws SQLException { + if (!this.serverCapability.streamingLoad()) { + throw new SQLException("please upgrade databend-query to >1.2.781 to use loadStreamToTable, current version=" + this.serverVersion); + } + + if (!sql.contains("@_databend_load")) { + throw new SQLException("invalid sql: must contain @_databend_load when used in loadStreamToTable "); + } + + if (loadMethod.equals(LoadMethod.STREAMING)) { + return streamingLoad(sql, inputStream, fileSize); + } else { + Instant now = Instant.now(); + long nanoTimestamp = now.getEpochSecond() * 1_000_000_000 + now.getNano(); + String fileName = String.valueOf(nanoTimestamp); + String location = "~/_databend_load/" + fileName; + sql = sql.replace("_databend_load", location); + uploadStream("~", "_databend_load", inputStream, fileName, fileSize, false); + Statement statement = this.createStatement(); + statement.execute(sql); + ResultSet rs = statement.getResultSet(); + while (rs.next()) { + } + return statement.getUpdateCount(); + } + } + + MultipartBody buildMultiPart(InputStream inputStream, long fileSize) { + RequestBody requestBody = new RequestBody() { + @Override + public MediaType contentType() { + return MediaType.parse("application/octet-stream"); + } + + @Override + public long contentLength() { + return fileSize; + } + + @Override + public void writeTo(BufferedSink sink) throws IOException { + try (Source source = Okio.source(inputStream)) { + sink.writeAll(source); + } + } + }; + return new MultipartBody.Builder() + .setType(MultipartBody.FORM) + .addFormDataPart( + "upload", + "java.io.InputStream", + requestBody + ).build(); + } + + int streamingLoad(String sql, InputStream inputStream, long fileSize) throws SQLException { + RetryPolicy retryPolicy = new RetryPolicy(true, true); + + try { + HashMap headers = new HashMap<>(); + DatabendSession session = this.session.get(); + if (session != null) { + String sessionString = objectMapper.writeValueAsString(session); + headers.put(DatabendQueryContextHeader, sessionString); + } + headers.put(DatabendSQLHeader, sql); + headers.put("Accept", "application/json"); + RequestBody requestBody = buildMultiPart(inputStream, fileSize); + ResponseWithBody response = requestHelper(STREAMING_LOAD_PATH, "put", requestBody, headers, retryPolicy); + JsonNode json = objectMapper.readTree(response.body); + JsonNode error = json.get("error"); + if (error != null) { + throw new SQLException("streaming load fail: code = " + error.get("code").asText() + ", message=" + error.get("message").asText()); + } + String base64 = response.response.headers().get(DatabendQueryContextHeader); + if (base64 != null) { + byte[] bytes = Base64.getUrlDecoder().decode(base64); + String str = new String(bytes, StandardCharsets.UTF_8); + try { + session = SESSION_JSON_CODEC.fromJson(str); + } catch(Exception e) { + throw new RuntimeException(e); + } + if (session != null) { + this.session.set(session); + } + } + JsonNode stats = json.get("stats"); + if (stats != null) { + int rows = stats.get("rows").asInt(-1); + if (rows != -1) { + return rows; + } + } + throw new SQLException("invalid response for " + STREAMING_LOAD_PATH + ": " + response.body); + } catch(JsonProcessingException e) { + throw new RuntimeException(e); + } + } + + void logout() throws SQLException { + DatabendSession session = this.session.get(); + if (session == null || !session.getNeedKeepAlive()) { + return; + } + RetryPolicy retryPolicy = new RetryPolicy(false, false); + RequestBody body = RequestBody.create(MEDIA_TYPE_JSON, "{}"); + requestHelper(LOGOUT_PATH, "post", body, new HashMap<>(), retryPolicy); + } + + + HttpUrl getUrl(String path) { + String host = this.driverUri.getUri().toString(); + HttpUrl url = HttpUrl.get(host); + return url.newBuilder().encodedPath(path).build(); + } + + ResponseWithBody sendRequestWithRetry(Request request, RetryPolicy retryPolicy, String path) throws SQLException { + String failReason = null; + + for (int j = 1; j <= 3; j++) { + try (Response response = httpClient.newCall(request).execute()) { + int code = response.code(); + if (code != 200) { + if (retryPolicy.shouldIgnore(code)) { + return new ResponseWithBody(response, ""); + } else { + failReason = "status code =" + response.code() + ", body = " + response.body().string(); + if (!retryPolicy.shouldRetry(code)) + break; + } + } else { + String body = response.body().string(); + return new ResponseWithBody(response, body); + } + } catch (IOException e) { + if (retryPolicy.shouldRetry(e)) { + if (failReason == null) { + failReason = e.getMessage(); + } + } else { + break; + } + } + if (j < 3) { + try { + MILLISECONDS.sleep(j * 100); + } catch (InterruptedException e2) { + Thread.currentThread().interrupt(); + return null; + } + } + } + throw new SQLException("Error accessing " + path + ": " + failReason); + } + + ResponseWithBody requestHelper(String path, String method, RequestBody body, Map headers, RetryPolicy retryPolicy) throws SQLException { + DatabendSession session = this.session.get(); + HttpUrl url = getUrl(path); + + Request.Builder builder = new Request.Builder().url(url); + this.setAdditionalHeaders().forEach(builder::addHeader); + if (headers != null) { + headers.forEach(builder::addHeader); + } + if (session.getNeedSticky()) { + builder.addHeader(ClientSettings.X_DATABEND_ROUTE_HINT, url.host()); + String lastNodeID = this.lastNodeID.get(); + if (lastNodeID != null) { + builder.addHeader(ClientSettings.X_DATABEND_STICKY_NODE, lastNodeID); + } + } + if ("post".equals(method)) { + builder = builder.post(body); + } else if ("put".equals(method)) { + builder = builder.put(body); + } else { + builder = builder.get(); + } + Request request = builder.build(); + return sendRequestWithRetry(request, retryPolicy, path); + } + + class HeartbeatManager implements Runnable { + private ScheduledFuture heartbeatFuture; + private long heartbeatIntervalMillis = 30000; + private long lastHeartbeatStartTimeMillis = 0; + private ScheduledExecutorService getScheduler() { + if (heartbeatScheduler == null) { + synchronized (HeartbeatManager.class) { + if (heartbeatScheduler == null) { + // create daemon thread so that it will not block JVM from exiting. + heartbeatScheduler = + Executors.newScheduledThreadPool( + 1, + runnable -> { + Thread thread = Executors.defaultThreadFactory().newThread(runnable); + thread.setName("heartbeat (" + thread.getId() + ")"); + thread.setDaemon(true); + return thread; + }); + } + } + } + return (ScheduledExecutorService) heartbeatScheduler; + } + + private void scheduleHeartbeat() { + long delay = Math.max(heartbeatIntervalMillis - (System.currentTimeMillis() - lastHeartbeatStartTimeMillis), 0); + heartbeatFuture = getScheduler().schedule(this, delay, MILLISECONDS); + } + + private ArrayList queryLiveness() { + ArrayList arr = new ArrayList<>(); + for (DatabendStatement stmt : statements.keySet()) { + QueryLiveness ql = stmt.queryLiveness(); + if (ql != null && !ql.stopped && ql.serverSupportHeartBeat) { + arr.add(ql); + } + } + return arr; + } + + private void doHeartbeat(ArrayList queryLivenesses ) { + long now = System.currentTimeMillis(); + lastHeartbeatStartTimeMillis = now; + Map> nodeToQueryID = new HashMap<>(); + Map queries = new HashMap<>(); + + for (QueryLiveness ql: queryLivenesses) { + if (now - ql.lastRequestTime.get() >= ql.resultTimeoutSecs * 1000 / 2) { + nodeToQueryID.computeIfAbsent(ql.nodeID, k -> new ArrayList<>()).add(ql.queryID); + queries.put(ql.queryID, ql); + } + } + if (nodeToQueryID.isEmpty()) { + return; + } + + Map map = new HashMap<>(); + map.put("node_to_queries", nodeToQueryID); + + try { + String body = objectMapper.writeValueAsString(map); + RequestBody requestBody = RequestBody.create(MEDIA_TYPE_JSON, body); + RetryPolicy retryPolicy = new RetryPolicy(true, false); + body = requestHelper(HEARTBEAT_PATH, "post", requestBody, null, retryPolicy).body; + JsonNode toRemove = objectMapper.readTree(body).get("queries_to_remove"); + if (toRemove.isArray()) { + for (JsonNode element : toRemove) { + String queryId = element.asText(); + queries.get(queryId).stopped = true; + } + } + } catch (JsonProcessingException e) { + logger.warning("fail to encode heartbeat body: " + e); + } catch (SQLException e) { + logger.warning("fail to send heartbeat: " + e); + } catch (Exception e) { + logger.warning("fail to send heartbeat: " + e); + throw new RuntimeException(e); + } + } + + public void onStartQuery(Long timeoutSecs) { + synchronized (DatabendConnectionImpl.this) { + if (timeoutSecs * 1000 / 4 < heartbeatIntervalMillis) { + heartbeatIntervalMillis = timeoutSecs * 1000 / 4; + if (heartbeatFuture != null) { + heartbeatFuture.cancel(false); + heartbeatFuture = null; + } + } + if (heartbeatFuture == null) { + scheduleHeartbeat(); + } + } + } + + @Override + public void run() { + ArrayList arr = queryLiveness(); + doHeartbeat(arr); + + heartbeatFuture = null; + synchronized (DatabendConnectionImpl.this) { + if (arr.size() > 0) { + if (heartbeatFuture == null) { + scheduleHeartbeat(); + } + } else { + heartbeatFuture = null; + } + } + } + } + + boolean isHeartbeatStopped() { + return heartbeatManager.heartbeatFuture == null; + } + + static class RetryPolicy { + boolean ignore404; + boolean retry503; + RetryPolicy(boolean ignore404, boolean retry503) { + this.ignore404 = ignore404; + this.retry503 = retry503; + } + + boolean shouldIgnore(int code) { + return ignore404 && code == 404; + } + + boolean shouldRetry(int code) { + return retry503 && (code == 502 || code == 503); + } + + boolean shouldRetry(IOException e) { + return e.getCause() instanceof ConnectException; + } + } + + static class ResponseWithBody { + public Response response; + public String body; + + ResponseWithBody(Response response, String body) { + this.response = response; + this.body = body; + } + } +} diff --git a/databend-jdbc/src/main/java/com/databend/jdbc/constant/DatabendConstant.java b/databend-jdbc/src/main/java/com/databend/jdbc/DatabendConstant.java similarity index 91% rename from databend-jdbc/src/main/java/com/databend/jdbc/constant/DatabendConstant.java rename to databend-jdbc/src/main/java/com/databend/jdbc/DatabendConstant.java index 09c89cca..7cb849df 100644 --- a/databend-jdbc/src/main/java/com/databend/jdbc/constant/DatabendConstant.java +++ b/databend-jdbc/src/main/java/com/databend/jdbc/DatabendConstant.java @@ -1,4 +1,4 @@ -package com.databend.jdbc.constant; +package com.databend.jdbc; import java.util.regex.Pattern; diff --git a/databend-jdbc/src/main/java/com/databend/jdbc/DatabendDatabaseMetaData.java b/databend-jdbc/src/main/java/com/databend/jdbc/DatabendDatabaseMetaData.java index ada6d6ff..0a07f6e4 100644 --- a/databend-jdbc/src/main/java/com/databend/jdbc/DatabendDatabaseMetaData.java +++ b/databend-jdbc/src/main/java/com/databend/jdbc/DatabendDatabaseMetaData.java @@ -21,11 +21,11 @@ import static com.databend.jdbc.DriverInfo.*; import static java.util.Objects.requireNonNull; -public class DatabendDatabaseMetaData implements DatabaseMetaData { +class DatabendDatabaseMetaData implements DatabaseMetaData { private static final String SEARCH_STRING_ESCAPE = "\\"; - private final DatabendConnection connection; + private final DatabendConnectionImpl connection; - public DatabendDatabaseMetaData(DatabendConnection connection) + public DatabendDatabaseMetaData(DatabendConnectionImpl connection) throws SQLException { requireNonNull(connection, "connection is null"); this.connection = connection; diff --git a/databend-jdbc/src/main/java/com/databend/jdbc/DatabendDriverUri.java b/databend-jdbc/src/main/java/com/databend/jdbc/DatabendDriverUri.java index b726deaf..8d4ff88b 100644 --- a/databend-jdbc/src/main/java/com/databend/jdbc/DatabendDriverUri.java +++ b/databend-jdbc/src/main/java/com/databend/jdbc/DatabendDriverUri.java @@ -26,7 +26,7 @@ import static com.databend.client.OkHttpUtils.*; import static com.databend.jdbc.ConnectionProperties.*; -import static com.databend.jdbc.constant.DatabendConstant.ENABLE_STR; +import static com.databend.jdbc.DatabendConstant.ENABLE_STR; import static com.google.common.base.MoreObjects.firstNonNull; import static com.google.common.base.Strings.isNullOrEmpty; import static java.lang.String.format; @@ -35,7 +35,7 @@ /** * Parses and extracts parameters from a databend JDBC URL */ -public final class DatabendDriverUri { +final class DatabendDriverUri { private static final String JDBC_URL_PREFIX = "jdbc:"; private static final String JDBC_URL_START = JDBC_URL_PREFIX + "databend://"; private static final Splitter QUERY_SPLITTER = Splitter.on('&').omitEmptyStrings(); diff --git a/databend-jdbc/src/main/java/com/databend/jdbc/DatabendNodeRouter.java b/databend-jdbc/src/main/java/com/databend/jdbc/DatabendNodeRouter.java index dfa31bda..0a8f9f2d 100644 --- a/databend-jdbc/src/main/java/com/databend/jdbc/DatabendNodeRouter.java +++ b/databend-jdbc/src/main/java/com/databend/jdbc/DatabendNodeRouter.java @@ -9,7 +9,7 @@ /** * Node manager manage a list of hosts */ -public interface DatabendNodeRouter { +interface DatabendNodeRouter { /** * Gets a copy of all possible query uris * diff --git a/databend-jdbc/src/main/java/com/databend/jdbc/DatabendNodes.java b/databend-jdbc/src/main/java/com/databend/jdbc/DatabendNodes.java index 86757251..351209f5 100644 --- a/databend-jdbc/src/main/java/com/databend/jdbc/DatabendNodes.java +++ b/databend-jdbc/src/main/java/com/databend/jdbc/DatabendNodes.java @@ -16,7 +16,7 @@ import java.util.concurrent.atomic.AtomicReference; import java.util.logging.Logger; -public class DatabendNodes implements DatabendNodeRouter { +class DatabendNodes implements DatabendNodeRouter { private AtomicReference> query_nodes_uris; protected final AtomicInteger index; @@ -110,7 +110,7 @@ public boolean needDiscovery() { } public List parseURI(List nodes) throws RuntimeException { - String host = null; + String host; List uris = new ArrayList<>(); try { for (DiscoveryNode node : nodes) { diff --git a/databend-jdbc/src/main/java/com/databend/jdbc/DatabendParameterMetaData.java b/databend-jdbc/src/main/java/com/databend/jdbc/DatabendParameterMetaData.java index f7622a62..b9c4a3ae 100644 --- a/databend-jdbc/src/main/java/com/databend/jdbc/DatabendParameterMetaData.java +++ b/databend-jdbc/src/main/java/com/databend/jdbc/DatabendParameterMetaData.java @@ -9,7 +9,7 @@ import static com.databend.jdbc.DatabendResultSetMetaData.getTypeClassName; import static java.util.Objects.requireNonNull; -public class DatabendParameterMetaData extends JdbcWrapper implements ParameterMetaData { +class DatabendParameterMetaData extends JdbcWrapper implements ParameterMetaData { protected final List params; protected final JdbcTypeMapping mapper; diff --git a/databend-jdbc/src/main/java/com/databend/jdbc/DatabendPreparedStatement.java b/databend-jdbc/src/main/java/com/databend/jdbc/DatabendPreparedStatement.java index e093eaf5..77a1e69d 100644 --- a/databend-jdbc/src/main/java/com/databend/jdbc/DatabendPreparedStatement.java +++ b/databend-jdbc/src/main/java/com/databend/jdbc/DatabendPreparedStatement.java @@ -59,13 +59,13 @@ import static com.databend.jdbc.ObjectCasts.*; import static com.databend.jdbc.StatementUtil.replaceParameterMarksWithValues; -import static com.databend.jdbc.constant.DatabendConstant.*; +import static com.databend.jdbc.DatabendConstant.*; import static java.lang.String.format; import static java.time.format.DateTimeFormatter.ISO_LOCAL_DATE; import static java.time.format.DateTimeFormatter.ISO_LOCAL_TIME; import static java.util.Objects.requireNonNull; -public class DatabendPreparedStatement extends DatabendStatement implements PreparedStatement { +class DatabendPreparedStatement extends DatabendStatement implements PreparedStatement { private static final Logger logger = Logger.getLogger(DatabendPreparedStatement.class.getPackage().getName()); static final DateTimeFormatter DATE_FORMATTER = ISODateTimeFormat.date(); private final RawStatementWrapper rawStatement; @@ -87,7 +87,7 @@ public class DatabendPreparedStatement extends DatabendStatement implements Prep private static final ObjectMapper objectMapper = new ObjectMapper(); - DatabendPreparedStatement(DatabendConnection connection, Consumer onClose, String sql) throws SQLException { + DatabendPreparedStatement(DatabendConnectionImpl connection, Consumer onClose, String sql) throws SQLException { super(connection, onClose); this.batchValues = new ArrayList<>(); this.batchValuesCSV = new ArrayList<>(); @@ -162,7 +162,7 @@ private StageAttachment uploadBatches() throws SQLException { } File saved = batchInsertUtils.saveBatchToCSV(batchValuesCSV); try (FileInputStream fis = new FileInputStream(saved)) { - DatabendConnection c = (DatabendConnection) getConnection(); + DatabendConnectionImpl c = (DatabendConnectionImpl) getConnection(); String uuid = UUID.randomUUID().toString().replace("-", ""); // format %Y/%m/%d/%H/%M/%S/fileName.csv String stagePrefix = String.format("%s/%s/%s/%s/%s/%s/%s/", @@ -202,7 +202,7 @@ private StageAttachment uploadBatches() throws SQLException { * @param stagePath The path of the stage in the Databend database. * @return A StageAttachment object which contains the details of the stage. */ - public static StageAttachment buildStateAttachment(DatabendConnection connection, String stagePath) { + static StageAttachment buildStateAttachment(DatabendConnectionImpl connection, String stagePath) { Map fileFormatOptions = new HashMap<>(); if (!Objects.equals(connection.binaryFormat(), "")) { fileFormatOptions.put("binary_format", String.valueOf(connection.binaryFormat())); @@ -231,7 +231,7 @@ public static StageAttachment buildStateAttachment(DatabendConnection connection * * @return true if delete success or resource not found */ - private boolean dropStageAttachment(StageAttachment attachment) { + boolean dropStageAttachment(StageAttachment attachment) { if (attachment == null) { return true; } @@ -244,7 +244,7 @@ private boolean dropStageAttachment(StageAttachment attachment) { } } - public int[] executeBatchByAttachment() throws SQLException { + int[] executeBatchByAttachment() throws SQLException { int[] batchUpdateCounts = new int[batchValues.size()]; if (batchValues.isEmpty()) { return batchUpdateCounts; diff --git a/databend-jdbc/src/main/java/com/databend/jdbc/DatabendResultSet.java b/databend-jdbc/src/main/java/com/databend/jdbc/DatabendResultSet.java index f65659f6..e34dada3 100644 --- a/databend-jdbc/src/main/java/com/databend/jdbc/DatabendResultSet.java +++ b/databend-jdbc/src/main/java/com/databend/jdbc/DatabendResultSet.java @@ -27,7 +27,7 @@ import static java.util.Objects.requireNonNull; import static java.util.concurrent.Executors.newCachedThreadPool; -public class DatabendResultSet extends AbstractDatabendResultSet { +class DatabendResultSet extends AbstractDatabendResultSet { private final Statement statement; private final DatabendClient client; @GuardedBy("this") @@ -65,7 +65,7 @@ private static Iterator flatten(Iterator> iterator, long maxR } - public QueryLiveness getLiveness() { + QueryLiveness getLiveness() { if (closed) { return null; } diff --git a/databend-jdbc/src/main/java/com/databend/jdbc/DatabendResultSetMetaData.java b/databend-jdbc/src/main/java/com/databend/jdbc/DatabendResultSetMetaData.java index 9dfe335e..dbc92dd1 100644 --- a/databend-jdbc/src/main/java/com/databend/jdbc/DatabendResultSetMetaData.java +++ b/databend-jdbc/src/main/java/com/databend/jdbc/DatabendResultSetMetaData.java @@ -12,7 +12,7 @@ import java.sql.Types; import java.util.List; -public class DatabendResultSetMetaData implements ResultSetMetaData { +class DatabendResultSetMetaData implements ResultSetMetaData { private final List databendColumnInfo; DatabendResultSetMetaData(List databendColumnInfo) { diff --git a/databend-jdbc/src/main/java/com/databend/jdbc/DatabendStatement.java b/databend-jdbc/src/main/java/com/databend/jdbc/DatabendStatement.java index d7fbbe44..5abdef74 100644 --- a/databend-jdbc/src/main/java/com/databend/jdbc/DatabendStatement.java +++ b/databend-jdbc/src/main/java/com/databend/jdbc/DatabendStatement.java @@ -21,8 +21,8 @@ import static java.lang.Math.toIntExact; import static java.util.Objects.requireNonNull; -public class DatabendStatement implements Statement { - private final AtomicReference connection; +class DatabendStatement implements Statement { + private final AtomicReference connection; private final Consumer onClose; private int currentUpdateCount = -1; private final AtomicReference currentResult = new AtomicReference<>(); @@ -30,7 +30,7 @@ public class DatabendStatement implements Statement { private final AtomicLong maxRows = new AtomicLong(); private final AtomicBoolean closeOnCompletion = new AtomicBoolean(); - DatabendStatement(DatabendConnection connection, Consumer onClose) { + DatabendStatement(DatabendConnectionImpl connection, Consumer onClose) { this.connection = new AtomicReference<>(requireNonNull(connection, "connection is null")); this.onClose = requireNonNull(onClose, "onClose is null"); } @@ -51,7 +51,7 @@ public int executeUpdate(String s) @Override public void close() throws SQLException { - DatabendConnection connection = this.connection.getAndSet(null); + DatabendConnectionImpl connection = this.connection.getAndSet(null); if (connection == null) { return; } @@ -438,9 +438,9 @@ protected final void checkOpen() connection(); } - protected final DatabendConnection connection() + protected final DatabendConnectionImpl connection() throws SQLException { - DatabendConnection connection = this.connection.get(); + DatabendConnectionImpl connection = this.connection.get(); if (connection == null) { throw new SQLException("Statement is closed"); } @@ -450,7 +450,7 @@ protected final DatabendConnection connection() return connection; } - public QueryLiveness queryLiveness() { + QueryLiveness queryLiveness() { DatabendResultSet r = currentResult.get(); if (r != null) { diff --git a/databend-jdbc/src/main/java/com/databend/jdbc/DatabendUnboundQueryResultSet.java b/databend-jdbc/src/main/java/com/databend/jdbc/DatabendUnboundQueryResultSet.java index 306a3dc1..95f9f3b7 100644 --- a/databend-jdbc/src/main/java/com/databend/jdbc/DatabendUnboundQueryResultSet.java +++ b/databend-jdbc/src/main/java/com/databend/jdbc/DatabendUnboundQueryResultSet.java @@ -16,7 +16,7 @@ * so that we need to make some modifications to the types, values, etc. after the returned results. * The actual returned datas from the interface is a new ResultSet. */ -public class DatabendUnboundQueryResultSet extends AbstractDatabendResultSet { +class DatabendUnboundQueryResultSet extends AbstractDatabendResultSet { private boolean closed = false; diff --git a/databend-jdbc/src/main/java/com/databend/jdbc/FileTransferAPI.java b/databend-jdbc/src/main/java/com/databend/jdbc/FileTransferAPI.java index 343a5758..b6b9255a 100644 --- a/databend-jdbc/src/main/java/com/databend/jdbc/FileTransferAPI.java +++ b/databend-jdbc/src/main/java/com/databend/jdbc/FileTransferAPI.java @@ -5,8 +5,16 @@ import java.io.InputStream; import java.sql.SQLException; +/** + * Deprecated. This interface has been replaced by {@link DatabendConnection}. + *

+ * Deprecated since version 4.0.1. Scheduled for removal in a future release. + * Please migrate to {@link DatabendConnection} for equivalent functionality. + */ +@Deprecated public interface FileTransferAPI { /** + * Deprecated. Use {@link DatabendConnection#uploadStream(String, String, InputStream, String, long, boolean)} instead. * Upload inputStream to the databend internal stage, the data would be uploaded as one file with no split. * Caller should close the input stream after the upload is done. * @@ -17,10 +25,13 @@ public interface FileTransferAPI { * @param fileSize the file size in the stage * @param compressData whether to compress the data * @throws SQLException failed to upload input stream + * + * @deprecated Replaced by DatabendConnection.uploadStream() since version 4.0.1 */ void uploadStream(String stageName, String destPrefix, InputStream inputStream, String destFileName, long fileSize, boolean compressData) throws SQLException; /** + * Deprecated. Use {@link DatabendConnection#downloadStream(String, String)} instead. * Download a file from the databend internal stage, the data would be downloaded as one file with no split. * * @param stageName the stage which contains the file @@ -28,6 +39,8 @@ public interface FileTransferAPI { * @param decompress whether to decompress the data * @return the input stream of the file * @throws SQLException failed to download input stream + * + * @deprecated Replaced by DatabendConnection.downloadStream() since version 4.0.1 */ InputStream downloadStream(String stageName, String sourceFileName, boolean decompress) throws SQLException; @@ -39,17 +52,7 @@ public interface FileTransferAPI { * @param tableName the target table name * @param params copy options and file options * @throws SQLException fail to copy into table + * @deprecated execute the Copy SQL directly */ void copyIntoTable(String database, String tableName, DatabendCopyParams params) throws SQLException; - - /** - * Upload inputStream into the target table - * - * @param sql the sql with syntax `Insert into

[() [values (?, ...)]] from @_databend_load [file_format=(...)]` - * @param inputStream the input stream of the file - * @param loadMethod one of "stage" or "streaming" - * @return num of rows loaded - * @throws SQLException fail to load file into table - */ - int loadStreamToTable(String sql, InputStream inputStream, long fileSize, String loadMethod) throws SQLException; } diff --git a/databend-jdbc/src/main/java/com/databend/jdbc/JdbcTypeMapping.java b/databend-jdbc/src/main/java/com/databend/jdbc/JdbcTypeMapping.java index 82aa5430..421a0cc0 100644 --- a/databend-jdbc/src/main/java/com/databend/jdbc/JdbcTypeMapping.java +++ b/databend-jdbc/src/main/java/com/databend/jdbc/JdbcTypeMapping.java @@ -4,7 +4,7 @@ import java.sql.Types; -public class JdbcTypeMapping { +class JdbcTypeMapping { /** * Converts {@link DatabendColumnInfo} to generic SQL type defined in JDBC. * diff --git a/databend-jdbc/src/main/java/com/databend/jdbc/JdbcWrapper.java b/databend-jdbc/src/main/java/com/databend/jdbc/JdbcWrapper.java index 70fb383b..540c37d9 100644 --- a/databend-jdbc/src/main/java/com/databend/jdbc/JdbcWrapper.java +++ b/databend-jdbc/src/main/java/com/databend/jdbc/JdbcWrapper.java @@ -3,7 +3,7 @@ import java.sql.SQLException; import java.sql.SQLFeatureNotSupportedException; -public abstract class JdbcWrapper { +abstract class JdbcWrapper { public T unwrap(Class iface) throws SQLException { if (iface.isAssignableFrom(getClass())) { return iface.cast(this); diff --git a/databend-jdbc/src/main/java/com/databend/jdbc/LoggerUtil.java b/databend-jdbc/src/main/java/com/databend/jdbc/LoggerUtil.java index 9c25f3e5..3d3864c6 100644 --- a/databend-jdbc/src/main/java/com/databend/jdbc/LoggerUtil.java +++ b/databend-jdbc/src/main/java/com/databend/jdbc/LoggerUtil.java @@ -16,7 +16,7 @@ @UtilityClass @CustomLog -public class LoggerUtil { +class LoggerUtil { private static Boolean slf4jAvailable; diff --git a/databend-jdbc/src/main/java/com/databend/jdbc/LoginRequest.java b/databend-jdbc/src/main/java/com/databend/jdbc/LoginRequest.java index b489fbaa..36ad338d 100644 --- a/databend-jdbc/src/main/java/com/databend/jdbc/LoginRequest.java +++ b/databend-jdbc/src/main/java/com/databend/jdbc/LoginRequest.java @@ -3,7 +3,7 @@ import java.util.Map; -public class LoginRequest { +class LoginRequest { public String database; public Map settings; } diff --git a/databend-jdbc/src/main/java/com/databend/jdbc/NonQueryRawStatement.java b/databend-jdbc/src/main/java/com/databend/jdbc/NonQueryRawStatement.java index e6236740..30d3be22 100644 --- a/databend-jdbc/src/main/java/com/databend/jdbc/NonQueryRawStatement.java +++ b/databend-jdbc/src/main/java/com/databend/jdbc/NonQueryRawStatement.java @@ -11,7 +11,7 @@ * INSERT) */ @EqualsAndHashCode(callSuper = true) -public class NonQueryRawStatement extends RawStatement { +class NonQueryRawStatement extends RawStatement { public NonQueryRawStatement(String sql, String cleanSql, List paramPositions) { super(sql, cleanSql, paramPositions); diff --git a/databend-jdbc/src/main/java/com/databend/jdbc/NonRegisteringDatabendDriver.java b/databend-jdbc/src/main/java/com/databend/jdbc/NonRegisteringDatabendDriver.java index 7bd86d34..46947be4 100644 --- a/databend-jdbc/src/main/java/com/databend/jdbc/NonRegisteringDatabendDriver.java +++ b/databend-jdbc/src/main/java/com/databend/jdbc/NonRegisteringDatabendDriver.java @@ -16,7 +16,7 @@ import static com.databend.client.OkHttpUtils.userAgentInterceptor; import static com.databend.jdbc.DriverInfo.*; -public class NonRegisteringDatabendDriver implements Driver, Closeable { +class NonRegisteringDatabendDriver implements Driver, Closeable { private final OkHttpClient httpClient = newHttpClient(); private static Properties urlProperties(String url, Properties info) { @@ -62,11 +62,11 @@ public Connection connect(String url, Properties info) OkHttpClient.Builder builder = httpClient.newBuilder(); uri.setupClient(builder); - DatabendConnection connection = new DatabendConnection(uri, builder.build()); + DatabendConnectionImpl connection = new DatabendConnectionImpl(uri, builder.build()); // ping the server host if (connection.useVerify()) { try { - connection.PingDatabendClientV1(); + connection.pingDatabendClientV1(); } catch (SQLException e) { throw new RuntimeException(e); } diff --git a/databend-jdbc/src/main/java/com/databend/jdbc/ParamMarker.java b/databend-jdbc/src/main/java/com/databend/jdbc/ParamMarker.java index 99c56c08..14de0ac4 100644 --- a/databend-jdbc/src/main/java/com/databend/jdbc/ParamMarker.java +++ b/databend-jdbc/src/main/java/com/databend/jdbc/ParamMarker.java @@ -5,7 +5,7 @@ @AllArgsConstructor @Value -public class ParamMarker { +class ParamMarker { // Id / index of the param marker in the SQL statement int id; // Position in the SQL subStatement diff --git a/databend-jdbc/src/main/java/com/databend/jdbc/PresignContext.java b/databend-jdbc/src/main/java/com/databend/jdbc/PresignContext.java index e3956ae9..0f506f4f 100644 --- a/databend-jdbc/src/main/java/com/databend/jdbc/PresignContext.java +++ b/databend-jdbc/src/main/java/com/databend/jdbc/PresignContext.java @@ -4,6 +4,7 @@ import com.fasterxml.jackson.databind.ObjectMapper; import okhttp3.Headers; +import java.sql.Connection; import java.sql.SQLException; import java.sql.Statement; import java.util.HashMap; @@ -12,7 +13,7 @@ import static com.google.common.base.MoreObjects.toStringHelper; import static java.util.Objects.requireNonNull; -public final class PresignContext { +final class PresignContext { private final PresignMethod method; private final String stageName; private final String fileName; @@ -27,19 +28,19 @@ private PresignContext(PresignMethod method, String stageName, String fileName, this.url = url; } - public static void createStageIfNotExists(DatabendConnection connection, String stageName) throws SQLException { + public static void createStageIfNotExists(Connection connection, String stageName) throws SQLException { String sql = String.format("CREATE STAGE IF NOT EXISTS %s", stageName); Statement statement = connection.createStatement(); statement.execute(sql); } - public static void dropStageIfExists(DatabendConnection connection, String stageName) throws SQLException { + public static void dropStageIfExists(DatabendConnectionImpl connection, String stageName) throws SQLException { String sql = String.format("DROP STAGE IF EXISTS %s", stageName); Statement statement = connection.createStatement(); statement.execute(sql); } - public static PresignContext getPresignContext(DatabendConnection connection, PresignMethod method, String stageName, String fileName) + public static PresignContext getPresignContext(DatabendConnectionImpl connection, PresignMethod method, String stageName, String fileName) throws SQLException { requireNonNull(connection, "connection is null"); requireNonNull(method, "method is null"); diff --git a/databend-jdbc/src/main/java/com/databend/jdbc/QueryLiveness.java b/databend-jdbc/src/main/java/com/databend/jdbc/QueryLiveness.java index cfdf2920..ce9189bc 100644 --- a/databend-jdbc/src/main/java/com/databend/jdbc/QueryLiveness.java +++ b/databend-jdbc/src/main/java/com/databend/jdbc/QueryLiveness.java @@ -2,7 +2,7 @@ import java.util.concurrent.atomic.AtomicLong; -public class QueryLiveness { +class QueryLiveness { String queryID; String nodeID; boolean serverSupportHeartBeat; diff --git a/databend-jdbc/src/main/java/com/databend/jdbc/QueryRawStatement.java b/databend-jdbc/src/main/java/com/databend/jdbc/QueryRawStatement.java index aefab2e6..a71d10de 100644 --- a/databend-jdbc/src/main/java/com/databend/jdbc/QueryRawStatement.java +++ b/databend-jdbc/src/main/java/com/databend/jdbc/QueryRawStatement.java @@ -15,7 +15,7 @@ */ @Getter @EqualsAndHashCode(callSuper = true) -public class QueryRawStatement extends RawStatement { +class QueryRawStatement extends RawStatement { private final String database; diff --git a/databend-jdbc/src/main/java/com/databend/jdbc/RawStatement.java b/databend-jdbc/src/main/java/com/databend/jdbc/RawStatement.java index 15719c1e..5064c772 100644 --- a/databend-jdbc/src/main/java/com/databend/jdbc/RawStatement.java +++ b/databend-jdbc/src/main/java/com/databend/jdbc/RawStatement.java @@ -8,7 +8,7 @@ import java.util.Optional; @Data -public abstract class RawStatement { +abstract class RawStatement { private final String sql; private final String cleanSql; diff --git a/databend-jdbc/src/main/java/com/databend/jdbc/RawStatementWrapper.java b/databend-jdbc/src/main/java/com/databend/jdbc/RawStatementWrapper.java index d74d8908..07b0fb5a 100644 --- a/databend-jdbc/src/main/java/com/databend/jdbc/RawStatementWrapper.java +++ b/databend-jdbc/src/main/java/com/databend/jdbc/RawStatementWrapper.java @@ -9,7 +9,7 @@ @CustomLog @Value -public class RawStatementWrapper { +class RawStatementWrapper { List subStatements; diff --git a/databend-jdbc/src/main/java/com/databend/jdbc/SetParamRawStatement.java b/databend-jdbc/src/main/java/com/databend/jdbc/SetParamRawStatement.java index 4b72e45a..0a1d1444 100644 --- a/databend-jdbc/src/main/java/com/databend/jdbc/SetParamRawStatement.java +++ b/databend-jdbc/src/main/java/com/databend/jdbc/SetParamRawStatement.java @@ -14,7 +14,7 @@ */ @Getter @EqualsAndHashCode(callSuper = true) -public class SetParamRawStatement extends RawStatement { +class SetParamRawStatement extends RawStatement { private final Pair additionalProperty; diff --git a/databend-jdbc/src/main/java/com/databend/jdbc/StatementInfoWrapper.java b/databend-jdbc/src/main/java/com/databend/jdbc/StatementInfoWrapper.java index fbd992b3..141c71b3 100644 --- a/databend-jdbc/src/main/java/com/databend/jdbc/StatementInfoWrapper.java +++ b/databend-jdbc/src/main/java/com/databend/jdbc/StatementInfoWrapper.java @@ -16,7 +16,7 @@ */ @Data @AllArgsConstructor -public class StatementInfoWrapper { +class StatementInfoWrapper { private String sql; private String id; private StatementType type; diff --git a/databend-jdbc/src/main/java/com/databend/jdbc/StatementType.java b/databend-jdbc/src/main/java/com/databend/jdbc/StatementType.java index 15593ab5..4af889a6 100644 --- a/databend-jdbc/src/main/java/com/databend/jdbc/StatementType.java +++ b/databend-jdbc/src/main/java/com/databend/jdbc/StatementType.java @@ -1,6 +1,6 @@ package com.databend.jdbc; -public enum StatementType { +enum StatementType { // SET PARAM_SETTING, // eg: SELECT, SHOW diff --git a/databend-jdbc/src/main/java/com/databend/jdbc/StatementUtil.java b/databend-jdbc/src/main/java/com/databend/jdbc/StatementUtil.java index d7d9e3f6..f85b20bc 100644 --- a/databend-jdbc/src/main/java/com/databend/jdbc/StatementUtil.java +++ b/databend-jdbc/src/main/java/com/databend/jdbc/StatementUtil.java @@ -21,7 +21,7 @@ @UtilityClass @CustomLog -public class StatementUtil { +class StatementUtil { private static final String SET_PREFIX = "set"; private static final Pattern SET_WITH_SPACE_REGEX = Pattern.compile(SET_PREFIX + " ", Pattern.CASE_INSENSITIVE); diff --git a/databend-jdbc/src/main/java/com/databend/jdbc/cloud/DatabendPresignClient.java b/databend-jdbc/src/main/java/com/databend/jdbc/cloud/DatabendPresignClient.java index 540463f1..023efe90 100644 --- a/databend-jdbc/src/main/java/com/databend/jdbc/cloud/DatabendPresignClient.java +++ b/databend-jdbc/src/main/java/com/databend/jdbc/cloud/DatabendPresignClient.java @@ -7,11 +7,11 @@ import java.io.InputStream; public interface DatabendPresignClient { - public void presignUpload(File srcFile, InputStream inputStream, Headers headers, String presignedUrl, long fileSize, boolean uploadFromStream) throws IOException; + void presignUpload(File srcFile, InputStream inputStream, Headers headers, String presignedUrl, long fileSize, boolean uploadFromStream) throws IOException; - public void presignDownload(String destFileName, Headers headers, String presignedUrl); + void presignDownload(String destFileName, Headers headers, String presignedUrl); - public InputStream presignDownloadStream(Headers headers, String presignedUrl); + InputStream presignDownloadStream(Headers headers, String presignedUrl); /** * presignUpload file through databend api instead of presigned url, it should only be adopted if presigned url is not available @@ -21,5 +21,5 @@ public interface DatabendPresignClient { * @param uploadFromStream whether the upload is from stream * @throws IOException */ - public void presignUpload(File srcFile, InputStream inputStream, String stageName, String relativePath, String fileName, long fileSize, boolean uploadFromStream) throws IOException; + void presignUpload(File srcFile, InputStream inputStream, String stageName, String relativePath, String fileName, long fileSize, boolean uploadFromStream) throws IOException; } diff --git a/databend-jdbc/src/main/java/com/databend/jdbc/examples/DatabendConnectionFactory.java b/databend-jdbc/src/main/java/com/databend/jdbc/examples/DatabendConnectionFactory.java index 2bfd1ffd..5e9653c1 100644 --- a/databend-jdbc/src/main/java/com/databend/jdbc/examples/DatabendConnectionFactory.java +++ b/databend-jdbc/src/main/java/com/databend/jdbc/examples/DatabendConnectionFactory.java @@ -1,6 +1,6 @@ package com.databend.jdbc.examples; -import com.databend.jdbc.DatabendConnection; +import com.databend.jdbc.DatabendConnectionImpl; import org.apache.commons.pool2.PooledObject; import org.apache.commons.pool2.PooledObjectFactory; import org.apache.commons.pool2.impl.DefaultPooledObject; @@ -10,7 +10,7 @@ import java.sql.SQLException; import java.util.Properties; -public class DatabendConnectionFactory implements PooledObjectFactory { +public class DatabendConnectionFactory implements PooledObjectFactory { private String url; private Properties properties; @@ -25,18 +25,18 @@ private Connection createConnection(String url, Properties p) throws SQLExceptio } @Override - public PooledObject makeObject() throws Exception { - DatabendConnection connection = (DatabendConnection) createConnection(url, properties); + public PooledObject makeObject() throws Exception { + DatabendConnectionImpl connection = (DatabendConnectionImpl) createConnection(url, properties); return new DefaultPooledObject<>(connection); } @Override - public void destroyObject(PooledObject p) throws Exception { + public void destroyObject(PooledObject p) throws Exception { p.getObject().close(); } @Override - public boolean validateObject(PooledObject p) { + public boolean validateObject(PooledObject p) { try { return !p.getObject().isClosed(); } catch (SQLException e) { @@ -45,11 +45,11 @@ public boolean validateObject(PooledObject p) { } @Override - public void activateObject(PooledObject p) throws Exception { + public void activateObject(PooledObject p) throws Exception { } @Override - public void passivateObject(PooledObject p) throws Exception { + public void passivateObject(PooledObject p) throws Exception { } } diff --git a/databend-jdbc/src/main/java/com/databend/jdbc/examples/DatabendConnectionPool.java b/databend-jdbc/src/main/java/com/databend/jdbc/examples/DatabendConnectionPool.java index 31ec78e9..ec40a6f3 100644 --- a/databend-jdbc/src/main/java/com/databend/jdbc/examples/DatabendConnectionPool.java +++ b/databend-jdbc/src/main/java/com/databend/jdbc/examples/DatabendConnectionPool.java @@ -1,18 +1,18 @@ package com.databend.jdbc.examples; -import com.databend.jdbc.DatabendConnection; +import com.databend.jdbc.DatabendConnectionImpl; import org.apache.commons.pool2.impl.GenericObjectPool; import org.apache.commons.pool2.impl.GenericObjectPoolConfig; import java.util.Properties; -public class DatabendConnectionPool extends GenericObjectPool { - public DatabendConnectionPool(DatabendConnectionFactory factory, GenericObjectPoolConfig config) { +public class DatabendConnectionPool extends GenericObjectPool { + public DatabendConnectionPool(DatabendConnectionFactory factory, GenericObjectPoolConfig config) { super(factory, config); } public void testDemo() throws Exception { - GenericObjectPoolConfig config = new GenericObjectPoolConfig<>(); + GenericObjectPoolConfig config = new GenericObjectPoolConfig<>(); // set max total connection config.setMaxTotal(10); // set min idle connection @@ -26,7 +26,7 @@ public void testDemo() throws Exception { DatabendConnectionPool pool = new DatabendConnectionPool(factory, config); // Get a connection from the pool - DatabendConnection connection = pool.borrowObject(); + DatabendConnectionImpl connection = pool.borrowObject(); // connection.uploadStream(); // Use the connection diff --git a/databend-jdbc/src/test/java/com/databend/jdbc/TestBasicDriver.java b/databend-jdbc/src/test/java/com/databend/jdbc/TestBasicDriver.java index 897a1cfa..6b372038 100644 --- a/databend-jdbc/src/test/java/com/databend/jdbc/TestBasicDriver.java +++ b/databend-jdbc/src/test/java/com/databend/jdbc/TestBasicDriver.java @@ -39,7 +39,7 @@ public void setUp() public void testBasic() throws SQLException { try (Connection connection = Utils.createConnection()) { - PaginationOptions p = connection.unwrap(DatabendConnection.class).getPaginationOptions(); + PaginationOptions p = connection.unwrap(DatabendConnectionImpl.class).getPaginationOptions(); Assert.assertEquals(p.getWaitTimeSecs(), PaginationOptions.getDefaultWaitTimeSec()); Assert.assertEquals(p.getMaxRowsInBuffer(), PaginationOptions.getDefaultMaxRowsInBuffer()); Assert.assertEquals(p.getMaxRowsPerPage(), PaginationOptions.getDefaultMaxRowsPerPage()); @@ -67,7 +67,7 @@ public void testExecuteInvalidSql() { @Test(groups = {"IT"}) public void testSchema() throws SQLException { try (Connection connection = Utils.createConnection()) { - PaginationOptions p = connection.unwrap(DatabendConnection.class).getPaginationOptions(); + PaginationOptions p = connection.unwrap(DatabendConnectionImpl.class).getPaginationOptions(); Assert.assertEquals(p.getWaitTimeSecs(), PaginationOptions.getDefaultWaitTimeSec()); Assert.assertEquals(p.getMaxRowsInBuffer(), PaginationOptions.getDefaultMaxRowsInBuffer()); Assert.assertEquals(p.getMaxRowsPerPage(), PaginationOptions.getDefaultMaxRowsPerPage()); @@ -226,7 +226,7 @@ public void testBasicWithProperties() throws SQLException { //INFO databend_query::servers::http::v1::http_query_handlers: receive http query: HttpQueryRequest { session_id: None, session: Some(HttpSessionConf { database: Some("test_basic_driver"), keep_server_session_secs: None, settings: None }), sql: "SELECT 1", pagination: PaginationConf { wait_time_secs: 10, max_rows_in_buffer: 100, max_rows_per_page: 100 }, string_fields: true, stage_attachment: None } try (Connection connection = Utils.createConnection("test_basic_driver", p)) { - PaginationOptions options = connection.unwrap(DatabendConnection.class).getPaginationOptions(); + PaginationOptions options = connection.unwrap(DatabendConnectionImpl.class).getPaginationOptions(); Assert.assertEquals(options.getWaitTimeSecs(), 10); Assert.assertEquals(options.getMaxRowsInBuffer(), 100); Assert.assertEquals(options.getMaxRowsPerPage(), 100); @@ -279,7 +279,7 @@ public void testUpdateSession() try (Connection connection = Utils.createConnection("test_basic_driver")) { connection.createStatement().execute("set max_threads=1"); connection.createStatement().execute("use test_basic_driver_2"); - DatabendSession session = connection.unwrap(DatabendConnection.class).getSession(); + DatabendSession session = connection.unwrap(DatabendConnectionImpl.class).getSession(); Assert.assertEquals(session.getDatabase(), "test_basic_driver_2"); Assert.assertEquals(session.getSettings().get("max_threads"), "1"); } diff --git a/databend-jdbc/src/test/java/com/databend/jdbc/TestCopyInto.java b/databend-jdbc/src/test/java/com/databend/jdbc/TestCopyInto.java index 9e8ec3b9..91fdb271 100644 --- a/databend-jdbc/src/test/java/com/databend/jdbc/TestCopyInto.java +++ b/databend-jdbc/src/test/java/com/databend/jdbc/TestCopyInto.java @@ -11,13 +11,13 @@ public class TestCopyInto { @Test(groups = {"UNIT"}) - public void TestParseSql() { + public void TestGenSQL() { DatabendStage s = DatabendStage.builder().stageName("~").path("a/b/c").build(); List files = new ArrayList<>(); files.add("file.csv"); - String sql = DatabendConnection.getCopyIntoSql("db1", DatabendCopyParams.builder().setFiles(files).setDatabendStage(s).setDatabaseTableName("tb1").build()); + String sql = DatabendConnectionImpl.getCopyIntoSql("db1", DatabendCopyParams.builder().setFiles(files).setDatabendStage(s).setDatabaseTableName("tb1").build()); assertEquals(sql.trim(), "COPY INTO db1.tb1 FROM @~/a/b/c FILES = ('file.csv') FILE_FORMAT = ( type = 'CSV' )"); - sql = DatabendConnection.getCopyIntoSql(null, DatabendCopyParams.builder().setDatabendStage(s).setDatabaseTableName("tb1").build()); + sql = DatabendConnectionImpl.getCopyIntoSql(null, DatabendCopyParams.builder().setDatabendStage(s).setDatabaseTableName("tb1").build()); assertEquals(sql.trim(), "COPY INTO tb1 FROM @~/a/b/c FILE_FORMAT = ( type = 'CSV' )"); } } diff --git a/databend-jdbc/src/test/java/com/databend/jdbc/TestDatabendDatabaseMetaData.java b/databend-jdbc/src/test/java/com/databend/jdbc/TestDatabendDatabaseMetaData.java index b2ec8ea9..1038d65d 100644 --- a/databend-jdbc/src/test/java/com/databend/jdbc/TestDatabendDatabaseMetaData.java +++ b/databend-jdbc/src/test/java/com/databend/jdbc/TestDatabendDatabaseMetaData.java @@ -117,7 +117,7 @@ public void testGetDatabaseProductVersion() Assert.assertTrue(metaData.getDatabaseProductVersion().contains(checkVersion)); if (Compatibility.serverCapability.streamingLoad && Compatibility.driverCapability.streamingLoad) { - DatabendConnection conn = connection.unwrap(DatabendConnection.class); + DatabendConnectionImpl conn = connection.unwrap(DatabendConnectionImpl.class); if (conn.getServerVersion() != null) { String semver = "v" + conn.getServerVersion().toString(); Assert.assertTrue(semver.startsWith(checkVersion), semver); diff --git a/databend-jdbc/src/test/java/com/databend/jdbc/TestDatabendDriverUri.java b/databend-jdbc/src/test/java/com/databend/jdbc/TestDatabendDriverUri.java index b6c0f7f0..82f55bd1 100644 --- a/databend-jdbc/src/test/java/com/databend/jdbc/TestDatabendDriverUri.java +++ b/databend-jdbc/src/test/java/com/databend/jdbc/TestDatabendDriverUri.java @@ -4,6 +4,7 @@ import org.testng.Assert; import org.testng.annotations.Test; +import java.sql.Connection; import java.sql.ResultSet; import java.sql.SQLException; import java.sql.Statement; @@ -176,9 +177,9 @@ public void testUserDatabaseProp() throws SQLException { Assert.assertEquals(uri.getMaxRowsPerPage().intValue(), PaginationOptions.getDefaultMaxRowsPerPage()); Assert.assertFalse(uri.presignedUrlDisabled().booleanValue()); Assert.assertTrue(uri.copyPurge().booleanValue()); - Assert.assertEquals("\\N", uri.nullDisplay().toString()); - Assert.assertEquals("base64", uri.binaryFormat().toString()); - Assert.assertEquals("enable", uri.getSslmode().toString()); + Assert.assertEquals("\\N", uri.nullDisplay()); + Assert.assertEquals("base64", uri.binaryFormat()); + Assert.assertEquals("enable", uri.getSslmode()); } @Test(groups = {"UNIT"}) @@ -235,12 +236,12 @@ public void testFull() throws SQLException { Assert.assertEquals(uri.getMaxRowsPerPage().intValue(), 7); Assert.assertFalse(uri.presignedUrlDisabled().booleanValue()); Assert.assertEquals("null", uri.nullDisplay().toString()); - Assert.assertEquals(false, uri.getStrNullAsNull()); + Assert.assertFalse(uri.getStrNullAsNull()); } @Test(groups = "IT") public void TestSetSchema() throws SQLException { - try (DatabendConnection connection = (DatabendConnection) Utils.createConnection()) { + try (DatabendConnectionImpl connection = (DatabendConnectionImpl) Utils.createConnection()) { connection.createStatement().execute("create or replace database test2"); connection.createStatement().execute("create or replace table test2.test2(id int)"); connection.setSchema("test2"); @@ -256,7 +257,7 @@ public void TestSetSessionSettings() throws SQLException { props.setProperty("session_settings", "max_threads=1,query_tag=tag1"); props.setProperty("user", "databend"); props.setProperty("password", "databend"); - try (DatabendConnection connection = (DatabendConnection) Utils.createConnection("default", props)) { + try (Connection connection = Utils.createConnection("default", props)) { Statement statement = connection.createStatement(); statement.execute("show settings"); ResultSet r = statement.getResultSet(); diff --git a/databend-jdbc/src/test/java/com/databend/jdbc/cloud/TestDatabendPresignClient.java b/databend-jdbc/src/test/java/com/databend/jdbc/TestDatabendPresignClient.java similarity index 85% rename from databend-jdbc/src/test/java/com/databend/jdbc/cloud/TestDatabendPresignClient.java rename to databend-jdbc/src/test/java/com/databend/jdbc/TestDatabendPresignClient.java index b37534ff..c67e5b69 100644 --- a/databend-jdbc/src/test/java/com/databend/jdbc/cloud/TestDatabendPresignClient.java +++ b/databend-jdbc/src/test/java/com/databend/jdbc/TestDatabendPresignClient.java @@ -1,7 +1,7 @@ -package com.databend.jdbc.cloud; +package com.databend.jdbc; -import com.databend.jdbc.DatabendConnection; -import com.databend.jdbc.Utils; +import com.databend.jdbc.cloud.DatabendPresignClient; +import com.databend.jdbc.cloud.DatabendPresignClientV1; import okhttp3.OkHttpClient; import org.testng.annotations.Test; @@ -28,7 +28,7 @@ private String generateRandomCSV(int lines) throws IOException { @Test(groups = {"LOCAL"}) public void uploadFileAPI() throws IOException, SQLException { String filePath = null; - try (DatabendConnection connection = Utils.createConnection().unwrap(DatabendConnection.class)) { + try (DatabendConnectionImpl connection = Utils.createConnection().unwrap(DatabendConnectionImpl.class)) { OkHttpClient client = connection.getHttpClient(); DatabendPresignClient presignClient = new DatabendPresignClientV1(client, connection.getURI().toString()); filePath = generateRandomCSV(10); diff --git a/databend-jdbc/src/test/java/com/databend/jdbc/TestFileTransfer.java b/databend-jdbc/src/test/java/com/databend/jdbc/TestFileTransfer.java index b53b19c9..8c1b946b 100644 --- a/databend-jdbc/src/test/java/com/databend/jdbc/TestFileTransfer.java +++ b/databend-jdbc/src/test/java/com/databend/jdbc/TestFileTransfer.java @@ -112,9 +112,9 @@ public void testFileTransfer() Connection connection = Utils.createConnection()) { String stageName = "test_stage"; DatabendConnection databendConnection = connection.unwrap(DatabendConnection.class); - PresignContext.createStageIfNotExists(databendConnection, stageName); + PresignContext.createStageIfNotExists(connection, stageName); databendConnection.uploadStream(stageName, "jdbc/test/", fileInputStream, "test.csv", f.length(), false); - downloaded = databendConnection.downloadStream(stageName, "jdbc/test/test.csv", false); + downloaded = databendConnection.downloadStream(stageName, "jdbc/test/test.csv"); byte[] arr = streamToByteArray(downloaded); Assert.assertEquals(arr.length, f.length()); } finally { @@ -135,9 +135,9 @@ public void testFileTransferThroughAPI() throws SQLException, IOException { String stageName = "test_stage_np"; DatabendConnection databendConnection = connection.unwrap(DatabendConnection.class); - PresignContext.createStageIfNotExists(databendConnection, stageName); + PresignContext.createStageIfNotExists(connection, stageName); databendConnection.uploadStream(stageName, "jdbc/test/", fileInputStream, "test.csv", f.length(), false); - InputStream downloaded = databendConnection.downloadStream(stageName, "jdbc/test/test.csv", false); + InputStream downloaded = databendConnection.downloadStream(stageName, "jdbc/test/test.csv"); byte[] arr = streamToByteArray(downloaded); Assert.assertEquals(arr.length, f.length()); } finally { @@ -152,8 +152,8 @@ public void testCopyInto() throws IOException, SQLException { try (FileInputStream fileInputStream = new FileInputStream(f)) { Connection connection = Utils.createConnection(); String stageName = "test_stage"; - DatabendConnection databendConnection = connection.unwrap(DatabendConnection.class); - PresignContext.createStageIfNotExists(databendConnection, stageName); + FileTransferAPI databendConnection = connection.unwrap(FileTransferAPI.class); + PresignContext.createStageIfNotExists(connection, stageName); databendConnection.uploadStream(stageName, "jdbc/c2/", fileInputStream, "complex.csv", f.length(), false); fileInputStream.close(); DatabendStage s = DatabendStage.builder().stageName(stageName).path("jdbc/c2/").build(); @@ -170,8 +170,8 @@ public void testCopyInto() throws IOException, SQLException { @DataProvider(name = "streamingLoad") private Object[][] provideTestData() { return new Object[][] { - {"streaming"}, - {"stage"} + {"STREAMING"}, + {"STAGE"} }; } @@ -197,7 +197,7 @@ public void testLoadStreamToTable(String method) throws IOException, SQLExceptio statement.execute("create or replace table test_load(i int, a Variant, b string)"); DatabendConnection databendConnection = connection.unwrap(DatabendConnection.class); String sql = "insert into test_load from @_databend_load file_format=(type=csv)"; - int nUpdate = databendConnection.loadStreamToTable(sql, fileInputStream, f.length(), method); + int nUpdate = databendConnection.loadStreamToTable(sql, fileInputStream, f.length(), DatabendConnection.LoadMethod.valueOf(method)); Assert.assertEquals(nUpdate, 10); fileInputStream.close(); ResultSet r = statement.executeQuery("SELECT * FROM test_load"); diff --git a/databend-jdbc/src/test/java/com/databend/jdbc/TestHeartbeat.java b/databend-jdbc/src/test/java/com/databend/jdbc/TestHeartbeat.java index 893e899e..e89f3116 100644 --- a/databend-jdbc/src/test/java/com/databend/jdbc/TestHeartbeat.java +++ b/databend-jdbc/src/test/java/com/databend/jdbc/TestHeartbeat.java @@ -48,7 +48,7 @@ public void testHeartbeat() throws SQLException, InterruptedException { statement.close(); Thread.sleep(5000); - Assert.assertTrue(c1.unwrap(DatabendConnection.class).isHeartbeatStopped()); + Assert.assertTrue(c1.unwrap(DatabendConnectionImpl.class).isHeartbeatStopped()); } } } diff --git a/databend-jdbc/src/test/java/com/databend/jdbc/TestMultiHost.java b/databend-jdbc/src/test/java/com/databend/jdbc/TestMultiHost.java index 85cee222..8bacca10 100644 --- a/databend-jdbc/src/test/java/com/databend/jdbc/TestMultiHost.java +++ b/databend-jdbc/src/test/java/com/databend/jdbc/TestMultiHost.java @@ -128,7 +128,7 @@ public void testUnSupportedAutoDiscovery() statement.execute("select value from system.configs where name = 'http_handler_port';"); ResultSet r = statement.getResultSet(); r.next(); - DatabendConnection dbc = (DatabendConnection) connection; + DatabendConnectionImpl dbc = (DatabendConnectionImpl) connection; // automatically Assert.assertFalse(dbc.isAutoDiscovery()); } diff --git a/databend-jdbc/src/test/java/com/databend/jdbc/TestPrepareStatement.java b/databend-jdbc/src/test/java/com/databend/jdbc/TestPrepareStatement.java index 150a62b7..4c53949a 100644 --- a/databend-jdbc/src/test/java/com/databend/jdbc/TestPrepareStatement.java +++ b/databend-jdbc/src/test/java/com/databend/jdbc/TestPrepareStatement.java @@ -465,8 +465,8 @@ public void testAllPreparedStatement() throws SQLException { @Test(groups = "IT") public void shouldBuildStageAttachmentWithFileFormatOptions() throws SQLException { Connection conn = Utils.createConnection(); - Assert.assertEquals("", conn.unwrap(DatabendConnection.class).binaryFormat()); - StageAttachment stageAttachment = DatabendPreparedStatement.buildStateAttachment((DatabendConnection) conn, + Assert.assertEquals("", conn.unwrap(DatabendConnectionImpl.class).binaryFormat()); + StageAttachment stageAttachment = DatabendPreparedStatement.buildStateAttachment((DatabendConnectionImpl) conn, "stagePath"); Assert.assertFalse(stageAttachment.getFileFormatOptions().containsKey("binary_format")); diff --git a/databend-jdbc/src/test/java/com/databend/jdbc/TestPresignContext.java b/databend-jdbc/src/test/java/com/databend/jdbc/TestPresignContext.java index c49e69b4..e10abfb9 100644 --- a/databend-jdbc/src/test/java/com/databend/jdbc/TestPresignContext.java +++ b/databend-jdbc/src/test/java/com/databend/jdbc/TestPresignContext.java @@ -18,16 +18,17 @@ public void TestPreisgnUrlBuild() { @Test(groups = {"IT"}) public void TestGetPresignUrl() throws SQLException { - DatabendConnection connection = (DatabendConnection) Utils.createConnection(); - PresignContext ctx = PresignContext.getPresignContext(connection, PresignContext.PresignMethod.UPLOAD, null, "test.csv"); - Assert.assertNotNull(ctx); - Assert.assertNotNull(ctx.getUrl()); - Assert.assertNotNull(ctx.getHeaders()); + try (DatabendConnectionImpl connection = (DatabendConnectionImpl) Utils.createConnection()) { + PresignContext ctx = PresignContext.getPresignContext(connection, PresignContext.PresignMethod.UPLOAD, null, "test.csv"); + Assert.assertNotNull(ctx); + Assert.assertNotNull(ctx.getUrl()); + Assert.assertNotNull(ctx.getHeaders()); + } } @Test(groups = {"IT"}) public void TestGetPresignUrlCase2() throws SQLException { - try (DatabendConnection connection = (DatabendConnection) Utils.createConnection()) { + try (DatabendConnectionImpl connection = (DatabendConnectionImpl) Utils.createConnection()) { String stageName = "test_stage"; PresignContext.createStageIfNotExists(connection, stageName); PresignContext ctx = PresignContext.getPresignContext(connection, PresignContext.PresignMethod.UPLOAD, stageName, "a/b/d/test.csv"); From 1c32acd1e50c26dc4f0e12533f3c84e733b73ff3 Mon Sep 17 00:00:00 2001 From: Yang Xiufeng Date: Mon, 15 Sep 2025 00:01:33 +0800 Subject: [PATCH 4/5] fix --- .../databend/jdbc/DatabendConnectionImpl.java | 4 +- .../jdbc/DatabendPreparedStatement.java | 61 ++++++------------- .../com/databend/jdbc/PresignContext.java | 13 ++-- .../java/com/databend/jdbc/Compatibility.java | 25 ++++++++ .../com/databend/jdbc/TestBasicDriver.java | 11 ++-- .../jdbc/TestDatabendDatabaseMetaData.java | 10 +-- .../databend/jdbc/TestDatabendDriverUri.java | 2 +- .../jdbc/TestDatabendPresignClient.java | 10 ++- .../com/databend/jdbc/TestFileTransfer.java | 23 ++++--- .../java/com/databend/jdbc/TestHeartbeat.java | 2 +- .../java/com/databend/jdbc/TestMultiHost.java | 7 +-- .../databend/jdbc/TestPrepareStatement.java | 3 +- .../com/databend/jdbc/TestPresignContext.java | 16 +++-- tests/compatibility/testng.xml | 1 + 14 files changed, 101 insertions(+), 87 deletions(-) diff --git a/databend-jdbc/src/main/java/com/databend/jdbc/DatabendConnectionImpl.java b/databend-jdbc/src/main/java/com/databend/jdbc/DatabendConnectionImpl.java index 0c5db399..f0be436b 100644 --- a/databend-jdbc/src/main/java/com/databend/jdbc/DatabendConnectionImpl.java +++ b/databend-jdbc/src/main/java/com/databend/jdbc/DatabendConnectionImpl.java @@ -967,7 +967,7 @@ public void uploadStream(String stageName, String destPrefix, InputStream inputS } else { // logger.log(Level.FINE, "presign to @" + s + "/" + dest); long presignStartTime = System.nanoTime(); - PresignContext ctx = PresignContext.getPresignContext(this, PresignContext.PresignMethod.UPLOAD, s, dest); + PresignContext ctx = PresignContext.newPresignContext(this, PresignContext.PresignMethod.UPLOAD, s, dest); long presignEndTime = System.nanoTime(); if (this.debug()) { logger.info("presign cost time: " + (presignEndTime - presignStartTime) / 1000000.0 + "ms"); @@ -994,7 +994,7 @@ public InputStream downloadStream(String stageName, String path) String s = stageName.replaceAll("/$", ""); DatabendPresignClient cli = new DatabendPresignClientV1(httpClient, this.httpUri.toString()); try { - PresignContext ctx = PresignContext.getPresignContext(this, PresignContext.PresignMethod.DOWNLOAD, s, path); + PresignContext ctx = PresignContext.newPresignContext(this, PresignContext.PresignMethod.DOWNLOAD, s, path); Headers h = ctx.getHeaders(); String presignUrl = ctx.getUrl(); return cli.presignDownloadStream(h, presignUrl); diff --git a/databend-jdbc/src/main/java/com/databend/jdbc/DatabendPreparedStatement.java b/databend-jdbc/src/main/java/com/databend/jdbc/DatabendPreparedStatement.java index 77a1e69d..4e87e3d8 100644 --- a/databend-jdbc/src/main/java/com/databend/jdbc/DatabendPreparedStatement.java +++ b/databend-jdbc/src/main/java/com/databend/jdbc/DatabendPreparedStatement.java @@ -9,57 +9,23 @@ import org.joda.time.format.DateTimeFormatter; import org.joda.time.format.ISODateTimeFormat; -import java.io.ByteArrayInputStream; -import java.io.ByteArrayOutputStream; -import java.io.File; -import java.io.FileInputStream; -import java.io.IOException; -import java.io.InputStream; -import java.io.Reader; +import java.io.*; import java.math.BigDecimal; import java.net.URL; import java.nio.charset.StandardCharsets; -import java.sql.Array; -import java.sql.Blob; -import java.sql.Clob; import java.sql.Date; -import java.sql.NClob; -import java.sql.ParameterMetaData; -import java.sql.PreparedStatement; -import java.sql.Ref; -import java.sql.ResultSet; -import java.sql.ResultSetMetaData; -import java.sql.RowId; -import java.sql.SQLException; -import java.sql.SQLFeatureNotSupportedException; -import java.sql.SQLXML; -import java.sql.Time; -import java.sql.Timestamp; -import java.sql.Types; -import java.time.LocalDate; -import java.time.LocalDateTime; -import java.time.LocalTime; -import java.time.OffsetDateTime; -import java.time.OffsetTime; +import java.sql.*; +import java.time.*; import java.time.format.DateTimeFormatterBuilder; -import java.util.ArrayList; -import java.util.Arrays; -import java.util.Base64; -import java.util.Calendar; -import java.util.Collections; -import java.util.HashMap; -import java.util.List; -import java.util.Map; -import java.util.Objects; -import java.util.UUID; +import java.util.*; import java.util.function.Consumer; import java.util.logging.Logger; import java.util.regex.Matcher; import java.util.stream.Collectors; +import static com.databend.jdbc.DatabendConstant.*; import static com.databend.jdbc.ObjectCasts.*; import static com.databend.jdbc.StatementUtil.replaceParameterMarksWithValues; -import static com.databend.jdbc.DatabendConstant.*; import static java.lang.String.format; import static java.time.format.DateTimeFormatter.ISO_LOCAL_DATE; import static java.time.format.DateTimeFormatter.ISO_LOCAL_TIME; @@ -162,7 +128,7 @@ private StageAttachment uploadBatches() throws SQLException { } File saved = batchInsertUtils.saveBatchToCSV(batchValuesCSV); try (FileInputStream fis = new FileInputStream(saved)) { - DatabendConnectionImpl c = (DatabendConnectionImpl) getConnection(); + Connection c = getConnection(); String uuid = UUID.randomUUID().toString().replace("-", ""); // format %Y/%m/%d/%H/%M/%S/fileName.csv String stagePrefix = String.format("%s/%s/%s/%s/%s/%s/%s/", @@ -175,9 +141,9 @@ private StageAttachment uploadBatches() throws SQLException { uuid); String fileName = saved.getName(); // upload to stage - c.uploadStream(null, stagePrefix, fis, fileName, saved.length(), false); + c.unwrap(DatabendConnection.class).uploadStream(null, stagePrefix, fis, fileName, saved.length(), false); String stagePath = "@~/" + stagePrefix + fileName; - return buildStateAttachment(c, stagePath); + return buildStateAttachment((DatabendConnection) c, stagePath); } catch (Exception e) { throw new SQLException(e); } finally { @@ -191,18 +157,25 @@ private StageAttachment uploadBatches() throws SQLException { } } + // only for compat test + static StageAttachment buildStateAttachment(DatabendConnection conn, String stagePath) { + return buildStateAttachment((Connection) conn, stagePath); + } + /** * This method is used to build a StageAttachment object which represents a * stage in Databend. * A stage in Databend is a temporary storage area where data files are stored * before being loaded into the Databend database. * - * @param connection The DatabendConnection object which contains the connection + * @param conn The DatabendConnection object which contains the connection * details to the Databend database. * @param stagePath The path of the stage in the Databend database. * @return A StageAttachment object which contains the details of the stage. */ - static StageAttachment buildStateAttachment(DatabendConnectionImpl connection, String stagePath) { + static StageAttachment buildStateAttachment(Connection conn, String stagePath) { + DatabendConnectionImpl connection = (DatabendConnectionImpl) conn; + Map fileFormatOptions = new HashMap<>(); if (!Objects.equals(connection.binaryFormat(), "")) { fileFormatOptions.put("binary_format", String.valueOf(connection.binaryFormat())); diff --git a/databend-jdbc/src/main/java/com/databend/jdbc/PresignContext.java b/databend-jdbc/src/main/java/com/databend/jdbc/PresignContext.java index 0f506f4f..15f0ef0d 100644 --- a/databend-jdbc/src/main/java/com/databend/jdbc/PresignContext.java +++ b/databend-jdbc/src/main/java/com/databend/jdbc/PresignContext.java @@ -28,19 +28,18 @@ private PresignContext(PresignMethod method, String stageName, String fileName, this.url = url; } - public static void createStageIfNotExists(Connection connection, String stageName) throws SQLException { + static void createStageIfNotExists(Connection connection, String stageName) throws SQLException { String sql = String.format("CREATE STAGE IF NOT EXISTS %s", stageName); Statement statement = connection.createStatement(); statement.execute(sql); } - public static void dropStageIfExists(DatabendConnectionImpl connection, String stageName) throws SQLException { - String sql = String.format("DROP STAGE IF EXISTS %s", stageName); - Statement statement = connection.createStatement(); - statement.execute(sql); + // only for compat test + static PresignContext getPresignContext(DatabendConnection connection, PresignMethod method, String stageName, String fileName) + throws SQLException { + return newPresignContext((Connection) connection, method, stageName, fileName); } - - public static PresignContext getPresignContext(DatabendConnectionImpl connection, PresignMethod method, String stageName, String fileName) + static PresignContext newPresignContext(Connection connection, PresignMethod method, String stageName, String fileName) throws SQLException { requireNonNull(connection, "connection is null"); requireNonNull(method, "method is null"); diff --git a/databend-jdbc/src/test/java/com/databend/jdbc/Compatibility.java b/databend-jdbc/src/test/java/com/databend/jdbc/Compatibility.java index b5fa8df0..aaa9427a 100644 --- a/databend-jdbc/src/test/java/com/databend/jdbc/Compatibility.java +++ b/databend-jdbc/src/test/java/com/databend/jdbc/Compatibility.java @@ -2,6 +2,10 @@ import com.vdurmont.semver4j.Semver; +import java.lang.reflect.InvocationTargetException; +import java.lang.reflect.Method; +import java.sql.Connection; + public class Compatibility { public static class Capability { boolean streamingLoad; @@ -65,4 +69,25 @@ public static boolean skipBugLowerThenOrEqualTo(String serverVersionBug, String } return false; } + + boolean isNewInterface() { + return driverVersion == null || driverVersion.isGreaterThanOrEqualTo(new Semver("4.0.1")); + } + + static Object invokeMethod(Object target, String methodName, + Class[] parameterTypes, Object[] args) { + Class targetClass = target.getClass(); + Method method = null; + try { + method = targetClass.getDeclaredMethod(methodName, parameterTypes); + method.setAccessible(true); + return method.invoke(target, args); + } catch (NoSuchMethodException | IllegalAccessException | InvocationTargetException e) { + throw new RuntimeException(e); + } + } + + static Object invokeMethodNoArg(Object target, String methodName) { + return invokeMethod(target, methodName, new Class[0], new Object[0]); + } } diff --git a/databend-jdbc/src/test/java/com/databend/jdbc/TestBasicDriver.java b/databend-jdbc/src/test/java/com/databend/jdbc/TestBasicDriver.java index 6b372038..ae11f87e 100644 --- a/databend-jdbc/src/test/java/com/databend/jdbc/TestBasicDriver.java +++ b/databend-jdbc/src/test/java/com/databend/jdbc/TestBasicDriver.java @@ -2,6 +2,7 @@ import com.databend.client.DatabendSession; import com.databend.client.PaginationOptions; +import com.vdurmont.semver4j.Semver; import org.testng.Assert; import org.testng.annotations.BeforeTest; import org.testng.annotations.Test; @@ -39,7 +40,7 @@ public void setUp() public void testBasic() throws SQLException { try (Connection connection = Utils.createConnection()) { - PaginationOptions p = connection.unwrap(DatabendConnectionImpl.class).getPaginationOptions(); + PaginationOptions p = (PaginationOptions) Compatibility.invokeMethodNoArg(connection, "getPaginationOptions"); Assert.assertEquals(p.getWaitTimeSecs(), PaginationOptions.getDefaultWaitTimeSec()); Assert.assertEquals(p.getMaxRowsInBuffer(), PaginationOptions.getDefaultMaxRowsInBuffer()); Assert.assertEquals(p.getMaxRowsPerPage(), PaginationOptions.getDefaultMaxRowsPerPage()); @@ -67,10 +68,6 @@ public void testExecuteInvalidSql() { @Test(groups = {"IT"}) public void testSchema() throws SQLException { try (Connection connection = Utils.createConnection()) { - PaginationOptions p = connection.unwrap(DatabendConnectionImpl.class).getPaginationOptions(); - Assert.assertEquals(p.getWaitTimeSecs(), PaginationOptions.getDefaultWaitTimeSec()); - Assert.assertEquals(p.getMaxRowsInBuffer(), PaginationOptions.getDefaultMaxRowsInBuffer()); - Assert.assertEquals(p.getMaxRowsPerPage(), PaginationOptions.getDefaultMaxRowsPerPage()); DatabendStatement statement = (DatabendStatement) connection.createStatement(); statement.execute("set global timezone='Asia/Shanghai';"); statement.execute("SELEcT schema_name as TABLE_SCHEM, catalog_name as TABLE_CATALOG FROM information_schema.schemata where schema_name = 'default' order by catalog_name, schema_name"); @@ -226,7 +223,7 @@ public void testBasicWithProperties() throws SQLException { //INFO databend_query::servers::http::v1::http_query_handlers: receive http query: HttpQueryRequest { session_id: None, session: Some(HttpSessionConf { database: Some("test_basic_driver"), keep_server_session_secs: None, settings: None }), sql: "SELECT 1", pagination: PaginationConf { wait_time_secs: 10, max_rows_in_buffer: 100, max_rows_per_page: 100 }, string_fields: true, stage_attachment: None } try (Connection connection = Utils.createConnection("test_basic_driver", p)) { - PaginationOptions options = connection.unwrap(DatabendConnectionImpl.class).getPaginationOptions(); + PaginationOptions options = (PaginationOptions) Compatibility.invokeMethodNoArg(connection, "getPaginationOptions"); Assert.assertEquals(options.getWaitTimeSecs(), 10); Assert.assertEquals(options.getMaxRowsInBuffer(), 100); Assert.assertEquals(options.getMaxRowsPerPage(), 100); @@ -279,7 +276,7 @@ public void testUpdateSession() try (Connection connection = Utils.createConnection("test_basic_driver")) { connection.createStatement().execute("set max_threads=1"); connection.createStatement().execute("use test_basic_driver_2"); - DatabendSession session = connection.unwrap(DatabendConnectionImpl.class).getSession(); + DatabendSession session = (DatabendSession) Compatibility.invokeMethodNoArg(connection, "getSession"); Assert.assertEquals(session.getDatabase(), "test_basic_driver_2"); Assert.assertEquals(session.getSettings().get("max_threads"), "1"); } diff --git a/databend-jdbc/src/test/java/com/databend/jdbc/TestDatabendDatabaseMetaData.java b/databend-jdbc/src/test/java/com/databend/jdbc/TestDatabendDatabaseMetaData.java index 1038d65d..072ada10 100644 --- a/databend-jdbc/src/test/java/com/databend/jdbc/TestDatabendDatabaseMetaData.java +++ b/databend-jdbc/src/test/java/com/databend/jdbc/TestDatabendDatabaseMetaData.java @@ -1,11 +1,13 @@ package com.databend.jdbc; +import com.vdurmont.semver4j.Semver; import org.testng.Assert; import org.testng.annotations.BeforeTest; import org.testng.annotations.Test; import java.math.BigDecimal; import java.math.RoundingMode; +import java.net.URI; import java.sql.*; import java.util.ArrayList; import java.util.List; @@ -117,11 +119,11 @@ public void testGetDatabaseProductVersion() Assert.assertTrue(metaData.getDatabaseProductVersion().contains(checkVersion)); if (Compatibility.serverCapability.streamingLoad && Compatibility.driverCapability.streamingLoad) { - DatabendConnectionImpl conn = connection.unwrap(DatabendConnectionImpl.class); - if (conn.getServerVersion() != null) { - String semver = "v" + conn.getServerVersion().toString(); + Semver ver = (Semver) Compatibility.invokeMethodNoArg(connection, "getServerVersion"); + if (ver != null) { + String semver = "v" + ver.toString(); Assert.assertTrue(semver.startsWith(checkVersion), semver); - Assert.assertNotNull(conn.getServerCapability()); + Assert.assertNotNull(Compatibility.invokeMethodNoArg(connection, "getServerCapability")); } } } diff --git a/databend-jdbc/src/test/java/com/databend/jdbc/TestDatabendDriverUri.java b/databend-jdbc/src/test/java/com/databend/jdbc/TestDatabendDriverUri.java index 82f55bd1..01536301 100644 --- a/databend-jdbc/src/test/java/com/databend/jdbc/TestDatabendDriverUri.java +++ b/databend-jdbc/src/test/java/com/databend/jdbc/TestDatabendDriverUri.java @@ -241,7 +241,7 @@ public void testFull() throws SQLException { @Test(groups = "IT") public void TestSetSchema() throws SQLException { - try (DatabendConnectionImpl connection = (DatabendConnectionImpl) Utils.createConnection()) { + try (Connection connection = Utils.createConnection()) { connection.createStatement().execute("create or replace database test2"); connection.createStatement().execute("create or replace table test2.test2(id int)"); connection.setSchema("test2"); diff --git a/databend-jdbc/src/test/java/com/databend/jdbc/TestDatabendPresignClient.java b/databend-jdbc/src/test/java/com/databend/jdbc/TestDatabendPresignClient.java index c67e5b69..62697278 100644 --- a/databend-jdbc/src/test/java/com/databend/jdbc/TestDatabendPresignClient.java +++ b/databend-jdbc/src/test/java/com/databend/jdbc/TestDatabendPresignClient.java @@ -1,11 +1,14 @@ package com.databend.jdbc; +import com.databend.client.PaginationOptions; import com.databend.jdbc.cloud.DatabendPresignClient; import com.databend.jdbc.cloud.DatabendPresignClientV1; import okhttp3.OkHttpClient; import org.testng.annotations.Test; import java.io.*; +import java.net.URI; +import java.sql.Connection; import java.sql.SQLException; public class TestDatabendPresignClient { @@ -28,9 +31,10 @@ private String generateRandomCSV(int lines) throws IOException { @Test(groups = {"LOCAL"}) public void uploadFileAPI() throws IOException, SQLException { String filePath = null; - try (DatabendConnectionImpl connection = Utils.createConnection().unwrap(DatabendConnectionImpl.class)) { - OkHttpClient client = connection.getHttpClient(); - DatabendPresignClient presignClient = new DatabendPresignClientV1(client, connection.getURI().toString()); + try (Connection connection = Utils.createConnection()) { + URI uri = (URI) Compatibility.invokeMethodNoArg(connection, "getURI"); + OkHttpClient client = (OkHttpClient) Compatibility.invokeMethodNoArg(connection, "getHttpClient"); + DatabendPresignClient presignClient = new DatabendPresignClientV1(client,uri.toString()); filePath = generateRandomCSV(10); File file = new File(filePath); InputStream inputStream = new FileInputStream(file); diff --git a/databend-jdbc/src/test/java/com/databend/jdbc/TestFileTransfer.java b/databend-jdbc/src/test/java/com/databend/jdbc/TestFileTransfer.java index 8c1b946b..0d9f99f3 100644 --- a/databend-jdbc/src/test/java/com/databend/jdbc/TestFileTransfer.java +++ b/databend-jdbc/src/test/java/com/databend/jdbc/TestFileTransfer.java @@ -111,10 +111,13 @@ public void testFileTransfer() try (FileInputStream fileInputStream = new FileInputStream(f); Connection connection = Utils.createConnection()) { String stageName = "test_stage"; - DatabendConnection databendConnection = connection.unwrap(DatabendConnection.class); - PresignContext.createStageIfNotExists(connection, stageName); + // use FileTransferAPI for compat test + FileTransferAPI databendConnection = connection.unwrap(FileTransferAPI.class); + String sql = String.format("CREATE STAGE IF NOT EXISTS %s", stageName); + Statement statement = connection.createStatement(); + statement.execute(sql); databendConnection.uploadStream(stageName, "jdbc/test/", fileInputStream, "test.csv", f.length(), false); - downloaded = databendConnection.downloadStream(stageName, "jdbc/test/test.csv"); + downloaded = databendConnection.downloadStream(stageName, "jdbc/test/test.csv", false); byte[] arr = streamToByteArray(downloaded); Assert.assertEquals(arr.length, f.length()); } finally { @@ -134,10 +137,13 @@ public void testFileTransferThroughAPI() throws SQLException, IOException { String stageName = "test_stage_np"; - DatabendConnection databendConnection = connection.unwrap(DatabendConnection.class); - PresignContext.createStageIfNotExists(connection, stageName); + // use FileTransferAPI for compat test + FileTransferAPI databendConnection = connection.unwrap(FileTransferAPI.class); + String sql = String.format("CREATE STAGE IF NOT EXISTS %s", stageName); + Statement statement = connection.createStatement(); + statement.execute(sql); databendConnection.uploadStream(stageName, "jdbc/test/", fileInputStream, "test.csv", f.length(), false); - InputStream downloaded = databendConnection.downloadStream(stageName, "jdbc/test/test.csv"); + InputStream downloaded = databendConnection.downloadStream(stageName, "jdbc/test/test.csv", false); byte[] arr = streamToByteArray(downloaded); Assert.assertEquals(arr.length, f.length()); } finally { @@ -152,8 +158,11 @@ public void testCopyInto() throws IOException, SQLException { try (FileInputStream fileInputStream = new FileInputStream(f)) { Connection connection = Utils.createConnection(); String stageName = "test_stage"; + // use FileTransferAPI for compat test FileTransferAPI databendConnection = connection.unwrap(FileTransferAPI.class); - PresignContext.createStageIfNotExists(connection, stageName); + String sql = String.format("CREATE STAGE IF NOT EXISTS %s", stageName); + Statement statement = connection.createStatement(); + statement.execute(sql); databendConnection.uploadStream(stageName, "jdbc/c2/", fileInputStream, "complex.csv", f.length(), false); fileInputStream.close(); DatabendStage s = DatabendStage.builder().stageName(stageName).path("jdbc/c2/").build(); diff --git a/databend-jdbc/src/test/java/com/databend/jdbc/TestHeartbeat.java b/databend-jdbc/src/test/java/com/databend/jdbc/TestHeartbeat.java index e89f3116..98caebc9 100644 --- a/databend-jdbc/src/test/java/com/databend/jdbc/TestHeartbeat.java +++ b/databend-jdbc/src/test/java/com/databend/jdbc/TestHeartbeat.java @@ -48,7 +48,7 @@ public void testHeartbeat() throws SQLException, InterruptedException { statement.close(); Thread.sleep(5000); - Assert.assertTrue(c1.unwrap(DatabendConnectionImpl.class).isHeartbeatStopped()); + Assert.assertTrue((boolean) Compatibility.invokeMethodNoArg(c1, "isHeartbeatStopped")); } } } diff --git a/databend-jdbc/src/test/java/com/databend/jdbc/TestMultiHost.java b/databend-jdbc/src/test/java/com/databend/jdbc/TestMultiHost.java index 8bacca10..8a24fc4b 100644 --- a/databend-jdbc/src/test/java/com/databend/jdbc/TestMultiHost.java +++ b/databend-jdbc/src/test/java/com/databend/jdbc/TestMultiHost.java @@ -122,15 +122,14 @@ public void testAutoDiscovery() @Test(groups = {"IT", "MULTI_HOST"}) public void testUnSupportedAutoDiscovery() - throws SQLException { + throws Exception { try (Connection connection = createConnection(UNSUPPORT_AUTO_DISCOVERY_JDBC_URL)) { DatabendStatement statement = (DatabendStatement) connection.createStatement(); statement.execute("select value from system.configs where name = 'http_handler_port';"); ResultSet r = statement.getResultSet(); r.next(); - DatabendConnectionImpl dbc = (DatabendConnectionImpl) connection; - // automatically - Assert.assertFalse(dbc.isAutoDiscovery()); + + Assert.assertFalse((boolean) Compatibility.invokeMethodNoArg(connection, "isAutoDiscovery")); } } diff --git a/databend-jdbc/src/test/java/com/databend/jdbc/TestPrepareStatement.java b/databend-jdbc/src/test/java/com/databend/jdbc/TestPrepareStatement.java index 4c53949a..74de2e55 100644 --- a/databend-jdbc/src/test/java/com/databend/jdbc/TestPrepareStatement.java +++ b/databend-jdbc/src/test/java/com/databend/jdbc/TestPrepareStatement.java @@ -465,8 +465,7 @@ public void testAllPreparedStatement() throws SQLException { @Test(groups = "IT") public void shouldBuildStageAttachmentWithFileFormatOptions() throws SQLException { Connection conn = Utils.createConnection(); - Assert.assertEquals("", conn.unwrap(DatabendConnectionImpl.class).binaryFormat()); - StageAttachment stageAttachment = DatabendPreparedStatement.buildStateAttachment((DatabendConnectionImpl) conn, + StageAttachment stageAttachment = DatabendPreparedStatement.buildStateAttachment((DatabendConnection) conn, "stagePath"); Assert.assertFalse(stageAttachment.getFileFormatOptions().containsKey("binary_format")); diff --git a/databend-jdbc/src/test/java/com/databend/jdbc/TestPresignContext.java b/databend-jdbc/src/test/java/com/databend/jdbc/TestPresignContext.java index e10abfb9..d222d8f5 100644 --- a/databend-jdbc/src/test/java/com/databend/jdbc/TestPresignContext.java +++ b/databend-jdbc/src/test/java/com/databend/jdbc/TestPresignContext.java @@ -3,7 +3,9 @@ import org.testng.Assert; import org.testng.annotations.Test; +import java.sql.Connection; import java.sql.SQLException; +import java.sql.Statement; public class TestPresignContext { @Test(groups = {"UNIT"}) @@ -18,8 +20,8 @@ public void TestPreisgnUrlBuild() { @Test(groups = {"IT"}) public void TestGetPresignUrl() throws SQLException { - try (DatabendConnectionImpl connection = (DatabendConnectionImpl) Utils.createConnection()) { - PresignContext ctx = PresignContext.getPresignContext(connection, PresignContext.PresignMethod.UPLOAD, null, "test.csv"); + try (Connection connection = Utils.createConnection()) { + PresignContext ctx = PresignContext.getPresignContext((DatabendConnection) connection, PresignContext.PresignMethod.UPLOAD, null, "test.csv"); Assert.assertNotNull(ctx); Assert.assertNotNull(ctx.getUrl()); Assert.assertNotNull(ctx.getHeaders()); @@ -28,10 +30,14 @@ public void TestGetPresignUrl() throws SQLException { @Test(groups = {"IT"}) public void TestGetPresignUrlCase2() throws SQLException { - try (DatabendConnectionImpl connection = (DatabendConnectionImpl) Utils.createConnection()) { + try (Connection connection = Utils.createConnection()) { String stageName = "test_stage"; - PresignContext.createStageIfNotExists(connection, stageName); - PresignContext ctx = PresignContext.getPresignContext(connection, PresignContext.PresignMethod.UPLOAD, stageName, "a/b/d/test.csv"); + + String sql = String.format("CREATE STAGE IF NOT EXISTS %s", stageName); + Statement statement = connection.createStatement(); + statement.execute(sql); + + PresignContext ctx = PresignContext.getPresignContext((DatabendConnection) connection, PresignContext.PresignMethod.UPLOAD, stageName, "a/b/d/test.csv"); Assert.assertNotNull(ctx); Assert.assertNotNull(ctx.getUrl()); Assert.assertNotNull(ctx.getHeaders()); diff --git a/tests/compatibility/testng.xml b/tests/compatibility/testng.xml index 282f2ea3..347006d3 100644 --- a/tests/compatibility/testng.xml +++ b/tests/compatibility/testng.xml @@ -3,6 +3,7 @@ + From 67043ca73ff95d3f9e30db5bb3712f679210f3a2 Mon Sep 17 00:00:00 2001 From: Yang Xiufeng Date: Mon, 15 Sep 2025 11:23:13 +0800 Subject: [PATCH 5/5] hide DatabendConnectionImpl --- .../databend/jdbc/DatabendConnectionImpl.java | 6 +++--- .../examples/DatabendConnectionFactory.java | 19 +++++++++---------- .../jdbc/examples/DatabendConnectionPool.java | 10 +++++----- 3 files changed, 17 insertions(+), 18 deletions(-) diff --git a/databend-jdbc/src/main/java/com/databend/jdbc/DatabendConnectionImpl.java b/databend-jdbc/src/main/java/com/databend/jdbc/DatabendConnectionImpl.java index f0be436b..6f8a45f1 100644 --- a/databend-jdbc/src/main/java/com/databend/jdbc/DatabendConnectionImpl.java +++ b/databend-jdbc/src/main/java/com/databend/jdbc/DatabendConnectionImpl.java @@ -69,7 +69,7 @@ import static java.util.concurrent.TimeUnit.MILLISECONDS; -public class DatabendConnectionImpl implements Connection, DatabendConnection, FileTransferAPI, Consumer { +class DatabendConnectionImpl implements Connection, DatabendConnection, FileTransferAPI, Consumer { private static final Logger logger = Logger.getLogger(DatabendConnectionImpl.class.getPackage().getName()); private static final String STREAMING_LOAD_PATH = "/v1/streaming_load"; private static final String LOGIN_PATH = "/v1/session/login"; @@ -88,7 +88,7 @@ public class DatabendConnectionImpl implements Connection, DatabendConnection, F private boolean autoDiscovery; private final AtomicReference session = new AtomicReference<>(); - private String routeHint = ""; + private String routeHint; private final AtomicReference lastNodeID = new AtomicReference<>(); private Semver serverVersion = null; private Capability serverCapability = null; @@ -259,7 +259,7 @@ private void setSession(DatabendSession session) { this.session.set(session); } - public OkHttpClient getHttpClient() { + private OkHttpClient getHttpClient() { return httpClient; } diff --git a/databend-jdbc/src/main/java/com/databend/jdbc/examples/DatabendConnectionFactory.java b/databend-jdbc/src/main/java/com/databend/jdbc/examples/DatabendConnectionFactory.java index 5e9653c1..1db8acd9 100644 --- a/databend-jdbc/src/main/java/com/databend/jdbc/examples/DatabendConnectionFactory.java +++ b/databend-jdbc/src/main/java/com/databend/jdbc/examples/DatabendConnectionFactory.java @@ -1,6 +1,5 @@ package com.databend.jdbc.examples; -import com.databend.jdbc.DatabendConnectionImpl; import org.apache.commons.pool2.PooledObject; import org.apache.commons.pool2.PooledObjectFactory; import org.apache.commons.pool2.impl.DefaultPooledObject; @@ -10,10 +9,10 @@ import java.sql.SQLException; import java.util.Properties; -public class DatabendConnectionFactory implements PooledObjectFactory { +public class DatabendConnectionFactory implements PooledObjectFactory { - private String url; - private Properties properties; + final private String url; + final private Properties properties; public DatabendConnectionFactory(String url, Properties properties) { this.url = url; @@ -25,18 +24,18 @@ private Connection createConnection(String url, Properties p) throws SQLExceptio } @Override - public PooledObject makeObject() throws Exception { - DatabendConnectionImpl connection = (DatabendConnectionImpl) createConnection(url, properties); + public PooledObject makeObject() throws Exception { + Connection connection = createConnection(url, properties); return new DefaultPooledObject<>(connection); } @Override - public void destroyObject(PooledObject p) throws Exception { + public void destroyObject(PooledObject p) throws Exception { p.getObject().close(); } @Override - public boolean validateObject(PooledObject p) { + public boolean validateObject(PooledObject p) { try { return !p.getObject().isClosed(); } catch (SQLException e) { @@ -45,11 +44,11 @@ public boolean validateObject(PooledObject p) { } @Override - public void activateObject(PooledObject p) throws Exception { + public void activateObject(PooledObject p) throws Exception { } @Override - public void passivateObject(PooledObject p) throws Exception { + public void passivateObject(PooledObject p) throws Exception { } } diff --git a/databend-jdbc/src/main/java/com/databend/jdbc/examples/DatabendConnectionPool.java b/databend-jdbc/src/main/java/com/databend/jdbc/examples/DatabendConnectionPool.java index ec40a6f3..8944ce2d 100644 --- a/databend-jdbc/src/main/java/com/databend/jdbc/examples/DatabendConnectionPool.java +++ b/databend-jdbc/src/main/java/com/databend/jdbc/examples/DatabendConnectionPool.java @@ -1,18 +1,18 @@ package com.databend.jdbc.examples; -import com.databend.jdbc.DatabendConnectionImpl; import org.apache.commons.pool2.impl.GenericObjectPool; import org.apache.commons.pool2.impl.GenericObjectPoolConfig; +import java.sql.Connection; import java.util.Properties; -public class DatabendConnectionPool extends GenericObjectPool { - public DatabendConnectionPool(DatabendConnectionFactory factory, GenericObjectPoolConfig config) { +public class DatabendConnectionPool extends GenericObjectPool { + public DatabendConnectionPool(DatabendConnectionFactory factory, GenericObjectPoolConfig config) { super(factory, config); } public void testDemo() throws Exception { - GenericObjectPoolConfig config = new GenericObjectPoolConfig<>(); + GenericObjectPoolConfig config = new GenericObjectPoolConfig<>(); // set max total connection config.setMaxTotal(10); // set min idle connection @@ -26,7 +26,7 @@ public void testDemo() throws Exception { DatabendConnectionPool pool = new DatabendConnectionPool(factory, config); // Get a connection from the pool - DatabendConnectionImpl connection = pool.borrowObject(); + Connection connection = pool.borrowObject(); // connection.uploadStream(); // Use the connection