datachain 0.19.1__py3-none-any.whl → 0.20.0__py3-none-any.whl

datachain/cli/commands/datasets.py

@@ -8,7 +8,8 @@ if TYPE_CHECKING:
 
 from datachain.cli.utils import determine_flavors
 from datachain.config import Config
-from datachain.error import DatasetNotFoundError
+from datachain.dataset import parse_dataset_name
+from datachain.error import DataChainError, DatasetNotFoundError
 from datachain.studio import list_datasets as list_datasets_studio
 
 
@@ -101,11 +102,15 @@ def list_datasets_local(catalog: "Catalog", name: Optional[str] = None):
 
     for d in catalog.ls_datasets():
         for v in d.versions:
-            yield (d.name, v.version)
+            yield (d.full_name, v.version)
 
 
 def list_datasets_local_versions(catalog: "Catalog", name: str):
-    ds = catalog.get_dataset(name)
+    namespace_name, project_name, name = parse_dataset_name(name)
+    namespace_name = namespace_name or catalog.metastore.default_namespace_name
+    project_name = project_name or catalog.metastore.default_project_name
+    project = catalog.metastore.get_project(project_name, namespace_name)
+    ds = catalog.get_dataset(name, project)
     for v in ds.versions:
         yield (name, v.version)
 
@@ -129,25 +134,29 @@ def rm_dataset(
     name: str,
     version: Optional[str] = None,
     force: Optional[bool] = False,
-    studio: bool = False,
-    local: bool = False,
-    all: bool = True,
+    studio: Optional[bool] = False,
     team: Optional[str] = None,
 ):
-    from datachain.studio import remove_studio_dataset
-
-    token = Config().read().get("studio", {}).get("token")
-    all, local, studio = determine_flavors(studio, local, all, token)
-
-    if all or local:
+    namespace_name, project_name, name = parse_dataset_name(name)
+    namespace_name = namespace_name or catalog.metastore.default_namespace_name
+    project_name = project_name or catalog.metastore.default_project_name
+
+    if not catalog.metastore.is_local_dataset(namespace_name) and studio:
+        from datachain.studio import remove_studio_dataset
+
+        token = Config().read().get("studio", {}).get("token")
+        if not token:
+            raise DataChainError(
+                "Not logged in to Studio. Log in with 'datachain auth login'."
+            )
+        remove_studio_dataset(team, name, namespace_name, project_name, version, force)
+    else:
         try:
-            catalog.remove_dataset(name, version=version, force=force)
+            project = catalog.metastore.get_project(project_name, namespace_name)
+            catalog.remove_dataset(name, project, version=version, force=force)
         except DatasetNotFoundError:
             print("Dataset not found in local", file=sys.stderr)
 
-    if (all or studio) and token:
-        remove_studio_dataset(team, name, version, force)
-
 
 def edit_dataset(
     catalog: "Catalog",
@@ -155,21 +164,27 @@ def edit_dataset(
     new_name: Optional[str] = None,
     description: Optional[str] = None,
     attrs: Optional[list[str]] = None,
-    studio: bool = False,
-    local: bool = False,
-    all: bool = True,
     team: Optional[str] = None,
 ):
-    from datachain.studio import edit_studio_dataset
+    namespace_name, project_name, name = parse_dataset_name(name)
+    namespace_name = namespace_name or catalog.metastore.default_namespace_name
+    project_name = project_name or catalog.metastore.default_project_name
 
-    token = Config().read().get("studio", {}).get("token")
-    all, local, studio = determine_flavors(studio, local, all, token)
-
-    if all or local:
+    if catalog.metastore.is_local_dataset(namespace_name):
         try:
-            catalog.edit_dataset(name, new_name, description, attrs)
+            catalog.edit_dataset(
+                name, catalog.metastore.default_project, new_name, description, attrs
+            )
         except DatasetNotFoundError:
             print("Dataset not found in local", file=sys.stderr)
-
-    if (all or studio) and token:
-        edit_studio_dataset(team, name, new_name, description, attrs)
+    else:
+        from datachain.studio import edit_studio_dataset
+
+        token = Config().read().get("studio", {}).get("token")
+        if not token:
+            raise DataChainError(
+                "Not logged in to Studio. Log in with 'datachain auth login'."
+            )
+        edit_studio_dataset(
+            team, name, namespace_name, project_name, new_name, description, attrs
+        )

datachain/cli/parser/__init__.py

@@ -221,26 +221,6 @@ def get_parser() -> ArgumentParser:  # noqa: PLR0915
         nargs="+",
         help="Dataset attributes",
     )
