atc-agents/deploy/openmetadata/enrich/om_enrich.py

"""Enrich OpenMetadata directly via REST: real sample data + table/column
profiles + full end-to-end lineage. Bypasses the flaky profiler/connectors.

Run inside the openmetadata_ingestion container (has psycopg2/pymysql/pymongo/
trino and network access to openmetadata-server).
"""
import datetime
import decimal
import json
import os
import time
import urllib.error
import urllib.parse
import urllib.request

OM = "http://openmetadata-server:8585/api"
TOKEN = os.environ["OM_TOKEN"]
SRC_PASS = os.environ.get("SRC_PASS", "Dell2026!")  # source DB password
H = {"Authorization": "Bearer " + TOKEN}
NOW_MS = int(time.time() * 1000)


def req(method, path, body=None):
    data = json.dumps(body).encode() if body is not None else None
    h = dict(H)
    if data is not None:
        h["Content-Type"] = "application/json"
    r = urllib.request.Request(OM + path, data=data, headers=h, method=method)
    try:
        with urllib.request.urlopen(r, timeout=60) as resp:
            return resp.status, json.loads(resp.read().decode() or "{}")
    except urllib.error.HTTPError as e:
        return e.code, e.read().decode()[:300]


def jval(v):
    if v is None or isinstance(v, (int, float, bool)):
        return v
    if isinstance(v, decimal.Decimal):
        return float(v)
    if isinstance(v, (datetime.datetime, datetime.date)):
        return v.isoformat()
    return str(v)


# ---- entity id maps -------------------------------------------------------
def list_map(path, key="fullyQualifiedName"):
    st, d = req("GET", path)
    out = {}
    if isinstance(d, dict):
        for x in d.get("data", []):
            out[x[key]] = x["id"]
    return out


TOPICS = list_map("/v1/topics?service=atc_kafka&limit=200")
CONTAINERS = list_map("/v1/containers?service=atc_s3&limit=200")
PIPELINES = list_map("/v1/pipelines?service=atc_airflow&limit=200")
_table_cache = {}


def table_id(fqn):
    if fqn in _table_cache:
        return _table_cache[fqn]
    st, d = req("GET", "/v1/tables/name/" + urllib.parse.quote(fqn, safe=""))
    tid = d["id"] if st == 200 else None
    _table_cache[fqn] = tid
    if tid is None:
        print("  ! table not found:", fqn)
    return tid


# ---- sample data + profile ------------------------------------------------
def push_sample_and_profile(fqn, columns, rows, row_count_estimate):
    tid = table_id(fqn)
    if not tid:
        return
    cols = list(columns)
    jrows = [[jval(v) for v in r] for r in rows]
    st, _ = req("PUT", "/v1/tables/%s/sampleData" % tid,
                {"columns": cols, "rows": jrows})
    # column profiles from the sample
    n = len(jrows)
    colprof = []
    for i, c in enumerate(cols):
        vals = [r[i] for r in jrows]
        nonnull = [v for v in vals if v is not None]
        colprof.append({
            "timestamp": NOW_MS, "name": c,
            "valuesCount": len(nonnull),
            "nullCount": n - len(nonnull),
            "distinctCount": len(set(map(str, nonnull))),
        })
    body = {
        "tableProfile": {"timestamp": NOW_MS, "rowCount": row_count_estimate,
                         "columnCount": len(cols)},
        "columnProfile": colprof,
    }
    st2, r2 = req("PUT", "/v1/tables/%s/tableProfile" % tid, body)
    print("  sample+profile %-45s rows=%-3d rowCount=%-12s -> %s/%s" % (
        fqn.split(".")[-1], n, row_count_estimate, st, st2))


# ---- lineage --------------------------------------------------------------
def edge(frm, to, pipeline_fqn=None):
    body = {"edge": {"fromEntity": frm, "toEntity": to}}
    if pipeline_fqn and pipeline_fqn in PIPELINES:
        body["edge"]["lineageDetails"] = {
            "pipeline": {"id": PIPELINES[pipeline_fqn], "type": "pipeline"}}
    st, r = req("PUT", "/v1/lineage", body)
    return st in (200, 201), (st, r)


def E_table(fqn):
    tid = table_id(fqn)
    return {"id": tid, "type": "table"} if tid else None


def E_topic(fqn):
    return {"id": TOPICS[fqn], "type": "topic"} if fqn in TOPICS else None


def E_container(fqn):
    return {"id": CONTAINERS[fqn], "type": "container"} if fqn in CONTAINERS else None


def lineage(frm, to, label, pipeline=None):
    if not frm or not to:
        print("  ! skip lineage (missing entity):", label)
        return
    ok, info = edge(frm, to, pipeline)
    print("  lineage %-50s -> %s" % (label, "OK" if ok else "ERR %s" % (info,)))


# ---- pull sample data from the real sources -------------------------------
def pg_sample(table, limit=50):
    import psycopg2
    c = psycopg2.connect(host="10.0.21.51", port=5432, user="mo",
                         password=SRC_PASS, dbname="postgres", connect_timeout=10)
    c.set_session(autocommit=True)
    cur = c.cursor()
    cur.execute("SELECT reltuples::bigint FROM pg_class WHERE oid=%s::regclass", (table,))
    est = cur.fetchone()[0]
    cur.execute("SELECT * FROM %s LIMIT %s" % (table, limit))
    cols = [d[0] for d in cur.description]
    rows = cur.fetchall()
    c.close()
    return cols, rows, int(est)


