805 lines
18 KiB
Python
805 lines
18 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
|
|
import time
|
|
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")
|
|
|
|
|
|
#GET ACTIVE CONNECTIONS
|
|
|
|
def get_active_connections(
|
|
data: DatabaseRequest
|
|
):
|
|
|
|
conn = None
|
|
cursor = 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 COUNT(*)
|
|
FROM pg_stat_activity
|
|
WHERE state = 'active'
|
|
AND pid <> pg_backend_pid();
|
|
""")
|
|
|
|
active_connections = cursor.fetchone()[0]
|
|
|
|
return {
|
|
"database_name": data.dbname,
|
|
"active_connections": active_connections
|
|
}
|
|
|
|
except psycopg2.Error as e:
|
|
|
|
raise HTTPException(
|
|
status_code=500,
|
|
detail=f"Failed to fetch active connections: {str(e)}"
|
|
)
|
|
|
|
finally:
|
|
if cursor:
|
|
cursor.close()
|
|
|
|
if conn:
|
|
conn.close()
|
|
|
|
|
|
# HOW MANY CONNECTIONS ARE BEING USED
|
|
def get_connection_usage(data: DatabaseRequest):
|
|
|
|
conn = None
|
|
cursor = 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
|
|
COUNT(*) AS current_connections,
|
|
current_setting('max_connections')::int AS max_connections
|
|
FROM pg_stat_activity
|
|
GROUP BY current_setting('max_connections');
|
|
""")
|
|
|
|
row = cursor.fetchone()
|
|
|
|
current_connections = row[0]
|
|
max_connections = row[1]
|
|
|
|
connection_usage_percent = round(
|
|
(current_connections / max_connections) * 100,
|
|
2
|
|
)
|
|
|
|
return {
|
|
"current_connections": current_connections,
|
|
"max_connections": max_connections,
|
|
"connection_usage_percent": connection_usage_percent
|
|
}
|
|
|
|
except psycopg2.Error as e:
|
|
|
|
return {
|
|
"current_connections": None,
|
|
"max_connections": None,
|
|
"connection_usage_percent": None,
|
|
"error": str(e)
|
|
}
|
|
|
|
finally:
|
|
|
|
if cursor:
|
|
cursor.close()
|
|
|
|
if conn:
|
|
conn.close()
|
|
|
|
# WEB SOCKETS FOR ACTIVE CONNECTIONS MONITORING
|
|
@app.websocket("/ws/database/active-connections")
|
|
async def database_active_connections_monitor(websocket: WebSocket):
|
|
|
|
await websocket.accept()
|
|
|
|
try:
|
|
|
|
raw_data = await websocket.receive_json()
|
|
|
|
data = DatabaseRequest(**raw_data)
|
|
|
|
while True:
|
|
|
|
try:
|
|
|
|
(
|
|
active_connections_result,
|
|
connection_usage_result
|
|
) = await asyncio.wait_for(
|
|
|
|
asyncio.gather(
|
|
|
|
asyncio.to_thread(
|
|
get_active_connections,
|
|
data
|
|
),
|
|
|
|
asyncio.to_thread(
|
|
get_connection_usage,
|
|
data
|
|
)
|
|
|
|
),
|
|
|
|
timeout=5
|
|
)
|
|
|
|
result = {
|
|
"database_name": data.dbname,
|
|
|
|
"active_connections":
|
|
active_connections_result[
|
|
"active_connections"
|
|
],
|
|
|
|
"current_connections":
|
|
connection_usage_result[
|
|
"current_connections"
|
|
],
|
|
|
|
"max_connections":
|
|
connection_usage_result[
|
|
"max_connections"
|
|
],
|
|
|
|
"connection_usage_percent":
|
|
connection_usage_result[
|
|
"connection_usage_percent"
|
|
]
|
|
}
|
|
|
|
except asyncio.TimeoutError:
|
|
|
|
result = {
|
|
"database_name": data.dbname,
|
|
"active_connections": None,
|
|
"current_connections": None,
|
|
"max_connections": None,
|
|
"connection_usage_percent": None
|
|
}
|
|
|
|
await websocket.send_json(result)
|
|
|
|
await asyncio.sleep(2)
|
|
|
|
except WebSocketDisconnect:
|
|
|
|
print(
|
|
"Database active connections WebSocket disconnected"
|
|
)
|
|
|
|
|
|
#TPS
|
|
def get_database_tps(data: DatabaseRequest):
|
|
|
|
conn = None
|
|
cursor = None
|
|
|
|
try:
|
|
|
|
conn = psycopg2.connect(
|
|
host=data.host,
|
|
port=data.port,
|
|
dbname=data.dbname,
|
|
user=data.user,
|
|
password=data.password,
|
|
connect_timeout=3
|
|
)
|
|
|
|
conn.autocommit = True
|
|
cursor = conn.cursor()
|
|
|
|
# First sample
|
|
cursor.execute("""
|
|
SELECT
|
|
xact_commit + xact_rollback
|
|
FROM pg_stat_database
|
|
WHERE datname = %s;
|
|
""", (data.dbname,))
|
|
|
|
first_count = cursor.fetchone()[0]
|
|
|
|
start_time = time.monotonic()
|
|
|
|
time.sleep(1)
|
|
|
|
# Second sample
|
|
cursor.execute("""
|
|
SELECT
|
|
xact_commit + xact_rollback
|
|
FROM pg_stat_database
|
|
WHERE datname = %s;
|
|
""", (data.dbname,))
|
|
|
|
second_count = cursor.fetchone()[0]
|
|
|
|
end_time = time.monotonic()
|
|
|
|
elapsed_time = end_time - start_time
|
|
|
|
transactions = second_count - first_count
|
|
|
|
tps = round(
|
|
transactions / elapsed_time,
|
|
2
|
|
)
|
|
|
|
return {
|
|
"database_name": data.dbname,
|
|
"tps": tps
|
|
}
|
|
|
|
except psycopg2.Error as e:
|
|
|
|
return {
|
|
"database_name": data.dbname,
|
|
"tps": None,
|
|
"error": str(e)
|
|
}
|
|
|
|
finally:
|
|
|
|
if cursor:
|
|
cursor.close()
|
|
|
|
if conn:
|
|
conn.close()
|
|
|
|
# CACHE HIT RATIO
|
|
def get_cache_hit_ratio(data: DatabaseRequest):
|
|
|
|
conn = None
|
|
cursor = 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
|
|
CASE
|
|
WHEN (blks_hit + blks_read) = 0 THEN 0
|
|
ELSE ROUND(
|
|
(blks_hit::numeric /
|
|
(blks_hit + blks_read)) * 100,
|
|
2
|
|
)
|
|
END AS cache_hit_ratio
|
|
FROM pg_stat_database
|
|
WHERE datname = %s;
|
|
""", (data.dbname,))
|
|
|
|
row = cursor.fetchone()
|
|
|
|
cache_hit_ratio = float(row[0]) if row else 0.0
|
|
|
|
return {
|
|
"database_name": data.dbname,
|
|
"cache_hit_ratio": cache_hit_ratio
|
|
}
|
|
|
|
except psycopg2.Error as e:
|
|
|
|
return {
|
|
"database_name": data.dbname,
|
|
"cache_hit_ratio": None,
|
|
"error": str(e)
|
|
}
|
|
|
|
finally:
|
|
|
|
if cursor:
|
|
cursor.close()
|
|
|
|
if conn:
|
|
conn.close()
|
|
|
|
#TPS AND CACHE HIT RATIO WEB SOCKET
|
|
@app.websocket("/ws/database/performance")
|
|
async def database_performance_monitor(websocket: WebSocket):
|
|
|
|
await websocket.accept()
|
|
|
|
try:
|
|
|
|
raw_data = await websocket.receive_json()
|
|
|
|
data = DatabaseRequest(**raw_data)
|
|
|
|
while True:
|
|
|
|
try:
|
|
|
|
(
|
|
tps_result,
|
|
cache_hit_result
|
|
) = await asyncio.wait_for(
|
|
|
|
asyncio.gather(
|
|
|
|
asyncio.to_thread(
|
|
get_database_tps,
|
|
data
|
|
),
|
|
|
|
asyncio.to_thread(
|
|
get_cache_hit_ratio,
|
|
data
|
|
)
|
|
|
|
),
|
|
|
|
timeout=5
|
|
)
|
|
|
|
result = {
|
|
"database_name": data.dbname,
|
|
|
|
"tps": tps_result[
|
|
"tps"
|
|
],
|
|
|
|
"cache_hit_ratio": cache_hit_result[
|
|
"cache_hit_ratio"
|
|
]
|
|
}
|
|
|
|
except asyncio.TimeoutError:
|
|
|
|
result = {
|
|
"database_name": data.dbname,
|
|
"tps": None,
|
|
"cache_hit_ratio": None
|
|
}
|
|
|
|
await websocket.send_json(result)
|
|
|
|
await asyncio.sleep(2)
|
|
|
|
except WebSocketDisconnect:
|
|
|
|
print("Database performance 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
|
|
)
|
|
|