def test_get_query_with_new_limit_upper(self):
sql = "SELECT * FROM birth_names LIMIT 1555"
parsed = sql_parse.ParsedQuery(sql)
newsql = parsed.set_or_update_query_limit(1000)
# applied as new limit is lower
expected = "SELECT * FROM birth_names LIMIT 1000"
self.assertEqual(newsql, expected)
def apply_limit_to_sql(cls, sql: str, limit: int, database) -> str:
"""
Alters the SQL statement to apply a LIMIT clause
:param sql: SQL query
:param limit: Maximum number of rows to be returned by the query
:param database: Database instance
:return: SQL query with limit clause
"""
# TODO: Fix circular import caused by importing Database
if cls.limit_method == LimitMethod.WRAP_SQL:
sql = sql.strip("\t\n ;")
if database.backend == 'db2':
qry = (
select("*")
.select_from(TextAsFrom(text(sql), ["*"]).alias("inner_qry"))
)
else:
qry = (
select("*")
.select_from(TextAsFrom(text(sql), ["*"]).alias("inner_qry"))
.limit(limit)
)
return database.compile_sqla_query(qry)
elif LimitMethod.FORCE_LIMIT:
if database.backend != 'db2':
parsed_query = sql_parse.ParsedQuery(sql)
sql = parsed_query.get_query_with_new_limit(limit)
return sql
def estimate_query_cost(
cls,
database: "Database",
schema: str,
sql: str,
source: Optional[str] = None) -> List[Dict[str, Any]]:
"""
Estimate the cost of a multiple statement SQL query.
:param database: Database instance
:param schema: Database schema
:param sql: SQL query with possibly multiple statements
:param source: Source of the query (eg, "sql_lab")
"""
extra = database.get_extra() or {}
if not cls.get_allow_cost_estimate(extra):
raise Exception("Database does not support cost estimation")
user_name = g.user.username if g.user else None
parsed_query = sql_parse.ParsedQuery(sql)
statements = parsed_query.get_statements()
engine = cls.get_engine(database, schema=schema, source=source)
costs = []
with closing(engine.raw_connection()) as conn:
cursor = conn.cursor()
for statement in statements:
processed_statement = cls.process_statement(
statement, database, user_name)
costs.append(
cls.estimate_statement_cost(processed_statement, cursor))
return costs
def get_limit_from_sql(cls, sql: str) -> Optional[int]:
"""
Extract limit from SQL query
:param sql: SQL query
:return: Value of limit clause in query
"""
parsed_query = sql_parse.ParsedQuery(sql)
return parsed_query.limit
def set_or_update_query_limit(cls, sql: str, limit: int) -> str:
"""
Create a query based on original query but with new limit clause
:param sql: SQL query
:param limit: New limit to insert/replace into query
:return: Query with new limit
"""
parsed_query = sql_parse.ParsedQuery(sql)
return parsed_query.set_or_update_query_limit(limit)
def test_basic_breakdown_statements(self):
multi_sql = """
SELECT * FROM ab_user;
SELECT * FROM ab_user LIMIT 1;
"""
parsed = sql_parse.ParsedQuery(multi_sql)
statements = parsed.get_statements()
self.assertEquals(len(statements), 2)
expected = ["SELECT * FROM ab_user", "SELECT * FROM ab_user LIMIT 1"]
self.assertEquals(statements, expected)
def apply_limit_to_sql(cls, sql, limit, database):
"""Alters the SQL statement to apply a LIMIT clause"""
if cls.limit_method == LimitMethod.WRAP_SQL:
sql = sql.strip("\t\n ;")
qry = (select("*").select_from(
TextAsFrom(text(sql), ["*"]).alias("inner_qry")).limit(limit))
return database.compile_sqla_query(qry)
elif LimitMethod.FORCE_LIMIT:
parsed_query = sql_parse.ParsedQuery(sql)
sql = parsed_query.get_query_with_new_limit(limit)
return sql
def test_messy_breakdown_statements(self):
multi_sql = """
SELECT 1;\t\n\n\n \t
\t\nSELECT 2;
SELECT * FROM birth_names;;;
SELECT * FROM birth_names LIMIT 1
"""
parsed = sql_parse.ParsedQuery(multi_sql)
statements = parsed.get_statements()
self.assertEquals(len(statements), 4)
expected = [
"SELECT 1",
"SELECT 2",
"SELECT * FROM birth_names",
"SELECT * FROM birth_names LIMIT 1",
]
self.assertEquals(statements, expected)
def rejected_tables(self, sql: str, database: "Database",
schema: str) -> Set["Table"]:
"""
Return the list of rejected SQL tables.
:param sql: The SQL statement
:param database: The SQL database
:param schema: The SQL database schema
:returns: The rejected tables
"""
query = sql_parse.ParsedQuery(sql)
return {
table
for table in query.tables
if not self.can_access_datasource(database, table, schema)
}
def test_messy_breakdown_statements(self):
multi_sql = """
SELECT 1;\t\n\n\n \t
\t\nSELECT 2;
SELECT * FROM ab_user;;;
SELECT * FROM ab_user LIMIT 1
"""
parsed = sql_parse.ParsedQuery(multi_sql)
statements = parsed.get_statements()
self.assertEquals(len(statements), 4)
expected = [
'SELECT 1',
'SELECT 2',
'SELECT * FROM ab_user',
'SELECT * FROM ab_user LIMIT 1',
]
self.assertEquals(statements, expected)
def rejected_tables(self, sql: str, database: "Database",
schema: str) -> Set["Table"]:
"""
Return the list of rejected SQL tables.
:param sql: The SQL statement
:param database: The SQL database
:param schema: The SQL database schema
:returns: The rejected tables
"""
from superset.sql_parse import Table
return {
table
for table in sql_parse.ParsedQuery(sql).tables
if not self.can_access_table(
database, Table(table.table, table.schema or schema))
}
def rejected_tables(self, sql: str, database: "Database", schema: str) -> List[str]:
"""
Return the list of rejected SQL table names.
Note the rejected table names conform to the [[cluster.]schema.]table construct.
:param sql: The SQL statement
:param database: The SQL database
:param schema: The SQL database schema
:returns: The rejected table names
"""
superset_query = sql_parse.ParsedQuery(sql)
return [
t
for t in superset_query.tables
if not self._datasource_access_by_fullname(database, t, schema)
]
def test_update_not_select(self):
sql = sql_parse.ParsedQuery("UPDATE t1 SET col1 = NULL")
self.assertEquals(False, sql.is_select())
self.assertEquals(False, sql.is_readonly())
def raise_for_access(
# pylint: disable=too-many-arguments,too-many-locals
self,
database: Optional["Database"] = None,
datasource: Optional["BaseDatasource"] = None,
query: Optional["Query"] = None,
query_context: Optional["QueryContext"] = None,
table: Optional["Table"] = None,
viz: Optional["BaseViz"] = None,
) -> None:
"""
Raise an exception if the user cannot access the resource.
:param database: The Superset database
:param datasource: The Superset datasource
:param query: The SQL Lab query
:param query_context: The query context
:param table: The Superset table (requires database)
:param viz: The visualization
:raises SupersetSecurityException: If the user cannot access the resource
"""
# pylint: disable=import-outside-toplevel
from superset.connectors.sqla.models import SqlaTable
from superset.extensions import feature_flag_manager
from superset.sql_parse import Table
if database and table or query:
if query:
database = query.database
database = cast("Database", database)
if self.can_access_database(database):
return
if query:
tables = {
Table(table_.table, table_.schema or query.schema)
for table_ in sql_parse.ParsedQuery(query.sql).tables
}
elif table:
tables = {table}
denied = set()
for table_ in tables:
schema_perm = self.get_schema_perm(database,
schema=table_.schema)
if not (schema_perm
and self.can_access("schema_access", schema_perm)):
datasources = SqlaTable.query_datasources_by_name(
self.get_session,
database,
table_.table,
schema=table_.schema)
# Access to any datasource is suffice.
for datasource_ in datasources:
if self.can_access("datasource_access",
datasource_.perm):
break
else:
denied.add(table_)
if denied:
raise SupersetSecurityException(
self.get_table_access_error_object(denied))
if datasource or query_context or viz:
if query_context:
datasource = query_context.datasource
elif viz:
datasource = viz.datasource
assert datasource
should_check_dashboard_access = (
feature_flag_manager.is_feature_enabled("DASHBOARD_RBAC")
or self.is_guest_user())
if not (self.can_access_schema(datasource) or self.can_access(
"datasource_access", datasource.perm or "") or
(should_check_dashboard_access
and self.can_access_based_on_dashboard(datasource))):
raise SupersetSecurityException(
self.get_datasource_access_error_object(datasource))
def get_query_with_new_limit(cls, sql, limit):
parsed_query = sql_parse.ParsedQuery(sql)
return parsed_query.get_query_with_new_limit(limit)
def rejected_datasources(self, sql, database, schema):
superset_query = sql_parse.ParsedQuery(sql)
return [
t for t in superset_query.tables
if not self.datasource_access_by_fullname(database, t, schema)
]
def test_get_query_with_new_limit(self):
sql = "SELECT * FROM birth_names LIMIT 555"
parsed = sql_parse.ParsedQuery(sql)
newsql = parsed.get_query_with_new_limit(1000)
expected = "SELECT * FROM birth_names LIMIT 1000"
self.assertEquals(newsql, expected)
def raise_for_access(
self,
database: Optional["Database"] = None,
datasource: Optional["BaseDatasource"] = None,
query: Optional["Query"] = None,
query_context: Optional["QueryContext"] = None,
table: Optional["Table"] = None,
viz: Optional["BaseViz"] = None,
) -> None:
"""
Raise an exception if the user cannot access the resource.
:param database: The Superset database
:param datasource: The Superset datasource
:param query: The SQL Lab query
:param query_context: The query context
:param table: The Superset table (requires database)
:param viz: The visualization
:raises SupersetSecurityException: If the user cannot access the resource
"""
from superset.connectors.sqla.models import SqlaTable
from superset.sql_parse import Table
if database and table or query:
if query:
database = query.database
database = cast("Database", database)
if self.can_access_database(database):
return
if query:
tables = {
Table(table_.table, table_.schema or query.schema)
for table_ in sql_parse.ParsedQuery(query.sql).tables
}
elif table:
tables = {table}
denied = set()
for table_ in tables:
schema_perm = self.get_schema_perm(database,
schema=table_.schema)
if not (schema_perm
and self.can_access("schema_access", schema_perm)):
datasources = SqlaTable.query_datasources_by_name(
self.get_session,
database,
table_.table,
schema=table_.schema)
# Access to any datasource is suffice.
for datasource in datasources:
if self.can_access("datasource_access",
datasource.perm):
break
else:
denied.add(table_)
if denied:
raise SupersetSecurityException(
self.get_table_access_error_object(denied), )
if datasource or query_context or viz:
if query_context:
datasource = query_context.datasource
elif viz:
datasource = viz.datasource
assert datasource
if not (self.can_access_schema(datasource) or self.can_access(
"datasource_access", datasource.perm or "")):
raise SupersetSecurityException(
self.get_datasource_access_error_object(datasource), )
def test_get_query_with_new_limit_comment_with_limit(self):
sql = "SELECT * FROM birth_names -- SOME COMMENT WITH LIMIT 555"
parsed = sql_parse.ParsedQuery(sql)
newsql = parsed.get_query_with_new_limit(1000)
self.assertEquals(newsql, sql + "\nLIMIT 1000")
def test_explain(self):
sql = sql_parse.ParsedQuery("EXPLAIN SELECT 1")
self.assertEquals(True, sql.is_explain())
self.assertEquals(False, sql.is_select())
self.assertEquals(True, sql.is_readonly())
def get_limit_from_sql(cls, sql):
parsed_query = sql_parse.ParsedQuery(sql)
return parsed_query.limit
def test_get_query_with_new_limit_comment_with_limit(self):
sql = 'SELECT * FROM ab_user --SOME COMMENT WITH LIMIT 555'
parsed = sql_parse.ParsedQuery(sql)
newsql = parsed.get_query_with_new_limit(1000)
self.assertEquals(newsql, sql + '\nLIMIT 1000')
def test_get_query_with_new_limit(self):
sql = 'SELECT * FROM ab_user LIMIT 555'
parsed = sql_parse.ParsedQuery(sql)
newsql = parsed.get_query_with_new_limit(1000)
expected = 'SELECT * FROM ab_user LIMIT 1000'
self.assertEquals(newsql, expected)
def extract_tables(self, query):
sq = sql_parse.ParsedQuery(query)
return sq.tables
def test_get_query_with_new_limit_comment(self):
sql = "SELECT * FROM birth_names -- SOME COMMENT"
parsed = sql_parse.ParsedQuery(sql)
newsql = parsed.set_or_update_query_limit(1000)
self.assertEqual(newsql, sql + "\nLIMIT 1000")