fix: transpiling sql to specific dialect (#1596)

* fix: transpiling sql to specific dialect

* fix: remove default mysql dialect

pandasai/data_loader/duck_db_connection_manager.py

@@ -2,6 +2,8 @@
 
 import duckdb
 
+from pandasai.query_builders.sql_parser import SQLParser
+
 
 class DuckDBConnectionManager:
     _instance = None
@@ -32,6 +34,7 @@ def register(self, name: str, df):
 
     def sql(self, query: str):
         """Executes an SQL query and returns the result as a Pandas DataFrame."""
+        query = SQLParser.transpile_sql_dialect(query, to_dialect="duckdb")
         return self.connection.sql(query)
 
     def close(self):

pandasai/data_loader/sql_loader.py

@@ -11,6 +11,7 @@
 from ..constants import (
     SUPPORTED_SOURCE_CONNECTORS,
 )
+from ..query_builders.sql_parser import SQLParser
 from .loader import DatasetLoader
 from .semantic_layer_schema import SemanticLayerSchema
 
@@ -40,6 +41,7 @@ def execute_query(self, query: str, params: Optional[list] = None) -> pd.DataFra
         connection_info = self.schema.source.connection
 
         load_function = self._get_loader_function(source_type)
+        query = SQLParser.transpile_sql_dialect(query, to_dialect=source_type)
 
         if not is_sql_query_safe(query):
             raise MaliciousQueryError(

pandasai/data_loader/view_loader.py

@@ -10,6 +10,7 @@
 from ..exceptions import MaliciousQueryError
 from ..helpers.sql_sanitizer import is_sql_query_safe
 from ..query_builders.base_query_builder import BaseQueryBuilder
+from ..query_builders.sql_parser import SQLParser
 from .duck_db_connection_manager import DuckDBConnectionManager
 from .loader import DatasetLoader
 from .local_loader import LocalDatasetLoader
@@ -89,6 +90,7 @@ def execute_query(self, query: str, params: Optional[list] = None) -> pd.DataFra
         if source_type in LOCAL_SOURCE_TYPES:
             return self.execute_local_query(query)
         load_function = self._get_loader_function(source_type)
+        query = SQLParser.transpile_sql_dialect(query, to_dialect=source_type)
 
         if not is_sql_query_safe(query):
             raise MaliciousQueryError(

pandasai/query_builders/sql_parser.py

@@ -4,13 +4,6 @@
 
 
 class SQLParser:
-    @staticmethod
-    def extract_table_names(query: str):
-        parsed = sqlglot.parse_one(query)
-        return {
-            table.alias_or_name for table in parsed.find_all(sqlglot.expressions.Table)
-        }
-
     @staticmethod
     def replace_table_and_column_names(query, table_mapping):
         """
@@ -57,3 +50,10 @@ def transform_node(node):
 
         # Convert back to SQL string
         return transformed.sql(pretty=True)
+
+    @staticmethod
+    def transpile_sql_dialect(query, to_dialect, from_dialect=None):
+        query = (
+            parse_one(query, read=from_dialect) if from_dialect else parse_one(query)
+        )
+        return query.sql(dialect=to_dialect, pretty=True)

tests/unit_tests/data_loader/test_sql_loader.py

@@ -153,7 +153,7 @@ def test_mysql_safe_query(self, mysql_schema):
             result = loader.execute_query("SELECT * FROM users")
 
             assert isinstance(result, DataFrame)
-            mock_sql_query.assert_called_once_with("SELECT * FROM users")
+            mock_sql_query.assert_called_once_with("SELECT\n  *\nFROM users")
 
     def test_mysql_malicious_with_no_import(self, mysql_schema):
         """Test loading data from a MySQL source creates a VirtualDataFrame and handles queries correctly."""

tests/unit_tests/query_builders/test_sql_parser.py

@@ -56,3 +56,9 @@ class TestSqlParser:
     def test_replace_table_names(query, table_mapping, expected):
         result = SQLParser.replace_table_and_column_names(query, table_mapping)
         assert result.strip() == expected.strip()
+
+    def test_mysql_transpilation(self):
+        query = '''SELECT COUNT(*) AS "total_rows"'''
+        expected = """SELECT\n  COUNT(*) AS `total_rows`"""
+        result = SQLParser.transpile_sql_dialect(query, to_dialect="mysql")
+        assert result.strip() == expected.strip()