codeapi-client 0.4.1__py3-none-any.whl

codeapi/client/_async/_jobs.py

@@ -0,0 +1,512 @@
+from __future__ import annotations
+
+import asyncio
+import time
+from pathlib import Path
+from typing import TYPE_CHECKING, AsyncGenerator, List
+
+import httpx
+
+from codeapi.types import (
+    CodeAPIError,
+    CodeAPIException,
+    CodeType,
+    Job,
+    JobResult,
+    JobStage,
+    JobStatus,
+    JobType,
+    StreamOutput,
+)
+
+if TYPE_CHECKING:
+    from . import AsyncClient
+
+
+class AsyncJobsClient:
+    def __init__(self, client: AsyncClient):
+        self._client = client
+
+    async def get(self, job_id: str) -> Job:
+        """Gets a job.
+
+        Args:
+            job_id (str): The ID of the job.
+
+        Returns:
+            Job: The job.
+
+        Raises:
+            HTTPException: If the request fails.
+        """
+        url = f"{self._client.base_url}/jobs/{job_id}"
+        async with httpx.AsyncClient() as client:
+            try:
+                response = await client.get(url, headers=self._client.api_key_header)
+                response.raise_for_status()
+                return Job(**response.json())
+            except httpx.HTTPStatusError as e:
+                raise self._client._get_http_exception(httpx_error=e)
+
+    async def list(
+        self,
+        job_type: JobType | None = None,
+        job_status: JobStatus | None = None,
+        job_stage: JobStage | None = None,
+    ) -> List[Job]:
+        """Gets a list of jobs.
+
+        Args:
+            job_type (JobType | None): Filter by job type.
+            job_status (JobStatus | None): Filter by job status.
+            job_stage (JobStage | None): Filter by job stage.
+
+        Returns:
+            list[Job]: List of jobs.
+
+        Raises:
+            HTTPException: If the request fails.
+        """
+        url = f"{self._client.base_url}/jobs"
+        params = {}
+        if job_type is not None:
+            params["job_type"] = str(job_type)
+        if job_status is not None:
+            params["job_status"] = str(job_status)
+        if job_stage is not None:
+            params["job_stage"] = str(job_stage)
+
+        async with httpx.AsyncClient() as client:
+            try:
+                response = await client.get(
+                    url, headers=self._client.api_key_header, params=params or None
+                )
+                response.raise_for_status()
+                return [Job(**item) for item in response.json()]
+            except httpx.HTTPStatusError as e:
+                raise self._client._get_http_exception(httpx_error=e)
+
+    async def get_url(self, job_id: str) -> str:
+        """Gets the URL for a job.
+
+        Args:
+            job_id (str): The job ID.
+
+        Returns:
+            str: The job URL.
+        """
+        return self._client.get_proxy_url(job_id)
+
+    async def get_status(self, job_id: str) -> JobStatus:
+        """Gets the status of a job.
+
+        Args:
+            job_id (str): The ID of the job.
+
+        Returns:
+            JobStatus: The status of the job.
+
+        Raises:
+            HTTPException: If the request fails.
+        """
+        url = f"{self._client.base_url}/jobs/{job_id}/status"
+        async with httpx.AsyncClient() as client:
+            try:
+                response = await client.get(url, headers=self._client.api_key_header)
+                response.raise_for_status()
+                return JobStatus(response.json())
+            except httpx.HTTPStatusError as e:
+                raise self._client._get_http_exception(httpx_error=e)
+
+    async def get_stage(self, job_id: str) -> JobStage:
+        """Gets the stage of a job.
+
+        Args:
+            job_id (str): The ID of the job.
+
+        Returns:
+            JobStage: The stage of the job.
+
+        Raises:
+            HTTPException: If the request fails.
+        """
+        url = f"{self._client.base_url}/jobs/{job_id}/stage"
+        async with httpx.AsyncClient() as client:
+            try:
+                response = await client.get(url, headers=self._client.api_key_header)
+                response.raise_for_status()
+                return JobStage(response.json())
+            except httpx.HTTPStatusError as e:
+                raise self._client._get_http_exception(httpx_error=e)
+
+    async def get_error(self, job_id: str) -> str | None:
+        """Gets the error message of a job.
+
+        Args:
+            job_id (str): The ID of the job.
+
+        Returns:
+            str | None: The error message of the job, or None if no error.
+
+        Raises:
+            HTTPException: If the request fails.
+        """
+        url = f"{self._client.base_url}/jobs/{job_id}/error"
+        async with httpx.AsyncClient() as client:
+            try:
+                response = await client.get(url, headers=self._client.api_key_header)
+                response.raise_for_status()
+                return response.json()
+            except httpx.HTTPStatusError as e:
+                raise self._client._get_http_exception(httpx_error=e)
+
+    async def get_code_id(self, job_id: str) -> str | None:
+        """Gets the code ID associated with a job.
+
+        Args:
+            job_id (str): The ID of the job.
+
+        Returns:
+            str | None: The code ID associated with the job, or None if no code ID.
+
+        Raises:
+            HTTPException: If the request fails.
+        """
+        url = f"{self._client.base_url}/jobs/{job_id}/code_id"
+        async with httpx.AsyncClient() as client:
+            try:
+                response = await client.get(url, headers=self._client.api_key_header)
+                response.raise_for_status()
+                return response.json()
+            except httpx.HTTPStatusError as e:
+                raise self._client._get_http_exception(httpx_error=e)
+
+    async def get_result(
+        self, job_id: str, zip_path: Path | str | None = None
+    ) -> JobResult:
+        """Gets the result of a job and optionally saves it to a file.
+
+        Args:
+            job_id (str): The ID of the job.
+            zip_path (Path | str | None): The path to save the result zip file.
+
+        Returns:
+            JobResult: The job result.
+
+        Raises:
+            HTTPException: If the request fails.
+        """
+        url = f"{self._client.base_url}/jobs/{job_id}/result"
+
+        async with httpx.AsyncClient() as client:
+            try:
+                async with client.stream(
+                    "GET", url, headers=self._client.api_key_header
+                ) as response:
+                    zip_bytes = await self._client._handle_stream_download_async(
+                        response, zip_path
+                    )
+            except httpx.HTTPStatusError as e:
+                raise self._client._get_http_exception(httpx_error=e)
+
+        return JobResult(zip_bytes=zip_bytes)
+
+    async def terminate(self, job_id: str) -> None:
+        """Terminates a job.
+
+        Args:
+            job_id (str): The ID of the job.
+
+        Raises:
+            HTTPException: If the request fails.
+        """
+        url = f"{self._client.base_url}/jobs/{job_id}/terminate"
+
+        async with httpx.AsyncClient() as client:
+            try:
+                response = await client.post(url, headers=self._client.api_key_header)
+                response.raise_for_status()
+            except httpx.HTTPStatusError as e:
+                raise self._client._get_http_exception(httpx_error=e)
+
+    async def stream_output(
+        self, job_id: str, read_timeout: float | None = 60
+    ) -> AsyncGenerator[StreamOutput, None]:
+        """Asynchronously yields stream output of a job line by line.
+
+        Args:
+            job_id (str): The ID of the job.
+            read_timeout (float | None): Read timeout for the stream request.
+
+        Yields:
+            StreamOutput: Stream output line by line.
+
+        Raises:
+            HTTPException: If the request fails.
+        """
+        url = f"{self._client.base_url}/jobs/{job_id}/stream"
+
+        timeout = httpx.Timeout(connect=5.0, read=read_timeout, write=5.0, pool=5.0)
+        async with httpx.AsyncClient() as client:
+            async with client.stream(
+                "GET", url, timeout=timeout, headers=self._client.api_key_header
+            ) as response:
+                async for line in response.aiter_lines():
+                    if line.strip():
+                        yield StreamOutput.model_validate_json(line)
+
+    async def await_stage(
+        self,
+        job_id: str,
+        timeout: float | None = None,
+        stage: JobStage = JobStage.POST_RUNNING,
+    ) -> Job:
+        """Waits for a job to reach a specific stage.
+
+        Args:
+            job_id (str): The ID of the job to wait for.
+            timeout (float | None): Maximum time to wait in seconds. If None, waits indefinitely.
+            stage (JobStage): The stage to wait for. Defaults to POST_RUNNING.
+
+        Returns:
+            Job: The final job object when the desired stage is reached.
+
+        Raises:
+            HTTPException: If the request fails.
+            APIException: If the job enters stage POST_RUNNING unexpectedly.
+            TimeoutError: If the timeout is exceeded.
+        """
+        start_time = time.time()
+
+        while True:
+            job = await self.get(job_id)
+            current_stage = job.job_stage
+            if stage == JobStage.RUNNING and current_stage == JobStage.POST_RUNNING:
+                raise CodeAPIException(
+                    error=CodeAPIError.JOB_POST_RUNNING,
+                    message=f"Job {job_id} is in stage POST_RUNNING",
+                )
+            if current_stage in [stage, JobStage.POST_RUNNING, JobStage.UNKNOWN]:
+                return job
+
+            if timeout is not None:
+                elapsed_time = time.time() - start_time
+                if elapsed_time >= timeout:
+                    raise TimeoutError(
+                        f"Job {job_id} did not reach stage {stage} "
+                        f"within {timeout} seconds"
+                    )
+
+            await asyncio.sleep(1)
+
+    async def is_healthy(self, job_id: str) -> bool:
+        """Checks whether a runner job is healthy.
+
+        Args:
+            job_id (str): The ID of the runner job.
+
+        Returns:
+            bool: True if runner job is healthy else False.
+
+        Raises:
+            HTTPException: If the request fails.
+        """
+        url = f"{self._client.base_url}/{job_id}/ping"
+        async with httpx.AsyncClient() as client:
+            try:
+                response = await client.get(url, headers=self._client.api_key_header)
+                response.raise_for_status()
+                return response.status_code == 200
+            except Exception:
+                return False
+
+    async def await_healthy(self, job_id: str, timeout: float | None = None) -> Job:
+        """Waits for a runner job to become healthy.
+
+        Args:
+            job_id (str): The ID of the runner job.
+            timeout (float | None): Maximum time to wait in seconds. If None, waits indefinitely.
+
+        Returns:
+            Job: The job object of the runner job.
+
+        Raises:
+            HTTPException: If the request fails.
+            APIException: If the job enters stage POST_RUNNING unexpectedly.
+            TimeoutError: If the timeout is exceeded.
+        """
+        start_time = time.time()
+
+        while not await self.is_healthy(job_id=job_id):
+            job = await self._client.jobs.get(job_id)
+            if job.job_stage == JobStage.POST_RUNNING:
+                raise CodeAPIException(
+                    error=CodeAPIError.JOB_POST_RUNNING,
+                    message=f"Job {job_id} unexpectedly entered stage POST_RUNNING",
+                )
+            if timeout is not None:
+                elapsed_time = time.time() - start_time
+                if elapsed_time >= timeout:
+                    raise TimeoutError(
+                        f"Job {job_id} did not become healthy within {timeout} seconds"
+                    )
+            await asyncio.sleep(1)
+
+        return await self._client.jobs.get(job_id)
+
+    async def get_latest(self, code_type: CodeType | None = None) -> Job | None:
+        """Get the most recent job.
+
+        Args:
+            code_type (CodeType | None): Filter by code type.
+
+        Returns:
+            Job | None: The most recent job, or None if no jobs exist.
+        """
+        job_type = JobType.from_code_type(code_type) if code_type else None
+        jobs = await self.list(job_type=job_type)
+        return jobs[0] if jobs else None
+
+    async def list_queued(self, code_type: CodeType | None = None) -> List[Job]:
+        """Get all queued jobs.
+
+        Args:
+            code_type (CodeType | None): Filter by code type.
+
+        Returns:
+            List[Job]: List of queued jobs.
+        """
+        job_type = JobType.from_code_type(code_type) if code_type else None
+        return await self.list(job_type=job_type, job_status=JobStatus.QUEUED)
+
+    async def list_scheduled(self, code_type: CodeType | None = None) -> List[Job]:
+        """Get all scheduled jobs.
+
+        Args:
+            code_type (CodeType | None): Filter by code type.
+
+        Returns:
+            List[Job]: List of scheduled jobs.
+        """
+        job_type = JobType.from_code_type(code_type) if code_type else None
+        return await self.list(job_type=job_type, job_status=JobStatus.SCHEDULED)
+
+    async def list_started(self, code_type: CodeType | None = None) -> List[Job]:
+        """Get all started jobs.
+
+        Args:
+            code_type (CodeType | None): Filter by code type.
+
+        Returns:
+            List[Job]: List of started jobs.
+        """
+        job_type = JobType.from_code_type(code_type) if code_type else None
+        return await self.list(job_type=job_type, job_status=JobStatus.STARTED)
+
+    async def list_deferred(self, code_type: CodeType | None = None) -> List[Job]:
+        """Get all deferred jobs.
+
+        Args:
+            code_type (CodeType | None): Filter by code type.
+
+        Returns:
+            List[Job]: List of deferred jobs.
+        """
+        job_type = JobType.from_code_type(code_type) if code_type else None
+        return await self.list(job_type=job_type, job_status=JobStatus.DEFERRED)
+
+    async def list_canceled(self, code_type: CodeType | None = None) -> List[Job]:
+        """Get all canceled jobs.
+
+        Args:
+            code_type (CodeType | None): Filter by code type.
+
+        Returns:
+            List[Job]: List of canceled jobs.
+        """
+        job_type = JobType.from_code_type(code_type) if code_type else None
+        return await self.list(job_type=job_type, job_status=JobStatus.CANCELED)
+
+    async def list_stopped(self, code_type: CodeType | None = None) -> List[Job]:
+        """Get all stopped jobs.
+
+        Args:
+            code_type (CodeType | None): Filter by code type.
+
+        Returns:
+            List[Job]: List of stopped jobs.
+        """
+        job_type = JobType.from_code_type(code_type) if code_type else None
+        return await self.list(job_type=job_type, job_status=JobStatus.STOPPED)
+
+    async def list_failed(self, code_type: CodeType | None = None) -> List[Job]:
+        """Get all failed jobs.
+
+        Args:
+            code_type (CodeType | None): Filter by code type.
+
+        Returns:
+            List[Job]: List of failed jobs.
+        """
+        job_type = JobType.from_code_type(code_type) if code_type else None
+        return await self.list(job_type=job_type, job_status=JobStatus.FAILED)
+
+    async def list_finished(self, code_type: CodeType | None = None) -> List[Job]:
+        """Get all finished jobs.
+
+        Args:
+            code_type (CodeType | None): Filter by code type.
+
+        Returns:
+            List[Job]: List of finished jobs.
+        """
+        job_type = JobType.from_code_type(code_type) if code_type else None
+        return await self.list(job_type=job_type, job_status=JobStatus.FINISHED)
+
+    async def list_timed_out(self, code_type: CodeType | None = None) -> List[Job]:
+        """Get all timed out jobs.
+
+        Args:
+            code_type (CodeType | None): Filter by code type.
+
+        Returns:
+            List[Job]: List of timed out jobs.
+        """
+        job_type = JobType.from_code_type(code_type) if code_type else None
+        return await self.list(job_type=job_type, job_status=JobStatus.TIMEOUT)
+
+    async def list_pre_running(self, code_type: CodeType | None = None) -> List[Job]:
+        """Get all pre-running jobs.
+
+        Args:
+            code_type (CodeType | None): Filter by code type.
+
+        Returns:
+            List[Job]: List of pre-running jobs.
+        """
+        job_type = JobType.from_code_type(code_type) if code_type else None
+        return await self.list(job_type=job_type, job_stage=JobStage.PRE_RUNNING)
+
+    async def list_running(self, code_type: CodeType | None = None) -> List[Job]:
+        """Get all running jobs.
+
+        Args:
+            code_type (CodeType | None): Filter by code type.
+
+        Returns:
+            List[Job]: List of running jobs.
+        """
+        job_type = JobType.from_code_type(code_type) if code_type else None
+        return await self.list(job_type=job_type, job_stage=JobStage.RUNNING)
+
+    async def list_post_running(self, code_type: CodeType | None = None) -> List[Job]:
+        """Get all post-running jobs.
+
+        Args:
+            code_type (CodeType | None): Filter by code type.
+
+        Returns:
+            List[Job]: List of post-running jobs.
+        """
+        job_type = JobType.from_code_type(code_type) if code_type else None
+        return await self.list(job_type=job_type, job_stage=JobStage.POST_RUNNING)