We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Issue when calculating lineage when a CTE has the same name as a table from the schema
**Fully reproducible code snippet**

The following code 
```sql
from sqlglot.lineage import lineage

selected_column = "col_a"
sql = """
with 

my_cte_name_also_a_table_name as (
    select * from raw.schema.my_table
)

, inter as (
    select * from my_cte_name_also_a_table_name
)

select * from inter
"""

schema = {"raw": {"schema" : {"my_cte_name_also_a_table_name": {"col_1": "int"}, "my_table": {"col_a": "int"}}}}
schema_without_table = {"my_table": {"col_a": "int"}}

l = lineage(column=selected_column, sql=sql, schema=schema, dialect="snowflake")
```

returns the error
```
Traceback (most recent call last):
  File "/xxx/short_issue_lineage.py", line 21, in <module>
    l = lineage(column=selected_column, sql=sql, schema=schema, dialect="snowflake")
  File "/xxx/lib/python3.9/site-packages/sqlglot/lineage.py", line 148, in lineage
    return to_node(column if isinstance(column, str) else column.name, scope)
  File "/xxx/lib/python3.9/site-packages/sqlglot/lineage.py", line 136, in to_node
    to_node(
  File "/xxx/lib/python3.9/site-packages/sqlglot/lineage.py", line 112, in to_node
    source = optimize(
  File "/xxx/lib/python3.9/site-packages/sqlglot/optimizer/optimizer.py", line 89, in optimize
    expression = rule(expression, **rule_kwargs)
  File "/xxx/lib/python3.9/site-packages/sqlglot/optimizer/qualify_columns.py", line 49, in qualify_columns
    _qualify_columns(scope, resolver)
  File "/xxx/lib/python3.9/site-packages/sqlglot/optimizer/qualify_columns.py", line 250, in _qualify_columns
    raise OptimizeError(f"Unknown column: {column_name}")
sqlglot.errors.OptimizeError: Unknown column: col_a
```


It looks like there is an issue with the logic being confused between `my_cte_name_also_a_table_name` the CTE (with a column called `col_a`) and `my_cte_name_also_a_table_name` the table from the schema with a column called `col_1`.

With the example above, if I remove the table from the schema and calculate the lineage, the error goes away.
```
l = lineage(column=selected_column, sql=sql, schema=schema_without_table, dialect="snowflake")
```


