rest.py 77.6 KB
Newer Older
1 2
# -*- coding: utf-8 -*-
# Copyright (C) Brian Moe, Branson Stephens (2015)
3
#
4
# This file is part of gracedb
5
#
6 7 8 9
# gracedb is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
10
#
11 12 13 14 15 16 17
# It is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with gracedb.  If not, see <http://www.gnu.org/licenses/>.
18

19
import six.moves.http_client, socket, ssl
20 21
import mimetypes
import urllib
22 23 24 25
import os
import sys
if os.name == 'posix':
    import pwd
26
import json
27
from six.moves.urllib.parse import urlparse, urlencode
28
from base64 import b64encode
29 30
import six
from six.moves import map
31

32
from .exceptions import HTTPError
33
from .version import __version__
34 35
from .utils import event_or_superevent, handle_str_or_list_arg, safe_netrc, \
    cleanListInput, get_dt_from_openssl_output, is_expired
36

37
DEFAULT_SERVICE_URL = "https://gracedb.ligo.org/api/"
38
DEFAULT_BASIC_SERVICE_URL = "https://gracedb.ligo.org/apibasic/"
Leo Pound Singer's avatar
Leo Pound Singer committed
39
KNOWN_TEST_HOSTS = ['moe.phys.uwm.edu', 'embb-dev.ligo.caltech.edu', 'simdb.phys.uwm.edu',]
40

41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59
#---------------------------------------------------------------------
# This monkey patch forces TLSv1 if the python version is 2.6.6.
# It was introduced because clients connection from CIT *occasionally*
# try to use SSLv3.  See:
# http://stackoverflow.com/questions/18669457/python-httplib-ssl23-get-server-hellounknown-protocol
#---------------------------------------------------------------------
if sys.hexversion < 0x20709f0:
    wrap_socket_orig = ssl.wrap_socket
    def wrap_socket_patched(sock, keyfile=None, certfile=None,
                            server_side=False, cert_reqs=ssl.CERT_NONE,
                            ssl_version=ssl.PROTOCOL_TLSv1, ca_certs=None,
                            do_handshake_on_connect=True,
                            suppress_ragged_eofs=True):
        return wrap_socket_orig(sock, keyfile, certfile, server_side,
                                cert_reqs, ssl_version, ca_certs,
                                do_handshake_on_connect,
                                suppress_ragged_eofs)
    ssl.wrap_socket = wrap_socket_patched

60
#-----------------------------------------------------------------
61
# HTTP/S Proxy classes
62 63
# Taken from: http://code.activestate.com/recipes/456195/

64
class ProxyHTTPConnection(six.moves.http_client.HTTPConnection):
65 66 67 68 69 70

    _ports = {'http' : 80, 'https' : 443}

    def request(self, method, url, body=None, headers={}):
        #request is called before connect, so can interpret url and get
        #real host/port to be used to make CONNECT request to proxy
71 72 73 74
        o = urlparse(url)
        proto = o.scheme
        port = o.port
        host = o.hostname
75
        if proto is None:
76
            raise ValueError("unknown URL type: %s" % url)
77 78 79 80
        if port is None:
            try:
                port = self._ports[proto]
            except KeyError:
81
                raise ValueError("unknown protocol for: %s" % url)
82 83
        self._real_host = host
        self._real_port = port
84
        six.moves.http_client.HTTPConnection.request(self, method, url, body, headers)
85 86 87


    def connect(self):
88
        six.moves.http_client.HTTPConnection.connect(self)
89 90 91 92 93 94 95 96 97
        #send proxy CONNECT request
        self.send("CONNECT %s:%d HTTP/1.0\r\n\r\n" % (self._real_host, self._real_port))
        #expect a HTTP/1.0 200 Connection established
        response = self.response_class(self.sock, strict=self.strict, method=self._method)
        (version, code, message) = response._read_status()
        #probably here we can handle auth requests...
        if code != 200:
            #proxy returned and error, abort connection, and raise exception
            self.close()
98
            raise socket.error("Proxy connection failed: %d %s" % (code, message.strip()))
99 100
        #eat up header block from proxy....
        while True:
101
            #should not use directly fp probably
102 103 104 105 106 107 108
            line = response.fp.readline()
            if line == '\r\n': break

class ProxyHTTPSConnection(ProxyHTTPConnection):

    default_port = 443

109 110
    def __init__(self, host, port = None, key_file = None, cert_file = None,
        strict = None, context = None):
111 112 113
        ProxyHTTPConnection.__init__(self, host, port)
        self.key_file = key_file
        self.cert_file = cert_file
114
        self.context = context
115 116 117 118

    def connect(self):
        ProxyHTTPConnection.connect(self)
        #make the sock ssl-aware
119 120
        if sys.hexversion < 0x20709f0:
            ssl = socket.ssl(self.sock, self.key_file, self.cert_file)
121
            self.sock = six.moves.http_client.FakeSocket(self.sock, ssl)
122 123
        else:
            self.sock = self.context.wrap_socket(self.sock)
124

125 126 127
# Hacky fake response type for kludging response object
# Has json() and read() methods
class FakeResponse(object):
128 129 130
    def __init__(self, json, status, *args, **kwargs):
        self._data = json
        self.status = status
131 132 133 134
    def json(self): return self._data
    def read(self): return json.dumps(self._data)


135 136 137 138 139
#-----------------------------------------------------------------
# Generic GSI REST

class GsiRest(object):
    """docstring for GracedbRest"""
140 141 142 143
    def __init__(self,
            url=DEFAULT_SERVICE_URL,
            proxy_host=None,
            proxy_port=3128,
144
            cred=None):
145 146
        if not cred:
            cred = findUserCredentials()
147 148 149 150
        if not cred:
            msg = "\nERROR: No certificate (or proxy) found. \n\n"
            msg += "Please run ligo-proxy-init or grid-proxy-init (as appropriate) "
            msg += "to generate one.\n\n"
151
            self.output_and_die(msg)
152 153 154 155 156
        if isinstance(cred, (list, tuple)):
            self.cert, self.key = cred
        elif cred:
            self.cert = self.key = cred

157 158 159 160 161
        o = urlparse(url)
        port = o.port
        host = o.hostname
        port = port or 443