-    parse_edit_dataset.add_argument(
-        "--studio",
-        action="store_true",
-        default=False,
-        help="Edit dataset from Studio",
-    )
-    parse_edit_dataset.add_argument(
-        "-L",
-        "--local",
-        action="store_true",
-        default=False,
-        help="Edit local dataset only",
-    )
-    parse_edit_dataset.add_argument(
-        "-a",
-        "--all",
-        action="store_true",
-        default=True,
-        help="Edit both datasets from studio and local",
-    )
     parse_edit_dataset.add_argument(
         "--team",
         action="store",
@@ -315,21 +295,7 @@ def get_parser() -> ArgumentParser:  # noqa: PLR0915
         "--studio",
         action="store_true",
         default=False,
-        help="Remove dataset from Studio",
-    )
-    rm_dataset_parser.add_argument(
-        "-L",
-        "--local",
-        action="store_true",
-        default=False,
-        help="Remove local datasets only",
-    )
-    rm_dataset_parser.add_argument(
-        "-a",
-        "--all",
-        action="store_true",
-        default=True,
-        help="Remove both local and studio",
+        help="Remove dataset from Studio only",
     )
     rm_dataset_parser.add_argument(
         "--team",

datachain/cli/parser/job.py

@@ -51,6 +51,13 @@ def add_jobs_parser(subparsers, parent_parser) -> None:
         help="Environment variables in KEY=VALUE format",
     )
 
+    studio_run_parser.add_argument(
+        "--cluster-id",
+        type=int,
+        action="store",
+        help="Compute cluster ID to run the job on",
+    )
+
     studio_run_parser.add_argument(
         "--workers",
         type=int,
@@ -165,3 +172,21 @@ def add_jobs_parser(subparsers, parent_parser) -> None:
         default=None,
         help="Team to check logs for (default: from config)",
     )
+
+    studio_clusters_help = "List compute clusters in Studio"
+    studio_clusters_description = "List compute clusters in Studio."
+
+    studio_clusters_parser = jobs_subparser.add_parser(
+        "clusters",
+        parents=[parent_parser],
+        description=studio_clusters_description,
+        help=studio_clusters_help,
+        formatter_class=CustomHelpFormatter,
+    )
+
+    studio_clusters_parser.add_argument(
+        "--team",
+        action="store",
+        default=None,
+        help="Team to list clusters for (default: from config)",
+    )

datachain/cli/parser/studio.py

@@ -89,8 +89,13 @@ def add_auth_parser(subparsers, parent_parser) -> None:
         help="Remove the token from the local project config",
     )
 
-    auth_team_help = "Set default team for Studio operations"
-    auth_team_description = "Set the default team for Studio operations."
+    auth_team_help = "Set or show default team for Studio operations"
+    auth_team_description = (
+        "Set or show the default team for Studio operations. "
+        "This will be used globally by default. "
+        "Use --local to set the team locally for the current project. "
+        "If no team name is provided, the default team will be shown."
+    )
 
     team_parser = auth_subparser.add_parser(
         "team",
@@ -102,13 +107,15 @@ def add_auth_parser(subparsers, parent_parser) -> None:
     team_parser.add_argument(
         "team_name",
         action="store",
+        default=None,
+        nargs="?",
         help="Name of the team to set as default",
     )
     team_parser.add_argument(
-        "--global",
+        "--local",
         action="store_true",
         default=False,
-        help="Set team globally for all projects",
+        help="Set team locally for the current project",
     )
 
     auth_token_help = "View Studio authentication token"  # noqa: S105