terrakio-core 0.3.9__py3-none-any.whl → 0.4.0__py3-none-any.whl

terrakio_core/endpoints/mass_stats.py

@@ -6,7 +6,7 @@ from pathlib import Path
 from urllib.parse import urlparse
 from ..helper.decorators import require_token, require_api_key, require_auth
 import aiohttp
-
+from typing import Dict, Any, Optional, List, Union
 class MassStats:
     def __init__(self, client):
         self._client = client
@@ -218,7 +218,7 @@ class MassStats:
             params["output"] = output
             
         return self._client._terrakio_request("GET", "mass_stats/download", params=params)
-    
+
     @require_api_key
     async def _upload_file(self, file_path: str, url: str, use_gzip: bool = False):
         """
@@ -237,6 +237,18 @@ class MassStats:
         except json.JSONDecodeError as e:
             raise ValueError(f"Invalid JSON in file {file_path}: {e}")
         
+        return await self._upload_json_data(json_data, url, use_gzip)
+
+    @require_api_key
+    async def _upload_json_data(self, json_data: Union[Dict, List], url: str, use_gzip: bool = False):
+        """
+        Helper method to upload JSON data directly to a signed URL.
+        
+        Args:
+            json_data: JSON data (dict or list) to upload
+            url: Signed URL to upload to
+            use_gzip: Whether to compress the data with gzip
+        """
         if hasattr(json, 'dumps') and 'ignore_nan' in json.dumps.__code__.co_varnames:
             dumps_kwargs = {'ignore_nan': True}
         else:
@@ -253,6 +265,7 @@ class MassStats:
             headers = {
                 'Content-Type': 'application/json'
             }
+        
         response = await self._client._regular_request("PUT", url, data=body, headers=headers)
         return response
     
@@ -386,36 +399,55 @@ class MassStats:
         region: str, 
         output: str, 
         config: Dict[str, Any],
-        request_json: Dict[str, Any], 
-        manifest_json: Dict[str, Any],
+        request_json: str,  # Path to request JSON file
         overwrite: bool = False, 
         skip_existing: bool = False,  
         location: str = None, 
         force_loc: bool = None, 
-        server: str = "dev-au.terrak.io"
+        server: str = None
     ) -> Dict[str, Any]:
         """
         Execute a mass stats job.
-
+        
         Args:
             name: The name of the job
             region: The region of the job
             output: The output of the job
             config: The config of the job
-            request_json: The request JSON
-            manifest_json: The manifest JSON
+            request_json: Path to the request JSON file
             overwrite: Whether to overwrite the job
             skip_existing: Whether to skip existing jobs
             location: The location of the job
             force_loc: Whether to force the location
             server: The server to use
-
+        
         Returns:
             API response as a dictionary
-
+        
         Raises:
             APIError: If the API request fails
         """
+        
+        def extract_manifest_from_request(request_data: List[Dict[str, Any]]) -> List[str]:
+            """Extract unique group names from request data to create manifest list."""
+            groups = []
+            seen_groups = set()
+            
+            for item in request_data:
+                if not isinstance(item, dict):
+                    raise ValueError("Each item in request JSON should be a dictionary")
+                    
+                if 'group' not in item:
+                    raise ValueError("Each item should have a 'group' field")
+                    
+                group = item['group']
+                if group not in seen_groups:
+                    groups.append(group)
+                    seen_groups.add(group)
+                    
+            return groups
+        
+        # Load and validate request JSON
         try:
             with open(request_json, 'r') as file:
                 request_data = json.load(file)
@@ -427,14 +459,35 @@ class MassStats:
             return e
         except json.JSONDecodeError as e:
             return e
-        upload_result = await self._upload_request(name = name, size = size, region = region, output = output, config = config, location = location, force_loc = force_loc, overwrite = overwrite, server = server, skip_existing = skip_existing)
+        
+        # Generate manifest from request data (kept in memory)
+        try:
+            manifest_groups = extract_manifest_from_request(request_data)
+        except Exception as e:
+            raise ValueError(f"Error extracting manifest from request JSON: {e}")
+        
+        # Get upload URLs
+        upload_result = await self._upload_request(
+            name=name, 
+            size=size, 
+            region=region, 
+            output=output, 
+            config=config, 
+            location=location, 
+            force_loc=force_loc, 
+            overwrite=overwrite, 
+            server=server, 
+            skip_existing=skip_existing
+        )
+        
         requests_url = upload_result.get('requests_url')
         manifest_url = upload_result.get('manifest_url')
+        
         if not requests_url:
             raise ValueError("No requests_url returned from server for request JSON upload")
         
+        # Upload request JSON file
         try:
-            # in this place we are uploading the request json file, we need to check whether the json is in the correct format or not
             self.validate_request(request_json)
             requests_response = await self._upload_file(request_json, requests_url, use_gzip=True)
             if requests_response.status not in [200, 201, 204]:
@@ -446,15 +499,17 @@ class MassStats:
         if not manifest_url:
             raise ValueError("No manifest_url returned from server for manifest JSON upload")
         
+        # Upload manifest JSON data directly (no temporary file needed)
         try:
-            manifest_response = await self._upload_file(manifest_json, manifest_url, use_gzip=False)
+            manifest_response = await self._upload_json_data(manifest_groups, manifest_url, use_gzip=False)
             if manifest_response.status not in [200, 201, 204]:
                 self._client.logger.error(f"Manifest upload error: {manifest_response.text()}")
                 raise Exception(f"Failed to upload manifest JSON: {manifest_response.text()}")
         except Exception as e:
-            raise Exception(f"Error uploading manifest JSON file {manifest_json}: {e}")
-
-        start_job_task_id =await self.start_job(upload_result.get("id"))
+            raise Exception(f"Error uploading manifest JSON: {e}")
+        
+        # Start the job
+        start_job_task_id = await self.start_job(upload_result.get("id"))
         return start_job_task_id
 
     @require_api_key