""" AWS Aurora Serverless v2 cluster lifecycle tasks. Flow: 1. AWSAuroraServerlessClusterCreateTask - Generates credentials and saves them immediately (before any AWS API call) - Creates DB cluster, handles idempotency (cluster may exist from previous failed attempt) - Creates DB instance, handles idempotency 2. AWSAuroraServerlessClusterReadyCheckTask - Waits for cluster status to be "available" - Waits for instance status to be "available" - Retrieves and saves the cluster endpoint (only available when cluster is ready) 3. AWSAuroraServerlessClusterDeleteTask - Deletes instance first (must be deleted before cluster) - Deletes cluster (with skip-final-snapshot) - Handles "not found" gracefully (idempotent) """ import logging import string import time from typing import Optional import boto3 from botocore.exceptions import ClientError from django.utils.crypto import get_random_string from labs.lab_manager_v2.tasks.base import BaseLabTask from labs.lab_manager_v2.exceptions import ( LabTaskRecoverableError, LabTaskUnrecoverableError, ) from labs.lab_manager_v2.schemas import ( LabResource, AWSAuroraServerlessClusterConfig, ) from labs.lab_manager_v2.consts import get_region_metadata logger = logging.getLogger(__name__) # AWS RDS error codes # Throttling errors (recoverable) THROTTLING_ERROR_CODES = frozenset([ "Throttling", "ThrottlingException", "RequestLimitExceeded", "ProvisionedThroughputExceededException", "RequestThrottled", "TooManyRequestsException", "BandwidthLimitExceeded", ]) # Transient errors (recoverable with retry) TRANSIENT_ERROR_CODES = frozenset([ "ServiceUnavailable", "InternalFailure", "InternalServerError", "InternalServiceError", ]) # Capacity errors (recoverable - might succeed later) CAPACITY_ERROR_CODES = frozenset([ "InsufficientDBClusterCapacityFault", "InsufficientDBInstanceCapacity", "StorageQuotaExceeded", "DBClusterQuotaExceededFault", "DBInstanceQuotaExceeded", "InstanceQuotaExceeded", ]) # Already exists errors (need special handling - might be from previous attempt) ALREADY_EXISTS_ERROR_CODES = frozenset([ "DBClusterAlreadyExistsFault", "DBInstanceAlreadyExistsFault", "DBInstanceAlreadyExists", ]) # Not found errors (expected during deletion or when checking non-existent resources) NOT_FOUND_ERROR_CODES = frozenset([ "DBClusterNotFoundFault", "DBInstanceNotFoundFault", "DBInstanceNotFound", ]) def _generate_aurora_admin_credentials() -> tuple[str, str]: """Generate random username and password for Aurora cluster admin.""" # Username: 8 lowercase letters (RDS requirement: must start with letter) username = get_random_string(length=8, allowed_chars=string.ascii_lowercase) # Password: 32 alphanumeric characters (RDS allows most printable chars, but we keep it simple) password = get_random_string(length=32, allowed_chars=string.ascii_letters + string.digits) return username, password def _get_rds_client(region_name: str): """Get boto3 RDS client for the specified region.""" lab_region = get_region_metadata(region_name) return boto3.client("rds", lab_region["region-code"]) def _is_recoverable_error(error_code: str) -> bool: """Check if the error is recoverable (should retry).""" return error_code in THROTTLING_ERROR_CODES or \ error_code in TRANSIENT_ERROR_CODES or \ error_code in CAPACITY_ERROR_CODES def _handle_client_error( e: ClientError, operation: str, resource_id: str, allow_already_exists: bool = False, allow_not_found: bool = False, ) -> Optional[str]: """ Handle boto3 ClientError with appropriate exception type. Args: e: The ClientError exception operation: Description of the operation (for logging) resource_id: ID of the resource being operated on allow_already_exists: If True, return "already_exists" instead of raising allow_not_found: If True, return "not_found" instead of raising Returns: "already_exists" or "not_found" if those conditions are allowed, None otherwise Raises: LabTaskRecoverableError: For throttling, transient, and capacity errors LabTaskUnrecoverableError: For all other errors """ error_code = e.response.get("Error", {}).get("Code", "Unknown") error_message = e.response.get("Error", {}).get("Message", str(e)) http_status = e.response.get("ResponseMetadata", {}).get("HTTPStatusCode", 0) logger.warning( f"{operation} failed for {resource_id}: " f"error_code={error_code}, http_status={http_status}, message={error_message}" ) # Handle "already exists" - might be from a previous attempt if error_code in ALREADY_EXISTS_ERROR_CODES: if allow_already_exists: logger.info(f"Resource {resource_id} already exists, treating as success") return "already_exists" raise LabTaskUnrecoverableError( f"{operation} failed: {resource_id} already exists. " f"This may indicate a naming conflict. Error: {error_message}" ) # Handle "not found" - expected during deletion if error_code in NOT_FOUND_ERROR_CODES: if allow_not_found: logger.info(f"Resource {resource_id} not found, treating as already deleted") return "not_found" raise LabTaskUnrecoverableError( f"{operation} failed: {resource_id} not found. Error: {error_message}" ) # Handle 429 (Too Many Requests) explicitly if http_status == 429 or error_code in THROTTLING_ERROR_CODES: raise LabTaskRecoverableError( f"{operation} throttled for {resource_id}: {error_message}. Will retry." ) # Handle other recoverable errors if _is_recoverable_error(error_code): raise LabTaskRecoverableError( f"{operation} failed with recoverable error for {resource_id}: " f"{error_code} - {error_message}. Will retry." ) # Handle 5xx server errors if 500 <= http_status < 600: raise LabTaskRecoverableError( f"{operation} failed with server error for {resource_id}: " f"HTTP {http_status} - {error_message}. Will retry." ) # All other errors are unrecoverable raise LabTaskUnrecoverableError( f"{operation} failed for {resource_id}: {error_code} - {error_message}" ) class AWSAuroraServerlessClusterCreateTask(BaseLabTask): """ Create AWS Aurora Serverless v2 cluster and instance. This task: 1. Generates and immediately persists credentials (before any AWS API call) 2. Creates the DB cluster (handles idempotency if cluster exists from retry) 3. Creates the DB instance (handles idempotency if instance exists from retry) Note: The cluster endpoint is NOT available until the cluster is ready. Use AWSAuroraServerlessClusterReadyCheckTask to wait and retrieve the endpoint. """ max_retries: int = 20 retry_interval_seconds: int = 30 # Cluster configuration DATABASE_NAME = "lab_demo" ENGINE = "aurora-mysql" ENGINE_MODE = "provisioned" PORT = 3306 BACKUP_RETENTION_PERIOD = 1 SERVERLESS_MIN_CAPACITY = 0 # Minimum ACU for serverless v2 SERVERLESS_MAX_CAPACITY = 1.5 # Maximum ACU for serverless v2 INSTANCE_CLASS = "db.serverless" STORAGE_TYPE = "aurora" def run(self): db_lab_instance = self.db_lab_instance # Generate identifiers cluster_id = f"{db_lab_instance.unique_tag}-c" instance_id = f"{db_lab_instance.unique_tag}-i" # Step 1: Check if we already have saved credentials (from a previous attempt) # Use select_for_update to prevent race conditions with other parallel tasks from django.db import transaction from labs.models import LabInstance as DBLabInstance with transaction.atomic(): locked_instance = DBLabInstance.objects.select_for_update().get(pk=db_lab_instance.pk) lab_resource = LabResource.from_db_lab_instance_resource(locked_instance.resources) if lab_resource.aws_aurora_serverless and lab_resource.aws_aurora_serverless.master_username: # Reuse existing credentials username = lab_resource.aws_aurora_serverless.master_username password = lab_resource.aws_aurora_serverless.master_password logger.info(f"Reusing existing Aurora credentials for {db_lab_instance.unique_tag}") else: # Generate new credentials and save immediately username, password = _generate_aurora_admin_credentials() aurora_config = AWSAuroraServerlessClusterConfig( cluster_identifier=cluster_id, instance_identifier=instance_id, port=self.PORT, endpoint=None, # Will be set when cluster is ready reader_endpoint=None, database_name=self.DATABASE_NAME, master_username=username, master_password=password, ) lab_resource.aws_aurora_serverless = aurora_config locked_instance.resources = lab_resource.to_db_lab_instance_resource() locked_instance.save(update_fields=["resources"]) logger.info(f"Saved Aurora credentials for {db_lab_instance.unique_tag}") # Step 2: Create the cluster rds = _get_rds_client(db_lab_instance.region_name) self._create_cluster(rds, cluster_id, username, password) # Step 3: Create the instance self._create_instance(rds, cluster_id, instance_id) logger.info( f"Aurora Serverless cluster creation initiated for {db_lab_instance.unique_tag}: " f"cluster={cluster_id}, instance={instance_id}" ) def _create_cluster(self, rds, cluster_id: str, username: str, password: str) -> None: """Create the Aurora DB cluster.""" lab_region = get_region_metadata(self.db_lab_instance.region_name) db_subnet_group = lab_region.get("aurora-db-subnet-group") security_group_id = lab_region.get("aurora-security-group-id") if not db_subnet_group or not security_group_id: raise LabTaskUnrecoverableError( f"Aurora network config incomplete for region {self.db_lab_instance.region_name}: " f"aurora-db-subnet-group={db_subnet_group}, aurora-security-group-id={security_group_id}" ) try: resp = rds.create_db_cluster( DBClusterIdentifier=cluster_id, Engine=self.ENGINE, EngineMode=self.ENGINE_MODE, Port=self.PORT, MasterUsername=username, MasterUserPassword=password, DatabaseName=self.DATABASE_NAME, BackupRetentionPeriod=self.BACKUP_RETENTION_PERIOD, ServerlessV2ScalingConfiguration={ "MinCapacity": self.SERVERLESS_MIN_CAPACITY, "MaxCapacity": self.SERVERLESS_MAX_CAPACITY, }, StorageType=self.STORAGE_TYPE, DBSubnetGroupName=db_subnet_group, VpcSecurityGroupIds=[security_group_id], Tags=[ {"Key": "trainer", "Value": self.db_lab_instance.unique_tag}, {"Key": "lab_instance_id", "Value": str(self.db_lab_instance.id)}, {"Key": "created_by", "Value": "labs.tidb.io"}, ], ) http_status = resp.get("ResponseMetadata", {}).get("HTTPStatusCode") if http_status != 200: raise LabTaskRecoverableError( f"Unexpected HTTP status {http_status} when creating cluster {cluster_id}" ) logger.info(f"Created Aurora cluster {cluster_id}, status: {resp.get('DBCluster', {}).get('Status')}") except ClientError as e: result = _handle_client_error( e, operation="Create Aurora cluster", resource_id=cluster_id, allow_already_exists=True, # May exist from previous attempt ) if result == "already_exists": # Verify the existing cluster matches our lab self._verify_existing_cluster(rds, cluster_id) def _verify_existing_cluster(self, rds, cluster_id: str) -> None: """Verify that an existing cluster belongs to this lab instance.""" try: resp = rds.describe_db_clusters(DBClusterIdentifier=cluster_id) clusters = resp.get("DBClusters", []) if not clusters: raise LabTaskUnrecoverableError( f"Cluster {cluster_id} reported as existing but describe returned empty" ) cluster = clusters[0] tags = {tag["Key"]: tag["Value"] for tag in cluster.get("TagList", [])} # Verify ownership if tags.get("trainer") != self.db_lab_instance.unique_tag: raise LabTaskUnrecoverableError( f"Cluster {cluster_id} exists but belongs to different lab: " f"expected {self.db_lab_instance.unique_tag}, got {tags.get('trainer')}" ) # Check status - if it's failed, we can't proceed status = cluster.get("Status") if status in ("failed", "deleting", "deleted"): raise LabTaskUnrecoverableError( f"Cluster {cluster_id} exists but is in terminal status: {status}" ) logger.info(f"Verified existing cluster {cluster_id} belongs to this lab, status: {status}") except ClientError as e: _handle_client_error(e, operation="Describe Aurora cluster", resource_id=cluster_id) def _create_instance(self, rds, cluster_id: str, instance_id: str) -> None: """Create the Aurora DB instance in the cluster.""" try: resp = rds.create_db_instance( DBInstanceIdentifier=instance_id, DBClusterIdentifier=cluster_id, Engine=self.ENGINE, DBInstanceClass=self.INSTANCE_CLASS, PubliclyAccessible=True, # Tags for cost tracking Tags=[ {"Key": "trainer", "Value": self.db_lab_instance.unique_tag}, {"Key": "lab_instance_id", "Value": str(self.db_lab_instance.id)}, {"Key": "created_by", "Value": "labs.tidb.io"}, ], ) http_status = resp.get("ResponseMetadata", {}).get("HTTPStatusCode") if http_status != 200: raise LabTaskRecoverableError( f"Unexpected HTTP status {http_status} when creating instance {instance_id}" ) logger.info( f"Created Aurora instance {instance_id}, status: {resp.get('DBInstance', {}).get('DBInstanceStatus')}" ) except ClientError as e: result = _handle_client_error( e, operation="Create Aurora instance", resource_id=instance_id, allow_already_exists=True, # May exist from previous attempt ) if result == "already_exists": self._verify_existing_instance(rds, cluster_id, instance_id) def _verify_existing_instance(self, rds, cluster_id: str, instance_id: str) -> None: """Verify that an existing instance belongs to this lab instance.""" try: resp = rds.describe_db_instances(DBInstanceIdentifier=instance_id) instances = resp.get("DBInstances", []) if not instances: raise LabTaskUnrecoverableError( f"Instance {instance_id} reported as existing but describe returned empty" ) instance = instances[0] # Verify it belongs to our cluster if instance.get("DBClusterIdentifier") != cluster_id: raise LabTaskUnrecoverableError( f"Instance {instance_id} exists but belongs to different cluster: " f"expected {cluster_id}, got {instance.get('DBClusterIdentifier')}" ) # Check status status = instance.get("DBInstanceStatus") if status in ("failed", "deleting", "deleted"): raise LabTaskUnrecoverableError( f"Instance {instance_id} exists but is in terminal status: {status}" ) logger.info(f"Verified existing instance {instance_id} belongs to this lab, status: {status}") except ClientError as e: _handle_client_error(e, operation="Describe Aurora instance", resource_id=instance_id) class AWSAuroraServerlessClusterReadyCheckTask(BaseLabTask): """ Wait for Aurora Serverless cluster and instance to be available. This task: 1. Checks cluster status until it's "available" 2. Checks instance status until it's "available" 3. Retrieves and saves the cluster endpoint This is a polling task - it will raise LabTaskRecoverableError until the resources are ready, allowing the task runner to retry. """ max_retries: int = 60 # Aurora can take 10-20 minutes to be available retry_interval_seconds: int = 30 def run(self): db_lab_instance = self.db_lab_instance cluster_id = f"{db_lab_instance.unique_tag}-c" instance_id = f"{db_lab_instance.unique_tag}-i" rds = _get_rds_client(db_lab_instance.region_name) # Check cluster status and get endpoint cluster_info = self._check_cluster_ready(rds, cluster_id) # Check instance status self._check_instance_ready(rds, instance_id) # Save the endpoint self._save_endpoint(cluster_info) logger.info( f"Aurora Serverless cluster ready for {db_lab_instance.unique_tag}: " f"endpoint={cluster_info['endpoint']}" ) def _check_cluster_ready(self, rds, cluster_id: str) -> dict: """Check if cluster is available and return cluster info.""" try: resp = rds.describe_db_clusters(DBClusterIdentifier=cluster_id) except ClientError as e: _handle_client_error(e, operation="Describe Aurora cluster", resource_id=cluster_id) clusters = resp.get("DBClusters", []) if not clusters: raise LabTaskRecoverableError( f"Aurora cluster {cluster_id} not found yet, will retry" ) cluster = clusters[0] status = cluster.get("Status") logger.info(f"Aurora cluster {cluster_id} status: {status}") # Terminal failure states if status in ("failed", "deleted", "deleting", "inaccessible-encryption-credentials"): raise LabTaskUnrecoverableError( f"Aurora cluster {cluster_id} is in terminal failed state: {status}" ) # Not ready yet if status != "available": raise LabTaskRecoverableError( f"Aurora cluster {cluster_id} is not ready yet, status: {status}" ) # Cluster is available - extract endpoint info endpoint = cluster.get("Endpoint") reader_endpoint = cluster.get("ReaderEndpoint") port = cluster.get("Port", 3306) if not endpoint: raise LabTaskRecoverableError( f"Aurora cluster {cluster_id} is available but endpoint not yet assigned" ) return { "endpoint": endpoint, "reader_endpoint": reader_endpoint, "port": port, } def _check_instance_ready(self, rds, instance_id: str) -> None: """Check if instance is available.""" try: resp = rds.describe_db_instances(DBInstanceIdentifier=instance_id) except ClientError as e: _handle_client_error(e, operation="Describe Aurora instance", resource_id=instance_id) instances = resp.get("DBInstances", []) if not instances: raise LabTaskRecoverableError( f"Aurora instance {instance_id} not found yet, will retry" ) instance = instances[0] status = instance.get("DBInstanceStatus") logger.info(f"Aurora instance {instance_id} status: {status}") # Terminal failure states if status in ("failed", "deleted", "deleting", "incompatible-restore", "incompatible-network"): raise LabTaskUnrecoverableError( f"Aurora instance {instance_id} is in terminal failed state: {status}" ) # Not ready yet if status != "available": raise LabTaskRecoverableError( f"Aurora instance {instance_id} is not ready yet, status: {status}" ) def _save_endpoint(self, cluster_info: dict) -> None: """Save the cluster endpoint to the lab instance resources.""" from django.db import transaction from labs.models import LabInstance as DBLabInstance db_lab_instance = self.db_lab_instance # Use select_for_update to prevent race conditions with other parallel tasks with transaction.atomic(): locked_instance = DBLabInstance.objects.select_for_update().get(pk=db_lab_instance.pk) lab_resource = LabResource.from_db_lab_instance_resource(locked_instance.resources) if not lab_resource.aws_aurora_serverless: raise LabTaskUnrecoverableError( f"Aurora config not found for lab {db_lab_instance.unique_tag}. " "AWSAuroraServerlessClusterCreateTask should have saved it." ) # Update with endpoint info lab_resource.aws_aurora_serverless.endpoint = cluster_info["endpoint"] lab_resource.aws_aurora_serverless.reader_endpoint = cluster_info["reader_endpoint"] lab_resource.aws_aurora_serverless.port = cluster_info["port"] locked_instance.resources = lab_resource.to_db_lab_instance_resource() locked_instance.save(update_fields=["resources"]) logger.info( f"Saved Aurora endpoint for {db_lab_instance.unique_tag}: {cluster_info['endpoint']}" ) class AWSAuroraServerlessClusterDeleteTask(BaseLabTask): """ Delete Aurora Serverless cluster and instance. This task: 1. Deletes the DB instance first (required before cluster deletion) 2. Waits for instance deletion to complete 3. Deletes the DB cluster (with skip-final-snapshot) 4. Handles "not found" gracefully (idempotent for retries) Note: This task may need multiple retries as deletion can take time. """ max_retries: int = 30 retry_interval_seconds: int = 30 def run(self): db_lab_instance = self.db_lab_instance db_lab_instance.refresh_from_db(fields=["resources"]) lab_resource = LabResource.from_db_lab_instance_resource(db_lab_instance.resources) if not lab_resource.aws_aurora_serverless: raise LabTaskUnrecoverableError( f"AWS Aurora Serverless config not found for lab {db_lab_instance.unique_tag}. " "AWSAuroraServerlessClusterCreateTask should have saved it." ) cluster_id = lab_resource.aws_aurora_serverless.cluster_identifier instance_id = lab_resource.aws_aurora_serverless.instance_identifier rds = _get_rds_client(db_lab_instance.region_name) self._delete_instance(rds, instance_id) self._delete_cluster(rds, cluster_id) logger.info( f"Aurora Serverless cluster and instance deletion initiated for {db_lab_instance.unique_tag}" ) def _delete_instance(self, rds, instance_id: str) -> str: """ Delete the DB instance. Returns: "deleted" if instance is gone "deleting" if deletion is in progress """ # First, check if instance exists try: resp = rds.describe_db_instances(DBInstanceIdentifier=instance_id) instances = resp.get("DBInstances", []) if not instances: logger.info(f"Aurora instance {instance_id} not found (already deleted)") return "deleted" status = instances[0].get("DBInstanceStatus") logger.info(f"Aurora instance {instance_id} current status: {status}") if status == "deleting": return "deleting" if status == "deleted": return "deleted" except ClientError as e: error_code = e.response.get("Error", {}).get("Code", "") if error_code in NOT_FOUND_ERROR_CODES: logger.info(f"Aurora instance {instance_id} not found (already deleted)") return "deleted" _handle_client_error(e, operation="Describe Aurora instance", resource_id=instance_id) # Instance exists and is not deleting - delete it try: resp = rds.delete_db_instance( DBInstanceIdentifier=instance_id, SkipFinalSnapshot=True, DeleteAutomatedBackups=True, ) http_status = resp.get("ResponseMetadata", {}).get("HTTPStatusCode") if http_status != 200: raise LabTaskRecoverableError( f"Unexpected HTTP status {http_status} when deleting instance {instance_id}" ) logger.info(f"Initiated deletion of Aurora instance {instance_id}") return "deleting" except ClientError as e: result = _handle_client_error( e, operation="Delete Aurora instance", resource_id=instance_id, allow_not_found=True, # May have been deleted by another process ) if result == "not_found": return "deleted" raise # Re-raise if _handle_client_error didn't handle it def _delete_cluster(self, rds, cluster_id: str) -> str: """ Delete the DB cluster. Returns: "deleted" if cluster is gone "deleting" if deletion is in progress """ # First, check if cluster exists try: resp = rds.describe_db_clusters(DBClusterIdentifier=cluster_id) clusters = resp.get("DBClusters", []) if not clusters: logger.info(f"Aurora cluster {cluster_id} not found (already deleted)") return "deleted" status = clusters[0].get("Status") logger.info(f"Aurora cluster {cluster_id} current status: {status}") if status in ("deleting", "deleted"): return status except ClientError as e: error_code = e.response.get("Error", {}).get("Code", "") if error_code in NOT_FOUND_ERROR_CODES: logger.info(f"Aurora cluster {cluster_id} not found (already deleted)") return "deleted" _handle_client_error(e, operation="Describe Aurora cluster", resource_id=cluster_id) # Cluster exists, has no members, and is not deleting - delete it try: resp = rds.delete_db_cluster( DBClusterIdentifier=cluster_id, SkipFinalSnapshot=True, DeleteAutomatedBackups=True, ) http_status = resp.get("ResponseMetadata", {}).get("HTTPStatusCode") if http_status != 200: raise LabTaskRecoverableError( f"Unexpected HTTP status {http_status} when deleting cluster {cluster_id}" ) logger.info(f"Initiated deletion of Aurora cluster {cluster_id}") return "deleting" except ClientError as e: result = _handle_client_error( e, operation="Delete Aurora cluster", resource_id=cluster_id, allow_not_found=True, # May have been deleted by another process ) if result == "not_found": return "deleted" raise # Re-raise if _handle_client_error didn't handle it