reata__sqllineage-224

{
  "status": "error",
  "error_type": "CommandError",
  "error_message": "Traceback (most recent call last):\n  File \"<stdin>\", line 5, in <module>\nAttributeError: 'LineageRunner' object has no attribute 'describe'\n",
  "details": {
    "stdout": "",
    "stderr": "Traceback (most recent call last):\n  File \"<stdin>\", line 5, in <module>\nAttributeError: 'LineageRunner' object has no attribute 'describe'\n",
    "exit_code": 1
  }
}

{
  "status": "ok",
  "stdout": "",
  "stderr": "",
  "exit_code": 0
}

{
  "status": "ok",
  "filepath": "sqllineage/cli.py",
  "content": "import argparse\nimport logging\n\nfrom sqllineage import DEFAULT_HOST, DEFAULT_PORT\nfrom sqllineage.drawing import draw_lineage_graph\nfrom sqllineage.runner import LineageRunner\nfrom sqllineage.utils.constant import LineageLevel\nfrom sqllineage.utils.helpers import extract_sql_from_args\n\nlogger = logging.getLogger(__name__)\n\n\ndef main(args=None) -> None:\n    \"\"\"\n    The command line interface entry point.\n\n    :param args: the command line arguments for sqllineage command\n    \"\"\"\n    parser = argparse.ArgumentParser(\n        prog=\"sqllineage\", description=\"SQL Lineage Parser.\"\n    )\n    parser.add_argument(\n        \"-e\", metavar=\"<quoted-query-string>\", help=\"SQL from command line\"\n    )\n    parser.add_argument(\"-f\", metavar=\"<filename>\", help=\"SQL from files\")\n    parser.add_argument(\n        \"-v\",\n        \"--verbose\",\n        help=\"increase output verbosity, show statement level lineage result\",\n        action=\"store_true\",\n    )\n    parser.add_argument(\n        \"-l\",\n        \"--level\",\n        help=\"lineage level, column or table, default at table level\",\n        choices=[LineageLevel.TABLE, LineageLevel.COLUMN],\n        default=LineageLevel.TABLE,\n    )\n    parser.add_argument(\n        \"-g\",\n        \"--graph-visualization\",\n        help=\"show graph visualization of the lineage in a webserver\",\n        action=\"store_true\",\n    )\n    parser.add_argument(\n        \"-H\",\n        \"--host\",\n        help=\"the host visualization webserver will be bind to\",\n        type=str,\n        default=DEFAULT_HOST,\n        metavar=\"<hostname>\",\n    )\n    parser.add_argument(\n        \"-p\",\n        \"--port\",\n        help=\"the port visualization webserver will be listening on\",\n        type=int,\n        default=DEFAULT_PORT,\n        metavar=\"<port_number>{0..65536}\",\n    )\n    args = parser.parse_args(args)\n    if args.e and args.f:\n        logging.warning(\n            \"Both -e and -f options are specified. -e option will be ignored\"\n        )\n    if args.f or args.e:\n        sql = extract_sql_from_args(args)\n        runner = LineageRunner(\n            sql,\n            verbose=args.verbose,\n            draw_options={\n                \"host\": args.host,\n                \"port\": args.port,\n                \"f\": args.f if args.f else None,\n            },\n        )\n        if args.graph_visualization:\n            runner.draw()\n        elif args.level == LineageLevel.COLUMN:\n            runner.print_column_lineage()\n        else:\n            runner.print_table_lineage()\n    elif args.graph_visualization:\n        return draw_lineage_graph(**{\"host\": args.host, \"port\": args.port})\n    else:\n        parser.print_help()\n\n\nif __name__ == \"__main__\":\n    main()",
  "start_line": 1,
  "end_line": 90,
  "total_lines": 90,
  "is_truncated": false
}

{
  "status": "ok",
  "filepath": "sqllineage/runner.py",
  "content": "import logging\nfrom typing import Dict, List, Tuple\n\nimport sqlparse\nfrom sqlparse.sql import Statement\n\nfrom sqllineage.core import LineageAnalyzer\nfrom sqllineage.core.holders import SQLLineageHolder\nfrom sqllineage.core.models import Column, Table\nfrom sqllineage.drawing import draw_lineage_graph\nfrom sqllineage.io import to_cytoscape\nfrom sqllineage.utils.constant import LineageLevel\n\n\nlogger = logging.getLogger(__name__)\n\n\ndef lazy_method(func):\n    def wrapper(*args, **kwargs):\n        self = args[0]\n        if not self._evaluated:\n            self._eval()\n        return func(*args, **kwargs)\n\n    return wrapper\n\n\ndef lazy_property(func):\n    return property(lazy_method(func))\n\n\nclass LineageRunner(object):\n    def __init__(\n        self,\n        sql: str,\n        encoding: str = None,\n        verbose: bool = False,\n        draw_options: Dict[str, str] = None,\n    ):\n        \"\"\"\n        The entry point of SQLLineage after command line options are parsed.\n\n        :param sql: a string representation of SQL statements.\n        :param encoding: the encoding for sql string\n        :param verbose: verbose flag indicate whether statement-wise lineage result will be shown\n        \"\"\"\n        self._encoding = encoding\n        self._sql = sql\n        self._verbose = verbose\n        self._draw_options = draw_options if draw_options else {}\n        self._evaluated = False\n        self._stmt: List[Statement] = []\n\n    @lazy_method\n    def __str__(self):\n        \"\"\"\n        print out the Lineage Summary.\n        \"\"\"\n        statements = self.statements(strip_comments=True)\n        source_tables = \"\\n    \".join(str(t) for t in self.source_tables)\n        target_tables = \"\\n    \".join(str(t) for t in self.target_tables)\n        combined = f\"\"\"Statements(#): {len(statements)}\nSource Tables:\n    {source_tables}\nTarget Tables:\n    {target_tables}\n\"\"\"\n        if self.intermediate_tables:\n            intermediate_tables = \"\\n    \".join(\n                str(t) for t in self.intermediate_tables\n            )\n            combined += f\"\"\"Intermediate Tables:\n    {intermediate_tables}\"\"\"\n        if self._verbose:\n            result = \"\"\n            for i, holder in enumerate(self._stmt_holders):\n                stmt_short = statements[i].replace(\"\\n\", \"\")\n                if len(stmt_short) > 50:\n                    stmt_short = stmt_short[:50] + \"...\"\n                content = str(holder).replace(\"\\n\", \"\\n    \")\n                result += f\"\"\"Statement #{i + 1}: {stmt_short}\n    {content}\n\"\"\"\n            combined = result + \"==========\\nSummary:\\n\" + combined\n        return combined\n\n    @lazy_method\n    def to_cytoscape(self, level=LineageLevel.TABLE) -> List[Dict[str, Dict[str, str]]]:\n        \"\"\"\n        to turn the DAG into cytoscape format.\n        \"\"\"\n        if level == LineageLevel.COLUMN:\n            return to_cytoscape(self._sql_holder.column_lineage_graph, compound=True)\n        else:\n            return to_cytoscape(self._sql_holder.table_lineage_graph)\n\n    def draw(self) -> None:\n        \"\"\"\n        to draw the lineage directed graph\n        \"\"\"\n        draw_options = self._draw_options\n        if draw_options.get(\"f\") is None:\n            draw_options.pop(\"f\", None)\n            draw_options[\"e\"] = self._sql\n        return draw_lineage_graph(**draw_options)\n\n    @lazy_method\n    def statements(self, **kwargs) -> List[str]:\n        \"\"\"\n        a list of statements.\n\n        :param kwargs: the key arguments that will be passed to `sqlparse.format`\n        \"\"\"\n        return [sqlparse.format(s.value, **kwargs) for s in self.statements_parsed]\n\n    @lazy_property\n    def statements_parsed(self) -> List[Statement]:\n        \"\"\"\n        a list of :class:`sqlparse.sql.Statement`\n        \"\"\"\n        return self._stmt\n\n    @lazy_property\n    def source_tables(self) -> List[Table]:\n        \"\"\"\n        a list of source :class:`sqllineage.models.Table`\n        \"\"\"\n        return sorted(self._sql_holder.source_tables, key=lambda x: str(x))\n\n    @lazy_property\n    def target_tables(self) -> List[Table]:\n        \"\"\"\n        a list of target :class:`sqllineage.models.Table`\n        \"\"\"\n        return sorted(self._sql_holder.target_tables, key=lambda x: str(x))\n\n    @lazy_property\n    def intermediate_tables(self) -> List[Table]:\n        \"\"\"\n        a list of intermediate :class:`sqllineage.models.Table`\n        \"\"\"\n        return sorted(self._sql_holder.intermediate_tables, key=lambda x: str(x))\n\n    @lazy_method\n    def get_column_lineage(self, exclude_subquery=True) -> List[Tuple[Column, Column]]:\n        \"\"\"\n        a list of column tuple :class:`sqllineage.models.Column`\n        \"\"\"\n        # sort by target column, and then source column\n        return sorted(",
  "start_line": 1,
  "end_line": 150,
  "total_lines": 181,
  "is_truncated": true
}

