-
Notifications
You must be signed in to change notification settings - Fork 62
/
entry_point.py
115 lines (98 loc) · 3.9 KB
/
entry_point.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
import logging
import os
import click
from io import StringIO
from github3 import GitHub
from github3.session import GitHubSession
from emmet.cli.admin import admin
from emmet.cli.calc import calc
from emmet.cli.tasks import tasks
from emmet.cli.utils import EmmetCliError, calcdb_from_mgrant
logger = logging.getLogger("")
CONTEXT_SETTINGS = dict(help_option_names=["-h", "--help"])
def opt_prompt():
return click.prompt("GitHub OPT", hide_input=True)
@click.group(context_settings=CONTEXT_SETTINGS)
@click.option(
"--spec-or-dbfile",
metavar="HOST/DB",
help="MongoGrant spec or path to db.json for DB to use.",
)
@click.option("--run", is_flag=True, help="Run DB/filesystem write operations.")
@click.option("--issue", type=int, help="Production tracker issue (required if --run).")
@click.option("--sbatch", is_flag=True, help="Switch to SBatch mode.")
@click.option(
"--ntries",
default=1,
show_default=True,
help="Number of jobs (for walltime > 48h).",
)
@click.option("--bb", is_flag=True, help="Use burst buffer.")
@click.option("--yes", is_flag=True, help="Automatic yes to all prompts.")
@click.option("--no-dupe-check", is_flag=True, help="Skip duplicate check(s).")
@click.option("--verbose", is_flag=True, help="Show debug messages.")
@click.version_option()
def emmet(spec_or_dbfile, run, issue, sbatch, ntries, bb, yes, no_dupe_check, verbose):
"""Command line interface for emmet"""
logger.setLevel(logging.DEBUG if verbose else logging.INFO)
ctx = click.get_current_context()
ctx.ensure_object(dict)
if not sbatch and bb:
raise EmmetCliError("Burst buffer only available in SBatch mode (--sbatch).")
if spec_or_dbfile:
client = calcdb_from_mgrant(spec_or_dbfile)
ctx.obj["CLIENT"] = client
# ctx.obj["MONGO_HANDLER"] = BufferedMongoHandler(
# host=client.host,
# port=client.port,
# database_name=client.db_name,
# username=client.user,
# password=client.password,
# level=logging.WARNING,
# authentication_db=client.db_name,
# collection="emmet_logs",
# buffer_periodical_flush_timing=False, # flush manually
# )
# logger.addHandler(ctx.obj["MONGO_HANDLER"])
# coll = ctx.obj["MONGO_HANDLER"].collection
# ensure_indexes(SETTINGS.log_fields, [coll])
if run:
if not issue:
raise EmmetCliError("Need issue number via --issue!")
ctx.obj["LOG_STREAM"] = StringIO()
memory_handler = logging.StreamHandler(ctx.obj["LOG_STREAM"])
formatter = logging.Formatter(
"%(asctime)s %(name)-12s %(levelname)-8s %(message)s"
)
memory_handler.setFormatter(formatter)
logger.addHandler(memory_handler)
CREDENTIALS = os.path.join(os.path.expanduser("~"), ".emmet_credentials")
if not os.path.exists(CREDENTIALS):
from github3 import authorize # TODO not supported anymore
user = click.prompt("GitHub Username")
password = click.prompt("GitHub Password", hide_input=True)
auth = authorize(
user,
password,
["user", "repo", "gist"],
"emmet CLI",
two_factor_callback=opt_prompt,
)
with open(CREDENTIALS, "w") as fd:
fd.write(auth.token)
with open(CREDENTIALS, "r") as fd:
token = fd.readline().strip()
ctx.obj["GH"] = gh = GitHub(session=GitHubSession(default_read_timeout=30))
gh.login(token=token)
else:
click.secho("DRY RUN! Add --run flag to execute changes.", fg="green")
def safe_entry_point():
try:
emmet()
except EmmetCliError as e:
click.secho(str(e), fg="red")
except Exception as e:
logger.info(e, exc_info=True)
emmet.add_command(admin)
emmet.add_command(calc)
emmet.add_command(tasks)