@altimateai/altimate-code 0.5.1 → 0.5.3

package/dbt-tools/dist/altimate_python_packages/altimate_packages/sqlglot/dialects/risingwave.py

@@ -0,0 +1,78 @@
+from __future__ import annotations
+from sqlglot.dialects.postgres import Postgres
+from sqlglot.generator import Generator
+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"}
+
+        def computedcolumnconstraint_sql(self, expression: exp.ComputedColumnConstraint) -> str:
+            return Generator.computedcolumnconstraint_sql(self, expression)