74 lines
2.6 KiB
Python
74 lines
2.6 KiB
Python
from __future__ import annotations
|
|
from sqlglot.dialects.postgres import Postgres
|
|
from sqlglot.tokens import TokenType
|
|
import typing as t
|
|
|
|
from sqlglot import exp
|
|
|
|
|
|
class RisingWave(Postgres):
|
|
class Tokenizer(Postgres.Tokenizer):
|
|
KEYWORDS = {
|
|
**Postgres.Tokenizer.KEYWORDS,
|
|
"SINK": TokenType.SINK,
|
|
"SOURCE": TokenType.SOURCE,
|
|
}
|
|
|
|
class Parser(Postgres.Parser):
|
|
WRAPPED_TRANSFORM_COLUMN_CONSTRAINT = False
|
|
|
|
PROPERTY_PARSERS = {
|
|
**Postgres.Parser.PROPERTY_PARSERS,
|
|
"ENCODE": lambda self: self._parse_encode_property(),
|
|
"INCLUDE": lambda self: self._parse_include_property(),
|
|
"KEY": lambda self: self._parse_encode_property(key=True),
|
|
}
|
|
|
|
def _parse_table_hints(self) -> t.Optional[t.List[exp.Expression]]:
|
|
# There is no hint in risingwave.
|
|
# Do nothing here to avoid WITH keywords conflict in CREATE SINK statement.
|
|
return None
|
|
|
|
def _parse_include_property(self) -> t.Optional[exp.Expression]:
|
|
header: t.Optional[exp.Expression] = None
|
|
coldef: t.Optional[exp.Expression] = None
|
|
|
|
this = self._parse_var_or_string()
|
|
|
|
if not self._match(TokenType.ALIAS):
|
|
header = self._parse_field()
|
|
if header:
|
|
coldef = self.expression(exp.ColumnDef, this=header, kind=self._parse_types())
|
|
|
|
self._match(TokenType.ALIAS)
|
|
alias = self._parse_id_var(tokens=self.ALIAS_TOKENS)
|
|
|
|
return self.expression(exp.IncludeProperty, this=this, alias=alias, column_def=coldef)
|
|
|
|
def _parse_encode_property(self, key: t.Optional[bool] = None) -> exp.EncodeProperty:
|
|
self._match_text_seq("ENCODE")
|
|
this = self._parse_var_or_string()
|
|
|
|
if self._match(TokenType.L_PAREN, advance=False):
|
|
properties = self.expression(
|
|
exp.Properties, expressions=self._parse_wrapped_properties()
|
|
)
|
|
else:
|
|
properties = None
|
|
|
|
return self.expression(exp.EncodeProperty, this=this, properties=properties, key=key)
|
|
|
|
class Generator(Postgres.Generator):
|
|
LOCKING_READS_SUPPORTED = False
|
|
|
|
TRANSFORMS = {
|
|
**Postgres.Generator.TRANSFORMS,
|
|
exp.FileFormatProperty: lambda self, e: f"FORMAT {self.sql(e, 'this')}",
|
|
}
|
|
|
|
PROPERTIES_LOCATION = {
|
|
**Postgres.Generator.PROPERTIES_LOCATION,
|
|
exp.FileFormatProperty: exp.Properties.Location.POST_EXPRESSION,
|
|
}
|
|
|
|
EXPRESSION_PRECEDES_PROPERTIES_CREATABLES = {"SINK"}
|