Add code to support incongruous MAST HST data

Author: cam72cam

astroquery/mast/core.py

@@ -12,6 +12,7 @@
 import warnings
 import json
 import time
+import string
 import os
 import re
 import keyring
@@ -822,6 +823,7 @@ def __init__(self, *args, **kwargs):
         super(ObservationsClass, self).__init__(*args, **kwargs)
 
         self._boto3 = None
+        self._botocore = None
 
     def list_missions(self):
         """
@@ -1303,7 +1305,9 @@ def enable_s3_hst_dataset(self):
         Requires the boto3 library to function.
         """
         import boto3
+        import botocore
         self._boto3 = boto3
+        self._botocore = botocore
 
         log.info("Using the S3 HST public dataset")
         log.warning("Your AWS account will be charged for access to the S3 bucket")
@@ -1316,6 +1320,7 @@ def disable_s3_hst_dataset(self):
         Disables downloading HST public files from S3 instead of MAST
         """
         self._boto3 = None
+        self._botocore = None
 
     def _download_from_s3(self, dataProduct, localPath, cache=True):
         # The following is a mishmash of BaseQuery._download_file and s3 access through boto
@@ -1328,18 +1333,19 @@ def _download_from_s3(self, dataProduct, localPath, cache=True):
         bkt = s3.Bucket(bkt_name)
 
         dataUri = dataProduct['dataURI']
+        filename = dataUri.split("/")[-1]
         obs_id = dataProduct['obs_id']
 
         obs_id = obs_id.lower()
-        
+
         # This next part is a bit funky.  Let me explain why:
         # We have 2 different possible URI schemes for HST:
         #   mast:HST/product/obs_id_filename.type (old style)
         #   mast:HST/product/obs_id/obs_id_filename.type (new style)
         # The first scheme was developed thinking that the obs_id in the filename
         # would *always* match the actual obs_id folder the file was placed in.
         # Unfortunately this assumption was false.
-        # We have been trying to switch to the new uri scheme as it specifies the 
+        # We have been trying to switch to the new uri scheme as it specifies the
         # obs_id used in the folder path correctly.
         # The cherry on top is that the obs_id in the new style URI is not always correct either!
         # When we are looking up files we have some code which iterates through all of
@@ -1348,23 +1354,38 @@ def _download_from_s3(self, dataProduct, localPath, cache=True):
         # So in conclusion we can't trust the last char obs_id from the file or from the database
         # So with that in mind, hold your nose when reading the following:
 
-        # magic associations logic per Brian - 0-9/a-e we convert to 0
-        magicValues = "123456789abcde"
-        if obs_id[-1] in magicValues:
-            # We only replace the first occurrence in the folder
-            # The filename remains with the original obs_id as part of the name
-            new_obs_id = obs_id[:-1] + "0"
-            dataUri = dataUri.replace(obs_id, new_obs_id, 1)
-            obs_id = new_obs_id
-            log.warning("This data product's path may not have been properly identified %s" % dataUri)
+        info_lookup = None
+        sane_path = os.path.join("hst", "public", obs_id[:4], obs_id, filename)
+        try:
+            info_lookup = s3_client.head_object(Bucket=bkt_name, Key=sane_path, RequestPayer='requester')
+            bucketPath = sane_path
+        except self._botocore.exceptions.ClientError as e:
+            if e.response['Error']['Code'] != "404":
+                raise
+
+        if info_lookup is None:
+            # Unfortunately our file placement logic is anything but sane
+            # We put files in folders that don't make sense
+            for ch in (string.digits + string.ascii_lowercase):
+                # The last char of the obs_folder (observation id) can be any lowercase or numeric char
+                insane_obs = obs_id[:-1] + ch
+                insane_path = os.path.join("hst", "public", insane_obs[:4], insane_obs, filename)
+
+                try:
+                    info_lookup = s3_client.head_object(Bucket=bkt_name, Key=insane_path, RequestPayer='requester')
+                    bucketPath = insane_path
+                    break
+                except self._botocore.exceptions.ClientError as e:
+                    if e.response['Error']['Code'] != "404":
+                        raise
 
-        bucketPath = "hst/public/" + obs_id[:4] + dataUri.replace("mast:HST/product", "")
+        if info_lookup is None:
+            raise Exception("Unable to locate file!")
 
         # Unfortunately, we can't use the reported file size in the reported product.  STScI's backing
         # archive database (CAOM) is frequently out of date and in many cases omits the required information.
         # length = dataProduct["size"]
         # Instead we ask the webserver (in this case S3) what the expected content length is and use that.
-        info_lookup = s3_client.head_object(Bucket=bkt_name, Key=bucketPath, RequestPayer='requester')
         length = info_lookup["ContentLength"]
 
         if cache and os.path.exists(localPath):