Fix review issues in migrate-users / create-workspaces PR

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/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 "

cometx/cli/migrate_users.py

@@ -20,16 +20,18 @@
 
 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 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 ``*<base64>``).
+    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):