{
  "status": "ok",
  "stdout": "Statements(#): 1\nSource Tables:\n    <default>.foo\n    <default>.foo.col\nTarget Tables:\n    <default>.foo\n    <default>.foo.col\n\n",
  "stderr": "",
  "exit_code": 0
}

{
  "status": "ok",
  "filepath": "sqllineage/core/holders.py",
  "content": "import itertools\nfrom typing import Set, Tuple, Union\n\nimport networkx as nx\nfrom networkx import DiGraph\n\nfrom sqllineage.core.models import Column, Path, SubQuery, Table\nfrom sqllineage.utils.constant import EdgeType, NodeTag\n\nDATASET_CLASSES = (Path, Table)\n\n\nclass ColumnLineageMixin:\n    def get_column_lineage(self, exclude_subquery=True) -> Set[Tuple[Column, ...]]:\n        self.graph: DiGraph  # For mypy attribute checking\n        # filter all the column node in the graph\n        column_nodes = [n for n in self.graph.nodes if isinstance(n, Column)]\n        column_graph = self.graph.subgraph(column_nodes)\n        source_columns = {column for column, deg in column_graph.in_degree if deg == 0}\n        # if a column lineage path ends at SubQuery, then it should be pruned\n        target_columns = {\n            node\n            for node, deg in column_graph.out_degree\n            if isinstance(node, Column) and deg == 0\n        }\n        if exclude_subquery:\n            target_columns = {\n                node for node in target_columns if isinstance(node.parent, Table)\n            }\n        columns = set()\n        for (source, target) in itertools.product(source_columns, target_columns):\n            simple_paths = list(nx.all_simple_paths(self.graph, source, target))\n            if len(simple_paths) == 1:\n                columns.add(tuple(simple_paths[0]))\n            # we can ignore when simple path doesn't exist, but could there be more than one simple path?\n        return columns\n\n\nclass SubQueryLineageHolder(ColumnLineageMixin):\n    \"\"\"\n    SubQuery/Query Level Lineage Result.\n\n    SubQueryLineageHolder will hold attributes like read, write, cte\n\n    Each of them is a Set[:class:`sqllineage.models.Table`].\n\n    This is the most atomic representation of lineage result.\n    \"\"\"\n\n    def __init__(self) -> None:\n        self.graph = nx.DiGraph()\n\n    def __or__(self, other):\n        self.graph = nx.compose(self.graph, other.graph)\n        return self\n\n    def _property_getter(self, prop) -> Set[Union[SubQuery, Table]]:\n        return {t for t, attr in self.graph.nodes(data=True) if attr.get(prop) is True}\n\n    def _property_setter(self, value, prop) -> None:\n        self.graph.add_node(value, **{prop: True})\n\n    @property\n    def read(self) -> Set[Union[SubQuery, Table]]:\n        return self._property_getter(NodeTag.READ)\n\n    def add_read(self, value) -> None:\n        self._property_setter(value, NodeTag.READ)\n        # the same table can be added (in SQL: joined) multiple times with different alias\n        if hasattr(value, \"alias\"):\n            self.graph.add_edge(value, value.alias, type=EdgeType.HAS_ALIAS)\n\n    @property\n    def write(self) -> Set[Union[SubQuery, Table]]:\n        return self._property_getter(NodeTag.WRITE)\n\n    def add_write(self, value) -> None:\n        self._property_setter(value, NodeTag.WRITE)\n\n    @property\n    def cte(self) -> Set[SubQuery]:\n        return self._property_getter(NodeTag.CTE)  # type: ignore\n\n    def add_cte(self, value) -> None:\n        self._property_setter(value, NodeTag.CTE)\n\n    def add_column_lineage(self, src: Column, tgt: Column) -> None:\n        self.graph.add_edge(src, tgt, type=EdgeType.LINEAGE)\n        self.graph.add_edge(tgt.parent, tgt, type=EdgeType.HAS_COLUMN)\n        if src.parent is not None:\n            # starting NetworkX v2.6, None is not allowed as node, see https://github.com/networkx/networkx/pull/4892\n            self.graph.add_edge(src.parent, src, type=EdgeType.HAS_COLUMN)\n\n\nclass StatementLineageHolder(SubQueryLineageHolder, ColumnLineageMixin):\n    \"\"\"\n    Statement Level Lineage Result.\n\n    Based on SubQueryLineageHolder, StatementLineageHolder holds extra attributes like drop and rename\n\n    For drop, it is a Set[:class:`sqllineage.models.Table`].\n\n    For rename, it a Set[Tuple[:class:`sqllineage.models.Table`, :class:`sqllineage.models.Table`]], with the first\n    table being original table before renaming and the latter after renaming.\n    \"\"\"\n\n    def __str__(self):\n        return \"\\n\".join(\n            f\"table {attr}: {sorted(getattr(self, attr), key=lambda x: str(x)) if getattr(self, attr) else '[]'}\"\n            for attr in [\"read\", \"write\", \"cte\", \"drop\", \"rename\"]\n        )\n\n    def __repr__(self):\n        return str(self)\n\n    @property\n    def read(self) -> Set[Table]:  # type: ignore\n        return {t for t in super().read if isinstance(t, DATASET_CLASSES)}\n\n    @property\n    def write(self) -> Set[Table]:  # type: ignore\n        return {t for t in super().write if isinstance(t, DATASET_CLASSES)}\n\n    @property\n    def drop(self) -> Set[Table]:\n        return self._property_getter(NodeTag.DROP)  # type: ignore\n\n    def add_drop(self, value) -> None:\n        self._property_setter(value, NodeTag.DROP)\n\n    @property\n    def rename(self) -> Set[Tuple[Table, Table]]:\n        return {\n            (src, tgt)\n            for src, tgt, attr in self.graph.edges(data=True)\n            if attr.get(\"type\") == EdgeType.RENAME\n        }\n\n    def add_rename(self, src: Table, tgt: Table) -> None:\n        self.graph.add_edge(src, tgt, type=EdgeType.RENAME)\n\n    @staticmethod\n    def of(holder: SubQueryLineageHolder):\n        stmt_holder = StatementLineageHolder()\n        stmt_holder.graph = holder.graph\n        return stmt_holder\n\n\nclass SQLLineageHolder(ColumnLineageMixin):\n    def __init__(self, graph: DiGraph):",
  "start_line": 1,
  "end_line": 150,
  "total_lines": 280,
  "is_truncated": true
}

