fix: ensure correct datatypes are fethed for RisingWave dialect

Author: MisterWheatley

sqlmesh/core/engine_adapter/risingwave.py

@@ -14,6 +14,7 @@
     CommentCreationTable,
 )
 
+from sqlmesh.utils.errors import SQLMeshError
 
 if t.TYPE_CHECKING:
     from sqlmesh.core._typing import TableName
@@ -32,5 +33,31 @@ class RisingwaveEngineAdapter(PostgresEngineAdapter):
     SUPPORTS_TRANSACTIONS = False
     MAX_IDENTIFIER_LENGTH = None
 
+    def columns(
+        self, table_name: TableName, include_pseudo_columns: bool = False
+    ) -> t.Dict[str, exp.DataType]:
+        """Fetches column names and types for the target_table"""
+        table = exp.to_table(table_name)
+
+        sql = (
+            exp.select("rw_columns.name AS column_name", "rw_columns.data_type AS data_type")
+            .from_("rw_catalog.rw_columns")
+            .join("rw_catalog.rw_relations", on="rw_relations.id=rw_columns.relation_id")
+            .join("rw_catalog.rw_schemas", on="rw_schemas.id=rw_relations.schema_id")
+            .where(exp.column("rw_relations.name", quoted=False).eq(table.alias_or_name))
+        )
+
+        if table.args.get("db"):
+            sql = sql.where(exp.column("rw_schemas.name").eq(table.args["db"].name))
+
+        self.execute(sql)
+        resp = self.cursor.fetchall()
+        if not resp:
+            raise SQLMeshError(f"Could not get columns for table {table_name}. Table not found.")
+        return {
+            column_name: exp.DataType.build(data_type, dialect=self.dialect, udt=True)
+            for column_name, data_type in resp
+        }
+
     def _truncate_table(self, table_name: TableName) -> None:
         return self.execute(exp.Delete(this=exp.to_table(table_name)))

tests/core/engine_adapter/test_risingwave.py

@@ -3,7 +3,7 @@
 from unittest.mock import call
 
 import pytest
-from sqlglot import parse_one
+from sqlglot import parse_one, exp
 from sqlmesh.core.engine_adapter.risingwave import RisingwaveEngineAdapter
 
 pytestmark = [pytest.mark.engine, pytest.mark.risingwave]
@@ -14,6 +14,47 @@ def adapter(make_mocked_engine_adapter):
     adapter = make_mocked_engine_adapter(RisingwaveEngineAdapter)
     return adapter
 
+def test_columns(adapter: t.Callable):
+    adapter.cursor.fetchall.return_value = [
+        ("smallint_col","smallint"),
+        ("int_col","integer"),
+        ("bigint_col","bigint"),
+        ("ts_col","timestamp without time zone"),
+        ("tstz_col","timestamp with time zone"),
+        ("int_array_col","integer[]"),
+        ("vchar_col","character varying"),
+        ("struct_col","struct<nested_col integer>")
+    ]
+    resp = adapter.columns("db.table")
+    assert resp == {
+        "smallint_col" : exp.DataType.build(exp.DataType.Type.SMALLINT,nested=False),
+        "int_col" : exp.DataType.build(exp.DataType.Type.INT,nested=False),
+        "bigint_col" : exp.DataType.build(exp.DataType.Type.BIGINT,nested=False),
+        "ts_col" : exp.DataType.build(exp.DataType.Type.TIMESTAMP,nested=False),
+        "tstz_col" : exp.DataType.build(exp.DataType.Type.TIMESTAMPTZ,nested=False),
+        "int_array_col" : exp.DataType.build(
+                              exp.DataType.Type.ARRAY,
+                              expressions = [exp.DataType.build(exp.DataType.Type.INT,nested=False)],
+                              nested = True
+                          ),
+        "vchar_col": exp.DataType.build(exp.DataType.Type.VARCHAR),
+        "struct_col": exp.DataType.build(
+                          exp.DataType.Type.STRUCT,
+                          expressions = [
+                              exp.ColumnDef(
+                                  this = exp.Identifier(
+                                             this = "nested_col",
+                                             quoted = False
+                                  ),
+                                  kind = exp.DataType.build(
+                                             exp.DataType.Type.INT,
+                                             nested = False
+                                  )
+                              )
+                          ],
+                          nested = True
+                      ),
+    }
 
 def test_create_view(adapter: t.Callable):
     adapter.create_view("db.view", parse_one("SELECT 1"), replace=True)