import argparse import os from dataclasses import asdict, dataclass, field from typing import List, Optional @dataclass class LogLocation: bucketUri: str roleArn: str = "" externalId: str = "" accountId: str = "" type: str = "TYPE_STANDARD" @dataclass class OPSAuditLogSetting: logLocation: LogLocation enabled: bool = True @dataclass class RegionWhiteList: servicePlan: str isPublic: bool = False openedTenantIds: List[str] = field(default_factory=list) @dataclass class RegionDetail: id: str cloudProviderId: str displayName: str area: str countryCode: str defaultBackupTime: str serviceAddress: str logicalName: str servicePlan: str sharedBucketUrl: str = "" state: str = "INACTIVE" regionWhiteLists: List[RegionWhiteList] = field(default_factory=list) sortPriority: int = 0 isPublic: bool = False ops_audit_log_setting: Optional[OPSAuditLogSetting] = None @dataclass class RegionProviderIdentity: identity: str @dataclass class Region: deployType: str region: RegionDetail regionProviderIdentities: List[RegionProviderIdentity] def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description="Register region to tidbcloud global") parser.add_argument( "--env", type=str, required=True, help="environment of the logical region", choices=["dev", "staging", "prod"], ) parser.add_argument( "--deploy-type", type=str, required=True, help="type of the logical region", choices=["DEPLOY_TYPE_DEDICATED", "DEPLOY_TYPE_NEXTGEN"], ) parser.add_argument( "--provider", type=str, required=True, help="type of cloud", choices=["aws", "gcp", "azure", "alicloud"], ) parser.add_argument( "--region", type=str, required=True, help="name of cloud region", ) parser.add_argument( "--logical-region-name", type=str, required=True, help="logical region name of tidb cloud region", ) parser.add_argument( "--display-name", type=str, required=True, help="display name of the region", ) parser.add_argument( "--area", type=str, required=True, help="area of the region", ) parser.add_argument( "--country-code", type=str, required=True, help="country code of the region", ) parser.add_argument( "--default-backup-time", type=str, required=True, help="default backup time for the region", ) parser.add_argument( "--service-plan", type=str, help="plan of the region", default="Premium", choices=["Starter", "Essential", "Premium", "BYOC"], ) parser.add_argument( "--service-address", type=str, required=True, help="address of regional server", ) parser.add_argument( "--dataplane-accounts", type=str, required=False, default=[], nargs="+", help="account of dataplane", ) parser.add_argument( "--state", type=str, required=False, default="INACTIVE", help="state of the region", choices=["ACTIVE", "INACTIVE"], ) parser.add_argument( "--audit-log-bucket", type=str, required=False, help="bucket of audit log for BYOC", ) parser.add_argument( "--audit-log-storage-account-name", type=str, required=False, help="storage account name of audit log for Azure BYOC", ) parser.add_argument( "--shared-bucket", type=str, required=False, help="shared bucket for sli purpose", ) parser.add_argument( "--tenant-ids", type=str, required=False, default="", help="Comma-separated TiDB Cloud tenant (org) ids to whitelist for this region; " "when non-empty, the ids are added to regionWhiteLists under the matching servicePlan entry. " "Empty string means no whitelist entry is written.", ) return parser.parse_args() def create_or_update_region(method: callable, url: str, args: argparse.Namespace): if args.service_plan == "BYOC" and args.provider == "aws": if not args.audit_log_bucket: raise ValueError("AWS BYOC requires --audit-log-bucket") if not args.shared_bucket: raise ValueError("AWS BYOC requires --shared-bucket") if args.service_plan == "BYOC" and args.provider == "azure": if not args.audit_log_storage_account_name: raise ValueError("Azure BYOC requires --audit-log-storage-account-name") if not args.audit_log_bucket: raise ValueError("Azure BYOC requires --audit-log-bucket") r = Region( deployType=args.deploy_type, region=RegionDetail( id=f"{args.provider}_{args.region}".upper().replace("-", "_"), cloudProviderId=args.provider.upper(), displayName=args.display_name, area=args.area, countryCode=args.country_code, defaultBackupTime=args.default_backup_time, serviceAddress=args.service_address, logicalName=args.logical_region_name, servicePlan=args.service_plan, state=args.state, ), regionProviderIdentities=[ RegionProviderIdentity(account_id) for account_id in args.dataplane_accounts ], ) if args.service_plan == "BYOC" and args.provider == "aws": daccount = args.dataplane_accounts[0] r.region.ops_audit_log_setting = OPSAuditLogSetting( logLocation=LogLocation( bucketUri=f"s3://{args.audit_log_bucket}", roleArn=f"arn:aws:iam::{daccount}:role/tidbcloud-audit-log-write-only-role", accountId=daccount, ) ) r.region.sharedBucketUrl = f"s3://{args.shared_bucket}" elif args.service_plan == "BYOC" and args.provider == "azure": daccount = args.dataplane_accounts[0] r.region.ops_audit_log_setting = OPSAuditLogSetting( logLocation=LogLocation( bucketUri=f"https://{args.audit_log_storage_account_name}.blob.core.windows.net/{args.audit_log_bucket}", accountId=daccount, ) ) tenant_ids = [tid.strip() for tid in args.tenant_ids.split(",") if tid.strip()] if tenant_ids: # Mirror .github/workflows/update-region.yml: tenants are whitelisted for # the region via regionWhiteLists (per-servicePlan), not via the legacy # top-level openedTenantIds field. Currently BYOC is the only consumer, # but the channel is generic so future service plans (Premium, etc.) # can reuse the same plumbing. r.region.regionWhiteLists = [ RegionWhiteList( servicePlan=args.service_plan, isPublic=False, openedTenantIds=tenant_ids, ) ] response = method( url, json=asdict(r), headers={ "Authorization": f"Basic {os.environ['GLOBAL_AUTH_KEY']}", }, ) if response.status_code >= 300: print(asdict(r)) raise Exception( f"register/update {args.logical_region_name} failed, global response: \n{response.text}" )