Mercurial > repos > greg > data_manager_gtdbtk_database_installer

comparison data_manager/gtdbtk_database_installer.py @ 0:3ab83cb7e2d2 draft

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
Uploaded
author greg
date Tue, 15 Mar 2022 15:32:31 +0000
parents
children 7093598fa300
comparison
equal deleted inserted replaced
-1:000000000000 0:3ab83cb7e2d2
1 #!/usr/bin/env python
2
3 import argparse
4 import json
5 import os
6 import sys
7 import tarfile
8 from urllib.request import Request, urlopen
9 from urllib.parse import urlparse
10
11
12 def url_download(url, work_dir):
13 url_parts = urlparse(url)
14 file_path = os.path.abspath(os.path.join(work_dir, os.path.basename(url_parts.path)))
15 src = None
16 dst = None
17 try:
18 req = Request(url)
19 src = urlopen(req)
20 with open(file_path, 'wb') as dst:
21 while True:
22 chunk = src.read(2**10)
23 if chunk:
24 dst.write(chunk)
25 else:
26 break
27 except Exception as e:
28 sys.exit(str(e))
29 finally:
30 if src:
31 src.close()
32 if tarfile.is_tarfile(file_path):
33 fh = tarfile.open(file_path, 'r:*')
34 else:
35 return file_path
36 fh.extractall(work_dir)
37 os.remove(file_path)
38 return work_dir
39
40
41 def download(database_id, database_name, url, out_file):
42
43 with open(out_file) as fh:
44 params = json.load(fh)
45
46 work_dir = params['output_data'][0]['extra_files_path']
47 os.makedirs(work_dir)
48 file_path = url_download(url, work_dir)
49
50 data_manager_json = {"data_tables": {}}
51 data_manager_entry = {}
52 data_manager_entry['value'] = database_id
53 data_manager_entry['name'] = database_name
54 data_manager_entry['path'] = file_path
55 data_manager_json["data_tables"]["gtdbtk_database"] = data_manager_entry
56
57 with open(out_file, 'w') as fh:
58 json.dump(data_manager_json, fh, sort_keys=True)
59
60
61 parser = argparse.ArgumentParser()
62
63 parser.add_argument('--database_name', dest='database_name', help='GTDB-Tk database display name')
64 parser.add_argument('--database_id', dest='database_id', help='Unique GTDB-Tk database id')
65 parser.add_argument('--url', dest='url', help='URL to download GTDB-Tk databse version')
66 parser.add_argument('--out_file', dest='out_file', help='JSON output file')
67
68 args = parser.parse_args()
69
70 download(args.database_id, args.database_name, args.url, args.out_file)