fix(duckdbConnection): make duckdb connection manager non singleton (#1710)

* fix(duckdbConnection): shared connection space between duckdb

* fix(dataset): add test case

Author: ArslanSaleem

pandasai/data_loader/duck_db_connection_manager.py

@@ -1,4 +1,3 @@
-import weakref
 from typing import Optional
 
 import duckdb
@@ -7,37 +6,34 @@
 
 
 class DuckDBConnectionManager:
-    _instance = None
-
-    def __new__(cls):
-        if cls._instance is None:
-            cls._instance = super(DuckDBConnectionManager, cls).__new__(cls)
-            cls._instance._init_connection()
-            weakref.finalize(cls._instance, cls._close_connection)
-        return cls._instance
-
-    def _init_connection(self):
+    def __init__(self):
         """Initialize a DuckDB connection."""
         self.connection = duckdb.connect()
         self._registered_tables = set()
 
-    @classmethod
-    def _close_connection(cls):
-        """Closes the DuckDB connection when the instance is deleted."""
-        if cls._instance and hasattr(cls._instance, "connection"):
-            cls._instance.connection.close()
-            cls._instance = None
+    def __del__(self):
+        """Destructor to ensure the DuckDB connection is closed."""
+        self.close()
 
     def register(self, name: str, df):
         """Registers a DataFrame as a DuckDB table."""
         self.connection.register(name, df)
         self._registered_tables.add(name)
 
+    def unregister(self, name: str):
+        """Unregister a previously registered DuckDB table."""
+        if name in self._registered_tables:
+            self.connection.unregister(name)
+            self._registered_tables.remove(name)
+
     def sql(self, query: str, params: Optional[list] = None):
         """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, params=params)
 
     def close(self):
-        """Manually close the connection if needed."""
-        self._close_connection()
+        """Closes the DuckDB connection."""
+        if hasattr(self, "connection") and self.connection:
+            self.connection.close()
+            self.connection = None
+            self._registered_tables.clear()

tests/unit_tests/data_loader/test_duckdbmanager.py

@@ -10,3 +10,12 @@ def duck_db_manager(self):
 
     def test_connection_correct_closing_doesnt_throw(self, duck_db_manager):
         duck_db_manager.close()
+
+    def test_unregister(self, duck_db_manager, sample_df):
+        duck_db_manager.register("test", sample_df)
+
+        assert "test" in duck_db_manager._registered_tables
+
+        duck_db_manager.unregister("test")
+
+        assert len(duck_db_manager._registered_tables) == 0