/
redis_init.py
98 lines (82 loc) · 2.65 KB
/
redis_init.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
import redis
import os
def main():
admin_pw = os.environ.get("REDIS_ADMIN_PW")
if admin_pw in (None, ""):
print("No ADMIN PW found.")
return
conn_kwargs = dict(
host=os.environ.get("REDIS_HOST", "127.0.0.1"),
port=os.environ.get("REDIS_PORT", 6379),
db=os.environ.get("REDIS_DB"),
)
try:
client = redis.Redis(username="admin", password=admin_pw, **conn_kwargs)
users = client.acl_users()
if not ({"admin", "scheduler", "executor"} - set(users)):
print("ACL users have already been set up.")
return
except redis.exceptions.ResponseError:
# no admin found.
client = redis.Redis(**conn_kwargs)
# initialize users.
print("No ACL users found. Initializing now.")
if client.acl_whoami() == "default":
nopass = admin_pw in (None, "")
client.acl_setuser(
"admin",
enabled=True,
nopass=nopass,
passwords=f"+{admin_pw}" if not nopass else None,
commands=["+@all"],
)
client.close()
del client
admin_client = redis.Redis(username="admin", password=admin_pw, **conn_kwargs)
else:
admin_client = client
assert admin_client.acl_whoami() == "admin"
admin_client.acl_setuser("default", enabled=False, commands=["-@all"])
sched_pw = os.environ.get("REDIS_SCHEDULER_PW")
nopass = sched_pw in (None, "")
admin_client.acl_setuser(
"scheduler",
enabled=True,
nopass=nopass,
passwords=f"+{sched_pw}" if not nopass else None,
commands=[
"-@all",
"+set",
"+get",
"+del",
"+@hash",
"+acl|whoami",
"+scan",
"+select",
],
keys=["job-*", "jobinfo-*", "projects", "users*"],
)
outputs_pw = os.environ.get("REDIS_OUTPUTS_PW")
nopass = outputs_pw in (None, "")
admin_client.acl_setuser(
"outputs",
enabled=True,
nopass=nopass,
passwords=f"+{outputs_pw}" if not nopass else None,
commands=["-@all", "+get", "+del", "+@hash", "+acl|whoami",],
keys=["job-*", "jobinfo-*", "users*"],
)
exec_pw = os.environ.get("REDIS_EXECUTOR_PW")
nopass = exec_pw in (None, "")
admin_client.acl_setuser(
"executor",
enabled=True,
nopass=nopass,
passwords=f"+{exec_pw}" if not nopass else None,
commands=["-@all", "+get", "+acl|whoami"],
keys=["job-*"],
)
admin_client.close()
print(f"Successfully created users: {admin_client.acl_users()}")
if __name__ == "__main__":
main()