feat: initial pass at hatchet function

mrkaye97 · mrkaye97 · commit 9dd8eec4704f · 2025-01-31T15:48:40.000-05:00
diff --git a/examples/simple/worker.py b/examples/simple/worker.py
@@ -8,20 +8,17 @@
 hatchet = Hatchet(debug=True)
 
 
-class MyWorkflow(BaseWorkflowImpl):
-    @hatchet.step(timeout="11s", retries=3)
-    def step1(self, context: Context) -> dict[str, str]:
-        print("executed step1")
-        return {
-            "step1": "step1",
-        }
+@hatchet.function(timeout="11s")
+def step1(context: Context) -> dict[str, str]:
+    print("executed step1")
+    return {
+        "step1": "step1",
+    }
 
 
 def main() -> None:
-    wf = MyWorkflow()
-
     worker = hatchet.worker("test-worker", max_runs=1)
-    worker.register_workflow(wf)
+    worker.register_function(step1)
     worker.start()
 
 
diff --git a/hatchet_sdk/loader.py b/hatchet_sdk/loader.py
@@ -1,168 +1,134 @@
 import json
-import os
 from logging import Logger, getLogger
-from typing import cast
-
-from pydantic import BaseModel, ConfigDict, ValidationInfo, field_validator
+from typing import Any
+
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    Field,
+    ValidationInfo,
+    field_validator,
+    model_validator,
+)
+from pydantic_settings import BaseSettings, SettingsConfigDict
 
 from hatchet_sdk.token import get_addresses_from_jwt, get_tenant_id_from_jwt
 
 
-class ClientTLSConfig(BaseModel):
-    tls_strategy: str
-    cert_file: str | None
-    key_file: str | None
-    ca_file: str | None
-    server_name: str
+def create_settings_config(env_prefix: str) -> SettingsConfigDict:
+    return SettingsConfigDict(
+        env_prefix=env_prefix,
+        env_file=(".env", ".env.hatchet", ".env.dev", ".env.local"),
+        extra="ignore",
+    )
 
 
-def _load_tls_config(host_port: str | None = None) -> ClientTLSConfig:
-    server_name = os.getenv("HATCHET_CLIENT_TLS_SERVER_NAME")
+class ClientTLSConfig(BaseSettings):
+    model_config = create_settings_config(
+        env_prefix="HATCHET_CLIENT_TLS_",
+    )
 
-    if not server_name and host_port:
-        server_name = host_port.split(":")[0]
+    strategy: str = "tls"
+    cert_file: str | None = None
+    key_file: str | None = None
+    root_ca_file: str | None = None
+    server_name: str = ""
 
-    if not server_name:
-        server_name = "localhost"
 
