From 8dc2a9ccbec63b9a2dec577547e301046a02a4d8 Mon Sep 17 00:00:00 2001 From: Toby Mao Date: Mon, 2 Oct 2023 08:52:04 -0700 Subject: [PATCH] Feat: add the ability to set meta in sql comments (#2351) --- sqlglot/expressions.py | 12 +++++++++++- tests/test_expressions.py | 5 +++++ 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py index 9a347f3f0c..b278258b97 100644 --- a/sqlglot/expressions.py +++ b/sqlglot/expressions.py @@ -52,6 +52,9 @@ def __new__(cls, clsname, bases, attrs): return klass +SQLGLOT_META = "sqlglot.meta" + + class Expression(metaclass=_Expression): """ The base class for all expressions in a syntax tree. Each Expression encapsulates any necessary @@ -266,7 +269,14 @@ def add_comments(self, comments: t.Optional[t.List[str]]) -> None: if self.comments is None: self.comments = [] if comments: - self.comments.extend(comments) + for comment in comments: + _, *meta = comment.split(SQLGLOT_META) + if meta: + for kv in "".join(meta).split(","): + k, *v = kv.split("=") + value = v[0].strip() if v else True + self.meta[k.strip()] = value + self.comments.append(comment) def append(self, arg_key: str, value: t.Any) -> None: """ diff --git a/tests/test_expressions.py b/tests/test_expressions.py index 832967c290..f8c8bcc069 100644 --- a/tests/test_expressions.py +++ b/tests/test_expressions.py @@ -951,3 +951,8 @@ def test_is_type(self): with self.assertRaises(ParseError): exp.DataType.build("foo") + + def test_set_meta(self): + query = parse_one("SELECT * FROM foo /* sqlglot.meta x = 1, y = a, z */") + self.assertEqual(query.find(exp.Table).meta, {"x": "1", "y": "a", "z": True}) + self.assertEqual(query.sql(), "SELECT * FROM foo /* sqlglot.meta x = 1, y = a, z */")