Add pkey, disabled_algorithms, transport_factory and auth_strategy parameters to paramiko.client.connect

Author: Ayush Patel

dlt/common/configuration/specs/sftp_credentials.py

@@ -1,12 +1,23 @@
-from typing import Any, Dict, Optional, Annotated
+from typing import Any, Dict, Optional, Annotated, TYPE_CHECKING
 
-from dlt.common.typing import TSecretStrValue, DictStrAny, SocketLike
+from dlt.common.typing import TSecretStrValue, SocketLike, SFTPTransportFactory, DictStrListStr
 from dlt.common.configuration.specs.base_configuration import (
     CredentialsConfiguration,
     configspec,
     NotResolved,
 )
 
+if TYPE_CHECKING:
+    try:
+        from paramiko import PKey
+        from paramiko.auth_strategy import AuthStrategy
+    except ImportError:
+        PKey = Any  # type: ignore[misc, assignment]
+        AuthStrategy = Any  # type: ignore[misc, assignment]
+else:
+    PKey = Any
+    AuthStrategy = Any
+
 
 @configspec
 class SFTPCredentials(CredentialsConfiguration):
@@ -32,6 +43,8 @@ class SFTPCredentials(CredentialsConfiguration):
     sftp_port: Optional[int] = 22
     sftp_username: Optional[str] = None
     sftp_password: Optional[TSecretStrValue] = None
+    # Runtime-only pkey; cannot be loaded from env var, skip configspec.
+    sftp_pkey: Annotated[Optional[PKey], NotResolved()] = None
     sftp_key_filename: Optional[str] = None
     sftp_key_passphrase: Optional[TSecretStrValue] = None
     sftp_timeout: Optional[float] = None
@@ -48,6 +61,10 @@ class SFTPCredentials(CredentialsConfiguration):
     sftp_gss_deleg_creds: Optional[bool] = True
     sftp_gss_host: Optional[str] = None
     sftp_gss_trust_dns: Optional[bool] = True
+    # Runtime-only vars below; cannot be loaded from env var, skip configspec.
+    sftp_disabled_algorithms: Annotated[Optional[DictStrListStr], NotResolved()] = None
+    sftp_transport_factory: Annotated[Optional[SFTPTransportFactory], NotResolved()] = None
+    sftp_auth_strategy: Annotated[Optional[AuthStrategy], NotResolved()] = None
 
     def to_fsspec_credentials(self) -> Dict[str, Any]:
         """Return a dict that can be passed to fsspec SFTP/SSHClient.connect method."""
@@ -56,6 +73,7 @@ def to_fsspec_credentials(self) -> Dict[str, Any]:
             "port": self.sftp_port,
             "username": self.sftp_username,
             "password": self.sftp_password,
+            "pkey": self.sftp_pkey,
             "key_filename": self.sftp_key_filename,
             "passphrase": self.sftp_key_passphrase,
             "timeout": self.sftp_timeout,
@@ -71,6 +89,9 @@ def to_fsspec_credentials(self) -> Dict[str, Any]:
             "gss_deleg_creds": self.sftp_gss_deleg_creds,
             "gss_host": self.sftp_gss_host,
             "gss_trust_dns": self.sftp_gss_trust_dns,
+            "disabled_algorithms": self.sftp_disabled_algorithms,
+            "transport_factory": self.sftp_transport_factory,
+            "auth_strategy": self.sftp_auth_strategy,
         }
 
         return credentials

dlt/common/typing.py

@@ -2,6 +2,7 @@
 from datetime import datetime, date  # noqa: I251
 import inspect
 import os
