terrakio-core 0.4.8__py3-none-any.whl → 0.4.94__py3-none-any.whl

terrakio_core/async_client.py

@@ -17,7 +17,9 @@ from .endpoints.group_management import GroupManagement
 from .endpoints.space_management import SpaceManagement
 from .endpoints.model_management import ModelManagement
 from .endpoints.auth import AuthClient
-from .convenience_functions.convenience_functions import zonal_stats as _zonal_stats, create_dataset_file as _create_dataset_file, request_geoquery_list as _request_geoquery_list
+from .convenience_functions.zonal_stats import zonal_stats as _zonal_stats
+from .convenience_functions.geoquries import request_geoquery_list as _request_geoquery_list
+from .convenience_functions.create_dataset_file import create_dataset_file as _create_dataset_file
 
 class AsyncClient(BaseClient):
     def __init__(self, url: Optional[str] = None, api_key: Optional[str] = None, verbose: bool = False, session: Optional[aiohttp.ClientSession] = None):
@@ -34,11 +36,15 @@ class AsyncClient(BaseClient):
 
     async def _terrakio_request(self, method: str, endpoint: str, **kwargs):
         if self.session is None:
+            # To this:
             headers = {
-                'Content-Type': 'application/json',
-                'x-api-key': self.key,
+                'x-api-key': self.key, 
                 'Authorization': self.token
             }
+
+            # Only add Content-Type if it's a JSON request
+            if 'json' in kwargs:
+                headers['Content-Type'] = 'application/json'
             clean_headers = {k: v for k, v in headers.items() if v is not None}
             async with aiohttp.ClientSession(headers=clean_headers, timeout=aiohttp.ClientTimeout(total=self.timeout)) as session:
                 return await self._make_request_with_retry(session, method, endpoint, **kwargs)
