hatchet-sdk 1.2.6__py3-none-any.whl → 1.3.1__py3-none-any.whl

hatchet_sdk/clients/rest/configuration.py

@@ -15,6 +15,7 @@
 import copy
 import http.client as httplib
 import logging
+import multiprocessing
 import sys
 from logging import FileHandler
 from typing import Any, ClassVar, Dict, List, Literal, Optional, TypedDict, Union
@@ -304,9 +305,12 @@ class Configuration:
            Set this to the SNI value expected by the server.
         """
 
-        self.connection_pool_maxsize = 100
-        """This value is passed to the aiohttp to limit simultaneous connections.
-           Default values is 100, None means no-limit.
+        self.connection_pool_maxsize = multiprocessing.cpu_count() * 5
+        """urllib3 connection pool's maximum number of connections saved
+           per pool. urllib3 uses 1 connection as default value, but this is
+           not the best value when you are making a lot of possibly parallel
+           requests to the same host, which is often the case here.
+           cpu_count * 5 is used as default value to increase performance.
         """
 
         self.proxy: Optional[str] = None

hatchet_sdk/clients/rest/models/tenant_step_run_queue_metrics.py

@@ -19,7 +19,7 @@ import pprint
 import re  # noqa: F401
 from typing import Any, ClassVar, Dict, List, Optional, Set
 
-from pydantic import BaseModel, ConfigDict, StrictInt
+from pydantic import BaseModel, ConfigDict
 from typing_extensions import Self
 
 
@@ -28,7 +28,7 @@ class TenantStepRunQueueMetrics(BaseModel):
     TenantStepRunQueueMetrics
     """  # noqa: E501
 
-    queues: Optional[Dict[str, StrictInt]] = None
+    queues: Optional[Dict[str, Any]] = None
     __properties: ClassVar[List[str]] = ["queues"]
 
     model_config = ConfigDict(

hatchet_sdk/clients/rest/models/v1_task_summary.py

@@ -34,6 +34,9 @@ class V1TaskSummary(BaseModel):
     """  # noqa: E501
 
     metadata: APIResourceMeta
+    action_id: Optional[StrictStr] = Field(
+        default=None, description="The action ID of the task.", alias="actionId"
+    )
     additional_metadata: Optional[Dict[str, Any]] = Field(
         default=None,
         description="Additional metadata for the task run.",
@@ -104,6 +107,7 @@ class V1TaskSummary(BaseModel):
     )
     __properties: ClassVar[List[str]] = [
         "metadata",
+        "actionId",
         "additionalMetadata",
         "children",
         "createdAt",
@@ -193,6 +197,7 @@ class V1TaskSummary(BaseModel):
                     if obj.get("metadata") is not None
                     else None
                 ),
+                "actionId": obj.get("actionId"),
                 "additionalMetadata": obj.get("additionalMetadata"),
                 "children": (
                     [V1TaskSummary.from_dict(_item) for _item in obj["children"]]

hatchet_sdk/clients/rest/models/v1_workflow_run.py

@@ -78,6 +78,9 @@ class V1WorkflowRun(BaseModel):
         description="The timestamp the task run was created.",
         alias="createdAt",
     )
+    parent_task_external_id: Optional[
+        Annotated[str, Field(min_length=36, strict=True, max_length=36)]
+    ] = Field(default=None, alias="parentTaskExternalId")
     __properties: ClassVar[List[str]] = [
         "metadata",
         "status",
@@ -93,6 +96,7 @@ class V1WorkflowRun(BaseModel):
         "workflowVersionId",
         "input",
         "createdAt",
+        "parentTaskExternalId",
     ]
 
     model_config = ConfigDict(
@@ -166,6 +170,7 @@ class V1WorkflowRun(BaseModel):
                 "workflowVersionId": obj.get("workflowVersionId"),
                 "input": obj.get("input"),
                 "createdAt": obj.get("createdAt"),
+                "parentTaskExternalId": obj.get("parentTaskExternalId"),
             }
         )
         return _obj

hatchet_sdk/clients/rest/rest.py

@@ -16,16 +16,23 @@ import io
 import json
 import re
 import ssl
-from typing import Optional, Union
 
-import aiohttp
-import aiohttp_retry
+import urllib3
 
 from hatchet_sdk.clients.rest.exceptions import ApiException, ApiValueError
 
-RESTResponseType = aiohttp.ClientResponse
+SUPPORTED_SOCKS_PROXIES = {"socks5", "socks5h", "socks4", "socks4a"}
+RESTResponseType = urllib3.HTTPResponse
 
-ALLOW_RETRY_METHODS = frozenset({"DELETE", "GET", "HEAD", "OPTIONS", "PUT", "TRACE"})
+
+def is_socks_proxy_url(url):
+    if url is None:
+        return False
+    split_section = url.split("://")
+    if len(split_section) < 2:
+        return False
+    else:
+        return split_section[0].lower() in SUPPORTED_SOCKS_PROXIES
 
 
 class RESTResponse(io.IOBase):
