from __future__ import annotations import json from datetime import datetime from typing import Any class ReviewRepositoryMixin: def list_watchlist(self, user_id: int) -> list[dict[str, Any]]: with self.connect() as connection: rows = connection.execute( """ SELECT code, name, sector, color, remark, created_at, updated_at FROM watchlist WHERE user_id = ? ORDER BY updated_at DESC, code """, (int(user_id),), ).fetchall() return [dict(row) for row in rows] def save_watchlist( self, user_id: int, code: str, name: str, sector: str, color: str, remark: str | None = None, ) -> None: now = datetime.now().astimezone().isoformat(timespec="seconds") with self.connect() as connection: existing = connection.execute( "SELECT remark FROM watchlist WHERE user_id = ? AND code = ?", (int(user_id), code), ).fetchone() saved_remark = ( str(existing["remark"] or "") if remark is None and existing else str(remark or "") ) connection.execute( """ INSERT INTO watchlist (user_id, code, name, sector, color, remark, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?) ON CONFLICT(user_id, code) DO UPDATE SET name = excluded.name, sector = excluded.sector, color = excluded.color, remark = excluded.remark, updated_at = excluded.updated_at """, (int(user_id), code, name, sector, color, saved_remark, now, now), ) def watchlist_price_history( self, codes: list[str], end_date: str, limit_per_code: int = 6 ) -> dict[str, list[dict[str, Any]]]: result: dict[str, list[dict[str, Any]]] = {} if not codes: return result with self.connect() as connection: for code in codes: rows = connection.execute( """ SELECT trade_date, ts_code, close, pct_chg FROM daily_bars WHERE substr(ts_code, 1, 6) = ? AND trade_date <= ? ORDER BY trade_date DESC LIMIT ? """, (str(code), end_date, int(limit_per_code)), ).fetchall() result[str(code)] = [dict(row) for row in reversed(rows)] return result def delete_watchlist(self, user_id: int, code: str) -> bool: with self.connect() as connection: cursor = connection.execute( "DELETE FROM watchlist WHERE user_id = ? AND code = ?", (int(user_id), code), ) return cursor.rowcount > 0 def list_notes( self, user_id: int, code: str = "", trade_date: str = "", scope: str = "all", ) -> list[dict[str, Any]]: clauses: list[str] = ["user_id = ?"] parameters: list[Any] = [int(user_id)] if scope == "daily": clauses.append("code = ''") elif scope == "stock": clauses.append("code <> ''") if code: clauses.append("code = ?") parameters.append(code) if trade_date: clauses.append("trade_date = ?") parameters.append(trade_date) where = f"WHERE {' AND '.join(clauses)}" if clauses else "" with self.connect() as connection: rows = connection.execute( f""" SELECT id, code, stock_name, trade_date, summary, content, plan, created_at, updated_at FROM review_notes {where} ORDER BY trade_date DESC, updated_at DESC, id DESC LIMIT 200 """, parameters, ).fetchall() return [dict(row) for row in rows] def save_note( self, user_id: int, code: str, stock_name: str, trade_date: str, content: str, plan: str, note_id: int | None = None, summary: str = "", ) -> int: now = datetime.now().astimezone().isoformat(timespec="seconds") with self.connect() as connection: if note_id: cursor = connection.execute( """ UPDATE review_notes SET code = ?, stock_name = ?, trade_date = ?, summary = ?, content = ?, plan = ?, updated_at = ? WHERE id = ? AND user_id = ? """, (code, stock_name, trade_date, summary, content, plan, now, note_id, int(user_id)), ) if cursor.rowcount == 0: raise ValueError("复盘笔记不存在。") return note_id cursor = connection.execute( """ INSERT INTO review_notes (user_id, code, stock_name, trade_date, summary, content, plan, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) """, (int(user_id), code, stock_name, trade_date, summary, content, plan, now, now), ) return int(cursor.lastrowid) def delete_note(self, user_id: int, note_id: int) -> bool: with self.connect() as connection: cursor = connection.execute( "DELETE FROM review_notes WHERE id = ? AND user_id = ?", (note_id, int(user_id)), ) return cursor.rowcount > 0 def save_trade_entry( self, user_id: int, trade_date: str, code: str, name: str, action: str, price: float, quantity: int, position_pct: float, pnl_amount: float | None, pnl_pct: float | None, thesis: str, execution: str, emotion: str, tags: list[str], trade_id: int | None = None, ) -> int: now = datetime.now().astimezone().isoformat(timespec="seconds") tags_json = json.dumps(tags, ensure_ascii=False, separators=(",", ":")) with self.connect() as connection: if trade_id: cursor = connection.execute( """ UPDATE trade_entries SET trade_date=?, code=?, name=?, action=?, price=?, quantity=?, position_pct=?, pnl_amount=?, pnl_pct=?, thesis=?, execution=?, emotion=?, tags=?, updated_at=? WHERE id=? AND user_id=? """, ( trade_date, code, name, action, price, quantity, position_pct, pnl_amount, pnl_pct, thesis, execution, emotion, tags_json, now, int(trade_id), int(user_id), ), ) if cursor.rowcount == 0: raise ValueError("交易记录不存在或无权修改。") return int(trade_id) cursor = connection.execute( """ INSERT INTO trade_entries (user_id, trade_date, code, name, action, price, quantity, position_pct, pnl_amount, pnl_pct, thesis, execution, emotion, tags, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( int(user_id), trade_date, code, name, action, price, quantity, position_pct, pnl_amount, pnl_pct, thesis, execution, emotion, tags_json, now, now, ), ) return int(cursor.lastrowid) def list_trade_entries( self, user_id: int, start_date: str = "", end_date: str = "", code: str = "", limit: int = 300, ) -> list[dict[str, Any]]: clauses = ["user_id = ?"] parameters: list[Any] = [int(user_id)] if start_date: clauses.append("trade_date >= ?") parameters.append(start_date) if end_date: clauses.append("trade_date <= ?") parameters.append(end_date) if code: clauses.append("code = ?") parameters.append(code) parameters.append(max(1, min(1000, int(limit)))) with self.connect() as connection: rows = connection.execute( f""" SELECT * FROM trade_entries WHERE {' AND '.join(clauses)} ORDER BY trade_date DESC, id DESC LIMIT ? """, parameters, ).fetchall() return [dict(row) for row in rows] def delete_trade_entry(self, user_id: int, trade_id: int) -> bool: with self.connect() as connection: cursor = connection.execute( "DELETE FROM trade_entries WHERE id = ? AND user_id = ?", (int(trade_id), int(user_id)), ) return cursor.rowcount > 0 def save_assistant_exchange( self, user_id: int, question: str, answer: str, context_date: str ) -> None: now = datetime.now().astimezone().isoformat(timespec="seconds") with self.connect() as connection: connection.executemany( """ INSERT INTO assistant_messages (user_id, role, content, context_date, created_at) VALUES (?, ?, ?, ?, ?) """, [ (int(user_id), "user", question, context_date, now), (int(user_id), "assistant", answer, context_date, now), ], ) connection.execute( """ DELETE FROM assistant_messages WHERE user_id = ? AND id NOT IN ( SELECT id FROM assistant_messages WHERE user_id = ? ORDER BY id DESC LIMIT 200 ) """, (int(user_id), int(user_id)), ) def list_assistant_messages(self, user_id: int, limit: int = 100) -> list[dict[str, Any]]: with self.connect() as connection: rows = connection.execute( """ SELECT role, content, context_date, created_at FROM assistant_messages WHERE user_id = ? ORDER BY id DESC LIMIT ? """, (int(user_id), max(1, min(200, int(limit)))), ).fetchall() return [dict(row) for row in reversed(rows)] def delete_assistant_messages(self, user_id: int) -> int: with self.connect() as connection: cursor = connection.execute( "DELETE FROM assistant_messages WHERE user_id = ?", (int(user_id),) ) return int(cursor.rowcount)