Allow direct code download

packages/bot/CHANGELOG.md

@@ -2,6 +2,12 @@
 
 ### Unreleased
 
+### 0.3.0
+
+#### Added
+
+- Support direct download method for code
+
 ### 0.2.3
 
 #### Added

packages/bot/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name    = "automa-bot"
-version = "0.2.3"
+version = "0.3.0"
 
 authors         = [{ name = "Sunkara, Inc.", email = "engineering@automa.app" }]
 description     = "Bot helpers for Automa"

packages/bot/src/automa/bot/resources/code.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import json
 import subprocess
 import tarfile
 from asyncio import to_thread
@@ -77,6 +78,42 @@ def _read_token(self, folder: str) -> str | None:
 
         return token
 
+    def _read_base_commit(self, folder: str) -> str | None:
+        base_commit = None
+
+        try:
+            with open(
+                f"{folder}/.git/automa_proposal_base_commit", "r", encoding="utf8"
+            ) as f:
+                base_commit = f.read().strip()
+        except FileNotFoundError:
+            pass
+
+        return base_commit
+
+    def _clone_code(self, folder: str, url: str) -> None:
+        path = Path(folder)
+
+        rmtree(folder, ignore_errors=True)
+        path.mkdir(parents=True, exist_ok=True)
+
+        subprocess.run(
+            ["git", "clone", "--depth=1", url, "."],
+            cwd=path,
+            check=True,
+        )
+
+        # Note down the base commit
+        result = subprocess.run(
+            ["git", "rev-parse", "HEAD"],
+            cwd=path,
+            capture_output=True,
+            text=True,
+            check=True,
+        )
+
+        self._write_base_commit(folder, result.stdout.strip())
+
     def _extract_download(self, folder: str) -> None:
         rmtree(folder, ignore_errors=True)
         Path(folder).mkdir(parents=True, exist_ok=True)
@@ -89,13 +126,24 @@ def _write_token(self, folder: str, token: str) -> None:
         with open(f"{folder}/.git/automa_proposal_token", "w", encoding="utf8") as f:
             f.write(token)
 
+    def _write_base_commit(self, folder: str, base_commit: str) -> None:
+        # Save the base commit for later use
+        with open(
+            f"{folder}/.git/automa_proposal_base_commit", "w", encoding="utf8"
+        ) as f:
+            f.write(base_commit)
+
 
 class CodeResource(SyncAPIResource, BaseCodeResource):
     def cleanup(self, body: CodeCleanupParams) -> None:
         folder = self._path(body["task"])
 
         rmtree(folder, ignore_errors=True)
-        remove(f"{folder}.tar.gz")
+
+        try:
+            remove(f"{folder}.tar.gz")
+        except FileNotFoundError:
+            pass
 
     def download(
         self, body: CodeDownloadParams, *, options: RequestOptions = {}
@@ -112,20 +160,31 @@ def download(
                 "json": body,
                 "headers": {
                     **options.get("headers", {}),
-                    "Accept": "application/gzip",
                 },
             },
         ) as response:
-            token = response.headers["x-automa-proposal-token"]
+            token = response.headers.get("x-automa-proposal-token")
+            content_type = response.headers.get("Content-Type", "")
+
+            if content_type.startswith("application/json"):
+                content = b"".join(response.iter_bytes())
 
-            rmtree(path, ignore_errors=True)
-            makedirs(path, exist_ok=True)
+                json_data = json.loads(content.decode("utf-8"))
 
-            with open(archive_path, "wb") as archive:
-                for chunk in response.iter_bytes(chunk_size=8192):
-                    archive.write(chunk)
+                self._clone_code(path, json_data["url"])
+            elif content_type.startswith("application/gzip"):
+                rmtree(path, ignore_errors=True)
+                makedirs(path, exist_ok=True)
 
