capsul-flask/capsulflask/db_model.py

from nanoid import generate


class DBModel:
  def __init__(self, connection, cursor):
    self.connection = connection
    self.cursor = cursor

  def login(self, email):
    self.cursor.execute("SELECT * FROM accounts WHERE email = %s", (email, ))
    if len(self.cursor.fetchall()) == 0:
      self.cursor.execute("INSERT INTO accounts (email) VALUES (%s)", (email, ))

    self.cursor.execute("SELECT token FROM login_tokens WHERE email = %s", (email, ))
    if len(self.cursor.fetchall()) > 2:
      return None

    token = generate()
    self.cursor.execute("INSERT INTO login_tokens (email, token) VALUES (%s, %s)", (email, token))
    self.connection.commit()

    return token
    
  def consumeToken(self, token):
    self.cursor.execute("SELECT email FROM login_tokens WHERE token = %s", (token, ))
    rows = self.cursor.fetchall()
    if len(rows) > 0:
      email = rows[0][0]
      self.cursor.execute("DELETE FROM login_tokens WHERE email = %s", (email, ))
      self.connection.commit()
      return email
    return None

  def allVmIds(self,):
    self.cursor.execute("SELECT id FROM vms")
    return map(lambda x: x[0], self.cursor.fetchall())