From 5ae5b6f61ab51ee4d37eb85ae55d2ee2fb7d21fd Mon Sep 17 00:00:00 2001 From: WilliamZhu Date: Thu, 25 Jun 2020 23:19:46 +0800 Subject: [PATCH] fix table alias problem --- .../mlsql/autosuggest/dsl/TokenMatcher.scala | 12 ++++----- .../autosuggest/statement/LexerUtils.scala | 16 +++++------ .../statement/TableExtractor.scala | 27 ++++++++++--------- .../autosuggest/AutoSuggestContextTest.scala | 12 +++++++++ .../autosuggest/SelectSuggesterTest.scala | 2 ++ .../antlr4/autosuggest/SelectTest.scala | 18 ------------- 6 files changed, 42 insertions(+), 45 deletions(-) delete mode 100644 src/test/java/com/intigua/antlr4/autosuggest/SelectTest.scala diff --git a/src/main/java/tech/mlsql/autosuggest/dsl/TokenMatcher.scala b/src/main/java/tech/mlsql/autosuggest/dsl/TokenMatcher.scala index 6457e74..6a93d99 100644 --- a/src/main/java/tech/mlsql/autosuggest/dsl/TokenMatcher.scala +++ b/src/main/java/tech/mlsql/autosuggest/dsl/TokenMatcher.scala @@ -159,7 +159,7 @@ class TokenMatcher(tokens: List[Token], val start: Int) { foods.foreach { foodw => - if (currentIndex >= tokens.size) { + if (currentIndex >= tokens.size && !foodw.optional) { isFail = true } else { val stepSize = foodw.foods.count @@ -184,10 +184,7 @@ class TokenMatcher(tokens: List[Token], val start: Int) { } } } - - if (!isFail && currentIndex == tokens.size) { - currentIndex = tokens.size - 1 - } + val targetIndex = if (isFail) -1 else currentIndex cacheResult = targetIndex this @@ -200,7 +197,7 @@ class TokenMatcher(tokens: List[Token], val start: Int) { foods.foreach { foodw => // if out of bound then mark fail - if (currentIndex <= -1) { + if (currentIndex <= -1 && !foodw.optional) { isFail = true } else { val stepSize = foodw.foods.count @@ -276,6 +273,9 @@ object TokenTypeWrapper { val LEFT_SQUARE_BRACKET = SqlBaseLexer.T__7 //[ val RIGHT_SQUARE_BRACKET = SqlBaseLexer.T__8 //] val COLON = SqlBaseLexer.T__9 //: + + val LIST = List(LEFT_BRACKET, RIGHT_BRACKET, COMMA, DOT, LEFT_SQUARE_BRACKET, RIGHT_SQUARE_BRACKET, COLON) + val MAP = LIST.map((_, 1)).toMap } object MLSQLTokenTypeWrapper { diff --git a/src/main/java/tech/mlsql/autosuggest/statement/LexerUtils.scala b/src/main/java/tech/mlsql/autosuggest/statement/LexerUtils.scala index 17b5c2b..f898f13 100644 --- a/src/main/java/tech/mlsql/autosuggest/statement/LexerUtils.scala +++ b/src/main/java/tech/mlsql/autosuggest/statement/LexerUtils.scala @@ -85,7 +85,7 @@ object LexerUtils { && colNum <= lastToken._1.getCharPositionInLine + lastToken._1.getText.size && (lastToken._1.getType != DSLSQLLexer.UNRECOGNIZED - && lastToken._1.getType != MLSQLTokenTypeWrapper.DOT) + && lastToken._1.getType != MLSQLTokenTypeWrapper.DOT) ) { return TokenPos(lastToken._2, TokenPosType.CURRENT, colNum - lastToken._1.getCharPositionInLine) } @@ -151,7 +151,11 @@ object LexerUtils { return TokenPos(lastToken._2, TokenPosType.NEXT, 0) } - if (colNum > lastToken._1.getCharPositionInLine && colNum <= lastToken._1.getCharPositionInLine + lastToken._1.getText.size) { + if (colNum > lastToken._1.getCharPositionInLine + && colNum <= lastToken._1.getCharPositionInLine + lastToken._1.getText.size + && !TokenTypeWrapper.MAP.contains(lastToken._1.getType) + + ) { return TokenPos(lastToken._2, TokenPosType.CURRENT, colNum - lastToken._1.getCharPositionInLine) } oneLineTokens.map { case (token, index) => @@ -161,13 +165,7 @@ object LexerUtils { //这个token是 [(,).]等，则不算 if (colNum == end && (1 <= token.getType) && ( - token.getType == DSLSQLLexer.UNRECOGNIZED - || token.getType == TokenTypeWrapper.COLON - || token.getType == TokenTypeWrapper.DOT - || token.getType == TokenTypeWrapper.LEFT_BRACKET - || token.getType == TokenTypeWrapper.RIGHT_BRACKET - || token.getType == TokenTypeWrapper.LEFT_SQUARE_BRACKET - || token.getType == TokenTypeWrapper.RIGHT_SQUARE_BRACKET + TokenTypeWrapper.MAP.contains(token.getType) )) { TokenPos(index, TokenPosType.NEXT, 0) } else if (start < colNum && colNum <= end) { diff --git a/src/main/java/tech/mlsql/autosuggest/statement/TableExtractor.scala b/src/main/java/tech/mlsql/autosuggest/statement/TableExtractor.scala index 0206005..b301b57 100644 --- a/src/main/java/tech/mlsql/autosuggest/statement/TableExtractor.scala +++ b/src/main/java/tech/mlsql/autosuggest/statement/TableExtractor.scala @@ -24,19 +24,22 @@ class TableExtractor(autoSuggestContext: AutoSuggestContext, ast: SingleStatemen override def extractor(start: Int, end: Int): List[MetaTableKeyWrapper] = { val dbTableTokens = tokens.slice(start, end) - val dbTable = if (dbTableTokens.length == 3) { - val List(dbToken, _, tableToken) = dbTableTokens - MetaTableKeyWrapper(MetaTableKey(None, Option(dbToken.getText), tableToken.getText), None) - } else if (dbTableTokens.length == 4) { - val List(dbToken, _, tableToken, aliasToken) = dbTableTokens - MetaTableKeyWrapper(MetaTableKey(None, Option(dbToken.getText), tableToken.getText), Option(aliasToken.getText)) - } else if (dbTableTokens.length == 5) { - val List(dbToken, _, tableToken, _, aliasToken) = dbTableTokens - MetaTableKeyWrapper(MetaTableKey(None, Option(dbToken.getText), tableToken.getText), Option(aliasToken.getText)) - } - else { - MetaTableKeyWrapper(MetaTableKey(None, None, dbTableTokens.head.getText), None) + val dbTable = dbTableTokens.length match { + case 2 => + val List(tableToken, aliasToken) = dbTableTokens + MetaTableKeyWrapper(MetaTableKey(None, None, tableToken.getText), Option(aliasToken.getText)) + case 3 => + val List(dbToken, _, tableToken) = dbTableTokens + MetaTableKeyWrapper(MetaTableKey(None, Option(dbToken.getText), tableToken.getText), None) + case 4 => + val List(dbToken, _, tableToken, aliasToken) = dbTableTokens + MetaTableKeyWrapper(MetaTableKey(None, Option(dbToken.getText), tableToken.getText), Option(aliasToken.getText)) + case 5 => + val List(dbToken, _, tableToken, _, aliasToken) = dbTableTokens + MetaTableKeyWrapper(MetaTableKey(None, Option(dbToken.getText), tableToken.getText), Option(aliasToken.getText)) + case _ => MetaTableKeyWrapper(MetaTableKey(None, None, dbTableTokens.head.getText), None) } + List(dbTable) } diff --git a/src/test/java/com/intigua/antlr4/autosuggest/AutoSuggestContextTest.scala b/src/test/java/com/intigua/antlr4/autosuggest/AutoSuggestContextTest.scala index f1365b5..592c360 100644 --- a/src/test/java/com/intigua/antlr4/autosuggest/AutoSuggestContextTest.scala +++ b/src/test/java/com/intigua/antlr4/autosuggest/AutoSuggestContextTest.scala @@ -181,6 +181,18 @@ class AutoSuggestContextTest extends BaseTest with BeforeAndAfterEach { println(items) } + test("table alias with temp table") { + val sql = + """ + |select a,b,c from table1 as table1; + |select aa,bb,cc from table2 as table2; + |select from table1 t1 left join table2 t2 on t1.a = t2. + |""".stripMargin + + val items = context.buildFromString(sql).suggest(4, 58) + assert(items.map(_.name) == List("aa", "bb", "cc")) + + } } diff --git a/src/test/java/com/intigua/antlr4/autosuggest/SelectSuggesterTest.scala b/src/test/java/com/intigua/antlr4/autosuggest/SelectSuggesterTest.scala index f1128ac..2db9364 100644 --- a/src/test/java/com/intigua/antlr4/autosuggest/SelectSuggesterTest.scala +++ b/src/test/java/com/intigua/antlr4/autosuggest/SelectSuggesterTest.scala @@ -255,4 +255,6 @@ class SelectSuggesterTest extends BaseTest { } + + } diff --git a/src/test/java/com/intigua/antlr4/autosuggest/SelectTest.scala b/src/test/java/com/intigua/antlr4/autosuggest/SelectTest.scala deleted file mode 100644 index 2c9ab0a..0000000 --- a/src/test/java/com/intigua/antlr4/autosuggest/SelectTest.scala +++ /dev/null @@ -1,18 +0,0 @@ -package com.intigua.antlr4.autosuggest - -/** - * 11/6/2020 WilliamZhu(allwefantasy@gmail.com) - */ -class SelectTest extends BaseTestWithoutSparkSession { - test("spark sql test") { - - val items = context.buildFromString( - """ - | -- yes - | select from (select a,b,c from table1) as table2; - |""".stripMargin).suggest(3, 8) - - println(items) - - } -}