507 lines
11 KiB
Python
507 lines
11 KiB
Python
from fastapi import (
|
|
FastAPI,
|
|
HTTPException,
|
|
WebSocket,
|
|
WebSocketDisconnect,
|
|
Depends,
|
|
Header,
|
|
Query,
|
|
status,
|
|
)
|
|
from datetime import datetime, timedelta, timezone
|
|
from pydantic import BaseModel
|
|
import jwt
|
|
from jwt import ExpiredSignatureError, InvalidTokenError
|
|
import psycopg2
|
|
from psycopg2.extras import RealDictCursor
|
|
import asyncio
|
|
|
|
|
|
app = FastAPI()
|
|
|
|
|
|
# ============================================================
|
|
# JWT CONFIG
|
|
# ============================================================
|
|
|
|
SECRET_KEY = "my-super-secret-key-change-this"
|
|
ALGORITHM = "HS256"
|
|
ACCESS_TOKEN_EXPIRE_MINUTES = 60
|
|
|
|
|
|
# ============================================================
|
|
# REQUEST MODELS
|
|
# ============================================================
|
|
|
|
class RegisterNodeRequest(BaseModel):
|
|
node_name: str
|
|
host: str
|
|
port: int = 5444
|
|
dbname: str = "postgres"
|
|
user: str = "postgres"
|
|
password: str
|
|
|
|
|
|
class LinkNodeRequest(BaseModel):
|
|
subscription_name: str
|
|
|
|
subscriber_host: str
|
|
subscriber_port: int = 5444
|
|
|
|
provider_host: str
|
|
provider_port: int = 5444
|
|
|
|
dbname: str = "postgres"
|
|
user: str = "postgres"
|
|
password: str
|
|
|
|
|
|
# ============================================================
|
|
# JWT TOKEN CREATION
|
|
# ============================================================
|
|
|
|
@app.post("/get-token")
|
|
def get_token():
|
|
|
|
payload = {
|
|
"sub": "frontend",
|
|
"iat": datetime.now(timezone.utc),
|
|
"exp": datetime.now(timezone.utc)
|
|
+ timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES),
|
|
}
|
|
|
|
token = jwt.encode(
|
|
payload,
|
|
SECRET_KEY,
|
|
algorithm=ALGORITHM,
|
|
)
|
|
|
|
return {
|
|
"access_token": token,
|
|
"token_type": "bearer",
|
|
"expires_in": ACCESS_TOKEN_EXPIRE_MINUTES * 60,
|
|
}
|
|
|
|
|
|
# ============================================================
|
|
# JWT VERIFICATION
|
|
# ============================================================
|
|
|
|
def verify_token(
|
|
authorization: str = Header(None),
|
|
):
|
|
|
|
if not authorization:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_401_UNAUTHORIZED,
|
|
detail="Authorization header missing",
|
|
headers={
|
|
"WWW-Authenticate": "Bearer"
|
|
},
|
|
)
|
|
|
|
parts = authorization.split()
|
|
|
|
if len(parts) != 2:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_401_UNAUTHORIZED,
|
|
detail="Invalid Authorization header",
|
|
)
|
|
|
|
scheme = parts[0]
|
|
token = parts[1]
|
|
|
|
if scheme.lower() != "bearer":
|
|
raise HTTPException(
|
|
status_code=status.HTTP_401_UNAUTHORIZED,
|
|
detail="Authorization scheme must be Bearer",
|
|
)
|
|
|
|
try:
|
|
|
|
payload = jwt.decode(
|
|
token,
|
|
SECRET_KEY,
|
|
algorithms=[ALGORITHM],
|
|
)
|
|
|
|
return payload
|
|
|
|
except ExpiredSignatureError:
|
|
|
|
raise HTTPException(
|
|
status_code=status.HTTP_401_UNAUTHORIZED,
|
|
detail="Token expired",
|
|
)
|
|
|
|
except InvalidTokenError:
|
|
|
|
raise HTTPException(
|
|
status_code=status.HTTP_401_UNAUTHORIZED,
|
|
detail="Invalid token",
|
|
)
|
|
|
|
|
|
# ============================================================
|
|
# REGISTER NODE
|
|
# JWT PROTECTED
|
|
# ============================================================
|
|
|
|
@app.post("/register-node")
|
|
def register_node(
|
|
data: RegisterNodeRequest,
|
|
token_data: dict = Depends(verify_token),
|
|
):
|
|
|
|
conn = None
|
|
cur = None
|
|
|
|
try:
|
|
|
|
conn = psycopg2.connect(
|
|
host=data.host,
|
|
port=data.port,
|
|
dbname=data.dbname,
|
|
user=data.user,
|
|
password=data.password,
|
|
)
|
|
|
|
cur = conn.cursor()
|
|
|
|
# Create Spock extension
|
|
cur.execute("""
|
|
CREATE EXTENSION IF NOT EXISTS spock;
|
|
""")
|
|
|
|
# Check version
|
|
cur.execute("""
|
|
SELECT extname, extversion
|
|
FROM pg_extension
|
|
WHERE extname = 'spock';
|
|
""")
|
|
|
|
spock = cur.fetchone()
|
|
|
|
dsn = (
|
|
f"host={data.host} "
|
|
f"port={data.port} "
|
|
f"dbname={data.dbname} "
|
|
f"user={data.user} "
|
|
f"password={data.password}"
|
|
)
|
|
|
|
# Register Spock node
|
|
cur.execute("""
|
|
SELECT spock.node_create(
|
|
node_name := %s,
|
|
dsn := %s
|
|
);
|
|
""", (
|
|
data.node_name,
|
|
dsn,
|
|
))
|
|
|
|
node_id = cur.fetchone()[0]
|
|
|
|
conn.commit()
|
|
|
|
cur.execute("""
|
|
SELECT node_id, node_name
|
|
FROM spock.node;
|
|
""")
|
|
|
|
rows = cur.fetchall()
|
|
|
|
nodes = [
|
|
{
|
|
"node_id": row[0],
|
|
"node_name": row[1],
|
|
}
|
|
for row in rows
|
|
]
|
|
|
|
return {
|
|
"status": "success",
|
|
"message": "Node registered",
|
|
"node_id": node_id,
|
|
"spock_version": spock[1],
|
|
"nodes": nodes,
|
|
"authenticated_as": token_data["sub"],
|
|
}
|
|
|
|
except Exception as e:
|
|
|
|
if conn:
|
|
conn.rollback()
|
|
|
|
raise HTTPException(
|
|
status_code=500,
|
|
detail=str(e),
|
|
)
|
|
|
|
finally:
|
|
|
|
if cur:
|
|
cur.close()
|
|
|
|
if conn:
|
|
conn.close()
|
|
|
|
|
|
# ============================================================
|
|
# LINK NODES
|
|
# JWT PROTECTED
|
|
# ============================================================
|
|
|
|
@app.post("/link-nodes")
|
|
def link_nodes(
|
|
data: LinkNodeRequest,
|
|
token_data: dict = Depends(verify_token),
|
|
):
|
|
|
|
conn = None
|
|
cur = None
|
|
|
|
try:
|
|
|
|
conn = psycopg2.connect(
|
|
host=data.subscriber_host,
|
|
port=data.subscriber_port,
|
|
dbname=data.dbname,
|
|
user=data.user,
|
|
password=data.password,
|
|
)
|
|
|
|
cur = conn.cursor()
|
|
|
|
provider_dsn = (
|
|
f"host={data.provider_host} "
|
|
f"port={data.provider_port} "
|
|
f"dbname={data.dbname} "
|
|
f"user={data.user} "
|
|
f"password={data.password}"
|
|
)
|
|
|
|
cur.execute("""
|
|
SELECT spock.sub_create(
|
|
subscription_name := %s,
|
|
provider_dsn := %s,
|
|
replication_sets := ARRAY[
|
|
'default',
|
|
'default_insert_only',
|
|
'ddl_sql'
|
|
],
|
|
synchronize_structure := false,
|
|
synchronize_data := false
|
|
);
|
|
""", (
|
|
data.subscription_name,
|
|
provider_dsn,
|
|
))
|
|
|
|
subscription_id = cur.fetchone()[0]
|
|
|
|
conn.commit()
|
|
|
|
return {
|
|
"status": "success",
|
|
"message": "Nodes linked successfully",
|
|
"subscription_id": subscription_id,
|
|
"subscription_name": data.subscription_name,
|
|
"subscriber": data.subscriber_host,
|
|
"provider": data.provider_host,
|
|
"authenticated_as": token_data["sub"],
|
|
}
|
|
|
|
except Exception as e:
|
|
|
|
if conn:
|
|
conn.rollback()
|
|
|
|
raise HTTPException(
|
|
status_code=500,
|
|
detail=str(e),
|
|
)
|
|
|
|
finally:
|
|
|
|
if cur:
|
|
cur.close()
|
|
|
|
if conn:
|
|
conn.close()
|
|
|
|
|
|
# ============================================================
|
|
# DATABASE CONFIG
|
|
# ============================================================
|
|
|
|
DB_CONFIG = {
|
|
"host": "172.16.10.123",
|
|
"port": 5444,
|
|
"dbname": "postgres",
|
|
"user": "postgres",
|
|
"password": "postgres",
|
|
}
|
|
|
|
|
|
# ============================================================
|
|
# SPOCK STATUS
|
|
# ============================================================
|
|
|
|
def get_spock_status():
|
|
|
|
conn = psycopg2.connect(**DB_CONFIG)
|
|
|
|
cursor = conn.cursor(
|
|
cursor_factory=RealDictCursor
|
|
)
|
|
|
|
cursor.execute("""
|
|
SELECT
|
|
subscription_name,
|
|
status,
|
|
provider_node
|
|
FROM spock.sub_show_status();
|
|
""")
|
|
|
|
rows = cursor.fetchall()
|
|
|
|
cursor.close()
|
|
conn.close()
|
|
|
|
return rows
|
|
|
|
|
|
# ============================================================
|
|
# WEBSOCKET JWT VERIFICATION
|
|
# ============================================================
|
|
|
|
def verify_websocket_token(token: str):
|
|
|
|
try:
|
|
|
|
payload = jwt.decode(
|
|
token,
|
|
SECRET_KEY,
|
|
algorithms=[ALGORITHM],
|
|
)
|
|
|
|
return payload
|
|
|
|
except ExpiredSignatureError:
|
|
return None
|
|
|
|
except InvalidTokenError:
|
|
return None
|
|
|
|
|
|
# ============================================================
|
|
# WEBSOCKET
|
|
#
|
|
# ws://localhost:8000/ws/spock-status?token=JWT_TOKEN
|
|
# ============================================================
|
|
|
|
@app.websocket("/ws/spock-status")
|
|
async def websocket_spock_status(
|
|
websocket: WebSocket,
|
|
token: str = Query(...),
|
|
):
|
|
|
|
token_data = verify_websocket_token(token)
|
|
|
|
if token_data is None:
|
|
|
|
await websocket.close(
|
|
code=1008,
|
|
reason="Invalid or expired JWT",
|
|
)
|
|
|
|
return
|
|
|
|
await websocket.accept()
|
|
|
|
try:
|
|
|
|
while True:
|
|
|
|
try:
|
|
|
|
subscription_status = (
|
|
await asyncio.to_thread(
|
|
get_spock_status
|
|
)
|
|
)
|
|
|
|
await websocket.send_json({
|
|
"database_reachable": True,
|
|
"subscriptions": subscription_status,
|
|
})
|
|
|
|
except Exception as e:
|
|
|
|
await websocket.send_json({
|
|
"database_reachable": False,
|
|
"subscriptions": [],
|
|
"error": str(e),
|
|
})
|
|
|
|
await asyncio.sleep(2)
|
|
|
|
except WebSocketDisconnect:
|
|
|
|
print(
|
|
f"WebSocket client disconnected: "
|
|
f"{token_data.get('sub')}"
|
|
)
|
|
|
|
|
|
# ============================================================
|
|
# GET SUBSCRIPTIONS
|
|
# JWT PROTECTED
|
|
# ============================================================
|
|
|
|
@app.get("/spock/subscriptions")
|
|
def get_spock_subscription_status(
|
|
token_data: dict = Depends(verify_token),
|
|
):
|
|
|
|
conn = None
|
|
cursor = None
|
|
|
|
try:
|
|
|
|
conn = psycopg2.connect(**DB_CONFIG)
|
|
|
|
cursor = conn.cursor(
|
|
cursor_factory=RealDictCursor
|
|
)
|
|
|
|
cursor.execute("""
|
|
SELECT *
|
|
FROM spock.sub_show_status();
|
|
""")
|
|
|
|
rows = cursor.fetchall()
|
|
|
|
return {
|
|
"success": True,
|
|
"count": len(rows),
|
|
"subscriptions": rows,
|
|
"authenticated_as": token_data["sub"],
|
|
}
|
|
|
|
except Exception as e:
|
|
|
|
raise HTTPException(
|
|
status_code=500,
|
|
detail=str(e),
|
|
)
|
|
|
|
finally:
|
|
|
|
if cursor:
|
|
cursor.close()
|
|
|
|
if conn:
|
|
conn.close() |