update version

Author: Chenglong-MS

.env.template

@@ -5,6 +5,8 @@
 DISABLE_DISPLAY_KEYS=false # if true, the display keys will not be shown in the frontend
 EXEC_PYTHON_IN_SUBPROCESS=false # if true, the python code will be executed in a subprocess to avoid crashing the main app, but it will increase the time of response
 
+LOCAL_DB_DIR= # the directory to store the local database, if not provided, the app will use the temp directory
+
 # External atabase connection settings
 # check https://duckdb.org/docs/stable/extensions/mysql.html 
 # and https://duckdb.org/docs/stable/extensions/postgres.html

py-src/data_formulator/db_manager.py

@@ -9,13 +9,14 @@
 from dotenv import load_dotenv
 
 class DuckDBManager:
-    def __init__(self, external_db_connections: Dict[str, Dict[str, Any]]):
+    def __init__(self, external_db_connections: Dict[str, Dict[str, Any]], local_db_dir: str):
         # Store session db file paths
         self._db_files: Dict[str, str] = {}
 
         # External db connections and tracking of installed extensions
         self._external_db_connections: Dict[str, Dict[str, Any]] = external_db_connections
         self._installed_extensions: Dict[str, List[str]] = {}
+        self._local_db_dir: str = local_db_dir
 
     @contextmanager
     def connection(self, session_id: str) -> ContextManager[duckdb.DuckDBPyConnection]:
@@ -33,7 +34,10 @@ def get_connection(self, session_id: str) -> duckdb.DuckDBPyConnection:
         """Internal method to get or create a DuckDB connection for a session"""
         # Get or create the db file path for this session
         if session_id not in self._db_files or self._db_files[session_id] is None:
-            db_file = os.path.join(tempfile.gettempdir(), f"df_{session_id}.duckdb")
+            db_dir = self._local_db_dir if self._local_db_dir else tempfile.gettempdir()
+            if not os.path.exists(db_dir):
+                db_dir = tempfile.gettempdir()
+            db_file = os.path.join(db_dir, f"df_{session_id}.duckdb")
             print(f"=== Creating new db file: {db_file}")
             self._db_files[session_id] = db_file
             # Initialize extension tracking for this file
@@ -73,12 +77,15 @@ def get_connection(self, session_id: str) -> duckdb.DuckDBPyConnection:
 env = load_dotenv()
 
 # Initialize the DB manager
-db_manager = DuckDBManager({
-    "db_name": os.getenv('DB_NAME'),
-    "db_type": os.getenv('DB_TYPE'),
-    "host": os.getenv('DB_HOST'),
-    "port": os.getenv('DB_PORT'),
-    "database": os.getenv('DB_DATABASE'),
-    "user": os.getenv('DB_USER'),
-    "password": os.getenv('DB_PASSWORD')
-} if os.getenv('USE_EXTERNAL_DB') == 'true' else None)
+db_manager = DuckDBManager(
+    external_db_connections={
+        "db_name": os.getenv('DB_NAME'),
+        "db_type": os.getenv('DB_TYPE'),
+        "host": os.getenv('DB_HOST'),
+        "port": os.getenv('DB_PORT'),
+        "database": os.getenv('DB_DATABASE'),
+        "user": os.getenv('DB_USER'),
+        "password": os.getenv('DB_PASSWORD')
+    } if os.getenv('USE_EXTERNAL_DB') == 'true' else None,
+    local_db_dir=os.getenv('LOCAL_DB_DIR')
+)

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "data_formulator"
-version = "0.1.7"
+version = "0.2"
 
 requires-python = ">=3.9"
 authors = [
@@ -34,7 +34,8 @@ dependencies = [
     "azure-keyvault-secrets",  
     "python-dotenv",  
     "vega_datasets",
-    "litellm"
+    "litellm",
+    "duckdb"
 ]
 
 [project.urls]