go

aiven_keep_alive_service.py

@@ -11,8 +11,10 @@ from pydantic import BaseModel
 # Load environment variables
 load_dotenv()
 
-sql_uri = os.getenv("AIVEN_MYSQL_URI")
-postgres_uri = os.getenv("AIVEN_POSTGRES_URI")
+# We expect these to be comma-separated strings if multiple URIs exist
+# Example: "mysql://user:pass@host1/db,mysql://user:pass@host2/db"
+mysql_uris_env = os.getenv("AIVEN_MYSQL_URIS", "")
+postgres_uris_env = os.getenv("AIVEN_POSTGRESQL_URIS", "")
 
 # Configure logging (console only)
 logging.basicConfig(
@@ -35,11 +37,21 @@ class AivenPingResponse(BaseModel):
 class AivenAllPingResponse(BaseModel):
     aiven_services: List[AivenPingResponse]
 
+# ---------- Helper to parse URI lists ----------
+def get_uris_from_env(env_str: str) -> List[str]:
+    """Splits a comma-separated string of URIs into a list."""
+    if not env_str:
+        return []
+    # Split by comma and strip whitespace
+    return [uri.strip() for uri in env_str.split(',') if uri.strip()]
+
 # ---------- MySQL ----------
-def ping_mysql() -> AivenPingResponse:
+def ping_mysql(uri: str) -> AivenPingResponse:
     now = datetime.utcnow().isoformat()
+    host = "unknown"
+    
     try:
-        parsed_uri = urlparse(sql_uri)
+        parsed_uri = urlparse(uri)
         host = parsed_uri.hostname
         port = parsed_uri.port or 3306
         user = parsed_uri.username
@@ -47,15 +59,18 @@ def ping_mysql() -> AivenPingResponse:
         database = parsed_uri.path.lstrip('/') or 'defaultdb'
 
         query_params = parse_qs(parsed_uri.query)
+        # Handle SSL mode
         ssl_mode = query_params.get('ssl-mode', ['REQUIRED'])[0]
+        ssl_config = {'mode': ssl_mode} if ssl_mode else None
 
+        # PyMySQL connection
         connection = pymysql.connect(
             host=host,
             port=port,
             user=user,
             password=password,
             database=database,
-            ssl={'mode': ssl_mode} if ssl_mode else None
+            ssl=ssl_config
         )
         with connection.cursor() as cursor:
             cursor.execute("SELECT 1")
@@ -73,26 +88,28 @@ def ping_mysql() -> AivenPingResponse:
         )
     except Exception as e:
         error_msg = str(e)
-        logger.error(f"MySQL ping failed: {error_msg}")
+        logger.error(f"MySQL ping failed for host {host}: {error_msg}")
         return AivenPingResponse(
             service_name="mysql_service",
             success=False,
             error=error_msg,
             time=now,
-            host=None,
+            host=host,
             db_type="mysql"
         )
 
 
 # ---------- PostgreSQL ----------
-def ping_postgres() -> AivenPingResponse:
+def ping_postgres(uri: str) -> AivenPingResponse:
     now = datetime.utcnow().isoformat()
+    host = "unknown"
+    
     try:
-        parsed_uri = urlparse(postgres_uri)
+        parsed_uri = urlparse(uri)
         host = parsed_uri.hostname
         db_name = parsed_uri.path.lstrip('/')
 
-        connection = psycopg2.connect(postgres_uri)
+        connection = psycopg2.connect(uri)
         cursor = connection.cursor()
         cursor.execute("SELECT 1")
         cursor.fetchone()
@@ -110,28 +127,35 @@ def ping_postgres() -> AivenPingResponse:
         )
     except Exception as e:
         error_msg = str(e)
-        logger.error(f"PostgreSQL ping failed: {error_msg}")
+        logger.error(f"PostgreSQL ping failed for host {host}: {error_msg}")
         return AivenPingResponse(
             service_name="postgres_service",
             success=False,
             error=error_msg,
             time=now,
-            host=None,
+            host=host,
             db_type="postgres"
         )
 
 
 # ---------- Run Aiven Pings ----------
-def ping_aiven_projects() -> list[AivenPingResponse]:
+def ping_aiven_projects() -> AivenAllPingResponse:
     logger.info("Starting Aiven service health check")
     results = []
-    if sql_uri:
-        results.append(ping_mysql())
-    if postgres_uri:
-        results.append(ping_postgres())
-    logger.info("Completed Aiven service health check")
-    final_result = AivenAllPingResponse(aiven_services=results)
-    return final_result
 
+    # Parse lists from environment variables
+    mysql_uris = get_uris_from_env(mysql_uris_env)
+    postgres_uris = get_uris_from_env(postgres_uris_env)
+
+    # Process MySQL URIs
+    for uri in mysql_uris:
+        results.append(ping_mysql(uri))
 
+    # Process PostgreSQL URIs
+    for uri in postgres_uris:
+        results.append(ping_postgres(uri))
 
+    logger.info(f"Completed health check. Checked {len(results)} services.")
+    
+    final_result = AivenAllPingResponse(aiven_services=results)
+    return final_result