models.py

from datetime import datetime, timedelta
from enum import Enum, IntEnum
from flask import current_app, url_for
from flask_hashids import HashidMixin
from flask_login import UserMixin
from time import sleep
from tqdm import tqdm
from werkzeug.security import generate_password_hash, check_password_hash
from werkzeug.utils import secure_filename
import json
import jwt
import os
import requests
import secrets
import shutil
import xml.etree.ElementTree as ET
import yaml
from app import db, hashids, login, mail, socketio
from app.converters.vrt import normalize_vrt_file
from app.email import create_message


##############################################################################
# enums                                                                      #
##############################################################################
# region enums
class CorpusStatus(IntEnum):
    UNPREPARED = 1
    SUBMITTED = 2
    QUEUED = 3
    BUILDING = 4
    BUILT = 5
    FAILED = 6
    STARTING_ANALYSIS_SESSION = 7
    RUNNING_ANALYSIS_SESSION = 8
    CANCELING_ANALYSIS_SESSION = 9


class JobStatus(IntEnum):
    INITIALIZING = 1
    SUBMITTED = 2
    QUEUED = 3
    RUNNING = 4
    CANCELING = 5
    CANCELED = 6
    COMPLETED = 7
    FAILED = 8


class Permission(IntEnum):
    '''
    Defines User permissions as integers by the power of 2. User permission
    can be evaluated using the bitwise operator &.
    '''
    ADMINISTRATE = 1
    CONTRIBUTE = 2
    USE_API = 4


class UserSettingJobStatusMailNotificationLevel(IntEnum):
    NONE = 1
    END = 2
    ALL = 3


class ProfilePrivacySettings(IntEnum):
    SHOW_EMAIL = 1
    SHOW_LAST_SEEN = 2
    SHOW_MEMBER_SINCE = 4
# endregion enums


##############################################################################
# mixins                                                                     #
##############################################################################
# region mixins
class FileMixin:
    '''
    Mixin for db.Model classes. All file related models should use this.
    '''
    creation_date = db.Column(db.DateTime, default=datetime.utcnow)
    filename = db.Column(db.String(255))
    mimetype = db.Column(db.String(255))

    def file_mixin_to_json_serializeable(self, backrefs=False, relationships=False):
        return {
            'creation_date': f'{self.creation_date.isoformat()}Z',
            'filename': self.filename,
            'mimetype': self.mimetype
        }
    
    @classmethod
    def create(cls, file_storage, **kwargs):
        filename = kwargs.pop('filename', file_storage.filename)
        mimetype = kwargs.pop('mimetype', file_storage.mimetype)
        obj = cls(
            filename=secure_filename(filename),
            mimetype=mimetype,
            **kwargs
        )
        db.session.add(obj)
        db.session.flush(objects=[obj])
        db.session.refresh(obj)
        try:
            file_storage.save(obj.path)
        except (AttributeError, OSError) as e:
            current_app.logger.error(e)
            db.session.rollback()
            raise e
        return obj
# endregion mixins


##############################################################################
# type_decorators                                                            #
##############################################################################
# region type_decorators
class IntEnumColumn(db.TypeDecorator):
    impl = db.Integer

    def __init__(self, enum_type, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.enum_type = enum_type

    def process_bind_param(self, value, dialect):
        if isinstance(value, self.enum_type) and isinstance(value.value, int):
            return value.value
        elif isinstance(value, int):
            return self.enum_type(value).value
        elif isinstance(value, str):
            return self.enum_type[value].value
        else:
            return TypeError()

    def process_result_value(self, value, dialect):
        return self.enum_type(value)


class ContainerColumn(db.TypeDecorator):
    impl = db.String

    def __init__(self, container_type, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.container_type = container_type

    def process_bind_param(self, value, dialect):
        if isinstance(value, self.container_type):
            return json.dumps(value)
        elif isinstance(value, str) and isinstance(json.loads(value), self.container_type):
            return value
        else:
            return TypeError()

    def process_result_value(self, value, dialect):
        return json.loads(value)
# endregion type_decorators


##############################################################################
# Models                                                                     #
##############################################################################
# region models
class Role(HashidMixin, db.Model):
    __tablename__ = 'roles'
    # Primary key
    id = db.Column(db.Integer, primary_key=True)
    # Fields
    name = db.Column(db.String(64), unique=True)
    default = db.Column(db.Boolean, default=False, index=True)
    permissions = db.Column(db.Integer, default=0)
    # Relationships
    users = db.relationship('User', backref='role', lazy='dynamic')

    def __repr__(self):
        return f'<Role {self.name}>'

    def add_permission(self, permission):
        if not self.has_permission(permission):
            self.permissions += permission

    def has_permission(self, permission):
        return self.permissions & permission == permission

    def remove_permission(self, permission):
        if self.has_permission(permission):
            self.permissions -= permission

    def reset_permissions(self):
        self.permissions = 0

    def to_json_serializeable(self, backrefs=False, relationships=False):
        json_serializeable = {
            'id': self.hashid,
            'default': self.default,
            'name': self.name,
            'permissions': self.permissions
        }
        if relationships:
            json_serializeable['users'] = {
                x.hashid: x.to_json_serializeable(relationships=True)