new site
This commit is contained in:
oliver
@@ -0,0 +1,242 @@
|
||||
"""
|
||||
Plan command for creating execution plans.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import List, Optional
|
||||
from uuid import UUID
|
||||
|
||||
import typer
|
||||
from rich.console import Console
|
||||
from rich.table import Table
|
||||
|
||||
from ..services.supabase import SupabaseService
|
||||
from ..models.schemas import ExecutionPlan, PlanStep, StepStatus
|
||||
from ..config import SEARCH_QUERIES
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
console = Console()
|
||||
|
||||
|
||||
def plan_command(
|
||||
country: str,
|
||||
cities: Optional[List[str]] = None,
|
||||
name: Optional[str] = None,
|
||||
) -> Optional[UUID]:
|
||||
"""
|
||||
Create an execution plan for scraping Odoo partners.
|
||||
|
||||
This generates a plan that specifies which cities to scrape and
|
||||
what steps to perform. The plan can then be executed with the
|
||||
execute command.
|
||||
|
||||
Args:
|
||||
country: Country code (e.g., DE for Germany)
|
||||
cities: List of city slugs to include (default: all cities in country)
|
||||
name: Optional name for the plan
|
||||
|
||||
Returns:
|
||||
Plan UUID if successful
|
||||
"""
|
||||
console.print(f"\n[bold blue]Creating Execution Plan for {country}[/bold blue]\n")
|
||||
|
||||
db = SupabaseService()
|
||||
|
||||
# Verify country exists
|
||||
country_obj = db.get_country_by_code(country.upper())
|
||||
if not country_obj:
|
||||
console.print(f"[red]Country not found: {country}[/red]")
|
||||
console.print("Run 'setup' command first to seed country data.")
|
||||
raise typer.Exit(1)
|
||||
|
||||
# Get cities
|
||||
all_cities = db.list_cities(country_id=country_obj.id)
|
||||
|
||||
if cities:
|
||||
# Filter to specified cities
|
||||
city_slugs = [c.lower() for c in cities]
|
||||
selected_cities = [c for c in all_cities if c.slug in city_slugs]
|
||||
|
||||
missing = set(city_slugs) - {c.slug for c in selected_cities}
|
||||
if missing:
|
||||
console.print(f"[yellow]Warning: Cities not found: {missing}[/yellow]")
|
||||
else:
|
||||
selected_cities = all_cities
|
||||
|
||||
if not selected_cities:
|
||||
console.print("[red]No cities found for this country.[/red]")
|
||||
raise typer.Exit(1)
|
||||
|
||||
# Generate plan name
|
||||
if not name:
|
||||
city_names = [c.slug for c in selected_cities[:3]]
|
||||
name = f"{country.upper()}-{'-'.join(city_names)}"
|
||||
if len(selected_cities) > 3:
|
||||
name += f"-and-{len(selected_cities) - 3}-more"
|
||||
|
||||
# Estimate company counts based on queries
|
||||
language = country_obj.language
|
||||
queries_per_city = len(SEARCH_QUERIES.get(language, SEARCH_QUERIES["en"]))
|
||||
estimated_companies = len(selected_cities) * queries_per_city * 15 # ~15 per query
|
||||
|
||||
# Show plan summary
|
||||
table = Table(title="Execution Plan Summary")
|
||||
table.add_column("Field", style="cyan")
|
||||
table.add_column("Value", style="green")
|
||||
|
||||
table.add_row("Plan Name", name)
|
||||
table.add_row("Country", f"{country_obj.name} ({country_obj.code})")
|
||||
table.add_row("Language", country_obj.language)
|
||||
table.add_row("Cities", str(len(selected_cities)))
|
||||
table.add_row("Queries/City", str(queries_per_city))
|
||||
table.add_row("Est. Companies", f"~{estimated_companies}")
|
||||
|
||||
console.print(table)
|
||||
console.print()
|
||||
|
||||
# Show city list
|
||||
console.print("[bold]Cities included:[/bold]")
|
||||
for city in selected_cities:
|
||||
console.print(f" - {city.name} ({city.slug})")
|
||||
|
||||
console.print()
|
||||
|
||||
# Confirm creation
|
||||
if not typer.confirm("Create this plan?"):
|
||||
console.print("[yellow]Plan creation cancelled.[/yellow]")
|
||||
raise typer.Exit(0)
|
||||
|
||||
# Create the plan
|
||||
plan = ExecutionPlan(
|
||||
name=name,
|
||||
country_code=country.upper(),
|
||||
cities=[c.slug for c in selected_cities],
|
||||
status=StepStatus.PENDING,
|
||||
)
|
||||
|
||||
plan = db.create_execution_plan(plan)
|
||||
console.print(f"\n[green]Plan created: {plan.id}[/green]")
|
||||
|
||||
# Create steps for each city
|
||||
step_types = ["scrape", "enrich", "content", "export"]
|
||||
|
||||
for city in selected_cities:
|
||||
for step_type in step_types:
|
||||
step = PlanStep(
|
||||
plan_id=plan.id,
|
||||
step_name=step_type,
|
||||
city_id=city.id,
|
||||
status=StepStatus.PENDING,
|
||||
)
|
||||
db.create_plan_step(step)
|
||||
|
||||
console.print(
|
||||
f"[green]Created {len(selected_cities) * len(step_types)} steps[/green]"
|
||||
)
|
||||
|
||||
# Print next steps
|
||||
console.print("\n[bold]Next steps:[/bold]")
|
||||
console.print(f" 1. Run: python -m directory.cli execute --plan {plan.id}")
|
||||
console.print(" 2. Or run step by step:")
|
||||
console.print(f" python -m directory.cli execute --plan {plan.id} --step scrape")
|
||||
console.print(
|
||||
f" python -m directory.cli execute --plan {plan.id} --step enrich"
|
||||
)
|
||||
console.print(
|
||||
f" python -m directory.cli execute --plan {plan.id} --step content"
|
||||
)
|
||||
console.print()
|
||||
|
||||
return plan.id
|
||||
|
||||
|
||||
def list_plans() -> None:
|
||||
"""List all execution plans."""
|
||||
db = SupabaseService()
|
||||
plans = db.list_execution_plans()
|
||||
|
||||
if not plans:
|
||||
console.print("[yellow]No execution plans found.[/yellow]")
|
||||
return
|
||||
|
||||
table = Table(title="Execution Plans")
|
||||
table.add_column("ID", style="cyan", max_width=36)
|
||||
table.add_column("Name", style="green")
|
||||
table.add_column("Country", style="blue")
|
||||
table.add_column("Cities", style="magenta")
|
||||
table.add_column("Status", style="yellow")
|
||||
table.add_column("Created", style="dim")
|
||||
|
||||
for plan in plans:
|
||||
table.add_row(
|
||||
str(plan.id),
|
||||
plan.name,
|
||||
plan.country_code,
|
||||
str(len(plan.cities)),
|
||||
plan.status.value if hasattr(plan.status, "value") else plan.status,
|
||||
str(plan.created_at)[:16] if plan.created_at else "-",
|
||||
)
|
||||
|
||||
console.print(table)
|
||||
|
||||
|
||||
def show_plan(plan_id: str) -> None:
|
||||
"""Show details of a specific plan."""
|
||||
db = SupabaseService()
|
||||
|
||||
try:
|
||||
plan_uuid = UUID(plan_id)
|
||||
except ValueError:
|
||||
console.print(f"[red]Invalid plan ID: {plan_id}[/red]")
|
||||
raise typer.Exit(1)
|
||||
|
||||
plan = db.get_execution_plan(plan_uuid)
|
||||
|
||||
if not plan:
|
||||
console.print(f"[red]Plan not found: {plan_id}[/red]")
|
||||
raise typer.Exit(1)
|
||||
|
||||
# Plan info
|
||||
console.print(f"\n[bold blue]Plan: {plan.name}[/bold blue]\n")
|
||||
|
||||
table = Table()
|
||||
table.add_column("Field", style="cyan")
|
||||
table.add_column("Value", style="green")
|
||||
|
||||
table.add_row("ID", str(plan.id))
|
||||
table.add_row("Country", plan.country_code)
|
||||
table.add_row("Cities", ", ".join(plan.cities))
|
||||
table.add_row("Status", plan.status.value if hasattr(plan.status, "value") else plan.status)
|
||||
table.add_row("Created", str(plan.created_at)[:19] if plan.created_at else "-")
|
||||
table.add_row("Started", str(plan.started_at)[:19] if plan.started_at else "-")
|
||||
table.add_row("Completed", str(plan.completed_at)[:19] if plan.completed_at else "-")
|
||||
|
||||
console.print(table)
|
||||
|
||||
# Steps
|
||||
if plan.steps:
|
||||
console.print("\n[bold]Steps:[/bold]")
|
||||
|
||||
steps_table = Table()
|
||||
steps_table.add_column("Step", style="cyan")
|
||||
steps_table.add_column("City", style="blue")
|
||||
steps_table.add_column("Status", style="yellow")
|
||||
steps_table.add_column("Progress", style="green")
|
||||
|
||||
for step in plan.steps:
|
||||
city_slug = "-"
|
||||
if step.city_id:
|
||||
city = db.get_city_by_slug("", step.city_id)
|
||||
# We don't have a get_city_by_id, so we'll just show the ID
|
||||
city_slug = str(step.city_id)[:8]
|
||||
|
||||
progress = f"{step.items_processed}/{step.items_total}" if step.items_total else "-"
|
||||
|
||||
steps_table.add_row(
|
||||
step.step_name,
|
||||
city_slug,
|
||||
step.status.value if hasattr(step.status, "value") else step.status,
|
||||
progress,
|
||||
)
|
||||
|
||||
console.print(steps_table)
|
||||
Reference in New Issue
Block a user