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,get_db from uuid import UUID # import models from sqlalchemy.orm import Session from models import NodeInterface, Cluster import os from dotenv import load_dotenv from sqlalchemy.exc import SQLAlchemyError from typing import List load_dotenv() origins = os.getenv("CORS_ORIGINS", "").split(",") 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 dbname: str user: str password: str class LinkNodeRequest(BaseModel): subscription_name: str # The node where sub_create() must run subscriber_host: str subscriber_port: int # The provider node provider_host: str provider_port: int dbname: str user: str password: str class NodeLinkRequest(BaseModel): from_node: str to_node: str class ClusterRequest(BaseModel): cluster_name: str links: list[tuple[str, str]] class NodeResponse(BaseModel): node_name: str uuid: UUID class ClusterNodesResponse(BaseModel): cluster_name: str nodes: List[NodeResponse] class ClusterResponse(BaseModel): uuid: UUID cluster_name: 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 = os.getenv("SECRET_KEY") ALGORITHM = os.getenv("ALGORITHM", "HS256") # 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" # } #!!!!!!!!!!!!INIT synchronize_structure := false, synchronize_data := false !!!!!!!!!!! # @app.post("/link-nodes") def link_nodes(data: LinkNodeRequest): 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, } 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() #!!!!!!!!!!!!THIS IS A TEMPORARY FIXSS!!!!!!!!!!!!! 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) ) @app.get("/spock/node-interfaces") def get_node_interfaces(db: Session = Depends(get_db),token_data: dict = Depends(verify_token)): try: rows = db.query(NodeInterface).all() return { "status": "success", "count": len(rows), "nodes": [ {"uuid": str(row.uuid), "if_id": row.if_id, "if_name": row.if_name, "if_nodeid": row.if_nodeid, "host": row.host, "port": row.port, "dbname": row.dbname, "user": row.user } for row in rows ] } except Exception as e: raise HTTPException( status_code=500, detail=str(e) ) @app.get("/spock/node-interfaces/{node_uuid}") def get_node_interface_by_uuid( node_uuid: UUID, db: Session = Depends(get_db), token_data: dict = Depends(verify_token) ): try: row = ( db.query(NodeInterface) .filter(NodeInterface.uuid == node_uuid) .first() ) if not row: raise HTTPException( status_code=404, detail="Node not found" ) return { "status": "success", "node": { "uuid": str(row.uuid), "if_id": row.if_id, "if_name": row.if_name, "if_nodeid": row.if_nodeid, "host": row.host, "port": row.port, "dbname": row.dbname, "user": row.user } } except HTTPException: raise except Exception as e: raise HTTPException( status_code=500, detail=str(e) ) # !!!!!!!!!!!!!!!!!PASSWORD MUST BE STORED IN MANAGEMENT TABLE!!!!!!!!!!!!!! @app.delete("/spock/del-node-interfaces") def delete_node_interface( uuid: UUID, db: Session = Depends(get_db), token_data: dict = Depends(verify_token) ): print("API called DEL API") print(f"uuid received: {uuid}") node = ( db.query(NodeInterface) .filter(NodeInterface.uuid == uuid) .first() ) if not node: raise HTTPException( status_code=404, detail="Node interface not found" ) conn = None cur = None try: conn = psycopg2.connect( host=node.host, port=node.port, dbname=node.dbname, user=node.user, password=os.getenv("DB_PASSWORD") ) cur = conn.cursor() cur.execute( "SELECT spock.node_drop(%s);", (node.if_name,) ) result = cur.fetchone() conn.commit() # 2. Delete management row db.delete(node) db.commit() return { "status": "success", "message": f"Spock node {node.if_name} unregistered successfully", "uuid": str(uuid), "node_name": node.if_name } except Exception as e: if conn: conn.rollback() print(f"Error: {e}") raise HTTPException( status_code=500, detail=f"Failed to unregister node: {str(e)}" ) finally: if cur: cur.close() if conn: conn.close() # @app.post("/spock/link-nodes-abstract") def link_nodes_abstract(data: NodeLinkRequest,db: Session): try: from_node = data.from_node to_node = data.to_node print("From Node:", from_node) from_node_data = ( db.query(NodeInterface) .filter(NodeInterface.if_name == from_node) .first() ) if not from_node_data: raise HTTPException( status_code=404, detail=f"Node '{from_node}' not found" ) print("fromFound:", from_node_data.if_name) print("fromHost:", from_node_data.host) print("fromPort:", from_node_data.port) print("fromDbname:", from_node_data.dbname) print("fromUser:", from_node_data.user) #================================================ print("===============================================") print("To Node:", to_node) to_node_data = ( db.query(NodeInterface) .filter(NodeInterface.if_name == to_node) .first() ) if not to_node_data: raise HTTPException( status_code=404, detail=f"Node '{to_node}' not found" ) print("toFound:", to_node_data.if_name) print("toHost:", to_node_data.host) print("toPort:", to_node_data.port) print("toDbname:", to_node_data.dbname) print("toUser:", to_node_data.user) #============================================================================= # Build internal function input link_data = LinkNodeRequest( subscription_name=f"{to_node}_from_{from_node}", subscriber_host=to_node_data.host, subscriber_port=to_node_data.port, provider_host=from_node_data.host, provider_port=from_node_data.port, dbname=to_node_data.dbname, user=to_node_data.user, password="postgres" ) # Call internal function result = link_nodes(link_data) print("Linking Result:", result["status"]) #============================================================================== return { "from_node": from_node, "to_node": to_node } except HTTPException: raise except SQLAlchemyError as e: print("Database error:", str(e)) raise HTTPException( status_code=500, detail="Management database error" ) except Exception as e: print("Unexpected error:", str(e)) raise HTTPException( status_code=500, detail="Internal server error" ) @app.post("/spock/create_cluster") # INPUT EXAMPLE # { # "cluster_name": "cluster1", # "links": [ # ["node1", "node2"], # ["node2", "node1"] # ] # } def create_cluster(data: ClusterRequest,db: Session = Depends(get_db),token_data: dict = Depends(verify_token)): try: print("Cluster Name:", data.cluster_name) print("Links:", data.links) results = [] node_names = set() for from_node, to_node in data.links: print(f"{from_node} -> {to_node}") #================= link_request = NodeLinkRequest( from_node=from_node, to_node=to_node ) result = link_nodes_abstract( link_request, db ) print( f"Successfully linked {from_node} -> {to_node}" ) node_names.add(from_node) node_names.add(to_node) results.append({ "from_node": from_node, "to_node": to_node, "status": "success" }) cluster = Cluster( cluster_name=data.cluster_name ) db.add(cluster) db.flush() nodes = ( db.query(NodeInterface) .filter(NodeInterface.if_name.in_(node_names)) .all() ) if len(nodes) != len(node_names): found_nodes = {node.if_name for node in nodes} missing_nodes = node_names - found_nodes raise HTTPException( status_code=404, detail=f"Nodes not found: {list(missing_nodes)}" ) for node in nodes: node.cluster = cluster db.commit() return { "cluster_name": data.cluster_name, "links": data.links, "links": results } except HTTPException: db.rollback() raise except Exception as e: db.rollback() print("Cluster creation failed:", str(e)) raise HTTPException( status_code=500, detail="Cluster creation failed" ) @app.get( "/spock/cluster-nodes/{cluster_uuid}", response_model=ClusterNodesResponse ) def get_cluster_nodes( cluster_uuid: UUID, db: Session = Depends(get_db), token_data=Depends(verify_token) ): cluster = ( db.query(Cluster) .filter(Cluster.uuid == cluster_uuid) .first() ) if not cluster: raise HTTPException( status_code=404, detail=f"Cluster with UUID '{cluster_uuid}' not found" ) nodes = ( db.query(NodeInterface) .filter(NodeInterface.cluster_id == cluster.uuid) .all() ) return { "cluster_name": cluster.cluster_name, "nodes": [ { "node_name": node.if_name, "uuid": node.uuid } for node in nodes ] } @app.get( "/spock/clusters", response_model=List[ClusterResponse] ) def get_clusters( db: Session = Depends(get_db), token_data=Depends(verify_token) ): clusters = db.query(Cluster).all() return [ { "uuid": cluster.uuid, "cluster_name": cluster.cluster_name } for cluster in clusters ]