Reference Node Name vs Source Name in Lineage

[johnpawleylr]

Hey all,

I've ran into something which I think is a small bug in the lineage code. A query like the following:

WITH dataset AS (
            SELECT *
            FROM catalog.db.table_a

            UNION

            SELECT *
            FROM catalog.db.table_b
        )

        SELECT aliased.x, aliased.created_at FROM dataset AS aliased;

gives a reference node name of dataset when tracing x or created_at. I believe it should be aliased as that is how dataset is aliased within the final select query. This is important to something I am working on, as if the reference_node_name is dataset then it is difficult to follow the chain of nodes in the lineage as the information about aliased is lost.

I've raised a change to fix this, the alteration is very minor reference_node_name=scope_name, rather than reference_node_name=reference_node_name on lineage.py:184. But it is causing an existing test to fail, test_lineage.py:427. This is because "dataset" is being assigned as the reference_node_name as well as the source_name.

So what i want to know, is what is the intention behind these two labels? Am I free to change the existing test? Or is reference_node_name being empty in this case a hard requirement?

[georgesittas]

@vchan seems like reference_node_name was introduced here. Any chance you remember the context behind that & how it interacts with SQLMesh? Perhaps that will help answer this question.

@johnpawleylr I need to understand the situation a bit more to decide whether this is something that should change. Can you provide a concrete, minimal example that demonstrates the issue with the current logic? Thanks.

[johnpawleylr]

@georgesittas, this is the test i was planning to add in my PR

    def test_lineage_aliased_cte_union_reference_node_name_is_correct(self) -> None:
        query = """
        WITH dataset AS (
            SELECT *
            FROM catalog.db.table_a

            UNION

            SELECT *
            FROM catalog.db.table_b
        )

        SELECT aliased.x, aliased.created_at FROM dataset AS aliased;
        """
        node = lineage("x", query)

        self.assertEqual(node.name, "x")

        self.assertEqual(node.downstream[0].reference_node_name, "aliased")
        self.assertEqual(node.downstream[1].reference_node_name, "aliased")

without the changes i have made to fix the test, the reference_node_name is equal to "dataset".

at least in my understanding of reference_node_name, "aliased" is what is being used to reference the CTE node named dataset, and therefore it should be the value of that attribute. This seems consistent with how subqueries are handled, e.g.

        query = """
        SELECT aliased.x, aliased.created_at FROM (
            SELECT *
            FROM catalog.db.table_a
        ) AS aliased;
        """
        node = lineage("x", query)

        self.assertEqual(node.name, "x")

        self.assertEqual(node.downstream[0].reference_node_name, "aliased")

although, for regular, non-Union CTEs the reference_node_name does currently equal to the cte name rather than the alias, so the following passes:

        query = """
        with dataset as (
            SELECT *
            FROM catalog.db.table_a
        )
        SELECT aliased.x, aliased.created_at FROM dataset AS aliased;
        """
        node = lineage("x", query)

        self.assertEqual(node.name, "x")

        self.assertEqual(node.downstream[0].reference_node_name, "dataset")

Though, in this case the name attribute of the node is aiased.x, so the existence of the alias is not lost in the node representation. In the union, the name is the index of the column in the results table, i.e. 0 for the first example, so the alias context is lost.

[georgesittas]

I'm not sure if reference_node_name is supposed to contain the CTE's alias, but rather the node's name. I think that the fact that it returns the alias for the subquery (derived table), is due to that being its name.

Regarding:

In the union, the name is the index of the column in the results table, i.e. 0 for the first example, so the alias context is lost.

I can see aliased in the node, see the expression attribute:

>>> from sqlglot.lineage import lineage
...
... query = """
... WITH dataset AS (
...     SELECT *
...     FROM catalog.db.table_a
...
...     UNION
...
...     SELECT *
...     FROM catalog.db.table_b
... )
...
... SELECT aliased.x, aliased.created_at FROM dataset AS aliased;
... """
...
... node = lineage("x", query)
...
>>> node
Node(name='x', expression=Alias(
  this=Column(
    this=Identifier(this=x, quoted=False),
    table=Identifier(this=aliased, quoted=False)),
  alias=Identifier(this=x, quoted=False)), source=Select(
  expressions=[
    Alias(
      this=Column(
        this=Identifier(this=x, quoted=False),
        table=Identifier(this=aliased, quoted=False)),
      alias=Identifier(this=x, quoted=False))],
  from=From(
    this=Table(
      this=Identifier(this=dataset, quoted=False),
      alias=TableAlias(
        this=Identifier(this=aliased, quoted=False)))),
...

I don't completely understand the implications of switching to alias_or_name for reference_node_name, so I'm reluctant to change this and risk at the moment. If this is something you can take care of at the application site, I'd encourage you to do that instead.

[johnpawleylr]

@georgesittas ah, so that's the top level Node in the structure returned by lineage(...). What I am looking at is the downstream nodes, these reference the two expressions which are combined in the Union statement. But you're right, the information on the alias is in that top-level expression.