Update src/utils/stac_client.py

src/utils/stac_client.py

@@ -4,110 +4,78 @@ from datetime import datetime, timedelta
 import random
 import requests
 from pystac_client import Client
-import os 
+import os
 
 from src.auth.auth import get_direct_access_token
 from src.utils.image import extract_url_after_filename
 
-class ProductDownloader:
+
+def get_product_content(s3_client, bucket_name, object_url):
     """
-    Class for downloading products from Copernicus Data Space Ecosystem.
-    
-    This class provides methods to download products using an S3 client connection,
-    either as in-memory content (bytes) or as files saved to disk.
+    Download the content of a product from S3 bucket.
+
+    Args:
+        s3_client: boto3 S3 client object
+        bucket_name (str): Name of the S3 bucket
+        object_url (str): Path to the object within the bucket
+
+    Returns:
+        bytes: Content of the downloaded file
     """
-    
-    def __init__(self, s3_client, bucket_name='eodata'):
-        """
-        Initialize the product downloader with an S3 client.
-        
-        Args:
-            s3_client: The boto3 S3 client to use for downloads
-            bucket_name (str): The S3 bucket name where products are stored (default: 'eodata')
-        """
-        
-        self.s3_client = s3_client
-        self.bucket_name = bucket_name
-    
-    
-    def get_product_content(self, product_path):
-        """
-        Download a Sentinel-2 product directly from Copernicus Data Space Ecosystem as a bytes object.
-        
-        Args:
-            product_path (str): S3 key or full S3 URI to the product
-            
-        Returns:
-            bytes: The product content as bytes
-            str: The filename of the product
-        """
-        # Extract S3 key if full URI is provided
-        
-        # Extract the filename from the path
-        _, filename = os.path.split(product_path)
-        
-        # Download the file to a bytes buffer
-        try:
-            # Create a bytes buffer
-            buffer = io.BytesIO()
-            
-            # Download the file to the buffer using the client
-            self.s3_client.download_fileobj(self.bucket_name, product_path, buffer)
-            
-            # Reset buffer position to the start
-            buffer.seek(0)
-            
-            # Get the bytes
-            product_content = buffer.getvalue()
-            
-            print(f"Successfully downloaded product: {filename}")
-            
-            # Return both the bytes and the filename
-            return product_content, filename
-        except Exception as e:
-            print(f"Error downloading product: {str(e)}")
-            raise
-    
-    async def download_product(self, product_path, output_filename=None):
-        """
-        Download a Sentinel-2 product directly from Copernicus Data Space Ecosystem to disk.
-        
-        Args:
-            product_path (str): S3 key or full S3 URI to the product
-            output_filename (str, optional): Filename to save the product to.
-                                              If None, uses the original filename.
-            
-        Returns:
-            str: Path to the downloaded file
-        """
-        # Extract S3 key if full URI is provided
-        if product_path.startswith('s3://'):
-            product_path = self.get_s3_key_from_href(product_path)
-        
-        # Extract the filename from the path
-        _, filename = os.path.split(product_path)
-        
-        # Use custom filename if provided, otherwise use the original
-        if output_filename is None:
-            output_filename = filename
-        
-        # Download the file using the client
-        try:
-            self.s3_client.download_file(self.bucket_name, product_path, output_filename)
-            print(f"Successfully downloaded product {filename} to {output_filename}")
-            return output_filename
-        except Exception as e:
-            print(f"Error downloading product: {str(e)}")
-            raise
+    print(f"Downloading {object_url}")
+
+    try:
+        # Download the file from S3
+        response = s3_client.get_object(Bucket=bucket_name, Key=object_url)
+        content = response['Body'].read()
+        print(f"Successfully downloaded {object_url}")
+    except Exception as e:
+        print(f"Error downloading file: {str(e)}")
+        raise
+
+    return content
+
+
+def get_product(s3_resource, bucket_name, object_url, output_path):
+    """
+    Download a product from S3 bucket and create output directory if it doesn't exist.
+
+    Args:
+        s3_resource: boto3 S3 resource object
+        bucket_name (str): Name of the S3 bucket
+        object_url (str): Path to the object within the bucket
+        output_path (str): Local directory to save the file
+
+    Returns:
+        str: Path to the downloaded file
+    """
+    # Create output directory if it doesn't exist
+    os.makedirs(output_path, exist_ok=True)
+
+    # Extract filename from the object URL
+    _, filename = os.path.split(object_url)
+
+    # Full path where the file will be saved
+    local_file_path = os.path.join(output_path, filename)
+
+    print(f"Downloading {object_url} to {local_file_path}...")
+
+    try:
+        # Download the file from S3
+        s3_resource.Bucket(bucket_name).download_file(object_url, local_file_path)
+        print(f"Successfully downloaded to {local_file_path}")
+    except Exception as e:
+        print(f"Error downloading file: {str(e)}")
+        raise
+
+    return local_file_path
 
