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

terrakio_core/endpoints/auth.py

@@ -1,8 +1,21 @@
-import os
 import json
-from typing import Dict, Any
-from ..exceptions import APIError
-from ..helper.decorators import require_token, require_api_key, require_auth
+import os
+from typing import Any, Dict
+
+from ..exceptions import (
+    APIKeyError,
+    AuthenticationExpireError,
+    InvalidUsernamePasswordError,
+    LoginError,
+    QuotaError,
+    RefreshAPIKeyError,
+    ResetPasswordError,
+    SignupError,
+    UserInfoError,
+    InvalidEmailFormatError,
+    EmailAlreadyExistsError,
+)
+from ..helper.decorators import require_api_key, require_auth, require_token
 
 class AuthClient:
     def __init__(self, client):
@@ -20,22 +33,23 @@ class AuthClient:
             Dict containing the authentication token
             
         Raises:
-            APIError: If the signup request fails
+            SignupError: If the signup request fails
         """
         payload = {
             "email": email,
             "password": password
         }
+        response, status = await self._client._terrakio_request("POST", "/users/signup", json=payload)
+        if status != 200:
+            if status == 422:
+                raise InvalidEmailFormatError(f"Invalid email format: {response}", status_code=status)
+            elif status == 409:
+                raise EmailAlreadyExistsError(f"Email already exists: {response}", status_code=status)
+            raise SignupError(f"Signup request failed: {response}", status_code=status)
+        else:
+            return response
 
-        
-        try:
-            result = await self._client._terrakio_request("POST", "/users/signup", json=payload)
-        except Exception as e:
-            self._client.logger.info(f"Signup failed: {str(e)}")
-            raise APIError(f"Signup request failed: {str(e)}")
-            
-
-    async def login(self, email: str, password: str) -> Dict[str, str]:
+    async def login(self, email: str, password: str) -> None:
         """
         Login a user with email and password.
         
@@ -44,7 +58,7 @@ class AuthClient:
             password: User's password
             
         Returns:
-            Dict containing the authentication token
+            None
             
         Raises:
             APIError: If the login request fails
@@ -53,10 +67,14 @@ class AuthClient:
             "email": email,
             "password": password
         }
-        
-        try:
-            result = await self._client._terrakio_request("POST", "/users/login", json=payload)
-            token_response = result.get("token")
+        response, status = await self._client._terrakio_request("POST", "/users/login", json=payload)
+        if status != 200:
+            if status == 401:
+                raise InvalidUsernamePasswordError(f"Invalid username or password: {response}", status_code=status)
+            else:
+                raise LoginError(f"Login request failed: {response}", status_code=status)
+        else:
+            token_response = response.get("token")
             
             if token_response:
                 self._client.token = token_response
@@ -72,11 +90,6 @@ class AuthClient:
                 self._client.logger.info(f"Successfully authenticated as: {email}")
                 self._client.logger.info(f"Using Terrakio API at: {self._client.url}")
                 
-            return {"token": token_response} if token_response else {"error": "Login failed"}
-            
-        except Exception as e:
-            self._client.logger.info(f"Login failed: {str(e)}")
-            raise APIError(f"Login request failed: {str(e)}")
 
     @require_token
     async def view_api_key(self) -> str:
@@ -87,11 +100,18 @@ class AuthClient:
             str: The API key
             
         Raises:
-            APIError: If the API request fails
+            AuthenticationExpireError: If authentication expired
+            APIKeyError: If the API key request fails
         """
-        result = await self._client._terrakio_request("GET", "/users/key")
-        api_key = result.get("apiKey")
-        return api_key
+        response, status = await self._client._terrakio_request("GET", "/users/key")
+        api_key = response.get("apiKey")
+        if status != 200:
+            if status == 400 and response.get("detail")["message"] == "Not authenticated":
+                raise AuthenticationExpireError(f"Authentication expired, please login again: {response}")
+            else:
+                raise APIKeyError(f"Error fetching API key: {response}", status_code=status)
+        else:
+            return api_key
 
     @require_api_key
     @require_token
@@ -103,17 +123,18 @@ class AuthClient:
             str: The new API key
             
         Raises:
-            APIError: If the API request fails
+            RefreshAPIKeyError: If the API key refresh request fails
         """
-        result = await self._client._terrakio_request("POST", "/users/refresh_key")
-        self._client.key = result.get("apiKey")
-                
-        self._update_config_key()
-        
-        return self._client.key
+        response, status = await self._client._terrakio_request("POST", "/users/refresh_key")
+        self._client.key = response.get("apiKey")
+        if status != 200:
+            raise RefreshAPIKeyError(f"Error refreshing API key: {response}", status_code=status)
+        else:
+            self._update_config_key()
+            return self._client.key
 
     @require_api_key
-    def get_user_info(self) -> Dict[str, Any]:
+    async def get_user_info(self) -> Dict[str, Any]:
         """
         Get information about the authenticated user.
         
@@ -121,9 +142,45 @@ class AuthClient:
             Dict[str, Any]: User information
             
         Raises:
-            APIError: If the API request fails
+            AuthenticationExpireError: If authentication expired
+            UserInfoError: If the user info request fails
         """
-        return self._client._terrakio_request("GET", "/users/info")
+        response, status = await self._client._terrakio_request("GET", "/users/info")
+        if status != 200:
+            if status == 400 and response.get("detail")["message"] == "Not authenticated":
+                raise AuthenticationExpireError(f"Authentication expired, please login again: {response}", status_code=status)
+            else:
+                raise UserInfoError(f"Error fetching user info: {response}", status_code=status)
+        else:
+            return response
+    
+    @require_api_key
+    async def reset_password(self, email : str) -> Dict[str, Any]:
+        """
+        Reset the password for a user by email.
+        """
+        response, status = await self._client._terrakio_request("GET", f"/users/reset-password?email={email}")
+        if status != 200:
+            raise ResetPasswordError(f"Error resetting password: {response}", status_code=status)
+        else:
+            return response['message']
+
+    @require_api_key
+    async def get_user_quota(self):
+        """
+        Get the user's quota.
+
+        Returns:
+            Dict: User's quota
+
+        Raises:
+            QuotaError: If the quota request fails
+        """
+        response, status = await self._client._terrakio_request("GET", "/users/quota")
+        if status != 200:
+            raise QuotaError(f"Error fetching quota: {response}", status_code = status)
+        else:
+            return response
 
     def _save_config(self, email: str, token: str):
         """
@@ -145,7 +202,6 @@ class AuthClient:
             config["TERRAKIO_API_KEY"] = self._client.key
             config["PERSONAL_TOKEN"] = token
             os.makedirs(os.path.dirname(config_path), exist_ok=True)
-            
             with open(config_path, 'w') as f:
                 json.dump(config, f, indent=4)
                 
@@ -177,11 +233,4 @@ class AuthClient:
             self._client.logger.info(f"API key updated in {config_path}")
             
         except Exception as e:
-            self._client.logger.info(f"Warning: Failed to update config file: {e}")
-
-
-# we have four different circumstances:
-# same expression, different region
-# same expression, same region
-# different expression, same region
-# different expression, different region
+            self._client.logger.info(f"Warning: Failed to update config file: {e}")

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']