manage.py 1.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475
  1. import argparse
  2. import sqlite3
  3. import uuid
  4. DB_FILENAME = "db.sqlite"
  5. SQL_CREATE_USERS = \
  6. """CREATE TABLE users (
  7. username TEXT PRIMARY KEY,
  8. api_key TEXT NOT NULL
  9. );"""
  10. SQL_CREATE_REQUESTS = \
  11. """CREATE TABLE requests (
  12. id TEXT PRIMARY KEY,
  13. username TEXT NOT NULL,
  14. depth INTEGER NOT NULL,
  15. width INTEGER NOT NULL,
  16. int_cycles INTEGER NOT NULL,
  17. ext_cycles INTEGER NOT NULL,
  18. status TEXT NOT NULL,
  19. FOREIGN KEY(username) REFERENCES users(username)
  20. );"""
  21. SQL_INSERT_USER = "INSERT INTO users VALUES('{username}', '{api_key}');"
  22. SQL_DELETE_USER = "DELETE FROM users WHERE username='{username}';"
  23. conn = sqlite3.connect(DB_FILENAME)
  24. c = conn.cursor()
  25. def init_db():
  26. c.execute(SQL_CREATE_USERS)
  27. c.execute(SQL_CREATE_REQUESTS)
  28. def add_user(username):
  29. api_key = uuid.uuid4().hex
  30. sql = SQL_INSERT_USER.format(username=username, api_key=api_key)
  31. c.execute(sql)
  32. conn.commit()
  33. print("User '%s' added succesfully (api_key: %s)" % (username, api_key))
  34. def remove_user(username):
  35. sql = SQL_DELETE_USER.format(username=username)
  36. c.execute(sql)
  37. conn.commit()
  38. print("User '%s' removed succesfully from DB" % username)
  39. def parse_args():
  40. parser = argparse.ArgumentParser(description='Script to manage the DB related to the DEVStone comparative web service.')
  41. parser.add_argument('-i', '--init', action='store_true', help='Init the DB')
  42. parser.add_argument('-a', '--add_user', type=str, help='Add a user to the DB')
  43. parser.add_argument('-r', '--remove_user', type=str, help='Remove a user from the DB')
  44. return parser.parse_args()
  45. if __name__ == '__main__':
  46. args = parse_args()
  47. if args.init:
  48. init_db()
  49. if args.add_user:
  50. add_user(args.add_user)
  51. if args.remove_user:
  52. remove_user(args.remove_user)