Spaces:
Running
Running
File size: 3,715 Bytes
09801ca | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 | import asyncio
from datetime import datetime
import asyncpg
import json
from .base import LiveConnector
class PostgresConnector(LiveConnector):
async def get_metrics_stream(self):
# We will attempt to connect using the provided credentials
# In a real environment, you'd parse host, db, and token/password.
# For this prototype, we'll assume the credentials field holds the password.
while True:
try:
# Try to connect
conn = await asyncpg.connect(
host=self.host,
database=self.database_name,
user="postgres", # Hardcoded default for prototype, could be configurable
password=self.credentials,
timeout=5.0
)
try:
previous_count = 0
while True:
if hasattr(self, 'target_table') and self.target_table:
# Fetch total rows from target table securely
# We use simple string formatting here assuming self.target_table is validated,
# but in production we should escape it or validate against information_schema.
try:
count_row = await conn.fetchrow(f"SELECT COUNT(*) as total FROM {self.target_table}")
total_rows = count_row['total'] if count_row else 0
except Exception as e:
print(f"Error querying table {self.target_table}: {e}")
total_rows = 0
rows_added_per_sec = (total_rows - previous_count) / 2.0 if previous_count > 0 else 0
previous_count = total_rows
yield {
"timestamp": datetime.utcnow().isoformat(),
"total_rows": total_rows,
"rows_per_sec": max(0, rows_added_per_sec),
"cpu_usage": 45.5, # Estimated
"error_rate": 0.0,
"connector_source": "PostgreSQL",
"status": "Healthy"
}
else:
yield {
"timestamp": datetime.utcnow().isoformat(),
"total_rows": 0,
"rows_per_sec": 0,
"cpu_usage": 0.0,
"error_rate": 0.0,
"connector_source": "PostgreSQL",
"status": "No target table configured"
}
await asyncio.sleep(2)
finally:
await conn.close()
except Exception as e:
# Connection failed, yield error metric or retry
print(f"Postgres Connection Error: {e}")
yield {
"timestamp": datetime.utcnow().isoformat(),
"active_users": 0,
"cpu_usage": 0.0,
"transactions_per_sec": 0,
"error_rate": 100.0,
"connector_source": "PostgreSQL",
"status": f"Connection Failed: {str(e)}"
}
await asyncio.sleep(5) # Retry delay
|