feat(protobuf): support loading protobuf schemas from a directory

ngjaying · ngjaying · commit 60dbe7bd555d · 2026-02-27T10:01:41.000+08:00
The protobuf schema parser now supports directories by passing the directory
path as the schema file. It will parse all '.proto' files within the
directory, allowing for schemas split across multiple files.

Signed-off-by: Jiyong Huang &lt;huangjy@emqx.io&gt;
diff --git a/internal/converter/protobuf/converter.go b/internal/converter/protobuf/converter.go
@@ -16,6 +16,9 @@ package protobuf
 
 import (
 	"fmt"
+	"os"
+	"path/filepath"
+	"strings"
 
 	"github.com/jhump/protoreflect/desc"            //nolint:staticcheck
 	"github.com/jhump/protoreflect/desc/protoparse" //nolint:staticcheck
@@ -44,19 +47,52 @@ func NewConverter(schemaFile string, soFile string, messageName string) (message
 	if soFile != "" {
 		return static.LoadStaticConverter(soFile, messageName)
 	} else {
-		if fds, err := protoParser.ParseFiles(schemaFile); err != nil {
-			return nil, fmt.Errorf("parse schema file %s failed: %s", schemaFile, err)
-		} else {
-			messageDescriptor := fds[0].FindMessage(messageName)
-			if messageDescriptor == nil {
-				return nil, fmt.Errorf("message type %s not found in schema file %s", messageName, schemaFile)
+		protoFiles, err := collectProtoFiles(schemaFile)
+		if err != nil {
+			return nil, fmt.Errorf("collect proto files from %s failed: %s", schemaFile, err)
+		}
+		fds, err := protoParser.ParseFiles(protoFiles...)
+		if err != nil {
+			return nil, fmt.Errorf("parse schema file(s) %s failed: %s", schemaFile, err)
+		}
+		for _, fd := range fds {
+			messageDescriptor := fd.FindMessage(messageName)
+			if messageDescriptor != nil {
+				return &Converter{
+					descriptor: messageDescriptor,
+					fc:         GetFieldConverter(),
+				}, nil
 			}
-			return &Converter{
-				descriptor: messageDescriptor,
-				fc:         GetFieldConverter(),
-			}, nil
 		}
+		return nil, fmt.Errorf("message type %s not found in schema path %s", messageName, schemaFile)
+	}
+}
+
+// collectProtoFiles returns a list of .proto file paths for the given path.
+// If the path is a directory, it returns full paths for directory entries.
+// If it is a single file, it returns the path as-is.
+func collectProtoFiles(path string) ([]string, error) {
+	info, err := os.Stat(path)
+	if err != nil {
+		return nil, err
+	}
+	if !info.IsDir() {
+		return []string{path}, nil
+	}
+	entries, err := os.ReadDir(path)
+	if err != nil {
+		return nil, err
+	}
+	var result []string
+	for _, e := range entries {
+		if !e.IsDir() && strings.HasSuffix(e.Name(), ".proto") {
+			result = append(result, filepath.Join(path, e.Name()))
+		}
+	}
+	if len(result) == 0 {
+		return nil, fmt.Errorf("no .proto files found in directory %s", path)
 	}
+	return result, nil
 }
 
 func (c *Converter) Encode(ctx api.StreamContext, d any) (b []byte, err error) {
diff --git a/internal/converter/protobuf/converter_test.go b/internal/converter/protobuf/converter_test.go
@@ -16,6 +16,7 @@ package protobuf
 
 import (
 	"fmt"
+	"path/filepath"
 	"reflect"
 	"testing"
 
@@ -300,3 +301,53 @@ func TestErr(t *testing.T) {
 	require.True(t, ok)
 	require.Equal(t, errorx.CovnerterErr, errWithCode.Code())
 }
+
+// ---- collectProtoFiles tests ----
+
+func TestCollectProtoFiles_SingleFile(t *testing.T) {
+	result, err := collectProtoFiles("../../schema/test/test1.proto")
+	require.NoError(t, err)
+	assert.Equal(t, []string{"../../schema/test/test1.proto"}, result)
+}
+
+func TestCollectProtoFiles_Directory(t *testing.T) {
+	result, err := collectProtoFiles("../../schema/test/multidir")
+	require.NoError(t, err)
+	assert.Len(t, result, 2)
+	assert.Contains(t, result, filepath.Join("../../schema/test/multidir", "msg_a.proto"))
+	assert.Contains(t, result, filepath.Join("../../schema/test/multidir", "msg_b.proto"))
+}
+
+func TestCollectProtoFiles_EmptyDir(t *testing.T) {
+	emptyDir := t.TempDir()
+	_, err := collectProtoFiles(emptyDir)
+	assert.Error(t, err)
+	assert.Contains(t, err.Error(), "no .proto files found")
+}
+
+func TestCollectProtoFiles_NotExist(t *testing.T) {
+	_, err := collectProtoFiles("../../schema/test/nonexistent")
+	assert.Error(t, err)
+}
+
+// ---- Directory-based NewConverter tests ----
+
+func TestNewConverter_FromDirectory(t *testing.T) {
+	// SensorData is defined in multidir/msg_a.proto
+	c, err := NewConverter("../../schema/test/multidir", "", "SensorData")
+	require.NoError(t, err)
+	require.NotNil(t, c)
+}
+
+func TestNewConverter_FromDirectory_SecondFile(t *testing.T) {
+	// VehicleStatus is defined in multidir/msg_b.proto
+	c, err := NewConverter("../../schema/test/multidir", "", "VehicleStatus")
+	require.NoError(t, err)
+	require.NotNil(t, c)
+}
+
+func TestNewConverter_FromDirectory_NotFound(t *testing.T) {
+	_, err := NewConverter("../../schema/test/multidir", "", "NonExistentMsg")
+	assert.Error(t, err)
+	assert.Contains(t, err.Error(), "not found")
+}
diff --git a/internal/schema/protobuf.go b/internal/schema/protobuf.go
@@ -36,6 +36,19 @@ func (p *PbType) Scan(logger api.Logger, schemaDir string) (map[string]*modules.
 	} else {
 		newSchemas = make(map[string]*modules.Files, len(files))
 		for _, file := range files {
+			// Subdirectory: treat as a single schema ID containing multiple .proto files
+			if file.IsDir() {
+				schemaId := file.Name()
+				ffs, ok := newSchemas[schemaId]
+				if !ok {
+					ffs = &modules.Files{}
+					newSchemas[schemaId] = ffs
+				}
+				// SchemaFile points to the directory itself
+				ffs.SchemaFile = filepath.Join(schemaDir, file.Name())
+				logger.Infof("schema directory %s/%s loaded", schemaDir, schemaId)
+				continue
+			}
 			fileName := filepath.Base(file.Name())
 			ext := filepath.Ext(fileName)
 			schemaId := strings.TrimSuffix(fileName, filepath.Ext(fileName))
@@ -59,15 +72,48 @@ func (p *PbType) Scan(logger api.Logger, schemaDir string) (map[string]*modules.
 }
 
 func (p *PbType) Infer(_ api.Logger, filePath string, messageId string) (ast.StreamFields, error) {
-	if fds, err := protoParser.ParseFiles(filePath); err != nil {
-		return nil, fmt.Errorf("parse schema file %s failed: %s", filePath, err)
-	} else {
-		messageDescriptor := fds[0].FindMessage(messageId)
-		if messageDescriptor == nil {
-			return nil, fmt.Errorf("message type %s not found in schema file %s", messageId, filePath)
+	protoFiles, err := collectProtoFiles(filePath)
+	if err != nil {
+		return nil, fmt.Errorf("collect proto files from %s failed: %s", filePath, err)
+	}
+	fds, err := protoParser.ParseFiles(protoFiles...)
+	if err != nil {
+		return nil, fmt.Errorf("parse schema file(s) %s failed: %s", filePath, err)
+	}
+	for _, fd := range fds {
+		messageDescriptor := fd.FindMessage(messageId)
+		if messageDescriptor != nil {
+			return convertMessage(messageDescriptor)
 		}
-		return convertMessage(messageDescriptor)
 	}
+	return nil, fmt.Errorf("message type %s not found in schema path %s", messageId, filePath)
+}
+
+// collectProtoFiles returns a list of .proto file paths for the given path.
+// If the path is a directory, it returns dir-relative paths (e.g. "multidir/msg_a.proto").
+// If it is a single file, it returns the path as-is.
+func collectProtoFiles(path string) ([]string, error) {
+	info, err := os.Stat(path)
+	if err != nil {
+		return nil, err
+	}
+	if !info.IsDir() {
+		return []string{path}, nil
+	}
+	entries, err := os.ReadDir(path)
+	if err != nil {
+		return nil, err
+	}
+	var result []string
+	for _, e := range entries {
+		if !e.IsDir() && strings.HasSuffix(e.Name(), ".proto") {
+			result = append(result, filepath.Join(path, e.Name()))
+		}
+	}
+	if len(result) == 0 {
+		return nil, fmt.Errorf("no .proto files found in directory %s", path)
+	}
+	return result, nil
 }
 
 func convertMessage(m *desc.MessageDescriptor) (ast.StreamFields, error) {
diff --git a/internal/schema/protobuf_test.go b/internal/schema/protobuf_test.go
@@ -1,6 +1,7 @@
 package schema
 
 import (
+	"path/filepath"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -51,3 +52,92 @@ func TestInferProtobufWithEmbedType(t *testing.T) {
 		t.Errorf("InferProtobuf result is not expected, got %v, expected %v", result, expected)
 	}
 }
+
+// ---- collectProtoFiles tests ----
+
+func TestCollectProtoFiles_SingleFile(t *testing.T) {
+	result, err := collectProtoFiles("test/test1.proto")
+	require.NoError(t, err)
+	assert.Equal(t, []string{"test/test1.proto"}, result)
+}
+
+func TestCollectProtoFiles_Directory(t *testing.T) {
+	result, err := collectProtoFiles("test/multidir")
+	require.NoError(t, err)
+	assert.Len(t, result, 2)
+	assert.Contains(t, result, filepath.Join("test/multidir", "msg_a.proto"))
+	assert.Contains(t, result, filepath.Join("test/multidir", "msg_b.proto"))
+}
+
+func TestCollectProtoFiles_EmptyDir(t *testing.T) {
+	emptyDir := t.TempDir()
+	_, err := collectProtoFiles(emptyDir)
+	assert.Error(t, err)
+	assert.Contains(t, err.Error(), "no .proto files found")
+}
+
+func TestCollectProtoFiles_NotExist(t *testing.T) {
+	_, err := collectProtoFiles("test/nonexistent")
+	assert.Error(t, err)
+}
+
+// ---- Directory-based Scan tests ----
+
+type mockLogger struct{}
+
+func (m *mockLogger) Debug(args ...interface{})                 {}
+func (m *mockLogger) Info(args ...interface{})                  {}
+func (m *mockLogger) Warn(args ...interface{})                  {}
+func (m *mockLogger) Error(args ...interface{})                 {}
+func (m *mockLogger) Debugln(args ...interface{})               {}
+func (m *mockLogger) Infoln(args ...interface{})                {}
+func (m *mockLogger) Warnln(args ...interface{})                {}
+func (m *mockLogger) Errorln(args ...interface{})               {}
+func (m *mockLogger) Debugf(format string, args ...interface{}) {}
+func (m *mockLogger) Infof(format string, args ...interface{})  {}
+func (m *mockLogger) Warnf(format string, args ...interface{})  {}
+func (m *mockLogger) Errorf(format string, args ...interface{}) {}
+
+func TestScan_WithSubdirectory(t *testing.T) {
+	pt := &PbType{}
+	schemas, err := pt.Scan(&mockLogger{}, "test")
+	require.NoError(t, err)
+	// Should find "multidir" as a subdirectory-based schema ID
+	assert.Contains(t, schemas, "multidir")
+	assert.NotEmpty(t, schemas["multidir"].SchemaFile)
+	// Should also find regular .proto files
+	assert.Contains(t, schemas, "test1")
+}
+
+// ---- Directory-based Infer tests ----
+
+func TestInfer_FromDirectory(t *testing.T) {
+	pt := &PbType{}
+	// Infer a message from directory containing multiple proto files
+	result, err := pt.Infer(nil, "test/multidir", "SensorData")
+	require.NoError(t, err)
+	expected := ast.StreamFields{
+		{Name: "temperature", FieldType: &ast.BasicType{Type: ast.FLOAT}},
+		{Name: "humidity", FieldType: &ast.BasicType{Type: ast.BIGINT}},
+	}
+	assert.Equal(t, expected, result)
+}
+
+func TestInfer_FromDirectory_SecondFile(t *testing.T) {
+	pt := &PbType{}
+	// Infer a message defined in the second proto file
+	result, err := pt.Infer(nil, "test/multidir", "VehicleStatus")
+	require.NoError(t, err)
+	// VehicleStatus has: speed (int32), vin (string), battery (BatteryInfo)
+	assert.Len(t, result, 3)
+	assert.Equal(t, "speed", result[0].Name)
+	assert.Equal(t, "vin", result[1].Name)
+	assert.Equal(t, "battery", result[2].Name)
+}
+
+func TestInfer_FromDirectory_MessageNotFound(t *testing.T) {
+	pt := &PbType{}
+	_, err := pt.Infer(nil, "test/multidir", "NonExistentMessage")
+	assert.Error(t, err)
+	assert.Contains(t, err.Error(), "not found")
+}
diff --git a/internal/schema/test/multidir/msg_a.proto b/internal/schema/test/multidir/msg_a.proto
@@ -0,0 +1,6 @@
+syntax = "proto3";
+
+message SensorData {
+  double temperature = 1;
+  int32 humidity = 2;
+}
diff --git a/internal/schema/test/multidir/msg_b.proto b/internal/schema/test/multidir/msg_b.proto
@@ -0,0 +1,12 @@
+syntax = "proto3";
+
+message VehicleStatus {
+  int32 speed = 1;
+  string vin = 2;
+  BatteryInfo battery = 3;
+}
+
+message BatteryInfo {
+  double voltage = 1;
+  int32 level = 2;
+}
