#!/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 USER_COLUMNS = [ "Id", "Username", "Email", "Password", "FirstName", "LastName", "Phone", "ProfileId", "UserRoleId", "ManagerId", "Title", "Department", "Division", "MiddleName", "Suffix", "Alias", "CommunityNickname", "MobilePhone", "TimeZoneSidKey", "LocaleSidKey", "LanguageLocaleKey", "EmailEncodingKey", "CurrencyIsoCode", "Street", "City", "State", "PostalCode", "Country", "IsActive", "CreatedById", "LastModifiedById", "LastLoginDate", "CreatedDate", "LastModifiedDate", ] def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description="Supplement missing Salesforce users into TiDB without truncating existing local/system users.") parser.add_argument("--snapshot-dir", required=True) parser.add_argument("--overrides-json", required=True) parser.add_argument("--user-ids", help="Optional comma-separated user ids to limit supplement scope.") 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 parse_user_ids(raw: str | None) -> set[str]: if not raw: return set() return {part.strip() for part in raw.split(",") if part.strip()} def load_overrides(path: Path) -> dict[str, Any]: payload = json.loads(path.read_text(encoding="utf-8")) if not isinstance(payload, dict): raise ValueError("overrides json must be an object") return payload def load_source_rows(path: Path) -> dict[str, dict[str, str]]: rows: dict[str, dict[str, str]] = {} with gzip.open(path, "rt", encoding="utf-8-sig", newline="") as src: for row in csv.DictReader(src): user_id = (row.get("Id") or "").strip() if user_id: rows[user_id] = row return rows def normalize_bool(value: str) -> int: return 1 if str(value).strip().lower() == "true" else 0 def normalize_datetime(value: str) -> str | None: raw = str(value or "").strip() if not raw: return None return raw.replace("T", " ").replace(".000+0000", "").replace(".000Z", "") def synthetic_email(user_id: str) -> str: return f"sf-user+{user_id.lower()}@example.invalid" def existing_id_set(cursor: pymysql.cursors.Cursor, table: str) -> set[str]: cursor.execute(f"SELECT Id FROM `{table}`") return {row[0] for row in cursor.fetchall()} def main() -> int: args = parse_args() snapshot_dir = Path(args.snapshot_dir).expanduser().resolve() overrides = load_overrides(Path(args.overrides_json).expanduser().resolve()) password_literal = str(overrides.get("password_literal") or "").strip() if not password_literal: raise ValueError("password_literal is required in overrides json") source_rows = load_source_rows(snapshot_dir / "User.csv.gz") requested_ids = parse_user_ids(args.user_ids) 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] = {"inserted": [], "updated_refs": [], "skipped": []} try: with conn.cursor() as cur: existing_user_ids = existing_id_set(cur, "User") existing_profile_ids = existing_id_set(cur, "Profile") existing_role_ids = existing_id_set(cur, "UserRole") missing_ids = sorted(set(source_rows) - existing_user_ids) if requested_ids: missing_ids = [user_id for user_id in missing_ids if user_id in requested_ids] email_owner: dict[str, str] = {} cur.execute("SELECT Id, Email FROM `User`") for user_id, email in cur.fetchall(): if email: email_owner[email] = user_id planned_rows: dict[str, dict[str, Any]] = {} for user_id in missing_ids: row = source_rows[user_id] user_override = ((overrides.get("users") or {}).get(user_id) or {}) profile_id = (user_override.get("profile_id") or row.get("ProfileId") or "").strip() if profile_id not in existing_profile_ids: report["skipped"].append( {"id": user_id, "reason": "missing_profile", "profile_id": profile_id} ) continue source_email = (row.get("Email") or "").strip() final_email = source_email owner = email_owner.get(final_email) if not final_email or (owner and owner != user_id): final_email = synthetic_email(user_id) email_owner[final_email] = user_id planned_rows[user_id] = { "Id": user_id, "Username": (row.get("Username") or "").strip(), "Email": final_email, "Password": password_literal, "FirstName": (row.get("FirstName") or None), "LastName": (row.get("LastName") or "").strip(), "Phone": (row.get("Phone") or None), "ProfileId": profile_id, "UserRoleId": (row.get("UserRoleId") or "").strip() or None, "ManagerId": (row.get("ManagerId") or "").strip() or None, "Title": (row.get("Title") or None), "Department": (row.get("Department") or None), "Division": (row.get("Division") or None), "MiddleName": (row.get("MiddleName") or None), "Suffix": (row.get("Suffix") or None), "Alias": (row.get("Alias") or None), "CommunityNickname": (row.get("CommunityNickname") or None), "MobilePhone": (row.get("MobilePhone") or None), "TimeZoneSidKey": (row.get("TimeZoneSidKey") or None), "LocaleSidKey": (row.get("LocaleSidKey") or None), "LanguageLocaleKey": (row.get("LanguageLocaleKey") or None), "EmailEncodingKey": (row.get("EmailEncodingKey") or None), "CurrencyIsoCode": (row.get("CurrencyIsoCode") or None), "Street": (row.get("Street") or None), "City": (row.get("City") or None), "State": (row.get("State") or None), "PostalCode": (row.get("PostalCode") or None), "Country": (row.get("Country") or None), "IsActive": normalize_bool(row.get("IsActive") or ""), "CreatedById": (row.get("CreatedById") or "").strip() or None, "LastModifiedById": (row.get("LastModifiedById") or "").strip() or None, "LastLoginDate": normalize_datetime(row.get("LastLoginDate") or ""), "CreatedDate": normalize_datetime(row.get("CreatedDate") or ""), "LastModifiedDate": normalize_datetime(row.get("LastModifiedDate") or ""), } placeholders = ", ".join(["%s"] * len(USER_COLUMNS)) insert_sql = ( f"INSERT INTO `User` ({', '.join(f'`{column}`' for column in USER_COLUMNS)}) " f"VALUES ({placeholders})" ) for user_id, row in planned_rows.items(): insert_row = dict(row) if insert_row["UserRoleId"] not in existing_role_ids: insert_row["UserRoleId"] = None for ref_key in ("ManagerId", "CreatedById", "LastModifiedById"): ref = insert_row[ref_key] if ref not in existing_user_ids: insert_row[ref_key] = None values = [insert_row[column] for column in USER_COLUMNS] cur.execute(insert_sql, values) existing_user_ids.add(user_id) report["inserted"].append( { "id": user_id, "username": insert_row["Username"], "email": insert_row["Email"], "source_email": row["Email"], "profile_id": insert_row["ProfileId"], } ) for user_id, row in planned_rows.items(): update_fields: list[str] = [] update_values: list[Any] = [] for ref_key in ("ManagerId", "CreatedById", "LastModifiedById"): ref = row[ref_key] if ref in existing_user_ids: update_fields.append(f"`{ref_key}` = %s") update_values.append(ref) if row["UserRoleId"] in existing_role_ids: update_fields.append("`UserRoleId` = %s") update_values.append(row["UserRoleId"]) if not update_fields: continue update_values.append(user_id) cur.execute(f"UPDATE `User` SET {', '.join(update_fields)} WHERE Id = %s", update_values) report["updated_refs"].append({"id": user_id, "fields": [field.split('`')[1] for field in update_fields]}) conn.commit() finally: conn.close() print(json.dumps(report, ensure_ascii=False, indent=2)) return 0 if __name__ == "__main__": raise SystemExit(main())