datachain 0.31.3__tar.gz → 0.32.0__tar.gz

{datachain-0.31.3 → datachain-0.32.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datachain
-Version: 0.31.3
+Version: 0.32.0
 Summary: Wrangle unstructured AI data at scale
 Author-email: Dmitry Petrov <support@dvc.org>
 License-Expression: Apache-2.0

{datachain-0.31.3 → datachain-0.32.0}/src/datachain/__init__.py

@@ -37,6 +37,7 @@ from datachain.lib.file import (
     VideoFrame,
 )
 from datachain.lib.model_store import ModelStore
+from datachain.lib.namespaces import delete as delete_namespace
 from datachain.lib.projects import create as create_project
 from datachain.lib.udf import Aggregator, Generator, Mapper
 from datachain.lib.utils import AbstractUDF, DataChainError
@@ -74,6 +75,7 @@ __all__ = [
     "create_project",
     "datasets",
     "delete_dataset",
+    "delete_namespace",
     "is_chain_type",
     "is_studio",
     "listings",

{datachain-0.31.3 → datachain-0.32.0}/src/datachain/catalog/catalog.py

@@ -144,26 +144,19 @@ def shutdown_process(
             return proc.wait()
 
 
-def process_output(stream: IO[bytes], callback: Callable[[str], None]) -> None:
+def _process_stream(stream: "IO[bytes]", callback: Callable[[str], None]) -> None:
     buffer = b""
+    while byt := stream.read(1):  # Read one byte at a time
+        buffer += byt
 
-    try:
-        while byt := stream.read(1):  # Read one byte at a time
-            buffer += byt
-
-            if byt in (b"\n", b"\r"):  # Check for newline or carriage return
-                line = buffer.decode("utf-8", errors="replace")
-                callback(line)
-                buffer = b""  # Clear buffer for the next line
-
-        if buffer:  # Handle any remaining data in the buffer
-            line = buffer.decode("utf-8", errors="replace")
+        if byt in (b"\n", b"\r"):  # Check for newline or carriage return
+            line = buffer.decode("utf-8")
             callback(line)
-    finally:
-        try:
-            stream.close()  # Ensure output is closed
-        except Exception:  # noqa: BLE001, S110
-            pass
+            buffer = b""  # Clear buffer for next line
+
+    if buffer:  # Handle any remaining data in the buffer
+        line = buffer.decode("utf-8")
+        callback(line)
 
 
 class DatasetRowsFetcher(NodesThreadPool):
@@ -1767,13 +1760,13 @@ class Catalog:
             recursive=recursive,
         )
 
-    @staticmethod
     def query(
+        self,
         query_script: str,
         env: Optional[Mapping[str, str]] = None,
         python_executable: str = sys.executable,
-        stdout_callback: Optional[Callable[[str], None]] = None,
-        stderr_callback: Optional[Callable[[str], None]] = None,
+        capture_output: bool = False,
+        output_hook: Callable[[str], None] = noop,
         params: Optional[dict[str, str]] = None,
         job_id: Optional[str] = None,
         interrupt_timeout: Optional[int] = None,
@@ -1788,18 +1781,13 @@ class Catalog:
             },
         )
         popen_kwargs: dict[str, Any] = {}
-
-        if stdout_callback is not None:
-            popen_kwargs = {"stdout": subprocess.PIPE}
-        if stderr_callback is not None:
-            popen_kwargs["stderr"] = subprocess.PIPE
+        if capture_output:
+            popen_kwargs = {"stdout": subprocess.PIPE, "stderr": subprocess.STDOUT}
 
         def raise_termination_signal(sig: int, _: Any) -> NoReturn:
             raise TerminationSignal(sig)
 
-        stdout_thread: Optional[Thread] = None
-        stderr_thread: Optional[Thread] = None
-
+        thread: Optional[Thread] = None
         with subprocess.Popen(cmd, env=env, **popen_kwargs) as proc:  # noqa: S603
             logger.info("Starting process %s", proc.pid)
 
@@ -1813,20 +1801,10 @@ class Catalog:
             orig_sigterm_handler = signal.getsignal(signal.SIGTERM)
             signal.signal(signal.SIGTERM, raise_termination_signal)
             try:
-                if stdout_callback is not None:
-                    stdout_thread = Thread(
-                        target=process_output,
-                        args=(proc.stdout, stdout_callback),
-                        daemon=True,
-                    )
-                    stdout_thread.start()
-                if stderr_callback is not None:
-                    stderr_thread = Thread(
-                        target=process_output,
-                        args=(proc.stderr, stderr_callback),
-                        daemon=True,
-                    )
-                    stderr_thread.start()
+                if capture_output:
+                    args = (proc.stdout, output_hook)
+                    thread = Thread(target=_process_stream, args=args, daemon=True)
+                    thread.start()
 
                 proc.wait()
             except TerminationSignal as exc:
@@ -1844,22 +1822,8 @@ class Catalog:
             finally:
                 signal.signal(signal.SIGTERM, orig_sigterm_handler)
                 signal.signal(signal.SIGINT, orig_sigint_handler)
-                # wait for the reader thread
-                thread_join_timeout_seconds = 30
-                if stdout_thread is not None:
-                    stdout_thread.join(timeout=thread_join_timeout_seconds)
-                    if stdout_thread.is_alive():
-                        logger.warning(
-                            "stdout thread is still alive after %s seconds",
-                            thread_join_timeout_seconds,
-                        )
-                if stderr_thread is not None:
-                    stderr_thread.join(timeout=thread_join_timeout_seconds)
-                    if stderr_thread.is_alive():
-                        logger.warning(
-                            "stderr thread is still alive after %s seconds",
-                            thread_join_timeout_seconds,
-                        )
+                if thread:
+                    thread.join()  # wait for the reader thread
 
         logger.info("Process %s exited with return code %s", proc.pid, proc.returncode)
         if proc.returncode in (

{datachain-0.31.3 → datachain-0.32.0}/src/datachain/data_storage/metastore.py

@@ -22,6 +22,7 @@ from sqlalchemy import (
     UniqueConstraint,
     select,
 )
+from sqlalchemy.sql import func as f
 
 from datachain.data_storage import JobQueryType, JobStatus
 from datachain.data_storage.serializer import Serializable
@@ -37,7 +38,9 @@ from datachain.dataset import (
 from datachain.error import (
     DatasetNotFoundError,
     DatasetVersionNotFoundError,
+    NamespaceDeleteNotAllowedError,
     NamespaceNotFoundError,
+    ProjectDeleteNotAllowedError,
     ProjectNotFoundError,
     TableMissingError,
 )
@@ -141,6 +144,10 @@ class AbstractMetastore(ABC, Serializable):
     def get_namespace(self, name: str, conn=None) -> Namespace:
         """Gets a single namespace by name"""
 
+    @abstractmethod
+    def remove_namespace(self, namespace_id: int, conn=None) -> None:
+        """Removes a single namespace by id"""
+
     @abstractmethod
     def list_namespaces(self, conn=None) -> list[Namespace]:
         """Gets a list of all namespaces"""
@@ -190,10 +197,30 @@ class AbstractMetastore(ABC, Serializable):
         It also creates project if not found and create flag is set to True.
         """
 
+    def is_default_project(self, project_name: str, namespace_name: str) -> bool:
+        return (
+            project_name == self.default_project_name
+            and namespace_name == self.default_namespace_name
+        )
+
+    def is_listing_project(self, project_name: str, namespace_name: str) -> bool:
+        return (
+            project_name == self.listing_project_name
+            and namespace_name == self.system_namespace_name
+        )
+
     @abstractmethod
     def get_project_by_id(self, project_id: int, conn=None) -> Project:
         """Gets a single project by id"""
 
+    @abstractmethod
+    def count_projects(self, namespace_id: Optional[int] = None) -> int:
+        """Counts projects in some namespace or in general."""
+
+    @abstractmethod
+    def remove_project(self, project_id: int, conn=None) -> None:
+        """Removes a single project by id"""
+
     @abstractmethod
     def list_projects(self, namespace_id: Optional[int], conn=None) -> list[Project]:
         """Gets list of projects in some namespace or in general (in all namespaces)"""
@@ -270,6 +297,10 @@ class AbstractMetastore(ABC, Serializable):
     ) -> Iterator[DatasetListRecord]:
         """Lists all datasets in some project or in all projects."""
 
+    @abstractmethod
+    def count_datasets(self, project_id: Optional[int] = None) -> int:
+        """Counts datasets in some project or in all projects."""
+
     @abstractmethod
     def list_datasets_by_prefix(
         self, prefix: str, project_id: Optional[int] = None
@@ -735,6 +766,18 @@ class AbstractDBMetastore(AbstractMetastore):
 
         return self.get_namespace(name)
 
+    def remove_namespace(self, namespace_id: int, conn=None) -> None:
+        num_projects = self.count_projects(namespace_id)
+        if num_projects > 0:
+            raise NamespaceDeleteNotAllowedError(
+                f"Namespace cannot be removed. It contains {num_projects} project(s). "
+                "Please remove the project(s) first."
+            )
+
+        n = self._namespaces
+        with self.db.transaction():
+            self.db.execute(self._namespaces_delete().where(n.c.id == namespace_id))
+
     def get_namespace(self, name: str, conn=None) -> Namespace:
         """Gets a single namespace by name"""
         n = self._namespaces
@@ -796,18 +839,6 @@ class AbstractDBMetastore(AbstractMetastore):
 
         return self.get_project(name, namespace.name)
 
-    def _is_listing_project(self, project_name: str, namespace_name: str) -> bool:
-        return (
-            project_name == self.listing_project_name
-            and namespace_name == self.system_namespace_name
-        )
-
-    def _is_default_project(self, project_name: str, namespace_name: str) -> bool:
-        return (
-            project_name == self.default_project_name
-            and namespace_name == self.default_namespace_name
-        )
-
     def get_project(
         self, name: str, namespace_name: str, create: bool = False, conn=None
     ) -> Project:
@@ -816,7 +847,7 @@ class AbstractDBMetastore(AbstractMetastore):
         p = self._projects
         validate = True
 
-        if self._is_listing_project(name, namespace_name) or self._is_default_project(
+        if self.is_listing_project(name, namespace_name) or self.is_default_project(
             name, namespace_name
         ):
             # we are always creating default and listing projects if they don't exist
@@ -858,7 +889,31 @@ class AbstractDBMetastore(AbstractMetastore):
             raise ProjectNotFoundError(f"Project with id {project_id} not found.")
         return self.project_class.parse(*rows[0])
 
-    def list_projects(self, namespace_id: Optional[int], conn=None) -> list[Project]:
+    def count_projects(self, namespace_id: Optional[int] = None) -> int:
+        p = self._projects
+        query = self._projects_select()
+        if namespace_id:
+            query = query.where(p.c.namespace_id == namespace_id)
+
+        query = select(f.count(1)).select_from(query.subquery())
+
+        return next(self.db.execute(query))[0]
+
+    def remove_project(self, project_id: int, conn=None) -> None:
+        num_datasets = self.count_datasets(project_id)
+        if num_datasets > 0:
+            raise ProjectDeleteNotAllowedError(
+                f"Project cannot be removed. It contains {num_datasets} dataset(s). "
+                "Please remove the dataset(s) first."
+            )
+
+        p = self._projects
+        with self.db.transaction():
+            self.db.execute(self._projects_delete().where(p.c.id == project_id))
+
+    def list_projects(
+        self, namespace_id: Optional[int] = None, conn=None
+    ) -> list[Project]:
         """
         Gets a list of projects inside some namespace, or in all namespaces
         """
@@ -1189,7 +1244,6 @@ class AbstractDBMetastore(AbstractMetastore):
     def list_datasets(
         self, project_id: Optional[int] = None
     ) -> Iterator["DatasetListRecord"]:
-        """Lists all datasets."""
         d = self._datasets
         query = self._base_list_datasets_query().order_by(
             self._datasets.c.name, self._datasets_versions.c.version
@@ -1198,6 +1252,16 @@ class AbstractDBMetastore(AbstractMetastore):
             query = query.where(d.c.project_id == project_id)
         yield from self._parse_dataset_list(self.db.execute(query))
 
+    def count_datasets(self, project_id: Optional[int] = None) -> int:
+        d = self._datasets
+        query = self._datasets_select()
+        if project_id:
+            query = query.where(d.c.project_id == project_id)
+
+        query = select(f.count(1)).select_from(query.subquery())
+
+        return next(self.db.execute(query))[0]
+
     def list_datasets_by_prefix(
         self, prefix: str, project_id: Optional[int] = None, conn=None
     ) -> Iterator["DatasetListRecord"]:

{datachain-0.31.3 → datachain-0.32.0}/src/datachain/error.py

@@ -34,6 +34,14 @@ class ProjectCreateNotAllowedError(NotAllowedError):
     pass
 
 
+class ProjectDeleteNotAllowedError(NotAllowedError):
+    pass
+
+
+class NamespaceDeleteNotAllowedError(NotAllowedError):
+    pass
+
+
 class ProjectNotFoundError(NotFoundError):
     pass
 

{datachain-0.31.3 → datachain-0.32.0}/src/datachain/lib/file.py

@@ -35,6 +35,7 @@ if TYPE_CHECKING:
     from datachain.catalog import Catalog
     from datachain.client.fsspec import Client
     from datachain.dataset import RowDict
+    from datachain.query.session import Session
 
 sha256 = partial(hashlib.sha256, usedforsecurity=False)
 
@@ -252,6 +253,15 @@ class File(DataModel):
         "last_modified",
     ]
 
+    # Allowed kwargs we forward to TextIOWrapper
+    _TEXT_WRAPPER_ALLOWED: ClassVar[tuple[str, ...]] = (
+        "encoding",
+        "errors",
+        "newline",
+        "line_buffering",
+        "write_through",
+    )
+
     @staticmethod
     def _validate_dict(
         v: Optional[Union[str, dict, list[dict]]],
@@ -328,7 +338,6 @@ class File(DataModel):
             from datachain.catalog.loader import get_catalog
 
             catalog = get_catalog()
-
         from datachain.client.fsspec import Client
 
         client_cls = Client.get_implementation(path)
@@ -341,6 +350,27 @@ class File(DataModel):
         file._set_stream(catalog)
         return file
 
+    @classmethod
+    def at(cls, uri: str, session: Optional["Session"] = None) -> "Self":
+        """Construct a File from a full URI in one call.
+
+        Example:
+            file = File.at("s3://bucket/path/to/output.png")
+            with file.open("wb") as f: ...
+        """
+        from datachain.client.fsspec import Client
+        from datachain.query.session import Session
+
+        if session is None:
+            session = Session.get()
+        catalog = session.catalog
+
+        client_cls = Client.get_implementation(uri)
+        source, rel_path = client_cls.split_url(uri)
+        file = cls(source=client_cls.get_uri(source), path=rel_path)
+        file._set_stream(catalog)
+        return file
+
     @classmethod
     def _from_row(cls, row: "RowDict") -> "Self":
         return cls(**{key: row[key] for key in cls._datachain_column_types})
@@ -354,28 +384,70 @@ class File(DataModel):
         return str(PurePosixPath(self.path).parent)
 
     @contextmanager
-    def open(self, mode: Literal["rb", "r"] = "rb") -> Iterator[Any]:
-        """Open the file and return a file object."""
-        if self.location:
-            with VFileRegistry.open(self, self.location) as f:  # type: ignore[arg-type]
-                yield f
+    def open(self, mode: str = "rb", **open_kwargs) -> Iterator[Any]:
+        """Open the file and return a file-like object.
 
-        else:
+        Supports both read ("rb", "r") and write modes (e.g. "wb", "w", "ab").
+        When opened in a write mode, metadata is refreshed after closing.
+        """
+        writing = any(ch in mode for ch in "wax+")
+        if self.location and writing:
+            raise VFileError(
+                "Writing to virtual file is not supported",
+                self.source,
+                self.path,
+            )
+
+        if self._catalog is None:
+            raise RuntimeError("Cannot open file: catalog is not set")
+
+        client: Client = self._catalog.get_client(self.source)
+
+        if not writing:
+            if self.location:
+                with VFileRegistry.open(self, self.location) as f:  # type: ignore[arg-type]
+                    yield self._wrap_text(f, mode, open_kwargs)
+                return
             if self._caching_enabled:
                 self.ensure_cached()
-            client: Client = self._catalog.get_client(self.source)
             with client.open_object(
                 self, use_cache=self._caching_enabled, cb=self._download_cb
             ) as f:
-                yield io.TextIOWrapper(f) if mode == "r" else f
+                yield self._wrap_text(f, mode, open_kwargs)
+            return
+
+        # write path
+        full_path = client.get_full_path(self.get_path_normalized())
+        with client.fs.open(full_path, mode, **open_kwargs) as f:
+            yield self._wrap_text(f, mode, open_kwargs)
+
+        # refresh metadata
+        info = client.fs.info(full_path)
+        refreshed = client.info_to_file(info, self.get_path_normalized())
+        for k, v in refreshed.model_dump().items():
+            setattr(self, k, v)
+
+    def _wrap_text(self, f: Any, mode: str, open_kwargs: dict[str, Any]) -> Any:
+        """Return stream possibly wrapped for text."""
+        if "b" in mode or isinstance(f, io.TextIOBase):
+            return f
+        filtered = {
+            k: open_kwargs[k] for k in self._TEXT_WRAPPER_ALLOWED if k in open_kwargs
+        }
+        return io.TextIOWrapper(f, **filtered)
 
     def read_bytes(self, length: int = -1):
         """Returns file contents as bytes."""
         with self.open() as stream:
             return stream.read(length)
 
-    def read_text(self):
-        """Returns file contents as text."""
+    def read_text(self, **open_kwargs):
+        """Return file contents decoded as text.
+
+        **open_kwargs : Any
+            Extra keyword arguments forwarded to ``open(mode="r", ...)``
+            (e.g. ``encoding="utf-8"``, ``errors="ignore"``)
+        """
         if self.location:
             raise VFileError(
                 "Reading text from virtual file is not supported",
@@ -383,7 +455,7 @@ class File(DataModel):
                 self.path,
             )
 
-        with self.open(mode="r") as stream:
+        with self.open(mode="r", **open_kwargs) as stream:
             return stream.read()
 
     def read(self, length: int = -1):
@@ -701,14 +773,19 @@ class TextFile(File):
     """`DataModel` for reading text files."""
 
     @contextmanager
-    def open(self, mode: Literal["rb", "r"] = "r"):
-        """Open the file and return a file object (default to text mode)."""
-        with super().open(mode=mode) as stream:
+    def open(self, mode: str = "r", **open_kwargs) -> Iterator[Any]:
+        """Open the file and return a file-like object.
+        Default to text mode"""
+        with super().open(mode=mode, **open_kwargs) as stream:
             yield stream
 
-    def read_text(self):
-        """Returns file contents as text."""
-        with self.open() as stream:
+    def read_text(self, **open_kwargs):
+        """Return file contents as text.
+
+        **open_kwargs : Any
+            Extra keyword arguments forwarded to ``open()`` (e.g. encoding).
+        """
+        with self.open(**open_kwargs) as stream:
             return stream.read()
 
     def save(self, destination: str, client_config: Optional[dict] = None):

{datachain-0.31.3 → datachain-0.32.0}/src/datachain/lib/namespaces.py

@@ -1,7 +1,11 @@
 from typing import Optional
 
-from datachain.error import NamespaceCreateNotAllowedError
-from datachain.namespace import Namespace
+from datachain.error import (
+    NamespaceCreateNotAllowedError,
+    NamespaceDeleteNotAllowedError,
+)
+from datachain.lib.projects import delete as delete_project
+from datachain.namespace import Namespace, parse_name
 from datachain.query import Session
 
 
@@ -71,3 +75,54 @@ def ls(session: Optional[Session] = None) -> list[Namespace]:
         ```
     """
     return Session.get(session).catalog.metastore.list_namespaces()
+
+
+def delete(name: str, session: Optional[Session]) -> None:
+    """
+    Removes a namespace by name.
+
+    Raises:
+        NamespaceNotFoundError: If the namespace does not exist.
+        NamespaceDeleteNotAllowedError: If the namespace is non-empty,
+            is the default namespace, or is a system namespace,
+            as these cannot be removed.
+
+    Parameters:
+        name : The name of the namespace.
+        session : Session to use for getting project.
+
+    Example:
+        ```py
+        import datachain as dc
+        from datachain.lib.namespace import delete as delete_namespace
+        delete_namespace("dev")
+        ```
+    """
+    session = Session.get(session)
+    metastore = session.catalog.metastore
+
+    namespace_name, project_name = parse_name(name)
+
+    if project_name:
+        return delete_project(project_name, namespace_name, session)
+
+    namespace = metastore.get_namespace(name)
+
+    if name == metastore.system_namespace_name:
+        raise NamespaceDeleteNotAllowedError(
+            f"Namespace {metastore.system_namespace_name} cannot be removed"
+        )
+
+    if name == metastore.default_namespace_name:
+        raise NamespaceDeleteNotAllowedError(
+            f"Namespace {metastore.default_namespace_name} cannot be removed"
+        )
+
+    num_projects = metastore.count_projects(namespace.id)
+    if num_projects > 0:
+        raise NamespaceDeleteNotAllowedError(
+            f"Namespace cannot be removed. It contains {num_projects} project(s). "
+            "Please remove the project(s) first."
+        )
+
+    metastore.remove_namespace(namespace.id)

{datachain-0.31.3 → datachain-0.32.0}/src/datachain/lib/projects.py

@@ -1,6 +1,6 @@
 from typing import Optional
 
-from datachain.error import ProjectCreateNotAllowedError
+from datachain.error import ProjectCreateNotAllowedError, ProjectDeleteNotAllowedError
 from datachain.project import Project
 from datachain.query import Session
 
@@ -86,3 +86,49 @@ def ls(
         namespace_id = session.catalog.metastore.get_namespace(namespace).id
 
     return session.catalog.metastore.list_projects(namespace_id)
+
+
+def delete(name: str, namespace: str, session: Optional[Session] = None) -> None:
+    """
+    Removes a project by name within a namespace.
+
+    Raises:
+        ProjectNotFoundError: If the project does not exist.
+        ProjectDeleteNotAllowedError: If the project is non-empty,
+            is the default project, or is a listing project,
+            as these cannot be removed.
+
+    Parameters:
+        name : The name of the project.
+        namespace : The name of the namespace.
+        session : Session to use for getting project.
+
+    Example:
+        ```py
+        import datachain as dc
+        dc.delete_project("my-project", "local")
+        ```
+    """
+    session = Session.get(session)
+    metastore = session.catalog.metastore
+
+    project = metastore.get_project(name, namespace)
+
+    if metastore.is_listing_project(name, namespace):
+        raise ProjectDeleteNotAllowedError(
+            f"Project {metastore.listing_project_name} cannot be removed"
+        )
+
+    if metastore.is_default_project(name, namespace):
+        raise ProjectDeleteNotAllowedError(
+            f"Project {metastore.default_project_name} cannot be removed"
+        )
+
+    num_datasets = metastore.count_datasets(project.id)
+    if num_datasets > 0:
+        raise ProjectDeleteNotAllowedError(
+            f"Project cannot be removed. It contains {num_datasets} dataset(s). "
+            "Please remove the dataset(s) first."
+        )
+
+    metastore.remove_project(project.id)

{datachain-0.31.3 → datachain-0.32.0}/src/datachain/namespace.py

@@ -9,6 +9,25 @@ N = TypeVar("N", bound="Namespace")
 NAMESPACE_NAME_RESERVED_CHARS = [".", "@"]
 
 
+def parse_name(name: str) -> tuple[str, Optional[str]]:
+    """
+    Parses namespace name into namespace and optional project name.
+    If both namespace and project are defined in name, they need to be split by dot
+    e.g dev.my-project
+    Valid inputs:
+        - dev.my-project
+        - dev
+    """
+    parts = name.split(".")
+    if len(parts) == 1:
+        return name, None
+    if len(parts) == 2:
+        return parts[0], parts[1]
+    raise InvalidNamespaceNameError(
+        f"Invalid namespace format: {name}. Expected 'namespace' or 'ns1.ns2'."
+    )
+
+
 @dataclass(frozen=True)
 class Namespace:
     id: int

{datachain-0.31.3 → datachain-0.32.0}/src/datachain.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datachain
-Version: 0.31.3
+Version: 0.32.0
 Summary: Wrangle unstructured AI data at scale
 Author-email: Dmitry Petrov <support@dvc.org>
 License-Expression: Apache-2.0