Improve task push

Author: vtm9

push_tasks.py

@@ -4,9 +4,12 @@
 """
 
 import argparse
+import base64
+import gzip
 import json
 import os
 import sys
+import time
 import urllib.error
 import urllib.request
 from pathlib import Path
@@ -15,7 +18,7 @@
 
 def read_json_files(directory: str) -> List[tuple]:
     """
-    Read all JSON files from the specified directory.
+    Read all JSON files from the specified directory recursively.
 
     Returns:
         List of tuples (filename, json_data)
@@ -27,7 +30,7 @@ def read_json_files(directory: str) -> List[tuple]:
         print(f"Error: Directory '{directory}' does not exist", file=sys.stderr)
         sys.exit(1)
 
-    for json_file in release_path.glob("*.json"):
+    for json_file in release_path.rglob("*.json"):
         try:
             with open(json_file, "r", encoding="utf-8") as f:
                 data = json.load(f)
@@ -41,61 +44,120 @@ def read_json_files(directory: str) -> List[tuple]:
     return json_files
 
 
-def publish_task(
-    url: str, task_data: dict, visibility: str, origin: str, auth_token: str
-) -> bool:
+def publish_task_batch(
+    url: str, tasks_batch: List[tuple], visibility: str, origin: str, auth_token: str
+) -> dict:
     """
-    Publish a single task to the remote URL.
+    Publish a batch of tasks to the remote URL with gzip compression and base64 encoding.
 
     Args:
         url: The target URL
-        task_data: The JSON data to publish
+        tasks_batch: List of tuples (filename, task_data)
         visibility: Visibility setting ('public' or 'hidden')
         origin: Origin of the tasks
         auth_token: Authorization token from environment
 
     Returns:
-        True if successful, False otherwise
+        Dictionary with 'success' count and 'failed' list of filenames
     """
-    # Prepare the payload
-    payload = {"task": task_data, "visibility": visibility, "origin": origin}
+    result = {"success": 0, "failed": []}
+
+    # Prepare the tasks list (just the task data, no wrapping)
+    tasks_list = [task_data for _, task_data in tasks_batch]
+
+    # Convert to JSON, compress with gzip, and encode with base64
+    json_data = json.dumps(tasks_list).encode("utf-8")
+    compressed_data = gzip.compress(json_data)
+    base64_payload = base64.b64encode(compressed_data).decode("ascii")
+
+    # Prepare the final payload structure matching the API
+    payload = {
+        "payload": base64_payload,
+        "visibility": visibility,
+        "origin": origin,
+    }
 
     # Prepare headers
     headers = {
         "Content-Type": "application/json",
+        "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36",
     }
 
     if auth_token:
         headers["X-AUTH-KEY"] = auth_token
 
     # Convert payload to JSON bytes
-    data = json.dumps(payload).encode("utf-8")
+    request_data = json.dumps(payload).encode("utf-8")
 
     # Create request
-    req = urllib.request.Request(url, data=data, headers=headers, method="POST")
+    req = urllib.request.Request(url, data=request_data, headers=headers, method="POST")
 
     try:
-        with urllib.request.urlopen(req) as response:
+        with urllib.request.urlopen(req, timeout=30) as response:
             status = response.status
             response_body = response.read().decode("utf-8")
             if status in (200, 201):
-                return True
+                result["success"] = len(tasks_batch)
+                return result
             else:
                 print(f"  Warning: Received status {status}", file=sys.stderr)
                 print(f"  Response: {response_body}", file=sys.stderr)
-                return False
+                result["failed"] = [filename for filename, _ in tasks_batch]
+                return result
     except urllib.error.HTTPError as e:
         error_body = e.read().decode("utf-8") if e.fp else ""
         print(f"  HTTP Error {e.code}: {e.reason}", file=sys.stderr)
         if error_body:
-            print(f"  Response: {error_body}", file=sys.stderr)
-        return False
+            # Truncate very long error messages
+            if len(error_body) > 500:
+                print(f"  Response: {error_body[:500]}...", file=sys.stderr)
+            else:
+                print(f"  Response: {error_body}", file=sys.stderr)
+        result["failed"] = [filename for filename, _ in tasks_batch]
+        return result
     except urllib.error.URLError as e:
         print(f"  URL Error: {e.reason}", file=sys.stderr)
