Add HTTPS file size retrieval to create_blast_db.py

Implemented a function to retrieve file sizes over HTTPS and integrated it with the existing FTP file size logic. This ensures support for both protocols and improves flexibility in handling different types of URIs.

src/create_blast_db.py

@@ -32,7 +32,7 @@
 
 from utils import (check_md5sum, check_output, edit_fasta, get_ftp_file_size,
                    get_mod_from_json, needs_parse_id, run_command, s3_sync,
-                   setup_logger, slack_message)
+                   setup_logger, slack_message, get_https_file_size)
 
 # Load environment variables
 load_dotenv()
@@ -129,7 +129,10 @@ def get_files_ftp(fasta_uri: str, md5sum: str) -> bool:
         return False
 
     try:
-        file_size = get_ftp_file_size(fasta_uri)
+        if fasta_uri.startswith('https'):
+            file_size = get_https_file_size(fasta_uri)
+        else:
+            file_size = get_ftp_file_size(fasta_uri)
         if file_size == 0:
             LOGGER.error(f"Failed to get file size for {fasta_uri}")
             return False

src/utils.py

@@ -425,3 +425,20 @@ def get_ftp_file_size(fasta_uri: str) -> int:
     except Exception as e:
         console.log(f"Error getting FTP file size: {e}")
         return 0
+
+
+import requests
+
+def get_https_file_size(https_uri: str) -> int:
+    try:
+        response = requests.head(https_uri, allow_redirects=True)
+        if response.status_code == 200:
+            size = int(response.headers.get('content-length', 0))
+            console.log(f"File size for {https_uri} is {size} bytes")
+            return size
+        else:
+            console.log(f"Couldn't determine size for {https_uri}. Status code: {response.status_code}")
+            return 0
+    except Exception as e:
+        console.log(f"Error getting HTTPS file size: {e}")
+        return 0