Merge lineage functionality into heading.py

- Move LineageTable class and lineage computation functions to heading.py
- Simplify LineageTable to use direct SQL instead of inheriting from Table
- Break circular import between heading and lineage modules
- Keep lineage.py as thin re-export for backward compatibility

Author: claude

datajoint/heading.py

@@ -17,6 +17,380 @@
 
 logger = logging.getLogger(__name__.split(".")[0])
 
+
+# =============================================================================
+# Lineage tracking for semantic matching in joins
+# =============================================================================
+
+
+def _parse_full_table_name(full_name):
+    """
+    Parse a full table name like `schema`.`table` into (schema, table).
+
+    :param full_name: full table name in format `schema`.`table`
+    :return: tuple (schema, table)
+    """
+    match = re.match(r"`(\w+)`\.`(\w+)`", full_name)
+    if not match:
+        raise DataJointError(f"Invalid table name format: {full_name}")
+    return match.group(1), match.group(2)
+
+
+def _get_primary_key_attrs(connection, schema, table_name):
+    """
+    Get the primary key attributes for a table by querying the database.
+
+    :param connection: database connection
+    :param schema: schema name
+    :param table_name: table name
+    :return: set of primary key attribute names
+    """
+    result = connection.query(
+        """
+        SELECT column_name
+        FROM information_schema.key_column_usage
+        WHERE table_schema = %s
+          AND table_name = %s
+          AND constraint_name = 'PRIMARY'
+        """,
+        args=(schema, table_name),
+    )
+    return {row[0] for row in result}
+
+
+def _compute_lineage_from_dependencies(connection, schema, table_name, attribute_name):
+    """
+    Compute lineage by traversing the FK graph.
+
+    Uses connection.dependencies which loads FK info from INFORMATION_SCHEMA.
+    This is the fallback when the ~lineage table doesn't exist.
+
+    :param connection: database connection
+    :param schema: schema name
+    :param table_name: table name
+    :param attribute_name: attribute name
+    :return: lineage string "schema.table.attribute" or None for native secondary attrs
+    """
+    connection.dependencies.load(force=False)
+
+    full_table_name = f"`{schema}`.`{table_name}`"
+
+    # Check if the table exists in the dependency graph
+    if full_table_name not in connection.dependencies:
+        # Table not in graph - compute lineage based on primary key status
+        pk_attrs = _get_primary_key_attrs(connection, schema, table_name)
+        if attribute_name in pk_attrs:
+            return f"{schema}.{table_name}.{attribute_name}"
+        else:
+            return None
+
+    # Check incoming edges (foreign keys TO this table's parents)
+    for parent, props in connection.dependencies.parents(full_table_name).items():
+        attr_map = props.get("attr_map", {})
+        if attribute_name in attr_map:
+            # This attribute is inherited from parent - recurse to find origin
+            parent_attr = attr_map[attribute_name]
+            # Handle alias nodes (numeric string nodes in the graph)
+            if parent.isdigit():
+                # Find the actual parent by traversing through alias
+                for grandparent, gprops in connection.dependencies.parents(
+                    parent
+                ).items():
+                    if not grandparent.isdigit():
+                        parent = grandparent
+                        parent_attr = gprops.get("attr_map", {}).get(
+                            attribute_name, parent_attr
+                        )
+                        break
+            parent_schema, parent_table = _parse_full_table_name(parent)
+            return _compute_lineage_from_dependencies(
+                connection, parent_schema, parent_table, parent_attr
+            )
+
+    # Not inherited - check if it's a primary key attribute
+    node_data = connection.dependencies.nodes.get(full_table_name, {})
+    pk_attrs = node_data.get("primary_key", set())
+
+    if attribute_name in pk_attrs:
+        # Native primary key attribute - has lineage to itself
+        return f"{schema}.{table_name}.{attribute_name}"
+    else:
+        # Native secondary attribute - no lineage
+        return None
+
+
+def _compute_all_lineage_for_table(connection, schema, table_name):
+    """
+    Compute lineage for all attributes in a table.
+
+    :param connection: database connection
+    :param schema: schema name
+    :param table_name: table name
+    :return: dict mapping attribute_name -> lineage (or None)
+    """
+    # Get all attributes using Heading
+    heading = Heading(
+        table_info=dict(
+            conn=connection,
+            database=schema,
+            table_name=table_name,
+            context=None,
+        )
+    )
+
+    # Compute lineage for each attribute
+    return {
+        attr: _compute_lineage_from_dependencies(connection, schema, table_name, attr)
+        for attr in heading.names
+    }
+
+
+def _get_lineage_for_heading(connection, schema, table_name):
+    """
+    Get lineage information for all attributes in a table.
+
+    First tries to load from ~lineage table, falls back to FK graph computation.
+
+    :param connection: database connection
+    :param schema: schema name
+    :param table_name: table name
+    :return: dict mapping attribute_name -> lineage (or None)
+    """
+    # Check if ~lineage table exists
+    lineage_table_exists = (
+        connection.query(
+            """
+        SELECT COUNT(*) FROM information_schema.tables
+        WHERE table_schema = %s AND table_name = '~lineage'
+        """,
+            args=(schema,),
+        ).fetchone()[0]
+        > 0
+    )
+
+    if lineage_table_exists:
+        # Load from ~lineage table
+        lineage_table = LineageTable(connection, schema)
+        return lineage_table.get_table_lineage(table_name)
+    else:
+        # Compute from FK graph
+        return _compute_all_lineage_for_table(connection, schema, table_name)
+
+
+class LineageTable:
+    """
+    Hidden table for storing attribute lineage information.
+
+    Each row maps (table_name, attribute_name) -> lineage string.
+    Only attributes with lineage are stored; absence means no lineage.
+    """
+
+    definition = """
+    # Attribute lineage tracking for semantic matching
+    table_name      : varchar(64)   # name of the table
+    attribute_name  : varchar(64)   # name of the attribute
+    ---
+    lineage         : varchar(200)  # "schema.table.attribute"
+    """
+
+    def __init__(self, connection, database):
+        # Lazy import to avoid circular dependency
+        from .table import Table
+
+        self._table_class = Table
+        self.database = database
+        self._connection = connection
+        self._heading = Heading(
+            table_info=dict(
+                conn=connection,
+                database=database,
+                table_name=self.table_name,
+                context=None,
+            )
+        )
+        self._support = [self.full_table_name]
+
+        if not self.is_declared:
+            self._declare()
+
+    @property
+    def table_name(self):
+        return "~lineage"
+
+    @property
+    def full_table_name(self):
+        return f"`{self.database}`.`{self.table_name}`"
+
+    @property
+    def is_declared(self):
+        return (
+            self._connection.query(
+                """
+            SELECT COUNT(*) FROM information_schema.tables
+            WHERE table_schema = %s AND table_name = %s
+            """,
+                args=(self.database, self.table_name),
+            ).fetchone()[0]
+            > 0
+        )
+
+    def _declare(self):
+        """Create the ~lineage table."""
+        self._connection.query(
+            f"""
+            CREATE TABLE IF NOT EXISTS {self.full_table_name} (
+                table_name VARCHAR(64) NOT NULL,
+                attribute_name VARCHAR(64) NOT NULL,
+                lineage VARCHAR(200) NOT NULL,
+                PRIMARY KEY (table_name, attribute_name)
+            ) ENGINE=InnoDB
+            """
+        )
+
+    def insert1(self, row, replace=False):
+        """Insert a single row."""
+        if replace:
+            self._connection.query(
+                f"""
+                REPLACE INTO {self.full_table_name}
+                (table_name, attribute_name, lineage)
+                VALUES (%s, %s, %s)
+                """,
+                args=(row["table_name"], row["attribute_name"], row["lineage"]),
+            )
+        else:
+            self._connection.query(
+                f"""
+                INSERT INTO {self.full_table_name}
+                (table_name, attribute_name, lineage)
+                VALUES (%s, %s, %s)
+                """,
+                args=(row["table_name"], row["attribute_name"], row["lineage"]),
+            )
+
+    def delete_quick(self, table_name=None, attribute_name=None):
+        """Delete rows without prompts."""
+        if table_name and attribute_name:
+            self._connection.query(
+                f"DELETE FROM {self.full_table_name} WHERE table_name=%s AND attribute_name=%s",
+                args=(table_name, attribute_name),
+            )
+        elif table_name:
+            self._connection.query(
+                f"DELETE FROM {self.full_table_name} WHERE table_name=%s",
+                args=(table_name,),
+            )
+        else:
+            self._connection.query(f"DELETE FROM {self.full_table_name}")
+
+    def store_lineage(self, table_name, attribute_name, lineage):
+        """
+        Store lineage for an attribute. Only stores if lineage is not None.
+
+        :param table_name: name of the table (without schema)
+        :param attribute_name: name of the attribute
+        :param lineage: lineage string "schema.table.attribute" or None
+        """
+        if lineage is None:
+            # No lineage - delete any existing entry
+            self.delete_quick(table_name, attribute_name)
+        else:
+            self.insert1(
+                dict(
+                    table_name=table_name,
+                    attribute_name=attribute_name,
+                    lineage=lineage,
+                ),
+                replace=True,
+            )
+
+    def get_lineage(self, table_name, attribute_name):
+        """
+        Get lineage for an attribute.
+
+        :param table_name: name of the table (without schema)
+        :param attribute_name: name of the attribute
+        :return: lineage string or None if no lineage
+        """
+        result = self._connection.query(
+            f"SELECT lineage FROM {self.full_table_name} WHERE table_name=%s AND attribute_name=%s",
+            args=(table_name, attribute_name),
+        ).fetchone()
+        return result[0] if result else None
+
+    def get_table_lineage(self, table_name):
+        """
+        Get lineage for all attributes in a table.
+
+        :param table_name: name of the table (without schema)
+        :return: dict mapping attribute_name -> lineage (only attributes with lineage)
+        """
+        result = self._connection.query(
+            f"SELECT attribute_name, lineage FROM {self.full_table_name} WHERE table_name=%s",
+            args=(table_name,),
+        ).fetchall()
+        return {row[0]: row[1] for row in result}
+
+    def delete_table_lineage(self, table_name):
+        """
+        Delete all lineage records for a table.
+
+        :param table_name: name of the table (without schema)
+        """
+        self.delete_quick(table_name)
+
+
+def compute_schema_lineage(connection, schema):
+    """
+    Compute and populate the ~lineage table for a schema.
+
+    Analyzes foreign key relationships to determine attribute origins.
+
+    :param connection: database connection
+    :param schema: schema object or schema name
+    """
+    from .schemas import Schema
+
+    if isinstance(schema, Schema):
+        schema_name = schema.database
+    else:
+        schema_name = schema
+
+    # Create or get the lineage table
+    lineage_table = LineageTable(connection, schema_name)
+
+    # Get all user tables in the schema (excluding hidden tables)
+    result = connection.query(
+        """
+        SELECT table_name
+        FROM information_schema.tables
+        WHERE table_schema = %s
+          AND table_name NOT LIKE '~%%'
+          AND table_type = 'BASE TABLE'
+        """,
+        args=(schema_name,),
+    )
+    tables = [row[0] for row in result]
+
+    # Ensure dependencies are loaded
+    connection.dependencies.load(force=True)
+
+    # Compute and store lineage for each table
+    for table_name in tables:
+        lineage_map = _compute_all_lineage_for_table(
+            connection, schema_name, table_name
+        )
+        for attr_name, lineage in lineage_map.items():
+            if lineage is not None:
+                lineage_table.store_lineage(table_name, attr_name, lineage)
+
+    logger.info(f"Computed lineage for schema `{schema_name}`: {len(tables)} tables")
+
+
+# =============================================================================
+# End of lineage tracking
+# =============================================================================
+
 default_attribute_properties = dict(  # these default values are set in computed attributes
     name=None,
     type="expression",
@@ -433,9 +807,7 @@ def _init_from_database(self):
 
         # Load lineage information for semantic matching
         try:
-            from .lineage import get_lineage_for_heading
-
-            lineage_map = get_lineage_for_heading(conn, database, table_name, None)
+            lineage_map = _get_lineage_for_heading(conn, database, table_name)
             for attr in attributes:
                 attr["lineage"] = lineage_map.get(attr["name"])
         except Exception as e: