Prod-1820 forma add sync library function to fetch cluster spec of latest submission in project (#110)

Author: Brandon Kaplan

sync/__init__.py

@@ -1,5 +1,5 @@
 """Library for leveraging the power of Sync"""
 
-__version__ = "1.4.1"
+__version__ = "1.5.0"
 
 TIME_FORMAT = "%Y-%m-%dT%H:%M:%SZ"

sync/api/projects.py

@@ -1,5 +1,6 @@
 """Project functions
 """
+
 import io
 import json
 import logging
@@ -416,6 +417,19 @@ def get_project_submission(project_id: str, submission_id: str) -> Response[dict
     return Response(result=response["result"])
 
 
+def get_submissions(project_id: str) -> Response[dict]:
+    """Get submissions for a project id
+
+    :param project_id: project ID
+    :type project_id: str
+    :return: List of Submission Configuration objects
+    :rtype: dict
+    """
+    recent_submissions = get_default_client().get_project_submissions(project_id)
+    if recent_submissions.get("result") and len(recent_submissions["result"]) > 0:
+        return Response(result=recent_submissions["result"])
+
+
 def get_latest_project_config_recommendation(
     project_id: str,
 ) -> Response[Union[AWSProjectConfiguration, AzureProjectConfiguration]]:
@@ -427,7 +441,7 @@ def get_latest_project_config_recommendation(
     :rtype: AWSProjectConfiguration or AzureProjectConfiguration
     """
     latest_recommendation = get_default_client().get_latest_project_recommendation(project_id)
-    if latest_recommendation.get("result"):
+    if latest_recommendation.get("result") and len(latest_recommendation["result"]) > 0:
         return Response(
             result=latest_recommendation["result"][0]["recommendation"]["configuration"]
         )
@@ -436,7 +450,7 @@ def get_latest_project_config_recommendation(
 def get_cluster_definition_and_recommendation(
     project_id: str, cluster_spec_str: str
 ) -> Response[dict]:
-    """Print Current Cluster Definition and Project Configuration Recommendatio.
+    """Print Current Cluster Definition and Project Configuration Recommendation.
     Throws error if no cluster recommendation found for project
 
     :param project_id: project ID

sync/cli/projects.py

@@ -7,6 +7,7 @@
     delete_project,
     get_project,
     get_projects,
+    get_submissions,
     reset_project,
     update_project,
 )
@@ -165,3 +166,47 @@ def delete(project: dict):
         click.echo(response.result)
     else:
         click.echo(str(response.error), err=True)
+
+
+@projects.command
+@click.argument("project", callback=validate_project)
+@click.option(
+    "--success-only",
+    is_flag=True,
+    default=False,
+    help="Only show the most recent successful submission",
+)
+def get_latest_submission_config(project: dict, success_only: bool = False):
+    """
+    Get the latest submission configuration for a project.
+    """
+
+    try:
+        submissions = get_submissions(project["id"]).result
+    except AttributeError as e:
+        click.echo(f"Failed to retrieve submissions. {e}", err=True)
+        return
+
+    if not submissions:
+        click.echo("No submissions found.", err=True)
+        return
+
+    if success_only:
+        latest_successful_submission = next(
+            (submission for submission in submissions if submission["state"] == "SUCCESS"), None
+        )
+
+        if not latest_successful_submission:
+            click.echo("No successful submissions found.", err=True)
+            return
+
+        submission_to_show = latest_successful_submission
+    else:
+        submission_to_show = submissions[0]
+    click.echo(
+        json.dumps(
+            submission_to_show.get("configuration", {}),
+            indent=2,
+            cls=DateTimeEncoderNaiveUTCDropMicroseconds,
+        )
+    )

sync/projects.py

@@ -0,0 +1,86 @@
+from uuid import UUID
+
+from sync.api.projects import get_project_by_app_id, get_submissions
+
+
+class SubmissionRetrievalError(Exception):
+    """Custom exception for errors retrieving submissions."""
+
+
+class NoSubmissionsFoundError(SubmissionRetrievalError):
+    """No submissions were found."""
+
+
+class NoSuccessfulSubmissionsFoundError(SubmissionRetrievalError):
+    """No successful submissions were found."""
+
+
+class InvalidUUIDError(ValueError):
+    """Raised when the provided UUID is invalid."""
+
+
+class ProjectResolutionError(Exception):
+    """Raised when there is an error resolving the project ID."""
+
+
+def validate_uuid(value: str) -> str:
+    try:
+        UUID(value)
+        return value
+    except ValueError:
+        raise InvalidUUIDError(f"Invalid UUID: {value}")
+
+
+def resolve_project_id(value: str) -> str:
+    try:
+        return validate_uuid(value)
+    except InvalidUUIDError:
+        try:
+            project_response = get_project_by_app_id(value)
+            if project_response.error:
+                raise ProjectResolutionError(
+                    f"Error resolving project ID: {project_response.error}"
+                )
+            return project_response.result.get("id")
+        except Exception as e:
+            raise ProjectResolutionError(
+                f"An error occurred while resolving the project ID: {str(e)}"
+            )
+
+
+def get_latest_submission_config(project_id: str, success_only: bool = False) -> dict:
+    """
+    Get the latest submission configuration for a project.
+
+    :param project_id: Sync project ID
+    :type project_id: str
+    :param success_only: Only show the most recent successful submission, if omitted shows the most
+        recent submission regardless of state
+    :type success_only: bool, optional
+    """
+
+    try:
+        submissions = get_submissions(project_id).result
+    except Exception as e:
+        raise SubmissionRetrievalError(
+            f"Failed to retrieve submissions for project '{project_id}'. {e}"
+        ) from e
+
+    if not submissions:
+        raise NoSubmissionsFoundError(f"No submissions found for project '{project_id}'.")
+
+    if success_only:
+        latest_successful_submission = next(
+            (submission for submission in submissions if submission["state"] == "SUCCESS"), None
+        )
+
+        if not latest_successful_submission:
+            raise NoSuccessfulSubmissionsFoundError(
+                f"No successful submissions found for project '{project_id}'."
+            )
+
+        submission_to_show = latest_successful_submission
+    else:
+        submission_to_show = submissions[0]
+    config = submission_to_show.get("configuration", {})
+    return config

tests/test_projects.py

@@ -0,0 +1,96 @@
+from unittest.mock import patch
+
+import pytest
+from pytest import raises
+
+from sync.projects import NoSuccessfulSubmissionsFoundError, get_latest_submission_config
+
+mock_base_configuration = {
+    "node_type_id": "Standard_DS3_v2",
+    "driver_node_type_id": "Standard_DS3_v2",
+    "custom_tags": {
+        "Vendor": "Databricks",
+        "Creator": "",
+        "ClusterName": "job-951738965563208-run-13053153698407-Job_cluster",
+        "ClusterId": "0315-161419-efa397m4",
+        "JobId": "951738965563208",
+    },
+    "num_workers": 2,
+    "spark_conf": {"spark.databricks.isv.product": "sync-gradient"},
+    "spark_version": "13.3.x-scala2.12",
+    "runtime_engine": "STANDARD",
+    "azure_attributes": {
+        "availability": "ON_DEMAND_AZURE",
+        "first_on_demand": 1,
+        "spot_bid_max_price": -1.0,
+    },
+}
+mock_success_configuration = {"state": "SUCCESS", **mock_base_configuration}
+mock_failed_configuration = {"state": "FAILED", **mock_base_configuration}
+
+mock_base_submission = {
+    "created_at": "2024-03-15T16:27:12Z",
+    "updated_at": "2024-03-15T16:27:21Z",
+    "id": "b7ee7431-3188-44ac-95fc-6f2068de39b6",
+    "project_id": "ff10d866-ebf5-46ff-84dc-b5430a5eea45",
+}
+mock_success_submission = {
+    **mock_base_submission,
+    "state": "SUCCESS",
+    "configuration": mock_success_configuration,
+}
+mock_failed_submission = {
+    **mock_base_submission,
+    "state": "FAILED",
+    "configuration": mock_failed_configuration,
+}
+
+
+@pytest.fixture
+def mock_mixed_submissions():
+    yield {"result": [mock_failed_submission, mock_success_submission]}
+
+
+@pytest.fixture
+def mock_no_success_submissions():
+    yield {"result": [mock_failed_submission]}
+
+
+@pytest.fixture
+def mock_get_submissions_empty():
+    """Mocks get_submissions function to return an empty list."""
+    yield {"result": []}
+
+
+@patch("sync.clients.sync.SyncClient._send")
+def test_get_latest_submission_config_success(mock_send, mock_mixed_submissions):
+    """Test get_latest_submission_config returns the latest successful submission configuration."""
+    mock_send.return_value = mock_mixed_submissions
+    with patch("sync.api.projects.get_submissions", lambda x: mock_mixed_submissions):
+        project_id = "ff10d866-ebf5-46ff-84dc-b5430a5eea45"
+        config = get_latest_submission_config(project_id, success_only=True)
+        assert "node_type_id" in config
+        assert "Vendor" in config.get("custom_tags", "")
+        assert config["state"] == "SUCCESS"
+
+
+@patch("sync.clients.sync.SyncClient._send")
+def test_get_latest_submission_config_failure_ok(mock_send, mock_mixed_submissions):
+    mock_send.return_value = mock_mixed_submissions
+    with patch("sync.api.projects.get_submissions", lambda x: mock_mixed_submissions):
+        project_id = "ff10d866-ebf5-46ff-84dc-b5430a5eea45"
+        config = get_latest_submission_config(project_id, success_only=False)
+        assert "node_type_id" in config
+        assert "Vendor" in config.get("custom_tags", "")
+        assert config["state"] == "FAILED"
+
+
+@patch("sync.clients.sync.SyncClient._send")
+def test_get_latest_submission_raises_no_successful_submissions_found_error(
+    mock_send, mock_no_success_submissions
+):
+    mock_send.return_value = mock_no_success_submissions
+    with patch("sync.api.projects.get_submissions", lambda x: mock_no_success_submissions):
+        project_id = "ff10d866-ebf5-46ff-84dc-b5430a5eea45"
+        with raises(NoSuccessfulSubmissionsFoundError):
+            get_latest_submission_config(project_id, success_only=True)