@@ -220,6 +226,7 @@ class AsyncClient(BaseClient):
             ValueError: If concurrency is too high or if data exceeds memory limit without streaming
             APIError: If the API request fails
         """
+        # the sync client didn't pass the self here, so the client is now async
         return await _zonal_stats(
             client=self,
             gdf=gdf,

terrakio_core/convenience_functions/create_dataset_file.py

@@ -0,0 +1,132 @@
+import asyncio
+import os
+import tempfile
+import time
+import uuid
+
+from ..helper.tiles import tiles
+
+async def create_dataset_file(
+    client,
+    aoi: str,
+    expression: str,
+    output: str,
+    download_path: str,
+    in_crs: str = "epsg:4326",
+    to_crs: str = "epsg:4326",
+    res: float = 0.0001,
+    region: str = None,
+    overwrite: bool = False,
+    skip_existing: bool = False,
+    non_interactive: bool = True,
+    name: str | None = None,
+    poll_interval: int = 30,
+    max_file_size_mb: int = 5120,
+    tile_size: int = 1024,
+    mask: bool = True
+) -> dict:
+    
+    if not name:
+        name = f"file-gen-{uuid.uuid4().hex[:8]}"
+    
+    body, reqs, groups = tiles(
+        name = name, 
+        aoi = aoi, 
+        expression = expression,
+        output = output,
+        tile_size = tile_size,
+        crs = in_crs,
+        res = res,
+        region = region,
+        to_crs = to_crs,
+        mask = mask,
+        overwrite = overwrite,
+        skip_existing = skip_existing,
+        non_interactive = non_interactive
+    )
+    with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False) as tempreq:
+        tempreq.write(reqs)
+        tempreqname = tempreq.name
+
+    task_id = await client.mass_stats.execute_job(
+        name=body["name"],
+        region=body["region"],
+        output=body["output"],
+        config = {},
+        overwrite=body["overwrite"],
+        skip_existing=body["skip_existing"],
+        request_json=tempreqname,
+    )
+
+    start_time = time.time()
+    status = None
+    client.logger.info(f"Tracking data generation job {task_id['task_id']}...")
+    while True:
+        try:
+            taskid = task_id['task_id']
+            trackinfo = await client.mass_stats.track_job([taskid])
+            status = trackinfo[taskid]['status']
+            if status == 'Completed':
+                client.logger.info('Data generated successfully!')
+                break
+            elif status in ['Failed', 'Cancelled', 'Error']:
+                raise RuntimeError(f"Job {taskid} failed with status: {status}")
+            else:
+                elapsed_time = time.time() - start_time
+                client.logger.info(f"Job status: {status} - Elapsed time: {elapsed_time:.1f}s")
+                await asyncio.sleep(poll_interval)
+                
+                
+        except KeyboardInterrupt:
+            client.logger.info(f"\nInterrupted! Job {taskid} is still running in the background.")
+            raise
+        except Exception as e:
+            client.logger.info(f"\nError tracking job: {e}")
+            raise
+
+    os.unlink(tempreqname)
+
+    combine_result = await client.mass_stats.combine_tiles(body["name"], body["overwrite"], body["output"], max_file_size_mb=max_file_size_mb)
+    combine_task_id = combine_result.get("task_id")
+
+    combine_start_time = time.time()
+    client.logger.info(f"Tracking file generation job {combine_task_id}...")
+    while True:
+        try:
+            trackinfo = await client.mass_stats.track_job([combine_task_id])
+            if body["output"] == "netcdf":
+                download_file_name = trackinfo[combine_task_id]['folder'] + '.nc'
+            elif body["output"] == "geotiff":
+                download_file_name = trackinfo[combine_task_id]['folder'] + '.tif'
+            bucket = trackinfo[combine_task_id]['bucket']
+            combine_status = trackinfo[combine_task_id]['status']
+            if combine_status == 'Completed':
+                client.logger.info('File/s generated successfully!')
+                break
+            elif combine_status in ['Failed', 'Cancelled', 'Error']:
+                raise RuntimeError(f"File generation job {combine_task_id} failed with status: {combine_status}")
+            else:
+                elapsed_time = time.time() - combine_start_time
+                client.logger.info(f"File generation job status: {combine_status} - Elapsed time: {elapsed_time:.1f}s")
+                time.sleep(poll_interval)
+        except KeyboardInterrupt:
+            client.logger.info(f"\nInterrupted! File generation job {combine_task_id} is still running in the background.")
+            raise
+        except Exception as e:
+            client.logger.info(f"\nError tracking file generation job: {e}")
+            raise
+
+    if download_path:
+        await client.mass_stats.download_file(
+            job_name=body["name"],
+            bucket=bucket,
+            file_type='processed',
+            folder='file-gen',
+            page_size=100,
+            output_path=download_path,
+        )
+    else:
+        path = f"{body['name']}/outputs/merged/{download_file_name}"
+        client.logger.info(f"Dataset file/s is available at {path}")
+
+    return {"generation_task_id": task_id, "combine_task_id": combine_task_id}

terrakio_core/convenience_functions/geoquries.py

@@ -0,0 +1,102 @@
+import asyncio
+
+import geopandas as gpd
+from shapely.geometry import shape
+
+from ..exceptions import APIError
+from ..helper.bounded_taskgroup import BoundedTaskGroup
+
+async def request_geoquery_list(
+        client,
+        quries: list[dict],
+        conc: int = 20,
+):
+    """
+    Execute multiple geo queries.
+    
+    Args:
+        client: The Terrakio client instance
+        quries: List of dictionaries containing query parameters
+        conc: The concurrency level for the requests
+        
+    Returns:
+        List of query results
+        
+    Raises:
+        ValueError: If the queries list is empty
+    """
+    if not quries:
+        raise ValueError("Queries list cannot be empty")
+    if conc > 100:
+        raise ValueError("Concurrency (conc) is too high. Please set conc to 100 or less.")
+
+    for i, query in enumerate(quries):
+        if 'expr' not in query:
+            raise ValueError(f"Query at index {i} is missing the required 'expr' key")
+        if 'feature' not in query:
+            raise ValueError(f"Query at index {i} is missing the required 'feature' key")
+        if 'in_crs' not in query:
+            raise ValueError(f"Query at index {i} is missing the required 'in_crs' key")
+    
+    completed_count = 0
+    lock = asyncio.Lock()
+    async def single_geo_query(query):
+        """
+        Execute multiple geo queries concurrently.
+        
+        Args:
+            quries: List of dictionaries containing query parameters
+        """
+        total_number_of_requests = len(quries)
+        nonlocal completed_count
+        try:
+            result = await client.geoquery(**query)
+            if isinstance(result, dict) and result.get("error"):
+                error_msg = f"Request failed: {result.get('error_message', 'Unknown error')}"
+                if result.get('status_code'):
+                    error_msg = f"Request failed with status {result['status_code']}: {result.get('error_message', 'Unknown error')}"
+                raise APIError(error_msg)
+            if isinstance(result, list):
+                result = result[0]
+                timestamp_number = result['request_count']
+                return timestamp_number
+            if not isinstance(result, xr.Dataset):
+                raise ValueError(f"Expected xarray Dataset, got {type(result)}")
+            
+            async with lock:
+                completed_count += 1
+                if completed_count % max(1, total_number_of_requests // 10) == 0:
+                    client.logger.info(f"Progress: {completed_count}/{total_number_of_requests} requests processed")
+            return result   
+        except Exception as e:
+            async with lock:
+                completed_count += 1
+            raise
+    
+    try:
+        async with BoundedTaskGroup(max_concurrency=conc) as tg:
+            tasks = [tg.create_task(single_geo_query(quries[idx])) for idx in range(len(quries))]
+        all_results = [task.result() for task in tasks]
+
+    except* Exception as eg:
+        for e in eg.exceptions:
+            if hasattr(e, 'response'):
+                raise APIError(f"API request failed: {e.response.text}")
+        raise
+    client.logger.info("All requests completed!")
+     
+    if not all_results:
+        raise ValueError("No valid results were returned for any geometry")
+    if isinstance(all_results, list) and type(all_results[0]) == int:
+        return sum(all_results)/len(all_results)
+    else:
+        geometries = []
+        for query in quries:
+            feature = query['feature']
+            geometry = shape(feature['geometry'])
+            geometries.append(geometry)
+        result_gdf = gpd.GeoDataFrame({
+            'geometry': geometries,
+            'dataset': all_results
+        })
+        return result_gdf