162 163 164 165 166 167 168
        # Versions of Python earlier than 2.7.9 don't use SSL Context
        # objects for this purpose, and do not do any server cert verification.
        ssl_context = None
        if sys.hexversion >= 0x20709f0:
            # Use the new method with SSL Context
            # Prepare SSL context
            ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLSv1)
169 170
            try:
                ssl_context.load_cert_chain(self.cert, self.key)
171
            except ssl.SSLError as e:
172 173 174
                msg = "\nERROR: Unable to load cert/key pair. \n\n"
                msg += "Please run ligo-proxy-init or grid-proxy-init again "
                msg += "or make sure your robot certificate is readable.\n\n"
175
                self.output_and_die(msg)
176 177 178 179 180 181 182 183
            # Generally speaking, test boxes use cheap/free certs from the LIGO CA.
            # These cannot be verified by the client.
            if host in KNOWN_TEST_HOSTS:
                ssl_context.verify_mode = ssl.CERT_NONE
            else:
                ssl_context.verify_mode = ssl.CERT_REQUIRED
                ssl_context.check_hostname = True
                # Find the various CA cert bundles stored on the system
184
                ssl_context.load_default_certs()
185 186 187 188

            if proxy_host:
                self.connector = lambda: ProxyHTTPSConnection(proxy_host, proxy_port, context=ssl_context)
            else:
189
                self.connector = lambda: six.moves.http_client.HTTPSConnection(host, port, context=ssl_context)
190
        else:
191 192
            # Using and older version of python. We'll pass in the cert and key files.
            if proxy_host:
193
                self.connector = lambda: ProxyHTTPSConnection(proxy_host, proxy_port,
194 195
                    key_file=self.key, cert_file=self.cert)
            else:
196
                self.connector = lambda: six.moves.http_client.HTTPSConnection(host, port,
197 198
                    key_file=self.key, cert_file=self.cert)

199 200 201

    def getConnection(self):
        return self.connector()
202

203 204
    # When there is a problem with the SSL connection or cert authentication,
    # either conn.request() or conn.getresponse() will throw an exception.
205
    # The following two wrappers are intended to catch these exceptions and
206 207 208 209 210
    # return an intelligible error message to the user.
    # A wrapper for getting the response:
    def get_response(self, conn):
        try:
            return conn.getresponse()
211
        except ssl.SSLError as e:
212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227
            # Check for a valid user proxy cert.
            expired, error = is_expired(self.cert)
            if expired is not None:
                if expired:
                    msg = "\nERROR: Your certificate or proxy has expired. \n\n"
                    msg += "Please run ligo-proxy-init or grid-proxy-init (as appropriate) "
                    msg += "to generate a fresh one.\n\n"
                else:
                    msg = "\nERROR \n\n"
                    msg += "Your certificate appears valid, but there was a problem "
                    msg += "establishing a secure connection: \n\n"
                    msg += "%s \n\n" % str(e)
            else:
                msg = "\nERROR \n\n"
                msg += "Unable to check certificate expiry date: %s \n\n" %error
                msg += "Problem establishing secure connection: %s \n\n" % str(e)
228
            self.output_and_die(msg)
229 230 231 232 233

    # A wrapper for making the request.
    def make_request(self, conn, *args, **kwargs):
        try:
            conn.request(*args, **kwargs)
234
        except ssl.SSLError as e:
235 236
            msg = "\nERROR \n\n"
            msg += "Problem establishing secure connection: %s \n\n" % str(e)
237
            self.output_and_die(msg)
238

239
    def request(self, method, url, body=None, headers=None, priming_url=None):
240 241 242 243 244 245 246 247 248 249 250
        # Bug in Python (versions < 2.7.1 (?))
        # http://bugs.python.org/issue11898
        # if the URL is unicode and the body of a request is binary,
        # the POST/PUT action fails because it tries to concatenate
        # the two which fails due to encoding problems.
        # Workaround is to cast all URLs to str.
        # This is probably bad in general,
        # but for our purposes, today, this will do.
        url = url and str(url)
        priming_url = priming_url and str(priming_url)

251 252
        # Add version string to user-agent header
        version_header = {'User-Agent': 'gracedb-client/{version}'.format(
253
            version=__version__)}
254 255 256 257 258
        if headers is None:
            headers = version_header
        else:
            headers.update(version_header)

259
        conn = self.getConnection()
260
        if priming_url:
261 262 263
            priming_header = {'connection': 'keep-alive'}
            priming_header.update(version_header)
            self.make_request(conn, "GET", priming_url, headers=priming_header)
264
            response = self.get_response(conn)
265 266 267 268 269
            if response.status != 200:
                response = self.adjustResponse(response)
            else:
                # Throw away the response and make sure to read the body.
                response = response.read()
270 271
        self.make_request(conn, method, url, body, headers or {})
        response = self.get_response(conn)
272 273 274 275 276
        return self.adjustResponse(response)

    def adjustResponse(self, response):
#       XXX WRONG.
        if response.status >= 400:
277 278 279 280 281 282 283
            response_content = response.read()
            if response.getheader('x-throttle-wait-seconds', None):
                try:
                    rdict = json.loads(response_content)
                    rdict['retry-after'] = response.getheader('x-throttle-wait-seconds')
                    response_content = json.dumps(rdict)
                except:
284
                    pass
285
            raise HTTPError(response.status, response.reason, response_content)
286
        response.json = lambda: self.load_json_or_die(response)
287 288 289 290 291 292 293 294 295 296 297 298 299 300 301
        return response

    def get(self, url, headers=None):
        return self.request("GET", url, headers=headers)

    def head(self, url, headers=None):
        return self.request("HEAD", url, headers=headers)

    def delete(self, url, headers=None):
        return self.request("DELETE", url, headers=headers)

    def options(self, url, headers=None):
        return self.request("OPTIONS", url, headers=headers)

    def post(self, *args, **kwargs):
302
        return self.post_or_put_or_patch("POST", *args, **kwargs)
303 304

    def put(self, *args, **kwargs):
305
        return self.post_or_put_or_patch("PUT", *args, **kwargs)
306

307 308 309 310 311
    def patch(self, *args, **kwargs):
        return self.post_or_put_or_patch("PATCH", *args, **kwargs)

    def post_or_put_or_patch(self, method, url, body=None, headers=None,
        files=None):
