terrakio-core 0.3.1__py3-none-any.whl → 0.3.3__py3-none-any.whl

terrakio_core/client.py

@@ -2,6 +2,7 @@ import json
 import asyncio
 from io import BytesIO
 from typing import Dict, Any, Optional, Union
+from functools import wraps
 
 import requests
 import aiohttp
@@ -18,8 +19,39 @@ import logging
 import textwrap
 
 
+def require_api_key(func):
+    """
+    Decorator to ensure an API key is available before a method can be executed.
+    Will check for the presence of an API key before allowing the function to be called.
+    """
+    @wraps(func)
+    def wrapper(self, *args, **kwargs):
+        # Check if the API key is set
+        if not self.key:
+            error_msg = "No API key found. Please provide an API key to use this client."
+            if not self.quiet:
+                print(error_msg)
+                print("API key is required for this function.")
+                print("You can set an API key by either:")
+                print("1. Loading from a config file")
+                print("2. Using the login() method: client.login(email='your-email@example.com', password='your-password')")
+            raise ConfigurationError(error_msg)
+            
+        # Check if URL is set (required for API calls)
+        if not self.url:
+            if not self.quiet:
+                print("API URL is not set. Using default API URL: https://api.terrak.io")
+            self.url = "https://api.terrak.io"
+            
+        return func(self, *args, **kwargs)
+    
+    # Mark as decorated to avoid double decoration
+    wrapper._is_decorated = True
+    return wrapper
+
+
 class BaseClient:
-    def __init__(self, url: Optional[str] = None, key: Optional[str] = None, 
+    def __init__(self, url: Optional[str] = None,
                 auth_url: Optional[str] = "https://dev-au.terrak.io",
                 quiet: bool = False, config_file: Optional[str] = None,
                 verify: bool = True, timeout: int = 300):
@@ -28,57 +60,96 @@ class BaseClient:
         self.verify = verify
         self.timeout = timeout
         self.auth_client = None
-        if auth_url:
-            from terrakio_core.auth import AuthClient
-            self.auth_client = AuthClient(
-                base_url=auth_url,
-                verify=verify,
-                timeout=timeout
-            )
+        # Initialize session early to avoid AttributeError
+        self.session = requests.Session()
+        self.user_management = None
+        self.dataset_management = None
+        self.mass_stats = None
+        self._aiohttp_session = None
+        
+        # if auth_url:
+        #     from terrakio_core.auth import AuthClient
+        #     self.auth_client = AuthClient(
+        #         base_url=auth_url,
+        #         verify=verify,
+        #         timeout=timeout
+        #     )
         self.url = url
-        self.key = key
-        if self.url is None or self.key is None:
-            from terrakio_core.config import read_config_file, DEFAULT_CONFIG_FILE
-            if config_file is None:
-                config_file = DEFAULT_CONFIG_FILE
-            try:
-                config = read_config_file(config_file)
-                if self.url is None:
-                    self.url = config.get('url')
-                if self.key is None:
-                    self.key = config.get('key')
-            except Exception as e:
-                raise ConfigurationError(
-                    f"Failed to read configuration: {e}\n\n"
-                    "To fix this issue:\n"
-                    "1. Create a file at ~/.terrakioapirc with:\n"
-                    "url: https://api.terrak.io\n"
-                    "key: your-api-key\n\n"
-                    "OR\n\n"
-                    "2. Initialize the client with explicit parameters:\n"
-                    "client = terrakio_api.Client(\n"
-                    "    url='https://api.terrak.io',\n"
-                    "    key='your-api-key'\n"
-                    ")"
-                )
-        if not self.url:
-            raise ConfigurationError("Missing API URL in configuration")
+        self.key = None
+        
+        # Load configuration from file
+        from terrakio_core.config import read_config_file, DEFAULT_CONFIG_FILE
+        
+        if config_file is None:
+            config_file = DEFAULT_CONFIG_FILE
+        
+        # Get config using the read_config_file function that now handles all cases
+        config = read_config_file(config_file, quiet=self.quiet)
+        
+        # If URL is not provided, try to get it from config
+        if self.url is None:
+            self.url = config.get('url')
+            
+        # Get API key from config file (never from parameters)
+        self.key = config.get('key')
+
+        self.token = config.get('token')
+        
+        # Update auth_client with token if it exists
+        if self.auth_client and self.token:
+            self.auth_client.token = self.token
+            self.auth_client.session.headers.update({
+                "Authorization": self.token
+            })
+        
+        # If we have a key, we're good to go even if not logged in
+        if self.key:
+            # If URL is missing but we have a key, use the default URL
+            if not self.url:
+                self.url = "https://api.terrak.io"
+            print("the token is! ", self.token)
+            print("the key is! ", self.key)
+            # Update session headers with the key
+            headers = {
+                'Content-Type': 'application/json',
+                'x-api-key': self.key
+            }
+            
+            # Only add Authorization header if token exists and is not None
+            if self.token:
+                headers['Authorization'] = self.token
+                
+            self.session.headers.update(headers)
+            print("the session headers are ", self.session.headers)
+
+            
+            if not self.quiet and config.get('user_email'):
+                print(f"Using API key for: {config.get('user_email')}")
+                
+            return
+
+        # Check if we have the required configuration
         if not self.key:
-            raise ConfigurationError("Missing API key in configuration")
+            # No API key available - inform the user
+            if not self.quiet:
+                print("API key is required to use this client.")
+                print("You can set an API key by either:")
+                print("1. Loading from a config file")
+                print("2. Using the login() method: client.login(email='your-email@example.com', password='your-password')")
+            return
+            
         self.url = self.url.rstrip('/')
         if not self.quiet:
             print(f"Using Terrakio API at: {self.url}")
-        self.session = requests.Session()
+            
+        # Update the session headers with API key
         self.session.headers.update({
             'Content-Type': 'application/json',
             'x-api-key': self.key
         })
-        self.user_management = None
-        self.dataset_management = None
-        self.mass_stats = None
-        self._aiohttp_session = None
 
     @property
+    @require_api_key
     async def aiohttp_session(self):
         if self._aiohttp_session is None or self._aiohttp_session.closed:
             self._aiohttp_session = aiohttp.ClientSession(
@@ -90,6 +161,7 @@ class BaseClient:
             )
         return self._aiohttp_session
 
+    @require_api_key
     async def wcs_async(self, expr: str, feature: Union[Dict[str, Any], ShapelyGeometry],
                         in_crs: str = "epsg:4326", out_crs: str = "epsg:4326",
                         output: str = "csv", resolution: int = -1, buffer: bool = False, 
@@ -130,6 +202,7 @@ class BaseClient:
             "resolution": resolution,
             **kwargs
         }
+        print("the payload is ", payload)
         request_url = f"{self.url}/geoquery"
         for attempt in range(retry + 1):
             try:
@@ -186,38 +259,67 @@ class BaseClient:
                     raise
                 continue
 
+    @require_api_key
     async def close_async(self):
         """Close the aiohttp session"""
         if self._aiohttp_session and not self._aiohttp_session.closed:
             await self._aiohttp_session.close()
             self._aiohttp_session = None
 
+    @require_api_key
     async def __aenter__(self):
         return self
 
+    @require_api_key
     async def __aexit__(self, exc_type, exc_val, exc_tb):
         await self.close_async()
 
     def signup(self, email: str, password: str) -> Dict[str, Any]:
         if not self.auth_client:
-            raise ConfigurationError("Authentication client not initialized. Please provide auth_url during client initialization.")
+            from terrakio_core.auth import AuthClient
+            self.auth_client = AuthClient(
+                base_url=self.url,
+                verify=self.verify,
+                timeout=self.timeout
+            )
+            self.auth_client.session = self.session
         return self.auth_client.signup(email, password)
 
     def login(self, email: str, password: str) -> Dict[str, str]:
+
         if not self.auth_client:
-            raise ConfigurationError("Authentication client not initialized. Please provide auth_url during client initialization.")
-        
+            from terrakio_core.auth import AuthClient
+            self.auth_client = AuthClient(
+                base_url=self.url,
+                verify=self.verify,
+                timeout=self.timeout
+            )
+            self.auth_client.session = self.session
         try:
             # First attempt to login
             token_response = self.auth_client.login(email, password)
             
-            print("the token response is ", token_response)
             # Only proceed with API key retrieval if login was successful
             if token_response:
                 # After successful login, get the API key
-                api_key_response = self.view_api_key()
+                api_key_response = self.auth_client.view_api_key()
                 self.key = api_key_response
                 
+                # Make sure URL is set
+                if not self.url:
+                    self.url = "https://api.terrak.io"
+                
+                # Make sure session headers are updated with the new key
+                self.session.headers.update({
+                    'Content-Type': 'application/json',
+                    'x-api-key': self.key
+                })
+                
+                # Set URL if not already set
+                if not self.url:
+                    self.url = "https://api.terrak.io"
+                    self.url = self.url.rstrip('/')
+                
                 # Save email and API key to config file
                 import os
                 import json
@@ -229,6 +331,7 @@ class BaseClient:
                             config = json.load(f)
                         config["EMAIL"] = email
                         config["TERRAKIO_API_KEY"] = self.key
+                        config["PERSONAL_TOKEN"] = token_response
                     
                     os.makedirs(os.path.dirname(config_path), exist_ok=True)
                     with open(config_path, 'w') as f:
@@ -236,6 +339,7 @@ class BaseClient:
                     
                     if not self.quiet:
                         print(f"Successfully authenticated as: {email}")
+                        print(f"Using Terrakio API at: {self.url}")
                         print(f"API key saved to {config_path}")
                 except Exception as e:
                     if not self.quiet:
@@ -247,47 +351,81 @@ class BaseClient:
                 print(f"Login failed: {str(e)}")
             raise
 
+    @require_api_key
     def refresh_api_key(self) -> str:
-        if not self.auth_client:
-            raise ConfigurationError("Authentication client not initialized. Please provide auth_url during client initialization.")
-        if not self.auth_client.token:
-            raise ConfigurationError("Not authenticated. Call login() first.")
-        self.key = self.auth_client.refresh_api_key()
-        self.session.headers.update({'x-api-key': self.key})
-        import os
-        config_path = os.path.join(os.environ.get("HOME", ""), ".tkio_config.json")
-        try:
-            config = {"EMAIL": "", "TERRAKIO_API_KEY": ""}
-            if os.path.exists(config_path):
-                with open(config_path, 'r') as f:
-                    config = json.load(f)
-            config["TERRAKIO_API_KEY"] = self.key
-            os.makedirs(os.path.dirname(config_path), exist_ok=True)
-            with open(config_path, 'w') as f:
-                json.dump(config, f, indent=4)
-            if not self.quiet:
-                print(f"API key generated successfully and updated in {config_path}")
-        except Exception as e:
-            if not self.quiet:
-                print(f"Warning: Failed to update config file: {e}")
+        # If we have auth_client and it has a token, refresh the key
+        if self.auth_client and self.auth_client.token:
+            self.key = self.auth_client.refresh_api_key()
+            self.session.headers.update({'x-api-key': self.key})
+            
+            # Update the config file with the new key
+            import os
+            config_path = os.path.join(os.environ.get("HOME", ""), ".tkio_config.json")
+            try:
+                config = {"EMAIL": "", "TERRAKIO_API_KEY": ""}
+                if os.path.exists(config_path):
+                    with open(config_path, 'r') as f:
+                        config = json.load(f)
+                config["TERRAKIO_API_KEY"] = self.key
+                os.makedirs(os.path.dirname(config_path), exist_ok=True)
+                with open(config_path, 'w') as f:
+                    json.dump(config, f, indent=4)
+                if not self.quiet:
+                    print(f"API key generated successfully and updated in {config_path}")
+            except Exception as e:
+                if not self.quiet:
+                    print(f"Warning: Failed to update config file: {e}")
+            return self.key
+        else:
+            # If we don't have auth_client with a token but have a key already, return it
+            if self.key:
+                if not self.quiet:
+                    print("Using existing API key from config file.")
+                return self.key
+            else:
+                raise ConfigurationError("No authentication token available. Please login first to refresh the API key.")
+
+    @require_api_key
+    def view_api_key(self) -> str:
+        # If we have the auth client and token, refresh the key
+        if self.auth_client and self.auth_client.token:
+            self.key = self.auth_client.view_api_key()
+            self.session.headers.update({'x-api-key': self.key})
+        
         return self.key
 
+    @require_api_key
     def view_api_key(self) -> str:
+        # If we have the auth client and token, refresh the key
         if not self.auth_client:
-            raise ConfigurationError("Authentication client not initialized. Please provide auth_url during client initialization.")
-        if not self.auth_client.token:
-            raise ConfigurationError("Not authenticated. Call login() first.")
-        self.key = self.auth_client.view_api_key()
-        self.session.headers.update({'x-api-key': self.key})
-        return self.key
+            from terrakio_core.auth import AuthClient
+            self.auth_client = AuthClient(
+                base_url=self.url,
+                verify=self.verify,
+                timeout=self.timeout
+            )
+            self.auth_client.session = self.session
+        return self.auth_client.view_api_key()
 
+    # @require_api_key
+    # def get_user_info(self) -> Dict[str, Any]:
+    #     return self.auth_client.get_user_info()
+    @require_api_key
     def get_user_info(self) -> Dict[str, Any]:
+        # Initialize auth_client if it doesn't exist
         if not self.auth_client:
-            raise ConfigurationError("Authentication client not initialized. Please provide auth_url during client initialization.")
-        if not self.auth_client.token:
-            raise ConfigurationError("Not authenticated. Call login() first.")
+            from terrakio_core.auth import AuthClient
+            self.auth_client = AuthClient(
+                base_url=self.url,
+                verify=self.verify,
+                timeout=self.timeout
+            )
+            # Use the same session as the base client
+            self.auth_client.session = self.session
+            
         return self.auth_client.get_user_info()
 
+    @require_api_key
     def wcs(self, expr: str, feature: Union[Dict[str, Any], ShapelyGeometry], in_crs: str = "epsg:4326",
             out_crs: str = "epsg:4326", output: str = "csv", resolution: int = -1,
             **kwargs):
@@ -340,6 +478,7 @@ class BaseClient:
             raise APIError(f"Request failed: {str(e)}")
 
     # Admin/protected methods
+    @require_api_key
     def _get_user_by_id(self, user_id: str):
         if not self.user_management:
             from terrakio_core.user_management import UserManagement
@@ -353,6 +492,7 @@ class BaseClient:
             )
         return self.user_management.get_user_by_id(user_id)
 
+    @require_api_key
     def _get_user_by_email(self, email: str):
         if not self.user_management:
             from terrakio_core.user_management import UserManagement
@@ -366,6 +506,7 @@ class BaseClient:
             )
         return self.user_management.get_user_by_email(email)
 
+    @require_api_key
     def _list_users(self, substring: str = None, uid: bool = False):
         if not self.user_management:
             from terrakio_core.user_management import UserManagement
@@ -379,6 +520,7 @@ class BaseClient:
             )
         return self.user_management.list_users(substring=substring, uid=uid)
 
+    @require_api_key
     def _edit_user(self, user_id: str, uid: str = None, email: str = None, role: str = None, apiKey: str = None, groups: list = None, quota: int = None):
         if not self.user_management:
             from terrakio_core.user_management import UserManagement
@@ -400,6 +542,7 @@ class BaseClient:
             quota=quota
         )
 
+    @require_api_key
     def _reset_quota(self, email: str, quota: int = None):
         if not self.user_management:
             from terrakio_core.user_management import UserManagement
@@ -413,6 +556,7 @@ class BaseClient:
             )
         return self.user_management.reset_quota(email=email, quota=quota)
 
+    @require_api_key
     def _delete_user(self, uid: str):
         if not self.user_management:
             from terrakio_core.user_management import UserManagement
