From 7a94d3997b896c86d6d2ba1141b4e5ac2f46c7e9 Mon Sep 17 00:00:00 2001 From: tobymao Date: Thu, 24 Nov 2022 20:45:08 -0800 Subject: [PATCH] refactor including/excluding --- sqlglot/dialects/mysql.py | 1 + sqlglot/expressions.py | 11 ++------ sqlglot/generator.py | 6 +++++ sqlglot/parser.py | 45 +++++++++++++++------------------ sqlglot/tokens.py | 4 --- tests/dialects/test_mysql.py | 2 ++ tests/dialects/test_postgres.py | 4 +++ 7 files changed, 35 insertions(+), 38 deletions(-) diff --git a/sqlglot/dialects/mysql.py b/sqlglot/dialects/mysql.py index 93a60f494a..7627b6e666 100644 --- a/sqlglot/dialects/mysql.py +++ b/sqlglot/dialects/mysql.py @@ -453,6 +453,7 @@ class Generator(generator.Generator): exp.CharacterSetProperty, exp.CollateProperty, exp.SchemaCommentProperty, + exp.LikeProperty, } WITH_PROPERTIES: t.Set[t.Type[exp.Property]] = set() diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py index 574cf28492..11893e8035 100644 --- a/sqlglot/expressions.py +++ b/sqlglot/expressions.py @@ -608,7 +608,6 @@ class Create(Expression): "this": True, "kind": True, "expression": False, - "like_properties": False, "exists": False, "properties": False, "temporary": False, @@ -1076,12 +1075,8 @@ class DistStyleProperty(Property): arg_types = {"this": True} -class IncludingProperty(Property): - arg_types = {"this": True} - - -class ExcludingProperty(Property): - arg_types = {"this": True} +class LikeProperty(Property): + arg_types = {"this": True, "expressions": False} class LocationProperty(Property): @@ -1137,8 +1132,6 @@ class Properties(Expression): "ENGINE": EngineProperty, "EXECUTE AS": ExecuteAsProperty, "FORMAT": FileFormatProperty, - "INCLUDING": IncludingProperty, - "EXCLUDING": ExcludingProperty, "LANGUAGE": LanguageProperty, "LOCATION": LocationProperty, "PARTITIONED_BY": PartitionedByProperty, diff --git a/sqlglot/generator.py b/sqlglot/generator.py index 2792dada09..cd2d448c33 100644 --- a/sqlglot/generator.py +++ b/sqlglot/generator.py @@ -94,6 +94,7 @@ class Generator: exp.DistStyleProperty, exp.DistKeyProperty, exp.SortKeyProperty, + exp.LikeProperty, } WITH_PROPERTIES = { @@ -568,6 +569,11 @@ def property_sql(self, expression): return f"{property_name}={self.sql(expression, 'this')}" + def likeproperty_sql(self, expression): + options = " ".join(f"{e.name} {self.sql(e, 'value')}" for e in expression.expressions) + options = f" {options}" if options else "" + return f"LIKE {self.sql(expression, 'this')}{options}" + def insert_sql(self, expression): overwrite = expression.args.get("overwrite") diff --git a/sqlglot/parser.py b/sqlglot/parser.py index 455a42a38b..ba136b86a9 100644 --- a/sqlglot/parser.py +++ b/sqlglot/parser.py @@ -175,7 +175,6 @@ class Parser(metaclass=_Parser): TokenType.GENERATED, TokenType.IDENTITY, TokenType.IF, - TokenType.INCLUDING, TokenType.INDEX, TokenType.ISNULL, TokenType.IMMUTABLE, @@ -471,8 +470,7 @@ class Parser(metaclass=_Parser): TokenType.DISTKEY: lambda self: self._parse_distkey(), TokenType.DISTSTYLE: lambda self: self._parse_property_assignment(exp.DistStyleProperty), TokenType.SORTKEY: lambda self: self._parse_sortkey(), - TokenType.INCLUDING: lambda self: self._parse_property_assignment(exp.IncludingProperty), - TokenType.EXCLUDING: lambda self: self._parse_property_assignment(exp.ExcludingProperty), + TokenType.LIKE: lambda self: self._parse_create_like(), TokenType.RETURNS: lambda self: self._parse_returns(), TokenType.COLLATE: lambda self: self._parse_property_assignment(exp.CollateProperty), TokenType.COMMENT: lambda self: self._parse_property_assignment(exp.SchemaCommentProperty), @@ -503,6 +501,7 @@ class Parser(metaclass=_Parser): ), TokenType.FOREIGN_KEY: lambda self: self._parse_foreign_key(), TokenType.UNIQUE: lambda self: self._parse_unique(), + TokenType.LIKE: lambda self: self._parse_create_like(), } NO_PAREN_FUNCTION_PARSERS = { @@ -815,7 +814,6 @@ def _parse_create(self): this = None expression = None properties = None - like_properties = None if create_token.token_type in (TokenType.FUNCTION, TokenType.PROCEDURE): this = self._parse_user_defined_function() @@ -833,27 +831,12 @@ def _parse_create(self): properties = self._parse_properties() if self._match(TokenType.ALIAS): expression = self._parse_select(nested=True) - elif self._match(TokenType.LIKE): - # create table a like b - expression = self.expression( - exp.Like, this=this, expression=self._parse_table(schema=True) - ) - elif self._match_pair(TokenType.L_PAREN, TokenType.LIKE): - # create table a (like b including constraints excluding statistics) - expression = self.expression( - exp.Like, this=this, expression=self._parse_table(schema=True) - ) - - like_properties = self._parse_properties() - - self._match_r_paren() return self.expression( exp.Create, this=this, kind=create_token.text, expression=expression, - like_properties=like_properties, exists=exists, properties=properties, temporary=temporary, @@ -899,6 +882,23 @@ def _parse_stored(self): def _parse_distkey(self): return self.expression(exp.DistKeyProperty, this=self._parse_wrapped(self._parse_var)) + def _parse_create_like(self): + table = self._parse_table(schema=True) + options = [] + while True: + if self._match_texts(("INCLUDING", "EXCLUDING")): + options.append( + self.expression( + exp.Property, + this=self._prev.text.upper(), + value=exp.Var(this=self._parse_id_var().this.upper()), + ) + ) + else: + break + + return self.expression(exp.LikeProperty, this=table, expressions=options) + def _parse_sortkey(self, compound=False): return self.expression( exp.SortKeyProperty, this=self._parse_wrapped_csv(self._parse_var), compound=compound @@ -2036,13 +2036,8 @@ def _parse_lambda(self): return self._parse_alias(self._parse_limit(self._parse_order(this))) def _parse_schema(self, this=None): - index = self._index - if ( - self._match_pair(TokenType.L_PAREN, TokenType.LIKE) - or not self._match(TokenType.L_PAREN) - or self._match(TokenType.SELECT) - ): + if not self._match(TokenType.L_PAREN) or self._match(TokenType.SELECT): self._retreat(index) return this diff --git a/sqlglot/tokens.py b/sqlglot/tokens.py index 00c6a4dcf5..dda0e012b6 100644 --- a/sqlglot/tokens.py +++ b/sqlglot/tokens.py @@ -159,7 +159,6 @@ class TokenType(AutoName): ENGINE = auto() ESCAPE = auto() EXCEPT = auto() - EXCLUDING = auto() EXECUTE = auto() EXISTS = auto() FALSE = auto() @@ -185,7 +184,6 @@ class TokenType(AutoName): ILIKE = auto() IMMUTABLE = auto() IN = auto() - INCLUDING = auto() INDEX = auto() INNER = auto() INSERT = auto() @@ -482,7 +480,6 @@ class Tokenizer(metaclass=_Tokenizer): "ENGINE": TokenType.ENGINE, "ESCAPE": TokenType.ESCAPE, "EXCEPT": TokenType.EXCEPT, - "EXCLUDING": TokenType.EXCLUDING, "EXECUTE": TokenType.EXECUTE, "EXISTS": TokenType.EXISTS, "FALSE": TokenType.FALSE, @@ -504,7 +501,6 @@ class Tokenizer(metaclass=_Tokenizer): "IMMUTABLE": TokenType.IMMUTABLE, "IGNORE NULLS": TokenType.IGNORE_NULLS, "IN": TokenType.IN, - "INCLUDING": TokenType.INCLUDING, "INDEX": TokenType.INDEX, "INNER": TokenType.INNER, "INSERT": TokenType.INSERT, diff --git a/tests/dialects/test_mysql.py b/tests/dialects/test_mysql.py index af98249058..79fe32a490 100644 --- a/tests/dialects/test_mysql.py +++ b/tests/dialects/test_mysql.py @@ -23,6 +23,8 @@ def test_identity(self): self.validate_identity("SELECT TRIM('bla' FROM ' XXX ')") self.validate_identity("@@GLOBAL.max_connections") + self.validate_identity("CREATE TABLE A LIKE B") + # SET Commands self.validate_identity("SET @var_name = expr") self.validate_identity("SET @name = 43") diff --git a/tests/dialects/test_postgres.py b/tests/dialects/test_postgres.py index c41fe01d80..ff16a96afd 100644 --- a/tests/dialects/test_postgres.py +++ b/tests/dialects/test_postgres.py @@ -98,6 +98,10 @@ def test_postgres(self): "spark": "CREATE TABLE x (a UUID, b BINARY)", }, ) + + self.validate_identity( + "CREATE TABLE A (LIKE B INCLUDING CONSTRAINT INCLUDING COMPRESSION EXCLUDING COMMENTS)" + ) self.validate_all( "SELECT SUM(x) OVER (PARTITION BY a ORDER BY d ROWS 1 PRECEDING)", write={