{
  "status": "ok",
  "filepath": "sqllineage/core/holders.py",
  "content": "        \"\"\"\n        The combined lineage result in representation of Directed Acyclic Graph.\n\n        :param graph: the Directed Acyclic Graph holding all the combined lineage result.\n        \"\"\"\n        self.graph = graph\n        self._selfloop_tables = self.__retrieve_tag_tables(NodeTag.SELFLOOP)\n        self._sourceonly_tables = self.__retrieve_tag_tables(NodeTag.SOURCE_ONLY)\n        self._targetonly_tables = self.__retrieve_tag_tables(NodeTag.TARGET_ONLY)\n\n    @property\n    def table_lineage_graph(self) -> DiGraph:\n        \"\"\"\n        The table level DiGraph held by SQLLineageHolder\n        \"\"\"\n        table_nodes = [n for n in self.graph.nodes if isinstance(n, DATASET_CLASSES)]\n        return self.graph.subgraph(table_nodes)\n\n    @property\n    def column_lineage_graph(self) -> DiGraph:\n        \"\"\"\n        The column level DiGraph held by SQLLineageHolder\n        \"\"\"\n        column_nodes = [n for n in self.graph.nodes if isinstance(n, Column)]\n        return self.graph.subgraph(column_nodes)\n\n    @property\n    def source_tables(self) -> Set[Table]:\n        \"\"\"\n        a list of source :class:`sqllineage.models.Table`\n        \"\"\"\n        source_tables = {\n            table for table, deg in self.table_lineage_graph.in_degree if deg == 0\n        }.intersection(\n            {table for table, deg in self.table_lineage_graph.out_degree if deg > 0}\n        )\n        source_tables |= self._selfloop_tables\n        source_tables |= self._sourceonly_tables\n        return source_tables\n\n    @property\n    def target_tables(self) -> Set[Table]:\n        \"\"\"\n        a list of target :class:`sqllineage.models.Table`\n        \"\"\"\n        target_tables = {\n            table for table, deg in self.table_lineage_graph.out_degree if deg == 0\n        }.intersection(\n            {table for table, deg in self.table_lineage_graph.in_degree if deg > 0}\n        )\n        target_tables |= self._selfloop_tables\n        target_tables |= self._targetonly_tables\n        return target_tables\n\n    @property\n    def intermediate_tables(self) -> Set[Table]:\n        \"\"\"\n        a list of intermediate :class:`sqllineage.models.Table`\n        \"\"\"\n        intermediate_tables = {\n            table for table, deg in self.table_lineage_graph.in_degree if deg > 0\n        }.intersection(\n            {table for table, deg in self.table_lineage_graph.out_degree if deg > 0}\n        )\n        intermediate_tables -= self.__retrieve_tag_tables(NodeTag.SELFLOOP)\n        return intermediate_tables\n\n    def __retrieve_tag_tables(self, tag) -> Set[Table]:\n        return {\n            table\n            for table, attr in self.graph.nodes(data=True)\n            if attr.get(tag) is True\n        }\n\n    @staticmethod\n    def of(*args: StatementLineageHolder):\n        \"\"\"\n        To assemble multiple :class:`sqllineage.holders.StatementLineageHolder` into\n        :class:`sqllineage.holders.SQLLineageHolder`\n        \"\"\"\n        g = DiGraph()\n        for holder in args:\n            g = nx.compose(g, holder.graph)\n            if holder.drop:\n                for table in holder.drop:\n                    if g.has_node(table) and g.degree[table] == 0:\n                        g.remove_node(table)\n            elif holder.rename:\n                for (table_old, table_new) in holder.rename:\n                    g = nx.relabel_nodes(g, {table_old: table_new})\n                    g.remove_edge(table_new, table_new)\n                    if g.degree[table_new] == 0:\n                        g.remove_node(table_new)\n            else:\n                read, write = holder.read, holder.write\n                if len(read) > 0 and len(write) == 0:\n                    # source only table comes from SELECT statement\n                    g.add_nodes_from(read, **{NodeTag.SOURCE_ONLY: True})\n                elif len(read) == 0 and len(write) > 0:\n                    # target only table comes from case like: 1) INSERT/UPDATE constant values; 2) CREATE TABLE\n                    g.add_nodes_from(write, **{NodeTag.TARGET_ONLY: True})\n                else:\n                    g.add_nodes_from(read)\n                    g.add_nodes_from(write)\n                    for source, target in itertools.product(read, write):\n                        g.add_edge(source, target, type=EdgeType.LINEAGE)\n        for table in {e[0] for e in nx.selfloop_edges(g)}:\n            g.nodes[table][NodeTag.SELFLOOP] = True\n        # find all the columns that we can't assign accurately to a parent table (with multiple parent candidates)\n        unresolved_cols = [\n            (s, t)\n            for s, t in g.edges\n            if isinstance(s, Column) and len(s.parent_candidates) > 1\n        ]\n        for unresolved_col, tgt_col in unresolved_cols:\n            # check if there's only one parent candidate contains the column with same name\n            src_cols = []\n            for parent in unresolved_col.parent_candidates:\n                src_col = Column(unresolved_col.raw_name)\n                src_col.parent = parent\n                if g.has_edge(parent, src_col):\n                    src_cols.append(src_col)\n            if len(src_cols) == 1:\n                g.add_edge(src_cols[0], tgt_col, type=EdgeType.LINEAGE)\n                g.remove_edge(unresolved_col, tgt_col)\n        # when unresolved column got resolved, it will be orphan node, and we can remove it\n        for node in [n for n, deg in g.degree if deg == 0]:\n            if isinstance(node, Column) and len(node.parent_candidates) > 1:\n                g.remove_node(node)\n        return SQLLineageHolder(g)",
  "start_line": 151,
  "end_line": 280,
  "total_lines": 280,
  "is_truncated": false
}

{
  "status": "ok",
  "filepath": "sqllineage/core/models.py",
  "content": "import warnings\nfrom typing import Dict, List, Optional, Set, Union\n\nfrom sqlparse import tokens as T\nfrom sqlparse.engine import grouping\nfrom sqlparse.sql import (\n    Case,\n    Comparison,\n    Function,\n    Identifier,\n    IdentifierList,\n    Operation,\n    Parenthesis,\n    Token,\n    TokenList,\n)\nfrom sqlparse.utils import imt\n\nfrom sqllineage.exceptions import SQLLineageException\nfrom sqllineage.utils.entities import ColumnQualifierTuple\nfrom sqllineage.utils.helpers import escape_identifier_name\nfrom sqllineage.utils.sqlparse import get_parameters, is_subquery\n\n\nclass Schema:\n    unknown = \"<default>\"\n\n    def __init__(self, name: str = unknown):\n        \"\"\"\n        Data Class for Schema\n\n        :param name: schema name\n        \"\"\"\n        self.raw_name = escape_identifier_name(name)\n\n    def __str__(self):\n        return self.raw_name.lower()\n\n    def __repr__(self):\n        return \"Schema: \" + str(self)\n\n    def __eq__(self, other):\n        return type(self) is type(other) and str(self) == str(other)\n\n    def __hash__(self):\n        return hash(str(self))\n\n    def __bool__(self):\n        return str(self) != self.unknown\n\n\nclass Table:\n    def __init__(self, name: str, schema: Schema = Schema(), **kwargs):\n        \"\"\"\n        Data Class for Table\n\n        :param name: table name\n        :param schema: schema as defined by :class:`Schema`\n        \"\"\"\n        if \".\" not in name:\n            self.schema = schema\n            self.raw_name = escape_identifier_name(name)\n        else:\n            schema_name, table_name = name.rsplit(\".\", 1)\n            if len(schema_name.split(\".\")) > 2:\n                # allow db.schema as schema_name, but a.b.c as schema_name is forbidden\n                raise SQLLineageException(\"Invalid format for table name: %s.\", name)\n            self.schema = Schema(schema_name)\n            self.raw_name = escape_identifier_name(table_name)\n            if schema:\n                warnings.warn(\"Name is in schema.table format, schema param is ignored\")\n        self.alias = kwargs.pop(\"alias\", self.raw_name)\n\n    def __str__(self):\n        return f\"{self.schema}.{self.raw_name.lower()}\"\n\n    def __repr__(self):\n        return \"Table: \" + str(self)\n\n    def __eq__(self, other):\n        return type(self) is type(other) and str(self) == str(other)\n\n    def __hash__(self):\n        return hash(str(self))\n\n    @staticmethod\n    def of(identifier: Identifier) -> \"Table\":\n        # rewrite identifier's get_real_name method, by matching the last dot instead of the first dot, so that the\n        # real name for a.b.c will be c instead of b\n        dot_idx, _ = identifier._token_matching(\n            lambda token: imt(token, m=(T.Punctuation, \".\")),\n            start=len(identifier.tokens),\n            reverse=True,\n        )\n        real_name = identifier._get_first_name(dot_idx, real_name=True)\n        # rewrite identifier's get_parent_name accordingly\n        parent_name = (\n            \"\".join(\n                [\n                    escape_identifier_name(token.value)\n                    for token in identifier.tokens[:dot_idx]\n                ]\n            )\n            if dot_idx\n            else None\n        )\n        schema = Schema(parent_name) if parent_name is not None else Schema()\n        alias = identifier.get_alias()\n        kwargs = {\"alias\": alias} if alias else {}\n        return Table(real_name, schema, **kwargs)\n\n\nclass Path:\n    def __init__(self, uri: str):\n        self.uri = escape_identifier_name(uri)\n\n    def __str__(self):\n        return self.uri\n\n    def __repr__(self):\n        return \"Path: \" + str(self)\n\n    def __eq__(self, other):\n        return type(self) is type(other) and self.uri == other.uri\n\n    def __hash__(self):\n        return hash(self.uri)\n\n\nclass Partition:\n    pass\n\n\nclass SubQuery:\n    def __init__(self, token: Parenthesis, alias: Optional[str]):\n        \"\"\"\n        Data Class for SubQuery\n\n        :param token: subquery token\n        :param alias: subquery name\n        \"\"\"\n        self.token = token\n        self._query = token.value\n        self.alias = alias if alias is not None else f\"subquery_{hash(self)}\"\n\n    def __str__(self):\n        return self.alias\n\n    def __repr__(self):\n        return \"SubQuery: \" + str(self)",
  "start_line": 1,
  "end_line": 150,
  "total_lines": 348,
  "is_truncated": true
}

