fix: Use original working dir for all paths except file uploads (#356)

* fix(file-management): ensure absolute path for notebook and requirements files are used in the entire application without getting affected by directory change

* fix(cli): Updated cli to remove unwanted code and optimized the flow to set path

* fix(cli): Removed unwanted spaces

* fix: Use original working dir for all paths except file uploads

---------

Co-authored-by: priya_kanabar-crest <priya.k@crestinfosystems.com>

.github/workflows/docker_build.yml

@@ -76,6 +76,6 @@ jobs:
         --dockerfile Dockerfile \
         --reproducible --cleanup --cache=true \
         --cache-repo ${{ secrets.ECR_REGISTRY }}/${{ secrets.ECR_REPOSITORY }}-cache \
-        --destination=${{ secrets.ECR_REGISTRY }}/${{ secrets.ECR_REPOSITORY }}:main-${{ github.sha }}-dev-python${{ matrix.python-version }} \
+        --destination=${{ secrets.ECR_REGISTRY }}/${{ secrets.ECR_REPOSITORY }}:zero-true-v${{ github.sha }}-python${{ matrix.python-version }} \
         --context dir:///workspace/ \
         --build-arg PYTHON_VERSION=${{ matrix.python-version }}

zt_backend/config.py

@@ -11,6 +11,7 @@ class Settings(BaseSettings):
     publish_url: str = "https://bxmm0wp9zk.execute-api.us-east-2.amazonaws.com/default/"
     comments_enabled: bool = False
     show_create_button: bool = False
+    zt_path: str = ""
 
 @lru_cache()
 def get_settings():

zt_backend/main.py

@@ -6,11 +6,11 @@
 from zt_backend.utils.dependencies import parse_dependencies, write_dependencies
 from copilot.copilot import copilot_app
 import zt_backend.router as router
+from pathlib import Path
 import os
 import webbrowser
 import logging
 import traceback
-import pkg_resources
 import matplotlib
 
 app = FastAPI()
@@ -40,15 +40,14 @@
 def open_project():
     try:
         matplotlib.use("Agg")
-        if not os.path.exists("notebook.ztnb"):
+        notebook_path = Path(settings.zt_path) / "notebook.ztnb"
+        if not notebook_path.exists():
             logger.info("No notebook file found, creating with empty notebook")
             write_notebook()
-        if not os.path.exists("requirements.txt"):
-            logger.info("No requirements file found, creating with base dependency")
-            with open("requirements.txt", "w") as file:
-                file.write(
-                    f"zero-true=={pkg_resources.get_distribution('zero-true').version}"
-                )
+        requirements_path = Path(settings.zt_path) / "requirements.txt"
+        if not requirements_path.exists():
+            logger.info("No requirements file found, creating empty file")
+            requirements_path.touch()
         else:
             write_dependencies(parse_dependencies())
         get_notebook()

zt_backend/router.py

@@ -484,9 +484,7 @@ def share_notebook(shareRequest: request.ShareRequest):
                 warning_message += python_warning
                 if zt_warning:
                     warning_message += "\n" + zt_warning
-                warning_message += (
-                    "\nWe recommend upgrading your versions before continuing. If you would like to continue, select confirm."
-                )
+                warning_message += "\nWe recommend upgrading your versions before continuing. If you would like to continue, select confirm."
                 upload_state.signed_url = signed_url
                 return {"warning": warning_message}
             if zt_warning:
@@ -514,7 +512,9 @@ def confirm_share(shareRequest: request.ShareRequest):
     if app_state.run_mode == "dev":
         if upload_state.signed_url:
             try:
-                publish_files(shareRequest.projectName.lower().strip(), upload_state.signed_url)
+                publish_files(
+                    shareRequest.projectName.lower().strip(), upload_state.signed_url
+                )
                 upload_state.signed_url = None
             except Exception as e:
                 raise HTTPException(
@@ -528,25 +528,44 @@ def confirm_share(shareRequest: request.ShareRequest):
 
 
 def publish_files(project_name, signed_url):
-    output_filename = f"{project_name}"
-    project_source = os.path.normpath(os.getcwd())
-    logger.info(project_source)
-    shutil.make_archive(
-        base_name=output_filename, format="gztar", root_dir=project_source
-    )
-
-    upload_files = {"file": open(f"{output_filename}.tar.gz", "rb")}
-    upload_response = requests.post(
-        signed_url["url"], data=signed_url["fields"], files=upload_files
-    )
-    if upload_response.status_code != 204:
-        return {
-            "Error": upload_response.json().get(
-                "message",
-                upload_response.json().get("Message", "Failed to upload files"),
+    try:
+        output_filename = Path(settings.zt_path) / f"{project_name}.tar.gz"
+        tar_base = str(Path(settings.zt_path) / project_name)
+
+        shutil.make_archive(
+            base_name=tar_base, format="gztar", root_dir=settings.zt_path
+        )
+
+        with output_filename.open("rb") as file:
+            upload_files = {"file": file}
+            upload_response = requests.post(
+                signed_url["url"], data=signed_url["fields"], files=upload_files
             )
-        }
-    os.remove(f"{output_filename}.tar.gz")
+
+        if upload_response.status_code != 204:
+            return {
+                "Error": upload_response.json().get(
+                    "message",
+                    upload_response.json().get("Message", "Failed to upload files"),
+                )
+            }
+
+        try:
+            output_filename.unlink()
+        except OSError as e:
+            logger.warning(f"Failed to remove temporary file: {e}")
+            try:
+                output_filename.unlink()
+            except OSError as e:
+                logger.error(f"Failed to remove temporary file after retry: {e}")
+
+    except Exception as e:
+        try:
+            if output_filename.exists():
+                output_filename.unlink()
+        except OSError:
+            pass
+        raise e
 
 
 @router.post("/api/upload_file")

zt_backend/utils/dependencies.py

@@ -6,13 +6,16 @@
 import traceback
 import pkg_resources
 import re
+from pathlib import Path
+from zt_backend.config import settings
 
 logger = logging.getLogger("__name__")
 
 
 def parse_dependencies():
     dependencies = []
-    with open("requirements.txt", "r", encoding="utf-8") as file:
+    requirements_path = Path(settings.zt_path) / "requirements.txt"
+    with requirements_path.open("r", encoding="utf-8") as file:
         for line in file:
             line = re.sub(r"#.*", "", line).strip()
             if not line:
@@ -42,23 +45,20 @@ def check_env(dependencies: notebook.Dependencies):
 
 
 def write_dependencies(dependencies: notebook.Dependencies):
-    with open("requirements.txt", "w", encoding="utf-8") as file:
-        file.seek(0)
-        file.write(
-            f"zero-true=={pkg_resources.get_distribution('zero-true').version}\n"
-        )
+    requirements_path = Path(settings.zt_path) / "requirements.txt"
+    with requirements_path.open("w", encoding="utf-8") as file:
         for dependency in dependencies.dependencies:
             if dependency.package:
                 file.write(f"{dependency.package}{dependency.version}\n")
-        file.truncate()
 
 
 async def dependency_update(
     dependency_request: request.DependencyRequest, websocket: WebSocket
 ):
     try:
         write_dependencies(dependency_request.dependencies)
-        command = ["pip", "install", "-r", "requirements.txt"]
+        requirements_path = Path(settings.zt_path) / "requirements.txt"
+        command = ["pip", "install", "-r", str(requirements_path)]
         with subprocess.Popen(
             command, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True
         ) as process:

zt_backend/utils/notebook.py

@@ -5,6 +5,7 @@
 from zt_backend.models import notebook
 from zt_backend.utils.debounce import debounce
 from zt_backend.config import settings
+from pathlib import Path
 import logging
 import duckdb
 import uuid
@@ -35,7 +36,9 @@ def get_notebook(id=""):
     try:
         logger.debug("Notebook id is empty")
         # If it doesn't exist in the database, load it from the TOML file
-        with open("notebook.ztnb", "r", encoding="utf-8") as project_file:
+        logger.info(f'{settings.zt_path}/notebook.ztnb')
+        notebook_path = Path(settings.zt_path) / "notebook.ztnb"
+        with notebook_path.open("r", encoding="utf-8") as project_file:
             toml_data = rtoml.loads(project_file.read().replace("\\", "\\\\"))
 
         try:
@@ -315,10 +318,11 @@ def save_notebook():
 
 
 def write_notebook():
-    tmp_uuid_file = f"notebook_{uuid.uuid4()}.ztnb"
+    tmp_uuid_file = Path(settings.zt_path) / f"notebook_{uuid.uuid4()}.ztnb"
+    notebook_path = Path(settings.zt_path) / "notebook.ztnb"
     logger.debug("Saving toml for notebook %s", notebook_state.zt_notebook.notebookId)
     try:
-        with open(tmp_uuid_file, "w", encoding="utf-8") as project_file:
+        with tmp_uuid_file.open("w", encoding="utf-8") as project_file:
             # Write notebookId
             project_file.write(
                 f'notebookId = "{notebook_state.zt_notebook.notebookId}"\nnotebookName = "{notebook_state.zt_notebook.notebookName}"\n\n'
@@ -358,7 +362,7 @@ def write_notebook():
 
                 project_file.write("\n")
 
-        os.replace(tmp_uuid_file, "notebook.ztnb")
+        tmp_uuid_file.replace(notebook_path)
     except Exception as e:
         logger.error(
             "Error saving notebook %s toml file: %s",
@@ -368,7 +372,7 @@ def write_notebook():
 
     finally:
         try:
-            os.remove(tmp_uuid_file)
+            tmp_uuid_file.unlink()
         except Exception as e:
             logger.debug(
                 "Error while deleting temporary toml file for notebook %s: %s",

zt_cli/cli.py

@@ -6,6 +6,7 @@
 import uvicorn
 from typing import Optional
 from rich import print
+from pathlib import Path
 import shutil
 import requests
 import pkg_resources
@@ -173,7 +174,7 @@ def app(
     """
     Start the Zero-True application.
     """
-
+    os.environ["ZT_PATH"] = str(Path.cwd())
     print_ascii_logo()
     os.environ["RUN_MODE"] = "app"
     log_path = os.path.normpath(
@@ -205,7 +206,7 @@ def notebook(
     """
     Start the Zero-True application.
     """
-
+    os.environ["ZT_PATH"] = str(Path.cwd())
     print_ascii_logo()
     os.environ["RUN_MODE"] = "dev"
     log_path = os.path.normpath(

zt_dev_cli/cli.py

@@ -4,6 +4,7 @@
 import os
 import shutil
 import typer
+from pathlib import Path
 from zt_backend.models.generate_schema import generate_schema
 from typing_extensions import Annotated
 from typing import Optional
@@ -59,7 +60,7 @@ def build():
 def app(
     port: Annotated[Optional[int], typer.Option(help="Port number to bind to.")] = 5173
 ):
-
+    os.environ["ZT_PATH"] = str(Path.cwd())
     print_ascii_logo()
 
     log_path = os.path.normpath(
@@ -90,7 +91,7 @@ def app(
 def notebook(
     port: Annotated[Optional[int], typer.Option(help="Port number to bind to.")] = 5173
 ):
-
+    os.environ["ZT_PATH"] = str(Path.cwd())
     print_ascii_logo()
 
     log_path = os.path.normpath(
@@ -118,4 +119,4 @@ def notebook(
 
 
 if __name__ == "__main__":
-    cli_app()
+    cli_app()