terrakio-core 0.3.0__py3-none-any.whl → 0.3.2__py3-none-any.whl

terrakio_core/__init__.py

@@ -4,4 +4,4 @@ Terrakio Core
 Core components for Terrakio API clients.
 """
 
-__version__ = "0.3.0"
+__version__ = "0.3.2"

terrakio_core/client.py

@@ -13,6 +13,7 @@ from shapely.geometry import shape, mapping
 from shapely.geometry.base import BaseGeometry as ShapelyGeometry
 from google.cloud import storage
 from .exceptions import APIError, ConfigurationError
+from .decorators import admin_only_params
 import logging
 import textwrap
 
@@ -129,6 +130,7 @@ class BaseClient:
             "resolution": resolution,
             **kwargs
         }
+        print("the payload is ", payload)
         request_url = f"{self.url}/geoquery"
         for attempt in range(retry + 1):
             try:
@@ -536,8 +538,8 @@ class BaseClient:
     def __exit__(self, exc_type, exc_val, exc_tb):
         self.close()
 
-    # Mass Stats methods
-    def upload_mass_stats(self, name, size, bucket, output, location=None, **kwargs):
+    @admin_only_params('location', 'force_loc', 'server')
+    def execute_job(self, name, region, output, config, overwrite=False, skip_existing=False, request_json=None, manifest_json=None, location=None, force_loc=None, server="dev-au.terrak.io"):
         if not self.mass_stats:
             from terrakio_core.mass_stats import MassStats
             if not self.url or not self.key:
@@ -548,20 +550,8 @@ class BaseClient:
                 verify=self.verify,
                 timeout=self.timeout
             )
-        return self.mass_stats.upload_request(name, size, bucket, output, location, **kwargs)
+        return self.mass_stats.execute_job(name, region, output, config, overwrite, skip_existing, request_json, manifest_json, location, force_loc, server)
 
-    def start_mass_stats_job(self, task_id):
-        if not self.mass_stats:
-            from terrakio_core.mass_stats import MassStats
-            if not self.url or not self.key:
-                raise ConfigurationError("Mass Stats client not initialized. Make sure API URL and key are set.")
-            self.mass_stats = MassStats(
-                base_url=self.url,
-                api_key=self.key,
-                verify=self.verify,
-                timeout=self.timeout
-            )
-        return self.mass_stats.start_job(task_id)
 
     def get_mass_stats_task_id(self, name, stage, uid=None):
         if not self.mass_stats:
@@ -576,7 +566,7 @@ class BaseClient:
             )
         return self.mass_stats.get_task_id(name, stage, uid)
 
-    def track_mass_stats_job(self, ids=None):
+    def track_mass_stats_job(self, ids: Optional[list] = None):
         if not self.mass_stats:
             from terrakio_core.mass_stats import MassStats
             if not self.url or not self.key:
@@ -1049,6 +1039,20 @@ class BaseClient:
             )
         return self.space_management.delete_data_in_path(path, region)
     
+    def start_mass_stats_job(self, task_id):
+        if not self.mass_stats:
+            from terrakio_core.mass_stats import MassStats
+            if not self.url or not self.key:
+                raise ConfigurationError("Mass Stats client not initialized. Make sure API URL and key are set.")
+            self.mass_stats = MassStats(
+                base_url=self.url,
+                api_key=self.key,
+                verify=self.verify,
+                timeout=self.timeout
+            )
+        return self.mass_stats.start_job(task_id)
+
+    
     def generate_ai_dataset(
         self,
         name: str,
@@ -1118,27 +1122,58 @@ class BaseClient:
             overwrite=True
         )["task_id"]
         print("the task id is ", task_id)
+
+        # Wait for job completion
+        import time
+        
+        while True:
+            result = self.track_mass_stats_job(ids=[task_id])
+            status = result[task_id]['status']
+            print(f"Job status: {status}")
+            
+            if status == "Completed":
+                break
+            elif status == "Error":
+                raise Exception(f"Job {task_id} encountered an error")
+                
+            # Wait 30 seconds before checking again
+            time.sleep(30)
+
+        # print("the result is ", result)
+        # after all the random sample jos are done, we then start the mass stats job
         task_id = self.start_mass_stats_job(task_id)
-        print("the task id is ", task_id)
-        return task_id
+        # now we hav ethe random sampel
 
+        # print("the task id is ", task_id)
+        return task_id
 
-    def train_model(self, model_name: str, training_data: dict) -> dict:
+    def train_model(self, model_name: str, training_dataset: str, task_type: str, model_category: str, architecture: str, region: str, hyperparameters: dict = None) -> dict:
         """
         Train a model using the external model training API.
-
+        
         Args:
             model_name (str): The name of the model to train.
