126 lines
3.8 KiB
Python
126 lines
3.8 KiB
Python
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_session(row):
|
|
access_token = decrypt_value(row[1]) if row[1] else None
|
|
request_token = decrypt_value(row[2]) if row[2] else None
|
|
return {
|
|
"api_key": row[0],
|
|
"access_token": access_token,
|
|
"request_token": request_token,
|
|
"user_name": row[3],
|
|
"broker_user_id": row[4],
|
|
"linked_at": row[5],
|
|
}
|
|
|
|
|
|
def get_session(user_id: str):
|
|
with db_transaction() as cur:
|
|
cur.execute(
|
|
"""
|
|
SELECT api_key, access_token, request_token, user_name, broker_user_id, linked_at
|
|
FROM zerodha_session
|
|
WHERE user_id = %s
|
|
ORDER BY linked_at DESC NULLS LAST, id DESC
|
|
LIMIT 1
|
|
""",
|
|
(user_id,),
|
|
)
|
|
row = cur.fetchone()
|
|
if row:
|
|
return _row_to_session(row)
|
|
|
|
with db_transaction() as cur:
|
|
cur.execute(
|
|
"""
|
|
SELECT broker, connected, access_token, api_key, user_name, broker_user_id, connected_at
|
|
FROM user_broker
|
|
WHERE user_id = %s
|
|
""",
|
|
(user_id,),
|
|
)
|
|
row = cur.fetchone()
|
|
if not row:
|
|
return None
|
|
broker, connected, access_token, api_key, user_name, broker_user_id, connected_at = row
|
|
if not connected or not access_token or not api_key:
|
|
return None
|
|
if (broker or "").strip().upper() != "ZERODHA":
|
|
return None
|
|
return {
|
|
"api_key": api_key,
|
|
"access_token": decrypt_value(access_token),
|
|
"request_token": None,
|
|
"user_name": user_name,
|
|
"broker_user_id": broker_user_id,
|
|
"linked_at": connected_at,
|
|
}
|
|
|
|
|
|
def set_session(user_id: str, data: dict):
|
|
access_token = data.get("access_token")
|
|
request_token = data.get("request_token")
|
|
linked_at = datetime.now(timezone.utc)
|
|
with db_transaction() as cur:
|
|
cur.execute(
|
|
"""
|
|
INSERT INTO zerodha_session (
|
|
user_id, linked_at, api_key, access_token, request_token, user_name, broker_user_id
|
|
)
|
|
VALUES (%s, %s, %s, %s, %s, %s, %s)
|
|
RETURNING linked_at
|
|
""",
|
|
(
|
|
user_id,
|
|
linked_at,
|
|
data.get("api_key"),
|
|
encrypt_value(access_token) if access_token else None,
|
|
encrypt_value(request_token) if request_token else None,
|
|
data.get("user_name"),
|
|
data.get("broker_user_id"),
|
|
),
|
|
)
|
|
linked_at = cur.fetchone()[0]
|
|
return {
|
|
**data,
|
|
"user_id": user_id,
|
|
"linked_at": linked_at,
|
|
"access_token": access_token,
|
|
"request_token": request_token,
|
|
}
|
|
|
|
|
|
def store_request_token(user_id: str, request_token: str):
|
|
with db_transaction() as cur:
|
|
cur.execute(
|
|
"""
|
|
INSERT INTO zerodha_request_token (user_id, request_token)
|
|
VALUES (%s, %s)
|
|
ON CONFLICT (user_id)
|
|
DO UPDATE SET request_token = EXCLUDED.request_token
|
|
""",
|
|
(user_id, encrypt_value(request_token)),
|
|
)
|
|
|
|
|
|
def consume_request_token(user_id: str):
|
|
with db_transaction() as cur:
|
|
cur.execute(
|
|
"SELECT request_token FROM zerodha_request_token WHERE user_id = %s",
|
|
(user_id,),
|
|
)
|
|
row = cur.fetchone()
|
|
if not row:
|
|
return None
|
|
cur.execute("DELETE FROM zerodha_request_token WHERE user_id = %s", (user_id,))
|
|
return decrypt_value(row[0])
|
|
|
|
|
|
def clear_session(user_id: str):
|
|
with db_transaction() as cur:
|
|
cur.execute("DELETE FROM zerodha_session WHERE user_id = %s", (user_id,))
|
|
cur.execute("DELETE FROM zerodha_request_token WHERE user_id = %s", (user_id,))
|