Added run code endpoint.

.coveragerc

@@ -21,3 +21,4 @@ exclude_lines =
     # Don't complain if tests don't hit defensive assertion code:
     raise AssertionError
     raise NotImplementedError
+    if TYPE_CHECKING

.pre-commit-config.yaml

@@ -8,7 +8,7 @@ repos:
         args: ["--branch", "main"]
 
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.7.4
+    rev: v0.8.0
     hooks:
       - id: ruff
         name: Run the ruff linter

CHANGELOG.md

@@ -5,13 +5,23 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
-## [Unreleased]
+## [0.1.0-rc.2] - 2024-11-26
+
+### Added
+
+- Added endpoint to run python code.
 
 ### Changed
 
 - Improved `README.md` to include required security configuration options to use `gVisor` as the container runtime.
 - Changed folder where runs are stored to `/runs` instead of `/tmp`.
-- Changed run commands to extract changed files even if the command fails.
+- Changed `execute_command` to extract changed files even if the command fails.
+- Changed `execute_command` to allow conditionally extracting changed files.
+- Renamed `ForbiddenError` to `ErrorMessage` to be more generic.
+- Updated dependencies:
+  - `ruff` from 0.7.4 to 0.8.0
+  - `pydantic` from 2.10.0 to 2.10.2
+  - `sentry-sdk` from 2.18.0 to 2.19.0
 
 ### Removed
 

daiv_sandbox/languages.py

@@ -0,0 +1,57 @@
+from __future__ import annotations
+
+import abc
+import io
+import tarfile
+from typing import TYPE_CHECKING, Literal
+
+if TYPE_CHECKING:
+    from daiv_sandbox.schemas import RunResult
+    from daiv_sandbox.sessions import SandboxDockerSession
+
+
+class LanguageManager(abc.ABC):
+    """
+    Abstract base class for language managers.
+    """
+
+    @abc.abstractmethod
+    def install_dependencies(self, session: SandboxDockerSession, dependencies: list[str]) -> RunResult:
+        pass
+
+    @abc.abstractmethod
+    def run_code(self, session: SandboxDockerSession, workdir: str, code: str) -> RunResult:
+        pass
+
+    @staticmethod
+    def factory(language: Literal["python"]) -> LanguageManager:
+        if language == "python":
+            return PythonLanguageManager()
+        raise ValueError(f"Unsupported language: {language}")
+
+
+class PythonLanguageManager(LanguageManager):
+    """
+    Language manager for Python.
+    """
+
+    def install_dependencies(self, session: SandboxDockerSession, dependencies: list[str]) -> RunResult:
+        """
+        Install dependencies.
+        """
+        return session.execute_command(f"pip install {' '.join(dependencies)}", workdir="/")
+
+    def run_code(self, session: SandboxDockerSession, workdir: str, code: str) -> RunResult:
+        """
+        Run code.
+        """
+        with io.BytesIO() as tar_file:
+            with tarfile.open(fileobj=tar_file, mode="w:gz") as tar:
+                tarinfo = tarfile.TarInfo(name="main.py")
+                tarinfo.size = len(code.encode())
+                tar.addfile(tarinfo, io.BytesIO(code.encode()))
+
+            tar_file.seek(0)
+            session.copy_to_runtime(workdir, tar_file)
+
+        return session.execute_command("python main.py", workdir=workdir)

daiv_sandbox/main.py

@@ -8,11 +8,13 @@
 from fastapi import Depends, FastAPI, HTTPException, Security
 from fastapi.security.api_key import APIKeyHeader
 from sentry_sdk.integrations.fastapi import FastApiIntegration
-from starlette.status import HTTP_403_FORBIDDEN
+from starlette.status import HTTP_400_BAD_REQUEST, HTTP_403_FORBIDDEN
+
+from daiv_sandbox.languages import LanguageManager
 
 from . import __version__
 from .config import settings
-from .schemas import ForbiddenError, RunRequest, RunResponse, RunResult
+from .schemas import ErrorMessage, RunCodeRequest, RunCodeResponse, RunRequest, RunResponse, RunResult
 from .sessions import SandboxDockerSession
 
 HEADER_API_KEY_NAME = "X-API-Key"
