bugfix/close SSHClient in sftp connector (#309)

* Add context manager for all fsspec clients, manually close SSHClient in sftp implementation

* fix s3 connector

* fix async upload

* Add get client for all child connector configs to also check requirements

* fix get client in child connector configs

Author: rbiseck3

CHANGELOG.md

@@ -1,9 +1,13 @@
-## 0.3.10-dev0
+## 0.3.10-dev1
 
 ### Enhancements
 
 * **Support more concrete FileData content for batch support**
 
+### Fixes
+
+* **Fix closing SSHClient in sftp connector**
+
 ## 0.3.9
 
 ### Enhancements

unstructured_ingest/__version__.py

@@ -1 +1 @@
-__version__ = "0.3.10-dev0"  # pragma: no cover
+__version__ = "0.3.10-dev1"  # pragma: no cover

unstructured_ingest/v2/processes/connectors/fsspec/azure.py

@@ -1,14 +1,14 @@
 from __future__ import annotations
 
+from contextlib import contextmanager
 from dataclasses import dataclass, field
-from pathlib import Path
 from time import time
-from typing import Any, Generator, Optional
+from typing import TYPE_CHECKING, Any, Generator, Optional
 
 from pydantic import Field, Secret
 
 from unstructured_ingest.utils.dep_check import requires_dependencies
-from unstructured_ingest.v2.interfaces import DownloadResponse, FileData, FileDataSourceMetadata
+from unstructured_ingest.v2.interfaces import FileDataSourceMetadata
 from unstructured_ingest.v2.processes.connector_registry import (
     DestinationRegistryEntry,
     SourceRegistryEntry,
@@ -25,6 +25,9 @@
 )
 from unstructured_ingest.v2.processes.connectors.fsspec.utils import json_serial, sterilize_dict
 
+if TYPE_CHECKING:
+    from adlfs import AzureBlobFileSystem
+
 CONNECTOR_TYPE = "azure"
 
 
@@ -89,24 +92,22 @@ def get_access_config(self) -> dict[str, Any]:
         }
         return access_configs
 
+    @requires_dependencies(["adlfs", "fsspec"], extras="azure")
+    @contextmanager
+    def get_client(self, protocol: str) -> Generator["AzureBlobFileSystem", None, None]:
+        with super().get_client(protocol=protocol) as client:
+            yield client
+
 
 @dataclass
 class AzureIndexer(FsspecIndexer):
     connection_config: AzureConnectionConfig
     index_config: AzureIndexerConfig
     connector_type: str = CONNECTOR_TYPE
 
-    @requires_dependencies(["adlfs", "fsspec"], extras="azure")
-    def precheck(self) -> None:
-        super().precheck()
-
     def sterilize_info(self, file_data: dict) -> dict:
         return sterilize_dict(data=file_data, default=azure_json_serial)
 
