sqlmesh.core.engine_adapter.risingwave

 1from __future__ import annotations
 2
 3import logging
 4import typing as t
 5
 6
 7from sqlglot import exp
 8
 9from sqlmesh.core.engine_adapter.postgres import PostgresEngineAdapter
10from sqlmesh.core.engine_adapter.shared import (
11    set_catalog,
12    CatalogSupport,
13    CommentCreationView,
14    CommentCreationTable,
15)
16
17from sqlmesh.utils.errors import SQLMeshError
18
19if t.TYPE_CHECKING:
20    from sqlmesh.core._typing import TableName
21
22logger = logging.getLogger(__name__)
23
24
25@set_catalog()
26class RisingwaveEngineAdapter(PostgresEngineAdapter):
27    DIALECT = "risingwave"
28    DEFAULT_BATCH_SIZE = 400
29    CATALOG_SUPPORT = CatalogSupport.SINGLE_CATALOG_ONLY
30    COMMENT_CREATION_TABLE = CommentCreationTable.COMMENT_COMMAND_ONLY
31    COMMENT_CREATION_VIEW = CommentCreationView.UNSUPPORTED
32    SUPPORTS_MATERIALIZED_VIEWS = True
33    SUPPORTS_TRANSACTIONS = False
34    MAX_IDENTIFIER_LENGTH = None
35    SUPPORTS_GRANTS = False
36
37    def columns(
38        self, table_name: TableName, include_pseudo_columns: bool = False
39    ) -> t.Dict[str, exp.DataType]:
40        """Fetches column names and types for the target_table"""
41        table = exp.to_table(table_name)
42
43        sql = (
44            exp.select("rw_columns.name AS column_name", "rw_columns.data_type AS data_type")
45            .from_("rw_catalog.rw_columns")
46            .join("rw_catalog.rw_relations", on="rw_relations.id=rw_columns.relation_id")
47            .join("rw_catalog.rw_schemas", on="rw_schemas.id=rw_relations.schema_id")
48            .where(
49                exp.and_(
50                    exp.column("name", table="rw_relations").eq(table.alias_or_name),
51                    exp.column("name", table="rw_columns").neq("_row_id"),
52                    exp.column("name", table="rw_columns").neq("_rw_timestamp"),
53                )
54            )
55        )
56
57        if table.db:
58            sql = sql.where(exp.column("name", table="rw_schemas").eq(table.db))
59
60        self.execute(sql)
61        resp = self.cursor.fetchall()
62        if not resp:
63            raise SQLMeshError(f"Could not get columns for table {table_name}. Table not found.")
64        return {
65            column_name: exp.DataType.build(data_type, dialect=self.dialect, udt=True)
66            for column_name, data_type in resp
67        }
68
69    def _truncate_table(self, table_name: TableName) -> None:
70        return self.execute(exp.Delete(this=exp.to_table(table_name)))