-            training_data (dict): Dictionary containing training data parameters.
-
+            training_dataset (str): The training dataset identifier.
+            task_type (str): The type of ML task (e.g., regression, classification).
+            model_category (str): The category of model (e.g., random_forest).
+            architecture (str): The model architecture.
+            region (str): The region identifier.
+            hyperparameters (dict, optional): Additional hyperparameters for training.
+            
         Returns:
             dict: The response from the model training API.
         """
-        endpoint = "https://modeltraining-573248941006.australia-southeast1.run.app/train_model"
         payload = {
             "model_name": model_name,
-            "training_data": training_data
+            "training_dataset": training_dataset,
+            "task_type": task_type,
+            "model_category": model_category,
+            "architecture": architecture,
+            "region": region,
+            "hyperparameters": hyperparameters
         }
+        endpoint = f"{self.url.rstrip('/')}/train_model"
+        print("the payload is ", payload)
         try:
             response = self.session.post(endpoint, json=payload, timeout=self.timeout, verify=self.verify)
             if not response.ok:
@@ -1155,35 +1190,163 @@ class BaseClient:
         except requests.RequestException as e:
             raise APIError(f"Model training request failed: {str(e)}")
 
-    def deploy_model(self, dataset: str, product:str, model_name:str, input_expression: str, model_training_job_name: str, uid: str, dates_iso8601: list):
-        # we have the dataset and we have the product, and we have the model name, we need to create a new json file and add that to the dataset as our virtual dataset
-        # upload the script to the bucket, the script should be able to download the model and do the inferencing
-        # we need to upload the the json to the to the dataset as our virtual dataset
-        # then we do nothing and wait for the user to make the request call to the explorer
-        # we should have a uniform script for the random forest deployment
-        # create a script for each model
-        # upload script to google bucket, 
-        # 
+    # Mass Stats methods
+    def combine_tiles(self, 
+        data_name: str,
+        usezarr: bool,
+        overwrite: bool,
+        output : str) -> dict:
+
+        if not self.mass_stats:
+            from terrakio_core.mass_stats import MassStats
+            if not self.url or not self.key:
+                raise ConfigurationError("Mass Stats client not initialized. Make sure API URL and key are set.")
+            self.mass_stats = MassStats(
+                base_url=self.url,
+                api_key=self.key,
+                verify=self.verify,
+                timeout=self.timeout
+            )
+        return self.mass_stats.combine_tiles(data_name, usezarr, overwrite, output)
+    
+
+
+    def create_dataset_file(
+        self,
+        name: str,
+        aoi: str,
+        expression: str,
+        output: str,
+        tile_size: float = 128.0,
+        crs: str = "epsg:4326",
+        res: float = 0.0001,
+        region: str = "aus",
+        to_crs: str = "epsg:4326",
+        overwrite: bool = True,
+        skip_existing: bool = False,
+        non_interactive: bool = True,
+        usezarr: bool = False,
+        poll_interval: int = 30   # seconds between job status checks
+    ) -> dict:
+        
+        from terrakio_core.generation.tiles import tiles
+        import tempfile
+        import time
+
+        body, reqs, groups = tiles(
+            name = name, 
+            aoi = aoi, 
+            expression = expression,
+            output = output,
+            tile_size = tile_size,
+            crs = crs,
+            res = res,
+            region = region,
+            to_crs = to_crs,
+            fully_cover = True,
+            overwrite = overwrite,
+            skip_existing = skip_existing,
+            non_interactive = non_interactive
+        )
+
+        # Create temp json files before upload
+        with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False) as tempreq:
+            tempreq.write(reqs)
+            tempreqname = tempreq.name
+        with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False) as tempmanifest:
+            tempmanifest.write(groups)
+            tempmanifestname = tempmanifest.name
+        
+        if not self.mass_stats:
+            from terrakio_core.mass_stats import MassStats
+            if not self.url or not self.key:
+                raise ConfigurationError("Mass Stats client not initialized. Make sure API URL and key are set.")
+            self.mass_stats = MassStats(
+                base_url=self.url,
+                api_key=self.key,
+                verify=self.verify,
+                timeout=self.timeout
+            )
 
+        task_id = self.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,
+            manifest_json=tempmanifestname,
+        )
+
+        ### Start combining tiles when generation-tiles job is done
+        start_time = time.time()
+        status = None
+        
+        while True:
+            try:
+                taskid = task_id['task_id']
+                trackinfo = self.mass_stats.track_job([taskid])
+                status = trackinfo[taskid]['status']
+                
+                # Check completion states
+                if status == 'Completed':
+                    print('Tiles generated successfully!')
+                    break
+                elif status in ['Failed', 'Cancelled', 'Error']:
+                    raise RuntimeError(f"Job {taskid} failed with status: {status}")
+                else:
+                    # Job is still running
+                    elapsed_time = time.time() - start_time
+                    print(f"Job status: {status} - Elapsed time: {elapsed_time:.1f}s", end='\r')
+                    
+                    # Sleep before next check
+                    time.sleep(poll_interval)
+                    
+                    
+            except KeyboardInterrupt:
+                print(f"\nInterrupted! Job {taskid} is still running in the background.")
+                raise
+            except Exception as e:
+                print(f"\nError tracking job: {e}")
+                raise
+
+        # Clean up temporary files
+        import os
+        os.unlink(tempreqname)
+        os.unlink(tempmanifestname)
+
+
+        # Start combining tiles
+        if not self.mass_stats:
+            from terrakio_core.mass_stats import MassStats
+            if not self.url or not self.key:
+                raise ConfigurationError("Mass Stats client not initialized. Make sure API URL and key are set.")
+            self.mass_stats = MassStats(
+                base_url=self.url,
+                api_key=self.key,
+                verify=self.verify,
+                timeout=self.timeout
+            )
+        
+        return self.mass_stats.combine_tiles(body["name"], usezarr, body["overwrite"], body["output"])
+
+    def deploy_model(self, dataset: str, product:str, model_name:str, input_expression: str, model_training_job_name: str, uid: str, dates_iso8601: list):
         script_content = self._generate_script(model_name, product, model_training_job_name, uid)
-        # self.create_dataset(collection = "terrakio-datasets", input = input, )
-        # we have the script, we need to upload it to the bucket
         script_name = f"{product}.py"
-        print("the script content is ", script_content)
-        print("the script name is ", script_name)
         self._upload_script_to_bucket(script_content, script_name, model_training_job_name, uid)
-        # after uploading the script, we need to create a new virtual dataset
         self._create_dataset(name = dataset, collection = "terrakio-datasets", products = [product], path = f"gs://terrakio-mass-requests/{uid}/{model_training_job_name}/inference_scripts", input = input_expression, dates_iso8601 = dates_iso8601, padding = 0)
 
     def _generate_script(self, model_name: str, product: str, model_training_job_name: str, uid: str) -> str:
         return textwrap.dedent(f'''
             import logging
             from io import BytesIO
-            from google.cloud import storage
-            from onnxruntime import InferenceSession
+
             import numpy as np
+            import pandas as pd
             import xarray as xr
-            import datetime
+            from google.cloud import storage
+            from onnxruntime import InferenceSession
 
             logging.basicConfig(
                 level=logging.INFO
@@ -1191,54 +1354,122 @@ class BaseClient:
 
             def get_model():
                 logging.info("Loading model for {model_name}...")
-                
+
                 client = storage.Client()
                 bucket = client.get_bucket('terrakio-mass-requests')
                 blob = bucket.blob('{uid}/{model_training_job_name}/models/{model_name}.onnx')
-                
+
                 model = BytesIO()
                 blob.download_to_file(model)
                 model.seek(0)
-                
+
                 session = InferenceSession(model.read(), providers=["CPUExecutionProvider"])
                 return session
 
             def {product}(*bands, model):
                 logging.info("start preparing data")
+                print("the bands are ", bands)
                 
-                original_shape = bands[0].shape
-                logging.info(f"Original shape: {{original_shape}}")
+                data_arrays = list(bands)
                 
-                transformed_bands = []
-                for band in bands:
-                    transformed_band = band.values.reshape(-1,1)
-                    transformed_bands.append(transformed_band)
+                print("the data arrays are ", [da.name for da in data_arrays])
                 
-                input_data = np.hstack(transformed_bands)
+                reference_array = data_arrays[0]
+                original_shape = reference_array.shape
+                logging.info(f"Original shape: {{original_shape}}")
                 
+                if 'time' in reference_array.dims:
+                    time_coords = reference_array.coords['time']
+                    if len(time_coords) == 1:
+                        output_timestamp = time_coords[0]
+                    else:
+                        years = [pd.to_datetime(t).year for t in time_coords.values]
+                        unique_years = set(years)
+                        
+                        if len(unique_years) == 1:
+                            year = list(unique_years)[0]
+                            output_timestamp = pd.Timestamp(f"{{year}}-01-01")
+                        else:
+                            latest_year = max(unique_years)
+                            output_timestamp = pd.Timestamp(f"{{latest_year}}-01-01")
+                else:
+                    output_timestamp = pd.Timestamp("1970-01-01")
+
+                averaged_bands = []
+                for data_array in data_arrays:
+                    if 'time' in data_array.dims:
+                        averaged_band = np.mean(data_array.values, axis=0)
+                        logging.info(f"Averaged band from {{data_array.shape}} to {{averaged_band.shape}}")
+                    else:
+                        averaged_band = data_array.values
+                        logging.info(f"No time dimension, shape: {{averaged_band.shape}}")
+
+                    flattened_band = averaged_band.reshape(-1, 1)
+                    averaged_bands.append(flattened_band)
+
+                input_data = np.hstack(averaged_bands)
+
                 logging.info(f"Final input shape: {{input_data.shape}}")
-                
+
                 output = model.run(None, {{"float_input": input_data.astype(np.float32)}})[0]
-                
+
                 logging.info(f"Model output shape: {{output.shape}}")
 
-                output_reshaped = output.reshape(original_shape)
+                if len(original_shape) >= 3:
+                    spatial_shape = original_shape[1:]
+                else:
+                    spatial_shape = original_shape
+
+                output_reshaped = output.reshape(spatial_shape)
+
+                output_with_time = np.expand_dims(output_reshaped, axis=0)
+
+                if 'time' in reference_array.dims:
+                    spatial_dims = [dim for dim in reference_array.dims if dim != 'time']
+                    spatial_coords = {{dim: reference_array.coords[dim] for dim in spatial_dims if dim in reference_array.coords}}
+                else:
+                    spatial_dims = list(reference_array.dims)
+                    spatial_coords = dict(reference_array.coords)
+
                 result = xr.DataArray(
-                    data=output_reshaped,
-                    dims=bands[0].dims,
-                    coords=bands[0].coords
+                    data=output_with_time.astype(np.float32),
+                    dims=['time'] + list(spatial_dims),
+                    coords={
+                        'time': [output_timestamp.values],
+                        'y': spatial_coords['y'].values,
+                        'x': spatial_coords['x'].values
+                    }
                 )
-                
                 return result
             ''').strip()
-
+    
     def _upload_script_to_bucket(self, script_content: str, script_name: str, model_training_job_name: str, uid: str):
         """Upload the generated script to Google Cloud Storage"""
 
         client = storage.Client()
         bucket = client.get_bucket('terrakio-mass-requests')
         blob = bucket.blob(f'{uid}/{model_training_job_name}/inference_scripts/{script_name}')
-        # the first layer is the uid, the second layer is the model training job name
         blob.upload_from_string(script_content, content_type='text/plain')
         logging.info(f"Script uploaded successfully to {uid}/{model_training_job_name}/inference_scripts/{script_name}")
 
+
+
+
+    def download_file_to_path(self, job_name, stage, file_name, output_path):
+        if not self.mass_stats:
+            from terrakio_core.mass_stats import MassStats
+            if not self.url or not self.key:
+                raise ConfigurationError("Mass Stats client not initialized. Make sure API URL and key are set.")
+            self.mass_stats = MassStats(
+                base_url=self.url,
+                api_key=self.key,
+                verify=self.verify,
+                timeout=self.timeout
+            )
+
+        # fetch bucket info based on job name and stage
+
+        taskid = self.mass_stats.get_task_id(job_name, stage).get('task_id')
+        trackinfo = self.mass_stats.track_job([taskid])
+        bucket = trackinfo[taskid]['bucket']
+        return self.mass_stats.download_file(job_name, bucket, file_name, output_path)

terrakio_core/dataset_management.py

@@ -83,10 +83,63 @@ class DatasetManagement:
         except requests.RequestException as e:
             raise APIError(f"Request failed: {str(e)}")
 
+    # def create_dataset(self, name: str, collection: str = "terrakio-datasets", **kwargs) -> Dict[str, Any]:
+    #     """
+    #     Create a new dataset.
+        
+    #     Args:
+    #         name: Name of the dataset (required)
+    #         collection: Dataset collection (default: 'terrakio-datasets')
+    #         **kwargs: Additional dataset parameters including:
+    #             - products: List of products
+    #             - dates_iso8601: List of dates
+    #             - bucket: Storage bucket
+    #             - path: Storage path
+    #             - data_type: Data type
+    #             - no_data: No data value
+    #             - l_max: Maximum level
+    #             - y_size: Y size
+    #             - x_size: X size
+    #             - proj4: Projection string
+    #             - abstract: Dataset abstract
+    #             - geotransform: Geotransform parameters
+                
+    #     Returns:
+    #         Created dataset information
+            
+    #     Raises:
+    #         APIError: If the API request fails
+    #     """
+    #     endpoint = f"{self.api_url}/datasets"
+    #     params = {"collection": collection}
+    #     # Create payload with required name parameter
+    #     payload = {"name": name}
+        
+    #     # Add optional parameters if provided
+    #     for param in ["products", "dates_iso8601", "bucket", "path", "data_type", 
+    #                  "no_data", "l_max", "y_size", "x_size", "proj4", "abstract", "geotransform", "input"]:
+    #         if param in kwargs:
+    #             payload[param] = kwargs[param]
+        
+    #     try:
+    #         response = self.session.post(
+    #             endpoint,
+    #             params=params,
+    #             json=payload,
+    #             timeout=self.timeout,
+    #             verify=self.verify
+    #         )
+            
+    #         if not response.ok:
+    #             raise APIError(f"API request failed: {response.status_code} {response.reason}")
+    #         return response.json()
+    #     except requests.RequestException as e:
+    #         raise APIError(f"Request failed: {str(e)}")
+
     def create_dataset(self, name: str, collection: str = "terrakio-datasets", **kwargs) -> Dict[str, Any]:
         """
         Create a new dataset.
-        
+                    
         Args:
             name: Name of the dataset (required)
             collection: Dataset collection (default: 'terrakio-datasets')
@@ -103,24 +156,23 @@ class DatasetManagement:
                 - proj4: Projection string
                 - abstract: Dataset abstract
                 - geotransform: Geotransform parameters
-                
+                - padding: Padding value
+                            
         Returns:
             Created dataset information
-            
+                        
         Raises:
             APIError: If the API request fails
         """
         endpoint = f"{self.api_url}/datasets"
         params = {"collection": collection}