312 313 314 315
        headers = headers or {}
        if not files:
            # Simple urlencoded body
            if isinstance(body, dict):
Tanner Prestegard's avatar
Tanner Prestegard committed
316
            # XXX What about the headers in the params?
317
                if 'content-type' not in headers:
318 319
                    headers['content-type'] = "application/json"
                body = json.dumps(body)
320 321 322
        else:
            body = body or {}
            if isinstance(body, dict):
323
                body = list(body.items())
324
            content_type, body = encode_multipart_formdata(body, files)
Tanner Prestegard's avatar
Tanner Prestegard committed
325
            # XXX What about the headers in the params?
326 327 328
            headers = {
                'content-type': content_type,
                'content-length': str(len(body)),
Tanner Prestegard's avatar
Tanner Prestegard committed
329
                #'connection': 'keep-alive',
330 331 332
            }
        return self.request(method, url, body, headers)

333 334 335 336 337 338 339
    # A utility for writing out an error message to the user and then stopping
    # execution. This seems to behave sensibly in both the interpreter and in
    # a script.
    @classmethod
    def output_and_die(cls, msg):
        raise RuntimeError(msg)

340
    # Given an HTTPResponse object, try to read its content and interpret as
341 342 343 344 345 346
    # JSON--or die trying.
    @classmethod
    def load_json_or_die(cls, response):

        # First check that the response object actually exists.
        if not response:
347
            raise ValueError("No response object")
348 349 350

        # Next, try to read the content of the response.
        response_content = response.read()
Leo Pound Singer's avatar
Leo Pound Singer committed
351
        response_content = response_content.decode('utf-8')
352 353
        if not response_content:
            response_content = '{}'
354 355 356 357 358

        # Finally, try to create a dict by decoding the response as JSON.
        rdict = None
        try:
            rdict = json.loads(response_content)
359 360 361
        except ValueError:
            msg = "ERROR: got unexpected content from the server:\n"
            msg += response_content
362
            raise ValueError(msg)
363 364 365

        return rdict

366 367 368 369 370 371
#------------------------------------------------------------------
# GraceDB
#
# Example Gracedb REST client.

class GraceDb(GsiRest):
372
    """Example GraceDb REST client
373

374 375 376 377 378 379 380 381 382 383
    The GraceDB service URL may be passed to the constructor
    if an alternate GraceDb instance is desired:

        >>> g = GraceDb("https://alternate.gracedb.edu/api/")
        >>> r = g.ping()

    The proxy_host and proxy_port may also be passed in if accessing
    GraceDB behind a proxy. For other kwargs accepted by the constructor,
    consult the source code.
    """
384 385 386 387
    def __init__(self, service_url=DEFAULT_SERVICE_URL, proxy_host=None,
            proxy_port=3128, api_version=None, *args, **kwargs):
        super(GraceDb, self).__init__(service_url, proxy_host,
            proxy_port, *args, **kwargs)
388

389 390 391 392 393 394 395 396 397 398 399 400 401 402
        # Check version type
        if api_version is not None and not isinstance(api_version,
            six.string_types):
            raise TypeError('api_version should be a string')

        # Sets default and versioned service URLs
        # (self._service_url, self._versioned_service_url)
        self._set_service_url(service_url, api_version)

        # Set version
        self._api_version = api_version

        # Set service_info to None, will be obtained from the server when
        # the user takes an action which needs this information.
403 404
        self._service_info = None

405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427
    def _set_service_url(self, service_url, api_version):
        """Sets versioned and unversioned service URLs"""
        # Make sure path ends with '/'
        if not service_url.endswith('/'):
            service_url += '/'

        # Default service url (unversioned)
        self._service_url = service_url

        # Versioned service url (if version provided)
        self._versioned_service_url = service_url
        if api_version and api_version != 'default':
            # If api_version is 'default', that is equivalent to not setting
            # the version and indicates that the user wants to use the
            # default/non-versioned API
            self._versioned_service_url += (api_version + '/')

    @property
    def service_url(self):
        # Will be removed in the future
        print("DEPRECATED: this attribute has been moved to '_service_url'")
        return self._service_url

428 429
    @property
    def service_info(self):
430
        """Info from root of API. Should be a dict"""
431
        if not self._service_info:
432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458
            # try-except block takes user-specified API version to use and
            # checks whether that version is available on the server
            try:
                r = self.request("GET", self._versioned_service_url)
            except HTTPError as e:
                # If we get a 404 error, that means that the versioned
                # service URL was not found. We assume that this happened
                # because the user requested an unavailable API version.
                if (e.status == 404):
                    # Get versions from unversioned API root
                    r = self.request("GET", self._service_url)
                    available_api_versions = r.json().get('API_VERSIONS', None)
                    if available_api_versions:
                        err_msg = ('Bad API version. Available versions for '
                            'this server are: {0}').format(
                            available_api_versions)
                    else:
                        # Case where server doesn't have versions, for some
                        # reason.
                        err_msg = ('This server does not have a versioned API.'
                            ' Reinstantiate your client without a version.')

                    # Raise error
                    raise ValueError(err_msg)
                else:
                    # Not a 404 error, must be something else
                    raise e
459
            self._service_info = r.json()
460 461
        return self._service_info

462 463 464 465
    @property
    def api_versions(self):
        return self.service_info.get('API_VERSIONS')

466 467 468 469 470 471 472 473 474 475 476
    @property
    def links(self):
        return self.service_info.get('links')

    @property
    def templates(self):
        return self.service_info.get('templates')

    @property
    def groups(self):
        return self.service_info.get('groups')
477

478 479 480 481 482 483 484
    @property
    def pipelines(self):
        return self.service_info.get('pipelines')

    @property
    def searches(self):
        return self.service_info.get('searches')
485

486 487 488 489 490 491
    # Would like to call this 'labels' to keep in line with how
    # other properties are named, but it's already used for a function.
    @property
    def allowed_labels(self):
        return self.service_info.get('labels')

492
    @property
493 494
    def em_groups(self):
        return self.service_info.get('em-groups')
495 496 497 498 499 500 501 502 503 504 505 506 507

    @property
    def wavebands(self):
        return self.service_info.get('wavebands')

    @property
    def eel_statuses(self):
        return self.service_info.get('eel-statuses')

    @property
    def obs_statuses(self):
        return self.service_info.get('obs-statuses')

