[PostgreSQL] Support more spots for parameters with the ANTLR parser

This commit is contained in:
aboyko
2024-07-30 16:48:31 -04:00
parent df9fbcdc0c
commit 2b7ac1a3c2
6 changed files with 10308 additions and 8800 deletions

View File

@@ -593,8 +593,8 @@ create_database_stmt
(LC_CTYPE EQUAL name_)?
(TABLESPACE EQUAL name_)?
(ALLOW_CONNECTIONS EQUAL name_)?
(CONNECTION LIMIT EQUAL INTEGER_LITERAL)?
(IS_TEMPLATE EQUAL INTEGER_LITERAL)?
(CONNECTION LIMIT EQUAL parameterOrIntegerLiteral)?
(IS_TEMPLATE EQUAL parameterOrIntegerLiteral)?
)
;
@@ -728,7 +728,7 @@ create_role_stmt
(SUPERUSER | NOSUPERUSER | CREATEDB | NOCREATEDB |
CREATEROLE | NOCREATEROLE | INHERIT | NOINHERIT | LOGIN | NOLOGIN |
REPLICATION | NOREPLICATION | BYPASSRLS | NOBYPASSRLS |
CONNECTION LIMIT INTEGER_LITERAL | ENCRYPTED? PASSWORD (SINGLEQ_STRING_LITERAL | NULL) |
CONNECTION LIMIT parameterOrIntegerLiteral | ENCRYPTED? PASSWORD (SINGLEQ_STRING_LITERAL | NULL) |
VALID UNTIL SINGLEQ_STRING_LITERAL | IN ROLE name_list | IN GROUP name_list | ROLE name_list |
ADMIN name_list | USER name_list | SYSID INTEGER_LITERAL)+)?
;
@@ -752,10 +752,10 @@ create_schema_stmt
;
create_sequence_stmt
: CREATE (TEMPORARY | TEMP)? SEQUENCE (IF NOT EXISTS)? name=identifier (INCREMENT BY? increment=INTEGER_LITERAL)?
(MINVALUE minvalue=INTEGER_LITERAL | NO MINVALUE)?
(MAXVALUE maxvalue=INTEGER_LITERAL | NO MAXVALUE)?
(START WITH? start=INTEGER_LITERAL)? (CACHE cache=INTEGER_LITERAL)? (NO? CYCLE)?
: CREATE (TEMPORARY | TEMP)? SEQUENCE (IF NOT EXISTS)? name=identifier (INCREMENT BY? increment=parameterOrIntegerLiteral)?
(MINVALUE minvalue=parameterOrIntegerLiteral | NO MINVALUE)?
(MAXVALUE maxvalue=parameterOrIntegerLiteral | NO MAXVALUE)?
(START WITH? start=parameterOrIntegerLiteral)? (CACHE cache=parameterOrIntegerLiteral)? (NO? CYCLE)?
(OWNED BY ((table_name=identifier DOT column_name_=identifier) | NONE))?
;
@@ -1106,8 +1106,8 @@ lock_stmt
;
move_stmt
: MOVE ((NEXT | PRIOR | FIRST | LAST | ABSOLUTE INTEGER | RELATIVE INTEGER_LITERAL | INTEGER_LITERAL |
ALL | FORWARD (INTEGER_LITERAL|ALL)? | BACKWARD (INTEGER_LITERAL|ALL)?) (FROM|IN)?)? cursor_name=name_
: MOVE ((NEXT | PRIOR | FIRST | LAST | ABSOLUTE INTEGER | RELATIVE parameterOrIntegerLiteral | parameterOrIntegerLiteral |
ALL | FORWARD (parameterOrIntegerLiteral|ALL)? | BACKWARD (parameterOrIntegerLiteral|ALL)?) (FROM|IN)?)? cursor_name=name_
;
notify_stmt
@@ -1366,9 +1366,9 @@ explain_parameter
frame
: UNBOUNDED PRECEDING
| INTEGER_LITERAL PRECEDING
| parameterOrIntegerLiteral PRECEDING
| CURRENT ROW
| INTEGER_LITERAL FOLLOWING
| parameterOrIntegerLiteral FOLLOWING
| UNBOUNDED FOLLOWING
;
@@ -1408,15 +1408,15 @@ order_by_item
;
limit_clause
: LIMIT (INTEGER_LITERAL | ALL | func_call)
: LIMIT (parameterOrIntegerLiteral | ALL | func_call)
;
offset_clause
: OFFSET INTEGER_LITERAL (ROW | ROWS)?
: OFFSET parameterOrIntegerLiteral (ROW | ROWS)?
;
fetch_clause
: FETCH (FIRST | NEXT) INTEGER_LITERAL? (ROW | ROWS) ONLY
: FETCH (FIRST | NEXT) parameterOrIntegerLiteral? (ROW | ROWS) ONLY
;
for_clause
@@ -1502,11 +1502,23 @@ expr
;
parameter
: prefix=':' SPEL
| prefix=':' identifier
: COLON SPEL
| COLON identifier
| COLON reserved_keyword
| COLON INTEGER_LITERAL
| prefix='?' SPEL
| prefix='?' INTEGER_LITERAL?
;
parameterOrIntegerLiteral
: parameter
| INTEGER_LITERAL
;
parameterOrNumericLiteral
: parameter
| NUMERIC_LITERAL
;
// TODO: is this necessary. can we just encapsulate within expr's operator precedence?
bool_expr
@@ -1777,11 +1789,15 @@ role_name_list
param_value
: ON | OFF | TRUE | FALSE | YES | NO | NONE
| SINGLEQ_STRING_LITERAL
| NUMERIC_LITERAL
| INTEGER_LITERAL
| parameterOrNumericLiteral
| parameterOrIntegerLiteral
| identifier
;
reserved_keyword
: CREATE | DROP | FROM | GROUP | LIMIT | ORDER | SELECT | UNION
;
// allow non-reserved keywords as identifiers
// TODO: is this necessary?
// easier to whitelist than blacklist

