113 lines
4.6 KiB
Python
113 lines
4.6 KiB
Python
"""Database connection management and schema initialisation."""
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import sqlite3
|
|
from contextlib import asynccontextmanager
|
|
from pathlib import Path
|
|
from typing import AsyncGenerator
|
|
|
|
import aiosqlite
|
|
|
|
DB_PATH: str = os.getenv("DB_PATH", "data/tracker.db")
|
|
|
|
_DDL = """
|
|
CREATE TABLE IF NOT EXISTS runs (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
created_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%SZ', 'now')),
|
|
scenario TEXT NOT NULL,
|
|
identifier TEXT NOT NULL,
|
|
success INTEGER NOT NULL,
|
|
failure_reason TEXT,
|
|
steps_json TEXT NOT NULL DEFAULT '[]',
|
|
total_ms INTEGER NOT NULL DEFAULT 0
|
|
);
|
|
|
|
CREATE INDEX IF NOT EXISTS idx_runs_scenario ON runs(scenario);
|
|
CREATE INDEX IF NOT EXISTS idx_runs_created_at ON runs(created_at);
|
|
CREATE INDEX IF NOT EXISTS idx_runs_identifier ON runs(identifier);
|
|
|
|
CREATE TABLE IF NOT EXISTS flow_configs (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
created_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%SZ', 'now')),
|
|
name TEXT NOT NULL,
|
|
is_active INTEGER NOT NULL DEFAULT 0,
|
|
search_texts_json TEXT NOT NULL DEFAULT '[]',
|
|
search_box_selector TEXT NOT NULL DEFAULT '',
|
|
search_box_selector_type TEXT NOT NULL DEFAULT 'css',
|
|
search_submit_selector TEXT NOT NULL DEFAULT '',
|
|
search_submit_selector_type TEXT NOT NULL DEFAULT 'css',
|
|
search_results_selector TEXT NOT NULL DEFAULT '',
|
|
search_results_selector_type TEXT NOT NULL DEFAULT 'css',
|
|
scroll_container_selector TEXT NOT NULL DEFAULT '',
|
|
scroll_container_selector_type TEXT NOT NULL DEFAULT 'css',
|
|
max_scrolls INTEGER NOT NULL DEFAULT 30,
|
|
scroll_pause_ms INTEGER NOT NULL DEFAULT 1200,
|
|
no_new_content_timeout_ms INTEGER NOT NULL DEFAULT 3000,
|
|
pagination_selector TEXT NOT NULL DEFAULT '',
|
|
pagination_selector_type TEXT NOT NULL DEFAULT 'css',
|
|
max_pages INTEGER NOT NULL DEFAULT 10,
|
|
pagination_wait_ms INTEGER NOT NULL DEFAULT 1500,
|
|
item_selector_template TEXT NOT NULL DEFAULT '',
|
|
item_selector_type TEXT NOT NULL DEFAULT 'css',
|
|
item_url_template TEXT NOT NULL DEFAULT '',
|
|
target_item_ids_json TEXT NOT NULL DEFAULT '[]',
|
|
target_url TEXT NOT NULL DEFAULT ''
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS batch_runs (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
batch_id TEXT NOT NULL UNIQUE,
|
|
config_id INTEGER,
|
|
config_name TEXT NOT NULL DEFAULT '',
|
|
workers INTEGER NOT NULL DEFAULT 1,
|
|
total_runs INTEGER NOT NULL DEFAULT 0,
|
|
stagger_ms INTEGER NOT NULL DEFAULT 0,
|
|
headless INTEGER NOT NULL DEFAULT 1,
|
|
started_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%SZ', 'now')),
|
|
finished_at TEXT,
|
|
succeeded INTEGER NOT NULL DEFAULT 0,
|
|
failed INTEGER NOT NULL DEFAULT 0,
|
|
stopped INTEGER NOT NULL DEFAULT 0,
|
|
error TEXT
|
|
);
|
|
"""
|
|
|
|
|
|
def init_db() -> None:
|
|
"""Create the DB file and apply the schema (sync, called once at startup)."""
|
|
path = Path(DB_PATH)
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
conn = sqlite3.connect(str(path))
|
|
conn.executescript(_DDL)
|
|
# Additive migrations for columns added after initial release
|
|
_migrate(conn)
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
|
|
def _migrate(conn: sqlite3.Connection) -> None:
|
|
existing = {row[1] for row in conn.execute("PRAGMA table_info(flow_configs)")}
|
|
additions = {
|
|
"search_results_selector_type": "TEXT NOT NULL DEFAULT 'css'",
|
|
"item_url_template": "TEXT NOT NULL DEFAULT ''",
|
|
"scroll_container_selector_type": "TEXT NOT NULL DEFAULT 'css'",
|
|
"pagination_selector": "TEXT NOT NULL DEFAULT ''",
|
|
"pagination_selector_type": "TEXT NOT NULL DEFAULT 'css'",
|
|
"max_pages": "INTEGER NOT NULL DEFAULT 10",
|
|
"pagination_wait_ms": "INTEGER NOT NULL DEFAULT 1500",
|
|
"target_url": "TEXT NOT NULL DEFAULT ''",
|
|
"scenario": "TEXT NOT NULL DEFAULT 'dynamic'",
|
|
}
|
|
for col, definition in additions.items():
|
|
if col not in existing:
|
|
conn.execute(f"ALTER TABLE flow_configs ADD COLUMN {col} {definition}")
|
|
|
|
|
|
@asynccontextmanager
|
|
async def get_db() -> AsyncGenerator[aiosqlite.Connection, None]:
|
|
"""Async context manager — yields an open, Row-factory-enabled connection."""
|
|
async with aiosqlite.connect(DB_PATH) as db:
|
|
db.row_factory = aiosqlite.Row
|
|
yield db
|