feat: implement ALTER SOURCE xx FORMAT xx ENCODE xx (...)

ci/scripts/e2e-source-test.sh

@@ -5,6 +5,15 @@ set -euo pipefail
 
 source ci/scripts/common.sh
 
+# Arguments:
+#   $1: subject name
+#   $2: schema file path
+function register_schema_registry() {
+    curl -X POST http://message_queue:8081/subjects/$1/versions \
+        -H ‘Content-Type: application/vnd.schemaregistry.v1+json’ \
+        --data-binary @<(jq -n --arg schema “$(cat $2)” ‘{schemaType: “PROTOBUF”, schema: $schema}’)
+}
+
 # prepare environment
 export CONNECTOR_LIBS_PATH="./connector-node/libs"
 
@@ -115,12 +124,13 @@ export RISINGWAVE_CI=true
 RUST_LOG="info,risingwave_stream=info,risingwave_batch=info,risingwave_storage=info" \
 cargo make ci-start ci-1cn-1fe
 python3 -m pip install requests protobuf confluent-kafka
-python3 e2e_test/schema_registry/pb.py "message_queue:29092" "http://message_queue:8081" "sr_pb_test" 20
+python3 e2e_test/schema_registry/pb.py "message_queue:29092" "http://message_queue:8081" "sr_pb_test" 20 user
 echo "make sure google/protobuf/source_context.proto is NOT in schema registry"
 curl --silent 'http://message_queue:8081/subjects'; echo
 # curl --silent --head -X GET 'http://message_queue:8081/subjects/google%2Fprotobuf%2Fsource_context.proto/versions' | grep 404
 curl --silent 'http://message_queue:8081/subjects' | grep -v 'google/protobuf/source_context.proto'
 sqllogictest -p 4566 -d dev './e2e_test/schema_registry/pb.slt'
+sqllogictest -p 4566 -d dev './e2e_test/schema_registry/alter_sr.slt'
 
 echo "--- Kill cluster"
 cargo make ci-kill

e2e_test/schema_registry/alter_sr.slt

@@ -0,0 +1,74 @@
+# Before running this test, seed data into kafka:
+#   python3 e2e_test/schema_registry/pb.py <brokerlist> <schema-registry-url> <topic> <num-records>
+
+statement ok
+CREATE SOURCE src_user WITH (
+    connector = 'kafka',
+    topic = 'sr_pb_test',
+    properties.bootstrap.server = 'message_queue:29092',
+    scan.startup.mode = 'earliest'
+)
+FORMAT PLAIN ENCODE PROTOBUF(
+    schema.registry = 'http://message_queue:8081',
+    message = 'test.User'
+);
+
+statement ok
+CREATE MATERIALIZED VIEW mv_user AS SELECT * FROM src_user;
+
+# Changing type is not allowed
+statement error Feature is not yet implemented: this altering statement will drop columns, which is not supported yet: \(city: character varying\)
+ALTER SOURCE src_user FORMAT PLAIN ENCODE PROTOBUF(
+    schema.registry = 'http://message_queue:8081',
+    message = 'test.UserWithNewType'
+);
+
+# Changing format/encode is not allowed
+statement error Feature is not yet implemented: the original definition is FORMAT Plain ENCODE Protobuf, and altering them is not supported yet
+ALTER SOURCE src_user FORMAT NATIVE ENCODE PROTOBUF(
+    schema.registry = 'http://message_queue:8081',
+    message = 'test.User'
+);
+
+statement ok
+ALTER SOURCE src_user FORMAT PLAIN ENCODE PROTOBUF(
+    schema.registry = 'http://message_queue:8081',
+    message = 'test.UserWithMoreFields'
+);
+
+# Dropping columns is not allowed
+statement error Feature is not yet implemented: this altering statement will drop columns, which is not supported yet: \(age: integer\)
+ALTER SOURCE src_user FORMAT PLAIN ENCODE PROTOBUF(
+    schema.registry = 'http://message_queue:8081',
+    message = 'test.User'
+);
+
+statement ok
+CREATE MATERIALIZED VIEW mv_more_fields AS SELECT * FROM src_user;
+
+system ok
+python3 e2e_test/schema_registry/pb.py "message_queue:29092" "http://message_queue:8081" "sr_pb_test" 5 user_with_more_fields
+
+sleep 10s
+
+query I
+SELECT COUNT(*) FROM mv_user;
+----
+25
+
+statement error
+SELECT SUM(age) FROM mv_user;
+
+query III
+SELECT COUNT(*), MAX(age), MIN(age) FROM mv_more_fields;
+----
+25 4 0
+
+statement ok
+DROP MATERIALIZED VIEW mv_user;
+
+statement ok
+DROP MATERIALIZED VIEW mv_more_fields;
+
+statement ok
+DROP SOURCE src_user;

