Files
portfolio-hub/app/repositories/postgres.py
T

684 lines
26 KiB
Python

"""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 LedgerEntry, Portfolio, StrategySet
_SELECT_PORTFOLIO = """
SELECT
portfolio_id,
broker,
status,
symbol,
cash_balance,
asset_quantity,
asset_price,
profit_buffer,
active_strategy_id,
avabile_strategies,
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 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',
symbol TEXT NOT NULL,
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(
"""
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("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_portfolio(row: Any) -> Portfolio:
if row is None:
raise ValueError("포트폴리오를 찾을 수 없습니다")
raw_available = row[9] if len(row) > 9 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",
symbol=row[3],
cash_balance=float(row[4]),
asset_quantity=float(row[5]),
price=float(row[6]),
profit_buffer=float(row[7]),
active_strategy_id=row[8] if len(row) > 8 else None,
avabile_strategies=avabile_strategies,
created_at=row[10] if len(row) > 10 else None,
updated_at=row[11] if len(row) > 11 else None,
last_trade_at=row[12] if len(row) > 12 else None,
)
@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 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()
return self._row_to_portfolio(row)
def get_or_create(
self,
portfolio_id: str,
symbol: str,
*,
broker: str = "kis",
status: str = "enabled",
active_strategy_id: str | None = None,
avabile_strategies: list[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, symbol,
cash_balance, asset_quantity, asset_price, profit_buffer,
active_strategy_id, avabile_strategies
)
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
ON CONFLICT (portfolio_id) DO NOTHING
""",
(
portfolio_id,
broker,
status,
symbol,
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,
*,
broker: str,
status: 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, symbol,
cash_balance, asset_quantity, asset_price, profit_buffer,
active_strategy_id, avabile_strategies
) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
RETURNING
portfolio_id, broker, status, symbol,
cash_balance, asset_quantity, asset_price, profit_buffer,
active_strategy_id, avabile_strategies, created_at, updated_at, last_trade_at
""",
(
portfolio_id,
broker,
status,
symbol,
float(cash_balance),
float(asset_quantity),
float(asset_price),
float(profit_buffer),
active_strategy_id,
avabile_strategies,
),
)
row = cur.fetchone()
return self._row_to_portfolio(row)
def update_portfolio(self, portfolio_id: str, **fields: object) -> Portfolio:
self.ensure_schema()
allowed = {
"broker": "broker",
"status": "status",
"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, symbol,
cash_balance, asset_quantity, asset_price, profit_buffer,
active_strategy_id, avabile_strategies, 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) -> 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 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()
return [self._row_to_portfolio(row) for row in rows]
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 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, symbol,
cash_balance, asset_quantity, asset_price, profit_buffer,
active_strategy_id, avabile_strategies, 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, symbol,
cash_balance, asset_quantity, asset_price, profit_buffer,
active_strategy_id, avabile_strategies, 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, symbol,
cash_balance, asset_quantity, asset_price, profit_buffer,
active_strategy_id, avabile_strategies, 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,))
cur.execute("DELETE FROM portfolios WHERE portfolio_id = %s", (portfolio_id,))
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