-        # Create payload with required name parameter
         payload = {"name": name}
-        
-        # Add optional parameters if provided
-        for param in ["products", "dates_iso8601", "bucket", "path", "data_type", 
-                     "no_data", "l_max", "y_size", "x_size", "proj4", "abstract", "geotransform", "input"]:
+                    
+        for param in ["products", "dates_iso8601", "bucket", "path", "data_type",
+                        "no_data", "l_max", "y_size", "x_size", "proj4", "abstract", "geotransform", "input", "padding"]:
             if param in kwargs:
                 payload[param] = kwargs[param]
-        
+                    
         try:
             response = self.session.post(
                 endpoint,
@@ -129,7 +181,7 @@ class DatasetManagement:
                 timeout=self.timeout,
                 verify=self.verify
             )
-            
+                            
             if not response.ok:
                 raise APIError(f"API request failed: {response.status_code} {response.reason}")
             return response.json()

terrakio_core/decorators.py

@@ -0,0 +1,18 @@
+# terrakio_core/decorators.py
+def admin_only_params(*restricted_params):
+    """
+    Decorator factory for restricting method parameters to admin users only.
+    """
+    def decorator(func):
+        def wrapper(self, *args, **kwargs):
+            if hasattr(self, '_is_admin') and self._is_admin:
+                return func(self, *args, **kwargs)
+            
+            admin_params_used = set(kwargs.keys()) & set(restricted_params)
+            if admin_params_used:
+                raise PermissionError(f"Parameters {admin_params_used} are only available to admin users")
+            
+            filtered_kwargs = {k: v for k, v in kwargs.items() if k not in restricted_params}
+            return func(self, *args, **filtered_kwargs)
+        return wrapper
+    return decorator

