87 lines
2.6 KiB
Python
87 lines
2.6 KiB
Python
"""SQLite 连接与建表。"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import sqlite3
|
|
from contextlib import contextmanager
|
|
from typing import Iterator
|
|
|
|
from . import config
|
|
|
|
_SCHEMA = """
|
|
CREATE TABLE IF NOT EXISTS providers (
|
|
id TEXT PRIMARY KEY,
|
|
name TEXT NOT NULL,
|
|
adapter TEXT NOT NULL,
|
|
base_url TEXT NOT NULL DEFAULT '',
|
|
api_key TEXT NOT NULL DEFAULT '',
|
|
default_model TEXT NOT NULL DEFAULT '',
|
|
enabled INTEGER NOT NULL DEFAULT 1,
|
|
note TEXT NOT NULL DEFAULT '',
|
|
created_at TEXT NOT NULL,
|
|
updated_at TEXT NOT NULL
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS models (
|
|
id TEXT PRIMARY KEY,
|
|
provider_id TEXT NOT NULL,
|
|
model_id TEXT NOT NULL,
|
|
display_name TEXT NOT NULL DEFAULT '',
|
|
enabled INTEGER NOT NULL DEFAULT 0,
|
|
kind TEXT NOT NULL DEFAULT 'image',
|
|
size_hint TEXT NOT NULL DEFAULT '',
|
|
supports_image_input INTEGER NOT NULL DEFAULT 0,
|
|
note TEXT NOT NULL DEFAULT '',
|
|
discovered_at TEXT NOT NULL DEFAULT '',
|
|
created_at TEXT NOT NULL,
|
|
updated_at TEXT NOT NULL,
|
|
FOREIGN KEY (provider_id) REFERENCES providers(id) ON DELETE CASCADE
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS generations (
|
|
id TEXT PRIMARY KEY,
|
|
provider_id TEXT NOT NULL DEFAULT '',
|
|
provider_name TEXT NOT NULL DEFAULT '',
|
|
model TEXT NOT NULL DEFAULT '',
|
|
profile_name TEXT NOT NULL DEFAULT 'user',
|
|
prompt TEXT NOT NULL DEFAULT '',
|
|
size TEXT NOT NULL DEFAULT '',
|
|
ratio TEXT NOT NULL DEFAULT '',
|
|
n INTEGER NOT NULL DEFAULT 1,
|
|
status TEXT NOT NULL DEFAULT 'success',
|
|
error TEXT NOT NULL DEFAULT '',
|
|
files TEXT NOT NULL DEFAULT '[]',
|
|
duration_ms INTEGER NOT NULL DEFAULT 0,
|
|
created_at TEXT NOT NULL
|
|
);
|
|
|
|
CREATE INDEX IF NOT EXISTS idx_models_provider ON models(provider_id);
|
|
CREATE INDEX IF NOT EXISTS idx_generations_created ON generations(created_at DESC);
|
|
"""
|
|
|
|
|
|
def connect() -> sqlite3.Connection:
|
|
config.ensure_dirs()
|
|
conn = sqlite3.connect(config.DB_PATH, timeout=30)
|
|
conn.row_factory = sqlite3.Row
|
|
conn.execute("PRAGMA foreign_keys = ON")
|
|
return conn
|
|
|
|
|
|
@contextmanager
|
|
def session() -> Iterator[sqlite3.Connection]:
|
|
conn = connect()
|
|
try:
|
|
yield conn
|
|
conn.commit()
|
|
except Exception:
|
|
conn.rollback()
|
|
raise
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
def init_db() -> None:
|
|
with session() as conn:
|
|
conn.executescript(_SCHEMA)
|