view TrustStoreGalaxyImport.py @ 31:6e2dba73eebd

Undo multi file ugh.

import sys
import shutil
import gzip
import tempfile
import os
from py_ts import TrustStoreClient, ts_utils
from galaxy.datatypes.checkers import *

def printNice(elem, f, depth):
    try:
        f.write('\t'*depth + elem.name + " (" + str(len(elem.fragments)) + " parts)\n")
    except AttributeError:
        f.write('\t'*depth + elem.name + "\n")
        for child in elem.children:
            printNice(child, f, depth+1)

def ungzip(download, outputFile):
    is_gzipped, is_valid = check_gzip(download)

    if is_gzipped and not is_valid:
        print "File is compressed (gzip) but not valid."
        sys.exit(4)
    elif is_gzipped and is_valid:
        # We need to uncompress the temp_name file, but BAM files must remain compressed in the BGZF format
        CHUNK_SIZE = 2**20 # 1Mb
        fd, uncompressed = tempfile.mkstemp(prefix='data_id_upload_gunzip_', dir=os.path.dirname(outputFile), text=False )
        gzipped_file = gzip.GzipFile(download, 'rb')
        while 1:
            try:
                chunk = gzipped_file.read(CHUNK_SIZE)
            except IOError:
                os.close(fd)
                os.remove(uncompressed)
                print 'Problem decompressing gzipped data', dataset, json_file
                sys.exit(4)
            if not chunk:
                break
            os.write(fd, chunk)
        os.close(fd)
        gzipped_file.close()

        shutil.copy(uncompressed, outputFile)
        try:
            os.remove(uncompressed)
            os.remove(download)
        except OSError:
                pass
    else:
        shutil.copy(download, outputFile)

if __name__ == '__main__':

    kms_url = sys.argv[1]
    ims_url = sys.argv[2]
    username = sys.argv[3]
    password = sys.argv[4]
    client_key = sys.argv[5]
    client_secret = sys.argv[6]
    storename = sys.argv[7]
    path = sys.argv[8]
    filename = ""
    outputFile = ""
    if len(sys.argv) > 10:
        filename = sys.argv[9]
        outputFile = sys.argv[10]
        # outputFileId = sys.argv[11]
        # outputFileType = sys.argv[12]
        # otherFilesDir = sys.argv[13]
    else:
        outputFile = sys.argv[9]

    config = TrustStoreClient.Config(ims_url, kms_url, client_key, client_secret)
    ts = TrustStoreClient.TrustStoreClient(headless=True, config=config)
    try:
        ts.authenticate(username, password)
    except TrustStoreClient.TrustStoreClientAuthenticationException as e:
        print e
        sys.exit(5)
    ts.getPrivateKey('privkey.pem')
    listing = ts.listStores()
    found = False
    for store in listing:
        if store.friendly_name == storename:
            found = True
            root = ts.listDirectory(store)
            location = None
            if path != "/":
                location = ts_utils.ts_utils.dirAtPath(root, path)
                if not location:
                    print "Path not found"
                    sys.exit(3)
            else:
                location = root
            if filename and filename != "":
                outputFileList = [outputFile]
                inputFileList = None
                # if "," in filename: # we have multiple files guys.
                #     inputFileList = filename.split(",")
                #     for inputFile in inputFileList:
                #         outName = "%s_%s_%s_%s_%s" % ('primary', outputFileId, inputFile.replace(".","-"), 'visible', outputFileType)
                #         outputFileList.append(os.path.join(otherFilesDir, outName))
                # else:
                inputFileList = [filename]
                for inFile, outFile in zip(inputFileList, outputFileList):
                    downloadMe = ts_utils.ts_utils.recurseToChildNamed(location, inFile)
                    if downloadMe:
                        download = ts.getFile(store, downloadMe)
                        ungzip(download, outFile)
                    else:
                        print "File %s not found" % inFile
                        sys.exit(4)
            else:
                with open(outputFile, 'w+') as f:
                    try:
                        for child in root.children:
                            printNice(child, f, 0)
                    except AttributeError as e:
                        print e
                        print root
    if not found:
        print "Store not found"
        sys.exit(2)