f
This commit is contained in:
145
SYS/config.py
145
SYS/config.py
@@ -8,6 +8,7 @@ from pathlib import Path
|
||||
from typing import Any, Dict, Optional, List
|
||||
from SYS.logger import log
|
||||
from SYS.utils import expand_path
|
||||
from SYS.database import db, get_config_all, save_config_value
|
||||
|
||||
DEFAULT_CONFIG_FILENAME = "config.conf"
|
||||
SCRIPT_DIR = Path(__file__).resolve().parent
|
||||
@@ -208,10 +209,9 @@ def parse_conf_text(text: str, *, base: Optional[Dict[str, Any]] = None) -> Dict
|
||||
|
||||
Supported patterns:
|
||||
- Top-level key/value: temp="./temp"
|
||||
- Sections: [store=folder] + name/path lines
|
||||
- Sections: [store=hydrusnetwork] + name/access key/url lines
|
||||
- Sections: [provider=OpenLibrary] + email/password lines
|
||||
- Dotted keys: store.folder.default.path="C:\\Media" (optional)
|
||||
- Dotted keys: store.hydrusnetwork.<name>.url="http://..." (optional)
|
||||
"""
|
||||
config: Dict[str, Any] = dict(base or {})
|
||||
|
||||
@@ -519,7 +519,6 @@ def get_local_storage_path(config: Dict[str, Any]) -> Optional[Path]:
|
||||
"""Get local storage path from config.
|
||||
|
||||
Supports multiple formats:
|
||||
- New: config["store"]["folder"]["any_name"]["path"]
|
||||
- Old: config["storage"]["local"]["path"]
|
||||
- Old: config["Local"]["path"]
|
||||
|
||||
@@ -529,17 +528,6 @@ def get_local_storage_path(config: Dict[str, Any]) -> Optional[Path]:
|
||||
Returns:
|
||||
Path object if found, None otherwise
|
||||
"""
|
||||
# Try new format: iterate all folder stores and use the first valid path found.
|
||||
store = config.get("store", {})
|
||||
if isinstance(store, dict):
|
||||
folder_config = store.get("folder", {})
|
||||
if isinstance(folder_config, dict):
|
||||
for name, inst_cfg in folder_config.items():
|
||||
if isinstance(inst_cfg, dict):
|
||||
p = inst_cfg.get("path") or inst_cfg.get("PATH")
|
||||
if p:
|
||||
return expand_path(p)
|
||||
|
||||
# Fall back to storage.local.path format
|
||||
storage = config.get("storage", {})
|
||||
if isinstance(storage, dict):
|
||||
@@ -686,32 +674,76 @@ def resolve_debug_log(config: Dict[str, Any]) -> Optional[Path]:
|
||||
return path
|
||||
|
||||
|
||||
def migrate_conf_to_db(config: Dict[str, Any]) -> None:
|
||||
"""Migrate the configuration dictionary to the database."""
|
||||
log("Migrating configuration from .conf to database...")
|
||||
for key, value in config.items():
|
||||
if key in ("store", "provider", "tool", "networking"):
|
||||
cat = key
|
||||
sub_dict = value
|
||||
if isinstance(sub_dict, dict):
|
||||
for subtype, subtype_items in sub_dict.items():
|
||||
if isinstance(subtype_items, dict):
|
||||
# For provider/tool/networking, subtype is the name (e.g. alldebrid)
|
||||
# but for store, it's the type (e.g. hydrusnetwork)
|
||||
if cat == "store" and str(subtype).strip().lower() == "folder":
|
||||
continue
|
||||
if cat != "store":
|
||||
for k, v in subtype_items.items():
|
||||
save_config_value(cat, subtype, "", k, v)
|
||||
else:
|
||||
for name, items in subtype_items.items():
|
||||
if isinstance(items, dict):
|
||||
for k, v in items.items():
|
||||
save_config_value(cat, subtype, name, k, v)
|
||||
else:
|
||||
# Global setting
|
||||
save_config_value("global", "", "", key, value)
|
||||
log("Configuration migration complete!")
|
||||
|
||||
|
||||
def load_config(
|
||||
config_dir: Optional[Path] = None, filename: str = DEFAULT_CONFIG_FILENAME
|
||||
) -> Dict[str, Any]:
|
||||
base_dir = config_dir or SCRIPT_DIR
|
||||
config_path = base_dir / filename
|
||||
cache_key = _make_cache_key(config_dir, filename, config_path)
|
||||
|
||||
if cache_key in _CONFIG_CACHE:
|
||||
return _CONFIG_CACHE[cache_key]
|
||||
|
||||
if config_path.suffix.lower() != ".conf":
|
||||
log(f"Unsupported config format: {config_path.name} (only .conf is supported)")
|
||||
_CONFIG_CACHE[cache_key] = {}
|
||||
return {}
|
||||
# 1. Try loading from database first
|
||||
db_config = get_config_all()
|
||||
if db_config:
|
||||
_CONFIG_CACHE[cache_key] = db_config
|
||||
return db_config
|
||||
|
||||
try:
|
||||
data = _load_conf_config(base_dir, config_path)
|
||||
except FileNotFoundError:
|
||||
_CONFIG_CACHE[cache_key] = {}
|
||||
return {}
|
||||
except OSError as exc:
|
||||
log(f"Failed to read {config_path}: {exc}")
|
||||
_CONFIG_CACHE[cache_key] = {}
|
||||
return {}
|
||||
# 2. If DB is empty, try loading from legacy config.conf
|
||||
if config_path.exists():
|
||||
if config_path.suffix.lower() != ".conf":
|
||||
log(f"Unsupported config format: {config_path.name} (only .conf is supported)")
|
||||
return {}
|
||||
|
||||
try:
|
||||
config = _load_conf_config(base_dir, config_path)
|
||||
# Migrate to database
|
||||
migrate_conf_to_db(config)
|
||||
|
||||
# Optional: Rename old config file to mark as migrated
|
||||
try:
|
||||
migrated_path = config_path.with_name(config_path.name + ".migrated")
|
||||
config_path.rename(migrated_path)
|
||||
log(f"Legacy config file renamed to {migrated_path.name}")
|
||||
except Exception as e:
|
||||
log(f"Could not rename legacy config file: {e}")
|
||||
|
||||
_CONFIG_CACHE[cache_key] = data
|
||||
return data
|
||||
_CONFIG_CACHE[cache_key] = config
|
||||
return config
|
||||
except Exception as e:
|
||||
log(f"Failed to load legacy config at {config_path}: {e}")
|
||||
return {}
|
||||
|
||||
return {}
|
||||
|
||||
|
||||
def reload_config(
|
||||
@@ -723,55 +755,12 @@ def reload_config(
|
||||
|
||||
|
||||
def _validate_config_safety(config: Dict[str, Any]) -> None:
|
||||
"""Check for dangerous configurations, like folder stores in non-empty dirs."""
|
||||
store = config.get("store")
|
||||
if not isinstance(store, dict):
|
||||
return
|
||||
"""Validate configuration safety.
|
||||
|
||||
folder_stores = store.get("folder")
|
||||
if not isinstance(folder_stores, dict):
|
||||
return
|
||||
|
||||
for name, cfg in folder_stores.items():
|
||||
if not isinstance(cfg, dict):
|
||||
continue
|
||||
|
||||
path_str = cfg.get("path") or cfg.get("PATH")
|
||||
if not path_str:
|
||||
continue
|
||||
|
||||
try:
|
||||
p = expand_path(path_str).resolve()
|
||||
# If the path doesn't exist yet, it's fine (will be created empty)
|
||||
if not p.exists():
|
||||
continue
|
||||
|
||||
if not p.is_dir():
|
||||
continue
|
||||
|
||||
# DB name from API/folder.py
|
||||
db_file = p / "medios-macina.db"
|
||||
if db_file.exists():
|
||||
# Existing portable library, allowed to re-attach
|
||||
continue
|
||||
|
||||
# Check if directory has any files (other than the DB we just checked)
|
||||
items = list(p.iterdir())
|
||||
if items:
|
||||
item_names = [i.name for i in items[:3]]
|
||||
if len(items) > 3:
|
||||
item_names.append("...")
|
||||
raise RuntimeError(
|
||||
f"Configuration Error: Local library '{name}' target directory is not empty.\n"
|
||||
f"Path: {p}\n"
|
||||
f"Found {len(items)} items: {', '.join(item_names)}\n"
|
||||
f"To prevent accidental mass-hashing, new libraries must be set to unique, empty folders."
|
||||
)
|
||||
except RuntimeError:
|
||||
raise
|
||||
except Exception:
|
||||
# We don't want to crash on invalid paths during validation if they aren't 'unsafe'
|
||||
pass
|
||||
Folder store validation has been removed because the folder store backend
|
||||
is no longer supported.
|
||||
"""
|
||||
return
|
||||
|
||||
|
||||
def save_config(
|
||||
@@ -787,7 +776,7 @@ def save_config(
|
||||
f"Unsupported config format: {config_path.name} (only .conf is supported)"
|
||||
)
|
||||
|
||||
# Safety Check: Validate folder stores are in empty dirs or existing libraries
|
||||
# Safety Check: placeholder (folder store validation removed)
|
||||
_validate_config_safety(config)
|
||||
|
||||
try:
|
||||
|
||||
253
SYS/database.py
Normal file
253
SYS/database.py
Normal file
@@ -0,0 +1,253 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import sqlite3
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
# The database is located in the project root
|
||||
ROOT_DIR = Path(__file__).resolve().parent.parent
|
||||
DB_PATH = ROOT_DIR / "medios.db"
|
||||
|
||||
class Database:
|
||||
_instance: Optional[Database] = None
|
||||
|
||||
def __new__(cls):
|
||||
if cls._instance is None:
|
||||
cls._instance = super(Database, cls).__new__(cls)
|
||||
cls._instance._init_db()
|
||||
return cls._instance
|
||||
|
||||
def _init_db(self):
|
||||
self.conn = sqlite3.connect(str(DB_PATH), check_same_thread=False)
|
||||
self.conn.row_factory = sqlite3.Row
|
||||
self._create_tables()
|
||||
|
||||
def _create_tables(self):
|
||||
cursor = self.conn.cursor()
|
||||
|
||||
# Config table: stores all settings previously in config.conf
|
||||
# category: global, store, provider, tool, networking
|
||||
# subtype: e.g., hydrusnetwork, folder, alldebrid
|
||||
# item_name: e.g., hn-local, default
|
||||
# key: the setting key
|
||||
# value: the setting value (serialized to string)
|
||||
cursor.execute("""
|
||||
CREATE TABLE IF NOT EXISTS config (
|
||||
category TEXT,
|
||||
subtype TEXT,
|
||||
item_name TEXT,
|
||||
key TEXT,
|
||||
value TEXT,
|
||||
PRIMARY KEY (category, subtype, item_name, key)
|
||||
)
|
||||
""")
|
||||
|
||||
# Logs table
|
||||
cursor.execute("""
|
||||
CREATE TABLE IF NOT EXISTS logs (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
timestamp DATETIME DEFAULT CURRENT_TIMESTAMP,
|
||||
level TEXT,
|
||||
module TEXT,
|
||||
message TEXT
|
||||
)
|
||||
""")
|
||||
|
||||
# Workers table (for background tasks)
|
||||
cursor.execute("""
|
||||
CREATE TABLE IF NOT EXISTS workers (
|
||||
id TEXT PRIMARY KEY,
|
||||
type TEXT,
|
||||
title TEXT,
|
||||
description TEXT,
|
||||
status TEXT DEFAULT 'running',
|
||||
progress REAL DEFAULT 0.0,
|
||||
details TEXT,
|
||||
result TEXT DEFAULT 'pending',
|
||||
error_message TEXT,
|
||||
created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
|
||||
updated_at DATETIME DEFAULT CURRENT_TIMESTAMP
|
||||
)
|
||||
""")
|
||||
|
||||
# Worker stdout/logs
|
||||
cursor.execute("""
|
||||
CREATE TABLE IF NOT EXISTS worker_stdout (
|
||||
worker_id TEXT,
|
||||
channel TEXT DEFAULT 'stdout',
|
||||
content TEXT,
|
||||
timestamp DATETIME DEFAULT CURRENT_TIMESTAMP,
|
||||
FOREIGN KEY(worker_id) REFERENCES workers(id)
|
||||
)
|
||||
""")
|
||||
|
||||
self.conn.commit()
|
||||
|
||||
def get_connection(self):
|
||||
return self.conn
|
||||
|
||||
def execute(self, query: str, params: tuple = ()):
|
||||
cursor = self.conn.cursor()
|
||||
cursor.execute(query, params)
|
||||
self.conn.commit()
|
||||
return cursor
|
||||
|
||||
def fetchall(self, query: str, params: tuple = ()):
|
||||
cursor = self.conn.cursor()
|
||||
cursor.execute(query, params)
|
||||
return cursor.fetchall()
|
||||
|
||||
def fetchone(self, query: str, params: tuple = ()):
|
||||
cursor = self.conn.cursor()
|
||||
cursor.execute(query, params)
|
||||
return cursor.fetchone()
|
||||
|
||||
# Singleton instance
|
||||
db = Database()
|
||||
|
||||
def get_db() -> Database:
|
||||
return db
|
||||
|
||||
def log_to_db(level: str, module: str, message: str):
|
||||
"""Log a message to the database."""
|
||||
try:
|
||||
db.execute(
|
||||
"INSERT INTO logs (level, module, message) VALUES (?, ?, ?)",
|
||||
(level, module, message)
|
||||
)
|
||||
except Exception:
|
||||
# Avoid recursive logging errors if DB is locked
|
||||
pass
|
||||
|
||||
# Initialize DB logger in the unified logger
|
||||
try:
|
||||
from SYS.logger import set_db_logger
|
||||
set_db_logger(log_to_db)
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
def save_config_value(category: str, subtype: str, item_name: str, key: str, value: Any):
|
||||
"""Save a configuration value to the database."""
|
||||
val_str = json.dumps(value) if not isinstance(value, str) else value
|
||||
db.execute(
|
||||
"INSERT OR REPLACE INTO config (category, subtype, item_name, key, value) VALUES (?, ?, ?, ?, ?)",
|
||||
(category, subtype, item_name, key, val_str)
|
||||
)
|
||||
|
||||
def get_config_all() -> Dict[str, Any]:
|
||||
"""Retrieve all configuration from the database in the legacy dict format."""
|
||||
try:
|
||||
db.execute("DELETE FROM config WHERE category='store' AND LOWER(subtype)='folder'")
|
||||
except Exception:
|
||||
pass
|
||||
rows = db.fetchall("SELECT category, subtype, item_name, key, value FROM config")
|
||||
config: Dict[str, Any] = {}
|
||||
|
||||
for row in rows:
|
||||
cat = row['category']
|
||||
sub = row['subtype']
|
||||
name = row['item_name']
|
||||
key = row['key']
|
||||
val = row['value']
|
||||
|
||||
# Drop legacy folder store entries (folder store is removed).
|
||||
if cat == 'store' and str(sub).strip().lower() == 'folder':
|
||||
continue
|
||||
|
||||
# Try to parse JSON value, fallback to string
|
||||
try:
|
||||
parsed_val = json.loads(val)
|
||||
except Exception:
|
||||
parsed_val = val
|
||||
|
||||
if cat == 'global':
|
||||
config[key] = parsed_val
|
||||
else:
|
||||
# Modular structure: config[cat][sub][name][key]
|
||||
if cat in ('provider', 'tool', 'networking'):
|
||||
cat_dict = config.setdefault(cat, {})
|
||||
sub_dict = cat_dict.setdefault(sub, {})
|
||||
sub_dict[key] = parsed_val
|
||||
elif cat == 'store':
|
||||
cat_dict = config.setdefault(cat, {})
|
||||
sub_dict = cat_dict.setdefault(sub, {})
|
||||
name_dict = sub_dict.setdefault(name, {})
|
||||
name_dict[key] = parsed_val
|
||||
else:
|
||||
config.setdefault(cat, {})[key] = parsed_val
|
||||
|
||||
return config
|
||||
|
||||
# Worker Management Methods for medios.db
|
||||
|
||||
def insert_worker(worker_id: str, worker_type: str, title: str = "", description: str = "") -> bool:
|
||||
try:
|
||||
db.execute(
|
||||
"INSERT INTO workers (id, type, title, description, status) VALUES (?, ?, ?, ?, 'running')",
|
||||
(worker_id, worker_type, title, description)
|
||||
)
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
def update_worker(worker_id: str, **kwargs) -> bool:
|
||||
if not kwargs:
|
||||
return True
|
||||
|
||||
# Filter valid columns
|
||||
valid_cols = {'type', 'title', 'description', 'status', 'progress', 'details', 'result', 'error_message'}
|
||||
cols = []
|
||||
vals = []
|
||||
for k, v in kwargs.items():
|
||||
if k in valid_cols:
|
||||
cols.append(f"{k} = ?")
|
||||
vals.append(v)
|
||||
|
||||
if not cols:
|
||||
return True
|
||||
|
||||
cols.append("updated_at = CURRENT_TIMESTAMP")
|
||||
query = f"UPDATE workers SET {', '.join(cols)} WHERE id = ?"
|
||||
vals.append(worker_id)
|
||||
|
||||
try:
|
||||
db.execute(query, tuple(vals))
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
def append_worker_stdout(worker_id: str, content: str, channel: str = 'stdout'):
|
||||
try:
|
||||
db.execute(
|
||||
"INSERT INTO worker_stdout (worker_id, channel, content) VALUES (?, ?, ?)",
|
||||
(worker_id, channel, content)
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def get_worker_stdout(worker_id: str, channel: Optional[str] = None) -> str:
|
||||
query = "SELECT content FROM worker_stdout WHERE worker_id = ?"
|
||||
params = [worker_id]
|
||||
if channel:
|
||||
query += " AND channel = ?"
|
||||
params.append(channel)
|
||||
query += " ORDER BY timestamp ASC"
|
||||
|
||||
rows = db.fetchall(query, tuple(params))
|
||||
return "\n".join(row['content'] for row in rows)
|
||||
|
||||
def get_active_workers() -> List[Dict[str, Any]]:
|
||||
rows = db.fetchall("SELECT * FROM workers WHERE status = 'running' ORDER BY created_at DESC")
|
||||
return [dict(row) for row in rows]
|
||||
|
||||
def get_worker(worker_id: str) -> Optional[Dict[str, Any]]:
|
||||
row = db.fetchone("SELECT * FROM workers WHERE id = ?", (worker_id,))
|
||||
return dict(row) if row else None
|
||||
|
||||
def expire_running_workers(older_than_seconds: int = 300, status: str = 'error', reason: str = 'timeout') -> int:
|
||||
# SQLITE doesn't have a simple way to do DATETIME - INTERVAL, so we'll use strftime/unixepoch if available
|
||||
# or just do regular update for all running ones for now as a simple fallback
|
||||
query = f"UPDATE workers SET status = ?, error_message = ? WHERE status = 'running'"
|
||||
db.execute(query, (status, reason))
|
||||
return 0 # We don't easily get the rowcount from db.execute right now
|
||||
@@ -4,9 +4,17 @@ import sys
|
||||
import inspect
|
||||
import threading
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
from SYS.rich_display import console_for
|
||||
|
||||
# Global DB logger set later to avoid circular imports
|
||||
_DB_LOGGER = None
|
||||
|
||||
def set_db_logger(func):
|
||||
global _DB_LOGGER
|
||||
_DB_LOGGER = func
|
||||
|
||||
_DEBUG_ENABLED = False
|
||||
_thread_local = threading.local()
|
||||
|
||||
@@ -207,6 +215,15 @@ def log(*args, **kwargs) -> None:
|
||||
console_for(file).print(prefix, *args, sep=sep, end=end)
|
||||
else:
|
||||
console_for(file).print(*args, sep=sep, end=end)
|
||||
|
||||
# Log to database if available
|
||||
if _DB_LOGGER:
|
||||
try:
|
||||
msg = sep.join(map(str, args))
|
||||
level = "DEBUG" if add_prefix else "INFO"
|
||||
_DB_LOGGER(level, f"{file_name}.{func_name}", msg)
|
||||
except Exception:
|
||||
pass
|
||||
finally:
|
||||
del frame
|
||||
del caller_frame
|
||||
|
||||
@@ -11,8 +11,17 @@ from datetime import datetime
|
||||
from threading import Thread, Lock
|
||||
import time
|
||||
|
||||
from API.folder import API_folder_store
|
||||
from SYS.logger import log
|
||||
from SYS.database import (
|
||||
db,
|
||||
insert_worker,
|
||||
update_worker,
|
||||
append_worker_stdout,
|
||||
get_worker_stdout as db_get_worker_stdout,
|
||||
get_active_workers as db_get_active_workers,
|
||||
get_worker as db_get_worker,
|
||||
expire_running_workers as db_expire_running_workers
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -157,7 +166,6 @@ class WorkerLoggingHandler(logging.StreamHandler):
|
||||
def __init__(
|
||||
self,
|
||||
worker_id: str,
|
||||
db: API_folder_store,
|
||||
manager: Optional["WorkerManager"] = None,
|
||||
buffer_size: int = 50,
|
||||
):
|
||||
@@ -165,12 +173,10 @@ class WorkerLoggingHandler(logging.StreamHandler):
|
||||
|
||||
Args:
|
||||
worker_id: ID of the worker to capture logs for
|
||||
db: Reference to LocalLibraryDB for storing logs
|
||||
buffer_size: Number of logs to buffer before flushing to DB
|
||||
"""
|
||||
super().__init__()
|
||||
self.worker_id = worker_id
|
||||
self.db = db
|
||||
self.manager = manager
|
||||
self.buffer_size = buffer_size
|
||||
self.buffer: list[str] = []
|
||||
@@ -232,7 +238,7 @@ class WorkerLoggingHandler(logging.StreamHandler):
|
||||
channel="log"
|
||||
)
|
||||
else:
|
||||
self.db.append_worker_stdout(
|
||||
append_worker_stdout(
|
||||
self.worker_id,
|
||||
log_text,
|
||||
channel="log"
|
||||
@@ -255,29 +261,22 @@ class WorkerLoggingHandler(logging.StreamHandler):
|
||||
|
||||
|
||||
class WorkerManager:
|
||||
"""Manages persistent worker tasks with auto-refresh capability."""
|
||||
"""Manages persistent worker tasks using the central medios.db."""
|
||||
|
||||
def __init__(self, library_root: Path, auto_refresh_interval: float = 2.0):
|
||||
def __init__(self, auto_refresh_interval: float = 2.0):
|
||||
"""Initialize the worker manager.
|
||||
|
||||
Args:
|
||||
library_root: Root directory for the local library database
|
||||
auto_refresh_interval: Seconds between auto-refresh checks (0 = disabled)
|
||||
"""
|
||||
self.library_root = Path(library_root)
|
||||
self.db = API_folder_store(library_root)
|
||||
self.auto_refresh_interval = auto_refresh_interval
|
||||
self.refresh_callbacks: List[Callable] = []
|
||||
self.refresh_thread: Optional[Thread] = None
|
||||
self._stop_refresh = False
|
||||
self._lock = Lock()
|
||||
# Reuse the DB's own lock so there is exactly one lock guarding the
|
||||
# sqlite connection (and it is safe for re-entrant/nested DB usage).
|
||||
self._db_lock = self.db._db_lock
|
||||
self.worker_handlers: Dict[str,
|
||||
WorkerLoggingHandler] = {} # Track active handlers
|
||||
self._worker_last_step: Dict[str,
|
||||
str] = {}
|
||||
self.worker_handlers: Dict[str, WorkerLoggingHandler] = {}
|
||||
self._worker_last_step: Dict[str, str] = {}
|
||||
|
||||
# Buffered stdout/log batching to reduce DB lock contention.
|
||||
self._stdout_buffers: Dict[Tuple[str, str], List[str]] = {}
|
||||
self._stdout_buffer_sizes: Dict[Tuple[str, str], int] = {}
|
||||
@@ -328,13 +327,11 @@ class WorkerManager:
|
||||
Count of workers updated.
|
||||
"""
|
||||
try:
|
||||
with self._db_lock:
|
||||
return self.db.expire_running_workers(
|
||||
older_than_seconds=older_than_seconds,
|
||||
status=status,
|
||||
reason=reason,
|
||||
worker_id_prefix=worker_id_prefix,
|
||||
)
|
||||
return db_expire_running_workers(
|
||||
older_than_seconds=older_than_seconds,
|
||||
status=status,
|
||||
reason=reason or "Stale worker expired"
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.error(f"Failed to expire stale workers: {exc}", exc_info=True)
|
||||
return 0
|
||||
@@ -362,7 +359,7 @@ class WorkerManager:
|
||||
The logging handler that was created, or None if there was an error
|
||||
"""
|
||||
try:
|
||||
handler = WorkerLoggingHandler(worker_id, self.db, manager=self)
|
||||
handler = WorkerLoggingHandler(worker_id, manager=self)
|
||||
with self._lock:
|
||||
self.worker_handlers[worker_id] = handler
|
||||
|
||||
@@ -437,16 +434,13 @@ class WorkerManager:
|
||||
True if worker was inserted successfully
|
||||
"""
|
||||
try:
|
||||
with self._db_lock:
|
||||
result = self.db.insert_worker(
|
||||
worker_id,
|
||||
worker_type,
|
||||
title,
|
||||
description,
|
||||
total_steps,
|
||||
pipe=pipe
|
||||
)
|
||||
if result > 0:
|
||||
success = insert_worker(
|
||||
worker_id,
|
||||
worker_type,
|
||||
title,
|
||||
description
|
||||
)
|
||||
if success:
|
||||
logger.debug(
|
||||
f"[WorkerManager] Tracking worker: {worker_id} ({worker_type})"
|
||||
)
|
||||
@@ -482,18 +476,16 @@ class WorkerManager:
|
||||
if progress > 0:
|
||||
kwargs["progress"] = progress
|
||||
if current_step:
|
||||
kwargs["current_step"] = current_step
|
||||
kwargs["details"] = current_step
|
||||
if details:
|
||||
kwargs["description"] = details
|
||||
if error:
|
||||
kwargs["error_message"] = error
|
||||
|
||||
if kwargs:
|
||||
kwargs["last_updated"] = datetime.now().isoformat()
|
||||
if "current_step" in kwargs and kwargs["current_step"]:
|
||||
self._worker_last_step[worker_id] = str(kwargs["current_step"])
|
||||
with self._db_lock:
|
||||
return self.db.update_worker(worker_id, **kwargs)
|
||||
if "details" in kwargs and kwargs["details"]:
|
||||
self._worker_last_step[worker_id] = str(kwargs["details"])
|
||||
return update_worker(worker_id, **kwargs)
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
@@ -515,7 +507,7 @@ class WorkerManager:
|
||||
worker_id: Unique identifier for the worker
|
||||
result: Result status ('completed', 'error', 'cancelled')
|
||||
error_msg: Error message if any
|
||||
result_data: Result data as JSON string
|
||||
result_data: Result data as JSON string (saved in details)
|
||||
|
||||
Returns:
|
||||
True if update was successful
|
||||
@@ -526,16 +518,15 @@ class WorkerManager:
|
||||
except Exception:
|
||||
pass
|
||||
kwargs = {
|
||||
"status": result,
|
||||
"completed_at": datetime.now().isoformat()
|
||||
"status": "finished",
|
||||
"result": result
|
||||
}
|
||||
if error_msg:
|
||||
kwargs["error_message"] = error_msg
|
||||
if result_data:
|
||||
kwargs["result_data"] = result_data
|
||||
kwargs["details"] = result_data
|
||||
|
||||
with self._db_lock:
|
||||
success = self.db.update_worker(worker_id, **kwargs)
|
||||
success = update_worker(worker_id, **kwargs)
|
||||
logger.info(f"[WorkerManager] Worker finished: {worker_id} ({result})")
|
||||
self._worker_last_step.pop(worker_id, None)
|
||||
return success
|
||||
@@ -553,8 +544,7 @@ class WorkerManager:
|
||||
List of active worker dictionaries
|
||||
"""
|
||||
try:
|
||||
with self._db_lock:
|
||||
return self.db.get_active_workers()
|
||||
return db_get_active_workers()
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"[WorkerManager] Error getting active workers: {e}",
|
||||
@@ -572,14 +562,9 @@ class WorkerManager:
|
||||
List of finished worker dictionaries
|
||||
"""
|
||||
try:
|
||||
with self._db_lock:
|
||||
all_workers = self.db.get_all_workers(limit=limit)
|
||||
# Filter to only finished workers
|
||||
finished = [
|
||||
w for w in all_workers
|
||||
if w.get("status") in ["completed", "error", "cancelled"]
|
||||
]
|
||||
return finished
|
||||
# We don't have a get_all_workers in database.py yet, but we'll use a local query
|
||||
rows = db.fetchall(f"SELECT * FROM workers WHERE status = 'finished' ORDER BY updated_at DESC LIMIT {limit}")
|
||||
return [dict(row) for row in rows]
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"[WorkerManager] Error getting finished workers: {e}",
|
||||
@@ -597,7 +582,13 @@ class WorkerManager:
|
||||
Worker data or None if not found
|
||||
"""
|
||||
try:
|
||||
with self._db_lock:
|
||||
return db_get_worker(worker_id)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"[WorkerManager] Error getting worker {worker_id}: {e}",
|
||||
exc_info=True
|
||||
)
|
||||
return None
|
||||
return self.db.get_worker(worker_id)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
@@ -815,14 +806,11 @@ class WorkerManager:
|
||||
ok = True
|
||||
for wid, ch, step, payload in pending_flush:
|
||||
try:
|
||||
with self._db_lock:
|
||||
result = self.db.append_worker_stdout(
|
||||
wid,
|
||||
payload,
|
||||
step=step,
|
||||
channel=ch
|
||||
)
|
||||
ok = ok and result
|
||||
append_worker_stdout(
|
||||
wid,
|
||||
payload,
|
||||
channel=ch
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"[WorkerManager] Error flushing stdout for {wid}: {e}",
|
||||
@@ -851,7 +839,6 @@ class WorkerManager:
|
||||
if not chunks:
|
||||
return True
|
||||
text = "".join(chunks)
|
||||
step = self._stdout_buffer_steps.get(key)
|
||||
self._stdout_buffers[key] = []
|
||||
self._stdout_buffer_sizes[key] = 0
|
||||
self._stdout_last_flush[key] = time.monotonic()
|
||||
@@ -860,13 +847,12 @@ class WorkerManager:
|
||||
if not text:
|
||||
return True
|
||||
try:
|
||||
with self._db_lock:
|
||||
return self.db.append_worker_stdout(
|
||||
worker_id,
|
||||
text,
|
||||
step=step,
|
||||
channel=channel,
|
||||
)
|
||||
append_worker_stdout(
|
||||
worker_id,
|
||||
text,
|
||||
channel=channel,
|
||||
)
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"[WorkerManager] Error flushing stdout for {worker_id}: {e}",
|
||||
@@ -884,8 +870,7 @@ class WorkerManager:
|
||||
Worker's stdout or empty string
|
||||
"""
|
||||
try:
|
||||
with self._db_lock:
|
||||
return self.db.get_worker_stdout(worker_id)
|
||||
return db_get_worker_stdout(worker_id)
|
||||
except Exception as e:
|
||||
logger.error(f"[WorkerManager] Error getting stdout: {e}", exc_info=True)
|
||||
return ""
|
||||
@@ -909,21 +894,20 @@ class WorkerManager:
|
||||
True if clear was successful
|
||||
"""
|
||||
try:
|
||||
with self._db_lock:
|
||||
return self.db.clear_worker_stdout(worker_id)
|
||||
# Not implemented in database.py yet, but we'll add it or skip it
|
||||
db.execute("DELETE FROM worker_stdout WHERE worker_id = ?", (worker_id,))
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.error(f"[WorkerManager] Error clearing stdout: {e}", exc_info=True)
|
||||
return False
|
||||
|
||||
def close(self) -> None:
|
||||
"""Close the worker manager and database connection."""
|
||||
"""Close the worker manager."""
|
||||
self.stop_auto_refresh()
|
||||
try:
|
||||
self._flush_all_stdout_buffers()
|
||||
except Exception:
|
||||
pass
|
||||
with self._db_lock:
|
||||
self.db.close()
|
||||
logger.info("[WorkerManager] Closed")
|
||||
|
||||
def _flush_all_stdout_buffers(self) -> None:
|
||||
|
||||
Reference in New Issue
Block a user