import logging import json import importlib.resources from pathlib import Path import yaml import enum from typing import ( List, Literal, Optional, Dict, Union, Annotated, Any, ) from pydantic import ( BaseModel, ConfigDict, Field, PositiveInt, model_validator, computed_field, ValidationError, ) from labs.lab_manager_v2.lab_guide import LabGuide, load_lab_guide from labs.lab_manager_v2.exceptions import ( LabNotFoundError, LabSpecValidationError, ) from labs.models import User logger = logging.getLogger(__name__) class DeliveryType(str, enum.Enum): SELF_PACED = "Self-Paced" INSTRUCTOR_LED = "Instructor-Led" class TiDBEnvironmentType(str, enum.Enum): TIDB_SELF_MANAGED = "TiDB Self-Managed" TIDB_CLOUD_STARTER = "TiDB Cloud Starter" TIDB_CLOUD_ESSENTIAL = "TiDB Cloud Essential" TIDB_CLOUD_DEDICATED = "TiDB Cloud Dedicated" class Region(str, enum.Enum): OREGON = "oregon" TOKYO = "tokyo" SINGAPORE = "singapore" FRANKFURT = "frankfurt" class RegionConfig(BaseModel): """Configuration for regions with cloud provider support.""" aws: Optional[List[Region]] = Field(default_factory=list) alicloud: Optional[List[Region]] = Field(default_factory=list) def get_all_regions(self) -> List[Region]: """Get all regions from all cloud providers.""" all_regions = [] if self.aws: all_regions.extend(self.aws) if self.alicloud: all_regions.extend(self.alicloud) return all_regions def get_supported_cloud_providers(self) -> List[str]: """Get all supported cloud providers.""" supported_cloud_providers = [] if self.aws: supported_cloud_providers.append("aws") if self.alicloud: supported_cloud_providers.append("alicloud") return supported_cloud_providers def get_supported_regions(self, cloud_provider: str) -> List[Region]: """Get all supported regions for the given cloud provider.""" if cloud_provider == "aws": return self.aws elif cloud_provider == "alicloud": return self.alicloud else: raise ValueError(f"Invalid cloud provider: {cloud_provider}") class ExternalMaterial(BaseModel): name: str url: str employee_only: bool = False LABS_SPEC_DIR = Path(__file__).parent.parent / "labs" class TiDBServerlessClusterRequirement(BaseModel): type: Literal["TiDBServerlessCluster"] _lab_spec: Optional["LabSpec"] = None class AWSAuroraServerlessClusterRequirement(BaseModel): type: Literal["AWSAuroraServerlessCluster"] _lab_spec: Optional["LabSpec"] = None class CloudStackRequirement(BaseModel): # "AWSCloudFormationStack" is kept for backwards compatibility with existing lab specs. type: Literal["CloudStack", "AWSCloudFormationStack"] _lab_spec: Optional["LabSpec"] = None def read_aws_cfn_template(self) -> str: if not self._lab_spec: raise ValueError( "LabSpec reference not set. This step must be part of a LabSpec." ) return (self._lab_spec._base_dir / "cloudformation.json").read_text( encoding="utf-8" ) def read_alicloud_ros_template(self) -> str: if not self._lab_spec: raise ValueError( "LabSpec reference not set. This step must be part of a LabSpec." ) return (self._lab_spec._base_dir / "alicloud_ros.json").read_text( encoding="utf-8" ) # Backwards-compat alias (deprecated): prefer CloudStackRequirement. AWSCloudFormationStackRequirement = CloudStackRequirement Requirement = Annotated[ Union[ TiDBServerlessClusterRequirement, AWSAuroraServerlessClusterRequirement, CloudStackRequirement, ], Field(discriminator="type"), ] class BootstrapShellScriptStep(BaseModel): name: str type: Literal["ShellScript"] target: List[str] shell: Optional[str] = None script: Optional[Path] = None _lab_spec: Optional["LabSpec"] = None @model_validator(mode="after") def ensure_one_of_shell_or_script(self) -> "BootstrapShellScriptStep": if not (self.shell or self.script): raise ValueError("shell step must provide either `shell:` or `script:`") if self.shell and self.script: raise ValueError( "shell step must provide only one of `shell:` or `script:`" ) return self @model_validator(mode="after") def ensure_at_least_one_target(self) -> "BootstrapShellScriptStep": if not self.target: raise ValueError("shell step must provide at least one `target:`") return self def get_shell_script(self) -> str: """Get the shell script content with proper path resolution.""" if self.shell: return self.shell if self.script: if not self._lab_spec: raise ValueError( "LabSpec reference not set. This step must be part of a LabSpec instance." ) return self._lab_spec.get_shell_script(self.script) BootstrapStep = Annotated[ Union[BootstrapShellScriptStep,], Field(discriminator="type"), ] class CleanupStep(BootstrapStep): pass class TopologyNode(BaseModel): role: str node_count: int web_ssh: bool = False https_reverse_proxy: bool = False ssh_user: Optional[str] = None @model_validator(mode="before") @classmethod def forbid_legacy_fields(cls, data: Any) -> Any: if isinstance(data, dict): legacy_fields = [k for k in ("entrypoint", "secure_access") if k in data] if legacy_fields: raise ValueError( f"legacy fields are not supported: {legacy_fields}. " "Please use `web_ssh` and `https_reverse_proxy` instead." ) return data @model_validator(mode="after") def ensure_ssh_user_exists_when_enable_web_ssh_or_https_reverse_proxy( self, ) -> "TopologyNode": if self.web_ssh or self.https_reverse_proxy: if not self.ssh_user: raise ValueError( "topology node must provide `ssh_user` when `web_ssh` or `https_reverse_proxy` is enabled" ) if self.https_reverse_proxy and self.ssh_user and self.ssh_user != "ubuntu": raise ValueError( "topology node must use `ssh_user: ubuntu` when `https_reverse_proxy` is enabled" ) return self class LabSpec(BaseModel): model_config = ConfigDict(use_enum_values=True) lab_code: str lab_level: str title_en: str title_ja: str title_zh: str title_ko: str description_en: str description_ja: str description_zh: str description_ko: str splash_image_en: str splash_image_ja: str splash_image_zh: str splash_image_ko: str category_en: List[str] = Field(default_factory=list) category_ja: List[str] = Field(default_factory=list) category_zh: List[str] = Field(default_factory=list) category_ko: List[str] = Field(default_factory=list) external_materials_en: List[ExternalMaterial] = Field(default_factory=list) external_materials_ja: List[ExternalMaterial] = Field(default_factory=list) external_materials_zh: List[ExternalMaterial] = Field(default_factory=list) external_materials_ko: List[ExternalMaterial] = Field(default_factory=list) authors: List[str] = Field(default_factory=list) target_audience: List[str] = Field(default_factory=list) delivery_type: DeliveryType estimated_duration_minutes: int = 60 features: Dict[str, Any] = Field(default_factory=dict) credits: PositiveInt credit_redeemable: bool employee_only: bool regions: RegionConfig tidb_versions: List[str] = Field(default_factory=list) tidb_environments: List[TiDBEnvironmentType] = Field(default_factory=list) topology: List[TopologyNode] requirements: List[Requirement] bootstrap: List[BootstrapStep] = Field(default_factory=list) # TODO: cleanup is not supported yet # cleanup: Optional[List[CleanupStep]] = None @computed_field @property def _base_dir(self) -> Path: """Get the base directory for this lab based on its name.""" current_file = Path(__file__) labs_root = current_file.parent.parent return labs_root / "labs" / self.lab_code def get_shell_script(self, script: Path) -> str: """Get the shell script content with proper path resolution.""" script_path = self._base_dir / script return script_path.read_text(encoding="utf-8") def model_post_init(self, __context) -> None: """Set LabSpec reference for all bootstrap steps after initialization.""" if self.bootstrap: for step in self.bootstrap: if isinstance(step, BootstrapShellScriptStep): step._lab_spec = self if self.requirements: for req in self.requirements: if isinstance(req, CloudStackRequirement): req._lab_spec = self elif isinstance(req, TiDBServerlessClusterRequirement): req._lab_spec = self elif isinstance(req, AWSAuroraServerlessClusterRequirement): req._lab_spec = self @model_validator(mode="after") def ensure_at_least_one_region(self) -> "LabSpec": if not self.regions.get_all_regions(): raise ValueError("please provide at least one region") return self @model_validator(mode="after") def ensure_at_least_one_tidb_environment(self) -> "LabSpec": if not self.tidb_environments: raise ValueError("please provide at least one TiDB environment") return self @model_validator(mode="after") def validate_bootstrap_steps(self) -> "LabSpec": ssh_nodes = set() for node in self.topology: if node.ssh_user: ssh_nodes.update( [f"{node.role}{i + 1}" for i in range(node.node_count)] ) for step in self.bootstrap: if isinstance(step, BootstrapShellScriptStep): step_targets = set(step.target) if not step_targets.issubset(ssh_nodes): raise ValueError( f"target {step_targets - ssh_nodes} not found in nodes with ssh_user" ) return self @model_validator(mode="after") def validate_estimated_duration_minutes(self) -> "LabSpec": if ( self.estimated_duration_minutes < 10 or self.estimated_duration_minutes > 300 ): raise ValueError( "estimated_duration_minutes must be between 10 and 300 minutes" ) return self @model_validator(mode="after") def validate_requirements(self) -> "LabSpec": aws_aurora_requirements = [ req for req in self.requirements if isinstance(req, AWSAuroraServerlessClusterRequirement) ] if aws_aurora_requirements: if not self.regions.aws: raise ValueError( "AWSAuroraServerlessCluster requires regions.aws to be non-empty" ) if self.regions.alicloud: raise ValueError( "AWSAuroraServerlessCluster does not support regions.alicloud" ) return self def get_lab_guide( self, language_code: str, fallback_to_other_language: bool = True, ) -> LabGuide: """Get the lab guide for the given language code.""" return load_lab_guide( lab_code=self.lab_code, language_code=language_code, lab_base_dir=self._base_dir, fallback_to_other_language=fallback_to_other_language, ) def to_frontend(self, user: Optional[User] = None) -> Dict: data = self.model_dump( exclude={ "_base_dir", "topology", "requirements", "bootstrap", }, ) employee = False if ( user and user.is_authenticated and user.email and user.email.strip().lower().endswith("@pingcap.com") ): employee = True if not employee: for key in [ "external_materials_en", "external_materials_ja", "external_materials_zh", "external_materials_ko", ]: materials = data.get(key) or [] # items are plain dicts after model_dump data[key] = [m for m in materials if not m.get("employee_only")] return data def load_lab_spec(lab_code: str) -> LabSpec: """Load a lab spec from a YAML file.""" lab_spec_path = LABS_SPEC_DIR / lab_code / "manifest.yaml" if not lab_spec_path.exists(): raise LabNotFoundError(f"Lab spec not found: {lab_code}") try: return LabSpec.model_validate( yaml.safe_load(lab_spec_path.read_text(encoding="utf-8")) ) except ValidationError as e: raise LabSpecValidationError(f"Lab spec validation error: {e}") from e # This is a temporary workaround to serve the lab list to the frontend, # we will migrate this static file to database & cache in the future ALL_LABS_JSON_PATH = "labs/lab_manager_v2/temp/all_labs.json" def generate_lab_list() -> None: labs = {} for dir in LABS_SPEC_DIR.iterdir(): if dir.is_dir() and (dir / "manifest.yaml").exists(): lab_spec = load_lab_spec(dir.name) _lab = lab_spec.to_frontend() _lab["migrate_to_v2"] = True _lab["ui_duration_before_lab_ready_seconds"] = 180 labs[lab_spec.lab_code] = _lab labs = dict(sorted(labs.items(), key=lambda x: x[0])) json.dump(labs, open(ALL_LABS_JSON_PATH, "w"), indent=2, ensure_ascii=False) def get_all_labs() -> dict: with open(ALL_LABS_JSON_PATH, "r") as f: return json.load(f)