Merge pull request #1285 from zhicwu/main

TCL support
ClickHouse · Mar 16, 2023 · c25103a · c25103a
2 parents d7fd089 + d39550c
commit c25103a
Show file tree

Hide file tree

Showing 11 changed files with 275 additions and 29 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,6 +1,7 @@
 ## 0.4.2
 
 ### New Features
+* support `BEGIN TRANSACTION`, `COMMIT`, and `ROLLBACK` in JDBC driver. [#975](https://github.com/ClickHouse/clickhouse-java/issues/975)
 * new options for JDBC driver
     * databaseTerm(catalog or schema, defaults to schema) [#1273](https://github.com/ClickHouse/clickhouse-java/issues/1273)
     * externalDatabase(true or false, defaults to true) [#1245](https://github.com/ClickHouse/clickhouse-java/issues/1245)

diff --git a/clickhouse-client/src/main/java/com/clickhouse/client/ClickHouseTransaction.java b/clickhouse-client/src/main/java/com/clickhouse/client/ClickHouseTransaction.java
@@ -119,6 +119,10 @@ public String toString() {
 
     static final String QUERY_SELECT_TX_ID = "SELECT transactionID()";
 
+    public static final String COMMAND_BEGIN = "BEGIN";
+    public static final String COMMAND_COMMIT = "COMMIT";
+    public static final String COMMAND_ROLLBACK = "ROLLBACK";
+
     // transaction state
     public static final int NEW = 0;
     public static final int ACTIVE = 1;
@@ -526,7 +530,7 @@ public void commit(Map<String, Serializable> settings) throws ClickHouseExceptio
                 boolean success = false;
                 try {
                     ensureTransactionId();
-                    issue("COMMIT", true, settings);
+                    issue(COMMAND_COMMIT, true, settings);
                     success = state.compareAndSet(currentState, COMMITTED);
                     return x;
                 } catch (ClickHouseException e) {
@@ -582,7 +586,7 @@ public void rollback(Map<String, Serializable> settings) throws ClickHouseExcept
                 boolean success = false;
                 try {
                     ensureTransactionId();
-                    issue("ROLLBACK", true, settings);
+                    issue(COMMAND_ROLLBACK, true, settings);
                     success = state.compareAndSet(currentState, ROLLED_BACK);
                     return x;
                 } catch (ClickHouseException e) {

diff --git a/clickhouse-jdbc/src/main/java/com/clickhouse/jdbc/ClickHouseConnection.java b/clickhouse-jdbc/src/main/java/com/clickhouse/jdbc/ClickHouseConnection.java
@@ -151,6 +151,14 @@ default PreparedStatement prepareStatement(String sql, int resultSetType, int re
         return prepareStatement(sql, resultSetType, resultSetConcurrency, ResultSet.HOLD_CURSORS_OVER_COMMIT);
     }
 
+    /**
+     * Starts a new transaction. It's no-op for a newly started transaction.
+     *
+     * @throws SQLException when current transaction is active state or not able to
+     *                      start new transaction
+     */
+    void begin() throws SQLException;
+
     /**
      * Gets configuration tied to this connection.
      *

diff --git a/clickhouse-jdbc/src/main/java/com/clickhouse/jdbc/internal/ClickHouseConnectionImpl.java b/clickhouse-jdbc/src/main/java/com/clickhouse/jdbc/internal/ClickHouseConnectionImpl.java
@@ -436,7 +436,7 @@ public void setAutoCommit(boolean autoCommit) throws SQLException {
             }
         } else { // start new transaction
             if (!txRef.compareAndSet(null, createTransaction())) {
-                log.warn("[JDBC Compliant Mode] not able to start a new transaction, reuse the exist one");
+                log.warn("Not able to start a new transaction, reuse the exist one: %s", txRef.get());
             }
         }
     }
@@ -449,9 +449,22 @@ public boolean getAutoCommit() throws SQLException {
     }
 
     @Override
-    public void commit() throws SQLException {
-        ensureOpen();
+    public void begin() throws SQLException {
+        if (getAutoCommit()) {
+            throw SqlExceptionUtils.clientError("Cannot start new transaction in auto-commit mode");
+        }
 
+        ensureTransactionSupport();
+
+        JdbcTransaction tx = txRef.get();
+        if (tx == null || !tx.isNew()) {
+            // invalid transaction state
+            throw new SQLException(JdbcTransaction.ERROR_TX_STARTED, SqlExceptionUtils.SQL_STATE_INVALID_TX_STATE);
+        }
+    }
+
+    @Override
+    public void commit() throws SQLException {
         if (getAutoCommit()) {
             throw SqlExceptionUtils.clientError("Cannot commit in auto-commit mode");
         }
@@ -475,8 +488,6 @@ public void commit() throws SQLException {
 
     @Override
     public void rollback() throws SQLException {
-        ensureOpen();
-
         if (getAutoCommit()) {
             throw SqlExceptionUtils.clientError("Cannot rollback in auto-commit mode");
         }

diff --git a/clickhouse-jdbc/src/main/java/com/clickhouse/jdbc/internal/ClickHouseStatementImpl.java b/clickhouse-jdbc/src/main/java/com/clickhouse/jdbc/internal/ClickHouseStatementImpl.java
@@ -5,6 +5,7 @@
 import java.io.Serializable;
 import java.sql.ResultSet;
 import java.sql.SQLException;
+import java.sql.SQLFeatureNotSupportedException;
 import java.sql.SQLWarning;
 import java.sql.Statement;
 import java.util.ArrayList;
@@ -22,6 +23,7 @@
 import com.clickhouse.client.ClickHouseRequest;
 import com.clickhouse.client.ClickHouseResponse;
 import com.clickhouse.client.ClickHouseResponseSummary;
+import com.clickhouse.client.ClickHouseTransaction;
 import com.clickhouse.client.ClickHouseRequest.Mutation;
 import com.clickhouse.client.config.ClickHouseClientOption;
 import com.clickhouse.config.ClickHouseConfigChangeListener;
@@ -93,6 +95,12 @@ private ClickHouseResponse getLastResponse(Map<ClickHouseOption, Serializable> o
         ClickHouseResponse response = null;
         for (int i = 0, len = parsedStmts.length; i < len; i++) {
             ClickHouseSqlStatement stmt = parsedStmts[i];
+            response = processSqlStatement(stmt);
+            if (response != null) {
+                updateResult(stmt, response);
+                continue;
+            }
+
             if (stmt.hasFormat()) {
                 request.format(ClickHouseFormat.valueOf(stmt.getFormat()));
             }
@@ -124,6 +132,23 @@ protected void ensureOpen() throws SQLException {
         }
     }
 
+    protected ClickHouseResponse processSqlStatement(ClickHouseSqlStatement stmt) throws SQLException {
+        if (stmt.isTCL()) {
+            if (stmt.containsKeyword(ClickHouseTransaction.COMMAND_BEGIN)) {
+                connection.begin();
+            } else if (stmt.containsKeyword(ClickHouseTransaction.COMMAND_COMMIT)) {
+                connection.commit();
+            } else if (stmt.containsKeyword(ClickHouseTransaction.COMMAND_ROLLBACK)) {
+                connection.rollback();
+            } else {
+                throw new SQLFeatureNotSupportedException("Unsupported TCL: " + stmt.getSQL());
+            }
+            return ClickHouseResponse.EMPTY;
+        }
+
+        return null;
+    }
+
     protected ClickHouseResponse executeStatement(String stmt, Map<ClickHouseOption, Serializable> options,
             List<ClickHouseExternalTable> tables, Map<String, String> settings) throws SQLException {
         boolean autoTx = connection.getAutoCommit() && connection.isTransactionSupported();
@@ -175,6 +200,10 @@ protected ClickHouseResponse executeStatement(String stmt, Map<ClickHouseOption,
     protected ClickHouseResponse executeStatement(ClickHouseSqlStatement stmt,
             Map<ClickHouseOption, Serializable> options, List<ClickHouseExternalTable> tables,
             Map<String, String> settings) throws SQLException {
+        ClickHouseResponse resp = processSqlStatement(stmt);
+        if (resp != null) {
+            return resp;
+        }
         return executeStatement(stmt.getSQL(), options, tables, settings);
     }
 
@@ -261,7 +290,7 @@ protected ResultSet updateResult(ClickHouseSqlStatement stmt, ClickHouseResponse
                     stmt.getTable(), this, response);
         } else {
             response.close();
-            currentUpdateCount = stmt.isDDL() ? 0L
+            currentUpdateCount = stmt.isDDL() || stmt.isTCL() ? 0L
                     : (response.getSummary().isEmpty() ? 1L : response.getSummary().getWrittenRows());
             currentResult = null;
         }

diff --git a/clickhouse-jdbc/src/main/java/com/clickhouse/jdbc/internal/JdbcTransaction.java b/clickhouse-jdbc/src/main/java/com/clickhouse/jdbc/internal/JdbcTransaction.java
@@ -19,6 +19,7 @@ public class JdbcTransaction {
     static final String ACTION_ROLLBACK = "rolled back";
 
     static final String ERROR_TX_NOT_STARTED = "Transaction not started";
+    static final String ERROR_TX_STARTED = "Transaction has been started";
 
     protected final ClickHouseTransaction tx;
     protected final String id;
@@ -36,6 +37,11 @@ public JdbcTransaction(ClickHouseTransaction tx) {
         this.savepoints = new LinkedList<>();
     }
 
+    public boolean isNew() {
+        return this.queries.isEmpty() && this.savepoints.isEmpty()
+                && (this.tx == null || this.tx.isNew() || this.tx.isActive());
+    }
+
     public void commit(Logger log) throws SQLException {
         if (this.tx != null) {
             try {

diff --git a/clickhouse-jdbc/src/main/java/com/clickhouse/jdbc/parser/ClickHouseSqlStatement.java b/clickhouse-jdbc/src/main/java/com/clickhouse/jdbc/parser/ClickHouseSqlStatement.java
@@ -130,6 +130,10 @@ public boolean isMutation() {
         return this.stmtType.getOperationType() == OperationType.WRITE || this.hasOutfile();
     }
 
+    public boolean isTCL() {
+        return this.stmtType.getLanguageType() == LanguageType.TCL;
+    }
+
     public boolean isIdemponent() {
         boolean result = this.stmtType.isIdempotent() && !this.hasOutfile();
 

diff --git a/clickhouse-jdbc/src/main/java/com/clickhouse/jdbc/parser/StatementType.java b/clickhouse-jdbc/src/main/java/com/clickhouse/jdbc/parser/StatementType.java
@@ -27,7 +27,8 @@ public enum StatementType {
     TRUNCATE(LanguageType.DDL, OperationType.UNKNOWN, true), // truncate statement
     UPDATE(LanguageType.DML, OperationType.WRITE, false), // the upcoming light-weight update statement
     USE(LanguageType.DDL, OperationType.UNKNOWN, true), // use statement
-    WATCH(LanguageType.DDL, OperationType.UNKNOWN, true); // watch statement
+    WATCH(LanguageType.DDL, OperationType.UNKNOWN, true), // watch statement
+    TRANSACTION(LanguageType.TCL, OperationType.WRITE, true); // TCL statement
 
     private LanguageType langType;
     private OperationType opType;

diff --git a/clickhouse-jdbc/src/main/javacc/ClickHouseSqlParser.jj b/clickhouse-jdbc/src/main/javacc/ClickHouseSqlParser.jj
@@ -382,6 +382,7 @@ void stmt(): {} {
     | updateStmt()   { token_source.stmtType = StatementType.UPDATE;   }
     | useStmt()      { token_source.stmtType = StatementType.USE;      }
     | watchStmt()    { token_source.stmtType = StatementType.WATCH;    }
+    | txStmt()       { token_source.stmtType = StatementType.TRANSACTION; }
 }
 
 // https://clickhouse.tech/docs/en/sql-reference/statements/alter/
@@ -641,6 +642,13 @@ void watchStmt(): {} { // not interested
     <WATCH> anyExprList()
 }
 
+// TCL
+void txStmt(): {} { // not interested
+    <COMMIT> { token_source.addPosition(token); }
+    | <ROLLBACK> { token_source.addPosition(token); }
+    | (<BEGIN> { token_source.addPosition(token); } <TRANSACTION>)
+}
+
 // columns
 void columnExprList(): {} {
     columnsExpr() (<COMMA> columnsExpr())*
@@ -995,6 +1003,9 @@ TOKEN: {
     | <UPDATE   : <U> <P> <D> <A> <T> <E>        >
     | <USE      : <U> <S> <E>                    >
     | <WATCH    : <W> <A> <T> <C> <H>            >
+    | <BEGIN    : <B> <E> <G> <I> <N>            >
+    | <COMMIT   : <C> <O> <M> <M> <I> <T>        >
+    | <ROLLBACK : <R> <O> <L> <L> <B> <A> <C> <K>>
 
     | <ALL      : <A> <L> <L>                        >
     | <AND      : <A> <N> <D>                        >
@@ -1060,6 +1071,7 @@ TOKEN: {
     | <TIMESTAMP: <T> <I> <M> <E> <S> <T> <A> <M> <P>>
     | <TOP      : <T> <O> <P>                        >
     | <TOTALS   : <T> <O> <T> <A> <L> <S>            >
+    | <TRANSACTION : <T> <R> <A> <N> <S> <A> <C> <T> <I> <O> <N>>
     | <UNION    : <U> <N> <I> <O> <N>                >
     | <USER     : <U> <S> <E> <R>                    >
     | <USING    : <U> <S> <I> <N> <G>                >