datachain 0.19.2__py3-none-any.whl → 0.20.1__py3-none-any.whl

datachain/namespace.py

@@ -0,0 +1,65 @@
+import builtins
+from dataclasses import dataclass, fields
+from datetime import datetime
+from typing import Any, Optional, TypeVar
+
+from datachain.error import InvalidNamespaceNameError
+
+N = TypeVar("N", bound="Namespace")
+NAMESPACE_NAME_RESERVED_CHARS = ["."]
+
+
+@dataclass(frozen=True)
+class Namespace:
+    id: int
+    uuid: str
+    name: str
+    description: Optional[str]
+    created_at: datetime
+
+    @staticmethod
+    def validate_name(name: str) -> None:
+        """Throws exception if name is invalid, otherwise returns None"""
+        if not name:
+            raise InvalidNamespaceNameError("Namespace name cannot be empty")
+
+        for c in NAMESPACE_NAME_RESERVED_CHARS:
+            if c in name:
+                raise InvalidNamespaceNameError(
+                    f"Character {c} is reserved and not allowed in namespace name"
+                )
+
+        if name in [Namespace.default(), Namespace.system()]:
+            raise InvalidNamespaceNameError(
+                f"Namespace name {name} is reserved and cannot be used."
+            )
+
+    @staticmethod
+    def default() -> str:
+        """Name of default namespace"""
+        return "local"
+
+    @staticmethod
+    def system() -> str:
+        """Name of the system namespace"""
+        return "system"
+
+    @property
+    def is_system(self):
+        return self.name == Namespace.system()
+
+    @classmethod
+    def parse(
+        cls: builtins.type[N],
+        id: int,
+        uuid: str,
+        name: str,
+        description: Optional[str],
+        created_at: datetime,
+    ) -> "Namespace":
+        return cls(id, uuid, name, description, created_at)
+
+    @classmethod
+    def from_dict(cls, d: dict[str, Any]) -> "Namespace":
+        kwargs = {f.name: d[f.name] for f in fields(cls) if f.name in d}
+        return cls(**kwargs)

datachain/project.py

@@ -0,0 +1,78 @@
+import builtins
+from dataclasses import dataclass, fields
+from datetime import datetime
+from typing import Any, Optional, TypeVar
+
+from datachain.error import InvalidProjectNameError
+from datachain.namespace import Namespace
+
+P = TypeVar("P", bound="Project")
+PROJECT_NAME_RESERVED_CHARS = ["."]
+
+
+@dataclass(frozen=True)
+class Project:
+    id: int
+    uuid: str
+    name: str
+    description: Optional[str]
+    created_at: datetime
+    namespace: Namespace
+
+    @staticmethod
+    def validate_name(name: str) -> None:
+        """Throws exception if name is invalid, otherwise returns None"""
+        if not name:
+            raise InvalidProjectNameError("Project name cannot be empty")
+
+        for c in PROJECT_NAME_RESERVED_CHARS:
+            if c in name:
+                raise InvalidProjectNameError(
+                    f"Character {c} is reserved and not allowed in project name."
+                )
+
+        if name in [Project.default(), Project.listing()]:
+            raise InvalidProjectNameError(
+                f"Project name {name} is reserved and cannot be used."
+            )
+
+    @staticmethod
+    def default() -> str:
+        """Name of default project"""
+        return "local"
+
+    @staticmethod
+    def listing() -> str:
+        """Name of listing project where all listing datasets will be saved"""
+        return "listing"
+
+    @classmethod
+    def parse(
+        cls: builtins.type[P],
+        namespace_id: int,
+        namespace_uuid: str,
+        namespace_name: str,
+        namespace_description: Optional[str],
+        namespace_created_at: datetime,
+        project_id: int,
+        uuid: str,
+        name: str,
+        description: Optional[str],
+        created_at: datetime,
+        project_namespace_id: int,
+    ) -> "Project":
+        namespace = Namespace.parse(
+            namespace_id,
+            namespace_uuid,
+            namespace_name,
+            namespace_description,
+            namespace_created_at,
+        )
+
+        return cls(project_id, uuid, name, description, created_at, namespace)
+
+    @classmethod
+    def from_dict(cls, d: dict[str, Any]) -> "Project":
+        namespace = Namespace.from_dict(d.pop("namespace"))
+        kwargs = {f.name: d[f.name] for f in fields(cls) if f.name in d}
+        return cls(**kwargs, namespace=namespace)

