mirror of https://github.com/apache/superset.git
chore: Set limit for a query in execute_sql_statement (#20066)
* test for query limit * fixed tests
This commit is contained in:
parent
55aef4db3c
commit
e69f6292c2
|
@ -159,6 +159,7 @@ class SqlJsonExecutionContext: # pylint: disable=too-many-instance-attributes
|
|||
start_time=start_time,
|
||||
tab_name=self.tab_name,
|
||||
status=self.status,
|
||||
limit=self.limit,
|
||||
sql_editor_id=self.sql_editor_id,
|
||||
tmp_table_name=self.create_table_as_select.target_table_name, # type: ignore
|
||||
tmp_schema_name=self.create_table_as_select.target_schema_name, # type: ignore
|
||||
|
@ -172,6 +173,7 @@ class SqlJsonExecutionContext: # pylint: disable=too-many-instance-attributes
|
|||
select_as_cta=False,
|
||||
start_time=start_time,
|
||||
tab_name=self.tab_name,
|
||||
limit=self.limit,
|
||||
status=self.status,
|
||||
sql_editor_id=self.sql_editor_id,
|
||||
user_id=self.user_id,
|
||||
|
|
|
@ -135,13 +135,13 @@ def cta_result(ctas_method: CtasMethod):
|
|||
|
||||
|
||||
# TODO(bkyryliuk): quote table and schema names for all databases
|
||||
def get_select_star(table: str, schema: Optional[str] = None):
|
||||
def get_select_star(table: str, limit: int, schema: Optional[str] = None):
|
||||
if backend() in {"presto", "hive"}:
|
||||
schema = quote_f(schema)
|
||||
table = quote_f(table)
|
||||
if schema:
|
||||
return f"SELECT *\nFROM {schema}.{table}"
|
||||
return f"SELECT *\nFROM {table}"
|
||||
return f"SELECT *\nFROM {schema}.{table}\nLIMIT {limit}"
|
||||
return f"SELECT *\nFROM {table}\nLIMIT {limit}"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("ctas_method", [CtasMethod.TABLE, CtasMethod.VIEW])
|
||||
|
@ -236,8 +236,9 @@ def test_run_sync_query_cta_config(setup_sqllab, ctas_method):
|
|||
f"CREATE {ctas_method} {CTAS_SCHEMA_NAME}.{tmp_table_name} AS \n{QUERY}"
|
||||
== query.executed_sql
|
||||
)
|
||||
|
||||
assert query.select_sql == get_select_star(tmp_table_name, schema=CTAS_SCHEMA_NAME)
|
||||
assert query.select_sql == get_select_star(
|
||||
tmp_table_name, limit=query.limit, schema=CTAS_SCHEMA_NAME
|
||||
)
|
||||
results = run_sql(query.select_sql)
|
||||
assert QueryStatus.SUCCESS == results["status"], result
|
||||
|
||||
|
@ -266,7 +267,10 @@ def test_run_async_query_cta_config(setup_sqllab, ctas_method):
|
|||
query = wait_for_success(result)
|
||||
|
||||
assert QueryStatus.SUCCESS == query.status
|
||||
assert get_select_star(tmp_table_name, schema=CTAS_SCHEMA_NAME) == query.select_sql
|
||||
assert (
|
||||
get_select_star(tmp_table_name, limit=query.limit, schema=CTAS_SCHEMA_NAME)
|
||||
== query.select_sql
|
||||
)
|
||||
assert (
|
||||
f"CREATE {ctas_method} {CTAS_SCHEMA_NAME}.{tmp_table_name} AS \n{QUERY}"
|
||||
== query.executed_sql
|
||||
|
@ -290,7 +294,7 @@ def test_run_async_cta_query(setup_sqllab, ctas_method):
|
|||
query = wait_for_success(result)
|
||||
|
||||
assert QueryStatus.SUCCESS == query.status
|
||||
assert get_select_star(table_name) in query.select_sql
|
||||
assert get_select_star(table_name, query.limit) in query.select_sql
|
||||
|
||||
assert f"CREATE {ctas_method} {table_name} AS \n{QUERY}" == query.executed_sql
|
||||
assert QUERY == query.sql
|
||||
|
@ -313,14 +317,20 @@ def test_run_async_cta_query_with_lower_limit(setup_sqllab, ctas_method):
|
|||
QUERY, cta=True, ctas_method=ctas_method, async_=True, tmp_table=tmp_table
|
||||
)
|
||||
query = wait_for_success(result)
|
||||
|
||||
assert QueryStatus.SUCCESS == query.status
|
||||
|
||||
assert get_select_star(tmp_table) == query.select_sql
|
||||
sqllite_select_sql = f"SELECT *\nFROM {tmp_table}\nLIMIT {query.limit}\nOFFSET 0"
|
||||
assert query.select_sql == (
|
||||
sqllite_select_sql
|
||||
if backend() == "sqlite"
|
||||
else get_select_star(tmp_table, query.limit)
|
||||
)
|
||||
|
||||
assert f"CREATE {ctas_method} {tmp_table} AS \n{QUERY}" == query.executed_sql
|
||||
assert QUERY == query.sql
|
||||
|
||||
assert query.rows == (1 if backend() == "presto" else 0)
|
||||
assert query.limit is None
|
||||
assert query.limit == 10000
|
||||
assert query.select_as_cta
|
||||
assert query.select_as_cta_used
|
||||
|
||||
|
|
Loading…
Reference in New Issue