{
  "status": "ok",
  "filepath": "sqllineage/core/models.py",
  "content": "\n    def __eq__(self, other):\n        return type(self) is type(other) and self._query == other._query\n\n    def __hash__(self):\n        return hash(self._query)\n\n    @staticmethod\n    def of(parenthesis: Parenthesis, alias: Optional[str]) -> \"SubQuery\":\n        return SubQuery(parenthesis, alias)\n\n\nclass Column:\n    def __init__(self, name: str, **kwargs):\n        \"\"\"\n        Data Class for Column\n\n        :param name: column name\n        :param parent: :class:`Table` or :class:`SubQuery`\n        :param kwargs:\n        \"\"\"\n        self._parent: Set[Union[Table, SubQuery]] = set()\n        self.raw_name = escape_identifier_name(name)\n        self.source_columns = kwargs.pop(\"source_columns\", ((self.raw_name, None),))\n\n    def __str__(self):\n        return (\n            f\"{self.parent}.{self.raw_name.lower()}\"\n            if self.parent is not None and not isinstance(self.parent, Path)\n            else f\"{self.raw_name.lower()}\"\n        )\n\n    def __repr__(self):\n        return \"Column: \" + str(self)\n\n    def __eq__(self, other):\n        return type(self) is type(other) and str(self) == str(other)\n\n    def __hash__(self):\n        return hash(str(self))\n\n    @property\n    def parent(self) -> Optional[Union[Table, SubQuery]]:\n        return list(self._parent)[0] if len(self._parent) == 1 else None\n\n    @parent.setter\n    def parent(self, value: Union[Table, SubQuery]):\n        self._parent.add(value)\n\n    @property\n    def parent_candidates(self) -> List[Union[Table, SubQuery]]:\n        return sorted(self._parent, key=lambda p: str(p))\n\n    @staticmethod\n    def of(token: Token):\n        if isinstance(token, Identifier):\n            alias = token.get_alias()\n            if alias:\n                # handle column alias, including alias for column name or Case, Function\n                kw_idx, kw = token.token_next_by(m=(T.Keyword, \"AS\"))\n                if kw_idx is None:\n                    # alias without AS\n                    kw_idx, _ = token.token_next_by(i=Identifier)\n                if kw_idx is None:\n                    # invalid syntax: col AS, without alias\n                    return Column(alias)\n                else:\n                    idx, _ = token.token_prev(kw_idx, skip_cm=True)\n                    expr = grouping.group(TokenList(token.tokens[: idx + 1]))[0]\n                    source_columns = Column._extract_source_columns(expr)\n                    return Column(alias, source_columns=source_columns)\n            else:\n                # select column name directly without alias\n                return Column(\n                    token.get_real_name(),\n                    source_columns=((token.get_real_name(), token.get_parent_name()),),\n                )\n        else:\n            # Wildcard, Case, Function without alias (thus not recognized as an Identifier)\n            source_columns = Column._extract_source_columns(token)\n            return Column(token.value, source_columns=source_columns)\n\n    @staticmethod\n    def _extract_source_columns(token: Token) -> List[ColumnQualifierTuple]:\n        if isinstance(token, Function):\n            # max(col1) AS col2\n            source_columns = [\n                cqt\n                for tk in get_parameters(token)\n                for cqt in Column._extract_source_columns(tk)\n            ]\n        elif isinstance(token, Parenthesis):\n            if is_subquery(token):\n                # This is to avoid circular import\n                from sqllineage.runner import LineageRunner\n\n                # (SELECT avg(col1) AS col1 FROM tab3), used after WHEN or THEN in CASE clause\n                src_cols = [\n                    lineage[0]\n                    for lineage in LineageRunner(token.value).get_column_lineage(\n                        exclude_subquery=False\n                    )\n                ]\n                source_columns = [\n                    ColumnQualifierTuple(src_col.raw_name, src_col.parent.raw_name)\n                    for src_col in src_cols\n                ]\n            else:\n                # (col1 + col2) AS col3\n                source_columns = [\n                    cqt\n                    for tk in token.tokens[1:-1]\n                    for cqt in Column._extract_source_columns(tk)\n                ]\n        elif isinstance(token, Operation):\n            # col1 + col2 AS col3\n            source_columns = [\n                cqt\n                for tk in token.get_sublists()\n                for cqt in Column._extract_source_columns(tk)\n            ]\n        elif isinstance(token, Case):\n            # CASE WHEN col1 = 2 THEN \"V1\" WHEN col1 = \"2\" THEN \"V2\" END AS col2\n            source_columns = [\n                cqt\n                for tk in token.get_sublists()\n                for cqt in Column._extract_source_columns(tk)\n            ]\n        elif isinstance(token, Comparison):\n            source_columns = Column._extract_source_columns(\n                token.left\n            ) + Column._extract_source_columns(token.right)\n        elif isinstance(token, IdentifierList):\n            source_columns = [\n                cqt\n                for tk in token.get_sublists()\n                for cqt in Column._extract_source_columns(tk)\n            ]\n        elif isinstance(token, Identifier):\n            if token.get_real_name():\n                # col1 AS col2\n                source_columns = [\n                    ColumnQualifierTuple(token.get_real_name(), token.get_parent_name())\n                ]\n            else:\n                # col1=1 AS int\n                source_columns = [\n                    cqt\n                    for tk in token.get_sublists()\n                    for cqt in Column._extract_source_columns(tk)",
  "start_line": 151,
  "end_line": 300,
  "total_lines": 348,
  "is_truncated": true
}

{
  "status": "ok",
  "filepath": "sqllineage/core/models.py",
  "content": "                ]\n        else:\n            # Handle literals other than *\n            if (\n                token.ttype is not None\n                and token.ttype[0] == T.Literal[0]\n                and token.value != \"*\"\n            ):\n                source_columns = []\n            else:\n                # select *\n                source_columns = [ColumnQualifierTuple(token.value, None)]\n        return source_columns\n\n    def to_source_columns(self, alias_mapping: Dict[str, Union[Table, SubQuery]]):\n        \"\"\"\n        Best guess for source table given all the possible table/subquery and their alias.\n        \"\"\"\n\n        def _to_src_col(name: str, parent: Union[Table, SubQuery] = None):\n            col = Column(name)\n            if parent:\n                col.parent = parent\n            return col\n\n        source_columns = set()\n        for (src_col, qualifier) in self.source_columns:\n            if qualifier is None:\n                if src_col == \"*\":\n                    # select *\n                    for table in set(alias_mapping.values()):\n                        source_columns.add(_to_src_col(src_col, table))\n                else:\n                    # select unqualified column\n                    src_col = _to_src_col(src_col, None)\n                    for table in set(alias_mapping.values()):\n                        # in case of only one table, we get the right answer\n                        # in case of multiple tables, a bunch of possible tables are set\n                        src_col.parent = table\n                    source_columns.add(src_col)\n            else:\n                if alias_mapping.get(qualifier):\n                    source_columns.add(\n                        _to_src_col(src_col, alias_mapping.get(qualifier))\n                    )\n                else:\n                    source_columns.add(_to_src_col(src_col, Table(qualifier)))\n        return source_columns",
  "start_line": 301,
  "end_line": 348,
  "total_lines": 348,
  "is_truncated": false
}

