sujiplr commented on a change in pull request #18746:
URL: https://github.com/apache/superset/pull/18746#discussion_r809731976
##########
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:
Corrected.
##########
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:
Corrected
--
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]
