"""Generate integration-test active_cluster_ids SQL from a backup protobuf. This module uses only the Python standard library. It parses the protobuf wire format directly, so generated bindings, protoc, pip packages, and compilation are not required. """ import random from datetime import datetime, timedelta from typing import Dict, Iterator, List, Tuple, Union UPDATED_AT_WINDOW_MICROSECONDS = 60 * 60 * 1_000_000 KEYSPACE_STATE_ENABLED = 0 ProtobufValue = Union[int, bytes] def read_varint(data: bytes, position: int) -> Tuple[int, int]: value = 0 for byte_index in range(10): if position >= len(data): raise ValueError("truncated protobuf varint") byte = data[position] position += 1 if byte_index == 9 and byte > 1: raise ValueError("protobuf varint exceeds 64 bits") value |= (byte & 0x7F) << (byte_index * 7) if byte < 0x80: return value, position raise ValueError("protobuf varint is too long") def iter_protobuf_fields( data: bytes, ) -> Iterator[Tuple[int, int, ProtobufValue]]: position = 0 while position < len(data): tag, position = read_varint(data, position) field_number = tag >> 3 wire_type = tag & 7 if field_number == 0: raise ValueError("invalid protobuf field number 0") if wire_type == 0: value, position = read_varint(data, position) elif wire_type == 1: end = position + 8 if end > len(data): raise ValueError(f"truncated protobuf field {field_number}") value, position = data[position:end], end elif wire_type == 2: size, position = read_varint(data, position) end = position + size if end > len(data): raise ValueError(f"truncated protobuf field {field_number}") value, position = data[position:end], end elif wire_type == 5: end = position + 4 if end > len(data): raise ValueError(f"truncated protobuf field {field_number}") value, position = data[position:end], end else: raise ValueError( f"unsupported protobuf wire type {wire_type} " f"for field {field_number}" ) yield field_number, wire_type, value def parse_bytes_map_entry(data: bytes) -> Tuple[bytes, bytes]: key = None value = None for field_number, wire_type, field_value in iter_protobuf_fields(data): if field_number == 1 and wire_type == 2: key = field_value elif field_number == 2 and wire_type == 2: value = field_value if not isinstance(key, bytes) or not isinstance(value, bytes): raise ValueError("invalid protobuf bytes map entry") return key, value def decode_utf8(value: bytes, description: str) -> str: try: return value.decode("utf-8") except UnicodeDecodeError as error: raise ValueError(f"{description} is not valid UTF-8") from error def parse_keyspace_meta(data: bytes) -> Tuple[int, int, Dict[str, str]]: # kvproto keyspacepb.KeyspaceMeta: # id = 1, state = 3, config = 7. ENABLED is the protobuf default (0), so # the state field is normally absent for keyspaces that Pack can select. keyspace_id = 0 keyspace_state = KEYSPACE_STATE_ENABLED config = {} for field_number, wire_type, value in iter_protobuf_fields(data): if field_number == 1 and wire_type == 0: if not isinstance(value, int): raise ValueError("invalid KeyspaceMeta.id") keyspace_id = value elif field_number == 3: if wire_type != 0 or not isinstance(value, int): raise ValueError("invalid KeyspaceMeta.state") keyspace_state = value elif field_number == 7 and wire_type == 2: if not isinstance(value, bytes): raise ValueError("invalid KeyspaceMeta.config") key, config_value = parse_bytes_map_entry(value) config_key = decode_utf8(key, "KeyspaceMeta config key") config[config_key] = decode_utf8( config_value, "KeyspaceMeta config value", ) return keyspace_id, keyspace_state, config def load_cluster_ids( backup_meta: bytes, source_label: str, ) -> Tuple[List[str], int, int, int, int]: # rfenginepb.ClusterBackupMeta: # cluster_id = 2, keyspace_meta = 6, is_lightweight = 8. backup_cluster_id = None is_lightweight = False raw_pd_meta = {} for field_number, wire_type, value in iter_protobuf_fields(backup_meta): if field_number == 2 and wire_type == 0: if not isinstance(value, int): raise ValueError("invalid ClusterBackupMeta.cluster_id") backup_cluster_id = value elif field_number == 6 and wire_type == 2: if not isinstance(value, bytes): raise ValueError("invalid ClusterBackupMeta.keyspace_meta") key, pd_value = parse_bytes_map_entry(value) if key in raw_pd_meta: raise ValueError(f"duplicate PD metadata key: {key!r}") raw_pd_meta[key] = pd_value elif field_number == 8 and wire_type == 0: if not isinstance(value, int): raise ValueError("invalid ClusterBackupMeta.is_lightweight") is_lightweight = bool(value) if backup_cluster_id is None: raise ValueError("ClusterBackupMeta has no cluster_id") if not is_lightweight: raise ValueError("ClusterBackupMeta is not marked as lightweight") keyspace_meta_prefix = f"/pd/{backup_cluster_id}/keyspaces/meta/".encode() cluster_ids = set() duplicate_count = 0 keyspace_count = 0 non_enabled_keyspace_count = 0 for key, value in sorted(raw_pd_meta.items()): if not key.startswith(keyspace_meta_prefix): continue path_id = key[len(keyspace_meta_prefix):] if len(path_id) != 8 or not path_id.isdigit(): raise ValueError(f"unexpected keyspace metadata key: {key!r}") keyspace_id, keyspace_state, config = parse_keyspace_meta(value) if keyspace_id != int(path_id): raise ValueError( f"keyspace ID mismatch: path={int(path_id)}, " f"protobuf={keyspace_id}" ) keyspace_count += 1 # cse-ctl intersects --keyspace-select.cluster-id-in with keyspaces # that are eligible for Pack. Mirroring its state filter here keeps # active_cluster_ids from selecting archived or otherwise inactive # clusters that cannot produce packed metadata. if keyspace_state != KEYSPACE_STATE_ENABLED: non_enabled_keyspace_count += 1 continue cluster_id = config.get("serverless_cluster_id") if cluster_id is None: continue if not cluster_id: raise ValueError( f"keyspace {keyspace_id} has an empty serverless_cluster_id" ) if len(cluster_id) > 255: raise ValueError( f"serverless_cluster_id for keyspace {keyspace_id} exceeds " "varchar(255)" ) if cluster_id in cluster_ids: duplicate_count += 1 else: cluster_ids.add(cluster_id) if keyspace_count == 0: raise ValueError( f"no state-bearing keyspace metadata found in {source_label}" ) if not cluster_ids: raise ValueError( f"no serverless_cluster_id values found in {source_label}" ) return ( sorted(cluster_ids), duplicate_count, keyspace_count, non_enabled_keyspace_count, backup_cluster_id, ) def mysql_string(value: str) -> str: replacements = { "\0": "\\0", "\n": "\\n", "\r": "\\r", "\\": "\\\\", "'": "\\'", '"': '\\"', "\x1a": "\\Z", } escaped = "".join( replacements.get(character, character) for character in value ) return f"'{escaped}'" def format_datetime(value: datetime) -> str: return value.strftime("%Y-%m-%d %H:%M:%S.%f") def format_datetime_seconds(value: datetime) -> str: return value.strftime("%Y-%m-%d %H:%M:%S") def build_sql( cluster_ids: List[str], created_at: datetime, s3_data_updated_at: datetime, rng: random.Random, source_label: str, backup_cluster_id: int, keyspace_count: int, non_enabled_keyspace_count: int, ) -> str: created_at_sql = mysql_string(format_datetime(created_at)) s3_data_updated_at_sql = mysql_string( format_datetime_seconds(s3_data_updated_at) ) rows = [] for cluster_id in cluster_ids: offset = rng.randrange( -UPDATED_AT_WINDOW_MICROSECONDS, UPDATED_AT_WINDOW_MICROSECONDS + 1, ) try: updated_at = s3_data_updated_at + timedelta(microseconds=offset) except OverflowError as error: raise ValueError( "s3_data_updated_at is too close to the supported " "datetime boundary" ) from error rows.append( " (" f"{mysql_string(cluster_id)}, " f"{created_at_sql}, " f"{mysql_string(format_datetime(updated_at))}, " f"{s3_data_updated_at_sql}" ")" ) return ( f"-- Generated from {source_label}\n" f"-- Backup cluster_id: {backup_cluster_id}; " f"parsed keyspaces: {keyspace_count}; " f"non-enabled keyspaces skipped: {non_enabled_keyspace_count}.\n" "-- All DATETIME values are UTC.\n" "-- update_count uses its table default.\n" "REPLACE INTO `active_cluster_ids` " "(`cluster_id`, `created_at`, `updated_at`, `s3_data_updated_at`)\n" "VALUES\n" + ",\n".join(rows) + ";\n" )