@@ -427,6 +571,7 @@ class BaseClient:
         return self.user_management.delete_user(uid=uid)
 
     # Dataset management protected methods
+    @require_api_key
     def _get_dataset(self, name: str, collection: str = "terrakio-datasets"):
         if not self.dataset_management:
             from terrakio_core.dataset_management import DatasetManagement
@@ -440,6 +585,7 @@ class BaseClient:
             )
         return self.dataset_management.get_dataset(name=name, collection=collection)
 
+    @require_api_key
     def _list_datasets(self, substring: str = None, collection: str = "terrakio-datasets"):
         if not self.dataset_management:
             from terrakio_core.dataset_management import DatasetManagement
@@ -453,6 +599,7 @@ class BaseClient:
             )
         return self.dataset_management.list_datasets(substring=substring, collection=collection)
 
+    @require_api_key
     def _create_dataset(self, name: str, collection: str = "terrakio-datasets", **kwargs):
         if not self.dataset_management:
             from terrakio_core.dataset_management import DatasetManagement
@@ -466,6 +613,7 @@ class BaseClient:
             )
         return self.dataset_management.create_dataset(name=name, collection=collection, **kwargs)
 
+    @require_api_key
     def _update_dataset(self, name: str, append: bool = True, collection: str = "terrakio-datasets", **kwargs):
         if not self.dataset_management:
             from terrakio_core.dataset_management import DatasetManagement
