refresh_token

Author: DanielePalaia

rabbitmq_amqp_python_client/common.py

@@ -18,6 +18,7 @@ class CommonValues(enum.Enum):
     key = "key"
     queue = "queues"
     bindings = "bindings"
+    path_tokens = "/auth/tokens"
 
 
 class ExchangeType(enum.Enum):

rabbitmq_amqp_python_client/connection.py

@@ -134,25 +134,24 @@ def _open_connections(self, reconnect_handlers: bool = False) -> None:
 
     def _create_connection(self) -> None:
 
-        user = None
-        password = None
-        mechs = None
+        self._user = None
+        self._password = None
+        self._mechs = None
 
         if self._oauth2_options is not None:
-            user = "no"
-            password = self._oauth2_options.token
-            mechs = "PLAIN"
-            print("password, mechs: " + user + " " + password)
+            self._user = "no"
+            self._password = self._oauth2_options.token
+            self._mechs = "PLAIN"
 
         if self._recovery_configuration.active_recovery is False:
             self._conn = BlockingConnection(
                 url=self._addr,
                 urls=self._addrs,
                 oauth2_options=self._oauth2_options,
                 ssl_domain=self._ssl_domain,
-                allowed_mechs=mechs,
-                user=user,
-                password=password,
+                allowed_mechs=self._mechs,
+                user=self._user,
+                password=self._password,
             )
         else:
             self._conn = BlockingConnection(
@@ -161,9 +160,9 @@ def _create_connection(self) -> None:
                 oauth2_options=self._oauth2_options,
                 ssl_domain=self._ssl_domain,
                 on_disconnection_handler=self._on_disconnection,
-                allowed_mechs=mechs,
-                user=user,
-                password=password,
+                allowed_mechs=self._mechs,
+                user=self._user,
+                password=self._password,
             )
 
     def dial(self) -> None:
@@ -383,3 +382,25 @@ def active_producers(self) -> int:
     def active_consumers(self) -> int:
         """Returns the number of active consumers"""
         return len(self._consumers)
+
+    def refresh_token(self, token: str) -> None:
+        """
+        Refresh the oauth token
+
+        Args:
+            token: the oauth token to refresh
+
+        Raises:
+            ValidationCodeException: If oauth is not enabled
+        """
+
+        if self._oauth2_options is None:
+            raise ValidationCodeException("the connection is not oauth enabled")
+
+        # update credentials (for reconnection)
+        self._user = "no"
+        self._password = self._oauth2_options.token
+        self._mechs = "PLAIN"
+
+        management = self.management()
+        management.refresh_token(token)

rabbitmq_amqp_python_client/management.py

@@ -573,3 +573,14 @@ def queue_info(self, name: str) -> QueueInfo:
             message_count=queue_info["message_count"],
             consumer_count=queue_info["consumer_count"],
         )
+
+    def refresh_token(self, token: str) -> None:
+
+        self.request(
+            token.encode(),
+            CommonValues.path_tokens.value,
+            CommonValues.command_put.value,
+            [
+                CommonValues.response_code_204.value,
+            ],
+        )

tests/test_connection.py

@@ -1,5 +1,5 @@
 import time
-from datetime import timedelta
+from datetime import datetime, timedelta
 
 from rabbitmq_amqp_python_client import (
     ConnectionClosed,
@@ -11,6 +11,7 @@
 )
 
 from .http_requests import delete_all_connections
+from .utils import token
 
 
 def on_disconnected():
@@ -44,7 +45,7 @@ def test_connection_ssl(ssl_context) -> None:
     environment.close()
 
 
-def test_connection_auth(environment_auth: Environment) -> None:
+def test_connection_oauth(environment_auth: Environment) -> None:
 
     connection = environment_auth.connection()
     connection.dial()
@@ -54,7 +55,7 @@ def test_connection_auth(environment_auth: Environment) -> None:
     connection.close()
 
 
-def test_connection_auth_with_timeout(environment_auth: Environment) -> None:
+def test_connection_oauth_with_timeout(environment_auth: Environment) -> None:
 
     connection = environment_auth.connection()
     connection.dial()
@@ -65,10 +66,32 @@ def test_connection_auth_with_timeout(environment_auth: Environment) -> None:
     try:
         management = connection.management()
         management.declare_queue(QuorumQueueSpecification(name="test-queue"))
+        management.close()
     except Exception:
         raised = True
 
     assert raised is True
+
+    connection.close()
+
+
+def test_connection_oauth_refresh_token(environment_auth: Environment) -> None:
+
+    connection = environment_auth.connection()
+    connection.dial()
+    # let the token expire
+    time.sleep(1)
+    raised = False
+    # token expired, refresh
+    connection.refresh_token(token(datetime.now() + timedelta(milliseconds=5000)))
+    time.sleep(3)
+    try:
+        management = connection.management()
+        management.declare_queue(QuorumQueueSpecification(name="test-queue"))
+    except Exception:
+        raised = True
+
+    assert raised is False
     connection.close()