refactor(cli): clean up parser and help text

* Remove shared args (--config, --site, --data-dir) from root parser.
* Replace if/elif chain in main() with cmd_* functions and argparse
* README.md: fix arg order in extract and process examples

Author: msherman64

README.md

@@ -60,7 +60,7 @@ The following command shows how to create a SQL user and set these permissions,
 and outputs a series of SQL commands to run manually, setting up this least-privilege user.
 
 ```bash
-chameleon-usage print-grant-sql --user ccusage
+chameleon-usage extract --print-grant-sql --grant-user ccusage
 ```
 
 Example, generated for user `ccusage` and default host `%`
@@ -132,18 +132,17 @@ site_key_name:
 The following command uses the provided config to fetch data from the OpenStack DB and write it to file for later analysis.
 
 ```bash
-chameleon-usage --config etc/site.yml --site site_key_name extract
+chameleon-usage extract --config etc/site.yml --site site_key_name
 ```
 
 ### Generating the usage timeline
 
 Finally, the `process` command takes the raw data and configuration, and outputs the usage timeline.
 
 ```bash
-chameleon-usage \
+chameleon-usage process \
   --config etc/site.yml \
   --site site_key_name \
-  process \
   --output output/usage \
   --start-date 2024-01-01 \
   --end-date 2024-12-31

src/chameleon_usage/cli.py

@@ -4,56 +4,149 @@
 from datetime import datetime
 from pathlib import Path
 
+import polars as pl
+
 from chameleon_usage.config import SiteConfig, load_config
+from chameleon_usage.exceptions import RawTableLoadError, log_raw_table_load_error
+from chameleon_usage.extract.dump_db import dump_to_parquet, generate_grant_sql
+from chameleon_usage.ingest import load_intervals
 from chameleon_usage.output import compat
+from chameleon_usage.pipeline import run_pipeline
+from chameleon_usage.schemas import PipelineSpec
 
 logger = logging.getLogger(__name__)
 
 
-def _add_shared_args(
-    parser: argparse.ArgumentParser, default: object | None = None
-) -> None:
-    parser.add_argument(
-        "--config",
-        help="Path to etc/site.yml (required for process command)",
-        default=default,
-    )
-    parser.add_argument(
-        "--site",
-        action="append",
-        help="Site key from etc/site.yml (repeatable). Defaults to all sites.",
-        default=default,
-    )
-    parser.add_argument(
-        "--data-dir",
-        help="Local directory or s3://. Overrides config.data_dir if set.",
-        default=default,
+def _resolve_sites(args) -> list[SiteConfig]:
+    """Load config, resolve site list and data_dir overrides."""
+    if not args.config:
+        raise SystemExit("Error: --config required")
+    sites_config = load_config(args.config)
+    site_keys = args.site or list(sites_config.keys())
+    configs = []
+    for key in site_keys:
+        config = sites_config[key]
+        if args.data_dir:
+            config.data_dir = args.data_dir.rstrip("/")
+        configs.append(config)
+    return configs
+
+
+# ── Subcommands ──────────────────────────────────────────────────────────
+
+
+def cmd_extract(args):
+    if args.print_grant_sql:
+        logger.info("%s", generate_grant_sql(args.grant_user, args.grant_host))
+        return
+
+    # Priority: --db-uri > $DATABASE_URI > config.db_uri
+    db_uri = args.db_uri or os.environ.get("DATABASE_URI")
+
+    if args.config:
+        for config in _resolve_sites(args):
+            site_db_uri = db_uri or config.db_uri
+            if not site_db_uri:
+                raise SystemExit(f"Error: no db_uri for site {config.key}")
+            if not config.data_dir:
+                raise SystemExit(f"Error: no output path for site {config.key}")
+            logger.info("Extracting %s...", config.key)
+            dump_to_parquet(site_db_uri, config.data_dir)
+    elif db_uri:
+        if not args.data_dir:
+            raise SystemExit("Error: --data-dir required when not using --config")
+        dump_to_parquet(db_uri, args.data_dir.rstrip("/"))
+    else:
+        raise SystemExit("Error: --db-uri, $DATABASE_URI, or --config required")
+
+
+def cmd_process(args):
+    start = datetime.fromisoformat(args.start_date)
+    end = datetime.fromisoformat(args.end_date)
+    spec = PipelineSpec(
+        group_cols=("metric", "resource", "site", "collector_type"),
+        time_range=(start, end),
     )
 
+    output_base = Path(args.output)
+    output_base.mkdir(parents=True, exist_ok=True)
+    usage_frames: list[pl.DataFrame] = []
+    export_uri = args.export_uri or os.environ.get("EXPORT_URI")
 
-def parse_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser()
-    _add_shared_args(parser)
+    for config in _resolve_sites(args):
+        if not config.data_dir:
+            raise SystemExit(f"Error: no data_dir for site {config.key}")
+
+        try:
+            intervals = load_intervals(config.data_dir, spec.time_range).with_columns(
+                pl.lit(config.key).alias("site")
+            )
+            site_usage = run_pipeline(intervals, spec, resample_interval=args.resample)
+        except RawTableLoadError as exc:
+            log_raw_table_load_error(logger, config.key, exc)
+            continue
+        except Exception:
+            logger.exception("[%s] unhandled exception", config.key)
+            raise
 
+        output_dir = output_base / config.key
+        output_dir.mkdir(parents=True, exist_ok=True)
+        site_usage_df = site_usage.collect()
+        site_usage_df.write_parquet(output_dir / "usage.parquet")
+        usage_frames.append(site_usage_df)
+
+    if export_uri and usage_frames:
+        combined_usage = pl.concat(usage_frames)
+        compat_output = compat.to_compat_format(combined_usage)
+        compat.write_compat_to_db(compat_output, db_uri=export_uri)
+
+
+# ── Parser ───────────────────────────────────────────────────────────────
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Chameleon cloud resource usage reporting.",
+    )
     subparsers = parser.add_subparsers(dest="command", required=True)
 
-    extract = subparsers.add_parser("extract")
-    _add_shared_args(extract, default=argparse.SUPPRESS)
+    extract = subparsers.add_parser(
+        "extract", help="Dump database tables to parquet files"
+    )
+    extract.add_argument("--config", help="Path to site config YAML.")
+    extract.add_argument(
+        "--site", action="append", help="Site key (repeatable). Defaults to all sites."
+    )
+    extract.add_argument(
+        "--data-dir", help="Local directory or s3://. Overrides config data_dir."
+    )
     extract.add_argument(
         "--db-uri",
         help="Database URI (mysql://user:pass@host:port). Falls back to $DATABASE_URI.",
     )
-
-    grant_sql = subparsers.add_parser(
-        "print-grant-sql", help="Print SQL to grant read access"
+    extract.add_argument(
+        "--print-grant-sql",
+        action="store_true",
+        help="Print SQL to grant read access and exit.",
     )
-    grant_sql.add_argument("--user", default="usage_exporter", help="MySQL username")
-    grant_sql.add_argument(
-        "--host", default="%", help="MySQL host patterns (default: %%)"
+    extract.add_argument(
+        "--grant-user", default="usage_exporter", help="MySQL username for grant SQL."
     )
+    extract.add_argument(
+        "--grant-host", default="%%", help="MySQL host pattern for grant SQL."
+    )
+    extract.set_defaults(func=cmd_extract)
 
-    process = subparsers.add_parser("process")
-    _add_shared_args(process, default=argparse.SUPPRESS)
+    process = subparsers.add_parser(
+        "process", help="Compute usage timelines from parquet data"
+    )
+    process.add_argument("--config", required=True, help="Path to site config YAML.")
+    process.add_argument(
+        "--site", action="append", help="Site key (repeatable). Defaults to all sites."
+    )
+    process.add_argument(
+        "--data-dir", help="Local directory or s3://. Overrides config data_dir."
+    )
     process.add_argument(
         "--output",
         required=True,
@@ -71,130 +164,21 @@ def parse_args() -> argparse.Namespace:
     )
     process.add_argument(
         "--resample",
-        help="Optional resample interval (e.g. 1d, 7d).",
+        help="Resample interval (e.g. 1h, 1d, 7d).",
     )
-
     process.add_argument(
         "--export-uri",
-        help="Optional DB URI to push output data into. Falls back to env var EXPORT_URI.",
+        help="DB URI to push output data into. Falls back to $EXPORT_URI.",
     )
+    process.set_defaults(func=cmd_process)
 
     return parser.parse_args()
 
 
-def process_site(config: SiteConfig, spec, resample: str):
-    """Process a site's data through the pipeline. Requires [pipeline] extras."""
-    import polars as pl
-
-    from chameleon_usage.ingest import load_intervals
-    from chameleon_usage.pipeline import run_pipeline
-
-    data_dir = config.data_dir
-    if data_dir is None:
-        raise SystemExit(f"Error: no data_dir for site {config.key}")
-
-    intervals = load_intervals(data_dir, spec.time_range).collect().lazy()
-
-    cols = set(intervals.collect_schema().names())
-    if "site" not in cols:
-        intervals = intervals.with_columns(pl.lit(config.key).alias("site"))
-
-    return run_pipeline(intervals, spec, resample_interval=resample)
-
-
 def main() -> None:
     logging.basicConfig(level=logging.INFO, format="%(message)s")
     args = parse_args()
-
-    if args.command == "print-grant-sql":
-        from chameleon_usage.extract.dump_db import generate_grant_sql
-
-        logger.info("%s", generate_grant_sql(args.user, args.host))
-        return
-
-    if args.command == "extract":
-        from chameleon_usage.extract.dump_db import dump_to_parquet
-
-        # Priority: --db-uri > $DATABASE_URI > config.db_uri
-        db_uri = args.db_uri or os.environ.get("DATABASE_URI")
-
-        if args.config:
-            sites_config = load_config(args.config)
-            site_keys = args.site or list(sites_config.keys())
-            for site_key in site_keys:
-                config = sites_config[site_key]
-                site_db_uri = db_uri or config.db_uri
-                if not site_db_uri:
-                    raise SystemExit(f"Error: no db_uri for site {site_key}")
-                # Priority: --data-dir > config.data_dir
-                output_path = args.data_dir or config.data_dir
-                if not output_path:
-                    raise SystemExit(f"Error: no output path for site {site_key}")
-                output_path = output_path.rstrip("/")
-                logger.info("Extracting %s...", site_key)
-                dump_to_parquet(site_db_uri, output_path)
-        elif db_uri:
-            if not args.data_dir:
-                raise SystemExit("Error: --data-dir required when not using --config")
-            dump_to_parquet(db_uri, args.data_dir.rstrip("/"))
-        else:
-            raise SystemExit("Error: --db-uri, $DATABASE_URI, or --config required")
-        return
-
-    if args.command == "process":
-        import polars as pl
-
-        if not args.config:
-            raise SystemExit("Error: --config required for process command")
-
-        from chameleon_usage.exceptions import (
-            RawTableLoadError,
-            log_raw_table_load_error,
-        )
-        from chameleon_usage.schemas import PipelineSpec
-
-        sites_config = load_config(args.config)
-        site_keys = args.site or list(sites_config.keys())
-
-        start = datetime.fromisoformat(args.start_date)
-        end = datetime.fromisoformat(args.end_date)
-        spec = PipelineSpec(
-            group_cols=("metric", "resource", "site", "collector_type"),
-            time_range=(start, end),
-        )
-
-        output_base = Path(args.output)
-        output_base.mkdir(parents=True, exist_ok=True)
-        usage_frames: list[pl.DataFrame] = []
-        # Priority: --export-uri > EXPORT_URI
-        export_uri = args.export_uri or os.environ.get("EXPORT_URI")
-
-        for site_key in site_keys:
-            config = sites_config[site_key]
-            if args.data_dir:
-                config.data_dir = args.data_dir.rstrip("/")
-            if not config.data_dir:
-                raise SystemExit(f"Error: no data_dir for site {site_key}")
-
-            try:
-                site_usage = process_site(config, spec, args.resample)
-            except RawTableLoadError as exc:
-                log_raw_table_load_error(logger, site_key, exc)
-                continue
-            except Exception:
-                logger.exception("[%s] unhandled exception", site_key)
-                raise
-
-            output_dir = output_base / site_key
-            output_dir.mkdir(parents=True, exist_ok=True)
-            site_usage_df = site_usage.collect()
-            site_usage_df.write_parquet(output_dir / "usage.parquet")
-            usage_frames.append(site_usage_df)
-
-        if export_uri and usage_frames:
-            combined_usage = pl.concat(usage_frames)
-            compat_output = compat.to_compat_format(combined_usage)
-            compat.write_compat_to_db(compat_output, db_uri=export_uri)
+    args.func(args)
 
 
 if __name__ == "__main__":