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 sqlalchemy.ext.associationproxy import association_proxy
from time import sleep
from tqdm import tqdm
from typing import Union
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

    @staticmethod
    def get(corpus_status: Union['CorpusStatus', int, str]) -> 'CorpusStatus':
        if isinstance(corpus_status, CorpusStatus):
            return corpus_status
        if isinstance(corpus_status, int):
            return CorpusStatus(corpus_status)
        if isinstance(corpus_status, str):
            return CorpusStatus[corpus_status]
        raise TypeError('corpus_status must be CorpusStatus, int, or str')


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

    @staticmethod
    def get(job_status: Union['JobStatus', int, str]) -> 'JobStatus':
        if isinstance(job_status, JobStatus):
            return job_status
        if isinstance(job_status, int):
            return JobStatus(job_status)
        if isinstance(job_status, str):
            return JobStatus[job_status]
        raise TypeError('job_status must be JobStatus, int, or str')


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

    @staticmethod
    def get(permission: Union['Permission', int, str]) -> 'Permission':
        if isinstance(permission, Permission):
            return permission
        if isinstance(permission, int):
            return Permission(permission)
        if isinstance(permission, str):
            return Permission[permission]
        raise TypeError('permission must be Permission, int, or str')


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


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


class CorpusFollowerPermission(IntEnum):
    VIEW = 1
    ADD_CORPUS_FILE = 2
    UPDATE_CORPUS_FILE = 4
    REMOVE_CORPUS_FILE = 8
    GENERATE_SHARE_LINK = 16
    REMOVE_FOLLOWER = 32
    UPDATE_FOLLOWER = 64

    @staticmethod
    def get(corpus_follower_permission: Union['CorpusFollowerPermission', int, str]) -> 'CorpusFollowerPermission':
        if isinstance(corpus_follower_permission, CorpusFollowerPermission):
            return corpus_follower_permission
        if isinstance(corpus_follower_permission, int):
            return CorpusFollowerPermission(corpus_follower_permission)
        if isinstance(corpus_follower_permission, str):
            return CorpusFollowerPermission[corpus_follower_permission]
        raise TypeError('corpus_follower_permission must be CorpusFollowerPermission, int, or str')
# 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)