datachain 0.14.2__py3-none-any.whl → 0.39.0__py3-none-any.whl

datachain/client/gcs.py

@@ -3,7 +3,7 @@ import json
 import os
 from collections.abc import Iterable
 from datetime import datetime
-from typing import Any, Optional, cast
+from typing import Any, cast
 
 from dateutil.parser import isoparse
 from gcsfs import GCSFileSystem
@@ -15,7 +15,7 @@ from .fsspec import DELIMITER, Client, ResultQueue
 
 # Patch gcsfs for consistency with s3fs
 GCSFileSystem.set_session = GCSFileSystem._set_session
-PageQueue = asyncio.Queue[Optional[Iterable[dict[str, Any]]]]
+PageQueue = asyncio.Queue[Iterable[dict[str, Any]] | None]
 
 
 class GCSClient(Client):
@@ -74,7 +74,7 @@ class GCSClient(Client):
             try:
                 await self._get_pages(prefix, page_queue)
                 found = await consumer
-                if not found:
+                if not found and prefix:
                     raise FileNotFoundError(f"Unable to resolve remote path: {prefix}")
             finally:
                 consumer.cancel()  # In case _get_pages() raised
@@ -115,7 +115,7 @@ class GCSClient(Client):
                     maxResults=page_size,
                     pageToken=next_page_token,
                     json_out=True,
-                    versions="true",
+                    versions="true" if self._is_version_aware() else "false",
                 )
                 assert page["kind"] == "storage#objects"
                 await page_queue.put(page.get("items", []))
@@ -134,12 +134,12 @@ class GCSClient(Client):
             source=self.uri,
             path=path,
             etag=v.get("etag", ""),
-            version=v.get("generation", ""),
+            version=v.get("generation", "") if self._is_version_aware() else "",
             is_latest=not v.get("timeDeleted"),
             last_modified=self.parse_timestamp(v["updated"]),
             size=v.get("size", ""),
         )
 
     @classmethod
-    def version_path(cls, path: str, version_id: Optional[str]) -> str:
+    def version_path(cls, path: str, version_id: str | None) -> str:
         return f"{path}#{version_id}" if version_id else path

datachain/client/hf.py

@@ -15,6 +15,34 @@ class classproperty:  # noqa: N801
         return self.fget(owner)
 
 
+def _wrap_class(sync_fs_class):
+    """
+    Analog of `AsyncFileSystemWrapper.wrap_class` from fsspec, but sets
+    asynchronous to False by default. This is similar to other Async FS
+    we initialize. E.g. it means we don't break things in Jupyter where code
+    run in async.
+
+    This also fixes write operations by ensuring they are properly forwarded
+    to the underlying filesystem without async buffering issues.
+    """
+    from fsspec.implementations.asyn_wrapper import AsyncFileSystemWrapper
+
+    class GeneratedAsyncFileSystemWrapper(AsyncFileSystemWrapper):
+        def __init__(self, *args, **kwargs):
+            sync_fs = sync_fs_class(*args, **kwargs)
+            super().__init__(sync_fs, asynchronous=False)
+
+        def open(self, path, mode="rb", **kwargs):
+            # Override open to ensure write operations work correctly.
+            # It seems to be a bug in the fsspec wrapper. It avoids
+            # wrapping open() explicitly but also doesn't redirect it to
+            # sync filesystem.
+            return self.sync_fs.open(path, mode, **kwargs)
+
+    GeneratedAsyncFileSystemWrapper.__name__ = f"Async{sync_fs_class.__name__}Wrapper"
+    return GeneratedAsyncFileSystemWrapper
+
+
 @functools.cache
 def get_hf_filesystem_cls():
     import fsspec
@@ -29,10 +57,9 @@ def get_hf_filesystem_cls():
             f"{fsspec_version} is installed."
         )
 
-    from fsspec.implementations.asyn_wrapper import AsyncFileSystemWrapper
     from huggingface_hub import HfFileSystem
 
