Spaces:
Running
Running
| 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 | |