fix(sql): use only added tables of connector

[ArslanSaleem]

Summary by CodeRabbit

• New Features

  • Added a method for checking allowed tables in SQL queries.
  • Enhanced security by introducing a new exception for malicious queries.

• Bug Fixes

  • Improved handling of table names in SQL queries to prevent potential misuse.

• Tests

  • Added new test methods to validate security checks on SQL queries and table name assignments.

[coderabbitai]

Walkthrough

The recent updates involve enhancing security and functionality in the handling of SQL queries within a data management library. A method for verifying the relevance of tables in SQL queries was added, alongside a new exception to handle malicious queries. There's also a refinement in the initialization of table names, ensuring they default to a connector's fallback name. Test cases have been expanded to cover these new features and changes.

Changes

File Path	Change Summary
pandasai/helpers/code_manager.py tests/test_codemanager.py	Added imports, a new exception, a method to check table relevance in SQL queries, and modified code cleaning to include SQL query checks. Also, added new test methods for checking query validity and table relevance.
pandasai/smart_dataframe/__init__.py	Altered table name setting logic to always use the connector's fallback name if available.
pandasai/helpers/sql.py	Introduced a new file with a function to extract table names from SQL queries.

🐇
To code, to guard, a rabbit's vow,
Through lines of logic, we hop and plow.
With each query checked, and tables named,
In data's burrow, our trust is claimed. 🌟
🐇

Thank you for using CodeRabbit. We offer it for free to the OSS community and would appreciate your support in helping us grow. If you find it useful, would you consider giving us a shout-out on your favorite social media?

Share

• X
• Mastodon
• Reddit
• LinkedIn

---

Tips

Chat

There are 3 ways to chat with CodeRabbit:

• Review comments: Directly reply to a review comment made by CodeRabbit. Example:
  • I pushed a fix in commit <commit_id>.
  • Generate unit-tests for this file.
• Files and specific lines of code (under the "Files changed" tab): Tag @coderabbitai in a new review comment at the desired location with your query. Examples:
  • @coderabbitai generate unit tests for this file.
  • @coderabbitai modularize this function.
• PR comments: Tag @coderabbitai in a new PR comment to ask questions about the PR branch. For the best results, please provide a very specific query, as very limited context is provided in this mode. Examples:
  • @coderabbitai generate interesting stats about this repository from git and render them as a table.
  • @coderabbitai show all the console.log statements in this repository.
  • @coderabbitai read src/utils.ts and generate unit tests.
  • @coderabbitai read the files in the src/scheduler package and generate a class diagram using mermaid and a README in the markdown format.

Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments.

CodeRabbit Commands (invoked as PR comments)

• @coderabbitai pause to pause the reviews on a PR.
• @coderabbitai resume to resume the paused reviews.
• @coderabbitai review to trigger a review. This is useful when automatic reviews are disabled for the repository.
• @coderabbitai resolve resolve all the CodeRabbit review comments.
• @coderabbitai help to get help.

Additionally, you can add @coderabbitai ignore anywhere in the PR description to prevent this PR from being reviewed.

CodeRabbit Configration File (.coderabbit.yaml)

• You can programmatically configure CodeRabbit by adding a .coderabbit.yaml file to the root of your repository.
• The JSON schema for the configuration file is available here.
• If your editor has YAML language server enabled, you can add the path at the top of this file to enable auto-completion and validation: # yaml-language-server: $schema=https://coderabbit.ai/integrations/coderabbit-overrides.v2.json

CodeRabbit Discord Community

Join our Discord Community to get help, request features, and share feedback.

[coderabbitai]

Review Status

Actionable comments generated: 4

Configuration used: CodeRabbit UI

Commits

Files that changed from the base of the PR and between 21819c6 and 0f807b7.

Files selected for processing (3)

• pandasai/helpers/code_manager.py (4 hunks)
• pandasai/smart_dataframe/init.py (1 hunks)
• tests/test_codemanager.py (4 hunks)

Additional comments: 8

tests/test_codemanager.py (4)

• 2-2: The import of the ast module is correctly added to support the new test cases that involve parsing Python code to check for SQL query validity.

• 12-12: The replacement of NoCodeFoundError with MaliciousQueryError in the import statement aligns with the new exception handling strategy for malicious SQL queries.

• 587-603: The test test_check_is_query_using_relevant_table_multiple_tables correctly checks that no exception is raised when a query involving multiple valid tables is passed. This test ensures that the method can handle complex queries with joins.

• 622-639: The test test_check_is_query_using_relevant_table_multiple_tables_one_unknown correctly raises a MaliciousQueryError when a query with at least one unknown table is passed. This test ensures that the method correctly identifies unauthorized table access even when other tables in the query are valid.

pandasai/smart_dataframe/__init__.py (1)

• 286-286: The logic for setting the table name in the __init__ method has been modified to always set the table name to the connector's fallback name if a connector is present. This change ensures consistency in how the table name is set and is in line with the PR objectives.

pandasai/helpers/code_manager.py (3)

• 11-11: The import of extract_table_names is consistent with the PR objectives to enhance SQL query security.

• 18-18: The addition of MaliciousQueryError to the exceptions list is appropriate for handling unauthorized table access in SQL queries.