-    fs_cls = AsyncFileSystemWrapper.wrap_class(HfFileSystem)
+    fs_cls = _wrap_class(HfFileSystem)
     # AsyncFileSystemWrapper does not set class properties, so we need to set them back.
     fs_cls.protocol = HfFileSystem.protocol
     return fs_cls

datachain/client/http.py

@@ -0,0 +1,157 @@
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING, Any, ClassVar, cast
+from urllib.parse import urlparse
+
+from fsspec.implementations.http import HTTPFileSystem
+
+from datachain.dataset import StorageURI
+from datachain.lib.file import File
+
+from .fsspec import Client
+
+if TYPE_CHECKING:
+    from datachain.cache import Cache
+
+
+class HTTPClient(Client):
+    FS_CLASS = HTTPFileSystem
+    PREFIX: ClassVar[str] = "http://"
+    protocol: ClassVar[str] = "http"
+
+    @classmethod
+    def create_fs(cls, **kwargs) -> HTTPFileSystem:
+        # Configure HTTPFileSystem options
+        kwargs.setdefault("simple_links", True)
+        kwargs.setdefault("same_scheme", True)
+        kwargs.setdefault("cache_type", "bytes")
+
+        kwargs.pop("version_aware", None)
+
+        fs = cls.FS_CLASS(**kwargs)
+        fs.invalidate_cache()
+        return cast("HTTPFileSystem", fs)
+
+    @classmethod
+    def from_name(
+        cls,
+        name: str,
+        cache: "Cache",
+        kwargs: dict[str, Any],
+    ) -> "HTTPClient":
+        parsed = urlparse(name)
+
+        if parsed.scheme:
+            name = parsed.netloc + parsed.path
+
+        return cls(name, kwargs, cache)
+
+    @classmethod
+    def split_url(cls, url: str) -> tuple[str, str]:
+        """Split HTTP/HTTPS URL into domain (bucket equivalent) and path."""
+        parsed = urlparse(url)
+        domain = parsed.netloc
+        path = parsed.path.lstrip("/")
+
+        if parsed.query:
+            path += f"?{parsed.query}"
+        if parsed.fragment:
+            path += f"#{parsed.fragment}"
+
+        return domain, path
+
+    @classmethod
+    def get_uri(cls, name: str) -> "StorageURI":
+        if not name.startswith(("http://", "https://")):
+            return StorageURI(f"{cls.PREFIX}{name}")
+        return StorageURI(name)
+
+    @classmethod
+    def is_root_url(cls, url: str) -> bool:
+        parsed = urlparse(url)
+        return parsed.path in ("", "/") and not parsed.query and not parsed.fragment
+
+    def get_full_path(self, rel_path: str, version_id: str | None = None) -> str:
+        if self.name.startswith(("http://", "https://")):
+            base_url = self.name
+        else:
+            if rel_path and "/" in rel_path:
+                first_part = rel_path.split("/")[0]
+                if "." in first_part and not first_part.startswith("."):
+                    return f"{self.protocol}://{rel_path}"
+
+            base_url = f"{self.protocol}://{self.name}"
+
+        if rel_path:
+            if not base_url.endswith("/") and not rel_path.startswith("/"):
+                base_url += "/"
+            full_url = base_url + rel_path
+        else:
+            full_url = base_url
+
+        return full_url
+
+    def url(self, path: str, expires: int = 3600, **kwargs) -> str:
+        """
+        Generate URL for the given path.
+        Note: HTTP URLs don't support signed/expiring URLs.
+        """
+        return self.get_full_path(path, kwargs.pop("version_id", None))
+
+    def info_to_file(self, v: dict[str, Any], path: str) -> File:
+        etag = v.get("ETag", "").strip('"')
+        last_modified = v.get("last_modified")
+        if last_modified:
+            if isinstance(last_modified, str):
+                try:
+                    from email.utils import parsedate_to_datetime
+
+                    last_modified = parsedate_to_datetime(last_modified)
+                except (ValueError, TypeError):
+                    last_modified = datetime.now(timezone.utc)
+            elif isinstance(last_modified, (int, float)):
+                last_modified = datetime.fromtimestamp(last_modified, timezone.utc)
+        else:
+            last_modified = datetime.now(timezone.utc)
+
+        return File(
+            source=self.uri,
+            path=path,
+            size=v.get("size", 0),
+            etag=etag,
+            version="",
+            is_latest=True,
+            last_modified=last_modified,
+        )
+
+    def upload(self, data: bytes, path: str) -> "File":
+        raise NotImplementedError(
+            "HTTP/HTTPS client is read-only. Upload operations are not supported."
+        )
+
+    def get_file_info(self, path: str, version_id: str | None = None) -> "File":
+        info = self.fs.info(self.get_full_path(path))
+        return self.info_to_file(info, path)
+
+    def open_object(self, file: "File", use_cache: bool = True, cb=None):
+        from datachain.client.fileslice import FileWrapper
+
+        if use_cache and (cache_path := self.cache.get_path(file)):
+            return open(cache_path, mode="rb")
+
+        assert not file.location
+        return FileWrapper(
+            self.fs.open(self.get_full_path(file.get_path_normalized())),
+            cb or (lambda x: None),
+        )
+
+    async def get_file(self, lpath, rpath, callback, version_id: str | None = None):
+        return await self.fs._get_file(lpath, rpath, callback=callback)
+
+    async def _fetch_dir(self, prefix: str, pbar, result_queue) -> set[str]:
+        full_url = self.get_full_path(prefix)
+        raise NotImplementedError(f"Cannot download file from {full_url}")
+
+
+class HTTPSClient(HTTPClient):
+    protocol = "https"
+    PREFIX = "https://"