e2e_test/schema_registry/pb.py

@@ -25,18 +25,38 @@ def get_user(i):
         sc=SourceContext(file_name="source/context_{:03}.proto".format(i)),
     )
 
+def get_user_with_more_fields(i):
+    return user_pb2.UserWithMoreFields(
+        id=i,
+        name="User_{}".format(i),
+        address="Address_{}".format(i),
+        city="City_{}".format(i),
+        gender=user_pb2.MALE if i % 2 == 0 else user_pb2.FEMALE,
+        sc=SourceContext(file_name="source/context_{:03}.proto".format(i)),
+        age=i,
+    )
+
+def get_user_with_new_type(i):
+    return user_pb2.UserWithNewType(
+        id=i,
+        name="User_{}".format(i),
+        address="Address_{}".format(i),
+        city=i,
+        gender=user_pb2.MALE if i % 2 == 0 else user_pb2.FEMALE,
+        sc=SourceContext(file_name="source/context_{:03}.proto".format(i)),
+    )
 
-def send_to_kafka(producer_conf, schema_registry_conf, topic, num_records):
+def send_to_kafka(producer_conf, schema_registry_conf, topic, num_records, get_user_fn, pb_message):
     schema_registry_client = SchemaRegistryClient(schema_registry_conf)
     serializer = ProtobufSerializer(
-        user_pb2.User,
+        pb_message,
         schema_registry_client,
         {"use.deprecated.format": False, 'skip.known.types': True},
     )
 
     producer = Producer(producer_conf)
     for i in range(num_records):
