Skip to content

Commit 91c0059

Browse files
zhicwuzhicwu
committed
Enhance parser to recognize idempotent operations
1 parent 4c51d65 commit 91c0059

File tree

9 files changed

+348
-151
lines changed

9 files changed

+348
-151
lines changed

‎src/main/java/ru/yandex/clickhouse/ClickHousePreparedStatementImpl.java‎

Lines changed: 6 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,6 @@
3636
import org.apache.http.entity.AbstractHttpEntity;
3737
import org.apache.http.impl.client.CloseableHttpClient;
3838

39-
import ru.yandex.clickhouse.jdbc.parser.ClickHouseSqlParser;
4039
import ru.yandex.clickhouse.jdbc.parser.ClickHouseSqlStatement;
4140
import ru.yandex.clickhouse.jdbc.parser.StatementType;
4241
import ru.yandex.clickhouse.response.ClickHouseResponse;
@@ -55,7 +54,6 @@ public class ClickHousePreparedStatementImpl extends ClickHouseStatementImpl imp
5554

5655
private final TimeZone dateTimeZone;
5756
private final TimeZone dateTimeTimeZone;
58-
private final ClickHouseSqlStatement parsedSql;
5957
private final String sql;
6058
private final List<String> sqlParts;
6159
private final ClickHousePreparedStatementParameter[] binds;
@@ -68,9 +66,11 @@ public ClickHousePreparedStatementImpl(CloseableHttpClient client,
6866
TimeZone serverTimeZone, int resultSetType) throws SQLException
6967
{
7068
super(client, connection, properties, resultSetType);
71-
this.parsedSql = ClickHouseSqlParser.parseSingleStatement(sql, properties);
69+
parseSingleStatement(sql);
70+
7271
this.sql = sql;
73-
PreparedStatementParser parser = PreparedStatementParser.parse(sql);
72+
PreparedStatementParser parser = PreparedStatementParser.parse(sql,
73+
parsedStmt.getEndPosition(ClickHouseSqlStatement.KEYWORD_VALUES));
7474
this.parameterList = parser.getParameters();
7575
this.insertBatchMode = parser.isValuesMode();
7676
this.sqlParts = parser.getParts();
@@ -352,8 +352,8 @@ public int[] executeBatch() throws SQLException {
352352
@Override
353353
public int[] executeBatch(Map<ClickHouseQueryParam, String> additionalDBParams) throws SQLException {
354354
int valuePosition = -1;
355-
if (parsedSql.getStatementType() == StatementType.INSERT && parsedSql.hasValues()) {
356-
valuePosition = parsedSql.getStartPosition(ClickHouseSqlStatement.KEYWORD_VALUES);
355+
if (parsedStmt.getStatementType() == StatementType.INSERT && parsedStmt.hasValues()) {
356+
valuePosition = parsedStmt.getStartPosition(ClickHouseSqlStatement.KEYWORD_VALUES);
357357
} else {
358358
Matcher matcher = VALUES.matcher(sql);
359359
if (matcher.find()) {
@@ -442,7 +442,6 @@ public ResultSetMetaData getMetaData() throws SQLException {
442442
return currentResult.getMetaData();
443443
}
444444

445-
ClickHouseSqlStatement parsedStmt = ClickHouseSqlParser.parseSingleStatement(sql, properties);
446445
if (!parsedStmt.isQuery() || (!parsedStmt.isRecognized() && !isSelect(sql))) {
447446
return null;
448447
}

‎src/main/java/ru/yandex/clickhouse/ClickHouseStatementImpl.java‎

Lines changed: 98 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -89,6 +89,8 @@ public class ClickHouseStatementImpl extends ConfigurableApi<ClickHouseStatement
8989

9090
private volatile String queryId;
9191

92+
protected ClickHouseSqlStatement parsedStmt;
93+
9294
/**
9395
* Current database name may be changed by {@link java.sql.Connection#setCatalog(String)}
9496
* between creation of this object and query execution, but javadoc does not allow
@@ -101,6 +103,43 @@ public class ClickHouseStatementImpl extends ConfigurableApi<ClickHouseStatement
101103
@Deprecated
102104
private static final String databaseKeyword = "CREATE DATABASE";
103105

106+
@Deprecated
107+
protected void parseSingleStatement(String sql) throws SQLException {
108+
this.parsedStmt = null;
109+
ClickHouseSqlStatement[] stmts = ClickHouseSqlParser.parse(sql, properties);
110+
111+
if (stmts.length == 1) {
112+
this.parsedStmt = stmts[0];
113+
} else {
114+
this.parsedStmt = new ClickHouseSqlStatement(sql, StatementType.UNKNOWN);
115+
// throw new SQLException("Multiple statements are not supported.");
116+
}
117+
118+
if (this.parsedStmt.isIdemponent()) {
119+
httpContext.setAttribute("is_idempotent", Boolean.TRUE);
120+
} else {
121+
httpContext.removeAttribute("is_idempotent");
122+
}
123+
}
124+
125+
@Deprecated
126+
private void parseSingleStatement(String sql, ClickHouseFormat preferredFormat) throws SQLException {
127+
parseSingleStatement(sql);
128+
129+
if (parsedStmt.isQuery() && !parsedStmt.hasFormat()) {
130+
String format = preferredFormat.name();
131+
Map<String, Integer> positions = new HashMap<>();
132+
positions.putAll(parsedStmt.getPositions());
133+
positions.put(ClickHouseSqlStatement.KEYWORD_FORMAT, sql.length());
134+
135+
sql = new StringBuilder(parsedStmt.getSQL()).append("\nFORMAT ").append(format).append(';')
136+
.toString();
137+
parsedStmt = new ClickHouseSqlStatement(sql, parsedStmt.getStatementType(),
138+
parsedStmt.getCluster(), parsedStmt.getDatabase(), parsedStmt.getTable(),
139+
format, parsedStmt.getOutfile(), parsedStmt.getParameters(), positions);
140+
}
141+
}
142+
104143
public ClickHouseStatementImpl(CloseableHttpClient client, ClickHouseConnection connection,
105144
ClickHouseProperties properties, int resultSetType) {
106145
super(null);
@@ -141,21 +180,23 @@ public ResultSet executeQuery(String sql,
141180
}
142181
additionalDBParams.put(ClickHouseQueryParam.EXTREMES, "0");
143182

144-
InputStream is = getInputStream(sql, additionalDBParams, externalData, additionalRequestParams);
183+
parseSingleStatement(sql, ClickHouseFormat.TabSeparatedWithNamesAndTypes);
184+
if (!parsedStmt.isRecognized() && isSelect(sql)) {
185+
Map<String, Integer> positions = new HashMap<>();
186+
String dbName = extractDBName(sql);
187+
String tableName = extractTableName(sql);
188+
if (extractWithTotals(sql)) {
189+
positions.put(ClickHouseSqlStatement.KEYWORD_TOTALS, 1);
190+
}
191+
parsedStmt = new ClickHouseSqlStatement(sql, StatementType.SELECT,
192+
null, dbName, tableName, null, null, null, positions);
193+
// httpContext.setAttribute("is_idempotent", Boolean.TRUE);
194+
}
145195

146-
ClickHouseSqlStatement parsedStmt = ClickHouseSqlParser.parseSingleStatement(sql, properties);
196+
InputStream is = getInputStream(sql, additionalDBParams, externalData, additionalRequestParams);
197+
147198
try {
148-
if (parsedStmt.isQuery() || (!parsedStmt.isRecognized() && isSelect(sql))) {
149-
if (!parsedStmt.isRecognized()) {
150-
Map<String, Integer> positions = new HashMap<>();
151-
String dbName = extractDBName(sql);
152-
String tableName = extractTableName(sql);
153-
if (extractWithTotals(sql)) {
154-
positions.put(ClickHouseSqlStatement.KEYWORD_TOTALS, 1);
155-
}
156-
parsedStmt = new ClickHouseSqlStatement(sql, StatementType.SELECT,
157-
null, dbName, tableName, null, null, positions);
158-
}
199+
if (parsedStmt.isQuery()) {
159200
currentUpdateCount = -1;
160201
currentResult = createResultSet(properties.isCompress()
161202
? new ClickHouseLZ4Stream(is) : is, properties.getBufferSize(),
@@ -193,8 +234,15 @@ public ClickHouseResponse executeQueryClickhouseResponse(String sql, Map<ClickHo
193234
public ClickHouseResponse executeQueryClickhouseResponse(String sql,
194235
Map<ClickHouseQueryParam, String> additionalDBParams,
195236
Map<String, String> additionalRequestParams) throws SQLException {
237+
parseSingleStatement(sql, ClickHouseFormat.JSONCompact);
238+
if (parsedStmt.isRecognized()) {
239+
sql = parsedStmt.getSQL();
240+
} else {
241+
sql = addFormatIfAbsent(sql, ClickHouseFormat.JSONCompact);
242+
}
243+
196244
InputStream is = getInputStream(
197-
addFormatIfAbsent(sql, properties, ClickHouseFormat.JSONCompact),
245+
sql,
198246
additionalDBParams,
199247
null,
200248
additionalRequestParams
@@ -223,15 +271,27 @@ public ClickHouseRowBinaryInputStream executeQueryClickhouseRowBinaryStream(Stri
223271

224272
@Override
225273
public ClickHouseRowBinaryInputStream executeQueryClickhouseRowBinaryStream(String sql, Map<ClickHouseQueryParam, String> additionalDBParams, Map<String, String> additionalRequestParams) throws SQLException {
274+
parseSingleStatement(sql, ClickHouseFormat.RowBinary);
275+
if (parsedStmt.isRecognized()) {
276+
sql = parsedStmt.getSQL();
277+
} else {
278+
sql = addFormatIfAbsent(sql, ClickHouseFormat.RowBinary);
279+
if (isSelect(sql)) {
280+
parsedStmt = new ClickHouseSqlStatement(sql, StatementType.SELECT);
281+
// httpContext.setAttribute("is_idempotent", Boolean.TRUE);
282+
} else {
283+
parsedStmt = new ClickHouseSqlStatement(sql, StatementType.UNKNOWN);
284+
}
285+
}
286+
226287
InputStream is = getInputStream(
227-
addFormatIfAbsent(sql, properties, ClickHouseFormat.RowBinary),
288+
sql,
228289
additionalDBParams,
229290
null,
230291
additionalRequestParams
231292
);
232-
ClickHouseSqlStatement parsedStmt = ClickHouseSqlParser.parseSingleStatement(sql, properties);
233293
try {
234-
if (parsedStmt.isQuery() || (!parsedStmt.isRecognized() && isSelect(sql))) {
294+
if (parsedStmt.isQuery()) {
235295
currentUpdateCount = -1;
236296
currentRowBinaryResult = new ClickHouseRowBinaryInputStream(properties.isCompress()
237297
? new ClickHouseLZ4Stream(is) : is, getConnection().getTimeZone(), properties);
@@ -249,6 +309,8 @@ public ClickHouseRowBinaryInputStream executeQueryClickhouseRowBinaryStream(Stri
249309

250310
@Override
251311
public int executeUpdate(String sql) throws SQLException {
312+
parseSingleStatement(sql, ClickHouseFormat.TabSeparatedWithNamesAndTypes);
313+
252314
InputStream is = null;
253315
try {
254316
is = getInputStream(sql, null, null, null);
@@ -490,24 +552,20 @@ public ClickHouseResponseSummary getResponseSummary() {
490552

491553
@Deprecated
492554
static String clickhousifySql(String sql) {
493-
return clickhousifySql(sql, null);
494-
}
495-
496-
static String clickhousifySql(String sql, ClickHouseProperties properties) {
497-
return addFormatIfAbsent(sql, properties, ClickHouseFormat.TabSeparatedWithNamesAndTypes);
555+
return addFormatIfAbsent(sql, ClickHouseFormat.TabSeparatedWithNamesAndTypes);
498556
}
499557

500558
/**
501559
* Adding FORMAT TabSeparatedWithNamesAndTypes if not added
502560
* adds format only to select queries
503561
*/
504-
private static String addFormatIfAbsent(final String sql, ClickHouseProperties properties, ClickHouseFormat format) {
562+
@Deprecated
563+
private static String addFormatIfAbsent(final String sql, ClickHouseFormat format) {
505564
String cleanSQL = sql.trim();
506-
ClickHouseSqlStatement parsedStmt = ClickHouseSqlParser.parseSingleStatement(cleanSQL, properties);
507-
if (!parsedStmt.isQuery() || (!parsedStmt.isRecognized() && !isSelect(cleanSQL))) {
565+
if (!isSelect(cleanSQL)) {
508566
return cleanSQL;
509567
}
510-
if (parsedStmt.hasFormat() || (!parsedStmt.isRecognized() && ClickHouseFormat.containsFormat(cleanSQL))) {
568+
if (ClickHouseFormat.containsFormat(cleanSQL)) {
511569
return cleanSQL;
512570
}
513571
StringBuilder sb = new StringBuilder();
@@ -542,6 +600,7 @@ static boolean isSelect(String sql) {
542600
return false;
543601
}
544602

603+
@Deprecated
545604
private String extractTableName(String sql) {
546605
String s = extractDBAndTableName(sql);
547606
if (s.contains(".")) {
@@ -551,6 +610,7 @@ private String extractTableName(String sql) {
551610
}
552611
}
553612

613+
@Deprecated
554614
private String extractDBName(String sql) {
555615
String s = extractDBAndTableName(sql);
556616
if (s.contains(".")) {
@@ -560,6 +620,7 @@ private String extractDBName(String sql) {
560620
}
561621
}
562622

623+
@Deprecated
563624
private String extractDBAndTableName(String sql) {
564625
if (Utils.startsWithIgnoreCase(sql, "select")) {
565626
String withoutStrings = Utils.retainUnquoted(sql, '\'');
@@ -582,6 +643,7 @@ private String extractDBAndTableName(String sql) {
582643
return "system.unknown";
583644
}
584645

646+
@Deprecated
585647
private boolean extractWithTotals(String sql) {
586648
if (Utils.startsWithIgnoreCase(sql, "select")) {
587649
String withoutStrings = Utils.retainUnquoted(sql, '\'');
@@ -596,15 +658,23 @@ private InputStream getInputStream(
596658
List<ClickHouseExternalData> externalData,
597659
Map<String, String> additionalRequestParams
598660
) throws ClickHouseException {
599-
sql = clickhousifySql(sql, properties);
661+
boolean ignoreDatabase = false;
662+
if (parsedStmt.isRecognized()) {
663+
sql = parsedStmt.getSQL();
664+
// TODO consider more scenarios like drop, show etc.
665+
ignoreDatabase = parsedStmt.getStatementType() == StatementType.CREATE
666+
&& parsedStmt.containsKeyword(ClickHouseSqlStatement.KEYWORD_DATABASE);
667+
} else {
668+
sql = clickhousifySql(sql);
669+
ignoreDatabase = sql.trim().regionMatches(true, 0, databaseKeyword, 0, databaseKeyword.length());
670+
}
600671
log.debug("Executing SQL: {}", sql);
601672

602673
additionalClickHouseDBParams = addQueryIdTo(
603674
additionalClickHouseDBParams == null
604675
? new EnumMap<ClickHouseQueryParam, String>(ClickHouseQueryParam.class)
605676
: additionalClickHouseDBParams);
606677

607-
boolean ignoreDatabase = sql.trim().regionMatches(true, 0, databaseKeyword, 0, databaseKeyword.length());
608678
URI uri;
609679
if (externalData == null || externalData.isEmpty()) {
610680
uri = buildRequestUri(

‎src/main/java/ru/yandex/clickhouse/PreparedStatementParser.java‎

Lines changed: 8 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -6,10 +6,6 @@
66
import java.util.regex.Matcher;
77
import java.util.regex.Pattern;
88

9-
import ru.yandex.clickhouse.jdbc.parser.ClickHouseSqlParser;
10-
import ru.yandex.clickhouse.jdbc.parser.ClickHouseSqlStatement;
11-
import ru.yandex.clickhouse.jdbc.parser.StatementType;
12-
import ru.yandex.clickhouse.settings.ClickHouseProperties;
139
import ru.yandex.clickhouse.util.apache.StringUtils;
1410

1511
/**
@@ -36,15 +32,16 @@ private PreparedStatementParser() {
3632

3733
@Deprecated
3834
static PreparedStatementParser parse(String sql) {
39-
return parse(sql, null);
35+
return parse(sql, -1);
4036
}
4137

42-
static PreparedStatementParser parse(String sql, ClickHouseProperties properties) {
38+
@Deprecated
39+
static PreparedStatementParser parse(String sql, int valuesEndPosition) {
4340
if (StringUtils.isBlank(sql)) {
4441
throw new IllegalArgumentException("SQL may not be blank");
4542
}
4643
PreparedStatementParser parser = new PreparedStatementParser();
47-
parser.parseSQL(sql, properties);
44+
parser.parseSQL(sql, valuesEndPosition);
4845
return parser;
4946
}
5047

@@ -66,7 +63,7 @@ private void reset() {
6663
valuesMode = false;
6764
}
6865

69-
private void parseSQL(String sql, ClickHouseProperties properties) {
66+
private void parseSQL(String sql, int valuesEndPosition) {
7067
reset();
7168
List<String> currentParamList = new ArrayList<String>();
7269
boolean afterBackSlash = false;
@@ -75,15 +72,10 @@ private void parseSQL(String sql, ClickHouseProperties properties) {
7572
boolean inSingleLineComment = false;
7673
boolean inMultiLineComment = false;
7774
boolean whiteSpace = false;
78-
ClickHouseSqlStatement parsedSql = ClickHouseSqlParser.parseSingleStatement(sql, properties);
7975
int endPosition = 0;
80-
if (parsedSql.getStatementType() == StatementType.INSERT) {
81-
endPosition = parsedSql.getEndPosition(ClickHouseSqlStatement.KEYWORD_VALUES) - 1;
82-
if (endPosition > 0) {
83-
valuesMode = true;
84-
} else {
85-
endPosition = 0;
86-
}
76+
if (valuesEndPosition > 0) {
77+
valuesMode = true;
78+
endPosition = valuesEndPosition;
8779
} else {
8880
Matcher matcher = VALUES.matcher(sql);
8981
if (matcher.find()) {

0 commit comments

Comments
 (0)