Support with Session() as sess: in chdb session

Author: auxten

chdb/session/state.py

@@ -1,18 +1,18 @@
 import tempfile
 import shutil
 
-from chdb import query
+from chdb import query, g_udf_path
 
 
-class Session():
+class Session:
     """
     Session will keep the state of query. All DDL and DML state will be kept in a dir.
     Dir path could be passed in as an argument. If not, a temporary dir will be created.
 
     If path is not specified, the temporary dir will be deleted when the Session object is deleted.
     Otherwise path will be kept.
-    
-    Note: The default database is "_local" and the default engine is "Memory" which means all data 
+
+    Note: The default database is "_local" and the default engine is "Memory" which means all data
     will be stored in memory. If you want to store data in disk, you should create another database.
     """
 
@@ -28,8 +28,17 @@ def __del__(self):
         if self._cleanup:
             self.cleanup()
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.cleanup()
+
     def cleanup(self):
-        shutil.rmtree(self._path)
+        try:
+            shutil.rmtree(self._path)
+        except:
+            pass
 
     def query(self, sql, fmt="CSV"):
         """