@@ -479,6 +627,7 @@ class BaseClient:
             )
         return self.dataset_management.update_dataset(name=name, append=append, collection=collection, **kwargs)
 
+    @require_api_key
     def _overwrite_dataset(self, name: str, collection: str = "terrakio-datasets", **kwargs):
         if not self.dataset_management:
             from terrakio_core.dataset_management import DatasetManagement
@@ -492,6 +641,7 @@ class BaseClient:
             )
         return self.dataset_management.overwrite_dataset(name=name, collection=collection, **kwargs)
 
+    @require_api_key
     def _delete_dataset(self, name: str, collection: str = "terrakio-datasets"):
         if not self.dataset_management:
             from terrakio_core.dataset_management import DatasetManagement
@@ -505,6 +655,7 @@ class BaseClient:
             )
         return self.dataset_management.delete_dataset(name=name, collection=collection)
 
+    @require_api_key
     def close(self):
         """Close all client sessions"""
         self.session.close()
@@ -531,13 +682,16 @@ class BaseClient:
                 # Event loop may already be closed, ignore
                 pass
 
+    @require_api_key
     def __enter__(self):
         return self
 
+    @require_api_key
     def __exit__(self, exc_type, exc_val, exc_tb):
         self.close()
 
     @admin_only_params('location', 'force_loc', 'server')
+    @require_api_key
     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
@@ -552,6 +706,7 @@ class BaseClient:
         return self.mass_stats.execute_job(name, region, output, config, overwrite, skip_existing, request_json, manifest_json, location, force_loc, server)
 
 
+    @require_api_key
     def get_mass_stats_task_id(self, name, stage, uid=None):
         if not self.mass_stats:
             from terrakio_core.mass_stats import MassStats
@@ -565,7 +720,8 @@ class BaseClient:
             )
         return self.mass_stats.get_task_id(name, stage, uid)
 
-    def track_mass_stats_job(self, ids=None):
+    @require_api_key
+    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:
@@ -578,6 +734,7 @@ class BaseClient:
             )
         return self.mass_stats.track_job(ids)
 
+    @require_api_key
     def get_mass_stats_history(self, limit=100):
         if not self.mass_stats:
             from terrakio_core.mass_stats import MassStats
@@ -591,6 +748,7 @@ class BaseClient:
             )
         return self.mass_stats.get_history(limit)
 
+    @require_api_key
     def start_mass_stats_post_processing(self, process_name, data_name, output, consumer_path, overwrite=False):
         if not self.mass_stats:
             from terrakio_core.mass_stats import MassStats
@@ -604,6 +762,7 @@ class BaseClient:
             )
         return self.mass_stats.start_post_processing(process_name, data_name, output, consumer_path, overwrite)
 
+    @require_api_key
     def download_mass_stats_results(self, id=None, force_loc=False, **kwargs):
         if not self.mass_stats:
             from terrakio_core.mass_stats import MassStats
@@ -617,6 +776,7 @@ class BaseClient:
             )
         return self.mass_stats.download_results(id, force_loc, **kwargs)
 
+    @require_api_key
     def cancel_mass_stats_job(self, id):
         if not self.mass_stats:
             from terrakio_core.mass_stats import MassStats
@@ -630,6 +790,7 @@ class BaseClient:
             )
         return self.mass_stats.cancel_job(id)
 
+    @require_api_key
     def cancel_all_mass_stats_jobs(self):
         if not self.mass_stats:
             from terrakio_core.mass_stats import MassStats