-from src.auth.auth import get_direct_access_token
-from src.utils.image import extract_url_after_filename
 
 def download_sentinel_image(username, password, start_date, end_date,
                             bbox=[-180, -90, 180, 90], limit=10):
     """
     Download a random Sentinel-2 image based on criteria.
-    
+
     Args:
         username (str): DESTINE username
         password (str): DESTINE password
@@ -115,7 +83,7 @@ def download_sentinel_image(username, password, start_date, end_date,
         cloud_cover (int, optional): Maximum cloud cover percentage
         bbox (list): Bounding box coordinates [west, south, east, north]
         limit (int): Maximum number of results to return
-        
+
     Returns:
         tuple: (image_content or error_message, metadata)
     """
@@ -130,14 +98,14 @@ def download_sentinel_image(username, password, start_date, end_date,
     stac_base_url = "https://cachea.destine.eu"
     stac_url = f"{stac_base_url}/stac/api"
     catalog = Client.open(stac_url)
-    
+
     start_date = datetime.strptime(start_date, "%Y-%m-%d")
     end_date = datetime.strptime(end_date, "%Y-%m-%d")
     days_between = (end_date - start_date).days
     random_start_day = random.randint(0, days_between - 7)  # Ensure we have 7 days
     random_start_date = start_date + timedelta(days=random_start_day)
     random_end_date = random_start_date + timedelta(days=1)
-    
+
     # Format dates for the API
     start_date_str = random_start_date.strftime("%Y-%m-%d")
     end_date_str = random_end_date.strftime("%Y-%m-%d")
@@ -150,49 +118,49 @@ def download_sentinel_image(username, password, start_date, end_date,
         "datetime": f"{start_date_str}/{end_date_str}",
         "limit": limit
     }
-    
+
 
     # Search for Sentinel-2 images
     search = catalog.search(**search_params)
-    
+
     # Get a list of items
     items = list(search.items())
     if not items:
         return "No Sentinel-2 images found", None
-    
+
     # Select a random item
     random_item = random.choice(items)
-    
+
     # Get metadata for the selected item
     metadata = {
         "id": random_item.id,
         "datetime": random_item.datetime.strftime("%Y-%m-%d %H:%M:%S"),
         "bbox": random_item.bbox,
     }
-    
-    
+
+
     # Get the assets of the random item
     assets = random_item.assets
     asset_keys = list(assets.keys())
-    
+
     # Filter the assets to get the one that ends with *_TCI_60m.jp2
     tci_assets = [assets[key].href for key in asset_keys if assets[key].href.endswith('_TCI_60m.jp2')]
-    
+
     if not tci_assets:
         return "No TCI assets found in the selected image", None
-    
+
     filepath = extract_url_after_filename(tci_assets[0])
     metadata["filename"] = os.path.basename(filepath)
-    
+
     # Download the file
     url = f"{stac_base_url}/stac/download?filename={filepath}"
-    
+
     headers = {
         'Authorization': f'Bearer {access_token}'
     }
-    
+
     response = requests.post(url, headers=headers, data={})
-    
+
     if response.status_code == 200:
         return response.content, metadata
     else: