Spaces:
Runtime error
Runtime error
File size: 5,431 Bytes
4815c43 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 | """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")
|