From 687dea2e65814606bece8041fe94fb2674d2590b Mon Sep 17 00:00:00 2001
From: Tobias Urdin <tobias.urdin@binero.se>
Date: Thu, 12 Jan 2023 07:19:46 +0000
Subject: [PATCH] Support overriding class for get_rpc_* helper functions

We currently do not support overriding the class being
instantiated in the RPC helper functions, this adds that
support so that projects that define their own classes
that inherit from oslo.messaging can use the helpers.

For example neutron utilizes code from neutron-lib that
has it's own RPCClient implementation that inherits from
oslo.messaging, in order for them to use for example
the get_rpc_client helper they need support to override
the class being returned. The alternative would be to
modify the internal _manual_load variable which seems
counter-productive to extending the API provided to
consumers.

Change-Id: Ie22f2ee47a4ca3f28a71272ee1ffdb88aaeb7758
---

diff --git a/oslo_messaging/rpc/client.py b/oslo_messaging/rpc/client.py
index 8e997e9..b96e30f 100644
--- a/oslo_messaging/rpc/client.py
+++ b/oslo_messaging/rpc/client.py
@@ -550,14 +550,16 @@
         return self.prepare(version=version).can_send_version()
 
 
-def get_rpc_client(transport, target, **kwargs):
+def get_rpc_client(transport, target, client_cls=RPCClient, **kwargs):
     """Construct an RPC client.
 
     :param transport: the messaging transport
     :type transport: Transport
     :param target: the exchange, topic and server to listen on
     :type target: Target
+    :param client_cls: The client class to instantiate
+    :type client_cls: class
     :param **kwargs: The kwargs will be passed down to the
-                     RPCClient constructor
+                     client_cls constructor
     """
-    return RPCClient(transport, target, _manual_load=False, **kwargs)
+    return client_cls(transport, target, _manual_load=False, **kwargs)
diff --git a/oslo_messaging/rpc/server.py b/oslo_messaging/rpc/server.py
index 78557e2..94d4888 100644
--- a/oslo_messaging/rpc/server.py
+++ b/oslo_messaging/rpc/server.py
@@ -200,7 +200,8 @@
 
 
 def get_rpc_server(transport, target, endpoints,
-                   executor=None, serializer=None, access_policy=None):
+                   executor=None, serializer=None, access_policy=None,
+                   server_cls=RPCServer):
     """Construct an RPC server.
 
     :param transport: the messaging transport
@@ -217,10 +218,12 @@
     :param access_policy: an optional access policy.
            Defaults to DefaultRPCAccessPolicy
     :type access_policy: RPCAccessPolicyBase
+    :param server_cls: The server class to instantiate
+    :type server_cls: class
     """
     dispatcher = rpc_dispatcher.RPCDispatcher(endpoints, serializer,
                                               access_policy)
-    return RPCServer(transport, target, dispatcher, executor)
+    return server_cls(transport, target, dispatcher, executor)
 
 
 def expected_exceptions(*exceptions):
diff --git a/oslo_messaging/rpc/transport.py b/oslo_messaging/rpc/transport.py
index 121d617..8f08db5 100644
--- a/oslo_messaging/rpc/transport.py
+++ b/oslo_messaging/rpc/transport.py
@@ -22,7 +22,8 @@
 
 
 def get_rpc_transport(conf, url=None,
-                      allowed_remote_exmods=None):
+                      allowed_remote_exmods=None,
+                      transport_cls=msg_transport.RPCTransport):
     """A factory method for Transport objects for RPCs.
 
     This method should be used to ensure the correct messaging functionality
@@ -43,7 +44,9 @@
                                   transport will deserialize remote exceptions
                                   from
     :type allowed_remote_exmods: list
+    :param transport_cls: the transport class to instantiate
+    :type transport_cls: class
     """
     return msg_transport._get_transport(
         conf, url, allowed_remote_exmods,
-        transport_cls=msg_transport.RPCTransport)
+        transport_cls=transport_cls)
diff --git a/releasenotes/notes/get-rpc-helpers-cls-8911826ac08aef2a.yaml b/releasenotes/notes/get-rpc-helpers-cls-8911826ac08aef2a.yaml
new file mode 100644
index 0000000..6ac9aab
--- /dev/null
+++ b/releasenotes/notes/get-rpc-helpers-cls-8911826ac08aef2a.yaml
@@ -0,0 +1,5 @@
+---
+features:
+  - |
+    The ``get_rpc_transport``, ``get_rpc_server`` and ``get_rpc_client`` helper
+    functions now have support for overriding the class that is instantiated.