@@ -643,6 +804,7 @@ class BaseClient:
             )
         return self.mass_stats.cancel_all_jobs()
 
+    @require_api_key
     def _create_pyramids(self, name, levels, config):
         if not self.mass_stats:
             from terrakio_core.mass_stats import MassStats
@@ -656,6 +818,7 @@ class BaseClient:
             )
         return self.mass_stats.create_pyramids(name, levels, config)
 
+    @require_api_key
     def random_sample(self, name, **kwargs):
         if not self.mass_stats:
             from terrakio_core.mass_stats import MassStats
@@ -669,6 +832,7 @@ class BaseClient:
             )
         return self.mass_stats.random_sample(name, **kwargs)
 
+    @require_api_key
     async def zonal_stats_async(self, gdb, expr, conc=20, inplace=False, output="csv", 
                                in_crs="epsg:4326", out_crs="epsg:4326", resolution=-1, buffer=False):
         """
@@ -857,6 +1021,7 @@ class BaseClient:
             else:
                 return result_gdf
 
+    @require_api_key
     def zonal_stats(self, gdb, expr, conc=20, inplace=False, output="csv", 
                    in_crs="epsg:4326", out_crs="epsg:4326", resolution=-1, buffer=False):
         """
@@ -919,75 +1084,80 @@ class BaseClient:
         return result
 
     # Group access management protected methods
+    @require_api_key
     def _get_group_users_and_datasets(self, group_name: str):
-        if not hasattr(self, "group_access_management") or self.group_access_management is None:
+        if not self.group_access:
             from terrakio_core.group_access_management import GroupAccessManagement
             if not self.url or not self.key:
                 raise ConfigurationError("Group access management client not initialized. Make sure API URL and key are set.")
