Fix review issues in migrate-users PR

- copy.py: move `import requests` to top-level; use `self.api._get_url_server()`
  instead of broken `self.api.server_url`; split HTTPError vs RequestException
  handlers for precise error messages
- migrate_users.py: add --url/--source-url args for self-hosted instances;
  fix _resolve_server_url (base64 padding, error handling, explicit-URL priority);
  require --source-api-key explicitly instead of silently reusing COMET_API_KEY;
  warn when source and dest URL+key are identical; handle dict-list API response
  in _get_existing_workspaces; show per-user output in dry-run; add
  --failures-output flag; use json= kwarg in _add_member; fix docstring examples
- tests/unit/test_migrate_users.py: 17 new unit tests covering _resolve_server_url,
  _get_existing_workspaces, _add_member, and dry-run behaviour

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: Douglas Blank

README.md

@@ -333,24 +333,25 @@ 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 ...]
 cometx migrate-users --api-key DEST_KEY --chargeback-report /path/to/report.json [FLAGS ...]
 ```
 
 **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
 

cometx/cli/__init__.py

@@ -108,6 +108,7 @@ def main(raw_args=sys.argv[1:]):
     add_subparser(subparsers, count, "count")
     add_subparser(subparsers, reproduce, "reproduce")
     add_subparser(subparsers, config, "config")
+    add_subparser(subparsers, migrate_users, "migrate-users")
     add_subparser(subparsers, rename_duplicates, "rename-duplicates")
     add_subparser(subparsers, smoke_test, "smoke-test")
     add_subparser(subparsers, migrate_users, "migrate-users")

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
@@ -834,10 +836,8 @@ def copy(self, source, destination, symlink, ignore, debug, sync, create_workspa
                 print(
                     f"Workspace {workspace_dst!r} does not exist, attempting to create it..."
                 )
+                url = f"{self.api._get_url_server()}/api/rest/v2/write/workspace/new"
                 try:
-                    import requests
-
-                    url = f"{self.api.server_url}/api/rest/v2/write/workspace/new"
                     response = requests.post(
                         url,
                         json={"name": workspace_dst},
@@ -848,12 +848,18 @@ def copy(self, source, destination, symlink, ignore, debug, sync, create_workspa
                     )
                     response.raise_for_status()
                     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 "

cometx/cli/migrate_users.py

@@ -20,16 +20,17 @@
 
 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
 import base64
 import json
 import os
 import sys
+
 import requests
 
 ADDITIONAL_ARGS = False
@@ -40,13 +41,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 +85,36 @@ 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 ``*<base64>``).
+    3. COMET_CLOUD_URL as a last resort.
+    """
+    if explicit_url:
+        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 default
+
+    return COMET_CLOUD_URL
 
 
 def _fetch_chargeback_report(server_url, source_api_key):
@@ -101,7 +139,11 @@ 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()
+    # The endpoint may return a list of strings or a list of dicts with a "name" key.
+    if data and isinstance(data[0], dict):
+        return {ws["name"] for ws in data}
+    return set(data)
 
 
 def _create_workspace(dest_url, headers, workspace_name):
@@ -126,7 +168,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 +205,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 +265,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 +302,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 +331,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 +360,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):