fix(sql): dialect sql parser fixed (#1778)

Author: ArslanSaleem

pandasai/core/code_generation/code_cleaning.py

@@ -55,7 +55,8 @@ def _clean_sql_query(self, sql_query: str) -> str:
         Clean the SQL query by trimming semicolons and validating table names.
         """
         sql_query = sql_query.rstrip(";")
-        table_names = SQLParser.extract_table_names(sql_query)
+        dialect = self.context.dfs[0].get_dialect()
+        table_names = SQLParser.extract_table_names(sql_query, dialect)
         allowed_table_names = {
             df.schema.name: df.schema.name for df in self.context.dfs
         } | {f'"{df.schema.name}"': df.schema.name for df in self.context.dfs}

pandasai/dataframe/base.py

@@ -132,20 +132,23 @@ def rows_count(self) -> int:
     def columns_count(self) -> int:
         return len(self.columns)
 
+    def get_dialect(self):
+        source = self.schema.source or None
+        if source:
+            dialect = "duckdb" if source.type in LOCAL_SOURCE_TYPES else source.type
+        else:
+            dialect = "postgres"
+
+        return dialect
+
     def serialize_dataframe(self) -> str:
         """
         Serialize DataFrame to string representation.
 
         Returns:
             str: Serialized string representation of the DataFrame
         """
-        source = self.schema.source or None
-
-        if source:
-            dialect = "duckdb" if source.type in LOCAL_SOURCE_TYPES else source.type
-        else:
-            dialect = "postgres"
-
+        dialect = self.get_dialect()
         return DataframeSerializer.serialize(self, dialect)
 
     def get_head(self):

tests/unit_tests/core/code_generation/test_code_cleaning.py

@@ -71,6 +71,7 @@ def test_clean_sql_query(self):
         mock_dataframe.schema = MagicMock()
         mock_dataframe.schema.name = "my_table"
         self.cleaner.context.dfs = [mock_dataframe]
+        mock_dataframe.get_dialect = MagicMock(return_value="duckdb")
         result = self.cleaner._clean_sql_query(sql_query)
         self.assertEqual(result, "SELECT * FROM my_table")
 
@@ -84,6 +85,7 @@ def test_validate_and_make_table_name_case_sensitive(self):
         self.cleaner.context.dfs = [mock_dataframe]
         mock_dataframe.schema = MagicMock()
         mock_dataframe.schema.name = "my_table"
+        mock_dataframe.get_dialect = MagicMock(return_value="duckdb")
         updated_node = self.cleaner._validate_and_make_table_name_case_sensitive(node)
         self.assertEqual(updated_node.value.value, "SELECT * FROM my_table")