Dealing with Larger Dataset and Optimizing SQL Exports to Avoid Application Crashes

[0xsharz]

I encountered an issue while using the default SQLTool. When attempting to list or export large datasets, the application would crash. This occurred because it first printed all the data in the chat before using FileTool to export it.
To resolve this problem, I added a new function to SQLTool called save_query_results. This function automatically executes queries for larger datasets that the user wants to export and saves the results directly to a specified file. This approach bypasses the need to print the data in the chat, preventing the application from breaking down when handling large amounts of data.

def save_query_results(self, query: str, file_name: str, chunk_size: int = 1000) -> str:
        """Executes the given SQL query and saves the results to a file.

        Args:
            query (str): The SQL query to execute.
            file_name (str): The name of the file to save the results to.
            chunk_size (int, optional): The number of rows to fetch and write at a time. Defaults to 1000.

        Returns:
            str: The file name if successful, otherwise returns an error message.
        """
        try:
            logger.info(f"Executing query and saving results to file: {file_name}")
            file_path = self.base_dir.joinpath(file_name)

            with self.Session() as sess, sess.begin():
                result = sess.execute(text(query))

                with open(file_path, 'w', newline='') as f:
                    writer = csv.writer(f)
                    writer.writerow(result.keys())  # Write header

                    while True:
                        chunk = result.fetchmany(chunk_size)
                        if not chunk:
                            break
                        writer.writerows(chunk)

            logger.info(f"Query results saved to {file_path}")
            return str(file_name)
        except Exception as e:
            logger.error(f"Error saving query results: {e}")
            return f"Error saving query results: {e}"
 

Hope you like it :)

[0xsharz]

sql.txt
Here is the full sql.py code.

[manthanguptaa]

Hey @0xsharz, are you still facing this issue?