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: 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, configuration_file TEXT NOT NULL, UNIQUE(root_folder, configuration_file) ) ''') 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, hierarchy_level TEXT, 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): """ close database file """ self.con.close() def copy(self): """ Create a copy of this database object. This can be used to access the database from another thread. Table initialization and model and label provider discovery is disabled to speedup this function. :return: Database """ return Database(self.path, initialization=False, discovery=False) def commit(self): """ commit changes """ 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 models(self) -> Iterator[Model]: """ get a list of all available models :return: iterator of models """ with closing(self.con.cursor()) as cursor: cursor.execute('SELECT * FROM models') for row in cursor: yield Model(self, row) 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) -> Iterator[LabelProvider]: """ get a list of all available label providers :return: iterator over label providers """ with closing(self.con.cursor()) as cursor: cursor.execute('SELECT * FROM label_providers') for row in cursor: yield LabelProvider(self, row) 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) -> Iterator[Project]: """ get a list of all available projects :return: iterator over projects """ with closing(self.con.cursor()) as cursor: cursor.execute('SELECT * FROM projects') for row in cursor: yield Project(self, row) 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