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)
Пример #2
0
    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
Пример #3
0
    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
Пример #4
0
    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
Пример #5
0
    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)
Пример #6
0
 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)
Пример #7
0
 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
Пример #8
0
 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)
Пример #9
0
    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)
        }
Пример #10
0
 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)
Пример #11
0
    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))
        }
Пример #12
0
    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)
        ]
Пример #13
0
 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())
Пример #14
0
    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))
Пример #15
0
 def get_query_with_new_limit(cls, sql, limit):
     parsed_query = sql_parse.ParsedQuery(sql)
     return parsed_query.get_query_with_new_limit(limit)
Пример #16
0
 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)
     ]
Пример #17
0
 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)
Пример #18
0
    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), )
Пример #19
0
 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")
Пример #20
0
    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())
Пример #21
0
 def get_limit_from_sql(cls, sql):
     parsed_query = sql_parse.ParsedQuery(sql)
     return parsed_query.limit
Пример #22
0
 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')
Пример #23
0
 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)
Пример #24
0
 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")