508 509 510 511
    @property
    def voevent_types(self):
        return self.service_info.get('voevent-types')

512 513 514 515
    @property
    def superevent_categories(self):
        return self.service_info.get('superevent-categories')

516 517
    def request(self, method, *args, **kwargs):
        if method.lower() in ['post', 'put']:
518
            kwargs['priming_url'] = self._service_url
519 520
        return GsiRest.request(self, method, *args, **kwargs)

521
    def _getCode(self, input_value, code_dict):
522 523 524 525 526 527
        """
        Check if input is valid and return coded version if it is
        code_dict is dict of { code : descriptive_name }
        """
        # Quick check for simple case where it's already coded
        if input_value in code_dict:
528
            return input_value
529 530 531 532 533 534 535 536 537 538 539

        # Loop over code_dict items, if we match either the key or
        # value (case-insensitive), return the code.
        input_lower = input_value.lower()
        for code, display in six.iteritems(code_dict):
            if (input_lower == code.lower() or
                input_lower == display.lower()):
                return code

        # Not found, return None
        return None
540

541
    # Search and filecontents are optional when creating an event.
542
    def createEvent(self, group, pipeline, filename, search=None, labels=None,
543
                    offline=False, filecontents=None, **kwargs):
544 545 546 547
        """Create a new GraceDB event

        Required args: group, pipeline, filename

548
        Optional args: search, labels, offline, filecontents
549

550 551 552 553 554
        The filename is the path to a file containing information about the
        event. The values for 'group', 'pipeline', and 'search' are restricted
        to those stored in the database. 'labels' should be a list of strings
        corresponding to labels (values restricted to those stored in the
        database).  'labels' can be a string if only a single label is
555 556
        being applied. 'offline' is a boolean which indicates whether the
        analysis is offline (True) or online/low-latency (False).
557 558 559 560

        Example:

            >>> g = GraceDb()
561 562
            >>> r = g.createEvent('CBC', 'gstlal', '/path/to/something.xml',
            ... labels='INJ', 'LowMass')
563 564 565 566
            >>> r.status
            201

        """
567 568 569
        errors = []
        if group not in self.groups:
            errors += ["bad group"]
570 571 572 573
        if pipeline not in self.pipelines:
            errors += ["bad pipeline"]
        if search and search not in self.searches:
            errors += ["bad search"]
574 575 576
        # Process offline arg
        if not isinstance(offline, bool):
            errors += ["offline should be True or False"]
577 578 579
        # Process label args - convert non-empty strings to list
        # to ensure consistent processing
        if labels:
580
            if isinstance(labels, six.string_types):
581 582 583 584 585 586 587 588 589 590 591 592 593 594 595
                # Convert to list
                labels = [labels]
            elif isinstance(labels, list):
                pass
            else:
                # Raise exception instead of adding errors. The next for loop
                # will break (before errors exception is raised) if labels
                # is of the wrong type
                raise TypeError("labels arg is {0}, should be str or list" \
                    .format(type(labels)))
            # Check labels against those in database
            for l in labels:
                if l not in self.allowed_labels:
                    raise NameError(("Label '{0}' does not exist in the "
                        "database").format(l))
596 597 598 599
        if errors:
            # XXX Terrible error messages / weak exception type
            raise Exception(str(errors))
        if filecontents is None:
600 601 602 603
            if filename == '-':
                filename = 'initial.data'
                filecontents = sys.stdin.read()
            else:
604
                filecontents = open(filename, 'rb').read()
605

606
        fields = [
607 608 609 610
            ('group', group),
            ('pipeline', pipeline),
            ('offline', offline),
        ]
611 612
        if search:
            fields.append(('search', search))
613 614 615
        if labels:
            for l in labels:
                fields.append(('labels', l))
616 617

        # Update fields with additional keyword arguments
618
        for key, value in six.iteritems(kwargs):
619 620
            fields.append((key, value))

621 622 623 624 625 626
        files = [('eventFile', filename, filecontents)]
        # Python httplib bug?  unicode link
        uri = str(self.links['events'])
        return self.post(uri, fields, files=files)

    def replaceEvent(self, graceid, filename, filecontents=None):
627 628 629 630 631
        """Replace an existing GraceDB event

        Required args: graceid, filename

        This function uploads a new event file, hence changing the basic details
632 633
        of an existing event.

634 635 636
        Example:

            >>> g = GraceDb()
637
            >>> r = g.replaceEvent('T101383', '/path/to/new/something.xml')
638 639

        """
640
        if filecontents is None:
641 642
            # Note: not allowing filename '-' here.  We want the event datafile
            # to be versioned.
643 644 645 646 647 648
            filecontents = open(filename, 'rb').read()
        return self.put(
                self.templates['event-detail-template'].format(graceid=graceid),
                files=[('eventFile', filename, filecontents)])

    def event(self, graceid):
649 650
        """Get information about a specific event

651
        Args: graceid
652 653 654 655 656 657 658

        Example:

            >>> g = GraceDb()
            >>> event_dict = g.event('T101383').json()

        """
659 660 661
        return self.get(
                self.templates['event-detail-template'].format(graceid=graceid))

662
    def events(self, query=None, orderby=None, count=None, columns=None):
663 664 665
        """Get a iterator of events in response to a query

        This function returns an iterator which yields event dictionaries.
666 667
        Optional arguments are query, orderby, count, and columns. The
        columns argument is a comma separated list of attributes that the
668 669 670 671 672 673 674 675 676 677 678
        user would like in each event dictionary. If columns are not specified,
        all attributes of the events are returned.

        Example:

            >>> g = GraceDb()
            >>> for event in g.events('ER5 submitter: "gstlalcbc"', columns='graceid,far,gpstime'):
            ...     print "%s %e %d" % (event['graceid'], event['far'], event['gpstime'])


        """
679 680 681 682
        uri = self.links['events']
        qdict = {}
        if query:   qdict['query'] = query
        if count:   qdict['count'] = count
683
        if orderby: qdict['sort'] = orderby
684
        if columns: qdict['columns'] = columns
685
        if qdict:
686
            uri += "?" + urlencode(qdict)
