498 lines
11 KiB
Python
498 lines
11 KiB
Python
from fastapi import FastAPI, HTTPException,WebSocket, WebSocketDisconnect, Header,Depends,Query,status
|
|
from datetime import datetime, timedelta, timezone
|
|
import jwt
|
|
from jwt import ExpiredSignatureError, InvalidTokenError
|
|
from pydantic import BaseModel
|
|
import psycopg2
|
|
from psycopg2.extras import RealDictCursor
|
|
import asyncio
|
|
from fastapi.middleware.cors import CORSMiddleware
|
|
|
|
from database import Base, engine,SessionLocal
|
|
# import models
|
|
from models import NodeInterface
|
|
origins = [
|
|
"http://localhost:5173",
|
|
"http://127.0.0.1:5173",
|
|
]
|
|
|
|
|
|
app = FastAPI()
|
|
|
|
Base.metadata.create_all(bind=engine)
|
|
|
|
app.add_middleware(
|
|
CORSMiddleware,
|
|
allow_origins=origins,
|
|
allow_credentials=True,
|
|
allow_methods=["*"],
|
|
allow_headers=["*"],
|
|
)
|
|
|
|
|
|
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
|
|
|
|
# The node where sub_create() must run
|
|
subscriber_host: str
|
|
subscriber_port: int = 5444
|
|
|
|
# The provider node
|
|
provider_host: str
|
|
provider_port: int = 5444
|
|
|
|
dbname: str = "postgres"
|
|
user: str = "postgres"
|
|
password: str
|
|
|
|
|
|
#POST /register-node
|
|
# {
|
|
# "node_name": "node2",
|
|
# "host": "172.16.10.110",
|
|
# "port": 5444,
|
|
# "dbname": "postgres",
|
|
# "user": "postgres",
|
|
# "password": "postgres"
|
|
# }
|
|
|
|
# postgres=# SELECT spock.node_add_interface(
|
|
# 'node2',
|
|
# 'node2_external',
|
|
# 'host=172.16.10.110 port=5444 dbname=postgres user=postgres password=postgres'
|
|
# );
|
|
|
|
|
|
# ============================================================
|
|
# JWT TOKEN CREATION
|
|
# ============================================================
|
|
SECRET_KEY = "my-super-secret-key-change-this"
|
|
ALGORITHM = "HS256"
|
|
|
|
|
|
@app.post("/get-token")
|
|
def get_token():
|
|
print("hio")
|
|
payload = {
|
|
"sub": "frontend",
|
|
"exp": datetime.now(timezone.utc) + timedelta(hours=1)
|
|
}
|
|
|
|
token = jwt.encode(
|
|
payload,
|
|
SECRET_KEY,
|
|
algorithm=ALGORITHM
|
|
)
|
|
|
|
return {
|
|
"access_token": token,
|
|
"token_type": "bearer"
|
|
}
|
|
|
|
|
|
|
|
|
|
# ============================================================
|
|
# 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 Nodes
|
|
@app.post("/register-node")
|
|
def register_node(data: RegisterNodeRequest, token_data: dict = Depends(verify_token)):
|
|
|
|
conn = None
|
|
cur = None
|
|
|
|
try:
|
|
# Connect to the node we want to register
|
|
conn = psycopg2.connect(
|
|
host=data.host,
|
|
port=data.port,
|
|
dbname=data.dbname,
|
|
user=data.user,
|
|
password=data.password
|
|
)
|
|
|
|
cur = conn.cursor()
|
|
|
|
# 1. Create Spock extension
|
|
cur.execute("""
|
|
CREATE EXTENSION IF NOT EXISTS spock;
|
|
""")
|
|
|
|
# 2. Check Spock version
|
|
cur.execute("""
|
|
SELECT extname, extversion
|
|
FROM pg_extension
|
|
WHERE extname = 'spock';
|
|
""")
|
|
|
|
spock = cur.fetchone()
|
|
|
|
# 3. Build node DSN
|
|
dsn = (
|
|
f"host={data.host} "
|
|
f"port={data.port} "
|
|
f"dbname={data.dbname} "
|
|
f"user={data.user} "
|
|
f"password={data.password}"
|
|
)
|
|
|
|
# 4. Create Spock node
|
|
cur.execute("""
|
|
SELECT spock.node_create(
|
|
node_name := %s,
|
|
dsn := %s
|
|
);
|
|
""", (
|
|
data.node_name,
|
|
dsn
|
|
))
|
|
|
|
node_id = cur.fetchone()[0]
|
|
|
|
conn.commit()
|
|
|
|
# 5. Show registered nodes
|
|
cur.execute("""
|
|
SELECT
|
|
n.node_id,
|
|
n.node_name,
|
|
ni.if_id
|
|
FROM spock.node n
|
|
LEFT JOIN spock.node_interface ni
|
|
ON ni.if_nodeid = n.node_id;
|
|
""")
|
|
|
|
rows = cur.fetchall()
|
|
|
|
nodes = [
|
|
{
|
|
"node_id": row[0],
|
|
"node_name": row[1],
|
|
"if_id": row[2]
|
|
}
|
|
for row in rows
|
|
]
|
|
registered_node = next(
|
|
node for node in nodes
|
|
if node["node_id"] == node_id
|
|
)
|
|
|
|
db = SessionLocal()
|
|
|
|
try:
|
|
node_interface = NodeInterface(
|
|
if_id=registered_node["if_id"],
|
|
if_name=registered_node["node_name"],
|
|
if_nodeid=registered_node["node_id"],
|
|
host=data.host,
|
|
port=data.port,
|
|
dbname=data.dbname,
|
|
user=data.user
|
|
)
|
|
|
|
db.add(node_interface)
|
|
db.commit()
|
|
|
|
finally:
|
|
db.close()
|
|
|
|
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()
|
|
|
|
|
|
|
|
# {
|
|
# "subscription_name": "node2_from_node1",
|
|
# "subscriber_host": "172.16.10.110",
|
|
# "subscriber_port": 5444,
|
|
# "provider_host": "172.16.10.123",
|
|
# "provider_port": 5444,
|
|
# "dbname": "postgres",
|
|
# "user": "postgres",
|
|
# "password": "postgres"
|
|
# }
|
|
@app.post("/link-nodes")
|
|
def link_nodes(data: LinkNodeRequest,token_data: dict = Depends(verify_token)):
|
|
|
|
conn = None
|
|
cur = None
|
|
|
|
try:
|
|
# IMPORTANT:
|
|
# Connect to the subscriber node.
|
|
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
|
|
provider_dsn = (
|
|
f"host={data.provider_host} "
|
|
f"port={data.provider_port} "
|
|
f"dbname={data.dbname} "
|
|
f"user={data.user} "
|
|
f"password={data.password}"
|
|
)
|
|
|
|
# Create subscription
|
|
# Create subscription
|
|
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()
|
|
|
|
|
|
DB_CONFIG = {
|
|
"host": "172.16.10.123",
|
|
"port": 5444,
|
|
"dbname": "postgres",
|
|
"user": "postgres",
|
|
"password": "postgres"
|
|
}
|
|
|
|
|
|
|
|
|
|
def get_spock_status():
|
|
conn = psycopg2.connect(**DB_CONFIG)
|
|
|
|
cursor = conn.cursor(cursor_factory=RealDictCursor)
|
|
|
|
# cursor.execute("""
|
|
# SELECT *
|
|
# FROM spock.sub_show_status();
|
|
# """)
|
|
|
|
cursor.execute("""
|
|
SELECT
|
|
subscription_name,
|
|
status,
|
|
provider_node
|
|
FROM spock.sub_show_status();
|
|
""")
|
|
|
|
rows = cursor.fetchall()
|
|
|
|
cursor.close()
|
|
conn.close()
|
|
|
|
return rows
|
|
|
|
|
|
@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:
|
|
status = await asyncio.to_thread(get_spock_status)
|
|
|
|
await websocket.send_json({
|
|
"database_reachable": True,
|
|
"subscriptions": status
|
|
})
|
|
|
|
except Exception as e:
|
|
|
|
await websocket.send_json({
|
|
"database_reachable": False,
|
|
"subscriptions": [],
|
|
"error": str(e)
|
|
})
|
|
return False
|
|
|
|
await asyncio.sleep(2)
|
|
|
|
except WebSocketDisconnect:
|
|
print("WebSocket client disconnected" f"{token_data.get('sub')}")
|
|
return True
|
|
|
|
|
|
|
|
@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()
|
|
|
|
cursor.close()
|
|
conn.close()
|
|
|
|
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)
|
|
)
|
|
|
|
|
|
|
|
|