[DH-5335] For SSH connection take ports if they are set (#362)

Dataherald · Jan 30, 2024 · 1a5a2be · 1a5a2be
1 parent 1c32ef4
commit 1a5a2be
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/dataherald/sql_database/base.py b/dataherald/sql_database/base.py
@@ -116,12 +116,15 @@ def from_uri_ssh(cls, database_info: DatabaseConnection):
         ssh = database_info.ssh_settings
 
         server = SSHTunnelForwarder(
-            (ssh.host, 22),
+            (ssh.host, 22 if not ssh.port else int(ssh.port)),
             ssh_username=ssh.username,
             ssh_password=fernet_encrypt.decrypt(ssh.password),
             ssh_pkey=file_path,
             ssh_private_key_password=fernet_encrypt.decrypt(ssh.private_key_password),
-            remote_bind_address=(db_uri_obj.hostname, 5432),
+            remote_bind_address=(
+                db_uri_obj.hostname,
+                5432 if not db_uri_obj.port else int(db_uri_obj.port),
+            ),
         )
         server.start()
         local_port = str(server.local_bind_port)

diff --git a/dataherald/sql_database/models/types.py b/dataherald/sql_database/models/types.py
@@ -51,6 +51,7 @@ class SSHSettings(BaseSettings):
     host: str | None
     username: str | None
     password: str | None
+    port: str | None = "22"
     private_key_password: str | None
 
     class Config: