terrakio-core 0.4.98__py3-none-any.whl → 0.4.98.1b1__py3-none-any.whl

terrakio_core/endpoints/dataset_management.py

@@ -1,51 +1,68 @@
-from typing import Dict, Any, List, Optional
-from ..helper.decorators import require_token, require_api_key, require_auth
+from typing import Any, Dict, List, Optional
+
+from ..exceptions import (
+    CommandPermissionError,
+    DatasetNotFoundError,
+    DatasetPermissionError,
+    GetDatasetError,
+    ListDatasetsError,
+    CreateDatasetError,
+    DatasetAlreadyExistsError,
+    DeleteDatasetError,
+    OverwriteDatasetError,
+)
+from ..helper.decorators import require_api_key, require_auth, require_token
 
 class DatasetManagement:
     def __init__(self, client):
         self._client = client
 
-
     @require_api_key
-    def list_datasets(self, substring: Optional[str] = None, collection: str = "terrakio-datasets") -> List[Dict[str, Any]]:
+    async def list_datasets(self, substring: Optional[str] = None) -> List[Dict[str, Any]]:
         """
         List datasets, optionally filtering by a substring and collection.
         
         Args:
             substring: Substring to filter by (optional)
-            collection: Dataset collection (default: 'terrakio-datasets')
             
         Returns:
             List of datasets matching the criteria
         """
-        params = {"collection": collection}
-        if substring:
-            params.update({"substring": substring})
-        return self._client._terrakio_request("GET", "/datasets", params = params)
+        params = {"substring": substring} if substring else None
+        response, status = await self._client._terrakio_request("GET", "/datasets", params=params)
+        if status != 200:
+            raise ListDatasetsError(f"List datasets failed with status {status}", status_code=status)
+        return response
 
     @require_api_key
-    def get_dataset(self, name: str, collection: str = "terrakio-datasets") -> Dict[str, Any]:
+    async def get_dataset(self, name: str) -> Dict[str, Any]:
         """
         Retrieve dataset info by dataset name.
         
         Args:
             name: The name of the dataset (required)
-            collection: The dataset collection (default: 'terrakio-datasets')
             
         Returns:
             Dataset information as a dictionary
             
         Raises:
-            APIError: If the API request fails
+            GetDatasetError: If the API request fails
+            DatasetNotFoundError: If the dataset is not found
+            DatasetPermissionError: If the user does not have permission to get the dataset
         """