-    @requires_dependencies(["adlfs", "fsspec"], extras="azure")
-    def run(self, **kwargs: Any) -> Generator[FileData, None, None]:
-        return super().run(**kwargs)
-
     def get_metadata(self, file_data: dict) -> FileDataSourceMetadata:
         path = file_data["name"]
         date_created = (
@@ -149,14 +150,6 @@ class AzureDownloader(FsspecDownloader):
     connector_type: str = CONNECTOR_TYPE
     download_config: Optional[AzureDownloaderConfig] = field(default_factory=AzureDownloaderConfig)
 
-    @requires_dependencies(["adlfs", "fsspec"], extras="azure")
-    def run(self, file_data: FileData, **kwargs: Any) -> DownloadResponse:
-        return super().run(file_data=file_data, **kwargs)
-
-    @requires_dependencies(["adlfs", "fsspec"], extras="azure")
-    async def run_async(self, file_data: FileData, **kwargs: Any) -> DownloadResponse:
-        return await super().run_async(file_data=file_data, **kwargs)
-
 
 class AzureUploaderConfig(FsspecUploaderConfig):
     pass
@@ -168,22 +161,6 @@ class AzureUploader(FsspecUploader):
     connection_config: AzureConnectionConfig
     upload_config: AzureUploaderConfig = field(default=None)
 
-    @requires_dependencies(["adlfs", "fsspec"], extras="azure")
-    def __post_init__(self):
-        super().__post_init__()
-
-    @requires_dependencies(["adlfs", "fsspec"], extras="azure")
-    def precheck(self) -> None:
-        super().precheck()
-
-    @requires_dependencies(["adlfs", "fsspec"], extras="azure")
-    def run(self, path: Path, file_data: FileData, **kwargs: Any) -> None:
-        return super().run(path=path, file_data=file_data, **kwargs)
-
-    @requires_dependencies(["adlfs", "fsspec"], extras="azure")
-    async def run_async(self, path: Path, file_data: FileData, **kwargs: Any) -> None:
-        return await super().run_async(path=path, file_data=file_data, **kwargs)
-
 
 azure_source_entry = SourceRegistryEntry(
     indexer=AzureIndexer,

unstructured_ingest/v2/processes/connectors/fsspec/box.py

@@ -1,16 +1,16 @@
 from __future__ import annotations
 
+from contextlib import contextmanager
 from dataclasses import dataclass, field
-from pathlib import Path
 from time import time
-from typing import Annotated, Any, Generator, Optional
+from typing import TYPE_CHECKING, Annotated, Any, Generator, Optional
 
 from dateutil import parser
 from pydantic import Field, Secret
 from pydantic.functional_validators import BeforeValidator
 
 from unstructured_ingest.utils.dep_check import requires_dependencies
-from unstructured_ingest.v2.interfaces import DownloadResponse, FileData, FileDataSourceMetadata
+from unstructured_ingest.v2.interfaces import FileDataSourceMetadata
 from unstructured_ingest.v2.processes.connector_registry import (
     DestinationRegistryEntry,
     SourceRegistryEntry,
@@ -28,6 +28,9 @@
 )
 from unstructured_ingest.v2.processes.connectors.utils import conform_string_to_dict
 
+if TYPE_CHECKING:
+    from boxfs import BoxFileSystem
+
 CONNECTOR_TYPE = "box"
 
 
@@ -72,21 +75,19 @@ def get_access_config(self) -> dict[str, Any]:
 
         return access_kwargs_with_oauth
 
+    @requires_dependencies(["boxfs"], extras="box")
+    @contextmanager
+    def get_client(self, protocol: str) -> Generator["BoxFileSystem", None, None]:
+        with super().get_client(protocol=protocol) as client:
+            yield client
+
 
 @dataclass
 class BoxIndexer(FsspecIndexer):
     connection_config: BoxConnectionConfig
     index_config: BoxIndexerConfig
     connector_type: str = CONNECTOR_TYPE
 
-    @requires_dependencies(["boxfs"], extras="box")
-    def run(self, **kwargs: Any) -> Generator[FileData, None, None]:
-        return super().run(**kwargs)
-
-    @requires_dependencies(["boxfs"], extras="box")
-    def precheck(self) -> None:
-        super().precheck()
-
     def get_metadata(self, file_data: dict) -> FileDataSourceMetadata:
         path = file_data["name"]
         date_created = None
@@ -126,14 +127,6 @@ class BoxDownloader(FsspecDownloader):
     connector_type: str = CONNECTOR_TYPE
     download_config: Optional[BoxDownloaderConfig] = field(default_factory=BoxDownloaderConfig)
 
-    @requires_dependencies(["boxfs"], extras="box")
-    def run(self, file_data: FileData, **kwargs: Any) -> DownloadResponse:
-        return super().run(file_data=file_data, **kwargs)
-
-    @requires_dependencies(["boxfs"], extras="box")
-    async def run_async(self, file_data: FileData, **kwargs: Any) -> DownloadResponse:
-        return await super().run_async(file_data=file_data, **kwargs)
-
 
 class BoxUploaderConfig(FsspecUploaderConfig):
     pass
@@ -145,22 +138,6 @@ class BoxUploader(FsspecUploader):
     connection_config: BoxConnectionConfig
     upload_config: BoxUploaderConfig = field(default=None)
 
-    @requires_dependencies(["boxfs"], extras="box")
-    def __post_init__(self):
-        super().__post_init__()
-
-    @requires_dependencies(["boxfs"], extras="box")
-    def precheck(self) -> None:
-        super().precheck()
-
-    @requires_dependencies(["boxfs"], extras="box")
-    def run(self, path: Path, file_data: FileData, **kwargs: Any) -> None:
-        return super().run(path=path, file_data=file_data, **kwargs)
-
-    @requires_dependencies(["boxfs"], extras="box")
-    async def run_async(self, path: Path, file_data: FileData, **kwargs: Any) -> None:
-        return await super().run_async(path=path, file_data=file_data, **kwargs)
-
 
 box_source_entry = SourceRegistryEntry(
     indexer=BoxIndexer,

unstructured_ingest/v2/processes/connectors/fsspec/dropbox.py

@@ -1,14 +1,14 @@
 from __future__ import annotations
 
+from contextlib import contextmanager
 from dataclasses import dataclass, field
-from pathlib import Path
 from time import time
-from typing import Any, Generator, Optional
+from typing import TYPE_CHECKING, Generator, Optional
 
 from pydantic import Field, Secret
 
 from unstructured_ingest.utils.dep_check import requires_dependencies
-from unstructured_ingest.v2.interfaces import DownloadResponse, FileData, FileDataSourceMetadata
+from unstructured_ingest.v2.interfaces import FileDataSourceMetadata
 from unstructured_ingest.v2.processes.connector_registry import (
     DestinationRegistryEntry,
     SourceRegistryEntry,
@@ -24,11 +24,16 @@
     FsspecUploaderConfig,
 )
 
+if TYPE_CHECKING:
+    from dropboxdrivefs import DropboxDriveFileSystem
+
 CONNECTOR_TYPE = "dropbox"
 
 
 class DropboxIndexerConfig(FsspecIndexerConfig):
-    pass
+    def model_post_init(self, __context):
+        if not self.path_without_protocol.startswith("/"):
+            self.path_without_protocol = "/" + self.path_without_protocol
 
 
 class DropboxAccessConfig(FsspecAccessConfig):
@@ -42,6 +47,12 @@ class DropboxConnectionConfig(FsspecConnectionConfig):
     )
     connector_type: str = Field(default=CONNECTOR_TYPE, init=False)
 
+    @requires_dependencies(["dropboxdrivefs", "fsspec"], extras="dropbox")
+    @contextmanager
+    def get_client(self, protocol: str) -> Generator["DropboxDriveFileSystem", None, None]:
+        with super().get_client(protocol=protocol) as client:
+            yield client
+
 
 @dataclass
 class DropboxIndexer(FsspecIndexer):
@@ -83,20 +94,6 @@ def get_metadata(self, file_data: dict) -> FileDataSourceMetadata:
             filesize_bytes=file_size,
         )
 
-    @requires_dependencies(["dropboxdrivefs", "fsspec"], extras="dropbox")
-    def __post_init__(self):
-        # dropbox expects the path to start with a /
-        if not self.index_config.path_without_protocol.startswith("/"):
-            self.index_config.path_without_protocol = "/" + self.index_config.path_without_protocol
-
-    @requires_dependencies(["dropboxdrivefs", "fsspec"], extras="dropbox")
-    def precheck(self) -> None:
-        super().precheck()
-
-    @requires_dependencies(["dropboxdrivefs", "fsspec"], extras="dropbox")
-    def run(self, **kwargs: Any) -> Generator[FileData, None, None]:
-        return super().run(**kwargs)
-
 
 class DropboxDownloaderConfig(FsspecDownloaderConfig):
     pass
@@ -111,14 +108,6 @@ class DropboxDownloader(FsspecDownloader):
         default_factory=DropboxDownloaderConfig
     )
 