-        self._extract_download(path)
+                with open(archive_path, "wb") as archive:
+                    for chunk in response.iter_bytes(chunk_size=8192):
+                        archive.write(chunk)
+
+                self._extract_download(path)
+            else:
+                raise ValueError(
+                    f"Unexpected Content-Type: {content_type} while downloading code."
+                )
 
         # Save the proposal token for later use
         self._write_token(path, token)
@@ -135,6 +194,7 @@ def download(
     def propose(self, body: CodeProposeParams, *, options: RequestOptions = {}):
         path = self._path(body["task"])
         token = self._read_token(path)
+        base_commit = self._read_base_commit(path)
 
         if not token:
             raise ValueError("Failed to read the stored proposal token")
@@ -149,6 +209,7 @@ def propose(self, body: CodeProposeParams, *, options: RequestOptions = {}):
                     **body.get("proposal", {}),
                     "token": token,
                     "diff": diff,
+                    **({"base_commit": base_commit} if base_commit else {}),
                 },
             },
             options=options,
@@ -160,7 +221,11 @@ async def cleanup(self, body: CodeCleanupParams) -> None:
         folder = self._path(body["task"])
 
         await to_thread(rmtree, folder, ignore_errors=True)
-        await to_thread(remove, f"{folder}.tar.gz")
+
+        try:
+            await to_thread(remove, f"{folder}.tar.gz")
+        except FileNotFoundError:
+            pass
 
     async def download(
         self, body: CodeDownloadParams, *, options: RequestOptions = {}
@@ -177,20 +242,34 @@ async def download(
                 "json": body,
                 "headers": {
                     **options.get("headers", {}),
-                    "Accept": "application/gzip",
                 },
             },
         ) as response:
-            token = response.headers["x-automa-proposal-token"]
+            token = response.headers.get("x-automa-proposal-token")
+            content_type = response.headers.get("Content-Type", "")
+
+            if content_type.startswith("application/json"):
+                content = b""
+
+                async for chunk in response.aiter_bytes():
+                    content += chunk
+
+                json_data = json.loads(content.decode("utf-8"))
 
-            await to_thread(rmtree, path, ignore_errors=True)
-            await to_thread(makedirs, path, exist_ok=True)
+                await to_thread(self._clone_code, path, json_data["url"])
+            elif content_type.startswith("application/gzip"):
+                await to_thread(rmtree, path, ignore_errors=True)
+                await to_thread(makedirs, path, exist_ok=True)
 
-            with open(archive_path, "wb") as archive:
-                async for chunk in response.aiter_bytes(chunk_size=8192):
-                    await to_thread(archive.write, chunk)
+                with open(archive_path, "wb") as archive:
+                    async for chunk in response.aiter_bytes(chunk_size=8192):
+                        await to_thread(archive.write, chunk)
 
-        await to_thread(self._extract_download, path)
+                await to_thread(self._extract_download, path)
+            else:
+                raise ValueError(
+                    f"Unexpected Content-Type: {content_type} while downloading code."
+                )
 
         # Save the proposal token for later use
         await to_thread(self._write_token, path, token)
@@ -200,6 +279,7 @@ async def download(
     async def propose(self, body: CodeProposeParams, *, options: RequestOptions = {}):
         path = self._path(body["task"])
         token = await to_thread(self._read_token, path)
+        base_commit = await to_thread(self._read_base_commit, path)
 
         if not token:
             raise ValueError("Failed to read the stored proposal token")
@@ -214,6 +294,7 @@ async def propose(self, body: CodeProposeParams, *, options: RequestOptions = {}
                     **body.get("proposal", {}),
                     "token": token,
                     "diff": diff,
+                    **({"base_commit": base_commit} if base_commit else {}),
                 },
             },
             options=options,

packages/bot/tests/fixtures/download_git/HEAD

@@ -0,0 +1 @@
+ref: refs/heads/master

packages/bot/tests/fixtures/download_git/refs/heads/master

@@ -0,0 +1 @@
+cc3f46ae7fdf71747b66b3e4272c0e5fe290d116