fix(pandasai-sql): dropping support for sqlite (#1558)

* fix(pandasai-sql): dropping support for sqlite

* fix(): standardize name in schema

* fix(): dropping name and description from dataframe

* fix: print leftover

* refactor(Loader): refactoring loader

* refactor(Loader): avoid creating temporary pandas dataframe

---------

Co-authored-by: scaliseraoul-sinaptik <raoul@sinaptik.ai>
Co-authored-by: Gabriele Venturi <lele.venturi@gmail.com>

Author: 3 people

extensions/connectors/sql/pandasai_sql/__init__.py

@@ -2,10 +2,7 @@
 
 import pandas as pd
 
-from pandasai.data_loader.semantic_layer_schema import (
-    SQLConnectionConfig,
-    SqliteConnectionConfig,
-)
+from pandasai.data_loader.semantic_layer_schema import SQLConnectionConfig
 
 
 def load_from_mysql(
@@ -38,15 +35,6 @@ def load_from_postgres(
     return pd.read_sql(query, conn, params=params)
 
 
-def load_from_sqlite(
-    connection_info: SqliteConnectionConfig, query: str, params: Optional[list] = None
-):
-    import sqlite3
-
-    conn = sqlite3.connect(connection_info.file_path)
-    return pd.read_sql(query, conn, params=params)
-
-
 def load_from_cockroachdb(
     connection_info: SQLConnectionConfig, query: str, params: Optional[list] = None
 ):

extensions/connectors/sql/poetry.lock

@@ -8,13 +8,9 @@
     load_from_cockroachdb,
     load_from_mysql,
     load_from_postgres,
-    load_from_sqlite,
 )
 
-from pandasai.data_loader.semantic_layer_schema import (
-    SQLConnectionConfig,
-    SqliteConnectionConfig,
-)
+from pandasai.data_loader.semantic_layer_schema import SQLConnectionConfig
 
 
 class TestDatabaseLoader(unittest.TestCase):
@@ -95,32 +91,6 @@ def test_load_from_postgres(self, mock_read_sql, mock_psycopg2_connect):
         self.assertIsInstance(result, pd.DataFrame)
         self.assertEqual(result.shape, (2, 2))
 
-    @patch("sqlite3.connect")
-    @patch("pandas.read_sql")
-    def test_load_from_sqlite(self, mock_read_sql, mock_sqlite3_connect):
-        # Setup the mock return values
-        mock_conn = MagicMock()
-        mock_sqlite3_connect.return_value = mock_conn
-        mock_read_sql.return_value = pd.DataFrame(
-            {"column1": [9, 10], "column2": [11, 12]}
-        )
-
-        # Test data
-        connection_info = {"file_path": "test_db.sqlite"}
-        query = "SELECT * FROM test_table"
-
-        connection_config = SqliteConnectionConfig(**connection_info)
-
-        result = load_from_sqlite(connection_config, query)
-
-        # Assert that the connection is made and SQL query is executed
-        mock_sqlite3_connect.assert_called_once_with("test_db.sqlite")
-        mock_read_sql.assert_called_once_with(query, mock_conn, params=None)
-
-        # Assert the result is a DataFrame
-        self.assertIsInstance(result, pd.DataFrame)
-        self.assertEqual(result.shape, (2, 2))
-
     @patch("psycopg2.connect")
     @patch("pandas.read_sql")
     def test_load_from_cockroachdb(self, mock_read_sql, mock_postgresql_connect):
@@ -229,29 +199,6 @@ def test_load_from_postgres_with_params(self, mock_read_sql, mock_psycopg2_conne
         self.assertIsInstance(result, pd.DataFrame)
         self.assertEqual(result.shape, (2, 2))
 
-    @patch("sqlite3.connect")
-    @patch("pandas.read_sql")
-    def test_load_from_sqlite_with_params(self, mock_read_sql, mock_sqlite3_connect):
-        mock_conn = MagicMock()
-        mock_sqlite3_connect.return_value = mock_conn
-        mock_read_sql.return_value = pd.DataFrame(
-            {"column1": [9, 10], "column2": [11, 12]}
-        )
-
-        connection_info = {"file_path": "test_db.sqlite"}
-        query = "SELECT * FROM test_table WHERE age > ?"
-        query_params = [30]
-
-        connection_config = SqliteConnectionConfig(**connection_info)
-
-        result = load_from_sqlite(connection_config, query, query_params)
-
-        mock_sqlite3_connect.assert_called_once_with("test_db.sqlite")
-        mock_read_sql.assert_called_once_with(query, mock_conn, params=query_params)
-
-        self.assertIsInstance(result, pd.DataFrame)
-        self.assertEqual(result.shape, (2, 2))
-
     @patch("psycopg2.connect")
     @patch("pandas.read_sql")
     def test_load_from_cockroachdb_with_params(

extensions/connectors/sql/tests/test_sql.py

@@ -4,7 +4,6 @@
 """
 
 import os
-import re
 from io import BytesIO
 from typing import List, Optional, Union
 from zipfile import ZipFile
@@ -27,6 +26,7 @@
 from .constants import LOCAL_SOURCE_TYPES, SQL_SOURCE_TYPES
 from .core.cache import Cache
 from .data_loader.loader import DatasetLoader
+from .data_loader.query_builder import QueryBuilder
 from .data_loader.semantic_layer_schema import (
     Column,
 )
@@ -39,11 +39,11 @@
 def create(
     path: str,
     df: Optional[DataFrame] = None,
-    name: Optional[str] = None,
     description: Optional[str] = None,
     columns: Optional[List[dict]] = None,
     source: Optional[dict] = None,
     relations: Optional[List[dict]] = None,
+    view: bool = False,
 ) -> Union[DataFrame, VirtualDataFrame]:
     """
     Creates a new dataset at the specified path with optional metadata, schema,
@@ -85,7 +85,6 @@ def create(
         >>> create(
         ...     path="my-org/my-dataset",
         ...     df=my_dataframe,
-        ...     name="My Dataset",
         ...     description="This is a sample dataset.",
         ...     columns=[
         ...         {"name": "id", "type": "integer", "description": "Primary key"},
@@ -103,54 +102,46 @@ def create(
         find_project_root(), "datasets", org_name, dataset_name
     )
 
+    schema_path = os.path.join(str(dataset_directory), "schema.yaml")
+    parquet_file_path = os.path.join(str(dataset_directory), "data.parquet")
     # Check if dataset already exists
-    if os.path.exists(dataset_directory):
-        schema_path = os.path.join(dataset_directory, "schema.yaml")
-        if os.path.exists(schema_path):
-            raise ValueError(f"Dataset already exists at path: {path}")
+    if os.path.exists(dataset_directory) and os.path.exists(schema_path):
+        raise ValueError(f"Dataset already exists at path: {path}")
 
     os.makedirs(dataset_directory, exist_ok=True)
 
-    # Save schema to yaml
-    schema_path = os.path.join(dataset_directory, "schema.yaml")
-
-    if df is None and source is None:
-        raise InvalidConfigError("Please provide either a DataFrame or a source")
+    if df is None and source is None and not view:
+        raise InvalidConfigError(
+            "Please provide either a DataFrame, a Source or a View"
+        )
 
     if df is not None:
         schema = df.schema
-        df.to_parquet(os.path.join(dataset_directory, "data.parquet"), index=False)
-    elif source.get("type") == "sqlite" and source.get("table"):
-        schema = SemanticLayerSchema(name=source.get("table"), source=Source(**source))
-        df = _dataset_loader.load(schema=schema)
-        df.to_parquet(os.path.join(dataset_directory, "data.parquet"), index=False)
-    elif source.get("table"):
-        schema = SemanticLayerSchema(name=source.get("table"), source=Source(**source))
-        df = _dataset_loader.load(schema=schema)
-    elif source.get("view"):
-        name = name or dataset_name
+        schema.name = sanitize_sql_table_name(dataset_name)
+        df.to_parquet(parquet_file_path, index=False)
+    elif view:
         _relation = [Relation(**relation) for relation in relations or ()]
-        schema = SemanticLayerSchema(
-            name=name, source=Source(**source), relations=_relation
+        schema: SemanticLayerSchema = SemanticLayerSchema(
+            name=sanitize_sql_table_name(dataset_name), relations=_relation, view=True
         )
-        df = _dataset_loader.load(schema=schema)
+    elif source.get("table"):
+        schema: SemanticLayerSchema = SemanticLayerSchema(
+            name=sanitize_sql_table_name(dataset_name), source=Source(**source)
+        )
+    else:
+        raise InvalidConfigError("Unable to create schema with the provided params")
 
-    schema.name = sanitize_sql_table_name(name or schema.name)
     schema.description = description or schema.description
     if columns:
         schema.columns = [Column(**column) for column in columns]
-    elif df is not None:
-        schema.columns = [
-            Column(name=str(name), type=DataFrame.get_column_type(dtype))
-            for name, dtype in df.dtypes.items()
-        ]
 
     with open(schema_path, "w") as yml_file:
         yml_file.write(schema.to_yaml())
 
     print(f"Dataset saved successfully to path: {dataset_directory}")
 
-    return _dataset_loader.load(path)
+    loader = DatasetLoader.create_loader_from_schema(schema, path)
+    return loader.load()
 
 
 # Global variable to store the current agent
@@ -206,9 +197,6 @@ def follow_up(query: str):
     return _current_agent.follow_up(query)
 
 
-_dataset_loader = DatasetLoader()
-
-
 def load(dataset_path: str) -> DataFrame:
     """
     Load data based on the provided dataset path.
@@ -223,7 +211,6 @@ def load(dataset_path: str) -> DataFrame:
     if len(path_parts) != 2:
         raise ValueError("The path must be in the format 'organization/dataset'.")
 
-    global _dataset_loader
     dataset_full_path = os.path.join(find_project_root(), "datasets", dataset_path)
     if not os.path.exists(dataset_full_path):
         api_key = os.environ.get("PANDABI_API_KEY", None)
@@ -244,13 +231,14 @@ def load(dataset_path: str) -> DataFrame:
         with ZipFile(BytesIO(file_data.content)) as zip_file:
             zip_file.extractall(dataset_full_path)
 
-    return _dataset_loader.load(dataset_path)
+    loader = DatasetLoader.create_loader_from_path(dataset_path)
+    return loader.load()
 
 
 def read_csv(filepath: str) -> DataFrame:
     data = pd.read_csv(filepath)
-    name = f"table_{sanitize_sql_table_name(filepath)}"
-    return DataFrame(data, name=name)
+    table = f"table_{sanitize_sql_table_name(filepath)}"
+    return DataFrame(data, _table_name=table)
 
 
 __all__ = [

pandasai/__init__.py

@@ -123,7 +123,7 @@ def _execute_local_sql_query(self, query: str) -> pd.DataFrame:
             with duckdb.connect() as con:
                 # Register all DataFrames in the state
                 for df in self._state.dfs:
-                    con.register(df.name, df)
+                    con.register(df.schema.source.table, df)
 
                 # Execute the query and fetch the result as a pandas DataFrame
                 result = con.sql(query).df()

pandasai/agent/base.py

@@ -29,29 +29,26 @@
     "mysql": "pandasai_sql",
     "postgres": "pandasai_sql",
     "cockroachdb": "pandasai_sql",
-    "sqlite": "pandasai_sql",
     "yahoo_finance": "pandasai_yfinance",
     "bigquery": "pandasai_bigquery",
     "snowflake": "pandasai_snowflake",
     "databricks": "pandasai_databricks",
     "oracle": "pandasai_oracle",
 }
 
-LOCAL_SOURCE_TYPES = ["csv", "parquet", "sqlite"]
+LOCAL_SOURCE_TYPES = ["csv", "parquet"]
 REMOTE_SOURCE_TYPES = [
     "mysql",
     "postgres",
     "cockroachdb",
-    "sqlite",
     "data",
     "yahoo_finance",
     "bigquery",
     "snowflake",
     "databricks",
     "oracle",
 ]
-SQL_SOURCE_TYPES = ["mysql", "postgres", "cockroachdb", "oracle", "sqlite"]
-
+SQL_SOURCE_TYPES = ["mysql", "postgres", "cockroachdb", "oracle"]
 VALID_COLUMN_TYPES = ["string", "integer", "float", "datetime", "boolean"]
 
 VALID_TRANSFORMATION_TYPES = [

pandasai/constants.py

@@ -29,7 +29,7 @@ def _check_direct_sql_func_def_exists(self, node: ast.AST) -> bool:
         return isinstance(node, ast.FunctionDef) and node.name == "execute_sql_query"
 
     def _replace_table_names(
-        self, sql_query: str, table_names: list, allowed_table_names: list
+        self, sql_query: str, table_names: list, allowed_table_names: dict
     ) -> str:
         """
         Replace table names in the SQL query with case-sensitive or authorized table names.
@@ -54,8 +54,11 @@ def _clean_sql_query(self, sql_query: str) -> str:
         """
         sql_query = sql_query.rstrip(";")
         table_names = extract_table_names(sql_query)
-        allowed_table_names = {df.name: df.name for df in self.context.dfs} | {
-            f'"{df.name}"': df.name for df in self.context.dfs
+        allowed_table_names = {
+            df.schema.source.table: df.schema.source.table for df in self.context.dfs
+        } | {
+            f'"{df.schema.source.table}"': df.schema.source.table
+            for df in self.context.dfs
         }
         return self._replace_table_names(sql_query, table_names, allowed_table_names)