-    @requires_dependencies(["dropboxdrivefs", "fsspec"], extras="dropbox")
-    def run(self, file_data: FileData, **kwargs: Any) -> DownloadResponse:
-        return super().run(file_data=file_data, **kwargs)
-
-    @requires_dependencies(["dropboxdrivefs", "fsspec"], extras="dropbox")
-    async def run_async(self, file_data: FileData, **kwargs: Any) -> DownloadResponse:
-        return await super().run_async(file_data=file_data, **kwargs)
-
 
 class DropboxUploaderConfig(FsspecUploaderConfig):
     pass
@@ -130,22 +119,6 @@ class DropboxUploader(FsspecUploader):
     connection_config: DropboxConnectionConfig
     upload_config: DropboxUploaderConfig = field(default=None)
 
-    @requires_dependencies(["dropboxdrivefs", "fsspec"], extras="dropbox")
-    def __post_init__(self):
-        super().__post_init__()
-
-    @requires_dependencies(["dropboxdrivefs", "fsspec"], extras="dropbox")
-    def precheck(self) -> None:
-        super().precheck()
-
-    @requires_dependencies(["dropboxdrivefs", "fsspec"], extras="dropbox")
-    def run(self, path: Path, file_data: FileData, **kwargs: Any) -> None:
-        return super().run(path=path, file_data=file_data, **kwargs)
-
-    @requires_dependencies(["dropboxdrivefs", "fsspec"], extras="dropbox")
-    async def run_async(self, path: Path, file_data: FileData, **kwargs: Any) -> None:
-        return await super().run_async(path=path, file_data=file_data, **kwargs)
-
 
 dropbox_source_entry = SourceRegistryEntry(
     indexer=DropboxIndexer,