diff options
Diffstat (limited to 'sqlglot/parser.py')
| -rw-r--r-- | sqlglot/parser.py | 89 |
1 files changed, 77 insertions, 12 deletions
diff --git a/sqlglot/parser.py b/sqlglot/parser.py index b7f91ab..1dab600 100644 --- a/sqlglot/parser.py +++ b/sqlglot/parser.py @@ -674,6 +674,7 @@ class Parser(metaclass=_Parser): "ON": lambda self: self._parse_on_property(), "ORDER BY": lambda self: self._parse_order(skip_order_token=True), "OUTPUT": lambda self: self.expression(exp.OutputModelProperty, this=self._parse_schema()), + "PARTITION": lambda self: self._parse_partitioned_of(), "PARTITION BY": lambda self: self._parse_partitioned_by(), "PARTITIONED BY": lambda self: self._parse_partitioned_by(), "PARTITIONED_BY": lambda self: self._parse_partitioned_by(), @@ -1743,6 +1744,58 @@ class Parser(metaclass=_Parser): return self._parse_csv(self._parse_conjunction) return [] + def _parse_partition_bound_spec(self) -> exp.PartitionBoundSpec: + def _parse_partition_bound_expr() -> t.Optional[exp.Expression]: + if self._match_text_seq("MINVALUE"): + return exp.var("MINVALUE") + if self._match_text_seq("MAXVALUE"): + return exp.var("MAXVALUE") + return self._parse_bitwise() + + this: t.Optional[exp.Expression | t.List[exp.Expression]] = None + expression = None + from_expressions = None + to_expressions = None + + if self._match(TokenType.IN): + this = self._parse_wrapped_csv(self._parse_bitwise) + elif self._match(TokenType.FROM): + from_expressions = self._parse_wrapped_csv(_parse_partition_bound_expr) + self._match_text_seq("TO") + to_expressions = self._parse_wrapped_csv(_parse_partition_bound_expr) + elif self._match_text_seq("WITH", "(", "MODULUS"): + this = self._parse_number() + self._match_text_seq(",", "REMAINDER") + expression = self._parse_number() + self._match_r_paren() + else: + self.raise_error("Failed to parse partition bound spec.") + + return self.expression( + exp.PartitionBoundSpec, + this=this, + expression=expression, + from_expressions=from_expressions, + to_expressions=to_expressions, + ) + + # https://www.postgresql.org/docs/current/sql-createtable.html + def _parse_partitioned_of(self) -> t.Optional[exp.PartitionedOfProperty]: + if not self._match_text_seq("OF"): + self._retreat(self._index - 1) + return None + + this = self._parse_table(schema=True) + + if self._match(TokenType.DEFAULT): + expression: exp.Var | exp.PartitionBoundSpec = exp.var("DEFAULT") + elif self._match_text_seq("FOR", "VALUES"): + expression = self._parse_partition_bound_spec() + else: + self.raise_error("Expecting either DEFAULT or FOR VALUES clause.") + + return self.expression(exp.PartitionedOfProperty, this=this, expression=expression) + def _parse_partitioned_by(self) -> exp.PartitionedByProperty: self._match(TokenType.EQ) return self.expression( @@ -2682,6 +2735,10 @@ class Parser(metaclass=_Parser): for join in iter(self._parse_join, None): this.append("joins", join) + if self._match_pair(TokenType.WITH, TokenType.ORDINALITY): + this.set("ordinality", True) + this.set("alias", self._parse_table_alias()) + return this def _parse_version(self) -> t.Optional[exp.Version]: @@ -4189,17 +4246,12 @@ class Parser(metaclass=_Parser): fmt = None to = self._parse_types() - if not to: - self.raise_error("Expected TYPE after CAST") - elif isinstance(to, exp.Identifier): - to = exp.DataType.build(to.name, udt=True) - elif to.this == exp.DataType.Type.CHAR: - if self._match(TokenType.CHARACTER_SET): - to = self.expression(exp.CharacterSet, this=self._parse_var_or_string()) - elif self._match(TokenType.FORMAT): + if self._match(TokenType.FORMAT): fmt_string = self._parse_string() fmt = self._parse_at_time_zone(fmt_string) + if not to: + to = exp.DataType.build(exp.DataType.Type.UNKNOWN) if to.this in exp.DataType.TEMPORAL_TYPES: this = self.expression( exp.StrToDate if to.this == exp.DataType.Type.DATE else exp.StrToTime, @@ -4215,8 +4267,14 @@ class Parser(metaclass=_Parser): if isinstance(fmt, exp.AtTimeZone) and isinstance(this, exp.StrToTime): this.set("zone", fmt.args["zone"]) - return this + elif not to: + self.raise_error("Expected TYPE after CAST") + elif isinstance(to, exp.Identifier): + to = exp.DataType.build(to.name, udt=True) + elif to.this == exp.DataType.Type.CHAR: + if self._match(TokenType.CHARACTER_SET): + to = self.expression(exp.CharacterSet, this=self._parse_var_or_string()) return self.expression( exp.Cast if strict else exp.TryCast, this=this, to=to, format=fmt, safe=safe @@ -4789,10 +4847,17 @@ class Parser(metaclass=_Parser): return self._parse_placeholder() def _parse_parameter(self) -> exp.Parameter: - wrapped = self._match(TokenType.L_BRACE) - this = self._parse_var() or self._parse_identifier() or self._parse_primary() + def _parse_parameter_part() -> t.Optional[exp.Expression]: + return ( + self._parse_identifier() or self._parse_primary() or self._parse_var(any_token=True) + ) + + self._match(TokenType.L_BRACE) + this = _parse_parameter_part() + expression = self._match(TokenType.COLON) and _parse_parameter_part() self._match(TokenType.R_BRACE) - return self.expression(exp.Parameter, this=this, wrapped=wrapped) + + return self.expression(exp.Parameter, this=this, expression=expression) def _parse_placeholder(self) -> t.Optional[exp.Expression]: if self._match_set(self.PLACEHOLDER_PARSERS): |