PyCS/pycs/database/Database.py

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

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:', 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")

        # create tables
        with 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                NOT NULL,
                    FOREIGN KEY (file) REFERENCES files(id)
                        ON UPDATE CASCADE ON DELETE CASCADE
                )
            ''')
            # cursor.execute('''
            #     CREATE INDEX IF NOT EXISTS idx_results_label ON results(label)
            # ''')

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

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

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

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

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

        :return: list of all available models
        """
        with closing(self.con.cursor()) as cursor:
            cursor.execute('SELECT * FROM models')
            return list(map(
                lambda row: Model(self, row),
                cursor.fetchall()
            ))

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

        :param identifier: unique identifier
        :return: model
        """
        with closing(self.con.cursor()) as cursor:
            cursor.execute('SELECT * FROM models WHERE id = ?', [identifier])
            row = cursor.fetchone()

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

            return None

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

        :return: list of all available label providers
        """
        with closing(self.con.cursor()) as cursor:
            cursor.execute('SELECT * FROM label_providers')
            return list(map(
                lambda row: LabelProvider(self, row),
                cursor.fetchall()
            ))

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

        :param identifier: unique identifier
        :return: label provider
        """
        with closing(self.con.cursor()) as cursor:
            cursor.execute('SELECT * FROM label_providers WHERE id = ?', [identifier])
            row = cursor.fetchone()

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

            return None

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

        :return: list of all available projects
        """
        with closing(self.con.cursor()) as cursor:
            cursor.execute('SELECT * FROM projects')
            return list(map(
                lambda row: Project(self, row),
                cursor.fetchall()
            ))

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

        :param identifier: unique identifier
        :return: project
        """
        with closing(self.con.cursor()) as cursor:
            cursor.execute('SELECT * FROM projects WHERE id = ?', [identifier])
            row = cursor.fetchone()

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

            return None

    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
        """
        with closing(self.con.cursor()) as cursor:
            cursor.execute('SELECT * FROM collections WHERE id = ?', [identifier])
            row = cursor.fetchone()

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

            return None

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

        :param identifier: unique identifier
        :return: file
        """
        with closing(self.con.cursor()) as cursor:
            cursor.execute('SELECT * FROM files WHERE id = ?', [identifier])
            row = cursor.fetchone()

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

            return None

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

        :param identifier: unique identifier
        :return: result
        """
        with closing(self.con.cursor()) as cursor:
            cursor.execute('SELECT * FROM results WHERE id = ?', [identifier])
            row = cursor.fetchone()

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

            return None