From d40c1932bfcb82f9dc1c33e57fee1ee50ee9a6a5 Mon Sep 17 00:00:00 2001
From: Mathieu Barbin <mathieu.barbin@gmail.com>
Date: Wed, 10 Jan 2024 12:46:59 +0100
Subject: [PATCH] Add helpers for advanced users

---
 lib/grpc-protoc-plugin/grpc_protoc_plugin.mli | 10 ++++++++++
 lib/grpc-protoc/grpc_protoc.ml                | 19 +++++++++++++++----
 lib/grpc-protoc/grpc_protoc.mli               | 12 ++++++++++++
 3 files changed, 37 insertions(+), 4 deletions(-)

diff --git a/lib/grpc-protoc-plugin/grpc_protoc_plugin.mli b/lib/grpc-protoc-plugin/grpc_protoc_plugin.mli
index b549936..7801bab 100644
--- a/lib/grpc-protoc-plugin/grpc_protoc_plugin.mli
+++ b/lib/grpc-protoc-plugin/grpc_protoc_plugin.mli
@@ -90,3 +90,13 @@ module Server_rpc : sig
 end
 
 val handlers : 'a list -> ('a, _) Grpc.Rpc.Handlers.t
+
+(** {1 Advanced users API}
+
+  The following functions are meant for advanced users only. *)
+
+val service_spec :
+  (module S with type Request.t = 'request and type Response.t = 'response) ->
+  Grpc.Rpc.Service_spec.t
+(** This is the function used by this module to extract and build a gRPC service
+    spec from the information generated by [ocaml_protoc_plugin]. *)
diff --git a/lib/grpc-protoc/grpc_protoc.ml b/lib/grpc-protoc/grpc_protoc.ml
index 5ba88aa..3ef7076 100644
--- a/lib/grpc-protoc/grpc_protoc.ml
+++ b/lib/grpc-protoc/grpc_protoc.ml
@@ -7,13 +7,18 @@ let decode (type a) (decode : Pbrt.Decoder.t -> a) buffer =
   let decoder = Pbrt.Decoder.of_string buffer in
   decode decoder
 
+let client_service_spec (rpc : _ Pbrt_services.Client.rpc) =
+  {
+    Grpc.Rpc.Service_spec.package = rpc.package;
+    service_name = rpc.service_name;
+  }
+
 module Client_rpc = struct
   let make (type request response)
       (rpc : (request, _, response, _) Pbrt_services.Client.rpc) ~request_mode
       ~response_mode =
     {
-      Grpc.Rpc.Client_rpc.service_spec =
-        { package = rpc.package; service_name = rpc.service_name };
+      Grpc.Rpc.Client_rpc.service_spec = client_service_spec rpc;
       rpc_name = rpc.rpc_name;
       encode_request = encode rpc.encode_pb_req;
       decode_response = decode rpc.decode_pb_res;
@@ -50,6 +55,12 @@ module Server_rpc = struct
     make rpc ~request_mode:Stream ~response_mode:Stream
 end
 
-let handlers { Pbrt_services.Server.package; service_name; handlers } =
+let server_service_spec
+    { Pbrt_services.Server.package; service_name; handlers = _ } =
+  { Grpc.Rpc.Service_spec.package; service_name }
+
+let handlers
+    ({ Pbrt_services.Server.package = _; service_name = _; handlers } as server)
+    =
   Grpc.Rpc.Handlers.With_service_spec
-    { service_spec = { package; service_name }; handlers }
+    { service_spec = server_service_spec server; handlers }
diff --git a/lib/grpc-protoc/grpc_protoc.mli b/lib/grpc-protoc/grpc_protoc.mli
index 8486af6..9c0c03d 100644
--- a/lib/grpc-protoc/grpc_protoc.mli
+++ b/lib/grpc-protoc/grpc_protoc.mli
@@ -114,3 +114,15 @@ module Server_rpc : sig
 end
 
 val handlers : 'a Pbrt_services.Server.t -> (_, 'a) Grpc.Rpc.Handlers.t
+
+(** {1 Advanced users API}
+
+  The following functions are meant for advanced users only. *)
+
+val client_service_spec : _ Pbrt_services.Client.rpc -> Grpc.Rpc.Service_spec.t
+(** This is the function used by this module to extract and build a gRPC service
+    spec from the information generated by [ocaml_protoc] on the client side. *)
+
+val server_service_spec : _ Pbrt_services.Server.t -> Grpc.Rpc.Service_spec.t
+(** This is the function used by this module to extract and build a gRPC service
+    spec from the information generated by [ocaml_protoc] on the server side. *)