Feat: Include token refereces in the meta of identifier expressions (#5022)

Author: izeigerman

sqlglot/dialects/bigquery.py

@@ -645,14 +645,16 @@ def _parse_table_part(self, schema: bool = False) -> t.Optional[exp.Expression]:
 
                     table_name += self._find_sql(start, self._prev)
 
-                this = exp.Identifier(this=table_name, quoted=this.args.get("quoted"))
+                this = exp.Identifier(
+                    this=table_name, quoted=this.args.get("quoted")
+                ).update_positions(this)
             elif isinstance(this, exp.Literal):
                 table_name = this.name
 
                 if self._is_connected() and self._parse_var(any_token=True):
                     table_name += self._prev.text
 
-                this = exp.Identifier(this=table_name, quoted=True)
+                this = exp.Identifier(this=table_name, quoted=True).update_positions(this)
 
             return this
 
@@ -666,15 +668,23 @@ def _parse_table_parts(
             # proj-1.db.tbl -- `1.` is tokenized as a float so we need to unravel it here
             if not table.catalog:
                 if table.db:
+                    previous_db = table.args["db"]
                     parts = table.db.split(".")
                     if len(parts) == 2 and not table.args["db"].quoted:
-                        table.set("catalog", exp.Identifier(this=parts[0]))
-                        table.set("db", exp.Identifier(this=parts[1]))
+                        table.set(
+                            "catalog", exp.Identifier(this=parts[0]).update_positions(previous_db)
+                        )
+                        table.set("db", exp.Identifier(this=parts[1]).update_positions(previous_db))
                 else:
+                    previous_this = table.this
                     parts = table.name.split(".")
                     if len(parts) == 2 and not table.this.quoted:
-                        table.set("db", exp.Identifier(this=parts[0]))
-                        table.set("this", exp.Identifier(this=parts[1]))
+                        table.set(
+                            "db", exp.Identifier(this=parts[0]).update_positions(previous_this)
+                        )
+                        table.set(
+                            "this", exp.Identifier(this=parts[1]).update_positions(previous_this)
+                        )
 
             if isinstance(table.this, exp.Identifier) and any("." in p.name for p in table.parts):
                 alias = table.this
@@ -683,6 +693,10 @@ def _parse_table_parts(
                     for p in split_num_words(".".join(p.name for p in table.parts), ".", 3)
                 )
 
+                for part in (catalog, db, this):
+                    if part:
+                        part.update_positions(table.this)
+
                 if rest and this:
                     this = exp.Dot.build([this, *rest])  # type: ignore
 
@@ -717,7 +731,13 @@ def _parse_table_parts(
                 )
 
                 info_schema_view = f"{table_parts[-2].name}.{table_parts[-1].name}"
-                table.set("this", exp.Identifier(this=info_schema_view, quoted=True))
+                new_this = exp.Identifier(this=info_schema_view, quoted=True).update_positions(
+                    line=table_parts[-2].meta.get("line"),
+                    col=table_parts[-1].meta.get("col"),
+                    start=table_parts[-2].meta.get("start"),
+                    end=table_parts[-1].meta.get("end"),
+                )
+                table.set("this", new_this)
                 table.set("db", seq_get(table_parts, -3))
                 table.set("catalog", seq_get(table_parts, -4))
 

sqlglot/expressions.py

@@ -64,6 +64,7 @@ def __new__(cls, clsname, bases, attrs):
 SQLGLOT_ANONYMOUS = "sqlglot.anonymous"
 TABLE_PARTS = ("this", "db", "catalog")
 COLUMN_PARTS = ("this", "table", "db", "catalog")
+POSITION_META_KEYS = ("line", "col", "start", "end")
 
 
 class Expression(metaclass=_Expression):
@@ -846,6 +847,32 @@ def not_(self, copy: bool = True):
         """
         return not_(self, copy=copy)
 
+    def update_positions(
+        self: E, other: t.Optional[Token | Expression] = None, **kwargs: t.Any
+    ) -> E:
+        """
+        Update this expression with positions from a token or other expression.
+
+        Args:
+            other: a token or expression to update this expression with.
+
+        Returns:
+            The updated expression.
+        """
+        if isinstance(other, Expression):
+            self.meta.update({k: v for k, v in other.meta.items() if k in POSITION_META_KEYS})
+        elif other is not None:
+            self.meta.update(
+                {
+                    "line": other.line,
+                    "col": other.col,
+                    "start": other.start,
+                    "end": other.end,
+                }
+            )
+        self.meta.update({k: v for k, v in kwargs.items() if k in POSITION_META_KEYS})
+        return self
+
     def as_(
         self,
         alias: str | Identifier,

sqlglot/parser.py

@@ -5618,6 +5618,7 @@ def _parse_function_call(
             return None
 
         comments = self._curr.comments
+        token = self._curr
         token_type = self._curr.token_type
         this = self._curr.text
         upper = this.upper()
@@ -5692,7 +5693,7 @@ def _parse_function_call(
                 this = func
             else:
                 if token_type == TokenType.IDENTIFIER:
-                    this = exp.Identifier(this=this, quoted=True)
+                    this = exp.Identifier(this=this, quoted=True).update_positions(token)
                 this = self.expression(exp.Anonymous, this=this, expressions=args)
 
         if isinstance(this, exp.Expression):
@@ -5751,7 +5752,7 @@ def _parse_introducer(self, token: Token) -> exp.Introducer | exp.Identifier:
         if literal:
             return self.expression(exp.Introducer, this=token.text, expression=literal)
 
-        return self.expression(exp.Identifier, this=token.text)
+        return self._identifier_expression(token)
 
     def _parse_session_parameter(self) -> exp.SessionParameter:
         kind = None
@@ -6981,7 +6982,7 @@ def _parse_id_var(
             (any_token and self._advance_any()) or self._match_set(tokens or self.ID_VAR_TOKENS)
         ):
             quoted = self._prev.token_type == TokenType.STRING
-            expression = self.expression(exp.Identifier, this=self._prev.text, quoted=quoted)
+            expression = self._identifier_expression(quoted=quoted)
 
         return expression
 
@@ -6991,7 +6992,10 @@ def _parse_string(self) -> t.Optional[exp.Expression]:
         return self._parse_placeholder()
 
     def _parse_string_as_identifier(self) -> t.Optional[exp.Identifier]:
-        return exp.to_identifier(self._match(TokenType.STRING) and self._prev.text, quoted=True)
+        output = exp.to_identifier(self._match(TokenType.STRING) and self._prev.text, quoted=True)
+        if output:
+            output.update_positions(self._prev)
+        return output
 
     def _parse_number(self) -> t.Optional[exp.Expression]:
         if self._match_set(self.NUMERIC_PARSERS):
@@ -7000,7 +7004,7 @@ def _parse_number(self) -> t.Optional[exp.Expression]:
 
     def _parse_identifier(self) -> t.Optional[exp.Expression]:
         if self._match(TokenType.IDENTIFIER):
-            return self.expression(exp.Identifier, this=self._prev.text, quoted=True)
+            return self._identifier_expression(quoted=True)
         return self._parse_placeholder()
 
     def _parse_var(
@@ -8246,3 +8250,11 @@ def _parse_max_min_by(self, expr_type: t.Type[exp.AggFunc]) -> exp.AggFunc:
         return self.expression(
             expr_type, this=seq_get(args, 0), expression=seq_get(args, 1), count=seq_get(args, 2)
         )
+
+    def _identifier_expression(
+        self, token: t.Optional[Token] = None, **kwargs: t.Any
+    ) -> exp.Identifier:
+        token = token or self._prev
+        expression = self.expression(exp.Identifier, this=token.text, **kwargs)
+        expression.update_positions(token)
+        return expression

tests/dialects/test_bigquery.py

@@ -2432,3 +2432,56 @@ def test_with_offset(self):
                     f"SELECT * FROM t1, UNNEST([1, 2]) AS hit WITH OFFSET {join_ops} JOIN foo",
                     f"SELECT * FROM t1, UNNEST([1, 2]) AS hit WITH OFFSET AS offset {join_ops} JOIN foo",
                 )
+
+    def test_identifier_meta(self):
+        ast = parse_one(
+            "SELECT a, b FROM test_schema.test_table_a UNION ALL SELECT c, d FROM test_catalog.test_schema.test_table_b",
+            dialect="bigquery",
+        )
+        for identifier in ast.find_all(exp.Identifier):
+            self.assertEqual(set(identifier.meta), {"line", "col", "start", "end"})
+
+        self.assertEqual(
+            ast.this.args["from"].this.args["this"].meta,
+            {"line": 1, "col": 41, "start": 29, "end": 40},
+        )
+        self.assertEqual(
+            ast.this.args["from"].this.args["db"].meta,
+            {"line": 1, "col": 28, "start": 17, "end": 27},
+        )
+        self.assertEqual(
+            ast.expression.args["from"].this.args["this"].meta,
+            {"line": 1, "col": 106, "start": 94, "end": 105},
+        )
+        self.assertEqual(
+            ast.expression.args["from"].this.args["db"].meta,
+            {"line": 1, "col": 93, "start": 82, "end": 92},
+        )
+        self.assertEqual(
+            ast.expression.args["from"].this.args["catalog"].meta,
+            {"line": 1, "col": 81, "start": 69, "end": 80},
+        )
+
+        information_schema_sql = "SELECT a, b FROM region.INFORMATION_SCHEMA.COLUMNS"
+        ast = parse_one(information_schema_sql, dialect="bigquery")
+        meta = ast.args["from"].this.this.meta
+        self.assertEqual(meta, {"line": 1, "col": 50, "start": 24, "end": 49})
+        assert (
+            information_schema_sql[meta["start"] : meta["end"] + 1] == "INFORMATION_SCHEMA.COLUMNS"
+        )
+
+    def test_quoted_identifier_meta(self):
+        sql = "SELECT `a` FROM `test_schema`.`test_table_a`"
+        ast = parse_one(sql, dialect="bigquery")
+        db_meta = ast.args["from"].this.args["db"].meta
+        self.assertEqual(sql[db_meta["start"] : db_meta["end"] + 1], "`test_schema`")
+        table_meta = ast.args["from"].this.this.meta
+        self.assertEqual(sql[table_meta["start"] : table_meta["end"] + 1], "`test_table_a`")
+
+        information_schema_sql = "SELECT a, b FROM `region.INFORMATION_SCHEMA.COLUMNS`"
+        ast = parse_one(information_schema_sql, dialect="bigquery")
+        table_meta = ast.args["from"].this.this.meta
+        assert (
+            information_schema_sql[table_meta["start"] : table_meta["end"] + 1]
+            == "`region.INFORMATION_SCHEMA.COLUMNS`"
+        )

tests/test_parser.py

@@ -956,3 +956,41 @@ def test_udf_meta(self):
         # Incomplete or incorrect anonymous meta comments are not registered
         ast = parse_one("YEAR(a) /* sqlglot.anon */")
         self.assertIsInstance(ast, exp.Year)
+
+    def test_identifier_meta(self):
+        ast = parse_one(
+            "SELECT a, b FROM test_schema.test_table_a UNION ALL SELECT c, d FROM test_catalog.test_schema.test_table_b"
+        )
+        for identifier in ast.find_all(exp.Identifier):
+            self.assertEqual(set(identifier.meta), {"line", "col", "start", "end"})
+
+        self.assertEqual(
+            ast.this.args["from"].this.args["this"].meta,
+            {"line": 1, "col": 41, "start": 29, "end": 40},
+        )
+        self.assertEqual(
+            ast.this.args["from"].this.args["db"].meta,
+            {"line": 1, "col": 28, "start": 17, "end": 27},
+        )
+        self.assertEqual(
+            ast.expression.args["from"].this.args["this"].meta,
+            {"line": 1, "col": 106, "start": 94, "end": 105},
+        )
+        self.assertEqual(
+            ast.expression.args["from"].this.args["db"].meta,
+            {"line": 1, "col": 93, "start": 82, "end": 92},
+        )
+        self.assertEqual(
+            ast.expression.args["from"].this.args["catalog"].meta,
+            {"line": 1, "col": 81, "start": 69, "end": 80},
+        )
+
+    def test_quoted_identifier_meta(self):
+        sql = 'SELECT "a" FROM "test_schema"."test_table_a"'
+        ast = parse_one(sql)
+
+        db_meta = ast.args["from"].this.args["db"].meta
+        self.assertEqual(sql[db_meta["start"] : db_meta["end"] + 1], '"test_schema"')
+
+        table_meta = ast.args["from"].this.this.meta
+        self.assertEqual(sql[table_meta["start"] : table_meta["end"] + 1], '"test_table_a"')