cleaning and fixing bug in the interaction with the cli

Author: miciav

src/maestro/cli_client.py

@@ -59,18 +59,31 @@ def create(
 
 @app.command()
 def run(
-    dag_id: str = typer.Argument(..., help="DAG ID to run"),
+    dag_input: str = typer.Argument(..., help="DAG ID or path to DAG YAML file"),
     resume: bool = typer.Option(False, "--resume", help="Resume from last checkpoint"),
     fail_fast: bool = typer.Option(True, "--fail-fast/--no-fail-fast", help="Stop on first failure"),
     server_url: str = typer.Option("http://localhost:8000", "--server", help="Maestro server URL")
 ):
-    """Run a previously created DAG"""
+    """Run a DAG - either by ID or by creating from a YAML file"""
     api_client.base_url = server_url
     check_server_connection()
 
     try:
+        # Check if the input is a file path or a DAG ID
+        if os.path.exists(dag_input) and dag_input.endswith(('.yaml', '.yml')):
+            # It's a file path - create the DAG first
+            dag_file_path = os.path.abspath(dag_input)
+            console.print(f"[blue]Creating DAG from file: {dag_file_path}[/blue]")
+            create_response = api_client.create_dag(dag_file_path)
+            dag_id = create_response['dag_id']
+            console.print(f"[bold green]✓ DAG created successfully with ID: {dag_id}[/bold green]")
+        else:
+            # It's a DAG ID
+            dag_id = dag_input
+        
+        # Now run the DAG
         response = api_client.run_dag(dag_id, resume, fail_fast)
-        console.print(f"[bold green]✓ {response['message']}[/bold green]")
+        console.print(f"[bold green]✓ DAG started successfully![/bold green]")
         console.print(f"[cyan]DAG ID:[/cyan] {response['dag_id']}")
         console.print(f"[cyan]Execution ID:[/cyan] {response['execution_id']}")
         console.print(f"[cyan]Status:[/cyan] [yellow]{response['status']}[/yellow]")
@@ -85,49 +98,6 @@ def run(
         console.print(f"[red]Error: {e}[/red]")
         raise typer.Exit(1)
 
-@app.command()
-def submit(
-    dag_file: str = typer.Argument(..., help="Path to the DAG YAML file"),
-    resume: bool = typer.Option(False, "--resume", help="Resume from last checkpoint"),
-    fail_fast: bool = typer.Option(True, "--fail-fast/--no-fail-fast", help="Stop on first failure"),
-    server_url: str = typer.Option("http://localhost:8000", "--server", help="Maestro server URL")
-):
-    """Submit a DAG for execution (legacy command)"""
-    console.print("[yellow]Warning: 'maestro submit' is a legacy command. Use 'maestro create' followed by 'maestro run' for better control.[/yellow]")
-    api_client.base_url = server_url
-    check_server_connection()
-    
-    try:
-        # Convert a relative path to absolute
-        dag_file_path = os.path.abspath(dag_file)
-        
-        if not os.path.exists(dag_file_path):
-            console.print(f"[red]Error: DAG file not found: {dag_file_path}[/red]")
-            raise typer.Exit(1)
-        
-        # First, create the DAG
-        create_response = api_client.create_dag(dag_file_path)
-        dag_id = create_response['dag_id']
-        console.print(f"[bold green]✓ DAG created successfully with ID: {dag_id}[/bold green]")
-
-        # Then, run the DAG
-        run_response = api_client.run_dag(dag_id, resume, fail_fast)
-        
-        console.print(f"[bold green]✓ DAG submitted successfully![/bold green]")
-        console.print(f"[cyan]DAG ID:[/cyan] {run_response['dag_id']}")
-        console.print(f"[cyan]Execution ID:[/cyan] {run_response['execution_id']}")
-        console.print(f"[cyan]Status:[/cyan] [yellow]{run_response['status']}[/yellow]")
-        console.print()
-        console.print("[bold blue]Available commands:[/bold blue]")
-        console.print(f"  • [bold]maestro status {run_response['dag_id']}[/bold] - Check DAG status")
-        console.print(f"  • [bold]maestro log {run_response['dag_id']}[/bold] - View logs")
-        console.print(f"  • [bold]maestro attach {run_response['dag_id']}[/bold] - Attach to live log stream")
-        console.print(f"  • [bold]maestro stop {run_response['dag_id']}[/bold] - Stop execution")
-        console.print()
-        
-    except Exception as e:
-        console.print(f"[red]Error: {e}[/red]")
-        raise typer.Exit(1)
 
 @app.command()
 def status(
@@ -150,7 +120,7 @@ def status(
         table.add_row("Execution ID", response["execution_id"])
         table.add_row("Status", response["status"])
         table.add_row("Started", response["started_at"] or "N/A")
-        table.add_row("Completed", response["completed_at"] or "Running...")
+        table.add_row("Completed", response["completed_at"] or "N/A")
         table.add_row("Thread ID", response["thread_id"] or "N/A")
 
         console.print(table)
@@ -281,17 +251,40 @@ def handle_exit(signum, frame):
 
 @app.command()
 def rm(
-    dag_id: str = typer.Argument(..., help="DAG ID to remove"),
-    force: bool = typer.Option(False, "--force", "-f", help="Force removal of a running DAG"),
+    dag_id: Optional[str] = typer.Argument(None, help="DAG ID to remove"),
+    all: bool = typer.Option(False, "--all", "-a", help="Remove all non-running DAGs (or all DAGs with --force)"),
+    force: bool = typer.Option(False, "--force", "-f", help="Force removal, even for running DAGs"),
     server_url: str = typer.Option("http://localhost:8000", "--server", help="Maestro server URL")
 ):
     """Remove a DAG and its executions"""
     api_client.base_url = server_url
     check_server_connection()
 
     try:
-        response = api_client.remove_dag(dag_id, force)
-        console.print(f"[bold green]✓ {response['message']}[/bold green]")
+        if all:
+            # Get all DAGs
+            dags = api_client.list_dags_v1()
+            removed_count = 0
+            
+            for dag in dags:
+                # Remove all non-running DAGs, or all if --force is used
+                if force or dag['status'] not in ['running']:
+                    response = api_client.remove_dag(dag['dag_id'], force)
+                    console.print(f"[bold green]✓ Removed DAG: {dag['dag_id']} (status: {dag['status']})[/bold green]")
+                    removed_count += 1
+                else:
+                    console.print(f"[yellow]⚠ Skipped running DAG: {dag['dag_id']}[/yellow]")
+            
+            if removed_count == 0:
+                console.print("[yellow]No DAGs to remove.[/yellow]")
+            else:
+                console.print(f"[bold green]✓ Removed {removed_count} DAG(s).[/bold green]")
+        else:
+            if not dag_id:
+                console.print("[red]Error: Please provide a DAG ID or use --all flag[/red]")
+                raise typer.Exit(1)
+            response = api_client.remove_dag(dag_id, force)
+            console.print(f"[bold green]✓ {response['message']}[/bold green]")
     except Exception as e:
         console.print(f"[red]Error: {e}[/red]")
         raise typer.Exit(1)
@@ -388,7 +381,7 @@ def list_dags(
                 dag["execution_id"][:8] + "..." if dag["execution_id"] else "N/A",
                 f"[{status_style}]{dag.get('status', 'N/A')}[/{status_style}]",
                 dag.get("started_at", "N/A"),
-                dag.get("completed_at", "N/A") or "Running...",
+                dag.get("completed_at", "N/A"),
                 str(dag.get("thread_id", "N/A"))
             )
 

src/maestro/core/dag.py

@@ -3,15 +3,27 @@
 from collections import deque
 from datetime import datetime
 from croniter import croniter
+from enum import Enum
 
 from maestro.core.task import Task
 
+
+class DAGStatus(str, Enum):
+    """Enumeration for the status of a DAG."""
+    PENDING = "pending"
+    RUNNING = "running"
+    COMPLETED = "completed"
+    FAILED = "failed"
+
+
 class DAG:
     def __init__(self, dag_id: str = "default_dag", start_time: Optional[datetime] = None, cron_schedule: Optional[str] = None):
         self.dag_id = dag_id
         self.start_time = start_time
         self.cron_schedule = cron_schedule
         self.tasks: Dict[str, Task] = {}
+        self.status: DAGStatus = DAGStatus.PENDING
+        self.execution_id: Optional[str] = None
 
         # Validate that only one scheduling method is used
         if start_time is not None and cron_schedule is not None:

src/maestro/core/orchestrator.py

@@ -7,7 +7,7 @@
 from rich import get_console
 from rich.logging import RichHandler
 
-from maestro.core.dag import DAG
+from maestro.core.dag import DAG, DAGStatus
 from maestro.core.task import TaskStatus
 from maestro.tasks.base import BaseTask
 from maestro.core.Task_Registry import TaskRegistry
@@ -101,15 +101,29 @@ def load_dag_from_file(self, filepath: str, dag_id: Optional[str] = None) -> DAG
     def run_dag_in_thread(
         self,
         dag: DAG,
+        execution_id: str = None,
         resume: bool = False,
         fail_fast: bool = True,
         status_callback=None
     ) -> str:
         """Execute DAG in a separate thread with concurrency."""
-        execution_id = str(uuid.uuid4())
-        # Create execution record in database synchronously
+        # Use provided execution_id or generate a new one
+        if execution_id is None:
+            execution_id = str(uuid.uuid4())
+        
+        dag.execution_id = execution_id
+        
+        # Create or update execution record in database synchronously
         with self.status_manager as sm:
-            sm.create_dag_execution(dag.dag_id, execution_id)
+            # Check if execution already exists (e.g., with 'created' status)
+            existing = sm.get_latest_execution(dag.dag_id)
+            if existing and existing["execution_id"] == execution_id:
+                # Update existing execution to 'running'
+                sm.update_dag_execution_status(dag.dag_id, execution_id, "running")
+            else:
+                # Create new execution
+                sm.create_dag_execution(dag.dag_id, execution_id)
+            
             # Initialize all tasks with pending status only if not resuming
             if not resume:
                 task_ids = list(dag.tasks.keys())
@@ -136,14 +150,18 @@ def execute():
                     # Determine overall DAG status
                     if has_running:
                         # Should not happen after execution completes, but handle it
+                        dag.status = DAGStatus.RUNNING
                         sm.update_dag_execution_status(dag.dag_id, execution_id, "running")
                     elif has_failed:
+                        dag.status = DAGStatus.FAILED
                         sm.update_dag_execution_status(dag.dag_id, execution_id, "failed")
                     else:
+                        dag.status = DAGStatus.COMPLETED
                         sm.update_dag_execution_status(dag.dag_id, execution_id, "completed")
 
             except Exception as e:
                 with self.status_manager as sm:
+                    dag.status = DAGStatus.FAILED
                     # Mark any running or pending tasks as failed when DAG execution fails
                     sm.mark_incomplete_tasks_as_failed(dag.dag_id, execution_id)
 
@@ -169,6 +187,7 @@ def run_dag(
     ):
         """Execute DAG with improved error handling and persistence."""
         dag_id = dag.dag_id
+        dag.status = DAGStatus.RUNNING
         execution_order = dag.get_execution_order()
 
         # Set up database logging if execution_id is provided
@@ -317,4 +336,4 @@ def get_dag_status(self, dag: DAG) -> Dict[str, Any]:
             "tasks": tasks_status,
             "summary": summary,
             "total_tasks": len(dag.tasks)
-        }
+        }

src/maestro/core/status_manager.py

@@ -259,6 +259,33 @@ def create_dag_execution(self, dag_id: str, execution_id: str) -> str:
 
             return execution_id
 
+    def create_dag_execution_with_status(self, dag_id: str, execution_id: str, status: str) -> str:
+        """Create a new DAG execution record with a specific status."""
+        with self._lock:
+            self._create_dag_if_not_exists(dag_id)
+            current_time = datetime.now().isoformat()
+            thread_id = threading.current_thread().ident
+            pid = threading.current_thread().ident
+            
+            with self._conn:
+                if status == "created":
+                    # For created status, don't set started_at or thread_id
+                    self._conn.execute("""
+                        INSERT OR REPLACE INTO executions 
+                        (id, dag_id, status, started_at, thread_id, pid)
+                        VALUES (?, ?, ?, NULL, NULL, ?)
+                    """, (execution_id, dag_id, status, pid))
+                else:
+                    # For other statuses, set started_at
+                    self._conn.execute("""
+                        INSERT OR REPLACE INTO executions 
+                        (id, dag_id, status, started_at, thread_id, pid)
+                        VALUES (?, ?, ?, ?, ?, ?)
+                    """, (execution_id, dag_id, status, current_time, str(thread_id), pid))
+                self._conn.commit()
+            
+            return execution_id
+    
     def initialize_tasks_for_execution(self, dag_id: str, execution_id: str, task_ids: List[str]):
         """Initialize all tasks for a DAG execution with pending status."""
         with self._lock:
@@ -278,6 +305,7 @@ def update_dag_execution_status(self, dag_id: str,
         """Update DAG execution status."""
         with self._lock:
             current_time = datetime.now().isoformat()
+            thread_id = threading.current_thread().ident
 
             with self._conn:
                 if status in ["completed", "failed", "cancelled"]:
@@ -286,6 +314,15 @@ def update_dag_execution_status(self, dag_id: str,
                         SET status = ?, completed_at = ?
                         WHERE dag_id = ? AND id = ?
                     """, (status, current_time, dag_id, execution_id))
+                elif status == "running":
+                    # When transitioning to running, set thread_id and started_at if not already set
+                    self._conn.execute("""
+                        UPDATE executions 
+                        SET status = ?, 
+                            thread_id = CASE WHEN thread_id IS NULL THEN ? ELSE thread_id END,
+                            started_at = CASE WHEN started_at IS NULL THEN ? ELSE started_at END
+                        WHERE dag_id = ? AND id = ?
+                    """, (status, str(thread_id), current_time, dag_id, execution_id))
                 else:
                     self._conn.execute("""
                         UPDATE executions 
@@ -586,6 +623,29 @@ def generate_unique_dag_id(self) -> str:
         suffix = ''.join(random.choices(string.ascii_lowercase + string.digits, k=6))
         return f"{base_name}_{suffix}"
 
+    def get_latest_execution(self, dag_id: str) -> Optional[Dict[str, Any]]:
+        """Get the latest execution for a DAG."""
+        with self._lock:
+            with self._conn:
+                cursor = self._conn.execute("""
+                    SELECT id, status, started_at, completed_at, thread_id, pid
+                    FROM executions 
+                    WHERE dag_id = ?
+                    ORDER BY COALESCE(started_at, datetime('now')) DESC
+                    LIMIT 1
+                """, (dag_id,))
+                row = cursor.fetchone()
+                if row:
+                    return {
+                        "execution_id": row[0],
+                        "status": row[1],
+                        "started_at": row[2],
+                        "completed_at": row[3],
+                        "thread_id": row[4],
+                        "pid": row[5]
+                    }
+                return None
+
     def save_dag_definition(self, dag):
         """Save the DAG definition to the database."""
         with self._lock:

src/maestro/server/app.py

@@ -79,7 +79,7 @@ def print_logo():
             return
 
         img = Image.open(logo_path).convert("RGBA")
-        img = img.resize((40, 20), Image.Resampling.LANCZOS)  # Resize for a reasonable terminal display
+        img = img.resize((37, 15), Image.Resampling.LANCZOS)  # Resize for a reasonable terminal display
 
         for y in range(img.height):
             for x in range(img.width):

src/maestro/server/docker_api.py

@@ -5,6 +5,7 @@
 from typing import Optional, List, Dict, Any
 import asyncio
 import json
+import uuid
 
 from maestro.core.orchestrator import Orchestrator
 from maestro.core.status_manager import StatusManager
@@ -83,8 +84,13 @@ async def create_dag(request: DAGCreateRequest):
         # Load and validate DAG
         dag = orchestrator.load_dag_from_file(request.dag_file_path, dag_id=dag_id)
 
+        # Create a new execution ID
+        execution_id = str(uuid.uuid4())
+        
         with orchestrator.status_manager as sm:
             sm.save_dag_definition(dag)
+            # Create initial execution with 'created' status
+            sm.create_dag_execution_with_status(dag_id, execution_id, "created")
 
         return DAGCreateResponse(
             dag_id=dag.dag_id,
@@ -103,11 +109,23 @@ async def run_dag(dag_id: str, request: DAGRunRequest, background_tasks: Backgro
             dag_definition = sm.get_dag_definition(dag_id)
             if not dag_definition:
                 raise HTTPException(status_code=404, detail=f"DAG '{dag_id}' not found.")
+            
+            # Get the latest execution for this DAG
+            latest_execution = sm.get_latest_execution(dag_id)
+            
+            if latest_execution and latest_execution["status"] == "created":
+                # Use the existing execution ID from the created DAG
+                execution_id = latest_execution["execution_id"]
+            else:
+                # Create a new execution ID for subsequent runs
+                execution_id = str(uuid.uuid4())
 
         dag = orchestrator.dag_loader.load_dag_from_dict(dag_definition)
-
-        execution_id = orchestrator.run_dag_in_thread(
+        
+        # Run the DAG with the determined execution ID
+        orchestrator.run_dag_in_thread(
             dag=dag,
+            execution_id=execution_id,
             resume=request.resume,
             fail_fast=request.fail_fast
         )