@@ -75,12 +77,13 @@ async def get_api_key(api_key_header: str | None = Security(api_key_header)) ->
     return api_key_header
 
 
-@app.post(
-    "/run/commands/",
-    responses={
-        403: {"content": {"application/json": {"example": {"detail": "Invalid API Key"}}}, "model": ForbiddenError}
-    },
-)
+common_responses = {
+    403: {"content": {"application/json": {"example": {"detail": "Invalid API Key"}}}, "model": ErrorMessage},
+    400: {"content": {"application/json": {"example": {"detail": "Error message"}}}, "model": ErrorMessage},
+}
+
+
+@app.post("/run/commands/", responses=common_responses)
 async def run_commands(request: RunRequest, api_key: str = Depends(get_api_key)) -> RunResponse:
     """
     Run a set of commands in a sandboxed container and return archive with changed files.
@@ -98,7 +101,10 @@ async def run_commands(request: RunRequest, api_key: str = Depends(get_api_key))
 
         command_workdir = Path(run_dir) / request.workdir if request.workdir else Path(run_dir)
 
-        results = [session.execute_command(command, workdir=command_workdir.as_posix()) for command in request.commands]
+        results = [
+            session.execute_command(command, workdir=command_workdir.as_posix(), extract_changed_files=True)
+            for command in request.commands
+        ]
 
         # Only create archive with changed files for the last command.
         if changed_files := results[-1].changed_files:
@@ -108,6 +114,36 @@ async def run_commands(request: RunRequest, api_key: str = Depends(get_api_key))
     return RunResponse(results=results, archive=archive)
 
 
+LANGUAGE_BASE_IMAGES = {"python": "python:3.12-slim"}
+
+
+@app.post("/run/code/", responses=common_responses)
+async def run_code(request: RunCodeRequest, api_key: str = Depends(get_api_key)) -> RunCodeResponse:
+    """
+    Run code in a sandboxed container and return the result.
+    """
+    run_dir = f"/runs/{request.run_id}"
+
+    with SandboxDockerSession(
+        image=LANGUAGE_BASE_IMAGES[request.language],
+        keep_template=True,
+        runtime=settings.RUNTIME,
+        run_id=request.run_id,
+    ) as session:
+        manager = LanguageManager.factory(request.language)
+
+        if request.dependencies:
+            install_result = manager.install_dependencies(session, request.dependencies)
+            if install_result.exit_code != 0:
+                raise HTTPException(status_code=HTTP_400_BAD_REQUEST, detail=install_result.output)
+
+        run_result = manager.run_code(session, run_dir, request.code)
+        if run_result.exit_code != 0:
+            raise HTTPException(status_code=HTTP_400_BAD_REQUEST, detail=run_result.output)
+
+    return RunCodeResponse(output=run_result.output)
+
+
 @app.get("/health/", responses={200: {"content": {"application/json": {"example": {"status": "ok"}}}}})
 async def health() -> dict[Literal["status"], Literal["ok"]]:
     """

daiv_sandbox/schemas.py

@@ -1,3 +1,5 @@
+from typing import Literal
+
 from pydantic import UUID4, Base64Bytes, BaseModel, Field
 
 
@@ -19,13 +21,26 @@ class RunResult(BaseModel):
     command: str = Field(..., description="Command that was executed.")
     output: str = Field(..., description="Output of the command.")
     exit_code: int = Field(..., description="Exit code of the command.")
-    changed_files: list[str] = Field(..., description="List of changed files.", exclude=True)
+    changed_files: list[str] = Field(default_factory=list, description="List of changed files.", exclude=True)
 
 
 class RunResponse(BaseModel):
     results: list[RunResult] = Field(..., description="List of results of each command.")
     archive: str | None = Field(..., description="Base64-encoded archive with the changed files.")
 
 
-class ForbiddenError(BaseModel):
+class ErrorMessage(BaseModel):
     detail: str = Field(..., description="Error message.")
+
+
+class RunCodeRequest(BaseModel):
+    run_id: UUID4 = Field(..., description="Unique identifier for the run.")
+    language: Literal["python"] = Field(..., description="Language to be used for the code execution.")
+    dependencies: list[str] = Field(
+        default_factory=list, description="List of dependencies to be installed in the sandbox."
+    )
+    code: str = Field(..., description="Code to be executed.")
+
+
+class RunCodeResponse(BaseModel):
+    output: str = Field(..., description="Output of the code execution.")