-            self.group_access_management = GroupAccessManagement(
+            self.group_access = GroupAccessManagement(
                 api_url=self.url,
                 api_key=self.key,
                 verify=self.verify,
                 timeout=self.timeout
             )
-        return self.group_access_management.get_group_users_and_datasets(group_name)
+        return self.group_access.get_group_users_and_datasets(group_name=group_name)
 
+    @require_api_key
     def _add_group_to_dataset(self, dataset: str, group: str):
-        if not hasattr(self, "group_access_management") or self.group_access_management is None:
+        if not self.group_access:
             from terrakio_core.group_access_management import GroupAccessManagement
             if not self.url or not self.key:
                 raise ConfigurationError("Group access management client not initialized. Make sure API URL and key are set.")
-            self.group_access_management = GroupAccessManagement(
+            self.group_access = GroupAccessManagement(
                 api_url=self.url,
                 api_key=self.key,
                 verify=self.verify,
                 timeout=self.timeout
             )
-        return self.group_access_management.add_group_to_dataset(dataset, group)
+        return self.group_access.add_group_to_dataset(dataset=dataset, group=group)
 
+    @require_api_key
     def _add_group_to_user(self, uid: str, group: str):
-        if not hasattr(self, "group_access_management") or self.group_access_management is None:
+        if not self.group_access:
             from terrakio_core.group_access_management import GroupAccessManagement
             if not self.url or not self.key:
                 raise ConfigurationError("Group access management client not initialized. Make sure API URL and key are set.")
-            self.group_access_management = GroupAccessManagement(
+            self.group_access = GroupAccessManagement(
                 api_url=self.url,
                 api_key=self.key,
                 verify=self.verify,
                 timeout=self.timeout
             )
-        print("the uid is and the group is ", uid, group)
-        return self.group_access_management.add_group_to_user(uid, group)
+        return self.group_access.add_group_to_user(uid=uid, group=group)
 
+    @require_api_key
     def _delete_group_from_user(self, uid: str, group: str):
-        if not hasattr(self, "group_access_management") or self.group_access_management is None:
+        if not self.group_access:
             from terrakio_core.group_access_management import GroupAccessManagement
             if not self.url or not self.key:
                 raise ConfigurationError("Group access management client not initialized. Make sure API URL and key are set.")
-            self.group_access_management = GroupAccessManagement(
+            self.group_access = GroupAccessManagement(
                 api_url=self.url,
                 api_key=self.key,
                 verify=self.verify,
                 timeout=self.timeout
             )
-        return self.group_access_management.delete_group_from_user(uid, group)
+        return self.group_access.delete_group_from_user(uid=uid, group=group)
 
+    @require_api_key
     def _delete_group_from_dataset(self, dataset: str, group: str):
-        if not hasattr(self, "group_access_management") or self.group_access_management is None:
+        if not self.group_access:
             from terrakio_core.group_access_management import GroupAccessManagement
             if not self.url or not self.key:
                 raise ConfigurationError("Group access management client not initialized. Make sure API URL and key are set.")
-            self.group_access_management = GroupAccessManagement(
+            self.group_access = GroupAccessManagement(
                 api_url=self.url,
                 api_key=self.key,
                 verify=self.verify,
                 timeout=self.timeout
             )
-        return self.group_access_management.delete_group_from_dataset(dataset, group)
+        return self.group_access.delete_group_from_dataset(dataset=dataset, group=group)
 
     # Space management protected methods
+    @require_api_key
     def _get_total_space_used(self):
-        if not hasattr(self, "space_management") or self.space_management is None:
+        if not self.space_management:
             from terrakio_core.space_management import SpaceManagement
             if not self.url or not self.key:
                 raise ConfigurationError("Space management client not initialized. Make sure API URL and key are set.")
@@ -999,8 +1169,9 @@ class BaseClient:
             )
         return self.space_management.get_total_space_used()
 
+    @require_api_key
     def _get_space_used_by_job(self, name: str, region: str = None):
-        if not hasattr(self, "space_management") or self.space_management is None:
+        if not self.space_management:
             from terrakio_core.space_management import SpaceManagement
             if not self.url or not self.key:
                 raise ConfigurationError("Space management client not initialized. Make sure API URL and key are set.")
@@ -1012,8 +1183,9 @@ class BaseClient:
             )
         return self.space_management.get_space_used_by_job(name, region)
 
+    @require_api_key
     def _delete_user_job(self, name: str, region: str = None):
-        if not hasattr(self, "space_management") or self.space_management is None:
+        if not self.space_management:
             from terrakio_core.space_management import SpaceManagement
             if not self.url or not self.key:
                 raise ConfigurationError("Space management client not initialized. Make sure API URL and key are set.")
@@ -1025,8 +1197,9 @@ class BaseClient:
             )
         return self.space_management.delete_user_job(name, region)
 
+    @require_api_key
     def _delete_data_in_path(self, path: str, region: str = None):
-        if not hasattr(self, "space_management") or self.space_management is None:
+        if not self.space_management:
             from terrakio_core.space_management import SpaceManagement
             if not self.url or not self.key:
                 raise ConfigurationError("Space management client not initialized. Make sure API URL and key are set.")
@@ -1038,6 +1211,22 @@ class BaseClient:
             )
         return self.space_management.delete_data_in_path(path, region)
     
+    @require_api_key
+    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)
+
+    
+    @require_api_key
     def generate_ai_dataset(
         self,
         name: str,
@@ -1107,44 +1296,32 @@ class BaseClient:
             overwrite=True
         )["task_id"]
         print("the task id is ", task_id)
-        task_id = self.start_mass_stats_job(task_id)
-        print("the task id is ", task_id)
-        return 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)
 
-    # def train_model(self, model_name: str, training_data: dict) -> 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.
-
-    #     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
-    #     }
-    #     try:
-    #         response = self.session.post(endpoint, json=payload, timeout=self.timeout, verify=self.verify)
-    #         if not response.ok:
-    #             error_msg = f"Model training request failed: {response.status_code} {response.reason}"
-    #             try:
-    #                 error_data = response.json()
-    #                 if "detail" in error_data:
-    #                     error_msg += f" - {error_data['detail']}"
-    #             except Exception:
-    #                 if response.text:
-    #                     error_msg += f" - {response.text}"
-    #             raise APIError(error_msg)
-    #         return response.json()
-    #     except requests.RequestException as e:
-    #         raise APIError(f"Model training request failed: {str(e)}")
+        # 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)
+        # now we hav ethe random sampel
 
+        # print("the task id is ", task_id)
+        return task_id
 
+    @require_api_key
     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.
@@ -1189,6 +1366,7 @@ class BaseClient:
             raise APIError(f"Model training request failed: {str(e)}")
 
     # Mass Stats methods
