ProtoGenerator.cs 8.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213
  1. using TEAMModelOS.SDK.Module.GrpcServer.Extensions;
  2. using TEAMModelOS.SDK.Module.GrpcServer.Model;
  3. using ProtoBuf;
  4. using System;
  5. using System.Collections.Generic;
  6. using System.IO;
  7. using System.Linq;
  8. using System.Text;
  9. using Grpc.Core;
  10. namespace TEAMModelOS.SDK.Module.GrpcServer.Internal
  11. {
  12. public class ProtoGenerator
  13. {
  14. /// <summary>
  15. /// 添加proto
  16. /// </summary>
  17. internal static void AddProto<TEntity>(string entityName)
  18. {
  19. if (!ProtoMethodInfo.Protos.ContainsKey(entityName))
  20. {
  21. ProtoMethodInfo.Protos.TryAdd(entityName, ProtoGenerator.FilterHead(Serializer.GetProto<TEntity>(ProtoBuf.Meta.ProtoSyntax.Proto3)));
  22. }
  23. }
  24. /// <summary>
  25. /// 获取实体对应的proto
  26. /// </summary>
  27. internal static string GetProto(string entityName)
  28. {
  29. var rst = ProtoMethodInfo.Protos.TryGetValue(entityName, out string proto);
  30. return rst ? proto : null;
  31. }
  32. /// <summary>
  33. /// 过滤头部 只保留message部分
  34. /// </summary>
  35. internal static string FilterHead(string proto)
  36. {
  37. var lines = new List<string>();
  38. using (var ms = new MemoryStream(Encoding.UTF8.GetBytes(proto)))
  39. using (var sr = new StreamReader(ms))
  40. {
  41. var readEnable = false;
  42. while (sr.Peek() > 0)
  43. {
  44. var line = sr.ReadLine();
  45. if (GrpcExtensionsOptions.Instance.ProtoMsgStartWithKeywords.Any(q => line.StartsWith(q)))
  46. {
  47. readEnable = true;
  48. }
  49. if (readEnable)
  50. {
  51. lines.Add(line);
  52. }
  53. }
  54. }
  55. return string.Join(Environment.NewLine, lines);
  56. }
  57. /// <summary>
  58. /// 生成grpc的message的proto内容
  59. /// </summary>
  60. private static string GenGrpcMessageProto(string pkgName, List<string> msgProtos)
  61. {
  62. var sb = new StringBuilder();
  63. sb.AppendLine("syntax = \"proto3\";");
  64. if (!string.IsNullOrWhiteSpace(GrpcExtensionsOptions.Instance.ProtoNameSpace))
  65. {
  66. sb.AppendLine("option csharp_namespace = \"" + GrpcExtensionsOptions.Instance.ProtoNameSpace.Trim() + "\";");
  67. }
  68. if (!string.IsNullOrWhiteSpace(pkgName))
  69. {
  70. sb.AppendLine($"package {pkgName.Trim()};");
  71. }
  72. sb.AppendLine();
  73. sb.AppendLine(Environment.NewLine);
  74. var sbMsg = new StringBuilder();
  75. foreach (var proto in msgProtos)
  76. {
  77. sbMsg.Append(proto);
  78. sbMsg.AppendLine(Environment.NewLine);
  79. }
  80. var msg = sbMsg.ToString();
  81. var msgMapProtos = new Dictionary<string, List<string>>();
  82. using (var ms = new MemoryStream(Encoding.UTF8.GetBytes(msg)))
  83. using (var sr = new StreamReader(ms))
  84. {
  85. bool readEnable = false;
  86. var msgName = "";
  87. var lines = new List<string>();
  88. while (sr.Peek() > 0)
  89. {
  90. var line = sr.ReadLine();
  91. if (GrpcExtensionsOptions.Instance.ProtoMsgStartWithKeywords.Any(q => line.StartsWith(q)))
  92. {
  93. readEnable = true;
  94. msgName = line.Split(new[] { " " }, StringSplitOptions.RemoveEmptyEntries)[1];
  95. lines = new List<string>();
  96. }
  97. if (readEnable)
  98. {
  99. lines.Add(line);
  100. }
  101. if (line.StartsWith("}"))
  102. {
  103. readEnable = false;
  104. if (!msgMapProtos.ContainsKey(msgName))
  105. {
  106. msgMapProtos.Add(msgName, lines.Select(q => q).ToList());
  107. }
  108. }
  109. }
  110. }
  111. msg = string.Join(Environment.NewLine + Environment.NewLine, msgMapProtos.Select(q => string.Join(Environment.NewLine, q.Value)));
  112. sb.Append(msg);
  113. return sb.ToString();
  114. }
  115. /// <summary>
  116. /// 生成grpc的service的proto内容
  117. /// </summary>
  118. private static string GenGrpcServiceProto(string msgProtoFile, string pkgName, string srvName, List<ProtoMethodInfo> methodInfo)
  119. {
  120. var sb = new StringBuilder();
  121. sb.AppendLine("syntax = \"proto3\";");
  122. if (!string.IsNullOrWhiteSpace(GrpcExtensionsOptions.Instance.ProtoNameSpace))
  123. {
  124. sb.AppendLine("option csharp_namespace = \"" + GrpcExtensionsOptions.Instance.ProtoNameSpace.Trim() + "\";");
  125. }
  126. if (!string.IsNullOrWhiteSpace(pkgName))
  127. {
  128. sb.AppendLine($"package {pkgName.Trim()};");
  129. }
  130. sb.AppendLine(string.Format("import \"{0}\";", msgProtoFile));
  131. sb.AppendLine(Environment.NewLine);
  132. sb.AppendLine("service " + srvName + " {");
  133. var template = @" rpc {0}({1}) returns({2})";
  134. methodInfo.ForEach(q => {
  135. var requestName = q.RequestName;
  136. var responseName = q.ResponseName;
  137. switch (q.MethodType)
  138. {
  139. case MethodType.Unary:
  140. break;
  141. case MethodType.ClientStreaming:
  142. requestName = "stream " + requestName;
  143. break;
  144. case MethodType.ServerStreaming:
  145. responseName = "stream " + responseName;
  146. break;
  147. case MethodType.DuplexStreaming:
  148. requestName = "stream " + requestName;
  149. responseName = "stream " + responseName;
  150. break;
  151. }
  152. sb.AppendLine(Environment.NewLine + string.Format(template, q.MethodName, requestName, responseName) + ";");
  153. });
  154. sb.AppendLine("}");
  155. return sb.ToString();
  156. }
  157. /// <summary>
  158. /// 生成proto文件
  159. /// </summary>
  160. public static void Gen(string dir)
  161. {
  162. if (ProtoInfo.Methods == null || ProtoInfo.Methods.Count == 0) return;
  163. if (!Directory.Exists(dir))
  164. {
  165. Directory.CreateDirectory(dir);
  166. }
  167. foreach (var grp in ProtoInfo.Methods.GroupBy(q => q.ServiceName))
  168. {
  169. var arr = grp.Key.Split(new[] { '.' }, StringSplitOptions.RemoveEmptyEntries);
  170. if (arr.Length > 2) continue;
  171. var pkg = arr.Length == 2 ? arr[0] : null;
  172. var srv = arr.Length == 2 ? arr[1] : arr[0];
  173. #region message
  174. var protoName = grp.Key;
  175. var msgProto = $"{protoName}.message.proto";
  176. var protoPath = Path.Combine(dir, msgProto);
  177. if (File.Exists(protoPath))
  178. {
  179. File.Delete(protoPath);
  180. }
  181. var msgProtos = new List<string>();
  182. var rqNames = grp.ToList().Select(q => q.RequestName).ToList();
  183. var rsNames = grp.ToList().Select(q => q.ResponseName).ToList();
  184. var msgNames = rqNames.Union(rsNames).Distinct().ToList();
  185. foreach (var n in msgNames)
  186. {
  187. msgProtos.Add(GetProto(n));
  188. }
  189. var msgProtoContent = GenGrpcMessageProto(pkg, msgProtos);
  190. File.AppendAllText(protoPath, msgProtoContent);
  191. #endregion
  192. #region service
  193. var srvProto = $"{protoName}.service.proto";
  194. protoPath = Path.Combine(dir, srvProto);
  195. if (File.Exists(protoPath))
  196. {
  197. File.Delete(protoPath);
  198. }
  199. var methodInfos = grp.ToList();
  200. var srvProtoContent = GenGrpcServiceProto(msgProto, pkg, srv, methodInfos);
  201. File.AppendAllText(protoPath, srvProtoContent);
  202. #endregion
  203. }
  204. }
  205. }
  206. }