handle multiple nested gateways in config

src/jobflow_remote/config/base.py

@@ -263,12 +263,57 @@ def cli_info(self) -> dict:
         )
 
 
+class ConnectionData(BaseModel):
+    """
+    The representation of a fabric connection.
+    Mainly used in case of nested gateways.
+    """
+
+    host: str = Field(description="The host to which to connect")
+    user: Optional[str] = Field(None, description="Login username")
+    port: Optional[int] = Field(None, description="Port number")
+    password: Optional[str] = Field(None, description="Login password")
+    key_filename: Optional[Union[str, list[str]]] = Field(
+        None,
+        description="The filename, or list of filenames, of optional private key(s) "
+        "and/or certs to try for authentication",
+    )
+    passphrase: Optional[str] = Field(
+        None, description="Passphrase used for decrypting private keys"
+    )
+    gateway: Optional[Union[str, "ConnectionData"]] = Field(
+        None, description="A shell command string to use as a proxy or gateway"
+    )
+    connect_kwargs: Optional[dict] = Field(
+        None,
+        description="Other keyword arguments passed to paramiko.client.SSHClient.connect",
+    )
+
+    def get_connect_kwargs(self) -> dict:
+        """
+        Return the fully filled connect_kwargs for Fabric.
+
+        Returns
+        -------
+        The RemoteHost.
+        """
+        connect_kwargs = dict(self.connect_kwargs) if self.connect_kwargs else {}
+        if self.password:
+            connect_kwargs["password"] = self.password
+        if self.key_filename:
+            connect_kwargs["key_filename"] = self.key_filename
+        if self.passphrase:
+            connect_kwargs["passphrase"] = self.passphrase
+
+        return connect_kwargs
+
+
 class RemoteWorker(WorkerBase):
     """
     Worker representing a remote host reached through an SSH connection.
 
-    Uses a Fabric Connection. Check Fabric documentation for more datails on the
-    options defininf a Connection.
+    Uses a Fabric Connection. Check Fabric documentation for more details on the
+    options defining a Connection.
     """
 
     type: Literal["remote"] = Field(
@@ -286,7 +331,7 @@ class RemoteWorker(WorkerBase):
     passphrase: Optional[str] = Field(
         None, description="Passphrase used for decrypting private keys"
     )
-    gateway: Optional[str] = Field(
+    gateway: Optional[Union[str, ConnectionData]] = Field(
         None, description="A shell command string to use as a proxy or gateway"
     )
     forward_agent: Optional[bool] = Field(

src/jobflow_remote/remote/host/remote.py

@@ -54,15 +54,48 @@ def __init__(
         self._create_connection()
 
     def _create_connection(self):
-        self._connection = fabric.Connection(
+        self._connection = self._get_single_connection(
             host=self.host,
             user=self.user,
             port=self.port,
             config=self.config,
             gateway=self.gateway,
+            connect_kwargs=self.connect_kwargs,
+        )
+
+    def _get_single_connection(
+        self,
+        host,
+        user,
+        port,
+        config,
+        gateway,
+        connect_kwargs,
+    ):
+        """
+        Helper method to generate a fabric Connection given standard parameters.
+        """
+        from jobflow_remote.config.base import ConnectionData
+
+        if isinstance(gateway, ConnectionData):
+            gateway = self._get_single_connection(
+                host=gateway.host,
+                user=gateway.user,
+                port=gateway.port,
+                config=None,
+                gateway=gateway.gateway,
+                connect_kwargs=gateway.get_connect_kwargs(),
+            )
+
+        return fabric.Connection(
+            host=host,
+            user=user,
+            port=port,
+            config=config,
+            gateway=gateway,
             forward_agent=self.forward_agent,
             connect_timeout=self.connect_timeout,
-            connect_kwargs=self.connect_kwargs,
+            connect_kwargs=connect_kwargs,
             inline_ssh_env=self.inline_ssh_env,
         )
 
@@ -160,14 +193,24 @@ def write_text_file(self, filepath: str | Path, content: str):
     def connect(self):
         self.connection.open()
         if self.keepalive:
-            self.connection.transport.set_keepalive(self.keepalive)
+            # create all the nested connections for all the gateways.
+            connection = self.connection
+            while connection:
+                if isinstance(connection, fabric.Connection):
+                    connection.transport.set_keepalive(self.keepalive)
+                    connection = connection.gateway
 
     def close(self) -> bool:
-        try:
-            self.connection.close()
-        except Exception:
-            return False
-        return True
+        connection = self.connection
+        all_closed = True
+        while connection:
+            try:
+                if isinstance(connection, fabric.Connection):
+                    connection.close()
+            except Exception:
+                all_closed = False
+            connection = connection.gateway
+        return all_closed
 
     @property
     def is_connected(self) -> bool: