feat: accept str|enum region/runtime, omit when unset for org default (#57)

* feat: accept str|enum region/runtime, omit when unset for org default

Widen the region/runtime parameters on WherobotsRunOperator,
WherobotsSqlOperator, WherobotsSqlHook, and WherobotsRestAPIHook.create_run to
`Optional[Union[str, Region/Runtime]]`. Strings are passed to the API untouched
(BYOC regions work without a release); enums are normalized to their value.

When region/runtime are omitted they are dropped from the request (region from
the /runs query string, runtime from the run payload, and both from the SQL
connect() call) so the API applies the organization's configured default.
`warn_for_default_region` no longer injects DEFAULT_REGION — it normalizes and
returns None when unset.

Adds tests for create_run (region omitted / string passthrough) and
warn_for_default_region.

Note: the SQL path's "omit -> org default" behavior also requires the
wherobots-python-dbapi release that makes connect() region/runtime optional;
bump the dependency constraint to that version at release time. The run-operator
path depends only on the studio-backend API change.

* fix: don't normalize region to str in warn_for_default_region; pin runtime in prod smoke tests

CI integration tests (which hit prod with the released dbapi) failed:

- test_sql crashed with "'str' object has no attribute 'value'": warn_for_default_region
  returned a normalized string, and the released enum-only connect() calls
  region.value on it. Normalization belongs at the request boundary
  (create_run / connect), so warn now returns the value unchanged
  (Optional[Union[str, Region]]); the enum keeps working with the current dbapi
  and the new dbapi normalizes internally.
- test_run got a 422 ("runtime field required") because the operator now omits
  runtime by default and the (not-yet-deployed) prod API still requires it. The
  prod smoke tests are meant to verify a run submits successfully, not the
  org-default mechanism (covered by unit tests), so they now pin
  runtime=Runtime.TINY for deterministic, transition-robust behavior.

* fix: handle TIMED_OUT run status from the API

The prod API returns a TIMED_OUT run status that the provider's RunStatus enum
didn't include, so polling a timed-out run raised a Pydantic ValidationError
(surfaced by the timeout smoke test once it reaches the real run path). Add
TIMED_OUT to RunStatus, treat it as a timeout in Run.is_timeout, and raise the
"failed due to timeout" RuntimeError for it in WherobotsRunOperator.execute.
Adds a TIMED_OUT case to the execute-state unit test.

Author: ClayMav

README.md

@@ -76,13 +76,17 @@ operator = WherobotsRunOperator(
 
 The arguments for the `WherobotsRunOperator` constructor:
 
-* `region: Region`: The Wherobots region where runs are hosted.
-  The values available can be found in `wherobots.db.region.Region`.
+* `region: str | Region`: The Wherobots region where runs are hosted. Accepts a
+  `Region` enum value (see `wherobots.db.region.Region`) or a raw string (a BYOC
+  region is passed through as-is). Optional — when omitted, your organization's
+  configured default region is used. Only set it to override that default.
 * `name: str`: The name of the run. If not specified, a default name will be
   generated.
-* `runtime: Runtime`: The runtime dictates the size and amount of resources
-  powering the run. The default value is `Runtime.TINY`; see available values
-  [here](https://github.com/wherobots/wherobots-python-dbapi/blob/main/wherobots/db/runtime.py).
+* `runtime: str | Runtime`: The runtime dictates the size and amount of resources
+  powering the run. Accepts a `Runtime` enum value (see available values
+  [here](https://github.com/wherobots/wherobots-python-dbapi/blob/main/wherobots/db/runtime.py))
+  or a raw string. Optional — when omitted, your organization's configured
+  default runtime is used.
 * `version: str`: The WherobotsDB version to use. Defaults to `latest`.
 * `poll_logs: bool`: If `True`, the operator will poll and `Logger.info()` the run logs
   until the run finishes. If `False`, the operator will not poll the logs, only track

airflow_providers_wherobots/hooks/rest_api.py

@@ -4,7 +4,7 @@
 
 import platform
 from functools import cached_property
-from typing import Any, Optional, Dict
+from typing import Any, Optional, Dict, Union
 
 import requests
 from importlib import metadata
@@ -112,12 +112,18 @@ def get_run(self, run_id: str) -> Run:
         resp_json = self._api_call("GET", f"/runs/{run_id}").json()
         return Run.model_validate(resp_json)
 
-    def create_run(self, payload: Dict[str, Any], region: Region) -> Run:
+    def create_run(
+        self, payload: Dict[str, Any], region: Optional[Union[str, Region]] = None
+    ) -> Run:
+        # Normalize enum -> value, pass strings through, and omit region
+        # entirely when unset so the API applies the org default.
+        region_value = region.value if isinstance(region, Region) else region
+        params = {"region": region_value} if region_value else {}
         resp_json = self._api_call(
             "POST",
             "/runs",
             payload=payload,
-            params={"region": region.value},
+            params=params,
         ).json()
         return Run.model_validate(resp_json)
 

airflow_providers_wherobots/hooks/sql.py

@@ -2,14 +2,13 @@
 Hook for Wherobots' Spatial SQL API interface.
 """
 
-from typing import Optional
+from typing import Optional, Union
 
 from airflow.providers.common.sql.hooks.sql import DbApiHook
 from wherobots.db import Connection as WDBConnection, connect
 from wherobots.db.constants import (
     DEFAULT_SESSION_WAIT_TIMEOUT_SECONDS,
     DEFAULT_READ_TIMEOUT_SECONDS,
-    DEFAULT_RUNTIME,
     DEFAULT_SESSION_TYPE,
 )
 from wherobots.db.region import Region
@@ -24,9 +23,9 @@ class WherobotsSqlHook(DbApiHook):  # type: ignore[misc]
 
     def __init__(  # type: ignore[no-untyped-def]
         self,
-        region: Optional[Region] = None,
+        region: Optional[Union[str, Region]] = None,
         wherobots_conn_id: str = DEFAULT_CONN_ID,
-        runtime: Runtime = DEFAULT_RUNTIME,
+        runtime: Optional[Union[str, Runtime]] = None,
         version: Optional[str] = None,
         session_wait_timeout: int = DEFAULT_SESSION_WAIT_TIMEOUT_SECONDS,
         read_timeout: int = DEFAULT_READ_TIMEOUT_SECONDS,
@@ -49,7 +48,7 @@ def __init__(  # type: ignore[no-untyped-def]
 
     def _create_or_get_sql_session(
         self,
-        runtime: Runtime = DEFAULT_RUNTIME,
+        runtime: Optional[Union[str, Runtime]] = None,
     ) -> WDBConnection:
         return connect(
             host=self._conn.host,

airflow_providers_wherobots/operators/__init__.py

@@ -1,16 +1,27 @@
 import logging
-from typing import Optional
+from typing import Optional, Union
 
 from wherobots.db import Region
-from wherobots.db.constants import DEFAULT_REGION
 
 logger = logging.getLogger(__name__)
 
 
-def warn_for_default_region(region: Optional[Region]) -> Region:
+def warn_for_default_region(
+    region: Optional[Union[str, Region]],
+) -> Optional[Union[str, Region]]:
+    """Resolve the region argument for a Wherobots operator.
+
+    The value is returned unchanged — a ``Region`` enum or a raw string (e.g. a
+    BYOC region). Normalization (enum -> value) happens at the request boundary
+    (``WherobotsRestAPIHook.create_run`` / ``wherobots.db.connect``), so the enum
+    keeps working with older ``wherobots-python-dbapi`` releases. When no region
+    is provided, returns ``None`` so the API applies the organization's
+    configured default region — only set ``region`` to override that default.
+    """
     if not region:
-        logger.warning(""""Parameter region was not specified, it will be required by Wherobots API in the near future, please specify it in the operator.
-        If you don't know your Wherobots Compute Region, please contact Wherobots support.""")
-        logger.warning(f"Using default region: {DEFAULT_REGION.value}")
-        region = DEFAULT_REGION
+        logger.info(
+            "No region specified; the Wherobots API will use your organization's "
+            "configured default region. Pass `region` to override it."
+        )
+        return None
     return region

airflow_providers_wherobots/operators/run.py

@@ -5,7 +5,7 @@
 import time
 from enum import auto
 from time import sleep
-from typing import Optional, Sequence, Any, Dict
+from typing import Optional, Sequence, Any, Dict, Union
 
 from airflow.models import BaseOperator
 from strenum import StrEnum
@@ -20,7 +20,6 @@
 )
 
 from wherobots.db import Runtime, Region
-from wherobots.db.constants import DEFAULT_RUNTIME
 
 
 class XComKey(StrEnum):
@@ -39,9 +38,9 @@ class WherobotsRunOperator(BaseOperator):
 
     def __init__(
         self,
-        region: Optional[Region] = None,
+        region: Optional[Union[str, Region]] = None,
         name: Optional[str] = None,
-        runtime: Runtime = DEFAULT_RUNTIME,
+        runtime: Optional[Union[str, Runtime]] = None,
         version: Optional[str] = None,
         run_python: Optional[Dict[str, Any]] = None,
         run_jar: Optional[Dict[str, Any]] = None,
@@ -56,10 +55,14 @@ def __init__(
         super().__init__(**kwargs)
         # If the user specifies the name, we will use it and rely on the server to validate the name
         self.run_payload: Dict[str, Any] = {
-            "runtime": runtime.value,
             "name": name or self.default_run_name,
             "timeoutSeconds": timeout_seconds,
         }
+        # Only include runtime when set; otherwise the API uses the org default.
+        if runtime is not None:
+            self.run_payload["runtime"] = (
+                runtime.value if isinstance(runtime, Runtime) else runtime
+            )
         self.region = region
         if version is not None:
             self.run_payload["version"] = version
@@ -189,10 +192,13 @@ def execute(self, context) -> Any:
                 run = self._wait_run_simple(rest_api_hook, run)
             # loop end, means run is in terminal state
             self._log_run_status(run)
+            # The API may report a timeout either as an explicit TIMED_OUT
+            # status or as FAILED with a timeout event.
+            if run.status == RunStatus.TIMED_OUT or (
+                run.status == RunStatus.FAILED and run.is_timeout
+            ):
+                raise RuntimeError(f"Run {run.ext_id} failed due to timeout")
             if run.status == RunStatus.FAILED:
-                # check events, see if the run is timeout
-                if run.is_timeout:
-                    raise RuntimeError(f"Run {run.ext_id} failed due to timeout")
                 raise RuntimeError(f"Run {run.ext_id} failed, please check the logs")
             if run.status == RunStatus.CANCELLED:
                 raise RuntimeError(f"Run {run.ext_id} was cancelled by user")

airflow_providers_wherobots/operators/sql.py

@@ -4,14 +4,13 @@
 
 from __future__ import annotations
 
-from typing import Sequence, Optional
+from typing import Sequence, Optional, Union
 
 from airflow.providers.common.sql.hooks.sql import DbApiHook
 from airflow.providers.common.sql.operators.sql import SQLExecuteQueryOperator
 from wherobots.db.constants import (
     DEFAULT_SESSION_WAIT_TIMEOUT_SECONDS,
     DEFAULT_READ_TIMEOUT_SECONDS,
-    DEFAULT_RUNTIME,
     DEFAULT_SESSION_TYPE,
 )
 from wherobots.db import Cursor as WDbCursor
@@ -43,9 +42,9 @@ class WherobotsSqlOperator(SQLExecuteQueryOperator):  # type: ignore[misc]
     def __init__(  # type: ignore[no-untyped-def]
         self,
         *,
-        region: Optional[Region] = None,
+        region: Optional[Union[str, Region]] = None,
         wherobots_conn_id: str = DEFAULT_CONN_ID,
-        runtime: Runtime = DEFAULT_RUNTIME,
+        runtime: Optional[Union[str, Runtime]] = None,
         version: Optional[str] = None,
         session_wait_timeout: int = DEFAULT_SESSION_WAIT_TIMEOUT_SECONDS,
         read_timeout: int = DEFAULT_READ_TIMEOUT_SECONDS,

airflow_providers_wherobots/wherobots/models.py

@@ -19,6 +19,7 @@ class RunStatus(StrEnum):
     FAILED = auto()
     COMPLETED = auto()
     CANCELLED = auto()
+    TIMED_OUT = auto()
 
     def is_active(self) -> bool:
         return self in [self.PENDING, self.RUNNING]
@@ -49,6 +50,8 @@ class Run(WherobotsModel):
 
     @property
     def is_timeout(self) -> bool:
+        if self.status == RunStatus.TIMED_OUT:
+            return True
         if not self.kube_app or not self.kube_app.events:
             return False
         return any(

tests/integration_tests/operators/test_run.py

@@ -9,7 +9,7 @@
 from airflow import DAG
 from airflow.models import Connection
 from airflow.utils.state import TaskInstanceState
-from wherobots.db import Region
+from wherobots.db import Region, Runtime
 
 from airflow_providers_wherobots.operators.run import WherobotsRunOperator
 from tests.unit_tests.operators.test_run import build_ti
@@ -25,6 +25,9 @@
 def test_prod_run_success(prod_conn: Connection, dag: DAG) -> None:
     operator = WherobotsRunOperator(
         region=Region.AWS_US_WEST_2,
+        # Pin runtime so this prod smoke test is deterministic; the
+        # omit-runtime -> org-default behavior is covered by unit tests.
+        runtime=Runtime.TINY,
         wherobots_conn_id=prod_conn.conn_id,
         task_id="test_run_smoke",
         name="airflow_operator_test_run_{{ ts_nodash }}",
@@ -44,6 +47,7 @@ def test_prod_run_success(prod_conn: Connection, dag: DAG) -> None:
 def test_prod_run_timeout(prod_conn: Connection, dag: DAG) -> None:
     operator = WherobotsRunOperator(
         region=Region.AWS_US_WEST_2,
+        runtime=Runtime.TINY,
         wherobots_conn_id=prod_conn.conn_id,
         task_id="test_run_smoke",
         name="airflow_operator_test_run_{{ ts_nodash }}",

tests/integration_tests/operators/test_sql.py

@@ -8,7 +8,7 @@
 import pytest
 from airflow import DAG
 from airflow.models import Connection
-from wherobots.db import Region
+from wherobots.db import Region, Runtime
 
 from airflow_providers_wherobots.operators.sql import WherobotsSqlOperator
 
@@ -23,6 +23,9 @@
 def test_prod_run_success(prod_conn: Connection, dag: DAG) -> None:
     operator = WherobotsSqlOperator(
         region=Region.AWS_US_WEST_2,
+        # Pin runtime so this prod smoke test is deterministic; the
+        # omit-runtime -> org-default behavior is covered by unit tests.
+        runtime=Runtime.TINY,
         task_id=TEST_TASK_ID,
         sql="SELECT pickup_datetime FROM wherobots_pro_data.nyc_taxi.yellow_2009_2010 LIMIT 10",
         wherobots_conn_id=prod_conn.conn_id,

tests/unit_tests/hooks/test_rest_api.py

@@ -129,6 +129,38 @@ def test_create_run(self, test_default_conn) -> None:
         with WherobotsRestAPIHook() as hook:
             hook.create_run(payload=create_payload, region=Region.AWS_US_WEST_2)
 
+    @responses.activate
+    def test_create_run_omits_region_when_none(self, test_default_conn) -> None:
+        """When region is None, no region query param is sent (API uses org default)."""
+        test_run: Run = helpers.run_factory.build()
+        url = f"https://{test_default_conn.host}/runs"
+        create_payload = {"name": test_run.name, "timeoutSeconds": 5000}
+        responses.add(
+            responses.POST,
+            url,
+            json=test_run.model_dump(mode="json"),
+            match=[matchers.query_string_matcher("")],
+            status=HTTPStatus.OK,
+        )
+        with WherobotsRestAPIHook() as hook:
+            hook.create_run(payload=create_payload, region=None)
+
+    @responses.activate
+    def test_create_run_passes_string_region(self, test_default_conn) -> None:
+        """A raw region string (e.g. a BYOC region) is passed through as-is."""
+        test_run: Run = helpers.run_factory.build()
+        url = f"https://{test_default_conn.host}/runs"
+        create_payload = {"name": test_run.name, "timeoutSeconds": 5000}
+        responses.add(
+            responses.POST,
+            url,
+            json=test_run.model_dump(mode="json"),
+            match=[matchers.query_param_matcher({"region": "byoc-acme-us-east-1"})],
+            status=HTTPStatus.OK,
+        )
+        with WherobotsRestAPIHook() as hook:
+            hook.create_run(payload=create_payload, region="byoc-acme-us-east-1")
+
     @responses.activate
     def test_get_run_logs(self, test_default_conn) -> None:
         """