{
  "status": "ok",
  "stdout": "74:            holder.add_read(tables[1])\n",
  "stderr": "",
  "exit_code": 0
}

{
  "status": "ok",
  "filepath": "sqllineage/core/analyzer.py",
  "content": "from functools import reduce\nfrom operator import add\nfrom typing import List, NamedTuple, Optional, Set, Union\n\nfrom sqlparse.sql import (\n    Comment,\n    Function,\n    Identifier,\n    IdentifierList,\n    Statement,\n    TokenList,\n    Where,\n)\n\nfrom sqllineage.core.handlers.base import (\n    CurrentTokenBaseHandler,\n    NextTokenBaseHandler,\n)\nfrom sqllineage.core.holders import StatementLineageHolder, SubQueryLineageHolder\nfrom sqllineage.core.models import SubQuery, Table\nfrom sqllineage.utils.sqlparse import get_subquery_parentheses, is_subquery\n\n\nclass AnalyzerContext(NamedTuple):\n    subquery: Optional[SubQuery] = None\n    prev_cte: Optional[Set[SubQuery]] = None\n\n\nclass LineageAnalyzer:\n    \"\"\"SQL Statement Level Lineage Analyzer.\"\"\"\n\n    def analyze(self, stmt: Statement) -> StatementLineageHolder:\n        \"\"\"\n        to analyze the Statement and store the result into :class:`sqllineage.holders.StatementLineageHolder`.\n\n        :param stmt: a SQL statement parsed by `sqlparse`\n        \"\"\"\n        if (\n            stmt.get_type() == \"DELETE\"\n            or stmt.token_first(skip_cm=True).normalized == \"TRUNCATE\"\n            or stmt.token_first(skip_cm=True).normalized.upper() == \"REFRESH\"\n            or stmt.token_first(skip_cm=True).normalized == \"CACHE\"\n            or stmt.token_first(skip_cm=True).normalized.upper() == \"UNCACHE\"\n            or stmt.token_first(skip_cm=True).normalized == \"SHOW\"\n        ):\n            holder = StatementLineageHolder()\n        elif stmt.get_type() == \"DROP\":\n            holder = self._extract_from_ddl_drop(stmt)\n        elif stmt.get_type() == \"ALTER\":\n            holder = self._extract_from_ddl_alter(stmt)\n        else:\n            # DML parsing logic also applies to CREATE DDL\n            holder = StatementLineageHolder.of(\n                self._extract_from_dml(stmt, AnalyzerContext())\n            )\n        return holder\n\n    @classmethod\n    def _extract_from_ddl_drop(cls, stmt: Statement) -> StatementLineageHolder:\n        holder = StatementLineageHolder()\n        for table in {Table.of(t) for t in stmt.tokens if isinstance(t, Identifier)}:\n            holder.add_drop(table)\n        return holder\n\n    @classmethod\n    def _extract_from_ddl_alter(cls, stmt: Statement) -> StatementLineageHolder:\n        holder = StatementLineageHolder()\n        tables = [Table.of(t) for t in stmt.tokens if isinstance(t, Identifier)]\n        keywords = [t for t in stmt.tokens if t.is_keyword]\n        if any(k.normalized == \"RENAME\" for k in keywords) and len(tables) == 2:\n            holder.add_rename(tables[0], tables[1])\n        if any(k.normalized == \"EXCHANGE\" for k in keywords) and len(tables) == 2:\n            holder.add_write(tables[0])\n            holder.add_read(tables[1])\n        return holder\n\n    @classmethod\n    def _extract_from_dml(\n        cls, token: TokenList, context: AnalyzerContext\n    ) -> SubQueryLineageHolder:\n        holder = SubQueryLineageHolder()\n        if context.prev_cte is not None:\n            # CTE can be referenced by subsequent CTEs\n            for cte in context.prev_cte:\n                holder.add_cte(cte)\n        if context.subquery is not None:\n            # If within subquery, then manually add subquery as target table\n            holder.add_write(context.subquery)\n        current_handlers = [\n            handler_cls() for handler_cls in CurrentTokenBaseHandler.__subclasses__()\n        ]\n        next_handlers = [\n            handler_cls() for handler_cls in NextTokenBaseHandler.__subclasses__()\n        ]\n\n        subqueries = []\n        for sub_token in token.tokens:\n            if cls.__token_negligible_before_tablename(sub_token):\n                continue\n\n            for sq in cls.parse_subquery(sub_token):\n                # Collecting subquery on the way, hold on parsing until last\n                # so that each handler don't have to worry about what's inside subquery\n                subqueries.append(sq)\n\n            for current_handler in current_handlers:\n                current_handler.handle(sub_token, holder)\n\n            if sub_token.is_keyword:\n                for next_handler in next_handlers:\n                    next_handler.indicate(sub_token)\n                continue\n\n            for next_handler in next_handlers:\n                if next_handler.indicator:\n                    next_handler.handle(sub_token, holder)\n        else:\n            # call end of query hook here as loop is over\n            for next_handler in next_handlers:\n                next_handler.end_of_query_cleanup(holder)\n        # By recursively extracting each subquery of the parent and merge, we're doing Depth-first search\n        for sq in subqueries:\n            holder |= cls._extract_from_dml(sq.token, AnalyzerContext(sq, holder.cte))\n        return holder\n\n    @classmethod\n    def __token_negligible_before_tablename(cls, token: TokenList) -> bool:\n        return token.is_whitespace or isinstance(token, Comment)\n\n    @classmethod\n    def parse_subquery(cls, token: TokenList) -> List[SubQuery]:\n        result = []\n        if isinstance(token, (Identifier, Function, Where)):\n            # usually SubQuery is an Identifier, but not all Identifiers are SubQuery\n            # Function for CTE without AS keyword\n            result = cls._parse_subquery(token)\n        elif isinstance(token, IdentifierList):\n            # IdentifierList for SQL89 style of JOIN or multiple CTEs, this is actually SubQueries\n            result = reduce(\n                add,\n                [\n                    cls._parse_subquery(identifier)\n                    for identifier in token.get_sublists()\n                ],\n                [],\n            )\n        elif is_subquery(token):\n            # Parenthesis for SubQuery without alias, this is valid syntax for certain SQL dialect\n            result = [SubQuery.of(token, None)]\n        return result",
  "start_line": 1,
  "end_line": 150,
  "total_lines": 162,
  "is_truncated": true
}

{
  "status": "ok",
  "stdout": "sqllineage/core/handlers/target.py:40:            holder.add_read(Table.of(token.right))\nsqllineage/core/handlers/swap_partition.py:22:            holder.add_read(Table(escape_identifier_name(identifiers[0].normalized)))\nsqllineage/core/handlers/source.py:97:            holder.add_read(tbl)\nsqllineage/core/handlers/target.py:27:            holder.add_write(Table.of(token.token_first(skip_cm=True)))\nsqllineage/core/handlers/target.py:39:            holder.add_write(Table.of(token.left))\nsqllineage/core/handlers/target.py:42:            holder.add_write(Path(token.value))\nsqllineage/core/handlers/target.py:53:                holder.add_write(Table.of(token))\nsqllineage/core/handlers/swap_partition.py:23:            holder.add_write(Table(escape_identifier_name(identifiers[3].normalized)))\n",
  "stderr": "",
  "exit_code": 0
}