-        user = get_user(i)
+        user = get_user_fn(i)
 
         producer.produce(
             topic=topic,
@@ -49,20 +69,29 @@ def send_to_kafka(producer_conf, schema_registry_conf, topic, num_records):
 
 
 if __name__ == "__main__":
-    if len(sys.argv) < 4:
-        print("pb.py <brokerlist> <schema-registry-url> <topic> <num-records>")
+    if len(sys.argv) < 5:
+        print("pb.py <brokerlist> <schema-registry-url> <topic> <num-records> <pb_message>")
         exit(1)
 
     broker_list = sys.argv[1]
     schema_registry_url = sys.argv[2]
     topic = sys.argv[3]
     num_records = int(sys.argv[4])
+    pb_message = sys.argv[5]
+
+    all_pb_messages = {
+        'user': (get_user, user_pb2.User),
+        'user_with_more_fields': (get_user_with_more_fields, user_pb2.UserWithMoreFields),
+        'user_with_new_type': (get_user_with_new_type, user_pb2.UserWithNewType),
+    }
+
+    assert pb_message in all_pb_messages, f'pb_message must be one of {list(all_pb_messages.keys())}'
 
     schema_registry_conf = {"url": schema_registry_url}
     producer_conf = {"bootstrap.servers": broker_list}
 
     try:
-        send_to_kafka(producer_conf, schema_registry_conf, topic, num_records)
+        send_to_kafka(producer_conf, schema_registry_conf, topic, num_records, *all_pb_messages[pb_message])
     except Exception as e:
         print("Send Protobuf data to schema registry and kafka failed {}", e)
         exit(1)

e2e_test/schema_registry/pb.slt

@@ -1,5 +1,5 @@
 # Before running this test, seed data into kafka:
-#   python3 e2e_test/schema_registry/pb.py
+#   python3 e2e_test/schema_registry/pb.py <brokerlist> <schema-registry-url> <topic> <num-records> <pb_message>
 
 # Create a table.
 statement ok

e2e_test/schema_registry/protobuf/user.proto

@@ -17,3 +17,22 @@ enum Gender {
   MALE = 0;
   FEMALE = 1;
 }
+
+message UserWithMoreFields {
+  int32 id = 1;
+  string name = 2;
+  string address = 3;
+  string city = 4;
+  Gender gender = 5;
+  google.protobuf.SourceContext sc = 6;
+  int32 age = 7; // new field here
+}
+
+message UserWithNewType {
+  int32 id = 1;
+  string name = 2;
+  string address = 3;
+  int32 city = 4; // change the type from string to int32
+  Gender gender = 5;
+  google.protobuf.SourceContext sc = 6;
+}

e2e_test/source/basic/kafka.slt

@@ -680,10 +680,10 @@ select count(*) from s16
 ----
 0
 
-statement error Feature is not yet implemented: Alter source with schema registry
+statement error Not supported: alter source with schema registry
 alter source s18 add column v10 int;
 
-statement error Feature is not yet implemented: Alter source with schema registry
+statement error Not supported: alter source with schema registry
 alter source s17 add column v10 int;
 
 query III rowsort

src/connector/src/parser/protobuf/parser.rs

@@ -525,19 +525,24 @@ fn protobuf_type_mapping(
     Ok(t)
 }
 
+/// Reference: <https://docs.confluent.io/platform/current/schema-registry/fundamentals/serdes-develop/index.html#wire-format>
+/// Wire format for Confluent pb header is:
+/// | 0          | 1-4        | 5-x             | x+1-end
+/// | magic-byte | schema-id  | message-indexes | protobuf-payload
 pub(crate) fn resolve_pb_header(payload: &[u8]) -> Result<&[u8]> {
     // there's a message index array at the front of payload
     // if it is the first message in proto def, the array is just and `0`
     // TODO: support parsing more complex index array
     let (_, remained) = extract_schema_id(payload)?;
+    // The message indexes are encoded as int using variable-length zig-zag encoding,
+    // prefixed by the length of the array.
+    // Note that if the first byte is 0, it is equivalent to (1, 0) as an optimization.
     match remained.first() {
         Some(0) => Ok(&remained[1..]),
-        Some(i) => {
-            Err(RwError::from(ProtocolError(format!("The payload message must be the first message in protobuf schema def, but the message index is {}", i))))
-        }
-        None => {
-            Err(RwError::from(ProtocolError("The proto payload is empty".to_owned())))
-        }
+        Some(i) => Ok(&remained[(*i as usize)..]),
+        None => Err(RwError::from(ProtocolError(
+            "The proto payload is empty".to_owned(),
+        ))),
     }
 }
 

src/frontend/src/catalog/catalog_service.rs

@@ -176,6 +176,8 @@ pub trait CatalogWriter: Send + Sync {
 
     async fn alter_owner(&self, object: Object, owner_id: u32) -> Result<()>;
 
+    async fn alter_source_with_sr(&self, source: PbSource) -> Result<()>;
+
     async fn alter_parallelism(&self, table_id: u32, parallelism: PbTableParallelism)
         -> Result<()>;
 
@@ -495,6 +497,11 @@ impl CatalogWriter for CatalogWriterImpl {
         self.wait_version(version).await
     }
 
+    async fn alter_source_with_sr(&self, source: PbSource) -> Result<()> {
+        let version = self.meta_client.alter_source_with_sr(source).await?;
+        self.wait_version(version).await
+    }
+
     async fn alter_parallelism(
         &self,
         table_id: u32,

src/frontend/src/handler/alter_source_column.rs

@@ -14,7 +14,6 @@
 
 use itertools::Itertools;
 use pgwire::pg_response::{PgResponse, StatementType};
-use risingwave_common::bail_not_implemented;
 use risingwave_common::catalog::ColumnId;
 use risingwave_common::error::{ErrorCode, Result, RwError};
 use risingwave_connector::source::{extract_source_struct, SourceEncode, SourceStruct};
@@ -63,15 +62,23 @@ pub async fn handle_alter_source_column(
     let SourceStruct { encode, .. } = extract_source_struct(&catalog.info)?;
     match encode {
         SourceEncode::Avro | SourceEncode::Protobuf => {
-            bail_not_implemented!("Alter source with schema registry")
+            return Err(ErrorCode::NotSupported(
+                "alter source with schema registry".to_string(),
+                "try `ALTER SOURCE .. FORMAT .. ENCODE .. (...)` instead".to_string(),
+            )
+            .into());
         }
         SourceEncode::Json if catalog.info.use_schema_registry => {
-            bail_not_implemented!("Alter source with schema registry")
+            return Err(ErrorCode::NotSupported(
+                "alter source with schema registry".to_string(),
+                "try `ALTER SOURCE .. FORMAT .. ENCODE .. (...)` instead".to_string(),
+            )
+            .into());
         }
         SourceEncode::Invalid | SourceEncode::Native => {
             return Err(RwError::from(ErrorCode::NotSupported(
-                format!("Alter source with encode {:?}", encode),
-                "Alter source with encode JSON | BYTES | CSV".into(),
+                format!("alter source with encode {:?}", encode),
+                "alter source with encode JSON | BYTES | CSV".into(),
             )));
         }
         _ => {}