|
| 1 | +"""Click-based CLI for ``unstract.clone``. |
| 2 | +
|
| 3 | +Single ``clone`` command. Platform keys can be passed via flags |
| 4 | +(``--source-key`` / ``--target-key``) or env vars |
| 5 | +(``UNSTRACT_SRC_PLATFORM_KEY`` / ``UNSTRACT_TGT_PLATFORM_KEY``) — env vars |
| 6 | +are preferred so the key never lands in shell history. |
| 7 | +""" |
| 8 | + |
| 9 | +from __future__ import annotations |
| 10 | + |
| 11 | +import logging |
| 12 | +import re |
| 13 | +import sys |
| 14 | +from typing import Any |
| 15 | + |
| 16 | +import click |
| 17 | + |
| 18 | +from unstract.clone.context import ( |
| 19 | + DEFAULT_CONCURRENCY, |
| 20 | + DEFAULT_MAX_FILE_SIZE, |
| 21 | + CloneOptions, |
| 22 | + OrgEndpoint, |
| 23 | +) |
| 24 | +from unstract.clone.exceptions import CloneError |
| 25 | +from unstract.clone.orchestrator import clone as run_clone |
| 26 | + |
| 27 | +_SIZE_UNITS: dict[str, int] = { |
| 28 | + "B": 1, |
| 29 | + "K": 1024, |
| 30 | + "KB": 1024, |
| 31 | + "M": 1024 * 1024, |
| 32 | + "MB": 1024 * 1024, |
| 33 | + "G": 1024 * 1024 * 1024, |
| 34 | + "GB": 1024 * 1024 * 1024, |
| 35 | +} |
| 36 | +_SIZE_RE = re.compile(r"^\s*(\d+(?:\.\d+)?)\s*([A-Za-z]*)\s*$") |
| 37 | + |
| 38 | + |
| 39 | +def _parse_size(value: str) -> int: |
| 40 | + """Accept ``25``, ``25MB``, ``1.5GB`` etc. Returns bytes.""" |
| 41 | + m = _SIZE_RE.match(value) |
| 42 | + if not m: |
| 43 | + raise click.BadParameter(f"can't parse size '{value}'") |
| 44 | + num, unit = m.group(1), m.group(2).upper() or "B" |
| 45 | + if unit not in _SIZE_UNITS: |
| 46 | + raise click.BadParameter( |
| 47 | + f"unknown size unit '{unit}'; use one of {sorted(_SIZE_UNITS)}" |
| 48 | + ) |
| 49 | + return int(float(num) * _SIZE_UNITS[unit]) |
| 50 | + |
| 51 | + |
| 52 | +def _configure_logging(verbose: bool) -> None: |
| 53 | + level = logging.DEBUG if verbose else logging.INFO |
| 54 | + logging.basicConfig( |
| 55 | + level=level, |
| 56 | + format="%(asctime)s %(levelname)-7s %(name)s: %(message)s", |
| 57 | + datefmt="%H:%M:%S", |
| 58 | + ) |
| 59 | + |
| 60 | + |
| 61 | +def _split_csv(value: str | None) -> tuple[str, ...] | None: |
| 62 | + if not value: |
| 63 | + return None |
| 64 | + return tuple(p.strip() for p in value.split(",") if p.strip()) |
| 65 | + |
| 66 | + |
| 67 | +@click.group() |
| 68 | +def cli() -> None: |
| 69 | + """Cloning organizations over the Platform API.""" |
| 70 | + |
| 71 | + |
| 72 | +@cli.command("clone") |
| 73 | +@click.option("--source-url", required=True, help="Base URL of the source deployment") |
| 74 | +@click.option( |
| 75 | + "--source-org", required=True, help="Source organization_id (slug in the URL path)" |
| 76 | +) |
| 77 | +@click.option( |
| 78 | + "--source-key", |
| 79 | + envvar="UNSTRACT_SRC_PLATFORM_KEY", |
| 80 | + required=True, |
| 81 | + help="Source admin's Platform API key (or env UNSTRACT_SRC_PLATFORM_KEY)", |
| 82 | +) |
| 83 | +@click.option("--target-url", required=True, help="Base URL of the target deployment") |
| 84 | +@click.option( |
| 85 | + "--target-org", required=True, help="Target organization_id (slug in the URL path)" |
| 86 | +) |
| 87 | +@click.option( |
| 88 | + "--target-key", |
| 89 | + envvar="UNSTRACT_TGT_PLATFORM_KEY", |
| 90 | + required=True, |
| 91 | + help="Target admin's Platform API key (or env UNSTRACT_TGT_PLATFORM_KEY)", |
| 92 | +) |
| 93 | +@click.option( |
| 94 | + "--dry-run", is_flag=True, help="Plan only — do not POST anything to target" |
| 95 | +) |
| 96 | +@click.option( |
| 97 | + "--include", |
| 98 | + default=None, |
| 99 | + help="Comma-separated phase names to include (default: all)", |
| 100 | +) |
| 101 | +@click.option( |
| 102 | + "--exclude", |
| 103 | + default=None, |
| 104 | + help="Comma-separated phase names to exclude", |
| 105 | +) |
| 106 | +@click.option( |
| 107 | + "--on-name-conflict", |
| 108 | + type=click.Choice(["adopt", "abort"]), |
| 109 | + default="adopt", |
| 110 | + show_default=True, |
| 111 | + help="What to do when a like-named entity exists in target", |
| 112 | +) |
| 113 | +@click.option( |
| 114 | + "--api-prefix", |
| 115 | + default="api/v1", |
| 116 | + show_default=True, |
| 117 | + help="Backend URL prefix (matches deployment's PATH_PREFIX env)", |
| 118 | +) |
| 119 | +@click.option( |
| 120 | + "--file-strategy", |
| 121 | + type=click.Choice(["platform_api", "skip"]), |
| 122 | + default="platform_api", |
| 123 | + show_default=True, |
| 124 | + help="How to move Prompt Studio document files. 'skip' = metadata only.", |
| 125 | +) |
| 126 | +@click.option( |
| 127 | + "--max-file-size", |
| 128 | + default="25MB", |
| 129 | + show_default=True, |
| 130 | + help="Per-file cap for the files phase. Oversize → reported, not aborted.", |
| 131 | +) |
| 132 | +@click.option( |
| 133 | + "--skip-files", |
| 134 | + is_flag=True, |
| 135 | + help="Alias for --file-strategy=skip.", |
| 136 | +) |
| 137 | +@click.option( |
| 138 | + "--concurrency", |
| 139 | + type=click.IntRange(min=1, max=32), |
| 140 | + default=DEFAULT_CONCURRENCY, |
| 141 | + show_default=True, |
| 142 | + help="Per-phase worker count. 1 = strictly sequential.", |
| 143 | +) |
| 144 | +@click.option("-v", "--verbose", is_flag=True, help="Debug logging") |
| 145 | +def clone_cmd( |
| 146 | + source_url: str, |
| 147 | + source_org: str, |
| 148 | + source_key: str, |
| 149 | + target_url: str, |
| 150 | + target_org: str, |
| 151 | + target_key: str, |
| 152 | + dry_run: bool, |
| 153 | + include: str | None, |
| 154 | + exclude: str | None, |
| 155 | + on_name_conflict: str, |
| 156 | + api_prefix: str, |
| 157 | + file_strategy: str, |
| 158 | + max_file_size: str, |
| 159 | + skip_files: bool, |
| 160 | + concurrency: int, |
| 161 | + verbose: bool, |
| 162 | +) -> None: |
| 163 | + """Clone configured resources from one org to another.""" |
| 164 | + _configure_logging(verbose) |
| 165 | + |
| 166 | + effective_strategy = "skip" if skip_files else file_strategy |
| 167 | + try: |
| 168 | + cap_bytes = _parse_size(max_file_size) |
| 169 | + except click.BadParameter as e: |
| 170 | + raise click.UsageError(str(e)) from e |
| 171 | + |
| 172 | + options = CloneOptions( |
| 173 | + dry_run=dry_run, |
| 174 | + include=_split_csv(include), |
| 175 | + exclude=_split_csv(exclude) or (), |
| 176 | + on_name_conflict=on_name_conflict, |
| 177 | + verbose=verbose, |
| 178 | + file_strategy=effective_strategy, |
| 179 | + max_file_size=cap_bytes if cap_bytes is not None else DEFAULT_MAX_FILE_SIZE, |
| 180 | + concurrency=concurrency, |
| 181 | + ) |
| 182 | + |
| 183 | + source = OrgEndpoint( |
| 184 | + base_url=source_url, |
| 185 | + organization_id=source_org, |
| 186 | + platform_key=source_key, |
| 187 | + api_path_prefix=api_prefix, |
| 188 | + ) |
| 189 | + target = OrgEndpoint( |
| 190 | + base_url=target_url, |
| 191 | + organization_id=target_org, |
| 192 | + platform_key=target_key, |
| 193 | + api_path_prefix=api_prefix, |
| 194 | + ) |
| 195 | + |
| 196 | + try: |
| 197 | + report = run_clone(source, target, options) |
| 198 | + except CloneError as e: |
| 199 | + click.echo(f"Clone failed: {e}", err=True) |
| 200 | + sys.exit(2) |
| 201 | + |
| 202 | + click.echo(report.render()) |
| 203 | + if report.aborted or any(p.failed for p in report.phases): |
| 204 | + sys.exit(1) |
| 205 | + |
| 206 | + |
| 207 | +def main(argv: list[str] | None = None) -> Any: |
| 208 | + return cli(args=argv, standalone_mode=True) |
| 209 | + |
| 210 | + |
| 211 | +if __name__ == "__main__": |
| 212 | + main() |
0 commit comments