def my_sample(table, schema="hr", limit=50):
    import pymysql
    c = pymysql.connect(host="10.0.21.51", port=3306, user="mo",
                        password=SRC_PASS, database=schema, connect_timeout=10)
    cur = c.cursor()
    cur.execute("SELECT table_rows FROM information_schema.tables "
                "WHERE table_schema=%s AND table_name=%s", (schema, table))
    est = cur.fetchone()[0]
    cur.execute("SELECT * FROM %s LIMIT %d" % (table, limit))
    cols = [d[0] for d in cur.description]
    rows = cur.fetchall()
    c.close()
    return cols, rows, int(est)


def mongo_sample(db, coll, limit=50):
    import pymongo
    mc = pymongo.MongoClient(
        "mongodb://mo:%s@10.0.21.51:27017/?authSource=admin" % SRC_PASS,
        serverSelectionTimeoutMS=8000)
    d = mc[db]
    est = d[coll].estimated_document_count()
    docs = list(d[coll].find({}, limit=limit))
    cols, seen = [], set()
    for doc in docs:
        for k in doc:
            if k not in seen:
                seen.add(k); cols.append(k)
    rows = [[doc.get(c) for c in cols] for doc in docs]
    mc.close()
    return cols, rows, int(est)


def trino_sample(fq, limit=50, exact_count=True):
    import trino
    conn = trino.dbapi.connect(host="10.0.21.50", port=8089, user="mo", catalog="iceberg")
    cur = conn.cursor()
    cur.execute("SELECT * FROM " + fq + " LIMIT %d" % limit)
    cols = [d[0] for d in cur.description]
    rows = cur.fetchall()
    est = len(rows)
    if exact_count:
        try:
            cur.execute("SELECT count(*) FROM " + fq)
            est = cur.fetchone()[0]
        except Exception:
            pass
    return cols, rows, int(est)


def main():
    print("== SAMPLE DATA + PROFILES ==")
    # source databases
    try:
        c, r, e = pg_sample("public.sales_orders")
        push_sample_and_profile("atc_postgres.postgres.public.sales_orders", c, r, e)
    except Exception as ex:
        print("  pg err", str(ex)[:120])
    try:
        c, r, e = my_sample("employee_events")
        push_sample_and_profile("atc_mysql.default.hr.employee_events", c, r, e)
    except Exception as ex:
        print("  mysql err", str(ex)[:120])
    try:
        c, r, e = mongo_sample("supplychain", "events")
        push_sample_and_profile("atc_mongodb.supplychain.supplychain.events", c, r, e)
    except Exception as ex:
        print("  mongo err", str(ex)[:120])

    # iceberg lakehouse (small, exact counts)
    iceberg = {
        'iceberg."curated_masked"."sales_orders_masked"':
            "atc_trino.iceberg.curated_masked.sales_orders_masked",
        'iceberg."curated_masked"."employee_events_masked"':
            "atc_trino.iceberg.curated_masked.employee_events_masked",
        'iceberg."hadoop"."historical_sales"':
            "atc_trino.iceberg.hadoop.historical_sales",
        'iceberg."hadoop"."historical_sales_hdfs"':
            "atc_trino.iceberg.hadoop.historical_sales_hdfs",
    }
    for fq, omfqn in iceberg.items():
        try:
            c, r, e = trino_sample(fq)
            push_sample_and_profile(omfqn, c, r, e)
        except Exception as ex:
            print("  trino err", omfqn, str(ex)[:120])

    print("== LINEAGE (end-to-end) ==")
    PG = "atc_postgres.postgres.public.sales_orders"
    MY = "atc_mysql.default.hr.employee_events"
    MO = "atc_mongodb.supplychain.supplychain.events"
    S3 = "atc_s3.data"
    # Source -> Debezium/Kafka CDC topic
    lineage(E_table(PG), E_topic('atc_kafka."postgres-sales.public.sales_orders"'),
            "sales_orders -> kafka(postgres-sales)")
    lineage(E_table(MY), E_topic('atc_kafka."mysql-hr.hr.employee_events"'),
            "employee_events -> kafka(mysql-hr)")
    lineage(E_table(MO), E_topic('atc_kafka."mongodb-supplychain.supplychain.events"'),
            "events -> kafka(mongodb-supplychain)")
    # Kafka CDC topic -> S3 archive
    lineage(E_topic('atc_kafka."postgres-sales.public.sales_orders"'), E_container(S3),
            "kafka(postgres-sales) -> s3")
    lineage(E_topic('atc_kafka."mysql-hr.hr.employee_events"'), E_container(S3),
            "kafka(mysql-hr) -> s3")
    lineage(E_topic('atc_kafka."mongodb-supplychain.supplychain.events"'), E_container(S3),
            "kafka(mongodb-supplychain) -> s3")
    # Source -> Trino query layer (read-through catalogs)
    lineage(E_table(PG), E_table("atc_trino.postgres_sales.public.sales_orders"),
            "sales_orders -> trino.postgres_sales")
    lineage(E_table(MY), E_table("atc_trino.mysql_hr.hr.employee_events"),
            "employee_events -> trino.mysql_hr")
    lineage(E_table(MO), E_table("atc_trino.mongodb_supplychain.supplychain.events"),
            "events -> trino.mongodb_supplychain")
    # Hadoop historical transform
    lineage(E_table("atc_trino.iceberg.hadoop.historical_sales"),
            E_table("atc_trino.iceberg.hadoop.historical_sales_hdfs"),
            "historical_sales -> historical_sales_hdfs", pipeline="atc_airflow.hadoop_to_trino")


if __name__ == "__main__":
    main()