diff --git a/pyproject.toml b/pyproject.toml index 98a2301..b269292 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -31,6 +31,8 @@ dependencies = [ "sentence-transformers>=2.2.0", "rich>=13.0.0", "httpx>=0.25.0", + "tomli_w>=1.0.0", + "tomli>=2.0.0; python_version < '3.11'", ] [project.optional-dependencies] diff --git a/src/tpuff/cli.py b/src/tpuff/cli.py index 531437b..fecb837 100644 --- a/src/tpuff/cli.py +++ b/src/tpuff/cli.py @@ -5,9 +5,11 @@ from tpuff import __version__ from tpuff.commands.delete import delete from tpuff.commands.edit import edit +from tpuff.commands.env import env from tpuff.commands.export import export from tpuff.commands.get import get from tpuff.commands.list import list_cmd +from tpuff.commands.scan import scan from tpuff.commands.schema import schema from tpuff.commands.search import search from tpuff.utils.output import resolve_output_mode @@ -26,12 +28,19 @@ default=None, help="Output format: human (rich tables) or plain (pipe-delimited). Auto-detects TTY if omitted.", ) +@click.option( + "--env", + "env_name", + default=None, + help="Use a specific named environment from ~/.tpuff/config.toml.", +) @click.pass_context -def cli(ctx: click.Context, debug: bool, output: str | None) -> None: +def cli(ctx: click.Context, debug: bool, output: str | None, env_name: str | None) -> None: """tpuff - CLI tool for Turbopuffer vector database.""" ctx.ensure_object(dict) ctx.obj["debug"] = debug ctx.obj["output_mode"] = resolve_output_mode(output) + ctx.obj["env"] = env_name # Register commands @@ -44,7 +53,9 @@ def cli(ctx: click.Context, debug: bool, output: str | None) -> None: cli.add_command(get) cli.add_command(export) cli.add_command(export, name="metrics") # alias +cli.add_command(scan) cli.add_command(schema) +cli.add_command(env) def main() -> None: diff --git a/src/tpuff/client.py b/src/tpuff/client.py index ac837d6..8685260 100644 --- a/src/tpuff/client.py +++ b/src/tpuff/client.py @@ -3,8 +3,10 @@ import os import sys +import click from turbopuffer import Turbopuffer +from tpuff.config import get_active_env, get_env from tpuff.utils.debug import debug_log from tpuff.utils.regions import DEFAULT_REGION @@ -12,24 +14,67 @@ _client_cache: dict[str, Turbopuffer] = {} +def _resolve_config(env_name: str | None = None) -> tuple[str | None, str | None, str | None]: + """Resolve api_key, region, base_url from config file. + + Args: + env_name: Optional specific environment name to use. + + Returns: + (api_key, region, base_url) from config, any may be None. + """ + if env_name: + env_dict = get_env(env_name) + if not env_dict: + print(f"Error: environment '{env_name}' not found in config", file=sys.stderr) + print("Run 'tpuff env list' to see available environments.", file=sys.stderr) + sys.exit(1) + return env_dict.get("api_key"), env_dict.get("region"), env_dict.get("base_url") + + active = get_active_env() + if active: + _, env_dict = active + return env_dict.get("api_key"), env_dict.get("region"), env_dict.get("base_url") + + return None, None, None + + def get_turbopuffer_client(region_override: str | None = None) -> Turbopuffer: """Get a Turbopuffer client configured with the appropriate region. + Resolution order: + 1. Environment variables (always win if set) + 2. Config file (~/.tpuff/config.toml) active or --env specified environment + 3. Error with helpful message + Args: region_override: Optional region to use instead of environment variable. Returns: Configured Turbopuffer client. """ - api_key = os.environ.get("TURBOPUFFER_API_KEY") + # Get --env flag from Click context if available + env_name = None + ctx = click.get_current_context(silent=True) + if ctx and ctx.obj: + env_name = ctx.obj.get("env") + + # Resolve from config file + cfg_api_key, cfg_region, cfg_base_url = _resolve_config(env_name) + + # Environment variables take priority, then config + api_key = os.environ.get("TURBOPUFFER_API_KEY") or cfg_api_key if not api_key: - print("Error: TURBOPUFFER_API_KEY environment variable is not set", file=sys.stderr) + print( + "Error: No API key found. Set TURBOPUFFER_API_KEY or run 'tpuff env add '.", + file=sys.stderr, + ) sys.exit(1) - # Determine the region - base_url = os.environ.get("TURBOPUFFER_BASE_URL") - region = region_override or os.environ.get("TURBOPUFFER_REGION") or DEFAULT_REGION + # Determine the region and base_url + base_url = os.environ.get("TURBOPUFFER_BASE_URL") or cfg_base_url or None + region = region_override or os.environ.get("TURBOPUFFER_REGION") or cfg_region or DEFAULT_REGION # Cache key is the region or base_url cache_key = base_url or region diff --git a/src/tpuff/commands/delete.py b/src/tpuff/commands/delete.py index 15e029d..8ee73ce 100644 --- a/src/tpuff/commands/delete.py +++ b/src/tpuff/commands/delete.py @@ -6,6 +6,7 @@ from rich.console import Console from tpuff.client import get_turbopuffer_client +from tpuff.config import CONFIG_FILE, deletes_allowed from tpuff.utils.debug import debug_log console = Console() @@ -30,6 +31,11 @@ def delete( region: str | None, ) -> None: """Delete namespace(s).""" + if not deletes_allowed(): + console.print("[red]Error: Deletes are disabled.[/red]") + console.print(f"[dim]To enable, add [bold]allow_deletes = true[/bold] to {CONFIG_FILE}[/dim]") + sys.exit(1) + client = get_turbopuffer_client(region) try: diff --git a/src/tpuff/commands/env.py b/src/tpuff/commands/env.py new file mode 100644 index 0000000..6a892be --- /dev/null +++ b/src/tpuff/commands/env.py @@ -0,0 +1,127 @@ +"""Environment management commands for tpuff.""" + +import click +from rich.console import Console +from rich.table import Table + +from tpuff.config import add_env, get_active_env, list_envs, remove_env, set_active +from tpuff.utils.output import is_plain, print_table_plain +from tpuff.utils.regions import DEFAULT_REGION, TURBOPUFFER_REGIONS, is_valid_region + +console = Console() + + +def mask_key(api_key: str) -> str: + """Mask an API key, showing only the first 8 characters.""" + if len(api_key) <= 8: + return api_key + return api_key[:8] + "..." + + +@click.group("env", context_settings={"help_option_names": ["-h", "--help"]}) +def env(): + """Manage tpuff environments.""" + pass + + +@env.command("add") +@click.argument("name") +@click.pass_context +def env_add(ctx: click.Context, name: str) -> None: + """Add a new environment (interactive setup).""" + api_key = click.prompt("API key") + region = click.prompt("Region", default=DEFAULT_REGION) + if not is_valid_region(region): + raise click.BadParameter( + f"Unknown region '{region}'. Valid regions:\n " + + "\n ".join(TURBOPUFFER_REGIONS), + param_hint="'region'", + ) + base_url = click.prompt("Base URL (optional, press Enter to skip)", default="") + + add_env(name, api_key, region, base_url) + click.echo(f"Environment '{name}' added.") + + # Show if it was set as active + active = get_active_env() + if active and active[0] == name: + click.echo("Set as active environment.") + + +@env.command("use") +@click.argument("name") +def env_use(name: str) -> None: + """Switch active environment.""" + set_active(name) + click.echo(f"Switched to environment '{name}'.") + + +@env.command("list") +@click.pass_context +def env_list(ctx: click.Context) -> None: + """List all environments.""" + envs = list_envs() + if not envs: + click.echo("No environments configured. Run 'tpuff env add ' to add one.") + return + + plain = is_plain(ctx) + headers = ["", "Name", "Region", "API Key"] + rows = [] + for name, env_dict, is_active in envs: + marker = "*" if is_active else "" + rows.append([ + marker, + name, + env_dict.get("region", DEFAULT_REGION), + mask_key(env_dict.get("api_key", "")), + ]) + + if plain: + print_table_plain(headers, rows) + else: + table = Table(show_header=True, header_style="cyan") + for h in headers: + table.add_column(h) + for r in rows: + # Highlight active env + if r[0] == "*": + table.add_row( + "[green]*[/green]", + f"[bold]{r[1]}[/bold]", + r[2], + f"[dim]{r[3]}[/dim]", + ) + else: + table.add_row("", r[1], r[2], f"[dim]{r[3]}[/dim]") + console.print(table) + + +# Alias +env.add_command(env_list, name="ls") + + +@env.command("rm") +@click.argument("name") +@click.confirmation_option(prompt="Are you sure you want to remove this environment?") +def env_rm(name: str) -> None: + """Remove an environment.""" + remove_env(name) + click.echo(f"Environment '{name}' removed.") + + +@env.command("show") +def env_show() -> None: + """Show the current active environment.""" + active = get_active_env() + if not active: + click.echo("No active environment. Run 'tpuff env add ' to add one.") + return + + name, env_dict = active + click.echo(f"Active environment: {name}") + click.echo(f" Region: {env_dict.get('region', DEFAULT_REGION)}") + click.echo(f" API Key: {mask_key(env_dict.get('api_key', ''))}") + base_url = env_dict.get("base_url", "") + if base_url: + click.echo(f" Base URL: {base_url}") diff --git a/src/tpuff/commands/list.py b/src/tpuff/commands/list.py index bb8b027..c10abc9 100644 --- a/src/tpuff/commands/list.py +++ b/src/tpuff/commands/list.py @@ -13,8 +13,6 @@ from tpuff.utils.metadata_fetcher import ( NamespaceWithMetadata, fetch_namespaces_with_metadata, - get_index_status, - get_unindexed_bytes, ) from tpuff.utils.output import is_plain, print_table_plain, status_print @@ -228,12 +226,6 @@ def display_namespaces( console.print("No namespaces found") return - status_print( - ctx, - f"\n[bold]Found {len(namespaces_with_metadata)} namespace(s):[/bold]\n", - console, - ) - # Sort by updated_at in descending order (most recent first) def sort_key(item: NamespaceWithMetadata): if not item.metadata: @@ -248,90 +240,61 @@ def sort_key(item: NamespaceWithMetadata): namespaces_with_metadata.sort(key=sort_key, reverse=True) - # Build headers - headers = ["Namespace"] - if all_regions: - headers.append("Region") - headers.extend(["Rows", "Logical Bytes", "Index Status", "Unindexed Bytes"]) - if include_recall: - headers.append("Recall") - headers.append("Updated") - - # Collect row data - table_rows = [] - for item in namespaces_with_metadata: - if item.metadata: - index_status = get_index_status(item.metadata) - unindexed = get_unindexed_bytes(item.metadata) - - if plain: + if plain: + # Plain mode: pipe-delimited for scripts + headers = ["Namespace"] + if all_regions: + headers.append("Region") + headers.extend(["Rows", "Size", "Updated"]) + if include_recall: + headers.append("Recall") + + table_rows = [] + for item in namespaces_with_metadata: + if item.metadata: row = [item.namespace_id] if all_regions: row.append(item.region or "") row.extend([ f"{item.metadata.approx_row_count:,}", format_bytes(item.metadata.approx_logical_bytes), - index_status, - format_bytes(unindexed), + format_updated_at(item.metadata.updated_at, plain=True), ]) if include_recall: row.append(format_recall(item.recall, plain=True)) - row.append(format_updated_at(item.metadata.updated_at, plain=True)) else: - index_status_display = ( - "[green]up-to-date[/green]" - if index_status == "up-to-date" - else "[red]updating[/red]" - ) - unindexed_display = ( - f"[red]{format_bytes(unindexed)}[/red]" - if unindexed > 0 - else format_bytes(0) - ) - - row = [f"[bold]{item.namespace_id}[/bold]"] - if all_regions and item.region: - row.append(f"[dim]{item.region}[/dim]") - row.extend([ - f"{item.metadata.approx_row_count:,}", - format_bytes(item.metadata.approx_logical_bytes), - index_status_display, - unindexed_display, - ]) - if include_recall: - row.append(format_recall(item.recall)) - row.append(format_updated_at(item.metadata.updated_at)) - - table_rows.append(row) - else: - if plain: row = [item.namespace_id] if all_regions: row.append(item.region or "") - row.extend(["N/A"] * 4) + row.extend(["N/A", "N/A", "N/A"]) if include_recall: row.append("N/A") - row.append("N/A") - else: - row = [f"[bold]{item.namespace_id}[/bold]"] - if all_regions and item.region: - row.append(f"[dim]{item.region}[/dim]") - row.extend(["[dim]N/A[/dim]"] * 4) - if include_recall: - row.append("[dim]N/A[/dim]") - row.append("[dim]N/A[/dim]") - table_rows.append(row) - if plain: print_table_plain(headers, table_rows) else: - table = Table(show_header=True, header_style="cyan") - for h in headers: - table.add_column(h) - for r in table_rows: - table.add_row(*r) - console.print(table) + # Human mode: lean ll-style output + for item in namespaces_with_metadata: + if item.metadata: + rows = f"{item.metadata.approx_row_count:,} rows" + size = format_bytes(item.metadata.approx_logical_bytes) + date = format_updated_at(item.metadata.updated_at) + region_part = f" [dim]{item.region}[/dim]" if all_regions and item.region else "" + recall_part = f" {format_recall(item.recall)}" if include_recall else "" + console.print( + f"[bold]{item.namespace_id}[/bold] " + f"[cyan]{rows:>12}[/cyan] " + f"[dim]{size:>10}[/dim]{region_part}{recall_part} " + f"[dim]{date}[/dim]" + ) + else: + region_part = f" [dim]{item.region}[/dim]" if all_regions and item.region else "" + console.print( + f"[bold]{item.namespace_id}[/bold] " + f"[dim]{'N/A':>12}[/dim] " + f"[dim]{'N/A':>10}[/dim]{region_part} " + f"[dim]N/A[/dim]" + ) @click.command("list", context_settings={"help_option_names": ["-h", "--help"]}) diff --git a/src/tpuff/commands/scan.py b/src/tpuff/commands/scan.py new file mode 100644 index 0000000..4502fd9 --- /dev/null +++ b/src/tpuff/commands/scan.py @@ -0,0 +1,126 @@ +"""Scan command for tpuff CLI — extract unique field values from a namespace.""" + +import json +import sys + +import click +from rich.console import Console + +from tpuff.client import get_namespace +from tpuff.utils.debug import debug_log +from tpuff.utils.output import is_plain + +# stderr console so progress doesn't pollute JSON output +err_console = Console(stderr=True) + + +@click.command("scan", context_settings={"help_option_names": ["-h", "--help"]}) +@click.option("-n", "--namespace", required=True, help="Namespace to scan") +@click.option("--field", required=True, help="Field name to extract unique values from") +@click.option("-r", "--region", help="Override the region (e.g., aws-us-east-1, gcp-us-central1)") +@click.option("--page-size", default=1000, type=int, help="Batch size per query (default: 1000)") +@click.pass_context +def scan( + ctx: click.Context, + namespace: str, + field: str, + region: str | None, + page_size: int, +) -> None: + """Scan a namespace and extract all unique values of a field.""" + plain = is_plain(ctx) + + try: + ns = get_namespace(namespace, region) + + # Get total document count from metadata for progress bar + metadata = ns.metadata() + meta_dict = metadata.model_dump() if hasattr(metadata, "model_dump") else {} + total_docs = meta_dict.get("approx_row_count") or None + + unique_values: set[str] = set() + last_id: str | None = None + total_scanned = 0 + + if not plain: + from rich.progress import ( + BarColumn, + MofNCompleteColumn, + Progress, + SpinnerColumn, + TextColumn, + ) + + progress = Progress( + SpinnerColumn(), + TextColumn("[progress.description]{task.description}"), + BarColumn(bar_width=None), + MofNCompleteColumn(), + console=err_console, + ) + task = progress.add_task("Scanning...", total=total_docs) + progress.start() + + while True: + query_params: dict = { + "rank_by": ["id", "asc"], + "include_attributes": [field], + "top_k": page_size, + } + + if last_id is not None: + query_params["filters"] = ["id", "Gt", last_id] + + debug_log("Scan query", query_params) + result = ns.query(**query_params) + rows = result.rows if hasattr(result, "rows") else [] + + for row in rows: + if hasattr(row, "model_dump"): + row_dict = row.model_dump() + else: + row_dict = {"id": getattr(row, "id", None)} + + value = row_dict.get(field) + if value is not None: + unique_values.add(value) + + last_id = row_dict.get("id", getattr(row, "id", None)) + + total_scanned += len(rows) + + if not plain: + progress.update( + task, + completed=total_scanned, + description=f"{len(unique_values)} unique values", + ) + + debug_log("Scan page", { + "rows": len(rows), + "total_scanned": total_scanned, + "unique_values": len(unique_values), + "last_id": last_id, + }) + + if len(rows) < page_size: + break + + if not plain: + progress.stop() + err_console.print( + f"[green]Done.[/green] Scanned {total_scanned} documents, " + f"found {len(unique_values)} unique values." + ) + + # Output the unique values as a sorted JSON array to stdout + click.echo(json.dumps(sorted(unique_values))) + + except Exception as e: + if not plain: + try: + progress.stop() + except Exception: + pass + err_console.print(f"[red]Error: {e}[/red]") + sys.exit(1) diff --git a/src/tpuff/config.py b/src/tpuff/config.py new file mode 100644 index 0000000..686a062 --- /dev/null +++ b/src/tpuff/config.py @@ -0,0 +1,116 @@ +"""Config management for tpuff environments. + +Stores named environments in ~/.tpuff/config.toml with API keys, regions, and base URLs. +""" + +import os +import stat +import sys +from pathlib import Path + +import tomli_w + +CONFIG_DIR = Path.home() / ".tpuff" +CONFIG_FILE = CONFIG_DIR / "config.toml" + + +def _read_toml(path: Path) -> dict: + """Read a TOML file, using tomllib (3.11+) or tomli fallback.""" + try: + import tomllib + except ModuleNotFoundError: + import tomli as tomllib + + with open(path, "rb") as f: + return tomllib.load(f) + + +def load_config() -> dict: + """Load config from ~/.tpuff/config.toml. Returns empty dict if not found.""" + if not CONFIG_FILE.exists(): + return {} + return _read_toml(CONFIG_FILE) + + +def save_config(config: dict) -> None: + """Write config to ~/.tpuff/config.toml with secure permissions.""" + CONFIG_DIR.mkdir(parents=True, exist_ok=True) + os.chmod(CONFIG_DIR, stat.S_IRWXU) # 0o700 + + CONFIG_FILE.write_bytes(tomli_w.dumps(config).encode()) + os.chmod(CONFIG_FILE, stat.S_IRUSR | stat.S_IWUSR) # 0o600 + + +def get_active_env() -> tuple[str, dict] | None: + """Return (name, env_dict) for the active environment, or None.""" + config = load_config() + active = config.get("active") + if not active: + return None + envs = config.get("envs", {}) + if active not in envs: + return None + return active, envs[active] + + +def get_env(name: str) -> dict | None: + """Get a specific named environment, or None if not found.""" + config = load_config() + return config.get("envs", {}).get(name) + + +def list_envs() -> list[tuple[str, dict, bool]]: + """Return list of (name, env_dict, is_active) for all environments.""" + config = load_config() + active = config.get("active") + envs = config.get("envs", {}) + return [(name, env_dict, name == active) for name, env_dict in envs.items()] + + +def add_env(name: str, api_key: str, region: str, base_url: str = "") -> None: + """Add or overwrite an environment.""" + config = load_config() + if "envs" not in config: + config["envs"] = {} + env = {"api_key": api_key, "region": region} + if base_url: + env["base_url"] = base_url + config["envs"][name] = env + # Set as active if it's the first env + if "active" not in config or not config["active"]: + config["active"] = name + save_config(config) + + +def remove_env(name: str) -> None: + """Remove an environment. Errors if not found.""" + config = load_config() + envs = config.get("envs", {}) + if name not in envs: + print(f"Error: environment '{name}' not found", file=sys.stderr) + sys.exit(1) + del envs[name] + # If we removed the active env, clear it or pick another + if config.get("active") == name: + if envs: + config["active"] = next(iter(envs)) + else: + config["active"] = "" + save_config(config) + + +def set_active(name: str) -> None: + """Set the active environment.""" + config = load_config() + envs = config.get("envs", {}) + if name not in envs: + print(f"Error: environment '{name}' not found", file=sys.stderr) + sys.exit(1) + config["active"] = name + save_config(config) + + +def deletes_allowed() -> bool: + """Check if deletes are allowed. Defaults to False.""" + config = load_config() + return config.get("allow_deletes", False) diff --git a/src/tpuff/utils/output.py b/src/tpuff/utils/output.py index e1d04ec..ecc6bf1 100644 --- a/src/tpuff/utils/output.py +++ b/src/tpuff/utils/output.py @@ -33,7 +33,7 @@ def is_plain(ctx: click.Context) -> bool: Returns: True if output mode is "plain". """ - return ctx.obj.get("output_mode") == "plain" + return (ctx.obj or {}).get("output_mode") == "plain" def print_table_plain(headers: list[str], rows: list[list[str]]) -> None: