lateral.py

import mysql.connector
from mysql.connector import errorcode
import os

# Database configuration (adjust according to your local environment)
config = {
    'user': 'root',
    'password': os.getenv('MYSQL_PASSWORD'),
    'host': '127.0.0.1',
    'database': 'test',
    'raise_on_warnings': False
}

def setup_test_data():
    try:
        # 1. Establish connection
        cnx = mysql.connector.connect(**config)
        cursor = cnx.cursor()

        # 2. Create test tables
        # nodes: node table (e.g. users)
        # edges: edge table (e.g. relationships)
        print("Creating tables...")
        cursor.execute("DROP TABLE IF EXISTS edges;")
        cursor.execute("DROP TABLE IF EXISTS nodes;")
        
        cursor.execute("""
            CREATE TABLE nodes (
                id INT PRIMARY KEY,
                name VARCHAR(50)
            );
        """)
        
        cursor.execute("""
            CREATE TABLE edges (
                from_id INT,
                to_id INT,
                weight INT,
                PRIMARY KEY (from_id, to_id),
                FOREIGN KEY (from_id) REFERENCES nodes(id)
            );
        """)

        # 3. Insert test data
        print("Inserting test data...")
        # Insert parent nodes
        nodes_data = [(1, 'Node_1'), (10, 'Node_2'), (11, 'Node_3'),(100, 'Node_4')]
        cursor.executemany("INSERT INTO nodes (id, name) VALUES (%s, %s)", nodes_data)

        # Insert multiple edges per node to test LIMIT behavior
        # Node 1 has 5 edges, Node 2 has 2 edges, Node 3 has 0 edges
        edges_data = [
            # Layer 1: 从 1 出发
            (1, 10, 100), (1, 11, 95), (1, 12, 90), (1, 13, 85),
            # Layer 2: 从 10, 11 出发 (每个节点有很多分支)
            (10, 100, 50), (10, 101, 40), (10, 102, 30),
            (11, 110, 55), (11, 111, 45), (11, 112, 35),
            # Layer 3: 从 100 出发
            (100, 1000, 10), (100, 1001, 5)
        ]
        cursor.executemany("INSERT INTO edges (from_id, to_id, weight) VALUES (%s, %s, %s)", edges_data)
        
        cnx.commit()

        # 4. Execute LATERAL + LIMIT query
        # Scenario: fetch top 2 outgoing edges by weight for each node
        print("\nVerifying LATERAL + LIMIT query results (Top 2 per node):")
        query = """
WITH RECURSIVE walk AS (
    -- Anchor: start from root node 1
    SELECT 
        from_id, 
        to_id, 
        weight, 
        1 AS depth
    FROM edges
    WHERE from_id = 1
    
    UNION ALL
    
    -- Recursive part: use LATERAL + LIMIT to restrict fan-out per level
    SELECT 
        t.from_id, 
        t.to_id, 
        t.weight, 
        w.depth + 1
    FROM walk w,
    LATERAL (
        SELECT from_id, to_id, weight
        FROM edges
        WHERE from_id = w.to_id  -- w.to_id is the destination from the previous level
        ORDER BY weight DESC
        LIMIT 2                  -- Key point: only follow 2 branches per node per level
    ) AS t
    WHERE w.depth < 3            -- Limit maximum depth to 3
)
SELECT * FROM walk;
        """
        
        cursor.execute(query)
        
        for (from_id, to_id, weight, depth) in cursor:
            print(f"from_id: {from_id} -> to_id: {to_id}, weight: {weight}, depth: {depth}")

    except mysql.connector.Error as err:
        if err.errno == errorcode.ER_ACCESS_DENIED_ERROR:
            print("Invalid username or password")
        elif err.errno == errorcode.ER_BAD_DB_ERROR:
            print("Database does not exist")
        else:
            print(err)
    finally:
        if 'cnx' in locals() and cnx.is_connected():
            cursor.close()
            cnx.close()

if __name__ == "__main__":
    setup_test_data()