import sqlite3 from contextlib import closing from time import time from typing import Optional, List 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, created INTEGER NOT NULL, reference TEXT, name TEXT 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, 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, 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 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