villebro commented on a change in pull request #18746:
URL: https://github.com/apache/superset/pull/18746#discussion_r809721602
##########
File path: superset/db_engine_specs/base.py
##########
@@ -649,6 +658,69 @@ def apply_limit_to_sql(
return sql
+ @classmethod
+ def apply_top_to_sql(
+ cls, sql: str, _limit: int, database: str = "Database", force: bool =
False
+ ) -> str:
+ """
+ Alters the SQL statement to apply a TOP clause
+ :param _limit: Maximum number of rows to be returned by the query
+ :param sql: SQL query
+ :param database: Database instance
+ :return: SQL query with top clause
+ """
+
+ cte = None
+ sql_remainder = None
+ sql_statement = sqlparse.format(sql, strip_comments=True)
+ query_limit: Optional[int] =
sql_parse._extract_top_from_query(sql_statement,
+
cls.top_keywords)
+ if not _limit:
+ final_limit = query_limit
+ elif int(query_limit or 0) < _limit and query_limit is not None:
+ final_limit = query_limit
+ else:
+ final_limit = _limit
+ if not cls.allows_cte_in_subquery:
+ cte, sql_remainder =
sql_parse.get_cte_reminder_query(sql_statement)
Review comment:
typo
```suggestion
cte, sql_remainder =
sql_parse.get_cte_remainder_query(sql_statement)
```
##########
File path: superset/db_engine_specs/base.py
##########
@@ -80,7 +80,6 @@
logger = logging.getLogger()
-
Review comment:
Let's avoid adding unrelated changes
##########
File path: superset/db_engine_specs/base.py
##########
@@ -300,6 +299,16 @@ class BaseEngineSpec: # pylint:
disable=too-many-public-methods
# If True, then it will allow in subquery ,
# if False it will allow as regular CTE
allows_cte_in_subquery = True
+ # Whether allow LIMIT clause in the SQL
+ # If True, then the database engine is allowed for LIMIT clause
+ # If False, then the database engine is allowed for TOP clause
+ allow_limit_clause = True
+ # This list will give keywords for select statements
+ # to consider for the engines with TOP SQl parsing
+ select_keywords: Set[str] = None
+ # This list will give the keywords for data limit statements
+ # to consider for the engines with TOP SQl parsing
+ top_keywords: Set[str] = None
Review comment:
I think it's fine to set the defaults here as I assume the majority of
TOP-syntax engines won't need to override these. Also note the comment typo
"SQl" (should be "SQL") and "list" replaced by "set".
```suggestion
# This set will give keywords for select statements
# to consider for the engines with TOP SQL parsing
select_keywords: Set[str] = {"SELECT"}
# This set will give the keywords for data limit statements
# to consider for the engines with TOP SQL parsing
top_keywords: Set[str] = {"TOP"}
```
##########
File path: superset/sql_lab.py
##########
@@ -287,15 +287,27 @@ def execute_sql_statement( # pylint:
disable=too-many-arguments,too-many-locals
return SupersetResultSet(data, cursor_description, db_engine_spec)
+def apply_cte(
+ database: Database, sql: str
+) -> str:
+ # Use the Basengine function to parse the SQL for CTE
+ sql = database.db_engine_spec.get_cte_query(sql)
+ return sql
+
+
Review comment:
This function doesn't appear to be used anywhere - I assume we can
delete it:
```suggestion
```
##########
File path: superset/sql_parse.py
##########
@@ -30,7 +30,8 @@
Token,
TokenList,
)
-from sqlparse.tokens import DDL, DML, Keyword, Name, Punctuation, String,
Whitespace
+from sqlparse.tokens import DDL, DML, Keyword, Name, Punctuation, String,
Whitespace, \
+ CTE
Review comment:
Please keep these in alphabetical order (`CTE` first in the list). This
will also most likely need to be delinted. Check `CONTRIBUTING.md` on how to
use `pre-commit` to automatically fix linting issues in your commits:
https://github.com/apache/superset/blob/master/CONTRIBUTING.md#git-hooks
```suggestion
from sqlparse.tokens import (
CTE,
DDL,
DML,
Keyword,
Name,
Punctuation,
String,
Whitespace,
)
```
##########
File path: superset/sql_parse.py
##########
@@ -78,6 +78,56 @@ def _extract_limit_from_query(statement: TokenList) ->
Optional[int]:
return None
+def _extract_top_from_query(statement: TokenList, top_keyword: Set[str]) ->
Optional[
+ int]:
+ """
+ Extract top clause value from SQL statement.
+
+ :param statement: SQL statement
+ :return: top value extracted from query, None if no top value present in
statement
+ """
+
+ str_statement = str(statement)
+ str_statement = str_statement.replace("\n", " ").replace("\r", "")
+ token = str_statement.rstrip().split(" ")
+ token = [part for part in token if part]
+ top = None
+ for i, _ in enumerate(token):
+ if token[i].upper() in top_keyword and len(token) - 1 > i:
+ try:
+ top = int(token[i + 1])
+ except ValueError:
+ top = None
+ break
+ return top
+
+
+def get_cte_reminder_query(sql: str) -> Optional[str]:
Review comment:
typo:
```suggestion
def get_cte_remainder_query(sql: str) -> Optional[str]:
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
