Implement missed token renewal step in Cluster client (#236)

## What is the goal of this PR?

We have fixed the bug where some RPC methods in Cluster client does not perform token renewal step.

## What are the changes implemented in this PR?

- Implement token renewal step for database, session, and transaction RPC calls

Author: Ganeshwara Hananda

typedb/common/rpc/stub.py

@@ -69,7 +69,8 @@ def channel(self) -> Channel:
     def stub(self) -> core_service_proto.TypeDBStub:
         pass
 
-    def resilient_call(self, function: Callable[[], T]) -> T:
+    @staticmethod
+    def resilient_call(function: Callable[[], T]) -> T:
         try:
             # TODO actually implement forced gRPC to reconnected rapidly, which provides resilience
             return function()

typedb/connection/cluster/database_manager.py

@@ -59,7 +59,7 @@ def all(self) -> List[_ClusterDatabase]:
         errors = []
         for address in self._database_mgrs:
             try:
-                res = self._client._stub(address).databases_all(cluster_database_manager_all_req())
+                res = self._client._stub(address).cluster_databases_all(cluster_database_manager_all_req())
                 return [_ClusterDatabase.of(db, self._client) for db in res.databases]
             except TypeDBClientException as e:
                 errors.append("- %s: %s\n" % (address, e))

typedb/connection/cluster/stub.py

@@ -18,18 +18,22 @@
 #   specific language governing permissions and limitations
 #   under the License.
 #
+from typing import Iterator
 from typing import TypeVar, Callable
 
 import typedb_protocol.cluster.cluster_database_pb2 as cluster_database_proto
 import typedb_protocol.cluster.cluster_server_pb2 as cluster_server_proto
 import typedb_protocol.cluster.cluster_service_pb2_grpc as cluster_service_proto
 import typedb_protocol.cluster.cluster_user_pb2 as cluster_user_proto
+import typedb_protocol.common.session_pb2 as session_proto
+import typedb_protocol.common.transaction_pb2 as transaction_proto
+import typedb_protocol.core.core_database_pb2 as core_database_proto
 import typedb_protocol.core.core_service_pb2_grpc as core_service_proto
 from grpc import Channel, RpcError
 
 from typedb.api.connection.credential import TypeDBCredential
-from typedb.common.rpc.stub import TypeDBStub
 from typedb.common.exception import CLUSTER_TOKEN_CREDENTIAL_INVALID, TypeDBClientException, UNABLE_TO_CONNECT
+from typedb.common.rpc.stub import TypeDBStub
 
 T = TypeVar('T')
 
@@ -51,28 +55,55 @@ def __init__(self, channel: Channel, credential: TypeDBCredential):
                 raise e2
 
     def servers_all(self, req: cluster_server_proto.ServerManager.All.Req) -> cluster_server_proto.ServerManager.All.Res:
-        return self.resilient_authenticated_call(lambda: self._cluster_stub.servers_all(req))
-
-    def databases_get(self, req: cluster_database_proto.ClusterDatabaseManager.Get.Req) -> cluster_database_proto.ClusterDatabaseManager.Get.Res:
-        return self.resilient_authenticated_call(lambda: self._cluster_stub.databases_get(req))
-
-    def databases_all(self, req: cluster_database_proto.ClusterDatabaseManager.All.Req) -> cluster_database_proto.ClusterDatabaseManager.All.Res:
-        return self.resilient_authenticated_call(lambda: self._cluster_stub.databases_all(req))
+        return self.may_renew_token(lambda: self._cluster_stub.servers_all(req))
 
     def users_all(self, req: cluster_user_proto.ClusterUserManager.All.Req) -> cluster_user_proto.ClusterUserManager.All.Res:
-        return self.resilient_authenticated_call(lambda: self._cluster_stub.users_all(req))
+        return self.may_renew_token(lambda: self._cluster_stub.users_all(req))
 
     def users_contains(self, req: cluster_user_proto.ClusterUserManager.Contains.Req) -> cluster_user_proto.ClusterUserManager.Contains.Res:
-        return self.resilient_authenticated_call(lambda: self._cluster_stub.users_contains(req))
+        return self.may_renew_token(lambda: self._cluster_stub.users_contains(req))
 
     def users_create(self, req: cluster_user_proto.ClusterUserManager.Create.Req) -> cluster_user_proto.ClusterUserManager.Create.Res:
-        return self.resilient_authenticated_call(lambda: self._cluster_stub.users_create(req))
+        return self.may_renew_token(lambda: self._cluster_stub.users_create(req))
 
     def user_password(self, req: cluster_user_proto.ClusterUser.Delete.Req) -> cluster_user_proto.ClusterUser.Delete.Res:
-        return self.resilient_authenticated_call(lambda: self._cluster_stub.user_password(req))
+        return self.may_renew_token(lambda: self._cluster_stub.user_password(req))
 
     def user_delete(self, req: cluster_user_proto.ClusterUser.Delete.Req) -> cluster_user_proto.ClusterUser.Delete.Res:
-        return self.resilient_authenticated_call(lambda: self._cluster_stub.user_delete(req))
+        return self.may_renew_token(lambda: self._cluster_stub.user_delete(req))
+
+    def cluster_databases_all(self, req: cluster_database_proto.ClusterDatabaseManager.All.Req) -> cluster_database_proto.ClusterDatabaseManager.All.Res:
+        return self.may_renew_token(lambda: self._cluster_stub.databases_all(req))
+
+    def databases_all(self, req: core_database_proto.CoreDatabaseManager.All.Req) -> core_database_proto.CoreDatabaseManager.All.Res:
+        return self.may_renew_token(self.resilient_call(lambda: self.stub().databases_all(req)))
+
+    def databases_get(self, req: cluster_database_proto.ClusterDatabaseManager.Get.Req) -> cluster_database_proto.ClusterDatabaseManager.Get.Res:
+        return self.may_renew_token(lambda: self._cluster_stub.databases_get(req))
+
+    def databases_contains(self, req: core_database_proto.CoreDatabaseManager.Contains.Req) -> core_database_proto.CoreDatabaseManager.Contains.Res:
+        return self.may_renew_token(lambda: super(_ClusterServerStub, self).databases_contains(req))
+
+    def databases_create(self, req: core_database_proto.CoreDatabaseManager.Create.Req) -> core_database_proto.CoreDatabaseManager.Create.Res:
+        return self.may_renew_token(lambda: super(_ClusterServerStub, self).databases_create(req))
+
+    def database_schema(self, req: core_database_proto.CoreDatabase.Schema.Req) -> core_database_proto.CoreDatabase.Schema.Res:
+        return self.may_renew_token(lambda: super(_ClusterServerStub, self).database_schema(req))
+
+    def database_delete(self, req: core_database_proto.CoreDatabase.Delete.Req) -> core_database_proto.CoreDatabase.Delete.Res:
+        return self.may_renew_token(lambda: super(_ClusterServerStub, self).database_delete(req))
+
+    def session_open(self, req: session_proto.Session.Open.Req) -> session_proto.Session.Open.Res:
+        return self.may_renew_token(lambda: super(_ClusterServerStub, self).session_open(req))
+
+    def session_close(self, req: session_proto.Session.Close.Req) -> session_proto.Session.Close.Res:
+        return self.may_renew_token(lambda: super(_ClusterServerStub, self).session_close(req))
+
+    def session_pulse(self, req: session_proto.Session.Pulse.Req) -> session_proto.Session.Pulse.Res:
+        return self.may_renew_token(lambda: super(_ClusterServerStub, self).session_pulse(req))
+
+    def transaction(self, request_iterator: Iterator[transaction_proto.Transaction.Client]) -> Iterator[transaction_proto.Transaction.Server]:
+        return self.may_renew_token(lambda: super(_ClusterServerStub, self).transaction(request_iterator))
 
     def channel(self) -> Channel:
         return self._channel
@@ -83,7 +114,7 @@ def stub(self) -> TypeDBStub:
     def token(self):
         return self._token
 
-    def resilient_authenticated_call(self, function: Callable[[], T]) -> T:
+    def may_renew_token(self, function: Callable[[], T]) -> T:
         try:
             return self.resilient_call(function)
         except TypeDBClientException as e:

typedb/connection/database_manager.py

@@ -21,8 +21,6 @@
 
 from typing import List
 
-from grpc import Channel
-
 from typedb.api.connection.database import DatabaseManager
 from typedb.common.exception import TypeDBClientException, DB_DOES_NOT_EXIST, MISSING_DB_NAME
 from typedb.common.rpc.request_builder import core_database_manager_contains_req, core_database_manager_create_req, \