Coverage for yaptide/admin/db_manage.py: 0%
154 statements
« prev ^ index » next coverage.py v7.4.4, created at 2024-07-01 12:55 +0000
1#! /usr/bin/env python
2"""
3This script is used to manage the database via a command line application.
4It is supposed to be used outside of the docker container, therefore by purpose it has
5no dependency on the yaptide package. All the tables used by ORM are defined here again.
6We do not use ORM model classes here, because we do not want to import the yaptide package.
7"""
8from enum import Enum, auto
9import os
11import click
12import sqlalchemy as db
13from werkzeug.security import generate_password_hash
16class TableTypes(Enum):
17 """Enum for table names used in ORM"""
19 User = auto()
20 YaptideUser = auto()
21 KeycloakUser = auto()
22 Simulation = auto()
23 Cluster = auto()
24 Task = auto()
25 Result = auto()
26 Input = auto()
27 Estimator = auto()
28 Page = auto()
31def connect_to_db(verbose: int = 0) -> tuple[db.Connection, db.MetaData, db.Engine]:
32 """Connects to the db"""
33 db_uri = os.environ.get('FLASK_SQLALCHEMY_DATABASE_URI')
34 if verbose > 1:
35 click.echo(f'Connecting to URI: {db_uri}')
36 if not db_uri:
37 click.echo(f'Database URI: {db_uri} not set - aborting', err=True)
38 raise click.Abort()
39 echo: bool = verbose > 1
40 engine = db.create_engine(db_uri, echo=echo)
41 try:
42 con = engine.connect()
43 metadata = db.MetaData()
44 metadata.reflect(bind=engine)
45 except db.exc.OperationalError:
46 click.echo(f'Connection to db {db_uri} failed', err=True)
47 raise click.Abort()
48 return con, metadata, engine
51def user_exists(name: str, auth_provider: str, users: db.Table, con) -> bool:
52 """Check if user already exists"""
53 stmt = db.select(users).filter_by(username=name, auth_provider=auth_provider)
54 users_found = con.execute(stmt).all()
55 if len(users_found) > 0:
56 click.echo(f'User: {name} exists')
57 return True
58 return False
61@click.group()
62def run():
63 """Manage database"""
66@run.command
67@click.option('-v', '--verbose', count=True)
68def list_users(verbose):
69 """List users"""
70 con, metadata, _ = connect_to_db(verbose=verbose)
71 users = metadata.tables[TableTypes.User.name]
72 stmt = db.select(users.c.username, users.c.auth_provider)
73 all_users = con.execute(stmt).all()
75 click.echo(f"{len(all_users)} users in DB:")
76 for user in all_users:
77 click.echo(f"Login {user.username}; Auth provider {user.auth_provider}")
80@run.command
81@click.argument('name')
82@click.option('password', '--password', default='')
83@click.option('-v', '--verbose', count=True)
84def add_user(name, password, verbose):
85 """Add yaptide user to database"""
86 con, metadata, _ = connect_to_db(verbose=verbose)
87 click.echo(f'Adding user: {name}')
88 if verbose > 2:
89 click.echo(f'Password: {password}')
91 users = metadata.tables[TableTypes.User.name]
93 if user_exists(name=name, auth_provider="YaptideUser", users=users, con=con):
94 click.echo(f'YaptideUser: {name} already exists, aborting add')
95 raise click.Abort()
97 click.echo(f'YaptideUser: {name} does not exist, adding')
98 click.echo(f'Adding user: {name}')
99 query = db.insert(users).values(username=name, auth_provider="YaptideUser")
100 result = con.execute(query)
101 user_id = result.inserted_primary_key[0]
102 password_hash = generate_password_hash(password)
104 yaptide_users = metadata.tables[TableTypes.YaptideUser.name]
105 query = db.insert(yaptide_users).values(id=user_id, password_hash=password_hash)
106 con.execute(query)
107 con.commit()
110@run.command
111@click.argument('name')
112@click.option('password', '--password', default='')
113@click.option('-v', '--verbose', count=True)
114def update_user(name, password, verbose):
115 """Update user in database"""
116 con, metadata, _ = connect_to_db(verbose=verbose)
117 click.echo(f'Updating user: {name}')
118 users = metadata.tables[TableTypes.User.name]
120 if not user_exists(name, "YaptideUser", users, con):
121 click.echo(f'YaptideUser: {name} does not exist, aborting update')
122 raise click.Abort()
124 # update password
125 if password:
126 pwd_hash: str = generate_password_hash(password)
127 yaptide_users = metadata.tables[TableTypes.YaptideUser.name]
128 stmt = db.update(yaptide_users).\
129 where(yaptide_users.c.id == users.c.id).\
130 where(users.c.username == name).\
131 where(users.c.auth_provider == "YaptideUser").\
132 values(password_hash=pwd_hash)
133 con.execute(stmt)
134 con.commit()
135 if verbose > 2:
136 click.echo(f'Updating password: {password}')
137 click.echo(f'Successfully updated user: {name}')
140@run.command
141@click.argument('name')
142@click.argument('auth_provider')
143def remove_user(name, auth_provider):
144 """Delete user"""
145 con, metadata, _ = connect_to_db()
146 click.echo(f'Deleting user: {name}')
147 users = metadata.tables[TableTypes.User.name]
149 # abort if user does not exist
150 if not user_exists(name, auth_provider, users, con):
151 click.echo(f"Aborting, user {name} does not exist")
152 raise click.Abort()
154 query = db.delete(users).where(users.c.username == name, users.c.auth_provider == auth_provider)
155 con.execute(query)
156 con.commit()
157 click.echo(f'Successfully deleted user: {name}')
160@run.command
161def list_tasks():
162 """List tasks"""
163 con, metadata, _ = connect_to_db()
164 tasks = metadata.tables[TableTypes.Task.name]
165 stmt = db.select(tasks.c.simulation_id, tasks.c.task_id)
166 all_tasks = con.execute(stmt).all()
168 click.echo(f"{len(all_tasks)} tasks in DB:")
169 for task in all_tasks:
170 click.echo(f"Simulation id {task.simulation_id}; Task id ...{task.task_id}")
173@run.command
174@click.option('-v', '--verbose', count=True)
175def list_simulations(verbose):
176 """List simulations"""
177 con, metadata, _ = connect_to_db(verbose=verbose)
178 simulations = metadata.tables[TableTypes.Simulation.name]
179 stmt = db.select(simulations.c.id, simulations.c.job_id, simulations.c.start_time, simulations.c.end_time)
180 sims = con.execute(stmt).all()
182 click.echo(f"{len(sims)} simulations in DB:")
183 for sim in sims:
184 click.echo(f"id {sim.id}; job id {sim.job_id}; start_time {sim.start_time}; end_time {sim.end_time};")
187@run.command
188@click.argument('cluster_name')
189@click.option('-v', '--verbose', count=True)
190def add_cluster(cluster_name, verbose):
191 """Adds cluster with provided name to database if it does not exist"""
192 con, metadata, _ = connect_to_db(verbose=verbose)
193 cluster_table = metadata.tables[TableTypes.Cluster.name]
195 # check if cluster already exists
196 stmt = db.select(cluster_table).filter_by(cluster_name=cluster_name)
197 clusters = con.execute(stmt).all()
198 if len(clusters) > 0:
199 click.echo(f"Cluster {cluster_name} already exists in DB")
200 raise click.Abort()
202 stmt = db.insert(cluster_table).values(cluster_name=cluster_name)
203 con.execute(stmt)
204 con.commit()
205 click.echo(f"Cluster {cluster_name} added to DB")
208@run.command
209def list_clusters():
210 """List clusters"""
211 con, metadata, _ = connect_to_db()
212 stmt = db.select(metadata.tables[TableTypes.Cluster.name])
213 clusters = con.execute(stmt).all()
215 click.echo(f"{len(clusters)} clusters in DB:")
216 for cluster in clusters:
217 click.echo(f"id {cluster.id}; cluster name {cluster.cluster_name};")
220if __name__ == "__main__":
221 run()