-        params = {"collection": collection}
-        return self._client._terrakio_request("GET", f"/datasets/{name}", params = params)
+        response, status = await self._client._terrakio_request("GET", f"/datasets/{name}")
+        if status != 200:
+            if status == 404:
+                raise DatasetNotFoundError(f"Dataset {name} not found", status_code = status)
+            if status == 403:
+                raise DatasetPermissionError(f"You do not have permission to get dataset {name}", status_code = status)
+            raise GetDatasetError(f"Get dataset failed with status {status}", status_code = status)
+        return response
 
     @require_api_key
     async def create_dataset(
         self, 
         name: str, 
-        collection: str = "terrakio-datasets",
         products: Optional[List[str]] = None,
         dates_iso8601: Optional[List[str]] = None,
         bucket: Optional[str] = None,
@@ -67,7 +84,6 @@ class DatasetManagement:
                             
         Args:
             name: Name of the dataset (required)
-            collection: Dataset collection (default: 'terrakio-datasets')
             products: List of products
             dates_iso8601: List of dates (will be automatically sorted chronologically)
             bucket: Storage bucket
@@ -89,7 +105,6 @@ class DatasetManagement:
         Raises:
             APIError: If the API request fails
         """
-        params = {"collection": collection}
         payload = {"name": name}
         param_mapping = {
             "products": products,
@@ -111,14 +126,42 @@ class DatasetManagement:
         for param, value in param_mapping.items():
             if value is not None:
                 payload[param] = value
-        return await self._client._terrakio_request("POST", "/datasets", params = params, json = payload)
-    
+        response, status = await self._client._terrakio_request("POST", "/datasets", json = payload)
+        if status != 200:
+            if status ==403:
+                raise CommandPermissionError(f"You do not have permission to create dataset {name}", status_code = status)
+            elif status == 409:
+                raise DatasetAlreadyExistsError(f"Dataset {name} already exists", status_code = status)
+            raise CreateDatasetError(f"Create dataset failed with status {status}", status_code = status)
+        return response
+        
+    @require_api_key
+    async def delete_dataset(self, name: str) -> Dict[str, Any]:
+        """
+        Delete a dataset by name.
+        
+        Args:
+            name: The name of the dataset (required)
+            
+        Returns:
+            Deleted dataset information
+            
+        Raises:
+            CommandPermissionError: If the user does not have permission to delete the dataset
+            DeleteDatasetError: If the API request fails
+        """
+        response, status = await self._client._terrakio_request("DELETE", f"/datasets/{name}")
+        if status != 200:
+            if status == 403:
+                raise CommandPermissionError(f"You do not have permission to delete dataset {name}", status_code = status)
+            raise DeleteDatasetError(f"Delete dataset failed with status {status}", status_code = status)
+        return response
+
     @require_api_key
     def update_dataset(
         self, 
         name: str, 
         append: bool = True, 
-        collection: str = "terrakio-datasets",
         products: Optional[List[str]] = None,
         dates_iso8601: Optional[List[str]] = None,
         bucket: Optional[str] = None,
@@ -140,7 +183,6 @@ class DatasetManagement:
         Args:
             name: Name of the dataset (required)
             append: Whether to append data or replace (default: True)
-            collection: Dataset collection (default: 'terrakio-datasets')
             products: List of products
             dates_iso8601: List of dates (will be automatically sorted chronologically)
             bucket: Storage bucket
@@ -162,11 +204,10 @@ class DatasetManagement:
         Raises:
             APIError: If the API request fails
         """
-        # Sort dates_iso8601 chronologically if provided
         if dates_iso8601 is not None:
             dates_iso8601 = sorted(dates_iso8601)
 
-        params = {"collection": collection, "append": str(append).lower()}
+        params = {"append": str(append).lower()}
         payload = {"name": name}
         param_mapping = {
             "products": products,
@@ -194,7 +235,6 @@ class DatasetManagement:
         self, 
         name: str, 
         append: bool = True, 
-        collection: str = "terrakio-datasets",
         products: Optional[List[str]] = None,
         dates_iso8601: Optional[List[str]] = None,
         bucket: Optional[str] = None,
@@ -217,7 +257,6 @@ class DatasetManagement:
         Args:
             name: Name of the dataset (required)
             append: Whether to append data or replace (default: True)
-            collection: Dataset collection (default: 'terrakio-datasets')
             products: List of products
             dates_iso8601: List of dates (will be automatically sorted chronologically)
             bucket: Storage bucket
@@ -244,7 +283,7 @@ class DatasetManagement:
         if dates_iso8601 is not None:
             dates_iso8601 = sorted(dates_iso8601)
 
-        params = {"collection": collection, "append": str(append).lower()}
+        params = {"append": str(append).lower()}
         payload = {"name": name}
         param_mapping = {
             "products": products,
@@ -267,14 +306,55 @@ class DatasetManagement:
             if value is not None:
                 payload[param] = value
         return self._client._terrakio_request("PATCH", "/datasets", params = params, json = payload)
-    
 
 
     @require_api_key
-    def overwrite_dataset(
+    async def _get_url_for_upload_inference_script(self, script_path: str) -> str:
+        """
+        Get the url for the upload of the inference script
+        """
+        # we have the path, and we just need to get the bucket name
+        payload = {
+            "script_path": script_path,
+        }
+        return await self._client._terrakio_request("POST", "models/update_inference_script", json = payload)
+    
+    @require_api_key
+    async def update_virtual_dataset_inference(
+        self,
+        name: str,
+        inference_script_path: str,
+        append: bool = True,
+    ):
+        """
+        Update the inference script for a virtual dataset.
+        """
+        params = {"append": str(append).lower()}
+        dataset_info = await self.get_dataset(name)
+        print("the current dataset info is: ", dataset_info)
+        script_path = dataset_info["path"]
+        product_name = dataset_info["products"][0]
+        script_path = script_path + f"/{product_name}.py"
+        upload_url_dict = await self._get_url_for_upload_inference_script(script_path)
+        upload_url = upload_url_dict["script_upload_url"]
+        try:
+            with open(inference_script_path, "r", encoding="utf-8") as f:
+                script_content = f.read()
+                script_bytes = script_content.encode('utf-8')
+                headers = {
+                    "Content-Type": "text/x-python",
+                    "Content-Length": str(len(script_bytes))
+                }
+                response = await self._client._regular_request("PUT", endpoint=upload_url, data=script_bytes, headers=headers)
+        except FileNotFoundError:
+            raise FileNotFoundError(f"Inference script file not found: {inference_script_path}")
+        except Exception as e:
+            raise Exception(f"Failed to upload inference script: {str(e)}")
+    
+    @require_api_key
+    async def overwrite_dataset(
         self, 
         name: str, 
-        collection: str = "terrakio-datasets",
         products: Optional[List[str]] = None,
         dates_iso8601: Optional[List[str]] = None,
         bucket: Optional[str] = None,
@@ -295,7 +375,6 @@ class DatasetManagement:
                             
         Args:
             name: Name of the dataset (required)
-            collection: Dataset collection (default: 'terrakio-datasets')
             products: List of products
             dates_iso8601: List of dates (will be automatically sorted chronologically)
             bucket: Storage bucket
@@ -315,9 +394,10 @@ class DatasetManagement:
             Overwritten dataset information
                                 
         Raises:
-            APIError: If the API request fails
+            CommandPermissionError: If the user does not have permission to overwrite the dataset
+            DatasetNotFoundError: If the dataset is not found
+            OverwriteDatasetError: If the API request fails
         """
-        params = {"collection": collection}
         payload = {"name": name}
         param_mapping = {
             "products": products,
@@ -338,25 +418,14 @@ class DatasetManagement:
         for param, value in param_mapping.items():
             if value is not None:
                 payload[param] = value
-        return self._client._terrakio_request("PUT", "/datasets", params = params, json = payload)
-    
-    @require_api_key
-    def delete_dataset(self, name: str, collection: str = "terrakio-datasets") -> Dict[str, Any]:
-        """
-        Delete a dataset by name.
-        
-        Args:
-            name: The name of the dataset (required)
-            collection: The dataset collection (default: 'terrakio-datasets')
-            
-        Returns:
-            Deleted dataset information
-            
-        Raises:
-            APIError: If the API request fails
-        """
-        params = {"collection": collection}
-        return self._client._terrakio_request("DELETE", f"/datasets/{name}", params = params)
+        response, status = await self._client._terrakio_request("PUT", "/datasets", json = payload)
+        if status != 200:
+            if status == 403:
+                raise CommandPermissionError(f"You do not have permission to overwrite dataset {name}", status_code = status)
+            elif status == 404:
+                raise DatasetNotFoundError(f"Dataset {name} not found", status_code = status)
+            raise OverwriteDatasetError(f"Failed to overwrite dataset: {response}", status_code = status)
+        return response
 
     @require_api_key
     def download_file_to_path(self, job_name, stage, file_name, output_path):
@@ -370,9 +439,6 @@ class DatasetManagement:
                 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']