mirror of
https://github.com/simon987/sist2.git
synced 2025-04-04 07:52:59 +00:00
137 lines
3.5 KiB
Python
137 lines
3.5 KiB
Python
from typing import Dict
|
|
import os
|
|
import shutil
|
|
|
|
from hexlib.db import Table, PersistentState
|
|
import pickle
|
|
|
|
from tesseract import get_tesseract_langs
|
|
import sqlite3
|
|
from config import LOG_FOLDER, logger
|
|
from sist2 import SearchBackendType, Sist2SearchBackend
|
|
|
|
RUNNING_FRONTENDS: Dict[str, int] = {}
|
|
|
|
TESSERACT_LANGS = get_tesseract_langs()
|
|
|
|
DB_SCHEMA_VERSION = "5"
|
|
|
|
from pydantic import BaseModel
|
|
|
|
|
|
def _serialize(item):
|
|
if isinstance(item, BaseModel):
|
|
return pickle.dumps(item)
|
|
if isinstance(item, bytes):
|
|
raise Exception("FIXME: bytes in PickleTable")
|
|
return item
|
|
|
|
|
|
def _deserialize(item):
|
|
if isinstance(item, bytes):
|
|
return pickle.loads(item)
|
|
return item
|
|
|
|
|
|
class PickleTable(Table):
|
|
|
|
def __getitem__(self, item):
|
|
row = super().__getitem__(item)
|
|
if row:
|
|
return dict((k, _deserialize(v)) for k, v in row.items())
|
|
return row
|
|
|
|
def __setitem__(self, key, value):
|
|
value = dict((k, _serialize(v)) for k, v in value.items())
|
|
super().__setitem__(key, value)
|
|
|
|
def __iter__(self):
|
|
for row in super().__iter__():
|
|
yield dict((k, _deserialize(v)) for k, v in row.items())
|
|
|
|
def sql(self, where_clause, *params):
|
|
for row in super().sql(where_clause, *params):
|
|
yield dict((k, _deserialize(v)) for k, v in row.items())
|
|
|
|
|
|
def get_log_files_to_remove(db: PersistentState, job_name: str, n: int):
|
|
if n < 0:
|
|
return []
|
|
|
|
counter = 0
|
|
to_remove = []
|
|
|
|
for row in db["task_done"].sql("WHERE has_logs=1 ORDER BY started DESC"):
|
|
if row["name"].endswith(f"[{job_name}]"):
|
|
counter += 1
|
|
|
|
if counter > n:
|
|
to_remove.append(row)
|
|
|
|
return to_remove
|
|
|
|
|
|
def delete_log_file(db: PersistentState, task_id: str):
|
|
db["task_done"][task_id] = {
|
|
"has_logs": 0
|
|
}
|
|
|
|
try:
|
|
os.remove(os.path.join(LOG_FOLDER, f"sist2-{task_id}.log"))
|
|
except:
|
|
pass
|
|
|
|
|
|
def migrate_v1_to_v2(db: PersistentState):
|
|
shutil.copy(db.dbfile, db.dbfile + "-before-migrate-v2.bak")
|
|
|
|
# Frontends
|
|
db._table_factory = PickleTable
|
|
frontends = [row["frontend"] for row in db["frontends"]]
|
|
del db["frontends"]
|
|
|
|
db._table_factory = Table
|
|
for frontend in frontends:
|
|
db["frontends"][frontend.name] = frontend
|
|
list(db["frontends"])
|
|
|
|
# Jobs
|
|
db._table_factory = PickleTable
|
|
jobs = [row["job"] for row in db["jobs"]]
|
|
del db["jobs"]
|
|
|
|
db._table_factory = Table
|
|
for job in jobs:
|
|
db["jobs"][job.name] = job
|
|
list(db["jobs"])
|
|
|
|
db["sist2_admin"]["info"] = {
|
|
"version": "2"
|
|
}
|
|
|
|
|
|
def create_default_search_backends(db: PersistentState):
|
|
es_backend = Sist2SearchBackend.create_default(name="elasticsearch",
|
|
backend_type=SearchBackendType("elasticsearch"))
|
|
db["search_backends"]["elasticsearch"] = es_backend
|
|
sqlite_backend = Sist2SearchBackend.create_default(name="sqlite", backend_type=SearchBackendType("sqlite"))
|
|
db["search_backends"]["sqlite"] = sqlite_backend
|
|
|
|
|
|
def migrate_v3_to_v4(db: PersistentState):
|
|
shutil.copy(db.dbfile, db.dbfile + "-before-migrate-v4.bak")
|
|
|
|
create_default_search_backends(db)
|
|
|
|
try:
|
|
conn = sqlite3.connect(db.dbfile)
|
|
conn.execute("ALTER TABLE task_done ADD COLUMN has_logs INTEGER DEFAULT 1")
|
|
conn.commit()
|
|
conn.close()
|
|
except Exception as e:
|
|
logger.exception(e)
|
|
|
|
db["sist2_admin"]["info"] = {
|
|
"version": "4"
|
|
}
|