terrakio_core/generation/tiles.py

@@ -0,0 +1,95 @@
+### implementing generation-tiles in python api
+### function should just generate the json file for mass_stats to pick up.
+
+import geopandas as gpd
+import shapely.geometry
+import json
+from rich import print
+
+def escape_newline(string):
+    if isinstance(string, list):
+        return [s.replace('\\n', '\n') for s in string]
+    else:
+        return string.replace('\\n', '\n')
+
+def get_bounds(aoi, crs, to_crs = None):
+    aoi : gpd.GeoDataFrame = gpd.read_file(aoi)
+    aoi = aoi.set_crs(crs, allow_override=True)
+    if to_crs:
+        aoi = aoi.to_crs(to_crs)
+    bounds = aoi.geometry[0].bounds
+    return *bounds, aoi
+
+def tile_generator(x_min, y_min, x_max, y_max, aoi, crs, res, tile_size, expression, output, fully_cover=True):
+    i_max = int((x_max-x_min)/(tile_size*res))
+    j_max = int((y_max-y_min)/(tile_size*res))
+    if fully_cover:
+        i_max += 1
+        j_max += 1
+    for j in range(0, int(j_max)):
+        for i in range(0, int(i_max)):
+            #print(f"Processing tile {i} {j}")
+            x = x_min + i*(tile_size*res)
+            y = y_max - j*(tile_size*res)
+            bbox = shapely.geometry.box(x, y-(tile_size*res), x + (tile_size*res), y)
+            if not aoi.geometry[0].intersects(bbox):
+                continue
+            feat  = {"type": "Feature", "geometry": bbox.__geo_interface__}
+            data = {
+                "feature": feat,
+                "in_crs": crs,
+                "out_crs": crs,
+                "resolution": res,
+                "expr" : expression,
+                "output" : output,
+            }
+            yield data, i , j
+
+
+def tiles(
+    name: str,
+    aoi : str, 
+    expression: str = "red=S2v2#(year,median).red@(year =2024) \n red",
+    output: str = "netcdf",
+    tile_size : float = 512,
+    crs : str = "epsg:3577",
+    res: float = 10,
+    region : str = "eu",
+    to_crs: str = None,
+    fully_cover: bool = True,
+    overwrite: bool = False,
+    skip_existing: bool = False,
+    non_interactive: bool = False,
+):
+    
+    # Create requests for each tile
+    reqs = []
+    x_min, y_min, x_max, y_max, aoi = get_bounds(aoi, crs, to_crs)
+    #print(f"Bounds: {x_min}, {y_min}, {x_max}, {y_max}")
+
+    if to_crs is None:
+        to_crs = crs
+    for tile_req, i, j in tile_generator(x_min, y_min, x_max, y_max, aoi, to_crs, res, tile_size, expression, output, fully_cover):
+        req_name = f"{name}_{i:02d}_{j:02d}"
+        reqs.append({"group": "tiles", "file": req_name, "request": tile_req})
+
+    #print(f"Generated {len(reqs)} tile requests.")
+
+
+    count = len(reqs)
+    groups = list(set(dic["group"] for dic in reqs))
+
+    body = {
+        "name" : name,
+        "output" : output,
+        "region" : region,
+        "size" : count,
+        "overwrite" : overwrite,
+        "non_interactive": non_interactive,
+        "skip_existing" : skip_existing,
+    }
+    request_json = json.dumps(reqs)
+    manifest_json = json.dumps(groups)
+
+    return body, request_json, manifest_json
+    