datachain/query/dataset.py

@@ -41,12 +41,13 @@ from datachain.data_storage.schema import (
     partition_col_names,
     partition_columns,
 )
-from datachain.dataset import DATASET_PREFIX, DatasetDependency, DatasetStatus, RowDict
+from datachain.dataset import DatasetDependency, DatasetStatus, RowDict
 from datachain.error import DatasetNotFoundError, QueryScriptCancelError
 from datachain.func.base import Function
 from datachain.lib.listing import is_listing_dataset, listing_dataset_expired
 from datachain.lib.udf import UDFAdapter, _get_cache
 from datachain.progress import CombinedDownloadCallback, TqdmCombinedDownloadCallback
+from datachain.project import Project
 from datachain.query.schema import C, UDFParamSpec, normalize_param
 from datachain.query.session import Session
 from datachain.query.udf import UdfInfo
@@ -83,7 +84,7 @@ PartitionByType = Union[
     Function, ColumnElement, Sequence[Union[Function, ColumnElement]]
 ]
 JoinPredicateType = Union[str, ColumnClause, ColumnElement]
-DatasetDependencyType = tuple[str, str]
+DatasetDependencyType = tuple["DatasetRecord", str]
 
 logger = logging.getLogger("datachain")
 
@@ -169,18 +170,17 @@ class QueryStep:
     """A query that returns all rows from specific dataset version"""
 
     catalog: "Catalog"
-    dataset_name: str
+    dataset: "DatasetRecord"
     dataset_version: str
 
     def apply(self) -> "StepResult":
         def q(*columns):
             return sqlalchemy.select(*columns)
 
-        dataset = self.catalog.get_dataset(self.dataset_name)
-        dr = self.catalog.warehouse.dataset_rows(dataset, self.dataset_version)
+        dr = self.catalog.warehouse.dataset_rows(self.dataset, self.dataset_version)
 
         return step_result(
-            q, dr.columns, dependencies=[(self.dataset_name, self.dataset_version)]
+            q, dr.columns, dependencies=[(self.dataset, self.dataset_version)]
         )
 
 
@@ -1095,6 +1095,8 @@ class DatasetQuery:
         self,
         name: str,
         version: Optional[str] = None,
+        project_name: Optional[str] = None,
+        namespace_name: Optional[str] = None,
         catalog: Optional["Catalog"] = None,
         session: Optional[Session] = None,
         indexing_column_types: Optional[dict[str, Any]] = None,
@@ -1128,33 +1130,38 @@ class DatasetQuery:
         if version:
             self.version = version
 
-        if is_listing_dataset(name):
-            if version:
-                # this listing dataset should already be listed as we specify
-                # exact version
-                self._set_starting_step(self.catalog.get_dataset(name))
-            else:
-                # not setting query step yet as listing dataset might not exist at
-                # this point
-                self.list_ds_name = name
+        namespace_name = namespace_name or self.catalog.metastore.default_namespace_name
+        project_name = project_name or self.catalog.metastore.default_project_name
+
+        if is_listing_dataset(name) and not version:
+            # not setting query step yet as listing dataset might not exist at
+            # this point
+            self.list_ds_name = name
         elif fallback_to_studio and is_token_set():
             self._set_starting_step(
-                self.catalog.get_dataset_with_remote_fallback(name, version)
+                self.catalog.get_dataset_with_remote_fallback(
+                    name,
+                    namespace_name=namespace_name,
+                    project_name=project_name,
+                    version=version,
+                )
             )
         else:
-            self._set_starting_step(self.catalog.get_dataset(name))
+            project = self.catalog.metastore.get_project(project_name, namespace_name)
+            self._set_starting_step(self.catalog.get_dataset(name, project=project))
 
     def _set_starting_step(self, ds: "DatasetRecord") -> None:
         if not self.version:
             self.version = ds.latest_version
 
-        self.starting_step = QueryStep(self.catalog, ds.name, self.version)
+        self.starting_step = QueryStep(self.catalog, ds, self.version)
 
         # at this point we know our starting dataset so setting up schemas
         self.feature_schema = ds.get_version(self.version).feature_schema
         self.column_types = copy(ds.schema)
         if "sys__id" in self.column_types:
             self.column_types.pop("sys__id")
+        self.project = ds.project
 
     def __iter__(self):
         return iter(self.db_results())
@@ -1162,21 +1169,6 @@ class DatasetQuery:
     def __or__(self, other):
         return self.union(other)
 
-    def pull_dataset(self, name: str, version: Optional[str] = None) -> "DatasetRecord":
-        print("Dataset not found in local catalog, trying to get from studio")
-
-        remote_ds_uri = f"{DATASET_PREFIX}{name}"
-        if version:
-            remote_ds_uri += f"@v{version}"
-
-        self.catalog.pull_dataset(
-            remote_ds_uri=remote_ds_uri,
-            local_ds_name=name,
-            local_ds_version=version,
-        )
-
-        return self.catalog.get_dataset(name)
-
     @staticmethod
     def get_table() -> "TableClause":
         table_name = "".join(
@@ -1657,6 +1649,8 @@ class DatasetQuery:
         workers: Union[bool, int] = False,
         min_task_size: Optional[int] = None,
         partition_by: Optional[PartitionByType] = None,
+        namespace: Optional[str] = None,
+        project: Optional[str] = None,
         cache: bool = False,
     ) -> "Self":
         query = self.clone()
@@ -1676,26 +1670,36 @@ class DatasetQuery:
 
     def _add_dependencies(self, dataset: "DatasetRecord", version: str):
         dependencies: set[DatasetDependencyType] = set()
-        for dep_name, dep_version in self.dependencies:
-            if Session.is_temp_dataset(dep_name):
+        for dep_dataset, dep_dataset_version in self.dependencies:
+            if Session.is_temp_dataset(dep_dataset.name):
                 # temp dataset are created for optimization and they will be removed
                 # afterwards. Therefore, we should not put them as dependencies, but
                 # their own direct dependencies
                 for dep in self.catalog.get_dataset_dependencies(
-                    dep_name, dep_version, indirect=False
+                    dep_dataset.name,
+                    dep_dataset_version,
+                    dep_dataset.project,
+                    indirect=False,
                 ):
                     if dep:
-                        dependencies.add((dep.name, dep.version))
+                        dep_project = self.catalog.metastore.get_project(
+                            dep.project, dep.namespace
+                        )
+                        dependencies.add(
+                            (
+                                self.catalog.get_dataset(dep.name, dep_project),
+                                dep.version,
+                            )
+                        )
             else:
-                dependencies.add((dep_name, dep_version))
+                dependencies.add((dep_dataset, dep_dataset_version))
 
-        for dep_name, dep_version in dependencies:
-            # ds_dependency_name, ds_dependency_version = dependency
+        for dep_dataset, dep_dataset_version in dependencies:
             self.catalog.metastore.add_dataset_dependency(
-                dataset.name,
+                dataset,
                 version,
-                dep_name,
-                dep_version,
+                dep_dataset,
+                dep_dataset_version,
             )
 
     def exec(self) -> "Self":
@@ -1711,6 +1715,7 @@ class DatasetQuery:
         self,
         name: Optional[str] = None,
         version: Optional[str] = None,
+        project: Optional[Project] = None,
         feature_schema: Optional[dict] = None,
         dependencies: Optional[list[DatasetDependency]] = None,
         description: Optional[str] = None,
@@ -1719,8 +1724,13 @@ class DatasetQuery:
         **kwargs,
     ) -> "Self":
         """Save the query as a dataset."""
+        project = project or self.catalog.metastore.default_project
         try:
-            if name and version and self.catalog.get_dataset(name).has_version(version):
+            if (
+                name
+                and version
+                and self.catalog.get_dataset(name, project).has_version(version)
+            ):
                 raise RuntimeError(f"Dataset {name} already has version {version}")
         except DatasetNotFoundError:
             pass
