SIP_GoldBees_Backend/backend/app/broker_store.py

from datetime import datetime, timezone

from app.services.crypto_service import decrypt_value, encrypt_value
from app.services.db import db_transaction


def _row_to_entry(row):
    (
        user_id,
        broker,
        connected,
        access_token,
        connected_at,
        api_key,
        api_secret,
        user_name,
        broker_user_id,
        auth_state,
        pending_broker,
        pending_api_key,
        pending_api_secret,
        pending_started_at,
    ) = row
    entry = {
        "broker": broker,
        "connected": bool(connected),
        "connected_at": connected_at,
        "api_key": api_key,
        "auth_state": auth_state,
        "user_name": user_name,
        "broker_user_id": broker_user_id,
    }
    if pending_broker or pending_api_key or pending_api_secret or pending_started_at:
        pending = {
            "broker": pending_broker,
            "api_key": pending_api_key,
            "api_secret": decrypt_value(pending_api_secret)
            if pending_api_secret
            else None,
            "started_at": pending_started_at,
        }
        entry["pending"] = pending
    return entry


def load_user_brokers():
    with db_transaction() as cur:
        cur.execute(
            """
            SELECT user_id, broker, connected, access_token, connected_at,
                   api_key, api_secret, user_name, broker_user_id, auth_state,
                   pending_broker, pending_api_key, pending_api_secret, pending_started_at
            FROM user_broker
            """
        )
        rows = cur.fetchall()
    return {row[0]: _row_to_entry(row) for row in rows}


def save_user_brokers(data):
    with db_transaction() as cur:
        for user_id, entry in data.items():
            cur.execute(
                """
                INSERT INTO user_broker (
                    user_id, broker, connected, access_token, connected_at,
                    api_key, api_secret, user_name, broker_user_id, auth_state,
                    pending_broker, pending_api_key, pending_api_secret, pending_started_at
                )
                VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
                ON CONFLICT (user_id)
                DO UPDATE SET
                    broker = EXCLUDED.broker,
                    connected = EXCLUDED.connected,
                    access_token = EXCLUDED.access_token,
                    connected_at = EXCLUDED.connected_at,
                    api_key = EXCLUDED.api_key,
                    api_secret = EXCLUDED.api_secret,
                    user_name = EXCLUDED.user_name,
                    broker_user_id = EXCLUDED.broker_user_id,
                    auth_state = EXCLUDED.auth_state,
                    pending_broker = EXCLUDED.pending_broker,
                    pending_api_key = EXCLUDED.pending_api_key,
                    pending_api_secret = EXCLUDED.pending_api_secret,
                    pending_started_at = EXCLUDED.pending_started_at
                """,
                (
                    user_id,
                    entry.get("broker"),
                    bool(entry.get("connected")),
                    encrypt_value(entry.get("access_token"))
                    if entry.get("access_token")
                    else None,
                    entry.get("connected_at"),
                    entry.get("api_key"),
                    encrypt_value(entry.get("api_secret"))
                    if entry.get("api_secret")
                    else None,
                    entry.get("user_name"),
                    entry.get("broker_user_id"),
                    entry.get("auth_state"),
                    (entry.get("pending") or {}).get("broker"),
                    (entry.get("pending") or {}).get("api_key"),
                    encrypt_value((entry.get("pending") or {}).get("api_secret"))
                    if (entry.get("pending") or {}).get("api_secret")
                    else None,
                    (entry.get("pending") or {}).get("started_at"),
                ),
            )


def now_utc():
    return datetime.now(timezone.utc)


def get_user_broker(user_id: str):
    with db_transaction() as cur:
        cur.execute(
            """
            SELECT user_id, broker, connected, access_token, connected_at,
                   api_key, api_secret, user_name, broker_user_id, auth_state,
                   pending_broker, pending_api_key, pending_api_secret, pending_started_at
            FROM user_broker
            WHERE user_id = %s
            """,
            (user_id,),
        )
        row = cur.fetchone()
    if not row:
        return None
    return _row_to_entry(row)


def clear_user_broker(user_id: str):
    with db_transaction() as cur:
        cur.execute("DELETE FROM user_broker WHERE user_id = %s", (user_id,))


