Merge pull request #2374 from andamian/cadc_datalink

A number of CADC authentication and doc fixes

Author: bsipocz

astroquery/cadc/core.py

@@ -12,6 +12,7 @@
 import requests
 from numpy import ma
 from urllib.parse import urlencode
+from urllib.error import HTTPError
 
 from ..utils.class_or_instance import class_or_instance
 from ..utils import async_to_sync, commons
@@ -104,15 +105,15 @@ def __init__(self, url=None, tap_plus_handler=None, verbose=None,
 
         super(CadcClass, self).__init__()
         self.baseurl = url
+        # _auth_session contains the credentials that are used by both
+        # the cadc tap and cadc datalink services
         if auth_session:
             self._auth_session = auth_session
         else:
-            self._auth_session = None
+            self._auth_session = authsession.AuthSession()
 
     @property
     def cadctap(self):
-        if not self._auth_session:
-            self._auth_session = authsession.AuthSession()
         if not hasattr(self, '_cadctap'):
             if self.baseurl is None:
                 self.baseurl = get_access_url(self.CADCTAP_SERVICE_URI)
@@ -125,6 +126,13 @@ def cadctap(self):
                     self.baseurl, session=self._auth_session)
         return self._cadctap
 
+    @property
+    def cadcdatalink(self):
+        if not hasattr(self, '_datalink'):
+            self._datalink = pyvo.dal.adhoc.DatalinkService(
+                self.data_link_url, session=self._auth_session)
+        return self._datalink
+
     @property
     def data_link_url(self):
         if not hasattr(self, '_data_link_url'):
@@ -154,9 +162,9 @@ def login(self, user=None, password=None, certificate_file=None):
         # start with a new session
         if not isinstance(self.cadctap._session, (requests.Session,
                                                   authsession.AuthSession)):
-            raise TypeError('Cannot login with user provided session that is '
-                            'not an pyvo.authsession.AuthSession or '
-                            'requests.Session')
+            raise AttributeError('Cannot login with user provided session that is '
+                                 'not an pyvo.authsession.AuthSession or '
+                                 'requests.Session')
         if not certificate_file and not (user and password):
             raise AttributeError('login credentials missing (user/password '
                                  'or certificate)')
@@ -216,13 +224,25 @@ def logout(self, verbose=None):
         if verbose is not None:
             warnings.warn('verbose deprecated since 0.4.0')
 
-        # the only way to ensure complete logout is to start with a new
-        # session. This is mainly because of certificates. Adding cert
-        # argument to a session already in use does not force it to
-        # re-do the HTTPS hand shake
-        self.cadctap._session = authsession.AuthSession()
-        self.cadctap._session.update_from_capabilities(
-            self.cadctap.capabilities)
+        if isinstance(self._auth_session, pyvo.auth.AuthSession):
+            # Remove the existing credentials (if any)
+            # PyVO should provide this reset credentials functionality
+            # TODO - this should be implemented in PyVO to avoid this deep
+            # intrusion into that package
+            self._auth_session.credentials.credentials = \
+                {key: value for (key, value) in self._auth_session.credentials.credentials.items()
+                    if key == pyvo.auth.securitymethods.ANONYMOUS}
+        elif isinstance(self._auth_session, requests.Session):
+            # the only way to ensure complete logout is to start with a new
+            # session. This is mainly because of certificates. Removing cert
+            # argument to a session already in use does not force it to
+            # re-do the HTTPS hand shake
+            self._auth_session = requests.Session()
+            self.cadctap._session = self._auth_session
+            self.cadcdatalink._session = self._auth_session
+        else:
+            raise RuntimeError(
+                'Do not know how to log out from custom session')
 
     @class_or_instance
     def query_region_async(self, coordinates, radius=0.016666666666667*u.deg,
@@ -271,7 +291,7 @@ def query_name_async(self, name):
 
         Parameters
         ----------
-        name: str
+        name : str
                 name of object to query for
 
         Returns
@@ -353,7 +373,7 @@ def get_images(self, coordinates, radius,
         for fn in filenames:
             try:
                 images.append(fn.get_fits())
-            except requests.exceptions.HTTPError as err:
+            except (requests.exceptions.HTTPError, HTTPError) as err:
                 # Catch HTTPError if user is unauthorized to access file
                 log.debug(
                     "{} - Problem retrieving the file: {}".
@@ -456,7 +476,8 @@ def get_image_list(self, query_result, coordinates, radius):
                             range(0, len(publisher_ids), batch_size)):
             datalink = pyvo.dal.adhoc.DatalinkResults.from_result_url(
                 '{}?{}'.format(self.data_link_url,
-                               urlencode({'ID': pid_sublist}, True)))
+                               urlencode({'ID': pid_sublist}, True),
+                               session=self.cadcdatalink._session))
             for service_def in datalink.bysemantics('#cutout'):
                 access_url = service_def.access_url
                 if isinstance(access_url, bytes):  # ASTROPY_LT_4_1
