Data lineage you can defend, part 2: AI-assisted lineage on AWS

import json, os, uuid, datetime
import boto3
import sqlglot
from sqlglot import exp
from sqlglot.lineage import lineage as sg_lineage

s3 = boto3.client("s3")
bedrock = boto3.client("bedrock-runtime")
datazone = boto3.client("datazone")

DOMAIN_ID = os.environ["DATAZONE_DOMAIN_ID"]
MODEL_ID = os.environ["LINEAGE_MODEL_ID"]   # e.g. a Claude inference profile ARN
CONFIDENCE_FLOOR = float(os.environ.get("CONFIDENCE_FLOOR", "0.85"))
MAX_EVENT_BYTES = 290_000                   # stay under the 300 KB lineage event limit
EVIDENCE_FACET = "myorg_lineageEvidence"    # custom facets are prefixed with your producer name

# OpenLineage transformation subtypes, grouped by type
DIRECT_SUBTYPES = {"IDENTITY", "TRANSFORMATION", "AGGREGATION"}
INDIRECT_SUBTYPES = {"JOIN", "GROUP_BY", "FILTER", "SORT", "WINDOW", "CONDITIONAL"}

LINEAGE_TOOL = {
    "toolSpec": {
        "name": "emit_lineage",
        "description": "Emit column-level lineage for the target table defined by this code.",
        "inputSchema": {"json": {
            "type": "object",
            "required": ["target_table", "columns"],
            "properties": {
                "target_table": {"type": "string"},
                "columns": {"type": "array", "minItems": 1, "items": {
                    "type": "object",
                    "required": ["target_column", "sources", "subtype",
                                 "summary", "confidence", "evidence_lines"],
                    "properties": {
                        "target_column": {"type": "string"},
                        "sources": {"type": "array", "items": {
                            "type": "object",
                            "required": ["table", "column"],
                            "properties": {"table": {"type": "string"},
                                           "column": {"type": "string"}}}},
                        "subtype": {"type": "string",
                                    "enum": sorted(DIRECT_SUBTYPES | INDIRECT_SUBTYPES) + ["UNKNOWN"]},
                        "summary": {"type": "string"},
                        "confidence": {"type": "number", "minimum": 0, "maximum": 1},
                        "evidence_lines": {"type": "array", "items": {"type": "integer"}}
                    }}}
            }}}
    }
}

SYSTEM = [{"text": (
    "You are a data lineage analyst. Derive column-level lineage ONLY from the code "
    "provided. Never invent tables or columns that do not appear in the code. If a "
    "source cannot be determined, use subtype UNKNOWN and confidence below 0.5. "
    "Cite the line numbers that justify each mapping in evidence_lines.")}]

def table_name(t: exp.Table) -> str:
    """Fully qualified table name without its alias, e.g. 'db.schema.orders'."""
    return ".".join(p for p in (t.catalog, t.db, t.name) if p)

def infer_subtype(expression: exp.Expression) -> str:
    """Classify a projected expression into an OpenLineage subtype."""
    e = expression.unalias() if isinstance(expression, exp.Alias) else expression
    if isinstance(e, exp.Column):
        return "IDENTITY"
    if e.find(exp.AggFunc):
        return "AGGREGATION"
    if e.find(exp.Case, exp.If):
        return "CONDITIONAL"
    return "TRANSFORMATION"

def try_sqlglot(sql: str, dialect: str, target_cols: list[str], schema: dict | None):
    """Deterministic path. Returns column mappings, or None if anything is unresolved."""
    if not target_cols:
        return None
    mappings = []
    try:
        for col in target_cols:
            node = sg_lineage(col, sql, schema=schema, dialect=dialect)
            sources = [{"table": table_name(leaf.source), "column": leaf.name.split(".")[-1]}
                       for leaf in node.walk()
                       if not leaf.downstream and isinstance(leaf.source, exp.Table)]
            if not sources:
                return None
            mappings.append({
                "target_column": col, "sources": sources,
                "subtype": infer_subtype(node.expression),
                "summary": "", "confidence": 0.98, "evidence_lines": [],
                # deterministic evidence: the exact expression the parser resolved
                "evidence_expr": node.expression.sql(dialect=dialect)[:500],
            })
        return mappings
    except Exception:
        return None

def ask_bedrock(code: str, hint: str) -> dict:
    numbered = "\n".join(f"{i+1:>5}: {l}" for i, l in enumerate(code.splitlines()))
    resp = bedrock.converse(
        modelId=MODEL_ID,
        system=SYSTEM,
        messages=[{"role": "user", "content": [{"text": f"{hint}\n\n<code>\n{numbered}\n</code>"}]}],
        toolConfig={"tools": [LINEAGE_TOOL],
                    "toolChoice": {"tool": {"name": "emit_lineage"}}},
        inferenceConfig={"temperature": 0},
    )
    for block in resp["output"]["message"]["content"]:
        if "toolUse" in block:
            return block["toolUse"]["input"]
    raise ValueError("Model did not return structured lineage")

def transformation(c: dict) -> dict:
    sub = c["subtype"]
    if sub in DIRECT_SUBTYPES:
        t = {"type": "DIRECT", "subtype": sub}
    elif sub in INDIRECT_SUBTYPES:
        t = {"type": "INDIRECT", "subtype": sub}
    else:  # UNKNOWN: keep the edge, omit the subtype, and force review downstream
        t = {"type": "INDIRECT"}
    if c.get("summary"):
        t["description"] = c["summary"][:500]
    return t

def to_openlineage(columns: list[dict], target_table: str, job_name: str,
                   namespace: str, evidence: dict) -> dict:
    """Build an OpenLineage RunEvent with a columnLineage facet."""
    fields, inputs = {}, {}
    for c in columns:
        fields[c["target_column"]] = {"inputFields": [
            {"namespace": namespace, "name": s["table"], "field": s["column"],
             "transformations": [transformation(c)]} for s in c["sources"]]}
        for s in c["sources"]:
            inputs[s["table"]] = {"namespace": namespace, "name": s["table"]}
    producer = f"https://example.com/lineage-detectors/{evidence['detector']}"
    return {
        "eventType": "COMPLETE",
        "eventTime": datetime.datetime.now(datetime.timezone.utc).isoformat(),
        "producer": producer,
        "schemaURL": "https://openlineage.io/spec/2-0-2/OpenLineage.json#/definitions/RunEvent",
        "run": {"runId": str(uuid.uuid4()),
                "facets": {EVIDENCE_FACET: {"_producer": producer,
                                            "_schemaURL": "https://example.com/schemas/evidence.json",
                                            **evidence}}},
        "job": {"namespace": namespace, "name": job_name},
        "inputs": list(inputs.values()),
        "outputs": [{"namespace": namespace, "name": target_table, "facets": {"columnLineage": {
            "_producer": producer,
            "_schemaURL": "https://openlineage.io/spec/facets/1-2-0/ColumnLineageDatasetFacet.json",
            "fields": fields}}}],
    }

def handler(event, _ctx):
    """Invoked per file by Distributed Map. Returns scored results; publishing happens
    downstream so that low-confidence results can be routed to review first."""
    code = s3.get_object(Bucket=event["bucket"], Key=event["key"])["Body"].read().decode()
    target_table = event.get("target_table")
    mappings = None
    if target_table:  # the deterministic path needs to know which table it is resolving
        mappings = try_sqlglot(code, event.get("dialect", "tsql"),
                               event.get("target_columns", []), event.get("schema"))
    if mappings:
        columns, detector = mappings, "sqlglot"
        evidence_meta = {"parser_version": sqlglot.__version__}
    else:
        result = ask_bedrock(code, f"Dialect hint: {event.get('dialect', 'unknown')}")
        columns, detector = result.get("columns") or [], "bedrock-claude"
        target_table = target_table or result.get("target_table")
        evidence_meta = {"model_id": MODEL_ID}

    if not columns or not target_table:
        return {"events": [], "detector": detector, "min_confidence": 0.0,
                "needs_review": True, "reason": "no lineage extracted", "key": event["key"]}

    min_conf = min(c["confidence"] for c in columns)
    has_unknown = any(c["subtype"] == "UNKNOWN" for c in columns)
    evidence = {"detector": detector, "source_uri": f"s3://{event['bucket']}/{event['key']}",
                "min_confidence": min_conf, **evidence_meta,
                "columns": {c["target_column"]: {"lines": c["evidence_lines"],
                                                 "expr": c.get("evidence_expr", "")}
                            for c in columns}}
    events = split_to_fit(columns, target_table, event["key"], event["namespace"], evidence)
    return {"events": events, "detector": detector, "min_confidence": min_conf,
            "needs_review": detector != "sqlglot" and (min_conf < CONFIDENCE_FLOOR or has_unknown)}

def split_to_fit(columns, target_table, job_name, namespace, evidence):
    """Split wide tables into several events so each stays under the size limit."""
    ev = to_openlineage(columns, target_table, job_name, namespace, evidence)
    if len(json.dumps(ev).encode("utf-8")) <= MAX_EVENT_BYTES or len(columns) == 1:
        return [ev]
    mid = len(columns) // 2
    trim = lambda cols: {**evidence, "columns": {c["target_column"]: evidence["columns"][c["target_column"]]
                                                 for c in cols}}
    return (split_to_fit(columns[:mid], target_table, job_name, namespace, trim(columns[:mid])) +
            split_to_fit(columns[mid:], target_table, job_name, namespace, trim(columns[mid:])))

def publish(ol_event: dict):
    """Called only for auto-approved or human-approved edges."""
    payload = json.dumps(ol_event).encode("utf-8")
    if len(payload) > MAX_EVENT_BYTES:
        raise ValueError(f"Lineage event is {len(payload)} bytes; split it before publishing")
    datazone.post_lineage_event(domainIdentifier=DOMAIN_ID, event=payload,
                                clientToken=ol_event["run"]["runId"])