{
  "status": "ok",
  "filepath": "sqllineage/core/handlers/target.py",
  "content": "from sqlparse.sql import Comparison, Function, Identifier, Token\nfrom sqlparse.tokens import Literal, Number\n\nfrom sqllineage.core.handlers.base import NextTokenBaseHandler\nfrom sqllineage.core.holders import SubQueryLineageHolder\nfrom sqllineage.core.models import Path, Table\nfrom sqllineage.exceptions import SQLLineageException\n\n\nclass TargetHandler(NextTokenBaseHandler):\n    \"\"\"Target Table Handler.\"\"\"\n\n    TARGET_TABLE_TOKENS = (\"INTO\", \"OVERWRITE\", \"TABLE\", \"VIEW\", \"UPDATE\", \"COPY\")\n\n    def _indicate(self, token: Token) -> bool:\n        return token.normalized in self.TARGET_TABLE_TOKENS\n\n    def _handle(self, token: Token, holder: SubQueryLineageHolder) -> None:\n        if isinstance(token, Function):\n            # insert into tab (col1, col2) values (val1, val2); Here tab (col1, col2) will be parsed as Function\n            # referring https://github.com/andialbrecht/sqlparse/issues/483 for further information\n            if not isinstance(token.token_first(skip_cm=True), Identifier):\n                raise SQLLineageException(\n                    \"An Identifier is expected, got %s[value: %s] instead.\"\n                    % (type(token).__name__, token)\n                )\n            holder.add_write(Table.of(token.token_first(skip_cm=True)))\n        elif isinstance(token, Comparison):\n            # create table tab1 like tab2, tab1 like tab2 will be parsed as Comparison\n            # referring https://github.com/andialbrecht/sqlparse/issues/543 for further information\n            if not (\n                isinstance(token.left, Identifier)\n                and isinstance(token.right, Identifier)\n            ):\n                raise SQLLineageException(\n                    \"An Identifier is expected, got %s[value: %s] instead.\"\n                    % (type(token).__name__, token)\n                )\n            holder.add_write(Table.of(token.left))\n            holder.add_read(Table.of(token.right))\n        elif token.ttype == Literal.String.Single:\n            holder.add_write(Path(token.value))\n        else:\n            if not isinstance(token, Identifier):\n                raise SQLLineageException(\n                    \"An Identifier is expected, got %s[value: %s] instead.\"\n                    % (type(token).__name__, token)\n                )\n            if token.token_first(skip_cm=True).ttype is Number.Integer:\n                # Special Handling for Spark Bucket Table DDL\n                pass\n            else:\n                holder.add_write(Table.of(token))",
  "start_line": 1,
  "end_line": 53,
  "total_lines": 53,
  "is_truncated": false
}

{
  "status": "ok",
  "stdout": "./sqllineage/core/handlers/source.py:116:                        holder.add_column_lineage(src_col, tgt_col)\n./sqllineage/core/holders.py:87:    def add_column_lineage(self, src: Column, tgt: Column) -> None:\n",
  "stderr": "",
  "exit_code": 0
}

{
  "status": "ok",
  "filepath": "sqllineage/core/handlers/source.py",
  "content": "import re\nfrom typing import Dict, List, Union\n\nfrom sqlparse.sql import (\n    Case,\n    Function,\n    Identifier,\n    IdentifierList,\n    Operation,\n    Parenthesis,\n    Token,\n)\nfrom sqlparse.tokens import Literal, Wildcard\n\nfrom sqllineage.core.handlers.base import NextTokenBaseHandler\nfrom sqllineage.core.holders import SubQueryLineageHolder\nfrom sqllineage.core.models import Column, Path, SubQuery, Table\nfrom sqllineage.exceptions import SQLLineageException\nfrom sqllineage.utils.constant import EdgeType\nfrom sqllineage.utils.sqlparse import get_subquery_parentheses\n\n\nclass SourceHandler(NextTokenBaseHandler):\n    \"\"\"Source Table & Column Handler.\"\"\"\n\n    SOURCE_TABLE_TOKENS = (\n        r\"FROM\",\n        # inspired by https://github.com/andialbrecht/sqlparse/blob/master/sqlparse/keywords.py\n        r\"((LEFT\\s+|RIGHT\\s+|FULL\\s+)?(INNER\\s+|OUTER\\s+|STRAIGHT\\s+)?|(CROSS\\s+|NATURAL\\s+)?)?JOIN\",\n    )\n\n    def __init__(self):\n        self.column_flag = False\n        self.columns = []\n        self.tables = []\n        self.union_barriers = []\n        super().__init__()\n\n    def _indicate(self, token: Token) -> bool:\n        if token.normalized in (\"UNION\", \"UNION ALL\"):\n            self.union_barriers.append((len(self.columns), len(self.tables)))\n\n        if any(re.match(regex, token.normalized) for regex in self.SOURCE_TABLE_TOKENS):\n            self.column_flag = False\n            return True\n        elif bool(token.normalized == \"SELECT\"):\n            self.column_flag = True\n            return True\n        else:\n            return False\n\n    def _handle(self, token: Token, holder: SubQueryLineageHolder) -> None:\n        if self.column_flag:\n            self._handle_column(token)\n        else:\n            self._handle_table(token, holder)\n\n    def _handle_table(self, token: Token, holder: SubQueryLineageHolder) -> None:\n        if isinstance(token, Identifier):\n            self.tables.append(self._get_dataset_from_identifier(token, holder))\n        elif isinstance(token, IdentifierList):\n            # This is to support join in ANSI-89 syntax\n            for identifier in token.get_sublists():\n                self.tables.append(\n                    self._get_dataset_from_identifier(identifier, holder)\n                )\n        elif isinstance(token, Parenthesis):\n            # SELECT col1 FROM (SELECT col2 FROM tab1), the subquery will be parsed as Parenthesis\n            # This syntax without alias for subquery is invalid in MySQL, while valid for SparkSQL\n            self.tables.append(SubQuery.of(token, None))\n        elif token.ttype == Literal.String.Single:\n            self.tables.append(Path(token.value))\n        else:\n            raise SQLLineageException(\n                \"An Identifier is expected, got %s[value: %s] instead.\"\n                % (type(token).__name__, token)\n            )\n\n    def _handle_column(self, token: Token) -> None:\n        column_token_types = (Identifier, Function, Operation, Case)\n        if isinstance(token, column_token_types) or token.ttype is Wildcard:\n            column_tokens = [token]\n        elif isinstance(token, IdentifierList):\n            column_tokens = [\n                sub_token\n                for sub_token in token.tokens\n                if isinstance(sub_token, column_token_types)\n            ]\n        else:\n            # SELECT constant value will end up here\n            column_tokens = []\n        for token in column_tokens:\n            self.columns.append(Column.of(token))\n\n    def end_of_query_cleanup(self, holder: SubQueryLineageHolder) -> None:\n        for i, tbl in enumerate(self.tables):\n            holder.add_read(tbl)\n        self.union_barriers.append((len(self.columns), len(self.tables)))\n        for i, (col_barrier, tbl_barrier) in enumerate(self.union_barriers):\n            prev_col_barrier, prev_tbl_barrier = (\n                (0, 0) if i == 0 else self.union_barriers[i - 1]\n            )\n            col_grp = self.columns[prev_col_barrier:col_barrier]\n            tbl_grp = self.tables[prev_tbl_barrier:tbl_barrier]\n            tgt_tbl = None\n            if holder.write:\n                if len(holder.write) > 1:\n                    raise SQLLineageException\n                tgt_tbl = list(holder.write)[0]\n            if tgt_tbl:\n                for tgt_col in col_grp:\n                    tgt_col.parent = tgt_tbl\n                    for src_col in tgt_col.to_source_columns(\n                        self._get_alias_mapping_from_table_group(tbl_grp, holder)\n                    ):\n                        holder.add_column_lineage(src_col, tgt_col)\n\n    @classmethod\n    def _get_dataset_from_identifier(\n        cls, identifier: Identifier, holder: SubQueryLineageHolder\n    ) -> Union[Path, Table, SubQuery]:\n        path_match = re.match(r\"(parquet|csv|json)\\.`(.*)`\", identifier.value)\n        if path_match is not None:\n            return Path(path_match.groups()[1])\n        else:\n            read: Union[Table, SubQuery, None] = None\n            subqueries = get_subquery_parentheses(identifier)\n            if len(subqueries) > 0:\n                # SELECT col1 FROM (SELECT col2 FROM tab1) dt, the subquery will be parsed as Identifier\n                # referring https://github.com/andialbrecht/sqlparse/issues/218 for further information\n                parenthesis, alias = subqueries[0]\n                read = SubQuery.of(parenthesis, alias)\n            else:\n                cte_dict = {s.alias: s for s in holder.cte}\n                if \".\" not in identifier.value:\n                    cte = cte_dict.get(identifier.get_real_name())\n                    if cte is not None:\n                        # could reference CTE with or without alias\n                        read = SubQuery.of(\n                            cte.token,\n                            identifier.get_alias() or identifier.get_real_name(),\n                        )\n                if read is None:\n                    read = Table.of(identifier)\n            return read\n\n    @classmethod\n    def _get_alias_mapping_from_table_group(\n        cls,\n        table_group: List[Union[Path, Table, SubQuery]],",
  "start_line": 1,
  "end_line": 150,
  "total_lines": 169,
  "is_truncated": true
}