@@ -523,7 +544,8 @@ def get_data_urls(self, query_result, include_auxiliaries=False):
             datalink = pyvo.dal.adhoc.DatalinkResults.from_result_url(
                 '{}?{}'.format(self.data_link_url,
                                urlencode({'ID': pid_sublist,
-                                          'REQUEST': 'downloads-only'}, True)))
+                                          'REQUEST': 'downloads-only'}, True)),
+                session=self.cadcdatalink._session)
             for service_def in datalink:
                 if service_def.semantics in ['http://www.opencadc.org/caom2#pkg', '#package']:
                     # TODO http://www.openadc.org/caom2#pkg has been replaced
@@ -582,7 +604,8 @@ def get_table(self, table, verbose=None):
             if table == t.name:
                 return t
 
-    def exec_sync(self, query, maxrec=None, uploads=None, output_file=None):
+    def exec_sync(self, query, maxrec=None, uploads=None, output_file=None,
+                  output_format='votable'):
         """
         Run a query and return the results or save them in an output_file
 
@@ -592,10 +615,13 @@ def exec_sync(self, query, maxrec=None, uploads=None, output_file=None):
             SQL to execute
         maxrec : int
             the maximum records to return. defaults to the service default
-        uploads:
+        uploads :
             Temporary tables to upload and run with the queries
-        output_file: str or file handler:
+        output_file : str or file handler
             File to save the results to
+        output_format :
+            Format of the output (default is basic). Must be one
+            of the formats supported by `astropy.table`
 
         Returns
         -------
@@ -611,10 +637,12 @@ def exec_sync(self, query, maxrec=None, uploads=None, output_file=None):
         result = response.to_table()
         if output_file:
             if isinstance(output_file, str):
-                with open(output_file, 'bw') as f:
-                    f.write(result)
+                fname = output_file
+            elif hasattr(output_file, 'name'):
+                fname = output_file.name
             else:
-                output_file.write(result)
+                raise AttributeError('Not a valid file name or file handler')
+            result.write(fname, format=output_format, overwrite=True)
         return result
 
     def create_async(self, query, maxrec=None, uploads=None):
@@ -681,9 +709,9 @@ def run_query(self, query, operation, output_file=None,
             when the job is executed in asynchronous mode,
             this flag specifies whether the execution will wait until results
             are available
-        upload_resource: str, optional, default None
+        upload_resource : str, optional, default None
             resource to be uploaded to UPLOAD_SCHEMA
-        upload_table_name: str, required if uploadResource is provided,
+        upload_table_name : str, required if uploadResource is provided,
             default None
             resource temporary table name associated to the uploaded resource
 
@@ -712,7 +740,7 @@ def load_async_job(self, jobid, verbose=None):
             warnings.warn('verbose deprecated since 0.4.0')
 
         return pyvo.dal.AsyncTAPJob('{}/async/{}'.format(
-            self.cadctap.baseurl, jobid))
+            self.cadctap.baseurl, jobid), session=self._auth_session)
 
     def list_async_jobs(self, phases=None, after=None, last=None,
                         short_description=True, verbose=None):
@@ -721,13 +749,13 @@ def list_async_jobs(self, phases=None, after=None, last=None,
 
         Parameters
         ----------
-        phases: list of str
+        phases : list of str
             Union of job phases to filter the results by.
-        after: datetime
+        after : datetime
             Return only jobs created after this datetime
-        last: int
+        last : int
             Return only the most recent number of jobs
-        short_description: flag - True or False
+        short_description : flag - True or False
             If True, the jobs in the list will contain only the information
             corresponding to the TAP ShortJobDescription object (job ID, phase,
             run ID, owner ID and creation ID) whereas if False, a separate GET
@@ -783,12 +811,19 @@ def get_access_url(service, capability=None):
     """
     Returns the URL corresponding to a service by doing a lookup in the cadc
     registry. It returns the access URL corresponding to cookie authentication.
