pgMultiMasterAgent/main.py

import os
import uvicorn
from dotenv import load_dotenv
from fastapi import FastAPI,Depends, Header, HTTPException, status, WebSocket, WebSocketDisconnect
import psycopg2
from jwt import ExpiredSignatureError, InvalidTokenError
import jwt
from pydantic import BaseModel
from psycopg2.extras import RealDictCursor
import asyncio

load_dotenv()
app = FastAPI()

SECRET_KEY = os.getenv("SECRET_KEY")
ALGORITHM = os.getenv("ALGORITHM", "HS256")


class NodeStatusRequest(BaseModel):
    node_name: str
    host: str
    port: int
    dbname: str
    user: str
    password: str



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",
        )
 





@app.get("/")
def root():
    return {"message": "PostgreSQL Agent is running"}



#{
#     "host": "172.16.10.178",
#     "node_name = "node3",
#     "port": 5444,
#     "dbname": "postgres",
#     "user": "postgres",
#     "password": "postgres",
#     "connect_timeout": 3,
# }


@app.get("/node/status")
def node_status(data: NodeStatusRequest,token_data=Depends(verify_token)):

    try:
        db_config = {
            "host": data.host,
            "port": data.port,
            "dbname": data.dbname,
            "user": data.user,
            "password": data.password,
            "connect_timeout": 3,
        }
        conn = psycopg2.connect(**db_config)

        conn.close()

        return {
            "node": data.node_name,
            "online": True,
            "status": "online"
        }

    except Exception as e:

        return {
            "node": data.node_name,
            "online": False,
            "status": "offline",
            "error": str(e)
        }



latest_replication_health = None
monitor_task = None
def get_replication_health(data: NodeStatusRequest):

    db_config = {
        "host": data.host,
        "port": data.port,
        "dbname": data.dbname,
        "user": data.user,
        "password": data.password,
        "connect_timeout": 2,
        "options": "-c statement_timeout=2000"
    }

    conn = None

    try:
        conn = psycopg2.connect(**db_config)

        cur = conn.cursor(cursor_factory=RealDictCursor)

        cur.execute("""
            SELECT
                subscription_name,
                status,
                provider_node
            FROM spock.sub_show_status();
        """)

        subscriptions = cur.fetchall()

        healthy = all(
            sub["status"] == "replicating"
            for sub in subscriptions
        )

        return {
            "node": data.node_name,
            "connection_healthy": healthy,
            "subscriptions": subscriptions
        }
    except psycopg2.errors.QueryCanceled:

        print("Query cancelled because it exceeded 2 seconds")

        return {
            "status": "timeout",
            "message": "PostgreSQL query exceeded timelimit and was cancelled"
        }

    except Exception as e:

        return {
            "node": data.node_name,
            "connection_healthy": False,
            "error": str(e)
        }

    finally:
        if conn:
            conn.close()


@app.websocket("/ws/node/replication-health")
async def replication_monitor(websocket: WebSocket):
    await websocket.accept()
    try:
        raw_data = await websocket.receive_json()
        data = NodeStatusRequest(**raw_data)
        global latest_replication_health

        while True:

            try:

                result = await asyncio.wait_for(
                    asyncio.to_thread(
                    get_replication_health,
                    data
                    ),
                    timeout=6
                    )

                latest_replication_health = result

                print("Replication health:")
                print(result)

            except asyncio.TimeoutError:

                latest_replication_health = {
                    "node": data.node_name,
                    "connection_healthy": False,
                    "status": "timeout",
                    "message": "PostgreSQL operation exceeded maximum time"
                }

            print(latest_replication_health)
            
            await websocket.send_json(
                latest_replication_health
            )

            # Wait 2 seconds before checking again
            await asyncio.sleep(2)
    except WebSocketDisconnect:

        print("WebSocket client disconnected")





















if __name__ == "__main__":
    host = os.getenv("HOST", "127.0.0.1")
    port = int(os.getenv("PORT", 9000))

    uvicorn.run(
        "main:app",
        host=host,
        port=port,
        reload=False
    )