Add option to set team in config for studio

src/datachain/cli.py

@@ -174,6 +174,30 @@ def add_studio_parser(subparsers, parent_parser) -> None:
         help=studio_logout_help,
     )
 
+    studio_team_help = "Set the default team for Datachain"
+    studio_team_description = (
+        "Set the default team for Datachain to use when interacting with Studio."
+    )
+
+    team_parser = studio_subparser.add_parser(
+        "team",
+        parents=[parent_parser],
+        description=studio_team_description,
+        help=studio_team_help,
+    )
+    # Add a positional argument for the team name
+    team_parser.add_argument(
+        "team_name",
+        action="store",
+        help="The name of the team to set as the default.",
+    )
+    team_parser.add_argument(
+        "--global",
+        action="store_true",
+        default=False,
+        help="Set the team globally for all Datachain projects.",
+    )
+
     studio_token_help = "View the token datachain uses to contact Studio"  # noqa: S105 # nosec B105
 
     studio_subparser.add_parser(

src/datachain/studio.py

@@ -21,9 +21,22 @@ def process_studio_cli_args(args: "Namespace"):
         return logout()
     if args.cmd == "token":
         return token()
+    if args.cmd == "team":
+        return set_team(args)
     raise DataChainError(f"Unknown command '{args.cmd}'.")
 
 
+def set_team(args: "Namespace"):
+    level = ConfigLevel.GLOBAL if args.__dict__.get("global") else ConfigLevel.LOCAL
+    config = Config(level)
+    with config.edit() as conf:
+        studio_conf = conf.get("studio", {})
+        studio_conf["team"] = args.team_name
+        conf["studio"] = studio_conf
+
+    print(f"Set default team to '{args.team_name}' in {config.config_file()}")
+
+
 def login(args: "Namespace"):
     from dvc_studio_client.auth import StudioAuthError, get_access_token
 

tests/test_cli_studio.py

@@ -82,3 +82,15 @@ def test_studio_token(capsys):
         del conf["studio"]["token"]
 
     assert main(["studio", "token"]) == 1
+
+
+def test_studio_team_local():
+    assert main(["studio", "team", "team_name"]) == 0
+    config = Config(ConfigLevel.LOCAL).read()
+    assert config["studio"]["team"] == "team_name"
+
+
+def test_studio_team_global():
+    assert main(["studio", "team", "team_name", "--global"]) == 0
+    config = Config(ConfigLevel.GLOBAL).read()
+    assert config["studio"]["team"] == "team_name"