• 388-390: The check for allowed tables in SQL queries within the _clean_code method is a good security practice. Ensure that this logic is thoroughly tested with various SQL query inputs.

Verification successful

---

The search results confirm that there are tests in tests/test_codemanager.py designed to verify the functionality of the _check_is_query_using_relevant_table method. The tests cover various scenarios, including invalid queries, valid queries, multiple tables, unknown tables, and multiple tables with at least one unknown. This indicates that the logic for checking allowed tables in SQL queries within the _clean_code method has been tested thoroughly.

Scripts Executed

The following scripts were executed for the analysis:

---

Script:

#!/bin/bash
# Search for tests that verify the allowed tables check in SQL queries.
rg --type python 'test_check_is_query_using_relevant_table'

Length of output: 90

---

---

Script:

#!/bin/bash
# Search for tests that verify the allowed tables check in SQL queries without specifying file type.
rg 'test_check_is_query_using_relevant_table'

Length of output: 509

[coderabbitai]

The test test_check_is_query_using_relevant_table_invalid_query correctly raises a MaliciousQueryError when an invalid query is passed. However, the error message "Query uses tables you don't provided, add new datatable or update query" contains grammatical errors.

- "Query uses tables you don't provided, add new datatable or update query"
+ "Query uses tables you haven't provided. Add a new datatable or update the query."

---

Committable suggestion

❗ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation.

Suggested change

	def test_check_is_query_using_relevant_table_invalid_query(
	self, code_manager: CodeManager
	):
	mock_node = ast.parse("sql_query = 'SELECT * FROM your_table'").body[0]
	with pytest.raises(MaliciousQueryError) as excinfo:
	code_manager._check_is_query_using_relevant_table(mock_node)
	assert str(excinfo.value) == (
	"Query uses tables you don't provided, add new datatable or update query"
	)
	def test_check_is_query_using_relevant_table_invalid_query(
	self, code_manager: CodeManager
	):
	mock_node = ast.parse("sql_query = 'SELECT * FROM your_table'").body[0]
	with pytest.raises(MaliciousQueryError) as excinfo:
	code_manager._check_is_query_using_relevant_table(mock_node)
	assert str(excinfo.value) == (
	"Query uses tables you haven't provided. Add a new datatable or update the query."
	)

[coderabbitai]

The test test_check_is_query_using_relevant_table_valid_query seems to correctly handle a valid query by not raising an exception. However, it lacks an assertion to confirm that the method under test behaves as expected when a valid query is passed.

[coderabbitai]

The test test_check_is_query_using_relevant_table_unknown_table properly raises a MaliciousQueryError when a query with an unknown table is passed. The error message should be corrected for grammatical accuracy as mentioned earlier.

[coderabbitai]

Review Status

Actionable comments generated: 0

Configuration used: CodeRabbit UI

Commits

Files that changed from the base of the PR and between 0f807b7 and 4ed51bd.

Files selected for processing (1)

• pandasai/helpers/sql.py (1 hunks)

Additional comments: 1

pandasai/helpers/sql.py (1)

• 4-12: The extract_table_names function uses a regex pattern to extract table names from SQL queries. The pattern looks for words following FROM or JOIN clauses, which is a common way to reference tables in SQL. The function then flattens the list of matches and filters out any empty strings to return a list of table names. This implementation seems correct and efficient for its intended purpose.

[coderabbitai]

Review Status

Actionable comments generated: 0

Configuration used: CodeRabbit UI

Commits

Files that changed from the base of the PR and between 4ed51bd and 0e54ebe.

Files selected for processing (1)

• pandasai/helpers/code_manager.py (4 hunks)

Files skipped from review as they are similar to previous changes (1)

• pandasai/helpers/code_manager.py

[gventuri]

@ArslanSaleem could you fix the tests? I've just updated a error message based on CodeRabbit suggestion!

[mspronesti]

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.

[gventuri]

@mspronesti good catch, I agree!

[gventuri]

Also, since the logic for the "direct sql" is becoming more complex, we should probably move it all outside of the code manager! But we can probably create a separate PR for this!

[coderabbitai]

Review Status

Actionable comments generated: 0

Configuration used: CodeRabbit UI

Commits

Files that changed from the base of the PR and between 0e54ebe and 831c939.

Files selected for processing (2)

• pandasai/helpers/code_manager.py (4 hunks)
• tests/test_codemanager.py (4 hunks)

Files skipped from review as they are similar to previous changes (2)

• pandasai/helpers/code_manager.py
• tests/test_codemanager.py

[codecov-commenter]

Codecov Report

All modified and coverable lines are covered by tests ✅

Comparison is base (21819c6) 85.15% compared to head (831c939) 85.21%.
Report is 2 commits behind head on main.

❗ Your organization needs to install the Codecov GitHub app to enable full functionality.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #869      +/-   ##
==========================================
+ Coverage   85.15%   85.21%   +0.05%     
==========================================
  Files          89       90       +1     
  Lines        3866     3881      +15     
==========================================
+ Hits         3292     3307      +15     
  Misses        574      574              

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[ArslanSaleem]

Also, since the logic for the "direct sql" is becoming more complex, we should probably move it all outside of the code manager! But we can probably create a separate PR for this!

@gventuri Yes, absolutely and more scenarios will come