unstructured-ingest 0.3.9__py3-none-any.whl → 0.3.10__py3-none-any.whl

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 @@ from unstructured_ingest.v2.processes.connectors.fsspec.fsspec import (
     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 @@ class DropboxIndexer(FsspecIndexer):
             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,

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

@@ -4,6 +4,7 @@ import os
 import random
 import shutil
 import tempfile
+from contextlib import contextmanager
 from dataclasses import dataclass, field
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Generator, Optional, TypeVar
@@ -78,6 +79,15 @@ class FsspecConnectionConfig(ConnectionConfig):
     access_config: Secret[FsspecAccessConfig]
     connector_type: str = Field(default=CONNECTOR_TYPE, init=False)
 
+    @contextmanager
+    def get_client(self, protocol: str) -> Generator["AbstractFileSystem", None, None]:
+        from fsspec import get_filesystem_class
+
+        client = get_filesystem_class(protocol)(
+            **self.get_access_config(),
+        )
+        yield client
+
 
 FsspecIndexerConfigT = TypeVar("FsspecIndexerConfigT", bound=FsspecIndexerConfig)
 FsspecConnectionConfigT = TypeVar("FsspecConnectionConfigT", bound=FsspecConnectionConfig)
@@ -89,14 +99,6 @@ class FsspecIndexer(Indexer):
     index_config: FsspecIndexerConfigT
     connector_type: str = Field(default=CONNECTOR_TYPE, init=False)
 
-    @property
-    def fs(self) -> "AbstractFileSystem":
-        from fsspec import get_filesystem_class
-
-        return get_filesystem_class(self.index_config.protocol)(
-            **self.connection_config.get_access_config(),
-        )
-
     def precheck(self) -> None:
         from fsspec import get_filesystem_class
 
@@ -110,7 +112,8 @@ class FsspecIndexer(Indexer):
                 return
             file_to_sample = valid_files[0]
             logger.debug(f"attempting to make HEAD request for file: {file_to_sample}")
-            self.fs.head(path=file_to_sample)
+            with self.connection_config.get_client(protocol=self.index_config.protocol) as client:
+                client.head(path=file_to_sample)
         except Exception as e:
             logger.error(f"failed to validate connection: {e}", exc_info=True)
             raise SourceConnectionError(f"failed to validate connection: {e}")
@@ -120,16 +123,18 @@ class FsspecIndexer(Indexer):
             # fs.ls does not walk directories
             # directories that are listed in cloud storage can cause problems
             # because they are seen as 0 byte files
-            files = self.fs.ls(self.index_config.path_without_protocol, detail=True)
+            with self.connection_config.get_client(protocol=self.index_config.protocol) as client:
+                files = client.ls(self.index_config.path_without_protocol, detail=True)
 
         else:
             # fs.find will recursively walk directories
             # "size" is a common key for all the cloud protocols with fs
-            found = self.fs.find(
-                self.index_config.path_without_protocol,
-                detail=True,
-            )
-            files = found.values()
+            with self.connection_config.get_client(protocol=self.index_config.protocol) as client:
+                found = client.find(
+                    self.index_config.path_without_protocol,
+                    detail=True,
+                )
+                files = found.values()
         filtered_files = [
             file for file in files if file.get("size") > 0 and file.get("type") == "file"
         ]
@@ -200,15 +205,8 @@ class FsspecDownloader(Downloader):
     )
 
     def is_async(self) -> bool:
-        return self.fs.async_impl
-
-    @property
-    def fs(self) -> "AbstractFileSystem":
-        from fsspec import get_filesystem_class
-
-        return get_filesystem_class(self.protocol)(
-            **self.connection_config.get_access_config(),
-        )
+        with self.connection_config.get_client(protocol=self.protocol) as client:
+            return client.async_impl
 
     def handle_directory_download(self, lpath: Path) -> None:
         # If the object's name contains certain characters (i.e. '?'), it
@@ -237,7 +235,8 @@ class FsspecDownloader(Downloader):
         download_path.parent.mkdir(parents=True, exist_ok=True)
         try:
             rpath = file_data.additional_metadata["original_file_path"]
-            self.fs.get(rpath=rpath, lpath=download_path.as_posix())
+            with self.connection_config.get_client(protocol=self.protocol) as client:
+                client.get(rpath=rpath, lpath=download_path.as_posix())
             self.handle_directory_download(lpath=download_path)
         except Exception as e:
             logger.error(f"failed to download file {file_data.identifier}: {e}", exc_info=True)
@@ -249,7 +248,8 @@ class FsspecDownloader(Downloader):
         download_path.parent.mkdir(parents=True, exist_ok=True)
         try:
             rpath = file_data.additional_metadata["original_file_path"]