{
  "status": "ok",
  "filepath": "sqllineage/core/models.py",
  "content": "import warnings\nfrom typing import Dict, List, Optional, Set, Union\n\nfrom sqlparse import tokens as T\nfrom sqlparse.engine import grouping\nfrom sqlparse.sql import (\n    Case,\n    Comparison,\n    Function,\n    Identifier,\n    IdentifierList,\n    Operation,\n    Parenthesis,\n    Token,\n    TokenList,\n)\nfrom sqlparse.utils import imt\n\nfrom sqllineage.exceptions import SQLLineageException\nfrom sqllineage.utils.entities import ColumnQualifierTuple\nfrom sqllineage.utils.helpers import escape_identifier_name\nfrom sqllineage.utils.sqlparse import get_parameters, is_subquery\n\n\nclass Schema:\n    unknown = \"<default>\"\n\n    def __init__(self, name: str = unknown):\n        \"\"\"\n        Data Class for Schema\n\n        :param name: schema name\n        \"\"\"\n        self.raw_name = escape_identifier_name(name)\n\n    def __str__(self):\n        return self.raw_name.lower()\n\n    def __repr__(self):\n        return \"Schema: \" + str(self)\n\n    def __eq__(self, other):\n        return type(self) is type(other) and str(self) == str(other)\n\n    def __hash__(self):\n        return hash(str(self))\n\n    def __bool__(self):\n        return str(self) != self.unknown\n\n\nclass Table:\n    def __init__(self, name: str, schema: Schema = Schema(), **kwargs):\n        \"\"\"\n        Data Class for Table\n\n        :param name: table name\n        :param schema: schema as defined by :class:`Schema`\n        \"\"\"\n        if \".\" not in name:\n            self.schema = schema\n            self.raw_name = escape_identifier_name(name)\n        else:\n            schema_name, table_name = name.rsplit(\".\", 1)\n            if len(schema_name.split(\".\")) > 2:\n                # allow db.schema as schema_name, but a.b.c as schema_name is forbidden\n                raise SQLLineageException(\"Invalid format for table name: %s.\", name)\n            self.schema = Schema(schema_name)\n            self.raw_name = escape_identifier_name(table_name)\n            if schema:\n                warnings.warn(\"Name is in schema.table format, schema param is ignored\")\n        self.alias = kwargs.pop(\"alias\", self.raw_name)\n\n    def __str__(self):\n        return f\"{self.schema}.{self.raw_name.lower()}\"\n\n    def __repr__(self):\n        return \"Table: \" + str(self)\n\n    def __eq__(self, other):\n        return type(self) is type(other) and str(self) == str(other)\n\n    def __hash__(self):\n        return hash(str(self))\n\n    @staticmethod\n    def of(identifier: Identifier) -> \"Table\":\n        # rewrite identifier's get_real_name method, by matching the last dot instead of the first dot, so that the\n        # real name for a.b.c will be c instead of b\n        dot_idx, _ = identifier._token_matching(\n            lambda token: imt(token, m=(T.Punctuation, \".\")),\n            start=len(identifier.tokens),\n            reverse=True,\n        )\n        real_name = identifier._get_first_name(dot_idx, real_name=True)\n        # rewrite identifier's get_parent_name accordingly\n        parent_name = (\n            \"\".join(\n                [\n                    escape_identifier_name(token.value)\n                    for token in identifier.tokens[:dot_idx]\n                ]\n            )\n            if dot_idx\n            else None\n        )\n        schema = Schema(parent_name) if parent_name is not None else Schema()\n        alias = identifier.get_alias()\n        kwargs = {\"alias\": alias} if alias else {}\n        return Table(real_name, schema, **kwargs)\n\n\nclass Path:\n    def __init__(self, uri: str):\n        self.uri = escape_identifier_name(uri)\n\n    def __str__(self):\n        return self.uri\n\n    def __repr__(self):\n        return \"Path: \" + str(self)\n\n    def __eq__(self, other):\n        return type(self) is type(other) and self.uri == other.uri\n\n    def __hash__(self):\n        return hash(self.uri)\n\n\nclass Partition:\n    pass\n\n\nclass SubQuery:\n    def __init__(self, token: Parenthesis, alias: Optional[str]):\n        \"\"\"\n        Data Class for SubQuery\n\n        :param token: subquery token\n        :param alias: subquery name\n        \"\"\"\n        self.token = token\n        self._query = token.value\n        self.alias = alias if alias is not None else f\"subquery_{hash(self)}\"\n\n    def __str__(self):\n        return self.alias\n\n    def __repr__(self):\n        return \"SubQuery: \" + str(self)",
  "start_line": 1,
  "end_line": 150,
  "total_lines": 348,
  "is_truncated": true
}

{
  "status": "ok",
  "filepath": "sqllineage/core/models.py",
  "content": "\n    def __eq__(self, other):\n        return type(self) is type(other) and self._query == other._query\n\n    def __hash__(self):\n        return hash(self._query)\n\n    @staticmethod\n    def of(parenthesis: Parenthesis, alias: Optional[str]) -> \"SubQuery\":\n        return SubQuery(parenthesis, alias)\n\n\nclass Column:\n    def __init__(self, name: str, **kwargs):\n        \"\"\"\n        Data Class for Column\n\n        :param name: column name\n        :param parent: :class:`Table` or :class:`SubQuery`\n        :param kwargs:\n        \"\"\"\n        self._parent: Set[Union[Table, SubQuery]] = set()\n        self.raw_name = escape_identifier_name(name)\n        self.source_columns = kwargs.pop(\"source_columns\", ((self.raw_name, None),))\n\n    def __str__(self):\n        return (\n            f\"{self.parent}.{self.raw_name.lower()}\"\n            if self.parent is not None and not isinstance(self.parent, Path)\n            else f\"{self.raw_name.lower()}\"\n        )\n\n    def __repr__(self):\n        return \"Column: \" + str(self)\n\n    def __eq__(self, other):\n        return type(self) is type(other) and str(self) == str(other)\n\n    def __hash__(self):\n        return hash(str(self))\n\n    @property\n    def parent(self) -> Optional[Union[Table, SubQuery]]:\n        return list(self._parent)[0] if len(self._parent) == 1 else None\n\n    @parent.setter\n    def parent(self, value: Union[Table, SubQuery]):\n        self._parent.add(value)\n\n    @property\n    def parent_candidates(self) -> List[Union[Table, SubQuery]]:\n        return sorted(self._parent, key=lambda p: str(p))\n\n    @staticmethod\n    def of(token: Token):\n        if isinstance(token, Identifier):\n            alias = token.get_alias()\n            if alias:\n                # handle column alias, including alias for column name or Case, Function\n                kw_idx, kw = token.token_next_by(m=(T.Keyword, \"AS\"))\n                if kw_idx is None:\n                    # alias without AS\n                    kw_idx, _ = token.token_next_by(i=Identifier)\n                if kw_idx is None:\n                    # invalid syntax: col AS, without alias\n                    return Column(alias)\n                else:\n                    idx, _ = token.token_prev(kw_idx, skip_cm=True)\n                    expr = grouping.group(TokenList(token.tokens[: idx + 1]))[0]\n                    source_columns = Column._extract_source_columns(expr)\n                    return Column(alias, source_columns=source_columns)\n            else:\n                # select column name directly without alias\n                return Column(\n                    token.get_real_name(),\n                    source_columns=((token.get_real_name(), token.get_parent_name()),),\n                )\n        else:\n            # Wildcard, Case, Function without alias (thus not recognized as an Identifier)\n            source_columns = Column._extract_source_columns(token)\n            return Column(token.value, source_columns=source_columns)\n\n    @staticmethod\n    def _extract_source_columns(token: Token) -> List[ColumnQualifierTuple]:\n        if isinstance(token, Function):\n            # max(col1) AS col2\n            source_columns = [\n                cqt\n                for tk in get_parameters(token)\n                for cqt in Column._extract_source_columns(tk)\n            ]\n        elif isinstance(token, Parenthesis):\n            if is_subquery(token):\n                # This is to avoid circular import\n                from sqllineage.runner import LineageRunner\n\n                # (SELECT avg(col1) AS col1 FROM tab3), used after WHEN or THEN in CASE clause\n                src_cols = [\n                    lineage[0]\n                    for lineage in LineageRunner(token.value).get_column_lineage(\n                        exclude_subquery=False\n                    )\n                ]\n                source_columns = [\n                    ColumnQualifierTuple(src_col.raw_name, src_col.parent.raw_name)\n                    for src_col in src_cols\n                ]\n            else:\n                # (col1 + col2) AS col3\n                source_columns = [\n                    cqt\n                    for tk in token.tokens[1:-1]\n                    for cqt in Column._extract_source_columns(tk)\n                ]\n        elif isinstance(token, Operation):\n            # col1 + col2 AS col3\n            source_columns = [\n                cqt\n                for tk in token.get_sublists()\n                for cqt in Column._extract_source_columns(tk)\n            ]\n        elif isinstance(token, Case):\n            # CASE WHEN col1 = 2 THEN \"V1\" WHEN col1 = \"2\" THEN \"V2\" END AS col2\n            source_columns = [\n                cqt\n                for tk in token.get_sublists()\n                for cqt in Column._extract_source_columns(tk)\n            ]\n        elif isinstance(token, Comparison):\n            source_columns = Column._extract_source_columns(\n                token.left\n            ) + Column._extract_source_columns(token.right)\n        elif isinstance(token, IdentifierList):\n            source_columns = [\n                cqt\n                for tk in token.get_sublists()\n                for cqt in Column._extract_source_columns(tk)\n            ]\n        elif isinstance(token, Identifier):\n            if token.get_real_name():\n                # col1 AS col2\n                source_columns = [\n                    ColumnQualifierTuple(token.get_real_name(), token.get_parent_name())\n                ]\n            else:\n                # col1=1 AS int\n                source_columns = [\n                    cqt\n                    for tk in token.get_sublists()\n                    for cqt in Column._extract_source_columns(tk)",
  "start_line": 151,
  "end_line": 300,
  "total_lines": 348,
  "is_truncated": true
}