+import socket
 from re import Pattern as _REPattern
 from types import FunctionType
 from typing import (
@@ -65,6 +66,15 @@
 
 typingGenericAlias: Tuple[Any, ...] = (_GenericAlias, _SpecialGenericAlias, GenericAlias)
 
+if TYPE_CHECKING:
+    try:
+        from paramiko import Transport
+    except ImportError:
+        Transport = Any  # type: ignore[misc, assignment]
+else:
+    Transport = Any
+
+SFTPTransportFactory: TypeAlias = Callable[[socket.socket], Transport]
 
 from dlt.common.pendulum import timedelta, pendulum
 
@@ -89,6 +99,7 @@
 DictStrAny: TypeAlias = Dict[str, Any]
 DictStrStr: TypeAlias = Dict[str, str]
 DictStrOptionalStr: TypeAlias = Dict[str, Optional[str]]
+DictStrListStr: TypeAlias = Dict[str, List[str]]
 StrAny: TypeAlias = Mapping[str, Any]  # immutable, covariant entity
 StrStr: TypeAlias = Mapping[str, str]  # immutable, covariant entity
 StrStrStr: TypeAlias = Mapping[str, Mapping[str, str]]  # immutable, covariant entity

pyproject.toml

@@ -250,6 +250,7 @@ dev = [
     "pydoclint>=0.6.5,<0.7",
     # limit the pyarrow version not to test on too new one
     "pyarrow>=16.0.0,<19.0.0",
+    "types-paramiko>=3.5.0.20250516",
 ]
 sources = [
     "connectorx>=0.3.3 ; python_version >= '3.9'",

tests/load/filesystem_sftp/test_filesystem_sftp.py

@@ -3,7 +3,6 @@
 import fsspec
 import socket
 import dlt
-from dlt.common.configuration.specs import SFTPCredentials
 
 from dlt.common.json import json
 from dlt.common.configuration.inject import with_config
@@ -12,6 +11,10 @@
 
 from tests.load.utils import ALL_FILESYSTEM_DRIVERS
 
+from paramiko.auth_strategy import Password
+from paramiko import RSAKey, Transport
+from paramiko.ssh_exception import SSHException
+
 if "sftp" not in ALL_FILESYSTEM_DRIVERS:
     pytest.skip("sftp filesystem driver not configured", allow_module_level=True)
 
@@ -126,7 +129,17 @@ def states():
         assert sorted(result_states) == sorted(expected_states)
 
 
-def run_sftp_auth(user, password=None, key=None, passphrase=None, sock=None):
+def run_sftp_auth(
+    user,
+    password=None,
+    pkey=None,
+    key=None,
+    passphrase=None,
+    sock=None,
+    disabled_algorithms=None,
+    transport_factory=None,
+    auth_strategy=None,
+):
     env_vars = {
         "SOURCES__FILESYSTEM__BUCKET_URL": "sftp://localhost",
         "SOURCES__FILESYSTEM__CREDENTIALS__SFTP_PORT": "2222",
@@ -144,6 +157,14 @@ def run_sftp_auth(user, password=None, key=None, passphrase=None, sock=None):
 
     config = get_config()
 
+    if disabled_algorithms:
+        config.credentials.sftp_disabled_algorithms = disabled_algorithms  # type: ignore[union-attr]
+    if transport_factory:
+        config.credentials.sftp_transport_factory = transport_factory  # type: ignore[union-attr]
+    if auth_strategy:
+        config.credentials.sftp_auth_strategy = auth_strategy  # type: ignore[union-attr]
+    if pkey:
+        config.credentials.sftp_pkey = pkey  # type: ignore[union-attr]
     if sock:
         config.credentials.sftp_sock = sock  # type: ignore[union-attr]
 
@@ -163,6 +184,17 @@ def test_filesystem_sftp_auth_private_key_protected():
     run_sftp_auth("bobby", key=get_key_path("bobby"), passphrase="passphrase123")
 
 
+def test_filesystem_sftp_auth_pkey():
+    run_sftp_auth("foo", pkey=RSAKey.from_private_key_file(get_key_path("foo")))
+
+
+def test_filesystem_sftp_pkey_auth_pkey_protected():
+    run_sftp_auth(
+        "bobby",
+        pkey=RSAKey.from_private_key_file(filename=get_key_path("bobby"), password="passphrase123"),
+    )
+
+
 # Test requires - ssh_agent with user's bobby key loaded. The commands and file names required are:
 # eval "$(ssh-agent -s)"
 # cp /path/to/tests/load/filesystem_sftp/bobby_rsa* ~/.ssh/id_rsa
@@ -190,3 +222,52 @@ def test_filesystem_sftp_with_socket():
     sock.close()
     with pytest.raises(OSError):
         run_sftp_auth("billy", key=get_key_path("billy"), sock=sock)
+
+
+class TaggedTransport(Transport):
+    """A Transport class that tags itself so we can detect it in tests."""
+
+    def __init__(self, sock, **kwargs):
+        super().__init__(sock, **kwargs)
+        # Add any custom state or markers you like:
+        self.factory_tag = "used-tagged-transport"
+
+
+def test_filesystem_sftp_with_tagged_transport():
+    created = []
+
+    def factory(sock, **kwargs):
+        t = TaggedTransport(sock, **kwargs)
+        created.append(t)
+        return t
+
+    run_sftp_auth("foo", key=get_key_path("foo"), transport_factory=factory)
+
+    # Verify it was used
+    assert len(created) == 1, "Custom transport factory never ran"
+    transport = created[0]
+    assert isinstance(transport, TaggedTransport)
+    assert transport.factory_tag == "used-tagged-transport"
+
+    # And ensure it's still functional
+    sftp = transport.open_sftp_client()
+    assert hasattr(sftp, "listdir")
+    sftp.close()
+
+
+def test_filesystem_sftp_disabled_algorithms():
+    # we know foo’s server uses rsa keys
+    with pytest.raises(SSHException):
+        run_sftp_auth(
+            "foo",
+            key=get_key_path("foo"),
+            disabled_algorithms={"pubkeys": ["ssh-rsa", "rsa-sha2-256", "rsa-sha2-512"]},
+        )
+
+
+def test_filesystem_sftp_auth_strategy():
+    # Verify that passing an alternate auth_strategy makes it through config.
+    run_sftp_auth(
+        "foo",
+        auth_strategy=Password("foo", lambda: "pass"),
+    )