1 files changed, 25 insertions, 46 deletions

diff --git a/fbin/db.py b/fbin/db.py
index be79c76..876efeb 100644
--- a/fbin/db.py
+++ b/fbin/db.py
@@ -4,57 +4,52 @@ import mimetypes
 import os
 
 from flask import current_app
-from sqlalchemy import create_engine, Column, Integer, String, DateTime, Text, Index, ForeignKey, Boolean, JSON, BigInteger
-from sqlalchemy.ext.declarative import declarative_base
-from sqlalchemy.orm import sessionmaker, relation, backref
 from sqlalchemy.orm.exc import NoResultFound
 from sqlalchemy.exc import IntegrityError
-from sqlalchemy.sql import and_
+from flask_sqlalchemy import SQLAlchemy
 
-engine = create_engine(current_app.config['DB_URI'])
+db = SQLAlchemy()
 
-Base = declarative_base(bind = engine)
-
-class User(Base):
+class User(db.Model):
     __tablename__ = 'users'
 
-    id = Column(Integer, primary_key = True)
-    username = Column(String, unique = True, index = True)
-    jab_id = Column(String(24), unique = True, index = True)
-    api_key_date = Column(DateTime, default = datetime.datetime.utcnow)
-    files = relation('File', backref = 'user', order_by = 'File.date.desc()')
+    id = db.Column(db.Integer, primary_key=True)
+    username = db.Column(db.String, unique=True, index=True)
+    jab_id = db.Column(db.String(24), unique=True, index=True)
+    api_key_date = db.Column(db.DateTime, default=datetime.datetime.utcnow)
+    files = db.relation('File', backref='user', order_by='File.date.desc()')
 
     def __init__(self, username, jab_id):
         self.username = username
         self.jab_id = jab_id
 
-class UserSession(Base):
+class UserSession(db.Model):
     __tablename__ = 'sessions'
 
-    id = Column(Integer, primary_key = True)
-    user_id = Column(Integer, ForeignKey('users.id'), index = True)
-    access_token = Column(String)
-    refresh_token = Column(String)
-    updated = Column(DateTime)
+    id = db.Column(db.Integer, primary_key=True)
+    user_id = db.Column(db.Integer, db.ForeignKey('users.id'), index=True)
+    access_token = db.Column(db.String)
+    refresh_token = db.Column(db.String)
+    updated = db.Column(db.DateTime)
 
     def __init__(self, user_id, access_token, refresh_token):
         self.user_id = user_id
         self.access_token = access_token
         self.refresh_token = refresh_token
 
-class File(Base):
+class File(db.Model):
     __tablename__ = 'files'
 
-    id = Column(Integer, primary_key = True)
-    hash = Column(String, unique = True, index = True)
-    filename = Column(String)
-    size = Column(BigInteger)
-    date = Column(DateTime)
-    user_id = Column(Integer, ForeignKey('users.id'), nullable = True)
-    ip = Column(String)
-    accessed = Column(DateTime)
-    scanned = Column(Boolean, nullable=False, default=False)
-    blocked_reason = Column(JSON)
+    id = db.Column(db.Integer, primary_key=True)
+    hash = db.Column(db.String, unique=True, index=True)
+    filename = db.Column(db.String)
+    size = db.Column(db.BigInteger)
+    date = db.Column(db.DateTime)
+    user_id = db.Column(db.Integer, db.ForeignKey('users.id'), nullable=True)
+    ip = db.Column(db.String)
+    accessed = db.Column(db.DateTime)
+    scanned = db.Column(db.Boolean, nullable=False, default=False)
+    blocked_reason = db.Column(db.JSON)
 
     def __init__(self, hash, filename, size, date, user_id = None, ip = None):
         self.hash = hash
@@ -114,19 +109,3 @@ class File(Base):
     @property
     def exists(self):
         return os.path.exists(self.get_path())
-
-Base.metadata.create_all()
-Session = sessionmaker(bind = engine, autoflush = True, autocommit = False)
-
-@contextmanager
-def session_scope():
-    session = Session()
-    try:
-        session.expire_on_commit = False
-        yield session
-        session.commit()
-    except:
-        session.rollback()
-        raise
-    finally:
-        session.close()