Interestingly, the code without the intermediate cte `inter` calculates the lineage correctly when the table with the CTE name is provided in the schema.
```
sql = """
with 

my_cte_name_also_a_table_name as (
    select * from raw.schema.my_table
)

select * from my_cte_name_also_a_table_name
"""
l = lineage(column=selected_column, sql=sql, schema=schema, dialect="snowflake")
```

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `sqlglot/lineage.py`
Content:
```
1 from __future__ import annotations
2 
3 import json
4 import typing as t
5 from dataclasses import dataclass, field
6 
7 from sqlglot import Schema, exp, maybe_parse
8 from sqlglot.optimizer import Scope, build_scope, optimize
9 from sqlglot.optimizer.lower_identities import lower_identities
10 from sqlglot.optimizer.qualify_columns import qualify_columns
11 from sqlglot.optimizer.qualify_tables import qualify_tables
12 
13 if t.TYPE_CHECKING:
14     from sqlglot.dialects.dialect import DialectType
15 
16 
17 @dataclass(frozen=True)
18 class Node:
19     name: str
20     expression: exp.Expression
21     source: exp.Expression
22     downstream: t.List[Node] = field(default_factory=list)
23     alias: str = ""
24 
25     def walk(self) -> t.Iterator[Node]:
26         yield self
27 
28         for d in self.downstream:
29             if isinstance(d, Node):
30                 yield from d.walk()
31             else:
32                 yield d
33 
34     def to_html(self, **opts) -> LineageHTML:
35         return LineageHTML(self, **opts)
36 
37 
38 def lineage(
39     column: str | exp.Column,
40     sql: str | exp.Expression,
41     schema: t.Optional[t.Dict | Schema] = None,
42     sources: t.Optional[t.Dict[str, str | exp.Subqueryable]] = None,
43     rules: t.Sequence[t.Callable] = (lower_identities, qualify_tables, qualify_columns),
44     dialect: DialectType = None,
45 ) -> Node:
46     """Build the lineage graph for a column of a SQL query.
47 
48     Args:
49         column: The column to build the lineage for.
50         sql: The SQL string or expression.
51         schema: The schema of tables.
52         sources: A mapping of queries which will be used to continue building lineage.
53         rules: Optimizer rules to apply, by default only qualifying tables and columns.
54         dialect: The dialect of input SQL.
55 
56     Returns:
57         A lineage node.
58     """
59 
60     expression = maybe_parse(sql, dialect=dialect)
61 
62     if sources:
63         expression = exp.expand(
64             expression,
65             {
66                 k: t.cast(exp.Subqueryable, maybe_parse(v, dialect=dialect))
67                 for k, v in sources.items()
68             },
69         )
70 
71     optimized = optimize(expression, schema=schema, rules=rules)
72     scope = build_scope(optimized)
73 
74     def to_node(
75         column_name: str,
76         scope: Scope,
77         scope_name: t.Optional[str] = None,
78         upstream: t.Optional[Node] = None,
79         alias: t.Optional[str] = None,
80     ) -> Node:
81         aliases = {
82             dt.alias: dt.comments[0].split()[1]
83             for dt in scope.derived_tables
84             if dt.comments and dt.comments[0].startswith("source: ")
85         }
86         if isinstance(scope.expression, exp.Union):
87             for scope in scope.union_scopes:
88                 node = to_node(
89                     column_name,
90                     scope=scope,
91                     scope_name=scope_name,
92                     upstream=upstream,
93                     alias=aliases.get(scope_name),
94                 )
95             return node
96 
97         # Find the specific select clause that is the source of the column we want.
98         # This can either be a specific, named select or a generic `*` clause.
99         select = next(
100             (select for select in scope.selects if select.alias_or_name == column_name),
101             exp.Star() if scope.expression.is_star else None,
102         )
103 
104         if not select:
105             raise ValueError(f"Could not find {column_name} in {scope.expression}")
106 
107         if isinstance(scope.expression, exp.Select):
108             # For better ergonomics in our node labels, replace the full select with
109             # a version that has only the column we care about.
110             #   "x", SELECT x, y FROM foo
111             #     => "x", SELECT x FROM foo
112             source = optimize(
113                 scope.expression.select(select, append=False), schema=schema, rules=rules
114             )
115             select = source.selects[0]
116         else:
117             source = scope.expression
118 
119         # Create the node for this step in the lineage chain, and attach it to the previous one.
120         node = Node(
121             name=f"{scope_name}.{column_name}" if scope_name else column_name,
122             source=source,
123             expression=select,
124             alias=alias or "",
125         )
126         if upstream:
127             upstream.downstream.append(node)
128 
129         # Find all columns that went into creating this one to list their lineage nodes.
130         for c in set(select.find_all(exp.Column)):
131             table = c.table
132             source = scope.sources.get(table)
133 
134             if isinstance(source, Scope):
135                 # The table itself came from a more specific scope. Recurse into that one using the unaliased column name.
136                 to_node(
137                     c.name, scope=source, scope_name=table, upstream=node, alias=aliases.get(table)
138                 )
139             else:
140                 # The source is not a scope - we've reached the end of the line. At this point, if a source is not found
141                 # it means this column's lineage is unknown. This can happen if the definition of a source used in a query
142                 # is not passed into the `sources` map.
143                 source = source or exp.Placeholder()
144                 node.downstream.append(Node(name=c.sql(), source=source, expression=source))
145 
146         return node
147 
148     return to_node(column if isinstance(column, str) else column.name, scope)
149 
150 
151 class LineageHTML:
152     """Node to HTML generator using vis.js.
153 
154     https://visjs.github.io/vis-network/docs/network/
155     """
156 
157     def __init__(
158         self,
159         node: Node,
160         dialect: DialectType = None,
161         imports: bool = True,
162         **opts: t.Any,
163     ):
164         self.node = node
165         self.imports = imports
166 
167         self.options = {
168             "height": "500px",
169             "width": "100%",
170             "layout": {
171                 "hierarchical": {
172                     "enabled": True,
173                     "nodeSpacing": 200,
174                     "sortMethod": "directed",
175                 },
176             },
177             "interaction": {
178                 "dragNodes": False,
179                 "selectable": False,
180             },
181             "physics": {
182                 "enabled": False,
183             },
184             "edges": {
185                 "arrows": "to",
186             },
187             "nodes": {
188                 "font": "20px monaco",
189                 "shape": "box",
190                 "widthConstraint": {
191                     "maximum": 300,
192                 },
193             },
194             **opts,
195         }
196 
197         self.nodes = {}
198         self.edges = []
199 
200         for node in node.walk():
201             if isinstance(node.expression, exp.Table):
202                 label = f"FROM {node.expression.this}"
203                 title = f"<pre>SELECT {node.name} FROM {node.expression.this}</pre>"
204                 group = 1
205             else:
206                 label = node.expression.sql(pretty=True, dialect=dialect)
207                 source = node.source.transform(
208                     lambda n: exp.Tag(this=n, prefix="<b>", postfix="</b>")
209                     if n is node.expression
210                     else n,
211                     copy=False,
212                 ).sql(pretty=True, dialect=dialect)
213                 title = f"<pre>{source}</pre>"
214                 group = 0
215 
216             node_id = id(node)
217 
218             self.nodes[node_id] = {
219                 "id": node_id,
220                 "label": label,
221                 "title": title,
222                 "group": group,
223             }
224 
225             for d in node.downstream:
226                 self.edges.append({"from": node_id, "to": id(d)})
227 
228     def __str__(self):
229         nodes = json.dumps(list(self.nodes.values()))
230         edges = json.dumps(self.edges)
231         options = json.dumps(self.options)
232         imports = (
233             """<script type="text/javascript" src="https://unpkg.com/vis-data@latest/peer/umd/vis-data.min.js"></script>
234   <script type="text/javascript" src="https://unpkg.com/vis-network@latest/peer/umd/vis-network.min.js"></script>
235   <link rel="stylesheet" type="text/css" href="https://unpkg.com/vis-network/styles/vis-network.min.css" />"""
236             if self.imports
237             else ""
238         )
239 
240         return f"""<div>
241   <div id="sqlglot-lineage"></div>
242   {imports}
243   <script type="text/javascript">
244     var nodes = new vis.DataSet({nodes})
245     nodes.forEach(row => row["title"] = new DOMParser().parseFromString(row["title"], "text/html").body.childNodes[0])
246 
247     new vis.Network(
248         document.getElementById("sqlglot-lineage"),
249         {{
250             nodes: nodes,
251             edges: new vis.DataSet({edges})
252         }},
253         {options},
254     )
255   </script>
256 </div>"""
257 
258     def _repr_html_(self) -> str:
259         return self.__str__()
260
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/sqlglot/lineage.py b/sqlglot/lineage.py
--- a/sqlglot/lineage.py
+++ b/sqlglot/lineage.py
@@ -109,10 +109,7 @@
             # a version that has only the column we care about.
             #   "x", SELECT x, y FROM foo
             #     => "x", SELECT x FROM foo
-            source = optimize(
-                scope.expression.select(select, append=False), schema=schema, rules=rules
-            )
-            select = source.selects[0]
+            source = t.cast(exp.Expression, scope.expression.select(select, append=False))
         else:
             source = scope.expression

{"golden_diff": "diff --git a/sqlglot/lineage.py b/sqlglot/lineage.py\n--- a/sqlglot/lineage.py\n+++ b/sqlglot/lineage.py\n@@ -109,10 +109,7 @@\n             # a version that has only the column we care about.\n             #   \"x\", SELECT x, y FROM foo\n             #     => \"x\", SELECT x FROM foo\n-            source = optimize(\n-                scope.expression.select(select, append=False), schema=schema, rules=rules\n-            )\n-            select = source.selects[0]\n+            source = t.cast(exp.Expression, scope.expression.select(select, append=False))\n         else:\n             source = scope.expression\n", "issue": "Issue when calculating lineage when a CTE has the same name as a table from the schema\n**Fully reproducible code snippet**\r\n\r\nThe following code \r\n```sql\r\nfrom sqlglot.lineage import lineage\r\n\r\nselected_column = \"col_a\"\r\nsql = \"\"\"\r\nwith \r\n\r\nmy_cte_name_also_a_table_name as (\r\n    select * from raw.schema.my_table\r\n)\r\n\r\n, inter as (\r\n    select * from my_cte_name_also_a_table_name\r\n)\r\n\r\nselect * from inter\r\n\"\"\"\r\n\r\nschema = {\"raw\": {\"schema\" : {\"my_cte_name_also_a_table_name\": {\"col_1\": \"int\"}, \"my_table\": {\"col_a\": \"int\"}}}}\r\nschema_without_table = {\"my_table\": {\"col_a\": \"int\"}}\r\n\r\nl = lineage(column=selected_column, sql=sql, schema=schema, dialect=\"snowflake\")\r\n```\r\n\r\nreturns the error\r\n```\r\nTraceback (most recent call last):\r\n  File \"/xxx/short_issue_lineage.py\", line 21, in <module>\r\n    l = lineage(column=selected_column, sql=sql, schema=schema, dialect=\"snowflake\")\r\n  File \"/xxx/lib/python3.9/site-packages/sqlglot/lineage.py\", line 148, in lineage\r\n    return to_node(column if isinstance(column, str) else column.name, scope)\r\n  File \"/xxx/lib/python3.9/site-packages/sqlglot/lineage.py\", line 136, in to_node\r\n    to_node(\r\n  File \"/xxx/lib/python3.9/site-packages/sqlglot/lineage.py\", line 112, in to_node\r\n    source = optimize(\r\n  File \"/xxx/lib/python3.9/site-packages/sqlglot/optimizer/optimizer.py\", line 89, in optimize\r\n    expression = rule(expression, **rule_kwargs)\r\n  File \"/xxx/lib/python3.9/site-packages/sqlglot/optimizer/qualify_columns.py\", line 49, in qualify_columns\r\n    _qualify_columns(scope, resolver)\r\n  File \"/xxx/lib/python3.9/site-packages/sqlglot/optimizer/qualify_columns.py\", line 250, in _qualify_columns\r\n    raise OptimizeError(f\"Unknown column: {column_name}\")\r\nsqlglot.errors.OptimizeError: Unknown column: col_a\r\n```\r\n\r\n\r\nIt looks like there is an issue with the logic being confused between `my_cte_name_also_a_table_name` the CTE (with a column called `col_a`) and `my_cte_name_also_a_table_name` the table from the schema with a column called `col_1`.\r\n\r\nWith the example above, if I remove the table from the schema and calculate the lineage, the error goes away.\r\n```\r\nl = lineage(column=selected_column, sql=sql, schema=schema_without_table, dialect=\"snowflake\")\r\n```\r\n\r\n\r\nInterestingly, the code without the intermediate cte `inter` calculates the lineage correctly when the table with the CTE name is provided in the schema.\r\n```\r\nsql = \"\"\"\r\nwith \r\n\r\nmy_cte_name_also_a_table_name as (\r\n    select * from raw.schema.my_table\r\n)\r\n\r\nselect * from my_cte_name_also_a_table_name\r\n\"\"\"\r\nl = lineage(column=selected_column, sql=sql, schema=schema, dialect=\"snowflake\")\r\n```\n", "before_files": [{"content": "from __future__ import annotations\n\nimport json\nimport typing as t\nfrom dataclasses import dataclass, field\n\nfrom sqlglot import Schema, exp, maybe_parse\nfrom sqlglot.optimizer import Scope, build_scope, optimize\nfrom sqlglot.optimizer.lower_identities import lower_identities\nfrom sqlglot.optimizer.qualify_columns import qualify_columns\nfrom sqlglot.optimizer.qualify_tables import qualify_tables\n\nif t.TYPE_CHECKING:\n    from sqlglot.dialects.dialect import DialectType\n\n\n@dataclass(frozen=True)\nclass Node:\n    name: str\n    expression: exp.Expression\n    source: exp.Expression\n    downstream: t.List[Node] = field(default_factory=list)\n    alias: str = \"\"\n\n    def walk(self) -> t.Iterator[Node]:\n        yield self\n\n        for d in self.downstream:\n            if isinstance(d, Node):\n                yield from d.walk()\n            else:\n                yield d\n\n    def to_html(self, **opts) -> LineageHTML:\n        return LineageHTML(self, **opts)\n\n\ndef lineage(\n    column: str | exp.Column,\n    sql: str | exp.Expression,\n    schema: t.Optional[t.Dict | Schema] = None,\n    sources: t.Optional[t.Dict[str, str | exp.Subqueryable]] = None,\n    rules: t.Sequence[t.Callable] = (lower_identities, qualify_tables, qualify_columns),\n    dialect: DialectType = None,\n) -> Node:\n    \"\"\"Build the lineage graph for a column of a SQL query.\n\n    Args:\n        column: The column to build the lineage for.\n        sql: The SQL string or expression.\n        schema: The schema of tables.\n        sources: A mapping of queries which will be used to continue building lineage.\n        rules: Optimizer rules to apply, by default only qualifying tables and columns.\n        dialect: The dialect of input SQL.\n\n    Returns:\n        A lineage node.\n    \"\"\"\n\n    expression = maybe_parse(sql, dialect=dialect)\n\n    if sources:\n        expression = exp.expand(\n            expression,\n            {\n                k: t.cast(exp.Subqueryable, maybe_parse(v, dialect=dialect))\n                for k, v in sources.items()\n            },\n        )\n\n    optimized = optimize(expression, schema=schema, rules=rules)\n    scope = build_scope(optimized)\n\n    def to_node(\n        column_name: str,\n        scope: Scope,\n        scope_name: t.Optional[str] = None,\n        upstream: t.Optional[Node] = None,\n        alias: t.Optional[str] = None,\n    ) -> Node:\n        aliases = {\n            dt.alias: dt.comments[0].split()[1]\n            for dt in scope.derived_tables\n            if dt.comments and dt.comments[0].startswith(\"source: \")\n        }\n        if isinstance(scope.expression, exp.Union):\n            for scope in scope.union_scopes:\n                node = to_node(\n                    column_name,\n                    scope=scope,\n                    scope_name=scope_name,\n                    upstream=upstream,\n                    alias=aliases.get(scope_name),\n                )\n            return node\n\n        # Find the specific select clause that is the source of the column we want.\n        # This can either be a specific, named select or a generic `*` clause.\n        select = next(\n            (select for select in scope.selects if select.alias_or_name == column_name),\n            exp.Star() if scope.expression.is_star else None,\n        )\n\n        if not select:\n            raise ValueError(f\"Could not find {column_name} in {scope.expression}\")\n\n        if isinstance(scope.expression, exp.Select):\n            # For better ergonomics in our node labels, replace the full select with\n            # a version that has only the column we care about.\n            #   \"x\", SELECT x, y FROM foo\n            #     => \"x\", SELECT x FROM foo\n            source = optimize(\n                scope.expression.select(select, append=False), schema=schema, rules=rules\n            )\n            select = source.selects[0]\n        else:\n            source = scope.expression\n\n        # Create the node for this step in the lineage chain, and attach it to the previous one.\n        node = Node(\n            name=f\"{scope_name}.{column_name}\" if scope_name else column_name,\n            source=source,\n            expression=select,\n            alias=alias or \"\",\n        )\n        if upstream:\n            upstream.downstream.append(node)\n\n        # Find all columns that went into creating this one to list their lineage nodes.\n        for c in set(select.find_all(exp.Column)):\n            table = c.table\n            source = scope.sources.get(table)\n\n            if isinstance(source, Scope):\n                # The table itself came from a more specific scope. Recurse into that one using the unaliased column name.\n                to_node(\n                    c.name, scope=source, scope_name=table, upstream=node, alias=aliases.get(table)\n                )\n            else:\n                # The source is not a scope - we've reached the end of the line. At this point, if a source is not found\n                # it means this column's lineage is unknown. This can happen if the definition of a source used in a query\n                # is not passed into the `sources` map.\n                source = source or exp.Placeholder()\n                node.downstream.append(Node(name=c.sql(), source=source, expression=source))\n\n        return node\n\n    return to_node(column if isinstance(column, str) else column.name, scope)\n\n\nclass LineageHTML:\n    \"\"\"Node to HTML generator using vis.js.\n\n    https://visjs.github.io/vis-network/docs/network/\n    \"\"\"\n\n    def __init__(\n        self,\n        node: Node,\n        dialect: DialectType = None,\n        imports: bool = True,\n        **opts: t.Any,\n    ):\n        self.node = node\n        self.imports = imports\n\n        self.options = {\n            \"height\": \"500px\",\n            \"width\": \"100%\",\n            \"layout\": {\n                \"hierarchical\": {\n                    \"enabled\": True,\n                    \"nodeSpacing\": 200,\n                    \"sortMethod\": \"directed\",\n                },\n            },\n            \"interaction\": {\n                \"dragNodes\": False,\n                \"selectable\": False,\n            },\n            \"physics\": {\n                \"enabled\": False,\n            },\n            \"edges\": {\n                \"arrows\": \"to\",\n            },\n            \"nodes\": {\n                \"font\": \"20px monaco\",\n                \"shape\": \"box\",\n                \"widthConstraint\": {\n                    \"maximum\": 300,\n                },\n            },\n            **opts,\n        }\n\n        self.nodes = {}\n        self.edges = []\n\n        for node in node.walk():\n            if isinstance(node.expression, exp.Table):\n                label = f\"FROM {node.expression.this}\"\n                title = f\"<pre>SELECT {node.name} FROM {node.expression.this}</pre>\"\n                group = 1\n            else:\n                label = node.expression.sql(pretty=True, dialect=dialect)\n                source = node.source.transform(\n                    lambda n: exp.Tag(this=n, prefix=\"<b>\", postfix=\"</b>\")\n                    if n is node.expression\n                    else n,\n                    copy=False,\n                ).sql(pretty=True, dialect=dialect)\n                title = f\"<pre>{source}</pre>\"\n                group = 0\n\n            node_id = id(node)\n\n            self.nodes[node_id] = {\n                \"id\": node_id,\n                \"label\": label,\n                \"title\": title,\n                \"group\": group,\n            }\n\n            for d in node.downstream:\n                self.edges.append({\"from\": node_id, \"to\": id(d)})\n\n    def __str__(self):\n        nodes = json.dumps(list(self.nodes.values()))\n        edges = json.dumps(self.edges)\n        options = json.dumps(self.options)\n        imports = (\n            \"\"\"<script type=\"text/javascript\" src=\"https://unpkg.com/vis-data@latest/peer/umd/vis-data.min.js\"></script>\n  <script type=\"text/javascript\" src=\"https://unpkg.com/vis-network@latest/peer/umd/vis-network.min.js\"></script>\n  <link rel=\"stylesheet\" type=\"text/css\" href=\"https://unpkg.com/vis-network/styles/vis-network.min.css\" />\"\"\"\n            if self.imports\n            else \"\"\n        )\n\n        return f\"\"\"<div>\n  <div id=\"sqlglot-lineage\"></div>\n  {imports}\n  <script type=\"text/javascript\">\n    var nodes = new vis.DataSet({nodes})\n    nodes.forEach(row => row[\"title\"] = new DOMParser().parseFromString(row[\"title\"], \"text/html\").body.childNodes[0])\n\n    new vis.Network(\n        document.getElementById(\"sqlglot-lineage\"),\n        {{\n            nodes: nodes,\n            edges: new vis.DataSet({edges})\n        }},\n        {options},\n    )\n  </script>\n</div>\"\"\"\n\n    def _repr_html_(self) -> str:\n        return self.__str__()\n", "path": "sqlglot/lineage.py"}], "after_files": [{"content": "from __future__ import annotations\n\nimport json\nimport typing as t\nfrom dataclasses import dataclass, field\n\nfrom sqlglot import Schema, exp, maybe_parse\nfrom sqlglot.optimizer import Scope, build_scope, optimize\nfrom sqlglot.optimizer.lower_identities import lower_identities\nfrom sqlglot.optimizer.qualify_columns import qualify_columns\nfrom sqlglot.optimizer.qualify_tables import qualify_tables\n\nif t.TYPE_CHECKING:\n    from sqlglot.dialects.dialect import DialectType\n\n\n@dataclass(frozen=True)\nclass Node:\n    name: str\n    expression: exp.Expression\n    source: exp.Expression\n    downstream: t.List[Node] = field(default_factory=list)\n    alias: str = \"\"\n\n    def walk(self) -> t.Iterator[Node]:\n        yield self\n\n        for d in self.downstream:\n            if isinstance(d, Node):\n                yield from d.walk()\n            else:\n                yield d\n\n    def to_html(self, **opts) -> LineageHTML:\n        return LineageHTML(self, **opts)\n\n\ndef lineage(\n    column: str | exp.Column,\n    sql: str | exp.Expression,\n    schema: t.Optional[t.Dict | Schema] = None,\n    sources: t.Optional[t.Dict[str, str | exp.Subqueryable]] = None,\n    rules: t.Sequence[t.Callable] = (lower_identities, qualify_tables, qualify_columns),\n    dialect: DialectType = None,\n) -> Node:\n    \"\"\"Build the lineage graph for a column of a SQL query.\n\n    Args:\n        column: The column to build the lineage for.\n        sql: The SQL string or expression.\n        schema: The schema of tables.\n        sources: A mapping of queries which will be used to continue building lineage.\n        rules: Optimizer rules to apply, by default only qualifying tables and columns.\n        dialect: The dialect of input SQL.\n\n    Returns:\n        A lineage node.\n    \"\"\"\n\n    expression = maybe_parse(sql, dialect=dialect)\n\n    if sources:\n        expression = exp.expand(\n            expression,\n            {\n                k: t.cast(exp.Subqueryable, maybe_parse(v, dialect=dialect))\n                for k, v in sources.items()\n            },\n        )\n\n    optimized = optimize(expression, schema=schema, rules=rules)\n    scope = build_scope(optimized)\n\n    def to_node(\n        column_name: str,\n        scope: Scope,\n        scope_name: t.Optional[str] = None,\n        upstream: t.Optional[Node] = None,\n        alias: t.Optional[str] = None,\n    ) -> Node:\n        aliases = {\n            dt.alias: dt.comments[0].split()[1]\n            for dt in scope.derived_tables\n            if dt.comments and dt.comments[0].startswith(\"source: \")\n        }\n        if isinstance(scope.expression, exp.Union):\n            for scope in scope.union_scopes:\n                node = to_node(\n                    column_name,\n                    scope=scope,\n                    scope_name=scope_name,\n                    upstream=upstream,\n                    alias=aliases.get(scope_name),\n                )\n            return node\n\n        # Find the specific select clause that is the source of the column we want.\n        # This can either be a specific, named select or a generic `*` clause.\n        select = next(\n            (select for select in scope.selects if select.alias_or_name == column_name),\n            exp.Star() if scope.expression.is_star else None,\n        )\n\n        if not select:\n            raise ValueError(f\"Could not find {column_name} in {scope.expression}\")\n\n        if isinstance(scope.expression, exp.Select):\n            # For better ergonomics in our node labels, replace the full select with\n            # a version that has only the column we care about.\n            #   \"x\", SELECT x, y FROM foo\n            #     => \"x\", SELECT x FROM foo\n            source = t.cast(exp.Expression, scope.expression.select(select, append=False))\n        else:\n            source = scope.expression\n\n        # Create the node for this step in the lineage chain, and attach it to the previous one.\n        node = Node(\n            name=f\"{scope_name}.{column_name}\" if scope_name else column_name,\n            source=source,\n            expression=select,\n            alias=alias or \"\",\n        )\n        if upstream:\n            upstream.downstream.append(node)\n\n        # Find all columns that went into creating this one to list their lineage nodes.\n        for c in set(select.find_all(exp.Column)):\n            table = c.table\n            source = scope.sources.get(table)\n\n            if isinstance(source, Scope):\n                # The table itself came from a more specific scope. Recurse into that one using the unaliased column name.\n                to_node(\n                    c.name, scope=source, scope_name=table, upstream=node, alias=aliases.get(table)\n                )\n            else:\n                # The source is not a scope - we've reached the end of the line. At this point, if a source is not found\n                # it means this column's lineage is unknown. This can happen if the definition of a source used in a query\n                # is not passed into the `sources` map.\n                source = source or exp.Placeholder()\n                node.downstream.append(Node(name=c.sql(), source=source, expression=source))\n\n        return node\n\n    return to_node(column if isinstance(column, str) else column.name, scope)\n\n\nclass LineageHTML:\n    \"\"\"Node to HTML generator using vis.js.\n\n    https://visjs.github.io/vis-network/docs/network/\n    \"\"\"\n\n    def __init__(\n        self,\n        node: Node,\n        dialect: DialectType = None,\n        imports: bool = True,\n        **opts: t.Any,\n    ):\n        self.node = node\n        self.imports = imports\n\n        self.options = {\n            \"height\": \"500px\",\n            \"width\": \"100%\",\n            \"layout\": {\n                \"hierarchical\": {\n                    \"enabled\": True,\n                    \"nodeSpacing\": 200,\n                    \"sortMethod\": \"directed\",\n                },\n            },\n            \"interaction\": {\n                \"dragNodes\": False,\n                \"selectable\": False,\n            },\n            \"physics\": {\n                \"enabled\": False,\n            },\n            \"edges\": {\n                \"arrows\": \"to\",\n            },\n            \"nodes\": {\n                \"font\": \"20px monaco\",\n                \"shape\": \"box\",\n                \"widthConstraint\": {\n                    \"maximum\": 300,\n                },\n            },\n            **opts,\n        }\n\n        self.nodes = {}\n        self.edges = []\n\n        for node in node.walk():\n            if isinstance(node.expression, exp.Table):\n                label = f\"FROM {node.expression.this}\"\n                title = f\"<pre>SELECT {node.name} FROM {node.expression.this}</pre>\"\n                group = 1\n            else:\n                label = node.expression.sql(pretty=True, dialect=dialect)\n                source = node.source.transform(\n                    lambda n: exp.Tag(this=n, prefix=\"<b>\", postfix=\"</b>\")\n                    if n is node.expression\n                    else n,\n                    copy=False,\n                ).sql(pretty=True, dialect=dialect)\n                title = f\"<pre>{source}</pre>\"\n                group = 0\n\n            node_id = id(node)\n\n            self.nodes[node_id] = {\n                \"id\": node_id,\n                \"label\": label,\n                \"title\": title,\n                \"group\": group,\n            }\n\n            for d in node.downstream:\n                self.edges.append({\"from\": node_id, \"to\": id(d)})\n\n    def __str__(self):\n        nodes = json.dumps(list(self.nodes.values()))\n        edges = json.dumps(self.edges)\n        options = json.dumps(self.options)\n        imports = (\n            \"\"\"<script type=\"text/javascript\" src=\"https://unpkg.com/vis-data@latest/peer/umd/vis-data.min.js\"></script>\n  <script type=\"text/javascript\" src=\"https://unpkg.com/vis-network@latest/peer/umd/vis-network.min.js\"></script>\n  <link rel=\"stylesheet\" type=\"text/css\" href=\"https://unpkg.com/vis-network/styles/vis-network.min.css\" />\"\"\"\n            if self.imports\n            else \"\"\n        )\n\n        return f\"\"\"<div>\n  <div id=\"sqlglot-lineage\"></div>\n  {imports}\n  <script type=\"text/javascript\">\n    var nodes = new vis.DataSet({nodes})\n    nodes.forEach(row => row[\"title\"] = new DOMParser().parseFromString(row[\"title\"], \"text/html\").body.childNodes[0])\n\n    new vis.Network(\n        document.getElementById(\"sqlglot-lineage\"),\n        {{\n            nodes: nodes,\n            edges: new vis.DataSet({edges})\n        }},\n        {options},\n    )\n  </script>\n</div>\"\"\"\n\n    def _repr_html_(self) -> str:\n        return self.__str__()\n", "path": "sqlglot/lineage.py"}]}