687 688
        while uri:
            response = self.get(uri).json()
689 690
            events = response.get('events',[])
            uri = response.get('links',{}).get('next')
691 692 693 694
            for event in events:
                 yield event

    def numEvents(self, query=None):
695
        """Get the number of events satisfying a query
696 697 698 699 700 701 702 703

        Example:

            >>> g = GraceDb()
            >>> g.numEvents('ER5 submitter: "gstlalcbc"')
            213

        """
704 705
        uri = self.links['events']
        if query:
706
            uri += "?" + urlencode({'query': query})
707 708
        return self.get(uri).json()['numRows']

709
    def createSuperevent(self, t_start, t_0, t_end, preferred_event,
710
        category='production', events=[], labels=None):
711
        """
Tanner Prestegard's avatar
Tanner Prestegard committed
712
        Create a superevent.
713

714 715 716 717
        Signature:
        createSuperevent(t_start, t_0, t_end, preferred_event,
            events=[], labels=None)

Tanner Prestegard's avatar
Tanner Prestegard committed
718 719 720 721 722 723
        Arguments:
            t_start:         t_start of superevent
            t_0:             t_0 of superevent
            t_end:           t_end of superevent
            preferred_event: graceid corresponding to event which will be set
                             as the preferred event for this superevent
724
            category:        superevent category ('production', 'test', 'mdc')
Tanner Prestegard's avatar
Tanner Prestegard committed
725
            events:          list of graceids corresponding to events which
726
                             should be attached to this superevent (optional)
Tanner Prestegard's avatar
Tanner Prestegard committed
727
            labels:          list of labels which should be attached to this
728
                             superevent at creation (optional)
729 730 731 732

        Example:

            >>> g = GraceDb()
733 734 735
            >>> r = g.createSuperevent(1, 2, 3, 'G123456',
            ... category='production', events=['G100', 'G101'],
            ... labels=['EM_READY', 'DQV'])
736 737 738 739 740 741 742 743 744 745 746 747 748 749 750 751 752 753 754 755 756 757 758 759 760 761 762
            >>> r.status
            201
        """
        # Process label args - convert non-empty strings to list
        # to ensure consistent processing
        if labels:
            if isinstance(labels, six.string_types):
                labels = [labels]
            elif isinstance(labels, list):
                pass
            else:
                raise TypeError("labels arg is {0}, should be str or list" \
                    .format(type(labels)))
            # Check labels against those in database
            for l in labels:
                if l not in self.allowed_labels:
                    raise NameError(("Label '{0}' does not exist in the "
                        "database").format(l))
        if events:
            if isinstance(events, six.string_types):
                events = [events]
            elif isinstance(events, list):
                pass
            else:
                raise TypeError("events arg is {0}, should be str or list" \
                    .format(type(events)))

763 764 765 766 767 768 769
        # validate category, convert to short form if necessary
        category = self._getCode(category.lower(),
            self.superevent_categories)
        if not category:
            raise ValueError("category must be one of: {0}".format(
                list(six.itervalues(self.superevent_categories))))

770 771 772 773 774
        # Set up request body for POST
        request_body = {
            't_start': t_start,
            't_0': t_0,
            't_end': t_end,
775
            'preferred_event': preferred_event,
776
            'category': category,
777 778 779 780 781 782 783 784 785 786 787 788
        }
        if events:
            request_body['events'] = events
        if labels:
            request_body['labels'] = labels

        # Python httplib bug?  unicode link
        uri = self.links['superevents']
        return self.post(uri, body=request_body)

    def updateSuperevent(self, superevent_id, t_start=None, t_0=None,
        t_end=None, preferred_event=None):
Tanner Prestegard's avatar
Tanner Prestegard committed
789 790 791
        """
        Update a superevent's parameters.

792 793 794 795
        Signature:
        updateSuperevent(superevent_id, t_start=None, t_0=None,
            t_end=None, preferred_event=None)

Tanner Prestegard's avatar
Tanner Prestegard committed
796 797 798 799 800 801 802 803 804 805 806 807 808 809 810 811 812 813 814
        Arguments:
            superevent_id:   id of superevent to update
            t_start:         t_start of superevent
            t_0:             t_0 of superevent
            t_end:           t_end of superevent
            preferred_event: graceid corresponding to event which will be set
                             as the preferred event for this superevent

            Any combination of these parameters (other than superevent_id)
            may be used; none are required.

        Example:

            >>> g = GraceDb()
            >>> r = g.updateSuperevent('S0001', t_start=12, preferred_event=
            ... 'G654321')
            >>> r.status
            200
        """
815 816 817 818
        # Make sure that at least one parameter is provided
        if not (t_start or t_0 or t_end or preferred_event):
            raise ValueError('Provide at least one of t_start, t_0, t_end, or '
                'preferred_event')
819 820 821 822 823 824 825 826 827 828 829 830 831 832 833

        request_body = {}
        if t_start:
            request_body['t_start'] = t_start
        if t_0:
            request_body['t_0'] = t_0
        if t_end:
            request_body['t_end'] = t_end
        if preferred_event:
            request_body['preferred_event'] = preferred_event
        template = self.templates['superevent-detail-template']
        uri = template.format(superevent_id=superevent_id)
        return self.patch(uri, body=request_body)

    def addEventToSuperevent(self, superevent_id, graceid):
Tanner Prestegard's avatar
Tanner Prestegard committed
834 835 836 837 838
        """
        Add an event to a superevent. Events can only be part of one superevent
        so the server will throw an error if the event is part of another
        superevent already.

839 840 841
        Signature:
            addEventToSuperevent(superevent_id, graceid)

Tanner Prestegard's avatar
Tanner Prestegard committed
842 843 844 845 846 847 848 849 850 851 852
        Arguments:
            superevent_id: id of superevent to which the event will be added
            graceid:       graceid of event to add to superevent

        Example:

            >>> g = GraceDb()
            >>> r = addEventToSuperevent('S0001', 'G123456')
            >>> r.status
            201
        """
853 854 855 856 857 858
        request_body = {'event': graceid}
        template = self.templates['superevent-event-list-template']
        uri = template.format(superevent_id=superevent_id)
        return self.post(uri, body=request_body)

    def removeEventFromSuperevent(self, superevent_id, graceid):
Tanner Prestegard's avatar
Tanner Prestegard committed
859 860 861
        """
        Remove an event from a superevent.

862 863 864
        Signature:
            removeEventFromSuperevent(superevent_id, graceid)

Tanner Prestegard's avatar
Tanner Prestegard committed
865 866 867 868 869 870 871 872 873 874
        Arguments:
            superevent_id, graceid

        Example:

            >>> g = GraceDb()
            >>> r = removeEventFromSuperevent('S0001', 'G123456')
            >>> r.status
            204
        """
875 876 877 878
        template = self.templates['superevent-event-detail-template']
        uri = template.format(superevent_id=superevent_id, graceid=graceid)
        return self.delete(uri)

879 880
    def superevent(self, superevent_id):
        """
Tanner Prestegard's avatar
Tanner Prestegard committed
881 882
        Get information about a specific superevent.

883 884 885
        Signature:
            superevent(superevent_id)

Tanner Prestegard's avatar
Tanner Prestegard committed
886 887
        Arguments:
            superevent_id
888 889 890 891

        Example:

            >>> g = GraceDb()
Tanner Prestegard's avatar
Tanner Prestegard committed
892
            >>> superevent = g.superevent('S0001').json()
893 894 895 896
        """
        return self.get(self.templates['superevent-detail-template'].format(
            superevent_id=superevent_id))

897
    def superevents(self, query='', orderby=[], count=None, columns=[]):
898
        """
Tanner Prestegard's avatar
Tanner Prestegard committed
899
        Get an iterator of superevents in response to a query.
900

901 902 903
        Signature:
            superevents(query='', orderby=[], count=None, columns=[])

904 905 906 907 908 909 910 911 912 913 914 915 916 917 918 919 920 921
        Arguments:
            query: query string for filtering superevents (same as on the
                   web interface)
            orderby: list of strings corresponding to attribute(s) of the
                     superevents used to order the results (optional).
                     Available options: created, t_start, t_0, t_end, is_gw,
                     id, preferred_event. Default is ascending order, but
                     prefix any option with "-" to apply descending order.
            count: each generator iteration will yield this many objects
                   (optional; default determined by the server)
            columns: which attributes of the superevents to return
                     (default: all).

        Example:

            >>> g = GraceDb()
            >>> for s in g.superevents(query='is_gw=True', orderby=['-preferred_event'], columns='superevent_id,events'):
            ...     print(s['superevent_id'])
922
        """
923 924 925 926 927 928 929 930 931 932 933 934 935 936 937 938 939 940 941 942 943 944 945 946 947 948 949 950 951 952 953
        # If columns is a comma-separated string, split it to a list
        if isinstance(columns, six.string_types):
            columns = columns.split(',')

        # If orderby is a list (should be), convert it to a comma-separated
        # string (that's what the server expects)
        if isinstance(orderby, list):
            orderby = ",".join(orderby)

        # Get URI
        uri = self.links['superevents']

        # Compile URL parameters
        qdict = {}
        if query:   qdict['query'] = query
        if count:   qdict['count'] = count
        if orderby: qdict['sort'] = orderby
        if qdict:
            uri += "?" + urlencode(qdict)

        # Get superevent information and construct a generator
        while uri:
            response = self.get(uri).json()
            superevents = response.get('superevents',[])
            uri = response.get('links',{}).get('next')
            for superevent in superevents:
                # If columns are specified, only return specific values
                if columns:
                    yield {k: superevent[k] for k in columns}
                else:
                    yield superevent
954

955 956 957 958 959 960 961 962 963 964 965 966 967 968 969
    def confirm_superevent_as_gw(self, superevent_id):
        """
        Upgrade a superevent's state to 'confirmed GW'.
        Requires specific server-side permissions.

        Signature:
            confirm_superevent_as_gw(superevent_id)

        Arguments:
            superevent_id
        """
        template = self.templates['superevent-confirm-as-gw-template']
        uri = template.format(superevent_id=superevent_id)
        return self.post(uri)

970 971
    @event_or_superevent
    def files(self, object_id, filename="", *args, **kwargs):
Tanner Prestegard's avatar
Tanner Prestegard committed
972 973
        """
        Files for a particular event or superevent.
974

975 976 977
        Signature:
            files(object_id, filename="")

Tanner Prestegard's avatar
Tanner Prestegard committed
978 979 980
        Arguments:
            object_id: event graceid or superevent id
            filename:  name of file (optional)
981

Tanner Prestegard's avatar
Tanner Prestegard committed
982 983 984
        If a filename is not specified, a list of filenames associated with the
        event or superevent is returned.
        If a filename is specified, the file's content is returned.
985

Tanner Prestegard's avatar
Tanner Prestegard committed
986
        Example 1: get a list of files
987 988 989 990 991 992 993

            >>> g = GraceDb()
            >>> event_files = g.files('T101383').json()
            >>> for filename in event_files.keys():
            ...     # do something
            ...     pass

Tanner Prestegard's avatar
Tanner Prestegard committed
994
        Example 2: get a file's content
995 996 997 998 999 1000 1001

            >>> outfile = open('my_skymap.png', 'w')
            >>> r = g.files('T101383', 'skymap.png')
            >>> outfile.write(r.read())
            >>> outfile.close()

        """
1002 1003 1004 1005 1006 1007 1008 1009 1010 1011 1012 1013 1014 1015
        is_superevent = kwargs.pop('is_superevent', False)
        if is_superevent:
            uri_kwargs = {'superevent_id': object_id}
            if filename:
                # Get specific file
                uri_kwargs['file_name'] = filename
                template = self.templates['superevent-file-detail-template']
            else:
                # Get list of files
                template = self.templates['superevent-file-list-template']
        else:
            template = self.templates['files-template']
            uri_kwargs = {'graceid': object_id, 'filename': filename}
        uri = template.format(**uri_kwargs)
1016 1017
        return self.get(uri)

1018
    def writeFile(self, object_id, filename, filecontents=None):
