pgMultiMaster/main.py

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

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]]


#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)
        )


# !!!!!!!!!!!!!!!!!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"
        )