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 consume_token(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 all_vm_ids(self,): self.cursor.execute("SELECT id FROM vms") return map(lambda x: x[0], self.cursor.fetchall()) def operating_systems_dict(self,): self.cursor.execute("SELECT id, template_image_file_name, description FROM os_images") operatingSystems = dict() for row in self.cursor.fetchall(): operatingSystems[row[0]] = dict(template_image_file_name=row[1], description=row[2]) return operatingSystems def vm_sizes_dict(self,): self.cursor.execute("SELECT id, dollars_per_month, vcpus, memory_mb, bandwidth_gb_per_month FROM vm_sizes") vmSizes = dict() for row in self.cursor.fetchall(): vmSizes[row[0]] = dict(dollars_per_month=row[1], vcpus=row[2], memory_mb=row[3], bandwidth_gb_per_month=row[4]) return vmSizes def list_ssh_public_keys_for_account(self, email): self.cursor.execute("SELECT name, content FROM ssh_public_keys WHERE email = %s", (email, )) return map( lambda x: dict(name=x[0], content=x[1]), self.cursor.fetchall() ) def list_vms_for_account(self, email): self.cursor.execute(""" SELECT vms.id, vms.last_seen_ipv4, vms.last_seen_ipv6, vms.size, os_images.description, vms.created, vms.deleted FROM vms JOIN os_images on os_images.id = vms.os WHERE vms.email = %s""", (email, ) ) return map( lambda x: dict(id=x[0], ipv4=x[1], ipv6=x[2], size=x[3], os=x[4], created=x[5], deleted=x[6]), self.cursor.fetchall() ) def create_vm(self, email, id, size, os): self.cursor.execute(""" INSERT INTO vms (email, id, size, os) VALUES (%s, %s, %s, %s) """, (email, id, size, os) ) self.connection.commit()