Tanner Prestegard's avatar
Tanner Prestegard committed
1019 1020
        """
        Upload a file for an event or superevent.
1021

1022 1023 1024
        Signature:
            writeFile(object_id, filename, filecontents=None)

Tanner Prestegard's avatar
Tanner Prestegard committed
1025 1026 1027
        Required args:
            object_id: event graceid or superevent id
            filename:  path to file
1028

1029
        This method creates a new log message with your file attached. It is
1030
        strongly preferred to use writeLog() instead of writeFile() so that you
1031
        can add a more suitable comment. That will make it easier for other
1032 1033 1034 1035 1036 1037 1038 1039 1040
        users to know what your file contains.

        Example:

            >>> g = GraceDb()
            >>> r = g.writeFile('T101383', '/path/to/my_interesting_plot.png')
            >>> print r.status

        """
1041 1042 1043
        print("WARNING: the writeFile() method is deprecated in favor "
            "of writeLog() and will be removed in a future release.")
        return self.writeLog(object_id, "FILE UPLOAD", filename, filecontents)
1044

1045 1046
    @event_or_superevent
    def logs(self, object_id, log_number=None, *args, **kwargs):
Tanner Prestegard's avatar
Tanner Prestegard committed
1047 1048
        """
        Get log messages associated with an event or superevent
1049

1050 1051 1052
        Signature:
            logs(object_id, log_number=None)

Tanner Prestegard's avatar
Tanner Prestegard committed
1053 1054 1055 1056
        Arguments:
            object_id:  event graceid or superevent id
            log_number: log message number (N) of log message to retrieve
                        (optional)
1057

Tanner Prestegard's avatar
Tanner Prestegard committed
1058 1059 1060
        If a log_number is specified, only a single log message is returned.
        If a log_number is not specified, a list of all log messages attached
        to the event or superevent in questions is returned.
1061

Tanner Prestegard's avatar
Tanner Prestegard committed
1062
        Example 1: get all log messages
1063 1064 1065 1066 1067 1068 1069 1070

            >>> g = GraceDb()
            >>> response_dict = g.logs('T101383').json()
            >>> print "Num logs = %d" % response_dict['numRows']
            >>> log_list = response_dict['log']
            >>> for log in log_list:
            ...     print log['comment']

Tanner Prestegard's avatar
Tanner Prestegard committed
1071 1072 1073 1074
        Example 2: get a single log message

            >>> g = GraceDb()
            >>> log_info = g.logs('T101383', 10).json()
1075
        """
1076 1077 1078 1079
        if log_number and not isinstance(log_number, int):
            raise TypeError('log_number should be an int')

        # Set up template and object id
1080
        is_superevent = kwargs.pop('is_superevent', False)
1081 1082 1083 1084 1085 1086 1087 1088 1089 1090 1091 1092 1093 1094 1095 1096 1097
        if is_superevent:
            uri_kwargs = {'superevent_id': object_id}
            if log_number:
                template = self.templates['superevent-log-detail-template']
            else:
                template = self.templates['superevent-log-list-template']
        else:
            uri_kwargs = {'graceid': object_id}
            if log_number:
                template = self.templates['event-log-detail-template']
            else:
                template = self.templates['event-log-template']

        if log_number:
            uri_kwargs['N'] = log_number

        uri = template.format(**uri_kwargs)
1098 1099
        return self.get(uri)

1100 1101
    @event_or_superevent
    def writeLog(self, object_id, message, filename=None, filecontents=None,
Tanner Prestegard's avatar
Tanner Prestegard committed
1102
            tag_name=[], displayName=[], *args, **kwargs):
Tanner Prestegard's avatar
Tanner Prestegard committed
1103
        """
1104 1105 1106 1107 1108
        Create a new log message.

        Signature:
        writeLog(object_id, message, filename=None, filecontents=None,
            tag_name=[], displayName=[])
Tanner Prestegard's avatar
Tanner Prestegard committed
1109

Tanner Prestegard's avatar
Tanner Prestegard committed
1110 1111 1112 1113 1114 1115 1116 1117 1118 1119 1120 1121 1122 1123 1124 1125 1126
        Arguments:
            object_id:    event graceid or superevent id
            message:      comment to post on the event log
            filename:     path to file to be uploaded (optional)
            filecontents: handler pointing to a file to be read
                          (optional)
            tag_name:     tag name string or list of tag names to be attached
                          to the log message
            displayName:  tag display name string or list of display names for
                          tags. If provided, there should be one for each tag.
                          (optional)

        If only object_id and message are provided, a text comment will be
        created in the event or superevent log. If a filename is also
        specified, the file will be attached to the log message and displayed
        alongside the message text. If a tag_name is provided, the message will
        be tagged.
1127 1128 1129 1130

        Example:

            >>> g = GraceDb()
Tanner Prestegard's avatar
Tanner Prestegard committed
1131 1132
            >>> r = g.writeLog('T101383', 'Good stuff.', '/path/to/plot.png',
            ... tag_name='analyst_comments')
1133
            >>> print r.status
Tanner Prestegard's avatar
Tanner Prestegard committed
1134
            201
1135
        """
Tanner Prestegard's avatar
Tanner Prestegard committed
1136
        # Check displayName length - should be 0 or same as tag_name
Tanner Prestegard's avatar
Tanner Prestegard committed
1137 1138
        if displayName and isinstance(tag_name, list) and \
            len(displayName) != len(tag_name):
Tanner Prestegard's avatar
Tanner Prestegard committed
1139 1140 1141 1142
            raise ValueError("For a list of tags, either provide no display "
                "names or a display name for each tag")

        is_superevent = kwargs.pop('is_superevent', False)
1143
        if is_superevent:
Tanner Prestegard's avatar
Tanner Prestegard committed
1144
            template = self.templates['superevent-log-list-template']
1145
            uri_kwargs = {'superevent_id': object_id}
1146 1147
        else:
            template = self.templates['event-log-template']
1148 1149
            uri_kwargs = {'graceid': object_id}
        uri = template.format(**uri_kwargs)
1150 1151 1152 1153 1154 1155 1156 1157
        files = None
        if filename:
            if filecontents is None:
                if filename == '-':
                    filename = 'stdin'
                    filecontents = sys.stdin.read()
                else:
                    filecontents = open(filename, "rb").read()
1158
            elif hasattr(filecontents, 'read'):
1159 1160 1161 1162
                # XXX Does not scale well.
                filecontents = filecontents.read()
            files = [('upload', os.path.basename(filename), filecontents)]

