#!/usr/bin/python
"""Reads an input file and submits a query to the WSA archive.

Example input file:
    
select top 10 from lassource
where
ra>140 and ra<240
    
"""

# ADD YOUR USERNAME/PASSWORD/COMMUNITY BELOW
USERNAME="xxxxxx"
PASSWORD="xxxxxx"
COMMUNITY="xxxxxxx"
#############################################

import os, sys, re, string
import urllib, urllib2, cookielib

# import the optparse library to ease command line options
from optparse import OptionParser

parser = OptionParser()

parser.add_option("--format", dest="format", default='csv',
                   help="Output format (csv/fits)")
parser.add_option("-o", "--output", dest="output", help="Output filename")
parser.add_option("--db", "--database", dest="database",  default='UKIDSSDR3PLUS',
                   help="Database to connect to")

(options, args) = parser.parse_args()

if args==[]:
    parser.error("There is no input query file.")

# Get the input file name from disk
infile=args[0]
assert os.access(infile, os.F_OK), "File not found."

query = open(infile,'r').read()

# Set request to login to the archive
URLDBLOGIN="http://surveys.roe.ac.uk:8080/wsa/DBLogin?user=%s&passwd=%s&community=+&community2=+%s&submit=Login"
q = URLDBLOGIN % (USERNAME, PASSWORD, COMMUNITY)
response=urllib2.urlopen(q)
request=urllib2.Request(q)

# Extract the cookies from the response header and use them for future connections
cj=cookielib.CookieJar()
cj.extract_cookies(response, request)
opener = urllib2.build_opener(urllib2.HTTPCookieProcessor(cj))

# This is the base query.
if options.format=='fits':
    dd={'formaction': 'freeform', 'sqlstmt': query, 'emailAddress': '',
        'database': options.database, 'timeout': 1800,
	'format': 'FITS', 'compress': 'GZIP', 'rows': 30}
elif options.format=='csv':
    dd={'formaction': 'freeform', 'sqlstmt': query, 'emailAddress': '',
        'database': options.database, 'timeout': 1800,
	'format': 'CSV', 'compress': 'NONE', 'rows': 30}

basequery="http://surveys.roe.ac.uk:8080/wsa/WSASQL?"+urllib.urlencode(dd)
res=opener.open(basequery).read()

# Look for the csv/fits link
try:
    csvfile=re.compile("<a href=\"(\S+%s).+" % options.format).search(res).group(1)
except:
    print "ERROR: Query Failed."
    sys.exit()
    
# Request the csvfile
csvres=opener.open(csvfile).read()
# Save file to local disk
if options.output:
    localcsv = options.output
else:
    localcsv = "./"+csvfile.split('/')[-1]
    
f=open(localcsv, 'w')
f.write(csvres)
f.close()

print "%s file saved" % localcsv