Feature: adds connection_timeout argument to SpeckleClient (#337)

* Add connection_timeout argument to SpeckleClient

* Reformat code with black

* Set default timeout to 10s

* Make connection retries configurable

pyproject.toml

@@ -36,7 +36,7 @@ pylint = "^2.14.4"
 mypy = "^0.982"
 pre-commit = "^2.20.0"
 commitizen = "^2.38.0"
-ruff = "^0.0.187"
+ruff = "^0.4.4"
 types-deprecated = "^1.2.9"
 types-ujson = "^5.6.0.0"
 types-requests = "^2.28.11.5"

src/speckle_automate/automation_context.py

@@ -1,4 +1,5 @@
 """This module provides an abstraction layer above the Speckle Automate runtime."""
+
 import time
 from dataclasses import dataclass, field
 from pathlib import Path
@@ -17,6 +18,7 @@ from speckle_automate.schema import (
 )
 from specklepy.api import operations
 from specklepy.api.client import SpeckleClient
+from specklepy.core.api.models import Branch
 from specklepy.logging.exceptions import SpeckleException
 from specklepy.objects.base import Base
 from specklepy.transports.memory import MemoryTransport
@@ -94,8 +96,10 @@ class AutomationContext:
 
     def receive_version(self) -> Base:
         """Receive the Speckle project version that triggered this automation run."""
+        # TODO: this is a quick hack to keep implementation consistency. Move to proper receive many versions
+        version_id = self.automation_run_data.triggers[0].payload.version_id
         commit = self.speckle_client.commit.get(
-            self.automation_run_data.project_id, self.automation_run_data.version_id
+            self.automation_run_data.project_id, version_id
         )
         if not commit.referencedObject:
             raise ValueError("The commit has no referencedObject, cannot receive it.")
@@ -104,7 +108,7 @@ class AutomationContext:
         )
         print(
             f"It took {self.elapsed():.2f} seconds to receive",
-            f" the speckle version {self.automation_run_data.version_id}",
+            f" the speckle version {version_id}",
         )
         return base
 
@@ -119,19 +123,27 @@ class AutomationContext:
             version_message (str): The message for the new version.
         """
 
-        if model_name == self.automation_run_data.branch_name:
-            raise ValueError(
-                f"The target model: {model_name} cannot match the model"
-                f" that triggered this automation:"
-                f" {self.automation_run_data.model_id} /"
-                f" {self.automation_run_data.branch_name}"
-            )
-
         branch = self.speckle_client.branch.get(
             self.automation_run_data.project_id, model_name, 1
         )
-        # we just check if it exists
-        if (not branch) or isinstance(branch, SpeckleException):
+        if isinstance(branch, Branch):
+            if not branch.id:
+                raise ValueError("Cannot use the branch without its id")
+            matching_trigger = [
+                t
+                for t in self.automation_run_data.triggers
+                if t.payload.model_id == branch.id
+            ]
+            if matching_trigger:
+                raise ValueError(
+                    f"The target model: {model_name} cannot match the model"
+                    f" that triggered this automation:"
+                    f" {matching_trigger[0].payload.model_id}"
+                )
+            model_id = branch.id
+
+        else:
+            # we just check if it exists
             branch_create = self.speckle_client.branch.create(
                 self.automation_run_data.project_id,
                 model_name,
@@ -139,8 +151,6 @@ class AutomationContext:
             if isinstance(branch_create, Exception):
                 raise branch_create
             model_id = branch_create
-        else:
-            model_id = branch.id
 
         root_object_id = operations.send(
             root_object,
@@ -174,7 +184,8 @@ class AutomationContext:
     ) -> None:
         link_resources = (
             [
-                f"{self.automation_run_data.model_id}@{self.automation_run_data.version_id}"
+                f"{t.payload.model_id}@{t.payload.version_id}"
+                for t in self.automation_run_data.triggers
             ]
             if include_source_model_version
             else []
@@ -194,47 +205,26 @@ class AutomationContext:
         """Report the current run status to the project of this automation."""
         query = gql(
             """