@@ -1745,6 +1755,7 @@ class DatasetQuery:
 
             dataset = self.catalog.create_dataset(
                 name,
+                project,
                 version=version,
                 feature_schema=feature_schema,
                 columns=columns,
@@ -1770,11 +1781,25 @@ class DatasetQuery:
 
             if dependencies:
                 # overriding dependencies
-                self.dependencies = {(dep.name, dep.version) for dep in dependencies}
+                self.dependencies = set()
+                for dep in dependencies:
+                    dep_project = self.catalog.metastore.get_project(
+                        dep.project, dep.namespace
+                    )
+                    self.dependencies.add(
+                        (self.catalog.get_dataset(dep.name, dep_project), dep.version)
+                    )
+
             self._add_dependencies(dataset, version)  # type: ignore [arg-type]
         finally:
             self.cleanup()
-        return self.__class__(name=name, version=version, catalog=self.catalog)
+        return self.__class__(
+            name=name,
+            namespace_name=project.namespace.name,
+            project_name=project.name,
+            version=version,
+            catalog=self.catalog,
+        )
 
     @property
     def is_ordered(self) -> bool:

datachain/query/session.py

@@ -108,7 +108,7 @@ class Session:
         prefix = self.get_temp_prefix()
         try:
             for dataset in list(self.catalog.metastore.list_datasets_by_prefix(prefix)):
-                self.catalog.remove_dataset(dataset.name, force=True)
+                self.catalog.remove_dataset(dataset.name, dataset.project, force=True)
         # suppress error when metastore has been reset during testing
         except TableMissingError:
             pass

datachain/remote/studio.py

@@ -17,6 +17,7 @@ import websockets
 from requests.exceptions import HTTPError, Timeout
 
 from datachain.config import Config
+from datachain.dataset import DatasetRecord
 from datachain.error import DataChainError
 from datachain.utils import STUDIO_URL, retry_with_backoff
 
@@ -36,13 +37,33 @@ logger = logging.getLogger("datachain")
 DATASET_ROWS_CHUNK_SIZE = 8192
 
 
+def get_studio_env_variable(name: str) -> Any:
+    """
+    Get the value of a DataChain Studio environment variable.
+    It first checks for the variable prefixed with 'DATACHAIN_STUDIO_',
+    then checks for the deprecated 'DVC_STUDIO_' prefix.
+    If neither is set, it returns the provided default value.
+    """
+    if (value := os.environ.get(f"DATACHAIN_STUDIO_{name}")) is not None:
+        return value
+    if (value := os.environ.get(f"DVC_STUDIO_{name}")) is not None:  # deprecated
+        logger.warning(
+            "Environment variable 'DVC_STUDIO_%s' is deprecated, "
+            "use 'DATACHAIN_STUDIO_%s' instead.",
+            name,
+            name,
+        )
+        return value
+    return None
+
+
 def _is_server_error(status_code: int) -> bool:
     return str(status_code).startswith("5")
 
 
 def is_token_set() -> bool:
     return (
-        bool(os.environ.get("DVC_STUDIO_TOKEN"))
+        bool(get_studio_env_variable("TOKEN"))
         or Config().read().get("studio", {}).get("token") is not None
     )
 
@@ -78,12 +99,12 @@ class StudioClient:
 
     @property
     def token(self) -> str:
-        token = os.environ.get("DVC_STUDIO_TOKEN") or self.config.get("token")
+        token = get_studio_env_variable("TOKEN") or self.config.get("token")
 
         if not token:
             raise DataChainError(
                 "Studio token is not set. Use `datachain auth login` "
-                "or environment variable `DVC_STUDIO_TOKEN` to set it."
+                "or environment variable `DATACHAIN_STUDIO_TOKEN` to set it."
             )
 
         return token
@@ -91,8 +112,8 @@ class StudioClient:
     @property
     def url(self) -> str:
         return (
-            os.environ.get("DVC_STUDIO_URL") or self.config.get("url") or STUDIO_URL
-        ) + "/api"
+            get_studio_env_variable("URL") or self.config.get("url") or STUDIO_URL
+        ).rstrip("/") + "/api"
 
     @property
     def config(self) -> dict:
@@ -107,13 +128,13 @@ class StudioClient:
         return self._team
 
     def _get_team(self) -> str:
-        team = os.environ.get("DVC_STUDIO_TEAM") or self.config.get("team")
+        team = get_studio_env_variable("TEAM") or self.config.get("team")
 
         if not team:
             raise DataChainError(
                 "Studio team is not set. "
                 "Use `datachain auth team <team_name>` "
-                "or environment variable `DVC_STUDIO_TEAM` to set it. "
+                "or environment variable `DATACHAIN_STUDIO_TEAM` to set it. "
                 "You can also set `studio.team` in the config file."
             )
 
