#!/usr/bin/env python3 from __future__ import annotations import argparse import csv import gzip import json from pathlib import Path from typing import Any import pymysql GROUP_COLUMNS = [ "Id", "Name", "DeveloperName", "Label", "Type", "Email", "CreatedDate", "LastModifiedDate", ] GROUP_MEMBER_COLUMNS = [ "Id", "GroupId", "UserOrGroupId", "UserOrGroupId_type", "CreatedDate", "LastModifiedDate", ] def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description="Supplement missing Salesforce Group and GroupMember rows into TiDB.") parser.add_argument("--snapshot-dir", required=True) parser.add_argument("--db-host", required=True) parser.add_argument("--db-port", type=int, default=4000) parser.add_argument("--db-user", required=True) parser.add_argument("--db-password", required=True) parser.add_argument("--db-database", required=True) return parser.parse_args() def normalize_datetime(raw: str) -> str | None: value = str(raw or "").strip() if not value: return None return value.replace("T", " ").replace(".000+0000", "").replace(".000Z", "") def normalize_nullable(raw: str) -> str | None: value = str(raw or "").strip() return value or None def ref_type(ref_id: str) -> str | None: if ref_id.startswith("005"): return "User" if ref_id.startswith("00G"): return "Group" return None def load_group_rows(snapshot_dir: Path) -> dict[str, dict[str, Any]]: rows: dict[str, dict[str, Any]] = {} with gzip.open(snapshot_dir / "Group.csv.gz", "rt", encoding="utf-8-sig", newline="") as src: for row in csv.DictReader(src): row_id = (row.get("Id") or "").strip() if not row_id: continue rows[row_id] = { "Id": row_id, "Name": normalize_nullable(row.get("Name") or ""), "DeveloperName": normalize_nullable(row.get("DeveloperName") or ""), "Type": (row.get("Type") or "").strip() or "Regular", "Email": normalize_nullable(row.get("Email") or ""), "CreatedDate": normalize_datetime(row.get("CreatedDate") or ""), "LastModifiedDate": normalize_datetime(row.get("LastModifiedDate") or ""), } return rows def load_group_member_rows(snapshot_dir: Path) -> dict[str, dict[str, Any]]: rows: dict[str, dict[str, Any]] = {} with gzip.open(snapshot_dir / "GroupMember.csv.gz", "rt", encoding="utf-8-sig", newline="") as src: for row in csv.DictReader(src): row_id = (row.get("Id") or "").strip() if not row_id: continue user_or_group_id = (row.get("UserOrGroupId") or "").strip() rows[row_id] = { "Id": row_id, "GroupId": (row.get("GroupId") or "").strip(), "UserOrGroupId": user_or_group_id, "UserOrGroupId_type": ref_type(user_or_group_id), "CreatedDate": normalize_datetime(row.get("SystemModstamp") or ""), "LastModifiedDate": normalize_datetime(row.get("SystemModstamp") or ""), } return rows def synthesize_group_name(row: dict[str, Any], used_names: set[str]) -> str: source_name = (row.get("Name") or "").strip() if source_name: candidate = source_name else: developer_name = (row.get("DeveloperName") or "").strip() or row["Id"] candidate = f"{developer_name}__{row['Type']}" if candidate not in used_names: return candidate fallback = f"{candidate}__{row['Id']}" if fallback not in used_names: return fallback suffix = 2 while True: retry = f"{fallback}__{suffix}" if retry not in used_names: return retry suffix += 1 def main() -> int: args = parse_args() snapshot_dir = Path(args.snapshot_dir).expanduser().resolve() source_groups = load_group_rows(snapshot_dir) source_group_members = load_group_member_rows(snapshot_dir) conn = pymysql.connect( host=args.db_host, port=args.db_port, user=args.db_user, password=args.db_password, database=args.db_database, ssl={"ssl": {}}, charset="utf8mb4", autocommit=False, ) report: dict[str, Any] = { "groups_inserted": 0, "group_members_inserted": 0, "group_member_skipped": [], } try: with conn.cursor() as cur: cur.execute("SELECT Id, Name FROM `Group`") existing_groups = {row[0] for row in cur.fetchall()} cur.execute("SELECT Name FROM `Group`") used_names = {row[0] for row in cur.fetchall() if row[0]} group_insert_sql = ( f"INSERT INTO `Group` ({', '.join(f'`{column}`' for column in GROUP_COLUMNS)}) " f"VALUES ({', '.join(['%s'] * len(GROUP_COLUMNS))})" ) for row_id in sorted(source_groups): if row_id in existing_groups: continue row = dict(source_groups[row_id]) final_name = synthesize_group_name(row, used_names) used_names.add(final_name) values = [ row["Id"], final_name, row["DeveloperName"], final_name, # target Label cannot be blank; mirror resolved name row["Type"], row["Email"], row["CreatedDate"], row["LastModifiedDate"], ] cur.execute(group_insert_sql, values) existing_groups.add(row_id) report["groups_inserted"] += 1 cur.execute("SELECT Id FROM `User`") existing_users = {row[0] for row in cur.fetchall()} cur.execute("SELECT Id FROM `GroupMember`") existing_group_members = {row[0] for row in cur.fetchall()} group_member_insert_sql = ( f"INSERT INTO `GroupMember` ({', '.join(f'`{column}`' for column in GROUP_MEMBER_COLUMNS)}) " f"VALUES ({', '.join(['%s'] * len(GROUP_MEMBER_COLUMNS))})" ) for row_id in sorted(source_group_members): if row_id in existing_group_members: continue row = source_group_members[row_id] if row["GroupId"] not in existing_groups: report["group_member_skipped"].append({"id": row_id, "reason": "missing_group", "group_id": row["GroupId"]}) continue ref_id = row["UserOrGroupId"] if row["UserOrGroupId_type"] == "Group" and ref_id not in existing_groups: report["group_member_skipped"].append({"id": row_id, "reason": "missing_ref_group", "ref_id": ref_id}) continue if row["UserOrGroupId_type"] == "User" and ref_id not in existing_users: report["group_member_skipped"].append({"id": row_id, "reason": "missing_ref_user", "ref_id": ref_id}) continue values = [row[column] for column in GROUP_MEMBER_COLUMNS] cur.execute(group_member_insert_sql, values) existing_group_members.add(row_id) report["group_members_inserted"] += 1 conn.commit() finally: conn.close() print(json.dumps(report, ensure_ascii=False, indent=2)) return 0 if __name__ == "__main__": raise SystemExit(main())