atc-agents/deploy/mask_pii_setup.py

#!/usr/bin/env python3
"""Add realistic PII columns to the source tables and backfill a bounded sample.

Runs inside the atc-agents api container (has psycopg2/pymysql/pymongo + network
to 10.0.21.51). Idempotent: ADD COLUMN IF NOT EXISTS / duplicate-column tolerant.
Backfills the most recent N rows so OpenMetadata auto-classification samples real
PII values. New generator rows also get PII (see updated generator scripts).
"""
import os
import random

import psycopg2
import pymysql
import pymongo

DB_HOST = os.getenv("SRC_DB_HOST", "10.0.21.51")
DB_USER = os.getenv("SRC_DB_USER", "mo")
DB_PASS = os.getenv("SRC_DB_PASSWORD", "Dell2026!")
N = int(os.getenv("PII_BACKFILL_ROWS", "4000"))
# asc: backfill lowest ids (so OpenMetadata's first-rows sample sees PII);
# desc: newest rows; both: do both ends.
ORDER = os.getenv("PII_BACKFILL_ORDER", "both").lower()


def _orders():
    return (["ASC", "DESC"] if ORDER == "both" else [ORDER.upper()])

FIRST = ["Sophie", "Liam", "Emma", "Noah", "Julia", "Lucas", "Mila", "Daan", "Anna", "Sem",
         "Eva", "Finn", "Tess", "Bram", "Lotte", "Max", "Sara", "Thijs", "Nina", "Ruben"]
LAST = ["de Vries", "Jansen", "Bakker", "Visser", "Smit", "Meijer", "Mulder", "Bos",
        "Vos", "Peters", "Hendriks", "Dijkstra", "Kok", "Willems", "Maas", "Koning"]
STREETS = ["Kerkstraat", "Dorpsweg", "Molenpad", "Stationsplein", "Industrieweg", "Lindelaan",
           "Beukenhof", "Havenstraat", "Parkweg", "Schoolstraat"]
CITIES = ["Amsterdam", "Rotterdam", "Utrecht", "Eindhoven", "Den Haag", "Groningen", "Tilburg", "Almere"]
DOMAINS = ["example.com", "mail.nl", "acme-corp.eu", "shopmail.com", "contact.io"]


def full_name():
    return f"{random.choice(FIRST)} {random.choice(LAST)}"


def email(name):
    base = name.lower().replace(" ", ".")
    return f"{base}{random.randint(1, 999)}@{random.choice(DOMAINS)}"


def phone():
    return f"+31 6 {random.randint(10, 99)} {random.randint(100000, 999999)}"


def ip():
    return ".".join(str(random.randint(1, 254)) for _ in range(4))


def iban():
    return f"NL{random.randint(10, 99)}DELL{random.randint(10**9, 10**10 - 1)}"


def national_id():
    return str(random.randint(100000000, 999999999))


def address():
    return f"{random.choice(STREETS)} {random.randint(1, 320)}, {random.randint(1000, 9999)} {random.choice(CITIES)}"


def dob():
    return f"{random.randint(1960, 2002)}-{random.randint(1, 12):02d}-{random.randint(1, 28):02d}"


def do_postgres():
    print("== PostgreSQL sales_orders ==")
    conn = psycopg2.connect(host=DB_HOST, dbname="postgres", user=DB_USER, password=DB_PASS, port=5432, connect_timeout=10)
    conn.autocommit = True
    cur = conn.cursor()
    for col, typ in [("customer_name", "text"), ("customer_email", "text"), ("customer_phone", "text"),
                     ("customer_ip", "text"), ("billing_iban", "text"), ("shipping_address", "text")]:
        cur.execute(f"ALTER TABLE public.sales_orders ADD COLUMN IF NOT EXISTS {col} {typ}")
    print("  columns ensured")
    total = 0
    for direction in _orders():
        cur.execute(f"SELECT order_id FROM public.sales_orders ORDER BY order_id {direction} LIMIT %s", (N,))
        ids = [r[0] for r in cur.fetchall()]
        rows = []
        for oid in ids:
            nm = full_name()
            rows.append((nm, email(nm), phone(), ip(), iban(), address(), oid))
        cur.executemany(
            "UPDATE public.sales_orders SET customer_name=%s, customer_email=%s, customer_phone=%s, "
            "customer_ip=%s, billing_iban=%s, shipping_address=%s WHERE order_id=%s", rows)
        total += len(rows)
    print(f"  backfilled {total} rows with PII ({'+'.join(_orders())})")
    cur.close()
    conn.close()


def do_mysql():
    print("== MySQL employee_events ==")
    conn = pymysql.connect(host=DB_HOST, user=DB_USER, password=DB_PASS, database="hr", port=3306,
                           connect_timeout=10, autocommit=True)
    cur = conn.cursor()
    for col, typ in [("employee_name", "varchar(120)"), ("employee_email", "varchar(160)"),
                     ("employee_phone", "varchar(40)"), ("national_id", "varchar(20)"),
                     ("home_address", "varchar(255)"), ("date_of_birth", "date")]:
        try:
            cur.execute(f"ALTER TABLE employee_events ADD COLUMN {col} {typ}")
        except Exception as e:
            if "Duplicate column" not in str(e):
                raise
    print("  columns ensured")
    total = 0
    for direction in _orders():
        cur.execute(f"SELECT event_id FROM employee_events ORDER BY event_id {direction} LIMIT %s", (N,))
        ids = [r[0] for r in cur.fetchall()]
        rows = []
        for eid in ids:
            nm = full_name()
            rows.append((nm, email(nm), phone(), national_id(), address(), dob(), eid))
        cur.executemany(
            "UPDATE employee_events SET employee_name=%s, employee_email=%s, employee_phone=%s, "
            "national_id=%s, home_address=%s, date_of_birth=%s WHERE event_id=%s", rows)
        total += len(rows)
    print(f"  backfilled {total} rows with PII ({'+'.join(_orders())})")
    cur.close()
    conn.close()


def do_mongo():
    print("== MongoDB events ==")
    uri = os.getenv("SRC_MONGO_URI", f"mongodb://{DB_HOST}:27017/?replicaSet=rs0")
    client = pymongo.MongoClient(uri, serverSelectionTimeoutMS=10000)
    coll = client["supplychain"]["events"]
    docs = list(coll.find({}, {"_id": 1}).sort("_id", -1).limit(N))
    n = 0
    for d in docs:
        nm = full_name()
        coll.update_one({"_id": d["_id"]}, {"$set": {
            "contact_name": nm, "contact_email": email(nm), "contact_phone": phone(),
        }})
        n += 1
    print(f"  updated {n} docs with PII fields")
    client.close()


if __name__ == "__main__":
    for fn in (do_postgres, do_mysql, do_mongo):
        try:
            fn()
        except Exception as e:
            print(f"  ERROR in {fn.__name__}: {e}")
    print("done")