sujiplr commented on a change in pull request #18746:
URL: https://github.com/apache/superset/pull/18746#discussion_r808939909
##########
File path: superset/db_engine_specs/mssql.py
##########
@@ -48,6 +48,9 @@ class MssqlEngineSpec(BaseEngineSpec):
limit_method = LimitMethod.WRAP_SQL
max_column_name_length = 128
allows_cte_in_subquery = False
+ allow_limit_clause = False
+ sel_keywords = {"SELECT"}
+ top_keywords = {"TOP"}
Review comment:
Corrected this.
##########
File path: superset/db_engine_specs/base.py
##########
@@ -638,8 +647,8 @@ def apply_limit_to_sql(
sql = sql.strip("\t\n ;")
qry = (
select("*")
- .select_from(TextAsFrom(text(sql), ["*"]).alias("inner_qry"))
- .limit(limit)
+ .select_from(TextAsFrom(text(sql),
["*"]).alias("inner_qry"))
+ .limit(limit)
Review comment:
Corrected this
##########
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 the default list of keywords for select statements
+ # to consider in the limit/top SQl parsing
+ select_keywords = {"SELECT"}
+ # This list will give the default list of keywords for data limit
statements
+ # to consider in the limit/top SQl parsing
+ top_keywords = {"LIMIT"}
Review comment:
corrected this
##########
File path: superset/db_engine_specs/base.py
##########
@@ -649,6 +658,70 @@ 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
+ """
+
+ td_sel_keywords = cls.select_keywords
+ td_top_keywords = cls.top_keywords
+ 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,
+
td_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)
+ if cte:
+ str_statement = str(sql_remainder)
+ cte = cte + "\n"
+ else:
+ cte = ""
+ str_statement = str(sql)
+ str_statement = str_statement.replace("\n", " ").replace("\r", "")
+
+ tokens = str_statement.rstrip().split(" ")
+ tokens = [token for token in tokens if token]
+ if cls.top_not_in_sql(str_statement, td_top_keywords):
+ selects = [i for i, word in enumerate(tokens) if
+ word.upper() in td_sel_keywords]
+ first_select = selects[0]
+ tokens.insert(first_select + 1, "TOP")
+ tokens.insert(first_select + 2, str(final_limit))
+
+ next_is_limit_token = False
+ new_tokens = []
+
+ for token in tokens:
+ if token in td_top_keywords:
+ next_is_limit_token = True
+ elif next_is_limit_token:
+ if token.isdigit():
+ token = str(final_limit)
+ next_is_limit_token = False
+ new_tokens.append(token)
+ sql = " ".join(new_tokens)
+ return cte + sql
+
+ def top_not_in_sql(sql: str, top_words: Set[str]) -> bool:
+ for top_word in top_words:
+ if top_word.upper() in sql.upper():
+ return False
+ return True
Review comment:
Corrected this
##########
File path: superset/db_engine_specs/base.py
##########
@@ -649,6 +658,70 @@ 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
+ """
+
+ td_sel_keywords = cls.select_keywords
+ td_top_keywords = cls.top_keywords
+ 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,
+
td_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)
+ if cte:
+ str_statement = str(sql_remainder)
+ cte = cte + "\n"
+ else:
+ cte = ""
+ str_statement = str(sql)
+ str_statement = str_statement.replace("\n", " ").replace("\r", "")
+
+ tokens = str_statement.rstrip().split(" ")
+ tokens = [token for token in tokens if token]
+ if cls.top_not_in_sql(str_statement, td_top_keywords):
Review comment:
Corrected this
--
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]