View File

@@ -2628,6 +2628,30 @@ public class PostgreSqlParserBaseListener implements PostgreSqlParserListener {
* <p>The default implementation does nothing.</p>
*/
@Override public void exitParameter(PostgreSqlParser.ParameterContext ctx) { }
/**
* {@inheritDoc}
*
* <p>The default implementation does nothing.</p>
*/
@Override public void enterParameterOrIntegerLiteral(PostgreSqlParser.ParameterOrIntegerLiteralContext ctx) { }
/**
* {@inheritDoc}
*
* <p>The default implementation does nothing.</p>
*/
@Override public void exitParameterOrIntegerLiteral(PostgreSqlParser.ParameterOrIntegerLiteralContext ctx) { }
/**
* {@inheritDoc}
*
* <p>The default implementation does nothing.</p>
*/
@Override public void enterParameterOrNumericLiteral(PostgreSqlParser.ParameterOrNumericLiteralContext ctx) { }
/**
* {@inheritDoc}
*
* <p>The default implementation does nothing.</p>
*/
@Override public void exitParameterOrNumericLiteral(PostgreSqlParser.ParameterOrNumericLiteralContext ctx) { }
/**
* {@inheritDoc}
*
@@ -3084,6 +3108,18 @@ public class PostgreSqlParserBaseListener implements PostgreSqlParserListener {
* <p>The default implementation does nothing.</p>
*/
@Override public void exitParam_value(PostgreSqlParser.Param_valueContext ctx) { }
/**
* {@inheritDoc}
*
* <p>The default implementation does nothing.</p>
*/
@Override public void enterReserved_keyword(PostgreSqlParser.Reserved_keywordContext ctx) { }
/**
* {@inheritDoc}
*
* <p>The default implementation does nothing.</p>
*/
@Override public void exitReserved_keyword(PostgreSqlParser.Reserved_keywordContext ctx) { }
/**
* {@inheritDoc}
*

View File

@@ -2187,6 +2187,26 @@ public interface PostgreSqlParserListener extends ParseTreeListener {
* @param ctx the parse tree
*/
void exitParameter(PostgreSqlParser.ParameterContext ctx);
/**
* Enter a parse tree produced by {@link PostgreSqlParser#parameterOrIntegerLiteral}.
* @param ctx the parse tree
*/
void enterParameterOrIntegerLiteral(PostgreSqlParser.ParameterOrIntegerLiteralContext ctx);
/**
* Exit a parse tree produced by {@link PostgreSqlParser#parameterOrIntegerLiteral}.
* @param ctx the parse tree
*/
void exitParameterOrIntegerLiteral(PostgreSqlParser.ParameterOrIntegerLiteralContext ctx);
/**
* Enter a parse tree produced by {@link PostgreSqlParser#parameterOrNumericLiteral}.
* @param ctx the parse tree
*/
void enterParameterOrNumericLiteral(PostgreSqlParser.ParameterOrNumericLiteralContext ctx);
/**
* Exit a parse tree produced by {@link PostgreSqlParser#parameterOrNumericLiteral}.
* @param ctx the parse tree
*/
void exitParameterOrNumericLiteral(PostgreSqlParser.ParameterOrNumericLiteralContext ctx);
/**
* Enter a parse tree produced by {@link PostgreSqlParser#bool_expr}.
* @param ctx the parse tree
@@ -2567,6 +2587,16 @@ public interface PostgreSqlParserListener extends ParseTreeListener {
* @param ctx the parse tree
*/
void exitParam_value(PostgreSqlParser.Param_valueContext ctx);
/**
* Enter a parse tree produced by {@link PostgreSqlParser#reserved_keyword}.
* @param ctx the parse tree
*/
void enterReserved_keyword(PostgreSqlParser.Reserved_keywordContext ctx);
/**
* Exit a parse tree produced by {@link PostgreSqlParser#reserved_keyword}.
* @param ctx the parse tree
*/
void exitReserved_keyword(PostgreSqlParser.Reserved_keywordContext ctx);
/**
* Enter a parse tree produced by {@link PostgreSqlParser#non_reserved_keyword}.
* @param ctx the parse tree

View File

@@ -156,9 +156,10 @@ public class PostgreSqlSemanticTokens implements SemanticTokensDataProvider {
public void exitParameter(ParameterContext param) {
if (param.identifier() != null) {
MySqlSemanticTokens.getAllLeafs(param.identifier()).forEach(t -> semantics.put(t, "parameter"));
}
if (param.INTEGER_LITERAL() != null) {
} else if (param.INTEGER_LITERAL() != null) {
semantics.put(param.INTEGER_LITERAL().getSymbol(), "parameter");
} else if (param.reserved_keyword() != null) {
MySqlSemanticTokens.getAllLeafs(param.reserved_keyword()).forEach(t -> semantics.put(t, "parameter"));
}
}

View File

@@ -100,5 +100,82 @@ public class PostgreSqlSemanticTokensTest {
assertThat(tokens.size()).isEqualTo(13);
}
@Test
void parameterInLimitClause_1() {
List<SemanticTokenData> tokens = provider.computeTokens("SELECT * FROM cards ORDER BY random() LIMIT :2", 0);
assertThat(tokens.size()).isEqualTo(12);
assertThat(tokens.get(0)).isEqualTo(new SemanticTokenData(0, 6, "keyword", new String[0]));
assertThat(tokens.get(1)).isEqualTo(new SemanticTokenData(7, 8, "operator", new String[0]));
assertThat(tokens.get(2)).isEqualTo(new SemanticTokenData(9, 13, "keyword", new String[0]));
assertThat(tokens.get(3)).isEqualTo(new SemanticTokenData(14, 19, "variable", new String[0]));
assertThat(tokens.get(4)).isEqualTo(new SemanticTokenData(20, 25, "keyword", new String[0]));
assertThat(tokens.get(5)).isEqualTo(new SemanticTokenData(26, 28, "keyword", new String[0]));
assertThat(tokens.get(6)).isEqualTo(new SemanticTokenData(29, 35, "method", new String[0]));
assertThat(tokens.get(7)).isEqualTo(new SemanticTokenData(35, 36, "operator", new String[0]));
assertThat(tokens.get(8)).isEqualTo(new SemanticTokenData(36, 37, "operator", new String[0]));
assertThat(tokens.get(9)).isEqualTo(new SemanticTokenData(38, 43, "keyword", new String[0]));
assertThat(tokens.get(10)).isEqualTo(new SemanticTokenData(44, 45, "operator", new String[0]));
assertThat(tokens.get(11)).isEqualTo(new SemanticTokenData(45, 46, "parameter", new String[0]));
}
@Test
void parameterInLimitClause_2() {
List<SemanticTokenData> tokens = provider.computeTokens("SELECT * FROM cards ORDER BY random() LIMIT ?2", 0);
assertThat(tokens.size()).isEqualTo(12);
assertThat(tokens.get(0)).isEqualTo(new SemanticTokenData(0, 6, "keyword", new String[0]));
assertThat(tokens.get(1)).isEqualTo(new SemanticTokenData(7, 8, "operator", new String[0]));
assertThat(tokens.get(2)).isEqualTo(new SemanticTokenData(9, 13, "keyword", new String[0]));
assertThat(tokens.get(3)).isEqualTo(new SemanticTokenData(14, 19, "variable", new String[0]));
assertThat(tokens.get(4)).isEqualTo(new SemanticTokenData(20, 25, "keyword", new String[0]));
assertThat(tokens.get(5)).isEqualTo(new SemanticTokenData(26, 28, "keyword", new String[0]));
assertThat(tokens.get(6)).isEqualTo(new SemanticTokenData(29, 35, "method", new String[0]));
assertThat(tokens.get(7)).isEqualTo(new SemanticTokenData(35, 36, "operator", new String[0]));
assertThat(tokens.get(8)).isEqualTo(new SemanticTokenData(36, 37, "operator", new String[0]));
assertThat(tokens.get(9)).isEqualTo(new SemanticTokenData(38, 43, "keyword", new String[0]));
assertThat(tokens.get(10)).isEqualTo(new SemanticTokenData(44, 45, "operator", new String[0]));
assertThat(tokens.get(11)).isEqualTo(new SemanticTokenData(45, 46, "parameter", new String[0]));
}
@Test
void parameterInLimitClause_3() {
List<SemanticTokenData> tokens = provider.computeTokens("SELECT * FROM cards ORDER BY random() LIMIT :#{qq}", 0);
assertThat(tokens.size()).isEqualTo(14);
assertThat(tokens.get(0)).isEqualTo(new SemanticTokenData(0, 6, "keyword", new String[0]));
assertThat(tokens.get(1)).isEqualTo(new SemanticTokenData(7, 8, "operator", new String[0]));
assertThat(tokens.get(2)).isEqualTo(new SemanticTokenData(9, 13, "keyword", new String[0]));
assertThat(tokens.get(3)).isEqualTo(new SemanticTokenData(14, 19, "variable", new String[0]));
assertThat(tokens.get(4)).isEqualTo(new SemanticTokenData(20, 25, "keyword", new String[0]));
assertThat(tokens.get(5)).isEqualTo(new SemanticTokenData(26, 28, "keyword", new String[0]));
assertThat(tokens.get(6)).isEqualTo(new SemanticTokenData(29, 35, "method", new String[0]));
assertThat(tokens.get(7)).isEqualTo(new SemanticTokenData(35, 36, "operator", new String[0]));
assertThat(tokens.get(8)).isEqualTo(new SemanticTokenData(36, 37, "operator", new String[0]));
assertThat(tokens.get(9)).isEqualTo(new SemanticTokenData(38, 43, "keyword", new String[0]));
assertThat(tokens.get(10)).isEqualTo(new SemanticTokenData(44, 45, "operator", new String[0]));
assertThat(tokens.get(11)).isEqualTo(new SemanticTokenData(45, 47, "operator", new String[0]));
assertThat(tokens.get(12)).isEqualTo(new SemanticTokenData(47, 49, "variable", new String[0]));
assertThat(tokens.get(13)).isEqualTo(new SemanticTokenData(49, 50, "operator", new String[0]));
}
@Test
void parameterInLimitClause_4() {
List<SemanticTokenData> tokens = provider.computeTokens("SELECT * FROM cards ORDER BY random() LIMIT :limit", 0);
assertThat(tokens.size()).isEqualTo(12);
assertThat(tokens.get(0)).isEqualTo(new SemanticTokenData(0, 6, "keyword", new String[0]));
assertThat(tokens.get(1)).isEqualTo(new SemanticTokenData(7, 8, "operator", new String[0]));
assertThat(tokens.get(2)).isEqualTo(new SemanticTokenData(9, 13, "keyword", new String[0]));
assertThat(tokens.get(3)).isEqualTo(new SemanticTokenData(14, 19, "variable", new String[0]));
assertThat(tokens.get(4)).isEqualTo(new SemanticTokenData(20, 25, "keyword", new String[0]));
assertThat(tokens.get(5)).isEqualTo(new SemanticTokenData(26, 28, "keyword", new String[0]));
assertThat(tokens.get(6)).isEqualTo(new SemanticTokenData(29, 35, "method", new String[0]));
assertThat(tokens.get(7)).isEqualTo(new SemanticTokenData(35, 36, "operator", new String[0]));
assertThat(tokens.get(8)).isEqualTo(new SemanticTokenData(36, 37, "operator", new String[0]));
assertThat(tokens.get(9)).isEqualTo(new SemanticTokenData(38, 43, "keyword", new String[0]));
assertThat(tokens.get(10)).isEqualTo(new SemanticTokenData(44, 45, "operator", new String[0]));
assertThat(tokens.get(11)).isEqualTo(new SemanticTokenData(45, 50, "parameter", new String[0]));
}
}