-    :param service: the service the capability belongs to. It can be identified
-    by a CADC uri ('ivo://cadc.nrc.ca/) which is looked up in the CADC registry
-    or by the URL where the service capabilities is found.
-    :param capability: uri representing the capability for which the access
-    url is sought
-    :return: the access url
+
+    Parameters
+    ----------
+    service : str
+        the service the capability belongs to. It can be identified
+        by a CADC uri ('ivo://cadc.nrc.ca/) which is looked up in the CADC registry
+        or by the URL where the service capabilities is found.
+    capability : str
+        uri representing the capability for which the access url is sought
+
+    Returns
+    -------
+    The access url
 
     Note
     ------

astroquery/cadc/tests/test_cadctap.py

@@ -22,7 +22,7 @@
 import requests
 try:
     pyvo_OK = True
-    from pyvo.auth import authsession
+    from pyvo.auth import authsession, securitymethods
     from astroquery.cadc import Cadc, conf
     import astroquery.cadc.core as cadc_core
 except ImportError:
@@ -134,13 +134,14 @@ def test_list_async_jobs():
 
 @patch('astroquery.cadc.core.get_access_url',
        Mock(side_effect=lambda x, y=None: 'https://some.url'))
+@patch('astroquery.cadc.core.pyvo.dal.TAPService.capabilities', [])  # TAP capabilities not needed
+@patch('astroquery.cadc.core.pyvo.dal.adhoc.DatalinkService.capabilities', [])  # DL capabilities not needed
 @pytest.mark.skipif(not pyvo_OK, reason='not pyvo_OK')
 def test_auth():
     # the Cadc() will cause a remote data call to TAP service capabilities
     # To avoid this, use an anonymous session and replace it with an
     # auth session later
-    cadc = Cadc(auth_session=requests.Session())
-    cadc.cadctap._session = authsession.AuthSession()
+    cadc = Cadc()
     user = 'user'
     password = 'password'
     cert = 'cert'
@@ -153,8 +154,13 @@ def test_auth():
     cadc.login(certificate_file=cert)
     assert cadc.cadctap._session.credentials.get(
         'ivo://ivoa.net/sso#tls-with-certificate').cert == cert
+    assert cadc.cadcdatalink._session.credentials.get(
+        'ivo://ivoa.net/sso#tls-with-certificate').cert == cert
     # reset and try with user password/cookies
-    cadc.cadctap._session = authsession.AuthSession()
+    cadc.logout()
+    for service in (cadc.cadctap, cadc.cadcdatalink):
+        assert len(service._session.credentials.credentials) == 1
+        assert securitymethods.ANONYMOUS in service._session.credentials.credentials.keys()
     post_mock = Mock()
     cookie = 'ABC'
     mock_resp = Mock()
@@ -166,6 +172,13 @@ def test_auth():
     assert cadc.cadctap._session.credentials.get(
         'ivo://ivoa.net/sso#cookie').cookies[cadc_core.CADC_COOKIE_PREFIX] == \
         '"{}"'.format(cookie)
+    assert cadc.cadcdatalink._session.credentials.get(
+        'ivo://ivoa.net/sso#cookie').cookies[cadc_core.CADC_COOKIE_PREFIX] == \
+        '"{}"'.format(cookie)
+    cadc.logout()
+    for service in (cadc.cadctap, cadc.cadcdatalink):
+        assert len(service._session.credentials.credentials) == 1
+        assert securitymethods.ANONYMOUS in service._session.credentials.credentials.keys()
 
 
 # make sure that caps is reset at the end of the test
@@ -204,6 +217,8 @@ def raise_for_status(self):
 
 @patch('astroquery.cadc.core.get_access_url',
        Mock(side_effect=lambda x, y=None: 'https://some.url'))
+@patch('astroquery.cadc.core.pyvo.dal.adhoc.DatalinkService',
+       Mock(return_value=Mock(capabilities=[])))  # DL capabilities not needed
 @pytest.mark.skipif(not pyvo_OK, reason='not pyvo_OK')
 def test_get_data_urls():
 
@@ -283,6 +298,10 @@ def test_misc():
 
 @patch('astroquery.cadc.core.get_access_url',
        Mock(side_effect=lambda x, y=None: 'https://some.url'))
+@patch('astroquery.cadc.core.pyvo.dal.TAPService',
+       Mock(return_value=Mock(capabilities=[])))  # TAP capabilities not needed
+@patch('astroquery.cadc.core.pyvo.dal.adhoc.DatalinkService',
+       Mock(return_value=Mock(capabilities=[])))  # DL capabilities not needed
 @pytest.mark.skipif(not pyvo_OK, reason='not pyvo_OK')
 def test_get_image_list():
     def get(*args, **kwargs):
@@ -371,7 +390,7 @@ def test_exec_sync():
     votable.to_xml(buffer)
     cadc = Cadc(auth_session=requests.Session())
     response = Mock()
-    response.to_table.return_value = buffer.getvalue()
+    response.to_table.return_value = table.to_table()
     cadc.cadctap.search = Mock(return_value=response)
     output_file = '{}/test_vooutput.xml'.format(tempfile.tempdir)
     cadc.exec_sync('some query', output_file=output_file)

astroquery/cadc/tests/test_cadctap_remote.py

@@ -151,6 +151,12 @@ def test_login_with_user_password(self):
                 format(now.strftime('%Y-%m-%dT%H:%M:%S.%f'))
             result = cadc.exec_sync(query)
             assert len(result) == 0
+            # login in again
+            cadc.login(os.environ['CADC_USER'], os.environ['CADC_PASSWD'])
+            query = "select top 1 * from caom2.Plane where metaRelease>'{}'". \
+                format(now.strftime('%Y-%m-%dT%H:%M:%S.%f'))
+            result = cadc.exec_sync(query)
+            assert len(result) == 1
 
     @pytest.mark.skipif(one_test, reason='One test mode')
     @pytest.mark.skipif('CADC_CERT' not in os.environ,