-        return False
+        result["failed"] = [filename for filename, _ in tasks_batch]
+        return result
     except Exception as e:
         print(f"  Error: {e}", file=sys.stderr)
-        return False
+        result["failed"] = [filename for filename, _ in tasks_batch]
+        return result
+
+
+def load_env_file(env_path: str = ".env") -> dict:
+    """
+    Load environment variables from a .env file.
+
+    Args:
+        env_path: Path to the .env file
+
+    Returns:
+        Dictionary of environment variables
+    """
+    env_vars = {}
+    env_file = Path(env_path)
+
+    if not env_file.exists():
+        return env_vars
+
+    try:
+        with open(env_file, "r", encoding="utf-8") as f:
+            for line in f:
+                line = line.strip()
+                # Skip empty lines and comments
+                if not line or line.startswith("#"):
+                    continue
+                # Parse KEY=VALUE format
+                if "=" in line:
+                    key, value = line.split("=", 1)
+                    # Remove quotes if present
+                    value = value.strip().strip("\"'")
+                    env_vars[key.strip()] = value
+    except Exception as e:
+        print(f"Warning: Error reading .env file: {e}", file=sys.stderr)
+
+    return env_vars
 
 
 def main():
@@ -121,6 +183,12 @@ def main():
         default="release",
         help="Directory containing JSON files (default: release)",
     )
+    parser.add_argument(
+        "--batch-size",
+        type=int,
+        default=20,
+        help="Number of tasks to send in each batch (default: 20)",
+    )
 
     args = parser.parse_args()
 
@@ -131,11 +199,16 @@ def main():
 
     visibility = "public" if args.public else "hidden"
 
-    # Get auth token from environment
-    auth_token = os.environ.get("CODEBATTLE_AUTH_TOKEN", "")
+    # Load .env file
+    env_vars = load_env_file()
+
+    # Get auth token from .env file or environment variable
+    auth_token = env_vars.get("CODEBATTLE_AUTH_TOKEN") or os.environ.get(
+        "CODEBATTLE_AUTH_TOKEN", ""
+    )
     if not auth_token:
         print(
-            "Warning: CODEBATTLE_AUTH_TOKEN environment variable not set",
+            "Warning: CODEBATTLE_AUTH_TOKEN not found in .env file or environment",
             file=sys.stderr,
         )
 
@@ -150,20 +223,39 @@ def main():
     print(f"\nPublishing {len(json_files)} tasks to {args.url}...")
     print(f"Visibility: {visibility}")
     print(f"Origin: {args.origin}")
+    print(f"Batch size: {args.batch_size}")
     print()
 
-    # Publish each task
+    # Publish tasks in batches
     success_count = 0
     fail_count = 0
+    total_batches = (len(json_files) + args.batch_size - 1) // args.batch_size
+
+    for batch_num in range(0, len(json_files), args.batch_size):
+        batch = json_files[batch_num : batch_num + args.batch_size]
+        batch_index = batch_num // args.batch_size + 1
+
+        print(
+            f"Batch {batch_index}/{total_batches} ({len(batch)} tasks)... ",
+            end="",
+            flush=True,
+        )
+
+        result = publish_task_batch(
+            args.url, batch, visibility, args.origin, auth_token
+        )
+
+        if result["success"] > 0:
+            print(f"✓ ({result['success']} tasks)")
+            success_count += result["success"]
+
+        if result["failed"]:
+            print(f"✗ Failed tasks: {', '.join(result['failed'])}")
+            fail_count += len(result["failed"])
 
-    for filename, task_data in json_files:
-        print(f"Publishing: {filename}... ", end="", flush=True)
-        if publish_task(args.url, task_data, visibility, args.origin, auth_token):
-            print("✓")
-            success_count += 1
-        else:
-            print("✗")
-            fail_count += 1
+        # Add a small delay between batches to avoid rate limiting
+        if batch_index < total_batches:
+            time.sleep(1)
 
     # Summary
     print()