datachain/client/local.py

@@ -2,14 +2,14 @@ import os
 import posixpath
 from datetime import datetime, timezone
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING, Any
 from urllib.parse import urlparse
 
 from fsspec.implementations.local import LocalFileSystem
 
 from datachain.lib.file import File
 
-from .fsspec import Client
+from .fsspec import Client, is_win_local_path
 
 if TYPE_CHECKING:
     from datachain.cache import Cache
@@ -57,9 +57,13 @@ class FileClient(Client):
             /home/user/animals/ -> file:///home/user/animals/
             C:\\windows\animals -> file:///C:/windows/animals
         """
+        parsed = urlparse(path)
+        if parsed.scheme and not is_win_local_path(path):
+            return path
+
         uri = Path(path).expanduser().absolute().resolve().as_uri()
-        if path[-1] == os.sep:
-            # we should keep os separator from the end of the path
+        if path and path[-1] in (os.sep, "/"):
+            # keep trailing separator so directory URIs stay rooted
             uri += "/"  # in uri (file:///...) all separators are / regardless of os
 
         return uri
@@ -99,13 +103,13 @@ class FileClient(Client):
         )
 
     async def get_current_etag(self, file: "File") -> str:
-        info = self.fs.info(self.get_full_path(file.path))
+        info = self.fs.info(self.get_full_path(file.get_path_normalized()))
         return self.info_to_file(info, "").etag
 
-    async def get_size(self, path: str, version_id: Optional[str] = None) -> int:
+    async def get_size(self, path: str, version_id: str | None = None) -> int:
         return self.fs.size(path)
 
-    async def get_file(self, lpath, rpath, callback, version_id: Optional[str] = None):
+    async def get_file(self, lpath, rpath, callback, version_id: str | None = None):
         return self.fs.get_file(lpath, rpath, callback=callback)
 
     async def ls_dir(self, path):
@@ -114,7 +118,7 @@ class FileClient(Client):
     def rel_path(self, path):
         return posixpath.relpath(path, self.name)
 
-    def get_full_path(self, rel_path, version_id: Optional[str] = None):
+    def get_full_path(self, rel_path, version_id: str | None = None):
         full_path = Path(self.name, rel_path).as_posix()
         if rel_path.endswith("/") or not rel_path:
             full_path += "/"
@@ -138,8 +142,8 @@ class FileClient(Client):
         if not self.use_symlinks:
             super().fetch_nodes(nodes, shared_progress_bar)
 
-    def do_instantiate_object(self, uid, dst):
+    def do_instantiate_object(self, file: File, dst: str) -> None:
         if self.use_symlinks:
-            os.symlink(Path(self.name, uid.path), dst)
+            os.symlink(Path(self.name, file.path), dst)
         else:
-            super().do_instantiate_object(uid, dst)
+            super().do_instantiate_object(file, dst)

datachain/client/s3.py

@@ -1,6 +1,6 @@
 import asyncio
 import os
-from typing import Any, Optional, cast
+from typing import Any, cast
 from urllib.parse import parse_qs, urlsplit, urlunsplit
 
 from botocore.exceptions import NoCredentialsError
@@ -80,7 +80,7 @@ class ClientS3(Client):
             finally:
                 await page_queue.put(None)
 
-        async def process_pages(page_queue, result_queue):
+        async def process_pages(page_queue, result_queue, prefix):
             found = False
             with tqdm(desc=f"Listing {self.uri}", unit=" objects", leave=False) as pbar:
                 while (res := await page_queue.get()) is not None:
@@ -94,14 +94,14 @@ class ClientS3(Client):
                     if entries:
                         await result_queue.put(entries)
                         pbar.update(len(entries))
-            if not found:
+            if not found and prefix:
                 raise FileNotFoundError(f"Unable to resolve remote path: {prefix}")
 
         try:
             prefix = start_prefix
             if prefix:
                 prefix = prefix.lstrip(DELIMITER) + DELIMITER
-            versions = True
+            versions = self._is_version_aware()
             fs = self.fs
             await fs.set_session()
             s3 = await fs.get_s3(self.name)
@@ -118,7 +118,9 @@ class ClientS3(Client):
                 Delimiter="",
             )
             page_queue: asyncio.Queue[list] = asyncio.Queue(2)
-            consumer = asyncio.create_task(process_pages(page_queue, result_queue))
+            consumer = asyncio.create_task(
+                process_pages(page_queue, result_queue, prefix)
+            )
             try:
                 await get_pages(it, page_queue)
                 await consumer
@@ -137,14 +139,16 @@ class ClientS3(Client):
             source=self.uri,
             path=v["Key"],
             etag=v.get("ETag", "").strip('"'),
-            version=ClientS3.clean_s3_version(v.get("VersionId", "")),
+            version=(
+                ClientS3.clean_s3_version(v.get("VersionId", "")) if versions else ""
+            ),
             is_latest=v.get("IsLatest", True),
             last_modified=v.get("LastModified", ""),
             size=v["Size"],
         )
 
     @classmethod
-    def version_path(cls, path: str, version_id: Optional[str]) -> str:
+    def version_path(cls, path: str, version_id: str | None) -> str:
         parts = list(urlsplit(path))
         query = parse_qs(parts[3])
         if "versionId" in query:
@@ -183,7 +187,7 @@ class ClientS3(Client):
         return subdirs
 
     @staticmethod
-    def clean_s3_version(ver: Optional[str]) -> str:
+    def clean_s3_version(ver: str | None) -> str:
         return ver if (ver is not None and ver != "null") else ""
 
     def info_to_file(self, v: dict[str, Any], path: str) -> File:
@@ -191,7 +195,11 @@ class ClientS3(Client):
             source=self.uri,
             path=path,
             size=v["size"],
-            version=ClientS3.clean_s3_version(v.get("VersionId", "")),
+            version=(
+                ClientS3.clean_s3_version(v.get("VersionId", ""))
+                if self._is_version_aware()
+                else ""
+            ),
             etag=v.get("ETag", "").strip('"'),
             is_latest=v.get("IsLatest", True),
             last_modified=v.get("LastModified", ""),

datachain/config.py

@@ -1,7 +1,6 @@
 from collections.abc import Mapping
 from contextlib import contextmanager
 from enum import Enum
-from typing import Optional, Union
 
 from tomlkit import TOMLDocument, dump, load
 
@@ -22,16 +21,13 @@ class Config:
     # In the order of precedence
     LEVELS = SYSTEM_LEVELS + LOCAL_LEVELS
 
-    def __init__(
-        self,
-        level: Optional[ConfigLevel] = None,
-    ):
+    def __init__(self, level: ConfigLevel | None = None):
         self.level = level
 
         self.init()
 
     @classmethod
-    def get_dir(cls, level: Optional[ConfigLevel]) -> str:
+    def get_dir(cls, level: ConfigLevel | None) -> str:
         if level == ConfigLevel.SYSTEM:
             return system_config_dir()
         if level == ConfigLevel.GLOBAL:
@@ -43,7 +39,7 @@ class Config:
         d = DataChainDir(self.get_dir(self.level))
         d.init()
 
-    def load_one(self, level: Optional[ConfigLevel] = None) -> TOMLDocument:
+    def load_one(self, level: ConfigLevel | None = None) -> TOMLDocument:
         config_path = DataChainDir(self.get_dir(level)).config
 
         try:
@@ -128,7 +124,7 @@ class Config:
         return remote_conf
 
 
-def merge(into: Union[TOMLDocument, dict], update: Union[TOMLDocument, dict]):
+def merge(into: TOMLDocument | dict, update: TOMLDocument | dict):
     """Merges second dict into first recursively"""
     for key, val in update.items():
         if isinstance(into.get(key), dict) and isinstance(val, dict):

datachain/data_storage/db_engine.py

@@ -1,7 +1,7 @@
 import logging
 from abc import ABC, abstractmethod
 from collections.abc import Iterator
-from typing import TYPE_CHECKING, Any, ClassVar, Optional, Union
+from typing import TYPE_CHECKING, Any, ClassVar
 
 import sqlalchemy as sa
 from sqlalchemy.sql import FROM_LINTING
@@ -58,7 +58,7 @@ class DatabaseEngine(ABC, Serializable):
     @classmethod
     def compile_to_args(
         cls, statement: "ClauseElement", **kwargs
-    ) -> Union[tuple[str], tuple[str, dict[str, Any]]]:
+    ) -> tuple[str] | tuple[str, dict[str, Any]]:
         """
         Compile a sqlalchemy query or ddl object to an args tuple.
 
