426 lines
8.8 KiB
Python
426 lines
8.8 KiB
Python
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
|
|
|
|
|
|
|
|
class DatabaseRequest(BaseModel):
|
|
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()
|
|
|
|
|
|
|
|
# {
|
|
# "node_name": "node3",
|
|
# "host": "172.16.10.178",
|
|
# "port": 5444,
|
|
# "dbname": "postgres",
|
|
# "user": "postgres",
|
|
# "password": "postgres"
|
|
# }
|
|
|
|
@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")
|
|
|
|
|
|
|
|
def get_database_status(
|
|
data: DatabaseRequest
|
|
):
|
|
|
|
conn = None
|
|
|
|
try:
|
|
conn = psycopg2.connect(
|
|
host=data.host,
|
|
port=data.port,
|
|
dbname=data.dbname,
|
|
user=data.user,
|
|
password=data.password,
|
|
connect_timeout=3
|
|
)
|
|
|
|
return {
|
|
"database_name": data.dbname,
|
|
"status": "online"
|
|
}
|
|
|
|
except psycopg2.Error:
|
|
|
|
return {
|
|
"database_name": data.dbname,
|
|
"status": "offline"
|
|
}
|
|
|
|
finally:
|
|
if conn:
|
|
conn.close()
|
|
|
|
|
|
def get_database_uptime(
|
|
data: DatabaseRequest
|
|
):
|
|
|
|
conn = None
|
|
|
|
try:
|
|
conn = psycopg2.connect(
|
|
host=data.host,
|
|
port=data.port,
|
|
dbname=data.dbname,
|
|
user=data.user,
|
|
password=data.password,
|
|
connect_timeout=3
|
|
)
|
|
|
|
cursor = conn.cursor()
|
|
|
|
cursor.execute("""
|
|
SELECT EXTRACT(
|
|
EPOCH FROM (
|
|
CURRENT_TIMESTAMP - pg_postmaster_start_time()
|
|
)
|
|
)::BIGINT;
|
|
""")
|
|
|
|
uptime_seconds = cursor.fetchone()[0]
|
|
|
|
cursor.close()
|
|
|
|
return {
|
|
"database_name": data.dbname,
|
|
"uptime_seconds": uptime_seconds
|
|
}
|
|
|
|
except psycopg2.Error as e:
|
|
|
|
return {
|
|
"database_name": data.dbname,
|
|
"uptime_seconds": None,
|
|
"error": str(e)
|
|
}
|
|
|
|
finally:
|
|
if conn:
|
|
conn.close()
|
|
|
|
|
|
@app.websocket("/ws/database/status")
|
|
async def database_status_monitor(websocket: WebSocket):
|
|
|
|
await websocket.accept()
|
|
|
|
try:
|
|
|
|
raw_data = await websocket.receive_json()
|
|
|
|
data = DatabaseRequest(**raw_data)
|
|
|
|
while True:
|
|
|
|
try:
|
|
|
|
status_result, uptime_result = await asyncio.wait_for(
|
|
asyncio.gather(
|
|
asyncio.to_thread(
|
|
get_database_status,
|
|
data
|
|
),
|
|
asyncio.to_thread(
|
|
get_database_uptime,
|
|
data
|
|
)
|
|
),
|
|
timeout=5
|
|
)
|
|
|
|
result = {
|
|
"database_name": data.dbname,
|
|
"status": status_result["status"],
|
|
"uptime_seconds": uptime_result["uptime_seconds"]
|
|
}
|
|
|
|
except asyncio.TimeoutError:
|
|
|
|
result = {
|
|
"database_name": data.dbname,
|
|
"status": "offline",
|
|
"uptime_seconds": None
|
|
}
|
|
|
|
await websocket.send_json(result)
|
|
|
|
await asyncio.sleep(2)
|
|
|
|
except WebSocketDisconnect:
|
|
|
|
print("Database status WebSocket 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
|
|
)
|
|
|