migration: preserve market data and search slice
This commit is contained in:
leefer
@@ -0,0 +1,222 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from datetime import datetime
|
||||
from typing import Any
|
||||
|
||||
|
||||
class MarketRepositoryMixin:
|
||||
def get_snapshot(self, trade_date: str) -> dict[str, Any] | None:
|
||||
with self.connect() as connection:
|
||||
row = connection.execute(
|
||||
"SELECT payload FROM dashboard_snapshots WHERE trade_date = ?",
|
||||
(trade_date,),
|
||||
).fetchone()
|
||||
if not row:
|
||||
return None
|
||||
try:
|
||||
return json.loads(row["payload"])
|
||||
except json.JSONDecodeError:
|
||||
return None
|
||||
|
||||
def get_latest_real_snapshot(
|
||||
self, trade_date: str, strictly_before: bool = False
|
||||
) -> dict[str, Any] | None:
|
||||
operator = "<" if strictly_before else "<="
|
||||
with self.connect() as connection:
|
||||
row = connection.execute(
|
||||
f"""
|
||||
SELECT payload FROM dashboard_snapshots
|
||||
WHERE trade_date {operator} ? AND source != 'demo'
|
||||
ORDER BY trade_date DESC LIMIT 1
|
||||
""",
|
||||
(trade_date,),
|
||||
).fetchone()
|
||||
if not row:
|
||||
return None
|
||||
try:
|
||||
return json.loads(row["payload"])
|
||||
except json.JSONDecodeError:
|
||||
return None
|
||||
|
||||
def save_snapshot(self, trade_date: str, source: str, payload: dict[str, Any]) -> None:
|
||||
updated_at = datetime.now().astimezone().isoformat(timespec="seconds")
|
||||
record_count = sum(
|
||||
len(payload.get(key) or [])
|
||||
for key in ("limits", "broken", "down_limits", "yesterday_limits")
|
||||
)
|
||||
content = json.dumps(payload, ensure_ascii=False, separators=(",", ":"))
|
||||
with self.connect() as connection:
|
||||
connection.execute(
|
||||
"""
|
||||
INSERT INTO dashboard_snapshots
|
||||
(trade_date, source, payload, record_count, updated_at)
|
||||
VALUES (?, ?, ?, ?, ?)
|
||||
ON CONFLICT(trade_date) DO UPDATE SET
|
||||
source = excluded.source,
|
||||
payload = excluded.payload,
|
||||
record_count = excluded.record_count,
|
||||
updated_at = excluded.updated_at
|
||||
""",
|
||||
(trade_date, source, content, record_count, updated_at),
|
||||
)
|
||||
|
||||
def get_data_snapshot(self, kind: str, cache_key: str) -> dict[str, Any] | None:
|
||||
with self.connect() as connection:
|
||||
row = connection.execute(
|
||||
"SELECT payload FROM data_snapshots WHERE kind = ? AND cache_key = ?",
|
||||
(kind, cache_key),
|
||||
).fetchone()
|
||||
if not row:
|
||||
return None
|
||||
try:
|
||||
return json.loads(row["payload"])
|
||||
except json.JSONDecodeError:
|
||||
return None
|
||||
|
||||
def get_latest_data_snapshot(
|
||||
self,
|
||||
kind: str,
|
||||
cache_key_prefix: str,
|
||||
maximum_cache_key: str,
|
||||
exclude_source: str = "",
|
||||
) -> dict[str, Any] | None:
|
||||
source_clause = " AND source != ?" if exclude_source else ""
|
||||
parameters: list[Any] = [kind, f"{cache_key_prefix}%", maximum_cache_key]
|
||||
if exclude_source:
|
||||
parameters.append(exclude_source)
|
||||
with self.connect() as connection:
|
||||
row = connection.execute(
|
||||
f"""
|
||||
SELECT payload FROM data_snapshots
|
||||
WHERE kind = ? AND cache_key LIKE ? AND cache_key <= ?{source_clause}
|
||||
ORDER BY cache_key DESC LIMIT 1
|
||||
""",
|
||||
parameters,
|
||||
).fetchone()
|
||||
if not row:
|
||||
return None
|
||||
try:
|
||||
return json.loads(row["payload"])
|
||||
except json.JSONDecodeError:
|
||||
return None
|
||||
|
||||
def save_data_snapshot(
|
||||
self, kind: str, cache_key: str, source: str, payload: dict[str, Any]
|
||||
) -> None:
|
||||
updated_at = datetime.now().astimezone().isoformat(timespec="seconds")
|
||||
content = json.dumps(payload, ensure_ascii=False, separators=(",", ":"))
|
||||
with self.connect() as connection:
|
||||
connection.execute(
|
||||
"""
|
||||
INSERT INTO data_snapshots (kind, cache_key, source, payload, updated_at)
|
||||
VALUES (?, ?, ?, ?, ?)
|
||||
ON CONFLICT(kind, cache_key) DO UPDATE SET
|
||||
source = excluded.source,
|
||||
payload = excluded.payload,
|
||||
updated_at = excluded.updated_at
|
||||
""",
|
||||
(kind, cache_key, source, content, updated_at),
|
||||
)
|
||||
|
||||
def search_stock_master(self, query: str, limit: int = 12) -> list[dict[str, Any]]:
|
||||
text = str(query or "").strip()
|
||||
if not text:
|
||||
return []
|
||||
escaped = text.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_")
|
||||
with self.connect() as connection:
|
||||
rows = connection.execute(
|
||||
"""
|
||||
SELECT ts_code, code, name, industry, market, list_date
|
||||
FROM stock_master
|
||||
WHERE code = ? OR name = ? OR name LIKE ? ESCAPE '\\'
|
||||
ORDER BY
|
||||
CASE WHEN code = ? THEN 0 WHEN name = ? THEN 1 ELSE 2 END,
|
||||
list_date DESC,
|
||||
code
|
||||
LIMIT ?
|
||||
""",
|
||||
(text, text, f"%{escaped}%", text, text, max(1, min(30, int(limit)))),
|
||||
).fetchall()
|
||||
return [dict(row) for row in rows]
|
||||
|
||||
def list_snapshot_payloads(self, end_date: str, limit: int = 260) -> list[dict[str, Any]]:
|
||||
with self.connect() as connection:
|
||||
rows = connection.execute(
|
||||
"""
|
||||
SELECT trade_date, payload FROM dashboard_snapshots
|
||||
WHERE trade_date <= ? ORDER BY trade_date DESC LIMIT ?
|
||||
""",
|
||||
(end_date, limit),
|
||||
).fetchall()
|
||||
result: list[dict[str, Any]] = []
|
||||
for row in reversed(rows):
|
||||
try:
|
||||
payload = json.loads(row["payload"])
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
payload["_snapshot_date"] = row["trade_date"]
|
||||
result.append(payload)
|
||||
return result
|
||||
|
||||
def start_sync(self, trade_date: str, source: str) -> int:
|
||||
started_at = datetime.now().astimezone().isoformat(timespec="seconds")
|
||||
with self.connect() as connection:
|
||||
cursor = connection.execute(
|
||||
"""
|
||||
INSERT INTO sync_runs (trade_date, source, status, started_at)
|
||||
VALUES (?, ?, 'running', ?)
|
||||
""",
|
||||
(trade_date, source, started_at),
|
||||
)
|
||||
return int(cursor.lastrowid)
|
||||
|
||||
def finish_sync(
|
||||
self,
|
||||
sync_id: int,
|
||||
status: str,
|
||||
record_count: int = 0,
|
||||
message: str = "",
|
||||
source: str | None = None,
|
||||
) -> None:
|
||||
finished_at = datetime.now().astimezone().isoformat(timespec="seconds")
|
||||
with self.connect() as connection:
|
||||
connection.execute(
|
||||
"""
|
||||
UPDATE sync_runs
|
||||
SET status = ?, finished_at = ?, record_count = ?, message = ?,
|
||||
source = COALESCE(?, source)
|
||||
WHERE id = ?
|
||||
""",
|
||||
(status, finished_at, record_count, message[:1000], source, sync_id),
|
||||
)
|
||||
|
||||
def status(self) -> dict[str, Any]:
|
||||
with self.connect() as connection:
|
||||
last_sync = connection.execute(
|
||||
"""
|
||||
SELECT id, trade_date, source, status, started_at, finished_at,
|
||||
record_count, message
|
||||
FROM sync_runs ORDER BY id DESC LIMIT 1
|
||||
"""
|
||||
).fetchone()
|
||||
snapshot_stats = connection.execute(
|
||||
"""
|
||||
SELECT COUNT(*) AS dates, COALESCE(SUM(record_count), 0) AS records,
|
||||
MAX(updated_at) AS updated_at
|
||||
FROM dashboard_snapshots
|
||||
"""
|
||||
).fetchone()
|
||||
watchlist_count = connection.execute("SELECT COUNT(*) FROM watchlist").fetchone()[0]
|
||||
note_count = connection.execute("SELECT COUNT(*) FROM review_notes").fetchone()[0]
|
||||
|
||||
return {
|
||||
"database": str(self.path.name),
|
||||
"snapshot_dates": int(snapshot_stats["dates"]),
|
||||
"snapshot_records": int(snapshot_stats["records"]),
|
||||
"updated_at": snapshot_stats["updated_at"],
|
||||
"last_sync": dict(last_sync) if last_sync else None,
|
||||
"watchlist_count": int(watchlist_count),
|
||||
"note_count": int(note_count),
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user