@@ -36,13 +43,13 @@ class RESTResponse(io.IOBase):
         self.reason = resp.reason
         self.data = None
 
-    async def read(self):
+    def read(self):
         if self.data is None:
-            self.data = await self.response.read()
+            self.data = self.response.data
         return self.data
 
     def getheaders(self):
-        """Returns a CIMultiDictProxy of the response headers."""
+        """Returns a dictionary of the response headers."""
         return self.response.headers
 
     def getheader(self, name, default=None):
@@ -53,38 +60,57 @@ class RESTResponse(io.IOBase):
 class RESTClientObject:
 
     def __init__(self, configuration) -> None:
+        # urllib3.PoolManager will pass all kw parameters to connectionpool
+        # https://github.com/shazow/urllib3/blob/f9409436f83aeb79fbaf090181cd81b784f1b8ce/urllib3/poolmanager.py#L75  # noqa: E501
+        # https://github.com/shazow/urllib3/blob/f9409436f83aeb79fbaf090181cd81b784f1b8ce/urllib3/connectionpool.py#L680  # noqa: E501
+        # Custom SSL certificates and client certificates: http://urllib3.readthedocs.io/en/latest/advanced-usage.html  # noqa: E501
+
+        # cert_reqs
+        if configuration.verify_ssl:
+            cert_reqs = ssl.CERT_REQUIRED
+        else:
+            cert_reqs = ssl.CERT_NONE
+
+        pool_args = {
+            "cert_reqs": cert_reqs,
+            "ca_certs": configuration.ssl_ca_cert,
+            "cert_file": configuration.cert_file,
+            "key_file": configuration.key_file,
+            "ca_cert_data": configuration.ca_cert_data,
+        }
+        if configuration.assert_hostname is not None:
+            pool_args["assert_hostname"] = configuration.assert_hostname
+
+        if configuration.retries is not None:
+            pool_args["retries"] = configuration.retries
+
+        if configuration.tls_server_name:
+            pool_args["server_hostname"] = configuration.tls_server_name
+
+        if configuration.socket_options is not None:
+            pool_args["socket_options"] = configuration.socket_options
+
+        if configuration.connection_pool_maxsize is not None:
+            pool_args["maxsize"] = configuration.connection_pool_maxsize
 
-        # maxsize is number of requests to host that are allowed in parallel
-        self.maxsize = configuration.connection_pool_maxsize
-
-        self.ssl_context = ssl.create_default_context(
-            cafile=configuration.ssl_ca_cert,
-            cadata=configuration.ca_cert_data,
-        )
-        if configuration.cert_file:
-            self.ssl_context.load_cert_chain(
-                configuration.cert_file, keyfile=configuration.key_file
-            )
-
-        if not configuration.verify_ssl:
-            self.ssl_context.check_hostname = False
-            self.ssl_context.verify_mode = ssl.CERT_NONE
-
-        self.proxy = configuration.proxy
-        self.proxy_headers = configuration.proxy_headers
-
-        self.retries = configuration.retries
+        # https pool manager
+        self.pool_manager: urllib3.PoolManager
 
-        self.pool_manager: Optional[aiohttp.ClientSession] = None
-        self.retry_client: Optional[aiohttp_retry.RetryClient] = None
+        if configuration.proxy:
+            if is_socks_proxy_url(configuration.proxy):
+                from urllib3.contrib.socks import SOCKSProxyManager
 
-    async def close(self) -> None:
-        if self.pool_manager:
-            await self.pool_manager.close()
-        if self.retry_client is not None:
-            await self.retry_client.close()
+                pool_args["proxy_url"] = configuration.proxy
+                pool_args["headers"] = configuration.proxy_headers
+                self.pool_manager = SOCKSProxyManager(**pool_args)
+            else:
+                pool_args["proxy_url"] = configuration.proxy
+                pool_args["proxy_headers"] = configuration.proxy_headers
+                self.pool_manager = urllib3.ProxyManager(**pool_args)
+        else:
+            self.pool_manager = urllib3.PoolManager(**pool_args)
 
