Repository: gventuri/pandas-ai
File: pandasai/helpers/code_manager.py (Python)

Code before:

    def check_direct_sql_func_def_exists(self, node: ast.AST):
        return (
            self._validate_direct_sql(self._dfs)
            and isinstance(node, ast.FunctionDef)
            and node.name == "execute_sql_query"
        )

    def _check_is_query_using_relevant_table(self, node: ast.Assign):
        for target in node.targets:
            if (
                isinstance(target, ast.Name)
                and target.id in "sql_query"
                and isinstance(node.value, ast.Constant)
                and isinstance(node.value.value, str)
            ):
                sql_query = node.value.value
                table_names = extract_table_names(sql_query)
                allowed_table_names = [df.table_name for df in self._dfs]
                if any(
                    table_name not in allowed_table_names for table_name in table_names
                ):
                    raise MaliciousQueryError(
                        "Query uses unauthorized tables: {unauthorized_tables}. Please add them as new datatables or update the query."
                    )

    def _clean_code(self, code: str, context: CodeExecutionContext) -> str:
        """
        A method to clean the code to prevent malicious code execution.

        Args:
            code(str): A python code.

        Returns:
            str: A clean code string.

        """

        # Clear recent optional dependencies
        self._additional_dependencies = []

        tree = ast.parse(code)

        # Check for imports and the node where analyze_data is defined
        new_body = []
        for node in tree.body:
            if isinstance(node, (ast.Import, ast.ImportFrom)):
                self._check_imports(node)
                continue

            if (

Review comment (refactor):
May I suggest making this utility return a bool and raising `MaliciousQueryError` when you call it (if it returns False)? In this scenario, I'd rename it to `_is_query_using_relevant_table`.

Code after:

    def check_direct_sql_func_def_exists(self, node: ast.AST):
        return (
            self._validate_direct_sql(self._dfs)
            and isinstance(node, ast.FunctionDef)
            and node.name == "execute_sql_query"
        )

    def _get_sql_irrelevant_tables(self, node: ast.Assign):
        for target in node.targets:
            if (
                isinstance(target, ast.Name)
                and target.id in "sql_query"
                and isinstance(node.value, ast.Constant)
                and isinstance(node.value.value, str)
            ):
                sql_query = node.value.value
                table_names = extract_table_names(sql_query)
                allowed_table_names = [df.table_name for df in self._dfs]
                return [
                    table_name
                    for table_name in table_names
                    if table_name not in allowed_table_names
                ]

    def _clean_code(self, code: str, context: CodeExecutionContext) -> str:
        """
        A method to clean the code to prevent malicious code execution.

        Args:
            code(str): A python code.

        Returns:
            str: A clean code string.

        """

        # Clear recent optional dependencies
        self._additional_dependencies = []

        tree = ast.parse(code)

        # Check for imports and the node where analyze_data is defined
        new_body = []
        for node in tree.body:
            if isinstance(node, (ast.Import, ast.ImportFrom)):
                self._check_imports(node)
                continue

            if (
                self._is_df_overwrite(node)