Revert "Merge lineage functionality into heading.py"

This reverts commit 99bc1f5.

Author: claude

datajoint/heading.py

@@ -17,380 +17,6 @@
 
 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",
@@ -807,7 +433,9 @@ def _init_from_database(self):
 
         # Load lineage information for semantic matching
         try:
-            lineage_map = _get_lineage_for_heading(conn, database, table_name)
+            from .lineage import get_lineage_for_heading
+
+            lineage_map = get_lineage_for_heading(conn, database, table_name, None)
             for attr in attributes:
                 attr["lineage"] = lineage_map.get(attr["name"])
         except Exception as e: