Spaces:

autodrive
/

dq-simulator

Runtime error

App Files Files Community

dq-simulator / simulator /cli.py

autodrive

Upload 38 files

4815c43 verified 7 months ago

Raw History Blame Contribute Delete

5.43 kB

"""CLI entry point."""

import sys
import click

from . import generator, injector, writer
from .issues import REGISTRY, list_issues


@click.group()
def cli():
    """Data Quality Issue Simulator — generate synthetic user events with injected DQ issues."""


@cli.command("list-issues")
@click.option("--category", "-c", default=None, help="Filter by category")
@click.option("--severity", "-s", default=None, help="Filter by severity (low/medium/high/critical)")
def list_issues_cmd(category, severity):
    """List all available data quality issues."""
    issues = list_issues(category)
    if severity:
        issues = [i for i in issues if i.severity == severity]

    categories = {}
    for issue in issues:
        categories.setdefault(issue.category, []).append(issue)

    total = 0
    for cat, cat_issues in sorted(categories.items()):
        click.echo(f"\n{'─'*60}")
        click.echo(f"  {cat.upper().replace('_', ' ')}  ({len(cat_issues)} issues)")
        click.echo(f"{'─'*60}")
        for i in cat_issues:
            sev_color = {"low": "white", "medium": "yellow", "high": "red", "critical": "bright_red"}
            sev = click.style(f"[{i.severity.upper():8s}]", fg=sev_color.get(i.severity, "white"))
            click.echo(f"  {sev} {i.id}")
            click.echo(f"           {i.description}")
        total += len(cat_issues)

    click.echo(f"\nTotal: {total} issues")


@cli.command("generate")
@click.option("--rows", "-n", default=100_000, show_default=True, help="Number of base rows to generate")
@click.option("--issues", "-i", default="all", show_default=True,
              help="Comma-separated issue IDs, or 'all'")
@click.option("--exclude", "-e", default="", help="Comma-separated issue IDs to exclude")
@click.option("--category", "-c", default=None, help="Only inject issues from this category")
@click.option("--severity", default=None, help="Only inject issues at this severity (low/medium/high/critical)")
@click.option("--rate", "-r", default=None, type=float,
              help="Override injection rate (0.0–1.0). Defaults to each issue's built-in rate.")
@click.option("--output", "-o", default="./output/user_events.parquet", show_default=True,
              help="Output path: local file path or s3://bucket/prefix/")
