[FEAT] Create obstore store in fsspec on demand

obstore/python/obstore/fsspec.py

@@ -22,15 +22,17 @@
 
 from __future__ import annotations
 
+from urllib.parse import urlparse
 import asyncio
 from collections import defaultdict
+from functools import lru_cache
 from typing import Any, Coroutine, Dict, List, Tuple
 
 import fsspec.asyn
 import fsspec.spec
 
 import obstore as obs
-
+from obstore.store import from_url
 
 class AsyncFsspecStore(fsspec.asyn.AsyncFileSystem):
     """An fsspec implementation based on a obstore Store.
@@ -43,8 +45,10 @@ class AsyncFsspecStore(fsspec.asyn.AsyncFileSystem):
 
     def __init__(
         self,
-        store: obs.store.ObjectStore,
         *args,
+        config: dict[str, Any] = {},
+        client_options: dict[str, Any] = {},
+        retry_config: dict[str, Any] = {},
         asynchronous: bool = False,
         loop=None,
         batch_size: int | None = None,
@@ -75,26 +79,79 @@ def __init__(
         ```
         """
 
-        self.store = store
+        self.config = config
+        self.client_options = client_options
+        self.retry_config = retry_config
+
         super().__init__(
             *args, asynchronous=asynchronous, loop=loop, batch_size=batch_size
         )
 
+    def _split_path(self, path: str) -> Tuple[str, str]:
+        """
+        Split bucket and file path
+
+        Args:
+            path  (str): Input path, like `s3://mybucket/path/to/file`
+
+        Examples:
+            >>> split_path("s3://mybucket/path/to/file")
+            ['mybucket', 'path/to/file']
+        """
+
+        protocol_with_bucket = ["s3", "s3a", "gcs", "gs", "abfs"]
+
+        if not self.protocol in protocol_with_bucket:
+            # no bucket name in path
+            return "", path
+
+        res = urlparse(path)
+        if res.scheme:
+            if res.scheme != self.protocol:
+                raise ValueError(f"Expect protocol to be {self.protocol}. Got {res.schema}")
+            path = res.netloc + res.path
+
+        if "/" not in path:
+            return path, ""
+        else:
+            path_li = path.split("/")
+            bucket = path_li[0]
+            file_path = "/".join(path_li[1:])
+            return (bucket, file_path)
+
+    @lru_cache(maxsize=10)
+    def _construct_store(self, bucket: str):
+        return from_url(
+            url=f"{self.protocol}://{bucket}",
+            config=self.config,
+            client_options=self.client_options,
+            retry_config=self.retry_config if self.retry_config else None,
+        )
+
     async def _rm_file(self, path, **kwargs):
-        return await obs.delete_async(self.store, path)
+        bucket, path = self._split_path(path)
+        store = self._construct_store(bucket)
+        return await obs.delete_async(store, path)
 
     async def _cp_file(self, path1, path2, **kwargs):
-        return await obs.copy_async(self.store, path1, path2)
+        bucket, path = self._split_path(path)
+        store = self._construct_store(bucket)
+        return await obs.copy_async(store, path1, path2)
 
     async def _pipe_file(self, path, value, **kwargs):
-        return await obs.put_async(self.store, path, value)
+        bucket, path = self._split_path(path)
+        store = self._construct_store(bucket)
+        return await obs.put_async(store, path, value)
 
     async def _cat_file(self, path, start=None, end=None, **kwargs):
+        bucket, path = self._split_path(path)
+        store = self._construct_store(bucket)
+
         if start is None and end is None:
-            resp = await obs.get_async(self.store, path)
-            return await resp.bytes_async()
+            resp = await obs.get_async(store, path)
+            return (await resp.bytes_async()).to_bytes()
 
-        range_bytes = await obs.get_range_async(self.store, path, start=start, end=end)
+        range_bytes = await obs.get_range_async(store, path, start=start, end=end)
         return range_bytes.to_bytes()
 
     async def _cat_ranges(
@@ -107,6 +164,9 @@ async def _cat_ranges(
         on_error="return",
         **kwargs,
     ):
+        bucket, path = self._split_path(path)
+        store = self._construct_store(bucket)
+
         if isinstance(starts, int):
             starts = [starts] * len(paths)
         if isinstance(ends, int):
@@ -122,7 +182,7 @@ async def _cat_ranges(
         for path, ranges in per_file_requests.items():
             offsets = [r[0] for r in ranges]
             ends = [r[1] for r in ranges]
-            fut = obs.get_ranges_async(self.store, path, starts=offsets, ends=ends)
+            fut = obs.get_ranges_async(store, path, starts=offsets, ends=ends)
             futs.append(fut)
 
         result = await asyncio.gather(*futs)
@@ -137,17 +197,24 @@ async def _cat_ranges(
         return output_buffers
 
     async def _put_file(self, lpath, rpath, **kwargs):
+        bucket, path = self._split_path(path)
+        store = self._construct_store(bucket)
         with open(lpath, "rb") as f:
-            await obs.put_async(self.store, rpath, f)
+            await obs.put_async(store, rpath, f)
 
     async def _get_file(self, rpath, lpath, **kwargs):
+        bucket, path = self._split_path(path)
+        store = self._construct_store(bucket)
         with open(lpath, "wb") as f:
-            resp = await obs.get_async(self.store, rpath)
+            resp = await obs.get_async(store, rpath)
             async for buffer in resp.stream():
                 f.write(buffer)
 
     async def _info(self, path, **kwargs):
-        head = await obs.head_async(self.store, path)
+        bucket, path = self._split_path(path)
+        store = self._construct_store(bucket)
+
+        head = await obs.head_async(store, path)
         return {
             # Required of `info`: (?)
             "name": head["path"],
@@ -160,7 +227,10 @@ async def _info(self, path, **kwargs):
         }
 
     async def _ls(self, path, detail=True, **kwargs):
-        result = await obs.list_with_delimiter_async(self.store, path)
+        bucket, path = self._split_path(path)
+        store = self._construct_store(bucket)
+
+        result = await obs.list_with_delimiter_async(store, path)
         objects = result["objects"]
         prefs = result["common_prefixes"]
         if detail:
@@ -178,6 +248,9 @@ async def _ls(self, path, detail=True, **kwargs):
 
     def _open(self, path, mode="rb", **kwargs):
         """Return raw bytes-mode file-like from the file-system"""
+        bucket, path = self._split_path(path)
+        store = self._construct_store(bucket)
+
         return BufferedFileSimple(self, path, mode, **kwargs)
 
 
@@ -201,3 +274,13 @@ def read(self, length: int = -1):
             data = self.fs.cat_file(self.path, self.loc, self.loc + length)
             self.loc += length
         return data
+
+
+class S3FsspecStore(AsyncFsspecStore):
+    protocol = "s3"
+
+class GCSFsspecStore(AsyncFsspecStore):
+    protocol = "gs"
+
+class AzureFsspecStore(AsyncFsspecStore):
+    protocol = "abfs"