diff options
Diffstat (limited to 'sqlglot/parser.py')
| -rw-r--r-- | sqlglot/parser.py | 118 |
1 files changed, 86 insertions, 32 deletions
diff --git a/sqlglot/parser.py b/sqlglot/parser.py index f39bb39..894d68e 100644 --- a/sqlglot/parser.py +++ b/sqlglot/parser.py @@ -434,6 +434,7 @@ class Parser(metaclass=_Parser): exp.Having: lambda self: self._parse_having(), exp.With: lambda self: self._parse_with(), exp.Window: lambda self: self._parse_named_window(), + exp.Qualify: lambda self: self._parse_qualify(), "JOIN_TYPE": lambda self: self._parse_join_side_and_kind(), } @@ -688,6 +689,7 @@ class Parser(metaclass=_Parser): "limit": lambda self: self._parse_limit(), "offset": lambda self: self._parse_offset(), "lock": lambda self: self._parse_lock(), + "sample": lambda self: self._parse_table_sample(as_modifier=True), } SHOW_PARSERS: t.Dict[str, t.Callable] = {} @@ -953,7 +955,8 @@ class Parser(metaclass=_Parser): self._prev_comments = None def _retreat(self, index: int) -> None: - self._advance(index - self._index) + if index != self._index: + self._advance(index - self._index) def _parse_command(self) -> exp.Expression: return self.expression(exp.Command, this=self._prev.text, expression=self._parse_string()) @@ -1515,12 +1518,10 @@ class Parser(metaclass=_Parser): def _parse_insert(self) -> exp.Expression: overwrite = self._match(TokenType.OVERWRITE) local = self._match(TokenType.LOCAL) - - this: t.Optional[exp.Expression] - alternative = None + if self._match_text_seq("DIRECTORY"): - this = self.expression( + this: t.Optional[exp.Expression] = self.expression( exp.Directory, this=self._parse_var_or_string(), local=local, @@ -1540,10 +1541,17 @@ class Parser(metaclass=_Parser): exists=self._parse_exists(), partition=self._parse_partition(), expression=self._parse_ddl_select(), + returning=self._parse_returning(), overwrite=overwrite, alternative=alternative, ) + def _parse_returning(self) -> t.Optional[exp.Expression]: + if not self._match(TokenType.RETURNING): + return None + + return self.expression(exp.Returning, expressions=self._parse_csv(self._parse_column)) + def _parse_row(self) -> t.Optional[exp.Expression]: if not self._match(TokenType.FORMAT): return None @@ -1601,6 +1609,7 @@ class Parser(metaclass=_Parser): this=self._parse_table(schema=True), using=self._parse_csv(lambda: self._match(TokenType.USING) and self._parse_table()), where=self._parse_where(), + returning=self._parse_returning(), ) def _parse_update(self) -> exp.Expression: @@ -1611,6 +1620,7 @@ class Parser(metaclass=_Parser): "expressions": self._match(TokenType.SET) and self._parse_csv(self._parse_equality), "from": self._parse_from(), "where": self._parse_where(), + "returning": self._parse_returning(), }, ) @@ -2156,11 +2166,12 @@ class Parser(metaclass=_Parser): return self.expression(exp.Values, expressions=expressions, alias=self._parse_table_alias()) - def _parse_table_sample(self) -> t.Optional[exp.Expression]: - if not self._match(TokenType.TABLE_SAMPLE): + def _parse_table_sample(self, as_modifier: bool = False) -> t.Optional[exp.Expression]: + if not self._match(TokenType.TABLE_SAMPLE) and not ( + as_modifier and self._match_text_seq("USING", "SAMPLE") + ): return None - method = self._parse_var() bucket_numerator = None bucket_denominator = None bucket_field = None @@ -2169,7 +2180,12 @@ class Parser(metaclass=_Parser): size = None seed = None - self._match_l_paren() + kind = "TABLESAMPLE" if self._prev.token_type == TokenType.TABLE_SAMPLE else "USING SAMPLE" + method = self._parse_var(tokens=(TokenType.ROW,)) + + self._match(TokenType.L_PAREN) + + num = self._parse_number() if self._match(TokenType.BUCKET): bucket_numerator = self._parse_number() @@ -2177,19 +2193,20 @@ class Parser(metaclass=_Parser): bucket_denominator = bucket_denominator = self._parse_number() self._match(TokenType.ON) bucket_field = self._parse_field() + elif self._match_set((TokenType.PERCENT, TokenType.MOD)): + percent = num + elif self._match(TokenType.ROWS): + rows = num else: - num = self._parse_number() + size = num - if self._match(TokenType.PERCENT): - percent = num - elif self._match(TokenType.ROWS): - rows = num - else: - size = num + self._match(TokenType.R_PAREN) - self._match_r_paren() - - if self._match(TokenType.SEED): + if self._match(TokenType.L_PAREN): + method = self._parse_var() + seed = self._match(TokenType.COMMA) and self._parse_number() + self._match_r_paren() + elif self._match_texts(("SEED", "REPEATABLE")): seed = self._parse_wrapped(self._parse_number) return self.expression( @@ -2202,6 +2219,7 @@ class Parser(metaclass=_Parser): rows=rows, size=size, seed=seed, + kind=kind, ) def _parse_pivots(self) -> t.List[t.Optional[exp.Expression]]: @@ -2531,7 +2549,7 @@ class Parser(metaclass=_Parser): this = self._parse_column() if type_token: - if this and not isinstance(this, exp.Star): + if isinstance(this, exp.Literal): return self.expression(exp.Cast, this=this, to=type_token) if not type_token.args.get("expressions"): self._retreat(index) @@ -2626,7 +2644,12 @@ class Parser(metaclass=_Parser): if value is None: value = exp.DataType(this=exp.DataType.Type.TIMESTAMP, expressions=expressions) elif type_token == TokenType.INTERVAL: - value = self.expression(exp.Interval, unit=self._parse_var()) + unit = self._parse_var() + + if not unit: + value = self.expression(exp.DataType, this=exp.DataType.Type.INTERVAL) + else: + value = self.expression(exp.Interval, unit=unit) if maybe_func and check_func: index2 = self._index @@ -3495,8 +3518,14 @@ class Parser(metaclass=_Parser): return self.expression(exp.Identifier, this=self._prev.text, quoted=True) return self._parse_placeholder() - def _parse_var(self, any_token: bool = False) -> t.Optional[exp.Expression]: - if (any_token and self._advance_any()) or self._match(TokenType.VAR): + def _parse_var( + self, any_token: bool = False, tokens: t.Optional[t.Collection[TokenType]] = None + ) -> t.Optional[exp.Expression]: + if ( + (any_token and self._advance_any()) + or self._match(TokenType.VAR) + or (self._match_set(tokens) if tokens else False) + ): return self.expression(exp.Var, this=self._prev.text) return self._parse_placeholder() @@ -3732,19 +3761,26 @@ class Parser(metaclass=_Parser): return self.expression(exp.RenameTable, this=self._parse_table(schema=True)) def _parse_alter(self) -> t.Optional[exp.Expression]: + start = self._prev + if not self._match(TokenType.TABLE): - return self._parse_as_command(self._prev) + return self._parse_as_command(start) exists = self._parse_exists() this = self._parse_table(schema=True) - if not self._curr: - return None - - parser = self.ALTER_PARSERS.get(self._curr.text.upper()) - actions = ensure_list(self._advance() or parser(self)) if parser else [] # type: ignore + if self._next: + self._advance() + parser = self.ALTER_PARSERS.get(self._prev.text.upper()) if self._prev else None - return self.expression(exp.AlterTable, this=this, exists=exists, actions=actions) + if parser: + return self.expression( + exp.AlterTable, + this=this, + exists=exists, + actions=ensure_list(parser(self)), + ) + return self._parse_as_command(start) def _parse_show(self) -> t.Optional[exp.Expression]: parser = self._find_parser(self.SHOW_PARSERS, self._show_trie) # type: ignore @@ -3775,7 +3811,15 @@ class Parser(metaclass=_Parser): whens = [] while self._match(TokenType.WHEN): - this = self._parse_conjunction() + matched = not self._match(TokenType.NOT) + self._match_text_seq("MATCHED") + source = ( + False + if self._match_text_seq("BY", "TARGET") + else self._match_text_seq("BY", "SOURCE") + ) + condition = self._parse_conjunction() if self._match(TokenType.AND) else None + self._match(TokenType.THEN) if self._match(TokenType.INSERT): @@ -3800,8 +3844,18 @@ class Parser(metaclass=_Parser): ) elif self._match(TokenType.DELETE): then = self.expression(exp.Var, this=self._prev.text) + else: + then = None - whens.append(self.expression(exp.When, this=this, then=then)) + whens.append( + self.expression( + exp.When, + matched=matched, + source=source, + condition=condition, + then=then, + ) + ) return self.expression( exp.Merge, |