feat: add log streaming and environment management utilities; implement asynchronous log streaming and context retrieval functions

Author: akosyakov

examples/execute_command.py

@@ -1,87 +1,60 @@
 #!/usr/bin/env python
 
 from gitpod import Gitpod
-from utils import get_context_url, get_authenticated_user, get_environment_class_id, get_environment_class, create_environment, check_environment_status, stream_logs
-import asyncio
+from utils.context import get_context_url, get_authenticated_user
+from utils.environment import get_environment_class_id, get_environment_class, create_environment
+from utils.task import create_command, run_command, delete_command, stream_command_logs
 
 def main() -> None:
+    """
+    Main function to create an environment, start a task, and stream its logs.
+    """
+    # Initialize Gitpod client
     client: Gitpod = Gitpod()
+    
+    # Get context URL and authenticated user
     context_url: str = get_context_url()
     user = get_authenticated_user(client)
+    
+    # Get environment class ID and details
     environment_class_id: str = get_environment_class_id(client, user)
-
     if not environment_class_id:
-        raise ValueError("no environment class found")
-
+        raise ValueError("no environment class found, please set the ENVIRONMENT_CLASS_ID environment variable")
     environment_class = get_environment_class(client, environment_class_id)
-
     if not environment_class:
-        raise ValueError("no environment class found")
+        raise ValueError("no environment class found, please set the ENVIRONMENT_CLASS_ID environment variable")
 
     print(f"Using repository: {context_url}")
-    print(f"Using environment class: {environment_class.display_name} (ID: {environment_class.id})")
+    print(f"Using environment class: {environment_class.display_name} ({environment_class.description})")
 
+    # Create environment
     environment_id: str = create_environment(client, context_url, environment_class)
-    print(f"Created environment: {environment_id}")
+    print(f"\nCreated environment: {environment_id}")
 
     try:
+        # Define command
         reference = "hello-world"
-        task = client.environments.automations.tasks.create(
-            spec={
-                "command": "echo 'Hello, World!' && false",
-            },
-            environment_id=environment_id,
-            metadata={
-                "name": "Hello, World!",
-                "description": "Prints 'Hello, World!' to the console.",
-                "reference": reference,
-            }
-        ).task
-        print(f"Created task: {task.id}")
-
-        execution = client.environments.automations.tasks.start(
-            id=task.id
-        ).task_execution
-        print(f"Started task execution: {execution.id}")
-
-        logs_access_token = client.environments.create_logs_token(environment_id=environment_id).access_token
-        print(f"Waiting for task logs to be available...")
-        
-        event_stream = client.events.watch(
-            environment_id=environment_id,
-            timeout=None
-        )
-
-        log_url = None
-        for event in event_stream:
-            if event.resource_type == "RESOURCE_TYPE_ENVIRONMENT" and event.resource_id == environment_id:
-                check_environment_status(client, environment_id)
-            elif event.resource_type == "RESOURCE_TYPE_TASK_EXECUTION" and event.resource_id == execution.id:
-                execution = client.environments.automations.tasks.executions.retrieve(id=execution.id).task_execution
-                print(f"Task execution status: {execution.status.phase}")
-                if execution.status.log_url:
-                    log_url = execution.status.log_url
-                    break
-
-        event_stream.http_response.close()
-    
-        if not log_url:
-            raise RuntimeError("Task logs are not available.")
-        
-        print(f"Task logs are available at: {log_url}")
-        print("")
-        print("Streaming logs:")
-
-        asyncio.run(stream_logs(log_url, logs_access_token))
-
-        execution = client.environments.automations.tasks.executions.retrieve(id=execution.id).task_execution
-        print(f"Task execution status: {execution.status.phase}")
-        if execution.status.phase == "TASK_EXECUTION_PHASE_FAILED":
-            print(f"Task execution failed: {execution.status.failure_message}")
-
+        name = "Hello, World!"
+        description = "Prints 'Hello, World!' to the console."
+        task_id = create_command(client, environment_id, reference, name, description)    
+        try:
+            # Run command
+            task_execution = run_command(client, task_id)
+            
+            # Stream logs
+            stream_command_logs(client, environment_id, task_execution)
+
+            # Check command status
+            execution = client.environments.automations.tasks.executions.retrieve(id=execution_id).task_execution
+            print(f"\nTask execution status: {execution.status.phase}")
+            if execution.status.phase == "TASK_EXECUTION_PHASE_FAILED":
+                print(f"Task execution failed: {execution.status.failure_message}")
+        finally: 
+            delete_command(client, task_id)
     finally:
+        # Clean up environment
         client.environments.delete(environment_id=environment_id)
-        print(f"Deleted environment: {environment_id}")
+        print(f"\nDeleted environment: {environment_id}")
 
 if __name__ == "__main__":
     main()

examples/utils/context.py

@@ -0,0 +1,25 @@
+import os
+from typing import Optional
+from gitpod import Gitpod
+from gitpod.types.user_get_authenticated_user_response import User
+
+def get_context_url() -> str:
+    """
+    Retrieve the context URL from the environment variable or use a default value.
+
+    Returns:
+        str: The context URL.
+    """
+    return os.environ.get("CONTEXT_URL", "https://github.com/gitpod-io/empty")
+
+def get_authenticated_user(client: Gitpod) -> User:
+    """
+    Get the authenticated user using the Gitpod client.
+
+    Args:
+        client (Gitpod): The Gitpod client instance.
+
+    Returns:
+        User: The authenticated user.
+    """
+    return client.users.get_authenticated_user(body={}).user

examples/utils.py renamed to examples/utils/environment.py

@@ -1,18 +1,21 @@
 import os
 from typing import Optional
-from gitpod.types.environment_create_params import SpecContentInitializerSpecContextURL
 from gitpod import Gitpod
 from gitpod.types.user_get_authenticated_user_response import User
 from gitpod.types.environments.class_list_response import ClassListResponse
-import httpx
+from gitpod.types.environment_create_params import SpecContentInitializerSpecContextURL
 
-def get_context_url() -> str:
-    return os.environ.get("CONTEXT_URL", "https://github.com/gitpod-io/empty")
+def get_environment_class_id(client: Gitpod, user: User) -> Optional[str]:
+    """
+    Get the most used environment class ID for the given user.
 
-def get_authenticated_user(client: Gitpod) -> User:
-    return client.users.get_authenticated_user(body={}).user
+    Args:
+        client (Gitpod): The Gitpod client instance.
+        user (User): The authenticated user.
 
-def get_environment_class_id(client: Gitpod, user: User) -> Optional[str]:
+    Returns:
+        Optional[str]: The environment class ID or None if not found.
+    """
     environment_class_id = os.environ.get("ENVIRONMENT_CLASS_ID")
     if environment_class_id:
         return environment_class_id
@@ -34,6 +37,16 @@ def get_environment_class_id(client: Gitpod, user: User) -> Optional[str]:
     return sorted_classes[0][0] if sorted_classes else None
 
 def get_environment_class(client: Gitpod, environment_class_id: str) -> Optional[ClassListResponse]:
+    """
+    Get the environment class details for the given environment class ID.
+
+    Args:
+        client (Gitpod): The Gitpod client instance.
+        environment_class_id (str): The environment class ID.
+
+    Returns:
+        Optional[ClassListResponse]: The environment class details or None if not found.
+    """
     page = client.environments.classes.list(filter={"enabled": True})
     while page:
         for cls in page.environment_classes:
@@ -46,6 +59,17 @@ def get_environment_class(client: Gitpod, environment_class_id: str) -> Optional
     return None
 
 def create_environment(client: Gitpod, context_url: str, environment_class: ClassListResponse) -> str:
+    """
+    Create a new environment with the given context URL and environment class.
+
+    Args:
+        client (Gitpod): The Gitpod client instance.
+        context_url (str): The context URL.
+        environment_class (ClassListResponse): The environment class details.
+
+    Returns:
+        str: The created environment ID.
+    """
     resp = client.environments.create(
         spec={
             "desired_phase": "ENVIRONMENT_PHASE_RUNNING",
@@ -67,18 +91,20 @@ def create_environment(client: Gitpod, context_url: str, environment_class: Clas
     )
     return resp.environment.id
 
-def check_environment_status(client: Gitpod, environment_id: str) -> None:
+def ensure_environment_healthy(client: Gitpod, environment_id: str) -> None:
+    """
+    Ensure the environment is running or will be running and raise an error if it is in an unexpected phase.
+
+    Args:
+        client (Gitpod): The Gitpod client instance.
+        environment_id (str): The environment ID.
+
+    Raises:
+        RuntimeError: If the environment is in an unexpected phase or has a failure message.
+    """
     environment = client.environments.retrieve(environment_id=environment_id).environment
     print(f"Environment status: {environment.status.phase}")
     if environment.status.phase in ["ENVIRONMENT_PHASE_STOPPING", "ENVIRONMENT_PHASE_STOPPED", "ENVIRONMENT_PHASE_DELETING", "ENVIRONMENT_PHASE_DELETED"]:
         raise RuntimeError(f"Environment {environment_id} is in an unexpected phase: {environment.status.phase}")
     elif environment.status.failure_message and len(environment.status.failure_message) > 0:
         raise RuntimeError(f"Environment {environment_id} failed: {'; '.join(environment.status.failure_message)}")
-
-async def stream_logs(log_url: str, logs_access_token: str) -> None:
-    async with httpx.AsyncClient() as client:
-        async with client.stream("GET", log_url, headers={"Authorization": f"Bearer {logs_access_token}"}) as response:
-            async for line in response.aiter_lines():
-                if line:
-                    print(line)
-

examples/utils/logs.py

@@ -0,0 +1,15 @@
+import httpx
+
+async def stream_logs(log_url: str, logs_access_token: str) -> None:
+    """
+    Stream logs from the given log URL using the provided access token.
+
+    Args:
+        log_url (str): The log URL.
+        logs_access_token (str): The access token for the logs.
+    """
+    async with httpx.AsyncClient() as client:
+        async with client.stream("GET", log_url, headers={"Authorization": f"Bearer {logs_access_token}"}) as response:
+            async for line in response.aiter_lines():
+                if line:
+                    print(line)

examples/utils/task.py

@@ -0,0 +1,129 @@
+from gitpod import Gitpod
+import asyncio
+from utils.logs import stream_logs
+from utils.environment import ensure_environment_healthy
+
+
+def create_command(client: Gitpod, environment_id: str, reference: str, name: str, description: str) -> str:
+    """
+    Create a task in the given environment.
+
+    Args:
+        client (Gitpod): The Gitpod client instance.
+        environment_id (str): The environment ID.
+        reference (str): The task reference.
+        name (str): The task name.
+        description (str): The task description.
+
+    Returns:
+        str: The task ID.
+    """
+    task = client.environments.automations.tasks.create(
+        spec={
+            "command": "echo 'Hello, World!'",
+        },
+        environment_id=environment_id,
+        metadata={
+            "name": name,
+            "description": description,
+            "reference": reference,
+        }
+    ).task
+    print(f"\nCreated task: {task.id}")
+    return task.id
+
+
+def run_command(client: Gitpod, task_id: str) -> str:
+    """
+    Start a task execution.
+
+    Args:
+        client (Gitpod): The Gitpod client instance.
+        task_id (str): The task ID.
+
+    Returns:
+        str: The task execution ID.
+    """
+    execution = client.environments.automations.tasks.start(
+        id=task_id
+    ).task_execution
+    print(f"Started task execution: {execution.id}")
+    return execution.id
+
+
+def delete_command(client: Gitpod, task_id: str) -> None:
+    """
+    Delete a task.
+
+    Args:
+        client (Gitpod): The Gitpod client instance.
+        task_id (str): The task ID.
+    """
+    client.environments.automations.tasks.delete(id=task_id)
+    print(f"Deleted task: {task_id}")
+
+
+def wait_for_task_logs(client: Gitpod, environment_id: str, execution_id: str) -> str:
+    """
+    Wait for the task logs to be available.
+
+    Args:
+        client (Gitpod): The Gitpod client instance.
+        environment_id (str): The environment ID.
+        execution_id (str): The task execution ID.
+
+    Returns:
+        str: The log URL.
+
+    Raises:
+        RuntimeError: If the task logs are not available.
+    """
+    def get_log_url() -> str:
+        execution = client.environments.automations.tasks.executions.retrieve(
+            id=execution_id).task_execution
+        return execution.status.log_url
+
+    print(f"Waiting for task logs to be available...")
+
+    # 1. open the stream before the first request
+    event_stream = client.events.watch(
+        environment_id=environment_id,
+        timeout=None
+    )
+
+    try:
+        # 2. make the first request to ensure that we did not miss any events
+        ensure_environment_healthy(client, environment_id)
+        log_url = get_log_url()
+        if log_url:
+            return log_url
+        
+        for event in event_stream:
+            if event.resource_type == "RESOURCE_TYPE_ENVIRONMENT" and event.resource_id == environment_id:
+                ensure_environment_healthy(client, environment_id)
+            elif event.resource_type == "RESOURCE_TYPE_TASK_EXECUTION" and event.resource_id == execution_id:
+                log_url = get_log_url()
+                print(f"Task execution status: {client.environments.automations.tasks.executions.retrieve(id=execution_id).task_execution.status.phase}")
+                if log_url:
+                    return log_url
+    finally:
+        event_stream.http_response.close()
+
+    raise RuntimeError("Task logs are not available.")
+
+
+def stream_command_logs(client: Gitpod, environment_id: str, execution_id: str) -> None:
+    """
+    Stream the logs of a command execution.
+
+    Args:
+        client (Gitpod): The Gitpod client instance.
+        environment_id (str): The environment ID.
+        execution_id (str): The task execution ID.
+    """
+    log_url = wait_for_task_logs(client, environment_id, execution_id)
+    print(f"\nTask logs are available at: {log_url}")
+    print("\nStreaming logs:")
+
+    logsAccessToken = client.environments.create_logs_token(environment_id=environment_id).access_token
+    asyncio.run(stream_logs(log_url, logsAccessToken))