terrakio_core/mass_stats.py

@@ -1,5 +1,8 @@
 import requests
-from typing import Optional, Dict, Any
+from typing import Optional, Dict, Any, List
+import json
+import json as json_lib
+import gzip
 
 class MassStats:
     def __init__(self, base_url: str, api_key: str, verify: bool = True, timeout: int = 60):
@@ -12,61 +15,272 @@ class MassStats:
             'x-api-key': self.api_key
         })
 
+    def _upload_file(self, file_path: str, url: str, use_gzip: bool = False):
+        """
+        Helper method to upload a JSON file to a signed URL.
+        
+        Args:
+            file_path: Path to the JSON file
+            url: Signed URL to upload to
+            use_gzip: Whether to compress the file with gzip
+        """
+        try:
+            with open(file_path, 'r') as file:
+                json_data = json_lib.load(file)
+        except FileNotFoundError:
+            raise FileNotFoundError(f"JSON file not found: {file_path}")
+        except json.JSONDecodeError as e:
+            raise ValueError(f"Invalid JSON in file {file_path}: {e}")
+        
+        # Check if using simplejson and support ignore_nan
+        if hasattr(json_lib, 'dumps') and 'ignore_nan' in json_lib.dumps.__code__.co_varnames:
+            dumps_kwargs = {'ignore_nan': True}
+        else:
+            dumps_kwargs = {}
+        
+        if use_gzip:
+            # Serialize and compress the JSON data
+            body = gzip.compress(json_lib.dumps(json_data, **dumps_kwargs).encode('utf-8'))
+            headers = {
+                'Content-Type': 'application/json',
+                'Content-Encoding': 'gzip'
+            }
+        else:
+            body = json_lib.dumps(json_data, **dumps_kwargs).encode('utf-8')
+            headers = {
+                'Content-Type': 'application/json'
+            }
+        
+        # Make the PUT request to the signed URL
+        response = requests.put(
+            url,
+            data=body,
+            headers=headers
+        )
+        
+        return response
+    
+
+    def download_file(self, job_name: str, bucket:str, file_name: str, output_path: str) -> str:
+        """
+        Download a file from mass_stats using job name and file name.
+        
+        Args:
+            job_name: Name of the job
+            file_name: Name of the file to download
+            output_path: Path where the file should be saved
+            
+        Returns:
+            str: Path to the downloaded file
+        """
+        import os
+        from pathlib import Path
+        
+        endpoint_url = f"{self.base_url}/mass_stats/download_files"
+        request_body = {
+            "job_name": job_name,
+            "bucket": bucket,
+            "file_name": file_name
+        }
+        
+        try:
+            # Get signed URL
+            response = self.session.post(
+                endpoint_url, 
+                json=request_body, 
+                verify=self.verify, 
+                timeout=self.timeout
+            )
+            signed_url = response.json().get('download_url')
+            if not signed_url:
+                raise Exception("No download URL received from server")    
+            print(f"Generated signed URL for download")
+            
+            # Create output directory if it doesn't exist
+            output_dir = Path(output_path).parent
+            output_dir.mkdir(parents=True, exist_ok=True)
+            
+            # Download the file using the signed URL
+            download_response = self.session.get(
+                signed_url,
+                verify=self.verify,
+                timeout=self.timeout,
+                stream=True  # Stream for large files
+            )
+            download_response.raise_for_status()
+
+            # Check if file exists in the response (content-length header)
+            content_length = download_response.headers.get('content-length')
+            if content_length and int(content_length) == 0:
+                raise Exception("File appears to be empty")
+            
+            # Write the file
+            with open(output_path, 'wb') as file:
+                for chunk in download_response.iter_content(chunk_size=8192):
+                    if chunk:
+                        file.write(chunk)
+            
+            # Verify file was written
+            if not os.path.exists(output_path):
+                raise Exception(f"File was not written to {output_path}")
+                
+            file_size = os.path.getsize(output_path)
+            print(f"File downloaded successfully to {output_path} (size: {file_size / (1024 * 1024):.4f} mb)")
+            
+            return output_path
+            
+        except self.session.exceptions.RequestException as e:
+            if hasattr(e, 'response') and e.response is not None:
+                error_detail = e.response.text
+                raise Exception(f"Error getting signed URL: {e}. Details: {error_detail}")
+            raise Exception(f"Error in download process: {e}")
+        except IOError as e:
+            raise Exception(f"Error writing file to {output_path}: {e}")
+        except Exception as e:
+            # Clean up partial file if it exists
+            if os.path.exists(output_path):
+                try:
+                    os.remove(output_path)
+                except:
+                    pass
+            raise
+    
+
+
+
     def upload_request(
         self,
         name: str,
         size: int,
-        bucket: str,
+        region: List[str],
         output: str,
+        config: Dict[str, Any],
         location: Optional[str] = None,
-        force_loc: bool = False,
-        config: Optional[Dict[str, Any]] = None,
+        force_loc: Optional[bool] = None,
         overwrite: bool = False,
         server: Optional[str] = None,
-        skip_existing: bool = False
+        skip_existing: bool = False,
     ) -> Dict[str, Any]:
         """
         Initiate a mass stats upload job.
         
         Args:
             name: Name of the job
-            size: Size of the data
-            bucket: Storage bucket
-            output: Output path or identifier
+            size: Size of the job
+            region: Region to run job [aus, eu, us]
+            output: Output type
+            config: Configuration dictionary
             location: (Optional) Location for the upload
             force_loc: Force location usage
-            config: Optional configuration dictionary
             overwrite: Overwrite existing data
             server: Optional server
             skip_existing: Skip existing files
         """
