177 lines
4.4 KiB
Python
177 lines
4.4 KiB
Python
import threading
|
|
from datetime import datetime, timezone
|
|
from uuid import uuid4
|
|
|
|
from psycopg2.extras import Json
|
|
|
|
from app.services.db import run_with_retry
|
|
|
|
_DEFAULT_USER_ID = None
|
|
_DEFAULT_LOCK = threading.Lock()
|
|
|
|
|
|
def _utc_now():
|
|
return datetime.now(timezone.utc)
|
|
|
|
|
|
def get_default_user_id():
|
|
global _DEFAULT_USER_ID
|
|
if _DEFAULT_USER_ID:
|
|
return _DEFAULT_USER_ID
|
|
|
|
def _op(cur, _conn):
|
|
cur.execute("SELECT id FROM app_user ORDER BY username LIMIT 1")
|
|
row = cur.fetchone()
|
|
return row[0] if row else None
|
|
|
|
user_id = run_with_retry(_op)
|
|
if user_id:
|
|
with _DEFAULT_LOCK:
|
|
_DEFAULT_USER_ID = user_id
|
|
return user_id
|
|
|
|
|
|
def _default_run_id(user_id: str) -> str:
|
|
return f"default_{user_id}"
|
|
|
|
|
|
def ensure_default_run(user_id: str):
|
|
run_id = _default_run_id(user_id)
|
|
|
|
def _op(cur, _conn):
|
|
now = _utc_now()
|
|
cur.execute(
|
|
"""
|
|
INSERT INTO strategy_run (
|
|
run_id, user_id, created_at, started_at, stopped_at, status, strategy, mode, broker, meta
|
|
)
|
|
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
|
|
ON CONFLICT (run_id) DO NOTHING
|
|
""",
|
|
(
|
|
run_id,
|
|
user_id,
|
|
now,
|
|
None,
|
|
None,
|
|
"STOPPED",
|
|
None,
|
|
None,
|
|
None,
|
|
Json({}),
|
|
),
|
|
)
|
|
return run_id
|
|
|
|
return run_with_retry(_op)
|
|
|
|
|
|
def get_active_run_id(user_id: str):
|
|
def _op(cur, _conn):
|
|
cur.execute(
|
|
"""
|
|
SELECT run_id
|
|
FROM strategy_run
|
|
WHERE user_id = %s AND status = 'RUNNING'
|
|
ORDER BY created_at DESC
|
|
LIMIT 1
|
|
""",
|
|
(user_id,),
|
|
)
|
|
row = cur.fetchone()
|
|
if row:
|
|
return row[0]
|
|
cur.execute(
|
|
"""
|
|
SELECT run_id
|
|
FROM strategy_run
|
|
WHERE user_id = %s
|
|
ORDER BY created_at DESC
|
|
LIMIT 1
|
|
""",
|
|
(user_id,),
|
|
)
|
|
row = cur.fetchone()
|
|
if row:
|
|
return row[0]
|
|
return None
|
|
|
|
run_id = run_with_retry(_op)
|
|
if run_id:
|
|
return run_id
|
|
return ensure_default_run(user_id)
|
|
|
|
|
|
def get_running_run_id(user_id: str):
|
|
def _op(cur, _conn):
|
|
cur.execute(
|
|
"""
|
|
SELECT run_id
|
|
FROM strategy_run
|
|
WHERE user_id = %s AND status = 'RUNNING'
|
|
ORDER BY created_at DESC
|
|
LIMIT 1
|
|
""",
|
|
(user_id,),
|
|
)
|
|
row = cur.fetchone()
|
|
return row[0] if row else None
|
|
|
|
return run_with_retry(_op)
|
|
|
|
|
|
def create_strategy_run(user_id: str, strategy: str | None, mode: str | None, broker: str | None, meta: dict | None):
|
|
run_id = str(uuid4())
|
|
|
|
def _op(cur, _conn):
|
|
now = _utc_now()
|
|
cur.execute(
|
|
"""
|
|
INSERT INTO strategy_run (
|
|
run_id, user_id, created_at, started_at, stopped_at, status, strategy, mode, broker, meta
|
|
)
|
|
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
|
|
""",
|
|
(
|
|
run_id,
|
|
user_id,
|
|
now,
|
|
now,
|
|
None,
|
|
"RUNNING",
|
|
strategy,
|
|
mode,
|
|
broker,
|
|
Json(meta or {}),
|
|
),
|
|
)
|
|
return run_id
|
|
|
|
return run_with_retry(_op)
|
|
|
|
|
|
def update_run_status(user_id: str, run_id: str, status: str, meta: dict | None = None):
|
|
def _op(cur, _conn):
|
|
now = _utc_now()
|
|
if status == "RUNNING":
|
|
cur.execute(
|
|
"""
|
|
UPDATE strategy_run
|
|
SET status = %s, started_at = COALESCE(started_at, %s), meta = COALESCE(meta, '{}'::jsonb) || %s
|
|
WHERE run_id = %s AND user_id = %s
|
|
""",
|
|
(status, now, Json(meta or {}), run_id, user_id),
|
|
)
|
|
else:
|
|
cur.execute(
|
|
"""
|
|
UPDATE strategy_run
|
|
SET status = %s, stopped_at = %s, meta = COALESCE(meta, '{}'::jsonb) || %s
|
|
WHERE run_id = %s AND user_id = %s
|
|
""",
|
|
(status, now, Json(meta or {}), run_id, user_id),
|
|
)
|
|
return True
|
|
|
|
return run_with_retry(_op)
|