@@ -75,8 +75,8 @@ class DatabaseEngine(ABC, Serializable):
     def execute(
         self,
         query,
-        cursor: Optional[Any] = None,
-        conn: Optional[Any] = None,
+        cursor: Any | None = None,
+        conn: Any | None = None,
     ) -> Iterator[tuple[Any, ...]]: ...
 
     def get_table(self, name: str) -> "Table":
@@ -90,7 +90,7 @@ class DatabaseEngine(ABC, Serializable):
 
     @abstractmethod
     def executemany(
-        self, query, params, cursor: Optional[Any] = None
+        self, query, params, cursor: Any | None = None
     ) -> Iterator[tuple[Any, ...]]: ...
 
     @abstractmethod
@@ -112,7 +112,13 @@ class DatabaseEngine(ABC, Serializable):
         return sa.inspect(self.engine).has_table(name)
 
     @abstractmethod
-    def create_table(self, table: "Table", if_not_exists: bool = True) -> None: ...
+    def create_table(
+        self,
+        table: "Table",
+        if_not_exists: bool = True,
+        *,
+        kind: str | None = None,
+    ) -> None: ...
 
     @abstractmethod
     def drop_table(self, table: "Table", if_exists: bool = False) -> None: ...

datachain/data_storage/job.py

@@ -3,6 +3,8 @@ from enum import Enum
 
 class JobStatus(int, Enum):
     CREATED = 1
+    SCHEDULED = 10
+    PROVISIONING = 12
     QUEUED = 2
     INIT = 3
     RUNNING = 4
@@ -11,10 +13,12 @@ class JobStatus(int, Enum):
     CANCELING = 7
     CANCELED = 8
     CANCELING_SCHEDULED = 9
+    TASK = 11
+    PENDING = 13
 
     @classmethod
     def finished(cls) -> tuple[int, ...]:
-        return cls.COMPLETE, cls.FAILED, cls.CANCELED
+        return cls.COMPLETE, cls.FAILED, cls.CANCELED, cls.TASK
 
 
 class JobQueryType(int, Enum):