PyCS/pycs/database/Database.py

import sqlite3
from contextlib import closing
from time import time
from typing import Optional, Iterator

from pycs.database.Collection import Collection
from pycs.database.File import File
from pycs.database.LabelProvider import LabelProvider
from pycs.database.Model import Model
from pycs.database.Project import Project
from pycs.database.Result import Result
from pycs.database.discovery.LabelProviderDiscovery import discover as discover_label_providers
from pycs.database.discovery.ModelDiscovery import discover as discover_models


class Database:
    """
    opens an sqlite database and allows to access several objects
    """

    def __init__(self, path: str = ':memory:', initialization=True, discovery=True):
        """
        opens or creates a given sqlite database and creates all required tables

        :param path: path to sqlite database
        """
        # save properties
        self.path = path

        # initialize database connection
        self.con = sqlite3.connect(path)
        self.con.execute("PRAGMA foreign_keys = ON")

        if initialization:
            # create tables
            with self, closing(self.con.cursor()) as cursor:
                cursor.execute('''
                    CREATE TABLE IF NOT EXISTS models (
                        id          INTEGER PRIMARY KEY,
                        name        TEXT                NOT NULL,
                        description TEXT,
                        root_folder TEXT                NOT NULL UNIQUE,
                        supports    TEXT                NOT NULL
                    )
                ''')
                cursor.execute('''
                    CREATE TABLE IF NOT EXISTS label_providers (
                        id          INTEGER PRIMARY KEY,
                        name        TEXT                NOT NULL,
                        description TEXT,
                        root_folder TEXT                NOT NULL UNIQUE
                    )
                ''')

                cursor.execute('''
                    CREATE TABLE IF NOT EXISTS projects (
                        id             INTEGER PRIMARY KEY,
                        name           TEXT                NOT NULL,
                        description    TEXT,
                        created        INTEGER             NOT NULL,
                        model          INTEGER,
                        label_provider INTEGER,
                        root_folder    TEXT                NOT NULL UNIQUE,
                        external_data  BOOL                NOT NULL,
                        data_folder    TEXT                NOT NULL,
                        FOREIGN KEY (model) REFERENCES models(id)
                            ON UPDATE CASCADE ON DELETE SET NULL,
                        FOREIGN KEY (label_provider) REFERENCES label_providers(id)
                            ON UPDATE CASCADE ON DELETE SET NULL
                    )
                ''')
                cursor.execute('''
                    CREATE TABLE IF NOT EXISTS labels (
                        id        INTEGER PRIMARY KEY,
                        project   INTEGER             NOT NULL,
                        parent    INTEGER,
                        created   INTEGER             NOT NULL,
                        reference TEXT,
                        name      TEXT                NOT NULL,
                        FOREIGN KEY (project) REFERENCES projects(id)
                            ON UPDATE CASCADE ON DELETE CASCADE,
                        FOREIGN KEY (parent) REFERENCES labels(id)
                            ON UPDATE CASCADE ON DELETE SET NULL,
                        UNIQUE(project, reference)
                    )
                ''')
                cursor.execute('''
                    CREATE TABLE IF NOT EXISTS collections (
                        id          INTEGER          PRIMARY KEY,
                        project     INTEGER NOT NULL,
                        reference   TEXT    NOT NULL,
                        name        TEXT    NOT NULL,
                        description TEXT,
                        position    INTEGER NOT NULL,
                        autoselect  INTEGER NOT NULL,
                        FOREIGN KEY (project) REFERENCES projects(id)
                            ON UPDATE CASCADE ON DELETE CASCADE,
                        UNIQUE(project, reference)
                    )
                ''')
                cursor.execute('''
                    CREATE TABLE IF NOT EXISTS files (
                        id         INTEGER PRIMARY KEY,
                        uuid       TEXT                NOT NULL,
                        project    INTEGER             NOT NULL,
                        collection INTEGER,
                        type       TEXT                NOT NULL,
                        name       TEXT                NOT NULL,
                        extension  TEXT                NOT NULL,
                        size       INTEGER             NOT NULL,
                        created    INTEGER             NOT NULL,
                        path       TEXT                NOT NULL,
                        frames     INTEGER,
                        fps        FLOAT,
                        FOREIGN KEY (project) REFERENCES projects(id)
                            ON UPDATE CASCADE ON DELETE CASCADE,
                        FOREIGN KEY (collection) REFERENCES collections(id)
                            ON UPDATE CASCADE ON DELETE SET NULL,
                        UNIQUE(project, path)
                    )
                ''')
                cursor.execute('''
                    CREATE TABLE IF NOT EXISTS results (
                        id     INTEGER PRIMARY KEY,
                        file   INTEGER             NOT NULL,
                        origin TEXT                NOT NULL,
                        type   TEXT                NOT NULL,
                        label  INTEGER,
                        data   TEXT,
                        FOREIGN KEY (file) REFERENCES files(id)
                            ON UPDATE CASCADE ON DELETE CASCADE
                    )
                ''')

        if discovery:
            # run discovery modules
            with self:
                discover_models(self.con)
                discover_label_providers(self.con)

    def close(self):
        self.con.close()

    def copy(self):
        return Database(self.path, initialization=False, discovery=False)

    def commit(self):
        self.con.commit()

    def __enter__(self):
        self.con.__enter__()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.con.__exit__(exc_type, exc_val, exc_tb)

    def get_object_by_id(self, table_name: str, identifier: int, cls):
        """
        create an object from cls and a row fetched from table_name and
        identified by the identifier

        :param table_name: table name
        :param identifier: unique identifier
        :param cls: class that is used to create the object
        :return: object of type cls
        """
        with closing(self.con.cursor()) as cursor:
            cursor.execute(f'SELECT * FROM {table_name} WHERE id = ?', [identifier])
            row = cursor.fetchone()

            if row is not None:
                return cls(self, row)

            return None

    def get_objects(self, table_name: str, cls):
        """
        get a list of all available objects in the table

        :param table_name: table name
        :param cls: class that is used to create the objects
        :return: list of object of type cls
        """

        with closing(self.con.cursor()) as cursor:
            cursor.execute(f'SELECT * FROM {table_name}')
            for row in cursor:
                yield cls(self, row)


    def models(self) -> Iterator[Model]:
        """
        get a list of all available models

        :return: iterator of models
        """
        return self.get_objects("models", Model)

    def model(self, identifier: int) -> Optional[Model]:
        """
        get a model using its unique identifier

        :param identifier: unique identifier
        :return: model
        """
        return self.get_object_by_id("models", identifier, Model)

    def label_providers(self) -> Iterator[LabelProvider]:
        """
        get a list of all available label providers

        :return: iterator over label providers
        """
        return self.get_objects("label_providers", LabelProvider)

    def label_provider(self, identifier: int) -> Optional[LabelProvider]:
        """
        get a label provider using its unique identifier

        :param identifier: unique identifier
        :return: label provider
        """
        return self.get_object_by_id("label_providers", identifier, LabelProvider)

    def projects(self) -> Iterator[Project]:
        """
        get a list of all available projects

        :return: iterator over projects
        """
        return self.get_objects("projects", Project)

    def project(self, identifier: int) -> Optional[Project]:
        """
        get a project using its unique identifier

        :param identifier: unique identifier
        :return: project
        """
        return self.get_object_by_id("projects", identifier, Project)

    def create_project(self,
                       name: str,
                       description: str,
                       model: Model,
                       label_provider: Optional[LabelProvider],
                       root_folder: str,
                       external_data: bool,
                       data_folder: str):
        """
        insert a project into the database

        :param name: project name
        :param description: project description
        :param model: used model
        :param label_provider: used label provider (optional)
        :param root_folder: path to project folder
        :param external_data: whether an external data directory is used
        :param data_folder: path to data folder
        :return: created project
        """
        # prepare some values
        created = int(time())
        label_provider_id = label_provider.identifier if label_provider is not None else None

        # insert statement
        with closing(self.con.cursor()) as cursor:
            cursor.execute('''
                INSERT INTO projects (
                    name, description, created, model, label_provider, root_folder, external_data, 
                    data_folder
                )
                VALUES (?, ?, ?, ?, ?, ?, ?, ?)
            ''', (name, description, created, model.identifier, label_provider_id, root_folder,
                  external_data, data_folder))

            return self.project(cursor.lastrowid)

    def collection(self, identifier: int) -> Optional[Collection]:
        """
        get a collection using its unique identifier

        :param identifier: unique identifier
        :return: collection
        """
        return self.get_object_by_id("collections", identifier, Collection)

    def file(self, identifier) -> Optional[File]:
        """
        get a file using its unique identifier

        :param identifier: unique identifier
        :return: file
        """
        return self.get_object_by_id("files", identifier, File)

    def result(self, identifier) -> Optional[Result]:
        """
        get a result using its unique identifier

        :param identifier: unique identifier
        :return: result
        """
        return self.get_object_by_id("results", identifier, Result)