{
  "status": "ok",
  "filepath": "sqllineage/core/holders.py",
  "content": "import itertools\nfrom typing import Set, Tuple, Union\n\nimport networkx as nx\nfrom networkx import DiGraph\n\nfrom sqllineage.core.models import Column, Path, SubQuery, Table\nfrom sqllineage.utils.constant import EdgeType, NodeTag\n\nDATASET_CLASSES = (Path, Table)\n\n\nclass ColumnLineageMixin:\n    def get_column_lineage(self, exclude_subquery=True) -> Set[Tuple[Column, ...]]:\n        self.graph: DiGraph  # For mypy attribute checking\n        # filter all the column node in the graph\n        column_nodes = [n for n in self.graph.nodes if isinstance(n, Column)]\n        column_graph = self.graph.subgraph(column_nodes)\n        source_columns = {column for column, deg in column_graph.in_degree if deg == 0}\n        # if a column lineage path ends at SubQuery, then it should be pruned\n        target_columns = {\n            node\n            for node, deg in column_graph.out_degree\n            if isinstance(node, Column) and deg == 0\n        }\n        if exclude_subquery:\n            target_columns = {\n                node for node in target_columns if isinstance(node.parent, Table)\n            }\n        columns = set()\n        for (source, target) in itertools.product(source_columns, target_columns):\n            simple_paths = list(nx.all_simple_paths(self.graph, source, target))\n            if len(simple_paths) == 1:\n                columns.add(tuple(simple_paths[0]))\n            # we can ignore when simple path doesn't exist, but could there be more than one simple path?\n        return columns\n\n\nclass SubQueryLineageHolder(ColumnLineageMixin):\n    \"\"\"\n    SubQuery/Query Level Lineage Result.\n\n    SubQueryLineageHolder will hold attributes like read, write, cte\n\n    Each of them is a Set[:class:`sqllineage.models.Table`].\n\n    This is the most atomic representation of lineage result.\n    \"\"\"\n\n    def __init__(self) -> None:\n        self.graph = nx.DiGraph()\n\n    def __or__(self, other):\n        self.graph = nx.compose(self.graph, other.graph)\n        return self\n\n    def _property_getter(self, prop) -> Set[Union[SubQuery, Table]]:\n        return {t for t, attr in self.graph.nodes(data=True) if attr.get(prop) is True}\n\n    def _property_setter(self, value, prop) -> None:\n        self.graph.add_node(value, **{prop: True})\n\n    @property\n    def read(self) -> Set[Union[SubQuery, Table]]:\n        return self._property_getter(NodeTag.READ)\n\n    def add_read(self, value) -> None:\n        self._property_setter(value, NodeTag.READ)\n        # the same table can be added (in SQL: joined) multiple times with different alias\n        if hasattr(value, \"alias\"):\n            self.graph.add_edge(value, value.alias, type=EdgeType.HAS_ALIAS)\n\n    @property\n    def write(self) -> Set[Union[SubQuery, Table]]:\n        return self._property_getter(NodeTag.WRITE)\n\n    def add_write(self, value) -> None:\n        self._property_setter(value, NodeTag.WRITE)\n\n    @property\n    def cte(self) -> Set[SubQuery]:\n        return self._property_getter(NodeTag.CTE)  # type: ignore\n\n    def add_cte(self, value) -> None:\n        self._property_setter(value, NodeTag.CTE)\n\n    def add_column_lineage(self, src: Column, tgt: Column) -> None:\n        self.graph.add_edge(src, tgt, type=EdgeType.LINEAGE)\n        self.graph.add_edge(tgt.parent, tgt, type=EdgeType.HAS_COLUMN)\n        if src.parent is not None:\n            # starting NetworkX v2.6, None is not allowed as node, see https://github.com/networkx/networkx/pull/4892\n            self.graph.add_edge(src.parent, src, type=EdgeType.HAS_COLUMN)\n\n\nclass StatementLineageHolder(SubQueryLineageHolder, ColumnLineageMixin):\n    \"\"\"\n    Statement Level Lineage Result.\n\n    Based on SubQueryLineageHolder, StatementLineageHolder holds extra attributes like drop and rename\n\n    For drop, it is a Set[:class:`sqllineage.models.Table`].\n\n    For rename, it a Set[Tuple[:class:`sqllineage.models.Table`, :class:`sqllineage.models.Table`]], with the first\n    table being original table before renaming and the latter after renaming.\n    \"\"\"\n\n    def __str__(self):\n        return \"\\n\".join(\n            f\"table {attr}: {sorted(getattr(self, attr), key=lambda x: str(x)) if getattr(self, attr) else '[]'}\"\n            for attr in [\"read\", \"write\", \"cte\", \"drop\", \"rename\"]\n        )\n\n    def __repr__(self):\n        return str(self)\n\n    @property\n    def read(self) -> Set[Table]:  # type: ignore\n        return {t for t in super().read if isinstance(t, DATASET_CLASSES)}\n\n    @property\n    def write(self) -> Set[Table]:  # type: ignore\n        return {t for t in super().write if isinstance(t, DATASET_CLASSES)}\n\n    @property\n    def drop(self) -> Set[Table]:\n        return self._property_getter(NodeTag.DROP)  # type: ignore\n\n    def add_drop(self, value) -> None:\n        self._property_setter(value, NodeTag.DROP)\n\n    @property\n    def rename(self) -> Set[Tuple[Table, Table]]:\n        return {\n            (src, tgt)\n            for src, tgt, attr in self.graph.edges(data=True)\n            if attr.get(\"type\") == EdgeType.RENAME\n        }\n\n    def add_rename(self, src: Table, tgt: Table) -> None:\n        self.graph.add_edge(src, tgt, type=EdgeType.RENAME)\n\n    @staticmethod\n    def of(holder: SubQueryLineageHolder):\n        stmt_holder = StatementLineageHolder()\n        stmt_holder.graph = holder.graph\n        return stmt_holder\n\n\nclass SQLLineageHolder(ColumnLineageMixin):\n    def __init__(self, graph: DiGraph):",
  "start_line": 1,
  "end_line": 150,
  "total_lines": 280,
  "is_truncated": true
}