@click.option("--glue-database", default=None, help="AWS Glue database name (optional)")
@click.option("--glue-table", default="user_events", show_default=True, help="Glue table name")
@click.option("--region", default="us-east-1", show_default=True, help="AWS region")
@click.option("--seed", default=42, show_default=True, help="Random seed for reproducibility")
@click.option("--no-issues", is_flag=True, default=False, help="Skip issue injection (clean data only)")
@click.option("--quiet", "-q", is_flag=True, default=False, help="Suppress per-issue output")
def generate_cmd(rows, issues, exclude, category, severity, rate, output,
                 glue_database, glue_table, region, seed, no_issues, quiet):
    """Generate synthetic user events and inject data quality issues."""

    click.echo(f"\n{'='*60}")
    click.echo(f"  DQ Issue Simulator")
    click.echo(f"{'='*60}")

    # 1. Generate clean data
    click.echo(f"\n[1/3] Generating {rows:,} clean user events (seed={seed}) ...")
    df = generator.generate(rows, seed=seed)
    click.echo(f"      Shape: {df.shape[0]:,} rows × {df.shape[1]} columns")

    # 2. Inject issues
    if no_issues:
        applied = []
        click.echo("\n[2/3] Skipping issue injection (--no-issues)")
    else:
        # Resolve issue list
        if issues.lower() == "all":
            issue_ids = None  # let injector pick all
        else:
            issue_ids = [i.strip() for i in issues.split(",") if i.strip()]

        # Validate explicit issue IDs
        if issue_ids:
            for iid in issue_ids:
                if iid not in REGISTRY:
                    click.echo(f"Error: unknown issue '{iid}'. Run 'list-issues' to see all.", err=True)
                    sys.exit(1)

        exclude_ids = [i.strip() for i in exclude.split(",") if i.strip()]

        # Apply category/severity filters when using 'all'
        if issue_ids is None:
            filtered = list_issues(category)
            if severity:
                filtered = [i for i in filtered if i.severity == severity]
            issue_ids = [i.id for i in filtered]

        n_issues = len(issue_ids) - len([e for e in exclude_ids if e in (issue_ids or [])])
        click.echo(f"\n[2/3] Injecting {n_issues} issues ...")

        df, applied = injector.inject(
            df,
            issue_ids=issue_ids,
            exclude_ids=exclude_ids or None,
            rate_override=rate,
            seed=seed,
            verbose=not quiet,
        )
        click.echo(f"\n      Applied: {len(applied)} issues | Final shape: {df.shape[0]:,} rows × {df.shape[1]} columns")

    # 3. Write output
    click.echo(f"\n[3/3] Writing parquet to: {output}")
    out_path = writer.write(
        df,
        output=output,
        glue_database=glue_database,
        glue_table=glue_table,
        region=region,
    )
    click.echo(f"      Written: {out_path}")
    if glue_database:
        click.echo(f"      Glue table updated: {glue_database}.{glue_table}")

    click.echo(f"\n{'='*60}")
    click.echo(f"  Done. {len(applied)} issues injected into {df.shape[0]:,} rows.")
    click.echo(f"{'='*60}\n")

	"""CLI entry point."""

	import sys
	import click

	from . import generator, injector, writer
	from .issues import REGISTRY, list_issues


	@click.group()
	def cli():
	"""Data Quality Issue Simulator — generate synthetic user events with injected DQ issues."""


	@cli.command("list-issues")
	@click.option("--category", "-c", default=None, help="Filter by category")
	@click.option("--severity", "-s", default=None, help="Filter by severity (low/medium/high/critical)")
	def list_issues_cmd(category, severity):
	"""List all available data quality issues."""
	issues = list_issues(category)
	if severity:
	issues = [i for i in issues if i.severity == severity]

	categories = {}
	for issue in issues:
	categories.setdefault(issue.category, []).append(issue)

	total = 0
	for cat, cat_issues in sorted(categories.items()):
	click.echo(f"\n{'─'*60}")
	click.echo(f" {cat.upper().replace('_', ' ')} ({len(cat_issues)} issues)")
	click.echo(f"{'─'*60}")
	for i in cat_issues:
	sev_color = {"low": "white", "medium": "yellow", "high": "red", "critical": "bright_red"}
	sev = click.style(f"[{i.severity.upper():8s}]", fg=sev_color.get(i.severity, "white"))
	click.echo(f" {sev} {i.id}")
	click.echo(f" {i.description}")
	total += len(cat_issues)

	click.echo(f"\nTotal: {total} issues")


	@cli.command("generate")
	@click.option("--rows", "-n", default=100_000, show_default=True, help="Number of base rows to generate")
	@click.option("--issues", "-i", default="all", show_default=True,
	help="Comma-separated issue IDs, or 'all'")
	@click.option("--exclude", "-e", default="", help="Comma-separated issue IDs to exclude")
	@click.option("--category", "-c", default=None, help="Only inject issues from this category")
	@click.option("--severity", default=None, help="Only inject issues at this severity (low/medium/high/critical)")
	@click.option("--rate", "-r", default=None, type=float,
	help="Override injection rate (0.0–1.0). Defaults to each issue's built-in rate.")
	@click.option("--output", "-o", default="./output/user_events.parquet", show_default=True,
	help="Output path: local file path or s3://bucket/prefix/")
	@click.option("--glue-database", default=None, help="AWS Glue database name (optional)")
	@click.option("--glue-table", default="user_events", show_default=True, help="Glue table name")
	@click.option("--region", default="us-east-1", show_default=True, help="AWS region")
	@click.option("--seed", default=42, show_default=True, help="Random seed for reproducibility")
	@click.option("--no-issues", is_flag=True, default=False, help="Skip issue injection (clean data only)")
	@click.option("--quiet", "-q", is_flag=True, default=False, help="Suppress per-issue output")
	def generate_cmd(rows, issues, exclude, category, severity, rate, output,
	glue_database, glue_table, region, seed, no_issues, quiet):
	"""Generate synthetic user events and inject data quality issues."""

	click.echo(f"\n{'='*60}")
	click.echo(f" DQ Issue Simulator")
	click.echo(f"{'='*60}")

	# 1. Generate clean data
	click.echo(f"\n[1/3] Generating {rows:,} clean user events (seed={seed}) ...")
	df = generator.generate(rows, seed=seed)
	click.echo(f" Shape: {df.shape[0]:,} rows × {df.shape[1]} columns")

	# 2. Inject issues
	if no_issues:
	applied = []
	click.echo("\n[2/3] Skipping issue injection (--no-issues)")
	else:
	# Resolve issue list
	if issues.lower() == "all":
	issue_ids = None # let injector pick all
	else:
	issue_ids = [i.strip() for i in issues.split(",") if i.strip()]

	# Validate explicit issue IDs
	if issue_ids:
	for iid in issue_ids:
	if iid not in REGISTRY:
	click.echo(f"Error: unknown issue '{iid}'. Run 'list-issues' to see all.", err=True)
	sys.exit(1)

	exclude_ids = [i.strip() for i in exclude.split(",") if i.strip()]

	# Apply category/severity filters when using 'all'
	if issue_ids is None:
	filtered = list_issues(category)
	if severity:
	filtered = [i for i in filtered if i.severity == severity]
	issue_ids = [i.id for i in filtered]

	n_issues = len(issue_ids) - len([e for e in exclude_ids if e in (issue_ids or [])])
	click.echo(f"\n[2/3] Injecting {n_issues} issues ...")

	df, applied = injector.inject(
	df,
	issue_ids=issue_ids,
	exclude_ids=exclude_ids or None,
	rate_override=rate,
	seed=seed,
	verbose=not quiet,
	)
	click.echo(f"\n Applied: {len(applied)} issues \| Final shape: {df.shape[0]:,} rows × {df.shape[1]} columns")

	# 3. Write output
	click.echo(f"\n[3/3] Writing parquet to: {output}")
	out_path = writer.write(
	df,
	output=output,
	glue_database=glue_database,
	glue_table=glue_table,
	region=region,
	)
	click.echo(f" Written: {out_path}")
	if glue_database:
	click.echo(f" Glue table updated: {glue_database}.{glue_table}")

	click.echo(f"\n{'='*60}")
	click.echo(f" Done. {len(applied)} issues injected into {df.shape[0]:,} rows.")
	click.echo(f"{'='*60}\n")