341 lines
16 KiB
Python
341 lines
16 KiB
Python
from __future__ import annotations
|
|
|
|
import sqlite3
|
|
import threading
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
|
|
ACTIVE_QUEUE_STATES = {
|
|
"detected",
|
|
"waiting_for_sab",
|
|
"ready",
|
|
"importing",
|
|
"retrying",
|
|
}
|
|
|
|
TERMINAL_QUEUE_STATES = {"imported", "failed", "skipped"}
|
|
|
|
|
|
class State:
|
|
def __init__(self, path: Path):
|
|
self.path = path
|
|
self.path.parent.mkdir(parents=True, exist_ok=True)
|
|
self._lock = threading.RLock()
|
|
self.conn = sqlite3.connect(self.path, check_same_thread=False)
|
|
self.conn.row_factory = sqlite3.Row
|
|
self.conn.execute("pragma journal_mode=WAL")
|
|
self.conn.execute("pragma busy_timeout = 5000")
|
|
self.migrate()
|
|
|
|
def migrate(self) -> None:
|
|
with self._lock:
|
|
self.conn.executescript(
|
|
"""
|
|
create table if not exists manual_batches (
|
|
id integer primary key autoincrement,
|
|
path text not null unique,
|
|
status text not null default 'active',
|
|
created_at text not null default current_timestamp,
|
|
completed_at text
|
|
);
|
|
create table if not exists import_history (
|
|
id integer primary key autoincrement,
|
|
source text not null,
|
|
target text not null,
|
|
status text not null,
|
|
bytes integer not null default 0,
|
|
created_at text not null default current_timestamp,
|
|
completed_at text,
|
|
error text
|
|
);
|
|
create table if not exists app_state (key text primary key, value text not null);
|
|
create table if not exists import_queue_items (
|
|
id integer primary key autoincrement,
|
|
source_type text not null,
|
|
source_id text not null,
|
|
source_path text,
|
|
name text not null,
|
|
state text not null,
|
|
reason text,
|
|
relative_path text,
|
|
size integer not null default 0,
|
|
batch_id integer,
|
|
job_id text,
|
|
sab_category text,
|
|
attempt_count integer not null default 0,
|
|
next_retry_at text,
|
|
last_error text,
|
|
claimed_by text,
|
|
claimed_at text,
|
|
first_seen_at text not null default current_timestamp,
|
|
updated_at text not null default current_timestamp,
|
|
completed_at text,
|
|
unique(source_type, source_id)
|
|
);
|
|
"""
|
|
)
|
|
columns = {row["name"] for row in self.conn.execute("pragma table_info(import_queue_items)")}
|
|
if "sab_category" not in columns:
|
|
self.conn.execute("alter table import_queue_items add column sab_category text")
|
|
if "attempt_count" not in columns:
|
|
self.conn.execute("alter table import_queue_items add column attempt_count integer not null default 0")
|
|
if "next_retry_at" not in columns:
|
|
self.conn.execute("alter table import_queue_items add column next_retry_at text")
|
|
if "last_error" not in columns:
|
|
self.conn.execute("alter table import_queue_items add column last_error text")
|
|
if "claimed_by" not in columns:
|
|
self.conn.execute("alter table import_queue_items add column claimed_by text")
|
|
if "claimed_at" not in columns:
|
|
self.conn.execute("alter table import_queue_items add column claimed_at text")
|
|
self.conn.commit()
|
|
|
|
def get_app_state(self, key: str, default: str | None = None) -> str | None:
|
|
with self._lock:
|
|
row = self.conn.execute("select value from app_state where key = ?", (key,)).fetchone()
|
|
return row["value"] if row else default
|
|
|
|
def set_app_state(self, key: str, value: str) -> None:
|
|
with self._lock:
|
|
self.conn.execute(
|
|
"insert into app_state(key, value) values (?, ?) on conflict(key) do update set value=excluded.value",
|
|
(key, value),
|
|
)
|
|
self.conn.commit()
|
|
|
|
def add_manual_batch(self, path: Path) -> dict[str, Any]:
|
|
with self._lock:
|
|
self.conn.execute("insert or ignore into manual_batches(path) values (?)", (str(path),))
|
|
self.conn.commit()
|
|
return self.get_manual_batch_by_path(path)
|
|
|
|
def get_manual_batch_by_path(self, path: Path) -> dict[str, Any]:
|
|
with self._lock:
|
|
row = self.conn.execute("select * from manual_batches where path = ?", (str(path),)).fetchone()
|
|
return dict(row)
|
|
|
|
def list_manual_batches(self, active_only: bool = False) -> list[dict[str, Any]]:
|
|
sql = "select * from manual_batches"
|
|
if active_only:
|
|
sql += " where status = 'active'"
|
|
sql += " order by created_at desc"
|
|
with self._lock:
|
|
return [dict(row) for row in self.conn.execute(sql)]
|
|
|
|
def delete_manual_batch(self, batch_id: int) -> None:
|
|
with self._lock:
|
|
self.conn.execute("delete from manual_batches where id = ?", (batch_id,))
|
|
self.conn.execute("delete from import_queue_items where batch_id = ? and source_type = 'manual'", (batch_id,))
|
|
self.conn.commit()
|
|
|
|
def complete_manual_batch(self, batch_id: int) -> None:
|
|
with self._lock:
|
|
self.conn.execute("update manual_batches set status='completed', completed_at=current_timestamp where id=?", (batch_id,))
|
|
self.conn.commit()
|
|
|
|
def add_history(self, source: Path, target: Path, status: str, bytes_count: int = 0, error: str | None = None) -> None:
|
|
with self._lock:
|
|
self.conn.execute(
|
|
"insert into import_history(source,target,status,bytes,error,completed_at) values (?,?,?,?,?,case when ? in ('imported','failed') then current_timestamp else null end)",
|
|
(str(source), str(target), status, bytes_count, error, status),
|
|
)
|
|
self.conn.commit()
|
|
|
|
def upsert_queue_item(
|
|
self,
|
|
*,
|
|
source_type: str,
|
|
source_id: str,
|
|
name: str,
|
|
state: str,
|
|
source_path: Path | None = None,
|
|
reason: str | None = None,
|
|
relative_path: str | None = None,
|
|
size: int = 0,
|
|
batch_id: int | None = None,
|
|
job_id: str | None = None,
|
|
sab_category: str | None = None,
|
|
preserve_finished_state: bool = True,
|
|
) -> dict[str, Any]:
|
|
with self._lock:
|
|
self.conn.execute(
|
|
"""
|
|
insert into import_queue_items(source_type, source_id, source_path, name, state, reason, relative_path, size, batch_id, job_id, sab_category)
|
|
values (?,?,?,?,?,?,?,?,?,?,?)
|
|
on conflict(source_type, source_id) do update set
|
|
source_path=excluded.source_path,
|
|
name=excluded.name,
|
|
state=case
|
|
when ? and import_queue_items.state in ('imported','failed','skipped','importing','retrying') then import_queue_items.state
|
|
else excluded.state
|
|
end,
|
|
reason=case
|
|
when ? and import_queue_items.state in ('imported','failed','skipped','importing','retrying') then coalesce(import_queue_items.reason, excluded.reason)
|
|
else excluded.reason
|
|
end,
|
|
relative_path=excluded.relative_path,
|
|
size=excluded.size,
|
|
batch_id=excluded.batch_id,
|
|
job_id=excluded.job_id,
|
|
sab_category=excluded.sab_category,
|
|
updated_at=current_timestamp,
|
|
next_retry_at=case
|
|
when excluded.state = 'retrying' then coalesce(import_queue_items.next_retry_at, excluded.next_retry_at)
|
|
when excluded.state = 'ready' then null
|
|
else import_queue_items.next_retry_at
|
|
end,
|
|
completed_at=case
|
|
when ? and import_queue_items.state in ('imported','failed','skipped') then import_queue_items.completed_at
|
|
when excluded.state in ('imported','failed','skipped') then current_timestamp
|
|
else null
|
|
end
|
|
""",
|
|
(source_type, source_id, str(source_path) if source_path else None, name, state, reason, relative_path, size, batch_id, job_id, sab_category, preserve_finished_state, preserve_finished_state, preserve_finished_state),
|
|
)
|
|
self.conn.commit()
|
|
row = self.conn.execute("select * from import_queue_items where source_type = ? and source_id = ?", (source_type, source_id)).fetchone()
|
|
return dict(row)
|
|
|
|
def mark_queue_item(self, source_type: str, source_id: str, state: str, reason: str | None = None) -> None:
|
|
with self._lock:
|
|
self.conn.execute(
|
|
"update import_queue_items set state=?, reason=?, updated_at=current_timestamp, claimed_by=null, claimed_at=null, next_retry_at=case when ?='ready' then null else next_retry_at end, completed_at=case when ? in ('imported','failed','skipped') then current_timestamp else completed_at end where source_type=? and source_id=?",
|
|
(state, reason, state, state, source_type, source_id),
|
|
)
|
|
self.conn.commit()
|
|
|
|
def claim_next_queue_item(self, worker_id: str) -> dict[str, Any] | None:
|
|
with self._lock:
|
|
row = self.conn.execute(
|
|
"""
|
|
update import_queue_items
|
|
set state='importing', claimed_by=?, claimed_at=current_timestamp, updated_at=current_timestamp
|
|
where id = (
|
|
select id from import_queue_items
|
|
where claimed_by is null
|
|
and (state = 'ready' or (state = 'retrying' and (next_retry_at is null or next_retry_at <= current_timestamp)))
|
|
order by case state when 'ready' then 0 else 1 end, updated_at asc, id asc
|
|
limit 1
|
|
) and claimed_by is null
|
|
returning *
|
|
""",
|
|
(worker_id,),
|
|
).fetchone()
|
|
self.conn.commit()
|
|
return dict(row) if row else None
|
|
|
|
def release_stale_claims(self) -> int:
|
|
with self._lock:
|
|
cursor = self.conn.execute(
|
|
"update import_queue_items set state='retrying', claimed_by=null, claimed_at=null, updated_at=current_timestamp where state='importing'"
|
|
)
|
|
self.conn.commit()
|
|
return cursor.rowcount
|
|
|
|
def claim_queue_item(self, item_id: int, worker_id: str, allowed_states: set[str]) -> dict[str, Any] | None:
|
|
placeholders = ",".join("?" for _ in allowed_states)
|
|
with self._lock:
|
|
row = self.conn.execute(
|
|
f"update import_queue_items set state='importing', claimed_by=?, claimed_at=current_timestamp, updated_at=current_timestamp where id=? and state in ({placeholders}) and claimed_by is null returning *",
|
|
(worker_id, item_id, *allowed_states),
|
|
).fetchone()
|
|
self.conn.commit()
|
|
return dict(row) if row else None
|
|
|
|
def transition_queue_item_if_unclaimed(self, item_id: int, allowed_states: set[str], new_state: str, reason: str) -> bool:
|
|
placeholders = ",".join("?" for _ in allowed_states)
|
|
with self._lock:
|
|
cursor = self.conn.execute(
|
|
f"update import_queue_items set state=?, reason=?, updated_at=current_timestamp, next_retry_at=null, completed_at=case when ? in ('failed','skipped') then current_timestamp else null end where id=? and state in ({placeholders}) and claimed_by is null",
|
|
(new_state, reason, new_state, item_id, *allowed_states),
|
|
)
|
|
self.conn.commit()
|
|
return cursor.rowcount > 0
|
|
|
|
def mark_queue_item_result(
|
|
self,
|
|
source_type: str,
|
|
source_id: str,
|
|
state: str,
|
|
reason: str | None = None,
|
|
*,
|
|
increment_attempts: bool = False,
|
|
next_retry_seconds: int | None = None,
|
|
) -> None:
|
|
with self._lock:
|
|
self.conn.execute(
|
|
"""
|
|
update import_queue_items
|
|
set state=?,
|
|
reason=?,
|
|
last_error=case when ? in ('failed','retrying','skipped') then ? else null end,
|
|
attempt_count=attempt_count + ?,
|
|
next_retry_at=case
|
|
when ? = 'retrying' and ? is not null then datetime('now', '+' || ? || ' seconds')
|
|
when ? in ('ready','imported','failed','skipped') then null
|
|
else next_retry_at
|
|
end,
|
|
claimed_by=null,
|
|
claimed_at=null,
|
|
updated_at=current_timestamp,
|
|
completed_at=case when ? in ('imported','failed','skipped') then current_timestamp else null end
|
|
where source_type=? and source_id=?
|
|
""",
|
|
(state, reason, state, reason, 1 if increment_attempts else 0, state, next_retry_seconds, next_retry_seconds, state, state, source_type, source_id),
|
|
)
|
|
self.conn.commit()
|
|
|
|
def delete_queue_item(self, item_id: int) -> bool:
|
|
with self._lock:
|
|
cursor = self.conn.execute("delete from import_queue_items where id = ?", (item_id,))
|
|
self.conn.commit()
|
|
return cursor.rowcount > 0
|
|
|
|
def delete_queue_item_if_unclaimed(self, item_id: int) -> bool:
|
|
with self._lock:
|
|
cursor = self.conn.execute("delete from import_queue_items where id = ? and claimed_by is null", (item_id,))
|
|
self.conn.commit()
|
|
return cursor.rowcount > 0
|
|
|
|
def delete_queue_items_by_state(self, source_type: str, state: str, reason: str | None = None) -> int:
|
|
with self._lock:
|
|
if reason is None:
|
|
cursor = self.conn.execute("delete from import_queue_items where source_type = ? and state = ?", (source_type, state))
|
|
else:
|
|
cursor = self.conn.execute("delete from import_queue_items where source_type = ? and state = ? and reason = ?", (source_type, state, reason))
|
|
self.conn.commit()
|
|
return cursor.rowcount
|
|
|
|
def get_queue_item(self, item_id: int) -> dict[str, Any] | None:
|
|
with self._lock:
|
|
row = self.conn.execute("select * from import_queue_items where id = ?", (item_id,)).fetchone()
|
|
return dict(row) if row else None
|
|
|
|
def remove_missing_manual_items(self, batch_id: int, source_ids: set[str]) -> None:
|
|
with self._lock:
|
|
rows = self.conn.execute("select source_id, state from import_queue_items where source_type='manual' and batch_id=?", (batch_id,)).fetchall()
|
|
for row in rows:
|
|
if row["source_id"] not in source_ids and row["state"] not in TERMINAL_QUEUE_STATES:
|
|
self.conn.execute("delete from import_queue_items where source_type='manual' and source_id=?", (row["source_id"],))
|
|
self.conn.commit()
|
|
|
|
def batch_has_active_items(self, batch_id: int) -> bool:
|
|
with self._lock:
|
|
row = self.conn.execute(
|
|
"select 1 from import_queue_items where batch_id = ? and source_type='manual' and state not in ('imported','failed','skipped') limit 1",
|
|
(batch_id,),
|
|
).fetchone()
|
|
return row is not None
|
|
|
|
def list_queue_items(self, active_only: bool = True) -> list[dict[str, Any]]:
|
|
sql = "select * from import_queue_items"
|
|
if active_only:
|
|
sql += " where state not in ('imported','failed','skipped')"
|
|
sql += " order by updated_at desc, id desc"
|
|
with self._lock:
|
|
return [dict(row) for row in self.conn.execute(sql)]
|
|
|
|
def list_history(self, limit: int = 100) -> list[dict[str, Any]]:
|
|
with self._lock:
|
|
return [dict(row) for row in self.conn.execute("select * from import_history order by id desc limit ?", (limit,))]
|