import json, os, urllib.request, urllib.parse, urllib.error

OM = "http://openmetadata-server:8585/api"
H = {"Authorization": "Bearer " + os.environ["OM_TOKEN"]}


def get(path):
    r = urllib.request.Request(OM + path, headers=H)
    with urllib.request.urlopen(r, timeout=30) as resp:
        return json.loads(resp.read().decode())


def trace(fqn):
    enc = urllib.parse.quote(fqn, safe="")
    g = get("/v1/lineage/table/name/" + enc + "?upstreamDepth=3&downstreamDepth=3")
    nodes = {}
    for n in g.get("nodes", []) + ([g["entity"]] if "entity" in g else []):
        nodes[n["id"]] = (n.get("type"), n.get("fullyQualifiedName") or n.get("name"))
    base = g.get("entity", {})
    nodes[base["id"]] = (base.get("type"), base.get("fullyQualifiedName"))

    def nm(i):
        t, f = nodes.get(i, ("?", i))
        short = (f or "").split(".")[-1].strip('"')
        return "%s(%s)" % (short, t)

    print("\n### lineage around:", fqn)
    print("UP (sources feeding it):")
    for e in g.get("upstreamEdges", []):
        print("   %s --> %s" % (nm(e["fromEntity"]), nm(e["toEntity"])))
    print("DOWN (where it flows to):")
    for e in g.get("downstreamEdges", []):
        print("   %s --> %s" % (nm(e["fromEntity"]), nm(e["toEntity"])))


for f in [
    "atc_postgres.postgres.public.sales_orders",
    "atc_mysql.default.hr.employee_events",
]:
    try:
        trace(f)
    except Exception as e:
        print(f, "ERR", str(e)[:120])