Feat: add risingwave engine adapter support for sqlmesh. (#3436)

Author: lin0303-siyuan

.circleci/continue_config.yml

@@ -298,6 +298,7 @@ workflows:
                 - spark
                 - clickhouse
                 - clickhouse-cluster
+                - risingwave
       - engine_tests_cloud:
           name: cloud_engine_<< matrix.engine >>
           context:

.circleci/wait-for-db.sh

@@ -55,6 +55,10 @@ trino_ready() {
     docker compose -f tests/core/engine_adapter/integration/docker/compose.trino.yaml exec trino /bin/bash -c '/usr/lib/trino/bin/health-check'
 }
 
+risingwave_ready() {
+    probe_port 4566
+}
+
 echo "Waiting for $ENGINE to be ready..."
 
 READINESS_FUNC="${ENGINE}_ready"

Makefile

@@ -187,6 +187,9 @@ spark-test: engine-spark-up
 trino-test: engine-trino-up
 	pytest -n auto -x -m "trino or trino_iceberg or trino_delta or trino_nessie" --retries 3 --junitxml=test-results/junit-trino.xml
 
+risingwave-test: engine-risingwave-up
+	pytest -n auto -x -m "risingwave" --retries 3 --junitxml=test-results/junit-risingwave.xml
+
 #################
 # Cloud Engines #
 #################

docs/integrations/engines/risingwave.md

@@ -0,0 +1,53 @@
+# RisingWave
+
+This page provides information about how to use SQLMesh with the RisingWave streaming database engine.
+
+## Connection options
+
+| Option         | Description                                                  | Type   | Required |
+|----------------|--------------------------------------------------------------|:------:|:--------:|
+| `type`         | Engine type name - must be `risingwave`                           | string | Y        |
+| `host`         | The hostname of the RisingWave engine                            | string | Y        |
+| `user`         | The username to use for authentication with the RisingWave engine | string | Y        |
+| `password`     | The password to use for authentication with the RisingWave engine | string | N        |
+| `port`         | The port number of the RisingWave engine server                          | int    | Y        |
+| `database`        | The name of the database instance to connect to                                 | string | Y        |
+| `role`            | The role to use for authentication with the RisingWave server                     | string | N        |
+| `sslmode`         | The security of the connection to the RisingWave server                           | string | N        |
+
+## Usage
+RisingWave engine has some different features as streaming database. You can create a resource that RisingWave can read data from with `CREATE SOURCE`. You can also create an external target where you can send data processed in RisingWave with `CREATE SINK`.
+
+To use this in SQLMesh, you can refer to optional pre-statements and post-statements as [SQL models doc](https://sqlmesh.readthedocs.io/en/stable/concepts/models/sql_models/) here specify.
+
+Below is an example of creating sink in SQLMesh models as post-statement.
+
+```sql
+MODEL (
+    name sqlmesh_example.view_model,
+    kind VIEW (
+      materialized true
+    )
+);
+
+SELECT
+  item_id,
+  COUNT(DISTINCT id) AS num_orders,
+FROM
+  sqlmesh_example.incremental_model
+GROUP BY item_id;
+
+CREATE
+  SINK IF NOT EXISTS kafka_sink
+FROM
+  @this_model
+WITH (
+  connector='kafka',
+  "properties.bootstrap.server"='localhost:9092',
+  topic='test1',
+)
+FORMAT PLAIN
+ENCODE JSON (force_append_only=true);
+```
+
+here `@this_model` macro is used to represent "sqlmesh_example.view_model" model.

setup.py

@@ -160,6 +160,9 @@
             "sse-starlette>=0.2.2",
             "pyarrow",
         ],
+        "risingwave": [
+            "psycopg2",
+        ],
     },
     classifiers=[
         "Intended Audience :: Developers",

sqlmesh/core/config/connection.py

@@ -1770,6 +1770,52 @@ def get_catalog(self) -> t.Optional[str]:
         return self.catalog_name
 
 
+class RisingwaveConnectionConfig(ConnectionConfig):
+    host: str
+    user: str
+    password: t.Optional[str] = None
+    port: int
+    database: str
+    role: t.Optional[str] = None
+    sslmode: t.Optional[str] = None
+
+    concurrent_tasks: int = 4
+    register_comments: bool = True
+    pre_ping: bool = True
+
+    type_: t.Literal["risingwave"] = Field(alias="type", default="risingwave")
+
+    @property
+    def _connection_kwargs_keys(self) -> t.Set[str]:
+        return {
+            "host",
+            "user",
+            "password",
+            "port",
+            "database",
+            "role",
+            "sslmode",
+        }
+
+    @property
+    def _engine_adapter(self) -> t.Type[EngineAdapter]:
+        return engine_adapter.RisingwaveEngineAdapter
+
+    @property
+    def _connection_factory(self) -> t.Callable:
+        from psycopg2 import connect
+
+        return connect
+
+    @property
+    def _cursor_init(self) -> t.Optional[t.Callable[[t.Any], None]]:
+        def init(cursor: t.Any) -> None:
+            sql = "SET RW_IMPLICIT_FLUSH TO true;"
+            cursor.execute(sql)
+
+        return init
+
+
 CONNECTION_CONFIG_TO_TYPE = {
     # Map all subclasses of ConnectionConfig to the value of their `type_` field.
     tpe.all_field_infos()["type_"].default: tpe

sqlmesh/core/engine_adapter/__init__.py

@@ -18,6 +18,7 @@
 from sqlmesh.core.engine_adapter.spark import SparkEngineAdapter
 from sqlmesh.core.engine_adapter.trino import TrinoEngineAdapter
 from sqlmesh.core.engine_adapter.athena import AthenaEngineAdapter
+from sqlmesh.core.engine_adapter.risingwave import RisingwaveEngineAdapter
 
 DIALECT_TO_ENGINE_ADAPTER = {
     "hive": SparkEngineAdapter,
@@ -33,6 +34,7 @@
     "mssql": MSSQLEngineAdapter,
     "trino": TrinoEngineAdapter,
     "athena": AthenaEngineAdapter,
+    "risingwave": RisingwaveEngineAdapter,
 }
 
 DIALECT_ALIASES = {

sqlmesh/core/engine_adapter/risingwave.py

@@ -0,0 +1,35 @@
+from __future__ import annotations
+
+import logging
+import typing as t
+
+
+from sqlglot import exp
+
+from sqlmesh.core.engine_adapter.postgres import PostgresEngineAdapter
+from sqlmesh.core.engine_adapter.shared import (
+    set_catalog,
+    CatalogSupport,
+    CommentCreationView,
+    CommentCreationTable,
+)
+
+
+if t.TYPE_CHECKING:
+    from sqlmesh.core._typing import TableName
+
+logger = logging.getLogger(__name__)
+
+
+@set_catalog()
+class RisingwaveEngineAdapter(PostgresEngineAdapter):
+    DIALECT = "risingwave"
+    DEFAULT_BATCH_SIZE = 400
+    CATALOG_SUPPORT = CatalogSupport.SINGLE_CATALOG_ONLY
+    COMMENT_CREATION_TABLE = CommentCreationTable.COMMENT_COMMAND_ONLY
+    COMMENT_CREATION_VIEW = CommentCreationView.UNSUPPORTED
+    SUPPORTS_MATERIALIZED_VIEWS = True
+    SUPPORTS_TRANSACTIONS = False
+
+    def _truncate_table(self, table_name: TableName) -> None:
+        return self.execute(exp.Delete(this=exp.to_table(table_name)))

tests/core/engine_adapter/integration/__init__.py

@@ -160,6 +160,9 @@ def supports_merge(self) -> bool:
         if self.dialect == "athena":
             return "hive" not in self.mark
 
+        if self.dialect == "risingwave":
+            return False
+
         return True
 
     @property
@@ -348,6 +351,20 @@ def get_table_comment(
             """
         elif self.dialect == "clickhouse":
             query = f"SELECT name, comment FROM system.tables WHERE database = '{schema_name}' AND name = '{table_name}'"
+        elif self.dialect == "risingwave":
+            query = f"""
+                SELECT
+                    c.relname,
+                    d.description
+                FROM pg_class c
+                INNER JOIN pg_description d ON c.oid = d.objoid AND d.objsubid = 0
+                INNER JOIN pg_namespace n ON c.relnamespace = n.oid
+                WHERE 
+                    c.relname = '{table_name}'
+                    AND n.nspname= '{schema_name}'
+                    AND c.relkind = '{'v' if table_kind == "VIEW" else 'r'}'
+                ;
+            """
 
         result = self.engine_adapter.fetchall(query)
 
@@ -439,6 +456,24 @@ def get_column_comments(
                     schema_name = '{schema_name}'
                     AND table_name = '{table_name}'
             """
+        elif self.dialect == "risingwave":
+            query = f"""
+                SELECT
+                    a.attname AS column_name, d.description
+                FROM
+                    pg_class c
+                INNER JOIN pg_namespace n ON c.relnamespace = n.oid
+                INNER JOIN pg_attribute a ON c.oid = a.attrelid
+                INNER JOIN pg_description d
+                ON 
+                    a.attnum = d.objsubid
+                    AND d.objoid = c.oid
+                WHERE
+                    n.nspname = '{schema_name}'
+                    AND c.relname = '{table_name}' 
+                    AND c.relkind = '{'v' if table_kind == "VIEW" else 'r'}'
+                ;
+            """
 
         result = self.engine_adapter.fetchall(query)
 

tests/core/engine_adapter/integration/config.yaml

@@ -100,6 +100,13 @@ gateways:
       cluster: cluster1
     state_connection:
       type: duckdb
+  inttest_risingwave:
+    connection:
+      type: risingwave
+      user: root
+      database: dev
+      host: {{ env_var('DOCKER_HOSTNAME', 'localhost') }}
+      port: 4566
 
 
   # Cloud databases