def set_pending_broker(user_id: str, broker: str, api_key: str, api_secret: str):
    started_at = now_utc()
    with db_transaction() as cur:
        cur.execute(
            """
            INSERT INTO user_broker (
                user_id, pending_broker, pending_api_key, pending_api_secret, pending_started_at,
                api_key, api_secret, auth_state
            )
            VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
            ON CONFLICT (user_id)
            DO UPDATE SET
                pending_broker = EXCLUDED.pending_broker,
                pending_api_key = EXCLUDED.pending_api_key,
                pending_api_secret = EXCLUDED.pending_api_secret,
                pending_started_at = EXCLUDED.pending_started_at,
                api_key = EXCLUDED.api_key,
                api_secret = EXCLUDED.api_secret,
                auth_state = EXCLUDED.auth_state
            """,
            (
                user_id,
                broker,
                api_key,
                encrypt_value(api_secret),
                started_at,
                api_key,
                encrypt_value(api_secret),
                "PENDING",
            ),
        )
    return {
        "broker": broker,
        "api_key": api_key,
        "api_secret": api_secret,
        "started_at": started_at,
    }


def get_pending_broker(user_id: str):
    with db_transaction() as cur:
        cur.execute(
            """
            SELECT pending_broker, pending_api_key, pending_api_secret, pending_started_at
            FROM user_broker
            WHERE user_id = %s
            """,
            (user_id,),
        )
        row = cur.fetchone()
    if not row:
        return None
    if not row[0] or not row[1] or not row[2]:
        return None
    return {
        "broker": row[0],
        "api_key": row[1],
        "api_secret": decrypt_value(row[2]),
        "started_at": row[3],
    }


def get_broker_credentials(user_id: str):
    with db_transaction() as cur:
        cur.execute(
            """
            SELECT api_key, api_secret, pending_api_key, pending_api_secret
            FROM user_broker
            WHERE user_id = %s
            """,
            (user_id,),
        )
        row = cur.fetchone()
    if not row:
        return None
    api_key, api_secret, pending_key, pending_secret = row
    key = api_key or pending_key
    secret = api_secret or pending_secret
    if not key or not secret:
        return None
    return {
        "api_key": key,
        "api_secret": decrypt_value(secret),
    }


def set_broker_auth_state(user_id: str, auth_state: str):
    with db_transaction() as cur:
        cur.execute(
            """
            UPDATE user_broker
            SET auth_state = %s
            WHERE user_id = %s
            """,
            (auth_state, user_id),
        )


def set_connected_broker(
    user_id: str,
    broker: str,
    access_token: str,
    api_key: str | None = None,
    api_secret: str | None = None,
    user_name: str | None = None,
    broker_user_id: str | None = None,
    auth_state: str | None = None,
):
    connected_at = now_utc()
    with db_transaction() as cur:
        cur.execute(
            """
            INSERT INTO user_broker (
                user_id, broker, connected, access_token, connected_at,
                api_key, api_secret, user_name, broker_user_id, auth_state,
                pending_broker, pending_api_key, pending_api_secret, pending_started_at
            )
            VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, NULL, NULL, NULL, NULL)
            ON CONFLICT (user_id)
            DO UPDATE SET
                broker = EXCLUDED.broker,
                connected = EXCLUDED.connected,
                access_token = EXCLUDED.access_token,
                connected_at = EXCLUDED.connected_at,
                api_key = EXCLUDED.api_key,
                api_secret = EXCLUDED.api_secret,
                user_name = EXCLUDED.user_name,
                broker_user_id = EXCLUDED.broker_user_id,
                auth_state = EXCLUDED.auth_state,
                pending_broker = NULL,
                pending_api_key = NULL,
                pending_api_secret = NULL,
                pending_started_at = NULL
            """,
            (
                user_id,
                broker,
                True,
                encrypt_value(access_token),
                connected_at,
                api_key,
                encrypt_value(api_secret) if api_secret else None,
                user_name,
                broker_user_id,
                auth_state,
            ),
        )
    return {
        "broker": broker,
        "connected": True,
        "access_token": access_token,
        "connected_at": connected_at,
        "api_key": api_key,
        "api_secret": api_secret,
        "user_name": user_name,
        "broker_user_id": broker_user_id,
        "auth_state": auth_state,
    }