adding closing connections logic

Author: DanielePalaia

rabbitmq_amqp_python_client/connection.py

@@ -37,6 +37,11 @@ def __init__(
         self._on_disconnection_handler = on_disconnection_handler
         self._conf_ssl_context: Optional[SslConfigurationContext] = ssl_context
         self._ssl_domain = None
+        self._connections = []  # type: ignore
+        self._index: int = -1
+
+    def _set_environment_connection_list(self, connections: []):  # type: ignore
+        self._connections = connections
 
     def dial(self) -> None:
         logger.debug("Establishing a connection to the amqp server")
@@ -72,9 +77,11 @@ def management(self) -> Management:
         return self._management
 
     # closes the connection to the AMQP 1.0 server.
+    # This method should be called just from Environment and not from the user
     def _close(self) -> None:
         logger.debug("Closing connection")
         self._conn.close()
+        self._connections.remove(self)
 
     def publisher(self, destination: str) -> Publisher:
         if validate_address(destination) is False:

rabbitmq_amqp_python_client/environment.py

@@ -1,8 +1,12 @@
+# For the moment this is just a Connection pooler to keep compatibility with other clients
+import logging
 from typing import Annotated, Callable, Optional, TypeVar
 
 from .connection import Connection
 from .ssl_configuration import SslConfigurationContext
 
+logger = logging.getLogger(__name__)
+
 MT = TypeVar("MT")
 CB = Annotated[Callable[[MT], None], "Message callback type"]
 
@@ -11,7 +15,7 @@ class Environment:
 
     def __init__(self):  # type: ignore
 
-        self._connections = []
+        self._connections: list[Connection] = []
 
     def connection(
         self,
@@ -28,10 +32,15 @@ def connection(
             ssl_context=ssl_context,
             on_disconnection_handler=on_disconnection_handler,
         )
-
+        logger.debug("Environment: Creating and returning a new connection")
         self._connections.append(connection)
+        connection._set_environment_connection_list(self._connections)
         return connection
 
     def close(self) -> None:
+        logger.debug("Environment: Closing all pending connections")
         for connection in self._connections:
             connection._close()
+
+    def connections(self) -> list[Connection]:
+        return self._connections

rabbitmq_amqp_python_client/qpid/proton/_tracing.py

@@ -32,8 +32,10 @@
 
 import proton
 from proton import Sender as ProtonSender
-from proton.handlers import IncomingMessageHandler as ProtonIncomingMessageHandler
-from proton.handlers import OutgoingMessageHandler as ProtonOutgoingMessageHandler
+from proton.handlers import \
+    IncomingMessageHandler as ProtonIncomingMessageHandler
+from proton.handlers import \
+    OutgoingMessageHandler as ProtonOutgoingMessageHandler
 
 _tracer = None
 _trace_key = proton.symbol("x-opt-qpid-tracestate")

rabbitmq_amqp_python_client/qpid/proton/_transport.py

@@ -139,7 +139,8 @@
 
 if TYPE_CHECKING:
     from ._condition import Condition
-    from ._endpoints import Connection  # would produce circular import
+    from ._endpoints import \
+        Connection  # would produce circular import
 
 
 class TraceAdapter:

tests/test_connection.py

@@ -43,6 +43,34 @@ def test_connection_ssl() -> None:
     environment.close()
 
 
+def test_environment_connections_management() -> None:
+
+    environment = Environment()
+    connection = environment.connection("amqp://guest:guest@localhost:5672/")
+    connection.dial()
+    connection2 = environment.connection("amqp://guest:guest@localhost:5672/")
+    connection2.dial()
+    connection3 = environment.connection("amqp://guest:guest@localhost:5672/")
+    connection3.dial()
+
+    assert len(environment.connections()) == 3
+
+    # this shouldn't happen but we test it anyway
+    connection._close()
+
+    assert len(environment.connections()) == 2
+
+    connection2._close()
+
+    assert len(environment.connections()) == 1
+
+    connection3._close()
+
+    assert len(environment.connections()) == 0
+
+    environment.close()
+
+
 def test_connection_reconnection() -> None:
 
     reconnected = False