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