"""Postgres repository implementation for portfolios and ledger.""" from __future__ import annotations import hashlib import json import uuid from dataclasses import dataclass from datetime import datetime, timezone from typing import Any import psycopg from psycopg import sql from app.domain.models import AuthToken, BotPrincipal, LedgerEntry, Portfolio, PortfolioComment, PortfolioTag, StrategySet, TokenPrincipal, User _SELECT_USER = """ SELECT id, username, email, permission, is_active, password_hash, created_at, updated_at, last_login_at FROM users """ _SELECT_AUTH_TOKEN = """ SELECT id, user_id, name, token_prefix, permission, token_type, bot_principal_id, expires_at, last_used_at, revoked_at, created_at FROM auth_tokens """ _SELECT_BOT_PRINCIPAL = """ SELECT id, name, owner_user_id, permission, description, is_active, created_at, updated_at, deactivated_at FROM bot_principals """ _SELECT_PORTFOLIO = """ SELECT portfolio_id, broker, status, name, symbol, cash_balance, asset_quantity, asset_price, profit_buffer, active_strategy_id, avabile_strategies, owner_user_id, created_at, updated_at, last_trade_at FROM portfolios """ @dataclass(frozen=True) class PortfolioDBConfig: host: str port: int dbname: str user: str password: str sslmode: str connect_timeout_sec: int = 10 application_name: str = "portfoliohub" class PostgresPortfolioRepository: def __init__(self, config: PortfolioDBConfig) -> None: self._config = config self._connection: psycopg.Connection | None = None def _connect(self) -> psycopg.Connection: if self._connection is None: self._connection = psycopg.connect( host=self._config.host, port=self._config.port, dbname=self._config.dbname, user=self._config.user, password=self._config.password, sslmode=self._config.sslmode, connect_timeout=self._config.connect_timeout_sec, application_name=self._config.application_name, autocommit=True, ) return self._connection def close(self) -> None: if self._connection is not None: self._connection.close() self._connection = None def ping(self) -> None: with self._connect().cursor() as cur: cur.execute("SELECT 1") cur.fetchone() def ensure_schema(self) -> None: conn = self._connect() with conn.cursor() as cur: cur.execute( """ CREATE TABLE IF NOT EXISTS users ( id TEXT PRIMARY KEY, username TEXT NOT NULL UNIQUE, email TEXT UNIQUE, permission TEXT NOT NULL DEFAULT 'read', is_active BOOLEAN NOT NULL DEFAULT TRUE, password_hash TEXT, created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), last_login_at TIMESTAMPTZ ) """ ) cur.execute( """ CREATE TABLE IF NOT EXISTS auth_tokens ( id TEXT PRIMARY KEY, user_id TEXT NOT NULL, name TEXT NOT NULL, token_prefix TEXT NOT NULL, token_hash TEXT NOT NULL UNIQUE, permission TEXT NOT NULL, token_type TEXT NOT NULL, bot_principal_id TEXT, expires_at TIMESTAMPTZ, last_used_at TIMESTAMPTZ, revoked_at TIMESTAMPTZ, created_at TIMESTAMPTZ NOT NULL DEFAULT NOW() ) """ ) cur.execute( """ CREATE TABLE IF NOT EXISTS bot_principals ( id TEXT PRIMARY KEY, name TEXT NOT NULL UNIQUE, owner_user_id TEXT, permission TEXT NOT NULL, description TEXT, is_active BOOLEAN NOT NULL DEFAULT TRUE, created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), deactivated_at TIMESTAMPTZ ) """ ) cur.execute("ALTER TABLE auth_tokens ADD COLUMN IF NOT EXISTS bot_principal_id TEXT") cur.execute("ALTER TABLE bot_principals ADD COLUMN IF NOT EXISTS owner_user_id TEXT") cur.execute("ALTER TABLE bot_principals ADD COLUMN IF NOT EXISTS permission TEXT") cur.execute("UPDATE bot_principals SET permission = 'read' WHERE permission IS NULL") cur.execute("ALTER TABLE bot_principals ALTER COLUMN permission SET NOT NULL") cur.execute( """ UPDATE bot_principals SET owner_user_id = fallback.id FROM ( SELECT id FROM users ORDER BY created_at ASC, id ASC LIMIT 1 ) AS fallback WHERE owner_user_id IS NULL """ ) cur.execute("SELECT COUNT(*) FROM bot_principals WHERE owner_user_id IS NULL") unresolved_bot_owners = int(cur.fetchone()[0]) if unresolved_bot_owners > 0: raise RuntimeError("bot_principals.owner_user_id backfill failed: users 테이블에 소유주가 필요합니다") cur.execute("ALTER TABLE bot_principals ALTER COLUMN owner_user_id SET NOT NULL") cur.execute( """ DO $$ BEGIN IF NOT EXISTS ( SELECT 1 FROM pg_constraint WHERE conname = 'fk_bot_principals_owner_user' ) THEN ALTER TABLE bot_principals ADD CONSTRAINT fk_bot_principals_owner_user FOREIGN KEY (owner_user_id) REFERENCES users(id) ON DELETE RESTRICT; END IF; END $$ """ ) cur.execute( """ CREATE INDEX IF NOT EXISTS idx_auth_tokens_user_created ON auth_tokens (user_id, created_at DESC) """ ) cur.execute( """ CREATE INDEX IF NOT EXISTS idx_auth_tokens_hash_active ON auth_tokens (token_hash) """ ) cur.execute( """ CREATE INDEX IF NOT EXISTS idx_auth_tokens_bot_created ON auth_tokens (bot_principal_id, created_at DESC) """ ) cur.execute( """ CREATE INDEX IF NOT EXISTS idx_bot_principals_active_created ON bot_principals (is_active, created_at DESC) """ ) cur.execute( """ CREATE INDEX IF NOT EXISTS idx_bot_principals_owner_active_created ON bot_principals (owner_user_id, is_active, created_at DESC) """ ) cur.execute( """ CREATE TABLE IF NOT EXISTS strategy_sets ( id TEXT PRIMARY KEY, name TEXT NOT NULL, config TEXT NOT NULL DEFAULT '', nick TEXT NOT NULL UNIQUE, created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), UNIQUE (name, config) ) """ ) cur.execute( """ CREATE TABLE IF NOT EXISTS portfolios ( portfolio_id TEXT PRIMARY KEY, broker TEXT NOT NULL DEFAULT 'kis', status TEXT NOT NULL DEFAULT 'enabled', name TEXT NOT NULL, symbol TEXT NOT NULL, owner_user_id TEXT, cash_balance DOUBLE PRECISION NOT NULL, asset_quantity DOUBLE PRECISION NOT NULL, asset_price DOUBLE PRECISION NOT NULL DEFAULT 0, profit_buffer DOUBLE PRECISION NOT NULL DEFAULT 0, active_strategy_id TEXT DEFAULT NULL, avabile_strategies TEXT[] DEFAULT NULL, created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), last_trade_at TIMESTAMPTZ ) """ ) cur.execute("ALTER TABLE portfolios ADD COLUMN IF NOT EXISTS name TEXT") cur.execute("ALTER TABLE portfolios ADD COLUMN IF NOT EXISTS owner_user_id TEXT") cur.execute("UPDATE portfolios SET name = symbol WHERE name IS NULL OR name = ''") cur.execute("ALTER TABLE portfolios ALTER COLUMN name SET NOT NULL") cur.execute( """ CREATE TABLE IF NOT EXISTS portfolio_ledger ( portfolio_id TEXT NOT NULL, recorded_at TIMESTAMPTZ NOT NULL, action TEXT NOT NULL, price DOUBLE PRECISION NOT NULL, quantity DOUBLE PRECISION NOT NULL, fee_amount DOUBLE PRECISION NOT NULL DEFAULT 0, tax_amount DOUBLE PRECISION NOT NULL DEFAULT 0, cash_balance_after DOUBLE PRECISION NOT NULL, asset_quantity_after DOUBLE PRECISION NOT NULL, total_value_after DOUBLE PRECISION NOT NULL, profit_buffer_after DOUBLE PRECISION NOT NULL DEFAULT 0, strategy_set_id TEXT NOT NULL DEFAULT 'legacy', note TEXT NOT NULL DEFAULT '' ) """ ) cur.execute( """ CREATE INDEX IF NOT EXISTS idx_portfolio_ledger_portfolio_time ON portfolio_ledger (portfolio_id, recorded_at DESC) """ ) cur.execute( """ CREATE INDEX IF NOT EXISTS idx_portfolios_broker_status_updated ON portfolios (broker, status, updated_at DESC) """ ) cur.execute( """ CREATE TABLE IF NOT EXISTS portfolio_comments ( id TEXT PRIMARY KEY, portfolio_id TEXT NOT NULL, content TEXT NOT NULL, created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW() ) """ ) cur.execute( """ CREATE INDEX IF NOT EXISTS idx_portfolio_comments_portfolio_created ON portfolio_comments (portfolio_id, created_at DESC) """ ) cur.execute( """ CREATE TABLE IF NOT EXISTS portfolio_tags ( portfolio_id TEXT NOT NULL, tag_value TEXT NOT NULL, created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), PRIMARY KEY (portfolio_id, tag_value) ) """ ) cur.execute( """ CREATE INDEX IF NOT EXISTS idx_portfolio_tags_tag_portfolio ON portfolio_tags (tag_value, portfolio_id) """ ) cur.execute("SELECT EXISTS (SELECT 1 FROM pg_extension WHERE extname = 'timescaledb')") if cur.fetchone()[0]: cur.execute( """ SELECT create_hypertable('portfolio_ledger', 'recorded_at', if_not_exists => TRUE) """ ) @staticmethod def _row_to_user(row: Any) -> User: if row is None: raise ValueError("사용자를 찾을 수 없습니다") return User( id=row[0], username=row[1], email=row[2], permission=row[3] if len(row) > 3 else "read", is_active=bool(row[4]) if len(row) > 4 else True, password_hash=row[5] if len(row) > 5 else None, created_at=row[6] if len(row) > 6 else None, updated_at=row[7] if len(row) > 7 else None, last_login_at=row[8] if len(row) > 8 else None, ) @staticmethod def _row_to_auth_token(row: Any) -> AuthToken: if row is None: raise ValueError("토큰을 찾을 수 없습니다") return AuthToken( id=row[0], user_id=row[1], name=row[2], token_prefix=row[3], permission=row[4], token_type=row[5], bot_principal_id=row[6] if len(row) > 6 else None, expires_at=row[7] if len(row) > 7 else None, last_used_at=row[8] if len(row) > 8 else None, revoked_at=row[9] if len(row) > 9 else None, created_at=row[10] if len(row) > 10 else None, ) @staticmethod def _row_to_bot_principal(row: Any) -> BotPrincipal: if row is None: raise ValueError("봇을 찾을 수 없습니다") return BotPrincipal( id=row[0], name=row[1], owner_user_id=row[2], permission=row[3], description=row[4] if len(row) > 4 else None, is_active=bool(row[5]) if len(row) > 5 else True, created_at=row[6] if len(row) > 6 else None, updated_at=row[7] if len(row) > 7 else None, deactivated_at=row[8] if len(row) > 8 else None, ) @staticmethod def _row_to_portfolio(row: Any) -> Portfolio: if row is None: raise ValueError("포트폴리오를 찾을 수 없습니다") raw_available = row[10] if len(row) > 10 else None if raw_available is None: avabile_strategies = None elif isinstance(raw_available, list): avabile_strategies = [str(item) for item in raw_available] else: avabile_strategies = [str(raw_available)] return Portfolio( portfolio_id=row[0], broker=row[1] or "kis", status=row[2] or "enabled", name=row[3], symbol=row[4], tags=None, cash_balance=float(row[5]), asset_quantity=float(row[6]), price=float(row[7]), profit_buffer=float(row[8]), active_strategy_id=row[9] if len(row) > 9 else None, avabile_strategies=avabile_strategies, owner_user_id=row[11] if len(row) > 11 else None, created_at=row[12] if len(row) > 12 else None, updated_at=row[13] if len(row) > 13 else None, last_trade_at=row[14] if len(row) > 14 else None, ) @staticmethod def _row_to_portfolio_tag(row: Any) -> PortfolioTag: if row is None: raise ValueError("태그를 찾을 수 없습니다") return PortfolioTag(portfolio_id=row[0], tag=row[1], created_at=row[2] if len(row) > 2 else None) def _load_tags_for_portfolios(self, cur: psycopg.Cursor, portfolio_ids: list[str]) -> dict[str, list[str]]: if not portfolio_ids: return {} cur.execute( """ SELECT portfolio_id, tag_value FROM portfolio_tags WHERE portfolio_id = ANY(%s) ORDER BY tag_value ASC """, (portfolio_ids,), ) tags_by_portfolio: dict[str, list[str]] = {portfolio_id: [] for portfolio_id in portfolio_ids} for portfolio_id, tag_value in cur.fetchall(): tags_by_portfolio.setdefault(str(portfolio_id), []).append(str(tag_value)) return tags_by_portfolio def _attach_tags(self, cur: psycopg.Cursor, portfolios: list[Portfolio]) -> list[Portfolio]: tags_by_portfolio = self._load_tags_for_portfolios(cur, [portfolio.portfolio_id for portfolio in portfolios]) for portfolio in portfolios: portfolio.tags = tags_by_portfolio.get(portfolio.portfolio_id, []) return portfolios @staticmethod def _row_to_ledger_entry(row: Any) -> LedgerEntry: return LedgerEntry( portfolio_id=row[0], recorded_at=row[1], action=row[2], price=float(row[3]), quantity=float(row[4]), fee_amount=float(row[5]), tax_amount=float(row[6]), cash_balance_after=float(row[7]), asset_quantity_after=float(row[8]), total_value_after=float(row[9]), profit_buffer_after=float(row[10]), strategy_set_id=row[11] or "legacy", note=row[12] or "", ) def _ensure_user_uniqueness( self, cur: psycopg.Cursor, *, username: str | None = None, email: str | None = None, exclude_user_id: str | None = None, ) -> None: if username is not None: if exclude_user_id is None: cur.execute("SELECT 1 FROM users WHERE username = %s", (username,)) else: cur.execute("SELECT 1 FROM users WHERE username = %s AND id <> %s", (username, exclude_user_id)) if cur.fetchone() is not None: raise ValueError(f"이미 존재하는 username입니다: {username}") if email is not None: if exclude_user_id is None: cur.execute("SELECT 1 FROM users WHERE email = %s", (email,)) else: cur.execute("SELECT 1 FROM users WHERE email = %s AND id <> %s", (email, exclude_user_id)) if cur.fetchone() is not None: raise ValueError(f"이미 존재하는 email입니다: {email}") def get_user(self, user_id: str) -> User: self.ensure_schema() with self._connect().cursor() as cur: cur.execute(_SELECT_USER + "WHERE id = %s", (user_id,)) row = cur.fetchone() return self._row_to_user(row) def get_user_by_username(self, username: str) -> User | None: self.ensure_schema() with self._connect().cursor() as cur: cur.execute(_SELECT_USER + "WHERE username = %s", (username,)) row = cur.fetchone() if row is None: return None return self._row_to_user(row) def list_users(self) -> list[User]: self.ensure_schema() with self._connect().cursor() as cur: cur.execute(_SELECT_USER + "ORDER BY created_at DESC, username ASC") rows = cur.fetchall() return [self._row_to_user(row) for row in rows] def create_user( self, username: str, *, email: str | None = None, password_hash: str | None = None, permission: str = "read", is_active: bool = True, ) -> User: self.ensure_schema() with self._connect().cursor() as cur: self._ensure_user_uniqueness(cur, username=username, email=email) user_id = str(uuid.uuid4()) cur.execute( """ INSERT INTO users (id, username, email, permission, is_active, password_hash) VALUES (%s, %s, %s, %s, %s, %s) RETURNING id, username, email, permission, is_active, password_hash, created_at, updated_at, last_login_at """, (user_id, username, email, permission, is_active, password_hash), ) row = cur.fetchone() return self._row_to_user(row) def update_user( self, user_id: str, *, username: str | None = None, email: str | None = None, password_hash: str | None = None, permission: str | None = None, is_active: bool | None = None, ) -> User: self.ensure_schema() assignments: list[str] = [] params: list[object] = [] with self._connect().cursor() as cur: cur.execute(_SELECT_USER + "WHERE id = %s", (user_id,)) existing_row = cur.fetchone() if existing_row is None: raise ValueError("사용자를 찾을 수 없습니다") self._ensure_user_uniqueness(cur, username=username, email=email, exclude_user_id=user_id) if username is not None: assignments.append("username = %s") params.append(username) if email is not None: assignments.append("email = %s") params.append(email) if password_hash is not None: assignments.append("password_hash = %s") params.append(password_hash) if permission is not None: assignments.append("permission = %s") params.append(permission) if is_active is not None: assignments.append("is_active = %s") params.append(is_active) if not assignments: return self._row_to_user(existing_row) assignments.append("updated_at = NOW()") cur.execute( f""" UPDATE users SET {', '.join(assignments)} WHERE id = %s RETURNING id, username, email, permission, is_active, password_hash, created_at, updated_at, last_login_at """, (*params, user_id), ) row = cur.fetchone() return self._row_to_user(row) def delete_user(self, user_id: str) -> None: self.ensure_schema() with self._connect().cursor() as cur: cur.execute("SELECT 1 FROM users WHERE id = %s", (user_id,)) if cur.fetchone() is None: raise ValueError("사용자를 찾을 수 없습니다") cur.execute("DELETE FROM auth_tokens WHERE user_id = %s", (user_id,)) cur.execute("DELETE FROM users WHERE id = %s", (user_id,)) def update_user_last_login(self, user_id: str) -> None: self.ensure_schema() with self._connect().cursor() as cur: cur.execute( "UPDATE users SET last_login_at = NOW(), updated_at = NOW() WHERE id = %s", (user_id,), ) def create_auth_token( self, user_id: str, *, name: str, token_hash: str, token_prefix: str, permission: str, token_type: str, bot_principal_id: str | None = None, expires_at: datetime | None = None, ) -> AuthToken: self.ensure_schema() with self._connect().cursor() as cur: if token_type != "bot": cur.execute("SELECT 1 FROM users WHERE id = %s", (user_id,)) if cur.fetchone() is None: raise ValueError("사용자를 찾을 수 없습니다") if bot_principal_id is not None: cur.execute("SELECT 1 FROM bot_principals WHERE id = %s", (bot_principal_id,)) if cur.fetchone() is None: raise ValueError("봇을 찾을 수 없습니다") token_id = str(uuid.uuid4()) cur.execute( """ INSERT INTO auth_tokens ( id, user_id, name, token_prefix, token_hash, permission, token_type, bot_principal_id, expires_at ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s) RETURNING id, user_id, name, token_prefix, permission, token_type, bot_principal_id, expires_at, last_used_at, revoked_at, created_at """, (token_id, user_id, name, token_prefix, token_hash, permission, token_type, bot_principal_id, expires_at), ) row = cur.fetchone() return self._row_to_auth_token(row) def list_auth_tokens(self, user_id: str) -> list[AuthToken]: self.ensure_schema() with self._connect().cursor() as cur: cur.execute( _SELECT_AUTH_TOKEN + "WHERE user_id = %s ORDER BY created_at DESC, id ASC", (user_id,), ) rows = cur.fetchall() return [self._row_to_auth_token(row) for row in rows] def revoke_auth_token(self, user_id: str, token_id: str) -> None: self.ensure_schema() with self._connect().cursor() as cur: cur.execute( "UPDATE auth_tokens SET revoked_at = NOW() WHERE id = %s AND user_id = %s RETURNING id", (token_id, user_id), ) row = cur.fetchone() if row is None: raise ValueError("토큰을 찾을 수 없습니다") def create_bot_principal( self, *, name: str, owner_user_id: str, permission: str, description: str | None = None, ) -> BotPrincipal: self.ensure_schema() with self._connect().cursor() as cur: cur.execute("SELECT 1 FROM bot_principals WHERE name = %s", (name,)) if cur.fetchone() is not None: raise ValueError(f"이미 존재하는 bot name입니다: {name}") cur.execute("SELECT 1 FROM users WHERE id = %s", (owner_user_id,)) if cur.fetchone() is None: raise ValueError("사용자를 찾을 수 없습니다") bot_id = str(uuid.uuid4()) cur.execute( """ INSERT INTO bot_principals (id, name, owner_user_id, permission, description) VALUES (%s, %s, %s, %s, %s) RETURNING id, name, owner_user_id, permission, description, is_active, created_at, updated_at, deactivated_at """, (bot_id, name, owner_user_id, permission, description), ) row = cur.fetchone() return self._row_to_bot_principal(row) def get_bot_principal(self, bot_principal_id: str) -> BotPrincipal: self.ensure_schema() with self._connect().cursor() as cur: cur.execute(_SELECT_BOT_PRINCIPAL + "WHERE id = %s", (bot_principal_id,)) row = cur.fetchone() return self._row_to_bot_principal(row) def list_bot_principals( self, *, is_active: bool | None = None, owner_user_id: str | None = None, ) -> list[BotPrincipal]: self.ensure_schema() query = _SELECT_BOT_PRINCIPAL params: list[object] = [] clauses: list[str] = [] if is_active is not None: clauses.append("is_active = %s") params.append(is_active) if owner_user_id is not None: clauses.append("owner_user_id = %s") params.append(owner_user_id) if clauses: query += "WHERE " + " AND ".join(clauses) + " " query += "ORDER BY created_at DESC, name ASC" with self._connect().cursor() as cur: cur.execute(query, tuple(params)) rows = cur.fetchall() return [self._row_to_bot_principal(row) for row in rows] def deactivate_bot_principal(self, bot_principal_id: str) -> BotPrincipal: self.ensure_schema() with self._connect().cursor() as cur: cur.execute( """ UPDATE bot_principals SET is_active = FALSE, deactivated_at = NOW(), updated_at = NOW() WHERE id = %s RETURNING id, name, owner_user_id, permission, description, is_active, created_at, updated_at, deactivated_at """, (bot_principal_id,), ) row = cur.fetchone() if row is None: raise ValueError("봇을 찾을 수 없습니다") cur.execute( """ UPDATE auth_tokens SET revoked_at = NOW() WHERE bot_principal_id = %s AND revoked_at IS NULL """, (bot_principal_id,), ) return self._row_to_bot_principal(row) def list_bot_auth_tokens(self, bot_principal_id: str) -> list[AuthToken]: self.ensure_schema() with self._connect().cursor() as cur: cur.execute( _SELECT_AUTH_TOKEN + "WHERE bot_principal_id = %s ORDER BY created_at DESC, id ASC", (bot_principal_id,), ) rows = cur.fetchall() return [self._row_to_auth_token(row) for row in rows] def revoke_bot_auth_token(self, bot_principal_id: str, token_id: str) -> None: self.ensure_schema() with self._connect().cursor() as cur: cur.execute( """ UPDATE auth_tokens SET revoked_at = NOW() WHERE id = %s AND bot_principal_id = %s RETURNING id """, (token_id, bot_principal_id), ) row = cur.fetchone() if row is None: raise ValueError("토큰을 찾을 수 없습니다") def get_token_principal(self, token_hash: str) -> TokenPrincipal | None: self.ensure_schema() with self._connect().cursor() as cur: cur.execute( """ SELECT t.id, COALESCE(b.owner_user_id, t.user_id), t.permission, t.token_type, t.bot_principal_id FROM auth_tokens t LEFT JOIN bot_principals b ON b.id = t.bot_principal_id WHERE token_hash = %s AND t.revoked_at IS NULL AND (t.expires_at IS NULL OR t.expires_at > NOW()) AND (t.bot_principal_id IS NULL OR b.is_active = TRUE) """, (token_hash,), ) row = cur.fetchone() if row is None: return None cur.execute("UPDATE auth_tokens SET last_used_at = NOW() WHERE id = %s", (row[0],)) return TokenPrincipal( token_id=row[0], user_id=row[1], permission=row[2], token_type=row[3], bot_principal_id=row[4], ) def get(self, portfolio_id: str) -> Portfolio: self.ensure_schema() with self._connect().cursor() as cur: cur.execute(_SELECT_PORTFOLIO + "WHERE portfolio_id = %s", (portfolio_id,)) row = cur.fetchone() portfolio = self._row_to_portfolio(row) self._attach_tags(cur, [portfolio]) return portfolio def get_or_create( self, portfolio_id: str, symbol: str, *, name: str | None = None, broker: str = "kis", status: str = "enabled", active_strategy_id: str | None = None, avabile_strategies: list[str] | None = None, owner_user_id: str | None = None, initial_cash: float = 0.0, ) -> Portfolio: self.ensure_schema() with self._connect().cursor() as cur: cur.execute( """ INSERT INTO portfolios ( portfolio_id, broker, status, name, symbol, owner_user_id, cash_balance, asset_quantity, asset_price, profit_buffer, active_strategy_id, avabile_strategies ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) ON CONFLICT (portfolio_id) DO NOTHING """, ( portfolio_id, broker, status, name or symbol, symbol, owner_user_id, float(initial_cash), 0.0, 0.0, 0.0, active_strategy_id, avabile_strategies, ), ) cur.execute(_SELECT_PORTFOLIO + "WHERE portfolio_id = %s", (portfolio_id,)) row = cur.fetchone() return self._row_to_portfolio(row) def create_portfolio( self, portfolio_id: str, *, owner_user_id: str, broker: str, status: str, name: str, symbol: str, cash_balance: float, asset_quantity: float, asset_price: float, profit_buffer: float, active_strategy_id: str | None, avabile_strategies: list[str] | None, ) -> Portfolio: self.ensure_schema() with self._connect().cursor() as cur: cur.execute("SELECT 1 FROM portfolios WHERE portfolio_id = %s", (portfolio_id,)) if cur.fetchone() is not None: raise ValueError(f"이미 존재하는 포트폴리오입니다: {portfolio_id}") cur.execute( """ INSERT INTO portfolios ( portfolio_id, broker, status, name, symbol, owner_user_id, cash_balance, asset_quantity, asset_price, profit_buffer, active_strategy_id, avabile_strategies ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) RETURNING portfolio_id, broker, status, name, symbol, cash_balance, asset_quantity, asset_price, profit_buffer, active_strategy_id, avabile_strategies, owner_user_id, created_at, updated_at, last_trade_at """, ( portfolio_id, broker, status, name, symbol, owner_user_id, float(cash_balance), float(asset_quantity), float(asset_price), float(profit_buffer), active_strategy_id, avabile_strategies, ), ) row = cur.fetchone() portfolio = self._row_to_portfolio(row) portfolio.tags = [] return portfolio def update_portfolio(self, portfolio_id: str, **fields: object) -> Portfolio: self.ensure_schema() allowed = { "broker": "broker", "status": "status", "name": "name", "symbol": "symbol", "asset_price": "asset_price", "active_strategy_id": "active_strategy_id", "avabile_strategies": "avabile_strategies", } assignments: list[str] = [] params: list[object] = [] for key, value in fields.items(): if value is None or key not in allowed: continue if key == "avabile_strategies": assignments.append("avabile_strategies = %s") params.append(value) else: assignments.append(f"{allowed[key]} = %s") params.append(value) if not assignments: return self.get(portfolio_id) assignments.append("updated_at = NOW()") with self._connect().cursor() as cur: cur.execute( f""" UPDATE portfolios SET {', '.join(assignments)} WHERE portfolio_id = %s RETURNING portfolio_id, broker, status, name, symbol, cash_balance, asset_quantity, asset_price, profit_buffer, active_strategy_id, avabile_strategies, owner_user_id, created_at, updated_at, last_trade_at """, (*params, portfolio_id), ) row = cur.fetchone() return self._row_to_portfolio(row) def list_portfolios( self, *, broker: str | None = None, status: str | None = None, owner_user_id: str | None = None, include_tags: list[str] | None = None, exclude_tags: list[str] | None = None, match_mode: str | None = None, ) -> list[Portfolio]: self.ensure_schema() query = _SELECT_PORTFOLIO clauses: list[str] = [] params: list[object] = [] if broker is not None: clauses.append("broker = %s") params.append(broker) if status is not None: clauses.append("status = %s") params.append(status) if owner_user_id is not None: clauses.append("owner_user_id = %s") params.append(owner_user_id) normalized_mode = (match_mode or "OR").upper() if include_tags: if normalized_mode == "AND": clauses.append( """ portfolio_id IN ( SELECT portfolio_id FROM portfolio_tags WHERE tag_value = ANY(%s) GROUP BY portfolio_id HAVING COUNT(DISTINCT tag_value) = %s ) """ ) params.append(include_tags) params.append(len(set(include_tags))) else: clauses.append( "portfolio_id IN (SELECT DISTINCT portfolio_id FROM portfolio_tags WHERE tag_value = ANY(%s))" ) params.append(include_tags) if exclude_tags: clauses.append( "portfolio_id NOT IN (SELECT DISTINCT portfolio_id FROM portfolio_tags WHERE tag_value = ANY(%s))" ) params.append(exclude_tags) if clauses: query += "WHERE " + " AND ".join(clauses) + " " query += "ORDER BY updated_at DESC, portfolio_id ASC" with self._connect().cursor() as cur: cur.execute(query, tuple(params)) rows = cur.fetchall() portfolios = [self._row_to_portfolio(row) for row in rows] self._attach_tags(cur, portfolios) return portfolios def add_tag(self, portfolio_id: str, *, tag: str) -> PortfolioTag: self.ensure_schema() with self._connect().cursor() as cur: cur.execute("SELECT 1 FROM portfolios WHERE portfolio_id = %s", (portfolio_id,)) if cur.fetchone() is None: raise ValueError("포트폴리오를 찾을 수 없습니다") try: cur.execute( """ INSERT INTO portfolio_tags (portfolio_id, tag_value) VALUES (%s, %s) RETURNING portfolio_id, tag_value, created_at """, (portfolio_id, tag), ) except psycopg.errors.UniqueViolation as exc: raise ValueError(f"이미 연결된 태그입니다: {tag}") from exc row = cur.fetchone() return self._row_to_portfolio_tag(row) def list_tags(self, portfolio_id: str) -> list[PortfolioTag]: self.ensure_schema() with self._connect().cursor() as cur: cur.execute("SELECT 1 FROM portfolios WHERE portfolio_id = %s", (portfolio_id,)) if cur.fetchone() is None: raise ValueError("포트폴리오를 찾을 수 없습니다") cur.execute( """ SELECT portfolio_id, tag_value, created_at FROM portfolio_tags WHERE portfolio_id = %s ORDER BY tag_value ASC """, (portfolio_id,), ) rows = cur.fetchall() return [self._row_to_portfolio_tag(row) for row in rows] def remove_tag(self, portfolio_id: str, *, tag: str) -> None: self.ensure_schema() with self._connect().cursor() as cur: cur.execute("SELECT 1 FROM portfolios WHERE portfolio_id = %s", (portfolio_id,)) if cur.fetchone() is None: raise ValueError("포트폴리오를 찾을 수 없습니다") cur.execute( "DELETE FROM portfolio_tags WHERE portfolio_id = %s AND tag_value = %s RETURNING portfolio_id", (portfolio_id, tag), ) if cur.fetchone() is None: raise ValueError("태그를 찾을 수 없습니다") def get_or_create_strategy_set(self, name: str, config: str) -> StrategySet: self.ensure_schema() nick = f"{name}:{hashlib.sha1(config.encode('utf-8')).hexdigest()[:12]}" with self._connect().cursor() as cur: new_id = str(uuid.uuid4()) cur.execute( """ INSERT INTO strategy_sets (id, name, config, nick) VALUES (%s, %s, %s, %s) ON CONFLICT (name, config) DO NOTHING """, (new_id, name, config, nick), ) cur.execute( """ SELECT id, name, config, nick, created_at FROM strategy_sets WHERE name = %s AND config = %s """, (name, config), ) row = cur.fetchone() return StrategySet(id=row[0], name=row[1], config=row[2], nick=row[3], created_at=row[4]) def get_strategy_set(self, id: str) -> StrategySet | None: self.ensure_schema() with self._connect().cursor() as cur: cur.execute( "SELECT id, name, config, nick, created_at FROM strategy_sets WHERE id = %s", (id,), ) row = cur.fetchone() if row is None: return None return StrategySet(id=row[0], name=row[1], config=row[2], nick=row[3], created_at=row[4]) def list_strategy_sets(self) -> list[StrategySet]: self.ensure_schema() with self._connect().cursor() as cur: cur.execute( "SELECT id, name, config, nick, created_at FROM strategy_sets ORDER BY created_at DESC" ) rows = cur.fetchall() return [StrategySet(id=row[0], name=row[1], config=row[2], nick=row[3], created_at=row[4]) for row in rows] def delete_strategy_set(self, id: str) -> None: self.ensure_schema() with self._connect().cursor() as cur: # 1. Check if this strategy set is referenced by any portfolio cur.execute( """ SELECT COUNT(*) FROM portfolios WHERE active_strategy_id = %s OR avabile_strategies @> ARRAY[%s] """, (id, id), ) count = cur.fetchone()[0] if count > 0: raise ValueError(f"Strategy set is in use by {count} portfolio(s) and cannot be deleted") # 2. If not in use, proceed with deletion cur.execute("DELETE FROM strategy_sets WHERE id = %s", (id,)) def save( self, portfolio: Portfolio, *, current_price: float, action: str, trade_price: float, quantity: float, fee_amount: float, tax_amount: float, total_value_after: float, strategy_set_id: str = "manual", note: str = "", ) -> None: self.ensure_schema() trade_time = datetime.now(timezone.utc) with self._connect().cursor() as cur: cur.execute( """ UPDATE portfolios SET broker = %s, status = %s, symbol = %s, cash_balance = %s, asset_quantity = %s, asset_price = %s, profit_buffer = %s, active_strategy_id = %s, avabile_strategies = %s, updated_at = %s, last_trade_at = CASE WHEN %s IN ('BUY', 'SELL') THEN %s ELSE last_trade_at END WHERE portfolio_id = %s """, ( portfolio.broker, portfolio.status, portfolio.symbol, float(portfolio.cash_balance), float(portfolio.asset_quantity), float(current_price), float(portfolio.profit_buffer), portfolio.active_strategy_id, portfolio.avabile_strategies, trade_time, action, trade_time, portfolio.portfolio_id, ), ) if action in {"BUY", "SELL"}: cur.execute( """ INSERT INTO portfolio_ledger ( portfolio_id, recorded_at, action, price, quantity, fee_amount, tax_amount, cash_balance_after, asset_quantity_after, total_value_after, profit_buffer_after, strategy_set_id, note ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) """, ( portfolio.portfolio_id, trade_time, action, float(trade_price), float(quantity), float(fee_amount), float(tax_amount), float(portfolio.cash_balance), float(portfolio.asset_quantity), float(total_value_after), float(portfolio.profit_buffer), strategy_set_id, note, ), ) def list_recent_trades(self, portfolio_id: str, limit: int = 5) -> list[LedgerEntry]: self.ensure_schema() with self._connect().cursor() as cur: cur.execute( """ SELECT portfolio_id, recorded_at, action, price, quantity, fee_amount, tax_amount, cash_balance_after, asset_quantity_after, total_value_after, profit_buffer_after, strategy_set_id, note FROM portfolio_ledger WHERE portfolio_id = %s ORDER BY recorded_at DESC LIMIT %s """, (portfolio_id, int(limit)), ) rows = cur.fetchall() return [self._row_to_ledger_entry(row) for row in rows] def reset(self, portfolio_id: str) -> Portfolio: self.ensure_schema() with self._connect().cursor() as cur: cur.execute( """ UPDATE portfolios SET cash_balance = 0, asset_quantity = 0, asset_price = 0, profit_buffer = 0, active_strategy_id = NULL, updated_at = NOW(), last_trade_at = NULL WHERE portfolio_id = %s RETURNING portfolio_id, broker, status, name, symbol, cash_balance, asset_quantity, asset_price, profit_buffer, active_strategy_id, avabile_strategies, owner_user_id, created_at, updated_at, last_trade_at """, (portfolio_id,), ) row = cur.fetchone() if row is None: raise ValueError("포트폴리오를 찾을 수 없습니다") cur.execute("DELETE FROM portfolio_ledger WHERE portfolio_id = %s", (portfolio_id,)) return self._row_to_portfolio(row) def adjust_buffer(self, portfolio_id: str, amount: float, action: str, note: str = "") -> Portfolio: self.ensure_schema() normalized = str(action).upper() if normalized not in {"INVEST", "PROFIT", "WITHDRAW"}: raise ValueError("action은 INVEST, PROFIT, WITHDRAW 중 하나여야 합니다") if amount <= 0: raise ValueError("amount는 0보다 커야 합니다") with self._connect().cursor() as cur: cur.execute(_SELECT_PORTFOLIO + "WHERE portfolio_id = %s", (portfolio_id,)) row = cur.fetchone() if row is None: raise ValueError("포트폴리오를 찾을 수 없습니다") portfolio = self._row_to_portfolio(row) updated_buffer = float(portfolio.profit_buffer) if normalized == "INVEST": updated_buffer += float(amount) else: updated_buffer -= float(amount) if updated_buffer < 0: raise ValueError("수익버퍼 잔액이 부족합니다") cur.execute( """ UPDATE portfolios SET profit_buffer = %s, updated_at = NOW() WHERE portfolio_id = %s RETURNING portfolio_id, broker, status, name, symbol, cash_balance, asset_quantity, asset_price, profit_buffer, active_strategy_id, avabile_strategies, owner_user_id, created_at, updated_at, last_trade_at """, (updated_buffer, portfolio_id), ) updated_row = cur.fetchone() total_value = float(portfolio.cash_balance) + float(portfolio.asset_quantity) * float(portfolio.price) cur.execute( """ INSERT INTO portfolio_ledger ( portfolio_id, recorded_at, action, price, quantity, fee_amount, tax_amount, cash_balance_after, asset_quantity_after, total_value_after, profit_buffer_after, strategy_set_id, note ) VALUES (%s, NOW(), %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) """, ( portfolio_id, normalized, 0.0, 0.0, 0.0, 0.0, float(portfolio.cash_balance), float(portfolio.asset_quantity), float(total_value), float(updated_buffer), "manual", note, ), ) return self._row_to_portfolio(updated_row) def transfer_buffer(self, portfolio_id: str, amount: float, direction: str, note: str = "") -> Portfolio: self.ensure_schema() normalized = str(direction).lower() if normalized not in {"to-buffer", "from-buffer"}: raise ValueError("direction은 to-buffer 또는 from-buffer 여야 합니다") if amount <= 0: raise ValueError("amount는 0보다 커야 합니다") with self._connect().cursor() as cur: cur.execute(_SELECT_PORTFOLIO + "WHERE portfolio_id = %s", (portfolio_id,)) row = cur.fetchone() if row is None: raise ValueError("포트폴리오를 찾을 수 없습니다") portfolio = self._row_to_portfolio(row) cash = float(portfolio.cash_balance) buf = float(portfolio.profit_buffer) if normalized == "to-buffer": if cash < float(amount): raise ValueError("이체 금액이 현금 잔고를 초과합니다") cash -= float(amount) buf += float(amount) ledger_action = "BUFFER_IN" else: if buf < float(amount): raise ValueError("이체 금액이 수익버퍼 잔고를 초과합니다") buf -= float(amount) cash += float(amount) ledger_action = "BUFFER_OUT" cur.execute( """ UPDATE portfolios SET cash_balance = %s, profit_buffer = %s, updated_at = NOW() WHERE portfolio_id = %s RETURNING portfolio_id, broker, status, name, symbol, cash_balance, asset_quantity, asset_price, profit_buffer, active_strategy_id, avabile_strategies, owner_user_id, created_at, updated_at, last_trade_at """, (cash, buf, portfolio_id), ) updated_row = cur.fetchone() total_value = float(cash) + float(portfolio.asset_quantity) * float(portfolio.price) cur.execute( """ INSERT INTO portfolio_ledger ( portfolio_id, recorded_at, action, price, quantity, fee_amount, tax_amount, cash_balance_after, asset_quantity_after, total_value_after, profit_buffer_after, strategy_set_id, note ) VALUES (%s, NOW(), %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) """, ( portfolio_id, ledger_action, 0.0, 0.0, 0.0, 0.0, float(cash), float(portfolio.asset_quantity), float(total_value), float(buf), "manual", note, ), ) return self._row_to_portfolio(updated_row) def delete(self, portfolio_id: str, *, clear_ledger: bool = True) -> None: self.ensure_schema() with self._connect().cursor() as cur: cur.execute("SELECT portfolio_id FROM portfolios WHERE portfolio_id = %s", (portfolio_id,)) if cur.fetchone() is None: raise ValueError("포트폴리오를 찾을 수 없습니다") if clear_ledger: cur.execute("DELETE FROM portfolio_ledger WHERE portfolio_id = %s", (portfolio_id,)) # Comments are always removed with the portfolio (no FK constraint, handled in application layer) cur.execute("DELETE FROM portfolio_comments WHERE portfolio_id = %s", (portfolio_id,)) cur.execute("DELETE FROM portfolios WHERE portfolio_id = %s", (portfolio_id,)) @staticmethod def _row_to_comment(row: Any) -> PortfolioComment: if row is None: raise ValueError("댓글을 찾을 수 없습니다") return PortfolioComment( id=row[0], portfolio_id=row[1], content=row[2], created_at=row[3], updated_at=row[4], ) def create_comment(self, portfolio_id: str, *, content: str) -> PortfolioComment: self.ensure_schema() with self._connect().cursor() as cur: cur.execute("SELECT 1 FROM portfolios WHERE portfolio_id = %s", (portfolio_id,)) if cur.fetchone() is None: raise ValueError("포트폴리오를 찾을 수 없습니다") comment_id = str(uuid.uuid4()) cur.execute( """ INSERT INTO portfolio_comments (id, portfolio_id, content) VALUES (%s, %s, %s) RETURNING id, portfolio_id, content, created_at, updated_at """, (comment_id, portfolio_id, content), ) row = cur.fetchone() return self._row_to_comment(row) def list_comments(self, portfolio_id: str) -> list[PortfolioComment]: self.ensure_schema() with self._connect().cursor() as cur: cur.execute("SELECT 1 FROM portfolios WHERE portfolio_id = %s", (portfolio_id,)) if cur.fetchone() is None: raise ValueError("포트폴리오를 찾을 수 없습니다") cur.execute( """ SELECT id, portfolio_id, content, created_at, updated_at FROM portfolio_comments WHERE portfolio_id = %s ORDER BY created_at DESC """, (portfolio_id,), ) rows = cur.fetchall() return [self._row_to_comment(row) for row in rows] def update_comment(self, portfolio_id: str, comment_id: str, *, content: str) -> PortfolioComment: self.ensure_schema() with self._connect().cursor() as cur: cur.execute("SELECT 1 FROM portfolios WHERE portfolio_id = %s", (portfolio_id,)) if cur.fetchone() is None: raise ValueError("포트폴리오를 찾을 수 없습니다") cur.execute( """ UPDATE portfolio_comments SET content = %s, updated_at = NOW() WHERE id = %s AND portfolio_id = %s RETURNING id, portfolio_id, content, created_at, updated_at """, (content, comment_id, portfolio_id), ) row = cur.fetchone() if row is None: raise ValueError("댓글을 찾을 수 없습니다") return self._row_to_comment(row) def delete_comment(self, portfolio_id: str, comment_id: str) -> None: self.ensure_schema() with self._connect().cursor() as cur: cur.execute("SELECT 1 FROM portfolios WHERE portfolio_id = %s", (portfolio_id,)) if cur.fetchone() is None: raise ValueError("포트폴리오를 찾을 수 없습니다") cur.execute( "DELETE FROM portfolio_comments WHERE id = %s AND portfolio_id = %s RETURNING id", (comment_id, portfolio_id), ) if cur.fetchone() is None: raise ValueError("댓글을 찾을 수 없습니다") def database_exists(config: PortfolioDBConfig, *, maintenance_db: str = "postgres") -> bool: with psycopg.connect( host=config.host, port=config.port, dbname=maintenance_db, user=config.user, password=config.password, sslmode=config.sslmode, connect_timeout=config.connect_timeout_sec, application_name=config.application_name, autocommit=True, ) as conn: with conn.cursor() as cur: cur.execute("SELECT 1 FROM pg_database WHERE datname = %s", (config.dbname,)) return cur.fetchone() is not None def create_database_if_missing(config: PortfolioDBConfig, *, maintenance_db: str = "postgres") -> bool: with psycopg.connect( host=config.host, port=config.port, dbname=maintenance_db, user=config.user, password=config.password, sslmode=config.sslmode, connect_timeout=config.connect_timeout_sec, application_name=config.application_name, autocommit=True, ) as conn: with conn.cursor() as cur: cur.execute("SELECT 1 FROM pg_database WHERE datname = %s", (config.dbname,)) if cur.fetchone() is not None: return False cur.execute(sql.SQL("CREATE DATABASE {}") .format(sql.Identifier(config.dbname))) return True