-            mutation ReportFunctionRunStatus(
-                $automationId: String!,
-                $automationRevisionId: String!,
-                $automationRunId: String!,
-                $versionId: String!,
-                $functionId: String!,
-                $functionName: String!,
-                $functionLogo: String,
-                $runStatus: AutomationRunStatus!
-                $elapsed: Float!
-                $contextView: String
-                $resultVersionIds: [String!]!
+            mutation AutomateFunctionRunStatusReport(
+                $functionRunId: String!
+                $status: AutomateRunStatus!
                 $statusMessage: String
-                $objectResults: JSONObject
+                $results: JSONObject
+                $contextView: String
             ){
-                automationMutations {
-                    functionRunStatusReport(input: {
-                        automationId: $automationId
-                        automationRevisionId: $automationRevisionId
-                        automationRunId: $automationRunId
-                        versionId: $versionId
-                        functionRuns: [
-                        {
-                            functionId: $functionId
-                            functionName: $functionName
-                            functionLogo: $functionLogo
-                            status: $runStatus,
-                            contextView: $contextView,
-                            elapsed: $elapsed,
-                            resultVersionIds: $resultVersionIds,
-                            statusMessage: $statusMessage
-                            results: $objectResults
-                        }]
-                   })
-                }
+                automateFunctionRunStatusReport(input: {
+                    functionRunId: $functionRunId
+                    status: $status
+                    statusMessage: $statusMessage
+                    contextView: $contextView
+                    results: $results
+                })
             }
-            """
+        """
         )
         if self.run_status in [AutomationStatus.SUCCEEDED, AutomationStatus.FAILED]:
             object_results = {
-                "version": "1.0.0",
+                "version": 1,
                 "values": {
                     "objectResults": self._automation_result.model_dump(by_alias=True)[
                         "objectResults"
@@ -246,19 +236,11 @@ class AutomationContext:
             object_results = None
 
         params = {
-            "automationId": self.automation_run_data.automation_id,
-            "automationRevisionId": self.automation_run_data.automation_revision_id,
-            "automationRunId": self.automation_run_data.automation_run_id,
-            "versionId": self.automation_run_data.version_id,
-            "functionId": self.automation_run_data.function_id,
-            "functionName": self.automation_run_data.function_name,
-            "functionLogo": self.automation_run_data.function_logo,
-            "runStatus": self.run_status.value,
+            "functionRunId": self.automation_run_data.function_run_id,
+            "status": self.run_status.value,
             "statusMessage": self._automation_result.status_message,
+            "results": object_results,
             "contextView": self._automation_result.result_view,
-            "elapsed": self.elapsed(),
-            "resultVersionIds": self._automation_result.result_versions,
-            "objectResults": object_results,
         }
         print(f"Reporting run status with content: {params}")
         self.speckle_client.httpclient.execute(query, params)

src/speckle_automate/schema.py

@@ -1,6 +1,7 @@
 """"""
+
 from enum import Enum
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Literal, Optional
 
 from pydantic import BaseModel, ConfigDict, Field
 from stringcase import camelcase
@@ -12,22 +13,30 @@ class AutomateBase(BaseModel):
     model_config = ConfigDict(alias_generator=camelcase, populate_by_name=True)
 
 
+class VersionCreationTriggerPayload(AutomateBase):
+    """Represents the version creation trigger payload."""
+
+    model_id: str
+    version_id: str
+
+
+class VersionCreationTrigger(AutomateBase):
+    """Represents a single version creation trigger for the automation run."""
+
+    trigger_type: Literal["versionCreation"]
+    payload: VersionCreationTriggerPayload
+
+
 class AutomationRunData(BaseModel):
     """Values of the project / model that triggered the run of this function."""
 
     project_id: str
-    model_id: str
-    branch_name: str
-    version_id: str
     speckle_server_url: str
-
     automation_id: str
-    automation_revision_id: str
     automation_run_id: str
+    function_run_id: str
 
-    function_id: str
-    function_name: str
-    function_logo: Optional[str]
+    triggers: List[VersionCreationTrigger]
 
     model_config = ConfigDict(
         alias_generator=camelcase, populate_by_name=True, protected_namespaces=()

src/specklepy/api/resources/branch.py

@@ -1,8 +1,9 @@
-from typing import Optional
+from typing import Optional, Union
 
 from specklepy.api.models import Branch
 from specklepy.core.api.resources.branch import Resource as CoreResource
 from specklepy.logging import metrics
+from specklepy.logging.exceptions import SpeckleException
 
 
 class Resource(CoreResource):
@@ -31,7 +32,9 @@ class Resource(CoreResource):
         metrics.track(metrics.SDK, self.account, {"name": "Branch Create"})
         return super().create(stream_id, name, description)
 
-    def get(self, stream_id: str, name: str, commits_limit: int = 10):
+    def get(
+        self, stream_id: str, name: str, commits_limit: int = 10
+    ) -> Union[Branch, None, SpeckleException]:
         """Get a branch by name from a stream
 
         Arguments:

src/specklepy/core/api/client.py

@@ -64,6 +64,8 @@ class SpeckleClient:
         host: str = DEFAULT_HOST,
         use_ssl: bool = USE_SSL,
         verify_certificate: bool = True,
+        connection_retries: int = 3,
+        connection_timeout: int = 10,
     ) -> None:
         ws_protocol = "ws"
         http_protocol = "http"
@@ -80,10 +82,15 @@ class SpeckleClient:
         self.ws_url = f"{ws_protocol}://{host}/graphql"
         self.account = Account()
         self.verify_certificate = verify_certificate
+        self.connection_retries = connection_retries
+        self.connection_timeout = connection_timeout
 
         self.httpclient = Client(
             transport=RequestsHTTPTransport(
-                url=self.graphql, verify=self.verify_certificate, retries=3
+                url=self.graphql,
+                verify=self.verify_certificate,
+                retries=self.connection_retries,
+                timeout=self.connection_timeout,
             )
         )
         self.wsclient = None

tests/integration/speckle_automate/test_automation_context.py

@@ -189,6 +189,9 @@ def automate_function(
         automation_context.mark_run_success("No forbidden types found.")
 
 
+@pytest.mark.skip(
+    "currently the function run cannot be integration tested with the server"
+)
 def test_function_run(automation_context: AutomationContext) -> None:
     """Run an integration test for the automate function."""
     automation_context = run_function(
@@ -215,6 +218,9 @@ def test_file_path():
     os.remove(path)
 
 
+@pytest.mark.skip(
+    "currently the function run cannot be integration tested with the server"
+)
 def test_file_uploads(
     automation_run_data: AutomationRunData, speckle_token: str, test_file_path: Path
 ):
@@ -230,6 +236,9 @@ def test_file_uploads(
     assert len(automation_context._automation_result.blobs) == 1
 
 
+@pytest.mark.skip(
+    "currently the function run cannot be integration tested with the server"
+)
 def test_create_version_in_project_raises_error_for_same_model(
     automation_context: AutomationContext,
 ) -> None:
@@ -239,6 +248,9 @@ def test_create_version_in_project_raises_error_for_same_model(
         )
 
 
+@pytest.mark.skip(
+    "currently the function run cannot be integration tested with the server"
+)
 def test_create_version_in_project(
     automation_context: AutomationContext,
 ) -> None:
@@ -252,6 +264,9 @@ def test_create_version_in_project(
     assert version_id is not None
 
 
+@pytest.mark.skip(
+    "currently the function run cannot be integration tested with the server"
+)
 def test_set_context_view(automation_context: AutomationContext) -> None:
     automation_context.set_context_view()