-    async def request(
+    def request(
         self,
         method,
         url,
@@ -93,7 +119,7 @@ class RESTClientObject:
         post_params=None,
         _request_timeout=None,
     ):
-        """Execute request
+        """Perform requests.
 
         :param method: http request method
         :param url: http request url
@@ -117,82 +143,105 @@ class RESTClientObject:
 
         post_params = post_params or {}
         headers = headers or {}
-        # url already contains the URL query string
-        timeout = _request_timeout or 5 * 60
-
-        if "Content-Type" not in headers:
-            headers["Content-Type"] = "application/json"
-
-        args = {"method": method, "url": url, "timeout": timeout, "headers": headers}
-
-        if self.proxy:
-            args["proxy"] = self.proxy
-        if self.proxy_headers:
-            args["proxy_headers"] = self.proxy_headers
-
-        # For `POST`, `PUT`, `PATCH`, `OPTIONS`, `DELETE`
-        if method in ["POST", "PUT", "PATCH", "OPTIONS", "DELETE"]:
-            if re.search("json", headers["Content-Type"], re.IGNORECASE):
-                if body is not None:
-                    body = json.dumps(body)
-                args["data"] = body
-            elif headers["Content-Type"] == "application/x-www-form-urlencoded":
-                args["data"] = aiohttp.FormData(post_params)
-            elif headers["Content-Type"] == "multipart/form-data":
-                # must del headers['Content-Type'], or the correct
-                # Content-Type which generated by aiohttp
-                del headers["Content-Type"]
-                data = aiohttp.FormData()
-                for param in post_params:
-                    k, v = param
-                    if isinstance(v, tuple) and len(v) == 3:
-                        data.add_field(k, value=v[1], filename=v[0], content_type=v[2])
-                    else:
-                        # Ensures that dict objects are serialized
-                        if isinstance(v, dict):
-                            v = json.dumps(v)
-                        elif isinstance(v, int):
-                            v = str(v)
-                        data.add_field(k, v)
-                args["data"] = data
-
-            # Pass a `bytes` or `str` parameter directly in the body to support
-            # other content types than Json when `body` argument is provided
-            # in serialized form
-            elif isinstance(body, str) or isinstance(body, bytes):
-                args["data"] = body
-            else:
-                # Cannot generate the request from given parameters
-                msg = """Cannot prepare a request message for provided
-                         arguments. Please check that your arguments match
-                         declared content type."""
-                raise ApiException(status=0, reason=msg)
 
-        pool_manager: Union[aiohttp.ClientSession, aiohttp_retry.RetryClient]
-
-        # https pool manager
-        if self.pool_manager is None:
-            self.pool_manager = aiohttp.ClientSession(
-                connector=aiohttp.TCPConnector(
-                    limit=self.maxsize, ssl=self.ssl_context
-                ),
-                trust_env=True,
-            )
-        pool_manager = self.pool_manager
-
-        if self.retries is not None and method in ALLOW_RETRY_METHODS:
-            if self.retry_client is None:
-                self.retry_client = aiohttp_retry.RetryClient(
-                    client_session=self.pool_manager,
-                    retry_options=aiohttp_retry.ExponentialRetry(
-                        attempts=self.retries,
-                        factor=2.0,
-                        start_timeout=0.1,
-                        max_timeout=120.0,
-                    ),
+        timeout = None
+        if _request_timeout:
+            if isinstance(_request_timeout, (int, float)):
+                timeout = urllib3.Timeout(total=_request_timeout)
+            elif isinstance(_request_timeout, tuple) and len(_request_timeout) == 2:
+                timeout = urllib3.Timeout(
+                    connect=_request_timeout[0], read=_request_timeout[1]
                 )
-            pool_manager = self.retry_client
 
-        r = await pool_manager.request(**args)
+        try:
+            # For `POST`, `PUT`, `PATCH`, `OPTIONS`, `DELETE`
+            if method in ["POST", "PUT", "PATCH", "OPTIONS", "DELETE"]:
+
+                # no content type provided or payload is json
+                content_type = headers.get("Content-Type")
+                if not content_type or re.search("json", content_type, re.IGNORECASE):
+                    request_body = None
+                    if body is not None:
+                        request_body = json.dumps(body)
+                    r = self.pool_manager.request(
+                        method,
+                        url,
+                        body=request_body,
+                        timeout=timeout,
+                        headers=headers,
+                        preload_content=False,
+                    )
+                elif content_type == "application/x-www-form-urlencoded":
+                    r = self.pool_manager.request(
+                        method,
+                        url,
+                        fields=post_params,
+                        encode_multipart=False,
+                        timeout=timeout,
+                        headers=headers,
+                        preload_content=False,
+                    )
+                elif content_type == "multipart/form-data":
+                    # must del headers['Content-Type'], or the correct
+                    # Content-Type which generated by urllib3 will be
+                    # overwritten.
+                    del headers["Content-Type"]
+                    # Ensures that dict objects are serialized
+                    post_params = [
+                        (a, json.dumps(b)) if isinstance(b, dict) else (a, b)
+                        for a, b in post_params
+                    ]
+                    r = self.pool_manager.request(
+                        method,
+                        url,
+                        fields=post_params,
+                        encode_multipart=True,
+                        timeout=timeout,
+                        headers=headers,
+                        preload_content=False,
+                    )
+                # Pass a `string` parameter directly in the body to support
+                # other content types than JSON when `body` argument is
+                # provided in serialized form.
+                elif isinstance(body, str) or isinstance(body, bytes):
+                    r = self.pool_manager.request(
+                        method,
+                        url,
+                        body=body,
+                        timeout=timeout,
+                        headers=headers,
+                        preload_content=False,
+                    )
+                elif headers["Content-Type"].startswith("text/") and isinstance(
+                    body, bool
+                ):
+                    request_body = "true" if body else "false"
+                    r = self.pool_manager.request(
+                        method,
+                        url,
+                        body=request_body,
+                        preload_content=False,
+                        timeout=timeout,
+                        headers=headers,
+                    )
+                else:
+                    # Cannot generate the request from given parameters
+                    msg = """Cannot prepare a request message for provided
+                             arguments. Please check that your arguments match
+                             declared content type."""
+                    raise ApiException(status=0, reason=msg)
+            # For `GET`, `HEAD`
+            else:
+                r = self.pool_manager.request(
+                    method,
+                    url,
+                    fields={},
+                    timeout=timeout,
+                    headers=headers,
+                    preload_content=False,
+                )
+        except urllib3.exceptions.SSLError as e:
+            msg = "\n".join([type(e).__name__, str(e)])
+            raise ApiException(status=0, reason=msg)
 
         return RESTResponse(r)

hatchet_sdk/clients/v1/api_client.py

@@ -1,4 +1,4 @@
-from typing import AsyncContextManager, ParamSpec, TypeVar
+from typing import ParamSpec, TypeVar
 
 from hatchet_sdk.clients.rest.api_client import ApiClient
 from hatchet_sdk.clients.rest.configuration import Configuration
@@ -40,5 +40,5 @@ class BaseRestClient:
 
         self.api_config.datetime_format = "%Y-%m-%dT%H:%M:%S.%fZ"
 
-    def client(self) -> AsyncContextManager[ApiClient]:
+    def client(self) -> ApiClient:
         return ApiClient(self.api_config)

hatchet_sdk/context/context.py

@@ -1,22 +1,20 @@
-import inspect
 import json
 import traceback
 from concurrent.futures import Future, ThreadPoolExecutor
 from datetime import timedelta
 from typing import TYPE_CHECKING, Any, cast
-
-from pydantic import BaseModel
+from warnings import warn
 
 from hatchet_sdk.clients.admin import AdminClient
 from hatchet_sdk.clients.dispatcher.dispatcher import (  # type: ignore[attr-defined]
     Action,
     DispatcherClient,
 )
-from hatchet_sdk.clients.durable_event_listener import (
+from hatchet_sdk.clients.events import EventClient
+from hatchet_sdk.clients.listeners.durable_event_listener import (
     DurableEventListener,
     RegisterDurableEventRequest,
 )
-from hatchet_sdk.clients.events import EventClient
 from hatchet_sdk.context.worker_context import WorkerContext
 from hatchet_sdk.features.runs import RunsClient
 from hatchet_sdk.logger import logger
@@ -29,21 +27,6 @@ if TYPE_CHECKING:
     from hatchet_sdk.runnables.types import R, TWorkflowInput
 
 
-DEFAULT_WORKFLOW_POLLING_INTERVAL = 5  # Seconds
-
-
-def get_caller_file_path() -> str:
-    caller_frame = inspect.stack()[2]
-
-    return caller_frame.filename
-
-
-class StepRunError(BaseModel):
-    step_id: str
-    step_run_action_name: str
-    error: str
-
-
 class Context:
     def __init__(
         self,
@@ -85,7 +68,7 @@ class Context:
     def trigger_data(self) -> JSONSerializableMapping:
         return self.data.triggers
 
-    def task_output(self, task: "Task[TWorkflowInput, R]") -> "R":
+    def _task_output(self, task: "Task[TWorkflowInput, R]") -> "R":
         from hatchet_sdk.runnables.types import R
 
         if self.was_skipped(task):
@@ -116,7 +99,25 @@ class Context:
 
         return parent_step_data
 
+    def task_output(self, task: "Task[TWorkflowInput, R]") -> "R":
+        from hatchet_sdk.runnables.types import R
+
+        ## If the task is async, we need to wrap its output in a coroutine
+        ## so that the type checker behaves right
+        async def _aio_output() -> "R":
+            return self._task_output(task)
+
+        if task.is_async_function:
+            return cast(R, _aio_output())
+
+        return self._task_output(task)
+
     def aio_task_output(self, task: "Task[TWorkflowInput, R]") -> "R":
+        warn(
+            "`aio_task_output` is deprecated. Use `task_output` instead.",
+            DeprecationWarning,
+        )
+
         if task.is_async_function:
             return self.task_output(task)