change db to postgres, fix rtp functions, change how db is connected

.gitignore

@@ -1,3 +1,4 @@
 config.local.py
 __pycache__/*
+*/__pycache__/*
 rrd/*

config.py

@@ -2,7 +2,10 @@ from os import environ
 
 # Proxstar
 VM_EXPIRE_MONTHS = int(environ.get('PROXSTAR_VM_EXPIRE_MONTHS', '3'))
-IGNORED_POOLS = [pool.strip() for pool in environ.get('PROXSTAR_IGNORED_POOLS', '').split(',')]
+IGNORED_POOLS = [
+    pool.strip()
+    for pool in environ.get('PROXSTAR_IGNORED_POOLS', '').split(',')
+]
 
 # Flask
 IP = environ.get('PROXSTAR_IP', '0.0.0.0')
@@ -11,11 +14,17 @@ SERVER_NAME = environ.get('PROXSTAR_SERVER_NAME', 'proxstar.csh.rit.edu')
 SECRET_KEY = environ.get('PROXSTAR_SECRET_KEY', '')
 
 # OIDC
-OIDC_ISSUER = environ.get('PROXSTAR_OIDC_ISSUER', 'https://sso.csh.rit.edu/auth/realms/csh')
+OIDC_ISSUER = environ.get('PROXSTAR_OIDC_ISSUER',
+                          'https://sso.csh.rit.edu/auth/realms/csh')
 OIDC_CLIENT_CONFIG = {
-    'client_id': environ.get('PROXSTAR_CLIENT_ID', 'proxstar'),
-    'client_secret': environ.get('PROXSTAR_CLIENT_SECRET', ''),
-    'post_logout_redirect_uris': [environ.get('PROXSTAR_REDIRECT_URI', 'https://proxstar.csh.rit.edu/logout')]
+    'client_id':
+    environ.get('PROXSTAR_CLIENT_ID', 'proxstar'),
+    'client_secret':
+    environ.get('PROXSTAR_CLIENT_SECRET', ''),
+    'post_logout_redirect_uris': [
+        environ.get('PROXSTAR_REDIRECT_URI',
+                    'https://proxstar.csh.rit.edu/logout')
+    ]
 }
 
 # Proxmox

proxstar/__init__.py

@@ -3,6 +3,8 @@ import time
 import subprocess
 from rq import Queue
 from redis import Redis
+from sqlalchemy import create_engine
+from sqlalchemy.orm import sessionmaker
 from werkzeug.contrib.cache import SimpleCache
 from flask_pyoidc.flask_pyoidc import OIDCAuthentication
 from flask import Flask, render_template, request, redirect, send_from_directory, session
@@ -16,8 +18,11 @@ redis_conn = Redis()
 q = Queue(connection=redis_conn)
 
 app = Flask(__name__)
-if os.path.exists(os.path.join(app.config.get('ROOT_DIR', os.getcwd()), "config.local.py")):
-    config = os.path.join(app.config.get('ROOT_DIR', os.getcwd()), "config.local.py")
+if os.path.exists(
+        os.path.join(
+            app.config.get('ROOT_DIR', os.getcwd()), "config.local.py")):
+    config = os.path.join(
+        app.config.get('ROOT_DIR', os.getcwd()), "config.local.py")
 else:
     config = os.path.join(app.config.get('ROOT_DIR', os.getcwd()), "config.py")
 app.config.from_pyfile(config)
@@ -29,6 +34,11 @@ auth = OIDCAuthentication(
     client_registration_info=app.config['OIDC_CLIENT_CONFIG'])
 cache = SimpleCache()
 
+engine = create_engine(app.config['SQLALCHEMY_DATABASE_URI'])
+Base.metadata.bind = engine
+DBSession = sessionmaker(bind=engine)
+db = DBSession()
+
 starrs = psycopg2.connect(
     "dbname='{}' user='{}' host='{}' password='{}'".format(
         app.config['STARRS_DB_NAME'], app.config['STARRS_DB_USER'],
@@ -39,10 +49,6 @@ starrs = psycopg2.connect(
 @app.route("/user/<string:user>")
 @auth.oidc_auth
 def list_vms(user=None):
-    print(q.jobs)
-    for job_id in q.job_ids:
-        print(job_id)
-        print(q.fetch_job(job_id).result)
     rtp_view = False
     rtp = 'rtp' in session['userinfo']['groups']
     active = 'active' in session['userinfo']['groups']
@@ -57,7 +63,7 @@ def list_vms(user=None):
         user = session['userinfo']['preferred_username']
         vms = cache.get('vms')
         if vms is None:
-            vms = get_vms_for_rtp(proxmox)
+            vms = get_vms_for_rtp(proxmox, db)
             cache.set('vms', vms, timeout=5 * 60)
         rtp_view = True
     else:
@@ -87,7 +93,6 @@ def isos():
 @app.route("/hostname/<string:name>")
 @auth.oidc_auth
 def hostname(name):
-    starrs = connect_starrs()
     valid, available = check_hostname(starrs, name)
     if not valid:
         return 'invalid'
@@ -104,7 +109,6 @@ def vm_details(vmid):
     rtp = 'rtp' in session['userinfo']['groups']
     active = 'active' in session['userinfo']['groups']
     proxmox = connect_proxmox()
-    starrs = connect_starrs()
     if 'rtp' in session['userinfo']['groups'] or int(
             vmid) in get_user_allowed_vms(proxmox, user):
         vm = get_vm(proxmox, vmid)
@@ -119,10 +123,11 @@ def vm_details(vmid):
                 [interface[0],
                  get_ip_for_mac(starrs, interface[1])])
         vm['expire'] = get_vm_expire(
-            vmid, app.config['VM_EXPIRE_MONTHS']).strftime('%m/%d/%Y')
+            db, vmid, app.config['VM_EXPIRE_MONTHS']).strftime('%m/%d/%Y')
         usage = get_user_usage(proxmox, user)
-        limits = get_user_usage_limits(user)
-        usage_check = check_user_usage(proxmox, user, vm['config']['cores'],
+        limits = get_user_usage_limits(db, user)
+        usage_check = check_user_usage(proxmox, db, user,
+                                       vm['config']['cores'],
                                        vm['config']['memory'], 0)
         return render_template(
             'vm_details.html',
@@ -207,7 +212,6 @@ def vm_mem(vmid, mem):
 def vm_renew(vmid):
     user = session['userinfo']['preferred_username']
     proxmox = connect_proxmox()
-    starrs = connect_starrs()
     if int(vmid) in get_user_allowed_vms(
             proxmox, user) or 'rtp' in session['userinfo']['groups']:
         renew_vm_expire(vmid, app.config['VM_EXPIRE_MONTHS'])
@@ -249,9 +253,9 @@ def iso_mount(vmid, iso):
 @auth.oidc_auth
 def delete(vmid):
     user = session['userinfo']['preferred_username']
+    rtp = 'rtp' in session['userinfo']['groups']
     proxmox = connect_proxmox()
-    starrs = connect_starrs()
-    if int(vmid) in get_user_allowed_vms(
+    if rtp or int(vmid) in get_user_allowed_vms(
             proxmox, user) or 'rtp' in session['userinfo']['groups']:
         q.enqueue(delete_vm_task, vmid)
         return '', 200
@@ -266,11 +270,10 @@ def create():
     rtp = 'rtp' in session['userinfo']['groups']
     active = 'active' in session['userinfo']['groups']
     proxmox = connect_proxmox()
-    starrs = connect_starrs()
     if active:
         if request.method == 'GET':
             usage = get_user_usage(proxmox, user)
-            limits = get_user_usage_limits(user)
+            limits = get_user_usage_limits(db, user)
             percents = get_user_usage_percent(proxmox, user, usage, limits)
             isos = get_isos(proxmox, app.config['PROXMOX_ISO_STORAGE'])
             pools = get_pools(proxmox)

proxstar/db.py

@@ -1,73 +1,64 @@
 import datetime
-from sqlalchemy import create_engine, exists
-from sqlalchemy.orm import sessionmaker
+from sqlalchemy import exists
 from dateutil.relativedelta import relativedelta
 from proxstar.ldapdb import *
-from proxstar.db_init import VM_Expiration, Usage_Limit, Base
-
-engine = create_engine('sqlite:///proxstar.db')
-Base.metadata.bind = engine
-
-DBSession = sessionmaker(bind=engine)
-session = DBSession()
+from proxstar.models import VM_Expiration, Usage_Limit, Base
 
 
-def get_vm_expire(vmid, months):
-    if session.query(exists().where(VM_Expiration.id == vmid)).scalar():
-        expire = session.query(VM_Expiration).filter(
+def get_vm_expire(db, vmid, months):
+    if db.query(exists().where(VM_Expiration.id == vmid)).scalar():
+        expire = db.query(VM_Expiration).filter(
             VM_Expiration.id == vmid).one().expire_date
     else:
         expire = datetime.date.today() + relativedelta(months=months)
         new_expire = VM_Expiration(id=vmid, expire_date=expire)
-        session.add(new_expire)
-        session.commit()
+        db.add(new_expire)
+        db.commit()
     return expire
 
 
-def renew_vm_expire(vmid, months):
-    if session.query(exists().where(VM_Expiration.id == vmid)).scalar():
-        expire = session.query(VM_Expiration).filter(
-            VM_Expiration.id == vmid).one()
+def renew_vm_expire(db, vmid, months):
+    if db.query(exists().where(VM_Expiration.id == vmid)).scalar():
+        expire = db.query(VM_Expiration).filter(VM_Expiration.id == vmid).one()
         new_expire = datetime.date.today() + relativedelta(months=months)
         expire.expire_date = new_expire
-        session.commit()
+        db.commit()
     else:
         expire = datetime.date.today() + relativedelta(months=months)
         new_expire = VM_Expiration(id=vmid, expire_date=expire)
-        session.add(new_expire)
-        session.commit()
+        db.add(new_expire)
+        db.commit()
 
 
-def delete_vm_expire(vmid):
-    if session.query(exists().where(VM_Expiration.id == vmid)).scalar():
-        expire = session.query(VM_Expiration).filter(
-            VM_Expiration.id == vmid).one()
-        session.delete(expire)
-        session.commit()
+def delete_vm_expire(db, vmid):
+    if db.query(exists().where(VM_Expiration.id == vmid)).scalar():
+        expire = db.query(VM_Expiration).filter(VM_Expiration.id == vmid).one()
+        db.delete(expire)
+        db.commit()
 
 
-def get_expired_vms():
+def get_expired_vms(db):
     expired = []
     today = datetime.date.today().strftime('%Y-%m-%d')
-    expire = session.query(VM_Expiration).filter(
+    expire = db.query(VM_Expiration).filter(
         VM_Expiration.expire_date < today).all()
     for vm in expire:
         expired.append(vm.id)
     return expired
 
 
-def get_user_usage_limits(user):
+def get_user_usage_limits(db, user):
     limits = dict()
     if is_rtp(user):
         limits['cpu'] = 1000
         limits['mem'] = 1000
         limits['disk'] = 100000
-    elif session.query(exists().where(Usage_Limit.id == user)).scalar():
-        limits['cpu'] = session.query(Usage_Limit).filter(
+    elif db.query(exists().where(Usage_Limit.id == user)).scalar():
+        limits['cpu'] = db.query(Usage_Limit).filter(
             Usage_Limit.id == user).one().cpu
-        limits['mem'] = session.query(Usage_Limit).filter(
+        limits['mem'] = db.query(Usage_Limit).filter(
             Usage_Limit.id == user).one().mem
-        limits['disk'] = session.query(Usage_Limit).filter(
+        limits['disk'] = db.query(Usage_Limit).filter(
             Usage_Limit.id == user).one().disk
     else:
         limits['cpu'] = 4
@@ -76,23 +67,21 @@ def get_user_usage_limits(user):
     return limits
 
 
-def set_user_usage_limits(user, cpu, mem, disk):
-    if session.query(exists().where(Usage_Limit.id == user)).scalar():
-        limits = session.query(Usage_Limit).filter(
-            Usage_Limit.id == user).one()
+def set_user_usage_limits(db, user, cpu, mem, disk):
+    if db.query(exists().where(Usage_Limit.id == user)).scalar():
+        limits = db.query(Usage_Limit).filter(Usage_Limit.id == user).one()
         limits.cpu = cpu
         limits.mem = mem
         limits.disk = disk
-        session.commit()
+        db.commit()
     else:
         limits = Usage_Limit(id=user, cpu=cpu, mem=mem, disk=disk)
-        session.add(limits)
-        session.commit()
+        db.add(limits)
+        db.commit()
 
 
-def delete_user_usage_limits(user):
-    if session.query(exists().where(Usage_Limit.id == user)).scalar():
-        limits = session.query(Usage_Limit).filter(
-            Usage_Limit.id == user).one()
-        session.delete(limits)
-        session.commit()
+def delete_user_usage_limits(db, user):
+    if db.query(exists().where(Usage_Limit.id == user)).scalar():
+        limits = db.query(Usage_Limit).filter(Usage_Limit.id == user).one()
+        db.delete(limits)
+        db.commit()

proxstar/ldapdb.py

@@ -17,7 +17,7 @@ def is_rtp(user):
     return rtp_group.check_member(ldap.get_member(user, uid=True))
 
 
-def is_active(ldap, user):
+def is_active(user):
     ldap = connect_ldap()
     rtp_group = ldap.get_group('active')
     return rtp_group.check_member(ldap.get_member(user, uid=True))

proxstar/models.py

@@ -1,7 +1,5 @@
-from sqlalchemy import Column, ForeignKey, Integer, String, Date
+from sqlalchemy import Column, Integer, String, Date
 from sqlalchemy.ext.declarative import declarative_base
-from sqlalchemy.orm import relationship
-from sqlalchemy import create_engine
 
 Base = declarative_base()
 
@@ -18,8 +16,3 @@ class Usage_Limit(Base):
     cpu = Column(Integer, nullable=False)
     mem = Column(Integer, nullable=False)
     disk = Column(Integer, nullable=False)
-
-
-engine = create_engine('sqlite:///proxstar.db')
-
-Base.metadata.create_all(engine)

proxstar/proxmox.py

@@ -1,7 +1,5 @@
 import time
-from functools import lru_cache
 from proxmoxer import ProxmoxAPI
-from flask import current_app as app
 from proxstar.db import *
 from proxstar.ldapdb import *
 
@@ -41,7 +39,7 @@ def get_vms_for_user(proxmox, user):
     return vms
 
 
-def get_vms_for_rtp(proxmox):
+def get_vms_for_rtp(proxmox, db):
     pools = get_pools(proxmox)
     pool_vms = []
     for pool in pools:
@@ -50,7 +48,7 @@ def get_vms_for_rtp(proxmox):
         pool_dict['vms'] = get_vms_for_user(proxmox, pool)
         pool_dict['num_vms'] = len(pool_dict['vms'])
         pool_dict['usage'] = get_user_usage(proxmox, pool)
-        pool_dict['limits'] = get_user_usage_limits(pool)
+        pool_dict['limits'] = get_user_usage_limits(db, pool)
         pool_dict['percents'] = get_user_usage_percent(
             proxmox, pool, pool_dict['usage'], pool_dict['limits'])
         pool_vms.append(pool_dict)
@@ -177,8 +175,8 @@ def get_user_usage(proxmox, user):
     return usage
 
 
-def check_user_usage(proxmox, user, vm_cpu, vm_mem, vm_disk):
-    limits = get_user_usage_limits(user)
+def check_user_usage(proxmox, db, user, vm_cpu, vm_mem, vm_disk):
+    limits = get_user_usage_limits(db, user)
     cur_usage = get_user_usage(proxmox, user)
     if int(cur_usage['cpu']) + int(vm_cpu) > int(limits['cpu']):
         return 'exceeds_cpu_limit'
@@ -203,7 +201,7 @@ def get_user_usage_percent(proxmox, user, usage=None, limits=None):
     return percents
 
 
-def create_vm(proxmox, starrs, user, name, cores, memory, disk, iso):
+def create_vm(proxmox, user, name, cores, memory, disk, iso):
     node = proxmox.nodes(get_node_least_mem(proxmox))
     vmid = get_free_vmid(proxmox)
     node.qemu.create(
@@ -228,7 +226,7 @@ def create_vm(proxmox, starrs, user, name, cores, memory, disk, iso):
     return vmid, mac
 
 
-def delete_vm(proxmox, starrs, vmid):
+def delete_vm(proxmox, vmid):
     node = proxmox.nodes(get_vm_node(proxmox, vmid))
     node.qemu(vmid).delete()
 

proxstar/starrs.py

@@ -2,18 +2,6 @@ import psycopg2
 from flask import current_app as app
 
 
-def connect_starrs():
-    try:
-        starrs = psycopg2.connect(
-            "dbname='{}' user='{}' host='{}' password='{}'".format(
-                app.config['STARRS_DB_NAME'], app.config['STARRS_DB_USER'],
-                app.config['STARRS_DB_HOST'], app.config['STARRS_DB_PASS']))
-    except:
-        print("Unable to connect to STARRS database.")
-        raise
-    return starrs
-
-
 def get_next_ip(starrs, range_name):
     c = starrs.cursor()
     try:

proxstar/static/js/script.js

@@ -393,9 +393,6 @@ $("#create-vm").click(function(){
     const disk = document.getElementById('disk').value
     const iso = document.getElementById('iso').value
     const user = document.getElementById('user')
-    if (user) {
-        const user_value = user.value
-    }
     const max_disk = $(this).data('max_disk')
     if (name && disk) {
         if (disk > max_disk) {
@@ -433,7 +430,7 @@ $("#create-vm").click(function(){
                             data.append('disk', disk);
                             data.append('iso', iso);
                             if (user) {
-                                data.append('user', user_value);
+                                data.append('user', user.value);
                             }
                             fetch('/vm/create', {
                                 credentials: 'same-origin',
@@ -640,9 +637,6 @@ $(".edit-limit").click(function(){
     })
     .then((willChange) => {
         if (willChange) {
-            console.log($(cpu).val());
-            console.log($(mem).val());
-            console.log($(disk).val());
             var data  = new FormData();
             data.append('cpu', $(cpu).val());
             data.append('mem', $(mem).val());

proxstar/tasks.py

@@ -1,34 +1,59 @@
 import os
+from flask import Flask
+from sqlalchemy import create_engine
+from sqlalchemy.orm import sessionmaker
 from proxstar.db import *
 from proxstar.starrs import *
 from proxstar.proxmox import *
-from flask import Flask, current_app
 
 app = Flask(__name__)
-config = os.path.join(app.config.get('ROOT_DIR', os.getcwd()), "config.py")
+if os.path.exists(
+        os.path.join(
+            app.config.get('ROOT_DIR', os.getcwd()), "config.local.py")):
+    config = os.path.join(
+        app.config.get('ROOT_DIR', os.getcwd()), "config.local.py")
+else:
+    config = os.path.join(app.config.get('ROOT_DIR', os.getcwd()), "config.py")
 app.config.from_pyfile(config)
 
 
+def connect_db():
+    engine = create_engine(app.config['SQLALCHEMY_DATABASE_URI'])
+    Base.metadata.bind = engine
+    DBSession = sessionmaker(bind=engine)
+    db = DBSession()
+    return db
+
+
+def connect_starrs():
+    starrs = psycopg2.connect(
+        "dbname='{}' user='{}' host='{}' password='{}'".format(
+            app.config['STARRS_DB_NAME'], app.config['STARRS_DB_USER'],
+            app.config['STARRS_DB_HOST'], app.config['STARRS_DB_PASS']))
+    return starrs
+
+
 def create_vm_task(user, name, cores, memory, disk, iso):
     with app.app_context():
         proxmox = connect_proxmox()
+        db = connect_db()
         starrs = connect_starrs()
-        vmid, mac = create_vm(proxmox, starrs, user, name, cores, memory, disk,
-                              iso)
+        vmid, mac = create_vm(proxmox, user, name, cores, memory, disk, iso)
         register_starrs(starrs, name, app.config['STARRS_USER'], mac,
                         get_next_ip(starrs,
                                     app.config['STARRS_IP_RANGE'])[0][0])
-        get_vm_expire(vmid, app.config['VM_EXPIRE_MONTHS'])
+        get_vm_expire(db, vmid, app.config['VM_EXPIRE_MONTHS'])
 
 
 def delete_vm_task(vmid):
     with app.app_context():
         proxmox = connect_proxmox()
+        db = connect_db()
         starrs = connect_starrs()
         vmname = get_vm_config(proxmox, vmid)['name']
-        delete_vm(proxmox, starrs, vmid)
+        delete_vm(proxmox, vmid)
         delete_starrs(starrs, vmname)
-        delete_vm_expire(vmid)
+        delete_vm_expire(db, vmid)
 
 
 def process_expired_vms_task():
@@ -37,7 +62,7 @@ def process_expired_vms_task():
         starrs = connect_starrs()
         expired_vms = get_expired_vms()
         print(expired_vms)
-    
+
     #    for vmid in expired_vms:
     #        vmname = get_vm_config(proxmox, vmid)['name']
     #        delete_vm(proxmox, starrs, vmid)