+        
+
+        
+        # Step 2: Create the upload job and get signed URLs
         url = f"{self.base_url}/mass_stats/upload"
+        
         data = {
             "name": name,
             "size": size,
-            "bucket": bucket,
+            "region": region,
             "output": output,
-            "force_loc": force_loc,
+            "config": config,
             "overwrite": overwrite,
             "skip_existing": skip_existing
         }
+        
         if location is not None:
             data["location"] = location
-        if config is not None:
-            data["config"] = config
+        if force_loc is not None:
+            data["force_loc"] = force_loc
         if server is not None:
             data["server"] = server
-        response = self.session.post(url, json=data, verify=self.verify, timeout=self.timeout)
-        print("the response is ", response.text)
-        # response.raise_for_status()
+        response = self.session.post(
+            url, 
+            json=data, 
+            verify=self.verify, 
+            timeout=self.timeout
+        )      
         return response.json()
 
+
+
+    
+    def execute_job(
+        self,
+        name: str,
+        region: str,
+        output: str,
+        config: Dict[str, Any],
+        overwrite: bool = False,
+        skip_existing: bool = False,
+        request_json: Optional[str] = None,
+        manifest_json: Optional[str] = None,
+        location: Optional[str] = None,
+        force_loc: Optional[bool] = None,
+        server: Optional[str] = None
+    ) -> Dict[str, Any]:
+        # Step 1: Calculate size from request JSON file if provided
+        size = 0
+        if request_json is not None:
+            try:
+                with open(request_json, 'r') as file:
+                    request_data = json_lib.load(file)
+                
+                if isinstance(request_data, list):
+                    size = len(request_data)
+                else:
+                    raise ValueError(f"Request JSON file {request_json} should contain a list of dictionaries")
+                    
+            except FileNotFoundError:
+                raise FileNotFoundError(f"Request JSON file not found: {request_json}")
+            except json.JSONDecodeError as e:
+                raise ValueError(f"Invalid JSON in request file {request_json}: {e}")
+            
+        upload_result = self.upload_request(name, size, region, output, config, location, force_loc, overwrite, server, skip_existing)
+        
+        # Step 3: Upload JSON files if provided
+        if request_json is not None or manifest_json is not None:
+            requests_url = upload_result.get('requests_url')
+            manifest_url = upload_result.get('manifest_url')
+            
+            if request_json is not None:
+                if not requests_url:
+                    raise ValueError("No requests_url returned from server for request JSON upload")
+                
+                try:
+                    requests_response = self._upload_file(request_json, requests_url, use_gzip=True)
+                    if requests_response.status_code not in [200, 201, 204]:
+                        print(f"Requests upload error: {requests_response.text}")
+                        raise Exception(f"Failed to upload request JSON: {requests_response.text}")
+                except Exception as e:
+                    raise Exception(f"Error uploading request JSON file {request_json}: {e}")
+            
+            if manifest_json is not None:
+                if not manifest_url:
+                    raise ValueError("No manifest_url returned from server for manifest JSON upload")
+                
+                try:
+                    manifest_response = self._upload_file(manifest_json, manifest_url, use_gzip=False)
+                    if manifest_response.status_code not in [200, 201, 204]:
+                        print(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 =self.start_job(upload_result.get("id"))
+        return start_job_task_id
+
+
     def start_job(self, task_id: str) -> Dict[str, Any]:
         """
         Start a mass stats job by task ID.
         """
         url = f"{self.base_url}/mass_stats/start/{task_id}"
-        print("the self session header is ", self.session.headers)
         response = self.session.post(url, verify=self.verify, timeout=self.timeout)
         response.raise_for_status()
         return response.json()
@@ -79,7 +293,7 @@ class MassStats:
         if uid is not None:
             url += f"&uid={uid}"
         response = self.session.get(url, verify=self.verify, timeout=self.timeout)
-        print("response text is ", response.text)
+        #print("response text is ", response.text)
         return response.json()
 
     def track_job(self, ids: Optional[list] = None) -> Dict[str, Any]:
@@ -259,4 +473,32 @@ class MassStats:
         print("Response text:", response.text)
         # response.raise_for_status()
         return response.json() 
+
+
+    ### Adding the wrapper function to call endpoint /mass_stats/combine_tiles
+    def combine_tiles(
+        self, 
+        data_name: str,
+        usezarr: bool = False,
+        overwrite: bool = True,
+        output : str = "netcdf"
+    ) -> Dict[str, Any]:
+
+        url = f"{self.base_url}/mass_stats/combine_tiles"
+        request_body = {
+            'data_name': data_name,
+            'usezarr': str(usezarr).lower(),
+            'output': output,
+            'overwrite': str(overwrite).lower()
+        }
+        print(f"Request body: {json.dumps(request_body, indent=2)}")
+        response = self.session.post(url, json=request_body, verify=self.verify, timeout=self.timeout)
+        print(f"Response text: {response.text}")
+        return response.json()
+
+
+
+  
+
+
     

{terrakio_core-0.3.0.dist-info → terrakio_core-0.3.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: terrakio-core
-Version: 0.3.0
+Version: 0.3.2
 Summary: Core components for Terrakio API clients
 Author-email: Yupeng Chao <yupeng@haizea.com.au>
 Project-URL: Homepage, https://github.com/HaizeaAnalytics/terrakio-python-api
@@ -22,6 +22,7 @@ Requires-Dist: xarray>=2023.1.0
 Requires-Dist: shapely>=2.0.0
 Requires-Dist: geopandas>=0.13.0
 Requires-Dist: google-cloud-storage>=2.0.0
+Requires-Dist: nest_asyncio
 
 # Terrakio Core
 

terrakio_core-0.3.2.dist-info/RECORD

@@ -0,0 +1,16 @@
+terrakio_core/__init__.py,sha256=faOeYeL7Lmg3aTRVxVEBT6Dbhi62N4eeXK3mAluD3pA,88
+terrakio_core/auth.py,sha256=Nuj0_X3Hiy17svYgGxrSAR-LXpTlP0J0dSrfMnkPUbI,7717
+terrakio_core/client.py,sha256=oJPg4hcdCBaPwYN0SKQy57YUsJBTTA-SbQ4aaHaNq8E,65578
+terrakio_core/config.py,sha256=AwJ1VgR5K7N32XCU5k7_Dp1nIv_FYt8MBonq9yKlGzA,2658
+terrakio_core/dataset_management.py,sha256=Hdk3nkwd70jw3lBNEaGixrqNVhUWOmsIYktzm_8vXdc,10913
+terrakio_core/decorators.py,sha256=QeNOUX6WEAmdgBL5Igt5DXyYduh3jnmLbodttmwvXhE,785
+terrakio_core/exceptions.py,sha256=9S-I20-QiDRj1qgjFyYUwYM7BLic_bxurcDOIm2Fu_0,410
+terrakio_core/group_access_management.py,sha256=NJ7SX4keUzZAUENmJ5L6ynKf4eRlqtyir5uoKFyY17A,7315
+terrakio_core/mass_stats.py,sha256=UGZo8BH4hzWe3k7pevsYAdRwnVZl-08lXjTlHD4nMgQ,18212
+terrakio_core/space_management.py,sha256=wlUUQrlj_4U_Lpjn9lbF5oj0Rv3NPvvnrd5mWej5kmA,4211
+terrakio_core/user_management.py,sha256=MMNWkz0V_9X7ZYjjteuRU4H4W3F16iuQw1dpA2wVTGg,7400
+terrakio_core/generation/tiles.py,sha256=eiiMNzqaga-c42kG_7zHXTF2o8ZInCPUj0Vu4Ye30Ts,2980
+terrakio_core-0.3.2.dist-info/METADATA,sha256=QaQJrukRnVgy2jAqNr-BDCUaybrqFMbjOlr2KqLlhJI,1476
+terrakio_core-0.3.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+terrakio_core-0.3.2.dist-info/top_level.txt,sha256=5cBj6O7rNWyn97ND4YuvvXm0Crv4RxttT4JZvNdOG6Q,14
+terrakio_core-0.3.2.dist-info/RECORD,,

terrakio_core-0.3.0.dist-info/RECORD

@@ -1,14 +0,0 @@
-terrakio_core/__init__.py,sha256=iguSJomKouzVNPOB2_Ox-FGnQBUQ0ykx8CshjVzU1QM,88
-terrakio_core/auth.py,sha256=Nuj0_X3Hiy17svYgGxrSAR-LXpTlP0J0dSrfMnkPUbI,7717
-terrakio_core/client.py,sha256=CQ1qiR_8tWKEGX-UT2wLeatk8fYMpyo9KseMpCapw7c,56813
-terrakio_core/config.py,sha256=AwJ1VgR5K7N32XCU5k7_Dp1nIv_FYt8MBonq9yKlGzA,2658
-terrakio_core/dataset_management.py,sha256=LKUESSDPRu1JubQaQJWdPqHLGt-_Xv77Fpb4IM7vkzM,8751
-terrakio_core/exceptions.py,sha256=9S-I20-QiDRj1qgjFyYUwYM7BLic_bxurcDOIm2Fu_0,410
-terrakio_core/group_access_management.py,sha256=NJ7SX4keUzZAUENmJ5L6ynKf4eRlqtyir5uoKFyY17A,7315
-terrakio_core/mass_stats.py,sha256=AqYJsd6nqo2BDh4vEPUDgsv4T0UR1_TPDoXa3WO3gTU,9284
-terrakio_core/space_management.py,sha256=wlUUQrlj_4U_Lpjn9lbF5oj0Rv3NPvvnrd5mWej5kmA,4211
-terrakio_core/user_management.py,sha256=MMNWkz0V_9X7ZYjjteuRU4H4W3F16iuQw1dpA2wVTGg,7400
-terrakio_core-0.3.0.dist-info/METADATA,sha256=8mS_NJQUoFcr1lE3iUQXQi5VwSZo07t3XF0pCL7VNSI,1448
-terrakio_core-0.3.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-terrakio_core-0.3.0.dist-info/top_level.txt,sha256=5cBj6O7rNWyn97ND4YuvvXm0Crv4RxttT4JZvNdOG6Q,14
-terrakio_core-0.3.0.dist-info/RECORD,,