Adding async variants to various FusionFS methods. (#109)

Completing async variants of FFS methods with some tests and documentation

Author: dialectician4

Cargo.lock

@@ -1,6 +1,6 @@
 [package]
 name = "pyfusion"
-version = "2.0.8"
+version = "2.0.9-dev0"
 edition = "2021"
 
 

Cargo.toml

@@ -2,7 +2,7 @@
 
 __author__ = """Fusion Devs"""
 __email__ = "fusion_developers@jpmorgan.com"
-__version__ = "2.0.8"
+__version__ = "2.0.9-dev0"
 
 from fusion._fusion import FusionCredentials
 from fusion.fs_sync import fsync

docs/fusion_filesystem.ipynb

@@ -46,29 +46,29 @@ class FusionEmbeddingsConnection(Connection):  # type: ignore
     - Provides methods for modifying and validating URLs specific to the Fusion Embedding API.
 
     Args:
-        http_auth (str or tuple, optional): HTTP auth information as either ':' separated string or a tuple. 
+        http_auth (str or tuple, optional): HTTP auth information as either ':' separated string or a tuple.
             Any value will be passed into requests as `auth`.
         use_ssl (bool, optional): Use SSL for the connection if `True`.
         verify_certs (bool, optional): Whether to verify SSL certificates.
         ssl_show_warn (bool, optional): Show warning when verify certs is disabled.
-        ca_certs (str, optional): Path to CA bundle. Defaults to configured OpenSSL bundles from environment variables 
-            and then certifi before falling back to the standard requests bundle to improve consistency with other 
+        ca_certs (str, optional): Path to CA bundle. Defaults to configured OpenSSL bundles from environment variables
+            and then certifi before falling back to the standard requests bundle to improve consistency with other
             Connection implementations.
-        client_cert (str, optional): Path to the file containing the private key and the certificate, 
+        client_cert (str, optional): Path to the file containing the private key and the certificate,
             or cert only if using client_key.
-        client_key (str, optional): Path to the file containing the private key if using separate cert and key files 
+        client_key (str, optional): Path to the file containing the private key if using separate cert and key files
             (client_cert will contain only the cert).
         headers (dict, optional): Any custom HTTP headers to be added to requests.
         http_compress (bool, optional): Use gzip compression.
-        opaque_id (str, optional): Send this value in the 'X-Opaque-Id' HTTP header for tracing 
+        opaque_id (str, optional): Send this value in the 'X-Opaque-Id' HTTP header for tracing
             all requests made by this transport.
-        pool_maxsize (int, optional): Maximum connection pool size used by pool-manager 
+        pool_maxsize (int, optional): Maximum connection pool size used by pool-manager
             for custom connection-pooling on current session.
-        metrics (Metrics, optional): Instance of a subclass of the `opensearchpy.Metrics` class, 
+        metrics (Metrics, optional): Instance of a subclass of the `opensearchpy.Metrics` class,
             used for collecting and reporting metrics related to the client's operations.
 
     Keyword Args:
-        root_url (str, optional): Root URL for the Fusion API. Defaults to 
+        root_url (str, optional): Root URL for the Fusion API. Defaults to
             "https://fusion.jpmorgan.com/api/v1/".
         credentials (FusionCredentials or str, optional): Credentials for the Fusion API. Can be a `FusionCredentials`
             object or a path to a credentials file. Defaults to "config/client_credentials.json".

py_src/fusion/__init__.py

@@ -215,13 +215,19 @@ def _use_catalog(self, catalog: str | None) -> str:
 
         return catalog
 
-    def get_fusion_filesystem(self) -> FusionHTTPFileSystem:
-        """Creates Fusion Filesystem.
+    def get_fusion_filesystem(self, **kwargs: Any) -> FusionHTTPFileSystem:
+        """Retrieve Fusion file system instance.
+
+        Note: This function always returns a reference to the exact same FFS instance since
+        an FFS instance is based off the FusionCredentials object.
 
         Returns: Fusion Filesystem
 
         """
-        return FusionHTTPFileSystem(client_kwargs={"root_url": self.root_url, "credentials": self.credentials})
+        as_async = kwargs.get("asynchronous", False)
+        return FusionHTTPFileSystem(
+            asynchronous=as_async, client_kwargs={"root_url": self.root_url, "credentials": self.credentials}
+        )
 
     def list_catalogs(self, output: bool = False) -> pd.DataFrame:
         """Lists the catalogs available to the API account.

py_src/fusion/embeddings.py

@@ -86,8 +86,21 @@ async def _async_raise_not_found_for_status(self, response: Any, url: str) -> No
             finally:
                 self._raise_not_found_for_status(response, url)
 
+    def _check_session_open(self) -> bool:
+        # Check that _session is active. Expects that if _session is populated with .set_session, result
+        # result was already awaited
+        if self._session is None:
+            return False
+        return not self._session.closed
+
+    async def _async_startup(self) -> None:
+        await self.set_session()
+        if not self._check_session_open():
+            raise RuntimeError("FusionFS session closed before operation")
+
     async def _decorate_url_a(self, url: str) -> str:
         url = urljoin(f"{self.client_kwargs['root_url']}catalogs/", url) if "http" not in url else url
+        url = url[:-1] if url[-1] == "/" else url
         return url
 
     def _decorate_url(self, url: str) -> str:
@@ -105,6 +118,16 @@ async def _isdir(self, path: str) -> bool:
             return False
 
     async def _changes(self, url: str) -> dict[Any, Any]:
+        """Get from given url.
+
+        Currently called within the context of the /changes api endpoint.
+
+        Args:
+            url: str
+
+        Returns:
+            Dict containing json-ified return from endpoint.
+        """
         url = self._decorate_url(url)
         try:
             session = await self.set_session()
@@ -120,7 +143,7 @@ async def _changes(self, url: str) -> dict[Any, Any]:
             logger.log(VERBOSE_LVL, f"Artificial error, {ex}")
             raise ex
 
-    async def _ls_real(self, url: str, detail: bool = True, **kwargs: Any) -> Any:
+    async def _ls_real(self, url: str, detail: bool = False, **kwargs: Any) -> Any:
         # ignoring URL-encoded arguments
         clean_url = url
         if "http" not in url:
@@ -195,9 +218,21 @@ def info(self, path: str, **kwargs: Any) -> Any:
 
         return res
 
-    async def _ls(self, url: str, detail: bool = True, **kwargs: Any) -> Any:
-        url = await self._decorate_url_a(url)
-        return await super()._ls(url, detail, **kwargs)
+    async def _info(self, path: str, **kwargs: Any) -> Any:
+        await self._async_startup()
+        path = await self._decorate_url_a(path)
+        kwargs["keep_protocol"] = True
+        res = await super()._ls(path, detail=True, **kwargs)
+        if res[0]["type"] != "file":
+            res = await super()._info(path, **kwargs)
+            if path.split("/")[-2] == "datasets":
+                target = path.split("/")[-1]
+                args = ["/".join(path.split("/")[:-1]) + f"/changes?datasets={quote(target)}"]
+                res["changes"] = await self._changes(*args)
+        split_path = path.split("/")
+        if len(split_path) > 1 and split_path[-2] == "distributions":
+            res = res[0]
+        return res
 
     def ls(self, url: str, detail: bool = False, **kwargs: Any) -> Any:
         """List resources.
@@ -223,7 +258,21 @@ def ls(self, url: str, detail: bool = False, **kwargs: Any) -> Any:
 
         return ret
 
-    def exists(self, url: str, detail: bool = True, **kwargs: Any) -> Any:
+    async def _ls(self, url: str, detail: bool = False, **kwargs: Any) -> Any:
+        await self._async_startup()
+        url = await self._decorate_url_a(url)
+        ret = await super()._ls(url, detail, **kwargs)
+        keep_protocol = kwargs.pop("keep_protocol", False)
+        if detail:
+            if not keep_protocol:
+                for k in ret:
+                    k["name"] = k["name"].split(f"{self.client_kwargs['root_url']}catalogs/")[-1]
+        elif not keep_protocol:
+            return [x.split(f"{self.client_kwargs['root_url']}catalogs/")[-1] for x in ret]
+
+        return ret
+
+    def exists(self, url: str, **kwargs: Any) -> Any:
         """Check existence.
 
         Args:
@@ -235,7 +284,13 @@ def exists(self, url: str, detail: bool = True, **kwargs: Any) -> Any:
 
         """
         url = self._decorate_url(url)
-        return super().exists(url, detail, **kwargs)
+        return super().exists(url, **kwargs)
+
+    async def _exists(self, url: str, **kwargs: Any) -> Any:
+        await self._async_startup()
+        url = self._decorate_url(url)
+        out = await super()._exists(url, **kwargs)
+        return out
 
     def isfile(self, path: str) -> Union[bool, Any]:
         """Is path a file.
@@ -264,6 +319,12 @@ def cat(self, url: str, start: Optional[int] = None, end: Optional[int] = None,
         url = self._decorate_url(url)
         return super().cat(url, start=start, end=end, **kwargs)
 
+    async def _cat(self, url: str, start: Optional[int] = None, end: Optional[int] = None, **kwargs: Any) -> Any:
+        await self._async_startup()
+        url = self._decorate_url(url)
+        out = await super()._cat(url, start=start, end=end, **kwargs)
+        return out
+
     async def _fetch_range(
         self,
         session: aiohttp.ClientSession,
@@ -836,6 +897,12 @@ def find(self, path: str, maxdepth: Optional[int] = None, withdirs: bool = False
         path = self._decorate_url(path)
         return super().find(path, maxdepth=maxdepth, withdirs=withdirs, **kwargs)
 
+    async def _find(self, path: str, maxdepth: Optional[int] = None, withdirs: bool = False, **kwargs: Any) -> Any:
+        await self._async_startup()
+        path = self._decorate_url(path)
+        out = await super()._find(path, maxdepth=maxdepth, withdirs=withdirs, **kwargs)
+        return out
+
     def glob(self, path: str, **kwargs: Any) -> Any:
         """Glob.
 
@@ -849,6 +916,10 @@ def glob(self, path: str, **kwargs: Any) -> Any:
 
         return super().glob(path, **kwargs)
 
+    async def _glob(self, path: str, **kwargs: Any) -> Any:
+        out = await super()._glob(path, **kwargs)
+        return out
+
     def open(
         self,
         path: str,

py_src/fusion/fusion.py

@@ -101,19 +101,78 @@ async def test_decorate_url_with_http_async(http_fs_instance: FusionHTTPFileSyst
 @pytest.mark.asyncio
 async def test_isdir_true(http_fs_instance: FusionHTTPFileSystem) -> None:
     http_fs_instance._decorate_url = AsyncMock(return_value="decorated_path_dir")  # type: ignore
-    http_fs_instance._info = AsyncMock(return_value={"type": "directory"})
+    http_fs_instance._info = AsyncMock(return_value={"type": "directory"})  # type: ignore
     result = await http_fs_instance._isdir("path_dir")
     assert result
 
 
 @pytest.mark.asyncio
 async def test_isdir_false(http_fs_instance: FusionHTTPFileSystem) -> None:
     http_fs_instance._decorate_url = AsyncMock(return_value="decorated_path_file")  # type: ignore
-    http_fs_instance._info = AsyncMock(return_value={"type": "file"})
+    http_fs_instance._info = AsyncMock(return_value={"type": "file"})  # type: ignore
     result = await http_fs_instance._isdir("path_file")
     assert not result
 
 
+@pytest.mark.asyncio
+async def test_check_sess_open(http_fs_instance: FusionHTTPFileSystem) -> None:
+    new_fs_session_closed = not http_fs_instance._check_session_open()
+    assert new_fs_session_closed
+
+    # Corresponds to running .set_session()
+    session_mock = MagicMock()
+    session_mock.closed = False
+    http_fs_instance._session = session_mock
+    fs_session_open = http_fs_instance._check_session_open()
+    assert fs_session_open
+
+    # Corresponds to situation where session was closed
+    session_mock2 = MagicMock()
+    session_mock2.closed = True
+    http_fs_instance._session = session_mock2
+    fs_session_closed = not http_fs_instance._check_session_open()
+    assert fs_session_closed
+
+
+@pytest.mark.asyncio
+async def test_async_startup(http_fs_instance: FusionHTTPFileSystem) -> None:
+    http_fs_instance._session = None
+    with (
+        patch("fusion.fusion_filesystem.FusionHTTPFileSystem.set_session") as SetSessionMock,
+        pytest.raises(RuntimeError) as re,
+    ):
+        await http_fs_instance._async_startup()
+    SetSessionMock.assert_called_once()
+    assert re.match("FusionFS session closed before operation")
+
+    # Mock an open session
+    MockClient = MagicMock()
+    MockClient.closed = False
+    http_fs_instance._session = MockClient
+    with patch("fusion.fusion_filesystem.FusionHTTPFileSystem.set_session") as SetSessionMock2:
+        await http_fs_instance._async_startup()
+    SetSessionMock2.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_exists_methods(http_fs_instance: FusionHTTPFileSystem) -> None:
+    with patch("fusion.fusion_filesystem.HTTPFileSystem.exists") as MockExists:
+        MockExists.return_value = True
+        exists_out = http_fs_instance.exists("dummy_path")
+        MockExists.assert_called_once()
+    assert exists_out
+
+    with (
+        patch("fusion.fusion_filesystem.HTTPFileSystem._exists") as Mock_Exists,
+        patch("fusion.fusion_filesystem.FusionHTTPFileSystem._async_startup") as MockStartup,
+    ):
+        Mock_Exists.return_value = True
+        _exists_out = await http_fs_instance._exists("dummy_path")
+        Mock_Exists.assert_awaited_once()
+        MockStartup.assert_awaited_once()
+    assert _exists_out
+
+
 @patch("requests.Session")
 def test_stream_single_file(mock_session_class: MagicMock, example_creds_dict: dict[str, Any], tmp_path: Path) -> None:
     url = "http://example.com/data"

py_src/fusion/fusion_filesystem.py

@@ -1,6 +1,6 @@
 [project]
 name = "pyfusion"
-version = "2.0.8"
+version = "2.0.9-dev0"
 
 homepage = "https://github.com/jpmorganchase/fusion"
 description = "JPMC Fusion Developer Tools"
@@ -241,7 +241,7 @@ exclude_lines = [
 
 
 [tool.bumpversion]
-current_version = "2.0.8"
+current_version = "2.0.9-dev0"
 parse = '(?P<major>\d+)\.(?P<minor>\d+)\.(?P<patch>\d+)(?:-(?P<release>[a-z]+)(?P<candidate>\d+))?'
 serialize = [
     '{major}.{minor}.{patch}-{release}{candidate}',