243 lines
7.6 KiB
Python
243 lines
7.6 KiB
Python
"""
|
|
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)
|