Tanner Prestegard's avatar
Tanner Prestegard committed
1163 1164 1165
        # Handle cases where tag_name or displayName are strings
        if isinstance(tag_name, str):
            tag_name = [tag_name]
1166 1167 1168 1169
        elif isinstance(tag_name, (tuple, set)):
            tag_name = list(tag_name)
        elif tag_name is None:
            tag_name = []
Tanner Prestegard's avatar
Tanner Prestegard committed
1170 1171 1172

        if isinstance(displayName, str):
            displayName = [displayName]
1173 1174 1175 1176
        elif isinstance(displayName, (tuple, set)):
            displayName = list(displayName)
        elif displayName is None:
            displayName = []
Tanner Prestegard's avatar
Tanner Prestegard committed
1177

Tanner Prestegard's avatar
Tanner Prestegard committed
1178 1179
        # Set up body of request
        body = {
Tanner Prestegard's avatar
Tanner Prestegard committed
1180
            'comment': message,
Tanner Prestegard's avatar
Tanner Prestegard committed
1181 1182 1183 1184 1185 1186 1187 1188 1189 1190 1191 1192 1193 1194
            'tagname': tag_name,
            'displayName': displayName,
        }

        # If files are attached, we have to encode the request body
        # differently, so we convert from a dict to a list of tuples.
        if files:
            fields = []
            for k,v in body.items():
                if isinstance(v, list):
                    for item in v: fields.append((k, item))
                else:
                    fields.append((k,v))
            body = fields
1195

Tanner Prestegard's avatar
Tanner Prestegard committed
1196
        return self.post(uri, body, files=files)
1197

1198
    def eels(self, graceid):
1199
        """Given a GraceID, get a list of EMBB log entries
1200 1201

        Example:
1202 1203 1204

            >>> g = GraceDb()
            >>> r = g.eels('T101383')
1205 1206 1207
            >>> full_dictionary = r.json()            # Convert the response to a dictionary
            >>> eel_list = full_dictionary['embblog'] # Pull out a list of EEL dicts

1208
        """
1209

1210 1211 1212 1213
        template = self.templates['embb-event-log-template']
        uri = template.format(graceid=graceid)
        return self.get(uri)

1214
    def writeEel(self, graceid, group, waveband, eel_status,
1215
            obs_status, **kwargs):
1216 1217
        """Write an EMBB event log entry

1218 1219
        Required args: graceid, group, waveband, eel_status, obs_status

1220
        (Note that 'group' here is the name of the EM MOU group, not
1221
        the LVC data analysis group responsible for the original detection.)
1222 1223 1224 1225 1226 1227 1228 1229 1230 1231 1232 1233

        Additional keyword arguments may be passed in to be sent in the POST
        data. Only the following kwargs are recognized:
            ra
            dec
            raWidth
            decWidth
            gpstime
            duration
            comment
            extra_info_dict

1234 1235 1236 1237
        Most of these are self-explanatory, but the 'extra_info_dict' is meant
        to be a JSON string containing any additional information the user may
        wish to convey.

1238
        Any other kwargs will be ignored.
1239
        """
1240
        # validate facility, waveband, eel_status, and obs_status
1241 1242
        if not group in self.em_groups:
            raise ValueError("group must be one of %s" % self.em_groups)
1243

1244 1245
        if not waveband in list(self.wavebands.keys()):
            raise ValueError("waveband must be one of %s" % list(self.wavebands.keys()))
1246 1247 1248

        eel_status = self._getCode(eel_status, self.eel_statuses)
        if not eel_status:
1249
            raise ValueError("EEL status must be one of %s" % list(self.eel_statuses.values()))
1250 1251 1252

        obs_status = self._getCode(obs_status, self.obs_statuses)
        if not obs_status:
1253
            raise ValueError("Observation status must be one of %s" % list(self.obs_statuses.values()))
1254 1255 1256 1257 1258

        template = self.templates['embb-event-log-template']
        uri = template.format(graceid=graceid)

        body = {
1259
            'group' : group,
1260 1261 1262 1263 1264 1265 1266
            'waveband' : waveband,
            'eel_status' : eel_status,
            'obs_status' : obs_status,
        }
        body.update(**kwargs)
        return self.post(uri, body=body)

1267 1268 1269
    @event_or_superevent
    def emobservations(self, object_id, emobservation_num=None, *args,
        **kwargs):
Tanner Prestegard's avatar
Tanner Prestegard committed
1270 1271 1272
        """
        Given an event graceid or superevent id, get a list of EM observation
        entries or a specific EM observation.
1273

1274 1275 1276
        Signature:
            emobservations(object_id, emobservation_num=None)

Tanner Prestegard's avatar
Tanner Prestegard committed
1277 1278 1279 1280 1281
        Arguments:
            object_id: event graceid or superevent id
            emobservation_num: number of the EM observation (N) (optional)

        Example 1: get a list of EM observations
1282 1283 1284

            >>> g = GraceDb()
            >>> r = g.emobservations('T101383')
Tanner Prestegard's avatar
Tanner Prestegard committed
1285 1286 1287 1288
            >>> full_dictionary = r.json()
            >>> emo_list = full_dictionary['observations']

        Example 2: get a single EM observation
1289

Tanner Prestegard's avatar
Tanner Prestegard committed
1290 1291 1292
            >>> g = GraceDb()
            >>> r = g.emobservations('T101383', 2)
            >>> observation_dict = r.json()
1293
        """
1294 1295 1296 1297 1298 1299 1300 1301 1302 1303 1304 1305 1306
        is_superevent = kwargs.pop('is_superevent', False)
        if is_superevent:
            uri_kwargs = {'superevent_id': object_id}
            if emobservation_num:
                template = self.templates['superevent-emobservation-detail-template']
            else:
                template = self.templates['superevent-emobservation-list-template']
        else:
            uri_kwargs = {'graceid': object_id}
            if emobservation_num:
                template = self.templates['emobservation-detail-template']
            else:
                template = self.templates['emobservation-list-template']
1307

1308 1309 1310 1311
        if emobservation_num:
            uri_kwargs['N'] = emobservation_num

        uri = template.format(**uri_kwargs)
1312 1313
        return self.get(uri)

Tanner Prestegard's avatar
Tanner Prestegard committed
1314