diff --git a/README.md b/README.md index 3ced807..81ab650 100644 --- a/README.md +++ b/README.md @@ -333,9 +333,7 @@ For more information, `cometx log --help` ## cometx migrate-users -This command is used to migrate users into workspaces from a source Comet environment to a destination environment. It reads workspace membership from a chargeback report from the source environment and invites each user to the corresponding workspace in the destination environment by email. If a user with the email does not exist in the destination environment, they still will be provisioned access to these workspaces once they sign up using that email. - -The chargeback report can be fetched automatically from the source environment's admin API, or you can provide a pre-downloaded JSON file. +This command migrates users into workspaces from a source Comet environment to a destination environment. It reads workspace membership from a chargeback report fetched from the source environment's admin API (or a pre-downloaded local file) and invites each user to the corresponding workspace in the destination environment by email. If a user with that email does not yet exist in the destination, they will be provisioned access once they sign up. ``` cometx migrate-users --api-key DEST_KEY --source-api-key SOURCE_KEY [FLAGS ...] @@ -343,14 +341,17 @@ cometx migrate-users --api-key DEST_KEY --chargeback-report /path/to/report.json ``` **Arguments:** -* `--api-key API_KEY` - API key for the destination environment where users will be added. Falls back to `COMET_API_KEY` environment variable if not provided. -* `--source-api-key SOURCE_API_KEY` - API key for the source environment. Used to fetch the chargeback report. Falls back to `COMET_API_KEY` environment variable if not provided. Required unless `--chargeback-report` is given. -* `--chargeback-report PATH` - Path to a local chargeback report JSON file. When provided, the report is loaded from this file instead of being fetched from the source environment, and `--source-api-key` is not required. +* `--api-key API_KEY` - API key for the destination environment. Falls back to `COMET_API_KEY` if not provided. +* `--url URL` - Base URL of the destination Comet environment (e.g. `https://comet.example.com`). Required for self-hosted instances when the API key does not encode the server URL. +* `--source-api-key SOURCE_API_KEY` - API key for the source environment. Used to fetch the chargeback report. Required unless `--chargeback-report` is given. +* `--source-url SOURCE_URL` - Base URL of the source Comet environment. Required for self-hosted source instances when the source API key does not encode the server URL. +* `--chargeback-report PATH` - Path to a local chargeback report JSON file. When provided, `--source-api-key` is not required. ### Flags -* `--create-workspaces` - Create workspaces on the destination environment if they don't already exist (default: off) +* `--create-workspaces` - Create workspaces on the destination if they don't already exist (default: off) * `--dry-run` - Print what would happen without making any changes +* `--failures-output PATH` - Path to write failed operations JSON (default: `bulk_add_failures_by_email.json`) ### Examples @@ -364,6 +365,9 @@ cometx migrate-users --api-key DEST_KEY --source-api-key SOURCE_KEY # Use a local chargeback report file cometx migrate-users --api-key DEST_KEY --chargeback-report /tmp/chargeback_reports.json +# Self-hosted environments — provide explicit URLs +cometx migrate-users --api-key DEST_KEY --url https://comet.dest.example.com \ + --source-api-key SOURCE_KEY --source-url https://comet.src.example.com ``` For more information, `cometx migrate-users --help` @@ -518,12 +522,8 @@ cometx admin chargeback-report [YEAR-MONTH] **Examples:** ``` cometx admin chargeback-report -<<<<<<< Updated upstream -cometx admin usage-report -======= cometx admin chargeback-report 2024-09 ``` ->>>>>>> Stashed changes #### usage-report diff --git a/cometx/cli/copy.py b/cometx/cli/copy.py index 5a21c82..11357ad 100644 --- a/cometx/cli/copy.py +++ b/cometx/cli/copy.py @@ -78,6 +78,8 @@ import time import urllib.parse import zipfile + +import requests from datetime import datetime, timedelta from comet_ml import APIExperiment, Artifact, Experiment, OfflineExperiment @@ -102,6 +104,7 @@ from ..framework.comet.download_manager import sanitize_filename from ..utils import remove_extra_slashes from .copy_utils import upload_single_offline_experiment +from .migrate_users import _create_workspace ADDITIONAL_ARGS = False @@ -834,26 +837,26 @@ def copy(self, source, destination, symlink, ignore, debug, sync, create_workspa print( f"Workspace {workspace_dst!r} does not exist, attempting to create it..." ) + dest_url = self.api._get_url_server() + headers = { + "Authorization": self.api.api_key, + "Content-Type": "application/json", + } try: - import requests - - url = f"{self.api.server_url}/api/rest/v2/write/workspace/new" - response = requests.post( - url, - json={"name": workspace_dst}, - headers={ - "Authorization": self.api.api_key, - "Content-Type": "application/json", - }, - ) - response.raise_for_status() + _create_workspace(dest_url, headers, workspace_dst) print(f"Workspace {workspace_dst!r} created successfully.") - except Exception as exc: + except requests.exceptions.HTTPError as exc: + raise Exception( + f"Workspace {workspace_dst!r} does not exist and could not be " + f"created automatically (HTTP {exc.response.status_code}). " + f"Please create it via the Comet UI and try again." + ) from exc + except requests.exceptions.RequestException as exc: raise Exception( f"Workspace {workspace_dst!r} does not exist and could not be " f"created automatically: {exc}. " f"Please create it via the Comet UI and try again." - ) + ) from exc else: raise Exception( f"{workspace_dst} does not exist; use --create-workspaces to " diff --git a/cometx/cli/migrate_users.py b/cometx/cli/migrate_users.py index e7997de..ed7841d 100644 --- a/cometx/cli/migrate_users.py +++ b/cometx/cli/migrate_users.py @@ -20,9 +20,9 @@ Examples: -$ cometx --api-key DEST_KEY migrate-users --source-api-key SOURCE_KEY --dry-run -$ cometx --api-key DEST_KEY migrate-users --source-api-key SOURCE_KEY -$ cometx --api-key DEST_KEY migrate-users --chargeback-report /path/to/report.json +$ cometx migrate-users --api-key DEST_KEY --source-api-key SOURCE_KEY --dry-run +$ cometx migrate-users --api-key DEST_KEY --source-api-key SOURCE_KEY +$ cometx migrate-users --api-key DEST_KEY --chargeback-report /path/to/report.json """ import argparse @@ -30,6 +30,8 @@ import json import os import sys +import urllib.parse + import requests ADDITIONAL_ARGS = False @@ -40,13 +42,29 @@ def get_parser_arguments(parser): parser.add_argument( "--api-key", - help="API key for the destination environment (used to add workspace members)", + help="API key for the destination environment (used to add workspace members). " + "Falls back to COMET_API_KEY environment variable if not provided.", + type=str, + default=None, + ) + parser.add_argument( + "--url", + help="Base URL of the destination Comet environment (e.g. https://comet.example.com). " + "Required for self-hosted instances when the API key does not encode the server URL.", type=str, default=None, ) parser.add_argument( "--source-api-key", - help="API key for the source environment (used to fetch the chargeback report)", + help="API key for the source environment (used to fetch the chargeback report). " + "Required unless --chargeback-report is given.", + type=str, + default=None, + ) + parser.add_argument( + "--source-url", + help="Base URL of the source Comet environment. " + "Required for self-hosted instances when the source API key does not encode the server URL.", type=str, default=None, ) @@ -68,15 +86,45 @@ def get_parser_arguments(parser): default=False, action="store_true", ) + parser.add_argument( + "--failures-output", + help="Path to write failed operations JSON (default: bulk_add_failures_by_email.json)", + type=str, + default="bulk_add_failures_by_email.json", + ) -def _resolve_server_url(api_key): - if "*" not in api_key: - return COMET_CLOUD_URL +def _resolve_server_url(api_key, explicit_url=None): + """Return the server base URL. - _, encoded = api_key.split("*", 1) - payload = json.loads(base64.b64decode(encoded)) - return payload["baseUrl"].rstrip("/") + Priority: + 1. ``explicit_url`` (from --url / --source-url), if provided. + 2. URL encoded inside the API key (new-style keys contain ``*``). + 3. Error — no silent fallback to cloud URL. + """ + if explicit_url: + parsed = urllib.parse.urlparse(explicit_url) + if parsed.scheme != "https": + print("[ERROR] --url/--source-url must use https://.") + sys.exit(1) + return explicit_url.rstrip("/") + + if "*" in api_key: + try: + _, encoded = api_key.split("*", 1) + # base64 padding may be missing + padding = (4 - len(encoded) % 4) % 4 + payload = json.loads(base64.b64decode(encoded + "=" * padding)) + return payload["baseUrl"].rstrip("/") + except Exception: + pass # Fall through to error below + + print( + "[ERROR] Cannot determine the server URL from the API key. " + "Pass --url (destination) or --source-url (source) explicitly, " + "e.g. --url https://comet.example.com" + ) + sys.exit(1) def _fetch_chargeback_report(server_url, source_api_key): @@ -101,7 +149,18 @@ def _get_existing_workspaces(dest_url, headers): url = f"{dest_url}/api/rest/v2/workspaces" resp = requests.get(url, headers=headers, timeout=15) resp.raise_for_status() - return set(resp.json()) + data = resp.json() + # Handle {"workspaceNames": [...]} dict shape + if isinstance(data, dict): + names = data.get("workspaceNames", []) + if names and isinstance(names[0], dict): + return {ws["name"] for ws in names} + return set(names) + # Handle list of dicts with a "name" key + if data and isinstance(data[0], dict): + return {ws["name"] for ws in data} + # Handle flat list of strings + return set(data) def _create_workspace(dest_url, headers, workspace_name): @@ -126,7 +185,7 @@ def _add_member(url, headers, email, workspace_name): response = requests.post( url, headers=headers, - data=json.dumps(payload), + json=payload, timeout=15, ) @@ -163,22 +222,30 @@ def migrate_users(parsed_args): print("[ERROR] No API key found. Set COMET_API_KEY or pass --api-key.") sys.exit(1) - source_api_key = parsed_args.source_api_key or os.environ.get("COMET_API_KEY") + source_api_key = parsed_args.source_api_key create_workspaces = parsed_args.create_workspaces dry_run = parsed_args.dry_run + failures_output = parsed_args.failures_output if not parsed_args.chargeback_report and not source_api_key: - print("[ERROR] Provide either --source-api-key or --chargeback-report.") + print( + "[ERROR] --source-api-key is required when --chargeback-report is not provided." + ) sys.exit(1) - dest_url = _resolve_server_url(api_key) + dest_url = _resolve_server_url(api_key, parsed_args.url) print(f"Destination URL: {dest_url}") if parsed_args.chargeback_report: data = _load_chargeback_report(parsed_args.chargeback_report) else: - source_url = _resolve_server_url(source_api_key) + source_url = _resolve_server_url(source_api_key, parsed_args.source_url) print(f"Source URL: {source_url}") + if source_url == dest_url and source_api_key == api_key: + print( + "[WARNING] Source and destination URL and API key are identical. " + "Are you sure you want to migrate users to the same environment?" + ) try: data = _fetch_chargeback_report(source_url, source_api_key) except requests.exceptions.RequestException as e: @@ -215,7 +282,7 @@ def migrate_users(parsed_args): elif create_workspaces and dry_run: print("[DRY RUN] Would check/create workspaces on destination\n") - url = f"{dest_url}/api/rest/v2/write/add-workspace-member" + add_member_url = f"{dest_url}/api/rest/v2/write/add-workspace-member" total_added = 0 total_skipped = 0 @@ -252,11 +319,12 @@ def migrate_users(parsed_args): continue if dry_run: + print(f" [DRY RUN] Would add '{email}' to '{ws_name}'") ws_added += 1 total_added += 1 continue - status, error_info = _add_member(url, headers, email, ws_name) + status, error_info = _add_member(add_member_url, headers, email, ws_name) if status == "added": ws_added += 1 @@ -280,7 +348,7 @@ def migrate_users(parsed_args): ) if dry_run: - print(f" [DRY RUN] Would add {ws_added}/{len(members)} users") + print(f" [DRY RUN] Total: would add {ws_added}/{len(members)} users") else: parts = [f" Added {ws_added}/{len(members)} users successfully"] if ws_already_member: @@ -309,10 +377,9 @@ def migrate_users(parsed_args): print(f" *** Remove --dry-run to execute for real ***") if failures: - failures_file = "bulk_add_failures_by_email.json" - with open(failures_file, "w") as f: + with open(failures_output, "w") as f: json.dump(failures, f, indent=2) - print(f"\n Failed operations saved to {failures_file}") + print(f"\n Failed operations saved to {failures_output}") def main(args): diff --git a/tests/unit/test_migrate_users.py b/tests/unit/test_migrate_users.py new file mode 100644 index 0000000..e812ee9 --- /dev/null +++ b/tests/unit/test_migrate_users.py @@ -0,0 +1,229 @@ +# -*- coding: utf-8 -*- +# **************************************** +# __ +# _________ ____ ___ ___ / /__ __ +# / ___/ __ \/ __ `__ \/ _ \/ __/ |/_/ +# / /__/ /_/ / / / / / / __/ /__> < +# \___/\____/_/ /_/ /_/\___/\__/_/|_| +# +# +# Copyright (c) 2022 Cometx Development +# Team. All rights reserved. +# **************************************** + +import argparse +import base64 +import json +import unittest +from unittest.mock import MagicMock, call, patch + +import pytest + +from cometx.cli.migrate_users import ( + _add_member, + _get_existing_workspaces, + _resolve_server_url, + migrate_users, +) + + +def _make_new_style_key(base_url): + payload = json.dumps({"baseUrl": base_url}).encode() + encoded = base64.b64encode(payload).decode().rstrip("=") + return f"abc123*{encoded}" + + +class TestResolveServerUrl: + def test_explicit_url_takes_priority(self): + key = _make_new_style_key("https://encoded.example.com") + assert _resolve_server_url(key, "https://explicit.example.com") == "https://explicit.example.com" + + def test_explicit_url_trailing_slash_stripped(self): + assert _resolve_server_url("anykey", "https://example.com/") == "https://example.com" + + def test_explicit_url_non_https_exits(self): + with pytest.raises(SystemExit): + _resolve_server_url("anykey", "http://example.com") + + def test_new_style_key_decoded(self): + key = _make_new_style_key("https://self-hosted.example.com") + assert _resolve_server_url(key) == "https://self-hosted.example.com" + + def test_new_style_key_trailing_slash_stripped(self): + key = _make_new_style_key("https://self-hosted.example.com/") + assert _resolve_server_url(key) == "https://self-hosted.example.com" + + def test_old_style_key_exits(self): + with pytest.raises(SystemExit): + _resolve_server_url("oldstylekey") + + def test_malformed_new_style_key_exits(self): + with pytest.raises(SystemExit): + _resolve_server_url("abc*notvalidbase64!!!") + + def test_new_style_key_missing_base_url_exits(self): + payload = json.dumps({"other": "value"}).encode() + encoded = base64.b64encode(payload).decode() + key = f"abc*{encoded}" + with pytest.raises(SystemExit): + _resolve_server_url(key) + + +class TestGetExistingWorkspaces: + @patch("cometx.cli.migrate_users.requests.get") + def test_flat_list_response(self, mock_get): + mock_get.return_value.json.return_value = ["ws1", "ws2"] + mock_get.return_value.raise_for_status = MagicMock() + result = _get_existing_workspaces("https://example.com", {}) + assert result == {"ws1", "ws2"} + + @patch("cometx.cli.migrate_users.requests.get") + def test_dict_list_response(self, mock_get): + mock_get.return_value.json.return_value = [{"name": "ws1"}, {"name": "ws2"}] + mock_get.return_value.raise_for_status = MagicMock() + result = _get_existing_workspaces("https://example.com", {}) + assert result == {"ws1", "ws2"} + + @patch("cometx.cli.migrate_users.requests.get") + def test_workspace_names_dict_response(self, mock_get): + mock_get.return_value.json.return_value = {"workspaceNames": ["ws1", "ws2"]} + mock_get.return_value.raise_for_status = MagicMock() + result = _get_existing_workspaces("https://example.com", {}) + assert result == {"ws1", "ws2"} + + @patch("cometx.cli.migrate_users.requests.get") + def test_workspace_names_dict_with_dict_entries(self, mock_get): + mock_get.return_value.json.return_value = { + "workspaceNames": [{"name": "ws1"}, {"name": "ws2"}] + } + mock_get.return_value.raise_for_status = MagicMock() + result = _get_existing_workspaces("https://example.com", {}) + assert result == {"ws1", "ws2"} + + +class TestAddMember: + @patch("cometx.cli.migrate_users.requests.post") + def test_success(self, mock_post): + mock_post.return_value.status_code = 200 + status, error = _add_member("http://example.com/add", {}, "user@example.com", "ws1") + assert status == "added" + assert error is None + + @patch("cometx.cli.migrate_users.requests.post") + def test_already_member(self, mock_post): + mock_post.return_value.status_code = 400 + mock_post.return_value.json.return_value = {"msg": "already member of ws1"} + status, error = _add_member("http://example.com/add", {}, "user@example.com", "ws1") + assert status == "already_member" + assert error is None + + @patch("cometx.cli.migrate_users.requests.post") + def test_failure(self, mock_post): + mock_post.return_value.status_code = 403 + mock_post.return_value.json.return_value = {"msg": "forbidden"} + status, error = _add_member("http://example.com/add", {}, "user@example.com", "ws1") + assert status == "failed" + assert error["status"] == 403 + + @patch("cometx.cli.migrate_users.requests.post") + def test_request_exception(self, mock_post): + import requests as req + mock_post.side_effect = req.exceptions.ConnectionError("timeout") + status, error = _add_member("http://example.com/add", {}, "user@example.com", "ws1") + assert status == "failed" + assert error["status"] == "exception" + + +class TestMigrateUsersDryRun: + def _make_args(self, **kwargs): + defaults = dict( + api_key="dest-key", + url=None, + source_api_key="source-key", + source_url=None, + chargeback_report=None, + create_workspaces=False, + dry_run=True, + failures_output="failures.json", + ) + defaults.update(kwargs) + return argparse.Namespace(**defaults) + + @patch("cometx.cli.migrate_users._resolve_server_url", return_value="https://dest.example.com") + @patch("cometx.cli.migrate_users._fetch_chargeback_report") + @patch("builtins.print") + def test_dry_run_no_api_calls(self, mock_print, mock_fetch, mock_resolve): + mock_fetch.return_value = { + "workspaces": [ + { + "name": "ws1", + "members": [ + {"email": "alice@example.com", "userName": "alice"}, + {"email": "bob@example.com", "userName": "bob"}, + ], + } + ] + } + with patch("cometx.cli.migrate_users.requests.post") as mock_post: + migrate_users(self._make_args()) + mock_post.assert_not_called() + + @patch("cometx.cli.migrate_users._resolve_server_url", return_value="https://dest.example.com") + @patch("cometx.cli.migrate_users._fetch_chargeback_report") + @patch("builtins.print") + def test_dry_run_prints_per_user(self, mock_print, mock_fetch, mock_resolve): + mock_fetch.return_value = { + "workspaces": [ + { + "name": "ws1", + "members": [ + {"email": "alice@example.com", "userName": "alice"}, + ], + } + ] + } + migrate_users(self._make_args()) + printed = " ".join(str(c) for c in mock_print.call_args_list) + assert "alice@example.com" in printed + assert "ws1" in printed + + @patch("cometx.cli.migrate_users._resolve_server_url", return_value="https://dest.example.com") + @patch("cometx.cli.migrate_users._fetch_chargeback_report") + @patch("builtins.print") + def test_dry_run_skips_no_email(self, mock_print, mock_fetch, mock_resolve): + mock_fetch.return_value = { + "workspaces": [ + { + "name": "ws1", + "members": [ + {"email": None, "userName": "noemail"}, + ], + } + ] + } + migrate_users(self._make_args()) + printed = " ".join(str(c) for c in mock_print.call_args_list) + assert "noemail" in printed + assert "no email" in printed + + def test_no_source_key_without_chargeback_report_exits(self): + args = self._make_args(source_api_key=None, chargeback_report=None) + with self.assertRaises(SystemExit): + migrate_users(args) + + def assertRaises(self, exc): + import contextlib + + @contextlib.contextmanager + def ctx(): + try: + yield + raise AssertionError(f"{exc} was not raised") + except exc: + pass + + return ctx() + + +if __name__ == "__main__": + unittest.main()