+    @require_api_key
     def combine_tiles(self, 
         data_name: str,
         usezarr: bool,
@@ -1209,7 +1387,8 @@ class BaseClient:
     
 
 
-    def generate_combine_tiles(
+    @require_api_key
+    def create_dataset_file(
         self,
         name: str,
         aoi: str,
@@ -1224,9 +1403,22 @@ class BaseClient:
         skip_existing: bool = False,
         non_interactive: bool = True,
         usezarr: bool = False,
-        poll_interval: int = 30   # seconds between job status checks
+        poll_interval: int = 30,
+        download_path: str = "/home/user/Downloads",
     ) -> 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
+            )
+        
+        
         from terrakio_core.generation.tiles import tiles
         import tempfile
         import time
@@ -1315,44 +1507,70 @@ class BaseClient:
         os.unlink(tempmanifestname)
 
 
+        # return self.mass_stats.combine_tiles(body["name"], usezarr, body["overwrite"], body["output"])
+    
         # 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"])
-
-
+        combine_result = self.mass_stats.combine_tiles(body["name"], usezarr, body["overwrite"], body["output"])
+        combine_task_id = combine_result.get("task_id")
 
+        # Poll combine_tiles job status
+        combine_start_time = time.time()
+        while True:
+            try:
+                trackinfo = self.mass_stats.track_job([combine_task_id])
+                download_file_name = trackinfo[combine_task_id]['folder'] + '.nc'
+                bucket = trackinfo[combine_task_id]['bucket']
+                combine_status = trackinfo[combine_task_id]['status']
+                if combine_status == 'Completed':
+                    print('Tiles combined successfully!')
+                    break
+                elif combine_status in ['Failed', 'Cancelled', 'Error']:
+                    raise RuntimeError(f"Combine job {combine_task_id} failed with status: {combine_status}")
+                else:
+                    elapsed_time = time.time() - combine_start_time
+                    print(f"Combine job status: {combine_status} - Elapsed time: {elapsed_time:.1f}s", end='\r')
+                    time.sleep(poll_interval)
+            except KeyboardInterrupt:
+                print(f"\nInterrupted! Combine job {combine_task_id} is still running in the background.")
+                raise
+            except Exception as e:
+                print(f"\nError tracking combine job: {e}")
+                raise
 
+        # Download the resulting file
+        if download_path:
+            self.mass_stats.download_file(body["name"], bucket, download_file_name, download_path)
+        else:
+            path = f"{body['name']}/combinedOutput/{download_file_name}"
+            print(f"Combined file is available at {path}")
 
+        return {"generation_task_id": task_id, "combine_task_id": combine_task_id}
 
 
+        # 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)
 
 
+    @require_api_key
     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)
         script_name = f"{product}.py"
         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)
 
+    @require_api_key
     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
@@ -1360,47 +1578,96 @@ 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()
-
+    
+    @require_api_key
     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"""
 
@@ -1410,3 +1677,55 @@ class BaseClient:
         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}")
 
+
+
+
+    @require_api_key
+    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)
+
+
+# Apply the @require_api_key decorator to ALL methods in BaseClient
+# except only the absolute minimum that shouldn't require auth
+def _apply_api_key_decorator():
+    # Only these methods can be used without API key
+    skip_auth_methods = [
+        '__init__', 'login', 'signup'
+    ]
+    
+    # Get all attributes of BaseClient
+    for attr_name in dir(BaseClient):
+        # Skip special methods and methods in skip list
+        if attr_name.startswith('__') and attr_name not in ['__enter__', '__exit__', '__aenter__', '__aexit__'] or attr_name in skip_auth_methods:
+            continue
+            
+        # Get the attribute
+        attr = getattr(BaseClient, attr_name)
+        
+        # Skip if not callable (not a method) or already decorated
+        if not callable(attr) or hasattr(attr, '_is_decorated'):
+            continue
+            
+        # Apply decorator to EVERY method not in skip list
+        setattr(BaseClient, attr_name, require_api_key(attr))
+        # Mark as decorated to avoid double decoration
+        getattr(BaseClient, attr_name)._is_decorated = True
+
+# Run the decorator application
+_apply_api_key_decorator()