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

ci/scripts/e2e-source-test.sh

@@ -107,12 +107,15 @@ echo "--- e2e, ci-1cn-1fe, protobuf schema registry"
 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
+python3 e2e_test/schema_registry/pb.py "message_queue:29092" "http://message_queue:8081" "sr_pb_test_more" 1 user_more
+python3 e2e_test/schema_registry/pb.py "message_queue:29092" "http://message_queue:8081" "sr_pb_test_new_type" 1 user_new_type
 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,41 @@
+# Before running this test, seed data into kafka:
+#   python3 e2e_test/schema_registry/pb.py <brokerlist> <schema-registry-url> <topic> <num-records> <pb_type>
+
+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'
+);
+
+# Changing type is not allowed
+statement error
+ALTER SOURCE src_user FORMAT plain ENCODE protobuf(
+    schema.registry = 'http://message_queue:8081',
+    message = 'test.UserNewType'
+);
+
+statement ok
+ALTER SOURCE src_user FORMAT plain ENCODE protobuf(
+    schema.registry = 'http://message_queue:8081',
+    message = 'test.UserMore'
+);
+
+# Dropping columns is not allowed
+statement error
+ALTER SOURCE src_user FORMAT plain ENCODE protobuf(
+    schema.registry = 'http://message_queue:8081',
+    message = 'test.User'
+);
+
+# Changing format/encode is not allowed
+statement error
+ALTER SOURCE src_user FORMAT json ENCODE protobuf(
+    schema.registry = 'http://message_queue:8081',
+    message = 'test.UserMore'
+);

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_more(i):
+    return user_pb2.UserMore(
+        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_new_type(i):
+    return user_pb2.UserNewType(
+        id=i,
+        name="User_{}".format(i),
+        address="Address_{}".format(i),
+        city="City_{}".format(i),
+        gender="MALE" if i % 2 == 0 else "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_type):
     schema_registry_client = SchemaRegistryClient(schema_registry_conf)
     serializer = ProtobufSerializer(
-        user_pb2.User,
+        pb_type,
         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_type>")
         exit(1)
 
     broker_list = sys.argv[1]
     schema_registry_url = sys.argv[2]
     topic = sys.argv[3]
     num_records = int(sys.argv[4])
+    pb_type = sys.argv[5]
+
+    all_pb_types = {
+        'user': (get_user, user_pb2.User),
+        'user_more': (get_user_more, user_pb2.UserMore),
+        'user_new_type': (get_user_new_type, user_pb2.UserNewType),
+    }
+
+    assert pb_type in all_pb_types, f'pb_type must be one of {list(all_pb_types.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_types[pb_type])
     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_type>
 
 # Create a table.
 statement ok

e2e_test/schema_registry/protobuf/user.proto

@@ -17,3 +17,22 @@ enum Gender {
   MALE = 0;
   FEMALE = 1;
 }
+
+message UserMore {
+  int32 id = 1;
+  string name = 2;
+  string address = 3;
+  string city = 4;
+  Gender gender = 5;
+  google.protobuf.SourceContext sc = 6;
+  int32 age = 7;
+}
+
+message UserNewType {
+  int32 id = 1;
+  string name = 2;
+  string address = 3;
+  string city = 4;
+  string gender = 5;
+  google.protobuf.SourceContext sc = 6;
+}

src/frontend/src/catalog/catalog_service.rs

@@ -175,6 +175,8 @@ pub trait CatalogWriter: Send + Sync {
 
     async fn alter_owner(&self, object: Object, owner_id: u32) -> Result<()>;
 
+    async fn alter_source_format_encode(&self, source: PbSource) -> Result<()>;
+
     async fn alter_set_schema(
         &self,
         object: alter_set_schema_request::Object,
@@ -490,6 +492,11 @@ impl CatalogWriter for CatalogWriterImpl {
             .await?;
         self.wait_version(version).await
     }
+
+    async fn alter_source_format_encode(&self, source: PbSource) -> Result<()> {
+        let version = self.meta_client.alter_source_format_encode(source).await?;
+        self.wait_version(version).await
+    }
 }
 
 impl CatalogWriterImpl {

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(),
             )));
         }
         _ => {}