diff --git a/src/openlayer/lib/data/commit.py b/src/openlayer/lib/data/commit.py index e94e8ff7..b46ced99 100644 --- a/src/openlayer/lib/data/commit.py +++ b/src/openlayer/lib/data/commit.py @@ -3,11 +3,13 @@ import os import tarfile import tempfile +import time from typing import Optional from ... import Openlayer from . import StorageType, _upload +from ...types.commit_retrieve_response import CommitRetrieveResponse def push( @@ -16,10 +18,28 @@ def push( project_id: str, message: str = "New commit", storage_type: Optional[StorageType] = None, -) -> None: + wait_for_completion: bool = False, + verbose: bool = False, +) -> Optional[CommitRetrieveResponse]: """Push a new commit to the Openlayer platform. - This is equivalent to running `openlayer push` from the Openlayer CLI.""" + This is equivalent to running `openlayer push` from the Openlayer CLI. + + If `wait_for_completion` is True, the function will wait for the commit to be + completed and return the commit object. + + Args: + client: The Openlayer client. + directory: The directory to push. + project_id: The id of the project to push to. + message: The commit message. + storage_type: The storage type to use. + wait_for_completion: Whether to wait for the commit to be completed. + verbose: Whether to print verbose output. + + Returns: + The commit object if `wait_for_completion` is True, otherwise None. + """ if not os.path.exists(directory): raise ValueError(f"Directory {directory} does not exist.") @@ -41,8 +61,52 @@ def push( ) # Create the project version (commit) - client.projects.commits.create( + commit = client.projects.commits.create( project_id=project_id, commit={"message": message, "source": "cli"}, storage_uri=presigned_url_response.storage_uri, ) + + if wait_for_completion: + return wait_for_commit_completion( + client=client, + project_version_id=commit.id, + verbose=verbose, + ) + + return None + + +def wait_for_commit_completion( + client: Openlayer, project_version_id: str, verbose: bool = True +) -> CommitRetrieveResponse: + """Wait for a commit to be processed by the Openlayer platform. + + Waits until the commit status is "completed" or "failed". + + Args: + client: The Openlayer client. + project_version_id: The id of the project version (commit) to wait for. + verbose: Whether to print verbose output. + + Returns: + The commit object. + """ + while True: + commit = client.commits.retrieve(project_version_id=project_version_id) + if commit.status == "completed": + if verbose: + print(f"Commit {project_version_id} completed successfully.") + return commit + elif commit.status == "failed": + raise Exception( + f"Commit {project_version_id} failed with status message:" + f" {commit.status_message}" + ) + else: + if verbose: + print( + f"Commit {project_version_id} is still processing (status:" + f" {commit.status})..." + ) + time.sleep(1)