-    return ClientTLSConfig(
-        tls_strategy=os.getenv("HATCHET_CLIENT_TLS_STRATEGY", "tls"),
-        cert_file=os.getenv("HATCHET_CLIENT_TLS_CERT_FILE"),
-        key_file=os.getenv("HATCHET_CLIENT_TLS_KEY_FILE"),
-        ca_file=os.getenv("HATCHET_CLIENT_TLS_ROOT_CA_FILE"),
-        server_name=server_name,
+class OTELConfig(BaseSettings):
+    model_config = create_settings_config(
+        env_prefix="HATCHET_CLIENT_OTEL_",
     )
 
+    service_name: str | None = None
+    exporter_otlp_endpoint: str | None = None
+    exporter_otlp_headers: str | None = None
+    exporter_otlp_protocol: str | None = None
 
-def parse_listener_timeout(timeout: str | None) -> int | None:
-    if timeout is None:
-        return None
 
-    return int(timeout)
+class HealthcheckConfig(BaseSettings):
+    model_config = create_settings_config(
+        env_prefix="HATCHET_CLIENT_WORKER_HEALTHCHECK_",
+    )
+
+    port: int = 8001
+    enabled: bool = False
 
 
 DEFAULT_HOST_PORT = "localhost:7070"
 
 
-class ClientConfig(BaseModel):
-    model_config = ConfigDict(arbitrary_types_allowed=True, validate_default=True)
+class ClientConfig(BaseSettings):
+    model_config = create_settings_config(
+        env_prefix="HATCHET_CLIENT_",
+    )
 
-    token: str = os.getenv("HATCHET_CLIENT_TOKEN", "")
+    token: str = ""
     logger: Logger = getLogger()
-    tenant_id: str = os.getenv("HATCHET_CLIENT_TENANT_ID", "")
-
-    ## IMPORTANT: Order matters here. The validators run in the order that the
-    ## fields are defined in the model. So, we need to make sure that the
-    ## host_port is set before we try to load the tls_config and server_url
-    host_port: str = os.getenv("HATCHET_CLIENT_HOST_PORT", DEFAULT_HOST_PORT)
-    tls_config: ClientTLSConfig = _load_tls_config()
 
+    tenant_id: str = ""
+    host_port: str = DEFAULT_HOST_PORT
     server_url: str = "https://app.dev.hatchet-tools.com"
-    namespace: str = os.getenv("HATCHET_CLIENT_NAMESPACE", "")
-    listener_v2_timeout: int | None = parse_listener_timeout(
-        os.getenv("HATCHET_CLIENT_LISTENER_V2_TIMEOUT")
-    )
-    grpc_max_recv_message_length: int = int(
-        os.getenv("HATCHET_CLIENT_GRPC_MAX_RECV_MESSAGE_LENGTH", 4 * 1024 * 1024)
-    )  # 4MB
-    grpc_max_send_message_length: int = int(
-        os.getenv("HATCHET_CLIENT_GRPC_MAX_SEND_MESSAGE_LENGTH", 4 * 1024 * 1024)
-    )  # 4MB
-    otel_exporter_oltp_endpoint: str | None = os.getenv(
-        "HATCHET_CLIENT_OTEL_EXPORTER_OTLP_ENDPOINT"
-    )
-    otel_service_name: str | None = os.getenv("HATCHET_CLIENT_OTEL_SERVICE_NAME")
-    otel_exporter_oltp_headers: str | None = os.getenv(
-        "HATCHET_CLIENT_OTEL_EXPORTER_OTLP_HEADERS"
-    )
-    otel_exporter_oltp_protocol: str | None = os.getenv(
-        "HATCHET_CLIENT_OTEL_EXPORTER_OTLP_PROTOCOL"
-    )
-    worker_healthcheck_port: int = int(
-        os.getenv("HATCHET_CLIENT_WORKER_HEALTHCHECK_PORT", 8001)
-    )
-    worker_healthcheck_enabled: bool = (
-        os.getenv("HATCHET_CLIENT_WORKER_HEALTHCHECK_ENABLED", "False") == "True"
-    )
-
-    @field_validator("token", mode="after")
-    @classmethod
-    def validate_token(cls, token: str) -> str:
-        if not token:
-            raise ValueError("Token must be set")
-
-        return token
-
-    @field_validator("namespace", mode="after")
-    @classmethod
-    def validate_namespace(cls, namespace: str) -> str:
-        if not namespace:
-            return ""
-
-        if not namespace.endswith("_"):
-            namespace = f"{namespace}_"
+    namespace: str = ""
 
-        return namespace.lower()
+    tls_config: ClientTLSConfig = Field(default_factory=lambda: ClientTLSConfig())
+    otel: OTELConfig = Field(default_factory=lambda: OTELConfig())
+    healthcheck: HealthcheckConfig = Field(default_factory=lambda: HealthcheckConfig())
 
-    @field_validator("tenant_id", mode="after")
-    @classmethod
-    def validate_tenant_id(cls, tenant_id: str, info: ValidationInfo) -> str:
-        token = cast(str | None, info.data.get("token"))
+    listener_v2_timeout: int | None = None
+    grpc_max_recv_message_length: int = Field(
+        default=4 * 1024 * 1024, description="4MB default"
+    )
+    grpc_max_send_message_length: int = Field(
+        default=4 * 1024 * 1024, description="4MB default"
+    )
 
-        if not tenant_id:
-            if not token:
-                raise ValueError("Either the token or tenant_id must be set")
+    worker_preset_labels: dict[str, str] = Field(default_factory=dict)
 
-            return get_tenant_id_from_jwt(token)
+    @model_validator(mode="after")
+    def validate_token_and_tenant(self) -> "ClientConfig":
+        if not self.token:
+            raise ValueError("Token must be set")
 
-        return tenant_id
+        if not self.tenant_id:
+            self.tenant_id = get_tenant_id_from_jwt(self.token)
 
-    @field_validator("host_port", mode="after")
-    @classmethod
-    def validate_host_port(cls, host_port: str, info: ValidationInfo) -> str:
-        if host_port and host_port != DEFAULT_HOST_PORT:
-            return host_port
+        return self
 
-        token = cast(str, info.data.get("token"))
+    @model_validator(mode="after")
+    def validate_addresses(self) -> "ClientConfig":
+        if self.host_port == DEFAULT_HOST_PORT:
+            server_url, grpc_broadcast_address = get_addresses_from_jwt(self.token)
+            self.host_port = grpc_broadcast_address
+            self.server_url = server_url
 
-        if not token:
-            raise ValueError("Token must be set")
+        if not self.tls_config.server_name:
+            self.tls_config.server_name = self.host_port.split(":")[0]
 
-        _, grpc_broadcast_address = get_addresses_from_jwt(token)
+        if not self.tls_config.server_name:
+            self.tls_config.server_name = "localhost"
 
-        return grpc_broadcast_address
+        return self
 
-    @field_validator("server_url", mode="after")
+    @field_validator("listener_v2_timeout")
     @classmethod
-    def validate_server_url(cls, server_url: str, info: ValidationInfo) -> str:
-        ## IMPORTANT: Order matters here. The validators run in the order that the
-        ## fields are defined in the model. So, we need to make sure that the
-        ## host_port is set before we try to load the server_url
-        host_port = cast(str, info.data.get("host_port"))
+    def validate_listener_timeout(cls, value: int | None | str) -> int | None:
+        if value is None:
+            return None
 
-        if host_port and host_port != DEFAULT_HOST_PORT:
-            return host_port
+        if isinstance(value, int):
+            return value
 
-        token = cast(str, info.data.get("token"))
+        return int(value)
 
-        if not token:
-            raise ValueError("Token must be set")
-
-        _server_url, _ = get_addresses_from_jwt(token)
-
-        return _server_url
-
-    @field_validator("tls_config", mode="after")
+    @field_validator("namespace")
     @classmethod
-    def validate_tls_config(
-        cls, tls_config: ClientTLSConfig, info: ValidationInfo
-    ) -> ClientTLSConfig:
-        ## IMPORTANT: Order matters here. This validator runs in the order
-        ## that the fields are defined in the model. So, we need to make sure
-        ## that the host_port is set before we try to load the tls_config
-        host_port = cast(str, info.data.get("host_port"))
-
-        return _load_tls_config(host_port)
+    def validate_namespace(cls, namespace: str) -> str:
+        if not namespace:
+            return ""
+        if not namespace.endswith("_"):
+            namespace = f"{namespace}_"
+        return namespace.lower()
 
     def __hash__(self) -> int:
         return hash(json.dumps(self.model_dump(), default=str))
diff --git a/hatchet_sdk/v2/hatchet.py b/hatchet_sdk/v2/hatchet.py
@@ -1,6 +1,15 @@
 import asyncio
 import logging
-from typing import TYPE_CHECKING, Any, Callable, Optional, Type, TypeVar, cast
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Optional,
+    ParamSpec,
+    Type,
+    TypeVar,
+    cast,
+)
 
 from hatchet_sdk.client import Client, new_client, new_client_raw
 from hatchet_sdk.clients.admin import AdminClient
