mirror of
https://github.com/Comfy-Org/ComfyUI.git
synced 2026-08-24 18:10:27 +08:00
refactor, adding tests
This commit is contained in:
pythongosssss
0
app/database/__init__.py
Normal file
0
app/database/__init__.py
Normal file
126
app/database/db.py
Normal file
126
app/database/db.py
Normal file
@@ -0,0 +1,126 @@
|
||||
import logging
|
||||
import os
|
||||
import sqlite3
|
||||
from contextlib import contextmanager
|
||||
from queue import Queue, Empty, Full
|
||||
import threading
|
||||
from app.database.updater import DatabaseUpdater
|
||||
import folder_paths
|
||||
from comfy.cli_args import args
|
||||
|
||||
|
||||
class Database:
|
||||
def __init__(self, database_path=None, pool_size=1):
|
||||
if database_path is None:
|
||||
self.exists = False
|
||||
database_path = "file::memory:?cache=shared"
|
||||
else:
|
||||
self.exists = os.path.exists(database_path)
|
||||
|
||||
self.database_path = database_path
|
||||
self.pool_size = pool_size
|
||||
# Store connections in a pool, default to 1 as normal usage is going to be from a single thread at a time
|
||||
self.connection_pool: Queue = Queue(maxsize=pool_size)
|
||||
self._db_lock = threading.Lock()
|
||||
self._initialized = False
|
||||
self._closing = False
|
||||
self._after_update_callbacks = []
|
||||
|
||||
def _setup(self):
|
||||
if self._initialized:
|
||||
return
|
||||
|
||||
with self._db_lock:
|
||||
if not self._initialized:
|
||||
self._make_db()
|
||||
self._initialized = True
|
||||
|
||||
def _create_connection(self):
|
||||
# TODO: Catch error for sqlite lib missing on linux
|
||||
logging.info(f"Creating connection to {self.database_path}")
|
||||
conn = sqlite3.connect(
|
||||
self.database_path,
|
||||
check_same_thread=False,
|
||||
uri=self.database_path.startswith("file::"),
|
||||
)
|
||||
conn.execute("PRAGMA foreign_keys = ON")
|
||||
self.exists = True
|
||||
logging.info(f"Connected!")
|
||||
return conn
|
||||
|
||||
def _make_db(self):
|
||||
with self._get_connection() as con:
|
||||
updater = DatabaseUpdater(con, self.database_path)
|
||||
result = updater.update()
|
||||
if result is not None:
|
||||
old_version, new_version = result
|
||||
|
||||
for callback in self._after_update_callbacks:
|
||||
callback(old_version, new_version)
|
||||
|
||||
def _transform(self, row, columns):
|
||||
return {col.name: value for value, col in zip(row, columns)}
|
||||
|
||||
@contextmanager
|
||||
def _get_connection(self):
|
||||
if self._closing:
|
||||
raise Exception("Database is shutting down")
|
||||
|
||||
try:
|
||||
# Try to get connection from pool
|
||||
connection = self.connection_pool.get_nowait()
|
||||
except Empty:
|
||||
# Create new connection if pool is empty
|
||||
connection = self._create_connection()
|
||||
|
||||
try:
|
||||
yield connection
|
||||
finally:
|
||||
try:
|
||||
# Try to add to pool if it's empty
|
||||
self.connection_pool.put_nowait(connection)
|
||||
except Full:
|
||||
# Pool is full, close the connection
|
||||
connection.close()
|
||||
|
||||
@contextmanager
|
||||
def get_connection(self):
|
||||
# Setup the database if it's not already initialized
|
||||
self._setup()
|
||||
with self._get_connection() as connection:
|
||||
yield connection
|
||||
|
||||
def execute(self, sql, *args):
|
||||
with self.get_connection() as connection:
|
||||
cursor = connection.execute(sql, args)
|
||||
results = cursor.fetchall()
|
||||
return results
|
||||
|
||||
def register_after_update_callback(self, callback):
|
||||
self._after_update_callbacks.append(callback)
|
||||
|
||||
def close(self):
|
||||
if self._closing:
|
||||
return
|
||||
# Drain and close all connections in the pool
|
||||
self._closing = True
|
||||
while True:
|
||||
try:
|
||||
conn = self.connection_pool.get_nowait()
|
||||
conn.close()
|
||||
except Empty:
|
||||
break
|
||||
self._closing = False
|
||||
|
||||
def __del__(self):
|
||||
try:
|
||||
self.close()
|
||||
except:
|
||||
pass
|
||||
|
||||
|
||||
# Create a global instance
|
||||
db_path = None
|
||||
if not args.memory_database:
|
||||
db_path = folder_paths.get_user_directory() + "/comfyui.db"
|
||||
db = Database(db_path)
|
||||
301
app/database/entities.py
Normal file
301
app/database/entities.py
Normal file
@@ -0,0 +1,301 @@
|
||||
from typing import Optional, Any, Callable
|
||||
from dataclasses import dataclass
|
||||
from functools import wraps
|
||||
from aiohttp import web
|
||||
from app.database.db import db
|
||||
|
||||
primitives = (bool, str, int, float, type(None))
|
||||
|
||||
|
||||
def is_primitive(obj):
|
||||
return isinstance(obj, primitives)
|
||||
|
||||
|
||||
class ValidationError(Exception):
|
||||
def __init__(self, message: str, field: str = None, value: Any = None):
|
||||
self.message = message
|
||||
self.field = field
|
||||
self.value = value
|
||||
super().__init__(self.message)
|
||||
|
||||
def to_json(self):
|
||||
result = {"message": self.message}
|
||||
if self.field is not None:
|
||||
result["field"] = self.field
|
||||
if self.value is not None:
|
||||
result["value"] = self.value
|
||||
return result
|
||||
|
||||
def __str__(self) -> str:
|
||||
return f"{self.message} {self.field} {self.value}"
|
||||
|
||||
|
||||
class EntityCommon(dict):
|
||||
@classmethod
|
||||
def _get_route(cls, include_key: bool):
|
||||
route = f"/db/{cls.__table_name__}"
|
||||
if include_key:
|
||||
route += "".join([f"/{{{k}}}" for k in cls.__key_columns__])
|
||||
return route
|
||||
|
||||
@classmethod
|
||||
def _register_route(cls, routes, verb: str, include_key: bool, handler: Callable):
|
||||
route = cls._get_route(include_key)
|
||||
|
||||
@getattr(routes, verb)(route)
|
||||
async def _(request):
|
||||
try:
|
||||
data = await handler(request)
|
||||
return web.json_response(data)
|
||||
except ValidationError as e:
|
||||
return web.json_response(e.to_json(), status=400)
|
||||
|
||||
@classmethod
|
||||
def _transform(cls, row: list[Any]):
|
||||
return {col: value for col, value in zip(cls.__columns__, row)}
|
||||
|
||||
@classmethod
|
||||
def _transform_rows(cls, rows: list[list[Any]]):
|
||||
return [cls._transform(row) for row in rows]
|
||||
|
||||
@classmethod
|
||||
def _validate(cls, fields: list[str], data: dict, allow_missing: bool = False):
|
||||
result = {}
|
||||
|
||||
if not isinstance(data, dict):
|
||||
raise ValidationError("Invalid data")
|
||||
|
||||
# Ensure all required fields are present
|
||||
for field in data:
|
||||
if field not in fields:
|
||||
raise ValidationError("Unknown field", field)
|
||||
|
||||
for key in fields:
|
||||
col = cls.__columns__[key]
|
||||
if key not in data:
|
||||
if col.required and not allow_missing:
|
||||
raise ValidationError("Missing field", key)
|
||||
else:
|
||||
# e.g. for updates, we allow missing fields
|
||||
continue
|
||||
elif data[key] is None and col.required:
|
||||
# Dont allow None for required fields
|
||||
raise ValidationError("Required field", key)
|
||||
|
||||
# Validate data type
|
||||
value = data[key]
|
||||
|
||||
if value is not None and not is_primitive(value):
|
||||
raise ValidationError("Invalid value", key, value)
|
||||
|
||||
try:
|
||||
type = col.type
|
||||
if value is not None and not isinstance(value, type):
|
||||
value = type(value)
|
||||
result[key] = value
|
||||
except Exception:
|
||||
raise ValidationError("Invalid value", key, value)
|
||||
|
||||
return result
|
||||
|
||||
@classmethod
|
||||
def _validate_id(cls, id: dict):
|
||||
return cls._validate(cls.__key_columns__, id)
|
||||
|
||||
@classmethod
|
||||
def _validate_data(cls, data: dict):
|
||||
return cls._validate(cls.__columns__.keys(), data)
|
||||
|
||||
def __setattr__(self, name, value):
|
||||
if name in self.__columns__:
|
||||
self[name] = value
|
||||
super().__setattr__(name, value)
|
||||
|
||||
def __getattr__(self, name):
|
||||
if name in self:
|
||||
return self[name]
|
||||
raise AttributeError(f"'{self.__class__.__name__}' has no attribute '{name}'")
|
||||
|
||||
|
||||
class GetEntity(EntityCommon):
|
||||
@classmethod
|
||||
def get(cls, top: Optional[int] = None, where: Optional[str] = None):
|
||||
limit = ""
|
||||
if top is not None and isinstance(top, int):
|
||||
limit = f" LIMIT {top}"
|
||||
result = db.execute(
|
||||
f"SELECT * FROM {cls.__table_name__}{limit}{f' WHERE {where}' if where else ''}",
|
||||
)
|
||||
|
||||
# Map each row in result to an instance of the class
|
||||
return cls._transform_rows(result)
|
||||
|
||||
@classmethod
|
||||
def register_route(cls, routes):
|
||||
async def get_handler(request):
|
||||
top = request.rel_url.query.get("top", None)
|
||||
if top is not None:
|
||||
try:
|
||||
top = int(top)
|
||||
except Exception:
|
||||
raise ValidationError("Invalid top parameter", "top", top)
|
||||
return cls.get(top)
|
||||
|
||||
cls._register_route(routes, "get", False, get_handler)
|
||||
|
||||
|
||||
class GetEntityById(EntityCommon):
|
||||
@classmethod
|
||||
def get_by_id(cls, id: dict):
|
||||
id = cls._validate_id(id)
|
||||
|
||||
result = db.execute(
|
||||
f"SELECT * FROM {cls.__table_name__} WHERE {cls.__where_clause__}",
|
||||
*[id[key] for key in cls.__key_columns__],
|
||||
)
|
||||
|
||||
return cls._transform_rows(result)
|
||||
|
||||
@classmethod
|
||||
def register_route(cls, routes):
|
||||
async def get_by_id_handler(request):
|
||||
id = {key: request.match_info.get(key, None) for key in cls.__key_columns__}
|
||||
return cls.get_by_id(id)
|
||||
|
||||
cls._register_route(routes, "get", True, get_by_id_handler)
|
||||
|
||||
|
||||
class CreateEntity(EntityCommon):
|
||||
@classmethod
|
||||
def create(cls, data: dict, allow_upsert: bool = False):
|
||||
data = cls._validate_data(data)
|
||||
values = ", ".join(["?"] * len(data))
|
||||
on_conflict = ""
|
||||
|
||||
data_keys = ", ".join(list(data.keys()))
|
||||
if allow_upsert:
|
||||
# Remove key columns from data
|
||||
upsert_keys = [key for key in data if key not in cls.__key_columns__]
|
||||
|
||||
set_clause = ", ".join([f"{k} = excluded.{k}" for k in upsert_keys])
|
||||
on_conflict = f" ON CONFLICT ({', '.join(cls.__key_columns__)}) DO UPDATE SET {set_clause}"
|
||||
sql = f"INSERT INTO {cls.__table_name__} ({data_keys}) VALUES ({values}){on_conflict} RETURNING *"
|
||||
result = db.execute(
|
||||
sql,
|
||||
*[data[key] for key in data],
|
||||
)
|
||||
|
||||
if len(result) == 0:
|
||||
raise RuntimeError("Failed to create entity")
|
||||
|
||||
return cls._transform_rows(result)[0]
|
||||
|
||||
@classmethod
|
||||
def register_route(cls, routes):
|
||||
async def create_handler(request):
|
||||
data = await request.json()
|
||||
return cls.create(data)
|
||||
|
||||
cls._register_route(routes, "post", False, create_handler)
|
||||
|
||||
|
||||
class UpdateEntity(EntityCommon):
|
||||
@classmethod
|
||||
def update(cls, id: list, data: dict):
|
||||
pass
|
||||
|
||||
|
||||
class UpsertEntity(CreateEntity):
|
||||
@classmethod
|
||||
def upsert(cls, data: dict):
|
||||
return cls.create(data, allow_upsert=True)
|
||||
|
||||
@classmethod
|
||||
def register_route(cls, routes):
|
||||
async def upsert_handler(request):
|
||||
data = await request.json()
|
||||
return cls.upsert(data)
|
||||
|
||||
cls._register_route(routes, "put", False, upsert_handler)
|
||||
|
||||
|
||||
class DeleteEntity(EntityCommon):
|
||||
@classmethod
|
||||
def delete(cls, id: list):
|
||||
pass
|
||||
|
||||
|
||||
class BaseEntity(GetEntity, CreateEntity, UpdateEntity, DeleteEntity, GetEntityById):
|
||||
pass
|
||||
|
||||
|
||||
@dataclass
|
||||
class Column:
|
||||
type: Any
|
||||
required: bool = False
|
||||
key: bool = False
|
||||
default: Any = None
|
||||
|
||||
|
||||
def column(type_: Any, required: bool = False, key: bool = False, default: Any = None):
|
||||
return Column(type_, required, key, default)
|
||||
|
||||
|
||||
def table(table_name: str):
|
||||
def decorator(cls):
|
||||
# Store table name
|
||||
cls.__table_name__ = table_name
|
||||
|
||||
# Process column definitions
|
||||
columns: dict[str, Column] = {}
|
||||
for attr_name, attr_value in cls.__dict__.items():
|
||||
if isinstance(attr_value, Column):
|
||||
columns[attr_name] = attr_value
|
||||
|
||||
# Store columns metadata
|
||||
cls.__columns__ = columns
|
||||
cls.__key_columns__ = [col for col in columns if columns[col].key]
|
||||
cls.__column_csv__ = ", ".join([col for col in columns])
|
||||
cls.__where_clause__ = " AND ".join(
|
||||
[f"{col} = ?" for col in cls.__key_columns__]
|
||||
)
|
||||
|
||||
# Add initialization
|
||||
original_init = cls.__init__
|
||||
|
||||
@wraps(original_init)
|
||||
def new_init(self, *args, **kwargs):
|
||||
# Initialize columns with default values
|
||||
for col_name, col_def in cls.__columns__.items():
|
||||
setattr(self, col_name, col_def.default)
|
||||
# Call original init
|
||||
original_init(self, *args, **kwargs)
|
||||
|
||||
cls.__init__ = new_init
|
||||
return cls
|
||||
|
||||
return decorator
|
||||
|
||||
|
||||
def test():
|
||||
@table("models")
|
||||
class Model(BaseEntity):
|
||||
id: int = column(int, required=True, key=True)
|
||||
path: str = column(str, required=True)
|
||||
name: str = column(str, required=True)
|
||||
description: Optional[str] = column(str)
|
||||
architecture: Optional[str] = column(str)
|
||||
type: str = column(str, required=True)
|
||||
hash: Optional[str] = column(str)
|
||||
source_url: Optional[str] = column(str)
|
||||
|
||||
return Model
|
||||
|
||||
|
||||
@table("test")
|
||||
class Test(GetEntity, CreateEntity):
|
||||
id: int = column(int, required=True, key=True)
|
||||
test: str = column(str, required=True)
|
||||
|
||||
|
||||
Model = test()
|
||||
32
app/database/routes.py
Normal file
32
app/database/routes.py
Normal file
@@ -0,0 +1,32 @@
|
||||
from app.database.db import db
|
||||
from aiohttp import web
|
||||
|
||||
def create_routes(
|
||||
routes, prefix, entity, get=False, get_by_id=False, post=False, delete=False
|
||||
):
|
||||
if get:
|
||||
@routes.get(f"/{prefix}/{table}")
|
||||
async def get_table(request):
|
||||
connection = db.get_connection()
|
||||
cursor = connection.cursor()
|
||||
cursor.execute(f"SELECT * FROM {table}")
|
||||
rows = cursor.fetchall()
|
||||
return web.json_response(rows)
|
||||
|
||||
if get_by_id:
|
||||
@routes.get(f"/{prefix}/{table}/{id}")
|
||||
async def get_table_by_id(request):
|
||||
connection = db.get_connection()
|
||||
cursor = connection.cursor()
|
||||
cursor.execute(f"SELECT * FROM {table} WHERE id = {id}")
|
||||
row = cursor.fetchone()
|
||||
return web.json_response(row)
|
||||
|
||||
if post:
|
||||
@routes.post(f"/{prefix}/{table}")
|
||||
async def post_table(request):
|
||||
data = await request.json()
|
||||
connection = db.get_connection()
|
||||
cursor = connection.cursor()
|
||||
cursor.execute(f"INSERT INTO {table} ({data}) VALUES ({data})")
|
||||
return web.json_response({"status": "success"})
|
||||
79
app/database/updater.py
Normal file
79
app/database/updater.py
Normal file
@@ -0,0 +1,79 @@
|
||||
import logging
|
||||
import os
|
||||
import sqlite3
|
||||
from app.database.versions.v1 import v1
|
||||
|
||||
|
||||
class DatabaseUpdater:
|
||||
def __init__(self, connection, database_path):
|
||||
self.connection = connection
|
||||
self.database_path = database_path
|
||||
self.current_version = self.get_db_version()
|
||||
self.version_updates = {
|
||||
1: v1,
|
||||
}
|
||||
self.max_version = max(self.version_updates.keys())
|
||||
self.update_required = self.current_version < self.max_version
|
||||
logging.info(f"Database version: {self.current_version}")
|
||||
|
||||
def get_db_version(self):
|
||||
return self.connection.execute("PRAGMA user_version").fetchone()[0]
|
||||
|
||||
def backup(self):
|
||||
bkp_path = self.database_path + ".bkp"
|
||||
if os.path.exists(bkp_path):
|
||||
# TODO: auto-rollback failed upgrades
|
||||
raise Exception(
|
||||
f"Database backup already exists, this indicates that a previous upgrade failed. Please restore this backup before continuing. Backup location: {bkp_path}"
|
||||
)
|
||||
|
||||
bkp = sqlite3.connect(bkp_path)
|
||||
self.connection.backup(bkp)
|
||||
bkp.close()
|
||||
logging.info("Database backup taken pre-upgrade.")
|
||||
return bkp_path
|
||||
|
||||
def update(self):
|
||||
if not self.update_required:
|
||||
return None
|
||||
|
||||
bkp_version = self.current_version
|
||||
bkp_path = None
|
||||
if self.current_version > 0:
|
||||
bkp_path = self.backup()
|
||||
|
||||
logging.info(f"Updating database: {self.current_version} -> {self.max_version}")
|
||||
|
||||
dirname = os.path.dirname(__file__)
|
||||
cursor = self.connection.cursor()
|
||||
for version in range(self.current_version + 1, self.max_version + 1):
|
||||
filename = os.path.join(dirname, f"versions/v{version}.sql")
|
||||
if not os.path.exists(filename):
|
||||
raise Exception(
|
||||
f"Database update script for version {version} not found"
|
||||
)
|
||||
|
||||
try:
|
||||
with open(filename, "r") as file:
|
||||
sql = file.read()
|
||||
cursor.executescript(sql)
|
||||
except Exception as e:
|
||||
raise Exception(
|
||||
f"Failed to execute update script for version {version}: {e}"
|
||||
)
|
||||
|
||||
method = self.version_updates[version]
|
||||
if method is not None:
|
||||
method(cursor)
|
||||
|
||||
cursor.execute("PRAGMA user_version = %d" % self.max_version)
|
||||
self.connection.commit()
|
||||
cursor.close()
|
||||
self.current_version = self.get_db_version()
|
||||
|
||||
if bkp_path:
|
||||
# Keep a copy of the backup in case something goes wrong and we need to rollback
|
||||
os.rename(bkp_path, self.database_path + f".v{bkp_version}.bkp")
|
||||
logging.info(f"Upgrade to successful.")
|
||||
|
||||
return (bkp_version, self.current_version)
|
||||
17
app/database/versions/v1.py
Normal file
17
app/database/versions/v1.py
Normal file
@@ -0,0 +1,17 @@
|
||||
from folder_paths import folder_names_and_paths, get_filename_list, get_full_path
|
||||
|
||||
|
||||
def v1(cursor):
|
||||
print("Updating to v1")
|
||||
for folder_name in folder_names_and_paths.keys():
|
||||
if folder_name == "custom_nodes":
|
||||
continue
|
||||
|
||||
files = get_filename_list(folder_name)
|
||||
for file in files:
|
||||
file_path = get_full_path(folder_name, file)
|
||||
file_without_extension = file.rsplit(".", maxsplit=1)[0]
|
||||
cursor.execute(
|
||||
"INSERT INTO models (path, name, type) VALUES (?, ?, ?)",
|
||||
(file_path, file_without_extension, folder_name),
|
||||
)
|
||||
41
app/database/versions/v1.sql
Normal file
41
app/database/versions/v1.sql
Normal file
@@ -0,0 +1,41 @@
|
||||
CREATE TABLE IF NOT EXISTS
|
||||
models (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
path TEXT NOT NULL,
|
||||
name TEXT NOT NULL,
|
||||
description TEXT,
|
||||
architecture TEXT,
|
||||
type TEXT NOT NULL,
|
||||
hash TEXT,
|
||||
source_url TEXT
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS
|
||||
tags (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
name TEXT NOT NULL UNIQUE
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS
|
||||
model_tags (
|
||||
model_id INTEGER NOT NULL,
|
||||
tag_id INTEGER NOT NULL,
|
||||
PRIMARY KEY (model_id, tag_id),
|
||||
FOREIGN KEY (model_id) REFERENCES models (id) ON DELETE CASCADE,
|
||||
FOREIGN KEY (tag_id) REFERENCES tags (id) ON DELETE CASCADE
|
||||
);
|
||||
|
||||
INSERT INTO
|
||||
tags (name)
|
||||
VALUES
|
||||
('character'),
|
||||
('style'),
|
||||
('concept'),
|
||||
('clothing'),
|
||||
('poses'),
|
||||
('background'),
|
||||
('vehicle'),
|
||||
('buildings'),
|
||||
('objects'),
|
||||
('animal'),
|
||||
('action');
|
||||
Reference in New Issue
Block a user