feat(protobuf): implement service parsing and add RPC function type #31

Adds the capability to parse services within protobuf files and introduces a new FunctionType for RPCs.
phodal · Oct 25, 2024 · 96aecea · 96aecea
1 parent fe1fdfe
commit 96aecea
Show file tree

Hide file tree

Showing 3 changed files with 92 additions and 5 deletions.
diff --git a/chapi-ast-protobuf/src/main/kotlin/chapi/ast/protobuf/ProtobufFullIdentListener.kt b/chapi-ast-protobuf/src/main/kotlin/chapi/ast/protobuf/ProtobufFullIdentListener.kt
@@ -2,10 +2,7 @@ package chapi.ast.protobuf
 
 import chapi.ast.antlr.Protobuf3BaseListener
 import chapi.ast.antlr.Protobuf3Parser
-import chapi.domain.core.CodeContainer
-import chapi.domain.core.CodeDataStruct
-import chapi.domain.core.CodeField
-import chapi.domain.core.DataStructType
+import chapi.domain.core.*
 
 class ProtobufFullIdentListener(var fileName: String) : Protobuf3BaseListener() {
     private var codeContainer: CodeContainer = CodeContainer(FullName = fileName)
@@ -30,6 +27,34 @@ class ProtobufFullIdentListener(var fileName: String) : Protobuf3BaseListener()
         /// todo spike for scene
     }
 
+    override fun enterServiceDef(ctx: Protobuf3Parser.ServiceDefContext?) {
+        val dsName = ctx!!.serviceName().text
+        val functions: List<CodeFunction> = ctx.serviceElement()?.mapNotNull { context ->
+            if (context.rpc() == null) return@mapNotNull null
+
+            val messageTypes = context.rpc().messageType().map { it.text }
+
+            CodeFunction(
+                Name = context.rpc().rpcName().text,
+                FilePath = codeContainer.FullName,
+                Package = codeContainer.PackageName,
+                Type = FunctionType.RPC,
+                Parameters = messageTypes.first()?.let { listOf(CodeProperty(TypeType = it, TypeValue = it)) } ?: emptyList(),
+                ReturnType = messageTypes.last(),
+            )
+        } ?: emptyList()
+
+        val codeDataStruct = CodeDataStruct(
+            NodeName = dsName,
+            Module = codeContainer.PackageName,
+            FilePath = codeContainer.FullName,
+            Package = codeContainer.PackageName,
+            Functions = functions
+        )
+
+        codeContainer.DataStructures += codeDataStruct
+    }
+
     private fun constructMessageDef(ctx: Protobuf3Parser.MessageDefContext): CodeDataStruct {
         val messageName = ctx.messageName().text
         val codeDataStruct = CodeDataStruct(

diff --git a/chapi-ast-protobuf/src/test/kotlin/chapi/ast/protobuf/ProtobufAnalyserTest.kt b/chapi-ast-protobuf/src/test/kotlin/chapi/ast/protobuf/ProtobufAnalyserTest.kt
@@ -1,5 +1,6 @@
 package chapi.ast.protobuf
 
+import chapi.domain.core.FunctionType
 import org.intellij.lang.annotations.Language
 import org.junit.jupiter.api.Test
 import org.junit.jupiter.api.Assertions.*
@@ -119,4 +120,63 @@ message Person {
         assertEquals("city", nestedField2.TypeKey)
         assertEquals("2", nestedField2.TypeValue)
     }
+
+    /// should support for service code
+    @Test
+    fun `should parse valid protobuf code with service and return a CodeContainer`() {
+        // Given
+        @Language("protobuf")
+        val protobufCode = """syntax = "proto3";
+
+package service.user;
+            
+service User {
+    rpc Login (UserBase) returns (UserBase);
+    rpc PhoneCheck (PhoneCheckReq) returns (PhoneCheckReply);
+    rpc UserEdit (UserBase) returns (.google.protobuf.Empty);
+    rpc ListUserInfo (ListUserInfoReq) returns (ListUserInfoReply);
+
+    /////////////like/////////////
+    rpc AddLike (LikeReq) returns (LikeReply);
+    rpc CancelLike (LikeReq) returns (LikeReply);
+    rpc ListUserLike(ListUserLikeReq) returns (ListUserLikeReply);
+    rpc IsLike(IsLikeReq) returns (IsLikeReply);
+
+    ////////////Relation///////////
+    rpc ModifyRelation (ModifyRelationReq) returns (ModifyRelationReply);
+    // 返回UserInfo的列表，分页
+    rpc ListFollowUserInfo (ListRelationUserInfoReq) returns (ListUserInfoReply);
+    rpc ListFanUserInfo (ListRelationUserInfoReq) returns (ListUserInfoReply);
+    rpc ListBlackUserInfo (ListRelationUserInfoReq) returns (ListUserInfoReply);
+    // 仅仅返回全部mid列表，不包含UserInfo
+    rpc ListFollow (ListRelationReq) returns (ListRelationReply);
+    rpc ListBlack (ListRelationReq) returns (ListRelationReply);
+}
+  """
+
+        val filePath = "path/to/file.proto"
+        val analyser = ProtobufAnalyser()
+
+        // When
+        val codeContainer = analyser.analysis(protobufCode, filePath)
+
+        // Then
+        assertNotNull(codeContainer)
+        assertEquals("service.user", codeContainer.PackageName)
+        assertTrue(codeContainer.DataStructures.isNotEmpty())
+
+        val dataStruct = codeContainer.DataStructures.first()
+        assertEquals("User", dataStruct.NodeName)
+
+        val function = dataStruct.Functions
+        assertEquals(14, function.size)
+
+        val firstFunction = function.first()
+        assertEquals("Login", firstFunction.Name)
+        assertEquals("service.user", firstFunction.Package)
+        assertEquals("path/to/file.proto", firstFunction.FilePath)
+        assertEquals(FunctionType.RPC, firstFunction.Type)
+        assertEquals("UserBase", firstFunction.Parameters.first().TypeType)
+        assertEquals("UserBase", firstFunction.ReturnType)
+    }
 }
diff --git a/chapi-domain/src/main/kotlin/chapi/domain/core/CodeFunction.kt b/chapi-domain/src/main/kotlin/chapi/domain/core/CodeFunction.kt
@@ -1,6 +1,5 @@
 package chapi.domain.core
 
-import chapi.domain.core.CodeCall
 import kotlinx.serialization.Serializable
 import kotlinx.serialization.json.JsonElement
 import kotlinx.serialization.json.JsonObject
@@ -12,6 +11,9 @@ enum class FunctionType {
 
     // for Golang block
     Block,
+
+    // for RPC
+    RPC,
 }
 
 @Serializable