-            await self.fs.get(rpath=rpath, lpath=download_path.as_posix())
+            with self.connection_config.get_client(protocol=self.protocol) as client:
+                await client.get(rpath=rpath, lpath=download_path.as_posix())
             self.handle_directory_download(lpath=download_path)
         except Exception as e:
             logger.error(f"failed to download file {file_data.identifier}: {e}", exc_info=True)
@@ -268,9 +268,11 @@ FsspecUploaderConfigT = TypeVar("FsspecUploaderConfigT", bound=FsspecUploaderCon
 class FsspecUploader(Uploader):
     connector_type: str = CONNECTOR_TYPE
     upload_config: FsspecUploaderConfigT = field(default=None)
+    connection_config: FsspecConnectionConfigT
 
     def is_async(self) -> bool:
-        return self.fs.async_impl
+        with self.connection_config.get_client(protocol=self.upload_config.protocol) as client:
+            return client.async_impl
 
     @property
     def fs(self) -> "AbstractFileSystem":
@@ -314,11 +316,13 @@ class FsspecUploader(Uploader):
         path_str = str(path.resolve())
         upload_path = self.get_upload_path(file_data=file_data)
         logger.debug(f"writing local file {path_str} to {upload_path}")
-        self.fs.upload(lpath=path_str, rpath=upload_path.as_posix())
+        with self.connection_config.get_client(protocol=self.upload_config.protocol) as client:
+            client.upload(lpath=path_str, rpath=upload_path.as_posix())
 
     async def run_async(self, path: Path, file_data: FileData, **kwargs: Any) -> None:
         upload_path = self.get_upload_path(file_data=file_data)
         path_str = str(path.resolve())
         # Odd that fsspec doesn't run exists() as async even when client support async
         logger.debug(f"writing local file {path_str} to {upload_path}")
-        self.fs.upload(lpath=path_str, rpath=upload_path.as_posix())
+        with self.connection_config.get_client(protocol=self.upload_config.protocol) as client:
+            client.upload(lpath=path_str, rpath=upload_path.as_posix())

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

@@ -1,16 +1,17 @@
 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, Union
+from typing import TYPE_CHECKING, Any, Generator, Optional, Union
 
 from dateutil import parser
 from pydantic import Field, Secret
 
 from unstructured_ingest.utils.dep_check import requires_dependencies
 from unstructured_ingest.utils.string_and_date_utils import json_to_dict
-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,
@@ -26,6 +27,9 @@ from unstructured_ingest.v2.processes.connectors.fsspec.fsspec import (
     FsspecUploaderConfig,
 )
 
+if TYPE_CHECKING:
+    from gcsfs import GCSFileSystem
+
 CONNECTOR_TYPE = "gcs"
 
 
@@ -93,6 +97,12 @@ class GcsConnectionConfig(FsspecConnectionConfig):
     access_config: Secret[GcsAccessConfig] = Field(default=GcsAccessConfig(), validate_default=True)
     connector_type: str = Field(default=CONNECTOR_TYPE, init=False)
 
+    @requires_dependencies(["gcsfs", "fsspec"], extras="gcs")
+    @contextmanager
+    def get_client(self, protocol: str) -> Generator["GCSFileSystem", None, None]:
+        with super().get_client(protocol=protocol) as client:
+            yield client
+
 
 @dataclass
 class GcsIndexer(FsspecIndexer):
@@ -100,14 +110,6 @@ class GcsIndexer(FsspecIndexer):
     index_config: GcsIndexerConfig
     connector_type: str = CONNECTOR_TYPE
 
-    @requires_dependencies(["gcsfs", "fsspec"], extras="gcs")
-    def run(self, **kwargs: Any) -> Generator[FileData, None, None]:
-        return super().run(**kwargs)
-
-    @requires_dependencies(["gcsfs", "fsspec"], extras="gcs")
-    def precheck(self) -> None:
-        super().precheck()
-
     def get_metadata(self, file_data: dict) -> FileDataSourceMetadata:
         path = file_data["name"]
         date_created = None
@@ -147,14 +149,6 @@ class GcsDownloader(FsspecDownloader):
     connector_type: str = CONNECTOR_TYPE
     download_config: Optional[GcsDownloaderConfig] = field(default_factory=GcsDownloaderConfig)
 
-    @requires_dependencies(["gcsfs", "fsspec"], extras="gcs")
-    def run(self, file_data: FileData, **kwargs: Any) -> DownloadResponse:
-        return super().run(file_data=file_data, **kwargs)
-
-    @requires_dependencies(["gcsfs", "fsspec"], extras="gcs")
-    async def run_async(self, file_data: FileData, **kwargs: Any) -> DownloadResponse:
-        return await super().run_async(file_data=file_data, **kwargs)
-
 
 class GcsUploaderConfig(FsspecUploaderConfig):
     pass
@@ -166,22 +160,6 @@ class GcsUploader(FsspecUploader):
     connection_config: GcsConnectionConfig
     upload_config: GcsUploaderConfig = field(default=None)
 
-    @requires_dependencies(["gcsfs", "fsspec"], extras="gcs")
-    def __post_init__(self):
-        super().__post_init__()
-
-    @requires_dependencies(["gcsfs", "fsspec"], extras="gcs")
-    def precheck(self) -> None:
-        super().precheck()
-
-    @requires_dependencies(["gcsfs", "fsspec"], extras="gcs")
-    def run(self, path: Path, file_data: FileData, **kwargs: Any) -> None:
-        return super().run(path=path, file_data=file_data, **kwargs)
-
-    @requires_dependencies(["gcsfs", "fsspec"], extras="gcs")
-    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)
-
 
 gcs_source_entry = SourceRegistryEntry(
     indexer=GcsIndexer,

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

@@ -1,15 +1,13 @@
 import contextlib
+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.processes.connector_registry import (
@@ -29,6 +27,9 @@ from unstructured_ingest.v2.processes.connectors.fsspec.fsspec import (
 
 CONNECTOR_TYPE = "s3"
 
+if TYPE_CHECKING:
+    from s3fs import S3FileSystem
+
 
 class S3IndexerConfig(FsspecIndexerConfig):
     pass
@@ -72,6 +73,12 @@ class S3ConnectionConfig(FsspecConnectionConfig):
         )
         return access_configs
 
+    @requires_dependencies(["s3fs", "fsspec"], extras="s3")
+    @contextmanager
+    def get_client(self, protocol: str) -> Generator["S3FileSystem", None, None]:
+        with super().get_client(protocol=protocol) as client:
+            yield client
+
 
 @dataclass
 class S3Indexer(FsspecIndexer):
@@ -97,7 +104,8 @@ class S3Indexer(FsspecIndexer):
         version = file_data.get("ETag").rstrip('"').lstrip('"') if "ETag" in file_data else None
         metadata: dict[str, str] = {}
         with contextlib.suppress(AttributeError):
-            metadata = self.fs.metadata(path=path)
+            with self.connection_config.get_client(protocol=self.index_config.protocol) as client:
+                metadata = client.metadata(path=path)
         record_locator = {
             "protocol": self.index_config.protocol,
             "remote_file_path": self.index_config.remote_url,
@@ -114,14 +122,6 @@ class S3Indexer(FsspecIndexer):
             filesize_bytes=file_size,
         )
 
-    @requires_dependencies(["s3fs", "fsspec"], extras="s3")
-    def run(self, **kwargs: Any) -> Generator[FileData, None, None]:
-        return super().run(**kwargs)
-
-    @requires_dependencies(["s3fs", "fsspec"], extras="s3")
-    def precheck(self) -> None:
-        super().precheck()
-
 
 class S3DownloaderConfig(FsspecDownloaderConfig):
     pass
@@ -134,14 +134,6 @@ class S3Downloader(FsspecDownloader):
     connector_type: str = CONNECTOR_TYPE
     download_config: Optional[S3DownloaderConfig] = field(default_factory=S3DownloaderConfig)
 
-    @requires_dependencies(["s3fs", "fsspec"], extras="s3")
-    def run(self, file_data: FileData, **kwargs: Any) -> DownloadResponse:
-        return super().run(file_data=file_data, **kwargs)
-
-    @requires_dependencies(["s3fs", "fsspec"], extras="s3")
-    async def run_async(self, file_data: FileData, **kwargs: Any) -> DownloadResponse:
-        return await super().run_async(file_data=file_data, **kwargs)
-
 
 class S3UploaderConfig(FsspecUploaderConfig):
     pass
@@ -153,22 +145,6 @@ class S3Uploader(FsspecUploader):
     connection_config: S3ConnectionConfig
     upload_config: S3UploaderConfig = field(default=None)
 
-    @requires_dependencies(["s3fs", "fsspec"], extras="s3")
-    def precheck(self) -> None:
-        super().precheck()
-
-    @requires_dependencies(["s3fs", "fsspec"], extras="s3")
-    def __post_init__(self):
-        super().__post_init__()
-
-    @requires_dependencies(["s3fs", "fsspec"], extras="s3")
-    def run(self, path: Path, file_data: FileData, **kwargs: Any) -> None:
-        return super().run(path=path, file_data=file_data, **kwargs)
-
-    @requires_dependencies(["s3fs", "fsspec"], extras="s3")
-    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)
-
 
 s3_source_entry = SourceRegistryEntry(
     indexer=S3Indexer,

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

@@ -1,16 +1,17 @@
 from __future__ import annotations
 
 import os
+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 urllib.parse import urlparse
 
 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 FileData, FileDataSourceMetadata
 from unstructured_ingest.v2.processes.connector_registry import (
     DestinationRegistryEntry,
     SourceRegistryEntry,
@@ -26,6 +27,9 @@ from unstructured_ingest.v2.processes.connectors.fsspec.fsspec import (
     FsspecUploaderConfig,
 )
 
+if TYPE_CHECKING:
+    from fsspec.implementations.sftp import SFTPFileSystem
+
 CONNECTOR_TYPE = "sftp"
 
 
@@ -67,6 +71,19 @@ class SftpConnectionConfig(FsspecConnectionConfig):
         }
         return access_config
 
+    @contextmanager
+    @requires_dependencies(["paramiko", "fsspec"], extras="sftp")
+    def get_client(self, protocol: str) -> Generator["SFTPFileSystem", None, None]:
+        # The paramiko.SSHClient() client that's opened by the SFTPFileSystem
+        # never gets closed so explicitly adding that as part of this context manager
+        from fsspec import get_filesystem_class
+
+        client: SFTPFileSystem = get_filesystem_class(protocol)(
+            **self.get_access_config(),
+        )
+        yield client
+        client.client.close()
+
 
 @dataclass
 class SftpIndexer(FsspecIndexer):
@@ -74,13 +91,11 @@ class SftpIndexer(FsspecIndexer):
     index_config: SftpIndexerConfig
     connector_type: str = CONNECTOR_TYPE
 
-    @requires_dependencies(["paramiko", "fsspec"], extras="sftp")
     def __post_init__(self):
         parsed_url = urlparse(self.index_config.remote_url)
         self.connection_config.host = parsed_url.hostname or self.connection_config.host
         self.connection_config.port = parsed_url.port or self.connection_config.port
 
-    @requires_dependencies(["paramiko", "fsspec"], extras="sftp")
     def run(self, **kwargs: Any) -> Generator[FileData, None, None]:
         for file in super().run(**kwargs):
             new_identifier = (
@@ -92,10 +107,6 @@ class SftpIndexer(FsspecIndexer):
             file.identifier = new_identifier
             yield file
 
-    @requires_dependencies(["paramiko", "fsspec"], extras="sftp")
-    def precheck(self) -> None:
-        super().precheck()
-
     def get_metadata(self, file_data: dict) -> FileDataSourceMetadata:
         path = file_data["name"]
         date_created = str(file_data.get("time").timestamp()) if "time" in file_data else None
@@ -128,20 +139,11 @@ class SftpDownloader(FsspecDownloader):
     connector_type: str = Field(default=CONNECTOR_TYPE, init=False)
     download_config: Optional[SftpDownloaderConfig] = field(default_factory=SftpDownloaderConfig)
 
-    @requires_dependencies(["paramiko", "fsspec"], extras="sftp")
     def __post_init__(self):
         parsed_url = urlparse(self.download_config.remote_url)
         self.connection_config.host = parsed_url.hostname or self.connection_config.host
         self.connection_config.port = parsed_url.port or self.connection_config.port
 
-    @requires_dependencies(["paramiko", "fsspec"], extras="sftp")
-    def run(self, file_data: FileData, **kwargs: Any) -> DownloadResponse:
-        return super().run(file_data=file_data, **kwargs)
-
-    @requires_dependencies(["paramiko", "fsspec"], extras="sftp")
-    async def run_async(self, file_data: FileData, **kwargs: Any) -> DownloadResponse:
-        return await super().run_async(file_data=file_data, **kwargs)
-
 
 class SftpUploaderConfig(FsspecUploaderConfig):
     pass
@@ -153,22 +155,6 @@ class SftpUploader(FsspecUploader):
     connection_config: SftpConnectionConfig
     upload_config: SftpUploaderConfig = field(default=None)
 
-    @requires_dependencies(["paramiko", "fsspec"], extras="sftp")
-    def __post_init__(self):
-        super().__post_init__()
-
-    @requires_dependencies(["paramiko", "fsspec"], extras="sftp")
-    def precheck(self) -> None:
-        super().precheck()
-
-    @requires_dependencies(["paramiko", "fsspec"], extras="sftp")
-    def run(self, path: Path, file_data: FileData, **kwargs: Any) -> None:
-        return super().run(path=path, file_data=file_data, **kwargs)
-
-    @requires_dependencies(["paramiko", "fsspec"], extras="sftp")
-    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)
-
 
 sftp_source_entry = SourceRegistryEntry(
     indexer=SftpIndexer,