@@ -17,6 +26,7 @@
 from hatchet_sdk.logger import logger
 from hatchet_sdk.rate_limit import RateLimit
 from hatchet_sdk.v2.workflows import (
+    BaseWorkflowImpl,
     ConcurrencyExpression,
     EmptyModel,
     Step,
@@ -30,6 +40,7 @@
 if TYPE_CHECKING:
     from hatchet_sdk.worker.worker import Worker
 
+P = ParamSpec("P")
 R = TypeVar("R")
 
 
@@ -187,6 +198,55 @@ def inner(func: Callable[[Any, Context], R]) -> Step[R]:
 
         return inner
 
+    def function(
+        self,
+        name: str = "",
+        on_events: list[str] = [],
+        on_crons: list[str] = [],
+        version: str = "",
+        timeout: str = "60m",
+        schedule_timeout: str = "5m",
+        sticky: StickyStrategy | None = None,
+        default_priority: int = 1,
+        concurrency: ConcurrencyExpression | None = None,
+        input_validator: Type[TWorkflowInput] | None = None,
+    ) -> Callable[[Callable[[Context], R]], BaseWorkflowImpl]:
+        declaration = WorkflowDeclaration[TWorkflowInput](
+            WorkflowConfig(
+                name=name,
+                on_events=on_events,
+                on_crons=on_crons,
+                version=version,
+                timeout=timeout,
+                schedule_timeout=schedule_timeout,
+                sticky=sticky,
+                default_priority=default_priority,
+                concurrency=concurrency,
+                input_validator=input_validator
+                or cast(Type[TWorkflowInput], EmptyModel),
+            ),
+            self,
+        )
+
+        def inner(func: Callable[[Context], R]) -> BaseWorkflowImpl:
+            class Workflow(BaseWorkflowImpl):
+                config = declaration.config
+
+                @self.step(
+                    name=name,
+                    timeout=timeout,
+                    retries=0,
+                    rate_limits=[],
+                    backoff_factor=None,
+                    backoff_max_seconds=None,
+                )
+                def fn(self, context: Context) -> R:
+                    return func(context)
+
+            return Workflow()
+
+        return inner
+
     def worker(
         self, name: str, max_runs: int | None = None, labels: dict[str, str | int] = {}
     ) -> "Worker":
diff --git a/hatchet_sdk/worker/worker.py b/hatchet_sdk/worker/worker.py
@@ -110,6 +110,9 @@ def register_workflow_from_opts(
             logger.error(e)
             sys.exit(1)
 
+    def register_function(self, function: "BaseWorkflowImpl") -> None:
+        self.register_workflow(function)
+
     def register_workflow(self, workflow: Union["BaseWorkflowImpl", Any]) -> None:
         namespace = self.client.config.namespace
 
