pgMultiMaster/main.py

from fastapi import FastAPI, HTTPException,WebSocket, WebSocketDisconnect
from pydantic import BaseModel
import psycopg2
from psycopg2.extras import RealDictCursor
import asyncio
app = FastAPI()


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



@app.post("/register-node")
def register_node(data: RegisterNodeRequest):

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

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

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


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):
    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")
        return True



@app.get("/spock/subscriptions")
def get_spock_subscription_status():
    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
        }

    except Exception as e:
        raise HTTPException(
            status_code=500,
            detail=str(e)
        )