daiv_sandbox/sessions.py

@@ -47,7 +47,7 @@ def copy_to_runtime(self, dest: str, data: BinaryIO):
         raise NotImplementedError
 
     @abstractmethod
-    def execute_command(self, command: str, workdir: str) -> RunResult:
+    def execute_command(self, command: str, workdir: str, extract_changed_files: bool = False) -> RunResult:
         raise NotImplementedError
 
     def __enter__(self):
@@ -196,7 +196,7 @@ def copy_to_runtime(self, dest: str, data: BinaryIO):
         else:
             raise RuntimeError(f"Failed to copy archive to {self.container.short_id}:{dest}")
 
-    def execute_command(self, command: str, workdir: str) -> RunResult:
+    def execute_command(self, command: str, workdir: str, extract_changed_files: bool = False) -> RunResult:
         """
         Execute a command in the container.
         """
@@ -232,7 +232,7 @@ def execute_command(self, command: str, workdir: str) -> RunResult:
             command=command,
             output=result.output.decode(),
             exit_code=result.exit_code,
-            changed_files=self._extract_changed_file_names(workdir, before_run_date),
+            changed_files=self._extract_changed_file_names(workdir, before_run_date) if extract_changed_files else [],
         )
 
     def _extract_changed_file_names(self, workdir: str, modified_after: datetime.datetime) -> list[str]:

pyproject.toml

@@ -19,9 +19,9 @@ classifiers = [
 dependencies = [
   "docker==7.1",
   "fastapi[standard]==0.115.5",
-  "pydantic==2.10",
+  "pydantic==2.10.2",
   "pydantic-settings==2.6.1",
-  "sentry-sdk==2.18",
+  "sentry-sdk==2.19",
 ]
 
 urls."Bug Tracker" = "https://github.com/srtab/daiv-sandbox/issues"
@@ -112,5 +112,5 @@ dev-dependencies = [
   "pytest-env==1.1.5",
   "pytest-mock==3.14.0",
   "pytest-xdist==3.6.1",
-  "ruff==0.7.4",
+  "ruff==0.8.0",
 ]

tests/test_languages.py

@@ -0,0 +1,56 @@
+from unittest.mock import MagicMock
+
+import pytest
+
+from daiv_sandbox.languages import LanguageManager, PythonLanguageManager
+from daiv_sandbox.schemas import RunResult
+from daiv_sandbox.sessions import SandboxDockerSession
+
+
+@pytest.fixture
+def setup_manager():
+    session = MagicMock(spec=SandboxDockerSession)
+    manager = PythonLanguageManager()
+    return session, manager
+
+
+def test_factory():
+    manager = LanguageManager.factory("python")
+    assert isinstance(manager, PythonLanguageManager)
+
+
+def test_factory_unsupported_language():
+    with pytest.raises(ValueError, match="Unsupported language: unsupported"):
+        LanguageManager.factory("unsupported")
+
+
+def test_install_dependencies(setup_manager):
+    session, manager = setup_manager
+
+    # Mock the expected result
+    expected_result = RunResult(command="pip install numpy pandas", output="Dependencies installed", exit_code=0)
+    session.execute_command.return_value = expected_result
+
+    # Call the method
+    result = manager.install_dependencies(session, ["numpy", "pandas"])
+
+    # Assertions
+    session.execute_command.assert_called_once_with("pip install numpy pandas", workdir="/")
+    assert result == expected_result
+
+
+def test_run_code(setup_manager):
+    session, manager = setup_manager
+
+    # Mock the expected result
+    expected_result = RunResult(command="python main.py", output="Code executed", exit_code=0)
+    session.execute_command.return_value = expected_result
+
+    # Call the method
+    code = "print('Hello, World!')"
+    result = manager.run_code(session, "/workdir", code)
+
+    # Assertions
+    session.copy_to_runtime.assert_called_once()
+    session.execute_command.assert_called_once_with("python main.py", workdir="/workdir")
+    assert result == expected_result