diff --git a/init_user_state.py b/init_user_state.py new file mode 100644 index 0000000000000000000000000000000000000000..c2b7628ecbeaa0db389201018147d8d8e3e8412b --- /dev/null +++ b/init_user_state.py @@ -0,0 +1,58 @@ +#!/usr/bin/env python3 +import argparse +import dataset +import sys +import subprocess +import rabbit_config as rcfg +from datetime import datetime + +parser = argparse.ArgumentParser() +parser.add_argument("-f", "--force", action="store_true", help="force update") +parser.add_argument( + "--dry-run", action="store_true", help="enable dry run mode" +) +args = parser.parse_args() + +default_state = "ok" +today = datetime.now() + +# Chunk size for insert into db +size = 1000 + +db = dataset.connect(f"sqlite:///{rcfg.db_path}/user_reg.db") +table = db["user_state"] + +if table.__len__() > 0 and not args.force: + print("table user_state not empty, abort.") + sys.exit() + +# Getting user list +users = subprocess.run( + ["ls", "/data/user"], stdout=subprocess.PIPE, encoding="UTF-8" +).stdout.split() + +# Update user_state table +# Insert many +if len(users) > 50: + start = 0 + while start < len(users): + end = start + size if start + size < len(users) else len(users) + data = [ + dict(username=user, state=default_state, date=today) + for user in users[start:end] + ] + if args.dry_run: + print(f"Table insert many from {start} to {end - 1}") + else: + table.insert_many(data, chunk_size=size) + start = end + +# Insert one by one +else: + for user in users: + if args.dry_run: + print(f"Table insert user: {user}, state: {default_state}") + else: + table.insert( + {"username": user, "state": default_state, "date": today} + )