@@ -291,13 +312,17 @@ class StudioClient:
     def edit_dataset(
         self,
         name: str,
+        namespace: str,
+        project: str,
         new_name: Optional[str] = None,
         description: Optional[str] = None,
         attrs: Optional[list[str]] = None,
     ) -> Response[DatasetInfoData]:
         body = {
             "new_name": new_name,
-            "dataset_name": name,
+            "name": name,
+            "namespace": namespace,
+            "project": project,
             "description": description,
             "attrs": attrs,
         }
@@ -310,44 +335,44 @@ class StudioClient:
     def rm_dataset(
         self,
         name: str,
+        namespace: str,
+        project: str,
         version: Optional[str] = None,
         force: Optional[bool] = False,
     ) -> Response[DatasetInfoData]:
         return self._send_request(
             "datachain/datasets",
             {
-                "dataset_name": name,
-                "dataset_version": version,
+                "name": name,
+                "namespace": namespace,
+                "project": project,
+                "version": version,
                 "force": force,
             },
             method="DELETE",
         )
 
-    def dataset_info(self, name: str) -> Response[DatasetInfoData]:
+    def dataset_info(
+        self, namespace: str, project: str, name: str
+    ) -> Response[DatasetInfoData]:
         def _parse_dataset_info(dataset_info):
             _parse_dates(dataset_info, ["created_at", "finished_at"])
             for version in dataset_info.get("versions"):
                 _parse_dates(version, ["created_at"])
+            _parse_dates(dataset_info.get("project"), ["created_at"])
+            _parse_dates(dataset_info.get("project").get("namespace"), ["created_at"])
 
             return dataset_info
 
         response = self._send_request(
-            "datachain/datasets/info", {"dataset_name": name}, method="GET"
+            "datachain/datasets/info",
+            {"namespace": namespace, "project": project, "name": name},
+            method="GET",
         )
         if response.ok:
             response.data = _parse_dataset_info(response.data)
         return response
 
-    def dataset_rows_chunk(
-        self, name: str, version: str, offset: int
-    ) -> Response[DatasetRowsData]:
-        req_data = {"dataset_name": name, "dataset_version": version}
-        return self._send_request_msgpack(
-            "datachain/datasets/rows",
-            {**req_data, "offset": offset, "limit": DATASET_ROWS_CHUNK_SIZE},
-            method="GET",
-        )
-
     def dataset_job_versions(self, job_id: str) -> Response[DatasetJobVersionsData]:
         return self._send_request(
             "datachain/datasets/dataset_job_versions",
@@ -356,20 +381,30 @@ class StudioClient:
         )
 
     def export_dataset_table(
-        self, name: str, version: str
+        self, dataset: DatasetRecord, version: str
     ) -> Response[DatasetExportSignedUrls]:
         return self._send_request(
             "datachain/datasets/export",
-            {"dataset_name": name, "dataset_version": version},
+            {
+                "namespace": dataset.project.namespace.name,
+                "project": dataset.project.name,
+                "name": dataset.name,
+                "version": version,
+            },
             method="GET",
         )
 
     def dataset_export_status(
-        self, name: str, version: str
+        self, dataset: DatasetRecord, version: str
     ) -> Response[DatasetExportStatus]:
         return self._send_request(
             "datachain/datasets/export-status",
-            {"dataset_name": name, "dataset_version": version},
+            {
+                "namespace": dataset.project.namespace.name,
+                "project": dataset.project.name,
+                "name": dataset.name,
+                "version": version,
+            },
             method="GET",
         )
 
@@ -392,7 +427,7 @@ class StudioClient:
         requirements: Optional[str] = None,
         repository: Optional[str] = None,
         priority: Optional[int] = None,
-        cluster_id: Optional[int] = None,
+        cluster: Optional[str] = None,
     ) -> Response[JobData]:
         data = {
             "query": query,
@@ -405,7 +440,7 @@ class StudioClient:
             "requirements": requirements,
             "repository": repository,
             "priority": priority,
-            "compute_cluster_id": cluster_id,
+            "compute_cluster_name": cluster,
         }
         return self._send_request("datachain/job", data)